prop/rust/prop_grid_rs/src/planes.rs

//! Resolved plane ids for every field the derivation paths read.
//!
//! `FieldGrid` stores planes by `"VAR:LEVEL"` name, but hashing that
//! name once per cell per field is exactly the cost the struct-of-arrays
//! rewrite exists to remove. `GridPlanes` resolves the whole set **once
//! per grid**; after that every access is an array index.
//!
//! All three consumers share this table — `pipeline::cell_to_conditions`,
//! the profile writer, and `weather_scalar_file::derive_row` — so the
//! pressure-level walk that each of them used to redo independently
//! happens once per cell via [`GridPlanes::levels_at`].

use crate::fetcher;
use crate::field_grid::{FieldGrid, PlaneId};
use crate::sounding_params::Level;

/// Plane ids for one pressure level's three variables. `dpt` is
/// optional because HRRR occasionally publishes a level without
/// dewpoint, and the old code tolerated that via `Option`.
#[derive(Debug, Clone, Copy)]
pub struct LevelPlanes {
    pub pres_mb: f64,
    pub tmp: PlaneId,
    pub hgt: PlaneId,
    pub dpt: Option<PlaneId>,
}

#[derive(Debug, Clone)]
pub struct GridPlanes {
    // Surface / column fields.
    pub tmp_2m: Option<PlaneId>,
    pub dpt_2m: Option<PlaneId>,
    pub ugrd_10m: Option<PlaneId>,
    pub vgrd_10m: Option<PlaneId>,
    pub tcdc: Option<PlaneId>,
    pub pwat: Option<PlaneId>,
    pub pres_sfc: Option<PlaneId>,
    pub hpbl: Option<PlaneId>,
    pub apcp: Option<PlaneId>,

    // f00 enrichment planes.
    pub nexrad_dbz: Option<PlaneId>,
    pub native_min_gradient: Option<PlaneId>,
    pub best_duct_freq_ghz: Option<PlaneId>,
    pub max_duct_thickness_m: Option<PlaneId>,
    pub duct_count: Option<PlaneId>,
    pub commercial_degradation_db: Option<PlaneId>,
    pub commercial_baseline_dbm: Option<PlaneId>,
    pub commercial_current_dbm: Option<PlaneId>,
    pub commercial_n_links: Option<PlaneId>,

    /// One entry per `fetcher::GRID_PRESSURE_LEVELS` level that is
    /// actually present in the grid, in the same order.
    pub levels: Vec<LevelPlanes>,
}

impl GridPlanes {
    pub fn resolve(grid: &FieldGrid) -> Self {
        let levels = fetcher::grid_level_keys()
            .iter()
            .filter_map(|k| {
                // TMP and HGT are both required — the old per-cell
                // filter used `?` on each, so a level missing either
                // was skipped entirely.
                let tmp = grid.plane_id(&k.tmp)?;
                let hgt = grid.plane_id(&k.hgt)?;
                Some(LevelPlanes {
                    pres_mb: k.pres_mb,
                    tmp,
                    hgt,
                    dpt: grid.plane_id(&k.dpt),
                })
            })
            .collect();

        Self {
            tmp_2m: grid.plane_id("TMP:2 m above ground"),
            dpt_2m: grid.plane_id("DPT:2 m above ground"),
            ugrd_10m: grid.plane_id("UGRD:10 m above ground"),
            vgrd_10m: grid.plane_id("VGRD:10 m above ground"),
            tcdc: grid.plane_id("TCDC:entire atmosphere"),
            pwat: grid.plane_id("PWAT:entire atmosphere (considered as a single layer)"),
            pres_sfc: grid.plane_id("PRES:surface"),
            hpbl: grid.plane_id("HPBL:surface"),
            apcp: grid.plane_id("APCP:surface"),
            nexrad_dbz: grid.plane_id("nexrad_max_reflectivity_dbz"),
            native_min_gradient: grid.plane_id("native_min_gradient"),
            best_duct_freq_ghz: grid.plane_id("best_duct_freq_ghz"),
            max_duct_thickness_m: grid.plane_id("max_duct_thickness_m"),
            duct_count: grid.plane_id("duct_count"),
            commercial_degradation_db: grid.plane_id("commercial_degradation_db"),
            commercial_baseline_dbm: grid.plane_id("commercial_baseline_dbm"),
            commercial_current_dbm: grid.plane_id("commercial_current_dbm"),
            commercial_n_links: grid.plane_id("commercial_n_links"),
            levels,
        }
    }

    /// Fill `out` with this cell's pressure-level profile, ascending by
    /// pressure index exactly as `fetcher::grid_level_keys()` orders it.
    /// Levels whose TMP or HGT is missing at this cell are skipped, which
    /// matches the `filter_map` the old per-cell code used.
    ///
    /// Takes `&mut Vec` so callers reuse one buffer across all ~95 k
    /// cells instead of allocating a fresh `Vec<Level>` three times per
    /// cell (once each for conditions, profile, and scalars).
    #[inline]
    pub fn levels_at(&self, grid: &FieldGrid, cell: usize, out: &mut Vec<Level>) {
        out.clear();
        for lp in &self.levels {
            let (Some(t), Some(h)) = (grid.at(lp.tmp, cell), grid.at(lp.hgt, cell)) else {
                continue;
            };
            out.push(Level {
                pres_mb: lp.pres_mb,
                hght_m: h as f64,
                tmpc: (t as f64) - 273.15,
                dwpc: grid.at_opt(lp.dpt, cell).map(|v| (v as f64) - 273.15),
            });
        }
    }
}

