jaxdem.writers.vtk_writer#

Implementation of the high-level VTKWriter frontend.

Classes

VTKWriter(directory, save_every, clean, ...)

High-level front end for writing simulation data to VTK files.
class jaxdem.writers.vtk_writer.VTKWriter(directory: Path = PosixPath('frames'), save_every: int = 1, clean: bool = False, 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 treats trajectory_axis as time T and the other leading axes as batch axes, preserving their original order. Multiple batch axes are flattened in row-major order. The logical layout is (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.

Equivalent logical layouts are: - (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. A single System can be shared across batch snapshots; otherwise its leading axes must match the State. Trajectories require matching System axes to supply each frame’s time and step count. Mixed per-leaf System layouts are not inferred.

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 leading batch or trajectory axes and queues frames for background output. Independent grains are wrapped with their domain. A DeformableParticleModel (including plastic variants) retains its actual unwrapped coordinates: wrapping nodes independently would tear a boundary-crossing mesh and change its exported element measures and edge lengths. Such coherent meshes may extend beyond the primary cell.

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 leading axis in state/system to treat as time. Negative indices count from the last leading axis, before (N, dim).

  • 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 subclass’s 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.