Hello World: Training neural-lam on DANRA

This notebook provides a beginner-friendly, end-to-end walkthrough for running a minimal model training pipeline in neural-lam using a small, public DANRA dataset.

It covers:

1. Environment setup (CPU-safe)

2. Data preprocessing with mllam-data-prep

3. Graph generation (single-level, for speed)

4. Training for 1 epoch on CPU

5. Evaluation and example predictions

6. Scaling tips for bigger runs

---

Prerequisites and Context

Important: This notebook is designed to be run from inside a local clone of mllam/neural-lam. All paths are relative to the repository root.

Paper Reference: For an in-depth context on the models used here, please refer to the paper
“Building Machine Learning Limited Area Models: Kilometer-Scale Weather Forecasting in Realistic Settings”
(Adamov et al., 2025) available at arXiv:2504.09340.

Python Version: Make sure you are using a Python version supported by the project (e.g. 3.10-3.12) with ipykernel installed.

Note on Future Graph Updates: The graph generation step currently uses create_graph.py, but it will be migrated to use the upcoming weather-model-graphs package in the near future.

Quick Run Cheatsheet

Item	Details
Runtime	~15 min on a 6-12 core laptop CPU (data prep + 1 training epoch).
Downloads	~60-120 MB streamed from the ECMWF object store. Cached under tests/datastore_examples/mdp/danra_100m_winds/.
Artifacts	danra.datastore.zarr, graph/1level/, saved_models/, and wandb/offline-run-*.
Resetting	Delete the .zarr, graph/, or saved_models/ directories to rerun from scratch.

Kernel Verification: Ensure you select a Python runtime (e.g., Python 3.10-3.12) that matches your local pyproject.toml specs. Make sure ipykernel or jupyter are installed locally so the notebook connects correctly to the right environment dependencies.\n

import glob
import os
import sys

import matplotlib.pyplot as plt
import xarray as xr
from IPython.display import Image, display

import neural_lam

# Verify we are at the repo root (should see neural_lam/, docs/, tests/, etc.)
print("Current directory:", os.getcwd())
print("Repo contents:", [f for f in os.listdir(".") if not f.startswith(".")])

# The notebook is typically located in docs/notebooks/.
# Change the working directory to the repository root so that paths like 'tests/...' resolve correctly.
if os.getcwd().endswith("notebooks"):
    os.chdir("../../")
print("Current working directory:", os.getcwd())

1. Environment Setup

You have two options for installation: Pick exactly one of the installers below. After it finishes, restart the kernel so sys.executable points at the fresh environment.

Option A (recommended): uv - if you followed the README setup:

%%bash
uv venv --no-project
uv pip install torch --index-url https://download.pytorch.org/whl/cpu
uv pip install .

Option B: pip - run the cell below if you haven’t installed yet:

# CPU-safe torch install (skip if you already have torch installed)
!{sys.executable} -m pip install torch --index-url https://download.pytorch.org/whl/cpu --quiet

# Install neural-lam and mllam-data-prep from the repo root
!{sys.executable} -m pip install -e . --quiet
!{sys.executable} -m pip install mllam-data-prep xarray matplotlib networkx --quiet

# Note: Once released, `weather-model-graphs` can also be installed here.

# Verify the install
print("neural-lam version:", neural_lam.__version__)

2. Data Configuration & Preprocessing

We use the DANRA 100m winds example already in this repo at:

tests/datastore_examples/mdp/danra_100m_winds/
  ├── config.yaml          <- neural-lam configuration
  └── danra.datastore.yaml <- mllam-data-prep datastore config

Aspect	Details
Spatial crop	~100x100 nodes centered on Denmark at 100 m resolution.
Time span	73 hourly steps in April 2022 (10 days).
Variables	4 state channels + 1 forcing feature pre-selected in config.yaml.
On-disk size	Processed .zarr is ~23 MB and lives next to the configs.

This example uses a cropped DANRA dataset (~100x100 grid points, ~10 days in April 2022) served from a public ECMWF object store - no local data download is required.

The mllam-data-prep command reads the datastore config, fetches the data, and writes a processed .zarr archive to disk.

