jaxdem.writers#

Interface for defining data writers.

This module provides the high-level VTKWriter, H5Writer, 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.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.

Modules

`async_base`	Defines the base infrastructure for asynchronous data writing.
`checkpoints`	Orbax checkpoint writers and loaders for simulations and RL models.
`h5_writer`	Asynchronous HDF5 snapshots using the existing serialization format.
`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.