jaxdem.forces#

Force-law interfaces.

Classes

ForceModel([laws])

Abstract base class for defining inter-particle force laws and their corresponding potential energies.

class jaxdem.forces.WCA(laws: tuple[ForceModel, ...] = ())#

Bases: ForceModel

Weeks-Chandler-Andersen (WCA) purely repulsive Lennard-Jones interaction.

Uses material-pair parameter:

epsilon_eff[mi, mj]

The length scale \(\sigma_{ij}\) is derived from particle radii (like spring.py):

\[\sigma_{ij} = R_i + R_j\]

Potential (for r < r_c = 2^(1/6) sigma):: U(r) = 4 eps [(sigma/r)^12 - (sigma/r)^6] + eps
else:: U(r) = 0
Force:: F_vec = 24 eps (2 (sigma/r)^12 - (sigma/r)^6) * (1/r^2) * r_ij

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

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

property required_material_properties: tuple[str, ...][source]#

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

These properties (e.g., ‘young_eff’, ‘restitution’, …) must be present in the System.mat_table for the model to function correctly. This is used for validation.

class jaxdem.forces.CundallStrackForce(laws: tuple[ForceModel, ...] = ())#

Bases: ForceModel

Cundall-Strack linear spring-dashpot normal and tangential force model with rolling friction.

Computes the interaction between two spheres using a linear elastic assumption combined with viscous damping and Coulomb friction.

Effective Properties The effective mass \(m_{eff}\), restitution coefficient \(e_{eff}\), and friction coefficient \(\mu\) are taken as:

\[m_{eff} = \left( \frac{1}{m_i} + \frac{1}{m_j} \right)^{-1}, \quad e_{eff} = \min(e_i, e_j), \quad \mu = \min(\mu_i, \mu_j)\]

The Shear Modulus \(G\) is computed per particle from Young’s modulus \(E\) and Poisson’s ratio \(\nu\):

\[G = \frac{E}{2(1 + \nu)}\]

The effective stiffnesses for the normal (\(k_n\)) and tangential (\(k_t\)) directions are modelled as springs in series:

\[k_n = \frac{2 E_i R_i E_j R_j}{E_i R_i + E_j R_j}, \quad k_t = \frac{2 G_i R_i G_j R_j}{G_i R_i + G_j R_j}\]

The viscous damping coefficient \(\beta\) and resulting directional damping coefficients are:

\[\beta = \frac{-\ln(e_{eff})}{\sqrt{\pi^2 + \ln^2(e_{eff})}}\]

\[\gamma_n = 2 \beta \sqrt{k_n m_{eff}}, \quad \gamma_t = 2 \beta \sqrt{k_t m_{eff}}\]

Forces The normal force \(F_n\) includes spring repulsion and viscous damping, constrained to be strictly repulsive:

\[F_n = \max(0, k_n \delta_n - \gamma_n v_n)\]

Note on Tangential Force: True Cundall-Strack uses an integrated shear displacement history. Because this stateless implementation does not track tangential history per pair, the tangential force is approximated using the dynamic viscous dashpot strictly capped by the Coulomb sliding friction limit:

\[\mathbf{F}_{t, trial} = -\gamma_t \mathbf{v}_t\]

\[\mathbf{F}_t = \min(\|\mathbf{F}_{t, trial}\|, \mu F_n) \frac{\mathbf{v}_t}{\|\mathbf{v}_t\|}\]

Rolling Friction A resistive torque opposing relative angular velocity at the contact:

\[\boldsymbol{\tau}_{\text{roll}} = -\mu_r \, R_{\text{eff}} \, F_n \, \hat{\omega}_{\text{rel}}\]

where \(\mu_r = \min(\mu_{r,i}, \mu_{r,j})\) is the effective rolling friction coefficient, \(R_{\text{eff}} = R_i R_j / (R_i + R_j)\), and \(\hat{\omega}_{\text{rel}}\) is the unit relative angular velocity. Setting \(\mu_r = 0\) (the default) disables rolling friction.

References

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

Compute Cundall-Strack normal and tangential forces and torque.

Parameters:

i (int) – Particle indices.
j (int) – Particle indices.
pos (jax.Array) – Particle positions.
state (State) – Current simulation state.
system (System) – System configuration.

Returns:

(force, torque) with dimension-agnostic shapes.

Return type:

tuple[jax.Array, jax.Array]

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

Compute conservative potential energy for the interaction.

\[U_{ij} = \frac{1}{2} k_n \delta_n^2\]

Parameters:

i (int) – Particle indices.
j (int) – Particle indices.
pos (jax.Array) – Particle positions.
state (State) – Current simulation state.
system (System) – System configuration.

