diff --git a/docs/src/make_awkward.md b/docs/src/make_awkward.md index fae35119..cef824aa 100644 --- a/docs/src/make_awkward.md +++ b/docs/src/make_awkward.md @@ -27,7 +27,18 @@ and thus the `abs` function computes the magnitude of each record as `sqrt(x**2 It is not necessary to install Vector's behaviors globally. They could be installed in the `vec` array only by passing `behavior=vector.backends.awkward.behavior` to the [ak.Array](https://awkward-array.org/doc/main/reference/generated/ak.Array.html) constructor. -The records can contain more fields than those that specify coordinates, which can be useful for specifying properties of a particle other than its momentum. Only the coordinate names are considered when performing vector calculations. Coordinates must be numbers (not, for instance, lists of numbers). Be careful about field names that coincide with coordinates, such as `rho` (azimuthal magnitude) and `tau` (proper time). +The records can contain more fields than those that specify coordinates, which can be useful for specifying properties of a particle other than its momentum. Only the coordinate names are considered when performing vector calculations. Coordinates must be numbers (not, for instance, lists of numbers). + +Field names that mean a coordinate under another name, such as `rho` (a synonym of `pt`) or `energy` in a record that already has `mass`, would silently change what a vector means, so Vector rejects them: + +```python +>>> ak.Array([{"pt": 1.1, "phi": 2.2, "eta": 3.3, "mass": 4.4, "energy": 5.5}], with_name="Momentum4D") +Traceback (most recent call last): + ... +TypeError: specify t= or tau=, but not more than one +``` + +This runs whenever the behaviors are attached, not only in `vector.Array` and `vector.zip`, so `jets["rho"] = ...` is caught too. It is the `__awkward_validation__` hook of each behavior class and needs `awkward>=2.8.11`; older versions skip it. Subclassed behaviors (see [Advanced: subclassing Awkward-Vector behaviors](awkward.ipynb)) inherit the check and can extend it by overriding `__awkward_validation__` and calling `super().__awkward_validation__()`. The `vector.Array` function (`vector.awk` is a synonym) is an alternative to the [ak.Array](https://awkward-array.org/doc/main/reference/generated/ak.Array.html) constructor, which installs Vector's behavior in the new array (not globally in `ak.behavior`). diff --git a/src/vector/_methods.py b/src/vector/_methods.py index 91f10de5..4bdeb3b0 100644 --- a/src/vector/_methods.py +++ b/src/vector/_methods.py @@ -5,6 +5,7 @@ from __future__ import annotations +import functools import typing from contextlib import suppress @@ -4357,6 +4358,178 @@ def _compute_module_of( ] +_repr_all_to_generic = { + **{x: x for x in ("x", "y", "rho", "phi", "z", "theta", "eta", "t", "tau")}, + **_repr_momentum_to_generic, +} + + +_azimuthal_combinations = (("x", "y"), ("rho", "phi")) +_azimuthal_names = ("x", "y", "rho", "phi") +_longitudinal_names = ("z", "theta", "eta") +_temporal_names = ("t", "tau") + + +# Every name, including momentum-aliases, grouped by geometry tier. +_azimuthal_fields = frozenset( + name + for name, generic in _repr_all_to_generic.items() + if generic in _azimuthal_names +) +_longitudinal_fields = frozenset( + name + for name, generic in _repr_all_to_generic.items() + if generic in _longitudinal_names +) +_temporal_fields = frozenset( + name for name, generic in _repr_all_to_generic.items() if generic in _temporal_names +) + + +# The 2 + 6 + 12 combinations that describe a vector, in the order reported to users. +_allowed_coordinates = ( + *_azimuthal_combinations, + *( + (*azimuthal, longitudinal) + for azimuthal in _azimuthal_combinations + for longitudinal in _longitudinal_names + ), + *( + (*azimuthal, longitudinal, temporal) + for azimuthal in _azimuthal_combinations + for longitudinal in _longitudinal_names + for temporal in _temporal_names + ), +) + + +def _coordinate_complaint(dimension: int | None, momentum: bool | None) -> str: + """Lists the combinations a vector of this ``dimension`` may be built from.""" + complaint = "unrecognized combination of coordinates, allowed combinations are:\n\n" + complaint += "\n".join( + " " + + ("" if dimension is not None else f"({len(names)}D) ") + + " ".join(f"{name}=" for name in names) + for names in _allowed_coordinates + if dimension in (None, len(names)) + ) + if momentum is not False: + complaint += "\n\nor their momentum equivalents" + return complaint + + +# Awkward Array validates on every array it attaches a vector behavior to; bounded +# because the field names this is keyed on come from user data. +@functools.lru_cache(maxsize=4096) +def _check_coordinate_names( + fieldnames: tuple[str, ...], + dimension: int | None = None, + momentum: bool | None = None, + allow_extra: bool = False, +) -> tuple[bool, int, tuple[tuple[str, str], ...], tuple[str, ...]]: + """ + Determines the dimension and the momentum-ness of a set of coordinate names, + raising a ``TypeError`` if they do not describe exactly one vector. Every + backend validates through this function. + + Args: + fieldnames (tuple of str): Coordinate names, as given by the user. + dimension (int or None): Dimension that the names must describe, or + None to deduce it from the names. + momentum (bool or None): Whether momentum-aliases are allowed (True), + not allowed (False), or unconstrained (None). + allow_extra (bool): If True, names that are not coordinates are + returned instead of rejected. + + Returns: + tuple: ``(is_momentum, dimension, coordinates, extra)``, in which + ``coordinates`` is a tuple of ``(generic name, given name)`` pairs in + canonical order and ``extra`` holds the names that are not coordinates. + + Examples: + >>> from vector._methods import _check_coordinate_names + >>> _check_coordinate_names(("pt", "phi", "eta")) + (True, 3, (('rho', 'pt'), ('phi', 'phi'), ('eta', 'eta')), ()) + """ + given: dict[str, str] = {} + extra: list[str] = [] + is_momentum = False + + for name in fieldnames: + generic = _repr_all_to_generic.get(name) + if generic is None: + extra.append(name) + continue + if name in _repr_momentum_to_generic: + is_momentum = True + if generic in given: + raise TypeError( + "duplicate coordinates (through momentum-aliases): " + f"{given[generic]!r} and {name!r} both map to {generic!r}" + ) + given[generic] = name + + if is_momentum and momentum is False: + raise TypeError( + "momentum-aliases are not allowed in a generic vector: " + + ", ".join(repr(x) for x in fieldnames if x in _repr_momentum_to_generic) + ) + if extra and not allow_extra: + raise TypeError(_coordinate_complaint(dimension, momentum)) + + if ("x" in given or "y" in given) and ("rho" in given or "phi" in given): + raise TypeError("specify x= and y= or rho= and phi=, but not both") + if sum(name in given for name in _longitudinal_names) > 1: + raise TypeError("specify z= or theta= or eta=, but not more than one") + if sum(name in given for name in _temporal_names) > 1: + raise TypeError("specify t= or tau=, but not more than one") + + azimuthal = next( + (names for names in _azimuthal_combinations if all(x in given for x in names)), + None, + ) + longitudinal = next((x for x in _longitudinal_names if x in given), None) + temporal = next((x for x in _temporal_names if x in given), None) + + if ( + azimuthal is None + or (temporal is not None and longitudinal is None) + or len(given) != 2 + (longitudinal is not None) + (temporal is not None) + ): + raise TypeError(_coordinate_complaint(dimension, momentum)) + + names = ( + *azimuthal, + *(() if longitudinal is None else (longitudinal,)), + *(() if temporal is None else (temporal,)), + ) + + if dimension is not None and dimension != len(names): + raise TypeError(_coordinate_complaint(dimension, momentum)) + + return ( + is_momentum, + len(names), + tuple((name, given[name]) for name in names), + tuple(extra), + ) + + +_CoordinateT = typing.TypeVar("_CoordinateT") + + +def _generic_coordinates( + coordinates: dict[str, _CoordinateT], + dimension: int | None = None, + momentum: bool | None = None, +) -> dict[str, _CoordinateT]: + """Validates a constructor's keyword arguments, keyed by their generic names.""" + _, _, names, _ = _check_coordinate_names( + tuple(coordinates), dimension=dimension, momentum=momentum + ) + return {name: coordinates[given] for name, given in names} + + # Caches mapping a concrete coordinate class to its marker type. These are # keyed on the concrete ``type(...)`` of a coordinate object, which is fixed at # import time apart from rare third-party subclasses; caching by concrete type diff --git a/src/vector/backends/_numba_object.py b/src/vector/backends/_numba_object.py index 5e744c5e..ff0d466e 100644 --- a/src/vector/backends/_numba_object.py +++ b/src/vector/backends/_numba_object.py @@ -27,6 +27,7 @@ Momentum, TemporalT, TemporalTau, + _check_coordinate_names, _from_signature, ) from vector.backends._numba import numba_modules @@ -575,6 +576,34 @@ def vector_obj_Temporal_mass(t, E, e, energy, tau, M, m, mass): return TemporalObjectTau(mass) +_vector_obj_azimuthal = { + ("x", "y"): vector_obj_Azimuthal_xy, + ("x", "py"): vector_obj_Azimuthal_xpy, + ("px", "y"): vector_obj_Azimuthal_pxy, + ("px", "py"): vector_obj_Azimuthal_pxpy, + ("rho", "phi"): vector_obj_Azimuthal_rhophi, + ("pt", "phi"): vector_obj_Azimuthal_ptphi, +} + +_vector_obj_longitudinal = { + "z": vector_obj_Longitudinal_z, + "pz": vector_obj_Longitudinal_pz, + "theta": vector_obj_Longitudinal_theta, + "eta": vector_obj_Longitudinal_eta, +} + +_vector_obj_temporal = { + "t": vector_obj_Temporal_t, + "E": vector_obj_Temporal_E, + "e": vector_obj_Temporal_e, + "energy": vector_obj_Temporal_energy, + "tau": vector_obj_Temporal_tau, + "M": vector_obj_Temporal_M, + "m": vector_obj_Temporal_m, + "mass": vector_obj_Temporal_mass, +} + + @numba.extending.overload(vector.obj) def vector_obj( unrecognized_argument=None, @@ -603,342 +632,130 @@ def vector_obj( "only keyword arguments are allowed in vector.obj; no positional arguments" ) - has_x = x is not None - has_px = px is not None - has_y = y is not None - has_py = py is not None - has_rho = rho is not None - has_pt = pt is not None - has_phi = phi is not None - has_z = z is not None - has_pz = pz is not None - has_theta = theta is not None - has_eta = eta is not None - has_t = t is not None - has_E = E is not None - has_e = e is not None - has_energy = energy is not None - has_tau = tau is not None - has_M = M is not None - has_m = m is not None - has_mass = mass is not None - - is_momentum = False - azimuthal = None - longitudinal = None - temporal = None - - if (has_x and has_y) and not (has_rho or has_pt or has_phi): - if has_px or has_py: - raise numba.TypingError( - "duplicate coordinates (through momentum-aliases): x/px or y/py" - ) - azimuthal = vector_obj_Azimuthal_xy - elif (has_x and has_py) and not (has_rho or has_pt or has_phi): - is_momentum = True - if has_px or has_y: - raise numba.TypingError( - "duplicate coordinates (through momentum-aliases): x/px or y/py" - ) - azimuthal = vector_obj_Azimuthal_xpy - elif (has_px and has_y) and not (has_rho or has_pt or has_phi): - is_momentum = True - if has_x or has_py: - raise numba.TypingError( - "duplicate coordinates (through momentum-aliases): x/px or y/py" - ) - azimuthal = vector_obj_Azimuthal_pxy - elif (has_px and has_py) and not (has_rho or has_pt or has_phi): - is_momentum = True - if has_x or has_y: - raise numba.TypingError( - "duplicate coordinates (through momentum-aliases): x/px or y/py" - ) - azimuthal = vector_obj_Azimuthal_pxpy - elif (has_rho and has_phi) and not (has_x or has_px or has_y or has_py): - if has_pt: - raise numba.TypingError( - "duplicate coordinates (through momentum-aliases): rho/pt" - ) - azimuthal = vector_obj_Azimuthal_rhophi - elif (has_pt and has_phi) and not (has_x or has_px or has_y or has_py): - is_momentum = True - if has_rho: - raise numba.TypingError( - "duplicate coordinates (through momentum-aliases): rho/pt" - ) - azimuthal = vector_obj_Azimuthal_ptphi + given = tuple( + name + for name, value in ( + ("x", x), + ("px", px), + ("y", y), + ("py", py), + ("rho", rho), + ("pt", pt), + ("phi", phi), + ("z", z), + ("pz", pz), + ("theta", theta), + ("eta", eta), + ("t", t), + ("E", E), + ("e", e), + ("energy", energy), + ("tau", tau), + ("M", M), + ("m", m), + ("mass", mass), + ) + if value is not None + ) - if has_z and not (has_theta or has_eta): - if has_pz: - raise numba.TypingError( - "duplicate coordinates (through momentum-aliases): z/pz" - ) - longitudinal = vector_obj_Longitudinal_z - elif has_pz and not (has_theta or has_eta): - is_momentum = True - if has_z: - raise numba.TypingError( - "duplicate coordinates (through momentum-aliases): z/pz" - ) - longitudinal = vector_obj_Longitudinal_pz - elif has_theta and not (has_z or has_eta): - longitudinal = vector_obj_Longitudinal_theta - elif has_eta and not (has_z or has_theta): - longitudinal = vector_obj_Longitudinal_eta - - if has_t and not (has_tau or has_M or has_m or has_mass): - if has_E or has_e or has_energy: - raise numba.TypingError( - "duplicate coordinates (through momentum-aliases): t/E/e/energy" - ) - temporal = vector_obj_Temporal_t - elif has_E and not (has_tau or has_M or has_m or has_mass): - is_momentum = True - if has_t or has_e or has_energy: - raise numba.TypingError( - "duplicate coordinates (through momentum-aliases): t/E/e/energy" - ) - temporal = vector_obj_Temporal_E - elif has_e and not (has_tau or has_M or has_m or has_mass): - is_momentum = True - if has_t or has_E or has_energy: - raise numba.TypingError( - "duplicate coordinates (through momentum-aliases): t/E/e/energy" - ) - temporal = vector_obj_Temporal_e - elif has_energy and not (has_tau or has_M or has_m or has_mass): - is_momentum = True - if has_t or has_E or has_e: - raise numba.TypingError( - "duplicate coordinates (through momentum-aliases): t/E/e/energy" - ) - temporal = vector_obj_Temporal_energy - elif has_tau and not (has_t or has_E or has_e or has_energy): - if has_M or has_m or has_mass: - raise numba.TypingError( - "duplicate coordinates (through momentum-aliases): tau/M/m/mass" - ) - temporal = vector_obj_Temporal_tau - elif has_M and not (has_t or has_E or has_e or has_energy): - is_momentum = True - if has_tau or has_m or has_mass: - raise numba.TypingError( - "duplicate coordinates (through momentum-aliases): tau/M/m/mass" - ) - temporal = vector_obj_Temporal_M - elif has_m and not (has_t or has_E or has_e or has_energy): - is_momentum = True - if has_tau or has_M or has_mass: - raise numba.TypingError( - "duplicate coordinates (through momentum-aliases): tau/M/m/mass" - ) - temporal = vector_obj_Temporal_m - elif has_mass and not (has_t or has_E or has_e or has_energy): - is_momentum = True - if has_tau or has_M or has_m: - raise numba.TypingError( - "duplicate coordinates (through momentum-aliases): tau/M/m/mass" + try: + is_momentum, dimension, names, _ = _check_coordinate_names(given) + except TypeError as err: + raise numba.TypingError(str(err)) from err + + azimuthal = _vector_obj_azimuthal[names[0][1], names[1][1]] + if dimension >= 3: + longitudinal = _vector_obj_longitudinal[names[2][1]] + if dimension == 4: + temporal = _vector_obj_temporal[names[3][1]] + + if dimension == 4: + cls = MomentumObject4D if is_momentum else VectorObject4D + + def vector_obj_impl( + unrecognized_argument=None, + x=None, + px=None, + y=None, + py=None, + rho=None, + pt=None, + phi=None, + z=None, + pz=None, + theta=None, + eta=None, + t=None, + E=None, + e=None, + energy=None, + tau=None, + M=None, + m=None, + mass=None, + ): + return cls( + azimuthal(x, px, y, py, rho, pt, phi), + longitudinal(z, pz, theta, eta), + temporal(t, E, e, energy, tau, M, m, mass), ) - temporal = vector_obj_Temporal_mass - - if azimuthal is not None and longitudinal is not None and temporal is not None: - if is_momentum: - - def vector_obj_impl( - unrecognized_argument=None, - x=None, - px=None, - y=None, - py=None, - rho=None, - pt=None, - phi=None, - z=None, - pz=None, - theta=None, - eta=None, - t=None, - E=None, - e=None, - energy=None, - tau=None, - M=None, - m=None, - mass=None, - ): - return MomentumObject4D( - azimuthal(x, px, y, py, rho, pt, phi), - longitudinal(z, pz, theta, eta), - temporal(t, E, e, energy, tau, M, m, mass), - ) - - else: - - def vector_obj_impl( - unrecognized_argument=None, - x=None, - px=None, - y=None, - py=None, - rho=None, - pt=None, - phi=None, - z=None, - pz=None, - theta=None, - eta=None, - t=None, - E=None, - e=None, - energy=None, - tau=None, - M=None, - m=None, - mass=None, - ): - return VectorObject4D( - azimuthal(x, px, y, py, rho, pt, phi), - longitudinal(z, pz, theta, eta), - temporal(t, E, e, energy, tau, M, m, mass), - ) - - elif azimuthal is not None and longitudinal is not None and temporal is None: - if is_momentum: - - def vector_obj_impl( - unrecognized_argument=None, - x=None, - px=None, - y=None, - py=None, - rho=None, - pt=None, - phi=None, - z=None, - pz=None, - theta=None, - eta=None, - t=None, - E=None, - e=None, - energy=None, - tau=None, - M=None, - m=None, - mass=None, - ): - return MomentumObject3D( - azimuthal(x, px, y, py, rho, pt, phi), - longitudinal(z, pz, theta, eta), - ) - - else: - - def vector_obj_impl( - unrecognized_argument=None, - x=None, - px=None, - y=None, - py=None, - rho=None, - pt=None, - phi=None, - z=None, - pz=None, - theta=None, - eta=None, - t=None, - E=None, - e=None, - energy=None, - tau=None, - M=None, - m=None, - mass=None, - ): - return VectorObject3D( - azimuthal(x, px, y, py, rho, pt, phi), - longitudinal(z, pz, theta, eta), - ) - - elif azimuthal is not None and longitudinal is None and temporal is None: - if is_momentum: - - def vector_obj_impl( - unrecognized_argument=None, - x=None, - px=None, - y=None, - py=None, - rho=None, - pt=None, - phi=None, - z=None, - pz=None, - theta=None, - eta=None, - t=None, - E=None, - e=None, - energy=None, - tau=None, - M=None, - m=None, - mass=None, - ): - return MomentumObject2D(azimuthal(x, px, y, py, rho, pt, phi)) - else: - - def vector_obj_impl( - unrecognized_argument=None, - x=None, - px=None, - y=None, - py=None, - rho=None, - pt=None, - phi=None, - z=None, - pz=None, - theta=None, - eta=None, - t=None, - E=None, - e=None, - energy=None, - tau=None, - M=None, - m=None, - mass=None, - ): - return VectorObject2D(azimuthal(x, px, y, py, rho, pt, phi)) + elif dimension == 3: + cls = MomentumObject3D if is_momentum else VectorObject3D + + def vector_obj_impl( + unrecognized_argument=None, + x=None, + px=None, + y=None, + py=None, + rho=None, + pt=None, + phi=None, + z=None, + pz=None, + theta=None, + eta=None, + t=None, + E=None, + e=None, + energy=None, + tau=None, + M=None, + m=None, + mass=None, + ): + return cls( + azimuthal(x, px, y, py, rho, pt, phi), + longitudinal(z, pz, theta, eta), + ) else: - raise numba.TypingError( - "unrecognized combination of coordinates, allowed combinations are:\n\n" - " (2D) x= y=\n" - " (2D) rho= phi=\n" - " (3D) x= y= z=\n" - " (3D) x= y= theta=\n" - " (3D) x= y= eta=\n" - " (3D) rho= phi= z=\n" - " (3D) rho= phi= theta=\n" - " (3D) rho= phi= eta=\n" - " (4D) x= y= z= t=\n" - " (4D) x= y= z= tau=\n" - " (4D) x= y= theta= t=\n" - " (4D) x= y= theta= tau=\n" - " (4D) x= y= eta= t=\n" - " (4D) x= y= eta= tau=\n" - " (4D) rho= phi= z= t=\n" - " (4D) rho= phi= z= tau=\n" - " (4D) rho= phi= theta= t=\n" - " (4D) rho= phi= theta= tau=\n" - " (4D) rho= phi= eta= t=\n" - " (4D) rho= phi= eta= tau=" - ) + cls = MomentumObject2D if is_momentum else VectorObject2D + + def vector_obj_impl( + unrecognized_argument=None, + x=None, + px=None, + y=None, + py=None, + rho=None, + pt=None, + phi=None, + z=None, + pz=None, + theta=None, + eta=None, + t=None, + E=None, + e=None, + energy=None, + tau=None, + M=None, + m=None, + mass=None, + ): + return cls(azimuthal(x, px, y, py, rho, pt, phi)) return vector_obj_impl diff --git a/src/vector/backends/awkward.py b/src/vector/backends/awkward.py index 767e3e0d..ffe23d76 100644 --- a/src/vector/backends/awkward.py +++ b/src/vector/backends/awkward.py @@ -64,6 +64,10 @@ Vector3D, Vector4D, VectorProtocol, + _azimuthal_fields, + _check_coordinate_names, + _longitudinal_fields, + _temporal_fields, ) from vector._typeutils import BoolCollection, Protocol, ScalarCollection from vector.backends.numpy import VectorNumpy2D, VectorNumpy3D, VectorNumpy4D @@ -607,15 +611,10 @@ def _class_to_name(cls: type[VectorProtocol]) -> str: # the vector class ############################################################ -# Generic and momentum-alias coordinate field names, grouped by geometry tier. # Used by ``_wrap_result`` to exclude (already-recomputed) coordinate fields when # carrying along "extra" record fields, so that stale, pre-computation coordinates -# are not leaked into the output. Previously these tuples were hand-copied into -# each branch and omitted the ``px``/``py`` momentum aliases, leaking stale values. -_azimuthal_fields = frozenset({"x", "y", "rho", "phi", "px", "py", "pt"}) -_longitudinal_fields = frozenset({"z", "theta", "eta", "pz"}) -_temporal_fields = frozenset({"t", "tau", "E", "e", "energy", "M", "m", "mass"}) - +# are not leaked into the output. +# # Exclude only azimuthal coordinates (carry longitudinal/temporal as extras). _coordinate_fields_azimuthal = _azimuthal_fields # Exclude azimuthal + longitudinal coordinates (carry temporal as extras). @@ -624,6 +623,22 @@ def _class_to_name(cls: type[VectorProtocol]) -> str: _coordinate_fields_all = _azimuthal_fields | _longitudinal_fields | _temporal_fields +def _check_fields(array: typing.Any, dimension: int, momentum: bool) -> None: + """ + Validates the fields of an array a vector behavior is being attached to. Its + record name was chosen elsewhere, so the complaint has to name the array. + """ + fields = tuple(array.fields) + try: + _check_coordinate_names( + fields, dimension=dimension, momentum=momentum, allow_extra=True + ) + except TypeError as err: + raise TypeError( + f"{type(array).__name__} with fields {list(fields)}: {err}" + ) from err + + def _yes_record( x: ak.Array, ) -> float | ak.Record | None: @@ -636,6 +651,8 @@ def _no_record(x: ak.Array) -> ak.Array | None: # Type for mixing in Awkward later class AwkwardProtocol(Protocol): + fields: list[str] + def __getitem__(self, where: typing.Any) -> float | ak.Array | ak.Record | None: ... @@ -739,11 +756,9 @@ def _wrap_result( names.append(name) arrays.append(self[name]) - if any( - f in fields for f in ("t", "tau", "M", "m", "mass", "E", "e", "energy") - ): + if any(f in _temporal_fields for f in fields): cls = cls.ProjectionClass4D - elif any(f in fields for f in ("z", "pz", "theta", "eta")): + elif any(f in _longitudinal_fields for f in fields): cls = cls.ProjectionClass3D else: cls = cls.ProjectionClass2D @@ -832,9 +847,7 @@ def _wrap_result( names.append(name) arrays.append(self[name]) - if any( - f in fields for f in ("t", "tau", "M", "m", "mass", "E", "e", "energy") - ): + if any(f in _temporal_fields for f in fields): cls = cls.ProjectionClass4D else: cls = cls.ProjectionClass3D @@ -1064,6 +1077,10 @@ class VectorAwkward2D(VectorAwkward, Planar, Vector2D): See :class:`MomentumAwkward2D` for momentum vectors. """ + def __awkward_validation__(self: AwkwardProtocol) -> None: + """Raises a ``TypeError`` if these fields do not describe a vector.""" + _check_fields(self, 2, False) + @property def azimuthal(self) -> AzimuthalAwkward: """ @@ -1094,6 +1111,10 @@ class MomentumAwkward2D(PlanarMomentum, VectorAwkward2D): See :class:`VectorAwkward2D` for vectors. """ + def __awkward_validation__(self: AwkwardProtocol) -> None: + """Raises a ``TypeError`` if these fields do not describe a momentum vector.""" + _check_fields(self, 2, True) + @property def azimuthal(self) -> AzimuthalAwkward: """ @@ -1124,6 +1145,10 @@ class VectorAwkward3D(VectorAwkward, Spatial, Vector3D): See :class:`MomentumAwkward3D` for momentum vectors. """ + def __awkward_validation__(self: AwkwardProtocol) -> None: + """Raises a ``TypeError`` if these fields do not describe a vector.""" + _check_fields(self, 3, False) + @property def azimuthal(self) -> AzimuthalAwkward: """ @@ -1174,6 +1199,10 @@ class MomentumAwkward3D(SpatialMomentum, VectorAwkward3D): See :class:`VectorAwkward3D` for vectors. """ + def __awkward_validation__(self: AwkwardProtocol) -> None: + """Raises a ``TypeError`` if these fields do not describe a momentum vector.""" + _check_fields(self, 3, True) + @property def azimuthal(self) -> AzimuthalAwkward: """ @@ -1224,6 +1253,10 @@ class VectorAwkward4D(VectorAwkward, Lorentz, Vector4D): See :class:`MomentumAwkward4D` for momentum vectors. """ + def __awkward_validation__(self: AwkwardProtocol) -> None: + """Raises a ``TypeError`` if these fields do not describe a vector.""" + _check_fields(self, 4, False) + @property def azimuthal(self) -> AzimuthalAwkward: """ @@ -1294,6 +1327,10 @@ class MomentumAwkward4D(LorentzMomentum, VectorAwkward4D): See :class:`VectorAwkward4D` for vectors. """ + def __awkward_validation__(self: AwkwardProtocol) -> None: + """Raises a ``TypeError`` if these fields do not describe a momentum vector.""" + _check_fields(self, 4, True) + @property def azimuthal(self) -> AzimuthalAwkward: """ diff --git a/src/vector/backends/awkward_constructors.py b/src/vector/backends/awkward_constructors.py index c43b1fbf..e7d67184 100644 --- a/src/vector/backends/awkward_constructors.py +++ b/src/vector/backends/awkward_constructors.py @@ -10,6 +10,8 @@ import numpy +from vector._methods import _check_coordinate_names + def _recname(is_momentum: bool, dimension: int) -> str: name = "Momentum" if is_momentum else "Vector" @@ -19,190 +21,18 @@ def _recname(is_momentum: bool, dimension: int) -> str: def _check_names( projectable: typing.Any, fieldnames: list[str] ) -> tuple[bool, int, list[str], typing.Any]: - complaint1 = "duplicate coordinates (through momentum-aliases): " + ", ".join( - repr(x) for x in fieldnames - ) - complaint2 = ( - "unrecognized combination of coordinates, allowed combinations are:\n\n" - " (2D) x= y=\n" - " (2D) rho= phi=\n" - " (3D) x= y= z=\n" - " (3D) x= y= theta=\n" - " (3D) x= y= eta=\n" - " (3D) rho= phi= z=\n" - " (3D) rho= phi= theta=\n" - " (3D) rho= phi= eta=\n" - " (4D) x= y= z= t=\n" - " (4D) x= y= z= tau=\n" - " (4D) x= y= theta= t=\n" - " (4D) x= y= theta= tau=\n" - " (4D) x= y= eta= t=\n" - " (4D) x= y= eta= tau=\n" - " (4D) rho= phi= z= t=\n" - " (4D) rho= phi= z= tau=\n" - " (4D) rho= phi= theta= t=\n" - " (4D) rho= phi= theta= tau=\n" - " (4D) rho= phi= eta= t=\n" - " (4D) rho= phi= eta= tau=" + """ + Determines the record name and the columns of an array of vectors from its + field names, allowing fields that are not coordinates to be carried along. + """ + is_momentum, dimension, coordinates, extra = _check_coordinate_names( + tuple(fieldnames), allow_extra=True ) - is_momentum = False - dimension = 0 - names = [] - columns = [] - - if "x" in fieldnames and "y" in fieldnames: - if dimension != 0: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 2 - names.extend(["x", "y"]) - columns.extend([projectable["x"], projectable["y"]]) - fieldnames.remove("x") - fieldnames.remove("y") - if "rho" in fieldnames and "phi" in fieldnames: - if dimension != 0: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 2 - names.extend(["rho", "phi"]) - columns.extend([projectable["rho"], projectable["phi"]]) - fieldnames.remove("rho") - fieldnames.remove("phi") - if "x" in fieldnames and "py" in fieldnames: - is_momentum = True - if dimension != 0: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 2 - names.extend(["x", "y"]) - columns.extend([projectable["x"], projectable["py"]]) - fieldnames.remove("x") - fieldnames.remove("py") - if "px" in fieldnames and "y" in fieldnames: - is_momentum = True - if dimension != 0: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 2 - names.extend(["x", "y"]) - columns.extend([projectable["px"], projectable["y"]]) - fieldnames.remove("px") - fieldnames.remove("y") - if "px" in fieldnames and "py" in fieldnames: - is_momentum = True - if dimension != 0: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 2 - names.extend(["x", "y"]) - columns.extend([projectable["px"], projectable["py"]]) - fieldnames.remove("px") - fieldnames.remove("py") - if "pt" in fieldnames and "phi" in fieldnames: - is_momentum = True - if dimension != 0: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 2 - names.extend(["rho", "phi"]) - columns.extend([projectable["pt"], projectable["phi"]]) - fieldnames.remove("pt") - fieldnames.remove("phi") - - if "z" in fieldnames: - if dimension != 2: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 3 - names.append("z") - columns.append(projectable["z"]) - fieldnames.remove("z") - if "theta" in fieldnames: - if dimension != 2: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 3 - names.append("theta") - columns.append(projectable["theta"]) - fieldnames.remove("theta") - if "eta" in fieldnames: - if dimension != 2: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 3 - names.append("eta") - columns.append(projectable["eta"]) - fieldnames.remove("eta") - if "pz" in fieldnames: - is_momentum = True - if dimension != 2: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 3 - names.append("z") - columns.append(projectable["pz"]) - fieldnames.remove("pz") - - if "t" in fieldnames: - if dimension != 3: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 4 - names.append("t") - columns.append(projectable["t"]) - fieldnames.remove("t") - if "tau" in fieldnames: - if dimension != 3: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 4 - names.append("tau") - columns.append(projectable["tau"]) - fieldnames.remove("tau") - if "E" in fieldnames: - is_momentum = True - if dimension != 3: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 4 - names.append("t") - columns.append(projectable["E"]) - fieldnames.remove("E") - if "e" in fieldnames: - is_momentum = True - if dimension != 3: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 4 - names.append("t") - columns.append(projectable["e"]) - fieldnames.remove("e") - if "energy" in fieldnames: - is_momentum = True - if dimension != 3: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 4 - names.append("t") - columns.append(projectable["energy"]) - fieldnames.remove("energy") - if "M" in fieldnames: - is_momentum = True - if dimension != 3: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 4 - names.append("tau") - columns.append(projectable["M"]) - fieldnames.remove("M") - if "m" in fieldnames: - is_momentum = True - if dimension != 3: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 4 - names.append("tau") - columns.append(projectable["m"]) - fieldnames.remove("m") - if "mass" in fieldnames: - is_momentum = True - if dimension != 3: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 4 - names.append("tau") - columns.append(projectable["mass"]) - fieldnames.remove("mass") - - if dimension == 0: - raise TypeError(complaint1 if is_momentum else complaint2) - - for name in fieldnames: - names.append(name) - columns.append(projectable[name]) + names = [name for name, _ in coordinates] + list(extra) + columns = [projectable[given] for _, given in coordinates] + [ + projectable[name] for name in extra + ] return is_momentum, dimension, names, columns @@ -287,6 +117,11 @@ def Array(*args: typing.Any, **kwargs: typing.Any) -> typing.Any: to make the vector a momentum vector. + A coordinate may be given only once, whether by its generic name or through + a momentum-alias, and the names must form exactly one of the combinations + above; anything else raises a ``TypeError``. Names that are not coordinates + become extra fields of the records. + No constraints are placed on the types of the vector fields, though if they are not numbers, mathematical operations will fail. Usually, you want them to be integers or floating-point numbers. @@ -312,8 +147,6 @@ def Array(*args: typing.Any, **kwargs: typing.Any) -> typing.Any: is_momentum, dimension, names, arrays = _check_names(akarray, fields.copy()) - assert 2 <= dimension <= 4, f"Dimension must be between 2-4, not {dimension}" - return awkward.with_name( awkward.zip( dict(builtins.zip(names, arrays, strict=True)), @@ -377,6 +210,11 @@ def zip(arrays: dict[str, typing.Any], depth_limit: int | None = None) -> typing - ``mass`` may be substituted for ``tau`` to make the vector a momentum vector. + + A coordinate may be given only once, whether by its generic name or through + a momentum-alias, and the names must form exactly one of the combinations + above; anything else raises a ``TypeError``. Names that are not coordinates + become extra fields of the records. """ import awkward diff --git a/src/vector/backends/numpy.py b/src/vector/backends/numpy.py index 762212d1..cd556caf 100644 --- a/src/vector/backends/numpy.py +++ b/src/vector/backends/numpy.py @@ -46,6 +46,7 @@ Vector4D, VectorProtocol, _aztype, + _check_coordinate_names, _coordinate_class_to_names, _coordinate_order, _handler_of, @@ -865,6 +866,46 @@ def tau(self) -> FloatArray: return self["tau"] +_azimuthal_numpy_type: dict[ + str, type[AzimuthalNumpyXY] | type[AzimuthalNumpyRhoPhi] +] = {"x": AzimuthalNumpyXY, "rho": AzimuthalNumpyRhoPhi} +_longitudinal_numpy_type: dict[ + str, + type[LongitudinalNumpyZ] + | type[LongitudinalNumpyTheta] + | type[LongitudinalNumpyEta], +] = { + "z": LongitudinalNumpyZ, + "theta": LongitudinalNumpyTheta, + "eta": LongitudinalNumpyEta, +} +_temporal_numpy_type: dict[str, type[TemporalNumpyT] | type[TemporalNumpyTau]] = { + "t": TemporalNumpyT, + "tau": TemporalNumpyTau, +} + + +def _check_dtype_names( + array: VectorNumpy2D | VectorNumpy3D | VectorNumpy4D, dimension: int +) -> tuple[tuple[str, str], ...]: + """ + Validates a structured dtype against the coordinates that ``array``'s class + expects, returning the ``(generic name, given name)`` pairs it found. + """ + if array.dtype.names is None: + raise TypeError( + f"{type(array).__name__} must have a structured dtype containing " + "its coordinates as fields" + ) + _, _, names, _ = _check_coordinate_names( + array.dtype.names, + dimension=dimension, + momentum=type(array)._IS_MOMENTUM, + allow_extra=True, + ) + return names + + class VectorNumpy(Vector, GetItem): # noqa: PLW1641 """Mixin class for NumPy vectors.""" @@ -1193,15 +1234,8 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return - if _has(self, ("x", "y")): - self._azimuthal_type = AzimuthalNumpyXY - elif _has(self, ("rho", "phi")): - self._azimuthal_type = AzimuthalNumpyRhoPhi - else: - raise TypeError( - f"{type(self).__name__} must have a structured dtype containing " - 'fields ("x", "y") or ("rho", "phi")' - ) + names = _check_dtype_names(self, 2) + self._azimuthal_type = _azimuthal_numpy_type[names[0][0]] _is_type_safe(self) @@ -1366,19 +1400,13 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return + names = _check_dtype_names(self, 2) + # Install a fresh dtype on ``self`` rather than mutating the dtype object, # which is shared with the base array and would rename the caller's fields. self.dtype = _momentum_to_generic_dtype(self.dtype) - if _has(self, ("x", "y")): - self._azimuthal_type = AzimuthalNumpyXY - elif _has(self, ("rho", "phi")): - self._azimuthal_type = AzimuthalNumpyRhoPhi - else: - raise TypeError( - f"{type(self).__name__} must have a structured dtype containing " - 'fields ("x", "y") or ("rho", "phi") or ("px", "py") or ("pt", "phi")' - ) + self._azimuthal_type = _azimuthal_numpy_type[names[0][0]] _is_type_safe(self) @@ -1434,26 +1462,9 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return - if _has(self, ("x", "y")): - self._azimuthal_type = AzimuthalNumpyXY - elif _has(self, ("rho", "phi")): - self._azimuthal_type = AzimuthalNumpyRhoPhi - else: - raise TypeError( - f"{type(self).__name__} must have a structured dtype containing " - 'fields ("x", "y") or ("rho", "phi")' - ) - if _has(self, ("z",)): - self._longitudinal_type = LongitudinalNumpyZ - elif _has(self, ("theta",)): - self._longitudinal_type = LongitudinalNumpyTheta - elif _has(self, ("eta",)): - self._longitudinal_type = LongitudinalNumpyEta - else: - raise TypeError( - f"{type(self).__name__} must have a structured dtype containing " - 'field "z" or "theta" or "eta"' - ) + names = _check_dtype_names(self, 3) + self._azimuthal_type = _azimuthal_numpy_type[names[0][0]] + self._longitudinal_type = _longitudinal_numpy_type[names[2][0]] _is_type_safe(self) @@ -1671,29 +1682,14 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return + names = _check_dtype_names(self, 3) + # Install a fresh dtype on ``self`` rather than mutating the dtype object, # which is shared with the base array and would rename the caller's fields. self.dtype = _momentum_to_generic_dtype(self.dtype) - if _has(self, ("x", "y")): - self._azimuthal_type = AzimuthalNumpyXY - elif _has(self, ("rho", "phi")): - self._azimuthal_type = AzimuthalNumpyRhoPhi - else: - raise TypeError( - f"{type(self).__name__} must have a structured dtype containing " - 'fields ("x", "y") or ("rho", "phi") or ("px", "py") or ("pt", "phi")' - ) - if _has(self, ("z",)): - self._longitudinal_type = LongitudinalNumpyZ - elif _has(self, ("theta",)): - self._longitudinal_type = LongitudinalNumpyTheta - elif _has(self, ("eta",)): - self._longitudinal_type = LongitudinalNumpyEta - else: - raise TypeError( - f"{type(self).__name__} must have a structured dtype containing " - 'field "z" or "theta" or "eta" or "pz"' - ) + + self._azimuthal_type = _azimuthal_numpy_type[names[0][0]] + self._longitudinal_type = _longitudinal_numpy_type[names[2][0]] _is_type_safe(self) @@ -1752,37 +1748,10 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return - if _has(self, ("x", "y")): - self._azimuthal_type = AzimuthalNumpyXY - elif _has(self, ("rho", "phi")): - self._azimuthal_type = AzimuthalNumpyRhoPhi - else: - raise TypeError( - f"{type(self).__name__} must have a structured dtype containing " - 'fields ("x", "y") or ("rho", "phi")' - ) - - if _has(self, ("z",)): - self._longitudinal_type = LongitudinalNumpyZ - elif _has(self, ("theta",)): - self._longitudinal_type = LongitudinalNumpyTheta - elif _has(self, ("eta",)): - self._longitudinal_type = LongitudinalNumpyEta - else: - raise TypeError( - f"{type(self).__name__} must have a structured dtype containing " - 'field "z" or "theta" or "eta"' - ) - - if _has(self, ("t",)): - self._temporal_type = TemporalNumpyT - elif _has(self, ("tau",)): - self._temporal_type = TemporalNumpyTau - else: - raise TypeError( - f"{type(self).__name__} must have a structured dtype containing " - 'field "t" or "tau"' - ) + names = _check_dtype_names(self, 4) + self._azimuthal_type = _azimuthal_numpy_type[names[0][0]] + self._longitudinal_type = _longitudinal_numpy_type[names[2][0]] + self._temporal_type = _temporal_numpy_type[names[3][0]] _is_type_safe(self) @@ -2061,41 +2030,15 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return + names = _check_dtype_names(self, 4) + # Install a fresh dtype on ``self`` rather than mutating the dtype object, # which is shared with the base array and would rename the caller's fields. self.dtype = _momentum_to_generic_dtype(self.dtype) - if _has(self, ("x", "y")): - self._azimuthal_type = AzimuthalNumpyXY - elif _has(self, ("rho", "phi")): - self._azimuthal_type = AzimuthalNumpyRhoPhi - else: - raise TypeError( - f"{type(self).__name__} must have a structured dtype containing " - 'fields ("x", "y") or ("rho", "phi") or ("px", "py") or ("pt", "phi")' - ) - - if _has(self, ("z",)): - self._longitudinal_type = LongitudinalNumpyZ - elif _has(self, ("theta",)): - self._longitudinal_type = LongitudinalNumpyTheta - elif _has(self, ("eta",)): - self._longitudinal_type = LongitudinalNumpyEta - else: - raise TypeError( - f"{type(self).__name__} must have a structured dtype containing " - 'field "z" or "theta" or "eta" or "pz"' - ) - - if _has(self, ("t",)): - self._temporal_type = TemporalNumpyT - elif _has(self, ("tau",)): - self._temporal_type = TemporalNumpyTau - else: - raise TypeError( - f"{type(self).__name__} must have a structured dtype containing " - 'field "t" or "tau" or "E" or "e" or "energy" or "M" or "m" or "mass"' - ) + self._azimuthal_type = _azimuthal_numpy_type[names[0][0]] + self._longitudinal_type = _longitudinal_numpy_type[names[2][0]] + self._temporal_type = _temporal_numpy_type[names[3][0]] _is_type_safe(self) @@ -2158,6 +2101,13 @@ def array(*args: typing.Any, **kwargs: typing.Any) -> VectorNumpy: - ``mass`` may be substituted for ``tau`` to make the vector a momentum vector. + + A coordinate may be given only once, whether by its generic name or through + a momentum-alias, and the names must form exactly one of the combinations + above; anything else raises a ``TypeError``. + + Names that are not coordinates become extra fields of the array, which can be + used to carry properties of a particle other than its momentum. """ names: tuple[str, ...] if len(args) == 1 and len(kwargs) == 0 and isinstance(args[0], dict): @@ -2173,11 +2123,11 @@ def array(*args: typing.Any, **kwargs: typing.Any) -> VectorNumpy: cls: type[VectorNumpy] - is_momentum = any(x in _repr_momentum_to_generic for x in names) + is_momentum, dimension, _, _ = _check_coordinate_names(names, allow_extra=True) - if any(x in ("t", "E", "e", "energy", "tau", "M", "m", "mass") for x in names): + if dimension == 4: cls = MomentumNumpy4D if is_momentum else VectorNumpy4D - elif any(x in ("z", "pz", "theta", "eta") for x in names): + elif dimension == 3: cls = MomentumNumpy3D if is_momentum else VectorNumpy3D else: cls = MomentumNumpy2D if is_momentum else VectorNumpy2D diff --git a/src/vector/backends/object.py b/src/vector/backends/object.py index 86c0cc23..d9e56c40 100644 --- a/src/vector/backends/object.py +++ b/src/vector/backends/object.py @@ -44,6 +44,7 @@ LongitudinalZ, Lorentz, LorentzMomentum, + Momentum, Planar, PlanarMomentum, SameVectorType, @@ -58,11 +59,12 @@ Vector4D, VectorProtocol, _aztype, + _check_coordinate_names, _coordinate_class_to_names, + _generic_coordinates, _handler_of, _ltype, _repr_generic_to_momentum, - _repr_momentum_to_generic, _ttype, ) from vector._typeutils import FloatArray @@ -293,6 +295,29 @@ def elements(self) -> tuple[float]: } +def _azimuthal_object(coordinates: dict[str, float]) -> AzimuthalObject: + """Builds an azimuthal object from validated, generically named coordinates.""" + if "x" in coordinates: + return AzimuthalObjectXY(coordinates["x"], coordinates["y"]) + return AzimuthalObjectRhoPhi(coordinates["rho"], coordinates["phi"]) + + +def _longitudinal_object(coordinates: dict[str, float]) -> LongitudinalObject: + """Builds a longitudinal object from validated, generically named coordinates.""" + if "z" in coordinates: + return LongitudinalObjectZ(coordinates["z"]) + if "theta" in coordinates: + return LongitudinalObjectTheta(coordinates["theta"]) + return LongitudinalObjectEta(coordinates["eta"]) + + +def _temporal_object(coordinates: dict[str, float]) -> TemporalObject: + """Builds a temporal object from validated, generically named coordinates.""" + if "t" in coordinates: + return TemporalObjectT(coordinates["t"]) + return TemporalObjectTau(coordinates["tau"]) + + def _replace_data(obj: typing.Any, result: typing.Any) -> typing.Any: if not isinstance(result, VectorObject): raise TypeError(f"can only assign a single vector to {type(obj).__name__}") @@ -683,25 +708,11 @@ def __init__( ) -> None: _is_type_safe(kwargs) - for k, v in kwargs.copy().items(): - kwargs.pop(k) - kwargs[_repr_momentum_to_generic.get(k, k)] = v - if not kwargs and azimuthal is not None: self.azimuthal = azimuthal elif kwargs and azimuthal is None: - if set(kwargs) == {"x", "y"}: - self.azimuthal = AzimuthalObjectXY(kwargs["x"], kwargs["y"]) - elif set(kwargs) == {"rho", "phi"}: - self.azimuthal = AzimuthalObjectRhoPhi(kwargs["rho"], kwargs["phi"]) - else: - complaint = """unrecognized combination of coordinates, allowed combinations are:\n - x= y= - rho= phi=""".replace(" ", " ") - if type(self) is VectorObject2D: - raise TypeError(complaint) - else: - raise TypeError(f"{complaint}\n\nor their momentum equivalents") + coordinates = _generic_coordinates(kwargs, 2, isinstance(self, Momentum)) + self.azimuthal = _azimuthal_object(coordinates) else: raise TypeError("must give Azimuthal if not giving keyword arguments") @@ -1071,44 +1082,13 @@ def __init__( ) -> None: _is_type_safe(kwargs) - for k, v in kwargs.copy().items(): - kwargs.pop(k) - kwargs[_repr_momentum_to_generic.get(k, k)] = v - if not kwargs and azimuthal is not None and longitudinal is not None: self.azimuthal = azimuthal self.longitudinal = longitudinal elif kwargs and azimuthal is None and longitudinal is None: - if set(kwargs) == {"x", "y", "z"}: - self.azimuthal = AzimuthalObjectXY(kwargs["x"], kwargs["y"]) - self.longitudinal = LongitudinalObjectZ(kwargs["z"]) - elif set(kwargs) == {"x", "y", "eta"}: - self.azimuthal = AzimuthalObjectXY(kwargs["x"], kwargs["y"]) - self.longitudinal = LongitudinalObjectEta(kwargs["eta"]) - elif set(kwargs) == {"x", "y", "theta"}: - self.azimuthal = AzimuthalObjectXY(kwargs["x"], kwargs["y"]) - self.longitudinal = LongitudinalObjectTheta(kwargs["theta"]) - elif set(kwargs) == {"rho", "phi", "z"}: - self.azimuthal = AzimuthalObjectRhoPhi(kwargs["rho"], kwargs["phi"]) - self.longitudinal = LongitudinalObjectZ(kwargs["z"]) - elif set(kwargs) == {"rho", "phi", "eta"}: - self.azimuthal = AzimuthalObjectRhoPhi(kwargs["rho"], kwargs["phi"]) - self.longitudinal = LongitudinalObjectEta(kwargs["eta"]) - elif set(kwargs) == {"rho", "phi", "theta"}: - self.azimuthal = AzimuthalObjectRhoPhi(kwargs["rho"], kwargs["phi"]) - self.longitudinal = LongitudinalObjectTheta(kwargs["theta"]) - else: - complaint = """unrecognized combination of coordinates, allowed combinations are:\n - x= y= z= - x= y= theta= - x= y= eta= - rho= phi= z= - rho= phi= theta= - rho= phi= eta=""".replace(" ", " ") - if type(self) is VectorObject3D: - raise TypeError(complaint) - else: - raise TypeError(f"{complaint}\n\nor their momentum equivalents") + coordinates = _generic_coordinates(kwargs, 3, isinstance(self, Momentum)) + self.azimuthal = _azimuthal_object(coordinates) + self.longitudinal = _longitudinal_object(coordinates) else: raise TypeError( "must give Azimuthal and Longitudinal if not giving keyword arguments" @@ -1764,10 +1744,6 @@ def __init__( ) -> None: _is_type_safe(kwargs) - for k, v in kwargs.copy().items(): - kwargs.pop(k) - kwargs[_repr_momentum_to_generic.get(k, k)] = v - if ( not kwargs and azimuthal is not None @@ -1778,72 +1754,10 @@ def __init__( self.longitudinal = longitudinal self.temporal = temporal elif kwargs and azimuthal is None and longitudinal is None and temporal is None: - if set(kwargs) == {"x", "y", "z", "t"}: - self.azimuthal = AzimuthalObjectXY(kwargs["x"], kwargs["y"]) - self.longitudinal = LongitudinalObjectZ(kwargs["z"]) - self.temporal = TemporalObjectT(kwargs["t"]) - elif set(kwargs) == {"x", "y", "eta", "t"}: - self.azimuthal = AzimuthalObjectXY(kwargs["x"], kwargs["y"]) - self.longitudinal = LongitudinalObjectEta(kwargs["eta"]) - self.temporal = TemporalObjectT(kwargs["t"]) - elif set(kwargs) == {"x", "y", "theta", "t"}: - self.azimuthal = AzimuthalObjectXY(kwargs["x"], kwargs["y"]) - self.longitudinal = LongitudinalObjectTheta(kwargs["theta"]) - self.temporal = TemporalObjectT(kwargs["t"]) - elif set(kwargs) == {"rho", "phi", "z", "t"}: - self.azimuthal = AzimuthalObjectRhoPhi(kwargs["rho"], kwargs["phi"]) - self.longitudinal = LongitudinalObjectZ(kwargs["z"]) - self.temporal = TemporalObjectT(kwargs["t"]) - elif set(kwargs) == {"rho", "phi", "eta", "t"}: - self.azimuthal = AzimuthalObjectRhoPhi(kwargs["rho"], kwargs["phi"]) - self.longitudinal = LongitudinalObjectEta(kwargs["eta"]) - self.temporal = TemporalObjectT(kwargs["t"]) - elif set(kwargs) == {"rho", "phi", "theta", "t"}: - self.azimuthal = AzimuthalObjectRhoPhi(kwargs["rho"], kwargs["phi"]) - self.longitudinal = LongitudinalObjectTheta(kwargs["theta"]) - self.temporal = TemporalObjectT(kwargs["t"]) - elif set(kwargs) == {"x", "y", "z", "tau"}: - self.azimuthal = AzimuthalObjectXY(kwargs["x"], kwargs["y"]) - self.longitudinal = LongitudinalObjectZ(kwargs["z"]) - self.temporal = TemporalObjectTau(kwargs["tau"]) - elif set(kwargs) == {"x", "y", "eta", "tau"}: - self.azimuthal = AzimuthalObjectXY(kwargs["x"], kwargs["y"]) - self.longitudinal = LongitudinalObjectEta(kwargs["eta"]) - self.temporal = TemporalObjectTau(kwargs["tau"]) - elif set(kwargs) == {"x", "y", "theta", "tau"}: - self.azimuthal = AzimuthalObjectXY(kwargs["x"], kwargs["y"]) - self.longitudinal = LongitudinalObjectTheta(kwargs["theta"]) - self.temporal = TemporalObjectTau(kwargs["tau"]) - elif set(kwargs) == {"rho", "phi", "z", "tau"}: - self.azimuthal = AzimuthalObjectRhoPhi(kwargs["rho"], kwargs["phi"]) - self.longitudinal = LongitudinalObjectZ(kwargs["z"]) - self.temporal = TemporalObjectTau(kwargs["tau"]) - elif set(kwargs) == {"rho", "phi", "eta", "tau"}: - self.azimuthal = AzimuthalObjectRhoPhi(kwargs["rho"], kwargs["phi"]) - self.longitudinal = LongitudinalObjectEta(kwargs["eta"]) - self.temporal = TemporalObjectTau(kwargs["tau"]) - elif set(kwargs) == {"rho", "phi", "theta", "tau"}: - self.azimuthal = AzimuthalObjectRhoPhi(kwargs["rho"], kwargs["phi"]) - self.longitudinal = LongitudinalObjectTheta(kwargs["theta"]) - self.temporal = TemporalObjectTau(kwargs["tau"]) - else: - complaint = """unrecognized combination of coordinates, allowed combinations are:\n - x= y= z= tau= - x= y= theta= t= - x= y= theta= tau= - x= y= eta= t= - x= y= z= t= - x= y= eta= tau= - rho= phi= z= t= - rho= phi= z= tau= - rho= phi= theta= t= - rho= phi= theta= tau= - rho= phi= eta= t= - rho= phi= eta= tau=""".replace(" ", " ") - if type(self) is VectorObject4D: - raise TypeError(complaint) - else: - raise TypeError(f"{complaint}\n\nor their momentum equivalents") + coordinates = _generic_coordinates(kwargs, 4, isinstance(self, Momentum)) + self.azimuthal = _azimuthal_object(coordinates) + self.longitudinal = _longitudinal_object(coordinates) + self.temporal = _temporal_object(coordinates) else: raise TypeError( "must give Azimuthal, Longitudinal, and Temporal if not giving keyword arguments" @@ -2200,84 +2114,27 @@ def _gather_coordinates( spatial_class: type[VectorObject3D], lorentz_class: type[VectorObject4D], coordinates: dict[str, typing.Any], + dimension: int, ) -> typing.Any: """ Helper function for :func:`vector.backends.object.obj`. Constructs and returns a 2D, 3D, or 4D ``VectorObject`` or ``MomentumObject`` with - the provided coordinates (dictionary), planar (``VectorObject2D`` or ``MomentumObject2D``), - spatial (``VectorObject3D`` or ``MomentumObject3D``), and lorentz - (``VectorObject4D`` or ``MomentumObject4D``) classes. + the provided (validated, generically named) coordinates, planar (``VectorObject2D`` + or ``MomentumObject2D``), spatial (``VectorObject3D`` or ``MomentumObject3D``), and + lorentz (``VectorObject4D`` or ``MomentumObject4D``) classes. """ - azimuthal: None | (AzimuthalObjectXY | AzimuthalObjectRhoPhi) = None - - if "x" in coordinates and "y" in coordinates: - if "rho" in coordinates or "phi" in coordinates: - raise TypeError("specify x= and y= or rho= and phi=, but not both") - azimuthal = AzimuthalObjectXY(coordinates.pop("x"), coordinates.pop("y")) - elif "rho" in coordinates and "phi" in coordinates: - if "x" in coordinates or "y" in coordinates: - raise TypeError("specify x= and y= or rho= and phi=, but not both") - azimuthal = AzimuthalObjectRhoPhi( - coordinates.pop("rho"), coordinates.pop("phi") + if dimension == 2: + return planar_class(azimuthal=_azimuthal_object(coordinates)) + if dimension == 3: + return spatial_class( + azimuthal=_azimuthal_object(coordinates), + longitudinal=_longitudinal_object(coordinates), ) - - longitudinal: None | ( - LongitudinalObjectZ | LongitudinalObjectTheta | LongitudinalObjectEta - ) = None - - if "z" in coordinates: - if "theta" in coordinates or "eta" in coordinates: - raise TypeError("specify z= or theta= or eta=, but not more than one") - longitudinal = LongitudinalObjectZ(coordinates.pop("z")) - elif "theta" in coordinates: - if "eta" in coordinates: - raise TypeError("specify z= or theta= or eta=, but not more than one") - longitudinal = LongitudinalObjectTheta(coordinates.pop("theta")) - elif "eta" in coordinates: - longitudinal = LongitudinalObjectEta(coordinates.pop("eta")) - - temporal: TemporalObjectT | TemporalObjectTau | None = None - - if "t" in coordinates: - if "tau" in coordinates: - raise TypeError("specify t= or tau=, but not more than one") - temporal = TemporalObjectT(coordinates.pop("t")) - elif "tau" in coordinates: - temporal = TemporalObjectTau(coordinates.pop("tau")) - - if not coordinates: - if azimuthal is not None and longitudinal is None and temporal is None: - return planar_class(azimuthal=azimuthal) - if azimuthal is not None and longitudinal is not None and temporal is None: - return spatial_class(azimuthal=azimuthal, longitudinal=longitudinal) - if azimuthal is not None and longitudinal is not None and temporal is not None: - return lorentz_class( - azimuthal=azimuthal, longitudinal=longitudinal, temporal=temporal - ) - - raise TypeError( - "unrecognized combination of coordinates, allowed combinations are:\n\n" - " (2D) x= y=\n" - " (2D) rho= phi=\n" - " (3D) x= y= z=\n" - " (3D) x= y= theta=\n" - " (3D) x= y= eta=\n" - " (3D) rho= phi= z=\n" - " (3D) rho= phi= theta=\n" - " (3D) rho= phi= eta=\n" - " (4D) x= y= z= t=\n" - " (4D) x= y= z= tau=\n" - " (4D) x= y= theta= t=\n" - " (4D) x= y= theta= tau=\n" - " (4D) x= y= eta= t=\n" - " (4D) x= y= eta= tau=\n" - " (4D) rho= phi= z= t=\n" - " (4D) rho= phi= z= tau=\n" - " (4D) rho= phi= theta= t=\n" - " (4D) rho= phi= theta= tau=\n" - " (4D) rho= phi= eta= t=\n" - " (4D) rho= phi= eta= tau=" + return lorentz_class( + azimuthal=_azimuthal_object(coordinates), + longitudinal=_longitudinal_object(coordinates), + temporal=_temporal_object(coordinates), ) @@ -3212,6 +3069,10 @@ def obj(**coordinates: float) -> VectorObject: to make the vector a momentum vector. + A coordinate may be given only once, whether by its generic name or through + a momentum-alias, and the names must form exactly one of the combinations + above; anything else raises a ``TypeError``. + Alternatively, the :class:`vector.VectorObject2D`, :class:`vector.VectorObject3D`, and :class:`vector.VectorObject4D` classes (with momentum @@ -3240,64 +3101,26 @@ def obj(**coordinates: float) -> VectorObject: - :meth:`vector.VectorObject4D.from_rhophietat` - :meth:`vector.VectorObject4D.from_rhophietatau` """ - is_momentum = False - generic_coordinates = {} - _is_type_safe(coordinates) - if "px" in coordinates: - is_momentum = True - generic_coordinates["x"] = coordinates.pop("px") - if "py" in coordinates: - is_momentum = True - generic_coordinates["y"] = coordinates.pop("py") - if "pt" in coordinates: - is_momentum = True - generic_coordinates["rho"] = coordinates.pop("pt") - if "pz" in coordinates: - is_momentum = True - generic_coordinates["z"] = coordinates.pop("pz") - if "E" in coordinates: - is_momentum = True - generic_coordinates["t"] = coordinates.pop("E") - if "e" in coordinates: - is_momentum = True - if "t" in generic_coordinates: - raise TypeError( - "duplicate coordinates (through momentum-aliases): 'e' and 'E' both map to 't'" - ) - generic_coordinates["t"] = coordinates.pop("e") - if "energy" in coordinates and "t" not in generic_coordinates: - is_momentum = True - generic_coordinates["t"] = coordinates.pop("energy") - if "M" in coordinates: - is_momentum = True - generic_coordinates["tau"] = coordinates.pop("M") - if "m" in coordinates: - is_momentum = True - if "tau" in generic_coordinates: - raise TypeError( - "duplicate coordinates (through momentum-aliases): 'm' and 'M' both map to 'tau'" - ) - generic_coordinates["tau"] = coordinates.pop("m") - if "mass" in coordinates and "tau" not in generic_coordinates: - is_momentum = True - generic_coordinates["tau"] = coordinates.pop("mass") - for x in list(coordinates): - if x not in generic_coordinates: - generic_coordinates[x] = coordinates.pop(x) - if len(coordinates) != 0: - raise TypeError( - "duplicate coordinates (through momentum-aliases): " - + ", ".join(repr(x) for x in coordinates) - ) + is_momentum, dimension, names, _ = _check_coordinate_names(tuple(coordinates)) + generic_coordinates = {name: coordinates[given] for name, given in names} + if is_momentum: return _gather_coordinates( - MomentumObject2D, MomentumObject3D, MomentumObject4D, generic_coordinates + MomentumObject2D, + MomentumObject3D, + MomentumObject4D, + generic_coordinates, + dimension, ) else: return _gather_coordinates( - VectorObject2D, VectorObject3D, VectorObject4D, generic_coordinates + VectorObject2D, + VectorObject3D, + VectorObject4D, + generic_coordinates, + dimension, ) diff --git a/src/vector/backends/sympy.py b/src/vector/backends/sympy.py index 7fe95ab8..0cc91e00 100644 --- a/src/vector/backends/sympy.py +++ b/src/vector/backends/sympy.py @@ -32,6 +32,7 @@ LongitudinalZ, Lorentz, LorentzMomentum, + Momentum, Planar, PlanarMomentum, SameVectorType, @@ -47,10 +48,10 @@ VectorProtocol, _aztype, _coordinate_class_to_names, + _generic_coordinates, _handler_of, _ltype, _repr_generic_to_momentum, - _repr_momentum_to_generic, _ttype, ) @@ -416,6 +417,29 @@ def elements(self) -> tuple[sympy.Symbol]: } +def _azimuthal_sympy(coordinates: dict[str, sympy.Symbol]) -> AzimuthalSympy: + """Builds an azimuthal object from validated, generically named coordinates.""" + if "x" in coordinates: + return AzimuthalSympyXY(coordinates["x"], coordinates["y"]) + return AzimuthalSympyRhoPhi(coordinates["rho"], coordinates["phi"]) + + +def _longitudinal_sympy(coordinates: dict[str, sympy.Symbol]) -> LongitudinalSympy: + """Builds a longitudinal object from validated, generically named coordinates.""" + if "z" in coordinates: + return LongitudinalSympyZ(coordinates["z"]) + if "theta" in coordinates: + return LongitudinalSympyTheta(coordinates["theta"]) + return LongitudinalSympyEta(coordinates["eta"]) + + +def _temporal_sympy(coordinates: dict[str, sympy.Symbol]) -> TemporalSympy: + """Builds a temporal object from validated, generically named coordinates.""" + if "t" in coordinates: + return TemporalSympyT(coordinates["t"]) + return TemporalSympyTau(coordinates["tau"]) + + def _is_type_safe(coordinates: dict[str, typing.Any]) -> None: if not all(isinstance(coord, sympy.Expr) for coord in coordinates.values()): raise TypeError("coordinates must be a sympy expression") @@ -761,26 +785,12 @@ class VectorSympy2D(VectorSympy, Planar, Vector2D): azimuthal: AzimuthalSympy def __init__(self, azimuthal: AzimuthalSympy | None = None, **kwargs: sympy.Symbol): - for k, v in kwargs.copy().items(): - kwargs.pop(k) - kwargs[_repr_momentum_to_generic.get(k, k)] = v - if not kwargs and azimuthal is not None: self.azimuthal = azimuthal elif kwargs and azimuthal is None: _is_type_safe(kwargs) - if set(kwargs) == {"x", "y"}: - self.azimuthal = AzimuthalSympyXY(kwargs["x"], kwargs["y"]) - elif set(kwargs) == {"rho", "phi"}: - self.azimuthal = AzimuthalSympyRhoPhi(kwargs["rho"], kwargs["phi"]) - else: - complaint = """unrecognized combination of coordinates, allowed combinations are:\n - x= y= - rho= phi=""".replace(" ", " ") - if type(self) is VectorSympy2D: - raise TypeError(complaint) - else: - raise TypeError(f"{complaint}\n\nor their momentum equivalents") + coordinates = _generic_coordinates(kwargs, 2, isinstance(self, Momentum)) + self.azimuthal = _azimuthal_sympy(coordinates) else: raise TypeError("must give Azimuthal if not giving keyword arguments") @@ -962,45 +972,14 @@ def __init__( longitudinal: LongitudinalSympy | None = None, **kwargs: sympy.Symbol, ): - for k, v in kwargs.copy().items(): - kwargs.pop(k) - kwargs[_repr_momentum_to_generic.get(k, k)] = v - if not kwargs and azimuthal is not None and longitudinal is not None: self.azimuthal = azimuthal self.longitudinal = longitudinal elif kwargs and azimuthal is None and longitudinal is None: _is_type_safe(kwargs) - if set(kwargs) == {"x", "y", "z"}: - self.azimuthal = AzimuthalSympyXY(kwargs["x"], kwargs["y"]) - self.longitudinal = LongitudinalSympyZ(kwargs["z"]) - elif set(kwargs) == {"x", "y", "eta"}: - self.azimuthal = AzimuthalSympyXY(kwargs["x"], kwargs["y"]) - self.longitudinal = LongitudinalSympyEta(kwargs["eta"]) - elif set(kwargs) == {"x", "y", "theta"}: - self.azimuthal = AzimuthalSympyXY(kwargs["x"], kwargs["y"]) - self.longitudinal = LongitudinalSympyTheta(kwargs["theta"]) - elif set(kwargs) == {"rho", "phi", "z"}: - self.azimuthal = AzimuthalSympyRhoPhi(kwargs["rho"], kwargs["phi"]) - self.longitudinal = LongitudinalSympyZ(kwargs["z"]) - elif set(kwargs) == {"rho", "phi", "eta"}: - self.azimuthal = AzimuthalSympyRhoPhi(kwargs["rho"], kwargs["phi"]) - self.longitudinal = LongitudinalSympyEta(kwargs["eta"]) - elif set(kwargs) == {"rho", "phi", "theta"}: - self.azimuthal = AzimuthalSympyRhoPhi(kwargs["rho"], kwargs["phi"]) - self.longitudinal = LongitudinalSympyTheta(kwargs["theta"]) - else: - complaint = """unrecognized combination of coordinates, allowed combinations are:\n - x= y= z= - x= y= theta= - x= y= eta= - rho= phi= z= - rho= phi= theta= - rho= phi= eta=""".replace(" ", " ") - if type(self) is VectorSympy3D: - raise TypeError(complaint) - else: - raise TypeError(f"{complaint}\n\nor their momentum equivalents") + coordinates = _generic_coordinates(kwargs, 3, isinstance(self, Momentum)) + self.azimuthal = _azimuthal_sympy(coordinates) + self.longitudinal = _longitudinal_sympy(coordinates) else: raise TypeError( "must give Azimuthal and Longitudinal if not giving keyword arguments" @@ -1236,10 +1215,6 @@ def __init__( temporal: TemporalSympy | None = None, **kwargs: sympy.Symbol, ): - for k, v in kwargs.copy().items(): - kwargs.pop(k) - kwargs[_repr_momentum_to_generic.get(k, k)] = v - if ( not kwargs and azimuthal is not None @@ -1251,72 +1226,10 @@ def __init__( self.temporal = temporal elif kwargs and azimuthal is None and longitudinal is None and temporal is None: _is_type_safe(kwargs) - if set(kwargs) == {"x", "y", "z", "t"}: - self.azimuthal = AzimuthalSympyXY(kwargs["x"], kwargs["y"]) - self.longitudinal = LongitudinalSympyZ(kwargs["z"]) - self.temporal = TemporalSympyT(kwargs["t"]) - elif set(kwargs) == {"x", "y", "eta", "t"}: - self.azimuthal = AzimuthalSympyXY(kwargs["x"], kwargs["y"]) - self.longitudinal = LongitudinalSympyEta(kwargs["eta"]) - self.temporal = TemporalSympyT(kwargs["t"]) - elif set(kwargs) == {"x", "y", "theta", "t"}: - self.azimuthal = AzimuthalSympyXY(kwargs["x"], kwargs["y"]) - self.longitudinal = LongitudinalSympyTheta(kwargs["theta"]) - self.temporal = TemporalSympyT(kwargs["t"]) - elif set(kwargs) == {"rho", "phi", "z", "t"}: - self.azimuthal = AzimuthalSympyRhoPhi(kwargs["rho"], kwargs["phi"]) - self.longitudinal = LongitudinalSympyZ(kwargs["z"]) - self.temporal = TemporalSympyT(kwargs["t"]) - elif set(kwargs) == {"rho", "phi", "eta", "t"}: - self.azimuthal = AzimuthalSympyRhoPhi(kwargs["rho"], kwargs["phi"]) - self.longitudinal = LongitudinalSympyEta(kwargs["eta"]) - self.temporal = TemporalSympyT(kwargs["t"]) - elif set(kwargs) == {"rho", "phi", "theta", "t"}: - self.azimuthal = AzimuthalSympyRhoPhi(kwargs["rho"], kwargs["phi"]) - self.longitudinal = LongitudinalSympyTheta(kwargs["theta"]) - self.temporal = TemporalSympyT(kwargs["t"]) - elif set(kwargs) == {"x", "y", "z", "tau"}: - self.azimuthal = AzimuthalSympyXY(kwargs["x"], kwargs["y"]) - self.longitudinal = LongitudinalSympyZ(kwargs["z"]) - self.temporal = TemporalSympyTau(kwargs["tau"]) - elif set(kwargs) == {"x", "y", "eta", "tau"}: - self.azimuthal = AzimuthalSympyXY(kwargs["x"], kwargs["y"]) - self.longitudinal = LongitudinalSympyEta(kwargs["eta"]) - self.temporal = TemporalSympyTau(kwargs["tau"]) - elif set(kwargs) == {"x", "y", "theta", "tau"}: - self.azimuthal = AzimuthalSympyXY(kwargs["x"], kwargs["y"]) - self.longitudinal = LongitudinalSympyTheta(kwargs["theta"]) - self.temporal = TemporalSympyTau(kwargs["tau"]) - elif set(kwargs) == {"rho", "phi", "z", "tau"}: - self.azimuthal = AzimuthalSympyRhoPhi(kwargs["rho"], kwargs["phi"]) - self.longitudinal = LongitudinalSympyZ(kwargs["z"]) - self.temporal = TemporalSympyTau(kwargs["tau"]) - elif set(kwargs) == {"rho", "phi", "eta", "tau"}: - self.azimuthal = AzimuthalSympyRhoPhi(kwargs["rho"], kwargs["phi"]) - self.longitudinal = LongitudinalSympyEta(kwargs["eta"]) - self.temporal = TemporalSympyTau(kwargs["tau"]) - elif set(kwargs) == {"rho", "phi", "theta", "tau"}: - self.azimuthal = AzimuthalSympyRhoPhi(kwargs["rho"], kwargs["phi"]) - self.longitudinal = LongitudinalSympyTheta(kwargs["theta"]) - self.temporal = TemporalSympyTau(kwargs["tau"]) - else: - complaint = """unrecognized combination of coordinates, allowed combinations are:\n - x= y= z= tau= - x= y= theta= t= - x= y= theta= tau= - x= y= eta= t= - x= y= z= t= - x= y= eta= tau= - rho= phi= z= t= - rho= phi= z= tau= - rho= phi= theta= t= - rho= phi= theta= tau= - rho= phi= eta= t= - rho= phi= eta= tau=""".replace(" ", " ") - if type(self) is VectorSympy4D: - raise TypeError(complaint) - else: - raise TypeError(f"{complaint}\n\nor their momentum equivalents") + coordinates = _generic_coordinates(kwargs, 4, isinstance(self, Momentum)) + self.azimuthal = _azimuthal_sympy(coordinates) + self.longitudinal = _longitudinal_sympy(coordinates) + self.temporal = _temporal_sympy(coordinates) else: raise TypeError( "must give Azimuthal, Longitudinal, and Temporal if not giving keyword arguments" diff --git a/tests/test_pr_659.py b/tests/test_pr_659.py new file mode 100644 index 00000000..8f914802 --- /dev/null +++ b/tests/test_pr_659.py @@ -0,0 +1,612 @@ +# Copyright (c) 2019, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. +# +# Distributed under the 3-clause BSD license, see accompanying file LICENSE +# or https://github.com/scikit-hep/vector for details. + +""" +Exhaustive tests of the coordinate names that every backend accepts, added with +https://github.com/scikit-hep/vector/pull/659. + +The rules are implemented once, in ``vector._methods._check_coordinate_names``, +so the tests below check every backend against the same independently written +reference implementation (:func:`reference`) and against each other. +""" + +from __future__ import annotations + +import itertools + +import numpy as np +import pytest + +import vector + +# Momentum-aliases; ``phi``, ``theta``, and ``eta`` have none. +ALIASES = { + "px": "x", + "py": "y", + "pt": "rho", + "pz": "z", + "E": "t", + "e": "t", + "energy": "t", + "M": "tau", + "m": "tau", + "mass": "tau", +} + +GENERIC_NAMES = ("x", "y", "rho", "phi", "z", "theta", "eta", "t", "tau") + +ALIAS_CHOICES = { + "x": ("x", "px"), + "y": ("y", "py"), + "rho": ("rho", "pt"), + "phi": ("phi",), + "z": ("z", "pz"), + "theta": ("theta",), + "eta": ("eta",), + "t": ("t", "E", "e", "energy"), + "tau": ("tau", "M", "m", "mass"), +} + + +def reference(names: tuple[str, ...]) -> tuple[bool, int] | None: + """ + Independent implementation of the rules: returns ``(is_momentum, dimension)`` + for a combination of coordinate names that describes exactly one vector and + None for one that does not. + """ + generic = [ALIASES.get(name, name) for name in names] + if len(set(generic)) != len(generic): + return None + + remaining = set(generic) + for azimuthal in ({"x", "y"}, {"rho", "phi"}): + if azimuthal <= remaining: + remaining -= azimuthal + break + else: + return None + + dimension = 2 + for longitudinal in ("z", "theta", "eta"): + if longitudinal in remaining: + remaining.remove(longitudinal) + dimension = 3 + break + if dimension == 3: + for temporal in ("t", "tau"): + if temporal in remaining: + remaining.remove(temporal) + dimension = 4 + break + + if remaining: + return None + + return any(name in ALIASES for name in names), dimension + + +def is_valid_for( + names: tuple[str, ...], dimension: int | None = None, momentum: bool | None = None +) -> bool: + """ + Whether ``names`` may be given to a constructor that is restricted to a + ``dimension`` and/or to (non-)momentum coordinates. + """ + expected = reference(names) + if expected is None: + return False + is_momentum, expected_dimension = expected + if dimension is not None and dimension != expected_dimension: + return False + return not (is_momentum and momentum is False) + + +# All 511 non-empty subsets of the generic coordinate names, of which 20 are vectors. +GENERIC_COMBINATIONS = [ + names + for size in range(1, len(GENERIC_NAMES) + 1) + for names in itertools.combinations(GENERIC_NAMES, size) +] + +VECTOR_COMBINATIONS = [names for names in GENERIC_COMBINATIONS if reference(names)] + +# The 6 + 24 + 192 combinations of coordinate names, including momentum-aliases. +ALIAS_COMBINATIONS = [ + names + for generic in VECTOR_COMBINATIONS + for names in itertools.product(*(ALIAS_CHOICES[name] for name in generic)) +] + + +def values(names: tuple[str, ...]) -> dict[str, float]: + """A distinct value for each coordinate, so that mix-ups are visible.""" + return {name: 1.0 + i for i, name in enumerate(names)} + + +def test_combinations_are_exhaustive(): + assert len(GENERIC_COMBINATIONS) == 511 + assert len(VECTOR_COMBINATIONS) == 2 + 6 + 12 + assert len(ALIAS_COMBINATIONS) == 6 + 24 + 192 + + +object_classes = { + (2, False): vector.VectorObject2D, + (3, False): vector.VectorObject3D, + (4, False): vector.VectorObject4D, + (2, True): vector.MomentumObject2D, + (3, True): vector.MomentumObject3D, + (4, True): vector.MomentumObject4D, +} + +numpy_classes = { + (2, False): vector.VectorNumpy2D, + (3, False): vector.VectorNumpy3D, + (4, False): vector.VectorNumpy4D, + (2, True): vector.MomentumNumpy2D, + (3, True): vector.MomentumNumpy3D, + (4, True): vector.MomentumNumpy4D, +} + + +def numpy_array(coordinates: dict[str, float]) -> np.ndarray: + return np.array( + [tuple(coordinates.values())], + dtype=[(name, np.float64) for name in coordinates], + ) + + +def record_name(names: tuple[str, ...]) -> str: + is_momentum, dimension = reference(names) + return f"{'Momentum' if is_momentum else 'Vector'}{dimension}D" + + +@pytest.mark.parametrize("names", GENERIC_COMBINATIONS) +def test_obj_generic_combinations(names): + coordinates = values(names) + + if reference(names) is None: + with pytest.raises(TypeError): + vector.obj(**coordinates) + else: + vec = vector.obj(**coordinates) + assert isinstance(vec, object_classes[reference(names)[1], False]) + for name, value in coordinates.items(): + assert getattr(vec, name) == pytest.approx(value) + + +@pytest.mark.parametrize("names", GENERIC_COMBINATIONS) +def test_object_class_generic_combinations(names): + coordinates = values(names) + + for (dimension, momentum), cls in object_classes.items(): + if is_valid_for(names, dimension, momentum): + vec = cls(**coordinates) + for name, value in coordinates.items(): + assert getattr(vec, name) == pytest.approx(value) + else: + with pytest.raises(TypeError): + cls(**coordinates) + + +@pytest.mark.parametrize("names", GENERIC_COMBINATIONS) +def test_array_generic_combinations(names): + coordinates = values(names) + + if reference(names) is None: + with pytest.raises(TypeError): + vector.array( + {name: np.array([value]) for name, value in coordinates.items()} + ) + else: + vec = vector.array( + {name: np.array([value]) for name, value in coordinates.items()} + ) + assert isinstance(vec, numpy_classes[reference(names)[1], False]) + for name, value in coordinates.items(): + assert getattr(vec, name)[0] == pytest.approx(value) + + +@pytest.mark.parametrize("names", GENERIC_COMBINATIONS) +def test_numpy_class_generic_combinations(names): + array = numpy_array(values(names)) + + for (dimension, momentum), cls in numpy_classes.items(): + if is_valid_for(names, dimension, momentum): + vec = array.view(cls) + for name, value in values(names).items(): + assert vec[name][0] == pytest.approx(value) + else: + with pytest.raises(TypeError): + array.view(cls) + + +@pytest.mark.parametrize("names", GENERIC_COMBINATIONS) +def test_awkward_generic_combinations(names): + ak = pytest.importorskip("awkward") + coordinates = values(names) + columns = {name: [value] for name, value in coordinates.items()} + + if reference(names) is None: + with pytest.raises(TypeError): + vector.Array(ak.Array(columns)) + with pytest.raises(TypeError): + vector.zip(columns) + else: + for vec in (vector.Array(ak.Array(columns)), vector.zip(columns)): + assert vec.layout.purelist_parameter("__record__") == record_name(names) + for name, value in coordinates.items(): + assert getattr(vec, name)[0] == pytest.approx(value) + + +@pytest.mark.parametrize("names", GENERIC_COMBINATIONS) +def test_sympy_generic_combinations(names): + sympy = pytest.importorskip("sympy") + coordinates = {name: sympy.Symbol(name) for name in names} + + sympy_classes = { + (2, False): vector.VectorSympy2D, + (3, False): vector.VectorSympy3D, + (4, False): vector.VectorSympy4D, + (2, True): vector.MomentumSympy2D, + (3, True): vector.MomentumSympy3D, + (4, True): vector.MomentumSympy4D, + } + + for (dimension, momentum), cls in sympy_classes.items(): + if is_valid_for(names, dimension, momentum): + vec = cls(**coordinates) + for name, symbol in coordinates.items(): + assert getattr(vec, name) == symbol + else: + with pytest.raises(TypeError): + cls(**coordinates) + + +@pytest.mark.parametrize("names", ALIAS_COMBINATIONS) +def test_alias_combinations(names): + is_momentum, dimension = reference(names) + coordinates = values(names) + + vec = vector.obj(**coordinates) + assert isinstance(vec, object_classes[dimension, is_momentum]) + + array = vector.array( + {name: np.array([value]) for name, value in coordinates.items()} + ) + assert isinstance(array, numpy_classes[dimension, is_momentum]) + + for name, value in coordinates.items(): + assert getattr(vec, name) == pytest.approx(value) + assert getattr(array, name)[0] == pytest.approx(value) + # the generic name of an alias reads back the same coordinate + generic = ALIASES.get(name, name) + assert getattr(vec, generic) == pytest.approx(value) + assert getattr(array, generic)[0] == pytest.approx(value) + + +@pytest.mark.parametrize("names", ALIAS_COMBINATIONS) +def test_alias_combinations_awkward(names): + ak = pytest.importorskip("awkward") + coordinates = values(names) + columns = {name: [value] for name, value in coordinates.items()} + + for vec in (vector.Array(ak.Array(columns)), vector.zip(columns)): + assert vec.layout.purelist_parameter("__record__") == record_name(names) + for name, value in coordinates.items(): + assert getattr(vec, name)[0] == pytest.approx(value) + + +@pytest.mark.parametrize("names", ALIAS_COMBINATIONS) +def test_duplicate_aliases(names): + """Adding any name that maps to a coordinate already given is an error.""" + coordinates = values(names) + + for name in names: + for alias in ALIAS_CHOICES[ALIASES.get(name, name)]: + if alias == name: + continue + duplicated = {**coordinates, alias: 99.0} + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.obj(**duplicated) + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.array( + {k: np.array([v]) for k, v in duplicated.items()}, + ) + + +def test_duplicate_aliases_awkward(): + ak = pytest.importorskip("awkward") + + for names in ALIAS_COMBINATIONS: + coordinates = values(names) + for name in names: + for alias in ALIAS_CHOICES[ALIASES.get(name, name)]: + if alias == name: + continue + duplicated = {**coordinates, alias: 99.0} + columns = {k: [v] for k, v in duplicated.items()} + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.Array(ak.Array(columns)) + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.zip(columns) + + +def test_generic_vectors_reject_momentum_aliases(): + with pytest.raises(TypeError, match="momentum-aliases are not allowed"): + vector.VectorObject2D(px=1.0, py=2.0) + with pytest.raises(TypeError, match="momentum-aliases are not allowed"): + vector.VectorObject3D(x=1.0, y=2.0, pz=3.0) + with pytest.raises(TypeError, match="momentum-aliases are not allowed"): + vector.VectorObject4D(x=1.0, y=2.0, z=3.0, energy=4.0) + with pytest.raises(TypeError, match="momentum-aliases are not allowed"): + numpy_array({"px": 1.0, "py": 2.0}).view(vector.VectorNumpy2D) + + +def test_momentum_vectors_accept_generic_names(): + assert vector.MomentumObject2D(x=1.0, y=2.0).px == pytest.approx(1.0) + assert vector.MomentumObject3D(x=1.0, y=2.0, z=3.0).pz == pytest.approx(3.0) + momentum = vector.MomentumObject4D(x=1.0, y=2.0, z=3.0, t=4.0) + assert momentum.energy == pytest.approx(4.0) + + +def test_conflicting_coordinates(): + with pytest.raises(TypeError, match="specify x= and y= or rho= and phi="): + vector.obj(x=1.0, y=2.0, rho=3.0, phi=4.0) + with pytest.raises(TypeError, match="specify z= or theta= or eta="): + vector.obj(x=1.0, y=2.0, z=3.0, eta=4.0) + with pytest.raises(TypeError, match="specify t= or tau="): + vector.obj(x=1.0, y=2.0, z=3.0, t=4.0, tau=5.0) + with pytest.raises(TypeError, match="specify t= or tau="): + vector.obj(pt=1.0, phi=2.0, eta=3.0, mass=4.0, energy=5.0) + + +def test_extra_fields(): + """Non-coordinate fields are records' payload; keyword arguments are not.""" + array = vector.array( + {"x": np.array([1.0]), "y": np.array([2.0]), "wow": np.array([3.0])} + ) + assert array["wow"][0] == pytest.approx(3.0) + + with pytest.raises(TypeError, match="unrecognized combination"): + vector.obj(x=1.0, y=2.0, wow=3.0) + with pytest.raises(TypeError, match="unrecognized combination"): + vector.VectorObject2D(x=1.0, y=2.0, wow=3.0) + + +def test_extra_fields_awkward(): + ak = pytest.importorskip("awkward") + + for vec in ( + vector.Array(ak.Array({"x": [1.0], "y": [2.0], "wow": [3.0]})), + vector.zip({"x": [1.0], "y": [2.0], "wow": [3.0]}), + ): + assert vec.wow[0] == pytest.approx(3.0) + + +def test_unstructured_numpy_array(): + with pytest.raises(TypeError, match="must have a structured dtype"): + np.array([1.0, 2.0]).view(vector.VectorNumpy2D) + with pytest.raises(TypeError, match="unrecognized combination"): + vector.array([1.0, 2.0]) + + +def test_coordinate_objects_are_still_required(): + with pytest.raises(TypeError, match="must give Azimuthal"): + vector.VectorObject2D() + with pytest.raises(TypeError, match="must give Azimuthal and Longitudinal"): + vector.VectorObject3D() + with pytest.raises( + TypeError, match="must give Azimuthal, Longitudinal, and Temporal" + ): + vector.VectorObject4D() + + +def test_complaint_lists_the_allowed_combinations(): + with pytest.raises(TypeError) as excinfo: + vector.obj(x=1.0) + complaint = str(excinfo.value) + for names in VECTOR_COMBINATIONS: + assert ( + " ({}D) {}".format(len(names), " ".join(f"{x}=" for x in names)) + in complaint + ) + assert "or their momentum equivalents" in complaint + + # a constructor that is restricted to one dimension only lists that dimension + with pytest.raises(TypeError) as excinfo: + vector.VectorObject2D(x=1.0, y=2.0, z=3.0) + complaint = str(excinfo.value) + assert " x= y=" in complaint + assert " rho= phi=" in complaint + assert "z=" not in complaint + assert "or their momentum equivalents" not in complaint + + +def awkward_validates() -> bool: + """Whether the installed Awkward Array calls ``__awkward_validation__``.""" + ak = pytest.importorskip("awkward") + + validated = [] + + class Probe(ak.Array): # type: ignore[misc] + def __awkward_validation__(self) -> None: + validated.append(None) + + ak.Array( + [{"x": 1.1}], + behavior={("*", "probe"): Probe}, + with_name="probe", + ) + return bool(validated) + + +@pytest.mark.parametrize("names", ALIAS_COMBINATIONS) +def test_awkward_behavior_validation(names): + ak = pytest.importorskip("awkward") + if not awkward_validates(): + pytest.skip("awkward is too old to validate behaviors") + + behavior = vector.backends.awkward.behavior + coordinates = values(names) + columns = {name: [value] for name, value in coordinates.items()} + + vec = ak.zip(columns, with_name=record_name(names), behavior=behavior) + for name, value in coordinates.items(): + assert getattr(vec, name)[0] == pytest.approx(value) + + # records are validated in the same way as arrays + assert getattr(vec[0], names[0]) == pytest.approx(coordinates[names[0]]) + + # a record name that does not describe these coordinates is rejected + for dimension in (2, 3, 4): + for momentum in (False, True): + name = f"{'Momentum' if momentum else 'Vector'}{dimension}D" + if is_valid_for(names, dimension, momentum): + ak.zip(columns, with_name=name, behavior=behavior) + else: + with pytest.raises(TypeError): + ak.zip(columns, with_name=name, behavior=behavior) + + +def test_awkward_behavior_validation_extra_fields(): + ak = pytest.importorskip("awkward") + if not awkward_validates(): + pytest.skip("awkward is too old to validate behaviors") + + behavior = vector.backends.awkward.behavior + columns = {"pt": [1.0], "phi": [2.0], "eta": [3.0], "mass": [4.0], "charge": [1]} + vec = ak.zip(columns, with_name="Momentum4D", behavior=behavior) + assert vec.pt[0] == pytest.approx(1.0) + assert vec.charge[0] == 1 + + # ... but a field that is a coordinate under another name is not payload + with pytest.raises(TypeError, match="specify t= or tau="): + ak.zip({**columns, "energy": [5.0]}, with_name="Momentum4D", behavior=behavior) + with pytest.raises(TypeError, match="duplicate coordinates"): + ak.zip({**columns, "rho": [5.0]}, with_name="Momentum4D", behavior=behavior) + + # assigning such a field to an existing array is caught as well + with pytest.raises(TypeError, match="duplicate coordinates"): + vec["rho"] = np.array([5.0]) + + +def test_awkward_behavior_validation_names_the_array(): + """The record name is chosen elsewhere, so the complaint has to identify it.""" + ak = pytest.importorskip("awkward") + if not awkward_validates(): + pytest.skip("awkward is too old to validate behaviors") + + behavior = vector.backends.awkward.behavior + columns = {"pt": [1.0], "phi": [2.0]} + + with pytest.raises(TypeError, match=r"MomentumArray4D with fields \['pt', 'phi'\]"): + ak.zip(columns, with_name="Momentum4D", behavior=behavior) + + vec = ak.zip(columns, with_name="Momentum2D", behavior=behavior) + with pytest.raises(TypeError, match=r"MomentumArray2D with fields"): + del vec["phi"] + + +def test_awkward_behavior_validation_subclass(): + ak = pytest.importorskip("awkward") + if not awkward_validates(): + pytest.skip("awkward is too old to validate behaviors") + + behavior = dict(vector.backends.awkward.behavior) + + class PtEtaPhiMArray(vector.backends.awkward.MomentumArray4D): # type: ignore[misc] + pass + + behavior["*", "PtEtaPhiM"] = PtEtaPhiMArray + + vec = ak.zip( + {"pt": [1.0], "phi": [2.0], "eta": [3.0], "mass": [4.0]}, + with_name="PtEtaPhiM", + behavior=behavior, + ) + assert vec.pt[0] == pytest.approx(1.0) + + with pytest.raises(TypeError, match="specify t= or tau="): + ak.zip( + {"pt": [1.0], "phi": [2.0], "eta": [3.0], "mass": [4.0], "energy": [5.0]}, + with_name="PtEtaPhiM", + behavior=behavior, + ) + + +# Every combination of generic names, plus enough momentum-aliases to reach every +# entry of the tables that the Numba implementation of ``vector.obj`` dispatches on. +NUMBA_COMBINATIONS = [ + *VECTOR_COMBINATIONS, + ("px", "py"), + ("x", "py"), + ("px", "y"), + ("pt", "phi"), + ("px", "py", "pz"), + ("px", "py", "pz", "E"), + ("px", "py", "pz", "e"), + ("px", "py", "pz", "energy"), + ("px", "py", "pz", "M"), + ("px", "py", "pz", "m"), + ("px", "py", "pz", "mass"), +] + + +@pytest.mark.numba +def test_numba_obj_combinations(): + numba = pytest.importorskip("numba") + pytest.importorskip("vector.backends._numba_object") + + # every call site is typed separately, so they are compiled together to + # keep the (considerable) compilation time down + source = "def make_all():\n return (\n" + for names in NUMBA_COMBINATIONS: + arguments = ", ".join( + f"{name}={value}" for name, value in values(names).items() + ) + source += f" vector.obj({arguments}),\n" + source += " )\n" + + namespace = {"vector": vector} + exec(source, namespace) + + vectors = numba.njit(namespace["make_all"])() + + for names, vec in zip(NUMBA_COMBINATIONS, vectors, strict=True): + is_momentum, dimension = reference(names) + assert isinstance(vec, object_classes[dimension, is_momentum]) + for name, value in values(names).items(): + assert getattr(vec, name) == pytest.approx(value) + + +@pytest.mark.numba +def test_numba_obj_invalid_combinations(): + numba = pytest.importorskip("numba") + pytest.importorskip("vector.backends._numba_object") + + @numba.njit + def duplicate(): + return vector.obj(x=1.0, px=2.0, y=3.0) + + @numba.njit + def two_temporal(): + return vector.obj(pt=1.0, phi=2.0, eta=3.0, mass=4.0, energy=5.0) + + @numba.njit + def two_longitudinal(): + return vector.obj(x=1.0, y=2.0, theta=3.0, eta=4.0) + + @numba.njit + def two_azimuthal(): + return vector.obj(x=1.0, y=2.0, rho=3.0, phi=4.0) + + for function, complaint in ( + (duplicate, "duplicate coordinates"), + (two_temporal, "specify t= or tau="), + (two_longitudinal, "specify z= or theta= or eta="), + (two_azimuthal, "specify x= and y= or rho= and phi="), + ): + with pytest.raises(numba.TypingError, match=complaint): + function()