jaxdem.writers#

Interface for defining data writers.

This module provides the high-level VTKWriter and CheckpointWriter front ends. It also provides the VTKBaseWriter plugin interface and concrete writers (e.g., VTKSpheresWriter, VTKDomainWriter) that export JAX-based simulation snapshots to files.

Classes

VTKBaseWriter()

Abstract base class for writers that output simulation data.

class jaxdem.writers.CheckpointLoader(directory: Path | str = PosixPath('checkpoints'))#

Bases: BaseCheckpointManager

Thin wrapper around Orbax checkpoint restoring for jaxdem.state and jaxdem.system.

load(step: int | None = None, *, strict: bool = True) → tuple[State, System][source]#

Restore a checkpoint.

Parameters:

step (Optional[int]) –
- If None, load the latest checkpoint.
- Otherwise, load the specified step.
strict (bool, optional) – If True (default), raise a RuntimeError when the loader cannot re-import a custom force function recorded in the checkpoint. If False, skip force functions that do not load, with a warning.

Returns:

The restored State and System.

Return type:

Tuple[State, System]

latest_step() → int | None[source]#

class jaxdem.writers.CheckpointModelLoader(directory: Path | str = PosixPath('checkpoints'))#

Bases: BaseCheckpointManager

Thin wrapper around Orbax checkpoint restoring for jaxdem.rl.models.Model.

load(step: int | None = None) → Model[source]#: Load the model at the given step. If step is None, load the latest checkpoint.

latest_step() → int | None[source]#

class jaxdem.writers.CheckpointModelWriter(directory: Path | str = PosixPath('checkpoints'), max_to_keep: int | None = None, save_every: int = 1, clean: bool = False)#

Bases: BaseCheckpointManager

Thin wrapper around Orbax checkpoint saving for jaxdem.rl.models.Model.

max_to_keep: int | None = None#: Keep the last max_to_keep checkpoints. If None, keep all checkpoints.

save_every: int = 1#: How often to write. The writer saves on every save_every-th call to save().

clean: bool = False#: If True, erase and recreate the target directory on construction. If False (the default), keep existing checkpoints in the directory, so a resumed run does not destroy earlier checkpoints.

save(model: Model, step: int) → None[source]#

Save the model at the given step.

Stores model_state and JSON metadata. model.metadata must contain JSON-serializable fields. The writer adds model_type to the metadata.

class jaxdem.writers.CheckpointWriter(directory: Path | str = PosixPath('checkpoints'), max_to_keep: int | None = None, save_every: int = 1, clean: bool = False)#

Bases: BaseCheckpointManager

Thin wrapper around Orbax checkpoint saving.

Notes

The writer serializes custom force functions passed via force_manager_kw by their fully-qualified module path (e.g. mypackage.forces.trap). A different script cannot restore functions defined in the top-level script (__main__). The writer emits a warning at save time if any force function lives in __main__. To keep checkpoints portable, define force functions in an importable module.

max_to_keep: int | None = None#: Keep the last max_to_keep checkpoints. If None, keep all checkpoints.

save_every: int = 1#: How often to write. The writer saves on every save_every-th call to save().

clean: bool = False#: If True, erase and recreate the target directory on construction. If False (the default), keep existing checkpoints in the directory, so a resumed run does not destroy earlier checkpoints.

save(state: State, system: System) → None[source]#

Save a checkpoint of the given state and system at the current step.

Parameters:

state (State) – The current state of the simulation.
system (System) – The current system configuration.

class jaxdem.writers.VTKBaseWriter#

Bases: Factory, ABC

Abstract base class for writers that output simulation data.

Concrete subclasses implement the write method to convert a snapshot (jaxdem.State, jaxdem.System pair) into a specific file format.

Example:#

To define a custom VTK writer, inherit from VTKBaseWriter and implement its abstract methods:

>>> @VTKBaseWriter.register("my_custom_vtk_writer")
>>> @dataclass(slots=True)
>>> class MyCustomVTKWriter(VTKBaseWriter):
        ...

