| """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): |
| |
| 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() |
| |
| |
| 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 |
|
|
