triangulax

`simulation`: Simulation utilities

This module provides lightweight utilities for running simulations with jax.lax.scan. It wraps the common pattern of stepping through time and logging measurements, and adds support for chunked execution with checkpointing callbacks.

When to use this module

For simple dynamics where each time step is just an ODE or gradient-descent step, you can use diffrax or optimistix directly — they already handle time-stepping, adaptive step sizes, and solver state internally. There is no need for an extra wrapper.

This module is useful when each time step involves more than just advancing an ODE. For example:

Topological modifications (T1 edge flips) that must happen between ODE steps, with cooldown logic
Coupled dynamics where different state components evolve under different solvers (e.g. ODE for positions + SDE for orientations)
Custom logging of derived quantities (energy, flip counts, full mesh snapshots) at every step
Long simulations that need periodic checkpointing to disk, or progress reporting

In these cases, you write a custom step_fn and measure_fn, and this module handles the scan wiring and chunked execution.

Simulation state and measurements

The simulation state (State) and log (Log) are user-defined objects — for example, a simple array, a collection of arrays, or a custom @dataclass. Custom classes need to be registered with JAX via jax.tree_util.register_dataclass. This module imposes no base classes, inheritance or protocol; it only requires that step_fn and measure_fn have the right signatures (see below). This keeps things JAX-idiomatic and flexible enough for anything from a gradient descent to a coupled ODE+SDE system with topology changes.

jax.config.update("jax_enable_x64", True)

source

simulate


def simulate(
    step_fn:Callable, # Advances the simulation state from the current time to the next time.
Signature: ``step_fn(state, next_time) -> new_state``.
    measure_fn:Callable, # Extracts measurements from the current state.
Signature: ``measure_fn(state) -> log``.
The returned pytree is stacked along axis 0 across all time steps.
    init:State, # Initial simulation state (any JAX pytree).
    timepoints:Float[Array, 'n_steps'], # Time points to step through. `step_fn` receives consecutive pairs.
)->tuple: # The simulation state after the last time step.

Run a simulation by scanning step_fn over timepoints, logging via measure_fn.

Wraps jax.lax.scan. Fully JIT- and vmap-compatible.

Test: damped harmonic oscillator

As a simple test, simulate the damped harmonic oscillator \(\ddot x = -x - 0.1 \dot x\) via forward Euler, and verify that simulate gives identical results to a manual jax.lax.scan. For illustration purposes, we create a dataclass to describe the current state (position \(x\), velocity \(v\), and time \(t\)) of the oscillator.

import dataclasses

@jax.tree_util.register_dataclass
@dataclasses.dataclass
class OscState:
    x: Float[jax.Array, ""]
    v: Float[jax.Array, ""]
    t: Float[jax.Array, ""]

@jax.tree_util.register_dataclass
@dataclasses.dataclass
class OscLog:
    x: Float[jax.Array, ""]
    energy: Float[jax.Array, ""]

gamma = 0.1

def osc_step(state: OscState, t_next: jax.Array) -> OscState:
    dt = t_next - state.t
    a = -state.x - gamma * state.v
    return OscState(x=state.x + dt * state.v,
                    v=state.v + dt * a,
                    t=t_next)

def osc_measure(state: OscState) -> OscLog:
    return OscLog(x=state.x,
                  energy=0.5 * state.x**2 + 0.5 * state.v**2)

dt = 0.01
timepoints = jnp.arange(0.0, 50.0, dt)
init = OscState(x=jnp.array(1.0), v=jnp.array(0.0), t=timepoints[0])

final_state, logs = simulate(osc_step, osc_measure, init, timepoints)
print(f"Final x={final_state.x:.4f}, energy decayed from {logs.energy[0]:.4f} to {logs.energy[-1]:.4f}")

Final x=0.0985, energy decayed from 0.5000 to 0.0054

# verify against manual scan
def _manual_scan_fn(state, t_next):
    new_state = osc_step(state, t_next)
    log = osc_measure(new_state)
    return new_state, log

final_ref, logs_ref = jax.lax.scan(_manual_scan_fn, init, timepoints)

assert jnp.allclose(logs.x, logs_ref.x)
assert jnp.allclose(logs.energy, logs_ref.energy)
assert jnp.allclose(final_state.x, final_ref.x)

JIT and vmap compatibility

# JIT
simulate_jit = jax.jit(simulate, static_argnums=(0, 1))
final_jit, logs_jit = simulate_jit(osc_step, osc_measure, init, timepoints)
assert jnp.allclose(logs_jit.x, logs.x)

# vmap over batch of initial conditions
init_batch = OscState(x=jnp.array([1.0, 2.0, 3.0]),
                      v=jnp.zeros(3),
                      t=jnp.zeros(3))

final_batch, logs_batch = jax.vmap(simulate, in_axes=(None, None, 0, None))(
    osc_step, osc_measure, init_batch, timepoints)

assert logs_batch.x.shape == (3, len(timepoints))
# first batch element should match the single-run result
assert jnp.allclose(logs_batch.x[0], logs.x)

Chunked simulation with checkpointing

jax.lax.scan compiles the entire loop, so there is no opportunity to run Python code (save to disk, print progress, etc.) mid-simulation. chunked_simulate splits the time points into chunks and runs jax.lax.scan on each chunk in a Python loop. Between chunks, an optional on_chunk callback is called with the current state, chunk logs, and chunk index — useful for checkpointing, logging, or early stopping.

