DRAFT: Optimize Connectivity Construction

docs/api.rst

@@ -480,3 +480,13 @@ Accurate Computing
 
    utils.computing.cross_fma
    utils.computing.dot_fma
+
+
+Constants
+------------------
+
+.. autosummary::
+   :toctree: generated/
+
+   constants.INT_FILL_VALUE
+   constants.INT_DTYPE

docs/getting-started/overview.rst

@@ -36,30 +36,35 @@ other geometric faces.
 Core Data Structures
 ====================
 
-The functionality of UXarray is built around three core data structures which provide
-an Unstructured Grid aware implementation of many Xarray functions and use cases.
+The functionality of UXarray is built around three core data structures:
+
+* :class:`Grid`
+  Used to represent an Unstructured Grid, housing grid-specific methods and
+  topology variables.
+
+* :class:`UxDataset`
+  Inherits from :py:class:`xarray.Dataset`, providing the same functionality
+  but extended to operate directly on Unstructured Grids. An :class:`UxDataset`
+  is linked to a :class:`Grid` object via the :attr:`UxDataset.uxgrid` property.
+
+* :class:`UxDataArray`
+  Similarly inherits from :py:class:`xarray.DataArray` and contains a
+  :attr:`UxDataArray.uxgrid` property just like :class:`UxDataset`.
 
-* ``Grid`` is used to represent our Unstructured Grid, housing grid-specific methods
-  and topology variables.
-* ``UxDataset`` inherits from the ``xarray.Dataset`` class, providing much of the same
-  functionality but extended to operate on Unstructured Grids. Other than new and
-  overloaded methods, it is linked to a ``Grid`` object through the use of a class
-  property (``UxDataset.uxgrid``) to provide a grid-aware implementation. An instance
-  of ``UxDataset`` can be thought of as a collection of Data Variables that reside on
-  some Unstructured Grid as defined in the ``uxgrid`` property.
-* ``UxDataArray`` similarly inherits from the ``xarray.DataArray`` class and contains
-  a ``Grid`` property (``UxDataArray.uxgrid``) just like ``UxDataset``.
 
 Core Functionality
 ==================
 
 In addition to providing a way to load in and interface with Unstructured Grids, we
 also aim to provide computational and analysis operators that directly operate on
 Unstructured Grids. Some of these include:
-* Visualization
-* Remapping
-* Subsetting & Selection
-* Aggregations
+- Visualization
+- Remapping
+- Subsetting & Selection
+- Cross Sections
+- Aggregations
+- Calculus Operations
+- Zonal Averaging
 
 A more detailed overview of supported functionality can be found in our `API Reference <https://uxarray.readthedocs.io/en/latest/api.html>`_
 and `User Guide <https://uxarray.readthedocs.io/en/latest/userguide.html>`_ sections.

test/test_centroids.py

@@ -69,9 +69,10 @@ def test_edge_centroids_from_triangle():
     grid = ux.open_grid(test_triangle, latlon=False)
     _populate_edge_centroids(grid)
 
-    centroid_x = np.mean(grid.node_x[grid.edge_node_connectivity[0][0:]])
-    centroid_y = np.mean(grid.node_y[grid.edge_node_connectivity[0][0:]])
-    centroid_z = np.mean(grid.node_z[grid.edge_node_connectivity[0][0:]])
+
+    centroid_x = grid.node_x[grid.edge_node_connectivity].mean(axis=1)
+    centroid_y = grid.node_y[grid.edge_node_connectivity].mean(axis=1)
+    centroid_z = grid.node_z[grid.edge_node_connectivity].mean(axis=1)
 
     assert centroid_x == grid.edge_x[0]
     assert centroid_y == grid.edge_y[0]

test/test_gradient.py

@@ -53,6 +53,7 @@ def test_quad_hex():
         else:
             assert grad.values[i] != 0
 
