Introduction to Vector#

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, spherical, and any combination of these with time or proper time for Lorentz vectors. In all, there are 12 coordinate systems: {$x$-$y$ vs $\rho$-$\phi$ in the azimuthal plane} × {$z$ vs $\theta$ vs $\eta$ longitudinally} × {$t$ vs $\tau$ 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:
- pure Python objects
- SymPy vectors
- NumPy arrays of vectors (as a structured array subclass)
- Awkward Arrays of vectors
- potential for more: CuPy, TensorFlow, Torch…
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.

This notebook requires Vector, NumPy, Awkward Array, SymPy, and Numba to run.

[1]:

import awkward as ak  # at least version 1.2.0rc5
import numba as nb
import numpy as np
import sympy

import vector

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 (lowercase, like np.array),
vector.awk to make an Awkward Array of vectors (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 classes to construct object type vectors:

[2]:

vector.VectorObject2D(x=1.1, y=2.2)

[2]:

VectorObject2D(x=1.1, y=2.2)

[3]:

vector.MomentumObject3D(px=1.1, py=2.2, pz=3.3)

[3]:

MomentumObject3D(px=1.1, py=2.2, pz=3.3)

[4]:

vector.VectorObject4D(x=1.1, y=2.2, eta=3.3, tau=4.4)

[4]:

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:

[5]:

vector.obj(x=3, y=4)  # Cartesian 2D vector

[5]:

VectorObject2D(x=3, y=4)

[6]:

vector.obj(rho=5, phi=0.9273)  # same in polar coordinates

[6]:

VectorObject2D(rho=5, phi=0.9273)

[7]:

vector.obj(x=3, y=4).isclose(
    vector.obj(rho=5, phi=0.9273)
)  # use "isclose" unless they are exactly equal

[7]:

True

[8]:

vector.obj(x=3, y=4, z=-2)  # Cartesian 3D vector

[8]:

VectorObject3D(x=3, y=4, z=-2)

[9]:

vector.obj(x=3, y=4, z=-2, t=10)  # Cartesian 4D vector

[9]:

VectorObject4D(x=3, y=4, z=-2, t=10)

[10]:

vector.obj(
    rho=5, phi=0.9273, eta=-0.39, t=10
)  # in rho-phi-eta-t cylindrical coordinates

[10]:

VectorObject4D(rho=5, phi=0.9273, eta=-0.39, t=10)

[11]:

vector.obj(
    pt=5, phi=0.9273, eta=-0.39, E=10
)  # use momentum-synonyms to get a momentum vector

[11]:

MomentumObject4D(pt=5, phi=0.9273, eta=-0.39, E=10)

[12]:

vector.obj(rho=5, phi=0.9273, eta=-0.39, t=10) == vector.obj(
    pt=5, phi=0.9273, eta=-0.390035, E=10
)

[12]:

False

[13]:

vector.obj(
    rho=5, phi=0.9273, eta=-0.39, t=10
).tau  # geometrical vectors have to use geometrical names ("tau", not "mass")

[13]:

8.426194916448265

[14]:

vector.obj(
    pt=5, phi=0.9273, eta=-0.39, E=10
).mass  # momentum vectors can use momentum names (as well as geometrical ones)

[14]:

8.426194916448265

[15]:

vector.obj(
    pt=5, phi=0.9273, theta=1.9513, mass=8.4262
)  # any combination of azimuthal, longitudinal, and temporal coordinates is allowed

[15]:

MomentumObject4D(pt=5, phi=0.9273, theta=1.9513, mass=8.4262)

[16]:

vector.obj(x=3, y=4, z=-2, t=10).isclose(
    vector.obj(pt=5, phi=0.9273, theta=1.9513, mass=8.4262)
)

[16]:

True

[17]:

# Test instance type for any level of granularity.
(
    isinstance(
        vector.obj(x=1.1, y=2.2), vector.Vector
    ),  # is a vector or array of vectors
    isinstance(vector.obj(x=1.1, y=2.2), vector.Vector2D),  # is 2D (not 3D or 4D)
    isinstance(
        vector.obj(x=1.1, y=2.2), vector.VectorObject
    ),  # is a vector object (not an array)
    isinstance(vector.obj(px=1.1, py=2.2), vector.Momentum),  # has momentum synonyms
    isinstance(
        vector.obj(x=1.1, y=2.2, z=3.3, t=4.4), vector.Planar
    ),  # has transverse plane (2D, 3D, or 4D)
    isinstance(
        vector.obj(x=1.1, y=2.2, z=3.3, t=4.4), vector.Spatial
    ),  # has all spatial coordinates (3D or 4D)
    isinstance(
        vector.obj(x=1.1, y=2.2, z=3.3, t=4.4), vector.Lorentz
    ),  # has temporal coordinates (4D)
    isinstance(
        vector.obj(x=1.1, y=2.2, z=3.3, t=4.4).azimuthal, vector.AzimuthalXY
    ),  # azimuthal coordinate type
    isinstance(
        vector.obj(x=1.1, y=2.2, z=3.3, t=4.4).longitudinal, vector.LongitudinalZ
    ),  # longitudinal coordinate type
    isinstance(
        vector.obj(x=1.1, y=2.2, z=3.3, t=4.4).temporal, vector.TemporalT
    ),  # temporal coordinate type
)

[17]:

(True, True, True, True, True, True, True, True, True, True)

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:

[18]:

vector.VectorObject2D.from_xy(1.1, 2.2)

[18]:

VectorObject2D(x=1.1, y=2.2)

[19]:

vector.MomentumObject3D.from_rhophiz(1.1, 2.2, 3.3)

[19]:

MomentumObject3D(pt=1.1, phi=2.2, pz=3.3)

[20]:

vector.VectorObject4D.from_xyetatau(1.1, 2.2, 3.3, 4.4)

[20]:

VectorObject4D(x=1.1, y=2.2, eta=3.3, tau=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:

[21]:

x, y, z, t, px, py, pz, eta, tau = sympy.symbols(
    "x y z t px py pz eta tau",
    real=True,  # see sympy assumptions to add more restrictions on the symbols
)

[22]:

vector.VectorSympy2D(x=x, y=y)

[22]:

VectorSympy2D(x=x, y=y)

[23]:

vector.MomentumSympy3D(px=px, py=py, pz=pz)

[23]:

MomentumSympy3D(px=px, py=py, pz=pz)

[24]:

vector.VectorSympy4D(x=x, y=y, eta=eta, tau=tau)

[24]:

VectorSympy4D(x=x, y=y, eta=eta, tau=tau)

and so on for every class.

[25]:

# Test instance type for any level of granularity.
(
    # is a vector or array of vectors
    isinstance(vector.VectorSympy2D(x=x, y=y), vector.Vector),
    # is 2D (not 3D or 4D)
    isinstance(vector.VectorSympy2D(x=x, y=y), vector.Vector2D),
    # is a sympy vector (not an array)
    isinstance(vector.VectorSympy2D(x=x, y=y), vector.VectorSympy),
    # has momentum synonyms
    isinstance(vector.MomentumSympy2D(px=px, py=py), vector.Momentum),
    # has transverse plane (2D, 3D, or 4D)
    isinstance(vector.VectorSympy4D(x=x, y=y, z=z, t=t), vector.Planar),
    # has all spatial coordinates (3D or 4D)
    isinstance(vector.VectorSympy4D(x=x, y=y, z=z, t=t), vector.Spatial),
    # has temporal coordinates (4D)
    isinstance(vector.VectorSympy4D(x=x, y=y, z=z, t=t), vector.Lorentz),
    # azimuthal coordinate type
    isinstance(vector.VectorSympy4D(x=x, y=y, z=z, t=t).azimuthal, vector.AzimuthalXY),
    # longitudinal coordinate type
    isinstance(
        vector.VectorSympy4D(x=x, y=y, z=z, t=t).longitudinal, vector.LongitudinalZ
    ),
    # temporal coordinate type
    isinstance(vector.VectorSympy4D(x=x, y=y, z=z, t=t).temporal, vector.TemporalT),
)

[25]:

(True, True, True, True, True, True, True, True, True, True)

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.

[26]:

sympy.init_session()  # latex printing

IPython console for SymPy 1.12 (Python 3.11.5-64-bit) (ground types: python)

These commands were executed:
>>> from sympy import *
>>> x, y, z, t = symbols('x y z t')
>>> k, m, n = symbols('k m n', integer=True)
>>> f, g, h = symbols('f g h', cls=Function)
>>> init_printing()

Documentation can be found at https://docs.sympy.org/1.12/

[27]:

v1 = vector.VectorSympy2D(x=x, y=y)
sympy.Eq(v1.rho, sympy.sqrt(x**2 + y**2))

[27]:

$\displaystyle \text{True}$

[28]:

v2 = vector.VectorSympy4D(x=x, y=y, z=z, t=t)

[29]:

v2.to_rhophithetatau().tau

[29]:

$\displaystyle \sqrt{\left|{- t^{2} + x^{2} + y^{2} + z^{2}}\right|}$

[30]:

values = {x: 3, y: 2, z: 1, t: 10}  # t**2 > x**2 + y**2 + z**2

[31]:

v2.is_timelike()

[31]:

$\displaystyle t^{2} - x^{2} - y^{2} - z^{2} > 0$

[32]:

v2.is_timelike().subs(values)

[32]:

$\displaystyle \text{True}$

[33]:

v2.to_rhophithetatau().tau.subs(values).evalf()

[33]:

$\displaystyle 9.2736184954957$

[34]:

v2.boost(v2.to_beta3())

[34]:

VectorSympy4D(x=x*(1 + x**2/(t**2*(1 + 1/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))*(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))) + x/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2) + x*y**2/(t**2*(1 + 1/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))*(1 - x**2/t**2 - y**2/t**2 - z**2/t**2)) + x*z**2/(t**2*(1 + 1/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))*(1 - x**2/t**2 - y**2/t**2 - z**2/t**2)), y=y*(1 + y**2/(t**2*(1 + 1/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))*(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))) + y/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2) + x**2*y/(t**2*(1 + 1/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))*(1 - x**2/t**2 - y**2/t**2 - z**2/t**2)) + y*z**2/(t**2*(1 + 1/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))*(1 - x**2/t**2 - y**2/t**2 - z**2/t**2)), z=z*(1 + z**2/(t**2*(1 + 1/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))*(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))) + z/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2) + x**2*z/(t**2*(1 + 1/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))*(1 - x**2/t**2 - y**2/t**2 - z**2/t**2)) + y**2*z/(t**2*(1 + 1/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))*(1 - x**2/t**2 - y**2/t**2 - z**2/t**2)), t=t/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2) + x**2/(t*sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2)) + y**2/(t*sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2)) + z**2/(t*sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2)))

