Neural-LAM Graph Structure — Interactive 3D Visualisation

Neural-LAM converts the NWP grid into a graph for GNN message passing. This notebook shows that graph in 3D so you can see exactly what the model operates on.

Part 1 — Concept (synthetic 9-node graph): Understand the structure clearly. Part 2 — Real data (meps_example): See the actual neural-lam graph.

All visualisations are interactive — rotate, zoom, hover.

import subprocess, sys
subprocess.run([sys.executable, "-m", "pip", "install",
                "--quiet", "neural-lam", "plotly>=5.20.0", "gdown"],
               check=True)
print("✅ Ready")

Part 1 — The Concept (Synthetic Graph)

Before looking at the real graph with 63,784 nodes, here is the same structure at human scale: 9 mesh nodes, ~72 grid nodes, all edges visible.

This is how GraphLAM works:

1. Encode — each grid node sends its features to nearby mesh nodes (g2m)

2. Process — mesh nodes exchange information with neighbours (m2m)

3. Decode — each mesh node sends predictions back to nearby grid nodes (m2g)

import numpy as np
import torch

np.random.seed(42)

# 3×3 mesh grid — 9 nodes
mesh_xy = np.array([
    [0.2, 0.2], [0.5, 0.2], [0.8, 0.2],
    [0.2, 0.5], [0.5, 0.5], [0.8, 0.5],
    [0.2, 0.8], [0.5, 0.8], [0.8, 0.8],
], dtype=np.float32)

# ~72 grid nodes — 8 per mesh cell
grid_pts = []
for mx, my in mesh_xy:
    for _ in range(8):
        gx = np.clip(mx + np.random.uniform(-0.13, 0.13), 0.02, 0.98)
        gy = np.clip(my + np.random.uniform(-0.13, 0.13), 0.02, 0.98)
        grid_pts.append([gx, gy])
grid_xy = np.array(grid_pts, dtype=np.float32)

# g2m: each grid node → nearest mesh node
g2m_src = [np.argmin(np.linalg.norm(mesh_xy - g, axis=1))
           for g in grid_xy]
g2m = torch.tensor([list(range(len(grid_xy))), g2m_src], dtype=torch.long)

# m2m: each mesh node → neighbours within distance 0.35
m2m_src, m2m_dst = [], []
for i, a in enumerate(mesh_xy):
    for j, b in enumerate(mesh_xy):
        if i != j and np.linalg.norm(a - b) < 0.35:
            m2m_src.append(i); m2m_dst.append(j)
m2m = torch.tensor([m2m_src, m2m_dst], dtype=torch.long)

# m2g: reverse of g2m
m2g = g2m[[1, 0]]

SYN = {"grid_xy": grid_xy, "mesh_xy": mesh_xy,
       "g2m": g2m, "m2m": m2m, "m2g": m2g}

print(f"✅ Synthetic graph built")
print(f"   Grid : {len(grid_xy)} nodes")
print(f"   Mesh : {len(mesh_xy)} nodes")
print(f"   g2m  : {g2m.shape[1]} edges")
print(f"   m2m  : {m2m.shape[1]} edges")
print(f"   m2g  : {m2g.shape[1]} edges")

def lines(edges, src_xyz, dst_xyz):
    xl, yl, zl = [], [], []
    for s, d in zip(edges[0].numpy(), edges[1].numpy()):
        xl += [src_xyz[s,0], dst_xyz[d,0], None]
        yl += [src_xyz[s,1], dst_xyz[d,1], None]
        zl += [src_xyz[s,2], dst_xyz[d,2], None]
    return xl, yl, zl

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

pio.renderers.default = "colab"

g = SYN["grid_xy"]; m = SYN["mesh_xy"]
g_xyz = np.column_stack([g, np.zeros(len(g))])
m_xyz = np.column_stack([m, np.full(len(m), 0.5)])

