Publish EchoCast-3D reproduction

.gitattributes

@@ -1,35 +1,2 @@
-*.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
+*.npz filter=lfs diff=lfs merge=lfs -text

conf/config.yaml

@@ -0,0 +1,53 @@
+seed: 42
+data:
+  format_version: echocast_3d_v1
+  frames: 8
+  history: 3
+  target: 5
+  interval_minutes: 6
+  elevations_deg: [0.5, 1.5, 2.4, 3.4]
+  azimuth_counts: [366, 366, 363, 363]
+  range_bins: [180, 180, 120, 120]
+  storage_shape: [8, 4, 366, 180]
+  missing_ratio: 0.15
+  dbz_min: 0.0
+  dbz_max: 70.0
+  train_samples: 1
+  validation_samples: 1
+model:
+  patch_size: [3, 3]
+  hidden_dim: 32
+  depth: 2
+  heads: 4
+  latent_tokens: 16
+  mask_ratio: 0.75
+  paper_token_count: 24400
+diffusion:
+  steps: 4
+  beta_start: 0.0001
+  beta_end: 0.02
+  ensemble_size: 3
+training:
+  epochs: 1
+  batch_size: 1
+  learning_rate: 0.00001
+  weight_decay: 0.01
+  reconstruction_weight: 0.1
+runtime:
+  device: cpu
+paths:
+  checkpoint: result/checkpoints/echocast_3d.pt
+  training_metrics: result/training/metrics.json
+  predictions: result/output/predictions.npz
+  evaluation: result/evaluation/metrics.json
+  figure: result/evaluation/comparison.png
+evaluation:
+  thresholds_dbz: [20.0, 30.0, 40.0]
+paper_model:
+  hidden_dim: 1024
+  depth: 12
+  heads: 16
+  batch_size: 48
+  epochs: 200
+  learning_rate: 0.00001
+  ensemble_size: 30

config.json

