diff --git a/src/vector/_backends/numba_object.py b/src/vector/_backends/numba_object.py index cd57234d..3eb3f75b 100644 --- a/src/vector/_backends/numba_object.py +++ b/src/vector/_backends/numba_object.py @@ -1440,6 +1440,7 @@ def overloader_impl(v): planar_binary_methods = ["deltaphi"] spatial_binary_methods = ["deltaangle", "deltaeta", "deltaR", "deltaR2"] +lorentz_binary_methods = ["deltaRapidityPhi", "deltaRapidityPhi2"] general_binary_methods = ["dot", "add", "subtract", "equal", "not_equal"] @@ -1629,6 +1630,10 @@ def overloader_impl(v1, v2): for vectortype in (VectorObject3DType, VectorObject4DType): add_binary_method(vectortype, "spatial", methodname) +for methodname in lorentz_binary_methods: + for vectortype in (VectorObject4DType,): + add_binary_method(vectortype, "lorentz", methodname) + for methodname in general_binary_methods: add_binary_method(VectorObject2DType, None, methodname) diff --git a/src/vector/_compute/lorentz/__init__.py b/src/vector/_compute/lorentz/__init__.py index 1665b5a3..22e96f29 100644 --- a/src/vector/_compute/lorentz/__init__.py +++ b/src/vector/_compute/lorentz/__init__.py @@ -31,6 +31,8 @@ import vector._compute.lorentz.boostY_gamma # noqa: F401 import vector._compute.lorentz.boostZ_beta # noqa: F401 import vector._compute.lorentz.boostZ_gamma # noqa: F401 +import vector._compute.lorentz.deltaRapidityPhi # noqa: F401 +import vector._compute.lorentz.deltaRapidityPhi2 # noqa: F401 import vector._compute.lorentz.dot # noqa: F401 import vector._compute.lorentz.equal # noqa: F401 import vector._compute.lorentz.Et # noqa: F401 diff --git a/src/vector/_compute/lorentz/deltaRapidityPhi.py b/src/vector/_compute/lorentz/deltaRapidityPhi.py new file mode 100644 index 00000000..fd2f1b96 --- /dev/null +++ b/src/vector/_compute/lorentz/deltaRapidityPhi.py @@ -0,0 +1,124 @@ +# Copyright (c) 2019-2021, Jonas Eschle, Jim Pivarski, Eduardo Rodrigues, and Henry Schreiner. +# +# Distributed under the 3-clause BSD license, see accompanying file LICENSE +# or https://github.com/scikit-hep/vector for details. + +import typing + +""" +.. code-block:: python + + Spatial.deltaRapidityPhi(self, other) +""" + +import numpy + +from vector._compute.lorentz import deltaRapidityPhi2 +from vector._methods import ( + AzimuthalRhoPhi, + AzimuthalXY, + LongitudinalEta, + LongitudinalTheta, + LongitudinalZ, + TemporalT, + TemporalTau, + _aztype, + _flavor_of, + _from_signature, + _handler_of, + _lib_of, + _ltype, + _ttype, +) + +dispatch_map = {} + + +def make_conversion( + azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 +): + lorentz_deltaRapidityPhi2, _ = deltaRapidityPhi2.dispatch_map[ + azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 + ] + + def f( + lib, + coord11, + coord12, + coord13, + coord14, + coord21, + coord22, + coord23, + coord24, + ): + return lib.sqrt( + lorentz_deltaRapidityPhi2( + lib, + coord11, + coord12, + coord13, + coord14, + coord21, + coord22, + coord23, + coord24, + ) + ) + + dispatch_map[ + azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 + ] = (f, float) + + +for azimuthal1 in (AzimuthalXY, AzimuthalRhoPhi): + for longitudinal1 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): + for temporal1 in (TemporalT, TemporalTau): + for azimuthal2 in (AzimuthalXY, AzimuthalRhoPhi): + for longitudinal2 in ( + LongitudinalZ, + LongitudinalTheta, + LongitudinalEta, + ): + for temporal2 in (TemporalT, TemporalTau): + make_conversion( + azimuthal1, + longitudinal1, + temporal1, + azimuthal2, + longitudinal2, + temporal2, + ) + + +def dispatch( + v1: typing.Any, + v2: typing.Any, +) -> typing.Any: + function, *returns = _from_signature( + __name__, + dispatch_map, + ( + _aztype(v1), + _ltype(v1), + _ttype(v1), + _aztype(v2), + _ltype(v2), + _ttype(v2), + ), + ) + with numpy.errstate(all="ignore"): + return _handler_of(v1, v2)._wrap_result( + _flavor_of(v1, v2), + function( + _lib_of(v1, v2), + *v1.azimuthal.elements, + *v1.longitudinal.elements, + *v1.temporal.elements, + *v2.azimuthal.elements, + *v2.longitudinal.elements, + *v2.temporal.elements, + ), + returns, + 2, + ) diff --git a/src/vector/_compute/lorentz/deltaRapidityPhi2.py b/src/vector/_compute/lorentz/deltaRapidityPhi2.py new file mode 100644 index 00000000..bdb0877c --- /dev/null +++ b/src/vector/_compute/lorentz/deltaRapidityPhi2.py @@ -0,0 +1,120 @@ +# Copyright (c) 2019-2021, Jonas Eschle, Jim Pivarski, Eduardo Rodrigues, and Henry Schreiner. +# +# Distributed under the 3-clause BSD license, see accompanying file LICENSE +# or https://github.com/scikit-hep/vector for details. + +import typing + +""" +.. code-block:: python + + Spatial.deltaRapidityPhi2(self, other) +""" + +import numpy + +from vector._compute.lorentz import rapidity +from vector._compute.planar import deltaphi +from vector._methods import ( + AzimuthalRhoPhi, + AzimuthalXY, + LongitudinalEta, + LongitudinalTheta, + LongitudinalZ, + TemporalT, + TemporalTau, + _aztype, + _flavor_of, + _from_signature, + _handler_of, + _lib_of, + _ltype, + _ttype, +) + +dispatch_map = {} + + +def make_conversion( + azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 +): + planar_deltaphi, _ = deltaphi.dispatch_map[azimuthal1, azimuthal2] + lorentz_rapidity1, _ = rapidity.dispatch_map[azimuthal1, longitudinal1, temporal1] + lorentz_rapidity2, _ = rapidity.dispatch_map[azimuthal2, longitudinal2, temporal2] + + def f( + lib, + coord11, + coord12, + coord13, + coord14, + coord21, + coord22, + coord23, + coord24, + ): + return ( + planar_deltaphi(lib, coord11, coord12, coord21, coord22) ** 2 + + ( + lorentz_rapidity1(lib, coord11, coord12, coord13, coord14) + - lorentz_rapidity2(lib, coord21, coord22, coord23, coord24) + ) + ** 2 + ) + + dispatch_map[ + azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 + ] = (f, float) + + +for azimuthal1 in (AzimuthalXY, AzimuthalRhoPhi): + for longitudinal1 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): + for temporal1 in (TemporalT, TemporalTau): + for azimuthal2 in (AzimuthalXY, AzimuthalRhoPhi): + for longitudinal2 in ( + LongitudinalZ, + LongitudinalTheta, + LongitudinalEta, + ): + for temporal2 in (TemporalT, TemporalTau): + make_conversion( + azimuthal1, + longitudinal1, + temporal1, + azimuthal2, + longitudinal2, + temporal2, + ) + + +def dispatch( + v1: typing.Any, + v2: typing.Any, +) -> typing.Any: + function, *returns = _from_signature( + __name__, + dispatch_map, + ( + _aztype(v1), + _ltype(v1), + _ttype(v1), + _aztype(v2), + _ltype(v2), + _ttype(v2), + ), + ) + with numpy.errstate(all="ignore"): + return _handler_of(v1, v2)._wrap_result( + _flavor_of(v1, v2), + function( + _lib_of(v1, v2), + *v1.azimuthal.elements, + *v1.longitudinal.elements, + *v1.temporal.elements, + *v2.azimuthal.elements, + *v2.longitudinal.elements, + *v2.temporal.elements, + ), + returns, + 2, + ) diff --git a/src/vector/_methods.py b/src/vector/_methods.py index eb0948b6..65cd07ad 100644 --- a/src/vector/_methods.py +++ b/src/vector/_methods.py @@ -653,7 +653,7 @@ def deltaeta(self, other: "VectorProtocol") -> ScalarCollection: def deltaR(self, other: "VectorProtocol") -> ScalarCollection: r""" Sum in quadrature of :doc:`vector._methods.VectorProtocolPlanar.deltaphi` - and :doc:`vector._methods.VectorProtocolPlanar.deltaeta`: + and :doc:`vector._methods.VectorProtocolSpatial.deltaeta`: $$\Delta R = \sqrt{\Delta\phi^2 + \Delta\eta^2}$$ """ @@ -897,6 +897,26 @@ def rapidity(self) -> ScalarCollection: """ raise AssertionError + def deltaRapidityPhi(self, other: "VectorProtocol") -> ScalarCollection: + r""" + Sum in quadrature of :doc:`vector._methods.VectorProtocolPlanar.deltaphi` + and the difference in :doc:`vector._methods.VectorProtocolLorentz.rapidity` + of the two vectors: + + $$\Delta R_{\mbox{rapidity}} = \sqrt{\Delta\phi^2 + \Delta \mbox{rapidity}^2}$$ + """ + raise AssertionError + + def deltaRapidityPhi2(self, other: "VectorProtocol") -> ScalarCollection: + r""" + Square of the sum in quadrature of + :doc:`vector._methods.VectorProtocolPlanar.deltaphi` and the difference in + :doc:`vector._methods.VectorProtocolLorentz.rapidity` of the two vectors: + + $$\Delta R_{\mbox{rapidity}} = \Delta\phi^2 + \Delta \mbox{rapidity}^2$$ + """ + raise AssertionError + def scale4D(self: SameVectorType, factor: ScalarCollection) -> SameVectorType: """ Same as ``scale``. @@ -2186,6 +2206,16 @@ def rapidity(self) -> ScalarCollection: return rapidity.dispatch(self) + def deltaRapidityPhi(self, other: VectorProtocol) -> ScalarCollection: + from vector._compute.lorentz import deltaRapidityPhi + + return deltaRapidityPhi.dispatch(self, other) + + def deltaRapidityPhi2(self, other: VectorProtocol) -> ScalarCollection: + from vector._compute.lorentz import deltaRapidityPhi2 + + return deltaRapidityPhi2.dispatch(self, other) + def boost_p4(self: SameVectorType, p4: VectorProtocolLorentz) -> SameVectorType: from vector._compute.lorentz import boost_p4