o robust angle calculation and enhanced documentation

Author: rajeeja

docs/user-guide/zonal-average.ipynb

@@ -17,15 +17,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "\n",
     "import uxarray as ux\n",
     "\n",
     "uxds = ux.open_dataset(\n",
     "    \"../../test/meshfiles/ugrid/outCSne30/outCSne30.ug\",\n",
     "    \"../../test/meshfiles/ugrid/outCSne30/outCSne30_vortex.nc\",\n",
-    ")\n",
-    "uxds[\"psi\"].plot(cmap=\"inferno\", periodic_elements=\"split\")"
+    ")"
    ]
   },
   {
@@ -35,10 +35,10 @@
    "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",
+    "A zonal average (or zonal mean) is a statistical measure that represents the average of a variable along lines of constant latitude or over latitudinal bands.\n",
     "\n",
     "UXarray provides two types of zonal averaging:\n",
-    "- **Non-conservative**: Weights candidate faces by the length of intersection of a line of constant latitude\n",
+    "- **Non-conservative**: Samples values at specific latitude lines\n",
     "- **Conservative**: Preserves integral quantities by weighting faces by their area overlap with latitude bands\n",
     "\n",
     "```{seealso}\n",
@@ -53,7 +53,26 @@
    "source": [
     "## Non-Conservative Zonal Averaging\n",
     "\n",
-    "The non-conservative method samples values at specific lines of constant latitude. This is the default behavior and is suitable for visualization and general analysis where exact conservation is not required."
+    "The non-conservative method samples values at specific lines of constant latitude. This is the default behavior and is suitable for visualization and general analysis where exact conservation is not required.\n",
+    "\n",
+    "Let's first visualize our data field and then demonstrate zonal averaging:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "global_field_plot",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Display the global field\n",
+    "uxds[\"psi\"].plot(cmap=\"inferno\", periodic_elements=\"split\", title=\"Global Field\")\n",
+    "\n",
+    "# Show conceptual latitude bands for zonal averaging\n",
+    "print(\"Latitude bands for zonal averaging:\")\n",
+    "lat_bands = np.arange(-90, 91, 30)\n",
+    "for i, lat in enumerate(lat_bands[:-1]):\n",
+    "    print(f\"Band {i + 1}: {lat}° to {lat_bands[i + 1]}°\")"
    ]
   },
   {
@@ -222,13 +241,13 @@
    "id": "visual_comparison",
    "metadata": {},
    "source": [
-    "### Visual Comparison: Conservative vs Non-Conservative",
-    "",
-    "The differences between methods reflect their fundamental approaches:",
-    "",
-    "- **Conservative**: More accurate for physical quantities because it accounts for the actual area of each face within latitude bands",
-    "- **Non-conservative**: Faster but approximates by sampling at specific latitude lines",
-    "",
+    "### Visual Comparison: Conservative vs Non-Conservative\n",
+    "\n",
+    "The differences between methods reflect their fundamental approaches:\n",
+    "\n",
+    "- **Conservative**: More accurate for physical quantities because it accounts for the actual area of each face within latitude bands\n",
+    "- **Non-conservative**: Faster but approximates by sampling at specific latitude lines\n",
+    "\n",
     "The differences you see indicate how much area-weighting matters for your specific data and grid resolution."
    ]
   },
@@ -239,8 +258,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import matplotlib.pyplot as plt\n",
-    "\n",
     "# Compare with non-conservative at same latitudes\n",
     "band_centers = 0.5 * (bands[:-1] + bands[1:])\n",
     "non_conservative_comparison = uxds[\"psi\"].zonal_mean(lat=band_centers)\n",
@@ -338,29 +355,35 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Create combined plot with matplotlib\n",
-    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))\n",
+    "# Display the global field\n",
+    "print(\"Global Field:\")\n",
+    "uxds[\"psi\"].plot(cmap=\"inferno\", periodic_elements=\"split\", title=\"Global Field\")\n",
     "\n",
-    "# Left: Data field\n",
-    "uxds[\"psi\"].plot(ax=ax1, cmap=\"inferno\")\n",
-    "ax1.set_title(\"Global Field\")\n",
+    "# Create zonal average plot\n",
+    "fig, ax = plt.subplots(1, 1, figsize=(8, 6))\n",
     "\n",
-    "# Right: Zonal average\n",
-    "ax2.plot(\n",
+    "ax.plot(\n",
     "    zonal_mean_psi.values,\n",
     "    zonal_mean_psi.coords[\"latitudes\"],\n",
     "    \"o-\",\n",
     "    linewidth=2,\n",
-    "    markersize=4,\n",
+    "    markersize=6,\n",
+    "    color=\"blue\",\n",
+    "    label=\"Zonal Mean\",\n",
     ")\n",
-    "ax2.set_xlabel(\"Zonal Mean Value\")\n",
-    "ax2.set_ylabel(\"Latitude (degrees)\")\n",
-    "ax2.set_title(\"Zonal Average\")\n",
-    "ax2.grid(True, alpha=0.3)\n",
-    "ax2.set_ylim(-90, 90)\n",
-    "ax2.set_xlim(0.8, 1.2)\n",
-    "\n",
-    "plt.suptitle(\"Combined Zonal Average & Raster Plot\")\n",
+    "\n",
+    "ax.set_xlabel(\"Zonal Mean Value\")\n",
+    "ax.set_ylabel(\"Latitude (degrees)\")\n",
+    "ax.set_title(\"Zonal Average\")\n",
+    "ax.grid(True, alpha=0.3)\n",
+    "ax.set_ylim(-90, 90)\n",
+    "ax.legend()\n",
+    "\n",
+    "# Add reference latitude lines\n",
+    "sample_bands = [-60, -30, 0, 30, 60]\n",
+    "for lat_band in sample_bands:\n",
+    "    ax.axhline(y=lat_band, color=\"red\", linestyle=\"--\", alpha=0.5, linewidth=1)\n",
+    "\n",
     "plt.tight_layout()\n",
     "plt.show()"
    ]

uxarray/core/zonal.py

@@ -8,7 +8,10 @@
     gca_const_lat_intersection,
     get_number_of_intersections,
 )
-from uxarray.grid.utils import _get_cartesian_face_edge_nodes_array
+from uxarray.grid.utils import (
+    _get_cartesian_face_edge_nodes_array,
+    _small_angle_of_2_vectors,
+)
 
 
 def _compute_non_conservative_zonal_mean(uxda, latitudes, use_robust_weights=False):
@@ -89,8 +92,21 @@ def _sort_points_by_angle(points):
 
     # Calculate angles (longitude)
     angles = np.empty(n_points, dtype=np.float64)
+    x_axis = np.array([1.0, 0.0, 0.0])
     for i in range(n_points):
-        angles[i] = np.arctan2(points[i, 1], points[i, 0])
+        # Project point to xy plane for longitude calculation
+        point_xy = np.array([points[i, 0], points[i, 1], 0.0])
+        point_xy_norm = np.linalg.norm(point_xy)
+
+        if point_xy_norm < 1e-15:
+            angles[i] = 0.0  # Point at pole
+        else:
+            point_xy_unit = point_xy / point_xy_norm
+            angle = _small_angle_of_2_vectors(x_axis, point_xy_unit)
+            # Determine sign based on y coordinate
+            if points[i, 1] < 0:
+                angle = -angle
+            angles[i] = angle
 
     # Simple insertion sort (numba-friendly for small arrays)
     sorted_points = points.copy()