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 image candidate pairs per particle.
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
@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 image candidate pairs per particle.
        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.
        """
        sys = Table.join(systems, particles_per_system, cutoffs)
        total_candidates = total_edges = max_cells = 0
        for _, (s, n_p, c) in sys:
            n_bins = num_cells(s, c).prod()
            total_candidates += min(n_p / n_bins * (3**3), n_p)
            total_edges += _estimate_avg_num_edges(
                n_p, s.cell.volume, c, base, multiplier
            )
            max_cells = max(n_bins, max_cells)
        total_candidates = next_higher_power(
            jnp.array(total_candidates * multiplier / sys.size), base=base
        )
        return UniversalNeighborlistParameters(
            avg_edges=int(total_edges // sys.size),
            avg_candidates=int(total_candidates),
            avg_image_candidates=int(total_candidates),  # Image candidates ~ candidates
            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 
 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
@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.
    """
    sys = Table.join(systems, particles_per_system, cutoffs)
    total_candidates = total_edges = max_cells = 0
    for _, (s, n_p, c) in sys:
        n_bins = num_cells(s, c).prod()
        total_candidates += min(n_p / n_bins * (3**3), n_p)
        total_edges += _estimate_avg_num_edges(
            n_p, s.cell.volume, c, base, multiplier
        )
        max_cells = max(n_bins, max_cells)
    total_candidates = next_higher_power(
        jnp.array(total_candidates * multiplier / sys.size), base=base
    )
    return UniversalNeighborlistParameters(
        avg_edges=int(total_edges // sys.size),
        avg_candidates=int(total_candidates),
        avg_image_candidates=int(total_candidates),  # Image candidates ~ candidates
        cells=int(max_cells),
    )