skmap.parallel.utils.apply_along_axis#

apply_along_axis(worker, axis, arr, n_jobs=4, *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 (int) – Axis used to execute the worker.

  • arr (array) – The input array.

  • n_jobs (int) – Number of parallel jobs to run the worker function

  • 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

>>> import multiprocessing
>>> import numpy as np
>>> from skmap 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, 4, const)
>>> print(arr.shape, out.shape)