[35]:

v2.boost(v2.to_beta3()).t

[35]:

$\displaystyle \frac{t}{\sqrt{1 - \frac{x^{2}}{t^{2}} - \frac{y^{2}}{t^{2}} - \frac{z^{2}}{t^{2}}}} + \frac{x^{2}}{t \sqrt{1 - \frac{x^{2}}{t^{2}} - \frac{y^{2}}{t^{2}} - \frac{z^{2}}{t^{2}}}} + \frac{y^{2}}{t \sqrt{1 - \frac{x^{2}}{t^{2}} - \frac{y^{2}}{t^{2}} - \frac{z^{2}}{t^{2}}}} + \frac{z^{2}}{t \sqrt{1 - \frac{x^{2}}{t^{2}} - \frac{y^{2}}{t^{2}} - \frac{z^{2}}{t^{2}}}}$

[36]:

v2.boost(v2.to_beta3()).t.simplify()

[36]:

$\displaystyle \frac{t \sqrt{\frac{t^{2} - x^{2} - y^{2} - z^{2}}{t^{2}}} \left(t^{2} + x^{2} + y^{2} + z^{2}\right)}{t^{2} - x^{2} - y^{2} - z^{2}}$

[37]:

v2.boost(v2.to_beta3()).t.subs(values)

[37]:

$\displaystyle \frac{57 \sqrt{86}}{43}$

[38]:

v2.boost(v2.to_beta3()).t.subs(values).evalf()

[38]:

$\displaystyle 12.292936145192$

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:

[39]:

# NumPy-like arguments (literally passed through to NumPy)
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)],
)

[39]:

VectorNumpy2D([(1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5)],
              dtype=[('x', '<f8'), ('y', '<f8')])

[40]:

