Neural-LAM Data Explorer

Before training a weather model, understanding the input data is essential. This notebook explores the meps_example dataset — real MEPS NWP data over the Nordic region — the same dataset neural-lam’s CI pipeline uses.

What you will see:

• The NWP grid layout over the Nordic domain

• All 18 atmospheric variables at a real timestep

• Static features: geopotential, land mask, border mask

• Normalisation statistics used during training

• How all tensors connect into a single batch

No training required. Runs fully on CPU in ~2 minutes.

import subprocess, sys, zipfile
from pathlib import Path

subprocess.run([sys.executable, "-m", "pip", "install",
                "--quiet", "plotly>=5.20.0", "gdown"], check=True)

import gdown

DATA_ROOT  = Path("meps_example")
ZIP_PATH   = DATA_ROOT / "example_data.zip"
STATIC_DIR = DATA_ROOT / "data/meps_example/static"
SAMPLES_DIR= DATA_ROOT / "data/meps_example/samples"
GRAPH_DIR  = DATA_ROOT / "graphs/1level"

if not STATIC_DIR.exists():
    DATA_ROOT.mkdir(exist_ok=True)
    if not ZIP_PATH.exists():
        print("⬇️  Downloading meps_example (~30 MB metadata + 2.8 GB data)...")
        gdown.download_folder(
            id="1N6ZT_mkfbdVloVsNs9T5YOrMtxd3jG-j",
            output=str(DATA_ROOT), quiet=False, use_cookies=False
        )
    print("📦 Extracting...")
    with zipfile.ZipFile(ZIP_PATH, "r") as z:
        z.extractall(str(DATA_ROOT))

print("✅ Data ready")

import numpy as np
import torch
from pathlib import Path

DATA_ROOT   = Path("meps_example")
STATIC_DIR  = DATA_ROOT / "data/meps_example/static"
SAMPLES_DIR = DATA_ROOT / "data/meps_example/samples"
GRAPH_DIR   = DATA_ROOT / "graphs/1level"

def load_pt(path):
    try:    t = torch.load(path, map_location="cpu", weights_only=True)
    except: t = torch.load(path, map_location="cpu")
    return t[0] if isinstance(t, list) else t

print("=" * 65)
print("meps_example — Complete Data Inventory")
print("=" * 65)
print()
print("📂 STATIC FILES")
print(f"  {'File':<40} {'Shape':<25} {'Notes'}")
print("  " + "-" * 75)
static_notes = {
    "nwp_xy.npy":           "Grid (x,y) positions",
    "surface_geopotential.npy": "Terrain height",
    "border_mask.npy":      "Domain boundary mask",
    "grid_features.pt":     "Precomputed grid node features",
    "parameter_mean.pt":    "Per-variable mean (normalisation)",
    "parameter_std.pt":     "Per-variable std  (normalisation)",
    "diff_mean.pt":         "Per-variable diff mean",
    "diff_std.pt":          "Per-variable diff std",
    "flux_stats.pt":        "TOA flux normalisation stats",
    "parameter_weights.npy":"Per-variable loss weights",
}
for p in sorted(STATIC_DIR.rglob("*")):
    if not p.is_file(): continue
    if p.suffix == ".npy":
        shape = str(np.load(p).shape)
    elif p.suffix == ".pt":
        shape = str(tuple(load_pt(p).shape))
    else:
        shape = f"{p.stat().st_size//1024} KB"
    note = static_notes.get(p.name, "")
    print(f"  {p.name:<40} {shape:<25} {note}")
print()
print("📂 SAMPLE FILES (train split)")
print(f"  {'File':<55} {'Shape'}")
print("  " + "-" * 70)
for p in sorted((SAMPLES_DIR / "train").glob("*.npy")):
    shape = str(np.load(p).shape)
    print(f"  {p.name:<55} {shape}")
print()
print("📂 GRAPH TENSORS (1level)")
graph_notes = {
    "g2m_edge_index": "(2, E_g2m) — grid→mesh edge indices",
    "g2m_features":   "(E_g2m, F) — grid→mesh edge features",
    "m2m_edge_index": "(2, E_m2m) — mesh→mesh edge indices",
    "m2m_features":   "(E_m2m, F) — mesh→mesh edge features",
    "m2g_edge_index": "(2, E_m2g) — mesh→grid edge indices",
    "m2g_features":   "(E_m2g, F) — mesh→grid edge features",
    "mesh_features":  "(N_mesh, F) — mesh node features",
}
for p in sorted(GRAPH_DIR.glob("*.pt")):
    t = load_pt(p)
    print(f"  {p.name:<35} {str(tuple(t.shape)):<25} {graph_notes.get(p.stem,'')}")

