diff --git a/lib/microwaveprop/propagation/grid.ex b/lib/microwaveprop/propagation/grid.ex index 45539075..e6c7adb1 100644 --- a/lib/microwaveprop/propagation/grid.ex +++ b/lib/microwaveprop/propagation/grid.ex @@ -21,6 +21,21 @@ defmodule Microwaveprop.Propagation.Grid do @doc "Returns the CONUS bounding box as a map." def bounds, do: %{lat_min: @lat_min, lat_max: @lat_max, lon_min: @lon_min, lon_max: @lon_max} + @doc "Returns the grid specification for wgrib2 -lola extraction." + def wgrib2_grid_spec do + lon_count = round((@lon_max - @lon_min) / @step) + 1 + lat_count = round((@lat_max - @lat_min) / @step) + 1 + + %{ + lon_start: @lon_min, + lon_count: lon_count, + lon_step: @step, + lat_start: @lat_min, + lat_count: lat_count, + lat_step: @step + } + end + defp float_range(start, stop, step) do count = round((stop - start) / step) + 1 Enum.map(0..(count - 1), fn i -> start + i * step end) diff --git a/lib/microwaveprop/weather/grib2/extractor.ex b/lib/microwaveprop/weather/grib2/extractor.ex index 6cdefe86..b0a9b3e8 100644 --- a/lib/microwaveprop/weather/grib2/extractor.ex +++ b/lib/microwaveprop/weather/grib2/extractor.ex @@ -36,29 +36,46 @@ defmodule Microwaveprop.Weather.Grib2.Extractor do def extract_grid(binary, points) do messages = split_messages(binary) - result = - Enum.reduce(messages, {:ok, init_grid(points)}, fn msg, {:ok, acc} -> - case extract_single_grid(msg, points) do - {:ok, point_values} -> - merged = - Enum.reduce(point_values, acc, fn {point, key, value}, grid -> - Map.update!(grid, point, &Map.put(&1, key, value)) - end) + # Pre-compute grid indices once from the first message's grid params. + # All messages in the same GRIB2 file share the same grid. + grid_indices = + case messages do + [first_msg | _] -> + case Section.parse_message(first_msg) do + {:ok, %{grid_params: grid}} -> precompute_indices(grid, points) + _ -> [] + end - {:ok, merged} + [] -> + [] + end - {:error, _reason} -> - # Skip messages with missing/corrupt data sections - {:ok, acc} - end - end) + if grid_indices == [] do + {:ok, %{}} + else + result = + Enum.reduce(messages, {:ok, init_grid(points)}, fn msg, {:ok, acc} -> + case extract_grid_with_indices(msg, grid_indices) do + {:ok, key, values} -> + merged = + Enum.reduce(values, acc, fn {point, value}, grid -> + Map.update!(grid, point, &Map.put(&1, key, value)) + end) - case result do - {:ok, grid} -> - {:ok, Map.reject(grid, fn {_point, values} -> values == %{} end)} + {:ok, merged} - error -> - error + {:error, _reason} -> + {:ok, acc} + end + end) + + case result do + {:ok, grid} -> + {:ok, Map.reject(grid, fn {_point, values} -> values == %{} end)} + + error -> + error + end end end @@ -100,36 +117,56 @@ defmodule Microwaveprop.Weather.Grib2.Extractor do Map.new(points, fn point -> {point, %{}} end) end - defp extract_single_grid(msg, points) do + # Pre-compute Lambert Conformal grid indices for all points (done once, reused per message) + defp precompute_indices(grid, points) do + Enum.flat_map(points, fn {lat, lon} = point -> + case LambertConformal.to_grid_index(grid, lat, lon) do + {:ok, {i, j}} -> + [{point, linear_index({i, j}, grid.nx, grid.scan_mode)}] + + {:error, :outside_grid} -> + [] + end + end) + end + + # Extract values for all pre-computed indices from a single GRIB message + defp extract_grid_with_indices(msg, grid_indices) do with {:ok, parsed} <- Section.parse_message(msg) do - %{grid_params: grid, product: prod, packing_params: packing, data: data} = parsed + %{product: prod, packing_params: packing, data: data} = parsed key = "#{prod.var}:#{prod.level}" - results = - Enum.reduce_while(points, {:ok, []}, fn {lat, lon} = point, {:ok, acc} -> - case LambertConformal.to_grid_index(grid, lat, lon) do - {:ok, {i, j}} -> - index = linear_index({i, j}, grid.nx, grid.scan_mode) + # For complex packing, decode_all is called once; for simple, each index is independent + unique_indices = grid_indices |> Enum.map(&elem(&1, 1)) |> Enum.uniq() - case unpack_value(packing, data, index) do - {:ok, value} -> {:cont, {:ok, [{point, key, value} | acc]}} - {:error, reason} -> {:halt, {:error, reason}} + case batch_unpack(packing, data, unique_indices) do + {:ok, index_values} -> + values = + Enum.flat_map(grid_indices, fn {point, index} -> + case Map.get(index_values, index) do + nil -> [] + value -> [{point, value}] end + end) - {:error, :outside_grid} -> - {:cont, {:ok, acc}} - end - end) + {:ok, key, values} - case results do - {:ok, point_values} -> {:ok, point_values} - {:error, reason} -> {:error, reason} + {:error, reason} -> + {:error, reason} end end rescue e -> {:error, "GRIB2 grid extraction failed: #{inspect(e)}"} end + defp batch_unpack(%{template: 3} = params, data, indices) do + ComplexPacking.extract_values(params, data, indices) + end + + defp batch_unpack(params, data, indices) do + SimplePacking.extract_values(params, data, indices) + end + defp extract_single(msg, lat, lon) do with {:ok, parsed} <- Section.parse_message(msg), %{grid_params: grid, product: prod, packing_params: packing, data: data} = parsed, diff --git a/lib/microwaveprop/weather/grib2/wgrib2.ex b/lib/microwaveprop/weather/grib2/wgrib2.ex new file mode 100644 index 00000000..67e43b0d --- /dev/null +++ b/lib/microwaveprop/weather/grib2/wgrib2.ex @@ -0,0 +1,166 @@ +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 + + @wgrib2_path System.find_executable("wgrib2") + @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 "Check if wgrib2 is available on the system." + def available? do + @wgrib2_path != nil + 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}] + + _ -> + [] + 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 + + result = + messages + |> Enum.with_index() + |> Enum.reduce(%{}, fn {msg, msg_idx}, acc -> + offset = msg_idx * bytes_per_message + key = "#{msg.var}:#{msg.level}" + + if offset + bytes_per_message <= byte_size(bin_data) do + chunk = binary_part(bin_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 diff --git a/lib/microwaveprop/weather/hrrr_client.ex b/lib/microwaveprop/weather/hrrr_client.ex index 2a0a74be..12e3d8df 100644 --- a/lib/microwaveprop/weather/hrrr_client.ex +++ b/lib/microwaveprop/weather/hrrr_client.ex @@ -1,7 +1,9 @@ defmodule Microwaveprop.Weather.HrrrClient do @moduledoc false + alias Microwaveprop.Propagation.Grid alias Microwaveprop.Weather.Grib2.Extractor + alias Microwaveprop.Weather.Grib2.Wgrib2 require Logger @@ -210,7 +212,7 @@ defmodule Microwaveprop.Weather.HrrrClient do end end - defp fetch_product_grid(date, hour, product, points) do + defp fetch_product_grid(date, hour, product, _points) do url = hrrr_url(date, hour, product) idx_url = url <> ".idx" @@ -228,10 +230,28 @@ defmodule Microwaveprop.Weather.HrrrClient do _ = Logger.info("HRRR grid downloading #{length(ranges)} GRIB ranges for #{product}"), {:ok, grib_binary} <- download_grib_ranges(url, ranges) do Logger.info("HRRR grid #{product} downloaded #{byte_size(grib_binary)} bytes, extracting") + extract_grid(grib_binary, wanted, product) + end + end + + defp extract_grid(grib_binary, wanted, product) do + if Wgrib2.available?() do + match_pattern = wgrib2_match_pattern(wanted) + grid_spec = Grid.wgrib2_grid_spec() + Logger.info("HRRR grid using wgrib2 for #{product} extraction") + Wgrib2.extract_grid(grib_binary, match_pattern, grid_spec) + else + Logger.info("HRRR grid using Elixir decoder for #{product} extraction (slow)") + points = Grid.conus_points() Extractor.extract_grid(grib_binary, points) end end + defp wgrib2_match_pattern(wanted) do + vars = wanted |> Enum.map(& &1.var) |> Enum.uniq() |> Enum.join("|") + ":(#{vars}):" + end + defp maybe_fetch_pressure_grid(date, hour, points, opts) do if Keyword.get(opts, :pressure, true) do fetch_product_grid(date, hour, :pressure, points) @@ -262,18 +282,54 @@ defmodule Microwaveprop.Weather.HrrrClient do defp download_grib_ranges(_url, []), do: {:ok, <<>>} defp download_grib_ranges(url, ranges) do - Enum.reduce_while(ranges, {:ok, <<>>}, fn {start, stop}, {:ok, acc} -> - case Req.get(url, [{:headers, [{"Range", "bytes=#{start}-#{stop}"}]} | req_options()]) do - {:ok, %{status: 206, body: body}} -> - {:cont, {:ok, acc <> body}} + # Merge adjacent/overlapping ranges to reduce HTTP requests + merged = merge_ranges(ranges) - {:ok, %{status: status}} -> - {:halt, {:error, "HRRR grib HTTP #{status}"}} + # Download ranges in parallel + results = + merged + |> Task.async_stream( + fn {start, stop} -> + case Req.get(url, [{:headers, [{"Range", "bytes=#{start}-#{stop}"}]} | req_options()]) do + {:ok, %{status: 206, body: body}} -> {:ok, start, body} + {:ok, %{status: status}} -> {:error, "HRRR grib HTTP #{status}"} + {:error, reason} -> {:error, reason} + end + end, + max_concurrency: 8, + timeout: 60_000 + ) + |> Enum.map(fn + {:ok, result} -> result + {:exit, reason} -> {:error, reason} + end) - {:error, reason} -> - {:halt, {:error, reason}} + # Sort by offset and concatenate + case Enum.find(results, &match?({:error, _}, &1)) do + {:error, reason} -> + {:error, reason} + + nil -> + binary = + results + |> Enum.sort_by(fn {:ok, offset, _} -> offset end) + |> Enum.map(fn {:ok, _, body} -> body end) + |> IO.iodata_to_binary() + + {:ok, binary} + end + end + + defp merge_ranges(ranges) do + ranges + |> Enum.sort_by(&elem(&1, 0)) + |> Enum.reduce([], fn {s, e}, acc -> + case acc do + [{ps, pe} | rest] when s <= pe + 1 -> [{ps, max(pe, e)} | rest] + _ -> [{s, e} | acc] end end) + |> Enum.reverse() end defp req_options do