ENH: Extend the echo top impermeability condition to all engines

doc/source/user_guide/optimizing_wind_retrieval.rst

@@ -430,3 +430,77 @@ more importance on horizontal winds compared to updraft velocities, then you may
 to tolerate more errors in the vertical velocity field so that finer details of the horizontal
 wind field can be generated. The above parameters are examples that apply to a 1 km
 resolution grid from two NEXRADs and vary for given radar configurations and storm coverages.
+
+------------------------------------
+Choosing an upper boundary condition
+------------------------------------
+
+PyDDA constrains the vertical velocity at the boundaries of the analysis domain by
+zeroing the gradient of the cost function with respect to :math:`w` there, so that
+:math:`w` is held at its first guess. At the surface this impermeability condition is
+always applied. At the top of the domain it is controlled by the :code:`upper_bc`
+keyword of :meth:`pydda.retrieval.get_dd_wind_field`, which takes one of three values:
+
+.. list-table::
+    :header-rows: 1
+    :widths: 20 80
+
+    * - :code:`upper_bc`
+      - Condition applied
+    * - :code:`0`
+      - No condition at the top of the domain.
+    * - :code:`1`
+      - :math:`w = 0` at the top vertical level of the grid.
+    * - :code:`2`
+      - :math:`w = 0` above the echo top, i.e. wherever no radar reports an
+        observation and the point is higher than *above* km in the grid's
+        vertical coordinate.
+
+The classic choice is :code:`upper_bc=1`, which imposes :math:`w = 0` at the top of the
+analysis domain. That is only physically defensible when the domain top is genuinely
+above the storm. If the grid is truncated through the middle of deep convection, the
+condition forces the mass continuity constraint to close the divergence profile at an
+arbitrary height and pushes a spurious compensating signal down into the levels you
+actually care about.
+
+Setting :code:`upper_bc=2` instead applies the impermeability condition at the *echo
+top*: every grid point at which none of the radars report a valid radial velocity is
+treated as being outside the storm, and :math:`w` is held at its first guess there.
+Because the first guess for :math:`w` is normally zero, this is equivalent to requiring
+that no mass crosses the top of the observed echo. The method is described in Thompson
+et al. (2026), https://doi.org/10.5194/egusphere-2026-4631.
+
+The :code:`above` keyword sets the lowest altitude, in km, at which the echo top
+condition may be applied. It exists so that clear air at low levels -- gaps between
+cells, the cone of silence, the far edge of the Dual Doppler lobes -- does not pin
+:math:`w` to zero close to the surface, which would suppress the very updrafts you are
+trying to retrieve. The default of 2 km is a reasonable starting point; raise it if
+your radars have poor low-level coverage.
+
+.. code-block:: python
+
+    grids_out, _ = pydda.retrieval.get_dd_wind_field([grid_kict, grid_ktlx],
+                                                Cm=256.0, Co=1e-2, Cx=1, Cy=1,
+                                                Cz=1, Cmod=1e-5, model_fields=["hrrr"],
+                                                refl_field='DBZ', wind_tol=0.5,
+                                                max_iterations=50, filter_window=15,
+                                                filter_order=3, engine='scipy',
+                                                upper_bc=2, above=2.0)
+
+Both :code:`upper_bc` and :code:`above` are supported by all of PyDDA's engines
+(:code:`"scipy"`, :code:`"jax"`, :code:`"tensorflow"` and :code:`"auglag"`). The set of
+points at which the condition is applied is calculated once at the start of the
+retrieval by :meth:`pydda.cost_functions.calculate_echo_top_mask` and is returned on the
+:code:`upper_bc_mask` attribute of the parameters object, so it can be inspected
+afterwards:
+
+.. code-block:: python
+
+    grids_out, parameters = pydda.retrieval.get_dd_wind_field(
+        [grid_kict, grid_ktlx], upper_bc=2, above=2.0, ...)
+    print("Fraction of the domain held impermeable: %.2f"
+          % parameters.upper_bc_mask.mean())
+
+.. note::
+    For backwards compatibility :code:`upper_bc=True` and :code:`upper_bc=False` are
+    still accepted and mean the same as :code:`1` and :code:`0`.

