Skip to content

Commit

Permalink
Merge pull request #447 from jrueb/extendvector
Browse files Browse the repository at this point in the history 
Extend Nanoevents vector and add tests
  • Loading branch information
@nsmith-
nsmith- authored Feb 12, 2021
2 parents 5465597 + c8b4e86 commit aa0b2fb
Show file tree
Hide file tree
Showing 2 changed files with 605 additions and 9 deletions.
244 changes: 235 additions & 9 deletions coffea/nanoevents/methods/vector.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"""
Expand All @@ -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(
Expand All @@ -136,13 +152,29 @@ 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
Returns a value within [-pi, pi)
"""
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):
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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"""
Expand All @@ -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(
Expand All @@ -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):
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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
):
Expand Down Expand Up @@ -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",
)
Expand Down Expand Up @@ -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",
Expand Down
Loading

0 comments on commit aa0b2fb

Please sign in to comment.