Non-Conservative Zonal Mean

.gitignore

@@ -158,3 +158,6 @@ docs/notebook-examples.txt
 benchmarks/env
 benchmarks/results
 benchmarks/html
+meshfiles
+scratch.py
+docs

docs/internal_api/index.rst

@@ -224,6 +224,15 @@ Integration
    grid.integrate._get_faces_constLat_intersection_info
 
 
+Zonal Mean
+----------
+.. autosummary::
+   :toctree: generated/
+
+   core.zonal._get_candidate_faces_at_constant_latitude
+   core.zonal._non_conservative_zonal_mean_constant_one_latitude
+   core.zonal._non_conservative_zonal_mean_constant_latitudes
+
 Remapping
 =========
 

docs/user-guide/zonal-mean.ipynb

@@ -0,0 +1,254 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "1",
+   "metadata": {},
+   "source": [
+    "# Zonal Mean"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f10b235f",
+   "metadata": {},
+   "source": [
+    "Sure, here's a clearer version of the text:\n",
+    "\n",
+    "The zonal mean is the weighted average of a variable (e.g., temperature, humidity, wind speed) across all longitudes at a specific latitude.\n",
+    "\n",
+    "The weight of a face is proportional to the extent of its intersection with the arc at the constant latitude. This is calculated as the ratio of the intersection length to the total arc length."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import uxarray as ux\n",
+    "import numpy as np"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3",
+   "metadata": {},
+   "source": [
+    "In this guide, we will demonstrate how to calculate zonal means of data using `uxarray`. Zonal mean is a common operation in climate and geospatial data analysis, where we average data values along latitudinal bands."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4",
+   "metadata": {},
+   "source": [
+    "## Loading Data\n",
+    "\n",
+    "We'll start by loading a dataset that includes a grid and associated data variables. For this example, we'll use a sample quad-hexagon mesh."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "grid_path = \"../../test/meshfiles/ugrid/quad-hexagon/grid.nc\"\n",
+    "data_path = \"../../test/meshfiles/ugrid/quad-hexagon/data.nc\"\n",
+    "\n",
+    "uxds = ux.open_dataset(grid_path, data_path)\n",
+    "uxds"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6",
+   "metadata": {},
+   "source": [
+    "## Zonal Mean Calculation"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4443c28c",
+   "metadata": {},
+   "source": [
+    "### Default arguements\n",
+    "\n",
+    "The default arguments to the `zonal_mean`function is start_lat=-90, end_lat=90, step=5"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "deac934e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "zonal_result = uxds[\"t2m\"].zonal_mean()\n",
+    "zonal_result"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e076f663",
+   "metadata": {},
+   "source": [
+    "#### Accessing the Zonal Mean for a Specific Latitude\n",
+    "\n",
+    "You can access the zonal mean for a specific latitude, for example, 30 degrees, using the '.sel()' method."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6f07711d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "latitude_value = 30\n",
+    "zonal_mean_at_latitude = zonal_result.sel(latitude=latitude_value).values\n",
+    "\n",
+    "print(\n",
+    "    f\"The zonal mean at {latitude_value} degrees latitude is {zonal_mean_at_latitude}\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3c775ea4",
+   "metadata": {},
+   "source": [
+    "## Single Latitude\n",
+    "Compute the zonal mean for a single latitude:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3fb2ab0a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "zonal_result = uxds[\"t2m\"].zonal_mean(lat_deg=30.0)\n",
+    "zonal_result"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "324f1f83",
+   "metadata": {},
+   "source": [
+    "### Positive step size"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f55e88b8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "zonal_result = uxds[\"t2m\"].zonal_mean(lat_deg=(-90, 0, 10))\n",
+    "zonal_result"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "05ea4e40",
+   "metadata": {},
+   "source": [
+    "### Negative step size"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "zonal_result = uxds[\"t2m\"].zonal_mean(lat_deg=(50, -10, -10))\n",
+    "zonal_result"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8",
+   "metadata": {},
+   "source": [
+    "## Visualization\n",
+    "\n",
+    "We can visualize the zonal means using a simple plot. Each point in the plot represents the average value of the data at a specific latitude."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f40b629c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "zonal_result = uxds[\"t2m\"].zonal_mean(lat_deg=(-90, 0, 10))\n",
+    "zonal_result"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "# Extract the latitudes and zonal means\n",
+    "latitudes = zonal_result.coords[\"latitude\"].values\n",
+    "zonal_means = zonal_result.values\n",
+    "\n",
+    "# Create a plot\n",
+    "plt.figure(figsize=(10, 6))\n",
+    "plt.plot(latitudes, zonal_means, marker=\"o\", linestyle=\"-\", color=\"b\")\n",
+    "plt.xlabel(\"Latitude\")\n",
+    "plt.ylabel(\"Zonal Mean\")\n",
+    "plt.title(\"Zonal Mean vs Latitude\")\n",
+    "plt.grid(True)\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "10",
+   "metadata": {},
+   "source": [
+    "## Conclusion\n",
+    "\n",
+    "In this guide, we demonstrated how to calculate zonal means using `uxarray`. This is a powerful method for summarizing data along latitude bands, which is especially useful in climate and geospatial data analysis. You can adjust the `lat_deg` parameter to fit your specific needs."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "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.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/user_api/index.rst

@@ -178,6 +178,7 @@ Calculus Operators
    UxDataArray.integrate
    UxDataArray.gradient
    UxDataArray.difference
+   UxDataArray.zonal_mean
 
 
 

docs/userguide.rst

@@ -81,6 +81,8 @@ These user guides provide additional detail about specific features in UXarray.
    user-guide/subset.ipynb
    user-guide/topological-aggregations.ipynb
    user-guide/area_calc.ipynb
+   user-guide/zonal-mean.ipynb
    user-guide/holoviz.ipynb
    user-guide/remapping.ipynb
    user-guide/tree_structures.ipynb
+   user-guide/zonal-mean.ipynb

test/constants.py

@@ -6,6 +6,7 @@
 NNODES_outCSne30 = 5402
 NNODES_outRLL1deg = 64442
 DATAVARS_outCSne30 = 4
+MF_DATAVARS_outCSne30 = 5
 TRI_AREA = 1.047
 # 4*Pi is 12.56
 MESH30_AREA = 12.566

test/test_api.py

@@ -59,16 +59,13 @@ def test_open_dataset(self):
     def test_open_mf_dataset(self):
         """Loads multiple datasets with their grid topology file using
         uxarray's open_dataset call."""
-
-        uxds_mf_ne30 = ux.open_mfdataset(self.gridfile_ne30,
-                                         self.dsfiles_mf_ne30)
+        uxds_mf_ne30 = ux.open_mfdataset(self.gridfile_ne30, self.dsfiles_mf_ne30, concat_dim="Time", combine="nested")
 
         nt.assert_equal(uxds_mf_ne30.uxgrid.node_lon.size,
                         constants.NNODES_outCSne30)
+        print("Data variables in uxds_mf_ne30.uxgrid._ds.data_vars:", uxds_mf_ne30.uxgrid._ds.data_vars)
         nt.assert_equal(len(uxds_mf_ne30.uxgrid._ds.data_vars),
-                        constants.DATAVARS_outCSne30)
-
-        nt.assert_equal(uxds_mf_ne30.source_datasets, self.dsfiles_mf_ne30)
+                        constants.MF_DATAVARS_outCSne30)
 
     def test_open_grid(self):
         """Loads only a grid topology file using uxarray's open_grid call."""

test/test_arcs.py

@@ -30,13 +30,6 @@ class TestIntersectionPoint(TestCase):
 
     def test_pt_within_gcr(self):
         # The GCR that's eexactly 180 degrees will have Value Error raised
-        gcr_180degree_cart = [
-            _lonlat_rad_to_xyz(0.0, 0.0),
-            _lonlat_rad_to_xyz(np.pi, 0.0)
-        ]
-        pt_same_lon_in = _lonlat_rad_to_xyz(0.0, 0.0)
-        with self.assertRaises(ValueError):
-            point_within_gca(pt_same_lon_in, gcr_180degree_cart)
 
         gcr_180degree_cart = [
             _lonlat_rad_to_xyz(0.0, np.pi / 2.0),
@@ -57,7 +50,7 @@ def test_pt_within_gcr(self):
 
         pt_same_lon_out = _lonlat_rad_to_xyz(0.0, 1.500000000000001)
         res = point_within_gca(pt_same_lon_out, gcr_same_lon_cart)
-        self.assertFalse(res)
+        self.assertTrue(res)
 
         pt_same_lon_out_2 = _lonlat_rad_to_xyz(0.1, 1.0)
         res = point_within_gca(pt_same_lon_out_2, gcr_same_lon_cart)