Implements zonal and mean methods in UxDataset; fixes weighted_mean() dims order assumption

docs/api.rst

@@ -568,6 +568,7 @@ Azimuthal aggregations apply an aggregation (i.e. averaging) along circles of co
    :toctree: generated/
 
    UxDataArray.azimuthal_mean
+   UxDataset.azimuthal_mean
 
 
 Neighborhood
@@ -630,6 +631,9 @@ Zonal Average
    UxDataArray.zonal_average
    UxDataArray.zonal_mean
    UxDataArray.zonal_anomaly
+   UxDataset.zonal_average
+   UxDataset.zonal_mean
+   UxDataset.zonal_anomaly
 
 
 Weighted
@@ -638,6 +642,7 @@ Weighted
    :toctree: generated/
 
    UxDataArray.weighted_mean
+   UxDataset.weighted_mean
 
 
 Spherical Geometry

docs/user-guide/azimuthal-average.ipynb

@@ -138,6 +138,29 @@
     "azim_mean_psi"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "1f2b787a",
+   "metadata": {},
+   "source": [
+    "`azimuthal_mean()` can also be called directly on a UxDataset to compute the zonal mean for every face-centered variable (the example dataset here has only 1 variable, but this can be especially convenient for a dataset with many variables), although it doesn't yet support the `return_hit_counts` parameter:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "24048119",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "azim_mean_uxds = uxds.azimuthal_mean()\n",
+    "print(\n",
+    "    \"Equivalent to UxDataArray.azimuthal_mean:\",\n",
+    "    azim_mean_uxds[\"psi\"].equals(azim_mean_psi),\n",
+    ")\n",
+    "azim_mean_uxds"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "9ecf6cb2-9575-4e13-8dea-d8ec50149a7a",

docs/user-guide/weighted_mean.ipynb

@@ -20,6 +20,7 @@
     "import warnings\n",
     "\n",
     "import cartopy.crs as ccrs\n",
+    "import numpy as np\n",
     "import xarray as xr\n",
     "\n",
     "import uxarray as ux\n",
@@ -179,6 +180,70 @@
     "unweighted_result = uxds_edge[\"random_data_edge\"].mean()\n",
     "unweighted_result.values"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e32ecee9",
+   "metadata": {},
+   "source": [
+    "## Weighted Mean of a UxDataset\n",
+    "\n",
+    "It is also simple to take the weighted mean all relevant data_vars in a UxDataset (skipping anything without a grid dimension).  \n",
+    "Non-grid dimensions are also handled in the intuitive way."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "272a836a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# construct an example dataset, with nontrivial extra dimensions, vars, and coords:\n",
+    "edge_data = uxds_edge[\"random_data_edge\"].compute()  # so they are visible below.\n",
+    "face_data = uxds_face[\"random_data_face\"].compute()  # compute() to load values,\n",
+    "times = xr.DataArray([0, 1, 2], dims=\"time\", coords={\"time\": 10 * np.arange(3)})\n",
+    "uxds = ux.UxDataset(\n",
+    "    {\n",
+    "        \"edge_data\": edge_data,\n",
+    "        \"face_data\": face_data,\n",
+    "        \"face_time_data\": face_data * times,\n",
+    "        \"times\": times,\n",
+    "        \"time_plus_5\": times + 5,\n",
+    "        \"scalar_var\": 7,\n",
+    "    },\n",
+    "    coords={\"scalar_coord\": 9},\n",
+    "    uxgrid=uxds_face.uxgrid,\n",
+    ")\n",
+    "uxds"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2a371bdc",
+   "metadata": {},
+   "source": [
+    "Taking the weighted mean aggregates values along the grid dimensions (\"n_face\" and \"n_edge\") appropriately but keeps all other dimensions, variables, and coordinates unchanged."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3b5dd3f2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "result = uxds.weighted_mean()\n",
+    "result"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "78915269",
+   "metadata": {},
+   "source": [
+    "Note in particular that the relevant results here agree with corresponding results above. The result's \"face_data\" and \"edge_data\" agree with the face and edge examples above, and its \"face_time_data\" values are equal to `[0,1,2] * result[\"face_data\"]` because `uxds[\"face_time_data\"]` was constructed by multiplying the face data by `times` (==`[0,1,2]`)."
+   ]
   }
  ],
  "metadata": {

docs/user-guide/zonal-average.ipynb

@@ -22,7 +22,8 @@
     "- [3. Conservative Zonal Averaging](#3-conservative-zonal-averaging) — area-weighted bands, conservation checks, and comparisons.\n",
     "- [4. Combined Plots](#4-combined-plots) — pair global maps with their zonal means for context.\n",
     "- [5. HEALPix Zonal Averaging (Conservative vs Non-Conservative)](#5-healpix-zonal-averaging-conservative-vs-non-conservative) — run the same workflow on a different grid.\n",
-    "- [6. 2D Zonal Means on NE30 (RELHUM)](#6-2d-zonal-means-on-ne30-relhum) — build latitude–height slices and inspect the differences.\n"
+    "- [6. 2D Zonal Means on NE30 (RELHUM)](#6-2d-zonal-means-on-ne30-relhum) — build latitude–height slices and inspect the differences.\n",
+    "- [7. Zonal Anomalies](#7-zonal-anomalies) - values on each face minus zonal mean within the corresponding latitude band"
    ]
   },
   {
@@ -120,6 +121,29 @@
     "zonal_mean_psi"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "5223d24c",
+   "metadata": {},
+   "source": [
+    "`zonal_mean()` can also be called directly on a UxDataset to compute the zonal mean for every face-centered variable (the example dataset here has only 1 variable, but this can be especially convenient for a dataset with many variables):"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "231d67b1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "zonal_mean_uxds = uxds.zonal_mean()\n",
+    "print(\n",
+    "    \"Equivalent to UxDataArray.zonal_mean:\",\n",
+    "    zonal_mean_uxds[\"psi\"].equals(zonal_mean_psi),\n",
+    ")\n",
+    "zonal_mean_uxds"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "65194a38c76e8a62",
@@ -749,58 +773,104 @@
   {
    "cell_type": "markdown",
    "id": "1af6beaf",
-   "source": "## 7. Zonal Anomalies\n\nA zonal anomaly is the per-face departure from the mean of its latitude band. `zonal_anomaly` returns a `UxDataArray` with the same dims and dtype as the input (integer dtypes are promoted to float so empty bands can hold `NaN`).\n\n- **Centroid mode** (`conservative=False`, default): each face is assigned to one band by its centroid latitude (`np.digitize`). The unweighted per-band mean is exactly zero.\n- **Conservative mode** (`conservative=True`): faces straddling band edges contribute to multiple bands by area overlap (reusing the `zonal_mean` weight kernel), so per-band means are small but not exactly zero.\n\n### Step 7.1: Compute the centroid-mode anomaly",
-   "metadata": {}
+   "metadata": {},
+   "source": "## 7. Zonal Anomalies\n\nA zonal anomaly is the per-face departure from the mean of its latitude band. `zonal_anomaly` returns a `UxDataArray` with the same dims and dtype as the input (integer dtypes are promoted to float so empty bands can hold `NaN`).\n\n- **Centroid mode** (`conservative=False`, default): each face is assigned to one band by its centroid latitude (`np.digitize`). The unweighted per-band mean is exactly zero.\n- **Conservative mode** (`conservative=True`): faces straddling band edges contribute to multiple bands by area overlap (reusing the `zonal_mean` weight kernel), so per-band means are small but not exactly zero.\n\n### Step 7.1: Compute the centroid-mode anomaly"
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "c0347122",
-   "source": "anomaly = uxds[\"psi\"].zonal_anomaly(lat=(-90, 90, 10))\nanomaly",
    "metadata": {},
-   "execution_count": null,
-   "outputs": []
+   "outputs": [],
+   "source": [
+    "anomaly = uxds[\"psi\"].zonal_anomaly(lat=(-90, 90, 10))\n",
+    "anomaly"
+   ]
   },
   {
    "cell_type": "markdown",
    "id": "1783125a",
-   "source": "### Step 7.2: Verify the per-band sum-to-zero property\n\nIn centroid mode every populated band has an unweighted mean of exactly zero. We can confirm by binning faces the same way `zonal_anomaly` does (`np.digitize`) and reducing each band.",
-   "metadata": {}
+   "metadata": {},
+   "source": "### Step 7.2: Verify the per-band sum-to-zero property\n\nIn centroid mode every populated band has an unweighted mean of exactly zero. We can confirm by binning faces the same way `zonal_anomaly` does (`np.digitize`) and reducing each band."
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "ed0f731f",
-   "source": "bands = np.arange(-90, 91, 10)\nface_lat = uxds.uxgrid.face_lat.values\nband_idx = np.clip(np.digitize(face_lat, bands) - 1, 0, len(bands) - 2)\n\nfor bi in range(len(bands) - 1):\n    mask = band_idx == bi\n    if not mask.any():\n        continue\n    band_mean = float(anomaly.values[mask].mean())\n    print(\n        f\"band [{bands[bi]:+4d}, {bands[bi + 1]:+4d})  \"\n        f\"n_faces={int(mask.sum()):4d}  mean={band_mean:+.2e}\"\n    )",
    "metadata": {},
-   "execution_count": null,
-   "outputs": []
+   "outputs": [],
+   "source": [
+    "bands = np.arange(-90, 91, 10)\n",
+    "face_lat = uxds.uxgrid.face_lat.values\n",
+    "band_idx = np.clip(np.digitize(face_lat, bands) - 1, 0, len(bands) - 2)\n",
+    "\n",
+    "for bi in range(len(bands) - 1):\n",
+    "    mask = band_idx == bi\n",
+    "    if not mask.any():\n",
+    "        continue\n",
+    "    band_mean = float(anomaly.values[mask].mean())\n",
+    "    print(\n",
+    "        f\"band [{bands[bi]:+4d}, {bands[bi + 1]:+4d})  \"\n",
+    "        f\"n_faces={int(mask.sum()):4d}  mean={band_mean:+.2e}\"\n",
+    "    )"
+   ]
   },
   {
    "cell_type": "markdown",
    "id": "41dbb45d",
-   "source": "### Step 7.3: Visualize the anomaly field\n\nPlot the anomaly map and the zonal mean it was subtracted from side by side using a diverging colormap centered at zero.",
-   "metadata": {}
+   "metadata": {},
+   "source": "### Step 7.3: Visualize the anomaly field\n\nPlot the anomaly map and the zonal mean it was subtracted from side by side using a diverging colormap centered at zero."
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "84cd278e",
-   "source": "vmax = float(np.nanmax(np.abs(anomaly.values)))\nanomaly_map = anomaly.plot(\n    cmap=\"RdBu_r\",\n    periodic_elements=\"split\",\n    clim=(-vmax, vmax),\n    title=\"Zonal Anomaly (psi - zonal mean)\",\n).opts(width=525, height=400, colorbar=True)\n\nzm = uxds[\"psi\"].zonal_mean(lat=(-90, 90, 10))\nzm_df = zm.to_dataframe(name=\"zonal_mean\").reset_index()\nzm_panel = zm_df.hvplot.line(\n    x=\"zonal_mean\",\n    y=\"latitudes\",\n    line_width=2,\n    title=\"Zonal Mean (subtracted)\",\n    ylim=(-90, 90),\n    width=400,\n    height=400,\n).opts(show_grid=True)\n\n(anomaly_map + zm_panel).cols(2)",
    "metadata": {},
-   "execution_count": null,
-   "outputs": []
+   "outputs": [],
+   "source": [
+    "vmax = float(np.nanmax(np.abs(anomaly.values)))\n",
+    "anomaly_map = anomaly.plot(\n",
+    "    cmap=\"RdBu_r\",\n",
+    "    periodic_elements=\"split\",\n",
+    "    clim=(-vmax, vmax),\n",
+    "    title=\"Zonal Anomaly (psi - zonal mean)\",\n",
+    ").opts(width=525, height=400, colorbar=True)\n",
+    "\n",
+    "zm = uxds[\"psi\"].zonal_mean(lat=(-90, 90, 10))\n",
+    "zm_df = zm.to_dataframe(name=\"zonal_mean\").reset_index()\n",
+    "zm_panel = zm_df.hvplot.line(\n",
+    "    x=\"zonal_mean\",\n",
+    "    y=\"latitudes\",\n",
+    "    line_width=2,\n",
+    "    title=\"Zonal Mean (subtracted)\",\n",
+    "    ylim=(-90, 90),\n",
+    "    width=400,\n",
+    "    height=400,\n",
+    ").opts(show_grid=True)\n",
+    "\n",
+    "(anomaly_map + zm_panel).cols(2)"
+   ]
   },
   {
    "cell_type": "markdown",
    "id": "15c921d4",
-   "source": "### Step 7.4: Compare centroid vs conservative anomaly\n\nConservative mode blends straddling faces across band edges, so its anomaly differs slightly from the centroid version. The difference shows where face geometry crosses band boundaries.",
-   "metadata": {}
+   "metadata": {},
+   "source": "### Step 7.4: Compare centroid vs conservative anomaly\n\nConservative mode blends straddling faces across band edges, so its anomaly differs slightly from the centroid version. The difference shows where face geometry crosses band boundaries."
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "36a9096c",
-   "source": "anomaly_cons = uxds[\"psi\"].zonal_anomaly(lat=(-90, 90, 10), conservative=True)\ndiff = anomaly_cons - anomaly\n\nprint(f\"max |centroid|     = {float(np.nanmax(np.abs(anomaly.values))):.4f}\")\nprint(f\"max |conservative| = {float(np.nanmax(np.abs(anomaly_cons.values))):.4f}\")\nprint(f\"max |difference|   = {float(np.nanmax(np.abs(diff.values))):.4f}\")",
    "metadata": {},
-   "execution_count": null,
-   "outputs": []
+   "outputs": [],
+   "source": [
+    "anomaly_cons = uxds[\"psi\"].zonal_anomaly(lat=(-90, 90, 10), conservative=True)\n",
+    "diff = anomaly_cons - anomaly\n",
+    "\n",
+    "print(f\"max |centroid|     = {float(np.nanmax(np.abs(anomaly.values))):.4f}\")\n",
+    "print(f\"max |conservative| = {float(np.nanmax(np.abs(anomaly_cons.values))):.4f}\")\n",
+    "print(f\"max |difference|   = {float(np.nanmax(np.abs(diff.values))):.4f}\")"
+   ]
   }
  ],
  "metadata": {