Vectorize coordinates._construct_face_centroids() (#1117)

* Vectorize coordinates._construct_face_centroids()

* numba implementation

* add numba prange

* remove comment

* faster n_node_per_face construction

* Remove unused variable assignment

---------

Co-authored-by: Philip Chmielowiec <67855069+philipc2@users.noreply.github.com>
Co-authored-by: Aaron Zedwick <95507181+aaronzedwick@users.noreply.github.com>

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

@@ -328,23 +328,25 @@ 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)
-    n_face = n_nodes_per_face.shape[0]
-
-    for i_face in prange(n_face):
-        n_max_nodes = n_nodes_per_face[i_face]
 
-        x = np.mean(node_x[face_nodes[i_face, 0:n_max_nodes]])
-        y = np.mean(node_y[face_nodes[i_face, 0:n_max_nodes]])
-        z = np.mean(node_z[face_nodes[i_face, 0:n_max_nodes]])
+    for face_idx in prange(face_nodes.shape[0]):
+        n_max_nodes = n_nodes_per_face[face_idx]
+        # Compute Cartesian Average
+        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]])
 
+        # 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
 
-        centroid_x[i_face] = x
-        centroid_y[i_face] = y
-        centroid_z[i_face] = z
     return centroid_x, centroid_y, centroid_z