@@ -0,0 +1 @@
+{"model_name":"EchoCast-3D","model_type":"echocast_3d","architectures":["EchoCast3D"],"framework":"PyTorch","domain":"weather","task":"probabilistic-3d-radar-nowcasting","implementation":{"entry_point":"model/echocast_3d.py","scope":"core-method full-radar-geometry 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":{"storage_shape":["B",8,4,366,180],"history_frames":3,"target_frames":5,"patch_size":[3,3],"paper_hidden_dim":1024,"paper_depth":12,"paper_heads":16,"paper_tokens":24400},"data":{"dataset":"CMA volumetric radar","format_version":"echocast_3d_v1","azimuth_counts":[366,366,363,363],"range_bins":[180,180,120,120],"interval_minutes":6,"synthetic":true},"configuration_sources":["conf/config.yaml","model/echocast_3d.py","scripts/fake_data.py","scripts/train.py","scripts/inference.py","scripts/result.py"]}

model/echocast_3d.py

@@ -0,0 +1,199 @@
+"""MaskDiT-style denoiser for native packed radar wedges."""
+
+from __future__ import annotations
+
+import math
+import json
+from pathlib import Path
+import numpy as np
+import yaml
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.utils.data import Dataset
+
+
+def load_config(path):
+    with open(path, encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def geometry_mask(config):
+    data = config["data"]
+    mask = np.zeros((data["frames"], 4, 366, 180), dtype=bool)
+    for elevation, (azimuths, ranges) in enumerate(zip(data["azimuth_counts"], data["range_bins"])):
+        mask[:, elevation, :azimuths, :ranges] = True
+    return mask
+
+
+def make_sample(config, seed):
+    rng = np.random.default_rng(seed); data = config["data"]
+    volume = np.zeros(tuple(data["storage_shape"]), dtype=np.float32)
+    phase = rng.uniform(0, 2*np.pi)
+    for t in range(data["frames"]):
+        for e, (na, nr) in enumerate(zip(data["azimuth_counts"], data["range_bins"])):
+            az=np.arange(na,dtype=np.float32)[:,None]; radius=np.arange(nr,dtype=np.float32)[None]
+            center=(80+12*t+8*e+7*np.sin(phase+t/3))%na; center_r=min(nr-15,42+4*t-5*e)
+            da=np.minimum(abs(az-center),na-abs(az-center)); core=np.exp(-.5*((da/(14+2*e))**2+((radius-center_r)/18)**2))
+            volume[t,e,:na,:nr]=np.clip(58*np.exp(-.22*e)*core+2*np.sin(radius/8+az/30+phase+t/2),0,70)
+    valid=geometry_mask(config); observed=valid.copy()
+    for t in range(data["history"]):
+        for e,(na,nr) in enumerate(zip(data["azimuth_counts"],data["range_bins"])):
+            width=max(1,int(round(na*data["missing_ratio"]))); start=int(rng.integers(0,na-width+1)); observed[t,e,start:start+width,:nr]=False
+    packed=volume.copy();packed[~observed]=0
+    return {"values":packed,"truth":volume,"validity":valid,"observed":observed}
+
+
+class PackedRadarDataset(Dataset):
+    def __init__(self,directory,split):
+        self.files=sorted(Path(directory).glob(f"{split}_*.npz"))
+        if not self.files: raise FileNotFoundError(f"no {split} samples; run fake_data.py")
+    def __len__(self): return len(self.files)
+    def __getitem__(self,index):
+        with np.load(self.files[index]) as sample:
+            return {key:torch.from_numpy(sample[key].copy()) for key in ("values","truth","validity","observed")}
+
+
+class DiffusionSchedule:
+    def __init__(self,steps,beta_start,beta_end,device):
+        self.steps=steps;self.betas=torch.linspace(beta_start,beta_end,steps,device=device);self.alphas=1-self.betas;self.alpha_bars=torch.cumprod(self.alphas,0)
+    def add_noise(self,clean,timestep):
+        noise=torch.randn_like(clean);alpha=self.alpha_bars[timestep].view(-1,1,1,1,1);return alpha.sqrt()*clean+(1-alpha).sqrt()*noise,noise
+    @torch.no_grad()
+    def sample(self,model,history,observed):
+        target=torch.randn(history.shape[0],5,*history.shape[2:],device=history.device)
+        for step in reversed(range(self.steps)):
+            t=torch.full((history.shape[0],),step,device=history.device,dtype=torch.long);clean=model(history,target,observed,t,None)
+            target=clean if step==0 else self.alpha_bars[step-1].sqrt()*clean+(1-self.alpha_bars[step-1]).sqrt()*torch.randn_like(clean)
+        return target
+
+
+def evaluate_ensemble(ensemble,truth,valid,thresholds):
+    mean=ensemble.mean(0);error=mean-truth;result={"mae":float(abs(error)[valid].mean()),"rmse":float(np.sqrt((error**2)[valid].mean()))}
+    result["crps"]=float(abs(ensemble-truth[None])[:,valid].mean()-.5*abs(ensemble[:,None]-ensemble[None,:])[:,:,valid].mean())
+    result["categorical"]={}
+    for threshold in thresholds:
+        forecast,observed=mean>=threshold,truth>=threshold;hits=int((forecast&observed&valid).sum());miss=int((~forecast&observed&valid).sum());false=int((forecast&~observed&valid).sum())
+        result["categorical"][str(threshold)]={"csi":hits/max(hits+miss+false,1),"far":false/max(hits+false,1),"pod":hits/max(hits+miss,1)}
+    result["coverage_ratio"]=[[float(np.isfinite(ensemble).all(0)[lead,e][valid[lead,e]].mean()) for e in range(4)] for lead in range(5)]
+    return result
+
+
+class DiTBlock(nn.Module):
+    """Linear-cost latent attention block suitable for 24k native wedge tokens."""
+
+    def __init__(self, hidden: int, heads: int, latent_tokens: int):
+        super().__init__()
+        self.latents = nn.Parameter(torch.randn(1, latent_tokens, hidden) / math.sqrt(hidden))
+        self.read = nn.MultiheadAttention(hidden, heads, batch_first=True)
+        self.write = nn.MultiheadAttention(hidden, heads, batch_first=True)
+        self.norm1 = nn.LayerNorm(hidden)
+        self.norm2 = nn.LayerNorm(hidden)
+        self.mlp = nn.Sequential(nn.Linear(hidden, hidden * 4), nn.GELU(), nn.Linear(hidden * 4, hidden))
+
+    def forward(self, tokens: torch.Tensor, conditioning: torch.Tensor) -> torch.Tensor:
+        latent = self.latents.expand(tokens.shape[0], -1, -1) + conditioning[:, None]
+        latent = latent + self.read(latent, self.norm1(tokens), self.norm1(tokens), need_weights=False)[0]
+        tokens = tokens + self.write(self.norm1(tokens), latent, latent, need_weights=False)[0]
+        return tokens + self.mlp(self.norm2(tokens + conditioning[:, None]))
+
+
+class EchoCast3D(nn.Module):
+    def __init__(self, config: dict):
+        super().__init__()
+        self.config = config
+        model = config["model"]
+        self.patch = tuple(model["patch_size"])
+        patch_area = self.patch[0] * self.patch[1]
+        self.embed = nn.Linear(config["data"]["frames"] * patch_area, model["hidden_dim"])
+        self.time_embed = nn.Embedding(config["diffusion"]["steps"], model["hidden_dim"])
+        self.elevation_embed = nn.Embedding(4, model["hidden_dim"])
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, model["hidden_dim"]))
+        self.encoder = nn.ModuleList(
+            [DiTBlock(model["hidden_dim"], model["heads"], model["latent_tokens"]) for _ in range(model["depth"])]
+        )
+        self.decoder = nn.ModuleList(
+            [DiTBlock(model["hidden_dim"], model["heads"], model["latent_tokens"]) for _ in range(model["depth"])]
+        )
+        self.head = nn.Linear(model["hidden_dim"], config["data"]["target"] * patch_area)
+        self.layouts = []
+        offset = 0
+        for e, (na, nr) in enumerate(zip(config["data"]["azimuth_counts"], config["data"]["range_bins"])):
+            rows, cols = na // self.patch[0], nr // self.patch[1]
+            count = rows * cols
+            self.layouts.append((e, na, nr, rows, cols, offset, offset + count))
+            offset += count
+        self.token_count = offset
+
+    def patchify(self, volume: torch.Tensor) -> torch.Tensor:
+        parts = []
+        for e, na, nr, _, _, _, _ in self.layouts:
+            wedge = volume[:, :, e, :na, :nr]
+            patches = F.unfold(wedge, kernel_size=self.patch, stride=self.patch).transpose(1, 2)
+            parts.append(patches)
+        return torch.cat(parts, dim=1)
+
+    def token_validity(self, validity: torch.Tensor) -> torch.Tensor:
+        parts = []
+        for e, na, nr, _, _, _, _ in self.layouts:
+            wedge = validity[:, :, e, :na, :nr].float()
+            patch = F.unfold(wedge, kernel_size=self.patch, stride=self.patch)
+            parts.append((patch > 0.5).all(dim=1))
+        return torch.cat(parts, dim=1)
+
+    def unpatchify(self, patches: torch.Tensor) -> torch.Tensor:
+        batch = patches.shape[0]
+        output = patches.new_zeros((batch, 5, 4, 366, 180))
+        for e, na, nr, rows, cols, start, end in self.layouts:
+            wedge = F.fold(
+                patches[:, start:end].transpose(1, 2),
+                output_size=(na, nr),
+                kernel_size=self.patch,
+                stride=self.patch,
+            )
+            output[:, :, e, :na, :nr] = wedge
+        return output
+
+    def forward(
+        self,
+        history: torch.Tensor,
+        noisy_target: torch.Tensor,
+        observed: torch.Tensor,
+        timestep: torch.Tensor,
+        token_mask: torch.Tensor | None,
+    ) -> torch.Tensor:
+        full = torch.cat((history, noisy_target), dim=1)
+        tokens = self.embed(self.patchify(full))
+        elevation_ids = torch.cat(
+            [torch.full((end - start,), e, device=full.device) for e, _, _, _, _, start, end in self.layouts]
+        )
+        tokens = tokens + self.elevation_embed(elevation_ids)[None]
+        if token_mask is not None:
+            tokens = torch.where(token_mask[..., None], self.mask_token.expand_as(tokens), tokens)
+        condition = self.time_embed(timestep)
+        for block in self.encoder:
+            tokens = block(tokens, condition)
+        for block in self.decoder:
+            tokens = block(tokens, condition)
+        return self.unpatchify(self.head(tokens))
+
+
+def masked_objective(
+    model: EchoCast3D,
+    prediction: torch.Tensor,
+    truth: torch.Tensor,
+    validity: torch.Tensor,
+    token_mask: torch.Tensor,
+    reconstruction_weight: float,
+) -> tuple[torch.Tensor, dict[str, float]]:
+    target_valid = validity[:, 3:]
+    token_pixels = token_mask[:, :, None].expand(-1, -1, 5 * 9)
+    masked_pixels = model.unpatchify(token_pixels.float()).bool()
+    squared = (prediction - truth) ** 2
+    unmasked = target_valid & ~masked_pixels
+    masked = target_valid & masked_pixels
+    score = squared[unmasked].mean()
+    reconstruction = squared[masked].mean()
+    loss = score + reconstruction_weight * reconstruction
+    return loss, {"score": score.item(), "reconstruction": reconstruction.item()}

