Skip to content

kups.core.neighborlist.parameters

Capacity-hint parameters shared by neighbor list implementations.

UniversalNeighborlistParameters is the concrete dataclass every application state holds and threads into from_state for the neighbor list classes. The estimate classmethod derives conservative power-of-two capacities from system geometry so callers don't have to guess.

UniversalNeighborlistParameters

Concrete parameter dataclass satisfying IsUniversalNeighborlistParams.

Holds the capacity hints needed by every neighbor list implementation. Use the estimate() classmethod to compute reasonable initial values from system geometry rather than guessing manually.

Attributes:

Name Type Description
avg_edges int

Average number of edges per particle (for edge capacity).
avg_candidates int

Average number of candidate pairs per particle.
avg_image_candidates int

Average number of candidate pairs per particle after periodic-image replication (equals avg_candidates when every cutoff stays within the minimum-image regime).
cells int

Maximum number of spatial hash cells across all systems.
Source code in src/kups/core/neighborlist/parameters.py 
 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
 99
100
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
134
@dataclass
class UniversalNeighborlistParameters:
    """Concrete parameter dataclass satisfying ``IsUniversalNeighborlistParams``.

    Holds the capacity hints needed by every neighbor list implementation.
    Use the ``estimate()`` classmethod to compute reasonable initial values
    from system geometry rather than guessing manually.

    Attributes:
        avg_edges: Average number of edges per particle (for edge capacity).
        avg_candidates: Average number of candidate pairs per particle.
        avg_image_candidates: Average number of candidate pairs per particle after
            periodic-image replication (equals ``avg_candidates`` when every cutoff
            stays within the minimum-image regime).
        cells: Maximum number of spatial hash cells across all systems.
    """

    avg_edges: int = field(static=True)
    avg_candidates: int = field(static=True)
    avg_image_candidates: int = field(static=True)
    cells: int = field(static=True)

    @classmethod
    @no_jax_tracing
    def estimate(
        cls,
        particles_per_system: Table[SystemId, Array],
        systems: Table[SystemId, NeighborListSystems],
        cutoffs: Table[SystemId, Array],
        *,
        base: float = 2,
        multiplier: float = 1.0,
    ) -> UniversalNeighborlistParameters:
        """Estimate parameters for all neighbor list types from system geometry.

        Computes conservative initial capacities based on particle density
        and cutoff radii. The estimates are rounded up to the next power of
        ``base`` to amortize future resizing.

        Args:
            particles_per_system: Number of particles per system.
            systems: System data with cell information.
            cutoffs: Cutoff distance per system.
            base: Base for power-of rounding (default 2).
            multiplier: Safety factor applied to the estimate (default 1.0).

        Returns:
            A ``UniversalNeighborlistParameters`` instance with estimated values.
        """

        def _next_power(total: float | Array) -> int:
            return int(next_higher_power(jnp.array(total * multiplier), base=base))

        sys = Table.join(systems, particles_per_system, cutoffs)
        total_candidates = total_image_candidates = total_edges = max_cells = 0
        with no_post_init():
            for _, (s, n_p, c) in sys:
                n_bins = num_cells(s, c).prod()
                candidates = min(n_p / n_bins * (3**3), n_p)
                # A cutoff reaching past perp/2 replicates each candidate once per
                # periodic image (product of per-axis image counts). Summing per
                # system keeps the estimate tight for heterogeneous cutoffs instead
                # of assuming every system replicates at the maximum rate.
                images = candidate_image_counts(s.cell, c).prod()
                total_candidates += candidates
                total_image_candidates += _next_power(candidates) * images
                total_edges += _estimate_avg_num_edges(
                    n_p, s.cell.volume, c, base, multiplier
                )
                max_cells = max(n_bins, max_cells)

        return UniversalNeighborlistParameters(
            avg_edges=int(total_edges // sys.size),
            avg_candidates=_next_power(total_candidates / sys.size),
            avg_image_candidates=_next_power(total_image_candidates / sys.size),
            cells=int(max_cells),
        )

estimate(particles_per_system, systems, cutoffs, *, base=2, multiplier=1.0) classmethod

Estimate parameters for all neighbor list types from system geometry.

Computes conservative initial capacities based on particle density and cutoff radii. The estimates are rounded up to the next power of base to amortize future resizing.

Parameters:

Name Type Description Default
particles_per_system Table[SystemId, Array]

Number of particles per system.

 required
systems Table[SystemId, NeighborListSystems]

System data with cell information.

 required
cutoffs Table[SystemId, Array]

Cutoff distance per system.

 required
base float

Base for power-of rounding (default 2).

 2
multiplier float

Safety factor applied to the estimate (default 1.0).

 1.0

Returns:

Type Description
UniversalNeighborlistParameters

A UniversalNeighborlistParameters instance with estimated values.
Source code in src/kups/core/neighborlist/parameters.py 
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
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
134
@classmethod
@no_jax_tracing
def estimate(
    cls,
    particles_per_system: Table[SystemId, Array],
    systems: Table[SystemId, NeighborListSystems],
    cutoffs: Table[SystemId, Array],
    *,
    base: float = 2,
    multiplier: float = 1.0,
) -> UniversalNeighborlistParameters:
    """Estimate parameters for all neighbor list types from system geometry.

    Computes conservative initial capacities based on particle density
    and cutoff radii. The estimates are rounded up to the next power of
    ``base`` to amortize future resizing.

    Args:
        particles_per_system: Number of particles per system.
        systems: System data with cell information.
        cutoffs: Cutoff distance per system.
        base: Base for power-of rounding (default 2).
        multiplier: Safety factor applied to the estimate (default 1.0).

    Returns:
        A ``UniversalNeighborlistParameters`` instance with estimated values.
    """

    def _next_power(total: float | Array) -> int:
        return int(next_higher_power(jnp.array(total * multiplier), base=base))

    sys = Table.join(systems, particles_per_system, cutoffs)
    total_candidates = total_image_candidates = total_edges = max_cells = 0
    with no_post_init():
        for _, (s, n_p, c) in sys:
            n_bins = num_cells(s, c).prod()
            candidates = min(n_p / n_bins * (3**3), n_p)
            # A cutoff reaching past perp/2 replicates each candidate once per
            # periodic image (product of per-axis image counts). Summing per
            # system keeps the estimate tight for heterogeneous cutoffs instead
            # of assuming every system replicates at the maximum rate.
            images = candidate_image_counts(s.cell, c).prod()
            total_candidates += candidates
            total_image_candidates += _next_power(candidates) * images
            total_edges += _estimate_avg_num_edges(
                n_p, s.cell.volume, c, base, multiplier
            )
            max_cells = max(n_bins, max_cells)

    return UniversalNeighborlistParameters(
        avg_edges=int(total_edges // sys.size),
        avg_candidates=_next_power(total_candidates / sys.size),
        avg_image_candidates=_next_power(total_image_candidates / sys.size),
        cells=int(max_cells),
    )