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

Commit dcbba4e

Browse files
dcheriandcherian
committed
Support nanfirst, nanlast with simple combine algo
Closes #227
1 parent 4164712 commit dcbba4e

File tree

4 files changed

+160
-26
lines changed
8000

4 files changed

+160
-26
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: 30 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -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,13 @@ 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 >= 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+
result = combine(array, axis=axis_, keepdims=True)
964968
if is_aggregate:
965969
# squeeze out DUMMY_AXIS if this is the last step i.e. called from _aggregate
966970
result = result.squeeze(axis=DUMMY_AXIS)
@@ -1539,6 +1543,10 @@ def _validate_reindex(
15391543
if all_numpy:
15401544
return True
15411545

1546+
if func in ["first", "last"]:
1547+
# have to do the grouped_combine since there's no good fill_value
1548+
reindex = False
1549+
15421550
if method == "blockwise" or _is_arg_reduction(func):
15431551
reindex = False
15441552

@@ -1552,6 +1560,9 @@ def _validate_reindex(
15521560
reindex = True
15531561

15541562
assert isinstance(reindex, bool)
1563+
1564+
if reindex and func in ["first", "last"]:
1565+
raise ValueError("reindex must be None or False when func is 'first' or 'last.")
15551566
return reindex
15561567

15571568

@@ -1875,6 +1886,21 @@ def groupby_reduce(
18751886
axis_ = np.core.numeric.normalize_axis_tuple(axis, array.ndim) # type: ignore
18761887
nax = len(axis_)
18771888

1889+
has_dask = is_duck_dask_array(array) or is_duck_dask_array(by_)
1890+
1891+
if _is_first_last_reduction(func):
1892+
if has_dask and nax != 1:
1893+
raise ValueError(
1894+
"For dask arrays: first, last, nanfirst, nanlast reductions are "
1895+
"only supported along a single axis. Please reshape appropriately."
1896+
)
1897+
1898+
elif nax not in [1, by_.ndim]:
1899+
raise ValueError(
1900+
"first, last, nanfirst, nanlast reductions are only supported "
1901+
"along a single axis or when reducing across all dimensions of `by`."
1902+
)
1903+
18781904
# TODO: make sure expected_groups is unique
18791905
if nax == 1 and by_.ndim > 1 and expected_groups is None:
18801906
if not any_by_dask:
@@ -1898,8 +1924,6 @@ def groupby_reduce(
18981924
axis_ = tuple(array.ndim + np.arange(-nax, 0))
18991925
nax = len(axis_)
19001926

1901-
has_dask = is_duck_dask_array(array) or is_duck_dask_array(by_)
1902-
19031927
# When axis is a subset of possible values; then npg will
19041928
# apply it to groups that don't exist along a particular axis (for e.g.)
19051929
# since these count as a group that is absent. thoo!
@@ -1986,6 +2010,6 @@ def groupby_reduce(
19862010
).reshape(result.shape[:-1] + grp_shape)
19872011
groups = final_groups
19882012

1989-
if _is_minmax_reduction(func) and is_bool_array:
2013+
if is_bool_array and (_is_minmax_reduction(func) or _is_first_last_reduction(func)):
19902014
result = result.astype(bool)
19912015
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: 82 additions & 19 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)(array_, 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,20 @@ 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(ValueError):
513+
groupby_reduce(np.empty((2, 3, 2)), np.ones((2, 3, 2)), func=func, axis=axis)
514+
515+
516+
@pytest.mark.parametrize("func", ["first", "last", "nanfirst", "nanlast"])
517+
@pytest.mark.parametrize("axis", [None, (0, 1, 2)])
518+
def test_first_last_disallowed_dask(axis, func):
519+
with pytest.raises(ValueError):
520+
groupby_reduce(dask.array.empty((2, 3, 2)), np.ones((2, 3, 2)), func=func, axis=axis)
521+
522+
489523
@requires_dask
490524
@pytest.mark.parametrize("func", ALL_FUNCS)
491525
@pytest.mark.parametrize(
@@ -495,8 +529,12 @@ def test_groupby_reduce_axis_subset_against_numpy(func, axis, engine):
495529
if "arg" in func and engine == "flox":
496530
pytest.skip()
497531

498-
if not isinstance(axis, int) and "arg" in func and (axis is None or len(axis) > 1):
499-
pytest.skip()
532+
if not isinstance(axis, int):
533+
if "arg" in func and (axis is None or len(axis) > 1):
534+
pytest.skip()
535+
if ("first" in func or "last" in func) and (axis is not None and len(axis) not in [1, 3]):
536+
pytest.skip()
537+
500538
if func in ["all", "any"]:
501539
fill_value = False
502540
else:
@@ -513,17 +551,45 @@ def test_groupby_reduce_axis_subset_against_numpy(func, axis, engine):
513551
kwargs = dict(
514552
func=func, axis=axis, expected_groups=[0, 2], fill_value=fill_value, engine=engine
515553
)
554+
expected, _ = groupby_reduce(array, by, **kwargs)
555+
if engine == "flox":
556+
kwargs.pop("engine")
557+
expected_npg, _ = groupby_reduce(array, by, **kwargs, engine="numpy")
558+
assert_equal(expected_npg, expected)
559+
560+
if func in ["all", "any"]:
561+
fill_value = False
562+
else:
563+
fill_value = 123
564+
565+
if "var" in func or "std" in func:
566+
tolerance = {"rtol": 1e-14, "atol": 1e-16}
567+
else:
568+
tolerance = None
569+
# tests against the numpy output to make sure dask compute matches
570+
by = np.broadcast_to(labels2d, (3, *labels2d.shape))
571+
rng = np.random.default_rng(12345)
572+
array = rng.random(by.shape)
573+
kwargs = dict(
574+
func=func, axis=axis, expected_groups=[0, 2], fill_value=fill_value, engine=engine
575+
)
576+
expected, _ = groupby_reduce(array, by, **kwargs)
577+
if engine == "flox":
578+
kwargs.pop("engine")
579+
expected_npg, _ = groupby_reduce(array, by, **kwargs, engine="numpy")
580+
assert_equal(expected_npg, expected)
581+
582+
if ("first" in func or "last" in func) and (
583+
axis is None or (not isinstance(axis, int) and len(axis) != 1)
584+
):
585+
return
586+
516587
with raise_if_dask_computes():
517588
actual, _ = groupby_reduce(
518589
da.from_array(array, chunks=(-1, 2, 3)),
519590
da.from_array(by, chunks=(-1, 2, 2)),
520591
**kwargs,
521592
)
522-
expected, _ = groupby_reduce(array, by, **kwargs)
523-
if engine == "flox":
524-
kwargs.pop("engine")
525-
expected_npg, _ = groupby_reduce(array, by, **kwargs, engine="numpy")
526-
assert_equal(expected_npg, expected)
527593
assert_equal(actual, expected, tolerance)
528594

529595

@@ -751,23 +817,17 @@ def test_fill_value_behaviour(func, chunks, fill_value, engine):
751817
if chunks is not None and not has_dask:
752818
pytest.skip()
753819

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-
820+
npfunc = _get_array_func(func)
763821
by = np.array([1, 2, 3, 1, 2, 3])
764822
array = np.array([np.nan, 1, 1, np.nan, 1, 1])
765823
if chunks:
766824
array = dask.array.from_array(array, chunks)
767825
actual, _ = groupby_reduce(
768826
array, by, func=func, engine=engine, fill_value=fill_value, expected_groups=[0, 1, 2, 3]
769827
)
770-
expected = np.array([fill_value, fill_value, npfunc([1.0, 1.0]), npfunc([1.0, 1.0])])
828+
expected = np.array(
829+
[fill_value, fill_value, npfunc([1.0, 1.0], axis=0), npfunc([1.0, 1.0], axis=0)]
830+
)
771831
assert_equal(actual, expected)
772832

773833

@@ -832,6 +892,8 @@ def test_cohorts_nd_by(func, method, axis, engine):
832892

833893
if axis is not None and method != "map-reduce":
834894
pytest.xfail()
895+
if axis is None and ("first" in func or "last" in func):
896+
pytest.skip()
835897

836898
kwargs = dict(func=func, engine=engine, method=method, axis=axis, fill_value=fill_value)
837899
actual, groups = groupby_reduce(array, by, **kwargs)
@@ -897,7 +959,8 @@ def test_bool_reductions(func, engine):
897959
pytest.skip()
898960
groups = np.array([1, 1, 1])
899961
data = np.array([True, True, False])
900-
expected = np.expand_dims(getattr(np, func)(data), -1)
962+
npfunc = _get_array_func(func)
963+
expected = np.expand_dims(npfunc(data, axis=0), -1)
901964
actual, _ = groupby_reduce(data, groups, func=func, engine=engine)
902965
assert_equal(expected, actual)
903966

0 commit comments

Comments
 (0)
0