Note
Go to the end to download the full example code.
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).
LinearMinimizer and
RotationMinimizer are special integrators
that drive the system towards a potential-energy minimum instead of
performing physical time integration. Because they subclass
Integrator, they can be plugged into the
same slots on System.
Let’s see how to choose, configure, and swap them.
Linear vs Rotation Integrators#
Every integration step calls both integrators in sequence:
step_before_force()step_before_force()Force evaluation
step_after_force()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:
print("Linear integrators:", list(jdem.LinearIntegrator._registry.keys()))
# Available Rotation Integrators
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# The following rotation integrators are registered:
print("Rotation integrators:", list(jdem.RotationIntegrator._registry.keys()))
Linear integrators: ['', 'euler', 'langevin', 'verlet', 'verlet_rescaling', 'vicsek_extrinsic', 'vicsek_intrinsic', 'lineargradientdescent', 'linearfire', 'optax']
Rotation integrators: ['', 'spiral', 'verletspiral', 'rotationgradientdescent', 'rotationfire', 'optax']
See the API documentation of each class for constructor parameters and algorithmic details.
Deactivating an Integrator#
Passing an empty string "" as the integrator type 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="",
)
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="",
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="lineargradientdescent",
rotation_integrator_type="",
linear_integrator_kw={"learning_rate": 1e-4},
)
print("GD learning rate:", system.linear_integrator.learning_rate)
GD learning rate: 1e-04
Minimizers#
Minimizers are integrators that descend the potential energy landscape
instead of advancing physical time. They are registered in both the
minimizer and integrator registries, so you can select them in
create() the same way you select any
integrator.
system = jdem.System.create(
state.shape,
linear_integrator_type="linearfire",
rotation_integrator_type="rotationfire",
)
print("Linear minimizer:", type(system.linear_integrator).__name__)
print("Rotation minimizer:", type(system.rotation_integrator).__name__)
Linear minimizer: LinearFIRE
Rotation minimizer: RotationFIRE
The minimize() Routine#
JaxDEM provides a convenience function
minimize() that runs a
while_loop until the potential energy converges or a maximum step
count is reached.
from jaxdem.minimizers import minimize
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,
linear_integrator_type="linearfire",
rotation_integrator_type="",
)
state, system, steps, pe = 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 7 steps, PE = 0.000000e+00
The Integration Loop#
For reference, the full per-step sequence executed by
step() is:
apply()— enforce boundary conditionsstep_before_force()step_before_force()Collider + force manager — compute forces and torques
step_after_force()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.
Total running time of the script: (0 minutes 2.649 seconds)