Returns:

Scalar potential energy.

Return type:

jax.Array

property required_material_properties: tuple[str, ...][source]#

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

These properties (e.g., ‘young_eff’, ‘restitution’, …) must be present in the System.mat_table for the model to function correctly. This is used for validation.

class jaxdem.forces.ForceManager(gravity: jax.Array, external_force: jax.Array, external_force_com: jax.Array, external_torque: jax.Array, is_com_force: tuple[bool, ...] = (), force_functions: tuple[ForceFunction, ...] = (), energy_functions: tuple[EnergyFunction | None, ...] = ())#

Bases: object

Manage custom force contributions external to the collider. It also accumulates forces in the state after collider application, accounting for rigid bodies.

gravity: jax.Array#: Constant acceleration applied to all particles. Shape (dim,).

external_force: jax.Array#: Accumulated external force applied to all particles (at particle position). This buffer is cleared when apply() is invoked.

external_force_com: jax.Array#: Accumulated external force applied to Center of Mass (does not induce torque). This buffer is cleared when apply() is invoked.

external_torque: jax.Array#: Accumulated external torque applied to all particles. This buffer is cleared when apply() is invoked.

is_com_force: tuple[bool, ...]#: Boolean array corresponding to force_functions with shape (n_forces,). If True, the force is applied to the Center of Mass (no induced torque). If False, the force is applied to the constituent particle (induces torque via lever arm).

force_functions: tuple[ForceFunction, ...]#: Tuple of callables with signature (pos, state, system) returning per-particle force and torque arrays.

energy_functions: tuple[EnergyFunction | None, ...]#: Tuple of callables (or None) with signature (pos, state, system) returning per-particle potential energy arrays. Corresponds to force_functions.

Create a ForceManager for a state with the given shape.

Parameters:

state_shape – Shape of the state position array, typically (..., dim).
gravity – Optional initial gravitational acceleration. Defaults to zeros of shape (dim,).
force_functions –
Sequence of callables or tuples. Supported formats:
- ForceFunc: Applied at particle, no potential energy.
- (ForceFunc, bool): Boolean specifies if it is a COM force.
- (ForceFunc, EnergyFunc): Includes potential energy function.
- (ForceFunc, EnergyFunc, bool): Includes energy and COM specifier.
Signature of ForceFunc: (pos, state, system) -> (Force, Torque) Signature of EnergyFunc: (pos, state, system) -> Energy

Supported formats for force_functions items: - func -> (func, None, False) - (func,) -> (func, None, False) - (func, bool) -> (func, None, bool) - (func, energy) -> (func, energy, False) - (func, energy, bool) -> (func, energy, bool) - (func, None, bool) -> (func, None, bool)

static add_force(state: State, system: System, force: jax.Array, *, is_com: bool = False) → System[source]#

Accumulate an external force to be applied on the next apply call for all particles.

Only the system is returned because the state is not modified: the force is buffered in the system’s ForceManager until apply() runs.

Parameters:

state (State) – Current state of the simulation. Used to normalize COM forces by the clump member count.
system (System) – Simulation system configuration.
force (jax.Array) – External force to be added to every particle (in the state’s current particle order).
is_com (bool, optional) – If True, force is applied to the Center of Mass (no induced torque); since the force is written to every clump member, each clump receives force in total, not force per member. If False (default), force is applied to Particle Position (induces torque).

static add_force_at(state: State, system: System, force: jax.Array, idx: jax.Array, *, is_com: bool = False) → System[source]#

Add an external force to particles with array index idx.

Only the system is returned because the state is not modified: the force is buffered in the system’s ForceManager until apply() runs.

Parameters:

state (State) – Current state of the simulation.
system (System) – Simulation system configuration.
force (jax.Array) – External force to be added to particles with array index idx.
idx (jax.Array) – Array indices of the particles affected by the external force.
is_com (bool, optional) – If True, force is applied to Center of Mass (no induced torque). If False (default), force is applied to Particle Position (induces torque).

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

Accumulate an external torque to be applied on the next apply call for all particles.

Only the system is returned because the state is not modified: the torque is buffered in the system’s ForceManager until apply() runs.

Parameters:

state (State) – Current state of the simulation. Used to normalize the torque by the clump member count.
system (System) – Simulation system configuration.
torque (jax.Array) – External torque to be added to every particle (in the state’s current particle order); since the torque is written to every clump member, each clump receives torque in total, not torque per member.

static add_torque_at(state: State, system: System, torque: jax.Array, idx: jax.Array) → System[source]#

