jaxdem.integrators

Time-integration interfaces and implementations.

Classes

Integrator()	Abstract base class for defining the interface for time-stepping.
LinearIntegrator()	Namespace for translation/linear-time integrators.
RotationIntegrator()	Namespace for rotation/angular-time integrators.

class jaxdem.integrators.DirectEuler

Bases: LinearIntegrator

Implements the explicit (forward) Euler integration method.

static step_after_force(state: State, system: System) → tuple[State, System]

Advances the simulation state by one time step after the force calculation using the Direct Euler method.

The update equations are:

\[\begin{split}& v(t + \\Delta t) &= v(t) + \\Delta t a(t) \\\\ & r(t + \\Delta t) &= r(t) + \\Delta t v(t + \\Delta t)\end{split}\]

where:

• \(r\) is the particle position (jaxdem.State.pos)

• \(v\) is the particle velocity (jaxdem.State.vel)

• \(a\) is the particle acceleration computed from forces (jaxdem.State.force)

• \(\\Delta t\) is the time step (jaxdem.System.dt)

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]

class jaxdem.integrators.Langevin(gamma: Array, k_B: Array, temperature: Array)

Bases: LinearIntegrator

TODO: add details

Parameters:

• gamma (jax.Array) – Friction coefficient.

• k_B (jax.Array) – Boltzmann constant (set to 1.0 for reduced units).

• temperature (jax.Array) – Target temperature \(T\).

gamma: Array

k_B: Array

temperature: Array

static step_before_force(state: State, system: System) → tuple[State, System]

TODO: add details

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

TODO: add details

Parameters:

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

• system (System) – Simulation system configuration.

Returns:

The updated state and system.

Return type:

Tuple[State, System]

class jaxdem.integrators.LinearIntegrator

Bases: Integrator

Namespace for translation/linear-time integrators.

Purpose

Groups integrators that update linear state (e.g., position and velocity). Concrete methods (e.g., DirectEuler) should subclass this to register via the Factory and to signal they operate on linear kinematics.

class jaxdem.integrators.RotationIntegrator

Bases: Integrator

Namespace for rotation/angular-time integrators.

Purpose

Groups integrators that update angular state (e.g., orientation, angular velocity). Concrete methods (e.g., DirectEulerRotation) should subclass this to register via the Factory and to signal they operate on rotational kinematics.

class jaxdem.integrators.Spiral

Bases: RotationIntegrator

Non-leapfrog spiral integrator for angular velocities.

The implementation follows the velocity update described in del Valle et al. (2023).

static step_after_force(state: State, system: System) → tuple[State, System]

Advance angular velocities by a single time step.

A third-order Runge–Kutta scheme (SSPRK3) integrates the rigid-body angular momentum equations in the principal axis frame. The quaternion is updated based on the spiral non-leapfrog algorithm.

• SPIRAL algorithm:

\[q(t + \Delta t) = q(t) \cdot e^\left(\frac{\Delta t}{2}\omega\right) \cdot e^\left(\frac{\Delta t^2}{4}\dot{\omega}\right)\]

Where the angular velocity and its derivative are purely imaginary quaternions (scalar part is zero and the vector part is equal to the vector). The exponential map of a purely imaginary quaternion is

\[e^u = \cos(|u|) + \frac{\vec{u}}{|u|}\sin(|u|)\]

Angular velocity is then updated using SSPRK3:

\[\begin{split}& \vec{\omega}(t + \Delta t) = \vec{\omega}(t) + \frac{1}{6}(k_1 + k_2 + 4k_3) \\ & k_1 = \Delta t\; \dot{\vec{\omega}}(\vec{\omega}(t + \Delta t / 2), \vec{\tau}(t + \Delta t)) \\ & k_2 = \Delta t\; \dot{\vec{\omega}}(\vec{\omega}(t + \Delta t / 2) + k1, \vec{\tau}(t + \Delta t)) \\ & k_3 = \Delta t\; \dot{\vec{\omega}}(\vec{\omega}(t + \Delta t / 2) + (k1 + k2)/4, \vec{\tau}(t + \Delta t)) \\\end{split}\]

Where the angular velocity derivative is a function of the torque and angular velocity:

\[\dot{\vec{\omega}} = (\tau - \vec{\omega} \times (I \vec{\omega}))I^{-1}\]

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]

Reference

del Valle et. al, SPIRAL: An efficient algorithm for the integration of the equation of rotational motion, https://doi.org/10.1016/j.cpc.2023.109077.

