ENH: Add IDW initialization from IEM surface observations

pydda/initialization/__init__.py

@@ -22,10 +22,12 @@
     make_wind_field_from_profile
     make_background_from_wrf
     make_initialization_from_era5
+    make_initialization_from_iem_obs
 
 """
 
 from .wind_fields import make_constant_wind_field
 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

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

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))