Fixing CGrid_Velocity interpolation for multiple particles

Using vectorized version of dot-product calculation as suggested by @michaeldenes in #2152 (comment)

Author: erikvansebille

parcels/application_kernels/interpolation.py

@@ -149,16 +149,24 @@ def CGrid_Velocity(
         py = np.array([grid.lat[yi, xi], grid.lat[yi, xi + 1], grid.lat[yi + 1, xi + 1], grid.lat[yi + 1, xi]])
 
     if vectorfield._mesh_type == "spherical":
-        px[0] = px[0] + 360 if px[0] < x - 225 else px[0]
-        px[0] = px[0] - 360 if px[0] > x + 225 else px[0]
+        px[0] = np.where(px[0] < x - 225, px[0] + 360, px[0])
+        px[0] = np.where(px[0] > x + 225, px[0] - 360, px[0])
         px[1:] = np.where(px[1:] - px[0] > 180, px[1:] - 360, px[1:])
         px[1:] = np.where(-px[1:] + px[0] > 180, px[1:] + 360, px[1:])
     xx = (1 - xsi) * (1 - eta) * px[0] + xsi * (1 - eta) * px[1] + xsi * eta * px[2] + (1 - xsi) * eta * px[3]
     np.testing.assert_allclose(xx, x, atol=1e-4)
-    c1 = i_u._geodetic_distance(py[0], py[1], px[0], px[1], vectorfield._mesh_type, np.dot(i_u.phi2D_lin(0.0, xsi), py))
-    c2 = i_u._geodetic_distance(py[1], py[2], px[1], px[2], vectorfield._mesh_type, np.dot(i_u.phi2D_lin(eta, 1.0), py))
-    c3 = i_u._geodetic_distance(py[2], py[3], px[2], px[3], vectorfield._mesh_type, np.dot(i_u.phi2D_lin(1.0, xsi), py))
-    c4 = i_u._geodetic_distance(py[3], py[0], px[3], px[0], vectorfield._mesh_type, np.dot(i_u.phi2D_lin(eta, 0.0), py))
+    c1 = i_u._geodetic_distance(
+        py[0], py[1], px[0], px[1], vectorfield._mesh_type, np.einsum("ij,ji->i", i_u.phi2D_lin(0.0, xsi), py)
+    )
+    c2 = i_u._geodetic_distance(
+        py[1], py[2], px[1], px[2], vectorfield._mesh_type, np.einsum("ij,ji->i", i_u.phi2D_lin(eta, 1.0), py)
+    )
+    c3 = i_u._geodetic_distance(
+        py[2], py[3], px[2], px[3], vectorfield._mesh_type, np.einsum("ij,ji->i", i_u.phi2D_lin(1.0, xsi), py)
+    )
+    c4 = i_u._geodetic_distance(
+        py[3], py[0], px[3], px[0], vectorfield._mesh_type, np.einsum("ij,ji->i", i_u.phi2D_lin(eta, 0.0), py)
+    )
 
     lenT = 2 if np.any(tau > 0) else 1
     lenZ = 2 if np.any(zeta > 0) else 1

parcels/tools/interpolation_utils.py

@@ -1,4 +1,3 @@
-import math
 from collections.abc import Callable
 from typing import Literal
 
@@ -179,7 +178,7 @@ def _geodetic_distance(lat1: float, lat2: float, lon1: float, lon2: float, mesh:
     if mesh == "spherical":
         rad = np.pi / 180.0
         deg2m = 1852 * 60.0
-        return np.sqrt(((lon2 - lon1) * deg2m * math.cos(rad * lat)) ** 2 + ((lat2 - lat1) * deg2m) ** 2)
+        return np.sqrt(((lon2 - lon1) * deg2m * np.cos(rad * lat)) ** 2 + ((lat2 - lat1) * deg2m) ** 2)
     else:
         return np.sqrt((lon2 - lon1) ** 2 + (lat2 - lat1) ** 2)
 
@@ -188,10 +187,10 @@ def _compute_jacobian_determinant(py: np.ndarray, px: np.ndarray, eta: float, xs
     dphidxsi = np.column_stack([eta - 1, 1 - eta, eta, -eta])
     dphideta = np.column_stack([xsi - 1, -xsi, xsi, 1 - xsi])
 
-    dxdxsi = np.dot(dphidxsi, px)
-    dxdeta = np.dot(dphideta, px)
-    dydxsi = np.dot(dphidxsi, py)
-    dydeta = np.dot(dphideta, py)
-    jac = dxdxsi * dydeta - dxdeta * dydxsi
-    # TODO check how to properly vectorize this function (ok to return diagonal of the Jacobian?)
-    return jac.diagonal()
+    dxdxsi_diag = np.einsum("ij,ji->i", dphidxsi, px)
+    dxdeta_diag = np.einsum("ij,ji->i", dphideta, px)
+    dydxsi_diag = np.einsum("ij,ji->i", dphidxsi, py)
+    dydeta_diag = np.einsum("ij,ji->i", dphideta, py)
+
+    jac_diag = dxdxsi_diag * dydeta_diag - dxdeta_diag * dydxsi_diag
+    return jac_diag