/// Fill one `.pgrid` record (`pgrid::field_names()` order) for `cell`.
///
/// Writes straight into the caller's flat body slice, so the profile
/// artifact costs one `f32` store per field per cell — no `rmpv::Value`
/// tree, no per-level `String` keys, no compression pass. That is the
/// whole point of the format change: the old `.mp.gz` write measured
/// 3.96 s per CONUS grid against 0.039 s for the entire derivation.
pub fn fill_pgrid_record(
    grid: &FieldGrid,
    p: &GridPlanes,
    cell: usize,
    levels: &[Level],
    out: &mut [f32],
) {
    use crate::pgrid;

    debug_assert_eq!(out.len(), pgrid::n_fields());
    let idx = pgrid_scalar_indices();

    let mut put = |slot: usize, v: Option<f32>| {
        out[slot] = v.unwrap_or(f32::NAN);
    };

    put(
        idx.surface_temp_c,
        grid.at_opt(p.tmp_2m, cell).map(|v| v - 273.15),
    );
    put(
        idx.surface_dewpoint_c,
        grid.at_opt(p.dpt_2m, cell).map(|v| v - 273.15),
    );
    put(
        idx.surface_pressure_mb,
        grid.at_opt(p.pres_sfc, cell).map(|v| v / 100.0),
    );
    put(idx.hpbl_m, grid.at_opt(p.hpbl, cell));
    put(idx.pwat_mm, grid.at_opt(p.pwat, cell));
    put(idx.wind_u, grid.at_opt(p.ugrd_10m, cell));
    put(idx.wind_v, grid.at_opt(p.vgrd_10m, cell));
    put(idx.cloud_cover_pct, grid.at_opt(p.tcdc, cell));
    put(idx.precip_mm, grid.at_opt(p.apcp, cell));
    put(
        idx.native_min_gradient,
        grid.at_opt(p.native_min_gradient, cell),
    );
    put(
        idx.best_duct_freq_ghz,
        grid.at_opt(p.best_duct_freq_ghz, cell),
    );
    put(
        idx.max_duct_thickness_m,
        grid.at_opt(p.max_duct_thickness_m, cell),
    );
    put(idx.duct_count, grid.at_opt(p.duct_count, cell));
    put(
        idx.nexrad_max_reflectivity_dbz,
        grid.at_opt(p.nexrad_dbz, cell),
    );
    put(
        idx.commercial_degradation_db,
        grid.at_opt(p.commercial_degradation_db, cell),
    );
    put(
        idx.commercial_baseline_dbm,
        grid.at_opt(p.commercial_baseline_dbm, cell),
    );
    put(
        idx.commercial_current_dbm,
        grid.at_opt(p.commercial_current_dbm, cell),
    );
    put(
        idx.commercial_n_links,
        grid.at_opt(p.commercial_n_links, cell),
    );

    // Pre-computed refractivity scalars: the Skew-T page and contact
    // detail need a fallback for when SoundingParams.derive yields nil
    // because some level lacks dewpoint.
    put(
        idx.surface_refractivity,
        crate::sounding_params::surface_refractivity(levels).map(|v| v as f32),
    );
    put(
        idx.min_refractivity_gradient,
        crate::sounding_params::min_refractivity_gradient(levels.to_vec()).map(|v| v as f32),
    );

    // Pressure levels. `levels` is filtered (levels missing TMP/HGT at
    // this cell are absent), so match each back to its slot by pressure
    // rather than by position.
    for l in levels {
        let Some(base) = pgrid_level_base(l.pres_mb) else {
            continue;
        };
        out[base] = l.hght_m as f32;
        out[base + 1] = l.tmpc as f32;
        out[base + 2] = l.dwpc.map(|v| v as f32).unwrap_or(f32::NAN);
    }
}

/// Slot of `hght_m_<P>mb` for a pressure level, or `None` if that level
/// is not part of the on-disk schema.
fn pgrid_level_base(pres_mb: f64) -> Option<usize> {
    let p = pres_mb.round() as u16;
    let i = crate::fetcher::GRID_PRESSURE_LEVELS
        .iter()
        .position(|&lv| lv == p)?;
    Some(crate::pgrid::SCALAR_FIELDS.len() + i * 3)
}

