diff --git a/doc/source/user_guide/optimizing_wind_retrieval.rst b/doc/source/user_guide/optimizing_wind_retrieval.rst index f73e360a..8a41b84a 100644 --- a/doc/source/user_guide/optimizing_wind_retrieval.rst +++ b/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`. diff --git a/pydda/cost_functions/__init__.py b/pydda/cost_functions/__init__.py index 79174927..54c392df 100644 --- a/pydda/cost_functions/__init__.py +++ b/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 diff --git a/pydda/cost_functions/_cost_functions_auglag.py b/pydda/cost_functions/_cost_functions_auglag.py index b4c54e9d..dc0f5ead 100644 --- a/pydda/cost_functions/_cost_functions_auglag.py +++ b/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),)) diff --git a/pydda/cost_functions/_cost_functions_jax.py b/pydda/cost_functions/_cost_functions_jax.py index c9495abe..b451a7c2 100644 --- a/pydda/cost_functions/_cost_functions_jax.py +++ b/pydda/cost_functions/_cost_functions_jax.py @@ -10,6 +10,19 @@ JAX_AVAILABLE = False +def _apply_upper_bc(grad_w, 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: + grad_w = grad_w.at[-1, :, :].set(0) + elif upper_bc == 2 and upper_bc_mask is not None: + grad_w = jnp.where(upper_bc_mask, 0.0, grad_w) + return grad_w + + def calculate_radial_vel_cost_function( vrs, azs, els, u, v, w, wts, rmsVr, weights, coeff=1.0 ): @@ -78,7 +91,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 @@ -108,6 +132,15 @@ def calculate_grad_radial_vel( weights: n_radars x_bins x y_bins float array Data weights for each pair of radars + 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 ------- y: 1-D float array @@ -137,8 +170,7 @@ def calculate_grad_radial_vel( # Impermeability condition p_z1 = p_z1.at[0, :, :].set(0) - if upper_bc is True: - p_z1 = p_z1.at[-1, :, :].set(0) + p_z1 = _apply_upper_bc(p_z1, upper_bc, upper_bc_mask) y = jnp.stack((p_x1, p_y1, p_z1), axis=0) return y.flatten() @@ -218,7 +250,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 @@ -248,6 +280,15 @@ def calculate_smoothness_gradient( Cz: float Constant controlling smoothness in z-direction + 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 ------- y: float array @@ -260,8 +301,7 @@ def calculate_smoothness_gradient( # Impermeability condition grad_w = grad_w.at[0, :, :].set(0) - if upper_bc is True: - grad_w = grad_w.at[-1, :, :].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() @@ -426,7 +466,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 @@ -455,6 +495,15 @@ def calculate_mass_continuity_gradient( anel: int = 1 use anelastic approximation, 0=don't + 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 ------- y: float array @@ -472,8 +521,7 @@ def calculate_mass_continuity_gradient( # Impermeability condition grad_w = grad_w.at[0, :, :].set(0) - if upper_bc is True: - grad_w = grad_w.at[-1, :, :].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() @@ -622,7 +670,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 @@ -650,6 +698,15 @@ def calculate_vertical_vorticity_gradient( coeff: float Weighting coefficient + 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 ------- Jv: 1D float array @@ -675,9 +732,8 @@ def calculate_vertical_vorticity_gradient( ) u_grad, v_grad, w_grad, _, _, _, _, _, _ = fun_vjp(1.0) # Impermeability condition - w_grad.at[0, :, :].set(0) - if upper_bc is True: - w_grad.at[-1, :, :].set(0) + w_grad = w_grad.at[0, :, :].set(0) + 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 9baf9c07..45ecbf15 100644 --- a/pydda/cost_functions/_cost_functions_numpy.py +++ b/pydda/cost_functions/_cost_functions_numpy.py @@ -7,6 +7,51 @@ laplace_filter = np.asarray([1, -2, 1], dtype=np.float64) +def calculate_echo_top_mask(vrs, z, above=2.0): + """ + Finds the grid points that lie above the echo top, i.e. the points above + *above* km at which none of the radars report a valid radial velocity. + + Holding w fixed at these points imposes an impermeability condition at the + echo top rather than at the (arbitrary) top of the analysis domain. See + Thompson et al. (2026), https://doi.org/10.5194/egusphere-2026-4631. + + Parameters + ---------- + vrs: list of float arrays + List of radial velocities from each radar. Points without a valid + observation are expected to be filled with a large negative number + (-9999.0 in PyDDA's retrievals). + z: float array + Heights of the grid points in m. This must broadcast against the + members of *vrs*, i.e. the 3D *point_z* field of the Grid. + above: float + Minimum height of the impermeable layer in km. The condition is never + applied below this level, so that clear air near the surface does not + pin w to its first guess. + + Returns + ------- + mask: 3D bool array + True wherever the impermeability condition is to be applied. + """ + n_obs = np.sum(np.stack([np.asarray(x) for x in vrs]) > -1000, axis=0) + return np.logical_and(np.asarray(z) > above * 1000.0, n_obs == 0) + + +def _apply_upper_bc(grad_w, 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: + grad_w[-1, :, :] = 0 + elif upper_bc == 2 and upper_bc_mask is not None: + grad_w[upper_bc_mask] = 0 + return grad_w + + def calculate_radial_vel_cost_function( vrs, azs, els, u, v, w, wts, rmsVr, weights, coeff=1.0, parallel=False ): @@ -90,7 +135,8 @@ def calculate_grad_radial_vel( weights, rmsVr, coeff=1.0, - upper_bc=True, + upper_bc=1, + upper_bc_mask=None, parallel=False, ): """ @@ -118,6 +164,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: 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 ------- y: 1-D float array @@ -176,8 +230,7 @@ def calculate_grad_radial_vel( # Impermeability condition p_z1[0, :, :] = 0 - if upper_bc is True: - p_z1[-1, :, :] = 0 + p_z1 = _apply_upper_bc(p_z1, upper_bc, upper_bc_mask) y = np.stack((p_x1, p_y1, p_z1), axis=0) return y.flatten() @@ -248,7 +301,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 @@ -269,6 +322,14 @@ def calculate_smoothness_gradient( Constant controlling smoothness in y-direction Cz: float Constant controlling smoothness in z-direction + 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 ------- y: float array @@ -295,8 +356,7 @@ def calculate_smoothness_gradient( # Impermeability condition grad_w[0, :, :] = 0 - if upper_bc is True: - grad_w[-1, :, :] = 0 + grad_w = _apply_upper_bc(grad_w, upper_bc, upper_bc_mask) y = np.stack([grad_u, grad_v, grad_w], axis=0) @@ -453,7 +513,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, vrs=0, coeff=1500.0, anel=1, upper_bc=True, above=2.0 + 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 @@ -479,6 +539,14 @@ def calculate_mass_continuity_gradient( Constant controlling contribution of mass continuity to cost function anel: int = 1 use anelastic approximation, 0=don't + 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 ------- y: float array @@ -500,15 +568,10 @@ def calculate_mass_continuity_gradient( grad_v = -np.gradient(div, dy, axis=1) * coeff grad_w = -np.gradient(div, dz, axis=0) * coeff - # Impermeability conditions - grad_w[0, :, :] = 0 # surface is impermeable - if upper_bc == 1: # is True: # impermeable at the grid top - grad_w[-1, :, :] = 0 - if upper_bc == 2: # impermeable at cloud top - N = np.sum(np.array(vrs) > -1000, axis=0) - z_mask = z > above * 1000 - n_mask = N == 0 - grad_w[z_mask & n_mask] = 0 + # Impermeability condition + 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() @@ -694,7 +757,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 @@ -720,6 +783,14 @@ def calculate_vertical_vorticity_gradient( V component of storm motion coeff: float Weighting coefficient + 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 ------- Jv: 1D float array @@ -794,8 +865,7 @@ def calculate_vertical_vorticity_gradient( # Impermeability condition w_grad[0, :, :] = 0 - if upper_bc is True: - w_grad[-1, :, :] = 0 + 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_tensorflow.py b/pydda/cost_functions/_cost_functions_tensorflow.py index 76668aed..9416ae32 100644 --- a/pydda/cost_functions/_cost_functions_tensorflow.py +++ b/pydda/cost_functions/_cost_functions_tensorflow.py @@ -8,6 +8,25 @@ TENSORFLOW_AVAILABLE = False +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 calculate_radial_vel_cost_function( vrs, azs, els, u, v, w, wts, rmsVr, weights, coeff=1.0 ): @@ -77,7 +96,19 @@ 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, lower_bc=True + vrs, + els, + azs, + u, + v, + w, + wts, + weights, + rmsVr, + coeff=1.0, + upper_bc=1, + upper_bc_mask=None, + lower_bc=True, ): """ Calculates the gradient of the cost function due to difference of wind @@ -106,8 +137,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`. lower_bc: bool Set to true to impose w=0 at bottom of domain. Returns @@ -154,11 +191,7 @@ def calculate_grad_radial_vel( [tf.zeros((1, u.shape[1], u.shape[2]), dtype=tf.float32), 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.float32)], - 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),)) @@ -232,7 +265,18 @@ 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, lower_bc=True + u, + v, + w, + dx, + dy, + dz, + Cx=1e-5, + Cy=1e-5, + Cz=1e-5, + upper_bc=1, + upper_bc_mask=None, + lower_bc=True, ): """ Calculates the gradient of the smoothness cost function @@ -255,8 +299,14 @@ def calculate_smoothness_gradient( Constant controlling smoothness in y-direction Cz: float Constant controlling smoothness in z-direction - 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`. lower_bc: bool Set to true to impose w=0 at bottom of domain. Returns @@ -283,11 +333,7 @@ def calculate_smoothness_gradient( [tf.zeros((1, u.shape[1], u.shape[2]), dtype=tf.float32), 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.float32)], - 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),)) @@ -526,7 +572,18 @@ 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, lower_bc=True + u, + v, + w, + z, + dx, + dy, + dz, + coeff=1500.0, + anel=1, + upper_bc=1, + upper_bc_mask=None, + lower_bc=True, ): """ Calculates the gradient of mass continuity cost function. This is done by @@ -554,8 +611,14 @@ def calculate_mass_continuity_gradient( Constant controlling contribution of mass continuity to cost function anel: int = 1 use anelastic approximation, 0=don't - 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 ------- @@ -586,11 +649,7 @@ def calculate_mass_continuity_gradient( [tf.zeros((1, u.shape[1], u.shape[2]), dtype=tf.float32), 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.float32)], - 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),)) @@ -661,8 +720,6 @@ def calculate_background_gradient(u, v, weights, u_back, v_back, Cb=0.01): Meridional winds vs height from sounding Cb: float Weight of background constraint to total cost function - upper_bc: bool - Set to true to impose w=0 at top of domain. Returns ------- @@ -782,7 +839,18 @@ def calculate_vertical_vorticity_cost(u, v, w, dx, dy, dz, Ut, Vt, coeff=1): # Using Jax version of function def calculate_vertical_vorticity_gradient( - u, v, w, dx, dy, dz, Ut, Vt, coeff=1e-5, upper_bc=True, lower_bc=True + u, + v, + w, + dx, + dy, + dz, + Ut, + Vt, + coeff=1e-5, + upper_bc=1, + upper_bc_mask=None, + lower_bc=True, ): """ Calculates the gradient of the cost function due to deviance from vertical @@ -809,8 +877,14 @@ def calculate_vertical_vorticity_gradient( V component of storm motion coeff: float Weighting coefficient - 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 ------- @@ -848,11 +922,7 @@ def calculate_vertical_vorticity_gradient( [tf.zeros((1, u.shape[1], u.shape[2]), dtype=tf.float32), 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.float32)], - 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),)) @@ -901,7 +971,17 @@ 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, lower_bc=True + u, + v, + w, + weights, + u_model, + v_model, + w_model, + coeff=1.0, + upper_bc=1, + upper_bc_mask=None, + lower_bc=True, ): """ Calculates the cost function for the model constraint. @@ -929,6 +1009,14 @@ def calculate_model_gradient( Vertical wind field from model coeff: float Weight of background constraint to total cost function + 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 ------- @@ -956,10 +1044,6 @@ def calculate_model_gradient( [tf.zeros((1, u.shape[1], u.shape[2]), dtype=tf.float32), 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.float32)], - 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),)) diff --git a/pydda/cost_functions/cost_functions.py b/pydda/cost_functions/cost_functions.py index 7787ffaf..52618238 100644 --- a/pydda/cost_functions/cost_functions.py +++ b/pydda/cost_functions/cost_functions.py @@ -362,6 +362,7 @@ def grad_J(winds, parameters): parameters.rmsVr, coeff=parameters.Co, upper_bc=parameters.upper_bc, + upper_bc_mask=parameters.upper_bc_mask, lower_bc=parameters.lower_bc, ) @@ -376,6 +377,7 @@ def grad_J(winds, parameters): parameters.dz, coeff=parameters.Cm, upper_bc=parameters.upper_bc, + upper_bc_mask=parameters.upper_bc_mask, lower_bc=parameters.lower_bc, ) @@ -391,6 +393,7 @@ def grad_J(winds, parameters): Cy=parameters.Cy, Cz=parameters.Cz, upper_bc=parameters.upper_bc, + upper_bc_mask=parameters.upper_bc_mask, ) if parameters.Cb > 0: @@ -415,6 +418,7 @@ def grad_J(winds, parameters): parameters.Vt, coeff=parameters.Cv, upper_bc=parameters.upper_bc, + upper_bc_mask=parameters.upper_bc_mask, lower_bc=parameters.lower_bc, ).numpy() @@ -428,6 +432,8 @@ def grad_J(winds, parameters): parameters.v_model, parameters.w_model, coeff=parameters.Cmod, + upper_bc=parameters.upper_bc, + upper_bc_mask=parameters.upper_bc_mask, ) if parameters.Cpoint > 0: @@ -440,7 +446,6 @@ def grad_J(winds, parameters): parameters.point_list, Cp=parameters.Cpoint, roi=parameters.roi, - upper_bc=parameters.upper_bc, ) if parameters.const_boundary_cond is True: grad = tf.reshape( @@ -529,6 +534,7 @@ def grad_J(winds, parameters): parameters.rmsVr, parameters.Co, parameters.upper_bc, + parameters.upper_bc_mask, True, ) ) @@ -543,11 +549,10 @@ def grad_J(winds, parameters): parameters.dx, parameters.dy, parameters.dz, - parameters.vrs, parameters.Cm, 1, parameters.upper_bc, - above=parameters.above, + parameters.upper_bc_mask, ) ) if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: @@ -564,6 +569,7 @@ def grad_J(winds, parameters): parameters.Cy, parameters.Cz, parameters.upper_bc, + parameters.upper_bc_mask, ) ) if parameters.Cb > 0: @@ -593,6 +599,7 @@ def grad_J(winds, parameters): parameters.Vt, parameters.Cv, parameters.upper_bc, + parameters.upper_bc_mask, ) ) if parameters.Cmod > 0: @@ -637,6 +644,7 @@ def grad_J(winds, parameters): parameters.rmsVr, coeff=parameters.Co, upper_bc=parameters.upper_bc, + upper_bc_mask=parameters.upper_bc_mask, ) if parameters.Cm > 0: @@ -650,6 +658,7 @@ def grad_J(winds, parameters): parameters.dz, coeff=parameters.Cm, upper_bc=parameters.upper_bc, + upper_bc_mask=parameters.upper_bc_mask, ) if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: @@ -664,6 +673,7 @@ def grad_J(winds, parameters): Cy=parameters.Cy, Cz=parameters.Cz, upper_bc=parameters.upper_bc, + upper_bc_mask=parameters.upper_bc_mask, ) if parameters.Cb > 0: @@ -689,6 +699,7 @@ def grad_J(winds, parameters): parameters.Vt, coeff=parameters.Cv, upper_bc=parameters.upper_bc, + upper_bc_mask=parameters.upper_bc_mask, ) if parameters.Cmod > 0: @@ -713,7 +724,6 @@ def grad_J(winds, parameters): parameters.point_list, Cp=parameters.Cpoint, roi=parameters.roi, - upper_bc=parameters.upper_bc, ) # Let's see if we need to enforce strong boundary conditions @@ -895,6 +905,7 @@ def grad_jax(winds, parameters): parameters.rmsVr, coeff=parameters.Co, upper_bc=parameters.upper_bc, + upper_bc_mask=parameters.upper_bc_mask, ) if parameters.Cm > 0: @@ -908,6 +919,7 @@ def grad_jax(winds, parameters): parameters.dz, coeff=parameters.Cm, upper_bc=parameters.upper_bc, + upper_bc_mask=parameters.upper_bc_mask, ) if parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0: @@ -922,6 +934,7 @@ def grad_jax(winds, parameters): Cy=parameters.Cy, Cz=parameters.Cz, upper_bc=parameters.upper_bc, + upper_bc_mask=parameters.upper_bc_mask, ) if parameters.Cb > 0: @@ -947,6 +960,7 @@ def grad_jax(winds, parameters): parameters.Vt, coeff=parameters.Cv, upper_bc=parameters.upper_bc, + upper_bc_mask=parameters.upper_bc_mask, ).numpy() if parameters.Cmod > 0: diff --git a/pydda/retrieval/wind_retrieve.py b/pydda/retrieval/wind_retrieve.py index e757a8df..2ce10aad 100644 --- a/pydda/retrieval/wind_retrieve.py +++ b/pydda/retrieval/wind_retrieve.py @@ -66,6 +66,7 @@ def _apply_low_pass_filter( J_function, grad_J, calculate_fall_speed, + calculate_echo_top_mask, grad_jax, J_function_jax, ) @@ -159,12 +160,22 @@ class DDParameters(object): Cartesian coordinates. roi: float The radius of influence of each point observation in m. - above: float - The altitude below which the cloud top impermeability would not apply (minimum impermeable height) upper_bc: int - 0 to not enforce impermeability at top of domain - 1 to enforce w=0 at top of domain (impermeability condition), - 2 to enforce w=0 at cloud top (or "above" if higher) (impermeability condition), + The upper boundary (impermeability) condition: + + * 0 to not enforce impermeability at the top of the domain, + * 1 to enforce w=0 at the top of the domain, + * 2 to enforce w=0 above the echo top (or *above*, whichever is higher). + + The legacy booleans True and False are equivalent to 1 and 0. + above: float + The minimum height of the echo top impermeability condition in km. + The condition given by *upper_bc* = 2 is never applied below this + level. + 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`. + This is calculated by PyDDA at the start of the retrieval. """ def __init__(self): @@ -203,8 +214,10 @@ def __init__(self): self.Cpoint = 0.0 self.Ut = 0.0 self.Vt = 0.0 - self.upper_bc = True + self.upper_bc = 1 self.lower_bc = True + self.above = 2.0 + self.upper_bc_mask = None self.roi = 1000.0 self.above = 2.0 self.frz = 4500.0 @@ -259,10 +272,10 @@ def _get_dd_wind_field_scipy( min_bca=30.0, max_bca=150.0, upper_bc=1, + above=2.0, model_fields=None, output_cost_functions=True, roi=1000.0, - above=2.0, wind_tol=0.1, tolerance=1e-8, const_boundary_cond=False, @@ -603,6 +616,11 @@ def _get_dd_wind_field_scipy( parameters.Cpoint = Cpoint parameters.roi = roi parameters.upper_bc = upper_bc + parameters.above = above + if upper_bc == 2: + parameters.upper_bc_mask = calculate_echo_top_mask( + parameters.vrs, parameters.z, above=above + ) parameters.points = points parameters.point_list = points parameters.parallel = parallel @@ -720,6 +738,10 @@ def loss_and_gradient(x): parameters.z = tf.constant(Grids[0]["point_z"].values, dtype=tf.float32) parameters.x = tf.constant(Grids[0]["point_x"].values, dtype=tf.float32) parameters.y = tf.constant(Grids[0]["point_y"].values, dtype=tf.float32) + if parameters.upper_bc_mask is not None: + parameters.upper_bc_mask = tf.constant( + parameters.upper_bc_mask, dtype=tf.bool + ) bounds = [(-x, x) for x in max_wind_mag * np.ones(winds.shape, dtype="float32")] winds = winds.astype("float32") winds, mult, AL_Filter, funcalls = auglag(winds, parameters, bounds) @@ -847,7 +869,8 @@ def _get_dd_wind_field_tensorflow( leise_nstep=1, min_bca=30.0, max_bca=150.0, - upper_bc=True, + upper_bc=1, + above=2.0, model_fields=None, output_cost_functions=True, roi=1000.0, @@ -881,6 +904,7 @@ def _get_dd_wind_field_tensorflow( parameters.Vt = Vt parameters.upper_bc = upper_bc parameters.lower_bc = lower_bc + parameters.above = above parameters.engine = "tensorflow" parameters.const_boundary_cond = const_boundary_cond @@ -1133,6 +1157,16 @@ def _get_dd_wind_field_tensorflow( else: parameters.bg_weights = weights_bg + if upper_bc == 2: + parameters.upper_bc_mask = tf.constant( + calculate_echo_top_mask( + [x.filled(-9999) for x in parameters.vrs], + Grids[0]["point_z"].values, + above=above, + ), + dtype=tf.bool, + ) + parameters.vrs = [ tf.constant(x.filled(-9999), dtype=tf.float32) for x in parameters.vrs ] @@ -1466,9 +1500,23 @@ def get_dd_wind_field( max_bca: float Maximum beam crossing angle in degrees between two radars. 150.0 is the typical value used in many publications. - upper_bc: bool - Set this to true to enforce w = 0 at the top of the atmosphere. This is - commonly called the impermeability condition. + upper_bc: int + The upper boundary condition to apply to the vertical velocity, which is + commonly called the impermeability condition: + + * 0 -- do not enforce impermeability at the top of the domain. + * 1 -- enforce :math:`w = 0` at the top of the analysis domain. + * 2 -- enforce :math:`w = 0` above the echo top, i.e. at every point + above *above* km at which none of the radars report a valid + observation. See Thompson et al. (2026), + https://doi.org/10.5194/egusphere-2026-4631. + + The legacy booleans True and False are equivalent to 1 and 0. + above: float + The minimum height of the echo top impermeability condition in km. + :code:`upper_bc=2` is never applied below this level, so that clear air + near the surface does not pin the vertical velocity to its first guess. + This has no effect for the other values of *upper_bc*. model_fields: list of strings The list of fields in the first grid in Grids that contain the custom data interpolated to the Grid's grid specification. Helper functions diff --git a/pydda/tests/test_cost_functions.py b/pydda/tests/test_cost_functions.py index c234c720..2764be4e 100644 --- a/pydda/tests/test_cost_functions.py +++ b/pydda/tests/test_cost_functions.py @@ -873,3 +873,224 @@ def test_model_cost_tf(): u, v, w, weights, u - 1, v - 1, w ) assert cost2 > cost1 + + +def _make_upper_bc_inputs(): + """A synthetic wind field and a single radar whose echo top is at level 5.""" + nz, ny, nx = 10, 8, 8 + rng = np.random.default_rng(42) + u = rng.random((nz, ny, nx)) + v = rng.random((nz, ny, nx)) + w = rng.random((nz, ny, nx)) + z = np.tile(np.linspace(0.0, 9000.0, nz)[:, None, None], (1, ny, nx)) + + # One radar reporting valid velocities only in the lowest six levels of one + # quadrant of the domain. Everything else is filled the way PyDDA fills + # points without an observation. + vr = np.full((nz, ny, nx), -9999.0) + vr[:6, :4, :4] = 1.0 + return u, v, w, z, [vr] + + +def _w_component(grad, shape): + """Pull the w block out of a flattened (3, nz, ny, nx) gradient.""" + return np.asarray(grad).reshape((3,) + shape)[2] + + +def test_calculate_echo_top_mask(): + u, v, w, z, vrs = _make_upper_bc_inputs() + mask = pydda.cost_functions.calculate_echo_top_mask(vrs, z, above=2.0) + + # The condition never applies below the "above" level... + assert not mask[z <= 2000.0].any() + # ...nor where the radar sees an echo... + assert not mask[np.logical_and(vrs[0] > -1000, z > 2000.0)].any() + # ...but it does apply in the echo-free air aloft. + assert mask[6:, 6:, 6:].all() + + # Raising "above" can only shrink the masked volume. + higher = pydda.cost_functions.calculate_echo_top_mask(vrs, z, above=5.0) + assert higher.sum() < mask.sum() + assert not np.logical_and(higher, ~mask).any() + + +def test_mass_continuity_gradient_upper_bc(): + """upper_bc selects between no condition, a grid top condition and an + echo top condition.""" + u, v, w, z, vrs = _make_upper_bc_inputs() + shape = u.shape + mask = pydda.cost_functions.calculate_echo_top_mask(vrs, z, above=2.0) + args = (u, v, w, z, 1000.0, 1000.0, 1000.0) + + none_bc = _w_component( + pydda.cost_functions.calculate_mass_continuity_gradient(*args, upper_bc=0), + shape, + ) + grid_top = _w_component( + pydda.cost_functions.calculate_mass_continuity_gradient(*args, upper_bc=1), + shape, + ) + echo_top = _w_component( + pydda.cost_functions.calculate_mass_continuity_gradient( + *args, upper_bc=2, upper_bc_mask=mask + ), + shape, + ) + + # The surface is always impermeable, the domain top only for upper_bc=1. + assert np.all(none_bc[0] == 0) + assert not np.all(none_bc[-1] == 0) + assert np.all(grid_top[-1] == 0) + np.testing.assert_allclose(grid_top[:-1], none_bc[:-1]) + + # upper_bc=2 pins w exactly on the mask and leaves everything else alone. + assert np.all(echo_top[mask] == 0) + np.testing.assert_allclose(echo_top[~mask], none_bc[~mask]) + assert np.any(echo_top != grid_top) + + # The legacy booleans still mean what they used to. This is the regression + # guard for `upper_bc is True`, which silently skipped the condition for + # the integer modes. + np.testing.assert_allclose( + _w_component( + pydda.cost_functions.calculate_mass_continuity_gradient( + *args, upper_bc=True + ), + shape, + ), + grid_top, + ) + np.testing.assert_allclose( + _w_component( + pydda.cost_functions.calculate_mass_continuity_gradient( + *args, upper_bc=False + ), + shape, + ), + none_bc, + ) + + +def test_all_gradients_honor_echo_top_upper_bc(): + """Every gradient term that takes upper_bc applies the echo top mask.""" + u, v, w, z, vrs = _make_upper_bc_inputs() + shape = u.shape + mask = pydda.cost_functions.calculate_echo_top_mask(vrs, z, above=2.0) + els = [np.ones(shape) * 0.5] + azs = [np.ones(shape) * 0.5] + wts = [np.zeros(shape)] + weights = np.ones((1,) + shape) + + grads = { + "radial_vel": pydda.cost_functions.calculate_grad_radial_vel( + vrs, els, azs, u, v, w, wts, weights, 1.0, upper_bc=2, upper_bc_mask=mask + ), + "smoothness": pydda.cost_functions.calculate_smoothness_gradient( + u, + v, + w, + 1000.0, + 1000.0, + 1000.0, + Cx=1e-4, + Cy=1e-4, + Cz=1e-4, + upper_bc=2, + upper_bc_mask=mask, + ), + "mass_continuity": pydda.cost_functions.calculate_mass_continuity_gradient( + u, v, w, z, 1000.0, 1000.0, 1000.0, upper_bc=2, upper_bc_mask=mask + ), + "vorticity": pydda.cost_functions.calculate_vertical_vorticity_gradient( + u, + v, + w, + 1000.0, + 1000.0, + 1000.0, + 1.0, + 1.0, + coeff=1e-5, + upper_bc=2, + upper_bc_mask=mask, + ), + } + for name, grad in grads.items(): + grad_w = _w_component(grad, shape) + assert np.all(grad_w[mask] == 0), "%s ignored the echo top mask" % name + assert np.all(grad_w[0] == 0), "%s ignored the surface condition" % name + + +@pytest.mark.skipif(not JAX_AVAILABLE, reason="Jax not installed") +def test_mass_continuity_gradient_upper_bc_jax(): + u, v, w, z, vrs = _make_upper_bc_inputs() + shape = u.shape + mask = pydda.cost_functions.calculate_echo_top_mask(vrs, z, above=2.0) + args = ( + jnp.array(u), + jnp.array(v), + jnp.array(w), + jnp.array(z), + 1000.0, + 1000.0, + 1000.0, + ) + + none_bc = _w_component( + pydda.cost_functions.jax.calculate_mass_continuity_gradient(*args, upper_bc=0), + shape, + ) + grid_top = _w_component( + pydda.cost_functions.jax.calculate_mass_continuity_gradient(*args, upper_bc=1), + shape, + ) + echo_top = _w_component( + pydda.cost_functions.jax.calculate_mass_continuity_gradient( + *args, upper_bc=2, upper_bc_mask=mask + ), + shape, + ) + + assert np.all(none_bc[0] == 0) + assert not np.all(none_bc[-1] == 0) + assert np.all(grid_top[-1] == 0) + assert np.all(echo_top[mask] == 0) + np.testing.assert_allclose(echo_top[~mask], none_bc[~mask]) + + +@pytest.mark.skipif(not TENSORFLOW_AVAILABLE, reason="TensorFlow not installed") +def test_mass_continuity_gradient_upper_bc_tf(): + u, v, w, z, vrs = _make_upper_bc_inputs() + shape = u.shape + mask = pydda.cost_functions.calculate_echo_top_mask(vrs, z, above=2.0) + args = ( + tf.constant(u, dtype=tf.float32), + tf.constant(v, dtype=tf.float32), + tf.constant(w, dtype=tf.float32), + tf.constant(z, dtype=tf.float32), + 1000.0, + 1000.0, + 1000.0, + ) + tf_mask = tf.constant(mask, dtype=tf.bool) + + none_bc = _w_component( + pydda.cost_functions.tf.calculate_mass_continuity_gradient(*args, upper_bc=0), + shape, + ) + grid_top = _w_component( + pydda.cost_functions.tf.calculate_mass_continuity_gradient(*args, upper_bc=1), + shape, + ) + echo_top = _w_component( + pydda.cost_functions.tf.calculate_mass_continuity_gradient( + *args, upper_bc=2, upper_bc_mask=tf_mask + ), + shape, + ) + + assert np.all(none_bc[0] == 0) + assert not np.all(none_bc[-1] == 0) + assert np.all(grid_top[-1] == 0) + assert np.all(echo_top[mask] == 0) + np.testing.assert_allclose(echo_top[~mask], none_bc[~mask]) diff --git a/pydda/tests/test_retrieval.py b/pydda/tests/test_retrieval.py index f50b44ca..a2c5c01e 100644 --- a/pydda/tests/test_retrieval.py +++ b/pydda/tests/test_retrieval.py @@ -23,7 +23,9 @@ TF_AVAILABLE = False try: + # The Jax engine needs jaxopt for the solver, not just jax itself. import jax + import jaxopt JAX_AVAILABLE = True except ImportError: @@ -78,6 +80,129 @@ def test_make_updraft_from_convergence_field(): assert np.ma.max(new_w > 3) +def _twpice_grids(): + """The TWP-ICE grid pair used by the retrieval tests, with a first guess + taken from the sounding (which has w = 0 everywhere).""" + Grid0 = pydda.io.read_grid(pydda.tests.EXAMPLE_RADAR0) + Grid1 = pydda.io.read_grid(pydda.tests.EXAMPLE_RADAR1) + sounding = pyart.io.read_arm_sonde(pydda.tests.SOUNDING_PATH) + Grid0 = pydda.initialization.make_wind_field_from_profile( + Grid0, sounding[1], vel_field="corrected_velocity" + ) + return Grid0, Grid1 + + +# Common arguments for the echo top boundary condition tests. The low pass +# filter and the masking are turned off so that the retrieved w can be compared +# against the first guess point by point. +ECHO_TOP_KWARGS = dict( + Co=100, + Cm=1500.0, + max_iterations=20, + Cz=0, + Cmod=0.0, + vel_name="corrected_velocity", + wind_tol=0.1, + refl_field="reflectivity", + frz=5000.0, + upper_bc=2, + above=2.0, + low_pass_filter=False, + mask_outside_opt=False, + mask_w_outside_opt=False, +) + + +def _assert_impermeable_above_echo_top(Grids, parameters): + """w must be untouched (i.e. still zero) wherever the echo top condition + applies, and the retrieval must still be physically sensible elsewhere.""" + # The TensorFlow engines store these as tensors rather than numpy arrays. + mask = np.asarray(parameters.upper_bc_mask) + z = np.asarray(parameters.z) + assert mask.dtype == bool + # The condition should cover a substantial part of, but not all of, the grid. + assert 0.0 < mask.mean() < 1.0 + assert not mask[z <= 2000.0].any() + + w = Grids[0]["w"].values.squeeze() + # A permanently zero gradient component is never moved by L-BFGS-B, so the + # first guess of w = 0 survives exactly. + np.testing.assert_array_equal(w[mask], 0.0) + # ...but the retrieval still produces updrafts where the radars see echoes. + assert np.nanmax(w[~mask]) > 5 + + u_mean = np.nanmean(Grids[0]["u"].values) + v_mean = np.nanmean(Grids[0]["v"].values) + assert u_mean > 0 + assert v_mean < 0 + + +def test_twpice_case_upper_bc_echo_top(): + """The echo top impermeability condition holds w fixed above the echo top.""" + Grid0, Grid1 = _twpice_grids() + Grids, parameters = pydda.retrieval.get_dd_wind_field( + [Grid0, Grid1], engine="scipy", **ECHO_TOP_KWARGS + ) + _assert_impermeable_above_echo_top(Grids, parameters) + + +def test_twpice_case_upper_bc_modes_differ(): + """The three upper boundary conditions give three different wind fields.""" + kwargs = dict(ECHO_TOP_KWARGS) + del kwargs["upper_bc"] + # This test only needs the three fields to be distinguishable, not + # converged, so it runs for fewer iterations than the others. + kwargs["max_iterations"] = 10 + + results = {} + for upper_bc in (0, 1, 2): + Grid0, Grid1 = _twpice_grids() + Grids, parameters = pydda.retrieval.get_dd_wind_field( + [deepcopy(Grid0), deepcopy(Grid1)], + engine="scipy", + upper_bc=upper_bc, + **kwargs, + ) + results[upper_bc] = (Grids[0]["w"].values.squeeze(), parameters) + + w_none, params_none = results[0] + w_grid_top, params_grid_top = results[1] + w_echo_top, _ = results[2] + + # No mask is built unless it is needed. + assert params_none.upper_bc_mask is None + assert params_grid_top.upper_bc_mask is None + + # upper_bc=1 pins only the model top; this is the regression guard for + # `upper_bc is True`, which used to skip the condition for integer modes. + np.testing.assert_array_equal(w_grid_top[-1, :, :], 0.0) + assert np.any(w_none[-1, :, :] != 0) + + assert np.any(w_echo_top != w_grid_top) + assert np.any(w_echo_top != w_none) + + +@pytest.mark.skipif(not JAX_AVAILABLE, reason="Jax not installed") +def test_twpice_case_upper_bc_echo_top_jax(): + """The echo top impermeability condition also works with the Jax engine.""" + Grid0, Grid1 = _twpice_grids() + Grids, parameters = pydda.retrieval.get_dd_wind_field( + [Grid0, Grid1], engine="jax", **ECHO_TOP_KWARGS + ) + _assert_impermeable_above_echo_top(Grids, parameters) + + +@pytest.mark.skipif(not TF_AVAILABLE, reason="TensorFlow not installed") +def test_twpice_case_upper_bc_echo_top_tensorflow(): + """The echo top impermeability condition also works with the TensorFlow + engine.""" + Grid0, Grid1 = _twpice_grids() + Grids, parameters = pydda.retrieval.get_dd_wind_field( + [Grid0, Grid1], engine="tensorflow", **ECHO_TOP_KWARGS + ) + _assert_impermeable_above_echo_top(Grids, parameters) + + @pytest.mark.skipif(not JAX_AVAILABLE, reason="Jax not installed") def test_twpice_case_jax(): """Use a test case from TWP-ICE"""