diff --git a/chainladder/core/pandas.py b/chainladder/core/pandas.py index cd0defa5f..ecab025f2 100644 --- a/chainladder/core/pandas.py +++ b/chainladder/core/pandas.py @@ -8,6 +8,9 @@ from __future__ import annotations import warnings + +from typing import TYPE_CHECKING, cast + import numpy as np import pandas as pd @@ -18,24 +21,23 @@ ) from chainladder.utils.utility_functions import concat, num_to_nan -from typing import cast, TYPE_CHECKING - - try: import dask.bag as db except ImportError: db = None if TYPE_CHECKING: - from chainladder import Triangle - from chainladder.utils.sparse import COO - from chainladder.core.typing import BackendArray, TriangleProtocol - from collections.abc import Callable + from collections.abc import Callable, Mapping + from types import ModuleType + from typing import Any, Literal + from numpy import ndarray from pandas import DataFrame, Series - from types import ModuleType from pandas._typing import IndexLabel - from typing import Any, Literal, Type + + from chainladder import Triangle + from chainladder.core.typing import BackendArray, TriangleProtocol + from chainladder.utils.sparse import COO _TrianglePandasBase = TriangleProtocol else: @@ -48,7 +50,7 @@ def __init__(self, obj: Triangle, by, axis=0, **kwargs): self.axis = self.obj._get_axis(axis) self.by = by if kwargs.get("groups", None): - self.groups = kwargs.get("groups", None) + self.groups = kwargs["groups"] elif self.axis == 0: self.groups = obj.index.groupby(by) elif axis == 1: @@ -64,6 +66,46 @@ def __getitem__(self, key): groups=self.groups, ) + def sum(self, *args: Any, **kwargs: Any) -> Triangle: + """Refer to pandas for ``sum`` functionality.""" + return _groupby_aggregate(self, "sum", *args, **kwargs) + + def mean(self, *args: Any, **kwargs: Any) -> Triangle: + """Refer to pandas for ``mean`` functionality.""" + return _groupby_aggregate(self, "mean", *args, **kwargs) + + def median(self, *args: Any, **kwargs: Any) -> Triangle: + """Refer to pandas for ``median`` functionality.""" + return _groupby_aggregate(self, "median", *args, **kwargs) + + def max(self, *args: Any, **kwargs: Any) -> Triangle: + """Refer to pandas for ``max`` functionality.""" + return _groupby_aggregate(self, "max", *args, **kwargs) + + def min(self, *args: Any, **kwargs: Any) -> Triangle: + """Refer to pandas for ``min`` functionality.""" + return _groupby_aggregate(self, "min", *args, **kwargs) + + def prod(self, *args: Any, **kwargs: Any) -> Triangle: + """Refer to pandas for ``prod`` functionality.""" + return _groupby_aggregate(self, "prod", *args, **kwargs) + + def var(self, *args: Any, **kwargs: Any) -> Triangle: + """Refer to pandas for ``var`` functionality.""" + return _groupby_aggregate(self, "var", *args, **kwargs) + + def std(self, *args: Any, **kwargs: Any) -> Triangle: + """Refer to pandas for ``std`` functionality.""" + return _groupby_aggregate(self, "std", *args, **kwargs) + + def cumsum(self, *args: Any, **kwargs: Any) -> Triangle: + """Refer to pandas for ``cumsum`` functionality.""" + return _groupby_aggregate(self, "cumsum", *args, **kwargs) + + def quantile(self, *args: Any, **kwargs: Any) -> Triangle: + """Refer to pandas for ``quantile`` functionality.""" + return _groupby_aggregate(self, "quantile", *args, **kwargs) + class TrianglePandas(_TrianglePandasBase): if TYPE_CHECKING: @@ -78,10 +120,112 @@ def set_backend( **kwargs: Any, ) -> Triangle: ... - # Installed by add_triangle_agg_func at runtime. - def sum( - self, axis: str | int | None = None, *args: Any, **kwargs: Any - ) -> TriangleProtocol: ... + def sum(self, axis: str | int | None = None, *args: Any, **kwargs: Any) -> Triangle: + """Refer to pandas for ``sum`` functionality.""" + return _triangle_aggregate(self, "nansum", axis, *args, **kwargs) + + def mean( + self, axis: str | int | None = None, *args: Any, **kwargs: Any + ) -> Triangle: + """Refer to pandas for ``mean`` functionality.""" + return _triangle_aggregate(self, "nanmean", axis, *args, **kwargs) + + def median( + self, axis: str | int | None = None, *args: Any, **kwargs: Any + ) -> Triangle: + """Refer to pandas for ``median`` functionality.""" + return _triangle_aggregate(self, "nanmedian", axis, *args, **kwargs) + + def max(self, axis: str | int | None = None, *args: Any, **kwargs: Any) -> Triangle: + """Refer to pandas for ``max`` functionality.""" + return _triangle_aggregate(self, "nanmax", axis, *args, **kwargs) + + def min(self, axis: str | int | None = None, *args: Any, **kwargs: Any) -> Triangle: + """Refer to pandas for ``min`` functionality.""" + return _triangle_aggregate(self, "nanmin", axis, *args, **kwargs) + + def prod( + self, axis: str | int | None = None, *args: Any, **kwargs: Any + ) -> Triangle: + """Refer to pandas for ``prod`` functionality.""" + return _triangle_aggregate(self, "nanprod", axis, *args, **kwargs) + + def var(self, axis: str | int | None = None, *args: Any, **kwargs: Any) -> Triangle: + """Refer to pandas for ``var`` functionality.""" + return _triangle_aggregate(self, "nanvar", axis, *args, **kwargs) + + def std(self, axis: str | int | None = None, *args: Any, **kwargs: Any) -> Triangle: + """Refer to pandas for ``std`` functionality.""" + return _triangle_aggregate(self, "nanstd", axis, *args, **kwargs) + + def cumsum( + self, axis: str | int | None = None, *args: Any, **kwargs: Any + ) -> Triangle: + """Refer to pandas for ``cumsum`` functionality.""" + return _triangle_aggregate(self, "nancumsum", axis, *args, **kwargs) + + def quantile( + self, axis: str | int | None = None, *args: Any, **kwargs: Any + ) -> Triangle: + """Refer to pandas for ``quantile`` functionality.""" + return _triangle_aggregate(self, "nanquantile", axis, *args, **kwargs) + + def diff( + self, axis: str | int | None = None, *args: Any, **kwargs: Any + ) -> Triangle: + """Refer to pandas for ``diff`` functionality.""" + return _triangle_aggregate(self, "diff", axis, *args, **kwargs) + + def to_clipboard(self, *args: Any, **kwargs: Any) -> None: + """Refer to pandas for ``to_clipboard`` functionality.""" + return self._df_passthru("to_clipboard", *args, **kwargs) + + def to_csv(self, *args: Any, **kwargs: Any) -> str | None: + """Refer to pandas for ``to_csv`` functionality.""" + return self._df_passthru("to_csv", *args, **kwargs) + + def to_excel(self, *args: Any, **kwargs: Any) -> None: + """Refer to pandas for ``to_excel`` functionality.""" + return self._df_passthru("to_excel", *args, **kwargs) + + def to_json(self, *args: Any, **kwargs: Any) -> str | None: + """Refer to pandas for ``to_json`` functionality.""" + return self._df_passthru("to_json", *args, **kwargs) + + def to_html(self, *args: Any, **kwargs: Any) -> str | None: + """Refer to pandas for ``to_html`` functionality.""" + return self._df_passthru("to_html", *args, **kwargs) + + def to_dict(self, *args: Any, **kwargs: Any) -> Mapping | list[Mapping]: + """Refer to pandas for ``to_dict`` functionality.""" + return self._df_passthru("to_dict", *args, **kwargs) + + def unstack(self, *args: Any, **kwargs: Any) -> DataFrame | Series: + """Refer to pandas for ``unstack`` functionality.""" + return self._df_passthru("unstack", *args, **kwargs) + + def pivot(self, *args: Any, **kwargs: Any) -> DataFrame: + """Refer to pandas for ``pivot`` functionality.""" + return self._df_passthru("pivot", *args, **kwargs) + + def drop_duplicates(self, *args: Any, **kwargs: Any) -> DataFrame | None: + """Refer to pandas for ``drop_duplicates`` functionality.""" + return self._df_passthru("drop_duplicates", *args, **kwargs) + + def describe(self, *args: Any, **kwargs: Any) -> DataFrame: + """Refer to pandas for ``describe`` functionality.""" + return self._df_passthru("describe", *args, **kwargs) + + def melt(self, *args: Any, **kwargs: Any) -> DataFrame: + """Refer to pandas for ``melt`` functionality.""" + return self._df_passthru("melt", *args, **kwargs) + + def pct_change(self, *args: Any, **kwargs: Any) -> DataFrame: + """Refer to pandas for ``pct_change`` functionality.""" + return self._df_passthru("pct_change", *args, **kwargs) + + def _df_passthru(self, method: str, *args: Any, **kwargs: Any) -> Any: + return getattr(pd.DataFrame, method)(self.to_frame(), *args, **kwargs) def to_frame( self, @@ -1282,221 +1426,134 @@ def xs( return result -def add_triangle_agg_func(cls: Type[TrianglePandas], k: str, v: str): - """ - Aggregate Overrides in Triangle +def _triangle_aggregate( + self: TrianglePandas, + method: str, + axis: str | int | None = None, + *args: Any, + **kwargs: Any, +) -> Triangle: """ + Apply an array aggregation, preserving Triangle dimensions and metadata. - def agg_func( - self: Triangle, axis: str | int | None = None, *args, **kwargs - ) -> Triangle | ndarray: - """ - Applies the aggregation function specified by k from the outer function. - - Parameters - ---------- - - self: Triangle - The triangle to which the aggregation function will be applied. - - axis: str | int | None - The axis of the triangle to which the aggregation function will be applied. - - Returns - ------- - - An aggregated Triangle or ndarray. - """ - keepdims = kwargs.get("keepdims", None) - obj = self.copy() - auto_sparse: bool = kwargs.pop("auto_sparse", True) - if axis is None: - if max(obj.shape) == 1: - axis: int = 0 - else: - axis: int = min([num for num, _ in enumerate(obj.shape) if _ != 1]) - else: - axis = self._get_axis(axis) - xp: ModuleType = obj.get_array_module() - func: Callable = getattr(xp, v) - kwargs.update({"keepdims": True}) - # Apply the function on the requested axis. - if axis in [0, 1] and not obj.is_full: - with warnings.catch_warnings(): - # the lower right half of an undeveloped triangle will always emit an - # empty or All-NaN slice warning when aggregating by index or column - # (depending on `func`; e.g. nanmean warns "empty slice", - # nanmax warns All-NaN slice - warnings.filterwarnings( - "ignore", - category=RuntimeWarning, - message=".*empty slice.*", - ) - warnings.filterwarnings( - "ignore", - category=RuntimeWarning, - message=".*All-NaN slice.*", - ) - obj.values = func(obj.values, axis=axis, *args, **kwargs) + A singleton result is returned as a scalar unless ``keepdims=True``. + """ + keepdims = kwargs.get("keepdims", None) + obj = self.copy() + auto_sparse: bool = kwargs.pop("auto_sparse", True) + if axis is None: + if max(obj.shape) == 1: + axis = 0 else: - obj.values = func(obj.values, axis=axis, *args, **kwargs) - - # Aggregation function will collapse a dimension, so - # adjust the dimensions of the original object to match that of the aggregation. - if axis == 0 and obj.values.shape[axis] == 1 and len(obj.index) > 1: - obj.index = pd.DataFrame( - [["(All)"] * len(obj.key_labels)], columns=obj.key_labels + axis = min([num for num, _ in enumerate(obj.shape) if _ != 1]) + else: + axis = self._get_axis(axis) + xp: ModuleType = obj.get_array_module() + func: Callable = getattr(xp, method) + kwargs.update({"keepdims": True}) + # Apply the function on the requested axis. + if axis in [0, 1] and not obj.is_full: + with warnings.catch_warnings(): + # the lower right half of an undeveloped triangle will always emit an + # empty or All-NaN slice warning when aggregating by index or column + # (depending on `func`; e.g. nanmean warns "empty slice", + # nanmax warns All-NaN slice + warnings.filterwarnings( + "ignore", + category=RuntimeWarning, + message=".*empty slice.*", ) - if axis == 1 and obj.values.shape[axis] == 1 and len(obj.columns) > 1: - obj.columns = pd.Index([0], name="columns") - if axis == 2 and obj.values.shape[axis] == 1 and len(obj.odims) > 1: - obj.odims = obj.odims[0:1] - # If axis is development, set the ddims to be the valuation date. - if axis == 3 and obj.values.shape[axis] == 1 and len(obj.ddims) > 1: - obj.ddims = pd.DatetimeIndex( - [self.valuation_date], dtype=__dt64_dtype__, freq=None + warnings.filterwarnings( + "ignore", + category=RuntimeWarning, + message=".*All-NaN slice.*", ) - obj._set_slicers() - if auto_sparse: - obj = obj._auto_sparse() - obj.values = cast("BackendArray", num_to_nan(obj.values)) - if not keepdims and obj.shape == (1, 1, 1, 1): - return obj.values[0, 0, 0, 0] - else: - return obj - - set_method(cls, agg_func, k) - + obj.values = func(obj.values, axis=axis, *args, **kwargs) + else: + obj.values = func(obj.values, axis=axis, *args, **kwargs) + + # Aggregation function will collapse a dimension, so + # adjust the dimensions of the original object to match that of the aggregation. + if axis == 0 and obj.values.shape[axis] == 1 and len(obj.index) > 1: + obj.index = pd.DataFrame( + [["(All)"] * len(obj.key_labels)], columns=obj.key_labels + ) + if axis == 1 and obj.values.shape[axis] == 1 and len(obj.columns) > 1: + obj.columns = pd.Index([0], name="columns") + if axis == 2 and obj.values.shape[axis] == 1 and len(obj.odims) > 1: + obj.odims = obj.odims[0:1] + # If axis is development, set the ddims to be the valuation date. + if axis == 3 and obj.values.shape[axis] == 1 and len(obj.ddims) > 1: + obj.ddims = pd.DatetimeIndex( + [self.valuation_date], dtype=__dt64_dtype__, freq=None + ) + obj._set_slicers() + if auto_sparse: + obj = obj._auto_sparse() + obj.values = cast("BackendArray", num_to_nan(obj.values)) + if not keepdims and obj.shape == (1, 1, 1, 1): + return obj.values[0, 0, 0, 0] + else: + return obj -def add_groupby_agg_func(cls, k: str, v: str): - """Aggregate Overrides in GroupBy""" - def agg_func(self, *args, **kwargs): - from chainladder.utils import concat +def _groupby_aggregate( + self: TriangleGroupBy, method: str, *args: Any, **kwargs: Any +) -> Triangle: + """Aggregate each group and restore the grouped axis labels.""" + from chainladder.utils import concat - obj = self.obj.copy() - auto_sparse = kwargs.pop("auto_sparse", True) - if db and obj.array_backend == "sparse": - _warn_dask_parallel_deprecated() + obj = self.obj.copy() + auto_sparse = kwargs.pop("auto_sparse", True) + if db and obj.array_backend == "sparse": + _warn_dask_parallel_deprecated() - def aggregate(i, obj, axis, v): - return getattr( - obj.iloc.__getitem__(tuple([slice(None)] * axis + [i])), v - )(axis, auto_sparse=False, keepdims=True) + def aggregate(i, obj, axis, v): + return getattr(obj.iloc.__getitem__(tuple([slice(None)] * axis + [i])), v)( + axis, auto_sparse=False, keepdims=True + ) - bag = db.from_sequence(self.groups.indices.values()) - bag = bag.map(aggregate, obj, self.axis, v) - values = bag.compute(scheduler="threads") - else: - values = [ - getattr( - obj.iloc.__getitem__(tuple([slice(None)] * self.axis + [i])), v - )(self.axis, auto_sparse=False, keepdims=True) - for i in self.groups.indices.values() - ] - obj = concat(values, axis=self.axis, ignore_index=True) - group_index = self.groups.first().index - if self.axis == 0: - if isinstance(group_index, pd.MultiIndex): - index = ( - pd - .DataFrame( - np.zeros(len(group_index)), - index=group_index, - columns=["_"], - ) - .reset_index() - .iloc[:, :-1] + bag = db.from_sequence(self.groups.indices.values()) + bag = bag.map(aggregate, obj, self.axis, method) + values = bag.compute(scheduler="threads") + else: + values = [ + getattr( + obj.iloc.__getitem__(tuple([slice(None)] * self.axis + [i])), method + )(self.axis, auto_sparse=False, keepdims=True) + for i in self.groups.indices.values() + ] + obj = concat(values, axis=self.axis, ignore_index=True) + group_index = self.groups.first().index + if self.axis == 0: + if isinstance(group_index, pd.MultiIndex): + index = ( + pd + .DataFrame( + np.zeros(len(group_index)), + index=group_index, + columns=["_"], ) - obj.index = index - else: - index = pd.DataFrame(group_index) - obj.index = index - if self.axis == 1: - obj.columns = group_index - if self.axis == 2: - odims = self.obj._to_datetime( - pd.Series(self.groups.indices.keys()).to_frame(), [0] + .reset_index() + .iloc[:, :-1] ) - obj.origin_grain = self.obj._get_grain(odims) - split = obj.origin_grain.split("-") - obj.origin_grain = {"2Q": "S"}.get(split[0], split[0]) - obj.odims = odims.values - obj._set_slicers() - if auto_sparse: - obj = obj._auto_sparse() - return obj - - set_method(cls=cls, func=agg_func, k=k) - - -def add_df_passthru(cls, k): - """Pass Through of pandas functionality""" - - def df_passthru(self, *args, **kwargs): - return getattr(pd.DataFrame, k)(self.to_frame(), *args, **kwargs) - - set_method(cls, df_passthru, k) - - -def set_method( - cls: Type[TrianglePandas | TriangleGroupBy], func: Callable, k: str -) -> None: - """ - Assigns methods to a class. - - Parameters - ---------- - - cls: Type[TrianglePandas | TriangleGroupBy] - Class to be modified. - - func: Callable - Method to be added to the class supplied to parameter cls. - - k: str - Name of the method to be added to the class supplied to parameter cls. - - Returns - ------- - - None - """ - func.__doc__ = "Refer to pandas for ``{}`` functionality.".format(k) - func.__name__ = k - setattr(cls, func.__name__, func) - - -df_passthru = ( - [ - "to_clipboard", - "to_csv", - "to_excel", - "to_json", - "to_html", - ] - + [ - "to_dict", - "unstack", - "pivot", - "drop_duplicates", - "describe", - "melt", - ] - + [ - "pct_change", - ] -) -for method in df_passthru: - add_df_passthru(TrianglePandas, method) - - -agg_funcs = ["sum", "mean", "median", "max", "min", "prod", "var"] -agg_funcs = agg_funcs + ["std", "cumsum", "quantile"] -for func in agg_funcs: - add_groupby_agg_func(TriangleGroupBy, func, func) -agg_funcs = {item: "nan" + item for item in agg_funcs} -more_aggs = ["diff"] -agg_funcs = {**agg_funcs, **{item: item for item in more_aggs}} -for method, func in agg_funcs.items(): - add_triangle_agg_func(TrianglePandas, method, func) + obj.index = index + else: + index = pd.DataFrame(group_index) + obj.index = index + if self.axis == 1: + obj.columns = group_index + if self.axis == 2: + odims = self.obj._to_datetime( + pd.Series(self.groups.indices.keys()).to_frame(), [0] + ) + if odims is None: + raise ValueError("Group labels could not be converted to origin dates.") + obj.origin_grain = self.obj._get_grain(odims) + split = obj.origin_grain.split("-") + obj.origin_grain = {"2Q": "S"}.get(split[0], split[0]) + obj.odims = odims.values + obj._set_slicers() + if auto_sparse: + obj = obj._auto_sparse() + return obj diff --git a/chainladder/core/tests/test_pandas.py b/chainladder/core/tests/test_pandas.py new file mode 100644 index 000000000..cf874ade4 --- /dev/null +++ b/chainladder/core/tests/test_pandas.py @@ -0,0 +1,198 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +from unittest.mock import Mock + +import numpy as np +import pandas as pd +import pytest + +from chainladder.core import pandas as triangle_pandas +from chainladder.core.pandas import TriangleGroupBy, TrianglePandas + +if TYPE_CHECKING: + from chainladder import Triangle + + +AGGREGATIONS = [ + "sum", + "mean", + "median", + "max", + "min", + "prod", + "var", + "std", + "cumsum", + "quantile", +] +PASSTHRU_METHODS = [ + "to_clipboard", + "to_csv", + "to_excel", + "to_json", + "to_html", + "to_dict", + "unstack", + "pivot", + "drop_duplicates", + "describe", + "melt", + "pct_change", +] +# Covered by test_triangle.test_add_df_passthru. +PASSTHRU_METHODS_COVERED_BY_TEST_TRIANGLE = { + "to_csv", + "to_html", + "describe", + "drop_duplicates", + "melt", + "unstack", +} +PASSTHRU_METHODS_NOT_COVERED_BY_TEST_TRIANGLE = [ + method + for method in PASSTHRU_METHODS + if method not in PASSTHRU_METHODS_COVERED_BY_TEST_TRIANGLE +] + + +def test_pandas_methods_are_defined_on_classes(): + """ + Define pandas-style methods directly on the classes, not attached at runtime. + + This lets static type checkers such as Pyright/Pylance and ty see them. + """ + assert ( + set(AGGREGATIONS + ["diff"] + PASSTHRU_METHODS) + <= TrianglePandas.__dict__.keys() + ) + assert set(AGGREGATIONS) <= TriangleGroupBy.__dict__.keys() + + +@pytest.mark.parametrize("method", AGGREGATIONS + ["diff"]) +def test_triangle_aggregation_dispatch(raa, monkeypatch, method): + """ + Delegate each Triangle method to ``_triangle_aggregate`` unchanged. + + Aggregations pass the ``nan``-prefixed NumPy function name (e.g. ``nansum``); + ``diff`` passes ``diff``. Arguments and the return value pass through as is. + """ + aggregate = Mock() + monkeypatch.setattr(triangle_pandas, "_triangle_aggregate", aggregate) + + result = getattr(raa, method)("origin", "arg", keepdims=True, auto_sparse=False) + + backend_method = method if method == "diff" else "nan" + method + aggregate.assert_called_once_with( + raa, backend_method, "origin", "arg", keepdims=True, auto_sparse=False + ) + assert result is aggregate.return_value + + +@pytest.mark.parametrize("method", AGGREGATIONS) +def test_groupby_aggregation_dispatch(clrd, monkeypatch, method): + """ + Delegate each TriangleGroupBy aggregation to ``_groupby_aggregate``. + + The method name, arguments and return value pass through unchanged. + """ + aggregate = Mock() + monkeypatch.setattr(triangle_pandas, "_groupby_aggregate", aggregate) + grouped = clrd.groupby("LOB") + + result = getattr(grouped, method)("arg", auto_sparse=False) + + aggregate.assert_called_once_with(grouped, method, "arg", auto_sparse=False) + assert result is aggregate.return_value + + +@pytest.mark.parametrize("method", PASSTHRU_METHODS_NOT_COVERED_BY_TEST_TRIANGLE) +def test_dataframe_passthru_dispatch(raa, monkeypatch, method): + """ + Call the same-named DataFrame method on ``to_frame()``. + + Arguments and the return value pass through unchanged. Only methods not + already covered by ``test_triangle.test_add_df_passthru`` are checked. + """ + frame = raa.to_frame() + monkeypatch.setattr(raa, "to_frame", lambda: frame) + passthru = Mock() + monkeypatch.setattr(pd.DataFrame, method, passthru) + + result = getattr(TrianglePandas, method)(raa, "arg", option="value") + + passthru.assert_called_once_with(frame, "arg", option="value") + assert result is passthru.return_value + + +@pytest.mark.parametrize("method", [name for name in AGGREGATIONS if name != "cumsum"]) +@pytest.mark.parametrize("axis", ["index", "columns", "origin", "development"]) +@pytest.mark.filterwarnings("error::RuntimeWarning") +def test_triangle_aggregation_values(clrd, method, axis): + """ + Match the NumPy ``nan*`` reduction along each axis, with zeros as NaN. + + The result keeps the input's backend and the ``keepdims`` shape, no + RuntimeWarning is raised, and the input triangle is not modified. + """ + if clrd.array_backend == "sparse" and method in {"var", "std"}: + pytest.skip("The sparse backend does not implement nanvar or nanstd") + triangle = clrd.iloc[11:13, :2, :5, :5] + before = triangle.copy() + axis_number = triangle._get_axis(axis) + kwargs = {"q": 0.5} if method == "quantile" else {} + expected = getattr(np, "nan" + method)( + triangle.set_backend("numpy").values, + axis=axis_number, + keepdims=True, + **kwargs, + ) + expected[expected == 0] = np.nan + + result = getattr(triangle, method)( + axis=axis, keepdims=True, auto_sparse=False, **kwargs + ) + + assert result.shape == expected.shape + assert result.array_backend == triangle.array_backend + np.testing.assert_allclose( + result.set_backend("numpy").values, expected, equal_nan=True + ) + assert triangle.equals(before) + + +@pytest.mark.parametrize("axis", range(4)) +def test_aggregation_default_axis_first_non_singleton(clrd, axis): + """With ``axis=None``, ``sum`` reduces the first axis of length > 1.""" + indexer = [slice(11, 13), slice(0, 2), slice(0, 5), slice(0, 5)] + for preceding_axis in range(axis): + start = indexer[preceding_axis].start + indexer[preceding_axis] = slice(start, start + 1) + triangle = clrd.iloc[tuple(indexer)] + + result = triangle.sum(axis=None, keepdims=True, auto_sparse=False) + expected = triangle.sum(axis=axis, keepdims=True, auto_sparse=False) + + assert result.equals(expected) + assert result.shape[axis] == 1 + + +@pytest.mark.parametrize("backend", ["numpy", "sparse"]) +@pytest.mark.parametrize("keepdims", [False, True]) +def test_aggregation_default_axis_singleton_result(raa: Triangle, backend, keepdims): + """ + With ``axis=None`` and all axes of length 1, ``sum`` returns the single value. + + It is a NumPy scalar, or with ``keepdims=True`` a (1, 1, 1, 1) triangle + keeping the index. Checked on the numpy and sparse backends. + """ + triangle = raa.iloc[:, :, :1, :1].set_backend(backend) + result = triangle.sum(axis=None, keepdims=keepdims) + + if keepdims: + assert result.shape == (1, 1, 1, 1) + assert result.index.equals(triangle.index) + assert result.values[0, 0, 0, 0] == triangle.values[0, 0, 0, 0] + else: + assert isinstance(result, np.floating) + assert result == triangle.values[0, 0, 0, 0]