Graham McIntire
f2efdd4ece
Instead of holding ~530MB GRIB binary in BEAM memory, download ranges directly to a temp file and run wgrib2 on it. Peak memory drops from ~530MB to just HTTP chunk buffers.
360 lines
11 KiB
Elixir
360 lines
11 KiB
Elixir
defmodule Microwaveprop.Weather.Grib2.Wgrib2 do
|
|
@moduledoc """
|
|
Fast GRIB2 grid extraction using the wgrib2 binary.
|
|
|
|
Uses wgrib2's `-lola` option to interpolate HRRR Lambert Conformal data
|
|
onto a regular lat-lon grid, outputting IEEE 754 binary floats.
|
|
Falls back to the pure-Elixir decoder if wgrib2 is not available.
|
|
"""
|
|
|
|
require Logger
|
|
|
|
@undefined_value 9.999e20
|
|
|
|
@doc """
|
|
Extract values from a GRIB2 binary at a regular lat-lon grid.
|
|
|
|
Takes the raw GRIB2 binary data, a regex pattern to match desired messages
|
|
(e.g. ":(TMP|DPT|PRES):"), and grid specification.
|
|
|
|
Returns `{:ok, %{{lat, lon} => %{"VAR:LEVEL" => float}}}` or `{:error, reason}`.
|
|
"""
|
|
def extract_grid(grib_binary, match_pattern, grid_spec) do
|
|
if available?() do
|
|
extract_with_wgrib2(grib_binary, match_pattern, grid_spec)
|
|
else
|
|
{:error, :wgrib2_not_available}
|
|
end
|
|
end
|
|
|
|
@doc """
|
|
Like `extract_grid/3`, but takes a path to a GRIB2 file on disk instead
|
|
of a binary. Avoids loading the entire file into memory — useful for
|
|
large native-level HRRR files (~530 MB).
|
|
"""
|
|
def extract_grid_from_file(grib_path, match_pattern, grid_spec) do
|
|
if available?() do
|
|
extract_file_with_wgrib2(grib_path, match_pattern, grid_spec)
|
|
else
|
|
{:error, :wgrib2_not_available}
|
|
end
|
|
end
|
|
|
|
@doc "Check if wgrib2 is available on the system."
|
|
def available?, do: wgrib2_path() != nil
|
|
|
|
defp wgrib2_path, do: System.find_executable("wgrib2")
|
|
|
|
defp extract_file_with_wgrib2(grib_path, match_pattern, grid_spec) do
|
|
%{
|
|
lon_start: lon_start,
|
|
lon_count: lon_count,
|
|
lon_step: lon_step,
|
|
lat_start: lat_start,
|
|
lat_count: lat_count,
|
|
lat_step: lat_step
|
|
} = grid_spec
|
|
|
|
wgrib2_lon_start = normalize_lon(lon_start)
|
|
lon_spec = "#{wgrib2_lon_start}:#{lon_count}:#{lon_step}"
|
|
lat_spec = "#{lat_start}:#{lat_count}:#{lat_step}"
|
|
|
|
tmp_bin = grib_path <> ".lola.bin"
|
|
|
|
try do
|
|
args = [grib_path, "-match", match_pattern, "-lola", lon_spec, lat_spec, tmp_bin, "bin"]
|
|
|
|
case System.cmd(wgrib2_path(), args, stderr_to_stdout: true) do
|
|
{output, 0} ->
|
|
messages = parse_wgrib2_inventory(output)
|
|
|
|
case File.read(tmp_bin) do
|
|
{:ok, bin_data} -> parse_lola_binary(bin_data, messages, grid_spec)
|
|
{:error, :enoent} -> {:ok, %{}}
|
|
{:error, reason} -> {:error, "Failed to read wgrib2 output: #{inspect(reason)}"}
|
|
end
|
|
|
|
{output, exit_code} ->
|
|
{:error, "wgrib2 failed (exit #{exit_code}): #{String.slice(output, 0, 200)}"}
|
|
end
|
|
after
|
|
File.rm(tmp_bin)
|
|
end
|
|
end
|
|
|
|
defp extract_with_wgrib2(grib_binary, match_pattern, grid_spec) do
|
|
%{
|
|
lon_start: lon_start,
|
|
lon_count: lon_count,
|
|
lon_step: lon_step,
|
|
lat_start: lat_start,
|
|
lat_count: lat_count,
|
|
lat_step: lat_step
|
|
} = grid_spec
|
|
|
|
# wgrib2 uses 0-360 longitude convention
|
|
wgrib2_lon_start = normalize_lon(lon_start)
|
|
|
|
lon_spec = "#{wgrib2_lon_start}:#{lon_count}:#{lon_step}"
|
|
lat_spec = "#{lat_start}:#{lat_count}:#{lat_step}"
|
|
|
|
# Write GRIB to temp file
|
|
tmp_grib = Path.join(System.tmp_dir!(), "hrrr_#{System.unique_integer([:positive])}.grib2")
|
|
tmp_bin = tmp_grib <> ".lola.bin"
|
|
|
|
try do
|
|
File.write!(tmp_grib, grib_binary)
|
|
|
|
# Run wgrib2: match desired messages, extract to regular lat-lon grid as binary
|
|
args = [
|
|
tmp_grib,
|
|
"-match",
|
|
match_pattern,
|
|
"-lola",
|
|
lon_spec,
|
|
lat_spec,
|
|
tmp_bin,
|
|
"bin"
|
|
]
|
|
|
|
case System.cmd(wgrib2_path(), args, stderr_to_stdout: true) do
|
|
{output, 0} ->
|
|
# Parse message inventory from stdout to know which vars were extracted
|
|
messages = parse_wgrib2_inventory(output)
|
|
|
|
case File.read(tmp_bin) do
|
|
{:ok, bin_data} ->
|
|
parse_lola_binary(bin_data, messages, grid_spec)
|
|
|
|
{:error, :enoent} ->
|
|
# No output file means no matching messages
|
|
{:ok, %{}}
|
|
|
|
{:error, reason} ->
|
|
{:error, "Failed to read wgrib2 output: #{inspect(reason)}"}
|
|
end
|
|
|
|
{output, exit_code} ->
|
|
{:error, "wgrib2 failed (exit #{exit_code}): #{String.slice(output, 0, 200)}"}
|
|
end
|
|
after
|
|
File.rm(tmp_grib)
|
|
File.rm(tmp_bin)
|
|
end
|
|
end
|
|
|
|
defp parse_wgrib2_inventory(output) do
|
|
output
|
|
|> String.split("\n")
|
|
|> Enum.flat_map(fn line ->
|
|
case String.split(line, ":", parts: 8) do
|
|
[_n, _offset, date, var, level | _] ->
|
|
[%{var: var, level: level, datetime: parse_wgrib2_date(date)}]
|
|
|
|
_ ->
|
|
[]
|
|
end
|
|
end)
|
|
end
|
|
|
|
# wgrib2 inventory date looks like "d=20250101000000" or "d=2025010100".
|
|
# Returns a `DateTime` in UTC, or `nil` if it can't be parsed.
|
|
defp parse_wgrib2_date("d=" <> digits), do: parse_wgrib2_date_digits(digits)
|
|
defp parse_wgrib2_date(_), do: nil
|
|
|
|
defp parse_wgrib2_date_digits(<<y::binary-4, m::binary-2, d::binary-2, h::binary-2, rest::binary>>) do
|
|
{mi, s} =
|
|
case rest do
|
|
<<mm::binary-2, ss::binary-2>> -> {mm, ss}
|
|
<<mm::binary-2>> -> {mm, "00"}
|
|
_ -> {"00", "00"}
|
|
end
|
|
|
|
with {year, ""} <- Integer.parse(y),
|
|
{month, ""} <- Integer.parse(m),
|
|
{day, ""} <- Integer.parse(d),
|
|
{hour, ""} <- Integer.parse(h),
|
|
{minute, ""} <- Integer.parse(mi),
|
|
{second, ""} <- Integer.parse(s),
|
|
{:ok, naive} <- NaiveDateTime.new(year, month, day, hour, minute, second),
|
|
{:ok, dt} <- DateTime.from_naive(naive, "Etc/UTC") do
|
|
DateTime.truncate(dt, :second)
|
|
else
|
|
_ -> nil
|
|
end
|
|
end
|
|
|
|
defp parse_wgrib2_date_digits(_), do: nil
|
|
|
|
@doc """
|
|
Like `extract_grid/3`, but returns a flat list of per-message results so
|
|
the time dimension is preserved. Used when a single GRIB2 blob carries
|
|
many timesteps (e.g. a month of ERA5 data fetched in one CDS request).
|
|
|
|
Each entry is `%{datetime: DateTime.t() | nil, var: String.t(),
|
|
level: String.t(), values: %{{lat, lon} => float}}`.
|
|
"""
|
|
def extract_grid_messages(grib_binary, match_pattern, grid_spec) do
|
|
if available?() do
|
|
extract_messages_with_wgrib2(grib_binary, match_pattern, grid_spec)
|
|
else
|
|
{:error, :wgrib2_not_available}
|
|
end
|
|
end
|
|
|
|
defp extract_messages_with_wgrib2(grib_binary, match_pattern, grid_spec) do
|
|
%{
|
|
lon_start: lon_start,
|
|
lon_count: lon_count,
|
|
lon_step: lon_step,
|
|
lat_start: lat_start,
|
|
lat_count: lat_count,
|
|
lat_step: lat_step
|
|
} = grid_spec
|
|
|
|
wgrib2_lon_start = normalize_lon(lon_start)
|
|
lon_spec = "#{wgrib2_lon_start}:#{lon_count}:#{lon_step}"
|
|
lat_spec = "#{lat_start}:#{lat_count}:#{lat_step}"
|
|
|
|
tmp_grib = Path.join(System.tmp_dir!(), "era5_#{System.unique_integer([:positive])}.grib2")
|
|
tmp_bin = tmp_grib <> ".lola.bin"
|
|
|
|
try do
|
|
File.write!(tmp_grib, grib_binary)
|
|
|
|
args = [tmp_grib, "-match", match_pattern, "-lola", lon_spec, lat_spec, tmp_bin, "bin"]
|
|
|
|
case System.cmd(wgrib2_path(), args, stderr_to_stdout: true) do
|
|
{output, 0} ->
|
|
messages = parse_wgrib2_inventory(output)
|
|
|
|
case File.read(tmp_bin) do
|
|
{:ok, bin_data} ->
|
|
{:ok, build_messages_per_message(bin_data, messages, grid_spec)}
|
|
|
|
{:error, :enoent} ->
|
|
{:ok, []}
|
|
|
|
{:error, reason} ->
|
|
{:error, "Failed to read wgrib2 output: #{inspect(reason)}"}
|
|
end
|
|
|
|
{output, exit_code} ->
|
|
{:error, "wgrib2 failed (exit #{exit_code}): #{String.slice(output, 0, 200)}"}
|
|
end
|
|
after
|
|
File.rm(tmp_grib)
|
|
File.rm(tmp_bin)
|
|
end
|
|
end
|
|
|
|
defp build_messages_per_message(bin_data, messages, grid_spec) do
|
|
%{lon_count: nx, lat_count: ny} = grid_spec
|
|
points_per_message = nx * ny
|
|
bytes_per_message = points_per_message * 4
|
|
|
|
messages
|
|
|> Enum.with_index()
|
|
|> Enum.flat_map(fn {msg, msg_idx} ->
|
|
offset = msg_idx * bytes_per_message
|
|
|
|
if offset + bytes_per_message <= byte_size(bin_data) do
|
|
chunk = binary_part(bin_data, offset, bytes_per_message)
|
|
values = extract_message_values(chunk, grid_spec)
|
|
|
|
[Map.put(msg, :values, values)]
|
|
else
|
|
[]
|
|
end
|
|
end)
|
|
end
|
|
|
|
defp extract_message_values(chunk, %{
|
|
lon_count: nx,
|
|
lat_count: ny,
|
|
lon_start: lon_start,
|
|
lon_step: lon_step,
|
|
lat_start: lat_start,
|
|
lat_step: lat_step
|
|
}) do
|
|
Enum.reduce(0..(ny - 1), %{}, fn j, outer ->
|
|
lat = Float.round(lat_start + j * lat_step, 3)
|
|
|
|
Enum.reduce(0..(nx - 1), outer, fn i, inner ->
|
|
offset = (j * nx + i) * 4
|
|
<<_::binary-size(offset), value::float-little-32, _::binary>> = chunk
|
|
|
|
if value > @undefined_value / 2 do
|
|
inner
|
|
else
|
|
lon = Float.round(denormalize_lon(lon_start + i * lon_step), 3)
|
|
Map.put(inner, {lat, lon}, value)
|
|
end
|
|
end)
|
|
end)
|
|
end
|
|
|
|
defp parse_lola_binary(bin_data, messages, grid_spec) do
|
|
%{lon_count: nx, lat_count: ny, lon_start: lon_start, lon_step: lon_step, lat_start: lat_start, lat_step: lat_step} =
|
|
grid_spec
|
|
|
|
points_per_message = nx * ny
|
|
bytes_per_message = points_per_message * 4
|
|
|
|
# wgrib2 -lola ... bin writes Fortran unformatted binary: each
|
|
# message is preceded by a 4-byte little-endian record length
|
|
# and followed by a duplicate 4-byte record length. So the
|
|
# stride per message is 4 + data + 4 = data + 8.
|
|
record_overhead = 8
|
|
stride = bytes_per_message + record_overhead
|
|
|
|
result =
|
|
messages
|
|
|> Enum.with_index()
|
|
|> Enum.reduce(%{}, fn {msg, msg_idx}, acc ->
|
|
# Skip the 4-byte header to reach the data
|
|
data_offset = msg_idx * stride + 4
|
|
key = "#{msg.var}:#{msg.level}"
|
|
|
|
if data_offset + bytes_per_message <= byte_size(bin_data) do
|
|
chunk = binary_part(bin_data, data_offset, bytes_per_message)
|
|
merge_message_values(acc, key, chunk, nx, ny, lon_start, lon_step, lat_start, lat_step)
|
|
else
|
|
acc
|
|
end
|
|
end)
|
|
|
|
{:ok, result}
|
|
end
|
|
|
|
defp merge_message_values(acc, key, chunk, nx, ny, lon_start, lon_step, lat_start, lat_step) do
|
|
# Binary is row-major: lat varies slowest, lon varies fastest
|
|
# Each value is a 32-bit IEEE 754 little-endian float
|
|
Enum.reduce(0..(ny - 1), acc, fn j, acc_outer ->
|
|
lat = Float.round(lat_start + j * lat_step, 3)
|
|
|
|
Enum.reduce(0..(nx - 1), acc_outer, fn i, acc_inner ->
|
|
offset = (j * nx + i) * 4
|
|
<<_::binary-size(offset), value::float-little-32, _::binary>> = chunk
|
|
|
|
if value > @undefined_value / 2 do
|
|
# Skip undefined values (ocean/outside-domain points)
|
|
acc_inner
|
|
else
|
|
lon = Float.round(denormalize_lon(lon_start + i * lon_step), 3)
|
|
point = {lat, lon}
|
|
existing = Map.get(acc_inner, point, %{})
|
|
Map.put(acc_inner, point, Map.put(existing, key, value))
|
|
end
|
|
end)
|
|
end)
|
|
end
|
|
|
|
# Convert -125.0 to 235.0 for wgrib2
|
|
defp normalize_lon(lon) when lon < 0, do: lon + 360.0
|
|
defp normalize_lon(lon), do: lon
|
|
|
|
# Convert back from 0-360 to -180..180
|
|
defp denormalize_lon(lon) when lon > 180.0, do: lon - 360.0
|
|
defp denormalize_lon(lon), do: lon
|
|
end
|