diff --git a/src/vector/_methods.py b/src/vector/_methods.py index 4d278854..1450ae69 100644 --- a/src/vector/_methods.py +++ b/src/vector/_methods.py @@ -4400,7 +4400,7 @@ def _lib_of(*objects: VectorProtocol) -> Module: # NumPy-like module if isinstance(obj, Vector): if lib is None: lib = obj.lib - elif lib is not obj.lib: + elif lib != obj.lib: raise TypeError( f"cannot use {lib} and {obj.lib} in the same calculation" ) diff --git a/src/vector/backends/awkward.py b/src/vector/backends/awkward.py index fa453a3e..61c2fe90 100644 --- a/src/vector/backends/awkward.py +++ b/src/vector/backends/awkward.py @@ -622,10 +622,38 @@ class AwkwardProtocol(Protocol): def __getitem__(self, where: typing.Any) -> float | ak.Array | ak.Record | None: ... +class _lib(typing.NamedTuple): + """a wrapper that respects the numpy-like interface of awkward-array and the module interface of numpy""" + + module: types.ModuleType + nplike: ak._nplikes.numpy_like.NumpyLike + + def __eq__(self, other: typing.Any) -> bool: + if isinstance(other, _lib): + return self.module is other.module and self.nplike is other.nplike + else: + return self.module is other + + def __ne__(self, other: typing.Any) -> bool: + return not self.__eq__(other) + + def __getattr__(self, name: str) -> typing.Any: + if fun := getattr(self.nplike, name, False): + return fun + else: + return getattr(self.module, name) + + class VectorAwkward: """Mixin class for Awkward vectors.""" - lib: types.ModuleType = numpy + @property + def lib(self): # type:ignore[no-untyped-def] + if ( + nplike := self.layout.backend.nplike # type:ignore[attr-defined] + ) is ak._nplikes.typetracer.TypeTracer.instance(): + return _lib(module=numpy, nplike=nplike) + return numpy def _wrap_result( self: AwkwardProtocol, diff --git a/tests/backends/test_awkward.py b/tests/backends/test_awkward.py index d677a256..4437fd77 100644 --- a/tests/backends/test_awkward.py +++ b/tests/backends/test_awkward.py @@ -5,6 +5,7 @@ from __future__ import annotations +import functools import importlib.metadata import numbers @@ -26,97 +27,125 @@ ) < packaging.version.Version("2.2.3") -def test_dimension_conversion(): +def _assert(bool1, bool2, backend): + # this is a helper function that differentiates between backends: + # - 'typetracer' backend: + # this method does nothing (no-op). typetracers can't be compared as they do not have numerical values. + # this is useful, because the expression that goes into this function (bool1 & bool2) will still be computed, + # i.e. this makes sure that the computation does (at least) not raise an error. + # - other backends: + # convert awkward arrays to lists and compare their values with python's builtin list comparison + if backend != "typetracer": + if isinstance(bool1, ak.Array): + bool1 = bool1.tolist() + if isinstance(bool2, ak.Array): + bool2 = bool2.tolist() + else: + assert bool1 == bool2 + + +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_dimension_conversion(backend): + assert_backend = functools.partial(_assert, backend=backend) + # 2D -> 3D vec = vector.Array( [ [{"x": 1, "y": 1.1}, {"x": 2, "y": 2.1}], [], - ] + ], + backend=backend, ) + # test alias - assert ak.all(vec.to_3D(z=1).z == 1) - assert ak.all(vec.to_3D(eta=1).eta == 1) - assert ak.all(vec.to_3D(theta=1).theta == 1) + assert_backend(ak.all(vec.to_3D(z=1).z == 1), True) + assert_backend(ak.all(vec.to_3D(eta=1).eta == 1), True) + assert_backend(ak.all(vec.to_3D(theta=1).theta == 1), True) - assert ak.all(vec.to_Vector3D(z=1).x == vec.x) - assert ak.all(vec.to_Vector3D(z=1).y == vec.y) + assert_backend(vec.to_Vector3D(z=1).x, vec.x) + assert_backend(vec.to_Vector3D(z=1).y, vec.y) # 2D -> 4D - assert ak.all(vec.to_Vector4D(z=1, t=1).t == 1) - assert ak.all(vec.to_Vector4D(z=1, t=1).z == 1) - assert ak.all(vec.to_Vector4D(eta=1, t=1).eta == 1) - assert ak.all(vec.to_Vector4D(eta=1, t=1).t == 1) - assert ak.all(vec.to_Vector4D(theta=1, t=1).theta == 1) - assert ak.all(vec.to_Vector4D(theta=1, t=1).t == 1) - assert ak.all(vec.to_Vector4D(z=1, tau=1).z == 1) - assert ak.all(vec.to_Vector4D(z=1, tau=1).tau == 1) - assert ak.all(vec.to_Vector4D(eta=1, tau=1).eta == 1) - assert ak.all(vec.to_Vector4D(eta=1, tau=1).tau == 1) - assert ak.all(vec.to_Vector4D(theta=1, tau=1).theta == 1) - assert ak.all(vec.to_Vector4D(theta=1, tau=1).tau == 1) + assert_backend(ak.all(vec.to_Vector4D(z=1, t=1).t == 1), True) + assert_backend(ak.all(vec.to_Vector4D(z=1, t=1).z == 1), True) + assert_backend(ak.all(vec.to_Vector4D(eta=1, t=1).eta == 1), True) + assert_backend(ak.all(vec.to_Vector4D(eta=1, t=1).t == 1), True) + assert_backend(ak.all(vec.to_Vector4D(theta=1, t=1).theta == 1), True) + assert_backend(ak.all(vec.to_Vector4D(theta=1, t=1).t == 1), True) + assert_backend(ak.all(vec.to_Vector4D(z=1, tau=1).z == 1), True) + assert_backend(ak.all(vec.to_Vector4D(z=1, tau=1).tau == 1), True) + assert_backend(ak.all(vec.to_Vector4D(eta=1, tau=1).eta == 1), True) + assert_backend(ak.all(vec.to_Vector4D(eta=1, tau=1).tau == 1), True) + assert_backend(ak.all(vec.to_Vector4D(theta=1, tau=1).theta == 1), True) + assert_backend(ak.all(vec.to_Vector4D(theta=1, tau=1).tau == 1), True) # test alias - assert ak.all(vec.to_4D(z=1, t=1).x == vec.x) - assert ak.all(vec.to_4D(z=1, t=1).y == vec.y) + assert_backend(vec.to_4D(z=1, t=1).x, vec.x) + assert_backend(vec.to_4D(z=1, t=1).y, vec.y) # 3D -> 4D vec = vector.Array( [ [{"x": 1, "y": 1.1, "z": 1.2}, {"x": 2, "y": 2.1, "z": 2.2}], [], - ] + ], + backend=backend, ) - assert ak.all(vec.to_Vector4D(t=1).t == 1) - assert ak.all(vec.to_Vector4D(tau=1).tau == 1) + assert_backend(ak.all(vec.to_Vector4D(t=1).t == 1), True) + assert_backend(ak.all(vec.to_Vector4D(tau=1).tau == 1), True) - assert ak.all(vec.to_Vector4D(t=1).x == vec.x) - assert ak.all(vec.to_Vector4D(t=1).y == vec.y) - assert ak.all(vec.to_Vector4D(t=1).z == vec.z) + assert_backend(vec.to_Vector4D(t=1).x, vec.x) + assert_backend(vec.to_Vector4D(t=1).y, vec.y) + assert_backend(vec.to_Vector4D(t=1).z, vec.z) # check if momentum coords work vec = vector.Array( [ [{"px": 1, "py": 1.1}, {"px": 2, "py": 2.1}], [], - ] + ], + backend=backend, ) - assert ak.all(vec.to_Vector3D(pz=1).pz == 1) + assert_backend(ak.all(vec.to_Vector3D(pz=1).pz == 1), True) - assert ak.all(vec.to_Vector4D(pz=1, m=1).pz == 1) - assert ak.all(vec.to_Vector4D(pz=1, m=1).m == 1) - assert ak.all(vec.to_Vector4D(pz=1, mass=1).mass == 1) - assert ak.all(vec.to_Vector4D(pz=1, M=1).M == 1) - assert ak.all(vec.to_Vector4D(pz=1, e=1).e == 1) - assert ak.all(vec.to_Vector4D(pz=1, energy=1).energy == 1) - assert ak.all(vec.to_Vector4D(pz=1, E=1).E == 1) + assert_backend(ak.all(vec.to_Vector4D(pz=1, m=1).pz == 1), True) + assert_backend(ak.all(vec.to_Vector4D(pz=1, m=1).m == 1), True) + assert_backend(ak.all(vec.to_Vector4D(pz=1, mass=1).mass == 1), True) + assert_backend(ak.all(vec.to_Vector4D(pz=1, M=1).M == 1), True) + assert_backend(ak.all(vec.to_Vector4D(pz=1, e=1).e == 1), True) + assert_backend(ak.all(vec.to_Vector4D(pz=1, energy=1).energy == 1), True) + assert_backend(ak.all(vec.to_Vector4D(pz=1, E=1).E == 1), True) vec = vector.Array( [ [{"px": 1, "py": 1.1, "pz": 1.2}, {"px": 2, "py": 2.1, "pz": 2.2}], [], - ] + ], + backend=backend, ) - assert ak.all(vec.to_Vector4D(m=1).m == 1) - assert ak.all(vec.to_Vector4D(mass=1).mass == 1) - assert ak.all(vec.to_Vector4D(M=1).M == 1) - assert ak.all(vec.to_Vector4D(e=1).e == 1) - assert ak.all(vec.to_Vector4D(energy=1).energy == 1) - assert ak.all(vec.to_Vector4D(E=1).E == 1) + assert_backend(ak.all(vec.to_Vector4D(m=1).m == 1), True) + assert_backend(ak.all(vec.to_Vector4D(mass=1).mass == 1), True) + assert_backend(ak.all(vec.to_Vector4D(M=1).M == 1), True) + assert_backend(ak.all(vec.to_Vector4D(e=1).e == 1), True) + assert_backend(ak.all(vec.to_Vector4D(energy=1).energy == 1), True) + assert_backend(ak.all(vec.to_Vector4D(E=1).E == 1), True) -def test_type_checks(): +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_type_checks(backend): with pytest.raises(TypeError): vector.Array( [ [{"x": 1, "y": 1.1, "z": 0.1}, {"x": 2, "y": 2.2, "z": [0.2]}], [], - ] + ], + backend=backend, ) with pytest.raises(TypeError): vector.Array( - [{"x": 1, "y": 1.1, "z": 0.1}, {"x": 2, "y": 2.2, "z": complex(1, 2)}] + [{"x": 1, "y": 1.1, "z": 0.1}, {"x": 2, "y": 2.2, "z": complex(1, 2)}], + backend=backend, ) with pytest.raises(TypeError): @@ -128,186 +157,262 @@ def test_type_checks(): ) -def test_basic(): - array = vector.Array([[{"x": 1, "y": 2}], [], [{"x": 3, "y": 4}]]) +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_basic(backend): + assert_backend = functools.partial(_assert, backend=backend) + + array = vector.Array( + [[{"x": 1, "y": 2}], [], [{"x": 3, "y": 4}]], + backend=backend, + ) assert isinstance(array, vector.backends.awkward.VectorArray2D) - assert array.x.tolist() == [[1], [], [3]] - assert array.y.tolist() == [[2], [], [4]] - assert array.rho.tolist() == [[2.23606797749979], [], [5]] - (a,), (), (c,) = array.phi.tolist() - assert a == pytest.approx(1.1071487177940904) - assert c == pytest.approx(0.9272952180016122) + assert_backend(array.x, [[1], [], [3]]) + assert_backend(array.y, [[2], [], [4]]) + assert_backend(array.rho, [[2.23606797749979], [], [5]]) + assert isinstance(array[2, 0], vector.backends.awkward.VectorRecord2D) - assert array[2, 0].rho == 5 - assert array.deltaphi(array).tolist() == [[0], [], [0]] + if backend == "cpu": + (a,), (), (c,) = array.phi.tolist() + assert a == pytest.approx(1.1071487177940904) + assert c == pytest.approx(0.9272952180016122) + assert array[2, 0].rho == 5 + assert array.deltaphi(array).tolist() == [[0], [], [0]] + else: + # at least make sure they run + _ = array.phi + _ = array[2, 0].rho + _ = array.deltaphi(array) - array = vector.Array([[{"pt": 1, "phi": 2}], [], [{"pt": 3, "phi": 4}]]) + array = vector.Array( + [[{"pt": 1, "phi": 2}], [], [{"pt": 3, "phi": 4}]], + backend=backend, + ) assert isinstance(array, vector.backends.awkward.MomentumArray2D) - assert array.pt.tolist() == [[1], [], [3]] + assert_backend(array.pt, [[1], [], [3]]) array = vector.Array( [ [{"x": 1, "y": 2, "z": 3, "wow": 99}], [], [{"x": 4, "y": 5, "z": 6, "wow": 123}], - ] + ], + backend=backend, ) assert isinstance(array, vector.backends.awkward.VectorArray3D) - assert array.wow.tolist() == [[99], [], [123]] + assert_backend(array.wow, [[99], [], [123]]) + +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_rotateZ(backend): + assert_backend = functools.partial(_assert, backend=backend) -def test_rotateZ(): - array = vector.Array([[{"pt": 1, "phi": 0}], [], [{"pt": 2, "phi": 1}]]) + array = vector.Array( + [[{"pt": 1, "phi": 0}], [], [{"pt": 2, "phi": 1}]], + backend=backend, + ) out = array.rotateZ(1) assert isinstance(out, vector.backends.awkward.MomentumArray2D) - assert out.tolist() == [[{"rho": 1, "phi": 1}], [], [{"rho": 2, "phi": 2}]] + assert_backend(out, [[{"rho": 1, "phi": 1}], [], [{"rho": 2, "phi": 2}]]) array = vector.Array( - [[{"x": 1, "y": 0, "wow": 99}], [], [{"x": 2, "y": 1, "wow": 123}]] + [[{"x": 1, "y": 0, "wow": 99}], [], [{"x": 2, "y": 1, "wow": 123}]], + backend=backend, ) out = array.rotateZ(0.1) assert isinstance(out, vector.backends.awkward.VectorArray2D) - assert out.wow.tolist() == [[99], [], [123]] + assert_backend(out.wow, [[99], [], [123]]) + +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_projection(backend): + assert_backend = functools.partial(_assert, backend=backend) -def test_projection(): array = vector.Array( [ [{"x": 1, "y": 2, "z": 3, "wow": 99}], [], [{"x": 4, "y": 5, "z": 6, "wow": 123}], - ] + ], + backend=backend, ) out = array.to_Vector2D() assert isinstance(out, vector.backends.awkward.VectorArray2D) - assert out.tolist() == [ - [{"x": 1, "y": 2, "wow": 99}], - [], - [{"x": 4, "y": 5, "wow": 123}], - ] + assert_backend( + out, + [ + [{"x": 1, "y": 2, "wow": 99}], + [], + [{"x": 4, "y": 5, "wow": 123}], + ], + ) out = array.to_Vector4D() assert isinstance(out, vector.backends.awkward.VectorArray4D) - assert out.tolist() == [ - [{"x": 1, "y": 2, "z": 3, "t": 0, "wow": 99}], - [], - [{"x": 4, "y": 5, "z": 6, "t": 0, "wow": 123}], - ] + assert_backend( + out, + [ + [{"x": 1, "y": 2, "z": 3, "t": 0, "wow": 99}], + [], + [{"x": 4, "y": 5, "z": 6, "t": 0, "wow": 123}], + ], + ) out = array.to_rhophietatau() assert isinstance(out, vector.backends.awkward.VectorArray4D) - (a,), (), (c,) = out.tolist() - assert a == pytest.approx( - { - "rho": 2.23606797749979, - "phi": 1.1071487177940904, - "eta": 1.1035868415601453, - "tau": 0, - "wow": 99, - } - ) - assert c == pytest.approx( - { - "rho": 6.4031242374328485, - "phi": 0.8960553845713439, - "eta": 0.8361481196083127, - "tau": 0, - "wow": 123, - } - ) - - -def test_add(): + if backend == "cpu": + (a,), (), (c,) = out.tolist() + assert a == pytest.approx( + { + "rho": 2.23606797749979, + "phi": 1.1071487177940904, + "eta": 1.1035868415601453, + "tau": 0, + "wow": 99, + } + ) + assert c == pytest.approx( + { + "rho": 6.4031242374328485, + "phi": 0.8960553845713439, + "eta": 0.8361481196083127, + "tau": 0, + "wow": 123, + } + ) + + +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_add(backend): + assert_backend = functools.partial(_assert, backend=backend) + one = vector.Array( - [[{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}], [], [{"x": 3, "y": 3.3}]] + [[{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}], [], [{"x": 3, "y": 3.3}]], + backend=backend, ) two = vector.Array( - [{"x": 10, "y": 20}, {"x": 100, "y": 200}, {"x": 1000, "y": 2000}] + [{"x": 10, "y": 20}, {"x": 100, "y": 200}, {"x": 1000, "y": 2000}], + backend=backend, ) assert isinstance(one.add(two), vector.backends.awkward.VectorArray2D) assert isinstance(two.add(one), vector.backends.awkward.VectorArray2D) - assert one.add(two).tolist() == [ - [{"x": 11, "y": 21.1}, {"x": 12, "y": 22.2}], - [], - [{"x": 1003, "y": 2003.3}], - ] - assert two.add(one).tolist() == [ - [{"x": 11, "y": 21.1}, {"x": 12, "y": 22.2}], - [], - [{"x": 1003, "y": 2003.3}], - ] + assert_backend( + one.add(two), + [ + [{"x": 11, "y": 21.1}, {"x": 12, "y": 22.2}], + [], + [{"x": 1003, "y": 2003.3}], + ], + ) + assert_backend( + two.add(one), + [ + [{"x": 11, "y": 21.1}, {"x": 12, "y": 22.2}], + [], + [{"x": 1003, "y": 2003.3}], + ], + ) two = vector.array({"x": [10, 100, 1000], "y": [20, 200, 2000]}) assert isinstance(one.add(two), vector.backends.awkward.VectorArray2D) assert isinstance(two.add(one), vector.backends.awkward.VectorArray2D) - assert one.add(two).tolist() == [ - [{"x": 11, "y": 21.1}, {"x": 12, "y": 22.2}], - [], - [{"x": 1003, "y": 2003.3}], - ] - assert two.add(one).tolist() == [ - [{"x": 11, "y": 21.1}, {"x": 12, "y": 22.2}], - [], - [{"x": 1003, "y": 2003.3}], - ] + assert_backend( + one.add(two), + [ + [{"x": 11, "y": 21.1}, {"x": 12, "y": 22.2}], + [], + [{"x": 1003, "y": 2003.3}], + ], + ) + assert_backend( + two.add(one), + [ + [{"x": 11, "y": 21.1}, {"x": 12, "y": 22.2}], + [], + [{"x": 1003, "y": 2003.3}], + ], + ) two = vector.obj(x=10, y=20) assert isinstance(one.add(two), vector.backends.awkward.VectorArray2D) assert isinstance(two.add(one), vector.backends.awkward.VectorArray2D) - assert one.add(two).tolist() == [ - [{"x": 11, "y": 21.1}, {"x": 12, "y": 22.2}], - [], - [{"x": 13, "y": 23.3}], - ] - assert two.add(one).tolist() == [ - [{"x": 11, "y": 21.1}, {"x": 12, "y": 22.2}], - [], - [{"x": 13, "y": 23.3}], - ] - - -def test_ufuncs(): + assert_backend( + one.add(two), + [ + [{"x": 11, "y": 21.1}, {"x": 12, "y": 22.2}], + [], + [{"x": 13, "y": 23.3}], + ], + ) + assert_backend( + two.add(one), + [ + [{"x": 11, "y": 21.1}, {"x": 12, "y": 22.2}], + [], + [{"x": 13, "y": 23.3}], + ], + ) + + +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_ufuncs(backend): + assert_backend = functools.partial(_assert, backend=backend) + array = vector.Array( - [[{"x": 3, "y": 4}, {"x": 5, "y": 12}], [], [{"x": 8, "y": 15}]] + [[{"x": 3, "y": 4}, {"x": 5, "y": 12}], [], [{"x": 8, "y": 15}]], + backend=backend, ) - assert abs(array).tolist() == [[5, 13], [], [17]] - assert (array**2).tolist() == [[5**2, 13**2], [], [17**2]] + assert_backend(abs(array), [[5, 13], [], [17]]) + assert_backend((array**2), [[5**2, 13**2], [], [17**2]]) one = vector.Array( - [[{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}], [], [{"x": 3, "y": 3.3}]] + [[{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}], [], [{"x": 3, "y": 3.3}]], + backend=backend, ) two = vector.Array( - [{"x": 10, "y": 20}, {"x": 100, "y": 200}, {"x": 1000, "y": 2000}] + [{"x": 10, "y": 20}, {"x": 100, "y": 200}, {"x": 1000, "y": 2000}], + backend=backend, + ) + assert_backend( + one + two, + [ + [{"x": 11, "y": 21.1}, {"x": 12, "y": 22.2}], + [], + [{"x": 1003, "y": 2003.3}], + ], + ) + + assert_backend( + one * 10, + [ + [{"x": 10, "y": 11}, {"x": 20, "y": 22}], + [], + [{"x": 30, "y": 33}], + ], ) - assert (one + two).tolist() == [ - [{"x": 11, "y": 21.1}, {"x": 12, "y": 22.2}], - [], - [{"x": 1003, "y": 2003.3}], - ] - assert (one * 10).tolist() == [ - [{"x": 10, "y": 11}, {"x": 20, "y": 22}], - [], - [{"x": 30, "y": 33}], - ] +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_zip(backend): + assert_backend = functools.partial(_assert, backend=backend) -def test_zip(): - v = vector.zip({"x": [[], [1]], "y": [[], [1]]}) + v = ak.to_backend(vector.zip({"x": [[], [1]], "y": [[], [1]]}), backend=backend) assert isinstance(v, vector.backends.awkward.VectorArray2D) assert isinstance(v[1], vector.backends.awkward.VectorArray2D) assert isinstance(v[1, 0], vector.backends.awkward.VectorRecord2D) - assert v.tolist() == [[], [{"x": 1, "y": 1}]] - assert v.x.tolist() == [[], [1]] - assert v.y.tolist() == [[], [1]] + assert_backend(v, [[], [{"x": 1, "y": 1}]]) + assert_backend(v.x, [[], [1]]) + assert_backend(v.y, [[], [1]]) @pytest.mark.skipif( awkward_without_record_reducers, reason="record reducers were added before awkward==2.2.3, but had some bugs", ) -def test_sum_2d(): +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_sum_2d(backend): + assert_backend = functools.partial(_assert, backend=backend) + v = vector.Array( [ [ @@ -319,65 +424,86 @@ def test_sum_2d(): {"rho": 4.0, "phi": 0.4}, {"rho": 1.0, "phi": 0.1}, ], - ] - ) - assert ak.almost_equal( - ak.sum(v, axis=0, keepdims=True), - vector.Array( - [ - [ - {"x": 1.950340654403632, "y": 0.3953536233081677}, - {"x": 7.604510287376507, "y": 2.3523506924148467}, - {"x": 0.9950041652780258, "y": 0.09983341664682815}, - ] - ] - ), - ) - assert ak.almost_equal( - ak.sum(v, axis=0, keepdims=False), - vector.Array( - [ - {"x": 1.950340654403632, "y": 0.3953536233081677}, - {"x": 7.604510287376507, "y": 2.3523506924148467}, - {"x": 0.9950041652780258, "y": 0.09983341664682815}, - ] - ), - ) - assert ak.almost_equal( - ak.sum(v, axis=1, keepdims=True), - ak.to_regular( - vector.Array( - [ - [{"x": 4.915270, "y": 0.89451074}], - [{"x": 5.63458463, "y": 1.95302699}], - ] - ) - ), - ) - assert ak.almost_equal( - ak.sum(v, axis=1, keepdims=False), - vector.Array( - [ - {"x": 4.915270, "y": 0.89451074}, - {"x": 5.63458463, "y": 1.95302699}, - ] - ), - ) - assert ak.almost_equal( - ak.sum(v.mask[[False, True]], axis=1), - vector.Array( - [ - {"x": 5.63458463, "y": 1.95302699}, - ] - )[[None, 0]], + ], + backend=backend, ) + # typetracer backend does not implement ak.almost_equal + if backend != "typetracer": + assert_backend( + ak.almost_equal( + ak.sum(v, axis=0, keepdims=True), + vector.Array( + [ + [ + {"x": 1.950340654403632, "y": 0.3953536233081677}, + {"x": 7.604510287376507, "y": 2.3523506924148467}, + {"x": 0.9950041652780258, "y": 0.09983341664682815}, + ] + ] + ), + ), + True, + ) + assert_backend( + ak.almost_equal( + ak.sum(v, axis=0, keepdims=False), + vector.Array( + [ + {"x": 1.950340654403632, "y": 0.3953536233081677}, + {"x": 7.604510287376507, "y": 2.3523506924148467}, + {"x": 0.9950041652780258, "y": 0.09983341664682815}, + ] + ), + ), + True, + ) + assert_backend( + ak.almost_equal( + ak.sum(v, axis=1, keepdims=True), + ak.to_regular( + vector.Array( + [ + [{"x": 4.915270, "y": 0.89451074}], + [{"x": 5.63458463, "y": 1.95302699}], + ] + ) + ), + ), + True, + ) + assert_backend( + ak.almost_equal( + ak.sum(v, axis=1, keepdims=False), + vector.Array( + [ + {"x": 4.915270, "y": 0.89451074}, + {"x": 5.63458463, "y": 1.95302699}, + ] + ), + ), + True, + ) + assert_backend( + ak.almost_equal( + ak.sum(v.mask[[False, True]], axis=1), + vector.Array( + [ + {"x": 5.63458463, "y": 1.95302699}, + ] + )[[None, 0]], + ), + True, + ) @pytest.mark.skipif( awkward_without_record_reducers, reason="record reducers were added before awkward==2.2.3, but had some bugs", ) -def test_sum_3d(): +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_sum_3d(backend): + assert_backend = functools.partial(_assert, backend=backend) + v = vector.Array( [ [ @@ -389,35 +515,58 @@ def test_sum_3d(): {"x": 4, "y": 5, "z": 6}, {"x": 1, "y": 1, "z": 1}, ], - ] + ], + backend=backend, + ) + assert_backend( + ak.sum(v, axis=0, keepdims=True), + [ + [ + {"x": 2, "y": 4, "z": 6}, + {"x": 8, "y": 10, "z": 12}, + {"x": 1, "y": 1, "z": 1}, + ] + ], + ) + assert_backend( + ak.sum(v, axis=0, keepdims=False), + [ + {"x": 2, "y": 4, "z": 6}, + {"x": 8, "y": 10, "z": 12}, + {"x": 1, "y": 1, "z": 1}, + ], + ) + assert_backend( + ak.sum(v, axis=1, keepdims=True), + [ + [{"x": 5, "y": 7, "z": 9}], + [{"x": 6, "y": 8, "z": 10}], + ], + ) + assert_backend( + ak.sum(v, axis=1, keepdims=False), + [ + {"x": 5, "y": 7, "z": 9}, + {"x": 6, "y": 8, "z": 10}, + ], + ) + assert_backend( + ak.sum(v.mask[[False, True]], axis=1), + [ + None, + {"x": 6, "y": 8, "z": 10}, + ], ) - assert ak.sum(v, axis=0, keepdims=True).to_list() == [ - [{"x": 2, "y": 4, "z": 6}, {"x": 8, "y": 10, "z": 12}, {"x": 1, "y": 1, "z": 1}] - ] - assert ak.sum(v, axis=0, keepdims=False).to_list() == [ - {"x": 2, "y": 4, "z": 6}, - {"x": 8, "y": 10, "z": 12}, - {"x": 1, "y": 1, "z": 1}, - ] - assert ak.sum(v, axis=1, keepdims=True).to_list() == [ - [{"x": 5, "y": 7, "z": 9}], - [{"x": 6, "y": 8, "z": 10}], - ] - assert ak.sum(v, axis=1, keepdims=False).to_list() == [ - {"x": 5, "y": 7, "z": 9}, - {"x": 6, "y": 8, "z": 10}, - ] - assert ak.sum(v.mask[[False, True]], axis=1).tolist() == [ - None, - {"x": 6, "y": 8, "z": 10}, - ] @pytest.mark.skipif( awkward_without_record_reducers, reason="record reducers were added before awkward==2.2.3, but had some bugs", ) -def test_sum_4d(): +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_sum_4d(backend): + assert_backend = functools.partial(_assert, backend=backend) + v = vector.Array( [ [ @@ -430,39 +579,58 @@ def test_sum_4d(): {"x": 4, "y": 5, "z": 6, "t": 0}, {"x": 1, "y": 1, "z": 1, "t": 0}, ], - ] + ], + backend=backend, + ) + assert_backend( + ak.sum(v, axis=0, keepdims=True), + [ + [ + {"t": 12, "z": 6, "x": 2, "y": 4}, + {"t": 2, "z": 12, "x": 8, "y": 10}, + {"t": 3, "z": 1, "x": 1, "y": 1}, + ] + ], ) - assert ak.sum(v, axis=0, keepdims=True).to_list() == [ + assert_backend( + ak.sum(v, axis=0, keepdims=False), [ {"t": 12, "z": 6, "x": 2, "y": 4}, {"t": 2, "z": 12, "x": 8, "y": 10}, {"t": 3, "z": 1, "x": 1, "y": 1}, - ] - ] - assert ak.sum(v, axis=0, keepdims=False).to_list() == [ - {"t": 12, "z": 6, "x": 2, "y": 4}, - {"t": 2, "z": 12, "x": 8, "y": 10}, - {"t": 3, "z": 1, "x": 1, "y": 1}, - ] - assert ak.sum(v, axis=1, keepdims=True).to_list() == [ - [{"t": 9, "z": 9, "x": 5, "y": 7}], - [{"t": 8, "z": 10, "x": 6, "y": 8}], - ] - assert ak.sum(v, axis=1, keepdims=False).to_list() == [ - {"t": 9, "z": 9, "x": 5, "y": 7}, - {"t": 8, "z": 10, "x": 6, "y": 8}, - ] - assert ak.sum(v.mask[[False, True]], axis=1).tolist() == [ - None, - {"t": 8, "z": 10, "x": 6, "y": 8}, - ] + ], + ) + assert_backend( + ak.sum(v, axis=1, keepdims=True), + [ + [{"t": 9, "z": 9, "x": 5, "y": 7}], + [{"t": 8, "z": 10, "x": 6, "y": 8}], + ], + ) + assert_backend( + ak.sum(v, axis=1, keepdims=False), + [ + {"t": 9, "z": 9, "x": 5, "y": 7}, + {"t": 8, "z": 10, "x": 6, "y": 8}, + ], + ) + assert_backend( + ak.sum(v.mask[[False, True]], axis=1), + [ + None, + {"t": 8, "z": 10, "x": 6, "y": 8}, + ], + ) @pytest.mark.skipif( awkward_without_record_reducers, reason="record reducers were added before awkward==2.2.3, but had some bugs", ) -def test_count_nonzero_2d(): +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_count_nonzero_2d(backend): + assert_backend = functools.partial(_assert, backend=backend) + v = vector.Array( [ [ @@ -475,19 +643,23 @@ def test_count_nonzero_2d(): {"rho": 4.0, "phi": 0.4}, {"rho": 1.0, "phi": 0.1}, ], - ] + ], + backend=backend, ) - assert ak.count_nonzero(v, axis=1).tolist() == [2, 3] - assert ak.count_nonzero(v, axis=1, keepdims=True).tolist() == [[2], [3]] - assert ak.count_nonzero(v, axis=0).tolist() == [2, 2, 1] - assert ak.count_nonzero(v, axis=0, keepdims=True).tolist() == [[2, 2, 1]] + assert_backend(ak.count_nonzero(v, axis=1), [2, 3]) + assert_backend(ak.count_nonzero(v, axis=1, keepdims=True), [[2], [3]]) + assert_backend(ak.count_nonzero(v, axis=0), [2, 2, 1]) + assert_backend(ak.count_nonzero(v, axis=0, keepdims=True), [[2, 2, 1]]) @pytest.mark.skipif( awkward_without_record_reducers, reason="record reducers were added before awkward==2.2.3, but had some bugs", ) -def test_count_nonzero_3d(): +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_count_nonzero_3d(backend): + assert_backend = functools.partial(_assert, backend=backend) + v = vector.Array( [ [ @@ -500,19 +672,23 @@ def test_count_nonzero_3d(): {"x": 4.0, "y": 5.0, "theta": 1.3}, {"x": 1.0, "y": 1.0, "theta": 1.9}, ], - ] + ], + backend=backend, ) - assert ak.count_nonzero(v, axis=1).tolist() == [2, 3] - assert ak.count_nonzero(v, axis=1, keepdims=True).tolist() == [[2], [3]] - assert ak.count_nonzero(v, axis=0).tolist() == [2, 2, 1] - assert ak.count_nonzero(v, axis=0, keepdims=True).tolist() == [[2, 2, 1]] + assert_backend(ak.count_nonzero(v, axis=1), [2, 3]) + assert_backend(ak.count_nonzero(v, axis=1, keepdims=True), [[2], [3]]) + assert_backend(ak.count_nonzero(v, axis=0), [2, 2, 1]) + assert_backend(ak.count_nonzero(v, axis=0, keepdims=True), [[2, 2, 1]]) @pytest.mark.skipif( awkward_without_record_reducers, reason="record reducers were added before awkward==2.2.3, but had some bugs", ) -def test_count_nonzero_4d(): +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_count_nonzero_4d(backend): + assert_backend = functools.partial(_assert, backend=backend) + v = vector.Array( [ [ @@ -525,12 +701,13 @@ def test_count_nonzero_4d(): {"x": 4.0, "y": 5.0, "z": 6.0, "t": 0.0}, {"x": 1.0, "y": 1.0, "z": 1.0, "t": 0.0}, ], - ] + ], + backend=backend, ) - assert ak.count_nonzero(v, axis=1).tolist() == [3, 3] - assert ak.count_nonzero(v, axis=1, keepdims=True).tolist() == [[3], [3]] - assert ak.count_nonzero(v, axis=0).tolist() == [2, 2, 2] - assert ak.count_nonzero(v, axis=0, keepdims=True).tolist() == [[2, 2, 2]] + assert_backend(ak.count_nonzero(v, axis=1), [3, 3]) + assert_backend(ak.count_nonzero(v, axis=1, keepdims=True), [[3], [3]]) + assert_backend(ak.count_nonzero(v, axis=0), [2, 2, 2]) + assert_backend(ak.count_nonzero(v, axis=0, keepdims=True), [[2, 2, 2]]) v2 = vector.Array( [ @@ -544,19 +721,23 @@ def test_count_nonzero_4d(): {"x": 4, "y": 5, "z": 6, "t": 1}, {"x": 0, "y": 0, "z": 0, "t": 0}, ], - ] + ], + backend=backend, ) - assert ak.count_nonzero(v2, axis=1).tolist() == [3, 2] - assert ak.count_nonzero(v2, axis=1, keepdims=True).tolist() == [[3], [2]] - assert ak.count_nonzero(v2, axis=0).tolist() == [2, 2, 1] - assert ak.count_nonzero(v2, axis=0, keepdims=True).tolist() == [[2, 2, 1]] + assert_backend(ak.count_nonzero(v2, axis=1), [3, 2]) + assert_backend(ak.count_nonzero(v2, axis=1, keepdims=True), [[3], [2]]) + assert_backend(ak.count_nonzero(v2, axis=0), [2, 2, 1]) + assert_backend(ak.count_nonzero(v2, axis=0, keepdims=True), [[2, 2, 1]]) @pytest.mark.skipif( awkward_without_record_reducers, reason="record reducers were added before awkward==2.2.3, but had some bugs", ) -def test_count_2d(): +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_count_2d(backend): + assert_backend = functools.partial(_assert, backend=backend) + v = vector.Array( [ [ @@ -569,23 +750,24 @@ def test_count_2d(): {"x": 4, "y": 5}, {"x": 1, "y": 1}, ], - ] + ], + backend=backend, ) - assert ak.count(v, axis=1).to_list() == [3, 3] - assert ak.count(v, axis=1, keepdims=True).to_list() == [[3], [3]] - assert ak.count(v, axis=0).to_list() == [2, 2, 2] - assert ak.count(v, axis=0, keepdims=True).to_list() == [[2, 2, 2]] - assert ak.count(v.mask[[False, True]], axis=1).tolist() == [ - None, - 3, - ] + assert_backend(ak.count(v, axis=1), [3, 3]) + assert_backend(ak.count(v, axis=1, keepdims=True), [[3], [3]]) + assert_backend(ak.count(v, axis=0), [2, 2, 2]) + assert_backend(ak.count(v, axis=0, keepdims=True), [[2, 2, 2]]) + assert_backend(ak.count(v.mask[[False, True]], axis=1), [None, 3]) @pytest.mark.skipif( awkward_without_record_reducers, reason="record reducers were added before awkward==2.2.3, but had some bugs", ) -def test_count_3d(): +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_count_3d(backend): + assert_backend = functools.partial(_assert, backend=backend) + v = vector.Array( [ [ @@ -600,21 +782,21 @@ def test_count_3d(): ], ] ) - assert ak.count(v, axis=1).to_list() == [3, 3] - assert ak.count(v, axis=1, keepdims=True).to_list() == [[3], [3]] - assert ak.count(v, axis=0).to_list() == [2, 2, 2] - assert ak.count(v, axis=0, keepdims=True).to_list() == [[2, 2, 2]] - assert ak.count(v.mask[[False, True]], axis=1).tolist() == [ - None, - 3, - ] + assert_backend(ak.count(v, axis=1), [3, 3]) + assert_backend(ak.count(v, axis=1, keepdims=True), [[3], [3]]) + assert_backend(ak.count(v, axis=0), [2, 2, 2]) + assert_backend(ak.count(v, axis=0, keepdims=True), [[2, 2, 2]]) + assert_backend(ak.count(v.mask[[False, True]], axis=1), [None, 3]) @pytest.mark.skipif( awkward_without_record_reducers, reason="record reducers were added before awkward==2.2.3, but had some bugs", ) -def test_count_4d(): +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_count_4d(backend): + assert_backend = functools.partial(_assert, backend=backend) + v = vector.Array( [ [ @@ -629,74 +811,98 @@ def test_count_4d(): ], ] ) - assert ak.count(v, axis=1).to_list() == [3, 3] - assert ak.count(v, axis=1, keepdims=True).to_list() == [[3], [3]] - assert ak.count(v, axis=0).to_list() == [2, 2, 2] - assert ak.count(v, axis=0, keepdims=True).to_list() == [[2, 2, 2]] - assert ak.count(v.mask[[False, True]], axis=1).tolist() == [ - None, - 3, - ] - - -def test_like(): - v1 = vector.zip( - { - "x": [10.0, 20.0, 30.0], - "y": [-10.0, 20.0, 30.0], - }, - ) - v2 = vector.zip( - { - "x": [10.0, 20.0, 30.0], - "y": [-10.0, 20.0, 30.0], - "z": [5.0, 1.0, 1.0], - }, - ) - v3 = vector.zip( - { - "x": [10.0, 20.0, 30.0], - "y": [-10.0, 20.0, 30.0], - "z": [5.0, 1.0, 1.0], - "t": [16.0, 31.0, 46.0], - }, - ) - - v1_v2 = vector.zip( - { - "x": [20.0, 40.0, 60.0], - "y": [-20.0, 40.0, 60.0], - }, - ) - v2_v1 = vector.zip( - { - "x": [20.0, 40.0, 60.0], - "y": [-20.0, 40.0, 60.0], - "z": [5.0, 1.0, 1.0], - }, - ) - v2_v3 = vector.zip( - { - "x": [20.0, 40.0, 60.0], - "y": [-20.0, 40.0, 60.0], - "z": [10.0, 2.0, 2.0], - }, - ) - v3_v2 = vector.zip( - { - "x": [20.0, 40.0, 60.0], - "y": [-20.0, 40.0, 60.0], - "z": [10.0, 2.0, 2.0], - "t": [16.0, 31.0, 46.0], - }, - ) - v1_v3 = vector.zip( - { - "x": [20.0, 40.0, 60.0], - "y": [-20.0, 40.0, 60.0], - "z": [5.0, 1.0, 1.0], - "t": [16.0, 31.0, 46.0], - }, + assert_backend(ak.count(v, axis=1), [3, 3]) + assert_backend(ak.count(v, axis=1, keepdims=True), [[3], [3]]) + assert_backend(ak.count(v, axis=0), [2, 2, 2]) + assert_backend(ak.count(v, axis=0, keepdims=True), [[2, 2, 2]]) + assert_backend(ak.count(v.mask[[False, True]], axis=1), [None, 3]) + + +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_like(backend): + assert_backend = functools.partial(_assert, backend=backend) + + v1 = ak.to_backend( + vector.zip( + { + "x": [10.0, 20.0, 30.0], + "y": [-10.0, 20.0, 30.0], + }, + ), + backend=backend, + ) + v2 = ak.to_backend( + vector.zip( + { + "x": [10.0, 20.0, 30.0], + "y": [-10.0, 20.0, 30.0], + "z": [5.0, 1.0, 1.0], + }, + ), + backend=backend, + ) + v3 = ak.to_backend( + vector.zip( + { + "x": [10.0, 20.0, 30.0], + "y": [-10.0, 20.0, 30.0], + "z": [5.0, 1.0, 1.0], + "t": [16.0, 31.0, 46.0], + }, + ), + backend=backend, + ) + + v1_v2 = ak.to_backend( + vector.zip( + { + "x": [20.0, 40.0, 60.0], + "y": [-20.0, 40.0, 60.0], + }, + ), + backend=backend, + ) + v2_v1 = ak.to_backend( + vector.zip( + { + "x": [20.0, 40.0, 60.0], + "y": [-20.0, 40.0, 60.0], + "z": [5.0, 1.0, 1.0], + }, + ), + backend=backend, + ) + v2_v3 = ak.to_backend( + vector.zip( + { + "x": [20.0, 40.0, 60.0], + "y": [-20.0, 40.0, 60.0], + "z": [10.0, 2.0, 2.0], + }, + ), + backend=backend, + ) + v3_v2 = ak.to_backend( + vector.zip( + { + "x": [20.0, 40.0, 60.0], + "y": [-20.0, 40.0, 60.0], + "z": [10.0, 2.0, 2.0], + "t": [16.0, 31.0, 46.0], + }, + ), + backend=backend, + ) + v1_v3 = ak.to_backend( + vector.zip( + { + "x": [20.0, 40.0, 60.0], + "y": [-20.0, 40.0, 60.0], + "z": [5.0, 1.0, 1.0], + "t": [16.0, 31.0, 46.0], + }, + ), + backend=backend, ) with pytest.raises(TypeError): @@ -731,94 +937,116 @@ def test_like(): v1.dot(v3) # 2D + 3D.like(2D) = 2D - assert ak.all(v1 + v2.like(v1) == v1_v2) - assert ak.all(v2.like(v1) + v1 == v1_v2) + assert_backend(v1 + v2.like(v1), v1_v2) + assert_backend(v2.like(v1) + v1, v1_v2) # 2D + 4D.like(2D) = 2D - assert ak.all(v1 + v3.like(v1) == v1_v2) - assert ak.all(v3.like(v1) + v1 == v1_v2) + assert_backend(v1 + v3.like(v1), v1_v2) + assert_backend(v3.like(v1) + v1, v1_v2) # 3D + 2D.like(3D) = 3D - assert ak.all(v2 + v1.like(v2) == v2_v1) - assert ak.all(v1.like(v2) + v2 == v2_v1) + assert_backend(v2 + v1.like(v2), v2_v1) + assert_backend(v1.like(v2) + v2, v2_v1) # 3D + 4D.like(3D) = 3D - assert ak.all(v2 + v3.like(v2) == v2_v3) - assert ak.all(v3.like(v2) + v2 == v2_v3) + assert_backend(v2 + v3.like(v2), v2_v3) + assert_backend(v3.like(v2) + v2, v2_v3) # 4D + 2D.like(4D) = 4D - assert ak.all(v3 + v1.like(v3) == v1_v3) - assert ak.all(v1.like(v3) + v3 == v1_v3) + assert_backend(v3 + v1.like(v3), v1_v3) + assert_backend(v1.like(v3) + v3, v1_v3) # 4D + 3D.like(4D) = 4D - assert ak.all(v3 + v2.like(v3) == v3_v2) - assert ak.all(v2.like(v3) + v3 == v3_v2) - - v1 = vector.zip( - { - "px": [10.0, 20.0, 30.0], - "py": [-10.0, 20.0, 30.0], - }, - ) - v2 = vector.zip( - { - "px": [10.0, 20.0, 30.0], - "py": [-10.0, 20.0, 30.0], - "pz": [5.0, 1.0, 1.0], - }, - ) - v3 = vector.zip( - { - "px": [10.0, 20.0, 30.0], - "py": [-10.0, 20.0, 30.0], - "pz": [5.0, 1.0, 1.0], - "t": [16.0, 31.0, 46.0], - }, + assert_backend(v3 + v2.like(v3), v3_v2) + assert_backend(v2.like(v3) + v3, v3_v2) + + v1 = ak.to_backend( + vector.zip( + { + "px": [10.0, 20.0, 30.0], + "py": [-10.0, 20.0, 30.0], + }, + ), + backend=backend, + ) + v2 = ak.to_backend( + vector.zip( + { + "px": [10.0, 20.0, 30.0], + "py": [-10.0, 20.0, 30.0], + "pz": [5.0, 1.0, 1.0], + }, + ), + backend=backend, + ) + v3 = ak.to_backend( + vector.zip( + { + "px": [10.0, 20.0, 30.0], + "py": [-10.0, 20.0, 30.0], + "pz": [5.0, 1.0, 1.0], + "t": [16.0, 31.0, 46.0], + }, + ), + backend=backend, ) # 2D + 3D.like(2D) = 2D - assert ak.all(v1 + v2.like(v1) == v1_v2) - assert ak.all(v2.like(v1) + v1 == v1_v2) + assert_backend(v1 + v2.like(v1), v1_v2) + assert_backend(v2.like(v1) + v1, v1_v2) # 2D + 4D.like(2D) = 2D - assert ak.all(v1 + v3.like(v1) == v1_v2) - assert ak.all(v3.like(v1) + v1 == v1_v2) + assert_backend(v1 + v3.like(v1), v1_v2) + assert_backend(v3.like(v1) + v1, v1_v2) # 3D + 2D.like(3D) = 3D - assert ak.all(v2 + v1.like(v2) == v2_v1) - assert ak.all(v1.like(v2) + v2 == v2_v1) + assert_backend(v2 + v1.like(v2), v2_v1) + assert_backend(v1.like(v2) + v2, v2_v1) # 3D + 4D.like(3D) = 3D - assert ak.all(v2 + v3.like(v2) == v2_v3) - assert ak.all(v3.like(v2) + v2 == v2_v3) + assert_backend(v2 + v3.like(v2), v2_v3) + assert_backend(v3.like(v2) + v2, v2_v3) # 4D + 2D.like(4D) = 4D - assert ak.all(v3 + v1.like(v3) == v1_v3) - assert ak.all(v1.like(v3) + v3 == v1_v3) + assert_backend(v3 + v1.like(v3), v1_v3) + assert_backend(v1.like(v3) + v3, v1_v3) # 4D + 3D.like(4D) = 4D - assert ak.all(v3 + v2.like(v3) == v3_v2) - assert ak.all(v2.like(v3) + v3 == v3_v2) - - -def test_handler_of(): - numpy_vec = vector.array( - { - "x": [1.1, 1.2, 1.3, 1.4, 1.5], - "y": [2.1, 2.2, 2.3, 2.4, 2.5], - "z": [3.1, 3.2, 3.3, 3.4, 3.5], - } - ) - awkward_vec2 = vector.zip( - { - "x": [1.0, 2.0, 3.0], - "y": [-1.0, 2.0, 3.0], - "z": [5.0, 10.0, 15.0], - "t": [16.0, 31.0, 46.0], - }, + assert_backend(v3 + v2.like(v3), v3_v2) + assert_backend(v2.like(v3) + v3, v3_v2) + + +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_handler_of(backend): + assert_backend = functools.partial(_assert, backend=backend) + + numpy_vec = ak.to_backend( + vector.array( + { + "x": [1.1, 1.2, 1.3, 1.4, 1.5], + "y": [2.1, 2.2, 2.3, 2.4, 2.5], + "z": [3.1, 3.2, 3.3, 3.4, 3.5], + } + ), + backend=backend, + ) + awkward_vec2 = ak.to_backend( + vector.zip( + { + "x": [1.0, 2.0, 3.0], + "y": [-1.0, 2.0, 3.0], + "z": [5.0, 10.0, 15.0], + "t": [16.0, 31.0, 46.0], + }, + ), + backend=backend, ) object_vec = VectorObject2D.from_xy(1.0, 1.0) protocol = vector._methods._handler_of(object_vec, numpy_vec, awkward_vec2) # chooses awkward backend - assert all(protocol == awkward_vec2) + assert_backend(protocol, awkward_vec2) -def test_momentum_coordinate_transforms(): - awkward_vec = vector.zip( - { - "px": [1.0, 2.0, 3.0], - "py": [-1.0, 2.0, 3.0], - }, +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_momentum_coordinate_transforms(backend): + awkward_vec = ak.to_backend( + vector.zip( + { + "px": [1.0, 2.0, 3.0], + "py": [-1.0, 2.0, 3.0], + }, + ), + backend=backend, ) for t1 in "pxpy", "ptphi": @@ -849,28 +1077,38 @@ def test_momentum_coordinate_transforms(): assert hasattr(transformed_object, t3) -def test_momentum_preservation(): - v1 = vector.zip( - { - "px": [10.0, 20.0, 30.0], - "py": [-10.0, 20.0, 30.0], - }, +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_momentum_preservation(backend): + v1 = ak.to_backend( + vector.zip( + { + "px": [10.0, 20.0, 30.0], + "py": [-10.0, 20.0, 30.0], + }, + ), + backend=backend, ) - v2 = vector.zip( - { - "x": [10.0, 20.0, 30.0], - "y": [-10.0, 20.0, 30.0], - "z": [5.0, 1.0, 1.0], - }, + v2 = ak.to_backend( + vector.zip( + { + "x": [10.0, 20.0, 30.0], + "y": [-10.0, 20.0, 30.0], + "z": [5.0, 1.0, 1.0], + }, + ), + backend=backend, ) - v3 = vector.zip( - { - "px": [10.0, 20.0, 30.0], - "py": [-10.0, 20.0, 30.0], - "pz": [5.0, 1.0, 1.0], - "t": [16.0, 31.0, 46.0], - }, + v3 = ak.to_backend( + vector.zip( + { + "px": [10.0, 20.0, 30.0], + "py": [-10.0, 20.0, 30.0], + "pz": [5.0, 1.0, 1.0], + "t": [16.0, 31.0, 46.0], + }, + ), + backend=backend, ) # momentum + generic = momentum @@ -894,7 +1132,8 @@ def test_momentum_preservation(): assert isinstance(v2.like(v3) + v3, vector.backends.awkward.MomentumAwkward4D) -def test_subclass_fields(): +@pytest.mark.parametrize("backend", ["cpu", "typetracer"]) +def test_subclass_fields(backend): @ak.mixin_class(vector.backends.awkward.behavior) class TwoVector(vector.backends.awkward.MomentumAwkward2D): pass @@ -914,15 +1153,18 @@ def divide(self, factor): LorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 LorentzVectorArray.MomentumClass = LorentzVectorArray # noqa: F821 - vec = ak.zip( - { - "pt": [[1, 2], [], [3], [4]], - "eta": [[1.2, 1.4], [], [1.6], [3.4]], - "phi": [[0.3, 0.4], [], [0.5], [0.6]], - "energy": [[50, 51], [], [52], [60]], - }, - with_name="LorentzVector", - behavior=vector.backends.awkward.behavior, + vec = ak.to_backend( + ak.zip( + { + "pt": [[1, 2], [], [3], [4]], + "eta": [[1.2, 1.4], [], [1.6], [3.4]], + "phi": [[0.3, 0.4], [], [0.5], [0.6]], + "energy": [[50, 51], [], [52], [60]], + }, + with_name="LorentzVector", + behavior=vector.backends.awkward.behavior, + ), + backend=backend, ) assert vec.like(vector.obj(x=1, y=2)).fields == ["rho", "phi"]