eumap.parallel.utils.apply_along_axis

apply_along_axis(worker, axis, arr, *args, **kwargs)[source]

Execute a function through a numpy.array axis in parallel [1]. It uses joblib and backend=loky, so avoid to send shared memory objects as arguments.

Parameters
  • worker (Callable) – Function to execute in parallel. It needs to have at least one argument (numpy.array).

  • axis – Axis used to execute the worker.

  • arr (array) – The input array.

  • args (any) – Additional arguments to the worker.

  • kwargs (any) – Additional named arguments to the worker.

Returns

The output array with one dimension less than the input array.

Return type

numpy.array

References

[1] Best answer from Eric O Lebigot

Examples

>>> from eumap import parallel
>>>
>>> def fn(arr, const):
>>>   return np.sum(arr) + const
>>>
>>> const = 1
>>> arr = np.ones((100,100,100))
>>>
>>> out = parallel.apply_along_axis(fn, 0, arr, const)
>>> print(arr.shape, out.shape)