From 7a954a4f23b2f675e2adc8914d9fdc0d77318e68 Mon Sep 17 00:00:00 2001 From: Robert Jackson Date: Tue, 29 Sep 2026 17:06:38 -0500 Subject: [PATCH 1/5] FIX: restore SciPy cost-function consistency --- pydda/cost_functions/_cost_functions_numpy.py | 96 ++++++++++++------- pydda/tests/test_cost_functions.py | 41 ++++++++ 2 files changed, 101 insertions(+), 36 deletions(-) diff --git a/pydda/cost_functions/_cost_functions_numpy.py b/pydda/cost_functions/_cost_functions_numpy.py index 45ecbf15..981b7378 100644 --- a/pydda/cost_functions/_cost_functions_numpy.py +++ b/pydda/cost_functions/_cost_functions_numpy.py @@ -300,6 +300,26 @@ def calculate_smoothness_cost(u, v, w, dx, dy, dz, Cx=1e-5, Cy=1e-5, Cz=1e-5): return np.sum(np.nan_to_num(x_term + y_term + z_term)) +def _gradient_adjoint(values, spacing, axis): + """Apply the transpose of ``np.gradient(..., spacing, axis=axis)``. + + ``calculate_smoothness_cost`` uses first-order edge differences, so its + gradient must use the corresponding transpose operator rather than a + second call to ``np.gradient``. Keeping this here also avoids changing + the discretization used by the existing objective function. + """ + values = np.moveaxis(values, axis, -1) + result = np.zeros_like(values) + scale = 1.0 / spacing + result[..., 0] -= values[..., 0] * scale + result[..., 1] += values[..., 0] * scale + result[..., :-2] -= values[..., 1:-1] * (0.5 * scale) + result[..., 2:] += values[..., 1:-1] * (0.5 * scale) + result[..., -2] -= values[..., -1] * scale + result[..., -1] += values[..., -1] * scale + return np.moveaxis(result, -1, axis) + + def calculate_smoothness_gradient( u, v, w, dx, dy, dz, Cx=1e-5, Cy=1e-5, Cz=1e-5, upper_bc=1, upper_bc_mask=None ): @@ -335,24 +355,28 @@ def calculate_smoothness_gradient( y: float array value of gradient of smoothness cost function """ - du = np.zeros(w.shape) - dv = np.zeros(w.shape) - dw = np.zeros(w.shape) - grad_u = np.zeros(w.shape) - grad_v = np.zeros(w.shape) - grad_w = np.zeros(w.shape) - scipy.ndimage.laplace(u, du, mode="wrap") - scipy.ndimage.laplace(v, dv, mode="wrap") - scipy.ndimage.laplace(w, dw, mode="wrap") - du = du / dx - dv = dv / dy - dw = dw / dz - scipy.ndimage.laplace(du, grad_u, mode="wrap") - scipy.ndimage.laplace(dv, grad_v, mode="wrap") - scipy.ndimage.laplace(dw, grad_w, mode="wrap") - grad_u = grad_u / dx - grad_v = grad_v / dy - grad_w = grad_w / dz + dxx = lambda field: np.gradient( + np.gradient(field, dx, axis=2), dx, axis=2 + ) + dyy = lambda field: np.gradient( + np.gradient(field, dy, axis=1), dy, axis=1 + ) + dzz = lambda field: np.gradient( + np.gradient(field, dz, axis=0), dz, axis=0 + ) + + sx = dxx(u) + dxx(v) + dxx(w) + sy = dyy(u) + dyy(v) + dyy(w) + sz = dzz(u) + dzz(v) + dzz(w) + + base_gradient = 2 * ( + Cx * _gradient_adjoint(_gradient_adjoint(sx, dx, 2), dx, 2) + + Cy * _gradient_adjoint(_gradient_adjoint(sy, dy, 1), dy, 1) + + Cz * _gradient_adjoint(_gradient_adjoint(sz, dz, 0), dz, 0) + ) + grad_u = base_gradient + grad_v = base_gradient + grad_w = base_gradient.copy() # Impermeability condition grad_w[0, :, :] = 0 @@ -363,7 +387,21 @@ def calculate_smoothness_gradient( return y.flatten() -def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3, power=2): +def _point_weights(x, y, z, the_point, roi): + """Return normalized inverse-distance weights inside the point ROI.""" + distance = np.sqrt( + (x - the_point["x"]) ** 2 + + (y - the_point["y"]) ** 2 + + (z - the_point["z"]) ** 2 + ) + weight = np.where(distance <= roi, 1.0 / np.maximum(distance, 1.0) ** 2, 0.0) + maximum = np.max(weight) + return weight / maximum if maximum > 0 else weight + + +def calculate_point_cost( + u, v, x, y, z, point_list, Cp=1e-3, power=2, roi=500.0 +): """ Calculates the cost function related to point observations. A mean square error cost function term is applied to points that are within the sphere of influence @@ -400,14 +438,7 @@ def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3, power=2): # Instead of worrying about whole domain, just find points in radius of influence # Since we know that the weight will be zero outside the sphere of influence anyways - dist = np.sqrt( - (x - the_point["x"]) ** 2 - + (y - the_point["y"]) ** 2 - + (z - the_point["z"]) ** 2 - ) - dist = np.maximum(dist, 1.0) - weight = 1 / dist**2 - weight = weight / np.max(weight) + weight = _point_weights(x, y, z, the_point, roi) J += np.sum(weight * ((u - the_point["u"]) ** 2 + (v - the_point["v"]) ** 2)) @@ -451,14 +482,7 @@ def calculate_point_gradient(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0): gradJ_w = np.zeros_like(u) for the_point in point_list: - dist = np.sqrt( - (x - the_point["x"]) ** 2 - + (y - the_point["y"]) ** 2 - + (z - the_point["z"]) ** 2 - ) - dist = np.maximum(dist, 1.0) - weight = 1 / dist**2 - weight = weight / np.max(weight) + weight = _point_weights(x, y, z, the_point, roi) gradJ_u += 2 * weight * (u - the_point["u"]) gradJ_v += 2 * weight * (v - the_point["v"]) @@ -740,7 +764,7 @@ def calculate_vertical_vorticity_cost(u, v, w, dx, dy, dz, Ut, Vt, coeff=1e-5): dwdy = np.gradient(w, dy, axis=1) dwdx = np.gradient(w, dx, axis=2) dudx = np.gradient(u, dx, axis=2) - dvdy = np.gradient(v, dy, axis=2) + dvdy = np.gradient(v, dy, axis=1) dudy = np.gradient(u, dy, axis=1) zeta = dvdx - dudy dzeta_dx = np.gradient(zeta, dx, axis=2) diff --git a/pydda/tests/test_cost_functions.py b/pydda/tests/test_cost_functions.py index 2764be4e..75258348 100644 --- a/pydda/tests/test_cost_functions.py +++ b/pydda/tests/test_cost_functions.py @@ -268,6 +268,47 @@ def test_calculate_smoothness_cost(): assert cost > 0 +def test_calculate_smoothness_gradient_matches_cost(): + """The NumPy smoothness gradient must differentiate its objective.""" + rng = np.random.default_rng(4) + shape = (5, 6, 7) + winds = rng.normal(size=(3,) + shape) + spacing = (1.3, 2.1, 0.7) + coeffs = (0.2, 0.7, 1.1) + + def cost(x): + return pydda.cost_functions.calculate_smoothness_cost( + x[0], x[1], x[2], *spacing, *coeffs + ) + + gradient = pydda.cost_functions.calculate_smoothness_gradient( + *winds, *spacing, *coeffs, upper_bc=0 + ).reshape(winds.shape) + direction = rng.normal(size=winds.shape) + direction[2, 0, :, :] = 0.0 # w is fixed at the lower boundary + epsilon = 1e-6 + finite_difference = (cost(winds + epsilon * direction) - cost( + winds - epsilon * direction + )) / (2 * epsilon) + np.testing.assert_allclose( + np.sum(gradient * direction), finite_difference, rtol=2e-5, atol=2e-5 + ) + + +def test_numpy_point_cost_accepts_roi(): + """The SciPy objective passes roi through to the NumPy point cost.""" + shape = (3, 3, 3) + zeros = np.zeros(shape) + point = [{"u": 1.0, "v": -1.0, "x": 0.0, "y": 0.0, "z": 0.0}] + cost = pydda.cost_functions.calculate_point_cost( + zeros, zeros, zeros, zeros, zeros, point, roi=100.0 + ) + assert np.isfinite(cost) + assert pydda.cost_functions.calculate_point_cost( + zeros, zeros, zeros, zeros, zeros, point, roi=0.1 + ) > 0 + + @pytest.mark.skipif(not JAX_AVAILABLE, reason="Jax not installed") def test_calculate_smoothness_cost_jax(): """The Laplacian of a constant field is zero""" From aed316df694efe29a0ce135344ccc849c5829dfc Mon Sep 17 00:00:00 2001 From: Robert Jackson Date: Tue, 29 Sep 2026 17:10:51 -0500 Subject: [PATCH 2/5] Ruff ruff --- doc/source/user_guide/retrieving_winds.rst | 48 ++++ examples/plot_examples.py | 3 +- pydda/__init__.py | 1 + pydda/constraints/__init__.py | 6 + pydda/cost_functions/__init__.py | 2 + pydda/cost_functions/_cost_functions_jax.py | 204 +++++++++++++- pydda/cost_functions/_cost_functions_numpy.py | 257 +++++++++++++++++- pydda/cost_functions/cost_functions.py | 212 +++++++++++++-- pydda/retrieval/wind_retrieve.py | 106 ++++++++ pydda/tests/test_cost_functions.py | 133 ++++++++- pydda/tests/test_initialization.py | 151 ++++++++++ pydda/tests/test_retrieval.py | 136 +++++++++ 12 files changed, 1203 insertions(+), 56 deletions(-) diff --git a/doc/source/user_guide/retrieving_winds.rst b/doc/source/user_guide/retrieving_winds.rst index f068d349..a1719dfb 100644 --- a/doc/source/user_guide/retrieving_winds.rst +++ b/doc/source/user_guide/retrieving_winds.rst @@ -22,11 +22,59 @@ PyDDA minimizes a cost function :math:`J` that corresponds to various penalties +----------------------------------------------------------------+----------------------------+ | :math:`J_{s} = \nabla^2 V` | Wind field smoothness | +----------------------------------------------------------------+----------------------------+ +| :math:`J_{vad} = \sum_{i_{vad}} [V_{VVAD} - \textbf{V}]^2` | VVAD winds | ++----------------------------------------------------------------+----------------------------+ The cost function to be minimized is a weighted sum of the various cost functions in PyDDA and are represented in Equation (1): :math:`J = c_{m}J_{m} + c_{o}J_{o} + c_{b}J_{b} + c_{s}J_{s} + ...` (1) +--------------------------------------------- +Filling data voids with the VVAD constraint +--------------------------------------------- + +The dual-Doppler problem is only well posed where two radars observe the same +point from sufficiently different directions. Outside the dual-Doppler lobes, +and at the low levels that lie below the lowest radar gate, the horizontal wind +is underdetermined by the radial velocities alone. + +:math:`J_{vad}` addresses this following the SWIRL system of +`Protat et al. (2024) `_. A +variational velocity azimuth display (VVAD) fits vertical profiles of a linear +wind model - :math:`U_0`, :math:`V_0`, the divergence, and the stretching and +shearing deformations - to each radar's Doppler velocities. Those profiles +reconstruct a horizontal wind at *every* point within range of the radar, +including points with no observation at all, which is what lets the term fill +data voids. + +Following Eq. (19) of Protat et al. (2024), the term carries a switch +:math:`i_{vad}` that is zero wherever more than one radar contributes, so the +VVAD never competes with genuine multi-Doppler information. PyDDA computes the +switch automatically from the same radar coverage arrays it uses to weight the +radial velocity term. + +.. code-block:: python + + Grids = pydda.constraints.make_constraint_from_vvad(Grids) + Grids, parameters = pydda.retrieval.get_dd_wind_field(Grids, Cvad=1.0) + +Protat et al. (2024) use a weight of 1, equal to that of the radial velocity +constraint. The constraint is available for the ``"scipy"`` and ``"jax"`` +engines. + +Two caveats are worth keeping in mind. First, every VAD technique assumes the +horizontal wind is linear across the domain, which is reasonable in stratiform +precipitation but not in convection. Second, PyDDA applies two acceptance +criteria beyond the point count used in the paper - a minimum azimuthal +coverage and a maximum condition number - because on a PyDDA analysis grid a +height level can hold hundreds of valid velocities confined to a narrow sector +of azimuth, which cannot determine the fit. See +:py:func:`pydda.constraints.vvad_retrieval` for details. + +PyDDA does not implement the double VAD (DVAD) that Protat et al. (2024) use +where two suitably spaced radars overlap, nor the optical flow term of their +Eq. (19). + ------------------------------- Doing your first wind retrieval ------------------------------- diff --git a/examples/plot_examples.py b/examples/plot_examples.py index 3f5b0446..71103b89 100644 --- a/examples/plot_examples.py +++ b/examples/plot_examples.py @@ -33,8 +33,9 @@ frz=5000.0, filter_window=5, mask_outside_opt=True, - upper_bc=1, + upper_bc=2, wind_tol=0.5, + max_iterations=150, engine="scipy", parallel=False, ) diff --git a/pydda/__init__.py b/pydda/__init__.py index 0e6784fe..34955719 100644 --- a/pydda/__init__.py +++ b/pydda/__init__.py @@ -11,6 +11,7 @@ from . import tests from . import constraints from . import io +from . import terrain __version__ = "2.5.0" diff --git a/pydda/constraints/__init__.py b/pydda/constraints/__init__.py index 5d6db7b2..95d46ae2 100644 --- a/pydda/constraints/__init__.py +++ b/pydda/constraints/__init__.py @@ -23,6 +23,9 @@ make_constraint_from_era5 download_needed_era_data get_iem_obs + make_constraint_from_vvad + vvad_retrieval + vvad_horizontal_wind """ @@ -31,3 +34,6 @@ from .model_data import make_constraint_from_era5 from .model_data import download_needed_era_data from .station_data import get_iem_obs +from .vad_data import make_constraint_from_vvad +from .vad_data import vvad_retrieval +from .vad_data import vvad_horizontal_wind diff --git a/pydda/cost_functions/__init__.py b/pydda/cost_functions/__init__.py index 54c392df..44e908ee 100644 --- a/pydda/cost_functions/__init__.py +++ b/pydda/cost_functions/__init__.py @@ -75,6 +75,8 @@ calculate_point_cost calculate_point_gradient calculate_echo_top_mask + calculate_vad_cost + calculate_vad_gradient """ import cmweather diff --git a/pydda/cost_functions/_cost_functions_jax.py b/pydda/cost_functions/_cost_functions_jax.py index b451a7c2..8ac68a6e 100644 --- a/pydda/cost_functions/_cost_functions_jax.py +++ b/pydda/cost_functions/_cost_functions_jax.py @@ -414,7 +414,7 @@ def calculate_point_gradient(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0): return gradJ * Cp -def calculate_mass_continuity(u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1): +def calculate_mass_continuity(u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1, fluid=None): """ Calculates the mass continuity cost function by taking the divergence of the wind field. @@ -442,6 +442,11 @@ def calculate_mass_continuity(u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1): Constant controlling contribution of mass continuity to cost function anel: int = 1 use anelastic approximation, 0=don't + fluid: 3D bool array or None + Mask of grid points that lie above the terrain surface, as returned by + :func:`pydda.terrain.terrain_ghost_cells`. Points below the terrain do + not contribute to the divergence penalty. Set to None to apply the + constraint everywhere, which is the behavior without terrain. Returns ------- @@ -462,11 +467,27 @@ def calculate_mass_continuity(u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1): anel_term = w / rho * drho_dz else: anel_term = jnp.zeros(w.shape) - return coeff * jnp.sum(jnp.square(dudx + dvdy + dwdz + anel_term)) / 2.0 + + div2 = jnp.square(dudx + dvdy + dwdz + anel_term) + if fluid is not None: + div2 = fluid * div2 + return coeff * jnp.sum(div2) / 2.0 def calculate_mass_continuity_gradient( - u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1, upper_bc=1, upper_bc_mask=None + u, + v, + w, + z, + dx, + dy, + dz, + coeff=1500.0, + anel=1, + upper_bc=1, + upper_bc_mask=None, + lower_bc=1, + fluid=None, ): """ Calculates the gradient of mass continuity cost function. This is done by @@ -515,17 +536,98 @@ def calculate_mass_continuity_gradient( else: z_in = z primals, fun_vjp = jax.vjp( - calculate_mass_continuity, u, v, w, z_in, dx, dy, dz, coeff, anel + calculate_mass_continuity, u, v, w, z_in, dx, dy, dz, coeff, anel, fluid ) - grad_u, grad_v, grad_w, _, _, _, _, _, _ = fun_vjp(1.0) + grad_u, grad_v, grad_w, _, _, _, _, _, _, _ = fun_vjp(1.0) # Impermeability condition - grad_w = grad_w.at[0, :, :].set(0) + if lower_bc == 1: + grad_w = grad_w.at[0, :, :].set(0) grad_w = _apply_upper_bc(grad_w, upper_bc, upper_bc_mask) y = jnp.stack([grad_u, grad_v, grad_w], axis=0) return y.flatten() +def calculate_terrain_cost(u, v, w, terrain, coeff=1.0): + """ + Calculates the cost function enforcing no-penetration through terrain. + + This is the penalty form of the ghost-cell immersed boundary condition + built by :func:`pydda.terrain.terrain_ghost_cells`. It replaces the flat + impermeability condition ``w = 0`` at the bottom of the domain with + no-penetration through the actual terrain surface, so that the retrieved + flow rises over terrain rather than through it. + + Solid cells that are not ghost cells carry no boundary information and are + damped towards zero by the same coefficient, so that they do not wander + during the optimization. + + Parameters + ---------- + u: Float array + Float array with u component of wind field + v: Float array + Float array with v component of wind field + w: Float array + Float array with w component of wind field + terrain: dict + Terrain geometry as returned by + :func:`pydda.terrain.terrain_ghost_cells`. + coeff: float + Constant controlling contribution of the terrain condition to the cost + function. + + Returns + ------- + J: float + value of the terrain cost function + """ + U = jnp.stack([u.ravel(), v.ravel(), w.ravel()], axis=1) + u_ghost = U[terrain["ghost_flat"]] + u_image = jnp.einsum("gsc,gs->gc", U[terrain["stencil_flat"]], terrain["stencil_w"]) + c = u_ghost - jnp.einsum("gij,gj->gi", terrain["M"], u_image) + + J = jnp.sum(jnp.square(c)) + solid_free = terrain["solid_free_flat"] + if len(solid_free) > 0: + J = J + jnp.sum(jnp.square(U[solid_free])) + return coeff * J / 2.0 + + +def calculate_terrain_gradient(u, v, w, terrain, coeff=1.0): + """ + Calculates the gradient of the terrain cost function. + + Parameters + ---------- + u: Float array + Float array with u component of wind field + v: Float array + Float array with v component of wind field + w: Float array + Float array with w component of wind field + terrain: dict + Terrain geometry as returned by + :func:`pydda.terrain.terrain_ghost_cells`. + coeff: float + Constant controlling contribution of the terrain condition to the cost + function. + + Returns + ------- + y: float array + value of gradient of the terrain cost function + """ + primals, fun_vjp = jax.vjp( + lambda uu, vv, ww: calculate_terrain_cost(uu, vv, ww, terrain, coeff), + u, + v, + w, + ) + grad_u, grad_v, grad_w = fun_vjp(1.0) + return jnp.stack([grad_u, grad_v, grad_w], axis=0).flatten() + + def calculate_background_cost(u, v, w, weights, u_back, v_back, Cb=0.01): """ Calculates the background cost function. The background cost function is @@ -820,3 +922,93 @@ def calculate_model_gradient(u, v, w, weights, u_model, v_model, w_model, coeff= u_grad, v_grad, w_grad, _, _, _, _, _ = fun_vjp(1.0) y = jnp.stack([u_grad, v_grad, w_grad], axis=0) return y.flatten().copy() + + +def calculate_vad_cost(u, v, vad_weights, u_vad, v_vad, coeff=1.0): + """ + Calculates the cost function for the VVAD constraint of Protat et al. + (2024), the sixth term of their Eq. (19), + + .. math:: + + i_{vad} W_{vad} \\sum (U - U_{vad})^2 + (V - V_{vad})^2 + + Vertical velocities do not enter this cost function; the VVAD supplies + horizontal wind components only. + + Parameters + ---------- + u: 3D array + Float array with u component of wind field + v: 3D array + Float array with v component of wind field + vad_weights: 3D array + Float array showing how much each point from the VVAD weighs into the + constraint. This carries the :math:`i_{vad}` switch of Eq. (19), i.e. + it is zero wherever more than one radar observes the point. + u_vad: 3D array + Float array with u component of wind field from the VVAD + v_vad: 3D array + Float array with v component of wind field from the VVAD + coeff: float + Weighting coefficient + + Returns + ------- + Jvad: float + Value of the VVAD cost function + + References + ---------- + Protat, A., V. Louf, and J. P. Brook, 2024: SWIRL: The First Australian + Operational Radar-Based 3D Wind Analysis System. *J. Atmos. Oceanic + Technol.*, **41**, 891-910, https://doi.org/10.1175/JTECH-D-23-0155.1. + """ + return coeff * jnp.sum( + (jnp.square(u - u_vad) + jnp.square(v - v_vad)) * vad_weights + ) + + +def calculate_vad_gradient( + u, v, vad_weights, u_vad, v_vad, coeff=1.0, upper_bc=1, upper_bc_mask=None +): + """ + Calculates the gradient of the VVAD cost function. The gradient with + respect to w is zero everywhere, since the VVAD constrains the horizontal + components only. + + Parameters + ---------- + u: 3D array + Float array with u component of wind field + v: 3D array + Float array with v component of wind field + vad_weights: 3D array + Float array showing how much each point from the VVAD weighs into the + constraint. + u_vad: 3D array + Float array with u component of wind field from the VVAD + v_vad: 3D array + Float array with v component of wind field from the VVAD + coeff: float + Weighting coefficient + upper_bc: int + Set to 1 to impose w = 0 at the top of the domain, 2 to impose it + above the echo top, and 0 for no upper boundary condition. + upper_bc_mask: 3D bool array or None + The echo top mask used when *upper_bc* is 2. + + Returns + ------- + y: 1D float array + Value of the gradient of the VVAD cost function + """ + primals, fun_vjp = jax.vjp( + calculate_vad_cost, u, v, vad_weights, u_vad, v_vad, coeff + ) + u_grad, v_grad, _, _, _, _ = fun_vjp(1.0) + w_grad = jnp.zeros(u.shape) + w_grad = _apply_upper_bc(w_grad, upper_bc, upper_bc_mask) + + y = jnp.stack([u_grad, v_grad, w_grad], axis=0) + return y.flatten().copy() diff --git a/pydda/cost_functions/_cost_functions_numpy.py b/pydda/cost_functions/_cost_functions_numpy.py index 981b7378..3e62729f 100644 --- a/pydda/cost_functions/_cost_functions_numpy.py +++ b/pydda/cost_functions/_cost_functions_numpy.py @@ -355,15 +355,9 @@ def calculate_smoothness_gradient( y: float array value of gradient of smoothness cost function """ - dxx = lambda field: np.gradient( - np.gradient(field, dx, axis=2), dx, axis=2 - ) - dyy = lambda field: np.gradient( - np.gradient(field, dy, axis=1), dy, axis=1 - ) - dzz = lambda field: np.gradient( - np.gradient(field, dz, axis=0), dz, axis=0 - ) + dxx = lambda field: np.gradient(np.gradient(field, dx, axis=2), dx, axis=2) + dyy = lambda field: np.gradient(np.gradient(field, dy, axis=1), dy, axis=1) + dzz = lambda field: np.gradient(np.gradient(field, dz, axis=0), dz, axis=0) sx = dxx(u) + dxx(v) + dxx(w) sy = dyy(u) + dyy(v) + dyy(w) @@ -399,9 +393,7 @@ def _point_weights(x, y, z, the_point, roi): return weight / maximum if maximum > 0 else weight -def calculate_point_cost( - u, v, x, y, z, point_list, Cp=1e-3, power=2, roi=500.0 -): +def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3, power=2, roi=500.0): """ Calculates the cost function related to point observations. A mean square error cost function term is applied to points that are within the sphere of influence @@ -490,7 +482,7 @@ def calculate_point_gradient(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0): return gradJ * Cp -def calculate_mass_continuity(u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1): +def calculate_mass_continuity(u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1, fluid=None): """ Calculates the mass continuity cost function by taking the divergence of the wind field. @@ -516,6 +508,11 @@ def calculate_mass_continuity(u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1): Constant controlling contribution of mass continuity to cost function anel: int = 1 use anelastic approximation, 0=don't + fluid: 3D bool array or None + Mask of grid points that lie above the terrain surface, as returned by + :func:`pydda.terrain.terrain_ghost_cells`. Points below the terrain do + not contribute to the divergence penalty. Set to None to apply the + constraint everywhere, which is the behavior without terrain. Returns ------- J: float @@ -533,11 +530,26 @@ def calculate_mass_continuity(u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1): anel_term = np.zeros(w.shape) div = dudx + dvdy + dwdz + anel_term + if fluid is not None: + return coeff * np.sum(fluid * np.square(div)) / 2.0 + return coeff * np.sum(np.square(div)) / 2.0 def calculate_mass_continuity_gradient( - u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1, upper_bc=1, upper_bc_mask=None + u, + v, + w, + z, + dx, + dy, + dz, + coeff=1500.0, + anel=1, + upper_bc=1, + upper_bc_mask=None, + lower_bc=1, + fluid=None, ): """ Calculates the gradient of mass continuity cost function. This is done by @@ -571,6 +583,16 @@ def calculate_mass_continuity_gradient( upper_bc_mask: 3D bool array or None The grid points at which w is held fixed when *upper_bc* is 2, as returned by :func:`pydda.cost_functions.calculate_echo_top_mask`. + lower_bc: int + Lower boundary (impermeability) condition. 1 enforces w = 0 on the + bottom plane of the domain, which assumes flat ground there. 0 + disables it, which is what the terrain boundary condition needs, since + over terrain the bottom plane is underground in some columns and free + air in others. The legacy booleans True and False are equivalent to 1 + and 0. + fluid: 3D bool array or None + Mask of grid points above the terrain surface. See + :func:`calculate_mass_continuity`. Returns ------- y: float array @@ -588,18 +610,136 @@ def calculate_mass_continuity_gradient( div = dudx + dvdy + dwdz + anel_term + if fluid is not None: + div = fluid * div + grad_u = -np.gradient(div, dx, axis=2) * coeff grad_v = -np.gradient(div, dy, axis=1) * coeff grad_w = -np.gradient(div, dz, axis=0) * coeff # Impermeability condition - grad_w[0, :, :] = 0 + if lower_bc == 1: + grad_w[0, :, :] = 0 grad_w = _apply_upper_bc(grad_w, upper_bc, upper_bc_mask) y = np.stack([grad_u, grad_v, grad_w], axis=0) return y.flatten() +def _terrain_residual(u, v, w, terrain): + """ + Ghost-cell immersed-boundary residual c = u_G - M u_I. + + Zero exactly when the velocity interpolated to each image point, reflected + back through the terrain surface, matches the ghost cell value, i.e. when + the flow does not penetrate the terrain. Linear in (u, v, w). + + Returns the stacked velocity array (n_points, 3) alongside the residual + (n_ghost, 3), since the gradient needs both. + """ + U = np.stack([u.ravel(), v.ravel(), w.ravel()], axis=1) + u_ghost = U[terrain["ghost_flat"]] + u_image = np.einsum("gsc,gs->gc", U[terrain["stencil_flat"]], terrain["stencil_w"]) + c = u_ghost - np.einsum("gij,gj->gi", terrain["M"], u_image) + return U, c + + +def calculate_terrain_cost(u, v, w, terrain, coeff=1.0): + """ + Calculates the cost function enforcing no-penetration through terrain. + + This is the penalty form of the ghost-cell immersed boundary condition + built by :func:`pydda.terrain.terrain_ghost_cells`. It replaces the flat + impermeability condition ``w = 0`` at the bottom of the domain with + no-penetration through the actual terrain surface, so that the retrieved + flow rises over terrain rather than through it. + + Solid cells that are not ghost cells carry no boundary information and are + damped towards zero by the same coefficient, so that they do not wander + during the optimization. + + Parameters + ---------- + u: Float array + Float array with u component of wind field + v: Float array + Float array with v component of wind field + w: Float array + Float array with w component of wind field + terrain: dict + Terrain geometry as returned by + :func:`pydda.terrain.terrain_ghost_cells`. + coeff: float + Constant controlling contribution of the terrain condition to the cost + function. + + Returns + ------- + J: float + value of the terrain cost function + """ + U, c = _terrain_residual(u, v, w, terrain) + J = np.sum(np.square(c)) + solid_free = terrain["solid_free_flat"] + if len(solid_free) > 0: + J += np.sum(np.square(U[solid_free])) + return coeff * J / 2.0 + + +def calculate_terrain_gradient(u, v, w, terrain, coeff=1.0): + """ + Calculates the gradient of the terrain cost function. + + The residual is linear in the wind field, so this gradient is exact rather + than an approximation. + + Parameters + ---------- + u: Float array + Float array with u component of wind field + v: Float array + Float array with v component of wind field + w: Float array + Float array with w component of wind field + terrain: dict + Terrain geometry as returned by + :func:`pydda.terrain.terrain_ghost_cells`. + coeff: float + Constant controlling contribution of the terrain condition to the cost + function. + + Returns + ------- + y: float array + value of gradient of the terrain cost function + """ + U, c = _terrain_residual(u, v, w, terrain) + grad = np.zeros_like(U) + + # Ghost cell indices are unique, so this is a plain assignment + grad[terrain["ghost_flat"]] += coeff * c + + # Each image point scatters back onto its eight stencil corners, which are + # shared between neighboring ghost cells, so this has to accumulate. + stencil_flat = terrain["stencil_flat"] + contrib = ( + -coeff + * terrain["stencil_w"][..., None] + * np.einsum("gij,gi->gj", terrain["M"], c)[:, None, :] + ) + flat = stencil_flat.ravel() + for comp in range(3): + grad[:, comp] += np.bincount( + flat, weights=contrib[:, :, comp].ravel(), minlength=grad.shape[0] + ) + + solid_free = terrain["solid_free_flat"] + if len(solid_free) > 0: + grad[solid_free] += coeff * U[solid_free] + + return grad.T.reshape((3,) + u.shape).flatten() + + def calculate_fall_speed(grid, refl_field=None, frz=4500.0): """ Estimates fall speed based on reflectivity. @@ -976,3 +1116,90 @@ def calculate_model_gradient(u, v, w, weights, u_model, v_model, w_model, coeff= y = np.stack([u_grad, v_grad, w_grad], axis=0) return y.flatten() + + +def calculate_vad_cost(u, v, vad_weights, u_vad, v_vad, coeff=1.0): + """ + Calculates the cost function for the VVAD constraint of Protat et al. + (2024), the sixth term of their Eq. (19), + + .. math:: + + i_{vad} W_{vad} \\sum (U - U_{vad})^2 + (V - V_{vad})^2 + + Vertical velocities do not enter this cost function; the VVAD supplies + horizontal wind components only. + + Parameters + ---------- + u: 3D array + Float array with u component of wind field + v: 3D array + Float array with v component of wind field + vad_weights: 3D array + Float array showing how much each point from the VVAD weighs into the + constraint. This carries the :math:`i_{vad}` switch of Eq. (19), i.e. + it is zero wherever more than one radar observes the point. + u_vad: 3D array + Float array with u component of wind field from the VVAD + v_vad: 3D array + Float array with v component of wind field from the VVAD + coeff: float + Weighting coefficient + + Returns + ------- + Jvad: float + Value of the VVAD cost function + + References + ---------- + Protat, A., V. Louf, and J. P. Brook, 2024: SWIRL: The First Australian + Operational Radar-Based 3D Wind Analysis System. *J. Atmos. Oceanic + Technol.*, **41**, 891-910, https://doi.org/10.1175/JTECH-D-23-0155.1. + """ + return coeff * np.sum((np.square(u - u_vad) + np.square(v - v_vad)) * vad_weights) + + +def calculate_vad_gradient( + u, v, vad_weights, u_vad, v_vad, coeff=1.0, upper_bc=1, upper_bc_mask=None +): + """ + Calculates the gradient of the VVAD cost function, which is simply twice + the weighted difference between the analysis and the VVAD wind for each of + u and v. The gradient with respect to w is zero everywhere, since the VVAD + constrains the horizontal components only. + + Parameters + ---------- + u: 3D array + Float array with u component of wind field + v: 3D array + Float array with v component of wind field + vad_weights: 3D array + Float array showing how much each point from the VVAD weighs into the + constraint. + u_vad: 3D array + Float array with u component of wind field from the VVAD + v_vad: 3D array + Float array with v component of wind field from the VVAD + coeff: float + Weighting coefficient + upper_bc: int + Set to 1 to impose w = 0 at the top of the domain, 2 to impose it + above the echo top, and 0 for no upper boundary condition. + upper_bc_mask: 3D bool array or None + The echo top mask used when *upper_bc* is 2. + + Returns + ------- + y: 1D float array + Value of the gradient of the VVAD cost function + """ + u_grad = coeff * 2 * (u - u_vad) * vad_weights + v_grad = coeff * 2 * (v - v_vad) * vad_weights + w_grad = np.zeros(u.shape) + w_grad = _apply_upper_bc(w_grad, upper_bc, upper_bc_mask) + + y = np.stack([u_grad, v_grad, w_grad], axis=0) + return y.flatten() diff --git a/pydda/cost_functions/cost_functions.py b/pydda/cost_functions/cost_functions.py index 52618238..9ef5395f 100644 --- a/pydda/cost_functions/cost_functions.py +++ b/pydda/cost_functions/cost_functions.py @@ -24,6 +24,17 @@ from . import _cost_functions_tensorflow +def _fluid_mask(parameters): + """ + The above-terrain mask used to restrict the mass continuity constraint, + or None when no terrain boundary condition is in use. + """ + terrain = getattr(parameters, "terrain", None) + if terrain is None or getattr(parameters, "Cterrain", 0.0) <= 0: + return None + return terrain["fluid"] + + def J_function(winds, parameters): """ Calculates the total cost function. This typically does not need to be @@ -46,6 +57,10 @@ def J_function(winds, parameters): J: float The value of the cost function """ + # Only the scipy and jax engines support the terrain boundary condition + Jterrain = 0 + # Only the scipy and jax engines support the VVAD constraint + Jvad = 0 if parameters.engine == "tensorflow": if not TENSORFLOW_AVAILABLE: raise ImportError( @@ -199,10 +214,22 @@ def J_function(winds, parameters): parameters.dy, parameters.dz, coeff=parameters.Cm, + fluid=_fluid_mask(parameters), ) else: Jmass = 0 + if parameters.Cterrain > 0: + Jterrain = _cost_functions_numpy.calculate_terrain_cost( + winds[0], + winds[1], + winds[2], + parameters.terrain, + coeff=parameters.Cterrain, + ) + else: + Jterrain = 0 + if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: Jsmooth = _cost_functions_numpy.calculate_smoothness_cost( winds[0], @@ -274,42 +301,64 @@ def J_function(winds, parameters): ) else: Jpoint = 0 + + if parameters.Cvad > 0: + Jvad = _cost_functions_numpy.calculate_vad_cost( + winds[0], + winds[1], + parameters.vad_weights, + parameters.u_vad, + parameters.v_vad, + coeff=parameters.Cvad, + ) + else: + Jvad = 0 elif parameters.engine == "jax": return J_function_jax(winds, parameters) if parameters.Nfeval % 10 == 0: - print( - ( - "Nfeval | Jvel | Jmass | Jsmooth | Jbg | Jvort | Jmodel | Jpoint |" - + " Max w " - ) - ) - print( - ( - "{:7d}".format(int(parameters.Nfeval)) - + "|" - + "{:9.4f}".format(float(Jvel)) - + "|" - + "{:9.4f}".format(float(Jmass)) - + "|" - + "{:9.4f}".format(float(Jsmooth)) - + "|" - + "{:9.4f}".format(float(Jbackground)) - + "|" - + "{:9.4f}".format(float(Jvorticity)) - + "|" - + "{:9.4f}".format(float(Jmod)) - + "|" - + "{:9.4f}".format(float(Jpoint)) - + "|" - + "{:9.4f}".format(np.ma.max(np.ma.abs(winds[2]))) - ) + header = "Nfeval | Jvel | Jmass | Jsmooth | Jbg | Jvort | Jmodel | Jpoint |" + row = ( + "{:7d}".format(int(parameters.Nfeval)) + + "|" + + "{:9.4f}".format(float(Jvel)) + + "|" + + "{:9.4f}".format(float(Jmass)) + + "|" + + "{:9.4f}".format(float(Jsmooth)) + + "|" + + "{:9.4f}".format(float(Jbackground)) + + "|" + + "{:9.4f}".format(float(Jvorticity)) + + "|" + + "{:9.4f}".format(float(Jmod)) + + "|" + + "{:9.4f}".format(float(Jpoint)) + + "|" ) + if parameters.Cterrain > 0: + header += " Jterr |" + row += "{:9.4f}".format(float(Jterrain)) + "|" + if parameters.Cvad > 0: + header += " Jvad |" + row += "{:9.4f}".format(float(Jvad)) + "|" + print(header + " Max w ") + print(row + "{:9.4f}".format(np.ma.max(np.ma.abs(winds[2])))) parameters.Nfeval += 1 # print("The cost functions print", Jvel + Jmass) - return Jvel + Jmass + Jsmooth + Jbackground + Jvorticity + Jmod + Jpoint + return ( + Jvel + + Jmass + + Jsmooth + + Jbackground + + Jvorticity + + Jmod + + Jpoint + + Jterrain + + Jvad + ) def grad_J(winds, parameters): @@ -553,6 +602,19 @@ def grad_J(winds, parameters): 1, parameters.upper_bc, parameters.upper_bc_mask, + parameters.lower_bc, + _fluid_mask(parameters), + ) + ) + if parameters.Cterrain > 0: + futures.append( + pool.submit( + _cost_functions_numpy.calculate_terrain_gradient, + winds[0], + winds[1], + winds[2], + parameters.terrain, + parameters.Cterrain, ) ) if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: @@ -630,6 +692,20 @@ def grad_J(winds, parameters): parameters.roi, ) ) + if parameters.Cvad > 0: + futures.append( + pool.submit( + _cost_functions_numpy.calculate_vad_gradient, + winds[0], + winds[1], + parameters.vad_weights, + parameters.u_vad, + parameters.v_vad, + parameters.Cvad, + parameters.upper_bc, + parameters.upper_bc_mask, + ) + ) grad = sum(f.result() for f in futures) else: grad = _cost_functions_numpy.calculate_grad_radial_vel( @@ -659,6 +735,17 @@ def grad_J(winds, parameters): coeff=parameters.Cm, upper_bc=parameters.upper_bc, upper_bc_mask=parameters.upper_bc_mask, + lower_bc=parameters.lower_bc, + fluid=_fluid_mask(parameters), + ) + + if parameters.Cterrain > 0: + grad += _cost_functions_numpy.calculate_terrain_gradient( + winds[0], + winds[1], + winds[2], + parameters.terrain, + coeff=parameters.Cterrain, ) if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: @@ -726,6 +813,18 @@ def grad_J(winds, parameters): roi=parameters.roi, ) + if parameters.Cvad > 0: + grad += _cost_functions_numpy.calculate_vad_gradient( + winds[0], + winds[1], + parameters.vad_weights, + parameters.u_vad, + parameters.v_vad, + coeff=parameters.Cvad, + upper_bc=parameters.upper_bc, + upper_bc_mask=parameters.upper_bc_mask, + ) + # Let's see if we need to enforce strong boundary conditions if parameters.const_boundary_cond is True: grad = np.reshape( @@ -804,10 +903,22 @@ def J_function_jax(winds, parameters): parameters.dy, parameters.dz, coeff=parameters.Cm, + fluid=_fluid_mask(parameters), ) else: Jmass = 0 + if parameters.Cterrain > 0: + Jterrain = _cost_functions_jax.calculate_terrain_cost( + winds[0], + winds[1], + winds[2], + parameters.terrain, + coeff=parameters.Cterrain, + ) + else: + Jterrain = 0 + if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: Jsmooth = _cost_functions_jax.calculate_smoothness_cost( winds[0], @@ -880,7 +991,29 @@ def J_function_jax(winds, parameters): else: Jpoint = 0 - return Jvel + Jsmooth + Jmass + Jmod + Jpoint + Jvorticity + Jbackground + if parameters.Cvad > 0: + Jvad = _cost_functions_jax.calculate_vad_cost( + winds[0], + winds[1], + parameters.vad_weights, + parameters.u_vad, + parameters.v_vad, + coeff=parameters.Cvad, + ) + else: + Jvad = 0 + + return ( + Jvel + + Jsmooth + + Jmass + + Jmod + + Jpoint + + Jvorticity + + Jbackground + + Jterrain + + Jvad + ) def grad_jax(winds, parameters): @@ -920,6 +1053,17 @@ def grad_jax(winds, parameters): coeff=parameters.Cm, upper_bc=parameters.upper_bc, upper_bc_mask=parameters.upper_bc_mask, + lower_bc=parameters.lower_bc, + fluid=_fluid_mask(parameters), + ) + + if parameters.Cterrain > 0: + grad += _cost_functions_jax.calculate_terrain_gradient( + winds[0], + winds[1], + winds[2], + parameters.terrain, + coeff=parameters.Cterrain, ) if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: @@ -986,6 +1130,18 @@ def grad_jax(winds, parameters): Cp=parameters.Cpoint, roi=parameters.roi, ) + + if parameters.Cvad > 0: + grad += _cost_functions_jax.calculate_vad_gradient( + winds[0], + winds[1], + parameters.vad_weights, + parameters.u_vad, + parameters.v_vad, + coeff=parameters.Cvad, + upper_bc=parameters.upper_bc, + upper_bc_mask=parameters.upper_bc_mask, + ) return grad diff --git a/pydda/retrieval/wind_retrieve.py b/pydda/retrieval/wind_retrieve.py index 2ce10aad..0286e817 100644 --- a/pydda/retrieval/wind_retrieve.py +++ b/pydda/retrieval/wind_retrieve.py @@ -212,6 +212,12 @@ def __init__(self): self.Cv = 0.0 self.Cmod = 0.0 self.Cpoint = 0.0 + self.Cterrain = 0.0 + self.terrain = None + self.Cvad = 0.0 + self.u_vad = None + self.v_vad = None + self.vad_weights = None self.Ut = 0.0 self.Vt = 0.0 self.upper_bc = 1 @@ -253,6 +259,9 @@ def _get_dd_wind_field_scipy( Cv=0.0, Cmod=0.0, Cpoint=0.0, + Cterrain=0.0, + terrain=None, + Cvad=0.0, cvtol=1e-2, gtol=1e-2, Jveltol=100.0, @@ -272,6 +281,7 @@ def _get_dd_wind_field_scipy( min_bca=30.0, max_bca=150.0, upper_bc=1, + lower_bc=1, above=2.0, model_fields=None, output_cost_functions=True, @@ -324,6 +334,12 @@ def _get_dd_wind_field_scipy( prev_grid = g + if Cvad > 0 and engine.lower() not in ("scipy", "jax"): + raise NotImplementedError( + 'The VVAD constraint is only implemented for the "scipy" and ' + '"jax" engines, not "%s".' % engine + ) + if engine.lower() == "auglag" and not TENSORFLOW_AVAILABLE: raise ModuleNotFoundError( "Tensorflow 2.6+ needs to be installed for the Augmented Lagrangian solver." @@ -621,6 +637,58 @@ def _get_dd_wind_field_scipy( parameters.upper_bc_mask = calculate_echo_top_mask( parameters.vrs, parameters.z, above=above ) + + parameters.Cterrain = Cterrain + parameters.terrain = terrain + parameters.lower_bc = lower_bc + if Cterrain > 0: + if terrain is None: + raise ValueError( + "Cterrain is nonzero but no terrain was given. Build one with " + "pydda.terrain.terrain_ghost_cells and pass it as the terrain " + "keyword." + ) + if engine.lower() not in ("scipy", "jax"): + raise NotImplementedError( + "The terrain boundary condition is only implemented for the " + '"scipy" and "jax" engines, not "%s". The augmented Lagrangian ' + "engine would need it as an equality constraint rather than a " + "penalty." % engine + ) + if lower_bc == 1: + # The flat impermeability condition and the terrain condition are + # mutually exclusive: over terrain the bottom plane of the domain + # is underground in some columns and free air in others. + print( + "Terrain boundary condition enabled, disabling the flat " + "impermeability condition at the bottom of the domain." + ) + parameters.lower_bc = 0 + + parameters.Cvad = Cvad + if Cvad > 0: + if "U_vvad" not in Grids[0].variables or "V_vvad" not in Grids[0].variables: + raise ValueError( + "Cvad is nonzero but the Grid has no U_vvad/V_vvad fields. Add " + "them with pydda.constraints.make_constraint_from_vvad." + ) + u_vad = Grids[0]["U_vvad"].values.squeeze() + v_vad = Grids[0]["V_vvad"].values.squeeze() + + # i_vad of Eq. (19) of Protat et al. (2024): the VVAD constrains the + # horizontal wind only where multi-Doppler information is unavailable, + # i.e. at points seen by at most one radar. parameters.weights has + # already been binarized above, so summing over the radar axis counts + # the radars contributing at each point. + n_rad = np.sum(parameters.weights, axis=0) + has_vad = np.logical_and(np.isfinite(u_vad), np.isfinite(v_vad)) + parameters.vad_weights = np.logical_and(n_rad <= 1, has_vad).astype(float) + parameters.u_vad = np.nan_to_num(u_vad) + parameters.v_vad = np.nan_to_num(v_vad) + print( + "VVAD constraint active at %d points." % int(parameters.vad_weights.sum()) + ) + parameters.points = points parameters.point_list = points parameters.parallel = parallel @@ -1446,6 +1514,34 @@ def get_dd_wind_field( Weight for cost function related to custom constraints. Cpoint: float Weight for cost function related to point observations. + Cterrain: float + Weight for the terrain (no-penetration) boundary condition. Set to a + value greater than zero to replace the flat impermeability condition + at the bottom of the domain with no-penetration through the actual + terrain surface, so that the retrieved flow rises over terrain rather + than through it. Requires *terrain*, and is only available for the + "scipy" and "jax" engines. The residual it penalizes has units of + m/s, the same as the radial velocity cost function, so a value of + order *Co* is a reasonable starting point. + terrain: dict or None + Terrain geometry as returned by + :py:func:`pydda.terrain.terrain_ghost_cells`. Needed when *Cterrain* + is nonzero, ignored otherwise. + Cvad: float + Weight for the VVAD constraint of Protat et al. (2024), the sixth term + of their Eq. (19). Set to a value greater than zero to nudge the + horizontal wind toward the velocity azimuth display reconstruction + wherever multi-Doppler information is unavailable, which supplies + winds in data voids such as the grid points below the lowest radar + gate. Requires the *U_vvad* and *V_vvad* fields, which are added by + :py:func:`pydda.constraints.make_constraint_from_vvad`, and is only + available for the "scipy" and "jax" engines. Protat et al. (2024) use + a weight of 1, equal to that of the radial velocity constraint. + lower_bc: int + Lower boundary (impermeability) condition. 1 enforces w = 0 on the + bottom plane of the domain, which assumes flat ground there, and is + the default. 0 disables it. Supplying a terrain boundary condition + sets this to 0 automatically, since the two are mutually exclusive. weights_obs: list of floating point arrays or None List of weights for each point in grid from each radar in Grids. Set to None to let PyDDA determine this automatically. @@ -1588,6 +1684,16 @@ def get_dd_wind_field( new_grids, u_init, v_init, w_init, engine, **kwargs ) elif engine.lower() == "tensorflow": + if kwargs.get("Cterrain", 0.0) > 0: + raise NotImplementedError( + "The terrain boundary condition is only implemented for the " + '"scipy" and "jax" engines, not "tensorflow".' + ) + if kwargs.get("Cvad", 0.0) > 0: + raise NotImplementedError( + 'The VVAD constraint is only implemented for the "scipy" and ' + '"jax" engines, not "tensorflow".' + ) return _get_dd_wind_field_tensorflow( new_grids, u_init, v_init, w_init, **kwargs ) diff --git a/pydda/tests/test_cost_functions.py b/pydda/tests/test_cost_functions.py index 75258348..4631db0e 100644 --- a/pydda/tests/test_cost_functions.py +++ b/pydda/tests/test_cost_functions.py @@ -287,9 +287,9 @@ def cost(x): direction = rng.normal(size=winds.shape) direction[2, 0, :, :] = 0.0 # w is fixed at the lower boundary epsilon = 1e-6 - finite_difference = (cost(winds + epsilon * direction) - cost( - winds - epsilon * direction - )) / (2 * epsilon) + finite_difference = ( + cost(winds + epsilon * direction) - cost(winds - epsilon * direction) + ) / (2 * epsilon) np.testing.assert_allclose( np.sum(gradient * direction), finite_difference, rtol=2e-5, atol=2e-5 ) @@ -304,9 +304,12 @@ def test_numpy_point_cost_accepts_roi(): zeros, zeros, zeros, zeros, zeros, point, roi=100.0 ) assert np.isfinite(cost) - assert pydda.cost_functions.calculate_point_cost( - zeros, zeros, zeros, zeros, zeros, point, roi=0.1 - ) > 0 + assert ( + pydda.cost_functions.calculate_point_cost( + zeros, zeros, zeros, zeros, zeros, point, roi=0.1 + ) + > 0 + ) @pytest.mark.skipif(not JAX_AVAILABLE, reason="Jax not installed") @@ -1135,3 +1138,121 @@ def test_mass_continuity_gradient_upper_bc_tf(): assert np.all(grid_top[-1] == 0) assert np.all(echo_top[mask] == 0) np.testing.assert_allclose(echo_top[~mask], none_bc[~mask]) + + +from scipy.optimize import approx_fprime + + +def _make_vad_inputs(seed=0, shape=(5, 6, 7)): + """Random winds, VVAD winds, and a partially active i_vad mask.""" + rng = np.random.default_rng(seed) + u = rng.normal(size=shape) + v = rng.normal(size=shape) + w = rng.normal(size=shape) + u_vad = rng.normal(size=shape) + v_vad = rng.normal(size=shape) + vad_weights = (rng.random(shape) > 0.4).astype(float) + return u, v, w, u_vad, v_vad, vad_weights + + +def test_vad_cost(): + """The VVAD cost is the i_vad-weighted sum of squared horizontal + departures, and is zero when the analysis equals the VVAD.""" + u, v, w, u_vad, v_vad, vad_weights = _make_vad_inputs() + + expected = 2.5 * np.sum((np.square(u - u_vad) + np.square(v - v_vad)) * vad_weights) + actual = pydda.cost_functions.np.calculate_vad_cost( + u, v, vad_weights, u_vad, v_vad, coeff=2.5 + ) + np.testing.assert_allclose(actual, expected) + + # No cost where the analysis already matches the VVAD. + assert ( + pydda.cost_functions.np.calculate_vad_cost( + u_vad, v_vad, vad_weights, u_vad, v_vad, coeff=2.5 + ) + == 0.0 + ) + + # No cost at all where i_vad is zero everywhere. + assert ( + pydda.cost_functions.np.calculate_vad_cost( + u, v, np.zeros_like(vad_weights), u_vad, v_vad, coeff=2.5 + ) + == 0.0 + ) + + +def test_vad_gradient(): + """The analytic VVAD gradient matches finite differences, leaves w + untouched, and vanishes wherever i_vad is zero.""" + u, v, w, u_vad, v_vad, vad_weights = _make_vad_inputs() + shape = u.shape + n = int(np.prod(shape)) + coeff = 1.7 + + def J(x): + x = x.reshape((3,) + shape) + return pydda.cost_functions.np.calculate_vad_cost( + x[0], x[1], vad_weights, u_vad, v_vad, coeff=coeff + ) + + x0 = np.concatenate([u.ravel(), v.ravel(), w.ravel()]) + finite_diff = approx_fprime(x0, J, 1e-6) + analytic = pydda.cost_functions.np.calculate_vad_gradient( + u, v, vad_weights, u_vad, v_vad, coeff=coeff, upper_bc=0 + ) + np.testing.assert_allclose(analytic, finite_diff, atol=1e-4) + + # The VVAD constrains the horizontal wind only. + assert np.all(analytic[2 * n :] == 0) + + # Eq. (19) of Protat et al. (2024) switches the term off entirely where + # multi-Doppler information is available, so it must contribute no + # gradient at those points. + grad = analytic.reshape((3,) + shape) + assert np.all(grad[0][vad_weights == 0] == 0) + assert np.all(grad[1][vad_weights == 0] == 0) + + +@pytest.mark.skipif(not JAX_AVAILABLE, reason="Jax not installed") +def test_vad_cost_jax(): + """The jax VVAD cost and gradient agree with the numpy ones.""" + u, v, w, u_vad, v_vad, vad_weights = _make_vad_inputs() + + np_cost = pydda.cost_functions.np.calculate_vad_cost( + u, v, vad_weights, u_vad, v_vad, coeff=2.5 + ) + jax_cost = pydda.cost_functions.jax.calculate_vad_cost( + u, v, vad_weights, u_vad, v_vad, coeff=2.5 + ) + np.testing.assert_allclose(float(jax_cost), np_cost, rtol=1e-5) + + np_grad = pydda.cost_functions.np.calculate_vad_gradient( + u, v, vad_weights, u_vad, v_vad, coeff=2.5 + ) + jax_grad = pydda.cost_functions.jax.calculate_vad_gradient( + u, v, vad_weights, u_vad, v_vad, coeff=2.5 + ) + np.testing.assert_allclose(np.asarray(jax_grad), np_grad, atol=1e-5) + + +def test_vad_gradient_upper_bc(): + """The VVAD gradient honors the upper boundary condition modes even + though its own w gradient is identically zero.""" + u, v, w, u_vad, v_vad, vad_weights = _make_vad_inputs() + mask = np.zeros(u.shape, dtype=bool) + mask[-2:] = True + + for upper_bc in (0, 1, 2): + grad = pydda.cost_functions.np.calculate_vad_gradient( + u, + v, + vad_weights, + u_vad, + v_vad, + coeff=1.0, + upper_bc=upper_bc, + upper_bc_mask=mask, + ) + assert np.all(_w_component(grad, u.shape) == 0) diff --git a/pydda/tests/test_initialization.py b/pydda/tests/test_initialization.py index b2ed3b40..b6a28053 100644 --- a/pydda/tests/test_initialization.py +++ b/pydda/tests/test_initialization.py @@ -254,3 +254,154 @@ def make_grid(origin_lat, origin_lon): speed_true_north = np.sqrt(u_true_north**2 + v_true_north**2) speed_grid_relative = np.sqrt(u_grid_relative**2 + v_grid_relative**2) np.testing.assert_allclose(speed_true_north, speed_grid_relative, rtol=1e-4) + + +def _linear_wind_grid( + U0=7.0, + V0=-3.0, + Ux=1.2e-4, + Uy=3.0e-5, + Vy=-4.0e-4, + fall_speed=0.0, + shift_radar=False, +): + """ + Builds a Grid whose Doppler velocity is synthesized from an exactly linear + horizontal wind field, so the VVAD fit of Protat et al. (2024) must + recover the wind exactly. The flow is irrotational by construction + (Vx == Uy), matching the closure the reconstruction assumes. + + Returns the Grid together with the true u and v fields and the true + control variables. + """ + from pydda.constraints.vad_data import _geometry + + Grid = pyart.testing.make_empty_grid( + (12, 40, 40), ((500.0, 8000.0), (-40000.0, 40000.0), (-40000.0, 40000.0)) + ) + Grid.add_field("velocity", {"data": np.zeros((12, 40, 40)), "_FillValue": -9999.0}) + Grid = pydda.io.read_from_pyart_grid(Grid) + + if shift_radar: + # Move the radar off the grid origin so the reconstruction origin is + # genuinely exercised. + Grid["radar_latitude"].values[0] += 0.12 + Grid["radar_longitude"].values[0] += -0.10 + Grid = pydda.retrieval.angles.add_elevation_as_field(Grid) + + Vx = Uy # irrotational + dx, dy, rh, az, el = _geometry(Grid) + u_true = U0 + Ux * dx + Uy * dy + v_true = V0 + Vx * dx + Vy * dy + + el = np.asarray(el) + vr = ( + np.cos(el) * np.sin(az) * u_true + + np.cos(el) * np.cos(az) * v_true + - np.sin(el) * fall_speed + ) + Grid["velocity"] = Grid["velocity"].copy(data=np.expand_dims(vr, 0)) + + truth = { + "U0": U0, + "V0": V0, + "DIV": Ux + Vy, + "DET": Ux - Vy, + "DES": Uy + Vx, + "fall_speed": fall_speed, + } + return Grid, u_true, v_true, truth + + +@pytest.mark.parametrize("shift_radar", [False, True]) +def test_vvad_recovers_linear_wind(shift_radar): + """A wind field that exactly satisfies the VAD linearity assumption must + be recovered to machine precision at every level.""" + Grid, u_true, v_true, truth = _linear_wind_grid(shift_radar=shift_radar) + profiles = pydda.constraints.vvad_retrieval(Grid, vel_field="velocity") + + assert profiles["valid"].values.all() + ok = profiles["valid"].values + for name, expected in truth.items(): + np.testing.assert_allclose( + profiles[name].values[ok], expected, atol=1e-9, rtol=0 + ) + + +def test_vvad_separates_fall_speed_from_divergence(): + """A uniform fall speed must be recovered on its own rather than aliasing + into the divergence.""" + Grid, _, _, truth = _linear_wind_grid(fall_speed=2.5) + profiles = pydda.constraints.vvad_retrieval(Grid, vel_field="velocity") + + ok = profiles["valid"].values + np.testing.assert_allclose(profiles["fall_speed"].values[ok], 2.5, atol=1e-9) + np.testing.assert_allclose( + profiles["DIV"].values[ok], truth["DIV"], atol=1e-12, rtol=0 + ) + + +def test_vvad_horizontal_wind_round_trip(): + """Reconstructing from the fitted profiles must return the linear wind + field it was fitted to, and NaN beyond max_range.""" + Grid, u_true, v_true, _ = _linear_wind_grid() + profiles = pydda.constraints.vvad_retrieval(Grid, vel_field="velocity") + + u_vad, v_vad = pydda.constraints.vvad_horizontal_wind(profiles, Grid) + np.testing.assert_allclose(u_vad, u_true, atol=1e-7) + np.testing.assert_allclose(v_vad, v_true, atol=1e-7) + + rh = np.sqrt(Grid["point_x"].values ** 2 + Grid["point_y"].values ** 2) + u_clip, _ = pydda.constraints.vvad_horizontal_wind( + profiles, Grid, max_range=20000.0 + ) + assert np.all(np.isnan(u_clip[rh > 20000.0])) + assert np.all(np.isfinite(u_clip[rh <= 20000.0])) + + +def test_vvad_rejects_undersampled_levels(): + """Each acceptance criterion must reject a level it is meant to, and an + invalid level must reconstruct as NaN rather than as a wild extrapolation. + """ + Grid, _, _, _ = _linear_wind_grid() + + for kwargs in ( + {"min_points": 10**9}, + {"min_azimuth_coverage": 1.01}, + {"max_condition": 1e-6}, + ): + profiles = pydda.constraints.vvad_retrieval( + Grid, vel_field="velocity", **kwargs + ) + assert not profiles["valid"].values.any(), kwargs + assert np.all(np.isnan(profiles["U0"].values)), kwargs + + u_vad, v_vad = pydda.constraints.vvad_horizontal_wind(profiles, Grid) + assert np.all(np.isnan(u_vad)) and np.all(np.isnan(v_vad)), kwargs + + +def test_make_constraint_from_vvad_adds_fields(): + """The constraint builder writes U_vvad/V_vvad into every Grid, and the + two combine modes cover the same points by different means.""" + Grid, u_true, v_true, _ = _linear_wind_grid() + Grids = pydda.constraints.make_constraint_from_vvad( + [Grid, Grid.copy(deep=True)], vel_field="velocity" + ) + + for g in Grids: + assert "U_vvad" in g.variables and "V_vvad" in g.variables + assert g["U_vvad"].dims == ("time", "z", "y", "x") + # No W_vvad: the VVAD supplies horizontal components only. + assert "W_vvad" not in g.variables + + np.testing.assert_allclose(Grids[0]["U_vvad"].values.squeeze(), u_true, atol=1e-7) + + nearest = pydda.constraints.make_constraint_from_vvad( + [Grid.copy(deep=True)], vel_field="velocity", combine="nearest" + ) + np.testing.assert_allclose(nearest[0]["U_vvad"].values.squeeze(), u_true, atol=1e-7) + + with pytest.raises(ValueError): + pydda.constraints.make_constraint_from_vvad( + [Grid], vel_field="velocity", combine="not_a_mode" + ) diff --git a/pydda/tests/test_retrieval.py b/pydda/tests/test_retrieval.py index a2c5c01e..b0a470e5 100644 --- a/pydda/tests/test_retrieval.py +++ b/pydda/tests/test_retrieval.py @@ -515,3 +515,139 @@ def test_nested_retrieval(): quiverkey_loc="bottom_right", ) return fig + + +# Common arguments for the VVAD constraint tests. The filtering and masking are +# turned off so the retrieved winds can be compared point by point. +VVAD_KWARGS = dict( + Co=1.0, + Cm=256.0, + Cx=1e-2, + Cy=1e-2, + Cz=1e-2, + max_iterations=30, + vel_name="corrected_velocity", + refl_field="reflectivity", + low_pass_filter=False, + mask_outside_opt=False, + mask_w_outside_opt=False, +) + + +def _twpice_vvad_grids(): + """The TWP-ICE grid pair with the VVAD constraint fields attached.""" + Grid0, Grid1 = _twpice_grids() + return pydda.constraints.make_constraint_from_vvad( + [Grid0, Grid1], vel_field="corrected_velocity" + ) + + +def test_vvad_constraint_on_real_grids(): + """The VVAD reconstruction on real radar data must be physically + plausible. This is the regression guard on the acceptance criteria: a + level fit from a narrow sector of azimuth yields deformation terms orders + of magnitude too large, which the reconstruction turns into winds of + hundreds of m/s.""" + Grids = _twpice_vvad_grids() + u_vad = Grids[0]["U_vvad"].values.squeeze() + v_vad = Grids[0]["V_vvad"].values.squeeze() + + # The constraint must actually cover a useful part of the domain. + assert 0.25 < np.isfinite(u_vad).mean() < 1.0 + assert np.nanmax(np.abs(u_vad)) < 60.0 + assert np.nanmax(np.abs(v_vad)) < 60.0 + + for Grid in Grids: + profiles = pydda.constraints.vvad_retrieval( + Grid, vel_field="corrected_velocity" + ) + ok = profiles["valid"].values + assert ok.sum() > 10 + # A mean divergence over a 150 km domain is O(1e-4) per second; 1e-3 + # would extrapolate to winds of over 100 m/s at the domain edge. + assert np.all(np.abs(profiles["DIV"].values[ok]) < 1e-3) + assert np.all(np.abs(profiles["U0"].values[ok]) < 60.0) + + +def test_vvad_weights_match_ivad_switch(): + """parameters.vad_weights must implement i_vad of Eq. (19) of Protat et + al. (2024) exactly: on where at most one radar contributes and the VVAD + has a value, off everywhere else.""" + Grids = _twpice_vvad_grids() + u_vad = Grids[0]["U_vvad"].values.squeeze() + v_vad = Grids[0]["V_vvad"].values.squeeze() + + _, parameters = pydda.retrieval.get_dd_wind_field( + deepcopy(Grids), Cvad=1.0, engine="scipy", **VVAD_KWARGS + ) + + n_rad = np.sum(parameters.weights, axis=0) + expected = np.logical_and( + n_rad <= 1, np.logical_and(np.isfinite(u_vad), np.isfinite(v_vad)) + ) + np.testing.assert_array_equal(parameters.vad_weights > 0, expected) + + # The TWP-ICE pair genuinely has both regimes, so this is a real test. + assert (n_rad > 1).any() and (n_rad <= 1).any() + assert not np.any(parameters.vad_weights[n_rad > 1] > 0) + + +def test_vvad_constraint_pulls_analysis_toward_vad(): + """Turning the VVAD constraint on must move the analysis toward the VVAD + winds where the constraint is active.""" + Grids = _twpice_vvad_grids() + u_vad = np.nan_to_num(Grids[0]["U_vvad"].values.squeeze()) + v_vad = np.nan_to_num(Grids[0]["V_vvad"].values.squeeze()) + + off, _ = pydda.retrieval.get_dd_wind_field( + deepcopy(Grids), Cvad=0.0, engine="scipy", **VVAD_KWARGS + ) + on, parameters = pydda.retrieval.get_dd_wind_field( + deepcopy(Grids), Cvad=1.0, engine="scipy", **VVAD_KWARGS + ) + + active = parameters.vad_weights > 0 + assert active.any() + + def departure(Grids_out): + u = Grids_out[0]["u"].values.squeeze() + v = Grids_out[0]["v"].values.squeeze() + return np.mean( + np.sqrt((u[active] - u_vad[active]) ** 2 + (v[active] - v_vad[active]) ** 2) + ) + + assert departure(on) < departure(off) + + +def test_vvad_requires_constraint_fields(): + """Cvad without the VVAD fields is a clear error rather than a crash.""" + Grid0, Grid1 = _twpice_grids() + with pytest.raises(ValueError, match="U_vvad"): + pydda.retrieval.get_dd_wind_field( + [Grid0, Grid1], Cvad=1.0, engine="scipy", **VVAD_KWARGS + ) + + +@pytest.mark.parametrize("engine", ["tensorflow", "auglag"]) +def test_vvad_unsupported_engines(engine): + """The VVAD constraint is implemented for the scipy and jax engines only.""" + Grids = _twpice_vvad_grids() + with pytest.raises(NotImplementedError, match="VVAD"): + pydda.retrieval.get_dd_wind_field(Grids, Cvad=1.0, engine=engine, **VVAD_KWARGS) + + +@pytest.mark.skipif(not JAX_AVAILABLE, reason="Jax not installed") +def test_vvad_constraint_jax(): + """The jax engine must reach the same solution as scipy with the VVAD + constraint active.""" + Grids = _twpice_vvad_grids() + scipy_out, _ = pydda.retrieval.get_dd_wind_field( + deepcopy(Grids), Cvad=1.0, engine="scipy", **VVAD_KWARGS + ) + jax_out, _ = pydda.retrieval.get_dd_wind_field( + deepcopy(Grids), Cvad=1.0, engine="jax", **VVAD_KWARGS + ) + + u_scipy = np.asarray(scipy_out[0]["u"].values.squeeze()) + u_jax = np.asarray(jax_out[0]["u"].values.squeeze()) + assert np.nanmean(np.abs(u_scipy - u_jax)) < 0.5 From 0d915b113b52b7579febc4fe08fc200b3614eb25 Mon Sep 17 00:00:00 2001 From: Robert Jackson Date: Tue, 29 Sep 2026 17:21:37 -0500 Subject: [PATCH 3/5] No vad data. --- pydda/constraints/__init__.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/pydda/constraints/__init__.py b/pydda/constraints/__init__.py index 95d46ae2..95167df3 100644 --- a/pydda/constraints/__init__.py +++ b/pydda/constraints/__init__.py @@ -34,6 +34,3 @@ from .model_data import make_constraint_from_era5 from .model_data import download_needed_era_data from .station_data import get_iem_obs -from .vad_data import make_constraint_from_vvad -from .vad_data import vvad_retrieval -from .vad_data import vvad_horizontal_wind From 4ed290168f36003ee1ea2160be60e38fc5cbda85 Mon Sep 17 00:00:00 2001 From: Robert Jackson Date: Tue, 29 Sep 2026 17:42:39 -0500 Subject: [PATCH 4/5] FIX: remove unavailable terrain import --- pydda/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pydda/__init__.py b/pydda/__init__.py index 34955719..0e6784fe 100644 --- a/pydda/__init__.py +++ b/pydda/__init__.py @@ -11,7 +11,6 @@ from . import tests from . import constraints from . import io -from . import terrain __version__ = "2.5.0" From 6572b817cd861eb2c054b8b7e8dae274f8a082a2 Mon Sep 17 00:00:00 2001 From: Robert Jackson Date: Tue, 29 Sep 2026 17:52:19 -0500 Subject: [PATCH 5/5] TEST: remove VVAD tests --- pydda/tests/test_cost_functions.py | 118 ---------------------- pydda/tests/test_initialization.py | 151 ----------------------------- pydda/tests/test_retrieval.py | 136 -------------------------- 3 files changed, 405 deletions(-) diff --git a/pydda/tests/test_cost_functions.py b/pydda/tests/test_cost_functions.py index 4631db0e..840ea717 100644 --- a/pydda/tests/test_cost_functions.py +++ b/pydda/tests/test_cost_functions.py @@ -1138,121 +1138,3 @@ def test_mass_continuity_gradient_upper_bc_tf(): assert np.all(grid_top[-1] == 0) assert np.all(echo_top[mask] == 0) np.testing.assert_allclose(echo_top[~mask], none_bc[~mask]) - - -from scipy.optimize import approx_fprime - - -def _make_vad_inputs(seed=0, shape=(5, 6, 7)): - """Random winds, VVAD winds, and a partially active i_vad mask.""" - rng = np.random.default_rng(seed) - u = rng.normal(size=shape) - v = rng.normal(size=shape) - w = rng.normal(size=shape) - u_vad = rng.normal(size=shape) - v_vad = rng.normal(size=shape) - vad_weights = (rng.random(shape) > 0.4).astype(float) - return u, v, w, u_vad, v_vad, vad_weights - - -def test_vad_cost(): - """The VVAD cost is the i_vad-weighted sum of squared horizontal - departures, and is zero when the analysis equals the VVAD.""" - u, v, w, u_vad, v_vad, vad_weights = _make_vad_inputs() - - expected = 2.5 * np.sum((np.square(u - u_vad) + np.square(v - v_vad)) * vad_weights) - actual = pydda.cost_functions.np.calculate_vad_cost( - u, v, vad_weights, u_vad, v_vad, coeff=2.5 - ) - np.testing.assert_allclose(actual, expected) - - # No cost where the analysis already matches the VVAD. - assert ( - pydda.cost_functions.np.calculate_vad_cost( - u_vad, v_vad, vad_weights, u_vad, v_vad, coeff=2.5 - ) - == 0.0 - ) - - # No cost at all where i_vad is zero everywhere. - assert ( - pydda.cost_functions.np.calculate_vad_cost( - u, v, np.zeros_like(vad_weights), u_vad, v_vad, coeff=2.5 - ) - == 0.0 - ) - - -def test_vad_gradient(): - """The analytic VVAD gradient matches finite differences, leaves w - untouched, and vanishes wherever i_vad is zero.""" - u, v, w, u_vad, v_vad, vad_weights = _make_vad_inputs() - shape = u.shape - n = int(np.prod(shape)) - coeff = 1.7 - - def J(x): - x = x.reshape((3,) + shape) - return pydda.cost_functions.np.calculate_vad_cost( - x[0], x[1], vad_weights, u_vad, v_vad, coeff=coeff - ) - - x0 = np.concatenate([u.ravel(), v.ravel(), w.ravel()]) - finite_diff = approx_fprime(x0, J, 1e-6) - analytic = pydda.cost_functions.np.calculate_vad_gradient( - u, v, vad_weights, u_vad, v_vad, coeff=coeff, upper_bc=0 - ) - np.testing.assert_allclose(analytic, finite_diff, atol=1e-4) - - # The VVAD constrains the horizontal wind only. - assert np.all(analytic[2 * n :] == 0) - - # Eq. (19) of Protat et al. (2024) switches the term off entirely where - # multi-Doppler information is available, so it must contribute no - # gradient at those points. - grad = analytic.reshape((3,) + shape) - assert np.all(grad[0][vad_weights == 0] == 0) - assert np.all(grad[1][vad_weights == 0] == 0) - - -@pytest.mark.skipif(not JAX_AVAILABLE, reason="Jax not installed") -def test_vad_cost_jax(): - """The jax VVAD cost and gradient agree with the numpy ones.""" - u, v, w, u_vad, v_vad, vad_weights = _make_vad_inputs() - - np_cost = pydda.cost_functions.np.calculate_vad_cost( - u, v, vad_weights, u_vad, v_vad, coeff=2.5 - ) - jax_cost = pydda.cost_functions.jax.calculate_vad_cost( - u, v, vad_weights, u_vad, v_vad, coeff=2.5 - ) - np.testing.assert_allclose(float(jax_cost), np_cost, rtol=1e-5) - - np_grad = pydda.cost_functions.np.calculate_vad_gradient( - u, v, vad_weights, u_vad, v_vad, coeff=2.5 - ) - jax_grad = pydda.cost_functions.jax.calculate_vad_gradient( - u, v, vad_weights, u_vad, v_vad, coeff=2.5 - ) - np.testing.assert_allclose(np.asarray(jax_grad), np_grad, atol=1e-5) - - -def test_vad_gradient_upper_bc(): - """The VVAD gradient honors the upper boundary condition modes even - though its own w gradient is identically zero.""" - u, v, w, u_vad, v_vad, vad_weights = _make_vad_inputs() - mask = np.zeros(u.shape, dtype=bool) - mask[-2:] = True - - for upper_bc in (0, 1, 2): - grad = pydda.cost_functions.np.calculate_vad_gradient( - u, - v, - vad_weights, - u_vad, - v_vad, - coeff=1.0, - upper_bc=upper_bc, - upper_bc_mask=mask, - ) - assert np.all(_w_component(grad, u.shape) == 0) diff --git a/pydda/tests/test_initialization.py b/pydda/tests/test_initialization.py index b6a28053..b2ed3b40 100644 --- a/pydda/tests/test_initialization.py +++ b/pydda/tests/test_initialization.py @@ -254,154 +254,3 @@ def make_grid(origin_lat, origin_lon): speed_true_north = np.sqrt(u_true_north**2 + v_true_north**2) speed_grid_relative = np.sqrt(u_grid_relative**2 + v_grid_relative**2) np.testing.assert_allclose(speed_true_north, speed_grid_relative, rtol=1e-4) - - -def _linear_wind_grid( - U0=7.0, - V0=-3.0, - Ux=1.2e-4, - Uy=3.0e-5, - Vy=-4.0e-4, - fall_speed=0.0, - shift_radar=False, -): - """ - Builds a Grid whose Doppler velocity is synthesized from an exactly linear - horizontal wind field, so the VVAD fit of Protat et al. (2024) must - recover the wind exactly. The flow is irrotational by construction - (Vx == Uy), matching the closure the reconstruction assumes. - - Returns the Grid together with the true u and v fields and the true - control variables. - """ - from pydda.constraints.vad_data import _geometry - - Grid = pyart.testing.make_empty_grid( - (12, 40, 40), ((500.0, 8000.0), (-40000.0, 40000.0), (-40000.0, 40000.0)) - ) - Grid.add_field("velocity", {"data": np.zeros((12, 40, 40)), "_FillValue": -9999.0}) - Grid = pydda.io.read_from_pyart_grid(Grid) - - if shift_radar: - # Move the radar off the grid origin so the reconstruction origin is - # genuinely exercised. - Grid["radar_latitude"].values[0] += 0.12 - Grid["radar_longitude"].values[0] += -0.10 - Grid = pydda.retrieval.angles.add_elevation_as_field(Grid) - - Vx = Uy # irrotational - dx, dy, rh, az, el = _geometry(Grid) - u_true = U0 + Ux * dx + Uy * dy - v_true = V0 + Vx * dx + Vy * dy - - el = np.asarray(el) - vr = ( - np.cos(el) * np.sin(az) * u_true - + np.cos(el) * np.cos(az) * v_true - - np.sin(el) * fall_speed - ) - Grid["velocity"] = Grid["velocity"].copy(data=np.expand_dims(vr, 0)) - - truth = { - "U0": U0, - "V0": V0, - "DIV": Ux + Vy, - "DET": Ux - Vy, - "DES": Uy + Vx, - "fall_speed": fall_speed, - } - return Grid, u_true, v_true, truth - - -@pytest.mark.parametrize("shift_radar", [False, True]) -def test_vvad_recovers_linear_wind(shift_radar): - """A wind field that exactly satisfies the VAD linearity assumption must - be recovered to machine precision at every level.""" - Grid, u_true, v_true, truth = _linear_wind_grid(shift_radar=shift_radar) - profiles = pydda.constraints.vvad_retrieval(Grid, vel_field="velocity") - - assert profiles["valid"].values.all() - ok = profiles["valid"].values - for name, expected in truth.items(): - np.testing.assert_allclose( - profiles[name].values[ok], expected, atol=1e-9, rtol=0 - ) - - -def test_vvad_separates_fall_speed_from_divergence(): - """A uniform fall speed must be recovered on its own rather than aliasing - into the divergence.""" - Grid, _, _, truth = _linear_wind_grid(fall_speed=2.5) - profiles = pydda.constraints.vvad_retrieval(Grid, vel_field="velocity") - - ok = profiles["valid"].values - np.testing.assert_allclose(profiles["fall_speed"].values[ok], 2.5, atol=1e-9) - np.testing.assert_allclose( - profiles["DIV"].values[ok], truth["DIV"], atol=1e-12, rtol=0 - ) - - -def test_vvad_horizontal_wind_round_trip(): - """Reconstructing from the fitted profiles must return the linear wind - field it was fitted to, and NaN beyond max_range.""" - Grid, u_true, v_true, _ = _linear_wind_grid() - profiles = pydda.constraints.vvad_retrieval(Grid, vel_field="velocity") - - u_vad, v_vad = pydda.constraints.vvad_horizontal_wind(profiles, Grid) - np.testing.assert_allclose(u_vad, u_true, atol=1e-7) - np.testing.assert_allclose(v_vad, v_true, atol=1e-7) - - rh = np.sqrt(Grid["point_x"].values ** 2 + Grid["point_y"].values ** 2) - u_clip, _ = pydda.constraints.vvad_horizontal_wind( - profiles, Grid, max_range=20000.0 - ) - assert np.all(np.isnan(u_clip[rh > 20000.0])) - assert np.all(np.isfinite(u_clip[rh <= 20000.0])) - - -def test_vvad_rejects_undersampled_levels(): - """Each acceptance criterion must reject a level it is meant to, and an - invalid level must reconstruct as NaN rather than as a wild extrapolation. - """ - Grid, _, _, _ = _linear_wind_grid() - - for kwargs in ( - {"min_points": 10**9}, - {"min_azimuth_coverage": 1.01}, - {"max_condition": 1e-6}, - ): - profiles = pydda.constraints.vvad_retrieval( - Grid, vel_field="velocity", **kwargs - ) - assert not profiles["valid"].values.any(), kwargs - assert np.all(np.isnan(profiles["U0"].values)), kwargs - - u_vad, v_vad = pydda.constraints.vvad_horizontal_wind(profiles, Grid) - assert np.all(np.isnan(u_vad)) and np.all(np.isnan(v_vad)), kwargs - - -def test_make_constraint_from_vvad_adds_fields(): - """The constraint builder writes U_vvad/V_vvad into every Grid, and the - two combine modes cover the same points by different means.""" - Grid, u_true, v_true, _ = _linear_wind_grid() - Grids = pydda.constraints.make_constraint_from_vvad( - [Grid, Grid.copy(deep=True)], vel_field="velocity" - ) - - for g in Grids: - assert "U_vvad" in g.variables and "V_vvad" in g.variables - assert g["U_vvad"].dims == ("time", "z", "y", "x") - # No W_vvad: the VVAD supplies horizontal components only. - assert "W_vvad" not in g.variables - - np.testing.assert_allclose(Grids[0]["U_vvad"].values.squeeze(), u_true, atol=1e-7) - - nearest = pydda.constraints.make_constraint_from_vvad( - [Grid.copy(deep=True)], vel_field="velocity", combine="nearest" - ) - np.testing.assert_allclose(nearest[0]["U_vvad"].values.squeeze(), u_true, atol=1e-7) - - with pytest.raises(ValueError): - pydda.constraints.make_constraint_from_vvad( - [Grid], vel_field="velocity", combine="not_a_mode" - ) diff --git a/pydda/tests/test_retrieval.py b/pydda/tests/test_retrieval.py index b0a470e5..a2c5c01e 100644 --- a/pydda/tests/test_retrieval.py +++ b/pydda/tests/test_retrieval.py @@ -515,139 +515,3 @@ def test_nested_retrieval(): quiverkey_loc="bottom_right", ) return fig - - -# Common arguments for the VVAD constraint tests. The filtering and masking are -# turned off so the retrieved winds can be compared point by point. -VVAD_KWARGS = dict( - Co=1.0, - Cm=256.0, - Cx=1e-2, - Cy=1e-2, - Cz=1e-2, - max_iterations=30, - vel_name="corrected_velocity", - refl_field="reflectivity", - low_pass_filter=False, - mask_outside_opt=False, - mask_w_outside_opt=False, -) - - -def _twpice_vvad_grids(): - """The TWP-ICE grid pair with the VVAD constraint fields attached.""" - Grid0, Grid1 = _twpice_grids() - return pydda.constraints.make_constraint_from_vvad( - [Grid0, Grid1], vel_field="corrected_velocity" - ) - - -def test_vvad_constraint_on_real_grids(): - """The VVAD reconstruction on real radar data must be physically - plausible. This is the regression guard on the acceptance criteria: a - level fit from a narrow sector of azimuth yields deformation terms orders - of magnitude too large, which the reconstruction turns into winds of - hundreds of m/s.""" - Grids = _twpice_vvad_grids() - u_vad = Grids[0]["U_vvad"].values.squeeze() - v_vad = Grids[0]["V_vvad"].values.squeeze() - - # The constraint must actually cover a useful part of the domain. - assert 0.25 < np.isfinite(u_vad).mean() < 1.0 - assert np.nanmax(np.abs(u_vad)) < 60.0 - assert np.nanmax(np.abs(v_vad)) < 60.0 - - for Grid in Grids: - profiles = pydda.constraints.vvad_retrieval( - Grid, vel_field="corrected_velocity" - ) - ok = profiles["valid"].values - assert ok.sum() > 10 - # A mean divergence over a 150 km domain is O(1e-4) per second; 1e-3 - # would extrapolate to winds of over 100 m/s at the domain edge. - assert np.all(np.abs(profiles["DIV"].values[ok]) < 1e-3) - assert np.all(np.abs(profiles["U0"].values[ok]) < 60.0) - - -def test_vvad_weights_match_ivad_switch(): - """parameters.vad_weights must implement i_vad of Eq. (19) of Protat et - al. (2024) exactly: on where at most one radar contributes and the VVAD - has a value, off everywhere else.""" - Grids = _twpice_vvad_grids() - u_vad = Grids[0]["U_vvad"].values.squeeze() - v_vad = Grids[0]["V_vvad"].values.squeeze() - - _, parameters = pydda.retrieval.get_dd_wind_field( - deepcopy(Grids), Cvad=1.0, engine="scipy", **VVAD_KWARGS - ) - - n_rad = np.sum(parameters.weights, axis=0) - expected = np.logical_and( - n_rad <= 1, np.logical_and(np.isfinite(u_vad), np.isfinite(v_vad)) - ) - np.testing.assert_array_equal(parameters.vad_weights > 0, expected) - - # The TWP-ICE pair genuinely has both regimes, so this is a real test. - assert (n_rad > 1).any() and (n_rad <= 1).any() - assert not np.any(parameters.vad_weights[n_rad > 1] > 0) - - -def test_vvad_constraint_pulls_analysis_toward_vad(): - """Turning the VVAD constraint on must move the analysis toward the VVAD - winds where the constraint is active.""" - Grids = _twpice_vvad_grids() - u_vad = np.nan_to_num(Grids[0]["U_vvad"].values.squeeze()) - v_vad = np.nan_to_num(Grids[0]["V_vvad"].values.squeeze()) - - off, _ = pydda.retrieval.get_dd_wind_field( - deepcopy(Grids), Cvad=0.0, engine="scipy", **VVAD_KWARGS - ) - on, parameters = pydda.retrieval.get_dd_wind_field( - deepcopy(Grids), Cvad=1.0, engine="scipy", **VVAD_KWARGS - ) - - active = parameters.vad_weights > 0 - assert active.any() - - def departure(Grids_out): - u = Grids_out[0]["u"].values.squeeze() - v = Grids_out[0]["v"].values.squeeze() - return np.mean( - np.sqrt((u[active] - u_vad[active]) ** 2 + (v[active] - v_vad[active]) ** 2) - ) - - assert departure(on) < departure(off) - - -def test_vvad_requires_constraint_fields(): - """Cvad without the VVAD fields is a clear error rather than a crash.""" - Grid0, Grid1 = _twpice_grids() - with pytest.raises(ValueError, match="U_vvad"): - pydda.retrieval.get_dd_wind_field( - [Grid0, Grid1], Cvad=1.0, engine="scipy", **VVAD_KWARGS - ) - - -@pytest.mark.parametrize("engine", ["tensorflow", "auglag"]) -def test_vvad_unsupported_engines(engine): - """The VVAD constraint is implemented for the scipy and jax engines only.""" - Grids = _twpice_vvad_grids() - with pytest.raises(NotImplementedError, match="VVAD"): - pydda.retrieval.get_dd_wind_field(Grids, Cvad=1.0, engine=engine, **VVAD_KWARGS) - - -@pytest.mark.skipif(not JAX_AVAILABLE, reason="Jax not installed") -def test_vvad_constraint_jax(): - """The jax engine must reach the same solution as scipy with the VVAD - constraint active.""" - Grids = _twpice_vvad_grids() - scipy_out, _ = pydda.retrieval.get_dd_wind_field( - deepcopy(Grids), Cvad=1.0, engine="scipy", **VVAD_KWARGS - ) - jax_out, _ = pydda.retrieval.get_dd_wind_field( - deepcopy(Grids), Cvad=1.0, engine="jax", **VVAD_KWARGS - ) - - u_scipy = np.asarray(scipy_out[0]["u"].values.squeeze()) - u_jax = np.asarray(jax_out[0]["u"].values.squeeze()) - assert np.nanmean(np.abs(u_scipy - u_jax)) < 0.5