8000 Support nanfirst, nanlast with simple combine algo · xarray-contrib/flox@8ba1009 · GitHub
[go: up one dir, main page]
Skip to content

Commit 8ba1009

Browse files
dcheriandcherian
committed
Support nanfirst, nanlast with simple combine algo
Closes #227
1 parent 622ddb2 commit 8ba1009

File tree

4 files changed

+114
-20
lines changed

4 files changed

+114
-20
lines changed

‎flox/aggregations.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -184,6 +184,8 @@ def __init__(
184184
self.chunk: FuncTuple = _atleast_1d(chunk)
185185
# how to aggregate results after first round of reduction
186186
self.combine: FuncTuple = _atleast_1d(combine)
187+
# simpler reductions used with the "simple combine" algorithm
188+
self.simple_combine = None
187189
# final aggregation
188190
self.aggregate: Callable | str = aggregate if aggregate else self.combine[0]
189191
# finalize results (see mean)
@@ -577,4 +579,16 @@ def _initialize_aggregation(
577579
else:
578580
agg.min_count = 0
579581

582+
simple_combine = []
583+
for combine in agg.combine:
584+
if isinstance(combine, str):
585+
if combine in ["nanfirst", "nanlast"]:
586+
simple_combine.append(getattr(xrutils, combine))
587+
else:
588+
simple_combine.append(getattr(np, combine))
589+
else:
590+
simple_combine.append(combine)
591+
592+
agg.simple_combine = tuple(simple_combine)
593+
580594
return agg

‎flox/core.py

Lines changed: 26 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@
2626
import pandas as pd
2727
import toolz as tlz
2828

29-
from . import xrdtypes
29+
from . import xrdtypes, xrutils
3030
from .aggregate_flox import _prepare_for_flox
3131
from .aggregations import (
3232
Aggregation,
@@ -94,6 +94,10 @@ def _is_minmax_reduction(func: T_Agg) -> bool:
9494
)
9595

9696

97+
def _is_first_last_reduction(func: T_Agg) -> bool:
98+
return isinstance(func, str) and func in ["nanfirst", "nanlast", "first", "last"]
99+
100+
97101
def _get_expected_groups(by: T_By, sort: bool) -> pd.Index:
98102
if is_duck_dask_array(by):
99103
raise ValueError("Please provide expected_groups if not grouping by a numpy array.")
@@ -954,13 +958,16 @@ def _simple_combine(
954958
results: IntermediateDict = {"groups": unique_groups}
955959
results["intermediates"] = []
956960
axis_ = axis[:-1] + (DUMMY_AXIS,)
957-
for idx, combine in enumerate(agg.combine):
961+
for idx, combine in enumerate(agg.simple_combine):
958962
array = _conc2(x_chunk, key1="intermediates", key2=idx, axis=axis_)
959963
assert array.ndim &g 6DB6 t;= 2
960964
with warnings.catch_warnings():
961965
warnings.filterwarnings("ignore", r"All-NaN (slice|axis) encountered")
962-
assert isinstance(combine, str)
963-
result = getattr(np, combine)(array, axis=axis_, keepdims=True)
966+
assert callable(combine)
967+
if combine in [xrutils.nanfirst, xrutils.nanlast]:
968+
result = combine(array, axis=DUMMY_AXIS, keepdims=True)
969+
else:
970+
result = combine(array, axis=axis_, keepdims=True)
964971
if is_aggregate:
965972
# squeeze out DUMMY_AXIS if this is the last step i.e. called from _aggregate
966973
result = result.squeeze(axis=DUMMY_AXIS)
@@ -1539,6 +1546,10 @@ def _validate_reindex(
15391546
if all_numpy:
15401547
return True
15411548

1549+
if func in ["first", "last"]:
1550+
# have to do the grouped_combine since there's no good fill_value
1551+
reindex = False
1552+
15421553
if method == "blockwise" or _is_arg_reduction(func):
15431554
reindex = False
15441555

@@ -1552,6 +1563,9 @@ def _validate_reindex(
15521563
reindex = True
15531564

15541565
assert isinstance(reindex, bool)
1566+
1567+
if reindex and func in ["first", "last"]:
1568+
raise ValueError("reindex must be None or False when func is 'first' or 'last.")
15551569
return reindex
15561570

15571571

@@ -1806,6 +1820,13 @@ def groupby_reduce(
18061820
"Try engine='numpy' or engine='numba' instead."
18071821
)
18081822

1823+
if _is_first_last_reduction(func) and (
1824+
isinstance(axis, tuple) and len(axis) not in [array.ndim, 1]
1825+
):
1826+
raise NotImplementedError(
1827+
"first, last, nanfirst, nanlast reductions are only supported along a single axis."
1828+
)
1829+
18091830
bys: T_Bys = tuple(np.asarray(b) if not is_duck_array(b) else b for b in by)
18101831
nby = len(bys)
18111832
by_is_dask = tuple(is_duck_dask_array(b) for b in bys)
@@ -1986,6 +2007,6 @@ def groupby_reduce(
19862007
).reshape(result.shape[:-1] + grp_shape)
19872008
groups = final_groups
19882009

1989-
if _is_minmax_reduction(func) and is_bool_array:
2010+
if _is_minmax_reduction(func) or _is_first_last_reduction(func) and is_bool_array:
19902011
result = result.astype(bool)
19912012
return (result, *groups) # type: ignore[return-value] # Unpack not in mypy yet

‎flox/xrutils.py

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
# The functions defined here were copied based on the source code
22
# defined in xarray
33

4-
54
import datetime
65
from typing import Any, Iterable
76

87
import numpy as np
98
import pandas as pd
9+
from numpy.core.multiarray import normalize_axis_index # type: ignore[attr-defined]
1010

1111
try:
1212
import cftime
@@ -283,3 +283,36 @@ def _contains_cftime_datetimes(array) -> bool:
283283
return isinstance(sample, cftime.datetime)
284284
else:
285285
return False
286+
287+
288+
def _select_along_axis(values, idx, axis):
289+
other_ind = np.ix_(*[np.arange(s) for s in idx.shape])
290+
sl = other_ind[:axis] + (idx,) + other_ind[axis:]
291+
return values[sl]
292+
293+
294+
def nanfirst(values, axis, keepdims=False):
295+
if isinstance(axis, tuple):
296+
(axis,) = axis
297+
values = np.asarray(values)
298+
axis = normalize_axis_index(axis, values.ndim)
299+
idx_first = np.argmax(~pd.isnull(values), axis=axis)
300+
result = _select_along_axis(values, idx_first, axis)
301+
if keepdims:
302+
return np.expand_dims(result, axis=axis)
303+
else:
304+
return result
305+
306+
307+
def nanlast(values, axis, keepdims=False):
308+
if isinstance(axis, tuple):
309+
(axis,) = axis
310+
values = np.asarray(values)
311+
axis = normalize_axis_index(axis, values.ndim)
312+
rev = (slice(None),) * axis + (slice(None, None, -1),)
313+
idx_last = -1 - np.argmax(~pd.isnull(values)[rev], axis=axis)
314+
result = _select_along_axis(values, idx_last, axis)
315+
if keepdims:
316+
return np.expand_dims(result, axis=axis)
317+
else:
318+
return result

‎tests/test_core.py

Lines changed: 40 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,14 @@
33
import itertools
44
import warnings
55
from functools import partial, reduce
6-
from typing import TYPE_CHECKING
6+
from typing import TYPE_CHECKING, Callable
77

88
import numpy as np
99
import pandas as pd
1010
import pytest
1111
from numpy_groupies.aggregate_numpy import aggregate
1212

13+
from flox import xrutils
1314
from flox.aggregations import Aggregation
1415
from flox.core import (
1516
_convert_expected_groups_to_index,
@@ -53,6 +54,7 @@ def dask_array_ones(*args):
5354 "sum",
5455
"nansum",
5556
"argmax",
57+
"nanfirst",
5658
pytest.param("nanargmax", marks=(pytest.mark.skip,)),
5759
"prod",
5860
"nanprod",
@@ -70,6 +72,7 @@ def dask_array_ones(*args):
7072
pytest.param("nanargmin", marks=(pytest.mark.skip,)),
7173
"any",
7274
"all",
75+
"nanlast",
7376
pytest.param("median", marks=(pytest.mark.skip,)),
7477
pytest.param("nanmedian", marks=(pytest.mark.skip,)),
7578
)
@@ -78,6 +81,21 @@ def dask_array_ones(*args):
7881
from flox.core import T_Engine, T_ExpectedGroupsOpt, T_Func2
7982

8083

84+
def _get_array_func(func: str) -> Callable:
85+
if func == "count":
86+
87+
def npfunc(x):
88+
x = np.asarray(x)
89+
return (~np.isnan(x)).sum()
90+
91+
elif func in ["nanfirst", "nanlast"]:
92+
npfunc = getattr(xrutils, func)
93+
else:
94+
npfunc = getattr(np, func)
95+
96+
return npfunc
97+
98+
8199
def test_alignment_error():
82100
da = np.ones((12,))
83101
labels = np.ones((5,))
@@ -217,6 +235,8 @@ def test_groupby_reduce_all(nby, size, chunks, func, add_nan_by, engine):
217235
if "arg" in func and add_nan_by:
218236
array_[..., nanmask] = np.nan
219237
expected = getattr(np, "nan" + func)(ar 2851 ray_, axis=-1, **kwargs)
238+
elif func in ["nanfirst", "nanlast"]:
239+
expected = getattr(xrutils, func)(array_[..., ~nanmask], axis=-1, **kwargs)
220240
else:
221241
expected = getattr(np, func)(array_[..., ~nanmask], axis=-1, **kwargs)
222242
for _ in range(nby):
@@ -486,6 +506,13 @@ def test_dask_reduce_axis_subset():
486506
)
487507

488508

509+
@pytest.mark.parametrize("func", ["first", "last", "nanfirst", "nanlast"])
510+
@pytest.mark.parametrize("axis", [(0, 1)])
511+
def test_first_last_disallowed(axis, func):
512+
with pytest.raises(NotImplementedError):
513+
groupby_reduce(np.empty((2, 3, 2)), np.ones((2, 3, 2)), func=func, axis=axis)
514+
515+
489516
@requires_dask
490517
@pytest.mark.parametrize("func", ALL_FUNCS)
491518
@pytest.mark.parametrize(
@@ -495,8 +522,12 @@ def test_groupby_reduce_axis_subset_against_numpy(func, axis, engine):
495522
if "arg" in func and engine == "flox":
496523
pytest.skip()
497524

498-
if not isinstance(axis, int) and "arg" in func and (axis is None or len(axis) > 1):
499-
pytest.skip()
525+
if not isinstance(axis, int):
526+
if "arg" in func and (axis is None or len(axis) > 1):
527+
pytest.skip()
528+
if func in ["nanfirst", "nanlast"] and (axis is not None and len(axis) not in [1, 3]):
529+
pytest.skip()
530+
500531
if func in ["all", "any"]:
501532
fill_value = False
502533
else:
@@ -751,23 +782,17 @@ def test_fill_value_behaviour(func, chunks, fill_value, engine):
751782
if chunks is not None and not has_dask:
752783
pytest.skip()
753784

754-
if func == "count":
755-
756-
def npfunc(x):
757-
x = np.asarray(x)
758-
return (~np.isnan(x)).sum()
759-
760-
else:
761-
npfunc = getattr(np, func)
762-
785+
npfunc = _get_array_func(func)
763786
by = np.array([1, 2, 3, 1, 2, 3])
764787
array = np.array([np.nan, 1, 1, np.nan, 1, 1])
765788
if chunks:
766789
array = dask.array.from_array(array, chunks)
767790
actual, _ = groupby_reduce(
768791
array, by, func=func, engine=engine, fill_value=fill_value, expected_groups=[0, 1, 2, 3]
769792
)
770-
expected = np.array([fill_value, fill_value, npfunc([1.0, 1.0]), npfunc([1.0, 1.0])])
793+
expected = np.array(
794+
[fill_value, fill_value, npfunc([1.0, 1.0], axis=0), npfunc([1.0, 1.0], axis=0)]
795+
)
771796
assert_equal(actual, expected)
772797

773798

@@ -897,7 +922,8 @@ def test_bool_reductions(func, engine):
897922
pytest.skip()
898923
groups = np.array([1, 1, 1])
899924
data = np.array([True, True, False])
900-
expected = np.expand_dims(getattr(np, func)(data), -1)
925+
npfunc = _get_array_func(func)
926+
expected = np.expand_dims(npfunc(data, axis=0), -1)
901927
actual, _ = groupby_reduce(data, groups, func=func, engine=engine)
902928
assert_equal(expected, actual)
903929

0 commit comments

Comments
 (0)
0