jaxdem.integrators

Time-integration interfaces and implementations.

Functions

free_mask(state)

Compute the per-particle mask that selects the free (non-fixed) particles.

Classes

Integrator()	Abstract base class that defines 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

Semi-implicit (symplectic) Euler integration method.

The method updates the velocity first. The position update then uses the new velocity. This makes the scheme symplectic: first-order accurate with bounded energy error, unlike a true forward Euler step.

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

Advance the simulation state by one time step after the force calculation with the semi-implicit (symplectic) Euler method.

The update equations are (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.

Integrate the underdamped Langevin equation

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

The scheme is the BAOAB splitting (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. This drives the system toward 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. It controls how strongly the thermostat damps and rethermalizes the velocities.

• 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.

Apply 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. The method splits system.key 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.

Update the velocities of the free particles with the new 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) subclass this to register with the Factory and to show that they operate on linear kinematics.

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

Compute the per-particle mask that selects the free (non-fixed) 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. Their prescribed velocities stay untouched. Integrators use this mask 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) subclass this to register with the Factory and to show that 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 method updates the quaternion with 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)}\]

The angular velocity and its derivative are purely imaginary quaternions. The scalar part is zero and the vector part equals the vector. The exponential map of a purely imaginary quaternion is

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

SSPRK3 then updates the angular velocity:

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

The angular velocity derivative depends on the torque and the 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

Velocity Verlet integration method.

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

Advance the simulation state by one half-step before the force calculation with 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.

Return type:

Tuple[State, System]

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

Advance the simulation state by one half-step after the force calculation with 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.

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 integrator uniformly rescales the translational (and optionally rotational) velocities so the instantaneous kinetic temperature matches temperature. Between rescalings the dynamics are purely Newtonian (standard Velocity Verlet, inherited from VelocityVerlet).

The integrator applies the rescaling at the end of step_after_force, after the second Verlet half-kick. The terminal velocities on rescaling steps are then exactly at the target temperature.

The thermostat statistics (kinetic energy sums and drift mean) exclude fixed particles. The rescaling never modifies their prescribed velocities.

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, the thermostat also rescales the angular velocities and counts rotational kinetic energy toward the temperature.

• subtract_drift (jax.Array) – Whether to remove center-of-mass drift before rescaling (0 or 1). When 1, the integrator subtracts the center-of-mass velocity on rescaling steps before it measures the 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 center-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]

Advance the angular state by one half-step before the force calculation with the spiral leapfrog scheme.

A third-order Runge–Kutta scheme (SSPRK3) integrates the rigid-body angular momentum equations in the principal axis frame. The method updates the quaternion with the spiral leapfrog algorithm.

• SPIRAL algorithm:

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

The angular velocity and its derivative are purely imaginary quaternions. The scalar part is zero and the vector part equals the vector. The exponential map of a purely imaginary quaternion is

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

SSPRK3 then updates the angular velocity:

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

The angular velocity derivative depends on the torque and the 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.

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

Advance the angular velocity by one half-step after the force calculation with the spiral leapfrog scheme.

A third-order Runge–Kutta scheme (SSPRK3) integrates the rigid-body angular momentum equations in the principal axis frame. This half-kick uses the torques computed at the new positions. The quaternion does not change here.

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

The angular velocity derivative depends on the torque and the 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.

• 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.

Each step, the integrator sets the translational velocity magnitude to v0. The direction comes from 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

• The integrator draws the noise per clump (one sample per rigid body) and broadcasts it to all clump members, so clumps move coherently.

• A collider may cache the neighbor list (e.g., NeighborList collider) or 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. It applies a random rotation to the normalized base direction. It does not add a random vector in force space as the extrinsic (vectorial-noise) variant does.

The base direction comes from: - the current accumulated force vector (from colliders + force functions), - the average neighbor velocity direction (including self).

The integrator then applies the noise per clump and broadcasts it 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).