File: test_2373_unzip_touching.py

package info (click to toggle)
python-awkward 2.8.9-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 24,932 kB
  • sloc: python: 178,875; cpp: 33,828; sh: 432; makefile: 21; javascript: 8
file content (128 lines) | stat: -rw-r--r-- 5,376 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
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE

from __future__ import annotations

import awkward as ak


def test():
    form = ak.forms.from_dict(
        {
            "class": "RecordArray",
            "fields": ["muon", "jet"],
            "contents": [
                {
                    "class": "ListOffsetArray",
                    "offsets": "i64",
                    "content": {
                        "class": "RecordArray",
                        "fields": ["pt", "eta", "phi", "crossref"],
                        "contents": [
                            {
                                "class": "NumpyArray",
                                "primitive": "int64",
                                "inner_shape": [],
                                "parameters": {},
                                "form_key": "muon_pt!",
                            },
                            {
                                "class": "NumpyArray",
                                "primitive": "int64",
                                "inner_shape": [],
                                "parameters": {},
                                "form_key": "muon_eta!",
                            },
                            {
                                "class": "NumpyArray",
                                "primitive": "int64",
                                "inner_shape": [],
                                "parameters": {},
                                "form_key": "muon_phi!",
                            },
                            {
                                "class": "ListOffsetArray",
                                "offsets": "i64",
                                "content": {
                                    "class": "NumpyArray",
                                    "primitive": "int64",
                                    "inner_shape": [],
                                    "parameters": {},
                                    "form_key": "muon_crossref_content!",
                                },
                                "parameters": {},
                                "form_key": "muon_crossref_index!",
                            },
                        ],
                        "parameters": {},
                        "form_key": "muon_record!",
                    },
                    "parameters": {},
                    "form_key": "muon_list!",
                },
                {
                    "class": "ListOffsetArray",
                    "offsets": "i64",
                    "content": {
                        "class": "RecordArray",
                        "fields": ["pt", "eta", "phi", "crossref", "thing1"],
                        "contents": [
                            {
                                "class": "NumpyArray",
                                "primitive": "int64",
                                "inner_shape": [],
                                "parameters": {},
                                "form_key": "jet_pt!",
                            },
                            {
                                "class": "NumpyArray",
                                "primitive": "int64",
                                "inner_shape": [],
                                "parameters": {},
                                "form_key": "jet_eta!",
                            },
                            {
                                "class": "NumpyArray",
                                "primitive": "int64",
                                "inner_shape": [],
                                "parameters": {},
                                "form_key": "jet_phi!",
                            },
                            {
                                "class": "ListOffsetArray",
                                "offsets": "i64",
                                "content": {
                                    "class": "NumpyArray",
                                    "primitive": "int64",
                                    "inner_shape": [],
                                    "parameters": {},
                                    "form_key": "jet_crossref_content!",
                                },
                                "parameters": {},
                                "form_key": "jet_crossref_index!",
                            },
                            {
                                "class": "NumpyArray",
                                "primitive": "int64",
                                "inner_shape": [],
                                "parameters": {},
                                "form_key": "jet_thing1!",
                            },
                        ],
                        "parameters": {},
                        "form_key": "jet_record!",
                    },
                    "parameters": {},
                    "form_key": "jet_list!",
                },
            ],
            "parameters": {},
            "form_key": "outer!",
        }
    )

    ttlayout, report = ak.typetracer.typetracer_with_report(form)

    ttarray = ak.Array(ttlayout)
    pairs = ak.cartesian([ttarray.muon, ttarray.jet], axis=1, nested=True)
    a, b = ak.unzip(pairs)
    assert set(report.data_touched) == {"muon_list!", "jet_list!"}