From 9741941f5f75d6d9fd561b6e5f03aea9a2ae1d3f Mon Sep 17 00:00:00 2001 From: Robert Jackson Date: Mon, 27 Apr 2026 15:54:47 -0500 Subject: [PATCH] ENH: Add IDW initialization from IEM surface observations Adds make_initialization_from_iem_obs to pydda.initialization. The function performs horizontal inverse-distance-weighting of Iowa Environmental Mesonet surface station winds and optionally blends them with a vertical sounding profile as the background. Includes seven unit tests covering shape, dims, sounding-only, station-only, anomaly blending, midpoint averaging, and power parameter behavior. Co-Authored-By: Claude Sonnet 4.6 --- pydda/initialization/__init__.py | 2 + pydda/initialization/wind_fields.py | 157 ++++++++++++++++++++++++++++ pydda/tests/test_initialization.py | 119 +++++++++++++++++++++ 3 files changed, 278 insertions(+) diff --git a/pydda/initialization/__init__.py b/pydda/initialization/__init__.py index 647cfec6..16e35a5d 100644 --- a/pydda/initialization/__init__.py +++ b/pydda/initialization/__init__.py @@ -22,6 +22,7 @@ make_wind_field_from_profile make_background_from_wrf make_initialization_from_era5 + make_initialization_from_iem_obs """ @@ -29,3 +30,4 @@ from .wind_fields import make_wind_field_from_profile from .wind_fields import make_background_from_wrf from .wind_fields import make_initialization_from_era5 +from .wind_fields import make_initialization_from_iem_obs diff --git a/pydda/initialization/wind_fields.py b/pydda/initialization/wind_fields.py index 6790f103..faaaa397 100644 --- a/pydda/initialization/wind_fields.py +++ b/pydda/initialization/wind_fields.py @@ -504,6 +504,163 @@ def make_background_from_wrf(Grid, file_path, wrf_time, radar_loc, vel_field=Non return Grid +def make_initialization_from_iem_obs( + Grid, station_obs, profile=None, power=2, vel_field=None +): + """ + Create an initial 3D wind field by inverse distance weighting (IDW) + interpolation of Iowa Environmental Mesonet surface observations. + + Surface station winds are used to compute horizontal anomalies relative to + an optional sounding background. The anomalies are spread horizontally via + IDW and added uniformly at every vertical level. When no sounding is + supplied the background is zero (i.e. the IDW result stands alone). + + Parameters + ---------- + Grid : xarray.Dataset + The PyDDA analysis grid. Must contain ``x``, ``y``, ``z``, + ``point_x``, ``point_y``, and ``point_z`` coordinate variables. + station_obs : list of dict + Surface observations, typically the output of + :func:`pydda.constraints.get_iem_obs`. Each dict must contain the + keys ``x``, ``y``, ``z``, ``u``, and ``v`` in the Grid's Cartesian + coordinate system (metres). + profile : pyart.core.HorizontalWindProfile or None + Optional vertical sounding profile used as the background wind field. + When provided the IDW step corrects only the horizontal anomalies + (departure from the sounding) and the sounding provides the vertical + structure. When ``None`` the background is zero and the surface IDW + values are copied unchanged to every model level. + power : float + IDW distance-weighting exponent. Higher values give more weight to + the nearest station. Typical values are 1 or 2 (default ``2``). + vel_field : str or None + Name of the radar velocity field in *Grid* used to determine the + output grid shape. ``None`` will auto-detect via + ``Grid.attrs["first_grid_name"]``. + + Returns + ------- + Grid : xarray.Dataset + The input Grid with ``u``, ``v``, and ``w`` fields added or replaced. + ``w`` is set to zero everywhere because surface stations do not + observe vertical motion. + + Notes + ----- + The horizontal IDW anomaly is applied identically at every height level. + This is equivalent to assuming that horizontal wind anomalies measured at + the surface are representative of the full column (a reasonable first + guess for a well-mixed boundary layer). If a sounding *profile* is + given, that profile provides the height-varying mean, while the stations + supply the mesoscale horizontal structure. + + The function raises ``ValueError`` if both *station_obs* is empty and + *profile* is ``None``, since there is no information to build a field. + + Examples + -------- + >>> import pydda + >>> Grid = pydda.io.read_grid(grid_file) + >>> stn_obs = pydda.constraints.get_iem_obs(Grid) + >>> Grid = pydda.initialization.make_initialization_from_iem_obs( + ... Grid, stn_obs, profile=my_sounding + ... ) + """ + if len(station_obs) == 0 and profile is None: + raise ValueError( + "station_obs is empty and no sounding profile was provided. " + "At least one source of wind information is required." + ) + + if vel_field is None: + vel_field = Grid.attrs["first_grid_name"] + + grid_x = Grid["x"].values # (nx,) + grid_y = Grid["y"].values # (ny,) + grid_z = Grid["z"].values # (nz,) + nz, ny, nx = len(grid_z), len(grid_y), len(grid_x) + + # --- vertical background from sounding or zeros --- + if profile is not None: + u_interp_snd = interp1d( + profile.height, profile.u_wind, bounds_error=False, fill_value="extrapolate" + ) + v_interp_snd = interp1d( + profile.height, profile.v_wind, bounds_error=False, fill_value="extrapolate" + ) + u_back = u_interp_snd(grid_z) # (nz,) + v_back = v_interp_snd(grid_z) # (nz,) + else: + u_back = np.zeros(nz) + v_back = np.zeros(nz) + + # Broadcast background to full 3D shape + u_3d = np.broadcast_to(u_back[:, np.newaxis, np.newaxis], (nz, ny, nx)).copy() + v_3d = np.broadcast_to(v_back[:, np.newaxis, np.newaxis], (nz, ny, nx)).copy() + w_3d = np.zeros((nz, ny, nx)) + + if len(station_obs) == 0: + # Background only; w stays zero + pass + else: + station_x = np.array([s["x"] for s in station_obs], dtype=float) + station_y = np.array([s["y"] for s in station_obs], dtype=float) + station_z = np.array([s["z"] for s in station_obs], dtype=float) + station_u = np.array([s["u"] for s in station_obs], dtype=float) + station_v = np.array([s["v"] for s in station_obs], dtype=float) + + # Anomaly of each station relative to sounding background at its height + u_back_stn = np.interp(station_z, grid_z, u_back) # (N,) + v_back_stn = np.interp(station_z, grid_z, v_back) # (N,) + u_anom = station_u - u_back_stn # (N,) + v_anom = station_v - v_back_stn # (N,) + + # Horizontal distances: (N, ny, nx) + grid_xx, grid_yy = np.meshgrid(grid_x, grid_y) # (ny, nx) each + dx = grid_xx[np.newaxis, :, :] - station_x[:, np.newaxis, np.newaxis] + dy = grid_yy[np.newaxis, :, :] - station_y[:, np.newaxis, np.newaxis] + d_h = np.sqrt(dx**2 + dy**2) + + # Avoid division by zero when a grid point coincides with a station + d_h = np.maximum(d_h, 1.0) + + weights = 1.0 / d_h**power # (N, ny, nx) + weights_norm = weights / weights.sum(axis=0, keepdims=True) + + # IDW-interpolated anomaly on the 2-D horizontal plane + u_anom_2d = (weights_norm * u_anom[:, np.newaxis, np.newaxis]).sum(axis=0) + v_anom_2d = (weights_norm * v_anom[:, np.newaxis, np.newaxis]).sum(axis=0) + + # Add horizontal anomaly uniformly to every vertical level + u_3d += u_anom_2d[np.newaxis, :, :] + v_3d += v_anom_2d[np.newaxis, :, :] + + u_attrs = { + "standard_name": "u_wind", + "long_name": "meridional component of wind velocity", + } + v_attrs = { + "standard_name": "v_wind", + "long_name": "zonal component of wind velocity", + } + w_attrs = { + "standard_name": "w_wind", + "long_name": "vertical component of wind velocity", + } + Grid["u"] = xr.DataArray( + np.expand_dims(u_3d, 0), dims=("time", "z", "y", "x"), attrs=u_attrs + ) + Grid["v"] = xr.DataArray( + np.expand_dims(v_3d, 0), dims=("time", "z", "y", "x"), attrs=v_attrs + ) + Grid["w"] = xr.DataArray( + np.expand_dims(w_3d, 0), dims=("time", "z", "y", "x"), attrs=w_attrs + ) + return Grid + + def make_intialization_from_hrrr(Grid, file_path, method="linear"): """ This function will read an HRRR GRIB2 file and return initial guess diff --git a/pydda/tests/test_initialization.py b/pydda/tests/test_initialization.py index 6baf317d..574eddc4 100644 --- a/pydda/tests/test_initialization.py +++ b/pydda/tests/test_initialization.py @@ -9,6 +9,7 @@ import pydda import pyart import numpy as np +import pytest from datetime import datetime @@ -55,6 +56,124 @@ def test_make_wind_field_from_profile(): assert np.all(Grid["w"].values == 0.0) +def _make_small_grid(): + """5x5x5 Cartesian grid for fast IDW tests.""" + Grid = pyart.testing.make_empty_grid( + (5, 5, 5), ((0, 10000), (-5000, 5000), (-5000, 5000)) + ) + Grid.add_field("zero_field", {"data": np.zeros((5, 5, 5)), "_FillValue": -9999.0}) + return pydda.io.read_from_pyart_grid(Grid) + + +def test_iem_idw_single_station_no_sounding(): + """With one station and no sounding, all grid points inherit the station wind.""" + Grid = _make_small_grid() + station_obs = [{"x": 0.0, "y": 0.0, "z": 0.0, "u": 5.0, "v": 3.0, "w": 0.0}] + Grid = pydda.initialization.make_initialization_from_iem_obs(Grid, station_obs) + # One station → all normalised IDW weights = 1 → uniform result everywhere + assert np.allclose(Grid["u"].values, 5.0) + assert np.allclose(Grid["v"].values, 3.0) + assert np.all(Grid["w"].values == 0.0) + + +def test_iem_idw_sounding_only(): + """Profile only, no stations: output should match sounding at each z level.""" + Grid = _make_small_grid() + grid_z = Grid["z"].values + u_snd = grid_z / 1000.0 # linearly increasing with height + v_snd = np.full_like(grid_z, 2.0) + profile = pyart.core.HorizontalWindProfile.from_u_and_v(grid_z, u_snd, v_snd) + + Grid = pydda.initialization.make_initialization_from_iem_obs( + Grid, [], profile=profile + ) + + for iz in range(len(grid_z)): + assert np.allclose(Grid["u"].values[0, iz, :, :], u_snd[iz]) + assert np.allclose(Grid["v"].values[0, iz, :, :], v_snd[iz]) + assert np.all(Grid["w"].values == 0.0) + + +def test_iem_idw_sounding_plus_station_anomaly(): + """Station anomaly is added on top of the sounding background at every level.""" + Grid = _make_small_grid() + grid_z = Grid["z"].values + u_back = 3.0 + profile = pyart.core.HorizontalWindProfile.from_u_and_v( + grid_z, np.full_like(grid_z, u_back), np.zeros_like(grid_z) + ) + # Station u=7 at z=0; sounding at z=0 is 3; anomaly = +4 + station_obs = [{"x": 0.0, "y": 0.0, "z": 0.0, "u": 7.0, "v": 0.0, "w": 0.0}] + Grid = pydda.initialization.make_initialization_from_iem_obs( + Grid, station_obs, profile=profile + ) + # Single station → anomaly spread is uniform → u = 3 + 4 = 7 everywhere + assert np.allclose(Grid["u"].values, u_back + 4.0) + assert np.all(Grid["w"].values == 0.0) + + +def test_iem_idw_midpoint_average(): + """Two equidistant stations: the midpoint should receive their average.""" + Grid = _make_small_grid() + grid_x = Grid["x"].values # e.g. [-5000, -2500, 0, 2500, 5000] + xa, xb = float(grid_x[0]), float(grid_x[-1]) + assert np.isclose(abs(xa), abs(xb)), "grid must be symmetric about x=0" + ny2, nx2 = len(Grid["y"].values) // 2, len(Grid["x"].values) // 2 + + station_obs = [ + {"x": xa, "y": 0.0, "z": 0.0, "u": 10.0, "v": 0.0, "w": 0.0}, + {"x": xb, "y": 0.0, "z": 0.0, "u": 0.0, "v": 0.0, "w": 0.0}, + ] + Grid = pydda.initialization.make_initialization_from_iem_obs(Grid, station_obs) + # At (x=0, y=0): equal distances to both stations → average u = 5 + u_mid = Grid["u"].values[0, :, ny2, nx2] + assert np.allclose(u_mid, 5.0, atol=1e-6) + + +def test_iem_idw_higher_power_weights_nearest(): + """Higher IDW power concentrates weight on the nearest station.""" + Grid = _make_small_grid() + grid_x = Grid["x"].values + xa, xb = float(grid_x[0]), float(grid_x[-1]) + ny2 = len(Grid["y"].values) // 2 + + station_obs = [ + {"x": xa, "y": 0.0, "z": 0.0, "u": 10.0, "v": 0.0, "w": 0.0}, + {"x": xb, "y": 0.0, "z": 0.0, "u": 0.0, "v": 0.0, "w": 0.0}, + ] + G1 = pydda.initialization.make_initialization_from_iem_obs( + Grid.copy(deep=True), station_obs, power=1 + ) + G2 = pydda.initialization.make_initialization_from_iem_obs( + Grid.copy(deep=True), station_obs, power=2 + ) + # At x[1] (closer to station A than to B), power=2 should give higher u + u_p1 = G1["u"].values[0, 0, ny2, 1] + u_p2 = G2["u"].values[0, 0, ny2, 1] + assert u_p2 > u_p1 + + +def test_iem_idw_output_shape_and_dims(): + """u, v, w must have shape (1, nz, ny, nx) with dims (time, z, y, x).""" + Grid = _make_small_grid() + nz = len(Grid["z"].values) + ny = len(Grid["y"].values) + nx = len(Grid["x"].values) + station_obs = [{"x": 0.0, "y": 0.0, "z": 0.0, "u": 2.0, "v": 1.0, "w": 0.0}] + Grid = pydda.initialization.make_initialization_from_iem_obs(Grid, station_obs) + expected_shape = (1, nz, ny, nx) + for field in ("u", "v", "w"): + assert Grid[field].values.shape == expected_shape + assert Grid[field].dims == ("time", "z", "y", "x") + + +def test_iem_idw_raises_when_no_data(): + """Empty station_obs with no profile must raise ValueError.""" + Grid = _make_small_grid() + with pytest.raises(ValueError): + pydda.initialization.make_initialization_from_iem_obs(Grid, []) + + def test_get_iem_data(): Grid = pyart.testing.make_empty_grid( (20, 20, 20), ((0, 100000.0), (-100000.0, 100000.0), (-100000.0, 100000.0))