From 0af4ee0ced59cb4b8e20702348ca2ab0b9360152 Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Tue, 26 Mar 2024 16:18:37 +0100 Subject: [PATCH] wrapper around lib --- src/vector/_compute/lorentz/Et.py | 3 +- src/vector/_compute/lorentz/Et2.py | 3 +- src/vector/_compute/lorentz/Mt.py | 3 +- src/vector/_compute/lorentz/Mt2.py | 3 +- src/vector/_compute/lorentz/beta.py | 3 +- src/vector/_compute/lorentz/boostX_gamma.py | 3 +- src/vector/_compute/lorentz/boostY_gamma.py | 3 +- src/vector/_compute/lorentz/boostZ_gamma.py | 3 +- src/vector/_compute/lorentz/boost_beta3.py | 3 +- src/vector/_compute/lorentz/boost_p4.py | 3 +- .../_compute/lorentz/deltaRapidityPhi.py | 3 +- src/vector/_compute/lorentz/gamma.py | 3 +- src/vector/_compute/lorentz/is_lightlike.py | 3 +- src/vector/_compute/lorentz/is_spacelike.py | 3 +- src/vector/_compute/lorentz/is_timelike.py | 3 +- src/vector/_compute/lorentz/isclose.py | 3 +- src/vector/_compute/lorentz/rapidity.py | 3 +- src/vector/_compute/lorentz/t.py | 3 +- src/vector/_compute/lorentz/t2.py | 3 +- src/vector/_compute/lorentz/tau.py | 3 +- src/vector/_compute/lorentz/tau2.py | 3 +- src/vector/_compute/lorentz/unit.py | 3 +- src/vector/_compute/planar/add.py | 5 +- src/vector/_compute/planar/deltaphi.py | 5 +- src/vector/_compute/planar/dot.py | 5 +- src/vector/_compute/planar/is_antiparallel.py | 3 +- src/vector/_compute/planar/is_parallel.py | 3 +- .../_compute/planar/is_perpendicular.py | 3 +- src/vector/_compute/planar/isclose.py | 3 +- src/vector/_compute/planar/phi.py | 3 +- src/vector/_compute/planar/rho.py | 3 +- src/vector/_compute/planar/rotateZ.py | 6 +- src/vector/_compute/planar/scale.py | 6 +- src/vector/_compute/planar/subtract.py | 5 +- src/vector/_compute/planar/unit.py | 3 +- src/vector/_compute/planar/x.py | 3 +- src/vector/_compute/planar/y.py | 3 +- src/vector/_compute/spatial/costheta.py | 3 +- src/vector/_compute/spatial/cottheta.py | 3 +- src/vector/_compute/spatial/deltaR.py | 3 +- src/vector/_compute/spatial/deltaangle.py | 3 +- src/vector/_compute/spatial/dot.py | 3 +- src/vector/_compute/spatial/eta.py | 3 +- .../_compute/spatial/is_antiparallel.py | 3 +- src/vector/_compute/spatial/is_parallel.py | 3 +- .../_compute/spatial/is_perpendicular.py | 3 +- src/vector/_compute/spatial/isclose.py | 3 +- src/vector/_compute/spatial/mag.py | 3 +- src/vector/_compute/spatial/mag2.py | 3 +- src/vector/_compute/spatial/rotateX.py | 5 +- src/vector/_compute/spatial/rotateY.py | 5 +- src/vector/_compute/spatial/rotate_axis.py | 3 +- src/vector/_compute/spatial/rotate_euler.py | 8 +- src/vector/_compute/spatial/scale.py | 5 +- src/vector/_compute/spatial/theta.py | 3 +- src/vector/_compute/spatial/unit.py | 3 +- src/vector/_compute/spatial/z.py | 3 +- src/vector/_lib.py | 123 ++++++++++++++++++ 58 files changed, 262 insertions(+), 57 deletions(-) create mode 100644 src/vector/_lib.py diff --git a/src/vector/_compute/lorentz/Et.py b/src/vector/_compute/lorentz/Et.py index abd654e0..606c2e6e 100644 --- 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 index 5abb2891..1c48cf8f 100644 --- 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 index a92ec398..dc2da243 100644 --- 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 index 41fd04ed..28934091 100644 --- 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 index 9c39ee92..3ee6cd46 100644 --- 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 index 83fe6e1e..9ce5bb48 100644 --- 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 index 7cb6715d..50866f8a 100644 --- 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 index 2b74d393..6a69acf4 100644 --- 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 index 06fbc071..402e809e 100644 --- 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 index b016be15..4a9a5e40 100644 --- 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 index fb789a41..3aa9034a 100644 --- 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 index 04d8c549..d6e3196c 100644 --- 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 index 8988da4b..6c18c94e 100644 --- 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 index bf8d99ed..3aa6b007 100644 --- 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 index f515fa1b..0f4dbb8d 100644 --- 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 index 505a02bf..2ffad544 100644 --- 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 index 0d7c6806..0af4c198 100644 --- 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 index 91d25737..06f05cd2 100644 --- 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 index 84b8bac8..94324232 100644 --- 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 index 09cf92be..32e98d34 100644 --- 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 index 69122df3..6589a4b3 100644 --- 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 index e513e699..5e5f92bc 100644 --- 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 index 4e7faa25..99c79d95 100644 --- 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, ) diff --git a/src/vector/_compute/planar/deltaphi.py b/src/vector/_compute/planar/deltaphi.py index b74a4482..cd8a064a 100644 --- a/src/vector/_compute/planar/deltaphi.py +++ b/src/vector/_compute/planar/deltaphi.py @@ -15,6 +15,7 @@ import numpy from vector._compute.planar import phi +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -66,7 +67,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, ) diff --git a/src/vector/_compute/planar/dot.py b/src/vector/_compute/planar/dot.py index aeeab9f8..37842f93 100644 --- a/src/vector/_compute/planar/dot.py +++ b/src/vector/_compute/planar/dot.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, @@ -63,7 +64,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, ) diff --git a/src/vector/_compute/planar/is_antiparallel.py b/src/vector/_compute/planar/is_antiparallel.py index 31691b65..59c13b5d 100644 --- a/src/vector/_compute/planar/is_antiparallel.py +++ b/src/vector/_compute/planar/is_antiparallel.py @@ -16,6 +16,7 @@ import numpy from vector._compute.planar import dot, rho +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -60,7 +61,7 @@ def dispatch(tolerance: typing.Any, v1: typing.Any, v2: typing.Any) -> typing.An return _handler_of(v1, v2)._wrap_result( _flavor_of(v1, v2), function( - _lib_of(v1, v2), + Lib(_lib_of(v1, v2)), tolerance, *v1.azimuthal.elements, *v2.azimuthal.elements, diff --git a/src/vector/_compute/planar/is_parallel.py b/src/vector/_compute/planar/is_parallel.py index b7ab8c19..c9dda976 100644 --- a/src/vector/_compute/planar/is_parallel.py +++ b/src/vector/_compute/planar/is_parallel.py @@ -16,6 +16,7 @@ import numpy from vector._compute.planar import dot, rho +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -60,7 +61,7 @@ def dispatch(tolerance: typing.Any, v1: typing.Any, v2: typing.Any) -> typing.An return _handler_of(v1, v2)._wrap_result( _flavor_of(v1, v2), function( - _lib_of(v1, v2), + Lib(_lib_of(v1, v2)), tolerance, *v1.azimuthal.elements, *v2.azimuthal.elements, diff --git a/src/vector/_compute/planar/is_perpendicular.py b/src/vector/_compute/planar/is_perpendicular.py index 389265ba..6040525f 100644 --- a/src/vector/_compute/planar/is_perpendicular.py +++ b/src/vector/_compute/planar/is_perpendicular.py @@ -16,6 +16,7 @@ import numpy from vector._compute.planar import dot, rho +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -60,7 +61,7 @@ def dispatch(tolerance: typing.Any, v1: typing.Any, v2: typing.Any) -> typing.An return _handler_of(v1, v2)._wrap_result( _flavor_of(v1, v2), function( - _lib_of(v1, v2), + Lib(_lib_of(v1, v2)), tolerance, *v1.azimuthal.elements, *v2.azimuthal.elements, diff --git a/src/vector/_compute/planar/isclose.py b/src/vector/_compute/planar/isclose.py index 8e33b661..f864ac15 100644 --- a/src/vector/_compute/planar/isclose.py +++ b/src/vector/_compute/planar/isclose.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, @@ -97,7 +98,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/planar/phi.py b/src/vector/_compute/planar/phi.py index 550c4087..b31f3527 100644 --- a/src/vector/_compute/planar/phi.py +++ b/src/vector/_compute/planar/phi.py @@ -16,6 +16,7 @@ import numpy +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -43,5 +44,5 @@ def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature(__name__, dispatch_map, (_aztype(v),)) with numpy.errstate(all="ignore"): return v._wrap_result( - _flavor_of(v), function(v.lib, *v.azimuthal.elements), returns, 1 + _flavor_of(v), function(Lib(v.lib), *v.azimuthal.elements), returns, 1 ) diff --git a/src/vector/_compute/planar/rho.py b/src/vector/_compute/planar/rho.py index 7fa2a5d1..bfe6e864 100644 --- a/src/vector/_compute/planar/rho.py +++ b/src/vector/_compute/planar/rho.py @@ -17,6 +17,7 @@ import numpy from vector._compute.planar import rho2 +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -44,5 +45,5 @@ def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature(__name__, dispatch_map, (_aztype(v),)) with numpy.errstate(all="ignore"): return v._wrap_result( - _flavor_of(v), function(v.lib, *v.azimuthal.elements), returns, 1 + _flavor_of(v), function(Lib(v.lib), *v.azimuthal.elements), returns, 1 ) diff --git a/src/vector/_compute/planar/rotateZ.py b/src/vector/_compute/planar/rotateZ.py index 755d29db..2b4efbcf 100644 --- a/src/vector/_compute/planar/rotateZ.py +++ b/src/vector/_compute/planar/rotateZ.py @@ -15,6 +15,7 @@ import numpy +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -48,5 +49,8 @@ def dispatch(angle: typing.Any, v: typing.Any) -> typing.Any: function, *returns = _from_signature(__name__, dispatch_map, (_aztype(v),)) with numpy.errstate(all="ignore"): return v._wrap_result( - _flavor_of(v), function(v.lib, angle, *v.azimuthal.elements), returns, 1 + _flavor_of(v), + function(Lib(v.lib), angle, *v.azimuthal.elements), + returns, + 1, ) diff --git a/src/vector/_compute/planar/scale.py b/src/vector/_compute/planar/scale.py index 7dce62e8..a0b8ff07 100644 --- a/src/vector/_compute/planar/scale.py +++ b/src/vector/_compute/planar/scale.py @@ -15,6 +15,7 @@ import numpy +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -49,5 +50,8 @@ def dispatch(factor: typing.Any, v: typing.Any) -> typing.Any: function, *returns = _from_signature(__name__, dispatch_map, (_aztype(v),)) with numpy.errstate(all="ignore"): return v._wrap_result( - _flavor_of(v), function(v.lib, factor, *v.azimuthal.elements), returns, 1 + _flavor_of(v), + function(Lib(v.lib), factor, *v.azimuthal.elements), + returns, + 1, ) diff --git a/src/vector/_compute/planar/subtract.py b/src/vector/_compute/planar/subtract.py index f5118c2d..171d9b51 100644 --- a/src/vector/_compute/planar/subtract.py +++ b/src/vector/_compute/planar/subtract.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, ) diff --git a/src/vector/_compute/planar/unit.py b/src/vector/_compute/planar/unit.py index b53586be..8fc0af61 100644 --- a/src/vector/_compute/planar/unit.py +++ b/src/vector/_compute/planar/unit.py @@ -17,6 +17,7 @@ import numpy from vector._compute.planar import rho +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -48,5 +49,5 @@ def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature(__name__, dispatch_map, (_aztype(v),)) with numpy.errstate(all="ignore"): return v._wrap_result( - _flavor_of(v), function(v.lib, *v.azimuthal.elements), returns, 1 + _flavor_of(v), function(Lib(v.lib), *v.azimuthal.elements), returns, 1 ) diff --git a/src/vector/_compute/planar/x.py b/src/vector/_compute/planar/x.py index b40d738d..6137362a 100644 --- a/src/vector/_compute/planar/x.py +++ b/src/vector/_compute/planar/x.py @@ -16,6 +16,7 @@ import numpy +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -43,5 +44,5 @@ def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature(__name__, dispatch_map, (_aztype(v),)) with numpy.errstate(all="ignore"): return v._wrap_result( - _flavor_of(v), function(v.lib, *v.azimuthal.elements), returns, 1 + _flavor_of(v), function(Lib(v.lib), *v.azimuthal.elements), returns, 1 ) diff --git a/src/vector/_compute/planar/y.py b/src/vector/_compute/planar/y.py index 7165a9b1..ce496a1b 100644 --- a/src/vector/_compute/planar/y.py +++ b/src/vector/_compute/planar/y.py @@ -16,6 +16,7 @@ import numpy +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -43,5 +44,5 @@ def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature(__name__, dispatch_map, (_aztype(v),)) with numpy.errstate(all="ignore"): return v._wrap_result( - _flavor_of(v), function(v.lib, *v.azimuthal.elements), returns, 1 + _flavor_of(v), function(Lib(v.lib), *v.azimuthal.elements), returns, 1 ) diff --git a/src/vector/_compute/spatial/costheta.py b/src/vector/_compute/spatial/costheta.py index 29df48df..6f720fbd 100644 --- a/src/vector/_compute/spatial/costheta.py +++ b/src/vector/_compute/spatial/costheta.py @@ -18,6 +18,7 @@ import numpy from vector._compute.spatial import mag, theta +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -79,7 +80,7 @@ def dispatch(v: typing.Any) -> typing.Any: with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), - function(v.lib, *v.azimuthal.elements, *v.longitudinal.elements), + function(Lib(v.lib), *v.azimuthal.elements, *v.longitudinal.elements), returns, 1, ) diff --git a/src/vector/_compute/spatial/cottheta.py b/src/vector/_compute/spatial/cottheta.py index eaabdfd9..15b404c5 100644 --- a/src/vector/_compute/spatial/cottheta.py +++ b/src/vector/_compute/spatial/cottheta.py @@ -19,6 +19,7 @@ from vector._compute.planar import rho from vector._compute.spatial import theta +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -78,7 +79,7 @@ def dispatch(v: typing.Any) -> typing.Any: with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), - function(v.lib, *v.azimuthal.elements, *v.longitudinal.elements), + function(Lib(v.lib), *v.azimuthal.elements, *v.longitudinal.elements), returns, 1, ) diff --git a/src/vector/_compute/spatial/deltaR.py b/src/vector/_compute/spatial/deltaR.py index ca4451aa..44008ec2 100644 --- a/src/vector/_compute/spatial/deltaR.py +++ b/src/vector/_compute/spatial/deltaR.py @@ -16,6 +16,7 @@ import numpy from vector._compute.spatial import deltaR2 +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -345,7 +346,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, *v2.azimuthal.elements, diff --git a/src/vector/_compute/spatial/deltaangle.py b/src/vector/_compute/spatial/deltaangle.py index 09dea3d6..36bf8444 100644 --- a/src/vector/_compute/spatial/deltaangle.py +++ b/src/vector/_compute/spatial/deltaangle.py @@ -16,6 +16,7 @@ import numpy from vector._compute.spatial import dot, mag +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -463,7 +464,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, *v2.azimuthal.elements, diff --git a/src/vector/_compute/spatial/dot.py b/src/vector/_compute/spatial/dot.py index 523fe0c5..8ac026b3 100644 --- a/src/vector/_compute/spatial/dot.py +++ b/src/vector/_compute/spatial/dot.py @@ -17,6 +17,7 @@ from vector._compute.planar import x, y from vector._compute.spatial import theta, z +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -566,7 +567,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, *v2.azimuthal.elements, diff --git a/src/vector/_compute/spatial/eta.py b/src/vector/_compute/spatial/eta.py index 2f4007b6..3df50b1c 100644 --- a/src/vector/_compute/spatial/eta.py +++ b/src/vector/_compute/spatial/eta.py @@ -17,6 +17,7 @@ import numpy +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -88,7 +89,7 @@ def dispatch(v: typing.Any) -> typing.Any: with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), - function(v.lib, *v.azimuthal.elements, *v.longitudinal.elements), + function(Lib(v.lib), *v.azimuthal.elements, *v.longitudinal.elements), returns, 1, ) diff --git a/src/vector/_compute/spatial/is_antiparallel.py b/src/vector/_compute/spatial/is_antiparallel.py index 06b6dad0..70aa962e 100644 --- a/src/vector/_compute/spatial/is_antiparallel.py +++ b/src/vector/_compute/spatial/is_antiparallel.py @@ -16,6 +16,7 @@ import numpy from vector._compute.spatial import dot, mag +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -72,7 +73,7 @@ def dispatch(tolerance: typing.Any, v1: typing.Any, v2: typing.Any) -> typing.An return _handler_of(v1, v2)._wrap_result( _flavor_of(v1, v2), function( - _lib_of(v1, v2), + Lib(_lib_of(v1, v2)), tolerance, *v1.azimuthal.elements, *v1.longitudinal.elements, diff --git a/src/vector/_compute/spatial/is_parallel.py b/src/vector/_compute/spatial/is_parallel.py index 4041e1c1..4c2a0e99 100644 --- a/src/vector/_compute/spatial/is_parallel.py +++ b/src/vector/_compute/spatial/is_parallel.py @@ -16,6 +16,7 @@ import numpy from vector._compute.spatial import dot, mag +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -72,7 +73,7 @@ def dispatch(tolerance: typing.Any, v1: typing.Any, v2: typing.Any) -> typing.An return _handler_of(v1, v2)._wrap_result( _flavor_of(v1, v2), function( - _lib_of(v1, v2), + Lib(_lib_of(v1, v2)), tolerance, *v1.azimuthal.elements, *v1.longitudinal.elements, diff --git a/src/vector/_compute/spatial/is_perpendicular.py b/src/vector/_compute/spatial/is_perpendicular.py index ef7ed03b..ee088df0 100644 --- a/src/vector/_compute/spatial/is_perpendicular.py +++ b/src/vector/_compute/spatial/is_perpendicular.py @@ -16,6 +16,7 @@ import numpy from vector._compute.spatial import dot, mag +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -72,7 +73,7 @@ def dispatch(tolerance: typing.Any, v1: typing.Any, v2: typing.Any) -> typing.An return _handler_of(v1, v2)._wrap_result( _flavor_of(v1, v2), function( - _lib_of(v1, v2), + Lib(_lib_of(v1, v2)), tolerance, *v1.azimuthal.elements, *v1.longitudinal.elements, diff --git a/src/vector/_compute/spatial/isclose.py b/src/vector/_compute/spatial/isclose.py index 668b56b6..acdda5e2 100644 --- a/src/vector/_compute/spatial/isclose.py +++ b/src/vector/_compute/spatial/isclose.py @@ -17,6 +17,7 @@ from vector._compute.planar import x, y from vector._compute.spatial import eta, z +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -680,7 +681,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/spatial/mag.py b/src/vector/_compute/spatial/mag.py index 4de065cf..d2b189fc 100644 --- a/src/vector/_compute/spatial/mag.py +++ b/src/vector/_compute/spatial/mag.py @@ -17,6 +17,7 @@ import numpy from vector._compute.spatial import mag2 +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -80,7 +81,7 @@ def dispatch(v: typing.Any) -> typing.Any: with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), - function(v.lib, *v.azimuthal.elements, *v.longitudinal.elements), + function(Lib(v.lib), *v.azimuthal.elements, *v.longitudinal.elements), returns, 1, ) diff --git a/src/vector/_compute/spatial/mag2.py b/src/vector/_compute/spatial/mag2.py index 21a3104f..01423cfc 100644 --- a/src/vector/_compute/spatial/mag2.py +++ b/src/vector/_compute/spatial/mag2.py @@ -16,6 +16,7 @@ import numpy +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -79,7 +80,7 @@ def dispatch(v: typing.Any) -> typing.Any: with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), - function(v.lib, *v.azimuthal.elements, *v.longitudinal.elements), + function(Lib(v.lib), *v.azimuthal.elements, *v.longitudinal.elements), returns, 1, ) diff --git a/src/vector/_compute/spatial/rotateX.py b/src/vector/_compute/spatial/rotateX.py index 3dceb210..a5229b54 100644 --- a/src/vector/_compute/spatial/rotateX.py +++ b/src/vector/_compute/spatial/rotateX.py @@ -17,6 +17,7 @@ from vector._compute.planar import x, y from vector._compute.spatial import z +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -92,7 +93,9 @@ def dispatch(angle: typing.Any, v: typing.Any) -> typing.Any: with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), - function(v.lib, angle, *v.azimuthal.elements, *v.longitudinal.elements), + function( + Lib(v.lib), angle, *v.azimuthal.elements, *v.longitudinal.elements + ), returns, 1, ) diff --git a/src/vector/_compute/spatial/rotateY.py b/src/vector/_compute/spatial/rotateY.py index affe8a2c..bf718f71 100644 --- a/src/vector/_compute/spatial/rotateY.py +++ b/src/vector/_compute/spatial/rotateY.py @@ -17,6 +17,7 @@ from vector._compute.planar import x, y from vector._compute.spatial import z +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -92,7 +93,9 @@ def dispatch(angle: typing.Any, v: typing.Any) -> typing.Any: with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), - function(v.lib, angle, *v.azimuthal.elements, *v.longitudinal.elements), + function( + Lib(v.lib), angle, *v.azimuthal.elements, *v.longitudinal.elements + ), returns, 1, ) diff --git a/src/vector/_compute/spatial/rotate_axis.py b/src/vector/_compute/spatial/rotate_axis.py index 90207c31..3747cb0d 100644 --- a/src/vector/_compute/spatial/rotate_axis.py +++ b/src/vector/_compute/spatial/rotate_axis.py @@ -17,6 +17,7 @@ from vector._compute.planar import x, y from vector._compute.spatial import z +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -157,7 +158,7 @@ def dispatch(angle: typing.Any, v1: typing.Any, v2: typing.Any) -> typing.Any: return _handler_of(v2)._wrap_result( # note: _handler_of(v2) _flavor_of(v2), # note: _flavor_of(v2) function( - _lib_of(v1, v2), + Lib(_lib_of(v1, v2)), angle, *v1.azimuthal.elements, *v1.longitudinal.elements, diff --git a/src/vector/_compute/spatial/rotate_euler.py b/src/vector/_compute/spatial/rotate_euler.py index 3f1e76b4..2c603473 100644 --- a/src/vector/_compute/spatial/rotate_euler.py +++ b/src/vector/_compute/spatial/rotate_euler.py @@ -17,6 +17,7 @@ from vector._compute.planar import x, y from vector._compute.spatial import z +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -292,7 +293,12 @@ def dispatch( return v._wrap_result( _flavor_of(v), function( - v.lib, phi, theta, psi, *v.azimuthal.elements, *v.longitudinal.elements + Lib(v.lib), + phi, + theta, + psi, + *v.azimuthal.elements, + *v.longitudinal.elements, ), returns, 1, diff --git a/src/vector/_compute/spatial/scale.py b/src/vector/_compute/spatial/scale.py index 46a9fcb3..ad2d2d37 100644 --- a/src/vector/_compute/spatial/scale.py +++ b/src/vector/_compute/spatial/scale.py @@ -15,6 +15,7 @@ import numpy +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -94,7 +95,9 @@ def dispatch(factor: typing.Any, v: typing.Any) -> typing.Any: with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), - function(v.lib, factor, *v.azimuthal.elements, *v.longitudinal.elements), + function( + Lib(v.lib), factor, *v.azimuthal.elements, *v.longitudinal.elements + ), returns, 1, ) diff --git a/src/vector/_compute/spatial/theta.py b/src/vector/_compute/spatial/theta.py index c237868a..9beb5912 100644 --- a/src/vector/_compute/spatial/theta.py +++ b/src/vector/_compute/spatial/theta.py @@ -17,6 +17,7 @@ import numpy from vector._compute.spatial import costheta +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -76,7 +77,7 @@ def dispatch(v: typing.Any) -> typing.Any: with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), - function(v.lib, *v.azimuthal.elements, *v.longitudinal.elements), + function(Lib(v.lib), *v.azimuthal.elements, *v.longitudinal.elements), returns, 1, ) diff --git a/src/vector/_compute/spatial/unit.py b/src/vector/_compute/spatial/unit.py index 0c6e01c4..21b91a2c 100644 --- a/src/vector/_compute/spatial/unit.py +++ b/src/vector/_compute/spatial/unit.py @@ -17,6 +17,7 @@ import numpy from vector._compute.spatial import mag +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -102,7 +103,7 @@ def dispatch(v: typing.Any) -> typing.Any: with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), - function(v.lib, *v.azimuthal.elements, *v.longitudinal.elements), + function(Lib(v.lib), *v.azimuthal.elements, *v.longitudinal.elements), returns, 1, ) diff --git a/src/vector/_compute/spatial/z.py b/src/vector/_compute/spatial/z.py index a41c6632..0535a916 100644 --- a/src/vector/_compute/spatial/z.py +++ b/src/vector/_compute/spatial/z.py @@ -18,6 +18,7 @@ import numpy from vector._compute.planar import rho +from vector._lib import Lib from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, @@ -79,7 +80,7 @@ def dispatch(v: typing.Any) -> typing.Any: with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), - function(v.lib, *v.azimuthal.elements, *v.longitudinal.elements), + function(Lib(v.lib), *v.azimuthal.elements, *v.longitudinal.elements), returns, 1, ) diff --git a/src/vector/_lib.py b/src/vector/_lib.py new file mode 100644 index 00000000..3be4c305 --- /dev/null +++ b/src/vector/_lib.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +import typing + +from vector._typeutils import FloatArray + +if typing.TYPE_CHECKING: + import sympy + +import types + + +class Lib: + def __init__(self, lib: types.ModuleType): + if lib.__name__ == "sympy": + import sympy + + self.lib: types.ModuleType = sympy + elif lib.__name__ == "numpy": + import numpy + + self.lib = numpy + assert self.lib is not None + + def sqrt(self, val: FloatArray | sympy.Symbol) -> FloatArray | sympy.Symbol: + return self.lib.sqrt(val) + + def nan_to_num( + self, val: FloatArray | sympy.Symbol, **kwargs: int | float + ) -> FloatArray | sympy.Symbol: + if self.lib.__name__ == "sympy": + return val + else: + return self.lib.nan_to_num(val, **kwargs) + + def sin(self, val: FloatArray | sympy.Symbol) -> FloatArray | sympy.Symbol: + return self.lib.sin(val) + + def cos(self, val: FloatArray | sympy.Symbol) -> FloatArray | sympy.Symbol: + return self.lib.cos(val) + + def tan(self, val: FloatArray | sympy.Symbol) -> FloatArray | sympy.Symbol: + return self.lib.tan(val) + + def sinh(self, val: FloatArray | sympy.Symbol) -> FloatArray | sympy.Symbol: + return self.lib.sinh(val) + + def cosh(self, val: FloatArray | sympy.Symbol) -> FloatArray | sympy.Symbol: + return self.lib.cosh(val) + + def maximum( + self, val1: FloatArray | sympy.Symbol, val2: FloatArray | sympy.Symbol + ) -> FloatArray | sympy.Symbol: + return self.lib.maximum(val1, val2) + + def arctan2( + self, val1: FloatArray | sympy.Symbol, val2: FloatArray | sympy.Symbol + ) -> FloatArray | sympy.Symbol: + if self.lib.__name__ == "sympy": + return self.lib.atan2(val1, val2) + else: + return self.lib.arctan2(val1, val2) + + def arcsinh(self, val: FloatArray | sympy.Symbol) -> FloatArray | sympy.Symbol: + if self.lib.__name__ == "sympy": + return self.lib.asinh(val) + else: + return self.lib.arcsinh(val) + + def arccos(self, val: FloatArray | sympy.Symbol) -> FloatArray | sympy.Symbol: + if self.lib.__name__ == "sympy": + return self.lib.acos(val) + else: + return self.lib.arccos(val) + + def arctan(self, val: FloatArray | sympy.Symbol) -> FloatArray | sympy.Symbol: + if self.lib.__name__ == "sympy": + return self.lib.atan(val) + else: + return self.lib.arctan(val) + + def absolute(self, val: FloatArray | sympy.Symbol) -> FloatArray | sympy.Symbol: + if self.lib.__name__ == "sympy": + return self.lib.Abs(val) + else: + return self.lib.absolute(val) + + def log(self, val: FloatArray | sympy.Symbol) -> FloatArray | sympy.Symbol: + return self.lib.log(val) + + @property + def pi(self) -> float | sympy.Symbol: + return self.lib.pi + + def exp(self, val: FloatArray | sympy.Symbol) -> FloatArray | sympy.Symbol: + return self.lib.exp(val) + + @property + def inf(self) -> float | sympy.Symbol: + if self.lib.__name__ == "sympy": + return self.lib.oo + else: + return self.lib.inf + + def isclose( + self, + val1: FloatArray | sympy.Symbol, + val2: FloatArray | sympy.Symbol, + *args: float, + ) -> FloatArray | sympy.Equality: + if self.lib.__name__ == "sympy": + return self.lib.Eq(val1, val2) + else: + return self.lib.isclose(val1, val2, *args) + + # fix these + def copysign( + self, val1: FloatArray | sympy.Symbol, val2: FloatArray | sympy.Symbol + ) -> FloatArray | sympy.Symbol: + return self.lib.copysign(val1, val2) + + def sign(self, val: FloatArray | sympy.Symbol) -> FloatArray | sympy.Symbol: + return self.lib.sign(val)