From d79cebcb814409823a0007e67111f09ba5914b42 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jonas=20R=C3=BCbenach?= Date: Tue, 9 Feb 2021 18:58:48 +0100 Subject: [PATCH 1/3] Extend Nanoevents vector and add tests Closes #445 --- coffea/nanoevents/methods/vector.py | 326 +++++++++++++++++++++++- tests/test_nanoevents_vector.py | 370 ++++++++++++++++++++++++++++ 2 files changed, 687 insertions(+), 9 deletions(-) create mode 100644 tests/test_nanoevents_vector.py diff --git a/coffea/nanoevents/methods/vector.py b/coffea/nanoevents/methods/vector.py index a34c527d6..dd2c83659 100644 --- a/coffea/nanoevents/methods/vector.py +++ b/coffea/nanoevents/methods/vector.py @@ -108,6 +108,14 @@ def absolute(self): """ return self.r + @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", + ) + @awkward.mixin_class_method(numpy.add, {"TwoVector"}) def add(self, other): """Add two vectors together elementwise using `x` and `y` components""" @@ -116,6 +124,14 @@ def add(self, other): with_name="TwoVector", ) + @awkward.mixin_class_method(numpy.subtract, {"TwoVector"}) + def subtract(self, other): + """Substract a vector from another elementwise using `x` and `y` compontents""" + return awkward.zip( + {"x": self.x - other.x, "y": self.y - other.y}, + with_name="TwoVector", + ) + def sum(self, axis=-1): """Sum an array of vectors elementwise using `x` and `y` components""" out = awkward.zip( @@ -136,6 +152,14 @@ def multiply(self, other): with_name="TwoVector", ) + @awkward.mixin_class_method(numpy.divide, {numbers.Number}) + def divide(self, other): + """Divide 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", + ) + def delta_phi(self, other): """Compute difference in angle between two vectors @@ -143,6 +167,15 @@ def delta_phi(self, other): """ return (self.phi - other.phi + numpy.pi) % (2 * numpy.pi) - numpy.pi + def dot(self, other): + """Compute the dot product of two vectors""" + return self.x * other.x + self.y * other.y + + @property + def unit(self): + """Unit vector, a vector of length 1 pointing the same direction""" + return self / self.r + @awkward.mixin_class(behavior) class PolarTwoVector(TwoVector): @@ -189,6 +222,42 @@ def r2(self): """Squared `r`""" return self.r ** 2 + @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 + + In reality, this directly adjusts `r` and `phi` for performance + """ + return awkward.zip( + { + "r": self.r * abs(other), + "phi": self.phi % (2 * numpy.pi) - (numpy.pi * (other < 0)) + }, + with_name="PolarTwoVector", + ) + + @awkward.mixin_class_method(numpy.divide, {numbers.Number}) + def divide(self, other): + """Divide this vector by a scalar elementwise using using `x` and `y` components + + In reality, this directly adjusts `r` and `phi` for performance + """ + return awkward.zip( + { + "r": self.r / abs(other), + "phi": self.phi % (2 * numpy.pi) + (numpy.pi * (other < 0)) + }, + with_name="PolarTwoVector", + ) + + @awkward.mixin_class_method(numpy.negative) + def negative(self): + """Returns the negative of the vector""" + return awkward.zip( + {"r": self.r, "phi": self.phi % (2 * numpy.pi) - numpy.pi}, + with_name="PolarTwoVector", + ) + @awkward.mixin_class(behavior) class ThreeVector(TwoVector): @@ -242,6 +311,14 @@ def absolute(self): """ return self.p + @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", + ) + @awkward.mixin_class_method(numpy.add, {"ThreeVector"}) def add(self, other): """Add two vectors together elementwise using `x`, `y`, and `z` components""" @@ -250,6 +327,14 @@ def add(self, other): with_name="ThreeVector", ) + @awkward.mixin_class_method(numpy.subtract, {"ThreeVector"}) + 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", + ) + def sum(self, axis=-1): """Sum an array of vectors elementwise using `x`, `y`, and `z` components""" out = awkward.zip( @@ -271,6 +356,34 @@ def multiply(self, other): with_name="ThreeVector", ) + @awkward.mixin_class_method(numpy.divide, {numbers.Number}) + def divide(self, other): + """Divide 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", + ) + + 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" + ) + + @property + def unit(self): + """Unit vector, a vector of length 1 pointing the same direction""" + return self / self.rho + @awkward.mixin_class(behavior) class SphericalThreeVector(ThreeVector, PolarTwoVector): @@ -322,6 +435,48 @@ def p2(self): """Squared `p`""" return self.rho ** 2 + @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 + + 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", + ) + + @awkward.mixin_class_method(numpy.divide, {numbers.Number}) + def divide(self, other): + """Divide this vector by a scalar elementwise using `x`, `y`, and `z` components + + In reality, this directly adjusts `r`, `theta` `phi` for performance + """ + return awkward.zip( + { + "rho": self.rho / other, + "theta": (numpy.sign(other) * self.theta + numpy.pi) % numpy.pi, + "phi": self.phi % (2 * numpy.pi) - numpy.pi * (other < 0) + }, + with_name="SphericalThreeVector", + ) + + @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", + ) + @awkward.mixin_class(behavior) class LorentzVector(ThreeVector): @@ -379,6 +534,19 @@ def add(self, other): with_name="LorentzVector", ) + @awkward.mixin_class_method(numpy.subtract, {"LorentzVector"}) + 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", + ) + def sum(self, axis=-1): """Sum an array of vectors elementwise using `x`, `y`, `z`, and `t` components""" out = awkward.zip( @@ -406,6 +574,19 @@ def multiply(self, other): with_name="LorentzVector", ) + @awkward.mixin_class_method(numpy.divide, {numbers.Number}) + def divide(self, other): + """Divide 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", + ) + def delta_r2(self, other): """Squared `delta_r`""" return (self.eta - other.eta) ** 2 + self.delta_phi(other) ** 2 @@ -417,6 +598,75 @@ def delta_r(self, other): """ return numpy.sqrt(self.delta_r2(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", + ) + + @property + def pvec(self): + """The `x`, `y` and `z` compontents as a `ThreeVector`""" + return awkward.zip( + { + "x": self.x, + "y": self.y, + "z": self.z + }, + with_name="ThreeVector" + ) + + @property + def boostvec(self): + """The `x`, `y` and `z` compontents divided by `t` as a `ThreeVector` + + This can be used for boosting. + """ + pvec = self.pvec + # Allow self to be a zero vector by using awkward.where + mask = self.t == 0 + t = awkward.where(mask, 1, self.t) + out = awkward.where( + mask, + awkward.zeros_like(pvec), + self.pvec / t, + highlevel=False + ) + return awkward.Array(out, behavior=self.behavior) + + 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) + v = gamma2 * bp * other + self.t * gamma * other + + out = awkward.zip( + { + "x": self.x + v.x, + "y": self.y + v.y, + "z": self.z + v.z, + "t": gamma * (self.t + bp) + }, + with_name="LorentzVector" + ) + return out + def metric_table( self, other, axis=1, metric=lambda a, b: a.delta_r(b), return_combinations=False ): @@ -584,14 +834,43 @@ def mass2(self): def multiply(self, other): """Multiply this vector by a scalar elementwise using `x`, `y`, `z`, and `t` components - In reality, this multiplies `pt` and `mass` by the scalar quantity for performance + In reality, this directly adjusts `pt`, `eta`, `phi` and `mass` for performance """ return awkward.zip( { - "pt": self.pt * other, - "eta": self.eta, - "phi": self.phi, - "mass": self.mass * other, + "pt": self.pt * abs(other), + "eta": self.eta * numpy.sign(other), + "phi": self.phi % (2 * numpy.pi) - (numpy.pi * (other < 0)), + "mass": self.mass * abs(other), + }, + with_name="PtEtaPhiMLorentzVector", + ) + + @awkward.mixin_class_method(numpy.divide, {numbers.Number}) + def divide(self, other): + """Divides this vector by a scalar elementwise using `x`, `y`, `z`, and `t` components + + In reality, this directly adjusts `pt`, `eta`, `phi` and `mass` for performance + """ + return awkward.zip( + { + "pt": self.pt / abs(other), + "eta": self.eta * numpy.sign(other), + "phi": self.phi % (2 * numpy.pi) - (numpy.pi * (other < 0)), + "mass": self.mass / abs(other), + }, + with_name="PtEtaPhiMLorentzVector", + ) + + @awkward.mixin_class_method(numpy.negative) + def negative(self): + """Returns the negative of the vector""" + return awkward.zip( + { + "pt": self.pt, + "eta": -self.eta, + "phi": self.phi % (2 * numpy.pi) - numpy.pi, + "mass": self.mass, }, with_name="PtEtaPhiMLorentzVector", ) @@ -680,18 +959,47 @@ def rho2(self): def multiply(self, other): """Multiply this vector by a scalar elementwise using `x`, `y`, `z`, and `t` components - In reality, this multiplies `pt` and `energy` by the scalar quantity for performance + In reality, this directly adjusts `pt`, `eta`, `phi` and `energy` for performance """ return awkward.zip( { - "pt": self.pt * other, - "eta": self.eta, - "phi": self.phi, + "pt": self.pt * abs(other), + "eta": self.eta * numpy.sign(other), + "phi": self.phi % (2 * numpy.pi) - (numpy.pi * (other < 0)), "energy": self.energy * other, }, with_name="PtEtaPhiELorentzVector", ) + @awkward.mixin_class_method(numpy.divide, {numbers.Number}) + def divide(self, other): + """Divides this vector by a scalar elementwise using `x`, `y`, `z`, and `t` components + + In reality, this directly adjusts `pt`, `eta`, `phi` and `energy` for performance + """ + return awkward.zip( + { + "pt": self.pt / abs(other), + "eta": self.eta * numpy.sign(other), + "phi": self.phi % (2 * numpy.pi) - (numpy.pi * (other < 0)), + "energy": self.energy / other, + }, + with_name="PtEtaPhiELorentzVector", + ) + + @awkward.mixin_class_method(numpy.negative) + def negative(self): + """Returns the negative of the vector""" + return awkward.zip( + { + "pt": self.pt, + "eta": -self.eta, + "phi": self.phi % (2 * numpy.pi) - numpy.pi, + "energy": -self.energy, + }, + with_name="PtEtaPhiELorentzVector", + ) + __all__ = [ "TwoVector", diff --git a/tests/test_nanoevents_vector.py b/tests/test_nanoevents_vector.py new file mode 100644 index 000000000..1c7b1135e --- /dev/null +++ b/tests/test_nanoevents_vector.py @@ -0,0 +1,370 @@ +import awkward as ak +from coffea.nanoevents.methods import vector +import pytest + + +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 test_two_vector(): + a = ak.zip( + { + "x": [[1, 2], [], [3], [4]], + "y": [[5, 6], [], [7], [8]] + }, + with_name="TwoVector", + highlevel=False + ) + a = ak.Array(a, behavior=vector.behavior) + b = ak.zip( + { + "x": [[11, 12], [], [13], [14]], + "y": [[15, 16], [], [17], [18]] + }, + with_name="TwoVector", + highlevel=False + ) + b = ak.Array(b, behavior=vector.behavior) + + assert record_arrays_equal(- a, ak.zip( + { + "x": [[-1, -2], [], [-3], [-4]], + "y": [[-5, -6], [], [-7], [-8]] + } + )) + + assert record_arrays_equal(a + b, ak.zip( + { + "x": [[12, 14], [], [16], [18]], + "y": [[20, 22], [], [24], [26]] + } + )) + assert record_arrays_equal(a - b, ak.zip( + { + "x": [[-10, -10], [], [-10], [-10]], + "y": [[-10, -10], [], [-10], [-10]] + } + )) + + assert record_arrays_equal(a * 2, ak.zip( + { + "x": [[2, 4], [], [6], [8]], + "y": [[10, 12], [], [14], [16]] + } + )) + 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 ak.all(abs(a.unit.r - 1) < ATOL) + assert ak.all(abs(a.unit.phi - a.phi) < ATOL) + + +def test_polar_two_vector(): + a = ak.zip( + { + "r": [[1, 2], [], [3], [4]], + "phi": [[0.3, 0.4], [], [0.5], [0.6]], + }, + with_name="PolarTwoVector", + highlevel=False + ) + a = ak.Array(a, behavior=vector.behavior) + + assert record_arrays_equal(a * 2, ak.zip( + { + "r": [[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)).phi - ak.Array([ + [-2.8415926535, -2.7415926535], + [], + [-2.6415926535], + [-2.5415926535] + ]) < ATOL) + assert record_arrays_equal(a / 2, ak.zip( + { + "r": [[0.5, 1], [], [1.5], [2]], + "phi": [[0.3, 0.4], [], [0.5], [0.6]] + } + )) + + 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 ak.all(a.unit.phi == a.phi) + + +def test_three_vector(): + a = ak.zip( + { + "x": [[1, 2], [], [3], [4]], + "y": [[5, 6], [], [7], [8]], + "z": [[9, 10], [], [11], [12]] + }, + with_name="ThreeVector", + highlevel=False + ) + a = ak.Array(a, behavior=vector.behavior) + b = ak.zip( + { + "x": [[4, 1], [], [10], [11]], + "y": [[17, 7], [], [11], [6]], + "z": [[9, 11], [], [5], [16]] + }, + with_name="ThreeVector", + highlevel=False + ) + b = ak.Array(b, behavior=vector.behavior) + + assert record_arrays_equal(- a, ak.zip( + { + "x": [[-1, -2], [], [-3], [-4]], + "y": [[-5, -6], [], [-7], [-8]], + "z": [[-9, -10], [], [-11], [-12]] + } + )) + + assert record_arrays_equal(a + b, ak.zip( + { + "x": [[5, 3], [], [13], [15]], + "y": [[22, 13], [], [18], [14]], + "z": [[18, 21], [], [16], [28]] + } + )) + assert record_arrays_equal(a - b, ak.zip( + { + "x": [[-3, 1], [], [-7], [-7]], + "y": [[-12, -1], [], [-4], [2]], + "z": [[0, -1], [], [6], [-4]] + } + )) + + assert record_arrays_equal(a * 2, ak.zip( + { + "x": [[2, 4], [], [6], [8]], + "y": [[10, 12], [], [14], [16]], + "z": [[18, 20], [], [22], [24]] + } + )) + assert record_arrays_equal(a / 2, ak.zip( + { + "x": [[0.5, 1], [], [1.5], [2]], + "y": [[2.5, 3], [], [3.5], [4]], + "z": [[4.5, 5], [], [5.5], [6]] + } + )) + + 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(a.cross(b), ak.zip( + { + "x": [[-108, -4], [], [-86], [56]], + "y": [[27, -12], [], [95], [68]], + "z": [[-3, 8], [], [-37], [-64]] + } + )) + assert record_arrays_equal(b.cross(a), ak.zip( + { + "x": [[108, 4], [], [86], [-56]], + "y": [[-27, 12], [], [-95], [-68]], + "z": [[3, -8], [], [37], [64]] + } + )) + + assert ak.all(abs(a.unit.rho - 1) < ATOL) + assert ak.all(abs(a.unit.phi - a.phi) < ATOL) + + +def test_spherical_three_vector(): + a = ak.zip( + { + "rho": [[1.0, 2.0], [], [3.0], [4.0]], + "theta": [[1.2, 0.7], [], [1.8], [1.9]], + "phi": [[0.3, 0.4], [], [0.5], [0.6]], + }, + with_name="SphericalThreeVector", + highlevel=False + ) + a = ak.Array(a, behavior=vector.behavior) + + 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) + +def test_lorentz_vector(): + a = ak.zip( + { + "x": [[1, 2], [], [3], [4]], + "y": [[5, 6], [], [7], [8]], + "z": [[9, 10], [], [11], [12]], + "t": [[50, 51], [], [52], [53]] + }, + with_name="LorentzVector", + highlevel=False + ) + a = ak.Array(a, behavior=vector.behavior) + b = ak.zip( + { + "x": [[4, 1], [], [10], [11]], + "y": [[17, 7], [], [11], [6]], + "z": [[9, 11], [], [5], [16]], + "t": [[60, 61], [], [62], [63]] + }, + with_name="LorentzVector", + highlevel=False + ) + b = ak.Array(b, behavior=vector.behavior) + + assert record_arrays_equal(- a, ak.zip( + { + "x": [[-1, -2], [], [-3], [-4]], + "y": [[-5, -6], [], [-7], [-8]], + "z": [[-9, -10], [], [-11], [-12]], + "t": [[-50, -51], [], [-52], [-53]] + } + )) + + assert record_arrays_equal(a + b, ak.zip( + { + "x": [[5, 3], [], [13], [15]], + "y": [[22, 13], [], [18], [14]], + "z": [[18, 21], [], [16], [28]], + "t": [[110, 112], [], [114], [116]] + } + )) + assert record_arrays_equal(a - b, ak.zip( + { + "x": [[-3, 1], [], [-7], [-7]], + "y": [[-12, -1], [], [-4], [2]], + "z": [[0, -1], [], [6], [-4]], + "t": [[-10, -10], [], [-10], [-10]] + } + )) + + assert record_arrays_equal(a * 2, ak.zip( + { + "x": [[2, 4], [], [6], [8]], + "y": [[10, 12], [], [14], [16]], + "z": [[18, 20], [], [22], [24]], + "t": [[100, 102], [], [104], [106]] + } + )) + assert record_arrays_equal(a / 2, ak.zip( + { + "x": [[0.5, 1], [], [1.5], [2]], + "y": [[2.5, 3], [], [3.5], [4]], + "z": [[4.5, 5], [], [5.5], [6]], + "t": [[25, 25.5], [], [26], [26.5]] + } + )) + + assert record_arrays_equal(a.pvec, ak.zip( + { + "x": [[1, 2], [], [3], [4]], + "y": [[5, 6], [], [7], [8]], + "z": [[9, 10], [], [11], [12]], + } + )) + + boosted = a.boost(-a.boostvec) + assert ak.all(abs(boosted.x) < ATOL) + assert ak.all(abs(boosted.y) < ATOL) + assert ak.all(abs(boosted.z) < ATOL) + +def test_pt_eta_phi_m_lorentz_vector(): + a = ak.zip( + { + "pt": [[1, 2], [], [3], [4]], + "eta": [[1.2, 1.4], [], [1.6], [3.4]], + "phi": [[0.3, 0.4], [], [0.5], [0.6]], + "mass": [[0.5, 0.9], [], [1.3], [4.5]] + }, + with_name="PtEtaPhiMLorentzVector", + highlevel=False + ) + a = ak.Array(a, behavior=vector.behavior) + + assert ak.all((a * (-2)).pt == ak.Array([[2, 4], [], [6], [8]])) + assert ak.all((a * (-2)).theta - ak.Array([ + [2.556488570968, 2.65804615357], + [], + [2.74315571762], + [3.07487087733] + ]) < ATOL) + assert ak.all((a * (-2)).phi - ak.Array([ + [-2.8415926535, -2.7415926535], + [], + [-2.6415926535], + [-2.5415926535] + ]) < ATOL) + assert record_arrays_equal(a / 2, ak.zip( + { + "pt": [[0.5, 1], [], [1.5], [2]], + "eta": [[1.2, 1.4], [], [1.6], [3.4]], + "phi": [[0.3, 0.4], [], [0.5], [0.6]], + "mass": [[0.25, 0.45], [], [0.65], [2.25]] + } + )) + assert record_arrays_equal(a * (-1), -a) + + boosted = a.boost(-a.boostvec) + assert ak.all(abs(boosted.x) < ATOL) + assert ak.all(abs(boosted.y) < ATOL) + assert ak.all(abs(boosted.z) < ATOL) + +def test_pt_eta_phi_e_lorentz_vector(): + a = ak.zip( + { + "pt": [[1, 2], [], [3], [4]], + "eta": [[1.2, 1.4], [], [1.6], [3.4]], + "phi": [[0.3, 0.4], [], [0.5], [0.6]], + "energy": [[50, 51], [], [52], [60]] + }, + with_name="PtEtaPhiELorentzVector", + highlevel=False + ) + a = ak.Array(a, behavior=vector.behavior) + + assert ak.all((a * (-2)).pt == ak.Array([[2, 4], [], [6], [8]])) + assert ak.all((a * (-2)).theta - ak.Array([ + [2.556488570968, 2.65804615357], + [], + [2.74315571762], + [3.07487087733] + ]) < ATOL) + assert ak.all((a * (-2)).phi - ak.Array([ + [-2.8415926535, -2.7415926535], + [], + [-2.6415926535], + [-2.5415926535] + ]) < ATOL) + assert record_arrays_equal(a / 2, ak.zip( + { + "pt": [[0.5, 1], [], [1.5], [2]], + "eta": [[1.2, 1.4], [], [1.6], [3.4]], + "phi": [[0.3, 0.4], [], [0.5], [0.6]], + "energy": [[25, 25.5], [], [26], [30]] + } + )) + assert record_arrays_equal(a * (-1), -a) + + boosted = a.boost(-a.boostvec) + assert ak.all(abs(boosted.x) < ATOL) + assert ak.all(abs(boosted.y) < ATOL) + assert ak.all(abs(boosted.z) < ATOL) From 72437245a3122b4560901e5795176fb16a5dd8ae Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jonas=20R=C3=BCbenach?= Date: Thu, 11 Feb 2021 16:52:59 +0100 Subject: [PATCH 2/3] Fix typo in documentation --- coffea/nanoevents/methods/vector.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/coffea/nanoevents/methods/vector.py b/coffea/nanoevents/methods/vector.py index dd2c83659..ba1ace520 100644 --- a/coffea/nanoevents/methods/vector.py +++ b/coffea/nanoevents/methods/vector.py @@ -173,7 +173,7 @@ def dot(self, other): @property def unit(self): - """Unit vector, a vector of length 1 pointing the same direction""" + """Unit vector, a vector of length 1 pointing in the same direction""" return self / self.r @@ -381,7 +381,7 @@ def cross(self, other): @property def unit(self): - """Unit vector, a vector of length 1 pointing the same direction""" + """Unit vector, a vector of length 1 pointing in the same direction""" return self / self.rho From c8b4e86bef4b6950a9b87d30943c6cb46a8ee5ea Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jonas=20R=C3=BCbenach?= Date: Fri, 12 Feb 2021 13:56:16 +0100 Subject: [PATCH 3/3] Nanoevents vector changes as discussed in the PR Single implementation for division Handle t < rho in boostvec Use more temporary variables --- coffea/nanoevents/methods/vector.py | 126 +++++----------------------- 1 file changed, 22 insertions(+), 104 deletions(-) diff --git a/coffea/nanoevents/methods/vector.py b/coffea/nanoevents/methods/vector.py index ba1ace520..ee0c11ef7 100644 --- a/coffea/nanoevents/methods/vector.py +++ b/coffea/nanoevents/methods/vector.py @@ -154,11 +154,10 @@ def multiply(self, other): @awkward.mixin_class_method(numpy.divide, {numbers.Number}) def divide(self, other): - """Divide 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", - ) + """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_phi(self, other): """Compute difference in angle between two vectors @@ -236,20 +235,6 @@ def multiply(self, other): with_name="PolarTwoVector", ) - @awkward.mixin_class_method(numpy.divide, {numbers.Number}) - def divide(self, other): - """Divide this vector by a scalar elementwise using using `x` and `y` components - - In reality, this directly adjusts `r` and `phi` for performance - """ - return awkward.zip( - { - "r": self.r / abs(other), - "phi": self.phi % (2 * numpy.pi) + (numpy.pi * (other < 0)) - }, - with_name="PolarTwoVector", - ) - @awkward.mixin_class_method(numpy.negative) def negative(self): """Returns the negative of the vector""" @@ -356,14 +341,6 @@ def multiply(self, other): with_name="ThreeVector", ) - @awkward.mixin_class_method(numpy.divide, {numbers.Number}) - def divide(self, other): - """Divide 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", - ) - def dot(self, other): """Compute the dot product of two vectors""" return self.x * other.x + self.y * other.y + self.z * other.z @@ -450,21 +427,6 @@ def multiply(self, other): with_name="SphericalThreeVector", ) - @awkward.mixin_class_method(numpy.divide, {numbers.Number}) - def divide(self, other): - """Divide this vector by a scalar elementwise using `x`, `y`, and `z` components - - In reality, this directly adjusts `r`, `theta` `phi` for performance - """ - return awkward.zip( - { - "rho": self.rho / other, - "theta": (numpy.sign(other) * self.theta + numpy.pi) % numpy.pi, - "phi": self.phi % (2 * numpy.pi) - numpy.pi * (other < 0) - }, - with_name="SphericalThreeVector", - ) - @awkward.mixin_class_method(numpy.negative) def negative(self): """Returns the negative of the vector""" @@ -574,19 +536,6 @@ def multiply(self, other): with_name="LorentzVector", ) - @awkward.mixin_class_method(numpy.divide, {numbers.Number}) - def divide(self, other): - """Divide 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", - ) - def delta_r2(self, other): """Squared `delta_r`""" return (self.eta - other.eta) ** 2 + self.delta_phi(other) ** 2 @@ -627,19 +576,18 @@ def pvec(self): def boostvec(self): """The `x`, `y` and `z` compontents divided by `t` as a `ThreeVector` - This can be used for boosting. - """ - pvec = self.pvec - # Allow self to be a zero vector by using awkward.where - mask = self.t == 0 - t = awkward.where(mask, 1, self.t) - out = awkward.where( - mask, - awkward.zeros_like(pvec), - self.pvec / t, - highlevel=False - ) - return awkward.Array(out, behavior=self.behavior) + 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 @@ -654,14 +602,15 @@ def boost(self, other): gamma2 = awkward.where(mask, 0, (gamma - 1) / b2) bp = self.dot(other) - v = gamma2 * bp * other + self.t * gamma * other + t = self.t + v = gamma2 * bp * other + t * gamma * other out = awkward.zip( { "x": self.x + v.x, "y": self.y + v.y, "z": self.z + v.z, - "t": gamma * (self.t + bp) + "t": gamma * (t + bp) }, with_name="LorentzVector" ) @@ -836,28 +785,13 @@ def multiply(self, other): In reality, this directly adjusts `pt`, `eta`, `phi` and `mass` for performance """ + absother = abs(other) return awkward.zip( { - "pt": self.pt * abs(other), - "eta": self.eta * numpy.sign(other), - "phi": self.phi % (2 * numpy.pi) - (numpy.pi * (other < 0)), - "mass": self.mass * abs(other), - }, - with_name="PtEtaPhiMLorentzVector", - ) - - @awkward.mixin_class_method(numpy.divide, {numbers.Number}) - def divide(self, other): - """Divides this vector by a scalar elementwise using `x`, `y`, `z`, and `t` components - - In reality, this directly adjusts `pt`, `eta`, `phi` and `mass` for performance - """ - return awkward.zip( - { - "pt": self.pt / abs(other), + "pt": self.pt * absother, "eta": self.eta * numpy.sign(other), "phi": self.phi % (2 * numpy.pi) - (numpy.pi * (other < 0)), - "mass": self.mass / abs(other), + "mass": self.mass * absother, }, with_name="PtEtaPhiMLorentzVector", ) @@ -971,22 +905,6 @@ def multiply(self, other): with_name="PtEtaPhiELorentzVector", ) - @awkward.mixin_class_method(numpy.divide, {numbers.Number}) - def divide(self, other): - """Divides this vector by a scalar elementwise using `x`, `y`, `z`, and `t` components - - In reality, this directly adjusts `pt`, `eta`, `phi` and `energy` for performance - """ - return awkward.zip( - { - "pt": self.pt / abs(other), - "eta": self.eta * numpy.sign(other), - "phi": self.phi % (2 * numpy.pi) - (numpy.pi * (other < 0)), - "energy": self.energy / other, - }, - with_name="PtEtaPhiELorentzVector", - ) - @awkward.mixin_class_method(numpy.negative) def negative(self): """Returns the negative of the vector"""