prop/rust/prop_grid_rs/tests/fused_pass_perf.rs

//! Wall-time smoke test for the fused derive+score pass on a
//! production-sized CONUS grid (473 × 201 = 95,073 cells, 48 planes).
//!
//! Ignored by default — it is a timing observation, not an assertion
//! about the machine it runs on. Run with:
//!
//! ```text
//! cargo test --release --test fused_pass_perf -- --ignored --nocapture
//! ```
//!
//! Context: this pass replaced three 95 k-cell derivation passes plus 23
//! band-major scoring passes over a staged
//! `Vec<(f64, f64, Conditions, BandInvariants)>`. The number printed here
//! is the one to watch when tuning `FUSE_CHUNK` or the scorer.

use chrono::{TimeZone, Utc};
use prop_grid_rs::field_grid::FieldGrid;
use prop_grid_rs::{fetcher, grid};

#[test]
#[ignore = "timing observation, not a pass/fail assertion"]
fn fused_pass_on_full_conus_grid() {
    let spec = grid::wgrib2_grid_spec();
    let n = spec.lat_count * spec.lon_count;
    assert_eq!(n, 95_073, "CONUS grid size changed");

    let mut g = FieldGrid::new(spec);

    // Surface planes, varied per cell so nothing folds to a constant.
    let ramp = |base: f32, span: f32| -> Vec<f32> {
        (0..n)
            .map(|c| base + ((c % 97) as f32 / 97.0) * span)
            .collect()
    };
    g.push_plane("TMP:2 m above ground", ramp(275.0, 30.0));
    g.push_plane("DPT:2 m above ground", ramp(270.0, 25.0));
    g.push_plane("UGRD:10 m above ground", ramp(-8.0, 16.0));
    g.push_plane("VGRD:10 m above ground", ramp(-8.0, 16.0));
    g.push_plane("TCDC:entire atmosphere", ramp(0.0, 100.0));
    g.push_plane("PRES:surface", ramp(97_000.0, 6_000.0));
    g.push_plane("HPBL:surface", ramp(100.0, 2_000.0));
    g.push_plane(
        "PWAT:entire atmosphere (considered as a single layer)",
        ramp(5.0, 45.0),
    );
    g.push_plane("APCP:surface", ramp(0.0, 3.0));

    // 13 pressure levels × 3 variables = the 39 pressure planes.
    for (i, k) in fetcher::grid_level_keys().iter().enumerate() {
        let dz = i as f32 * 250.0;
        g.push_plane(&k.tmp, ramp(275.0 - dz * 0.0065, 20.0));
        g.push_plane(&k.dpt, ramp(268.0 - dz * 0.0065, 18.0));
        g.push_plane(&k.hgt, ramp(100.0 + dz, 50.0));
    }
    assert_eq!(g.n_planes(), 48);

    let valid_time = Utc.with_ymd_and_hms(2026, 6, 15, 18, 0, 0).unwrap();

    let started = std::time::Instant::now();
    let out = prop_grid_rs::pipeline::derive_and_score_for_bench(&g, valid_time, true, None);
    let elapsed = started.elapsed();

    println!(
        "fused pass: {n} cells × {} planes × {} bands in {:.3} s ({} scored, {} scalars)",
        g.n_planes(),
        out.band_bodies.len(),
        elapsed.as_secs_f64(),
        out.cells_scored,
        out.scalar_rows.len(),
    );

    assert_eq!(out.cells_scored, n as u32);
    assert_eq!(out.band_bodies.len(), 23);
    for (_, body) in &out.band_bodies {
        assert_eq!(body.len(), n);
    }

    // Sanity: the pgrid body must actually carry the 13-level profile,
    // otherwise the timing above is measuring an empty structure.
    let nf = prop_grid_rs::pgrid::n_fields();
    assert_eq!(out.pgrid_body.len(), n * nf);
    let names = prop_grid_rs::pgrid::field_names();
    let t850 = names.iter().position(|f| f == "tmpc_850mb").unwrap();
    assert!(
        !out.pgrid_body[t850].is_nan(),
        "cell 0 should carry an 850 mb temperature"
    );

    // Now time the on-disk artifacts, which is what remains of the step.
    let dir = tempfile::tempdir().unwrap();

    let t = std::time::Instant::now();
    let path =
        prop_grid_rs::pgrid::write_atomic(dir.path(), valid_time, &spec, &out.pgrid_body, false)
            .unwrap();
    let profile_write = t.elapsed();
    let profile_bytes = std::fs::metadata(&path).unwrap().len();

    let t = std::time::Instant::now();
    prop_grid_rs::weather_scalar_file::write_atomic(dir.path(), valid_time, &out.scalar_rows)
        .unwrap();
    let scalar_write = t.elapsed();

    let t = std::time::Instant::now();
    for (band_mhz, body) in &out.band_bodies {
        prop_grid_rs::scores_file::write_atomic_dense(
            dir.path(),
            *band_mhz,
            valid_time,
            body,
            spec,
            false,
        )
        .unwrap();
    }
    let scores_write = t.elapsed();

    println!(
        "pgrid write       : {:.3} s ({:.1} MB on disk)",
        profile_write.as_secs_f64(),
        profile_bytes as f64 / 1e6
    );

    // Single-cell read: what /map point-detail and Skew-T actually do.
    // The old `.mp.gz` path gunzipped and unpacked the whole file for this.
    let raw = std::fs::read(&path).unwrap();
    let header = prop_grid_rs::pgrid::decode_header(&raw).expect("header");
    let t = std::time::Instant::now();
    let mut sink = 0.0f32;
    for cell in (0..n).step_by(1_000) {
        let rec = prop_grid_rs::pgrid::read_cell(&raw, &header, cell).unwrap();
        sink += rec[0];
    }
    let reads = (n / 1_000) as f64;
    println!(
        "pgrid read_cell   : {:.1} µs/cell over {reads} reads (checksum {sink:.1})",
        t.elapsed().as_secs_f64() * 1e6 / reads
    );
    println!("scalar_file  write: {:.3} s", scalar_write.as_secs_f64());
    println!("23 score files    : {:.3} s", scores_write.as_secs_f64());
}