Graham McIntire
2fd88a94ea
Phase A — 7 pipeline bugs:
- retain_scores_window: fix timeline self-deletion (NOTIFY payload run_time|valid_time)
- Rust/Elixir weight divergence: Rust loads band_weights.json at startup
- Aurora boost ported to Rust (Kp query + per-cell boost for bands <= 432 MHz)
- Commercial-link boost applied in Rust per-cell scoring
- f00 native gradient preferred over pressure-level gradient
- HRDPS files: greedy regex fixed, now visible to timeline/prune/retain
- GEFS/HRRR collision: GEFS namespace as .gefs.prop, merge at read
Phase B — Validation harness:
- scripts/validate_algo.py: out-of-sample Spearman rho(score,distance) + baselines
- docs/algo-reports/validation-2026-08-01.{json,md}
Phase C — Forecast-skill evaluation:
- scripts/validate_forecast.py: skill degradation by lead time (0h-24h)
- docs/algo-reports/forecast-2026-08-01.{json,md}
Phase D — Calibration + model improvements:
- recalibrate.py: validation gate before weight deploy
- Recalibrator: Nx.max(0.0) replaces Nx.abs(), L2 regularization, val-set integrity
- ML train/serve defaults unified; prop_compare null-handling skew fixed
- Latitude-aware sunrise: solar-declination sunrise_hour(lat,month) (Elixir + Rust)
- Path scoring: wind/sky/rain from HRRR (was ~30% of composite weight silent)
- Region multiplier documented as unvalidated; PWAT/refractivity doc-vs-code noted
- algo.md: synced scoring sections, marked retired features, D5/D6 changelog
- Credo: path_compute cyclomatic complexity bumped (6->10 fields) — informational only
1577 lines
57 KiB
Rust
1577 lines
57 KiB
Rust
//! End-to-end f01..f48 chain step: fetch surface + pressure GRIBs, decode,
|
||
//! build per-cell `Conditions`, score every band, write score-grid files
|
||
//! (`<band>/<iso>.prop`) to the shared scores directory.
|
||
//!
|
||
//! This is the Rust equivalent of `PropagationGridWorker.process_forecast_hour/4`
|
||
//! minus the f00-only enrichment (native duct, NEXRAD, commercial link
|
||
//! boost, ProfilesFile write) — Elixir retains f00.
|
||
|
||
use std::path::Path;
|
||
|
||
use chrono::{DateTime, Datelike, Timelike, Utc};
|
||
|
||
use crate::band_config;
|
||
use crate::commercial::{self, LinkLookupEntry};
|
||
use crate::decoder;
|
||
use crate::fetcher::{self, HrrrClient, Product};
|
||
use crate::field_grid::FieldGrid;
|
||
use crate::grid::{hrdps_grid_spec, hrdps_only_points, wgrib2_grid_spec};
|
||
use crate::hrdps_fetcher::{self, HrdpsClient};
|
||
use crate::metrics;
|
||
use crate::native_duct;
|
||
use crate::nexrad::{self, NexradObservation};
|
||
use crate::pgrid;
|
||
use crate::planes::GridPlanes;
|
||
use crate::scorer::{self, Conditions};
|
||
use crate::scores_file;
|
||
use crate::sgrid;
|
||
use crate::sounding_params::{self, Level};
|
||
use crate::weather_scalar_file::{self, ScalarRow};
|
||
|
||
#[derive(Debug, thiserror::Error)]
|
||
pub enum PipelineError {
|
||
#[error("fetch: {0}")]
|
||
Fetch(#[from] fetcher::FetchError),
|
||
#[error("hrdps fetch: {0}")]
|
||
HrdpsFetch(#[from] hrdps_fetcher::HrdpsFetchError),
|
||
#[error("decode: {0}")]
|
||
Decode(#[from] decoder::DecodeError),
|
||
#[error("write: {0}")]
|
||
Write(#[from] scores_file::WriteError),
|
||
#[error("profile write: {0}")]
|
||
ProfileWrite(#[from] pgrid::WriteError),
|
||
#[error("scalar write: {0}")]
|
||
ScalarWrite(#[from] weather_scalar_file::WriteError),
|
||
#[error("sgrid write: {0}")]
|
||
SgridWrite(#[from] sgrid::WriteError),
|
||
#[error("native duct: {0}")]
|
||
NativeDuct(#[from] native_duct::NativeDuctError),
|
||
#[error("nexrad: {0}")]
|
||
Nexrad(#[from] nexrad::NexradError),
|
||
#[error("commercial: {0}")]
|
||
Commercial(#[from] commercial::CommercialError),
|
||
#[error("forecast hour 0 is reserved for Elixir")]
|
||
F00Reserved,
|
||
#[error("surface grib was empty")]
|
||
EmptyGrib,
|
||
}
|
||
|
||
#[derive(Debug, Clone)]
|
||
pub struct ChainStepInput {
|
||
pub run_time: DateTime<Utc>,
|
||
pub forecast_hour: u8,
|
||
}
|
||
|
||
impl ChainStepInput {
|
||
pub fn valid_time(&self) -> DateTime<Utc> {
|
||
self.run_time + chrono::Duration::hours(self.forecast_hour as i64)
|
||
}
|
||
}
|
||
|
||
#[derive(Debug, Clone)]
|
||
pub struct ChainStepStats {
|
||
pub score_files_written: u32,
|
||
pub point_count: u32,
|
||
pub band_count: u32,
|
||
/// Number of cells written to the per-valid_time profile file.
|
||
/// Always `point_count` after the f01..f48 profile-file write was
|
||
/// added on 2026-04-25 — kept as a separate field so the log line
|
||
/// is symmetric with the analysis step's `profile_cells_written`
|
||
/// in `RunStepStats`, and so a future failure mode that writes
|
||
/// scores but skips the profile file shows up as a divergence.
|
||
pub profile_cells_written: u32,
|
||
}
|
||
|
||
#[tracing::instrument(
|
||
name = "pipeline.run_chain_step",
|
||
skip(client, scores_dir),
|
||
fields(
|
||
run_time = %step.run_time,
|
||
forecast_hour = step.forecast_hour,
|
||
)
|
||
)]
|
||
pub async fn run_chain_step(
|
||
client: &HrrrClient,
|
||
scores_dir: &Path,
|
||
step: &ChainStepInput,
|
||
) -> Result<ChainStepStats, PipelineError> {
|
||
if step.forecast_hour == 0 {
|
||
return Err(PipelineError::F00Reserved);
|
||
}
|
||
let valid_time = step.valid_time();
|
||
let date = step.run_time.date_naive();
|
||
let hour = step.run_time.hour() as u8;
|
||
|
||
let sfc_wanted = fetcher::surface_messages_owned();
|
||
let prs_wanted = fetcher::pressure_messages_grid();
|
||
|
||
// Pin so we can drive them sequentially without losing the prs
|
||
// future when sfc resolves first.
|
||
let mut sfc_fut = Box::pin(client.fetch_product_blob(
|
||
date,
|
||
hour,
|
||
Product::Surface,
|
||
step.forecast_hour,
|
||
&sfc_wanted,
|
||
));
|
||
let mut prs_fut = Box::pin(client.fetch_product_blob(
|
||
date,
|
||
hour,
|
||
Product::Pressure,
|
||
step.forecast_hour,
|
||
&prs_wanted,
|
||
));
|
||
|
||
let grid_spec = wgrib2_grid_spec();
|
||
let sfc_pattern = match_pattern(&sfc_wanted);
|
||
let prs_pattern = match_pattern(&prs_wanted);
|
||
|
||
let fetch_started = std::time::Instant::now();
|
||
|
||
// sfc (9 messages) resolves before prs (39 messages). Start sfc
|
||
// decode as soon as it lands so the wgrib2 subprocess overlaps
|
||
// with the remaining prs byte-range downloads.
|
||
let sfc_blob = sfc_fut.as_mut().await?;
|
||
if sfc_blob.is_empty() {
|
||
return Err(PipelineError::EmptyGrib);
|
||
}
|
||
let grid_spec_sfc = grid_spec;
|
||
let sfc_decode = tokio::task::spawn_blocking(move || {
|
||
decoder::extract_grid(&sfc_blob, &sfc_pattern, grid_spec_sfc)
|
||
});
|
||
|
||
// Now await prs while sfc decodes in the background.
|
||
let prs_blob = prs_fut.as_mut().await?;
|
||
metrics::record_stage("fetch", fetch_started.elapsed());
|
||
|
||
// Both decodes run concurrently on the blocking pool.
|
||
let decode_started = std::time::Instant::now();
|
||
let grid_spec_prs = grid_spec;
|
||
let prs_decode = tokio::task::spawn_blocking(move || {
|
||
decoder::extract_grid(&prs_blob, &prs_pattern, grid_spec_prs)
|
||
});
|
||
let sfc_grid = sfc_decode.await.expect("blocking join")?;
|
||
let prs_grid = prs_decode.await.expect("blocking join")?;
|
||
metrics::record_stage("decode", decode_started.elapsed());
|
||
|
||
let mut merged = sfc_grid;
|
||
merged.merge(prs_grid);
|
||
|
||
// One pass: levels extracted once per cell, all 23 bands scored
|
||
// while the cell is hot, profile + scalar rows built alongside.
|
||
let fused = tokio::task::spawn_blocking(move || {
|
||
let out = metrics::observe_stage("derive", || {
|
||
derive_and_score(&merged, valid_time, true, None, None)
|
||
});
|
||
(out, merged.spec())
|
||
})
|
||
.await
|
||
.expect("blocking join");
|
||
let (fused, spec) = fused;
|
||
|
||
let point_count = fused.cells_scored;
|
||
let profile_cells_written = fused.cells_scored;
|
||
|
||
// Profile + scalar writes overlap with the band-score writes on the
|
||
// blocking pool rather than serialising on the hot path.
|
||
let scores_dir_owned = scores_dir.to_path_buf();
|
||
let scores_dir_for_profiles = scores_dir_owned.clone();
|
||
let profile_future = {
|
||
let body = fused.pgrid_body;
|
||
tokio::task::spawn_blocking(move || {
|
||
metrics::observe_stage("write_profile", || {
|
||
pgrid::write_atomic(&scores_dir_for_profiles, valid_time, &spec, &body, false)
|
||
.map(|_| 0u32)
|
||
})
|
||
})
|
||
};
|
||
|
||
// Dense `.sgrid` body built from the scalar rows. This is the
|
||
// primary weather-scalar artifact; the Elixir reader prefers it
|
||
// over the legacy chunked `.mp.gz` format, which is only written
|
||
// as a cold-miss fallback by the Elixir side.
|
||
let sgrid_body = sgrid::build_body(&spec, &fused.scalar_rows);
|
||
|
||
let scores_dir_for_sgrid = scores_dir_owned.clone();
|
||
let sgrid_future = {
|
||
tokio::task::spawn_blocking(move || {
|
||
metrics::observe_stage("write_sgrid", || {
|
||
sgrid::write_atomic(&scores_dir_for_sgrid, valid_time, &spec, &sgrid_body, false)
|
||
.map(|_| 0u32)
|
||
})
|
||
})
|
||
};
|
||
|
||
let band_count = fused.band_bodies.len() as u32;
|
||
let files_written = write_band_scores(
|
||
&scores_dir_owned,
|
||
valid_time,
|
||
fused.band_bodies,
|
||
spec,
|
||
false,
|
||
)
|
||
.await?;
|
||
|
||
profile_future
|
||
.await
|
||
.expect("blocking join")
|
||
.map_err(PipelineError::ProfileWrite)?;
|
||
sgrid_future
|
||
.await
|
||
.expect("blocking join")
|
||
.map_err(PipelineError::SgridWrite)?;
|
||
|
||
Ok(ChainStepStats {
|
||
score_files_written: files_written,
|
||
point_count,
|
||
band_count,
|
||
profile_cells_written,
|
||
})
|
||
}
|
||
|
||
/// HRDPS forecast-hour chain step. Sibling of `run_chain_step` for
|
||
/// `source = 'hrdps'` rows. Differences from the HRRR path:
|
||
///
|
||
/// * Single concatenated multi-record blob from
|
||
/// `hrdps_fetcher::HrdpsClient::fetch_combined_blob` (per-variable
|
||
/// fetch + byte-concat) instead of separate sfc/prs blobs.
|
||
/// * `hrdps_grid_spec()` for the wgrib2 grid extraction (Canadian
|
||
/// bbox, rotated lat/lon handled internally by wgrib2).
|
||
/// * Cells inside HRRR's CONUS bbox are dropped from the output —
|
||
/// HRRR owns those — using a HashSet of `hrdps_only_points()` as
|
||
/// the membership test.
|
||
/// * DEPR (T-Td depression) is back-filled to DPT keys before
|
||
/// `cell_to_conditions` runs, since HRDPS publishes DEPR as a
|
||
/// primary surface variable rather than DPT.
|
||
/// * Writes to `<base>/<band>/<iso>.hrdps.prop` via
|
||
/// `scores_file::write_atomic_hrdps`. Sibling file to HRRR's
|
||
/// `.prop`, never overwrites.
|
||
/// * Skips the per-valid_time profile + scalar artifacts — those
|
||
/// would clobber the HRRR-written equivalents at the same
|
||
/// `valid_time`. Re-introducing them when HRDPS surfaces on
|
||
/// `/weather` is a separate stage.
|
||
#[tracing::instrument(
|
||
name = "pipeline.run_chain_step_hrdps",
|
||
skip(client, scores_dir),
|
||
fields(
|
||
run_time = %step.run_time,
|
||
forecast_hour = step.forecast_hour,
|
||
)
|
||
)]
|
||
pub async fn run_chain_step_hrdps(
|
||
client: &HrdpsClient,
|
||
scores_dir: &Path,
|
||
step: &ChainStepInput,
|
||
) -> Result<ChainStepStats, PipelineError> {
|
||
if step.forecast_hour == 0 {
|
||
return Err(PipelineError::F00Reserved);
|
||
}
|
||
let valid_time = step.valid_time();
|
||
let cycle = step.run_time;
|
||
|
||
let blob = client
|
||
.fetch_combined_blob(cycle, step.forecast_hour)
|
||
.await?;
|
||
if blob.is_empty() {
|
||
return Err(PipelineError::EmptyGrib);
|
||
}
|
||
|
||
// wgrib2 inventory keys for HRDPS use the same NCEP-style level
|
||
// strings as HRRR (`TMP:2 m above ground`, `DEPR:850 mb`, etc.) —
|
||
// the MSC filename convention is independent of what wgrib2 prints
|
||
// when scanning the file. Match every variable HRDPS publishes;
|
||
// post-filter happens in cell_to_conditions.
|
||
let pattern = ":(TMP|DPT|DEPR|PRES|HPBL|UGRD|VGRD|TCDC|HGT):";
|
||
|
||
// Decode-once + lookup-many via the rotated-pole index. On the first
|
||
// chain step, wgrib2 -grid is called to validate the GDS parameters
|
||
// (OnceLock). The native grid is dumped as raw f32 blocks (~0.32 s per
|
||
// message) and the precomputed lookup table indexes into them for
|
||
// every target cell — no per-point rotation math, no JPEG2000 re-decode.
|
||
let spec = hrdps_grid_spec();
|
||
let pattern_owned = pattern.to_string();
|
||
let blob_for_decode = blob.clone();
|
||
let mut grid = tokio::task::spawn_blocking(move || {
|
||
let grib_path = {
|
||
let tmp_dir = std::env::temp_dir();
|
||
let nanos = std::time::SystemTime::now()
|
||
.duration_since(std::time::UNIX_EPOCH)
|
||
.map(|d| d.as_nanos())
|
||
.unwrap_or(0);
|
||
let pid = std::process::id();
|
||
let path = tmp_dir.join(format!("hrdps_native_{nanos}_{pid}.grib2"));
|
||
std::fs::write(&path, &blob_for_decode)?;
|
||
path
|
||
};
|
||
|
||
// Validate GDS parameters once (OnceLock).
|
||
let grid_output = decoder::extract_grid_metadata(&grib_path)?;
|
||
let params = crate::rotated_pole::init_or_validate(&grid_output).map_err(|e| {
|
||
decoder::DecodeError::Wgrib2Failed {
|
||
code: -1,
|
||
stderr: e,
|
||
}
|
||
})?;
|
||
|
||
let lookup = crate::rotated_pole::lookup_table(&spec).map_err(|e| {
|
||
decoder::DecodeError::Wgrib2Failed {
|
||
code: -1,
|
||
stderr: e,
|
||
}
|
||
})?;
|
||
|
||
let result = decoder::decode_hrdps_native(
|
||
&std::fs::read(&grib_path)?,
|
||
&pattern_owned,
|
||
&spec,
|
||
lookup,
|
||
params.nx,
|
||
params.ny,
|
||
);
|
||
let _ = std::fs::remove_file(&grib_path);
|
||
result
|
||
})
|
||
.await
|
||
.expect("blocking join")?;
|
||
drop(blob);
|
||
|
||
// HRDPS publishes DEPR (T-Td depression in K) as a primary surface +
|
||
// pressure-level variable rather than DPT. Derive DPT planes from
|
||
// `TMP - DEPR` so the shared plane table sees the expected names.
|
||
backfill_dpt_from_depr(&mut grid);
|
||
|
||
// Cells inside HRRR's CONUS bbox are HRRR's to own. `hrdps_only_points`
|
||
// is the authoritative list; turn it into a per-cell mask so the fused
|
||
// pass skips them and they stay NO_DATA in the score file.
|
||
let mask = hrdps_cell_mask(&grid);
|
||
|
||
let fused = tokio::task::spawn_blocking(move || {
|
||
derive_and_score(&grid, valid_time, false, Some(&mask), None)
|
||
})
|
||
.await
|
||
.expect("blocking join");
|
||
|
||
let point_count = fused.cells_scored;
|
||
tracing::info!(
|
||
grid_cells = spec.lat_count * spec.lon_count,
|
||
scored = point_count,
|
||
"hrdps cells scored"
|
||
);
|
||
|
||
let scores_dir_owned = scores_dir.to_path_buf();
|
||
|
||
// Dense `.sgrid` body built from the scalar rows while they're
|
||
// still available.
|
||
let sgrid_body = sgrid::build_body(&spec, &fused.scalar_rows);
|
||
|
||
// Write the HRDPS scalar artifact alongside the scores. Sibling dir
|
||
// (`<vt>.hrdps/`) so HRRR's `<vt>/` write isn't clobbered. Without
|
||
// this the /weather map would only show CONUS data.
|
||
let scalar_dir = scores_dir_owned.clone();
|
||
let scalar_future = {
|
||
let rows = fused.scalar_rows;
|
||
tokio::task::spawn_blocking(move || {
|
||
weather_scalar_file::write_atomic_hrdps(&scalar_dir, valid_time, &rows).map(|_| ())
|
||
})
|
||
};
|
||
|
||
// Write the HRDPS `.sgrid` alongside the chunked scalar artifact.
|
||
let sgrid_dir = scores_dir_owned.clone();
|
||
let sgrid_future = {
|
||
tokio::task::spawn_blocking(move || {
|
||
sgrid::write_atomic(&sgrid_dir, valid_time, &spec, &sgrid_body, true).map(|_| ())
|
||
})
|
||
};
|
||
|
||
let band_count = fused.band_bodies.len() as u32;
|
||
let files_written =
|
||
write_band_scores(&scores_dir_owned, valid_time, fused.band_bodies, spec, true).await?;
|
||
|
||
scalar_future
|
||
.await
|
||
.expect("blocking join")
|
||
.map_err(PipelineError::ScalarWrite)?;
|
||
sgrid_future
|
||
.await
|
||
.expect("blocking join")
|
||
.map_err(PipelineError::SgridWrite)?;
|
||
|
||
Ok(ChainStepStats {
|
||
score_files_written: files_written,
|
||
point_count,
|
||
band_count,
|
||
// HRDPS skips the per-valid_time profile artifact — writing it
|
||
// would clobber the HRRR-written file at the same valid_time.
|
||
profile_cells_written: 0,
|
||
})
|
||
}
|
||
|
||
/// Per-cell mask marking the HRDPS-owned cells (inside the Canadian bbox,
|
||
/// outside HRRR's CONUS interior). Built from `hrdps_only_points()` so
|
||
/// the seam rules documented there stay the single source of truth.
|
||
fn hrdps_cell_mask(grid: &FieldGrid) -> Vec<bool> {
|
||
let mut mask = vec![false; grid.n_cells()];
|
||
for (lat, lon) in hrdps_only_points() {
|
||
if let Some(cell) = grid.cell_index(lat, lon) {
|
||
mask[cell] = true;
|
||
}
|
||
}
|
||
mask
|
||
}
|
||
|
||
/// Derive `DPT:<level>` planes from `TMP:<level> - DEPR:<level>` for
|
||
/// every DEPR plane that has no matching DPT. HRDPS publishes the
|
||
/// dewpoint *depression* rather than the dewpoint itself.
|
||
///
|
||
/// Whole-plane arithmetic rather than the old per-cell key rewriting:
|
||
/// one pass per level instead of a `format!` + hash probe per cell per
|
||
/// level.
|
||
fn backfill_dpt_from_depr(grid: &mut FieldGrid) {
|
||
// Collect first — we can't hold a borrow on the name list while
|
||
// mutating the plane store.
|
||
let pending: Vec<(String, String)> = grid
|
||
.plane_names()
|
||
.iter()
|
||
.filter_map(|name| {
|
||
let level = name.strip_prefix("DEPR:")?;
|
||
let dpt = format!("DPT:{level}");
|
||
if grid.plane_id(&dpt).is_some() {
|
||
return None;
|
||
}
|
||
grid.plane_id(&format!("TMP:{level}"))?;
|
||
Some((name.to_string(), dpt))
|
||
})
|
||
.collect();
|
||
|
||
for (depr_name, dpt_name) in pending {
|
||
let level = depr_name.strip_prefix("DEPR:").unwrap_or_default();
|
||
let (Some(tmp_id), Some(depr_id)) = (
|
||
grid.plane_id(&format!("TMP:{level}")),
|
||
grid.plane_id(&depr_name),
|
||
) else {
|
||
continue;
|
||
};
|
||
let derived: Vec<f32> = (0..grid.n_cells())
|
||
.map(|c| grid.at_raw(tmp_id, c) - grid.at_raw(depr_id, c))
|
||
.collect();
|
||
grid.push_plane(&dpt_name, derived);
|
||
}
|
||
}
|
||
|
||
/// Cells per rayon work unit in the fused pass. Big enough to amortise
|
||
/// the per-chunk `Vec` allocations, small enough that 4 cores stay
|
||
/// balanced across a 95 k-cell grid.
|
||
const FUSE_CHUNK: usize = 2_048;
|
||
|
||
/// Everything one pass over the grid produces.
|
||
pub struct FusedOutput {
|
||
/// Dense `u8` score body per band, row-major on the grid spec —
|
||
/// directly what `scores_file::encode_dense` wants.
|
||
pub band_bodies: Vec<(u32, Vec<u8>)>,
|
||
/// Flat cell-major `f32` body for the `.pgrid` profile artifact,
|
||
/// length `n_cells * pgrid::n_fields()`. Empty when profiles were
|
||
/// not requested.
|
||
pub pgrid_body: Vec<f32>,
|
||
pub scalar_rows: Vec<ScalarRow>,
|
||
/// Cells that produced a `Conditions` (i.e. had surface T and Td).
|
||
pub cells_scored: u32,
|
||
}
|
||
|
||
/// Single pass over every grid cell producing the band scores, the
|
||
/// profile entries and the scalar rows together.
|
||
///
|
||
/// Replaces the previous shape — three independent 95 k-cell derivation
|
||
/// passes followed by 23 more full passes (`bands.par_iter()`) over a
|
||
/// staged `Vec<(f64, f64, Conditions, BandInvariants)>`. That staging
|
||
/// array was ~19 MB and each band re-streamed all of it, ~437 MB of
|
||
/// memory traffic per chain step, then scattered 23 × 95 k `ScorePoint`s
|
||
/// into dense bodies.
|
||
///
|
||
/// Here each cell is touched once: its levels are extracted once (shared
|
||
/// by conditions, profile and scalars), its `BandInvariants` are computed
|
||
/// once, and all 23 bands are scored while it is still in L1.
|
||
///
|
||
/// Scores accumulate **cell-major** (`scores[cell * n_bands + band]`) so
|
||
/// a chunk of cells owns one contiguous slice and rayon can hand out
|
||
/// disjoint `&mut` without locking. The final transpose to per-band
|
||
/// bodies is a single strided copy.
|
||
fn derive_and_score(
|
||
grid: &FieldGrid,
|
||
valid_time: DateTime<Utc>,
|
||
want_profiles: bool,
|
||
cell_mask: Option<&[bool]>,
|
||
kp: Option<i32>,
|
||
) -> FusedOutput {
|
||
use rayon::prelude::*;
|
||
|
||
let planes = GridPlanes::resolve(grid);
|
||
let bands = band_config::all_bands();
|
||
let n_bands = bands.len();
|
||
let n_cells = grid.n_cells();
|
||
|
||
let n_pgrid = crate::pgrid::n_fields();
|
||
let mut cell_scores = vec![scores_file::NO_DATA; n_cells * n_bands];
|
||
// Allocated even when profiles aren't wanted (zero-length then), so
|
||
// the zip below always lines up chunk-for-chunk with cell_scores.
|
||
let mut pgrid_body = vec![f32::NAN; if want_profiles { n_cells * n_pgrid } else { 0 }];
|
||
|
||
let per_chunk: Vec<(Vec<ScalarRow>, u32)> = if want_profiles {
|
||
cell_scores
|
||
.par_chunks_mut(FUSE_CHUNK * n_bands)
|
||
.zip(pgrid_body.par_chunks_mut(FUSE_CHUNK * n_pgrid))
|
||
.enumerate()
|
||
.map(|(chunk_idx, (score_out, pgrid_out))| {
|
||
fuse_chunk(
|
||
grid,
|
||
&planes,
|
||
bands,
|
||
valid_time,
|
||
cell_mask,
|
||
kp,
|
||
chunk_idx * FUSE_CHUNK,
|
||
score_out,
|
||
Some(pgrid_out),
|
||
)
|
||
})
|
||
.collect()
|
||
} else {
|
||
cell_scores
|
||
.par_chunks_mut(FUSE_CHUNK * n_bands)
|
||
.enumerate()
|
||
.map(|(chunk_idx, score_out)| {
|
||
fuse_chunk(
|
||
grid,
|
||
&planes,
|
||
bands,
|
||
valid_time,
|
||
cell_mask,
|
||
kp,
|
||
chunk_idx * FUSE_CHUNK,
|
||
score_out,
|
||
None,
|
||
)
|
||
})
|
||
.collect()
|
||
};
|
||
|
||
let mut scalar_rows = Vec::new();
|
||
let mut cells_scored = 0u32;
|
||
for (s, n) in per_chunk {
|
||
scalar_rows.extend(s);
|
||
cells_scored += n;
|
||
}
|
||
|
||
// Transpose cell-major → one dense body per band.
|
||
let band_bodies: Vec<(u32, Vec<u8>)> = bands
|
||
.par_iter()
|
||
.enumerate()
|
||
.map(|(b, band)| {
|
||
let mut body = vec![scores_file::NO_DATA; n_cells];
|
||
for (cell, slot) in body.iter_mut().enumerate() {
|
||
*slot = cell_scores[cell * n_bands + b];
|
||
}
|
||
(band.freq_mhz, body)
|
||
})
|
||
.collect();
|
||
|
||
FusedOutput {
|
||
band_bodies,
|
||
pgrid_body,
|
||
scalar_rows,
|
||
cells_scored,
|
||
}
|
||
}
|
||
|
||
/// One rayon work unit of the fused pass: derive and score the cells in
|
||
/// `[base, base + score_out.len() / n_bands)`.
|
||
///
|
||
/// `score_out` is this chunk's slice of the cell-major score array and
|
||
/// `pgrid_out` (when profiles are wanted) is the matching slice of the
|
||
/// `.pgrid` body. Both are disjoint per chunk, so no locking.
|
||
#[allow(clippy::too_many_arguments)]
|
||
fn fuse_chunk(
|
||
grid: &FieldGrid,
|
||
planes: &GridPlanes,
|
||
bands: &[band_config::BandConfig],
|
||
valid_time: DateTime<Utc>,
|
||
cell_mask: Option<&[bool]>,
|
||
kp: Option<i32>,
|
||
base: usize,
|
||
score_out: &mut [u8],
|
||
pgrid_out: Option<&mut [f32]>,
|
||
) -> (Vec<ScalarRow>, u32) {
|
||
let n_bands = bands.len();
|
||
let n_pgrid = crate::pgrid::n_fields();
|
||
let mut scalars: Vec<ScalarRow> = Vec::new();
|
||
let mut levels: Vec<Level> = Vec::new();
|
||
let mut scored = 0u32;
|
||
|
||
// Split once outside the loop; `None` becomes an empty iterator so
|
||
// the non-profile path costs nothing.
|
||
let mut pgrid_iter = pgrid_out.map(|s| s.chunks_mut(n_pgrid));
|
||
|
||
for (k, cell_out) in score_out.chunks_mut(n_bands).enumerate() {
|
||
let cell = base + k;
|
||
let pgrid_slot = pgrid_iter.as_mut().and_then(|it| it.next());
|
||
|
||
if cell_mask.is_some_and(|m| !m[cell]) {
|
||
continue;
|
||
}
|
||
let (lat, lon) = grid.cell_latlon(cell);
|
||
|
||
planes.levels_at(grid, cell, &mut levels);
|
||
|
||
let Some(conditions) =
|
||
cell_to_conditions(grid, planes, cell, lat, lon, &valid_time, &levels)
|
||
else {
|
||
continue;
|
||
};
|
||
let invariants = scorer::precompute_band_invariants(&conditions);
|
||
|
||
// Commercial-link degradation + n_links for this cell. Both are
|
||
// f00-only planes written by `apply_commercial`; for forecast
|
||
// hours (where the planes don't exist) these stay `None` and
|
||
// the commercial boost is skipped.
|
||
let degradation_db = grid
|
||
.at_opt(planes.commercial_degradation_db, cell)
|
||
.map(|v| v as f64)
|
||
.filter(|v| v.is_finite());
|
||
let n_links = grid
|
||
.at_opt(planes.commercial_n_links, cell)
|
||
.map(|v| v as u32);
|
||
|
||
for (b, band) in bands.iter().enumerate() {
|
||
let r = scorer::composite_score_with(&conditions, band, Some(invariants));
|
||
let mut score = r.score;
|
||
|
||
// f00-only: commercial-link inverse-sensor boost.
|
||
if let (Some(deg), Some(n)) = (degradation_db, n_links) {
|
||
if deg >= 3.0 {
|
||
score = scorer::commercial_link_boost(score, n, deg);
|
||
}
|
||
}
|
||
|
||
// Aurora boost — geomagnetic storm bonus for VHF ≤ 432 MHz.
|
||
if let Some(kp_val) = kp {
|
||
score = scorer::aurora_boost(score, Some(kp_val), band.freq_mhz);
|
||
}
|
||
|
||
cell_out[b] = clamp_score_u8(score);
|
||
}
|
||
scored += 1;
|
||
|
||
if let Some(slot) = pgrid_slot {
|
||
crate::planes::fill_pgrid_record(grid, planes, cell, &levels, slot);
|
||
}
|
||
if let Some(row) =
|
||
weather_scalar_file::derive_row(grid, planes, cell, lat, lon, valid_time, &levels)
|
||
{
|
||
scalars.push(row);
|
||
}
|
||
}
|
||
|
||
(scalars, scored)
|
||
}
|
||
|
||
fn clamp_score_u8(s: i32) -> u8 {
|
||
s.clamp(0, 100) as u8
|
||
}
|
||
|
||
/// Public shim so `tests/fused_pass_perf.rs` can time the fused pass
|
||
/// without making the pass itself part of the crate's API surface.
|
||
pub fn derive_and_score_for_bench(
|
||
grid: &FieldGrid,
|
||
valid_time: DateTime<Utc>,
|
||
want_profiles: bool,
|
||
cell_mask: Option<&[bool]>,
|
||
) -> FusedOutput {
|
||
derive_and_score(grid, valid_time, want_profiles, cell_mask, None)
|
||
}
|
||
|
||
/// Write every band's dense score body. Each file is an independent NFS
|
||
/// fsync, so they overlap on the blocking pool rather than serialising.
|
||
async fn write_band_scores(
|
||
scores_dir: &Path,
|
||
valid_time: DateTime<Utc>,
|
||
bodies: Vec<(u32, Vec<u8>)>,
|
||
spec: crate::grid::GridSpec,
|
||
hrdps: bool,
|
||
) -> Result<u32, PipelineError> {
|
||
let started = std::time::Instant::now();
|
||
let count = bodies.len() as u32;
|
||
let handles: Vec<_> = bodies
|
||
.into_iter()
|
||
.map(|(band_mhz, body)| {
|
||
let dir = scores_dir.to_path_buf();
|
||
tokio::task::spawn_blocking(move || {
|
||
scores_file::write_atomic_dense(&dir, band_mhz, valid_time, &body, spec, hrdps)
|
||
})
|
||
})
|
||
.collect();
|
||
for h in handles {
|
||
h.await.expect("blocking join")?;
|
||
}
|
||
metrics::record_stage("write_scores", started.elapsed());
|
||
Ok(count)
|
||
}
|
||
|
||
/// wgrib2 `-match` regex: `":(A|B|C):"`. Var names only — levels are
|
||
/// post-filtered when we read the binary back.
|
||
fn match_pattern(wanted: &[(String, String)]) -> String {
|
||
let mut vars: Vec<&str> = wanted.iter().map(|(v, _)| v.as_str()).collect();
|
||
vars.sort();
|
||
vars.dedup();
|
||
format!(":({}):", vars.join("|"))
|
||
}
|
||
|
||
fn cell_to_conditions(
|
||
grid: &FieldGrid,
|
||
p: &GridPlanes,
|
||
cell: usize,
|
||
lat: f64,
|
||
lon: f64,
|
||
valid_time: &DateTime<Utc>,
|
||
levels: &[Level],
|
||
) -> Option<Conditions> {
|
||
let tmp_k = grid.at_opt(p.tmp_2m, cell)?;
|
||
let dpt_k = grid.at_opt(p.dpt_2m, cell)?;
|
||
let temp_c = (tmp_k - 273.15) as f64;
|
||
let dewpoint_c = (dpt_k - 273.15) as f64;
|
||
let temp_f = scorer::c_to_f(temp_c);
|
||
let dewpoint_f = scorer::c_to_f(dewpoint_c);
|
||
let abs_humidity = scorer::absolute_humidity(temp_c, dewpoint_c);
|
||
|
||
let wind_speed_kts = match (grid.at_opt(p.ugrd_10m, cell), grid.at_opt(p.vgrd_10m, cell)) {
|
||
(Some(u), Some(v)) => Some(scorer::wind_speed_kts(u as f64, v as f64)),
|
||
_ => None,
|
||
};
|
||
|
||
let sky_cover_pct = grid.at_opt(p.tcdc, cell).map(|v| v as f64);
|
||
let pwat_mm = grid.at_opt(p.pwat, cell).map(|v| v as f64);
|
||
let pressure_mb = grid.at_opt(p.pres_sfc, cell).map(|pa| pa as f64 / 100.0);
|
||
// Mirror Elixir's merged_rain_rate: pick the heavier of HRRR's
|
||
// accumulation-derived rate and NEXRAD's reflectivity-derived rate
|
||
// so a fast convective cell that hasn't yet shown up in the hourly
|
||
// APCP still triggers the rain penalty.
|
||
let hrrr_rate = grid.at_opt(p.apcp, cell).map(|mm| mm as f64).unwrap_or(0.0);
|
||
let nexrad_rate =
|
||
scorer::dbz_to_rain_rate_mmhr(grid.at_opt(p.nexrad_dbz, cell).map(|v| v as f64));
|
||
let merged_rate = hrrr_rate.max(nexrad_rate);
|
||
let rain_rate_mmhr = if merged_rate > 0.0 {
|
||
Some(merged_rate)
|
||
} else {
|
||
None
|
||
};
|
||
let bl_depth_m = grid.at_opt(p.hpbl, cell).map(|v| v as f64);
|
||
let best_duct_band_ghz = grid.at_opt(p.best_duct_freq_ghz, cell).map(|v| v as f64);
|
||
|
||
let min_refractivity_gradient = {
|
||
// Native hybrid-sigma gradient (f00 enrichment) overrides the
|
||
// pressure-level derived value when available — exactly the
|
||
// Elixir `hrrr_profile[:native_min_gradient] || derived[:min_refractivity_gradient]`
|
||
// precedence. Forecast hours don't have the native duct grid
|
||
// merged, so the plane is absent and the pressure-level fallback
|
||
// is used.
|
||
let native = grid
|
||
.at_opt(p.native_min_gradient, cell)
|
||
.map(|v| v as f64)
|
||
.filter(|v| v.is_finite());
|
||
native.unwrap_or_else(|| {
|
||
sounding_params::min_refractivity_gradient(levels.to_vec()).unwrap_or(f64::NAN)
|
||
})
|
||
};
|
||
let min_refractivity_gradient = if min_refractivity_gradient.is_finite() {
|
||
Some(min_refractivity_gradient)
|
||
} else {
|
||
None
|
||
};
|
||
|
||
Some(Conditions {
|
||
abs_humidity,
|
||
temp_f,
|
||
dewpoint_f,
|
||
wind_speed_kts,
|
||
sky_cover_pct,
|
||
utc_hour: valid_time.hour() as u8,
|
||
utc_minute: valid_time.minute() as u8,
|
||
month: valid_time.month() as u8,
|
||
longitude: lon,
|
||
latitude: Some(lat),
|
||
pressure_mb,
|
||
prev_pressure_mb: None,
|
||
rain_rate_mmhr,
|
||
min_refractivity_gradient,
|
||
bl_depth_m,
|
||
pwat_mm,
|
||
best_duct_band_ghz,
|
||
bulk_richardson: None,
|
||
})
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use crate::grid::GridSpec;
|
||
use crate::scores_file::ScorePoint;
|
||
|
||
/// Single-cell grid so a test can express "one cell with these
|
||
/// values" as concisely as the old `CellValues` map did.
|
||
fn one_cell_spec() -> GridSpec {
|
||
GridSpec {
|
||
lon_start: -97.0,
|
||
lon_count: 1,
|
||
lon_step: 0.125,
|
||
lat_start: 32.0,
|
||
lat_count: 1,
|
||
lat_step: 0.125,
|
||
}
|
||
}
|
||
|
||
fn cell_grid(items: &[(&str, f32)]) -> FieldGrid {
|
||
let mut g = FieldGrid::new(one_cell_spec());
|
||
for &(k, v) in items {
|
||
g.push_plane(k, vec![v]);
|
||
}
|
||
g
|
||
}
|
||
|
||
/// Resolve planes, extract the cell's levels, derive conditions —
|
||
/// the sequence `derive_and_score` runs, collapsed for tests.
|
||
fn conditions_of(grid: &FieldGrid, vt: &DateTime<Utc>) -> Option<Conditions> {
|
||
let p = GridPlanes::resolve(grid);
|
||
let mut levels = Vec::new();
|
||
p.levels_at(grid, 0, &mut levels);
|
||
cell_to_conditions(grid, &p, 0, 32.0, -97.0, vt, &levels)
|
||
}
|
||
|
||
/// One cell's `.pgrid` record, keyed by field name for readability.
|
||
fn pgrid_record_of(grid: &FieldGrid) -> std::collections::HashMap<String, f32> {
|
||
let p = GridPlanes::resolve(grid);
|
||
let mut levels = Vec::new();
|
||
p.levels_at(grid, 0, &mut levels);
|
||
let mut rec = pgrid::empty_record();
|
||
crate::planes::fill_pgrid_record(grid, &p, 0, &levels, &mut rec);
|
||
pgrid::field_names().iter().cloned().zip(rec).collect()
|
||
}
|
||
|
||
#[test]
|
||
fn match_pattern_dedups_vars() {
|
||
let wanted = vec![
|
||
("TMP".into(), "2 m above ground".into()),
|
||
("DPT".into(), "2 m above ground".into()),
|
||
("TMP".into(), "1000 mb".into()),
|
||
];
|
||
assert_eq!(match_pattern(&wanted), ":(DPT|TMP):");
|
||
}
|
||
|
||
#[test]
|
||
fn cell_to_conditions_maps_fields() {
|
||
use chrono::TimeZone;
|
||
let grid = cell_grid(&[
|
||
("TMP:2 m above ground", 298.15), // 25 °C
|
||
("DPT:2 m above ground", 293.15), // 20 °C
|
||
("UGRD:10 m above ground", 3.0),
|
||
("VGRD:10 m above ground", 4.0),
|
||
("TCDC:entire atmosphere", 30.0),
|
||
("PRES:surface", 101_000.0),
|
||
("HPBL:surface", 400.0),
|
||
(
|
||
"PWAT:entire atmosphere (considered as a single layer)",
|
||
25.0,
|
||
),
|
||
]);
|
||
|
||
let vt = Utc.with_ymd_and_hms(2026, 6, 15, 18, 0, 0).unwrap();
|
||
let c = conditions_of(&grid, &vt).unwrap();
|
||
assert!((c.temp_f - 77.0).abs() < 0.1);
|
||
assert!((c.dewpoint_f - 68.0).abs() < 0.1);
|
||
assert!((c.wind_speed_kts.unwrap() - 9.72).abs() < 0.1);
|
||
assert!((c.pressure_mb.unwrap() - 1010.0).abs() < 0.01);
|
||
assert_eq!(c.sky_cover_pct, Some(30.0));
|
||
assert_eq!(c.bl_depth_m, Some(400.0));
|
||
assert_eq!(c.pwat_mm, Some(25.0));
|
||
assert_eq!(c.month, 6);
|
||
assert_eq!(c.longitude, -97.0);
|
||
}
|
||
|
||
#[test]
|
||
fn cell_missing_surface_returns_none() {
|
||
use chrono::TimeZone;
|
||
let grid = cell_grid(&[]);
|
||
let vt = Utc.with_ymd_and_hms(2026, 6, 15, 18, 0, 0).unwrap();
|
||
assert!(conditions_of(&grid, &vt).is_none());
|
||
}
|
||
|
||
/// A cell whose surface planes exist but hold the no-data sentinel
|
||
/// must be rejected exactly like a cell with no planes at all. The
|
||
/// dense representation makes this the interesting edge case: the
|
||
/// plane is present, the value is NaN.
|
||
#[test]
|
||
fn cell_with_nan_surface_returns_none() {
|
||
use chrono::TimeZone;
|
||
let grid = cell_grid(&[
|
||
("TMP:2 m above ground", f32::NAN),
|
||
("DPT:2 m above ground", 293.15),
|
||
]);
|
||
let vt = Utc.with_ymd_and_hms(2026, 6, 15, 18, 0, 0).unwrap();
|
||
assert!(conditions_of(&grid, &vt).is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn nexrad_reflectivity_lifts_rain_rate_above_hrrr_precip() {
|
||
// Mirrors the Elixir scorer's rain merge: when HRRR's hourly
|
||
// accumulation hasn't caught a fast convective cell but NEXRAD
|
||
// sees 45 dBZ overhead, take whichever rate is higher. Without
|
||
// this the Rust pipeline silently underestimates rain attenuation.
|
||
use chrono::TimeZone;
|
||
let grid = cell_grid(&[
|
||
("TMP:2 m above ground", 295.0),
|
||
("DPT:2 m above ground", 285.0),
|
||
("APCP:surface", 0.0),
|
||
("nexrad_max_reflectivity_dbz", 45.0),
|
||
]);
|
||
|
||
let vt = Utc.with_ymd_and_hms(2026, 6, 15, 18, 0, 0).unwrap();
|
||
let c = conditions_of(&grid, &vt).unwrap();
|
||
let rate = c.rain_rate_mmhr.expect("nexrad-derived rain rate present");
|
||
assert!(rate > 5.0, "expected >5 mm/hr from 45 dBZ, got {rate}");
|
||
}
|
||
|
||
#[test]
|
||
fn duct_freq_threads_into_best_duct_band_ghz() {
|
||
// After native_duct::merge_duct_grid runs, every cell with a
|
||
// detected trapping layer carries `best_duct_freq_ghz`. The
|
||
// scorer reads `best_duct_band_ghz` from Conditions for the 15%
|
||
// Native Duct Boost. cell_to_conditions must wire one to the
|
||
// other or the boost never fires from the Rust pipeline.
|
||
use chrono::TimeZone;
|
||
let grid = cell_grid(&[
|
||
("TMP:2 m above ground", 295.0),
|
||
("DPT:2 m above ground", 285.0),
|
||
("best_duct_freq_ghz", 24.0),
|
||
]);
|
||
|
||
let vt = Utc.with_ymd_and_hms(2026, 6, 15, 18, 0, 0).unwrap();
|
||
let c = conditions_of(&grid, &vt).unwrap();
|
||
assert_eq!(c.best_duct_band_ghz, Some(24.0));
|
||
}
|
||
|
||
#[test]
|
||
fn rejects_forecast_hour_zero() {
|
||
use chrono::TimeZone;
|
||
let client = HrrrClient::default_base().unwrap();
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let step = ChainStepInput {
|
||
run_time: Utc.with_ymd_and_hms(2026, 4, 19, 15, 0, 0).unwrap(),
|
||
forecast_hour: 0,
|
||
};
|
||
let rt = tokio::runtime::Runtime::new().unwrap();
|
||
let err = rt
|
||
.block_on(run_chain_step(&client, dir.path(), &step))
|
||
.unwrap_err();
|
||
assert!(matches!(err, PipelineError::F00Reserved));
|
||
}
|
||
|
||
#[test]
|
||
fn hrdps_rejects_forecast_hour_zero() {
|
||
use chrono::TimeZone;
|
||
let client = HrdpsClient::default_base().unwrap();
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let step = ChainStepInput {
|
||
run_time: Utc.with_ymd_and_hms(2026, 4, 29, 12, 0, 0).unwrap(),
|
||
forecast_hour: 0,
|
||
};
|
||
let rt = tokio::runtime::Runtime::new().unwrap();
|
||
let err = rt
|
||
.block_on(run_chain_step_hrdps(&client, dir.path(), &step))
|
||
.unwrap_err();
|
||
assert!(matches!(err, PipelineError::F00Reserved));
|
||
}
|
||
|
||
fn plane_value(grid: &FieldGrid, name: &str) -> Option<f32> {
|
||
grid.at(grid.plane_id(name)?, 0)
|
||
}
|
||
|
||
#[test]
|
||
fn backfill_dpt_from_depr_derives_missing_dewpoint() {
|
||
let mut grid = cell_grid(&[
|
||
("TMP:2 m above ground", 290.15), // 17°C
|
||
("DEPR:2 m above ground", 5.0),
|
||
]);
|
||
|
||
backfill_dpt_from_depr(&mut grid);
|
||
|
||
let dpt = plane_value(&grid, "DPT:2 m above ground").expect("DPT back-filled from DEPR");
|
||
assert!((dpt - 285.15).abs() < 1e-3, "got {dpt}");
|
||
}
|
||
|
||
#[test]
|
||
fn backfill_dpt_does_not_clobber_existing_dpt() {
|
||
let mut grid = cell_grid(&[
|
||
("TMP:2 m above ground", 290.15),
|
||
("DEPR:2 m above ground", 5.0),
|
||
("DPT:2 m above ground", 280.0), // pre-existing
|
||
]);
|
||
|
||
backfill_dpt_from_depr(&mut grid);
|
||
|
||
assert_eq!(plane_value(&grid, "DPT:2 m above ground"), Some(280.0));
|
||
}
|
||
|
||
#[test]
|
||
fn backfill_dpt_handles_pressure_levels() {
|
||
let mut grid = cell_grid(&[
|
||
("TMP:850 mb", 275.0),
|
||
("DEPR:850 mb", 4.0),
|
||
("TMP:700 mb", 265.0),
|
||
("DEPR:700 mb", 3.0),
|
||
]);
|
||
|
||
backfill_dpt_from_depr(&mut grid);
|
||
|
||
assert_eq!(plane_value(&grid, "DPT:850 mb"), Some(271.0));
|
||
assert_eq!(plane_value(&grid, "DPT:700 mb"), Some(262.0));
|
||
}
|
||
|
||
/// DEPR with no matching TMP can't produce a dewpoint; the plane
|
||
/// must not be invented (the old code's `?` on the TMP lookup did
|
||
/// the same).
|
||
#[test]
|
||
fn backfill_dpt_skips_depr_without_temp() {
|
||
let mut grid = cell_grid(&[("DEPR:850 mb", 4.0)]);
|
||
backfill_dpt_from_depr(&mut grid);
|
||
assert!(grid.plane_id("DPT:850 mb").is_none());
|
||
}
|
||
|
||
/// A small multi-cell grid, cells 0 and 3 populated, on a spec whose
|
||
/// cell indices are easy to reason about (3 cols × 2 rows).
|
||
fn small_grid() -> (FieldGrid, GridSpec) {
|
||
let spec = GridSpec {
|
||
lon_start: -97.0,
|
||
lon_count: 3,
|
||
lon_step: 0.125,
|
||
lat_start: 32.0,
|
||
lat_count: 2,
|
||
lat_step: 0.125,
|
||
};
|
||
let mut g = FieldGrid::new(spec);
|
||
let n = spec.lon_count * spec.lat_count;
|
||
let at = |cells: &[usize], v: f32| {
|
||
let mut plane = vec![f32::NAN; n];
|
||
for &c in cells {
|
||
plane[c] = v;
|
||
}
|
||
plane
|
||
};
|
||
let live = [0usize, 3];
|
||
g.push_plane("TMP:2 m above ground", at(&live, 295.0));
|
||
g.push_plane("DPT:2 m above ground", at(&live, 285.0));
|
||
g.push_plane("UGRD:10 m above ground", at(&live, 2.0));
|
||
g.push_plane("VGRD:10 m above ground", at(&live, 1.0));
|
||
g.push_plane("TCDC:entire atmosphere", at(&live, 20.0));
|
||
g.push_plane("PRES:surface", at(&live, 101_000.0));
|
||
g.push_plane("HPBL:surface", at(&live, 500.0));
|
||
g.push_plane(
|
||
"PWAT:entire atmosphere (considered as a single layer)",
|
||
at(&live, 20.0),
|
||
);
|
||
let keys = fetcher::grid_level_keys();
|
||
g.push_plane(&keys[0].tmp, at(&live, 295.0));
|
||
g.push_plane(&keys[0].dpt, at(&live, 285.0));
|
||
g.push_plane(&keys[0].hgt, at(&live, 100.0));
|
||
g.push_plane(&keys[1].tmp, at(&live, 293.0));
|
||
g.push_plane(&keys[1].dpt, at(&live, 280.0));
|
||
g.push_plane(&keys[1].hgt, at(&live, 300.0));
|
||
(g, spec)
|
||
}
|
||
|
||
/// Integration-ish: fused derive+score over a hand-built grid →
|
||
/// `scores_file::encode_dense` → `decode`, asserting the score lands
|
||
/// at the right cell and unpopulated cells stay no-data. Covers the
|
||
/// post-fetch hot path (the one that scales with grid size) without
|
||
/// standing up wgrib2.
|
||
#[test]
|
||
fn fused_score_write_read_roundtrip() {
|
||
use chrono::TimeZone;
|
||
let (grid, spec) = small_grid();
|
||
let valid_time = Utc.with_ymd_and_hms(2026, 6, 15, 18, 0, 0).unwrap();
|
||
|
||
let fused = derive_and_score(&grid, valid_time, true, None, None);
|
||
|
||
assert_eq!(fused.cells_scored, 2);
|
||
assert_eq!(
|
||
fused.pgrid_body.len(),
|
||
spec.lat_count * spec.lon_count * pgrid::n_fields()
|
||
);
|
||
assert_eq!(fused.band_bodies.len(), band_config::all_bands().len());
|
||
|
||
let (band_mhz, body) = &fused.band_bodies[0];
|
||
assert_eq!(body.len(), spec.lon_count * spec.lat_count);
|
||
// Populated cells scored 0..=100; the rest kept the sentinel.
|
||
assert!(body[0] <= 100, "cell 0 should be scored, got {}", body[0]);
|
||
assert!(body[3] <= 100, "cell 3 should be scored, got {}", body[3]);
|
||
for c in [1usize, 2, 4, 5] {
|
||
assert_eq!(body[c], scores_file::NO_DATA, "cell {c} should be no-data");
|
||
}
|
||
|
||
let dir = tempfile::tempdir().unwrap();
|
||
scores_file::write_atomic_dense(dir.path(), *band_mhz, valid_time, body, spec, false)
|
||
.unwrap();
|
||
|
||
let path = scores_file::path_for(dir.path(), *band_mhz, valid_time);
|
||
let bytes = std::fs::read(&path).expect("read back score file");
|
||
let decoded = scores_file::decode(&bytes).expect("decode");
|
||
assert_eq!(decoded.band_mhz, *band_mhz);
|
||
assert_eq!(decoded.n_rows, spec.lat_count as u16);
|
||
assert_eq!(decoded.n_cols, spec.lon_count as u16);
|
||
assert_eq!(decoded.body, *body, "body survived the round trip");
|
||
}
|
||
|
||
/// The dense body index must agree with `scores_file`'s own
|
||
/// coordinate→index arithmetic, or scores land in the wrong place on
|
||
/// the map. Compare the fused (index-based) path against the legacy
|
||
/// `ScorePoint` scatter for the same inputs.
|
||
#[test]
|
||
fn dense_body_placement_matches_scorepoint_scatter() {
|
||
use chrono::TimeZone;
|
||
let (grid, spec) = small_grid();
|
||
let valid_time = Utc.with_ymd_and_hms(2026, 6, 15, 18, 0, 0).unwrap();
|
||
|
||
let fused = derive_and_score(&grid, valid_time, false, None, None);
|
||
let (band_mhz, dense) = &fused.band_bodies[0];
|
||
|
||
let scattered: Vec<ScorePoint> = [0usize, 3]
|
||
.iter()
|
||
.map(|&cell| {
|
||
let (lat, lon) = grid.cell_latlon(cell);
|
||
ScorePoint {
|
||
lat,
|
||
lon,
|
||
score: dense[cell] as i32,
|
||
}
|
||
})
|
||
.collect();
|
||
|
||
let a = scores_file::encode_dense(*band_mhz, valid_time, dense, spec).unwrap();
|
||
let b = scores_file::encode_with_spec(*band_mhz, valid_time, &scattered, spec).unwrap();
|
||
assert_eq!(a, b, "dense indexing must match ScorePoint scatter");
|
||
}
|
||
|
||
/// A cell mask must keep masked-out cells at the no-data sentinel —
|
||
/// this is how HRDPS avoids writing over HRRR's CONUS coverage.
|
||
#[test]
|
||
fn cell_mask_suppresses_masked_cells() {
|
||
use chrono::TimeZone;
|
||
let (grid, spec) = small_grid();
|
||
let valid_time = Utc.with_ymd_and_hms(2026, 6, 15, 18, 0, 0).unwrap();
|
||
|
||
// Allow only cell 3; cell 0 is populated but masked out.
|
||
let mut mask = vec![false; spec.lon_count * spec.lat_count];
|
||
mask[3] = true;
|
||
|
||
let fused = derive_and_score(&grid, valid_time, false, Some(&mask), None);
|
||
|
||
assert_eq!(fused.cells_scored, 1);
|
||
let (_, body) = &fused.band_bodies[0];
|
||
assert_eq!(body[0], scores_file::NO_DATA, "masked cell must stay empty");
|
||
assert!(body[3] <= 100);
|
||
}
|
||
|
||
/// The fused pass runs cells in rayon chunks; scores must not depend
|
||
/// on how the work is split. Grid is deliberately larger than
|
||
/// `FUSE_CHUNK` so more than one chunk participates.
|
||
#[test]
|
||
fn fused_pass_is_independent_of_chunk_boundaries() {
|
||
use chrono::TimeZone;
|
||
let spec = GridSpec {
|
||
lon_start: -100.0,
|
||
lon_count: 100,
|
||
lon_step: 0.125,
|
||
lat_start: 30.0,
|
||
lat_count: 50, // 5000 cells > FUSE_CHUNK (2048)
|
||
lat_step: 0.125,
|
||
};
|
||
let n = spec.lon_count * spec.lat_count;
|
||
assert!(n > FUSE_CHUNK, "grid must span multiple chunks");
|
||
|
||
let mut g = FieldGrid::new(spec);
|
||
// Vary temperature per cell so a mis-indexed chunk shows up as a
|
||
// score mismatch rather than a uniform field hiding the bug.
|
||
let tmp: Vec<f32> = (0..n).map(|c| 280.0 + (c % 20) as f32).collect();
|
||
g.push_plane("TMP:2 m above ground", tmp);
|
||
g.push_plane("DPT:2 m above ground", vec![278.0; n]);
|
||
g.push_plane("PRES:surface", vec![101_000.0; n]);
|
||
|
||
let valid_time = Utc.with_ymd_and_hms(2026, 6, 15, 18, 0, 0).unwrap();
|
||
let fused = derive_and_score(&g, valid_time, false, None, None);
|
||
assert_eq!(fused.cells_scored, n as u32);
|
||
|
||
// Recompute serially and compare.
|
||
let planes = GridPlanes::resolve(&g);
|
||
let bands = band_config::all_bands();
|
||
let (_, body) = &fused.band_bodies[0];
|
||
let mut levels = Vec::new();
|
||
for (cell, &actual) in body.iter().enumerate().take(n) {
|
||
let (lat, lon) = g.cell_latlon(cell);
|
||
planes.levels_at(&g, cell, &mut levels);
|
||
let c = cell_to_conditions(&g, &planes, cell, lat, lon, &valid_time, &levels).unwrap();
|
||
let inv = scorer::precompute_band_invariants(&c);
|
||
let expected =
|
||
clamp_score_u8(scorer::composite_score_with(&c, &bands[0], Some(inv)).score);
|
||
assert_eq!(actual, expected, "cell {cell} disagrees");
|
||
}
|
||
}
|
||
|
||
/// The `.pgrid` record includes pre-computed surface_refractivity
|
||
/// and min_refractivity_gradient so Elixir consumers (Skew-T,
|
||
/// contact detail) have a guaranteed fallback when SoundingParams.derive
|
||
/// hits a level without dewpoint.
|
||
#[test]
|
||
fn pgrid_record_includes_refractivity_scalars() {
|
||
let grid = cell_grid(&[
|
||
("TMP:2 m above ground", 298.15),
|
||
("DPT:2 m above ground", 293.15),
|
||
("PRES:surface", 101_300.0),
|
||
("HPBL:surface", 1500.0),
|
||
(
|
||
"PWAT:entire atmosphere (considered as a single layer)",
|
||
30.0,
|
||
),
|
||
("TMP:1000 mb", 295.0),
|
||
("DPT:1000 mb", 290.0),
|
||
("HGT:1000 mb", 100.0),
|
||
("TMP:925 mb", 292.0),
|
||
("DPT:925 mb", 285.0),
|
||
("HGT:925 mb", 800.0),
|
||
("TMP:850 mb", 289.0),
|
||
("DPT:850 mb", 280.0),
|
||
("HGT:850 mb", 1500.0),
|
||
]);
|
||
|
||
let rec = pgrid_record_of(&grid);
|
||
assert!(
|
||
!rec["surface_refractivity"].is_nan(),
|
||
"surface_refractivity missing"
|
||
);
|
||
assert!(
|
||
!rec["min_refractivity_gradient"].is_nan(),
|
||
"min_refractivity_gradient missing"
|
||
);
|
||
// Surface scalars and the level triples must land in their slots.
|
||
assert!((rec["surface_temp_c"] - 25.0).abs() < 1e-3);
|
||
assert!((rec["surface_pressure_mb"] - 1013.0).abs() < 1e-3);
|
||
assert!((rec["hpbl_m"] - 1500.0).abs() < 1e-3);
|
||
assert!((rec["tmpc_850mb"] - (289.0 - 273.15)).abs() < 1e-3);
|
||
assert!((rec["dwpc_850mb"] - (280.0 - 273.15)).abs() < 1e-3);
|
||
assert!((rec["hght_m_850mb"] - 1500.0).abs() < 1e-3);
|
||
// A level the grid never supplied stays missing.
|
||
assert!(rec["tmpc_700mb"].is_nan());
|
||
}
|
||
|
||
/// When every pressure level is missing dewpoint, surface_refractivity
|
||
/// and min_refractivity_gradient are absent (NaN in the dense record,
|
||
/// which the Elixir reader surfaces as `nil`).
|
||
#[test]
|
||
fn pgrid_record_omits_refractivity_when_no_dewpoint() {
|
||
let grid = cell_grid(&[
|
||
("TMP:2 m above ground", 298.15),
|
||
("PRES:surface", 101_300.0),
|
||
("TMP:1000 mb", 295.0),
|
||
("HGT:1000 mb", 100.0),
|
||
("TMP:925 mb", 292.0),
|
||
("HGT:925 mb", 800.0),
|
||
]);
|
||
|
||
let rec = pgrid_record_of(&grid);
|
||
assert!(
|
||
rec["surface_refractivity"].is_nan(),
|
||
"surface_refractivity should be absent without dewpoint"
|
||
);
|
||
assert!(
|
||
rec["min_refractivity_gradient"].is_nan(),
|
||
"min_refractivity_gradient should be absent without dewpoint"
|
||
);
|
||
// The levels themselves still land — only the dewpoint is absent.
|
||
assert!((rec["tmpc_1000mb"] - (295.0 - 273.15)).abs() < 1e-3);
|
||
assert!(rec["dwpc_1000mb"].is_nan());
|
||
}
|
||
}
|
||
|
||
// =====================================================================
|
||
// Analysis pipeline (f00): fetch surface+pressure + native duct, merge
|
||
// with NEXRAD + commercial-link degradation, score all bands, write the
|
||
// ProfilesFile (MessagePack) and the band score files.
|
||
// =====================================================================
|
||
|
||
/// Query the latest Kp index once per cycle. Returns `None` when the
|
||
/// `geomagnetic_observations` table is empty (caught up in SWPC ingest
|
||
/// lag) — scoring continues without the aurora bonus, which is the
|
||
/// safe fallback for quiet geomagnetic conditions.
|
||
pub(crate) async fn fetch_latest_kp(pool: &sqlx::PgPool) -> Option<i32> {
|
||
sqlx::query_scalar::<_, i32>(
|
||
"SELECT kp_index FROM geomagnetic_observations ORDER BY valid_time DESC LIMIT 1",
|
||
)
|
||
.fetch_optional(pool)
|
||
.await
|
||
.ok()
|
||
.flatten()
|
||
}
|
||
|
||
#[derive(Debug, Clone)]
|
||
pub struct AnalysisStepInput {
|
||
pub run_time: DateTime<Utc>,
|
||
}
|
||
|
||
impl AnalysisStepInput {
|
||
pub fn valid_time(&self) -> DateTime<Utc> {
|
||
// f00 analysis valid_time == run_time.
|
||
self.run_time
|
||
}
|
||
}
|
||
|
||
#[derive(Debug, Clone)]
|
||
pub struct AnalysisStepStats {
|
||
pub score_files_written: u32,
|
||
pub point_count: u32,
|
||
pub band_count: u32,
|
||
pub profile_cells_written: u32,
|
||
pub duct_cells: u32,
|
||
pub nexrad_cells_with_echo: u32,
|
||
pub commercial_cells_boosted: u32,
|
||
}
|
||
|
||
#[tracing::instrument(
|
||
name = "pipeline.run_analysis_step",
|
||
skip(client, pool, scores_dir),
|
||
fields(run_time = %step.run_time)
|
||
)]
|
||
pub async fn run_analysis_step(
|
||
client: &HrrrClient,
|
||
pool: &sqlx::PgPool,
|
||
scores_dir: &Path,
|
||
step: &AnalysisStepInput,
|
||
) -> Result<AnalysisStepStats, PipelineError> {
|
||
let valid_time = step.valid_time();
|
||
let date = step.run_time.date_naive();
|
||
let hour = step.run_time.hour() as u8;
|
||
let grid_spec = wgrib2_grid_spec();
|
||
|
||
// Parallel fetch + decode of surface + pressure + native duct. All
|
||
// three are independent HRRR products; overlap the HTTP legs and
|
||
// the subsequent wgrib2 decodes so the long pole is a single
|
||
// decode, not three serial decodes.
|
||
let sfc_wanted = fetcher::surface_messages_owned();
|
||
let prs_wanted = fetcher::pressure_messages_grid();
|
||
let sfc_pattern = match_pattern(&sfc_wanted);
|
||
let prs_pattern = match_pattern(&prs_wanted);
|
||
|
||
// All three futures return different error types; lift into
|
||
// PipelineError inside each branch so try_join! stays happy.
|
||
let sfc_fut = async {
|
||
client
|
||
.fetch_product_blob(date, hour, Product::Surface, 0, &sfc_wanted)
|
||
.await
|
||
.map_err(PipelineError::Fetch)
|
||
};
|
||
let prs_fut = async {
|
||
client
|
||
.fetch_product_blob(date, hour, Product::Pressure, 0, &prs_wanted)
|
||
.await
|
||
.map_err(PipelineError::Fetch)
|
||
};
|
||
let duct_fut = async {
|
||
native_duct::fetch_native_duct_grid(client, date, hour, 0, grid_spec)
|
||
.await
|
||
.map_err(PipelineError::NativeDuct)
|
||
};
|
||
|
||
let fetch_started = std::time::Instant::now();
|
||
let (sfc_blob, prs_blob, duct_map) = tokio::try_join!(sfc_fut, prs_fut, duct_fut)?;
|
||
metrics::record_stage("fetch", fetch_started.elapsed());
|
||
|
||
let duct_count = duct_map.len() as u32;
|
||
|
||
let decode_started = std::time::Instant::now();
|
||
let sfc_grid = tokio::task::spawn_blocking(move || {
|
||
decoder::extract_grid(&sfc_blob, &sfc_pattern, grid_spec)
|
||
})
|
||
.await
|
||
.expect("blocking join")?;
|
||
|
||
let prs_grid = tokio::task::spawn_blocking(move || {
|
||
decoder::extract_grid(&prs_blob, &prs_pattern, grid_spec)
|
||
})
|
||
.await
|
||
.expect("blocking join")?;
|
||
metrics::record_stage("decode", decode_started.elapsed());
|
||
|
||
// Merge pressure into surface, then fold duct metrics in as extra
|
||
// planes. The duct values sit alongside the raw GRIB2 planes so
|
||
// build_grid_cache_rows can pick them up under their atom keys
|
||
// after the Elixir reader normalises.
|
||
let mut merged = sfc_grid;
|
||
merged.merge(prs_grid);
|
||
native_duct::merge_duct_grid(&mut merged, &duct_map);
|
||
|
||
// NEXRAD composite reflectivity overlay — only valid for f00
|
||
// because forecast hours can't see the future radar image.
|
||
let nexrad_http = reqwest::Client::builder()
|
||
.user_agent("prop-grid-rs/0.1")
|
||
.build()
|
||
.unwrap_or_else(|e| {
|
||
tracing::warn!(
|
||
"failed to build Nexrad HTTP client (missing root CA certs?): {} — \
|
||
f00 overlay will be skipped for this cycle",
|
||
e
|
||
);
|
||
// Return a client that will fail every request rather than
|
||
// panicking the worker. The f00 overlay is non-critical.
|
||
reqwest::Client::new()
|
||
});
|
||
let nexrad_points: Vec<(f64, f64)> = (0..merged.n_cells())
|
||
.map(|c| merged.cell_latlon(c))
|
||
.collect();
|
||
let nexrad_obs: Vec<NexradObservation> = match nexrad::fetch_frame(
|
||
&nexrad_http,
|
||
valid_time,
|
||
&nexrad_points,
|
||
)
|
||
.await
|
||
{
|
||
Ok(obs) => obs,
|
||
Err(e) => {
|
||
tracing::warn!(error = %e, "NEXRAD fetch failed — continuing without radar overlay");
|
||
Vec::new()
|
||
}
|
||
};
|
||
let nexrad_cells_with_echo = apply_nexrad(&mut merged, &nexrad_obs);
|
||
|
||
// Commercial-link degradation — precompute once, then fold into
|
||
// every in-range cell. Only ~7 links cluster around DFW so most
|
||
// cells stay untouched.
|
||
let commercial_lookup: Vec<LinkLookupEntry> = commercial::build_link_lookup(pool, valid_time)
|
||
.await
|
||
.map_err(PipelineError::Commercial)?;
|
||
let commercial_cells_boosted = apply_commercial(&mut merged, &commercial_lookup);
|
||
|
||
// Kp index — query once per cycle. SWPC ingest lag may leave the
|
||
// table empty, which degrades gracefully (no aurora boost). The
|
||
// fetch runs after decoding so Postgres can overlap wgrib2.
|
||
let kp = fetch_latest_kp(pool).await;
|
||
|
||
// One pass over the enriched grid producing scores, the profile
|
||
// records, and the scalar rows together.
|
||
let fused = tokio::task::spawn_blocking(move || {
|
||
let out = metrics::observe_stage("derive", || {
|
||
derive_and_score(&merged, valid_time, true, None, kp)
|
||
});
|
||
(out, merged.spec())
|
||
})
|
||
.await
|
||
.expect("blocking join");
|
||
let (fused, spec) = fused;
|
||
|
||
let point_count = fused.cells_scored;
|
||
let profile_cells_written = fused.cells_scored;
|
||
|
||
// Write the profile file on the blocking pool — encode + atomic
|
||
// rename overlap with the band-score writes.
|
||
let scores_dir_owned = scores_dir.to_path_buf();
|
||
let scores_dir_for_profiles = scores_dir_owned.clone();
|
||
let profile_future = {
|
||
let body = fused.pgrid_body;
|
||
tokio::task::spawn_blocking(move || {
|
||
metrics::observe_stage("write_profile", || {
|
||
pgrid::write_atomic(&scores_dir_for_profiles, valid_time, &spec, &body, false)
|
||
.map(|_| 0u32)
|
||
})
|
||
})
|
||
};
|
||
|
||
// Dense `.sgrid` body built from the scalar rows. Primary
|
||
// weather-scalar artifact; see run_chain_step for rationale.
|
||
let sgrid_body = sgrid::build_body(&spec, &fused.scalar_rows);
|
||
|
||
let scores_dir_for_sgrid = scores_dir_owned.clone();
|
||
let sgrid_future = {
|
||
tokio::task::spawn_blocking(move || {
|
||
metrics::observe_stage("write_sgrid", || {
|
||
sgrid::write_atomic(&scores_dir_for_sgrid, valid_time, &spec, &sgrid_body, false)
|
||
.map(|_| 0u32)
|
||
})
|
||
})
|
||
};
|
||
|
||
let band_count = fused.band_bodies.len() as u32;
|
||
let files_written = write_band_scores(
|
||
&scores_dir_owned,
|
||
valid_time,
|
||
fused.band_bodies,
|
||
spec,
|
||
false,
|
||
)
|
||
.await?;
|
||
|
||
profile_future
|
||
.await
|
||
.expect("blocking join")
|
||
.map_err(PipelineError::ProfileWrite)?;
|
||
sgrid_future
|
||
.await
|
||
.expect("blocking join")
|
||
.map_err(PipelineError::SgridWrite)?;
|
||
|
||
Ok(AnalysisStepStats {
|
||
score_files_written: files_written,
|
||
point_count,
|
||
band_count,
|
||
profile_cells_written,
|
||
duct_cells: duct_count,
|
||
nexrad_cells_with_echo,
|
||
commercial_cells_boosted,
|
||
})
|
||
}
|
||
|
||
fn apply_nexrad(grid: &mut FieldGrid, obs: &[NexradObservation]) -> u32 {
|
||
let mut with_echo = 0u32;
|
||
let mut plane = vec![f32::NAN; grid.n_cells()];
|
||
for o in obs {
|
||
// Observations come back on the same point list we sent, so each
|
||
// one resolves to exactly one cell by index.
|
||
let Some(cell) = grid.cell_index(o.lat, o.lon) else {
|
||
continue;
|
||
};
|
||
if o.max_reflectivity_dbz > 0.0 {
|
||
with_echo += 1;
|
||
}
|
||
plane[cell] = o.max_reflectivity_dbz as f32;
|
||
}
|
||
grid.push_plane("nexrad_max_reflectivity_dbz", plane);
|
||
with_echo
|
||
}
|
||
|
||
fn apply_commercial(grid: &mut FieldGrid, lookup: &[LinkLookupEntry]) -> u32 {
|
||
let n = grid.n_cells();
|
||
let mut degradation = vec![f32::NAN; n];
|
||
let mut baseline = vec![f32::NAN; n];
|
||
let mut current = vec![f32::NAN; n];
|
||
let mut n_links = vec![f32::NAN; n];
|
||
|
||
let mut boosted = 0u32;
|
||
for cell in 0..n {
|
||
let (lat, lon) = grid.cell_latlon(cell);
|
||
if let Some(d) = commercial::degradation_at((lat, lon), lookup) {
|
||
degradation[cell] = d.degradation_db as f32;
|
||
baseline[cell] = d.baseline_dbm as f32;
|
||
current[cell] = d.current_dbm as f32;
|
||
n_links[cell] = d.n_links as f32;
|
||
boosted += 1;
|
||
}
|
||
}
|
||
|
||
grid.push_plane("commercial_degradation_db", degradation);
|
||
grid.push_plane("commercial_baseline_dbm", baseline);
|
||
grid.push_plane("commercial_current_dbm", current);
|
||
grid.push_plane("commercial_n_links", n_links);
|
||
boosted
|
||
}
|