Add FireWx-FM training and data loader pipeline

training/README.md

@@ -0,0 +1,96 @@
+# FireWx-FM Training and Data Loader Pipeline
+
+This folder contains the original FireWx-FM cache builder and PyTorch training
+pipeline used for the released regional wildfire occupancy checkpoints. The raw
+provider datasets and local cache tensors are not redistributed here. The code
+shows how those sources are converted into the model tensor contract and how the
+training loader samples tiles from that cache.
+
+## Files
+
+| Path | Purpose |
+|---|---|
+| `build_phase1_cache_regional_hrrr.py` | Builds the regional HRRR/FIRMS/static cache on the California EPSG:5070 grid. |
+| `train_cold_tiled_mainline.py` | Contains the `ColdFeatureStore`, tile sampler, full-map loader, compact U-Net, losses, and training loop. |
+| `eval_metrics.py` | Metric utilities used by the training script for validation and test summaries. |
+| `cache_helpers.py` | Minimal helper functions used by the regional cache builder. |
+| `train_utils.py` | Reproducibility helper used by the training script. |
+| `configs/stage1_cache_regional_hrrr_ca_5km_l12_template.json` | Template for constructing the 5 km, 12-hour-lead regional cache. |
+| `configs/train_firewx_fm_seed7_template.json` | Template for the seeded FireWx-FM training run. |
+
+## Data boundary
+
+The scripts expect local copies of the public source datasets described in
+`../data_sources/DATA_SOURCES.md`: NOAA HRRR fields, NASA FIRMS active-fire
+detections, LANDFIRE fuel and canopy layers, Wildfire Risk to Communities
+housing density, and LandScan population. WFIGS and MTBS support event-level
+experiments, but they are not input channels for the released occupancy
+checkpoint.
+
+## Cache construction
+
+The cache builder writes:
+
+| File group | Arrays |
+|---|---|
+| `inputs/phase1_regional_*.npz` | `weather`, `firewx`, `firewx_valid`, `y_count`, `y_occ`, `lat`, `lon`, `weather_names`, `firewx_names` |
+| `static/static_regional_phase1_v1.npz` | `static`, `static_valid`, `lat`, `lon`, `coord_crs`, `static_names` |
+| `splits/{train,val,test}.csv` | Sample metadata and paths used by the training loader. |
+| `manifests/phase1_cache_summary.json` | Cache shape, split counts, channel names, and grid metadata. |
+
+Run it with a local copy of the template config:
+
+```bash
+python training/build_phase1_cache_regional_hrrr.py \
+  --config training/configs/stage1_cache_regional_hrrr_ca_5km_l12_template.json
+```
+
+The released configuration uses a California grid at 5 km resolution in
+`EPSG:5070`, 6-hourly HRRR issue times, and 12-hour occupancy labels derived
+from FIRMS detections.
+
+## Tensor contract
+
+`ColdFeatureStore` is the key data-loader class. It assembles each model input
+as:
+
+```python
+x = np.concatenate([weather, firewx, *extra, static_valid, static], axis=0)
+```
+
+For the released regional cache, `firewx` has zero feature channels and
+`extra = [firewx_valid]`. The resulting 16-channel tensor is:
+
+| Channels | Names | Source |
+|---:|---|---|
+| 0-9 | `t2m`, `d2m`, `u10`, `v10`, `cape`, `sp`, `blh`, `vis`, `prate`, `tp` | NOAA HRRR |
+| 10 | `firewx_valid` | Dynamic/input validity mask for this regional cache |
+| 11 | `static_valid` | Static reprojection validity mask |
+| 12-15 | `fuel_fbfm40`, `canopy_cover`, `housing_density`, `population` | LANDFIRE, WRC housing density, LandScan |
+
+The same contract is available in machine-readable form at
+`../models/wildfire_fm/input_channels.json`.
+
+## Training
+
+The training script reads `splits/train.csv`, `splits/val.csv`, and
+`splits/test.csv` from `index_root`. Training uses 32-by-32 tiles sampled from
+the cached time maps. Validation and test use full maps.
+
+```bash
+python training/train_cold_tiled_mainline.py \
+  --config training/configs/train_firewx_fm_seed7_template.json \
+  --run-name firewx_fm_seed7
+```
+
+Released checkpoints are stored on the Hub under
+`models/wildfire_fm/checkpoints/seed_*/best_firms_prauc.pt`. The released seed
+configs and training summaries are under `models/wildfire_fm/configs/` and
+`models/wildfire_fm/metrics/`.
+
+## Inference adaptation
+
+For inference, the important part is to reproduce the same 16-channel tensor
+order. If a downstream environment already constructs `[channel, y, x]` tensors
+in that order, it can use `models/wildfire_fm/modeling_unet.py` directly and
+load one of the released seeded checkpoints.

training/build_phase1_cache_regional_hrrr.py

