Publish PrecipDD reproduction

.gitattributes

@@ -1,35 +1,7 @@
-*.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
+*.py text eol=lf
+*.md text eol=lf
+*.json text eol=lf
+*.yaml text eol=lf
+*.npz binary
+*.pt binary
+*.png binary

conf/config.yaml

@@ -0,0 +1,64 @@
+project:
+  name: PrecipDD
+  format_version: precipdd_v1
+  seed: 23
+  license: Apache-2.0
+
+paths:
+  data: data/daily_precip_anomaly.npz
+  checkpoint: result/checkpoints/precipdd.pt
+  training_metrics: result/training/metrics.json
+  predictions: result/output/predictions.npz
+  evaluation_metrics: result/evaluation/metrics.json
+  comparison_figure: result/evaluation/comparison.png
+
+data:
+  layout: NCHW
+  input_shape: [1, 55, 160]
+  latitude_range: [-60.0, 75.0]
+  latitude_note: 55 engineering grid centers at 2.5 degree spacing; the paper text reports inconsistent endpoints alongside 55 points
+  longitude_range: [0.0, 397.5]
+  longitude_note: 0-360 E plus repeated 0-40 E extension
+  normalization: daily climatology anomaly divided by zonal-mean anomaly standard deviation at each latitude
+  train_samples: 96
+  validation_samples: 32
+  test_days_per_year: 6
+  test_years: [1980, 2020]
+  synthetic_members: 8
+  paper_members: 80
+  paper_period: [1850, 2100]
+  paper_train_days: 730000
+  paper_validation_days: 146000
+
+model:
+  filters: [8, 8, 16, 16, 16]
+  kernel_size: 3
+  dense_units: 32
+  conv_activation: tanh
+  dense_activation: sigmoid
+  output_activation: linear
+  expected_feature_shape: [16, 14, 40]
+  flattened_features: 8960
+
+training:
+  ensemble_members: 2
+  paper_ensemble_members: 5
+  epochs: 3
+  batch_size: 8
+  learning_rate: 0.001
+  l2_weight_decay: 0.0001
+  loss: MAE
+  optimizer: Adam
+  initialization: Xavier-uniform
+  num_workers: 0
+
+runtime:
+  device: auto
+  ddp_backend_gpu: nccl
+  ddp_backend_cpu: gloo
+
+evaluation:
+  emergence_threshold_c: 0.42
+  occlusion_patch: 7
+  occlusion_stride: 8
+  occlusion_max_days: 12

config.json

