fix bad values for high (abs) eta (#172)

* fix bad values for high (abs) eta

* fix awkward backend test

* avoid where function

* Allow these functions in compute methods.

* fix handling for z=0

* This solves the Numba issue; the Awkward one will need an Awkward fix (update ak.nan_to_num).

* another attempt at using where
- should not break numba
- should decay 0-rank array to scalar

* fix inverted condition

* remove unused imports

* revert to nan_to_num implementation

Co-authored-by: Jim Pivarski <[email protected]>

Author: bfis

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:
 

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):

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,
             }

tests/test_compute_features.py

@@ -371,6 +371,9 @@ def analyze_callable(node, context):
     "arctan2",
     "sinh",
     "cosh",
+    "tanh",
+    "arcsinh",
+    "arccosh",
     "arctanh",
     "isclose",
 ]