wrapper around lib

scikit-hep · Mar 26, 2024 · 0af4ee0 · 0af4ee0
1 parent 0ce85d4
commit 0af4ee0
Show file tree

Hide file tree

Showing 58 changed files with 262 additions and 57 deletions.
diff --git a/src/vector/_compute/lorentz/Et.py b/src/vector/_compute/lorentz/Et.py
@@ -17,6 +17,7 @@
 import numpy
 
 from vector._compute.lorentz import Et2, t
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -115,7 +116,7 @@ def dispatch(v: typing.Any) -> typing.Any:
         return v._wrap_result(
             _flavor_of(v),
             function(
-                v.lib,
+                Lib(v.lib),
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,
                 *v.temporal.elements,

diff --git a/src/vector/_compute/lorentz/Et2.py b/src/vector/_compute/lorentz/Et2.py
@@ -17,6 +17,7 @@
 import numpy
 
 from vector._compute.lorentz import t
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -117,7 +118,7 @@ def dispatch(v: typing.Any) -> typing.Any:
         return v._wrap_result(
             _flavor_of(v),
             function(
-                v.lib,
+                Lib(v.lib),
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,
                 *v.temporal.elements,

diff --git a/src/vector/_compute/lorentz/Mt.py b/src/vector/_compute/lorentz/Mt.py
@@ -17,6 +17,7 @@
 import numpy
 
 from vector._compute.lorentz import Mt2
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -111,7 +112,7 @@ def dispatch(v: typing.Any) -> typing.Any:
         return v._wrap_result(
             _flavor_of(v),
             function(
-                v.lib,
+                Lib(v.lib),
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,
                 *v.temporal.elements,

diff --git a/src/vector/_compute/lorentz/Mt2.py b/src/vector/_compute/lorentz/Mt2.py
@@ -18,6 +18,7 @@
 
 from vector._compute.lorentz import tau2
 from vector._compute.spatial import z
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -112,7 +113,7 @@ def dispatch(v: typing.Any) -> typing.Any:
         return v._wrap_result(
             _flavor_of(v),
             function(
-                v.lib,
+                Lib(v.lib),
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,
                 *v.temporal.elements,

diff --git a/src/vector/_compute/lorentz/beta.py b/src/vector/_compute/lorentz/beta.py
@@ -19,6 +19,7 @@
 
 from vector._compute.lorentz import t
 from vector._compute.spatial import mag
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -154,7 +155,7 @@ def dispatch(v: typing.Any) -> typing.Any:
         return v._wrap_result(
             _flavor_of(v),
             function(
-                v.lib,
+                Lib(v.lib),
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,
                 *v.temporal.elements,

diff --git a/src/vector/_compute/lorentz/boostX_gamma.py b/src/vector/_compute/lorentz/boostX_gamma.py
@@ -18,6 +18,7 @@
 from vector._compute.lorentz import t
 from vector._compute.planar import x, y
 from vector._compute.spatial import z
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -244,7 +245,7 @@ def dispatch(gamma: typing.Any, v: typing.Any) -> typing.Any:
         return v._wrap_result(
             _flavor_of(v),
             function(
-                v.lib,
+                Lib(v.lib),
                 gamma,
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,

diff --git a/src/vector/_compute/lorentz/boostY_gamma.py b/src/vector/_compute/lorentz/boostY_gamma.py
@@ -18,6 +18,7 @@
 from vector._compute.lorentz import t
 from vector._compute.planar import x, y
 from vector._compute.spatial import z
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -244,7 +245,7 @@ def dispatch(gamma: typing.Any, v: typing.Any) -> typing.Any:
         return v._wrap_result(
             _flavor_of(v),
             function(
-                v.lib,
+                Lib(v.lib),
                 gamma,
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,

diff --git a/src/vector/_compute/lorentz/boostZ_gamma.py b/src/vector/_compute/lorentz/boostZ_gamma.py
@@ -17,6 +17,7 @@
 
 from vector._compute.lorentz import t
 from vector._compute.spatial import z
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -219,7 +220,7 @@ def dispatch(gamma: typing.Any, v: typing.Any) -> typing.Any:
         return v._wrap_result(
             _flavor_of(v),
             function(
-                v.lib,
+                Lib(v.lib),
                 gamma,
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,

diff --git a/src/vector/_compute/lorentz/boost_beta3.py b/src/vector/_compute/lorentz/boost_beta3.py
@@ -24,6 +24,7 @@
 from vector._compute.lorentz import transform4D
 from vector._compute.planar import x, y
 from vector._compute.spatial import z
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -394,7 +395,7 @@ def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any:
         return _handler_of(v1, v2)._wrap_result(
             _flavor_of(v1, v2),
             function(
-                _lib_of(v1, v2),
+                Lib(_lib_of(v1, v2)),
                 *v1.azimuthal.elements,
                 *v1.longitudinal.elements,
                 *v1.temporal.elements,

diff --git a/src/vector/_compute/lorentz/boost_p4.py b/src/vector/_compute/lorentz/boost_p4.py
@@ -24,6 +24,7 @@
 from vector._compute.lorentz import transform4D
 from vector._compute.planar import x, y
 from vector._compute.spatial import mag2, z
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -784,7 +785,7 @@ def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any:
         return _handler_of(v1, v2)._wrap_result(
             _flavor_of(v1, v2),
             function(
-                _lib_of(v1, v2),
+                Lib(_lib_of(v1, v2)),
                 *v1.azimuthal.elements,
                 *v1.longitudinal.elements,
                 *v1.temporal.elements,

diff --git a/src/vector/_compute/lorentz/deltaRapidityPhi.py b/src/vector/_compute/lorentz/deltaRapidityPhi.py
@@ -16,6 +16,7 @@
 import numpy
 
 from vector._compute.lorentz import deltaRapidityPhi2
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -113,7 +114,7 @@ def dispatch(
         return _handler_of(v1, v2)._wrap_result(
             _flavor_of(v1, v2),
             function(
-                _lib_of(v1, v2),
+                Lib(_lib_of(v1, v2)),
                 *v1.azimuthal.elements,
                 *v1.longitudinal.elements,
                 *v1.temporal.elements,

diff --git a/src/vector/_compute/lorentz/gamma.py b/src/vector/_compute/lorentz/gamma.py
@@ -18,6 +18,7 @@
 import numpy
 
 from vector._compute.lorentz import t, tau
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -154,7 +155,7 @@ def dispatch(v: typing.Any) -> typing.Any:
         return v._wrap_result(
             _flavor_of(v),
             function(
-                v.lib,
+                Lib(v.lib),
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,
                 *v.temporal.elements,

diff --git a/src/vector/_compute/lorentz/is_lightlike.py b/src/vector/_compute/lorentz/is_lightlike.py
@@ -16,6 +16,7 @@
 import numpy
 
 from vector._compute.lorentz import dot
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -69,7 +70,7 @@ def dispatch(tolerance: typing.Any, v: typing.Any) -> typing.Any:
         return v._wrap_result(
             _flavor_of(v),
             function(
-                v.lib,
+                Lib(v.lib),
                 tolerance,
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,

diff --git a/src/vector/_compute/lorentz/is_spacelike.py b/src/vector/_compute/lorentz/is_spacelike.py
@@ -16,6 +16,7 @@
 import numpy
 
 from vector._compute.lorentz import dot
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -67,7 +68,7 @@ def dispatch(tolerance: typing.Any, v: typing.Any) -> typing.Any:
         return v._wrap_result(
             _flavor_of(v),
             function(
-                v.lib,
+                Lib(v.lib),
                 tolerance,
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,

diff --git a/src/vector/_compute/lorentz/is_timelike.py b/src/vector/_compute/lorentz/is_timelike.py
@@ -16,6 +16,7 @@
 import numpy
 
 from vector._compute.lorentz import dot
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -67,7 +68,7 @@ def dispatch(tolerance: typing.Any, v: typing.Any) -> typing.Any:
         return v._wrap_result(
             _flavor_of(v),
             function(
-                v.lib,
+                Lib(v.lib),
                 tolerance,
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,

diff --git a/src/vector/_compute/lorentz/isclose.py b/src/vector/_compute/lorentz/isclose.py
@@ -17,6 +17,7 @@
 
 from vector._compute.lorentz import t
 from vector._compute.spatial import isclose
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -224,7 +225,7 @@ def dispatch(
         return _handler_of(v1, v2)._wrap_result(
             _flavor_of(v1, v2),
             function(
-                _lib_of(v1, v2),
+                Lib(_lib_of(v1, v2)),
                 rtol,
                 atol,
                 equal_nan,

diff --git a/src/vector/_compute/lorentz/rapidity.py b/src/vector/_compute/lorentz/rapidity.py
@@ -18,6 +18,7 @@
 
 from vector._compute.lorentz import t
 from vector._compute.spatial import z
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -128,7 +129,7 @@ def dispatch(v: typing.Any) -> typing.Any:
         return v._wrap_result(
             _flavor_of(v),
             function(
-                v.lib,
+                Lib(v.lib),
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,
                 *v.temporal.elements,

diff --git a/src/vector/_compute/lorentz/t.py b/src/vector/_compute/lorentz/t.py
@@ -17,6 +17,7 @@
 import numpy
 
 from vector._compute.lorentz import t2
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -111,7 +112,7 @@ def dispatch(v: typing.Any) -> typing.Any:
         return v._wrap_result(
             _flavor_of(v),
             function(
-                v.lib,
+                Lib(v.lib),
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,
                 *v.temporal.elements,

diff --git a/src/vector/_compute/lorentz/t2.py b/src/vector/_compute/lorentz/t2.py
@@ -18,6 +18,7 @@
 
 from vector._compute.lorentz import tau2
 from vector._compute.spatial import mag2
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -126,7 +127,7 @@ def dispatch(v: typing.Any) -> typing.Any:
         return v._wrap_result(
             _flavor_of(v),
             function(
-                v.lib,
+                Lib(v.lib),
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,
                 *v.temporal.elements,

diff --git a/src/vector/_compute/lorentz/tau.py b/src/vector/_compute/lorentz/tau.py
@@ -17,6 +17,7 @@
 import numpy
 
 from vector._compute.lorentz import tau2
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -117,7 +118,7 @@ def dispatch(v: typing.Any) -> typing.Any:
         return v._wrap_result(
             _flavor_of(v),
             function(
-                v.lib,
+                Lib(v.lib),
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,
                 *v.temporal.elements,

diff --git a/src/vector/_compute/lorentz/tau2.py b/src/vector/_compute/lorentz/tau2.py
@@ -17,6 +17,7 @@
 import numpy
 
 from vector._compute.spatial import mag2
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -111,7 +112,7 @@ def dispatch(v: typing.Any) -> typing.Any:
         return v._wrap_result(
             _flavor_of(v),
             function(
-                v.lib,
+                Lib(v.lib),
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,
                 *v.temporal.elements,

diff --git a/src/vector/_compute/lorentz/unit.py b/src/vector/_compute/lorentz/unit.py
@@ -17,6 +17,7 @@
 import numpy
 
 from vector._compute.lorentz import tau2
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -249,7 +250,7 @@ def dispatch(v: typing.Any) -> typing.Any:
         return v._wrap_result(
             _flavor_of(v),
             function(
-                v.lib,
+                Lib(v.lib),
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,
                 *v.temporal.elements,

diff --git a/src/vector/_compute/planar/add.py b/src/vector/_compute/planar/add.py
@@ -16,6 +16,7 @@
 import numpy
 
 from vector._compute.planar import x, y
+from vector._lib import Lib
 from vector._methods import (
     AzimuthalRhoPhi,
     AzimuthalXY,
@@ -75,7 +76,9 @@ def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any:
     with numpy.errstate(all="ignore"):
         return _handler_of(v1, v2)._wrap_result(
             _flavor_of(v1, v2),
-            function(_lib_of(v1, v2), *v1.azimuthal.elements, *v2.azimuthal.elements),
+            function(
+                Lib(_lib_of(v1, v2)), *v1.azimuthal.elements, *v2.azimuthal.elements
+            ),
             returns,
             2,
         )