Apply the initial changes

Author: hongyuchen1030

uxarray/core/zonal.py

@@ -3,10 +3,6 @@
 
 
 from uxarray.grid.integrate import _zonal_face_weights, _zonal_face_weights_robust
-from uxarray.grid.utils import (
-    _get_cartesian_face_edge_nodes,
-    _get_lonlat_rad_face_edge_nodes,
-)
 
 
 
@@ -22,28 +18,8 @@ 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_edges_cartesian = _get_cartesian_face_edge_nodes(
-        uxgrid.face_node_connectivity.values,
-        uxgrid.n_face,
-        uxgrid.n_max_face_edges,
-        uxgrid.node_x.values,
-        uxgrid.node_y.values,
-        uxgrid.node_z.values,
-    )
-
-
-    faces_edges_lonlat_rad = _get_lonlat_rad_face_edge_nodes(
-        uxgrid.face_node_connectivity.values,
-        uxgrid.n_face,
-        uxgrid.n_max_face_edges,
-        uxgrid.node_lon.values,
-        uxgrid.node_lat.values,
-    )
-
-    # Check if the bounds are available
-    if "bounds" not in uxgrid._ds:
-        uxgrid._populate_bounds()
+    bounds = uxgrid.bounds.values
+    faces_edges_cartesian = uxgrid.face_edges_cartesian.values
 
     for i, lat in enumerate(latitudes):
         face_indices = uxda.uxgrid.get_faces_at_constant_latitude(lat)

uxarray/grid/geometry.py

