jaxdem.integrators

Time-integration interfaces and implementations.

Functions

free_mask(state)

Per-particle mask selecting non-fixed (free) particles.

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 semi-implicit (symplectic) Euler integration method.

The velocity is updated first and the position update then uses the new velocity, which makes the scheme symplectic (first-order accurate, but with bounded energy error), unlike a true forward Euler step.

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 semi-implicit (symplectic) Euler method.

The update equations are (note the position update uses the updated velocity):

\[\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

Langevin thermostat integrator for the translational degrees of freedom.

Integrates the underdamped Langevin equation

\[m\,\dot{\vec{v}} = \vec{F} - m \gamma \vec{v} + \sqrt{2 m \gamma k_B T}\, \vec{\eta}(t)\]

using the BAOAB splitting scheme (Leimkuhler & Matthews): half kick (B), half drift (A), exact Ornstein-Uhlenbeck update of the velocity (O), half drift (A), and a final half kick (B) after the force evaluation. The O-step samples the friction and Gaussian noise exactly, driving the system towards the canonical distribution at temperature \(T\).

Fixed particles keep their prescribed velocities and are not thermostatted.

Parameters:

• gamma (jax.Array) – Friction (collision) coefficient \(\gamma\) with units of inverse time; sets how strongly velocities are damped and rethermalized.

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

• temperature (jax.Array) – Target temperature \(T\) of the thermostat.

gamma: Array

k_B: Array

temperature: Array

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

Perform the BAOA part of the BAOAB step.

Applies a half kick with the current forces, a half drift, the exact Ornstein-Uhlenbeck velocity update

\[\vec{v} \leftarrow c_1 \vec{v} + c_2 \vec{\eta}, \qquad c_1 = e^{-\gamma \Delta t}, \qquad c_2 = \sqrt{\tfrac{k_B T}{m}\left(1 - e^{-2\gamma \Delta t}\right)}\]

with \(\vec{\eta} \sim \mathcal{N}(0, 1)\), and a second half drift. system.key is split to draw the noise. Fixed particles keep their prescribed velocities.

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]

Perform the final B (half kick) part of the BAOAB step.

Updates the velocities of free particles with the freshly computed forces: \(\vec{v} \leftarrow \vec{v} + \tfrac{\Delta t}{2} \vec{F}/m\).

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.

jaxdem.integrators.free_mask(state: State) → jax.Array

Per-particle mask selecting non-fixed (free) particles.

Returns a boolean array of shape (…, N, 1) that is True for free particles and False for fixed ones. Multiplying a (…, N, dim) update by this mask zeroes the update on fixed particles while leaving their prescribed velocities untouched. This is the single idiom integrators should use to respect jaxdem.State.fixed.

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), \vec{\tau}(t + \Delta t)) \\ & k_2 = \Delta t\; \dot{\vec{\omega}}(\vec{\omega}(t) + k_1, \vec{\tau}(t + \Delta t)) \\ & k_3 = \Delta t\; \dot{\vec{\omega}}(\vec{\omega}(t) + (k_1 + k_2)/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)\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:

\[v(t + \Delta t) = v(t + \Delta t / 2) + \Delta t a(t + \Delta t) / 2\]

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: VelocityVerlet

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, inherited from VelocityVerlet).

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.

Fixed particles are excluded from the thermostat statistics (kinetic energy sums and drift mean) and their prescribed velocities are never modified.

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

classmethod Create(temperature: float = 1.0, k_B: float = 1.0, rescale_every: int = 1, can_rotate: bool = False, subtract_drift: bool = False) → VelocityVerletRescaling

Create the thermostat from plain Python values.

Parameters:

• temperature (float, default 1.0) – Target temperature.

• k_B (float, default 1.0) – Boltzmann constant (1.0 for reduced units).

• rescale_every (int, default 1) – Rescale velocities every this many steps.

• can_rotate (bool, default False) – Include rotational DOF in the thermostat.

• subtract_drift (bool, default False) – Remove centre-of-mass drift before rescaling.

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:

\[\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) + k_1, \vec{\tau}(t)) \\ & k_3 = \Delta t/2\; \dot{\vec{\omega}}(\vec{\omega}(t) + (k_1 + k_2)/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 + \Delta t/2), \vec{\tau}(t + \Delta t)) \\ & k_2 = \Delta t/2\; \dot{\vec{\omega}}(\vec{\omega}(t + \Delta t/2) + k_1, \vec{\tau}(t + \Delta t)) \\ & k_3 = \Delta t/2\; \dot{\vec{\omega}}(\vec{\omega}(t + \Delta t/2) + (k_1 + k_2)/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:

• 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 (semi-implicit / symplectic Euler) 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).