jaxdem.forces.force_manager#

Utilities for managing per-particle force contributions.

Classes

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

Manage per-particle force contributions prior to pairwise interactions.
class jaxdem.forces.force_manager.ForceManager(gravity: jax.Array, external_force: jax.Array, external_torque: jax.Array, force_functions: Tuple[ForceFunction, ...] = ())[source]#

Bases: object

Manage per-particle force contributions prior to pairwise interactions.

gravity: jax.Array#

Constant acceleration applied to all particles. Shape (dim,).

external_force: jax.Array#

Accumulated external force applied to all particles. 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.

force_functions: Tuple[ForceFunction, ...]#

Optional tuple of callables with signature (state, system, i) returning per-particle force and torque contributions for particle i.

static create(dim: int, shape: Tuple, *, gravity: jax.Array | None = None, force_functions: Sequence[ForceFunction] = ()) → ForceManager[source][source]#

Create a ForceManager for a system with dim spatial dimensions.

Parameters:

  • dim – Spatial dimension of the managed system.

  • gravity – Optional initial gravitational acceleration. Defaults to zeros.

  • force_functions – Sequence of callables with signature (state, system, index) returning per-particle force/torque contributions.

static add_force(state: State, system: System, force: jax.Array, *, torque: jax.Array | None = None) → System[source][source]#

Accumulate an external force (and optionally 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.

  • force – Force contribution to accumulate. Must be a single (dim,) vector applied uniformly.

  • torque – Optional torque contribution to accumulate. Must be broadcast-compatible with external_torque.

Returns:

A new jaxdem.System instance with the updated accumulated forces and torques.

Return type:

System

static add_force_at(state: State, system: System, force: jax.Array, idx: jax.Array, *, torque: jax.Array | None = None) → System[source][source]#

Accumulate an external force (and optionally torque) to be applied on the next apply call over idx particles.

Parameters:

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

  • system (System) – Simulation system configuration.

  • force – Force contribution to accumulate. Must be an array broadcastable to external_force[idx].

  • idx – integer indices of the particles receiving the contribution.

  • torque – Optional torque contribution to accumulate. Must be an array broadcastable to external_torque[idx].

Returns:

A new jaxdem.System instance with the updated accumulated forces and torques.

Return type:

System

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

Overwrite state.force with managed per-particle contributions.

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]

Note

  • This method donates state and system