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/constraints/__init__.py b/pydda/constraints/__init__.py index 5d6db7b2..95167df3 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 """ 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 45ecbf15..3e62729f 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,22 @@ 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 +381,19 @@ 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 +430,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 +474,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"]) @@ -466,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. @@ -492,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 @@ -509,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 @@ -547,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 @@ -564,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. @@ -740,7 +904,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) @@ -952,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 2764be4e..840ea717 100644 --- a/pydda/tests/test_cost_functions.py +++ b/pydda/tests/test_cost_functions.py @@ -268,6 +268,50 @@ 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"""