def dark_layout(title, z_labels=None):
    z_axis = dict(title="Layer", backgroundcolor="#1e293b",
                  gridcolor="#334155", color="#94a3b8")
    if z_labels:
        z_axis.update(tickvals=z_labels[0], ticktext=z_labels[1])
    return dict(
        title=dict(text=title, font=dict(size=17, color="#ffffff")),
        height=600, paper_bgcolor="#0f172a", font=dict(color="#ffffff"),
        scene=dict(
            bgcolor="#0f172a",
            xaxis=dict(title="X", backgroundcolor="#1e293b",
                       gridcolor="#334155", zeroline=False,
                       color="#94a3b8", range=[-0.05, 1.05]),
            yaxis=dict(title="Y", backgroundcolor="#1e293b",
                       gridcolor="#334155", zeroline=False,
                       color="#94a3b8", range=[-0.05, 1.05]),
            zaxis=z_axis,
            camera=dict(eye=dict(x=1.6, y=-1.2, z=1.2)),
            aspectmode="manual", aspectratio=dict(x=1, y=1, z=0.55)
        ),
        legend=dict(bgcolor="rgba(15,23,42,0.9)", bordercolor="#334155",
                    borderwidth=1, orientation="h",
                    x=0.5, xanchor="center", y=-0.05),
        margin=dict(l=0, r=0, t=60, b=0)
    )

fig = go.Figure([
    go.Scatter3d(x=g_xyz[:,0], y=g_xyz[:,1], z=g_xyz[:,2],
                 mode="markers", name="Grid nodes (observations)",
                 marker=dict(size=5, color="#3b82f6", opacity=0.85),
                 hovertemplate="Grid #%{pointNumber}<extra></extra>"),
    go.Scatter3d(x=m_xyz[:,0], y=m_xyz[:,1], z=m_xyz[:,2],
                 mode="markers+text", name="Mesh nodes (latent)",
                 text=[f"M{i}" for i in range(len(m))],
                 textposition="top center",
                 textfont=dict(color="#ffffff", size=9),
                 marker=dict(size=14, color="#f97316", opacity=1.0,
                             line=dict(color="#ffffff", width=1)),
                 hovertemplate="Mesh #%{pointNumber}<extra></extra>"),
])
fig.update_layout(**dark_layout(
    "Step 1 — Two Layers: Grid (observations) and Mesh (latent)",
    z_labels=([0, 0.5], ["Grid — observations", "Mesh — latent space"])
))
fig.show(renderer="colab")
print("🔵 72 grid nodes carry raw NWP observations")
print("🟠  9 mesh nodes form the latent space for GNN processing")

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

pio.renderers.default = "colab"

xl, yl, zl = lines(SYN["g2m"], g_xyz, m_xyz)

fig = go.Figure([
    go.Scatter3d(x=xl, y=yl, z=zl, mode="lines",
                 name="g2m edges", line=dict(color="#a78bfa", width=1.5),
                 opacity=0.55, hoverinfo="skip"),
    go.Scatter3d(x=g_xyz[:,0], y=g_xyz[:,1], z=g_xyz[:,2],
                 mode="markers", name="Grid nodes",
                 marker=dict(size=5, color="#3b82f6", opacity=0.85),
                 hoverinfo="skip"),
    go.Scatter3d(x=m_xyz[:,0], y=m_xyz[:,1], z=m_xyz[:,2],
                 mode="markers+text", name="Mesh nodes",
                 text=[f"M{i}" for i in range(len(SYN["mesh_xy"]))],
                 textposition="top center",
                 textfont=dict(color="#ffffff", size=9),
                 marker=dict(size=14, color="#f97316", opacity=1.0,
                             line=dict(color="#ffffff", width=1)),
                 hoverinfo="skip"),
])
fig.update_layout(**dark_layout(
    "Step 2 — Encode: Grid → Mesh (g2m edges)",
    z_labels=([0, 0.5], ["Grid", "Mesh"])
))
fig.add_annotation(
    text="🟣 Each purple line = one grid observation aggregated into a mesh node",
    xref="paper", yref="paper", x=0.5, y=1.06,
    xanchor="center", showarrow=False,
    font=dict(color="#94a3b8", size=12)
)
fig.show(renderer="colab")
print(f"   {SYN['g2m'].shape[1]} g2m edges — every grid node feeds exactly 1 mesh node")

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

pio.renderers.default = "colab"

xl_mm, yl_mm, zl_mm = lines(SYN["m2m"], m_xyz, m_xyz)

