Skip to content

Commit b2c60b5

Browse files
Saransh-cpphenryiii
Saransh-cpp
and
henryiii
authored
Add custom reprs to awkward coordinate classes (#212)
* Add reprs to Awkward coordinate classes * Update doctests * Update src/vector/backends/awkward.py Co-authored-by: Henry Schreiner <[email protected]> * Use f-strings in reprs Co-authored-by: Henry Schreiner <[email protected]>
1 parent ba4e8b4 commit b2c60b5

File tree

1 file changed

+59
-0
lines changed

1 file changed

+59
-0
lines changed

‎src/vector/backends/awkward.py

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,9 @@ class AzimuthalAwkward(CoordinatesAwkward, Azimuthal):
100100
to construct azimuthal type objects.
101101
"""
102102

103+
def __repr__(self) -> str:
104+
return f"{type(self).__name__}{self.elements}"
105+
103106
@classmethod
104107
def from_fields(cls, array: ak.Array) -> "AzimuthalAwkward":
105108
"""
@@ -112,6 +115,8 @@ def from_fields(cls, array: ak.Array) -> "AzimuthalAwkward":
112115
>>> import awkward as ak
113116
>>> a = ak.Array([{"x": [1, 2]}, {"y": [1]}])
114117
>>> az = vector.backends.awkward.AzimuthalAwkward.from_fields(a)
118+
>>> az
119+
AzimuthalAwkwardXY(<Array [[1, 2], None] type='2 * option[var * int64]'>, <Array [None, [1]] type='2 * option[var * int64]'>)
115120
>>> az.elements
116121
(<Array [[1, 2], None] type='2 * option[var * int64]'>, <Array [None, [1]] type='2 * option[var
117122
* int64]'>)
@@ -139,6 +144,8 @@ def from_momentum_fields(cls, array: ak.Array) -> "AzimuthalAwkward":
139144
>>> import awkward as ak
140145
>>> a = ak.Array([{"px": [1, 2]}, {"py": [1]}])
141146
>>> az = vector.backends.awkward.AzimuthalAwkward.from_momentum_fields(a)
147+
>>> az
148+
AzimuthalAwkwardXY(<Array [[1, 2], None] type='2 * option[var * int64]'>, <Array [None, [1]] type='2 * option[var * int64]'>)
142149
>>> az.elements
143150
(<Array [[1, 2], None] type='2 * option[var * int64]'>, <Array [None, [1]] type='2 * option[var
144151
* int64]'>)
@@ -174,6 +181,9 @@ class LongitudinalAwkward(CoordinatesAwkward, Longitudinal):
174181
to construct longitudinal type objects.
175182
"""
176183

184+
def __repr__(self) -> str:
185+
return f"{type(self).__name__}{self.elements}"
186+
177187
@classmethod
178188
def from_fields(cls, array: ak.Array) -> "LongitudinalAwkward":
179189
"""
@@ -187,6 +197,8 @@ def from_fields(cls, array: ak.Array) -> "LongitudinalAwkward":
187197
>>> import awkward as ak
188198
>>> a = ak.Array([{"theta": [1, 0]}])
189199
>>> l = vector.backends.awkward.LongitudinalAwkward.from_fields(a)
200+
>>> l
201+
LongitudinalAwkwardTheta(<Array [[1, 0]] type='1 * var * int64'>,)
190202
>>> l.elements
191203
(<Array [[1, 0]] type='1 * var * int64'>,)
192204
"""
@@ -216,6 +228,8 @@ def from_momentum_fields(cls, array: ak.Array) -> "LongitudinalAwkward":
216228
>>> import awkward as ak
217229
>>> a = ak.Array([{"theta": [1, 0]}])
218230
>>> l = vector.backends.awkward.LongitudinalAwkward.from_momentum_fields(a)
231+
>>> l
232+
LongitudinalAwkwardTheta(<Array [[1, 0]] type='1 * var * int64'>,)
219233
>>> l.elements
220234
(<Array [[1, 0]] type='1 * var * int64'>,)
221235
"""
@@ -246,6 +260,9 @@ class TemporalAwkward(CoordinatesAwkward, Temporal):
246260
to construct longitudinal type objects.
247261
"""
248262

263+
def __repr__(self) -> str:
264+
return f"{type(self).__name__}{self.elements}"
265+
249266
@classmethod
250267
def from_fields(cls, array: ak.Array) -> "TemporalAwkward":
251268
"""
@@ -258,6 +275,8 @@ def from_fields(cls, array: ak.Array) -> "TemporalAwkward":
258275
>>> import awkward as ak
259276
>>> a = ak.Array([{"tau": [1, 0]}])
260277
>>> t = vector.backends.awkward.TemporalAwkward.from_fields(a)
278+
>>> t
279+
TemporalAwkwardTau(<Array [[1, 0]] type='1 * var * int64'>,)
261280
>>> t.elements
262281
(<Array [[1, 0]] type='1 * var * int64'>,)
263282
"""
@@ -284,6 +303,8 @@ def from_momentum_fields(cls, array: ak.Array) -> "TemporalAwkward":
284303
>>> import awkward as ak
285304
>>> a = ak.Array([{"mass": [1, 0]}])
286305
>>> t = vector.backends.awkward.TemporalAwkward.from_momentum_fields(a)
306+
>>> t
307+
TemporalAwkwardTau(<Array [[1, 0]] type='1 * var * int64'>,)
287308
>>> t.elements
288309
(<Array [[1, 0]] type='1 * var * int64'>,)
289310
"""
@@ -320,6 +341,8 @@ class AzimuthalAwkwardXY(AzimuthalAwkward, AzimuthalXY):
320341
>>> import awkward as ak
321342
>>> a = ak.Array([{"x": [1, 2]}, {"y": [1]}])
322343
>>> az = vector.backends.awkward.AzimuthalAwkwardXY(a["x"], a["y"])
344+
>>> az
345+
AzimuthalAwkwardXY(<Array [[1, 2], None] type='2 * option[var * int64]'>, <Array [None, [1]] type='2 * option[var * int64]'>)
323346
>>> az.elements
324347
(<Array [[1, 2], None] type='2 * option[var * int64]'>, <Array [None, [1]] type='2 * option[var * int64]'>)"""
325348

@@ -352,6 +375,8 @@ class AzimuthalAwkwardRhoPhi(AzimuthalAwkward, AzimuthalRhoPhi):
352375
>>> import awkward as ak
353376
>>> a = ak.Array([{"rho": [1, 2]}, {"phi": [1]}])
354377
>>> az = vector.backends.awkward.AzimuthalAwkwardRhoPhi(a["rho"], a["phi"])
378+
>>> az
379+
AzimuthalAwkwardRhoPhi(<Array [[1, 2], None] type='2 * option[var * int64]'>, <Array [None, [1]] type='2 * option[var * int64]'>)
355380
>>> az.elements
356381
(<Array [[1, 2], None] type='2 * option[var * int64]'>, <Array [None, [1]] type='2 * option[var * int64]'>)
357382
"""
@@ -385,6 +410,8 @@ class LongitudinalAwkwardZ(LongitudinalAwkward, LongitudinalZ):
385410
>>> import awkward as ak
386411
>>> a = ak.Array([{"z": [1, 2]}])
387412
>>> l = vector.backends.awkward.LongitudinalAwkwardZ(a["z"])
413+
>>> l
414+
LongitudinalAwkwardZ(<Array [[1, 2]] type='1 * var * int64'>,)
388415
>>> l.elements
389416
(<Array [[1, 2]] type='1 * var * int64'>,)
390417
"""
@@ -417,6 +444,8 @@ class LongitudinalAwkwardTheta(LongitudinalAwkward, LongitudinalTheta):
417444
>>> import awkward as ak
418445
>>> a = ak.Array([{"theta": [1, 2]}])
419446
>>> l = vector.backends.awkward.LongitudinalAwkwardTheta(a["theta"])
447+
>>> l
448+
LongitudinalAwkwardTheta(<Array [[1, 2]] type='1 * var * int64'>,)
420449
>>> l.elements
421450
(<Array [[1, 2]] type='1 * var * int64'>,)
422451
"""
@@ -449,6 +478,8 @@ class LongitudinalAwkwardEta(LongitudinalAwkward, LongitudinalEta):
449478
>>> import awkward as ak
450479
>>> a = ak.Array([{"eta": [1, 2]}])
451480
>>> l = vector.backends.awkward.LongitudinalAwkwardEta(a["eta"])
481+
>>> l
482+
LongitudinalAwkwardEta(<Array [[1, 2]] type='1 * var * int64'>,)
452483
>>> l.elements
453484
(<Array [[1, 2]] type='1 * var * int64'>,)
454485
"""
@@ -481,6 +512,8 @@ class TemporalAwkwardT(TemporalAwkward, TemporalT):
481512
>>> import awkward as ak
482513
>>> a = ak.Array([{"t": [1, 2]}])
483514
>>> t = vector.backends.awkward.TemporalAwkwardT(a["t"])
515+
>>> t
516+
TemporalAwkwardT(<Array [[1, 2]] type='1 * var * int64'>,)
484517
>>> t.elements
485518
(<Array [[1, 2]] type='1 * var * int64'>,)
486519
"""
@@ -513,6 +546,8 @@ class TemporalAwkwardTau(TemporalAwkward, TemporalTau):
513546
>>> import awkward as ak
514547
>>> a = ak.Array([{"tau": [1, 2]}])
515548
>>> t = vector.backends.awkward.TemporalAwkwardTau(a["tau"])
549+
>>> t
550+
TemporalAwkwardTau(<Array [[1, 2]] type='1 * var * int64'>,)
516551
>>> t.elements
517552
(<Array [[1, 2]] type='1 * var * int64'>,)
518553
"""
@@ -904,6 +939,8 @@ def azimuthal(self) -> AzimuthalAwkward:
904939
>>> a = vector.Array(
905940
... [{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}],
906941
... )
942+
>>> a.azimuthal
943+
AzimuthalAwkwardXY(<Array [1, 2] type='2 * int64'>, <Array [1.1, 2.2] type='2 * float64'>)
907944
>>> a.azimuthal.elements
908945
(<Array [1, 2] type='2 * int64'>, <Array [1.1, 2.2] type='2 * float64'>)
909946
"""
@@ -932,6 +969,8 @@ def azimuthal(self) -> AzimuthalAwkward:
932969
>>> a = vector.Array(
933970
... [{"px": 1, "py": 1.1}, {"px": 2, "py": 2.2}],
934971
... )
972+
>>> a.azimuthal
973+
AzimuthalAwkwardXY(<Array [1, 2] type='2 * int64'>, <Array [1.1, 2.2] type='2 * float64'>)
935974
>>> a.azimuthal.elements
936975
(<Array [1, 2] type='2 * int64'>, <Array [1.1, 2.2] type='2 * float64'>)
937976
"""
@@ -960,6 +999,8 @@ def azimuthal(self) -> AzimuthalAwkward:
960999
>>> a = vector.Array(
9611000
... [{"x": 1, "y": 1.1, "z": 0.1}, {"x": 2, "y": 2.2, "z": 0.2}],
9621001
... )
1002+
>>> a.azimuthal
1003+
AzimuthalAwkwardXY(<Array [1, 2] type='2 * int64'>, <Array [1.1, 2.2] type='2 * float64'>)
9631004
>>> a.azimuthal.elements
9641005
(<Array [1, 2] type='2 * int64'>, <Array [1.1, 2.2] type='2 * float64'>)
9651006
"""
@@ -978,6 +1019,8 @@ def longitudinal(self) -> LongitudinalAwkward:
9781019
>>> a = vector.Array(
9791020
... [{"x": 1, "y": 1.1, "z": 0.1}, {"x": 2, "y": 2.2, "z": 0.2}],
9801021
... )
1022+
>>> a.longitudinal
1023+
LongitudinalAwkwardZ(<Array [0.1, 0.2] type='2 * float64'>,)
9811024
>>> a.longitudinal.elements
9821025
(<Array [0.1, 0.2] type='2 * float64'>,)
9831026
"""
@@ -1006,6 +1049,8 @@ def azimuthal(self) -> AzimuthalAwkward:
10061049
>>> a = vector.Array(
10071050
... [{"px": 1, "py": 1.1, "z": 0.1}, {"px": 2, "py": 2.2, "z": 0.2}],
10081051
... )
1052+
>>> a.azimuthal
1053+
AzimuthalAwkwardXY(<Array [1, 2] type='2 * int64'>, <Array [1.1, 2.2] type='2 * float64'>)
10091054
>>> a.azimuthal.elements
10101055
(<Array [1, 2] type='2 * int64'>, <Array [1.1, 2.2] type='2 * float64'>)
10111056
"""
@@ -1024,6 +1069,8 @@ def longitudinal(self) -> LongitudinalAwkward:
10241069
>>> a = vector.Array(
10251070
... [{"px": 1, "py": 1.1, "z": 0.1}, {"px": 2, "py": 2.2, "z": 0.2}],
10261071
... )
1072+
>>> a.longitudinal
1073+
LongitudinalAwkwardZ(<Array [0.1, 0.2] type='2 * float64'>,)
10271074
>>> a.longitudinal.elements
10281075
(<Array [0.1, 0.2] type='2 * float64'>,)
10291076
"""
@@ -1052,6 +1099,8 @@ def azimuthal(self) -> AzimuthalAwkward:
10521099
>>> a = vector.Array(
10531100
... [{"x": 1, "y": 1.1, "z": 0.1, "tau": 1}, {"x": 2, "y": 2.2, "z": 0.2, "tau": 3}],
10541101
... )
1102+
>>> a.azimuthal
1103+
AzimuthalAwkwardXY(<Array [1, 2] type='2 * int64'>, <Array [1.1, 2.2] type='2 * float64'>)
10551104
>>> a.azimuthal.elements
10561105
(<Array [1, 2] type='2 * int64'>, <Array [1.1, 2.2] type='2 * float64'>)
10571106
"""
@@ -1070,6 +1119,8 @@ def longitudinal(self) -> LongitudinalAwkward:
10701119
>>> a = vector.Array(
10711120
... [{"x": 1, "y": 1.1, "z": 0.1, "tau": 1}, {"x": 2, "y": 2.2, "z": 0.2, "tau": 3}],
10721121
... )
1122+
>>> a.longitudinal
1123+
LongitudinalAwkwardZ(<Array [0.1, 0.2] type='2 * float64'>,)
10731124
>>> a.longitudinal.elements
10741125
(<Array [0.1, 0.2] type='2 * float64'>,)
10751126
"""
@@ -1088,6 +1139,8 @@ def temporal(self) -> TemporalAwkward:
10881139
>>> a = vector.Array(
10891140
... [{"x": 1, "y": 1.1, "z": 0.1, "tau": 1}, {"x": 2, "y": 2.2, "z": 0.2, "tau": 3}],
10901141
... )
1142+
>>> a.temporal
1143+
TemporalAwkwardTau(<Array [1, 3] type='2 * int64'>,)
10911144
>>> a.temporal.elements
10921145
(<Array [1, 3] type='2 * int64'>,)
10931146
"""
@@ -1116,6 +1169,8 @@ def azimuthal(self) -> AzimuthalAwkward:
11161169
>>> a = vector.Array(
11171170
... [{"px": 1, "py": 1.1, "z": 0.1, "m": 1}, {"px": 2, "py": 2.2, "z": 0.2, "m": 3}],
11181171
... )
1172+
>>> a.azimuthal
1173+
AzimuthalAwkwardXY(<Array [1, 2] type='2 * int64'>, <Array [1.1, 2.2] type='2 * float64'>)
11191174
>>> a.azimuthal.elements
11201175
(<Array [1, 2] type='2 * int64'>, <Array [1.1, 2.2] type='2 * float64'>)
11211176
"""
@@ -1134,6 +1189,8 @@ def longitudinal(self) -> LongitudinalAwkward:
11341189
>>> a = vector.Array(
11351190
... [{"px": 1, "py": 1.1, "z": 0.1, "m": 1}, {"px": 2, "py": 2.2, "z": 0.2, "m": 3}],
11361191
... )
1192+
>>> a.longitudinal
1193+
LongitudinalAwkwardZ(<Array [0.1, 0.2] type='2 * float64'>,)
11371194
>>> a.longitudinal.elements
11381195
(<Array [0.1, 0.2] type='2 * float64'>,)
11391196
"""
@@ -1151,6 +1208,8 @@ def temporal(self) -> TemporalAwkward:
11511208
>>> a = vector.Array(
11521209
... [{"px": 1, "py": 1.1, "z": 0.1, "m": 1}, {"px": 2, "py": 2.2, "z": 0.2, "m": 3}],
11531210
... )
1211+
>>> a.temporal
1212+
TemporalAwkwardTau(<Array [1, 3] type='2 * int64'>,)
11541213
>>> a.temporal.elements
11551214
(<Array [1, 3] type='2 * int64'>,)
11561215
"""

0 commit comments

Comments
 (0)