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