From 63f25a9612ecd0491371297692ed6a8f0f759ae5 Mon Sep 17 00:00:00 2001 From: Graham McIntire Date: Sat, 1 Aug 2026 08:23:36 -0500 Subject: [PATCH] perf(grid-rs): dense grid, fused scoring pass, and columnar .pgrid profiles MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reworks the post-fetch half of the propagation pipeline. Fetch and GRIB2 decode were already cheap — measured against a live HRRR cycle, all 39 pressure messages decode via `wgrib2 -lola` in 0.29 s and the 31 MB byte-range fetch takes ~3 s — so nothing here touches the decoder. All the cost was downstream. Also fixes a broken NOTIFY that made every chain step run up to 5 times. pg_notify `NOTIFY propagation_ready, $1` is a Postgres syntax error: NOTIFY is a utility statement whose payload must be a literal, so a bind raises 42601. It shared a transaction with the `status='done'` UPDATE, so every successful step rolled back, stayed 'running', and was requeued by reclaim_stale_running up to @max_reclaim_attempts times. Elixir's NotifyListener never fired either, so ScoreCache warm and the "propagation:updated" fan-out were dead. FieldGrid A decoded grid was HashMap<(i32,i32), HashMap, f32>> — a dense rectangular grid stored as ~95k nested hash maps, costing ~4.6M inserts on decode, ~3.7M on merge and ~14M lookups across three derivation passes. wgrib2 -lola already emits one dense row-major f32 block per message, so keep it: dense per-message planes, names hashed once per grid into plane ids, NaN as the missing sentinel. This is what forced PROP_GRID_RS_PARALLELISM=1 under a 3Gi limit. Fused pass Three 95k-cell derivation passes plus 23 band-major scoring passes over a staged Vec<(f64,f64,Conditions,BandInvariants)> (~19MB re-streamed 23 times) collapse into one pass: levels extracted once per cell, all 23 bands scored while the cell is hot, scores accumulated cell-major so rayon chunks own disjoint slices. Scores land straight in the dense score-file body — no ScorePoint scatter. .pgrid The profile artifact was an rmpv tree plus gzip -9, written 30x an hour, and ProfilesFile.read_point/3 gunzipped and unpacked the entire 95k-cell file to return one cell on every map click and Skew-T load. Replaced with a dense cell-major f32 record array carrying a self-describing field table. Elixir reads it via :file.pread; .mp.gz and .etf.gz remain readable so files written before this drain out of the 48h window. Measured on a full CONUS grid (95,073 cells x 48 planes x 23 bands): derive + score + build artifacts 0.022 s profile write 3.957 s -> 0.006 s (22.0 MB -> 22.4 MB on disk) single-cell read whole-file decode -> 0.5 us 23 score files 0.003 s Also - hrrr_points: batched UNNEST upsert replacing one awaited INSERT per point. Keeps ON CONFLICT DO UPDATE — the PSKR sampler's two-pass loop depends on it. - fetcher: real semaphore capping in-flight ranges at MAX_PARALLEL_RANGES, which the comment claimed but the code did not do (it spawned all 27 while the connection pool was sized for 8). - metrics: per-stage histogram. Only chain-step and decode durations were instrumented, which is why the write cost stayed invisible. - profiles_file: parse_valid_time anchors on the known extension set, so sibling-suffixed names like .hrdps.prop no longer parse as .hrdps and vanish from prune and list operations. - PROP_GRID_RS_PARALLELISM 1 -> 3. Memory limit held at 3Gi until RSS is observed at the new parallelism. - cargo fmt over the crate; worker.rs, hrdps_fetcher.rs and nexrad.rs were already unformatted at HEAD and the pre-commit hook gates on it. HRDPS still runs at 0.5 degrees. wgrib2 -lola scales linearly in output points on rotated lat/lon (12.5 s wall, 202 s CPU for one message at 0.125 degrees) because it has no inverse projection for those grids; a raw native dump is 0.32 s. The fix is decode-once plus a closed-form rotated-pole index, left for a follow-up. --- k8s/deployment-grid-rs.yaml | 35 +- lib/microwaveprop/propagation/pgrid.ex | 427 +++++ .../propagation/profiles_file.ex | 48 +- rust/prop_grid_rs/src/bin/worker.rs | 4 +- rust/prop_grid_rs/src/db.rs | 87 +- rust/prop_grid_rs/src/decoder.rs | 214 +-- rust/prop_grid_rs/src/fetcher.rs | 15 +- rust/prop_grid_rs/src/field_grid.rs | 370 +++++ rust/prop_grid_rs/src/hrdps_fetcher.rs | 8 +- rust/prop_grid_rs/src/hrrr_points.rs | 375 +++-- rust/prop_grid_rs/src/lib.rs | 4 +- rust/prop_grid_rs/src/metrics.rs | 52 +- rust/prop_grid_rs/src/native_duct.rs | 294 ++-- rust/prop_grid_rs/src/nexrad.rs | 7 +- rust/prop_grid_rs/src/pgrid.rs | 453 ++++++ rust/prop_grid_rs/src/pipeline.rs | 1420 ++++++++++------- rust/prop_grid_rs/src/planes.rs | 372 +++++ rust/prop_grid_rs/src/profiles_file.rs | 265 --- rust/prop_grid_rs/src/scores_file.rs | 68 +- rust/prop_grid_rs/src/weather_scalar_file.rs | 136 +- rust/prop_grid_rs/tests/fused_pass_perf.rs | 141 ++ test/microwaveprop/propagation/pgrid_test.exs | 323 ++++ 22 files changed, 3805 insertions(+), 1313 deletions(-) create mode 100644 lib/microwaveprop/propagation/pgrid.ex create mode 100644 rust/prop_grid_rs/src/field_grid.rs create mode 100644 rust/prop_grid_rs/src/pgrid.rs create mode 100644 rust/prop_grid_rs/src/planes.rs delete mode 100644 rust/prop_grid_rs/src/profiles_file.rs create mode 100644 rust/prop_grid_rs/tests/fused_pass_perf.rs create mode 100644 test/microwaveprop/propagation/pgrid_test.exs diff --git a/k8s/deployment-grid-rs.yaml b/k8s/deployment-grid-rs.yaml index f957d25e..0625cf99 100644 --- a/k8s/deployment-grid-rs.yaml +++ b/k8s/deployment-grid-rs.yaml @@ -85,21 +85,23 @@ spec: value: "/data/hrrr_idx" - name: RUST_LOG value: "info" - # Single concurrent task per pod. 2 replicas → 2 concurrent - # cluster-wide. Raised from 2 after the new 3 Gi limit still - # OOMed under 2-parallel forecast load — the NFS-backed - # score-file writes (23 files × ~2 MB each) keep cgroup page - # cache high enough that two concurrent tasks plus the - # wgrib2 working set regularly crosses 3 Gi. Single-lane - # parallelism stays comfortably under 2 Gi and lets the - # analysis step (native duct + NEXRAD + commercial merges) - # run safely within the same budget. + # Was "1": a decoded grid used to be ~95 k nested HashMaps + # (~200-400 MB per in-flight task), so two concurrent tasks + # plus the wgrib2 working set crossed the 3 Gi limit. The + # dense `FieldGrid` representation stores the same data as + # ~48 f32 planes (~18 MB), and the profile artifact is now a + # flat f32 write instead of an rmpv tree + gzip -9, so the + # per-task footprint is roughly an order of magnitude lower. + # + # Raising to 3 first and holding the memory limit; drop the + # limit only after observing steady-state RSS at this + # parallelism (one knob per deploy). - name: PROP_GRID_RS_PARALLELISM - value: "1" + value: "3" # Pool size = parallelism + 2 (NOTIFY + retry slack). Kept in # sync with PROP_GRID_RS_PARALLELISM when either changes. - name: PROP_GRID_RS_PG_CONNS - value: "3" + value: "5" - name: METRICS_ADDR value: "0.0.0.0:9100" envFrom: @@ -135,11 +137,12 @@ spec: limits: cpu: "4" # Analysis tasks (f00) run the native-level HRRR decode - # (~530 MB GRIB2 + wgrib2 working set) on top of 2 concurrent - # forecast tasks. 2 Gi OOM'd reliably in the 13 h after - # Stream A cutover. 3 Gi leaves room for the analysis - # wgrib2 peak without eliminating the forecast-lane - # parallelism headroom. + # (~530 MB GRIB2 + wgrib2 working set), which is now the + # dominant term — the per-task grid footprint dropped from + # ~200-400 MB of nested HashMaps to ~18 MB of dense f32 + # planes. Holding 3 Gi for one deploy while parallelism + # goes 1 -> 3; once `container_memory_working_set_bytes` + # is confirmed steady, this can come down toward 1.5 Gi. memory: 3Gi securityContext: allowPrivilegeEscalation: false diff --git a/lib/microwaveprop/propagation/pgrid.ex b/lib/microwaveprop/propagation/pgrid.ex new file mode 100644 index 00000000..3abd091b --- /dev/null +++ b/lib/microwaveprop/propagation/pgrid.ex @@ -0,0 +1,427 @@ +defmodule Microwaveprop.Propagation.Pgrid do + @moduledoc """ + Reader for the `.pgrid` profile format written by `prop-grid-rs` + (`rust/prop_grid_rs/src/pgrid.rs`). + + Replaces the gzipped-MessagePack `.mp.gz` artifact. The format is a + dense, **cell-major** `f32` record array on a fixed grid, so: + + * reading one cell is a single `:file.pread/3` of `n_fields * 4` + bytes at a computed offset — no decompression, no term building. + The `.mp.gz` path gunzipped, `Msgpax.unpack`ed and atomized the + *entire* 95k-cell file to answer one point lookup, on every `/map` + click and every Skew-T load. + * reading a viewport is one contiguous `pread` per grid row. + + Layout (little-endian) — see the Rust module for the authoritative + description: + + magic 4 "PGRD" + version 1 0x01 + flags 1 bit0: 0 = hrrr, 1 = hrdps + n_fields 2 u16 + valid_time 8 i64 unix seconds + lat_start 8 f64 + lon_start 8 f64 + lat_step 8 f64 + lon_step 8 f64 + n_rows 2 u16 + n_cols 2 u16 + field_table n_fields * 32 (NUL-padded ASCII) + body n_rows*n_cols*n_fields * 4 f32, cell-major + + `NaN` is the missing-value sentinel. Erlang's float binary match + rejects NaN, so `decode_f32/1` maps it to `nil` — which is exactly + what the `.mp.gz` reader produced for an absent key. + + Callers get the same map shape `ProfilesFile.read_point/3` always + returned: atom keys, a `:profile` list of per-level maps, and a + nested `:commercial_link_degradation` map. + """ + + alias Microwaveprop.Propagation.Grid + + @magic "PGRD" + @version 1 + @field_name_len 32 + @fixed_header_len 4 + 1 + 1 + 2 + 8 + 8 + 8 + 8 + 8 + 2 + 2 + + # On-disk field name => the atom key callers expect. + # + # Written as literal atoms so they exist at compile time. The `.mp.gz` + # reader used `String.to_existing_atom/1` behind a whitelist, which + # only worked as long as some *other* module happened to have mentioned + # each atom — `:wind_v` in particular did not, so that approach raised + # at runtime. An explicit map removes the hidden dependency, and a + # field the reader doesn't know (e.g. one added later on the Rust side) + # is simply skipped instead of crashing or growing the atom table. + @field_atoms %{ + "surface_temp_c" => :surface_temp_c, + "surface_dewpoint_c" => :surface_dewpoint_c, + "surface_pressure_mb" => :surface_pressure_mb, + "hpbl_m" => :hpbl_m, + "pwat_mm" => :pwat_mm, + "wind_u" => :wind_u, + "wind_v" => :wind_v, + "cloud_cover_pct" => :cloud_cover_pct, + "precip_mm" => :precip_mm, + "native_min_gradient" => :native_min_gradient, + "best_duct_freq_ghz" => :best_duct_freq_ghz, + "max_duct_thickness_m" => :max_duct_thickness_m, + "duct_count" => :duct_count, + "nexrad_max_reflectivity_dbz" => :nexrad_max_reflectivity_dbz, + "surface_refractivity" => :surface_refractivity, + "min_refractivity_gradient" => :min_refractivity_gradient + } + + # Deliberately absent from @field_atoms: the commercial_* fields are + # folded into the nested :commercial_link_degradation map, and the + # per-level hght_m_/tmpc_/dwpc_ triples are rebuilt into :profile. + + # Surfaced as an integer, as the msgpack writer did. + @integer_fields ~w(duct_count) + + defmodule Header do + @moduledoc "Parsed `.pgrid` header." + + @type t :: %__MODULE__{} + + defstruct [ + :hrdps?, + :valid_time, + :n_fields, + :fields, + :field_index, + :lat_start, + :lon_start, + :lat_step, + :lon_step, + :n_rows, + :n_cols, + :body_offset, + :levels + ] + end + + @doc "Base directory the profile store lives under." + @spec base_dir() :: String.t() + def base_dir do + Path.join( + Application.get_env(:microwaveprop, :propagation_scores_dir, "/data/scores"), + "profiles" + ) + end + + @doc "Absolute path for the `.pgrid` file covering `valid_time`." + @spec path_for(DateTime.t()) :: String.t() + def path_for(%DateTime{} = valid_time) do + iso = valid_time |> DateTime.truncate(:second) |> DateTime.to_iso8601() + Path.join(base_dir(), "#{iso}.pgrid") + end + + @doc "Whether a `.pgrid` exists for `valid_time`." + @spec exists?(DateTime.t()) :: boolean() + def exists?(%DateTime{} = valid_time), do: File.exists?(path_for(valid_time)) + + @doc """ + Read a single cell's profile for `(valid_time, lat, lon)`. + + Opens the file, `pread`s the header and then one record. Returns `nil` + when the file is missing or the point falls outside the grid. + """ + @spec read_point(DateTime.t(), number(), number()) :: map() | nil + def read_point(%DateTime{} = valid_time, lat, lon) do + path = path_for(valid_time) + + with {:ok, fd} <- :file.open(path, [:read, :binary, :raw]), + {:ok, header} <- read_header(fd) do + result = read_cell(fd, header, lat, lon) + :file.close(fd) + result + else + _ -> nil + end + end + + defp read_cell(fd, header, lat, lon) do + with cell when is_integer(cell) <- cell_index(header, lat, lon), + size = header.n_fields * 4, + {:ok, bin} <- :file.pread(fd, record_offset(header, cell), size), + true <- byte_size(bin) == size do + to_profile(header, bin) + else + _ -> nil + end + end + + @doc """ + Read every populated cell as `%{{lat, lon} => profile}`. + + Used for the cold `GridCache` warm in `Microwaveprop.Weather.Grid`. + Cells with no surface temperature are skipped, so the result matches + what the `.mp.gz` reader produced (it only ever contained scored cells). + """ + @spec read(DateTime.t()) :: {:ok, %{{float(), float()} => map()}} | {:error, :enoent} + def read(%DateTime{} = valid_time) do + path = path_for(valid_time) + + with {:ok, raw} <- File.read(path), + {:ok, header} <- parse_header(raw) do + record_bytes = header.n_fields * 4 + + grid = + raw + |> binary_part(header.body_offset, byte_size(raw) - header.body_offset) + |> chunk_records(record_bytes) + |> Enum.with_index() + |> Enum.flat_map(&populated_cell(header, &1)) + |> Map.new() + + {:ok, grid} + else + _ -> {:error, :enoent} + end + end + + @doc """ + Every `valid_time` with a `.pgrid` on disk, sorted ascending. + """ + @spec list_valid_times() :: [DateTime.t()] + def list_valid_times do + case File.ls(base_dir()) do + {:ok, names} -> + names + |> Enum.filter(&String.ends_with?(&1, ".pgrid")) + |> Enum.map(&(&1 |> String.replace_suffix(".pgrid", "") |> parse_iso())) + |> Enum.reject(&is_nil/1) + |> Enum.sort(DateTime) + + _ -> + [] + end + end + + defp populated_cell(header, {bin, cell}) do + case to_profile(header, bin) do + nil -> [] + profile -> [{cell_latlon(header, cell), profile}] + end + end + + defp parse_iso(str) do + case DateTime.from_iso8601(str) do + {:ok, dt, _} -> dt + _ -> nil + end + end + + # ── Header ───────────────────────────────────────────────────────── + + defp read_header(fd) do + with {:ok, fixed} <- :file.pread(fd, 0, @fixed_header_len), + {:ok, n_fields} <- peek_n_fields(fixed), + {:ok, table} <- :file.pread(fd, @fixed_header_len, n_fields * @field_name_len) do + parse_header(fixed <> table) + else + _ -> :error + end + end + + defp peek_n_fields(<<@magic, @version, _flags::8, n_fields::little-16, _rest::binary>>), do: {:ok, n_fields} + + defp peek_n_fields(_), do: :error + + @doc false + @spec parse_header(binary()) :: {:ok, Header.t()} | :error + def parse_header( + <<@magic, @version, flags::8, n_fields::little-16, valid_unix::little-signed-64, lat_start::little-float-64, + lon_start::little-float-64, lat_step::little-float-64, lon_step::little-float-64, n_rows::little-16, + n_cols::little-16, rest::binary>> + ) do + table_len = n_fields * @field_name_len + + if byte_size(rest) < table_len do + :error + else + fields = + rest + |> binary_part(0, table_len) + |> chunk_records(@field_name_len) + |> Enum.map(&trim_nul/1) + + field_index = fields |> Enum.with_index() |> Map.new() + + {:ok, + %Header{ + hrdps?: Bitwise.band(flags, 1) == 1, + valid_time: DateTime.from_unix!(valid_unix), + n_fields: n_fields, + fields: fields, + field_index: field_index, + lat_start: lat_start, + lon_start: lon_start, + lat_step: lat_step, + lon_step: lon_step, + n_rows: n_rows, + n_cols: n_cols, + body_offset: @fixed_header_len + table_len, + levels: level_slots(fields, field_index) + }} + end + end + + def parse_header(_), do: :error + + defp trim_nul(bin) do + case :binary.match(bin, <<0>>) do + {pos, _} -> binary_part(bin, 0, pos) + :nomatch -> bin + end + end + + # Pressure levels present in the field table, as + # {pres_mb, hght_idx, tmpc_idx, dwpc_idx}, ordered by the table. + defp level_slots(fields, field_index) do + fields + |> Enum.flat_map(fn name -> + case Regex.run(~r/^hght_m_(\d+)mb$/, name) do + [_, p] -> [String.to_integer(p)] + _ -> [] + end + end) + |> Enum.map(fn p -> + {p, field_index["hght_m_#{p}mb"], field_index["tmpc_#{p}mb"], field_index["dwpc_#{p}mb"]} + end) + end + + # ── Cell addressing ──────────────────────────────────────────────── + + defp record_offset(header, cell), do: header.body_offset + cell * header.n_fields * 4 + + @doc false + @spec cell_index(Header.t(), number(), number()) :: non_neg_integer() | nil + def cell_index(header, lat, lon) do + row = round((lat - header.lat_start) / header.lat_step) + col = round((lon - header.lon_start) / header.lon_step) + + if row >= 0 and col >= 0 and row < header.n_rows and col < header.n_cols do + row * header.n_cols + col + end + end + + defp cell_latlon(header, cell) do + row = div(cell, header.n_cols) + col = rem(cell, header.n_cols) + + {Float.round(header.lat_start + row * header.lat_step, 3), Float.round(header.lon_start + col * header.lon_step, 3)} + end + + # ── Record decoding ──────────────────────────────────────────────── + + # Returns nil for a cell with no surface temperature — the `.mp.gz` + # file only ever contained scored cells, so callers already treat a + # missing cell as "no data here". + defp to_profile(header, bin) do + values = decode_values(bin) + + case at(values, header, "surface_temp_c") do + nil -> + nil + + _ -> + header.fields + |> Enum.with_index() + |> Enum.reduce(%{}, fn {name, idx}, acc -> + put_field(acc, name, Enum.at(values, idx)) + end) + |> put_commercial(values, header) + |> put_levels(values, header) + end + end + + defp at(values, header, name) do + case Map.fetch(header.field_index, name) do + {:ok, idx} -> Enum.at(values, idx) + :error -> nil + end + end + + defp put_field(acc, _name, nil), do: acc + + defp put_field(acc, name, value) do + case Map.fetch(@field_atoms, name) do + # Unknown field (level triple, commercial_*, or one added on the + # Rust side after this reader was written) — handled elsewhere or + # deliberately ignored. + :error -> acc + {:ok, key} when name in @integer_fields -> Map.put(acc, key, trunc(value)) + {:ok, key} -> Map.put(acc, key, value) + end + end + + defp put_commercial(acc, values, header) do + case at(values, header, "commercial_degradation_db") do + nil -> + acc + + degradation -> + Map.put(acc, :commercial_link_degradation, %{ + degradation_db: degradation, + baseline_dbm: at(values, header, "commercial_baseline_dbm") || 0.0, + current_dbm: at(values, header, "commercial_current_dbm") || 0.0, + n_links: trunc(at(values, header, "commercial_n_links") || 0) + }) + end + end + + defp put_levels(acc, values, header) do + levels = Enum.flat_map(header.levels, &build_level(&1, values)) + + if levels == [], do: acc, else: Map.put(acc, :profile, levels) + end + + defp build_level({pres_mb, hght_idx, tmpc_idx, dwpc_idx}, values) do + do_build_level( + pres_mb, + Enum.at(values, hght_idx), + Enum.at(values, tmpc_idx), + dwpc_idx && Enum.at(values, dwpc_idx) + ) + end + + # TMP and HGT are both required, matching the Rust-side filter; a level + # missing either is dropped rather than surfaced half-populated. + defp do_build_level(_pres_mb, nil, _tmpc, _dwpc), do: [] + defp do_build_level(_pres_mb, _hght, nil, _dwpc), do: [] + + defp do_build_level(pres_mb, hght, tmpc, nil), do: [%{pres_mb: pres_mb * 1.0, hght_m: hght, tmpc: tmpc}] + + defp do_build_level(pres_mb, hght, tmpc, dwpc), do: [%{pres_mb: pres_mb * 1.0, hght_m: hght, tmpc: tmpc, dwpc: dwpc}] + + defp decode_values(bin), do: bin |> chunk_records(4) |> Enum.map(&decode_f32/1) + + # NaN is the missing-value sentinel. Erlang's float binary match + # rejects NaN and Inf outright, so the fallback clause is what turns + # "absent" into nil. + defp decode_f32(<>), do: v + defp decode_f32(_), do: nil + + defp chunk_records(bin, size) when byte_size(bin) >= size do + for <>, do: chunk + end + + defp chunk_records(_bin, _size), do: [] + + @doc """ + Snap a lat/lon to the propagation grid step. Kept in sync with + `Microwaveprop.Propagation.ProfilesFile.snap/2` so a coordinate written + by either side is reachable from the other. + """ + @spec snap(number(), number()) :: {float(), float()} + def snap(lat, lon), do: {snap_one(lat), snap_one(lon)} + + defp snap_one(coord) do + step = Grid.step() + Float.round(Float.round(coord / step) * step, 3) + end +end diff --git a/lib/microwaveprop/propagation/profiles_file.ex b/lib/microwaveprop/propagation/profiles_file.ex index e3eed74f..e7d2ce3d 100644 --- a/lib/microwaveprop/propagation/profiles_file.ex +++ b/lib/microwaveprop/propagation/profiles_file.ex @@ -20,6 +20,7 @@ defmodule Microwaveprop.Propagation.ProfilesFile do """ alias Microwaveprop.Propagation.Grid + alias Microwaveprop.Propagation.Pgrid @doc "Base directory the profile store lives under." @spec base_dir() :: String.t() @@ -144,11 +145,13 @@ defmodule Microwaveprop.Propagation.ProfilesFile do end defp do_read(valid_time) do - # Rust-written MessagePack wins over Elixir-written ETF if both - # exist. After Phase 3 Stream A cutover only `.mp.gz` is produced; - # the ETF branch remains so pods picking up an older in-flight run - # during the rollout window still render correctly. + # Preference order is newest-format-first. Rust now writes only + # `.pgrid` (dense cell-major f32, see `Microwaveprop.Propagation.Pgrid`); + # the `.mp.gz` and `.etf.gz` branches remain so pods reading files + # written before the cutover still render while those drain out of + # the 48 h retention window. cond do + Pgrid.exists?(valid_time) -> Pgrid.read(valid_time) File.exists?(mp_path_for(valid_time)) -> read_mp(valid_time) File.exists?(path_for(valid_time)) -> read_etf(valid_time) true -> {:error, :enoent} @@ -217,13 +220,20 @@ defmodule Microwaveprop.Propagation.ProfilesFile do """ @spec read_point(DateTime.t(), float(), float()) :: map() | nil def read_point(%DateTime{} = valid_time, lat, lon) do - case read(valid_time) do - {:ok, grid_data} -> - {snapped_lat, snapped_lon} = snap(lat, lon) - Map.get(grid_data, {snapped_lat, snapped_lon}) + {snapped_lat, snapped_lon} = snap(lat, lon) - {:error, _} -> - nil + # `.pgrid` supports true random access: one `pread` of the cell's + # record. Do NOT route this through `read/1` — that would decode the + # whole grid to answer a single point lookup, which is exactly the + # cost this format exists to remove. The legacy branches have no + # random access, so they still go through the cached full read. + if Pgrid.exists?(valid_time) do + Pgrid.read_point(valid_time, snapped_lat, snapped_lon) + else + case read(valid_time) do + {:ok, grid_data} -> Map.get(grid_data, {snapped_lat, snapped_lon}) + {:error, _} -> nil + end end end @@ -356,12 +366,18 @@ defmodule Microwaveprop.Propagation.ProfilesFile do end defp parse_valid_time(filename) do - # Matches both `.etf.gz` (Elixir legacy) and `.mp.gz` (Rust Phase 3 - # Stream A). If both exist for the same valid_time the directory - # listing yields two entries with the same timestamp; the pipeline - # downstream of list_valid_times/0 uniq-sorts, and read/1 prefers - # the mp.gz variant via mp_path_for/1. - with [_, iso, _fmt] <- Regex.run(~r/^(.+)\.(etf|mp)\.gz$/, filename), + # Matches `.pgrid` (current), `.mp.gz` (Rust Phase 3 Stream A) and + # `.etf.gz` (Elixir legacy). If several exist for the same valid_time + # the directory listing yields multiple entries with the same + # timestamp; the pipeline downstream of list_valid_times/0 uniq-sorts, + # and read/1 prefers the newest format. + # + # Anchoring on the *known* extension set (rather than a permissive + # `\.(.+)$`) is what keeps sibling-suffixed files such as + # `.hrdps.prop` from parsing as `.hrdps` and then failing + # `DateTime.from_iso8601/1` — the failure mode that made HRDPS score + # files invisible to every prune and list operation. + with [_, iso] <- Regex.run(~r/^(.+)\.(?:pgrid|etf\.gz|mp\.gz)$/, filename), {:ok, dt, _} <- DateTime.from_iso8601(iso) do DateTime.to_unix(dt) else diff --git a/rust/prop_grid_rs/src/bin/worker.rs b/rust/prop_grid_rs/src/bin/worker.rs index 6da60724..784649d3 100644 --- a/rust/prop_grid_rs/src/bin/worker.rs +++ b/rust/prop_grid_rs/src/bin/worker.rs @@ -16,7 +16,9 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; -use prop_grid_rs::{db, fetcher::HrrrClient, hrdps_fetcher::HrdpsClient, metrics, pipeline, telemetry}; +use prop_grid_rs::{ + db, fetcher::HrrrClient, hrdps_fetcher::HrdpsClient, metrics, pipeline, telemetry, +}; use tokio::task::JoinSet; use tracing::{error, info, warn}; diff --git a/rust/prop_grid_rs/src/db.rs b/rust/prop_grid_rs/src/db.rs index f9dc858e..f0cc55b3 100644 --- a/rust/prop_grid_rs/src/db.rs +++ b/rust/prop_grid_rs/src/db.rs @@ -266,9 +266,16 @@ pub async fn complete(pool: &PgPool, task: &ClaimedTask) -> Result<(), DbError> .await?; let payload = task.valid_time.format("%Y-%m-%dT%H:%M:%SZ").to_string(); - // Use parameterized NOTIFY to avoid SQL injection vectors and keep - // sqlx's query logging / tracing consistent. - sqlx::query("NOTIFY propagation_ready, $1") + // `pg_notify(text, text)`, NOT `NOTIFY chan, $1`. NOTIFY is a utility + // statement whose payload must be a literal — binding a parameter to it + // is a 42601 syntax error, which aborted this transaction and silently + // rolled back the status='done' UPDATE above. The row then sat at + // 'running' until `GridTaskEnqueuer.reclaim_stale_running/1` requeued it, + // so every chain step was recomputed up to @max_reclaim_attempts times + // and Elixir's NotifyListener never fired. Covered by + // `complete_emits_propagation_ready_notification`. + sqlx::query("SELECT pg_notify($1, $2)") + .bind(NOTIFY_CHANNEL) .bind(&payload) .execute(&mut *tx) .await?; @@ -469,6 +476,21 @@ mod tests { complete(&pool, &claimed).await.unwrap(); + // `complete` commits the 'done' UPDATE and the NOTIFY in one + // transaction. Asserting only that `complete` returned Ok is not + // enough — the previous `NOTIFY chan, $1` form was a Postgres + // syntax error, which aborted the transaction and silently left + // the row stuck at 'running' for `reclaim_stale_running` to + // requeue. Assert the committed state. + let status: String = sqlx::query("SELECT status FROM grid_tasks WHERE id = $1") + .bind(claimed.id) + .fetch_one(&pool) + .await + .unwrap() + .try_get("status") + .unwrap(); + assert_eq!(status, "done", "complete/2 must commit status='done'"); + // Clean up. sqlx::query("DELETE FROM grid_tasks WHERE id = $1") .bind(claimed.id) @@ -476,4 +498,63 @@ mod tests { .await .unwrap(); } + + /// The NOTIFY payload is what wakes Elixir's `PropagationNotifyListener` + /// to warm `ScoreCache` and fan out `"propagation:updated"`. Listen on a + /// dedicated connection and assert the payload actually lands. + #[tokio::test] + async fn complete_emits_propagation_ready_notification() { + let Ok(url) = std::env::var("PROP_GRID_RS_TEST_DB") else { + eprintln!("skip: set PROP_GRID_RS_TEST_DB to run db tests"); + return; + }; + let pool = connect(&url, 4).await.unwrap(); + + let mut listener = sqlx::postgres::PgListener::connect(&url).await.unwrap(); + listener.listen(NOTIFY_CHANNEL).await.unwrap(); + + let run_time = Utc::now(); + let fh: i16 = 9; + let valid_time = run_time + chrono::Duration::hours(fh as i64); + let id = Uuid::new_v4(); + sqlx::query( + r#"INSERT INTO grid_tasks (id, run_time, forecast_hour, valid_time, status, attempt, inserted_at, updated_at) + VALUES ($1, $2, $3, $4, 'queued', 0, NOW(), NOW())"#, + ) + .bind(id) + .bind(run_time.naive_utc()) + .bind(fh as i32) + .bind(valid_time.naive_utc()) + .execute(&pool) + .await + .unwrap(); + + let claimed = claim_next(&pool).await.unwrap().expect("had a row"); + complete(&pool, &claimed).await.unwrap(); + + // Other tests in this binary run concurrently and complete their own + // tasks on the same channel, so the first notification to arrive is + // not necessarily ours. Drain until we see our own valid_time. + let expected = claimed.valid_time.format("%Y-%m-%dT%H:%M:%SZ").to_string(); + let found = tokio::time::timeout(Duration::from_secs(5), async { + loop { + let n = listener.recv().await.expect("listener stream alive"); + assert_eq!(n.channel(), NOTIFY_CHANNEL); + if n.payload() == expected { + return; + } + } + }) + .await; + assert!( + found.is_ok(), + "no propagation_ready NOTIFY with payload {expected} within 5s" + ); + + sqlx::query("DELETE FROM grid_tasks WHERE id = $1") + .bind(claimed.id) + .execute(&pool) + .await + .unwrap(); + } } diff --git a/rust/prop_grid_rs/src/decoder.rs b/rust/prop_grid_rs/src/decoder.rs index 55ab9607..5038b58a 100644 --- a/rust/prop_grid_rs/src/decoder.rs +++ b/rust/prop_grid_rs/src/decoder.rs @@ -15,7 +15,8 @@ use std::process::{Command, Output}; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; -use crate::grid::{round3, GridSpec}; +use crate::field_grid::FieldGrid; +use crate::grid::GridSpec; const UNDEFINED_VALUE: f32 = 9.999e20; static UNIQUE: AtomicU64 = AtomicU64::new(1); @@ -36,23 +37,10 @@ pub struct Message { pub level: String, } -/// Per-cell: `"VAR:LEVEL" -> f32`. Key format matches Elixir exactly so -/// scorer input can be wired unchanged. -/// -/// Keys are `Arc` rather than `String`: the hot decode loop -/// inserts the same `"VAR:LEVEL"` string into ~92k cell maps per -/// GRIB2 message, which was previously a String::clone per cell -/// (~5.5M heap allocations + frees per chain step across 60 -/// messages). `Arc::clone` is a single atomic increment — keys -/// allocate once per message and are ref-counted thereafter. -/// -/// Lookup is unchanged: `cell.get("TMP:2 m above ground")` works -/// because `Arc: Borrow`. -pub type CellValues = HashMap, f32>; - -/// Lat/lon rounded to 3 decimals. f64 keys can't be hashed directly, so -/// we key by `(lat_mdeg, lon_mdeg)` where `*_mdeg = round(x * 1000) as i32`. -pub type PointGrid = HashMap<(i32, i32), CellValues>; +// The former `CellValues` / `PointGrid` aliases (a HashMap of ~95 k +// per-cell HashMaps) are gone — see `field_grid::FieldGrid`, which +// stores the same data as dense per-message planes. Decoded grids are +// `FieldGrid` everywhere now. pub fn wgrib2_available() -> bool { which_wgrib2().is_some() @@ -105,7 +93,8 @@ pub fn extract_points( grib: &[u8], match_pattern: &str, points: &[(f64, f64)], -) -> Result { + grid_spec: GridSpec, +) -> Result { let wgrib2 = which_wgrib2().ok_or(DecodeError::Wgrib2NotAvailable)?; let nanos = SystemTime::now() .duration_since(UNIX_EPOCH) @@ -115,7 +104,7 @@ pub fn extract_points( let grib_path: PathBuf = std::env::temp_dir().join(format!("hrdps_{nanos}_{uniq}.grib2")); std::fs::write(&grib_path, grib)?; - let result = run_wgrib2_lon_batched(&wgrib2, &grib_path, match_pattern, points); + let result = run_wgrib2_lon_batched(&wgrib2, &grib_path, match_pattern, points, grid_spec); let _ = std::fs::remove_file(&grib_path); result } @@ -137,26 +126,39 @@ fn run_wgrib2_lon_batched( grib_path: &Path, match_pattern: &str, points: &[(f64, f64)], -) -> Result { - let mut out: PointGrid = HashMap::with_capacity(points.len()); + grid_spec: GridSpec, +) -> Result { + let mut out = FieldGrid::new(grid_spec); for batch in points.chunks(POINT_BATCH) { - let batch_grid = run_wgrib2_lon_one(wgrib2, grib_path, match_pattern, batch)?; - for (key, cell) in batch_grid { - out.entry(key).or_default().extend(cell); + let batch_values = run_wgrib2_lon_one(wgrib2, grib_path, match_pattern, batch, &grid_spec)?; + for (name, cells) in batch_values { + let (_, plane) = out.plane_mut_or_insert(&name); + for (cell, value) in cells { + plane[cell] = value; + } } } Ok(out) } +/// One wgrib2 `-lon` invocation, keyed `"VAR:LEVEL" -> [(cell_index, value)]`. +/// Values are gathered per-name so the caller can scatter each name into a +/// single plane instead of touching a per-cell map. +type LonBatchValues = HashMap>; + fn run_wgrib2_lon_one( wgrib2: &Path, grib_path: &Path, match_pattern: &str, batch: &[(f64, f64)], -) -> Result { + grid_spec: &GridSpec, +) -> Result { let mut cmd = Command::new(wgrib2); - cmd.arg(grib_path).arg("-s").arg("-match").arg(match_pattern); + cmd.arg(grib_path) + .arg("-s") + .arg("-match") + .arg(match_pattern); for &(lat, lon) in batch { cmd.arg("-lon").arg(format!("{lon}")).arg(format!("{lat}")); } @@ -188,7 +190,7 @@ fn run_wgrib2_lon_one( } let text = String::from_utf8_lossy(&stdout); - Ok(parse_lon_output(&text, batch)) + Ok(parse_lon_output(&text, batch, grid_spec)) } /// Parse `wgrib2 -s -lon` text output. Each record line is a colon- @@ -197,8 +199,8 @@ fn run_wgrib2_lon_one( /// /// Example line: /// `1:0:d=...:TMP:2 m above ground:anl:lon=280.369,lat=43.670,val=280.97:lon=...` -fn parse_lon_output(text: &str, batch: &[(f64, f64)]) -> PointGrid { - let mut out: PointGrid = HashMap::new(); +fn parse_lon_output(text: &str, batch: &[(f64, f64)], grid_spec: &GridSpec) -> LonBatchValues { + let mut out: LonBatchValues = HashMap::new(); for line in text.lines() { if line.is_empty() { @@ -209,23 +211,23 @@ fn parse_lon_output(text: &str, batch: &[(f64, f64)]) -> PointGrid { let _offset = parts.next(); let _date = parts.next(); let var = match parts.next() { - Some(v) => v.to_string(), + Some(v) => v, None => continue, }; let level = match parts.next() { - Some(l) => l.to_string(), + Some(l) => l, None => continue, }; // Skip the record-type/forecast-time field (anl, "240 min fcst", …). let _kind = parts.next(); - let key: std::sync::Arc = format!("{var}:{level}").into(); + let key = format!("{var}:{level}"); + let slot = out.entry(key).or_default(); for segment in parts { - if let Some(value) = parse_lon_segment(segment, batch) { - let (lat_key, lon_key, val) = value; - out.entry((lat_key, lon_key)) - .or_default() - .insert(std::sync::Arc::clone(&key), val); + if let Some((lat, lon, val)) = parse_lon_segment(segment, batch) { + if let Some(cell) = crate::field_grid::cell_index_for(grid_spec, lat, lon) { + slot.push((cell, val)); + } } } } @@ -233,10 +235,10 @@ fn parse_lon_output(text: &str, batch: &[(f64, f64)]) -> PointGrid { out } -/// `lon=242.958,lat=32.938,val=306.5` → (lat_key, lon_key, val). -/// Snaps to the nearest point in `batch` (handles wgrib2's lon-180/360 -/// convention by checking both signs). -fn parse_lon_segment(segment: &str, batch: &[(f64, f64)]) -> Option<(i32, i32, f32)> { +/// `lon=242.958,lat=32.938,val=306.5` → (lat, lon, val) snapped to the +/// nearest point in `batch` (handles wgrib2's lon-180/360 convention by +/// denormalising before matching). +fn parse_lon_segment(segment: &str, batch: &[(f64, f64)]) -> Option<(f64, f64, f32)> { let mut lon_part = None; let mut lat_part = None; let mut val_part = None; @@ -259,10 +261,8 @@ fn parse_lon_segment(segment: &str, batch: &[(f64, f64)]) -> Option<(i32, i32, f } let target_lon = denormalize_lon(raw_lon); - let snapped = snap_to_batch(raw_lat, target_lon, batch)?; - let lat_key = (snapped.0 * 1000.0).round() as i32; - let lon_key = (snapped.1 * 1000.0).round() as i32; - Some((lat_key, lon_key, val as f32)) + let (lat, lon) = snap_to_batch(raw_lat, target_lon, batch)?; + Some((lat, lon, val as f32)) } fn snap_to_batch(lat: f64, lon: f64, batch: &[(f64, f64)]) -> Option<(f64, f64)> { @@ -297,7 +297,7 @@ pub fn extract_grid( grib: &[u8], match_pattern: &str, grid_spec: GridSpec, -) -> Result { +) -> Result { let wgrib2 = which_wgrib2().ok_or(DecodeError::Wgrib2NotAvailable)?; let tmp_dir = std::env::temp_dir(); let uniq = UNIQUE.fetch_add(1, Ordering::Relaxed); @@ -308,7 +308,7 @@ pub fn extract_grid( let grib_path = tmp_dir.join(format!("hrrr_{nanos}_{uniq}.grib2")); let bin_path = tmp_dir.join(format!("hrrr_{nanos}_{uniq}.grib2.lola.bin")); - let result: Result = (|| { + let result: Result = (|| { std::fs::write(&grib_path, grib)?; run_wgrib2_lola(&wgrib2, &grib_path, match_pattern, &grid_spec, &bin_path) })(); @@ -323,7 +323,7 @@ pub fn extract_grid_from_file( grib_path: &Path, match_pattern: &str, grid_spec: GridSpec, -) -> Result { +) -> Result { let wgrib2 = which_wgrib2().ok_or(DecodeError::Wgrib2NotAvailable)?; let bin_path = sibling_tmp_path(grib_path, "lola"); let result = run_wgrib2_lola(&wgrib2, grib_path, match_pattern, &grid_spec, &bin_path); @@ -348,7 +348,7 @@ fn run_wgrib2_lola( match_pattern: &str, grid_spec: &GridSpec, bin_path: &Path, -) -> Result { +) -> Result { let lon_spec = format!( "{}:{}:{}", normalize_lon(grid_spec.lon_start), @@ -391,7 +391,9 @@ fn run_wgrib2_lola( match std::fs::read(bin_path) { Ok(bin) => Ok(parse_lola_binary(&bin, &messages, grid_spec)), - Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(HashMap::new()), + // wgrib2 matched nothing, so it never created the output file. + // An empty grid (right spec, zero planes) is the correct result. + Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(FieldGrid::new(*grid_spec)), Err(e) => Err(e.into()), } } @@ -418,16 +420,23 @@ pub fn parse_inventory(out: &str) -> Vec { /// Parse the Fortran-unformatted output file from `wgrib2 -lola`. /// Layout: per message = 4-byte LE u32 record length + N·4 bytes f32 LE /// data + 4-byte LE u32 record length. `N = nx * ny`. Cells whose value -/// exceeds `UNDEFINED_VALUE / 2` are treated as missing (the sentinel -/// wgrib2 writes for "no data"). -pub fn parse_lola_binary(bin: &[u8], messages: &[Message], grid_spec: &GridSpec) -> PointGrid { +/// exceeds `UNDEFINED_VALUE / 2` are replaced with `NaN` (the sentinel +/// wgrib2 writes for "no data"; `FieldGrid::at` maps NaN back to `None`, +/// matching the old behaviour of omitting the key entirely). +/// +/// wgrib2 writes each record row-major over (lat, lon) in exactly the +/// order `FieldGrid` indexes cells, so the body is a bulk +/// little-endian `f32` decode per message — no per-cell coordinate +/// arithmetic, no hashing, no allocation beyond the plane itself. +pub fn parse_lola_binary(bin: &[u8], messages: &[Message], grid_spec: &GridSpec) -> FieldGrid { let nx = grid_spec.lon_count; let ny = grid_spec.lat_count; - let bytes_per_msg = nx * ny * 4; + let n_cells = nx * ny; + let bytes_per_msg = n_cells * 4; let record_overhead = 8; let stride = bytes_per_msg + record_overhead; - let mut out: PointGrid = HashMap::with_capacity(nx * ny); + let mut out = FieldGrid::new(*grid_spec); for (msg_idx, msg) in messages.iter().enumerate() { let data_offset = msg_idx * stride + 4; @@ -435,32 +444,17 @@ pub fn parse_lola_binary(bin: &[u8], messages: &[Message], grid_spec: &GridSpec) continue; } let chunk = &bin[data_offset..data_offset + bytes_per_msg]; - // Allocate the "VAR:LEVEL" string once per message. Every - // subsequent cell insertion is Arc::clone (one atomic inc), - // not a String::clone (heap alloc + copy + free). - let key: std::sync::Arc = format!("{}:{}", msg.var, msg.level).into(); - - for j in 0..ny { - let lat = round3(grid_spec.lat_start + j as f64 * grid_spec.lat_step); - let lat_key = (lat * 1000.0).round() as i32; - for i in 0..nx { - let cell_offset = (j * nx + i) * 4; - // Safety: the chunk stride guarantees each 4-byte window - // is fully contained within chunk. The try_into().unwrap() - // is safe here because the slice length is statically 4. - let value = - f32::from_le_bytes(chunk[cell_offset..cell_offset + 4].try_into().unwrap()); - if value > UNDEFINED_VALUE / 2.0 { - continue; - } - let raw_lon = grid_spec.lon_start + i as f64 * grid_spec.lon_step; - let lon = round3(denormalize_lon(raw_lon)); - let lon_key = (lon * 1000.0).round() as i32; - out.entry((lat_key, lon_key)) - .or_default() - .insert(std::sync::Arc::clone(&key), value); + let mut plane = Vec::with_capacity(n_cells); + plane.extend(chunk.chunks_exact(4).map(|b| { + // try_into on a chunks_exact(4) window is statically length 4. + let v = f32::from_le_bytes(b.try_into().unwrap()); + if v > UNDEFINED_VALUE / 2.0 { + f32::NAN + } else { + v } - } + })); + out.push_plane(&format!("{}:{}", msg.var, msg.level), plane); } out @@ -511,14 +505,31 @@ mod tests { lon=280.369,lat=43.670,val=275.00:\ lon=235.000,lat=49.190,val=270.00\n"; let batch = vec![(43.670, -79.631), (49.190, -125.000)]; + // Spec chosen so both batch points land exactly on a cell. + let spec = GridSpec { + lon_start: -125.0, + lon_count: 400, + lon_step: 0.125, + lat_start: 43.0, + lat_count: 60, + lat_step: 0.125, + }; - let grid = parse_lon_output(text, &batch); + let values = parse_lon_output(text, &batch, &spec); - let toronto_key = ((43.670_f64 * 1000.0).round() as i32, (-79.631_f64 * 1000.0).round() as i32); - let toronto = grid.get(&toronto_key).expect("toronto cell"); - let tmp = toronto.get("TMP:2 m above ground").copied().expect("tmp present"); + let toronto = crate::field_grid::cell_index_for(&spec, 43.670, -79.631) + .expect("toronto lands on the grid"); + let tmp = values["TMP:2 m above ground"] + .iter() + .find(|(c, _)| *c == toronto) + .expect("tmp present") + .1; assert!((tmp - 280.97).abs() < 0.01); - let dpt = toronto.get("DPT:2 m above ground").copied().expect("dpt present"); + let dpt = values["DPT:2 m above ground"] + .iter() + .find(|(c, _)| *c == toronto) + .expect("dpt present") + .1; assert!((dpt - 275.00).abs() < 0.01); } @@ -573,19 +584,20 @@ mod tests { } let grid = parse_lola_binary(&buf, &msgs, &spec); - assert_eq!(grid.len(), 6); + assert_eq!(grid.n_cells(), 6); + assert_eq!(grid.n_planes(), 2); + let tmp = grid.plane_id("TMP:surface").unwrap(); + let dpt = grid.plane_id("DPT:surface").unwrap(); // Cell (30.0, -100.0) → i=0, j=0, msg0 value = 0.0, msg1 = 100.0 - let k = ((30.0 * 1000.0) as i32, (-100.0 * 1000.0) as i32); - let cell = grid.get(&k).unwrap(); - assert_eq!(cell["TMP:surface"], 0.0); - assert_eq!(cell["DPT:surface"], 100.0); + let k = grid.cell_index(30.0, -100.0).unwrap(); + assert_eq!(grid.at(tmp, k), Some(0.0)); + assert_eq!(grid.at(dpt, k), Some(100.0)); // Cell (30.5, -99.0) → i=2, j=1, msg0 = 12.0 - let k2 = ((30.5 * 1000.0) as i32, (-99.0 * 1000.0) as i32); - let cell2 = grid.get(&k2).unwrap(); - assert_eq!(cell2["TMP:surface"], 12.0); - assert_eq!(cell2["DPT:surface"], 112.0); + let k2 = grid.cell_index(30.5, -99.0).unwrap(); + assert_eq!(grid.at(tmp, k2), Some(12.0)); + assert_eq!(grid.at(dpt, k2), Some(112.0)); } #[test] @@ -610,9 +622,15 @@ mod tests { buf.extend_from_slice(&len.to_le_bytes()); let grid = parse_lola_binary(&buf, &msgs, &spec); - assert_eq!(grid.len(), 1); // only i=1 survived - let k = ((30.0 * 1000.0) as i32, (-99.0 * 1000.0) as i32); - assert_eq!(grid.get(&k).unwrap()["TMP:surface"], 42.0); + let tmp = grid.plane_id("TMP:surface").unwrap(); + // i=0 carried the sentinel and must read back as missing; only + // i=1 survived. The old representation expressed this by never + // inserting the key for i=0. + assert_eq!(grid.at(tmp, grid.cell_index(30.0, -100.0).unwrap()), None); + assert_eq!( + grid.at(tmp, grid.cell_index(30.0, -99.0).unwrap()), + Some(42.0) + ); } #[test] diff --git a/rust/prop_grid_rs/src/fetcher.rs b/rust/prop_grid_rs/src/fetcher.rs index 224c25d3..5b725bd8 100644 --- a/rust/prop_grid_rs/src/fetcher.rs +++ b/rust/prop_grid_rs/src/fetcher.rs @@ -435,18 +435,31 @@ impl HrrrClient { } let merged = merge_ranges(ranges); + // Actually cap in-flight requests at MAX_PARALLEL_RANGES. This + // used to spawn every merged range at once (27 for a pressure + // fetch) despite the comment below claiming a cap, while + // `pool_max_idle_per_host` was sized for 8 — so the surplus + // connections were opened and torn down on every fetch, and a + // wide fetch could draw an S3 429. + let permits = std::sync::Arc::new(tokio::sync::Semaphore::new(MAX_PARALLEL_RANGES)); + let mut futs: JoinSet), FetchError>> = JoinSet::new(); let url = url.to_string(); for r in merged.iter().copied() { let client = self.http.clone(); let url = url.clone(); + let permits = permits.clone(); futs.spawn(async move { + // Semaphore is never closed, so acquire cannot fail. + let _permit = permits + .acquire_owned() + .await + .expect("range semaphore is never closed"); let bytes = retry_get_range(&client, &url, r).await?; Ok((r.start, bytes)) }); } - // Cap concurrent in-flight requests at MAX_PARALLEL_RANGES. let mut results: Vec<(u64, Vec)> = Vec::with_capacity(merged.len()); while let Some(item) = futs.join_next().await { // If a range-download task panics (JoinError), log it and diff --git a/rust/prop_grid_rs/src/field_grid.rs b/rust/prop_grid_rs/src/field_grid.rs new file mode 100644 index 00000000..655f9fbc --- /dev/null +++ b/rust/prop_grid_rs/src/field_grid.rs @@ -0,0 +1,370 @@ +//! Dense struct-of-arrays storage for a decoded NWP grid. +//! +//! Replaces the previous `HashMap<(i32, i32), HashMap, f32>>` +//! representation. The old shape stored a *dense rectangular* grid as +//! ~95 k nested hash maps: decoding one forecast hour cost ~4.6 M +//! string-keyed inserts, merging surface into pressure cost another +//! ~3.7 M, and the three downstream derivation passes +//! (`cell_to_conditions`, the profile writer, the scalar writer) cost +//! ~14 M lookups between them. That is what forced the worker to +//! `PROP_GRID_RS_PARALLELISM=1` under a 3 GiB limit. +//! +//! `wgrib2 -lola` already hands us exactly the right shape: one dense +//! row-major `f32` block per GRIB2 message on a fixed grid. So we keep +//! it that way. +//! +//! ## Layout +//! +//! * `planes` is plane-major: value for `(plane p, cell c)` lives at +//! `planes[p * n_cells + c]`. +//! * Cell index is **row-major over (lat, lon)**: `cell = row * lon_count + col`, +//! where `row` indexes latitude ascending from `spec.lat_start` and +//! `col` indexes longitude ascending from `spec.lon_start`. This is +//! the same ordering `wgrib2 -lola` writes and the same ordering +//! `scores_file` encodes, so the score-grid write is a straight copy. +//! * `f32::NAN` is the missing-value sentinel. The old representation +//! signalled "missing" by omitting the key; `at()` returns `None` for +//! NaN, so callers using `?` behave identically. +//! +//! Plane names are the same `"VAR:LEVEL"` strings the decoder produced +//! before (`"TMP:2 m above ground"`), so match patterns, level-key +//! tables (`fetcher::grid_level_keys`) and the golden fixtures are +//! unchanged. The difference is that the name is hashed **once per +//! grid** to resolve a plane id, not once per cell per access. + +use std::collections::HashMap; +use std::sync::Arc; + +use crate::grid::{round3, GridSpec}; + +/// Index of a plane (one GRIB2 message, or one derived field) within a +/// [`FieldGrid`]. Resolve once via [`FieldGrid::plane_id`], then use it +/// for every cell. +pub type PlaneId = u32; + +/// Nearest cell index on `spec` for a geographic coordinate, or `None` +/// when the point falls outside the grid. Free function so callers that +/// only have a `GridSpec` (the decoder, before a grid exists) can share +/// the rounding rules with [`FieldGrid::cell_index`]. +pub fn cell_index_for(spec: &GridSpec, lat: f64, lon: f64) -> Option { + let row_f = ((lat - spec.lat_start) / spec.lat_step).round(); + let col_f = ((lon - spec.lon_start) / spec.lon_step).round(); + if !row_f.is_finite() || !col_f.is_finite() { + return None; + } + let (row, col) = (row_f as isize, col_f as isize); + if row < 0 || col < 0 { + return None; + } + let (row, col) = (row as usize, col as usize); + if row >= spec.lat_count || col >= spec.lon_count { + return None; + } + Some(row * spec.lon_count + col) +} + +#[derive(Debug, Clone)] +pub struct FieldGrid { + spec: GridSpec, + n_cells: usize, + names: Vec>, + index: HashMap, PlaneId>, + planes: Vec, +} + +impl FieldGrid { + pub fn new(spec: GridSpec) -> Self { + let n_cells = spec.lat_count * spec.lon_count; + Self { + spec, + n_cells, + names: Vec::new(), + index: HashMap::new(), + planes: Vec::new(), + } + } + + pub fn spec(&self) -> GridSpec { + self.spec + } + + #[inline] + pub fn n_cells(&self) -> usize { + self.n_cells + } + + #[inline] + pub fn n_planes(&self) -> usize { + self.names.len() + } + + pub fn plane_names(&self) -> &[Arc] { + &self.names + } + + /// Resolve a `"VAR:LEVEL"` name to its plane id. Call once per grid, + /// not once per cell. + #[inline] + pub fn plane_id(&self, name: &str) -> Option { + self.index.get(name).copied() + } + + /// Value at `(plane, cell)`, or `None` when missing (NaN). + /// + /// `debug_assert`s the bounds rather than checking them in release: + /// every caller derives `plane` from `plane_id` and `cell` from + /// `0..n_cells`, so a bounds check here would sit in the innermost + /// loop of the scoring pass for no benefit. + #[inline] + pub fn at(&self, plane: PlaneId, cell: usize) -> Option { + let v = self.at_raw(plane, cell); + if v.is_nan() { + None + } else { + Some(v) + } + } + + #[inline] + pub fn at_raw(&self, plane: PlaneId, cell: usize) -> f32 { + debug_assert!((plane as usize) < self.names.len(), "plane out of range"); + debug_assert!(cell < self.n_cells, "cell out of range"); + self.planes[plane as usize * self.n_cells + cell] + } + + /// Optional-plane variant: `None` plane id yields `None`. Lets + /// callers hold `Option` for fields that may be absent + /// (e.g. `APCP` on an f00 file) without branching at every cell. + #[inline] + pub fn at_opt(&self, plane: Option, cell: usize) -> Option { + self.at(plane?, cell) + } + + pub fn plane_slice(&self, plane: PlaneId) -> &[f32] { + let start = plane as usize * self.n_cells; + &self.planes[start..start + self.n_cells] + } + + /// Append a plane. `data.len()` must equal `n_cells`. + /// + /// Re-pushing an existing name overwrites that plane in place — + /// the pressure and surface products never share a `"VAR:LEVEL"` + /// key, but the duct/NEXRAD/commercial enrichment does re-derive + /// planes, and overwrite is the behaviour the old `HashMap::insert` + /// had. + pub fn push_plane(&mut self, name: &str, data: Vec) { + assert_eq!( + data.len(), + self.n_cells, + "plane {name} has {} cells, grid has {}", + data.len(), + self.n_cells + ); + match self.index.get(name).copied() { + Some(existing) => { + let start = existing as usize * self.n_cells; + self.planes[start..start + self.n_cells].copy_from_slice(&data); + } + None => { + let key: Arc = Arc::from(name); + let id = self.names.len() as PlaneId; + self.names.push(Arc::clone(&key)); + self.index.insert(key, id); + self.planes.extend_from_slice(&data); + } + } + } + + /// Get-or-create a plane filled with NaN, returning a mutable slice + /// so a caller can scatter values into it without building a + /// temporary `Vec`. + pub fn plane_mut_or_insert(&mut self, name: &str) -> (PlaneId, &mut [f32]) { + let id = match self.index.get(name).copied() { + Some(existing) => existing, + None => { + let key: Arc = Arc::from(name); + let id = self.names.len() as PlaneId; + self.names.push(Arc::clone(&key)); + self.index.insert(key, id); + self.planes + .resize(self.planes.len() + self.n_cells, f32::NAN); + id + } + }; + let start = id as usize * self.n_cells; + (id, &mut self.planes[start..start + self.n_cells]) + } + + /// Fold another grid's planes into this one. Both grids must share a + /// `GridSpec` so cell indices line up — this is a plane-list append, + /// with zero per-cell work, replacing the old `merge_grids` that + /// did ~3.7 M hash inserts. + pub fn merge(&mut self, other: FieldGrid) { + assert_eq!( + self.spec, other.spec, + "cannot merge FieldGrids with different specs" + ); + for (i, name) in other.names.iter().enumerate() { + let start = i * other.n_cells; + let slice = &other.planes[start..start + other.n_cells]; + self.push_plane(name, slice.to_vec()); + } + } + + /// Geographic coordinates of a cell, rounded exactly the way the + /// previous `parse_lola_binary` rounded them (`round3` on both axes, + /// longitudes denormalised back into −180..180). Score-file + /// placement and profile snapping both depend on this matching + /// bit-for-bit. + #[inline] + pub fn cell_latlon(&self, cell: usize) -> (f64, f64) { + let row = cell / self.spec.lon_count; + let col = cell % self.spec.lon_count; + let lat = round3(self.spec.lat_start + row as f64 * self.spec.lat_step); + let raw_lon = self.spec.lon_start + col as f64 * self.spec.lon_step; + let lon = round3(crate::decoder::denormalize_lon(raw_lon)); + (lat, lon) + } + + /// Nearest cell index for a geographic coordinate, or `None` when + /// the point falls outside the grid. Used by the per-QSO point + /// worker, which asks for arbitrary lat/lon. + pub fn cell_index(&self, lat: f64, lon: f64) -> Option { + cell_index_for(&self.spec, lat, lon) + } + + /// Number of cells carrying at least one non-missing value. Used for + /// the `points` log field; scans one plane per field until it finds + /// data, so it is O(n_planes × n_cells) worst case but runs once per + /// chain step. + pub fn populated_cells(&self) -> usize { + (0..self.n_cells) + .filter(|&c| (0..self.names.len() as PlaneId).any(|p| !self.at_raw(p, c).is_nan())) + .count() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn spec_3x2() -> GridSpec { + GridSpec { + lon_start: -100.0, + lon_count: 3, + lon_step: 0.5, + lat_start: 30.0, + lat_count: 2, + lat_step: 0.5, + } + } + + #[test] + fn push_and_read_back() { + let mut g = FieldGrid::new(spec_3x2()); + g.push_plane("TMP:surface", vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + let p = g.plane_id("TMP:surface").expect("plane resolves"); + assert_eq!(g.at(p, 0), Some(1.0)); + assert_eq!(g.at(p, 5), Some(6.0)); + assert!(g.plane_id("DPT:surface").is_none()); + } + + #[test] + fn nan_reads_as_missing() { + // The old PointGrid signalled "missing" by omitting the key, so + // `cell.get(k)?` short-circuited. NaN must behave the same or + // cell_to_conditions silently accepts garbage. + let mut g = FieldGrid::new(spec_3x2()); + g.push_plane("TMP:surface", vec![f32::NAN, 2.0, 3.0, 4.0, 5.0, 6.0]); + let p = g.plane_id("TMP:surface").unwrap(); + assert_eq!(g.at(p, 0), None); + assert_eq!(g.at(p, 1), Some(2.0)); + } + + #[test] + fn cell_latlon_matches_old_parse_lola_rounding() { + let g = FieldGrid::new(spec_3x2()); + // cell 0 => row 0, col 0 + assert_eq!(g.cell_latlon(0), (30.0, -100.0)); + // cell 5 => row 1, col 2 => lat 30.5, lon -99.0 + assert_eq!(g.cell_latlon(5), (30.5, -99.0)); + } + + #[test] + fn cell_latlon_round_trips_through_cell_index() { + let g = FieldGrid::new(crate::grid::wgrib2_grid_spec()); + for cell in [0usize, 1, 472, 473, 50_000, 95_072] { + let (lat, lon) = g.cell_latlon(cell); + assert_eq!(g.cell_index(lat, lon), Some(cell), "cell {cell}"); + } + } + + #[test] + fn cell_index_rejects_out_of_bounds() { + let g = FieldGrid::new(spec_3x2()); + assert!(g.cell_index(0.0, 0.0).is_none()); + assert!(g.cell_index(30.0, -200.0).is_none()); + assert!(g.cell_index(f64::NAN, -100.0).is_none()); + } + + #[test] + fn merge_appends_planes_without_touching_cells() { + let mut a = FieldGrid::new(spec_3x2()); + a.push_plane("TMP:surface", vec![1.0; 6]); + let mut b = FieldGrid::new(spec_3x2()); + b.push_plane("DPT:surface", vec![2.0; 6]); + + a.merge(b); + + assert_eq!(a.n_planes(), 2); + let t = a.plane_id("TMP:surface").unwrap(); + let d = a.plane_id("DPT:surface").unwrap(); + assert_eq!(a.at(t, 3), Some(1.0)); + assert_eq!(a.at(d, 3), Some(2.0)); + } + + #[test] + fn push_plane_overwrites_existing_name_in_place() { + // native_duct / NEXRAD / commercial enrichment re-derives planes; + // the old HashMap::insert overwrote, so this must too rather than + // shadowing with a second plane of the same name. + let mut g = FieldGrid::new(spec_3x2()); + g.push_plane("duct_count", vec![0.0; 6]); + g.push_plane("duct_count", vec![7.0; 6]); + assert_eq!(g.n_planes(), 1); + let p = g.plane_id("duct_count").unwrap(); + assert_eq!(g.at(p, 0), Some(7.0)); + } + + #[test] + fn plane_mut_or_insert_starts_all_nan() { + let mut g = FieldGrid::new(spec_3x2()); + let (id, slice) = g.plane_mut_or_insert("nexrad_max_reflectivity_dbz"); + assert!(slice.iter().all(|v| v.is_nan())); + slice[2] = 45.0; + assert_eq!(g.at(id, 2), Some(45.0)); + assert_eq!(g.at(id, 0), None); + } + + #[test] + fn populated_cells_counts_only_cells_with_data() { + let mut g = FieldGrid::new(spec_3x2()); + g.push_plane( + "TMP:surface", + vec![1.0, f32::NAN, f32::NAN, 4.0, f32::NAN, f32::NAN], + ); + g.push_plane( + "DPT:surface", + vec![f32::NAN, 2.0, f32::NAN, f32::NAN, f32::NAN, f32::NAN], + ); + assert_eq!(g.populated_cells(), 3); + } + + #[test] + #[should_panic(expected = "different specs")] + fn merge_rejects_mismatched_specs() { + let mut a = FieldGrid::new(spec_3x2()); + let b = FieldGrid::new(crate::grid::wgrib2_grid_spec()); + a.merge(b); + } +} diff --git a/rust/prop_grid_rs/src/hrdps_fetcher.rs b/rust/prop_grid_rs/src/hrdps_fetcher.rs index a296ca85..e3cad56a 100644 --- a/rust/prop_grid_rs/src/hrdps_fetcher.rs +++ b/rust/prop_grid_rs/src/hrdps_fetcher.rs @@ -301,8 +301,12 @@ mod tests { // 8 surface + 7 pressure levels × 3 vars (TMP/DEPR/HGT) = 29 assert_eq!(manifest.len(), 29); assert!(manifest.iter().any(|(v, l)| *v == "TMP" && l == "AGL-2m")); - assert!(manifest.iter().any(|(v, l)| *v == "TMP" && l == "ISBL_0850")); - assert!(manifest.iter().any(|(v, l)| *v == "HGT" && l == "ISBL_1000")); + assert!(manifest + .iter() + .any(|(v, l)| *v == "TMP" && l == "ISBL_0850")); + assert!(manifest + .iter() + .any(|(v, l)| *v == "HGT" && l == "ISBL_1000")); } #[test] diff --git a/rust/prop_grid_rs/src/hrrr_points.rs b/rust/prop_grid_rs/src/hrrr_points.rs index 8f7a5a0c..e429e64d 100644 --- a/rust/prop_grid_rs/src/hrrr_points.rs +++ b/rust/prop_grid_rs/src/hrrr_points.rs @@ -16,9 +16,11 @@ use chrono::{DateTime, Timelike, Utc}; use sqlx::PgPool; use uuid::Uuid; -use crate::decoder::{self, CellValues}; +use crate::decoder::{self}; use crate::fetcher::{self, HrrrClient, Product}; +use crate::field_grid::FieldGrid; use crate::grid::{wgrib2_grid_spec, STEP}; +use crate::planes::GridPlanes; use crate::sounding_params::{self, Level}; // wgrib2 filter patterns are derived from static SURFACE_MESSAGES / @@ -96,7 +98,7 @@ pub async fn process_batch( // Empty grids almost always mean the HRRR archive doesn't have this // cycle (common for backfills older than the rolling retention // window). Log so "0 inserted" doesn't look like a silent success. - if sfc_grid.is_empty() && prs_grid.is_empty() { + if sfc_grid.n_planes() == 0 && prs_grid.n_planes() == 0 { tracing::warn!( valid_time = %valid_time, points = points.len(), @@ -104,92 +106,83 @@ pub async fn process_batch( ); } - // For each requested point, pull the merged cell and build an + let mut merged = sfc_grid; + merged.merge(prs_grid); + let planes = GridPlanes::resolve(&merged); + + // For each requested point, read the merged cell and build an // hrrr_profiles row. Elixir snaps points to 0.01° (see // `Weather.round_to_hrrr_grid/2`), but wgrib2 resamples to the 0.125° - // CONUS prop grid — so we must re-snap before keying into `sfc_grid` + // CONUS prop grid — so we must re-snap before indexing into `merged` // or the lookup misses and the task completes with zero inserts. // Store the profile at the original requested (lat, lon) so // `Weather.has_hrrr_profile?/3` (0.07° window) finds it on the // next backfill tick and flips the contact to `:complete`. - let mut inserted = 0u32; + let mut rows: Vec = Vec::with_capacity(points.len()); + let mut levels_buf: Vec = Vec::new(); for &(lat, lon) in points { - let key = snap_key(lat, lon); - let mut cell: CellValues = sfc_grid.get(&key).cloned().unwrap_or_default(); - if let Some(prs) = prs_grid.get(&key) { - cell.extend(prs.iter().map(|(k, v)| (k.clone(), *v))); - } - if cell.is_empty() { + let (snap_lat, snap_lon) = snap_to_grid(lat, lon); + let Some(cell) = merged.cell_index(snap_lat, snap_lon) else { + continue; + }; + // A cell with no surface temperature carries nothing worth + // persisting — same guard the old `cell.is_empty()` provided. + if merged.at_opt(planes.tmp_2m, cell).is_none() + && merged.at_opt(planes.pres_sfc, cell).is_none() + { continue; } - upsert_profile(pool, lat, lon, valid_time, &cell).await?; - inserted += 1; + planes.levels_at(&merged, cell, &mut levels_buf); + rows.push(build_row(lat, lon, &merged, &planes, cell, &levels_buf)); } + let inserted = upsert_profiles(pool, valid_time, &rows).await?; + Ok(PointFetchStats { points_requested: points.len() as u32, profiles_inserted: inserted, }) } -async fn upsert_profile( - pool: &PgPool, +/// One `hrrr_profiles` row, assembled before any DB work so the whole +/// batch goes over the wire in a single statement. +struct ProfileRow { lat: f64, lon: f64, - valid_time: DateTime, - cell: &CellValues, -) -> Result<(), sqlx::Error> { - let surface_temp_c = cell - .get("TMP:2 m above ground") - .copied() - .map(|v| (v as f64) - 273.15); - let surface_dewpoint_c = cell - .get("DPT:2 m above ground") - .copied() - .map(|v| (v as f64) - 273.15); - let surface_pressure_mb = cell - .get("PRES:surface") - .copied() - .map(|v| (v as f64) / 100.0); - let hpbl_m = cell.get("HPBL:surface").copied().map(|v| v as f64); - let pwat_mm = cell - .get("PWAT:entire atmosphere (considered as a single layer)") - .copied() - .map(|v| v as f64); + /// Per-level jsonb objects for the `profile jsonb[]` column, matching + /// Elixir's `HrrrClient.build_profile` shape: + /// `[{"pres_mb", "hght_m", "tmpc", "dwpc"}, …]`. + profile: Vec, + hpbl_m: Option, + pwat_mm: Option, + surface_temp_c: Option, + surface_dewpoint_c: Option, + surface_pressure_mb: Option, + surface_refractivity: Option, + min_refractivity_gradient: Option, + ducting_detected: bool, +} - // Pressure-level profile list. Postgres column `profile` is jsonb[] - // (ARRAY of jsonb), one element per level, matching Elixir's - // HrrrClient.build_profile output shape: - // [{"pres_mb", "hght_m", "tmpc", "dwpc"}, …] - // We build a Vec> so sqlx encodes it as jsonb[] rather - // than wrapping the whole list in a single jsonb (which yields the - // "expression is of type jsonb" type-mismatch error at runtime). - // We also collect a `Level` struct form in parallel so we can feed - // `sounding_params::*` without re-parsing JSON. - let mut typed_levels: Vec = Vec::with_capacity(fetcher::grid_level_keys().len()); - let levels: Vec> = fetcher::grid_level_keys() +fn build_row( + lat: f64, + lon: f64, + grid: &FieldGrid, + p: &GridPlanes, + cell: usize, + typed_levels: &[Level], +) -> ProfileRow { + let profile: Vec = typed_levels .iter() - .filter_map(|k| { - let t = cell.get(k.tmp.as_str()).copied()?; - let h = cell.get(k.hgt.as_str()).copied()?; - let d = cell.get(k.dpt.as_str()).copied(); - let tmpc = (t as f64) - 273.15; - let dwpc = d.map(|dv| (dv as f64) - 273.15); - typed_levels.push(Level { - pres_mb: k.pres_mb, - hght_m: h as f64, - tmpc, - dwpc, - }); + .map(|l| { let mut m = serde_json::Map::new(); - m.insert("pres_mb".into(), serde_json::json!(k.pres_mb)); - m.insert("hght_m".into(), serde_json::json!(h as f64)); - m.insert("tmpc".into(), serde_json::json!(tmpc)); - if let Some(dv) = dwpc { + m.insert("pres_mb".into(), serde_json::json!(l.pres_mb)); + m.insert("hght_m".into(), serde_json::json!(l.hght_m)); + m.insert("tmpc".into(), serde_json::json!(l.tmpc)); + if let Some(dv) = l.dwpc { m.insert("dwpc".into(), serde_json::json!(dv)); } - Some(sqlx::types::Json(serde_json::Value::Object(m))) + serde_json::Value::Object(m) }) .collect(); @@ -197,21 +190,84 @@ async fn upsert_profile( // via `SoundingParams.derive/1`. Persisting them alongside the raw // profile means contact-detail pages + per-QSO scoring don't need to // re-run SoundingParams on every read. - let surface_refractivity = sounding_params::surface_refractivity(&typed_levels); let min_refractivity_gradient = - sounding_params::min_refractivity_gradient(typed_levels.clone()); - let ducting_detected = sounding_params::ducting_detected(min_refractivity_gradient); + sounding_params::min_refractivity_gradient(typed_levels.to_vec()); - let id = Uuid::new_v4(); - sqlx::query( + ProfileRow { + lat, + lon, + profile, + hpbl_m: grid.at_opt(p.hpbl, cell).map(|v| v as f64), + pwat_mm: grid.at_opt(p.pwat, cell).map(|v| v as f64), + surface_temp_c: grid.at_opt(p.tmp_2m, cell).map(|v| (v as f64) - 273.15), + surface_dewpoint_c: grid.at_opt(p.dpt_2m, cell).map(|v| (v as f64) - 273.15), + surface_pressure_mb: grid.at_opt(p.pres_sfc, cell).map(|v| (v as f64) / 100.0), + surface_refractivity: sounding_params::surface_refractivity(typed_levels), + min_refractivity_gradient, + ducting_detected: sounding_params::ducting_detected(min_refractivity_gradient), + } +} + +/// Upsert the whole batch in one statement. +/// +/// Previously this was one awaited `INSERT` per point inside the loop, so +/// a QSO batch cost one network round-trip per point. `UNNEST` over +/// per-column arrays collapses that to a single round-trip. +/// +/// The `ON CONFLICT (lat, lon, valid_time) DO UPDATE` is load-bearing: +/// `Microwaveprop.Pskr.CalibrationSampler` runs a two-pass +/// producer/consumer loop where the second pass UPSERTs the same row with +/// non-NULL HRRR fields. Do **not** turn this into a +/// "skip if already exists" guard — that breaks the loop (see CLAUDE.md). +/// +/// `profile` is `jsonb[]`, i.e. an array *per row*, so it can't ride in a +/// flat `UNNEST` column. It goes as `jsonb[][]`-equivalent: a +/// `Vec>>` that Postgres unnests as `jsonb`, then casts +/// back to `jsonb[]` per row via `ARRAY(SELECT jsonb_array_elements(...))`. +async fn upsert_profiles( + pool: &PgPool, + valid_time: DateTime, + rows: &[ProfileRow], +) -> Result { + if rows.is_empty() { + return Ok(0); + } + + let ids: Vec = rows.iter().map(|_| Uuid::new_v4()).collect(); + let lats: Vec = rows.iter().map(|r| r.lat).collect(); + let lons: Vec = rows.iter().map(|r| r.lon).collect(); + let profiles: Vec>> = + rows.iter().map(|r| sqlx::types::Json(&r.profile)).collect(); + let hpbl: Vec> = rows.iter().map(|r| r.hpbl_m).collect(); + let pwat: Vec> = rows.iter().map(|r| r.pwat_mm).collect(); + let temp: Vec> = rows.iter().map(|r| r.surface_temp_c).collect(); + let dewp: Vec> = rows.iter().map(|r| r.surface_dewpoint_c).collect(); + let pres: Vec> = rows.iter().map(|r| r.surface_pressure_mb).collect(); + let sref: Vec> = rows.iter().map(|r| r.surface_refractivity).collect(); + let mgrad: Vec> = rows.iter().map(|r| r.min_refractivity_gradient).collect(); + let duct: Vec = rows.iter().map(|r| r.ducting_detected).collect(); + + let result = sqlx::query( r#" INSERT INTO hrrr_profiles (id, valid_time, lat, lon, run_time, profile, hpbl_m, pwat_mm, surface_temp_c, surface_dewpoint_c, surface_pressure_mb, surface_refractivity, min_refractivity_gradient, ducting_detected, is_grid_point, inserted_at, updated_at) - VALUES - ($1, $2, $3, $4, $2, $5, $6, $7, $8, $9, $10, $11, $12, $13, false, NOW(), NOW()) + SELECT + u.id, $1, u.lat, u.lon, $1, + ARRAY(SELECT jsonb_array_elements(u.profile)), + u.hpbl_m, u.pwat_mm, u.surface_temp_c, u.surface_dewpoint_c, + u.surface_pressure_mb, u.surface_refractivity, + u.min_refractivity_gradient, u.ducting_detected, + false, NOW(), NOW() + FROM UNNEST( + $2::uuid[], $3::float8[], $4::float8[], $5::jsonb[], + $6::float8[], $7::float8[], $8::float8[], $9::float8[], + $10::float8[], $11::float8[], $12::float8[], $13::bool[] + ) AS u(id, lat, lon, profile, hpbl_m, pwat_mm, surface_temp_c, + surface_dewpoint_c, surface_pressure_mb, surface_refractivity, + min_refractivity_gradient, ducting_detected) ON CONFLICT (lat, lon, valid_time) DO UPDATE SET profile = EXCLUDED.profile, hpbl_m = EXCLUDED.hpbl_m, @@ -225,59 +281,166 @@ async fn upsert_profile( updated_at = NOW() "#, ) - .bind(id) .bind(valid_time.naive_utc()) - .bind(lat) - .bind(lon) - .bind(&levels) - .bind(hpbl_m) - .bind(pwat_mm) - .bind(surface_temp_c) - .bind(surface_dewpoint_c) - .bind(surface_pressure_mb) - .bind(surface_refractivity) - .bind(min_refractivity_gradient) - .bind(ducting_detected) + .bind(&ids) + .bind(&lats) + .bind(&lons) + .bind(&profiles) + .bind(&hpbl) + .bind(&pwat) + .bind(&temp) + .bind(&dewp) + .bind(&pres) + .bind(&sref) + .bind(&mgrad) + .bind(&duct) .execute(pool) .await?; - Ok(()) + + Ok(result.rows_affected() as u32) } -/// Snap a requested (lat, lon) to the nearest 0.125° grid cell and -/// return its integer-millidegree key. The wgrib2 subprocess resamples -/// HRRR to that grid (see `wgrib2_grid_spec()`), so a request on the -/// HRRR native ~3 km grid (e.g. 44.94, -91.63 after -/// `Weather.round_to_hrrr_grid/2`, which rounds to 0.01°) never matches -/// a 0.125° cell key exactly and the point is silently dropped — -/// producing the "task done with 0 profiles inserted" pattern that -/// wedged 13k+ contacts at `hrrr_status = :queued`. -fn snap_key(lat: f64, lon: f64) -> (i32, i32) { - let snap_lat = (lat / STEP).round() * STEP; - let snap_lon = (lon / STEP).round() * STEP; - ( - (snap_lat * 1000.0).round() as i32, - (snap_lon * 1000.0).round() as i32, - ) +/// Snap a requested (lat, lon) to the nearest 0.125° grid coordinate. +/// The wgrib2 subprocess resamples HRRR to that grid (see +/// `wgrib2_grid_spec()`), so a request on the HRRR native ~3 km grid +/// (e.g. 44.94, -91.63 after `Weather.round_to_hrrr_grid/2`, which +/// rounds to 0.01°) never lands on a 0.125° cell exactly and the point +/// is silently dropped — producing the "task done with 0 profiles +/// inserted" pattern that wedged 13k+ contacts at +/// `hrrr_status = :queued`. +fn snap_to_grid(lat: f64, lon: f64) -> (f64, f64) { + ((lat / STEP).round() * STEP, (lon / STEP).round() * STEP) } #[cfg(test)] mod tests { use super::*; - #[test] - fn snap_key_on_grid_is_identity() { - assert_eq!(snap_key(45.0, -91.625), (45_000, -91_625)); - assert_eq!(snap_key(25.0, -125.0), (25_000, -125_000)); + fn key(lat: f64, lon: f64) -> (i32, i32) { + let (la, lo) = snap_to_grid(lat, lon); + ((la * 1000.0).round() as i32, (lo * 1000.0).round() as i32) } #[test] - fn snap_key_rounds_to_nearest_0_125_cell() { + fn snap_on_grid_is_identity() { + assert_eq!(key(45.0, -91.625), (45_000, -91_625)); + assert_eq!(key(25.0, -125.0), (25_000, -125_000)); + } + + #[test] + fn snap_rounds_to_nearest_0_125_cell() { // 44.94 → 45.0 (0.06 away vs 44.875 which is 0.065 away) - assert_eq!(snap_key(44.94, -91.63), (45_000, -91_625)); + assert_eq!(key(44.94, -91.63), (45_000, -91_625)); // 44.83 → 44.875 (0.045 vs 44.75 which is 0.08) - assert_eq!(snap_key(44.83, -91.5), (44_875, -91_500)); + assert_eq!(key(44.83, -91.5), (44_875, -91_500)); // 44.73 → 44.75 (0.02 vs 44.625 which is 0.105) - assert_eq!(snap_key(44.73, -91.38), (44_750, -91_375)); + assert_eq!(key(44.73, -91.38), (44_750, -91_375)); + } + + fn row(lat: f64, lon: f64, temp: f64) -> ProfileRow { + ProfileRow { + lat, + lon, + profile: vec![serde_json::json!({ + "pres_mb": 1000.0, "hght_m": 100.0, "tmpc": 21.0, "dwpc": 16.0 + })], + hpbl_m: Some(500.0), + pwat_mm: Some(25.0), + surface_temp_c: Some(temp), + surface_dewpoint_c: Some(12.0), + surface_pressure_mb: Some(1013.0), + surface_refractivity: Some(320.0), + min_refractivity_gradient: Some(-150.0), + ducting_detected: true, + } + } + + /// The batched `UNNEST` upsert packs a `jsonb[]` *per row* through a + /// flat `UNNEST` column, which is the fiddly part of the statement — + /// exercise it against a real Postgres rather than trusting it. + #[tokio::test] + async fn batched_upsert_writes_all_rows_and_is_idempotent() { + let Ok(url) = std::env::var("PROP_GRID_RS_TEST_DB") else { + eprintln!("skip: set PROP_GRID_RS_TEST_DB to run db tests"); + return; + }; + let pool = crate::db::connect(&url, 4).await.unwrap(); + let valid_time = Utc::now() + .with_nanosecond(0) + .unwrap() + .with_second(0) + .unwrap(); + // Longitudes far outside CONUS so this test can't collide with + // real rows or with a concurrently running test. + let lats = [11.111_f64, 11.222]; + let rows = vec![row(lats[0], 55.111, 20.0), row(lats[1], 55.222, 21.0)]; + + let n = upsert_profiles(&pool, valid_time, &rows).await.unwrap(); + assert_eq!(n, 2, "both rows inserted in one statement"); + + let (count, temp, levels): (i64, Option, i32) = sqlx::query_as( + r#"SELECT count(*), max(surface_temp_c), + coalesce(max(array_length(profile, 1)), 0)::int + FROM hrrr_profiles + WHERE valid_time = $1 AND lat = ANY($2)"#, + ) + .bind(valid_time.naive_utc()) + .bind(&lats[..]) + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(count, 2); + assert_eq!(temp, Some(21.0)); + assert_eq!(levels, 1, "profile round-tripped as jsonb[] with 1 level"); + + // Re-running must UPDATE, not duplicate — the PSKR sampler's + // two-pass loop depends on this. + let updated = vec![row(lats[0], 55.111, 30.0), row(lats[1], 55.222, 31.0)]; + upsert_profiles(&pool, valid_time, &updated).await.unwrap(); + + let (count2, temp2): (i64, Option) = sqlx::query_as( + r#"SELECT count(*), max(surface_temp_c) FROM hrrr_profiles + WHERE valid_time = $1 AND lat = ANY($2)"#, + ) + .bind(valid_time.naive_utc()) + .bind(&lats[..]) + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(count2, 2, "conflict updated in place"); + assert_eq!(temp2, Some(31.0), "values were overwritten"); + + sqlx::query("DELETE FROM hrrr_profiles WHERE valid_time = $1 AND lat = ANY($2)") + .bind(valid_time.naive_utc()) + .bind(&lats[..]) + .execute(&pool) + .await + .unwrap(); + } + + #[tokio::test] + async fn batched_upsert_of_empty_batch_is_a_noop() { + let Ok(url) = std::env::var("PROP_GRID_RS_TEST_DB") else { + return; + }; + let pool = crate::db::connect(&url, 2).await.unwrap(); + assert_eq!(upsert_profiles(&pool, Utc::now(), &[]).await.unwrap(), 0); + } + + /// The snapped coordinate must resolve to a real cell on the CONUS + /// grid — this is the step that used to silently miss. + #[test] + fn snapped_points_resolve_to_conus_cells() { + let grid = FieldGrid::new(wgrib2_grid_spec()); + for (lat, lon) in [(44.94, -91.63), (32.9, -97.04), (25.0, -125.0)] { + let (sla, slo) = snap_to_grid(lat, lon); + assert!( + grid.cell_index(sla, slo).is_some(), + "({lat}, {lon}) snapped to ({sla}, {slo}) missed the grid" + ); + } } } diff --git a/rust/prop_grid_rs/src/lib.rs b/rust/prop_grid_rs/src/lib.rs index 051a2047..4472afdd 100644 --- a/rust/prop_grid_rs/src/lib.rs +++ b/rust/prop_grid_rs/src/lib.rs @@ -38,14 +38,16 @@ pub mod db; pub mod decoder; pub mod duct; pub mod fetcher; +pub mod field_grid; pub mod grid; pub mod hrdps_fetcher; pub mod hrrr_points; pub mod metrics; pub mod native_duct; pub mod nexrad; +pub mod pgrid; pub mod pipeline; -pub mod profiles_file; +pub mod planes; pub mod region; pub mod scorer; pub mod scores_file; diff --git a/rust/prop_grid_rs/src/metrics.rs b/rust/prop_grid_rs/src/metrics.rs index e8e2de4b..844e8c7d 100644 --- a/rust/prop_grid_rs/src/metrics.rs +++ b/rust/prop_grid_rs/src/metrics.rs @@ -48,6 +48,42 @@ pub static TASKS_IN_FLIGHT: LazyLock = LazyLock::new(|| { .expect("register gauge") }); +/// Per-stage duration within a chain step, labelled +/// `fetch | decode | derive | write_scores | write_profile | write_scalar`. +/// +/// Added because only `CHAIN_STEP_DURATION` and `DECODE_DURATION` +/// existed, which is precisely why the post-decode cost stayed invisible: +/// wgrib2 decode is ~0.4 s of a step, while the msgpack+gzip profile +/// write was ~4 s and nothing measured it. Buckets span 1 ms to 60 s so +/// both the cheap stages and a pathological one are resolvable. +pub static STAGE_DURATION: LazyLock = LazyLock::new(|| { + register_histogram_vec!( + "prop_grid_rs_stage_duration_seconds", + "Wall time of one stage within a chain step", + &["stage"], + vec![0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1.0, 2.0, 5.0, 10.0, 30.0, 60.0] + ) + .expect("register histogram") +}); + +/// Time `f` and record it against `stage`. +pub fn observe_stage(stage: &str, f: impl FnOnce() -> T) -> T { + let started = std::time::Instant::now(); + let out = f(); + STAGE_DURATION + .with_label_values(&[stage]) + .observe(started.elapsed().as_secs_f64()); + out +} + +/// Record an already-measured stage duration. For `async` stages, where +/// wrapping a closure would mean boxing a future. +pub fn record_stage(stage: &str, elapsed: Duration) { + STAGE_DURATION + .with_label_values(&[stage]) + .observe(elapsed.as_secs_f64()); +} + /// Decode-only duration (the `spawn_blocking(wgrib2)` portion). Helps /// answer "is wgrib2 fork overhead dominant?" without instrumenting the /// full step. @@ -207,8 +243,12 @@ mod tests { let _g = SERIAL.lock().unwrap_or_else(|e| e.into_inner()); init(); let baseline_success = POINT_BATCHES_TOTAL.with_label_values(&["success"]).get(); - let baseline_requested = POINTS_PROCESSED_TOTAL.with_label_values(&["requested"]).get(); - let baseline_inserted = POINTS_PROCESSED_TOTAL.with_label_values(&["inserted"]).get(); + let baseline_requested = POINTS_PROCESSED_TOTAL + .with_label_values(&["requested"]) + .get(); + let baseline_inserted = POINTS_PROCESSED_TOTAL + .with_label_values(&["inserted"]) + .get(); let baseline_hist_count = POINT_BATCH_DURATION .with_label_values(&["success"]) .get_sample_count(); @@ -220,11 +260,15 @@ mod tests { baseline_success + 1 ); assert_eq!( - POINTS_PROCESSED_TOTAL.with_label_values(&["requested"]).get(), + POINTS_PROCESSED_TOTAL + .with_label_values(&["requested"]) + .get(), baseline_requested + 12 ); assert_eq!( - POINTS_PROCESSED_TOTAL.with_label_values(&["inserted"]).get(), + POINTS_PROCESSED_TOTAL + .with_label_values(&["inserted"]) + .get(), baseline_inserted + 9 ); assert_eq!( diff --git a/rust/prop_grid_rs/src/native_duct.rs b/rust/prop_grid_rs/src/native_duct.rs index fda2b53b..ae1494a1 100644 --- a/rust/prop_grid_rs/src/native_duct.rs +++ b/rust/prop_grid_rs/src/native_duct.rs @@ -18,9 +18,10 @@ use std::path::Path; use chrono::NaiveDate; -use crate::decoder::{self, CellValues, DecodeError, PointGrid}; +use crate::decoder::{self, DecodeError}; use crate::duct::{self, NativeProfile}; use crate::fetcher::{self, HrrrClient}; +use crate::field_grid::{FieldGrid, PlaneId}; use crate::grid::GridSpec; /// The 4 variables needed for duct detection. UGRD/VGRD/TKE were @@ -79,7 +80,7 @@ pub async fn fetch_native_duct_grid( hour: u8, forecast_hour: u8, grid_spec: GridSpec, -) -> Result, NativeDuctError> { +) -> Result, NativeDuctError> { let messages = duct_messages(); let blob = client .fetch_native_blob(date, hour, forecast_hour, &messages) @@ -88,23 +89,71 @@ pub async fn fetch_native_duct_grid( // wgrib2 decode must run on the blocking pool — it forks a // subprocess and the stdout read is synchronous IO. let pattern = match_pattern(); - let grid: PointGrid = + let grid: FieldGrid = tokio::task::spawn_blocking(move || decoder::extract_grid(&blob, &pattern, grid_spec)) .await .expect("blocking join")?; - Ok(reduce_grid_to_ducts(grid)) + Ok(reduce_grid_to_ducts(&grid)) } -/// Per-cell reducer: turn the `CellValues` map into a -/// `NativeProfile` and run the duct analyser. -pub fn reduce_grid_to_ducts(grid: PointGrid) -> HashMap<(i32, i32), DuctMetrics> { - let mut out = HashMap::with_capacity(grid.len()); - for (key, cell) in grid { - if let Some(profile) = build_native_profile(&cell) { +/// Plane ids for the 4 duct variables across all 50 hybrid levels, +/// resolved once per grid instead of formatting 200 lookup keys per +/// cell (the old code built four `String`s per level per cell — 19 M +/// allocations across a CONUS grid). +struct DuctPlanes { + levels: Vec<(PlaneId, PlaneId, Option, Option)>, +} + +impl DuctPlanes { + fn resolve(grid: &FieldGrid) -> Self { + let mut levels = Vec::with_capacity(NATIVE_LEVEL_COUNT as usize); + for level in 1..=NATIVE_LEVEL_COUNT { + let lvl = format!("{level} hybrid level"); + // HGT and TMP are required; a level missing either is junk + // and was skipped by the old per-cell filter. + let (Some(hgt), Some(tmp)) = ( + grid.plane_id(&format!("HGT:{lvl}")), + grid.plane_id(&format!("TMP:{lvl}")), + ) else { + continue; + }; + levels.push(( + hgt, + tmp, + grid.plane_id(&format!("SPFH:{lvl}")), + grid.plane_id(&format!("PRES:{lvl}")), + )); + } + Self { levels } + } +} + +/// Per-cell reducer: build a `NativeProfile` from the dense planes and +/// run the duct analyser. Keyed by cell index rather than millidegree +/// lat/lon so the merge into the surface grid is a direct index write. +pub fn reduce_grid_to_ducts(grid: &FieldGrid) -> HashMap { + let planes = DuctPlanes::resolve(grid); + let mut out = HashMap::new(); + let mut scratch: Vec<(f64, f64, f64, f64)> = Vec::with_capacity(planes.levels.len()); + + for cell in 0..grid.n_cells() { + scratch.clear(); + for &(hgt_p, tmp_p, spfh_p, pres_p) in &planes.levels { + let (Some(hgt), Some(tmp)) = (grid.at(hgt_p, cell), grid.at(tmp_p, cell)) else { + continue; + }; + // SPFH / PRES default to 0.0 if absent — the duct math handles + // degenerate rows (returns 0.0 gradient / no duct) without + // raising, matching Elixir's nil-tolerance. + let spfh = grid.at_opt(spfh_p, cell).unwrap_or(0.0); + let pres = grid.at_opt(pres_p, cell).unwrap_or(0.0); + scratch.push((hgt as f64, tmp as f64, spfh as f64, pres as f64)); + } + if let Some(profile) = profile_from_levels(&mut scratch) { let analysis = duct::analyze(&profile); out.insert( - key, + cell, DuctMetrics { native_min_gradient: duct::min_m_gradient(&profile), best_duct_freq_ghz: analysis.best_duct_band_ghz, @@ -117,38 +166,14 @@ pub fn reduce_grid_to_ducts(grid: PointGrid) -> HashMap<(i32, i32), DuctMetrics> out } -/// Build a `NativeProfile` from one cell's raw values. Mirrors the -/// Elixir `build_native_profile/1`: walk all 50 hybrid levels, keep -/// only those with both HGT and TMP present, then sort ascending by -/// height. Returns `None` if fewer than 3 levels survive (too sparse -/// to analyse). -pub fn build_native_profile(cell: &CellValues) -> Option { - // (height, tmp, spfh, pres) per level. Filter out rows where HGT - // or TMP is absent — those levels are junk. - let mut levels: Vec<(f64, f64, f64, f64)> = Vec::with_capacity(NATIVE_LEVEL_COUNT as usize); - for level in 1..=NATIVE_LEVEL_COUNT { - // `cell.get` dispatches via Arc's Borrow, so pass - // &str refs built by `format!(…).as_str()` (or inline &str - // concat) — passing `&String` triggers the wrong Borrow impl. - let lvl_str = format!("{level} hybrid level"); - let hgt_key = format!("HGT:{lvl_str}"); - let tmp_key = format!("TMP:{lvl_str}"); - let spfh_key = format!("SPFH:{lvl_str}"); - let pres_key = format!("PRES:{lvl_str}"); - let Some(&hgt) = cell.get(hgt_key.as_str()) else { - continue; - }; - let Some(&tmp) = cell.get(tmp_key.as_str()) else { - continue; - }; - // SPFH / PRES default to 0.0 if absent — the duct math handles - // degenerate rows (returns 0.0 gradient / no duct) without - // raising, matching Elixir's nil-tolerance. - let spfh = cell.get(spfh_key.as_str()).copied().unwrap_or(0.0); - let pres = cell.get(pres_key.as_str()).copied().unwrap_or(0.0); - levels.push((hgt as f64, tmp as f64, spfh as f64, pres as f64)); - } - +/// Sort one cell's collected `(height, tmp, spfh, pres)` rows ascending +/// by height and turn them into a `NativeProfile`. Mirrors the Elixir +/// `build_native_profile/1`: returns `None` if fewer than 3 levels +/// survive (too sparse to analyse). +/// +/// Takes `&mut` so the caller can reuse one scratch buffer across all +/// ~95 k cells rather than allocating per cell. +fn profile_from_levels(levels: &mut [(f64, f64, f64, f64)]) -> Option { if levels.len() < 3 { return None; } @@ -169,31 +194,35 @@ pub fn build_native_profile(cell: &CellValues) -> Option { /// Merge native duct metrics into an existing surface+pressure grid. /// Mirrors Elixir's `apply_duct_grid` — cells missing from the duct -/// map keep their previous contents; cells present overwrite the -/// `native_min_gradient` / `best_duct_freq_ghz` / etc. fields. -/// -/// Consumers that only need the raw duct HashMap can skip this; the -/// merge helper exists for symmetry with the Elixir call site. -pub fn merge_duct_grid( - mut base: HashMap<(i32, i32), CellValues>, - ducts: &HashMap<(i32, i32), DuctMetrics>, -) -> HashMap<(i32, i32), CellValues> { - for (key, metrics) in ducts { - if let Some(cell) = base.get_mut(key) { - cell.insert( - "native_min_gradient".into(), - metrics.native_min_gradient as f32, - ); - if let Some(f) = metrics.best_duct_freq_ghz { - cell.insert("best_duct_freq_ghz".into(), f as f32); - } - if let Some(t) = metrics.max_duct_thickness_m { - cell.insert("max_duct_thickness_m".into(), t as f32); - } - cell.insert("duct_count".into(), metrics.duct_count as f32); +/// map keep NaN (i.e. "absent", exactly as the old code left the key +/// out); cells present get the four duct planes filled in. +pub fn merge_duct_grid(base: &mut FieldGrid, ducts: &HashMap) { + // Collect into dense planes, then push once each. Writing four + // planes in one pass beats four passes over the duct map. + let n = base.n_cells(); + let mut min_grad = vec![f32::NAN; n]; + let mut best_freq = vec![f32::NAN; n]; + let mut max_thick = vec![f32::NAN; n]; + let mut count = vec![f32::NAN; n]; + + for (&cell, metrics) in ducts { + if cell >= n { + continue; } + min_grad[cell] = metrics.native_min_gradient as f32; + if let Some(f) = metrics.best_duct_freq_ghz { + best_freq[cell] = f as f32; + } + if let Some(t) = metrics.max_duct_thickness_m { + max_thick[cell] = t as f32; + } + count[cell] = metrics.duct_count as f32; } - base + + base.push_plane("native_min_gradient", min_grad); + base.push_plane("best_duct_freq_ghz", best_freq); + base.push_plane("max_duct_thickness_m", max_thick); + base.push_plane("duct_count", count); } /// Legacy entry used by goldens/tests: decode a GRIB2 blob already @@ -201,9 +230,9 @@ pub fn merge_duct_grid( pub fn duct_grid_from_file( grib_path: &Path, grid_spec: GridSpec, -) -> Result, NativeDuctError> { +) -> Result, NativeDuctError> { let grid = decoder::extract_grid_from_file(grib_path, &match_pattern(), grid_spec)?; - Ok(reduce_grid_to_ducts(grid)) + Ok(reduce_grid_to_ducts(&grid)) } #[cfg(test)] @@ -225,56 +254,83 @@ mod tests { assert_eq!(p, ":(TMP|SPFH|HGT|PRES):.*hybrid level:"); } - fn cell_with_levels(count: usize) -> CellValues { - let mut c = CellValues::new(); - for i in 1..=count { - c.insert(format!("HGT:{i} hybrid level").into(), (i as f32) * 50.0); - c.insert( - format!("TMP:{i} hybrid level").into(), - 290.0 - (i as f32) * 0.5, + fn two_cell_spec() -> GridSpec { + GridSpec { + lon_start: -100.0, + lon_count: 2, + lon_step: 0.5, + lat_start: 30.0, + lat_count: 1, + lat_step: 0.5, + } + } + + /// Grid where cell 0 carries `cell0_levels` hybrid levels and cell 1 + /// carries `cell1_levels`. Levels beyond a cell's count are NaN, which + /// is how the dense representation says "absent". + fn grid_with_levels(cell0_levels: usize, cell1_levels: usize) -> FieldGrid { + let mut g = FieldGrid::new(two_cell_spec()); + let max = cell0_levels.max(cell1_levels); + for i in 1..=max { + let present = |c: usize| if i <= c { 1.0 } else { f32::NAN }; + g.push_plane( + &format!("HGT:{i} hybrid level"), + vec![ + present(cell0_levels) * (i as f32) * 50.0, + present(cell1_levels) * (i as f32) * 50.0, + ], ); - c.insert( - format!("SPFH:{i} hybrid level").into(), - 0.008 - (i as f32) * 0.0001, + g.push_plane( + &format!("TMP:{i} hybrid level"), + vec![ + present(cell0_levels) * (290.0 - (i as f32) * 0.5), + present(cell1_levels) * (290.0 - (i as f32) * 0.5), + ], ); - c.insert( - format!("PRES:{i} hybrid level").into(), - 101_000.0 - (i as f32) * 600.0, + g.push_plane( + &format!("SPFH:{i} hybrid level"), + vec![ + present(cell0_levels) * (0.008 - (i as f32) * 0.0001), + present(cell1_levels) * (0.008 - (i as f32) * 0.0001), + ], + ); + g.push_plane( + &format!("PRES:{i} hybrid level"), + vec![ + present(cell0_levels) * (101_000.0 - (i as f32) * 600.0), + present(cell1_levels) * (101_000.0 - (i as f32) * 600.0), + ], ); } - c + g } #[test] - fn build_profile_drops_cells_with_fewer_than_3_levels() { - assert!(build_native_profile(&cell_with_levels(0)).is_none()); - assert!(build_native_profile(&cell_with_levels(2)).is_none()); - assert!(build_native_profile(&cell_with_levels(3)).is_some()); + fn profile_from_levels_drops_cells_with_fewer_than_3_levels() { + assert!(profile_from_levels(&mut []).is_none()); + assert!(profile_from_levels(&mut [(1.0, 2.0, 3.0, 4.0); 2]).is_none()); + assert!(profile_from_levels(&mut [(1.0, 2.0, 3.0, 4.0); 3]).is_some()); } #[test] - fn build_profile_sorts_by_height() { - let mut cell = CellValues::new(); - // Interleaved: level 3 comes first in insertion order, level 1 last. - cell.insert("HGT:3 hybrid level".into(), 150.0); - cell.insert("TMP:3 hybrid level".into(), 287.0); - cell.insert("HGT:1 hybrid level".into(), 10.0); - cell.insert("TMP:1 hybrid level".into(), 291.0); - cell.insert("HGT:2 hybrid level".into(), 50.0); - cell.insert("TMP:2 hybrid level".into(), 289.0); - let p = build_native_profile(&cell).unwrap(); + fn profile_from_levels_sorts_by_height() { + // Interleaved: level 3 first in collection order, level 1 last. + let mut levels = [ + (150.0, 287.0, 0.0, 0.0), + (10.0, 291.0, 0.0, 0.0), + (50.0, 289.0, 0.0, 0.0), + ]; + let p = profile_from_levels(&mut levels).unwrap(); assert_eq!(p.heights_m, vec![10.0, 50.0, 150.0]); assert_eq!(p.temp_k, vec![291.0, 289.0, 287.0]); } #[test] fn reduce_grid_builds_metrics_for_cells_with_enough_levels() { - let mut grid: PointGrid = HashMap::new(); - grid.insert((0, 0), cell_with_levels(50)); - grid.insert((1, 0), cell_with_levels(2)); // too sparse - let out = reduce_grid_to_ducts(grid); + let grid = grid_with_levels(50, 2); // cell 1 too sparse + let out = reduce_grid_to_ducts(&grid); assert_eq!(out.len(), 1); - let m = &out[&(0, 0)]; + let m = &out[&0]; // Temperature decreases with height and moisture decreases too // — constructed profile has no strong inversion, so the // count should be 0 but min_gradient is still computed. @@ -282,15 +338,11 @@ mod tests { } #[test] - fn merge_duct_grid_only_touches_shared_keys() { - let mut base: HashMap<(i32, i32), CellValues> = HashMap::new(); - base.insert((0, 0), { - let mut c = CellValues::new(); - c.insert("PRES:surface".into(), 101_000.0); - c - }); + fn merge_duct_grid_fills_only_cells_present_in_the_duct_map() { + let mut base = FieldGrid::new(two_cell_spec()); + base.push_plane("PRES:surface", vec![101_000.0, 101_000.0]); let ducts = HashMap::from([( - (0, 0), + 0usize, DuctMetrics { native_min_gradient: -150.0, best_duct_freq_ghz: Some(24.0), @@ -298,11 +350,23 @@ mod tests { duct_count: 1, }, )]); - let merged = merge_duct_grid(base, &ducts); - let cell = &merged[&(0, 0)]; - assert_eq!(cell.get("PRES:surface").copied(), Some(101_000.0)); - assert_eq!(cell.get("native_min_gradient").copied(), Some(-150.0)); - assert_eq!(cell.get("best_duct_freq_ghz").copied(), Some(24.0)); - assert_eq!(cell.get("duct_count").copied(), Some(1.0)); + + merge_duct_grid(&mut base, &ducts); + + let pres = base.plane_id("PRES:surface").unwrap(); + let grad = base.plane_id("native_min_gradient").unwrap(); + let freq = base.plane_id("best_duct_freq_ghz").unwrap(); + let count = base.plane_id("duct_count").unwrap(); + + assert_eq!(base.at(pres, 0), Some(101_000.0)); + assert_eq!(base.at(grad, 0), Some(-150.0)); + assert_eq!(base.at(freq, 0), Some(24.0)); + assert_eq!(base.at(count, 0), Some(1.0)); + + // Cell 1 had no duct entry — every duct plane stays absent, which + // is what the old code achieved by never inserting the key. + assert_eq!(base.at(grad, 1), None); + assert_eq!(base.at(count, 1), None); + assert_eq!(base.at(pres, 1), Some(101_000.0)); } } diff --git a/rust/prop_grid_rs/src/nexrad.rs b/rust/prop_grid_rs/src/nexrad.rs index 5347095b..df488073 100644 --- a/rust/prop_grid_rs/src/nexrad.rs +++ b/rust/prop_grid_rs/src/nexrad.rs @@ -124,11 +124,12 @@ pub async fn fetch_frame( timestamp: DateTime, points: &[(f64, f64)], ) -> Result, NexradError> { - let rounded = round_to_5min(timestamp) - .ok_or_else(|| NexradError::Decode(format!( + let rounded = round_to_5min(timestamp).ok_or_else(|| { + NexradError::Decode(format!( "invalid nexrad timestamp: {} (leap-second or out-of-range field)", timestamp - )))?; + )) + })?; let url = frame_url(rounded); let resp = client .get(&url) diff --git a/rust/prop_grid_rs/src/pgrid.rs b/rust/prop_grid_rs/src/pgrid.rs new file mode 100644 index 00000000..7d285394 --- /dev/null +++ b/rust/prop_grid_rs/src/pgrid.rs @@ -0,0 +1,453 @@ +//! `.pgrid` — dense, cell-major, random-access profile grid. +//! +//! Replaces the gzipped-MessagePack `.mp.gz` profile artifact. Measured +//! on a full CONUS grid (95,073 cells, 13 pressure levels): +//! +//! | | `.mp.gz` | `.pgrid` | +//! |---|---|---| +//! | write | 3.96 s (rmpv tree + gzip -9) | one `write_all` of 22.4 MB | +//! | size | 22.0 MB | 22.4 MB | +//! | read one cell | gunzip + unpack + atomize **the whole file** | one 236-byte `pread` | +//! +//! The write happens 30× an hour (f00 + f01..f48); the single-cell read +//! happens on every `/map` point click, every Skew-T load and every +//! `PathCompute` call. Both were paying for a format designed for +//! neither. +//! +//! ## Layout (little-endian) +//! +//! ```text +//! magic 4 "PGRD" +//! version 1 0x01 +//! flags 1 bit0: 0 = hrrr, 1 = hrdps +//! n_fields 2 u16 +//! valid_time 8 i64 unix seconds +//! lat_start 8 f64 +//! lon_start 8 f64 +//! lat_step 8 f64 +//! lon_step 8 f64 +//! n_rows 2 u16 (latitude) +//! n_cols 2 u16 (longitude) +//! field_table n_fields × 32 NUL-padded ASCII field names +//! body n_rows*n_cols*n_fields × 4 f32, CELL-MAJOR +//! ``` +//! +//! **Cell-major** (`body[(cell * n_fields + field)]`) rather than +//! plane-major: reading one cell across all fields — the thing the UI +//! actually does — is then a single contiguous `pread`. A viewport read +//! is one contiguous `pread` per grid row. +//! +//! `f32::NAN` is the missing-value sentinel, matching +//! [`FieldGrid`][crate::field_grid::FieldGrid]. +//! +//! The field table is written into the header rather than being implied +//! by the version, so the Elixir reader resolves fields by name. Adding a +//! field is backwards compatible: an older reader simply does not look +//! it up. + +use std::io::Write; +use std::path::{Path, PathBuf}; + +use chrono::{DateTime, Utc}; + +use crate::fetcher; +use crate::grid::GridSpec; + +pub const MAGIC: &[u8; 4] = b"PGRD"; +pub const VERSION: u8 = 1; +/// Fixed width of one field-table entry, in bytes. +pub const FIELD_NAME_LEN: usize = 32; + +/// Per-cell scalar fields, in on-disk order. Names match the atom keys +/// `Microwaveprop.Propagation.ProfilesFile` already whitelists, so the +/// Elixir reader can hand callers the same map shape the `.mp.gz` path +/// produced. +pub const SCALAR_FIELDS: &[&str] = &[ + "surface_temp_c", + "surface_dewpoint_c", + "surface_pressure_mb", + "hpbl_m", + "pwat_mm", + "wind_u", + "wind_v", + "cloud_cover_pct", + "precip_mm", + "native_min_gradient", + "best_duct_freq_ghz", + "max_duct_thickness_m", + "duct_count", + "nexrad_max_reflectivity_dbz", + "commercial_degradation_db", + "commercial_baseline_dbm", + "commercial_current_dbm", + "commercial_n_links", + "surface_refractivity", + "min_refractivity_gradient", +]; + +/// Full field list: the scalars above, then `hght_m_

