jaxdem.forces.force_manager#

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 outside 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 outside the collider.

After the collider runs, apply() adds these contributions to the state forces and aggregates them over 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). apply() clears this buffer.

external_force_com: jax.Array#

Accumulated external force applied to the center of mass (induces no torque). apply() clears this buffer.

external_torque: jax.Array#

Accumulated external torque applied to all particles. apply() clears this buffer.

is_com_force: tuple[bool, ...]#

Boolean array corresponding to force_functions with shape (n_forces,). If True, apply the force to the center of mass (no induced torque). If False, apply the force at the particle position (induces torque through the 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. Signature of ForceFunc: (pos, state, system) -> (Force, Torque). Signature of EnergyFunc: (pos, state, system) -> Energy. Supported formats:

    • 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]#

Buffer an external force on all particles for the next apply call.

The method returns only system. The state does not change because the force waits in the ForceManager buffer 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 add to every particle (in the state’s current particle order).

  • is_com (bool, optional) – If True, apply the force to the center of mass (no induced torque). The force goes to every clump member, so each clump receives force in total, not force per member. If False (default), apply the force at the particle position (induces torque).

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

Buffer an external force on the particles with array index idx for the next apply call.

The method returns only system. The state does not change because the force waits in the ForceManager buffer until apply() runs.

Parameters:

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

  • system (System) – Simulation system configuration.

  • force (jax.Array) – External force to add to the particles with array index idx.

  • idx (jax.Array) – Array indices of the particles the external force acts on.

  • is_com (bool, optional) – If True, apply the force to the center of mass (no induced torque). If False (default), apply the force at the particle position (induces torque).

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

Buffer an external torque on all particles for the next apply call.

The method returns only system. The state does not change because the torque waits in the ForceManager buffer 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 add to every particle (in the state’s current particle order). The torque goes to every clump member, so 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]#

Buffer an external torque on the particles with array index idx for the next apply call.

The method returns only system. The state does not change because the torque waits in the ForceManager buffer until apply() runs.

Parameters:

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

  • system (System) – Simulation system configuration.

  • torque (jax.Array) – External torque to add to the particles with array index idx.

  • idx (jax.Array) – Array indices of the particles the external torque acts on.

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

Add the managed per-particle contributions to the collider forces, then aggregate over clumps and broadcast back to the members.

Parameters:

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

  • system (System) – Simulation system configuration.

Returns:

The updated state and system.

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) – Current state of the simulation.

  • system (System) – Simulation system configuration.

Returns:

Scalar total potential energy.

Return type:

jax.Array