classmethod is_active(state: State, system: System) → bool[source]#: Check whether this writer has data to write for the given state and system.

abstractmethod classmethod write(state: State, system: System, filename: Path, binary: bool) → None[source]#

Write information from a simulation snapshot to a VTK PolyData file.

Concrete writers implement this method. The caller converts all JAX arrays to NumPy arrays before it calls write.

Parameters:

state (State) – The simulation jaxdem.State snapshot to write.
system (System) – The simulation jaxdem.System configuration.
filename (Path) – Target path of the VTK file. The caller guarantees that the parent directory exists.
binary (bool) – If True, write the VTK file in binary mode. If False, write it in ASCII (human-readable) mode.

class jaxdem.writers.VTKDeformableEdgeAdjacenciesWriter#

Bases: VTKBaseWriter

A VTKBaseWriter that writes the shared hinges between adjacent elements.

In 3D the writer draws each shared edge as a VTK line. In 2D it draws each shared vertex as a VTK point. It attaches the rest and current bending angles and the bending stiffness eb as CellData arrays.

classmethod is_active(state: State, system: System) → bool[source]#: Check whether this writer has data to write for the given state and system.

classmethod write(state: State, system: System, filename: Path, binary: bool) → None[source]#

Write information from a simulation snapshot to a VTK PolyData file.

Concrete writers implement this method. The caller converts all JAX arrays to NumPy arrays before it calls write.

Parameters:

state (State) – The simulation jaxdem.State snapshot to write.
system (System) – The simulation jaxdem.System configuration.
filename (Path) – Target path of the VTK file. The caller guarantees that the parent directory exists.
binary (bool) – If True, write the VTK file in binary mode. If False, write it in ASCII (human-readable) mode.

class jaxdem.writers.VTKDeformableEdgesWriter#

Bases: VTKBaseWriter

A VTKBaseWriter that writes deformable particle edges as VTK lines.

It attaches the rest and current edge lengths and the edge stiffness el as CellData arrays.

classmethod is_active(state: State, system: System) → bool[source]#: Check whether this writer has data to write for the given state and system.

classmethod write(state: State, system: System, filename: Path, binary: bool) → None[source]#

Write information from a simulation snapshot to a VTK PolyData file.

Concrete writers implement this method. The caller converts all JAX arrays to NumPy arrays before it calls write.

Parameters:

state (State) – The simulation jaxdem.State snapshot to write.
system (System) – The simulation jaxdem.System configuration.
filename (Path) – Target path of the VTK file. The caller guarantees that the parent directory exists.
binary (bool) – If True, write the VTK file in binary mode. If False, write it in ASCII (human-readable) mode.

class jaxdem.writers.VTKDeformableElementsWriter#

Bases: VTKBaseWriter

A VTKBaseWriter that writes deformable particle elements.

The writer draws 3-vertex elements as VTK triangles and 2-vertex elements as VTK lines. It attaches per-element properties (elements_id, ec, gamma, rest and current element measures, partial_content, and element_normals) as CellData arrays.

classmethod is_active(state: State, system: System) → bool[source]#: Check whether this writer has data to write for the given state and system.

classmethod write(state: State, system: System, filename: Path, binary: bool) → None[source]#

Write information from a simulation snapshot to a VTK PolyData file.

Concrete writers implement this method. The caller converts all JAX arrays to NumPy arrays before it calls write.

Parameters:

state (State) – The simulation jaxdem.State snapshot to write.
system (System) – The simulation jaxdem.System configuration.
filename (Path) – Target path of the VTK file. The caller guarantees that the parent directory exists.
binary (bool) – If True, write the VTK file in binary mode. If False, write it in ASCII (human-readable) mode.

class jaxdem.writers.VTKDomainWriter#

Bases: VTKBaseWriter

A VTKBaseWriter that writes the simulation domain as a VTK geometric primitive.

The writer draws the domain as an axis-aligned cuboid (3D) or rectangle (2D) with a vtkCubeSource. It pads 2D input arrays to 3D as VTK requires.