class jaxdem.integrators.VelocityVerlet

Bases: LinearIntegrator

Implements the Velocity Verlet integration method.

static step_before_force(state: State, system: System) → tuple[State, System]

Advances the simulation state by one half-step before the force calculation using the Velocity Verlet scheme.

The update equations are:

\[\begin{split}& v(t + \\Delta t / 2) &= v(t) + \\Delta t a(t) / 2 \\\\ & r(t + \\Delta t) &= r(t) + \\Delta t v(t + \\Delta t / 2) \\Delta t\end{split}\]

where:

• \(r\) is the particle position (jaxdem.State.pos)

• \(v\) is the particle velocity (jaxdem.State.vel)

• \(a\) is the particle acceleration computed from forces (jaxdem.State.force)

• \(\\Delta t\) is the time step (jaxdem.System.dt)

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

Advances the simulation state by one half-step after the force calculation using the Velocity Verlet scheme.

The update equations are:

\[\begin{split}& v(t + \\Delta t) &= v(t + \\Delta t / 2) + \\Delta t a(t) / 2 \\\\\end{split}\]

where:

• \(v\) is the particle velocity (jaxdem.State.vel)

• \(a\) is the particle acceleration computed from forces (jaxdem.State.force)

• \(\\Delta t\) is the time step (jaxdem.System.dt)

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]

class jaxdem.integrators.VelocityVerletRescaling(k_B: Array, temperature: Array, rescale_every: Array, can_rotate: Array, subtract_drift: Array)

Bases: LinearIntegrator

Velocity Verlet with periodic velocity-rescaling thermostat.

Every rescale_every steps, the translational (and optionally rotational) velocities are uniformly rescaled so that the instantaneous kinetic temperature matches temperature. Between rescalings the dynamics are purely Newtonian (standard Velocity Verlet).

The rescaling is applied at the end of step_after_force, after the second Verlet half-kick, so that the terminal velocities on rescaling steps are exactly at the target temperature.

Parameters:

• k_B (jax.Array) – Boltzmann constant (set to 1.0 for reduced units).

• temperature (jax.Array) – Target temperature \(T\).

• rescale_every (jax.Array) – Rescale velocities every this many steps (scalar integer).

• can_rotate (jax.Array) – Whether to include rotational DOF in the thermostat (0 or 1). When 1, angular velocities are rescaled together with linear velocities and rotational KE counts toward the temperature.

• subtract_drift (jax.Array) – Whether to remove centre-of-mass drift before rescaling (0 or 1). When 1, the COM velocity is subtracted on rescaling steps before measuring temperature. Mainly relevant for small systems.

k_B: Array

temperature: Array

rescale_every: Array

can_rotate: Array

subtract_drift: Array

static step_before_force(state: State, system: System) → tuple[State, System]

static step_after_force(state: State, system: System) → tuple[State, System]

class jaxdem.integrators.VelocityVerletSpiral

Bases: RotationIntegrator

Leapfrog spiral integrator for angular velocities adapted to Velocity Verlet.

The implementation follows the velocity update described in del Valle et al. (2023).

static step_before_force(state: State, system: System) → tuple[State, System]

Advances the simulation state by one half-step before the force calculation using the Velocity Verlet scheme.

A third-order Runge–Kutta scheme (SSPRK3) integrates the rigid-body angular momentum equations in the principal axis frame. The quaternion is updated with the spiral leapfrog algorithm to implement a Velocity Verlet-like method.

• SPIRAL algorithm:

\[q(t + \Delta t) = q(t) \cdot e^\left(\frac{\Delta t}{2}\omega(t + \Delta t/2)\right)\]

Where the angular velocity and its derivative are purely imaginary quaternions (scalar part is zero and the vector part is equal to the vector). The exponential map of a purely imaginary quaternion is

\[e^u = \cos(|u|) + \frac{\vec{u}}{|u|}\sin(|u|)\]

Angular velocity is then updated using SSPRK3 which we:

\[\begin{split}& \vec{\omega}(t + \Delta t/2) = \vec{\omega}(t) + \frac{1}{6}(k_1 + k_2 + 4k_3) \\ & k_1 = \Delta t/2\; \dot{\vec{\omega}}(\vec{\omega}(t), \vec{\tau}(t)) \\ & k_2 = \Delta t/2\; \dot{\vec{\omega}}(\vec{\omega}(t) + k1, \vec{\tau}(t)) \\ & k_3 = \Delta t/2\; \dot{\vec{\omega}}(\vec{\omega}(t) + (k1 + k2)/4, \vec{\tau}(t)) \\\end{split}\]

