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 
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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
@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]:
        batch, keep, ctx = select_candidates(
            keys,
            systems,
            self.selector,
            self.masks,
            queries=queries,
            queried_keys=queried_keys,
        )
        edges = self.compactor(keep, batch, ctx)
        for postprocessor in self.postprocessors:
            edges = postprocessor(edges, ctx)
        return edges

build_radius_graph(neighborlist, keys, systems, *, queries, queried_keys)

Share cutoff, group masks and update symmetry across radius selectors.

Source code in src/kups/core/neighborlist/pipeline.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
def build_radius_graph(
    neighborlist: SelectableNeighborList[Literal[2]],
    keys: Table[ParticleId, NeighborListPoints],
    systems: Table[SystemId, NeighborListSystems],
    *,
    queries: Table[ParticleId, NeighborListPoints] | None,
    queried_keys: Index[ParticleId] | None,
) -> Edges[Literal[2]]:
    """Share cutoff, group masks and update symmetry across radius selectors."""
    assert queries is None or queried_keys is None, (
        "Neighbor-list calls cannot combine queries with queried_keys. "
        "Use queried_keys for self-graph updates, or queries for bipartite queries."
    )
    query_size = (
        queried_keys.size
        if queried_keys is not None
        else (queries.size if queries is not None else keys.size)
    )
    pipeline = Pipeline[Literal[2]](
        selector=neighborlist.selector(query_size, systems),
        masks=(
            InBoundsMask(),
            InclusionMatchMask(),
            QueriedKeysDedupMask(),
            DistanceCutoffMask(neighborlist.cutoffs),
            ExclusionMask(),
        ),
        compactor=ReduceCompactor(neighborlist.avg_edges.multiply(query_size)),
        postprocessors=(MirrorPairEdges(),),
    )
    if queries is not None:
        return pipeline(keys, systems, queries=queries)
    return pipeline(keys, systems, queried_keys=queried_keys)

select_candidates(keys, systems, selector, masks, *, queries=None, queried_keys=None)

Prepare and mask candidates, returning the batch, keep mask and context.

Both graph pipelines and pair evaluators use this selection phase; the latter retain image flags when compacting the result.

Source code in src/kups/core/neighborlist/pipeline.py 
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
def select_candidates[D: int](
    keys: Table[ParticleId, NeighborListPoints],
    systems: Table[SystemId, NeighborListSystems],
    selector: CandidateSelector[D],
    masks: tuple[Mask[D], ...],
    *,
    queries: Table[ParticleId, NeighborListPoints] | None = None,
    queried_keys: Index[ParticleId] | None = None,
) -> tuple[CandidateBatch[D], Array, PipelineContext]:
    """Prepare and mask candidates, returning the batch, keep mask and context.

    Both graph pipelines and pair evaluators use this selection phase; the
    latter retain image flags when compacting the result.
    """
    ctx = _prepare(keys, queries, systems, queried_keys)
    batch = selector(ctx)
    keep = jnp.ones((len(batch.edges),), dtype=bool)
    for mask in masks:
        keep &= mask(batch, ctx)
    return batch, keep, ctx