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