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 0f7d14e0e..d8e8773b6 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: @@ -179,6 +198,15 @@ 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) + + 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) + 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)),