training/README.md

# 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 released California 5 km, 12-hour-lead cache. |
| `configs/train_firewx_fm_seed7_template.json` | Template for the seeded California FireWx-FM training run. |
| `configs/stage1_cache_conus_hrrr_us_5km_l12_template.json` | Template for constructing a CONUS 5 km, 12-hour-lead cache. |
| `configs/train_firewx_fm_conus_seed7_template.json` | Template for CONUS retraining after local data preparation. |
| `NATIONWIDE_RETRAINING.md` | Scope note and commands for nationwide retraining. |

## 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.

The released checkpoints are California regional weights. For nationwide
training, use the CONUS templates and notes in `NATIONWIDE_RETRAINING.md`.

## 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.

For serving-oriented retraining, set `positive_tile_placement` to
`random_containing`. This samples positive tiles so the fire cell can appear at
different positions within the tile instead of always being centered. The
training script also supports `input_normalization`; use it to compute
train-split z-score statistics for continuous weather and static channels while
leaving masks and categorical fuel codes unscaled.

```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.

Avoid non-overlapping 32-by-32 serving tiles. If full-map inference is not
available, use overlap-tiled inference with halo cropping or overlap blending.
The helper `models/wildfire_fm/tiled_inference.py` provides
`predict_probability_tiled` for this stitching pattern.