jaxdem.forces#
Force-law interfaces.
Classes
|
Abstract base class for inter-particle force laws and their potential energies. |
- class jaxdem.forces.WCA(laws: tuple[ForceModel, ...] = ())#
Bases:
ForceModelWeeks-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
- class jaxdem.forces.CundallStrackForce(laws: tuple[ForceModel, ...] = ())#
Bases:
ForceModelCundall-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.
- 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:
objectManage 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_functionswith 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 toforce_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
ForceManagerfor 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
applycall.The method returns only
system. The state does not change because the force waits in theForceManagerbuffer untilapply()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
forcein total, notforceper 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
idxfor the nextapplycall.The method returns only
system. The state does not change because the force waits in theForceManagerbuffer untilapply()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
applycall.The method returns only
system. The state does not change because the torque waits in theForceManagerbuffer untilapply()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
torquein total, nottorqueper 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
idxfor the nextapplycall.The method returns only
system. The state does not change because the torque waits in theForceManagerbuffer untilapply()runs.
- 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.
- class jaxdem.forces.ForceModel(laws: tuple[ForceModel, ...] = ())#
Bases:
Factory,ABCAbstract 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:#
Example:#
To define a custom force model, inherit from
ForceModeland 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
ForceModelinstances 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:
- Returns:
A tuple
(force, torque)whereforcehas shape(dim,)andtorquehas 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:
- Returns:
Scalar potential energy of the interaction between particles \(i\) and \(j\).
- Return type:
jax.Array
- 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
- class jaxdem.forces.ForceRouter(laws: tuple[ForceModel, ...] = (), table: tuple[tuple[ForceModel, ...], ...] = ())#
Bases:
ForceModelA 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 emptyLawCombiner, 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_propertiesis 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 theForceModelthat governs interactions between speciesaandb.
- 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
ForceRouterfrom 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 emptyLawCombiner(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:
- 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:
- Returns:
A tuple
(force, torque)computed by the law attable[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:
- 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:
ForceModelHertzian 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
youngandpoissonproperties per particle fromSystem.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:
- Returns:
(force, torque)with shapes(dim,)and(ang_dim,).- Return type:
tuple[jax.Array, jax.Array]
- class jaxdem.forces.LawCombiner(laws: tuple[ForceModel, ...] = ())#
Bases:
ForceModelA 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_modelis the sub-law itself. Laws that read their own configuration fromjaxdem.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_propertiesis the union of the requirements of all contained laws.
- 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:
- 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:
- 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:
ForceModelLennard-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#
- class jaxdem.forces.SpringForce(laws: tuple[ForceModel, ...] = ())#
Bases:
ForceModelLinear spring-like interaction between particles.
Notes
The model reads the ‘effective Young’s modulus’ (\(k_{eff,\; ij}\)) from the
jaxdem.System.mat_tableusing 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-
wherehelpers injaxdem.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:
- 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:
- Returns:
Scalar potential energy of the interaction between particles \(i\) and \(j\).
- Return type:
jax.Array
- class jaxdem.forces.WCAShifted(laws: tuple[ForceModel, ...] = ())#
Bases:
ForceModelContact-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].
- class jaxdem.forces.SphereFacetSpringForce(laws: tuple[ForceModel, ...] = ())#
Bases:
ForceModelLinear 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 usecutoff >= max facet circumradius + thickness.The contact thickness comes from the facet vertices’
state.rad(set throughState.add_facet(thickness=...)). The force model itself has no thickness parameter.
- class jaxdem.forces.FacetFacetSpringForce(laws: tuple[ForceModel, ...] = ())#
Bases:
ForceModelLinear 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 usecutoff >= max facet circumradius + thickness.The contact thickness comes from the facet vertices’
state.rad(set throughState.add_facet(thickness=...)). The force model itself has no thickness parameter.
Modules
Cundall-Strack linear spring-dashpot contact force model. |
|
Facet contact force model. |
|
External and custom force contributions that do not depend on the collider. |
|
Hertzian (nonlinear) normal contact force model. |
|
Composite force model that sums multiple force laws. |
|
Force model router selecting laws based on species pairs. |
|
Linear spring force model. |
|