@@ -0,0 +1,484 @@
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import warnings
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List, Tuple
+
+import numpy as np
+import pandas as pd
+import rasterio
+import xarray as xr
+from rasterio.enums import Resampling
+from rasterio.transform import from_bounds
+from rasterio.warp import reproject, transform, transform_bounds
+from sklearn.neighbors import NearestNeighbors
+
+from cache_helpers import ensure_dirs, resampling_from_name, split_for_timestamp
+
+
+HRRR_RE = re.compile(r"hrrr\.(\d{8})\.t(\d{2})z\.wrfsfcf00\.grib2$")
+
+VAR_GROUPS = {
+    "level2": {
+        "filter_by_keys": {"typeOfLevel": "heightAboveGround", "level": 2},
+        "vars": ["t2m", "d2m", "sh2", "r2", "pt"],
+    },
+    "level10": {
+        "filter_by_keys": {"typeOfLevel": "heightAboveGround", "level": 10},
+        "vars": ["u10", "v10", "max_10si"],
+    },
+    "surface_instant": {
+        "filter_by_keys": {"typeOfLevel": "surface", "stepType": "instant"},
+        "vars": ["cape", "sp", "blh", "vis", "prate", "gust", "t"],
+    },
+    "surface_accum": {
+        "filter_by_keys": {"typeOfLevel": "surface", "stepType": "accum"},
+        "vars": ["tp", "ssrun", "bgrun"],
+    },
+}
+
+
+@dataclass
+class SampleRecord:
+    issue_ts: pd.Timestamp
+    target_ts: pd.Timestamp
+    hrrr_path: Path
+    label_path: Path
+
+
+@dataclass
+class RegionalContext:
+    config: Dict[str, object]
+    cache_root: Path
+    firms_dir: Path
+    hrrr_roots: List[Path]
+    static_layers: Dict[str, Dict[str, str]]
+    sample_dir: Path
+    static_dir: Path
+    manifest_dir: Path
+    split_dir: Path
+
+
+def load_json(path: Path) -> Dict[str, object]:
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def build_context(config: Dict[str, object]) -> RegionalContext:
+    cache_root = Path(config["cache_root"])
+    sample_dir = cache_root / "inputs"
+    static_dir = cache_root / "static"
+    manifest_dir = cache_root / "manifests"
+    split_dir = cache_root / "splits"
+    ensure_dirs([sample_dir, static_dir, manifest_dir, split_dir])
+    return RegionalContext(
+        config=config,
+        cache_root=cache_root,
+        firms_dir=Path(config["firms_dir"]),
+        hrrr_roots=[Path(p) for p in config["hrrr_roots"]],
+        static_layers=config["static_layers"],
+        sample_dir=sample_dir,
+        static_dir=static_dir,
+        manifest_dir=manifest_dir,
+        split_dir=split_dir,
+    )
+
+
+def build_projected_grid(bounds: Dict[str, float], target_crs: str, resolution_m: float) -> Tuple[np.ndarray, np.ndarray]:
+    west, south, east, north = transform_bounds(
+        "EPSG:4326",
+        target_crs,
+        float(bounds["lon_min"]),
+        float(bounds["lat_min"]),
+        float(bounds["lon_max"]),
+        float(bounds["lat_max"]),
+        densify_pts=21,
+    )
+    res = float(resolution_m)
+    x_asc = np.arange(west + 0.5 * res, east, res, dtype=np.float32)
+    y_asc = np.arange(south + 0.5 * res, north, res, dtype=np.float32)
+    if x_asc.size == 0 or y_asc.size == 0:
+        raise RuntimeError(f"Empty projected grid for resolution_m={resolution_m}")
+    return y_asc[::-1].copy(), x_asc.copy()
+
+
+def projected_transform(y_desc: np.ndarray, x_asc: np.ndarray) -> rasterio.Affine:
+    res_y = float(np.median(np.abs(np.diff(y_desc))))
+    res_x = float(np.median(np.diff(x_asc)))
+    west = float(x_asc[0] - 0.5 * res_x)
+    east = float(x_asc[-1] + 0.5 * res_x)
+    south = float(y_desc[-1] - 0.5 * res_y)
+    north = float(y_desc[0] + 0.5 * res_y)
+    return from_bounds(west, south, east, north, len(x_asc), len(y_desc))
+
+
+def projected_edges(values_asc: np.ndarray) -> np.ndarray:
+    step = float(np.median(np.diff(values_asc)))
+    edges = np.empty(values_asc.size + 1, dtype=np.float64)
+    edges[1:-1] = 0.5 * (values_asc[:-1] + values_asc[1:])
+    edges[0] = values_asc[0] - 0.5 * step
+    edges[-1] = values_asc[-1] + 0.5 * step
+    return edges
+
+
+def list_hrrr_issue_paths(roots: List[Path]) -> Dict[pd.Timestamp, Path]:
+    out: Dict[pd.Timestamp, Path] = {}
+    for root in roots:
+        if not root.exists():
+            continue
+        for path in root.rglob("*.grib2"):
+            m = HRRR_RE.match(path.name)
+            if not m:
+                continue
+            date_part, hour_part = m.groups()
+            ts = pd.to_datetime(f"{date_part}{hour_part}", format="%Y%m%d%H")
+            out.setdefault(ts, path)
+    return dict(sorted(out.items()))
+
+
+def build_sample_records(ctx: RegionalContext) -> List[SampleRecord]:
+    hours = {int(v) for v in ctx.config.get("issue_hours", [0, 6, 12, 18])}
+    lead_hours = int(ctx.config["target_offset_hours"])
+    hrrr_index = list_hrrr_issue_paths(ctx.hrrr_roots)
+    records: List[SampleRecord] = []
+    for item in ctx.config["input_ranges"]:
+        for day in pd.date_range(item["start"], item["end"], freq="D"):
+            for hour in sorted(hours):
+                issue_ts = pd.Timestamp(year=day.year, month=day.month, day=day.day, hour=hour)
+                hrrr_path = hrrr_index.get(issue_ts)
+                if hrrr_path is None:
+                    continue
+                target_ts = issue_ts + pd.Timedelta(hours=lead_hours)
+                label_path = ctx.firms_dir / f"{target_ts.strftime('%Y%m%d_%H')}.csv"
+                if not label_path.exists():
+                    continue
+                records.append(
+                    SampleRecord(
+                        issue_ts=issue_ts,
+                        target_ts=target_ts,
+                        hrrr_path=hrrr_path,
+                        label_path=label_path,
+                    )
+                )
+    if not records:
+        raise RuntimeError("No regional HRRR/FIRMS sample records found.")
+    return records
+
+
+def open_hrrr_dataset(path: Path, group_name: str) -> xr.Dataset:
+    warnings.filterwarnings("ignore", category=FutureWarning)
+    # Ignore shared cfgrib sidecar indexes under raw HRRR roots. Some existing
+    # `.idx` files are truncated/corrupted and can poison concurrent regional
+    # cache jobs. Using an empty indexpath forces a fresh in-process scan.
+    kwargs = {
+        "filter_by_keys": dict(VAR_GROUPS[group_name]["filter_by_keys"]),
+        "indexpath": "",
+    }
+    return xr.open_dataset(path, engine="cfgrib", backend_kwargs=kwargs)
+
+
+def build_source_resampler(sample_grib: Path, bounds: Dict[str, float], target_crs: str, y_desc: np.ndarray, x_asc: np.ndarray) -> Dict[str, np.ndarray]:
+    ds = open_hrrr_dataset(sample_grib, "level2")
+    lat2d = np.asarray(ds["latitude"].values, dtype=np.float64)
+    lon2d = np.asarray(ds["longitude"].values, dtype=np.float64)
+    ds.close()
+    lon2d = np.where(lon2d > 180.0, lon2d - 360.0, lon2d)
+    buffer_deg = float(bounds.get("buffer_deg", 1.0))
+    mask = (
+        (lat2d >= float(bounds["lat_min"]) - buffer_deg)
+        & (lat2d <= float(bounds["lat_max"]) + buffer_deg)
+        & (lon2d >= float(bounds["lon_min"]) - buffer_deg)
+        & (lon2d <= float(bounds["lon_max"]) + buffer_deg)
+    )
+    flat_idx = np.flatnonzero(mask.ravel())
+    if flat_idx.size == 0:
+        raise RuntimeError("No HRRR source points remain after California bbox crop.")
+    src_lat = lat2d.ravel()[flat_idx]
+    src_lon = lon2d.ravel()[flat_idx]
+    src_x, src_y = transform("EPSG:4326", target_crs, src_lon.tolist(), src_lat.tolist())
+    src_xy = np.column_stack([np.asarray(src_x, dtype=np.float32), np.asarray(src_y, dtype=np.float32)])
+
+    yy, xx = np.meshgrid(y_desc, x_asc, indexing="ij")
+    tgt_xy = np.column_stack([xx.ravel().astype(np.float32), yy.ravel().astype(np.float32)])
+    nn = NearestNeighbors(n_neighbors=1, algorithm="kd_tree")
+    nn.fit(src_xy)
+    _, gather = nn.kneighbors(tgt_xy, return_distance=True)
+    return {
+        "flat_idx": flat_idx.astype(np.int64),
+        "gather_idx": gather[:, 0].astype(np.int64),
+        "target_h": np.array([len(y_desc)], dtype=np.int64),
+        "target_w": np.array([len(x_asc)], dtype=np.int64),
+    }
+
+
+def resample_field(field: np.ndarray, resampler: Dict[str, np.ndarray]) -> np.ndarray:
+    flat = np.asarray(field, dtype=np.float32).ravel()[resampler["flat_idx"]]
+    gathered = flat[resampler["gather_idx"]]
+    h = int(resampler["target_h"][0])
+    w = int(resampler["target_w"][0])
+    return np.nan_to_num(gathered.reshape(h, w), nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32)
+
+
+def load_hrrr_cube(ctx: RegionalContext, records: List[SampleRecord], y_desc: np.ndarray, x_asc: np.ndarray) -> Tuple[np.ndarray, List[str]]:
+    weather_names = list(ctx.config["hrrr_vars"])
+    weather = np.zeros((len(records), len(weather_names), len(y_desc), len(x_asc)), dtype=np.float32)
+    resampler = build_source_resampler(records[0].hrrr_path, ctx.config, str(ctx.config["target_crs"]), y_desc, x_asc)
+
+    var_to_group = {}
+    for group_name, spec in VAR_GROUPS.items():
+        for name in spec["vars"]:
+            var_to_group[name] = group_name
+
+    for i, record in enumerate(records):
+        opened: Dict[str, xr.Dataset] = {}
+        try:
+            for j, name in enumerate(weather_names):
+                group = var_to_group.get(name)
+                if group is None:
+                    raise KeyError(f"No HRRR group mapping for variable '{name}'")
+                ds = opened.get(group)
+                if ds is None:
+                    ds = open_hrrr_dataset(record.hrrr_path, group)
+                    opened[group] = ds
+                if name not in ds.data_vars:
+                    raise KeyError(f"Variable '{name}' missing from {record.hrrr_path} in group '{group}'")
+                weather[i, j] = resample_field(np.asarray(ds[name].values, dtype=np.float32), resampler)
+        finally:
+            for ds in opened.values():
+                ds.close()
+        if (i + 1) % 20 == 0 or (i + 1) == len(records):
+            print(
+                json.dumps(
+                    {
+                        "stage": "load_hrrr_cube_progress",
+                        "loaded_samples": i + 1,
+                        "total_samples": len(records),
+                        "latest_issue_timestamp": record.issue_ts.isoformat(),
+                    }
+                ),
+                flush=True,
+            )
+    return weather, weather_names
+
+
+def load_static_cube_projected(ctx: RegionalContext, y_desc: np.ndarray, x_asc: np.ndarray) -> Tuple[np.ndarray, List[str], np.ndarray]:
+    transform_dst = projected_transform(y_desc, x_asc)
+    shape = (len(y_desc), len(x_asc))
+    res_y = float(np.median(np.abs(np.diff(y_desc))))
+    res_x = float(np.median(np.diff(x_asc)))
+    west = float(x_asc[0] - 0.5 * res_x)
+    east = float(x_asc[-1] + 0.5 * res_x)
+    south = float(y_desc[-1] - 0.5 * res_y)
+    north = float(y_desc[0] + 0.5 * res_y)
+    target_crs = str(ctx.config["target_crs"])
+
+    arrays: List[np.ndarray] = []
+    valid_arrays: List[np.ndarray] = []
+    names: List[str] = []
+    for name, spec in ctx.static_layers.items():
+        names.append(name)
+        dst = np.full(shape, np.nan, dtype=np.float32)
+        with rasterio.open(spec["path"]) as src:
+            src_bounds = transform_bounds(target_crs, src.crs, west, south, east, north, densify_pts=21)
+            window = src.window(*src_bounds).round_offsets().round_lengths()
+            source = src.read(1, window=window)
+            source_transform = src.window_transform(window)
+            reproject(
+                source=source,
+                destination=dst,
+                src_transform=source_transform,
+                src_crs=src.crs,
+                dst_transform=transform_dst,
+                dst_crs=target_crs,
+                src_nodata=src.nodata,
+                dst_nodata=np.nan,
+                resampling=resampling_from_name(spec["resampling"]),
+            )
+        valid = np.isfinite(dst) & (dst > -9000.0)
+        valid_arrays.append(valid.astype(np.float32))
+        dst = np.where(dst <= -9000.0, np.nan, dst)
+        dst = np.nan_to_num(dst, nan=0.0, posinf=0.0, neginf=0.0)
+        arrays.append(dst.astype(np.float32))
+    valid_fraction = np.stack(valid_arrays, axis=0).mean(axis=0, keepdims=True).astype(np.float32)
+    return np.stack(arrays, axis=0), names, valid_fraction
+
+
+def rasterize_firms_counts_projected(firms_path: Path, y_desc: np.ndarray, x_asc: np.ndarray, target_crs: str, bounds: Dict[str, float]) -> np.ndarray:
+    if not firms_path.exists():
+        return np.zeros((len(y_desc), len(x_asc)), dtype=np.float32)
+    try:
+        df = pd.read_csv(firms_path, usecols=["latitude", "longitude", "type"])
+    except ValueError:
+        df = pd.read_csv(firms_path)
+    if "type" in df.columns:
+        df = df[df["type"] == 0]
+    if df.empty:
+        return np.zeros((len(y_desc), len(x_asc)), dtype=np.float32)
+    df = df[
+        (df["latitude"] >= float(bounds["lat_min"]))
+        & (df["latitude"] <= float(bounds["lat_max"]))
+        & (df["longitude"] >= float(bounds["lon_min"]))
+        & (df["longitude"] <= float(bounds["lon_max"]))
+    ]
+    if df.empty:
+        return np.zeros((len(y_desc), len(x_asc)), dtype=np.float32)
+
+    xs, ys = transform(
+        "EPSG:4326",
+        target_crs,
+        df["longitude"].astype(float).tolist(),
+        df["latitude"].astype(float).tolist(),
+    )
+    y_asc = y_desc[::-1]
+    y_edges = projected_edges(y_asc)
+    x_edges = projected_edges(x_asc)
+    counts, _, _ = np.histogram2d(np.asarray(ys, dtype=np.float64), np.asarray(xs, dtype=np.float64), bins=[y_edges, x_edges])
+    return counts.astype(np.float32)[::-1, :]
+
+
+def write_cache(
+    ctx: RegionalContext,
+    records: List[SampleRecord],
+    weather: np.ndarray,
+    static: np.ndarray,
+    static_valid: np.ndarray,
+    y_desc: np.ndarray,
+    x_asc: np.ndarray,
+    weather_names: List[str],
+    static_names: List[str],
+) -> pd.DataFrame:
+    static_path = ctx.static_dir / "static_regional_phase1_v1.npz"
+    np.savez_compressed(
+        static_path,
+        static=static,
+        static_valid=static_valid,
+        lat=y_desc,
+        lon=x_asc,
+        coord_crs=np.array([str(ctx.config["target_crs"])], dtype=object),
+        static_names=np.array(static_names, dtype=object),
+    )
+
+    h = int(len(y_desc))
+    w = int(len(x_asc))
+    empty_firewx = np.zeros((0, h, w), dtype=np.float32)
+    firewx_valid = np.ones((1, h, w), dtype=np.float32)
+    rows: List[Dict[str, object]] = []
+    target_crs = str(ctx.config["target_crs"])
+
+    for i, record in enumerate(records):
+        y_count = rasterize_firms_counts_projected(record.label_path, y_desc, x_asc, target_crs, ctx.config)
+        y_occ = (y_count > 0).astype(np.float32)
+        sample_id = record.issue_ts.strftime("%Y%m%d_%H")
+        sample_path = ctx.sample_dir / f"phase1_regional_{sample_id}.npz"
+        np.savez_compressed(
+            sample_path,
+            weather=weather[i],
+            firewx=empty_firewx,
+            firewx_valid=firewx_valid,
+            y_count=y_count[None, ...],
+            y_occ=y_occ[None, ...],
+            lat=y_desc,
+            lon=x_asc,
+            coord_crs=np.array([target_crs], dtype=object),
+            weather_names=np.array(weather_names, dtype=object),
+            firewx_names=np.array([], dtype=object),
+        )
+        split = split_for_timestamp(record.issue_ts, ctx.config)
+        rows.append(
+            {
+                "sample_id": sample_id,
+                "input_date": str(record.issue_ts.date()),
+                "target_date": str(record.target_ts.date()),
+                "input_timestamp": record.issue_ts.isoformat(),
+                "target_timestamp": record.target_ts.isoformat(),
+                "split": split,
+                "sample_path": str(sample_path),
+                "static_path": str(static_path),
+                "pos_cells": int(y_occ.sum()),
+                "fire_points": float(y_count.sum()),
+            }
+        )
+        if (i + 1) % 20 == 0 or (i + 1) == len(records):
+            print(
+                json.dumps(
+                    {
+                        "stage": "write_cache_progress",
+                        "written_samples": i + 1,
+                        "total_samples": len(records),
+                        "latest_sample_id": sample_id,
+                    }
+                ),
+                flush=True,
+            )
+
+    manifest = pd.DataFrame(rows).sort_values("input_timestamp").reset_index(drop=True)
+    manifest.to_csv(ctx.manifest_dir / "phase1_manifest.csv", index=False)
+    for split in ["train", "val", "test"]:
+        manifest[manifest["split"] == split].copy().to_csv(ctx.split_dir / f"{split}.csv", index=False)
+
+    summary = {
+        "num_samples": int(len(manifest)),
+        "weather_shape": list(weather.shape),
+        "firewx_shape": [len(records), 0, h, w],
+        "firewx_valid_shape": [len(records), 1, h, w],
+        "static_shape": list(static.shape),
+        "static_valid_shape": list(static_valid.shape),
+        "splits": {k: int((manifest["split"] == k).sum()) for k in ["train", "val", "test"]},
+        "total_positive_cells": int(manifest["pos_cells"].sum()),
+        "total_fire_points": float(manifest["fire_points"].sum()),
+        "grid_height": h,
+        "grid_width": w,
+        "target_crs": str(ctx.config["target_crs"]),
+        "target_resolution_m": int(ctx.config["target_resolution_m"]),
+        "target_offset_hours": int(ctx.config["target_offset_hours"]),
+        "weather_names": weather_names,
+        "static_names": static_names,
+    }
+    (ctx.manifest_dir / "phase1_cache_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
+    return manifest
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=Path, required=True)
+    args = parser.parse_args()
+
+    config = load_json(args.config)
+    ctx = build_context(config)
+    records = build_sample_records(ctx)
+    y_desc, x_asc = build_projected_grid(ctx.config, str(ctx.config["target_crs"]), float(ctx.config["target_resolution_m"]))
+
+    print(
+        f"[stage] regional_samples={len(records)} first={records[0].issue_ts.isoformat()} last={records[-1].issue_ts.isoformat()}",
+        flush=True,
+    )
+    print(
+        f"[stage] grid_crs={ctx.config['target_crs']} resolution_m={ctx.config['target_resolution_m']} "
+        f"shape=({len(y_desc)}, {len(x_asc)})",
+        flush=True,
+    )
+    print("[stage] load_hrrr_cube", flush=True)
+    weather, weather_names = load_hrrr_cube(ctx, records, y_desc, x_asc)
+    print("[stage] load_static_cube", flush=True)
+    static, static_names, static_valid = load_static_cube_projected(ctx, y_desc, x_asc)
+    print("[stage] write_cache", flush=True)
+    manifest = write_cache(
+        ctx=ctx,
+        records=records,
+        weather=weather,
+        static=static,
+        static_valid=static_valid,
+        y_desc=y_desc,
+        x_asc=x_asc,
+        weather_names=weather_names,
+        static_names=static_names,
+    )
+    print("[stage] done", flush=True)
+    print(f"Built regional HRRR cache with {len(manifest)} samples.")
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

training/cache_helpers.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Dict, List
+
+import pandas as pd
+from rasterio.enums import Resampling
+
+
+def ensure_dirs(paths: List[Path]) -> None:
+    for path in paths:
+        path.mkdir(parents=True, exist_ok=True)
+
+
+def resampling_from_name(name: str) -> Resampling:
+    mapping = {
+        "nearest": Resampling.nearest,
+        "bilinear": Resampling.bilinear,
+        "average": Resampling.average,
+    }
+    return mapping[name]
+
+
+def split_for_timestamp(ts: pd.Timestamp, config: Dict[str, object]) -> str:
+    if "split_years" in config:
+        for split, years in config["split_years"].items():
+            if int(ts.year) in [int(v) for v in years]:
+                return split
+        raise RuntimeError(f"Year {ts.year} does not map to any split.")
+
+    split_months = config["split_months"]
+    for split, months in split_months.items():
+        if int(ts.month) in months:
+            return split
+    raise RuntimeError(f"Month {ts.month} does not map to any split.")

training/configs/stage1_cache_regional_hrrr_ca_5km_l12_template.json