import numpy as np
import plotly.graph_objects as go
import plotly.io as pio

pio.renderers.default = "colab"

# nwp_xy shape: (2, 268, 238) — first dim is (x, y)
xy  = np.load(STATIC_DIR / "nwp_xy.npy")   # (2, H, W)
H, W = xy.shape[1], xy.shape[2]
N_GRID = H * W

x_flat = xy[0].flatten()
y_flat = xy[1].flatten()

# Normalize to [0,1]
xn = (x_flat - x_flat.min()) / (x_flat.max() - x_flat.min())
yn = (y_flat - y_flat.min()) / (y_flat.max() - y_flat.min())

# Water mask for colouring
wtr = np.load(list((SAMPLES_DIR / "train").glob("wtr_*.npy"))[0])  # (268,238)
wtr_flat = wtr.flatten().astype(float)

MAX = 6000
idx = np.random.default_rng(42).choice(N_GRID, MAX, replace=False)

fig = go.Figure(go.Scatter(
    x=xn[idx], y=yn[idx],
    mode="markers",
    marker=dict(
        size=2.5,
        color=wtr_flat[idx],
        colorscale=[[0,"#1e40af"], [1,"#16a34a"]],
        opacity=0.75,
        showscale=True,
        colorbar=dict(
            title="Land (1) / Sea (0)",
            tickvals=[0, 1], ticktext=["Sea", "Land"],
            tickfont=dict(color="#94a3b8"),
            titlefont=dict(color="#94a3b8"),
            thickness=12, len=0.5
        )
    ),
    hovertemplate="x=%{x:.3f}  y=%{y:.3f}<extra></extra>"
))
fig.update_layout(
    title=dict(
        text=f"MEPS NWP Grid — {N_GRID:,} observation points  ({H}×{W} grid)",
        font=dict(size=16, color="#ffffff")
    ),
    height=520, paper_bgcolor="#0f172a", plot_bgcolor="#0f172a",
    font=dict(color="#ffffff"),
    xaxis=dict(title="X (normalized)", gridcolor="#334155",
               zeroline=False, color="#94a3b8", scaleanchor="y"),
    yaxis=dict(title="Y (normalized)", gridcolor="#334155",
               zeroline=False, color="#94a3b8"),
    margin=dict(l=60, r=20, t=60, b=60)
)
fig.show(renderer="colab")

print()
print(f"   Grid size  : {H} × {W} = {N_GRID:,} nodes")
print(f"   X range    : {x_flat.min():.1f} → {x_flat.max():.1f}  (raw projection units)")
print(f"   Y range    : {y_flat.min():.1f} → {y_flat.max():.1f}")
print(f"   🟦 Blue = sea grid points")
print(f"   🟩 Green = land grid points")

NWP State Variables

Each .npy sample file has shape (65, 268, 238, 18):

• 65 — timesteps (every 6 hours)

• 268 × 238 — spatial grid (the MEPS Nordic domain)

• 18 — atmospheric variables (17 state + 1 extra channel)

The 17 core state variables (from parameter_mean.pt):

Idx	Variable	Physical meaning
0	u10	10m eastward wind
1	v10	10m northward wind
2	t2m	2m temperature
3–16	Pressure level vars	Wind, temp, humidity at multiple levels

The 18th channel and nwp_toa_downwelling_shortwave_flux are forcing variables — known future inputs (solar radiation, boundary conditions) fed to the model at each step.

import numpy as np
import plotly.subplots as sp
import plotly.graph_objects as go
import plotly.io as pio

pio.renderers.default = "colab"

# Load one NWP sample — shape (65, 268, 238, 18)
nwp_path = sorted((SAMPLES_DIR / "train").glob("nwp_*_mbr000.npy"))[0]
nwp = np.load(nwp_path)  # (65, 268, 238, 18)

T_IDX = 0  # first timestep
H, W  = 268, 238

