-
Notifications
You must be signed in to change notification settings - Fork 33
fix: vectors should be demoted to the lowest dimension #413
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
c6a2304
58caade
2441f86
968f4f7
86ca566
1aaa2d1
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -575,19 +575,21 @@ def elements(self) -> tuple[ArrayOrRecord]: | |
|
||
|
||
def _class_to_name(cls: type[VectorProtocol]) -> str: | ||
# respect the type of classes inheriting VectorAwkward classes | ||
is_vector = "vector.backends" in cls.__module__ | ||
if issubclass(cls, Momentum): | ||
if issubclass(cls, Vector2D): | ||
return "Momentum2D" | ||
return "Momentum2D" if is_vector else cls.__name__[:-5] | ||
if issubclass(cls, Vector3D): | ||
return "Momentum3D" | ||
return "Momentum3D" if is_vector else cls.__name__[:-5] | ||
if issubclass(cls, Vector4D): | ||
return "Momentum4D" | ||
return "Momentum4D" if is_vector else cls.__name__[:-5] | ||
if issubclass(cls, Vector2D): | ||
return "Vector2D" | ||
return "Vector2D" if is_vector else cls.__name__[:-5] | ||
if issubclass(cls, Vector3D): | ||
return "Vector3D" | ||
return "Vector3D" if is_vector else cls.__name__[:-5] | ||
if issubclass(cls, Vector4D): | ||
return "Vector4D" | ||
return "Vector4D" if is_vector else cls.__name__[:-5] | ||
Comment on lines
+578
to
+592
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This function was always returning Vector classes; hence, performing vector methods on subclasses of vector mixins always returned a scikit-hep vector - In [5]: a = ak.zip(
...: {
...: "pt": [10.0, 20.0, 30.0],
...: "eta": [0.0, 1.1, 2.2],
...: "phi": [0.1, 0.9, -1.1],
...: "mass": [1.0, 1.0, 1.0],
...: },
...: with_name="PtEtaPhiMLorentzVector",
...: behavior=vector.behavior,
...: )
In [6]: a
Out[6]: <PtEtaPhiMLorentzVectorArray [{pt: 10, eta: 0, phi: 0.1, ...}, ...] type='3...'>
In [7]: a + a
Out[7]: <LorentzVectorArray [{x: 19.9, y: 2, z: 0, ...}, ..., {...}] type='3 * Lore...'>
In [8]: a - a
Out[8]: <LorentzVectorArray [{x: 0, y: 0, z: 0, t: 0}, ..., {...}] type='3 * Lorent...'>
In [9]: a.to_Vector4D()
Out[9]: <PtEtaPhiMLorentzVectorArray [{pt: 10, eta: 0, phi: 0.1, ...}, ...] type='3...'>
In [10]: a.to_Vector3D()
Out[10]: <MomentumArray3D [{rho: 10, phi: 0.1, ...}, ..., {...}] type='3 * Momentum3...'>
In [11]: a.to_beta3()
Out[11]: <MomentumArray3D [{rho: 0.995, phi: 0.1, ...}, ...] type='3 * Momentum3D[rh...'> Now, everything works well, but only if a user names the classes as "{anything}Array" - In [2]: from coffea.nanoevents.methods import vector
...: import awkward as ak
...: a = ak.zip(
...: {
...: "x": [[1, 2], [], [3], [4]],
...: "y": [[5, 6], [], [7], [8]],
...: "z": [[9, 10], [], [11], [12]],
...: "t": [[50, 51], [], [52], [53]],
...: },
...: with_name="LorentzVector",
...: behavior=vector.behavior,
...: )
...: a.boost(-a.boostvec)
Out[2]: <LorentzVectorArray [[{x: 0, y: -8.88e-16, ...}, ...], ...] type='4 * var *...'>
In [3]: a + a
Out[3]: <LorentzVectorArray [[{x: 2, y: 10, z: 18, ...}, ...], ...] type='4 * var *...'>
In [4]: a.to_Vector2D()
Out[4]: <TwoVectorArray [[{x: 1, y: 5}, {...}], ..., [{...}]] type='4 * var * TwoVe...'>
In [5]: a.to_Vector3D()
Out[5]: <ThreeVectorArray [[{x: 1, y: 5, z: 9}, {...}], ...] type='4 * var * ThreeV...'> Maybe I should mention this somewhere in the docs. |
||
|
||
raise AssertionError(repr(cls)) | ||
|
||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This now generalises to classes inheriting vector mixins.