+    # TODO:
     expected_values = np.array([27.95, 20.79, 28.96, 0, 0, 0, 0, 60.64, 0, 86.45, 0, 0, 0, 0, 0, 0, 0, 0, 0])
     nt.assert_almost_equal(grad.values, expected_values, 1e-2)
 

test/test_grid.py

@@ -500,15 +500,16 @@ def test_connectivity_build_face_edges_connectivity_mpas():
     edge_nodes_expected.sort(axis=1)
     edge_nodes_expected = np.unique(edge_nodes_expected, axis=0)
 
-    edge_nodes_output, _, _ = _build_edge_node_connectivity(mpas_grid_ux.face_node_connectivity.values,
-                                                            mpas_grid_ux.n_face,
-                                                            mpas_grid_ux.n_max_face_nodes)
 
-    assert np.array_equal(edge_nodes_expected, edge_nodes_output)
+    edge_node_connectivity, _ = _build_edge_node_connectivity(mpas_grid_ux.face_node_connectivity.values,
+                                                              mpas_grid_ux.n_nodes_per_face.values)
+
+    print()
+    assert np.array_equal(edge_nodes_expected, edge_node_connectivity)
 
     n_face = mpas_grid_ux.n_node
     n_node = mpas_grid_ux.n_face
-    n_edge = edge_nodes_output.shape[0]
+    n_edge = edge_node_connectivity.shape[0]
 
     assert (n_face == n_edge - n_node + 2)
 

uxarray/core/zonal.py

@@ -3,7 +3,7 @@
 
 
 from uxarray.grid.integrate import _zonal_face_weights, _zonal_face_weights_robust
-from uxarray.grid.utils import _get_cartesian_face_edge_nodes
+# from uxarray.grid.utils import _get_cartesian_face_edge_nodes
 
 
 def _compute_non_conservative_zonal_mean(uxda, latitudes, use_robust_weights=False):