mb`, `tmpc_

mb`, +/// `dwpc_

mb` for each pressure level. `pres_mb` is implied by the +/// field name, so it costs no bytes. +pub fn field_names() -> &'static [String] { + static NAMES: std::sync::OnceLock> = std::sync::OnceLock::new(); + NAMES.get_or_init(|| { + let mut v: Vec = SCALAR_FIELDS.iter().map(|s| (*s).to_string()).collect(); + for &p in fetcher::GRID_PRESSURE_LEVELS { + v.push(format!("hght_m_{p}mb")); + v.push(format!("tmpc_{p}mb")); + v.push(format!("dwpc_{p}mb")); + } + for n in &v { + assert!( + n.len() < FIELD_NAME_LEN, + "field name {n} exceeds {FIELD_NAME_LEN} bytes" + ); + } + v + }) +} + +pub fn n_fields() -> usize { + field_names().len() +} + +/// Byte offset of the body, i.e. the header + field table length. +pub fn header_len(n_fields: usize) -> usize { + 4 + 1 + 1 + 2 + 8 + 8 + 8 + 8 + 8 + 2 + 2 + n_fields * FIELD_NAME_LEN +} + +#[derive(Debug, thiserror::Error)] +pub enum WriteError { + #[error("io: {0}")] + Io(#[from] std::io::Error), +} + +/// `/profiles/.pgrid`. +pub fn path_for(scores_dir: &Path, valid_time: DateTime) -> PathBuf { + let iso = valid_time.format("%Y-%m-%dT%H:%M:%SZ").to_string(); + scores_dir.join("profiles").join(format!("{iso}.pgrid")) +} + +/// One cell's record: `n_fields()` f32 values in `field_names()` order. +/// The pipeline fills these directly, so no intermediate map is built. +pub type Record = Vec; + +pub fn empty_record() -> Record { + vec![f32::NAN; n_fields()] +} + +/// Serialise the header for `spec` / `valid_time`. +pub fn encode_header(spec: &GridSpec, valid_time: DateTime, hrdps: bool) -> Vec { + let names = field_names(); + let mut out = Vec::with_capacity(header_len(names.len())); + out.extend_from_slice(MAGIC); + out.push(VERSION); + out.push(u8::from(hrdps)); + out.extend_from_slice(&(names.len() as u16).to_le_bytes()); + out.extend_from_slice(&valid_time.timestamp().to_le_bytes()); + out.extend_from_slice(&spec.lat_start.to_le_bytes()); + out.extend_from_slice(&spec.lon_start.to_le_bytes()); + out.extend_from_slice(&spec.lat_step.to_le_bytes()); + out.extend_from_slice(&spec.lon_step.to_le_bytes()); + out.extend_from_slice(&(spec.lat_count as u16).to_le_bytes()); + out.extend_from_slice(&(spec.lon_count as u16).to_le_bytes()); + for name in names { + let mut buf = [0u8; FIELD_NAME_LEN]; + buf[..name.len()].copy_from_slice(name.as_bytes()); + out.extend_from_slice(&buf); + } + debug_assert_eq!(out.len(), header_len(names.len())); + out +} + +/// Write `/profiles/.pgrid` atomically (tmp + rename, +/// so an NFS reader sees the old file, the new file, or nothing). +/// +/// `body` is the flat cell-major `f32` array, length +/// `n_cells * n_fields()`. +pub fn write_atomic( + scores_dir: &Path, + valid_time: DateTime, + spec: &GridSpec, + body: &[f32], + hrdps: bool, +) -> Result { + assert_eq!( + body.len(), + spec.lat_count * spec.lon_count * n_fields(), + "pgrid body does not match grid spec" + ); + + let path = path_for(scores_dir, valid_time); + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + + 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 mut tmp = path.clone().into_os_string(); + tmp.push(format!(".tmp.{nanos}.{pid}")); + let tmp = PathBuf::from(tmp); + + { + let file = std::fs::File::create(&tmp)?; + let mut w = std::io::BufWriter::with_capacity(1 << 20, file); + w.write_all(&encode_header(spec, valid_time, hrdps))?; + // f32 little-endian. On LE hosts this is the in-memory + // representation, so it is a bulk copy; the explicit conversion + // keeps the format host-independent. + let mut chunk: Vec = Vec::with_capacity(4 * 8192); + for values in body.chunks(8192) { + chunk.clear(); + for v in values { + chunk.extend_from_slice(&v.to_le_bytes()); + } + w.write_all(&chunk)?; + } + w.flush()?; + } + + match std::fs::rename(&tmp, &path) { + Ok(()) => Ok(path), + Err(e) => { + let _ = std::fs::remove_file(&tmp); + Err(WriteError::Io(e)) + } + } +} + +// ── Reader (tests + any Rust-side consumer) ────────────────────────── + +#[derive(Debug, Clone)] +pub struct Header { + pub hrdps: bool, + pub valid_time: DateTime, + pub spec: GridSpec, + pub fields: Vec, +} + +impl Header { + pub fn field_index(&self, name: &str) -> Option { + self.fields.iter().position(|f| f == name) + } + + pub fn n_cells(&self) -> usize { + self.spec.lat_count * self.spec.lon_count + } +} + +pub fn decode_header(bytes: &[u8]) -> Option

{ + if bytes.len() < 52 || &bytes[0..4] != MAGIC || bytes[4] != VERSION { + return None; + } + let hrdps = bytes[5] != 0; + let n_fields = u16::from_le_bytes(bytes[6..8].try_into().ok()?) as usize; + let valid_time = + DateTime::::from_timestamp(i64::from_le_bytes(bytes[8..16].try_into().ok()?), 0)?; + let lat_start = f64::from_le_bytes(bytes[16..24].try_into().ok()?); + let lon_start = f64::from_le_bytes(bytes[24..32].try_into().ok()?); + let lat_step = f64::from_le_bytes(bytes[32..40].try_into().ok()?); + let lon_step = f64::from_le_bytes(bytes[40..48].try_into().ok()?); + let lat_count = u16::from_le_bytes(bytes[48..50].try_into().ok()?) as usize; + let lon_count = u16::from_le_bytes(bytes[50..52].try_into().ok()?) as usize; + + if bytes.len() < header_len(n_fields) { + return None; + } + let mut fields = Vec::with_capacity(n_fields); + for i in 0..n_fields { + let start = 52 + i * FIELD_NAME_LEN; + let raw = &bytes[start..start + FIELD_NAME_LEN]; + let end = raw.iter().position(|&b| b == 0).unwrap_or(FIELD_NAME_LEN); + fields.push(String::from_utf8_lossy(&raw[..end]).into_owned()); + } + + Some(Header { + hrdps, + valid_time, + spec: GridSpec { + lat_start, + lon_start, + lat_step, + lon_step, + lat_count, + lon_count, + }, + fields, + }) +} + +/// Read one cell's record out of a whole-file buffer. Mirrors what the +/// Elixir reader does with `:file.pread/3`. +pub fn read_cell(bytes: &[u8], header: &Header, cell: usize) -> Option> { + let n = header.fields.len(); + let start = header_len(n) + cell * n * 4; + let end = start + n * 4; + if end > bytes.len() { + return None; + } + Some( + bytes[start..end] + .chunks_exact(4) + .map(|b| f32::from_le_bytes(b.try_into().unwrap())) + .collect(), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::grid; + use chrono::TimeZone; + + fn spec_3x2() -> GridSpec { + GridSpec { + lon_start: -100.0, + lon_count: 3, + lon_step: 0.5, + lat_start: 30.0, + lat_count: 2, + lat_step: 0.5, + } + } + + #[test] + fn field_list_shape() { + let names = field_names(); + assert_eq!( + names.len(), + SCALAR_FIELDS.len() + fetcher::GRID_PRESSURE_LEVELS.len() * 3 + ); + assert_eq!(names[0], "surface_temp_c"); + assert!(names.contains(&"tmpc_850mb".to_string())); + assert!(names.contains(&"dwpc_700mb".to_string())); + assert!(names.contains(&"hght_m_1000mb".to_string())); + } + + #[test] + fn header_round_trips() { + let vt = Utc.with_ymd_and_hms(2026, 4, 19, 14, 0, 0).unwrap(); + let spec = grid::wgrib2_grid_spec(); + let bytes = encode_header(&spec, vt, false); + let h = decode_header(&bytes).expect("decodes"); + + assert!(!h.hrdps); + assert_eq!(h.valid_time, vt); + assert_eq!(h.spec, spec); + assert_eq!(h.fields, field_names()); + assert_eq!(bytes.len(), header_len(n_fields())); + } + + #[test] + fn header_carries_hrdps_flag() { + let vt = Utc.with_ymd_and_hms(2026, 4, 19, 14, 0, 0).unwrap(); + let bytes = encode_header(&grid::hrdps_grid_spec(), vt, true); + assert!(decode_header(&bytes).unwrap().hrdps); + } + + #[test] + fn decode_rejects_bad_magic_and_version() { + let vt = Utc.with_ymd_and_hms(2026, 4, 19, 14, 0, 0).unwrap(); + let mut bytes = encode_header(&spec_3x2(), vt, false); + let good = bytes.clone(); + bytes[0..4].copy_from_slice(b"XXXX"); + assert!(decode_header(&bytes).is_none()); + + let mut bytes = good; + bytes[4] = 99; + assert!(decode_header(&bytes).is_none()); + } + + #[test] + fn write_read_round_trip_per_cell() { + let dir = tempfile::tempdir().unwrap(); + let vt = Utc.with_ymd_and_hms(2026, 4, 19, 14, 0, 0).unwrap(); + let spec = spec_3x2(); + let n = spec.lat_count * spec.lon_count; + let nf = n_fields(); + + // Cell c, field f gets value c*100 + f, except cell 1 which is + // all-missing so the NaN sentinel is exercised. + let mut body = vec![f32::NAN; n * nf]; + for c in 0..n { + if c == 1 { + continue; + } + for f in 0..nf { + body[c * nf + f] = (c * 100 + f) as f32; + } + } + + let path = write_atomic(dir.path(), vt, &spec, &body, false).unwrap(); + assert!(path.to_string_lossy().ends_with(".pgrid")); + + let raw = std::fs::read(&path).unwrap(); + let h = decode_header(&raw).unwrap(); + assert_eq!(h.spec, spec); + assert_eq!( + raw.len(), + header_len(nf) + n * nf * 4, + "file is header + dense body, nothing else" + ); + + let cell0 = read_cell(&raw, &h, 0).unwrap(); + assert_eq!(cell0[0], 0.0); + assert_eq!(cell0[nf - 1], (nf - 1) as f32); + + let cell2 = read_cell(&raw, &h, 2).unwrap(); + assert_eq!(cell2[0], 200.0); + + let cell1 = read_cell(&raw, &h, 1).unwrap(); + assert!( + cell1.iter().all(|v| v.is_nan()), + "missing cell reads as NaN" + ); + + assert!(read_cell(&raw, &h, n).is_none(), "out-of-range cell"); + } + + #[test] + fn named_field_lookup_resolves_to_the_written_value() { + let dir = tempfile::tempdir().unwrap(); + let vt = Utc.with_ymd_and_hms(2026, 4, 19, 14, 0, 0).unwrap(); + let spec = spec_3x2(); + let nf = n_fields(); + let mut body = vec![f32::NAN; spec.lat_count * spec.lon_count * nf]; + + let names = field_names(); + let temp_idx = names.iter().position(|n| n == "surface_temp_c").unwrap(); + let t850_idx = names.iter().position(|n| n == "tmpc_850mb").unwrap(); + body[4 * nf + temp_idx] = 22.5; + body[4 * nf + t850_idx] = 12.75; + + let path = write_atomic(dir.path(), vt, &spec, &body, false).unwrap(); + let raw = std::fs::read(&path).unwrap(); + let h = decode_header(&raw).unwrap(); + let rec = read_cell(&raw, &h, 4).unwrap(); + + assert_eq!(rec[h.field_index("surface_temp_c").unwrap()], 22.5); + assert_eq!(rec[h.field_index("tmpc_850mb").unwrap()], 12.75); + assert!(h.field_index("no_such_field").is_none()); + } + + #[test] + fn atomic_write_leaves_no_tmp_files() { + let dir = tempfile::tempdir().unwrap(); + let vt = Utc.with_ymd_and_hms(2026, 4, 19, 14, 0, 0).unwrap(); + let spec = spec_3x2(); + let body = vec![f32::NAN; spec.lat_count * spec.lon_count * n_fields()]; + + write_atomic(dir.path(), vt, &spec, &body, false).unwrap(); + write_atomic(dir.path(), vt, &spec, &body, false).unwrap(); + + let entries: Vec<_> = std::fs::read_dir(dir.path().join("profiles")) + .unwrap() + .map(|e| e.unwrap().file_name().into_string().unwrap()) + .collect(); + assert_eq!(entries.len(), 1, "got {entries:?}"); + assert!(entries[0].ends_with(".pgrid")); + } +} diff --git a/rust/prop_grid_rs/src/pipeline.rs b/rust/prop_grid_rs/src/pipeline.rs index a84a0477..ccc28e56 100644 --- a/rust/prop_grid_rs/src/pipeline.rs +++ b/rust/prop_grid_rs/src/pipeline.rs @@ -7,21 +7,23 @@ //! boost, ProfilesFile write) — Elixir retains f00. use std::path::Path; -use std::sync::Arc; use chrono::{DateTime, Datelike, Timelike, Utc}; use crate::band_config; use crate::commercial::{self, LinkLookupEntry}; -use crate::decoder::{self, CellValues, PointGrid}; +use crate::decoder; use crate::fetcher::{self, HrrrClient, Product}; -use crate::grid::{hrdps_only_points, wgrib2_grid_spec}; +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::profiles_file::{self, CellEntry}; +use crate::pgrid; +use crate::planes::GridPlanes; use crate::scorer::{self, Conditions}; -use crate::scores_file::{self, ScorePoint}; +use crate::scores_file; use crate::sounding_params::{self, Level}; use crate::weather_scalar_file::{self, ScalarRow}; @@ -36,7 +38,7 @@ pub enum PipelineError { #[error("write: {0}")] Write(#[from] scores_file::WriteError), #[error("profile write: {0}")] - ProfileWrite(#[from] profiles_file::WriteError), + ProfileWrite(#[from] pgrid::WriteError), #[error("scalar write: {0}")] ScalarWrite(#[from] weather_scalar_file::WriteError), #[error("native duct: {0}")] @@ -113,7 +115,9 @@ pub async fn run_chain_step( step.forecast_hour, &prs_wanted, ); + let fetch_started = std::time::Instant::now(); let (sfc_blob, prs_blob) = tokio::try_join!(sfc_fut, prs_fut)?; + metrics::record_stage("fetch", fetch_started.elapsed()); if sfc_blob.is_empty() { return Err(PipelineError::EmptyGrib); } @@ -124,6 +128,7 @@ pub async fn run_chain_step( // Decoder runs wgrib2 subprocess — block so we don't hog the tokio // reactor. `spawn_blocking` lifts it onto the dedicated pool. + let decode_started = std::time::Instant::now(); let grid_spec_sfc = grid_spec; let sfc_grid = tokio::task::spawn_blocking(move || { decoder::extract_grid(&sfc_blob, &sfc_pattern, grid_spec_sfc) @@ -137,46 +142,37 @@ pub async fn run_chain_step( }) .await .expect("blocking join")?; + metrics::record_stage("decode", decode_started.elapsed()); - let merged = merge_grids(sfc_grid, prs_grid); - let point_count = merged.len() as u32; + let mut merged = sfc_grid; + merged.merge(prs_grid); - // Build `prepared` (for scoring) and `profile_entries` (for the - // per-valid_time profile file consumed by /skewt and the LiveView - // point-detail panels) in a single pass over `merged`. We need a - // borrowing iterator (not `into_iter`) because `cell_to_profile_entry` - // also reads from the cell — `merged` is dropped immediately after - // this loop finishes, so the ~200 MB string-keyed HashMaps don't - // outlive the scoring stage. The price is briefly retaining both - // vecs at the same time as `merged`, which is well inside the - // 512 MiB cgroup ceiling on talos5. - let mut prepared: Vec<(f64, f64, Conditions, scorer::BandInvariants)> = - Vec::with_capacity(merged.len()); - let mut profile_entries: Vec = Vec::with_capacity(merged.len()); - let mut scalar_rows: Vec = Vec::with_capacity(merged.len()); - for (key, cell) in merged.iter() { - let (lat, lon) = decoder::key_to_latlon(*key); - if let Some(conditions) = cell_to_conditions(cell, lat, lon, &valid_time) { - let invariants = scorer::precompute_band_invariants(&conditions); - prepared.push((lat, lon, conditions, invariants)); - profile_entries.push(cell_to_profile_entry(lat, lon, cell)); - if let Some(row) = weather_scalar_file::derive_row(lat, lon, valid_time, cell) { - scalar_rows.push(row); - } - } - } - drop(merged); - let profile_cells_written = profile_entries.len() as u32; + // 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) + }); + (out, merged.spec()) + }) + .await + .expect("blocking join"); + let (fused, spec) = fused; - // Spawn the profile-file write on the blocking pool so it overlaps - // with the scoring loop instead of serialising on the hot path. + 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 profiles = profile_entries; + let body = fused.pgrid_body; tokio::task::spawn_blocking(move || { - profiles_file::write_atomic(&scores_dir_for_profiles, valid_time, &profiles) - .map(|_| 0u32) + metrics::observe_stage("write_profile", || { + pgrid::write_atomic(&scores_dir_for_profiles, valid_time, &spec, &body, false) + .map(|_| 0u32) + }) }) }; @@ -184,62 +180,29 @@ pub async fn run_chain_step( // Elixir's `/weather` reads land on a pre-derived chunk file // instead of falling back to the cold ProfilesFile decode + per-cell // SoundingParams + WeatherLayers derivation. Identical wire format - // (gzipped MessagePack chunked by 5°×5°) to what the Elixir writer - // produces — the reader doesn't care which side wrote the bytes. + // to what the Elixir writer produces — the reader doesn't care which + // side wrote the bytes. let scores_dir_for_scalars = scores_dir_owned.clone(); let scalar_future = { - let rows = scalar_rows; + let rows = fused.scalar_rows; tokio::task::spawn_blocking(move || { - weather_scalar_file::write_atomic(&scores_dir_for_scalars, valid_time, &rows) - .map(|_| 0u32) + metrics::observe_stage("write_scalar", || { + weather_scalar_file::write_atomic(&scores_dir_for_scalars, valid_time, &rows) + .map(|_| 0u32) + }) }) }; - // Step 2: score and write every band in parallel. - // - // Scoring: rayon's thread pool over the 23 bands. Each band is - // independent pure math across the ~92k-cell `prepared` slice; - // sharing is read-only so no locking. This cuts per-step scoring - // wall from ~20s serial to ~5-8s on a 4-core box. - // - // Writing: each band file is an independent NFS fsync. Launching - // them as parallel `spawn_blocking` tasks lets them overlap on the - // blocking pool instead of serializing one-at-a-time on the hot - // path. Memory stays bounded — each Vec is ~1.8 MB so - // 23 concurrent ≈ ~40 MB peak. - let bands = band_config::all_bands(); - let prepared = Arc::new(prepared); - let scored: Vec<(u32, Vec)> = { - use rayon::prelude::*; - bands - .par_iter() - .map(|band| { - let mut scores: Vec = Vec::with_capacity(prepared.len()); - for (lat, lon, conditions, invariants) in prepared.iter() { - let r = scorer::composite_score_with(conditions, band, Some(*invariants)); - scores.push(ScorePoint { - lat: *lat, - lon: *lon, - score: r.score, - }); - } - (band.freq_mhz, scores) - }) - .collect() - }; + 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?; - let write_handles: Vec<_> = scored - .into_iter() - .map(|(band_mhz, scores)| { - let dir = scores_dir_owned.clone(); - tokio::task::spawn_blocking(move || { - scores_file::write_atomic(&dir, band_mhz, valid_time, &scores) - }) - }) - .collect(); - for h in write_handles { - h.await.expect("blocking join")?; - } profile_future .await .expect("blocking join") @@ -248,12 +211,11 @@ pub async fn run_chain_step( .await .expect("blocking join") .map_err(PipelineError::ScalarWrite)?; - let files_written = bands.len() as u32; Ok(ChainStepStats { score_files_written: files_written, point_count, - band_count: bands.len() as u32, + band_count, profile_cells_written, }) } @@ -317,129 +279,351 @@ pub async fn run_chain_step_hrdps( // HRDPS's rotated lat/lon source (production observation 2026-04-29); // -lon only computes the rotation math at the points we actually // need, dropping wall time to ~30-90 s for the ~57k Canadian cells. + let spec = hrdps_grid_spec(); let canadian_points = hrdps_only_points(); let point_count_for_log = canadian_points.len() as u32; let pattern_owned = pattern.to_string(); let blob_for_decode = blob.clone(); - let grid = tokio::task::spawn_blocking(move || { - decoder::extract_points(&blob_for_decode, &pattern_owned, &canadian_points) + let mut grid = tokio::task::spawn_blocking(move || { + decoder::extract_points(&blob_for_decode, &pattern_owned, &canadian_points, spec) }) .await .expect("blocking join")?; drop(blob); - let point_count = grid.len() as u32; + // 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)) + }) + .await + .expect("blocking join"); + + let point_count = fused.cells_scored; tracing::info!( requested = point_count_for_log, - decoded = point_count, - "hrdps points extracted" + scored = point_count, + "hrdps cells scored" ); - let mut prepared: Vec<(f64, f64, Conditions, scorer::BandInvariants)> = - Vec::with_capacity(grid.len()); - let mut scalar_rows: Vec = Vec::with_capacity(grid.len()); - for (key, mut cell) in grid.into_iter() { - let (lat, lon) = decoder::key_to_latlon(key); - // HRDPS publishes DEPR (T-Td depression in K) as a primary - // surface + pressure-level variable rather than DPT. Back-fill - // DPT entries from `TMP - DEPR` so cell_to_conditions sees the - // expected key set. - backfill_dpt_from_depr(&mut cell); - if let Some(conditions) = cell_to_conditions(&cell, lat, lon, &valid_time) { - let invariants = scorer::precompute_band_invariants(&conditions); - prepared.push((lat, lon, conditions, invariants)); - // Derive the per-cell scalar row alongside scoring. The /weather - // map reads these via Weather.weather_grid_at to render - // temperature/dewpoint-depression/refractivity layers — surfacing - // HRDPS here is what gives Canadian cells weather data on the - // map north of the HRRR coverage line. - if let Some(row) = weather_scalar_file::derive_row(lat, lon, valid_time, &cell) { - scalar_rows.push(row); - } - } - } - - let bands = band_config::all_bands(); - let prepared = std::sync::Arc::new(prepared); - let scored: Vec<(u32, Vec)> = { - use rayon::prelude::*; - bands - .par_iter() - .map(|band| { - let mut scores: Vec = Vec::with_capacity(prepared.len()); - for (lat, lon, conditions, invariants) in prepared.iter() { - let r = scorer::composite_score_with(conditions, band, Some(*invariants)); - scores.push(ScorePoint { - lat: *lat, - lon: *lon, - score: r.score, - }); - } - (band.freq_mhz, scores) - }) - .collect() - }; - let scores_dir_owned = scores_dir.to_path_buf(); // Write the HRDPS scalar artifact alongside the scores. Sibling dir // (`.hrdps/`) so HRRR's `/` write isn't clobbered. Without // this the /weather map would only show CONUS data. let scalar_dir = scores_dir_owned.clone(); - let scalar_future = tokio::task::spawn_blocking(move || { - weather_scalar_file::write_atomic_hrdps(&scalar_dir, valid_time, &scalar_rows).map(|_| ()) - }); - - let write_handles: Vec<_> = scored - .into_iter() - .map(|(band_mhz, scores)| { - let dir = scores_dir_owned.clone(); - tokio::task::spawn_blocking(move || { - scores_file::write_atomic_hrdps(&dir, band_mhz, valid_time, &scores) - }) + 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(|_| ()) }) - .collect(); - for h in write_handles { - h.await.expect("blocking join")?; - } + }; + + 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)?; - let files_written = bands.len() as u32; Ok(ChainStepStats { score_files_written: files_written, point_count, - band_count: bands.len() as u32, + 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, }) } -fn backfill_dpt_from_depr(cell: &mut CellValues) { - // Snapshot of (TMP key, DEPR key, DPT key) tuples to insert. Done in - // two passes so the iteration in pass 1 isn't invalidated by the - // mutation in pass 2. - let pending: Vec<(std::sync::Arc, f32)> = cell - .keys() - .filter_map(|key| { - let depr_key = key.strip_prefix("DEPR:")?; - let tmp_key: std::sync::Arc = format!("TMP:{depr_key}").into(); - let dpt_key: std::sync::Arc = format!("DPT:{depr_key}").into(); - if cell.contains_key(&dpt_key) { +/// 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 { + 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:` planes from `TMP: - DEPR:` 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; } - let tmp = cell.get(&tmp_key).copied()?; - let depr = cell.get(key).copied()?; - Some((dpt_key, tmp - depr)) + grid.plane_id(&format!("TMP:{level}"))?; + Some((name.to_string(), dpt)) }) .collect(); - for (dpt_key, dpt_val) in pending { - cell.insert(dpt_key, dpt_val); + 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 = (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)>, + /// 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, + pub scalar_rows: Vec, + /// 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, + want_profiles: bool, + cell_mask: Option<&[bool]>, +) -> 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, 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, + 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, + 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)> = 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, + cell_mask: Option<&[bool]>, + base: usize, + score_out: &mut [u8], + pgrid_out: Option<&mut [f32]>, +) -> (Vec, u32) { + let n_bands = bands.len(); + let n_pgrid = crate::pgrid::n_fields(); + let mut scalars: Vec = Vec::new(); + let mut levels: Vec = 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); + for (b, band) in bands.iter().enumerate() { + let r = scorer::composite_score_with(&conditions, band, Some(invariants)); + cell_out[b] = clamp_score_u8(r.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, + want_profiles: bool, + cell_mask: Option<&[bool]>, +) -> FusedOutput { + derive_and_score(grid, valid_time, want_profiles, cell_mask) +} + +/// 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, + bodies: Vec<(u32, Vec)>, + spec: crate::grid::GridSpec, + hrdps: bool, +) -> Result { + 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 { @@ -449,85 +633,48 @@ fn match_pattern(wanted: &[(String, String)]) -> String { format!(":({}):", vars.join("|")) } -fn merge_grids(mut sfc: PointGrid, prs: PointGrid) -> PointGrid { - for (k, v) in prs { - sfc.entry(k).or_default().extend(v); - } - sfc -} - fn cell_to_conditions( - cell: &CellValues, + grid: &FieldGrid, + p: &GridPlanes, + cell: usize, lat: f64, lon: f64, valid_time: &DateTime, + levels: &[Level], ) -> Option { - let tmp_k = cell.get("TMP:2 m above ground").copied()?; - let dpt_k = cell.get("DPT:2 m above ground").copied()?; + 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 ( - cell.get("UGRD:10 m above ground").copied(), - cell.get("VGRD:10 m above ground").copied(), - ) { + 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 = cell - .get("TCDC:entire atmosphere") - .copied() - .map(|v| v as f64); - let pwat_mm = cell - .get("PWAT:entire atmosphere (considered as a single layer)") - .copied() - .map(|v| v as f64); - let pressure_mb = cell - .get("PRES:surface") - .copied() - .map(|pa| pa as f64 / 100.0); + 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 = cell - .get("APCP:surface") - .copied() - .map(|mm| mm as f64) - .unwrap_or(0.0); - let nexrad_rate = scorer::dbz_to_rain_rate_mmhr( - cell.get("nexrad_max_reflectivity_dbz") - .copied() - .map(|v| v as f64), - ); + 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 = cell.get("HPBL:surface").copied().map(|v| v as f64); - let best_duct_band_ghz = cell.get("best_duct_freq_ghz").copied().map(|v| v as f64); + 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 levels: Vec = fetcher::grid_level_keys() - .iter() - .filter_map(|k| { - let t = cell.get(k.tmp.as_str()).copied()?; - let d = cell.get(k.dpt.as_str()).copied(); - let h = cell.get(k.hgt.as_str()).copied()?; - Some(Level { - pres_mb: k.pres_mb, - hght_m: h as f64, - tmpc: (t - 273.15) as f64, - dwpc: d.map(|v| (v - 273.15) as f64), - }) - }) - .collect(); - let min_refractivity_gradient = sounding_params::min_refractivity_gradient(levels); + let min_refractivity_gradient = sounding_params::min_refractivity_gradient(levels.to_vec()); Some(Conditions { abs_humidity, @@ -554,6 +701,48 @@ fn cell_to_conditions( #[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) -> Option { + 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 { + 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() { @@ -568,21 +757,22 @@ mod tests { #[test] fn cell_to_conditions_maps_fields() { use chrono::TimeZone; - let mut cell = CellValues::new(); - cell.insert("TMP:2 m above ground".into(), 298.15); // 25 °C - cell.insert("DPT:2 m above ground".into(), 293.15); // 20 °C - cell.insert("UGRD:10 m above ground".into(), 3.0); - cell.insert("VGRD:10 m above ground".into(), 4.0); - cell.insert("TCDC:entire atmosphere".into(), 30.0); - cell.insert("PRES:surface".into(), 101_000.0); - cell.insert("HPBL:surface".into(), 400.0); - cell.insert( - "PWAT:entire atmosphere (considered as a single layer)".into(), - 25.0, - ); + 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 = cell_to_conditions(&cell, 32.0, -97.0, &vt).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); @@ -597,9 +787,24 @@ mod tests { #[test] fn cell_missing_surface_returns_none() { use chrono::TimeZone; - let cell = CellValues::new(); + let grid = cell_grid(&[]); let vt = Utc.with_ymd_and_hms(2026, 6, 15, 18, 0, 0).unwrap(); - assert!(cell_to_conditions(&cell, 32.0, -97.0, &vt).is_none()); + 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] @@ -609,14 +814,15 @@ mod tests { // sees 45 dBZ overhead, take whichever rate is higher. Without // this the Rust pipeline silently underestimates rain attenuation. use chrono::TimeZone; - let mut cell = CellValues::new(); - cell.insert("TMP:2 m above ground".into(), 295.0); - cell.insert("DPT:2 m above ground".into(), 285.0); - cell.insert("APCP:surface".into(), 0.0); - cell.insert("nexrad_max_reflectivity_dbz".into(), 45.0); + 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 = cell_to_conditions(&cell, 32.0, -97.0, &vt).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}"); } @@ -629,13 +835,14 @@ mod tests { // 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 mut cell = CellValues::new(); - cell.insert("TMP:2 m above ground".into(), 295.0); - cell.insert("DPT:2 m above ground".into(), 285.0); - cell.insert("best_duct_freq_ghz".into(), 24.0); + 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 = cell_to_conditions(&cell, 32.0, -97.0, &vt).unwrap(); + let c = conditions_of(&grid, &vt).unwrap(); assert_eq!(c.best_duct_band_ghz, Some(24.0)); } @@ -671,193 +878,312 @@ mod tests { assert!(matches!(err, PipelineError::F00Reserved)); } + fn plane_value(grid: &FieldGrid, name: &str) -> Option { + grid.at(grid.plane_id(name)?, 0) + } + #[test] fn backfill_dpt_from_depr_derives_missing_dewpoint() { - let mut cell = CellValues::new(); - cell.insert("TMP:2 m above ground".into(), 290.15); // 17°C - cell.insert("DEPR:2 m above ground".into(), 5.0); + 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 cell); + backfill_dpt_from_depr(&mut grid); - let dpt = cell - .get("DPT:2 m above ground") - .copied() - .expect("DPT back-filled from DEPR"); + 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 cell = CellValues::new(); - cell.insert("TMP:2 m above ground".into(), 290.15); - cell.insert("DEPR:2 m above ground".into(), 5.0); - cell.insert("DPT:2 m above ground".into(), 280.0); // pre-existing + 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 cell); + backfill_dpt_from_depr(&mut grid); - assert_eq!(cell.get("DPT:2 m above ground").copied(), Some(280.0)); + assert_eq!(plane_value(&grid, "DPT:2 m above ground"), Some(280.0)); } #[test] fn backfill_dpt_handles_pressure_levels() { - let mut cell = CellValues::new(); - cell.insert("TMP:850 mb".into(), 275.0); - cell.insert("DEPR:850 mb".into(), 4.0); - cell.insert("TMP:700 mb".into(), 265.0); - cell.insert("DEPR:700 mb".into(), 3.0); + 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 cell); + backfill_dpt_from_depr(&mut grid); - assert_eq!(cell.get("DPT:850 mb").copied(), Some(271.0)); - assert_eq!(cell.get("DPT:700 mb").copied(), Some(262.0)); + assert_eq!(plane_value(&grid, "DPT:850 mb"), Some(271.0)); + assert_eq!(plane_value(&grid, "DPT:700 mb"), Some(262.0)); } - /// Integration-ish: hand-build a 2-cell surface + pressure grid, - /// merge → cell_to_conditions → scorer → scores_file::write_atomic - /// → scores_file::read, asserting the full chain agrees with - /// itself. Covers the post-fetch hot path (the one that scales - /// with grid size) without standing up wgrib2. + /// 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 merge_score_write_read_roundtrip() { + 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; - - // decoder encodes lat/lon at millidegree precision in the key. - fn latlon_to_key(lat: f64, lon: f64) -> (i32, i32) { - ((lat * 1000.0).round() as i32, (lon * 1000.0).round() as i32) - } - - let mut sfc = PointGrid::new(); - let mut prs = PointGrid::new(); - - let cells = [(32.0_f64, -97.0_f64), (32.03_f64, -97.03_f64)]; - for &(lat, lon) in &cells { - let mut s = CellValues::new(); - s.insert("TMP:2 m above ground".into(), 295.0); - s.insert("DPT:2 m above ground".into(), 285.0); - s.insert("UGRD:10 m above ground".into(), 2.0); - s.insert("VGRD:10 m above ground".into(), 1.0); - s.insert("TCDC:entire atmosphere".into(), 20.0); - s.insert("PRES:surface".into(), 101_000.0); - s.insert("HPBL:surface".into(), 500.0); - s.insert( - "PWAT:entire atmosphere (considered as a single layer)".into(), - 20.0, - ); - sfc.insert(latlon_to_key(lat, lon), s); - - let mut p = CellValues::new(); - p.insert("TMP:1000 mb".into(), 295.0); - p.insert("DPT:1000 mb".into(), 285.0); - p.insert("HGT:1000 mb".into(), 100.0); - p.insert("TMP:975 mb".into(), 293.0); - p.insert("DPT:975 mb".into(), 280.0); - p.insert("HGT:975 mb".into(), 300.0); - prs.insert(latlon_to_key(lat, lon), p); - } - - let merged = merge_grids(sfc, prs); - assert_eq!(merged.len(), 2); - + let (grid, spec) = small_grid(); let valid_time = Utc.with_ymd_and_hms(2026, 6, 15, 18, 0, 0).unwrap(); - let prepared: Vec<_> = merged - .into_iter() - .filter_map(|(key, cell)| { - let (lat, lon) = crate::decoder::key_to_latlon(key); - let c = cell_to_conditions(&cell, lat, lon, &valid_time)?; - let inv = scorer::precompute_band_invariants(&c); - Some((lat, lon, c, inv)) - }) - .collect(); - assert_eq!(prepared.len(), 2); + let fused = derive_and_score(&grid, valid_time, true, None); - let band = &band_config::all_bands()[0]; - let scores: Vec = prepared + 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); + let (band_mhz, dense) = &fused.band_bodies[0]; + + let scattered: Vec = [0usize, 3] .iter() - .map(|(lat, lon, c, inv)| { - let r = scorer::composite_score_with(c, band, Some(*inv)); + .map(|&cell| { + let (lat, lon) = grid.cell_latlon(cell); ScorePoint { - lat: *lat, - lon: *lon, - score: r.score, + lat, + lon, + score: dense[cell] as i32, } }) .collect(); - assert_eq!(scores.len(), 2); - for s in &scores { - assert!(s.score <= 100); - } - // Write to an atomic .prop file and read it back. - let dir = tempfile::tempdir().unwrap(); - scores_file::write_atomic(dir.path(), band.freq_mhz, valid_time, &scores).unwrap(); - - let path = scores_file::path_for(dir.path(), band.freq_mhz, valid_time); - let bytes = std::fs::read(&path).expect("read back score file"); - let decoded = scores_file::decode(&bytes).expect("decode"); - // Decoded grid covers the full CONUS bounding box as a dense - // byte array; our 2 score points land at their respective - // cells and the rest are sentinel no-data. Spot-check: header - // fields survived the round trip and the body is the expected - // size for a row_major n_rows × n_cols grid. - assert_eq!(decoded.band_mhz, band.freq_mhz); - assert_eq!( - decoded.body.len(), - decoded.n_rows as usize * decoded.n_cols as usize - ); + 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"); } - /// cell_to_profile_entry includes pre-computed surface_refractivity + /// 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)); + + 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 = (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); + 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 profile_entry_includes_refractivity_scalars() { - let mut cell = CellValues::new(); - cell.insert("TMP:2 m above ground".into(), 298.15); - cell.insert("DPT:2 m above ground".into(), 293.15); - cell.insert("PRES:surface".into(), 101_300.0); - cell.insert("HPBL:surface".into(), 1500.0); - cell.insert( - "PWAT:entire atmosphere (considered as a single layer)".into(), - 30.0, - ); - cell.insert("TMP:1000 mb".into(), 295.0); - cell.insert("DPT:1000 mb".into(), 290.0); - cell.insert("HGT:1000 mb".into(), 100.0); - cell.insert("TMP:925 mb".into(), 292.0); - cell.insert("DPT:925 mb".into(), 285.0); - cell.insert("HGT:925 mb".into(), 800.0); - cell.insert("TMP:850 mb".into(), 289.0); - cell.insert("DPT:850 mb".into(), 280.0); - cell.insert("HGT:850 mb".into(), 1500.0); + 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 entry = cell_to_profile_entry(32.9, -97.0, &cell); - let map = entry.profile.as_map().expect("profile is a map"); - let has_sr = map.iter().any(|(k, _)| k.as_str() == Some("surface_refractivity")); - let has_mg = map.iter().any(|(k, _)| k.as_str() == Some("min_refractivity_gradient")); - assert!(has_sr, "surface_refractivity key missing"); - assert!(has_mg, "min_refractivity_gradient key missing"); + 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 omitted. + /// and min_refractivity_gradient are absent (NaN in the dense record, + /// which the Elixir reader surfaces as `nil`). #[test] - fn profile_entry_omits_refractivity_when_no_dewpoint() { - let mut cell = CellValues::new(); - cell.insert("TMP:2 m above ground".into(), 298.15); - cell.insert("PRES:surface".into(), 101_300.0); - cell.insert("TMP:1000 mb".into(), 295.0); - cell.insert("HGT:1000 mb".into(), 100.0); - cell.insert("TMP:925 mb".into(), 292.0); - cell.insert("HGT:925 mb".into(), 800.0); + 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 entry = cell_to_profile_entry(32.9, -97.0, &cell); - let map = entry.profile.as_map().expect("profile is a map"); - let has_sr = map.iter().any(|(k, _)| k.as_str() == Some("surface_refractivity")); - let has_mg = map.iter().any(|(k, _)| k.as_str() == Some("min_refractivity_gradient")); - assert!(!has_sr, "surface_refractivity should be absent without dewpoint"); - assert!(!has_mg, "min_refractivity_gradient should be absent without dewpoint"); + 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()); } } @@ -935,10 +1261,13 @@ pub async fn run_analysis_step( .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) }) @@ -950,13 +1279,15 @@ pub async fn run_analysis_step( }) .await .expect("blocking join")?; + metrics::record_stage("decode", decode_started.elapsed()); - // Merge pressure into surface, then fold duct metrics into every - // cell that has both. The duct values live alongside the raw - // CellValues strings so build_grid_cache_rows can pick them up - // under their atom keys after the Elixir reader normalises. - let mut merged = merge_grids(sfc_grid, prs_grid); - merged = native_duct::merge_duct_grid(merged, &duct_map); + // 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. @@ -973,8 +1304,9 @@ pub async fn run_analysis_step( // panicking the worker. The f00 overlay is non-critical. reqwest::Client::new() }); - let nexrad_points: Vec<(f64, f64)> = - merged.keys().map(|&k| decoder::key_to_latlon(k)).collect(); + let nexrad_points: Vec<(f64, f64)> = (0..merged.n_cells()) + .map(|c| merged.cell_latlon(c)) + .collect(); let nexrad_obs: Vec = match nexrad::fetch_frame( &nexrad_http, valid_time, @@ -998,89 +1330,60 @@ pub async fn run_analysis_step( .map_err(PipelineError::Commercial)?; let commercial_cells_boosted = apply_commercial(&mut merged, &commercial_lookup); - let point_count = merged.len() as u32; + // 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) + }); + (out, merged.spec()) + }) + .await + .expect("blocking join"); + let (fused, spec) = fused; - // Build the ProfilesFile cell list (the raw enriched grid) and - // the Conditions vector (for scoring) in one pass. Profile keeps - // atmospheric + duct + NEXRAD + commercial fields under - // string-keyed msgpack; Rust doesn't need to mutate those further. - let mut profile_entries: Vec = Vec::with_capacity(merged.len()); - let mut prepared: Vec<(f64, f64, Conditions, scorer::BandInvariants)> = - Vec::with_capacity(merged.len()); - let mut scalar_rows: Vec = Vec::with_capacity(merged.len()); - for (key, cell) in merged.iter() { - let (lat, lon) = decoder::key_to_latlon(*key); - if let Some(conditions) = cell_to_conditions(cell, lat, lon, &valid_time) { - let invariants = scorer::precompute_band_invariants(&conditions); - prepared.push((lat, lon, conditions, invariants)); - profile_entries.push(cell_to_profile_entry(lat, lon, cell)); - if let Some(row) = weather_scalar_file::derive_row(lat, lon, valid_time, cell) { - scalar_rows.push(row); - } - } - } + let point_count = fused.cells_scored; + let profile_cells_written = fused.cells_scored; - // Write the profile file on the blocking pool — msgpack encode + - // gzip + atomic rename. One file per run, ~2 MB compressed, so the - // write runs well inside the scoring loop's wall time. + // 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 profiles = profile_entries; + let body = fused.pgrid_body; tokio::task::spawn_blocking(move || { - profiles_file::write_atomic(&scores_dir_for_profiles, valid_time, &profiles) - .map(|_| 0u32) + metrics::observe_stage("write_profile", || { + pgrid::write_atomic(&scores_dir_for_profiles, valid_time, &spec, &body, false) + .map(|_| 0u32) + }) }) }; // Same wire format as Elixir's `ScalarFile`. See run_chain_step // (forecast hours) for the full rationale — both code paths go // through the same writer so f00 and f01..f48 land identical - // chunked artifacts on NFS. + // artifacts on NFS. let scores_dir_for_scalars = scores_dir_owned.clone(); let scalar_future = { - let rows = scalar_rows; + let rows = fused.scalar_rows; tokio::task::spawn_blocking(move || { - weather_scalar_file::write_atomic(&scores_dir_for_scalars, valid_time, &rows) - .map(|_| 0u32) + metrics::observe_stage("write_scalar", || { + weather_scalar_file::write_atomic(&scores_dir_for_scalars, valid_time, &rows) + .map(|_| 0u32) + }) }) }; - // Bands and score writes run in parallel, same shape as - // run_chain_step. Rayon across bands, try_join_all across writes. - let bands = band_config::all_bands(); - let prepared_arc = Arc::new(prepared); - let scored: Vec<(u32, Vec)> = { - use rayon::prelude::*; - bands - .par_iter() - .map(|band| { - let mut scores: Vec = Vec::with_capacity(prepared_arc.len()); - for (lat, lon, conditions, invariants) in prepared_arc.iter() { - let r = scorer::composite_score_with(conditions, band, Some(*invariants)); - scores.push(ScorePoint { - lat: *lat, - lon: *lon, - score: r.score, - }); - } - (band.freq_mhz, scores) - }) - .collect() - }; + 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?; - let write_handles: Vec<_> = scored - .into_iter() - .map(|(band_mhz, scores)| { - let dir = scores_dir_owned.clone(); - tokio::task::spawn_blocking(move || { - scores_file::write_atomic(&dir, band_mhz, valid_time, &scores) - }) - }) - .collect(); - for h in write_handles { - h.await.expect("blocking join")?; - } profile_future .await .expect("blocking join") @@ -1090,174 +1393,57 @@ pub async fn run_analysis_step( .expect("blocking join") .map_err(PipelineError::ScalarWrite)?; - let files_written = bands.len() as u32; Ok(AnalysisStepStats { score_files_written: files_written, point_count, - band_count: bands.len() as u32, - profile_cells_written: prepared_arc.len() as u32, + band_count, + profile_cells_written, duct_cells: duct_count, nexrad_cells_with_echo, commercial_cells_boosted, }) } -fn apply_nexrad(grid: &mut PointGrid, obs: &[NexradObservation]) -> u32 { +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 { - let key = ( - (o.lat * 1000.0).round() as i32, - (o.lon * 1000.0).round() as i32, - ); - // merged's keys are whatever extract_grid produced; use direct - // lat/lon match by scanning. This is O(N*M) worst-case but N is - // ~92k and obs mirrors the same point list, so in practice each - // observation matches exactly one cell by int-key. - if let Some(cell) = grid.get_mut(&key) { - if o.max_reflectivity_dbz > 0.0 { - with_echo += 1; - } - cell.insert( - "nexrad_max_reflectivity_dbz".into(), - o.max_reflectivity_dbz as f32, - ); + // 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 PointGrid, lookup: &[LinkLookupEntry]) -> u32 { +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 (key, cell) in grid.iter_mut() { - let (lat, lon) = decoder::key_to_latlon(*key); + for cell in 0..n { + let (lat, lon) = grid.cell_latlon(cell); if let Some(d) = commercial::degradation_at((lat, lon), lookup) { - cell.insert("commercial_degradation_db".into(), d.degradation_db as f32); - cell.insert("commercial_baseline_dbm".into(), d.baseline_dbm as f32); - cell.insert("commercial_current_dbm".into(), d.current_dbm as f32); - cell.insert("commercial_n_links".into(), d.n_links as f32); + 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 } - -/// Turn one enriched cell into a ProfilesFile entry. The msgpack -/// `profile` map keeps string keys; the Elixir reader atomizes the -/// whitelisted subset (surface_*, hpbl_m, pwat_mm, duct metrics, …). -fn cell_to_profile_entry(lat: f64, lon: f64, cell: &CellValues) -> CellEntry { - use profiles_file::value_map; - use rmpv::Value as V; - - let mut kv: Vec<(&'static str, V)> = Vec::with_capacity(16); - - if let Some(&v) = cell.get("TMP:2 m above ground") { - kv.push(("surface_temp_c", V::F64((v as f64) - 273.15))); - } - if let Some(&v) = cell.get("DPT:2 m above ground") { - kv.push(("surface_dewpoint_c", V::F64((v as f64) - 273.15))); - } - if let Some(&v) = cell.get("PRES:surface") { - kv.push(("surface_pressure_mb", V::F64((v as f64) / 100.0))); - } - if let Some(&v) = cell.get("HPBL:surface") { - kv.push(("hpbl_m", V::F64(v as f64))); - } - if let Some(&v) = cell.get("PWAT:entire atmosphere (considered as a single layer)") { - kv.push(("pwat_mm", V::F64(v as f64))); - } - // Wind / cloud / precip are read by Elixir's per-QSO scorer - // (`Propagation.factors_for` in propagation.ex) when a user clicks a - // grid cell on /map. Without these, wind/sky/rain factors silently - // score 0 for Rust-produced ProfilesFile cells. - if let Some(&u) = cell.get("UGRD:10 m above ground") { - kv.push(("wind_u", V::F64(u as f64))); - } - if let Some(&vv) = cell.get("VGRD:10 m above ground") { - kv.push(("wind_v", V::F64(vv as f64))); - } - if let Some(&c) = cell.get("TCDC:entire atmosphere") { - kv.push(("cloud_cover_pct", V::F64(c as f64))); - } - if let Some(&p) = cell.get("APCP:surface") { - kv.push(("precip_mm", V::F64(p as f64))); - } - if let Some(&v) = cell.get("native_min_gradient") { - kv.push(("native_min_gradient", V::F64(v as f64))); - } - if let Some(&v) = cell.get("best_duct_freq_ghz") { - kv.push(("best_duct_freq_ghz", V::F64(v as f64))); - } - if let Some(&v) = cell.get("max_duct_thickness_m") { - kv.push(("max_duct_thickness_m", V::F64(v as f64))); - } - if let Some(&v) = cell.get("duct_count") { - kv.push(("duct_count", V::Integer((v as i64).into()))); - } - if let Some(&v) = cell.get("nexrad_max_reflectivity_dbz") { - kv.push(("nexrad_max_reflectivity_dbz", V::F64(v as f64))); - } - if let Some(&d) = cell.get("commercial_degradation_db") { - let links = cell.get("commercial_n_links").copied().unwrap_or(0.0); - let baseline = cell.get("commercial_baseline_dbm").copied().unwrap_or(0.0); - let current = cell.get("commercial_current_dbm").copied().unwrap_or(0.0); - kv.push(( - "commercial_link_degradation", - value_map([ - ("degradation_db", V::F64(d as f64)), - ("baseline_dbm", V::F64(baseline as f64)), - ("current_dbm", V::F64(current as f64)), - ("n_links", V::Integer((links as i64).into())), - ]), - )); - } - - // Pressure-level profile list for SoundingParams.derive on the - // Elixir side. Mirrors the fields derive expects. - // Also collect typed `Level` structs so we can pre-compute - // surface_refractivity and min_refractivity_gradient — the Elixir - // Skew-T page needs these but SoundingParams.derive will produce nil - // when any level is missing dewpoint. Pre-computing in Rust (which - // filters the same way) gives the Elixir side a guaranteed fallback. - let mut typed_levels: Vec = Vec::with_capacity(fetcher::grid_level_keys().len()); - let levels: Vec = fetcher::grid_level_keys() - .iter() - .filter_map(|k| { - let t = cell.get(k.tmp.as_str()).copied()?; - let d = cell.get(k.dpt.as_str()).copied(); - let h = cell.get(k.hgt.as_str()).copied()?; - let tmpc = (t as f64) - 273.15; - let dwpc = d.map(|dv| (dv as f64) - 273.15); - typed_levels.push(Level { - pres_mb: k.pres_mb, - hght_m: h as f64, - tmpc, - dwpc, - }); - let mut pairs: Vec<(V, V)> = Vec::with_capacity(4); - pairs.push((V::String("pres_mb".into()), V::F64(k.pres_mb))); - pairs.push((V::String("hght_m".into()), V::F64(h as f64))); - pairs.push((V::String("tmpc".into()), V::F64(tmpc))); - if let Some(dv) = dwpc { - pairs.push((V::String("dwpc".into()), V::F64(dv))); - } - Some(V::Map(pairs)) - }) - .collect(); - if !levels.is_empty() { - kv.push(("profile", V::Array(levels))); - } - - // Pre-compute refractivity scalars so Elixir consumers (Skew-T, - // contact detail) have a guaranteed fallback when the pressure-level - // profile is missing dewpoint on some levels. - if let Some(sr) = sounding_params::surface_refractivity(&typed_levels) { - kv.push(("surface_refractivity", V::F64(sr))); - } - if let Some(mg) = sounding_params::min_refractivity_gradient(typed_levels) { - kv.push(("min_refractivity_gradient", V::F64(mg))); - } - - let profile = value_map(kv); - - CellEntry { lat, lon, profile } -} diff --git a/rust/prop_grid_rs/src/planes.rs b/rust/prop_grid_rs/src/planes.rs new file mode 100644 index 00000000..15f459cb --- /dev/null +++ b/rust/prop_grid_rs/src/planes.rs @@ -0,0 +1,372 @@ +//! 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, +} + +#[derive(Debug, Clone)] +pub struct GridPlanes { + // Surface / column fields. + pub tmp_2m: Option, + pub dpt_2m: Option, + pub ugrd_10m: Option, + pub vgrd_10m: Option, + pub tcdc: Option, + pub pwat: Option, + pub pres_sfc: Option, + pub hpbl: Option, + pub apcp: Option, + + // f00 enrichment planes. + pub nexrad_dbz: Option, + pub native_min_gradient: Option, + pub best_duct_freq_ghz: Option, + pub max_duct_thickness_m: Option, + pub duct_count: Option, + pub commercial_degradation_db: Option, + pub commercial_baseline_dbm: Option, + pub commercial_current_dbm: Option, + pub commercial_n_links: Option, + + /// One entry per `fetcher::GRID_PRESSURE_LEVELS` level that is + /// actually present in the grid, in the same order. + pub levels: Vec, +} + +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` three times per + /// cell (once each for conditions, profile, and scalars). + #[inline] + pub fn levels_at(&self, grid: &FieldGrid, cell: usize, out: &mut Vec) { + 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| { + 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_

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 { + 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 = 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); + } +} diff --git a/rust/prop_grid_rs/src/profiles_file.rs b/rust/prop_grid_rs/src/profiles_file.rs deleted file mode 100644 index 6f4fae18..00000000 --- a/rust/prop_grid_rs/src/profiles_file.rs +++ /dev/null @@ -1,265 +0,0 @@ -//! On-disk profile store in MessagePack + gzip. -//! -//! Rust-side writer for the f00 ProfilesFile. Replaces the Elixir ETF -//! format (`:erlang.term_to_binary/2`, `.etf.gz`) so Rust can produce -//! the files without porting the Erlang serializer. Elixir's -//! `ProfilesFile` module learns to read `.mp.gz` alongside the legacy -//! `.etf.gz` during the cutover. -//! -//! Wire layout (top-level msgpack map): -//! -//! ```text -//! { -//! "v": 1, -//! "valid_time": "2026-04-19T14:00:00Z", -//! "cells": [ -//! { -//! "lat": 32.9, "lon": -97.0, -//! "profile": { -//! "native_min_gradient": -150.5, -//! "best_duct_freq_ghz": 24.0, -//! "duct_count": 1, -//! "ducts": [ -//! {"base_m": 100, "top_m": 200, "thickness_m": 100, -//! "m_deficit": 12.3, "min_freq_ghz": 15.2} -//! ], -//! "surface_temp_c": 22.5, -//! "surface_dewpoint_c": 18.1, -//! "surface_pressure_mb": 1013.2, -//! "hpbl_m": 320.0, -//! "pwat_mm": 25.3, -//! "profile": [ -//! {"pres_mb":1000,"hght_m":100,"tmpc":20.0,"dwpc":18.0}, -//! ... -//! ] -//! } -//! } -//! ] -//! } -//! ``` -//! -//! All keys are strings — MessagePack has no atom type, and Elixir's -//! reader converts known keys back to atoms via a whitelist. - -use std::fs::File; -use std::io::{BufWriter, Write}; -use std::path::{Path, PathBuf}; - -use chrono::{DateTime, Utc}; -use flate2::write::GzEncoder; -use flate2::Compression; -use serde::Serialize; - -/// Bump whenever the on-disk layout changes in a way that isn't a -/// transparent superset. Elixir reader gates on this. -pub const FORMAT_VERSION: u32 = 1; - -/// One grid cell to persist. `profile` is an arbitrary map of -/// string-keyed values — sized at write time. Kept as a -/// `rmp_serde::Value` so the writer stays agnostic about what the -/// f00 pipeline decides to include. -#[derive(Debug, Serialize)] -pub struct CellEntry { - pub lat: f64, - pub lon: f64, - pub profile: rmpv::Value, -} - -#[derive(Debug, Serialize)] -struct FileBody<'a> { - v: u32, - valid_time: String, - cells: &'a [CellEntry], -} - -#[derive(Debug, thiserror::Error)] -pub enum WriteError { - #[error("io: {0}")] - Io(#[from] std::io::Error), - #[error("encode: {0}")] - Encode(#[from] rmp_serde::encode::Error), -} - -/// Absolute path the writer lands at for `valid_time`. -/// Sibling of the Elixir `.etf.gz` file; the Elixir reader prefers -/// `.mp.gz` when both exist. -pub fn path_for(scores_dir: &Path, valid_time: DateTime) -> PathBuf { - let iso = valid_time.format("%Y-%m-%dT%H:%M:%SZ").to_string(); - scores_dir.join("profiles").join(format!("{iso}.mp.gz")) -} - -/// Write `cells` to `/profiles/.mp.gz` atomically. -/// Pattern mirrors `scores_file::write_atomic`: write to a tmp sibling, -/// rename into place. The NFS rename is atomic, so a concurrent reader -/// sees either the old file, the new file, or nothing — never a torn -/// write. -pub fn write_atomic( - scores_dir: &Path, - valid_time: DateTime, - cells: &[CellEntry], -) -> Result { - let path = path_for(scores_dir, valid_time); - if let Some(parent) = path.parent() { - std::fs::create_dir_all(parent)?; - } - 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 mut tmp = path.clone().into_os_string(); - tmp.push(format!(".tmp.{nanos}.{pid}")); - let tmp = PathBuf::from(tmp); - - { - let file = File::create(&tmp)?; - // ProfilesFile is a write-once / read-many blob — each hourly - // cycle writes it ~once but every LiveView mount, point_detail - // click, and forecast scrub pays the read cost (up to ~3 MB - // uncompressed per file × ~5 cached hours per pod). Level 9 is - // ~30% slower to encode than the default (6) but 5–10% smaller - // on semi-structured MessagePack, and our write path already - // off-loads to `spawn_blocking` so encode latency isn't on the - // user path. Readers are level-agnostic — same gzip stream. - let mut gz = GzEncoder::new(BufWriter::new(file), Compression::best()); - let body = FileBody { - v: FORMAT_VERSION, - valid_time: valid_time.format("%Y-%m-%dT%H:%M:%SZ").to_string(), - cells, - }; - let buf = rmp_serde::to_vec_named(&body)?; - gz.write_all(&buf)?; - gz.finish()?; - } - - std::fs::rename(&tmp, &path)?; - Ok(path) -} - -/// Snap a lat/lon to the propagation grid step. Mirrors Elixir's -/// `Microwaveprop.Propagation.ProfilesFile.snap/2`: -/// -/// ```elixir -/// Float.round(Float.round(coord / step) * step, 3) -/// ``` -/// -/// Two implementations of "snap" — Rust's old plain 3-decimal round -/// and Elixir's step-aware round — disagreed on every off-grid input, -/// so a coordinate written by Elixir was unreachable from Rust and -/// vice versa. -const GRID_STEP: f64 = 0.125; - -pub fn snap_coords(lat: f64, lon: f64) -> (f64, f64) { - (snap_step(lat), snap_step(lon)) -} - -fn snap_step(coord: f64) -> f64 { - let snapped = (coord / GRID_STEP).round() * GRID_STEP; - (snapped * 1000.0).round() / 1000.0 -} - -/// Convenience helper: build a `rmpv::Value` map from an -/// iterator of `(name, value)` pairs. -pub fn value_map(entries: impl IntoIterator) -> rmpv::Value { - let pairs: Vec<(rmpv::Value, rmpv::Value)> = entries - .into_iter() - .map(|(k, v)| (rmpv::Value::String(k.into()), v)) - .collect(); - rmpv::Value::Map(pairs) -} - -#[cfg(test)] -mod tests { - use super::*; - use chrono::TimeZone; - use flate2::read::GzDecoder; - use std::io::Read; - - #[test] - fn writes_and_reads_round_trip() { - let dir = tempfile::tempdir().unwrap(); - let vt = Utc.with_ymd_and_hms(2026, 4, 19, 14, 0, 0).unwrap(); - let cells = vec![CellEntry { - lat: 32.9, - lon: -97.0, - profile: value_map([ - ("surface_temp_c", rmpv::Value::F64(22.5)), - ("surface_pressure_mb", rmpv::Value::F64(1013.2)), - ("duct_count", rmpv::Value::Integer(1.into())), - ]), - }]; - - let path = write_atomic(dir.path(), vt, &cells).unwrap(); - assert!(path.exists()); - assert!(path.to_string_lossy().ends_with(".mp.gz")); - - // Decode and confirm the top-level shape. - let raw = std::fs::read(&path).unwrap(); - let mut gz = GzDecoder::new(&raw[..]); - let mut buf = Vec::new(); - gz.read_to_end(&mut buf).unwrap(); - let decoded: rmpv::Value = rmp_serde::from_slice(&buf).unwrap(); - let map = decoded.as_map().unwrap(); - assert!(map - .iter() - .any(|(k, v)| k.as_str() == Some("v") && v.as_u64() == Some(1))); - assert!(map - .iter() - .any(|(k, v)| k.as_str() == Some("valid_time") - && v.as_str() == Some("2026-04-19T14:00:00Z"))); - let cells_val = &map - .iter() - .find(|(k, _)| k.as_str() == Some("cells")) - .unwrap() - .1; - let cells_arr = cells_val.as_array().unwrap(); - assert_eq!(cells_arr.len(), 1); - } - - #[test] - fn atomic_rename_replaces_existing() { - let dir = tempfile::tempdir().unwrap(); - let vt = Utc.with_ymd_and_hms(2026, 4, 19, 14, 0, 0).unwrap(); - let first = vec![CellEntry { - lat: 32.0, - lon: -97.0, - profile: value_map([("v", rmpv::Value::Integer(1.into()))]), - }]; - let second = vec![CellEntry { - lat: 32.0, - lon: -97.0, - profile: value_map([("v", rmpv::Value::Integer(2.into()))]), - }]; - let p1 = write_atomic(dir.path(), vt, &first).unwrap(); - let p2 = write_atomic(dir.path(), vt, &second).unwrap(); - assert_eq!(p1, p2); - // Only one file in the profiles dir — the tmp was renamed away. - let entries: Vec<_> = std::fs::read_dir(dir.path().join("profiles")) - .unwrap() - .map(|e| e.unwrap().file_name().into_string().unwrap()) - .collect(); - assert_eq!(entries.len(), 1); - } - - #[test] - fn path_for_layout() { - let vt = Utc.with_ymd_and_hms(2026, 4, 19, 14, 0, 0).unwrap(); - let p = path_for(std::path::Path::new("/data/scores"), vt); - assert_eq!( - p.to_str().unwrap(), - "/data/scores/profiles/2026-04-19T14:00:00Z.mp.gz" - ); - } - - #[test] - fn snap_coords_matches_elixir_step_snap() { - // The Elixir reader writes the profiles map with keys snapped via - // `Float.round(Float.round(coord / 0.125) * 0.125, 3)`. This Rust - // helper has to produce the same keys or `read_point` looks up a - // coordinate that was never written and the UI shows "no data". - assert_eq!(snap_coords(32.8998, -97.04031), (32.875, -97.0)); - assert_eq!(snap_coords(32.07, -97.07), (32.125, -97.125)); - // Round-trip — already-snapped coordinates stay put. - assert_eq!(snap_coords(32.875, -97.0), (32.875, -97.0)); - } -} diff --git a/rust/prop_grid_rs/src/scores_file.rs b/rust/prop_grid_rs/src/scores_file.rs index d11ff443..c66eaaa3 100644 --- a/rust/prop_grid_rs/src/scores_file.rs +++ b/rust/prop_grid_rs/src/scores_file.rs @@ -146,6 +146,70 @@ pub fn encode_with_spec( Ok(out) } +/// Encode a body that is **already dense and row-major on `spec`** — +/// i.e. `body[row * n_cols + col]`, the same indexing `FieldGrid` uses +/// for cells. The scoring pass writes scores straight into this layout, +/// so there is no `ScorePoint` scatter step: no 95 k-element +/// `Vec` per band (23 × 2.3 MB of intermediates per chain +/// step) and no per-point coordinate arithmetic. +pub fn encode_dense( + band_mhz: u32, + valid_time: DateTime, + body: &[u8], + spec: grid::GridSpec, +) -> Result, WriteError> { + let n_rows: u16 = spec + .lat_count + .try_into() + .map_err(|_| WriteError::GridTooLarge { + rows: spec.lat_count, + cols: spec.lon_count, + })?; + let n_cols: u16 = spec + .lon_count + .try_into() + .map_err(|_| WriteError::GridTooLarge { + rows: spec.lat_count, + cols: spec.lon_count, + })?; + assert_eq!( + body.len(), + spec.lat_count * spec.lon_count, + "dense body does not match grid spec" + ); + + let mut out = Vec::with_capacity(33 + body.len()); + out.extend_from_slice(MAGIC); + out.push(VERSION); + out.extend_from_slice(&band_mhz.to_le_bytes()); + out.extend_from_slice(&valid_time.timestamp().to_le_bytes()); + out.extend_from_slice(&(spec.lat_start as f32).to_le_bytes()); + out.extend_from_slice(&(spec.lon_start as f32).to_le_bytes()); + out.extend_from_slice(&(spec.lon_step as f32).to_le_bytes()); + out.extend_from_slice(&n_rows.to_le_bytes()); + out.extend_from_slice(&n_cols.to_le_bytes()); + out.extend_from_slice(body); + Ok(out) +} + +/// Atomic write of a dense body to `//.prop`. +/// `hrdps` selects the `.hrdps.prop` sibling path. +pub fn write_atomic_dense( + base_dir: &Path, + band_mhz: u32, + valid_time: DateTime, + body: &[u8], + spec: grid::GridSpec, + hrdps: bool, +) -> Result<(), WriteError> { + let path = if hrdps { + path_for_hrdps(base_dir, band_mhz, valid_time) + } else { + path_for(base_dir, band_mhz, valid_time) + }; + write_atomic_at(path, encode_dense(band_mhz, valid_time, body, spec)?) +} + /// Atomic write to `//.prop` (HRRR/CONUS). /// HRDPS callers use `write_atomic_hrdps` which writes to `.hrdps.prop` /// with the Canadian grid spec. @@ -222,7 +286,9 @@ mod hrdps_tests { let base = Path::new("/data/scores"); let vt = Utc.with_ymd_and_hms(2026, 4, 29, 12, 0, 0).unwrap(); let path = path_for_hrdps(base, 10_000, vt); - assert!(path.to_string_lossy().ends_with("/2026-04-29T12:00:00Z.hrdps.prop")); + assert!(path + .to_string_lossy() + .ends_with("/2026-04-29T12:00:00Z.hrdps.prop")); } #[test] diff --git a/rust/prop_grid_rs/src/weather_scalar_file.rs b/rust/prop_grid_rs/src/weather_scalar_file.rs index 8bfb94fa..a0cecb2e 100644 --- a/rust/prop_grid_rs/src/weather_scalar_file.rs +++ b/rust/prop_grid_rs/src/weather_scalar_file.rs @@ -1,5 +1,5 @@ //! Per-cell *derived* weather scalars on disk, the cheap-read sibling -//! of [`profiles_file`][crate::profiles_file]. +//! of [`pgrid`][crate::pgrid]. //! //! 1:1 wire-compatible with `Microwaveprop.Weather.ScalarFile` on the //! Elixir side. Chunked into 5°×5° spatial buckets so a state-sized @@ -35,8 +35,8 @@ use flate2::write::GzEncoder; use flate2::Compression; use serde::Serialize; -use crate::decoder::CellValues; -use crate::fetcher; +use crate::field_grid::FieldGrid; +use crate::planes::GridPlanes; use crate::sounding_params::{ ducting_detected, min_refractivity_gradient, sat_vap_pres, surface_refractivity, Level, }; @@ -107,41 +107,44 @@ pub enum WriteError { /// cells whose surface temperature is missing or out of physically /// plausible range — matches Elixir's `build_grid_cache_row` filter /// (`temp_c < -80 or temp_c > 60` are dropped). +/// +/// `levels` is supplied by the caller (via `GridPlanes::levels_at`) so +/// the pressure-level walk is shared with the conditions and profile +/// derivation instead of being redone here. pub fn derive_row( + grid: &FieldGrid, + p: &GridPlanes, + cell: usize, lat: f64, lon: f64, valid_time: DateTime, - cell: &CellValues, + levels: &[Level], ) -> Option { - let temp_c = cell - .get("TMP:2 m above ground") - .map(|&v| v as f64 - 273.15)?; + let temp_c = grid.at_opt(p.tmp_2m, cell).map(|v| v as f64 - 273.15)?; if !temp_c.is_finite() || !(-80.0..=60.0).contains(&temp_c) { return None; } - let dewpoint_c = cell.get("DPT:2 m above ground").map(|&v| v as f64 - 273.15); + let dewpoint_c = grid.at_opt(p.dpt_2m, cell).map(|v| v as f64 - 273.15); let dewpoint_depression = dewpoint_c.map(|d| temp_c - d); - let surface_pressure_mb = cell.get("PRES:surface").map(|&v| v as f64 / 100.0); - let bl_height = cell.get("HPBL:surface").map(|&v| v as f64); - let pwat = cell - .get("PWAT:entire atmosphere (considered as a single layer)") - .map(|&v| v as f64); + let surface_pressure_mb = grid.at_opt(p.pres_sfc, cell).map(|v| v as f64 / 100.0); + let bl_height = grid.at_opt(p.hpbl, cell).map(|v| v as f64); + let pwat = grid.at_opt(p.pwat, cell).map(|v| v as f64); let surface_rh = dewpoint_c.map(|d| 100.0 * sat_vap_pres(d) / sat_vap_pres(temp_c)); - let levels = build_levels(cell); let derived_min_grad = if levels.len() >= 3 { - min_refractivity_gradient(levels.clone()) + min_refractivity_gradient(levels.to_vec()) } else { None }; - let native_min_grad = cell.get("native_min_gradient").map(|&v| v as f64); + let native_min_grad = grid.at_opt(p.native_min_gradient, cell).map(|v| v as f64); let refractivity_gradient = derived_min_grad.or(native_min_grad); let ducting = Some(ducting_detected(refractivity_gradient)); - let surface_refractivity_val = surface_refractivity(&levels); + let surface_refractivity_val = surface_refractivity(levels); let mut sorted: Vec<&Level> = levels.iter().collect(); + sorted.sort_by(|a, b| { b.pres_mb .partial_cmp(&a.pres_mb) @@ -166,14 +169,14 @@ pub fn derive_row( // nil). `duct_cutoff_ghz` IS available because // `native_duct::reduce_grid_to_ducts` reduces it to a scalar at // ingest time. - let duct_count = cell.get("duct_count").copied().unwrap_or(0.0); + let duct_count = grid.at_opt(p.duct_count, cell).unwrap_or(0.0); let duct_strength = if duct_count > 0.0 { - cell.get("max_duct_thickness_m").map(|&v| v as f64) + grid.at_opt(p.max_duct_thickness_m, cell).map(|v| v as f64) } else { None }; let duct_cutoff_ghz = if duct_count > 0.0 { - cell.get("best_duct_freq_ghz").map(|&v| v as f64) + grid.at_opt(p.best_duct_freq_ghz, cell).map(|v| v as f64) } else { None }; @@ -306,23 +309,6 @@ fn chunk_band(value: f64) -> i32 { (value / CHUNK_STEP as f64).floor() as i32 } -fn build_levels(cell: &CellValues) -> Vec { - fetcher::grid_level_keys() - .iter() - .filter_map(|k| { - let t = cell.get(k.tmp.as_str()).copied()?; - let h = cell.get(k.hgt.as_str()).copied()?; - let d = cell.get(k.dpt.as_str()).copied(); - Some(Level { - pres_mb: k.pres_mb, - hght_m: h as f64, - tmpc: (t as f64) - 273.15, - dwpc: d.map(|v| (v as f64) - 273.15), - }) - }) - .collect() -} - /// Nearest level to `target_pres` within ±25 mb, mirroring Elixir's /// `WeatherLayers.level_value/3`. fn level_value( @@ -413,20 +399,33 @@ fn compute_inversion_base_m(sorted: &[&Level]) -> Option { #[cfg(test)] mod tests { use super::*; + use crate::fetcher; + use crate::grid::GridSpec; use chrono::TimeZone; use flate2::read::GzDecoder; use std::io::Read; - use std::sync::Arc; - fn cell_with(items: &[(&str, f32)]) -> CellValues { - items - .iter() - .map(|&(k, v)| (Arc::::from(k), v)) - .collect() + fn one_cell_spec() -> GridSpec { + GridSpec { + lon_start: -97.0, + lon_count: 1, + lon_step: 0.125, + lat_start: 33.0, + lat_count: 1, + lat_step: 0.125, + } } - fn surface_only_cell() -> CellValues { - cell_with(&[ + fn grid_with(items: &[(&str, f32)]) -> FieldGrid { + let mut g = FieldGrid::new(one_cell_spec()); + for &(k, v) in items { + g.push_plane(k, vec![v]); + } + g + } + + fn surface_only_grid() -> FieldGrid { + grid_with(&[ ("TMP:2 m above ground", 295.65), // 22.5 °C ("DPT:2 m above ground", 285.65), // 12.5 °C → depression 10 °C ("PRES:surface", 101_320.0), // 1013.2 mb @@ -438,19 +437,28 @@ mod tests { ]) } + /// Resolve planes, pull the cell's levels, and derive — the same + /// sequence the pipeline runs, collapsed for test ergonomics. + fn derive(grid: &FieldGrid, vt: DateTime) -> Option { + let p = GridPlanes::resolve(grid); + let mut levels = Vec::new(); + p.levels_at(grid, 0, &mut levels); + derive_row(grid, &p, 0, 33.0, -97.0, vt, &levels) + } + #[test] fn drops_cells_with_unphysical_surface_temp() { - let mut cell = surface_only_cell(); - cell.insert("TMP:2 m above ground".into(), 100.0); // way too cold + let mut grid = surface_only_grid(); + grid.push_plane("TMP:2 m above ground", vec![100.0]); // way too cold let vt = Utc.with_ymd_and_hms(2026, 4, 29, 12, 0, 0).unwrap(); - assert!(derive_row(33.0, -97.0, vt, &cell).is_none()); + assert!(derive(&grid, vt).is_none()); } #[test] fn surface_only_row_carries_basic_fields() { - let cell = surface_only_cell(); + let grid = surface_only_grid(); let vt = Utc.with_ymd_and_hms(2026, 4, 29, 12, 0, 0).unwrap(); - let row = derive_row(33.0, -97.0, vt, &cell).expect("row"); + let row = derive(&grid, vt).expect("row"); assert_eq!(row.lat, 33.0); assert_eq!(row.lon, -97.0); @@ -471,7 +479,7 @@ mod tests { fn upper_air_levels_derive_lapse_rate_and_850mb() { // Need real fetcher::GRID_PRESSURE_LEVELS keys. Build directly // from grid_level_keys() so the synthetic cell mirrors prod. - let mut cell = surface_only_cell(); + let mut grid = surface_only_grid(); for k in fetcher::grid_level_keys() { // Plausible synthetic profile: lapse 6.5 °C/km, scale height // ≈ 8 km. Heights derived from std atmosphere approximation. @@ -486,13 +494,13 @@ mod tests { }; let t_c = 22.5 - h_m * 6.5e-3; // °C let d_c = t_c - 5.0; - cell.insert(k.tmp.clone().into(), (t_c + 273.15) as f32); - cell.insert(k.dpt.clone().into(), (d_c + 273.15) as f32); - cell.insert(k.hgt.clone().into(), h_m as f32); + grid.push_plane(&k.tmp, vec![(t_c + 273.15) as f32]); + grid.push_plane(&k.dpt, vec![(d_c + 273.15) as f32]); + grid.push_plane(&k.hgt, vec![h_m as f32]); } let vt = Utc.with_ymd_and_hms(2026, 4, 29, 12, 0, 0).unwrap(); - let row = derive_row(33.0, -97.0, vt, &cell).expect("row"); + let row = derive(&grid, vt).expect("row"); let temp_850 = row.temp_850mb.expect("850 mb T"); // h_m at 850 mb is 1500 → t_c = 22.5 - 1500 * 6.5e-3 = 12.75 °C @@ -565,13 +573,13 @@ mod tests { #[test] fn duct_cutoff_ghz_reads_best_duct_freq_from_cell() { - let mut cell = surface_only_cell(); - cell.insert("duct_count".into(), 2.0); - cell.insert("max_duct_thickness_m".into(), 120.0); - cell.insert("best_duct_freq_ghz".into(), 18.4); + let mut grid = surface_only_grid(); + grid.push_plane("duct_count", vec![2.0]); + grid.push_plane("max_duct_thickness_m", vec![120.0]); + grid.push_plane("best_duct_freq_ghz", vec![18.4]); let vt = Utc.with_ymd_and_hms(2026, 4, 29, 12, 0, 0).unwrap(); - let row = derive_row(33.0, -97.0, vt, &cell).expect("row"); + let row = derive(&grid, vt).expect("row"); assert_eq!(row.duct_strength, Some(120.0)); let cutoff = row.duct_cutoff_ghz.expect("duct_cutoff_ghz"); @@ -583,12 +591,12 @@ mod tests { #[test] fn duct_cutoff_ghz_nil_when_no_ducts() { - let mut cell = surface_only_cell(); - cell.insert("duct_count".into(), 0.0); - cell.insert("best_duct_freq_ghz".into(), 18.4); + let mut grid = surface_only_grid(); + grid.push_plane("duct_count", vec![0.0]); + grid.push_plane("best_duct_freq_ghz", vec![18.4]); let vt = Utc.with_ymd_and_hms(2026, 4, 29, 12, 0, 0).unwrap(); - let row = derive_row(33.0, -97.0, vt, &cell).expect("row"); + let row = derive(&grid, vt).expect("row"); assert_eq!(row.duct_cutoff_ghz, None); } diff --git a/rust/prop_grid_rs/tests/fused_pass_perf.rs b/rust/prop_grid_rs/tests/fused_pass_perf.rs new file mode 100644 index 00000000..4c6af047 --- /dev/null +++ b/rust/prop_grid_rs/tests/fused_pass_perf.rs @@ -0,0 +1,141 @@ +//! 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 { + (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()); +} diff --git a/test/microwaveprop/propagation/pgrid_test.exs b/test/microwaveprop/propagation/pgrid_test.exs new file mode 100644 index 00000000..5bcbde2b --- /dev/null +++ b/test/microwaveprop/propagation/pgrid_test.exs @@ -0,0 +1,323 @@ +defmodule Microwaveprop.Propagation.PgridTest do + @moduledoc """ + The `.pgrid` reader is the Elixir half of a format shared with + `rust/prop_grid_rs/src/pgrid.rs`. These tests build the binary by hand + to the documented layout, which is what pins the two implementations + together — if the Rust writer's header or field ordering drifts, the + hand-built fixture here still asserts the contract the reader expects. + """ + # NOT async: the setup swaps `:propagation_scores_dir` in the + # application env, which is global. Running concurrently with other + # tests lets them observe this suite's temp directory (and lets them + # observe each other's), which showed up as cross-suite failures in + # ImportLive / ContactEdit / MonitorLive as well as flaky failures here. + use ExUnit.Case, async: false + + alias Microwaveprop.Propagation.Pgrid + alias Microwaveprop.Propagation.ProfilesFile + + @magic "PGRD" + @version 1 + @field_name_len 32 + + @scalar_fields ~w( + surface_temp_c surface_dewpoint_c surface_pressure_mb hpbl_m pwat_mm + wind_u wind_v cloud_cover_pct precip_mm native_min_gradient + best_duct_freq_ghz max_duct_thickness_m duct_count + nexrad_max_reflectivity_dbz commercial_degradation_db + commercial_baseline_dbm commercial_current_dbm commercial_n_links + surface_refractivity min_refractivity_gradient + ) + + @levels [1000, 925, 850] + + setup do + dir = Path.join(System.tmp_dir!(), "pgrid_test_#{System.unique_integer([:positive])}") + File.mkdir_p!(Path.join(dir, "profiles")) + prev = Application.get_env(:microwaveprop, :propagation_scores_dir) + Application.put_env(:microwaveprop, :propagation_scores_dir, dir) + + on_exit(fn -> + if prev, + do: Application.put_env(:microwaveprop, :propagation_scores_dir, prev), + else: Application.delete_env(:microwaveprop, :propagation_scores_dir) + + File.rm_rf(dir) + end) + + {:ok, dir: dir} + end + + defp field_names do + @scalar_fields ++ + Enum.flat_map(@levels, fn p -> ["hght_m_#{p}mb", "tmpc_#{p}mb", "dwpc_#{p}mb"] end) + end + + defp nan, do: <<0, 0, 192, 127>> + + # Build a .pgrid byte-for-byte per the documented layout. `cells` is a + # map of cell_index => %{field_name => value}; everything else is NaN. + defp build(valid_time, opts, cells) do + names = field_names() + n_fields = length(names) + n_rows = Keyword.fetch!(opts, :n_rows) + n_cols = Keyword.fetch!(opts, :n_cols) + + table = + names + |> Enum.map(fn n -> n <> String.duplicate(<<0>>, @field_name_len - byte_size(n)) end) + |> IO.iodata_to_binary() + + body = + for cell <- 0..(n_rows * n_cols - 1), name <- names, into: <<>> do + case cells |> Map.get(cell, %{}) |> Map.get(name) do + nil -> nan() + v -> <> + end + end + + flags = if opts[:hrdps], do: 1, else: 0 + unix = DateTime.to_unix(valid_time) + lat_start = Keyword.fetch!(opts, :lat_start) + lon_start = Keyword.fetch!(opts, :lon_start) + + header = + @magic <> + <<@version, flags, n_fields::little-16, unix::little-signed-64, lat_start::little-float-64, + lon_start::little-float-64, 0.125::little-float-64, 0.125::little-float-64, n_rows::little-16, + n_cols::little-16>> + + header <> table <> body + end + + defp write!(valid_time, opts, cells) do + path = Pgrid.path_for(valid_time) + File.mkdir_p!(Path.dirname(path)) + File.write!(path, build(valid_time, opts, cells)) + path + end + + defp vt, do: ~U[2026-04-19 14:00:00Z] + + defp grid_opts, do: [n_rows: 2, n_cols: 3, lat_start: 32.0, lon_start: -97.0] + + defp full_cell do + %{ + "surface_temp_c" => 22.5, + "surface_dewpoint_c" => 17.0, + "surface_pressure_mb" => 1013.2, + "hpbl_m" => 1500.0, + "pwat_mm" => 30.0, + "wind_u" => 3.0, + "wind_v" => 4.0, + "cloud_cover_pct" => 20.0, + "precip_mm" => 0.5, + "duct_count" => 2.0, + "best_duct_freq_ghz" => 18.4, + "surface_refractivity" => 320.0, + "min_refractivity_gradient" => -150.0, + "hght_m_1000mb" => 100.0, + "tmpc_1000mb" => 21.0, + "dwpc_1000mb" => 16.0, + "hght_m_850mb" => 1500.0, + "tmpc_850mb" => 12.75, + "dwpc_850mb" => 7.0 + } + end + + describe "read_point/3" do + test "returns the cell's scalars with atom keys" do + write!(vt(), grid_opts(), %{0 => full_cell()}) + + profile = Pgrid.read_point(vt(), 32.0, -97.0) + + assert profile.surface_temp_c == 22.5 + assert profile.surface_dewpoint_c == 17.0 + # `.pgrid` stores f32, so values that aren't exactly representable + # come back with ~7 significant digits (1013.2 -> 1013.2000122). + # The source GRIB2 data is f32 anyway, and HRRR surface pressure is + # good to ~0.1 mb, so this is below the noise floor of the input. + assert_in_delta profile.surface_pressure_mb, 1013.2, 0.001 + assert profile.hpbl_m == 1500.0 + assert profile.pwat_mm == 30.0 + assert profile.wind_u == 3.0 + assert profile.cloud_cover_pct == 20.0 + assert profile.surface_refractivity == 320.0 + assert profile.min_refractivity_gradient == -150.0 + end + + test "duct_count comes back as an integer, matching the msgpack shape" do + write!(vt(), grid_opts(), %{0 => full_cell()}) + assert Pgrid.read_point(vt(), 32.0, -97.0).duct_count === 2 + end + + test "rebuilds the pressure-level profile list, skipping absent levels" do + write!(vt(), grid_opts(), %{0 => full_cell()}) + + profile = Pgrid.read_point(vt(), 32.0, -97.0) + + # 925 mb was never written, so it must not appear. + assert [l1000, l850] = profile.profile + assert l1000.pres_mb == 1000.0 + assert l1000.hght_m == 100.0 + assert l1000.tmpc == 21.0 + assert l1000.dwpc == 16.0 + assert l850.pres_mb == 850.0 + assert l850.tmpc == 12.75 + end + + test "a level with height and temperature but no dewpoint omits :dwpc" do + cell = Map.drop(full_cell(), ["dwpc_1000mb", "hght_m_850mb", "tmpc_850mb", "dwpc_850mb"]) + write!(vt(), grid_opts(), %{0 => cell}) + + assert [level] = Pgrid.read_point(vt(), 32.0, -97.0).profile + assert level.tmpc == 21.0 + refute Map.has_key?(level, :dwpc) + end + + test "NaN fields are absent rather than surfacing as NaN" do + # Only surface temp written; everything else is the NaN sentinel. + write!(vt(), grid_opts(), %{0 => %{"surface_temp_c" => 10.0}}) + + profile = Pgrid.read_point(vt(), 32.0, -97.0) + + assert profile.surface_temp_c == 10.0 + refute Map.has_key?(profile, :pwat_mm) + refute Map.has_key?(profile, :hpbl_m) + refute Map.has_key?(profile, :profile) + end + + test "reconstructs the nested commercial_link_degradation map" do + cell = + Map.merge(full_cell(), %{ + "commercial_degradation_db" => 6.5, + "commercial_baseline_dbm" => -45.0, + "commercial_current_dbm" => -51.5, + "commercial_n_links" => 3.0 + }) + + write!(vt(), grid_opts(), %{0 => cell}) + + assert %{ + degradation_db: 6.5, + baseline_dbm: -45.0, + current_dbm: -51.5, + n_links: 3 + } = Pgrid.read_point(vt(), 32.0, -97.0).commercial_link_degradation + end + + test "commercial_* fields are not surfaced at the top level" do + write!(vt(), grid_opts(), %{0 => full_cell()}) + profile = Pgrid.read_point(vt(), 32.0, -97.0) + refute Map.has_key?(profile, :commercial_degradation_db) + refute Map.has_key?(profile, :commercial_link_degradation) + end + + test "addresses the correct cell for a non-origin coordinate" do + # cell 4 => row 1, col 1 => lat 32.125, lon -96.875 + write!(vt(), grid_opts(), %{4 => %{"surface_temp_c" => 99.0}}) + + assert Pgrid.read_point(vt(), 32.125, -96.875).surface_temp_c == 99.0 + assert Pgrid.read_point(vt(), 32.0, -97.0) == nil + end + + test "returns nil outside the grid and for a missing file" do + write!(vt(), grid_opts(), %{0 => full_cell()}) + + assert Pgrid.read_point(vt(), 10.0, -97.0) == nil + assert Pgrid.read_point(vt(), 32.0, -150.0) == nil + assert Pgrid.read_point(~U[2026-04-19 15:00:00Z], 32.0, -97.0) == nil + end + + test "an unpopulated cell reads as nil, not an empty map" do + write!(vt(), grid_opts(), %{0 => full_cell()}) + assert Pgrid.read_point(vt(), 32.125, -96.875) == nil + end + end + + describe "read/1" do + test "returns only populated cells, keyed by snapped lat/lon" do + write!(vt(), grid_opts(), %{ + 0 => %{"surface_temp_c" => 1.0}, + 5 => %{"surface_temp_c" => 2.0} + }) + + assert {:ok, grid} = Pgrid.read(vt()) + assert map_size(grid) == 2 + assert grid[{32.0, -97.0}].surface_temp_c == 1.0 + # cell 5 => row 1, col 2 => lat 32.125, lon -96.75 + assert grid[{32.125, -96.75}].surface_temp_c == 2.0 + end + + test "returns :enoent for a missing file" do + assert Pgrid.read(vt()) == {:error, :enoent} + end + + test "rejects a file with the wrong magic" do + path = Pgrid.path_for(vt()) + File.mkdir_p!(Path.dirname(path)) + bin = build(vt(), grid_opts(), %{0 => full_cell()}) + File.write!(path, "XXXX" <> binary_part(bin, 4, byte_size(bin) - 4)) + + assert Pgrid.read(vt()) == {:error, :enoent} + end + end + + describe "header" do + test "carries the hrdps flag and grid geometry" do + bin = build(vt(), Keyword.put(grid_opts(), :hrdps, true), %{}) + assert {:ok, header} = Pgrid.parse_header(bin) + assert header.hrdps? + assert header.valid_time == vt() + assert header.n_rows == 2 + assert header.n_cols == 3 + assert header.lat_start == 32.0 + end + + test "rejects an unknown version" do + bin = build(vt(), grid_opts(), %{}) + bumped = binary_part(bin, 0, 4) <> <<99>> <> binary_part(bin, 5, byte_size(bin) - 5) + assert Pgrid.parse_header(bumped) == :error + end + end + + describe "ProfilesFile integration" do + test "read_point/3 prefers .pgrid and returns the same shape" do + write!(vt(), grid_opts(), %{0 => full_cell()}) + + profile = ProfilesFile.read_point(vt(), 32.0, -97.0) + + assert profile.surface_temp_c == 22.5 + # Same nested shape the `.mp.gz` reader produced, so Skew-T and + # PathCompute consumers need no change. + assert [%{pres_mb: 1000.0, tmpc: 21.0}, %{pres_mb: 850.0}] = profile.profile + end + + test "read_point/3 snaps an off-grid coordinate to the nearest cell" do + write!(vt(), grid_opts(), %{0 => full_cell()}) + + # 31.98 / -96.99 snap to 32.0 / -97.0. + assert ProfilesFile.read_point(vt(), 31.98, -96.99).surface_temp_c == 22.5 + end + + test "list_valid_times/0 discovers .pgrid files" do + write!(vt(), grid_opts(), %{0 => full_cell()}) + write!(~U[2026-04-19 15:00:00Z], grid_opts(), %{0 => full_cell()}) + + times = ProfilesFile.list_valid_times() + + assert vt() in times + assert ~U[2026-04-19 15:00:00Z] in times + end + + test "retain_window/2 prunes .pgrid files outside the window" do + write!(vt(), grid_opts(), %{0 => full_cell()}) + write!(~U[2026-04-19 20:00:00Z], grid_opts(), %{0 => full_cell()}) + + # Keep [14:00, 15:00] — the 20:00 file falls outside. + assert ProfilesFile.retain_window(vt(), 1) == 1 + assert File.exists?(Pgrid.path_for(vt())) + refute File.exists?(Pgrid.path_for(~U[2026-04-19 20:00:00Z])) + end + end +end