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.

static create(state_shape: tuple[int, ...], *, gravity: jax.Array | None = None, force_functions: Sequence[ForceFunction | tuple[ForceFunction, bool] | tuple[ForceFunction, EnergyFunction] | tuple[ForceFunction, EnergyFunction, bool]] = ()) → ForceManager[source]#

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.

Parameters:

state (State) – Current state of the simulation.
system (System) – Simulation system configuration.
force (jax.Array) – External force to be added to all particles in the current order.
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_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 ID=idx.

Parameters:

state (State) – Current state of the simulation.
system (System) – Simulation system configuration.
force (jax.Array) – External force to be added to particles with ID=idx.
idx (jax.Array) – ID 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.

Parameters:

state (State) – Current state of the simulation.
system (System) – Simulation system configuration.
torque (jax.Array) – External torque to be added to all particles in the current order..

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

Add an external torque to particles with ID=idx.

Parameters:

state (State) – Current state of the simulation.
system (System) – Simulation system configuration.
torque (jax.Array) – External torque to be added to particles with ID=idx.
idx (jax.Array) – ID 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 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

Static species-to-force lookup table.

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

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.

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

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.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

Sum a tuple of elementary force laws.

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.

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.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\).
A small epsilon is added to the squared distance (\(r^2\)) before taking the square root to prevent division by zero or NaN issues when particles are perfectly co-located.

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

\[\delta = (R_i + R_j) - r\]

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 scalar overlap \(s\) is defined as:

\[s = \max \left(0, \frac{R_i + R_j}{r} - 1 \right)\]

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

\[F_{ij} = k_{eff,\; ij} s r_{ij}\]

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

\[E_{ij} = \frac{1}{2} k_{eff,\; ij} s^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 vector acting on particle \(i\) due to particle \(j\).

Return type:

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.

Modules

`cundall_strack`	Cundall-Strack linear spring-dashpot 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`