perf(grid-rs): dense grid, fused scoring pass, and columnar .pgrid profiles
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<Arc<str>, 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 <iso>.hrdps.prop no longer parse as
<iso>.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.
This commit is contained in:
parent
ed8bb33acb
commit
63f25a9612
22 changed files with 3805 additions and 1313 deletions
|
|
@ -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
|
||||
|
|
|
|||
427
lib/microwaveprop/propagation/pgrid.ex
Normal file
427
lib/microwaveprop/propagation/pgrid.ex
Normal file
|
|
@ -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(<<v::float-little-32>>), do: v
|
||||
defp decode_f32(_), do: nil
|
||||
|
||||
defp chunk_records(bin, size) when byte_size(bin) >= size do
|
||||
for <<chunk::binary-size(^size) <- bin>>, 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
|
||||
|
|
@ -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
|
||||
# `<iso>.hrdps.prop` from parsing as `<iso>.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
|
||||
|
|
|
|||
|
|
@ -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};
|
||||
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<str>` 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<str>: Borrow<str>`.
|
||||
pub type CellValues = HashMap<std::sync::Arc<str>, 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<PointGrid, DecodeError> {
|
||||
grid_spec: GridSpec,
|
||||
) -> Result<FieldGrid, DecodeError> {
|
||||
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<PointGrid, DecodeError> {
|
||||
let mut out: PointGrid = HashMap::with_capacity(points.len());
|
||||
grid_spec: GridSpec,
|
||||
) -> Result<FieldGrid, DecodeError> {
|
||||
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<String, Vec<(usize, f32)>>;
|
||||
|
||||
fn run_wgrib2_lon_one(
|
||||
wgrib2: &Path,
|
||||
grib_path: &Path,
|
||||
match_pattern: &str,
|
||||
batch: &[(f64, f64)],
|
||||
) -> Result<PointGrid, DecodeError> {
|
||||
grid_spec: &GridSpec,
|
||||
) -> Result<LonBatchValues, DecodeError> {
|
||||
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<str> = 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<PointGrid, DecodeError> {
|
||||
) -> Result<FieldGrid, DecodeError> {
|
||||
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<PointGrid, DecodeError> = (|| {
|
||||
let result: Result<FieldGrid, DecodeError> = (|| {
|
||||
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<PointGrid, DecodeError> {
|
||||
) -> Result<FieldGrid, DecodeError> {
|
||||
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<PointGrid, DecodeError> {
|
||||
) -> Result<FieldGrid, DecodeError> {
|
||||
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<Message> {
|
|||
/// 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<str> = 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]
|
||||
|
|
|
|||
|
|
@ -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<Result<(u64, Vec<u8>), 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<u8>)> = Vec::with_capacity(merged.len());
|
||||
while let Some(item) = futs.join_next().await {
|
||||
// If a range-download task panics (JoinError), log it and
|
||||
|
|
|
|||
370
rust/prop_grid_rs/src/field_grid.rs
Normal file
370
rust/prop_grid_rs/src/field_grid.rs
Normal file
|
|
@ -0,0 +1,370 @@
|
|||
//! Dense struct-of-arrays storage for a decoded NWP grid.
|
||||
//!
|
||||
//! Replaces the previous `HashMap<(i32, i32), HashMap<Arc<str>, 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<usize> {
|
||||
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<Arc<str>>,
|
||||
index: HashMap<Arc<str>, PlaneId>,
|
||||
planes: Vec<f32>,
|
||||
}
|
||||
|
||||
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<str>] {
|
||||
&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<PlaneId> {
|
||||
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<f32> {
|
||||
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<PlaneId>` 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<PlaneId>, cell: usize) -> Option<f32> {
|
||||
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<f32>) {
|
||||
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<str> = 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<str> = 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<usize> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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<ProfileRow> = Vec::with_capacity(points.len());
|
||||
let mut levels_buf: Vec<Level> = 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<Utc>,
|
||||
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<serde_json::Value>,
|
||||
hpbl_m: Option<f64>,
|
||||
pwat_mm: Option<f64>,
|
||||
surface_temp_c: Option<f64>,
|
||||
surface_dewpoint_c: Option<f64>,
|
||||
surface_pressure_mb: Option<f64>,
|
||||
surface_refractivity: Option<f64>,
|
||||
min_refractivity_gradient: Option<f64>,
|
||||
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<Json<Value>> 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<Level> = Vec::with_capacity(fetcher::grid_level_keys().len());
|
||||
let levels: Vec<sqlx::types::Json<serde_json::Value>> = fetcher::grid_level_keys()
|
||||
fn build_row(
|
||||
lat: f64,
|
||||
lon: f64,
|
||||
grid: &FieldGrid,
|
||||
p: &GridPlanes,
|
||||
cell: usize,
|
||||
typed_levels: &[Level],
|
||||
) -> ProfileRow {
|
||||
let profile: Vec<serde_json::Value> = 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<Json<Vec<Value>>>` 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<Utc>,
|
||||
rows: &[ProfileRow],
|
||||
) -> Result<u32, sqlx::Error> {
|
||||
if rows.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let ids: Vec<Uuid> = rows.iter().map(|_| Uuid::new_v4()).collect();
|
||||
let lats: Vec<f64> = rows.iter().map(|r| r.lat).collect();
|
||||
let lons: Vec<f64> = rows.iter().map(|r| r.lon).collect();
|
||||
let profiles: Vec<sqlx::types::Json<&Vec<serde_json::Value>>> =
|
||||
rows.iter().map(|r| sqlx::types::Json(&r.profile)).collect();
|
||||
let hpbl: Vec<Option<f64>> = rows.iter().map(|r| r.hpbl_m).collect();
|
||||
let pwat: Vec<Option<f64>> = rows.iter().map(|r| r.pwat_mm).collect();
|
||||
let temp: Vec<Option<f64>> = rows.iter().map(|r| r.surface_temp_c).collect();
|
||||
let dewp: Vec<Option<f64>> = rows.iter().map(|r| r.surface_dewpoint_c).collect();
|
||||
let pres: Vec<Option<f64>> = rows.iter().map(|r| r.surface_pressure_mb).collect();
|
||||
let sref: Vec<Option<f64>> = rows.iter().map(|r| r.surface_refractivity).collect();
|
||||
let mgrad: Vec<Option<f64>> = rows.iter().map(|r| r.min_refractivity_gradient).collect();
|
||||
let duct: Vec<bool> = 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<f64>, 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<f64>) = 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"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -48,6 +48,42 @@ pub static TASKS_IN_FLIGHT: LazyLock<IntGauge> = 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<HistogramVec> = 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<T>(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!(
|
||||
|
|
|
|||
|
|
@ -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<HashMap<(i32, i32), DuctMetrics>, NativeDuctError> {
|
||||
) -> Result<HashMap<usize, DuctMetrics>, 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<PlaneId>, Option<PlaneId>)>,
|
||||
}
|
||||
|
||||
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<usize, DuctMetrics> {
|
||||
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<NativeProfile> {
|
||||
// (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<str>'s Borrow<str>, 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<NativeProfile> {
|
||||
if levels.len() < 3 {
|
||||
return None;
|
||||
}
|
||||
|
|
@ -169,31 +194,35 @@ pub fn build_native_profile(cell: &CellValues) -> Option<NativeProfile> {
|
|||
|
||||
/// 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<usize, DuctMetrics>) {
|
||||
// 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<HashMap<(i32, i32), DuctMetrics>, NativeDuctError> {
|
||||
) -> Result<HashMap<usize, DuctMetrics>, 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));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -124,11 +124,12 @@ pub async fn fetch_frame(
|
|||
timestamp: DateTime<Utc>,
|
||||
points: &[(f64, f64)],
|
||||
) -> Result<Vec<NexradObservation>, 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)
|
||||
|
|
|
|||
453
rust/prop_grid_rs/src/pgrid.rs
Normal file
453
rust/prop_grid_rs/src/pgrid.rs
Normal file
|
|
@ -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_<P>mb`, `tmpc_<P>mb`,
|
||||
/// `dwpc_<P>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<Vec<String>> = std::sync::OnceLock::new();
|
||||
NAMES.get_or_init(|| {
|
||||
let mut v: Vec<String> = 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),
|
||||
}
|
||||
|
||||
/// `<scores_dir>/profiles/<iso>.pgrid`.
|
||||
pub fn path_for(scores_dir: &Path, valid_time: DateTime<Utc>) -> 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<f32>;
|
||||
|
||||
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<Utc>, hrdps: bool) -> Vec<u8> {
|
||||
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 `<scores_dir>/profiles/<iso>.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<Utc>,
|
||||
spec: &GridSpec,
|
||||
body: &[f32],
|
||||
hrdps: bool,
|
||||
) -> Result<PathBuf, WriteError> {
|
||||
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<u8> = 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<Utc>,
|
||||
pub spec: GridSpec,
|
||||
pub fields: Vec<String>,
|
||||
}
|
||||
|
||||
impl Header {
|
||||
pub fn field_index(&self, name: &str) -> Option<usize> {
|
||||
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<Header> {
|
||||
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::<Utc>::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<Vec<f32>> {
|
||||
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"));
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load diff
372
rust/prop_grid_rs/src/planes.rs
Normal file
372
rust/prop_grid_rs/src/planes.rs
Normal file
|
|
@ -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<PlaneId>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct GridPlanes {
|
||||
// Surface / column fields.
|
||||
pub tmp_2m: Option<PlaneId>,
|
||||
pub dpt_2m: Option<PlaneId>,
|
||||
pub ugrd_10m: Option<PlaneId>,
|
||||
pub vgrd_10m: Option<PlaneId>,
|
||||
pub tcdc: Option<PlaneId>,
|
||||
pub pwat: Option<PlaneId>,
|
||||
pub pres_sfc: Option<PlaneId>,
|
||||
pub hpbl: Option<PlaneId>,
|
||||
pub apcp: Option<PlaneId>,
|
||||
|
||||
// f00 enrichment planes.
|
||||
pub nexrad_dbz: Option<PlaneId>,
|
||||
pub native_min_gradient: Option<PlaneId>,
|
||||
pub best_duct_freq_ghz: Option<PlaneId>,
|
||||
pub max_duct_thickness_m: Option<PlaneId>,
|
||||
pub duct_count: Option<PlaneId>,
|
||||
pub commercial_degradation_db: Option<PlaneId>,
|
||||
pub commercial_baseline_dbm: Option<PlaneId>,
|
||||
pub commercial_current_dbm: Option<PlaneId>,
|
||||
pub commercial_n_links: Option<PlaneId>,
|
||||
|
||||
/// One entry per `fetcher::GRID_PRESSURE_LEVELS` level that is
|
||||
/// actually present in the grid, in the same order.
|
||||
pub levels: Vec<LevelPlanes>,
|
||||
}
|
||||
|
||||
impl GridPlanes {
|
||||
pub fn resolve(grid: &FieldGrid) -> Self {
|
||||
let levels = fetcher::grid_level_keys()
|
||||
.iter()
|
||||
.filter_map(|k| {
|
||||
// TMP and HGT are both required — the old per-cell
|
||||
// filter used `?` on each, so a level missing either
|
||||
// was skipped entirely.
|
||||
let tmp = grid.plane_id(&k.tmp)?;
|
||||
let hgt = grid.plane_id(&k.hgt)?;
|
||||
Some(LevelPlanes {
|
||||
pres_mb: k.pres_mb,
|
||||
tmp,
|
||||
hgt,
|
||||
dpt: grid.plane_id(&k.dpt),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
Self {
|
||||
tmp_2m: grid.plane_id("TMP:2 m above ground"),
|
||||
dpt_2m: grid.plane_id("DPT:2 m above ground"),
|
||||
ugrd_10m: grid.plane_id("UGRD:10 m above ground"),
|
||||
vgrd_10m: grid.plane_id("VGRD:10 m above ground"),
|
||||
tcdc: grid.plane_id("TCDC:entire atmosphere"),
|
||||
pwat: grid.plane_id("PWAT:entire atmosphere (considered as a single layer)"),
|
||||
pres_sfc: grid.plane_id("PRES:surface"),
|
||||
hpbl: grid.plane_id("HPBL:surface"),
|
||||
apcp: grid.plane_id("APCP:surface"),
|
||||
nexrad_dbz: grid.plane_id("nexrad_max_reflectivity_dbz"),
|
||||
native_min_gradient: grid.plane_id("native_min_gradient"),
|
||||
best_duct_freq_ghz: grid.plane_id("best_duct_freq_ghz"),
|
||||
max_duct_thickness_m: grid.plane_id("max_duct_thickness_m"),
|
||||
duct_count: grid.plane_id("duct_count"),
|
||||
commercial_degradation_db: grid.plane_id("commercial_degradation_db"),
|
||||
commercial_baseline_dbm: grid.plane_id("commercial_baseline_dbm"),
|
||||
commercial_current_dbm: grid.plane_id("commercial_current_dbm"),
|
||||
commercial_n_links: grid.plane_id("commercial_n_links"),
|
||||
levels,
|
||||
}
|
||||
}
|
||||
|
||||
/// Fill `out` with this cell's pressure-level profile, ascending by
|
||||
/// pressure index exactly as `fetcher::grid_level_keys()` orders it.
|
||||
/// Levels whose TMP or HGT is missing at this cell are skipped, which
|
||||
/// matches the `filter_map` the old per-cell code used.
|
||||
///
|
||||
/// Takes `&mut Vec` so callers reuse one buffer across all ~95 k
|
||||
/// cells instead of allocating a fresh `Vec<Level>` three times per
|
||||
/// cell (once each for conditions, profile, and scalars).
|
||||
#[inline]
|
||||
pub fn levels_at(&self, grid: &FieldGrid, cell: usize, out: &mut Vec<Level>) {
|
||||
out.clear();
|
||||
for lp in &self.levels {
|
||||
let (Some(t), Some(h)) = (grid.at(lp.tmp, cell), grid.at(lp.hgt, cell)) else {
|
||||
continue;
|
||||
};
|
||||
out.push(Level {
|
||||
pres_mb: lp.pres_mb,
|
||||
hght_m: h as f64,
|
||||
tmpc: (t as f64) - 273.15,
|
||||
dwpc: grid.at_opt(lp.dpt, cell).map(|v| (v as f64) - 273.15),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Fill one `.pgrid` record (`pgrid::field_names()` order) for `cell`.
|
||||
///
|
||||
/// Writes straight into the caller's flat body slice, so the profile
|
||||
/// artifact costs one `f32` store per field per cell — no `rmpv::Value`
|
||||
/// tree, no per-level `String` keys, no compression pass. That is the
|
||||
/// whole point of the format change: the old `.mp.gz` write measured
|
||||
/// 3.96 s per CONUS grid against 0.039 s for the entire derivation.
|
||||
pub fn fill_pgrid_record(
|
||||
grid: &FieldGrid,
|
||||
p: &GridPlanes,
|
||||
cell: usize,
|
||||
levels: &[Level],
|
||||
out: &mut [f32],
|
||||
) {
|
||||
use crate::pgrid;
|
||||
|
||||
debug_assert_eq!(out.len(), pgrid::n_fields());
|
||||
let idx = pgrid_scalar_indices();
|
||||
|
||||
let mut put = |slot: usize, v: Option<f32>| {
|
||||
out[slot] = v.unwrap_or(f32::NAN);
|
||||
};
|
||||
|
||||
put(
|
||||
idx.surface_temp_c,
|
||||
grid.at_opt(p.tmp_2m, cell).map(|v| v - 273.15),
|
||||
);
|
||||
put(
|
||||
idx.surface_dewpoint_c,
|
||||
grid.at_opt(p.dpt_2m, cell).map(|v| v - 273.15),
|
||||
);
|
||||
put(
|
||||
idx.surface_pressure_mb,
|
||||
grid.at_opt(p.pres_sfc, cell).map(|v| v / 100.0),
|
||||
);
|
||||
put(idx.hpbl_m, grid.at_opt(p.hpbl, cell));
|
||||
put(idx.pwat_mm, grid.at_opt(p.pwat, cell));
|
||||
put(idx.wind_u, grid.at_opt(p.ugrd_10m, cell));
|
||||
put(idx.wind_v, grid.at_opt(p.vgrd_10m, cell));
|
||||
put(idx.cloud_cover_pct, grid.at_opt(p.tcdc, cell));
|
||||
put(idx.precip_mm, grid.at_opt(p.apcp, cell));
|
||||
put(
|
||||
idx.native_min_gradient,
|
||||
grid.at_opt(p.native_min_gradient, cell),
|
||||
);
|
||||
put(
|
||||
idx.best_duct_freq_ghz,
|
||||
grid.at_opt(p.best_duct_freq_ghz, cell),
|
||||
);
|
||||
put(
|
||||
idx.max_duct_thickness_m,
|
||||
grid.at_opt(p.max_duct_thickness_m, cell),
|
||||
);
|
||||
put(idx.duct_count, grid.at_opt(p.duct_count, cell));
|
||||
put(
|
||||
idx.nexrad_max_reflectivity_dbz,
|
||||
grid.at_opt(p.nexrad_dbz, cell),
|
||||
);
|
||||
put(
|
||||
idx.commercial_degradation_db,
|
||||
grid.at_opt(p.commercial_degradation_db, cell),
|
||||
);
|
||||
put(
|
||||
idx.commercial_baseline_dbm,
|
||||
grid.at_opt(p.commercial_baseline_dbm, cell),
|
||||
);
|
||||
put(
|
||||
idx.commercial_current_dbm,
|
||||
grid.at_opt(p.commercial_current_dbm, cell),
|
||||
);
|
||||
put(
|
||||
idx.commercial_n_links,
|
||||
grid.at_opt(p.commercial_n_links, cell),
|
||||
);
|
||||
|
||||
// Pre-computed refractivity scalars: the Skew-T page and contact
|
||||
// detail need a fallback for when SoundingParams.derive yields nil
|
||||
// because some level lacks dewpoint.
|
||||
put(
|
||||
idx.surface_refractivity,
|
||||
crate::sounding_params::surface_refractivity(levels).map(|v| v as f32),
|
||||
);
|
||||
put(
|
||||
idx.min_refractivity_gradient,
|
||||
crate::sounding_params::min_refractivity_gradient(levels.to_vec()).map(|v| v as f32),
|
||||
);
|
||||
|
||||
// Pressure levels. `levels` is filtered (levels missing TMP/HGT at
|
||||
// this cell are absent), so match each back to its slot by pressure
|
||||
// rather than by position.
|
||||
for l in levels {
|
||||
let Some(base) = pgrid_level_base(l.pres_mb) else {
|
||||
continue;
|
||||
};
|
||||
out[base] = l.hght_m as f32;
|
||||
out[base + 1] = l.tmpc as f32;
|
||||
out[base + 2] = l.dwpc.map(|v| v as f32).unwrap_or(f32::NAN);
|
||||
}
|
||||
}
|
||||
|
||||
/// Slot of `hght_m_<P>mb` for a pressure level, or `None` if that level
|
||||
/// is not part of the on-disk schema.
|
||||
fn pgrid_level_base(pres_mb: f64) -> Option<usize> {
|
||||
let p = pres_mb.round() as u16;
|
||||
let i = crate::fetcher::GRID_PRESSURE_LEVELS
|
||||
.iter()
|
||||
.position(|&lv| lv == p)?;
|
||||
Some(crate::pgrid::SCALAR_FIELDS.len() + i * 3)
|
||||
}
|
||||
|
||||
/// Resolved slot for each scalar field, computed once.
|
||||
struct PgridScalarIndices {
|
||||
surface_temp_c: usize,
|
||||
surface_dewpoint_c: usize,
|
||||
surface_pressure_mb: usize,
|
||||
hpbl_m: usize,
|
||||
pwat_mm: usize,
|
||||
wind_u: usize,
|
||||
wind_v: usize,
|
||||
cloud_cover_pct: usize,
|
||||
precip_mm: usize,
|
||||
native_min_gradient: usize,
|
||||
best_duct_freq_ghz: usize,
|
||||
max_duct_thickness_m: usize,
|
||||
duct_count: usize,
|
||||
nexrad_max_reflectivity_dbz: usize,
|
||||
commercial_degradation_db: usize,
|
||||
commercial_baseline_dbm: usize,
|
||||
commercial_current_dbm: usize,
|
||||
commercial_n_links: usize,
|
||||
surface_refractivity: usize,
|
||||
min_refractivity_gradient: usize,
|
||||
}
|
||||
|
||||
fn pgrid_scalar_indices() -> &'static PgridScalarIndices {
|
||||
static IDX: std::sync::OnceLock<PgridScalarIndices> = std::sync::OnceLock::new();
|
||||
IDX.get_or_init(|| {
|
||||
let at = |name: &str| {
|
||||
crate::pgrid::SCALAR_FIELDS
|
||||
.iter()
|
||||
.position(|f| *f == name)
|
||||
.unwrap_or_else(|| panic!("pgrid scalar field {name} missing from SCALAR_FIELDS"))
|
||||
};
|
||||
PgridScalarIndices {
|
||||
surface_temp_c: at("surface_temp_c"),
|
||||
surface_dewpoint_c: at("surface_dewpoint_c"),
|
||||
surface_pressure_mb: at("surface_pressure_mb"),
|
||||
hpbl_m: at("hpbl_m"),
|
||||
pwat_mm: at("pwat_mm"),
|
||||
wind_u: at("wind_u"),
|
||||
wind_v: at("wind_v"),
|
||||
cloud_cover_pct: at("cloud_cover_pct"),
|
||||
precip_mm: at("precip_mm"),
|
||||
native_min_gradient: at("native_min_gradient"),
|
||||
best_duct_freq_ghz: at("best_duct_freq_ghz"),
|
||||
max_duct_thickness_m: at("max_duct_thickness_m"),
|
||||
duct_count: at("duct_count"),
|
||||
nexrad_max_reflectivity_dbz: at("nexrad_max_reflectivity_dbz"),
|
||||
commercial_degradation_db: at("commercial_degradation_db"),
|
||||
commercial_baseline_dbm: at("commercial_baseline_dbm"),
|
||||
commercial_current_dbm: at("commercial_current_dbm"),
|
||||
commercial_n_links: at("commercial_n_links"),
|
||||
surface_refractivity: at("surface_refractivity"),
|
||||
min_refractivity_gradient: at("min_refractivity_gradient"),
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::grid::GridSpec;
|
||||
|
||||
fn one_cell_spec() -> GridSpec {
|
||||
GridSpec {
|
||||
lon_start: -100.0,
|
||||
lon_count: 1,
|
||||
lon_step: 0.125,
|
||||
lat_start: 30.0,
|
||||
lat_count: 1,
|
||||
lat_step: 0.125,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn absent_fields_resolve_to_none() {
|
||||
let grid = FieldGrid::new(one_cell_spec());
|
||||
let p = GridPlanes::resolve(&grid);
|
||||
assert!(p.tmp_2m.is_none());
|
||||
assert!(p.apcp.is_none());
|
||||
assert!(p.levels.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn levels_at_skips_levels_missing_temp_or_height() {
|
||||
let mut grid = FieldGrid::new(one_cell_spec());
|
||||
let keys = fetcher::grid_level_keys();
|
||||
// Level 0 complete, level 1 has TMP but no HGT plane at all.
|
||||
grid.push_plane(&keys[0].tmp, vec![295.15]);
|
||||
grid.push_plane(&keys[0].hgt, vec![100.0]);
|
||||
grid.push_plane(&keys[0].dpt, vec![290.15]);
|
||||
grid.push_plane(&keys[1].tmp, vec![293.15]);
|
||||
|
||||
let p = GridPlanes::resolve(&grid);
|
||||
let mut out = Vec::new();
|
||||
p.levels_at(&grid, 0, &mut out);
|
||||
|
||||
assert_eq!(out.len(), 1, "level without HGT must be dropped");
|
||||
assert_eq!(out[0].pres_mb, keys[0].pres_mb);
|
||||
assert!((out[0].tmpc - 22.0).abs() < 1e-3);
|
||||
assert!((out[0].dwpc.unwrap() - 17.0).abs() < 1e-3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn levels_at_tolerates_missing_dewpoint_at_a_cell() {
|
||||
let mut grid = FieldGrid::new(one_cell_spec());
|
||||
let keys = fetcher::grid_level_keys();
|
||||
grid.push_plane(&keys[0].tmp, vec![295.15]);
|
||||
grid.push_plane(&keys[0].hgt, vec![100.0]);
|
||||
// DPT plane exists but this cell has no value.
|
||||
grid.push_plane(&keys[0].dpt, vec![f32::NAN]);
|
||||
|
||||
let p = GridPlanes::resolve(&grid);
|
||||
let mut out = Vec::new();
|
||||
p.levels_at(&grid, 0, &mut out);
|
||||
|
||||
assert_eq!(out.len(), 1);
|
||||
assert_eq!(out[0].dwpc, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn levels_at_clears_the_buffer_between_cells() {
|
||||
let mut grid = FieldGrid::new(one_cell_spec());
|
||||
let keys = fetcher::grid_level_keys();
|
||||
grid.push_plane(&keys[0].tmp, vec![295.15]);
|
||||
grid.push_plane(&keys[0].hgt, vec![100.0]);
|
||||
let p = GridPlanes::resolve(&grid);
|
||||
|
||||
let mut out = vec![Level {
|
||||
pres_mb: 1.0,
|
||||
hght_m: 1.0,
|
||||
tmpc: 1.0,
|
||||
dwpc: None,
|
||||
}];
|
||||
p.levels_at(&grid, 0, &mut out);
|
||||
assert_eq!(out.len(), 1);
|
||||
assert_eq!(out[0].pres_mb, keys[0].pres_mb);
|
||||
}
|
||||
}
|
||||
|
|
@ -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<Utc>) -> 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 `<scores_dir>/profiles/<iso>.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<Utc>,
|
||||
cells: &[CellEntry],
|
||||
) -> Result<PathBuf, WriteError> {
|
||||
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<Item = (&'static str, rmpv::Value)>) -> 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));
|
||||
}
|
||||
}
|
||||
|
|
@ -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<ScorePoint>` 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<Utc>,
|
||||
body: &[u8],
|
||||
spec: grid::GridSpec,
|
||||
) -> Result<Vec<u8>, 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 `<base>/<band_mhz>/<iso>.prop`.
|
||||
/// `hrdps` selects the `.hrdps.prop` sibling path.
|
||||
pub fn write_atomic_dense(
|
||||
base_dir: &Path,
|
||||
band_mhz: u32,
|
||||
valid_time: DateTime<Utc>,
|
||||
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 `<base>/<band_mhz>/<iso>.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]
|
||||
|
|
|
|||
|
|
@ -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<Utc>,
|
||||
cell: &CellValues,
|
||||
levels: &[Level],
|
||||
) -> Option<ScalarRow> {
|
||||
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<Level> {
|
||||
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<f64> {
|
|||
#[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::<str>::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<Utc>) -> Option<ScalarRow> {
|
||||
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);
|
||||
}
|
||||
|
|
|
|||
141
rust/prop_grid_rs/tests/fused_pass_perf.rs
Normal file
141
rust/prop_grid_rs/tests/fused_pass_perf.rs
Normal file
|
|
@ -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<f32> {
|
||||
(0..n)
|
||||
.map(|c| base + ((c % 97) as f32 / 97.0) * span)
|
||||
.collect()
|
||||
};
|
||||
g.push_plane("TMP:2 m above ground", ramp(275.0, 30.0));
|
||||
g.push_plane("DPT:2 m above ground", ramp(270.0, 25.0));
|
||||
g.push_plane("UGRD:10 m above ground", ramp(-8.0, 16.0));
|
||||
g.push_plane("VGRD:10 m above ground", ramp(-8.0, 16.0));
|
||||
g.push_plane("TCDC:entire atmosphere", ramp(0.0, 100.0));
|
||||
g.push_plane("PRES:surface", ramp(97_000.0, 6_000.0));
|
||||
g.push_plane("HPBL:surface", ramp(100.0, 2_000.0));
|
||||
g.push_plane(
|
||||
"PWAT:entire atmosphere (considered as a single layer)",
|
||||
ramp(5.0, 45.0),
|
||||
);
|
||||
g.push_plane("APCP:surface", ramp(0.0, 3.0));
|
||||
|
||||
// 13 pressure levels × 3 variables = the 39 pressure planes.
|
||||
for (i, k) in fetcher::grid_level_keys().iter().enumerate() {
|
||||
let dz = i as f32 * 250.0;
|
||||
g.push_plane(&k.tmp, ramp(275.0 - dz * 0.0065, 20.0));
|
||||
g.push_plane(&k.dpt, ramp(268.0 - dz * 0.0065, 18.0));
|
||||
g.push_plane(&k.hgt, ramp(100.0 + dz, 50.0));
|
||||
}
|
||||
assert_eq!(g.n_planes(), 48);
|
||||
|
||||
let valid_time = Utc.with_ymd_and_hms(2026, 6, 15, 18, 0, 0).unwrap();
|
||||
|
||||
let started = std::time::Instant::now();
|
||||
let out = prop_grid_rs::pipeline::derive_and_score_for_bench(&g, valid_time, true, None);
|
||||
let elapsed = started.elapsed();
|
||||
|
||||
println!(
|
||||
"fused pass: {n} cells × {} planes × {} bands in {:.3} s ({} scored, {} scalars)",
|
||||
g.n_planes(),
|
||||
out.band_bodies.len(),
|
||||
elapsed.as_secs_f64(),
|
||||
out.cells_scored,
|
||||
out.scalar_rows.len(),
|
||||
);
|
||||
|
||||
assert_eq!(out.cells_scored, n as u32);
|
||||
assert_eq!(out.band_bodies.len(), 23);
|
||||
for (_, body) in &out.band_bodies {
|
||||
assert_eq!(body.len(), n);
|
||||
}
|
||||
|
||||
// Sanity: the pgrid body must actually carry the 13-level profile,
|
||||
// otherwise the timing above is measuring an empty structure.
|
||||
let nf = prop_grid_rs::pgrid::n_fields();
|
||||
assert_eq!(out.pgrid_body.len(), n * nf);
|
||||
let names = prop_grid_rs::pgrid::field_names();
|
||||
let t850 = names.iter().position(|f| f == "tmpc_850mb").unwrap();
|
||||
assert!(
|
||||
!out.pgrid_body[t850].is_nan(),
|
||||
"cell 0 should carry an 850 mb temperature"
|
||||
);
|
||||
|
||||
// Now time the on-disk artifacts, which is what remains of the step.
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
||||
let t = std::time::Instant::now();
|
||||
let path =
|
||||
prop_grid_rs::pgrid::write_atomic(dir.path(), valid_time, &spec, &out.pgrid_body, false)
|
||||
.unwrap();
|
||||
let profile_write = t.elapsed();
|
||||
let profile_bytes = std::fs::metadata(&path).unwrap().len();
|
||||
|
||||
let t = std::time::Instant::now();
|
||||
prop_grid_rs::weather_scalar_file::write_atomic(dir.path(), valid_time, &out.scalar_rows)
|
||||
.unwrap();
|
||||
let scalar_write = t.elapsed();
|
||||
|
||||
let t = std::time::Instant::now();
|
||||
for (band_mhz, body) in &out.band_bodies {
|
||||
prop_grid_rs::scores_file::write_atomic_dense(
|
||||
dir.path(),
|
||||
*band_mhz,
|
||||
valid_time,
|
||||
body,
|
||||
spec,
|
||||
false,
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
let scores_write = t.elapsed();
|
||||
|
||||
println!(
|
||||
"pgrid write : {:.3} s ({:.1} MB on disk)",
|
||||
profile_write.as_secs_f64(),
|
||||
profile_bytes as f64 / 1e6
|
||||
);
|
||||
|
||||
// Single-cell read: what /map point-detail and Skew-T actually do.
|
||||
// The old `.mp.gz` path gunzipped and unpacked the whole file for this.
|
||||
let raw = std::fs::read(&path).unwrap();
|
||||
let header = prop_grid_rs::pgrid::decode_header(&raw).expect("header");
|
||||
let t = std::time::Instant::now();
|
||||
let mut sink = 0.0f32;
|
||||
for cell in (0..n).step_by(1_000) {
|
||||
let rec = prop_grid_rs::pgrid::read_cell(&raw, &header, cell).unwrap();
|
||||
sink += rec[0];
|
||||
}
|
||||
let reads = (n / 1_000) as f64;
|
||||
println!(
|
||||
"pgrid read_cell : {:.1} µs/cell over {reads} reads (checksum {sink:.1})",
|
||||
t.elapsed().as_secs_f64() * 1e6 / reads
|
||||
);
|
||||
println!("scalar_file write: {:.3} s", scalar_write.as_secs_f64());
|
||||
println!("23 score files : {:.3} s", scores_write.as_secs_f64());
|
||||
}
|
||||
323
test/microwaveprop/propagation/pgrid_test.exs
Normal file
323
test/microwaveprop/propagation/pgrid_test.exs
Normal file
|
|
@ -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 -> <<v::float-little-32>>
|
||||
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
|
||||
Loading…
Add table
Reference in a new issue