classmethod write(state: State, system: System, filename: Path, binary: bool) → None[source]#

Write information from a simulation snapshot to a VTK PolyData file.

Concrete writers implement this method. The caller converts all JAX arrays to NumPy arrays before it calls write.

Parameters:

state (State) – The simulation jaxdem.State snapshot to write.
system (System) – The simulation jaxdem.System configuration.
filename (Path) – Target path of the VTK file. The caller guarantees that the parent directory exists.
binary (bool) – If True, write the VTK file in binary mode. If False, write it in ASCII (human-readable) mode.

class jaxdem.writers.VTKSpheresWriter#

Bases: VTKBaseWriter

A VTKBaseWriter that writes particle centers as VTK points.

The writer skips facet vertices (facet_id != -1). It attaches per-particle State fields (e.g., vel, rad, mass) and quaternion components as PointData arrays. It pads positions and 2-component vectors to 3D as VTK requires.

classmethod is_active(state: State, system: System) → bool[source]#: Check whether this writer has data to write for the given state and system.

classmethod write(state: State, system: System, filename: Path, binary: bool) → None[source]#

Write information from a simulation snapshot to a VTK PolyData file.

Concrete writers implement this method. The caller converts all JAX arrays to NumPy arrays before it calls write.

Parameters:

state (State) – The simulation jaxdem.State snapshot to write.
system (System) – The simulation jaxdem.System configuration.
filename (Path) – Target path of the VTK file. The caller guarantees that the parent directory exists.
binary (bool) – If True, write the VTK file in binary mode. If False, write it in ASCII (human-readable) mode.

class jaxdem.writers.VTKFacetSpheresWriter#

Bases: VTKBaseWriter

A VTKBaseWriter that writes facet vertices (facet_id != -1) as VTK points.

classmethod is_active(state: State, system: System) → bool[source]#: Check whether this writer has data to write for the given state and system.

classmethod write(state: State, system: System, filename: Path, binary: bool) → None[source]#

Write information from a simulation snapshot to a VTK PolyData file.

Concrete writers implement this method. The caller converts all JAX arrays to NumPy arrays before it calls write.

Parameters:

state (State) – The simulation jaxdem.State snapshot to write.
system (System) – The simulation jaxdem.System configuration.
filename (Path) – Target path of the VTK file. The caller guarantees that the parent directory exists.
binary (bool) – If True, write the VTK file in binary mode. If False, write it in ASCII (human-readable) mode.

class jaxdem.writers.VTKFacetsWriter#

Bases: VTKBaseWriter

A VTKBaseWriter that writes facets as thickened VTK cells.

The writer groups vertices by facet_id. In 2D it expands each 2-vertex facet into a convex-hull polygon that covers the facet plus its thickness, or a quad when the hull is not available. In 3D it expands each 3-vertex facet into a convex-hull triangle mesh, or a flat prism when the hull is not available. It also attaches per-vertex State fields and quaternion components as PointData arrays.

classmethod is_active(state: State, system: System) → bool[source]#: Check whether this writer has data to write for the given state and system.

classmethod write(state: State, system: System, filename: Path, binary: bool) → None[source]#

Write information from a simulation snapshot to a VTK PolyData file.

Concrete writers implement this method. The caller converts all JAX arrays to NumPy arrays before it calls write.

Parameters:

state (State) – The simulation jaxdem.State snapshot to write.
system (System) – The simulation jaxdem.System configuration.
filename (Path) – Target path of the VTK file. The caller guarantees that the parent directory exists.
binary (bool) – If True, write the VTK file in binary mode. If False, write it in ASCII (human-readable) mode.

class jaxdem.writers.VTKWriter(directory: Path = PosixPath('frames'), save_every: int = 1, clean: bool = True, max_workers: int = 8, max_queue_size: int = 512, writers: list[str] = <factory>, binary: bool = True)#

Bases: BaseAsyncWriter

High-level front end for writing simulation data to VTK files.

