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_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 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_cbc(self): distribution = Distribution(comm=self.comm, dim_data=self.ddpr[self.comm.Get_rank()]) a = LocalArray(distribution, dtype=np.int_) a.fill(12) b = a.copy() assert_localarrays_equal(a, b, check_dtype=True)
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_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_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_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_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_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_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_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_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_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_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 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=(53, 77), dist={ 0: 'c', 1: 'b' }, grid_shape=(2, 2)) self.larr = LocalArray(d)
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): # 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 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_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_cbc(self): new_dtype = np.int8 d = Distribution(comm=self.comm, dim_data=self.ddpr[self.comm.Get_rank()]) a = LocalArray(d, dtype=np.int32) a.fill(12) 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 _local_call(comm, local_func_name, ddpr, kwargs): import distarray.localapi.random as local_random from distarray.localapi.maps import Distribution local_func = getattr(local_random, local_func_name) if len(ddpr): dim_data = ddpr[comm.Get_rank()] else: dim_data = () dist = Distribution(dim_data=dim_data, comm=comm) return proxyize(local_func(distribution=dist, **kwargs))
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_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_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 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)
def test_create_comm_with_list(self): """ Test that create_comm_with_list() works correctly. """ nodes = list(range(self.max_size)) comm = create_comm_with_list(nodes) if comm == MPI.COMM_NULL: # Only proceed when not COMM_NULL. return # Run a simple test to confirm this comm works. size = len(nodes) nrows = size * 3 d = Distribution.from_shape(comm=comm, shape=(nrows, 20)) a = zeros(d) expected = numpy.zeros((nrows // size, 20)) assert_array_equal(a.ndarray, expected) # Cleanup. comm.Free()
def test_block_complex(self): """Can we compute local indices for a block distribution?""" distribution = Distribution.from_shape(comm=self.comm, shape=(8, 2)) la = LocalArray(distribution) self.assertEqual(la.global_shape, (8, 2)) self.assertEqual(la.grid_shape, (4, 1)) self.assertEqual(la.local_shape, (2, 2)) expected_lis = [(0, 0), (0, 1), (1, 0), (1, 1)] if la.comm_rank == 0: gis = [(0, 0), (0, 1), (1, 0), (1, 1)] elif la.comm_rank == 1: gis = [(2, 0), (2, 1), (3, 0), (3, 1)] elif la.comm_rank == 2: gis = [(4, 0), (4, 1), (5, 0), (5, 1)] elif la.comm_rank == 3: gis = [(6, 0), (6, 1), (7, 0), (7, 1)] result = [la.local_from_global(gi) for gi in gis] self.assertEqual(result, expected_lis)
def test_empty(self): size = self.comm_size nrows = size * 3 d = Distribution.from_shape(comm=self.comm, shape=(nrows, 20)) a = localarray.empty(d) self.assertEqual(a.global_shape, (nrows, 20))
def test_create_localarray(self): # regression test for issue #144 dist = Distribution.from_shape(comm=self.comm, shape=(16, 16), dist=('n', 'b')) la = LocalArray(dist)
def test_bad_global_limits(self): """ Test that invalid global_limits fails as expected. """ d = Distribution.from_shape(comm=self.comm, shape=(4, 4)) a = LocalArray(d) with self.assertRaises(InvalidDimensionError): a.global_limits(-1)