## `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. ``` python jax.config.update("jax_enable_x64", True) ``` ------------------------------------------------------------------------ source ### simulate ``` python 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`](https://nikolas-claussen.github.io/triangulax/src/simulation.html#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. ``` python import dataclasses ``` ``` python @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) ``` ``` python 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 ``` python # 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 ``` python # 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`](https://nikolas-claussen.github.io/triangulax/src/simulation.html#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 ``` python 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 ``` python # 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 ``` ``` python # 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](https://orbax.readthedocs.io/en/latest/) library (maintained by the JAX team) handles serialization of arbitrary JAX pytrees, including registered dataclasses like [`HeMesh`](https://nikolas-claussen.github.io/triangulax/src/halfedge_datastructure.html#hemesh) or `SimState` Alternatively, the [`HeMesh`](https://nikolas-claussen.github.io/triangulax/src/halfedge_datastructure.html#hemesh) and [`GeomMesh`](https://nikolas-claussen.github.io/triangulax/src/halfedge_datastructure.html#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. ``` python import tempfile, os import orbax.checkpoint as ocp ``` ``` python # 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 ``` python # 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: ``` python from triangulax.triangular import TriMesh from triangulax.mesh import HeMesh ``` ``` python 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