prop/rust/prop_grid_rs/src/grid.rs

//! Grid definitions at 0.125° resolution. 1:1 port of
//! `lib/microwaveprop/propagation/grid.ex`. Two regions, disjoint by
//! construction:
//!
//!   * `conus_points()` — HRRR coverage (lat 25-50, lon -125 to -66).
//!   * `hrdps_only_points()` — Canadian extent covered by HRDPS but not
//!     HRRR (lat 49-60 minus the HRRR overlap, lon -141 to -52).
//!
//! Coverage stops at 60°N for v1 (SRTM elevation cuts off there; Arctic
//! CDEM coverage is deferred to a follow-up plan).

pub const LAT_MIN: f64 = 25.0;
pub const LAT_MAX: f64 = 50.0;
pub const LON_MIN: f64 = -125.0;
pub const LON_MAX: f64 = -66.0;
pub const STEP: f64 = 0.125;

pub const HRDPS_LAT_MIN: f64 = 49.0;
pub const HRDPS_LAT_MAX: f64 = 60.0;
pub const HRDPS_LON_MIN: f64 = -141.0;
pub const HRDPS_LON_MAX: f64 = -52.0;

// HRDPS uses a coarser grid than HRRR. Reason: each `wgrib2 -lon` point
// extraction redundantly JPEG2000-decodes every matched record, and the
// HRDPS rotated lat/lon source amplifies the per-record cost. At 0.125°
// (matching HRRR) production observed ~5 min/batch × 57 batches = ~5 h
// per chain step, far slower than the dev-bench claim of 30-90 s. Until
// the decoder is rewritten to decode-once + lookup-many, 0.5° (~55 km
// cells) keeps Canadian coverage visible on /weather without choking the
// pipeline. ~3.5 k cells fit in 4 batches → ~20 min/chain step.
pub const HRDPS_STEP: f64 = 0.5;

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct GridSpec {
    pub lon_start: f64,
    pub lon_count: usize,
    pub lon_step: f64,
    pub lat_start: f64,
    pub lat_count: usize,
    pub lat_step: f64,
}

pub fn wgrib2_grid_spec() -> GridSpec {
    let lon_count = ((LON_MAX - LON_MIN) / STEP).round() as usize + 1;
    let lat_count = ((LAT_MAX - LAT_MIN) / STEP).round() as usize + 1;
    GridSpec {
        lon_start: LON_MIN,
        lon_count,
        lon_step: STEP,
        lat_start: LAT_MIN,
        lat_count,
        lat_step: STEP,
    }
}

pub fn conus_points() -> Vec<(f64, f64)> {
    let spec = wgrib2_grid_spec();
    let mut out = Vec::with_capacity(spec.lat_count * spec.lon_count);
    for j in 0..spec.lat_count {
        let lat = round3(spec.lat_start + j as f64 * spec.lat_step);
        for i in 0..spec.lon_count {
            let lon = round3(spec.lon_start + i as f64 * spec.lon_step);
            out.push((lat, lon));
        }
    }
    out
}

/// HRDPS grid spec covering the full Canadian bbox. The `hrdps_only_points()`
/// list is a subset (excludes the HRRR overlap), but the underlying wgrib2
/// extraction needs the full bbox so the score-file's lat_start/lon_start
/// align with the cells the worker actually wrote.
pub fn hrdps_grid_spec() -> GridSpec {
    let lon_count = ((HRDPS_LON_MAX - HRDPS_LON_MIN) / HRDPS_STEP).round() as usize + 1;
    let lat_count = ((HRDPS_LAT_MAX - HRDPS_LAT_MIN) / HRDPS_STEP).round() as usize + 1;
    GridSpec {
        lon_start: HRDPS_LON_MIN,
        lon_count,
        lon_step: HRDPS_STEP,
        lat_start: HRDPS_LAT_MIN,
        lat_count,
        lat_step: HRDPS_STEP,
    }
}

