jaxdem.writers.vtkWriter#

Implementation of the high-level VTKWriter frontend.

Classes

VTKWriter(writers, directory, binary, clean, ...)

High-level front end for writing simulation data to VTK files.
class jaxdem.writers.vtkWriter.VTKWriter(writers: ~typing.List[str] = <factory>, directory: ~pathlib.Path = PosixPath('frames'), binary: bool = True, clean: bool = True, save_every: int = 1, max_queue_size: int = 512, max_workers: int | None = None, _counter: int = 0, _writer_classes: ~typing.List = <factory>, _manifest: ~typing.Dict = <factory>)[source]#

Bases: object

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.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 None, all available VTKBaseWriter subclasses will be used.

directory: Path#

The base directory where output VTK files will be saved. Subdirectories might be created within this path for batched outputs. Defaults to “frames”.

binary: bool#

If True, VTK files will be written in binary format. If False, files will be written in ASCII format. Defaults to True.

clean: bool#

If True, the directory will be completely emptied before any files are written. Defaults to True. This is useful for starting a fresh set of output frames.

save_every: int#

How often to write; writes on every save_every-th call to save().

max_queue_size: int#

The maximum number of scheduled writes allowed. 0 means unbounded.

max_workers: int | None#

Maximum number of worker threads for the internal thread pool.

close()[source][source]#

Flush all pending tasks and shut down the internal thread pool. Safe to call multiple times.

block_until_ready()[source][source]#

Wait until all scheduled writer tasks complete.

This will wait for all pending futures, propagate exceptions (if any), and clear the pending set.

save(state: State, system: System, *, trajectory: bool = False, trajectory_axis: int = 0, batch0: int = 0)[source][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 writes the resulting frames using the registered writers. It also creates per-batch ParaView .pvd collections referencing the generated files.

Parameters:

  • state (State) – The simulation jaxdem.State object to be saved. Its array leaves must end with (N, dim).

  • system (System) – The jaxdem.System object corresponding to state. Leading axes must be compatible (or broadcastable) with those of state.

  • trajectory (bool, optional) – If True, interpret trajectory_axis as time and write a trajectory; if False, interpret the leading axis as batch.

  • trajectory_axis (int, optional) – The axis in state/system to treat as the trajectory (time) axis when trajectory=True. This axis is swapped to the front prior to writing.

  • batch0 (in) – Initial value of batch from where to start counting the batches.