Repository: scikit-hep/vector

Issue:
Preserve field alias under addition
A user [asked a question](https://stackoverflow.com/questions/72906846/can-i-recalculate-the-energy-of-an-awkward-array-of-vectors-by-declaring-a-new-m/72929486?noredirect=1#comment128831169_72929486) about using Vector, and in answering I realised that field aliases are not preserved under addition. I assume this generalises to all operations given [the source that I think is responsible](
https://github.com/scikit-hep/vector/blob/4e756fe2cd8a3430260e1cf4cf008874cfc1bb3c/src/vector/backends/awkward.py#L886-L888) currently hard-codes the name.

I think we should try to preserve the existing aliasing where possible. At present, users need to know that a particular alias is "canonical", and that their array could change fields under any operation.

Comments:

Patch:
diff --git a/src/vector/backends/awkward.py b/src/vector/backends/awkward.py
index eefb1e1..f29bc03 100644
--- a/src/vector/backends/awkward.py
+++ b/src/vector/backends/awkward.py
@@ -673,7 +673,13 @@ class VectorAwkward:
             fields = ak.fields(self)
             if num_vecargs == 1:
                 for name in fields:
-                    if name not in ("x", "y", "rho", "phi"):
+                    if name not in (
+                        "x",
+                        "y",
+                        "rho",
+                        "pt",
+                        "phi",
+                    ):
                         names.append(name)
                         arrays.append(self[name])
 
@@ -720,12 +726,20 @@ class VectorAwkward:
                         "x",
                         "y",
                         "rho",
+                        "pt",
                         "phi",
                         "z",
+                        "pz",
                         "theta",
                         "eta",
                         "t",
                         "tau",
+                        "m",
+                        "M",
+                        "mass",
+                        "e",
+                        "E",
+                        "energy",
                     ):
                         names.append(name)
                         arrays.append(self[name])
@@ -774,7 +788,17 @@ class VectorAwkward:
             fields = ak.fields(self)
             if num_vecargs == 1:
                 for name in fields:
-                    if name not in ("x", "y", "rho", "phi", "z", "theta", "eta"):
+                    if name not in (
+                        "x",
+                        "y",
+                        "rho",
+                        "pt",
+                        "phi",
+                        "z",
+                        "pz",
+                        "theta",
+                        "eta",
+                    ):
                         names.append(name)
                         arrays.append(self[name])
 
@@ -831,12 +855,20 @@ class VectorAwkward:
                         "x",
                         "y",
                         "rho",
+                        "pt",
                         "phi",
                         "z",
+                        "pz",
                         "theta",
                         "eta",
                         "t",
                         "tau",
+                        "m",
+                        "M",
+                        "mass",
+                        "e",
+                        "E",
+                        "energy",
                     ):
                         names.append(name)
                         arrays.append(self[name])
@@ -897,12 +929,20 @@ class VectorAwkward:
                         "x",
                         "y",
                         "rho",
+                        "pt",
                         "phi",
                         "z",
+                        "pz",
                         "theta",
                         "eta",
                         "t",
                         "tau",
+                        "m",
+                        "M",
+                        "mass",
+                        "e",
+                        "E",
+                        "energy",
                     ):
                         names.append(name)
                         arrays.append(self[name])


Test patch:
diff --git a/tests/backends/test_awkward.py b/tests/backends/test_awkward.py
index 85caa2c..0ad09c4 100644
--- a/tests/backends/test_awkward.py
+++ b/tests/backends/test_awkward.py
@@ -6,7 +6,9 @@
 from __future__ import annotations
 
 import importlib.metadata
+import numbers
 
+import numpy as np
 import packaging.version
 import pytest
 
@@ -895,3 +897,39 @@ def test_momentum_preservation():
     # 4D + 3D.like(4D) = 4D
     assert isinstance(v3 + v2.like(v3), MomentumAwkward4D)
     assert isinstance(v2.like(v3) + v3, MomentumAwkward4D)
+
+
+def test_subclass_fields():
+    @ak.mixin_class(vector.backends.awkward.behavior)
+    class TwoVector(MomentumAwkward2D):
+        pass
+
+    @ak.mixin_class(vector.backends.awkward.behavior)
+    class ThreeVector(MomentumAwkward3D):
+        pass
+
+    @ak.mixin_class(vector.backends.awkward.behavior)
+    class LorentzVector(MomentumAwkward4D):
+        @ak.mixin_class_method(np.divide, {numbers.Number})
+        def divide(self, factor):
+            return self.scale(1 / factor)
+
+    LorentzVectorArray.ProjectionClass2D = TwoVectorArray  # noqa: F821
+    LorentzVectorArray.ProjectionClass3D = ThreeVectorArray  # noqa: F821
+    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,
+    )
+
+    assert vec.like(vector.obj(x=1, y=2)).fields == ["rho", "phi"]
+    assert vec.like(vector.obj(x=1, y=2, z=3)).fields == ["rho", "phi", "eta"]
+    assert (vec / 2).fields == ["rho", "phi", "eta", "t"]