jaxdem.forces.router#

Force model router selecting laws based on species pairs.

Classes

ForceRouter([laws, table])

Static species-to-force lookup table.
class jaxdem.forces.router.ForceRouter(laws: Tuple[ForceModel, ...] = (), table: Tuple[Tuple[ForceModel, ...], ...] = ())[source]#

Bases: ForceModel

Static species-to-force lookup table.

table: Tuple[Tuple[ForceModel, ...], ...]#
static from_dict(S: int, mapping: dict[Tuple[int, int], ForceModel]) → ForceRouter[source][source]#
static force(i: int, j: int, pos: jax.Array, state: State, system: System) → jax.Array[source][source]#
static energy(i: int, j: int, pos: jax.Array, state: State, system: System) → jax.Array[source][source]#