Where the angular velocity derivative is a function of the torque and angular velocity:

\[\dot{\vec{\omega}} = (\tau - \vec{\omega} \times (I \vec{\omega}))I^{-1}\]

Parameters:

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

• system (System) – Simulation system configuration.

Returns:

• Tuple[State, System] – The updated state and system after one time step.

• Reference

• ———–

• del Valle et. al, SPIRAL (An efficient algorithm for the integration of the equation of rotational motion, https://doi.org/10.1016/j.cpc.2023.109077.)

static step_after_force(state: State, system: System) → tuple[State, System]

Advances the simulation state by one half-step after the force calculation using the Velocity Verlet scheme.

A third-order Runge–Kutta scheme (SSPRK3) integrates the rigid-body angular momentum equations in the principal axis frame. The quaternion is updated with the spiral leapfrog algorithm to implement a Velocity Verlet-like method.

\[\begin{split}& \vec{\omega}(t + \Delta t) = \vec{\omega}(t + \Delta t/2) + \frac{1}{6}(k_1 + k_2 + 4k_3) \\ & k_1 = \Delta t/2\; \dot{\vec{\omega}}(\vec{\omega}(t), \vec{\tau}(t)) \\ & k_2 = \Delta t/2\; \dot{\vec{\omega}}(\vec{\omega}(t) + k1, \vec{\tau}(t)) \\ & k_3 = \Delta t/2\; \dot{\vec{\omega}}(\vec{\omega}(t) + (k1 + k2)/4, \vec{\tau}(t)) \\\end{split}\]

Where the angular velocity derivative is a function of the torque and angular velocity:

\[\dot{\vec{\omega}} = (\tau - \vec{\omega} \times (I \vec{\omega}))I^{-1}\]

Parameters:

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

• system (System) – Simulation system configuration.

Returns:

• Tuple[State, System] – The updated state and system after one time step.

• Reference

• ———–

• del Valle et. al, SPIRAL (An efficient algorithm for the integration of the equation of rotational motion, https://doi.org/10.1016/j.cpc.2023.109077.)

class jaxdem.integrators.VicsekExtrinsic(neighbor_radius: Array, eta: Array, v0: Array, max_neighbors: int)

Bases: LinearIntegrator

Vicsek-model integrator with extrinsic (vectorial) noise.

This integrator implements a Vicsek-like update rule by directly setting the translational velocity magnitude to v0 each step, based on the direction of a vector that combines:

• the current accumulated force vector (from colliders + force functions),

• the average neighbor velocity direction (including self),

• an additive random unit vector scaled by eta (extrinsic noise).

Notes

• Noise is generated per clump (one sample per rigid body) and then broadcast to all clump members so clumps move coherently.

• Neighbor lists may be cached (e.g., NeighborList collider) or may sort the state (e.g., some cell-list builders). This integrator uses the returned state from create_neighbor_list for consistency.

neighbor_radius: Array

eta: Array

v0: Array

max_neighbors: int

static step_after_force(state: State, system: System) → tuple[State, System]

class jaxdem.integrators.VicsekIntrinsic(neighbor_radius: Array, eta: Array, v0: Array, max_neighbors: int)

Bases: LinearIntegrator

Vicsek-model integrator with intrinsic noise.

This variant perturbs the direction of the desired motion by applying a random rotation to the normalized base direction (rather than adding a random vector in force space as in the extrinsic / vectorial-noise variant).

The base direction is computed from: - the current accumulated force vector (from colliders + force functions), - the average neighbor velocity direction (including self), then noise is applied per clump and broadcast to all clump members so clumps move coherently.

neighbor_radius: Array

eta: Array

v0: Array

max_neighbors: int

static step_after_force(state: State, system: System) → tuple[State, System]

Modules

direct_euler	Direct (forward) Integrator.
langevin	Langevin Integrator
spiral	Angular-velocity integrator based on the spiral scheme.
velocity_verlet	Velocity Verlet Integrator.
velocity_verlet_rescaling	Velocity Verlet integrator with periodic velocity-rescaling thermostat.
velocity_verlet_spiral	Angular-velocity integrator based on the spiral scheme.
vicsek	Vicsek-style integrators (extrinsic and intrinsic noise).