scripts/common.py

"""Data, provenance, and metrics shared by the independent reproduction."""
import argparse
import hashlib
import json
import os
import platform
import random
import sys
from pathlib import Path

import h5py
import numpy as np
import torch
import yaml
from torch.utils.data import Dataset

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model import FluxNetShallowWater

FIELDS = ("h", "mx", "my")


def load_config(path=None):
    with open(path or ROOT / "config/config.yaml", encoding="utf-8") as stream:
        cfg = yaml.safe_load(stream)
    for name, expected in {"epochs": 300, "pushforward_steps": 5}.items():
        if cfg["training"][name] != expected:
            raise ValueError(f"Formal experiment requires {name}={expected}")
    if cfg["training"]["early_stopping"] or cfg["data"]["normalization"] != "none":
        raise ValueError("Formal run requires 300 full epochs and physical states")
    return cfg


def checkpoint_config_matches(saved, current):
    # The published dataset can be relocated; all experiment settings must still match.
    saved = {**saved, "data": {**saved["data"], "root": None}}
    current = {**current, "data": {**current["data"], "root": None}}
    return saved == current


def setup(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.set_num_threads(min(8, int(os.environ.get("SLURM_CPUS_PER_TASK", 8))))
    torch.backends.cudnn.benchmark = True
    torch.backends.cuda.matmul.allow_tf32 = False
    torch.backends.cudnn.allow_tf32 = False


def device():
    if not os.environ.get("SLURM_JOB_ID"):
        raise RuntimeError("Run training/inference inside a SCNet DCU SLURM allocation")
    if not torch.cuda.is_available() or not torch.version.hip:
        raise RuntimeError("A HIP/DCU PyTorch device is required; CPU fallback is disabled")
    return torch.device("cuda")


def output_dir(cfg):
    path = ROOT / cfg["paths"]["results"]
    path.mkdir(parents=True, exist_ok=True)
    return path


def save_json(path, obj):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(obj, indent=2, ensure_ascii=False, allow_nan=False) + "\n")
    temporary.replace(path)


