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:

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:

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.

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.