Version requirement: This notebook requires mllam-data-prep >= 0.6.0. Check with python -m mllam_data_prep --version or upgrade with pip install --upgrade mllam-data-prep.\n \n Note: This will download and process approximately 60-120 MB of data. It may take a few minutes on the first run.

# Run mllam-data-prep to create the processed DANRA .zarr dataset
# Output will appear next to danra.datastore.yaml as danra.datastore.zarr
!{sys.executable} -m mllam_data_prep \
    tests/datastore_examples/mdp/danra_100m_winds/danra.datastore.yaml

# Confirm the .zarr dataset was created
zarr_path = "tests/datastore_examples/mdp/danra_100m_winds/danra.datastore.zarr"
if os.path.exists(zarr_path):
    print("✅ Data preprocessing complete. zarr dataset found at:", zarr_path)
else:
    print("❌ zarr dataset not found - check the output above for errors.")

# Optional: Visualize the output datastore
ds = xr.open_zarr("tests/datastore_examples/mdp/danra_100m_winds/danra.datastore.zarr")
print(ds)

# Plot a single temporal snapshot of a state feature (e.g. T2M)
# Note: This is a simplified visualization using flat grid_indexes
if "state" in ds:
    first_time_idx = 0
    feature_idx = 0
    field = ds["state"].isel(time=first_time_idx, state_feature=feature_idx).values
    plt.figure(figsize=(6, 4))
    import matplotlib.colors as mcolors
    # Calculate grid dimensions dynamically (approx square)
    grid_side = int(len(field)**0.5)
    field_2d = field[:grid_side*grid_side].reshape((grid_side, grid_side))
    im = plt.imshow(field_2d, cmap="viridis")
    plt.colorbar(im, label="Value")
    plt.title(f"State Feature {feature_idx} at timestep {first_time_idx} (2D Projection)")
    plt.xlabel("Grid X")
    plt.ylabel("Grid Y")
    plt.show()\n

3. Graph Generation

neural-lam uses a graph to define the message-passing structure of the GNN. Different graph types exist:

Graph type	Command flag	Use case
L1-LAM	--name 1level --levels 1	Quick demo (this notebook)
GC-LAM	--name multiscale	Standard multi-scale model
Hi-LAM	--name hierarchical --hierarchical	Hierarchical model (production)

For this Hello World example we use the L1 (single-level) graph - the lightest option, ideal for CPU.

# Generate the single-level graph for fast CPU execution
# Graph files are stored in tests/datastore_examples/mdp/danra_100m_winds/graph/1level/
!{sys.executable} -m neural_lam.create_graph \
    --config_path tests/datastore_examples/mdp/danra_100m_winds/config.yaml \
    --name 1level \
    --levels 1

# Confirm the graph was created
graph_files = glob.glob(
    "tests/datastore_examples/mdp/danra_100m_winds/graph/1level/**", recursive=True
)
print(f"Graph files created ({len(graph_files)}):", graph_files)

graph_dir = "tests/datastore_examples/mdp/danra_100m_winds/graph/1level"
graph_files = glob.glob(os.path.join(graph_dir, "**", "*.pt"), recursive=True)

if graph_files:
    print(f"✅ Graph created with {len(graph_files)} tensor file(s):")
    for f in sorted(graph_files):
        size_kb = os.path.getsize(f) / 1024
        print(f"  {os.path.basename(f):40s} {size_kb:.1f} KB")
else:
    print("❌ No graph .pt files found - check the create_graph output above.")

# Visualize the generated graph topology directly onto a 2D plot.\n!{sys.executable} -m neural_lam.plot_graph \\\n    --config_path tests/datastore_examples/mdp/danra_100m_winds/config.yaml \\\n    --graph 1level\n

4. Training on CPU

We train the graph_lam model for 1 epoch using the L1 graph. Key flags used here:

Flag	Value	Why
--model	graph_lam	Simplest model - compatible with non-hierarchical graphs
--graph	1level	Must match the graph name used in the previous step
--epochs	1	Minimal run - just to verify the pipeline works end-to-end
--processor_layers	2	Reduced from default (4) for CPU
--ar_steps_train	1	Unroll 1 time-step during training - reduces memory and compute
--ar_steps_eval	1	Also 1 step for the val pass within this demo run

