From 8b5a43cc92aa557278466dd65567133531581210 Mon Sep 17 00:00:00 2001 From: Jan-Lukas Wynen Date: Mon, 21 Nov 2022 15:10:40 +0100 Subject: [PATCH 1/3] Try using scipp.elemwise_func --- src/scippneutron/conversion/tof.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/src/scippneutron/conversion/tof.py b/src/scippneutron/conversion/tof.py index 0f7d14e0e..875762a19 100644 --- a/src/scippneutron/conversion/tof.py +++ b/src/scippneutron/conversion/tof.py @@ -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? + 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? 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)), From 6af1be70ea97b61f59a68f59583b8a9fda93e9fc Mon Sep 17 00:00:00 2001 From: Simon Heybrock Date: Thu, 8 Dec 2022 10:02:02 +0100 Subject: [PATCH 2/3] Make numba optional --- lib/CMakeLists.txt | 4 +-- src/scippneutron/conversion/tof.py | 43 ++++++++++++++++-------------- 2 files changed, 25 insertions(+), 22 deletions(-) diff --git a/lib/CMakeLists.txt b/lib/CMakeLists.txt index d1c201994..820201c8c 100644 --- a/lib/CMakeLists.txt +++ b/lib/CMakeLists.txt @@ -11,10 +11,10 @@ set(CMAKE_CXX_EXTENSIONS OFF) set(CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake") find_package(Python 3.8 REQUIRED COMPONENTS Interpreter Development) -find_package(scipp 0.17 REQUIRED COMPONENTS conan-config) +find_package(scipp 22.11 REQUIRED COMPONENTS conan-config) find_package(GTest CONFIG REQUIRED) find_package(pybind11 CONFIG REQUIRED) -find_package(scipp 0.17 REQUIRED) +find_package(scipp 22.11 REQUIRED) find_program(CCACHE ccache) if(CCACHE) diff --git a/src/scippneutron/conversion/tof.py b/src/scippneutron/conversion/tof.py index 875762a19..1afccf6b1 100644 --- a/src/scippneutron/conversion/tof.py +++ b/src/scippneutron/conversion/tof.py @@ -133,6 +133,25 @@ def _energy_transfer_t0(energy, tof, length): return length.astype(dtype, copy=False) * sc.sqrt(c / energy) +try: + import numba +except ImportError: + energy_transfer_direct_from_tof_numba_float64 = None +else: + sig = numba.double(numba.double, numba.double, numba.double, numba.double) + + @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) + + energy_transfer_direct_from_tof_numba_float64 = sc.elemwise_func( + dE_or_nan, unit_func=dE, auto_convert_dtypes=True) + + def energy_transfer_direct_from_tof(*, tof: VariableLike, L1: VariableLike, L2: VariableLike, incident_energy: VariableLike) -> VariableLike: @@ -180,27 +199,11 @@ def energy_transfer_direct_from_tof(*, tof: VariableLike, L1: VariableLike, 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? - func = sc.elemwise_func(dE_or_nan, unit_func=dE, auto_convert_dtypes=True) - return func(tof, t0, scale, incident_energy) + if (energy_transfer_direct_from_tof_numba_float64 is not None) and (tof.dtype + == 'float64'): + return energy_transfer_direct_from_tof_numba_float64(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? 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)), From ff332714311398adf6e450488d7795460a4213df Mon Sep 17 00:00:00 2001 From: Simon Heybrock Date: Thu, 8 Dec 2022 10:03:57 +0100 Subject: [PATCH 3/3] Comment on dtype check --- src/scippneutron/conversion/tof.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/scippneutron/conversion/tof.py b/src/scippneutron/conversion/tof.py index 1afccf6b1..d8e8773b6 100644 --- a/src/scippneutron/conversion/tof.py +++ b/src/scippneutron/conversion/tof.py @@ -201,6 +201,9 @@ def energy_transfer_direct_from_tof(*, tof: VariableLike, L1: VariableLike, if (energy_transfer_direct_from_tof_numba_float64 is not None) and (tof.dtype == 'float64'): + # We could use this even of tof.dtype is not float64, but then there would be + # an allocation of significant size for the dtype conversion, so using the + # Numba branch is probably not with it. return energy_transfer_direct_from_tof_numba_float64(tof, t0, scale, incident_energy)