Vectorize coordinates._construct_face_centroids()

uxarray/grid/connectivity.py

@@ -146,13 +146,14 @@ def _build_n_nodes_per_face(face_nodes, n_face, n_max_face_nodes):
     """Constructs ``n_nodes_per_face``, which contains the number of non-fill-
     value nodes for each face in ``face_node_connectivity``"""
 
-    # padding to shape [n_face, n_max_face_nodes + 1]
-    closed = np.ones((n_face, n_max_face_nodes + 1), dtype=INT_DTYPE) * INT_FILL_VALUE
-
-    closed[:, :-1] = face_nodes.copy()
-
-    n_nodes_per_face = np.argmax(closed == INT_FILL_VALUE, axis=1)
-
+    n_face, n_max_face_nodes = face_nodes.shape
+    n_nodes_per_face = np.empty(n_face, dtype=INT_DTYPE)
+    for i in range(n_face):
+        c = 0
+        for j in range(n_max_face_nodes):
+            if face_nodes[i, j] != INT_FILL_VALUE:
+                c += 1
+        n_nodes_per_face[i] = c
     return n_nodes_per_face
 
 

uxarray/grid/coordinates.py

@@ -5,7 +5,7 @@
 
 from uxarray.conventions import ugrid
 
-from numba import njit
+from numba import njit, prange
 from uxarray.constants import ERROR_TOLERANCE
 from typing import Union
 
@@ -305,6 +305,7 @@ def _populate_face_centroids(grid, repopulate=False):
         )
 
 
+@njit(cache=True, parallel=True)
 def _construct_face_centroids(node_x, node_y, node_z, face_nodes, n_nodes_per_face):
     """Constructs the xyz centroid coordinate for each face using Cartesian
     Averaging.
@@ -327,17 +328,26 @@ def _construct_face_centroids(node_x, node_y, node_z, face_nodes, n_nodes_per_fa
     tuple
         The x, y, and z coordinates of the centroids.
     """
+
     centroid_x = np.zeros((face_nodes.shape[0]), dtype=np.float64)
     centroid_y = np.zeros((face_nodes.shape[0]), dtype=np.float64)
     centroid_z = np.zeros((face_nodes.shape[0]), dtype=np.float64)
 
-    for face_idx, n_max_nodes in enumerate(n_nodes_per_face):
+    for face_idx in prange(face_nodes.shape[0]):
+        n_max_nodes = n_nodes_per_face[face_idx]
         # Compute Cartesian Average
-        centroid_x[face_idx] = np.mean(node_x[face_nodes[face_idx, 0:n_max_nodes]])
-        centroid_y[face_idx] = np.mean(node_y[face_nodes[face_idx, 0:n_max_nodes]])
-        centroid_z[face_idx] = np.mean(node_z[face_nodes[face_idx, 0:n_max_nodes]])
+        x = np.mean(node_x[face_nodes[face_idx, 0:n_max_nodes]])
+        y = np.mean(node_y[face_nodes[face_idx, 0:n_max_nodes]])
+        z = np.mean(node_z[face_nodes[face_idx, 0:n_max_nodes]])
 
-    return _normalize_xyz(centroid_x, centroid_y, centroid_z)
+        # Normalize coordinates
+        x, y, z = _normalize_xyz_scalar(x, y, z)
+        # Store coordinates
+        centroid_x[face_idx] = x
+        centroid_y[face_idx] = y
+        centroid_z[face_idx] = z
+
+    return centroid_x, centroid_y, centroid_z
 
 
 def _welzl_recursive(points, boundary, R):