Fixes in the JAX engine, parallel support for scipy engine, and background weight fixes

examples/plot_examples.py

@@ -36,6 +36,7 @@
     upper_bc=1,
     wind_tol=0.5,
     engine="scipy",
+    parallel=False,
 )
 
 # Plot a horizontal cross section

pydda/__init__.py

@@ -12,7 +12,7 @@
 from . import constraints
 from . import io
 
-__version__ = "2.3.0"
+__version__ = "2.4.0"
 
 print("Welcome to PyDDA %s" % __version__)
 print("If you are using PyDDA in your publications, please cite:")

pydda/cost_functions/_cost_functions_jax.py

@@ -548,7 +548,7 @@ def calculate_background_gradient(u, v, w, weights, u_back, v_back, Cb=0.01):
         calculate_background_cost, u, v, w, weights, u_back, v_back, Cb
     )
     u_grad, v_grad, w_grad, _, _, _, _ = fun_vjp(1.0)
-    y = np.stack([u_grad, v_grad, w_grad], axis=0)
+    y = jnp.stack([u_grad, v_grad, w_grad], axis=0)
     return y.flatten().copy()
 
 

pydda/cost_functions/_cost_functions_numpy.py

@@ -8,7 +8,7 @@
 
 
 def calculate_radial_vel_cost_function(
-    vrs, azs, els, u, v, w, wts, rmsVr, weights, coeff=1.0
+    vrs, azs, els, u, v, w, wts, rmsVr, weights, coeff=1.0, parallel=False
 ):
     """
     Calculates the cost function due to difference of the wind field from
@@ -54,22 +54,44 @@ def calculate_radial_vel_cost_function(
     Technol., 26, 2089–2106, https://doi.org/10.1175/2009JTECHA1256.1
     """
 
-    J_o = 0
     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)
+        v_ar = (
+            np.cos(els_arr) * np.sin(azs_arr) * u[np.newaxis]
+            + np.cos(els_arr) * np.cos(azs_arr) * v[np.newaxis]
+            + np.sin(els_arr) * (w[np.newaxis] - np.abs(wts_arr))
+        )
+        return lambda_o * np.sum(np.square(vrs_arr - v_ar) * weights)
+
+    J_o = 0
     for i in range(len(vrs)):
         v_ar = (
             np.cos(els[i]) * np.sin(azs[i]) * u
             + np.cos(els[i]) * np.cos(azs[i]) * v
             + np.sin(els[i]) * (w - np.abs(wts[i]))
         )
-        the_weight = weights[i]
-        J_o += lambda_o * np.sum(np.square(vrs[i] - v_ar) * the_weight)
+        J_o += lambda_o * np.sum(np.square(vrs[i] - v_ar) * weights[i])
 
     return J_o
 
 
 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=True,
+    parallel=False,
 ):
     """
     Calculates the gradient of the cost function due to difference of wind
@@ -112,29 +134,45 @@ def calculate_grad_radial_vel(
     # Use zero for all masked values since we don't want to add them into
     # the cost function
 
-    p_x1 = np.zeros(vrs[0].shape)
-    p_y1 = np.zeros(vrs[0].shape)
-    p_z1 = np.zeros(vrs[0].shape)
     lambda_o = coeff / (rmsVr * rmsVr)
 
-    for i in range(len(vrs)):
+    if parallel:
+        vrs_arr = np.stack(vrs)
+        els_arr = np.stack(els)
+        azs_arr = np.stack(azs)
+        wts_arr = np.stack(wts)
         v_ar = (
-            np.cos(els[i]) * np.sin(azs[i]) * u
-            + np.cos(els[i]) * np.cos(azs[i]) * v
-            + np.sin(els[i]) * (w - np.abs(wts[i]))
+            np.cos(els_arr) * np.sin(azs_arr) * u[np.newaxis]
+            + np.cos(els_arr) * np.cos(azs_arr) * v[np.newaxis]
+            + np.sin(els_arr) * (w[np.newaxis] - np.abs(wts_arr))
         )
+        residual = 2 * (v_ar - vrs_arr) * lambda_o
+        p_x1 = np.sum(residual * np.cos(els_arr) * np.sin(azs_arr) * weights, axis=0)
+        p_y1 = np.sum(residual * np.cos(els_arr) * np.cos(azs_arr) * weights, axis=0)
+        p_z1 = np.sum(residual * np.sin(els_arr) * weights, axis=0)
+    else:
+        p_x1 = np.zeros(vrs[0].shape)
+        p_y1 = np.zeros(vrs[0].shape)
+        p_z1 = np.zeros(vrs[0].shape)
+
+        for i in range(len(vrs)):
+            v_ar = (
+                np.cos(els[i]) * np.sin(azs[i]) * u
+                + np.cos(els[i]) * np.cos(azs[i]) * v
+                + np.sin(els[i]) * (w - np.abs(wts[i]))
+            )
 
-        x_grad = (
-            2 * (v_ar - vrs[i]) * np.cos(els[i]) * np.sin(azs[i]) * weights[i]
-        ) * lambda_o
-        y_grad = (
-            2 * (v_ar - vrs[i]) * np.cos(els[i]) * np.cos(azs[i]) * weights[i]
-        ) * lambda_o
-        z_grad = (2 * (v_ar - vrs[i]) * np.sin(els[i]) * weights[i]) * lambda_o
-
-        p_x1 += x_grad
-        p_y1 += y_grad
-        p_z1 += z_grad
+            x_grad = (
+                2 * (v_ar - vrs[i]) * np.cos(els[i]) * np.sin(azs[i]) * weights[i]
+            ) * lambda_o
+            y_grad = (
+                2 * (v_ar - vrs[i]) * np.cos(els[i]) * np.cos(azs[i]) * weights[i]
+            ) * lambda_o
+            z_grad = (2 * (v_ar - vrs[i]) * np.sin(els[i]) * weights[i]) * lambda_o
+
+            p_x1 += x_grad
+            p_y1 += y_grad
+            p_z1 += z_grad
 
     # Impermeability condition
     p_z1[0, :, :] = 0