jaxdem.writers.async_base#

Defines the base infrastructure for asynchronous data writing.

Classes

BaseAsyncWriter([directory, save_every, ...])

Infrastructure for non-blocking JAX data writing.

class jaxdem.writers.async_base.BaseAsyncWriter(directory: Path = PosixPath('frames'), save_every: int = 1, clean: bool = True, max_workers: int = 4)#

Bases: object

Infrastructure for non-blocking JAX data writing.

This class provides a pool of background worker threads and a task queue to ensure that slow disk I/O operations and device-to-host transfers do not block the main simulation loop.

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.

submit(func: Callable[[...], Any], *args: Any, **kwargs: Any) → None[source]#: Pushes a task to the background worker queue.

close() → None[source]#: Blocks the main thread until all pending writes are finished and shuts down the background threads.

block_until_ready() → None[source]#: Waits until all currently pending tasks in the queue are completed.