Integrate all data variables in UxDataset.integrate()

docs/user-guide/area_calc.ipynb

@@ -758,6 +758,54 @@
     "    f\"Percentage difference between large face and sum of small faces: {percentage_diff:.4f}%\"\n",
     ")"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 8. Integrating Data Variables Over the Mesh\n",
+    "\n",
+    "Face areas are the weights used to integrate a field over the unstructured mesh. For a face-centered variable, the integral is the dot product of the face values with the face areas.\n",
+    "\n",
+    "`UxDataset.integrate()` integrates **every** grid-mapped data variable in the dataset and returns a new `UxDataset` holding each variable's integral. Variables that are not mapped to the grid are skipped."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "uxds = ux.tutorial.open_dataset(\"mpas-dyamond-30km-gradient\")\n",
+    "uxds.data_vars"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Integrate all data variables at once\n",
+    "integrals = uxds.integrate()\n",
+    "integrals"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "To integrate a single variable instead, use `UxDataArray.integrate()`:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "uxds[\"gaussian\"].integrate()"
+   ]
   }
  ],
  "metadata": {

test/core/test_dataset.py

@@ -38,7 +38,40 @@ def test_integrate(gridpath, datasetpath, mesh_constants):
     """Load a dataset and calculate integrate()."""
     uxds_var2_ne30 = ux.open_dataset(gridpath("ugrid", "outCSne30", "outCSne30.ug"), datasetpath("ugrid", "outCSne30", "outCSne30_var2.nc"))
     integrate_var2 = uxds_var2_ne30.integrate()
-    nt.assert_almost_equal(integrate_var2, mesh_constants['VAR2_INTG'], decimal=3)
+    # integrate() now returns a UxDataset with the integral of each variable
+    assert isinstance(integrate_var2, UxDataset)
+    nt.assert_almost_equal(integrate_var2["var2"].values, mesh_constants['VAR2_INTG'], decimal=3)
+
+
+def test_integrate_multiple_data_arrays(gridpath, datasetpath, mesh_constants):
+    """integrate() integrates every data variable into a new UxDataset."""
+    uxds = ux.open_dataset(gridpath("ugrid", "outCSne30", "outCSne30.ug"), datasetpath("ugrid", "outCSne30", "outCSne30_var2.nc"))
+
+    # Add a second face-centered variable: doubling the data doubles the integral
+    uxds["var2_doubled"] = uxds["var2"] * 2.0
+
+    result = uxds.integrate()
+    assert isinstance(result, UxDataset)
+    assert set(result.data_vars) == {"var2", "var2_doubled"}
+
+    nt.assert_almost_equal(result["var2"].values, mesh_constants['VAR2_INTG'], decimal=3)
+    nt.assert_almost_equal(
+        result["var2_doubled"].values, 2.0 * result["var2"].values, decimal=10
+    )
+
+
+def test_integrate_skips_non_grid_variables(gridpath, datasetpath):
+    """Variables not mapped to the grid are skipped with a warning."""
+    uxds = ux.open_dataset(gridpath("ugrid", "outCSne30", "outCSne30.ug"), datasetpath("ugrid", "outCSne30", "outCSne30_var2.nc"))
+
+    # A variable whose final dimension does not map to the grid
+    uxds["not_on_grid"] = xr.DataArray(np.arange(5.0), dims=["other_dim"])
+
+    with pytest.warns(UserWarning, match="skipped during integration"):
+        result = uxds.integrate()
+
+    assert "not_on_grid" not in result.data_vars
+    assert "var2" in result.data_vars
 
 def test_info(gridpath, datasetpath):
     """Tests custom info containing grid information."""

uxarray/core/dataset.py

@@ -570,7 +570,13 @@ def info(self, buf: IO = None, show_attrs=False) -> None:
         buf.write("\n".join(lines))
 
     def integrate(self, quadrature_rule="triangular", order=4):
-        """Integrates over all the faces of the givfen mesh.
+        """Integrates over all the faces of the given mesh, for every data
+        variable in the dataset.
+
+        Each data variable is integrated independently using
+        :meth:`UxDataArray.integrate`, and the results are collected into a new
+        :class:`UxDataset`. Variables whose final dimension does not map to the
+        grid (i.e. that are not face/node/edge centered) are skipped.
 
         Parameters
         ----------
@@ -581,45 +587,45 @@ def integrate(self, quadrature_rule="triangular", order=4):
 
         Returns
         -------
-        Calculated integral : float
+        uxds : UxDataset
+            Dataset containing the integrated value of each (grid-mapped) data
+            variable.
 
         Examples
         --------
-        Open a UXarray dataset
+        Open a UXarray dataset and integrate every data variable
 
         >>> import uxarray as ux
         >>> uxds = ux.open_dataset("grid.ug", "centroid_pressure_data_ug")
-
-        # Compute the integral
         >>> integral = uxds.integrate()
-        """
 
-        # TODO: Deprecation Warning
-        warn(
-            "This method currently only works when there is a single DataArray in this Dataset. For integration of a "
-            "single data variable, use the UxDataArray.integrate() method instead. This function will be deprecated and "
-            "replaced with one that can perform a Dataset-wide integration in a future release.",
-            DeprecationWarning,
-        )
+        Access the integral of a single variable
 
-        integral = 0.0
+        >>> integral["psi"]
+        """
+        integrated_vars = {}
+        skipped = []
 
-        face_areas = self.uxgrid.face_areas.values
+        for name, uxda in self.items():
+            try:
+                integrated_vars[name] = uxda.integrate(
+                    quadrature_rule=quadrature_rule, order=order
+                )
+            except ValueError:
+                # Variable is not mapped to the grid (or its location is not
+                # yet supported for integration); skip it rather than failing
+                # the whole dataset-wide integration.
+                skipped.append(name)
 
-        # TODO: Should we fix this requirement? Shouldn't it be applicable to
-        # TODO: all variables of dataset or a dataarray instead?
-        var_key = list(self.keys())
-        if len(var_key) > 1:
-            # warning: print message
-            print(
-                "WARNING: The dataset has more than one variable, using the first variable for integration"
+        if skipped:
+            warn(
+                "The following variables were skipped during integration because "
+                "their final dimension does not map to the grid (n_face, n_node, "
+                f"or n_edge) or is not yet supported: {skipped}.",
+                UserWarning,
             )
 
-        var_key = var_key[0]
-        face_vals = self[var_key].to_numpy()
-        integral = np.dot(face_areas, face_vals)
-
-        return integral
+        return UxDataset(integrated_vars, uxgrid=self.uxgrid)
 
     def to_array(
         self,