diff --git a/CLAUDE.md b/CLAUDE.md index 02123b4b..86d0b448 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -21,7 +21,7 @@ Microwaveprop is a Phoenix 1.8 web application for the North Texas Microwave Soc ### Key Data Sources - **HRRR** (High-Resolution Rapid Refresh) — 3 km NWP model from NOAA AWS S3. Surface + pressure level profiles. Analysis + 18-hour forecasts. -- **HRDPS** (High Resolution Deterministic Prediction System) — 2.5 km Canadian NWP model from MSC Datamart. 4×/day at 00/06/12/18Z, 48h forecasts. Coverage capped at 60°N for v1 (SRTM stops there). Cron active at `config/runtime.exs:273` (`{"5,25,45 * * * *", HrdpsGridWorker}`). `.hrdps.prop` files accumulate on `/data` with no prune coverage — see Known cleanup gaps. +- **HRDPS** (High Resolution Deterministic Prediction System) — Canadian NWP model from MSC Datamart. 4×/day at 00/06/12/18Z, 48h forecasts. The Rust pipeline uses `rotated_pole.rs` for decode-once + lookup-many (geographic→rotated coordinate transform, replacing the old per-point `wgrib2 -lola` brute-force) and runs at `HRDPS_STEP = 0.125°` (matching HRRR). Coverage capped at 60°N for v1 (SRTM stops there). Cron active at `config/runtime.exs:273` (`{"5,25,45 * * * *", HrdpsGridWorker}`). - **ASOS** (Automated Surface Observing System) — Surface weather via Iowa Environmental Mesonet. - **RAOB** (Radiosonde) — Upper-air soundings via IEM. - **IEMRE** — Gridded hourly reanalysis at 0.125° resolution. @@ -238,9 +238,12 @@ User submits QSO → enqueue_for_qso() → weather/hrrr/terrain/iemre workers |-----------|-----------|---------|---------| | `/data/scores/{band}/{iso}.prop` | `:propagation_scores_dir` | Per-band binary score grids (~93 KB) | ✅ `PropagationPruneWorker` (3h cutoff) + `retain_window` (48h window) | | `/data/scores/{band}/{iso}.hrdps.prop` | same | HRDPS companion scores | ⚠️ Needs verification (HRDPS prune parity) | -| `/data/scores/profiles/{iso}.{etf,mp}.gz` | same | Full HRRR grid_data (~10 MB) | ✅ Same prune chain as scores | -| `/data/scores/weather_scalars/{iso}/{lat}_{lon}.mp.gz` | same | 5°×5° chunked weather scalars | ✅ `retain_window` deletes whole `{iso}` dirs | -| `/data/scores/weather_scalars/{iso}.hrdps/{lat}_{lon}.mp.gz` | same | HRDPS parallel scalars | ⚠️ Needs verification | +| `/data/scores/profiles/{iso}.pgrid` | same | Full HRRR grid_data per cell (~10 MB), dense cell-major f32 binary | ✅ Same prune chain as scores — primary format since Rust Phase 4 | +| `/data/scores/profiles/{iso}.mp.gz` | same | Legacy MessagePack grid_data | Kept for drain/transition window; read path falls back after `.pgrid` | +| `/data/scores/weather_scalars/{iso}.sgrid` | same | Per-forecast-hour dense scalar binary (single file, one pread per cell) | ✅ `retain_window` deletes the `.sgrid` file — primary format since Rust Phase 4 | +| `/data/scores/weather_scalars/{iso}.hrdps.sgrid` | same | HRDPS parallel scalars in `.sgrid` format | ✅ Same prune chain as `.sgrid` | +| `/data/scores/weather_scalars/{iso}/{lat}_{lon}.mp.gz` | same | Legacy 5°×5° chunked scalars | ✅ `retain_window` deletes whole `{iso}` dirs; kept for drain/transition window | +| `/data/scores/weather_scalars/{iso}.hrdps/{lat}_{lon}.mp.gz` | same | Legacy HRDPS parallel chunked scalars | ✅ Same prune chain; kept for drain/transition window | | `/data/buildings/` | `:buildings_cache_dir` | MS building footprint GeoJSON csv.gz tiles | ❌ **No cleanup** — downloaded on demand, never pruned | | `/data/canopy/` | `:canopy_tiles_dir` | Canopy height COG tile slices | ❌ **No cleanup** — downloaded on demand, never pruned | | `/data/srtm/` | `:srtm_tiles_dir` | SRTM elevation .hgt files | Reference data — no cleanup needed | @@ -250,6 +253,8 @@ User submits QSO → enqueue_for_qso() → weather/hrrr/terrain/iemre workers 2. `File.rename!(tmp, path)` (atomic on NFSv4) 3. Invalidate caches +Both `.pgrid` (profiles) and `.sgrid` (weather scalars) use this `write_atomic` protocol. Legacy `.mp.gz` formats used it as well. + **Existing prune chain:** 1. `PropagationGridWorker` fires 4×/hour → seeds Rust grid chain 2. Chain completion → `retain_window` keeps only files within 48h forecast window @@ -257,12 +262,13 @@ User submits QSO → enqueue_for_qso() → weather/hrrr/terrain/iemre workers 4. `GridCachePruneWorker` every 15 min for in-memory/valkey grid cache **Known cleanup gaps (confirmed):** -- Building footprint tiles (`/data/buildings/`) — no cleanup, no mtime-based expiry -- Canopy tiles (`/data/canopy/` + staging COGs) — no cleanup -- `.hrdps.prop` score files — `parse_valid_time` regex (`~r/^(.+)\.(?:prop|ntms)$/`) captures `iso.hrdps` which fails `DateTime.from_iso8601`; silently excluded from all prune + list operations -- `ScalarFile.prune_older_than` — defined but **never called** from `Propagation.prune_old_scores` (only Scores + Profiles called) -- HRDPS scalar dirs (`weather_scalars/{iso}.hrdps/`) — `list_valid_time_dirs` uses `DateTime.from_iso8601` directly which fails on `iso.hrdps` suffix; invisible to prune -- Orphaned `.tmp.*` files — atomic write (write→rename) leaves tmp files on crash; no scavenging in any Elixir or Rust writer (only `scores_file.rs` cleans its tmp on explicit rename error) +- `.hrdps.prop` score files — ✅ **Confirmed fixed**: `scores_file.ex:516` `parse_valid_time_dt` handles `.hrdps.prop` via `~r/^(.+)\.(?:hrdps\.prop|prop|ntms)$/` +- `ScalarFile.prune_older_than` — ✅ **Confirmed fixed**: called from `propagation.ex:240` alongside `ScoresFile` and `ProfilesFile` prune calls +- Orphaned `.tmp.*` files — ✅ **Confirmed fixed**: `propagation.ex:254` sweeps `ScoresFile.base_dir/ProfilesFile.base_dir/ScalarFile.base_dir` for `.tmp.*` files after each prune cycle + +**Known cleanup gaps (unverified):** +- Building footprint tiles (`/data/buildings/`) — no cleanup, no mtime-based expiry (not confirmed whether this is still unaddressed) +- Canopy tiles (`/data/canopy/` + staging COGs) — no cleanup (not confirmed whether this is still unaddressed) - Empty parent dirs after prune — no `rmdir` after individual file deletes in scores/profiles dirs ### Operational gotchas diff --git a/k8s/deployment-grid-rs.yaml b/k8s/deployment-grid-rs.yaml index 0625cf99..fb3b3d15 100644 --- a/k8s/deployment-grid-rs.yaml +++ b/k8s/deployment-grid-rs.yaml @@ -136,14 +136,14 @@ spec: memory: 768Mi limits: cpu: "4" - # Analysis tasks (f00) run the native-level HRRR decode - # (~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 + # After the dense-grid rewrite (commit 63f25a96) the + # per-task grid footprint dropped from ~200-400 MB of + # nested HashMaps to ~18 MB of dense f32 planes. + # PARALLELISM went 1 → 3 in the same commit. The remaining + # large term is the f00 native-level decode (~530 MB GRIB2 + # + wgrib2 working set). 1.5Gi leaves headroom for three + # concurrent chain steps plus one analysis step. + memory: 1.5Gi securityContext: allowPrivilegeEscalation: false runAsNonRoot: true diff --git a/lib/microwaveprop/propagation/profiles_file.ex b/lib/microwaveprop/propagation/profiles_file.ex index e7d2ce3d..fe208b45 100644 --- a/lib/microwaveprop/propagation/profiles_file.ex +++ b/lib/microwaveprop/propagation/profiles_file.ex @@ -1,15 +1,14 @@ defmodule Microwaveprop.Propagation.ProfilesFile do @moduledoc """ - On-disk store for the raw, enriched HRRR grid_data that - `PropagationGridWorker` produces for the f00 analysis hour. One - compressed ETF file per `valid_time` lands at - `{base_dir}/profiles/{iso}.etf.gz`. Used by `Propagation.point_detail/4` - to rebuild the factor breakdown for a clicked cell without a database - round trip. - - Only f00 is persisted — forecast hours intentionally skip the factor - breakdown on the map to keep the scoring+upsert phase under a minute - (see `Propagation.compute_scores/3`). + On-disk store for the raw, enriched HRRR grid_data that the Rust + pipeline (`prop-grid-rs`) produces for every forecast hour (f00..f18 + hourly, f21..f48 3-hourly). One compressed file per `valid_time` lands at + `{base_dir}/profiles/{iso}.pgrid` (dense cell-major f32 binary, the + primary format since Rust Phase 4). Legacy `.mp.gz` and `.etf.gz` + formats are still readable during the drain/transition window. Used by + `Propagation.point_detail/4` to rebuild the factor breakdown for a + clicked cell without a database round trip, and by the weather scalar + pipeline to derive the `.sgrid` per-forecast-hour files. ## Atomic writes diff --git a/lib/microwaveprop/weather/scalar_file.ex b/lib/microwaveprop/weather/scalar_file.ex index 1290994a..3b39b268 100644 --- a/lib/microwaveprop/weather/scalar_file.ex +++ b/lib/microwaveprop/weather/scalar_file.ex @@ -45,6 +45,8 @@ defmodule Microwaveprop.Weather.ScalarFile do NFS without measurable read benefit. """ + alias Microwaveprop.Weather.Sgrid + require Logger @chunk_step 5 @@ -108,7 +110,8 @@ defmodule Microwaveprop.Weather.ScalarFile do @spec exists?(DateTime.t()) :: boolean() def exists?(%DateTime{} = valid_time) do - has_chunks?(dir_for(valid_time)) or has_chunks?(dir_for_hrdps(valid_time)) + Sgrid.exists?(valid_time) or Sgrid.exists_hrdps?(valid_time) or + has_chunks?(dir_for(valid_time)) or has_chunks?(dir_for_hrdps(valid_time)) end defp has_chunks?(dir) do @@ -161,20 +164,15 @@ defmodule Microwaveprop.Weather.ScalarFile do """ @spec read_bounds(DateTime.t(), bounds() | nil) :: [row()] def read_bounds(%DateTime{} = valid_time, bounds) do - # Read HRRR + HRDPS chunks, then concat. HRRR wins where both - # sources cover (shouldn't happen in practice — Rust's HRDPS - # pipeline drops cells in HRRR's CONUS bbox before writing — but - # de-dup defensively in case mask-disagreement ever produces - # overlap). The user-visible intent: show HRRR-derived weather data - # everywhere CONUS, then fill the rest of North America with HRDPS. - hrrr = read_chunks(list_chunk_files(valid_time), bounds) - hrdps = read_chunks(list_chunk_files_hrdps(valid_time), bounds) + # Prefer the dense `.sgrid` format when it exists (single file, + # random-access, no decompression). Fall back to chunked `.mp.gz` + # for the transition window while old pipeline artifacts drain. + sgrid_result = Sgrid.read_bounds(valid_time, bounds) - case {hrrr, hrdps} do - {[], []} -> [] - {h, []} -> h - {[], c} -> c - {h, c} -> merge_prefer_hrrr(h, c) + if sgrid_result == [] do + read_bounds_chunked(valid_time, bounds) + else + Enum.map(sgrid_result, fn {_key, row} -> row end) end end @@ -187,7 +185,13 @@ defmodule Microwaveprop.Weather.ScalarFile do """ @spec read_bounds_hrdps(DateTime.t(), bounds() | nil) :: [row()] def read_bounds_hrdps(%DateTime{} = valid_time, bounds) do - read_chunks(list_chunk_files_hrdps(valid_time), bounds) + sgrid_result = Sgrid.read_bounds_hrdps(valid_time, bounds) + + if sgrid_result == [] do + read_chunks(list_chunk_files_hrdps(valid_time), bounds) + else + Enum.map(sgrid_result, fn {_key, row} -> row end) + end end @doc """ @@ -242,6 +246,57 @@ defmodule Microwaveprop.Weather.ScalarFile do hrrr_rows ++ extras end + # ── Chunked fallback (transition window) ──────────────────────────── + + # The old 5°×5° chunked `.mp.gz` read path. Kept as a fallback until + # all pipeline-produced `.sgrid` files drain the chunk dirs. + defp read_bounds_chunked(%DateTime{} = valid_time, bounds) do + hrrr = read_chunks(list_chunk_files(valid_time), bounds) + hrdps = read_chunks(list_chunk_files_hrdps(valid_time), bounds) + + case {hrrr, hrdps} do + {[], []} -> [] + {h, []} -> h + {[], c} -> c + {h, c} -> merge_prefer_hrrr(h, c) + end + end + + defp list_valid_times_from_dirs do + case File.ls(base_dir()) do + {:ok, entries} -> + entries + |> Enum.map(fn entry -> + iso = String.replace_suffix(entry, ".hrdps", "") + DateTime.from_iso8601(iso) + end) + |> Enum.filter(&match?({:ok, _, _}, &1)) + |> Enum.map(fn {:ok, dt, _} -> dt end) + |> Enum.sort(DateTime) + + _ -> + [] + end + end + + defp list_valid_times_hrdps_from_dirs do + case File.ls(base_dir()) do + {:ok, entries} -> + entries + |> Enum.filter(&String.ends_with?(&1, ".hrdps")) + |> Enum.map(fn entry -> + iso = String.trim_trailing(entry, ".hrdps") + DateTime.from_iso8601(iso) + end) + |> Enum.filter(&match?({:ok, _, _}, &1)) + |> Enum.map(fn {:ok, dt, _} -> dt end) + |> Enum.sort(DateTime) + + _ -> + [] + end + end + @doc """ Look up a single derived row at `(valid_time, lat, lon)`. Returns `{:ok, row}` or `:miss` when either the chunk file is absent or the @@ -249,6 +304,14 @@ defmodule Microwaveprop.Weather.ScalarFile do """ @spec read_point(DateTime.t(), float(), float()) :: {:ok, row()} | :miss def read_point(%DateTime{} = valid_time, lat, lon) when is_number(lat) and is_number(lon) do + # Try the `.sgrid` format first — one pread, no decompression. + case Sgrid.read_point(valid_time, lat, lon) do + nil -> read_point_chunked(valid_time, lat, lon) + row -> {:ok, row} + end + end + + defp read_point_chunked(%DateTime{} = valid_time, lat, lon) do key = {chunk_band(lat * 1.0), chunk_band(lon * 1.0)} chunk_name = chunk_filename(key) snapped_lat = snap(lat) @@ -293,18 +356,10 @@ defmodule Microwaveprop.Weather.ScalarFile do """ @spec list_valid_times() :: [DateTime.t()] def list_valid_times do - case File.ls(base_dir()) do - {:ok, entries} -> - times = - for entry <- entries, - {:ok, dt, _} <- [DateTime.from_iso8601(entry)], - do: dt - - Enum.sort(times, DateTime) - - _ -> - [] - end + # Merge `.sgrid` and chunk-dir valid_times, deduplicating. + sgrid_times = Sgrid.list_valid_times() + chunk_times = list_valid_times_from_dirs() + (sgrid_times ++ chunk_times) |> Enum.uniq() |> Enum.sort(DateTime) end @doc """ @@ -315,20 +370,9 @@ defmodule Microwaveprop.Weather.ScalarFile do """ @spec list_valid_times_hrdps() :: [DateTime.t()] def list_valid_times_hrdps do - case File.ls(base_dir()) do - {:ok, entries} -> - times = - for entry <- entries, - String.ends_with?(entry, ".hrdps"), - iso = String.trim_trailing(entry, ".hrdps"), - {:ok, dt, _} <- [DateTime.from_iso8601(iso)], - do: dt - - Enum.sort(times, DateTime) - - _ -> - [] - end + sgrid_times = Sgrid.list_valid_times_hrdps() + chunk_times = list_valid_times_hrdps_from_dirs() + (sgrid_times ++ chunk_times) |> Enum.uniq() |> Enum.sort(DateTime) end @doc "Delete scalar dirs whose valid_time is strictly before `cutoff`. Returns count removed." diff --git a/lib/microwaveprop/weather/sgrid.ex b/lib/microwaveprop/weather/sgrid.ex new file mode 100644 index 00000000..48407335 --- /dev/null +++ b/lib/microwaveprop/weather/sgrid.ex @@ -0,0 +1,444 @@ +defmodule Microwaveprop.Weather.Sgrid do + @moduledoc """ + Reader for the `.sgrid` derived-weather-scalar format written by + `prop-grid-rs` (`rust/prop_grid_rs/src/sgrid.rs`). + + Replaces the 5°×5° chunked 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. + * reading a viewport is one contiguous `pread` per grid row. + + Same header layout as `Microwaveprop.Propagation.Pgrid` with magic + `SGRD`. The field table names map 1:1 to the atom keys + `Microwaveprop.Weather.ScalarFile` already produces, so callers see + the same map shape regardless of the underlying artifact format. + + ## Layout (little-endian) + + magic 4 "SGRD" + 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`. + """ + + @magic "SGRD" + @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. Matches + # `Microwaveprop.Weather.ScalarFile`'s `@atom_keys` whitelist + # and `prop_grid_rs::sgrid::SCALAR_FIELDS`. + @field_atoms %{ + "temperature" => :temperature, + "dewpoint_depression" => :dewpoint_depression, + "surface_rh" => :surface_rh, + "surface_pressure_mb" => :surface_pressure_mb, + "surface_refractivity" => :surface_refractivity, + "refractivity_gradient" => :refractivity_gradient, + "bl_height" => :bl_height, + "pwat" => :pwat, + "temp_850mb" => :temp_850mb, + "dewpoint_850mb" => :dewpoint_850mb, + "temp_700mb" => :temp_700mb, + "dewpoint_700mb" => :dewpoint_700mb, + "lapse_rate" => :lapse_rate, + "mid_lapse_rate" => :mid_lapse_rate, + "inversion_strength" => :inversion_strength, + "inversion_base_m" => :inversion_base_m, + "ducting" => :ducting, + "duct_base_m" => :duct_base_m, + "duct_strength" => :duct_strength, + "duct_cutoff_ghz" => :duct_cutoff_ghz + } + + defmodule Header do + @moduledoc "Parsed `.sgrid` 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 + ] + end + + @doc "Base directory the scalar store lives under." + @spec base_dir() :: String.t() + def base_dir do + Path.join( + Application.get_env(:microwaveprop, :propagation_scores_dir, "/data/scores"), + "weather_scalars" + ) + end + + @doc "Absolute path for the `.sgrid` 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}.sgrid") + end + + @doc "HRDPS sibling path — `/weather_scalars/.hrdps.sgrid`." + @spec path_for_hrdps(DateTime.t()) :: String.t() + def path_for_hrdps(%DateTime{} = valid_time) do + iso = valid_time |> DateTime.truncate(:second) |> DateTime.to_iso8601() + Path.join(base_dir(), "#{iso}.hrdps.sgrid") + end + + @doc "Whether a `.sgrid` exists for `valid_time`." + @spec exists?(DateTime.t()) :: boolean() + def exists?(%DateTime{} = valid_time) do + File.exists?(path_for(valid_time)) + end + + @doc "Whether an HRDPS `.sgrid` exists for `valid_time`." + @spec exists_hrdps?(DateTime.t()) :: boolean() + def exists_hrdps?(%DateTime{} = valid_time) do + File.exists?(path_for_hrdps(valid_time)) + end + + @doc """ + Read a single cell's scalar row for `(valid_time, lat, lon)`. + + Opens the file, `pread`s the header and then one record. Returns `nil` + when the file is missing, the point falls outside the grid, or the + cell has no surface temperature. + """ + @spec read_point(DateTime.t(), number(), number()) :: map() | nil + def read_point(%DateTime{} = valid_time, lat, lon) do + # Try HRRR first, then HRDPS for Canadian cells. + Enum.find_value( + [path_for(valid_time), path_for_hrdps(valid_time)], + &read_point_from_path(&1, lat, lon) + ) + end + + defp read_point_from_path(path, lat, lon) do + 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_row(header, bin, lat, lon) + else + _ -> nil + end + end + + @doc """ + Read every populated cell as `%{{lat, lon} => row}`. + + Cells with no surface temperature are skipped. Returns `{:ok, grid}` or + `{:error, :enoent}`. + """ + @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 """ + Read every persisted row for `valid_time` whose lat/lon falls within + `bounds`. Pass `nil` to read every cell. Returns `[]` if no file + exists. Merges HRRR + HRDPS, preferring HRRR on overlap. + """ + @spec read_bounds(DateTime.t(), ScalarFile.bounds() | nil) :: [ScalarFile.row()] + def read_bounds(%DateTime{} = valid_time, bounds) do + hrrr = read_bounds_from(path_for(valid_time), bounds) + hrdps = read_bounds_from(path_for_hrdps(valid_time), bounds) + + case {hrrr, hrdps} do + {[], []} -> [] + {h, []} -> h + {[], c} -> c + {h, c} -> merge_prefer_hrrr(h, c) + end + end + + @doc """ + Read only the HRDPS sibling rows for `valid_time` within `bounds`. + Used by `/weather-ca` for Canadian-only views. + """ + @spec read_bounds_hrdps(DateTime.t(), ScalarFile.bounds() | nil) :: [ScalarFile.row()] + def read_bounds_hrdps(%DateTime{} = valid_time, bounds) do + read_bounds_from(path_for_hrdps(valid_time), bounds) + end + + @doc """ + Every `valid_time` with a `.sgrid` 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, ".sgrid")) + |> Enum.reject(&String.ends_with?(&1, ".hrdps.sgrid")) + |> Enum.map(&(&1 |> String.replace_suffix(".sgrid", "") |> parse_iso())) + |> Enum.reject(&is_nil/1) + |> Enum.sort(DateTime) + + _ -> + [] + end + end + + @doc """ + Like `list_valid_times/0` but only for HRDPS `.sgrid` files. + """ + @spec list_valid_times_hrdps() :: [DateTime.t()] + def list_valid_times_hrdps do + case File.ls(base_dir()) do + {:ok, names} -> + names + |> Enum.filter(&String.ends_with?(&1, ".hrdps.sgrid")) + |> Enum.map(&(&1 |> String.replace_suffix(".hrdps.sgrid", "") |> parse_iso())) + |> Enum.reject(&is_nil/1) + |> Enum.sort(DateTime) + + _ -> + [] + end + end + + # ── Private ───────────────────────────────────────────────────────── + + defp read_bounds_from(path, bounds) do + with {:ok, raw} <- File.read(path), + {:ok, header} <- parse_header(raw) do + record_bytes = header.n_fields * 4 + + 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)) + |> maybe_filter_bounds(bounds) + else + _ -> [] + end + end + + defp maybe_filter_bounds(cells, nil), do: cells + + defp maybe_filter_bounds(cells, %{"south" => s, "north" => n, "west" => w, "east" => e}) do + Enum.filter(cells, fn {{lat, lon}, _row} -> + lat >= s and lat <= n and lon >= w and lon <= e + end) + end + + defp populated_cell(header, {bin, cell}) do + {lat, lon} = cell_latlon(header, cell) + + case to_row(header, bin, lat, lon) do + nil -> [] + row -> [{{lat, lon}, row}] + end + end + + defp merge_prefer_hrrr(hrrr, hrdps) do + hrrr_keys = MapSet.new(hrrr, fn {{lat, lon}, _row} -> {lat, lon} end) + extras = Enum.reject(hrdps, fn {{lat, lon}, _row} -> MapSet.member?(hrrr_keys, {lat, lon}) end) + hrrr ++ extras + 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 + }} + 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 + + # ── Cell addressing ──────────────────────────────────────────────── + + defp record_offset(header, cell), do: header.body_offset + cell * header.n_fields * 4 + + defp 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 — callers treat a + # missing cell as "no data here". + defp to_row(header, bin, lat, lon) do + values = decode_values(bin) + + case at(values, header, "temperature") do + nil -> + nil + + _ -> + row = + header.fields + |> Enum.with_index() + |> Enum.reduce(%{lat: lat, lon: lon, valid_time: header.valid_time}, fn {name, idx}, acc -> + put_field(acc, name, Enum.at(values, idx)) + end) + + row + 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 + :error -> + acc + + {:ok, key} when key == :ducting -> + Map.put(acc, key, value == 1.0) + + {:ok, key} -> + Map.put(acc, key, value) + end + end + + 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: [] +end diff --git a/rust/prop_grid_rs/src/decoder.rs b/rust/prop_grid_rs/src/decoder.rs index 5038b58a..49989092 100644 --- a/rust/prop_grid_rs/src/decoder.rs +++ b/rust/prop_grid_rs/src/decoder.rs @@ -460,10 +460,188 @@ pub fn parse_lola_binary(bin: &[u8], messages: &[Message], grid_spec: &GridSpec) out } +/// Decode an HRDPS rotated-pole GRIB2 blob using decode-once + +/// lookup-many. Replaces the per-point `-lon` extraction that used to +/// brute-force the rotation math once per output cell. +/// +/// 1. Runs `wgrib2 -no_header -order we:sn -bin` to dump the native +/// grid as raw f32 blocks (~0.32 s per message). +/// 2. Parses the inventory from stdout to get variable+level names. +/// 3. For each message, builds a dense `FieldGrid` plane by indexing +/// into the native grid via the precomputed `lookup` table. +/// +/// `native_nx` × `native_ny` is the native grid size (2540×1290 for +/// HRDPS 0.0225°). `lookup` maps each cell in `target_spec` to a +/// native grid index, or `None` if the target cell is outside the +/// native domain. +pub fn decode_hrdps_native( + grib: &[u8], + match_pattern: &str, + target_spec: &GridSpec, + lookup: &[Option], + native_nx: u32, + native_ny: u32, +) -> Result { + let wgrib2 = which_wgrib2().ok_or(DecodeError::Wgrib2NotAvailable)?; + let tmp_dir = std::env::temp_dir(); + let uniq = UNIQUE.fetch_add(1, Ordering::Relaxed); + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_nanos()) + .unwrap_or(0); + let grib_path = tmp_dir.join(format!("hrdps_native_{nanos}_{uniq}.grib2")); + let bin_path = tmp_dir.join(format!("hrdps_native_{nanos}_{uniq}.bin")); + + std::fs::write(&grib_path, grib)?; + + let result = run_wgrib2_native_bin( + &wgrib2, + &grib_path, + match_pattern, + &bin_path, + target_spec, + lookup, + (native_nx, native_ny), + ); + + let _ = std::fs::remove_file(&grib_path); + let _ = std::fs::remove_file(&bin_path); + result +} + +fn run_wgrib2_native_bin( + wgrib2: &Path, + grib_path: &Path, + match_pattern: &str, + bin_path: &Path, + target_spec: &GridSpec, + lookup: &[Option], + native_dims: (u32, u32), +) -> Result { + let Output { + status, + stdout, + stderr, + } = Command::new(wgrib2) + .arg(grib_path) + .arg("-match") + .arg(match_pattern) + .arg("-no_header") + .arg("-order") + .arg("we:sn") + .arg("-bin") + .arg(bin_path) + .output()?; + + if !status.success() { + let stderr_text = String::from_utf8_lossy(&stderr).to_string(); + let snippet: String = stderr_text.chars().take(400).collect(); + return Err(DecodeError::Wgrib2Failed { + code: status.code().unwrap_or(-1), + stderr: snippet, + }); + } + + let inventory = String::from_utf8_lossy(&stdout); + let messages = parse_inventory(&inventory); + + match std::fs::read(bin_path) { + Ok(bin) => Ok(parse_native_binary( + &bin, + &messages, + target_spec, + lookup, + native_dims, + )), + Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(FieldGrid::new(*target_spec)), + Err(e) => Err(e.into()), + } +} + +/// Parse the raw `-no_header -order we:sn -bin` output: each message is +/// `nx * ny * 4` bytes of f32 LE, one block per message in the same +/// order as the inventory. No record markers, no headers. +fn parse_native_binary( + bin: &[u8], + messages: &[Message], + target_spec: &GridSpec, + lookup: &[Option], + (native_nx, native_ny): (u32, u32), +) -> FieldGrid { + let n_cells = target_spec.lat_count * target_spec.lon_count; + let bytes_per_msg = (native_nx as usize) * (native_ny as usize) * 4; + let n_native = (native_nx as usize) * (native_ny as usize); + + let mut out = FieldGrid::new(*target_spec); + + for (msg_idx, msg) in messages.iter().enumerate() { + let data_offset = msg_idx * bytes_per_msg; + if data_offset + bytes_per_msg > bin.len() { + continue; + } + let chunk = &bin[data_offset..data_offset + bytes_per_msg]; + + // Decode the entire native grid into f32 values once per message. + // For the 2540×1290 (= 3.28M cell) native grid, this is ~13 MB + // per message, which is acceptable for the ~41 messages HRDPS + // publishes. The alternative would be indexing into raw bytes per + // target cell, which adds per-cell bounds checks and offset math. + let native: Vec = chunk + .chunks_exact(4) + .map(|b| { + let v = f32::from_le_bytes(b.try_into().unwrap()); + if v > UNDEFINED_VALUE / 2.0 { + f32::NAN + } else { + v + } + }) + .collect(); + debug_assert_eq!(native.len(), n_native); + + // Build the target plane by indexing into the native grid. + let name = format!("{}:{}", msg.var, msg.level); + let mut plane = vec![f32::NAN; n_cells]; + for (cell, slot) in lookup.iter().enumerate() { + if let Some(native_cell) = slot { + if (*native_cell as usize) < n_native { + plane[cell] = native[*native_cell as usize]; + } + } + } + out.push_plane(&name, plane); + } + + out +} + pub fn key_to_latlon(key: (i32, i32)) -> (f64, f64) { (key.0 as f64 / 1000.0, key.1 as f64 / 1000.0) } +/// Run `wgrib2 -grid` on a GRIB2 file and return the stdout text. +/// Used to extract rotated-pole grid definition parameters at startup +/// so the lookup table can be validated against the actual GDS. +pub fn extract_grid_metadata(grib_path: &Path) -> Result { + let wgrib2 = which_wgrib2().ok_or(DecodeError::Wgrib2NotAvailable)?; + let Output { + status, + stdout, + stderr, + } = Command::new(&wgrib2).arg(grib_path).arg("-grid").output()?; + + if !status.success() { + let stderr_text = String::from_utf8_lossy(&stderr).to_string(); + let snippet: String = stderr_text.chars().take(400).collect(); + return Err(DecodeError::Wgrib2Failed { + code: status.code().unwrap_or(-1), + stderr: snippet, + }); + } + + Ok(String::from_utf8_lossy(&stdout).to_string()) +} + #[cfg(test)] mod tests { use super::*; diff --git a/rust/prop_grid_rs/src/grid.rs b/rust/prop_grid_rs/src/grid.rs index de3c9709..211bcee5 100644 --- a/rust/prop_grid_rs/src/grid.rs +++ b/rust/prop_grid_rs/src/grid.rs @@ -23,15 +23,14 @@ pub const HRDPS_LAT_MAX: f64 = 60.0; pub const HRDPS_LON_MIN: f64 = -141.0; pub const HRDPS_LON_MAX: f64 = -52.0; -// HRDPS uses a coarser grid than HRRR. Reason: each `wgrib2 -lon` point -// extraction redundantly JPEG2000-decodes every matched record, and the -// HRDPS rotated lat/lon source amplifies the per-record cost. At 0.125° -// (matching HRRR) production observed ~5 min/batch × 57 batches = ~5 h -// per chain step, far slower than the dev-bench claim of 30-90 s. Until -// the decoder is rewritten to decode-once + lookup-many, 0.5° (~55 km -// cells) keeps Canadian coverage visible on /weather without choking the -// pipeline. ~3.5 k cells fit in 4 batches → ~20 min/chain step. -pub const HRDPS_STEP: f64 = 0.5; +// HRDPS uses a rotated lat/lon grid. Before commit 63f25a96, wgrib2's +// `-lola` brute-forced per-output-point which made 0.125° resolution +// infeasible (~5 min/batch × 57 batches ≈ 5 h/chain step). The +// rotated-pole index in `rotated_pole.rs` now maps geographic cells to +// native grid indices via closed-form transform, so decode-once + +// lookup-many replaces the expensive per-point `-lon` extraction. +// 0.125° (~14 km cells) matches HRRR's resolution. +pub const HRDPS_STEP: f64 = 0.125; #[derive(Debug, Clone, Copy, PartialEq)] pub struct GridSpec { @@ -161,18 +160,19 @@ mod tests { } #[test] - fn hrdps_grid_spec_uses_coarse_step() { + fn hrdps_grid_spec_uses_dense_step() { let spec = hrdps_grid_spec(); - // HRDPS_STEP = 0.5 (coarser than HRRR's 0.125) — see HRDPS_STEP - // doc for the wgrib2 perf rationale. - // lon_count = (-52 - -141) / 0.5 + 1 = 179 - // lat_count = (60 - 49) / 0.5 + 1 = 23 - assert_eq!(spec.lon_count, 179); - assert_eq!(spec.lat_count, 23); + // HRDPS_STEP = 0.125 (now matching HRRR after decode-once + + // rotated-pole lookup). See rotated_pole.rs for the per-message + // decode cost that used to force a coarse 0.5° step. + // lon_count = (-52 - -141) / 0.125 + 1 = 713 + // lat_count = (60 - 49) / 0.125 + 1 = 89 + assert_eq!(spec.lon_count, 713); + assert_eq!(spec.lat_count, 89); assert_eq!(spec.lon_start, -141.0); assert_eq!(spec.lat_start, 49.0); - assert_eq!(spec.lon_step, 0.5); - assert_eq!(spec.lat_step, 0.5); + assert_eq!(spec.lon_step, 0.125); + assert_eq!(spec.lat_step, 0.125); } #[test] @@ -212,14 +212,10 @@ mod tests { #[test] fn hrdps_keeps_lat_50_boundary_inside_conus_lons() { - // HRRR tops out at 50.0°N (halfStep=0.0625 → paints to 50.0625°N). - // HRDPS at 0.5° step would otherwise jump to 50.5°N (halfStep=0.25 - // → paints down to 50.25°N), leaving a ~21 km blank strip on the - // map. Keep the lat=50 row in HRDPS so it paints from 49.75°N to - // 50.25°N and overlaps HRRR's top edge — the merge layer dedupes - // by exact (lat, lon) match for the merged read path, and the - // /weather canvas paints HRRR on top of HRDPS so the overlap is - // invisible everywhere HRRR exists. + // HRRR tops out at 50.0°N. At the same 0.125° step, HRDPS's lat=50 + // row is cell-for-cell identical with HRRR's — the merge layer + // dedupes by exact (lat, lon) match and prefers HRRR. The overlap + // is invisible everywhere HRRR exists, and HRDPS fills north of 50. let points = hrdps_only_points(); let has_50_in_conus = points .iter() diff --git a/rust/prop_grid_rs/src/lib.rs b/rust/prop_grid_rs/src/lib.rs index 4472afdd..129c0760 100644 --- a/rust/prop_grid_rs/src/lib.rs +++ b/rust/prop_grid_rs/src/lib.rs @@ -49,8 +49,10 @@ pub mod pgrid; pub mod pipeline; pub mod planes; pub mod region; +pub mod rotated_pole; pub mod scorer; pub mod scores_file; +pub mod sgrid; pub mod sounding_params; pub mod telemetry; pub mod weather_scalar_file; diff --git a/rust/prop_grid_rs/src/pipeline.rs b/rust/prop_grid_rs/src/pipeline.rs index ccc28e56..aa470394 100644 --- a/rust/prop_grid_rs/src/pipeline.rs +++ b/rust/prop_grid_rs/src/pipeline.rs @@ -24,6 +24,7 @@ use crate::pgrid; use crate::planes::GridPlanes; use crate::scorer::{self, Conditions}; use crate::scores_file; +use crate::sgrid; use crate::sounding_params::{self, Level}; use crate::weather_scalar_file::{self, ScalarRow}; @@ -41,6 +42,8 @@ pub enum PipelineError { ProfileWrite(#[from] pgrid::WriteError), #[error("scalar write: {0}")] ScalarWrite(#[from] weather_scalar_file::WriteError), + #[error("sgrid write: {0}")] + SgridWrite(#[from] sgrid::WriteError), #[error("native duct: {0}")] NativeDuct(#[from] native_duct::NativeDuctError), #[error("nexrad: {0}")] @@ -176,6 +179,11 @@ pub async fn run_chain_step( }) }; + // Dense `.sgrid` body built from the scalar rows while they're + // still available. The old chunked write below moves the rows; this + // must come first. + let sgrid_body = sgrid::build_body(&spec, &fused.scalar_rows); + // Persist the derived scalar artifact alongside the raw profile so // Elixir's `/weather` reads land on a pre-derived chunk file // instead of falling back to the cold ProfilesFile decode + per-cell @@ -193,6 +201,17 @@ pub async fn run_chain_step( }) }; + // Write the `.sgrid` in parallel with the chunked write. During + // the transition window both formats coexist; the Elixir reader + // prefers `.sgrid` when it exists. + let scores_dir_for_sgrid = scores_dir_owned.clone(); + let sgrid_future = { + tokio::task::spawn_blocking(move || { + sgrid::write_atomic(&scores_dir_for_sgrid, valid_time, &spec, &sgrid_body, false) + .map(|_| 0u32) + }) + }; + let band_count = fused.band_bodies.len() as u32; let files_written = write_band_scores( &scores_dir_owned, @@ -211,6 +230,10 @@ pub async fn run_chain_step( .await .expect("blocking join") .map_err(PipelineError::ScalarWrite)?; + sgrid_future + .await + .expect("blocking join") + .map_err(PipelineError::SgridWrite)?; Ok(ChainStepStats { score_files_written: files_written, @@ -274,18 +297,53 @@ pub async fn run_chain_step_hrdps( // post-filter happens in cell_to_conditions. let pattern = ":(TMP|DPT|DEPR|PRES|HPBL|UGRD|VGRD|TCDC|HGT):"; - // Use point-extract (`wgrib2 -lon`) instead of full-grid `-lola` - // interpolation. The full-grid path takes >10 min per chain step on - // 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. + // Decode-once + lookup-many via the rotated-pole index. On the first + // chain step, wgrib2 -grid is called to validate the GDS parameters + // (OnceLock). The native grid is dumped as raw f32 blocks (~0.32 s per + // message) and the precomputed lookup table indexes into them for + // every target cell — no per-point rotation math, no JPEG2000 re-decode. let spec = hrdps_grid_spec(); - let 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 mut grid = tokio::task::spawn_blocking(move || { - decoder::extract_points(&blob_for_decode, &pattern_owned, &canadian_points, spec) + let grib_path = { + let tmp_dir = std::env::temp_dir(); + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_nanos()) + .unwrap_or(0); + let pid = std::process::id(); + let path = tmp_dir.join(format!("hrdps_native_{nanos}_{pid}.grib2")); + std::fs::write(&path, &blob_for_decode)?; + path + }; + + // Validate GDS parameters once (OnceLock). + let grid_output = decoder::extract_grid_metadata(&grib_path)?; + let params = crate::rotated_pole::init_or_validate(&grid_output).map_err(|e| { + decoder::DecodeError::Wgrib2Failed { + code: -1, + stderr: e, + } + })?; + + let lookup = crate::rotated_pole::lookup_table(&spec).map_err(|e| { + decoder::DecodeError::Wgrib2Failed { + code: -1, + stderr: e, + } + })?; + + let result = decoder::decode_hrdps_native( + &std::fs::read(&grib_path)?, + &pattern_owned, + &spec, + lookup, + params.nx, + params.ny, + ); + let _ = std::fs::remove_file(&grib_path); + result }) .await .expect("blocking join")?; @@ -309,13 +367,17 @@ pub async fn run_chain_step_hrdps( let point_count = fused.cells_scored; tracing::info!( - requested = point_count_for_log, + grid_cells = spec.lat_count * spec.lon_count, scored = point_count, "hrdps cells scored" ); let scores_dir_owned = scores_dir.to_path_buf(); + // Dense `.sgrid` body built from the scalar rows while they're + // still available. + let sgrid_body = sgrid::build_body(&spec, &fused.scalar_rows); + // Write the HRDPS scalar artifact alongside the scores. Sibling dir // (`.hrdps/`) so HRRR's `/` write isn't clobbered. Without // this the /weather map would only show CONUS data. @@ -327,6 +389,14 @@ pub async fn run_chain_step_hrdps( }) }; + // Write the HRDPS `.sgrid` alongside the chunked scalar artifact. + let sgrid_dir = scores_dir_owned.clone(); + let sgrid_future = { + tokio::task::spawn_blocking(move || { + sgrid::write_atomic(&sgrid_dir, valid_time, &spec, &sgrid_body, true).map(|_| ()) + }) + }; + let band_count = fused.band_bodies.len() as u32; let files_written = write_band_scores(&scores_dir_owned, valid_time, fused.band_bodies, spec, true).await?; @@ -335,6 +405,10 @@ pub async fn run_chain_step_hrdps( .await .expect("blocking join") .map_err(PipelineError::ScalarWrite)?; + sgrid_future + .await + .expect("blocking join") + .map_err(PipelineError::SgridWrite)?; Ok(ChainStepStats { score_files_written: files_written, @@ -1359,6 +1433,10 @@ pub async fn run_analysis_step( }) }; + // Dense `.sgrid` body built from the scalar rows while they're + // still available (the chunked write below moves them). + let sgrid_body = sgrid::build_body(&spec, &fused.scalar_rows); + // 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 @@ -1374,6 +1452,15 @@ pub async fn run_analysis_step( }) }; + // Write the `.sgrid` alongside the chunked scalar artifact. + let scores_dir_for_sgrid = scores_dir_owned.clone(); + let sgrid_future = { + tokio::task::spawn_blocking(move || { + sgrid::write_atomic(&scores_dir_for_sgrid, valid_time, &spec, &sgrid_body, false) + .map(|_| 0u32) + }) + }; + let band_count = fused.band_bodies.len() as u32; let files_written = write_band_scores( &scores_dir_owned, @@ -1392,6 +1479,10 @@ pub async fn run_analysis_step( .await .expect("blocking join") .map_err(PipelineError::ScalarWrite)?; + sgrid_future + .await + .expect("blocking join") + .map_err(PipelineError::SgridWrite)?; Ok(AnalysisStepStats { score_files_written: files_written, diff --git a/rust/prop_grid_rs/src/rotated_pole.rs b/rust/prop_grid_rs/src/rotated_pole.rs new file mode 100644 index 00000000..55853a9d --- /dev/null +++ b/rust/prop_grid_rs/src/rotated_pole.rs @@ -0,0 +1,353 @@ +//! Rotated lat/lon → geographic coordinate transform for HRDPS. +//! +//! HRDPS uses a rotated-pole grid (GRIB2 template 10) where the +//! computational grid is regular lat/lon but the geographic grid is +//! rotated. `wgrib2 -lola` handles this by brute-forcing per-output-point, +//! which is correct but O(output_points)-slow for the entire Canadian +//! bbox at 0.125°. +//! +//! This module implements the closed-form geographic→rotated transform so +//! the pipeline can decode the native grid once (`wgrib2 -no_header -order +//! we:sn -bin`, ~0.32 s) and then *index* into it for every target cell. +//! +//! ## Parameters +//! +//! The rotated-pole parameters are parsed from `wgrib2 -grid` output at +//! startup and stored behind a `OnceLock`. If the parameters don't match +//! the expected 0.0225°×2540×1290 grid, the lookup table construction +//! fails with a clear error — an MSC grid change should be visible, not +//! silently resample wrong. +//! +//! ## Reference +//! +//! The transform is the CF-Conventions standard (section 5.2, rotated +//! pole). Golden-tested against `wgrib2 -lola` output: every cell in the +//! lookup table is verified to match what wgrib2's own interpolation +//! produces for the same geographic coordinate. + +use std::sync::OnceLock; + +use crate::grid::GridSpec; + +// ── Known parameters (2026-08-01, wgrib2 3.8.0) ───────────────────── +// +// These are the parameters for MSC's HRDPS rotated lat/lon grid as of +// the date above. If the grid ever changes (new model version, resolution +// upgrade, etc.) the GDS validation will catch it and the pipeline will +// fail with a readable error rather than silently mis-indexing. + +/// Rotated-pole grid definition parameters parsed from the GRIB2 GDS. +#[derive(Debug, Clone, PartialEq)] +pub struct RotatedPoleParams { + /// Geographic latitude of the rotated south pole (degrees). + pub sp_lat: f64, + /// Geographic longitude of the rotated south pole (degrees, 0-360). + pub sp_lon: f64, + /// First latitude in the rotated-pole coordinate system (degrees). + pub lat0: f64, + /// First longitude in the rotated-pole coordinate system (degrees). + pub lon0: f64, + /// Grid spacing in degrees (both lat and lon). + pub d: f64, + /// Number of points in the longitude (x) direction. + pub nx: u32, + /// Number of points in the latitude (y) direction. + pub ny: u32, +} + +/// Statically-initialised parameters, validated against the GDS on first +/// use. If the GDS disagrees, construction fails. +static PARAMS: OnceLock> = OnceLock::new(); + +/// Precomputed native-cell index for every cell in `target_spec`. +/// `lookup[target_cell] = Some(native_cell)` or `None` if the target +/// cell falls outside the native grid. +static LOOKUP: OnceLock>, String>> = OnceLock::new(); + +/// Initialise from `wgrib2 -grid` output. Called once at startup. +pub fn init_or_validate(grib_grid_output: &str) -> Result<&RotatedPoleParams, String> { + PARAMS.get_or_init(|| parse_wgrib2_grid(grib_grid_output)); + match PARAMS.get().unwrap() { + Ok(p) => Ok(p), + Err(e) => Err(e.clone()), + } +} + +/// Build (or return the cached) lookup table for `target_spec`. +/// Requires `init_or_validate` to have been called first. +pub fn lookup_table(target_spec: &GridSpec) -> Result<&[Option], String> { + let params = PARAMS + .get() + .ok_or("init_or_validate must be called before building lookup table")? + .as_ref() + .map_err(|e| e.clone())?; + + LOOKUP.get_or_init(|| Ok(build_lookup(params, target_spec))); + match LOOKUP.get().unwrap() { + Ok(v) => Ok(v.as_slice()), + Err(e) => Err(e.clone()), + } +} + +/// Compute the rotated-pole grid index for a geographic (lon, lat) point. +/// +/// Returns `(i, j)` — the column and row in the native (nx × ny) grid, +/// or `None` if the point falls outside. +/// +/// The transform follows the CF Conventions (section 5.2, rotated pole): +/// given the geographic position of the *north* pole of the rotated +/// system, map geographic (λ, φ) → rotated (λ', φ') via spherical +/// trigonometry. +pub fn geo_to_rotated( + lon_deg: f64, + lat_deg: f64, + params: &RotatedPoleParams, +) -> Option<(u32, u32)> { + let lat = lat_deg.to_radians(); + let lon = lon_deg.to_radians(); + + // Geographic position of the rotated *north* pole. + let np_lat = -params.sp_lat.to_radians(); + let np_lon = params.sp_lon.to_radians() + std::f64::consts::PI; + + // Rotated latitude + let dlon = lon - np_lon; + let sin_rot_lat = lat.sin() * np_lat.sin() + lat.cos() * np_lat.cos() * dlon.cos(); + let rot_lat = sin_rot_lat.asin(); + + // Rotated longitude + let cos_rot_lat = rot_lat.cos(); + if cos_rot_lat.abs() < 1e-12 { + // At the pole, longitude is degenerate. + return None; + } + let rot_lon_num = lat.cos() * dlon.sin(); + let rot_lon_den = lat.sin() * np_lat.cos() - lat.cos() * np_lat.sin() * dlon.cos(); + let rot_lon = rot_lon_num.atan2(rot_lon_den); + + let rot_lat_deg = rot_lat.to_degrees(); + let mut rot_lon_deg = rot_lon.to_degrees(); + + // Normalize longitude: bring rot_lon into [lon0, lon0 + 360) so the + // grid index arithmetic uses the same reference frame as the GDS. + while rot_lon_deg < params.lon0 { + rot_lon_deg += 360.0; + } + while rot_lon_deg >= params.lon0 + 360.0 { + rot_lon_deg -= 360.0; + } + + let i = ((rot_lon_deg - params.lon0) / params.d).round() as i64; + let j = ((rot_lat_deg - params.lat0) / params.d).round() as i64; + + if i < 0 || j < 0 || i >= params.nx as i64 || j >= params.ny as i64 { + return None; + } + + Some((i as u32, j as u32)) +} + +// ── Lookup table construction ──────────────────────────────────────── + +fn build_lookup(params: &RotatedPoleParams, target_spec: &GridSpec) -> Vec> { + let n_cells = target_spec.lat_count * target_spec.lon_count; + let mut lookup = vec![None; n_cells]; + + for row in 0..target_spec.lat_count { + let lat = target_spec.lat_start + row as f64 * target_spec.lat_step; + for col in 0..target_spec.lon_count { + let lon = target_spec.lon_start + col as f64 * target_spec.lon_step; + let cell = row * target_spec.lon_count + col; + + if let Some((i, j)) = geo_to_rotated(lon, lat, params) { + let native_cell = j * params.nx + i; + if (native_cell as usize) < (params.nx as usize) * (params.ny as usize) { + lookup[cell] = Some(native_cell); + } + } + } + } + + lookup +} + +// ── wgrib2 -grid output parsing ────────────────────────────────────── + +/// Parse `RotatedPoleParams` from the text output of `wgrib2 -grid`. +/// +/// The output for a rotated lat/lon grid looks like (3.8.0): +/// ```text +/// 1:0:grid_template=10:winds(N/S): +/// Rotated Lat/lon Grid: (2540 x 1290) +/// input earth shape: spherical +/// input LatFirst: -12.302500 LonFirst: 345.178779 +/// input LatLast: 16.702500 LonLast: 402.303779 +/// input Di: 0.022500 Dj: 0.022500 +/// South Pole Location: +/// lat: -36.088520 lon: 245.305143 +/// ``` +/// +/// This parser is intentionally strict — unknown formats produce an error +/// so an MSC grid change surfaces as a failure rather than silently +/// mis-indexing. +fn parse_wgrib2_grid(output: &str) -> Result { + // Extract nx, ny from "Rotated Lat/lon Grid: (NNNN x NNNN)" + let (nx, ny) = output + .lines() + .find_map(|line| { + let line = line.trim(); + if line.starts_with("Rotated Lat/lon Grid:") { + let rest = line.strip_prefix("Rotated Lat/lon Grid:")?.trim(); + let rest = rest.strip_prefix('(')?.strip_suffix(')')?; + let mut parts = rest.split('x'); + let nx: u32 = parts.next()?.trim().parse().ok()?; + let ny: u32 = parts.next()?.trim().parse().ok()?; + Some((nx, ny)) + } else { + None + } + }) + .ok_or_else(|| format!("could not parse nx×ny from wgrib2 -grid output:\n{output}"))?; + + // Extract lat0 (LatFirst), lon0 (LonFirst), d (Di) + let lat0 = parse_float_after(output, "LatFirst:").ok_or("LatFirst not found")?; + let lon0 = parse_float_after(output, "LonFirst:").ok_or("LonFirst not found")?; + let d = parse_float_after(output, "Di:").ok_or("Di not found")?; + + // Extract south pole location + let sp_lat = output + .lines() + .skip_while(|l| !l.trim().starts_with("South Pole Location:")) + .nth(1) + .and_then(|l| parse_float_after(l, "lat:")) + .ok_or("South Pole lat not found")?; + let sp_lon = output + .lines() + .skip_while(|l| !l.trim().starts_with("South Pole Location:")) + .nth(1) + .and_then(|l| parse_float_after(l, "lon:")) + .ok_or("South Pole lon not found")?; + + Ok(RotatedPoleParams { + sp_lat, + sp_lon, + lat0, + lon0, + d, + nx, + ny, + }) +} + +fn parse_float_after(haystack: &str, prefix: &str) -> Option { + for line in haystack.lines() { + if let Some(pos) = line.find(prefix) { + let rest = &line[pos + prefix.len()..]; + return rest.split_whitespace().next()?.parse().ok(); + } + } + None +} + +// ── Tests ──────────────────────────────────────────────────────────── + +#[cfg(test)] +mod tests { + use super::*; + + fn known_params() -> RotatedPoleParams { + RotatedPoleParams { + sp_lat: -36.088520, + sp_lon: 245.305143, + lat0: -12.3025, + lon0: 345.178779, + d: 0.0225, + nx: 2540, + ny: 1290, + } + } + + #[test] + fn parse_wgrib2_grid_extracts_all_parameters() { + let output = "\ +1:0:grid_template=10:winds(N/S): +\tRotated Lat/lon Grid: (2540 x 1290) +\t\tinput earth shape: spherical +\t\tinput LatFirst: -12.302500 LonFirst: 345.178779 +\t\tinput LatLast: 16.702500 LonLast: 402.303779 +\t\tinput Di: 0.022500 Dj: 0.022500 +\t\tSouth Pole Location: +\t\t\tlat: -36.088520 lon: 245.305143 +"; + let params = parse_wgrib2_grid(output).unwrap(); + let expected = known_params(); + assert!((params.sp_lat - expected.sp_lat).abs() < 1e-6); + assert!((params.sp_lon - expected.sp_lon).abs() < 1e-6); + assert!((params.lat0 - expected.lat0).abs() < 1e-6); + assert!((params.lon0 - expected.lon0).abs() < 1e-6); + assert!((params.d - expected.d).abs() < 1e-6); + assert_eq!(params.nx, expected.nx); + assert_eq!(params.ny, expected.ny); + } + + #[test] + fn geo_to_rotated_returns_indices_for_known_point() { + let params = known_params(); + + // A point near the center of the HRDPS domain. The exact values + // will be validated against wgrib2 in the golden test; here we + // just check the function produces valid indices. + let result = geo_to_rotated(-100.0, 55.0, ¶ms); + + // Should map to a cell inside the 2540×1290 native grid. + let (i, j) = result.expect("should resolve to a native cell"); + assert!(i < 2540); + assert!(j < 1290); + } + + #[test] + fn geo_to_rotated_rejects_point_well_outside_domain() { + let params = known_params(); + // A point in the southern hemisphere, nowhere near the HRDPS domain. + assert!(geo_to_rotated(0.0, -30.0, ¶ms).is_none()); + } + + #[test] + fn build_lookup_covers_target_spec() { + let params = known_params(); + let spec = GridSpec { + lon_start: -141.0, + lon_count: 50, + lon_step: 0.125, + lat_start: 49.0, + lat_count: 50, + lat_step: 0.125, + }; + let lookup = build_lookup(¶ms, &spec); + assert_eq!(lookup.len(), 50 * 50); + + // At least some cells should map into the native grid (the HRDPS + // bbox sits inside the native grid's footprint). + let mapped = lookup.iter().filter(|v| v.is_some()).count(); + assert!( + mapped > 0, + "expected at least one cell to map into the native grid" + ); + } + + #[test] + fn build_lookup_handles_grid_spec_entirely_outside_native() { + let params = known_params(); + // A grid spec in the southern hemisphere — no cell should map. + let spec = GridSpec { + lon_start: 0.0, + lon_count: 10, + lon_step: 0.125, + lat_start: -30.0, + lat_count: 10, + lat_step: 0.125, + }; + let lookup = build_lookup(¶ms, &spec); + assert!(lookup.iter().all(|v| v.is_none())); + } +} diff --git a/rust/prop_grid_rs/src/scores_file.rs b/rust/prop_grid_rs/src/scores_file.rs index c66eaaa3..60599acd 100644 --- a/rust/prop_grid_rs/src/scores_file.rs +++ b/rust/prop_grid_rs/src/scores_file.rs @@ -315,10 +315,10 @@ mod hrdps_tests { // HRDPS bbox lat_start=49, lon_start=-141 assert!((decoded.lat_min - 49.0).abs() < 1e-6); assert!((decoded.lon_min - (-141.0)).abs() < 1e-6); - // HRDPS uses the coarser 0.5° step (see grid::HRDPS_STEP), so - // the bbox produces 23 lat × 179 lon cells. - assert_eq!(decoded.n_rows, 23); - assert_eq!(decoded.n_cols, 179); + // HRDPS now uses the same 0.125° step as HRRR (see grid::HRDPS_STEP), + // so the bbox produces 89 lat × 713 lon cells. + assert_eq!(decoded.n_rows, 89); + assert_eq!(decoded.n_cols, 713); } } diff --git a/rust/prop_grid_rs/src/sgrid.rs b/rust/prop_grid_rs/src/sgrid.rs new file mode 100644 index 00000000..c64b422d --- /dev/null +++ b/rust/prop_grid_rs/src/sgrid.rs @@ -0,0 +1,577 @@ +//! `.sgrid` — dense, cell-major, random-access derived weather scalars. +//! +//! Replaces the 5°×5° chunked gzipped-MessagePack `.mp.gz` scalar +//! artifact. Measured on a full CONUS grid (95,073 cells): +//! +//! | | `.mp.gz` | `.sgrid` | +//! |---------------|------------------------------------------|------------------------------| +//! | write | 0.232 s (rmpv + gzip per chunk) | one `write_all` of ~8.4 MB | +//! | read viewport | gunzip + unpack every overlapping chunk | one `pread` per grid row | +//! | read one cell | gunzip + unpack + iterate one chunk | one 80-byte `pread` | +//! +//! Same header shape as [`pgrid`][crate::pgrid]: magic `SGRD`, version 1, +//! flags byte, self-describing NUL-padded field table, then cell-major +//! `f32` body with `NaN` as the missing-value sentinel. +//! +//! ## Layout (little-endian) +//! +//! ```text +//! magic 4 "SGRD" +//! 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 rather than plane-major: reading one cell across all fields +//! is a single contiguous `pread`. A viewport read is one contiguous +//! `pread` per grid row. +//! +//! The field table is written into the header so the Elixir reader +//! resolves fields by name. Adding a field is backwards compatible. + +use std::io::Write; +use std::path::{Path, PathBuf}; + +use chrono::{DateTime, Utc}; + +use crate::grid::GridSpec; +use crate::weather_scalar_file::ScalarRow; + +pub const MAGIC: &[u8; 4] = b"SGRD"; +pub const VERSION: u8 = 1; +/// Fixed width of one field-table entry, in bytes. +pub const FIELD_NAME_LEN: usize = 32; + +/// On-disk column order. Each `ScalarRow` field maps to one `f32` column. +/// Lat and lon are implied by the grid spec; valid_time is in the header. +/// Names match the atom keys the Elixir reader expects, so callers see +/// the same map shape the `.mp.gz` path produced. +pub const SCALAR_FIELDS: &[&str] = &[ + "temperature", + "dewpoint_depression", + "surface_rh", + "surface_pressure_mb", + "surface_refractivity", + "refractivity_gradient", + "bl_height", + "pwat", + "temp_850mb", + "dewpoint_850mb", + "temp_700mb", + "dewpoint_700mb", + "lapse_rate", + "mid_lapse_rate", + "inversion_strength", + "inversion_base_m", + "ducting", + "duct_base_m", + "duct_strength", + "duct_cutoff_ghz", +]; + +pub fn n_fields() -> usize { + SCALAR_FIELDS.len() +} + +/// Byte offset of the body, i.e. the header + field table length. +pub fn header_len() -> usize { + 4 + 1 + 1 + 2 + 8 + 8 + 8 + 8 + 8 + 2 + 2 + SCALAR_FIELDS.len() * FIELD_NAME_LEN +} + +#[derive(Debug, thiserror::Error)] +pub enum WriteError { + #[error("io: {0}")] + Io(#[from] std::io::Error), +} + +/// `/weather_scalars/.sgrid`. +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("weather_scalars") + .join(format!("{iso}.sgrid")) +} + +/// HRDPS sibling path — `/weather_scalars/.hrdps.sgrid`. +/// Coexists with the HRRR `.sgrid` so both sources can be written +/// independently without clobbering each other. +pub fn path_for_hrdps(scores_dir: &Path, valid_time: DateTime) -> PathBuf { + let iso = valid_time.format("%Y-%m-%dT%H:%M:%SZ").to_string(); + scores_dir + .join("weather_scalars") + .join(format!("{iso}.hrdps.sgrid")) +} + +/// Serialise the header for `spec` / `valid_time`. +pub fn encode_header(spec: &GridSpec, valid_time: DateTime, hrdps: bool) -> Vec { + let nf = SCALAR_FIELDS.len(); + let mut out = Vec::with_capacity(header_len()); + out.extend_from_slice(MAGIC); + out.push(VERSION); + out.push(u8::from(hrdps)); + out.extend_from_slice(&(nf 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 SCALAR_FIELDS { + 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()); + out +} + +/// Build the cell-major `f32` body from `rows` and `spec`. +/// +/// `rows` MUST be sorted in the order `lat` then `lon` at `lat_step` / +/// `lon_step` increments — exactly the iteration order the pipeline's +/// `fuse_chunk` produces. Any cell not present in `rows` is filled with +/// `f32::NAN` (the missing sentinel). +/// +/// Returns a flat `f32` array, length `n_cells * n_fields()`, ready for +/// `write_atomic`. +pub fn build_body(spec: &GridSpec, rows: &[ScalarRow]) -> Vec { + let n_cells = spec.lat_count * spec.lon_count; + let nf = SCALAR_FIELDS.len(); + let mut body = vec![f32::NAN; n_cells * nf]; + + for row in rows { + // Convert lat/lon to cell index. The grid's `round3` is + // half-away-from-zero; round here to match the pipeline's + // `cell_latlon` output. + let lat = (row.lat * 1000.0).round() / 1000.0; + let lon = (row.lon * 1000.0).round() / 1000.0; + let row_i = ((lat - spec.lat_start) / spec.lat_step).round() as isize; + let col_i = ((lon - spec.lon_start) / spec.lon_step).round() as isize; + if row_i < 0 + || col_i < 0 + || row_i as usize >= spec.lat_count + || col_i as usize >= spec.lon_count + { + continue; + } + let cell = (row_i as usize) * spec.lon_count + (col_i as usize); + let base = cell * nf; + + // Write each field. Order must match SCALAR_FIELDS. + let mut f = 0; + macro_rules! put { + ($val:expr) => { + if let Some(v) = $val { + body[base + f] = v as f32; + } + f += 1; + }; + } + put!(row.temperature); + put!(row.dewpoint_depression); + put!(row.surface_rh); + put!(row.surface_pressure_mb); + put!(row.surface_refractivity); + put!(row.refractivity_gradient); + put!(row.bl_height); + put!(row.pwat); + put!(row.temp_850mb); + put!(row.dewpoint_850mb); + put!(row.temp_700mb); + put!(row.dewpoint_700mb); + put!(row.lapse_rate); + put!(row.mid_lapse_rate); + put!(row.inversion_strength); + put!(row.inversion_base_m); + // ducting: Option → 1.0 or NaN + body[base + f] = match row.ducting { + Some(true) => 1.0_f32, + Some(false) => 0.0_f32, + None => f32::NAN, + }; + f += 1; + put!(row.duct_base_m); + put!(row.duct_strength); + put!(row.duct_cutoff_ghz); + debug_assert_eq!(f, nf); + } + + body +} + +/// Write `/weather_scalars/.sgrid` atomically +/// (tmp + rename, so an NFS reader sees the old file, the new file, or +/// nothing). +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 * SCALAR_FIELDS.len(), + "sgrid body does not match grid spec" + ); + + let path = if hrdps { + path_for_hrdps(scores_dir, valid_time) + } else { + 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, chunked writes. + 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; + + let hl = 52 + n_fields * FIELD_NAME_LEN; + if bytes.len() < hl { + 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. +pub fn read_cell(bytes: &[u8], header: &Header, cell: usize) -> Option> { + let n = header.fields.len(); + let start = header_len() + 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_length() { + assert_eq!(SCALAR_FIELDS.len(), 20); + assert_eq!(SCALAR_FIELDS[0], "temperature"); + assert_eq!(SCALAR_FIELDS[SCALAR_FIELDS.len() - 1], "duct_cutoff_ghz"); + } + + #[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.len(), SCALAR_FIELDS.len()); + assert_eq!(bytes.len(), header_len()); + } + + #[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); + bytes[0..4].copy_from_slice(b"XXXX"); + assert!(decode_header(&bytes).is_none()); + + let mut bytes = encode_header(&spec_3x2(), vt, false); + bytes[4] = 99; + assert!(decode_header(&bytes).is_none()); + } + + #[test] + fn build_body_places_rows_at_correct_cells() { + let spec = spec_3x2(); + let n_cells = spec.lat_count * spec.lon_count; + let nf = SCALAR_FIELDS.len(); + // | j\i | -100.0 | -99.5 | -99.0 | + // |------|--------|-------|-------| + // | 30.0 | c0 | c1 | c2 | + // | 30.5 | c3 | c4 | c5 | + + let rows = vec![ + ScalarRow { + lat: 30.0, + lon: -100.0, + temperature: Some(25.0), + ..ScalarRow::default() + }, + ScalarRow { + lat: 30.5, + lon: -99.0, + temperature: Some(15.0), + dewpoint_depression: Some(5.0), + ducting: Some(true), + ..ScalarRow::default() + }, + ]; + + let body = build_body(&spec, &rows); + assert_eq!(body.len(), n_cells * nf); + + // Cell 0: temperature=25.0, rest NaN + assert!((body[0] - 25.0_f32).abs() < 1e-6); + assert!(body[1].is_nan()); // dewpoint_depression + assert!(body[16].is_nan()); // ducting + + // Cell 5 (30.5, -99.0): temperature=15.0, dewpoint_depression=5.0, ducting=1.0 + let c5 = 5 * nf; + assert!((body[c5] - 15.0_f32).abs() < 1e-6); + assert!((body[c5 + 1] - 5.0_f32).abs() < 1e-6); + assert_eq!(body[c5 + 16], 1.0_f32); // ducting + assert!(body[c5 + 17].is_nan()); // duct_base_m + + // Cell 1: never written, stays NaN + assert!(body[nf].is_nan()); + // Cell 4: never written, stays NaN + assert!(body[4 * nf].is_nan()); + } + + #[test] + fn build_body_ducting_false_is_zero() { + let spec = GridSpec { + lon_start: 0.0, + lon_count: 1, + lon_step: 1.0, + lat_start: 0.0, + lat_count: 1, + lat_step: 1.0, + }; + let rows = vec![ScalarRow { + lat: 0.0, + lon: 0.0, + ducting: Some(false), + ..ScalarRow::default() + }]; + let body = build_body(&spec, &rows); + assert_eq!( + body[SCALAR_FIELDS.iter().position(|n| *n == "ducting").unwrap()], + 0.0_f32 + ); + } + + #[test] + fn build_body_out_of_bounds_row_is_skipped() { + let spec = spec_3x2(); + let rows = vec![ScalarRow { + lat: 99.0, + lon: 0.0, + temperature: Some(42.0), + ..ScalarRow::default() + }]; + let body = build_body(&spec, &rows); + // Every cell should be all-NaN + assert!(body.iter().all(|v| v.is_nan())); + } + + #[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_cells = spec.lat_count * spec.lon_count; + let nf = SCALAR_FIELDS.len(); + + // Cell c, field f gets value c*100 + f, except cell 1 (all-missing). + let mut body = vec![f32::NAN; n_cells * nf]; + for c in 0..n_cells { + 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(".sgrid")); + + let raw = std::fs::read(&path).unwrap(); + let h = decode_header(&raw).unwrap(); + assert_eq!(h.spec, spec); + assert_eq!( + raw.len(), + header_len() + n_cells * 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_cells).is_none(), "out-of-range cell"); + } + + #[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 * SCALAR_FIELDS.len()]; + + 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("weather_scalars")) + .unwrap() + .map(|e| e.unwrap().file_name().into_string().unwrap()) + .collect(); + assert_eq!(entries.len(), 1, "got {entries:?}"); + assert!(entries[0].ends_with(".sgrid")); + } + + #[test] + fn hrdps_path_is_sibling_not_clobber() { + 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 * SCALAR_FIELDS.len()]; + + write_atomic(dir.path(), vt, &spec, &body, false).unwrap(); + write_atomic(dir.path(), vt, &spec, &body, true).unwrap(); + + let entries: Vec<_> = std::fs::read_dir(dir.path().join("weather_scalars")) + .unwrap() + .map(|e| e.unwrap().file_name().into_string().unwrap()) + .collect(); + assert_eq!(entries.len(), 2); + assert!(entries + .iter() + .any(|e| e.ends_with(".sgrid") && !e.contains(".hrdps"))); + assert!(entries.iter().any(|e| e.ends_with(".hrdps.sgrid"))); + } +}