kups.core.neighborlist.common ¶

def candidate_image_counts(cells: Cell[AnyPeriodicity], cutoffs: Array) -> Array:
    """Return per-system, per-axis periodic-image window widths.

    Each periodic axis needs ``ceil(2 * cutoff / perpendicular_length)``
    consecutive integer shifts -- the tight count of lattice planes a sphere of
    radius ``cutoff`` can reach along that axis under the strict ``< cutoff``
    distance mask (the max number of integers in an open interval of that
    width). At ``ratio <= 0.5`` this collapses to one image, recovering the
    minimum-image convention. Open axes and non-finite ratios use one image.
    ``perpendicular_lengths`` is the correct per-axis measure for arbitrary
    skew (it equals ``1 / |column of inverse_vectors|``).
    """
    ratio = cutoffs[..., None] / cells.perpendicular_lengths
    images = jnp.maximum(jnp.ceil(2 * ratio), 1).astype(int)
    images = jnp.where(jnp.isfinite(ratio), images, 1)
    return jnp.where(jnp.array(cells.periodic), images, 1)

def candidates_to_batch(
    candidates: Candidates,
    shifts: Array,
    is_minimum_image: Array,
) -> CandidateBatch[Literal[2]]:
    """Pack ``(candidates, flat shifts, is_min)`` into a ``CandidateBatch[2]``."""
    indices_2d = jnp.stack(
        [candidates.key_idx.indices, candidates.query_idx.indices], axis=-1
    )
    edges: Edges[Literal[2]] = Edges(
        Index(candidates.key_idx.keys, indices_2d),
        jnp.expand_dims(shifts, axis=-2),
    )
    return CandidateBatch(
        edges=edges,
        is_minimum_image=is_minimum_image,
        query_keys=candidates.query_idx.keys,
    )

def lift_query_candidates(candidates: Candidates, ctx: PipelineContext) -> Candidates:
    """Convert query-local self-update candidates to ``ctx.keys`` positions."""
    if ctx.queried_keys is None:
        return candidates
    oob = ctx.keys.size
    query_idx = ctx.queried_keys.at[candidates.query_idx.indices].get(
        mode="fill", fill_value=oob
    )
    return Candidates(
        key_idx=candidates.key_idx, query_idx=Index(ctx.keys.keys, query_idx)
    )

def make_batch_with_mic(
    candidates: Candidates,
    keys: Table[ParticleId, NeighborListPoints],
    queries: Table[ParticleId, NeighborListPoints],
    systems: Table[SystemId, NeighborListSystems],
) -> CandidateBatch[Literal[2]]:
    """Pack raw candidates with minimum-image shifts; ``is_minimum_image=all-True``."""
    min_shifts = _minimum_image_shifts(candidates, keys, queries, systems)
    return candidates_to_batch(
        candidates,
        min_shifts,
        jnp.ones((candidates.key_idx.size,), dtype=bool),
    )

def real_distance_sq(
    key_positions: Array,
    query_positions: Array,
    frames: MaterializedFrame,
    shifts: Array,
) -> Array:
    """Squared real-space distance between already-broadcast candidate pairs.

    Args:
        key_positions: Fractional left endpoint positions, shape ``(n, 3)``.
        query_positions: Fractional right endpoint positions, shape ``(n, 3)``.
        frames: Materialized cell frames broadcast to the candidate key system.
        shifts: ``(n, 3)`` fractional shifts.

    Returns:
        ``(n,)`` array of squared distances in real coordinates.
    """
    deltas = key_positions - query_positions - shifts
    real_deltas = frames.to_real(deltas)
    return jnp.einsum("...d,...d->...", real_deltas, real_deltas)

def replicate_for_images(
    candidates: Candidates,
    keys: Table[ParticleId, NeighborListPoints],
    queries: Table[ParticleId, NeighborListPoints],
    systems: Table[SystemId, NeighborListSystems],
    cutoffs: Table[SystemId, Array],
    max_image_candidates: Capacity[int] | None,
) -> CandidateBatch[Literal[2]]:
    """Replicate candidates across their periodic-image windows.

    For each candidate pair:
    - If ``cutoff[sys] / perp_axes <= 0.5`` on every axis: emit 1 copy with MIC
      shifts (the minimum image is the only image in range).
    - Otherwise: emit the per-axis window of integer shifts anchored at the pair's
      separation (see ``_get_candidate_images``); ``is_minimum_image`` flags the
      closest copy per pair (via real-distance argmin) so ``ExclusionMask`` keeps
      non-minimum image periodic copies of excluded pairs. Over-emitted copies are
      pruned downstream by ``DistanceCutoffMask``.

    Args:
        candidates: Raw candidate pair indices.
        keys, queries, systems: Pipeline tables (fractional coords).
        cutoffs: Per-system cutoff.
        max_image_candidates: Capacity for replicated-candidates buffer.
            When ``None``, falls back to ``FixedCapacity(candidates.key_idx.size)``
            with an error message — pass an editable capacity if image
            replication is expected.

    Returns:
        ``CandidateBatch`` with shifts populated and ``is_minimum_image`` set.
    """
    cutoffs_t = Table.broadcast_to(cutoffs, systems)
    if max_image_candidates is None:
        max_image_candidates = FixedCapacity(
            candidates.key_idx.size,
            "Cutoff is larger than half the cell length, "
            "we need to generate additional images. "
            "Please provide a editable max_candidates.",
        )

    idx, offsets = _get_candidate_images(
        candidates, keys, queries, systems, cutoffs_t.data, max_image_candidates
    )

    if idx.size == candidates.key_idx.size:
        # No replication needed — MIC shifts cover everything.
        min_shifts = _minimum_image_shifts(candidates, keys, queries, systems)
        return candidates_to_batch(
            candidates, min_shifts, jnp.ones((candidates.key_idx.size,), dtype=bool)
        )

    replicated = bind(candidates).at(idx).get()
    is_min = _minimum_image_mask(
        replicated, offsets, idx, keys, queries, systems, candidates.key_idx.size
    )
    return candidates_to_batch(replicated, offsets, is_min)