kups.core.neighborlist.common ¶

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.lhs.indices, candidates.rhs.indices], axis=-1)
    edges: Edges[Literal[2]] = Edges(
        Index(candidates.lhs.keys, indices_2d),
        jnp.expand_dims(shifts, axis=-2),
    )
    return CandidateBatch(edges=edges, is_minimum_image=is_minimum_image)

def make_batch_with_mic(
    candidates: Candidates,
    lh: Table[ParticleId, NeighborListPoints],
    rh: 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, lh, rh, systems)
    return candidates_to_batch(
        candidates,
        min_shifts,
        jnp.ones((candidates.lhs.size,), dtype=bool),
    )

def real_distance_sq(
    lh: Table[ParticleId, NeighborListPoints],
    rh: Table[ParticleId, NeighborListPoints],
    systems: Table[SystemId, NeighborListSystems],
    lh_idx: Array,
    rh_idx: Array,
    shifts: Array,
) -> Array:
    """Squared real-space distance between candidate pairs.

    Args:
        lh, rh, systems: Pipeline tables (positions in fractional coords).
        lh_idx, rh_idx: ``(n,)`` raw index arrays.
        shifts: ``(n, 3)`` fractional shifts.

    Returns:
        ``(n,)`` array of squared distances in real coordinates.
    """
    frames = systems.map_data(lambda s: s.cell.frame.materialize())
    frames = frames[lh.data.system[lh_idx]]
    deltas = lh.data.positions[lh_idx] - rh.data.positions[rh_idx] - shifts
    real_deltas = frames.to_real(deltas)
    return jnp.einsum("...d,...d->...", real_deltas, real_deltas)

def replicate_for_images(
    candidates: Candidates,
    lh: Table[ParticleId, NeighborListPoints],
    rh: Table[ParticleId, NeighborListPoints],
    systems: Table[SystemId, NeighborListSystems],
    cutoffs: Table[SystemId, Array],
    max_image_candidates: Capacity[int] | None,
) -> CandidateBatch[Literal[2]]:
    """Replicate candidates by image multiplicity, attaching shifts and is-min flag.

    For each candidate pair:
    - If ``max(cutoff[sys] / perp_axes) <= 0.5``: emit 1 copy with MIC shifts.
    - Otherwise: emit per-image copies with replicated shifts; the
      ``is_minimum_image`` flag is set per copy so ``ExclusionMask`` can keep
      non-minimum image periodic copies of excluded pairs.

    Args:
        candidates: Raw candidate pair indices.
        lh, rh, systems: Pipeline tables (fractional coords).
        cutoffs: Per-system cutoff.
        max_image_candidates: Capacity for replicated-candidates buffer.
            When ``None``, falls back to ``FixedCapacity(candidates.lhs.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.lhs.size,
            "Cutoff is larger than half the cell length, "
            "we need to generate additional images. "
            "Please provide a editable max_candidates.",
        )

    idx, image_shifts, has_been_replicated = _get_candidate_images(
        candidates, lh, systems, cutoffs_t.data, max_image_candidates
    )
    min_shifts = _minimum_image_shifts(candidates, lh, rh, systems)

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

    replicated = bind(candidates).at(idx).get()
    final_shifts = jnp.where(
        has_been_replicated[:, None], image_shifts, min_shifts[idx]
    )
    is_min = (min_shifts[idx] == final_shifts).all(axis=-1)
    return candidates_to_batch(replicated, final_shifts, is_min)