Guide for adding a new sink to PhysicsNeMo Curator. Covers discovery questions, implementation patterns (simple writer, append-based, split-based), output naming, parallel partitioning, testing, and registration.
Installer avec Codex ou Claude Copiez ce prompt, collez-le dans Codex, Claude ou un autre assistant, puis laissez-le vérifier la page du skill et l'installer pour vous.
Une commande directe contourne le prompt de vérification. Examinez la source avant de l'exécuter.
Guide for adding a new sink to PhysicsNeMo Curator. Covers discovery questions, implementation patterns (simple writer, append-based, split-based), output naming, parallel partitioning, testing, and registration.
Adding a New Sink
This skill walks through adding a new Sink to PhysicsNeMo Curator,
from initial design through implementation, testing, and registration.
Step 0: Gather Requirements
Before writing any code, answer these questions. Ask the user if anything
is unclear.
Domain
Which submodule does this sink belong to?
Domain
Type parameter
Submodule
Dependency group
mesh
Sink["Mesh"]
src/physicsnemo_curator/domains/mesh/
mesh (physicsnemo, pyvista, pyarrow, torch)
da
Sink["xr.DataArray"]
src/physicsnemo_curator/domains/da/
da (xarray, earth2studio, zarr)
atm
Sink["AtomicData"]
src/physicsnemo_curator/domains/atm/
mesh (nvalchemi, torch)
Sink Design
Output format — What file format will the sink produce? (tensordict
memmap, Zarr, NetCDF4, HDF5, Parquet, VTK, NumPy, custom)
Naming strategy — How are output files/directories named?
Index-based: Use pipeline index (e.g. ). Simplest approach.
mesh_0001_0
Template-based: Use naming_template with placeholders like
{index}, {seq}, {relpath}, {stem}, {run_id}, {mesh_name}.
Data-driven: Use coordinate metadata from the data itself (e.g.
variable names, time values). Used when data carries its own identity.
Append or overwrite? — When the same output path is hit twice, should
the sink append to the existing file or overwrite it?
Splitting? — Should the output be split along a dimension? (e.g.
one file per variable, one file per year, one file per run)
Chunking/compression? — Does the format support chunking or
compression? Should these be configurable?
Parameters — What user-configurable parameters does the sink need?
(output directory, chunk sizes, compression level, format options, etc.)
Dependencies — Does it require additional imports beyond the domain
group? (e.g. netCDF4, h5py, zarr, specific I/O libraries)
Streaming? — Can items be written one-at-a-time as the iterator is
consumed, or does the sink need to buffer all items first?
Parallel safety — Does the sink need to coordinate concurrent writes?
(e.g. multiple workers writing to the same Zarr chunk). If so, implement
partition_indices().
Source integration — Does the sink need metadata from the source for
output naming? (directory mirroring, run IDs, etc.) If so, implement
set_source().
# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES.# SPDX-FileCopyrightText: All rights reserved.# SPDX-License-Identifier: Apache-2.0## Licensed under the Apache License, Version 2.0 (the "License");# you may not use this file except in compliance with the License.# You may obtain a copy of the License at## http://www.apache.org/licenses/LICENSE-2.0## Unless required by applicable law or agreed to in writing, software# distributed under the License is distributed on an "AS IS" BASIS,# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.# See the License for the specific language governing permissions and# limitations under the License.
Import Block Template
"""Brief description of what this sink writes."""from __future__ import annotations
import pathlib
import time
from typing import TYPE_CHECKING, ClassVar
from physicsnemo_curator.core.base import Param, Sink
from physicsnemo_curator.core.logging import get_logger
if TYPE_CHECKING:
from collections.abc import Iterator
from physicsnemo.mesh import Mesh # or: import xarray as xr# or: from physicsnemo_curator.domains.atm.types import AtomicData
Sink Class Template
class <ClassName>(Sink["Mesh"]):
"""Write :class:`~physicsnemo.mesh.Mesh` objects to <format>.
<2-3 sentence description of what the sink writes, how it names
outputs, and any append/split behaviour.>
Parameters
----------
output_dir : str
Directory where outputs will be written.
naming_template : str or None, optional
Format string for output names. Placeholders: ``{index}``,
``{seq}``, ``{relpath}``, ``{stem}``. Default: ``mesh_{index:04d}_{seq}``.
Examples
--------
>>> sink = <ClassName>(output_dir="./output/") # doctest: +SKIP
>>> paths = sink(mesh_iterator, index=0) # doctest: +SKIP
"""
name: ClassVar[str] = "<Display Name>"
description: ClassVar[str] = "<Short one-line description for CLI>" @classmethoddefparams(cls) -> list[Param]:
"""Return parameter descriptors for this sink.
Returns
-------
list[Param]
Parameter descriptors.
"""return [
Param(
name="output_dir",
description="Output directory for files",
type=str,
),
Param(
name="naming_template",
description=(
"Format string for output names. ""Placeholders: {index}, {seq}, {relpath}, {stem}. ""Default: mesh_{index:04d}_{seq}"
),
type=str,
default=None,
),
]
def__init__(
self,
output_dir: str,
naming_template: str | None = None,
) -> None:
"""Initialize the sink.
Parameters
----------
output_dir : str
Output directory for files.
naming_template : str or None, optional
Format string for output names.
"""self._output_dir = pathlib.Path(output_dir)
self._naming_template = naming_template
self._source: Source[Mesh] | None = Noneself._log = get_logger(self)
defset_source(self, source: Source[Mesh]) -> None:
"""Inject source reference for placeholder resolution.
Called automatically by :meth:`Pipeline.write`.
Parameters
----------
source : Source[Mesh]
The upstream source providing items.
"""self._source = source
def__call__(self, items: Iterator[Mesh], index: int) -> list[str]:
"""Consume items and write each to storage.
Parameters
----------
items : Iterator[Mesh]
Stream of data items to persist.
index : int
Source index (used for naming output files).
Returns
-------
list[str]
Paths of the files written.
"""self._log.info("idx_%d: Starting write", index)
t0 = time.perf_counter()
self._output_dir.mkdir(parents=True, exist_ok=True)
paths: list[str] = []
for seq, item inenumerate(items):
out_path = self._resolve_path(index, seq)
self._write_item(item, out_path)
paths.append(str(out_path))
elapsed = time.perf_counter() - t0
self._log.info("idx_%d: Wrote %d items (%.2fs)", index, len(paths), elapsed)
return paths
def_resolve_path(self, index: int, seq: int) -> pathlib.Path:
"""Resolve output path from naming template.
Parameters
----------
index : int
Pipeline source index.
seq : int
Sequence number within the index.
Returns
-------
pathlib.Path
Resolved output path.
"""ifself._naming_template isNone:
name = f"mesh_{index:04d}_{seq}"else:
kwargs: dict[str, object] = {"index": index, "seq": seq}
ifself._source isnotNone:
ifhasattr(self._source, "relative_path"):
rel = self._source.relative_path(index)
kwargs["relpath"] = str(pathlib.PurePosixPath(rel).parent)
kwargs["stem"] = pathlib.PurePosixPath(rel).stem
name = self._naming_template.format(**kwargs)
returnself._output_dir / name
def_write_item(self, item: Mesh, path: pathlib.Path) -> None:
"""Write a single item to disk.
Parameters
----------
item : Mesh
Data item to write.
path : pathlib.Path
Output file/directory path.
"""# Implementation here
...
@propertydefoutput_dir(self) -> pathlib.Path:
"""Return the output directory path."""returnself._output_dir
Key Design Decisions
__call__ receives Iterator[T] and index: int, returns list[str].
This is the Sink[T] ABC contract. Key differences from filters:
Sinks consume the iterator (no yielding back)
Sinks return paths of written files as list[str]
The index parameter comes from the source (pipeline position)
Sinks create output directories as needed (mkdir(parents=True, exist_ok=True))
Logging
Use the structured component logger from physicsnemo_curator.core.logging:
The most sophisticated pattern. Supports two execution modes:
Sequential mode (default): uses a third-party writer library with
batched append semantics. One store per index.
Pre-allocated parallel mode: allocates the full Zarr store upfront
at construction time, then workers write to non-overlapping offsets
with zero locking.
If the sink supports appending to existing files, implement a two-branch
write method:
def_append_to_store(self, da: xr.DataArray, path: pathlib.Path) -> None:
"""Append data to existing store, or create new one."""
ds = da.to_dataset(name="data")
if path.exists():
# Append along the appropriate dimension
ds.to_zarr(store=str(path), mode="a", append_dim="time")
else:
path.parent.mkdir(parents=True, exist_ok=True)
ds.to_zarr(store=str(path), mode="w", encoding=self._encoding)
For NetCDF4 (which doesn't support native append), use load + concat + rewrite:
Override this method when the sink has constraints on concurrent writes.
The pipeline runner uses it to group indices that must be processed by
the same worker.
defpartition_indices(self, indices: list[int]) -> list[list[int]] | None:
"""Group indices into partitions for same-worker processing.
Each returned group is a list of indices that must be handled
sequentially by a single worker. The runner never splits a
group across workers.
Parameters
----------
indices : list[int]
The indices to partition.
Returns
-------
list[list[int]] | None
Partitioned groups, or ``None`` if no partitioning required.
"""returnNone# Default: no constraints
Chunk-Aligned Partitioning (ZarrSink pattern)
When multiple indices map to the same Zarr chunk, they must be written
by the same worker to avoid corruption:
from collections import defaultdict
defpartition_indices(self, indices: list[int]) -> list[list[int]] | None:
"""Group indices by Zarr chunk alignment."""
chunk_size = self._chunks.get(self._append_dim, 1)
if chunk_size <= 1:
returnNone# One index per chunk = no constraints
groups: dict[int, list[int]] = defaultdict(list)
for idx in indices:
chunk_id = idx // chunk_size
groups[chunk_id].append(idx)
return [sorted(group) for _, group insorted(groups.items())]
When the store is pre-allocated and the partition map is computed at
construction time:
defpartition_indices(self, indices: list[int] | None = None) -> list[list[int]] | None:
"""Return pre-computed chunk groups for parallel writes."""ifnotself._parallel:
returnNoneif indices isNone:
returnself._chunk_groups # Computed in _preallocate()# Filter to requested indices
idx_set = set(indices)
return [
[i for i in group if i in idx_set]
for group inself._chunk_groups
ifany(i in idx_set for i in group)
]
Step 1c: Source Integration (Optional)
set_source() Method
Implement this when the sink needs metadata from the source for output
path resolution (directory mirroring, run IDs, etc.). It is called
automatically by Pipeline.write().
defset_source(self, source: Source[Mesh]) -> None:
"""Inject source reference for placeholder resolution.
Called automatically by :meth:`Pipeline.write` when the sink is
attached to a pipeline. Use this to resolve naming template
placeholders that depend on source metadata.
Parameters
----------
source : Source[Mesh]
The upstream source providing items.
"""self._source = source
Source methods available for naming:
Method
Returns
Use case
source.relative_path(index)
str
Directory mirroring (relpath, stem)
source.run_id()
str
Group outputs by run identifier
source.mesh_name(index, seq)
str
Per-mesh naming from source metadata
Param Declaration Patterns
Required parameter (no default — user must provide):
Param(name="output_dir", description="Output directory for files", type=str)
Param(
name="flip_triangle_normals",
description="Reverse triangle vertex order for VTK normal convention",
type=bool,
default=True,
)
Property Accessors
Expose key configuration as read-only properties for testing and
introspection:
@propertydefoutput_dir(self) -> pathlib.Path:
"""Return the output directory path."""returnself._output_dir
@propertydefcompression_level(self) -> int:
"""Return the configured compression level."""returnself._compression_level
@propertydefnaming_template(self) -> str | None:
"""Return the configured naming template."""returnself._naming_template
Iterator Consumption Patterns
All sinks fully consume their input iterator before returning.
Pattern 1: Sequential Enumeration (most common):
for seq, mesh inenumerate(items):
# Process each item immediately
paths.append(str(self._write(mesh, index, seq)))
Pattern 2: Batch Collection (buffered writes):
batch: list[T] = []
for item in items:
batch.append(item)
iflen(batch) >= self._batch_size:
self._flush_batch(batch)
batch = []
if batch:
self._flush_batch(batch) # Don't forget remainder
Pattern 3: Stateful Accumulation (data-driven):
for da in items:
written = self._write_dataarray(da)
paths.extend(written)
return paths
Step 2: Write Tests
Create the test file at test/domains/<domain>/test_<name>.py.
Test File Structure
"""Tests for <ClassName>."""# SPDX header (same as source files)from __future__ import annotations
import pathlib
import pytest
pytestmark = pytest.mark.requires("<domain>")
# ---------------------------------------------------------------------------# Test helpers# ---------------------------------------------------------------------------def_create_test_data(...) -> ...:
"""Create minimal test data for unit tests.
Use small sizes (10-50 points, 3-5 time steps).
"""
...
Test Classes
Unit Tests (Metadata)
classTest<ClassName>Unit:
"""Metadata tests."""deftest_params_list(self) -> None:
from <module> import <ClassName>
params = <ClassName>.params()
assertlen(params) > 0
names = [p.name for p in params]
assert"output_dir"in names # or "output_path"deftest_name_and_description(self) -> None:
from <module> import <ClassName>
assertisinstance(<ClassName>.name, str)
assertlen(<ClassName>.name) > 0assertlen(<ClassName>.description) > 0deftest_properties(self) -> None:
from <module> import <ClassName>
sink = <ClassName>(output_dir="/tmp/test")
assert sink.output_dir == pathlib.Path("/tmp/test")
deftest_roundtrip(self, tmp_path: pathlib.Path) -> None:
"""Written data can be read back and matches original."""
original = _create_test_data()
sink = <ClassName>(output_dir=str(tmp_path / "out"))
paths = sink(iter([original]), index=0)
# Read back and verify
loaded = _read_output(paths[0])
# Assert data matches original
Append Tests (if applicable)
deftest_append_to_existing(self, tmp_path: pathlib.Path) -> None:
"""Sink appends to existing file rather than overwriting."""
sink = <ClassName>(output_dir=str(tmp_path / "out"))
# Write first batch
paths_1 = sink(iter([_create_test_data(time_start=0)]), index=0)
# Write second batch (same output path)
paths_2 = sink(iter([_create_test_data(time_start=1)]), index=1)
# Verify the file contains both batches
...
Split Tests (if applicable)
deftest_splits_by_variable(self, tmp_path: pathlib.Path) -> None:
"""Data with variable dimension is split into separate outputs."""
da = _create_test_data(variables=["temperature", "pressure"])
sink = <ClassName>(output_dir=str(tmp_path / "out"))
paths = sink(iter([da]), index=0)
assertlen(paths) == 2# Verify separate outputs for each variable
@pytest.mark.e2eclassTest<ClassName>Pipeline:
"""End-to-end pipeline tests."""deftest_in_pipeline(self, tmp_path: pathlib.Path) -> None:
"""Sink works correctly at the end of a pipeline."""from physicsnemo_curator.domains.mesh.sources.vtk import VTKSource
# or: from physicsnemo_curator.domains.da.sources.era5 import ERA5Source# Create source data
...
sink = <ClassName>(output_dir=str(tmp_path / "output"))
pipeline = source.write(sink)
for i inrange(len(pipeline)):
paths = pipeline[i]
assertlen(paths) >= 1for p in paths:
assert pathlib.Path(p).exists()
Registry Tests
classTest<ClassName>Registry:
"""Test that the sink is registered."""deftest_sink_registered(self) -> None:
from physicsnemo_curator.core.registry import registry
names = [s.name for s in registry.list_sinks("<domain>")]
assert"<Display Name>"in names
Test Markers Reference
Marker
Purpose
pytestmark = pytest.mark.requires("mesh")
Module-level: skip all if mesh deps missing
@pytest.mark.requires("da")
Skip if da dependencies not installed
@pytest.mark.requires("atm")
Skip if atm dependencies not installed
@pytest.mark.integration
Tests that touch filesystem
@pytest.mark.e2e
End-to-end pipeline tests
@pytest.mark.slow
Slow tests, excluded from quick CI
Running Tests
# Unit and write tests (fast, no network)
uv run pytest test/domains/<domain>/test_<name>.py -v
# Pipeline integration test
uv run pytest test/domains/<domain>/test_<name>.py -v -k "Pipeline"# Full domain test suite for regressions
uv run pytest test/domains/<domain>/ -v -k "not slow"
Step 3: Register the Sink
Edit the domain __init__.py
For mesh sinks, edit src/physicsnemo_curator/domains/mesh/__init__.py:
For da sinks, edit src/physicsnemo_curator/domains/da/__init__.py with
the same pattern using "da" as the submodule name.
For atm sinks, edit src/physicsnemo_curator/domains/atm/__init__.py
with "atm" as the submodule name.
Lazy Loading Pattern (Optional Heavy Dependencies)
If the sink requires heavy optional dependencies (pyvista, zarr, etc.),
use lazy loading in the domain sinks __init__.py:
__all__ = ["MeshSink", "MeshVTUSink", "MeshZarrSink"]
def__getattr__(name: str):
"""Lazy-load optional sinks to avoid import at module load."""if name == "MeshZarrSink":
from physicsnemo_curator.domains.mesh.sinks.mesh_zarr import MeshZarrSink
return MeshZarrSink
if name == "MeshVTUSink":
from physicsnemo_curator.domains.mesh.sinks.mesh_vtu import MeshVTUSink
return MeshVTUSink
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
Step 4: Quality Checks
Run all checks before committing:
# Format
uv run ruff format \
src/physicsnemo_curator/domains/<domain>/sinks/<name>.py \
test/domains/<domain>/test_<name>.py
# Lint
uv run ruff check --fix \
src/physicsnemo_curator/domains/<domain>/sinks/<name>.py \
test/domains/<domain>/test_<name>.py
# Docstring coverage (must be >= 99%)
uv run interrogate
# Type checking
uv run ty check
# Run sink tests
uv run pytest test/domains/<domain>/test_<name>.py -v
# Run full domain test suite for regressions
uv run pytest test/domains/<domain>/ -v -k "not slow"
Step 5: Commit
Use the commit tool (or follow Conventional Commits format):
feat(<domain>): add <ClassName> for <format> output
<Optional body describing the output format, naming, and features.>
Checklist
Before considering the sink complete, verify:
Core
SPDX license headers on all new files
Sink inherits from Sink["Mesh"], Sink["xr.DataArray"], or Sink["AtomicData"]
name and description ClassVars are set
params() returns all configurable parameters
__call__() receives Iterator[T] and index: int, returns list[str]
Iterator is fully consumed (no partial reads)
Output directory is created automatically (mkdir(parents=True, exist_ok=True))
Empty iterator returns [] (no crash, no empty files)
Logging via get_logger(self) (not logging.getLogger(__name__))
Timing logged for write operations
Naming and Source Integration
naming_template param if index-based naming is used
set_source() implemented if template uses source-backed placeholders
Placeholder validation in __init__ (fail fast on invalid templates)
Parallel Safety (if applicable)
partition_indices() returns correct groups or None
Concurrent writes to same file are impossible within a partition
Atomic writes (temp + rename) where format doesn't support append
Documentation
NumPy-style docstrings on class, __init__, __call__, and helpers
from __future__ import annotations at top
TYPE_CHECKING block for Iterator and domain type imports
Read-only @property accessors for key config values
Registration
Registered in domain __init__.py with registry.register_sink()
Added to __all__ in domain __init__.py
Lazy loading if sink has heavy optional dependencies
Roundtrip test: written data can be read back correctly
Append tests (if applicable): data accumulates rather than overwrites
Split tests (if applicable): data is split by variable/coordinate
Partition tests (if applicable): correct grouping, None when no constraint
Pipeline test: sink works at end of a pipeline (marked @pytest.mark.e2e)
Registry test: sink is discoverable
Quality Gates
ruff format clean
ruff check clean
interrogate >= 99%
ty check clean
All tests pass
No regressions in existing tests
Reference: Sink[T] ABC
From src/physicsnemo_curator/core/base.py:
classSink[T](ABC):
"""Abstract sink that persists items and returns output file paths.
The sink consumes a generator of items and writes each one to storage,
returning the file paths of the written outputs.
Subclasses must set :attr:`name` and :attr:`description` and implement
:meth:`params` and :meth:`__call__`.
"""
name: ClassVar[str]
"""Human-readable display name for the interactive CLI."""
description: ClassVar[str]
"""Short description shown in the interactive CLI.""" @classmethod @abstractmethoddefparams(cls) -> list[Param]:
"""Declare the configurable parameters for this sink.
Returns
-------
list[Param]
Ordered list of parameter descriptors.
"""
...
@abstractmethoddef__call__(self, items: Iterator[T], index: int) -> list[str]:
"""Consume items and persist them to storage.
Parameters
----------
items : Iterator[T]
Stream of data items to write.
index : int
Source index being processed (useful for naming output files).
Returns
-------
list[str]
Paths of the files written.
"""
...
defpartition_indices(self, indices: list[int]) -> list[list[int]] | None:
"""Group indices into partitions that MUST be processed by the same worker.
Each returned group is a list of indices that must be handled
sequentially by a single worker. The runner will never split a
group across workers.
Override this method when the sink has constraints on concurrent
writes (e.g., multiple indices writing to the same Zarr chunk must
go through the same worker).
Parameters
----------
indices : list[int]
The indices to partition.
Returns
-------
list[list[int]] | None
Partitioned groups, or ``None`` if no partitioning is required
(the default).
"""returnNone
Key differences from Source[T] and Filter[T]:
Sources have __len__ and __getitem__ — sinks do not
Filters receive and return Generator[T] — sinks consume Iterator[T]
and return list[str] (file paths)
Sinks are the terminal stage: they consume items, they do not yield
The index parameter is the source-level index from the pipeline
Return value is a list of all file paths written (can be empty)
Sinks do NOT have dashboard_panel(), artifacts(), or merge()
(those are filter-only methods)
partition_indices() is unique to sinks (not on sources or filters)
set_source() is an optional protocol method (not in the ABC)
Reference: Param Dataclass
@dataclass(frozen=True)classParam:
"""Descriptor for a configurable parameter."""
name: str
description: strtype: type
default: Any = REQUIRED # Sentinel: no default = required param
choices: list[Any] | None = None
Use REQUIRED sentinel (or omit default) for mandatory parameters