# Pandas-like arguments (dict from names to column arrays)
vector.VectorNumpy2D({"x": [1.1, 1.2, 1.3, 1.4, 1.5], "y": [2.1, 2.2, 2.3, 2.4, 2.5]})

[40]:

VectorNumpy2D([(1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5)],
              dtype=[('x', '<f8'), ('y', '<f8')])

[41]:

# As with objects, the coordinate system and dimension is taken from the names of the fields.
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],
    }
)

[41]:

VectorNumpy4D([(1.1, 2.1, 3.1, 4.1), (1.2, 2.2, 3.2, 4.2), (1.3, 2.3, 3.3, 4.3),
               (1.4, 2.4, 3.4, 4.4), (1.5, 2.5, 3.5, 4.5)],
              dtype=[('x', '<f8'), ('y', '<f8'), ('z', '<f8'), ('t', '<f8')])

and so on for every class.

Or, you can use a single wrapper function to construct all possible combinations of NumPy type vectors:

[42]:

# NumPy-like arguments (literally passed through to NumPy)
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)],
)

[42]:

VectorNumpy2D([(1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5)],
              dtype=[('x', '<f8'), ('y', '<f8')])

[43]:

# Pandas-like arguments (dict from names to column arrays)
vector.array({"x": [1.1, 1.2, 1.3, 1.4, 1.5], "y": [2.1, 2.2, 2.3, 2.4, 2.5]})

[43]:

VectorNumpy2D([(1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5)],
              dtype=[('x', '<f8'), ('y', '<f8')])

[44]:

# As with objects, the coordinate system and dimension is taken from the names of the fields.
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],
    }
)

[44]:

VectorNumpy4D([(1.1, 2.1, 3.1, 4.1), (1.2, 2.2, 3.2, 4.2), (1.3, 2.3, 3.3, 4.3),
               (1.4, 2.4, 3.4, 4.4), (1.5, 2.5, 3.5, 4.5)],
              dtype=[('x', '<f8'), ('y', '<f8'), ('z', '<f8'), ('t', '<f8')])

[45]:

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],
    }
)

[45]:

MomentumNumpy4D([(1.1, 2.1, 3.1, 4.1), (1.2, 2.2, 3.2, 4.2), (1.3, 2.3, 3.3, 4.3),
                 (1.4, 2.4, 3.4, 4.4), (1.5, 2.5, 3.5, 4.5)],
                dtype=[('rho', '<f8'), ('phi', '<f8'), ('eta', '<f8'), ('tau', '<f8')])

Existing NumPy arrays can be viewed as arrays of vectors, but it needs to be a structured array with recognized field names.

[46]:

# NumPy array         # interpret groups of four values as named fields              # give it vector properties and methods
np.arange(0, 24, 0.1).view(
    [("x", float), ("y", float), ("z", float), ("t", float)]
).view(vector.VectorNumpy4D)

[46]:

VectorNumpy4D([( 0. ,  0.1,  0.2,  0.3), ( 0.4,  0.5,  0.6,  0.7),
               ( 0.8,  0.9,  1. ,  1.1), ( 1.2,  1.3,  1.4,  1.5),
               ( 1.6,  1.7,  1.8,  1.9), ( 2. ,  2.1,  2.2,  2.3),
               ( 2.4,  2.5,  2.6,  2.7), ( 2.8,  2.9,  3. ,  3.1),
               ( 3.2,  3.3,  3.4,  3.5), ( 3.6,  3.7,  3.8,  3.9),
               ( 4. ,  4.1,  4.2,  4.3), ( 4.4,  4.5,  4.6,  4.7),
               ( 4.8,  4.9,  5. ,  5.1), ( 5.2,  5.3,  5.4,  5.5),
               ( 5.6,  5.7,  5.8,  5.9), ( 6. ,  6.1,  6.2,  6.3),
               ( 6.4,  6.5,  6.6,  6.7), ( 6.8,  6.9,  7. ,  7.1),
               ( 7.2,  7.3,  7.4,  7.5), ( 7.6,  7.7,  7.8,  7.9),
               ( 8. ,  8.1,  8.2,  8.3), ( 8.4,  8.5,  8.6,  8.7),
               ( 8.8,  8.9,  9. ,  9.1), ( 9.2,  9.3,  9.4,  9.5),
               ( 9.6,  9.7,  9.8,  9.9), (10. , 10.1, 10.2, 10.3),
               (10.4, 10.5, 10.6, 10.7), (10.8, 10.9, 11. , 11.1),
               (11.2, 11.3, 11.4, 11.5), (11.6, 11.7, 11.8, 11.9),
               (12. , 12.1, 12.2, 12.3), (12.4, 12.5, 12.6, 12.7),
               (12.8, 12.9, 13. , 13.1), (13.2, 13.3, 13.4, 13.5),
               (13.6, 13.7, 13.8, 13.9), (14. , 14.1, 14.2, 14.3),
               (14.4, 14.5, 14.6, 14.7), (14.8, 14.9, 15. , 15.1),
               (15.2, 15.3, 15.4, 15.5), (15.6, 15.7, 15.8, 15.9),
               (16. , 16.1, 16.2, 16.3), (16.4, 16.5, 16.6, 16.7),
               (16.8, 16.9, 17. , 17.1), (17.2, 17.3, 17.4, 17.5),
               (17.6, 17.7, 17.8, 17.9), (18. , 18.1, 18.2, 18.3),
               (18.4, 18.5, 18.6, 18.7), (18.8, 18.9, 19. , 19.1),
               (19.2, 19.3, 19.4, 19.5), (19.6, 19.7, 19.8, 19.9),
               (20. , 20.1, 20.2, 20.3), (20.4, 20.5, 20.6, 20.7),
               (20.8, 20.9, 21. , 21.1), (21.2, 21.3, 21.4, 21.5),
               (21.6, 21.7, 21.8, 21.9), (22. , 22.1, 22.2, 22.3),
               (22.4, 22.5, 22.6, 22.7), (22.8, 22.9, 23. , 23.1),
               (23.2, 23.3, 23.4, 23.5), (23.6, 23.7, 23.8, 23.9)],
              dtype=[('x', '<f8'), ('y', '<f8'), ('z', '<f8'), ('t', '<f8')])

Since VectorNumpy2D, VectorNumpy3D, VectorNumpy4D, and their momentum equivalents are NumPy array subclasses, all of the normal NumPy methods and functions work on them.

