diff --git a/src/awkward/_v2/_typetracer.py b/src/awkward/_v2/_typetracer.py index 4ebb0e940d..65b3632e6f 100644 --- a/src/awkward/_v2/_typetracer.py +++ b/src/awkward/_v2/_typetracer.py @@ -549,9 +549,9 @@ def searchsorted(self, *args, **kwargs): # haystack, needle, side="right" raise NotImplementedError - def argsort(self, *args, **kwargs): + def argsort(self, array, *args, **kwargs): # array - raise NotImplementedError + return TypeTracerArray(np.int64, array.shape) ############################ manipulation diff --git a/src/awkward/_v2/contents/listoffsetarray.py b/src/awkward/_v2/contents/listoffsetarray.py index 66da0b2ef1..7bad0a97b6 100644 --- a/src/awkward/_v2/contents/listoffsetarray.py +++ b/src/awkward/_v2/contents/listoffsetarray.py @@ -1656,14 +1656,12 @@ def _reduce_next( ) distincts_length = outlength * maxcount[0] - nextcarry = ak._v2.index.Index64.empty(nextlen, self._nplike) - nextparents = ak._v2.index.Index64.empty(nextlen, self._nplike) + np_nextcarry = self._nplike.empty(nextlen, dtype=np.int64) + np_nextparents = self._nplike.empty(nextlen, dtype=np.int64) maxnextparents = ak._v2.index.Index64.empty(1, self._nplike) distincts = ak._v2.index.Index64.empty(distincts_length, self._nplike) assert ( - nextcarry.nplike is self._nplike - and nextparents.nplike is self._nplike - and maxnextparents.nplike is self._nplike + maxnextparents.nplike is self._nplike and distincts.nplike is self._nplike and self._offsets.nplike is self._nplike and offsetscopy.nplike is self._nplike @@ -1672,16 +1670,16 @@ def _reduce_next( self._handle_error( self._nplike[ "awkward_ListOffsetArray_reduce_nonlocal_preparenext_64", - nextcarry.dtype.type, - nextparents.dtype.type, + np_nextcarry.dtype.type, + np_nextparents.dtype.type, maxnextparents.dtype.type, distincts.dtype.type, self._offsets.dtype.type, offsetscopy.dtype.type, parents.dtype.type, ]( - nextcarry.data, - nextparents.data, + np_nextcarry, + np_nextparents, nextlen, maxnextparents.data, distincts.data, @@ -1694,6 +1692,11 @@ def _reduce_next( ) ) + # A "stable" sort is essential for the subsequent steps. + reorder = self._nplike.argsort(np_nextparents, kind="stable") + nextcarry = ak._v2.index.Index64(np_nextcarry[reorder]) + nextparents = ak._v2.index.Index64(np_nextparents[reorder]) + nextstarts = ak._v2.index.Index64.empty(maxnextparents[0] + 1, self._nplike) assert ( nextstarts.nplike is self._nplike and nextparents.nplike is self._nplike diff --git a/studies/reducers.py b/studies/reducers.py index 75dd84adff..072300b41a 100644 --- a/studies/reducers.py +++ b/studies/reducers.py @@ -1506,3 +1506,78 @@ def IndexedOptionArray_reduce_next(self, negaxis, parents, length): [101*31, 103*37, 107*41, 109*43, 113*47], [ 1], [ 53*2, 59*3, 61*5, 67*7, 71*11]] + +content = RawArray([1,2,3, 4,5,6, 5,6,7, 8,9,10, 9,10,11, 12,13,14, 13,14,15, 16,17,18]) +assert content.tolist() == [1,2,3, 4,5,6, 5,6,7, 8,9,10, 9,10,11, 12,13,14, 13,14,15, 16,17,18] + +depth1 = RegularArray(content, 3) +depth2 = RegularArray(depth1, 2) +depth3 = RegularArray(depth2, 2) +depth4 = RegularArray(depth3, 2) +print(list(depth4)) +assert list(depth4) == [[[[[1, 2, 3],[4, 5, 6]],[[5, 6, 7],[8, 9, 10]]],[[[9, 10, 11],[12, 13, 14]],[[13, 14, 15],[16, 17, 18]]]]] +print("depth1", list(depth1.reduce(-1))) +print("depth2", list(depth2.reduce(-1))) +print("depth3", list(depth3.reduce(-1))) +print("depth4 (-1)", list(depth4.reduce(-1))) +# >>> np.prod(array, axis=-1) +# array([[[ 6, 120], +# [ 210, 720]], +# +# [[ 990, 2184], +# [2730, 4896]]]) +assert list(depth4.reduce(-1)) == [[[[6,120],[210,720]],[[990,2184],[2730,4896]]]] + +print("depth4 (-2)", list(depth4.reduce(-2))) +# >>> np.prod(array, axis=-2) +# array([[[ 4, 10, 18], +# [ 40, 54, 70]], +# +# [[108, 130, 154], +# [208, 238, 270]]]) +assert list(depth4.reduce(-2)) == [[[[4,10,18],[40,54,70]],[[108,130,154],[208,238,270]]]] + +print("depth4 (-3)", list(depth4.reduce(-3))) +# >>> np.prod(array, axis=-3) +# array([[[ 5, 12, 21], +# [ 32, 45, 60]], +# +# [[117, 140, 165], +# [192, 221, 252]]]) +assert list(depth4.reduce(-3)) == [[[[5,12,21],[32,45,60]],[[117,140,165],[192,221,252]]]] + +print("depth4 (-4)", list(depth4.reduce(-4))) +# >>> np.prod(array, axis=-4) +# array([[[ 9, 20, 33], +# [ 48, 65, 84]], +# +# [[ 65, 84, 105], +# [128, 153, 180]]]) +assert list(depth4.reduce(-4)) == [[[[6,120],[210,720]],[[990,2184],[2730,4896]]]] + + + +# >>> np.prod(array, axis=0) +# array([[[ 9, 20, 33], +# [ 48, 65, 84]], +# +# [[ 65, 84, 105], +# [128, 153, 180]]]) +# >>> np.prod(array, axis=1) +# array([[[ 5, 12, 21], +# [ 32, 45, 60]], +# +# [[117, 140, 165], +# [192, 221, 252]]]) +# >>> np.prod(array, axis=2) +# array([[[ 4, 10, 18], +# [ 40, 54, 70]], +# +# [[108, 130, 154], +# [208, 238, 270]]]) +# >>> np.prod(array, axis=3) +# array([[[ 6, 120], +# [ 210, 720]], +# +# [[ 990, 2184], +# [2730, 4896]]]) diff --git a/tests/v2/test_1271-fix-4D-reducers.py b/tests/v2/test_1271-fix-4D-reducers.py new file mode 100644 index 0000000000..31552fc17d --- /dev/null +++ b/tests/v2/test_1271-fix-4D-reducers.py @@ -0,0 +1,17 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE + +import pytest # noqa: F401 +import numpy as np # noqa: F401 +import awkward as ak # noqa: F401 + + +def test(): + akarray = ak._v2.highlevel.Array( + [[[[1], [4]], [[5], [8]]], [[[9], [12]], [[13], [16]]]] + ) + nparray = np.array([[[[1], [4]], [[5], [8]]], [[[9], [12]], [[13], [16]]]]) + + assert ak._v2.sum(akarray, axis=3).tolist() == np.sum(nparray, axis=3).tolist() + assert ak._v2.sum(akarray, axis=2).tolist() == np.sum(nparray, axis=2).tolist() + assert ak._v2.sum(akarray, axis=1).tolist() == np.sum(nparray, axis=1).tolist() + assert ak._v2.sum(akarray, axis=0).tolist() == np.sum(nparray, axis=0).tolist()