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_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_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_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_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_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_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