fig = go.Figure([
    go.Scatter3d(x=g_xyz[:,0], y=g_xyz[:,1], z=g_xyz[:,2],
                 mode="markers", name="Grid nodes (faded)",
                 marker=dict(size=4, color="#3b82f6", opacity=0.2),
                 hoverinfo="skip"),
    go.Scatter3d(x=xl_mm, y=yl_mm, z=zl_mm,
                 mode="lines", name="m2m edges",
                 line=dict(color="#10b981", width=3),
                 opacity=0.8, hoverinfo="skip"),
    go.Scatter3d(x=m_xyz[:,0], y=m_xyz[:,1], z=m_xyz[:,2],
                 mode="markers+text", name="Mesh nodes",
                 text=[f"M{i}" for i in range(len(SYN["mesh_xy"]))],
                 textposition="top center",
                 textfont=dict(color="#ffffff", size=9),
                 marker=dict(size=14, color="#f97316", opacity=1.0,
                             line=dict(color="#ffffff", width=1)),
                 hovertemplate="Mesh #%{pointNumber}<extra></extra>"),
])
fig.update_layout(**dark_layout(
    "Step 3 — Process: Mesh ↔ Mesh (m2m edges)",
    z_labels=([0, 0.5], ["Grid", "Mesh"])
))
fig.add_annotation(
    text="🟢 Mesh nodes exchange information with their spatial neighbours",
    xref="paper", yref="paper", x=0.5, y=1.06,
    xanchor="center", showarrow=False,
    font=dict(color="#94a3b8", size=12)
)
fig.show(renderer="colab")
print(f"   {SYN['m2m'].shape[1]} m2m edges — spatial information propagates across the domain")

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

pio.renderers.default = "colab"

xl_g2m, yl_g2m, zl_g2m = lines(SYN["g2m"], g_xyz, m_xyz)
xl_m2m, yl_m2m, zl_m2m = lines(SYN["m2m"], m_xyz, m_xyz)
xl_m2g, yl_m2g, zl_m2g = lines(SYN["m2g"], m_xyz, g_xyz)

fig = go.Figure([
    go.Scatter3d(x=xl_g2m, y=yl_g2m, z=zl_g2m, mode="lines",
                 name="① Encode  (g2m)",
                 line=dict(color="#a78bfa", width=1.5),
                 opacity=0.45, hoverinfo="skip"),
    go.Scatter3d(x=xl_m2m, y=yl_m2m, z=zl_m2m, mode="lines",
                 name="② Process (m2m)",
                 line=dict(color="#10b981", width=3),
                 opacity=0.7, hoverinfo="skip"),
    go.Scatter3d(x=xl_m2g, y=yl_m2g, z=zl_m2g, mode="lines",
                 name="③ Decode  (m2g)",
                 line=dict(color="#fbbf24", width=1.5),
                 opacity=0.45, hoverinfo="skip"),
    go.Scatter3d(x=g_xyz[:,0], y=g_xyz[:,1], z=g_xyz[:,2],
                 mode="markers", name="Grid nodes",
                 marker=dict(size=5, color="#3b82f6", opacity=0.8),
                 hoverinfo="skip"),
    go.Scatter3d(x=m_xyz[:,0], y=m_xyz[:,1], z=m_xyz[:,2],
                 mode="markers+text", name="Mesh nodes",
                 text=[f"M{i}" for i in range(len(SYN["mesh_xy"]))],
                 textposition="top center",
                 textfont=dict(color="#ffffff", size=9),
                 marker=dict(size=14, color="#f97316", opacity=1.0,
                             line=dict(color="#ffffff", width=1)),
                 hoverinfo="skip"),
])
fig.update_layout(**dark_layout(
    "Complete Pipeline — Encode → Process → Decode",
    z_labels=([0, 0.5], ["Grid", "Mesh"])
))
fig.add_annotation(
    text="🟣 Encode     🟢 Process     🟡 Decode",
    xref="paper", yref="paper", x=0.5, y=1.06,
    xanchor="center", showarrow=False,
    font=dict(color="#94a3b8", size=13)
)
fig.show(renderer="colab")
for name, tensor, desc in [
    ("g2m", SYN["g2m"], "grid → mesh (encode)"),
    ("m2m", SYN["m2m"], "mesh ↔ mesh (process)"),
    ("m2g", SYN["m2g"], "mesh → grid (decode)"),
]:
    print(f"   {name}: {tensor.shape[1]:>4} edges — {desc}")

