Datastore Explorer

Datastore Explorer#

This notebook walks through the structure of the DANRA zarr datastore used in neural-lam. It covers dimensions, coordinates, time splits, and 2D visualizations of each feature.

Prerequisites: Run the Hello World notebook first to generate danra.datastore.zarr.

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

1. Setup#

%matplotlib inline
import os

import matplotlib.pyplot as plt
import numpy as np
import xarray as xr

# Resolve repo root
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. Load the Datastore#

We open the zarr archive directly for inspection, and separately load the datastore object (needed to recover the true 2-D grid shape).

ZARR_PATH   = "tests/datastore_examples/mdp/danra_100m_winds/danra.datastore.zarr"
CONFIG_PATH = "tests/datastore_examples/mdp/danra_100m_winds/config.yaml"

if not os.path.exists(ZARR_PATH):
    raise FileNotFoundError(
        f"Zarr dataset not found at {ZARR_PATH}.\n"
        "Run the Hello World notebook first to generate it."
    )

ds = xr.open_zarr(ZARR_PATH, consolidated=True)
print("✅ Zarr datastore loaded")
print(ds)

from neural_lam.config import load_config_and_datastore

_, datastore = load_config_and_datastore(config_path=CONFIG_PATH)

# Recover true 2-D grid dimensions from the datastore.
# Never derive grid shape by taking the square root of grid_index — the grid
# is not guaranteed to be square.
grid_shape = datastore.grid_shape_state
grid_y, grid_x = grid_shape.y, grid_shape.x
assert grid_y * grid_x == ds.sizes["grid_index"], (
    f"Grid shape mismatch: {grid_y}×{grid_x}={grid_y*grid_x} "
    f"!= zarr grid_index size {ds.sizes['grid_index']}"
)
print(f"Grid shape: {grid_y} rows × {grid_x} cols = {grid_y * grid_x} points")

3. Dimensions and Coordinates#

The datastore uses a flat grid_index dimension for spatial data — a 1-D representation of the 2-D grid obtained by stacking the x and y axes.

print("=== Dimensions ===")
for dim, size in ds.sizes.items():
    print(f"  {dim:<25s}: {size}")

print("\n=== Coordinates ===")
for coord in ds.coords:
    print(f"  {coord:<35s}: {ds.coords[coord].dtype}")

print("\n=== Data Variables ===")
for var in ds.data_vars:
    print(f"  {var:<40s}: {list(ds[var].dims)}")

4. Time Splits#

The dataset is pre-split into train / val / test partitions along the time dimension by mllam-data-prep.

The split boundaries are stored in the splits data variable, indexed by split_name and split_part.

print("=== Time Splits ===")
for split in ["train", "val", "test"]:
    da = ds.splits.sel(split_name=split)
    start = da.sel(split_part="start").values.item()
    end   = da.sel(split_part="end").values.item()
    print(f"  {split:<8s}: {start}  →  {end}")

print(f"\n  Total timesteps: {ds.sizes['time']}")

5. State Features#

State variables are the atmospheric fields the model predicts (wind components, temperature, relative humidity).

feature_names = ds.state_feature.values.tolist()
has_units     = "state_feature_units"     in ds.coords
has_lname     = "state_feature_long_name" in ds.coords

print("=== State Features ===")
for i, feat in enumerate(feature_names):
    units = ds.state_feature_units.values[i]     if has_units else "—"
    lname = ds.state_feature_long_name.values[i] if has_lname else feat
    print(f"  [{i}] {feat:<10s}  units={units:<10s}  ({lname})")

6. Visualize State Features#

We reshape the flat grid_index dimension back into the true 2-D grid shape using grid_shape_state from the datastore (loaded in Section 2).

n_features = ds.sizes["state_feature"]

fig, axes = plt.subplots(1, n_features, figsize=(5 * n_features, 4))
if n_features == 1:
    axes = [axes]

for idx, (ax, feat_name) in enumerate(zip(axes, feature_names)):
    field    = ds["state"].isel(time=0, state_feature=idx).values
    field_2d = field.reshape(grid_y, grid_x)
    im = ax.imshow(field_2d, cmap="viridis", origin="lower")
    plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    ax.set_title(f"{feat_name}\n(t=0)")
    ax.set_xlabel("Grid X")
    ax.set_ylabel("Grid Y")

plt.suptitle("State Features — First Timestep (2022-04-01 00:00)", fontsize=13)
plt.tight_layout()
plt.show()

7. Forcing and Static Variables#

print("=== Forcing Features ===")
for feat in ds.forcing_feature.values:
    print(f"  {feat}")

print("\n=== Static Features ===")
for feat in ds.static_feature.values:
    print(f"  {feat}")

n_static = ds.sizes["static_feature"]
fig, axes = plt.subplots(1, n_static, figsize=(5 * n_static, 4))
if n_static == 1:
    axes = [axes]
for idx, (ax, feat) in enumerate(zip(axes, ds.static_feature.values)):
    field = ds["static"].isel(static_feature=idx).values.reshape(grid_y, grid_x)
    im = ax.imshow(field, cmap="terrain", origin="lower")
    plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    ax.set_title(feat)
plt.suptitle("Static Features", fontsize=13)
plt.tight_layout()
plt.show()

8. Normalisation Statistics#

mllam-data-prep computes per-feature mean and standard deviation over the training split. These are stored in the zarr as data variables with the naming convention <category>__train__<stat>.

print("=== State Feature — Training Statistics ===")
means = ds["state__train__mean"].values
stds  = ds["state__train__std"].values
for feat, mean, std in zip(feature_names, means, stds):
    print(f"  {feat:<10s}  mean={mean:+.4f}  std={std:.4f}")
Contents