jaxdem.forces.force_manager#

Utilities for managing external and custom force contributions that do not depend on the collider.

Functions

default_energy_func(pos, state, system)

Classes

ForceManager(gravity, external_force, ...[, ...])

Manage custom force contributions external to the collider.

class jaxdem.forces.force_manager.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, ...] = ())[source]#

Bases: object

Manage custom force contributions external to the collider. It also accumulates forces in the state after collider application, accounting for 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). This buffer is cleared when apply() is invoked.

external_force_com: jax.Array#: Accumulated external force applied to Center of Mass (does not induce torque). 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.

is_com_force: Tuple[bool, ...]#: Boolean array corresponding to force_functions with shape (n_forces,). If True, the force is applied to the Center of Mass (no induced torque). If False, the force is applied to the constituent particle (induces torque via 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] | Tuple[ForceFunction, EnergyFunction, bool]] = ()) → ForceManager[source][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. Supported formats:
- ForceFunc: Applied at particle, no potential energy.
- (ForceFunc, bool): Boolean specifies if it is a COM force.
- (ForceFunc, EnergyFunc): Includes potential energy function.
- (ForceFunc, EnergyFunc, bool): Includes energy and COM specifier.
Signature of ForceFunc: (pos, state, system) -> (Force, Torque) Signature of EnergyFunc: (pos, state, system) -> Energy

Supported formats for force_functions items: - 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][source]#

Accumulate an external force 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 (jax.Array) – External force to be added to all particles in the current order.
is_com (bool, optional) – If True, force is applied to Center of Mass (no induced torque). If False (default), force is applied to Particle Position (induces torque).

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

Add an external force to particles with ID=idx.

Parameters:

state (State) – Current state of the simulation.
system (System) – Simulation system configuration.
force (jax.Array) – External force to be added to particles with ID=idx.
idx (jax.Array) – ID of the particles affected by the external force.
is_com (bool, optional) – If True, force is applied to Center of Mass (no induced torque). If False (default), force is applied to Particle Position (induces torque).

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

Accumulate an external 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.
torque (jax.Array) – External torque to be added to all particles in the current order..

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

Add an external torque to particles with ID=idx.

Parameters:

state (State) – Current state of the simulation.
system (System) – Simulation system configuration.
torque (jax.Array) – External torque to be added to particles with ID=idx.
idx (jax.Array) – ID of the particles affected by the external force.

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

Accumulate managed per-particle contributions on top of collider/contact forces, then perform final clump aggregation + broadcast.

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]

static compute_potential_energy(state: State, system: System) → jax.Array[source][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) – The current state of the simulation.
system (System) – The configuration of the simulation.

Returns:

A scalar JAX array representing the total potential energy of each particle.

Return type:

jax.Array