OPT: Combine cost/gradient into one

pydda/cost_functions/__init__.py

@@ -87,6 +87,7 @@
 from ._cost_functions_numpy import calculate_radial_vel_cost_function
 from ._cost_functions_numpy import calculate_fall_speed
 from ._cost_functions_numpy import calculate_grad_radial_vel
+from ._cost_functions_numpy import calculate_radial_vel_cost_and_gradient
 from ._cost_functions_numpy import calculate_mass_continuity
 from ._cost_functions_numpy import calculate_mass_continuity_gradient
 from ._cost_functions_numpy import calculate_smoothness_cost
@@ -99,4 +100,4 @@
 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
+from .cost_functions import J_function, J_and_grad, grad_J, grad_jax, J_function_jax

pydda/cost_functions/_cost_functions_numpy.py

@@ -235,6 +235,70 @@ def calculate_grad_radial_vel(
     return y.flatten()
 
 
+def calculate_radial_vel_cost_and_gradient(
+    vrs,
+    els,
+    azs,
+    u,
+    v,
+    w,
+    wts,
+    weights,
+    rmsVr,
+    coeff=1.0,
+    upper_bc=1,
+    upper_bc_mask=None,
+    parallel=False,
+):
+    """Calculate radar cost and gradient while sharing the radar projection.
+
+    This is used by the combined SciPy optimizer callback.  The existing
+    cost-only and gradient-only functions remain available for callers that
+    need just one result.
+    """
+    lambda_o = coeff / (rmsVr * rmsVr)
+
+    if parallel:
+        vrs_arr = np.stack(vrs)
+        els_arr = np.stack(els)
+        azs_arr = np.stack(azs)
+        wts_arr = np.stack(wts)
+        cos_els = np.cos(els_arr)
+        sin_els = np.sin(els_arr)
+        x_projection = cos_els * np.sin(azs_arr)
+        y_projection = cos_els * np.cos(azs_arr)
+        v_ar = (
+            x_projection * u[np.newaxis]
+            + y_projection * v[np.newaxis]
+            + sin_els * (w[np.newaxis] - np.abs(wts_arr))
+        )
+        residual = v_ar - vrs_arr
+        cost = lambda_o * np.sum(np.square(residual) * weights)
+        p_x1 = np.sum(2 * residual * x_projection * weights, axis=0) * lambda_o
+        p_y1 = np.sum(2 * residual * y_projection * weights, axis=0) * lambda_o
+        p_z1 = np.sum(2 * residual * sin_els * weights, axis=0) * lambda_o
+    else:
+        cost = 0.0
+        p_x1 = np.zeros(vrs[0].shape)
+        p_y1 = np.zeros(vrs[0].shape)
+        p_z1 = np.zeros(vrs[0].shape)
+        for vr, el, az, wt, weight in zip(vrs, els, azs, wts, weights):
+            cos_el = np.cos(el)
+            sin_el = np.sin(el)
+            x_projection = cos_el * np.sin(az)
+            y_projection = cos_el * np.cos(az)
+            v_ar = x_projection * u + y_projection * v + sin_el * (w - np.abs(wt))
+            residual = v_ar - vr
+            cost += lambda_o * np.sum(np.square(residual) * weight)
+            p_x1 += 2 * residual * x_projection * weight * lambda_o
+            p_y1 += 2 * residual * y_projection * weight * lambda_o
+            p_z1 += 2 * residual * sin_el * weight * lambda_o
+
+    p_z1[0, :, :] = 0
+    p_z1 = _apply_upper_bc(p_z1, upper_bc, upper_bc_mask)
+    return cost, np.stack((p_x1, p_y1, p_z1), axis=0).flatten()
+
+
 def calculate_smoothness_cost(u, v, w, dx, dy, dz, Cx=1e-5, Cy=1e-5, Cz=1e-5):
     """
     Calculates the smoothness cost function by taking the Laplacian of the

pydda/cost_functions/_cost_functions_tensorflow.py

@@ -95,6 +95,51 @@ def calculate_radial_vel_cost_function(
     return J_o
 
 
+def calculate_radial_vel_cost_and_gradient(
+    vrs,
+    els,
+    azs,
+    u,
+    v,
+    w,
+    wts,
+    weights,
+    rmsVr,
+    coeff=1.0,
+    upper_bc=1,
+    upper_bc_mask=None,
+    lower_bc=True,
+):
+    """Calculate radar cost and gradient while sharing each radar residual."""
+    lambda_o = coeff / (rmsVr * rmsVr)
+    costs = []
+    p_x1 = tf.zeros_like(u)
+    p_y1 = tf.zeros_like(v)
+    p_z1 = tf.zeros_like(w)
+
+    for i in range(len(vrs)):
+        cos_el = tf.math.cos(els[i])
+        sin_el = tf.math.sin(els[i])
+        v_ar = (
+            cos_el * tf.math.sin(azs[i]) * u
+            + cos_el * tf.math.cos(azs[i]) * v
+            + sin_el * (w - tf.math.abs(wts[i]))
+        )
+        residual = v_ar - vrs[i]
+        costs.append(lambda_o * tf.reduce_sum(tf.math.square(residual) * weights[i]))
+        p_x1 += 2 * residual * cos_el * tf.math.sin(azs[i]) * weights[i] * lambda_o
+        p_y1 += 2 * residual * cos_el * tf.math.cos(azs[i]) * weights[i] * lambda_o
+        p_z1 += 2 * residual * sin_el * weights[i] * lambda_o
+
+    if lower_bc:
+        p_z1 = tf.concat(
+            [tf.zeros_like(p_z1[0:1]), p_z1[1:]],
+            axis=0,
+        )
+    p_z1 = _apply_upper_bc(p_z1, upper_bc, upper_bc_mask)
+    return tf.add_n(costs), tf.reshape(tf.stack((p_x1, p_y1, p_z1)), [-1])
+
+
 def calculate_grad_radial_vel(
     vrs,
     els,