Skip to content

Storage format plugins

Register a BaseStorage subclass with polarism.results.storage.register_storage. Select it with cfg.result.formats = ("example-scalar-csv",). The tuple is combined with the built-in save_hdf5, save_npy, and save_json flags; each selected name is constructed once. Unknown names fail configuration validation and list the registered formats. The built-in names are hdf5, npy, and json.

The constructor contract is StorageClass(output_dir, batch_size, config, x, y). output_dir is the run's Path; config is the full Config; x and y are host coordinate arrays. Subclass BaseStorage and implement dump_batch, write_static, and finalize. Its add_to_batch(t, nodes, **context) buffers saved fields and scalars and calls dump_batch at the configured batch size. node_metadata holds each saved node's label and optional units. Override add_to_batch if a format supports only some node types.

Per-frame storage methods run on the background writer thread; construction and finalize run on the simulation thread. cached values and scalar_groups passed to add_to_batch are host NumPy arrays, including 0-D arrays for scalars. Flush the final batch in finalize. Storage errors propagate to the simulation thread.

"""Small CSV backend for scalar result nodes."""
from __future__ import annotations

import csv

from polarism.results.storage import BaseStorage, register_storage


@register_storage("example-scalar-csv")
class ScalarCSVStorage(BaseStorage):
    """Write scalar nodes to a single CSV file."""

    def __init__(self, output_dir, batch_size, config, x, y):
        """Set up scalar CSV storage."""
        super().__init__(output_dir, batch_size)
        self.path = output_dir / "scalars.csv"
        self.columns = None

    def add_to_batch(self, t, nodes, **context):
        """Keep only scalar nodes for the shared batch implementation."""
        scalars = [node for node in nodes if node.save and not node.is_field]
        super().add_to_batch(t, scalars, **context)

    def dump_batch(self):
        """Append buffered scalar rows."""
        if not self.time_buffer:
            return
        columns = list(self.scalar_buffers)
        if self.columns is None:
            self.columns = columns
        with self.path.open("a", newline="") as stream:
            writer = csv.writer(stream)
            if self.total_saved == 0:
                writer.writerow(["time", *self.columns])
            for index, t in enumerate(self.time_buffer):
                writer.writerow([t, *(self.scalar_buffers[name][index] for name in self.columns)])
        self.total_saved += len(self.time_buffer)
        self._clear_buffers()

    def write_static(self, name, field):
        """Ignore static fields in this scalar-only format."""

    def finalize(self):
        """Flush the remaining rows."""
        self.dump_batch()

This small example writes only saved scalar nodes, including built-in norm, to scalars.csv. The integration test loads it through plugin discovery and checks the output.