def test_add(self): """See if binary ufunc works for a LocalArray.""" d = Distribution.from_shape(comm=self.comm, shape=(16, 16)) a = LocalArray(d, dtype='int32') b = LocalArray(d, dtype='int32') a.fill(1) b.fill(1) c = localarray.add(a, b) self.assertTrue(np.all(c.ndarray == 2)) c = localarray.empty_like(a) c = localarray.add(a, b, c) self.assertTrue(np.all(c.ndarray == 2)) d0 = Distribution.from_shape(comm=self.comm, shape=(16, 16)) d1 = Distribution.from_shape(comm=self.comm, shape=(20, 20)) a = LocalArray(d0, dtype='int32') b = LocalArray(d1, dtype='int32') self.assertRaises(IncompatibleArrayError, localarray.add, a, b) d0 = Distribution.from_shape(comm=self.comm, shape=(16, 16)) d1 = Distribution.from_shape(comm=self.comm, shape=(20, 20)) a = LocalArray(d0, dtype='int32') b = LocalArray(d0, dtype='int32') c = LocalArray(d1, dtype='int32') self.assertRaises(IncompatibleArrayError, localarray.add, a, b, c)
def test_add(self): """See if binary ufunc works for a LocalArray.""" d = Distribution.from_shape(comm=self.comm, shape=(16, 16)) a = LocalArray(d, dtype='int32') b = LocalArray(d, dtype='int32') a.fill(1) b.fill(1) c = localarray.add(a, b) self.assertTrue(np.all(c.ndarray == 2)) c = localarray.empty_like(a) c = localarray.add(a, b, c) self.assertTrue(np.all(c.ndarray == 2)) d0 = Distribution.from_shape(comm=self.comm, shape=(16, 16)) d1 = Distribution.from_shape(comm=self.comm, shape=(20, 20)) a = LocalArray(d0, dtype='int32') b = LocalArray(d1, dtype='int32') self.assertRaises(IncompatibleArrayError, localarray.add, a, b) d0 = Distribution.from_shape(comm=self.comm, shape=(16, 16)) d1 = Distribution.from_shape(comm=self.comm, shape=(20, 20)) a = LocalArray(d0, dtype='int32') b = LocalArray(d0, dtype='int32') c = LocalArray(d1, dtype='int32') self.assertRaises(IncompatibleArrayError, localarray.add, a, b, c)
def test_bad_distribution(self): """Test that invalid distribution type fails as expected.""" with self.assertRaises(ValueError): # Invalid distribution type 'x'. Distribution.from_shape(comm=self.comm, shape=(7, ), dist={0: 'x'}, grid_shape=(4, ))
def setUp(self): self.dist_1d = Distribution.from_shape(comm=self.comm, shape=(7,), grid_shape=(4,)) self.larr_1d = LocalArray(self.dist_1d, buf=None) self.dist_2d = Distribution.from_shape(comm=self.comm, shape=(16, 16), grid_shape=(4, 1)) self.larr_2d = LocalArray(self.dist_2d, buf=None)
def setUp(self): self.dist_1d = Distribution.from_shape(comm=self.comm, shape=(7, ), grid_shape=(4, )) self.larr_1d = LocalArray(self.dist_1d, buf=None) self.dist_2d = Distribution.from_shape(comm=self.comm, shape=(16, 16), grid_shape=(4, 1)) self.larr_2d = LocalArray(self.dist_2d, buf=None)
def test_arecompatible(self): """Test if two DistArrays are compatible.""" d0 = Distribution.from_shape(comm=self.comm, shape=(16,16)) a = LocalArray(d0, dtype='int64') b = LocalArray(d0, dtype='float32') self.assertTrue(arecompatible(a,b)) da = Distribution.from_shape(comm=self.comm, shape=(16, 16), dist='c') a = LocalArray(da, dtype='int64') db = Distribution.from_shape(comm=self.comm, shape=(16, 16), dist='b') b = LocalArray(db, dtype='float32') self.assertFalse(arecompatible(a,b))
def test_asdist_like(self): """Test asdist_like for success and failure.""" d = Distribution.from_shape(comm=self.comm, shape=(16, 16), dist=('b', 'n')) a = LocalArray(d) b = LocalArray(d) new_a = a.asdist_like(b) self.assertEqual(id(a), id(new_a)) d2 = Distribution.from_shape(comm=self.comm, shape=(16, 16), dist=('n', 'b')) a = LocalArray(d) b = LocalArray(d2) self.assertRaises(IncompatibleArrayError, a.asdist_like, b)
def test_copy_bn(self): distribution = Distribution.from_shape(comm=self.comm, shape=(16, 16), dist=('b', 'n')) a = LocalArray(distribution, dtype=np.int_) a.fill(11) b = a.copy() assert_localarrays_equal(a, b, check_dtype=True)
def test_cyclic_simple(self): """Can we compute local indices for a cyclic distribution?""" distribution = Distribution.from_shape(comm=self.comm, shape=(10,), dist={0: 'c'}) la = LocalArray(distribution) self.assertEqual(la.global_shape, (10,)) self.assertEqual(la.grid_shape, (4,)) if la.comm_rank == 0: gis = (0, 4, 8) self.assertEqual(la.local_shape, (3,)) calc_result = [la.local_from_global((gi,)) for gi in gis] result = [(0,), (1,), (2,)] elif la.comm_rank == 1: gis = (1, 5, 9) self.assertEqual(la.local_shape, (3,)) calc_result = [la.local_from_global((gi,)) for gi in gis] result = [(0,), (1,), (2,)] elif la.comm_rank == 2: gis = (2, 6) self.assertEqual(la.local_shape, (2,)) calc_result = [la.local_from_global((gi,)) for gi in gis] result = [(0,), (1,)] elif la.comm_rank == 3: gis = (3, 7) self.assertEqual(la.local_shape, (2,)) calc_result = [la.local_from_global((gi,)) for gi in gis] result = [(0,), (1,)] self.assertEqual(result, calc_result)
def test_asdist_like(self): """Test asdist_like for success and failure.""" d = Distribution.from_shape(comm=self.comm, shape=(16, 16), dist=('b', 'n')) a = LocalArray(d) b = LocalArray(d) new_a = a.asdist_like(b) self.assertEqual(id(a), id(new_a)) d2 = Distribution.from_shape(comm=self.comm, shape=(16, 16), dist=('n', 'b')) a = LocalArray(d) b = LocalArray(d2) self.assertRaises(IncompatibleArrayError, a.asdist_like, b)
def test_rank_coords(self): """ Test that we can go from rank to coords and back. """ d = Distribution.from_shape(comm=self.comm, shape=(4, 4)) la = LocalArray(d) max_size = self.comm.Get_size() for rank in range(max_size): self.round_trip(la, rank=rank)
def test_fromndarray_like(self): """Can we build an array using fromndarray_like?""" d = Distribution.from_shape(comm=self.comm, shape=(16, 16), dist=('c', 'b')) la0 = LocalArray(d, dtype='int8') la1 = localarray.fromndarray_like(la0.ndarray, la0) assert_localarrays_equal(la0, la1, check_dtype=True)
def test_ones(self): size = self.comm_size nrows = size * 3 d = Distribution.from_shape(comm=self.comm, shape=(nrows, 20)) a = localarray.ones(d) expected = np.ones((nrows // size, 20)) assert_array_equal(a.ndarray, expected)
def test_rank_coords(self): """ Test that we can go from rank to coords and back. """ d = Distribution.from_shape(comm=self.comm, shape=(4, 4)) la = LocalArray(d) max_size = self.comm.Get_size() for rank in range(max_size): self.round_trip(la, rank=rank)
def test_crazy(self): """Can we go from global to local indices and back for a complex case?""" distribution = Distribution.from_shape(comm=self.comm, shape=(10, 100, 20), dist=('b', 'c', 'n')) la = LocalArray(distribution) self.round_trip(la)
def test_slicing(self): distribution = Distribution.from_shape(self.comm, (16, 16), dist=('b', 'n')) a = ones(distribution) if self.comm.Get_rank() == 0: dd00 = {"dist_type": 'b', "size": 5, "start": 0, "stop": 3, "proc_grid_size": 2, "proc_grid_rank": 0} dd01 = {"dist_type": 'n', "size": 16} new_distribution = Distribution(self.comm, [dd00, dd01]) rvals = a.global_index.get_slice((slice(5, None), slice(None)), new_distribution=new_distribution) assert_array_equal(rvals, np.ones((3, 16))) elif self.comm.Get_rank() == 1: dd10 = {"dist_type": 'b', "size": 5, "start": 3, "stop": 5, "proc_grid_size": 2, "proc_grid_rank": 1} dd11 = {"dist_type": 'n', "size": 16} new_distribution = Distribution(self.comm, [dd10, dd11]) rvals = a.global_index.get_slice((slice(None, 10), slice(None)), new_distribution=new_distribution) assert_array_equal(rvals, np.ones((2, 16)))
def test_pack_unpack_index(self): distribution = Distribution.from_shape(comm=self.comm, shape=(16, 16, 2), dist=('c', 'b', 'n')) a = LocalArray(distribution) for global_inds, value in ndenumerate(a): packed_ind = a.pack_index(global_inds) self.assertEqual(global_inds, a.unpack_index(packed_ind))
def test_slicing(self): distribution = Distribution.from_shape(self.comm, (16, 16), dist=('b', 'n')) a = ones(distribution) if self.comm.Get_rank() == 0: dd00 = { "dist_type": 'b', "size": 5, "start": 0, "stop": 3, "proc_grid_size": 2, "proc_grid_rank": 0 } dd01 = {"dist_type": 'n', "size": 16} new_distribution = Distribution(self.comm, [dd00, dd01]) rvals = a.global_index.get_slice((slice(5, None), slice(None)), new_distribution=new_distribution) assert_array_equal(rvals, np.ones((3, 16))) elif self.comm.Get_rank() == 1: dd10 = { "dist_type": 'b', "size": 5, "start": 3, "stop": 5, "proc_grid_size": 2, "proc_grid_rank": 1 } dd11 = {"dist_type": 'n', "size": 16} new_distribution = Distribution(self.comm, [dd10, dd11]) rvals = a.global_index.get_slice((slice(None, 10), slice(None)), new_distribution=new_distribution) assert_array_equal(rvals, np.ones((2, 16)))
def test_negative(self): """See if unary ufunc works for a LocalArray.""" d = Distribution.from_shape(comm=self.comm, shape=(16, 16)) a = LocalArray(d, dtype='int32') a.fill(1) b = localarray.negative(a) self.assertTrue(np.all(a.ndarray == -b.ndarray)) b = localarray.empty_like(a) b = localarray.negative(a, b) self.assertTrue(np.all(a.ndarray == -b.ndarray)) d2 = Distribution.from_shape(comm=self.comm, shape=(20, 20)) a = LocalArray(d, dtype='int32') b = LocalArray(d2, dtype='int32') self.assertRaises(IncompatibleArrayError, localarray.negative, b, a)
def test_ndenumerate(self): d = Distribution.from_shape(comm=self.comm, shape=(16, 16, 2), dist=('c', 'b', 'n')) a = LocalArray(d) for global_inds, value in ndenumerate(a): a.global_index[global_inds] = 0.0
def test_crazy(self): """Can we go from global to local indices and back for a complex case?""" distribution = Distribution.from_shape(comm=self.comm, shape=(10, 100, 20), dist=('b', 'c', 'n')) la = LocalArray(distribution) self.round_trip(la)
def test_cyclic_simple(self): """Can we compute local indices for a cyclic distribution?""" distribution = Distribution.from_shape(comm=self.comm, shape=(10, ), dist={0: 'c'}) la = LocalArray(distribution) self.assertEqual(la.global_shape, (10, )) self.assertEqual(la.grid_shape, (4, )) if la.comm_rank == 0: gis = (0, 4, 8) self.assertEqual(la.local_shape, (3, )) calc_result = [la.local_from_global((gi, )) for gi in gis] result = [(0, ), (1, ), (2, )] elif la.comm_rank == 1: gis = (1, 5, 9) self.assertEqual(la.local_shape, (3, )) calc_result = [la.local_from_global((gi, )) for gi in gis] result = [(0, ), (1, ), (2, )] elif la.comm_rank == 2: gis = (2, 6) self.assertEqual(la.local_shape, (2, )) calc_result = [la.local_from_global((gi, )) for gi in gis] result = [(0, ), (1, )] elif la.comm_rank == 3: gis = (3, 7) self.assertEqual(la.local_shape, (2, )) calc_result = [la.local_from_global((gi, )) for gi in gis] result = [(0, ), (1, )] self.assertEqual(result, calc_result)
def test_cyclic(self): """Can we go from global to local indices and back for cyclic?""" distribution = Distribution.from_shape(comm=self.comm, shape=(8, 8), dist=('c', 'n')) la = LocalArray(distribution) self.round_trip(la)
def test_negative(self): """See if unary ufunc works for a LocalArray.""" d = Distribution.from_shape(comm=self.comm, shape=(16, 16)) a = LocalArray(d, dtype='int32') a.fill(1) b = localarray.negative(a) self.assertTrue(np.all(a.ndarray == -b.ndarray)) b = localarray.empty_like(a) b = localarray.negative(a, b) self.assertTrue(np.all(a.ndarray == -b.ndarray)) d2 = Distribution.from_shape(comm=self.comm, shape=(20, 20)) a = LocalArray(d, dtype='int32') b = LocalArray(d2, dtype='int32') self.assertRaises(IncompatibleArrayError, localarray.negative, b, a)
def test_pack_unpack_index(self): distribution = Distribution.from_shape(comm=self.comm, shape=(16, 16, 2), dist=('c', 'b', 'n')) a = LocalArray(distribution) for global_inds, value in ndenumerate(a): packed_ind = a.pack_index(global_inds) self.assertEqual(global_inds, a.unpack_index(packed_ind))
def test_copy_bn(self): distribution = Distribution.from_shape(comm=self.comm, shape=(16, 16), dist=('b', 'n')) a = LocalArray(distribution, dtype=np.int_) a.fill(11) b = a.copy() assert_localarrays_equal(a, b, check_dtype=True)
def test_abs(self): """See if unary ufunc works for a LocalArray.""" d = Distribution.from_shape(comm=self.comm, shape=(16, 16)) a = LocalArray(d, dtype='int32') a.fill(-5) a[2, 3] = 11 b = abs(a) self.assertTrue(np.all(abs(a.ndarray) == b.ndarray))
def test_abs(self): """See if unary ufunc works for a LocalArray.""" d = Distribution.from_shape(comm=self.comm, shape=(16, 16)) a = LocalArray(d, dtype='int32') a.fill(-5) a[2, 3] = 11 b = abs(a) self.assertTrue(np.all(abs(a.ndarray) == b.ndarray))
def test_cart_coords(self): """Test getting the cart_coords attribute.""" self.dist_1d = Distribution.from_shape(comm=self.comm, shape=(7, ), grid_shape=(4, )) actual_1d = self.dist_1d.cart_coords expected_1d = tuple( self.dist_1d.comm.Get_coords(self.dist_1d.comm_rank)) self.assertEqual(actual_1d, expected_1d) self.dist_2d = Distribution.from_shape(comm=self.comm, shape=(16, 16), grid_shape=(4, 1)) actual_2d = self.dist_2d.cart_coords expected_2d = tuple( self.dist_2d.comm.Get_coords(self.dist_2d.comm_rank)) self.assertEqual(actual_2d, expected_2d)
def setUp(self): # On Python3, an 'int' gets converted to 'np.int64' on copy, # so we force the numpy type to start with so we get back # the same thing. self.int_type = np.int64 self.distribution = Distribution.from_shape(comm=self.comm, shape=(4, )) self.int_larr = LocalArray(self.distribution, dtype=self.int_type) self.int_larr.fill(3)
def setUp(self): d = Distribution.from_shape(comm=self.comm, shape=(30, 60), dist={ 0: 'n', 1: 'b' }, grid_shape=(1, 4)) self.larr = LocalArray(d)
def setUp(self): d = Distribution.from_shape(comm=self.comm, shape=(53, 77), dist={ 0: 'c', 1: 'b' }, grid_shape=(2, 2)) self.larr = LocalArray(d)
def test_no_grid_shape(self): """Create array init when passing no grid_shape.""" dist_nogrid = Distribution.from_shape(comm=self.comm, shape=(7, )) grid_shape = dist_nogrid.grid_shape # For 1D array as created here, we expect grid_shape # to just be the number of engines as a tuple. max_size = self.comm.Get_size() expected_grid_shape = (max_size, ) self.assertEqual(grid_shape, expected_grid_shape)
def setUp(self): # On Python3, an 'int' gets converted to 'np.int64' on copy, # so we force the numpy type to start with so we get back # the same thing. self.int_type = np.int64 self.distribution = Distribution.from_shape(comm=self.comm, shape=(4,)) self.int_larr = LocalArray(self.distribution, dtype=self.int_type) self.int_larr.fill(3)
def setUp(self): d = Distribution.from_shape(comm=self.comm, shape=(16, 16), dist={ 0: 'b', 1: 'n' }, grid_shape=(4, 1)) self.larr = LocalArray(d)
def test_save_2d(self): d = Distribution.from_shape(comm=self.comm, shape=(11, 15)) la = LocalArray(d) np_arr = numpy.random.random(la.local_shape) la.ndarray = np_arr save_hdf5(self.output_path, la, key=self.key, mode='w') with self.h5py.File(self.output_path, 'r', driver='mpio', comm=self.comm) as fp: for i, v in ndenumerate(la): self.assertEqual(v, fp[self.key][i])
def setUp(self): d = Distribution.from_shape(comm=self.comm, shape=(53, 77, 99), dist={ 0: 'b', 1: 'n', 2: 'c' }, grid_shape=(2, 1, 2)) self.larr = LocalArray(d, dtype='float64')
def test_astype_bn(self): new_dtype = np.float32 d = Distribution.from_shape(comm=self.comm, shape=(16, 16), dist=('b', 'n')) a = LocalArray(d, dtype=np.int_) a.fill(11) b = a.astype(new_dtype) assert_localarrays_allclose(a, b, check_dtype=False) self.assertEqual(b.dtype, new_dtype) self.assertEqual(b.ndarray.dtype, new_dtype)
def test_astype_bn(self): new_dtype = np.float32 d = Distribution.from_shape(comm=self.comm, shape=(16, 16), dist=('b', 'n')) a = LocalArray(d, dtype=np.int_) a.fill(11) b = a.astype(new_dtype) assert_localarrays_allclose(a, b, check_dtype=False) self.assertEqual(b.dtype, new_dtype) self.assertEqual(b.ndarray.dtype, new_dtype)
def test_global_limits_cyclic(self): """Find the boundaries of a cyclic distribution""" d = Distribution.from_shape(comm=self.comm, shape=(16, 16), dist=('c', 'n')) a = LocalArray(d) answers = [(0, 12), (1, 13), (2, 14), (3, 15)] limits = a.global_limits(0) self.assertEqual(limits, answers[a.comm_rank]) answers = 4 * [(0, 15)] limits = a.global_limits(1) self.assertEqual(limits, answers[a.comm_rank])
def test_fromfunction_complicated(self): """Can we build an array using fromfunction and a nontrivial function.""" def f(*global_inds): return sum(global_inds) d = Distribution.from_shape(comm=self.comm, shape=(16, 16), dist=('b', 'c')) a = localarray.fromfunction(f, d, dtype='int64') self.assertEqual(a.global_shape, (16,16)) self.assertEqual(a.dtype, np.dtype('int64')) for global_inds, value in localarray.ndenumerate(a): self.assertEqual(sum(global_inds), value)
def test_save_2d(self): d = Distribution.from_shape(comm=self.comm, shape=(11, 15)) la = LocalArray(d) np_arr = numpy.random.random(la.local_shape) la.ndarray = np_arr save_hdf5(self.output_path, la, key=self.key, mode='w') with self.h5py.File(self.output_path, 'r', driver='mpio', comm=self.comm) as fp: for i, v in ndenumerate(la): self.assertEqual(v, fp[self.key][i])
def test_global_limits_cyclic(self): """Find the boundaries of a cyclic distribution""" d = Distribution.from_shape(comm=self.comm, shape=(16, 16), dist=('c', 'n')) a = LocalArray(d) answers = [(0, 12), (1, 13), (2, 14), (3, 15)] limits = a.global_limits(0) self.assertEqual(limits, answers[a.comm_rank]) answers = 4 * [(0, 15)] limits = a.global_limits(1) self.assertEqual(limits, answers[a.comm_rank])
def test_grid_shape(self): """Test various ways of setting the grid_shape.""" dist = Distribution.from_shape(comm=self.comm, shape=(20, 20), dist='b') self.assertEqual(dist.grid_shape, (12, 1)) dist = Distribution.from_shape(comm=self.comm, shape=(2 * 10, 6 * 10), dist=('b', 'b')) self.assertEqual(dist.grid_shape, (2, 6)) dist = Distribution.from_shape(comm=self.comm, shape=(6 * 10, 2 * 10), dist='bb') self.assertEqual(dist.grid_shape, (6, 2)) dist = Distribution.from_shape(comm=self.comm, shape=(100, 10, 300), dist='bnc') self.assertEqual(dist.grid_shape, (2, 1, 6)) dist = Distribution.from_shape(comm=self.comm, shape=(100, 50, 300), dist='bbb') self.assertEqual(dist.grid_shape, (2, 1, 6)) dist = Distribution.from_shape(comm=self.comm, shape=(100, 100, 150), dist='bbb') self.assertEqual(dist.grid_shape, (2, 2, 3))
def test_block_simple(self): """Can we compute local indices for a block distribution?""" distribution = Distribution.from_shape(comm=self.comm, shape=(4, 4)) la = LocalArray(distribution) self.assertEqual(la.global_shape, (4, 4)) self.assertEqual(la.grid_shape, (4, 1)) self.assertEqual(la.local_shape, (1, 4)) row_result = [(0, 0), (0, 1), (0, 2), (0, 3)] row = la.comm_rank calc_row_result = [la.local_from_global((row, col)) for col in range(la.global_shape[1])] self.assertEqual(row_result, calc_row_result)
def test_indexing_1(self): """Can we get and set local elements for a complex dist?""" distribution = Distribution.from_shape(comm=self.comm, shape=(16, 16, 2), dist=('c', 'b', 'n')) a = LocalArray(distribution) b = LocalArray(distribution) for i, value in ndenumerate(a): a.global_index[i] = 0.0 for i, value in ndenumerate(a): b.global_index[i] = a.global_index[i] for i, value in ndenumerate(a): self.assertEqual(b.global_index[i], a.global_index[i]) self.assertEqual(a.global_index[i], 0.0)
def test_global_limits_block(self): """Find the boundaries of a block distribution""" d = Distribution.from_shape(comm=self.comm, shape=(16, 16), dist=('b', 'n')) a = LocalArray(d) answers = [(0, 3), (4, 7), (8, 11), (12, 15)] limits = a.global_limits(0) self.assertEqual(limits, answers[a.comm_rank]) answers = 4 * [(0, 15)] limits = a.global_limits(1) self.assertEqual(limits, answers[a.comm_rank])
def test_indexing_1(self): """Can we get and set local elements for a complex dist?""" distribution = Distribution.from_shape(comm=self.comm, shape=(16, 16, 2), dist=('c', 'b', 'n')) a = LocalArray(distribution) b = LocalArray(distribution) for i, value in ndenumerate(a): a.global_index[i] = 0.0 for i, value in ndenumerate(a): b.global_index[i] = a.global_index[i] for i, value in ndenumerate(a): self.assertEqual(b.global_index[i], a.global_index[i]) self.assertEqual(a.global_index[i], 0.0)
def check_op(self, op): """Check unary operation for success. Check the one- and two-arg ufunc versions as well as the method version attached to a LocalArray. """ d = Distribution.from_shape(comm=self.comm, shape=(16, 16), dist=('b', 'n')) x = localarray.ones(d) y = localarray.ones(d) with warnings.catch_warnings(): warnings.simplefilter("ignore", category=RuntimeWarning) result0 = op(x, casting='unsafe') # standard form op(x, y=y, casting='unsafe') # two-arg form assert_array_equal(result0.ndarray, y.ndarray)