jaxdem.forces#

Force-law interfaces and concrete implementations.

Classes

ForceModel([laws])

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

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

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, frozen=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, state: State, system: System) → Tuple[jax.Array, jax.Array][source][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, state: State, system: System) → jax.Array[source][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.LawCombiner(laws: Tuple[ForceModel, ...] = ())[source]#

Bases: ForceModel

Sum a tuple of elementary force laws.

static energy(i, j, state, system)[source][source]#

static force(i, j, state, system)[source][source]#

class jaxdem.forces.ForceRouter(laws: Tuple[ForceModel, ...] = (), table: Tuple[Tuple[ForceModel, ...], ...] = ())[source]#

Bases: ForceModel

Static species-to-force lookup table.

table: Tuple[Tuple['ForceModel', ...], ...]#

static energy(i, j, state, system)[source][source]#

static force(i, j, state, system)[source][source]#

static from_dict(S: int, mapping: dict[Tuple[int, int], ForceModel])[source][source]#

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

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 energy(i: int, j: int, state: State, system: System) → jax.Array[source][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

static force(i: int, j: int, state: State, system: System) → Tuple[jax.Array, jax.Array][source][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

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_torque: jax.Array, force_functions: Tuple[ForceFunction, ...] = ())[source]#

Bases: object

Manage per-particle force contributions prior to pairwise interactions.

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

external_force: jax.Array#: Accumulated external force applied to all particles. 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.

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

static add_force(state: State, system: System, force: jax.Array, *, torque: jax.Array | None = None) → System[source][source]#

Accumulate an external force (and optionally 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.
force – Force contribution to accumulate. Must be a single (dim,) vector applied uniformly.
torque – Optional torque contribution to accumulate. Must be broadcast-compatible with external_torque.

Returns:

A new jaxdem.System instance with the updated accumulated forces and torques.

Return type:

System

static add_force_at(state: State, system: System, force: jax.Array, idx: jax.Array, *, torque: jax.Array | None = None) → System[source][source]#

Accumulate an external force (and optionally torque) to be applied on the next apply call over idx particles.

Parameters:

state (State) – Current state of the simulation.
system (System) – Simulation system configuration.
force – Force contribution to accumulate. Must be an array broadcastable to external_force[idx].
idx – integer indices of the particles receiving the contribution.
torque – Optional torque contribution to accumulate. Must be an array broadcastable to external_torque[idx].

Returns:

A new jaxdem.System instance with the updated accumulated forces and torques.

Return type:

System

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

Overwrite state.force with managed per-particle contributions.

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]

Note

This method donates state and system

static create(dim: int, shape: Tuple, *, gravity: jax.Array | None = None, force_functions: Sequence[ForceFunction] = ()) → ForceManager[source][source]#

Create a ForceManager for a system with dim spatial dimensions.

Parameters:

dim – Spatial dimension of the managed system.
gravity – Optional initial gravitational acceleration. Defaults to zeros.
force_functions – Sequence of callables with signature (state, system, index) returning per-particle force/torque contributions.

Modules

`force_manager`	Utilities for managing per-particle force contributions.
`law_combiner`	Composite force model that sums multiple force laws.
`router`	Force model router selecting laws based on species pairs.
`spring`	Linear spring force model.