jaxdem.forces#

Force-law interfaces.

Classes

ForceModel([laws])

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

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

Bases: ForceModel

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

The model reads the material-pair parameter epsilon_eff[mi, mj].

The model derives the length scale \(\sigma_{ij}\) from the particle radii (as in 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]#

Names of the material properties this force model needs.

Each name (for example ‘young_eff’ or ‘restitution’) must be present in System.mat_table. 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 with a linear elastic assumption, viscous damping, and Coulomb friction.

Effective Properties The model computes the effective mass \(m_{eff}\), restitution coefficient \(e_{eff}\), and friction coefficient \(\mu\) 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 model computes the shear modulus \(G\) per particle from Young’s modulus \(E\) and Poisson’s ratio \(\nu\):

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

The model treats the effective stiffnesses for the normal (\(k_n\)) and tangential (\(k_t\)) directions 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 the 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, limited to repulsive values:

\[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. This stateless implementation does not track shear history per pair. It approximates the tangential force with the viscous dashpot, 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 The rolling friction torque resists the 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 the conservative potential energy of 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]#

Names of the material properties this force model needs.

Each name (for example ‘young_eff’ or ‘restitution’) must be present in System.mat_table. 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 outside the collider.

After the collider runs, apply() adds these contributions to the state forces and aggregates them over 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). apply() clears this buffer.

external_force_com: jax.Array#

Accumulated external force applied to the center of mass (induces no torque). apply() clears this buffer.

external_torque: jax.Array#

Accumulated external torque applied to all particles. apply() clears this buffer.

is_com_force: tuple[bool, ...]#

Boolean array corresponding to force_functions with shape (n_forces,). If True, apply the force to the center of mass (no induced torque). If False, apply the force at the particle position (induces torque through the 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 | None] | tuple[ForceFunction, EnergyFunction | None, 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. Signature of ForceFunc: (pos, state, system) -> (Force, Torque). Signature of EnergyFunc: (pos, state, system) -> Energy. Supported formats:

    • 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]#

Buffer an external force on all particles for the next apply call.

The method returns only system. The state does not change because the force waits in the ForceManager buffer 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 add to every particle (in the state’s current particle order).

  • is_com (bool, optional) – If True, apply the force to the center of mass (no induced torque). The force goes to every clump member, so each clump receives force in total, not force per member. If False (default), apply the force at the particle position (induces torque).

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

Buffer an external force on the particles with array index idx for the next apply call.

The method returns only system. The state does not change because the force waits in the ForceManager buffer until apply() runs.

Parameters:
  • state (State) – Current state of the simulation.

  • system (System) – Simulation system configuration.

  • force (jax.Array) – External force to add to the particles with array index idx.

  • idx (jax.Array) – Array indices of the particles the external force acts on.

  • is_com (bool, optional) – If True, apply the force to the center of mass (no induced torque). If False (default), apply the force at the particle position (induces torque).

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

Buffer an external torque on all particles for the next apply call.

The method returns only system. The state does not change because the torque waits in the ForceManager buffer 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 add to every particle (in the state’s current particle order). The torque goes to every clump member, so 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]#

Buffer an external torque on the particles with array index idx for the next apply call.

The method returns only system. The state does not change because the torque waits in the ForceManager buffer until apply() runs.

Parameters:
  • state (State) – Current state of the simulation.

  • system (System) – Simulation system configuration.

  • torque (jax.Array) – External torque to add to the particles with array index idx.

  • idx (jax.Array) – Array indices of the particles the external torque acts on.

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

Add the managed per-particle contributions to the collider forces, then aggregate over clumps and broadcast back to the members.

Parameters:
  • state (State) – Current state of the simulation.

  • system (System) – Simulation system configuration.

Returns:

The updated state and system.

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) – Current state of the simulation.

  • system (System) – Simulation system configuration.

Returns:

Scalar total potential energy.

Return type:

jax.Array

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

Bases: Factory, ABC

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

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

