Vector: arrays of 2D, 3D, and Lorentz vectors
Vector is a Python 3.8+ library (Python 3.6 and 3.7 supported till v0.9.0
and v1.0.0
, respectively) for 2D, 3D, and Lorentz vectors, especially arrays of vectors, to solve common physics problems in a NumPy-like way.
Main features of Vector:
- Pure Python with NumPy as its only dependency. This makes it easier to install.
- Vectors may be represented in a variety of coordinate systems: Cartesian, cylindrical, pseudorapidity, and any combination of these with time or proper time for Lorentz vectors. In all, there are 12 coordinate systems: {x-y vs ρ-φ in the azimuthal plane} × {z vs θ vs η longitudinally} × {t vs τ temporally}.
- Uses names and conventions set by ROOT's TLorentzVector and Math::LorentzVector, as well as scikit-hep/math, uproot-methods TLorentzVector, henryiii/hepvector, and coffea.nanoevents.methods.vector.
- Implemented on a variety of backends:
- Awkward backend also implemented in Numba for JIT-compiled calculations on vectors.
- JAX and Dask support through Awkward Arrays.
- Distinction between geometrical vectors, which have a minimum of attribute and method names, and vectors representing momentum, which have synonyms like
pt
= rho
, energy
= t
, mass
= tau
.
Installation
To install, use pip install vector
or your favorite way to install in an environment.
Overview
This overview is based on the documentation here.
import vector
import numpy as np
import awkward as ak
import numba as nb
Constructing a vector or an array of vectors
The easiest way to create one or many vectors is with a helper function:
vector.obj
to make a pure Python vector object,vector.arr
to make a NumPy array of vectors (or array
, lowercase, like np.array
),vector.awk
to make an Awkward Array of vectors (or Array
, uppercase, like ak.Array
).vector.zip
to make an Awkward Array of vectors (similar to ak.zip
)
Pure Python vectors
You can directly use the VectorObject
and MomentumObject
classes to construct object type vectors:
vector.VectorObject2D(x=1.1, y=2.2)
vector.MomentumObject3D(px=1.1, py=2.2, pz=3.3)
vector.VectorObject4D(x=1.1, y=2.2, eta=3.3, tau=4.4)
and so on for every class.
Or, you can use a single wrapper function to construct all possible combinations of
object type vectors:
vector.obj(x=3, y=4)
vector.obj(rho=5, phi=0.9273)
vector.obj(x=3, y=4).isclose(vector.obj(rho=5, phi=0.9273))
vector.obj(x=3, y=4, z=-2)
vector.obj(x=3, y=4, z=-2, t=10)
vector.obj(rho=5, phi=0.9273, eta=-0.39, t=10)
vector.obj(pt=5, phi=0.9273, eta=-0.39, E=10)
vector.obj(rho=5, phi=0.9273, eta=-0.39, t=10) == vector.obj(
pt=5, phi=0.9273, eta=-0.390035, E=10
)
vector.obj(rho=5, phi=0.9273, eta=-0.39, t=10).tau
vector.obj(pt=5, phi=0.9273, eta=-0.39, E=10).mass
vector.obj(pt=5, phi=0.9273, theta=1.9513, mass=8.4262)
vector.obj(x=3, y=4, z=-2, t=10).isclose(
vector.obj(pt=5, phi=0.9273, theta=1.9513, mass=8.4262)
)
(
isinstance(vector.obj(x=1.1, y=2.2), vector.Vector),
isinstance(vector.obj(x=1.1, y=2.2), vector.Vector2D),
isinstance(vector.obj(x=1.1, y=2.2), vector.VectorObject),
isinstance(vector.obj(px=1.1, py=2.2), vector.Momentum),
isinstance(vector.obj(x=1.1, y=2.2, z=3.3, t=4.4), vector.Planar),
isinstance(vector.obj(x=1.1, y=2.2, z=3.3, t=4.4), vector.Spatial),
isinstance(vector.obj(x=1.1, y=2.2, z=3.3, t=4.4), vector.Lorentz),
isinstance(vector.obj(x=1.1, y=2.2, z=3.3, t=4.4).azimuthal, vector.AzimuthalXY),
isinstance(
vector.obj(x=1.1, y=2.2, z=3.3, t=4.4).longitudinal, vector.LongitudinalZ
),
isinstance(vector.obj(x=1.1, y=2.2, z=3.3, t=4.4).temporal, vector.TemporalT),
)
The allowed keyword arguments for 2D vectors are:
x
and y
for Cartesian azimuthal coordinates,px
and py
for momentum,rho
and phi
for polar azimuthal coordinates,pt
and phi
for momentum.
For 3D vectors, you need the above and:
z
for the Cartesian longitudinal coordinate,pz
for momentum,theta
for the spherical polar angle (from $0$ to $\pi$, inclusive),eta
for pseudorapidity, which is a kind of spherical polar angle.
For 4D vectors, you need the above and:
t
for the Cartesian temporal coordinate,E
or energy
to get four-momentum,tau
for the "proper time" (temporal coordinate in the vector's rest coordinate system),M
or mass
to get four-momentum.
Since momentum vectors have momentum-synonyms in addition to the geometrical names, any momentum-synonym will make the whole vector a momentum vector.
If you want to bypass the dimension and coordinate system inference through keyword arguments (e.g. for static typing), you can use specialized constructors:
vector.VectorObject2D.from_xy(1.1, 2.2)
vector.MomentumObject3D.from_rhophiz(1.1, 2.2, 3.3)
vector.VectorObject4D.from_xyetatau(1.1, 2.2, 3.3, 4.4)
and so on, for all combinations of azimuthal, longitudinal, and temporal coordinates, geometric and momentum-flavored.
SymPy vectors
Note: Operations on SymPy vectors are only 100% compatible with numeric vectors (Python, NumPy, and Awkward backends) if the vectors are positive time-like, that is, if t**2 > x**2 + y**2 + z**2
. The space-like and negative time-like cases have different sign conventions.
You can directly use the VectorSympy
and MomentumSympy
classes to construct object type vectors:
import sympy
x, y, z, t, px, py, pz, eta, tau = sympy.symbols(
"x y z t px py pz eta tau",
real=True,
)
vector.VectorSympy2D(x=x, y=y)
vector.MomentumSympy3D(px=px, py=py, pz=pz)
vector.VectorSympy4D(x=x, y=y, eta=eta, tau=tau)
and so on for every class.
(
isinstance(vector.VectorSympy2D(x=x, y=y), vector.Vector),
isinstance(vector.VectorSympy2D(x=x, y=y), vector.Vector2D),
isinstance(vector.VectorSympy2D(x=x, y=y), vector.VectorSympy),
isinstance(vector.MomentumSympy2D(px=px, py=py), vector.Momentum),
isinstance(vector.VectorSympy4D(x=x, y=y, z=z, t=t), vector.Planar),
isinstance(vector.VectorSympy4D(x=x, y=y, z=z, t=t), vector.Spatial),
isinstance(vector.VectorSympy4D(x=x, y=y, z=z, t=t), vector.Lorentz),
isinstance(vector.VectorSympy4D(x=x, y=y, z=z, t=t).azimuthal, vector.AzimuthalXY),
isinstance(
vector.VectorSympy4D(x=x, y=y, z=z, t=t).longitudinal, vector.LongitudinalZ
),
isinstance(vector.VectorSympy4D(x=x, y=y, z=z, t=t).temporal, vector.TemporalT),
)
Since VectorSympy2D
, VectorSympy3D
, VectorSympy4D
, and their momentum equivalents operate on SymPy expressions, all of the normal SymPy methods and functions work on the results, coordinates, and the vectors.
sympy.init_session()
v1 = vector.VectorSympy2D(x=x, y=y)
sympy.Eq(v1.rho, sympy.sqrt(x**2 + y**2))
v2 = vector.VectorSympy4D(x=x, y=y, z=z, t=t)
v2.to_rhophithetatau().tau
values = {x: 3, y: 2, z: 1, t: 10}
v2.is_timelike()
v2.is_timelike().subs(values)
v2.to_rhophithetatau().tau.subs(values).evalf()
v2.boost(v2.to_beta3())
v2.boost(v2.to_beta3()).t
v2.boost(v2.to_beta3()).t.simplify()
v2.boost(v2.to_beta3()).t.subs(values)
v2.boost(v2.to_beta3()).t.subs(values).evalf()
All of the keyword arguments and rules that apply to vector.obj
construction apply to vector.VectorSympyND
and vector.MomentumObjectND
objects.
NumPy arrays of vectors
You can directly use the VectorNumpy
classes to construct object type vectors:
vector.VectorNumpy2D(
[(1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5)],
dtype=[("x", float), ("y", float)],
)
vector.VectorNumpy2D({"x": [1.1, 1.2, 1.3, 1.4, 1.5], "y": [2.1, 2.2, 2.3, 2.4, 2.5]})
vector.VectorNumpy4D(
{
"x": [1.1, 1.2, 1.3, 1.4, 1.5],
"y": [2.1, 2.2, 2.3, 2.4, 2.5],
"z": [3.1, 3.2, 3.3, 3.4, 3.5],
"t": [4.1, 4.2, 4.3, 4.4, 4.5],
}
)
and so on for every class.
Or, you can use a single wrapper function to construct all possible combinations of
NumPy type vectors:
vector.array(
[(1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5)],
dtype=[("x", float), ("y", float)],
)
vector.array({"x": [1.1, 1.2, 1.3, 1.4, 1.5], "y": [2.1, 2.2, 2.3, 2.4, 2.5]})
vector.array(
{
"x": [1.1, 1.2, 1.3, 1.4, 1.5],
"y": [2.1, 2.2, 2.3, 2.4, 2.5],
"z": [3.1, 3.2, 3.3, 3.4, 3.5],
"t": [4.1, 4.2, 4.3, 4.4, 4.5],
}
)
vector.array(
{
"pt": [1.1, 1.2, 1.3, 1.4, 1.5],
"phi": [2.1, 2.2, 2.3, 2.4, 2.5],
"eta": [3.1, 3.2, 3.3, 3.4, 3.5],
"M": [4.1, 4.2, 4.3, 4.4, 4.5],
}
)
Existing NumPy arrays can be viewed as arrays of vectors, but it needs to be a structured array with recognized field names.
np.arange(0, 24, 0.1).view(
[
("x", float),
("y", float),
("z", float),
("t", float),
]
).view(
vector.VectorNumpy4D
)
Since VectorNumpy2D
, VectorNumpy3D
, VectorNumpy4D
, and their momentum equivalents are NumPy array subclasses, all of the normal NumPy methods and functions work on them.
np.arange(0, 24, 0.1).view(
[("x", float), ("y", float), ("z", float), ("t", float)]
).view(vector.VectorNumpy4D).reshape(6, 5, 2)
All of the keyword arguments and rules that apply to vector.obj
construction apply to vector.arr
dtypes.
Geometrical names are used in the dtype, even if momentum-synonyms are used in construction.
vector.arr({"px": [1, 2, 3, 4], "py": [1.1, 2.2, 3.3, 4.4], "pz": [0.1, 0.2, 0.3, 0.4]})
Awkward Arrays of vectors
Awkward Arrays are arrays with more complex data structures than NumPy allows, such as variable-length lists, nested records, missing and even heterogeneous data (multiple data types: use sparingly).
The vector.awk
function behaves exactly like the ak.Array constructor, except that it makes arrays of vectors.
vector.awk(
[
[{"x": 1, "y": 1.1, "z": 0.1}, {"x": 2, "y": 2.2, "z": 0.2}],
[],
[{"x": 3, "y": 3.3, "z": 0.3}],
[
{"x": 4, "y": 4.4, "z": 0.4},
{"x": 5, "y": 5.5, "z": 0.5},
{"x": 6, "y": 6.6, "z": 0.6},
],
]
)
If you want any records named "Vector2D
", "Vector3D
", "Vector4D
", "Momentum2D
", "Momentum3D
", or "Momentum4D
" to be interpreted as vectors, register the behaviors globally.
vector.register_awkward()
ak.Array(
[
[{"x": 1, "y": 1.1, "z": 0.1}, {"x": 2, "y": 2.2, "z": 0.2}],
[],
[{"x": 3, "y": 3.3, "z": 0.3}],
[
{"x": 4, "y": 4.4, "z": 0.4},
{"x": 5, "y": 5.5, "z": 0.5},
{"x": 6, "y": 6.6, "z": 0.6},
],
],
with_name="Vector3D",
)
All of the keyword arguments and rules that apply to vector.obj
construction apply to vector.awk
field names.
Finally, the VectorAwkward
mixins can be subclassed to create custom vector classes. The awkward behavior classes and projections must be named as *Array
. For example, coffea
uses the following names - TwoVectorArray
, ThreeVectorArray
, PolarTwoVectorArray
, SphericalThreeVectorArray
, ...
Vector properties
Any geometrical coordinate can be computed from vectors in any coordinate system; they'll be provided or computed as needed.
vector.obj(x=3, y=4).rho
vector.obj(rho=5, phi=0.9273).x
vector.obj(rho=5, phi=0.9273).y
vector.obj(x=1, y=2, z=3).theta
vector.obj(x=1, y=2, z=3).eta
Some properties are not coordinates, but derived from them.
vector.obj(x=1, y=2, z=3).costheta
vector.obj(x=1, y=2, z=3).mag
vector.obj(x=1, y=2, z=3).mag2
These properties are provided because they can be computed faster or with more numerical stability in different coordinate systems. For instance, the magnitude ignores phi
in polar coordinates.
vector.obj(rho=3, phi=0.123456789, z=4).mag2
Momentum vectors have geometrical properties as well as their momentum-synonyms.
vector.obj(px=3, py=4).rho
vector.obj(px=3, py=4).pt
vector.obj(x=1, y=2, z=3, E=4).tau
vector.obj(x=1, y=2, z=3, E=4).mass
Here's the key thing: arrays of vectors return arrays of coordinates.
vector.arr(
{
"x": [1.0, 2.0, 3.0, 4.0, 5.0],
"y": [1.1, 2.2, 3.3, 4.4, 5.5],
"z": [0.1, 0.2, 0.3, 0.4, 0.5],
}
).theta
vector.awk(
[
[{"x": 1, "y": 1.1, "z": 0.1}, {"x": 2, "y": 2.2, "z": 0.2}],
[],
[{"x": 3, "y": 3.3, "z": 0.3}],
[{"x": 4, "y": 4.4, "z": 0.4}, {"x": 5, "y": 5.5, "z": 0.5}],
]
).theta
array = (
np.random.normal(0, 1, 150)
.view([(x, float) for x in ("x", "y", "z")])
.view(vector.MomentumNumpy3D)
.reshape(5, 5, 2)
)
array.pt
array.shape
array.pt.shape
array = vector.awk(
[
[
{x: np.random.normal(0, 1) for x in ("px", "py", "pz")}
for inner in range(np.random.poisson(1.5))
]
for outer in range(50)
]
)
array.pt
ak.num(array)
ak.num(array.pt)
Vector methods
Vector methods require arguments (in parentheses), which may be scalars or other vectors, depending on the calculation.
vector.obj(x=3, y=4).rotateZ(0.1)
vector.obj(rho=5, phi=0.4).rotateZ(0.1)
print(
vector.arr({"rho": [1, 2, 3, 4, 5], "phi": [0.1, 0.2, 0.3, 0.4, 0.5]}).rotateZ(0.5)
)
print(
vector.arr({"rho": [1, 2, 3, 4, 5], "phi": [0.1, 0.2, 0.3, 0.4, 0.5]}).rotateZ(
np.array([0.1, 0.2, 0.3, 0.4, 0.5])
)
)
print(
vector.awk(
[[{"rho": 1, "phi": 0.1}, {"rho": 2, "phi": 0.2}], [], [{"rho": 3, "phi": 0.3}]]
).rotateZ(0.5)
)
print(
vector.awk(
[[{"rho": 1, "phi": 0.1}, {"rho": 2, "phi": 0.2}], [], [{"rho": 3, "phi": 0.3}]]
).rotateZ([0.1, 0.2, 0.3])
)
print(
vector.awk(
[[{"rho": 1, "phi": 0.1}, {"rho": 2, "phi": 0.2}], [], [{"rho": 3, "phi": 0.3}]]
).rotateZ([[0.1, 0.2], [], [0.3]])
)
Some methods are equivalent to binary operators.
vector.obj(x=3, y=4).scale(10)
vector.obj(x=3, y=4) * 10
10 * vector.obj(x=3, y=4)
vector.obj(rho=5, phi=0.5) * 10
Some methods involve more than one vector.
vector.obj(x=1, y=2).add(vector.obj(x=5, y=5))
vector.obj(x=1, y=2) + vector.obj(x=5, y=5)
vector.obj(x=1, y=2).dot(vector.obj(x=5, y=5))
vector.obj(x=1, y=2) @ vector.obj(x=5, y=5)
The vectors can use different coordinate systems. Conversions are necessary, but minimized for speed and numeric stability.
vector.obj(x=3, y=4) @ vector.obj(x=6, y=8)
vector.obj(rho=5, phi=0.9273) @ vector.obj(x=6, y=8)
vector.obj(x=3, y=4) @ vector.obj(rho=10, phi=0.9273)
vector.obj(rho=5, phi=0.9273) @ vector.obj(rho=10, phi=0.9273)
In Python, some "operators" are actually built-in functions, such as abs
.
abs(vector.obj(x=3, y=4))
Note that abs
returns
rho
for 2D vectorsmag
for 3D vectorstau
(mass
) for 4D vectors
Use the named properties when you want magnitude in a specific number of dimensions; use abs
when you want the magnitude for any number of dimensions.
The vectors can be from different backends. Normal rules for broadcasting Python numbers, NumPy arrays, and Awkward Arrays apply.
vector.arr({"x": [1, 2, 3, 4, 5], "y": [0.1, 0.2, 0.3, 0.4, 0.5]}) + vector.obj(
x=10, y=5
)
(
vector.awk(
[
[{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}],
[],
[{"x": 3, "y": 3.3}],
[{"x": 4, "y": 4.4}, {"x": 5, "y": 5.5}],
]
)
+ vector.obj(x=10, y=5)
)
(
vector.awk(
[
[{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}],
[],
[{"x": 3, "y": 3.3}],
[{"x": 4, "y": 4.4}, {"x": 5, "y": 5.5}],
]
)
+ vector.arr(
{"x": [4, 3, 2, 1], "y": [0.1, 0.1, 0.1, 0.1]}
)
)
Some operations are defined for 2D or 3D vectors, but are usable on higher-dimensional vectors because the additional components can be ignored or are passed through unaffected.
vector.obj(rho=1, phi=0.5).deltaphi(vector.obj(rho=2, phi=0.3))
vector.obj(rho=1, phi=0.5, z=10).deltaphi(vector.obj(rho=2, phi=0.3, theta=1.4))
vector.obj(rho=1, phi=0.5, z=10, t=100).deltaphi(
vector.obj(rho=2, phi=0.3, theta=1.4, tau=1000)
)
vector.obj(rho=1, phi=0.5).deltaphi(vector.obj(rho=2, phi=0.3, theta=1.4, tau=1000))
This is especially useful for giving 4D vectors all the capabilities of 3D vectors.
vector.obj(x=1, y=2, z=3).rotateX(np.pi / 4)
vector.obj(x=1, y=2, z=3, tau=10).rotateX(np.pi / 4)
vector.obj(pt=1, phi=1.3, eta=2).deltaR(vector.obj(pt=2, phi=0.3, eta=1))
vector.obj(pt=1, phi=1.3, eta=2, mass=5).deltaR(
vector.obj(pt=2, phi=0.3, eta=1, mass=10)
)
For a few operations - +
, -
, ==
, !=
, ... - the dimension of the vectors should be equal. This can be achieved by using the like
method, to_{coordinate_name}
methods, to_Vector*D
methods. The to_Vector*D
methods provide more flexibility to the users, that is, new coordinate values can be passed into the methods as named arguments.
v1 = vector.obj(x=1, y=2, z=3)
v2 = vector.obj(x=1, y=2)
v1 - v2.like(v1)
v1.like(v2) - v2
v1 - v2.to_xyz()
v1.to_xy() - v2
v1 - v2.to_Vector3D(z=3)
v1.to_Vector2D() - v2
Similarly, for a few vector methods, the dimension of the input vectors are type checked strictly.
For instance, a cross-product is only defined for 3D and 7D vectors; hence, running the method on a 4D vector will error out.
vector.obj(x=0.1, y=0.2, z=0.3).cross(vector.obj(x=0.4, y=0.5, z=0.6))
The (current) list of properties and methods is:
Planar (2D, 3D, 4D):
x
(px
)y
(py
)rho
(pt
): two-dimensional magnituderho2
(pt2
): two-dimensional magnitude squaredphi
deltaphi(vector)
: difference in phi
(signed and rectified to $-\pi$ through $\pi$)rotateZ(angle)
transform2D(obj)
: the obj
must supply components through obj["xx"]
, obj["xy"]
, obj["yx"]
, obj["yy"]
is_parallel(vector, tolerance=1e-5)
: only true if they're pointing in the same directionis_antiparallel(vector, tolerance=1e-5)
: only true if they're pointing in opposite directionsis_perpendicular(vector, tolerance=1e-5)
Spatial (3D, 4D):
z
(pz
)theta
eta
costheta
cottheta
mag
(p
): three-dimensional magnitude, does not include temporal componentmag2
(p2
): three-dimensional magnitude squaredcross
: cross-product (strictly 3D)deltaangle(vector)
: difference in angle (always non-negative)deltaeta(vector)
: difference in eta
(signed)deltaR(vector)
: $\Delta R = \sqrt{\Delta\phi^2 + \Delta\eta^2}$deltaR2(vector)
: the above, squaredrotateX(angle)
rotateY(angle)
rotate_axis(axis, angle)
: the magnitude of axis
is ignored, but it must be at least 3Drotate_euler(phi, theta, psi, order="zxz")
: the arguments are in the same order as ROOT::Math::EulerAngles, and order="zxz"
agrees with ROOT's choice of conventionsrotate_nautical(yaw, pitch, roll)
rotate_quaternion(u, i, j, k)
: again, the conventions match ROOT::Math::Quaternion.transform3D(obj)
: the obj
must supply components through obj["xx"]
, obj["xy"]
, etc.is_parallel(vector, tolerance=1e-5)
: only true if they're pointing in the same directionis_antiparallel(vector, tolerance=1e-5)
: only true if they're pointing in opposite directionsis_perpendicular(vector, tolerance=1e-5)
Lorentz (4D only):
t
(E
, energy
): follows the ROOT::Math::LorentzVector behavior of treating spacelike vectors as negative t
and negative tau
and truncating wrong-direction timelike vectorst2
(E2
, energy2
)tau
(M
, mass
): see note abovetau2
(M2
, mass2
)beta
: scalar(s) between $0$ (inclusive) and $1$ (exclusive, unless the vector components are infinite)deltaRapidityPhi
: $\Delta R_{\mbox{rapidity}} = \Delta\phi^2 + \Delta \mbox{rapidity}^2$deltaRapidityPhi2
: the above, squaredgamma
: scalar(s) between $1$ (inclusive) and $\infty$rapidity
: scalar(s) between $0$ (inclusive) and $\infty$boost_p4(four_vector)
: change coordinate system using another 4D vector as the differenceboost_beta(three_vector)
: change coordinate system using a 3D beta vector (all components between $-1$ and $+1$)boost(vector)
: uses the dimension of the given vector
to determine behaviorboostX(beta=None, gamma=None)
: supply beta
xor gamma
, but not bothboostY(beta=None, gamma=None)
: supply beta
xor gamma
, but not bothboostZ(beta=None, gamma=None)
: supply beta
xor gamma
, but not bothtransform4D(obj)
: the obj
must supply components through obj["xx"]
, obj["xy"]
, etc.to_beta3()
: turns a four_vector
(for boost_p4
) into a three_vector
(for boost_beta3
)is_timelike(tolerance=0)
is_spacelike(tolerance=0)
is_lightlike(tolerance=1e-5)
: note the different tolerance
All numbers of dimensions:
unit()
: note the parenthesesdot(vector)
: can also use the @
operatoradd(vector)
: can also use the +
operatorsubtract(vector)
: can also use the -
operatorscale(factor)
: can also use the *
operatorequal(vector)
: can also use the ==
operator, but consider isclose
insteadnot_equal(vector)
: can also use the !=
operator, but consider isclose
insteadsum()
: can also use the numpy.sum
or awkward.sum
, only for NumPy and Awkward vectorscount_nonzero()
: can also use numpy.count_nonzero
or awkward.count_nonzero
, only for NumPy and Awkward vectorscount()
: can also use awkward.count
, only for Awkward vectorsisclose(vector, rtol=1e-5, atol=1e-8, equal_nan=False)
: works like np.isclose; arrays also have an allclose methodto_VectorND(coordinates)
/to_ND(coordinates)
: replace N
with the required vector dimensionto_{coordinate-names}
: for example - to_rhophietatau
like(other)
: projects the vector into the dimensions of other
, for example - two_d_vector.like(three_d_vector)
Compiling your Python with Numba
Numba is a just-in-time (JIT) compiler for a mathematically relevant subset of NumPy and Python. It allows you to write fast code without leaving the Python environment. The drawback of Numba is that it can only compile code blocks involving objects and functions that it recognizes.
The Vector library includes extensions to inform Numba about vector objects, vector NumPy arrays, and vector Awkward Arrays. At the time of writing, the implementation of vector NumPy arrays is incomplete due to numba/numba#6148.
For instance, consider the following function:
@nb.njit
def compute_mass(v1, v2):
return (v1 + v2).mass
compute_mass(vector.obj(px=1, py=2, pz=3, E=4), vector.obj(px=-1, py=-2, pz=-3, E=4))
When the two MomentumObject4D
objects are passed as arguments, Numba recognizes them and replaces the Python objects with low-level structs. When it compiles the function, it recognizes +
as the 4D add
function and recognizes .mass
as the tau
component of the result.
Although this demonstrates that Numba can manipulate vector objects, there is no performance advantage (and a likely disadvantage) to compiling a calculation on just a few vectors. The advantage comes when many vectors are involved, in arrays.
array = vector.awk(
[
[
dict(
{x: np.random.normal(0, 1) for x in ("px", "py", "pz")},
E=np.random.normal(10, 1),
)
for inner in range(np.random.poisson(1.5))
]
for outer in range(50)
]
)
@nb.njit
def compute_masses(array):
out = np.empty(len(array), np.float64)
for i, event in enumerate(array):
total = vector.obj(px=0.0, py=0.0, pz=0.0, E=0.0)
for vec in event:
total = total + vec
out[i] = total.mass
return out
compute_masses(array)
Talks about vector
Status as of November 17, 2023
First major release of vector is out and the package has reached a stable position. The work is spearheaded by bug reports and feature requests created on GitHub. It can only be improved by your feedback!
Contributors ✨
Thanks goes to these wonderful people (emoji key):
This project follows the
all-contributors
specification. Contributions of any kind welcome! See
CONTRIBUTING.md for information on setting up a
development environment.
Acknowledgements
This library was primarily developed by Jim Pivarski, Henry Schreiner, and Eduardo Rodrigues.
Support for this work was provided by the National Science Foundation cooperative agreement OAC-1836650 (IRIS-HEP) and OAC-1450377 (DIANA/HEP). Any opinions, findings, conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the National Science Foundation.