feat: add constructors for VectorObject4D and MomentumObject4D (#232

) * feat: useful constructor * Reprs + PEP * Some tests * Constructor for MomentumObject2D * Update docs * Create overloads and a constructor for `VectorObject3D` * Docs for the new constructor * Reprs + docs * Make `VectorObject2D` constructor more general * Add `MomentumObject2D` overloads * Modify constructor for `MomentumObject3D` + typing overloads + docs * Add tests * Make TypeError more clear * Make TypeError dynamic * Make TypeError dynamic * Update the existing code with azimuthal= longitudinal= * More azimuthal= longitudinal= * Add constructors to `VectorObject4D` and `MomentumObject4D` * Add type overloads * Rebase --------- Co-authored-by: Henry Schreiner <henryschreineriii@gmail.com>
scikit-hep · Feb 17, 2023 · 9310cd4 · 9310cd4
1 parent bcd6ddb
commit 9310cd4
Show file tree

Hide file tree

Showing 34 changed files with 1,738 additions and 426 deletions.
diff --git a/src/vector/_methods.py b/src/vector/_methods.py
diff --git a/src/vector/backends/numpy.py b/src/vector/backends/numpy.py
@@ -166,11 +166,11 @@ def _getitem(
                 *(out[x] for x in _coordinate_class_to_names[_ttype(array)])  # type: ignore[arg-type]
             )
         if temporal is not None:
-            return array.ObjectClass(azimuthal, longitudinal, temporal)  # type: ignore[call-arg, arg-type, return-value]
+            return array.ObjectClass(azimuthal=azimuthal, longitudinal=longitudinal, temporal=temporal)  # type: ignore[arg-type, return-value]
         elif longitudinal is not None:
-            return array.ObjectClass(azimuthal, longitudinal)  # type: ignore[call-arg, arg-type, return-value]
+            return array.ObjectClass(azimuthal=azimuthal, longitudinal=longitudinal)  # type: ignore[arg-type, return-value]
         elif azimuthal is not None:
-            return array.ObjectClass(azimuthal)  # type: ignore[call-arg, return-value]
+            return array.ObjectClass(azimuthal=azimuthal)  # type: ignore[return-value]
         else:
             return array.ObjectClass(*out.view(numpy.ndarray))  # type: ignore[misc, return-value]
 

