Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 49 additions & 2 deletions src/vector/_methods.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,10 +22,12 @@


class Coordinates:
pass
__slots__ = ()


class Azimuthal(Coordinates):
__slots__ = ()

@property
def elements(self) -> tuple[ScalarCollection, ScalarCollection]:
"""
Expand All @@ -38,6 +40,8 @@ def elements(self) -> tuple[ScalarCollection, ScalarCollection]:


class Longitudinal(Coordinates):
__slots__ = ()

@property
def elements(self) -> tuple[ScalarCollection]:
"""
Expand All @@ -50,6 +54,8 @@ def elements(self) -> tuple[ScalarCollection]:


class Temporal(Coordinates):
__slots__ = ()

@property
def elements(self) -> tuple[ScalarCollection]:
"""
Expand All @@ -68,6 +74,7 @@ class AzimuthalXY(Azimuthal):
y (scalar, ``np.ndarray``, ``ak.Array``, etc.): The $y$ coordinate(s).
"""

__slots__ = ()
x: ScalarCollection
y: ScalarCollection

Expand All @@ -79,6 +86,7 @@ class AzimuthalRhoPhi(Azimuthal):
phi (scalar, ``np.ndarray``, ``ak.Array``, etc.): The $\phi$ coordinate(s).
"""

__slots__ = ()
rho: ScalarCollection
phi: ScalarCollection

Expand All @@ -89,6 +97,7 @@ class LongitudinalZ(Longitudinal):
z (scalar, ``np.ndarray``, ``ak.Array``, etc.): The $z$ coordinate(s).
"""

__slots__ = ()
z: ScalarCollection


Expand All @@ -98,6 +107,7 @@ class LongitudinalTheta(Longitudinal):
theta (scalar, ``np.ndarray``, ``ak.Array``, etc.): The $\theta$ coordinate(s).
"""

__slots__ = ()
theta: ScalarCollection


Expand All @@ -107,6 +117,7 @@ class LongitudinalEta(Longitudinal):
eta (scalar, ``np.ndarray``, ``ak.Array``, etc.): The $\eta$ coordinate(s).
"""

__slots__ = ()
eta: ScalarCollection


Expand All @@ -116,6 +127,7 @@ class TemporalT(Temporal):
t (scalar, ``np.ndarray``, ``ak.Array``, etc.): The $t$ coordinate(s).
"""

__slots__ = ()
t: ScalarCollection


Expand All @@ -125,6 +137,7 @@ class TemporalTau(Temporal):
tau (scalar, ``np.ndarray``, ``ak.Array``, etc.): The $\tau$ coordinate(s).
"""

__slots__ = ()
tau: ScalarCollection


