return mask, not saving

Author: adebardo

doc/source/stats.md

@@ -105,16 +105,38 @@ In this example, we will create different altitude classes from a chosen interva
 Once these bins are created, we reapply them and compute the mean, minimum, and maximum values of the same raster
 for each sub-interval. It is also possible to use a reference other than the raster itself for the group_by.
 
+A dictionary containing the masks that have been created during the computation will also be returned by the function.
+Using geoUtils functions makes it very easy to visualise them.
+
 ```{code-cell} ipython3
 from geoutils.stats import grouped_stats
+import math
+import matplotlib.pyplot as plt
 
 group_by = {"rast": rast}
 bins = {"rast": [400, 1000, 2000, 3000, np.inf]}
 to_aggregate = {"rast": rast}
 statistics = ["mean", "min", "max"]
 
-df = grouped_stats.grouped_stats(group_by, bins, to_aggregate, statistics)
-df
+df, masks = grouped_stats.grouped_stats(group_by, bins, to_aggregate, statistics)
+
+groups = list(masks["groupby_rast"].keys())
+n = len(groups)
+
+ncols = 3
+nrows = math.ceil(n / ncols)
+
+fig, axes = plt.subplots(nrows, ncols, figsize=(5*ncols, 5*nrows))
+axes = axes.flatten()
+
+for ax, group in zip(axes, groups):
+    masks["groupby_rast"][group].plot(ax=ax)
+    ax.set_title(group)
+
+for ax in axes[n:]:
+    ax.axis("off")
+
+plt.tight_layout()
 ```
 
 ```{warning}
@@ -135,10 +157,14 @@ from geoutils.stats import grouped_stats
 group_by = {"rast": rast}
 elev_mask = rast > 2000
 bins = {"rast": elev_mask}
+
+elev_mask.plot()
+plt.show()
+
 to_aggregate = {"rast": rast}
 statistics = ["mean", "min", "max"]
 
-df = grouped_stats.grouped_stats(group_by, bins, to_aggregate, statistics, save_masks="")
+df, _ = grouped_stats.grouped_stats(group_by, bins, to_aggregate, statistics)
 df
 ```
 

geoutils/stats/grouped_stats.py

@@ -24,7 +24,6 @@
 from __future__ import annotations
 
 import logging
-from pathlib import Path
 from typing import Any
 
 import numpy as np
@@ -87,23 +86,20 @@ def grouped_stats(
     bins: dict[str, RasterType | list[int] | object],
     aggregated_vars: dict[str, NDArrayNum],
     statistics: list[str],
-    save_csv: str | Path | None = None,
-    save_masks: str | Path | None = None,
-) -> pd.DataFrame:
+) -> tuple[Any, dict[str, dict[str, NDArrayNum] | None]]:
     """
     Get statistics grouped (=binned) by other variables, whether categorical or continuous.
     :param groupby_vars: Dictionary containing Raster to group by.
     :param bins: Bins to use. Can be a list of interval, binary mask or segmentation map.
     :param aggregated_vars: Values to group, can be a dictionary .
     :param statistics: List or dict of statistics to compute, e.g. ["mean", "std"].
-    :param save_csv: Path to save CSV file.
-    :param save_masks: Path to save masks.
     """
 
     # Check that groupby_vars and bins as the same values
     if set(groupby_vars) != set(bins):
         raise ValueError("One bins/mask/segmentation entry per input array required.")
 
+    # Load bins
     new_bins = {}
     for key, value in bins.items():
         if isinstance(value, str):