diff --git a/src/vector/backends/object.py b/src/vector/backends/object.py
diff --git a/tests/compute/lorentz/test_beta.py b/tests/compute/lorentz/test_beta.py
@@ -12,107 +12,107 @@
 
 def test_xy_z_t():
     vec = vector.backends.object.VectorObject4D(
-        vector.backends.object.AzimuthalObjectXY(3, 4),
-        vector.backends.object.LongitudinalObjectZ(10),
-        vector.backends.object.TemporalObjectT(20),
+        azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4),
+        longitudinal=vector.backends.object.LongitudinalObjectZ(10),
+        temporal=vector.backends.object.TemporalObjectT(20),
     )
     assert vec.beta == pytest.approx(0.5590169943749475)
 
 
 def test_xy_z_tau():
     vec = vector.backends.object.VectorObject4D(
-        vector.backends.object.AzimuthalObjectXY(3, 4),
-        vector.backends.object.LongitudinalObjectZ(10),
-        vector.backends.object.TemporalObjectTau(16.583123951777),
+        azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4),
+        longitudinal=vector.backends.object.LongitudinalObjectZ(10),
+        temporal=vector.backends.object.TemporalObjectTau(16.583123951777),
     )
     assert vec.beta == pytest.approx(0.5590169943749475)
 
 
 def test_xy_theta_t():
     vec = vector.backends.object.VectorObject4D(
-        vector.backends.object.AzimuthalObjectXY(3, 4),
-        vector.backends.object.LongitudinalObjectTheta(0.4636476090008061),
-        vector.backends.object.TemporalObjectT(20),
+        azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4),
+        longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061),
+        temporal=vector.backends.object.TemporalObjectT(20),
     )
     assert vec.beta == pytest.approx(0.5590169943749475)
 
 
 def test_xy_theta_tau():
     vec = vector.backends.object.VectorObject4D(
-        vector.backends.object.AzimuthalObjectXY(3, 4),
-        vector.backends.object.LongitudinalObjectTheta(0.4636476090008061),
-        vector.backends.object.TemporalObjectTau(16.583123951777),
+        azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4),
+        longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061),
+        temporal=vector.backends.object.TemporalObjectTau(16.583123951777),
     )
     assert vec.beta == pytest.approx(0.5590169943749475)
 
 
 def test_xy_eta_t():
     vec = vector.backends.object.VectorObject4D(
-        vector.backends.object.AzimuthalObjectXY(3, 4),
-        vector.backends.object.LongitudinalObjectEta(1.4436354751788103),
-        vector.backends.object.TemporalObjectT(20),
+        azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4),
+        longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103),
+        temporal=vector.backends.object.TemporalObjectT(20),
     )
     assert vec.beta == pytest.approx(0.5590169943749475)
 
 
 def test_xy_eta_tau():
     vec = vector.backends.object.VectorObject4D(
-        vector.backends.object.AzimuthalObjectXY(3, 4),
-        vector.backends.object.LongitudinalObjectEta(1.4436354751788103),
-        vector.backends.object.TemporalObjectTau(16.583123951777),
+        azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4),
+        longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103),
+        temporal=vector.backends.object.TemporalObjectTau(16.583123951777),
     )
     assert vec.beta == pytest.approx(0.5590169943749475)
 
 
 def test_rhophi_z_t():
     vec = vector.backends.object.VectorObject4D(
-        vector.backends.object.AzimuthalObjectRhoPhi(5, 0),
-        vector.backends.object.LongitudinalObjectZ(10),
-        vector.backends.object.TemporalObjectT(20),
+        azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0),
+        longitudinal=vector.backends.object.LongitudinalObjectZ(10),
+        temporal=vector.backends.object.TemporalObjectT(20),
     )
     assert vec.beta == pytest.approx(0.5590169943749475)
 
 
 def test_rhophi_z_tau():
     vec = vector.backends.object.VectorObject4D(
-        vector.backends.object.AzimuthalObjectRhoPhi(5, 0),
-        vector.backends.object.LongitudinalObjectZ(10),
-        vector.backends.object.TemporalObjectTau(16.583123951777),
+        azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0),
+        longitudinal=vector.backends.object.LongitudinalObjectZ(10),
+        temporal=vector.backends.object.TemporalObjectTau(16.583123951777),
     )
     assert vec.beta == pytest.approx(0.5590169943749475)
 
 
 def test_rhophi_theta_t():
     vec = vector.backends.object.VectorObject4D(
-        vector.backends.object.AzimuthalObjectRhoPhi(5, 0),
-        vector.backends.object.LongitudinalObjectTheta(0.4636476090008061),
-        vector.backends.object.TemporalObjectT(20),
+        azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0),
+        longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061),
+        temporal=vector.backends.object.TemporalObjectT(20),
     )
     assert vec.beta == pytest.approx(0.5590169943749475)
 
 
 def test_rhophi_theta_tau():
     vec = vector.backends.object.VectorObject4D(
-        vector.backends.object.AzimuthalObjectRhoPhi(5, 0),
-        vector.backends.object.LongitudinalObjectTheta(0.4636476090008061),
-        vector.backends.object.TemporalObjectTau(16.583123951777),
+        azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0),
+        longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061),
+        temporal=vector.backends.object.TemporalObjectTau(16.583123951777),
     )
     assert vec.beta == pytest.approx(0.5590169943749475)
 
 
 def test_rhophi_eta_t():
     vec = vector.backends.object.VectorObject4D(
-        vector.backends.object.AzimuthalObjectRhoPhi(5, 0),
-        vector.backends.object.LongitudinalObjectEta(1.4436354751788103),
-        vector.backends.object.TemporalObjectT(20),
+        azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0),
+        longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103),
+        temporal=vector.backends.object.TemporalObjectT(20),
     )
     assert vec.beta == pytest.approx(0.5590169943749475)
 
 
 def test_rhophi_eta_tau():
     vec = vector.backends.object.VectorObject4D(
-        vector.backends.object.AzimuthalObjectRhoPhi(5, 0),
-        vector.backends.object.LongitudinalObjectEta(1.4436354751788103),
-        vector.backends.object.TemporalObjectTau(16.583123951777),
+        azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0),
+        longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103),
+        temporal=vector.backends.object.TemporalObjectTau(16.583123951777),
     )
     assert vec.beta == pytest.approx(0.5590169943749475)
diff --git a/tests/compute/lorentz/test_boostX_beta.py b/tests/compute/lorentz/test_boostX_beta.py
@@ -13,9 +13,9 @@
 
 def test():
     vec = vector.backends.object.VectorObject4D(
-        vector.backends.object.AzimuthalObjectXY(3, 2),
-        vector.backends.object.LongitudinalObjectZ(1),
-        vector.backends.object.TemporalObjectT(4),
+        azimuthal=vector.backends.object.AzimuthalObjectXY(3, 2),
+        longitudinal=vector.backends.object.LongitudinalObjectZ(1),
+        temporal=vector.backends.object.TemporalObjectT(4),
     )
     out = vec.boostX(beta=-0.9428090415820634)
     assert out.x == pytest.approx(-2.313708498984761)

diff --git a/tests/compute/lorentz/test_boostX_gamma.py b/tests/compute/lorentz/test_boostX_gamma.py
@@ -13,9 +13,9 @@
 
 def test():
     vec = vector.backends.object.VectorObject4D(
-        vector.backends.object.AzimuthalObjectXY(3, 2),
-        vector.backends.object.LongitudinalObjectZ(1),
-        vector.backends.object.TemporalObjectT(4),
+        azimuthal=vector.backends.object.AzimuthalObjectXY(3, 2),
+        longitudinal=vector.backends.object.LongitudinalObjectZ(1),
+        temporal=vector.backends.object.TemporalObjectT(4),
     )
     out = vec.boostX(gamma=-3)
     assert out.x == pytest.approx(-2.313708498984761)

