diff --git a/src/vector/_backends/numba_object.py b/src/vector/_backends/numba_object.py index 1918e692..cd57234d 100644 --- a/src/vector/_backends/numba_object.py +++ b/src/vector/_backends/numba_object.py @@ -52,16 +52,31 @@ def nan_to_num(x, copy=True, nan=0.0, posinf=None, neginf=None): if isinstance(x, numba.types.Array): - def nan_to_num_impl(x, copy=True, nan=0.0, posinf=None, neginf=None): - out = numpy.copy(x).reshape(-1) if copy else x.reshape(-1) - for i in range(len(out)): - if numpy.isnan(out[i]): - out[i] = nan - if posinf is not None and numpy.isinf(out[i]) and out[i] > 0: - out[i] = posinf - if neginf is not None and numpy.isinf(out[i]) and out[i] < 0: - out[i] = neginf - return out.reshape(x.shape) + if isinstance(nan, numba.types.Array): + + def nan_to_num_impl(x, copy=True, nan=0.0, posinf=None, neginf=None): + out = numpy.copy(x).reshape(-1) if copy else x.reshape(-1) + for i in range(len(out)): + if numpy.isnan(out[i]): + out[i] = nan[i] + if posinf is not None and numpy.isinf(out[i]) and out[i] > 0: + out[i] = posinf + if neginf is not None and numpy.isinf(out[i]) and out[i] < 0: + out[i] = neginf + return out.reshape(x.shape) + + else: + + def nan_to_num_impl(x, copy=True, nan=0.0, posinf=None, neginf=None): + out = numpy.copy(x).reshape(-1) if copy else x.reshape(-1) + for i in range(len(out)): + if numpy.isnan(out[i]): + out[i] = nan + if posinf is not None and numpy.isinf(out[i]) and out[i] > 0: + out[i] = posinf + if neginf is not None and numpy.isinf(out[i]) and out[i] < 0: + out[i] = neginf + return out.reshape(x.shape) else: diff --git a/src/vector/_compute/spatial/eta.py b/src/vector/_compute/spatial/eta.py index 35a55777..b44a947b 100644 --- a/src/vector/_compute/spatial/eta.py +++ b/src/vector/_compute/spatial/eta.py @@ -4,6 +4,7 @@ # or https://github.com/scikit-hep/vector for details. import typing +from math import inf, nan """ .. code-block:: python @@ -29,8 +30,11 @@ def xy_z(lib, x, y, z): return lib.nan_to_num( - lib.arctanh(z / lib.sqrt(x**2 + y**2 + z**2)), nan=0.0 - ) * lib.absolute(lib.sign(z)) + lib.arcsinh(z / lib.sqrt(x**2 + y**2)), + nan=lib.nan_to_num((z != 0) * inf, posinf=nan), + posinf=inf, + neginf=-inf, + ) def xy_theta(lib, x, y, theta): @@ -43,8 +47,11 @@ def xy_eta(lib, x, y, eta): def rhophi_z(lib, rho, phi, z): return lib.nan_to_num( - lib.arctanh(z / lib.sqrt(rho**2 + z**2)), nan=0.0 - ) * lib.absolute(lib.sign(z)) + lib.arcsinh(z / rho), + nan=lib.nan_to_num((z != 0) * inf, posinf=nan), + posinf=inf, + neginf=-inf, + ) def rhophi_theta(lib, rho, phi, theta): diff --git a/tests/backends/test_awkward.py b/tests/backends/test_awkward.py index f17fb944..8ac37d5f 100644 --- a/tests/backends/test_awkward.py +++ b/tests/backends/test_awkward.py @@ -93,7 +93,7 @@ def test_projection(): { "rho": 6.4031242374328485, "phi": 0.8960553845713439, - "eta": 0.8361481196083127, + "eta": 0.8361481196083128, "tau": 0, "wow": 123, } diff --git a/tests/test_compute_features.py b/tests/test_compute_features.py index 0a446c78..d66b70a6 100644 --- a/tests/test_compute_features.py +++ b/tests/test_compute_features.py @@ -371,6 +371,9 @@ def analyze_callable(node, context): "arctan2", "sinh", "cosh", + "tanh", + "arcsinh", + "arccosh", "arctanh", "isclose", ]