diff --git a/coffea/nanoevents/methods/vector.py b/coffea/nanoevents/methods/vector.py index a34c527d6..ee0c11ef7 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,13 @@ 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 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 @@ -143,6 +166,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 in the same direction""" + return self / self.r + @awkward.mixin_class(behavior) class PolarTwoVector(TwoVector): @@ -189,6 +221,28 @@ 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.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 +296,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 +312,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 +341,26 @@ def multiply(self, 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 in the same direction""" + return self / self.rho + @awkward.mixin_class(behavior) class SphericalThreeVector(ThreeVector, PolarTwoVector): @@ -322,6 +412,33 @@ 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.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 +496,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( @@ -417,6 +547,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. 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 + + out = awkward.zip( + { + "x": self.x + v.x, + "y": self.y + v.y, + "z": self.z + v.z, + "t": gamma * (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 +783,28 @@ 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 """ + absother = abs(other) return awkward.zip( { - "pt": self.pt * other, - "eta": self.eta, - "phi": self.phi, - "mass": self.mass * other, + "pt": self.pt * absother, + "eta": self.eta * numpy.sign(other), + "phi": self.phi % (2 * numpy.pi) - (numpy.pi * (other < 0)), + "mass": self.mass * absother, + }, + 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 +893,31 @@ 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.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)