Non-Conservative Zonal Mean

benchmarks/mpas_ocean.py

@@ -2,10 +2,10 @@
 import urllib.request
 from pathlib import Path
 
-import uxarray as ux
-
 import numpy as np
 
+import uxarray as ux
+
 current_path = Path(os.path.dirname(os.path.realpath(__file__)))
 
 data_var = 'bottomDepth'
@@ -33,15 +33,16 @@
 class DatasetBenchmark:
     """Class used as a template for benchmarks requiring a ``UxDataset`` in
     this module across both resolutions."""
-    param_names = ['resolution',]
-    params = [['480km', '120km'],]
+    param_names = ['resolution', ]
+    params = [['480km', '120km'], ]
 
     def setup(self, resolution, *args, **kwargs):
         self.uxds = ux.open_dataset(file_path_dict[resolution][0], file_path_dict[resolution][1])
 
     def teardown(self, resolution, *args, **kwargs):
         del self.uxds
 
+
 class GridBenchmark:
     """Class used as a template for benchmarks requiring a ``Grid`` in this
     module across both resolutions."""
@@ -62,6 +63,7 @@ def time_gradient(self, resolution):
     def peakmem_gradient(self, resolution):
         grad = self.uxds[data_var].gradient()
 
+
 class Integrate(DatasetBenchmark):
     def time_integrate(self, resolution):
         self.uxds[data_var].integrate()
@@ -74,12 +76,10 @@ class GeoDataFrame(DatasetBenchmark):
     param_names = DatasetBenchmark.param_names + ['exclude_antimeridian']
     params = DatasetBenchmark.params + [[True, False]]
 
-
     def time_to_geodataframe(self, resolution, exclude_antimeridian):
         self.uxds[data_var].to_geodataframe(exclude_antimeridian=exclude_antimeridian)
 
 
-
 class ConnectivityConstruction(DatasetBenchmark):
     def time_n_nodes_per_face(self, resolution):
         self.uxds.uxgrid.n_nodes_per_face
@@ -136,6 +136,7 @@ def time_nearest_neighbor_remapping(self):
     def time_inverse_distance_weighted_remapping(self):
         self.uxds_480["bottomDepth"].remap.inverse_distance_weighted(self.uxds_120.uxgrid)
 
+
 class HoleEdgeIndices(DatasetBenchmark):
     def time_construct_hole_edge_indices(self, resolution):
         ux.grid.geometry._construct_boundary_edge_indices(self.uxds.uxgrid.edge_face_connectivity)
@@ -145,6 +146,7 @@ class DualMesh(DatasetBenchmark):
     def time_dual_mesh_construction(self, resolution):
         self.uxds.uxgrid.get_dual()
 
+
 class ConstructFaceLatLon(GridBenchmark):
     def time_welzl(self, resolution):
         self.uxgrid.construct_face_centers(method='welzl')
@@ -177,9 +179,26 @@ def setup(self, resolution, lat_step):
         self.lats = np.arange(-45, 45, lat_step)
         _ = self.uxgrid.bounds
 
+class CrossSections(DatasetBenchmark):
+    param_names = DatasetBenchmark.param_names + ['n_lat']
+    params = DatasetBenchmark.params + [[1, 2, 4, 8]]
+
+    def time_constant_lat_fast(self, resolution, n_lat):
+        for lat in np.linspace(-89, 89, n_lat):
+            self.uxds.uxgrid.constant_latitude_cross_section(lat, method='fast')
     def teardown(self, resolution, lat_step):
         del self.uxgrid
 
     def time_const_lat(self, resolution, lat_step):
         for lat in self.lats:
             self.uxgrid.cross_section.constant_latitude(lat)
+
+
+class ZonalAverage(DatasetBenchmark):
+    def setup(self, resolution, *args, **kwargs):
+        self.uxds = ux.open_dataset(file_path_dict[resolution][0], file_path_dict[resolution][1])
+        bounds = self.uxds.uxgrid.bounds
+
+    def time_zonal_average(self, resolution):
+        lat_step = 10
+        self.uxds['bottomDepth'].zonal_mean(lat=(-45, 45, lat_step))