Add an external torque to particles with array index idx.

Only the system is returned because the state is not modified: the torque is buffered in the system’s ForceManager until apply() runs.

Parameters:

state (State) – Current state of the simulation.
system (System) – Simulation system configuration.
torque (jax.Array) – External torque to be added to particles with array index idx.
idx (jax.Array) – Array indices of the particles affected by the external force.

static apply(state: State, system: System) → tuple[State, System][source]#

Accumulate managed per-particle contributions on top of collider/contact forces, then perform final clump aggregation + broadcast.

Parameters:

state (State) – Current state of the simulation.
system (System) – Simulation system configuration.

Returns:

The updated state and system after one time step.

Return type:

Tuple[State, System]

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

Compute the total potential energy of the system.

Notes

The energy of clump members is divided by the number of spheres in the clump.

Parameters:

state (State) – The current state of the simulation.
system (System) – The configuration of the simulation.

Returns:

Scalar JAX array representing the total potential energy.

Return type:

jax.Array

class jaxdem.forces.ForceModel(laws: tuple[ForceModel, ...] = ())#

Bases: Factory, ABC

Abstract base class for defining inter-particle force laws and their corresponding potential energies.

Concrete subclasses implement specific force and energy models, such as linear springs, Hertzian contacts, etc.

Notes:#

The force() and energy() methods should correctly handle the case where i and j refer to the same particle (i.e., i == j). There is no guarantee that self-interaction calls will not occur.

Example:#

To define a custom force model, inherit from ForceModel and implement its abstract methods:

>>> @ForceModel.register("myCustomForce")
>>> @jax.tree_util.register_dataclass
>>> @dataclass(slots=True)
>>> class MyCustomForce(ForceModel):
        ...

laws: tuple[ForceModel, ...]#

A static tuple of other ForceModel instances that compose this force model.

This allows for creating composite force models (e.g., a total force being the sum of a spring force and a damping force).

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

Compute the force and torque vector acting on particle \(i\) due to particle \(j\).

Parameters:

i (int) – Index of the first particle (on which the interaction acts).
j (int) – Index of the second particle (which is exerting the interaction).
state (State) – Current state of the simulation.
system (System) – Simulation system configuration.

Returns:

A tuple (force, torque) where force has shape (dim,) and torque has shape (1,) in 2D or (3,) in 3D.

Return type:

Tuple[jax.Array, jax.Array]

abstractmethod 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\).

Parameters:

i (int) – Index of the first particle.
j (int) – Index of the second particle.
state (State) – Current state of the simulation.
system (System) – Simulation system configuration.

Returns:

Scalar JAX array representing the potential energy of the interaction between particles \(i\) and \(j\).

Return type:

jax.Array

property requires_history: bool[source]#: Indicates whether this force model requires persistent pair history.

init_history(shape: tuple[int, ...]) → Any[source]#

Initialize history variables for this force model.

Parameters:: shape (tuple[int, ...]) – The expected shape for pair-wise quantities, typically (…, N, max_neighbors).
Returns:: A PyTree of initialized JAX arrays.
Return type:: Any

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

Compute the force and torque, and update history.

By default, this calls force and returns history unchanged.

property required_material_properties: tuple[str, ...][source]#

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

These properties (e.g., ‘young_eff’, ‘restitution’, …) must be present in the System.mat_table for the model to function correctly. This is used for validation.

class jaxdem.forces.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

class jaxdem.forces.HertzianForce(laws: tuple[ForceModel, ...] = ())#

Bases: ForceModel

Hertzian nonlinear normal contact force between elastic spheres.

The effective Young’s modulus \(E^*\) is computed directly from the per-particle Young’s modulus \(E\) and Poisson’s ratio \(\nu\):

\[\frac{1}{E^*} = \frac{1 - \nu_i^2}{E_i} + \frac{1 - \nu_j^2}{E_j}\]

The effective radius is:

\[\frac{1}{R^*} = \frac{1}{R_i} + \frac{1}{R_j}\]

The Hertzian stiffness combines both:

\[k = \tfrac{4}{3}\, E^* \sqrt{R^*}\]

The penetration depth \(\delta\) between particles \(i\) and \(j\) is:

\[\delta = \max(0,\; R_i + R_j - r)\]

where \(r = \|r_{ij}\|\).

The Hertzian normal force and contact energy are:

\[\mathbf{F}_{ij} = k \; \delta^{3/2} \; \hat{n}_{ij}, \qquad U_{ij} = \tfrac{2}{5}\, k \; \delta^{5/2}\]

where \(\hat{n}_{ij} = \mathbf{r}_{ij} / r\).

