add-source

# 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: 1. **Name and URL** — What is the dataset called? Where is it hosted? (HuggingFace `hf://datasets/org/repo`, S3, HTTP, local path) 2. **File format** — What format are the files? (VTK/VTP/VTU, Parquet, HDF5, NetCDF, NumPy, Zarr, LMDB, CSV, custom) 3. **Schema / fields** — What fields does the data contain? (coordinates, velocity, pressure, temperature, connectivity, forces, energies, etc.) 4. **Spatial dimensions** — Is the data 1D, 2D, 3D? Structured grid or unstructured mesh? Molecular graph? 5. **Dataset size** — How large is the full dataset? How many samples? 6. **Organization** — How are files organized? Single file per sample? Run directories (`run_0/`, `run_1/`)? Splits (train/val/test)? Multi-file databases? 7. **License** — What license? Is it permissive enough for inclusion? 8. **Documentation** — Is there a paper, README, or data card? 9. **Parameters** — What varies between samples? (geometry, physics parameters, initial conditions, time steps) 10. **Authentication** — Does accessing the data require tokens or credentials? 11. **Parallel constraints** — Does the data backend have concurrency limitations? (e.g. LMDB single-env-per-process, HDF5 locking) 12. **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: ```python from physicsnemo.mesh import Mesh mesh = Mesh( points=points_tensor, # torch.Tensor, shape (n_points, 3) cells=cells_tensor, # torch.Tensor, shape (n_cells, nodes_per_cell), dtype=long point_data=point_data_td, # TensorDict with batch_size=[n_points], or None cell_data=cell_data_td, # TensorDict with batch_size=[n_cells], or None global_data=global_data_td,# TensorDict with batch_size=[], or None ) ``` 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`: ```python from nvalchemi.data import AtomicData data = AtomicData( positions=positions, # torch.Tensor, shape (n_atoms, 3) atomic_numbers=z, # torch.Tensor, shape (n_atoms,) cell=cell, # torch.Tensor, shape (3, 3) or None pbc=pbc, # torch.Tensor, shape (3,) bool energy=energy, # torch.Tensor, shape (1,) or None forces=forces, # torch.Tensor, shape (n_atoms, 3) or None ) ``` ### 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` | `da/sources/random.py` | Synthetic | da | Testing, minimal DA source | | `RandomAtomicSource` | `atm/sources/random.py` | Synthetic | atm | Testing, minimal atomic source | ## 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: ```python # 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 ```python """<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 ```python _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="", ), # Add dataset-specific params (mesh_type, split, backend, etc.) ] def __init__( self, url: str = _DEFAULT_URL, storage_options: dict[str, object] | None = None, cache_storage: str = "", ) -> None: """Initialize the source. Parameters ---------- url : str Base dataset URL. storage_options : dict[str, object] | None Extra keyword arguments for the fsspec filesystem. cache_storage : str Local cache directory. """ self._url = url self._storage_options = storage_options or {} self._cache_storage = cache_storage or tempfile.mkdtemp( prefix="curator_<name>_" ) self._log = get_logger(self) # Initialize file discovery and data access here # Load lightweight metadata eagerly (e.g. parameter tables) # Defer heavy data loading (geometry, fields) to __getitem__ self._discover_items() def _discover_items(self) -> None: """Eagerly discover available items (files, runs, etc.).""" ... def __len__(self) -> int: """Return the number of items in this source.""" return self._count def __getitem__(self, index: int) -> Generator[Mesh]: """Yield mesh(es) for the given index. Parameters ---------- index : int Zero-based index. Negative indices are supported. Yields ------ Mesh One or more mesh objects for this index. Raises ------ IndexError If *index* is out of range. """ if index < 0: index += len(self) if index < 0 or index >= len(self): msg = f"Index {index} out of range for source with {len(self)} items." raise IndexError(msg) # Load data for this index and construct Mesh yield mesh ``` ### Logging Use the structured component logger: ```python 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): ```python 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