Part 2 — Real Data (meps_example)

The synthetic graph above had 9 mesh nodes and 72 grid nodes. The real neural-lam graph for MEPS has:

• 63,784 grid nodes

• 6,561 mesh nodes

• 100,656 g2m edges

Showing all nodes at once creates a solid wall. Instead, the next cell picks a 10% spatial patch of the real domain and shows only the nodes and edges inside it — real data, readable scale.

from pathlib import Path
import zipfile
import gdown

DATA_ROOT = Path("meps_example")
ZIP_PATH  = DATA_ROOT / "example_data.zip"
STATIC_DIR = DATA_ROOT / "data/meps_example/static"

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 archive...")
    with zipfile.ZipFile(ZIP_PATH, "r") as zf:
        zf.extractall(str(DATA_ROOT))
else:
    print("✅ meps_example already present")

import torch, numpy as np, plotly.graph_objects as go, plotly.io as pio
from pathlib import Path

pio.renderers.default = "colab"

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

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

raw = np.load(STATIC_DIR / "nwp_xy.npy")
print(f"nwp_xy raw shape: {raw.shape}")
if raw.ndim == 3:
    raw = np.moveaxis(raw, 0, -1) if raw.shape[0] == 2 else raw
    raw = raw.reshape(-1, raw.shape[-1])
if raw.ndim == 2 and raw.shape[0] == 2 and raw.shape[1] != 2:
    raw = raw.T
raw     = raw.astype(np.float32)
grid_xy = (raw - raw.min(0)) / (raw.max(0) - raw.min(0) + 1e-8)

g2m = unwrap(load_pt(GRAPH_DIR / "g2m_edge_index.pt")).long()
m2m = unwrap(load_pt(GRAPH_DIR / "m2m_edge_index.pt")).long()
m2g = unwrap(load_pt(GRAPH_DIR / "m2g_edge_index.pt")).long()

print(f"grid_xy size      : {len(grid_xy)}")
print(f"g2m[0] max index  : {g2m[0].max().item()}")
print(f"g2m[1] max index  : {g2m[1].max().item()}")

# Reconstruct mesh XY and drop edges that reference grid nodes outside the array
n_mesh  = int(g2m[1].max().item()) + 1
mesh_xy = np.zeros((n_mesh, 2), dtype=np.float32)
counts  = np.zeros(n_mesh,      dtype=np.float32)

g_arr   = g2m[0].numpy()
m_arr   = g2m[1].numpy()
valid   = g_arr < len(grid_xy)  # ignore edges beyond available grid nodes
print(f"Valid g2m edges   : {valid.sum():,} of {len(valid):,}")

for g, m in zip(g_arr[valid], m_arr[valid]):
    mesh_xy[m] += grid_xy[g]
    counts[m]  += 1
mesh_xy /= np.maximum(counts, 1)[:, None]

print()
print(f"✅ Real graph loaded")
print(f"   Grid : {len(grid_xy):,} nodes")
print(f"   Mesh : {n_mesh:,} nodes")
print(f"   g2m  : {g2m.shape[1]:,} edges  ({(~valid).sum()} skipped)")
print(f"   Mesh x: {mesh_xy[:,0].min():.3f} → {mesh_xy[:,0].max():.3f}")
print(f"   Mesh y: {mesh_xy[:,1].min():.3f} → {mesh_xy[:,1].max():.3f}")

# Focus on a central 30% spatial patch
CX, CY, R = 0.5, 0.5, 0.15
g_mask = (np.abs(grid_xy[:,0]-CX) < R) & (np.abs(grid_xy[:,1]-CY) < R)
g_idx  = np.where(g_mask)[0]

# Only use valid g2m edges within patch
patch_mask = valid & np.isin(g_arr, g_idx)
patch_g2m  = g2m[:, patch_mask]
m_idx_set  = np.unique(patch_g2m[1].numpy())