@@ -0,0 +1,12 @@
+{
+  "model_name": "PrecipDD",
+  "model_type": "precipdd",
+  "architectures": ["PrecipDD"],
+  "framework": "PyTorch",
+  "domain": "climate",
+  "task": "daily-precipitation-anomaly-to-agmt-regression",
+  "implementation": {"entry_point": "model/precipdd.py", "scope": "dimension-faithful reduced-sample 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": {"input_shape": [1,55,160], "conv_layers": 5, "filters": [8,8,16,16,16], "pooled_feature_shape": [16,14,40], "flattened_features": 8960, "dense": [32,1], "output_activation": "linear"},
+  "data": {"datasets": ["CESM2 Large Ensemble"], "format_version": "precipdd_v1", "layout": "NCHW", "paper_members": 80, "paper_period": [1850,2100], "synthetic": true},
+  "configuration_sources": ["conf/config.yaml","model/precipdd.py","scripts/fake_data.py","scripts/train.py","scripts/inference.py","scripts/result.py"]
+}

model/precipdd.py

@@ -0,0 +1,134 @@
+"""Dimension-faithful PyTorch implementation of the precipitation DD CNN."""
+
+from __future__ import annotations
+
+import json
+import math
+import random
+from pathlib import Path
+from typing import Iterable
+
+import numpy as np
+import torch
+import yaml
+from torch import Tensor, nn
+
+
+DATA_FORMAT_VERSION = "precipdd_v1"
+INPUT_SHAPE = (1, 55, 160)
+
+
+def load_config(path: str | Path) -> dict:
+    with open(path, "r", encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def seed_all(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+class PrecipDD(nn.Module):
+    """Five-convolution scalar regressor with the paper's 8,960 features."""
+
+    def __init__(self, filters: Iterable[int] = (8, 8, 16, 16, 16), dense_units: int = 32):
+        super().__init__()
+        filters = tuple(int(value) for value in filters)
+        if len(filters) != 5 or filters[-1] != 16:
+            raise ValueError("DD requires five convolution layers and 16 final filters")
+        layers: list[nn.Module] = []
+        channels = 1
+        for index, width in enumerate(filters):
+            layers.extend((nn.Conv2d(channels, width, 3, padding=1), nn.Tanh()))
+            if index < 2:
+                # TensorFlow SAME pooling is required for 55 -> 28 -> 14 latitude points.
+                layers.append(nn.MaxPool2d(2, stride=2, ceil_mode=True))
+            channels = width
+        self.features = nn.Sequential(*layers)
+        self.hidden = nn.Linear(16 * 14 * 40, dense_units)
+        self.output = nn.Linear(dense_units, 1)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        for module in self.modules():
+            if isinstance(module, (nn.Conv2d, nn.Linear)):
+                nn.init.xavier_uniform_(module.weight)
+                fan_in, fan_out = nn.init._calculate_fan_in_and_fan_out(module.weight)
+                bound = math.sqrt(6.0 / (fan_in + fan_out))
+                nn.init.uniform_(module.bias, -bound, bound)
+
+    def forward_features(self, precipitation: Tensor) -> Tensor:
+        if precipitation.ndim != 4 or tuple(precipitation.shape[1:]) != INPUT_SHAPE:
+            raise ValueError(f"input must have shape [B,1,55,160], got {tuple(precipitation.shape)}")
+        features = self.features(precipitation)
+        if tuple(features.shape[1:]) != (16, 14, 40):
+            raise RuntimeError(f"feature shape must be [B,16,14,40], got {tuple(features.shape)}")
+        return features
+
+    def forward(self, precipitation: Tensor) -> Tensor:
+        features = self.forward_features(precipitation).flatten(1)
+        return self.output(torch.sigmoid(self.hidden(features))).squeeze(-1)
+
+
+def validate_archive(archive: np.lib.npyio.NpzFile) -> None:
+    required = {"precipitation", "agmt", "split", "year", "day_of_year", "latitude", "longitude", "format_version"}
+    if missing := required.difference(archive.files):
+        raise ValueError(f"dataset missing fields: {sorted(missing)}")
+    if str(archive["format_version"]) != DATA_FORMAT_VERSION:
+        raise ValueError("dataset format_version mismatch")
+    if archive["precipitation"].ndim != 4 or tuple(archive["precipitation"].shape[1:]) != INPUT_SHAPE:
+        raise ValueError("precipitation must be float data with shape [N,1,55,160]")
+    if archive["agmt"].shape != (len(archive["precipitation"]),):
+        raise ValueError("AGMT must contain one scalar for every daily map")
+    if archive["latitude"].shape != (55,) or archive["longitude"].shape != (160,):
+        raise ValueError("coordinates must contain 55 latitudes and 160 extended longitudes")
+    if not np.all(np.isfinite(archive["precipitation"])) or not np.all(np.isfinite(archive["agmt"])):
+        raise ValueError("dataset contains non-finite values")
+
+
+def load_ensemble(checkpoint_path: str | Path, device: torch.device) -> tuple[list[PrecipDD], dict]:
+    try:
+        checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    except TypeError:
+        checkpoint = torch.load(checkpoint_path, map_location=device)
+    if checkpoint.get("format_version") != DATA_FORMAT_VERSION:
+        raise ValueError("checkpoint format_version mismatch")
+    settings = checkpoint["model_config"]
+    models = []
+    for state in checkpoint["ensemble_states"]:
+        model = PrecipDD(settings["filters"], settings["dense_units"]).to(device)
+        model.load_state_dict(state)
+        model.eval()
+        models.append(model)
+    return models, checkpoint
+
+
+@torch.no_grad()
+def ensemble_predict(models: list[PrecipDD], values: Tensor, batch_size: int = 32) -> Tensor:
+    predictions = []
+    for start in range(0, len(values), batch_size):
+        batch = values[start:start + batch_size]
+        predictions.append(torch.stack([model(batch) for model in models]).mean(0))
+    return torch.cat(predictions) if predictions else torch.empty(0, device=values.device)
+
+
+def linear_trend(values: np.ndarray, years: np.ndarray) -> float:
+    valid = np.isfinite(values) & np.isfinite(years)
+    if valid.sum() < 2 or np.ptp(years[valid]) == 0:
+        return float("nan")
+    return float(np.polyfit(years[valid], values[valid], 1)[0] * 10.0)
+
+
+def correlation(target: np.ndarray, prediction: np.ndarray) -> float:
+    if len(target) < 2 or np.std(target) == 0 or np.std(prediction) == 0:
+        return float("nan")
+    return float(np.corrcoef(target, prediction)[0, 1])
+
+
+def write_json(path: str | Path, payload: dict) -> None:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(payload, indent=2, allow_nan=False) + "\n", encoding="utf-8")

scripts/fake_data.py

@@ -0,0 +1,72 @@
+"""Create spatially coherent daily anomalies with weather and warming signals."""
+
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.precipdd import DATA_FORMAT_VERSION, load_config
+
+
+def agmt_for_year(year, member, rng):
+    forced = -0.45 + 0.0042 * (year - 1850) + 0.000030 * max(year - 1970, 0) ** 2
+    return forced + 0.08 * np.sin((year - 1850) / 8.0 + member) + rng.normal(0, 0.045)
+
+
+def main():
+    config = load_config(ROOT / "conf/config.yaml")
+    settings = config["data"]
+    rng = np.random.default_rng(config["project"]["seed"])
+    train_n, val_n = settings["train_samples"], settings["validation_samples"]
+    first, last = settings["test_years"]
+    test_year = np.repeat(np.arange(first, last + 1), settings["test_days_per_year"])
+    train_year = rng.integers(1850, 2101, train_n)
+    val_year = rng.integers(1850, 2101, val_n)
+    year = np.concatenate((train_year, val_year, test_year)).astype(np.int16)
+    split = np.concatenate((np.zeros(train_n, np.int8), np.ones(val_n, np.int8), np.full(len(test_year), 2, np.int8)))
+    day = np.concatenate((rng.integers(1, 366, train_n + val_n), np.tile(np.linspace(15, 350, settings["test_days_per_year"], dtype=int), last - first + 1))).astype(np.int16)
+    member = rng.integers(0, settings["synthetic_members"], len(year), dtype=np.int16)
+    agmt = np.array([agmt_for_year(int(y), int(m), rng) for y, m in zip(year, member)], dtype=np.float32)
+
+    latitude = np.arange(-60.0, 77.5, 2.5, dtype=np.float32)
+    base_longitude = np.arange(0.0, 360.0, 2.5, dtype=np.float32)
+    longitude = np.arange(0.0, 400.0, 2.5, dtype=np.float32)
+    lat2d, lon2d = np.meshgrid(latitude, base_longitude, indexing="ij")
+    east_pacific = np.exp(-((lat2d / 16) ** 2 + ((lon2d - 245) / 35) ** 2))
+    storm_tracks = np.exp(-((np.abs(lat2d) - 45) / 12) ** 2) * (0.65 + 0.35 * np.cos(np.deg2rad(lon2d * 2)))
+    fingerprint = (east_pacific + storm_tracks).astype(np.float32)
+    fingerprint /= fingerprint.max()
+    coarse = torch.from_numpy(rng.normal(size=(len(year), 1, 14, 36)).astype(np.float32))
+    weather = F.interpolate(coarse, size=(55, 144), mode="bilinear", align_corners=False).numpy()
+    synoptic = np.sin(2 * np.pi * day[:, None, None] / 7.0 + np.deg2rad(lon2d)[None])
+    noise = rng.normal(0, 0.28, size=weather.shape).astype(np.float32)
+    amplitude = 1.0 + 0.34 * np.maximum(agmt, -0.5)[:, None, None] * fingerprint[None]
+    fields = weather[:, 0] * amplitude + 0.32 * synoptic * (0.3 + fingerprint[None]) + noise[:, 0]
+    fields += 0.22 * agmt[:, None, None] * (fingerprint[None] - 0.35)
+    fields = fields.astype(np.float32)
+    fields = np.concatenate((fields, fields[:, :, :16]), axis=2)
+    fingerprint = np.concatenate((fingerprint, fingerprint[:, :16]), axis=1)
+    fields -= fields.mean(axis=2, keepdims=True)
+    zonal_std = fields.std(axis=2, keepdims=True).mean(axis=0, keepdims=True)
+    fields /= np.maximum(zonal_std, 1e-5)
+
+    output = ROOT / config["paths"]["data"]
+    output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(output, format_version=np.array(DATA_FORMAT_VERSION), precipitation=fields[:, None], agmt=agmt,
+                        split=split, year=year, day_of_year=day, member=member, latitude=latitude,
+                        longitude=longitude, synthetic=np.array(True), fingerprint=fingerprint)
+    metadata = {"synthetic": True, "samples": len(year), "splits": {"train": train_n, "validation": val_n, "test": len(test_year)},
+                "precipitation_shape": [len(year), 1, 55, 160], "target_shape": [len(year)],
+                "science_note": "Structured engineering data with daily weather variability and an AGMT-dependent spatial variance signal; not CESM2 LE."}
+    (output.parent / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8")
+    print(f"data={output.relative_to(ROOT)} shape={fields[:, None].shape} target={agmt.shape}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,34 @@
+"""Run ensemble inference for all held-out synthetic daily maps."""
+
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.precipdd import ensemble_predict, load_config, load_ensemble, validate_archive
+
+
+def main():
+    config = load_config(ROOT / "conf/config.yaml")
+    device = torch.device("cuda" if torch.cuda.is_available() and config["runtime"]["device"] != "cpu" else "cpu")
+    data = np.load(ROOT / config["paths"]["data"])
+    validate_archive(data)
+    models, checkpoint = load_ensemble(ROOT / config["paths"]["checkpoint"], device)
+    mask = data["split"] == 2
+    fields = torch.from_numpy(data["precipitation"][mask]).float().to(device)
+    prediction = ensemble_predict(models, fields, config["training"]["batch_size"]).cpu().numpy()
+    output = ROOT / config["paths"]["predictions"]
+    output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(output, format_version=data["format_version"], prediction=prediction, target=data["agmt"][mask],
+                        precipitation=data["precipitation"][mask], year=data["year"][mask], day_of_year=data["day_of_year"][mask],
+                        latitude=data["latitude"], longitude=data["longitude"], ensemble_members=np.array(len(models)),
+                        checkpoint_world_size=np.array(checkpoint["world_size"]))
+    print(f"predictions={output.relative_to(ROOT)} days={len(prediction)} members={len(models)}")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,83 @@
+"""Compute daily/annual detection metrics and a 7x7 occlusion trend map."""
+
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.precipdd import correlation, ensemble_predict, linear_trend, load_config, load_ensemble, write_json
+
+
+def main():
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    config = load_config(ROOT / "conf/config.yaml")
+    values = np.load(ROOT / config["paths"]["predictions"])
+    prediction, target, year = values["prediction"], values["target"], values["year"].astype(float)
+    if prediction.shape != target.shape or values["precipitation"].shape[1:] != (1, 55, 160):
+        raise ValueError("inference output violates scalar target or [N,1,55,160] input contract")
+    years = np.unique(year).astype(int)
+    annual_prediction = np.array([prediction[year == current].mean() for current in years])
+    annual_target = np.array([target[year == current].mean() for current in years])
+    threshold = config["evaluation"]["emergence_threshold_c"]
+    em_fraction = np.array([(prediction[year == current] > threshold).mean() for current in years])
+    metrics = {
+        "daily": {"correlation": correlation(target, prediction), "rmse_c": float(np.sqrt(np.mean((prediction - target) ** 2)))},
+        "annual": {"correlation": correlation(annual_target, annual_prediction), "rmse_c": float(np.sqrt(np.mean((annual_prediction - annual_target) ** 2)))},
+        "emergence": {"threshold_c": threshold, "fraction_all_days": float((prediction > threshold).mean()),
+                      "fraction_trend_per_decade": linear_trend(em_fraction, years.astype(float))},
+        "trend_c_per_decade": {"daily_prediction": linear_trend(prediction, year), "annual_prediction": linear_trend(annual_prediction, years.astype(float)),
+                               "annual_target": linear_trend(annual_target, years.astype(float))},
+        "samples": {"daily": len(prediction), "years": len(years)}
+    }
+
+    device = torch.device("cuda" if torch.cuda.is_available() and config["runtime"]["device"] != "cpu" else "cpu")
+    models, _ = load_ensemble(ROOT / config["paths"]["checkpoint"], device)
+    count = min(config["evaluation"]["occlusion_max_days"], len(prediction))
+    indices = np.linspace(0, len(prediction) - 1, count, dtype=int)
+    selected = torch.from_numpy(values["precipitation"][indices]).float().to(device)
+    baseline = ensemble_predict(models, selected).cpu().numpy()
+    selected_year = year[indices]
+    patch, stride = config["evaluation"]["occlusion_patch"], config["evaluation"]["occlusion_stride"]
+    half = patch // 2
+    sensitivity = np.zeros((55, 160), dtype=np.float32)
+    for lat_start in range(0, 55, stride):
+        for lon_start in range(0, 160, stride):
+            masked = selected.clone()
+            lat_stop, lon_stop = min(lat_start + patch, 55), min(lon_start + patch, 160)
+            masked[:, :, lat_start:lat_stop, lon_start:lon_stop] = 0.0
+            delta = baseline - ensemble_predict(models, masked).cpu().numpy()
+            score = linear_trend(delta, selected_year) if np.ptp(selected_year) else float(delta.mean())
+            sensitivity[lat_start:min(lat_start + stride, 55), lon_start:min(lon_start + stride, 160)] = score
+    metrics["occlusion"] = {"patch": [patch, patch], "map_shape": [55, 160], "stride": stride, "sampled_days": count,
+                            "quantity": "AGMT occlusion-sensitivity trend in degC per decade"}
+    write_json(ROOT / config["paths"]["evaluation_metrics"], metrics)
+
+    fig, axes = plt.subplots(3, 1, figsize=(11, 11), constrained_layout=True)
+    axes[0].plot(years, annual_target, color="#202020", label="target AGMT")
+    axes[0].plot(years, annual_prediction, color="#c84c32", label="DD estimate")
+    axes[0].axhline(threshold, color="#777777", linestyle="--", label="0.42 C EM threshold")
+    axes[0].set(ylabel="AGMT anomaly (C)", title="Annual mean of daily estimates")
+    axes[0].legend(ncol=3)
+    axes[1].plot(years, em_fraction, color="#196f82")
+    axes[1].set(xlabel="Year", ylabel="Fraction", ylim=(-0.03, 1.03), title="Emergence days (estimated AGMT > 0.42 C)")
+    limit = float(np.max(np.abs(sensitivity))) or 1e-6
+    image = axes[2].imshow(sensitivity, origin="lower", aspect="auto", extent=(0, 400, values["latitude"][0], values["latitude"][-1]),
+                           cmap="RdBu_r", vmin=-limit, vmax=limit)
+    axes[2].set(xlabel="Longitude (degrees E, extended)", ylabel="Latitude", title="7x7 occlusion-sensitivity trend")
+    fig.colorbar(image, ax=axes[2], label="C decade-1")
+    figure = ROOT / config["paths"]["comparison_figure"]
+    figure.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(figure, dpi=160)
+    plt.close(fig)
+    print(f"metrics={config['paths']['evaluation_metrics']} figure={config['paths']['comparison_figure']}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,93 @@
+"""Train independently initialized DD members with optional DDP."""
+
+import copy
+import json
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler, TensorDataset
+
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.precipdd import DATA_FORMAT_VERSION, PrecipDD, load_config, seed_all, validate_archive
+
+
+def main():
+    config = load_config(ROOT / "conf/config.yaml")
+    runtime, settings = config["runtime"], config["training"]
+    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", "1"))
+    use_cuda = (torch.cuda.is_available() and runtime["device"] != "cpu"
+                and torch.cuda.device_count() >= local_world_size)
+    if distributed:
+        torch.distributed.init_process_group(runtime["ddp_backend_gpu"] if use_cuda else runtime["ddp_backend_cpu"])
+    rank = torch.distributed.get_rank() if distributed else 0
+    device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu")
+    if use_cuda:
+        torch.cuda.set_device(device)
+    data = np.load(ROOT / config["paths"]["data"])
+    validate_archive(data)
+    train_mask, val_mask = data["split"] == 0, data["split"] == 1
+    dataset = TensorDataset(torch.from_numpy(data["precipitation"][train_mask]).float(), torch.from_numpy(data["agmt"][train_mask]).float())
+    val_x = torch.from_numpy(data["precipitation"][val_mask]).float().to(device)
+    val_y = torch.from_numpy(data["agmt"][val_mask]).float().to(device)
+    ensemble_states, histories = [], []
+    for member in range(settings["ensemble_members"]):
+        seed_all(config["project"]["seed"] + member)
+        sampler = DistributedSampler(dataset, shuffle=True, seed=member) if distributed else None
+        loader = DataLoader(dataset, batch_size=settings["batch_size"], shuffle=sampler is None, sampler=sampler,
+                            num_workers=settings["num_workers"])
+        model = PrecipDD(config["model"]["filters"], config["model"]["dense_units"]).to(device)
+        wrapped = DistributedDataParallel(model, device_ids=[local_rank] if use_cuda else None) if distributed else model
+        optimizer = torch.optim.Adam(wrapped.parameters(), lr=settings["learning_rate"], weight_decay=settings["l2_weight_decay"])
+        history, best_loss, best_state = [], float("inf"), None
+        for epoch in range(settings["epochs"]):
+            if sampler is not None:
+                sampler.set_epoch(member * settings["epochs"] + epoch)
+            wrapped.train()
+            total, count = 0.0, 0
+            for fields, target in loader:
+                fields, target = fields.to(device), target.to(device)
+                optimizer.zero_grad(set_to_none=True)
+                loss = nn.functional.l1_loss(wrapped(fields), target)
+                loss.backward()
+                optimizer.step()
+                total += float(loss.detach()) * len(fields)
+                count += len(fields)
+            summary = torch.tensor([total, count], dtype=torch.float64, device=device)
+            if distributed:
+                torch.distributed.all_reduce(summary)
+            wrapped.eval()
+            with torch.no_grad():
+                val_loss = float(nn.functional.l1_loss(model(val_x), val_y))
+            history.append({"epoch": epoch + 1, "train_mae": float(summary[0] / summary[1]), "validation_mae": val_loss})
+            if val_loss < best_loss:
+                best_loss = val_loss
+                best_state = {key: value.detach().cpu().clone() for key, value in model.state_dict().items()}
+        ensemble_states.append(best_state)
+        histories.append(history)
+    if rank == 0:
+        checkpoint_path = ROOT / config["paths"]["checkpoint"]
+        checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
+        torch.save({"format_version": DATA_FORMAT_VERSION, "ensemble_states": ensemble_states,
+                    "model_config": {"filters": config["model"]["filters"], "dense_units": config["model"]["dense_units"],
+                                     "input_shape": [1, 55, 160], "feature_shape": [16, 14, 40], "flattened_features": 8960},
+                    "training_config": settings, "histories": histories, "world_size": torch.distributed.get_world_size() if distributed else 1}, checkpoint_path)
+        metrics_path = ROOT / config["paths"]["training_metrics"]
+        metrics_path.parent.mkdir(parents=True, exist_ok=True)
+        metrics_path.write_text(json.dumps({"ensemble_members": len(histories), "history": histories}, indent=2) + "\n", encoding="utf-8")
+        print(f"checkpoint={checkpoint_path.relative_to(ROOT)} members={len(ensemble_states)}")
+    if distributed:
+        torch.distributed.barrier()
+        torch.distributed.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()