From 0f2c1c0980c0de74f12ecda06ac7f4ee8b25ee71 Mon Sep 17 00:00:00 2001 From: Iason Krommydas Date: Thu, 30 Oct 2025 17:07:53 -0500 Subject: [PATCH 01/15] first try --- src/vector/backends/awkward_constructors.py | 20 +++ src/vector/backends/numpy.py | 165 ++++++++++++++++++++ 2 files changed, 185 insertions(+) diff --git a/src/vector/backends/awkward_constructors.py b/src/vector/backends/awkward_constructors.py index 625b1fae..14c53aa8 100644 --- a/src/vector/backends/awkward_constructors.py +++ b/src/vector/backends/awkward_constructors.py @@ -199,6 +199,26 @@ def _check_names( if dimension == 0: raise TypeError(complaint1 if is_momentum else complaint2) + # Check if any remaining fieldnames would conflict with already-processed coordinates + # when mapped to generic names (e.g., pt was processed, rho shouldn't remain) + if fieldnames: + from vector._methods import _repr_momentum_to_generic + + # Reconstruct original fieldnames from names already processed + # to check against remaining fieldnames + original_fieldnames = list(names) # Start with processed generic names + original_fieldnames.extend(fieldnames) # Add remaining fieldnames + + # Map to generic names - but we need the original input fieldnames + # Actually, we need to check if remaining fieldnames conflict with processed ones + # The processed ones are in 'names' (generic form), remaining are in 'fieldnames' + + # Check each remaining fieldname to see if its generic form was already used + for fname in fieldnames: + generic = _repr_momentum_to_generic.get(fname, fname) + if generic in names: + raise TypeError(complaint1 if is_momentum else complaint2) + for name in fieldnames: names.append(name) columns.append(projectable[name]) diff --git a/src/vector/backends/numpy.py b/src/vector/backends/numpy.py index d05588b9..6f91300d 100644 --- a/src/vector/backends/numpy.py +++ b/src/vector/backends/numpy.py @@ -2071,6 +2071,168 @@ def __setitem__(self, where: typing.Any, what: typing.Any) -> None: return _setitem(self, where, what, True) +def _validate_numpy_coordinates(fieldnames: tuple[str, ...]) -> None: + """ + Validate coordinate field names using dimension-guard pattern. + + This follows the same logic as _check_names in awkward_constructors to ensure + consistent validation across backends. + + Raises TypeError if duplicate or conflicting coordinates are detected. + """ + 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=" + ) + + is_momentum = False + dimension = 0 + fieldnames_copy = list(fieldnames) + + # 2D azimuthal coordinates + if "x" in fieldnames_copy and "y" in fieldnames_copy: + if dimension != 0: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 2 + fieldnames_copy.remove("x") + fieldnames_copy.remove("y") + if "rho" in fieldnames_copy and "phi" in fieldnames_copy: + if dimension != 0: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 2 + fieldnames_copy.remove("rho") + fieldnames_copy.remove("phi") + if "x" in fieldnames_copy and "py" in fieldnames_copy: + is_momentum = True + if dimension != 0: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 2 + fieldnames_copy.remove("x") + fieldnames_copy.remove("py") + if "px" in fieldnames_copy and "y" in fieldnames_copy: + is_momentum = True + if dimension != 0: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 2 + fieldnames_copy.remove("px") + fieldnames_copy.remove("y") + if "px" in fieldnames_copy and "py" in fieldnames_copy: + is_momentum = True + if dimension != 0: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 2 + fieldnames_copy.remove("px") + fieldnames_copy.remove("py") + if "pt" in fieldnames_copy and "phi" in fieldnames_copy: + is_momentum = True + if dimension != 0: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 2 + fieldnames_copy.remove("pt") + fieldnames_copy.remove("phi") + + # 3D longitudinal coordinates + if "z" in fieldnames_copy: + if dimension != 2: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 3 + fieldnames_copy.remove("z") + if "theta" in fieldnames_copy: + if dimension != 2: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 3 + fieldnames_copy.remove("theta") + if "eta" in fieldnames_copy: + if dimension != 2: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 3 + fieldnames_copy.remove("eta") + if "pz" in fieldnames_copy: + is_momentum = True + if dimension != 2: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 3 + fieldnames_copy.remove("pz") + + # 4D temporal coordinates + if "t" in fieldnames_copy: + if dimension != 3: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 4 + fieldnames_copy.remove("t") + if "tau" in fieldnames_copy: + if dimension != 3: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 4 + fieldnames_copy.remove("tau") + if "E" in fieldnames_copy: + is_momentum = True + if dimension != 3: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 4 + fieldnames_copy.remove("E") + if "e" in fieldnames_copy: + is_momentum = True + if dimension != 3: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 4 + fieldnames_copy.remove("e") + if "energy" in fieldnames_copy: + is_momentum = True + if dimension != 3: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 4 + fieldnames_copy.remove("energy") + if "M" in fieldnames_copy: + is_momentum = True + if dimension != 3: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 4 + fieldnames_copy.remove("M") + if "m" in fieldnames_copy: + is_momentum = True + if dimension != 3: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 4 + fieldnames_copy.remove("m") + if "mass" in fieldnames_copy: + is_momentum = True + if dimension != 3: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 4 + fieldnames_copy.remove("mass") + + # Check if any remaining fieldnames would conflict with already-processed coordinates + # when mapped to generic names (e.g., pt was processed, rho shouldn't remain) + if fieldnames_copy: + # Map all original fieldnames to generic names to detect conflicts + generic_names = [_repr_momentum_to_generic.get(x, x) for x in fieldnames] + if len(generic_names) != len(set(generic_names)): + raise TypeError(complaint1 if is_momentum else complaint2) + + def array(*args: typing.Any, **kwargs: typing.Any) -> VectorNumpy: """ Constructs a NumPy array of vectors, whose type is determined by the dtype @@ -2138,6 +2300,9 @@ def array(*args: typing.Any, **kwargs: typing.Any) -> VectorNumpy: is_momentum = any(x in _repr_momentum_to_generic for x in names) + # Validate coordinates using dimension-guard pattern (same as awkward _check_names) + _validate_numpy_coordinates(names) + if any(x in ("t", "E", "e", "energy", "tau", "M", "m", "mass") for x in names): cls = MomentumNumpy4D if is_momentum else VectorNumpy4D elif any(x in ("z", "pz", "theta", "eta") for x in names): From b04d5cfde0c1bb2ff3beb76fc8593c02913901fd Mon Sep 17 00:00:00 2001 From: Iason Krommydas Date: Thu, 30 Oct 2025 17:27:16 -0500 Subject: [PATCH 02/15] import at the top --- src/vector/backends/awkward_constructors.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/vector/backends/awkward_constructors.py b/src/vector/backends/awkward_constructors.py index 14c53aa8..95fadea2 100644 --- a/src/vector/backends/awkward_constructors.py +++ b/src/vector/backends/awkward_constructors.py @@ -9,6 +9,8 @@ import numpy +from vector._methods import _repr_momentum_to_generic + def _recname(is_momentum: bool, dimension: int) -> str: name = "Momentum" if is_momentum else "Vector" @@ -202,8 +204,6 @@ def _check_names( # Check if any remaining fieldnames would conflict with already-processed coordinates # when mapped to generic names (e.g., pt was processed, rho shouldn't remain) if fieldnames: - from vector._methods import _repr_momentum_to_generic - # Reconstruct original fieldnames from names already processed # to check against remaining fieldnames original_fieldnames = list(names) # Start with processed generic names From 96d9c7ce80453a94050d75d9c4036ae0da0e7d66 Mon Sep 17 00:00:00 2001 From: Iason Krommydas Date: Thu, 30 Oct 2025 19:04:03 -0500 Subject: [PATCH 03/15] check leftovers --- src/vector/backends/awkward_constructors.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/src/vector/backends/awkward_constructors.py b/src/vector/backends/awkward_constructors.py index 95fadea2..3c2409d8 100644 --- a/src/vector/backends/awkward_constructors.py +++ b/src/vector/backends/awkward_constructors.py @@ -202,23 +202,19 @@ def _check_names( raise TypeError(complaint1 if is_momentum else complaint2) # Check if any remaining fieldnames would conflict with already-processed coordinates - # when mapped to generic names (e.g., pt was processed, rho shouldn't remain) + # or with each other when mapped to generic names (e.g., "x" and "px" both map to "x") if fieldnames: - # Reconstruct original fieldnames from names already processed - # to check against remaining fieldnames - original_fieldnames = list(names) # Start with processed generic names - original_fieldnames.extend(fieldnames) # Add remaining fieldnames - - # Map to generic names - but we need the original input fieldnames - # Actually, we need to check if remaining fieldnames conflict with processed ones - # The processed ones are in 'names' (generic form), remaining are in 'fieldnames' - - # Check each remaining fieldname to see if its generic form was already used + # Check leftovers against already-processed coordinates for fname in fieldnames: generic = _repr_momentum_to_generic.get(fname, fname) if generic in names: raise TypeError(complaint1 if is_momentum else complaint2) + # Check leftovers against each other for duplicates + leftover_generics = [_repr_momentum_to_generic.get(x, x) for x in fieldnames] + if len(leftover_generics) != len(set(leftover_generics)): + raise TypeError(complaint1 if is_momentum else complaint2) + for name in fieldnames: names.append(name) columns.append(projectable[name]) From a500b6fca19becf48cfff0ceff2eee1def67b3ac Mon Sep 17 00:00:00 2001 From: Iason Krommydas Date: Thu, 30 Oct 2025 20:35:15 -0500 Subject: [PATCH 04/15] add test --- src/vector/backends/object.py | 4 +- tests/backends/test_coordinate_validation.py | 1414 ++++++++++++++++++ 2 files changed, 1416 insertions(+), 2 deletions(-) create mode 100644 tests/backends/test_coordinate_validation.py diff --git a/src/vector/backends/object.py b/src/vector/backends/object.py index 16cb8101..d1bde4b7 100644 --- a/src/vector/backends/object.py +++ b/src/vector/backends/object.py @@ -3211,7 +3211,7 @@ def obj(**coordinates: float) -> VectorObject: if "E" in coordinates: is_momentum = True generic_coordinates["t"] = coordinates.pop("E") - if "e" in coordinates: + if "e" in coordinates and "t" not in generic_coordinates: is_momentum = True generic_coordinates["t"] = coordinates.pop("e") if "energy" in coordinates and "t" not in generic_coordinates: @@ -3220,7 +3220,7 @@ def obj(**coordinates: float) -> VectorObject: if "M" in coordinates: is_momentum = True generic_coordinates["tau"] = coordinates.pop("M") - if "m" in coordinates: + if "m" in coordinates and "tau" not in generic_coordinates: is_momentum = True generic_coordinates["tau"] = coordinates.pop("m") if "mass" in coordinates and "tau" not in generic_coordinates: diff --git a/tests/backends/test_coordinate_validation.py b/tests/backends/test_coordinate_validation.py new file mode 100644 index 00000000..3f524b7e --- /dev/null +++ b/tests/backends/test_coordinate_validation.py @@ -0,0 +1,1414 @@ +# Copyright (c) 2019-2025, 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. + +from __future__ import annotations + +import numpy as np +import pytest + +import vector + +ak = pytest.importorskip("awkward") + +pytestmark = pytest.mark.awkward + + +# ============================================================================ +# Duplicate temporal coordinates (t-like vs tau-like) +# ============================================================================ +# Temporal coordinates: t, E, e, energy (all map to 't') +# tau, M, m, mass (all map to 'tau') +# These are mutually exclusive + + +def test_duplicate_E_e_object(): + """vector.obj should reject E + e""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.obj(pt=1.0, phi=0.5, eta=1.0, E=5.0, e=5.0) + + +def test_duplicate_E_e_numpy(): + """vector.array should reject E + e""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.array( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "E": np.array([5.0, 6.0]), + "e": np.array([5.0, 6.0]), + } + ) + + +def test_duplicate_E_e_awkward(): + """vector.Array should reject E + e""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.Array( + ak.Array( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "E": np.array([5.0, 6.0]), + "e": np.array([5.0, 6.0]), + } + ) + ) + + +def test_duplicate_E_e_zip(): + """vector.zip should reject E + e""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.zip( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "E": np.array([5.0, 6.0]), + "e": np.array([5.0, 6.0]), + } + ) + + +def test_duplicate_E_energy_object(): + """vector.obj should reject E + energy""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.obj(pt=1.0, phi=0.5, eta=1.0, E=5.0, energy=5.0) + + +def test_duplicate_E_energy_numpy(): + """vector.array should reject E + energy""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.array( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "E": np.array([5.0, 6.0]), + "energy": np.array([5.0, 6.0]), + } + ) + + +def test_duplicate_E_energy_awkward(): + """vector.Array should reject E + energy""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.Array( + ak.Array( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "E": np.array([5.0, 6.0]), + "energy": np.array([5.0, 6.0]), + } + ) + ) + + +def test_duplicate_E_energy_zip(): + """vector.zip should reject E + energy""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.zip( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "E": np.array([5.0, 6.0]), + "energy": np.array([5.0, 6.0]), + } + ) + + +def test_duplicate_e_energy_object(): + """vector.obj should reject e + energy""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.obj(pt=1.0, phi=0.5, eta=1.0, e=5.0, energy=5.0) + + +def test_duplicate_e_energy_numpy(): + """vector.array should reject e + energy""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.array( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "e": np.array([5.0, 6.0]), + "energy": np.array([5.0, 6.0]), + } + ) + + +def test_duplicate_e_energy_awkward(): + """vector.Array should reject e + energy""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.Array( + ak.Array( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "e": np.array([5.0, 6.0]), + "energy": np.array([5.0, 6.0]), + } + ) + ) + + +def test_duplicate_e_energy_zip(): + """vector.zip should reject e + energy""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.zip( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "e": np.array([5.0, 6.0]), + "energy": np.array([5.0, 6.0]), + } + ) + + +def test_duplicate_M_m_object(): + """vector.obj should reject M + m""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.obj(pt=1.0, phi=0.5, eta=1.0, M=0.5, m=0.5) + + +def test_duplicate_M_m_numpy(): + """vector.array should reject M + m""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.array( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "M": np.array([0.5, 0.5]), + "m": np.array([0.5, 0.5]), + } + ) + + +def test_duplicate_M_m_awkward(): + """vector.Array should reject M + m""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.Array( + ak.Array( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "M": np.array([0.5, 0.5]), + "m": np.array([0.5, 0.5]), + } + ) + ) + + +def test_duplicate_M_m_zip(): + """vector.zip should reject M + m""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.zip( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "M": np.array([0.5, 0.5]), + "m": np.array([0.5, 0.5]), + } + ) + + +def test_duplicate_M_mass_object(): + """vector.obj should reject M + mass""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.obj(pt=1.0, phi=0.5, eta=1.0, M=0.5, mass=0.5) + + +def test_duplicate_M_mass_numpy(): + """vector.array should reject M + mass""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.array( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "M": np.array([0.5, 0.5]), + "mass": np.array([0.5, 0.5]), + } + ) + + +def test_duplicate_M_mass_awkward(): + """vector.Array should reject M + mass""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.Array( + ak.Array( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "M": np.array([0.5, 0.5]), + "mass": np.array([0.5, 0.5]), + } + ) + ) + + +def test_duplicate_M_mass_zip(): + """vector.zip should reject M + mass""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.zip( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "M": np.array([0.5, 0.5]), + "mass": np.array([0.5, 0.5]), + } + ) + + +def test_duplicate_m_mass_object(): + """vector.obj should reject m + mass""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.obj(pt=1.0, phi=0.5, eta=1.0, m=0.5, mass=0.5) + + +def test_duplicate_m_mass_numpy(): + """vector.array should reject m + mass""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.array( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "m": np.array([0.5, 0.5]), + "mass": np.array([0.5, 0.5]), + } + ) + + +def test_duplicate_m_mass_awkward(): + """vector.Array should reject m + mass""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.Array( + ak.Array( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "m": np.array([0.5, 0.5]), + "mass": np.array([0.5, 0.5]), + } + ) + ) + + +def test_duplicate_m_mass_zip(): + """vector.zip should reject m + mass""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.zip( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "m": np.array([0.5, 0.5]), + "mass": np.array([0.5, 0.5]), + } + ) + + +def test_conflicting_energy_mass_object(): + """vector.obj should reject energy + mass (t-like + tau-like)""" + with pytest.raises(TypeError, match="specify t= or tau=, but not more than one"): + vector.obj(pt=1.0, phi=0.5, eta=1.0, energy=5.0, mass=0.5) + + +def test_conflicting_energy_mass_numpy(): + """vector.array should reject energy + mass (t-like + tau-like)""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.array( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "energy": np.array([5.0, 6.0]), + "mass": np.array([0.5, 0.5]), + } + ) + + +def test_conflicting_energy_mass_awkward(): + """vector.Array should reject energy + mass (t-like + tau-like)""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.Array( + ak.Array( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "energy": np.array([5.0, 6.0]), + "mass": np.array([0.5, 0.5]), + } + ) + ) + + +def test_conflicting_energy_mass_zip(): + """vector.zip should reject energy + mass (t-like + tau-like)""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.zip( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "energy": np.array([5.0, 6.0]), + "mass": np.array([0.5, 0.5]), + } + ) + + +def test_conflicting_t_tau_object(): + """vector.obj should reject t + tau""" + with pytest.raises(TypeError, match="specify t= or tau="): + vector.obj(x=1.0, y=2.0, z=3.0, t=5.0, tau=0.5) + + +def test_conflicting_t_tau_numpy(): + """vector.array should reject t + tau""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.array( + { + "x": np.array([1.0, 2.0]), + "y": np.array([2.0, 3.0]), + "z": np.array([3.0, 4.0]), + "t": np.array([5.0, 6.0]), + "tau": np.array([0.5, 0.5]), + } + ) + + +def test_conflicting_t_tau_awkward(): + """vector.Array should reject t + tau""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.Array( + ak.Array( + { + "x": np.array([1.0, 2.0]), + "y": np.array([2.0, 3.0]), + "z": np.array([3.0, 4.0]), + "t": np.array([5.0, 6.0]), + "tau": np.array([0.5, 0.5]), + } + ) + ) + + +def test_conflicting_t_tau_zip(): + """vector.zip should reject t + tau""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.zip( + { + "x": np.array([1.0, 2.0]), + "y": np.array([2.0, 3.0]), + "z": np.array([3.0, 4.0]), + "t": np.array([5.0, 6.0]), + "tau": np.array([0.5, 0.5]), + } + ) + + +def test_conflicting_E_mass_object(): + """vector.obj should reject E + mass""" + with pytest.raises(TypeError, match="specify t= or tau=, but not more than one"): + vector.obj(pt=1.0, phi=0.5, eta=1.0, E=5.0, mass=0.5) + + +def test_conflicting_E_mass_numpy(): + """vector.array should reject E + mass""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.array( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "E": np.array([5.0, 6.0]), + "mass": np.array([0.5, 0.5]), + } + ) + + +def test_conflicting_E_mass_awkward(): + """vector.Array should reject E + mass""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.Array( + ak.Array( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "E": np.array([5.0, 6.0]), + "mass": np.array([0.5, 0.5]), + } + ) + ) + + +def test_conflicting_E_mass_zip(): + """vector.zip should reject E + mass""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.zip( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "E": np.array([5.0, 6.0]), + "mass": np.array([0.5, 0.5]), + } + ) + + +def test_conflicting_t_mass_object(): + """vector.obj should reject t + mass""" + with pytest.raises(TypeError, match="specify t= or tau=, but not more than one"): + vector.obj(x=1.0, y=2.0, z=3.0, t=5.0, mass=0.5) + + +def test_conflicting_t_mass_numpy(): + """vector.array should reject t + mass""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.array( + { + "x": np.array([1.0, 2.0]), + "y": np.array([2.0, 3.0]), + "z": np.array([3.0, 4.0]), + "t": np.array([5.0, 6.0]), + "mass": np.array([0.5, 0.5]), + } + ) + + +def test_conflicting_t_mass_awkward(): + """vector.Array should reject t + mass""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.Array( + ak.Array( + { + "x": np.array([1.0, 2.0]), + "y": np.array([2.0, 3.0]), + "z": np.array([3.0, 4.0]), + "t": np.array([5.0, 6.0]), + "mass": np.array([0.5, 0.5]), + } + ) + ) + + +def test_conflicting_t_mass_zip(): + """vector.zip should reject t + mass""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.zip( + { + "x": np.array([1.0, 2.0]), + "y": np.array([2.0, 3.0]), + "z": np.array([3.0, 4.0]), + "t": np.array([5.0, 6.0]), + "mass": np.array([0.5, 0.5]), + } + ) + + +def test_conflicting_energy_tau_object(): + """vector.obj should reject energy + tau""" + with pytest.raises(TypeError, match="specify t= or tau=, but not more than one"): + vector.obj(x=1.0, y=2.0, z=3.0, energy=5.0, tau=0.5) + + +def test_conflicting_energy_tau_numpy(): + """vector.array should reject energy + tau""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.array( + { + "x": np.array([1.0, 2.0]), + "y": np.array([2.0, 3.0]), + "z": np.array([3.0, 4.0]), + "energy": np.array([5.0, 6.0]), + "tau": np.array([0.5, 0.5]), + } + ) + + +def test_conflicting_energy_tau_awkward(): + """vector.Array should reject energy + tau""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.Array( + ak.Array( + { + "x": np.array([1.0, 2.0]), + "y": np.array([2.0, 3.0]), + "z": np.array([3.0, 4.0]), + "energy": np.array([5.0, 6.0]), + "tau": np.array([0.5, 0.5]), + } + ) + ) + + +def test_conflicting_energy_tau_zip(): + """vector.zip should reject energy + tau""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.zip( + { + "x": np.array([1.0, 2.0]), + "y": np.array([2.0, 3.0]), + "z": np.array([3.0, 4.0]), + "energy": np.array([5.0, 6.0]), + "tau": np.array([0.5, 0.5]), + } + ) + + +# ============================================================================ +# Duplicate azimuthal coordinates +# ============================================================================ +# x <-> px, y <-> py, rho <-> pt + + +def test_duplicate_px_x_object(): + """vector.obj should reject px + x""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.obj(px=1.0, x=1.0, y=2.0, z=3.0, t=5.0) + + +def test_duplicate_px_x_numpy(): + """vector.array should reject px + x""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.array( + { + "px": np.array([1.0, 2.0]), + "x": np.array([1.0, 2.0]), + "y": np.array([2.0, 3.0]), + "z": np.array([3.0, 4.0]), + "t": np.array([5.0, 6.0]), + } + ) + + +def test_duplicate_px_x_awkward(): + """vector.Array should reject px + x""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.Array( + ak.Array( + { + "px": np.array([1.0, 2.0]), + "x": np.array([1.0, 2.0]), + "y": np.array([2.0, 3.0]), + "z": np.array([3.0, 4.0]), + "t": np.array([5.0, 6.0]), + } + ) + ) + + +def test_duplicate_px_x_zip(): + """vector.zip should reject px + x""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.zip( + { + "px": np.array([1.0, 2.0]), + "x": np.array([1.0, 2.0]), + "y": np.array([2.0, 3.0]), + "z": np.array([3.0, 4.0]), + "t": np.array([5.0, 6.0]), + } + ) + + +def test_duplicate_py_y_object(): + """vector.obj should reject py + y""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.obj(x=1.0, py=2.0, y=2.0, z=3.0, t=5.0) + + +def test_duplicate_py_y_numpy(): + """vector.array should reject py + y""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.array( + { + "x": np.array([1.0, 2.0]), + "py": np.array([2.0, 3.0]), + "y": np.array([2.0, 3.0]), + "z": np.array([3.0, 4.0]), + "t": np.array([5.0, 6.0]), + } + ) + + +def test_duplicate_py_y_awkward(): + """vector.Array should reject py + y""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.Array( + ak.Array( + { + "x": np.array([1.0, 2.0]), + "py": np.array([2.0, 3.0]), + "y": np.array([2.0, 3.0]), + "z": np.array([3.0, 4.0]), + "t": np.array([5.0, 6.0]), + } + ) + ) + + +def test_duplicate_py_y_zip(): + """vector.zip should reject py + y""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.zip( + { + "x": np.array([1.0, 2.0]), + "py": np.array([2.0, 3.0]), + "y": np.array([2.0, 3.0]), + "z": np.array([3.0, 4.0]), + "t": np.array([5.0, 6.0]), + } + ) + + +def test_duplicate_pt_rho_object(): + """vector.obj should reject pt + rho""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.obj(pt=1.0, rho=1.0, phi=0.5, eta=1.0, mass=0.5) + + +def test_duplicate_pt_rho_numpy(): + """vector.array should reject pt + rho""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.array( + { + "pt": np.array([1.0, 2.0]), + "rho": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "mass": np.array([0.5, 0.5]), + } + ) + + +def test_duplicate_pt_rho_awkward(): + """vector.Array should reject pt + rho""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.Array( + ak.Array( + { + "pt": np.array([1.0, 2.0]), + "rho": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "mass": np.array([0.5, 0.5]), + } + ) + ) + + +def test_duplicate_pt_rho_zip(): + """vector.zip should reject pt + rho""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.zip( + { + "pt": np.array([1.0, 2.0]), + "rho": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "mass": np.array([0.5, 0.5]), + } + ) + + +# ============================================================================ +# Duplicate longitudinal coordinates +# ============================================================================ +# z <-> pz + + +def test_duplicate_pz_z_object(): + """vector.obj should reject pz + z""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.obj(x=1.0, y=2.0, pz=3.0, z=3.0, t=5.0) + + +def test_duplicate_pz_z_numpy(): + """vector.array should reject pz + z""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.array( + { + "x": np.array([1.0, 2.0]), + "y": np.array([2.0, 3.0]), + "pz": np.array([3.0, 4.0]), + "z": np.array([3.0, 4.0]), + "t": np.array([5.0, 6.0]), + } + ) + + +def test_duplicate_pz_z_awkward(): + """vector.Array should reject pz + z""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.Array( + ak.Array( + { + "x": np.array([1.0, 2.0]), + "y": np.array([2.0, 3.0]), + "pz": np.array([3.0, 4.0]), + "z": np.array([3.0, 4.0]), + "t": np.array([5.0, 6.0]), + } + ) + ) + + +def test_duplicate_pz_z_zip(): + """vector.zip should reject pz + z""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.zip( + { + "x": np.array([1.0, 2.0]), + "y": np.array([2.0, 3.0]), + "pz": np.array([3.0, 4.0]), + "z": np.array([3.0, 4.0]), + "t": np.array([5.0, 6.0]), + } + ) + + +# ============================================================================ +# Mixed azimuthal coordinate systems (from _gather_coordinates) +# ============================================================================ + + +def test_mixed_xy_with_rho_object(): + """vector.obj should reject x+y with rho""" + with pytest.raises(TypeError, match="specify x= and y= or rho= and phi="): + vector.obj(x=1.0, y=2.0, rho=1.0, z=3.0) + + +def test_mixed_xy_with_phi_object(): + """vector.obj should reject x+y with phi""" + with pytest.raises(TypeError, match="specify x= and y= or rho= and phi="): + vector.obj(x=1.0, y=2.0, phi=0.5, z=3.0) + + +def test_mixed_rhophi_with_x_object(): + """vector.obj should reject rho+phi with x""" + with pytest.raises(TypeError, match="specify x= and y= or rho= and phi="): + vector.obj(rho=1.0, phi=0.5, x=1.0, z=3.0) + + +def test_mixed_rhophi_with_y_object(): + """vector.obj should reject rho+phi with y""" + with pytest.raises(TypeError, match="specify x= and y= or rho= and phi="): + vector.obj(rho=1.0, phi=0.5, y=2.0, z=3.0) + + +# ============================================================================ +# Mixed longitudinal coordinates (from _gather_coordinates) +# ============================================================================ + + +def test_mixed_z_theta_object(): + """vector.obj should reject z with theta""" + with pytest.raises(TypeError, match="specify z= or theta= or eta="): + vector.obj(x=1.0, y=2.0, z=3.0, theta=1.0) + + +def test_mixed_z_eta_object(): + """vector.obj should reject z with eta""" + with pytest.raises(TypeError, match="specify z= or theta= or eta="): + vector.obj(x=1.0, y=2.0, z=3.0, eta=1.0) + + +def test_mixed_theta_eta_object(): + """vector.obj should reject theta with eta""" + with pytest.raises(TypeError, match="specify z= or theta= or eta="): + vector.obj(x=1.0, y=2.0, theta=1.0, eta=1.0) + + +# ============================================================================ +# Valid combinations (ensure validation doesn't reject valid inputs) +# ============================================================================ + + +def test_valid_pt_phi_eta_mass_object(): + """vector.obj should accept pt, phi, eta, mass""" + vec = vector.obj(pt=1.0, phi=0.5, eta=1.0, mass=0.5) + assert vec.pt == 1.0 + assert vec.phi == 0.5 + assert vec.eta == 1.0 + assert vec.mass == 0.5 + + +def test_valid_pt_phi_eta_mass_numpy(): + """vector.array should accept pt, phi, eta, mass""" + arr = vector.array( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "mass": np.array([0.5, 0.5]), + } + ) + assert np.allclose(arr.pt, [1.0, 2.0]) + assert np.allclose(arr.phi, [0.5, 1.0]) + assert np.allclose(arr.eta, [1.0, 1.5]) + assert np.allclose(arr.mass, [0.5, 0.5]) + + +def test_valid_pt_phi_eta_mass_awkward(): + """vector.Array should accept pt, phi, eta, mass""" + arr = vector.Array( + ak.Array( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "mass": np.array([0.5, 0.5]), + } + ) + ) + assert ak.all(arr.pt == ak.Array([1.0, 2.0])) + assert ak.all(arr.phi == ak.Array([0.5, 1.0])) + assert ak.all(arr.eta == ak.Array([1.0, 1.5])) + assert ak.all(arr.mass == ak.Array([0.5, 0.5])) + + +def test_valid_pt_phi_eta_mass_zip(): + """vector.zip should accept pt, phi, eta, mass""" + arr = vector.zip( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "eta": np.array([1.0, 1.5]), + "mass": np.array([0.5, 0.5]), + } + ) + assert ak.all(arr.pt == ak.Array([1.0, 2.0])) + assert ak.all(arr.phi == ak.Array([0.5, 1.0])) + assert ak.all(arr.eta == ak.Array([1.0, 1.5])) + assert ak.all(arr.mass == ak.Array([0.5, 0.5])) + + +def test_valid_x_y_z_energy_object(): + """vector.obj should accept x, y, z, energy""" + vec = vector.obj(x=1.0, y=2.0, z=3.0, energy=5.0) + assert vec.x == 1.0 + assert vec.y == 2.0 + assert vec.z == 3.0 + assert vec.energy == 5.0 + + +def test_valid_x_y_z_energy_numpy(): + """vector.array should accept x, y, z, energy""" + arr = vector.array( + { + "x": np.array([1.0, 2.0]), + "y": np.array([2.0, 3.0]), + "z": np.array([3.0, 4.0]), + "energy": np.array([5.0, 6.0]), + } + ) + assert np.allclose(arr.x, [1.0, 2.0]) + assert np.allclose(arr.y, [2.0, 3.0]) + assert np.allclose(arr.z, [3.0, 4.0]) + assert np.allclose(arr.energy, [5.0, 6.0]) + + +def test_valid_x_y_z_energy_awkward(): + """vector.Array should accept x, y, z, energy""" + arr = vector.Array( + ak.Array( + { + "x": np.array([1.0, 2.0]), + "y": np.array([2.0, 3.0]), + "z": np.array([3.0, 4.0]), + "energy": np.array([5.0, 6.0]), + } + ) + ) + assert ak.all(arr.x == ak.Array([1.0, 2.0])) + assert ak.all(arr.y == ak.Array([2.0, 3.0])) + assert ak.all(arr.z == ak.Array([3.0, 4.0])) + assert ak.all(arr.energy == ak.Array([5.0, 6.0])) + + +def test_valid_x_y_z_energy_zip(): + """vector.zip should accept x, y, z, energy""" + arr = vector.zip( + { + "x": np.array([1.0, 2.0]), + "y": np.array([2.0, 3.0]), + "z": np.array([3.0, 4.0]), + "energy": np.array([5.0, 6.0]), + } + ) + assert ak.all(arr.x == ak.Array([1.0, 2.0])) + assert ak.all(arr.y == ak.Array([2.0, 3.0])) + assert ak.all(arr.z == ak.Array([3.0, 4.0])) + assert ak.all(arr.energy == ak.Array([5.0, 6.0])) + + +def test_valid_px_py_pz_E_object(): + """vector.obj should accept px, py, pz, E""" + vec = vector.obj(px=1.0, py=2.0, pz=3.0, E=5.0) + assert vec.px == 1.0 + assert vec.py == 2.0 + assert vec.pz == 3.0 + assert vec.E == 5.0 + + +def test_valid_px_py_pz_E_numpy(): + """vector.array should accept px, py, pz, E""" + arr = vector.array( + { + "px": np.array([1.0, 2.0]), + "py": np.array([2.0, 3.0]), + "pz": np.array([3.0, 4.0]), + "E": np.array([5.0, 6.0]), + } + ) + assert np.allclose(arr.px, [1.0, 2.0]) + assert np.allclose(arr.py, [2.0, 3.0]) + assert np.allclose(arr.pz, [3.0, 4.0]) + assert np.allclose(arr.E, [5.0, 6.0]) + + +def test_valid_px_py_pz_E_awkward(): + """vector.Array should accept px, py, pz, E""" + arr = vector.Array( + ak.Array( + { + "px": np.array([1.0, 2.0]), + "py": np.array([2.0, 3.0]), + "pz": np.array([3.0, 4.0]), + "E": np.array([5.0, 6.0]), + } + ) + ) + assert ak.all(arr.px == ak.Array([1.0, 2.0])) + assert ak.all(arr.py == ak.Array([2.0, 3.0])) + assert ak.all(arr.pz == ak.Array([3.0, 4.0])) + assert ak.all(ak.Array([5.0, 6.0]) == arr.E) + + +def test_valid_px_py_pz_E_zip(): + """vector.zip should accept px, py, pz, E""" + arr = vector.zip( + { + "px": np.array([1.0, 2.0]), + "py": np.array([2.0, 3.0]), + "pz": np.array([3.0, 4.0]), + "E": np.array([5.0, 6.0]), + } + ) + assert ak.all(arr.px == ak.Array([1.0, 2.0])) + assert ak.all(arr.py == ak.Array([2.0, 3.0])) + assert ak.all(arr.pz == ak.Array([3.0, 4.0])) + assert ak.all(ak.Array([5.0, 6.0]) == arr.E) + + +# ============================================================================ +# Incomplete azimuthal coordinate pairs +# ============================================================================ + + +def test_incomplete_x_without_y_object(): + """vector.obj should reject x without y""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.obj(x=1.0, z=3.0) + + +def test_incomplete_x_without_y_numpy(): + """vector.array should reject x without y""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.array( + { + "x": np.array([1.0, 2.0]), + "z": np.array([3.0, 4.0]), + } + ) + + +def test_incomplete_x_without_y_awkward(): + """vector.Array should reject x without y""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.Array( + ak.Array( + { + "x": np.array([1.0, 2.0]), + "z": np.array([3.0, 4.0]), + } + ) + ) + + +def test_incomplete_x_without_y_zip(): + """vector.zip should reject x without y""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.zip( + { + "x": np.array([1.0, 2.0]), + "z": np.array([3.0, 4.0]), + } + ) + + +def test_incomplete_rho_without_phi_object(): + """vector.obj should reject rho without phi""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.obj(rho=1.0, z=3.0) + + +def test_incomplete_rho_without_phi_numpy(): + """vector.array should reject rho without phi""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.array( + { + "rho": np.array([1.0, 2.0]), + "z": np.array([3.0, 4.0]), + } + ) + + +def test_incomplete_rho_without_phi_awkward(): + """vector.Array should reject rho without phi""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.Array( + ak.Array( + { + "rho": np.array([1.0, 2.0]), + "z": np.array([3.0, 4.0]), + } + ) + ) + + +def test_incomplete_rho_without_phi_zip(): + """vector.zip should reject rho without phi""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.zip( + { + "rho": np.array([1.0, 2.0]), + "z": np.array([3.0, 4.0]), + } + ) + + +# ============================================================================ +# Mixed azimuthal coordinate components +# ============================================================================ + + +def test_mixed_x_phi_object(): + """vector.obj should reject x with phi (mixed systems)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.obj(x=1.0, phi=0.5, z=3.0) + + +def test_mixed_x_phi_numpy(): + """vector.array should reject x with phi (mixed systems)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.array( + { + "x": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "z": np.array([3.0, 4.0]), + } + ) + + +def test_mixed_x_phi_awkward(): + """vector.Array should reject x with phi (mixed systems)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.Array( + ak.Array( + { + "x": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "z": np.array([3.0, 4.0]), + } + ) + ) + + +def test_mixed_x_phi_zip(): + """vector.zip should reject x with phi (mixed systems)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.zip( + { + "x": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "z": np.array([3.0, 4.0]), + } + ) + + +def test_mixed_y_rho_object(): + """vector.obj should reject y with rho (mixed systems)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.obj(y=2.0, rho=1.0, z=3.0) + + +def test_mixed_y_rho_numpy(): + """vector.array should reject y with rho (mixed systems)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.array( + { + "y": np.array([2.0, 3.0]), + "rho": np.array([1.0, 2.0]), + "z": np.array([3.0, 4.0]), + } + ) + + +def test_mixed_y_rho_awkward(): + """vector.Array should reject y with rho (mixed systems)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.Array( + ak.Array( + { + "y": np.array([2.0, 3.0]), + "rho": np.array([1.0, 2.0]), + "z": np.array([3.0, 4.0]), + } + ) + ) + + +def test_mixed_y_rho_zip(): + """vector.zip should reject y with rho (mixed systems)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.zip( + { + "y": np.array([2.0, 3.0]), + "rho": np.array([1.0, 2.0]), + "z": np.array([3.0, 4.0]), + } + ) + + +# ============================================================================ +# Temporal without proper 3D base +# ============================================================================ + + +def test_temporal_without_longitudinal_object(): + """vector.obj should reject x+y+t (temporal without longitudinal)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.obj(x=1.0, y=2.0, t=5.0) + + +def test_temporal_without_longitudinal_numpy(): + """vector.array should reject x+y+t (temporal without longitudinal)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.array( + { + "x": np.array([1.0, 2.0]), + "y": np.array([2.0, 3.0]), + "t": np.array([5.0, 6.0]), + } + ) + + +def test_temporal_without_longitudinal_awkward(): + """vector.Array should reject x+y+t (temporal without longitudinal)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.Array( + ak.Array( + { + "x": np.array([1.0, 2.0]), + "y": np.array([2.0, 3.0]), + "t": np.array([5.0, 6.0]), + } + ) + ) + + +def test_temporal_without_longitudinal_zip(): + """vector.zip should reject x+y+t (temporal without longitudinal)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.zip( + { + "x": np.array([1.0, 2.0]), + "y": np.array([2.0, 3.0]), + "t": np.array([5.0, 6.0]), + } + ) + + +def test_mass_without_longitudinal_object(): + """vector.obj should reject pt+phi+mass (temporal without longitudinal)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.obj(pt=1.0, phi=0.5, mass=0.5) + + +def test_mass_without_longitudinal_numpy(): + """vector.array should reject pt+phi+mass (temporal without longitudinal)""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.array( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "mass": np.array([0.5, 0.5]), + } + ) + + +def test_mass_without_longitudinal_awkward(): + """vector.Array should reject pt+phi+mass (temporal without longitudinal)""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.Array( + ak.Array( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "mass": np.array([0.5, 0.5]), + } + ) + ) + + +def test_mass_without_longitudinal_zip(): + """vector.zip should reject pt+phi+mass (temporal without longitudinal)""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.zip( + { + "pt": np.array([1.0, 2.0]), + "phi": np.array([0.5, 1.0]), + "mass": np.array([0.5, 0.5]), + } + ) + + +# ============================================================================ +# Missing required coordinates +# ============================================================================ + + +def test_only_temporal_object(): + """vector.obj should reject only t (missing spatial coords)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.obj(t=5.0) + + +def test_only_temporal_numpy(): + """vector.array should reject only t (missing spatial coords)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.array( + { + "t": np.array([5.0, 6.0]), + } + ) + + +def test_only_temporal_awkward(): + """vector.Array should reject only t (missing spatial coords)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.Array( + ak.Array( + { + "t": np.array([5.0, 6.0]), + } + ) + ) + + +def test_only_temporal_zip(): + """vector.zip should reject only t (missing spatial coords)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.zip( + { + "t": np.array([5.0, 6.0]), + } + ) + + +def test_only_longitudinal_object(): + """vector.obj should reject only z (missing azimuthal coords)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.obj(z=3.0) + + +def test_only_longitudinal_numpy(): + """vector.array should reject only z (missing azimuthal coords)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.array( + { + "z": np.array([3.0, 4.0]), + } + ) + + +def test_only_longitudinal_awkward(): + """vector.Array should reject only z (missing azimuthal coords)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.Array( + ak.Array( + { + "z": np.array([3.0, 4.0]), + } + ) + ) + + +def test_only_longitudinal_zip(): + """vector.zip should reject only z (missing azimuthal coords)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.zip( + { + "z": np.array([3.0, 4.0]), + } + ) + + +def test_longitudinal_temporal_without_azimuthal_object(): + """vector.obj should reject z+t without azimuthal coords""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.obj(z=3.0, t=5.0) + + +def test_longitudinal_temporal_without_azimuthal_numpy(): + """vector.array should reject z+t without azimuthal coords""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.array( + { + "z": np.array([3.0, 4.0]), + "t": np.array([5.0, 6.0]), + } + ) + + +def test_longitudinal_temporal_without_azimuthal_awkward(): + """vector.Array should reject z+t without azimuthal coords""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.Array( + ak.Array( + { + "z": np.array([3.0, 4.0]), + "t": np.array([5.0, 6.0]), + } + ) + ) + + +def test_longitudinal_temporal_without_azimuthal_zip(): + """vector.zip should reject z+t without azimuthal coords""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.zip( + { + "z": np.array([3.0, 4.0]), + "t": np.array([5.0, 6.0]), + } + ) From d1edbe5436b4657c0393cdb2b84f25721226363a Mon Sep 17 00:00:00 2001 From: Iason Krommydas Date: Wed, 4 Feb 2026 15:07:21 -0600 Subject: [PATCH 05/15] call validation function in __array_finalize__ --- src/vector/backends/numpy.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/src/vector/backends/numpy.py b/src/vector/backends/numpy.py index bf394f8a..1abcb037 100644 --- a/src/vector/backends/numpy.py +++ b/src/vector/backends/numpy.py @@ -1190,6 +1190,8 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return + _validate_numpy_coordinates(self.dtype.names or ()) + if _has(self, ("x", "y")): self._azimuthal_type = AzimuthalNumpyXY elif _has(self, ("rho", "phi")): @@ -1363,6 +1365,8 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return + _validate_numpy_coordinates(self.dtype.names or ()) + self.dtype.names = tuple( _repr_momentum_to_generic.get(x, x) for x in (self.dtype.names or ()) ) @@ -1431,6 +1435,8 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return + _validate_numpy_coordinates(self.dtype.names or ()) + if _has(self, ("x", "y")): self._azimuthal_type = AzimuthalNumpyXY elif _has(self, ("rho", "phi")): @@ -1664,6 +1670,8 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return + _validate_numpy_coordinates(self.dtype.names or ()) + self.dtype.names = tuple( _repr_momentum_to_generic.get(x, x) for x in (self.dtype.names or ()) ) @@ -1745,6 +1753,8 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return + _validate_numpy_coordinates(self.dtype.names or ()) + if _has(self, ("x", "y")): self._azimuthal_type = AzimuthalNumpyXY elif _has(self, ("rho", "phi")): @@ -2047,6 +2057,8 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return + _validate_numpy_coordinates(self.dtype.names or ()) + self.dtype.names = tuple( _repr_momentum_to_generic.get(x, x) for x in (self.dtype.names or ()) ) @@ -2321,9 +2333,6 @@ def array(*args: typing.Any, **kwargs: typing.Any) -> VectorNumpy: is_momentum = any(x in _repr_momentum_to_generic for x in names) - # Validate coordinates using dimension-guard pattern (same as awkward _check_names) - _validate_numpy_coordinates(names) - if any(x in ("t", "E", "e", "energy", "tau", "M", "m", "mass") for x in names): cls = MomentumNumpy4D if is_momentum else VectorNumpy4D elif any(x in ("z", "pz", "theta", "eta") for x in names): From 0fa491e18148d9b08b55bce3b6175d380d3df403 Mon Sep 17 00:00:00 2001 From: Iason Krommydas Date: Wed, 4 Feb 2026 15:09:33 -0600 Subject: [PATCH 06/15] move test file --- tests/{backends/test_coordinate_validation.py => test_pr_659.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/{backends/test_coordinate_validation.py => test_pr_659.py} (100%) diff --git a/tests/backends/test_coordinate_validation.py b/tests/test_pr_659.py similarity index 100% rename from tests/backends/test_coordinate_validation.py rename to tests/test_pr_659.py From bfe3f0b1db7b5b96613ab033c5d14f19bc39f633 Mon Sep 17 00:00:00 2001 From: Iason Krommydas Date: Wed, 4 Feb 2026 15:32:02 -0600 Subject: [PATCH 07/15] Revert "call validation function in __array_finalize__" This reverts commit d1edbe5436b4657c0393cdb2b84f25721226363a. --- src/vector/backends/numpy.py | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) diff --git a/src/vector/backends/numpy.py b/src/vector/backends/numpy.py index 1abcb037..bf394f8a 100644 --- a/src/vector/backends/numpy.py +++ b/src/vector/backends/numpy.py @@ -1190,8 +1190,6 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return - _validate_numpy_coordinates(self.dtype.names or ()) - if _has(self, ("x", "y")): self._azimuthal_type = AzimuthalNumpyXY elif _has(self, ("rho", "phi")): @@ -1365,8 +1363,6 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return - _validate_numpy_coordinates(self.dtype.names or ()) - self.dtype.names = tuple( _repr_momentum_to_generic.get(x, x) for x in (self.dtype.names or ()) ) @@ -1435,8 +1431,6 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return - _validate_numpy_coordinates(self.dtype.names or ()) - if _has(self, ("x", "y")): self._azimuthal_type = AzimuthalNumpyXY elif _has(self, ("rho", "phi")): @@ -1670,8 +1664,6 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return - _validate_numpy_coordinates(self.dtype.names or ()) - self.dtype.names = tuple( _repr_momentum_to_generic.get(x, x) for x in (self.dtype.names or ()) ) @@ -1753,8 +1745,6 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return - _validate_numpy_coordinates(self.dtype.names or ()) - if _has(self, ("x", "y")): self._azimuthal_type = AzimuthalNumpyXY elif _has(self, ("rho", "phi")): @@ -2057,8 +2047,6 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return - _validate_numpy_coordinates(self.dtype.names or ()) - self.dtype.names = tuple( _repr_momentum_to_generic.get(x, x) for x in (self.dtype.names or ()) ) @@ -2333,6 +2321,9 @@ def array(*args: typing.Any, **kwargs: typing.Any) -> VectorNumpy: is_momentum = any(x in _repr_momentum_to_generic for x in names) + # Validate coordinates using dimension-guard pattern (same as awkward _check_names) + _validate_numpy_coordinates(names) + if any(x in ("t", "E", "e", "energy", "tau", "M", "m", "mass") for x in names): cls = MomentumNumpy4D if is_momentum else VectorNumpy4D elif any(x in ("z", "pz", "theta", "eta") for x in names): From 489dccf715089ef7ba4a3891f86c66c46a511b7a Mon Sep 17 00:00:00 2001 From: Iason Krommydas Date: Wed, 4 Feb 2026 16:16:00 -0600 Subject: [PATCH 08/15] sympy validation and testing --- src/vector/backends/sympy.py | 23 +++++ tests/test_pr_659.py | 184 ++++++++++++++++++++++++++++++++++- 2 files changed, 206 insertions(+), 1 deletion(-) diff --git a/src/vector/backends/sympy.py b/src/vector/backends/sympy.py index fe1e001e..6063a67f 100644 --- a/src/vector/backends/sympy.py +++ b/src/vector/backends/sympy.py @@ -437,6 +437,23 @@ def _replace_data(obj: typing.Any, result: typing.Any) -> typing.Any: return obj +def _validate_sympy_coordinates(coordinates: dict[str, typing.Any]) -> None: + """ + Validate coordinate names for duplicate/conflicting coordinates. + + Checks that no two coordinate names map to the same generic coordinate + (e.g., E and e both map to t, so having both is invalid). + + Raises TypeError if duplicate or conflicting coordinates are detected. + """ + generic_keys = [_repr_momentum_to_generic.get(k, k) for k in coordinates] + if len(generic_keys) != len(set(generic_keys)): + raise TypeError( + "duplicate coordinates (through momentum-aliases): " + + ", ".join(repr(x) for x in coordinates) + ) + + class VectorSympy(Vector): # noqa: PLW1641 """Mixin class for Sympy vectors.""" @@ -744,6 +761,8 @@ class VectorSympy2D(VectorSympy, Planar, Vector2D): azimuthal: AzimuthalSympy def __init__(self, azimuthal: AzimuthalSympy | None = None, **kwargs: sympy.Symbol): + _validate_sympy_coordinates(kwargs) + for k, v in kwargs.copy().items(): kwargs.pop(k) kwargs[_repr_momentum_to_generic.get(k, k)] = v @@ -945,6 +964,8 @@ def __init__( longitudinal: LongitudinalSympy | None = None, **kwargs: sympy.Symbol, ): + _validate_sympy_coordinates(kwargs) + for k, v in kwargs.copy().items(): kwargs.pop(k) kwargs[_repr_momentum_to_generic.get(k, k)] = v @@ -1219,6 +1240,8 @@ def __init__( temporal: TemporalSympy | None = None, **kwargs: sympy.Symbol, ): + _validate_sympy_coordinates(kwargs) + for k, v in kwargs.copy().items(): kwargs.pop(k) kwargs[_repr_momentum_to_generic.get(k, k)] = v diff --git a/tests/test_pr_659.py b/tests/test_pr_659.py index 3f524b7e..b763dc05 100644 --- a/tests/test_pr_659.py +++ b/tests/test_pr_659.py @@ -11,8 +11,19 @@ import vector ak = pytest.importorskip("awkward") +sympy = pytest.importorskip("sympy") -pytestmark = pytest.mark.awkward +pytestmark = [pytest.mark.awkward, pytest.mark.sympy] + +# Define sympy symbols for tests +_x, _y = sympy.symbols("x y") +_rho, _phi = sympy.symbols("rho phi") +_z, _eta, _theta = sympy.symbols("z eta theta") +_t, _tau = sympy.symbols("t tau") +_px, _py = sympy.symbols("px py") +_pt = sympy.symbols("pt") +_pz = sympy.symbols("pz") +_M, _m, _mass, _E, _e, _energy = sympy.symbols("M m mass E e energy") # ============================================================================ @@ -73,6 +84,12 @@ def test_duplicate_E_e_zip(): ) +def test_duplicate_E_e_sympy(): + """MomentumSympy4D should reject E + e""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.MomentumSympy4D(pt=_pt, phi=_phi, eta=_eta, E=_E, e=_e) + + def test_duplicate_E_energy_object(): """vector.obj should reject E + energy""" with pytest.raises(TypeError, match="unrecognized combination"): @@ -123,6 +140,12 @@ def test_duplicate_E_energy_zip(): ) +def test_duplicate_E_energy_sympy(): + """MomentumSympy4D should reject E + energy""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.MomentumSympy4D(pt=_pt, phi=_phi, eta=_eta, E=_E, energy=_energy) + + def test_duplicate_e_energy_object(): """vector.obj should reject e + energy""" with pytest.raises(TypeError, match="unrecognized combination"): @@ -173,6 +196,12 @@ def test_duplicate_e_energy_zip(): ) +def test_duplicate_e_energy_sympy(): + """MomentumSympy4D should reject e + energy""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.MomentumSympy4D(pt=_pt, phi=_phi, eta=_eta, e=_e, energy=_energy) + + def test_duplicate_M_m_object(): """vector.obj should reject M + m""" with pytest.raises(TypeError, match="unrecognized combination"): @@ -223,6 +252,12 @@ def test_duplicate_M_m_zip(): ) +def test_duplicate_M_m_sympy(): + """MomentumSympy4D should reject M + m""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.MomentumSympy4D(pt=_pt, phi=_phi, eta=_eta, M=_M, m=_m) + + def test_duplicate_M_mass_object(): """vector.obj should reject M + mass""" with pytest.raises(TypeError, match="unrecognized combination"): @@ -273,6 +308,12 @@ def test_duplicate_M_mass_zip(): ) +def test_duplicate_M_mass_sympy(): + """MomentumSympy4D should reject M + mass""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.MomentumSympy4D(pt=_pt, phi=_phi, eta=_eta, M=_M, mass=_mass) + + def test_duplicate_m_mass_object(): """vector.obj should reject m + mass""" with pytest.raises(TypeError, match="unrecognized combination"): @@ -323,6 +364,12 @@ def test_duplicate_m_mass_zip(): ) +def test_duplicate_m_mass_sympy(): + """MomentumSympy4D should reject m + mass""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.MomentumSympy4D(pt=_pt, phi=_phi, eta=_eta, m=_m, mass=_mass) + + def test_conflicting_energy_mass_object(): """vector.obj should reject energy + mass (t-like + tau-like)""" with pytest.raises(TypeError, match="specify t= or tau=, but not more than one"): @@ -373,6 +420,12 @@ def test_conflicting_energy_mass_zip(): ) +def test_conflicting_energy_mass_sympy(): + """MomentumSympy4D should reject energy + mass (t-like + tau-like)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.MomentumSympy4D(pt=_pt, phi=_phi, eta=_eta, energy=_energy, mass=_mass) + + def test_conflicting_t_tau_object(): """vector.obj should reject t + tau""" with pytest.raises(TypeError, match="specify t= or tau="): @@ -423,6 +476,12 @@ def test_conflicting_t_tau_zip(): ) +def test_conflicting_t_tau_sympy(): + """VectorSympy4D should reject t + tau""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.VectorSympy4D(x=_x, y=_y, z=_z, t=_t, tau=_tau) + + def test_conflicting_E_mass_object(): """vector.obj should reject E + mass""" with pytest.raises(TypeError, match="specify t= or tau=, but not more than one"): @@ -473,6 +532,12 @@ def test_conflicting_E_mass_zip(): ) +def test_conflicting_E_mass_sympy(): + """MomentumSympy4D should reject E + mass""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.MomentumSympy4D(pt=_pt, phi=_phi, eta=_eta, E=_E, mass=_mass) + + def test_conflicting_t_mass_object(): """vector.obj should reject t + mass""" with pytest.raises(TypeError, match="specify t= or tau=, but not more than one"): @@ -523,6 +588,12 @@ def test_conflicting_t_mass_zip(): ) +def test_conflicting_t_mass_sympy(): + """MomentumSympy4D should reject t + mass""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.MomentumSympy4D(x=_x, y=_y, z=_z, t=_t, mass=_mass) + + def test_conflicting_energy_tau_object(): """vector.obj should reject energy + tau""" with pytest.raises(TypeError, match="specify t= or tau=, but not more than one"): @@ -573,6 +644,12 @@ def test_conflicting_energy_tau_zip(): ) +def test_conflicting_energy_tau_sympy(): + """MomentumSympy4D should reject energy + tau""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.MomentumSympy4D(x=_x, y=_y, z=_z, energy=_energy, tau=_tau) + + # ============================================================================ # Duplicate azimuthal coordinates # ============================================================================ @@ -629,6 +706,12 @@ def test_duplicate_px_x_zip(): ) +def test_duplicate_px_x_sympy(): + """MomentumSympy4D should reject px + x""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.MomentumSympy4D(px=_px, x=_x, y=_y, z=_z, t=_t) + + def test_duplicate_py_y_object(): """vector.obj should reject py + y""" with pytest.raises(TypeError, match="duplicate coordinates"): @@ -679,6 +762,12 @@ def test_duplicate_py_y_zip(): ) +def test_duplicate_py_y_sympy(): + """MomentumSympy4D should reject py + y""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.MomentumSympy4D(x=_x, py=_py, y=_y, z=_z, t=_t) + + def test_duplicate_pt_rho_object(): """vector.obj should reject pt + rho""" with pytest.raises(TypeError, match="duplicate coordinates"): @@ -729,6 +818,12 @@ def test_duplicate_pt_rho_zip(): ) +def test_duplicate_pt_rho_sympy(): + """MomentumSympy4D should reject pt + rho""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.MomentumSympy4D(pt=_pt, rho=_rho, phi=_phi, eta=_eta, mass=_mass) + + # ============================================================================ # Duplicate longitudinal coordinates # ============================================================================ @@ -785,6 +880,12 @@ def test_duplicate_pz_z_zip(): ) +def test_duplicate_pz_z_sympy(): + """MomentumSympy4D should reject pz + z""" + with pytest.raises(TypeError, match="duplicate coordinates"): + vector.MomentumSympy4D(x=_x, y=_y, pz=_pz, z=_z, t=_t) + + # ============================================================================ # Mixed azimuthal coordinate systems (from _gather_coordinates) # ============================================================================ @@ -901,6 +1002,15 @@ def test_valid_pt_phi_eta_mass_zip(): assert ak.all(arr.mass == ak.Array([0.5, 0.5])) +def test_valid_pt_phi_eta_mass_sympy(): + """MomentumSympy4D should accept pt, phi, eta, mass""" + vec = vector.MomentumSympy4D(pt=_pt, phi=_phi, eta=_eta, mass=_mass) + assert vec.pt == _pt + assert vec.phi == _phi + assert vec.eta == _eta + assert vec.mass == _mass + + def test_valid_x_y_z_energy_object(): """vector.obj should accept x, y, z, energy""" vec = vector.obj(x=1.0, y=2.0, z=3.0, energy=5.0) @@ -960,6 +1070,15 @@ def test_valid_x_y_z_energy_zip(): assert ak.all(arr.energy == ak.Array([5.0, 6.0])) +def test_valid_x_y_z_energy_sympy(): + """MomentumSympy4D should accept x, y, z, energy""" + vec = vector.MomentumSympy4D(x=_x, y=_y, z=_z, energy=_energy) + assert vec.x == _x + assert vec.y == _y + assert vec.z == _z + assert vec.energy == _energy + + def test_valid_px_py_pz_E_object(): """vector.obj should accept px, py, pz, E""" vec = vector.obj(px=1.0, py=2.0, pz=3.0, E=5.0) @@ -1019,6 +1138,15 @@ def test_valid_px_py_pz_E_zip(): assert ak.all(ak.Array([5.0, 6.0]) == arr.E) +def test_valid_px_py_pz_E_sympy(): + """MomentumSympy4D should accept px, py, pz, E""" + vec = vector.MomentumSympy4D(px=_px, py=_py, pz=_pz, E=_E) + assert vec.px == _px + assert vec.py == _py + assert vec.pz == _pz + assert vec.E == _E + + # ============================================================================ # Incomplete azimuthal coordinate pairs # ============================================================================ @@ -1065,6 +1193,12 @@ def test_incomplete_x_without_y_zip(): ) +def test_incomplete_x_without_y_sympy(): + """VectorSympy3D should reject x without y""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.VectorSympy3D(x=_x, z=_z) + + def test_incomplete_rho_without_phi_object(): """vector.obj should reject rho without phi""" with pytest.raises(TypeError, match="unrecognized combination"): @@ -1106,6 +1240,12 @@ def test_incomplete_rho_without_phi_zip(): ) +def test_incomplete_rho_without_phi_sympy(): + """VectorSympy3D should reject rho without phi""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.VectorSympy3D(rho=_rho, z=_z) + + # ============================================================================ # Mixed azimuthal coordinate components # ============================================================================ @@ -1155,6 +1295,12 @@ def test_mixed_x_phi_zip(): ) +def test_mixed_x_phi_sympy(): + """VectorSympy3D should reject x with phi (mixed systems)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.VectorSympy3D(x=_x, phi=_phi, z=_z) + + def test_mixed_y_rho_object(): """vector.obj should reject y with rho (mixed systems)""" with pytest.raises(TypeError, match="unrecognized combination"): @@ -1199,6 +1345,12 @@ def test_mixed_y_rho_zip(): ) +def test_mixed_y_rho_sympy(): + """VectorSympy3D should reject y with rho (mixed systems)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.VectorSympy3D(y=_y, rho=_rho, z=_z) + + # ============================================================================ # Temporal without proper 3D base # ============================================================================ @@ -1248,6 +1400,12 @@ def test_temporal_without_longitudinal_zip(): ) +def test_temporal_without_longitudinal_sympy(): + """VectorSympy4D should reject x+y+t (temporal without longitudinal)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.VectorSympy4D(x=_x, y=_y, t=_t) + + def test_mass_without_longitudinal_object(): """vector.obj should reject pt+phi+mass (temporal without longitudinal)""" with pytest.raises(TypeError, match="unrecognized combination"): @@ -1292,6 +1450,12 @@ def test_mass_without_longitudinal_zip(): ) +def test_mass_without_longitudinal_sympy(): + """MomentumSympy4D should reject pt+phi+mass (temporal without longitudinal)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.MomentumSympy4D(pt=_pt, phi=_phi, mass=_mass) + + # ============================================================================ # Missing required coordinates # ============================================================================ @@ -1335,6 +1499,12 @@ def test_only_temporal_zip(): ) +def test_only_temporal_sympy(): + """VectorSympy4D should reject only t (missing spatial coords)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.VectorSympy4D(t=_t) + + def test_only_longitudinal_object(): """vector.obj should reject only z (missing azimuthal coords)""" with pytest.raises(TypeError, match="unrecognized combination"): @@ -1373,6 +1543,12 @@ def test_only_longitudinal_zip(): ) +def test_only_longitudinal_sympy(): + """VectorSympy3D should reject only z (missing azimuthal coords)""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.VectorSympy3D(z=_z) + + def test_longitudinal_temporal_without_azimuthal_object(): """vector.obj should reject z+t without azimuthal coords""" with pytest.raises(TypeError, match="unrecognized combination"): @@ -1412,3 +1588,9 @@ def test_longitudinal_temporal_without_azimuthal_zip(): "t": np.array([5.0, 6.0]), } ) + + +def test_longitudinal_temporal_without_azimuthal_sympy(): + """VectorSympy4D should reject z+t without azimuthal coords""" + with pytest.raises(TypeError, match="unrecognized combination"): + vector.VectorSympy4D(z=_z, t=_t) From a25bec31e57cc69cba0e9fd30de815a1a9968dae Mon Sep 17 00:00:00 2001 From: Iason Krommydas Date: Wed, 4 Feb 2026 17:14:29 -0600 Subject: [PATCH 09/15] share common validation function between numpy and sympy --- src/vector/_methods.py | 162 +++++++++++++++++++++++++++++++++ src/vector/backends/numpy.py | 165 +--------------------------------- src/vector/backends/object.py | 1 - src/vector/backends/sympy.py | 24 +---- tests/test_pr_659.py | 14 +-- 5 files changed, 175 insertions(+), 191 deletions(-) diff --git a/src/vector/_methods.py b/src/vector/_methods.py index 9a668e3f..42442189 100644 --- a/src/vector/_methods.py +++ b/src/vector/_methods.py @@ -4496,3 +4496,165 @@ def _flavor_of(*objects: VectorProtocol) -> type[VectorProtocol]: return handler.MomentumClass else: return handler.GenericClass + + +def _validate_coordinates(fieldnames: tuple[str, ...]) -> None: + """ + Validate coordinate field names for constructing vectors. + + This follows the same logic as _check_names in awkward_constructors to ensure + consistent validation across backends. + + Raises TypeError if duplicate or conflicting coordinates are detected. + """ + 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=" + ) + + is_momentum = False + dimension = 0 + fieldnames_copy = list(fieldnames) + + # 2D azimuthal coordinates + if "x" in fieldnames_copy and "y" in fieldnames_copy: + if dimension != 0: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 2 + fieldnames_copy.remove("x") + fieldnames_copy.remove("y") + if "rho" in fieldnames_copy and "phi" in fieldnames_copy: + if dimension != 0: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 2 + fieldnames_copy.remove("rho") + fieldnames_copy.remove("phi") + if "x" in fieldnames_copy and "py" in fieldnames_copy: + is_momentum = True + if dimension != 0: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 2 + fieldnames_copy.remove("x") + fieldnames_copy.remove("py") + if "px" in fieldnames_copy and "y" in fieldnames_copy: + is_momentum = True + if dimension != 0: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 2 + fieldnames_copy.remove("px") + fieldnames_copy.remove("y") + if "px" in fieldnames_copy and "py" in fieldnames_copy: + is_momentum = True + if dimension != 0: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 2 + fieldnames_copy.remove("px") + fieldnames_copy.remove("py") + if "pt" in fieldnames_copy and "phi" in fieldnames_copy: + is_momentum = True + if dimension != 0: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 2 + fieldnames_copy.remove("pt") + fieldnames_copy.remove("phi") + + # 3D longitudinal coordinates + if "z" in fieldnames_copy: + if dimension != 2: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 3 + fieldnames_copy.remove("z") + if "theta" in fieldnames_copy: + if dimension != 2: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 3 + fieldnames_copy.remove("theta") + if "eta" in fieldnames_copy: + if dimension != 2: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 3 + fieldnames_copy.remove("eta") + if "pz" in fieldnames_copy: + is_momentum = True + if dimension != 2: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 3 + fieldnames_copy.remove("pz") + + # 4D temporal coordinates + if "t" in fieldnames_copy: + if dimension != 3: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 4 + fieldnames_copy.remove("t") + if "tau" in fieldnames_copy: + if dimension != 3: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 4 + fieldnames_copy.remove("tau") + if "E" in fieldnames_copy: + is_momentum = True + if dimension != 3: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 4 + fieldnames_copy.remove("E") + if "e" in fieldnames_copy: + is_momentum = True + if dimension != 3: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 4 + fieldnames_copy.remove("e") + if "energy" in fieldnames_copy: + is_momentum = True + if dimension != 3: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 4 + fieldnames_copy.remove("energy") + if "M" in fieldnames_copy: + is_momentum = True + if dimension != 3: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 4 + fieldnames_copy.remove("M") + if "m" in fieldnames_copy: + is_momentum = True + if dimension != 3: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 4 + fieldnames_copy.remove("m") + if "mass" in fieldnames_copy: + is_momentum = True + if dimension != 3: + raise TypeError(complaint1 if is_momentum else complaint2) + dimension = 4 + fieldnames_copy.remove("mass") + + # Check if any remaining fieldnames would conflict with already-processed coordinates + # when mapped to generic names (e.g., pt was processed, rho shouldn't remain) + if fieldnames_copy: + # Map all original fieldnames to generic names to detect conflicts + generic_names = [_repr_momentum_to_generic.get(x, x) for x in fieldnames] + if len(generic_names) != len(set(generic_names)): + raise TypeError(complaint1 if is_momentum else complaint2) diff --git a/src/vector/backends/numpy.py b/src/vector/backends/numpy.py index bf394f8a..088d8487 100644 --- a/src/vector/backends/numpy.py +++ b/src/vector/backends/numpy.py @@ -52,6 +52,7 @@ _ltype, _repr_momentum_to_generic, _ttype, + _validate_coordinates, ) from vector._typeutils import BoolCollection, FloatArray, ScalarCollection @@ -2092,168 +2093,6 @@ def __setitem__(self, where: typing.Any, what: typing.Any) -> None: return _setitem(self, where, what, True) -def _validate_numpy_coordinates(fieldnames: tuple[str, ...]) -> None: - """ - Validate coordinate field names using dimension-guard pattern. - - This follows the same logic as _check_names in awkward_constructors to ensure - consistent validation across backends. - - Raises TypeError if duplicate or conflicting coordinates are detected. - """ - 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=" - ) - - is_momentum = False - dimension = 0 - fieldnames_copy = list(fieldnames) - - # 2D azimuthal coordinates - if "x" in fieldnames_copy and "y" in fieldnames_copy: - if dimension != 0: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 2 - fieldnames_copy.remove("x") - fieldnames_copy.remove("y") - if "rho" in fieldnames_copy and "phi" in fieldnames_copy: - if dimension != 0: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 2 - fieldnames_copy.remove("rho") - fieldnames_copy.remove("phi") - if "x" in fieldnames_copy and "py" in fieldnames_copy: - is_momentum = True - if dimension != 0: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 2 - fieldnames_copy.remove("x") - fieldnames_copy.remove("py") - if "px" in fieldnames_copy and "y" in fieldnames_copy: - is_momentum = True - if dimension != 0: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 2 - fieldnames_copy.remove("px") - fieldnames_copy.remove("y") - if "px" in fieldnames_copy and "py" in fieldnames_copy: - is_momentum = True - if dimension != 0: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 2 - fieldnames_copy.remove("px") - fieldnames_copy.remove("py") - if "pt" in fieldnames_copy and "phi" in fieldnames_copy: - is_momentum = True - if dimension != 0: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 2 - fieldnames_copy.remove("pt") - fieldnames_copy.remove("phi") - - # 3D longitudinal coordinates - if "z" in fieldnames_copy: - if dimension != 2: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 3 - fieldnames_copy.remove("z") - if "theta" in fieldnames_copy: - if dimension != 2: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 3 - fieldnames_copy.remove("theta") - if "eta" in fieldnames_copy: - if dimension != 2: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 3 - fieldnames_copy.remove("eta") - if "pz" in fieldnames_copy: - is_momentum = True - if dimension != 2: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 3 - fieldnames_copy.remove("pz") - - # 4D temporal coordinates - if "t" in fieldnames_copy: - if dimension != 3: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 4 - fieldnames_copy.remove("t") - if "tau" in fieldnames_copy: - if dimension != 3: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 4 - fieldnames_copy.remove("tau") - if "E" in fieldnames_copy: - is_momentum = True - if dimension != 3: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 4 - fieldnames_copy.remove("E") - if "e" in fieldnames_copy: - is_momentum = True - if dimension != 3: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 4 - fieldnames_copy.remove("e") - if "energy" in fieldnames_copy: - is_momentum = True - if dimension != 3: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 4 - fieldnames_copy.remove("energy") - if "M" in fieldnames_copy: - is_momentum = True - if dimension != 3: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 4 - fieldnames_copy.remove("M") - if "m" in fieldnames_copy: - is_momentum = True - if dimension != 3: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 4 - fieldnames_copy.remove("m") - if "mass" in fieldnames_copy: - is_momentum = True - if dimension != 3: - raise TypeError(complaint1 if is_momentum else complaint2) - dimension = 4 - fieldnames_copy.remove("mass") - - # Check if any remaining fieldnames would conflict with already-processed coordinates - # when mapped to generic names (e.g., pt was processed, rho shouldn't remain) - if fieldnames_copy: - # Map all original fieldnames to generic names to detect conflicts - generic_names = [_repr_momentum_to_generic.get(x, x) for x in fieldnames] - if len(generic_names) != len(set(generic_names)): - raise TypeError(complaint1 if is_momentum else complaint2) - - def array(*args: typing.Any, **kwargs: typing.Any) -> VectorNumpy: """ Constructs a NumPy array of vectors, whose type is determined by the dtype @@ -2322,7 +2161,7 @@ def array(*args: typing.Any, **kwargs: typing.Any) -> VectorNumpy: is_momentum = any(x in _repr_momentum_to_generic for x in names) # Validate coordinates using dimension-guard pattern (same as awkward _check_names) - _validate_numpy_coordinates(names) + _validate_coordinates(names) if any(x in ("t", "E", "e", "energy", "tau", "M", "m", "mass") for x in names): cls = MomentumNumpy4D if is_momentum else VectorNumpy4D diff --git a/src/vector/backends/object.py b/src/vector/backends/object.py index dd1e040d..2b87e1b9 100644 --- a/src/vector/backends/object.py +++ b/src/vector/backends/object.py @@ -3214,7 +3214,6 @@ def obj(**coordinates: float) -> VectorObject: generic_coordinates = {} _is_type_safe(coordinates) - if "px" in coordinates: is_momentum = True generic_coordinates["x"] = coordinates.pop("px") diff --git a/src/vector/backends/sympy.py b/src/vector/backends/sympy.py index 6063a67f..dd568c53 100644 --- a/src/vector/backends/sympy.py +++ b/src/vector/backends/sympy.py @@ -36,6 +36,7 @@ _repr_generic_to_momentum, _repr_momentum_to_generic, _ttype, + _validate_coordinates, ) @@ -437,23 +438,6 @@ def _replace_data(obj: typing.Any, result: typing.Any) -> typing.Any: return obj -def _validate_sympy_coordinates(coordinates: dict[str, typing.Any]) -> None: - """ - Validate coordinate names for duplicate/conflicting coordinates. - - Checks that no two coordinate names map to the same generic coordinate - (e.g., E and e both map to t, so having both is invalid). - - Raises TypeError if duplicate or conflicting coordinates are detected. - """ - generic_keys = [_repr_momentum_to_generic.get(k, k) for k in coordinates] - if len(generic_keys) != len(set(generic_keys)): - raise TypeError( - "duplicate coordinates (through momentum-aliases): " - + ", ".join(repr(x) for x in coordinates) - ) - - class VectorSympy(Vector): # noqa: PLW1641 """Mixin class for Sympy vectors.""" @@ -761,7 +745,7 @@ class VectorSympy2D(VectorSympy, Planar, Vector2D): azimuthal: AzimuthalSympy def __init__(self, azimuthal: AzimuthalSympy | None = None, **kwargs: sympy.Symbol): - _validate_sympy_coordinates(kwargs) + _validate_coordinates(tuple(kwargs)) for k, v in kwargs.copy().items(): kwargs.pop(k) @@ -964,7 +948,7 @@ def __init__( longitudinal: LongitudinalSympy | None = None, **kwargs: sympy.Symbol, ): - _validate_sympy_coordinates(kwargs) + _validate_coordinates(tuple(kwargs)) for k, v in kwargs.copy().items(): kwargs.pop(k) @@ -1240,7 +1224,7 @@ def __init__( temporal: TemporalSympy | None = None, **kwargs: sympy.Symbol, ): - _validate_sympy_coordinates(kwargs) + _validate_coordinates(tuple(kwargs)) for k, v in kwargs.copy().items(): kwargs.pop(k) diff --git a/tests/test_pr_659.py b/tests/test_pr_659.py index b763dc05..1451d7c4 100644 --- a/tests/test_pr_659.py +++ b/tests/test_pr_659.py @@ -422,7 +422,7 @@ def test_conflicting_energy_mass_zip(): def test_conflicting_energy_mass_sympy(): """MomentumSympy4D should reject energy + mass (t-like + tau-like)""" - with pytest.raises(TypeError, match="unrecognized combination"): + with pytest.raises(TypeError, match="duplicate coordinates"): vector.MomentumSympy4D(pt=_pt, phi=_phi, eta=_eta, energy=_energy, mass=_mass) @@ -534,7 +534,7 @@ def test_conflicting_E_mass_zip(): def test_conflicting_E_mass_sympy(): """MomentumSympy4D should reject E + mass""" - with pytest.raises(TypeError, match="unrecognized combination"): + with pytest.raises(TypeError, match="duplicate coordinates"): vector.MomentumSympy4D(pt=_pt, phi=_phi, eta=_eta, E=_E, mass=_mass) @@ -590,7 +590,7 @@ def test_conflicting_t_mass_zip(): def test_conflicting_t_mass_sympy(): """MomentumSympy4D should reject t + mass""" - with pytest.raises(TypeError, match="unrecognized combination"): + with pytest.raises(TypeError, match="duplicate coordinates"): vector.MomentumSympy4D(x=_x, y=_y, z=_z, t=_t, mass=_mass) @@ -646,7 +646,7 @@ def test_conflicting_energy_tau_zip(): def test_conflicting_energy_tau_sympy(): """MomentumSympy4D should reject energy + tau""" - with pytest.raises(TypeError, match="unrecognized combination"): + with pytest.raises(TypeError, match="duplicate coordinates"): vector.MomentumSympy4D(x=_x, y=_y, z=_z, energy=_energy, tau=_tau) @@ -708,7 +708,7 @@ def test_duplicate_px_x_zip(): def test_duplicate_px_x_sympy(): """MomentumSympy4D should reject px + x""" - with pytest.raises(TypeError, match="duplicate coordinates"): + with pytest.raises(TypeError, match="unrecognized combination"): vector.MomentumSympy4D(px=_px, x=_x, y=_y, z=_z, t=_t) @@ -764,7 +764,7 @@ def test_duplicate_py_y_zip(): def test_duplicate_py_y_sympy(): """MomentumSympy4D should reject py + y""" - with pytest.raises(TypeError, match="duplicate coordinates"): + with pytest.raises(TypeError, match="unrecognized combination"): vector.MomentumSympy4D(x=_x, py=_py, y=_y, z=_z, t=_t) @@ -1452,7 +1452,7 @@ def test_mass_without_longitudinal_zip(): def test_mass_without_longitudinal_sympy(): """MomentumSympy4D should reject pt+phi+mass (temporal without longitudinal)""" - with pytest.raises(TypeError, match="unrecognized combination"): + with pytest.raises(TypeError, match="duplicate coordinates"): vector.MomentumSympy4D(pt=_pt, phi=_phi, mass=_mass) From 368f655d95e9ebbfe13d25d43ce66a20bacd471e Mon Sep 17 00:00:00 2001 From: Iason Krommydas Date: Wed, 4 Feb 2026 17:15:52 -0600 Subject: [PATCH 10/15] remove comment --- src/vector/backends/numpy.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/vector/backends/numpy.py b/src/vector/backends/numpy.py index 088d8487..29e28885 100644 --- a/src/vector/backends/numpy.py +++ b/src/vector/backends/numpy.py @@ -2160,7 +2160,6 @@ def array(*args: typing.Any, **kwargs: typing.Any) -> VectorNumpy: is_momentum = any(x in _repr_momentum_to_generic for x in names) - # Validate coordinates using dimension-guard pattern (same as awkward _check_names) _validate_coordinates(names) if any(x in ("t", "E", "e", "energy", "tau", "M", "m", "mass") for x in names): From 840da50e5c2fe6e0adf3b135c8517b5207349de5 Mon Sep 17 00:00:00 2001 From: Iason Krommydas Date: Wed, 4 Feb 2026 18:20:28 -0600 Subject: [PATCH 11/15] should check here too --- src/vector/backends/numpy.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/src/vector/backends/numpy.py b/src/vector/backends/numpy.py index 29e28885..db190cf6 100644 --- a/src/vector/backends/numpy.py +++ b/src/vector/backends/numpy.py @@ -1191,6 +1191,8 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return + _validate_coordinates(self.dtype.names or ()) + if _has(self, ("x", "y")): self._azimuthal_type = AzimuthalNumpyXY elif _has(self, ("rho", "phi")): @@ -1364,6 +1366,8 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return + _validate_coordinates(self.dtype.names or ()) + self.dtype.names = tuple( _repr_momentum_to_generic.get(x, x) for x in (self.dtype.names or ()) ) @@ -1432,6 +1436,8 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return + _validate_coordinates(self.dtype.names or ()) + if _has(self, ("x", "y")): self._azimuthal_type = AzimuthalNumpyXY elif _has(self, ("rho", "phi")): @@ -1665,6 +1671,8 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return + _validate_coordinates(self.dtype.names or ()) + self.dtype.names = tuple( _repr_momentum_to_generic.get(x, x) for x in (self.dtype.names or ()) ) @@ -1746,6 +1754,8 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return + _validate_coordinates(self.dtype.names or ()) + if _has(self, ("x", "y")): self._azimuthal_type = AzimuthalNumpyXY elif _has(self, ("rho", "phi")): @@ -2048,6 +2058,8 @@ def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return + _validate_coordinates(self.dtype.names or ()) + self.dtype.names = tuple( _repr_momentum_to_generic.get(x, x) for x in (self.dtype.names or ()) ) From f0b25eb8dab9c4e22f4c84863ea85144b661ec2d Mon Sep 17 00:00:00 2001 From: Iason Krommydas Date: Wed, 4 Feb 2026 18:53:36 -0600 Subject: [PATCH 12/15] accidentally deleted this --- src/vector/backends/object.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/vector/backends/object.py b/src/vector/backends/object.py index 2b87e1b9..dd1e040d 100644 --- a/src/vector/backends/object.py +++ b/src/vector/backends/object.py @@ -3214,6 +3214,7 @@ def obj(**coordinates: float) -> VectorObject: generic_coordinates = {} _is_type_safe(coordinates) + if "px" in coordinates: is_momentum = True generic_coordinates["x"] = coordinates.pop("px") From f1d3d94a6cd45f81600440f1f204eea1e5691220 Mon Sep 17 00:00:00 2001 From: Iason Krommydas Date: Wed, 4 Feb 2026 18:58:35 -0600 Subject: [PATCH 13/15] this is better? --- src/vector/backends/awkward_constructors.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/src/vector/backends/awkward_constructors.py b/src/vector/backends/awkward_constructors.py index fdf41cca..dc98b344 100644 --- a/src/vector/backends/awkward_constructors.py +++ b/src/vector/backends/awkward_constructors.py @@ -51,6 +51,7 @@ def _check_names( dimension = 0 names = [] columns = [] + fieldnames_orig = list(fieldnames) if "x" in fieldnames and "y" in fieldnames: if dimension != 0: @@ -202,17 +203,11 @@ def _check_names( raise TypeError(complaint1 if is_momentum else complaint2) # Check if any remaining fieldnames would conflict with already-processed coordinates - # or with each other when mapped to generic names (e.g., "x" and "px" both map to "x") + # when mapped to generic names (e.g., pt was processed, rho shouldn't remain) if fieldnames: - # Check leftovers against already-processed coordinates - for fname in fieldnames: - generic = _repr_momentum_to_generic.get(fname, fname) - if generic in names: - raise TypeError(complaint1 if is_momentum else complaint2) - - # Check leftovers against each other for duplicates - leftover_generics = [_repr_momentum_to_generic.get(x, x) for x in fieldnames] - if len(leftover_generics) != len(set(leftover_generics)): + # Map all original fieldnames to generic names to detect conflicts + generic_names = [_repr_momentum_to_generic.get(x, x) for x in fieldnames_orig] + if len(generic_names) != len(set(generic_names)): raise TypeError(complaint1 if is_momentum else complaint2) for name in fieldnames: From 75e5d54afafc19bf909fd091b9eb32d17e615a72 Mon Sep 17 00:00:00 2001 From: Iason Krommydas Date: Wed, 4 Feb 2026 20:31:17 -0600 Subject: [PATCH 14/15] claude assisted testing insanity --- tests/test_pr_659.py | 1902 ++++++++---------------------------------- 1 file changed, 330 insertions(+), 1572 deletions(-) diff --git a/tests/test_pr_659.py b/tests/test_pr_659.py index 1451d7c4..ce37e144 100644 --- a/tests/test_pr_659.py +++ b/tests/test_pr_659.py @@ -5,1592 +5,350 @@ from __future__ import annotations +import itertools + import numpy as np import pytest import vector +from vector._methods import Momentum ak = pytest.importorskip("awkward") sympy = pytest.importorskip("sympy") pytestmark = [pytest.mark.awkward, pytest.mark.sympy] -# Define sympy symbols for tests -_x, _y = sympy.symbols("x y") -_rho, _phi = sympy.symbols("rho phi") -_z, _eta, _theta = sympy.symbols("z eta theta") -_t, _tau = sympy.symbols("t tau") -_px, _py = sympy.symbols("px py") -_pt = sympy.symbols("pt") -_pz = sympy.symbols("pz") -_M, _m, _mass, _E, _e, _energy = sympy.symbols("M m mass E e energy") - - -# ============================================================================ -# Duplicate temporal coordinates (t-like vs tau-like) -# ============================================================================ -# Temporal coordinates: t, E, e, energy (all map to 't') -# tau, M, m, mass (all map to 'tau') -# These are mutually exclusive - - -def test_duplicate_E_e_object(): - """vector.obj should reject E + e""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.obj(pt=1.0, phi=0.5, eta=1.0, E=5.0, e=5.0) - - -def test_duplicate_E_e_numpy(): - """vector.array should reject E + e""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.array( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "E": np.array([5.0, 6.0]), - "e": np.array([5.0, 6.0]), - } - ) - - -def test_duplicate_E_e_awkward(): - """vector.Array should reject E + e""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.Array( - ak.Array( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "E": np.array([5.0, 6.0]), - "e": np.array([5.0, 6.0]), - } - ) - ) - - -def test_duplicate_E_e_zip(): - """vector.zip should reject E + e""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.zip( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "E": np.array([5.0, 6.0]), - "e": np.array([5.0, 6.0]), - } - ) - - -def test_duplicate_E_e_sympy(): - """MomentumSympy4D should reject E + e""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.MomentumSympy4D(pt=_pt, phi=_phi, eta=_eta, E=_E, e=_e) - - -def test_duplicate_E_energy_object(): - """vector.obj should reject E + energy""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.obj(pt=1.0, phi=0.5, eta=1.0, E=5.0, energy=5.0) - - -def test_duplicate_E_energy_numpy(): - """vector.array should reject E + energy""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.array( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "E": np.array([5.0, 6.0]), - "energy": np.array([5.0, 6.0]), - } - ) - - -def test_duplicate_E_energy_awkward(): - """vector.Array should reject E + energy""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.Array( - ak.Array( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "E": np.array([5.0, 6.0]), - "energy": np.array([5.0, 6.0]), - } - ) - ) - - -def test_duplicate_E_energy_zip(): - """vector.zip should reject E + energy""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.zip( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "E": np.array([5.0, 6.0]), - "energy": np.array([5.0, 6.0]), - } - ) - - -def test_duplicate_E_energy_sympy(): - """MomentumSympy4D should reject E + energy""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.MomentumSympy4D(pt=_pt, phi=_phi, eta=_eta, E=_E, energy=_energy) - - -def test_duplicate_e_energy_object(): - """vector.obj should reject e + energy""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.obj(pt=1.0, phi=0.5, eta=1.0, e=5.0, energy=5.0) - - -def test_duplicate_e_energy_numpy(): - """vector.array should reject e + energy""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.array( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "e": np.array([5.0, 6.0]), - "energy": np.array([5.0, 6.0]), - } - ) - - -def test_duplicate_e_energy_awkward(): - """vector.Array should reject e + energy""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.Array( - ak.Array( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "e": np.array([5.0, 6.0]), - "energy": np.array([5.0, 6.0]), - } - ) - ) - - -def test_duplicate_e_energy_zip(): - """vector.zip should reject e + energy""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.zip( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "e": np.array([5.0, 6.0]), - "energy": np.array([5.0, 6.0]), - } - ) - - -def test_duplicate_e_energy_sympy(): - """MomentumSympy4D should reject e + energy""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.MomentumSympy4D(pt=_pt, phi=_phi, eta=_eta, e=_e, energy=_energy) - - -def test_duplicate_M_m_object(): - """vector.obj should reject M + m""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.obj(pt=1.0, phi=0.5, eta=1.0, M=0.5, m=0.5) - - -def test_duplicate_M_m_numpy(): - """vector.array should reject M + m""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.array( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "M": np.array([0.5, 0.5]), - "m": np.array([0.5, 0.5]), - } - ) - - -def test_duplicate_M_m_awkward(): - """vector.Array should reject M + m""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.Array( - ak.Array( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "M": np.array([0.5, 0.5]), - "m": np.array([0.5, 0.5]), - } - ) - ) - - -def test_duplicate_M_m_zip(): - """vector.zip should reject M + m""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.zip( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "M": np.array([0.5, 0.5]), - "m": np.array([0.5, 0.5]), - } - ) - - -def test_duplicate_M_m_sympy(): - """MomentumSympy4D should reject M + m""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.MomentumSympy4D(pt=_pt, phi=_phi, eta=_eta, M=_M, m=_m) - - -def test_duplicate_M_mass_object(): - """vector.obj should reject M + mass""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.obj(pt=1.0, phi=0.5, eta=1.0, M=0.5, mass=0.5) - - -def test_duplicate_M_mass_numpy(): - """vector.array should reject M + mass""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.array( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "M": np.array([0.5, 0.5]), - "mass": np.array([0.5, 0.5]), - } - ) - - -def test_duplicate_M_mass_awkward(): - """vector.Array should reject M + mass""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.Array( - ak.Array( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "M": np.array([0.5, 0.5]), - "mass": np.array([0.5, 0.5]), - } - ) - ) - - -def test_duplicate_M_mass_zip(): - """vector.zip should reject M + mass""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.zip( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "M": np.array([0.5, 0.5]), - "mass": np.array([0.5, 0.5]), - } - ) - - -def test_duplicate_M_mass_sympy(): - """MomentumSympy4D should reject M + mass""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.MomentumSympy4D(pt=_pt, phi=_phi, eta=_eta, M=_M, mass=_mass) - - -def test_duplicate_m_mass_object(): - """vector.obj should reject m + mass""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.obj(pt=1.0, phi=0.5, eta=1.0, m=0.5, mass=0.5) - - -def test_duplicate_m_mass_numpy(): - """vector.array should reject m + mass""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.array( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "m": np.array([0.5, 0.5]), - "mass": np.array([0.5, 0.5]), - } - ) - - -def test_duplicate_m_mass_awkward(): - """vector.Array should reject m + mass""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.Array( - ak.Array( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "m": np.array([0.5, 0.5]), - "mass": np.array([0.5, 0.5]), - } - ) - ) - - -def test_duplicate_m_mass_zip(): - """vector.zip should reject m + mass""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.zip( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "m": np.array([0.5, 0.5]), - "mass": np.array([0.5, 0.5]), - } - ) - - -def test_duplicate_m_mass_sympy(): - """MomentumSympy4D should reject m + mass""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.MomentumSympy4D(pt=_pt, phi=_phi, eta=_eta, m=_m, mass=_mass) - - -def test_conflicting_energy_mass_object(): - """vector.obj should reject energy + mass (t-like + tau-like)""" - with pytest.raises(TypeError, match="specify t= or tau=, but not more than one"): - vector.obj(pt=1.0, phi=0.5, eta=1.0, energy=5.0, mass=0.5) - - -def test_conflicting_energy_mass_numpy(): - """vector.array should reject energy + mass (t-like + tau-like)""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.array( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "energy": np.array([5.0, 6.0]), - "mass": np.array([0.5, 0.5]), - } - ) - - -def test_conflicting_energy_mass_awkward(): - """vector.Array should reject energy + mass (t-like + tau-like)""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.Array( - ak.Array( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "energy": np.array([5.0, 6.0]), - "mass": np.array([0.5, 0.5]), - } - ) - ) - - -def test_conflicting_energy_mass_zip(): - """vector.zip should reject energy + mass (t-like + tau-like)""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.zip( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "energy": np.array([5.0, 6.0]), - "mass": np.array([0.5, 0.5]), - } - ) - - -def test_conflicting_energy_mass_sympy(): - """MomentumSympy4D should reject energy + mass (t-like + tau-like)""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.MomentumSympy4D(pt=_pt, phi=_phi, eta=_eta, energy=_energy, mass=_mass) - - -def test_conflicting_t_tau_object(): - """vector.obj should reject t + tau""" - with pytest.raises(TypeError, match="specify t= or tau="): - vector.obj(x=1.0, y=2.0, z=3.0, t=5.0, tau=0.5) - - -def test_conflicting_t_tau_numpy(): - """vector.array should reject t + tau""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.array( - { - "x": np.array([1.0, 2.0]), - "y": np.array([2.0, 3.0]), - "z": np.array([3.0, 4.0]), - "t": np.array([5.0, 6.0]), - "tau": np.array([0.5, 0.5]), - } - ) - - -def test_conflicting_t_tau_awkward(): - """vector.Array should reject t + tau""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.Array( - ak.Array( - { - "x": np.array([1.0, 2.0]), - "y": np.array([2.0, 3.0]), - "z": np.array([3.0, 4.0]), - "t": np.array([5.0, 6.0]), - "tau": np.array([0.5, 0.5]), - } - ) - ) - - -def test_conflicting_t_tau_zip(): - """vector.zip should reject t + tau""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.zip( - { - "x": np.array([1.0, 2.0]), - "y": np.array([2.0, 3.0]), - "z": np.array([3.0, 4.0]), - "t": np.array([5.0, 6.0]), - "tau": np.array([0.5, 0.5]), - } - ) - - -def test_conflicting_t_tau_sympy(): - """VectorSympy4D should reject t + tau""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.VectorSympy4D(x=_x, y=_y, z=_z, t=_t, tau=_tau) - - -def test_conflicting_E_mass_object(): - """vector.obj should reject E + mass""" - with pytest.raises(TypeError, match="specify t= or tau=, but not more than one"): - vector.obj(pt=1.0, phi=0.5, eta=1.0, E=5.0, mass=0.5) - - -def test_conflicting_E_mass_numpy(): - """vector.array should reject E + mass""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.array( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "E": np.array([5.0, 6.0]), - "mass": np.array([0.5, 0.5]), - } - ) - - -def test_conflicting_E_mass_awkward(): - """vector.Array should reject E + mass""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.Array( - ak.Array( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "E": np.array([5.0, 6.0]), - "mass": np.array([0.5, 0.5]), - } - ) - ) - - -def test_conflicting_E_mass_zip(): - """vector.zip should reject E + mass""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.zip( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "E": np.array([5.0, 6.0]), - "mass": np.array([0.5, 0.5]), - } - ) - - -def test_conflicting_E_mass_sympy(): - """MomentumSympy4D should reject E + mass""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.MomentumSympy4D(pt=_pt, phi=_phi, eta=_eta, E=_E, mass=_mass) - - -def test_conflicting_t_mass_object(): - """vector.obj should reject t + mass""" - with pytest.raises(TypeError, match="specify t= or tau=, but not more than one"): - vector.obj(x=1.0, y=2.0, z=3.0, t=5.0, mass=0.5) - - -def test_conflicting_t_mass_numpy(): - """vector.array should reject t + mass""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.array( - { - "x": np.array([1.0, 2.0]), - "y": np.array([2.0, 3.0]), - "z": np.array([3.0, 4.0]), - "t": np.array([5.0, 6.0]), - "mass": np.array([0.5, 0.5]), - } - ) - - -def test_conflicting_t_mass_awkward(): - """vector.Array should reject t + mass""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.Array( - ak.Array( - { - "x": np.array([1.0, 2.0]), - "y": np.array([2.0, 3.0]), - "z": np.array([3.0, 4.0]), - "t": np.array([5.0, 6.0]), - "mass": np.array([0.5, 0.5]), - } - ) - ) - - -def test_conflicting_t_mass_zip(): - """vector.zip should reject t + mass""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.zip( - { - "x": np.array([1.0, 2.0]), - "y": np.array([2.0, 3.0]), - "z": np.array([3.0, 4.0]), - "t": np.array([5.0, 6.0]), - "mass": np.array([0.5, 0.5]), - } - ) - - -def test_conflicting_t_mass_sympy(): - """MomentumSympy4D should reject t + mass""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.MomentumSympy4D(x=_x, y=_y, z=_z, t=_t, mass=_mass) - - -def test_conflicting_energy_tau_object(): - """vector.obj should reject energy + tau""" - with pytest.raises(TypeError, match="specify t= or tau=, but not more than one"): - vector.obj(x=1.0, y=2.0, z=3.0, energy=5.0, tau=0.5) - - -def test_conflicting_energy_tau_numpy(): - """vector.array should reject energy + tau""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.array( - { - "x": np.array([1.0, 2.0]), - "y": np.array([2.0, 3.0]), - "z": np.array([3.0, 4.0]), - "energy": np.array([5.0, 6.0]), - "tau": np.array([0.5, 0.5]), - } - ) - - -def test_conflicting_energy_tau_awkward(): - """vector.Array should reject energy + tau""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.Array( - ak.Array( - { - "x": np.array([1.0, 2.0]), - "y": np.array([2.0, 3.0]), - "z": np.array([3.0, 4.0]), - "energy": np.array([5.0, 6.0]), - "tau": np.array([0.5, 0.5]), - } - ) - ) - - -def test_conflicting_energy_tau_zip(): - """vector.zip should reject energy + tau""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.zip( - { - "x": np.array([1.0, 2.0]), - "y": np.array([2.0, 3.0]), - "z": np.array([3.0, 4.0]), - "energy": np.array([5.0, 6.0]), - "tau": np.array([0.5, 0.5]), - } - ) - - -def test_conflicting_energy_tau_sympy(): - """MomentumSympy4D should reject energy + tau""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.MomentumSympy4D(x=_x, y=_y, z=_z, energy=_energy, tau=_tau) - - -# ============================================================================ -# Duplicate azimuthal coordinates -# ============================================================================ -# x <-> px, y <-> py, rho <-> pt - - -def test_duplicate_px_x_object(): - """vector.obj should reject px + x""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.obj(px=1.0, x=1.0, y=2.0, z=3.0, t=5.0) - - -def test_duplicate_px_x_numpy(): - """vector.array should reject px + x""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.array( - { - "px": np.array([1.0, 2.0]), - "x": np.array([1.0, 2.0]), - "y": np.array([2.0, 3.0]), - "z": np.array([3.0, 4.0]), - "t": np.array([5.0, 6.0]), - } - ) - - -def test_duplicate_px_x_awkward(): - """vector.Array should reject px + x""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.Array( - ak.Array( - { - "px": np.array([1.0, 2.0]), - "x": np.array([1.0, 2.0]), - "y": np.array([2.0, 3.0]), - "z": np.array([3.0, 4.0]), - "t": np.array([5.0, 6.0]), - } - ) - ) - - -def test_duplicate_px_x_zip(): - """vector.zip should reject px + x""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.zip( - { - "px": np.array([1.0, 2.0]), - "x": np.array([1.0, 2.0]), - "y": np.array([2.0, 3.0]), - "z": np.array([3.0, 4.0]), - "t": np.array([5.0, 6.0]), - } - ) - - -def test_duplicate_px_x_sympy(): - """MomentumSympy4D should reject px + x""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.MomentumSympy4D(px=_px, x=_x, y=_y, z=_z, t=_t) - - -def test_duplicate_py_y_object(): - """vector.obj should reject py + y""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.obj(x=1.0, py=2.0, y=2.0, z=3.0, t=5.0) - - -def test_duplicate_py_y_numpy(): - """vector.array should reject py + y""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.array( - { - "x": np.array([1.0, 2.0]), - "py": np.array([2.0, 3.0]), - "y": np.array([2.0, 3.0]), - "z": np.array([3.0, 4.0]), - "t": np.array([5.0, 6.0]), - } - ) - - -def test_duplicate_py_y_awkward(): - """vector.Array should reject py + y""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.Array( - ak.Array( - { - "x": np.array([1.0, 2.0]), - "py": np.array([2.0, 3.0]), - "y": np.array([2.0, 3.0]), - "z": np.array([3.0, 4.0]), - "t": np.array([5.0, 6.0]), - } - ) - ) - - -def test_duplicate_py_y_zip(): - """vector.zip should reject py + y""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.zip( - { - "x": np.array([1.0, 2.0]), - "py": np.array([2.0, 3.0]), - "y": np.array([2.0, 3.0]), - "z": np.array([3.0, 4.0]), - "t": np.array([5.0, 6.0]), - } - ) - - -def test_duplicate_py_y_sympy(): - """MomentumSympy4D should reject py + y""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.MomentumSympy4D(x=_x, py=_py, y=_y, z=_z, t=_t) - - -def test_duplicate_pt_rho_object(): - """vector.obj should reject pt + rho""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.obj(pt=1.0, rho=1.0, phi=0.5, eta=1.0, mass=0.5) - - -def test_duplicate_pt_rho_numpy(): - """vector.array should reject pt + rho""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.array( - { - "pt": np.array([1.0, 2.0]), - "rho": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "mass": np.array([0.5, 0.5]), - } - ) - - -def test_duplicate_pt_rho_awkward(): - """vector.Array should reject pt + rho""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.Array( - ak.Array( - { - "pt": np.array([1.0, 2.0]), - "rho": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "mass": np.array([0.5, 0.5]), - } - ) - ) - - -def test_duplicate_pt_rho_zip(): - """vector.zip should reject pt + rho""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.zip( - { - "pt": np.array([1.0, 2.0]), - "rho": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "mass": np.array([0.5, 0.5]), - } - ) - - -def test_duplicate_pt_rho_sympy(): - """MomentumSympy4D should reject pt + rho""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.MomentumSympy4D(pt=_pt, rho=_rho, phi=_phi, eta=_eta, mass=_mass) - - -# ============================================================================ -# Duplicate longitudinal coordinates -# ============================================================================ -# z <-> pz - - -def test_duplicate_pz_z_object(): - """vector.obj should reject pz + z""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.obj(x=1.0, y=2.0, pz=3.0, z=3.0, t=5.0) - - -def test_duplicate_pz_z_numpy(): - """vector.array should reject pz + z""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.array( - { - "x": np.array([1.0, 2.0]), - "y": np.array([2.0, 3.0]), - "pz": np.array([3.0, 4.0]), - "z": np.array([3.0, 4.0]), - "t": np.array([5.0, 6.0]), - } - ) - - -def test_duplicate_pz_z_awkward(): - """vector.Array should reject pz + z""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.Array( - ak.Array( - { - "x": np.array([1.0, 2.0]), - "y": np.array([2.0, 3.0]), - "pz": np.array([3.0, 4.0]), - "z": np.array([3.0, 4.0]), - "t": np.array([5.0, 6.0]), - } - ) - ) - - -def test_duplicate_pz_z_zip(): - """vector.zip should reject pz + z""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.zip( - { - "x": np.array([1.0, 2.0]), - "y": np.array([2.0, 3.0]), - "pz": np.array([3.0, 4.0]), - "z": np.array([3.0, 4.0]), - "t": np.array([5.0, 6.0]), - } - ) - - -def test_duplicate_pz_z_sympy(): - """MomentumSympy4D should reject pz + z""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.MomentumSympy4D(x=_x, y=_y, pz=_pz, z=_z, t=_t) - - -# ============================================================================ -# Mixed azimuthal coordinate systems (from _gather_coordinates) -# ============================================================================ - - -def test_mixed_xy_with_rho_object(): - """vector.obj should reject x+y with rho""" - with pytest.raises(TypeError, match="specify x= and y= or rho= and phi="): - vector.obj(x=1.0, y=2.0, rho=1.0, z=3.0) - - -def test_mixed_xy_with_phi_object(): - """vector.obj should reject x+y with phi""" - with pytest.raises(TypeError, match="specify x= and y= or rho= and phi="): - vector.obj(x=1.0, y=2.0, phi=0.5, z=3.0) - - -def test_mixed_rhophi_with_x_object(): - """vector.obj should reject rho+phi with x""" - with pytest.raises(TypeError, match="specify x= and y= or rho= and phi="): - vector.obj(rho=1.0, phi=0.5, x=1.0, z=3.0) - - -def test_mixed_rhophi_with_y_object(): - """vector.obj should reject rho+phi with y""" - with pytest.raises(TypeError, match="specify x= and y= or rho= and phi="): - vector.obj(rho=1.0, phi=0.5, y=2.0, z=3.0) - - -# ============================================================================ -# Mixed longitudinal coordinates (from _gather_coordinates) -# ============================================================================ - - -def test_mixed_z_theta_object(): - """vector.obj should reject z with theta""" - with pytest.raises(TypeError, match="specify z= or theta= or eta="): - vector.obj(x=1.0, y=2.0, z=3.0, theta=1.0) - +ALL_COORDINATES = [ + "x", + "y", + "rho", + "phi", + "px", + "py", + "pt", + "z", + "theta", + "eta", + "pz", + "t", + "tau", + "E", + "e", + "energy", + "M", + "m", + "mass", +] -def test_mixed_z_eta_object(): - """vector.obj should reject z with eta""" - with pytest.raises(TypeError, match="specify z= or theta= or eta="): - vector.obj(x=1.0, y=2.0, z=3.0, eta=1.0) +MOMENTUM_COORDINATES = {"px", "py", "pz", "pt", "E", "e", "energy", "M", "m", "mass"} +TEMPORAL_COORDS = {"t", "tau", "E", "e", "energy", "M", "m", "mass"} +LONGITUDINAL_COORDS = {"z", "theta", "eta", "pz"} + +COORD_ALIASES = { + "px": "x", + "py": "y", + "pz": "z", + "pt": "rho", + "E": "t", + "e": "t", + "energy": "t", + "M": "tau", + "m": "tau", + "mass": "tau", +} + +VALID_2_COMBINATIONS = [ + {"x", "y"}, + {"rho", "phi"}, + {"px", "py"}, + {"pt", "phi"}, + {"x", "py"}, + {"px", "y"}, +] + +ALL_2_COMBINATIONS = list(itertools.combinations(ALL_COORDINATES, 2)) + +VALID_3_COMBINATIONS = [ + az | {lon} for az in VALID_2_COMBINATIONS for lon in ("z", "theta", "eta", "pz") +] + +ALL_3_COMBINATIONS = list(itertools.combinations(ALL_COORDINATES, 3)) + +VALID_4_COMBINATIONS = [ + az | {lon} | {temp} + for az in VALID_2_COMBINATIONS + for lon in ("z", "theta", "eta", "pz") + for temp in ("t", "tau", "E", "e", "energy", "M", "m", "mass") +] + +ALL_4_COMBINATIONS = list(itertools.combinations(ALL_COORDINATES, 4)) + + +def _is_valid_2(combo): + return set(combo) in VALID_2_COMBINATIONS + + +def _is_valid_3(combo): + return set(combo) in VALID_3_COMBINATIONS + + +def _is_valid_4(combo): + return set(combo) in VALID_4_COMBINATIONS + + +def _has_valid_3_subset(combo): + for triple in itertools.combinations(combo, 3): + if set(triple) in VALID_3_COMBINATIONS: + return True + return False + + +def _has_valid_2_subset(combo): + for pair in itertools.combinations(combo, 2): + if set(pair) in VALID_2_COMBINATIONS: + return True + return False + + +def _is_momentum(combo): + return any(c in MOMENTUM_COORDINATES for c in combo) + + +def _to_canonical(coord): + return COORD_ALIASES.get(coord, coord) + + +def _has_duplicate(combo): + canonicals = [_to_canonical(c) for c in combo] + return len(canonicals) != len(set(canonicals)) + + +def _will_error_for_non_obj(combo): + """Check if combo with valid subset will error for non-obj backends.""" + # Check for canonical duplicates (e.g., x and px both map to x) + if _has_duplicate(combo): + return True -def test_mixed_theta_eta_object(): - """vector.obj should reject theta with eta""" - with pytest.raises(TypeError, match="specify z= or theta= or eta="): - vector.obj(x=1.0, y=2.0, theta=1.0, eta=1.0) - - -# ============================================================================ -# Valid combinations (ensure validation doesn't reject valid inputs) -# ============================================================================ + # Temporal without longitudinal + has_temporal = any(c in TEMPORAL_COORDS for c in combo) + has_longitudinal = any(c in LONGITUDINAL_COORDS for c in combo) + if has_temporal and not has_longitudinal: + return True + # Multiple longitudinal coords (z, theta, eta, pz) - only one allowed + longitudinal_count = sum(1 for c in combo if c in LONGITUDINAL_COORDS) + if longitudinal_count > 1: + return True -def test_valid_pt_phi_eta_mass_object(): - """vector.obj should accept pt, phi, eta, mass""" - vec = vector.obj(pt=1.0, phi=0.5, eta=1.0, mass=0.5) - assert vec.pt == 1.0 - assert vec.phi == 0.5 - assert vec.eta == 1.0 - assert vec.mass == 0.5 + # Multiple temporal coords - only one allowed + temporal_count = sum(1 for c in combo if c in TEMPORAL_COORDS) + if temporal_count > 1: + return True - -def test_valid_pt_phi_eta_mass_numpy(): - """vector.array should accept pt, phi, eta, mass""" - arr = vector.array( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "mass": np.array([0.5, 0.5]), - } - ) - assert np.allclose(arr.pt, [1.0, 2.0]) - assert np.allclose(arr.phi, [0.5, 1.0]) - assert np.allclose(arr.eta, [1.0, 1.5]) - assert np.allclose(arr.mass, [0.5, 0.5]) - - -def test_valid_pt_phi_eta_mass_awkward(): - """vector.Array should accept pt, phi, eta, mass""" - arr = vector.Array( - ak.Array( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "mass": np.array([0.5, 0.5]), - } - ) - ) - assert ak.all(arr.pt == ak.Array([1.0, 2.0])) - assert ak.all(arr.phi == ak.Array([0.5, 1.0])) - assert ak.all(arr.eta == ak.Array([1.0, 1.5])) - assert ak.all(arr.mass == ak.Array([0.5, 0.5])) - - -def test_valid_pt_phi_eta_mass_zip(): - """vector.zip should accept pt, phi, eta, mass""" - arr = vector.zip( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "eta": np.array([1.0, 1.5]), - "mass": np.array([0.5, 0.5]), - } - ) - assert ak.all(arr.pt == ak.Array([1.0, 2.0])) - assert ak.all(arr.phi == ak.Array([0.5, 1.0])) - assert ak.all(arr.eta == ak.Array([1.0, 1.5])) - assert ak.all(arr.mass == ak.Array([0.5, 0.5])) - - -def test_valid_pt_phi_eta_mass_sympy(): - """MomentumSympy4D should accept pt, phi, eta, mass""" - vec = vector.MomentumSympy4D(pt=_pt, phi=_phi, eta=_eta, mass=_mass) - assert vec.pt == _pt - assert vec.phi == _phi - assert vec.eta == _eta - assert vec.mass == _mass - - -def test_valid_x_y_z_energy_object(): - """vector.obj should accept x, y, z, energy""" - vec = vector.obj(x=1.0, y=2.0, z=3.0, energy=5.0) - assert vec.x == 1.0 - assert vec.y == 2.0 - assert vec.z == 3.0 - assert vec.energy == 5.0 - - -def test_valid_x_y_z_energy_numpy(): - """vector.array should accept x, y, z, energy""" - arr = vector.array( - { - "x": np.array([1.0, 2.0]), - "y": np.array([2.0, 3.0]), - "z": np.array([3.0, 4.0]), - "energy": np.array([5.0, 6.0]), - } - ) - assert np.allclose(arr.x, [1.0, 2.0]) - assert np.allclose(arr.y, [2.0, 3.0]) - assert np.allclose(arr.z, [3.0, 4.0]) - assert np.allclose(arr.energy, [5.0, 6.0]) - - -def test_valid_x_y_z_energy_awkward(): - """vector.Array should accept x, y, z, energy""" - arr = vector.Array( - ak.Array( - { - "x": np.array([1.0, 2.0]), - "y": np.array([2.0, 3.0]), - "z": np.array([3.0, 4.0]), - "energy": np.array([5.0, 6.0]), - } - ) - ) - assert ak.all(arr.x == ak.Array([1.0, 2.0])) - assert ak.all(arr.y == ak.Array([2.0, 3.0])) - assert ak.all(arr.z == ak.Array([3.0, 4.0])) - assert ak.all(arr.energy == ak.Array([5.0, 6.0])) - - -def test_valid_x_y_z_energy_zip(): - """vector.zip should accept x, y, z, energy""" - arr = vector.zip( - { - "x": np.array([1.0, 2.0]), - "y": np.array([2.0, 3.0]), - "z": np.array([3.0, 4.0]), - "energy": np.array([5.0, 6.0]), - } - ) - assert ak.all(arr.x == ak.Array([1.0, 2.0])) - assert ak.all(arr.y == ak.Array([2.0, 3.0])) - assert ak.all(arr.z == ak.Array([3.0, 4.0])) - assert ak.all(arr.energy == ak.Array([5.0, 6.0])) - - -def test_valid_x_y_z_energy_sympy(): - """MomentumSympy4D should accept x, y, z, energy""" - vec = vector.MomentumSympy4D(x=_x, y=_y, z=_z, energy=_energy) - assert vec.x == _x - assert vec.y == _y - assert vec.z == _z - assert vec.energy == _energy - - -def test_valid_px_py_pz_E_object(): - """vector.obj should accept px, py, pz, E""" - vec = vector.obj(px=1.0, py=2.0, pz=3.0, E=5.0) - assert vec.px == 1.0 - assert vec.py == 2.0 - assert vec.pz == 3.0 - assert vec.E == 5.0 - - -def test_valid_px_py_pz_E_numpy(): - """vector.array should accept px, py, pz, E""" - arr = vector.array( - { - "px": np.array([1.0, 2.0]), - "py": np.array([2.0, 3.0]), - "pz": np.array([3.0, 4.0]), - "E": np.array([5.0, 6.0]), - } - ) - assert np.allclose(arr.px, [1.0, 2.0]) - assert np.allclose(arr.py, [2.0, 3.0]) - assert np.allclose(arr.pz, [3.0, 4.0]) - assert np.allclose(arr.E, [5.0, 6.0]) - - -def test_valid_px_py_pz_E_awkward(): - """vector.Array should accept px, py, pz, E""" - arr = vector.Array( - ak.Array( - { - "px": np.array([1.0, 2.0]), - "py": np.array([2.0, 3.0]), - "pz": np.array([3.0, 4.0]), - "E": np.array([5.0, 6.0]), - } - ) - ) - assert ak.all(arr.px == ak.Array([1.0, 2.0])) - assert ak.all(arr.py == ak.Array([2.0, 3.0])) - assert ak.all(arr.pz == ak.Array([3.0, 4.0])) - assert ak.all(ak.Array([5.0, 6.0]) == arr.E) - - -def test_valid_px_py_pz_E_zip(): - """vector.zip should accept px, py, pz, E""" - arr = vector.zip( - { - "px": np.array([1.0, 2.0]), - "py": np.array([2.0, 3.0]), - "pz": np.array([3.0, 4.0]), - "E": np.array([5.0, 6.0]), - } + # Multiple azimuthal pairs (e.g., x,y and rho,phi both present) + valid_2_count = sum( + 1 + for pair in itertools.combinations(combo, 2) + if set(pair) in VALID_2_COMBINATIONS ) - assert ak.all(arr.px == ak.Array([1.0, 2.0])) - assert ak.all(arr.py == ak.Array([2.0, 3.0])) - assert ak.all(arr.pz == ak.Array([3.0, 4.0])) - assert ak.all(ak.Array([5.0, 6.0]) == arr.E) - - -def test_valid_px_py_pz_E_sympy(): - """MomentumSympy4D should accept px, py, pz, E""" - vec = vector.MomentumSympy4D(px=_px, py=_py, pz=_pz, E=_E) - assert vec.px == _px - assert vec.py == _py - assert vec.pz == _pz - assert vec.E == _E - - -# ============================================================================ -# Incomplete azimuthal coordinate pairs -# ============================================================================ - - -def test_incomplete_x_without_y_object(): - """vector.obj should reject x without y""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.obj(x=1.0, z=3.0) - - -def test_incomplete_x_without_y_numpy(): - """vector.array should reject x without y""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.array( - { - "x": np.array([1.0, 2.0]), - "z": np.array([3.0, 4.0]), - } - ) - - -def test_incomplete_x_without_y_awkward(): - """vector.Array should reject x without y""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.Array( - ak.Array( - { - "x": np.array([1.0, 2.0]), - "z": np.array([3.0, 4.0]), - } - ) - ) - - -def test_incomplete_x_without_y_zip(): - """vector.zip should reject x without y""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.zip( - { - "x": np.array([1.0, 2.0]), - "z": np.array([3.0, 4.0]), - } - ) - - -def test_incomplete_x_without_y_sympy(): - """VectorSympy3D should reject x without y""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.VectorSympy3D(x=_x, z=_z) - - -def test_incomplete_rho_without_phi_object(): - """vector.obj should reject rho without phi""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.obj(rho=1.0, z=3.0) - - -def test_incomplete_rho_without_phi_numpy(): - """vector.array should reject rho without phi""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.array( - { - "rho": np.array([1.0, 2.0]), - "z": np.array([3.0, 4.0]), - } - ) - - -def test_incomplete_rho_without_phi_awkward(): - """vector.Array should reject rho without phi""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.Array( - ak.Array( - { - "rho": np.array([1.0, 2.0]), - "z": np.array([3.0, 4.0]), - } - ) - ) - - -def test_incomplete_rho_without_phi_zip(): - """vector.zip should reject rho without phi""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.zip( - { - "rho": np.array([1.0, 2.0]), - "z": np.array([3.0, 4.0]), - } - ) - - -def test_incomplete_rho_without_phi_sympy(): - """VectorSympy3D should reject rho without phi""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.VectorSympy3D(rho=_rho, z=_z) - - -# ============================================================================ -# Mixed azimuthal coordinate components -# ============================================================================ - - -def test_mixed_x_phi_object(): - """vector.obj should reject x with phi (mixed systems)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.obj(x=1.0, phi=0.5, z=3.0) - - -def test_mixed_x_phi_numpy(): - """vector.array should reject x with phi (mixed systems)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.array( - { - "x": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "z": np.array([3.0, 4.0]), - } - ) - - -def test_mixed_x_phi_awkward(): - """vector.Array should reject x with phi (mixed systems)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.Array( - ak.Array( - { - "x": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "z": np.array([3.0, 4.0]), - } - ) - ) - - -def test_mixed_x_phi_zip(): - """vector.zip should reject x with phi (mixed systems)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.zip( - { - "x": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "z": np.array([3.0, 4.0]), - } - ) - - -def test_mixed_x_phi_sympy(): - """VectorSympy3D should reject x with phi (mixed systems)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.VectorSympy3D(x=_x, phi=_phi, z=_z) - - -def test_mixed_y_rho_object(): - """vector.obj should reject y with rho (mixed systems)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.obj(y=2.0, rho=1.0, z=3.0) - - -def test_mixed_y_rho_numpy(): - """vector.array should reject y with rho (mixed systems)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.array( - { - "y": np.array([2.0, 3.0]), - "rho": np.array([1.0, 2.0]), - "z": np.array([3.0, 4.0]), - } - ) - - -def test_mixed_y_rho_awkward(): - """vector.Array should reject y with rho (mixed systems)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.Array( - ak.Array( - { - "y": np.array([2.0, 3.0]), - "rho": np.array([1.0, 2.0]), - "z": np.array([3.0, 4.0]), - } - ) - ) - - -def test_mixed_y_rho_zip(): - """vector.zip should reject y with rho (mixed systems)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.zip( - { - "y": np.array([2.0, 3.0]), - "rho": np.array([1.0, 2.0]), - "z": np.array([3.0, 4.0]), - } - ) - - -def test_mixed_y_rho_sympy(): - """VectorSympy3D should reject y with rho (mixed systems)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.VectorSympy3D(y=_y, rho=_rho, z=_z) - - -# ============================================================================ -# Temporal without proper 3D base -# ============================================================================ - - -def test_temporal_without_longitudinal_object(): - """vector.obj should reject x+y+t (temporal without longitudinal)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.obj(x=1.0, y=2.0, t=5.0) - - -def test_temporal_without_longitudinal_numpy(): - """vector.array should reject x+y+t (temporal without longitudinal)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.array( - { - "x": np.array([1.0, 2.0]), - "y": np.array([2.0, 3.0]), - "t": np.array([5.0, 6.0]), - } - ) - - -def test_temporal_without_longitudinal_awkward(): - """vector.Array should reject x+y+t (temporal without longitudinal)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.Array( - ak.Array( - { - "x": np.array([1.0, 2.0]), - "y": np.array([2.0, 3.0]), - "t": np.array([5.0, 6.0]), - } - ) - ) - - -def test_temporal_without_longitudinal_zip(): - """vector.zip should reject x+y+t (temporal without longitudinal)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.zip( - { - "x": np.array([1.0, 2.0]), - "y": np.array([2.0, 3.0]), - "t": np.array([5.0, 6.0]), - } - ) - - -def test_temporal_without_longitudinal_sympy(): - """VectorSympy4D should reject x+y+t (temporal without longitudinal)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.VectorSympy4D(x=_x, y=_y, t=_t) - - -def test_mass_without_longitudinal_object(): - """vector.obj should reject pt+phi+mass (temporal without longitudinal)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.obj(pt=1.0, phi=0.5, mass=0.5) - - -def test_mass_without_longitudinal_numpy(): - """vector.array should reject pt+phi+mass (temporal without longitudinal)""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.array( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "mass": np.array([0.5, 0.5]), - } - ) - - -def test_mass_without_longitudinal_awkward(): - """vector.Array should reject pt+phi+mass (temporal without longitudinal)""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.Array( - ak.Array( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "mass": np.array([0.5, 0.5]), - } - ) - ) - - -def test_mass_without_longitudinal_zip(): - """vector.zip should reject pt+phi+mass (temporal without longitudinal)""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.zip( - { - "pt": np.array([1.0, 2.0]), - "phi": np.array([0.5, 1.0]), - "mass": np.array([0.5, 0.5]), - } - ) - - -def test_mass_without_longitudinal_sympy(): - """MomentumSympy4D should reject pt+phi+mass (temporal without longitudinal)""" - with pytest.raises(TypeError, match="duplicate coordinates"): - vector.MomentumSympy4D(pt=_pt, phi=_phi, mass=_mass) - - -# ============================================================================ -# Missing required coordinates -# ============================================================================ - - -def test_only_temporal_object(): - """vector.obj should reject only t (missing spatial coords)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.obj(t=5.0) - - -def test_only_temporal_numpy(): - """vector.array should reject only t (missing spatial coords)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.array( - { - "t": np.array([5.0, 6.0]), - } - ) - - -def test_only_temporal_awkward(): - """vector.Array should reject only t (missing spatial coords)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.Array( - ak.Array( - { - "t": np.array([5.0, 6.0]), - } - ) - ) - - -def test_only_temporal_zip(): - """vector.zip should reject only t (missing spatial coords)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.zip( - { - "t": np.array([5.0, 6.0]), - } - ) - - -def test_only_temporal_sympy(): - """VectorSympy4D should reject only t (missing spatial coords)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.VectorSympy4D(t=_t) - - -def test_only_longitudinal_object(): - """vector.obj should reject only z (missing azimuthal coords)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.obj(z=3.0) - - -def test_only_longitudinal_numpy(): - """vector.array should reject only z (missing azimuthal coords)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.array( - { - "z": np.array([3.0, 4.0]), - } - ) - - -def test_only_longitudinal_awkward(): - """vector.Array should reject only z (missing azimuthal coords)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.Array( - ak.Array( - { - "z": np.array([3.0, 4.0]), - } - ) - ) - - -def test_only_longitudinal_zip(): - """vector.zip should reject only z (missing azimuthal coords)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.zip( - { - "z": np.array([3.0, 4.0]), - } - ) - - -def test_only_longitudinal_sympy(): - """VectorSympy3D should reject only z (missing azimuthal coords)""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.VectorSympy3D(z=_z) - - -def test_longitudinal_temporal_without_azimuthal_object(): - """vector.obj should reject z+t without azimuthal coords""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.obj(z=3.0, t=5.0) - - -def test_longitudinal_temporal_without_azimuthal_numpy(): - """vector.array should reject z+t without azimuthal coords""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.array( - { - "z": np.array([3.0, 4.0]), - "t": np.array([5.0, 6.0]), - } - ) - - -def test_longitudinal_temporal_without_azimuthal_awkward(): - """vector.Array should reject z+t without azimuthal coords""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.Array( - ak.Array( - { - "z": np.array([3.0, 4.0]), - "t": np.array([5.0, 6.0]), - } - ) - ) - - -def test_longitudinal_temporal_without_azimuthal_zip(): - """vector.zip should reject z+t without azimuthal coords""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.zip( - { - "z": np.array([3.0, 4.0]), - "t": np.array([5.0, 6.0]), - } - ) - - -def test_longitudinal_temporal_without_azimuthal_sympy(): - """VectorSympy4D should reject z+t without azimuthal coords""" - with pytest.raises(TypeError, match="unrecognized combination"): - vector.VectorSympy4D(z=_z, t=_t) + return valid_2_count > 1 + + +def _get_first_valid_2_subset(combo): + """Get the first valid 2-subset from a combo.""" + for pair in itertools.combinations(combo, 2): + if set(pair) in VALID_2_COMBINATIONS: + return set(pair) + return None + + +def _get_first_valid_3_subset(combo): + """Get the first valid 3-subset from a combo.""" + for triple in itertools.combinations(combo, 3): + if set(triple) in VALID_3_COMBINATIONS: + return set(triple) + return None + + +def _is_momentum_numpy(combo): + """Numpy checks if ANY field name is a momentum coordinate.""" + return any(c in MOMENTUM_COORDINATES for c in combo) + + +def _is_momentum_awkward(combo): + """Awkward only sets is_momentum when momentum coords are consumed as vector coords.""" + # Check for valid 4-subset first + for quad in itertools.combinations(combo, 4): + if set(quad) in VALID_4_COMBINATIONS: + return _is_momentum(quad) + # Check for valid 3-subset + valid_3 = _get_first_valid_3_subset(combo) + if valid_3 is not None: + return _is_momentum(valid_3) + # Check for valid 2-subset + valid_2 = _get_first_valid_2_subset(combo) + if valid_2 is not None: + return _is_momentum(valid_2) + return False + + +def _get_sympy_class(coords): + has_momentum = _is_momentum(coords) + n = len(coords) + if n <= 2: + return vector.MomentumSympy2D if has_momentum else vector.VectorSympy2D + elif n == 3: + return vector.MomentumSympy3D if has_momentum else vector.VectorSympy3D + else: + return vector.MomentumSympy4D if has_momentum else vector.VectorSympy4D + + +@pytest.mark.parametrize( + "combo", + ALL_2_COMBINATIONS, + ids=[f"{a}_{b}" for a, b in ALL_2_COMBINATIONS], +) +def test_2_combinations(combo): + is_valid = _is_valid_2(combo) + is_momentum = _is_momentum(combo) + error_pattern = "duplicate coordinates|unrecognized combination|must have a structured dtype|specify" + + if is_valid: + v_obj = vector.obj(**dict.fromkeys(combo, 1.0)) + v_numpy = vector.array({c: np.array([1.0, 2.0]) for c in combo}) + v_awkward = vector.Array(ak.Array({c: [1.0, 2.0] for c in combo})) + v_zip = vector.zip({c: np.array([1.0, 2.0]) for c in combo}) + v_sympy = _get_sympy_class(combo)(**{c: sympy.Symbol(c) for c in combo}) + + assert isinstance(v_obj, Momentum) == is_momentum + assert isinstance(v_numpy, Momentum) == is_momentum + assert isinstance(v_awkward, Momentum) == is_momentum + assert isinstance(v_zip, Momentum) == is_momentum + assert isinstance(v_sympy, Momentum) == is_momentum + else: + with pytest.raises(TypeError, match=error_pattern): + vector.obj(**dict.fromkeys(combo, 1.0)) + + with pytest.raises(TypeError, match=error_pattern): + vector.array({c: np.array([1.0, 2.0]) for c in combo}) + + with pytest.raises(TypeError, match=error_pattern): + vector.Array(ak.Array({c: [1.0, 2.0] for c in combo})) + + with pytest.raises(TypeError, match=error_pattern): + vector.zip({c: np.array([1.0, 2.0]) for c in combo}) + + with pytest.raises(TypeError, match=error_pattern): + _get_sympy_class(combo)(**{c: sympy.Symbol(c) for c in combo}) + + +@pytest.mark.parametrize( + "combo", + ALL_3_COMBINATIONS, + ids=[f"{a}_{b}_{c}" for a, b, c in ALL_3_COMBINATIONS], +) +def test_3_combinations(combo): + is_valid = _is_valid_3(combo) + has_valid_2 = _has_valid_2_subset(combo) + is_momentum = _is_momentum(combo) + error_pattern = "duplicate coordinates|unrecognized combination|must have a structured dtype|specify" + + if is_valid: + v_obj = vector.obj(**dict.fromkeys(combo, 1.0)) + v_numpy = vector.array({c: np.array([1.0, 2.0]) for c in combo}) + v_awkward = vector.Array(ak.Array({c: [1.0, 2.0] for c in combo})) + v_zip = vector.zip({c: np.array([1.0, 2.0]) for c in combo}) + v_sympy = _get_sympy_class(combo)(**{c: sympy.Symbol(c) for c in combo}) + + assert isinstance(v_obj, Momentum) == is_momentum + assert isinstance(v_numpy, Momentum) == is_momentum + assert isinstance(v_awkward, Momentum) == is_momentum + assert isinstance(v_zip, Momentum) == is_momentum + assert isinstance(v_sympy, Momentum) == is_momentum + else: + with pytest.raises(TypeError, match=error_pattern): + vector.obj(**dict.fromkeys(combo, 1.0)) + + with pytest.raises(TypeError, match=error_pattern): + _get_sympy_class(combo)(**{c: sympy.Symbol(c) for c in combo}) + + if has_valid_2 and not _will_error_for_non_obj(combo): + # numpy/awkward/zip create a 2D vector with extra fields + v_numpy = vector.array({c: np.array([1.0, 2.0]) for c in combo}) + v_awkward = vector.Array(ak.Array({c: [1.0, 2.0] for c in combo})) + v_zip = vector.zip({c: np.array([1.0, 2.0]) for c in combo}) + + assert isinstance(v_numpy, Momentum) == _is_momentum_numpy(combo) + assert isinstance(v_awkward, Momentum) == _is_momentum_awkward(combo) + assert isinstance(v_zip, Momentum) == _is_momentum_awkward(combo) + else: + with pytest.raises(TypeError, match=error_pattern): + vector.array({c: np.array([1.0, 2.0]) for c in combo}) + + with pytest.raises(TypeError, match=error_pattern): + vector.Array(ak.Array({c: [1.0, 2.0] for c in combo})) + + with pytest.raises(TypeError, match=error_pattern): + vector.zip({c: np.array([1.0, 2.0]) for c in combo}) + + +@pytest.mark.parametrize( + "combo", + ALL_4_COMBINATIONS, + ids=[f"{a}_{b}_{c}_{d}" for a, b, c, d in ALL_4_COMBINATIONS], +) +def test_4_combinations(combo): + is_valid = _is_valid_4(combo) + has_valid_3 = _has_valid_3_subset(combo) + has_valid_2 = _has_valid_2_subset(combo) + is_momentum = _is_momentum(combo) + error_pattern = "duplicate coordinates|unrecognized combination|must have a structured dtype|specify" + + if is_valid: + v_obj = vector.obj(**dict.fromkeys(combo, 1.0)) + v_numpy = vector.array({c: np.array([1.0, 2.0]) for c in combo}) + v_awkward = vector.Array(ak.Array({c: [1.0, 2.0] for c in combo})) + v_zip = vector.zip({c: np.array([1.0, 2.0]) for c in combo}) + v_sympy = _get_sympy_class(combo)(**{c: sympy.Symbol(c) for c in combo}) + + assert isinstance(v_obj, Momentum) == is_momentum + assert isinstance(v_numpy, Momentum) == is_momentum + assert isinstance(v_awkward, Momentum) == is_momentum + assert isinstance(v_zip, Momentum) == is_momentum + assert isinstance(v_sympy, Momentum) == is_momentum + else: + # obj and sympy are strict - always error for invalid combos + with pytest.raises(TypeError, match=error_pattern): + vector.obj(**dict.fromkeys(combo, 1.0)) + + with pytest.raises(TypeError, match=error_pattern): + _get_sympy_class(combo)(**{c: sympy.Symbol(c) for c in combo}) + + if has_valid_3 and not _will_error_for_non_obj(combo): + # numpy/awkward/zip create a 3D vector with extra fields + v_numpy = vector.array({c: np.array([1.0, 2.0]) for c in combo}) + v_awkward = vector.Array(ak.Array({c: [1.0, 2.0] for c in combo})) + v_zip = vector.zip({c: np.array([1.0, 2.0]) for c in combo}) + + assert isinstance(v_numpy, Momentum) == _is_momentum_numpy(combo) + assert isinstance(v_awkward, Momentum) == _is_momentum_awkward(combo) + assert isinstance(v_zip, Momentum) == _is_momentum_awkward(combo) + elif has_valid_2 and not _will_error_for_non_obj(combo): + # numpy/awkward/zip create a 2D vector with extra fields + v_numpy = vector.array({c: np.array([1.0, 2.0]) for c in combo}) + v_awkward = vector.Array(ak.Array({c: [1.0, 2.0] for c in combo})) + v_zip = vector.zip({c: np.array([1.0, 2.0]) for c in combo}) + + assert isinstance(v_numpy, Momentum) == _is_momentum_numpy(combo) + assert isinstance(v_awkward, Momentum) == _is_momentum_awkward(combo) + assert isinstance(v_zip, Momentum) == _is_momentum_awkward(combo) + else: + with pytest.raises(TypeError, match=error_pattern): + vector.array({c: np.array([1.0, 2.0]) for c in combo}) + + with pytest.raises(TypeError, match=error_pattern): + vector.Array(ak.Array({c: [1.0, 2.0] for c in combo})) + + with pytest.raises(TypeError, match=error_pattern): + vector.zip({c: np.array([1.0, 2.0]) for c in combo}) From 84af22ad857fcf860c9b360dc9b181af5465fde4 Mon Sep 17 00:00:00 2001 From: Iason Krommydas Date: Wed, 4 Feb 2026 21:19:42 -0600 Subject: [PATCH 15/15] add numba testing now too, probably will revert as it is painfully slow --- tests/test_pr_659.py | 53 ++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 51 insertions(+), 2 deletions(-) diff --git a/tests/test_pr_659.py b/tests/test_pr_659.py index ce37e144..763e6e6c 100644 --- a/tests/test_pr_659.py +++ b/tests/test_pr_659.py @@ -15,8 +15,10 @@ ak = pytest.importorskip("awkward") sympy = pytest.importorskip("sympy") +numba = pytest.importorskip("numba") + +pytestmark = [pytest.mark.awkward, pytest.mark.sympy, pytest.mark.numba] -pytestmark = [pytest.mark.awkward, pytest.mark.sympy] ALL_COORDINATES = [ "x", @@ -41,7 +43,7 @@ ] MOMENTUM_COORDINATES = {"px", "py", "pz", "pt", "E", "e", "energy", "M", "m", "mass"} - +AZIMUTHAL_COORDS = {"x", "y", "rho", "phi", "px", "py", "pt"} TEMPORAL_COORDS = {"t", "tau", "E", "e", "energy", "M", "m", "mass"} LONGITUDINAL_COORDS = {"z", "theta", "eta", "pz"} @@ -204,6 +206,25 @@ def _get_sympy_class(coords): return vector.MomentumSympy4D if has_momentum else vector.VectorSympy4D +def _numba_obj(combo): + """Create vector.obj inside a jitted function with the given coordinates.""" + kwargs = ", ".join(f"{c}=1.0" for c in combo) + local_ns = {"vector": vector, "numba": numba} + exec(f"@numba.njit\ndef f():\n return vector.obj({kwargs})", local_ns) + return local_ns["f"] + + +def _will_numba_error(combo): + """Check if numba will error. Numba errors on duplicates, no valid azimuthal, or extra azimuthal coords.""" + if _has_duplicate(combo): + return True + if not _has_valid_2_subset(combo): + return True + # Numba errors if there are more than 2 azimuthal coordinates (canonical form) + canonical_azimuthal = {_to_canonical(c) for c in combo if c in AZIMUTHAL_COORDS} + return len(canonical_azimuthal) > 2 + + @pytest.mark.parametrize( "combo", ALL_2_COMBINATIONS, @@ -214,14 +235,18 @@ def test_2_combinations(combo): is_momentum = _is_momentum(combo) error_pattern = "duplicate coordinates|unrecognized combination|must have a structured dtype|specify" + numba_error_pattern = "duplicate coordinates|unrecognized combination" + if is_valid: v_obj = vector.obj(**dict.fromkeys(combo, 1.0)) + v_numba = _numba_obj(combo)() v_numpy = vector.array({c: np.array([1.0, 2.0]) for c in combo}) v_awkward = vector.Array(ak.Array({c: [1.0, 2.0] for c in combo})) v_zip = vector.zip({c: np.array([1.0, 2.0]) for c in combo}) v_sympy = _get_sympy_class(combo)(**{c: sympy.Symbol(c) for c in combo}) assert isinstance(v_obj, Momentum) == is_momentum + assert isinstance(v_numba, Momentum) == is_momentum assert isinstance(v_numpy, Momentum) == is_momentum assert isinstance(v_awkward, Momentum) == is_momentum assert isinstance(v_zip, Momentum) == is_momentum @@ -230,6 +255,9 @@ def test_2_combinations(combo): with pytest.raises(TypeError, match=error_pattern): vector.obj(**dict.fromkeys(combo, 1.0)) + with pytest.raises(numba.TypingError, match=numba_error_pattern): + _numba_obj(combo)() + with pytest.raises(TypeError, match=error_pattern): vector.array({c: np.array([1.0, 2.0]) for c in combo}) @@ -253,26 +281,37 @@ def test_3_combinations(combo): has_valid_2 = _has_valid_2_subset(combo) is_momentum = _is_momentum(combo) error_pattern = "duplicate coordinates|unrecognized combination|must have a structured dtype|specify" + numba_error_pattern = "duplicate coordinates|unrecognized combination" if is_valid: v_obj = vector.obj(**dict.fromkeys(combo, 1.0)) + v_numba = _numba_obj(combo)() v_numpy = vector.array({c: np.array([1.0, 2.0]) for c in combo}) v_awkward = vector.Array(ak.Array({c: [1.0, 2.0] for c in combo})) v_zip = vector.zip({c: np.array([1.0, 2.0]) for c in combo}) v_sympy = _get_sympy_class(combo)(**{c: sympy.Symbol(c) for c in combo}) assert isinstance(v_obj, Momentum) == is_momentum + assert isinstance(v_numba, Momentum) == is_momentum assert isinstance(v_numpy, Momentum) == is_momentum assert isinstance(v_awkward, Momentum) == is_momentum assert isinstance(v_zip, Momentum) == is_momentum assert isinstance(v_sympy, Momentum) == is_momentum else: + # obj and sympy are strict - always error for invalid combos with pytest.raises(TypeError, match=error_pattern): vector.obj(**dict.fromkeys(combo, 1.0)) with pytest.raises(TypeError, match=error_pattern): _get_sympy_class(combo)(**{c: sympy.Symbol(c) for c in combo}) + # numba is permissive like numpy/awkward/zip + if _will_numba_error(combo): + with pytest.raises(numba.TypingError, match=numba_error_pattern): + _numba_obj(combo)() + else: + _numba_obj(combo)() + if has_valid_2 and not _will_error_for_non_obj(combo): # numpy/awkward/zip create a 2D vector with extra fields v_numpy = vector.array({c: np.array([1.0, 2.0]) for c in combo}) @@ -304,15 +343,18 @@ def test_4_combinations(combo): has_valid_2 = _has_valid_2_subset(combo) is_momentum = _is_momentum(combo) error_pattern = "duplicate coordinates|unrecognized combination|must have a structured dtype|specify" + numba_error_pattern = "duplicate coordinates|unrecognized combination" if is_valid: v_obj = vector.obj(**dict.fromkeys(combo, 1.0)) + v_numba = _numba_obj(combo)() v_numpy = vector.array({c: np.array([1.0, 2.0]) for c in combo}) v_awkward = vector.Array(ak.Array({c: [1.0, 2.0] for c in combo})) v_zip = vector.zip({c: np.array([1.0, 2.0]) for c in combo}) v_sympy = _get_sympy_class(combo)(**{c: sympy.Symbol(c) for c in combo}) assert isinstance(v_obj, Momentum) == is_momentum + assert isinstance(v_numba, Momentum) == is_momentum assert isinstance(v_numpy, Momentum) == is_momentum assert isinstance(v_awkward, Momentum) == is_momentum assert isinstance(v_zip, Momentum) == is_momentum @@ -325,6 +367,13 @@ def test_4_combinations(combo): with pytest.raises(TypeError, match=error_pattern): _get_sympy_class(combo)(**{c: sympy.Symbol(c) for c in combo}) + # numba is permissive like numpy/awkward/zip + if _will_numba_error(combo): + with pytest.raises(numba.TypingError, match=numba_error_pattern): + _numba_obj(combo)() + else: + _numba_obj(combo)() + if has_valid_3 and not _will_error_for_non_obj(combo): # numpy/awkward/zip create a 3D vector with extra fields v_numpy = vector.array({c: np.array([1.0, 2.0]) for c in combo})