Publish ACE2 reproduction

.gitattributes

@@ -1,35 +1,5 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.json text eol=lf
+*.md text eol=lf
+*.py text eol=lf
+*.yaml text eol=lf
+weight/** filter=lfs diff=lfs merge=lfs -text

conf/config.yaml

@@ -0,0 +1,33 @@
+seed: 17
+data:
+  format_version: ace2-full-grid-v1
+  path: data/ace2_tiny.npz
+  samples: 4
+  time_steps: 6
+  channels: 50
+  height: 180
+  width: 360
+  levels: 8
+  dt_hours: 6
+model:
+  width: 4
+  depth: 8
+  modes_lat: 4
+  modes_lon: 4
+  forcing_channels: 4
+paper_model:
+  embedding_dim: 384
+  parameters: 450000000
+  autoregressive_training_steps: 2
+train:
+  epochs: 1
+  batch_size: 1
+  learning_rate: 0.001
+  two_step_weight: 0.5
+  checkpoint: result/checkpoints/ace2.pt
+inference:
+  steps: 3
+  output: result/output/predictions.npz
+evaluation:
+  output: result/evaluation/metrics.json
+  figure: result/evaluation/comparison.png

config.json

@@ -0,0 +1,12 @@
+{
+  "model_name": "ACE2",
+  "model_type": "ace2",
+  "architectures": ["CompactSFNO"],
+  "framework": "PyTorch",
+  "domain": "climate",
+  "task": "global-atmospheric-emulation",
+  "implementation": {"entry_point": "model/ace2.py", "scope": "core-method full-grid reduced-model engineering reproduction", "train_script": "scripts/train.py", "inference_script": "scripts/inference.py", "evaluation_script": "scripts/result.py", "synthetic_data_script": "scripts/fake_data.py"},
+  "architecture": {"grid": [180,360], "state_channels": 50, "forcing_channels": 4, "vertical_layers": 8, "engineering_width": 4, "engineering_depth": 8, "paper_embedding": 384, "paper_parameters": 450000000},
+  "data": {"datasets": ["ERA5","SHiELD"], "format_version": "ace2-full-grid-v1", "time_step_hours": 6, "layout": "BCHW", "synthetic": true},
+  "configuration_sources": ["conf/config.yaml","model/ace2.py","scripts/fake_data.py","scripts/train.py","scripts/inference.py","scripts/result.py"]
+}

model/ace2.py

@@ -0,0 +1,177 @@
+import os
+import random
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+from torch import nn
+
+
+LEVELS_HPA = (1000, 850, 700, 500, 300, 200, 100, 50)
+CHANNELS = (
+    [f"temperature_{p}" for p in LEVELS_HPA]
+    + [f"specific_humidity_{p}" for p in LEVELS_HPA]
+    + [f"u_wind_{p}" for p in LEVELS_HPA]
+    + [f"v_wind_{p}" for p in LEVELS_HPA]
+    + [f"geopotential_{p}" for p in LEVELS_HPA]
+    + [
+        "surface_pressure",
+        "air_temperature_2m",
+        "specific_humidity_2m",
+        "eastward_wind_10m",
+        "northward_wind_10m",
+        "sea_surface_temperature",
+        "total_precipitation_6h",
+        "surface_downward_shortwave",
+        "surface_downward_longwave",
+        "toa_outgoing_longwave",
+    ]
+)
+assert len(CHANNELS) == 50
+Q_INDICES = tuple(range(8, 16)) + (42,)
+SURFACE_PRESSURE = 40
+PRECIPITATION = 46
+RADIATION_INDICES = (47, 48, 49)
+
+
+class SpectralConv2d(nn.Module):
+    def __init__(self, width, modes_lat, modes_lon):
+        super().__init__()
+        self.modes_lat, self.modes_lon = modes_lat, modes_lon
+        scale = 1.0 / width
+        self.weight = nn.Parameter(
+            scale * torch.randn(width, width, modes_lat, modes_lon, dtype=torch.cfloat)
+        )
+
+    def forward(self, x):
+        spectrum = torch.fft.rfft2(x, norm="ortho")
+        out = torch.zeros_like(spectrum)
+        ml = min(self.modes_lat, spectrum.shape[-2])
+        mn = min(self.modes_lon, spectrum.shape[-1])
+        out[:, :, :ml, :mn] = torch.einsum(
+            "bixy,ioxy->boxy", spectrum[:, :, :ml, :mn], self.weight[:, :, :ml, :mn]
+        )
+        return torch.fft.irfft2(out, s=x.shape[-2:], norm="ortho")
+
+
+class SFNOBlock(nn.Module):
+    def __init__(self, width, modes_lat, modes_lon):
+        super().__init__()
+        self.spectral = SpectralConv2d(width, modes_lat, modes_lon)
+        self.mlp = nn.Sequential(
+            nn.Conv2d(width, width * 2, 1), nn.GELU(), nn.Conv2d(width * 2, width, 1)
+        )
+        self.norm = nn.GroupNorm(1, width)
+
+    def forward(self, x):
+        return x + self.mlp(self.norm(self.spectral(x)))
+
+
+class CompactSFNO(nn.Module):
+    def __init__(self, channels=50, forcing_channels=4, width=4, depth=1,
+                 modes_lat=4, modes_lon=4):
+        super().__init__()
+        self.lift = nn.Conv2d(channels + forcing_channels, width, 1)
+        self.blocks = nn.Sequential(
+            *[SFNOBlock(width, modes_lat, modes_lon) for _ in range(depth)]
+        )
+        self.project = nn.Sequential(nn.GELU(), nn.Conv2d(width, channels, 1))
+
+    def forward(self, state, forcing):
+        features = self.blocks(self.lift(torch.cat((state, forcing), dim=1)))
+        return state + self.project(features)
+
+
+def area_weights(height, device, dtype):
+    lat = torch.linspace(-89.5, 89.5, height, device=device, dtype=dtype)
+    return torch.cos(torch.deg2rad(lat)).view(1, 1, height, 1)
+
+
+def weighted_mean(x, weights):
+    return (x * weights).sum(dim=(-2, -1), keepdim=True) / (
+        weights.sum(dim=(-2, -1), keepdim=True) * x.shape[-1]
+    )
+
+
+def hard_correct(previous, predicted):
+    """Apply differentiable positivity, dry-mass, and global-water constraints."""
+    out = predicted.clone()
+    positive = list(Q_INDICES) + [PRECIPITATION] + list(RADIATION_INDICES)
+    out[:, positive] = torch.clamp_min(out[:, positive], 0.0)
+    weights = area_weights(out.shape[-2], out.device, out.dtype)
+
+    q_prev = previous[:, Q_INDICES].sum(dim=1, keepdim=True)
+    water_target = weighted_mean(q_prev, weights)
+    precip = weighted_mean(out[:, PRECIPITATION:PRECIPITATION + 1], weights)
+    precip_scale = torch.clamp(
+        0.5 * water_target / torch.clamp_min(precip, 1e-8), max=1.0
+    )
+    out[:, PRECIPITATION:PRECIPITATION + 1] *= precip_scale
+    precip = weighted_mean(out[:, PRECIPITATION:PRECIPITATION + 1], weights)
+    q_target = torch.clamp_min(water_target - precip, 0.0)
+    q_now = weighted_mean(out[:, Q_INDICES].sum(dim=1, keepdim=True), weights)
+    out[:, Q_INDICES] *= q_target / torch.clamp_min(q_now, 1e-8)
+
+    q_new = out[:, Q_INDICES].sum(dim=1, keepdim=True)
+    dry_target = weighted_mean(
+        previous[:, SURFACE_PRESSURE:SURFACE_PRESSURE + 1] - q_prev, weights
+    )
+    dry_now = weighted_mean(
+        out[:, SURFACE_PRESSURE:SURFACE_PRESSURE + 1] - q_new, weights
+    )
+    out[:, SURFACE_PRESSURE:SURFACE_PRESSURE + 1] += dry_target - dry_now
+    return out
+
+
+def load_config(root=None):
+    root = Path(root) if root is not None else Path(__file__).resolve().parents[1]
+    with (root / "conf" / "config.yaml").open(encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def seed_all(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def forcing_for_hours(hours, height=180, width=360):
+    hours = np.asarray(hours, dtype=np.float32)
+    phase = 2 * np.pi * hours / (365.25 * 24)
+    lat = np.deg2rad(np.linspace(-89.5, 89.5, height, dtype=np.float32))
+    lon = np.deg2rad(np.linspace(0.5, 359.5, width, dtype=np.float32))
+    solar = np.maximum(
+        0,
+        np.cos(lat)[None, :, None]
+        * np.cos(lon[None, None, :] + phase[:, None, None]),
+    )
+    fields = np.empty((len(hours), 4, height, width), dtype=np.float32)
+    fields[:, 0] = np.sin(phase)[:, None, None]
+    fields[:, 1] = np.cos(phase)[:, None, None]
+    fields[:, 2] = (400.0 + 0.01 * hours)[:, None, None] / 500.0
+    fields[:, 3] = solar
+    return fields
+
+
+def init_distributed():
+    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    use_cuda = torch.cuda.is_available() and torch.cuda.device_count() >= world_size
+    if distributed:
+        backend = "nccl" if use_cuda else "gloo"
+        torch.distributed.init_process_group(backend=backend)
+        rank = torch.distributed.get_rank()
+        local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    else:
+        rank = local_rank = 0
+    device = torch.device(
+        f"cuda:{local_rank}" if use_cuda else "cpu"
+    )
+    if device.type == "cuda":
+        torch.cuda.set_device(device)
+    return distributed, rank, device
+
+
+def build_model(config):
+    return CompactSFNO(channels=config["data"]["channels"], **config["model"])

scripts/fake_data.py

@@ -0,0 +1,50 @@
+from pathlib import Path
+import sys
+
+import numpy as np
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from model.ace2 import (
+    PRECIPITATION,
+    Q_INDICES,
+    RADIATION_INDICES,
+    SURFACE_PRESSURE,
+    forcing_for_hours,
+    load_config,
+    seed_all,
+)
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def main():
+    cfg = load_config(ROOT)
+    seed_all(cfg["seed"])
+    d = cfg["data"]
+    shape = (d["samples"], d["time_steps"], d["channels"], d["height"], d["width"])
+    rng = np.random.default_rng(cfg["seed"])
+    lat = np.deg2rad(np.linspace(-89.5, 89.5, d["height"], dtype=np.float32))[None, None, None, :, None]
+    lon = np.deg2rad(np.linspace(0.5, 359.5, d["width"], dtype=np.float32))[None, None, None, None, :]
+    time = np.arange(d["time_steps"], dtype=np.float32)[None, :, None, None, None]
+    channel = np.arange(d["channels"], dtype=np.float32)[None, None, :, None, None]
+    state = (0.2 * np.sin(lat * (1 + channel % 3)) + 0.1 * np.cos(lon + time / 4)
+             + 0.002 * channel + 0.003 * time).astype(np.float32)
+    state = np.broadcast_to(state, shape).copy()
+    state += rng.normal(0, 0.005, (d["samples"], 1, d["channels"], 1, 1)).astype(np.float32)
+    state[:, :, Q_INDICES] = np.maximum(state[:, :, Q_INDICES] * 0.01 + 0.003, 0)
+    state[:, :, SURFACE_PRESSURE] = 1.0 + 0.01 * np.cos(lat[:, :, 0])
+    state[:, :, PRECIPITATION] = np.maximum(0, 0.001 * (np.sin(lon[:, :, 0] + time[:, :, 0]) + 1))
+    state[:, :, RADIATION_INDICES] = np.maximum(state[:, :, RADIATION_INDICES] + 0.5, 0)
+    hours = (np.arange(d["samples"])[:, None] * d["time_steps"] + np.arange(d["time_steps"])[None]) * d["dt_hours"]
+    forcing = forcing_for_hours(hours.reshape(-1), d["height"], d["width"]).reshape(
+        d["samples"], d["time_steps"], 4, d["height"], d["width"]
+    )
+    path = ROOT / d["path"]
+    path.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(path, state=state.astype(np.float16), forcing=forcing.astype(np.float16), hours=hours)
+    print(f"created {path}: state={state.shape}, forcing={forcing.shape}, dt={d['dt_hours']}h")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,38 @@
+from pathlib import Path
+import sys
+
+import numpy as np
+import torch
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.ace2 import build_model, hard_correct, load_config
+
+
+def main():
+    cfg = load_config(ROOT)
+    data = np.load(ROOT / cfg["data"]["path"])
+    model = build_model(cfg)
+    checkpoint = torch.load(
+        ROOT / cfg["train"]["checkpoint"], map_location="cpu", weights_only=True
+    )
+    if checkpoint["format_version"] != cfg["data"]["format_version"]:
+        raise ValueError("checkpoint format_version mismatch")
+    model.load_state_dict(checkpoint["model"])
+    model.eval()
+    current = torch.from_numpy(data["state"][0, 0].astype(np.float32)).unsqueeze(0)
+    forecast = []
+    with torch.no_grad():
+        for step in range(cfg["inference"]["steps"]):
+            forcing = torch.from_numpy(data["forcing"][0, step + 1].astype(np.float32)).unsqueeze(0)
+            current = hard_correct(current, model(current, forcing))
+            forecast.append(current.squeeze(0).numpy().astype(np.float16))
+    output = ROOT / cfg["inference"]["output"]
+    output.parent.mkdir(parents=True, exist_ok=True)
+    leads = np.arange(1, cfg["inference"]["steps"] + 1) * cfg["data"]["dt_hours"]
+    np.savez_compressed(output, forecast=np.stack(forecast), lead_hours=leads)
+    print(f"saved {output}: forecast={tuple(np.stack(forecast).shape)}, leads={leads.tolist()}h")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,68 @@
+import json
+from pathlib import Path
+import sys
+
+import numpy as np
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.ace2 import PRECIPITATION, Q_INDICES, SURFACE_PRESSURE, load_config
+
+
+def weighted_mean(x, w):
+    return np.sum(x * w, axis=(-2, -1)) / np.sum(np.broadcast_to(w, x.shape), axis=(-2, -1))
+
+
+def rmse(pred, target, w):
+    return float(np.sqrt(np.sum((pred - target) ** 2 * w) / np.sum(np.broadcast_to(w, pred.shape))))
+
+
+def main():
+    cfg = load_config(ROOT)
+    data = np.load(ROOT / cfg["data"]["path"])
+    steps = cfg["inference"]["steps"]
+    truth = data["state"][0, 1:steps + 1].astype(np.float32)
+    pred = np.load(ROOT / cfg["inference"]["output"])["forecast"].astype(np.float32)
+    initial = data["state"][0, 0].astype(np.float32)
+    persistence = np.broadcast_to(initial, truth.shape)
+    w = np.cos(np.deg2rad(np.linspace(-89.5, 89.5, 180, dtype=np.float32)))[None, None, :, None]
+    pg, tg = weighted_mean(pred, w), weighted_mean(truth, w)
+    denom = np.sum((tg - tg.mean(axis=0, keepdims=True)) ** 2)
+    r2 = float(1 - np.sum((pg - tg) ** 2) / max(float(denom), 1e-12))
+    dry0 = weighted_mean(initial[SURFACE_PRESSURE] - initial[list(Q_INDICES)].sum(0), w[0, 0])
+    dry = weighted_mean(pred[:, SURFACE_PRESSURE] - pred[:, list(Q_INDICES)].sum(1), w[0, 0])
+    previous = np.concatenate((initial[None], pred[:-1]), axis=0)
+    water_previous = weighted_mean(previous[:, list(Q_INDICES)].sum(1), w[0, 0])
+    water = weighted_mean(pred[:, list(Q_INDICES)].sum(1) + pred[:, PRECIPITATION], w[0, 0])
+    model_rmse, baseline_rmse = rmse(pred, truth, w), rmse(persistence, truth, w)
+    metrics = {
+        "area_weighted_rmse": model_rmse,
+        "global_mean_r2": r2,
+        "conservation": {
+            "max_abs_global_dry_mass_error": float(np.max(np.abs(dry - dry0))),
+            "max_abs_global_moisture_closure_error": float(np.max(np.abs(water - water_previous))),
+        },
+        "comparison": {
+            "persistence_area_weighted_rmse": baseline_rmse,
+            "rmse_skill_vs_persistence": float(1 - model_rmse / baseline_rmse),
+        },
+    }
+    output = ROOT / cfg["evaluation"]["output"]
+    output.parent.mkdir(parents=True, exist_ok=True)
+    output.write_text(json.dumps(metrics, indent=2) + "\n", encoding="utf-8")
+    figure = ROOT / cfg["evaluation"]["figure"]
+    figure.parent.mkdir(parents=True, exist_ok=True)
+    fig, axes = plt.subplots(1, 2, figsize=(9, 3.8), constrained_layout=True)
+    axes[0].bar(["ACE2", "Persistence"], [model_rmse, baseline_rmse], color=["#287f71", "#8c96a8"])
+    axes[0].set(ylabel="Area-weighted RMSE", title="Forecast error")
+    axes[1].bar(["Dry mass", "Moisture"], [metrics["conservation"]["max_abs_global_dry_mass_error"], metrics["conservation"]["max_abs_global_moisture_closure_error"]], color="#d69b36")
+    axes[1].set_yscale("log"); axes[1].set(title="Conservation residual", ylabel="Maximum absolute error")
+    fig.savefig(figure, dpi=150); plt.close(fig)
+    print(json.dumps(metrics, indent=2))
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,72 @@
+from pathlib import Path
+import sys
+
+import numpy as np
+import torch
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, Dataset, DistributedSampler
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.ace2 import build_model, hard_correct, init_distributed, load_config, seed_all
+
+
+class Windows(Dataset):
+    def __init__(self, path):
+        data = np.load(path)
+        self.state = data["state"]
+        self.forcing = data["forcing"]
+        self.windows = [(n, t) for n in range(len(self.state)) for t in range(self.state.shape[1] - 2)]
+
+    def __len__(self):
+        return len(self.windows)
+
+    def __getitem__(self, index):
+        n, t = self.windows[index]
+        return tuple(torch.from_numpy(x.astype(np.float32)) for x in (
+            self.state[n, t], self.state[n, t + 1], self.state[n, t + 2],
+            self.forcing[n, t + 1], self.forcing[n, t + 2]))
+
+
+def main():
+    cfg = load_config(ROOT)
+    seed_all(cfg["seed"])
+    distributed, rank, device = init_distributed()
+    dataset = Windows(ROOT / cfg["data"]["path"])
+    sampler = DistributedSampler(dataset, shuffle=True) if distributed else None
+    loader = DataLoader(dataset, batch_size=cfg["train"]["batch_size"], sampler=sampler,
+                        shuffle=sampler is None, num_workers=0)
+    model = build_model(cfg).to(device)
+    if distributed:
+        model = DistributedDataParallel(model, device_ids=[device.index] if device.type == "cuda" else None)
+    optimizer = torch.optim.Adam(model.parameters(), lr=cfg["train"]["learning_rate"])
+    for epoch in range(cfg["train"]["epochs"]):
+        if sampler:
+            sampler.set_epoch(epoch)
+        total = 0.0
+        for x0, y1, y2, f1, f2 in loader:
+            x0, y1, y2, f1, f2 = (x.to(device) for x in (x0, y1, y2, f1, f2))
+            p1 = hard_correct(x0, model(x0, f1))
+            p2 = hard_correct(p1, model(p1, f2))
+            loss = torch.mean((p1 - y1) ** 2) + cfg["train"]["two_step_weight"] * torch.mean((p2 - y2) ** 2)
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            optimizer.step()
+            total += loss.item()
+        if rank == 0:
+            print(f"epoch={epoch + 1} two_step_loss={total / len(loader):.7f}")
+    if rank == 0:
+        path = ROOT / cfg["train"]["checkpoint"]
+        path.parent.mkdir(parents=True, exist_ok=True)
+        raw_model = model.module if distributed else model
+        torch.save({"model": raw_model.state_dict(), "model_config": cfg["model"], "format_version": cfg["data"]["format_version"]}, path)
+        metrics = ROOT / "result/training/metrics.json"
+        metrics.parent.mkdir(parents=True, exist_ok=True)
+        metrics.write_text(__import__("json").dumps({"history": [{"epoch": cfg["train"]["epochs"], "loss": total / len(loader)}], "world_size": int(__import__("os").environ.get("WORLD_SIZE", "1"))}, indent=2))
+        print(f"saved {path}")
+    if distributed:
+        torch.distributed.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()