| name | add-source |
| description | Guide for adding a new data source to PhysicsNeMo Curator. Use when adding a remote dataset (HuggingFace, S3, etc.) or local file-based source. Covers discovery questions, file format handling, Mesh/DataArray/AtomicData construction, parallel partitioning, testing, and registration.
|
Adding a New Data Source
This skill walks through adding a new Source to PhysicsNeMo Curator,
from initial research 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 source belong to?
| Domain | Type parameter | Submodule | Dependency group |
|---|
| mesh | Source[Mesh] | src/physicsnemo_curator/domains/mesh/ | mesh (physicsnemo, pyvista, pyarrow, torch) |
| da | Source[xr.DataArray] | src/physicsnemo_curator/domains/da/ | da (xarray, earth2studio, zarr) |
| atm | Source[AtomicData] | src/physicsnemo_curator/domains/atm/ | mesh (nvalchemi, torch, lmdb) |
Dataset Information
Gather all of the following before proceeding:
- Name and URL — What is the dataset called? Where is it hosted?
(HuggingFace
hf://datasets/org/repo, S3, HTTP, local path)
- File format — What format are the files? (VTK/VTP/VTU, Parquet,
HDF5, NetCDF, NumPy, Zarr, LMDB, CSV, custom)
- Schema / fields — What fields does the data contain? (coordinates,
velocity, pressure, temperature, connectivity, forces, energies, etc.)
- Spatial dimensions — Is the data 1D, 2D, 3D? Structured grid or
unstructured mesh? Molecular graph?
- Dataset size — How large is the full dataset? How many samples?
- Organization — How are files organized? Single file per sample?
Run directories (
run_0/, run_1/)? Splits (train/val/test)?
Multi-file databases?
- License — What license? Is it permissive enough for inclusion?
- Documentation — Is there a paper, README, or data card?
- Parameters — What varies between samples? (geometry, physics
parameters, initial conditions, time steps)
- Authentication — Does accessing the data require tokens or credentials?
- Parallel constraints — Does the data backend have concurrency
limitations? (e.g. LMDB single-env-per-process, HDF5 locking)
- Backend options — Is there a fast native reader (Rust) available
as well as a Python fallback?
Output Type Construction
For mesh sources, determine how to construct physicsnemo.mesh.Mesh:
- VTK-based: Use
pyvista.read() then from_pyvista() (see existing
VTK sources)
- Non-VTK (Parquet, HDF5, NumPy): Construct
Mesh directly:
from physicsnemo.mesh import Mesh
mesh = Mesh(
points=points_tensor,
cells=cells_tensor,
point_data=point_data_td,
cell_data=cell_data_td,
global_data=global_data_td,
)
For da sources, determine how to construct xarray.DataArray with
appropriate dimensions, coordinates, and attributes.
For atm sources, determine how to construct
nvalchemi.data.AtomicData:
from nvalchemi.data import AtomicData
data = AtomicData(
positions=positions,
atomic_numbers=z,
cell=cell,
pbc=pbc,
energy=energy,
forces=forces,
)
File Discovery Strategy
Each source handles its own file discovery and caching internally.
Choose the approach based on dataset organization:
| Pattern | Approach | When to use |
|---|
| Flat directory of files | pathlib.Path.glob() | Local-only sources (VTKSource, ASELMDBSource) |
| Remote flat files | fsspec + fs.glob() | Generic remote file glob |
Run-indexed dirs (run_0/, run_1/) | fsspec + fs.ls() + regex | Benchmark datasets with numbered runs (DrivAerML, AhmedML) |
| API/backend-driven | Domain library fetcher | Weather reanalysis (ERA5, HRRR via earth2studio) |
| Single-file datasets | Direct fsspec/pyarrow | Parquet tables, Zarr stores |
Reference: Existing Sources
| Source | File | Format | Domain | Good example for |
|---|
VTKSource | mesh/sources/vtk.py | VTK | mesh | Local file glob, pyvista/rust backend, relative_path |
DrivAerMLSource | mesh/sources/drivaerml.py | VTP/VTU | mesh | Remote fsspec, multi-mesh, run_id, mesh_name |
AhmedMLSource | mesh/sources/ahmedml.py | VTP | mesh | HF remote, CSV metadata, multi-mode, run_id |
NavierStokesCylinderSource | mesh/sources/ns_cylinder.py | Parquet | mesh | Direct Mesh construction, non-VTK |
D3PlotSource | mesh/sources/d3plot.py | d3plot | mesh | Binary format, python/rust dual backend |
AnsysRSTSource | mesh/sources/ansys_rst.py | .rst | mesh | Auto-discovery of results, complex extraction |
OpenRadiossSource | mesh/sources/openradioss.py | VTK | mesh | Per-timestep multi-file, run_id |
ERA5Source | da/sources/era5.py | API | da | Multi-backend weather, DataArray, earth2studio |
HRRRSource | da/sources/hrrr.py | API | da | HRRR 3km analysis, multiple providers |
ASELMDBSource | atm/sources/aselmdb.py | LMDB | atm | partition_indices, rust fallback, relative_path |
RandomMeshSource | mesh/sources/random.py | Synthetic | mesh | Testing, minimal source |
RandomDataArraySource | |
Step 1: Create the Source Module
Create the source file at the appropriate location:
- mesh:
src/physicsnemo_curator/domains/mesh/sources/<name>.py
- da:
src/physicsnemo_curator/domains/da/sources/<name>.py
- atm:
src/physicsnemo_curator/domains/atm/sources/<name>.py
Required SPDX Header
Every source file MUST start with:
Import Block Template
"""<Dataset name> data source for PhysicsNeMo Curator.
Reads <format> data from <location> and yields :class:`~physicsnemo.mesh.Mesh`
objects for use in curator pipelines.
Examples
--------
>>> source = <ClassName>(...) # doctest: +SKIP
>>> len(source) # doctest: +SKIP
<expected_count>
>>> mesh = next(source[0]) # doctest: +SKIP
"""
from __future__ import annotations
import pathlib
import tempfile
from typing import TYPE_CHECKING, ClassVar, Literal
from physicsnemo_curator.core.base import Param, Source
from physicsnemo_curator.core.logging import get_logger
if TYPE_CHECKING:
from collections.abc import Generator
from physicsnemo.mesh import Mesh
Source Class Template
_DEFAULT_URL = "<dataset_url>"
class <ClassName>(Source[Mesh]):
"""Read meshes from <Dataset Name>.
<2-3 sentence description of the dataset, what it contains,
and where it comes from. Include a link to the paper or dataset page.>
Parameters
----------
url : str
Base URL for the dataset.
storage_options : dict[str, object] | None
Extra keyword arguments for the fsspec filesystem (e.g. HF token).
cache_storage : str
Local directory for caching downloaded files.
Note
----
- Dataset: `<org/repo> <<url>>`_
- Paper: `arXiv:<id> <https://arxiv.org/abs/<id>>`_
- License: `<license_name> <<license_url>>`_
"""
name: ClassVar[str] = "<Display Name>"
description: ClassVar[str] = "<Short description for CLI>"
@classmethod
def params(cls) -> list[Param]:
"""Return configurable parameters for this source.
Returns
-------
list[Param]
Parameter list for CLI configuration.
"""
return [
Param(
name="url",
description="Base dataset URL",
type=str,
default=_DEFAULT_URL,
),
Param(
name="cache_storage",
description="Local cache directory for downloaded files",
type=str,
default="",
),
]
def __init__(
self,
url: str = _DEFAULT_URL,
storage_options: dict[str, object] | None = ,
cache_storage: = ,
) -> :
._url = url
._storage_options = storage_options {}
._cache_storage = cache_storage tempfile.mkdtemp(
prefix=
)
._log = get_logger()
._discover_items()
() -> :
...
() -> :
._count
() -> Generator[Mesh]:
index < :
index += ()
index < index >= ():
msg =
IndexError(msg)
mesh
Logging
Use the structured component logger:
from physicsnemo_curator.core.logging import get_logger
class MySource(Source[Mesh]):
def __init__(self, ...) -> None:
self._log = get_logger(self)
self._log.info("Discovered %d items in %s", count, url)
def __getitem__(self, index: int) -> Generator[Mesh]:
self._log.debug("Reading index %d", index)
...
For module-level logging (outside class scope):
from physicsnemo_curator.core.logging import get_logger
logger = get_logger("MySourceModule")
Do NOT use logging.getLogger(__name__) — always use get_logger.
Key Implementation Patterns
Pattern A: Local File-Based (VTKSource, ASELMDBSource)
Discovers files at construction time via pathlib.Path.glob(), then
reads them lazily in __getitem__:
def __init__(self, input_path: str, file_pattern: str = "**/*") -> None:
root = pathlib.Path(input_path)
if not root.is_dir():
msg = f"Path {root} is not a directory."
raise FileNotFoundError(msg)
self._root = root.resolve()
self._files: list[pathlib.Path] = sorted(
p.resolve()
for p in root.glob(file_pattern)
if p.is_file() and p.suffix.lower() in self._VALID_EXTENSIONS
)
if not self._files:
msg = f"No valid files found in {root} with pattern {file_pattern!r}"
raise ValueError(msg)
self._log.info("Discovered %d files in %s", len(self._files), root)
def __len__(self) -> int:
return len(self._files)
def __getitem__(self, index: int) -> Generator[Mesh]:
path = self._files[index]
mesh = self._read_file(path)
mesh
Pattern B: Remote HuggingFace Hub (DrivAerML, AhmedML)
Uses fsspec for file discovery and WholeFileCacheFileSystem for local
caching of downloaded files:
import fsspec
def __init__(self, url: str, cache_storage: str = "", cache: bool = True, ...) -> None:
self._url = url
self._cache = cache
self._protocol, self._path = fsspec.utils.split_protocol(url)
target_fs = fsspec.filesystem(self._protocol, **storage_options)
if cache:
self._fs = fsspec.filesystem(
"filecache",
target_protocol=self._protocol,
target_options=storage_options,
cache_storage=cache_storage or tempfile.mkdtemp(),
)
else:
self._fs = target_fs
self._run_indices = self._discover_runs()
self._log.info("Discovered %d runs", len(self._run_indices))
def _discover_runs(self) -> list[int]:
"""Find run directories via regex pattern matching."""
import re
pattern = re.compile(r"run_(\d+)/?$")
entries = self._fs.ls(self._path, detail=False)
runs = sorted((m.group()) e entries (m := pattern.search(e)))
runs
() -> Generator[Mesh]:
run_id = ._run_indices[index]
remote_path =
local_path = ._fs.(remote_path).name
mesh = ._read_vtk(local_path)
mesh
Pattern C: API/Backend-Driven (ERA5Source, HRRRSource)
Uses a domain library (earth2studio) to fetch data, routing requests
to different backends:
def __init__(
self,
variables: list[str],
time_range: tuple[str, str],
backend: str = "arco",
...
) -> None:
self._variables = variables
self._backend = backend
self._log = get_logger(self)
self._times = self._build_time_index(time_range)
def __getitem__(self, index: int) -> Generator[xr.DataArray]:
time = self._times[index]
da = self._fetch(time)
yield da
Pattern D: Multi-File Database (ASELMDBSource)
Multiple files each containing many items, with cumulative indexing:
def __init__(self, data_dir: str, file_pattern: str = "**/*.aselmdb") -> None:
self._root = pathlib.Path(data_dir).resolve()
self._db_files = sorted(
p.resolve() for p in self._root.glob(file_pattern)
if p.suffix == ".aselmdb"
)
self._row_counts = [self._count_rows(f) for f in self._db_files]
self._cumulative_counts = list(itertools.accumulate(self._row_counts))
self._total = self._cumulative_counts[-1] if self._cumulative_counts else 0
def __len__(self) -> int:
return self._total
def __getitem__(self, index: int) -> Generator[AtomicData]:
file_idx = self._find_file_for_index(index)
local_idx = index - (self._cumulative_counts[file_idx - 1] if file_idx > 0 else )
data = ._read_single(._db_files[file_idx], local_idx)
data
() -> :
bisect
bisect.bisect_right(._cumulative_counts, index) -
Pattern E: Multi-Output Per Index (DrivAerML "multi" mode)
A single index can yield multiple items (e.g. multiple mesh parts):
def __getitem__(self, index: int) -> Generator[Mesh | DomainMesh]:
run_id = self._run_indices[index]
for part in self._mesh_parts:
remote_path = self._resolve_part_path(run_id, part)
local_path = self._download(remote_path)
mesh = self._read_vtk(local_path)
yield mesh
Dual Backend Pattern (Python + Rust)
Many sources support both a Python backend and a faster Rust backend:
def __init__(self, ..., backend: Literal["python", "rust"] = "python") -> None:
self._backend: Literal["python", "rust"] = backend
if backend == "rust":
try:
from physicsnemo_curator._lib.some_module import reader as _
except (ImportError, ModuleNotFoundError):
self._log.warning("Rust backend unavailable; falling back to 'python'.")
self._backend = "python"
def __getitem__(self, index: int) -> Generator[Mesh]:
if self._backend == "rust":
mesh = self._read_with_rust(index)
else:
mesh = self._read_with_python(index)
yield mesh
Lazy Loading
Follow this pattern for efficiency:
- Eager (in
__init__): lightweight metadata, parameter tables, file
counts, directory listings
- Lazy (in
__getitem__): geometry, field data, large arrays
- Cached: data that is shared across indices (e.g. a single geometry
used by all snapshots)
def __init__(self, ...):
self._geometry_loaded = False
self._points: torch.Tensor | None = None
self._cells: torch.Tensor | None = None
def _load_geometry(self) -> None:
"""Load shared geometry (called once, result cached)."""
if self._geometry_loaded:
return
self._geometry_loaded = True
def __getitem__(self, index):
self._load_geometry()
Step 1b: Sink Integration Protocol (Optional)
Sources can optionally provide methods that sinks use for output naming.
These are NOT part of the ABC but are checked via hasattr() by
Pipeline.write() and sinks.
relative_path() Method
For file-based sources, enables directory mirroring in output:
@property
def root(self) -> pathlib.Path:
"""Return the root directory of this source.
Returns
-------
pathlib.Path
The root directory containing the discovered files.
"""
return self._root
def relative_path(self, index: int) -> str:
"""Return the path of the *index*-th file relative to the root.
Used by sinks to resolve ``{relpath}`` and ``{stem}`` naming
placeholders, enabling output layouts that mirror the input.
Parameters
----------
index : int
Zero-based file index.
Returns
-------
str
POSIX-style relative path (e.g. ``"subdir/mesh.vtu"``).
"""
return self._files[index].relative_to(self._root).as_posix()
run_id() Method
For run-based datasets, provides the dataset run identifier:
def run_id(self, index: int) -> int:
"""Return the dataset run ID for the given source index.
Used by sinks to resolve the ``{run_id}`` naming placeholder.
Parameters
----------
index : int
Zero-based index into the sorted run list.
Returns
-------
int
The dataset run ID (e.g. 1, 5, 12).
"""
return self._run_indices[index]
mesh_name() Method
For multi-output sources, provides canonical names for each part:
def mesh_name(self, index: int, seq: int) -> str:
"""Return the canonical output name for mesh at (index, seq).
Used by sinks to resolve the ``{mesh_name}`` naming placeholder.
Parameters
----------
index : int
Source index (which run).
seq : int
Sequence number within this index (which part).
Returns
-------
str
Resolved name like ``"domain_1"`` or ``"drivaer_5.stl"``.
"""
run_id = self._run_indices[index]
part = self._mesh_parts[seq]
return self._TEMPLATES[part].format(run_id=run_id)
Step 1c: Parallel Partitioning (Optional)
partition_indices() Method
Override when the source has concurrency constraints:
def partition_indices(self, indices: list[int]) -> list[list[int]] | None:
"""Group indices so all from same file go to same worker.
Improves data locality and avoids concurrent access issues
(e.g. LMDB single-env-per-process constraint).
Parameters
----------
indices : list[int]
The indices to partition.
Returns
-------
list[list[int]] | None
Groups of indices (one per file), or ``None`` if all from one file.
"""
from collections import defaultdict
file_groups: dict[int, list[int]] = defaultdict(list)
for idx in indices:
file_idx = self._find_file_for_index(idx)
file_groups[file_idx].append(idx)
if len(file_groups) <= 1:
return None
return [sorted(file_groups[k]) for k in sorted(file_groups)]
Param Declaration Patterns
Required parameter (no default — user must provide):
Param(name="input_path", description="Path to data directory", type=str)
URL with default (remote dataset):
Param(name="url", description="Base HuggingFace Hub URL", type=str, default=_DEFAULT_URL)
Backend choice:
Param(
name="backend",
description="Reading backend: pyvista (default) or rust (faster)",
type=str,
default="pyvista",
choices=["pyvista", "rust"],
)
Glob pattern:
Param(name="file_pattern", description="Glob pattern for file discovery", type=str, default="**/*")
Cache toggle:
Param(
name="cache",
description="Persist downloaded files across sessions",
type=bool,
default=True,
)
Mesh conversion options (for VTK-based sources):
Param(
name="manifold_dim",
description="Target manifold dimension (auto, 0, 1, 2, 3)",
type=str,
default="auto",
choices=["auto", "0", "1", "2", "3"],
),
Param(
name="point_source",
description="Point source mode: vertices or cell_centroids",
type=str,
default="vertices",
choices=["vertices", "cell_centroids"],
),
Param(
name="warn_on_lost_data",
description="Warn when data arrays are discarded during conversion",
type=bool,
default=True,
),
Metadata path (optional supplementary file):
Param(
name="metadata_path",
description="Path to metadata file (empty = auto-detect)",
type=str,
default="",
)
Step 2: Write Tests
Create the test file at test/domains/<domain>/test_<name>.py.
Test File Structure
"""Tests for <ClassName>."""
from __future__ import annotations
import pathlib
import numpy as np
import pytest
pytestmark = pytest.mark.requires("<domain>")
def _write_mock_dataset(root: pathlib.Path, n_samples: int = 3) -> None:
"""Create a minimal mock dataset for unit tests.
Write mock files that mirror the real dataset structure.
Use small sizes (10-50 points, 3-5 samples).
"""
...
class Test<ClassName>Unit:
"""Metadata and parameter tests (no data access)."""
def test_params_list(self) -> None:
from <module> import <ClassName>
params = <ClassName>.params()
assert len(params) > 0
names = [p.name for p in params]
assert "url" in names
def test_name_and_description() -> :
<module> <ClassName>
(<ClassName>.name, )
(<ClassName>.name) >
(<ClassName>.description, )
(<ClassName>.description) >
Local Mock Tests
class Test<ClassName>Local:
"""Tests against local mock data."""
@pytest.fixture(autouse=True)
def _setup(self, tmp_path: pathlib.Path) -> None:
self.mock_root = tmp_path / "mock_dataset"
_write_mock_dataset(self.mock_root)
def test_len(self) -> None:
from <module> import <ClassName>
source = <ClassName>(url=str(self.mock_root))
assert len(source) == 3
def test_getitem_returns_correct_type(self) -> None:
from physicsnemo.mesh import Mesh
from <module> import <ClassName>
source = <ClassName>(url=str(self.mock_root))
item = next(source[0])
assert isinstance(item, Mesh)
assert item.n_points > 0
def test_negative_index(self) -> None:
from <module> import <ClassName>
source = <ClassName>(url=(.mock_root))
item_neg = (source[-])
item_pos = (source[(source) - ])
item_neg.n_points == item_pos.n_points
() -> :
<module> <ClassName>
source = <ClassName>(url=(.mock_root))
pytest.raises(IndexError):
(source[(source)])
() -> :
<module> <ClassName>
source = <ClassName>(url=(.mock_root))
item = (source[])
item.point_data.keys()
() -> :
<module> <ClassName>
source = <ClassName>(url=(.mock_root))
item_0 = (source[])
item_1 = (source[])
...
() -> :
<module> <ClassName>
empty = .mock_root /
empty.mkdir()
pytest.raises((ValueError, FileNotFoundError)):
<ClassName>(url=(empty))
Sink Integration Tests
def test_relative_path(self) -> None:
"""relative_path returns POSIX path relative to root."""
from <module> import <ClassName>
source = <ClassName>(input_path=str(self.mock_root))
rel = source.relative_path(0)
assert "/" not in rel or not rel.startswith("/")
assert pathlib.PurePosixPath(rel).suffix != ""
def test_run_id(self) -> None:
"""run_id returns integer run identifier."""
from <module> import <ClassName>
source = <ClassName>(url=str(self.mock_root))
rid = source.run_id(0)
assert isinstance(rid, (int, str))
Partition Tests (if applicable)
def test_partition_indices_groups_by_file(self) -> None:
"""Items from same file go to same group."""
from <module> import <ClassName>
source = <ClassName>(data_dir=str(self.mock_root))
groups = source.partition_indices(list(range(len(source))))
if groups is not None:
flat = [idx for g in groups for idx in g]
assert sorted(flat) == list(range(len(source)))
def test_partition_single_file_returns_none(self) -> None:
"""Single file = no partitioning needed."""
from <module> import <ClassName>
source = <ClassName>(data_dir=str(self.single_file_root))
assert source.partition_indices(list(range(len(source)))) is None
Backend Tests
def test_backend_fallback(self) -> None:
"""Rust backend falls back gracefully if unavailable."""
from <module> import <ClassName>
source = <ClassName>(input_path=str(self.mock_root), backend="rust")
item = next(source[0])
assert item.n_points > 0
Pipeline Integration Tests
@pytest.mark.e2e
class Test<ClassName>Pipeline:
"""End-to-end pipeline tests."""
@pytest.fixture(autouse=True)
def _setup(self, tmp_path: pathlib.Path) -> None:
self.mock_root = tmp_path / "mock"
_write_mock_dataset(self.mock_root)
def test_source_in_pipeline(self, tmp_path: pathlib.Path) -> None:
"""Source works as pipeline head."""
from physicsnemo_curator.domains.mesh.sinks.mesh_writer import MeshSink
from <module> import <ClassName>
source = <ClassName>(url=str(self.mock_root))
sink = MeshSink(output_dir=str(tmp_path / "output"))
pipeline = source.write(sink)
assert len(pipeline) == len(source)
paths = pipeline[0]
assert len(paths) >= 1
assert pathlib.Path(paths[0]).exists()
Registry Tests
class Test<ClassName>Registry:
"""Test that the source is registered."""
def test_source_registered(self) -> None:
from physicsnemo_curator.core.registry import registry
names = [s.name for s in registry.list_sources("<domain>")]
assert "<Display Name>" in names
E2E Tests (Remote)
@pytest.mark.e2e
@pytest.mark.slow
class Test<ClassName>E2E:
"""End-to-end tests against live dataset (requires network)."""
@pytest.fixture(autouse=True)
def _setup(self, tmp_path: pathlib.Path) -> None:
from <module> import <ClassName>
self.source = <ClassName>(cache_storage=str(tmp_path / "cache"))
def test_discovers_items(self) -> None:
assert len(self.source) == <expected_count>
def test_reads_first_item(self) -> None:
from physicsnemo.mesh import Mesh
item = next(self.source[0])
assert isinstance(item, Mesh)
assert item.n_points == <expected_points>
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 test |
@pytest.mark.slow | Slow test, excluded from quick CI runs |
Running Tests
uv run pytest test/domains/<domain>/test_<name>.py -v -k "not E2E"
uv run pytest test/domains/<domain>/test_<name>.py -v
uv run pytest test/domains/<domain>/test_<name>.py -v -k "E2E"
Step 3: Register the Source
Edit the domain __init__.py
For mesh sources, edit src/physicsnemo_curator/domains/mesh/__init__.py:
from physicsnemo_curator.domains.mesh.sources.<module> import <ClassName>
registry.register_source("mesh", <ClassName>)
__all__ = [
...,
"<ClassName>",
...,
]
For da sources, edit src/physicsnemo_curator/domains/da/__init__.py with
the same pattern using "da" as the submodule name.
For atm sources, edit src/physicsnemo_curator/domains/atm/__init__.py
with "atm" as the submodule name.
Step 4: Quality Checks
Run all checks before committing:
uv run ruff format \
src/physicsnemo_curator/domains/<domain>/sources/<name>.py \
test/domains/<domain>/test_<name>.py
uv run ruff check --fix \
src/physicsnemo_curator/domains/<domain>/sources/<name>.py \
test/domains/<domain>/test_<name>.py
uv run interrogate
uv run ty check
uv run pytest test/domains/<domain>/test_<name>.py -v -k "not E2E"
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 <dataset> datasets
<Optional body describing the dataset, format, and what it provides.>
Checklist
Before considering the source complete, verify:
Core
Data Access
Sink Integration (if applicable)
Parallel Safety (if applicable)
Documentation
Registration
Testing
Quality Gates
Reference: Source[T] ABC
From src/physicsnemo_curator/core/base.py:
class Source[T](ABC):
"""Abstract data source that yields items of type *T*.
A source represents a collection of data items (e.g. files on disk).
Each item is accessed by integer index and may yield one or more *T*
objects (generator semantics allow a single source item to expand into
multiple outputs).
Subclasses must set the class-level :attr:`name` and :attr:`description`
attributes and implement :meth:`params`, :meth:`__len__`, and
:meth:`__getitem__`.
"""
name: ClassVar[str]
"""Human-readable display name for the interactive CLI."""
description: ClassVar[str]
"""Short description shown in the interactive CLI."""
@classmethod
@abstractmethod
def params(cls) -> list[Param]: ...
@abstractmethod
def __len__(self) -> int: ...
@abstractmethod
def __getitem__(self, index: int) -> Generator[T]: ...
def partition_indices(self, indices: list[int]) -> list[list[int]] | None:
"""Group indices into partitions that MUST be processed by the same worker.
Override when the source has constraints on concurrent access
(e.g. LMDB allows only one environment open per file per process).
Returns ``None`` by default (no partitioning required).
"""
return None
def filter() -> Pipeline[T]: ...
() -> Pipeline[T]: ...
Key points:
__getitem__ returns Generator[T] (can yield multiple items per index)
partition_indices() is optional (same as on Sink)
write() automatically calls s.set_source(self) if sink has that method
- Sources do NOT have
dashboard_panel(), artifacts(), or flush()
Reference: Param Dataclass
@dataclass(frozen=True)
class Param:
"""Descriptor for a configurable parameter."""
name: str
description: str
type: type
default: Any = REQUIRED
choices: list[Any] | None = None
- Use
REQUIRED sentinel (or omit default) for mandatory parameters
choices constrains values to a finite set
type is used for CLI parsing and validation