Repository: scikit-hep/vector

Issue:
Record type is not correct for downcasting awkward vectors
### Vector Version

1.1.1.post1

### Python Version

3.10.12

### OS / Environment

OS/x

### Describe the bug

```python
import awkward as ak
import vector.backends.awkward

a = ak.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],
    },
    with_name="Momentum4D",
    behavior=vector.backends.awkward.behavior,
)
b = ak.zip(
    {
        "x": [-10.0, 20.0, -30.0],
        "y": [-10.0, -20.0, 30.0],
        "z": [5.0, -10.0, 15.0],
    },
    with_name="Momentum3D",
    behavior=vector.backends.awkward.behavior,
)
c = ak.zip(
    {"x": [-10.0, 13.0, 15.0], "y": [12.0, -4.0, 41.0]},
    with_name="Momentum2D",
    behavior=vector.backends.awkward.behavior,
)

print("a+b")
print(a+b)
print(ak.parameters(a+b))
print("b+a")
print(b+a)
print(ak.parameters(b+a))
```

returns 
```
a+b
[{x: 0, y: -20, z: 10}, {x: 40, y: 0, z: 0}, {x: 0, y: 60, z: 30}]
{'__record__': 'Momentum4D'}
b+a
[{x: 0, y: -20, z: 10}, {x: 40, y: 0, z: 0}, {x: 0, y: 60, z: 30}]
{'__record__': 'Momentum3D'}
```
So, although the fields available are the ones expected, the record type is incorrect in the first case. Similar issues are there for `a+c` and `b+c`.

### Any additional but relevant log output

_No response_

Discussion:
codecov-commenter: ## [Codecov](https://app.codecov.io/gh/scikit-hep/vector/pull/413?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scikit-hep) Report
Attention: `8 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`76cc674`)](https://app.codecov.io/gh/scikit-hep/vector/commit/76cc674e42ffc0121841dac2ab891b6bdc3c5b1c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scikit-hep) 83.21% compared to head [(`6afedf3`)](https://app.codecov.io/gh/scikit-hep/vector/pull/413?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scikit-hep) 83.43%.
> Report is 1 commits behind head on main.

| [Files](https://app.codecov.io/gh/scikit-hep/vector/pull/413?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scikit-hep) | Patch % | Lines |
|---|---|---|
| [src/vector/\_methods.py](https://app.codecov.io/gh/scikit-hep/vector/pull/413?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scikit-hep#diff-c3JjL3ZlY3Rvci9fbWV0aG9kcy5weQ==) | 78.37% | [8 Missing :warning: ](https://app.codecov.io/gh/scikit-hep/vector/pull/413?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scikit-hep) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #413      +/-   ##
==========================================
+ Coverage   83.21%   83.43%   +0.21%     
==========================================
  Files          96       96              
  Lines       11429    11449      +20     
==========================================
+ Hits         9511     9552      +41     
+ Misses       1918     1897      -21     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/scikit-hep/vector/pull/413?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scikit-hep).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scikit-hep).

Saransh-cpp: @nsmith- could you check if this is working well? The new tests and the old tests are passing.

The code is a bit dirty; I will refactor it tomorrow and make this ready for review.
Saransh-cpp: I think this should be considered a breaking change in vector. @henryiii, how should the version be incremented once this is in? Should it go to `v2.0.0` or just `v1.2.0`?
nsmith-: I checked out the branch, things look ok interactively, and the problematic coffea tests pass without the explicit cast now. (though some other tests fail, but I think this is on coffea side... checking)
nsmith-: Actually no, there is some issue now with the fact that we are subclassing the vector array method mixins:
```
>       boosted = a.boost(-a.boostvec)

tests/test_nanoevents_vector.py:313:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
.env/lib/python3.10/site-packages/vector/_methods.py:3578: in boost
    return boost_beta3.dispatch(self, booster)
