def test_basic(self):
# Basic check
chunks = generate_chunks(self.shape, self.dtype, 3e6)
assert_equal(chunks, (10 * (1,), 4 * (2048,), (144,)))
# Uneven chunks in the final split
chunks = generate_chunks(self.shape, self.dtype, 1e6)
assert_equal(chunks, (10 * (1,), 10 * (819,) + (2,), (144,)))
def test_power_of_two(self):
# Check power_of_two
chunks = generate_chunks(self.shape,
self.dtype,
1e6,
dims_to_split=[1],
power_of_two=True)
assert_equal(chunks, ((10, ), 128 * (64, ), (144, )))
chunks = generate_chunks(self.shape,
self.dtype,
self.nbytes / 16,
dims_to_split=[1],
power_of_two=True)
assert_equal(chunks, ((10, ), 16 * (512, ), (144, )))
# Check power_of_two when dimension is not a power-of-two itself
shape = (10, 32768 - 2048, 144)
chunks = generate_chunks(shape,
self.dtype,
self.nbytes / 10,
dims_to_split=(0, 1),
power_of_two=True)
assert_equal(chunks, (10 * (1, ), 3 * (8192, ) + (6144, ), (144, )))
# Check swapping the order of dims_to_split
chunks = generate_chunks(shape,
self.dtype,
self.nbytes / 10,
dims_to_split=(1, 0),
power_of_two=True)
assert_equal(chunks, ((10, ), 60 * (512, ), (144, )))
def test_corner_cases(self):
# Corner case: don't select any dimensions to split -> one chunk
chunks = generate_chunks(self.shape, self.dtype, 1e6, ())
assert_equal(chunks, ((10,), (8192,), (144,)))
# Corner case: all bytes results in one chunk
chunks = generate_chunks(self.shape, self.dtype, self.nbytes)
assert_equal(chunks, ((10,), (8192,), (144,)))
# Corner case: one byte less than the full size results in a split
chunks = generate_chunks(self.shape, self.dtype, self.nbytes - 1)
assert_equal(chunks, ((5, 5), (8192,), (144,)))
def test_basic(self):
# Basic check
chunks = generate_chunks(self.shape, self.dtype, 3e6)
assert_equal(chunks, (10 * (1,), 4 * (2048,), (144,)))
# Check that bad dims_to_split are ignored
chunks = generate_chunks(self.shape, self.dtype, 3e6, (0, 10))
assert_equal(chunks, (10 * (1,), (8192,), (144,)))
# Uneven chunks in the final split
chunks = generate_chunks(self.shape, self.dtype, 1e6)
assert_equal(chunks, (10 * (1,), 2 * (820,) + 8 * (819,), (144,)))
def test_max_dim_elements(self):
chunks = generate_chunks(self.shape, self.dtype, 150000,
dims_to_split=(0, 1), power_of_two=True,
max_dim_elements={1: 50})
assert_equal(chunks, ((4, 4, 2), 256 * (32,), (144,)))
# Case where max_dim_elements forces chunks to be smaller than
# max_chunk_size.
chunks = generate_chunks(self.shape, self.dtype, 1e6,
dims_to_split=(0, 1), power_of_two=True,
max_dim_elements={0: 4, 1: 50})
assert_equal(chunks, ((4, 4, 2), 256 * (32,), (144,)))
def make_dask_array(self, var_name, slices=()):
"""Turn (part of) an existing ndarray into a dask array."""
array_name = self.array_name(var_name)
array = getattr(self, var_name)
# The chunking is determined by full array to keep things predictable
chunks = generate_chunks(array.shape, array.dtype, array.nbytes / 10.)
dask_array = da.from_array(array, chunks)[slices]
offset = tuple(s.start for s in slices)
return array_name, dask_array, offset
def to_dask_array(x, chunks=None):
"""Turn ndarray `x` into dask array with the standard vis-like chunking."""
if chunks is None:
itemsize = np.dtype('complex64').itemsize
# Special case for 2-D weights_channel array ensures one chunk per dump
n_corrprods = x.shape[2] if x.ndim >= 3 else x.shape[1] // itemsize
# This contrives to have a vis array with 1 dump and 4 channels per chunk
chunk_size = 4 * n_corrprods * itemsize
chunks = generate_chunks(x.shape, x.dtype, chunk_size,
dims_to_split=(0, 1), power_of_two=True)
return da.from_array(x, chunks)
def test_max_dim_elements_ignore(self):
"""Elements not in `dims_to_split` are ignored"""
chunks = generate_chunks(self.shape,
self.dtype,
150000,
dims_to_split=(1, 17),
power_of_two=True,
max_dim_elements={
0: 2,
1: 50
})
assert_equal(chunks, ((10, ), 1024 * (8, ), (144, )))