jaxdem.forces.router#

Force model router selecting laws based on species pairs.

Classes

ForceRouter([laws, table])

A ForceModel implementation that dispatches to different force laws based on the species of the interacting particles.
class jaxdem.forces.router.ForceRouter(laws: tuple[ForceModel, ...] = (), table: tuple[tuple[ForceModel, ...], ...] = ())#

Bases: ForceModel

A ForceModel implementation that dispatches to different force laws based on 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 law at table[species_id[i]][species_id[j]] is evaluated.

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.

  • For per-pair scalar calls, dispatch uses jax.lax.switch(), so only the selected law is evaluated at runtime.

  • For batched calls, all \(S^2\) laws are evaluated for every pair and the result is selected afterwards, so the cost grows quadratically with the number of species.

  • 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 governing interactions between species a and b.

property requires_history: bool[source]#

Indicates whether this force model requires persistent pair history.

init_history(shape: tuple[int, ...]) → Any[source]#
static force_and_history(i: int, j: int, pos: jax.Array, state: State, system: System, history: Any) → tuple[jax.Array, jax.Array, Any][source]#
property required_material_properties: tuple[str, ...][source]#

A static tuple of strings specifying the material properties required by this force model.

The sorted union of the material properties required by all laws in the table. These properties must be present in the System.mat_table for the model to function correctly. This is 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 mapping is symmetrized: entry (a, b) also populates (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 governing interactions between those species.

Returns:

A router with the fully populated, symmetric lookup table.

Return type:

ForceRouter

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

Compute the force and torque acting on particle \(i\) due to particle \(j\) using the law selected by their species.

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.

Returns:

A tuple (force, torque) computed by the law at table[species_id[i]][species_id[j]].

Return type:

Tuple[jax.Array, jax.Array]

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

Compute the potential energy of the interaction between particle \(i\) and particle \(j\) using the law selected by their species.

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.

Returns:

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

Return type:

jax.Array