About logging: By default, neural-lam logs via Weights & Biases. To suppress W&B upload during this demo, we set WANDB_MODE=offline so metrics are saved locally without requiring a W&B account or login.

# Set W&B to offline mode for this demo (no account required)
# Remove or change to WANDB_MODE=online to enable cloud logging
os.environ["WANDB_MODE"] = "offline"

!{sys.executable} -m neural_lam.train_model \
    --config_path tests/datastore_examples/mdp/danra_100m_winds/config.yaml \
    --model graph_lam \
    --graph 1level \
    --epochs 1 \
    --processor_layers 2 \
    --ar_steps_train 1 \
    --ar_steps_eval 1 \
    --num_workers 2 \
    --batch_size 2 \
    --val_steps_to_log 1\n

# Find the checkpoint saved during training
ckpts = glob.glob("saved_models/**/*.ckpt", recursive=True)
print("Checkpoints found:", ckpts)

5. Evaluation & Visualization

Evaluation reuses the same train_model command with --eval test and the --load flag pointing to the checkpoint from training.

Example predictions are plotted automatically and saved by the logger. Set --n_example_pred to control how many prediction plots to produce.

ckpts = glob.glob("saved_models/**/*.ckpt", recursive=True)

if not ckpts:
    print("❌ No checkpoint found - make sure the training cell above completed without errors.")
    ckpt_path = None
else:
    ckpt_path = max(ckpts, key=os.path.getmtime)
    print(f"✅ Using checkpoint: {ckpt_path}")

if not ckpt_path:
    print("Skipping evaluation because no checkpoint was found in saved_models/. Run the training cell above first.")
else:
    os.environ["WANDB_MODE"] = "offline"

    !{sys.executable} -m neural_lam.train_model \
        --config_path tests/datastore_examples/mdp/danra_100m_winds/config.yaml \
        --model graph_lam \
        --graph 1level \
        --eval test \
        --load {ckpt_path} \
        --n_example_pred 2 \
        --ar_steps_eval 4 \
        --num_workers 2 \
        --batch_size 2 \
        --val_steps_to_log 1

pred_images = sorted(glob.glob("**/*pred*.png", recursive=True)) + sorted(
    glob.glob("**/*example*.png", recursive=True)
)

if pred_images:
    print(f"Found {len(pred_images)} prediction plot(s):")
    for img_path in pred_images[:4]:
        print(f"  {img_path}")
        display(Image(filename=img_path))
else:
    print("No prediction images found yet.")
    print("Check wandb/offline-run-*/ or saved_models/ for saved plots.")

6. Scaling Tips & Next Steps

🔁 Use a larger / hierarchical graph

For a production-quality run, switch to the hierarchical Hi-LAM graph:

python -m neural_lam.create_graph \
    --config_path <your_config.yaml> \
    --name hierarchical \
    --hierarchical

python -m neural_lam.train_model \
    --config_path <your_config.yaml> \
    --model hi_lam \
    --graph hierarchical \
    --epochs 200

⚡ Enable GPU training

Simply replace the CPU-only torch install with the CUDA variant and --devices will auto-detect your GPU:

pip install torch --index-url https://download.pytorch.org/whl/cu121

📦 Use larger / full DANRA data

Modify danra.datastore.yaml to extend the coord_ranges.time window, or point inputs[*].path at a larger dataset. See the mllam-data-prep README for full configuration options.

For large datasets (>=10 GB), use parallel preprocessing:

python -m mllam_data_prep \
    <your_datastore.yaml> \
    --dask-distributed-local-core-fraction 0.5

📊 Enable cloud logging with W&B

Remove the WANDB_MODE=offline line (or set it to online) and optionally set --logger wandb --logger-project <your-project-name>.

🔗 Further reading

• neural-lam README

• mllam-data-prep

• Graph-based Neural Weather Prediction (NeurIPS 2023)

• Probabilistic Weather Forecasting with Hierarchical GNNs (NeurIPS 2024)\n