def sha256(path):
    digest = hashlib.sha256()
    with open(path, "rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def provenance():
    return {"hostname": platform.node(), "slurm_job_id": os.environ.get("SLURM_JOB_ID"),
            "torch": torch.__version__, "hip": torch.version.hip,
            "python": sys.version, "device": torch.cuda.get_device_name(),
            "device_count": torch.cuda.device_count(),
            "source_sha256": {str(p.relative_to(ROOT)): sha256(p)
                              for p in sorted(ROOT.rglob("*.py")) if "__pycache__" not in str(p)}}


class Trajectories:
    def __init__(self, cfg, split, validate=True):
        self.split = split
        self.files = sorted((Path(cfg["data"]["root"]) / split).glob("*.h5"))
        if len(self.files) != cfg["data"]["expected_counts"][split]:
            raise ValueError(f"Unexpected {split} trajectory count: {len(self.files)}")
        self.values = []
        self.metadata = []
        expected = (cfg["data"]["expected_frames"][split], *cfg["data"]["shape"])
        for path in self.files:
            with h5py.File(path, "r") as source:
                fields = []
                for field in FIELDS:
                    array = source[field][...]
                    if array.dtype != np.float32 or array.shape != expected:
                        raise ValueError(f"{path}:{field}: {array.dtype} {array.shape}")
                    if validate and not np.isfinite(array).all():
                        raise ValueError(f"Nonfinite data: {path}:{field}")
                    fields.append(array)
                state = np.stack(fields, axis=1)
                if validate and state[:, 0].min() < -cfg["data"]["reference_negative_tolerance"]:
                    raise ValueError(f"Substantive negative reference depth: {path}")
                if validate and state[0, 0].min() < 0:
                    raise ValueError(f"Negative initial reference depth: {path}")
                times = source["t"][...].astype(np.float64)
                if not np.allclose(np.diff(times), cfg["data"]["dt"], atol=2e-6):
                    raise ValueError(f"Unexpected snapshot interval: {path}")
                self.times = times
                self.metadata.append({"file": str(path), "bytes": path.stat().st_size,
                                      "frames": len(times), "t_start": float(times[0]),
                                      "t_end": float(times[-1]), "min_h": float(state[:, 0].min()),
                                      "reference_negative_h_count": int((state[:, 0] < 0).sum()),
                                      "min": state.min(axis=(0, 2, 3)).tolist(),
                                      "max": state.max(axis=(0, 2, 3)).tolist()})
                self.values.append(torch.from_numpy(state))
        self.values = torch.stack(self.values)
        reference_negatives = sum(m["reference_negative_h_count"] for m in self.metadata)
        print(f"DATA {split}: {tuple(self.values.shape)}; dt={cfg['data']['dt']} "
              f"t_end={self.times[-1]:.2f}; normalization=none; "
              f"raw_reference_negative_cells={reference_negatives} (roundoff retained, no clipping)", flush=True)


class Windows(Dataset):
    def __init__(self, trajectories, steps):
        self.trajectories = trajectories
        self.steps = steps
        self.starts = trajectories.values.shape[1] - steps

    def __len__(self):
        return self.trajectories.values.shape[0] * self.starts

    def __getitem__(self, index):
        trajectory, start = divmod(index, self.starts)
        return self.trajectories.values[trajectory, start:start + self.steps + 1]


def build_model(cfg, target_device):
    return FluxNetShallowWater(**cfg["model"]).to(target_device)


@torch.no_grad()
def validate(model, loader, target_device):
    model.eval()
    square = torch.zeros(3, device=target_device, dtype=torch.float64)
    absolute = torch.zeros_like(square)
    count = 0
    negative = 0
    terminal_square = torch.zeros_like(square)
    for window in loader:
        window = window.to(target_device, non_blocking=True)
        prediction = model(window[:, 0])
        error = prediction.double() - window[:, 1].double()
        square += error.square().sum(dim=(0, 2, 3))
        absolute += error.abs().sum(dim=(0, 2, 3))
        current = prediction
        for _ in range(1, window.shape[1] - 1):
            current = model(current)
        terminal_error = current.double() - window[:, -1].double()
        terminal_square += terminal_error.square().sum(dim=(0, 2, 3))
        count += prediction.shape[0] * prediction.shape[2] * prediction.shape[3]
        negative += int((prediction[:, 0] < 0).sum())
    return {"mse": float(((square + terminal_square) / count).mean()),
            "step1_mse": float((square / count).mean()),
            "step5_mse": float((terminal_square / count).mean()),
            "mae": (absolute / count).tolist(), "negative_h_ratio": negative / count}


@torch.no_grad()
def rollout_metrics(model, trajectories, target_device, batch_size=10,
                    save_path=None, example_cases=("A1", "A2", "B1", "B2")):
    model.eval()
    values = trajectories.values
    n, frames, channels, height, width = values.shape
    steps = frames - 1
    mae = np.empty((n, steps, channels), dtype=np.float64)
    cons_relative = np.empty_like(mae)
    cons_abs = np.empty_like(mae)
    cons_scaled = np.empty_like(mae)
    negatives = np.empty((n, steps), dtype=np.float64)
    minima = np.empty((n, steps), dtype=np.float64)
    examples = {}
    example_indices = {}
    for case in example_cases:
        matches = [i for i, p in enumerate(trajectories.files) if p.name.startswith("Case" + case + "_")]
        if matches:
            example_indices[matches[0]] = case
    for begin in range(0, n, batch_size):
        end = min(begin + batch_size, n)
        current = values[begin:end, 0].to(target_device)
        initial_sum = current.double().sum(dim=(-2, -1))
        initial_abs_sum = current.double().abs().sum(dim=(-2, -1))
        traces = []
        keep = any(begin <= idx < end for idx in example_indices)
        if keep:
            traces.append(current.cpu().numpy())
        for step in range(steps):
            current = model(current)
            if not torch.isfinite(current).all():
                raise FloatingPointError(f"Nonfinite {trajectories.split} rollout at step {step + 1}")
            target = values[begin:end, step + 1].to(target_device)
            mae[begin:end, step] = (current.double() - target.double()).abs().mean((-2, -1)).cpu().numpy()
            drift = (current.double().sum((-2, -1)) - initial_sum).abs()
            cons_abs[begin:end, step] = drift.cpu().numpy()
            # Paper's signed-total denominator is ill-conditioned for zero-net
            # momentum; preserve it and also report absolute and L1-scaled drift.
            cons_relative[begin:end, step] = (drift / (initial_sum.abs() + 1e-8)).cpu().numpy()
            cons_scaled[begin:end, step] = (drift / (initial_abs_sum + 1e-8)).cpu().numpy()
            negatives[begin:end, step] = (current[:, 0] < 0).double().mean((-2, -1)).cpu().numpy()
            minima[begin:end, step] = current[:, 0].amin((-2, -1)).cpu().numpy()
            if keep:
                traces.append(current.cpu().numpy())
        if keep:
            traces = np.stack(traces, axis=1)
            for idx, case in example_indices.items():
                if begin <= idx < end:
                    examples[case + "_prediction"] = traces[idx - begin]
                    examples[case + "_target"] = values[idx].numpy()
                    examples[case + "_file"] = np.array(trajectories.files[idx].name)
        print(f"ROLLOUT {trajectories.split} {end}/{n}: "
              f"MAE={mae[:end].mean((0, 1)).tolist()} "
              f"cons_abs_max={cons_abs[:end].max((0, 1)).tolist()} "
              f"negative_h_ratio={negatives[:end].mean():.6e}", flush=True)
    summary = {"split": trajectories.split, "trajectories": n, "steps": steps,
               "t_end": float(trajectories.times[-1]), "fields": list(FIELDS),
               "mae": mae.mean((0, 1)).tolist(),
               "mae_std_across_trajectories": mae.mean(1).std(0).tolist(),
               "mae_final": mae[:, -1].mean(0).tolist(),
               "conservation_relative_mean": cons_relative.mean((0, 1)).tolist(),
               "conservation_relative_max": cons_relative.max((0, 1)).tolist(),
               "conservation_absolute_sum_mean": cons_abs.mean((0, 1)).tolist(),
               "conservation_absolute_sum_max": cons_abs.max((0, 1)).tolist(),
               "conservation_initial_l1_scaled_max": cons_scaled.max((0, 1)).tolist(),
               "negative_h_ratio": float(negatives.mean()),
               "negative_h_percent": float(100 * negatives.mean()),
               "min_h": float(minima.min()),
               "reduction": "all trajectories, forecast steps 1..T, all grid cells; initial frame excluded",
               "accumulation_dtype": "float64 (model states float32)",
               "zero_momentum_note": "Relative drift uses |sum(m0)|+1e-8; inspect absolute drift and initial-L1-scaled drift for zero-total cases."}
    for horizon in (2.0, 2.4, 4.8):
        count = int(np.searchsorted(trajectories.times[1:], horizon + 1e-6, side="right"))
        if count and horizon <= trajectories.times[-1] + 1e-6:
            summary[f"mae_through_t{horizon:.1f}"] = mae[:, :count].mean((0, 1)).tolist()
    if save_path:
        save_path = Path(save_path)
        np.savez_compressed(save_path.with_suffix(".npz"), time=trajectories.times[1:],
                            mae=mae, conservation_relative=cons_relative,
                            conservation_absolute_sum=cons_abs,
                            conservation_l1_scaled=cons_scaled,
                            negative_h_ratio=negatives, min_h=minima,
                            files=np.array([p.name for p in trajectories.files]))
        np.savez_compressed(save_path.parent / (save_path.stem + "_examples.npz"),
                            time=trajectories.times, **examples)
        save_json(save_path.with_suffix(".json"), summary)
    return summary


def parser(description):
    arg = argparse.ArgumentParser(description=description)
    arg.add_argument("--config", type=Path, default=ROOT / "config/config.yaml")
    return arg