Non-Conservative Zonal Mean

docs/userguide.rst

@@ -16,3 +16,4 @@ common tasks that you can accomplish with UXarray.
    user-guide/grid-formats.rst
    user-guide/data-structures.ipynb
    user-guide/area_calc.ipynb
+   user-guide/zonal-mean.ipynb

test/test_grid.py

@@ -913,6 +913,7 @@ def test_from_face_vertices(self):
 
 
 class TestLatlonBounds(TestCase):
+    gridfile_mpas = current_path / "meshfiles" / "mpas" / "QU" / "oQU480.231010.nc"
     def test_populate_bounds_GCA_mix(self):
         face_1 = [[10.0, 60.0], [10.0, 10.0], [50.0, 10.0], [50.0, 60.0]]
         face_2 = [[350, 60.0], [350, 10.0], [50.0, 10.0], [50.0, 60.0]]
@@ -932,3 +933,9 @@ def test_populate_bounds_GCA_mix(self):
         bounds_xarray = grid.bounds
         face_bounds = bounds_xarray.values
         nt.assert_allclose(grid.bounds.values, expected_bounds, atol=ERROR_TOLERANCE)
+
+    def test_populate_bounds_MPAS(self):
+        xrds = xr.open_dataset(self.gridfile_mpas)
+        uxgrid = ux.Grid.from_dataset(xrds, use_dual=True)
+        bounds_xarray = uxgrid.bounds
+        pass

test/test_helpers.py

@@ -241,34 +241,6 @@ def test_angle_of_2_vectors(self):
 
 class TestFaceEdgeConnectivityHelper(TestCase):
 
-    def test_get_cartesian_face_edge_nodes(self):
-        # Load the dataset
-
-        uxds = ux.open_dataset(gridfile_geoflowsmall_grid,
-                               gridfile_geoflowsmall_var)
-
-        # Initialize array
-
-        face_edges_connectivity_cartesian = []
-
-        # Get the connectivity
-
-        for i in range(len(uxds.uxgrid.face_node_connectivity)):
-            face_edges_connectivity_cartesian.append(
-                _get_cartesian_face_edge_nodes(
-                    uxds.uxgrid.face_node_connectivity.values[i],
-                    uxds.uxgrid.face_edge_connectivity.values[i],
-                    uxds.uxgrid.edge_node_connectivity.values,
-                    uxds.uxgrid.node_x.values, uxds.uxgrid.node_y.values,
-                    uxds.uxgrid.node_z.values))
-
-        # Stack the arrays to get the desired (3,3) array
-
-        face_edges_connectivity_cartesian = np.vstack(
-            face_edges_connectivity_cartesian)
-
-        assert (face_edges_connectivity_cartesian.ndim == 3)
-
     def test_get_cartesian_face_edge_nodes_pipeline(self):
         # Create the vertices for the grid, based around the North Pole
 
@@ -281,14 +253,13 @@ def test_get_cartesian_face_edge_nodes_pipeline(self):
         # Construct the grid from the vertices
         grid = ux.Grid.from_face_vertices(vertices, latlon=False)
         face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes(
-            grid.face_node_connectivity.values[0],
-            grid.face_edge_connectivity.values[0],
-            grid.edge_node_connectivity.values, grid.node_x.values,
+            grid.face_node_connectivity.values,
+            grid.node_x.values,
             grid.node_y.values, grid.node_z.values)
 
         # Check that the face_edges_connectivity_cartesian works as an input to _pole_point_inside_polygon
         result = ux.grid.geometry._pole_point_inside_polygon(
-            'North', face_edges_connectivity_cartesian)
+            'North', face_edges_connectivity_cartesian[0])
 
         # Assert that the result is True
 
