From e0af53e595156078fcf83164f40dc1788a9f4718 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Sun, 14 Jan 2024 14:07:08 -0600 Subject: [PATCH 01/46] first try at backing nanoevents.methods.vector with scikit-hep vector --- pyproject.toml | 1 + src/coffea/nanoevents/methods/vector.py | 359 +++++++++++------------- tests/test_nanoevents_vector.py | 34 ++- 3 files changed, 181 insertions(+), 213 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 97753217b..6e017b38b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,6 +42,7 @@ dependencies = [ "dask[array]>=2023.4.0", "dask-awkward>=2024.3.0", "dask-histogram>=2024.2.0", + "vector>=1.1.1", "correctionlib>=2.3.3", "pyarrow>=6.0.0", "fsspec-xrootd>=0.2.3", diff --git a/src/coffea/nanoevents/methods/vector.py b/src/coffea/nanoevents/methods/vector.py index 515c5ef7c..d1fe10e35 100644 --- a/src/coffea/nanoevents/methods/vector.py +++ b/src/coffea/nanoevents/methods/vector.py @@ -50,7 +50,13 @@ import numba import numpy import pytz +import vector from dask_awkward import dask_method +from vector.backends.awkward import ( + MomentumAwkward2D, + MomentumAwkward3D, + MomentumAwkward4D, +) from coffea.util import deprecate @@ -109,10 +115,39 @@ def delta_r(eta1, phi1, eta2, phi2): behavior = {} +behavior.update(vector.backends.awkward.behavior) + + +class TwoVectorArray: + pass + + +class PolarTwoVectorArray: + pass + + +class ThreeVectorArray: + pass + + +class SphericalThreeVectorArray: + pass + + +class LorentzVectorArray: + pass + + +class PtEtaPhiMLorentzVectorArray: + pass + + +class PtEtaPhiELorentzVectorArray: + pass @awkward.mixin_class(behavior) -class TwoVector: +class TwoVector(MomentumAwkward2D): """A cartesian 2-dimensional vector A heavy emphasis towards a momentum vector interpretation is assumed, hence @@ -127,40 +162,12 @@ def r(self): :math:`\sqrt{x^2+y^2}` """ - return numpy.sqrt(self.r2) - - @property - def phi(self): - r"""Polar angle relative to X axis - - :math:`\text{arctan2}(y, x)` - """ - return numpy.arctan2(self.y, self.x) - - @property - def px(self): - """Alias for `x`""" - return self.x - - @property - def py(self): - """Alias for `y`""" - return self.y + return self.rho @property def r2(self): """Squared `r`""" - return self.x * self.x + self.y * self.y - - @property - def pt2(self): - """Alias for `r2`""" - return self.r2 - - @property - def pt(self): - """Alias for `r`""" - return self.r + return self.rho2 @awkward.mixin_class_method(numpy.absolute) def absolute(self): @@ -240,11 +247,7 @@ def delta_phi(self, other): Returns a value within [-pi, pi) """ - return delta_phi(self.phi, other.phi) - - def dot(self, other): - """Compute the dot product of two vectors""" - return self.x * other.x + self.y * other.y + return self.deltaphi(other) @property def unit(self): @@ -256,47 +259,10 @@ def unit(self): class PolarTwoVector(TwoVector): """A polar coordinate 2-dimensional vector - This mixin class requires the parent class to provide items `r` and `phi`. + This mixin class requires the parent class to provide items `rho` and `phi`. Some additional properties are overridden for performance """ - @property - def x(self): - r"""Cartesian x value - - :math:`r \cos{\phi}` - """ - return self.r * numpy.cos(self.phi) - - @property - def y(self): - r"""Cartesian y value - - :math:`r \sin{\phi}` - """ - return self.r * numpy.sin(self.phi) - - @property - def r(self): - r"""Distance from origin in XY plane - - :math:`\sqrt{x^2+y^2}` - """ - return self["r"] - - @property - def phi(self): - r"""Azimuthal angle relative to X axis in XY plane - - :math:`\text{arctan2}(y, x)` - """ - return self["phi"] - - @property - def r2(self): - """Squared `r`""" - return self.r * self.r - @awkward.mixin_class_method(numpy.multiply, {numbers.Number}) def multiply(self, other): """Multiply this vector by a scalar elementwise using using `x` and `y` components @@ -305,7 +271,7 @@ def multiply(self, other): """ return awkward.zip( { - "r": self.r * abs(other), + "rho": self.r * abs(other), "phi": self.phi % (2 * numpy.pi) - (numpy.pi * (other < 0)), }, with_name="PolarTwoVector", @@ -316,14 +282,14 @@ def multiply(self, other): def negative(self): """Returns the negative of the vector""" return awkward.zip( - {"r": self.r, "phi": self.phi % (2 * numpy.pi) - numpy.pi}, + {"rho": self.r, "phi": self.phi % (2 * numpy.pi) - numpy.pi}, with_name="PolarTwoVector", behavior=self.behavior, ) @awkward.mixin_class(behavior) -class ThreeVector(TwoVector): +class ThreeVector(MomentumAwkward3D): """A cartesian 3-dimensional vector A heavy emphasis towards a momentum vector interpretation is assumed. @@ -331,41 +297,18 @@ class ThreeVector(TwoVector): """ @property - def pz(self): - """Alias for `z`""" - return self.z - - @property - def rho2(self): - """Squared `rho`""" - return self.r2 + self.z * self.z - - @property - def rho(self): - r"""Distance from origin in 3D - - :math:`\sqrt{x^2+y^2+z^2} = \sqrt{r^2+z^2}` - """ - return numpy.sqrt(self.rho2) - - @property - def theta(self): - r"""Inclination angle from XY plane + def r(self): + r"""Distance from origin in XY plane - :math:`\text{arctan2}(r, z)` + :math:`\sqrt{x^2+y^2}` """ - return numpy.arctan2(self.r, self.z) + return self.rho @property - def p2(self): - """Squared `p`""" + def r2(self): + """Squared `r`""" return self.rho2 - @property - def p(self): - """Alias for `rho`""" - return self.rho - @awkward.mixin_class_method(numpy.absolute) def absolute(self): """Returns magnitude of the 3D vector @@ -383,18 +326,33 @@ def negative(self): behavior=self.behavior, ) - @awkward.mixin_class_method(numpy.add, {"ThreeVector"}) + @awkward.mixin_class_method( + numpy.add, {"ThreeVector", "TwoVector", "PolarTwoVector"} + ) def add(self, other): """Add two vectors together elementwise using `x`, `y`, and `z` components""" + other_3d = other.to_Vector3D() return awkward.zip( - {"x": self.x + other.x, "y": self.y + other.y, "z": self.z + other.z}, + { + "x": self.x + other_3d.x, + "y": self.y + other_3d.y, + "z": self.z + other_3d.z, + }, with_name="ThreeVector", behavior=self.behavior, ) + @awkward.mixin_class_method(numpy.divide, {numbers.Number}) + def divide(self, other): + """Divide this vector by a scalar elementwise using its cartesian components + This is realized by using the multiplication functionality""" + return self.multiply(1 / other) + @awkward.mixin_class_method( numpy.subtract, { + "TwoVector", + "PolarTwoVector", "ThreeVector", "SphericalThreeVector", "LorentzVector", @@ -405,8 +363,14 @@ def add(self, other): ) def subtract(self, other): """Subtract a vector from another elementwise using `x`, `y`, and `z` components""" + other_3d = other.to_Vector3D() + return awkward.zip( - {"x": self.x - other.x, "y": self.y - other.y, "z": self.z - other.z}, + { + "x": self.x - other_3d.x, + "y": self.y - other_3d.y, + "z": self.z - other_3d.z, + }, with_name="ThreeVector", behavior=self.behavior, ) @@ -448,6 +412,13 @@ def cross(self, other): behavior=self.behavior, ) + def delta_phi(self, other): + """Compute difference in angle between two vectors + + Returns a value within [-pi, pi) + """ + return self.deltaphi(other) + @property def unit(self): """Unit vector, a vector of length 1 pointing in the same direction""" @@ -455,7 +426,7 @@ def unit(self): @awkward.mixin_class(behavior) -class SphericalThreeVector(ThreeVector, PolarTwoVector): +class SphericalThreeVector(ThreeVector): """A spherical coordinate 3-dimensional vector This mixin class requires the parent class to provide items `rho`, `theta`, and `phi`. @@ -468,42 +439,8 @@ def r(self): :math:`\sqrt{x^2+y^2} = \rho \sin(\theta)` """ - return self.rho * numpy.sin(self.theta) - - @property - def z(self): - r"""Cartesian z value - - :math:`\rho \cos(\theta)` - """ - return self.rho * numpy.cos(self.theta) - - @property - def rho(self): - r"""Distance from origin in 3D - - :math:`\sqrt{x^2+y^2+z^2} = \sqrt{r^2+z^2}` - """ - return self["rho"] - - @property - def theta(self): - r"""Inclination angle from XY plane - - :math:`\text{arctan2}(r, z)` - """ - return self["theta"] - - @property - def p(self): - """Alias for `rho`""" return self.rho - @property - def p2(self): - """Squared `p`""" - return self.rho * self.rho - @awkward.mixin_class_method(numpy.multiply, {numbers.Number}) def multiply(self, other): """Multiply this vector by a scalar elementwise using `x`, `y`, and `z` components @@ -561,7 +498,7 @@ def _nearest_core(x, y, axis, metric, return_metric, threshold): @awkward.mixin_class(behavior) -class LorentzVector(ThreeVector): +class LorentzVector(MomentumAwkward4D): """A cartesian Lorentz vector A heavy emphasis towards a momentum vector interpretation is assumed. @@ -603,29 +540,51 @@ def absolute(self): """ return self.mass - @awkward.mixin_class_method(numpy.add, {"LorentzVector"}) + @awkward.mixin_class_method( + numpy.add, + { + "LorentzVector", + "ThreeVector", + "SphericalThreeVector", + "TwoVector", + "PolarTwoVector", + }, + ) def add(self, other): """Add two vectors together elementwise using `x`, `y`, `z`, and `t` components""" + other_4d = other.to_Vector4D() return awkward.zip( { - "x": self.x + other.x, - "y": self.y + other.y, - "z": self.z + other.z, - "t": self.t + other.t, + "x": self.x + other_4d.x, + "y": self.y + other_4d.y, + "z": self.z + other_4d.z, + "t": self.t + other_4d.t, }, with_name="LorentzVector", behavior=self.behavior, ) - @awkward.mixin_class_method(numpy.subtract, {"LorentzVector"}, transpose=False) + @awkward.mixin_class_method( + numpy.subtract, + { + "LorentzVector", + "ThreeVector", + "SphericalThreeVector", + "TwoVector", + "PolarTwoVector", + }, + transpose=False, + ) def subtract(self, other): """Subtract a vector from another elementwise using `x`, `y`, `z`, and `t` components""" + other_4d = other.to_Vector4D() + return awkward.zip( { - "x": self.x - other.x, - "y": self.y - other.y, - "z": self.z - other.z, - "t": self.t - other.t, + "x": self.x - other_4d.x, + "y": self.y - other_4d.y, + "z": self.z - other_4d.z, + "t": self.t - other_4d.t, }, with_name="LorentzVector", behavior=self.behavior, @@ -658,16 +617,29 @@ def multiply(self, other): behavior=self.behavior, ) + @awkward.mixin_class_method(numpy.divide, {numbers.Number}) + def divide(self, other): + """Divide this vector by a scalar elementwise using its cartesian components + This is realized by using the multiplication functionality""" + return self.multiply(1 / other) + def delta_r2(self, other): """Squared `delta_r`""" - return delta_r(self.eta, self.phi, other.eta, other.phi) ** 2 + return self.deltaR2(other) def delta_r(self, other): r"""Distance between two Lorentz vectors in (eta,phi) plane :math:`\sqrt{\Delta\eta^2 + \Delta\phi^2}` """ - return delta_r(self.eta, self.phi, other.eta, other.phi) + return self.deltaR(other) + + def delta_phi(self, other): + """Compute difference in angle between two vectors + + Returns a value within [-pi, pi) + """ + return self.deltaphi(other) @awkward.mixin_class_method(numpy.negative) def negative(self): @@ -687,12 +659,6 @@ def pvec(self): behavior=self.behavior, ) - @property - def rapidity(self): - pz = self.z - e = self.energy - return 0.5 * (numpy.log(e + pz) - numpy.log(e - pz)) - @property def boostvec(self): """The `x`, `y` and `z` components divided by `t` as a `ThreeVector` @@ -700,40 +666,7 @@ def boostvec(self): This can be used for boosting. For cases where `|t| <= rho`, this returns the unit vector. """ - rho = self.rho - t = self.t - with numpy.errstate(divide="ignore"): - out = self.pvec * awkward.where( - rho == 0, 0, awkward.where(abs(t) <= rho, 1 / rho, 1 / t) - ) - return out - - def boost(self, other): - """Apply a Lorentz boost given by the `ThreeVector` `other` and return it - - Note that this follows the convention that, for example in order to boost - a vector into its own rest frame, one needs to use the negative of its `boostvec` - """ - b2 = other.rho2 - gamma = (1 - b2) ** (-0.5) - mask = b2 == 0 - b2 = awkward.where(mask, 1, b2) - gamma2 = awkward.where(mask, 0, (gamma - 1) / b2) - - bp = self.dot(other) - t = self.t - v = gamma2 * bp * other + t * gamma * other - - return awkward.zip( - { - "x": self.x + v.x, - "y": self.y + v.y, - "z": self.z + v.z, - "t": gamma * (t + bp), - }, - with_name="LorentzVector", - behavior=self.behavior, - ) + return self.to_beta3() @dask_method def metric_table( @@ -821,7 +754,7 @@ def nearest( @awkward.mixin_class(behavior) -class PtEtaPhiMLorentzVector(LorentzVector, SphericalThreeVector): +class PtEtaPhiMLorentzVector(LorentzVector): """A Lorentz vector using pseudorapidity and mass This mixin class requires the parent class to provide items `pt`, `eta`, `phi`, and `mass`. @@ -946,7 +879,7 @@ def negative(self): @awkward.mixin_class(behavior) -class PtEtaPhiELorentzVector(LorentzVector, SphericalThreeVector): +class PtEtaPhiELorentzVector(LorentzVector): """A Lorentz vector using pseudorapidity and energy This mixin class requires the parent class to provide items `pt`, `eta`, `phi`, and `energy`. @@ -1061,6 +994,34 @@ def negative(self): ) +TwoVectorArray.ProjectionClass2D = TwoVectorArray +TwoVectorArray.ProjectionClass3D = ThreeVectorArray +TwoVectorArray.ProjectionClass4D = LorentzVectorArray + +PolarTwoVectorArray.ProjectionClass2D = PolarTwoVectorArray +PolarTwoVectorArray.ProjectionClass3D = SphericalThreeVectorArray +PolarTwoVectorArray.ProjectionClass4D = LorentzVectorArray + +ThreeVectorArray.ProjectionClass2D = TwoVectorArray +ThreeVectorArray.ProjectionClass3D = ThreeVectorArray +ThreeVectorArray.ProjectionClass4D = LorentzVectorArray + +SphericalThreeVectorArray.ProjectionClass2D = TwoVectorArray +SphericalThreeVectorArray.ProjectionClass3D = SphericalThreeVectorArray +SphericalThreeVectorArray.ProjectionClass4D = LorentzVectorArray + +LorentzVectorArray.ProjectionClass2D = TwoVectorArray +LorentzVectorArray.ProjectionClass3D = ThreeVectorArray +LorentzVectorArray.ProjectionClass4D = LorentzVectorArray + +PtEtaPhiMLorentzVectorArray.ProjectionClass2D = TwoVectorArray +PtEtaPhiMLorentzVectorArray.ProjectionClass3D = ThreeVectorArray +PtEtaPhiMLorentzVectorArray.ProjectionClass4D = LorentzVectorArray + +PtEtaPhiELorentzVectorArray.ProjectionClass2D = TwoVectorArray +PtEtaPhiELorentzVectorArray.ProjectionClass3D = ThreeVectorArray +PtEtaPhiELorentzVectorArray.ProjectionClass4D = LorentzVectorArray + __all__ = [ "TwoVector", "PolarTwoVector", diff --git a/tests/test_nanoevents_vector.py b/tests/test_nanoevents_vector.py index 1a83b9960..e65f5d741 100644 --- a/tests/test_nanoevents_vector.py +++ b/tests/test_nanoevents_vector.py @@ -58,7 +58,7 @@ def test_two_vector(): def test_polar_two_vector(): a = ak.zip( { - "r": [[1, 2], [], [3], [4]], + "rho": [[1, 2], [], [3], [4]], "phi": [[0.3, 0.4], [], [0.5], [0.6]], }, with_name="PolarTwoVector", @@ -67,9 +67,9 @@ def test_polar_two_vector(): assert record_arrays_equal( a * 2, - ak.zip({"r": [[2, 4], [], [6], [8]], "phi": [[0.3, 0.4], [], [0.5], [0.6]]}), + ak.zip({"rho": [[2, 4], [], [6], [8]], "phi": [[0.3, 0.4], [], [0.5], [0.6]]}), ) - assert ak.all((a * (-2)).r == [[2, 4], [], [6], [8]]) + assert ak.all((a * (-2)).rho == [[2, 4], [], [6], [8]]) assert ak.all( (a * (-2)).phi - ak.Array( @@ -80,7 +80,7 @@ def test_polar_two_vector(): assert record_arrays_equal( a / 2, ak.zip( - {"r": [[0.5, 1], [], [1.5], [2]], "phi": [[0.3, 0.4], [], [0.5], [0.6]]} + {"rho": [[0.5, 1], [], [1.5], [2]], "phi": [[0.3, 0.4], [], [0.5], [0.6]]} ), ) @@ -444,12 +444,18 @@ def test_lorentz_vector_numba(a_dtype, b_dtype): 110.66616465749593, 110.68423555321688, ] - assert pytest.approx(a.delta_phi(b), abs=1e-7) == [ - 0.03369510734601633, - -0.1798534997924781, - 0.33292327383538156, - 0.6078019961139605, - ] + + computed_dphi = a.delta_phi(b).to_numpy() + + assert pytest.approx(computed_dphi, abs=1e-6) == np.array( + [ + 0.03369510734601633, + -0.1798534997924781, + 0.33292327383538156, + 0.6078019961139605, + ], + dtype=computed_dphi.dtype, + ) @pytest.mark.parametrize( @@ -519,7 +525,7 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): ) elif twocoord == "PolarTwoVector": c = ak.zip( - {"r": [-10.0, 13.0, 15.0], "phi": [1.22, -1.0, 1.0]}, + {"rho": [-10.0, 13.0, 15.0], "phi": [1.22, -1.0, 1.0]}, with_name=twocoord, behavior=vector.behavior, ) @@ -532,9 +538,9 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): assert record_arrays_equal(a.delta_phi(c), c.delta_phi(a)) assert record_arrays_equal(b.delta_phi(c), c.delta_phi(b)) - assert record_arrays_equal(a - b, -(b - a)) - assert record_arrays_equal(a - c, -(c - a)) - assert record_arrays_equal(b - c, -(c - b)) + assert record_arrays_equal((a - b).to_Vector3D(), -(b - a)) + assert record_arrays_equal((a - c).to_Vector2D(), -(c - a)) + assert record_arrays_equal((b - c).to_Vector2D(), -(c - b)) @pytest.mark.parametrize("optimization_enabled", [True, False]) From 3d5ad8d7e3f035eae9ae08667c551f5c2dbdc8fb Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Thu, 18 Jan 2024 13:52:06 -0600 Subject: [PATCH 02/46] fix math in doc --- src/coffea/nanoevents/methods/vector.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/coffea/nanoevents/methods/vector.py b/src/coffea/nanoevents/methods/vector.py index d1fe10e35..7b5293b79 100644 --- a/src/coffea/nanoevents/methods/vector.py +++ b/src/coffea/nanoevents/methods/vector.py @@ -663,7 +663,7 @@ def pvec(self): def boostvec(self): """The `x`, `y` and `z` components divided by `t` as a `ThreeVector` - This can be used for boosting. For cases where `|t| <= rho`, this + This can be used for boosting. For cases where `|t| <= r`, this returns the unit vector. """ return self.to_beta3() From a24ace44deef3e8c04b7a9589d325207edb8f606 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Sun, 21 Jan 2024 12:20:06 -0600 Subject: [PATCH 03/46] use noqa instead of actually satisfying flake - more readable --- src/coffea/nanoevents/methods/vector.py | 70 ++++++++----------------- 1 file changed, 21 insertions(+), 49 deletions(-) diff --git a/src/coffea/nanoevents/methods/vector.py b/src/coffea/nanoevents/methods/vector.py index 7b5293b79..3ca68b3fc 100644 --- a/src/coffea/nanoevents/methods/vector.py +++ b/src/coffea/nanoevents/methods/vector.py @@ -118,34 +118,6 @@ def delta_r(eta1, phi1, eta2, phi2): behavior.update(vector.backends.awkward.behavior) -class TwoVectorArray: - pass - - -class PolarTwoVectorArray: - pass - - -class ThreeVectorArray: - pass - - -class SphericalThreeVectorArray: - pass - - -class LorentzVectorArray: - pass - - -class PtEtaPhiMLorentzVectorArray: - pass - - -class PtEtaPhiELorentzVectorArray: - pass - - @awkward.mixin_class(behavior) class TwoVector(MomentumAwkward2D): """A cartesian 2-dimensional vector @@ -994,33 +966,33 @@ def negative(self): ) -TwoVectorArray.ProjectionClass2D = TwoVectorArray -TwoVectorArray.ProjectionClass3D = ThreeVectorArray -TwoVectorArray.ProjectionClass4D = LorentzVectorArray +TwoVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +TwoVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +TwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -PolarTwoVectorArray.ProjectionClass2D = PolarTwoVectorArray -PolarTwoVectorArray.ProjectionClass3D = SphericalThreeVectorArray -PolarTwoVectorArray.ProjectionClass4D = LorentzVectorArray +PolarTwoVectorArray.ProjectionClass2D = PolarTwoVectorArray # noqa: F821 +PolarTwoVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 +PolarTwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -ThreeVectorArray.ProjectionClass2D = TwoVectorArray -ThreeVectorArray.ProjectionClass3D = ThreeVectorArray -ThreeVectorArray.ProjectionClass4D = LorentzVectorArray +ThreeVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +ThreeVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +ThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -SphericalThreeVectorArray.ProjectionClass2D = TwoVectorArray -SphericalThreeVectorArray.ProjectionClass3D = SphericalThreeVectorArray -SphericalThreeVectorArray.ProjectionClass4D = LorentzVectorArray +SphericalThreeVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +SphericalThreeVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 +SphericalThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -LorentzVectorArray.ProjectionClass2D = TwoVectorArray -LorentzVectorArray.ProjectionClass3D = ThreeVectorArray -LorentzVectorArray.ProjectionClass4D = LorentzVectorArray +LorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +LorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +LorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -PtEtaPhiMLorentzVectorArray.ProjectionClass2D = TwoVectorArray -PtEtaPhiMLorentzVectorArray.ProjectionClass3D = ThreeVectorArray -PtEtaPhiMLorentzVectorArray.ProjectionClass4D = LorentzVectorArray +PtEtaPhiMLorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +PtEtaPhiMLorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +PtEtaPhiMLorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -PtEtaPhiELorentzVectorArray.ProjectionClass2D = TwoVectorArray -PtEtaPhiELorentzVectorArray.ProjectionClass3D = ThreeVectorArray -PtEtaPhiELorentzVectorArray.ProjectionClass4D = LorentzVectorArray +PtEtaPhiELorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +PtEtaPhiELorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +PtEtaPhiELorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 __all__ = [ "TwoVector", From 6cc680090e2b41bad696411153b9d812c7df190c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 21 Jan 2024 18:20:21 +0000 Subject: [PATCH 04/46] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/coffea/nanoevents/methods/vector.py | 40 ++++++++++++------------- 1 file changed, 20 insertions(+), 20 deletions(-) diff --git a/src/coffea/nanoevents/methods/vector.py b/src/coffea/nanoevents/methods/vector.py index 3ca68b3fc..88d505dfc 100644 --- a/src/coffea/nanoevents/methods/vector.py +++ b/src/coffea/nanoevents/methods/vector.py @@ -966,32 +966,32 @@ def negative(self): ) -TwoVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 -TwoVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 -TwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +TwoVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +TwoVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +TwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -PolarTwoVectorArray.ProjectionClass2D = PolarTwoVectorArray # noqa: F821 -PolarTwoVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 -PolarTwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +PolarTwoVectorArray.ProjectionClass2D = PolarTwoVectorArray # noqa: F821 +PolarTwoVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 +PolarTwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -ThreeVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 -ThreeVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 -ThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +ThreeVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +ThreeVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +ThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -SphericalThreeVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 -SphericalThreeVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 -SphericalThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +SphericalThreeVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +SphericalThreeVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 +SphericalThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -LorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 -LorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 -LorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +LorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +LorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +LorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -PtEtaPhiMLorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 -PtEtaPhiMLorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 -PtEtaPhiMLorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +PtEtaPhiMLorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +PtEtaPhiMLorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +PtEtaPhiMLorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -PtEtaPhiELorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 -PtEtaPhiELorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +PtEtaPhiELorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +PtEtaPhiELorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 PtEtaPhiELorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 __all__ = [ From 1226f968efb4f2c3bfc5942ff23280a392b0b7ea Mon Sep 17 00:00:00 2001 From: Nick Smith Date: Fri, 26 Jan 2024 15:35:11 -0600 Subject: [PATCH 05/46] Improve tests of vector --- tests/test_nanoevents_vector.py | 93 ++++++++++++++++++--------------- 1 file changed, 51 insertions(+), 42 deletions(-) diff --git a/tests/test_nanoevents_vector.py b/tests/test_nanoevents_vector.py index e65f5d741..6d933fe35 100644 --- a/tests/test_nanoevents_vector.py +++ b/tests/test_nanoevents_vector.py @@ -2,15 +2,24 @@ import numpy as np import pytest +from numpy.testing import assert_allclose from coffea.nanoevents.methods import vector ATOL = 1e-8 -def record_arrays_equal(a, b): - return (ak.fields(a) == ak.fields(b)) and all( - ak.all(a[f] == b[f]) for f in ak.fields(a) - ) +def assert_record_arrays_equal(a, b, check_type=False): + if check_type: + assert type(a) == type(b) + assert ak.fields(a) == ak.fields(b) + assert all(ak.all(a[f] == b[f]) for f in ak.fields(a)) + + +def assert_awkward_allclose(actual, desired): + flat_actual = ak.flatten(actual, axis=None) + flat_desired = ak.flatten(desired, axis=None) + # we should check None values, but not used in these tests + assert_allclose(flat_actual, flat_desired) def test_two_vector(): @@ -25,31 +34,31 @@ def test_two_vector(): behavior=vector.behavior, ) - assert record_arrays_equal( + assert_record_arrays_equal( -a, ak.zip({"x": [[-1, -2], [], [-3], [-4]], "y": [[-5, -6], [], [-7], [-8]]}) ) - assert record_arrays_equal( + assert_record_arrays_equal( a + b, ak.zip({"x": [[12, 14], [], [16], [18]], "y": [[20, 22], [], [24], [26]]}), ) - assert record_arrays_equal( + assert_record_arrays_equal( a - b, ak.zip( {"x": [[-10, -10], [], [-10], [-10]], "y": [[-10, -10], [], [-10], [-10]]} ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a * 2, ak.zip({"x": [[2, 4], [], [6], [8]], "y": [[10, 12], [], [14], [16]]}) ) - assert record_arrays_equal( + assert_record_arrays_equal( a / 2, ak.zip({"x": [[0.5, 1], [], [1.5], [2]], "y": [[2.5, 3], [], [3.5], [4]]}), ) - assert record_arrays_equal(a.dot(b), ak.Array([[86, 120], [], [158], [200]])) - assert record_arrays_equal(b.dot(a), ak.Array([[86, 120], [], [158], [200]])) + assert_awkward_allclose(a.dot(b), ak.Array([[86, 120], [], [158], [200]])) + assert_awkward_allclose(b.dot(a), ak.Array([[86, 120], [], [158], [200]])) assert ak.all(abs(a.unit.r - 1) < ATOL) assert ak.all(abs(a.unit.phi - a.phi) < ATOL) @@ -65,7 +74,7 @@ def test_polar_two_vector(): behavior=vector.behavior, ) - assert record_arrays_equal( + assert_record_arrays_equal( a * 2, ak.zip({"rho": [[2, 4], [], [6], [8]], "phi": [[0.3, 0.4], [], [0.5], [0.6]]}), ) @@ -77,7 +86,7 @@ def test_polar_two_vector(): ) < ATOL ) - assert record_arrays_equal( + assert_record_arrays_equal( a / 2, ak.zip( {"rho": [[0.5, 1], [], [1.5], [2]], "phi": [[0.3, 0.4], [], [0.5], [0.6]]} @@ -86,7 +95,7 @@ def test_polar_two_vector(): assert ak.all(abs((-a).x + a.x) < ATOL) assert ak.all(abs((-a).y + a.y) < ATOL) - assert record_arrays_equal(a * (-1), -a) + assert_record_arrays_equal(a * (-1), -a) assert ak.all(a.unit.phi == a.phi) @@ -111,7 +120,7 @@ def test_three_vector(): behavior=vector.behavior, ) - assert record_arrays_equal( + assert_record_arrays_equal( -a, ak.zip( { @@ -122,7 +131,7 @@ def test_three_vector(): ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a + b, ak.zip( { @@ -132,7 +141,7 @@ def test_three_vector(): } ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a - b, ak.zip( { @@ -142,7 +151,7 @@ def test_three_vector(): } ), ) - assert record_arrays_equal( + assert_record_arrays_equal( b - a, ak.zip( { @@ -153,7 +162,7 @@ def test_three_vector(): ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a * 2, ak.zip( { @@ -163,7 +172,7 @@ def test_three_vector(): } ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a / 2, ak.zip( { @@ -177,7 +186,7 @@ def test_three_vector(): assert ak.all(a.dot(b) == ak.Array([[170, 154], [], [162], [284]])) assert ak.all(b.dot(a) == ak.Array([[170, 154], [], [162], [284]])) - assert record_arrays_equal( + assert_record_arrays_equal( a.cross(b), ak.zip( { @@ -187,7 +196,7 @@ def test_three_vector(): } ), ) - assert record_arrays_equal( + assert_record_arrays_equal( b.cross(a), ak.zip( { @@ -216,7 +225,7 @@ def test_spherical_three_vector(): assert ak.all(abs((-a).x + a.x) < ATOL) assert ak.all(abs((-a).y + a.y) < ATOL) assert ak.all(abs((-a).z + a.z) < ATOL) - assert record_arrays_equal(a * (-1), -a) + assert_record_arrays_equal(a * (-1), -a, check_type=True) def test_lorentz_vector(): @@ -241,7 +250,7 @@ def test_lorentz_vector(): behavior=vector.behavior, ) - assert record_arrays_equal( + assert_record_arrays_equal( -a, ak.zip( { @@ -253,7 +262,7 @@ def test_lorentz_vector(): ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a + b, ak.zip( { @@ -264,7 +273,7 @@ def test_lorentz_vector(): } ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a - b, ak.zip( { @@ -276,7 +285,7 @@ def test_lorentz_vector(): ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a * 2, ak.zip( { @@ -287,7 +296,7 @@ def test_lorentz_vector(): } ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a / 2, ak.zip( { @@ -299,7 +308,7 @@ def test_lorentz_vector(): ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a.pvec, ak.zip( { @@ -344,7 +353,7 @@ def test_pt_eta_phi_m_lorentz_vector(): ) < ATOL ) - assert record_arrays_equal( + assert_record_arrays_equal( a / 2, ak.zip( { @@ -355,7 +364,7 @@ def test_pt_eta_phi_m_lorentz_vector(): } ), ) - assert record_arrays_equal(a * (-1), -a) + assert_record_arrays_equal(a * (-1), -a, check_type=True) boosted = a.boost(-a.boostvec) assert ak.all(abs(boosted.x) < ATOL) @@ -390,7 +399,7 @@ def test_pt_eta_phi_e_lorentz_vector(): ) < ATOL ) - assert record_arrays_equal( + assert_record_arrays_equal( a / 2, ak.zip( { @@ -401,7 +410,7 @@ def test_pt_eta_phi_e_lorentz_vector(): } ), ) - assert record_arrays_equal(a * (-1), -a) + assert_record_arrays_equal(a * (-1), -a, check_type=True) boosted = a.boost(-a.boostvec) assert ak.all(abs(boosted.x) < ATOL) @@ -530,17 +539,17 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): behavior=vector.behavior, ) - assert record_arrays_equal(a + b, b + a) - assert record_arrays_equal(a + c, c + a) - assert record_arrays_equal(b + c, c + b) + assert_record_arrays_equal(a + b, b + a, check_type=True) + assert_record_arrays_equal(a + c, c + a, check_type=True) + assert_record_arrays_equal(b + c, c + b, check_type=True) - assert record_arrays_equal(a.delta_phi(b), b.delta_phi(a)) - assert record_arrays_equal(a.delta_phi(c), c.delta_phi(a)) - assert record_arrays_equal(b.delta_phi(c), c.delta_phi(b)) + assert_allclose(a.delta_phi(b), -b.delta_phi(a)) + assert_allclose(a.delta_phi(c), -c.delta_phi(a)) + assert_allclose(b.delta_phi(c), -c.delta_phi(b)) - assert record_arrays_equal((a - b).to_Vector3D(), -(b - a)) - assert record_arrays_equal((a - c).to_Vector2D(), -(c - a)) - assert record_arrays_equal((b - c).to_Vector2D(), -(c - b)) + assert_record_arrays_equal((a - b), -(b - a), check_type=True) + assert_record_arrays_equal((a - c), -(c - a), check_type=True) + assert_record_arrays_equal((b - c), -(c - b), check_type=True) @pytest.mark.parametrize("optimization_enabled", [True, False]) From 653f8b759e8b01bec1271e09697feac46d4ec7c9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 26 Jan 2024 21:40:14 +0000 Subject: [PATCH 06/46] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/test_nanoevents_vector.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_nanoevents_vector.py b/tests/test_nanoevents_vector.py index 6d933fe35..7f2e69cd8 100644 --- a/tests/test_nanoevents_vector.py +++ b/tests/test_nanoevents_vector.py @@ -1,8 +1,8 @@ import awkward as ak import numpy as np import pytest - from numpy.testing import assert_allclose + from coffea.nanoevents.methods import vector ATOL = 1e-8 From c239da9c02da49573913bf0573ca9cae7c887a26 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Fri, 26 Jan 2024 15:42:38 -0600 Subject: [PATCH 07/46] appease linter --- tests/test_nanoevents_vector.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_nanoevents_vector.py b/tests/test_nanoevents_vector.py index 7f2e69cd8..5d925aa58 100644 --- a/tests/test_nanoevents_vector.py +++ b/tests/test_nanoevents_vector.py @@ -10,7 +10,7 @@ def assert_record_arrays_equal(a, b, check_type=False): if check_type: - assert type(a) == type(b) + assert type(a) is type(b) assert ak.fields(a) == ak.fields(b) assert all(ak.all(a[f] == b[f]) for f in ak.fields(a)) From f82641cd797afc8ee905e5f31d095bef30afce04 Mon Sep 17 00:00:00 2001 From: Nick Smith Date: Mon, 29 Jan 2024 17:08:40 -0600 Subject: [PATCH 08/46] Manually build binary ufunc dispatch tables Eventually we might also want to pawn these off on vector --- src/coffea/nanoevents/methods/vector.py | 100 ++++++++---------------- 1 file changed, 32 insertions(+), 68 deletions(-) diff --git a/src/coffea/nanoevents/methods/vector.py b/src/coffea/nanoevents/methods/vector.py index 88d505dfc..f9676b198 100644 --- a/src/coffea/nanoevents/methods/vector.py +++ b/src/coffea/nanoevents/methods/vector.py @@ -158,7 +158,6 @@ def negative(self): behavior=self.behavior, ) - @awkward.mixin_class_method(numpy.add, {"TwoVector"}) def add(self, other): """Add two vectors together elementwise using `x` and `y` components""" return awkward.zip( @@ -167,18 +166,6 @@ def add(self, other): behavior=self.behavior, ) - @awkward.mixin_class_method( - numpy.subtract, - { - "TwoVector", - "ThreeVector", - "SphericalThreeVector", - "LorentzVector", - "PtEtaPhiMLorentzVector", - "PtEtaPhiELorentzVector", - }, - transpose=False, - ) def subtract(self, other): """Subtract a vector from another elementwise using `x` and `y` components""" return awkward.zip( @@ -298,17 +285,13 @@ def negative(self): behavior=self.behavior, ) - @awkward.mixin_class_method( - numpy.add, {"ThreeVector", "TwoVector", "PolarTwoVector"} - ) def add(self, other): """Add two vectors together elementwise using `x`, `y`, and `z` components""" - other_3d = other.to_Vector3D() return awkward.zip( { - "x": self.x + other_3d.x, - "y": self.y + other_3d.y, - "z": self.z + other_3d.z, + "x": self.x + other.x, + "y": self.y + other.y, + "z": self.z + other.z, }, with_name="ThreeVector", behavior=self.behavior, @@ -320,28 +303,14 @@ def divide(self, other): This is realized by using the multiplication functionality""" return self.multiply(1 / other) - @awkward.mixin_class_method( - numpy.subtract, - { - "TwoVector", - "PolarTwoVector", - "ThreeVector", - "SphericalThreeVector", - "LorentzVector", - "PtEtaPhiMLorentzVector", - "PtEtaPhiELorentzVector", - }, - transpose=False, - ) def subtract(self, other): """Subtract a vector from another elementwise using `x`, `y`, and `z` components""" - other_3d = other.to_Vector3D() return awkward.zip( { - "x": self.x - other_3d.x, - "y": self.y - other_3d.y, - "z": self.z - other_3d.z, + "x": self.x - other.x, + "y": self.y - other.y, + "z": self.z - other.z, }, with_name="ThreeVector", behavior=self.behavior, @@ -512,51 +481,28 @@ def absolute(self): """ return self.mass - @awkward.mixin_class_method( - numpy.add, - { - "LorentzVector", - "ThreeVector", - "SphericalThreeVector", - "TwoVector", - "PolarTwoVector", - }, - ) def add(self, other): """Add two vectors together elementwise using `x`, `y`, `z`, and `t` components""" - other_4d = other.to_Vector4D() return awkward.zip( { - "x": self.x + other_4d.x, - "y": self.y + other_4d.y, - "z": self.z + other_4d.z, - "t": self.t + other_4d.t, + "x": self.x + other.x, + "y": self.y + other.y, + "z": self.z + other.z, + "t": self.t + other.t, }, with_name="LorentzVector", behavior=self.behavior, ) - @awkward.mixin_class_method( - numpy.subtract, - { - "LorentzVector", - "ThreeVector", - "SphericalThreeVector", - "TwoVector", - "PolarTwoVector", - }, - transpose=False, - ) def subtract(self, other): """Subtract a vector from another elementwise using `x`, `y`, `z`, and `t` components""" - other_4d = other.to_Vector4D() return awkward.zip( { - "x": self.x - other_4d.x, - "y": self.y - other_4d.y, - "z": self.z - other_4d.z, - "t": self.t - other_4d.t, + "x": self.x - other.x, + "y": self.y - other.y, + "z": self.z - other.z, + "t": self.t - other.t, }, with_name="LorentzVector", behavior=self.behavior, @@ -966,6 +912,24 @@ def negative(self): ) +_binary_dispatch_cls = { + "TwoVector": TwoVector, + "PolarTwoVector": TwoVector, + "ThreeVector": ThreeVector, + "SphericalThreeVector": ThreeVector, + "LorentzVector": LorentzVector, + "PtEtaPhiMLorentzVector": LorentzVector, + "PtEtaPhiELorentzVector": LorentzVector, +} +_rank = [TwoVector, ThreeVector, LorentzVector] + +for lhs, lhs_to in _binary_dispatch_cls.items(): + for rhs, rhs_to in _binary_dispatch_cls.items(): + out_to = min(lhs_to, rhs_to, key=_rank.index) + behavior[(numpy.add, lhs, rhs)] = out_to.add + behavior[(numpy.subtract, lhs, rhs)] = out_to.subtract + + TwoVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 TwoVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 TwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 From c1926b99309eba758a048a2af46954aa8667be5d Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Tue, 30 Jan 2024 09:28:26 -0600 Subject: [PATCH 09/46] skooch pin for vector --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 6e017b38b..73cde2d52 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,7 +42,7 @@ dependencies = [ "dask[array]>=2023.4.0", "dask-awkward>=2024.3.0", "dask-histogram>=2024.2.0", - "vector>=1.1.1", + "vector>=1.2.0", "correctionlib>=2.3.3", "pyarrow>=6.0.0", "fsspec-xrootd>=0.2.3", From 491507c2dcb454bbe8d059a68ae11440ca60afe7 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Wed, 6 Mar 2024 11:42:03 -0600 Subject: [PATCH 10/46] skooch pin for vector --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 73cde2d52..2a53a9959 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,7 +42,7 @@ dependencies = [ "dask[array]>=2023.4.0", "dask-awkward>=2024.3.0", "dask-histogram>=2024.2.0", - "vector>=1.2.0", + "vector>=1.3.0", "correctionlib>=2.3.3", "pyarrow>=6.0.0", "fsspec-xrootd>=0.2.3", From f27e7bf7b72b8375d0ab7cfc11fce7fb60cb666e Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Thu, 14 Mar 2024 14:07:48 +0100 Subject: [PATCH 11/46] Depend on vector v1.3 + cleanup --- src/coffea/nanoevents/methods/vector.py | 183 ++++-------------------- tests/test_nanoevents_vector.py | 16 +-- 2 files changed, 39 insertions(+), 160 deletions(-) diff --git a/src/coffea/nanoevents/methods/vector.py b/src/coffea/nanoevents/methods/vector.py index f9676b198..406ef5476 100644 --- a/src/coffea/nanoevents/methods/vector.py +++ b/src/coffea/nanoevents/methods/vector.py @@ -152,27 +152,7 @@ def absolute(self): @awkward.mixin_class_method(numpy.negative) def negative(self): """Returns the negative of the vector""" - return awkward.zip( - {"x": -self.x, "y": -self.y}, - with_name="TwoVector", - behavior=self.behavior, - ) - - def add(self, other): - """Add two vectors together elementwise using `x` and `y` components""" - return awkward.zip( - {"x": self.x + other.x, "y": self.y + other.y}, - with_name="TwoVector", - behavior=self.behavior, - ) - - def subtract(self, other): - """Subtract a vector from another elementwise using `x` and `y` components""" - return awkward.zip( - {"x": self.x - other.x, "y": self.y - other.y}, - with_name="TwoVector", - behavior=self.behavior, - ) + return self.scale(-1) def sum(self, axis=-1): """Sum an array of vectors elementwise using `x` and `y` components""" @@ -188,11 +168,7 @@ def sum(self, axis=-1): @awkward.mixin_class_method(numpy.multiply, {numbers.Number}) def multiply(self, other): """Multiply this vector by a scalar elementwise using `x` and `y` components""" - return awkward.zip( - {"x": self.x * other, "y": self.y * other}, - with_name="TwoVector", - behavior=self.behavior, - ) + return self.scale(other) @awkward.mixin_class_method(numpy.divide, {numbers.Number}) def divide(self, other): @@ -228,23 +204,12 @@ def multiply(self, other): In reality, this directly adjusts `r` and `phi` for performance """ - return awkward.zip( - { - "rho": self.r * abs(other), - "phi": self.phi % (2 * numpy.pi) - (numpy.pi * (other < 0)), - }, - with_name="PolarTwoVector", - behavior=self.behavior, - ) + return self.scale(other) @awkward.mixin_class_method(numpy.negative) def negative(self): """Returns the negative of the vector""" - return awkward.zip( - {"rho": self.r, "phi": self.phi % (2 * numpy.pi) - numpy.pi}, - with_name="PolarTwoVector", - behavior=self.behavior, - ) + return self.scale(-1) @awkward.mixin_class(behavior) @@ -279,23 +244,7 @@ def absolute(self): @awkward.mixin_class_method(numpy.negative) def negative(self): """Returns the negative of the vector""" - return awkward.zip( - {"x": -self.x, "y": -self.y, "z": -self.z}, - with_name="ThreeVector", - behavior=self.behavior, - ) - - def add(self, other): - """Add two vectors together elementwise using `x`, `y`, and `z` components""" - return awkward.zip( - { - "x": self.x + other.x, - "y": self.y + other.y, - "z": self.z + other.z, - }, - with_name="ThreeVector", - behavior=self.behavior, - ) + return self.scale(-1) @awkward.mixin_class_method(numpy.divide, {numbers.Number}) def divide(self, other): @@ -303,19 +252,6 @@ def divide(self, other): This is realized by using the multiplication functionality""" return self.multiply(1 / other) - def subtract(self, other): - """Subtract a vector from another elementwise using `x`, `y`, and `z` components""" - - return awkward.zip( - { - "x": self.x - other.x, - "y": self.y - other.y, - "z": self.z - other.z, - }, - with_name="ThreeVector", - behavior=self.behavior, - ) - def sum(self, axis=-1): """Sum an array of vectors elementwise using `x`, `y`, and `z` components""" return awkward.zip( @@ -331,27 +267,7 @@ def sum(self, axis=-1): @awkward.mixin_class_method(numpy.multiply, {numbers.Number}) def multiply(self, other): """Multiply this vector by a scalar elementwise using `x`, `y`, and `z` components""" - return awkward.zip( - {"x": self.x * other, "y": self.y * other, "z": self.z * other}, - with_name="ThreeVector", - behavior=self.behavior, - ) - - def dot(self, other): - """Compute the dot product of two vectors""" - return self.x * other.x + self.y * other.y + self.z * other.z - - def cross(self, other): - """Compute the cross product of two vectors""" - return awkward.zip( - { - "x": self.y * other.z - self.z * other.y, - "y": self.z * other.x - self.x * other.z, - "z": self.x * other.y - self.y * other.x, - }, - with_name="ThreeVector", - behavior=self.behavior, - ) + return self.scale(other) def delta_phi(self, other): """Compute difference in angle between two vectors @@ -388,28 +304,12 @@ def multiply(self, other): In reality, this directly adjusts `r`, `theta` and `phi` for performance """ - return awkward.zip( - { - "rho": self.rho * abs(other), - "theta": (numpy.sign(other) * self.theta + numpy.pi) % numpy.pi, - "phi": self.phi % (2 * numpy.pi) - numpy.pi * (other < 0), - }, - with_name="SphericalThreeVector", - behavior=self.behavior, - ) + return self.scale(other) @awkward.mixin_class_method(numpy.negative) def negative(self): """Returns the negative of the vector""" - return awkward.zip( - { - "rho": self.rho, - "theta": (-self.theta + numpy.pi) % numpy.pi, - "phi": self.phi % (2 * numpy.pi) - numpy.pi, - }, - with_name="SphericalThreeVector", - behavior=self.behavior, - ) + return self.scale(-1) def _metric_table_core(a, b, axis, metric, return_combinations): @@ -481,33 +381,6 @@ def absolute(self): """ return self.mass - def add(self, other): - """Add two vectors together elementwise using `x`, `y`, `z`, and `t` components""" - return awkward.zip( - { - "x": self.x + other.x, - "y": self.y + other.y, - "z": self.z + other.z, - "t": self.t + other.t, - }, - with_name="LorentzVector", - behavior=self.behavior, - ) - - def subtract(self, other): - """Subtract a vector from another elementwise using `x`, `y`, `z`, and `t` components""" - - return awkward.zip( - { - "x": self.x - other.x, - "y": self.y - other.y, - "z": self.z - other.z, - "t": self.t - other.t, - }, - with_name="LorentzVector", - behavior=self.behavior, - ) - def sum(self, axis=-1): """Sum an array of vectors elementwise using `x`, `y`, `z`, and `t` components""" return awkward.zip( @@ -524,22 +397,13 @@ def sum(self, axis=-1): @awkward.mixin_class_method(numpy.multiply, {numbers.Number}) def multiply(self, other): """Multiply this vector by a scalar elementwise using `x`, `y`, `z`, and `t` components""" - return awkward.zip( - { - "x": self.x * other, - "y": self.y * other, - "z": self.z * other, - "t": self.t * other, - }, - with_name="LorentzVector", - behavior=self.behavior, - ) + return self.scale(other) @awkward.mixin_class_method(numpy.divide, {numbers.Number}) def divide(self, other): """Divide this vector by a scalar elementwise using its cartesian components This is realized by using the multiplication functionality""" - return self.multiply(1 / other) + return self.scale(1 / other) def delta_r2(self, other): """Squared `delta_r`""" @@ -562,11 +426,7 @@ def delta_phi(self, other): @awkward.mixin_class_method(numpy.negative) def negative(self): """Returns the negative of the vector""" - return awkward.zip( - {"x": -self.x, "y": -self.y, "z": -self.z, "t": -self.t}, - with_name="LorentzVector", - behavior=self.behavior, - ) + return self.scale(-1) @property def pvec(self): @@ -795,6 +655,12 @@ def negative(self): behavior=self.behavior, ) + @awkward.mixin_class_method(numpy.divide, {numbers.Number}) + def divide(self, other): + """Divide this vector by a scalar elementwise using its cartesian components + This is realized by using the multiplication functionality""" + return self.multiply(1 / other) + @awkward.mixin_class(behavior) class PtEtaPhiELorentzVector(LorentzVector): @@ -911,6 +777,12 @@ def negative(self): behavior=self.behavior, ) + @awkward.mixin_class_method(numpy.divide, {numbers.Number}) + def divide(self, other): + """Divide this vector by a scalar elementwise using its cartesian components + This is realized by using the multiplication functionality""" + return self.multiply(1 / other) + _binary_dispatch_cls = { "TwoVector": TwoVector, @@ -933,30 +805,37 @@ def negative(self): TwoVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 TwoVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 TwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +TwoVectorArray.MomentumClass = PolarTwoVectorArray # noqa: F821 PolarTwoVectorArray.ProjectionClass2D = PolarTwoVectorArray # noqa: F821 PolarTwoVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 PolarTwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +PolarTwoVectorArray.MomentumClass = PolarTwoVectorArray # noqa: F821 ThreeVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 ThreeVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 ThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +ThreeVectorArray.MomentumClass = SphericalThreeVectorArray # noqa: F821 -SphericalThreeVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +SphericalThreeVectorArray.ProjectionClass2D = PolarTwoVectorArray # noqa: F821 SphericalThreeVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 SphericalThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +SphericalThreeVectorArray.MomentumClass = SphericalThreeVectorArray # noqa: F821 LorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 LorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 LorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +LorentzVectorArray.MomentumClass = LorentzVectorArray # noqa: F821 PtEtaPhiMLorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 PtEtaPhiMLorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 PtEtaPhiMLorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +PtEtaPhiMLorentzVectorArray.MomentumClass = LorentzVectorArray # noqa: F821 PtEtaPhiELorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 PtEtaPhiELorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 PtEtaPhiELorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +PtEtaPhiELorentzVectorArray.MomentumClass = LorentzVectorArray # noqa: F821 __all__ = [ "TwoVector", diff --git a/tests/test_nanoevents_vector.py b/tests/test_nanoevents_vector.py index 5d925aa58..12dbcec77 100644 --- a/tests/test_nanoevents_vector.py +++ b/tests/test_nanoevents_vector.py @@ -12,7 +12,7 @@ def assert_record_arrays_equal(a, b, check_type=False): if check_type: assert type(a) is type(b) assert ak.fields(a) == ak.fields(b) - assert all(ak.all(a[f] == b[f]) for f in ak.fields(a)) + assert all(ak.all(ak.isclose(a[f], b[f])) for f in ak.fields(a)) def assert_awkward_allclose(actual, desired): @@ -97,7 +97,7 @@ def test_polar_two_vector(): assert ak.all(abs((-a).y + a.y) < ATOL) assert_record_arrays_equal(a * (-1), -a) - assert ak.all(a.unit.phi == a.phi) + assert ak.all(ak.isclose(a.unit.phi, a.phi)) def test_three_vector(): @@ -539,17 +539,17 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): behavior=vector.behavior, ) - assert_record_arrays_equal(a + b, b + a, check_type=True) - assert_record_arrays_equal(a + c, c + a, check_type=True) - assert_record_arrays_equal(b + c, c + b, check_type=True) + assert_record_arrays_equal(a.like(b) + b, b + a.like(b), check_type=True) + assert_record_arrays_equal(a.like(c) + c, c + a.like(c), check_type=True) + assert_record_arrays_equal(b.like(c) + c, c + b.like(c), check_type=True) assert_allclose(a.delta_phi(b), -b.delta_phi(a)) assert_allclose(a.delta_phi(c), -c.delta_phi(a)) assert_allclose(b.delta_phi(c), -c.delta_phi(b)) - assert_record_arrays_equal((a - b), -(b - a), check_type=True) - assert_record_arrays_equal((a - c), -(c - a), check_type=True) - assert_record_arrays_equal((b - c), -(c - b), check_type=True) + assert_record_arrays_equal((a.like(b) - b), -(b - a.like(b)), check_type=True) + assert_record_arrays_equal((a.like(c) - c), -(c - a.like(c)), check_type=True) + assert_record_arrays_equal((b.like(c) - c), -(c - b.like(c)), check_type=True) @pytest.mark.parametrize("optimization_enabled", [True, False]) From 4ac77ec4791ff5deff23eccd4e9dd2e34a6f61e3 Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Thu, 14 Mar 2024 14:29:43 +0100 Subject: [PATCH 12/46] Candidate classes need MmomentumClass too --- src/coffea/nanoevents/methods/candidate.py | 3 +++ src/coffea/nanoevents/methods/vector.py | 4 ++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/src/coffea/nanoevents/methods/candidate.py b/src/coffea/nanoevents/methods/candidate.py index 2291a744a..19ab3a457 100644 --- a/src/coffea/nanoevents/methods/candidate.py +++ b/src/coffea/nanoevents/methods/candidate.py @@ -70,4 +70,7 @@ class PtEtaPhiECandidate(Candidate, vector.PtEtaPhiELorentzVector): pass +PtEtaPhiMCandidate.MomentumClass = PtEtaPhiMCandidate +PtEtaPhiECandidate.MomentumClass = PtEtaPhiECandidate + __all__ = ["Candidate", "PtEtaPhiMCandidate", "PtEtaPhiECandidate"] diff --git a/src/coffea/nanoevents/methods/vector.py b/src/coffea/nanoevents/methods/vector.py index 406ef5476..2944aa72d 100644 --- a/src/coffea/nanoevents/methods/vector.py +++ b/src/coffea/nanoevents/methods/vector.py @@ -175,7 +175,7 @@ def divide(self, other): """Divide this vector by a scalar elementwise using its cartesian components This is realized by using the multiplication functionality""" - return self.multiply(1 / other) + return self.scale(1 / other) def delta_phi(self, other): """Compute difference in angle between two vectors @@ -250,7 +250,7 @@ def negative(self): def divide(self, other): """Divide this vector by a scalar elementwise using its cartesian components This is realized by using the multiplication functionality""" - return self.multiply(1 / other) + return self.scale(1 / other) def sum(self, axis=-1): """Sum an array of vectors elementwise using `x`, `y`, and `z` components""" From 768de15645a1bd0491895c11566cb4ed18434b63 Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Fri, 15 Mar 2024 14:23:50 +0100 Subject: [PATCH 13/46] Depend on vector v1.3.1 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index dc48ba11c..0930a8440 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,7 +42,7 @@ dependencies = [ "dask[array]>=2023.4.0", "dask-awkward>=2024.3.0", "dask-histogram>=2024.3.0", - "vector>=1.3.0", + "vector>=1.3.1", "correctionlib>=2.3.3", "pyarrow>=6.0.0", "fsspec-xrootd>=0.2.3", From c0887231ef16612dcb2fb198ce82e8edb26c702e Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Fri, 15 Mar 2024 15:17:14 +0100 Subject: [PATCH 14/46] MomentumClass for Jet and Candidate --- src/coffea/nanoevents/methods/candidate.py | 1 + src/coffea/nanoevents/methods/nanoaod.py | 3 +++ 2 files changed, 4 insertions(+) diff --git a/src/coffea/nanoevents/methods/candidate.py b/src/coffea/nanoevents/methods/candidate.py index 19ab3a457..c1ed00d39 100644 --- a/src/coffea/nanoevents/methods/candidate.py +++ b/src/coffea/nanoevents/methods/candidate.py @@ -70,6 +70,7 @@ class PtEtaPhiECandidate(Candidate, vector.PtEtaPhiELorentzVector): pass +Candidate.MomentumClass = Candidate PtEtaPhiMCandidate.MomentumClass = PtEtaPhiMCandidate PtEtaPhiECandidate.MomentumClass = PtEtaPhiECandidate diff --git a/src/coffea/nanoevents/methods/nanoaod.py b/src/coffea/nanoevents/methods/nanoaod.py index 394d1058a..4e75f4497 100644 --- a/src/coffea/nanoevents/methods/nanoaod.py +++ b/src/coffea/nanoevents/methods/nanoaod.py @@ -431,6 +431,9 @@ def constituents(self, dask_array): ) +JetArray.MomentumClass = JetArray # noqa: F821 + + _set_repr_name("Jet") From ef5561b55039da7a9e3c71eacb65cba3d68d6995 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Sun, 14 Jan 2024 14:07:08 -0600 Subject: [PATCH 15/46] first try at backing nanoevents.methods.vector with scikit-hep vector --- pyproject.toml | 1 + src/coffea/nanoevents/methods/vector.py | 359 +++++++++++------------- tests/test_nanoevents_vector.py | 34 ++- 3 files changed, 181 insertions(+), 213 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index aed1ce8b7..98dee3fbf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,7 @@ dependencies = [ "dask[array]>=2023.4.0;python_version<'3.9'", "dask-awkward>=2024.3.0", "dask-histogram>=2024.3.0", + "vector>=1.1.1", "correctionlib>=2.3.3", "pyarrow>=6.0.0", "fsspec-xrootd>=0.2.3", diff --git a/src/coffea/nanoevents/methods/vector.py b/src/coffea/nanoevents/methods/vector.py index 515c5ef7c..d1fe10e35 100644 --- a/src/coffea/nanoevents/methods/vector.py +++ b/src/coffea/nanoevents/methods/vector.py @@ -50,7 +50,13 @@ import numba import numpy import pytz +import vector from dask_awkward import dask_method +from vector.backends.awkward import ( + MomentumAwkward2D, + MomentumAwkward3D, + MomentumAwkward4D, +) from coffea.util import deprecate @@ -109,10 +115,39 @@ def delta_r(eta1, phi1, eta2, phi2): behavior = {} +behavior.update(vector.backends.awkward.behavior) + + +class TwoVectorArray: + pass + + +class PolarTwoVectorArray: + pass + + +class ThreeVectorArray: + pass + + +class SphericalThreeVectorArray: + pass + + +class LorentzVectorArray: + pass + + +class PtEtaPhiMLorentzVectorArray: + pass + + +class PtEtaPhiELorentzVectorArray: + pass @awkward.mixin_class(behavior) -class TwoVector: +class TwoVector(MomentumAwkward2D): """A cartesian 2-dimensional vector A heavy emphasis towards a momentum vector interpretation is assumed, hence @@ -127,40 +162,12 @@ def r(self): :math:`\sqrt{x^2+y^2}` """ - return numpy.sqrt(self.r2) - - @property - def phi(self): - r"""Polar angle relative to X axis - - :math:`\text{arctan2}(y, x)` - """ - return numpy.arctan2(self.y, self.x) - - @property - def px(self): - """Alias for `x`""" - return self.x - - @property - def py(self): - """Alias for `y`""" - return self.y + return self.rho @property def r2(self): """Squared `r`""" - return self.x * self.x + self.y * self.y - - @property - def pt2(self): - """Alias for `r2`""" - return self.r2 - - @property - def pt(self): - """Alias for `r`""" - return self.r + return self.rho2 @awkward.mixin_class_method(numpy.absolute) def absolute(self): @@ -240,11 +247,7 @@ def delta_phi(self, other): Returns a value within [-pi, pi) """ - return delta_phi(self.phi, other.phi) - - def dot(self, other): - """Compute the dot product of two vectors""" - return self.x * other.x + self.y * other.y + return self.deltaphi(other) @property def unit(self): @@ -256,47 +259,10 @@ def unit(self): class PolarTwoVector(TwoVector): """A polar coordinate 2-dimensional vector - This mixin class requires the parent class to provide items `r` and `phi`. + This mixin class requires the parent class to provide items `rho` and `phi`. Some additional properties are overridden for performance """ - @property - def x(self): - r"""Cartesian x value - - :math:`r \cos{\phi}` - """ - return self.r * numpy.cos(self.phi) - - @property - def y(self): - r"""Cartesian y value - - :math:`r \sin{\phi}` - """ - return self.r * numpy.sin(self.phi) - - @property - def r(self): - r"""Distance from origin in XY plane - - :math:`\sqrt{x^2+y^2}` - """ - return self["r"] - - @property - def phi(self): - r"""Azimuthal angle relative to X axis in XY plane - - :math:`\text{arctan2}(y, x)` - """ - return self["phi"] - - @property - def r2(self): - """Squared `r`""" - return self.r * self.r - @awkward.mixin_class_method(numpy.multiply, {numbers.Number}) def multiply(self, other): """Multiply this vector by a scalar elementwise using using `x` and `y` components @@ -305,7 +271,7 @@ def multiply(self, other): """ return awkward.zip( { - "r": self.r * abs(other), + "rho": self.r * abs(other), "phi": self.phi % (2 * numpy.pi) - (numpy.pi * (other < 0)), }, with_name="PolarTwoVector", @@ -316,14 +282,14 @@ def multiply(self, other): def negative(self): """Returns the negative of the vector""" return awkward.zip( - {"r": self.r, "phi": self.phi % (2 * numpy.pi) - numpy.pi}, + {"rho": self.r, "phi": self.phi % (2 * numpy.pi) - numpy.pi}, with_name="PolarTwoVector", behavior=self.behavior, ) @awkward.mixin_class(behavior) -class ThreeVector(TwoVector): +class ThreeVector(MomentumAwkward3D): """A cartesian 3-dimensional vector A heavy emphasis towards a momentum vector interpretation is assumed. @@ -331,41 +297,18 @@ class ThreeVector(TwoVector): """ @property - def pz(self): - """Alias for `z`""" - return self.z - - @property - def rho2(self): - """Squared `rho`""" - return self.r2 + self.z * self.z - - @property - def rho(self): - r"""Distance from origin in 3D - - :math:`\sqrt{x^2+y^2+z^2} = \sqrt{r^2+z^2}` - """ - return numpy.sqrt(self.rho2) - - @property - def theta(self): - r"""Inclination angle from XY plane + def r(self): + r"""Distance from origin in XY plane - :math:`\text{arctan2}(r, z)` + :math:`\sqrt{x^2+y^2}` """ - return numpy.arctan2(self.r, self.z) + return self.rho @property - def p2(self): - """Squared `p`""" + def r2(self): + """Squared `r`""" return self.rho2 - @property - def p(self): - """Alias for `rho`""" - return self.rho - @awkward.mixin_class_method(numpy.absolute) def absolute(self): """Returns magnitude of the 3D vector @@ -383,18 +326,33 @@ def negative(self): behavior=self.behavior, ) - @awkward.mixin_class_method(numpy.add, {"ThreeVector"}) + @awkward.mixin_class_method( + numpy.add, {"ThreeVector", "TwoVector", "PolarTwoVector"} + ) def add(self, other): """Add two vectors together elementwise using `x`, `y`, and `z` components""" + other_3d = other.to_Vector3D() return awkward.zip( - {"x": self.x + other.x, "y": self.y + other.y, "z": self.z + other.z}, + { + "x": self.x + other_3d.x, + "y": self.y + other_3d.y, + "z": self.z + other_3d.z, + }, with_name="ThreeVector", behavior=self.behavior, ) + @awkward.mixin_class_method(numpy.divide, {numbers.Number}) + def divide(self, other): + """Divide this vector by a scalar elementwise using its cartesian components + This is realized by using the multiplication functionality""" + return self.multiply(1 / other) + @awkward.mixin_class_method( numpy.subtract, { + "TwoVector", + "PolarTwoVector", "ThreeVector", "SphericalThreeVector", "LorentzVector", @@ -405,8 +363,14 @@ def add(self, other): ) def subtract(self, other): """Subtract a vector from another elementwise using `x`, `y`, and `z` components""" + other_3d = other.to_Vector3D() + return awkward.zip( - {"x": self.x - other.x, "y": self.y - other.y, "z": self.z - other.z}, + { + "x": self.x - other_3d.x, + "y": self.y - other_3d.y, + "z": self.z - other_3d.z, + }, with_name="ThreeVector", behavior=self.behavior, ) @@ -448,6 +412,13 @@ def cross(self, other): behavior=self.behavior, ) + def delta_phi(self, other): + """Compute difference in angle between two vectors + + Returns a value within [-pi, pi) + """ + return self.deltaphi(other) + @property def unit(self): """Unit vector, a vector of length 1 pointing in the same direction""" @@ -455,7 +426,7 @@ def unit(self): @awkward.mixin_class(behavior) -class SphericalThreeVector(ThreeVector, PolarTwoVector): +class SphericalThreeVector(ThreeVector): """A spherical coordinate 3-dimensional vector This mixin class requires the parent class to provide items `rho`, `theta`, and `phi`. @@ -468,42 +439,8 @@ def r(self): :math:`\sqrt{x^2+y^2} = \rho \sin(\theta)` """ - return self.rho * numpy.sin(self.theta) - - @property - def z(self): - r"""Cartesian z value - - :math:`\rho \cos(\theta)` - """ - return self.rho * numpy.cos(self.theta) - - @property - def rho(self): - r"""Distance from origin in 3D - - :math:`\sqrt{x^2+y^2+z^2} = \sqrt{r^2+z^2}` - """ - return self["rho"] - - @property - def theta(self): - r"""Inclination angle from XY plane - - :math:`\text{arctan2}(r, z)` - """ - return self["theta"] - - @property - def p(self): - """Alias for `rho`""" return self.rho - @property - def p2(self): - """Squared `p`""" - return self.rho * self.rho - @awkward.mixin_class_method(numpy.multiply, {numbers.Number}) def multiply(self, other): """Multiply this vector by a scalar elementwise using `x`, `y`, and `z` components @@ -561,7 +498,7 @@ def _nearest_core(x, y, axis, metric, return_metric, threshold): @awkward.mixin_class(behavior) -class LorentzVector(ThreeVector): +class LorentzVector(MomentumAwkward4D): """A cartesian Lorentz vector A heavy emphasis towards a momentum vector interpretation is assumed. @@ -603,29 +540,51 @@ def absolute(self): """ return self.mass - @awkward.mixin_class_method(numpy.add, {"LorentzVector"}) + @awkward.mixin_class_method( + numpy.add, + { + "LorentzVector", + "ThreeVector", + "SphericalThreeVector", + "TwoVector", + "PolarTwoVector", + }, + ) def add(self, other): """Add two vectors together elementwise using `x`, `y`, `z`, and `t` components""" + other_4d = other.to_Vector4D() return awkward.zip( { - "x": self.x + other.x, - "y": self.y + other.y, - "z": self.z + other.z, - "t": self.t + other.t, + "x": self.x + other_4d.x, + "y": self.y + other_4d.y, + "z": self.z + other_4d.z, + "t": self.t + other_4d.t, }, with_name="LorentzVector", behavior=self.behavior, ) - @awkward.mixin_class_method(numpy.subtract, {"LorentzVector"}, transpose=False) + @awkward.mixin_class_method( + numpy.subtract, + { + "LorentzVector", + "ThreeVector", + "SphericalThreeVector", + "TwoVector", + "PolarTwoVector", + }, + transpose=False, + ) def subtract(self, other): """Subtract a vector from another elementwise using `x`, `y`, `z`, and `t` components""" + other_4d = other.to_Vector4D() + return awkward.zip( { - "x": self.x - other.x, - "y": self.y - other.y, - "z": self.z - other.z, - "t": self.t - other.t, + "x": self.x - other_4d.x, + "y": self.y - other_4d.y, + "z": self.z - other_4d.z, + "t": self.t - other_4d.t, }, with_name="LorentzVector", behavior=self.behavior, @@ -658,16 +617,29 @@ def multiply(self, other): behavior=self.behavior, ) + @awkward.mixin_class_method(numpy.divide, {numbers.Number}) + def divide(self, other): + """Divide this vector by a scalar elementwise using its cartesian components + This is realized by using the multiplication functionality""" + return self.multiply(1 / other) + def delta_r2(self, other): """Squared `delta_r`""" - return delta_r(self.eta, self.phi, other.eta, other.phi) ** 2 + return self.deltaR2(other) def delta_r(self, other): r"""Distance between two Lorentz vectors in (eta,phi) plane :math:`\sqrt{\Delta\eta^2 + \Delta\phi^2}` """ - return delta_r(self.eta, self.phi, other.eta, other.phi) + return self.deltaR(other) + + def delta_phi(self, other): + """Compute difference in angle between two vectors + + Returns a value within [-pi, pi) + """ + return self.deltaphi(other) @awkward.mixin_class_method(numpy.negative) def negative(self): @@ -687,12 +659,6 @@ def pvec(self): behavior=self.behavior, ) - @property - def rapidity(self): - pz = self.z - e = self.energy - return 0.5 * (numpy.log(e + pz) - numpy.log(e - pz)) - @property def boostvec(self): """The `x`, `y` and `z` components divided by `t` as a `ThreeVector` @@ -700,40 +666,7 @@ def boostvec(self): This can be used for boosting. For cases where `|t| <= rho`, this returns the unit vector. """ - rho = self.rho - t = self.t - with numpy.errstate(divide="ignore"): - out = self.pvec * awkward.where( - rho == 0, 0, awkward.where(abs(t) <= rho, 1 / rho, 1 / t) - ) - return out - - def boost(self, other): - """Apply a Lorentz boost given by the `ThreeVector` `other` and return it - - Note that this follows the convention that, for example in order to boost - a vector into its own rest frame, one needs to use the negative of its `boostvec` - """ - b2 = other.rho2 - gamma = (1 - b2) ** (-0.5) - mask = b2 == 0 - b2 = awkward.where(mask, 1, b2) - gamma2 = awkward.where(mask, 0, (gamma - 1) / b2) - - bp = self.dot(other) - t = self.t - v = gamma2 * bp * other + t * gamma * other - - return awkward.zip( - { - "x": self.x + v.x, - "y": self.y + v.y, - "z": self.z + v.z, - "t": gamma * (t + bp), - }, - with_name="LorentzVector", - behavior=self.behavior, - ) + return self.to_beta3() @dask_method def metric_table( @@ -821,7 +754,7 @@ def nearest( @awkward.mixin_class(behavior) -class PtEtaPhiMLorentzVector(LorentzVector, SphericalThreeVector): +class PtEtaPhiMLorentzVector(LorentzVector): """A Lorentz vector using pseudorapidity and mass This mixin class requires the parent class to provide items `pt`, `eta`, `phi`, and `mass`. @@ -946,7 +879,7 @@ def negative(self): @awkward.mixin_class(behavior) -class PtEtaPhiELorentzVector(LorentzVector, SphericalThreeVector): +class PtEtaPhiELorentzVector(LorentzVector): """A Lorentz vector using pseudorapidity and energy This mixin class requires the parent class to provide items `pt`, `eta`, `phi`, and `energy`. @@ -1061,6 +994,34 @@ def negative(self): ) +TwoVectorArray.ProjectionClass2D = TwoVectorArray +TwoVectorArray.ProjectionClass3D = ThreeVectorArray +TwoVectorArray.ProjectionClass4D = LorentzVectorArray + +PolarTwoVectorArray.ProjectionClass2D = PolarTwoVectorArray +PolarTwoVectorArray.ProjectionClass3D = SphericalThreeVectorArray +PolarTwoVectorArray.ProjectionClass4D = LorentzVectorArray + +ThreeVectorArray.ProjectionClass2D = TwoVectorArray +ThreeVectorArray.ProjectionClass3D = ThreeVectorArray +ThreeVectorArray.ProjectionClass4D = LorentzVectorArray + +SphericalThreeVectorArray.ProjectionClass2D = TwoVectorArray +SphericalThreeVectorArray.ProjectionClass3D = SphericalThreeVectorArray +SphericalThreeVectorArray.ProjectionClass4D = LorentzVectorArray + +LorentzVectorArray.ProjectionClass2D = TwoVectorArray +LorentzVectorArray.ProjectionClass3D = ThreeVectorArray +LorentzVectorArray.ProjectionClass4D = LorentzVectorArray + +PtEtaPhiMLorentzVectorArray.ProjectionClass2D = TwoVectorArray +PtEtaPhiMLorentzVectorArray.ProjectionClass3D = ThreeVectorArray +PtEtaPhiMLorentzVectorArray.ProjectionClass4D = LorentzVectorArray + +PtEtaPhiELorentzVectorArray.ProjectionClass2D = TwoVectorArray +PtEtaPhiELorentzVectorArray.ProjectionClass3D = ThreeVectorArray +PtEtaPhiELorentzVectorArray.ProjectionClass4D = LorentzVectorArray + __all__ = [ "TwoVector", "PolarTwoVector", diff --git a/tests/test_nanoevents_vector.py b/tests/test_nanoevents_vector.py index 7fcb7177c..754307ec3 100644 --- a/tests/test_nanoevents_vector.py +++ b/tests/test_nanoevents_vector.py @@ -58,7 +58,7 @@ def test_two_vector(): def test_polar_two_vector(): a = ak.zip( { - "r": [[1, 2], [], [3], [4]], + "rho": [[1, 2], [], [3], [4]], "phi": [[0.3, 0.4], [], [0.5], [0.6]], }, with_name="PolarTwoVector", @@ -67,9 +67,9 @@ def test_polar_two_vector(): assert record_arrays_equal( a * 2, - ak.zip({"r": [[2, 4], [], [6], [8]], "phi": [[0.3, 0.4], [], [0.5], [0.6]]}), + ak.zip({"rho": [[2, 4], [], [6], [8]], "phi": [[0.3, 0.4], [], [0.5], [0.6]]}), ) - assert ak.all((a * (-2)).r == [[2, 4], [], [6], [8]]) + assert ak.all((a * (-2)).rho == [[2, 4], [], [6], [8]]) assert ak.all( (a * (-2)).phi - ak.Array( @@ -80,7 +80,7 @@ def test_polar_two_vector(): assert record_arrays_equal( a / 2, ak.zip( - {"r": [[0.5, 1], [], [1.5], [2]], "phi": [[0.3, 0.4], [], [0.5], [0.6]]} + {"rho": [[0.5, 1], [], [1.5], [2]], "phi": [[0.3, 0.4], [], [0.5], [0.6]]} ), ) @@ -444,12 +444,18 @@ def test_lorentz_vector_numba(a_dtype, b_dtype): 110.66616465749593, 110.68423555321688, ] - assert pytest.approx(a.delta_phi(b), abs=1e-7) == [ - 0.03369510734601633, - -0.1798534997924781, - 0.33292327383538156, - 0.6078019961139605, - ] + + computed_dphi = a.delta_phi(b).to_numpy() + + assert pytest.approx(computed_dphi, abs=1e-6) == np.array( + [ + 0.03369510734601633, + -0.1798534997924781, + 0.33292327383538156, + 0.6078019961139605, + ], + dtype=computed_dphi.dtype, + ) @pytest.mark.parametrize( @@ -519,7 +525,7 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): ) elif twocoord == "PolarTwoVector": c = ak.zip( - {"r": [-10.0, 13.0, 15.0], "phi": [1.22, -1.0, 1.0]}, + {"rho": [-10.0, 13.0, 15.0], "phi": [1.22, -1.0, 1.0]}, with_name=twocoord, behavior=vector.behavior, ) @@ -532,9 +538,9 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): assert record_arrays_equal(a.delta_phi(c), c.delta_phi(a)) assert record_arrays_equal(b.delta_phi(c), c.delta_phi(b)) - assert record_arrays_equal(a - b, -(b - a)) - assert record_arrays_equal(a - c, -(c - a)) - assert record_arrays_equal(b - c, -(c - b)) + assert record_arrays_equal((a - b).to_Vector3D(), -(b - a)) + assert record_arrays_equal((a - c).to_Vector2D(), -(c - a)) + assert record_arrays_equal((b - c).to_Vector2D(), -(c - b)) @pytest.mark.parametrize("optimization_enabled", [True, False]) From 11da182d9a725471902ac72c1b6c240322883a55 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Thu, 18 Jan 2024 13:52:06 -0600 Subject: [PATCH 16/46] fix math in doc --- src/coffea/nanoevents/methods/vector.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/coffea/nanoevents/methods/vector.py b/src/coffea/nanoevents/methods/vector.py index d1fe10e35..7b5293b79 100644 --- a/src/coffea/nanoevents/methods/vector.py +++ b/src/coffea/nanoevents/methods/vector.py @@ -663,7 +663,7 @@ def pvec(self): def boostvec(self): """The `x`, `y` and `z` components divided by `t` as a `ThreeVector` - This can be used for boosting. For cases where `|t| <= rho`, this + This can be used for boosting. For cases where `|t| <= r`, this returns the unit vector. """ return self.to_beta3() From 914ed4e486ab2991c96328652c8291fc821acd28 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Sun, 21 Jan 2024 12:20:06 -0600 Subject: [PATCH 17/46] use noqa instead of actually satisfying flake - more readable --- src/coffea/nanoevents/methods/vector.py | 70 ++++++++----------------- 1 file changed, 21 insertions(+), 49 deletions(-) diff --git a/src/coffea/nanoevents/methods/vector.py b/src/coffea/nanoevents/methods/vector.py index 7b5293b79..3ca68b3fc 100644 --- a/src/coffea/nanoevents/methods/vector.py +++ b/src/coffea/nanoevents/methods/vector.py @@ -118,34 +118,6 @@ def delta_r(eta1, phi1, eta2, phi2): behavior.update(vector.backends.awkward.behavior) -class TwoVectorArray: - pass - - -class PolarTwoVectorArray: - pass - - -class ThreeVectorArray: - pass - - -class SphericalThreeVectorArray: - pass - - -class LorentzVectorArray: - pass - - -class PtEtaPhiMLorentzVectorArray: - pass - - -class PtEtaPhiELorentzVectorArray: - pass - - @awkward.mixin_class(behavior) class TwoVector(MomentumAwkward2D): """A cartesian 2-dimensional vector @@ -994,33 +966,33 @@ def negative(self): ) -TwoVectorArray.ProjectionClass2D = TwoVectorArray -TwoVectorArray.ProjectionClass3D = ThreeVectorArray -TwoVectorArray.ProjectionClass4D = LorentzVectorArray +TwoVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +TwoVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +TwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -PolarTwoVectorArray.ProjectionClass2D = PolarTwoVectorArray -PolarTwoVectorArray.ProjectionClass3D = SphericalThreeVectorArray -PolarTwoVectorArray.ProjectionClass4D = LorentzVectorArray +PolarTwoVectorArray.ProjectionClass2D = PolarTwoVectorArray # noqa: F821 +PolarTwoVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 +PolarTwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -ThreeVectorArray.ProjectionClass2D = TwoVectorArray -ThreeVectorArray.ProjectionClass3D = ThreeVectorArray -ThreeVectorArray.ProjectionClass4D = LorentzVectorArray +ThreeVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +ThreeVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +ThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -SphericalThreeVectorArray.ProjectionClass2D = TwoVectorArray -SphericalThreeVectorArray.ProjectionClass3D = SphericalThreeVectorArray -SphericalThreeVectorArray.ProjectionClass4D = LorentzVectorArray +SphericalThreeVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +SphericalThreeVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 +SphericalThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -LorentzVectorArray.ProjectionClass2D = TwoVectorArray -LorentzVectorArray.ProjectionClass3D = ThreeVectorArray -LorentzVectorArray.ProjectionClass4D = LorentzVectorArray +LorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +LorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +LorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -PtEtaPhiMLorentzVectorArray.ProjectionClass2D = TwoVectorArray -PtEtaPhiMLorentzVectorArray.ProjectionClass3D = ThreeVectorArray -PtEtaPhiMLorentzVectorArray.ProjectionClass4D = LorentzVectorArray +PtEtaPhiMLorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +PtEtaPhiMLorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +PtEtaPhiMLorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -PtEtaPhiELorentzVectorArray.ProjectionClass2D = TwoVectorArray -PtEtaPhiELorentzVectorArray.ProjectionClass3D = ThreeVectorArray -PtEtaPhiELorentzVectorArray.ProjectionClass4D = LorentzVectorArray +PtEtaPhiELorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +PtEtaPhiELorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +PtEtaPhiELorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 __all__ = [ "TwoVector", From 2d54aedbec49602c9317f3197008f128dad791dd Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 21 Jan 2024 18:20:21 +0000 Subject: [PATCH 18/46] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/coffea/nanoevents/methods/vector.py | 40 ++++++++++++------------- 1 file changed, 20 insertions(+), 20 deletions(-) diff --git a/src/coffea/nanoevents/methods/vector.py b/src/coffea/nanoevents/methods/vector.py index 3ca68b3fc..88d505dfc 100644 --- a/src/coffea/nanoevents/methods/vector.py +++ b/src/coffea/nanoevents/methods/vector.py @@ -966,32 +966,32 @@ def negative(self): ) -TwoVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 -TwoVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 -TwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +TwoVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +TwoVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +TwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -PolarTwoVectorArray.ProjectionClass2D = PolarTwoVectorArray # noqa: F821 -PolarTwoVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 -PolarTwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +PolarTwoVectorArray.ProjectionClass2D = PolarTwoVectorArray # noqa: F821 +PolarTwoVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 +PolarTwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -ThreeVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 -ThreeVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 -ThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +ThreeVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +ThreeVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +ThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -SphericalThreeVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 -SphericalThreeVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 -SphericalThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +SphericalThreeVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +SphericalThreeVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 +SphericalThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -LorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 -LorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 -LorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +LorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +LorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +LorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -PtEtaPhiMLorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 -PtEtaPhiMLorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 -PtEtaPhiMLorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +PtEtaPhiMLorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +PtEtaPhiMLorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +PtEtaPhiMLorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -PtEtaPhiELorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 -PtEtaPhiELorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +PtEtaPhiELorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +PtEtaPhiELorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 PtEtaPhiELorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 __all__ = [ From d7cd6102564b9e8716b7a4bfeb939a5b120dca37 Mon Sep 17 00:00:00 2001 From: Nick Smith Date: Fri, 26 Jan 2024 15:35:11 -0600 Subject: [PATCH 19/46] Improve tests of vector --- tests/test_nanoevents_vector.py | 93 ++++++++++++++++++--------------- 1 file changed, 51 insertions(+), 42 deletions(-) diff --git a/tests/test_nanoevents_vector.py b/tests/test_nanoevents_vector.py index 754307ec3..f7372f883 100644 --- a/tests/test_nanoevents_vector.py +++ b/tests/test_nanoevents_vector.py @@ -2,15 +2,24 @@ import numpy as np import pytest +from numpy.testing import assert_allclose from coffea.nanoevents.methods import vector ATOL = 1e-8 -def record_arrays_equal(a, b): - return (ak.fields(a) == ak.fields(b)) and all( - ak.all(a[f] == b[f]) for f in ak.fields(a) - ) +def assert_record_arrays_equal(a, b, check_type=False): + if check_type: + assert type(a) == type(b) + assert ak.fields(a) == ak.fields(b) + assert all(ak.all(a[f] == b[f]) for f in ak.fields(a)) + + +def assert_awkward_allclose(actual, desired): + flat_actual = ak.flatten(actual, axis=None) + flat_desired = ak.flatten(desired, axis=None) + # we should check None values, but not used in these tests + assert_allclose(flat_actual, flat_desired) def test_two_vector(): @@ -25,31 +34,31 @@ def test_two_vector(): behavior=vector.behavior, ) - assert record_arrays_equal( + assert_record_arrays_equal( -a, ak.zip({"x": [[-1, -2], [], [-3], [-4]], "y": [[-5, -6], [], [-7], [-8]]}) ) - assert record_arrays_equal( + assert_record_arrays_equal( a + b, ak.zip({"x": [[12, 14], [], [16], [18]], "y": [[20, 22], [], [24], [26]]}), ) - assert record_arrays_equal( + assert_record_arrays_equal( a - b, ak.zip( {"x": [[-10, -10], [], [-10], [-10]], "y": [[-10, -10], [], [-10], [-10]]} ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a * 2, ak.zip({"x": [[2, 4], [], [6], [8]], "y": [[10, 12], [], [14], [16]]}) ) - assert record_arrays_equal( + assert_record_arrays_equal( a / 2, ak.zip({"x": [[0.5, 1], [], [1.5], [2]], "y": [[2.5, 3], [], [3.5], [4]]}), ) - assert record_arrays_equal(a.dot(b), ak.Array([[86, 120], [], [158], [200]])) - assert record_arrays_equal(b.dot(a), ak.Array([[86, 120], [], [158], [200]])) + assert_awkward_allclose(a.dot(b), ak.Array([[86, 120], [], [158], [200]])) + assert_awkward_allclose(b.dot(a), ak.Array([[86, 120], [], [158], [200]])) assert ak.all(abs(a.unit.r - 1) < ATOL) assert ak.all(abs(a.unit.phi - a.phi) < ATOL) @@ -65,7 +74,7 @@ def test_polar_two_vector(): behavior=vector.behavior, ) - assert record_arrays_equal( + assert_record_arrays_equal( a * 2, ak.zip({"rho": [[2, 4], [], [6], [8]], "phi": [[0.3, 0.4], [], [0.5], [0.6]]}), ) @@ -77,7 +86,7 @@ def test_polar_two_vector(): ) < ATOL ) - assert record_arrays_equal( + assert_record_arrays_equal( a / 2, ak.zip( {"rho": [[0.5, 1], [], [1.5], [2]], "phi": [[0.3, 0.4], [], [0.5], [0.6]]} @@ -86,7 +95,7 @@ def test_polar_two_vector(): assert ak.all(abs((-a).x + a.x) < ATOL) assert ak.all(abs((-a).y + a.y) < ATOL) - assert record_arrays_equal(a * (-1), -a) + assert_record_arrays_equal(a * (-1), -a) assert ak.all(a.unit.phi == a.phi) @@ -111,7 +120,7 @@ def test_three_vector(): behavior=vector.behavior, ) - assert record_arrays_equal( + assert_record_arrays_equal( -a, ak.zip( { @@ -122,7 +131,7 @@ def test_three_vector(): ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a + b, ak.zip( { @@ -132,7 +141,7 @@ def test_three_vector(): } ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a - b, ak.zip( { @@ -142,7 +151,7 @@ def test_three_vector(): } ), ) - assert record_arrays_equal( + assert_record_arrays_equal( b - a, ak.zip( { @@ -153,7 +162,7 @@ def test_three_vector(): ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a * 2, ak.zip( { @@ -163,7 +172,7 @@ def test_three_vector(): } ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a / 2, ak.zip( { @@ -177,7 +186,7 @@ def test_three_vector(): assert ak.all(a.dot(b) == ak.Array([[170, 154], [], [162], [284]])) assert ak.all(b.dot(a) == ak.Array([[170, 154], [], [162], [284]])) - assert record_arrays_equal( + assert_record_arrays_equal( a.cross(b), ak.zip( { @@ -187,7 +196,7 @@ def test_three_vector(): } ), ) - assert record_arrays_equal( + assert_record_arrays_equal( b.cross(a), ak.zip( { @@ -216,7 +225,7 @@ def test_spherical_three_vector(): assert ak.all(abs((-a).x + a.x) < ATOL) assert ak.all(abs((-a).y + a.y) < ATOL) assert ak.all(abs((-a).z + a.z) < ATOL) - assert record_arrays_equal(a * (-1), -a) + assert_record_arrays_equal(a * (-1), -a, check_type=True) def test_lorentz_vector(): @@ -241,7 +250,7 @@ def test_lorentz_vector(): behavior=vector.behavior, ) - assert record_arrays_equal( + assert_record_arrays_equal( -a, ak.zip( { @@ -253,7 +262,7 @@ def test_lorentz_vector(): ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a + b, ak.zip( { @@ -264,7 +273,7 @@ def test_lorentz_vector(): } ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a - b, ak.zip( { @@ -276,7 +285,7 @@ def test_lorentz_vector(): ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a * 2, ak.zip( { @@ -287,7 +296,7 @@ def test_lorentz_vector(): } ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a / 2, ak.zip( { @@ -299,7 +308,7 @@ def test_lorentz_vector(): ), ) - assert record_arrays_equal( + assert_record_arrays_equal( a.pvec, ak.zip( { @@ -344,7 +353,7 @@ def test_pt_eta_phi_m_lorentz_vector(): ) < ATOL ) - assert record_arrays_equal( + assert_record_arrays_equal( a / 2, ak.zip( { @@ -355,7 +364,7 @@ def test_pt_eta_phi_m_lorentz_vector(): } ), ) - assert record_arrays_equal(a * (-1), -a) + assert_record_arrays_equal(a * (-1), -a, check_type=True) boosted = a.boost(-a.boostvec) assert ak.all(abs(boosted.x) < ATOL) @@ -390,7 +399,7 @@ def test_pt_eta_phi_e_lorentz_vector(): ) < ATOL ) - assert record_arrays_equal( + assert_record_arrays_equal( a / 2, ak.zip( { @@ -401,7 +410,7 @@ def test_pt_eta_phi_e_lorentz_vector(): } ), ) - assert record_arrays_equal(a * (-1), -a) + assert_record_arrays_equal(a * (-1), -a, check_type=True) boosted = a.boost(-a.boostvec) assert ak.all(abs(boosted.x) < ATOL) @@ -530,17 +539,17 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): behavior=vector.behavior, ) - assert record_arrays_equal(a + b, b + a) - assert record_arrays_equal(a + c, c + a) - assert record_arrays_equal(b + c, c + b) + assert_record_arrays_equal(a + b, b + a, check_type=True) + assert_record_arrays_equal(a + c, c + a, check_type=True) + assert_record_arrays_equal(b + c, c + b, check_type=True) - assert record_arrays_equal(a.delta_phi(b), b.delta_phi(a)) - assert record_arrays_equal(a.delta_phi(c), c.delta_phi(a)) - assert record_arrays_equal(b.delta_phi(c), c.delta_phi(b)) + assert_allclose(a.delta_phi(b), -b.delta_phi(a)) + assert_allclose(a.delta_phi(c), -c.delta_phi(a)) + assert_allclose(b.delta_phi(c), -c.delta_phi(b)) - assert record_arrays_equal((a - b).to_Vector3D(), -(b - a)) - assert record_arrays_equal((a - c).to_Vector2D(), -(c - a)) - assert record_arrays_equal((b - c).to_Vector2D(), -(c - b)) + assert_record_arrays_equal((a - b), -(b - a), check_type=True) + assert_record_arrays_equal((a - c), -(c - a), check_type=True) + assert_record_arrays_equal((b - c), -(c - b), check_type=True) @pytest.mark.parametrize("optimization_enabled", [True, False]) From d698ef9d21728f3c076a214fbf6167d11611fb7a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 26 Jan 2024 21:40:14 +0000 Subject: [PATCH 20/46] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/test_nanoevents_vector.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_nanoevents_vector.py b/tests/test_nanoevents_vector.py index f7372f883..1f05e97a9 100644 --- a/tests/test_nanoevents_vector.py +++ b/tests/test_nanoevents_vector.py @@ -1,8 +1,8 @@ import awkward as ak import numpy as np import pytest - from numpy.testing import assert_allclose + from coffea.nanoevents.methods import vector ATOL = 1e-8 From afa0d4922ebf01b95aa6d32239f29931ae2bc6ff Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Fri, 26 Jan 2024 15:42:38 -0600 Subject: [PATCH 21/46] appease linter --- tests/test_nanoevents_vector.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_nanoevents_vector.py b/tests/test_nanoevents_vector.py index 1f05e97a9..8fdb8857b 100644 --- a/tests/test_nanoevents_vector.py +++ b/tests/test_nanoevents_vector.py @@ -10,7 +10,7 @@ def assert_record_arrays_equal(a, b, check_type=False): if check_type: - assert type(a) == type(b) + assert type(a) is type(b) assert ak.fields(a) == ak.fields(b) assert all(ak.all(a[f] == b[f]) for f in ak.fields(a)) From a85fcdc06ed98d096528643f9c7c62f4fe2db9bc Mon Sep 17 00:00:00 2001 From: Nick Smith Date: Mon, 29 Jan 2024 17:08:40 -0600 Subject: [PATCH 22/46] Manually build binary ufunc dispatch tables Eventually we might also want to pawn these off on vector --- src/coffea/nanoevents/methods/vector.py | 100 ++++++++---------------- 1 file changed, 32 insertions(+), 68 deletions(-) diff --git a/src/coffea/nanoevents/methods/vector.py b/src/coffea/nanoevents/methods/vector.py index 88d505dfc..f9676b198 100644 --- a/src/coffea/nanoevents/methods/vector.py +++ b/src/coffea/nanoevents/methods/vector.py @@ -158,7 +158,6 @@ def negative(self): behavior=self.behavior, ) - @awkward.mixin_class_method(numpy.add, {"TwoVector"}) def add(self, other): """Add two vectors together elementwise using `x` and `y` components""" return awkward.zip( @@ -167,18 +166,6 @@ def add(self, other): behavior=self.behavior, ) - @awkward.mixin_class_method( - numpy.subtract, - { - "TwoVector", - "ThreeVector", - "SphericalThreeVector", - "LorentzVector", - "PtEtaPhiMLorentzVector", - "PtEtaPhiELorentzVector", - }, - transpose=False, - ) def subtract(self, other): """Subtract a vector from another elementwise using `x` and `y` components""" return awkward.zip( @@ -298,17 +285,13 @@ def negative(self): behavior=self.behavior, ) - @awkward.mixin_class_method( - numpy.add, {"ThreeVector", "TwoVector", "PolarTwoVector"} - ) def add(self, other): """Add two vectors together elementwise using `x`, `y`, and `z` components""" - other_3d = other.to_Vector3D() return awkward.zip( { - "x": self.x + other_3d.x, - "y": self.y + other_3d.y, - "z": self.z + other_3d.z, + "x": self.x + other.x, + "y": self.y + other.y, + "z": self.z + other.z, }, with_name="ThreeVector", behavior=self.behavior, @@ -320,28 +303,14 @@ def divide(self, other): This is realized by using the multiplication functionality""" return self.multiply(1 / other) - @awkward.mixin_class_method( - numpy.subtract, - { - "TwoVector", - "PolarTwoVector", - "ThreeVector", - "SphericalThreeVector", - "LorentzVector", - "PtEtaPhiMLorentzVector", - "PtEtaPhiELorentzVector", - }, - transpose=False, - ) def subtract(self, other): """Subtract a vector from another elementwise using `x`, `y`, and `z` components""" - other_3d = other.to_Vector3D() return awkward.zip( { - "x": self.x - other_3d.x, - "y": self.y - other_3d.y, - "z": self.z - other_3d.z, + "x": self.x - other.x, + "y": self.y - other.y, + "z": self.z - other.z, }, with_name="ThreeVector", behavior=self.behavior, @@ -512,51 +481,28 @@ def absolute(self): """ return self.mass - @awkward.mixin_class_method( - numpy.add, - { - "LorentzVector", - "ThreeVector", - "SphericalThreeVector", - "TwoVector", - "PolarTwoVector", - }, - ) def add(self, other): """Add two vectors together elementwise using `x`, `y`, `z`, and `t` components""" - other_4d = other.to_Vector4D() return awkward.zip( { - "x": self.x + other_4d.x, - "y": self.y + other_4d.y, - "z": self.z + other_4d.z, - "t": self.t + other_4d.t, + "x": self.x + other.x, + "y": self.y + other.y, + "z": self.z + other.z, + "t": self.t + other.t, }, with_name="LorentzVector", behavior=self.behavior, ) - @awkward.mixin_class_method( - numpy.subtract, - { - "LorentzVector", - "ThreeVector", - "SphericalThreeVector", - "TwoVector", - "PolarTwoVector", - }, - transpose=False, - ) def subtract(self, other): """Subtract a vector from another elementwise using `x`, `y`, `z`, and `t` components""" - other_4d = other.to_Vector4D() return awkward.zip( { - "x": self.x - other_4d.x, - "y": self.y - other_4d.y, - "z": self.z - other_4d.z, - "t": self.t - other_4d.t, + "x": self.x - other.x, + "y": self.y - other.y, + "z": self.z - other.z, + "t": self.t - other.t, }, with_name="LorentzVector", behavior=self.behavior, @@ -966,6 +912,24 @@ def negative(self): ) +_binary_dispatch_cls = { + "TwoVector": TwoVector, + "PolarTwoVector": TwoVector, + "ThreeVector": ThreeVector, + "SphericalThreeVector": ThreeVector, + "LorentzVector": LorentzVector, + "PtEtaPhiMLorentzVector": LorentzVector, + "PtEtaPhiELorentzVector": LorentzVector, +} +_rank = [TwoVector, ThreeVector, LorentzVector] + +for lhs, lhs_to in _binary_dispatch_cls.items(): + for rhs, rhs_to in _binary_dispatch_cls.items(): + out_to = min(lhs_to, rhs_to, key=_rank.index) + behavior[(numpy.add, lhs, rhs)] = out_to.add + behavior[(numpy.subtract, lhs, rhs)] = out_to.subtract + + TwoVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 TwoVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 TwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 From fcbabdc826d794e447eebabffa1d194aad489b71 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Tue, 30 Jan 2024 09:28:26 -0600 Subject: [PATCH 23/46] skooch pin for vector --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 98dee3fbf..20d01a290 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,7 @@ dependencies = [ "dask[array]>=2023.4.0;python_version<'3.9'", "dask-awkward>=2024.3.0", "dask-histogram>=2024.3.0", - "vector>=1.1.1", + "vector>=1.2.0", "correctionlib>=2.3.3", "pyarrow>=6.0.0", "fsspec-xrootd>=0.2.3", From dc2d51f1599fdb5e57afbf755abb88b9eea3350c Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Wed, 6 Mar 2024 11:42:03 -0600 Subject: [PATCH 24/46] skooch pin for vector --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 20d01a290..693c5ac1d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,7 @@ dependencies = [ "dask[array]>=2023.4.0;python_version<'3.9'", "dask-awkward>=2024.3.0", "dask-histogram>=2024.3.0", - "vector>=1.2.0", + "vector>=1.3.0", "correctionlib>=2.3.3", "pyarrow>=6.0.0", "fsspec-xrootd>=0.2.3", From 0fa3691d45c4af830cf05d26b2c4553d7d8d607b Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Thu, 14 Mar 2024 14:07:48 +0100 Subject: [PATCH 25/46] Depend on vector v1.3 + cleanup --- src/coffea/nanoevents/methods/vector.py | 183 ++++-------------------- tests/test_nanoevents_vector.py | 16 +-- 2 files changed, 39 insertions(+), 160 deletions(-) diff --git a/src/coffea/nanoevents/methods/vector.py b/src/coffea/nanoevents/methods/vector.py index f9676b198..406ef5476 100644 --- a/src/coffea/nanoevents/methods/vector.py +++ b/src/coffea/nanoevents/methods/vector.py @@ -152,27 +152,7 @@ def absolute(self): @awkward.mixin_class_method(numpy.negative) def negative(self): """Returns the negative of the vector""" - return awkward.zip( - {"x": -self.x, "y": -self.y}, - with_name="TwoVector", - behavior=self.behavior, - ) - - def add(self, other): - """Add two vectors together elementwise using `x` and `y` components""" - return awkward.zip( - {"x": self.x + other.x, "y": self.y + other.y}, - with_name="TwoVector", - behavior=self.behavior, - ) - - def subtract(self, other): - """Subtract a vector from another elementwise using `x` and `y` components""" - return awkward.zip( - {"x": self.x - other.x, "y": self.y - other.y}, - with_name="TwoVector", - behavior=self.behavior, - ) + return self.scale(-1) def sum(self, axis=-1): """Sum an array of vectors elementwise using `x` and `y` components""" @@ -188,11 +168,7 @@ def sum(self, axis=-1): @awkward.mixin_class_method(numpy.multiply, {numbers.Number}) def multiply(self, other): """Multiply this vector by a scalar elementwise using `x` and `y` components""" - return awkward.zip( - {"x": self.x * other, "y": self.y * other}, - with_name="TwoVector", - behavior=self.behavior, - ) + return self.scale(other) @awkward.mixin_class_method(numpy.divide, {numbers.Number}) def divide(self, other): @@ -228,23 +204,12 @@ def multiply(self, other): In reality, this directly adjusts `r` and `phi` for performance """ - return awkward.zip( - { - "rho": self.r * abs(other), - "phi": self.phi % (2 * numpy.pi) - (numpy.pi * (other < 0)), - }, - with_name="PolarTwoVector", - behavior=self.behavior, - ) + return self.scale(other) @awkward.mixin_class_method(numpy.negative) def negative(self): """Returns the negative of the vector""" - return awkward.zip( - {"rho": self.r, "phi": self.phi % (2 * numpy.pi) - numpy.pi}, - with_name="PolarTwoVector", - behavior=self.behavior, - ) + return self.scale(-1) @awkward.mixin_class(behavior) @@ -279,23 +244,7 @@ def absolute(self): @awkward.mixin_class_method(numpy.negative) def negative(self): """Returns the negative of the vector""" - return awkward.zip( - {"x": -self.x, "y": -self.y, "z": -self.z}, - with_name="ThreeVector", - behavior=self.behavior, - ) - - def add(self, other): - """Add two vectors together elementwise using `x`, `y`, and `z` components""" - return awkward.zip( - { - "x": self.x + other.x, - "y": self.y + other.y, - "z": self.z + other.z, - }, - with_name="ThreeVector", - behavior=self.behavior, - ) + return self.scale(-1) @awkward.mixin_class_method(numpy.divide, {numbers.Number}) def divide(self, other): @@ -303,19 +252,6 @@ def divide(self, other): This is realized by using the multiplication functionality""" return self.multiply(1 / other) - def subtract(self, other): - """Subtract a vector from another elementwise using `x`, `y`, and `z` components""" - - return awkward.zip( - { - "x": self.x - other.x, - "y": self.y - other.y, - "z": self.z - other.z, - }, - with_name="ThreeVector", - behavior=self.behavior, - ) - def sum(self, axis=-1): """Sum an array of vectors elementwise using `x`, `y`, and `z` components""" return awkward.zip( @@ -331,27 +267,7 @@ def sum(self, axis=-1): @awkward.mixin_class_method(numpy.multiply, {numbers.Number}) def multiply(self, other): """Multiply this vector by a scalar elementwise using `x`, `y`, and `z` components""" - return awkward.zip( - {"x": self.x * other, "y": self.y * other, "z": self.z * other}, - with_name="ThreeVector", - behavior=self.behavior, - ) - - def dot(self, other): - """Compute the dot product of two vectors""" - return self.x * other.x + self.y * other.y + self.z * other.z - - def cross(self, other): - """Compute the cross product of two vectors""" - return awkward.zip( - { - "x": self.y * other.z - self.z * other.y, - "y": self.z * other.x - self.x * other.z, - "z": self.x * other.y - self.y * other.x, - }, - with_name="ThreeVector", - behavior=self.behavior, - ) + return self.scale(other) def delta_phi(self, other): """Compute difference in angle between two vectors @@ -388,28 +304,12 @@ def multiply(self, other): In reality, this directly adjusts `r`, `theta` and `phi` for performance """ - return awkward.zip( - { - "rho": self.rho * abs(other), - "theta": (numpy.sign(other) * self.theta + numpy.pi) % numpy.pi, - "phi": self.phi % (2 * numpy.pi) - numpy.pi * (other < 0), - }, - with_name="SphericalThreeVector", - behavior=self.behavior, - ) + return self.scale(other) @awkward.mixin_class_method(numpy.negative) def negative(self): """Returns the negative of the vector""" - return awkward.zip( - { - "rho": self.rho, - "theta": (-self.theta + numpy.pi) % numpy.pi, - "phi": self.phi % (2 * numpy.pi) - numpy.pi, - }, - with_name="SphericalThreeVector", - behavior=self.behavior, - ) + return self.scale(-1) def _metric_table_core(a, b, axis, metric, return_combinations): @@ -481,33 +381,6 @@ def absolute(self): """ return self.mass - def add(self, other): - """Add two vectors together elementwise using `x`, `y`, `z`, and `t` components""" - return awkward.zip( - { - "x": self.x + other.x, - "y": self.y + other.y, - "z": self.z + other.z, - "t": self.t + other.t, - }, - with_name="LorentzVector", - behavior=self.behavior, - ) - - def subtract(self, other): - """Subtract a vector from another elementwise using `x`, `y`, `z`, and `t` components""" - - return awkward.zip( - { - "x": self.x - other.x, - "y": self.y - other.y, - "z": self.z - other.z, - "t": self.t - other.t, - }, - with_name="LorentzVector", - behavior=self.behavior, - ) - def sum(self, axis=-1): """Sum an array of vectors elementwise using `x`, `y`, `z`, and `t` components""" return awkward.zip( @@ -524,22 +397,13 @@ def sum(self, axis=-1): @awkward.mixin_class_method(numpy.multiply, {numbers.Number}) def multiply(self, other): """Multiply this vector by a scalar elementwise using `x`, `y`, `z`, and `t` components""" - return awkward.zip( - { - "x": self.x * other, - "y": self.y * other, - "z": self.z * other, - "t": self.t * other, - }, - with_name="LorentzVector", - behavior=self.behavior, - ) + return self.scale(other) @awkward.mixin_class_method(numpy.divide, {numbers.Number}) def divide(self, other): """Divide this vector by a scalar elementwise using its cartesian components This is realized by using the multiplication functionality""" - return self.multiply(1 / other) + return self.scale(1 / other) def delta_r2(self, other): """Squared `delta_r`""" @@ -562,11 +426,7 @@ def delta_phi(self, other): @awkward.mixin_class_method(numpy.negative) def negative(self): """Returns the negative of the vector""" - return awkward.zip( - {"x": -self.x, "y": -self.y, "z": -self.z, "t": -self.t}, - with_name="LorentzVector", - behavior=self.behavior, - ) + return self.scale(-1) @property def pvec(self): @@ -795,6 +655,12 @@ def negative(self): behavior=self.behavior, ) + @awkward.mixin_class_method(numpy.divide, {numbers.Number}) + def divide(self, other): + """Divide this vector by a scalar elementwise using its cartesian components + This is realized by using the multiplication functionality""" + return self.multiply(1 / other) + @awkward.mixin_class(behavior) class PtEtaPhiELorentzVector(LorentzVector): @@ -911,6 +777,12 @@ def negative(self): behavior=self.behavior, ) + @awkward.mixin_class_method(numpy.divide, {numbers.Number}) + def divide(self, other): + """Divide this vector by a scalar elementwise using its cartesian components + This is realized by using the multiplication functionality""" + return self.multiply(1 / other) + _binary_dispatch_cls = { "TwoVector": TwoVector, @@ -933,30 +805,37 @@ def negative(self): TwoVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 TwoVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 TwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +TwoVectorArray.MomentumClass = PolarTwoVectorArray # noqa: F821 PolarTwoVectorArray.ProjectionClass2D = PolarTwoVectorArray # noqa: F821 PolarTwoVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 PolarTwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +PolarTwoVectorArray.MomentumClass = PolarTwoVectorArray # noqa: F821 ThreeVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 ThreeVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 ThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +ThreeVectorArray.MomentumClass = SphericalThreeVectorArray # noqa: F821 -SphericalThreeVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +SphericalThreeVectorArray.ProjectionClass2D = PolarTwoVectorArray # noqa: F821 SphericalThreeVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 SphericalThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +SphericalThreeVectorArray.MomentumClass = SphericalThreeVectorArray # noqa: F821 LorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 LorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 LorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +LorentzVectorArray.MomentumClass = LorentzVectorArray # noqa: F821 PtEtaPhiMLorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 PtEtaPhiMLorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 PtEtaPhiMLorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +PtEtaPhiMLorentzVectorArray.MomentumClass = LorentzVectorArray # noqa: F821 PtEtaPhiELorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 PtEtaPhiELorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 PtEtaPhiELorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +PtEtaPhiELorentzVectorArray.MomentumClass = LorentzVectorArray # noqa: F821 __all__ = [ "TwoVector", diff --git a/tests/test_nanoevents_vector.py b/tests/test_nanoevents_vector.py index 8fdb8857b..6aa13d164 100644 --- a/tests/test_nanoevents_vector.py +++ b/tests/test_nanoevents_vector.py @@ -12,7 +12,7 @@ def assert_record_arrays_equal(a, b, check_type=False): if check_type: assert type(a) is type(b) assert ak.fields(a) == ak.fields(b) - assert all(ak.all(a[f] == b[f]) for f in ak.fields(a)) + assert all(ak.all(ak.isclose(a[f], b[f])) for f in ak.fields(a)) def assert_awkward_allclose(actual, desired): @@ -97,7 +97,7 @@ def test_polar_two_vector(): assert ak.all(abs((-a).y + a.y) < ATOL) assert_record_arrays_equal(a * (-1), -a) - assert ak.all(a.unit.phi == a.phi) + assert ak.all(ak.isclose(a.unit.phi, a.phi)) def test_three_vector(): @@ -539,17 +539,17 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): behavior=vector.behavior, ) - assert_record_arrays_equal(a + b, b + a, check_type=True) - assert_record_arrays_equal(a + c, c + a, check_type=True) - assert_record_arrays_equal(b + c, c + b, check_type=True) + assert_record_arrays_equal(a.like(b) + b, b + a.like(b), check_type=True) + assert_record_arrays_equal(a.like(c) + c, c + a.like(c), check_type=True) + assert_record_arrays_equal(b.like(c) + c, c + b.like(c), check_type=True) assert_allclose(a.delta_phi(b), -b.delta_phi(a)) assert_allclose(a.delta_phi(c), -c.delta_phi(a)) assert_allclose(b.delta_phi(c), -c.delta_phi(b)) - assert_record_arrays_equal((a - b), -(b - a), check_type=True) - assert_record_arrays_equal((a - c), -(c - a), check_type=True) - assert_record_arrays_equal((b - c), -(c - b), check_type=True) + assert_record_arrays_equal((a.like(b) - b), -(b - a.like(b)), check_type=True) + assert_record_arrays_equal((a.like(c) - c), -(c - a.like(c)), check_type=True) + assert_record_arrays_equal((b.like(c) - c), -(c - b.like(c)), check_type=True) @pytest.mark.parametrize("optimization_enabled", [True, False]) From c0d3342dc05b6ac746034daf808493853f9b676d Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Thu, 14 Mar 2024 14:29:43 +0100 Subject: [PATCH 26/46] Candidate classes need MmomentumClass too --- src/coffea/nanoevents/methods/candidate.py | 3 +++ src/coffea/nanoevents/methods/vector.py | 4 ++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/src/coffea/nanoevents/methods/candidate.py b/src/coffea/nanoevents/methods/candidate.py index 2291a744a..19ab3a457 100644 --- a/src/coffea/nanoevents/methods/candidate.py +++ b/src/coffea/nanoevents/methods/candidate.py @@ -70,4 +70,7 @@ class PtEtaPhiECandidate(Candidate, vector.PtEtaPhiELorentzVector): pass +PtEtaPhiMCandidate.MomentumClass = PtEtaPhiMCandidate +PtEtaPhiECandidate.MomentumClass = PtEtaPhiECandidate + __all__ = ["Candidate", "PtEtaPhiMCandidate", "PtEtaPhiECandidate"] diff --git a/src/coffea/nanoevents/methods/vector.py b/src/coffea/nanoevents/methods/vector.py index 406ef5476..2944aa72d 100644 --- a/src/coffea/nanoevents/methods/vector.py +++ b/src/coffea/nanoevents/methods/vector.py @@ -175,7 +175,7 @@ def divide(self, other): """Divide this vector by a scalar elementwise using its cartesian components This is realized by using the multiplication functionality""" - return self.multiply(1 / other) + return self.scale(1 / other) def delta_phi(self, other): """Compute difference in angle between two vectors @@ -250,7 +250,7 @@ def negative(self): def divide(self, other): """Divide this vector by a scalar elementwise using its cartesian components This is realized by using the multiplication functionality""" - return self.multiply(1 / other) + return self.scale(1 / other) def sum(self, axis=-1): """Sum an array of vectors elementwise using `x`, `y`, and `z` components""" From 7e83eeacd388e1fb2d6b70e59c696d1690cdcc69 Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Fri, 15 Mar 2024 14:23:50 +0100 Subject: [PATCH 27/46] Depend on vector v1.3.1 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 693c5ac1d..61ddb17d5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,7 @@ dependencies = [ "dask[array]>=2023.4.0;python_version<'3.9'", "dask-awkward>=2024.3.0", "dask-histogram>=2024.3.0", - "vector>=1.3.0", + "vector>=1.3.1", "correctionlib>=2.3.3", "pyarrow>=6.0.0", "fsspec-xrootd>=0.2.3", From 4f90c2b3db1023d34b92a864c456aa8b70c9ce29 Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Fri, 15 Mar 2024 15:17:14 +0100 Subject: [PATCH 28/46] MomentumClass for Jet and Candidate --- src/coffea/nanoevents/methods/candidate.py | 1 + src/coffea/nanoevents/methods/nanoaod.py | 3 +++ 2 files changed, 4 insertions(+) diff --git a/src/coffea/nanoevents/methods/candidate.py b/src/coffea/nanoevents/methods/candidate.py index 19ab3a457..c1ed00d39 100644 --- a/src/coffea/nanoevents/methods/candidate.py +++ b/src/coffea/nanoevents/methods/candidate.py @@ -70,6 +70,7 @@ class PtEtaPhiECandidate(Candidate, vector.PtEtaPhiELorentzVector): pass +Candidate.MomentumClass = Candidate PtEtaPhiMCandidate.MomentumClass = PtEtaPhiMCandidate PtEtaPhiECandidate.MomentumClass = PtEtaPhiECandidate diff --git a/src/coffea/nanoevents/methods/nanoaod.py b/src/coffea/nanoevents/methods/nanoaod.py index 3d17d9280..fcd44a086 100644 --- a/src/coffea/nanoevents/methods/nanoaod.py +++ b/src/coffea/nanoevents/methods/nanoaod.py @@ -446,6 +446,9 @@ def constituents(self, dask_array): ) +JetArray.MomentumClass = JetArray # noqa: F821 + + _set_repr_name("Jet") From ad3d6c3575fd0b9fcac198f1bb10d33d3eb08d19 Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Wed, 20 Mar 2024 20:53:45 +0100 Subject: [PATCH 29/46] fix: add correct behaviors and projection classes for candidate classes --- src/coffea/nanoevents/methods/candidate.py | 41 +++++++++++++++++----- 1 file changed, 32 insertions(+), 9 deletions(-) diff --git a/src/coffea/nanoevents/methods/candidate.py b/src/coffea/nanoevents/methods/candidate.py index c1ed00d39..99531509f 100644 --- a/src/coffea/nanoevents/methods/candidate.py +++ b/src/coffea/nanoevents/methods/candidate.py @@ -10,10 +10,8 @@ from coffea.nanoevents.methods import vector -behavior = dict(vector.behavior) - -@awkward.mixin_class(behavior) +@awkward.mixin_class(vector.behavior) class Candidate(vector.LorentzVector): """A Lorentz vector with charge @@ -50,8 +48,8 @@ def sum(self, axis=-1): ) -@awkward.mixin_class(behavior) -class PtEtaPhiMCandidate(Candidate, vector.PtEtaPhiMLorentzVector): +@awkward.mixin_class(vector.behavior) +class PtEtaPhiMCandidate(Candidate, vector.PtEtaPhiMCandidate): """A Lorentz vector in eta, mass coordinates with charge This mixin class requires the parent class to provide items `pt`, `eta`, `phi`, `mass`, and `charge`. @@ -60,7 +58,7 @@ class PtEtaPhiMCandidate(Candidate, vector.PtEtaPhiMLorentzVector): pass -@awkward.mixin_class(behavior) +@awkward.mixin_class(vector.behavior) class PtEtaPhiECandidate(Candidate, vector.PtEtaPhiELorentzVector): """A Lorentz vector in eta, energy coordinates with charge @@ -70,8 +68,33 @@ class PtEtaPhiECandidate(Candidate, vector.PtEtaPhiELorentzVector): pass -Candidate.MomentumClass = Candidate -PtEtaPhiMCandidate.MomentumClass = PtEtaPhiMCandidate -PtEtaPhiECandidate.MomentumClass = PtEtaPhiECandidate +vector._binary_dispatch_cls.update( + { + "Candidate": Candidate, + } +) +vector._rank += [Candidate] + +for lhs, lhs_to in vector._binary_dispatch_cls.items(): + for rhs, rhs_to in vector._binary_dispatch_cls.items(): + if lhs == "Candidate" or rhs == "Candidate": + out_to = min(lhs_to, rhs_to, key=vector._rank.index) + vector.behavior[(numpy.add, lhs, rhs)] = out_to.add + vector.behavior[(numpy.subtract, lhs, rhs)] = out_to.subtract + +CandidateArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +CandidateArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +CandidateArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821 +CandidateArray.MomentumClass = CandidateArray # noqa: F821 + +PtEtaPhiMCandidateArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +PtEtaPhiMCandidateArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +PtEtaPhiMCandidateArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821 +PtEtaPhiMCandidateArray.MomentumClass = PtEtaPhiMCandidateArray # noqa: F821 + +PtEtaPhiECandidateArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +PtEtaPhiECandidateArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +PtEtaPhiECandidateArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821 +PtEtaPhiECandidate.MomentumClass = PtEtaPhiECandidateArray # noqa: F821 __all__ = ["Candidate", "PtEtaPhiMCandidate", "PtEtaPhiECandidate"] From 001abe0d90a78f70ca22bf01b0c27c36de786fe3 Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Wed, 20 Mar 2024 20:56:28 +0100 Subject: [PATCH 30/46] typo --- src/coffea/nanoevents/methods/candidate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/coffea/nanoevents/methods/candidate.py b/src/coffea/nanoevents/methods/candidate.py index 99531509f..c1db112b9 100644 --- a/src/coffea/nanoevents/methods/candidate.py +++ b/src/coffea/nanoevents/methods/candidate.py @@ -49,7 +49,7 @@ def sum(self, axis=-1): @awkward.mixin_class(vector.behavior) -class PtEtaPhiMCandidate(Candidate, vector.PtEtaPhiMCandidate): +class PtEtaPhiMCandidate(Candidate, vector.PtEtaPhiMLorentzVector): """A Lorentz vector in eta, mass coordinates with charge This mixin class requires the parent class to provide items `pt`, `eta`, `phi`, `mass`, and `charge`. From e96fe5b9e701cde6d33eebe53bb5d1aceff208e9 Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Wed, 20 Mar 2024 21:08:30 +0100 Subject: [PATCH 31/46] fix behavior --- src/coffea/nanoevents/methods/candidate.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/coffea/nanoevents/methods/candidate.py b/src/coffea/nanoevents/methods/candidate.py index c1db112b9..cfbe5c18d 100644 --- a/src/coffea/nanoevents/methods/candidate.py +++ b/src/coffea/nanoevents/methods/candidate.py @@ -10,8 +10,10 @@ from coffea.nanoevents.methods import vector +behavior = dict(vector.behavior) -@awkward.mixin_class(vector.behavior) + +@awkward.mixin_class(behavior) class Candidate(vector.LorentzVector): """A Lorentz vector with charge @@ -48,7 +50,7 @@ def sum(self, axis=-1): ) -@awkward.mixin_class(vector.behavior) +@awkward.mixin_class(behavior) class PtEtaPhiMCandidate(Candidate, vector.PtEtaPhiMLorentzVector): """A Lorentz vector in eta, mass coordinates with charge @@ -58,7 +60,7 @@ class PtEtaPhiMCandidate(Candidate, vector.PtEtaPhiMLorentzVector): pass -@awkward.mixin_class(vector.behavior) +@awkward.mixin_class(behavior) class PtEtaPhiECandidate(Candidate, vector.PtEtaPhiELorentzVector): """A Lorentz vector in eta, energy coordinates with charge @@ -79,8 +81,8 @@ class PtEtaPhiECandidate(Candidate, vector.PtEtaPhiELorentzVector): for rhs, rhs_to in vector._binary_dispatch_cls.items(): if lhs == "Candidate" or rhs == "Candidate": out_to = min(lhs_to, rhs_to, key=vector._rank.index) - vector.behavior[(numpy.add, lhs, rhs)] = out_to.add - vector.behavior[(numpy.subtract, lhs, rhs)] = out_to.subtract + behavior[(numpy.add, lhs, rhs)] = out_to.add + behavior[(numpy.subtract, lhs, rhs)] = out_to.subtract CandidateArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 CandidateArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 From b4ae25105c3394a16f15f7e3d0472788bb27b1b7 Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Thu, 30 May 2024 22:25:09 +0200 Subject: [PATCH 32/46] use copy_behaviors to fix ufuncs --- src/coffea/nanoevents/methods/delphes.py | 54 ++++++++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/src/coffea/nanoevents/methods/delphes.py b/src/coffea/nanoevents/methods/delphes.py index 1f22e814e..aa56ab37c 100644 --- a/src/coffea/nanoevents/methods/delphes.py +++ b/src/coffea/nanoevents/methods/delphes.py @@ -130,6 +130,12 @@ def z(self): _set_repr_name("Vertex") +awkward._util.copy_behaviors(vector.LorentzVector, Vertex, behavior) + +VertexArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +VertexArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +VertexArray.ProjectionClass4D = VertexArray # noqa: F821 +VertexArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -170,6 +176,12 @@ def mass(self): _set_repr_name("Particle") +awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Particle, behavior) + +ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ParticleArray.ProjectionClass4D = ParticleArray # noqa: F821 +ParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -180,6 +192,12 @@ def mass(self): _set_repr_name("MasslessParticle") +awkward._util.copy_behaviors(Particle, MasslessParticle, behavior) + +MasslessParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +MasslessParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +MasslessParticleArray.ProjectionClass4D = MasslessParticleArray # noqa: F821 +MasslessParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -187,6 +205,12 @@ class Photon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Photon") +awkward._util.copy_behaviors(MasslessParticle, Photon, behavior) + +PhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +PhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +PhotonArray.ProjectionClass4D = PhotonArray # noqa: F821 +PhotonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -194,6 +218,12 @@ class Electron(MasslessParticle, base.NanoCollection): ... _set_repr_name("Electron") +awkward._util.copy_behaviors(MasslessParticle, Electron, behavior) + +ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 +ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -201,6 +231,12 @@ class Muon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Muon") +awkward._util.copy_behaviors(MasslessParticle, Muon, behavior) + +MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +MuonArray.ProjectionClass4D = MuonArray # noqa: F821 +MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -208,6 +244,12 @@ class Jet(Particle, base.NanoCollection): ... _set_repr_name("Jet") +awkward._util.copy_behaviors(Particle, Jet, behavior) + +JetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +JetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +JetArray.ProjectionClass4D = JetArray # noqa: F821 +JetArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -215,6 +257,12 @@ class Track(Particle, base.NanoCollection): ... _set_repr_name("Track") +awkward._util.copy_behaviors(Particle, Track, behavior) + +TrackArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TrackArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TrackArray.ProjectionClass4D = TrackArray # noqa: F821 +TrackArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -225,6 +273,12 @@ def pt(self): _set_repr_name("Tower") +awkward._util.copy_behaviors(MasslessParticle, Tower, behavior) + +TowerArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TowerArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TowerArray.ProjectionClass4D = TowerArray # noqa: F821 +TowerArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 __all__ = [ From b736e8e4b1c54f29a53955a2887772bbcfcf5ff0 Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Fri, 7 Jun 2024 13:47:58 +0200 Subject: [PATCH 33/46] Typo --- src/coffea/nanoevents/methods/candidate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/coffea/nanoevents/methods/candidate.py b/src/coffea/nanoevents/methods/candidate.py index cfbe5c18d..fc2afba3c 100644 --- a/src/coffea/nanoevents/methods/candidate.py +++ b/src/coffea/nanoevents/methods/candidate.py @@ -97,6 +97,6 @@ class PtEtaPhiECandidate(Candidate, vector.PtEtaPhiELorentzVector): PtEtaPhiECandidateArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 PtEtaPhiECandidateArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 PtEtaPhiECandidateArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821 -PtEtaPhiECandidate.MomentumClass = PtEtaPhiECandidateArray # noqa: F821 +PtEtaPhiECandidateArray.MomentumClass = PtEtaPhiECandidateArray # noqa: F821 __all__ = ["Candidate", "PtEtaPhiMCandidate", "PtEtaPhiECandidate"] From 539d886eb1bd34b2bed84c5668c3fcd294dc3f05 Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Sat, 29 Jun 2024 00:55:42 +0530 Subject: [PATCH 34/46] fix: copy behaviors for the remaining vector sub-classes --- src/coffea/nanoevents/methods/candidate.py | 14 +--- src/coffea/nanoevents/methods/delphes.py | 28 +++++-- src/coffea/nanoevents/methods/nanoaod.py | 95 +++++++++++++++++++++- src/coffea/nanoevents/methods/pdune.py | 36 ++++++++ src/coffea/nanoevents/methods/physlite.py | 36 ++++++++ 5 files changed, 185 insertions(+), 24 deletions(-) diff --git a/src/coffea/nanoevents/methods/candidate.py b/src/coffea/nanoevents/methods/candidate.py index fc2afba3c..e3c427276 100644 --- a/src/coffea/nanoevents/methods/candidate.py +++ b/src/coffea/nanoevents/methods/candidate.py @@ -70,19 +70,9 @@ class PtEtaPhiECandidate(Candidate, vector.PtEtaPhiELorentzVector): pass -vector._binary_dispatch_cls.update( - { - "Candidate": Candidate, - } +awkward.behavior.update( + awkward._util.copy_behaviors(vector.LorentzVector, Candidate, behavior) ) -vector._rank += [Candidate] - -for lhs, lhs_to in vector._binary_dispatch_cls.items(): - for rhs, rhs_to in vector._binary_dispatch_cls.items(): - if lhs == "Candidate" or rhs == "Candidate": - out_to = min(lhs_to, rhs_to, key=vector._rank.index) - behavior[(numpy.add, lhs, rhs)] = out_to.add - behavior[(numpy.subtract, lhs, rhs)] = out_to.subtract CandidateArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 CandidateArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 diff --git a/src/coffea/nanoevents/methods/delphes.py b/src/coffea/nanoevents/methods/delphes.py index aa56ab37c..c232fedbe 100644 --- a/src/coffea/nanoevents/methods/delphes.py +++ b/src/coffea/nanoevents/methods/delphes.py @@ -130,7 +130,9 @@ def z(self): _set_repr_name("Vertex") -awkward._util.copy_behaviors(vector.LorentzVector, Vertex, behavior) +awkward.behavior.update( + awkward._util.copy_behaviors(vector.LorentzVector, Vertex, behavior) +) VertexArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 VertexArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -176,7 +178,9 @@ def mass(self): _set_repr_name("Particle") -awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Particle, behavior) +awkward.behavior.update( + awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Particle, behavior) +) ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -192,7 +196,9 @@ def mass(self): _set_repr_name("MasslessParticle") -awkward._util.copy_behaviors(Particle, MasslessParticle, behavior) +awkward.behavior.update( + awkward._util.copy_behaviors(Particle, MasslessParticle, behavior) +) MasslessParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 MasslessParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -205,7 +211,9 @@ class Photon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Photon") -awkward._util.copy_behaviors(MasslessParticle, Photon, behavior) +awkward.behavior.update( + awkward._util.copy_behaviors(MasslessParticle, Photon, behavior) +) PhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 PhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -218,7 +226,9 @@ class Electron(MasslessParticle, base.NanoCollection): ... _set_repr_name("Electron") -awkward._util.copy_behaviors(MasslessParticle, Electron, behavior) +awkward.behavior.update( + awkward._util.copy_behaviors(MasslessParticle, Electron, behavior) +) ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -231,7 +241,7 @@ class Muon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Muon") -awkward._util.copy_behaviors(MasslessParticle, Muon, behavior) +awkward.behavior.update(awkward._util.copy_behaviors(MasslessParticle, Muon, behavior)) MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -244,7 +254,7 @@ class Jet(Particle, base.NanoCollection): ... _set_repr_name("Jet") -awkward._util.copy_behaviors(Particle, Jet, behavior) +awkward.behavior.update(awkward._util.copy_behaviors(Particle, Jet, behavior)) JetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 JetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -257,7 +267,7 @@ class Track(Particle, base.NanoCollection): ... _set_repr_name("Track") -awkward._util.copy_behaviors(Particle, Track, behavior) +awkward.behavior.update(awkward._util.copy_behaviors(Particle, Track, behavior)) TrackArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 TrackArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -273,7 +283,7 @@ def pt(self): _set_repr_name("Tower") -awkward._util.copy_behaviors(MasslessParticle, Tower, behavior) +awkward.behavior.update(awkward._util.copy_behaviors(MasslessParticle, Tower, behavior)) TowerArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 TowerArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 diff --git a/src/coffea/nanoevents/methods/nanoaod.py b/src/coffea/nanoevents/methods/nanoaod.py index fcd44a086..f9c699eb4 100644 --- a/src/coffea/nanoevents/methods/nanoaod.py +++ b/src/coffea/nanoevents/methods/nanoaod.py @@ -38,6 +38,18 @@ class PtEtaPhiMCollection(vector.PtEtaPhiMLorentzVector, base.NanoCollection): pass +awkward.behavior.update( + awkward._util.copy_behaviors( + vector.PtEtaPhiMLorentzVector, PtEtaPhiMCollection, behavior + ) +) + +PtEtaPhiMCollectionArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +PtEtaPhiMCollectionArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +PtEtaPhiMCollectionArray.ProjectionClass4D = PtEtaPhiMCollectionArray # noqa: F821 +PtEtaPhiMCollectionArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + + @awkward.mixin_class(behavior) class GenParticle(vector.PtEtaPhiMLorentzVector, base.NanoCollection): """NanoAOD generator-level particle object, including parent and child self-references @@ -144,6 +156,14 @@ def distinctChildrenDeep(self, dask_array): _set_repr_name("GenParticle") +awkward.behavior.update( + awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, GenParticle, behavior) +) + +GenParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +GenParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +GenParticleArray.ProjectionClass4D = GenParticleArray # noqa: F821 +GenParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -164,6 +184,14 @@ def parent(self, dask_array): _set_repr_name("GenVisTau") +awkward.behavior.update( + awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, GenVisTau, behavior) +) + +GenVisTauArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +GenVisTauArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +GenVisTauArray.ProjectionClass4D = GenVisTauArray # noqa: F821 +GenVisTauArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -228,6 +256,14 @@ def matched_photon(self, dask_array): _set_repr_name("Electron") +awkward.behavior.update( + awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Electron, behavior) +) + +ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 +ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -262,6 +298,14 @@ def matched_jet(self, dask_array): _set_repr_name("Muon") +awkward.behavior.update( + awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Muon, behavior) +) + +MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +MuonArray.ProjectionClass4D = MuonArray # noqa: F821 +MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -286,6 +330,14 @@ def matched_jet(self, dask_array): _set_repr_name("Tau") +awkward.behavior.update( + awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Tau, behavior) +) + +TauArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TauArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TauArray.ProjectionClass4D = TauArray # noqa: F821 +TauArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -361,6 +413,14 @@ def matched_jet(self, dask_array): _set_repr_name("Photon") +awkward.behavior.update( + awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Photon, behavior) +) + +PhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +PhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +PhotonArray.ProjectionClass4D = PhotonArray # noqa: F821 +PhotonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -377,6 +437,14 @@ def matched_muon(self, dask_array): _set_repr_name("FsrPhoton") +awkward.behavior.update( + awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, FsrPhoton, behavior) +) + +FsrPhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +FsrPhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +FsrPhotonArray.ProjectionClass4D = FsrPhotonArray # noqa: F821 +FsrPhotonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -446,10 +514,15 @@ def constituents(self, dask_array): ) -JetArray.MomentumClass = JetArray # noqa: F821 - - _set_repr_name("Jet") +awkward.behavior.update( + awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Jet, behavior) +) + +JetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +JetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +JetArray.ProjectionClass4D = JetArray # noqa: F821 +JetArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -512,6 +585,14 @@ def constituents(self, dask_array): _set_repr_name("FatJet") +awkward.behavior.update( + awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, FatJet, behavior) +) + +FatJetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +FatJetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +FatJetArray.ProjectionClass4D = FatJetArray # noqa: F821 +FatJetArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -524,6 +605,14 @@ def r(self): _set_repr_name("MissingET") +awkward.behavior.update( + awkward._util.copy_behaviors(vector.PolarTwoVector, MissingET, behavior) +) + +MissingETArray.ProjectionClass2D = MissingETArray # noqa: F821 +MissingETArray.ProjectionClass3D = vector.SphericalThreeVectorArray # noqa: F821 +MissingETArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821 +MissingETArray.MomentumClass = MissingETArray # noqa: F821 @awkward.mixin_class(behavior) diff --git a/src/coffea/nanoevents/methods/pdune.py b/src/coffea/nanoevents/methods/pdune.py index aff8e7730..4c4296008 100644 --- a/src/coffea/nanoevents/methods/pdune.py +++ b/src/coffea/nanoevents/methods/pdune.py @@ -92,6 +92,14 @@ def mass(self): _set_repr_name("Particle") +awkward.behavior.update( + awkward._util.copy_behaviors(vector.LorentzVector, Particle, behavior) +) + +ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ParticleArray.ProjectionClass4D = ParticleArray # noqa: F821 +ParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -130,6 +138,14 @@ def t(self): _set_repr_name("TrackParticle") +awkward.behavior.update( + awkward._util.copy_behaviors(vector.LorentzVector, TrackParticle, behavior) +) + +TrackParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TrackParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TrackParticleArray.ProjectionClass4D = TrackParticleArray # noqa: F821 +TrackParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -149,6 +165,12 @@ def trackParticle(self): _set_repr_name("Muon") +awkward.behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) + +MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +MuonArray.ProjectionClass4D = MuonArray # noqa: F821 +MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -176,6 +198,12 @@ def trackParticle(self): _set_repr_name("Electron") +awkward.behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) + +ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 +ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -218,3 +246,11 @@ def parents(self): _set_repr_name("TruthParticle") +awkward.behavior.update( + awkward._util.copy_behaviors(vector.LorentzVector, TruthParticle, behavior) +) + +TruthParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TruthParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TruthParticleArray.ProjectionClass4D = TruthParticleArray # noqa: F821 +TruthParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 diff --git a/src/coffea/nanoevents/methods/physlite.py b/src/coffea/nanoevents/methods/physlite.py index a20b7e8ec..28554a6c9 100644 --- a/src/coffea/nanoevents/methods/physlite.py +++ b/src/coffea/nanoevents/methods/physlite.py @@ -143,6 +143,14 @@ def mass(self): _set_repr_name("Particle") +awkward.behavior.update( + awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Particle, behavior) +) + +ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ParticleArray.ProjectionClass4D = ParticleArray # noqa: F821 +ParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -181,6 +189,14 @@ def t(self): _set_repr_name("TrackParticle") +awkward.behavior.update( + awkward._util.copy_behaviors(vector.LorentzVector, TrackParticle, behavior) +) + +TrackParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TrackParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TrackParticleArray.ProjectionClass4D = TrackParticleArray # noqa: F821 +TrackParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -209,6 +225,12 @@ def trackParticle(self, dask_array): _set_repr_name("Muon") +awkward.behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) + +MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +MuonArray.ProjectionClass4D = MuonArray # noqa: F821 +MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -261,6 +283,12 @@ def caloClusters(self, dask_array): _set_repr_name("Electron") +awkward.behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) + +ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 +ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -303,3 +331,11 @@ def parents(self): _set_repr_name("TruthParticle") +awkward.behavior.update( + awkward._util.copy_behaviors(vector.LorentzVector, TruthParticle, behavior) +) + +TruthParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TruthParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TruthParticleArray.ProjectionClass4D = TruthParticleArray # noqa: F821 +TruthParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 From 45d25981d6cc5427431290593336fdb2c1832b26 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Sun, 30 Jun 2024 10:19:31 -0500 Subject: [PATCH 35/46] awkward 2.6.6 is required --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 14b2155b5..64f979547 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,7 @@ classifiers = [ "Topic :: Utilities", ] dependencies = [ - "awkward>=2.6.3", + "awkward>=2.6.6", "uproot>=5.3.0,!=5.3.3,!=5.3.4,!=5.3.5", "dask[array]>=2024.3.0;python_version>'3.8'", "dask[array]>=2023.4.0;python_version<'3.9'", From deb8e28553217b512757b3219a397304d897d595 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Mon, 1 Jul 2024 07:57:35 -0500 Subject: [PATCH 36/46] update LowPtElectron --- src/coffea/nanoevents/methods/nanoaod.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/coffea/nanoevents/methods/nanoaod.py b/src/coffea/nanoevents/methods/nanoaod.py index 748e3dada..8fa587410 100644 --- a/src/coffea/nanoevents/methods/nanoaod.py +++ b/src/coffea/nanoevents/methods/nanoaod.py @@ -298,6 +298,14 @@ def matched_photon(self, dask_array): _set_repr_name("LowPtElectron") +awkward.behavior.update( + awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, LowPtElectron, behavior) +) + +LowPtElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +LowPtElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +LowPtElectronArray.ProjectionClass4D = LowPtElectronArray # noqa: F821 +LowPtElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) From 39611c7da85cda2aaa7d75024208c2255ed06cea Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Thu, 11 Jul 2024 18:32:35 +0200 Subject: [PATCH 37/46] Review changes - a new util function to register projection classes - do not modify the global awkward behavior - tests for geometric dimension errors --- src/coffea/nanoevents/methods/candidate.py | 39 +++--- src/coffea/nanoevents/methods/delphes.py | 134 +++++++++++-------- src/coffea/nanoevents/methods/nanoaod.py | 144 +++++++++++++-------- src/coffea/nanoevents/methods/pdune.py | 72 ++++++----- src/coffea/nanoevents/methods/physlite.py | 70 ++++++---- src/coffea/nanoevents/methods/vector.py | 86 +++++++----- src/coffea/util.py | 7 + tests/test_nanoevents_vector.py | 14 ++ 8 files changed, 347 insertions(+), 219 deletions(-) diff --git a/src/coffea/nanoevents/methods/candidate.py b/src/coffea/nanoevents/methods/candidate.py index e3c427276..43eca9ea7 100644 --- a/src/coffea/nanoevents/methods/candidate.py +++ b/src/coffea/nanoevents/methods/candidate.py @@ -9,6 +9,7 @@ import numpy from coffea.nanoevents.methods import vector +from coffea.util import register_projection_classes behavior = dict(vector.behavior) @@ -70,23 +71,29 @@ class PtEtaPhiECandidate(Candidate, vector.PtEtaPhiELorentzVector): pass -awkward.behavior.update( - awkward._util.copy_behaviors(vector.LorentzVector, Candidate, behavior) -) - -CandidateArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -CandidateArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -CandidateArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821 -CandidateArray.MomentumClass = CandidateArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors(vector.LorentzVector, Candidate, behavior)) -PtEtaPhiMCandidateArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -PtEtaPhiMCandidateArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -PtEtaPhiMCandidateArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821 -PtEtaPhiMCandidateArray.MomentumClass = PtEtaPhiMCandidateArray # noqa: F821 +register_projection_classes( + CandidateArray, # noqa: F821 + vector.PolarTwoVectorArray, + vector.SphericalThreeVectorArray, + vector.LorentzVectorArray, + CandidateArray, # noqa: F821 +) +register_projection_classes( + PtEtaPhiMCandidateArray, # noqa: F821 + vector.PolarTwoVectorArray, + vector.SphericalThreeVectorArray, + vector.LorentzVectorArray, + PtEtaPhiMCandidateArray, # noqa: F821 +) +register_projection_classes( + PtEtaPhiECandidateArray, # noqa: F821 + vector.PolarTwoVectorArray, + vector.SphericalThreeVectorArray, + vector.LorentzVectorArray, + PtEtaPhiECandidateArray, # noqa: F821 +) -PtEtaPhiECandidateArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -PtEtaPhiECandidateArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -PtEtaPhiECandidateArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821 -PtEtaPhiECandidateArray.MomentumClass = PtEtaPhiECandidateArray # noqa: F821 __all__ = ["Candidate", "PtEtaPhiMCandidate", "PtEtaPhiECandidate"] diff --git a/src/coffea/nanoevents/methods/delphes.py b/src/coffea/nanoevents/methods/delphes.py index c232fedbe..cc65121f6 100644 --- a/src/coffea/nanoevents/methods/delphes.py +++ b/src/coffea/nanoevents/methods/delphes.py @@ -7,6 +7,7 @@ import numpy from coffea.nanoevents.methods import base, candidate, vector +from coffea.util import register_projection_classes behavior = {} behavior.update(base.behavior) @@ -130,15 +131,16 @@ def z(self): _set_repr_name("Vertex") -awkward.behavior.update( - awkward._util.copy_behaviors(vector.LorentzVector, Vertex, behavior) +behavior.update(awkward._util.copy_behaviors(vector.LorentzVector, Vertex, behavior)) + +register_projection_classes( + VertexArray, # noqa: F821 + vector.TwoVectorArray, # noqa: F821 + vector.ThreeVectorArray, # noqa: F821 + VertexArray, # noqa: F821 + vector.LorentzVectorArray, # noqa: F821 ) -VertexArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -VertexArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -VertexArray.ProjectionClass4D = VertexArray # noqa: F821 -VertexArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 - @awkward.mixin_class(behavior) class Particle(vector.PtEtaPhiMLorentzVector): @@ -178,14 +180,17 @@ def mass(self): _set_repr_name("Particle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Particle, behavior) ) -ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -ParticleArray.ProjectionClass4D = ParticleArray # noqa: F821 -ParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +register_projection_classes( + ParticleArray, # noqa: F821 + vector.TwoVectorArray, # noqa: F821 + vector.ThreeVectorArray, # noqa: F821 + ParticleArray, # noqa: F821 + vector.LorentzVectorArray, # noqa: F821 +) @awkward.mixin_class(behavior) @@ -196,57 +201,63 @@ def mass(self): _set_repr_name("MasslessParticle") -awkward.behavior.update( - awkward._util.copy_behaviors(Particle, MasslessParticle, behavior) +behavior.update(awkward._util.copy_behaviors(Particle, MasslessParticle, behavior)) + +register_projection_classes( + MasslessParticleArray, # noqa: F821 + vector.TwoVectorArray, # noqa: F821 + vector.ThreeVectorArray, # noqa: F821 + MasslessParticleArray, # noqa: F821 + vector.LorentzVectorArray, # noqa: F821 ) -MasslessParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -MasslessParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -MasslessParticleArray.ProjectionClass4D = MasslessParticleArray # noqa: F821 -MasslessParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 - @awkward.mixin_class(behavior) class Photon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Photon") -awkward.behavior.update( - awkward._util.copy_behaviors(MasslessParticle, Photon, behavior) +behavior.update(awkward._util.copy_behaviors(MasslessParticle, Photon, behavior)) + +register_projection_classes( + PhotonArray, # noqa: F821 + vector.TwoVectorArray, # noqa: F821 + vector.ThreeVectorArray, # noqa: F821 + PhotonArray, # noqa: F821 + vector.LorentzVectorArray, # noqa: F821 ) -PhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -PhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -PhotonArray.ProjectionClass4D = PhotonArray # noqa: F821 -PhotonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 - @awkward.mixin_class(behavior) class Electron(MasslessParticle, base.NanoCollection): ... _set_repr_name("Electron") -awkward.behavior.update( - awkward._util.copy_behaviors(MasslessParticle, Electron, behavior) +behavior.update(awkward._util.copy_behaviors(MasslessParticle, Electron, behavior)) + +register_projection_classes( + ElectronArray, # noqa: F821 + vector.TwoVectorArray, # noqa: F821 + vector.ThreeVectorArray, # noqa: F821 + ElectronArray, # noqa: F821 + vector.LorentzVectorArray, # noqa: F821 ) -ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 -ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 - @awkward.mixin_class(behavior) class Muon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Muon") -awkward.behavior.update(awkward._util.copy_behaviors(MasslessParticle, Muon, behavior)) - -MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -MuonArray.ProjectionClass4D = MuonArray # noqa: F821 -MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors(MasslessParticle, Muon, behavior)) + +register_projection_classes( + MuonArray, # noqa: F821 + vector.TwoVectorArray, # noqa: F821 + vector.ThreeVectorArray, # noqa: F821 + MuonArray, # noqa: F821 + vector.LorentzVectorArray, # noqa: F821 +) @awkward.mixin_class(behavior) @@ -254,12 +265,15 @@ class Jet(Particle, base.NanoCollection): ... _set_repr_name("Jet") -awkward.behavior.update(awkward._util.copy_behaviors(Particle, Jet, behavior)) - -JetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -JetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -JetArray.ProjectionClass4D = JetArray # noqa: F821 -JetArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors(Particle, Jet, behavior)) + +register_projection_classes( + JetArray, # noqa: F821 + vector.TwoVectorArray, # noqa: F821 + vector.ThreeVectorArray, # noqa: F821 + JetArray, # noqa: F821 + vector.LorentzVectorArray, # noqa: F821 +) @awkward.mixin_class(behavior) @@ -267,12 +281,15 @@ class Track(Particle, base.NanoCollection): ... _set_repr_name("Track") -awkward.behavior.update(awkward._util.copy_behaviors(Particle, Track, behavior)) - -TrackArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -TrackArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -TrackArray.ProjectionClass4D = TrackArray # noqa: F821 -TrackArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors(Particle, Track, behavior)) + +register_projection_classes( + TrackArray, # noqa: F821 + vector.TwoVectorArray, # noqa: F821 + vector.ThreeVectorArray, # noqa: F821 + TrackArray, # noqa: F821 + vector.LorentzVectorArray, # noqa: F821 +) @awkward.mixin_class(behavior) @@ -283,12 +300,15 @@ def pt(self): _set_repr_name("Tower") -awkward.behavior.update(awkward._util.copy_behaviors(MasslessParticle, Tower, behavior)) - -TowerArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -TowerArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -TowerArray.ProjectionClass4D = TowerArray # noqa: F821 -TowerArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors(MasslessParticle, Tower, behavior)) + +register_projection_classes( + TowerArray, # noqa: F821 + vector.TwoVectorArray, # noqa: F821 + vector.ThreeVectorArray, # noqa: F821 + TowerArray, # noqa: F821 + vector.LorentzVectorArray, # noqa: F821 +) __all__ = [ diff --git a/src/coffea/nanoevents/methods/nanoaod.py b/src/coffea/nanoevents/methods/nanoaod.py index f9c699eb4..bea4f9a54 100644 --- a/src/coffea/nanoevents/methods/nanoaod.py +++ b/src/coffea/nanoevents/methods/nanoaod.py @@ -6,6 +6,7 @@ from dask_awkward import dask_property from coffea.nanoevents.methods import base, candidate, vector +from coffea.util import register_projection_classes behavior = {} behavior.update(base.behavior) @@ -38,16 +39,19 @@ class PtEtaPhiMCollection(vector.PtEtaPhiMLorentzVector, base.NanoCollection): pass -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors( vector.PtEtaPhiMLorentzVector, PtEtaPhiMCollection, behavior ) ) -PtEtaPhiMCollectionArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -PtEtaPhiMCollectionArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -PtEtaPhiMCollectionArray.ProjectionClass4D = PtEtaPhiMCollectionArray # noqa: F821 -PtEtaPhiMCollectionArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +register_projection_classes( + PtEtaPhiMCollectionArray, # noqa: F821 + vector.TwoVectorArray, # noqa: F821 + vector.ThreeVectorArray, # noqa: F821 + PtEtaPhiMCollectionArray, # noqa: F821 + vector.LorentzVectorArray, # noqa: F821 +) @awkward.mixin_class(behavior) @@ -156,14 +160,17 @@ def distinctChildrenDeep(self, dask_array): _set_repr_name("GenParticle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, GenParticle, behavior) ) -GenParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -GenParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -GenParticleArray.ProjectionClass4D = GenParticleArray # noqa: F821 -GenParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +register_projection_classes( + GenParticleArray, # noqa: F821 + vector.TwoVectorArray, # noqa: F821 + vector.ThreeVectorArray, # noqa: F821 + GenParticleArray, # noqa: F821 + vector.LorentzVectorArray, # noqa: F821 +) @awkward.mixin_class(behavior) @@ -184,14 +191,17 @@ def parent(self, dask_array): _set_repr_name("GenVisTau") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, GenVisTau, behavior) ) -GenVisTauArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -GenVisTauArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -GenVisTauArray.ProjectionClass4D = GenVisTauArray # noqa: F821 -GenVisTauArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +register_projection_classes( + GenVisTauArray, # noqa: F821 + vector.TwoVectorArray, # noqa: F821 + vector.ThreeVectorArray, # noqa: F821 + GenVisTauArray, # noqa: F821 + vector.LorentzVectorArray, # noqa: F821 +) @awkward.mixin_class(behavior) @@ -256,14 +266,17 @@ def matched_photon(self, dask_array): _set_repr_name("Electron") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Electron, behavior) ) -ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 -ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +register_projection_classes( + ElectronArray, # noqa: F821 + vector.TwoVectorArray, # noqa: F821 + vector.ThreeVectorArray, # noqa: F821 + ElectronArray, # noqa: F821 + vector.LorentzVectorArray, # noqa: F821 +) @awkward.mixin_class(behavior) @@ -298,14 +311,17 @@ def matched_jet(self, dask_array): _set_repr_name("Muon") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Muon, behavior) ) -MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -MuonArray.ProjectionClass4D = MuonArray # noqa: F821 -MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +register_projection_classes( + MuonArray, # noqa: F821 + vector.TwoVectorArray, # noqa: F821 + vector.ThreeVectorArray, # noqa: F821 + MuonArray, # noqa: F821 + vector.LorentzVectorArray, # noqa: F821 +) @awkward.mixin_class(behavior) @@ -330,14 +346,17 @@ def matched_jet(self, dask_array): _set_repr_name("Tau") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Tau, behavior) ) -TauArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -TauArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -TauArray.ProjectionClass4D = TauArray # noqa: F821 -TauArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +register_projection_classes( + TauArray, # noqa: F821 + vector.TwoVectorArray, # noqa: F821 + vector.ThreeVectorArray, # noqa: F821 + TauArray, # noqa: F821 + vector.LorentzVectorArray, # noqa: F821 +) @awkward.mixin_class(behavior) @@ -413,14 +432,17 @@ def matched_jet(self, dask_array): _set_repr_name("Photon") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Photon, behavior) ) -PhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -PhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -PhotonArray.ProjectionClass4D = PhotonArray # noqa: F821 -PhotonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +register_projection_classes( + PhotonArray, # noqa: F821 + vector.TwoVectorArray, # noqa: F821 + vector.ThreeVectorArray, # noqa: F821 + PhotonArray, # noqa: F821 + vector.LorentzVectorArray, # noqa: F821 +) @awkward.mixin_class(behavior) @@ -437,14 +459,17 @@ def matched_muon(self, dask_array): _set_repr_name("FsrPhoton") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, FsrPhoton, behavior) ) -FsrPhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -FsrPhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -FsrPhotonArray.ProjectionClass4D = FsrPhotonArray # noqa: F821 -FsrPhotonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +register_projection_classes( + FsrPhotonArray, # noqa: F821 + vector.TwoVectorArray, # noqa: F821 + vector.ThreeVectorArray, # noqa: F821 + FsrPhotonArray, # noqa: F821 + vector.LorentzVectorArray, # noqa: F821 +) @awkward.mixin_class(behavior) @@ -515,14 +540,17 @@ def constituents(self, dask_array): _set_repr_name("Jet") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Jet, behavior) ) -JetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -JetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -JetArray.ProjectionClass4D = JetArray # noqa: F821 -JetArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +register_projection_classes( + JetArray, # noqa: F821 + vector.TwoVectorArray, # noqa: F821 + vector.ThreeVectorArray, # noqa: F821 + JetArray, # noqa: F821 + vector.LorentzVectorArray, # noqa: F821 +) @awkward.mixin_class(behavior) @@ -585,14 +613,17 @@ def constituents(self, dask_array): _set_repr_name("FatJet") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, FatJet, behavior) ) -FatJetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -FatJetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -FatJetArray.ProjectionClass4D = FatJetArray # noqa: F821 -FatJetArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +register_projection_classes( + FatJetArray, # noqa: F821 + vector.TwoVectorArray, # noqa: F821 + vector.ThreeVectorArray, # noqa: F821 + FatJetArray, # noqa: F821 + vector.LorentzVectorArray, # noqa: F821 +) @awkward.mixin_class(behavior) @@ -605,14 +636,17 @@ def r(self): _set_repr_name("MissingET") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.PolarTwoVector, MissingET, behavior) ) -MissingETArray.ProjectionClass2D = MissingETArray # noqa: F821 -MissingETArray.ProjectionClass3D = vector.SphericalThreeVectorArray # noqa: F821 -MissingETArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821 -MissingETArray.MomentumClass = MissingETArray # noqa: F821 +register_projection_classes( + MissingETArray, # noqa: F821 + MissingETArray, # noqa: F821 + vector.ThreeVectorArray, # noqa: F821 + vector.LorentzVectorArray, # noqa: F821 + MissingETArray, # noqa: F821 +) @awkward.mixin_class(behavior) diff --git a/src/coffea/nanoevents/methods/pdune.py b/src/coffea/nanoevents/methods/pdune.py index 4c4296008..088855850 100644 --- a/src/coffea/nanoevents/methods/pdune.py +++ b/src/coffea/nanoevents/methods/pdune.py @@ -4,6 +4,7 @@ import numpy from coffea.nanoevents.methods import base, vector +from coffea.util import register_projection_classes behavior = {} behavior.update(base.behavior) @@ -92,15 +93,16 @@ def mass(self): _set_repr_name("Particle") -awkward.behavior.update( - awkward._util.copy_behaviors(vector.LorentzVector, Particle, behavior) +behavior.update(awkward._util.copy_behaviors(vector.LorentzVector, Particle, behavior)) + +register_projection_classes( + ParticleArray, # noqa: F821 + vector.TwoVectorArray, + vector.ThreeVectorArray, + ParticleArray, # noqa: F821 + vector.LorentzVectorArray, ) -ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -ParticleArray.ProjectionClass4D = ParticleArray # noqa: F821 -ParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 - @awkward.mixin_class(behavior) class TrackParticle(vector.LorentzVector, base.NanoCollection): @@ -138,14 +140,17 @@ def t(self): _set_repr_name("TrackParticle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.LorentzVector, TrackParticle, behavior) ) -TrackParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -TrackParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -TrackParticleArray.ProjectionClass4D = TrackParticleArray # noqa: F821 -TrackParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +register_projection_classes( + TrackParticleArray, # noqa: F821 + vector.TwoVectorArray, + vector.ThreeVectorArray, + TrackParticleArray, # noqa: F821 + vector.LorentzVectorArray, +) @awkward.mixin_class(behavior) @@ -165,12 +170,15 @@ def trackParticle(self): _set_repr_name("Muon") -awkward.behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) - -MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -MuonArray.ProjectionClass4D = MuonArray # noqa: F821 -MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) + +register_projection_classes( + MuonArray, # noqa: F821 + vector.TwoVectorArray, + vector.ThreeVectorArray, + MuonArray, # noqa: F821 + vector.LorentzVectorArray, +) @awkward.mixin_class(behavior) @@ -198,12 +206,15 @@ def trackParticle(self): _set_repr_name("Electron") -awkward.behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) - -ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 -ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) + +register_projection_classes( + ElectronArray, # noqa: F821 + vector.TwoVectorArray, + vector.ThreeVectorArray, + ElectronArray, # noqa: F821 + vector.LorentzVectorArray, +) @awkward.mixin_class(behavior) @@ -246,11 +257,14 @@ def parents(self): _set_repr_name("TruthParticle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.LorentzVector, TruthParticle, behavior) ) -TruthParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -TruthParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -TruthParticleArray.ProjectionClass4D = TruthParticleArray # noqa: F821 -TruthParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +register_projection_classes( + TruthParticleArray, # noqa: F821 + vector.TwoVectorArray, + vector.ThreeVectorArray, + TruthParticleArray, # noqa: F821 + vector.LorentzVectorArray, +) diff --git a/src/coffea/nanoevents/methods/physlite.py b/src/coffea/nanoevents/methods/physlite.py index 28554a6c9..d04787208 100644 --- a/src/coffea/nanoevents/methods/physlite.py +++ b/src/coffea/nanoevents/methods/physlite.py @@ -8,6 +8,7 @@ from dask_awkward import dask_property from coffea.nanoevents.methods import base, vector +from coffea.util import register_projection_classes behavior = {} behavior.update(base.behavior) @@ -143,14 +144,17 @@ def mass(self): _set_repr_name("Particle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Particle, behavior) ) -ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -ParticleArray.ProjectionClass4D = ParticleArray # noqa: F821 -ParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +register_projection_classes( + ParticleArray, # noqa: F821 + vector.TwoVectorArray, + vector.ThreeVectorArray, + ParticleArray, # noqa: F821 + vector.LorentzVectorArray, +) @awkward.mixin_class(behavior) @@ -189,14 +193,17 @@ def t(self): _set_repr_name("TrackParticle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.LorentzVector, TrackParticle, behavior) ) -TrackParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -TrackParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -TrackParticleArray.ProjectionClass4D = TrackParticleArray # noqa: F821 -TrackParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +register_projection_classes( + TrackParticleArray, # noqa: F821 + vector.TwoVectorArray, + vector.ThreeVectorArray, + TrackParticleArray, # noqa: F821 + vector.LorentzVectorArray, +) @awkward.mixin_class(behavior) @@ -225,12 +232,15 @@ def trackParticle(self, dask_array): _set_repr_name("Muon") -awkward.behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) - -MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -MuonArray.ProjectionClass4D = MuonArray # noqa: F821 -MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) + +register_projection_classes( + MuonArray, # noqa: F821 + vector.TwoVectorArray, + vector.ThreeVectorArray, + MuonArray, # noqa: F821 + vector.LorentzVectorArray, +) @awkward.mixin_class(behavior) @@ -283,12 +293,15 @@ def caloClusters(self, dask_array): _set_repr_name("Electron") -awkward.behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) - -ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 -ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) + +register_projection_classes( + ElectronArray, # noqa: F821 + vector.TwoVectorArray, + vector.ThreeVectorArray, + ElectronArray, # noqa: F821 + vector.LorentzVectorArray, +) @awkward.mixin_class(behavior) @@ -331,11 +344,14 @@ def parents(self): _set_repr_name("TruthParticle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.LorentzVector, TruthParticle, behavior) ) -TruthParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 -TruthParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 -TruthParticleArray.ProjectionClass4D = TruthParticleArray # noqa: F821 -TruthParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +register_projection_classes( + TruthParticleArray, # noqa: F821 + vector.TwoVectorArray, + vector.ThreeVectorArray, + TruthParticleArray, # noqa: F821 + vector.LorentzVectorArray, +) diff --git a/src/coffea/nanoevents/methods/vector.py b/src/coffea/nanoevents/methods/vector.py index 2944aa72d..79e334024 100644 --- a/src/coffea/nanoevents/methods/vector.py +++ b/src/coffea/nanoevents/methods/vector.py @@ -58,7 +58,7 @@ MomentumAwkward4D, ) -from coffea.util import deprecate +from coffea.util import deprecate, register_projection_classes _cst = pytz.timezone("US/Central") _depttime = _cst.localize(datetime(2024, 6, 30, 11, 59, 59)) @@ -802,40 +802,56 @@ def divide(self, other): behavior[(numpy.subtract, lhs, rhs)] = out_to.subtract -TwoVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 -TwoVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 -TwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -TwoVectorArray.MomentumClass = PolarTwoVectorArray # noqa: F821 - -PolarTwoVectorArray.ProjectionClass2D = PolarTwoVectorArray # noqa: F821 -PolarTwoVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 -PolarTwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -PolarTwoVectorArray.MomentumClass = PolarTwoVectorArray # noqa: F821 - -ThreeVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 -ThreeVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 -ThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -ThreeVectorArray.MomentumClass = SphericalThreeVectorArray # noqa: F821 - -SphericalThreeVectorArray.ProjectionClass2D = PolarTwoVectorArray # noqa: F821 -SphericalThreeVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 -SphericalThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -SphericalThreeVectorArray.MomentumClass = SphericalThreeVectorArray # noqa: F821 - -LorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 -LorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 -LorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -LorentzVectorArray.MomentumClass = LorentzVectorArray # noqa: F821 - -PtEtaPhiMLorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 -PtEtaPhiMLorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 -PtEtaPhiMLorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -PtEtaPhiMLorentzVectorArray.MomentumClass = LorentzVectorArray # noqa: F821 - -PtEtaPhiELorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 -PtEtaPhiELorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 -PtEtaPhiELorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 -PtEtaPhiELorentzVectorArray.MomentumClass = LorentzVectorArray # noqa: F821 +register_projection_classes( + TwoVectorArray, # noqa: F821 + TwoVectorArray, # noqa: F821 + ThreeVectorArray, # noqa: F821 + LorentzVectorArray, # noqa: F821 + PolarTwoVectorArray, # noqa: F821 +) +register_projection_classes( + PolarTwoVectorArray, # noqa: F821 + PolarTwoVectorArray, # noqa: F821 + SphericalThreeVectorArray, # noqa: F821 + LorentzVectorArray, # noqa: F821 + PolarTwoVectorArray, # noqa: F821 +) +register_projection_classes( + ThreeVectorArray, # noqa: F821 + TwoVectorArray, # noqa: F821 + ThreeVectorArray, # noqa: F821 + LorentzVectorArray, # noqa: F821 + SphericalThreeVectorArray, # noqa: F821 +) +register_projection_classes( + SphericalThreeVectorArray, # noqa: F821 + PolarTwoVectorArray, # noqa: F821 + SphericalThreeVectorArray, # noqa: F821 + LorentzVectorArray, # noqa: F821 + SphericalThreeVectorArray, # noqa: F821 +) +register_projection_classes( + LorentzVectorArray, # noqa: F821 + TwoVectorArray, # noqa: F821 + ThreeVectorArray, # noqa: F821 + LorentzVectorArray, # noqa: F821 + LorentzVectorArray, # noqa: F821 +) +register_projection_classes( + PtEtaPhiMLorentzVectorArray, # noqa: F821 + TwoVectorArray, # noqa: F821 + ThreeVectorArray, # noqa: F821 + LorentzVectorArray, # noqa: F821 + LorentzVectorArray, # noqa: F821 +) +register_projection_classes( + PtEtaPhiELorentzVectorArray, # noqa: F821 + TwoVectorArray, # noqa: F821 + ThreeVectorArray, # noqa: F821 + LorentzVectorArray, # noqa: F821 + LorentzVectorArray, # noqa: F821 +) + __all__ = [ "TwoVector", diff --git a/src/coffea/util.py b/src/coffea/util.py index bfb0b5119..662d28638 100644 --- a/src/coffea/util.py +++ b/src/coffea/util.py @@ -246,3 +246,10 @@ def _remove_not_interpretable(branch, emit_warning=True): return False else: return True + + +def register_projection_classes(vec_class, proj_2d, proj_3d, proj_4d, momentum_class): + vec_class.ProjectionClass2D = proj_2d + vec_class.ProjectionClass3D = proj_3d + vec_class.ProjectionClass4D = proj_4d + vec_class.MomentumClass = momentum_class diff --git a/tests/test_nanoevents_vector.py b/tests/test_nanoevents_vector.py index 6aa13d164..665970ab1 100644 --- a/tests/test_nanoevents_vector.py +++ b/tests/test_nanoevents_vector.py @@ -543,6 +543,13 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): assert_record_arrays_equal(a.like(c) + c, c + a.like(c), check_type=True) assert_record_arrays_equal(b.like(c) + c, c + b.like(c), check_type=True) + with pytest.raises(TypeError): + a + b == b + a + with pytest.raises(TypeError): + a + c == c + a + with pytest.raises(TypeError): + b + c == c + b + assert_allclose(a.delta_phi(b), -b.delta_phi(a)) assert_allclose(a.delta_phi(c), -c.delta_phi(a)) assert_allclose(b.delta_phi(c), -c.delta_phi(b)) @@ -551,6 +558,13 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): assert_record_arrays_equal((a.like(c) - c), -(c - a.like(c)), check_type=True) assert_record_arrays_equal((b.like(c) - c), -(c - b.like(c)), check_type=True) + with pytest.raises(TypeError): + a - b == -(b - a) + with pytest.raises(TypeError): + a - c == -(c - a) + with pytest.raises(TypeError): + b - c == -(c - b) + @pytest.mark.parametrize("optimization_enabled", [True, False]) def test_dask_metric_table_and_nearest(optimization_enabled): From d4b4d6cc68ce80cc8642ca47cf7e9f78e616e2d1 Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Fri, 19 Jul 2024 11:32:24 +0200 Subject: [PATCH 38/46] Revert "fix: updates from review" --- src/coffea/nanoevents/methods/candidate.py | 39 +++--- src/coffea/nanoevents/methods/delphes.py | 134 ++++++++----------- src/coffea/nanoevents/methods/nanoaod.py | 144 ++++++++------------- src/coffea/nanoevents/methods/pdune.py | 72 +++++------ src/coffea/nanoevents/methods/physlite.py | 70 ++++------ src/coffea/nanoevents/methods/vector.py | 86 +++++------- src/coffea/util.py | 7 - tests/test_nanoevents_vector.py | 14 -- 8 files changed, 219 insertions(+), 347 deletions(-) diff --git a/src/coffea/nanoevents/methods/candidate.py b/src/coffea/nanoevents/methods/candidate.py index 43eca9ea7..e3c427276 100644 --- a/src/coffea/nanoevents/methods/candidate.py +++ b/src/coffea/nanoevents/methods/candidate.py @@ -9,7 +9,6 @@ import numpy from coffea.nanoevents.methods import vector -from coffea.util import register_projection_classes behavior = dict(vector.behavior) @@ -71,29 +70,23 @@ class PtEtaPhiECandidate(Candidate, vector.PtEtaPhiELorentzVector): pass -behavior.update(awkward._util.copy_behaviors(vector.LorentzVector, Candidate, behavior)) - -register_projection_classes( - CandidateArray, # noqa: F821 - vector.PolarTwoVectorArray, - vector.SphericalThreeVectorArray, - vector.LorentzVectorArray, - CandidateArray, # noqa: F821 -) -register_projection_classes( - PtEtaPhiMCandidateArray, # noqa: F821 - vector.PolarTwoVectorArray, - vector.SphericalThreeVectorArray, - vector.LorentzVectorArray, - PtEtaPhiMCandidateArray, # noqa: F821 -) -register_projection_classes( - PtEtaPhiECandidateArray, # noqa: F821 - vector.PolarTwoVectorArray, - vector.SphericalThreeVectorArray, - vector.LorentzVectorArray, - PtEtaPhiECandidateArray, # noqa: F821 +awkward.behavior.update( + awkward._util.copy_behaviors(vector.LorentzVector, Candidate, behavior) ) +CandidateArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +CandidateArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +CandidateArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821 +CandidateArray.MomentumClass = CandidateArray # noqa: F821 + +PtEtaPhiMCandidateArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +PtEtaPhiMCandidateArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +PtEtaPhiMCandidateArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821 +PtEtaPhiMCandidateArray.MomentumClass = PtEtaPhiMCandidateArray # noqa: F821 + +PtEtaPhiECandidateArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +PtEtaPhiECandidateArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +PtEtaPhiECandidateArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821 +PtEtaPhiECandidateArray.MomentumClass = PtEtaPhiECandidateArray # noqa: F821 __all__ = ["Candidate", "PtEtaPhiMCandidate", "PtEtaPhiECandidate"] diff --git a/src/coffea/nanoevents/methods/delphes.py b/src/coffea/nanoevents/methods/delphes.py index cc65121f6..c232fedbe 100644 --- a/src/coffea/nanoevents/methods/delphes.py +++ b/src/coffea/nanoevents/methods/delphes.py @@ -7,7 +7,6 @@ import numpy from coffea.nanoevents.methods import base, candidate, vector -from coffea.util import register_projection_classes behavior = {} behavior.update(base.behavior) @@ -131,16 +130,15 @@ def z(self): _set_repr_name("Vertex") -behavior.update(awkward._util.copy_behaviors(vector.LorentzVector, Vertex, behavior)) - -register_projection_classes( - VertexArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - VertexArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 +awkward.behavior.update( + awkward._util.copy_behaviors(vector.LorentzVector, Vertex, behavior) ) +VertexArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +VertexArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +VertexArray.ProjectionClass4D = VertexArray # noqa: F821 +VertexArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + @awkward.mixin_class(behavior) class Particle(vector.PtEtaPhiMLorentzVector): @@ -180,17 +178,14 @@ def mass(self): _set_repr_name("Particle") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Particle, behavior) ) -register_projection_classes( - ParticleArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - ParticleArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ParticleArray.ProjectionClass4D = ParticleArray # noqa: F821 +ParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -201,63 +196,57 @@ def mass(self): _set_repr_name("MasslessParticle") -behavior.update(awkward._util.copy_behaviors(Particle, MasslessParticle, behavior)) - -register_projection_classes( - MasslessParticleArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - MasslessParticleArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 +awkward.behavior.update( + awkward._util.copy_behaviors(Particle, MasslessParticle, behavior) ) +MasslessParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +MasslessParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +MasslessParticleArray.ProjectionClass4D = MasslessParticleArray # noqa: F821 +MasslessParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + @awkward.mixin_class(behavior) class Photon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Photon") -behavior.update(awkward._util.copy_behaviors(MasslessParticle, Photon, behavior)) - -register_projection_classes( - PhotonArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - PhotonArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 +awkward.behavior.update( + awkward._util.copy_behaviors(MasslessParticle, Photon, behavior) ) +PhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +PhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +PhotonArray.ProjectionClass4D = PhotonArray # noqa: F821 +PhotonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + @awkward.mixin_class(behavior) class Electron(MasslessParticle, base.NanoCollection): ... _set_repr_name("Electron") -behavior.update(awkward._util.copy_behaviors(MasslessParticle, Electron, behavior)) - -register_projection_classes( - ElectronArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - ElectronArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 +awkward.behavior.update( + awkward._util.copy_behaviors(MasslessParticle, Electron, behavior) ) +ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 +ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + @awkward.mixin_class(behavior) class Muon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Muon") -behavior.update(awkward._util.copy_behaviors(MasslessParticle, Muon, behavior)) - -register_projection_classes( - MuonArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - MuonArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +awkward.behavior.update(awkward._util.copy_behaviors(MasslessParticle, Muon, behavior)) + +MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +MuonArray.ProjectionClass4D = MuonArray # noqa: F821 +MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -265,15 +254,12 @@ class Jet(Particle, base.NanoCollection): ... _set_repr_name("Jet") -behavior.update(awkward._util.copy_behaviors(Particle, Jet, behavior)) - -register_projection_classes( - JetArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - JetArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +awkward.behavior.update(awkward._util.copy_behaviors(Particle, Jet, behavior)) + +JetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +JetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +JetArray.ProjectionClass4D = JetArray # noqa: F821 +JetArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -281,15 +267,12 @@ class Track(Particle, base.NanoCollection): ... _set_repr_name("Track") -behavior.update(awkward._util.copy_behaviors(Particle, Track, behavior)) - -register_projection_classes( - TrackArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - TrackArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +awkward.behavior.update(awkward._util.copy_behaviors(Particle, Track, behavior)) + +TrackArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TrackArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TrackArray.ProjectionClass4D = TrackArray # noqa: F821 +TrackArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -300,15 +283,12 @@ def pt(self): _set_repr_name("Tower") -behavior.update(awkward._util.copy_behaviors(MasslessParticle, Tower, behavior)) - -register_projection_classes( - TowerArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - TowerArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +awkward.behavior.update(awkward._util.copy_behaviors(MasslessParticle, Tower, behavior)) + +TowerArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TowerArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TowerArray.ProjectionClass4D = TowerArray # noqa: F821 +TowerArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 __all__ = [ diff --git a/src/coffea/nanoevents/methods/nanoaod.py b/src/coffea/nanoevents/methods/nanoaod.py index 141d329d0..8fa587410 100644 --- a/src/coffea/nanoevents/methods/nanoaod.py +++ b/src/coffea/nanoevents/methods/nanoaod.py @@ -6,7 +6,6 @@ from dask_awkward import dask_property from coffea.nanoevents.methods import base, candidate, vector -from coffea.util import register_projection_classes behavior = {} behavior.update(base.behavior) @@ -39,19 +38,16 @@ class PtEtaPhiMCollection(vector.PtEtaPhiMLorentzVector, base.NanoCollection): pass -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors( vector.PtEtaPhiMLorentzVector, PtEtaPhiMCollection, behavior ) ) -register_projection_classes( - PtEtaPhiMCollectionArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - PtEtaPhiMCollectionArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +PtEtaPhiMCollectionArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +PtEtaPhiMCollectionArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +PtEtaPhiMCollectionArray.ProjectionClass4D = PtEtaPhiMCollectionArray # noqa: F821 +PtEtaPhiMCollectionArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -160,17 +156,14 @@ def distinctChildrenDeep(self, dask_array): _set_repr_name("GenParticle") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, GenParticle, behavior) ) -register_projection_classes( - GenParticleArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - GenParticleArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +GenParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +GenParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +GenParticleArray.ProjectionClass4D = GenParticleArray # noqa: F821 +GenParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -191,17 +184,14 @@ def parent(self, dask_array): _set_repr_name("GenVisTau") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, GenVisTau, behavior) ) -register_projection_classes( - GenVisTauArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - GenVisTauArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +GenVisTauArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +GenVisTauArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +GenVisTauArray.ProjectionClass4D = GenVisTauArray # noqa: F821 +GenVisTauArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -266,17 +256,14 @@ def matched_photon(self, dask_array): _set_repr_name("Electron") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Electron, behavior) ) -register_projection_classes( - ElectronArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - ElectronArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 +ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -353,17 +340,14 @@ def matched_jet(self, dask_array): _set_repr_name("Muon") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Muon, behavior) ) -register_projection_classes( - MuonArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - MuonArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +MuonArray.ProjectionClass4D = MuonArray # noqa: F821 +MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -388,17 +372,14 @@ def matched_jet(self, dask_array): _set_repr_name("Tau") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Tau, behavior) ) -register_projection_classes( - TauArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - TauArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +TauArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TauArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TauArray.ProjectionClass4D = TauArray # noqa: F821 +TauArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -474,17 +455,14 @@ def matched_jet(self, dask_array): _set_repr_name("Photon") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Photon, behavior) ) -register_projection_classes( - PhotonArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - PhotonArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +PhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +PhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +PhotonArray.ProjectionClass4D = PhotonArray # noqa: F821 +PhotonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -501,17 +479,14 @@ def matched_muon(self, dask_array): _set_repr_name("FsrPhoton") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, FsrPhoton, behavior) ) -register_projection_classes( - FsrPhotonArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - FsrPhotonArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +FsrPhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +FsrPhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +FsrPhotonArray.ProjectionClass4D = FsrPhotonArray # noqa: F821 +FsrPhotonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -582,17 +557,14 @@ def constituents(self, dask_array): _set_repr_name("Jet") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Jet, behavior) ) -register_projection_classes( - JetArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - JetArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +JetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +JetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +JetArray.ProjectionClass4D = JetArray # noqa: F821 +JetArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -655,17 +627,14 @@ def constituents(self, dask_array): _set_repr_name("FatJet") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, FatJet, behavior) ) -register_projection_classes( - FatJetArray, # noqa: F821 - vector.TwoVectorArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - FatJetArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 -) +FatJetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +FatJetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +FatJetArray.ProjectionClass4D = FatJetArray # noqa: F821 +FatJetArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -678,17 +647,14 @@ def r(self): _set_repr_name("MissingET") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.PolarTwoVector, MissingET, behavior) ) -register_projection_classes( - MissingETArray, # noqa: F821 - MissingETArray, # noqa: F821 - vector.ThreeVectorArray, # noqa: F821 - vector.LorentzVectorArray, # noqa: F821 - MissingETArray, # noqa: F821 -) +MissingETArray.ProjectionClass2D = MissingETArray # noqa: F821 +MissingETArray.ProjectionClass3D = vector.SphericalThreeVectorArray # noqa: F821 +MissingETArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821 +MissingETArray.MomentumClass = MissingETArray # noqa: F821 @awkward.mixin_class(behavior) diff --git a/src/coffea/nanoevents/methods/pdune.py b/src/coffea/nanoevents/methods/pdune.py index 088855850..4c4296008 100644 --- a/src/coffea/nanoevents/methods/pdune.py +++ b/src/coffea/nanoevents/methods/pdune.py @@ -4,7 +4,6 @@ import numpy from coffea.nanoevents.methods import base, vector -from coffea.util import register_projection_classes behavior = {} behavior.update(base.behavior) @@ -93,16 +92,15 @@ def mass(self): _set_repr_name("Particle") -behavior.update(awkward._util.copy_behaviors(vector.LorentzVector, Particle, behavior)) - -register_projection_classes( - ParticleArray, # noqa: F821 - vector.TwoVectorArray, - vector.ThreeVectorArray, - ParticleArray, # noqa: F821 - vector.LorentzVectorArray, +awkward.behavior.update( + awkward._util.copy_behaviors(vector.LorentzVector, Particle, behavior) ) +ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ParticleArray.ProjectionClass4D = ParticleArray # noqa: F821 +ParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 + @awkward.mixin_class(behavior) class TrackParticle(vector.LorentzVector, base.NanoCollection): @@ -140,17 +138,14 @@ def t(self): _set_repr_name("TrackParticle") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.LorentzVector, TrackParticle, behavior) ) -register_projection_classes( - TrackParticleArray, # noqa: F821 - vector.TwoVectorArray, - vector.ThreeVectorArray, - TrackParticleArray, # noqa: F821 - vector.LorentzVectorArray, -) +TrackParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TrackParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TrackParticleArray.ProjectionClass4D = TrackParticleArray # noqa: F821 +TrackParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -170,15 +165,12 @@ def trackParticle(self): _set_repr_name("Muon") -behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) - -register_projection_classes( - MuonArray, # noqa: F821 - vector.TwoVectorArray, - vector.ThreeVectorArray, - MuonArray, # noqa: F821 - vector.LorentzVectorArray, -) +awkward.behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) + +MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +MuonArray.ProjectionClass4D = MuonArray # noqa: F821 +MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -206,15 +198,12 @@ def trackParticle(self): _set_repr_name("Electron") -behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) - -register_projection_classes( - ElectronArray, # noqa: F821 - vector.TwoVectorArray, - vector.ThreeVectorArray, - ElectronArray, # noqa: F821 - vector.LorentzVectorArray, -) +awkward.behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) + +ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 +ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -257,14 +246,11 @@ def parents(self): _set_repr_name("TruthParticle") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.LorentzVector, TruthParticle, behavior) ) -register_projection_classes( - TruthParticleArray, # noqa: F821 - vector.TwoVectorArray, - vector.ThreeVectorArray, - TruthParticleArray, # noqa: F821 - vector.LorentzVectorArray, -) +TruthParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TruthParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TruthParticleArray.ProjectionClass4D = TruthParticleArray # noqa: F821 +TruthParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 diff --git a/src/coffea/nanoevents/methods/physlite.py b/src/coffea/nanoevents/methods/physlite.py index d04787208..28554a6c9 100644 --- a/src/coffea/nanoevents/methods/physlite.py +++ b/src/coffea/nanoevents/methods/physlite.py @@ -8,7 +8,6 @@ from dask_awkward import dask_property from coffea.nanoevents.methods import base, vector -from coffea.util import register_projection_classes behavior = {} behavior.update(base.behavior) @@ -144,17 +143,14 @@ def mass(self): _set_repr_name("Particle") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Particle, behavior) ) -register_projection_classes( - ParticleArray, # noqa: F821 - vector.TwoVectorArray, - vector.ThreeVectorArray, - ParticleArray, # noqa: F821 - vector.LorentzVectorArray, -) +ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ParticleArray.ProjectionClass4D = ParticleArray # noqa: F821 +ParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -193,17 +189,14 @@ def t(self): _set_repr_name("TrackParticle") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.LorentzVector, TrackParticle, behavior) ) -register_projection_classes( - TrackParticleArray, # noqa: F821 - vector.TwoVectorArray, - vector.ThreeVectorArray, - TrackParticleArray, # noqa: F821 - vector.LorentzVectorArray, -) +TrackParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TrackParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TrackParticleArray.ProjectionClass4D = TrackParticleArray # noqa: F821 +TrackParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -232,15 +225,12 @@ def trackParticle(self, dask_array): _set_repr_name("Muon") -behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) - -register_projection_classes( - MuonArray, # noqa: F821 - vector.TwoVectorArray, - vector.ThreeVectorArray, - MuonArray, # noqa: F821 - vector.LorentzVectorArray, -) +awkward.behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) + +MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +MuonArray.ProjectionClass4D = MuonArray # noqa: F821 +MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -293,15 +283,12 @@ def caloClusters(self, dask_array): _set_repr_name("Electron") -behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) - -register_projection_classes( - ElectronArray, # noqa: F821 - vector.TwoVectorArray, - vector.ThreeVectorArray, - ElectronArray, # noqa: F821 - vector.LorentzVectorArray, -) +awkward.behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) + +ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 +ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 @awkward.mixin_class(behavior) @@ -344,14 +331,11 @@ def parents(self): _set_repr_name("TruthParticle") -behavior.update( +awkward.behavior.update( awkward._util.copy_behaviors(vector.LorentzVector, TruthParticle, behavior) ) -register_projection_classes( - TruthParticleArray, # noqa: F821 - vector.TwoVectorArray, - vector.ThreeVectorArray, - TruthParticleArray, # noqa: F821 - vector.LorentzVectorArray, -) +TruthParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 +TruthParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 +TruthParticleArray.ProjectionClass4D = TruthParticleArray # noqa: F821 +TruthParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 diff --git a/src/coffea/nanoevents/methods/vector.py b/src/coffea/nanoevents/methods/vector.py index 79e334024..2944aa72d 100644 --- a/src/coffea/nanoevents/methods/vector.py +++ b/src/coffea/nanoevents/methods/vector.py @@ -58,7 +58,7 @@ MomentumAwkward4D, ) -from coffea.util import deprecate, register_projection_classes +from coffea.util import deprecate _cst = pytz.timezone("US/Central") _depttime = _cst.localize(datetime(2024, 6, 30, 11, 59, 59)) @@ -802,56 +802,40 @@ def divide(self, other): behavior[(numpy.subtract, lhs, rhs)] = out_to.subtract -register_projection_classes( - TwoVectorArray, # noqa: F821 - TwoVectorArray, # noqa: F821 - ThreeVectorArray, # noqa: F821 - LorentzVectorArray, # noqa: F821 - PolarTwoVectorArray, # noqa: F821 -) -register_projection_classes( - PolarTwoVectorArray, # noqa: F821 - PolarTwoVectorArray, # noqa: F821 - SphericalThreeVectorArray, # noqa: F821 - LorentzVectorArray, # noqa: F821 - PolarTwoVectorArray, # noqa: F821 -) -register_projection_classes( - ThreeVectorArray, # noqa: F821 - TwoVectorArray, # noqa: F821 - ThreeVectorArray, # noqa: F821 - LorentzVectorArray, # noqa: F821 - SphericalThreeVectorArray, # noqa: F821 -) -register_projection_classes( - SphericalThreeVectorArray, # noqa: F821 - PolarTwoVectorArray, # noqa: F821 - SphericalThreeVectorArray, # noqa: F821 - LorentzVectorArray, # noqa: F821 - SphericalThreeVectorArray, # noqa: F821 -) -register_projection_classes( - LorentzVectorArray, # noqa: F821 - TwoVectorArray, # noqa: F821 - ThreeVectorArray, # noqa: F821 - LorentzVectorArray, # noqa: F821 - LorentzVectorArray, # noqa: F821 -) -register_projection_classes( - PtEtaPhiMLorentzVectorArray, # noqa: F821 - TwoVectorArray, # noqa: F821 - ThreeVectorArray, # noqa: F821 - LorentzVectorArray, # noqa: F821 - LorentzVectorArray, # noqa: F821 -) -register_projection_classes( - PtEtaPhiELorentzVectorArray, # noqa: F821 - TwoVectorArray, # noqa: F821 - ThreeVectorArray, # noqa: F821 - LorentzVectorArray, # noqa: F821 - LorentzVectorArray, # noqa: F821 -) - +TwoVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +TwoVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +TwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +TwoVectorArray.MomentumClass = PolarTwoVectorArray # noqa: F821 + +PolarTwoVectorArray.ProjectionClass2D = PolarTwoVectorArray # noqa: F821 +PolarTwoVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 +PolarTwoVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +PolarTwoVectorArray.MomentumClass = PolarTwoVectorArray # noqa: F821 + +ThreeVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +ThreeVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +ThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +ThreeVectorArray.MomentumClass = SphericalThreeVectorArray # noqa: F821 + +SphericalThreeVectorArray.ProjectionClass2D = PolarTwoVectorArray # noqa: F821 +SphericalThreeVectorArray.ProjectionClass3D = SphericalThreeVectorArray # noqa: F821 +SphericalThreeVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +SphericalThreeVectorArray.MomentumClass = SphericalThreeVectorArray # noqa: F821 + +LorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +LorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +LorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +LorentzVectorArray.MomentumClass = LorentzVectorArray # noqa: F821 + +PtEtaPhiMLorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +PtEtaPhiMLorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +PtEtaPhiMLorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +PtEtaPhiMLorentzVectorArray.MomentumClass = LorentzVectorArray # noqa: F821 + +PtEtaPhiELorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 +PtEtaPhiELorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 +PtEtaPhiELorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 +PtEtaPhiELorentzVectorArray.MomentumClass = LorentzVectorArray # noqa: F821 __all__ = [ "TwoVector", diff --git a/src/coffea/util.py b/src/coffea/util.py index 662d28638..bfb0b5119 100644 --- a/src/coffea/util.py +++ b/src/coffea/util.py @@ -246,10 +246,3 @@ def _remove_not_interpretable(branch, emit_warning=True): return False else: return True - - -def register_projection_classes(vec_class, proj_2d, proj_3d, proj_4d, momentum_class): - vec_class.ProjectionClass2D = proj_2d - vec_class.ProjectionClass3D = proj_3d - vec_class.ProjectionClass4D = proj_4d - vec_class.MomentumClass = momentum_class diff --git a/tests/test_nanoevents_vector.py b/tests/test_nanoevents_vector.py index 665970ab1..6aa13d164 100644 --- a/tests/test_nanoevents_vector.py +++ b/tests/test_nanoevents_vector.py @@ -543,13 +543,6 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): assert_record_arrays_equal(a.like(c) + c, c + a.like(c), check_type=True) assert_record_arrays_equal(b.like(c) + c, c + b.like(c), check_type=True) - with pytest.raises(TypeError): - a + b == b + a - with pytest.raises(TypeError): - a + c == c + a - with pytest.raises(TypeError): - b + c == c + b - assert_allclose(a.delta_phi(b), -b.delta_phi(a)) assert_allclose(a.delta_phi(c), -c.delta_phi(a)) assert_allclose(b.delta_phi(c), -c.delta_phi(b)) @@ -558,13 +551,6 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): assert_record_arrays_equal((a.like(c) - c), -(c - a.like(c)), check_type=True) assert_record_arrays_equal((b.like(c) - c), -(c - b.like(c)), check_type=True) - with pytest.raises(TypeError): - a - b == -(b - a) - with pytest.raises(TypeError): - a - c == -(c - a) - with pytest.raises(TypeError): - b - c == -(c - b) - @pytest.mark.parametrize("optimization_enabled", [True, False]) def test_dask_metric_table_and_nearest(optimization_enabled): From dc80e611278ea30b77fd64d675aaf25c7a553375 Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Fri, 19 Jul 2024 11:44:18 +0200 Subject: [PATCH 39/46] Add back useful changes * More tests * Don't modify the global awkward behavior --- src/coffea/nanoevents/methods/candidate.py | 4 +--- src/coffea/nanoevents/methods/delphes.py | 20 +++++++++----------- src/coffea/nanoevents/methods/nanoaod.py | 22 +++++++++++----------- src/coffea/nanoevents/methods/pdune.py | 12 +++++------- src/coffea/nanoevents/methods/physlite.py | 10 +++++----- tests/test_nanoevents_vector.py | 14 ++++++++++++++ 6 files changed, 45 insertions(+), 37 deletions(-) diff --git a/src/coffea/nanoevents/methods/candidate.py b/src/coffea/nanoevents/methods/candidate.py index e3c427276..f4f417b63 100644 --- a/src/coffea/nanoevents/methods/candidate.py +++ b/src/coffea/nanoevents/methods/candidate.py @@ -70,9 +70,7 @@ class PtEtaPhiECandidate(Candidate, vector.PtEtaPhiELorentzVector): pass -awkward.behavior.update( - awkward._util.copy_behaviors(vector.LorentzVector, Candidate, behavior) -) +behavior.update(awkward._util.copy_behaviors(vector.LorentzVector, Candidate, behavior)) CandidateArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 CandidateArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 diff --git a/src/coffea/nanoevents/methods/delphes.py b/src/coffea/nanoevents/methods/delphes.py index c232fedbe..cc734824b 100644 --- a/src/coffea/nanoevents/methods/delphes.py +++ b/src/coffea/nanoevents/methods/delphes.py @@ -130,9 +130,7 @@ def z(self): _set_repr_name("Vertex") -awkward.behavior.update( - awkward._util.copy_behaviors(vector.LorentzVector, Vertex, behavior) -) +behavior.update(awkward._util.copy_behaviors(vector.LorentzVector, Vertex, behavior)) VertexArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 VertexArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -178,7 +176,7 @@ def mass(self): _set_repr_name("Particle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Particle, behavior) ) @@ -196,7 +194,7 @@ def mass(self): _set_repr_name("MasslessParticle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(Particle, MasslessParticle, behavior) ) @@ -211,7 +209,7 @@ class Photon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Photon") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(MasslessParticle, Photon, behavior) ) @@ -226,7 +224,7 @@ class Electron(MasslessParticle, base.NanoCollection): ... _set_repr_name("Electron") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(MasslessParticle, Electron, behavior) ) @@ -241,7 +239,7 @@ class Muon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Muon") -awkward.behavior.update(awkward._util.copy_behaviors(MasslessParticle, Muon, behavior)) +behavior.update(awkward._util.copy_behaviors(MasslessParticle, Muon, behavior)) MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -254,7 +252,7 @@ class Jet(Particle, base.NanoCollection): ... _set_repr_name("Jet") -awkward.behavior.update(awkward._util.copy_behaviors(Particle, Jet, behavior)) +behavior.update(awkward._util.copy_behaviors(Particle, Jet, behavior)) JetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 JetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -267,7 +265,7 @@ class Track(Particle, base.NanoCollection): ... _set_repr_name("Track") -awkward.behavior.update(awkward._util.copy_behaviors(Particle, Track, behavior)) +behavior.update(awkward._util.copy_behaviors(Particle, Track, behavior)) TrackArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 TrackArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -283,7 +281,7 @@ def pt(self): _set_repr_name("Tower") -awkward.behavior.update(awkward._util.copy_behaviors(MasslessParticle, Tower, behavior)) +behavior.update(awkward._util.copy_behaviors(MasslessParticle, Tower, behavior)) TowerArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 TowerArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 diff --git a/src/coffea/nanoevents/methods/nanoaod.py b/src/coffea/nanoevents/methods/nanoaod.py index 8fa587410..e46b3c287 100644 --- a/src/coffea/nanoevents/methods/nanoaod.py +++ b/src/coffea/nanoevents/methods/nanoaod.py @@ -38,7 +38,7 @@ class PtEtaPhiMCollection(vector.PtEtaPhiMLorentzVector, base.NanoCollection): pass -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors( vector.PtEtaPhiMLorentzVector, PtEtaPhiMCollection, behavior ) @@ -156,7 +156,7 @@ def distinctChildrenDeep(self, dask_array): _set_repr_name("GenParticle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, GenParticle, behavior) ) @@ -184,7 +184,7 @@ def parent(self, dask_array): _set_repr_name("GenVisTau") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, GenVisTau, behavior) ) @@ -256,7 +256,7 @@ def matched_photon(self, dask_array): _set_repr_name("Electron") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Electron, behavior) ) @@ -340,7 +340,7 @@ def matched_jet(self, dask_array): _set_repr_name("Muon") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Muon, behavior) ) @@ -372,7 +372,7 @@ def matched_jet(self, dask_array): _set_repr_name("Tau") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Tau, behavior) ) @@ -455,7 +455,7 @@ def matched_jet(self, dask_array): _set_repr_name("Photon") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Photon, behavior) ) @@ -479,7 +479,7 @@ def matched_muon(self, dask_array): _set_repr_name("FsrPhoton") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, FsrPhoton, behavior) ) @@ -557,7 +557,7 @@ def constituents(self, dask_array): _set_repr_name("Jet") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Jet, behavior) ) @@ -627,7 +627,7 @@ def constituents(self, dask_array): _set_repr_name("FatJet") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, FatJet, behavior) ) @@ -647,7 +647,7 @@ def r(self): _set_repr_name("MissingET") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.PolarTwoVector, MissingET, behavior) ) diff --git a/src/coffea/nanoevents/methods/pdune.py b/src/coffea/nanoevents/methods/pdune.py index 4c4296008..43e2649b5 100644 --- a/src/coffea/nanoevents/methods/pdune.py +++ b/src/coffea/nanoevents/methods/pdune.py @@ -92,9 +92,7 @@ def mass(self): _set_repr_name("Particle") -awkward.behavior.update( - awkward._util.copy_behaviors(vector.LorentzVector, Particle, behavior) -) +behavior.update(awkward._util.copy_behaviors(vector.LorentzVector, Particle, behavior)) ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -138,7 +136,7 @@ def t(self): _set_repr_name("TrackParticle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.LorentzVector, TrackParticle, behavior) ) @@ -165,7 +163,7 @@ def trackParticle(self): _set_repr_name("Muon") -awkward.behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) +behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -198,7 +196,7 @@ def trackParticle(self): _set_repr_name("Electron") -awkward.behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) +behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -246,7 +244,7 @@ def parents(self): _set_repr_name("TruthParticle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.LorentzVector, TruthParticle, behavior) ) diff --git a/src/coffea/nanoevents/methods/physlite.py b/src/coffea/nanoevents/methods/physlite.py index 28554a6c9..301ad3b08 100644 --- a/src/coffea/nanoevents/methods/physlite.py +++ b/src/coffea/nanoevents/methods/physlite.py @@ -143,7 +143,7 @@ def mass(self): _set_repr_name("Particle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Particle, behavior) ) @@ -189,7 +189,7 @@ def t(self): _set_repr_name("TrackParticle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.LorentzVector, TrackParticle, behavior) ) @@ -225,7 +225,7 @@ def trackParticle(self, dask_array): _set_repr_name("Muon") -awkward.behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) +behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -283,7 +283,7 @@ def caloClusters(self, dask_array): _set_repr_name("Electron") -awkward.behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) +behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -331,7 +331,7 @@ def parents(self): _set_repr_name("TruthParticle") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors(vector.LorentzVector, TruthParticle, behavior) ) diff --git a/tests/test_nanoevents_vector.py b/tests/test_nanoevents_vector.py index 6aa13d164..46b9cbd1e 100644 --- a/tests/test_nanoevents_vector.py +++ b/tests/test_nanoevents_vector.py @@ -543,6 +543,13 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): assert_record_arrays_equal(a.like(c) + c, c + a.like(c), check_type=True) assert_record_arrays_equal(b.like(c) + c, c + b.like(c), check_type=True) +with pytest.raises(TypeError): + a + b == b + a + with pytest.raises(TypeError): + a + c == c + a + with pytest.raises(TypeError): + b + c == c + b + assert_allclose(a.delta_phi(b), -b.delta_phi(a)) assert_allclose(a.delta_phi(c), -c.delta_phi(a)) assert_allclose(b.delta_phi(c), -c.delta_phi(b)) @@ -551,6 +558,13 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): assert_record_arrays_equal((a.like(c) - c), -(c - a.like(c)), check_type=True) assert_record_arrays_equal((b.like(c) - c), -(c - b.like(c)), check_type=True) + with pytest.raises(TypeError): + a - b == -(b - a) + with pytest.raises(TypeError): + a - c == -(c - a) + with pytest.raises(TypeError): + b - c == -(c - b) + @pytest.mark.parametrize("optimization_enabled", [True, False]) def test_dask_metric_table_and_nearest(optimization_enabled): From 8349641317f21ee0ab0784d55f27e52397a1a865 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 19 Jul 2024 09:45:44 +0000 Subject: [PATCH 40/46] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/coffea/nanoevents/methods/delphes.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/src/coffea/nanoevents/methods/delphes.py b/src/coffea/nanoevents/methods/delphes.py index cc734824b..5a5b0979a 100644 --- a/src/coffea/nanoevents/methods/delphes.py +++ b/src/coffea/nanoevents/methods/delphes.py @@ -194,9 +194,7 @@ def mass(self): _set_repr_name("MasslessParticle") -behavior.update( - awkward._util.copy_behaviors(Particle, MasslessParticle, behavior) -) +behavior.update(awkward._util.copy_behaviors(Particle, MasslessParticle, behavior)) MasslessParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 MasslessParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -209,9 +207,7 @@ class Photon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Photon") -behavior.update( - awkward._util.copy_behaviors(MasslessParticle, Photon, behavior) -) +behavior.update(awkward._util.copy_behaviors(MasslessParticle, Photon, behavior)) PhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 PhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 @@ -224,9 +220,7 @@ class Electron(MasslessParticle, base.NanoCollection): ... _set_repr_name("Electron") -behavior.update( - awkward._util.copy_behaviors(MasslessParticle, Electron, behavior) -) +behavior.update(awkward._util.copy_behaviors(MasslessParticle, Electron, behavior)) ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 From fed625b100e6e02e903e1a9cddef61d6c63a58f3 Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Fri, 19 Jul 2024 11:46:57 +0200 Subject: [PATCH 41/46] Indentation --- tests/test_nanoevents_vector.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/tests/test_nanoevents_vector.py b/tests/test_nanoevents_vector.py index 46b9cbd1e..665970ab1 100644 --- a/tests/test_nanoevents_vector.py +++ b/tests/test_nanoevents_vector.py @@ -543,12 +543,12 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): assert_record_arrays_equal(a.like(c) + c, c + a.like(c), check_type=True) assert_record_arrays_equal(b.like(c) + c, c + b.like(c), check_type=True) -with pytest.raises(TypeError): - a + b == b + a - with pytest.raises(TypeError): - a + c == c + a - with pytest.raises(TypeError): - b + c == c + b + with pytest.raises(TypeError): + a + b == b + a + with pytest.raises(TypeError): + a + c == c + a + with pytest.raises(TypeError): + b + c == c + b assert_allclose(a.delta_phi(b), -b.delta_phi(a)) assert_allclose(a.delta_phi(c), -c.delta_phi(a)) @@ -559,11 +559,11 @@ def test_inherited_method_transpose(lcoord, threecoord, twocoord): assert_record_arrays_equal((b.like(c) - c), -(c - b.like(c)), check_type=True) with pytest.raises(TypeError): - a - b == -(b - a) - with pytest.raises(TypeError): - a - c == -(c - a) - with pytest.raises(TypeError): - b - c == -(c - b) + a - b == -(b - a) + with pytest.raises(TypeError): + a - c == -(c - a) + with pytest.raises(TypeError): + b - c == -(c - b) @pytest.mark.parametrize("optimization_enabled", [True, False]) From 6e48c1771aa0aefffb9e23309ad3f27b0dce6595 Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Mon, 22 Jul 2024 12:07:47 +0200 Subject: [PATCH 42/46] fix: use the new type signature for copy_behaviors --- src/coffea/nanoevents/methods/candidate.py | 2 +- src/coffea/nanoevents/methods/delphes.py | 18 ++++++------ src/coffea/nanoevents/methods/nanoaod.py | 34 ++++++++-------------- src/coffea/nanoevents/methods/pdune.py | 10 +++---- src/coffea/nanoevents/methods/physlite.py | 10 +++---- 5 files changed, 32 insertions(+), 42 deletions(-) diff --git a/src/coffea/nanoevents/methods/candidate.py b/src/coffea/nanoevents/methods/candidate.py index 43eca9ea7..d335349ea 100644 --- a/src/coffea/nanoevents/methods/candidate.py +++ b/src/coffea/nanoevents/methods/candidate.py @@ -71,7 +71,7 @@ class PtEtaPhiECandidate(Candidate, vector.PtEtaPhiELorentzVector): pass -behavior.update(awkward._util.copy_behaviors(vector.LorentzVector, Candidate, behavior)) +behavior.update(awkward._util.copy_behaviors("LorentzVector", "Candidate", behavior)) register_projection_classes( CandidateArray, # noqa: F821 diff --git a/src/coffea/nanoevents/methods/delphes.py b/src/coffea/nanoevents/methods/delphes.py index cc65121f6..69c3015d5 100644 --- a/src/coffea/nanoevents/methods/delphes.py +++ b/src/coffea/nanoevents/methods/delphes.py @@ -131,7 +131,7 @@ def z(self): _set_repr_name("Vertex") -behavior.update(awkward._util.copy_behaviors(vector.LorentzVector, Vertex, behavior)) +behavior.update(awkward._util.copy_behaviors("LorentzVector", "Vertex", behavior)) register_projection_classes( VertexArray, # noqa: F821 @@ -181,7 +181,7 @@ def mass(self): _set_repr_name("Particle") behavior.update( - awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Particle, behavior) + awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "Particle", behavior) ) register_projection_classes( @@ -201,7 +201,7 @@ def mass(self): _set_repr_name("MasslessParticle") -behavior.update(awkward._util.copy_behaviors(Particle, MasslessParticle, behavior)) +behavior.update(awkward._util.copy_behaviors("Particle", "MasslessParticle", behavior)) register_projection_classes( MasslessParticleArray, # noqa: F821 @@ -217,7 +217,7 @@ class Photon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Photon") -behavior.update(awkward._util.copy_behaviors(MasslessParticle, Photon, behavior)) +behavior.update(awkward._util.copy_behaviors("MasslessParticle", "Photon", behavior)) register_projection_classes( PhotonArray, # noqa: F821 @@ -233,7 +233,7 @@ class Electron(MasslessParticle, base.NanoCollection): ... _set_repr_name("Electron") -behavior.update(awkward._util.copy_behaviors(MasslessParticle, Electron, behavior)) +behavior.update(awkward._util.copy_behaviors("MasslessParticle", "Electron", behavior)) register_projection_classes( ElectronArray, # noqa: F821 @@ -249,7 +249,7 @@ class Muon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Muon") -behavior.update(awkward._util.copy_behaviors(MasslessParticle, Muon, behavior)) +behavior.update(awkward._util.copy_behaviors("MasslessParticle", "Muon", behavior)) register_projection_classes( MuonArray, # noqa: F821 @@ -265,7 +265,7 @@ class Jet(Particle, base.NanoCollection): ... _set_repr_name("Jet") -behavior.update(awkward._util.copy_behaviors(Particle, Jet, behavior)) +behavior.update(awkward._util.copy_behaviors("Particle", "Jet", behavior)) register_projection_classes( JetArray, # noqa: F821 @@ -281,7 +281,7 @@ class Track(Particle, base.NanoCollection): ... _set_repr_name("Track") -behavior.update(awkward._util.copy_behaviors(Particle, Track, behavior)) +behavior.update(awkward._util.copy_behaviors("Particle", "Track", behavior)) register_projection_classes( TrackArray, # noqa: F821 @@ -300,7 +300,7 @@ def pt(self): _set_repr_name("Tower") -behavior.update(awkward._util.copy_behaviors(MasslessParticle, Tower, behavior)) +behavior.update(awkward._util.copy_behaviors("MasslessParticle", "Tower", behavior)) register_projection_classes( TowerArray, # noqa: F821 diff --git a/src/coffea/nanoevents/methods/nanoaod.py b/src/coffea/nanoevents/methods/nanoaod.py index 141d329d0..43be888fa 100644 --- a/src/coffea/nanoevents/methods/nanoaod.py +++ b/src/coffea/nanoevents/methods/nanoaod.py @@ -41,7 +41,7 @@ class PtEtaPhiMCollection(vector.PtEtaPhiMLorentzVector, base.NanoCollection): behavior.update( awkward._util.copy_behaviors( - vector.PtEtaPhiMLorentzVector, PtEtaPhiMCollection, behavior + "PtEtaPhiMLorentzVector", "PtEtaPhiMCollection", behavior ) ) @@ -161,7 +161,7 @@ def distinctChildrenDeep(self, dask_array): _set_repr_name("GenParticle") behavior.update( - awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, GenParticle, behavior) + awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "GenParticle", behavior) ) register_projection_classes( @@ -192,7 +192,7 @@ def parent(self, dask_array): _set_repr_name("GenVisTau") behavior.update( - awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, GenVisTau, behavior) + awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "GenVisTau", behavior) ) register_projection_classes( @@ -267,7 +267,7 @@ def matched_photon(self, dask_array): _set_repr_name("Electron") behavior.update( - awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Electron, behavior) + awkward._util.copy_behaviors("PtEtaPhiMCandidate", "Electron", behavior) ) register_projection_classes( @@ -312,7 +312,7 @@ def matched_photon(self, dask_array): _set_repr_name("LowPtElectron") awkward.behavior.update( - awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, LowPtElectron, behavior) + awkward._util.copy_behaviors("PtEtaPhiMCandidate", "LowPtElectron", behavior) ) LowPtElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 @@ -353,9 +353,7 @@ def matched_jet(self, dask_array): _set_repr_name("Muon") -behavior.update( - awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Muon, behavior) -) +behavior.update(awkward._util.copy_behaviors("PtEtaPhiMCandidate", "Muon", behavior)) register_projection_classes( MuonArray, # noqa: F821 @@ -388,9 +386,7 @@ def matched_jet(self, dask_array): _set_repr_name("Tau") -behavior.update( - awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Tau, behavior) -) +behavior.update(awkward._util.copy_behaviors("PtEtaPhiMCandidate", "Tau", behavior)) register_projection_classes( TauArray, # noqa: F821 @@ -474,9 +470,7 @@ def matched_jet(self, dask_array): _set_repr_name("Photon") -behavior.update( - awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, Photon, behavior) -) +behavior.update(awkward._util.copy_behaviors("PtEtaPhiMCandidate", "Photon", behavior)) register_projection_classes( PhotonArray, # noqa: F821 @@ -502,7 +496,7 @@ def matched_muon(self, dask_array): _set_repr_name("FsrPhoton") behavior.update( - awkward._util.copy_behaviors(candidate.PtEtaPhiMCandidate, FsrPhoton, behavior) + awkward._util.copy_behaviors("PtEtaPhiMCandidate", "FsrPhoton", behavior) ) register_projection_classes( @@ -582,9 +576,7 @@ def constituents(self, dask_array): _set_repr_name("Jet") -behavior.update( - awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Jet, behavior) -) +behavior.update(awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "Jet", behavior)) register_projection_classes( JetArray, # noqa: F821 @@ -656,7 +648,7 @@ def constituents(self, dask_array): _set_repr_name("FatJet") behavior.update( - awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, FatJet, behavior) + awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "FatJet", behavior) ) register_projection_classes( @@ -678,9 +670,7 @@ def r(self): _set_repr_name("MissingET") -behavior.update( - awkward._util.copy_behaviors(vector.PolarTwoVector, MissingET, behavior) -) +behavior.update(awkward._util.copy_behaviors("PolarTwoVector", "MissingET", behavior)) register_projection_classes( MissingETArray, # noqa: F821 diff --git a/src/coffea/nanoevents/methods/pdune.py b/src/coffea/nanoevents/methods/pdune.py index 088855850..55e082d74 100644 --- a/src/coffea/nanoevents/methods/pdune.py +++ b/src/coffea/nanoevents/methods/pdune.py @@ -93,7 +93,7 @@ def mass(self): _set_repr_name("Particle") -behavior.update(awkward._util.copy_behaviors(vector.LorentzVector, Particle, behavior)) +behavior.update(awkward._util.copy_behaviors("LorentzVector", "Particle", behavior)) register_projection_classes( ParticleArray, # noqa: F821 @@ -141,7 +141,7 @@ def t(self): _set_repr_name("TrackParticle") behavior.update( - awkward._util.copy_behaviors(vector.LorentzVector, TrackParticle, behavior) + awkward._util.copy_behaviors("LorentzVector", "TrackParticle", behavior) ) register_projection_classes( @@ -170,7 +170,7 @@ def trackParticle(self): _set_repr_name("Muon") -behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) +behavior.update(awkward._util.copy_behaviors("Particle", "Muon", behavior)) register_projection_classes( MuonArray, # noqa: F821 @@ -206,7 +206,7 @@ def trackParticle(self): _set_repr_name("Electron") -behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) +behavior.update(awkward._util.copy_behaviors("Particle", "Electron", behavior)) register_projection_classes( ElectronArray, # noqa: F821 @@ -258,7 +258,7 @@ def parents(self): _set_repr_name("TruthParticle") behavior.update( - awkward._util.copy_behaviors(vector.LorentzVector, TruthParticle, behavior) + awkward._util.copy_behaviors("LorentzVector", "TruthParticle", behavior) ) register_projection_classes( diff --git a/src/coffea/nanoevents/methods/physlite.py b/src/coffea/nanoevents/methods/physlite.py index d04787208..732198c36 100644 --- a/src/coffea/nanoevents/methods/physlite.py +++ b/src/coffea/nanoevents/methods/physlite.py @@ -145,7 +145,7 @@ def mass(self): _set_repr_name("Particle") behavior.update( - awkward._util.copy_behaviors(vector.PtEtaPhiMLorentzVector, Particle, behavior) + awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "Particle", behavior) ) register_projection_classes( @@ -194,7 +194,7 @@ def t(self): _set_repr_name("TrackParticle") behavior.update( - awkward._util.copy_behaviors(vector.LorentzVector, TrackParticle, behavior) + awkward._util.copy_behaviors("LorentzVector", "TrackParticle", behavior) ) register_projection_classes( @@ -232,7 +232,7 @@ def trackParticle(self, dask_array): _set_repr_name("Muon") -behavior.update(awkward._util.copy_behaviors(Particle, Muon, behavior)) +behavior.update(awkward._util.copy_behaviors("Particle", "Muon", behavior)) register_projection_classes( MuonArray, # noqa: F821 @@ -293,7 +293,7 @@ def caloClusters(self, dask_array): _set_repr_name("Electron") -behavior.update(awkward._util.copy_behaviors(Particle, Electron, behavior)) +behavior.update(awkward._util.copy_behaviors("Particle", "Electron", behavior)) register_projection_classes( ElectronArray, # noqa: F821 @@ -345,7 +345,7 @@ def parents(self): _set_repr_name("TruthParticle") behavior.update( - awkward._util.copy_behaviors(vector.LorentzVector, TruthParticle, behavior) + awkward._util.copy_behaviors("LorentzVector", "TruthParticle", behavior) ) register_projection_classes( From 01f1c041261b4c301df903984d22198fe905ea61 Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Tue, 23 Jul 2024 15:35:04 +0200 Subject: [PATCH 43/46] Update src/coffea/nanoevents/methods/nanoaod.py --- src/coffea/nanoevents/methods/nanoaod.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/coffea/nanoevents/methods/nanoaod.py b/src/coffea/nanoevents/methods/nanoaod.py index 43be888fa..c3dff1695 100644 --- a/src/coffea/nanoevents/methods/nanoaod.py +++ b/src/coffea/nanoevents/methods/nanoaod.py @@ -311,7 +311,7 @@ def matched_photon(self, dask_array): _set_repr_name("LowPtElectron") -awkward.behavior.update( +behavior.update( awkward._util.copy_behaviors("PtEtaPhiMCandidate", "LowPtElectron", behavior) ) From 955f769a968ddaf38566e684a41fbc7e3c18f9b5 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Mon, 5 Aug 2024 10:28:31 -0500 Subject: [PATCH 44/46] use python -m pytest (windows path change!?) --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 11b3cd4db..5727d2d9b 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -108,7 +108,7 @@ jobs: - name: Test with pytest run: | - pytest --cov-report=xml --cov=coffea --deselect=test_taskvine + python -m pytest --cov-report=xml --cov=coffea --deselect=test_taskvine - name: Upload codecov if: matrix.os == 'ubuntu-latest' && matrix.python-version == 3.12 uses: codecov/codecov-action@v4 From 94210ef996110b37642f5e152b638ee25653f7c0 Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Thu, 8 Aug 2024 17:45:15 +0100 Subject: [PATCH 45/46] fix: copy_behaviors should be called before class definition --- src/coffea/nanoevents/methods/base.py | 6 +- src/coffea/nanoevents/methods/candidate.py | 4 +- src/coffea/nanoevents/methods/delphes.py | 31 ++++++---- src/coffea/nanoevents/methods/nanoaod.py | 71 +++++++++++++--------- src/coffea/nanoevents/methods/pdune.py | 24 +++++--- src/coffea/nanoevents/methods/physlite.py | 28 +++++---- 6 files changed, 99 insertions(+), 65 deletions(-) diff --git a/src/coffea/nanoevents/methods/base.py b/src/coffea/nanoevents/methods/base.py index ec9ae2d6b..896f53a25 100644 --- a/src/coffea/nanoevents/methods/base.py +++ b/src/coffea/nanoevents/methods/base.py @@ -154,6 +154,9 @@ def add_systematic( Systematic.add_kind(kind) +behavior[("__typestr__", "NanoEvents")] = "event" + + @awkward.mixin_class(behavior) class NanoEvents(Systematic): """NanoEvents mixin class @@ -167,9 +170,6 @@ def metadata(self): return self.layout.purelist_parameter("metadata") -behavior[("__typestr__", "NanoEvents")] = "event" - - @awkward.mixin_class(behavior) class NanoCollection: """A NanoEvents collection diff --git a/src/coffea/nanoevents/methods/candidate.py b/src/coffea/nanoevents/methods/candidate.py index 53f1e6789..31a6ac0e5 100644 --- a/src/coffea/nanoevents/methods/candidate.py +++ b/src/coffea/nanoevents/methods/candidate.py @@ -12,6 +12,8 @@ behavior = dict(vector.behavior) +behavior.update(awkward._util.copy_behaviors("LorentzVector", "Candidate", behavior)) + @awkward.mixin_class(behavior) class Candidate(vector.LorentzVector): @@ -70,8 +72,6 @@ class PtEtaPhiECandidate(Candidate, vector.PtEtaPhiELorentzVector): pass -behavior.update(awkward._util.copy_behaviors("LorentzVector", "Candidate", behavior)) - CandidateArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 CandidateArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 CandidateArray.ProjectionClass4D = vector.LorentzVectorArray # noqa: F821 diff --git a/src/coffea/nanoevents/methods/delphes.py b/src/coffea/nanoevents/methods/delphes.py index 1f090a61f..6b061d369 100644 --- a/src/coffea/nanoevents/methods/delphes.py +++ b/src/coffea/nanoevents/methods/delphes.py @@ -107,6 +107,8 @@ def eta(self): _set_repr_name("MissingET") +behavior.update(awkward._util.copy_behaviors("LorentzVector", "Vertex", behavior)) + @awkward.mixin_class(behavior) class Vertex(vector.LorentzVector): @@ -130,13 +132,16 @@ def z(self): _set_repr_name("Vertex") -behavior.update(awkward._util.copy_behaviors("LorentzVector", "Vertex", behavior)) VertexArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 VertexArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 VertexArray.ProjectionClass4D = VertexArray # noqa: F821 VertexArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update( + awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "Particle", behavior) +) + @awkward.mixin_class(behavior) class Particle(vector.PtEtaPhiMLorentzVector): @@ -176,15 +181,14 @@ def mass(self): _set_repr_name("Particle") -behavior.update( - awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "Particle", behavior) -) ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 ParticleArray.ProjectionClass4D = ParticleArray # noqa: F821 ParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors("Particle", "MasslessParticle", behavior)) + @awkward.mixin_class(behavior) class MasslessParticle(Particle, base.NanoCollection): @@ -194,78 +198,84 @@ def mass(self): _set_repr_name("MasslessParticle") -behavior.update(awkward._util.copy_behaviors("Particle", "MasslessParticle", behavior)) MasslessParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 MasslessParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 MasslessParticleArray.ProjectionClass4D = MasslessParticleArray # noqa: F821 MasslessParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors("MasslessParticle", "Photon", behavior)) + @awkward.mixin_class(behavior) class Photon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Photon") -behavior.update(awkward._util.copy_behaviors("MasslessParticle", "Photon", behavior)) PhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 PhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 PhotonArray.ProjectionClass4D = PhotonArray # noqa: F821 PhotonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors("MasslessParticle", "Electron", behavior)) + @awkward.mixin_class(behavior) class Electron(MasslessParticle, base.NanoCollection): ... _set_repr_name("Electron") -behavior.update(awkward._util.copy_behaviors("MasslessParticle", "Electron", behavior)) ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors("MasslessParticle", "Muon", behavior)) + @awkward.mixin_class(behavior) class Muon(MasslessParticle, base.NanoCollection): ... _set_repr_name("Muon") -behavior.update(awkward._util.copy_behaviors("MasslessParticle", "Muon", behavior)) MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 MuonArray.ProjectionClass4D = MuonArray # noqa: F821 MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors("Particle", "Jet", behavior)) + @awkward.mixin_class(behavior) class Jet(Particle, base.NanoCollection): ... _set_repr_name("Jet") -behavior.update(awkward._util.copy_behaviors("Particle", "Jet", behavior)) JetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 JetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 JetArray.ProjectionClass4D = JetArray # noqa: F821 JetArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors("Particle", "Track", behavior)) + @awkward.mixin_class(behavior) class Track(Particle, base.NanoCollection): ... _set_repr_name("Track") -behavior.update(awkward._util.copy_behaviors("Particle", "Track", behavior)) TrackArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 TrackArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 TrackArray.ProjectionClass4D = TrackArray # noqa: F821 TrackArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors("MasslessParticle", "Tower", behavior)) + @awkward.mixin_class(behavior) class Tower(MasslessParticle, base.NanoCollection): @@ -275,7 +285,6 @@ def pt(self): _set_repr_name("Tower") -behavior.update(awkward._util.copy_behaviors("MasslessParticle", "Tower", behavior)) TowerArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 TowerArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 diff --git a/src/coffea/nanoevents/methods/nanoaod.py b/src/coffea/nanoevents/methods/nanoaod.py index bfba663d6..f6517a599 100644 --- a/src/coffea/nanoevents/methods/nanoaod.py +++ b/src/coffea/nanoevents/methods/nanoaod.py @@ -31,6 +31,13 @@ def namefcn(self): behavior[classname].__repr__ = namefcn +behavior.update( + awkward._util.copy_behaviors( + "PtEtaPhiMLorentzVector", "PtEtaPhiMCollection", behavior + ) +) + + @awkward.mixin_class(behavior) class PtEtaPhiMCollection(vector.PtEtaPhiMLorentzVector, base.NanoCollection): """Generic collection that has Lorentz vector properties""" @@ -38,18 +45,17 @@ class PtEtaPhiMCollection(vector.PtEtaPhiMLorentzVector, base.NanoCollection): pass -behavior.update( - awkward._util.copy_behaviors( - "PtEtaPhiMLorentzVector", "PtEtaPhiMCollection", behavior - ) -) - PtEtaPhiMCollectionArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 PtEtaPhiMCollectionArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 PtEtaPhiMCollectionArray.ProjectionClass4D = PtEtaPhiMCollectionArray # noqa: F821 PtEtaPhiMCollectionArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update( + awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "GenParticle", behavior) +) + + @awkward.mixin_class(behavior) class GenParticle(vector.PtEtaPhiMLorentzVector, base.NanoCollection): """NanoAOD generator-level particle object, including parent and child self-references @@ -156,15 +162,16 @@ def distinctChildrenDeep(self, dask_array): _set_repr_name("GenParticle") -behavior.update( - awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "GenParticle", behavior) -) GenParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 GenParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 GenParticleArray.ProjectionClass4D = GenParticleArray # noqa: F821 GenParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update( + awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "GenVisTau", behavior) +) + @awkward.mixin_class(behavior) class GenVisTau(candidate.PtEtaPhiMCandidate, base.NanoCollection): @@ -184,15 +191,16 @@ def parent(self, dask_array): _set_repr_name("GenVisTau") -behavior.update( - awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "GenVisTau", behavior) -) GenVisTauArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 GenVisTauArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 GenVisTauArray.ProjectionClass4D = GenVisTauArray # noqa: F821 GenVisTauArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update( + awkward._util.copy_behaviors("PtEtaPhiMCandidate", "Electron", behavior) +) + @awkward.mixin_class(behavior) class Electron(candidate.PtEtaPhiMCandidate, base.NanoCollection, base.Systematic): @@ -256,15 +264,16 @@ def matched_photon(self, dask_array): _set_repr_name("Electron") -behavior.update( - awkward._util.copy_behaviors("PtEtaPhiMCandidate", "Electron", behavior) -) ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update( + awkward._util.copy_behaviors("PtEtaPhiMCandidate", "LowPtElectron", behavior) +) + @awkward.mixin_class(behavior) class LowPtElectron(candidate.PtEtaPhiMCandidate, base.NanoCollection, base.Systematic): @@ -298,15 +307,14 @@ def matched_photon(self, dask_array): _set_repr_name("LowPtElectron") -behavior.update( - awkward._util.copy_behaviors("PtEtaPhiMCandidate", "LowPtElectron", behavior) -) LowPtElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 LowPtElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 LowPtElectronArray.ProjectionClass4D = LowPtElectronArray # noqa: F821 LowPtElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors("PtEtaPhiMCandidate", "Muon", behavior)) + @awkward.mixin_class(behavior) class Muon(candidate.PtEtaPhiMCandidate, base.NanoCollection, base.Systematic): @@ -340,13 +348,14 @@ def matched_jet(self, dask_array): _set_repr_name("Muon") -behavior.update(awkward._util.copy_behaviors("PtEtaPhiMCandidate", "Muon", behavior)) MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 MuonArray.ProjectionClass4D = MuonArray # noqa: F821 MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors("PtEtaPhiMCandidate", "Tau", behavior)) + @awkward.mixin_class(behavior) class Tau(candidate.PtEtaPhiMCandidate, base.NanoCollection, base.Systematic): @@ -370,13 +379,14 @@ def matched_jet(self, dask_array): _set_repr_name("Tau") -behavior.update(awkward._util.copy_behaviors("PtEtaPhiMCandidate", "Tau", behavior)) TauArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 TauArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 TauArray.ProjectionClass4D = TauArray # noqa: F821 TauArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors("PtEtaPhiMCandidate", "Photon", behavior)) + @awkward.mixin_class(behavior) class Photon(candidate.PtEtaPhiMCandidate, base.NanoCollection, base.Systematic): @@ -451,13 +461,16 @@ def matched_jet(self, dask_array): _set_repr_name("Photon") -behavior.update(awkward._util.copy_behaviors("PtEtaPhiMCandidate", "Photon", behavior)) PhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 PhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 PhotonArray.ProjectionClass4D = PhotonArray # noqa: F821 PhotonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update( + awkward._util.copy_behaviors("PtEtaPhiMCandidate", "FsrPhoton", behavior) +) + @awkward.mixin_class(behavior) class FsrPhoton(candidate.PtEtaPhiMCandidate, base.NanoCollection): @@ -473,15 +486,14 @@ def matched_muon(self, dask_array): _set_repr_name("FsrPhoton") -behavior.update( - awkward._util.copy_behaviors("PtEtaPhiMCandidate", "FsrPhoton", behavior) -) FsrPhotonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 FsrPhotonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 FsrPhotonArray.ProjectionClass4D = FsrPhotonArray # noqa: F821 FsrPhotonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "Jet", behavior)) + @awkward.mixin_class(behavior) class Jet(vector.PtEtaPhiMLorentzVector, base.NanoCollection, base.Systematic): @@ -551,13 +563,16 @@ def constituents(self, dask_array): _set_repr_name("Jet") -behavior.update(awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "Jet", behavior)) JetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 JetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 JetArray.ProjectionClass4D = JetArray # noqa: F821 JetArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update( + awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "FatJet", behavior) +) + @awkward.mixin_class(behavior) class FatJet(vector.PtEtaPhiMLorentzVector, base.NanoCollection, base.Systematic): @@ -619,15 +634,14 @@ def constituents(self, dask_array): _set_repr_name("FatJet") -behavior.update( - awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "FatJet", behavior) -) FatJetArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 FatJetArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 FatJetArray.ProjectionClass4D = FatJetArray # noqa: F821 FatJetArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors("PolarTwoVector", "MissingET", behavior)) + @awkward.mixin_class(behavior) class MissingET(vector.PolarTwoVector, base.NanoCollection, base.Systematic): @@ -639,7 +653,6 @@ def r(self): _set_repr_name("MissingET") -behavior.update(awkward._util.copy_behaviors("PolarTwoVector", "MissingET", behavior)) MissingETArray.ProjectionClass2D = MissingETArray # noqa: F821 MissingETArray.ProjectionClass3D = vector.SphericalThreeVectorArray # noqa: F821 diff --git a/src/coffea/nanoevents/methods/pdune.py b/src/coffea/nanoevents/methods/pdune.py index 5e19e2a7a..a32e6c0e0 100644 --- a/src/coffea/nanoevents/methods/pdune.py +++ b/src/coffea/nanoevents/methods/pdune.py @@ -82,6 +82,9 @@ def _get_global_index(target, eventindex, index): return target_offsets + index +behavior.update(awkward._util.copy_behaviors("LorentzVector", "Particle", behavior)) + + @awkward.mixin_class(behavior) class Particle(vector.LorentzVector, base.NanoCollection): """Generic particle collection that has Lorentz vector properties""" @@ -92,13 +95,16 @@ def mass(self): _set_repr_name("Particle") -behavior.update(awkward._util.copy_behaviors("LorentzVector", "Particle", behavior)) ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 ParticleArray.ProjectionClass4D = ParticleArray # noqa: F821 ParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update( + awkward._util.copy_behaviors("LorentzVector", "TrackParticle", behavior) +) + @awkward.mixin_class(behavior) class TrackParticle(vector.LorentzVector, base.NanoCollection): @@ -136,15 +142,14 @@ def t(self): _set_repr_name("TrackParticle") -behavior.update( - awkward._util.copy_behaviors("LorentzVector", "TrackParticle", behavior) -) TrackParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 TrackParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 TrackParticleArray.ProjectionClass4D = TrackParticleArray # noqa: F821 TrackParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors("Particle", "Muon", behavior)) + @awkward.mixin_class(behavior) class Muon(Particle): @@ -163,13 +168,14 @@ def trackParticle(self): _set_repr_name("Muon") -behavior.update(awkward._util.copy_behaviors("Particle", "Muon", behavior)) MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 MuonArray.ProjectionClass4D = MuonArray # noqa: F821 MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors("Particle", "Electron", behavior)) + @awkward.mixin_class(behavior) class Electron(Particle): @@ -196,13 +202,16 @@ def trackParticle(self): _set_repr_name("Electron") -behavior.update(awkward._util.copy_behaviors("Particle", "Electron", behavior)) ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update( + awkward._util.copy_behaviors("LorentzVector", "TruthParticle", behavior) +) + @awkward.mixin_class(behavior) class TruthParticle(vector.LorentzVector, base.NanoCollection): @@ -244,9 +253,6 @@ def parents(self): _set_repr_name("TruthParticle") -behavior.update( - awkward._util.copy_behaviors("LorentzVector", "TruthParticle", behavior) -) TruthParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 TruthParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 diff --git a/src/coffea/nanoevents/methods/physlite.py b/src/coffea/nanoevents/methods/physlite.py index 01298b3e9..61165ee56 100644 --- a/src/coffea/nanoevents/methods/physlite.py +++ b/src/coffea/nanoevents/methods/physlite.py @@ -133,6 +133,11 @@ def _get_global_index(target, eventindex, index): return target_offsets + index +behavior.update( + awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "Particle", behavior) +) + + @awkward.mixin_class(behavior) class Particle(vector.PtEtaPhiMLorentzVector, base.NanoCollection): """Generic particle collection that has Lorentz vector properties""" @@ -143,15 +148,16 @@ def mass(self): _set_repr_name("Particle") -behavior.update( - awkward._util.copy_behaviors("PtEtaPhiMLorentzVector", "Particle", behavior) -) ParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 ParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 ParticleArray.ProjectionClass4D = ParticleArray # noqa: F821 ParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update( + awkward._util.copy_behaviors("LorentzVector", "TrackParticle", behavior) +) + @awkward.mixin_class(behavior) class TrackParticle(vector.LorentzVector, base.NanoCollection): @@ -189,15 +195,14 @@ def t(self): _set_repr_name("TrackParticle") -behavior.update( - awkward._util.copy_behaviors("LorentzVector", "TrackParticle", behavior) -) TrackParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 TrackParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 TrackParticleArray.ProjectionClass4D = TrackParticleArray # noqa: F821 TrackParticleArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors("Particle", "Muon", behavior)) + @awkward.mixin_class(behavior) class Muon(Particle): @@ -225,13 +230,14 @@ def trackParticle(self, dask_array): _set_repr_name("Muon") -behavior.update(awkward._util.copy_behaviors("Particle", "Muon", behavior)) MuonArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 MuonArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 MuonArray.ProjectionClass4D = MuonArray # noqa: F821 MuonArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update(awkward._util.copy_behaviors("Particle", "Electron", behavior)) + @awkward.mixin_class(behavior) class Electron(Particle): @@ -283,13 +289,16 @@ def caloClusters(self, dask_array): _set_repr_name("Electron") -behavior.update(awkward._util.copy_behaviors("Particle", "Electron", behavior)) ElectronArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 ElectronArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 ElectronArray.ProjectionClass4D = ElectronArray # noqa: F821 ElectronArray.MomentumClass = vector.LorentzVectorArray # noqa: F821 +behavior.update( + awkward._util.copy_behaviors("LorentzVector", "TruthParticle", behavior) +) + @awkward.mixin_class(behavior) class TruthParticle(vector.LorentzVector, base.NanoCollection): @@ -331,9 +340,6 @@ def parents(self): _set_repr_name("TruthParticle") -behavior.update( - awkward._util.copy_behaviors("LorentzVector", "TruthParticle", behavior) -) TruthParticleArray.ProjectionClass2D = vector.TwoVectorArray # noqa: F821 TruthParticleArray.ProjectionClass3D = vector.ThreeVectorArray # noqa: F821 From 72f117d61d8432072f052cb9cdb9db9dffefcece Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Thu, 8 Aug 2024 13:11:06 -0500 Subject: [PATCH 46/46] bump date of complete vector deprecation --- src/coffea/nanoevents/methods/vector.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/coffea/nanoevents/methods/vector.py b/src/coffea/nanoevents/methods/vector.py index 2944aa72d..2f111373b 100644 --- a/src/coffea/nanoevents/methods/vector.py +++ b/src/coffea/nanoevents/methods/vector.py @@ -61,14 +61,14 @@ from coffea.util import deprecate _cst = pytz.timezone("US/Central") -_depttime = _cst.localize(datetime(2024, 6, 30, 11, 59, 59)) +_depttime = _cst.localize(datetime(2024, 12, 31, 11, 59, 59)) deprecate( ( "coffea.nanoevents.methods.vector will be removed and replaced with scikit-hep vector. " "Nanoevents schemas internal to coffea will be migrated. " "Otherwise please consider using that package!" ), - version="2024.7.0", + version="2025.1.0", date=str(_depttime), category=FutureWarning, )