mm_mask   = np.isin(m2m[0].numpy(), m_idx_set) & np.isin(m2m[1].numpy(), m_idx_set)
patch_m2m = m2m[:, mm_mask]

m2g_g = m2g[1].numpy()
mg_mask   = np.isin(m2g[0].numpy(), m_idx_set) &             np.isin(m2g_g, g_idx) & (m2g_g < len(grid_xy))
patch_m2g = m2g[:, mg_mask]

g_xyz = np.column_stack([grid_xy[g_idx],    np.zeros(len(g_idx))])
m_xyz = np.column_stack([mesh_xy[m_idx_set], np.full(len(m_idx_set), 0.5)])

g_local = {v: i for i, v in enumerate(g_idx)}
m_local = {v: i for i, v in enumerate(m_idx_set)}

def patch_lines(edges, src_xyz, dst_xyz, sm, dm):
    xl, yl, zl = [], [], []
    for s, d in zip(edges[0].numpy(), edges[1].numpy()):
        if s in sm and d in dm:
            xl += [src_xyz[sm[s],0], dst_xyz[dm[d],0], None]
            yl += [src_xyz[sm[s],1], dst_xyz[dm[d],1], None]
            zl += [src_xyz[sm[s],2], dst_xyz[dm[d],2], None]
    return xl, yl, zl

xl_g2m, yl_g2m, zl_g2m = patch_lines(patch_g2m, g_xyz, m_xyz, g_local, m_local)
xl_m2m, yl_m2m, zl_m2m = patch_lines(patch_m2m, m_xyz, m_xyz, m_local, m_local)
xl_m2g, yl_m2g, zl_m2g = patch_lines(patch_m2g, m_xyz, g_xyz, m_local, g_local)

# Assemble the Plotly figure for this patch
fig = go.Figure([
    go.Scatter3d(x=xl_g2m, y=yl_g2m, z=zl_g2m, mode="lines",
                 name="g2m (encode)", line=dict(color="#a78bfa", width=1),
                 opacity=0.4, hoverinfo="skip"),
    go.Scatter3d(x=xl_m2m, y=yl_m2m, z=zl_m2m, mode="lines",
                 name="m2m (process)", line=dict(color="#10b981", width=2),
                 opacity=0.7, hoverinfo="skip"),
    go.Scatter3d(x=xl_m2g, y=yl_m2g, z=zl_m2g, mode="lines",
                 name="m2g (decode)", line=dict(color="#fbbf24", width=1),
                 opacity=0.4, hoverinfo="skip"),
    go.Scatter3d(x=g_xyz[:,0], y=g_xyz[:,1], z=g_xyz[:,2],
                 mode="markers",
                 name=f"Grid patch ({len(g_idx):,} nodes)",
                 marker=dict(size=3, color="#3b82f6", opacity=0.7),
                 hovertemplate="Grid #%{pointNumber}<extra></extra>"),
    go.Scatter3d(x=m_xyz[:,0], y=m_xyz[:,1], z=m_xyz[:,2],
                 mode="markers",
                 name=f"Mesh patch ({len(m_idx_set):,} nodes)",
                 marker=dict(size=8, color="#f97316", opacity=1.0,
                             line=dict(color="#ffffff", width=1)),
                 hovertemplate="Mesh #%{pointNumber}<extra></extra>"),
])
fig.update_layout(
    title=dict(
        text=f"Real meps_example Graph — Centre Patch\n<sup>Full graph: {len(grid_xy):,} grid · {n_mesh:,} mesh · {g2m.shape[1]:,} g2m edges</sup>",
        font=dict(size=16, color="#ffffff")
    ),
    height=650, paper_bgcolor="#0f172a", font=dict(color="#ffffff"),
    scene=dict(
        bgcolor="#0f172a",
        xaxis=dict(title="X", backgroundcolor="#1e293b",
                   gridcolor="#334155", zeroline=False, color="#94a3b8"),
        yaxis=dict(title="Y", backgroundcolor="#1e293b",
                   gridcolor="#334155", zeroline=False, color="#94a3b8"),
        zaxis=dict(title="Layer", tickvals=[0, 0.5],
                   ticktext=["Grid", "Mesh"], range=[-0.05, 0.8],
                   backgroundcolor="#1e293b",
                   gridcolor="#334155", color="#94a3b8"),
        camera=dict(eye=dict(x=1.6, y=-1.2, z=1.2)),
        aspectmode="manual", aspectratio=dict(x=1, y=1, z=0.55)
    ),
    legend=dict(bgcolor="rgba(15,23,42,0.9)", bordercolor="#334155",
                borderwidth=1, orientation="h",
                x=0.5, xanchor="center", y=-0.05),
    margin=dict(l=0, r=0, t=80, b=0)
)
fig.add_annotation(
    text="🟣 Encode &nbsp;&nbsp;&nbsp; 🟢 Process &nbsp;&nbsp;&nbsp; 🟡 Decode",
    xref="paper", yref="paper", x=0.5, y=1.08,
    xanchor="center", showarrow=False,
    font=dict(color="#94a3b8", size=12)
)
fig.show(renderer="colab")

