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 = 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.