# Pick 4 most interpretable variables
VARS = {
    "u10 — 10m East Wind (m/s)":     (nwp[T_IDX, :, :, 0],  "RdBu_r"),
    "v10 — 10m North Wind (m/s)":    (nwp[T_IDX, :, :, 1],  "RdBu_r"),
    "t2m — 2m Temperature (K)":      (nwp[T_IDX, :, :, 2],  "thermal"),
    "Var 3 — Pressure Level Field":  (nwp[T_IDX, :, :, 3],  "Viridis"),
}

fig = sp.make_subplots(
    rows=2, cols=2,
    subplot_titles=list(VARS.keys()),
    horizontal_spacing=0.12,
    vertical_spacing=0.18
)

for i, (title, (field, cmap)) in enumerate(VARS.items()):
    row, col = i // 2 + 1, i % 2 + 1
    fig.add_trace(
        go.Heatmap(
            z=field,
            colorscale=cmap,
            showscale=True,
            colorbar=dict(
                len=0.42, thickness=10,
                x=0.46 if col == 1 else 1.02,
                y=0.78 if row == 1 else 0.22,
                tickfont=dict(color="#94a3b8", size=8),
            ),
            hovertemplate="row=%{y}  col=%{x}<br>value=%{z:.2f}<extra></extra>"
        ),
        row=row, col=col
    )
    fig.update_xaxes(showticklabels=False, row=row, col=col)
    fig.update_yaxes(showticklabels=False, row=row, col=col)

fig.update_layout(
    title=dict(
        text=f"NWP State Variables — Timestep {T_IDX}  ({nwp_path.name})",
        font=dict(size=15, color="#ffffff")
    ),
    height=600,
    paper_bgcolor="#0f172a",
    plot_bgcolor="#0f172a",
    font=dict(color="#ffffff"),
    margin=dict(l=20, r=60, t=80, b=20)
)
for ann in fig.layout.annotations:
    ann.font.color = "#94a3b8"
    ann.font.size  = 11

fig.show(renderer="colab")
print(f"   NWP array shape : {nwp.shape}")
print(f"   Showing timestep: {T_IDX} of {nwp.shape[0]}")
print(f"   Variable range  : {nwp.min():.1f} → {nwp.max():.1f}  (raw, unnormalised)")

import numpy as np
import plotly.subplots as sp
import plotly.graph_objects as go
import plotly.io as pio

pio.renderers.default = "colab"

geo  = np.load(STATIC_DIR / "surface_geopotential.npy")  # (268, 238)
brd  = np.load(STATIC_DIR / "border_mask.npy")           # (268, 238)
wtr  = np.load(list((SAMPLES_DIR/"train").glob("wtr_*.npy"))[0])  # (268,238)

static_fields = {
    "Surface Geopotential (terrain)": (geo,  "earth"),
    "Border Mask":                    (brd,  "Greys"),
    "Water Mask (land/sea)":          (wtr,  "Blues_r"),
}

fig = sp.make_subplots(
    rows=1, cols=3,
    subplot_titles=list(static_fields.keys()),
    horizontal_spacing=0.08
)

for i, (title, (field, cmap)) in enumerate(static_fields.items()):
    fig.add_trace(
        go.Heatmap(
            z=field,
            colorscale=cmap,
            showscale=True,
            colorbar=dict(
                len=0.7, thickness=10,
                x=0.30*i + 0.28,
                tickfont=dict(color="#94a3b8", size=8),
            ),
            hovertemplate="row=%{y}  col=%{x}<br>%{z:.3f}<extra></extra>"
        ),
        row=1, col=i+1
    )
    fig.update_xaxes(showticklabels=False, row=1, col=i+1)
    fig.update_yaxes(showticklabels=False, row=1, col=i+1)

fig.update_layout(
    title=dict(text="Static Features — Fixed inputs to the model",
               font=dict(size=15, color="#ffffff")),
    height=380,
    paper_bgcolor="#0f172a", plot_bgcolor="#0f172a",
    font=dict(color="#ffffff"),
    margin=dict(l=20, r=60, t=60, b=20)
)
for ann in fig.layout.annotations:
    ann.font.color = "#94a3b8"
    ann.font.size  = 11