Trade-off: Smaller chunks = more frequent callbacks, but each chunk incurs a small dispatch overhead.

source

chunked_simulate


def chunked_simulate(
    step_fn:Callable, # Advances the simulation state. Same as :func:[`simulate`](https://nikolas-claussen.github.io/triangulax/src/simulation.html#simulate).
    measure_fn:Callable, # Extracts measurements. Same as :func:[`simulate`](https://nikolas-claussen.github.io/triangulax/src/simulation.html#simulate).
    init:State, # Initial simulation state.
    timepoints:Float[Array, 'n_steps'], # Time points to step through.
    chunk_size:int, # Number of time steps per chunk.
    on_chunk:Union=None, # Callback invoked after each chunk with ``(state, chunk_logs, chunk_index)``.
Runs in Python (not JIT'd), so it can do I/O.
)->tuple: # The simulation state after the last time step.

Run a simulation in chunks, with an optional callback between chunks.

Each chunk is executed as a single jax.lax.scan. Between chunks, the on_chunk callback is called in Python, enabling checkpointing, progress reporting, or early stopping.

Test: chunked simulation matches single-pass

# chunked simulation with chunk_size=1000 should match single-pass exactly

chunk_log = []
def on_chunk_cb(state, logs, chunk_idx):
    chunk_log.append(chunk_idx)

final_chunked, logs_chunked = chunked_simulate(
    osc_step, osc_measure, init, timepoints, chunk_size=1000, on_chunk=on_chunk_cb)

assert jnp.allclose(logs_chunked.x, logs.x)
assert jnp.allclose(logs_chunked.energy, logs.energy)
assert jnp.allclose(final_chunked.x, final_state.x)
assert len(chunk_log) == 5  # 5000 steps / 1000

# test with non-divisible chunk size

final_chunked2, logs_chunked2 = chunked_simulate(
    osc_step, osc_measure, init, timepoints, chunk_size=1500)

assert logs_chunked2.x.shape == logs.x.shape
assert jnp.allclose(logs_chunked2.x, logs.x)
print(f"Non-divisible chunks: shape {logs_chunked2.x.shape}, results match.")

Non-divisible chunks: shape (5000,), results match.

Example: checkpointing to disk with orbax

A common use case for on_chunk is saving the simulation state to disk between chunks, so a long simulation can be resumed after a crash. The orbax-checkpoint library (maintained by the JAX team) handles serialization of arbitrary JAX pytrees, including registered dataclasses like HeMesh or SimState Alternatively, the HeMesh and GeomMesh classes have their own .save() / .load() methods that write to .npz archives.

Below is a minimal example using orbax. Note: orbax is an optional dependency — it is not required by triangulax itself.

import tempfile, os
import orbax.checkpoint as ocp

# Create a temporary directory for checkpoints
ckpt_dir = tempfile.mkdtemp()

# Define an on_chunk callback that saves the state after each chunk
checkpointer = ocp.StandardCheckpointer()

def checkpoint_callback(state, chunk_logs, chunk_idx):
    path = os.path.join(ckpt_dir, f"chunk_{chunk_idx:04d}")
    checkpointer.save(path, state)
    print(f"  Chunk {chunk_idx}: saved checkpoint to {path}")

# Run chunked simulation with checkpointing
final_ckpt, logs_ckpt = chunked_simulate(
    osc_step, osc_measure, init, timepoints, chunk_size=2500,
    on_chunk=checkpoint_callback)

assert jnp.allclose(logs_ckpt.x, logs.x)

  Chunk 0: saved checkpoint to /var/folders/vm/1jl6rjln6n9cjt54vsr9n4800000gr/T/tmpfj_ii7hd/chunk_0000
  Chunk 1: saved checkpoint to /var/folders/vm/1jl6rjln6n9cjt54vsr9n4800000gr/T/tmpfj_ii7hd/chunk_0001

# Restore from the last checkpoint and verify it matches the final state
restored = checkpointer.restore(os.path.join(ckpt_dir, "chunk_0001"), init)

assert jnp.allclose(restored.x, final_ckpt.x)
assert jnp.allclose(restored.v, final_ckpt.v)
print(f"Restored state matches final state: x={restored.x:.4f}")

# To resume a simulation from a checkpoint, pass the restored state as `init`:
# final_state, logs = simulate(step_fn, measure_fn, restored_state, remaining_timepoints)

Restored state matches final state: x=0.0985

You can also use orbax to save meshes, like so:

from triangulax.triangular import TriMesh
from triangulax.mesh import HeMesh

mesh = TriMesh.read_obj("../test_meshes/disk.obj")
hemesh = HeMesh.from_triangles(mesh.vertices.shape[0], mesh.faces)


ckpt_dir = tempfile.mkdtemp()

path_hemesh = os.path.join(ckpt_dir, "mesh-orbax")
checkpointer.save(path_hemesh, hemesh)

restored_hemesh = HeMesh(**checkpointer.restore(path_hemesh,))

hemesh == restored_hemesh

Warning: readOBJ() ignored non-comment line 3:
  o flat_tri_ecmc

True

simulation: Simulation utilities