Skip to content

kups.core.neighborlist.pipeline

Pipeline runner: selector → mask sequence → compactor.

A Pipeline[D] is the modularized form of a neighbor list. Concrete public NL classes (CellListNeighborList, DenseNearestNeighborList, etc.) build and run a Pipeline internally inside their __call__.

Pipeline

Selector → mask sequence → compactor.

Attributes:

Name Type Description
selector CandidateSelector[D]

Produces a CandidateBatch[D] (handles PBC replication).
masks tuple[Mask, ...]

Tuple of mask criteria over CandidateBatch[D]; results are conjuncted via &.
compactor Compactor[D]

Produces the final Edges[D] from the accumulated mask.
Source code in src/kups/core/neighborlist/pipeline.py 
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
@dataclass
class Pipeline[D: int]:
    """Selector → mask sequence → compactor.

    Attributes:
        selector: Produces a ``CandidateBatch[D]`` (handles PBC replication).
        masks: Tuple of mask criteria over ``CandidateBatch[D]``; results
            are conjuncted via ``&``.
        compactor: Produces the final ``Edges[D]`` from the accumulated mask.
    """

    selector: CandidateSelector[D]
    masks: tuple[Mask, ...] = field(static=True)
    compactor: Compactor[D]

    def __call__(
        self,
        lh: Table[ParticleId, NeighborListPoints],
        rh: Table[ParticleId, NeighborListPoints] | None,
        systems: Table[SystemId, NeighborListSystems],
        rh_index_remap: Index[ParticleId] | None = None,
    ) -> Edges[D]:
        ctx = _prepare(lh, rh, systems, rh_index_remap)
        batch = self.selector(ctx)
        keep = jnp.ones((batch.lh_idx.size,), dtype=bool)
        for mask in self.masks:
            keep &= mask(batch, ctx)
        return self.compactor(keep, batch, ctx)