diff --git a/tests/compute/lorentz/test_boostY_beta.py b/tests/compute/lorentz/test_boostY_beta.py
@@ -13,9 +13,9 @@
 
 def test():
     vec = vector.backends.object.VectorObject4D(
-        vector.backends.object.AzimuthalObjectXY(2, 3),
-        vector.backends.object.LongitudinalObjectZ(1),
-        vector.backends.object.TemporalObjectT(4),
+        azimuthal=vector.backends.object.AzimuthalObjectXY(2, 3),
+        longitudinal=vector.backends.object.LongitudinalObjectZ(1),
+        temporal=vector.backends.object.TemporalObjectT(4),
     )
     out = vec.boostY(beta=-0.9428090415820634)
     assert out.x == pytest.approx(2)

diff --git a/tests/compute/lorentz/test_boostY_gamma.py b/tests/compute/lorentz/test_boostY_gamma.py
@@ -13,9 +13,9 @@
 
 def test():
     vec = vector.backends.object.VectorObject4D(
-        vector.backends.object.AzimuthalObjectXY(2, 3),
-        vector.backends.object.LongitudinalObjectZ(1),
-        vector.backends.object.TemporalObjectT(4),
+        azimuthal=vector.backends.object.AzimuthalObjectXY(2, 3),
+        longitudinal=vector.backends.object.LongitudinalObjectZ(1),
+        temporal=vector.backends.object.TemporalObjectT(4),
     )
     out = vec.boostY(gamma=-3)
     assert out.x == pytest.approx(2)

diff --git a/tests/compute/lorentz/test_boostZ_beta.py b/tests/compute/lorentz/test_boostZ_beta.py
@@ -13,9 +13,9 @@
 
 def test():
     vec = vector.backends.object.VectorObject4D(
-        vector.backends.object.AzimuthalObjectXY(1, 2),
-        vector.backends.object.LongitudinalObjectZ(3),
-        vector.backends.object.TemporalObjectT(4),
+        azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2),
+        longitudinal=vector.backends.object.LongitudinalObjectZ(3),
+        temporal=vector.backends.object.TemporalObjectT(4),
     )
     out = vec.boostZ(beta=-0.9428090415820634)
     assert out.x == pytest.approx(1)

diff --git a/tests/compute/lorentz/test_boostZ_gamma.py b/tests/compute/lorentz/test_boostZ_gamma.py
@@ -13,9 +13,9 @@
 
 def test():
     vec = vector.backends.object.VectorObject4D(
-        vector.backends.object.AzimuthalObjectXY(1, 2),
-        vector.backends.object.LongitudinalObjectZ(3),
-        vector.backends.object.TemporalObjectT(4),
+        azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2),
+        longitudinal=vector.backends.object.LongitudinalObjectZ(3),
+        temporal=vector.backends.object.TemporalObjectT(4),
     )
     out = vec.boostZ(gamma=-3)
     assert out.x == pytest.approx(1)

diff --git a/tests/compute/lorentz/test_boost_beta3.py b/tests/compute/lorentz/test_boost_beta3.py
@@ -13,9 +13,9 @@
 
 def test():
     vec = vector.backends.object.VectorObject4D(
-        vector.backends.object.AzimuthalObjectXY(1, 2),
-        vector.backends.object.LongitudinalObjectZ(3),
-        vector.backends.object.TemporalObjectT(4),
+        azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2),
+        longitudinal=vector.backends.object.LongitudinalObjectZ(3),
+        temporal=vector.backends.object.TemporalObjectT(4),
     )
     beta = vector.backends.object.VectorObject3D(
         azimuthal=vector.backends.object.AzimuthalObjectXY(5 / 15, 6 / 15),

diff --git a/tests/compute/lorentz/test_boost_p4.py b/tests/compute/lorentz/test_boost_p4.py
@@ -13,14 +13,14 @@
 
 def test():
     vec = vector.backends.object.VectorObject4D(
-        vector.backends.object.AzimuthalObjectXY(1, 2),
-        vector.backends.object.LongitudinalObjectZ(3),
-        vector.backends.object.TemporalObjectT(4),
+        azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2),
+        longitudinal=vector.backends.object.LongitudinalObjectZ(3),
+        temporal=vector.backends.object.TemporalObjectT(4),
     )
     p4 = vector.backends.object.VectorObject4D(
-        vector.backends.object.AzimuthalObjectXY(5, 6),
-        vector.backends.object.LongitudinalObjectZ(7),
-        vector.backends.object.TemporalObjectT(15),
+        azimuthal=vector.backends.object.AzimuthalObjectXY(5, 6),
+        longitudinal=vector.backends.object.LongitudinalObjectZ(7),
+        temporal=vector.backends.object.TemporalObjectT(15),
     )
     out = vec.boost_p4(p4)
     assert out.x == pytest.approx(3.5537720741941676)