pydda/cost_functions/__init__.py

@@ -74,6 +74,7 @@
     calculate_fall_speed
     calculate_point_cost
     calculate_point_gradient
+    calculate_echo_top_mask
 """
 
 import cmweather
@@ -95,4 +96,5 @@
 from ._cost_functions_numpy import calculate_model_cost
 from ._cost_functions_numpy import calculate_model_gradient
 from ._cost_functions_numpy import calculate_point_cost, calculate_point_gradient
+from ._cost_functions_numpy import calculate_echo_top_mask
 from .cost_functions import J_function, grad_J, grad_jax, J_function_jax

pydda/cost_functions/_cost_functions_auglag.py

@@ -14,6 +14,25 @@
 from ._cost_functions_tensorflow import _tf_gradient
 
 
+def _apply_upper_bc(p_z1, upper_bc, upper_bc_mask=None):
+    """
+    Zeroes the vertical velocity gradient where the upper impermeability
+    condition applies. *upper_bc* is compared with == rather than *is* so that
+    both the integer modes and the legacy booleans are honored.
+    """
+    if upper_bc == 1:
+        p_z1 = tf.concat(
+            [
+                p_z1[:-1, :, :],
+                tf.zeros((1, p_z1.shape[1], p_z1.shape[2]), dtype=p_z1.dtype),
+            ],
+            axis=0,
+        )
+    elif upper_bc == 2 and upper_bc_mask is not None:
+        p_z1 = tf.where(upper_bc_mask, tf.zeros_like(p_z1), p_z1)
+    return p_z1
+
+
 def radial_velocity_function(winds, parameters):
     """
     Calculates the total cost function. This typically does not need to be
@@ -347,6 +366,7 @@ def grad_radial_velocity(winds, parameters):
         parameters.rmsVr,
         coeff=parameters.Co,
         upper_bc=parameters.upper_bc,
+        upper_bc_mask=parameters.upper_bc_mask,
     )
 
     return grad
@@ -372,6 +392,7 @@ def grad_mass_cont(winds, parameters):
         parameters.dy,
         parameters.dz,
         upper_bc=parameters.upper_bc,
+        upper_bc_mask=parameters.upper_bc_mask,
         coeff=1.0,
     )
 
@@ -393,10 +414,14 @@ def grad_smooth_cost(winds, parameters):
         winds[0],
         winds[1],
         winds[2],
+        parameters.dx,
+        parameters.dy,
+        parameters.dz,
         Cx=parameters.Cx,
         Cy=parameters.Cy,
         Cz=parameters.Cz,
         upper_bc=parameters.upper_bc,
+        upper_bc_mask=parameters.upper_bc_mask,
     )
 
     return grad
@@ -421,7 +446,6 @@ def grad_background_cost(winds, parameters):
         parameters.u_back,
         parameters.v_back,
         parameters.Cb,
-        upper_bc=parameters.upper_bc,
     )
 
     return grad
@@ -448,6 +472,8 @@ def grad_vertical_vorticity_cost(winds, parameters):
         parameters.Ut,
         parameters.Vt,
         coeff=parameters.Cv,
+        upper_bc=parameters.upper_bc,
+        upper_bc_mask=parameters.upper_bc_mask,
     )
 
     return grad
@@ -473,6 +499,8 @@ def grad_model_cost(winds, parameters):
         parameters.v_model,
         parameters.w_model,
         coeff=parameters.Cmod,
+        upper_bc=parameters.upper_bc,
+        upper_bc_mask=parameters.upper_bc_mask,
     )
 
     return grad
