dim arg in measures_improvement

AUTHORS.rst

@@ -23,3 +23,4 @@ Contributors
 * Yannick Rousseau
 * Marco Braun <[email protected]> `@vindelico <https://github.com/vindelico>`_
 * Sarah-Claude Bourdeau-Goulet <[email protected]> `@sarahclaude <https://github.com/sarahclaude>`_
+* Éric Dupuis <[email protected]> `@coxipi <https://github.com/coxipi>`_

CHANGELOG.rst

@@ -4,14 +4,15 @@ Changelog
 
 v0.9.2 (Unreleased)
 -------------------
-Contributors to this version: Juliette Lavoie (:user:`juliettelavoie`).
+Contributors to this version: Juliette Lavoie (:user:`juliettelavoie`), Éric Dupuis (:user:`coxipi`)
 
 
 
 Internal changes
 ^^^^^^^^^^^^^^^^
 * Include domain in `weight_location` in ``regrid_dataset``. (:pull:`414`).
 * Added pins to xarray, xclim,  h5py, and netcdf4. (:pull:`414`).
+* `xs.diagnostics.measures_improvement` now accepts `dim`, which specifies `dimension(s)` on which the proportion of improved pixels are computed.
 
 Bug fixes
 ^^^^^^^^^

src/xscen/diagnostics.py

@@ -525,7 +525,9 @@ def measures_heatmap(
 
 
 def measures_improvement(
-    meas_datasets: Union[list[xr.Dataset], dict], to_level: str = "diag-improved"
+    meas_datasets: Union[list[xr.Dataset], dict],
+    dim: str | Sequence[str] | None = None,
+    to_level: str = "diag-improved",
 ) -> xr.Dataset:
     """
     Calculate the fraction of improved grid points for each property between two datasets of measures.
@@ -536,6 +538,8 @@ def measures_improvement(
         List of 2 datasets: Initial dataset of measures and final (improved) dataset of measures.
         Both datasets must have the same variables.
         It is also possible to pass a dictionary where the values are the datasets and the key are not used.
+    dim : str or sequence of str, optional
+        Dimension(s) on which to compute the percentage of improved grid points.
     to_level: str
         processing_level to assign to the output dataset
 
@@ -554,24 +558,22 @@ def measures_improvement(
         )
     ds1 = meas_datasets[0]
     ds2 = meas_datasets[1]
+    if dim is None:
+        dims = ds1.dims
+    dims = [dim] if isinstance(dim, str) else dim
+
     percent_better = []
     for var in ds2.data_vars:
         if "xclim.sdba.measures.RATIO" in ds1[var].attrs["history"]:
             diff_bias = abs(ds1[var] - 1) - abs(ds2[var] - 1)
         else:
             diff_bias = abs(ds1[var]) - abs(ds2[var])
-        diff_bias = diff_bias.values.ravel()
-        diff_bias = diff_bias[~np.isnan(diff_bias)]
-
-        total = ds2[var].values.ravel()
-        total = total[~np.isnan(total)]
+        total_improved = (diff_bias >= 0).sum(dim=dims)
+        total_notnull = diff_bias.notnull().sum(dim=dims)
+        percent_better_var = total_improved / total_notnull
+        percent_better.append(percent_better_var.expand_dims({"properties": [var]}))
 
-        improved = diff_bias >= 0
-        percent_better.append(np.sum(improved) / len(total))
-
-    ds_better = xr.DataArray(
-        percent_better, coords={"properties": list(ds2.data_vars)}, dims="properties"
-    )
+    ds_better = xr.concat(percent_better, dim="properties")
 
     ds_better = ds_better.to_dataset(name="improved_grid_points")