OneScience-Group/GraphDOP

1# coding=utf-82#3# SPDX-License-Identifier: Apache-2.04#5# Minimal reproduction of GraphDOP (ECMWF, "Towards skilful medium-range6# forecasts learnt directly from observations", 2025 preprint) following the7# encoder -- processor -- decoder design:8#9#   * Encoder: a GNN that projects gridded "observations" inside the input10#     window onto a latent mesh (a coarse regular lat/lon grid), using graph11#     edges with (forward bearing, haversine distance) features.12#   * Processor: a transformer that advances the latent atmospheric state13#     forward in time, once per output frame (latent-space rollout).14#   * Decoder: a GNN that maps the latent mesh back onto the target grid and15#     predicts per-channel observations with instrument-like output MLPs.16#17# Differences from the paper (documented in README.md): the paper consumes18# irregular, instrument-specific Level-1 observations with dynamic graphs built19# per batch (PyTorch Geometric); here the OneScience ERA5-h5 gridded pipeline is20# used as the observation placeholder, and the graphs are fixed regular-grid21# meshes. The weighted MSE objective is kept (per-channel weights).22import math23 24import torch25import torch.nn as nn26import torch.nn.functional as F27 28 29def _latlon_grid(shape):30    """Regular lat/lon coordinates for a (H, W) grid, North-to-South rows."""31    H, W = shape32    lat = torch.linspace(90.0, -90.0, H)33    lon = torch.linspace(0.0, 360.0 - 360.0 / W, W)34    return lat, lon35 36 37def _haversine(lat1, lon1, lat2, lon2):38    """Haversine distance in metres between points given in degrees."""39    R = 6371000.040    p1 = torch.deg2rad(lat1)41    p2 = torch.deg2rad(lat2)42    dp = torch.deg2rad(lat2 - lat1)43    dl = torch.deg2rad(lon2 - lon1)44    a = torch.sin(dp / 2) ** 2 + torch.cos(p1) * torch.cos(p2) * torch.sin(dl / 2) ** 245    return 2 * R * torch.asin(torch.sqrt(a.clamp(0, 1)))46 47 48def _bearing(lat1, lon1, lat2, lon2):49    """Initial forward bearing in radians from point 1 to point 2."""50    p1 = torch.deg2rad(lat1)51    p2 = torch.deg2rad(lat2)52    dl = torch.deg2rad(lon2 - lon1)53    y = torch.sin(dl) * torch.cos(p2)54    x = torch.cos(p1) * torch.sin(p2) - torch.sin(p1) * torch.cos(p2) * torch.cos(dl)55    return torch.atan2(y, x)56 57 58def build_mesh_graph(mesh_shape):59    """60    Build a fixed 8-neighbourhood graph over a regular latent mesh.61    Longitude wraps around; edge features are (forward bearing [rad],62    haversine distance [km]).63    """64    H, W = mesh_shape65    lat, lon = _latlon_grid(mesh_shape)66    lat = lat.view(-1, 1).expand(H, W)67    lon = lon.view(1, -1).expand(H, W)68 69    src_list, dst_list, feat_list = [], [], []70    for i in range(H):71        for j in range(W):72            for di, dj in ((-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1)):73                ni, nj = i + di, (j + dj) % W74                if not (0 <= ni < H):75                    continue76                s = i * W + j77                d = ni * W + nj78                dist_km = _haversine(lat[i, j], lon[i, j], lat[ni, nj], lon[ni, nj]) / 1000.079                bear = _bearing(lat[i, j], lon[i, j], lat[ni, nj], lon[ni, nj])80                src_list.append(s)81                dst_list.append(d)82                feat_list.append(torch.stack([bear / math.pi, dist_km / 1000.0]))83    edge_index = torch.stack([torch.as_tensor(src_list), torch.as_tensor(dst_list)], dim=0)84    edge_attr = torch.stack(feat_list)85    return edge_index, edge_attr86 87 88def _mlp(in_dim, out_dim, hidden_dim, n_layers=2):89    dims = [in_dim] + [hidden_dim] * (n_layers - 1) + [out_dim]90    layers = []91    for i in range(len(dims) - 1):92        layers.append(nn.Linear(dims[i], dims[i + 1]))93        if i < len(dims) - 2:94            layers.append(nn.GELU())95    return nn.Sequential(*layers)96 97 98class GNNLayer(nn.Module):99    """Message-passing layer with edge features (mean-aggregate, residual)."""100 101    def __init__(self, dim, edge_dim=2, hidden_dim=64):102        super().__init__()103        self.edge_mlp = _mlp(2 * dim + edge_dim, dim, hidden_dim)104        self.node_mlp = _mlp(dim, dim, hidden_dim)105        self.norm = nn.LayerNorm(dim)106 107    def forward(self, x, edge_index, edge_attr):108        B, N, D = x.shape109        src, dst = edge_index110        offsets = torch.arange(B, device=x.device) * N111        src_b = (src.unsqueeze(0) + offsets.view(B, 1)).reshape(-1)112        dst_b = (dst.unsqueeze(0) + offsets.view(B, 1)).reshape(-1)113        edge_attr_b = edge_attr.unsqueeze(0).expand(B, -1, -1).reshape(-1, edge_attr.size(1))114        xb = x.reshape(B * N, D)115        msg = self.edge_mlp(torch.cat([xb[src_b], xb[dst_b], edge_attr_b], dim=1))116        agg = torch.zeros_like(xb)117        agg.index_add_(0, dst_b, msg)118        cnt = torch.bincount(dst_b, minlength=B * N).clamp(min=1).unsqueeze(1)119        agg = agg / cnt120        agg = agg.reshape(B, N, D)121        return self.norm(x + self.node_mlp(agg))122 123 124class ObsEncoder(nn.Module):125    """126    Maps the observation grid onto the latent mesh with a per-cell input MLP,127    an adaptive pooling to the mesh resolution, and graph message passing.128    """129 130    def __init__(self, in_channels, latent_dim, mesh_shape, num_layers=2, hidden_dim=64):131        super().__init__()132        self.in_channels = in_channels133        self.input_mlp = _mlp(in_channels, latent_dim, hidden_dim)134        self.gnn = nn.ModuleList([GNNLayer(latent_dim, hidden_dim=hidden_dim) for _ in range(num_layers)])135        self.mesh_shape = mesh_shape136        self.edge_index, self.edge_attr = build_mesh_graph(mesh_shape)137 138    def forward(self, x):139        B, C, H, W = x.shape140        feat = x.permute(0, 2, 3, 1).reshape(-1, C)141        feat = self.input_mlp(feat).reshape(B, H, W, -1).permute(0, 3, 1, 2)142        mesh = F.adaptive_avg_pool2d(feat, self.mesh_shape)143        mesh = mesh.permute(0, 2, 3, 1).reshape(B, -1, mesh.size(1))144        edge_index, edge_attr = self.edge_index.to(x.device), self.edge_attr.to(x.device)145        for layer in self.gnn:146            mesh = layer(mesh, edge_index, edge_attr)147        return mesh148 149 150class LatentProcessor(nn.Module):151    """Transformer over latent mesh tokens that advances the state in time."""152 153    def __init__(self, latent_dim, mesh_shape, num_blocks=1, n_heads=4, hidden_dim=128):154        super().__init__()155        n_nodes = mesh_shape[0] * mesh_shape[1]156        self.pos_emb = nn.Parameter(torch.zeros(1, n_nodes, latent_dim))157        nn.init.trunc_normal_(self.pos_emb, std=0.02)158        block = nn.TransformerEncoderLayer(159            d_model=latent_dim, nhead=n_heads, dim_feedforward=hidden_dim,160            dropout=0.0, activation="gelu", batch_first=True, norm_first=True,161        )162        self.blocks = nn.ModuleList([block for _ in range(num_blocks)])163 164    def forward(self, mesh):165        tokens = mesh + self.pos_emb166        for block in self.blocks:167            tokens = block(tokens)168        return tokens169 170 171class ObsDecoder(nn.Module):172    """173    Maps the latent mesh back onto the target grid (bilinear upsample) and174    predicts per-channel observations with an output MLP.175    """176 177    def __init__(self, latent_dim, out_channels, grid_shape, mesh_shape, num_layers=2, hidden_dim=64):178        super().__init__()179        self.gnn = nn.ModuleList([GNNLayer(latent_dim, hidden_dim=hidden_dim) for _ in range(num_layers)])180        self.grid_shape = grid_shape181        self.mesh_shape = mesh_shape182        self.edge_index, self.edge_attr = build_mesh_graph(mesh_shape)183        self.output_mlp = _mlp(latent_dim, out_channels, hidden_dim)184 185    def forward(self, mesh):186        B, N, D = mesh.shape187        edge_index = self.edge_index.to(mesh.device)188        edge_attr = self.edge_attr.to(mesh.device)189        for layer in self.gnn:190            mesh = layer(mesh, edge_index, edge_attr)191        H, W = self.grid_shape192        Hm, Wm = self.mesh_shape193        mesh = mesh.transpose(1, 2).reshape(B, D, Hm, Wm)194        grid = F.interpolate(mesh, size=self.grid_shape, mode="bilinear", align_corners=False)195        grid = grid.permute(0, 2, 3, 1).reshape(B, H * W, D)196        return self.output_mlp(grid).reshape(B, H, W, -1).permute(0, 3, 1, 2)197 198 199class GraphDOP(nn.Module):200    """201    Config-driven GraphDOP wrapper.202 203    Args:204        in_channels: Number of observation channels per frame.205        out_channels: Number of forecast channels per frame.206        input_steps: Number of input (observation window) frames.207        output_steps: Number of forecast frames.208        grid_shape: Spatial shape of the (gridded) observation field.209        mesh_shape: Latent mesh resolution (each dimension, powers of two fine).210        latent_dim: Feature dimension of latent mesh tokens.211        num_encoder_layers / num_decoder_layers: GNN message-passing layers.212        num_processor_blocks: Transformer blocks in the processor.213        n_heads: Attention heads of the processor.214        channel_weights: Per-channel weights for the weighted MSE objective.215    """216 217    def __init__(218        self,219        in_channels=6,220        out_channels=6,221        input_steps=2,222        output_steps=2,223        grid_shape=(32, 32),224        mesh_shape=(8, 8),225        latent_dim=64,226        num_encoder_layers=2,227        num_decoder_layers=2,228        num_processor_blocks=1,229        n_heads=4,230        hidden_dim=64,231        channel_weights=None,232    ):233        super().__init__()234        self.in_channels = int(in_channels)235        self.out_channels = int(out_channels)236        self.input_steps = int(input_steps)237        self.output_steps = int(output_steps)238        self.grid_shape = (int(grid_shape[0]), int(grid_shape[1]))239        self.mesh_shape = (int(mesh_shape[0]), int(mesh_shape[1]))240 241        self.encoder = ObsEncoder(242            self.in_channels, int(latent_dim), self.mesh_shape, num_layers=int(num_encoder_layers), hidden_dim=int(hidden_dim)243        )244        self.processor = LatentProcessor(245            int(latent_dim), self.mesh_shape, num_blocks=int(num_processor_blocks), n_heads=int(n_heads), hidden_dim=int(hidden_dim)246        )247        self.decoder = ObsDecoder(248            int(latent_dim), self.out_channels, self.grid_shape, self.mesh_shape,249            num_layers=int(num_decoder_layers), hidden_dim=int(hidden_dim),250        )251 252        if channel_weights is None:253            channel_weights = torch.ones(self.out_channels)254        self.register_buffer("channel_weights", torch.as_tensor(channel_weights, dtype=torch.float32))255 256    def forward(self, x):257        """258        Args:259            x: Observation frames, shape [batch, input_steps, C, H, W].260        Returns:261            Forecast frames, shape [batch, output_steps, C, H, W].262        """263        latents = torch.stack([self.encoder(x[:, t]) for t in range(self.input_steps)], dim=0)264        latent = latents.mean(dim=0)265        outs = []266        for _ in range(self.output_steps):267            latent = self.processor(latent)268            outs.append(self.decoder(latent))269        return torch.stack(outs, dim=1)270 271    def wmse_loss(self, pred, target):272        """Weighted mean squared error objective (Eq. 1 of the paper)."""273        diff = (pred - target) ** 2274        w = self.channel_weights.view(1, 1, self.out_channels, 1, 1)275        return (diff * w).mean()276