def test_key_and_push_args(self):
context = Context()
da = context.ones((2, 2))
db = da*2
def dummy_func(*args, **kwargs):
fn = lambda x: x
db = DecoratorBase(fn)
return db.key_and_push_args(args, kwargs)
# Push some distarrays
arg_keys1, kw_keys1 = dummy_func(da, db, foo=da, bar=db)
# with some other data too
arg_keys2, kw_keys2 = dummy_func(da, 'question', answer=42, foo=db)
self.assertEqual(arg_keys1, "(%s, %s,)" % (da.key, db.key))
# assert we pushed the right key, keystr pair
self.assertTrue("'foo': %s" % (da.key) in kw_keys1)
self.assertTrue("'bar': %s" % (db.key) in kw_keys1)
# lots of string manipulation to parse out the relevant pieces
# of the python commands.
self.assertEqual(arg_keys2[1: -2].split(', ')[0], da.key)
_key = arg_keys2[1: -2].split(', ')[1]
self.assertEqual(context._pull0(_key), 'question')
self.assertTrue("'answer'" in kw_keys2)
self.assertTrue("'foo'" in kw_keys2)
self.assertTrue(db.key in kw_keys2)
class TestDistArrayCreation(IpclusterTestCase):
"""Test distarray creation methods"""
def setUp(self):
self.context = Context(self.client)
def test_zeros(self):
shape = (16, 16)
zero_distarray = self.context.zeros(shape)
zero_ndarray = numpy.zeros(shape)
assert_array_equal(zero_distarray.tondarray(), zero_ndarray)
def test_ones(self):
shape = (16, 16)
one_distarray = self.context.ones(shape)
one_ndarray = numpy.ones(shape)
assert_array_equal(one_distarray.tondarray(), one_ndarray)
def test_empty(self):
shape = (16, 16)
empty_distarray = self.context.empty(shape)
self.assertEqual(empty_distarray.shape, shape)
def test_fromndarray(self):
ndarr = numpy.arange(16).reshape(4, 4)
distarr = self.context.fromndarray(ndarr)
for (i, j), val in numpy.ndenumerate(ndarr):
self.assertEqual(distarr[i, j], ndarr[i, j])
def test_writing_two_datasets(self):
h5py = import_or_skip('h5py')
datalen = 33
dac = Context(self.client)
da = dac.empty((datalen,), dist={0: 'b'})
for i in range(datalen):
da[i] = i
output_path = temp_filepath('.hdf5')
try:
# make a file, and write to dataset 'foo'
with h5py.File(output_path, 'w') as fp:
fp['foo'] = np.arange(10)
# try saving to a different dataset
dac.save_hdf5(output_path, da, key='bar', mode='a')
with h5py.File(output_path, 'r') as fp:
self.assertTrue("foo" in fp)
self.assertTrue("bar" in fp)
finally:
if os.path.exists(output_path):
os.remove(output_path)
def test_write_3d(self):
h5py = import_or_skip('h5py')
shape = (4, 5, 3)
source = np.random.random(shape)
dac = Context(self.client)
dist = {0: 'b', 1: 'c', 2: 'n'}
da = dac.empty(shape, dist=dist)
for i in range(shape[0]):
for j in range(shape[1]):
for k in range(shape[2]):
da[i, j, k] = source[i, j, k]
output_path = temp_filepath('.hdf5')
try:
dac.save_hdf5(output_path, da, mode='w')
self.assertTrue(os.path.exists(output_path))
with h5py.File(output_path, 'r') as fp:
self.assertTrue("buffer" in fp)
assert_allclose(source, fp["buffer"])
finally:
if os.path.exists(output_path):
os.remove(output_path)
def test_different_contexts(self):
ctx1 = Context(targets=range(4))
ctx2 = Context(targets=range(3))
da1 = ctx1.ones((10,))
da2 = ctx2.ones((10,))
db1 = self.local_sin(da1)
db2 = self.local_sin(da2)
ndarr1 = db1.toarray()
ndarr2 = db2.toarray()
assert_array_equal(ndarr1, ndarr2)
def test_save_load_with_filenames(self):
dac = Context(self.client)
da = dac.empty((100,), dist={0: 'b'})
output_paths = [temp_filepath() for target in dac.targets]
try:
dac.save(output_paths, da)
db = dac.load(output_paths)
self.assertTrue(isinstance(db, DistArray))
self.assertEqual(da, db)
finally:
for filepath in output_paths:
if os.path.exists(filepath):
os.remove(filepath)
def test_save_load_with_prefix(self):
dac = Context(self.client)
da = dac.empty((100,), dist={0: 'b'})
output_path = temp_filepath()
try:
dac.save(output_path, da)
db = dac.load(output_path)
self.assertTrue(isinstance(db, DistArray))
self.assertEqual(da, db)
finally:
for rank in dac.targets:
filepath = output_path + "_" + str(rank) + ".dnpy"
if os.path.exists(filepath):
os.remove(filepath)
def test_determine_context(self):
context = Context()
context2 = Context() # for cross Context checking
da = context.ones((2, 2))
def dummy_func(*args, **kwargs):
fn = lambda x: x
db = DecoratorBase(fn)
return db.determine_context(args, kwargs)
self.assertEqual(dummy_func(6, 7, context), context)
self.assertEqual(dummy_func('ab', da), context)
self.assertEqual(dummy_func(a=da), context)
self.assertEqual(dummy_func(context, a=da), context)
self.assertRaises(TypeError, dummy_func, 'foo')
self.assertRaises(ContextError, dummy_func, context, context2)
class TestDistArrayCreation(unittest.TestCase):
"""Test distarray creation methods"""
def setUp(self):
self.context = Context()
def tearDown(self):
self.context.close()
def test___init__(self):
shape = (100, 100)
distribution = Distribution.from_shape(self.context, shape, ('b', 'c'))
da = DistArray(distribution, dtype=int)
da.fill(42)
nda = numpy.empty(shape, dtype=int)
nda.fill(42)
assert_array_equal(da.tondarray(), nda)
def test_zeros(self):
shape = (16, 16)
zero_distarray = self.context.zeros(shape)
zero_ndarray = numpy.zeros(shape)
assert_array_equal(zero_distarray.tondarray(), zero_ndarray)
def test_ones(self):
shape = (16, 16)
one_distarray = self.context.ones(shape)
one_ndarray = numpy.ones(shape)
assert_array_equal(one_distarray.tondarray(), one_ndarray)
def test_empty(self):
shape = (16, 16)
empty_distarray = self.context.empty(shape)
self.assertEqual(empty_distarray.shape, shape)
def test_fromndarray(self):
ndarr = numpy.arange(16).reshape(4, 4)
distarr = self.context.fromndarray(ndarr)
for (i, j), val in numpy.ndenumerate(ndarr):
self.assertEqual(distarr[i, j], ndarr[i, j])
def test_grid_rank(self):
# regression test for issue #235
a = self.context.empty((4, 4, 4), dist=('b', 'n', 'b'),
grid_shape=(1, 1, 4))
self.assertEqual(a.grid_shape, (1, 1, 4))
def test_fromfunction(self):
fn = lambda i, j: i + j
shape = (7, 9)
expected = numpy.fromfunction(fn, shape, dtype=int)
result = self.context.fromfunction(fn, shape, dtype=int)
assert_array_equal(expected, result.tondarray())
def test_vectorize(self):
"""Test the @vectorize decorator for parity with NumPy's"""
context = Context()
a = numpy.arange(16).reshape(4, 4)
da = context.fromndarray(a)
@vectorize
def da_fn(a, b, c):
return a**2 + b + c
@numpy.vectorize
def a_fn(a, b, c):
return a**2 + b + c
a = a_fn(a, a, 6)
db = da_fn(da, da, 6)
assert_array_equal(db.toarray(), a)
def test_local(self):
"""Test the @local decorator"""
context = Context()
da = context.empty((4, 4))
a = numpy.empty((4, 4))
def fill_a(a):
for (i, j), _ in numpy.ndenumerate(a):
a[i, j] = i + j
return a
@local
def fill_da(da):
for i in da.distribution[0].global_iter:
for j in da.distribution[1].global_iter:
da.global_index[i, j] = i + j
return da
da = fill_da(da)
a = fill_a(a)
assert_array_equal(da.toarray(), a)
def test_write_block(self):
h5py = import_or_skip('h5py')
datalen = 33
dac = Context(self.client)
da = dac.empty((datalen,), dist={0: 'b'})
for i in range(datalen):
da[i] = i
output_path = temp_filepath('.hdf5')
try:
dac.save_hdf5(output_path, da, mode='w')
self.assertTrue(os.path.exists(output_path))
with h5py.File(output_path, 'r') as fp:
self.assertTrue("buffer" in fp)
expected = np.arange(datalen)
assert_equal(expected, fp["buffer"])
finally:
if os.path.exists(output_path):
os.remove(output_path)
class TestLocalDecorator(TestCase):
# Functions for @local decorator tests. These are here so we can
# guarantee they are pushed to the engines before we try to use them.
@local
def local_add50(da):
return da + 50
@local
def local_add_num(da, num):
return da + num
@local
def assert_allclose(da, db):
assert numpy.allclose(da, db), "Arrays not equal within tolerance."
@local
def local_sin(da):
return numpy.sin(da)
@local
def local_sum(da):
return numpy.sum(da.get_localarray())
@local
def call_barrier(da):
from mpi4py import MPI
MPI.COMM_WORLD.Barrier()
return da
@local
def local_add_nums(da, num1, num2, num3):
return da + num1 + num2 + num3
@local
def local_add_distarrayproxies(da, dg):
return da + dg
@local
def local_add_mixed(da, num1, dg, num2):
return da + num1 + dg + num2
@local
def local_add_ndarray(da, num, ndarr):
return da + num + ndarr
@local
def local_add_kwargs(da, num1, num2=55):
return da + num1 + num2
@local
def local_add_supermix(da, num1, db, num2, dc, num3=99, num4=66):
return da + num1 + db + num2 + dc + num3 + num4
@local
def local_none(da):
return None
@local
def parameterless():
"""This is a parameterless function."""
return None
@classmethod
def setUpClass(self):
self.context = Context()
self.da = self.context.empty((5, 5))
self.da.fill(2 * numpy.pi)
def test_local(self):
context = Context()
"""Test the @local decorator"""
da = context.empty((4, 4))
a = numpy.empty((4, 4))
def fill_a(a):
for (i, j), _ in numpy.ndenumerate(a):
a[i, j] = i + j
return a
@local
def fill_da(da):
for i in da.maps[0].global_index:
for j in da.maps[1].global_index:
da[i, j] = i + j
return da
da = fill_da(da)
a = fill_a(a)
assert_array_equal(da.toarray(), a)
def test_local_sin(self):
db = self.local_sin(self.da)
self.assert_allclose(db, 0)
def test_local_add50(self):
dc = self.local_add50(self.da)
self.assert_allclose(dc, 2 * numpy.pi + 50)
def test_local_sum(self):
dd = self.local_sum(self.da)
lshapes = self.da.get_localshapes()
expected = []
for lshape in lshapes:
expected.append(lshape[0] * lshape[1] * (2 * numpy.pi))
for (v, e) in zip(dd, expected):
self.assertAlmostEqual(v, e, places=5)
def test_local_add_num(self):
de = self.local_add_num(self.da, 11)
self.assert_allclose(de, 2 * numpy.pi + 11)
def test_local_add_nums(self):
df = self.local_add_nums(self.da, 11, 12, 13)
self.assert_allclose(df, 2 * numpy.pi + 11 + 12 + 13)
def test_local_add_distarrayproxies(self):
dg = self.context.empty((5, 5))
dg.fill(33)
dh = self.local_add_distarrayproxies(self.da, dg)
self.assert_allclose(dh, 33 + 2 * numpy.pi)
def test_local_add_mixed(self):
di = self.context.empty((5, 5))
di.fill(33)
dj = self.local_add_mixed(self.da, 11, di, 12)
self.assert_allclose(dj, 2 * numpy.pi + 11 + 33 + 12)
@unittest.skip('Locally adding ndarrays not supported.')
def test_local_add_ndarray(self):
shp = self.da.get_localshapes()[0]
ndarr = numpy.empty(shp)
ndarr.fill(33)
dk = self.local_add_ndarray(self.da, 11, ndarr)
self.assert_allclose(dk, 2 * numpy.pi + 11 + 33)
def test_local_add_kwargs(self):
dl = self.local_add_kwargs(self.da, 11, num2=12)
self.assert_allclose(dl, 2 * numpy.pi + 11 + 12)
def test_local_add_supermix(self):
dm = self.context.empty((5, 5))
dm.fill(22)
dn = self.context.empty((5, 5))
dn.fill(44)
do = self.local_add_supermix(self.da, 11, dm, 33, dc=dn, num3=55)
expected = 2 * numpy.pi + 11 + 22 + 33 + 44 + 55 + 66
self.assert_allclose(do, expected)
def test_local_none(self):
dp = self.local_none(self.da)
self.assertTrue(dp is None)
def test_barrier(self):
self.call_barrier(self.da)
def test_parameterless(self):
self.assertRaises(TypeError, self.parameterless)
def test_function_metadata(self):
name = "parameterless"
docstring = """This is a parameterless function."""
self.assertEqual(self.parameterless.__name__, name)
self.assertEqual(self.parameterless.__doc__, docstring)
def setUpClass(self):
self.context = Context()
self.da = self.context.empty((5, 5))
self.da.fill(2 * numpy.pi)
def setUp(self):
self.context = Context()
def setUp(self):
self.dac = Context(self.client)
class TestDistArrayCreationFromGlobalDimData(unittest.TestCase):
def setUp(self):
self.context = Context()
def tearDown(self):
self.context.close()
def test_from_global_dim_data_irregular_block(self):
global_size = 10
bounds = (0, 2, 3, 4, 10)
glb_dim_data = (
{'dist_type': 'b',
'bounds': bounds},
)
distribution = Distribution(self.context, glb_dim_data)
distarr = DistArray(distribution, dtype=int)
for i in range(global_size):
distarr[i] = i
def test_from_global_dim_data_1d(self):
total_size = 40
list_of_indices = [
[29, 38, 18, 19, 11, 33, 10, 1, 22, 25],
[5, 15, 34, 12, 16, 24, 23, 39, 6, 36],
[0, 7, 27, 4, 32, 37, 21, 26, 9, 17],
[35, 14, 20, 13, 3, 30, 2, 8, 28, 31],
]
glb_dim_data = (
{'dist_type': 'u',
'indices': list_of_indices,
},
)
distribution = Distribution(self.context, glb_dim_data)
distarr = DistArray(distribution, dtype=int)
for i in range(total_size):
distarr[i] = i
localarrays = distarr.get_localarrays()
for i, arr in enumerate(localarrays):
assert_allclose(arr, list_of_indices[i])
def test_from_global_dim_data_bu(self):
rows = 9
row_break_point = rows // 2
cols = 10
col_indices = numpy.random.permutation(range(cols))
col_break_point = len(col_indices) // 3
indices = [col_indices[:col_break_point], col_indices[col_break_point:]]
glb_dim_data = (
{
'dist_type': 'b',
'bounds': (0, row_break_point, rows)
},
{
'dist_type': 'u',
'indices' : indices
},
)
distribution = Distribution(self.context, glb_dim_data)
distarr = DistArray(distribution, dtype=int)
for i in range(rows):
for j in range(cols):
distarr[i, j] = i*cols + j
def test_from_global_dim_data_bc(self):
""" Test creation of a block-cyclic array. """
rows, cols = 5, 9
global_dim_data = (
# dim 0
{
'dist_type': 'c',
'proc_grid_size': 2,
'size': rows,
'block_size': 2,
},
# dim 1
{
'dist_type': 'c',
'proc_grid_size': 2,
'size': cols,
'block_size': 2,
},)
distribution = Distribution(self.context, global_dim_data)
distarr = DistArray(distribution, dtype=int)
for i in range(rows):
for j in range(cols):
distarr[i, j] = i*cols + j
las = distarr.get_localarrays()
local_shapes = [la.local_shape for la in las]
self.assertSequenceEqual(local_shapes, [(3,5), (3,4), (2,5), (2,4)])
def test_from_global_dim_data_uu(self):
rows = 6
cols = 20
row_ixs = numpy.random.permutation(range(rows))
col_ixs = numpy.random.permutation(range(cols))
row_indices = [row_ixs[:rows//2], row_ixs[rows//2:]]
col_indices = [col_ixs[:cols//4], col_ixs[cols//4:]]
glb_dim_data = (
{'dist_type': 'u',
'indices': row_indices},
{'dist_type': 'u',
'indices' : col_indices},
)
distribution = Distribution(self.context, glb_dim_data)
distarr = DistArray(distribution, dtype=int)
for i in range(rows):
for j in range(cols):
distarr[i, j] = i*cols + j
def test_global_dim_data_local_dim_data_equivalence(self):
rows, cols = 5, 9
glb_dim_data = (
{'dist_type': 'c',
'block_size': 2,
'size': rows,
'proc_grid_size': 2,
},
{'dist_type': 'c',
'block_size': 2,
'proc_grid_size': 2,
'size': cols,
},
)
distribution = Distribution(self.context, glb_dim_data)
actual = distribution.get_dim_data_per_rank()
expected = [
({'block_size': 2,
'dist_type': 'c',
'proc_grid_rank': 0,
'proc_grid_size': 2,
'size': rows,
'start': 0},
{'block_size': 2,
'dist_type': 'c',
'proc_grid_rank': 0,
'proc_grid_size': 2,
'size': cols,
'start': 0}),
({'block_size': 2,
'dist_type': 'c',
'proc_grid_rank': 0,
'proc_grid_size': 2,
'size': rows,
'start': 0},
{'block_size': 2,
'dist_type': 'c',
'proc_grid_rank': 1,
'proc_grid_size': 2,
'size': cols,
'start': 2}),
({'block_size': 2,
'dist_type': 'c',
'proc_grid_rank': 1,
'proc_grid_size': 2,
'size': rows,
'start': 2},
{'block_size': 2,
'dist_type': 'c',
'proc_grid_rank': 0,
'proc_grid_size': 2,
'size': cols,
'start': 0}),
({'block_size': 2,
'dist_type': 'c',
'proc_grid_rank': 1,
'proc_grid_size': 2,
'size': rows,
'start': 2},
{'block_size': 2,
'dist_type': 'c',
'proc_grid_rank': 1,
'proc_grid_size': 2,
'size': cols,
'start': 2}),
]
self.assertSequenceEqual(actual, expected)
def test_irregular_block_assignment(self):
global_shape = (5, 9)
global_dim_data = (
{
'dist_type': 'b',
'bounds': (0, 5),
},
{
'dist_type': 'b',
'bounds': (0, 2, 6, 7, 9),
}
)
distribution = Distribution(self.context, global_dim_data)
distarr = DistArray(distribution, dtype=int)
for i in range(global_shape[0]):
for j in range(global_shape[1]):
distarr[i, j] = i + j
def setUp(self):
self.context = Context(self.client)
def setUp(self):
self.dac = Context()
class TestDistArray(unittest.TestCase):
def setUp(self):
self.dac = Context()
def tearDown(self):
self.dac.close()
def test_set_and_getitem_block_dist(self):
size = 10
dap = self.dac.empty((size,), dist={0: 'b'})
for val in range(size):
dap[val] = val
for val in range(size):
self.assertEqual(dap[val], val)
for i in range(1, size + 1):
dap[-i] = i
self.assertEqual(dap[-i], i)
def test_set_and_getitem_nd_block_dist(self):
size = 5
dap = self.dac.empty((size, size), dist={0: 'b', 1: 'b'})
for row in range(size):
for col in range(size):
val = size*row + col
dap[row, col] = val
self.assertEqual(dap[row, col], val)
for row in range(1 ,size + 1):
for col in range(1, size + 1):
dap[-row, -col] = row + col
self.assertEqual(dap[-row, -col], row + col)
def test_set_and_getitem_cyclic_dist(self):
size = 10
dap = self.dac.empty((size,), dist={0: 'c'})
for val in range(size):
dap[val] = val
self.assertEqual(dap[val], val)
for i in range(1, size + 1):
dap[-i] = i
self.assertEqual(dap[-i], i)
def test_get_index_error(self):
dap = self.dac.empty((10,), dist={0: 'c'})
with self.assertRaises(IndexError):
dap[11]
with self.assertRaises(IndexError):
dap[-11]
def test_set_index_error(self):
dap = self.dac.empty((10,), dist={0: 'c'})
with self.assertRaises(IndexError):
dap[11] = 55
with self.assertRaises(IndexError):
dap[-11] = 55
def test_iteration(self):
size = 10
dap = self.dac.empty((size,), dist={0: 'c'})
dap.fill(10)
for val in dap:
self.assertEqual(val, 10)
def test_tondarray(self):
dap = self.dac.empty((3, 3))
ndarr = numpy.arange(9).reshape(3, 3)
for (i, j), val in numpy.ndenumerate(ndarr):
dap[i, j] = ndarr[i, j]
numpy.testing.assert_array_equal(dap.tondarray(), ndarr)
def test_global_tolocal_bug(self):
# gh-issue #154
dap = self.dac.zeros((3, 3), dist=('n', 'b'))
ndarr = numpy.zeros((3, 3))
numpy.testing.assert_array_equal(dap.tondarray(), ndarr)
class TestDistArray(IpclusterTestCase):
def setUp(self):
self.dac = Context(self.client)
def test_set_and_getitem_block_dist(self):
size = 10
dap = self.dac.empty((size,), dist={0: 'b'})
for val in range(size):
dap[val] = val
for val in range(size):
self.assertEqual(dap[val], val)
def test_set_and_getitem_nd_block_dist(self):
size = 5
dap = self.dac.empty((size, size), dist={0: 'b', 1: 'b'})
for row in range(size):
for col in range(size):
val = size*row + col
dap[row, col] = val
for row in range(size):
for col in range(size):
val = size*row + col
self.assertEqual(dap[row, col], val)
def test_set_and_getitem_cyclic_dist(self):
size = 10
dap = self.dac.empty((size,), dist={0: 'c'})
for val in range(size):
dap[val] = val
for val in range(size):
self.assertEqual(dap[val], val)
@unittest.skip("Slicing not yet implemented.")
def test_slice_in_getitem_block_dist(self):
dap = self.dac.empty((100,), dist={0: 'b'})
self.assertIsInstance(dap[20:40], DistArray)
@unittest.skip("Slicing not yet implemented.")
def test_slice_in_setitem_raises_valueerror(self):
dap = self.dac.empty((100,), dist={0: 'b'})
vals = numpy.random.random(20)
with self.assertRaises(NotImplementedError):
dap[20:40] = vals
@unittest.skip('Slice assignment not yet implemented.')
def test_slice_size_error(self):
dap = self.dac.empty((100,), dist={0: 'c'})
with self.assertRaises(NotImplementedError):
dap[20:40] = (11, 12)
def test_get_index_error(self):
dap = self.dac.empty((10,), dist={0: 'c'})
with self.assertRaises(IndexError):
dap[11]
def test_set_index_error(self):
dap = self.dac.empty((10,), dist={0: 'c'})
with self.assertRaises(IndexError):
dap[11] = 55
def test_iteration(self):
size = 10
dap = self.dac.empty((size,), dist={0: 'c'})
dap.fill(10)
for val in dap:
self.assertEqual(val, 10)
def test_tondarray(self):
dap = self.dac.empty((3, 3))
ndarr = numpy.arange(9).reshape(3, 3)
for (i, j), val in numpy.ndenumerate(ndarr):
dap[i, j] = ndarr[i, j]
numpy.testing.assert_array_equal(dap.tondarray(), ndarr)
def test_global_tolocal_bug(self):
# gh-issue #154
dap = self.dac.zeros((3, 3), dist=('n', 'b'))
ndarr = numpy.zeros((3, 3))
numpy.testing.assert_array_equal(dap.tondarray(), ndarr)