@@ -307,14 +278,12 @@ def test_get_cartesian_face_edge_nodes_filled_value(self):
         # Construct the grid from the vertices
         grid = ux.Grid.from_face_vertices(vertices, latlon=False)
         face_edges_connectivity_cartesian = _get_cartesian_face_edge_nodes(
-            grid.face_node_connectivity.values[0],
-            grid.face_edge_connectivity.values[0],
-            grid.edge_node_connectivity.values, grid.node_x.values,
+            grid.face_node_connectivity.values, grid.node_x.values,
             grid.node_y.values, grid.node_z.values)
 
         # Check that the face_edges_connectivity_cartesian works as an input to _pole_point_inside_polygon
         result = ux.grid.geometry._pole_point_inside_polygon(
-            'North', face_edges_connectivity_cartesian)
+            'North', face_edges_connectivity_cartesian[0])
 
         # Assert that the result is True
         self.assertTrue(result)
@@ -331,10 +300,9 @@ def test_get_lonlat_face_edge_nodes_pipeline(self):
         # Construct the grid from the vertices
         grid = ux.Grid.from_face_vertices(vertices, latlon=False)
         face_edges_connectivity_lonlat = _get_lonlat_rad_face_edge_nodes(
-            grid.face_node_connectivity.values[0],
-            grid.face_edge_connectivity.values[0],
-            grid.edge_node_connectivity.values, grid.node_lon.values,
+            grid.face_node_connectivity.values,grid.node_lon.values,
             grid.node_lat.values)
+        face_edges_connectivity_lonlat = face_edges_connectivity_lonlat[0]
 
         # Convert all the values into cartesian coordinates
         face_edges_connectivity_cartesian = []

uxarray/core/dataarray.py

@@ -7,6 +7,7 @@
 
 from uxarray.grid import Grid
 import uxarray.core.dataset
+from uxarray.grid.utils import _get_cartesian_face_edge_nodes
 
 if TYPE_CHECKING:
     from uxarray.core.dataset import UxDataset
@@ -21,6 +22,10 @@
     _calculate_edge_node_difference,
 )
 
+from uxarray.core.zonal import (
+    _non_conservative_zonal_mean_constant_latitudes,
+)
+
 from uxarray.plot.accessor import UxDataArrayPlotAccessor
 from uxarray.subset import DataArraySubsetAccessor
 from uxarray.remap import UxDataArrayRemapAccessor
@@ -576,7 +581,66 @@ def difference(self, destination: Optional[str] = "edge"):
 
         return uxda
 