print()
print(f"   Patch stats:")
print(f"   🔵 Grid : {len(g_idx):,}  of {len(grid_xy):,}")
print(f"   🟠 Mesh : {len(m_idx_set):,}  of {n_mesh:,}")
print(f"   🟣 g2m  : {patch_g2m.shape[1]:,}")
print(f"   🟢 m2m  : {patch_m2m.shape[1]:,}")
print(f"   🟡 m2g  : {patch_m2g.shape[1]:,}")

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

pio.renderers.default = "colab"
np.random.seed(42)

# 3 levels: 9 → 4 → 1 node
L = [
    np.array([[0.2,0.2],[0.5,0.2],[0.8,0.2],
              [0.2,0.5],[0.5,0.5],[0.8,0.5],
              [0.2,0.8],[0.5,0.8],[0.8,0.8]], dtype=np.float32),
    np.array([[0.3,0.3],[0.7,0.3],
              [0.3,0.7],[0.7,0.7]], dtype=np.float32),
    np.array([[0.5, 0.5]], dtype=np.float32),
]
z_heights = [0.2, 0.55, 0.9]
colors    = ["#3b82f6", "#f97316", "#10b981"]
sizes     = [10, 16, 22]
level_labels = ["L0 — Fine (9 nodes)", "L1 — Mid (4 nodes)", "L2 — Coarse (1 node)"]

def nearest_edges(src, dst):
    return np.array(
        [(si, int(np.argmin(np.linalg.norm(dst - sp, axis=1))))
         for si, sp in enumerate(src)],
        dtype=np.int64
    ).T

up_edges = [nearest_edges(L[0], L[1]),
            nearest_edges(L[1], L[2])]

xyz = [np.column_stack([lv, np.full(len(lv), z)])
       for lv, z in zip(L, z_heights)]

traces = []

# Up-edge lines
for i, (edges, col) in enumerate(zip(up_edges, ["#f472b6", "#fb923c"])):
    xl, yl, zl = [], [], []
    for s, d in zip(edges[0], edges[1]):
        xl += [xyz[i][s,0],   xyz[i+1][d,0],   None]
        yl += [xyz[i][s,1],   xyz[i+1][d,1],   None]
        zl += [xyz[i][s,2],   xyz[i+1][d,2],   None]
    traces.append(go.Scatter3d(
        x=xl, y=yl, z=zl, mode="lines",
        name=f"L{i}→L{i+1} up-edges",
        line=dict(color=col, width=2),
        opacity=0.7, hoverinfo="skip"
    ))

# Node dots + labels
for i, (xyzl, lv) in enumerate(zip(xyz, L)):
    traces.append(go.Scatter3d(
        x=xyzl[:,0], y=xyzl[:,1], z=xyzl[:,2],
        mode="markers+text",
        name=level_labels[i],
        text=[f"L{i}-{j}" for j in range(len(lv))],
        textposition="top center",
        textfont=dict(color="#ffffff", size=9),
        marker=dict(size=sizes[i], color=colors[i], opacity=1.0,
                    line=dict(color="#ffffff", width=1)),
        hovertemplate=f"Level {i} Node %{{pointNumber}}<extra></extra>"
    ))

