File: test_2501_positional_record_reducer.py

package info (click to toggle)
python-awkward 2.6.5-1
  • links: PTS, VCS
  • area: main
  • in suites: sid
  • size: 23,088 kB
  • sloc: python: 148,689; cpp: 33,562; sh: 432; makefile: 21; javascript: 8
file content (153 lines) | stat: -rw-r--r-- 3,832 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE

from __future__ import annotations

import pytest

import awkward as ak


def _min_pair(array, mask):
    array = ak.typetracer.length_zero_if_typetracer(array)

    # Find location of minimum 0 slot
    i_min = ak.argmin(array["0"], axis=-1, keepdims=True, mask_identity=True)
    # Index into array
    pair_min = ak.ravel(array[ak.from_regular(i_min)], highlevel=False)
    if mask:
        return pair_min
    else:
        form = pair_min.content.form
        length_one_content = form.length_one_array(
            backend=pair_min.backend, highlevel=False
        )
        identity_content = ak.fill_none(length_one_content, 0, highlevel=False)
        identity = ak.record.Record(identity_content, 0)
        return ak.fill_none(pair_min, identity, highlevel=False)


def _argmin_pair(array, mask):
    array = ak.typetracer.length_zero_if_typetracer(array)

    assert not mask
    # Find location of minimum 0 slot
    return ak.argmin(array["0"], axis=-1, keepdims=False, mask_identity=mask)


def _argmin_pair_bad(array, mask):
    array = ak.typetracer.length_zero_if_typetracer(array)

    assert not mask
    # Find location of minimum 0 slot
    return ak.argmin(array["0"], axis=-1, keepdims=False, mask_identity=True)


def test_non_positional():
    behavior = {(ak.min, "pair"): _min_pair}

    x = ak.Array(
        [
            [
                [1, 2, 3],
                [5, 4, 3],
                [2],
            ],
            [
                [8],
                [],
                [10, 4, 4],
            ],
        ]
    )
    y = 2 * x - x**2
    z = ak.zip((x, y), with_name="pair", behavior=behavior)

    assert ak.almost_equal(
        ak.min(z, axis=-1, mask_identity=True),
        ak.Array(
            [
                [(1, 1), (3, -3), (2, 0)],
                [
                    (8, -48),
                    None,
                    (4, -8),
                ],
            ],
            with_name="pair",
        ),
    )

    assert ak.almost_equal(
        ak.min(z, axis=-1, mask_identity=False),
        ak.Array(
            [
                [(1, 1), (3, -3), (2, 0)],
                [
                    (8, -48),
                    (0, 0),
                    (4, -8),
                ],
            ],
            with_name="pair",
        ),
    )


def test_positional_bad():
    behavior = {(ak.argmin, "pair"): _argmin_pair_bad}

    x = ak.Array(
        [
            [
                [1, 2, 3],
                [5, 4, 3],
                [2],
            ],
            [
                [8],
                [],
                [10, 4, 4],
            ],
        ]
    )
    y = 2 * x - x**2
    z = ak.zip((x, y), with_name="pair", behavior=behavior)

    with pytest.raises(TypeError, match=r"'pair' returned an option"):
        assert ak.almost_equal(
            ak.argmin(z, axis=-1, mask_identity=True), [[0, 2, 0], [0, None, 1]]
        )

    with pytest.raises(TypeError, match=r"'pair' returned an option"):
        assert ak.almost_equal(
            ak.argmin(z, axis=-1, mask_identity=False), [[0, 2, 0], [0, -1, 1]]
        )


def test_positional_good():
    behavior = {(ak.argmin, "pair"): _argmin_pair}

    x = ak.Array(
        [
            [
                [1, 2, 3],
                [5, 4, 3],
                [2],
            ],
            [
                [8],
                [],
                [10, 4, 4],
            ],
        ]
    )
    y = 2 * x - x**2
    z = ak.zip((x, y), with_name="pair", behavior=behavior)

    assert ak.almost_equal(
        ak.argmin(z, axis=-1, mask_identity=True), [[0, 2, 0], [0, None, 1]]
    )

    assert ak.almost_equal(
        ak.argmin(z, axis=-1, mask_identity=False), [[0, 2, 0], [0, -1, 1]]
    )