Skip to content

kups.core.neighborlist.compact

Compactors for the neighbor list pipeline.

Compactor variants:

  • ReduceCompactor — compresses surviving candidate rows via jnp.where(keep, size=k) with a capacity assertion, preserving the emitted edge degree.
  • MaskOnlyCompactor — preserves the candidate count, replacing failing entries with OOB indices and zero shifts. Used by RefineMaskNeighborList.

Graph-level output shaping, such as restoring undirected symmetry after queried_keys deduplication, lives in pipeline postprocessors.

MaskOnlyCompactor

Bases: Compactor[D]

In-place compaction: failing entries become OOB indices and zero shifts.

No size change; preserves the candidate count from the selector. Pair candidates are already in their output index space.

Source code in src/kups/core/neighborlist/compact.py 
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
@dataclass
class MaskOnlyCompactor[D: int](Compactor[D]):
    """In-place compaction: failing entries become OOB indices and zero shifts.

    No size change; preserves the candidate count from the selector. Pair
    candidates are already in their output index space.
    """

    def __call__(
        self,
        keep: Array,
        batch: CandidateBatch[D],
        ctx: PipelineContext,
    ) -> Edges[D]:
        oob = max(ctx.keys.size, ctx.edge_query_table.size)
        indices_in = batch.edges.indices.indices
        indices = where_broadcast_last(keep, indices_in, oob)
        shifts = where_broadcast_last(keep, batch.edges.shifts, 0)
        return Edges(Index(batch.edges.indices.keys, indices), shifts)

ReduceCompactor

Bases: Compactor[D]

Compact surviving candidates to a size-bounded Edges[D].

Compacts whole candidate rows, so pair neighbor lists and fixed higher-degree topology share the same implementation.

Source code in src/kups/core/neighborlist/compact.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
@dataclass
class ReduceCompactor[D: int](Compactor[D]):
    """Compact surviving candidates to a size-bounded ``Edges[D]``.

    Compacts whole candidate rows, so pair neighbor lists and fixed higher-degree
    topology share the same implementation.
    """

    avg_edges: Capacity[int]

    def __call__(
        self,
        keep: Array,
        batch: CandidateBatch[D],
        ctx: PipelineContext,
    ) -> Edges[D]:
        oob = max(ctx.keys.size, ctx.edge_query_table.size)
        max_edges = self.avg_edges.generate_assertion(keep.sum())
        sort_idxs = jnp.where(keep, size=max_edges.size, fill_value=keep.size)[0]
        shifts = batch.edges.shifts.at[sort_idxs].get(
            mode="fill", fill_value=0, indices_are_sorted=True
        )

        indices = batch.edges.indices.indices.at[sort_idxs].get(
            mode="fill", fill_value=oob, indices_are_sorted=True
        )
        return Edges(Index(batch.edges.indices.keys, indices), shifts)