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

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 | None] | tuple[ForceFunction, EnergyFunction | None, bool]] = ()) → ForceManager[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]#

Accumulate an external force to be applied on the next apply call for all particles.

Only the system is returned because the state is not modified: the force is buffered in the system’s ForceManager until apply() 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 be added to every particle (in the state’s current particle order).

  • is_com (bool, optional) – If True, force is applied to the Center of Mass (no induced torque); since the force is written to every clump member, each clump receives force in total, not force per member. 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]#

Add an external force to particles with array index idx.

Only the system is returned because the state is not modified: the force is buffered in the system’s ForceManager until apply() runs.

Parameters:

  • state (State) – Current state of the simulation.

  • system (System) – Simulation system configuration.

  • force (jax.Array) – External force to be added to particles with array index idx.

  • idx (jax.Array) – Array indices 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]#

Accumulate an external torque to be applied on the next apply call for all particles.

Only the system is returned because the state is not modified: the torque is buffered in the system’s ForceManager until apply() 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 be added to every particle (in the state’s current particle order); since the torque is written to every clump member, each clump receives torque in total, not torque per member.

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

Add an external torque to particles with array index idx.

Only the system is returned because the state is not modified: the torque is buffered in the system’s ForceManager until apply() runs.

Parameters:

  • state (State) – Current state of the simulation.

  • system (System) – Simulation system configuration.

  • torque (jax.Array) – External torque to be added to particles with array index idx.

  • idx (jax.Array) – Array indices of the particles affected by the external force.

static apply(state: State, system: System) → tuple[State, System][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]#

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:

Scalar JAX array representing the total potential energy.

Return type:

jax.Array