Notes:#

  • The force() and energy() methods must handle the case where i and j refer to the same particle (i == j). Self-interaction calls can 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, ...]#

Static tuple of other ForceModel instances that compose this force model.

Use it to build composite force models, for example a spring force plus 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 on particle \(i\) from particle \(j\).

Parameters:
  • i (int) – Index of the first particle (on which the interaction acts).

  • j (int) – Index of the second particle (which exerts the interaction).

  • pos (jax.Array) – Particle positions.

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

  • pos (jax.Array) – Particle positions.

  • state (State) – Current state of the simulation.

  • system (System) – Simulation system configuration.

Returns:

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

Return type:

jax.Array

property requires_history: bool[source]#

Whether this force model needs persistent pair history.

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

Initialize the 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, or None by default.

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]#

Names of the material properties this force model needs.

Each name (for example ‘young_eff’ or ‘restitution’) must be present in System.mat_table. Used for validation.

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

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

Bases: ForceModel

Hertzian nonlinear normal contact force between elastic spheres.

The model computes the effective Young’s modulus \(E^*\) 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 model reads the young and poisson properties per particle from System.mat_table. It does not use a matchmaker effective value.

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

Compute the Hertzian normal contact force on particle i from particle 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]#

Compute the 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]#

Names of the material properties this force model needs.

Each name (for example ‘young_eff’ or ‘restitution’) must be present in System.mat_table. Used for validation.

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

Bases: ForceModel

A ForceModel 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

  • The combiner evaluates each sub-law with a system whose force_model is the sub-law itself. 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]#

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 contained laws. Each name must be present in System.mat_table. 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 on particle \(i\) from 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\) from 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 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.

The model reads the material-pair parameter epsilon_eff[mi, mj].

The model derives the length scale \(\sigma_{ij}\) from the particle radii (as in 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]#

Names of the material properties this force model needs.

Each name (for example ‘young_eff’ or ‘restitution’) must be present in System.mat_table. Used for validation.

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

Bases: ForceModel

Linear spring-like interaction between particles.

Notes

  • The model reads the ‘effective Young’s modulus’ (\(k_{eff,\; ij}\)) from the jaxdem.System.mat_table using the material IDs of the interacting particles.

  • The force is zero if \(i == j\).

  • The model computes distances and normals with the zero-safe double-where helpers in jaxdem.utils.linalg. The force and its gradients stay 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 the linear spring force on particle \(i\) from particle \(j\).

Returns zero when \(i = j\).

Parameters:
  • i (int) – Index of the first particle.

  • j (int) – Index of the second particle.

  • pos (jax.Array) – Particle positions.

  • 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 the linear spring 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.

  • pos (jax.Array) – Particle positions.

  • state (State) – Current state of the simulation.

  • system (System) – Simulation system configuration.

Returns:

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

Return type:

jax.Array

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

Names of the material properties this force model needs.

Each name (for example ‘young_eff’ or ‘restitution’) must be present in System.mat_table. Used for validation.

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

Bases: ForceModel

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

The interaction starts 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)

The model reads the 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]#

Names of the material properties this force model needs.

Each name (for example ‘young_eff’ or ‘restitution’) must be present in System.mat_table. Used for validation.

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

Bases: ForceModel

Linear spring contact between spheres and facets.

Warning

The model detects facet contacts through the facet’s vertex spheres, in particular the facet’s primary vertex. The collider’s neighbor cutoff must cover 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 model misses the contact or applies it asymmetrically. Cell-list based colliders must use cutoff >= max facet circumradius + thickness.

The contact thickness comes from the facet vertices’ state.rad (set through 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

The model detects facet contacts through the facets’ vertex spheres, in particular each facet’s primary vertex. The collider’s neighbor cutoff must cover 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 model misses the contact or applies it asymmetrically. Cell-list based colliders must use cutoff >= max facet circumradius + thickness.

The contact thickness comes from the facet vertices’ state.rad (set through 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

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