@@ -0,0 +1,71 @@
+{
+  "cache_root": "/path/to/local/firewx_fm_cache/regional_hrrr_ca_5km_l12",
+  "firms_dir": "/path/to/firms_three_hourly_csv",
+  "hrrr_roots": [
+    "/path/to/hrrr/2024",
+    "/path/to/optional/hrrr_backfill"
+  ],
+  "input_ranges": [
+    {
+      "start": "2024-06-01",
+      "end": "2024-10-30"
+    }
+  ],
+  "issue_hours": [
+    0,
+    6,
+    12,
+    18
+  ],
+  "target_offset_hours": 12,
+  "lat_min": 32.0,
+  "lat_max": 42.5,
+  "lon_min": -125.0,
+  "lon_max": -114.0,
+  "buffer_deg": 1.0,
+  "target_crs": "EPSG:5070",
+  "target_resolution_m": 5000,
+  "hrrr_vars": [
+    "t2m",
+    "d2m",
+    "u10",
+    "v10",
+    "cape",
+    "sp",
+    "blh",
+    "vis",
+    "prate",
+    "tp"
+  ],
+  "static_layers": {
+    "fuel_fbfm40": {
+      "path": "/path/to/LANDFIRE_FBFM40_CONUS.tif",
+      "resampling": "nearest"
+    },
+    "canopy_cover": {
+      "path": "/path/to/LANDFIRE_CC_CONUS.tif",
+      "resampling": "nearest"
+    },
+    "housing_density": {
+      "path": "/path/to/WRC_HUDen_CONUS.tif",
+      "resampling": "bilinear"
+    },
+    "population": {
+      "path": "/path/to/landscan-global-2024.tif",
+      "resampling": "bilinear"
+    }
+  },
+  "split_months": {
+    "train": [
+      6,
+      7,
+      8
+    ],
+    "val": [
+      9
+    ],
+    "test": [
+      10
+    ]
+  }
+}

training/configs/train_firewx_fm_seed7_template.json

@@ -0,0 +1,62 @@
+{
+  "index_root": "/path/to/local/firewx_fm_cache/regional_hrrr_ca_5km_l12",
+  "run_root": "/path/to/local/firewx_fm_runs",
+  "tile_size": 32,
+  "max_positive_tiles_per_sample": 64,
+  "min_negative_tiles_per_sample": 4,
+  "negative_to_positive_ratio": 3.0,
+  "positive_rate_source": "tiles",
+  "pos_weight_source": "tiles",
+  "pos_weight_cap": 50.0,
+  "batch_size": 32,
+  "eval_batch_size": 2,
+  "epochs": 10,
+  "learning_rate": 0.0002,
+  "weight_decay": 0.0001,
+  "base_channels": 32,
+  "dropout": 0.1,
+  "norm_type": "group",
+  "norm_groups": 8,
+  "init_head_bias_from_positive_rate": true,
+  "augment_flip": true,
+  "threshold": 0.5,
+  "metric_thresholds": [
+    0.1,
+    0.2,
+    0.3,
+    0.5,
+    0.7,
+    0.9
+  ],
+  "topk_area_fractions": [
+    0.01,
+    0.05,
+    0.1
+  ],
+  "fss_radii": [
+    1,
+    2,
+    4,
+    8
+  ],
+  "reliability_bins": 10,
+  "amp": true,
+  "seed": 7,
+  "boundary_radii": [
+    1,
+    2,
+    4
+  ],
+  "coarsen_factors": [
+    2,
+    4,
+    8
+  ],
+  "use_aux_spatial_head": true,
+  "aux_spatial_radius": 2,
+  "aux_spatial_loss_weight": 0.5,
+  "train_target_mode": "dilate_max",
+  "train_target_radius": 2,
+  "loss_type": "bce",
+  "init_checkpoint": ""
+}

training/eval_metrics.py

