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.

jaxdem.integrators.free_mask(state: State) → jax.Array[source]#

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.Integrator#

Bases: Factory, ABC

Abstract base class that defines the interface for time-stepping.

Example

To define a custom integrator, inherit from Integrator and implement its abstract methods:

>>> @Integrator.register("myCustomIntegrator")
... @jax.tree_util.register_dataclass
... @dataclass(slots=True)
... class MyCustomIntegrator(Integrator):
...     ...

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

Advance the simulation state before the force evaluation.

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][source]#

Advance the simulation state after the force evaluation.

System runs the rotational hook before the linear hook, so linear thermostats can use both terminal 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 initialize(state: State, system: System) → tuple[State, System][source]#

Initialize the integrator.

Called by System.initialize after initial forces and torques have been computed. Override this hook for integrator-specific setup, such as a backward half-kick for staggered velocities. The default returns the state and system unchanged. Initialization must not advance the simulation clock or perform an ordinary position/orientation step.

Parameters:

state (State) – Current state of the simulation.
system (System) – Simulation system configuration.

Returns:

The state and system after initialization.

Return type:

Tuple[State, System]

Example

>>> state, system = system.initialize(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.

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.

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.