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.