@@ -0,0 +1,826 @@
+from __future__ import annotations
+
+import math
+from typing import Dict, Iterable, List
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from sklearn.metrics import average_precision_score, roc_auc_score
+
+try:
+    from scipy import ndimage as _scipy_ndimage
+except Exception:  # pragma: no cover - optional on login nodes, present in Slurm env.
+    _scipy_ndimage = None
+
+
+def _safe_div(num: float, den: float) -> float:
+    return float(num / den) if den else 0.0
+
+
+def threshold_metrics(prob: np.ndarray, target: np.ndarray, threshold: float) -> Dict[str, float]:
+    pred = prob >= threshold
+    pos = target > 0.5
+    tp = int(np.logical_and(pred, pos).sum())
+    fp = int(np.logical_and(pred, ~pos).sum())
+    fn = int(np.logical_and(~pred, pos).sum())
+    precision = _safe_div(tp, tp + fp)
+    recall = _safe_div(tp, tp + fn)
+    far = _safe_div(fp, tp + fp)
+    csi = _safe_div(tp, tp + fp + fn)
+    f1 = _safe_div(2.0 * precision * recall, precision + recall)
+    beta_sq = 4.0
+    f2 = _safe_div((1.0 + beta_sq) * precision * recall, beta_sq * precision + recall)
+    freq_bias = _safe_div(tp + fp, tp + fn)
+    return {
+        "threshold": float(threshold),
+        "tp": tp,
+        "fp": fp,
+        "fn": fn,
+        "precision": precision,
+        "recall": recall,
+        "far": far,
+        "csi": csi,
+        "f1": f1,
+        "f2": f2,
+        "frequency_bias": freq_bias,
+        "predicted_positive_rate": float(pred.mean()),
+    }
+
+
+def log_score(prob: np.ndarray, target: np.ndarray, eps: float = 1e-6) -> float:
+    prob_clip = np.clip(np.asarray(prob, dtype=np.float32), eps, 1.0 - eps)
+    target_arr = np.asarray(target, dtype=np.float32)
+    loss = -(target_arr * np.log(prob_clip) + (1.0 - target_arr) * np.log(1.0 - prob_clip))
+    return float(np.mean(loss))
+
+
+def _spatial_dilate(mask: torch.Tensor, radius: int) -> torch.Tensor:
+    if radius <= 0:
+        return mask > 0.5
+    pooled = F.max_pool2d(
+        mask.float().unsqueeze(1),
+        kernel_size=radius * 2 + 1,
+        stride=1,
+        padding=radius,
+    ).squeeze(1)
+    return pooled > 0.5
+
+
+def _spatial_erode(mask: torch.Tensor, radius: int) -> torch.Tensor:
+    if radius <= 0:
+        return mask > 0.5
+    inv = 1.0 - (mask > 0.5).float()
+    pooled = F.max_pool2d(
+        inv.unsqueeze(1),
+        kernel_size=radius * 2 + 1,
+        stride=1,
+        padding=radius,
+    ).squeeze(1)
+    return pooled < 0.5
+
+
+def _boundary_mask(mask: torch.Tensor, width: int = 1) -> torch.Tensor:
+    mask_bool = mask > 0.5
+    if width <= 0:
+        return mask_bool
+    eroded = _spatial_erode(mask_bool.float(), width)
+    return mask_bool & ~eroded
+
+
+def _sample_times_to_hours(sample_times: np.ndarray) -> np.ndarray:
+    times = np.asarray(sample_times)
+    if np.issubdtype(times.dtype, np.datetime64):
+        return times.astype("datetime64[h]").astype(np.int64)
+    return times.astype(np.int64)
+
+
+def _build_tolerance_support(
+    pred_bool: torch.Tensor,
+    target_bool: torch.Tensor,
+    sample_times: np.ndarray,
+    temporal_tolerance_steps: int,
+    spatial_tolerance_radius: int,
+    time_step_hours: int,
+) -> tuple[torch.Tensor, torch.Tensor, int]:
+    times_h = _sample_times_to_hours(sample_times)
+    target_support = torch.zeros_like(target_bool, dtype=torch.bool)
+    pred_support = torch.zeros_like(pred_bool, dtype=torch.bool)
+    tolerance_hours = int(temporal_tolerance_steps) * int(time_step_hours)
+
+    for idx, current_time in enumerate(times_h):
+        window = np.abs(times_h - current_time) <= tolerance_hours
+        target_union = target_bool[window].float().amax(dim=0, keepdim=True)
+        pred_union = pred_bool[window].float().amax(dim=0, keepdim=True)
+        target_support[idx] = _spatial_dilate(target_union, spatial_tolerance_radius)[0]
+        pred_support[idx] = _spatial_dilate(pred_union, spatial_tolerance_radius)[0]
+    return target_support, pred_support, tolerance_hours
+
+
+def tolerant_threshold_metrics(
+    prob_maps: np.ndarray,
+    target_maps: np.ndarray,
+    sample_times: np.ndarray,
+    threshold: float,
+    temporal_tolerance_steps: int,
+    spatial_tolerance_radius: int,
+    region_mask: np.ndarray | None = None,
+    time_step_hours: int = 24,
+) -> Dict[str, float]:
+    pred = torch.from_numpy((prob_maps >= threshold).astype(np.float32))
+    target = torch.from_numpy((target_maps > 0.5).astype(np.float32))
+    pred_bool = pred > 0.5
+    target_bool = target > 0.5
+    target_support, pred_support, tolerance_hours = _build_tolerance_support(
+        pred_bool=pred_bool,
+        target_bool=target_bool,
+        sample_times=sample_times,
+        temporal_tolerance_steps=temporal_tolerance_steps,
+        spatial_tolerance_radius=spatial_tolerance_radius,
+        time_step_hours=time_step_hours,
+    )
+
+    matched_pred = pred_bool & target_support
+    matched_target = target_bool & pred_support
+
+    if region_mask is not None:
+        region = torch.from_numpy(_region_mask_to_bool(region_mask, prob_maps.shape[1:])).to(dtype=torch.bool)
+        pred_bool = pred_bool & region.unsqueeze(0)
+        target_bool = target_bool & region.unsqueeze(0)
+        matched_pred = matched_pred & region.unsqueeze(0)
+        matched_target = matched_target & region.unsqueeze(0)
+        total_cells = int(region.sum().item()) * int(pred_bool.shape[0])
+    else:
+        total_cells = int(pred_bool.numel())
+
+    pred_total = int(pred_bool.sum().item())
+    target_total = int(target_bool.sum().item())
+    matched_pred_total = int(matched_pred.sum().item())
+    matched_target_total = int(matched_target.sum().item())
+    precision = _safe_div(matched_pred_total, pred_total)
+    recall = _safe_div(matched_target_total, target_total)
+    f1 = _safe_div(2.0 * precision * recall, precision + recall)
+    return {
+        "threshold": float(threshold),
+        "temporal_tolerance_steps": int(temporal_tolerance_steps),
+        "temporal_tolerance_hours": int(tolerance_hours),
+        "time_step_hours": int(time_step_hours),
+        "spatial_tolerance_radius": int(spatial_tolerance_radius),
+        "predicted_positive_cells": pred_total,
+        "target_positive_cells": target_total,
+        "matched_predicted_cells": matched_pred_total,
+        "matched_target_cells": matched_target_total,
+        "precision": precision,
+        "recall": recall,
+        "f1": f1,
+        "predicted_positive_rate": _safe_div(pred_total, total_cells),
+    }
+
+
+def neighborhood_contingency(
+    prob_maps: np.ndarray,
+    target_maps: np.ndarray,
+    sample_times: np.ndarray,
+    threshold: float,
+    temporal_tolerance_steps: int,
+    spatial_tolerance_radius: int,
+    region_mask: np.ndarray | None = None,
+    time_step_hours: int = 24,
+) -> Dict[str, float]:
+    pred_bool = torch.from_numpy((prob_maps >= threshold).astype(np.float32)) > 0.5
+    target_bool = torch.from_numpy((target_maps > 0.5).astype(np.float32)) > 0.5
+    target_support, pred_support, tolerance_hours = _build_tolerance_support(
+        pred_bool=pred_bool,
+        target_bool=target_bool,
+        sample_times=sample_times,
+        temporal_tolerance_steps=temporal_tolerance_steps,
+        spatial_tolerance_radius=spatial_tolerance_radius,
+        time_step_hours=time_step_hours,
+    )
+
+    if region_mask is not None:
+        region = torch.from_numpy(_region_mask_to_bool(region_mask, prob_maps.shape[1:])).to(dtype=torch.bool)
+        pred_bool = pred_bool & region.unsqueeze(0)
+        target_bool = target_bool & region.unsqueeze(0)
+        target_support = target_support & region.unsqueeze(0)
+        pred_support = pred_support & region.unsqueeze(0)
+
+    hits = 0
+    false_alarms = 0
+    misses = 0
+    true_negatives = 0
+    for idx in range(int(pred_bool.shape[0])):
+        pred_event = bool(pred_bool[idx].any().item())
+        target_event = bool(target_bool[idx].any().item())
+        pred_match = bool((pred_bool[idx] & target_support[idx]).any().item()) if pred_event else False
+        target_match = bool((target_bool[idx] & pred_support[idx]).any().item()) if target_event else False
+        if target_event and target_match:
+            hits += 1
+        elif target_event:
+            misses += 1
+        elif not target_event and not pred_event:
+            true_negatives += 1
+        if pred_event and not pred_match:
+            false_alarms += 1
+
+    precision = _safe_div(hits, hits + false_alarms)
+    recall = _safe_div(hits, hits + misses)
+    f1 = _safe_div(2.0 * precision * recall, precision + recall)
+    far = _safe_div(false_alarms, hits + false_alarms)
+    csi = _safe_div(hits, hits + false_alarms + misses)
+    return {
+        "threshold": float(threshold),
+        "temporal_tolerance_steps": int(temporal_tolerance_steps),
+        "temporal_tolerance_hours": int(tolerance_hours),
+        "time_step_hours": int(time_step_hours),
+        "spatial_tolerance_radius": int(spatial_tolerance_radius),
+        "hits": int(hits),
+        "false_alarms": int(false_alarms),
+        "misses": int(misses),
+        "true_negatives": int(true_negatives),
+        "precision": precision,
+        "recall": recall,
+        "f1": f1,
+        "far": far,
+        "csi": csi,
+        "predicted_event_rate": _safe_div(hits + false_alarms, pred_bool.shape[0]),
+        "target_event_rate": _safe_div(hits + misses, pred_bool.shape[0]),
+    }
+
+
+def neighborhood_contingency_metrics(
+    prob_maps: np.ndarray,
+    target_maps: np.ndarray,
+    sample_times: np.ndarray,
+    thresholds: Iterable[float],
+    temporal_tolerances_steps: Iterable[int],
+    spatial_tolerances_radii: Iterable[int],
+    time_step_hours: int = 24,
+) -> Dict[str, Dict[str, Dict[str, float]]]:
+    out: Dict[str, Dict[str, Dict[str, float]]] = {}
+    temporal_values = [int(v) for v in temporal_tolerances_steps]
+    spatial_values = [int(v) for v in spatial_tolerances_radii]
+    for temporal_steps in temporal_values:
+        for spatial_radius in spatial_values:
+            if temporal_steps == 0 and spatial_radius == 0:
+                continue
+            combo_key = f"t{temporal_steps}_s{spatial_radius}"
+            out[combo_key] = {
+                f"{float(t):.4f}": neighborhood_contingency(
+                    prob_maps=prob_maps,
+                    target_maps=target_maps,
+                    sample_times=sample_times,
+                    threshold=float(t),
+                    temporal_tolerance_steps=temporal_steps,
+                    spatial_tolerance_radius=spatial_radius,
+                    time_step_hours=time_step_hours,
+                )
+                for t in thresholds
+            }
+    return out
+
+
+def reliability_bins(prob: np.ndarray, target: np.ndarray, n_bins: int) -> List[Dict[str, float]]:
+    edges = np.linspace(0.0, 1.0, n_bins + 1)
+    rows: List[Dict[str, float]] = []
+    for idx in range(n_bins):
+        lo = edges[idx]
+        hi = edges[idx + 1]
+        if idx == n_bins - 1:
+            mask = (prob >= lo) & (prob <= hi)
+        else:
+            mask = (prob >= lo) & (prob < hi)
+        count = int(mask.sum())
+        if count == 0:
+            rows.append(
+                {
+                    "bin": idx,
+                    "lo": float(lo),
+                    "hi": float(hi),
+                    "count": 0,
+                    "mean_confidence": 0.0,
+                    "empirical_accuracy": 0.0,
+                }
+            )
+            continue
+        mean_conf = float(prob[mask].mean())
+        acc = float(target[mask].mean())
+        rows.append(
+            {
+                "bin": idx,
+                "lo": float(lo),
+                "hi": float(hi),
+                "count": count,
+                "mean_confidence": mean_conf,
+                "empirical_accuracy": acc,
+            }
+        )
+    return rows
+
+
+def expected_calibration_error(prob: np.ndarray, target: np.ndarray, n_bins: int) -> float:
+    bins = reliability_bins(prob, target, n_bins)
+    total = max(int(prob.size), 1)
+    return float(
+        sum(abs(row["empirical_accuracy"] - row["mean_confidence"]) * row["count"] for row in bins) / total
+    )
+
+
+def topk_area_metrics(prob: np.ndarray, target: np.ndarray, fractions: Iterable[float]) -> Dict[str, Dict[str, float]]:
+    order = np.argsort(prob)[::-1]
+    total = prob.size
+    positive_rate = float(target.mean())
+    pos_total = max(float(target.sum()), 1.0)
+    out: Dict[str, Dict[str, float]] = {}
+    for frac in fractions:
+        frac = float(frac)
+        keep = max(int(math.ceil(total * frac)), 1)
+        idx = order[:keep]
+        top_target = target[idx]
+        precision = float(top_target.mean())
+        recall = float(top_target.sum() / pos_total)
+        lift = float(precision / positive_rate) if positive_rate > 0 else 0.0
+        out[f"{frac:.4f}"] = {
+            "fraction": frac,
+            "cells": keep,
+            "precision": precision,
+            "recall": recall,
+            "lift": lift,
+        }
+    return out
+
+
+def _crop_for_factor(arr: np.ndarray, factor: int) -> np.ndarray:
+    if factor <= 1:
+        return arr
+    h = int(arr.shape[-2])
+    w = int(arr.shape[-1])
+    new_h = (h // factor) * factor
+    new_w = (w // factor) * factor
+    if new_h <= 0 or new_w <= 0:
+        raise ValueError(f"Cannot coarsen shape {(h, w)} by factor={factor}")
+    return arr[..., :new_h, :new_w]
+
+
+def _avg_pool_maps(arr: np.ndarray, factor: int) -> np.ndarray:
+    if factor <= 1:
+        return arr.astype(np.float32, copy=False)
+    cropped = _crop_for_factor(arr, factor).astype(np.float32, copy=False)
+    tensor = torch.from_numpy(cropped).unsqueeze(1)
+    pooled = F.avg_pool2d(tensor, kernel_size=factor, stride=factor)
+    return pooled.squeeze(1).cpu().numpy().astype(np.float32, copy=False)
+
+
+def _max_pool_binary_maps(arr: np.ndarray, factor: int) -> np.ndarray:
+    if factor <= 1:
+        return (arr > 0.5).astype(np.float32, copy=False)
+    cropped = _crop_for_factor(arr, factor).astype(np.float32, copy=False)
+    tensor = torch.from_numpy((cropped > 0.5).astype(np.float32)).unsqueeze(1)
+    pooled = F.max_pool2d(tensor, kernel_size=factor, stride=factor)
+    return (pooled.squeeze(1).cpu().numpy() > 0.5).astype(np.float32)
+
+
+def coarsened_metrics(
+    prob_maps: np.ndarray,
+    target_maps: np.ndarray,
+    thresholds: Iterable[float],
+    factors: Iterable[int],
+    reference_positive_rate: float | None = None,
+) -> Dict[str, Dict[str, object]]:
+    out: Dict[str, Dict[str, object]] = {}
+    target_binary = (target_maps > 0.5).astype(np.float32, copy=False)
+    for factor in factors:
+        factor = int(factor)
+        if factor <= 1:
+            continue
+        prob_coarse = _avg_pool_maps(prob_maps, factor)
+        target_fraction = _avg_pool_maps(target_binary, factor)
+        target_any = _max_pool_binary_maps(target_binary, factor)
+
+        prob_flat = prob_coarse.reshape(-1)
+        any_flat = target_any.reshape(-1)
+        fraction_flat = target_fraction.reshape(-1)
+
+        any_metrics: Dict[str, object] = {
+            "positive_rate": float(any_flat.mean()) if any_flat.size > 0 else 0.0,
+            "positive_cells": int(any_flat.sum()) if any_flat.size > 0 else 0,
+            "total_cells": int(any_flat.size),
+            "pr_auc": float(average_precision_score(any_flat, prob_flat)) if float(any_flat.sum()) > 0 else 0.0,
+            "brier": float(np.mean((prob_flat - any_flat) ** 2)) if any_flat.size > 0 else 0.0,
+            "log_score": log_score(prob_flat, any_flat) if any_flat.size > 0 else 0.0,
+            "threshold_metrics": {
+                f"{float(t):.4f}": threshold_metrics(prob_flat, any_flat, float(t)) for t in thresholds
+            }
+            if any_flat.size > 0
+            else {},
+        }
+        if any_flat.size > 0 and float(np.unique(any_flat).size) > 1:
+            any_metrics["auroc"] = float(roc_auc_score(any_flat, prob_flat))
+        else:
+            any_metrics["auroc"] = 0.0
+        ref_rate = reference_positive_rate if reference_positive_rate is not None else float(any_flat.mean())
+        brier_ref = float(np.mean((ref_rate - any_flat) ** 2)) if any_flat.size > 0 else 0.0
+        any_metrics["reference_positive_rate"] = float(ref_rate)
+        any_metrics["brier_skill_score"] = float(1.0 - any_metrics["brier"] / brier_ref) if brier_ref > 0 else 0.0
+
+        out[f"x{factor}"] = {
+            "grid_shape": {"lat": int(prob_coarse.shape[-2]), "lon": int(prob_coarse.shape[-1])},
+            "any": any_metrics,
+            "fraction": {
+                "mean_target_fraction": float(fraction_flat.mean()) if fraction_flat.size > 0 else 0.0,
+                "mean_predicted_probability": float(prob_flat.mean()) if prob_flat.size > 0 else 0.0,
+                "mae": float(np.mean(np.abs(prob_flat - fraction_flat))) if fraction_flat.size > 0 else 0.0,
+                "rmse": float(np.sqrt(np.mean((prob_flat - fraction_flat) ** 2))) if fraction_flat.size > 0 else 0.0,
+                "brier": float(np.mean((prob_flat - fraction_flat) ** 2)) if fraction_flat.size > 0 else 0.0,
+            },
+        }
+    return out
+
+
+def _region_mask_to_bool(mask: np.ndarray, sample_shape: tuple[int, ...]) -> np.ndarray:
+    region = np.asarray(mask).astype(bool)
+    if region.shape == sample_shape:
+        return region
+    raise ValueError(f"Region mask shape {region.shape} does not match sample shape {sample_shape}")
+
+
+def region_metric_bundle(
+    prob_maps: np.ndarray,
+    target_maps: np.ndarray,
+    mask: np.ndarray,
+    thresholds: Iterable[float],
+    topk_fractions: Iterable[float],
+    n_bins: int,
+    reference_positive_rate: float | None = None,
+    sample_times: np.ndarray | None = None,
+    temporal_tolerances_steps: Iterable[int] | None = None,
+    spatial_tolerances_radii: Iterable[int] | None = None,
+    time_step_hours: int = 24,
+) -> Dict[str, object]:
+    region = _region_mask_to_bool(mask, prob_maps.shape[1:])
+    prob = prob_maps[:, region].reshape(-1)
+    target = target_maps[:, region].reshape(-1)
+    positive_rate = float(target.mean()) if target.size > 0 else 0.0
+    metrics: Dict[str, object] = {
+        "mask_cells": int(region.sum()),
+        "mask_fraction": float(region.mean()),
+        "positive_rate": positive_rate,
+        "positive_cells": int(target.sum()) if target.size > 0 else 0,
+        "total_cells": int(target.size),
+        "pr_auc": float(average_precision_score(target, prob)) if float(target.sum()) > 0 else 0.0,
+        "brier": float(np.mean((prob - target) ** 2)) if target.size > 0 else 0.0,
+        "log_score": log_score(prob, target) if target.size > 0 else 0.0,
+        "ece": expected_calibration_error(prob, target, n_bins) if target.size > 0 else 0.0,
+        "reliability_bins": reliability_bins(prob, target, n_bins) if target.size > 0 else [],
+        "threshold_metrics": {f"{float(t):.4f}": threshold_metrics(prob, target, float(t)) for t in thresholds}
+        if target.size > 0
+        else {},
+        "topk_area_metrics": topk_area_metrics(prob, target, topk_fractions) if target.size > 0 else {},
+    }
+    if target.size > 0 and float(np.unique(target).size) > 1:
+        metrics["auroc"] = float(roc_auc_score(target, prob))
+    else:
+        metrics["auroc"] = 0.0
+    ref_rate = reference_positive_rate if reference_positive_rate is not None else positive_rate
+    brier_ref = float(np.mean((ref_rate - target) ** 2)) if target.size > 0 else 0.0
+    metrics["reference_positive_rate"] = float(ref_rate)
+    metrics["brier_skill_score"] = float(1.0 - metrics["brier"] / brier_ref) if brier_ref > 0 else 0.0
+    if sample_times is not None:
+        temporal_values = [int(v) for v in (temporal_tolerances_steps or [])]
+        spatial_values = [int(v) for v in (spatial_tolerances_radii or [])]
+        tolerant_metrics: Dict[str, Dict[str, Dict[str, float]]] = {}
+        for temporal_steps in temporal_values:
+            for spatial_radius in spatial_values:
+                if temporal_steps == 0 and spatial_radius == 0:
+                    continue
+                combo_key = f"t{temporal_steps}_s{spatial_radius}"
+                tolerant_metrics[combo_key] = {
+                    f"{float(t):.4f}": tolerant_threshold_metrics(
+                        prob_maps=prob_maps,
+                        target_maps=target_maps,
+                        sample_times=sample_times,
+                        threshold=float(t),
+                        temporal_tolerance_steps=temporal_steps,
+                        spatial_tolerance_radius=spatial_radius,
+                        region_mask=region,
+                        time_step_hours=time_step_hours,
+                    )
+                    for t in thresholds
+                }
+        if tolerant_metrics:
+            metrics["tolerant_threshold_metrics"] = tolerant_metrics
+    return metrics
+
+
+def fss_metrics(
+    prob_maps: np.ndarray,
+    target_maps: np.ndarray,
+    thresholds: Iterable[float],
+    radii: Iterable[int],
+) -> Dict[str, Dict[str, float]]:
+    prob_t = torch.from_numpy(prob_maps.astype(np.float32))
+    tgt_t = torch.from_numpy((target_maps > 0.5).astype(np.float32))
+    out: Dict[str, Dict[str, float]] = {}
+    for threshold in thresholds:
+        pred = (prob_t >= float(threshold)).float()
+        row: Dict[str, float] = {}
+        for radius in radii:
+            radius = int(radius)
+            kernel = radius * 2 + 1
+            frac_pred = F.avg_pool2d(pred.unsqueeze(1), kernel_size=kernel, stride=1, padding=radius).squeeze(1)
+            frac_tgt = F.avg_pool2d(tgt_t.unsqueeze(1), kernel_size=kernel, stride=1, padding=radius).squeeze(1)
+            mse = torch.mean((frac_pred - frac_tgt) ** 2).item()
+            ref = torch.mean(frac_pred**2 + frac_tgt**2).item()
+            score = 1.0 - (mse / ref) if ref > 0 else 1.0
+            row[str(radius)] = float(score)
+        out[f"{float(threshold):.4f}"] = row
+    return out
+
+
+def boundary_metrics(
+    prob_maps: np.ndarray,
+    target_maps: np.ndarray,
+    thresholds: Iterable[float],
+    radii: Iterable[int],
+    boundary_width: int = 1,
+) -> Dict[str, Dict[str, Dict[str, float]]]:
+    prob_t = torch.from_numpy(prob_maps.astype(np.float32))
+    tgt_boundary = _boundary_mask(torch.from_numpy((target_maps > 0.5).astype(np.float32)), width=boundary_width)
+    out: Dict[str, Dict[str, Dict[str, float]]] = {}
+    for threshold in thresholds:
+        pred_boundary = _boundary_mask((prob_t >= float(threshold)).float(), width=boundary_width)
+        row: Dict[str, Dict[str, float]] = {}
+        pred_boundary_cells = int(pred_boundary.sum().item())
+        tgt_boundary_cells = int(tgt_boundary.sum().item())
+        for radius in radii:
+            radius = int(radius)
+            pred_band = _spatial_dilate(pred_boundary.float(), radius)
+            tgt_band = _spatial_dilate(tgt_boundary.float(), radius)
+            band_intersection = int((pred_band & tgt_band).sum().item())
+            band_union = int((pred_band | tgt_band).sum().item())
+            pred_match = int((pred_boundary & tgt_band).sum().item())
+            tgt_match = int((tgt_boundary & pred_band).sum().item())
+            denom = pred_boundary_cells + tgt_boundary_cells
+            row[str(radius)] = {
+                "boundary_iou": _safe_div(band_intersection, band_union) if band_union > 0 else 1.0,
+                "surface_dice": _safe_div(pred_match + tgt_match, denom) if denom > 0 else 1.0,
+                "pred_boundary_cells": pred_boundary_cells,
+                "target_boundary_cells": tgt_boundary_cells,
+            }
+        out[f"{float(threshold):.4f}"] = row
+    return out
+
+
+def buffered_overlap_metrics(
+    prob_maps: np.ndarray,
+    target_maps: np.ndarray,
+    thresholds: Iterable[float],
+    radii: Iterable[int],
+) -> Dict[str, Dict[str, Dict[str, float]]]:
+    prob_t = torch.from_numpy(prob_maps.astype(np.float32))
+    target_bool = torch.from_numpy((target_maps > 0.5).astype(np.float32)) > 0.5
+    target_cells = int(target_bool.sum().item())
+    out: Dict[str, Dict[str, Dict[str, float]]] = {}
+    for threshold in thresholds:
+        pred_bool = prob_t >= float(threshold)
+        pred_cells = int(pred_bool.sum().item())
+        row: Dict[str, Dict[str, float]] = {}
+        for radius in radii:
+            radius = int(radius)
+            pred_dilated = _spatial_dilate(pred_bool.float(), radius)
+            target_dilated = _spatial_dilate(target_bool.float(), radius)
+            pred_match = int((pred_bool & target_dilated).sum().item())
+            target_match = int((target_bool & pred_dilated).sum().item())
+            buffered_precision = _safe_div(pred_match, pred_cells)
+            buffered_recall = _safe_div(target_match, target_cells)
+            buffered_f1 = _safe_div(
+                2.0 * buffered_precision * buffered_recall,
+                buffered_precision + buffered_recall,
+            )
+            pred_d_target_intersection = int((pred_dilated & target_bool).sum().item())
+            pred_d_target_union = int((pred_dilated | target_bool).sum().item())
+            pred_target_d_intersection = int((pred_bool & target_dilated).sum().item())
+            pred_target_d_union = int((pred_bool | target_dilated).sum().item())
+            sym_intersection = int((pred_dilated & target_dilated).sum().item())
+            sym_union = int((pred_dilated | target_dilated).sum().item())
+            row[str(radius)] = {
+                "predicted_positive_cells": pred_cells,
+                "target_positive_cells": target_cells,
+                "buffered_precision": buffered_precision,
+                "buffered_recall": buffered_recall,
+                "buffered_f1": buffered_f1,
+                "pred_dilated_iou": _safe_div(pred_d_target_intersection, pred_d_target_union),
+                "target_dilated_iou": _safe_div(pred_target_d_intersection, pred_target_d_union),
+                "symmetric_dilated_iou": _safe_div(sym_intersection, sym_union),
+            }
+        out[f"{float(threshold):.4f}"] = row
+    return out
+
+
+def distance_transform_metrics(
+    prob_maps: np.ndarray,
+    target_maps: np.ndarray,
+    thresholds: Iterable[float],
+    cutoffs: Iterable[int],
+) -> Dict[str, Dict[str, Dict[str, float]]]:
+    if _scipy_ndimage is None:
+        return {}
+    target_bool = target_maps > 0.5
+    cutoff_values = [int(v) for v in cutoffs if int(v) > 0]
+    if not cutoff_values:
+        return {}
+    out: Dict[str, Dict[str, Dict[str, float]]] = {}
+    for threshold in thresholds:
+        pred_bool = prob_maps >= float(threshold)
+        per_cutoff: Dict[int, Dict[str, object]] = {
+            cutoff: {
+                "pred_to_target": [],
+                "target_to_pred": [],
+                "baddeley_delta": [],
+                "empty_empty": 0,
+                "pred_empty_target_nonempty": 0,
+                "pred_nonempty_target_empty": 0,
+            }
+            for cutoff in cutoff_values
+        }
+        for idx in range(int(pred_bool.shape[0])):
+            pred_i = pred_bool[idx]
+            target_i = target_bool[idx]
+            pred_any = bool(pred_i.any())
+            target_any = bool(target_i.any())
+            if pred_any:
+                dt_pred = _scipy_ndimage.distance_transform_edt(~pred_i)
+            else:
+                dt_pred = None
+            if target_any:
+                dt_target = _scipy_ndimage.distance_transform_edt(~target_i)
+            else:
+                dt_target = None
+            for cutoff in cutoff_values:
+                bucket = per_cutoff[cutoff]
+                pred_to_target: List[float] = bucket["pred_to_target"]  # type: ignore[assignment]
+                target_to_pred: List[float] = bucket["target_to_pred"]  # type: ignore[assignment]
+                baddeley_delta: List[float] = bucket["baddeley_delta"]  # type: ignore[assignment]
+                if pred_any and target_any and dt_pred is not None and dt_target is not None:
+                    pred_to_target.extend(np.minimum(dt_target[pred_i], cutoff).astype(np.float32).tolist())
+                    target_to_pred.extend(np.minimum(dt_pred[target_i], cutoff).astype(np.float32).tolist())
+                    pred_dt_clip = np.minimum(dt_pred, cutoff)
+                    target_dt_clip = np.minimum(dt_target, cutoff)
+                    baddeley_delta.append(float(np.sqrt(np.mean((pred_dt_clip - target_dt_clip) ** 2))))
+                elif pred_any and not target_any:
+                    pred_to_target.extend([float(cutoff)] * int(pred_i.sum()))
+                    baddeley_delta.append(float(cutoff))
+                    bucket["pred_nonempty_target_empty"] = int(bucket["pred_nonempty_target_empty"]) + 1
+                elif target_any and not pred_any:
+                    target_to_pred.extend([float(cutoff)] * int(target_i.sum()))
+                    baddeley_delta.append(float(cutoff))
+                    bucket["pred_empty_target_nonempty"] = int(bucket["pred_empty_target_nonempty"]) + 1
+                else:
+                    baddeley_delta.append(0.0)
+                    bucket["empty_empty"] = int(bucket["empty_empty"]) + 1
+        row: Dict[str, Dict[str, float]] = {}
+        for cutoff in cutoff_values:
+            bucket = per_cutoff[cutoff]
+            pred_to_target_arr = np.asarray(bucket["pred_to_target"], dtype=np.float32)
+            target_to_pred_arr = np.asarray(bucket["target_to_pred"], dtype=np.float32)
+            if pred_to_target_arr.size and target_to_pred_arr.size:
+                symmetric = np.concatenate([pred_to_target_arr, target_to_pred_arr])
+            elif pred_to_target_arr.size:
+                symmetric = pred_to_target_arr
+            else:
+                symmetric = target_to_pred_arr
+            baddeley_arr = np.asarray(bucket["baddeley_delta"], dtype=np.float32)
+            row[str(cutoff)] = {
+                "distance_cutoff": float(cutoff),
+                "mean_pred_to_target_distance": float(pred_to_target_arr.mean()) if pred_to_target_arr.size else 0.0,
+                "mean_target_to_pred_distance": float(target_to_pred_arr.mean()) if target_to_pred_arr.size else 0.0,
+                "mean_symmetric_surface_distance": float(symmetric.mean()) if symmetric.size else 0.0,
+                "hausdorff95_distance": float(np.percentile(symmetric, 95)) if symmetric.size else 0.0,
+                "baddeley_delta_p2": float(baddeley_arr.mean()) if baddeley_arr.size else 0.0,
+                "empty_empty_samples": float(bucket["empty_empty"]),
+                "pred_empty_target_nonempty_samples": float(bucket["pred_empty_target_nonempty"]),
+                "pred_nonempty_target_empty_samples": float(bucket["pred_nonempty_target_empty"]),
+            }
+        out[f"{float(threshold):.4f}"] = row
+    return out
+
+
+def metric_bundle(
+    prob_maps: np.ndarray,
+    target_maps: np.ndarray,
+    thresholds: Iterable[float],
+    topk_fractions: Iterable[float],
+    fss_radii: Iterable[int],
+    n_bins: int,
+    boundary_radii: Iterable[int] | None = None,
+    coarsen_factors: Iterable[int] | None = None,
+    distance_cutoffs: Iterable[int] | None = None,
+    reference_positive_rate: float | None = None,
+    sample_times: np.ndarray | None = None,
+    temporal_tolerances_steps: Iterable[int] | None = None,
+    spatial_tolerances_radii: Iterable[int] | None = None,
+    region_masks: Dict[str, np.ndarray] | None = None,
+    time_step_hours: int = 24,
+) -> Dict[str, object]:
+    prob = prob_maps.reshape(-1)
+    target = target_maps.reshape(-1)
+    positive_rate = float(target.mean())
+    metrics: Dict[str, object] = {
+        "positive_rate": positive_rate,
+        "positive_cells": int(target.sum()),
+        "total_cells": int(target.size),
+        "pr_auc": float(average_precision_score(target, prob)) if float(target.sum()) > 0 else 0.0,
+        "brier": float(np.mean((prob - target) ** 2)),
+        "log_score": log_score(prob, target),
+        "ece": expected_calibration_error(prob, target, n_bins),
+        "reliability_bins": reliability_bins(prob, target, n_bins),
+        "threshold_metrics": {f"{float(t):.4f}": threshold_metrics(prob, target, float(t)) for t in thresholds},
+        "topk_area_metrics": topk_area_metrics(prob, target, topk_fractions),
+        "fss": fss_metrics(prob_maps, target_maps, thresholds, fss_radii),
+        "boundary_metrics": boundary_metrics(
+            prob_maps,
+            target_maps,
+            thresholds,
+            boundary_radii if boundary_radii is not None else [1, 2, 4],
+        ),
+        "buffered_overlap_metrics": buffered_overlap_metrics(
+            prob_maps,
+            target_maps,
+            thresholds,
+            boundary_radii if boundary_radii is not None else [1, 2, 4],
+        ),
+        "coarsened_metrics": coarsened_metrics(
+            prob_maps,
+            target_maps,
+            thresholds,
+            coarsen_factors if coarsen_factors is not None else [2, 4, 8],
+            reference_positive_rate=reference_positive_rate,
+        ),
+    }
+    if distance_cutoffs is not None:
+        metrics["distance_transform_metrics"] = distance_transform_metrics(
+            prob_maps,
+            target_maps,
+            thresholds,
+            distance_cutoffs,
+        )
+    if float(np.unique(target).size) > 1:
+        metrics["auroc"] = float(roc_auc_score(target, prob))
+    else:
+        metrics["auroc"] = 0.0
+    ref_rate = reference_positive_rate if reference_positive_rate is not None else positive_rate
+    brier_ref = float(np.mean((ref_rate - target) ** 2))
+    metrics["reference_positive_rate"] = float(ref_rate)
+    metrics["brier_skill_score"] = float(1.0 - metrics["brier"] / brier_ref) if brier_ref > 0 else 0.0
+    if sample_times is not None:
+        temporal_values = [int(v) for v in (temporal_tolerances_steps or [])]
+        spatial_values = [int(v) for v in (spatial_tolerances_radii or [])]
+        tolerant_metrics: Dict[str, Dict[str, Dict[str, float]]] = {}
+        for temporal_steps in temporal_values:
+            for spatial_radius in spatial_values:
+                if temporal_steps == 0 and spatial_radius == 0:
+                    continue
+                combo_key = f"t{temporal_steps}_s{spatial_radius}"
+                tolerant_metrics[combo_key] = {
+                    f"{float(t):.4f}": tolerant_threshold_metrics(
+                        prob_maps=prob_maps,
+                        target_maps=target_maps,
+                        sample_times=sample_times,
+                        threshold=float(t),
+                        temporal_tolerance_steps=temporal_steps,
+                        spatial_tolerance_radius=spatial_radius,
+                        time_step_hours=time_step_hours,
+                    )
+                    for t in thresholds
+                }
+        if tolerant_metrics:
+            metrics["tolerant_threshold_metrics"] = tolerant_metrics
+        neighborhood_metrics = neighborhood_contingency_metrics(
+            prob_maps=prob_maps,
+            target_maps=target_maps,
+            sample_times=sample_times,
+            thresholds=thresholds,
+            temporal_tolerances_steps=temporal_values,
+            spatial_tolerances_radii=spatial_values,
+            time_step_hours=time_step_hours,
+        )
+        if neighborhood_metrics:
+            metrics["neighborhood_contingency_metrics"] = neighborhood_metrics
+    if region_masks:
+        metrics["region_metrics"] = {
+            name: region_metric_bundle(
+                prob_maps=prob_maps,
+                target_maps=target_maps,
+                mask=mask,
+                thresholds=thresholds,
+                topk_fractions=topk_fractions,
+                n_bins=n_bins,
+                reference_positive_rate=reference_positive_rate,
+                sample_times=sample_times,
+                temporal_tolerances_steps=temporal_tolerances_steps,
+                spatial_tolerances_radii=spatial_tolerances_radii,
+                time_step_hours=time_step_hours,
+            )
+            for name, mask in region_masks.items()
+        }
+    return metrics

training/requirements-training.txt

@@ -0,0 +1,9 @@
+numpy
+pandas
+torch
+scikit-learn
+scipy
+rasterio
+xarray
+cfgrib
+eccodes

training/train_cold_tiled_mainline.py

@@ -0,0 +1,901 @@
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import math
+import random
+import time
+from pathlib import Path
+from typing import Dict, Iterable, List, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset
+
+from eval_metrics import metric_bundle
+from train_utils import set_seed
+
+
+def load_json(path: Path) -> Dict[str, object]:
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def read_rows(path: Path) -> List[Dict[str, str]]:
+    with path.open("r", encoding="utf-8", newline="") as fh:
+        return list(csv.DictReader(fh))
+
+
+def integral_image(mask: np.ndarray) -> np.ndarray:
+    return np.pad(mask.astype(np.int32), ((1, 0), (1, 0)), mode="constant").cumsum(0).cumsum(1)
+
+
+def rect_sum(ii: np.ndarray, top: int, left: int, height: int, width: int) -> int:
+    bottom = top + height
+    right = left + width
+    return int(ii[bottom, right] - ii[top, right] - ii[bottom, left] + ii[top, left])
+
+
+def centered_tile_origin(y: int, x: int, tile_h: int, tile_w: int, h: int, w: int) -> Tuple[int, int]:
+    top = min(max(y - tile_h // 2, 0), h - tile_h)
+    left = min(max(x - tile_w // 2, 0), w - tile_w)
+    return int(top), int(left)
+
+
+def positive_tile_origins(mask: np.ndarray, tile_h: int, tile_w: int, max_tiles: int, rng: random.Random) -> List[Tuple[int, int]]:
+    ys, xs = np.where(mask > 0.5)
+    if ys.size == 0:
+        return []
+    origins = list(
+        {
+            centered_tile_origin(int(y), int(x), tile_h, tile_w, mask.shape[0], mask.shape[1])
+            for y, x in zip(ys, xs)
+        }
+    )
+    rng.shuffle(origins)
+    if max_tiles > 0:
+        origins = origins[:max_tiles]
+    return origins
+
+
+def negative_tile_origins(
+    mask: np.ndarray,
+    tile_h: int,
+    tile_w: int,
+    desired: int,
+    rng: random.Random,
+    forbidden: Iterable[Tuple[int, int]] = (),
+) -> List[Tuple[int, int]]:
+    h, w = mask.shape
+    ii = integral_image(mask)
+    forbidden_set = set(forbidden)
+    out: List[Tuple[int, int]] = []
+    seen = set(forbidden_set)
+    max_top = max(h - tile_h, 0)
+    max_left = max(w - tile_w, 0)
+    max_attempts = max(2000, desired * 50)
+    attempts = 0
+    while len(out) < desired and attempts < max_attempts:
+        attempts += 1
+        top = rng.randint(0, max_top)
+        left = rng.randint(0, max_left)
+        origin = (top, left)
+        if origin in seen:
+            continue
+        if rect_sum(ii, top, left, tile_h, tile_w) != 0:
+            continue
+        seen.add(origin)
+        out.append(origin)
+    return out
+
+
+def make_norm(norm_type: str, num_channels: int, norm_groups: int) -> nn.Module:
+    if norm_type == "batch":
+        return nn.BatchNorm2d(num_channels)
+    if norm_type == "group":
+        groups = max(1, min(int(norm_groups), num_channels))
+        while num_channels % groups != 0 and groups > 1:
+            groups -= 1
+        return nn.GroupNorm(groups, num_channels)
+    if norm_type == "instance":
+        return nn.InstanceNorm2d(num_channels, affine=True)
+    if norm_type in {"none", "identity"}:
+        return nn.Identity()
+    raise ValueError(f"Unsupported norm_type: {norm_type}")
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_ch: int, out_ch: int, norm_type: str, norm_groups: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False),
+            make_norm(norm_type, out_ch, norm_groups),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False),
+            make_norm(norm_type, out_ch, norm_groups),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+class UNetSmallFlex(nn.Module):
+    def __init__(
+        self,
+        in_ch: int,
+        base: int = 32,
+        dropout: float = 0.1,
+        norm_type: str = "group",
+        norm_groups: int = 8,
+        prior_prob: float | None = None,
+        use_aux_spatial_head: bool = False,
+        aux_prior_prob: float | None = None,
+    ):
+        super().__init__()
+        self.enc1 = ConvBlock(in_ch, base, norm_type, norm_groups)
+        self.enc2 = ConvBlock(base, base * 2, norm_type, norm_groups)
+        self.enc3 = ConvBlock(base * 2, base * 4, norm_type, norm_groups)
+        self.enc4 = ConvBlock(base * 4, base * 8, norm_type, norm_groups)
+        self.pool = nn.MaxPool2d(2)
+        self.bottleneck = ConvBlock(base * 8, base * 16, norm_type, norm_groups)
+        self.up4 = nn.ConvTranspose2d(base * 16, base * 8, 2, stride=2)
+        self.dec4 = ConvBlock(base * 16, base * 8, norm_type, norm_groups)
+        self.up3 = nn.ConvTranspose2d(base * 8, base * 4, 2, stride=2)
+        self.dec3 = ConvBlock(base * 8, base * 4, norm_type, norm_groups)
+        self.up2 = nn.ConvTranspose2d(base * 4, base * 2, 2, stride=2)
+        self.dec2 = ConvBlock(base * 4, base * 2, norm_type, norm_groups)
+        self.up1 = nn.ConvTranspose2d(base * 2, base, 2, stride=2)
+        self.dec1 = ConvBlock(base * 2, base, norm_type, norm_groups)
+        self.drop = nn.Dropout2d(p=dropout)
+        self.head = nn.Conv2d(base, 1, kernel_size=1)
+        self.use_aux_spatial_head = bool(use_aux_spatial_head)
+        self.aux_head = nn.Conv2d(base, 1, kernel_size=1) if self.use_aux_spatial_head else None
+        if prior_prob is not None:
+            prior_prob = float(min(max(prior_prob, 1e-6), 1.0 - 1e-6))
+            nn.init.constant_(self.head.bias, math.log(prior_prob / (1.0 - prior_prob)))
+        if self.aux_head is not None and aux_prior_prob is not None:
+            aux_prior_prob = float(min(max(aux_prior_prob, 1e-6), 1.0 - 1e-6))
+            nn.init.constant_(self.aux_head.bias, math.log(aux_prior_prob / (1.0 - aux_prior_prob)))
+
+    @staticmethod
+    def _match_hw(x: torch.Tensor, ref: torch.Tensor) -> torch.Tensor:
+        diff_y = ref.size(2) - x.size(2)
+        diff_x = ref.size(3) - x.size(3)
+        if diff_y > 0 or diff_x > 0:
+            x = F.pad(x, [diff_x // 2, diff_x - diff_x // 2, diff_y // 2, diff_y - diff_y // 2])
+        if diff_y < 0:
+            y0 = (-diff_y) // 2
+            x = x[:, :, y0 : y0 + ref.size(2), :]
+        if diff_x < 0:
+            x0 = (-diff_x) // 2
+            x = x[:, :, :, x0 : x0 + ref.size(3)]
+        return x
+
+    def forward(self, x: torch.Tensor, return_aux: bool = False):
+        e1 = self.enc1(x)
+        e2 = self.enc2(self.pool(e1))
+        e3 = self.enc3(self.pool(e2))
+        e4 = self.enc4(self.pool(e3))
+        b = self.bottleneck(self.pool(e4))
+        d4 = self.dec4(torch.cat([self._match_hw(self.up4(b), e4), e4], dim=1))
+        d3 = self.dec3(torch.cat([self._match_hw(self.up3(d4), e3), e3], dim=1))
+        d2 = self.dec2(torch.cat([self._match_hw(self.up2(d3), e2), e2], dim=1))
+        d1 = self.dec1(torch.cat([self._match_hw(self.up1(d2), e1), e1], dim=1))
+        features = self.drop(d1)
+        logits = self.head(features)
+        if return_aux and self.aux_head is not None:
+            return logits, self.aux_head(features)
+        return logits
+
+
+class ColdFeatureStore:
+    def __init__(self, rows: List[Dict[str, str]]):
+        self.cache: Dict[str, Dict[str, np.ndarray]] = {}
+        static_path = Path(rows[0]["static_path"])
+        static_npz = np.load(static_path, allow_pickle=True)
+        static = self._sanitize(static_npz["static"].astype(np.float32))
+        static_valid = static_npz["static_valid"].astype(np.float32) if "static_valid" in static_npz else None
+        static_parts = []
+        if static_valid is not None:
+            static_parts.append(static_valid.astype(np.float32))
+        static_parts.append(static)
+        static_x = np.concatenate(static_parts, axis=0).astype(np.float32)
+        for row in rows:
+            sample = np.load(row["sample_path"], allow_pickle=True)
+            weather = self._sanitize(sample["weather"].astype(np.float32))
+            firewx = self._sanitize(sample["firewx"].astype(np.float32))
+            extra = []
+            if "firewx_valid" in sample:
+                extra.append(sample["firewx_valid"].astype(np.float32))
+            x = np.concatenate([weather, firewx, *extra, static_x], axis=0).astype(np.float32)
+            y = np.nan_to_num(sample["y_occ"].astype(np.float32), nan=0.0, posinf=0.0, neginf=0.0)
+            self.cache[str(row["sample_id"])] = {"x": x, "y": y}
+
+    @staticmethod
+    def _sanitize(x: np.ndarray) -> np.ndarray:
+        x = np.where(x <= -9000.0, np.nan, x)
+        x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)
+        return x.astype(np.float32, copy=False)
+
+    def get(self, sample_id: str) -> Dict[str, np.ndarray]:
+        return self.cache[sample_id]
+
+
+def pooled_positive_rate(rows: List[Dict[str, str]], store: ColdFeatureStore, radius: int) -> float:
+    if radius <= 0:
+        return float(full_map_stats(rows, store)["positive_rate"])
+    pos = 0.0
+    total = 0.0
+    for row in rows:
+        sample = store.get(str(row["sample_id"]))
+        y = torch.from_numpy(sample["y"].astype(np.float32)).unsqueeze(0)
+        pooled = F.max_pool2d(y, kernel_size=radius * 2 + 1, stride=1, padding=radius)
+        pos += float((pooled > 0.5).sum().item())
+        total += float(pooled.numel())
+    return float(pos / total) if total > 0 else 0.0
+
+
+def transform_spatial_aux_target(target: torch.Tensor, radius: int) -> torch.Tensor:
+    if radius <= 0:
+        return target
+    return F.max_pool2d(target, kernel_size=radius * 2 + 1, stride=1, padding=radius)
+
+
+def transform_train_target(target: torch.Tensor, config: Dict[str, object]) -> torch.Tensor:
+    mode = str(config.get("train_target_mode", "hard"))
+    radius = int(config.get("train_target_radius", 0))
+    if mode == "hard" or radius <= 0:
+        return target
+    kernel = radius * 2 + 1
+    if mode == "dilate_max":
+        return F.max_pool2d(target, kernel_size=kernel, stride=1, padding=radius)
+    if mode == "soft_pool":
+        return F.avg_pool2d(target, kernel_size=kernel, stride=1, padding=radius)
+    raise ValueError(f"Unsupported train_target_mode: {mode}")
+
+
+def maybe_unpack_logits(output):
+    if isinstance(output, tuple):
+        return output
+    return output, None
+
+
+class FocalBCEWithLogitsLoss(nn.Module):
+    def __init__(self, gamma: float = 2.0, alpha: float = 0.75):
+        super().__init__()
+        self.gamma = float(gamma)
+        self.alpha = float(alpha)
+
+    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        bce = F.binary_cross_entropy_with_logits(logits, targets, reduction="none")
+        probs = torch.sigmoid(logits)
+        pt = torch.where(targets > 0.5, probs, 1.0 - probs)
+        alpha_t = torch.where(
+            targets > 0.5,
+            torch.full_like(targets, self.alpha),
+            torch.full_like(targets, 1.0 - self.alpha),
+        )
+        loss = alpha_t * ((1.0 - pt) ** self.gamma) * bce
+        return loss.mean()
+
+
+class OHEMBCEWithLogitsLoss(nn.Module):
+    def __init__(self, pos_weight: float, neg_pos_ratio: float = 8.0, min_negatives: int = 64):
+        super().__init__()
+        self.pos_weight = float(pos_weight)
+        self.neg_pos_ratio = float(neg_pos_ratio)
+        self.min_negatives = int(min_negatives)
+
+    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        losses = F.binary_cross_entropy_with_logits(
+            logits,
+            targets,
+            reduction="none",
+            pos_weight=torch.tensor([self.pos_weight], device=logits.device, dtype=logits.dtype),
+        )
+        batch_losses: List[torch.Tensor] = []
+        batch_size = int(logits.shape[0])
+        for i in range(batch_size):
+            loss_i = losses[i].reshape(-1)
+            tgt_i = targets[i].reshape(-1) > 0.5
+            pos_loss = loss_i[tgt_i]
+            neg_loss = loss_i[~tgt_i]
+            num_pos = int(pos_loss.numel())
+            keep_neg = max(self.min_negatives, int(self.neg_pos_ratio * max(num_pos, 1)))
+            keep_neg = min(keep_neg, int(neg_loss.numel()))
+            if keep_neg > 0:
+                neg_loss = torch.topk(neg_loss, k=keep_neg, largest=True).values
+            else:
+                neg_loss = neg_loss[:0]
+            denom = max(num_pos + int(neg_loss.numel()), 1)
+            batch_losses.append((pos_loss.sum() + neg_loss.sum()) / denom)
+        return torch.stack(batch_losses).mean()
+
+
+class SoftFbetaBCEWithLogitsLoss(nn.Module):
+    def __init__(self, pos_weight: float, beta: float = 2.0, bce_weight: float = 0.4, smooth: float = 1.0):
+        super().__init__()
+        self.pos_weight = float(pos_weight)
+        self.beta = float(beta)
+        self.bce_weight = float(bce_weight)
+        self.smooth = float(smooth)
+
+    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        targets = targets.clamp(0.0, 1.0)
+        bce = F.binary_cross_entropy_with_logits(
+            logits,
+            targets,
+            reduction="mean",
+            pos_weight=torch.tensor([self.pos_weight], device=logits.device, dtype=logits.dtype),
+        )
+        probs = torch.sigmoid(logits.float())
+        targets_f = targets.float()
+        dims = tuple(range(1, probs.ndim))
+        tp = (probs * targets_f).sum(dim=dims)
+        fp = (probs * (1.0 - targets_f)).sum(dim=dims)
+        fn = ((1.0 - probs) * targets_f).sum(dim=dims)
+        beta_sq = self.beta * self.beta
+        fbeta = ((1.0 + beta_sq) * tp + self.smooth) / ((1.0 + beta_sq) * tp + beta_sq * fn + fp + self.smooth)
+        soft_loss = 1.0 - fbeta.mean()
+        return self.bce_weight * bce + (1.0 - self.bce_weight) * soft_loss
+
+
+class TverskyBCEWithLogitsLoss(nn.Module):
+    def __init__(
+        self,
+        pos_weight: float,
+        alpha: float = 0.3,
+        beta: float = 0.7,
+        bce_weight: float = 0.4,
+        smooth: float = 1.0,
+    ):
+        super().__init__()
+        self.pos_weight = float(pos_weight)
+        self.alpha = float(alpha)
+        self.beta = float(beta)
+        self.bce_weight = float(bce_weight)
+        self.smooth = float(smooth)
+
+    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        targets = targets.clamp(0.0, 1.0)
+        bce = F.binary_cross_entropy_with_logits(
+            logits,
+            targets,
+            reduction="mean",
+            pos_weight=torch.tensor([self.pos_weight], device=logits.device, dtype=logits.dtype),
+        )
+        probs = torch.sigmoid(logits.float())
+        targets_f = targets.float()
+        dims = tuple(range(1, probs.ndim))
+        tp = (probs * targets_f).sum(dim=dims)
+        fp = (probs * (1.0 - targets_f)).sum(dim=dims)
+        fn = ((1.0 - probs) * targets_f).sum(dim=dims)
+        tversky = (tp + self.smooth) / (tp + self.alpha * fp + self.beta * fn + self.smooth)
+        return self.bce_weight * bce + (1.0 - self.bce_weight) * (1.0 - tversky.mean())
+
+
+class FSSBCEWithLogitsLoss(nn.Module):
+    def __init__(
+        self,
+        pos_weight: float,
+        radii: Iterable[int] = (1, 2),
+        bce_weight: float = 0.5,
+        eps: float = 1e-6,
+    ):
+        super().__init__()
+        self.pos_weight = float(pos_weight)
+        self.radii = [int(v) for v in radii if int(v) > 0]
+        if not self.radii:
+            self.radii = [1]
+        self.bce_weight = float(bce_weight)
+        self.eps = float(eps)
+
+    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        targets = targets.clamp(0.0, 1.0)
+        bce = F.binary_cross_entropy_with_logits(
+            logits,
+            targets,
+            reduction="mean",
+            pos_weight=torch.tensor([self.pos_weight], device=logits.device, dtype=logits.dtype),
+        )
+        probs = torch.sigmoid(logits.float())
+        targets_f = targets.float()
+        fss_losses: List[torch.Tensor] = []
+        for radius in self.radii:
+            kernel = radius * 2 + 1
+            pred_frac = F.avg_pool2d(probs, kernel_size=kernel, stride=1, padding=radius)
+            target_frac = F.avg_pool2d(targets_f, kernel_size=kernel, stride=1, padding=radius)
+            mse = torch.mean((pred_frac - target_frac) ** 2)
+            reference = torch.mean(pred_frac**2 + target_frac**2)
+            fss = 1.0 - mse / (reference + self.eps)
+            fss_losses.append(1.0 - fss)
+        fss_loss = torch.stack(fss_losses).mean()
+        return self.bce_weight * bce + (1.0 - self.bce_weight) * fss_loss
+
+
+def build_loss(config: Dict[str, object], pos_weight: float, device: torch.device) -> nn.Module:
+    loss_type = str(config.get("loss_type", "bce"))
+    if loss_type == "bce":
+        return nn.BCEWithLogitsLoss(pos_weight=torch.tensor([pos_weight], dtype=torch.float32, device=device))
+    if loss_type == "focal_bce":
+        return FocalBCEWithLogitsLoss(
+            gamma=float(config.get("focal_gamma", 2.0)),
+            alpha=float(config.get("focal_alpha", 0.75)),
+        )
+    if loss_type == "ohem_bce":
+        return OHEMBCEWithLogitsLoss(
+            pos_weight=pos_weight,
+            neg_pos_ratio=float(config.get("ohem_neg_pos_ratio", 8.0)),
+            min_negatives=int(config.get("ohem_min_negatives", 64)),
+        )
+    if loss_type == "soft_fbeta_bce":
+        return SoftFbetaBCEWithLogitsLoss(
+            pos_weight=pos_weight,
+            beta=float(config.get("soft_fbeta_beta", 2.0)),
+            bce_weight=float(config.get("soft_metric_bce_weight", 0.4)),
+            smooth=float(config.get("soft_metric_smooth", 1.0)),
+        )
+    if loss_type == "tversky_bce":
+        return TverskyBCEWithLogitsLoss(
+            pos_weight=pos_weight,
+            alpha=float(config.get("tversky_alpha", 0.3)),
+            beta=float(config.get("tversky_beta", 0.7)),
+            bce_weight=float(config.get("soft_metric_bce_weight", 0.4)),
+            smooth=float(config.get("soft_metric_smooth", 1.0)),
+        )
+    if loss_type == "fss_bce":
+        return FSSBCEWithLogitsLoss(
+            pos_weight=pos_weight,
+            radii=[int(v) for v in config.get("fss_loss_radii", [1, 2])],
+            bce_weight=float(config.get("soft_metric_bce_weight", 0.5)),
+        )
+    raise ValueError(f"Unsupported loss_type: {loss_type}")
+
+
+def full_map_stats(rows: List[Dict[str, str]], store: ColdFeatureStore) -> Dict[str, float]:
+    pos = 0.0
+    total = 0.0
+    for row in rows:
+        sample = store.get(str(row["sample_id"]))
+        y = sample["y"]
+        pos += float((y > 0.5).sum())
+        total += float(y.size)
+    neg = max(total - pos, 0.0)
+    raw_pos_weight = float(max(neg / max(pos, 1.0), 1.0)) if total > 0 else 1.0
+    return {
+        "source": "full_map",
+        "positive_cells": pos,
+        "total_cells": total,
+        "positive_rate": float(pos / total) if total > 0 else 0.0,
+        "raw_pos_weight": raw_pos_weight,
+    }
+
+
+def tile_stats(tile_rows: List[Dict[str, object]], store: ColdFeatureStore) -> Dict[str, float]:
+    pos = 0.0
+    total = 0.0
+    for row in tile_rows:
+        sample = store.get(str(row["sample_id"]))
+        top = int(row["top"])
+        left = int(row["left"])
+        tile_size = int(row["tile_size"])
+        target = sample["y"][0, top : top + tile_size, left : left + tile_size]
+        pos += float((target > 0.5).sum())
+        total += float(target.size)
+    neg = max(total - pos, 0.0)
+    raw_pos_weight = float(max(neg / max(pos, 1.0), 1.0)) if total > 0 else 1.0
+    return {
+        "source": "tiles",
+        "positive_cells": pos,
+        "total_cells": total,
+        "positive_rate": float(pos / total) if total > 0 else 0.0,
+        "raw_pos_weight": raw_pos_weight,
+    }
+
+
+def select_stats(source: str, train_rows: List[Dict[str, str]], tile_rows: List[Dict[str, object]], store: ColdFeatureStore) -> Dict[str, float]:
+    if source == "full_map":
+        return full_map_stats(train_rows, store)
+    if source == "tiles":
+        return tile_stats(tile_rows, store)
+    raise ValueError(f"Unsupported stats source: {source}")
+
+
+def build_train_tiles(
+    train_rows: List[Dict[str, str]],
+    store: ColdFeatureStore,
+    tile_size: int,
+    max_positive_tiles_per_sample: int,
+    min_negative_tiles_per_sample: int,
+    neg_pos_ratio: float,
+    rng: random.Random,
+) -> List[Dict[str, object]]:
+    tile_rows: List[Dict[str, object]] = []
+    for row in train_rows:
+        sample_id = str(row["sample_id"])
+        mask = store.get(sample_id)["y"][0]
+        pos_origins = positive_tile_origins(
+            mask=mask,
+            tile_h=tile_size,
+            tile_w=tile_size,
+            max_tiles=max_positive_tiles_per_sample,
+            rng=rng,
+        )
+        neg_count = max(min_negative_tiles_per_sample, int(math.ceil(len(pos_origins) * neg_pos_ratio)))
+        neg_origins = negative_tile_origins(
+            mask=mask,
+            tile_h=tile_size,
+            tile_w=tile_size,
+            desired=neg_count,
+            rng=rng,
+            forbidden=pos_origins,
+        )
+        for top, left in pos_origins:
+            tile_rows.append(
+                {
+                    "sample_id": sample_id,
+                    "tile_type": "positive",
+                    "top": int(top),
+                    "left": int(left),
+                    "tile_size": int(tile_size),
+                }
+            )
+        for top, left in neg_origins:
+            tile_rows.append(
+                {
+                    "sample_id": sample_id,
+                    "tile_type": "negative",
+                    "top": int(top),
+                    "left": int(left),
+                    "tile_size": int(tile_size),
+                }
+            )
+    return tile_rows
+
+
+class TrainTileDataset(Dataset):
+    def __init__(self, tile_rows: List[Dict[str, object]], store: ColdFeatureStore, augment_flip: bool, seed: int):
+        self.rows = tile_rows
+        self.store = store
+        self.augment_flip = bool(augment_flip)
+        self.rng = random.Random(seed)
+
+    def __len__(self) -> int:
+        return len(self.rows)
+
+    def __getitem__(self, idx: int):
+        row = self.rows[idx]
+        sample = self.store.get(str(row["sample_id"]))
+        top = int(row["top"])
+        left = int(row["left"])
+        tile_size = int(row["tile_size"])
+        x = sample["x"][:, top : top + tile_size, left : left + tile_size]
+        y = sample["y"][:, top : top + tile_size, left : left + tile_size]
+        if self.augment_flip:
+            if self.rng.random() < 0.5:
+                x = x[:, :, ::-1].copy()
+                y = y[:, :, ::-1].copy()
+            if self.rng.random() < 0.5:
+                x = x[:, ::-1, :].copy()
+                y = y[:, ::-1, :].copy()
+        return {"x": torch.from_numpy(x), "y": torch.from_numpy(y)}
+
+
+class FullMapDataset(Dataset):
+    def __init__(self, rows: List[Dict[str, str]], store: ColdFeatureStore):
+        self.rows = rows
+        self.store = store
+
+    def __len__(self) -> int:
+        return len(self.rows)
+
+    def __getitem__(self, idx: int):
+        row = self.rows[idx]
+        sample = self.store.get(str(row["sample_id"]))
+        return {"x": torch.from_numpy(sample["x"]), "y": torch.from_numpy(sample["y"]), "sample_id": str(row["sample_id"])}
+
+
+def evaluate(
+    model: nn.Module,
+    loader: DataLoader,
+    device: torch.device,
+    thresholds: List[float],
+    topk_area_fractions: List[float],
+    fss_radii: List[int],
+    n_bins: int,
+    reference_positive_rate: float,
+    criterion: nn.Module,
+    amp: bool,
+) -> Dict[str, object]:
+    model.eval()
+    total_loss = 0.0
+    total_items = 0
+    all_prob_maps: List[np.ndarray] = []
+    all_target_maps: List[np.ndarray] = []
+    with torch.no_grad():
+        for batch in loader:
+            x = torch.nan_to_num(batch["x"], nan=0.0, posinf=0.0, neginf=0.0).to(device, non_blocking=True)
+            y = torch.nan_to_num(batch["y"], nan=0.0, posinf=0.0, neginf=0.0).to(device, non_blocking=True)
+            with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=amp and device.type == "cuda"):
+                logits, _ = maybe_unpack_logits(model(x))
+                loss = criterion(logits, y)
+            prob = torch.sigmoid(logits.float()).detach().cpu().numpy()[:, 0, :, :]
+            target = y.float().detach().cpu().numpy()[:, 0, :, :]
+            all_prob_maps.append(prob)
+            all_target_maps.append(target)
+            total_loss += float(loss.item()) * x.size(0)
+            total_items += int(x.size(0))
+    prob_maps = np.concatenate(all_prob_maps, axis=0)
+    target_maps = np.concatenate(all_target_maps, axis=0)
+    metrics = metric_bundle(
+        prob_maps=prob_maps,
+        target_maps=target_maps,
+        thresholds=thresholds,
+        topk_fractions=topk_area_fractions,
+        fss_radii=fss_radii,
+        n_bins=n_bins,
+        reference_positive_rate=reference_positive_rate,
+    )
+    metrics.update({"loss": total_loss / max(total_items, 1), "num_samples": int(prob_maps.shape[0])})
+    return metrics
+
+
+def save_checkpoint(path: Path, model: nn.Module, optimizer: torch.optim.Optimizer, epoch: int, metrics: Dict[str, float], config: Dict[str, object]) -> None:
+    torch.save(
+        {
+            "epoch": epoch,
+            "model": model.state_dict(),
+            "optimizer": optimizer.state_dict(),
+            "metrics": metrics,
+            "config": config,
+        },
+        path,
+    )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=Path, required=True)
+    parser.add_argument("--run-name", type=str, required=True)
+    args = parser.parse_args()
+
+    config = load_json(args.config)
+    set_seed(int(config.get("seed", 7)))
+    rng = random.Random(int(config.get("seed", 7)))
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+    index_root = Path(config["index_root"])
+    run_root = Path(config["run_root"])
+    ckpt_dir = run_root / "checkpoints" / args.run_name
+    metric_dir = run_root / "metrics" / args.run_name
+    ckpt_dir.mkdir(parents=True, exist_ok=True)
+    metric_dir.mkdir(parents=True, exist_ok=True)
+
+    train_rows = read_rows(index_root / "splits" / "train.csv")
+    val_rows = read_rows(index_root / "splits" / "val.csv")
+    test_rows = read_rows(index_root / "splits" / "test.csv")
+    store = ColdFeatureStore(train_rows + val_rows + test_rows)
+
+    tile_rows = build_train_tiles(
+        train_rows=train_rows,
+        store=store,
+        tile_size=int(config.get("tile_size", 16)),
+        max_positive_tiles_per_sample=int(config.get("max_positive_tiles_per_sample", 64)),
+        min_negative_tiles_per_sample=int(config.get("min_negative_tiles_per_sample", 4)),
+        neg_pos_ratio=float(config.get("negative_to_positive_ratio", 2.0)),
+        rng=rng,
+    )
+    train_ds = TrainTileDataset(tile_rows=tile_rows, store=store, augment_flip=bool(config.get("augment_flip", True)), seed=int(config.get("seed", 7)))
+    val_ds = FullMapDataset(val_rows, store)
+    test_ds = FullMapDataset(test_rows, store)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    train_loader = DataLoader(train_ds, batch_size=int(config.get("batch_size", 64)), shuffle=True, num_workers=0, pin_memory=device.type == "cuda")
+    val_loader = DataLoader(val_ds, batch_size=int(config.get("eval_batch_size", 8)), shuffle=False, num_workers=0, pin_memory=device.type == "cuda")
+    test_loader = DataLoader(test_ds, batch_size=int(config.get("eval_batch_size", 8)), shuffle=False, num_workers=0, pin_memory=device.type == "cuda")
+
+    positive_rate_source = str(config.get("positive_rate_source", "full_map"))
+    pos_weight_source = str(config.get("pos_weight_source", "full_map"))
+    metric_reference_positive_rate = full_map_stats(train_rows, store)["positive_rate"]
+    positive_rate_stats = select_stats(positive_rate_source, train_rows, tile_rows, store)
+    pos_weight_stats = select_stats(pos_weight_source, train_rows, tile_rows, store)
+    train_positive_rate = float(positive_rate_stats["positive_rate"])
+    raw_pos_weight = float(pos_weight_stats["raw_pos_weight"])
+    pos_weight_cap = float(config.get("pos_weight_cap", 300.0))
+    pos_weight = float(min(max(raw_pos_weight, 1.0), pos_weight_cap))
+    use_aux_spatial_head = bool(config.get("use_aux_spatial_head", False))
+    aux_spatial_radius = int(config.get("aux_spatial_radius", 1))
+    aux_positive_rate = pooled_positive_rate(train_rows, store, aux_spatial_radius) if use_aux_spatial_head else None
+
+    in_ch = int(store.get(str(train_rows[0]["sample_id"]))["x"].shape[0])
+    model = UNetSmallFlex(
+        in_ch=in_ch,
+        base=int(config.get("base_channels", 32)),
+        dropout=float(config.get("dropout", 0.1)),
+        norm_type=str(config.get("norm_type", "group")),
+        norm_groups=int(config.get("norm_groups", 8)),
+        prior_prob=train_positive_rate if bool(config.get("init_head_bias_from_positive_rate", True)) else None,
+        use_aux_spatial_head=use_aux_spatial_head,
+        aux_prior_prob=aux_positive_rate if bool(config.get("init_head_bias_from_positive_rate", True)) else None,
+    ).to(device)
+    init_checkpoint = config.get("init_checkpoint")
+    if init_checkpoint:
+        checkpoint = torch.load(str(init_checkpoint), map_location="cpu")
+        state = checkpoint.get("model", checkpoint)
+        if not isinstance(state, dict):
+            raise RuntimeError(f"Unexpected init_checkpoint format: {init_checkpoint}")
+        model.load_state_dict(state, strict=False)
+    criterion = build_loss(config=config, pos_weight=pos_weight, device=device)
+    aux_criterion = nn.BCEWithLogitsLoss()
+    optimizer = torch.optim.AdamW(model.parameters(), lr=float(config.get("learning_rate", 3e-4)), weight_decay=float(config.get("weight_decay", 1e-4)))
+    scaler = torch.amp.GradScaler("cuda", enabled=bool(config.get("amp", True)) and device.type == "cuda")
+
+    thresholds = [float(v) for v in config.get("metric_thresholds", [0.1, 0.2, 0.3, 0.5])]
+    topk_area_fractions = [float(v) for v in config.get("topk_area_fractions", [0.01, 0.05, 0.1])]
+    fss_radii = [int(v) for v in config.get("fss_radii", [1, 2, 4, 8])]
+    n_bins = int(config.get("reliability_bins", 10))
+    default_threshold_key = f"{float(config.get('threshold', 0.3)):.4f}"
+
+    summary_seed = {
+        "index_root": str(index_root),
+        "tile_size": int(config.get("tile_size", 16)),
+        "num_train_tiles": len(tile_rows),
+        "positive_train_tiles": int(sum(1 for row in tile_rows if row["tile_type"] == "positive")),
+        "negative_train_tiles": int(sum(1 for row in tile_rows if row["tile_type"] == "negative")),
+        "input_channels": in_ch,
+        "train_positive_rate": train_positive_rate,
+        "metric_reference_positive_rate": metric_reference_positive_rate,
+        "pos_weight_source": pos_weight_source,
+        "positive_rate_source": positive_rate_source,
+        "pos_weight": pos_weight,
+        "loss_type": str(config.get("loss_type", "bce")),
+        "train_target_mode": str(config.get("train_target_mode", "hard")),
+        "train_target_radius": int(config.get("train_target_radius", 0)),
+    }
+    (metric_dir / "tile_summary.json").write_text(json.dumps(summary_seed, indent=2), encoding="utf-8")
+
+    best_pr = -1.0
+    best_state = None
+    history: List[Dict[str, float]] = []
+    aux_spatial_loss_weight = float(config.get("aux_spatial_loss_weight", 0.0))
+    for epoch in range(1, int(config.get("epochs", 30)) + 1):
+        model.train()
+        epoch_start = time.time()
+        train_loss = 0.0
+        train_main_loss = 0.0
+        train_aux_spatial_loss = 0.0
+        train_items = 0
+        for batch in train_loader:
+            x = torch.nan_to_num(batch["x"], nan=0.0, posinf=0.0, neginf=0.0).to(device, non_blocking=True)
+            y = torch.nan_to_num(batch["y"], nan=0.0, posinf=0.0, neginf=0.0).to(device, non_blocking=True)
+            train_target = transform_train_target(y.float(), config)
+            optimizer.zero_grad(set_to_none=True)
+            with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=bool(config.get("amp", True)) and device.type == "cuda"):
+                logits, aux_logits = maybe_unpack_logits(model(x, return_aux=use_aux_spatial_head))
+                main_loss = criterion(logits.float(), train_target.float())
+            aux_loss = torch.zeros((), device=device, dtype=torch.float32)
+            if aux_logits is not None and aux_spatial_loss_weight > 0.0:
+                aux_target = transform_spatial_aux_target(y.float(), aux_spatial_radius)
+                aux_loss = aux_criterion(aux_logits.float(), aux_target)
+            loss = main_loss + aux_spatial_loss_weight * aux_loss
+            if not torch.isfinite(loss):
+                raise RuntimeError("Non-finite tiled cold-start loss detected.")
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+            train_loss += float(loss.item()) * x.size(0)
+            train_main_loss += float(main_loss.item()) * x.size(0)
+            train_aux_spatial_loss += float(aux_loss.item()) * x.size(0)
+            train_items += int(x.size(0))
+
+        metrics = evaluate(
+            model=model,
+            loader=val_loader,
+            device=device,
+            thresholds=thresholds,
+            topk_area_fractions=topk_area_fractions,
+            fss_radii=fss_radii,
+            n_bins=n_bins,
+            reference_positive_rate=metric_reference_positive_rate,
+            criterion=criterion,
+            amp=bool(config.get("amp", True)),
+        )
+        threshold_metrics = metrics["threshold_metrics"][default_threshold_key]
+        record = {
+            "epoch": epoch,
+            "train_loss": train_loss / max(train_items, 1),
+            "train_main_loss": train_main_loss / max(train_items, 1),
+            "train_aux_spatial_loss": train_aux_spatial_loss / max(train_items, 1),
+            "val_loss": metrics["loss"],
+            "val_pr_auc": metrics["pr_auc"],
+            "val_auroc": metrics["auroc"],
+            "val_brier": metrics["brier"],
+            "val_brier_skill_score": metrics["brier_skill_score"],
+            "val_ece": metrics["ece"],
+            "val_positive_rate": metrics["positive_rate"],
+            "val_precision": threshold_metrics["precision"],
+            "val_positive_recall": threshold_metrics["recall"],
+            "val_far": threshold_metrics["far"],
+            "val_csi": threshold_metrics["csi"],
+            "val_f1": threshold_metrics["f1"],
+            "val_f2": threshold_metrics["f2"],
+            "val_frequency_bias": threshold_metrics["frequency_bias"],
+            "val_threshold_metrics": metrics["threshold_metrics"],
+            "val_topk_area_metrics": metrics["topk_area_metrics"],
+            "val_fss": metrics["fss"],
+            "minutes": (time.time() - epoch_start) / 60.0,
+        }
+        history.append(record)
+        (metric_dir / "history.json").write_text(json.dumps(history, indent=2), encoding="utf-8")
+        save_checkpoint(ckpt_dir / "latest.pt", model, optimizer, epoch, record, config)
+        if record["val_pr_auc"] > best_pr:
+            best_pr = record["val_pr_auc"]
+            best_state = {k: v.detach().cpu() for k, v in model.state_dict().items()}
+            torch.save({"model": best_state, "epoch": epoch, "metrics": record, "config": config}, ckpt_dir / "best_firms_prauc.pt")
+        print(json.dumps(record), flush=True)
+
+    if best_state is None:
+        best_state = {k: v.detach().cpu() for k, v in model.state_dict().items()}
+    model.load_state_dict(best_state, strict=True)
+    val_best = evaluate(
+        model=model,
+        loader=val_loader,
+        device=device,
+        thresholds=thresholds,
+        topk_area_fractions=topk_area_fractions,
+        fss_radii=fss_radii,
+        n_bins=n_bins,
+        reference_positive_rate=metric_reference_positive_rate,
+        criterion=criterion,
+        amp=bool(config.get("amp", True)),
+    )
+    test_best = evaluate(
+        model=model,
+        loader=test_loader,
+        device=device,
+        thresholds=thresholds,
+        topk_area_fractions=topk_area_fractions,
+        fss_radii=fss_radii,
+        n_bins=n_bins,
+        reference_positive_rate=metric_reference_positive_rate,
+        criterion=criterion,
+        amp=bool(config.get("amp", True)),
+    )
+    summary = {
+        "run_name": args.run_name,
+        "device": str(device),
+        "index_root": str(index_root),
+        "input_channels": in_ch,
+        "tile_size": int(config.get("tile_size", 16)),
+        "num_train_tiles": len(tile_rows),
+        "positive_train_tiles": int(sum(1 for row in tile_rows if row["tile_type"] == "positive")),
+        "negative_train_tiles": int(sum(1 for row in tile_rows if row["tile_type"] == "negative")),
+        "train_positive_rate": train_positive_rate,
+        "metric_reference_positive_rate": metric_reference_positive_rate,
+        "pos_weight": pos_weight,
+        "use_aux_spatial_head": use_aux_spatial_head,
+        "aux_spatial_radius": aux_spatial_radius,
+        "aux_spatial_loss_weight": aux_spatial_loss_weight,
+        "aux_positive_rate": aux_positive_rate,
+        "init_checkpoint": str(init_checkpoint) if init_checkpoint else "",
+        "best_val_pr_auc": float(val_best["pr_auc"]),
+        "best_val_auroc": float(val_best["auroc"]),
+        "best_test_pr_auc": float(test_best["pr_auc"]),
+        "best_test_auroc": float(test_best["auroc"]),
+        "best_test_brier": float(test_best["brier"]),
+        "best_test_ece": float(test_best["ece"]),
+        "best_test_topk_area_metrics": test_best["topk_area_metrics"],
+        "best_test_threshold_metrics": test_best["threshold_metrics"],
+        "best_test_fss": test_best["fss"],
+        "epochs": int(config.get("epochs", 30)),
+    }
+    (metric_dir / "run_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
+    print(json.dumps(summary, indent=2), flush=True)
+
+
+if __name__ == "__main__":
+    main()

training/train_utils.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+import random
+
+import numpy as np
+import torch
+
+
+def set_seed(seed: int) -> None:
+    """Set the RNG state used by the original FireWx-FM training script."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)