@@ -114,72 +110,69 @@ def grouped_stats(
 
     crs, shape, transform = init_binnings_attributes(next(iter(groupby_vars.values())))
 
-    # panda dataframe works on arrays
+    # Flatten arrays for pandas
     groupby_vars_array = from_raster_to_flattened(groupby_vars)
     aggregated_vars_arrays = from_raster_to_flattened(aggregated_vars)
 
-    # Init panda data frame
     df = pd.DataFrame({**groupby_vars_array, **aggregated_vars_arrays})
+
     groupby_keys = []
-    # Analyze type of binnings
+    returned_masks = {}
+
     for groupby_key in groupby_vars.keys():
         group_col = f"groupby_{groupby_key}"
-
         bins_array = new_bins[groupby_key]
 
-        # is interval
+        if isinstance(bins[groupby_key], (str, Raster)):
+            returned_masks[group_col] = None
+        else:
+            returned_masks[group_col] = {}
+
+        # ---------- Interval bins ----------
         if is_interval(bins_array):
             cut = pd.cut(df[groupby_key], bins=bins_array)
             df[group_col] = cut
 
-            if save_masks and not isinstance(bins[groupby_key], str):
+            if returned_masks[group_col] is not None:
                 for interval in cut.cat.categories:
                     mask_flat = (cut == interval).to_numpy()
-                    mask = mask_flat.reshape(shape)
-                    outdir = Path(save_masks) / group_col
-                    outdir.mkdir(parents=True, exist_ok=True)
-                    fname = f"{groupby_key}_{interval.left}_{interval.right}.tif"
-                    Raster.from_array(mask.reshape(shape), transform=transform, crs=crs).save(outdir / fname)
+                    mask_to_raster = Raster.from_array(mask_flat.reshape(shape), transform, crs)
+
+                    returned_masks[group_col][f"{groupby_key}_{interval.left}_{interval.right}"] = mask_to_raster
 
+        # ---------- Raster (mask or segmentation) ----------
         elif isinstance(bins_array, Raster):
             data = np.asarray(bins_array.data.data).flatten()
+
             if len(data) != len(df):
                 raise ValueError(f"Mask has invalid length {len(data)}, expected {len(df)}")
-            # binary mask
+
+            # Binary mask
             if bins_array.is_mask:
                 df[group_col] = pd.Categorical(data, categories=[False, True])
-                if save_masks and not isinstance(bins[groupby_key], str):
-                    mask = data.reshape(shape).astype("uint8")
-                    fname = f"{groupby_key}_mask.tif"
-                    outdir = Path(save_masks) / group_col
-                    outdir.mkdir(parents=True, exist_ok=True)
-                    Raster.from_array(mask.reshape(shape), transform=transform, crs=crs).save(outdir / fname)
-            # segmentation
+
+                if returned_masks[group_col] is not None:
+                    mask_to_raster = Raster.from_array(data.reshape(shape).astype("uint8"), transform, crs)
+
+                    returned_masks[group_col][f"{groupby_key}_mask"] = mask_to_raster
+
+            # Segmentation
             else:
                 df[group_col] = pd.Categorical(data)
 
-                if save_masks and not isinstance(bins[groupby_key], str):
-                    labels = np.unique(data[~np.isnan(data)])
-                    for lbl in labels:
-                        mask_flat = data == lbl
-                        mask = mask_flat.reshape(shape)
+                if returned_masks[group_col] is not None:
+                    segm_names = np.unique(data[~np.isnan(data)])
+                    for segm_name in segm_names:
+                        mask_to_raster = Raster.from_array((data == segm_name).reshape(shape), transform, crs)
+
+                        returned_masks[group_col][f"{groupby_key}_seg_{int(segm_name)}"] = mask_to_raster
 
-                        fname = f"{groupby_key}_seg_{int(lbl)}.tif"
-                        outdir = Path(save_masks) / group_col
-                        outdir.mkdir(parents=True, exist_ok=True)
-                        Raster.from_array(mask.reshape(shape), transform=transform, crs=crs).save(outdir / fname)
         else:
             raise NotImplementedError("This type of bins does not yet work")
 
-        # Compute the stats
         groupby_keys.append(group_col)
 
-    # Use aggregation directly on categorical columns (fast)
-    result = df.groupby(groupby_keys)[list(aggregated_vars.keys())].agg(statistics)
-
-    if save_csv:
-        csv_file = Path(save_csv) / "grouped_stats.csv"
-        csv_file.parent.mkdir(parents=True, exist_ok=True)
-        result.reset_index().to_csv(csv_file, index=False)
+    # Compute statistics
+    result = df.groupby(groupby_keys, observed=True)[list(aggregated_vars.keys())].agg(statistics)
 
-    return result
+    return result, returned_masks

tests/test_stats/test_grouped_stats.py

@@ -1,4 +1,3 @@
-from pathlib import PosixPath
 from typing import Any, Dict, List
 
 import numpy as np
@@ -91,7 +90,7 @@ def test_is_interval() -> None:
         ([0, 5, -2, 40, np.inf], ValueError, "Values must be strictly increasing."),
     ],
 )  # type: ignore