This class converts JAX-based jaxdem.State and jaxdem.System pytrees into VTK files. It handles batches, trajectories, and dispatch to registered jaxdem.VTKBaseWriter subclasses.

How leading axes are interpreted#

Let particle positions have shape (..., N, dim), where N is the number of particles and dim is 2 or 3. Define L = state.pos_c.ndim - 2, i.e., the number of leading axes before (N, dim).

L == 0 — single snapshot
The input is one frame. It is written directly into frames/batch_00000000/ (no batching, no trajectory).
trajectory=False (default)
The writer treats all leading axes as batch axes (not time). If multiple batch axes exist, the writer flattens them into a single batch axis: (B, N, dim) with B = prod(shape[:L]). The writer writes each batch b as a single snapshot under its own subdirectory frames/batch_XXXXXXXX/. No trajectory is implied.
- Example: (B, N, dim) → B separate directories with one frame each.
- Example: (B1, B2, N, dim) → flatten to (B1*B2, N, dim) and treat as above.
trajectory=True
The writer swaps the axis given by trajectory_axis to the front (axis 0) and treats it as time T. Any remaining leading axes are batch axes. If more than one non-time leading axis exists, the writer flattens them into a single batch axis. The data becomes (T, B, N, dim) with B = prod(other leading axes).
- If there is only time (L == 1): (T, N, dim) — a single batch
  directory frames/batch_00000000/ contains a time series with T frames.
- If there is time plus batching (L >= 2): (T, B, N, dim) — each
  batch b gets its own directory frames/batch_XXXXXXXX/ containing a time series (T frames) for that batch.

After these swaps/reshapes, dispatch is: - (N, dim) → single snapshot - (B, N, dim) → batches (no time) - (T, N, dim) → single batch with a trajectory - (T, B, N, dim) → per-batch trajectories

Concrete writers receive per-frame NumPy arrays. The writer slices and broadcasts System leaves to match the current frame and batch.

writers: list[str]#: Names of the registered VTKBaseWriter subclasses to use for writing. If empty, use all registered subclasses. Name matching follows registry keys: case-insensitive, and spaces, underscores, and hyphens are ignored. The spelling given here sets the output file and .pvd names.

binary: bool = True#: If True, write VTK files in binary format. If False, write files in ASCII format.

save(state: State, system: System, *, trajectory: bool = False, trajectory_axis: int = 0, batch0: int = 0) → None[source]#

Schedule writing of a jaxdem.State / jaxdem.System pair to VTK files.

This public entry point interprets the leading axes as batch or trajectory axes, swaps and flattens axes as needed, and pushes the data to the background writer queue.

Parameters:

state (State) – The simulation jaxdem.State object to save.
system (System) – The jaxdem.System object corresponding to state.
trajectory (bool, optional) – If True, interpret trajectory_axis as time.
trajectory_axis (int, optional) – The axis in state/system to treat as the trajectory axis.
batch0 (int, optional) – The starting batch index for the input data.

directory: Path#: The root directory where the writer saves simulation frames.

save_every: int#: Save frequency. The writer pushes a frame to the queue on the first call and on every save_every-th call to the save() method.

clean: bool#: If True, the writer deletes and recreates directory on initialization. Safety checks prevent deleting the current working directory or the system root.

max_workers: int#: The number of background worker threads to use for parallel I/O.

max_queue_size: int#: Maximum number of pending tasks in the background queue. When the queue is full, submit() blocks until a worker frees a slot. This backpressure keeps memory bounded when the simulation outruns disk I/O. Set to 0 for an unbounded queue.

Modules

`async_base`	Defines the base infrastructure for asynchronous data writing.
`checkpoints`	Orbax checkpoint writers and loaders for simulations and RL models.
`vtk_deformable_particle_writer`	VTK writers for deformable particles.
`vtk_domain_writer`	VTK writer for domain geometry.
`vtk_facets_writer`	VTK writer for facet geometry.
`vtk_spheres_writer`	VTK writers that export particle centers as VTK points.
`vtk_writer`	Implementation of the high-level VTKWriter frontend.