fig = go.Figure(traces)
fig.update_layout(
    title=dict(
        text="HiLAM — Hierarchical Mesh: Fine → Mid → Coarse",
        font=dict(size=17, color="#ffffff")
    ),
    height=650, paper_bgcolor="#0f172a", font=dict(color="#ffffff"),
    scene=dict(
        bgcolor="#0f172a",
        xaxis=dict(title="X", backgroundcolor="#1e293b",
                   gridcolor="#334155", zeroline=False, color="#94a3b8",
                   range=[-0.05, 1.05]),
        yaxis=dict(title="Y", backgroundcolor="#1e293b",
                   gridcolor="#334155", zeroline=False, color="#94a3b8",
                   range=[-0.05, 1.05]),
        zaxis=dict(title="Level",
                   tickvals=z_heights,
                   ticktext=["L0 Fine", "L1 Mid", "L2 Coarse"],
                   range=[0.0, 1.1],
                   backgroundcolor="#1e293b",
                   gridcolor="#334155", color="#94a3b8"),
        camera=dict(eye=dict(x=1.5, y=-1.5, z=1.4)),
        aspectmode="manual", aspectratio=dict(x=1, y=1, z=0.8)
    ),
    legend=dict(bgcolor="rgba(15,23,42,0.9)", bordercolor="#334155",
                borderwidth=1, orientation="h",
                x=0.5, xanchor="center", y=-0.05),
    margin=dict(l=0, r=0, t=60, b=0)
)
fig.add_annotation(
    text="🔵 L0 fine &nbsp;&nbsp; 🟠 L1 mid &nbsp;&nbsp; 🟢 L2 coarse &nbsp;&nbsp;🩷 up-edges (aggregate local→global) &nbsp;&nbsp;🟠 down-edges (broadcast global→local)",
    xref="paper", yref="paper", x=0.5, y=1.06,
    xanchor="center", showarrow=False,
    font=dict(color="#94a3b8", size=11)
)
fig.show(renderer="colab")

print("   L0 — 9 nodes  : fine spatial detail (one per mesh cell)")
print("   L1 — 4 nodes  : regional patterns")
print("   L2 — 1 node   : global context (sees entire domain)")
print()
print("   Up sweep   L0→L1→L2 : aggregate local → global")
print("   Down sweep L2→L1→L0 : broadcast global → local predictions")
print()
print("   GraphLAM uses only L0.")
print("   HiLAM uses all 3 levels — better for long-range dependencies.")

What You’ve Seen

Visualisation	What it shows
Nodes only	Two layers — grid observations below, latent mesh above
g2m edges	How observations flow into the mesh (encode)
m2m edges	How mesh nodes talk to each other (process)
Full pipeline	All three steps together
Real data patch	Actual meps_example graph — 63,784 grid, 6,561 mesh
HiLAM hierarchy	How multiple mesh levels enable global context

Next Steps

I want to…	Go here
Run a full training example	hello_world.ipynb
Understand the data	data_exploration.ipynb
Understand the model code	model_architecture.ipynb
Build my own graph	python -m neural_lam.create_graph --help

import numpy as np, torch
from pathlib import Path

ROOT   = Path("meps_example/data/meps_example")
STATIC = ROOT / "static"
TRAIN  = ROOT / "samples/train"

nwp = np.load(TRAIN / "nwp_2022040100_mbr000.npy")
print(f"nwp shape  : {nwp.shape}   dtype={nwp.dtype}")
print(f"nwp min    : {nwp.min():.3f}  max={nwp.max():.3f}")

for p in sorted(TRAIN.glob("*")):
    a = np.load(p)
    print(f"  {p.name:<55} {a.shape}")

print()
for p in sorted(STATIC.rglob("*")):
    if not p.is_file(): continue
    if p.suffix == ".npy":
        a = np.load(p)
        print(f"static  {p.name:<35} {a.shape}")
    elif p.suffix == ".pt":
        try:    t = torch.load(p, map_location="cpu", weights_only=True)
        except: t = torch.load(p, map_location="cpu")
        if isinstance(t, list): t = t[0]
        print(f"static  {p.name:<35} {tuple(t.shape)}")