Notes

The material young and poisson properties are read directly from System.mat_table per particle; no matchmaker effective value is used.

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

Compute Hertzian normal contact force on particle i from j.

\[\mathbf{F}_{ij} = \tfrac{4}{3}\, E^*\, \sqrt{R^*}\; \delta^{3/2}\; \hat{n}_{ij}\]

Parameters:

i (int) – Particle indices.
j (int) – Particle indices.
pos (jax.Array) – Particle positions (rotated to lab frame).
state (State) – Current simulation state.
system (System) – System configuration.

Returns:

(force, torque) with shapes (dim,) and (ang_dim,).

Return type:

tuple[jax.Array, jax.Array]

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

Hertzian contact energy.

\[U_{ij} = \tfrac{2}{5} \cdot \tfrac{4}{3}\, E^*\, \sqrt{R^*}\, \delta^{5/2}\]

Parameters:

i (int) – Particle indices.
j (int) – Particle indices.
pos (jax.Array) – Particle positions.
state (State) – Current simulation state.
system (System) – System configuration.

Returns:

Scalar potential energy.

Return type:

jax.Array

property required_material_properties: tuple[str, ...][source]#

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

These properties (e.g., ‘young_eff’, ‘restitution’, …) must be present in the System.mat_table for the model to function correctly. This is used for validation.

class jaxdem.forces.LawCombiner(laws: tuple[ForceModel, ...] = ())#

Bases: ForceModel

A ForceModel implementation that sums a tuple of elementary force laws.

The total force, torque, and potential energy of the interaction between particles \(i\) and \(j\) are the sums over the contained laws:

\[F_{ij} = \sum_k F^{(k)}_{ij}, \qquad \tau_{ij} = \sum_k \tau^{(k)}_{ij}, \qquad E_{ij} = \sum_k E^{(k)}_{ij}\]

Notes

Each sub-law is evaluated with a system whose force_model is the sub-law itself, so laws that read their own configuration from jaxdem.System.force_model (including nested combiners) work correctly.
An empty combiner (laws=()) returns zero force, torque, and energy. ForceRouter.from_dict() uses it as the default no-interaction law.
required_material_properties is the union of the requirements of all contained laws.

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 contained laws. These properties must be present in the System.mat_table for the model to function correctly. This is used for validation.

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