scripts/fake_data.py

@@ -0,0 +1,32 @@
+"""Create realistic packed-wedge synthetic data when licensed observations are unavailable."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import sys
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+
+from model.echocast_3d import load_config, make_sample
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="conf/config.yaml")
+    parser.add_argument("--output", default="data")
+    args = parser.parse_args()
+    config = load_config(ROOT / args.config)
+    output = ROOT / args.output
+    output.mkdir(parents=True, exist_ok=True)
+    for split, count in (("train", config["data"]["train_samples"]), ("validation", config["data"]["validation_samples"])):
+        for index in range(count):
+            sample = make_sample(config, config["seed"] + index + (10000 if split == "validation" else 0))
+            path = output / f"{split}_{index:04d}.npz"
+            np.savez(path, **sample)
+            print(f"wrote {path} ({path.stat().st_size / 1024**2:.2f} MiB)")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,50 @@
+"""Generate native packed-wedge ensemble forecasts."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import torch
+import sys
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+
+from model.echocast_3d import PackedRadarDataset, load_config, DiffusionSchedule, EchoCast3D
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="conf/config.yaml")
+    parser.add_argument("--data", default="data")
+    parser.add_argument("--checkpoint", default="result/checkpoints/echocast_3d.pt")
+    parser.add_argument("--output", default="result/output/predictions.npz")
+    args = parser.parse_args()
+    config = load_config(ROOT / args.config)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = EchoCast3D(config).to(device)
+    checkpoint = torch.load(ROOT / args.checkpoint, map_location=device, weights_only=False)
+    model.load_state_dict(checkpoint["model"])
+    model.eval()
+    sample = PackedRadarDataset(ROOT / args.data, "validation")[0]
+    history = sample["values"][None, :3].float().to(device)
+    observed = sample["observed"][None, :3].bool().to(device)
+    schedule = DiffusionSchedule(device=device, **{k: config["diffusion"][k] for k in ("steps", "beta_start", "beta_end")})
+    ensemble = []
+    for member in range(config["diffusion"]["ensemble_size"]):
+        torch.manual_seed(config["seed"] + member)
+        ensemble.append(schedule.sample(model, history, observed)[0].cpu().numpy())
+    output = ROOT / args.output
+    output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez(
+        output,
+        ensemble=np.stack(ensemble),
+        truth=sample["truth"][3:].numpy(),
+        validity=sample["validity"][3:].numpy(),
+        observed_history=sample["observed"][:3].numpy(),
+    )
+    print(f"saved {output}: ensemble={len(ensemble)}, steps={schedule.steps}")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,50 @@
+"""Evaluate the ensemble and render a 2D composite-reflectivity summary."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import sys
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+
+from model.echocast_3d import load_config, evaluate_ensemble
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="conf/config.yaml")
+    parser.add_argument("--predictions", default="result/output/predictions.npz")
+    parser.add_argument("--metrics", default="result/evaluation/metrics.json")
+    parser.add_argument("--figure", default="result/evaluation/comparison.png")
+    args = parser.parse_args()
+    config = load_config(ROOT / args.config)
+    with np.load(ROOT / args.predictions) as data:
+        ensemble = data["ensemble"]
+        truth = data["truth"]
+        valid = data["validity"].astype(bool)
+    result = evaluate_ensemble(ensemble, truth, valid, config["evaluation"]["thresholds_dbz"])
+    metrics_path = ROOT / args.metrics
+    metrics_path.parent.mkdir(parents=True, exist_ok=True)
+    metrics_path.write_text(json.dumps(result, indent=2), encoding="utf-8")
+    prediction = ensemble.mean(axis=0)
+    fig, axes = plt.subplots(2, 5, figsize=(15, 6), constrained_layout=True)
+    for lead in range(5):
+        for row, field in enumerate((truth, prediction)):
+            composite = np.where(valid[lead], field[lead], np.nan).max(axis=0)
+            image = axes[row, lead].imshow(composite.T, origin="lower", vmin=0, vmax=60, cmap="turbo", aspect="auto")
+            axes[row, lead].set_title(f"{'Truth' if row == 0 else 'Ensemble mean'} +{(lead + 1) * 6} min")
+            axes[row, lead].set_xlabel("azimuth index")
+            if lead == 0:
+                axes[row, lead].set_ylabel("range bin")
+    fig.colorbar(image, ax=axes, label="composite reflectivity (dBZ)", shrink=0.8)
+    figure_path=ROOT/args.figure;figure_path.parent.mkdir(parents=True,exist_ok=True);fig.savefig(figure_path,dpi=120)
+    plt.close(fig)
+    print(json.dumps(result, indent=2))
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,77 @@
+"""DDP-capable EchoCast-3D training entry point."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler
+import sys
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+
+from model.echocast_3d import PackedRadarDataset, load_config, DiffusionSchedule, EchoCast3D, masked_objective
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="conf/config.yaml")
+    parser.add_argument("--data", default="data")
+    parser.add_argument("--output", default="result/checkpoints/echocast_3d.pt")
+    args = parser.parse_args()
+    config = load_config(ROOT / args.config)
+    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    if distributed:
+        dist.init_process_group("nccl" if torch.cuda.is_available() else "gloo")
+    use_cuda=torch.cuda.is_available() and torch.cuda.device_count()>=int(os.environ.get("WORLD_SIZE","1")) and config["runtime"]["device"]!="cpu"
+    device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu")
+    if device.type == "cuda":
+        torch.cuda.set_device(device)
+    torch.manual_seed(config["seed"] + local_rank)
+    dataset = PackedRadarDataset(ROOT / args.data, "train")
+    sampler = DistributedSampler(dataset, shuffle=True) if distributed else None
+    loader = DataLoader(dataset, batch_size=config["training"]["batch_size"], sampler=sampler, shuffle=sampler is None)
+    raw_model = EchoCast3D(config).to(device)
+    model = DistributedDataParallel(raw_model, device_ids=[local_rank]) if distributed else raw_model
+    optimizer = torch.optim.AdamW(
+        model.parameters(), lr=config["training"]["learning_rate"], weight_decay=config["training"]["weight_decay"]
+    )
+    schedule = DiffusionSchedule(device=device, **{k: config["diffusion"][k] for k in ("steps", "beta_start", "beta_end")})
+    for epoch in range(config["training"]["epochs"]):
+        if sampler:
+            sampler.set_epoch(epoch)
+        for batch in loader:
+            truth = batch["truth"].float().to(device)
+            values = batch["values"].float().to(device)
+            validity = batch["validity"].bool().to(device)
+            observed = batch["observed"].bool().to(device)
+            clean_target = truth[:, 3:]
+            timestep = torch.randint(schedule.steps, (truth.shape[0],), device=device)
+            noisy_target, _ = schedule.add_noise(clean_target, timestep)
+            token_mask = torch.rand(truth.shape[0], raw_model.token_count, device=device) < config["model"]["mask_ratio"]
+            prediction = model(values[:, :3], noisy_target, observed[:, :3], timestep, token_mask)
+            loss, parts = masked_objective(
+                raw_model, prediction, clean_target, validity, token_mask, config["training"]["reconstruction_weight"]
+            )
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            optimizer.step()
+            if local_rank == 0:
+                print(f"epoch={epoch + 1} loss={loss.item():.5f} score={parts['score']:.5f} recon={parts['reconstruction']:.5f}")
+    if local_rank == 0:
+        output = ROOT / args.output
+        output.parent.mkdir(parents=True, exist_ok=True)
+        torch.save({"model": raw_model.state_dict(), "model_config": config["model"], "format_version": config["data"]["format_version"], "config": config}, output)
+        metrics=ROOT/config["paths"]["training_metrics"];metrics.parent.mkdir(parents=True,exist_ok=True);metrics.write_text(json.dumps({"epochs":config["training"]["epochs"],"loss":float(loss),"world_size":int(os.environ.get("WORLD_SIZE","1"))},indent=2))
+        print(f"saved {output}; geometry tokens={raw_model.token_count}, paper states {config['model']['paper_token_count']}")
+    if distributed:
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()