[47]:

np.arange(0, 24, 0.1).view(
    [("x", float), ("y", float), ("z", float), ("t", float)]
).view(vector.VectorNumpy4D).reshape(6, 5, 2)

[47]:

VectorNumpy4D([[[( 0. ,  0.1,  0.2,  0.3), ( 0.4,  0.5,  0.6,  0.7)],
                [( 0.8,  0.9,  1. ,  1.1), ( 1.2,  1.3,  1.4,  1.5)],
                [( 1.6,  1.7,  1.8,  1.9), ( 2. ,  2.1,  2.2,  2.3)],
                [( 2.4,  2.5,  2.6,  2.7), ( 2.8,  2.9,  3. ,  3.1)],
                [( 3.2,  3.3,  3.4,  3.5), ( 3.6,  3.7,  3.8,  3.9)]],

               [[( 4. ,  4.1,  4.2,  4.3), ( 4.4,  4.5,  4.6,  4.7)],
                [( 4.8,  4.9,  5. ,  5.1), ( 5.2,  5.3,  5.4,  5.5)],
                [( 5.6,  5.7,  5.8,  5.9), ( 6. ,  6.1,  6.2,  6.3)],
                [( 6.4,  6.5,  6.6,  6.7), ( 6.8,  6.9,  7. ,  7.1)],
                [( 7.2,  7.3,  7.4,  7.5), ( 7.6,  7.7,  7.8,  7.9)]],

               [[( 8. ,  8.1,  8.2,  8.3), ( 8.4,  8.5,  8.6,  8.7)],
                [( 8.8,  8.9,  9. ,  9.1), ( 9.2,  9.3,  9.4,  9.5)],
                [( 9.6,  9.7,  9.8,  9.9), (10. , 10.1, 10.2, 10.3)],
                [(10.4, 10.5, 10.6, 10.7), (10.8, 10.9, 11. , 11.1)],
                [(11.2, 11.3, 11.4, 11.5), (11.6, 11.7, 11.8, 11.9)]],

               [[(12. , 12.1, 12.2, 12.3), (12.4, 12.5, 12.6, 12.7)],
                [(12.8, 12.9, 13. , 13.1), (13.2, 13.3, 13.4, 13.5)],
                [(13.6, 13.7, 13.8, 13.9), (14. , 14.1, 14.2, 14.3)],
                [(14.4, 14.5, 14.6, 14.7), (14.8, 14.9, 15. , 15.1)],
                [(15.2, 15.3, 15.4, 15.5), (15.6, 15.7, 15.8, 15.9)]],

               [[(16. , 16.1, 16.2, 16.3), (16.4, 16.5, 16.6, 16.7)],
                [(16.8, 16.9, 17. , 17.1), (17.2, 17.3, 17.4, 17.5)],
                [(17.6, 17.7, 17.8, 17.9), (18. , 18.1, 18.2, 18.3)],
                [(18.4, 18.5, 18.6, 18.7), (18.8, 18.9, 19. , 19.1)],
                [(19.2, 19.3, 19.4, 19.5), (19.6, 19.7, 19.8, 19.9)]],

               [[(20. , 20.1, 20.2, 20.3), (20.4, 20.5, 20.6, 20.7)],
                [(20.8, 20.9, 21. , 21.1), (21.2, 21.3, 21.4, 21.5)],
                [(21.6, 21.7, 21.8, 21.9), (22. , 22.1, 22.2, 22.3)],
                [(22.4, 22.5, 22.6, 22.7), (22.8, 22.9, 23. , 23.1)],
                [(23.2, 23.3, 23.4, 23.5), (23.6, 23.7, 23.8, 23.9)]]],
              dtype=[('x', '<f8'), ('y', '<f8'), ('z', '<f8'), ('t', '<f8')])

All of the keyword arguments and rules that apply to vector.obj construction apply to vector.array dtypes.

Geometrical names are used in the dtype, even if momentum-synonyms are used in construction.

[48]:

vector.array(
    {"px": [1, 2, 3, 4], "py": [1.1, 2.2, 3.3, 4.4], "pz": [0.1, 0.2, 0.3, 0.4]}
)

[48]:

MomentumNumpy3D([(1., 1.1, 0.1), (2., 2.2, 0.2), (3., 3.3, 0.3), (4., 4.4, 0.4)],
                dtype=[('x', '<f8'), ('y', '<f8'), ('z', '<f8')])

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.Array function behaves exactly like the ak.Array constructor, except that it makes arrays of vectors.

[49]:

vector.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},
        ],
    ]
)

[49]:

[[{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, ...}, {x: 6, y: 6.6, z: 0.6}]]
-----------------------------------------------------------------------
type: 4 * var * Vector3D[
    x: int64,
    y: float64,
    z: float64
]

If you want any records named “Vector2D”, “Vector3D”, “Vector4D”, “Momentum2D”, “Momentum3D”, or “Momentum4D” to be interpreted as vectors, register the behaviors globally.

[50]:

vector.register_awkward()

[51]:

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",
)

[51]:

[[{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, ...}, {x: 6, y: 6.6, z: 0.6}]]
-----------------------------------------------------------------------
type: 4 * var * Vector3D[
    x: int64,
    y: float64,
    z: float64
]

All of the keyword arguments and rules that apply to vector.obj construction apply to vector.Array 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.

[52]:

vector.obj(x=3, y=4).rho

[52]:

$\displaystyle 5.0$

[53]:

vector.obj(rho=5, phi=0.9273).x

[53]:

$\displaystyle 2.99998087197215$

[54]:

vector.obj(rho=5, phi=0.9273).y

[54]:

$\displaystyle 4.00001434594943$

[55]:

vector.obj(x=1, y=2, z=3).theta

[55]:

$\displaystyle 0.640522312679424$

[56]:

vector.obj(x=1, y=2, z=3).eta

[56]:

$\displaystyle 1.10358684156015$

Some properties are not coordinates, but derived from them.

[57]:

vector.obj(x=1, y=2, z=3).costheta

[57]:

$\displaystyle 0.801783725737273$

[58]:

vector.obj(x=1, y=2, z=3).mag  # spatial magnitude

[58]:

$\displaystyle 3.74165738677394$

[59]:

vector.obj(x=1, y=2, z=3).mag2  # spatial magnitude squared

[59]:

$\displaystyle 14$

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.

[60]:

vector.obj(rho=3, phi=0.123456789, z=4).mag2

[60]:

$\displaystyle 25$

Momentum vectors have geometrical properties as well as their momentum-synonyms.

[61]:

vector.obj(px=3, py=4).rho

[61]:

$\displaystyle 5.0$

[62]:

vector.obj(px=3, py=4).pt

[62]:

$\displaystyle 5.0$

[63]:

vector.obj(x=1, y=2, z=3, E=4).tau

[63]:

$\displaystyle 1.4142135623731$

[64]:

vector.obj(x=1, y=2, z=3, E=4).mass

[64]:

$\displaystyle 1.4142135623731$

Here’s the key thing: arrays of vectors return arrays of coordinates.

[65]:

vector.array(
    {
        "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

[65]:

array([1.50363023, 1.50363023, 1.50363023, 1.50363023, 1.50363023])

[66]:

vector.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}],
    ]
).theta

[66]:

[[1.5, 1.5],
 [],
 [1.5],
 [1.5, 1.5]]
-----------------------
type: 4 * var * float64

[67]:

# Make a large, random NumPy array of 3D momentum vectors.
array = (
    np.random.normal(0, 1, 150)
    .view([(x, float) for x in ("x", "y", "z")])
    .view(vector.MomentumNumpy3D)
    .reshape(5, 5, 2)
)
array

[67]:

MomentumNumpy3D([[[( 0.55677765, -2.39134767, -1.14009875e+00),
                   (-0.4099286 , -1.35111044,  1.22805909e+00)],
                  [( 2.17202778, -0.68150151,  1.67793905e-01),
                   (-0.41973476, -1.62000792, -2.62404832e-04)],
                  [(-1.51641767, -1.40147683,  2.47076578e+00),
                   ( 0.60257785, -1.30590029, -6.77753691e-01)],
                  [( 0.19640056,  0.20738716,  2.04561027e-02),
                   (-0.53479171,  0.94471872,  7.09228434e-02)],
                  [( 0.01832378,  1.12336773, -9.97689716e-01),
                   (-1.67101189,  0.62237498, -6.39758959e-02)]],

                 [[( 0.87762728, -1.98374752, -2.01469599e+00),
                   ( 0.17471933,  0.12161887, -9.42606081e-02)],
                  [(-0.5445472 , -0.8928596 , -6.90405216e-01),
                   ( 0.29222871, -0.14303945, -7.06930314e-01)],
                  [( 0.11103332, -0.81876775,  2.21759321e-01),
                   ( 0.6930197 ,  0.44511978,  3.74892308e-01)],
                  [( 0.47797932, -0.54189062, -1.06199035e+00),
                   ( 0.70067255, -0.01567734,  4.51905779e-01)],
                  [(-0.98069469, -0.67226787,  1.63755227e+00),
                   ( 1.33964903, -1.01568176,  9.33791315e-01)]],

                 [[(-0.79781737,  0.8604048 ,  1.75923001e-01),
                   (-1.27055185,  0.28886057,  8.56541197e-01)],
                  [(-1.05366515, -0.15041122,  7.55757703e-01),
                   ( 1.65815748,  1.32132905, -1.46950305e+00)],
                  [(-0.67879973, -1.95059887, -1.15938843e-01),
                   (-1.7877064 , -0.69352448,  3.20265686e-01)],
                  [( 0.28149666,  0.54879607, -7.82537808e-01),
                   ( 0.31366026,  0.36907726,  3.36517183e-01)],
                  [(-1.01885702, -0.66434627,  6.41372208e-01),
                   (-0.77216596,  0.08913472, -2.49368534e-01)]],

                 [[(-0.75470986, -0.15824692, -2.04355139e-01),
                   ( 0.24682491, -0.44316793, -2.55907965e-01)],
                  [(-1.14906251,  1.08234259, -4.37065854e-01),
                   ( 0.39882584, -0.80339686,  8.25692598e-01)],
                  [(-0.2832258 ,  1.04610965,  4.32860052e-01),
                   (-0.27051201, -1.54016307, -6.61530929e-01)],
                  [(-0.10674002,  0.14105683,  1.27337139e+00),
                   (-0.4728045 ,  0.22134932, -1.07465967e+00)],
                  [(-0.10652267,  0.00870547,  1.32620623e+00),
                   (-0.96161573,  0.4993631 ,  5.57284082e-01)]],

                 [[( 0.51468134,  0.48772909, -2.59407633e-01),
                   ( 0.61031019,  1.25211695, -5.04208824e-01)],
                  [(-0.31302188, -0.07605632, -1.93811954e+00),
                   ( 1.02092718,  0.26294964, -9.27338521e-01)],
                  [( 1.88041867, -0.91344964, -6.85620574e-01),
                   (-0.07825839, -0.14948375, -1.52505471e-01)],
                  [(-1.71248754,  0.25225344, -1.28889867e+00),
                   (-0.96428705,  0.87685202, -4.98042758e-01)],
                  [(-0.45582323,  0.40835731,  1.82438535e-01),
                   ( 0.58168782,  1.04561222, -7.12156826e-01)]]],
                dtype=[('x', '<f8'), ('y', '<f8'), ('z', '<f8')])

[68]:

# Get the transverse momentum of each one.
array.pt

[68]:

array([[[2.45530956, 1.41192807],
        [2.27643339, 1.67350021],
        [2.06486316, 1.43821961],
        [0.2856267 , 1.08558539],
        [1.12351716, 1.78315208]],

       [[2.16921272, 0.21288023],
        [1.04581543, 0.32535812],
        [0.82626208, 0.82365522],
        [0.72257157, 0.70084792],
        [1.18899376, 1.68115108]],

       [[1.17337504, 1.30297446],
        [1.06434664, 2.12023506],
        [2.06533412, 1.91751672],
        [0.61677994, 0.48435605],
        [1.21631641, 0.77729356]],

       [[0.77112195, 0.50726753],
        [1.57854684, 0.89694401],
        [1.08377223, 1.56373879],
        [0.1768911 , 0.52205327],
        [0.10687781, 1.08354433]],

       [[0.70906738, 1.39293768],
        [0.32212926, 1.05424609],
        [2.09054171, 0.16872986],
        [1.7309666 , 1.30334914],
        [0.61198898, 1.19652231]]])