/// Canadian-only grid points: cells inside the HRDPS bbox (49-60°N,
/// -141 to -52°W) but outside HRRR's CONUS bbox. Disjoint from
/// `conus_points()` by construction so the two grids never double-write
/// the same `(lat, lon)`.
pub fn hrdps_only_points() -> Vec<(f64, f64)> {
    let spec = hrdps_grid_spec();
    let mut out = Vec::with_capacity(spec.lat_count * spec.lon_count);
    for j in 0..spec.lat_count {
        let lat = round3(spec.lat_start + j as f64 * spec.lat_step);
        for i in 0..spec.lon_count {
            let lon = round3(spec.lon_start + i as f64 * spec.lon_step);
            // HRDPS now walks at 0.5° while HRRR walks at 0.125°, so the
            // coarse cells are no longer subset-aligned with HRRR's CONUS
            // grid. The disjointness rule is "if a coarse cell *center*
            // falls inside HRRR's lat/lon range, drop it" — HRRR will
            // cover that area at finer resolution anyway and the
            // /weather merge prefers HRRR rows on collision.
            if !in_conus_bbox(lat, lon) {
                out.push((lat, lon));
            }
        }
    }
    out
}

#[inline]
fn in_conus_bbox(lat: f64, lon: f64) -> bool {
    (LAT_MIN..=LAT_MAX).contains(&lat) && (LON_MIN..=LON_MAX).contains(&lon)
}

/// Matches Elixir's `Float.round(x, 3)` — banker's rounding isn't used,
/// half-away-from-zero is. For grid coords this matches BEAM's behavior
/// for the values we care about (no exact half cases).
pub fn round3(x: f64) -> f64 {
    (x * 1000.0).round() / 1000.0
}

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

    #[test]
    fn grid_spec_matches_elixir() {
        let spec = wgrib2_grid_spec();
        // Elixir: lon_count = round((-66 - -125) / 0.125) + 1 = 473
        //        lat_count = round((50 - 25) / 0.125) + 1 = 201
        assert_eq!(spec.lon_count, 473);
        assert_eq!(spec.lat_count, 201);
        assert_eq!(spec.lon_start, -125.0);
        assert_eq!(spec.lat_start, 25.0);
        assert_eq!(spec.lon_step, 0.125);
        assert_eq!(spec.lat_step, 0.125);
    }

    #[test]
    fn conus_points_count_matches_grid_spec() {
        let points = conus_points();
        assert_eq!(points.len(), 473 * 201);
    }

    #[test]
    fn conus_points_corner_values() {
        let points = conus_points();
        assert_eq!(points[0], (25.0, -125.0));
        assert_eq!(points[points.len() - 1], (50.0, -66.0));
    }

    #[test]
    fn hrdps_grid_spec_uses_coarse_step() {
        let spec = hrdps_grid_spec();
        // HRDPS_STEP = 0.5 (coarser than HRRR's 0.125) — see HRDPS_STEP
        // doc for the wgrib2 perf rationale.
        // lon_count = (-52 - -141) / 0.5 + 1 = 179
        // lat_count = (60 - 49) / 0.5 + 1 = 23
        assert_eq!(spec.lon_count, 179);
        assert_eq!(spec.lat_count, 23);
        assert_eq!(spec.lon_start, -141.0);
        assert_eq!(spec.lat_start, 49.0);
        assert_eq!(spec.lon_step, 0.5);
        assert_eq!(spec.lat_step, 0.5);
    }

    #[test]
    fn hrdps_only_points_are_disjoint_from_conus() {
        let conus: std::collections::HashSet<(u64, u64)> = conus_points()
            .into_iter()
            .map(|(la, lo)| (la.to_bits(), lo.to_bits()))
            .collect();

        for (la, lo) in hrdps_only_points() {
            assert!(
                !conus.contains(&(la.to_bits(), lo.to_bits())),
                "hrdps point {la},{lo} overlaps conus"
            );
        }
    }

    #[test]
    fn hrdps_only_points_all_within_bbox() {
        for (lat, lon) in hrdps_only_points() {
            assert!((49.0..=60.0).contains(&lat));
            assert!((-141.0..=-52.0).contains(&lon));
        }
    }

    #[test]
    fn hrdps_only_includes_cells_above_50n() {
        // At least one cell strictly north of HRRR's lat_max=50.
        let any_above_50 = hrdps_only_points().iter().any(|(la, _)| *la > 50.0);
        assert!(any_above_50);
    }
}