Adding a New Datastore

Adding a New Datastore#

This notebook walks through the steps required to plug a new dataset into neural-lam. neural-lam uses a datastore abstraction layer — implement a class that satisfies the BaseDatastore interface and the rest of the framework works automatically.

Note: Run from the repository root, or from docs/notebooks/ — the setup cell handles the working directory automatically.

1. Setup#

import inspect
import os

# Resolve repo root regardless of where the kernel was launched from
while not os.path.isdir(os.path.join(os.getcwd(), "neural_lam")):
    parent = os.path.dirname(os.getcwd())
    if parent == os.getcwd():
        raise RuntimeError("Could not find repo root (neural_lam/ not found)")
    os.chdir(parent)
print("Working directory:", os.getcwd())

2. Overview of the Datastore Abstraction#

All datastores inherit from BaseDatastore (neural_lam/datastore/base.py). The interface requires methods covering:

  • Variable names, units, and long names per category (state, forcing, static)

  • Standardisation statistics (mean, std) per category via get_standardization_dataarray

  • An xarray.DataArray of shape (time, grid_index, {category}_feature) per split via get_dataarray

  • Grid coordinates via get_xy and a boundary_mask cached property

  • coords_projection property returning a Cartopy CRS

For regular projected grids, inherit from BaseRegularGridDatastore instead — it provides helpers for grid shape and coordinate stacking.

3. Inspect the Abstract Interface#

from neural_lam.datastore.base import BaseDatastore

# Collect every abstract method/property a subclass must implement
abstract_members = sorted(
    name
    for name in dir(BaseDatastore)
    if getattr(getattr(BaseDatastore, name, None), "__isabstractmethod__", False)
)
print("Abstract members that must be implemented:")
for name in abstract_members:
    print(f"  {name}")

4. Existing Datastore Implementations#

neural-lam ships with two backends you can use as reference:

# 1. MDPDatastore — wraps mllam-data-prep zarr outputs (recommended for new datasets)
from neural_lam.datastore.mdp import MDPDatastore

# 2. NpyFilesDatastoreMEPS — reads raw .npy files (legacy MEPS format)
from neural_lam.datastore.npyfilesmeps.store import NpyFilesDatastoreMEPS

print("MDPDatastore          — config-driven zarr backend")
print("NpyFilesDatastoreMEPS — legacy numpy file backend (MEPS format)")
print()
print(f"MDPDatastore.SHORT_NAME          = {MDPDatastore.SHORT_NAME!r}")
print(f"NpyFilesDatastoreMEPS.SHORT_NAME = {NpyFilesDatastoreMEPS.SHORT_NAME!r}")

6. Option B — Custom Datastore Class#

If your data format doesn’t fit mllam-data-prep (custom binary formats, API sources, etc.), subclass BaseRegularGridDatastore and implement the required interface.

Key points to get right:

  • The standardisation method is get_standardization_dataarray (not get_normalization_dataarray)

  • The projection property is coords_projection (not projection)

  • The boundary mask is a @cached_property named boundary_mask (not a get_boundary_mask() method)

  • num_grid_points and state_feature_weights_values must also be implemented

# Template — do not run as-is; adapt the NotImplementedError stubs to your data

from datetime import timedelta
from functools import cached_property
from pathlib import Path
from typing import List, Optional

import cartopy.crs as ccrs
import numpy as np
import xarray as xr

from neural_lam.datastore.base import BaseRegularGridDatastore, CartesianGridShape