Compute the total force and torque acting on particle \(i\) due to particle \(j\) by summing all contained laws.

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) with the sums of the forces and torques of all contained laws acting on particle \(i\) due to particle \(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 total potential energy of the interaction between particle \(i\) and particle \(j\) by summing all contained laws.

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 total potential energy of the interaction between particles \(i\) and \(j\).

Return type:

jax.Array

class jaxdem.forces.LennardJones(laws: tuple[ForceModel, ...] = ())#

Bases: ForceModel

Lennard-Jones (LJ) 12-6 interaction with a per-pair cutoff and energy shift.

Uses material-pair parameter:

epsilon_eff[mi, mj]

The length scale \(\sigma_{ij}\) is derived from particle radii (like spring.py):

\[\sigma_{ij} = R_i + R_j\]

Potential (for \(r < r_c = 2.5 \sigma_{ij}\)):

\[U(r) = 4 \epsilon \left[\left(\frac{\sigma}{r}\right)^{12} - \left(\frac{\sigma}{r}\right)^6 \right] - U(r_c)\]

else:

\[U(r) = 0\]

Force (for \(r < r_c\)):

\[\mathbf{F} = 24 \epsilon \left(2 \left(\frac{\sigma}{r}\right)^{12} - \left(\frac{\sigma}{r}\right)^6\right) \frac{1}{r^2}\, \mathbf{r}_{ij}\]

RC_FACTOR: ClassVar[float] = 2.5#

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

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

property required_material_properties: tuple[str, ...][source]#

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

These properties (e.g., ‘young_eff’, ‘restitution’, …) must be present in the System.mat_table for the model to function correctly. This is used for validation.

class jaxdem.forces.SpringForce(laws: tuple[ForceModel, ...] = ())#

Bases: ForceModel

A ForceModel implementation for a linear spring-like interaction between particles.

Notes

The ‘effective Young’s modulus’ (\(k_{eff,\; ij}\)) is retrieved from the jaxdem.System.mat_table based on the material IDs of the interacting particles.
The force is zero if \(i == j\).
Distances and normals are computed with the zero-safe double-where helpers in jaxdem.utils.linalg, so the force and its gradients remain finite when particles are perfectly co-located.

The penetration \(\delta\) (overlap) between two particles \(i\) and \(j\) is:

\[\delta = \max\left(0, (R_i + R_j) - r\right)\]

where \(R_i\) and \(R_j\) are the radii of particles \(i\) and \(j\) respectively, and \(r = ||r_{ij}||\) is the distance between their centers.

The force \(F_{ij}\) acting on particle \(i\) due to particle \(j\) is:

\[F_{ij} = k_{eff,\; ij}\, \delta\, \hat{n}_{ij}\]

where \(\hat{n}_{ij} = \vec{r}_{ij} / r\) is the unit vector from particle \(j\) to particle \(i\).

The potential energy \(E_{ij}\) of the interaction is:

\[E_{ij} = \frac{1}{2} k_{eff,\; ij} \delta^2\]

where \(k_{eff,\; ij}\) is the effective Young’s modulus for the particle pair.

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

Compute linear spring-like interaction force acting on particle \(i\) due to particle \(j\).

Returns zero when \(i = j\).

Parameters:

i (int) – Index of the first particle.
j (int) – Index of the second particle.
state (State) – Current state of the simulation.
system (System) – Simulation system configuration.

Returns:

(force, torque) with shapes (dim,) and (ang_dim,). The torque is always zero for this model.

Return type:

tuple[jax.Array, jax.Array]

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

Compute linear spring-like interaction potential energy between particle \(i\) and particle \(j\).

Returns zero when \(i = j\).

Parameters:

i (int) – Index of the first particle.
j (int) – Index of the second particle.
state (State) – Current state of the simulation.
system (System) – Simulation system configuration.

Returns:

Scalar JAX array representing the potential energy of the interaction between particles \(i\) and \(j\).

Return type:

jax.Array

property required_material_properties: tuple[str, ...][source]#

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

These properties (e.g., ‘young_eff’, ‘restitution’, …) must be present in the System.mat_table for the model to function correctly. This is used for validation.

class jaxdem.forces.WCAShifted(laws: tuple[ForceModel, ...] = ())#

Bases: ForceModel

Contact-start, force-shifted WCA/LJ repulsion.

This model enforces that the interaction “begins” at contact:

cutoff at \(r_c = \sigma_{ij}\) where \(\sigma_{ij} = R_i + R_j\)
\(U(r_c) = 0\)
\(F(r_c) = 0\) (force-shifted; smooth turn-on at contact)

Uses material-pair parameter:

epsilon_eff[mi, mj]

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

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

property required_material_properties: tuple[str, ...][source]#

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

These properties (e.g., ‘young_eff’, ‘restitution’, …) must be present in the System.mat_table for the model to function correctly. This is used for validation.

class jaxdem.forces.SphereFacetSpringForce(laws: tuple[ForceModel, ...] = ())#

Bases: ForceModel

Linear spring contact between spheres and facets.

Warning

Facet contacts are detected through the facet’s vertex spheres (in particular the facet’s primary vertex). The collider’s neighbor cutoff must therefore be at least the facet circumradius (the largest vertex-to-contact-point distance) plus the contact thickness (state.rad). If the primary vertex lies outside the cutoff while the contact point is in range, the contact is silently missed or applied asymmetrically. Cell-list based colliders must be configured with cutoff >= max facet circumradius + thickness.

The contact thickness is taken from the facet vertices’ state.rad (set via State.add_facet(thickness=...)); the force model itself has no thickness parameter.

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

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

class jaxdem.forces.FacetFacetSpringForce(laws: tuple[ForceModel, ...] = ())#

Bases: ForceModel

Linear spring contact between facets.

Warning

Facet contacts are detected through the facets’ vertex spheres (in particular each facet’s primary vertex). The collider’s neighbor cutoff must therefore be at least the facet circumradius (the largest vertex-to-contact-point distance) plus the contact thickness (state.rad). If a primary vertex lies outside the cutoff while the contact point is in range, the contact is silently missed or applied asymmetrically. Cell-list based colliders must be configured with cutoff >= max facet circumradius + thickness.

The contact thickness is taken from the facet vertices’ state.rad (set via State.add_facet(thickness=...)); the force model itself has no thickness parameter.

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

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

Modules

`cundall_strack`	Cundall-Strack linear spring-dashpot contact force model.
`facet_contact`	Facet contact force model.
`force_manager`	Utilities for managing external and custom force contributions that do not depend on the collider.
`hertz`	Hertzian (nonlinear) normal contact force model.
`law_combiner`	Composite force model that sums multiple force laws.
`lennardjones`
`router`	Force model router selecting laws based on species pairs.
`spring`	Linear spring force model.
`wca`
`wca_shifted`