Integrators and Minimizers

LinearIntegrator and RotationIntegrator advance the simulation state in time. A System holds one of each — one for translational degrees of freedom (position, velocity) and one for rotational degrees of freedom (orientation, angular velocity).

JaxDEM also supports energy minimization: instead of integrating the equations of motion, you can drive the system towards a potential-energy minimum. Minimizers are not integrator subclasses — they are optax gradient-transformation constructor functions (such as jaxdem.minimizers.fire()) passed to create() via minimizer / minimizer_kw, as described in the Minimizers section below.

Let’s see how to choose, configure, and swap them.

Linear vs Rotation Integrators

Every integration step calls both integrators in sequence:

1. step_before_force()

2. step_before_force()

3. Force evaluation

4. step_after_force()

5. step_after_force()

This split lets you mix and match: you could pair a Velocity Verlet linear integrator with a SPIRAL rotation integrator, or disable rotation entirely while keeping translation active.

import jax.numpy as jnp
import jaxdem as jdem

state = jdem.State.create(pos=jnp.zeros((1, 3)))
system = jdem.System.create(
    state.shape,
    linear_integrator_type="verlet",
    rotation_integrator_type="verletspiral",
)
print("Linear integrator:", type(system.linear_integrator).__name__)
print("Rotation integrator:", type(system.rotation_integrator).__name__)

Linear integrator: VelocityVerlet
Rotation integrator: VelocityVerletSpiral

Choosing an Integrator

Integrators are selected by their registered name when calling create(). Some integrators accept additional keyword arguments through linear_integrator_kw or rotation_integrator_kw. Consult the API reference for the specific parameters of each integrator.

Available Linear Integrators

The following linear integrators are registered:

# (the empty key ``""`` is a registered no-op; we filter it out)
print("Linear integrators:", sorted(k for k in jdem.LinearIntegrator._registry if k))

Linear integrators: ['euler', 'langevin', 'verlet', 'verletrescaling', 'vicsekextrinsic', 'vicsekintrinsic']

Available Rotation Integrators

The following rotation integrators are registered:

print(
    "Rotation integrators:", sorted(k for k in jdem.RotationIntegrator._registry if k)
)

Rotation integrators: ['spiral', 'verletspiral']

See the API documentation of each class for constructor parameters and algorithmic details.

Deactivating an Integrator

Passing None as the integrator type (the empty string "" is equivalent) selects the base no-op integrator, which leaves the corresponding degrees of freedom untouched. This is useful when you want to freeze translation or rotation.

# Freeze rotation — only translate:
system = jdem.System.create(
    state.shape,
    linear_integrator_type="verlet",
    rotation_integrator_type=None,
)
print("Rotation integrator (deactivated):", type(system.rotation_integrator).__name__)

Rotation integrator (deactivated): RotationIntegrator

Freeze translation — only rotate:

system = jdem.System.create(
    state.shape,
    linear_integrator_type=None,
    rotation_integrator_type="verletspiral",
)
print("Linear integrator (deactivated):", type(system.linear_integrator).__name__)

Linear integrator (deactivated): LinearIntegrator

Passing Constructor Arguments

Some integrators take additional parameters. Pass them via linear_integrator_kw or rotation_integrator_kw:

system = jdem.System.create(
    state.shape,
    linear_integrator_type="langevin",
    rotation_integrator_type=None,
    linear_integrator_kw={"gamma": 0.5, "temperature": 0.1, "k_B": 1.0},
)
print("Langevin gamma:", system.linear_integrator.gamma)

Langevin gamma: 0.5

Minimizers

Minimizers in JaxDEM are standard optax optimizers that descend the potential energy landscape. You configure them by passing the constructor function (such as jaxdem.fire, jaxdem.damped_newtonian, or standard optax optimizers like optax.adam) to minimizer, along with any keyword arguments in minimizer_kw.

Inside System.create, the optimizer is constructed and wrapped in a custom wrapper (CustomGradientTransformation) that keeps track of the constructor function and arguments for serialization.

For checkpoint serialization, JaxDEM saves the import path of the constructor function (e.g. "jaxdem.minimizers.fire" or "optax.adam") and the dictionary of keyword parameters. Upon restoration, it resolves and calls the function to recreate the optimizer.

We therefore recommend defining custom constructor functions and custom target functions (target_fn) in an importable module (not in __main__), so that system snapshots can be restored later. This is a recommendation, not a requirement: minimization itself works exactly the same with functions defined in your main script or a notebook — only saving and reloading the system through a checkpoint needs the import path to be resolvable.

Note

Minimizers and integrators are independent.

The minimizer and the linear/rotation integrators live in separate fields of the system and never touch each other’s configuration or internal state: minimize() does not advance time or use the integrators, and step() ignores the minimizer. You can first minimize and then integrate, integrate and then minimize, or alternate between them in any order without one affecting the other.

Note

Using Composite Optimizers (e.g., `optax.chain`)

If you want to use a composite optimizer like optax.chain, you cannot pass it directly as an instantiated object because the checkpoint writer does not support arbitrary nested object serialization.

Instead, define a simple wrapper function that constructs the chain. As above, the minimizer works no matter where this function is defined; placing it in a separate importable module (rather than your main script or __main__) is only needed so the checkpoint writer can save its import path and restore the minimizer upon loading.

# In an importable module, e.g., my_optimizers.py
import optax

def my_chained_optimizer(learning_rate=1e-3, max_grad_norm=1.0):
    return optax.chain(
        optax.clip_by_global_norm(max_grad_norm),
        optax.adam(learning_rate)
    )

# Then pass it to the system creation:
system = jdem.System.create(
    ...,
    minimizer=my_optimizers.my_chained_optimizer,
    minimizer_kw={"learning_rate": 1e-4, "max_grad_norm": 0.5}
)

system = jdem.System.create(
    state.shape,
    minimizer=jdem.fire,
    minimizer_kw={"dt": 1e-2},
)
print("Minimizer:", system.minimizer.type_name)

Minimizer: fire

The minimize() Routine

JaxDEM provides a convenience method minimize() that runs a while_loop until the potential energy converges or a maximum step count is reached.

state = jdem.State.create(
    pos=jnp.array([[0.0, 0.0], [1.5, 0.0]]),
    rad=jnp.array([1.0, 1.0]),
)
system = jdem.System.create(
    state.shape,
    minimizer=jdem.fire,
    minimizer_kw={"dt": 1e-2},
)

state, system, steps, pe = system.minimize(
    state, system, max_steps=500, pe_tol=1e-12, pe_diff_tol=1e-12
)
print(f"Converged in {steps} steps, PE = {pe:.6e}")

Converged in 1 steps, PE = 0.000000e+00

The Integration Loop

For reference, the full per-step sequence executed by step() is:

1. apply() — enforce boundary conditions

2. step_before_force()

3. step_before_force()

4. Collider + force manager — compute forces and torques

5. step_after_force()

6. step_after_force()

The step_before_force() / step_after_force() split lets multi-stage schemes (such as Velocity Verlet) position their updates around the force evaluation correctly.

Fixed Particles

Particles with state.fixed = True are immobile: the integrator multiplies velocity updates by (1 - fixed) so their velocity stays constant. See The Simulation State for how to set this field.