class MyCustomDatastore(BaseRegularGridDatastore):
    """Custom datastore for MyDataset."""

    SHORT_NAME = "my_dataset"

    def __init__(self, config_path: str, n_boundary_points: int = 30):
        self._config_path = Path(config_path)
        self._root_path = self._config_path.parent
        self._n_boundary_points = n_boundary_points
        # Load your config/data here, e.g.:
        # import yaml
        # with open(config_path) as f:
        #     self._raw_config = yaml.safe_load(f)

    # ------------------------------------------------------------------ #
    # Required by BaseDatastore
    # ------------------------------------------------------------------ #

    @property
    def root_path(self) -> Path:
        return self._root_path

    @property
    def config(self):
        # Must return the raw config mapping loaded in __init__.
        # Do NOT return an empty dict — framework code inspects it.
        raise NotImplementedError("Return self._raw_config or equivalent")

    @property
    def step_length(self) -> timedelta:
        return timedelta(hours=3)  # adjust to match your dataset

    def get_vars_names(self, category: str) -> List[str]:
        mapping = {
            "state":   ["temperature_2m", "wind_u_10m"],
            "forcing": ["radiation"],
            "static":  ["land_sea_mask", "orography"],
        }
        return mapping.get(category, [])

    def get_vars_units(self, category: str) -> List[str]:
        return ["K", "m/s"] if category == "state" else []

    def get_vars_long_names(self, category: str) -> List[str]:
        return ["2m Temperature", "10m U-wind"] if category == "state" else []

    def get_num_data_vars(self, category: str) -> int:
        return len(self.get_vars_names(category))

    def get_standardization_dataarray(self, category: str) -> xr.Dataset:
        # Must return an xr.Dataset with variables:
        #   {category}_mean, {category}_std
        # and for category=="state" also:
        #   state_diff_mean_standardized, state_diff_std_standardized
        raise NotImplementedError

    def get_dataarray(
        self,
        category: str,
        split: Optional[str],
        standardize: bool = False,
    ) -> Optional[xr.DataArray]:
        # Must return DataArray with dims (time, grid_index, {category}_feature)
        # Return None if this category is not supported by your datastore.
        raise NotImplementedError

    @cached_property
    def boundary_mask(self) -> xr.DataArray:
        # Must return a DataArray with dim (grid_index,).
        # Value 1 = boundary point, 0 = interior point.
        raise NotImplementedError

    def get_xy(self, category: str, stacked: bool = True) -> np.ndarray:
        # stacked=True  -> return (N_grid, 2) projected x,y coordinates
        # stacked=False -> return (N_x, N_y, 2) for regular grids
        raise NotImplementedError

    @property
    def coords_projection(self) -> ccrs.Projection:
        # Must match the CRS of the x,y coordinates returned by get_xy.
        return ccrs.LambertConformal(
            central_longitude=25.0,
            central_latitude=56.7,
        )

    @property
    def num_grid_points(self) -> int:
        gs = self.grid_shape_state
        return gs.x * gs.y

    @cached_property
    def state_feature_weights_values(self) -> List[float]:
        # Return one weight per state feature, used to scale the loss.
        names = self.get_vars_names("state")
        return [1.0] * len(names)

    # ------------------------------------------------------------------ #
    # Required by BaseRegularGridDatastore
    # ------------------------------------------------------------------ #

    @cached_property
    def grid_shape_state(self) -> CartesianGridShape:
        return CartesianGridShape(y=128, x=128)  # adjust to your grid

    def get_lat_lon(self, category: str) -> np.ndarray:
        # Return (N_grid, 2) array of (lat, lon) coordinates
        raise NotImplementedError


print("✅ Template loaded — adapt the NotImplementedError stubs to your dataset")

7. Register Your Datastore#

neural-lam discovers datastores via the kind field in config.yaml. Add your class to the registry in neural_lam/datastore/__init__.py:

registry_example = """# In neural_lam/datastore/__init__.py

from .mdp import MDPDatastore
from .npyfilesmeps.store import NpyFilesDatastoreMEPS
from .my_dataset import MyCustomDatastore          # <-- add this

DATASTORES = {
    MDPDatastore.SHORT_NAME:           MDPDatastore,
    NpyFilesDatastoreMEPS.SHORT_NAME:  NpyFilesDatastoreMEPS,
    MyCustomDatastore.SHORT_NAME:      MyCustomDatastore,   # <-- and this
}
"""
print(registry_example)

8. Wire Into config.yaml#

config_example = """datastore:
  kind: my_dataset          # matches MyCustomDatastore.SHORT_NAME
  config_path: path/to/my_dataset.yaml

training:
  state_feature_weighting:
    __config_class__: ManualStateFeatureWeighting
    weights:
      temperature_2m: 1.0
      wind_u_10m: 1.0
"""
print(config_example)

9. Validate Your Implementation#

This helper checks the mandatory interface before you attempt a full training run.

Note: boundary_mask is a @cached_property on the datastore, not a method — access it as an attribute.

def validate_datastore(datastore):
    """Quick sanity check for a new datastore implementation."""
    print("=== Datastore Validation ===")

    # Variable names
    for category in ["state", "forcing", "static"]:
        names = datastore.get_vars_names(category)
        print(f"  {category:<8s}: {len(names)} features — {names}")

    # DataArray shape
    for split in ["train", "val", "test"]:
        da = datastore.get_dataarray("state", split)
        assert "grid_index" in da.dims, "Missing grid_index dimension"
        assert "time" in da.dims, "Missing time dimension"
        print(f"  {split:<8s}: state shape = {dict(da.sizes)}")

    # Grid positions
    xy = datastore.get_xy("state", stacked=True)
    assert xy.shape[1] == 2, "get_xy must return (N_grid, 2)"
    print(f"  Grid XY shape : {xy.shape}")

    # boundary_mask is a @cached_property, access it as an attribute (not a method call)
    mask = datastore.boundary_mask
    assert mask.dims == ("grid_index",), "boundary_mask must have dim (grid_index,)"
    assert mask.shape[0] == xy.shape[0], "boundary_mask length must equal N_grid"
    print(f"  Boundary mask : {int(mask.sum())} boundary points of {len(mask)}")

    # Grid shape
    gs = datastore.grid_shape_state
    assert gs.y * gs.x == xy.shape[0], "grid_shape_state does not match N_grid"
    print(f"  Grid shape    : {gs.y}×{gs.x}")

    print("✅ Validation passed")


# Usage:
# from neural_lam.config import load_config_and_datastore
# _, datastore = load_config_and_datastore(config_path="path/to/config.yaml")
# validate_datastore(datastore)
print("Call validate_datastore(datastore) after loading your custom datastore.")
Contents