jaxdem.forces.router#

Force model router selecting laws based on species pairs.

Classes

ForceRouter([laws, table])

A ForceModel that selects the force law from the species of the interacting particles.
class jaxdem.forces.router.ForceRouter(laws: tuple[ForceModel, ...] = (), table: tuple[tuple[ForceModel, ...], ...] = ())#

Bases: ForceModel

A ForceModel that selects the force law from the species of the interacting particles.

The router holds a symmetric \(S \times S\) lookup table of force laws, where \(S\) is the number of species. For a particle pair \((i, j)\), the router evaluates the law at table[species_id[i]][species_id[j]].

Notes

  • Use from_dict to build the table from a mapping of species pairs. Pairs not present in the mapping default to an empty LawCombiner, which produces zero force, torque, and energy.

  • Dispatch evaluates each upper-triangle entry and selects the result with jax.lax.select_n(). The cost grows quadratically with the number of species, for scalar and batched calls alike.

  • required_material_properties is the union of the requirements of all laws in the table.

table: tuple[tuple[ForceModel, ...], ...]#

A symmetric \(S \times S\) table where entry table[a][b] is the ForceModel that governs interactions between species a and b.

property supports_analytical_energy_gradient: bool[source]#

Whether every law reachable through the table supports minimization.

property species_capacity: int | None[source]#

Smallest capacity imposed by this table or a nested router.

history_shape(dim: int) → tuple[int, ...][source]#

Trailing shape of one pair’s history; stateless laws use (0,).

Composite laws flatten child histories internally and restore this shape before invoking the child law.

init_history(pair_shape: tuple[int, ...], dim: int) → Array[source]#

Initialize the history variables for this force model.

Parameters:

  • pair_shape (tuple[int, ...]) – Collider-specific pair axes, such as a flat contact pool or a dense (N, N) pair array.

  • dim (int) – Spatial dimension used by dimension-dependent history layouts.

Returns:

A zero-filled array with shape pair_shape + history_shape(dim).

Return type:

jax.Array

search_radii(state: State, system: System) → jax.Array[source]#

Conservative per-primitive reach across every routable law.

static force(i: int, j: int, pos: jax.Array, state: State, system: System, history: jax.Array, *, advance_history: bool = True, quasistatic: bool = False) → tuple[jax.Array, jax.Array, jax.Array][source]#
property required_material_properties: tuple[str, ...][source]#

Names of the material properties this force model needs.

The sorted union of the material properties required by all laws in the table. Each name must be present in System.mat_table. Used for validation.

static from_dict(S: int, mapping: dict[tuple[int, int], ForceModel]) → ForceRouter[source]#

Build a ForceRouter from a mapping of species pairs to force laws.

The router symmetrizes the mapping: entry (a, b) also fills (b, a). Pairs not present in the mapping default to an empty LawCombiner (zero force, torque, and energy).

Parameters:

  • S (int) – Number of species. The resulting table has shape S x S.

  • mapping (dict[tuple[int, int], ForceModel]) – Mapping from species-index pairs to the force law that governs interactions between those species.

Returns:

A router with the fully populated, symmetric lookup table.

Return type:

ForceRouter

static energy(i: int, j: int, pos: jax.Array, state: State, system: System, history: jax.Array) → jax.Array[source]#

Compute the pair energy with the law the species select.

Every law in the table is evaluated on its own slice of the concatenated history and the result is selected by species pair.

Parameters:

  • i (int) – Index of the first particle.

  • j (int) – Index of the second particle.

  • pos (jax.Array) – Particle positions used to evaluate the interaction.

  • state (State) – Current state of the simulation.

  • system (System) – Simulation system configuration.

  • history (jax.Array) – Concatenated pair history of the table’s laws.

Returns:

Scalar potential energy computed by the law at table[species_id[i]][species_id[j]].

Return type:

jax.Array