@@ -1346,13 +1346,91 @@ def compute_temp_latlon_array(
     return temp_latlon_array
 
 
+def _populate_face_edges_cartesian(grid, return_array=False):
+    faces_edges_cartesian = _get_cartesian_face_edge_nodes(
+        grid.face_node_connectivity.values,
+        grid.n_face,
+        grid.n_max_face_edges,
+        grid.node_x.values,
+        grid.node_y.values,
+        grid.node_z.values,
+    )
+
+    faces_edges_cartesian_xarray = xr.DataArray(
+        faces_edges_cartesian,
+        dims=["n_face", "n_max_face_edges", "two", "three"],
+        attrs={
+            "cf_role": "face_edges_cartesian",
+            "_FillValue": INT_FILL_VALUE,
+            "long_name": "Provide the Cartesian coordinates of the edge for each face",
+            "start_index": INT_DTYPE(0)
+        },
+    )
+
+    if return_array:
+        return faces_edges_cartesian_xarray
+    else:
+        grid._ds["face_edges_cartesian"] = faces_edges_cartesian_xarray
+
+
+def _populate_faces_edges_cartesian(grid, return_array=False):
+    faces_edges_cartesian = _get_cartesian_face_edge_nodes(
+        grid.face_node_connectivity.values,
+        grid.n_face,
+        grid.n_max_face_edges,
+        grid.node_x.values,
+        grid.node_y.values,
+        grid.node_z.values,
+    )
+
+    faces_edges_cartesian_xarray = xr.DataArray(
+        faces_edges_cartesian,
+        dims=["n_face", "n_max_face_edges", "two", "three"],
+        attrs={
+            "cf_role": "faces_edges_cartesian",
+            "_FillValue": INT_FILL_VALUE,
+            "long_name": "Provide the Cartesian coordinates of the edge for each face",
+            "start_index": INT_DTYPE(0)
+        },
+    )
+
+    if return_array:
+        return faces_edges_cartesian_xarray
+    else:
+        grid._ds["faces_edges_cartesian"] = faces_edges_cartesian_xarray
+
+
+def _populate_faces_edges_spherical(grid, return_array=False):
+    faces_edges_lonlat_rad = _get_lonlat_rad_face_edge_nodes(
+        grid.face_node_connectivity.values,
+        grid.n_face,
+        grid.n_max_face_edges,
+        grid.node_lon.values,
+        grid.node_lat.values,
+    )
+
+    faces_edges_lonlat_rad_xarray = xr.DataArray(
+        faces_edges_lonlat_rad,
+        dims=["n_face", "n_max_face_edges", "two", "two"],
+        attrs={
+            "cf_role": "faces_edges_spherical",
+            "_FillValue": INT_FILL_VALUE,
+            "long_name": "Provide the Spherical coordinates (lon lat) in radians of the edge for each face",
+            "start_index": INT_DTYPE(0)
+        },
+    )
+
+    if return_array:
+        return faces_edges_lonlat_rad_xarray
+    else:
+        grid._ds["faces_edges_spherical"] = faces_edges_lonlat_rad_xarray
+
+
 def _populate_bounds(
-    grid,
-    faces_edges_cartesian: np.ndarray = None,
-    faces_edges_lonlat_rad: np.ndarray = None,
-    is_latlonface: bool = False,
-    is_face_GCA_list=None,
-    return_array=False
+        grid,
+        is_latlonface: bool = False,
+        is_face_GCA_list=None,
+        return_array=False
 ):
     """Populates the bounds of the grid based on the geometry of its faces,
     taking into account special conditions such as faces crossing the
@@ -1430,26 +1508,11 @@ def _populate_bounds(
     # TODO: Is it possible to have a more flexible normalization? (Avoid duplicated normalizations)
     grid.normalize_cartesian_coordinates()
 
-    # Prepare data for Numba functions
-    # TODO: remove the duplicate call for this function in both zonal average and bounds
-    if faces_edges_cartesian is None:
-        faces_edges_cartesian = _get_cartesian_face_edge_nodes(
-            grid.face_node_connectivity.values,
-            grid.n_face,
-            grid.n_max_face_edges,
-            grid.node_x.values,
-            grid.node_y.values,
-            grid.node_z.values,
-        )
+    # If the face_edges_cartesian and face_edges_lonlat_rad didn't exist,
+    # The following call will automatically populate them
+    faces_edges_cartesian = grid.faces_edges_cartesian.values
 
-    if faces_edges_lonlat_rad is None:
-        faces_edges_lonlat_rad = _get_lonlat_rad_face_edge_nodes(
-            grid.face_node_connectivity.values,
-            grid.n_face,
-            grid.n_max_face_edges,
-            grid.node_lon.values,
-            grid.node_lat.values,
-        )
+    faces_edges_lonlat_rad = grid.faces_edges_spherical.values
 
     n_nodes_per_face = grid.n_nodes_per_face.values
 

uxarray/grid/grid.py

@@ -71,6 +71,8 @@
     _grid_to_matplotlib_polycollection,
     _grid_to_matplotlib_linecollection,
     _populate_bounds,
+    _populate_faces_edges_cartesian,
+    _populate_faces_edges_spherical,
     _construct_boundary_edge_indices,
     compute_temp_latlon_array,
     _find_faces,
@@ -1444,20 +1446,24 @@ def face_areas(self, value):
         self._ds["face_areas"] = value
 
     @property
-    def face_edges_cartesian(self):
+    def faces_edges_cartesian(self):
         """Cartesian Coordinates for each Face Edge.
 
         Dimensions ``(n_face, n_max_face_edges, 2, 3)``
         """
-        pass
+        if "faces_edges_cartesian" not in self._ds:
+            _populate_faces_edges_cartesian(self)
+        return self._ds["faces_edges_cartesian"]
 
     @property
-    def face_edges_spherical(self):
-        """Latitude Longitude Bounds for each Face in radians.
+    def faces_edges_spherical(self):
+        """Latitude Longitude Coordinates for each Face in radians.
 
         Dimensions ``(n_face, n_max_face_edges, 2, 2)``
         """
-        pass
+        if "faces_edges_spherical" not in self._ds:
+            _populate_faces_edges_spherical(self)
+        return self._ds["faces_edges_spherical"]
 
     @property
     def bounds(self):
@@ -1472,6 +1478,10 @@ def bounds(self):
                     "This initial execution will be significantly longer.",
                     RuntimeWarning,
                 )
+            if "faces_edges_cartesian" not in self._ds:
+                _populate_faces_edges_cartesian(self)
+            if "faces_edges_spherical" not in self._ds:
+                _populate_faces_edges_spherical(self)
             _populate_bounds(self)
         return self._ds["bounds"]