diff --git a/geoutils/stats/stats.py b/geoutils/stats/stats.py index 1ff6d6e67..1379e68a5 100644 --- a/geoutils/stats/stats.py +++ b/geoutils/stats/stats.py @@ -24,6 +24,7 @@ from typing import Any import numpy as np +import pandas as pd from scipy.stats import iqr from geoutils._typing import NDArrayNum @@ -120,7 +121,7 @@ def _statistics(data: NDArrayNum, counts: tuple[int, int] | None = None) -> dict def _get_single_stat( stats_dict: dict[str, np.floating[Any]], stats_aliases: dict[str, str], stat_name: str -) -> np.floating[Any]: +) -> pd.DataFrame: """ Retrieve a single statistic based on a flexible name or alias. @@ -138,3 +139,43 @@ def _get_single_stat( else: logging.warning("Statistic name '%s' is not recognized", stat_name) return np.float32(np.nan) + +def _grouped_stats(arrays: dict[str, NDArrayNum], bins: list[NDArrayNum], values: dict[str, NDArrayNum], + statistics: list[str]): + """ + Get statistics grouped (=binned) by other variables, whether categorical or continuous. + + :param arrays: Arrays to group by. + :param bins: Bins to use. + :param values: Values to group, can be a dictionary . + :param statistics: List or dict of statistics to compute, e.g. ["mean", "std"]. + """ + if len(arrays) != len(bins): + raise ValueError("One bins array must be provided per input array.") + + # Build dataframe of arrays to bin with, and values to bin + df = pd.DataFrame(data=arrays) + for k, v in values.items(): + df[k] = v + + # Apply binning + group_keys = [] + for i, b in enumerate(bins): + k = list(arrays.keys())[i] + df[f"bin_{k}"] = pd.cut(df[k], bins=b, include_lowest=True) + group_keys.append(f"bin_{k}") + + # Perform aggregation + grouped = df.groupby(group_keys) + result = grouped[list(values.keys())].agg(statistics) + + return result + + +arrays = {"slope": np.random.normal(size=100), "aspect": np.random.normal(size=100)} +values = {"band1": np.random.normal(size=100), "band2": np.random.normal(size=100)} +statistics = ["mean", "std"] +bins = [np.linspace(-2, 2, 10), 10] + +df = _grouped_stats(arrays, bins, values, statistics) +df \ No newline at end of file