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.

Sign up for GitHub

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

Try using scipp.elemwise_func #381

Draft
wants to merge 3 commits into
base: main
Choose a base branch
from
Draft

Try using scipp.elemwise_func #381

Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8b5a43c
Try using scipp.elemwise_func
jl-wynen Nov 21, 2022
6af1be7
Make numba optional
SimonHeybrock Dec 8, 2022
ff33271
Comment on dtype check
SimonHeybrock Dec 8, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 22 additions & 0 deletions src/scippneutron/conversion/tof.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -179,6 +179,28 @@ def energy_transfer_direct_from_tof(*, tof: VariableLike, L1: VariableLike,
c = _energy_constant(elem_unit(incident_energy), tof, L2)
dtype = _common_dtype(incident_energy, tof)
scale = (c * L2**2).astype(dtype, copy=False)

import numba
sig = numba.double(numba.double, numba.double, numba.double, numba.double)

# This is a way to avoid repeating the part that is needed for the unit computation,
# but unfortunately we need to manually use numba.cfunc for the nested call.
# I cannot think of a way to do this in sc.elemwise_func automatically.
@numba.cfunc(sig)
def dE(tof, t0, scale, incident_energy):
return incident_energy - scale / (tof - t0)**2

def dE_or_nan(tof, t0, scale, incident_energy):
delta_tof = tof - t0
return np.NAN if delta_tof <= 0.0 else dE(tof, t0, scale, incident_energy)

# Converting everything to float64...
# could make a different branch for float32?
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Branch? We would need to extend the C++ implementation, no?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, and branch here to compile two difference numba.cfuncs

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a practical need for single precision floats in this case?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could be, inelastic data can be huge.

func = sc.elemwise_func(dE_or_nan, unit_func=dE, auto_convert_dtypes=True)
return func(tof, t0, scale, incident_energy)

# Should we keep this to continue support if the user does not have numba?
# Or should numba become a mandatory dependency of scippneutron?
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For now we should keep both because numba is a pretty big dependency. But we need a way to select which implementation is used in order to test both.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One way would be to use tox, with two different environments (with and without numba)?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes. This would be expensive but thorough because it would ensure that everyhing works without numba.

delta_tof = tof - t0
return sc.where(delta_tof <= sc.scalar(0, unit=elem_unit(delta_tof)),
sc.scalar(np.nan, dtype=dtype, unit=elem_unit(incident_energy)),
Expand Down