docs/api.rst

@@ -196,6 +196,7 @@ I/O & Conversion
    UxDataArray.to_polycollection
    UxDataArray.to_dataset
 
+
 UxDataset
 -----------
 
@@ -263,6 +264,8 @@ UxDataArray
    UxDataArray.plot
    UxDataArray.plot.polygons
    UxDataArray.plot.points
+   UxDataArray.plot.line
+   UxDataArray.plot.scatter
 
 UxDataset
 ~~~~~~~~~
@@ -413,16 +416,13 @@ on each face.
    UxDataArray.topological_all
    UxDataArray.topological_any
 
-
-
-Intersections
+Zonal Average
 ~~~~~~~~~~~~~
 
 .. autosummary::
    :toctree: generated/
 
-   grid.intersections.gca_gca_intersection
-   grid.intersections.gca_const_lat_intersection
+   UxDataArray.zonal_mean
 
 
 Spherical Geometry

docs/user-guide/zonal-average.ipynb

@@ -0,0 +1,218 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "dbade3aaedad9ae7",
+   "metadata": {},
+   "source": [
+    "# Zonal Averaging\n",
+    "\n",
+    "This section demonstrates how to perform Zonal Averaging using UXarray.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "185e2061bc4c75b9",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-11-22T06:45:59.229482Z",
+     "start_time": "2024-11-22T06:45:59.221543Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import uxarray as ux\n",
+    "import numpy as np\n",
+    "\n",
+    "uxds = ux.open_dataset(\n",
+    "    \"../../test/meshfiles/ugrid/outCSne30/outCSne30.ug\",\n",
+    "    \"../../test/meshfiles/ugrid/outCSne30/outCSne30_vortex.nc\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8e67d1ce4aa97e6b",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-11-22T06:45:59.882964Z",
+     "start_time": "2024-11-22T06:45:59.619647Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# set periodic_elements='exclude' for larger datasets\n",
+    "uxds[\"psi\"].plot(cmap=\"inferno\", periodic_elements=\"split\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d938d659b89bc9cb",
+   "metadata": {},
+   "source": [
+    "## What is a Zonal Average/Mean?\n",
+    "\n",
+    "A zonal average (or zonal mean) is a statistical measure that represents the average of a variable along one or more lines of constant latitude. In other words, it's the mean value calculated around the sphere at constant latitudes. \n",
+    "\n",
+    "UXarray currently implements a conservative Zonal Mean, which weights candidate faces by the length of intersection by a given line of constant latitude. \n",
+    "\n",
+    "\n",
+    "```{seealso}\n",
+    "[NCL Zonal Average](https://www.ncl.ucar.edu/Applications/zonal.shtml)\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d342f8f449543994",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-11-22T06:46:00.346066Z",
+     "start_time": "2024-11-22T06:46:00.223218Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "zonal_mean_psi = uxds[\"psi\"].zonal_mean()\n",
+    "zonal_mean_psi"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "65194a38c76e8a62",
+   "metadata": {},
+   "source": "The default latitude range is between -90 and 90 degrees with a step size of 10 degrees. "
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b5933beaf2f598ab",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-11-22T06:46:01.797976Z",
+     "start_time": "2024-11-22T06:46:01.473015Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "(zonal_mean_psi.plot.line() * zonal_mean_psi.plot.scatter(color=\"red\")).opts(\n",
+    "    title=\"Zonal Average Plot (Default)\", xticks=np.arange(-90, 100, 20), xlim=(-95, 95)\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ba2d641a3076c692",
+   "metadata": {},
+   "source": [
+    "The range of latitudes can be modified by using the `lat` parameter. It accepts:\n",
+    "\n",
+    "* **Single scalar**: e.g., `lat=45`\n",
+    "* **List/array**: e.g., `lat=[10, 20]` or `lat=np.array([10, 20])`\n",
+    "* **Tuple**: e.g., `(min_lat, max_lat, step)`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4f665827daac1c46",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-11-22T06:46:02.884242Z",
+     "start_time": "2024-11-22T06:46:02.791434Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "zonal_mean_psi_large = uxds[\"psi\"].zonal_mean(lat=(-90, 90, 1))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1998f1a55b67100c",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-11-22T06:46:03.248676Z",
+     "start_time": "2024-11-22T06:46:03.216455Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "(\n",
+    "    zonal_mean_psi_large.plot.line()\n",
+    "    * zonal_mean_psi_large.plot.scatter(color=\"red\", s=1)\n",
+    ").opts(\n",
+    "    title=\"Zonal Average Plot (Larger Sample)\",\n",
+    "    xticks=np.arange(-90, 100, 20),\n",
+    "    xlim=(-95, 95),\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4b91ebb8f733a318",
+   "metadata": {},
+   "source": [
+    "## Combined Plots\n",
+    "\n",
+    "It is often desired to plot the zonal average along side other plots, such as color or contour plots. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cb2255761173d53e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "(\n",
+    "    uxds[\"psi\"].plot(\n",
+    "        cmap=\"inferno\",\n",
+    "        periodic_elements=\"split\",\n",
+    "        height=250,\n",
+    "        width=500,\n",
+    "        colorbar=False,\n",
+    "        ylim=(-90, 90),\n",
+    "    )\n",
+    "    + zonal_mean_psi.plot.line(\n",
+    "        x=\"psi_zonal_mean\",\n",
+    "        y=\"latitudes\",\n",
+    "        height=250,\n",
+    "        width=150,\n",
+    "        ylabel=\"\",\n",
+    "        ylim=(-90, 90),\n",
+    "        xlim=(0.8, 1.2),\n",
+    "        xticks=[0.8, 0.9, 1.0, 1.1, 1.2],\n",
+    "        yticks=[-90, -45, 0, 45, 90],\n",
+    "        grid=True,\n",
+    "    )\n",
+    ").opts(title=\"Combined Zonal Average & Raster Plot\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/userguide.rst

@@ -49,6 +49,9 @@ These user guides provide detailed explanations of the core functionality in UXa
 `Cross-Sections <user-guide/cross-sections.ipynb>`_
  Select cross-sections of a grid
 
+`Zonal Means <user-guide/zonal-average.ipynb>`_
+ Compute the zonal averages across lines of constant latitude
+
 `Remapping <user-guide/remapping.ipynb>`_
  Remap (a.k.a Regrid) between unstructured grids
 
@@ -98,6 +101,7 @@ These user guides provide additional details about specific features in UXarray.
    user-guide/advanced-plotting.ipynb
    user-guide/subset.ipynb
    user-guide/cross-sections.ipynb
+   user-guide/zonal-average.ipynb
    user-guide/remapping.ipynb
    user-guide/topological-aggregations.ipynb
    user-guide/calculus.ipynb

test/test_cross_sections.py

@@ -2,6 +2,9 @@
 import pytest
 import numpy as np
 from pathlib import Path
+import os
+
+import numpy.testing as nt
 
 # Define the current path and file paths for grid and data
 current_path = Path(__file__).resolve().parent
@@ -115,3 +118,20 @@ def test_constant_lat_out_of_bounds():
     )
 
     assert len(candidate_faces) == 0
+
+
+class TestArcs:
+    def test_latitude_along_arc(self):
+        node_lon = np.array([-40, -40, 40, 40])
+        node_lat = np.array([-20, 20, 20, -20])
+        face_node_connectivity = np.array([[0, 1, 2, 3]], dtype=np.int64)
+
+        uxgrid = ux.Grid.from_topology(node_lon, node_lat, face_node_connectivity)
+
+        # intersection at exactly 20 degrees latitude
+        out1 = uxgrid.get_faces_at_constant_latitude(lat=20)
+
+        # intersection at 25.41 degrees latitude (max along the great circle arc)
+        out2 = uxgrid.get_faces_at_constant_latitude(lat=25.41)
+
+        nt.assert_array_equal(out1, out2)