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