fig.show(renderer="colab")
print("   Geopotential — encodes terrain height (mountains slow the wind)")
print("   Border mask  — marks boundary nodes where forcing is applied")
print("   Water mask   — used by the model to distinguish land/sea physics")

import torch
import numpy as np
import plotly.graph_objects as go
import plotly.subplots as sp
import plotly.io as pio

pio.renderers.default = "colab"

param_mean = load_pt(STATIC_DIR / "parameter_mean.pt").numpy()  # (17,)
param_std  = load_pt(STATIC_DIR / "parameter_std.pt").numpy()   # (17,)
weights    = np.load(STATIC_DIR / "parameter_weights.npy")       # (17,)

VAR_NAMES = [
    "u10","v10","t2m",
    "u_500","v_500","z_500","t_500",
    "u_700","v_700","z_700","t_700",
    "u_850","v_850","z_850","t_850",
    "u_925","v_925",
]
labels = VAR_NAMES[:len(param_mean)]

fig = sp.make_subplots(
    rows=1, cols=2,
    subplot_titles=["Normalisation: Mean ± Std per variable",
                    "Loss Weights per variable"],
    horizontal_spacing=0.12
)

# Mean ± Std
fig.add_trace(go.Bar(
    x=labels, y=param_mean,
    error_y=dict(type="data", array=param_std,
                 color="#f97316", thickness=1.5, width=4),
    marker_color="#3b82f6", name="mean ± std",
    hovertemplate="%{x}<br>mean=%{y:.2f}<extra></extra>"
), row=1, col=1)

# Weights
fig.add_trace(go.Bar(
    x=labels, y=weights,
    marker_color=[
        "#ef4444" if w == weights.max() else "#10b981"
        for w in weights
    ],
    name="loss weight",
    hovertemplate="%{x}<br>weight=%{y:.4f}<extra></extra>"
), row=1, col=2)

for col in [1, 2]:
    fig.update_xaxes(tickangle=-45, color="#94a3b8",
                     gridcolor="#1e293b", row=1, col=col)
    fig.update_yaxes(color="#94a3b8",
                     gridcolor="#334155", row=1, col=col)

fig.update_layout(
    height=420,
    paper_bgcolor="#0f172a", plot_bgcolor="#0f172a",
    font=dict(color="#ffffff"),
    showlegend=False,
    margin=dict(l=60, r=20, t=60, b=80)
)
for ann in fig.layout.annotations:
    ann.font.color = "#94a3b8"

fig.show(renderer="colab")

print()
print(f"   17 variables · means range {param_mean.min():.1f} → {param_mean.max():.1f}")
print(f"   std range: {param_std.min():.2f} → {param_std.max():.2f}")
print()
print(f"   Higher loss weight = model penalised more for errors in that variable")
print(f"   Max weight variable: {labels[int(weights.argmax())]}  ({weights.max():.4f})")

Complete Data Flow

meps_example/
│
├── data/meps_example/
│   ├── static/
│   │   ├── nwp_xy.npy           (2, 268, 238)   — grid positions
│   │   ├── surface_geopotential (268, 238)       — terrain
│   │   ├── border_mask.npy      (268, 238)       — boundary
│   │   ├── grid_features.pt     (63784, 4)       — precomputed node feats
│   │   ├── parameter_mean.pt    (17,)            — normalisation
│   │   └── parameter_std.pt     (17,)            — normalisation
│   │
│   └── samples/train/
│       ├── nwp_*_mbr000.npy     (65, 268, 238, 18)  — state variables
│       ├── nwp_toa_*.npy        (65, 268, 238)       — solar forcing
│       └── wtr_*.npy            (268, 238)           — water mask
│
└── graphs/1level/
    ├── g2m_edge_index.pt   (2, 100656)  — grid→mesh
    ├── m2m_edge_index.pt   (2, 51520)   — mesh→mesh
    └── m2g_edge_index.pt   (2, 255136)  — mesh→grid

What WeatherDataset does with this:

1. Load nwp[t:t+2] as init states — shape (2, 63784, 17)

2. Normalise using parameter_mean / parameter_std

3. Append grid_features + surface_geopotential as static context

4. Load nwp_toa[t+2:t+T] as forcing — known future solar radiation

5. Return (init_states, target_states, forcing, static_features)

Next: See these tensors flow through the model → forecast_visualiser.ipynb