Expand All @@ -148,6 +161,8 @@ class VectorProtocol:
momentum-synonyms.
"""

__slots__ = ()

@property
def lib(self) -> Module: ...

Expand Down Expand Up @@ -710,6 +725,8 @@ def like(self, other: VectorProtocol) -> VectorProtocol:


class VectorProtocolPlanar(VectorProtocol):
__slots__ = ()

@property
def azimuthal(self) -> Azimuthal:
"""
Expand Down Expand Up @@ -833,6 +850,8 @@ def is_perpendicular(


class VectorProtocolSpatial(VectorProtocolPlanar):
__slots__ = ()

@property
def longitudinal(self) -> Longitudinal:
"""
Expand Down Expand Up @@ -1087,6 +1106,8 @@ def is_perpendicular(


class VectorProtocolLorentz(VectorProtocolSpatial):
__slots__ = ()

@property
def temporal(self) -> Temporal:
"""
Expand Down Expand Up @@ -1458,6 +1479,8 @@ def is_lightlike(self, tolerance: ScalarCollection = 1e-5) -> BoolCollection:


class MomentumProtocolPlanar(VectorProtocolPlanar):
__slots__ = ()

@property
def px(self) -> ScalarCollection:
"""Momentum-synonym for :attr:`vector._methods.VectorProtocolPlanar.x`."""
Expand All @@ -1480,6 +1503,8 @@ def pt2(self) -> ScalarCollection:


class MomentumProtocolSpatial(VectorProtocolSpatial, MomentumProtocolPlanar):
__slots__ = ()

@property
def pz(self) -> ScalarCollection:
"""Momentum-synonym for :attr:`vector._methods.VectorProtocolSpatial.z`."""
Expand All @@ -1502,6 +1527,8 @@ def p2(self) -> ScalarCollection:


class MomentumProtocolLorentz(VectorProtocolLorentz, MomentumProtocolSpatial):
__slots__ = ()

@property
def E(self) -> ScalarCollection:
"""Momentum-synonyor :attr:`vector._methods.VectorProtocolLorentz.t`."""
Expand Down Expand Up @@ -1648,6 +1675,8 @@ def transverse_mass2(self) -> ScalarCollection:


class Vector(VectorProtocol):
__slots__ = ()

@typing.overload
def __new__(cls, *, x: float, y: float) -> vector.VectorObject2D: ...

Expand Down Expand Up @@ -3170,6 +3199,8 @@ def like(self, other: VectorProtocol) -> VectorProtocol:


class Vector2D(Vector, VectorProtocolPlanar):
__slots__ = ()

def to_Vector2D(self) -> VectorProtocolPlanar:
return self

Expand Down Expand Up @@ -3306,6 +3337,8 @@ def to_4D(self, **kwargs: float | None) -> VectorProtocolLorentz:


class Vector3D(Vector, VectorProtocolSpatial):
__slots__ = ()

def to_Vector2D(self) -> VectorProtocolPlanar:
return self._wrap_result(
type(self),
Expand Down Expand Up @@ -3384,6 +3417,8 @@ def to_4D(self, **kwargs: float | None) -> VectorProtocolLorentz:


class Vector4D(Vector, VectorProtocolLorentz):
__slots__ = ()

def to_Vector2D(self) -> VectorProtocolPlanar:
return self._wrap_result(
type(self),
Expand Down Expand Up @@ -3423,6 +3458,8 @@ def to_4D(self) -> VectorProtocolLorentz:


class Planar(VectorProtocolPlanar):
__slots__ = ()

@property
def x(self) -> ScalarCollection:
from vector._compute.planar import x
Expand Down Expand Up @@ -3554,6 +3591,8 @@ def isclose(


class Spatial(Planar, VectorProtocolSpatial):
__slots__ = ()

@property
def z(self) -> ScalarCollection:
from vector._compute.spatial import z
Expand Down Expand Up @@ -3794,6 +3833,8 @@ def isclose(


class Lorentz(Spatial, VectorProtocolLorentz):
__slots__ = ()

@property
def t(self) -> ScalarCollection:
from vector._compute.lorentz import t
Expand Down Expand Up @@ -4065,10 +4106,12 @@ def isclose(


class Momentum:
pass
__slots__ = ()


class PlanarMomentum(Momentum, MomentumProtocolPlanar):
__slots__ = ()

@property
def px(self) -> ScalarCollection:
return self.x
Expand All @@ -4087,6 +4130,8 @@ def pt2(self) -> ScalarCollection:


class SpatialMomentum(PlanarMomentum, MomentumProtocolSpatial):
__slots__ = ()

@property
def pz(self) -> ScalarCollection:
return self.z
Expand All @@ -4105,6 +4150,8 @@ def p2(self) -> ScalarCollection:


class LorentzMomentum(SpatialMomentum, MomentumProtocolLorentz):
__slots__ = ()

@property
def E(self) -> ScalarCollection:
return self.t
Expand Down
23 changes: 23 additions & 0 deletions src/vector/backends/object.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,18 +71,26 @@
class CoordinatesObject:
"""Coordinates class for the Object backend."""

__slots__ = ()


class AzimuthalObject(CoordinatesObject, Azimuthal):
"""Azimuthal class for the Object backend."""

__slots__ = ()


class LongitudinalObject(CoordinatesObject, Longitudinal):
"""Longitudinal class for the Object backend."""

__slots__ = ()


class TemporalObject(CoordinatesObject, Temporal):
"""Temporal class for the Object backend."""

__slots__ = ()


class TupleXY(typing.NamedTuple):
"""``x`` and ``y`` coordinates as a ``NamedTuple``."""
Expand All @@ -97,6 +105,8 @@ class AzimuthalObjectXY(AzimuthalObject, AzimuthalXY, TupleXY):
Use the ``elements`` property to retrieve the coordinates.
"""

__slots__ = ()

@property
def elements(self) -> tuple[float, float]:
"""
Expand Down Expand Up @@ -126,6 +136,8 @@ class AzimuthalObjectRhoPhi(AzimuthalObject, AzimuthalRhoPhi, TupleRhoPhi):
Use the ``elements`` property to retrieve the coordinates.
"""

__slots__ = ()

@property
def elements(self) -> tuple[float, float]:
"""
Expand Down Expand Up @@ -154,6 +166,8 @@ class LongitudinalObjectZ(LongitudinalObject, LongitudinalZ, TupleZ):
Use the ``elements`` property to retrieve the coordinates.
"""

__slots__ = ()

@property
def elements(self) -> tuple[float]:
"""
Expand Down Expand Up @@ -182,6 +196,8 @@ class LongitudinalObjectTheta(LongitudinalObject, LongitudinalTheta, TupleTheta)
Use the ``elements`` property to retrieve the coordinates.
"""

__slots__ = ()

@property
def elements(self) -> tuple[float]:
"""
Expand Down Expand Up @@ -210,6 +226,8 @@ class LongitudinalObjectEta(LongitudinalObject, LongitudinalEta, TupleEta):
Use the ``elements`` property to retrieve the coordinates.
"""

__slots__ = ()

@property
def elements(self) -> tuple[float]:
"""
Expand Down Expand Up @@ -238,6 +256,8 @@ class TemporalObjectT(TemporalObject, TemporalT, TupleT):
Use the ``elements`` property to retrieve the coordinates.
"""

__slots__ = ()

@property
def elements(self) -> tuple[float]:
"""
Expand Down Expand Up @@ -266,6 +286,8 @@ class TemporalObjectTau(TemporalObject, TemporalTau, TupleTau):
Use the ``elements`` property to retrieve the coordinates.
"""

__slots__ = ()

@property
def elements(self) -> tuple[float]:
"""
Expand Down Expand Up @@ -329,6 +351,7 @@ def _replace_data(obj: typing.Any, result: typing.Any) -> typing.Any:
class VectorObject(Vector): # noqa: PLW1641
"""Mixin class for Object vectors."""

__slots__ = ()
lib = numpy

# The type ignore comments below cannot be removed because `VectorObject`
Expand Down
Loading