def test_indexing_1(self): """Can we get and set local elements for a complex dist?""" a = da.LocalArray((16,16,2), dist=('c', 'b', 'n'), comm=self.comm) b = da.LocalArray((16,16,2), dist=('c', 'b', 'n'), comm=self.comm) for global_inds, value in da.ndenumerate(a): a[global_inds] = 0.0 for global_inds, value in da.ndenumerate(a): b[global_inds] = a[global_inds] for global_inds, value in da.ndenumerate(a): self.assertEqual(b[global_inds],a[global_inds]) self.assertEqual(a[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) a = dla.fromfunction(f, (16, 16), dtype='int64', dist=('b', 'c'), comm=self.comm) self.assertEqual(a.shape, (16,16)) self.assertEqual(a.dtype, np.dtype('int64')) for global_inds, value in dla.ndenumerate(a): self.assertEqual(sum(global_inds), value)
def test_pack_unpack_index(self): a = da.LocalArray((16,16,2), dist=('c', 'b', 'n'), comm=self.comm) for global_inds, value in da.ndenumerate(a): packed_ind = a.pack_index(global_inds) self.assertEqual(global_inds, a.unpack_index(packed_ind))
def test_ndenumerate(self): a = dla.LocalArray((16, 16, 2), dist=('c', 'b', 'n'), comm=self.comm) for global_inds, value in dla.ndenumerate(a): a[global_inds] = 0.0