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.

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/scikit-hep/vector/pull/438?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scikit-hep) Report
Attention: Patch coverage is `50.00000%` with `1 lines` in your changes are missing coverage. Please review.
> Project coverage is 86.27%. Comparing base [(`b907615`)](https://app.codecov.io/gh/scikit-hep/vector/commit/b907615f5035b4bdad2335f2c8d319fff8fd3d9f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scikit-hep) to head [(`a587a75`)](https://app.codecov.io/gh/scikit-hep/vector/pull/438?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scikit-hep).

| [Files](https://app.codecov.io/gh/scikit-hep/vector/pull/438?dropdown=coverage&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/backends/awkward.py](https://app.codecov.io/gh/scikit-hep/vector/pull/438?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scikit-hep#diff-c3JjL3ZlY3Rvci9iYWNrZW5kcy9hd2t3YXJkLnB5) | 50.00% | [1 Missing :warning: ](https://app.codecov.io/gh/scikit-hep/vector/pull/438?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     #438   +/-   ##
=======================================
  Coverage   86.27%   86.27%           
=======================================
  Files          96       96           
  Lines       11139    11139           
=======================================
  Hits         9610     9610           
  Misses       1529     1529           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/scikit-hep/vector/pull/438?dropdown=coverage&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: Tagging @agoose77 here because I cannot request a review from you for some reason 🙂 
Saransh-cpp: Thanks! I noticed that vector in general does not test subclassing the awkward mixins, which is why a few bugs popped up as soon as coffea decided to inherit them. I'll add tests for this.

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"]