@@ -18,14 +18,7 @@ def _compute_non_conservative_zonal_mean(uxda, latitudes, use_robust_weights=Fal
         # Create a NumPy array for storing results
         result = np.zeros(shape, dtype=uxda.dtype)
 
-    faces_edge_nodes_xyz = _get_cartesian_face_edge_nodes(
-        uxgrid.face_node_connectivity.values,
-        uxgrid.n_face,
-        uxgrid.n_max_face_nodes,
-        uxgrid.node_x.values,
-        uxgrid.node_y.values,
-        uxgrid.node_z.values,
-    )
+    face_edge_nodes_cartesian = uxda.uxgrid.face_edge_nodes_cartesian
 
     bounds = uxgrid.bounds.values
 
@@ -34,19 +27,21 @@ def _compute_non_conservative_zonal_mean(uxda, latitudes, use_robust_weights=Fal
 
         z = np.sin(np.deg2rad(lat))
 
-        faces_edge_nodes_xyz_candidate = faces_edge_nodes_xyz[face_indices, :, :, :]
+        face_edge_nodes_cartesian_candidate = face_edge_nodes_cartesian[
+            face_indices, :, :, :
+        ]
 
         n_nodes_per_face_candidate = n_nodes_per_face[face_indices]
 
         bounds_candidate = bounds[face_indices]
 
         if use_robust_weights:
             weights = _zonal_face_weights_robust(
-                faces_edge_nodes_xyz_candidate, z, bounds_candidate
+                face_edge_nodes_cartesian_candidate, z, bounds_candidate
             )["weight"].to_numpy()
         else:
             weights = _zonal_face_weights(
-                faces_edge_nodes_xyz_candidate,
+                face_edge_nodes_cartesian_candidate,
                 bounds_candidate,
                 n_nodes_per_face_candidate,
                 z,

uxarray/geometry/face_edges.py

@@ -0,0 +1,193 @@
+import numpy as np
+
+from uxarray.constants import INT_FILL_VALUE
+from uxarray.grid.utils import _swap_first_fill_value_with_last
+
+
+def _construct_face_edge_nodes_cartesian(
+    face_node_conn, n_face, n_max_face_edges, node_x, node_y, node_z
+):
+    """Construct an array to hold the edge Cartesian coordinates connectivity
+    for multiple faces in a grid.
+
+    Parameters
+    ----------
+    face_node_conn : np.ndarray
+        An array of shape (n_face, n_max_face_edges) containing the node indices for each face. Accessed through `grid.face_node_connectivity.value`.
+    n_face : int
+        The number of faces in the grid. Accessed through `grid.n_face`.
+    n_max_face_edges : int
+        The maximum number of edges for any face in the grid. Accessed through `grid.n_max_face_edges`.
+    node_x : np.ndarray
+        An array of shape (n_nodes,) containing the x-coordinate values of the nodes. Accessed through `grid.node_x`.
+    node_y : np.ndarray
+        An array of shape (n_nodes,) containing the y-coordinate values of the nodes. Accessed through `grid.node_y`.
+    node_z : np.ndarray
+        An array of shape (n_nodes,) containing the z-coordinate values of the nodes. Accessed through `grid.node_z`.
+
+    Returns
+    -------
+    face_edges_cartesian : np.ndarray
+        An array of shape (n_face, n_max_face_edges, 2, 3) containing the Cartesian coordinates of the edges
+        for each face. It might contain dummy values if the grid has holes.
+
+    Examples
+    --------
+    >>> face_node_conn = np.array(
+    ...     [
+    ...         [0, 1, 2, 3, 4],
+    ...         [0, 1, 3, 4, INT_FILL_VALUE],
+    ...         [0, 1, 3, INT_FILL_VALUE, INT_FILL_VALUE],
+    ...     ]
+    ... )
+    >>> n_face = 3
+    >>> n_max_face_edges = 5
+    >>> node_x = np.array([0, 1, 1, 0, 1, 0])
+    >>> node_y = np.array([0, 0, 1, 1, 2, 2])
+    >>> node_z = np.array([0, 0, 0, 0, 1, 1])
+    >>> _get_cartesian_face_edge_nodes(
+    ...     face_node_conn, n_face, n_max_face_edges, node_x, node_y, node_z
+    ... )
+    array([[[[    0,     0,     0],
+         [    1,     0,     0]],
+
+        [[    1,     0,     0],
+         [    1,     1,     0]],
+
+        [[    1,     1,     0],
+         [    0,     1,     0]],
+
+        [[    0,     1,     0],
+         [    1,     2,     1]],
+
+        [[    1,     2,     1],
+         [    0,     0,     0]]],
+
+
+       [[[    0,     0,     0],
+         [    1,     0,     0]],
+
+        [[    1,     0,     0],
+         [    0,     1,     0]],
+
+        [[    0,     1,     0],
+         [    1,     2,     1]],
+
+        [[    1,     2,     1],
+         [    0,     0,     0]],
+
+        [[INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE],
+        [INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE]]],
+
+
+       [[[    0,     0,     0],
+         [    1,     0,     0]],
+
+        [[    1,     0,     0],
+         [    0,     1,     0]],
+
+        [[    0,     1,     0],
+         [    0,     0,     0]],
+
+        [[INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE],
+         [INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE]],
+
+        [[INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE],
+         [INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE]]]])
+    """
+
+    # face_edge_connectivity (n_face, n_edge)
+
+    # each edge should have a shape (2, 3)
+
+    # Shift node connections to create edge connections
+    face_node_conn_shift = np.roll(face_node_conn, -1, axis=1)
+
+    # Construct edge connections by combining original and shifted node connections
+    face_edge_conn = np.array([face_node_conn, face_node_conn_shift]).T.swapaxes(0, 1)
+
+    # swap the first occurrence of INT_FILL_VALUE with the last value in each sub-array
+    face_edge_conn = _swap_first_fill_value_with_last(face_edge_conn)
+
+    # Get the indices of the nodes from face_edge_conn
+    face_edge_conn_flat = face_edge_conn.reshape(-1)
+
+    valid_mask = face_edge_conn_flat != INT_FILL_VALUE
+
+    # Get the valid node indices
+    valid_edges = face_edge_conn_flat[valid_mask]
+
+    #  Create an array to hold the Cartesian coordinates of the edges
+    face_edges_cartesian = np.full(
+        (len(face_edge_conn_flat), 3), INT_FILL_VALUE, dtype=float
+    )
+
+    # Fill the array with the Cartesian coordinates of the edges
+    face_edges_cartesian[valid_mask, 0] = node_x[valid_edges]
+    face_edges_cartesian[valid_mask, 1] = node_y[valid_edges]
+    face_edges_cartesian[valid_mask, 2] = node_z[valid_edges]
+
+    return face_edges_cartesian.reshape(n_face, n_max_face_edges, 2, 3)
+
+
+def _construct_face_edge_nodes_spherical(
+    face_node_conn, n_face, n_max_face_edges, node_lon, node_lat
+):
+    """Construct an array to hold the edge latitude and longitude in radians
+    connectivity for multiple faces in a grid.
+
+    Parameters
+    ----------
+    face_node_conn : np.ndarray
+        An array of shape (n_face, n_max_face_edges) containing the node indices for each face. Accessed through `grid.face_node_connectivity.value`.
+    n_face : int
+        The number of faces in the grid. Accessed through `grid.n_face`.
+    n_max_face_edges : int
+        The maximum number of edges for any face in the grid. Accessed through `grid.n_max_face_edges`.
+    node_lon : np.ndarray
+        An array of shape (n_nodes,) containing the longitude values of the nodes in degrees. Accessed through `grid.node_lon`.
+    node_lat : np.ndarray
+        An array of shape (n_nodes,) containing the latitude values of the nodes in degrees. Accessed through `grid.node_lat`.
+
+    Returns
+    -------
+    face_edges_lonlat_rad : np.ndarray
+        An array of shape (n_face, n_max_face_edges, 2, 2) containing the latitude and longitude coordinates
+        in radians for the edges of each face. It might contain dummy values if the grid has holes.
+
+    Notes
+    -----
+    If the grid has holes, the function will return an entry of dummy value faces_edges_coordinates[i] filled with INT_FILL_VALUE.
+    """
+
+    # Convert node coordinates to radians
+    node_lon_rad = np.deg2rad(node_lon)
+    node_lat_rad = np.deg2rad(node_lat)
+
+    # Shift node connections to create edge connections
+    face_node_conn_shift = np.roll(face_node_conn, -1, axis=1)
+
+    # Construct edge connections by combining original and shifted node connections
+    face_edge_conn = np.array([face_node_conn, face_node_conn_shift]).T.swapaxes(0, 1)
+
+    # swap the first occurrence of INT_FILL_VALUE with the last value in each sub-array
+    face_edge_conn = _swap_first_fill_value_with_last(face_edge_conn)
+
+    # Get the indices of the nodes from face_edge_conn
+    face_edge_conn_flat = face_edge_conn.reshape(-1)
+
+    valid_mask = face_edge_conn_flat != INT_FILL_VALUE
+
+    # Get the valid node indices
+    valid_edges = face_edge_conn_flat[valid_mask]
+
+    # Create an array to hold the latitude and longitude in radians for the edges
+    face_edges_lonlat_rad = np.full(
+        (len(face_edge_conn_flat), 2), INT_FILL_VALUE, dtype=float
+    )
+
+    # Fill the array with the latitude and longitude in radians for the edges
+    face_edges_lonlat_rad[valid_mask, 0] = node_lon_rad[valid_edges]
+    face_edges_lonlat_rad[valid_mask, 1] = node_lat_rad[valid_edges]
+
+    return face_edges_lonlat_rad.reshape(n_face, n_max_face_edges, 2, 2)