diff --git a/README.md b/README.md index e9e0b1df..613d7115 100644 --- a/README.md +++ b/README.md @@ -293,6 +293,8 @@ ak.Array( All of the keyword arguments and rules that apply to `vector.obj` construction apply to `vector.awk` field names. +Finally, the `VectorAwkward` mixins can be subclassed to create custom vector classes. The awkward behavior classes and projections must be named as `*Array`. For example, `coffea` uses the following names - `TwoVectorArray`, `ThreeVectorArray`, `PolarTwoVectorArray`, `SphericalThreeVectorArray`, ... + ## Vector properties Any geometrical coordinate can be computed from vectors in any coordinate system; they'll be provided or computed as needed. diff --git a/docs/usage/intro.ipynb b/docs/usage/intro.ipynb index 83d6df58..3b2aa3be 100644 --- a/docs/usage/intro.ipynb +++ b/docs/usage/intro.ipynb @@ -1060,7 +1060,9 @@ "id": "beginning-expert", "metadata": {}, "source": [ - "All of the keyword arguments and rules that apply to `vector.obj` construction apply to `vector.Array` field names." + "All of the keyword arguments and rules that apply to `vector.obj` construction apply to `vector.Array` field names.\n", + "\n", + "Finally, the `VectorAwkward` mixins can be subclassed to create custom vector classes. The awkward behavior classes and projections must be named as `*Array`. For example, `coffea` uses the following names - `TwoVectorArray`, `ThreeVectorArray`, `PolarTwoVectorArray`, `SphericalThreeVectorArray`, ..." ] }, { diff --git a/src/vector/_methods.py b/src/vector/_methods.py index df41fb06..7511cda6 100644 --- a/src/vector/_methods.py +++ b/src/vector/_methods.py @@ -4120,6 +4120,58 @@ def _get_handler_index(obj: VectorProtocol) -> int: ) +def _check_instance( + any_or_all: typing.Callable[[typing.Iterable[bool]], bool], + objects: tuple[VectorProtocol, ...], + clas: type[VectorProtocol], +) -> bool: + return any_or_all(isinstance(v, clas) for v in objects) + + +def _demote_handler_vector( + handler: VectorProtocol, + objects: tuple[VectorProtocol, ...], + vector_class: type[VectorProtocol], + new_vector: VectorProtocol, +) -> VectorProtocol: + """ + Demotes the handler vector to the lowest possible dimension while respecting + the priority of backends. + """ + # if all the objects are not from the same backend + # choose the {X}D object of the backend with highest priority (if it exists) + # or demote the first encountered object of the backend with highest priority to {X}D + backends = [ + next( + x.__module__ + for x in type(obj).__mro__ + if "vector.backends." in x.__module__ + ) + for obj in objects + ] + if len({_handler_priority.index(backend) for backend in backends}) != 1: + new_type = type(new_vector) + flag = 0 + # if there is a {X}D object of the backend with highest priority + # make it the new handler + for obj in objects: + if type(obj) == new_type: + handler = obj + flag = 1 + break + # else, demote the dimension of the object of the backend with highest priority + if flag == 0: + handler = new_vector + # if all objects are from the same backend + # use the {X}D one as the handler + else: + for obj in objects: + if isinstance(obj, vector_class): + handler = obj + + return handler + + def _handler_of(*objects: VectorProtocol) -> VectorProtocol: """ Determines which vector should wrap the output of a dispatched function. @@ -4137,6 +4189,19 @@ def _handler_of(*objects: VectorProtocol) -> VectorProtocol: handler = obj assert handler is not None + + if _check_instance(all, objects, Vector): + # if there is a 2D vector in objects + if _check_instance(any, objects, Vector2D): + handler = _demote_handler_vector( + handler, objects, Vector2D, handler.to_Vector2D() + ) + # if there is no 2D vector but a 3D vector in objects + elif _check_instance(any, objects, Vector3D): + handler = _demote_handler_vector( + handler, objects, Vector3D, handler.to_Vector3D() + ) + return handler diff --git a/src/vector/backends/awkward.py b/src/vector/backends/awkward.py index 4b50f0a1..71d5b72d 100644 --- a/src/vector/backends/awkward.py +++ b/src/vector/backends/awkward.py @@ -575,19 +575,21 @@ def elements(self) -> tuple[ArrayOrRecord]: def _class_to_name(cls: type[VectorProtocol]) -> str: + # respect the type of classes inheriting VectorAwkward classes + is_vector = "vector.backends" in cls.__module__ if issubclass(cls, Momentum): if issubclass(cls, Vector2D): - return "Momentum2D" + return "Momentum2D" if is_vector else cls.__name__[:-5] if issubclass(cls, Vector3D): - return "Momentum3D" + return "Momentum3D" if is_vector else cls.__name__[:-5] if issubclass(cls, Vector4D): - return "Momentum4D" + return "Momentum4D" if is_vector else cls.__name__[:-5] if issubclass(cls, Vector2D): - return "Vector2D" + return "Vector2D" if is_vector else cls.__name__[:-5] if issubclass(cls, Vector3D): - return "Vector3D" + return "Vector3D" if is_vector else cls.__name__[:-5] if issubclass(cls, Vector4D): - return "Vector4D" + return "Vector4D" if is_vector else cls.__name__[:-5] raise AssertionError(repr(cls)) diff --git a/tests/backends/test_awkward.py b/tests/backends/test_awkward.py index 4786fbb6..9797be1b 100644 --- a/tests/backends/test_awkward.py +++ b/tests/backends/test_awkward.py @@ -602,3 +602,101 @@ def test_count_4d(): None, 3, ] + + +def test_demotion(): + 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_v3 = vector.zip( + { + "x": [20.0, 40.0, 60.0], + "y": [-20.0, 40.0, 60.0], + "z": [10.0, 2.0, 2.0], + }, + ) + v1_v3 = vector.zip( + { + "x": [20.0, 40.0, 60.0], + "y": [-20.0, 40.0, 60.0], + }, + ) + + # order should not matter + assert all(v1 + v2 == v1_v2) + assert all(v2 + v1 == v1_v2) + assert all(v1 + v3 == v1_v3) + assert all(v3 + v1 == v1_v3) + assert all(v2 + v3 == v2_v3) + assert all(v3 + v2 == v2_v3) + + 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], + }, + ) + + # order should not matter + assert all(v1 + v2 == v1_v2) + assert all(v2 + v1 == v1_v2) + assert all(v1 + v3 == v1_v3) + assert all(v3 + v1 == v1_v3) + assert all(v2 + v3 == v2_v3) + assert all(v3 + v2 == v2_v3) + + v2 = vector.zip( + { + "x": [10.0, 20.0, 30.0], + "y": [-10.0, 20.0, 30.0], + "z": [5.0, 1.0, 1.0], + }, + ) + + # momentum + generic = generic + assert all(v1 + v2 == v1_v2) + assert all(v2 + v1 == v1_v2) + assert all(v1 + v3 == v1_v3) + assert all(v3 + v1 == v1_v3) + assert all(v2 + v3 == v2_v3) + assert all(v3 + v2 == v2_v3) diff --git a/tests/backends/test_numpy.py b/tests/backends/test_numpy.py index 0eea3517..e1a389de 100644 --- a/tests/backends/test_numpy.py +++ b/tests/backends/test_numpy.py @@ -486,3 +486,121 @@ def test_count_nonzero_4d(): assert numpy.count_nonzero(v2, axis=1, keepdims=True).tolist() == [[3], [2]] assert numpy.count_nonzero(v2, axis=0).tolist() == [2, 2, 1] assert numpy.count_nonzero(v2, axis=0, keepdims=True).tolist() == [[2, 2, 1]] + + +def test_demotion(): + v1 = vector.array( + { + "x": [10.0, 20.0, 30.0], + "y": [-10.0, 20.0, 30.0], + }, + ) + v2 = vector.array( + { + "x": [10.0, 20.0, 30.0], + "y": [-10.0, 20.0, 30.0], + "z": [5.0, 1.0, 1.0], + }, + ) + v3 = vector.array( + { + "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.array( + { + "x": [20.0, 40.0, 60.0], + "y": [-20.0, 40.0, 60.0], + }, + ) + v2_v3 = vector.array( + { + "x": [20.0, 40.0, 60.0], + "y": [-20.0, 40.0, 60.0], + "z": [10.0, 2.0, 2.0], + }, + ) + v1_v3 = vector.array( + { + "x": [20.0, 40.0, 60.0], + "y": [-20.0, 40.0, 60.0], + }, + ) + + # order should not matter + assert all(v1 + v2 == v1_v2) + assert all(v2 + v1 == v1_v2) + assert all(v1 + v3 == v1_v3) + assert all(v3 + v1 == v1_v3) + assert all(v2 + v3 == v2_v3) + assert all(v3 + v2 == v2_v3) + + v1 = vector.array( + { + "px": [10.0, 20.0, 30.0], + "py": [-10.0, 20.0, 30.0], + }, + ) + v2 = vector.array( + { + "px": [10.0, 20.0, 30.0], + "py": [-10.0, 20.0, 30.0], + "pz": [5.0, 1.0, 1.0], + }, + ) + v3 = vector.array( + { + "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], + }, + ) + + p_v1_v2 = vector.array( + { + "px": [20.0, 40.0, 60.0], + "py": [-20.0, 40.0, 60.0], + }, + ) + p_v2_v3 = vector.array( + { + "px": [20.0, 40.0, 60.0], + "py": [-20.0, 40.0, 60.0], + "pz": [10.0, 2.0, 2.0], + }, + ) + p_v1_v3 = vector.array( + { + "px": [20.0, 40.0, 60.0], + "py": [-20.0, 40.0, 60.0], + }, + ) + + # order should not matter + assert all(v1 + v2 == p_v1_v2) + assert all(v2 + v1 == p_v1_v2) + assert all(v1 + v3 == p_v1_v3) + assert all(v3 + v1 == p_v1_v3) + assert all(v2 + v3 == p_v2_v3) + assert all(v3 + v2 == p_v2_v3) + + v2 = vector.array( + { + "x": [10.0, 20.0, 30.0], + "y": [-10.0, 20.0, 30.0], + "z": [5.0, 1.0, 1.0], + }, + ) + + # momentum + generic = generic + assert all(v1 + v2 == v1_v2) + assert all(v2 + v1 == v1_v2) + assert all(v1 + v3 == v1_v3) + assert all(v3 + v1 == v1_v3) + assert all(v2 + v3 == v2_v3) + assert all(v3 + v2 == v2_v3) diff --git a/tests/backends/test_object.py b/tests/backends/test_object.py index a976611f..c50598c2 100644 --- a/tests/backends/test_object.py +++ b/tests/backends/test_object.py @@ -257,3 +257,41 @@ def test_array_casting(): with pytest.raises(TypeError): vector.obj(x=1, y=False) + + +def test_demotion(): + v1 = vector.obj(x=0.1, y=0.2) + v2 = vector.obj(x=1, y=2, z=3) + v3 = vector.obj(x=10, y=20, z=30, t=40) + + # order should not matter + assert v1 + v2 == vector.obj(x=1.1, y=2.2) + assert v2 + v1 == vector.obj(x=1.1, y=2.2) + assert v1 + v3 == vector.obj(x=10.1, y=20.2) + assert v3 + v1 == vector.obj(x=10.1, y=20.2) + assert v2 + v3 == vector.obj(x=11, y=22, z=33) + assert v3 + v2 == vector.obj(x=11, y=22, z=33) + + v1 = vector.obj(px=0.1, py=0.2) + v2 = vector.obj(px=1, py=2, pz=3) + v3 = vector.obj(px=10, py=20, pz=30, t=40) + + # order should not matter + assert v1 + v2 == vector.obj(px=1.1, py=2.2) + assert v2 + v1 == vector.obj(px=1.1, py=2.2) + assert v1 + v3 == vector.obj(px=10.1, py=20.2) + assert v3 + v1 == vector.obj(px=10.1, py=20.2) + assert v2 + v3 == vector.obj(px=11, py=22, pz=33) + assert v3 + v2 == vector.obj(px=11, py=22, pz=33) + + v1 = vector.obj(px=0.1, py=0.2) + v2 = vector.obj(x=1, y=2, z=3) + v3 = vector.obj(px=10, py=20, pz=30, t=40) + + # momentum + generic = generic + assert v1 + v2 == vector.obj(x=1.1, y=2.2) + assert v2 + v1 == vector.obj(x=1.1, y=2.2) + assert v1 + v3 == vector.obj(px=10.1, py=20.2) + assert v3 + v1 == vector.obj(px=10.1, py=20.2) + assert v2 + v3 == vector.obj(x=11, y=22, z=33) + assert v3 + v2 == vector.obj(x=11, y=22, z=33) diff --git a/tests/test_methods.py b/tests/test_methods.py index a69d6815..f5b5762d 100644 --- a/tests/test_methods.py +++ b/tests/test_methods.py @@ -5,15 +5,180 @@ from __future__ import annotations -import vector +import pytest +import vector +from vector import VectorObject2D, VectorObject3D, VectorObject4D -class CustomVector(vector.VectorObject4D): - pass +awkward = pytest.importorskip("awkward") def test_handler_of(): - object_a = CustomVector.from_xyzt(0.0, 0.0, 0.0, 0.0) - object_b = CustomVector.from_xyzt(1.0, 1.0, 1.0, 1.0) + object_a = VectorObject4D.from_xyzt(0.0, 0.0, 0.0, 0.0) + object_b = VectorObject4D.from_xyzt(1.0, 1.0, 1.0, 1.0) + protocol = vector._methods._handler_of(object_a, object_b) + assert protocol == object_a + + object_a = VectorObject3D.from_xyz(0.0, 0.0, 0.0) + object_b = VectorObject4D.from_xyzt(1.0, 1.0, 1.0, 1.0) + protocol = vector._methods._handler_of(object_a, object_b) + assert protocol == object_a + + object_a = VectorObject4D.from_xyzt(0.0, 0.0, 0.0, 0.0) + object_b = VectorObject3D.from_xyz(1.0, 1.0, 1.0) + protocol = vector._methods._handler_of(object_a, object_b) + assert protocol == object_b + + object_a = VectorObject2D.from_xy(0.0, 0.0) + object_b = VectorObject4D.from_xyzt(1.0, 1.0, 1.0, 1.0) + protocol = vector._methods._handler_of(object_a, object_b) + assert protocol == object_a + + object_a = VectorObject4D.from_xyzt(0.0, 0.0, 0.0, 0.0) + object_b = VectorObject2D.from_xy(1.0, 1.0) + protocol = vector._methods._handler_of(object_a, object_b) + assert protocol == object_b + + object_a = VectorObject2D.from_xy(0.0, 0.0) + object_b = VectorObject3D.from_xyz(1.0, 1.0, 1.0) protocol = vector._methods._handler_of(object_a, object_b) assert protocol == object_a + + object_a = VectorObject3D.from_xyz(0.0, 0.0, 0.0) + object_b = VectorObject2D.from_xy(1.0, 1.0) + protocol = vector._methods._handler_of(object_a, object_b) + assert protocol == object_b + + awkward_a = vector.zip( + { + "x": [10.0, 20.0, 30.0], + "y": [-10.0, 20.0, 30.0], + "z": [5.0, 10.0, 15.0], + "t": [16.0, 31.0, 46.0], + }, + ) + object_b = VectorObject2D.from_xy(1.0, 1.0) + protocol = vector._methods._handler_of(awkward_a, object_b) + # chooses awkward backend and converts the vector to 2D + assert all(protocol == awkward_a.to_Vector2D()) + + awkward_a = vector.zip( + { + "x": [10.0, 20.0, 30.0], + "y": [-10.0, 20.0, 30.0], + }, + ) + object_b = VectorObject4D.from_xyzt(1.0, 1.0, 1.0, 1.0) + protocol = vector._methods._handler_of(object_b, awkward_a) + # chooses awkward backend and the vector is already of the + # lower dimension + assert all(protocol == awkward_a) + + awkward_a = vector.zip( + { + "x": [10.0, 20.0, 30.0], + "y": [-10.0, 20.0, 30.0], + }, + ) + awkward_b = 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], + }, + ) + object_b = VectorObject4D.from_xyzt(1.0, 1.0, 1.0, 1.0) + protocol = vector._methods._handler_of(object_b, awkward_a, awkward_b) + # chooses awkward backend and the 2D awkward vector + # (first encountered awkward vector) + assert all(protocol == awkward_a) + + awkward_a = vector.zip( + { + "x": [10.0, 20.0, 30.0], + "y": [-10.0, 20.0, 30.0], + "z": [-10.0, 20.0, 30.0], + }, + ) + awkward_b = 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], + }, + ) + object_b = VectorObject2D.from_xy(1.0, 1.0) + protocol = vector._methods._handler_of(awkward_b, object_b, awkward_a) + # chooses awkward backend and converts awkward_b to 2D + # (first encountered awkward vector) + assert all(protocol == awkward_b.to_Vector2D()) + + awkward_a = vector.zip( + { + "x": [10.0, 20.0, 30.0], + "y": [-10.0, 20.0, 30.0], + "z": [5.0, 1.0, 1.0], + }, + ) + awkward_b = 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], + }, + ) + object_b = VectorObject2D.from_xy(1.0, 1.0) + protocol = vector._methods._handler_of(object_b, awkward_a, awkward_b) + # chooses awkward backend and converts the vector to 2D + # (the first awkward vector encountered is used as the base) + assert all(protocol == awkward_a.to_Vector2D()) + + numpy_a = 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_b = 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], + }, + ) + object_b = VectorObject2D.from_xy(1.0, 1.0) + protocol = vector._methods._handler_of(object_b, numpy_a, awkward_b) + # chooses awkward backend and converts the vector to 2D + assert all(protocol == awkward_b.to_Vector2D()) + + awkward_a = vector.zip( + { + "x": [10.0, 20.0, 30.0], + "y": [-10.0, 20.0, 30.0], + "z": [5.0, 1.0, 1.0], + }, + ) + numpy_a = vector.array( + { + "x": [1.1, 1.2, 1.3, 1.4, 1.5], + "y": [2.1, 2.2, 2.3, 2.4, 2.5], + } + ) + awkward_b = 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], + }, + ) + object_b = VectorObject3D.from_xyz(1.0, 1.0, 1.0) + protocol = vector._methods._handler_of(object_b, awkward_a, awkward_b, numpy_a) + # chooses awkward backend and converts the vector to 2D + # (the first awkward vector encountered is used as the base) + assert all(protocol == awkward_a.to_Vector2D())