.env/lib/python3.10/site-packages/vector/_compute/lorentz/boost_beta3.py:393: in dispatch
    return _handler_of(v1, v2)._wrap_result(
.env/lib/python3.10/site-packages/vector/_methods.py:4192: in _handler_of
    handler = _demote_handler_vector(
.env/lib/python3.10/site-packages/vector/_methods.py:4144: in _demote_handler_vector
    if len({_handler_priority.index(obj.__module__) for obj in objects}) != 1:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

.0 = <tuple_iterator object at 0x12d673e20>

>   if len({_handler_priority.index(obj.__module__) for obj in objects}) != 1:
E   ValueError: 'coffea.nanoevents.methods.vector' is not in list

.env/lib/python3.10/site-packages/vector/_methods.py:4144: ValueError
```

(the `.boostvec` property is just calling `.to_beta3()`)
henryiii: IMO, I'd be fine with calling this v1.2.0, since to me this is a bug fix - a rather impactful one, but still a bug fix. Before it would make a Momentum4D vector with only three dimensions, now the dimensions line up.

But also feel free to do what seems best to you.
lgray: Ah - this is great! Just seeing the PR now. :-)
nsmith-: There is still an issue with `(a-b)` vs. `-(b-a)`:
```python
import awkward as ak
from coffea.nanoevents.methods import vector

a = ak.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],
    },
    with_name="LorentzVector",
    behavior=vector.behavior,
)
b = ak.zip(
    {
        "x": [-10.0, 20.0, -30.0],
        "y": [-10.0, -20.0, 30.0],
        "z": [5.0, -10.0, 15.0],
    },
    with_name="ThreeVector",
    behavior=vector.behavior,
)
c = ak.zip(
    {"rho": [-10.0, 13.0, 15.0], "phi": [1.22, -1.0, 1.0]},
    with_name="TwoVector",
    behavior=vector.behavior,
)

print("a-b")
print(a-b)
print(ak.parameters(a-b))
print("-(b-a)")
print(-(b-a))
print(ak.parameters(-(b-a)))
```
produces
```
a-b
[{x: 20, y: 0, z: 0, t: 16}, {x: 0, ...}, {x: 60, y: 0, z: 0, t: 46}]
{'__record__': 'LorentzVector'}
-(b-a)
[{x: 20, y: -0, z: -0}, {x: -0, y: 40, z: 20}, {x: 60, y: -0, z: -0}]
{'__record__': 'ThreeVector'}
```
They have a type matching their content, but the type demotion isn't working (there should be no `t` in the first case)
Though `a+b` and `b+a` seem to be working OK.
Saransh-cpp: I think this is on the coffea side (on the `use_scikithep_vector` branch)? I might be wrong -

https://github.com/CoffeaTeam/coffea/blob/894496fc9fe8d3be79263bf7b03e72c18aaf13de/src/coffea/nanoevents/methods/vector.py#L507-L545
nsmith-: @Saransh-cpp you're right, it's to do with how we are building the subclasses' behavior dispatch tables. I've made a fix that now passes tests on the coffea PR with this PR.
lgray: Nice - Since we're deprecating nanoevents.vector in this next release of coffea. I think a reasonable timeline for getting a new version of vector with this is and merging the corresponding coffea PR is the beginning of March. Any faster than that and people may get spooked.
lgray: @Saransh-cpp concerning the version number for the next release. This isn't a major design-changing reimplementation of the package, but it does change behavior significantly, so `1.2.0` makes since IMO.
Saransh-cpp: Thanks for the reviews! Given that we don't have any release schedule for vector, I'll just merge this in and create a new release. 

Patch:
diff --git a/README.md b/README.md
index e9e0b1d..613d711 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 83d6df5..3b2aa3b 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 df41fb0..7511cda 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 4b50f0a..71d5b72 100644
--- a/src/vector/backends/awkward.py
+++ b/src/vector/backends/awkward.py
@@ -575,19 +575,21 @@ class TemporalAwkwardTau(TemporalAwkward, TemporalTau):
 
 
 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))
 


Test patch:
diff --git a/tests/backends/test_awkward.py b/tests/backends/test_awkward.py
index 4786fbb..9797be1 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 0eea351..e1a389d 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 a976611..c50598c 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 a69d681..f5b5762 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())