-        pass
+    def zonal_mean(self, step_size=1):
+        """Computes the Zonal Mean for face-centered data. The zonal average is
+        computed from -90 to 90 degrees latitude, with a given step size.
+
+        Parameters
+        ----------
+        step_size : float, default=1
+            The step size for the zonal average.
+
+        Returns
+        -------
+        UxDataArray
+            UxDataArray containing the zonal average of the data variable
+
+        Example
+        -------
+        >>> uxds['var'].zonal_mean()
+        """
+
+        # Check if the data is face-centered
+        if not self._face_centered():
+            raise NotImplementedError(
+                "Zonal average computations are currently only supported for face-centered data variables."
+            )
+
+        # Get the data, face bounds, face node connectivity, and whether the faces are latlon
+        data = self.values
+        face_bounds = self.uxgrid.bounds.values
+        is_latlonface = False  # Currently not used, but may be useful in the future
+
+        # Get the list of face polygon represented by edges in Cartesian coordinates
+        face_edges_cart = _get_cartesian_face_edge_nodes(
+            self.uxgrid.face_node_connectivity.values,
+            self.uxgrid.node_x.values,
+            self.uxgrid.node_y.values,
+            self.uxgrid.node_z.values,
+        )
+
+        _zonal_avg_res = _non_conservative_zonal_mean_constant_latitudes(
+            face_edges_cart,
+            face_bounds,
+            data,
+            step_size,
+            is_latlonface=is_latlonface,
+        )
+
+        # Set Dimension of result
+        dims = list(self.dims[:-1]) + ["latitude"]
+
+        # Result is stored and returned as a UxDataArray
+        uxda = UxDataArray(
+            _zonal_avg_res,
+            uxgrid=self.uxgrid,
+            dims=dims,
+            name=self.name + "_zonal_average"
+            if self.name is not None
+            else "zonal_average",
+        )
+
+        return uxda
 
     def _face_centered(self) -> bool:
         """Returns whether the data stored is Face Centered (i.e. contains the

uxarray/core/zonal.py

@@ -0,0 +1,157 @@
+import numpy as np
+from uxarray.constants import INT_FILL_VALUE
+from uxarray.grid.integrate import _get_zonal_faces_weight_at_constLat
+
+
+def _get_candidate_faces_at_constant_latitude(bounds, constLat: float) -> np.ndarray:
+    """Return the indices of the faces whose latitude bounds contain the
+    constant latitude.
+
+    Parameters
+    ----------
+    bounds : np.ndarray, shape (n_face, 2, 2)
+        The latitude and longitude bounds of the faces.
+
+    constLat : float
+        The constant latitude to check against. Expected range is [-90, 90].
+
+    Returns
+    -------
+    np.ndarray, shape (n_candidate_faces, )
+        An array of indices of the faces whose latitude bounds contain the constant latitude `constLat`.
+    """
+
+    # Check if the constant latitude is within the range of [-90, 90]
+    if constLat < -90 or constLat > 90:
+        raise ValueError("The constant latitude must be within the range of [-90, 90].")
+
+    # Extract the latitude bounds
+    lat_bounds_min = bounds[:, 0, 0]  # Minimum latitude bound
+    lat_bounds_max = bounds[:, 0, 1]  # Maximum latitude bound
+
+    # Check if the constant latitude is within the bounds of each face
+    within_bounds = (lat_bounds_min <= constLat) & (lat_bounds_max >= constLat)
+
+    # Get the indices of faces where the condition is True
+    candidate_faces = np.where(within_bounds)[0]
+
+    return candidate_faces
+
+
+def _non_conservative_zonal_mean_constant_one_latitude(
+    face_edges_cart: np.ndarray,
+    face_bounds: np.ndarray,
+    face_data: np.ndarray,
+    constLat: float,
+    is_latlonface=False,
+) -> float:
+    """Helper function for _non_conservative_zonal_mean_constant_latitudes.
+    Calculate the zonal mean of the data at a constant latitude.
+
+    Parameters
+    ----------
+    face_edges_cart : np.ndarray, shape (n_face, n_edge, 2, 3)
+        The Cartesian coordinates of the face edges.
+    bounds : np.ndarray, shape (n_face, 2, 2)
+        The latitude and longitude bounds of the faces.
+    face_data : np.ndarray, shape (..., n_face)
+        The data on the faces.
+    constLat : float
+        The constant latitude in degrees. Expected range is [-90, 90].
+    is_latlonface : bool, optional
+        A flag indicating if the current face is a latitudinal/longitudinal (latlon) face,
+        meaning its edges align with lines of constant latitude or longitude. If `True`,
+        edges are treated as following constant latitudinal or longitudinal lines. If `False`,
+        edges are considered as great circle arcs (GCA). Default is `False`.
+
+    Returns
+    -------
+    float
+        The zonal mean of the data at the constant latitude.
+    """
+
+    # Get the indices of the faces whose latitude bounds contain the constant latitude
+    candidate_faces_indices = _get_candidate_faces_at_constant_latitude(
+        face_bounds, constLat
+    )
+
+    # Check if there are no candidate faces,
+    if len(candidate_faces_indices) == 0:
+        return INT_FILL_VALUE  # TODO: Determin an appropriate dummy value
+
+    # Get the face data of the candidate faces
+    candidate_face_data = face_data[..., candidate_faces_indices]
+
+    # Get the list of face polygon represented by edges in Cartesian coordinates
+    candidate_face_edges_cart = face_edges_cart[candidate_faces_indices]
+
+    weight_df = _get_zonal_faces_weight_at_constLat(
+        candidate_face_edges_cart,
+        np.sin(np.deg2rad(constLat)),  # Latitude in cartesian coordinates
+        face_bounds,
+        is_directed=False,
+        is_latlonface=is_latlonface,
+    )
+
+    # Compute the zonal mean(weighted average) of the candidate faces
+    weights = weight_df["weight"].values
+    zonal_mean = np.sum(candidate_face_data * weights) / np.sum(weights)
+
+    return zonal_mean
+
+
+def _non_conservative_zonal_mean_constant_latitudes(
+    face_edges_cart: np.ndarray,
+    face_bounds: np.ndarray,
+    face_data: np.ndarray,
+    step_size: float,
+    is_latlonface=False,
+) -> np.ndarray:
+    """Calculate the zonal mean of the data from -90 to 90 degrees latitude,
+    with a given step size.
+
+    Parameters
+    ----------
+    face_edges_cart : np.ndarray, shape (n_face, n_edge, 2, 3)
+        The Cartesian coordinates of the face edges.
+    bounds : np.ndarray, shape (n_face, 2, 2)
+        The latitude and longitude bounds of the faces.
+    face_data : np.ndarray, shape (..., n_face)
+        The data on the faces. It may have multiple non-grid dimensions (e.g., time, level).
+    step_size : float
+        The step size in degrees for the latitude.
+    is_latlonface : bool, optional
+        A flag indicating if the current face is a latitudinal/longitudinal (latlon) face,
+        meaning its edges align with lines of constant latitude or longitude. If `True`,
+        edges are treated as following constant latitudinal or longitudinal lines. If `False`,
+        edges are considered as great circle arcs (GCA). Default is `False`.
+
+    Returns
+    -------
+    np.ndarray
+        The zonal mean of the data from -90 to 90 degrees latitude. The shape of the output
+        is (..., n_latitudes), where n_latitudes is the number of latitude steps from -90 to 90.
+    """
+
+    # Generate latitudes from -90 to 90 with the given step size
+    latitudes = np.arange(-90, 90 + step_size, step_size)
+
+    # Initialize an empty list to store the zonal mean for each latitude
+    zonal_means = []
+
+    # Calculate the zonal mean for each latitude
+    for constLat in latitudes:
+        zonal_mean = _non_conservative_zonal_mean_constant_one_latitude(
+            face_edges_cart, face_bounds, face_data, constLat, is_latlonface
+        )
+        zonal_means.append(zonal_mean)
+
+    # Convert the list of zonal means to a NumPy array
+    zonal_means = np.array(zonal_means)
+
+    # Reshape the zonal mean array to have the same leading dimensions as the input data
+    # and an additional dimension for the latitudes
+    expected_shape = face_data.shape[:-1] + (len(latitudes),)
+    zonal_means = zonal_means.reshape(expected_shape)
+
+    return zonal_means

uxarray/grid/geometry.py

@@ -881,23 +881,31 @@ def _populate_bounds(
     intervals_tuple_list = []
     intervals_name_list = []
 
+    faces_edges_cartesian = _get_cartesian_face_edge_nodes(
+        grid.face_node_connectivity.values,
+        grid.node_x.values,
+        grid.node_y.values,
+        grid.node_z.values,
+    )
+
+    faces_edges_lonlat_rad = _get_lonlat_rad_face_edge_nodes(
+        grid.face_node_connectivity.values,
+        grid.node_lon.values,
+        grid.node_lat.values,
+    )
+
     for face_idx, face_nodes in enumerate(grid.face_node_connectivity):
-        face_edges_cartesian = _get_cartesian_face_edge_nodes(
-            grid.face_node_connectivity.values[face_idx],
-            grid.face_edge_connectivity.values[face_idx],
-            grid.edge_node_connectivity.values,
-            grid.node_x.values,
-            grid.node_y.values,
-            grid.node_z.values,
-        )
+        face_edges_cartesian = faces_edges_cartesian[face_idx]
 
-        face_edges_lonlat_rad = _get_lonlat_rad_face_edge_nodes(
-            grid.face_node_connectivity.values[face_idx],
-            grid.face_edge_connectivity.values[face_idx],
-            grid.edge_node_connectivity.values,
-            grid.node_lon.values,
-            grid.node_lat.values,
-        )
+        # Skip processing if the face is a dummy face
+        if np.any(face_edges_cartesian == INT_FILL_VALUE):
+            continue
+
+        face_edges_lonlat_rad = faces_edges_lonlat_rad[face_idx]
+
+        # Skip processing if the face is a dummy face
+        if np.any(face_edges_lonlat_rad == INT_FILL_VALUE):
+            continue
 
         is_GCA_list = (
             is_face_GCA_list[face_idx] if is_face_GCA_list is not None else None