Skip to content

kups.core.neighborlist.pipeline

Pipeline runner: selector → mask sequence → compactor → postprocessors.

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 → postprocessors.

Attributes:

Name Type Description
selector CandidateSelector[D]

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

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

Produces compacted Edges[D] from the accumulated mask.
postprocessors tuple[Postprocessor[D], ...]

Edge transforms applied sequentially after compaction.
Source code in src/kups/core/neighborlist/pipeline.py 
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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
@dataclass
class Pipeline[D: int]:
    """Selector → mask sequence → compactor → postprocessors.

    Attributes:
        selector: Produces a ``CandidateBatch[D]`` (handles PBC replication).
        masks: Tuple of mask criteria over ``CandidateBatch[D]``; results
            are conjuncted via ``&``.
        compactor: Produces compacted ``Edges[D]`` from the accumulated mask.
        postprocessors: Edge transforms applied sequentially after compaction.
    """

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

    @overload
    def __call__(
        self,
        keys: Table[ParticleId, NeighborListPoints],
        systems: Table[SystemId, NeighborListSystems],
        *,
        queries: Table[ParticleId, NeighborListPoints],
    ) -> Edges[D]: ...
    @overload
    def __call__(
        self,
        keys: Table[ParticleId, NeighborListPoints],
        systems: Table[SystemId, NeighborListSystems],
        *,
        queried_keys: Index[ParticleId] | None = None,
    ) -> Edges[D]: ...
    def __call__(
        self,
        keys: Table[ParticleId, NeighborListPoints],
        systems: Table[SystemId, NeighborListSystems],
        *,
        queries: Table[ParticleId, NeighborListPoints] | None = None,
        queried_keys: Index[ParticleId] | None = None,
    ) -> Edges[D]:
        ctx = _prepare(keys, queries, systems, queried_keys)
        batch = self.selector(ctx)
        keep = jnp.ones((len(batch.edges),), dtype=bool)
        for mask in self.masks:
            keep &= mask(batch, ctx)
        edges = self.compactor(keep, batch, ctx)
        for postprocessor in self.postprocessors:
            edges = postprocessor(edges, ctx)
        return edges