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_astype(self): """ Test that astype() works as expected. """ # Convert int array to float. float_larr = self.int_larr.astype(float) for global_inds, value in ndenumerate(float_larr): self.assertEqual(value, 3.0) self.assertTrue(isinstance(value, float)) # No type specification for a copy. # Should get same type as we started with. int_larr2 = self.int_larr.astype(None) for global_inds, value in ndenumerate(int_larr2): self.assertEqual(value, 3) self.assertTrue(isinstance(value, self.int_type))
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_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_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)