mirror of
https://github.com/wassname/scikit-image.git
synced 2026-07-04 01:24:37 +08:00
Blake Griffith
72 lines
2.2 KiB
Python
72 lines
2.2 KiB
Python
from multiprocessing import cpu_count
|
|
|
|
import dask.array as da
|
|
|
|
__all__ = ['process_chunks']
|
|
|
|
|
|
def _get_chunks(shape, ncpu):
|
|
"""
|
|
Split the array into equal sized chunks based on the number of
|
|
available processors. The last chunk in each dimension absorbs the
|
|
remainder array elements if the number of cpus does not divide evenly into
|
|
the number of array elements.
|
|
"""
|
|
chunks = []
|
|
|
|
for i in shape:
|
|
regular_chunk = i // ncpu
|
|
remainder_chunk = regular_chunk + (i % ncpu)
|
|
|
|
if regular_chunk == 0:
|
|
chunk_lens = (remainder_chunk,)
|
|
else:
|
|
chunk_lens = (regular_chunk,) * (ncpu - 1) + (remainder_chunk,)
|
|
|
|
chunks.append(chunk_lens)
|
|
return tuple(chunks)
|
|
|
|
|
|
def process_chunks(function, array, args=(), kwargs={} chunks=None, depth=0, mode=None):
|
|
"""Map a function in parallel across an array.
|
|
|
|
Split an array into possibly overlapping chunks of a given depth and
|
|
boundary type, call the given function in parallel on the chunks, combine
|
|
the chunks and return the resulting array.
|
|
|
|
Parameters
|
|
----------
|
|
function : function
|
|
Function to be mapped which takes an array as an argument.
|
|
array : numpy array
|
|
array which the function will be applied to.
|
|
chunks : int, tuple, or tuple of tuples
|
|
One tuple of length array.ndim or a list of tuples of length ndim.
|
|
Where each subtuple adds to the size of the array in the corresponding
|
|
dimension. If None, the array is broken up into chunks based on the
|
|
number of available cpus.
|
|
depth : int
|
|
integer equal to the depth of the internal external padding
|
|
mode : 'reflect', 'periodic', 'wrap', 'nearest'
|
|
type of external boundary padding
|
|
|
|
Notes
|
|
-----
|
|
Be careful choosing the depth so that it is never larger than the length of
|
|
a chunk.
|
|
|
|
"""
|
|
if chunks == None:
|
|
shape = array.shape
|
|
ncpu = cpu_count()
|
|
chunks = _get_chunks(shape, ncpu)
|
|
|
|
if mode == 'wrap':
|
|
mode = 'periodic'
|
|
|
|
def wrapped_func(arr):
|
|
return function(arr, *args, **kwargs)
|
|
|
|
darr = da.from_array(array, chunks=chunks)
|
|
return darr.map_overlap(wrapped_func, depth, boundary=mode).compute()
|