Python localarray.LocalArray类代码示例

    def test_add(self):
        """See if binary ufunc works for a LocalArray."""
        d = Distribution.from_shape(comm=self.comm, shape=(16, 16))
        a = LocalArray(d, dtype='int32')
        b = LocalArray(d, dtype='int32')
        a.fill(1)
        b.fill(1)
        c = localarray.add(a, b)
        self.assertTrue(np.all(c.ndarray == 2))

        c = localarray.empty_like(a)
        c = localarray.add(a, b, c)
        self.assertTrue(np.all(c.ndarray == 2))

        d0 = Distribution.from_shape(comm=self.comm, shape=(16, 16))
        d1 = Distribution.from_shape(comm=self.comm, shape=(20, 20))
        a = LocalArray(d0, dtype='int32')
        b = LocalArray(d1, dtype='int32')
        self.assertRaises(IncompatibleArrayError, localarray.add, a, b)

        d0 = Distribution.from_shape(comm=self.comm, shape=(16, 16))
        d1 = Distribution.from_shape(comm=self.comm, shape=(20, 20))
        a = LocalArray(d0, dtype='int32')
        b = LocalArray(d0, dtype='int32')
        c = LocalArray(d1, dtype='int32')
        self.assertRaises(IncompatibleArrayError, localarray.add, a, b, c)

 def test_view_bn(self):
     d = Distribution.from_shape((16, 16), dist=('b', 'n'), comm=self.comm)
     a = LocalArray(d, dtype=np.int32)
     a.fill(11)
     b = a.view()
     assert_localarrays_equal(a, b)
     self.assertEqual(id(a.local_data), id(b.local_data))

class TestArrayConversion(MpiTestCase):
    """ Test array conversion methods. """

    def setUp(self):
        # On Python3, an 'int' gets converted to 'np.int64' on copy,
        # so we force the numpy type to start with so we get back
        # the same thing.
        self.int_type = np.int64
        self.distribution = Distribution.from_shape(comm=self.comm,
                                                    shape=(4,))
        self.int_larr = LocalArray(self.distribution, dtype=self.int_type)
        self.int_larr.fill(3)

    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_copy_bn(self):
     distribution = Distribution.from_shape(comm=self.comm,
                                     shape=(16, 16), dist=('b', 'n'))
     a = LocalArray(distribution, dtype=np.int_)
     a.fill(11)
     b = a.copy()
     assert_localarrays_equal(a, b, check_dtype=True)

 def test_copy_cbc(self):
     distribution = Distribution(self.ddpr[self.comm.Get_rank()],
                                 comm=self.comm)
     a = LocalArray(distribution, dtype=np.int_)
     a.fill(12)
     b = a.copy()
     assert_localarrays_equal(a, b, check_dtype=True)

    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 randn(distribution=None):
    if distribution is None:
        return np.random.randn()
    else:
        dtype = np.random.randn(1).dtype
        la = LocalArray(distribution, dtype=dtype)
        la.ndarray[:] = np.random.randn(*la.local_shape)
        return la

    def setUp(self):
        self.dist_1d = Distribution.from_shape(comm=self.comm,
                                        shape=(7,), grid_shape=(4,))
        self.larr_1d = LocalArray(self.dist_1d, buf=None)

        self.dist_2d = Distribution.from_shape(comm=self.comm,
                                        shape=(16, 16), grid_shape=(4, 1))
        self.larr_2d = LocalArray(self.dist_2d, buf=None)

def normal(loc=0.0, scale=1.0, distribution=None):
    if distribution is None:
        return np.random.normal(loc, scale)
    else:
        dtype = np.random.normal(loc, scale, size=1).dtype
        la = LocalArray(distribution, dtype=dtype)
        la.ndarray[:] = np.random.normal(loc, scale, size=la.local_shape)
        return la

def beta(a, b, distribution=None):
    if distribution is None:
        return np.random.beta(a, b)
    else:
        dtype = np.random.beta(a, b, size=1).dtype
        la = LocalArray(distribution, dtype=dtype)
        la.ndarray[:] = np.random.beta(a, b, size=la.local_shape)
        return la

def randint(low, high=None, distribution=None):
    if distribution is None:
        return np.random.randint(low, high)
    else:
        dtype = np.random.randint(low, high, size=1).dtype
        la = LocalArray(distribution, dtype=dtype)
        la.ndarray[:] = np.random.randint(low, high, size=la.local_shape)
        return la

 def test_astype_bn(self):
     new_dtype = np.float32
     d = Distribution.from_shape((16, 16), dist=('b', 'n'),  comm=self.comm)
     a = LocalArray(d, dtype=np.int_)
     a.fill(11)
     b = a.astype(new_dtype)
     assert_localarrays_allclose(a, b, check_dtype=False)
     self.assertEqual(b.dtype, new_dtype)
     self.assertEqual(b.ndarray.dtype, new_dtype)

 def test_astype_cbc(self):
     new_dtype = np.int8
     d = Distribution(comm=self.comm, dim_data=self.ddpr[self.comm.Get_rank()])
     a = LocalArray(d, dtype=np.int32)
     a.fill(12)
     b = a.astype(new_dtype)
     assert_localarrays_allclose(a, b, check_dtype=False)
     self.assertEqual(b.dtype, new_dtype)
     self.assertEqual(b.ndarray.dtype, new_dtype)

 def test_global_limits_cyclic(self):
     """Find the boundaries of a cyclic distribution"""
     d = Distribution.from_shape((16, 16), dist=('c', 'n'), comm=self.comm)
     a = LocalArray(d)
     answers = [(0, 12), (1, 13), (2, 14), (3, 15)]
     limits = a.global_limits(0)
     self.assertEqual(limits, answers[a.comm_rank])
     answers = 4 * [(0, 15)]
     limits = a.global_limits(1)
     self.assertEqual(limits, answers[a.comm_rank])