jaxdem.forces#

Force-law interfaces.

Classes

ForceModel([laws])

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

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.

Every law has the same force and energy signatures. A law without pair memory returns its history argument unchanged and ignores advance_history; a law without velocity-dependent terms ignores quasistatic.

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

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.

property supports_analytical_energy_gradient: bool[source]#

Whether force is the negative analytical gradient of energy.

Translational components use body-center positions and rotational components use world angular increments. Force laws are assumed to satisfy this contract by default. Laws that do not must override this property with False; value-based minimizers then differentiate energy directly, with stored pair history held fixed.

property species_capacity: int | None[source]#

Smallest species-table capacity reachable through this law.

Ordinary laws do not constrain species identifiers. Composite laws override this capability so callers need not inspect concrete types.

abstractmethod static force(i: int, j: int, pos: jax.Array, state: State, system: System, history: jax.Array, *, advance_history: bool = True, quasistatic: bool = False) → tuple[jax.Array, 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.
history (jax.Array) – Stored pair memory with trailing shape history_shape(dim). Laws without memory return it unchanged.
advance_history (bool, optional) – Advance the stored memory over this step. When false, the law evaluates with the stored memory and returns it exactly.
quasistatic (bool, optional) – No time passes during this evaluation. state.vel, state.ang_vel, and the domain’s image rates then encode a finite displacement divided by system.dt (see jaxdem.colliders.advance_history()). The law must switch off every velocity-dependent force, and its memory update must depend on velocity only through velocity * system.dt, so the result is independent of dt.

Returns:

(force, torque, history).

Return type:

Tuple[jax.Array, jax.Array, jax.Array]

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

Compute the pair potential energy.

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.
history (jax.Array) – Stored pair memory. Laws without memory ignore it.

Returns:

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

Return type:

jax.Array

search_radii(state: State, system: System) → jax.Array[source]#

Return conservative radii for a single search snapshot.

A nonzero pair interaction must fit within the sum of the two bounds. Custom finite-range laws must override this method to use accelerated colliders; the naive all-pairs collider does not require a bound.

history_shape(dim: int) → tuple[int, ...][source]#

Trailing shape of one pair’s history; stateless laws use (0,).

Composite laws flatten child histories internally and restore this shape before invoking the child law.

init_history(pair_shape: tuple[int, ...], dim: int) → Array[source]#

Initialize the history variables for this force model.

Parameters:

pair_shape (tuple[int, ...]) – Collider-specific pair axes, such as a flat contact pool or a dense (N, N) pair array.
dim (int) – Spatial dimension used by dimension-dependent history layouts.

Returns:

A zero-filled array with shape pair_shape + history_shape(dim).

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.

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`	Finite-range Lennard-Jones interactions with shifted pair energy.
`router`	Force model router selecting laws based on species pairs.
`spring`	Linear spring force model.
`wca`	Purely repulsive Weeks-Chandler-Andersen pair interactions.
`wca_shifted`	Contact-start, force-shifted Weeks-Chandler-Andersen repulsion.