scipp.elemwise_func#

scipp.elemwise_func(func=None, *, unit_func=None, dtype='float64', auto_convert_dtypes=False)#

Create a function for transforming input variables based on element-wise operation.

This uses numba.cfunc to compile a kernel that Scipp can use for transforming the variable contents. Only variables with dtype=float64 are supported. Variances are not supported.

Custom kernels can reduce intermediate memory consumption and improve performance in multi-step operations with large input variables.

Parameters:

  • func (Optional[Callable], default: None) – Function to compute an output element from input element values.

  • unit_func (Optional[Callable], default: None) – Function to compute the output unit. If None, func will be used.

  • auto_convert_dtypes (bool, default: False) – Set to True to automatically convert all inputs to float64.

Returns:

Callable – A callable that applies func to the elements of the variables passed to it.

Examples

We can define a fused multiply-add operation as follows:

>>> def fmadd(a, b, c):
...     return a * b + c

>>> func = sc.elemwise_func(fmadd)

>>> x = sc.linspace('x', 0.0, 1.0, num=4, unit='m')
>>> y = x - 0.2 * x
>>> z = sc.scalar(1.2, unit='m**2')

>>> func(x, y, z)
<scipp.Variable> (x: 4)    float64            [m^2]  [1.2, 1.28889, 1.55556, 2]

Note that fmadd(x, y, z) would have the same effect in this case, but requires a potentially large intermediate allocation for the result of “a * b”.