/// Resolved slot for each scalar field, computed once.
struct PgridScalarIndices {
    surface_temp_c: usize,
    surface_dewpoint_c: usize,
    surface_pressure_mb: usize,
    hpbl_m: usize,
    pwat_mm: usize,
    wind_u: usize,
    wind_v: usize,
    cloud_cover_pct: usize,
    precip_mm: usize,
    native_min_gradient: usize,
    best_duct_freq_ghz: usize,
    max_duct_thickness_m: usize,
    duct_count: usize,
    nexrad_max_reflectivity_dbz: usize,
    commercial_degradation_db: usize,
    commercial_baseline_dbm: usize,
    commercial_current_dbm: usize,
    commercial_n_links: usize,
    surface_refractivity: usize,
    min_refractivity_gradient: usize,
}

fn pgrid_scalar_indices() -> &'static PgridScalarIndices {
    static IDX: std::sync::OnceLock<PgridScalarIndices> = std::sync::OnceLock::new();
    IDX.get_or_init(|| {
        let at = |name: &str| {
            crate::pgrid::SCALAR_FIELDS
                .iter()
                .position(|f| *f == name)
                .unwrap_or_else(|| panic!("pgrid scalar field {name} missing from SCALAR_FIELDS"))
        };
        PgridScalarIndices {
            surface_temp_c: at("surface_temp_c"),
            surface_dewpoint_c: at("surface_dewpoint_c"),
            surface_pressure_mb: at("surface_pressure_mb"),
            hpbl_m: at("hpbl_m"),
            pwat_mm: at("pwat_mm"),
            wind_u: at("wind_u"),
            wind_v: at("wind_v"),
            cloud_cover_pct: at("cloud_cover_pct"),
            precip_mm: at("precip_mm"),
            native_min_gradient: at("native_min_gradient"),
            best_duct_freq_ghz: at("best_duct_freq_ghz"),
            max_duct_thickness_m: at("max_duct_thickness_m"),
            duct_count: at("duct_count"),
            nexrad_max_reflectivity_dbz: at("nexrad_max_reflectivity_dbz"),
            commercial_degradation_db: at("commercial_degradation_db"),
            commercial_baseline_dbm: at("commercial_baseline_dbm"),
            commercial_current_dbm: at("commercial_current_dbm"),
            commercial_n_links: at("commercial_n_links"),
            surface_refractivity: at("surface_refractivity"),
            min_refractivity_gradient: at("min_refractivity_gradient"),
        }
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::grid::GridSpec;

    fn one_cell_spec() -> GridSpec {
        GridSpec {
            lon_start: -100.0,
            lon_count: 1,
            lon_step: 0.125,
            lat_start: 30.0,
            lat_count: 1,
            lat_step: 0.125,
        }
    }

    #[test]
    fn absent_fields_resolve_to_none() {
        let grid = FieldGrid::new(one_cell_spec());
        let p = GridPlanes::resolve(&grid);
        assert!(p.tmp_2m.is_none());
        assert!(p.apcp.is_none());
        assert!(p.levels.is_empty());
    }

    #[test]
    fn levels_at_skips_levels_missing_temp_or_height() {
        let mut grid = FieldGrid::new(one_cell_spec());
        let keys = fetcher::grid_level_keys();
        // Level 0 complete, level 1 has TMP but no HGT plane at all.
        grid.push_plane(&keys[0].tmp, vec![295.15]);
        grid.push_plane(&keys[0].hgt, vec![100.0]);
        grid.push_plane(&keys[0].dpt, vec![290.15]);
        grid.push_plane(&keys[1].tmp, vec![293.15]);

        let p = GridPlanes::resolve(&grid);
        let mut out = Vec::new();
        p.levels_at(&grid, 0, &mut out);

        assert_eq!(out.len(), 1, "level without HGT must be dropped");
        assert_eq!(out[0].pres_mb, keys[0].pres_mb);
        assert!((out[0].tmpc - 22.0).abs() < 1e-3);
        assert!((out[0].dwpc.unwrap() - 17.0).abs() < 1e-3);
    }

    #[test]
    fn levels_at_tolerates_missing_dewpoint_at_a_cell() {
        let mut grid = FieldGrid::new(one_cell_spec());
        let keys = fetcher::grid_level_keys();
        grid.push_plane(&keys[0].tmp, vec![295.15]);
        grid.push_plane(&keys[0].hgt, vec![100.0]);
        // DPT plane exists but this cell has no value.
        grid.push_plane(&keys[0].dpt, vec![f32::NAN]);

        let p = GridPlanes::resolve(&grid);
        let mut out = Vec::new();
        p.levels_at(&grid, 0, &mut out);

        assert_eq!(out.len(), 1);
        assert_eq!(out[0].dwpc, None);
    }

    #[test]
    fn levels_at_clears_the_buffer_between_cells() {
        let mut grid = FieldGrid::new(one_cell_spec());
        let keys = fetcher::grid_level_keys();
        grid.push_plane(&keys[0].tmp, vec![295.15]);
        grid.push_plane(&keys[0].hgt, vec![100.0]);
        let p = GridPlanes::resolve(&grid);

        let mut out = vec![Level {
            pres_mb: 1.0,
            hght_m: 1.0,
            tmpc: 1.0,
            dwpc: None,
        }];
        p.levels_at(&grid, 0, &mut out);
        assert_eq!(out.len(), 1);
        assert_eq!(out[0].pres_mb, keys[0].pres_mb);
    }
}