[69]:

# The array and its components have the same shape.
array.shape

[69]:

$\displaystyle \left( 5, \ 5, \ 2\right)$

[70]:

array.pt.shape

[70]:

$\displaystyle \left( 5, \ 5, \ 2\right)$

[71]:

# Make a large, random Awkward Array of 3D momentum vectors.
array = vector.Array(
    [
        [
            {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

[71]:

[[],
 [{x: 1.42, y: 1.83, z: -1.13}],
 [{x: -2.65, y: -0.374, z: 0.377}, {x: 0.659, y: -0.326, z: 0.924}],
 [{x: -0.0696, y: -0.258, z: 1.15}, {x: -1.3, y: 1.75, z: 0.486}],
 [],
 [{x: -0.745, y: 0.516, z: 0.74}, {x: -1.12, y: 0.194, z: 0.992}],
 [{x: -0.0514, y: 1.35, z: 3.13}],
 [],
 [],
 [{x: 0.797, y: 0.629, z: 0.761}, {...}, {x: 1.11, y: -1.32, z: 0.733}],
 ...,
 [{x: -0.136, y: 1.05, z: -0.978}],
 [],
 [{x: 0.199, y: -0.0224, z: -0.722}],
 [{x: -0.355, y: 0.835, z: -0.552}],
 [{x: 0.847, y: -1.1, z: -0.557}, {...}, ..., {x: -1.18, y: 1.25, z: -0.526}],
 [{x: 0.476, y: -0.506, z: 0.509}, {...}, ..., {x: 1.5, y: 0.96, z: 0.858}],
 [{x: 0.296, y: -0.65, z: -2.33}, {x: -0.717, y: 0.154, z: 0.364}],
 [{x: -0.235, y: 0.157, z: 0.866}],
 [{x: 0.463, y: -0.0896, z: -0.587}, {x: -2.03, y: -1.26, z: 0.28}]]
------------------------------------------------------------------------------
type: 50 * var * Momentum3D[
    x: float64,
    y: float64,
    z: float64
]

[72]:

# Get the transverse momentum of each one, in the same nested structure.
array.pt

[72]:

[[],
 [2.32],
 [2.68, 0.735],
 [0.268, 2.18],
 [],
 [0.906, 1.13],
 [1.35],
 [],
 [],
 [1.02, 1.23, 1.73],
 ...,
 [1.06],
 [],
 [0.2],
 [0.907],
 [1.39, 0.961, 0.944, 1.72],
 [0.695, 1.62, 0.846, 1.78],
 [0.714, 0.734],
 [0.283],
 [0.472, 2.39]]
----------------------------
type: 50 * var * float64

[73]:

# The array and its components have the same list lengths (and can therefore be used together in subsequent calculations).
ak.num(array)

[73]:

[0,
 1,
 2,
 2,
 0,
 2,
 1,
 0,
 0,
 3,
 ...,
 1,
 0,
 1,
 1,
 4,
 4,
 2,
 1,
 2]
----------------
type: 50 * int64

[74]:

ak.num(array.pt)

[74]:

[0,
 1,
 2,
 2,
 0,
 2,
 1,
 0,
 0,
 3,
 ...,
 1,
 0,
 1,
 1,
 4,
 4,
 2,
 1,
 2]
----------------
type: 50 * int64

Vector methods#

Vector methods require arguments (in parentheses), which may be scalars or other vectors, depending on the calculation.

[75]:

vector.obj(x=3, y=4).rotateZ(0.1)

[75]:

VectorObject2D(x=2.585678829246765, y=4.279516911052588)

[76]:

vector.obj(rho=5, phi=0.4).rotateZ(0.1)

[76]:

VectorObject2D(rho=5, phi=0.5)

[77]:

# Broadcasts a scalar rotation angle of 0.5 to all elements of the NumPy array.
print(
    vector.array({"rho": [1, 2, 3, 4, 5], "phi": [0.1, 0.2, 0.3, 0.4, 0.5]}).rotateZ(
        0.5
    )
)

[(1., 0.6) (2., 0.7) (3., 0.8) (4., 0.9) (5., 1. )]

[78]:

# Matches each rotation angle to an element of the NumPy array.
print(
    vector.array({"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])
    )
)

[(1., 0.2) (2., 0.4) (3., 0.6) (4., 0.8) (5., 1. )]

[79]:

# Broadcasts a scalar rotation angle of 0.5 to all elements of the Awkward Array.
print(
    vector.Array(
        [[{"rho": 1, "phi": 0.1}, {"rho": 2, "phi": 0.2}], [], [{"rho": 3, "phi": 0.3}]]
    ).rotateZ(0.5)
)

[[{rho: 1, phi: 0.6}, {rho: 2, phi: 0.7}], [], [{rho: 3, phi: 0.8}]]

[80]:

# Broadcasts a rotation angle of 0.1 to both elements of the first list, 0.2 to the empty list, and 0.3 to the only element of the last list.
print(
    vector.Array(
        [[{"rho": 1, "phi": 0.1}, {"rho": 2, "phi": 0.2}], [], [{"rho": 3, "phi": 0.3}]]
    ).rotateZ([0.1, 0.2, 0.3])
)

[[{rho: 1, phi: 0.2}, {rho: 2, phi: 0.3}], [], [{rho: 3, phi: 0.6}]]

[81]:

# Matches each rotation angle to an element of the Awkward Array.
print(
    vector.Array(
        [[{"rho": 1, "phi": 0.1}, {"rho": 2, "phi": 0.2}], [], [{"rho": 3, "phi": 0.3}]]
    ).rotateZ([[0.1, 0.2], [], [0.3]])
)

[[{rho: 1, phi: 0.2}, {rho: 2, phi: 0.4}], [], [{rho: 3, phi: 0.6}]]

Some methods are equivalent to binary operators.

[82]:

vector.obj(x=3, y=4).scale(10)

[82]:

VectorObject2D(x=30, y=40)

[83]:

vector.obj(x=3, y=4) * 10

[83]:

VectorObject2D(x=30, y=40)

[84]:

10 * vector.obj(x=3, y=4)

[84]:

VectorObject2D(x=30, y=40)

[85]:

vector.obj(rho=5, phi=0.5) * 10

[85]:

VectorObject2D(rho=50, phi=0.5)

Some methods involve more than one vector.

[86]:

vector.obj(x=1, y=2).add(vector.obj(x=5, y=5))

[86]:

VectorObject2D(x=6, y=7)

[87]:

vector.obj(x=1, y=2) + vector.obj(x=5, y=5)

[87]:

VectorObject2D(x=6, y=7)

[88]:

vector.obj(x=1, y=2).dot(vector.obj(x=5, y=5))

[88]:

$\displaystyle 15$

[89]:

vector.obj(x=1, y=2) @ vector.obj(x=5, y=5)

[89]:

$\displaystyle 15$

The vectors can use different coordinate systems. Conversions are necessary, but minimized for speed and numeric stability.

[90]:

vector.obj(x=3, y=4) @ vector.obj(x=6, y=8)  # both are Cartesian, dot product is exact

[90]:

$\displaystyle 50$

[91]:

vector.obj(rho=5, phi=0.9273) @ vector.obj(
    x=6, y=8
)  # one is polar, dot product is approximate

[91]:

$\displaystyle 49.9999999994283$

[92]:

vector.obj(x=3, y=4) @ vector.obj(
    rho=10, phi=0.9273
)  # one is polar, dot product is approximate

[92]:

$\displaystyle 49.9999999994283$

[93]:

vector.obj(rho=5, phi=0.9273) @ vector.obj(
    rho=10, phi=0.9273
)  # both are polar, a formula that depends on phi differences is used

[93]:

$\displaystyle 50.0$

In Python, some “operators” are actually built-in functions, such as abs.

[94]:

abs(vector.obj(x=3, y=4))

[94]:

$\displaystyle 5.0$

Note that abs returns

rho for 2D vectors
mag for 3D vectors
tau (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.

[95]:

vector.array({"x": [1, 2, 3, 4, 5], "y": [0.1, 0.2, 0.3, 0.4, 0.5]}) + vector.obj(
    x=10, y=5
)

[95]:

VectorNumpy2D([(11., 5.1), (12., 5.2), (13., 5.3), (14., 5.4), (15., 5.5)],
              dtype=[('x', '<f8'), ('y', '<f8')])

[96]:

(
    vector.Array(
        [  # an Awkward Array of vectors
            [{"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)  # and a single vector object
)

[96]:

[[{x: 11, y: 6.1}, {x: 12, y: 7.2}],
 [],
 [{x: 13, y: 8.3}],
 [{x: 14, y: 9.4}, {x: 15, y: 10.5}]]
-------------------------------------
type: 4 * var * Vector2D[
    x: int64,
    y: float64
]

[97]:

(
    vector.Array(
        [  # an Awkward Array of vectors
            [{"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.array(
        {"x": [4, 3, 2, 1], "y": [0.1, 0.1, 0.1, 0.1]}
    )  # and a NumPy array of vectors
)

[97]:

[[{x: 5, y: 1.2}, {x: 6, y: 2.3}],
 [],
 [{x: 5, y: 3.4}],
 [{x: 5, y: 4.5}, {x: 6, y: 5.6}]]
----------------------------------
type: 4 * var * Vector2D[
    x: float64,
    y: float64
]

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.

[98]:

vector.obj(rho=1, phi=0.5).deltaphi(
    vector.obj(rho=2, phi=0.3)
)  # deltaphi is a planar operation (defined on the transverse plane)

[98]:

$\displaystyle 0.2$

[99]:

vector.obj(rho=1, phi=0.5, z=10).deltaphi(
    vector.obj(rho=2, phi=0.3, theta=1.4)
)  # but we can use it on 3D vectors

[99]:

$\displaystyle 0.2$

[100]:

vector.obj(rho=1, phi=0.5, z=10, t=100).deltaphi(
    vector.obj(rho=2, phi=0.3, theta=1.4, tau=1000)
)  # and 4D vectors

[100]:

$\displaystyle 0.2$

[101]:

vector.obj(rho=1, phi=0.5).deltaphi(
    vector.obj(rho=2, phi=0.3, theta=1.4, tau=1000)
)  # and mixed dimensionality

[101]:

$\displaystyle 0.2$

This is especially useful for giving 4D vectors all the capabilities of 3D vectors.

[102]:

vector.obj(x=1, y=2, z=3).rotateX(np.pi / 4)

[102]:

VectorObject3D(x=1, y=-0.7071067811865472, z=3.5355339059327378)

[103]:

vector.obj(x=1, y=2, z=3, tau=10).rotateX(np.pi / 4)

[103]:

VectorObject4D(x=1, y=-0.7071067811865472, z=3.5355339059327378, tau=10)

[104]:

vector.obj(pt=1, phi=1.3, eta=2).deltaR(vector.obj(pt=2, phi=0.3, eta=1))

[104]:

$\displaystyle 1.4142135623731$

[105]:

vector.obj(pt=1, phi=1.3, eta=2, mass=5).deltaR(
    vector.obj(pt=2, phi=0.3, eta=1, mass=10)
)

[105]:

$\displaystyle 1.4142135623731$

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.

[106]:

v1 = vector.obj(x=1, y=2, z=3)
v2 = vector.obj(x=1, y=2)

print(v1 - v2.like(v1))           # transforms v2 to v1's coordinate system (imputes z=0)
print(v1.like(v2) - v2)           # transforms v1 to v2's coordinate system (removes z)
print(v1 - v2.to_xyz())           # transforms v2 to xyz coordinates (imputes z=0)
print(v1.to_xy() - v2)            # transforms v1 to xy coordinates (removes z)
print(v1 - v2.to_Vector3D(z=3))   # transforms v2 to 3D (imputes z=3)
print(v1.to_Vector2D() - v2)      # transforms v1 to 2D (removes z)

VectorObject3D(x=0, y=0, z=3.0)
VectorObject2D(x=0, y=0)
VectorObject3D(x=0, y=0, z=3.0)
VectorObject2D(x=0, y=0)
VectorObject3D(x=0, y=0, z=0)
VectorObject2D(x=0, y=0)

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.

[107]:

vector.obj(x=0.1, y=0.2, z=0.3).cross(vector.obj(x=0.4, y=0.5, z=0.6))

[107]:

VectorObject3D(x=-0.03, y=0.06, z=-0.030000000000000013)

The (current) list of properties and methods is:

Planar (2D, 3D, 4D):

x (px)
y (py)
rho (pt): two-dimensional magnitude
rho2 (pt2): two-dimensional magnitude squared
phi
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 direction
is_antiparallel(vector, tolerance=1e-5): only true if they’re pointing in opposite directions
is_perpendicular(vector, tolerance=1e-5)

Spatial (3D, 4D):

z (pz)
theta
eta
costheta
cottheta
mag (p): three-dimensional magnitude, does not include temporal component
mag2 (p2): three-dimensional magnitude squared
cross: 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, squared
rotateX(angle)
rotateY(angle)
rotate_axis(axis, angle): the magnitude of axis is ignored, but it must be at least 3D
rotate_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 conventions
rotate_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 direction
is_antiparallel(vector, tolerance=1e-5): only true if they’re pointing in opposite directions
is_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 vectors
t2 (E2, energy2)
tau (M, mass): see note above
tau2 (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, squared
gamma: 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 difference
boost_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 behavior
boostX(beta=None, gamma=None): supply beta xor gamma, but not both
boostY(beta=None, gamma=None): supply beta xor gamma, but not both
boostZ(beta=None, gamma=None): supply beta xor gamma, but not both
transform4D(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 parentheses
dot(vector): can also use the @ operator
add(vector): can also use the + operator
subtract(vector): can also use the - operator
scale(factor): can also use the * operator
equal(vector): can also use the == operator, but consider isclose instead
not_equal(vector): can also use the != operator, but consider isclose instead
sum(): can also use the numpy.sum or awkward.sum, only for NumPy and Awkward vectors
count_nonzero(): can also use numpy.count_nonzero or awkward.count_nonzero, only for NumPy and Awkward vectors
count(): can also use awkward.count, only for Awkward vectors
isclose(vector, rtol=1e-5, atol=1e-8, equal_nan=False): works like np.isclose; arrays also have an allclose method
to_VectorND(coordinates)/to_ND(coordinates): replace N with the required vector dimension
to_{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:

[108]:

@nb.njit
def compute_mass(v1, v2):
    return (v1 + v2).mass

[109]:

compute_mass(vector.obj(px=1, py=2, pz=3, E=4), vector.obj(px=-1, py=-2, pz=-3, E=4))

[109]:

$\displaystyle 8.0$

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.

[110]:

# This is still not a large number. You want millions.
array = vector.Array(
    [
        [
            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)
    ]
)
array

[110]:

[[{x: 0.315, y: 0.693, z: 2.09, t: 9.42}],
 [{x: -0.149, y: -0.385, z: -0.293, t: 8.9}, {x: -0.397, y: -0.674, ...}],
 [],
 [],
 [{x: -0.802, y: 2.2, z: -1.14, t: 9.28}, {x: 0.341, y: -0.135, ...}],
 [{x: -1.72, y: 0.922, z: 2.27, t: 11}],
 [{x: 1.25, y: 1.45, z: 0.832, t: 9.73}],
 [{x: 0.327, y: 1.17, z: -1.98, t: 10}, ..., {x: -0.171, y: 1.32, z: ..., ...}],
 [{x: -0.86, y: -0.585, z: -0.935, t: 10}, ..., {x: 0.366, y: -0.783, ...}],
 [{x: -0.0849, y: 0.189, z: -0.815, t: 8.42}, {x: 1.43, y: 1.94, ...}],
 ...,
 [{x: -0.412, y: -0.0528, z: -3.04, t: 10}, ..., {x: 0.406, y: -0.839, ...}],
 [],
 [{x: -0.537, y: 0.306, z: -0.446, t: 10.7}, ..., {x: -0.335, y: 0.516, ...}],
 [{x: 1.19, y: 1.42, z: 2.3, t: 8.32}, {...}, {x: -0.434, y: -0.756, ...}],
 [{x: -0.88, y: -0.763, z: -1.24, t: 9.08}, ..., {x: 1.85, y: 0.672, ...}],
 [{x: -1.19, y: 0.412, z: 0.985, t: 11.1}, {x: 0.132, y: -1.57, ...}],
 [],
 [{x: -0.843, y: 1.28, z: 0.147, t: 10.8}, ..., {x: -0.564, y: -2.19, ...}],
 [{x: 1.37, y: 1.31, z: -0.897, t: 8.11}]]
--------------------------------------------------------------------------------
type: 50 * var * Momentum4D[
    x: float64,
    y: float64,
    z: float64,
    t: float64
]

[111]:

@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

[112]:

compute_masses(array)

[112]:

array([ 9.1537046 , 16.76309287,  0.        ,  0.        , 19.91509641,
       10.54846375,  9.50567043, 50.02632217, 29.51323214, 15.98664711,
        9.45752217, 20.20656999, 38.98606898,  0.        ,  9.32749399,
        9.80117359, 10.70486904, 10.92982222, 35.85520701,  0.        ,
       29.04849791, 19.86279222, 10.31575156, 41.42782629, 11.75731099,
        9.81553407, 39.64420462, 20.78386006, 19.88365545,  8.76680915,
        9.54468769, 21.83630629,  8.90974542, 21.75954326, 19.27688599,
       10.06840737,  9.11771604, 10.78960379,  0.        , 10.55577844,
       19.37816397, 48.01216378,  0.        , 31.96752978, 28.30692309,
       28.78329415, 23.09265714,  0.        , 31.48155952,  7.83341747])

Talks about vector#

9th October 2023 - What’s new with Vector? First major release is out! - PyHEP 2023 (virtual) 🎥
13th September 2022 - Constructing HEP vectors and analyzing HEP data using Vector - PyHEP 2022 (virtual) 🎥
20th July 2022 - Analysis Grand Challenge / HEP Scientific Python Ecosystem - DANCE/CoDaS@Snowmass 2022 computational and data science software training
25th April 2022 - Foundation libraries (uproot, awkward, hist, mplhep) - IRIS-HEP AGC Tools 2022 Workshop 🎥
3rd November 2021 - Data handling: uproot, awkward & vector - IRIS-HEP AGC Tools 2021 Workshop 🎥

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!

[ ]:

Introduction to Vector

Contents

Introduction to Vector#

Constructing a vector or an array of vectors#

Pure Python vectors#

SymPy vectors#

NumPy arrays of vectors#

Awkward Arrays of vectors#

Vector properties#

Vector methods#

Compiling your Python with Numba#

Talks about vector#

Status as of November 17, 2023#