MCPcopy Create free account
hub / github.com/NSLS2/PyXRF / _chunk_numpy_array

Function _chunk_numpy_array

pyxrf/core/map_processing.py:243–289  ·  view source on GitHub ↗

Convert a numpy array into Dask array with chunks of given size. The function splits the array into chunks along axes 0 and 1. If the array has more than 2 dimensions, then the remaining dimensions are not chunked. Note, that `dask_array = da.array(data, chunks=...)` will set the ch

(data, chunk_size)

Source from the content-addressed store, hash-verified

241
242
243def _chunk_numpy_array(data, chunk_size):
244 """
245 Convert a numpy array into Dask array with chunks of given size. The function
246 splits the array into chunks along axes 0 and 1. If the array has more than 2 dimensions,
247 then the remaining dimensions are not chunked. Note, that
248 `dask_array = da.array(data, chunks=...)` will set the chunk size, but not split the
249 data into chunks, therefore the array can not be loaded block by block by workers
250 controlled by a distributed scheduler.
251
252 Parameters
253 ----------
254 data: ndarray(float), 2 or more dimensions
255 XRF map of the shape `(ny, nx, ne)`, where `ny` and `nx` represent the image size
256 and `ne` is the number of points in spectra
257 chunk_size: tuple(int, int) or list(int, int)
258 Chunk size for axis 0 and 1: `(chunk_y, chunk_x`). The function will accept
259 chunk size values that are larger then the respective `data` array dimensions.
260
261 Returns
262 -------
263 data_dask: dask.array
264 Dask array with the given chunk size
265 """
266
267 chunk_y, chunk_x = chunk_size
268 ny, nx = data.shape[0:2]
269 chunk_y, chunk_x = min(chunk_y, ny), min(chunk_x, nx)
270
271 def _get_slice(n1, n2):
272 data_slice = data[
273 slice(n1 * chunk_y, min(n1 * chunk_y + chunk_y, ny)),
274 slice(n2 * chunk_x, min(n2 * chunk_x + chunk_x, nx)),
275 ]
276 # Wrap the slice into a list wiht appropriate dimensions
277 for _ in range(2, data.ndim):
278 data_slice = [data_slice]
279 return data_slice
280
281 # Chunk the numpy array and assemble it as a dask array
282 data_dask = da.block(
283 [
284 [_get_slice(_1, _2) for _2 in range(int(math.ceil(nx / chunk_x)))]
285 for _1 in range(int(math.ceil(ny / chunk_y)))
286 ]
287 )
288
289 return data_dask
290
291
292def _array_numpy_to_dask(data, chunk_pixels, n_chunks_min=4):

Callers 3

test_chunk_numpy_arrayFunction · 0.90
_array_numpy_to_daskFunction · 0.85
_prepare_xrf_maskFunction · 0.85

Calls 1

_get_sliceFunction · 0.85

Tested by 1

test_chunk_numpy_arrayFunction · 0.72