Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix bad values for high (abs) eta #172

Merged
merged 10 commits into from
Mar 2, 2022
35 changes: 25 additions & 10 deletions src/vector/_backends/numba_object.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:

Expand Down
15 changes: 11 additions & 4 deletions src/vector/_compute/spatial/eta.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
# or https://github.com/scikit-hep/vector for details.

import typing
from math import inf, nan

"""
.. code-block:: python
Expand All @@ -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):
Expand All @@ -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):
Expand Down
2 changes: 1 addition & 1 deletion tests/backends/test_awkward.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@ def test_projection():
{
"rho": 6.4031242374328485,
"phi": 0.8960553845713439,
"eta": 0.8361481196083127,
"eta": 0.8361481196083128,
"tau": 0,
"wow": 123,
}
Expand Down
3 changes: 3 additions & 0 deletions tests/test_compute_features.py
Original file line number Diff line number Diff line change
Expand Up @@ -371,6 +371,9 @@ def analyze_callable(node, context):
"arctan2",
"sinh",
"cosh",
"tanh",
"arcsinh",
"arccosh",
"arctanh",
"isclose",
]