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 every law in the table 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 requires_history: bool[source]#

Whether this force model needs 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]#

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 force(i: int, j: int, pos: jax.Array, state: State, system: System) → tuple[jax.Array, jax.Array][source]#

Compute the force and torque on particle \(i\) from particle \(j\) with the law their species select.

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\) with the law their species select.

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 potential energy computed by the law at table[species_id[i]][species_id[j]].

Return type:

jax.Array