jaxdem.writers.vtkWriter#

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.vtkWriter.VTKWriter(directory: Path = PosixPath('frames'), save_every: int = 1, clean: bool = True, max_workers: int = 4, writers: list[str] = <factory>, binary: bool = True)#

Bases: BaseAsyncWriter

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

This class orchestrates the conversion of JAX-based jaxdem.State and jaxdem.System pytrees into VTK files, handling 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)

    All leading axes are treated as batch axes (not time). If multiple batch axes are present, they are flattened into a single batch axis: (B, N, dim) with B = prod(shape[:L]). Each batch b is written 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 axis given by trajectory_axis is swapped to the front (axis 0) and interpreted as time T. Any remaining leading axes are batch axes. If more than one non-time leading axis exists, they are flattened into a single batch axis so 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; leaves in System are sliced/broadcast consistently with the current frame/batch.

writers: list[str]#

A list of strings specifying which registered VTKBaseWriter subclasses should be used for writing. If empty, all available subclasses will be used.

binary: bool = True#

If True, VTK files will be written in binary format. If False, files will be written 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 axes (batch vs. trajectory), performs any required axis swapping and flattening, and then pushes the data to the background writer queue.

Parameters:

  • state (State) – The simulation jaxdem.State object to be saved.

  • 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 simulation frames will be saved.

save_every: int#

Frequency of saving. A frame is pushed to the queue every save_every calls to the save() method.

clean: bool#

If True, the directory is deleted and recreated upon initialization. Basic safety checks are performed to 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.