-def test_is_interval_with_error(interval, type_error, match) -> None:
+def test_is_interval_with_error(interval: List[Any], type_error: TypeError, match: str) -> None:
 
     with pytest.raises(type_error, match=match):
         _ = grouped_stats.is_interval(interval)
@@ -104,28 +103,86 @@ def test_is_interval_with_error(interval, type_error, match) -> None:
         ({"slope3": [0, 1, 2], "slope4": [0, 1, 2]}),  # different keys
     ],
 )  # type: ignore
-def test_grouped_stats_errors(slope_raster, aggregated_vars, bins_test, statistics) -> None:
+def test_grouped_stats_errors(
+    slope_raster: Raster, aggregated_vars: Dict[str, Any], bins_test: Dict[str, Any], statistics: List[Any]
+) -> None:
     group_by_test = {"slope1": slope_raster, "slope2": slope_raster}
     with pytest.raises(ValueError, match="One bins/mask/segmentation entry per input array required."):
         _ = grouped_stats.grouped_stats(group_by_test, bins_test, aggregated_vars, statistics)
 
 
-def test_grouped_stats_interval(slope_raster: Raster, aggregated_vars: Dict[str, Any], statistics: List[Any]) -> None:
+def test_grouped_stats_interval(
+    slope_raster: Raster, aggregated_vars: Dict[str, Any], statistics: List[Any], elev_raster: Raster
+) -> None:
     group_by_test = {"slope1": slope_raster}
     expected_df = pd.DataFrame(
         {
             ("raster", "mean"): [3, 8, 15.5],
             ("raster", "min"): [1, 6, 11],
         },
         index=pd.CategoricalIndex(
-            pd.IntervalIndex.from_breaks([-0.001, 5.0, 10.0, np.inf], closed="right"),
+            pd.IntervalIndex.from_breaks([0.0, 5.0, 10.0, np.inf], closed="right"),
             ordered=True,
             name="groupby_slope1",
         ),
     )
-    df_test = grouped_stats.grouped_stats(group_by_test, {"slope1": [0, 5, 10, np.inf]}, aggregated_vars, statistics)
+
+    crs, _, transform_test = grouped_stats.init_binnings_attributes(elev_raster)
+
+    expected_mask = {
+        "groupby_slope1": {
+            "slope1_0.0_5.0": Raster.from_array(
+                np.array(
+                    [
+                        [True, True, True, True, True],
+                        [False, False, False, False, False],
+                        [False, False, False, False, False],
+                        [False, False, False, False, False],
+                    ]
+                ),
+                transform=transform_test,
+                crs=crs,
+                nodata=None,
+            ),
+            "slope1_5.0_10.0": Raster.from_array(
+                np.array(
+                    [
+                        [False, False, False, False, False],
+                        [True, True, True, True, True],
+                        [False, False, False, False, False],
+                        [False, False, False, False, False],
+                    ]
+                ),
+                transform=transform_test,
+                crs=crs,
+                nodata=None,
+            ),
+            "slope1_10.0_inf": Raster.from_array(
+                np.array(
+                    [
+                        [False, False, False, False, False],
+                        [False, False, False, False, False],
+                        [True, True, True, True, True],
+                        [True, True, True, True, True],
+                    ]
+                ),
+                transform=transform_test,
+                crs=crs,
+                nodata=None,
+            ),
+        }
+    }
+
+    df_test, masks_test = grouped_stats.grouped_stats(
+        group_by_test, {"slope1": [0, 5, 10, np.inf]}, aggregated_vars, statistics
+    )
     assert_frame_equal(df_test, expected_df)
 
+    assert all(
+        expected_mask["groupby_slope1"][group] == masks_test["groupby_slope1"][group]  # type: ignore
+        for group in expected_mask["groupby_slope1"]
+    )
+
 
 def test_grouped_stats_mask(
     slope_raster: Raster, aggregated_vars: Dict[str, Any], mask_raster: Raster, statistics: List[Any]  # type: ignore
@@ -141,62 +198,51 @@ def test_grouped_stats_mask(
         ),
     )
 