@@ -571,7 +599,18 @@ def calculate_radial_vel_cost_function(
 
 
 def calculate_grad_radial_vel(
-    vrs, els, azs, u, v, w, wts, weights, rmsVr, coeff=1.0, upper_bc=True
+    vrs,
+    els,
+    azs,
+    u,
+    v,
+    w,
+    wts,
+    weights,
+    rmsVr,
+    coeff=1.0,
+    upper_bc=1,
+    upper_bc_mask=None,
 ):
     """
     Calculates the gradient of the cost function due to difference of wind
@@ -600,8 +639,14 @@ def calculate_grad_radial_vel(
         Background velocity field name
     weights: n_radars x_bins x y_bins float array
         Data weights for each pair of radars
-    upper_bc: bool
-        Set to true to impose w=0 at top of domain.
+    upper_bc: int
+        Upper boundary (impermeability) condition. 0 disables it, 1 enforces
+        w = 0 at the top of the domain, and 2 enforces w = 0 above the echo
+        top as given by *upper_bc_mask*. The legacy booleans True and False
+        are equivalent to 1 and 0.
+    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`.
 
     Returns
     -------
@@ -635,10 +680,7 @@ def calculate_grad_radial_vel(
         [tf.zeros((1, u.shape[1], u.shape[2]), dtype=tf.float64), p_z1[1:, :, :]],
         axis=0,
     )
-    if upper_bc is True:
-        p_z1 = tf.concat(
-            [p_z1[:-1, :, :], tf.zeros((1, u.shape[1], u.shape[2]))], axis=0
-        )
+    p_z1 = _apply_upper_bc(p_z1, upper_bc, upper_bc_mask)
     y = tf.stack((p_x1, p_y1, p_z1), axis=0)
     return tf.reshape(y, (3 * np.prod(u.shape),))
 
@@ -706,7 +748,7 @@ def calculate_smoothness_cost(u, v, w, dx, dy, dz, Cx=1e-5, Cy=1e-5, Cz=1e-5):
 
 
 def calculate_smoothness_gradient(
-    u, v, w, dx, dy, dz, Cx=1e-5, Cy=1e-5, Cz=1e-5, upper_bc=True
+    u, v, w, dx, dy, dz, Cx=1e-5, Cy=1e-5, Cz=1e-5, upper_bc=1, upper_bc_mask=None
 ):
     """
     Calculates the gradient of the smoothness cost function
@@ -755,11 +797,7 @@ def calculate_smoothness_gradient(
 
     # Impermeability condition
     p_z1 = tf.concat([tf.zeros((1, u.shape[1], u.shape[2])), p_z1[1:, :, :]], axis=0)
-    if upper_bc is True:
-        p_z1 = tf.concat(
-            [p_z1[:-1, :, :], tf.zeros((1, u.shape[1], u.shape[2]), dtype=tf.float64)],
-            axis=0,
-        )
+    p_z1 = _apply_upper_bc(p_z1, upper_bc, upper_bc_mask)
     y = tf.stack((p_x1, p_y1, p_z1), axis=0)
     return tf.reshape(y, (3 * np.prod(u.shape),))
 
@@ -928,7 +966,7 @@ def calculate_mass_continuity(u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1):
 
 
 def calculate_mass_continuity_gradient(
-    u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1, upper_bc=True
+    u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1, upper_bc=1, upper_bc_mask=None
 ):
     """
     Calculates the gradient of mass continuity cost function. This is done by
@@ -977,11 +1015,7 @@ def calculate_mass_continuity_gradient(
 
     # Impermeability condition
     p_z1 = tf.concat([tf.zeros((1, u.shape[1], u.shape[2])), p_z1[1:, :, :]], axis=0)
-    if upper_bc is True:
-        p_z1 = tf.concat(
-            [p_z1[:-1, :, :], tf.zeros((1, u.shape[1], u.shape[2]), dtype=tf.float64)],
-            axis=0,
-        )
+    p_z1 = _apply_upper_bc(p_z1, upper_bc, upper_bc_mask)
     y = tf.stack((p_x1, p_y1, p_z1), axis=0)
     return tf.reshape(y, (3 * np.prod(u.shape),))
 
@@ -1249,7 +1283,7 @@ def calculate_vertical_vorticity_cost(u, v, w, dx, dy, dz, Ut, Vt, coeff=1e-5):
 
 
 def calculate_vertical_vorticity_gradient(
-    u, v, w, dx, dy, dz, Ut, Vt, coeff=1e-5, upper_bc=True
+    u, v, w, dx, dy, dz, Ut, Vt, coeff=1e-5, upper_bc=1, upper_bc_mask=None
 ):
     """
     Calculates the gradient of the cost function due to deviance from vertical
@@ -1276,8 +1310,14 @@ def calculate_vertical_vorticity_gradient(
         V component of storm motion
     coeff: float
         Weighting coefficient
-    upper_bc: bool
-        If true, impose w=0 at top of domain as a boundary condition.
+    upper_bc: int
+        Upper boundary (impermeability) condition. 0 disables it, 1 enforces
+        w = 0 at the top of the domain, and 2 enforces w = 0 above the echo
+        top as given by *upper_bc_mask*. The legacy booleans True and False
+        are equivalent to 1 and 0.
+    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`.
 
     Returns
     -------
@@ -1310,10 +1350,7 @@ def calculate_vertical_vorticity_gradient(
 
     # Impermeability condition
     p_z1 = tf.concat([tf.zeros((1, u.shape[1], u.shape[2])), p_z1[1:, :, :]], axis=0)
-    if upper_bc is True:
-        p_z1 = tf.concat(
-            [p_z1[:-1, :, :], tf.zeros((1, u.shape[1], u.shape[2]))], axis=0
-        )
+    p_z1 = _apply_upper_bc(p_z1, upper_bc, upper_bc_mask)
     y = tf.stack((p_x1, p_y1, p_z1), axis=0)
     return tf.reshape(y, (3 * np.prod(u.shape),))
 
@@ -1362,7 +1399,16 @@ def calculate_model_cost(u, v, w, weights, u_model, v_model, w_model, coeff=1.0)
 
 
 def calculate_model_gradient(
-    u, v, w, weights, u_model, v_model, w_model, coeff=1.0, upper_bc=True
+    u,
+    v,
+    w,
+    weights,
+    u_model,
+    v_model,
+    w_model,
+    coeff=1.0,
+    upper_bc=1,
+    upper_bc_mask=None,
 ):
     """
     Calculates the cost function for the model constraint.
@@ -1390,8 +1436,14 @@ def calculate_model_gradient(
         Vertical wind field from model
     coeff: float
         Weight of background constraint to total cost function
-    upper_bc: bool
-        If true, impose w=0 at top of domain as boundary condition.
+    upper_bc: int
+        Upper boundary (impermeability) condition. 0 disables it, 1 enforces
+        w = 0 at the top of the domain, and 2 enforces w = 0 above the echo
+        top as given by *upper_bc_mask*. The legacy booleans True and False
+        are equivalent to 1 and 0.
+    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`.
 
     Returns
     -------
@@ -1415,9 +1467,6 @@ def calculate_model_gradient(
 
     # Impermeability condition
     p_z1 = tf.concat([tf.zeros((1, u.shape[1], u.shape[2])), p_z1[1:, :, :]], axis=0)
-    if upper_bc is True:
-        p_z1 = tf.concat(
-            [p_z1[:-1, :, :], tf.zeros((1, u.shape[1], u.shape[2]))], axis=0
-        )
+    p_z1 = _apply_upper_bc(p_z1, upper_bc, upper_bc_mask)
     y = tf.stack((p_x1, p_y1, p_z1), axis=0)
     return tf.reshape(y, (3 * np.prod(u.shape),))