jaxdem.forces#

Force-law interfaces and concrete implementations.

Classes

ForceModel([required_material_properties, laws])

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

class jaxdem.forces.ForceModel(required_material_properties: Tuple[str, ...] = (), 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):
        ...

required_material_properties: Tuple[str, ...]#

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.

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) → jax.Array[source][source]#

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

Parameters:

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

Returns:

Force vector acting on particle \(i\) due to particle \(j\). Shape (dim,).

Return type:

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

class jaxdem.forces.LawCombiner(required_material_properties: Tuple[str, ...] = (), 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]#

classmethod registry_name() → str[source]#

property type_name: str[source]#

class jaxdem.forces.ForceRouter(required_material_properties: Tuple[str, ...] = (), 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(required_material_properties: Tuple[str, ...] = ('young_eff',), 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) → 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

classmethod registry_name() → str[source]#

property type_name: str[source]#

Modules

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