-    df_test = grouped_stats.grouped_stats(group_by_test, {"slope1": mask_raster}, aggregated_vars, statistics)
+    df_test, masks_test = grouped_stats.grouped_stats(
+        group_by_test, {"slope1": mask_raster}, aggregated_vars, statistics
+    )
     assert_frame_equal(df_test, expected_df)
 
+    expected_mask = {"groupby_slope1": None}
+
+    assert expected_mask == masks_test
+
 
 def test_grouped_stats_segm(
     slope_raster: Raster, aggregated_vars: Dict[str, Any], segm_raster: Raster, statistics: List[Any]
 ) -> None:
-    group_by_test = {"slope1": slope_raster}
+    group_by_test = {"segm": slope_raster}
 
     expected_df = pd.DataFrame(
         {
             ("raster", "mean"): [8.125000, 11.285714, 13.200000],
             ("raster", "min"): [1, 3, 4],
         },
         index=pd.CategoricalIndex(
-            [1, 2, 3], categories=[1, 2, 3], ordered=False, dtype="category", name="groupby_slope1"
+            [1, 2, 3], categories=[1, 2, 3], ordered=False, dtype="category", name="groupby_segm"
         ),
     )
-    df_test = grouped_stats.grouped_stats(group_by_test, {"slope1": segm_raster}, aggregated_vars, statistics)
+    df_test, masks_test = grouped_stats.grouped_stats(group_by_test, {"segm": segm_raster}, aggregated_vars, statistics)
     assert_frame_equal(df_test, expected_df)
 
+    expected_mask = {"groupby_segm": None}
+
+    assert expected_mask == masks_test
+
 
 def test_intersection_stats(
     slope_raster: Raster, aggregated_vars: Dict[str, Any], segm_raster: Raster, statistics: List[Any]
 ) -> None:
     group_by_test = {"slope1": slope_raster, "slope2": slope_raster}
     bins_test = {"slope1": [0, 5, 10, np.inf], "slope2": segm_raster}
 
-    df_test = grouped_stats.grouped_stats(group_by_test, bins_test, aggregated_vars, statistics)
+    df_test, _ = grouped_stats.grouped_stats(group_by_test, bins_test, aggregated_vars, statistics)
 
     assert df_test.loc[(pd.Interval(10.0, np.inf), 3), ("raster", "mean")] == (15 + 19 + 20) / 3
     assert df_test.loc[(pd.Interval(10.0, np.inf), 3), ("raster", "min")] == 15
 
-    assert df_test.loc[(pd.Interval(-0.001, 5), 1), ("raster", "mean")] == (1 + 2 + 5) / 3
-    assert df_test.loc[(pd.Interval(-0.001, 5), 1), ("raster", "min")] == 1
+    assert df_test.loc[(pd.Interval(0.0, 5), 1), ("raster", "mean")] == (1 + 2 + 5) / 3
+    assert df_test.loc[(pd.Interval(0.0, 5), 1), ("raster", "min")] == 1
 
     assert df_test.loc[(pd.Interval(10.0, np.inf), 2), ("raster", "mean")] == (16 + 12 + 17 + 18) / 4
     assert df_test.loc[(pd.Interval(10.0, np.inf), 2), ("raster", "min")] == 12
-
-
-def test_savings_csv_and_tiff(
-    slope_raster: Raster,
-    aggregated_vars: Dict[str, Any],
-    segm_raster: Raster,
-    statistics: List[Any],
-    tmp_path: PosixPath,
-) -> None:
-    group_by_test = {"slope1": slope_raster, "slope2": slope_raster}
-    bins_test = {"slope1": [0, 5, 10, np.inf], "slope2": segm_raster}
-
-    _ = grouped_stats.grouped_stats(
-        group_by_test, bins_test, aggregated_vars, statistics, save_csv=tmp_path, save_masks=tmp_path
-    )
-
-    for interval in ["0_5.0", "10.0_inf", "5.0_10.0"]:
-        assert tmp_path / "groupbyslope1" / f"slope1_{interval}.tif"
-    for seg in [1, 2, 3]:
-        assert tmp_path / "groupbyslope2" / f"slope2_{seg}.tif"
-    assert tmp_path / "grouped_stats.csv"