xref: /petsc/src/binding/petsc4py/test/test_vec.py (revision dda1fc5685d49b41e9220fbfbd490f082712efdf)

from petsc4py import PETSc
import unittest

# --------------------------------------------------------------------

class BaseTestVec(object):

    COMM = None
    TYPE = None

    def setUp(self):
        v = PETSc.Vec()
        v.create(self.COMM)
        v.setSizes(100)
        v.setType(self.TYPE)
        self.vec = v

    def tearDown(self):
        self.vec.destroy()
        self.vec = None
        PETSc.garbage_cleanup()

    def testDuplicate(self):
        self.vec.set(1)
        vec = self.vec.duplicate()
        self.assertFalse(self.vec.equal(vec))
        self.assertEqual(self.vec.sizes, vec.sizes)
        del vec

    def testCopy(self):
        self.vec.set(1)
        vec = self.vec.duplicate()
        self.vec.copy(vec)
        self.assertTrue(self.vec.equal(vec))
        del vec

    def testDot(self):
        self.vec.set(1)
        d = self.vec.dot(self.vec)
        self.assertAlmostEqual(abs(d), self.vec.getSize())
        self.vec.dotBegin(self.vec)
        d = self.vec.dotEnd(self.vec)
        self.assertAlmostEqual(abs(d), self.vec.getSize())

    def testNorm(self):
        from math import sqrt
        self.vec.set(1)
        n1 = self.vec.norm(PETSc.NormType.NORM_1)
        n2 = self.vec.norm(PETSc.NormType.NORM_2)
        ni = self.vec.norm(PETSc.NormType.NORM_INFINITY)
        self.assertAlmostEqual(n1, self.vec.getSize())
        self.assertAlmostEqual(n2, sqrt(self.vec.getSize()))
        self.assertAlmostEqual(n2, self.vec.norm())
        self.assertAlmostEqual(ni, 1.0)
        self.vec.normBegin(PETSc.NormType.NORM_1)
        nn1 = self.vec.normEnd(PETSc.NormType.NORM_1)
        self.assertAlmostEqual(nn1, n1)
        self.vec.normBegin()
        nn2 = self.vec.normEnd()
        self.assertAlmostEqual(nn2, n2)
        self.vec.normBegin(PETSc.NormType.NORM_INFINITY)
        nni = self.vec.normEnd(PETSc.NormType.NORM_INFINITY)
        self.assertAlmostEqual(nni, ni)

    def testNormalize(self):
        from math import sqrt
        self.vec.set(1)
        n2 = self.vec.normalize()
        self.assertAlmostEqual(n2, sqrt(self.vec.getSize()))
        self.assertAlmostEqual(1, self.vec.norm())

    def testSumMinMax(self):
        self.vec.set(1)
        self.assertEqual(self.vec.sum(), self.vec.getSize())
        self.vec.set(-7)
        self.assertEqual(self.vec.min()[1], -7)
        self.vec.set(10)
        self.assertEqual(self.vec.max()[1], 10)

    def testSwap(self):
        v1 = self.vec
        v2 = v1.duplicate()
        v1.set(1)
        v2.set(2)
        v1.swap(v2)
        idx, _ = self.vec.getOwnershipRange()
        self.assertEqual(v1[idx], 2)
        self.assertEqual(v2[idx], 1)

    def testBsize(self):
        self.vec.setBlockSize(1)
        self.assertEqual(self.vec.getBlockSize(), 1)
        self.vec.setBlockSize(1)

    def testGetSetVals(self):
        start, end = self.vec.getOwnershipRange()
        self.vec[start] = -7
        self.vec[end-1]   = -7
        self.vec.assemble()
        self.assertEqual(self.vec[start], -7)
        self.assertEqual(self.vec[end-1], -7)
        for i in range(start, end): self.vec[i] = i
        self.vec.assemble()
        values = [self.vec[i] for i in range(start, end)]
        self.assertEqual(values, list(range(start, end)))
        sz = self.vec.getSize()
        self.assertEqual(self.vec.sum(), (sz-1)/2.0*sz)

    def testGetSetValsBlocked(self):
        return
        lsize, gsize = self.vec.getSizes()
        start, end = self.vec.getOwnershipRange()
        bsizes  = list(range(1, lsize+1))
        nblocks = list(range(1, lsize+1))
        compat = [(bs, nb)
                  for bs in bsizes  if not (gsize%bs or lsize % bs)
                  for nb in nblocks if bs*nb <= lsize]
        for bsize, nblock in compat:
            self.vec.setBlockSize(bsize)
            bindex = [start//bsize+i  for i in range(nblock)]
            bvalue = [float(i) for i in range(nblock*bsize)]
            self.vec.setValuesBlocked(bindex, bvalue)
            self.vec.assemble()
            index = [start+i for i in range(nblock*bsize)]
            value = self.vec.getValues(index)
            self.assertEqual(bvalue, list(value))

    def testGetSetArray(self):
        self.vec.set(1)
        arr0 = self.vec.getArray().copy()
        self.assertEqual(arr0.sum(), self.vec.getLocalSize())
        arr0 = self.vec.getArray().copy()
        self.vec.setRandom()
        arr1 = self.vec.getArray().copy()
        self.vec.setArray(arr1)
        arr1 = self.vec.getArray().copy()
        arr2 = self.vec.getArray().copy()
        self.assertTrue((arr1 == arr2).all())
        import numpy
        refs = self.vec.getRefCount()
        arr3 = numpy.asarray(self.vec)
        self.assertEqual(self.vec.getRefCount(), refs+1)
        self.assertTrue((arr1 == arr3).all())
        arr3[:] = 0
        self.assertAlmostEqual(abs(self.vec.sum()), 0)
        self.assertEqual(self.vec.max()[1], 0)
        self.assertEqual(self.vec.min()[1], 0)
        self.vec.set(1)
        self.assertAlmostEqual(abs(arr3.sum()), self.vec.getLocalSize())
        self.assertEqual(arr3.min(), 1)
        self.assertEqual(arr3.max(), 1)
        del arr3
        self.assertEqual(self.vec.getRefCount(), refs)

    def testPlaceArray(self):
        self.vec.set(1)
        array = self.vec.getArray().copy()
        self.vec.placeArray(array)
        array[:] = 2
        self.assertAlmostEqual(abs(self.vec.sum()), 2*self.vec.getSize())
        self.vec.resetArray()
        self.assertAlmostEqual(abs(self.vec.sum()), self.vec.getSize())

    def testLocalVector(self):
        rank = self.vec.getComm().Get_rank()
        self.vec.getArray()[:] = rank + 1
        ln = self.vec.getLocalSize()
        lvec = self.vec.createLocalVector()
        self.vec.getLocalVector(lvec)
        self.assertEqual(abs(lvec.sum()), (rank+1)*ln)
        self.vec.restoreLocalVector(lvec)
        self.vec.getLocalVector(lvec,readonly=True)
        self.assertEqual(abs(lvec.sum()), (rank+1)*ln)
        self.vec.restoreLocalVector(lvec,readonly=True)
        lvec.destroy()

    def testSetOption(self):
        opt1 = PETSc.Vec.Option.IGNORE_OFF_PROC_ENTRIES
        opt2 = PETSc.Vec.Option.IGNORE_NEGATIVE_INDICES
        for opt in [opt1, opt2]*2:
            for flag in [True,False]*2:
                self.vec.setOption(opt,flag)

    def testGetSetItem(self):
        v = self.vec
        w = v.duplicate()
        #
        v[...] = 7
        self.assertEqual(v.max()[1], 7)
        self.assertEqual(v.min()[1], 7)
        #
        v.setRandom()
        w[...] = v
        self.assertTrue(w.equal(v))
        #
        v.setRandom()
        w[...] = v.getArray()
        self.assertTrue(w.equal(v))
        #
        s, e = v.getOwnershipRange()
        v.setRandom()
        w[s:e] = v.getArray().copy()
        w.assemble()
        self.assertTrue(w.equal(v))
        w1, v1 = w[s],   v[s]
        w2, v2 = w[e-1], v[e-1]
        self.assertEqual(w1, v1)
        self.assertEqual(w2, v2)

    def testMAXPY(self):
        y = self.vec
        y.set(1)
        x = [y.copy() for _ in range(3)]
        a = [1]*len(x)
        y.maxpy(a, x)
        z = y.duplicate()
        z.set(len(x)+1)
        assert (y.equal(z))

    def testBinOp(self):
        x = self.vec
        x.set(1)
        n = x.getSize()
        y = 2 + 2*x + 1 - x*3 - 1
        self.assertEqual(y.min()[1], 1)
        self.assertEqual(y.max()[1], 1)
        z = (4*x)/(2*y)
        self.assertEqual(z.min()[1], 2)
        self.assertEqual(z.max()[1], 2)
        z = z/2
        self.assertEqual(z.min()[1], 1)
        self.assertEqual(z.max()[1], 1)
        s = (+x) @ (-y)
        self.assertEqual(s, -n)
        #
        M, N = n, 2*n
        A = PETSc.Mat().createDense((M, N), comm=self.COMM)
        A.setUp()
        rs, re = A.getOwnershipRange()
        cs, ce = A.getOwnershipRangeColumn()
        a, b = 3, 5
        for i in range(rs, re):
            for j in range(N):
                A[i, j] = a*i + b*j
        A.assemble()
        y = x @ A
        self.assertEqual(y.getSize(), N)
        for i in range(cs, ce):
            self.assertEqual(y[i], a*M*(M-1)/2 + b*i*M)
        y.set(1)
        z = A @ y
        self.assertEqual(z.getSize(), M)
        for i in range(rs, re):
            self.assertEqual(z[i], b*N*(N-1)/2 + a*i*N)

# --------------------------------------------------------------------

class TestVecSeq(BaseTestVec, unittest.TestCase):
    COMM = PETSc.COMM_SELF
    TYPE = PETSc.Vec.Type.SEQ

class TestVecMPI(BaseTestVec, unittest.TestCase):
    COMM  = PETSc.COMM_WORLD
    TYPE = PETSc.Vec.Type.MPI

class TestVecShared(BaseTestVec, unittest.TestCase):
    if PETSc.COMM_WORLD.getSize() == 1:
        TYPE = PETSc.Vec.Type.SHARED
    else:
        TYPE = PETSc.Vec.Type.MPI
    COMM  = PETSc.COMM_WORLD

#class TestVecSieve(BaseTestVec, unittest.TestCase):
#    CLASS = PETSc.VecSieve
#    TARGS = ([],)

#class TestVecGhost(BaseTestVec, unittest.TestCase):
#    CLASS = PETSc.VecGhost
#    TARGS = ([],)

# --------------------------------------------------------------------

class TestVecWithArray(unittest.TestCase):

    def testCreateSeq(self):
        import numpy
        a = numpy.zeros(5, dtype=PETSc.ScalarType)

        v1 = PETSc.Vec().createWithArray(a, comm=PETSc.COMM_SELF)
        v2 = PETSc.Vec().createWithArray(a, size=5, comm=PETSc.COMM_SELF)
        v3 = PETSc.Vec().createWithArray(a, size=3, comm=PETSc.COMM_SELF)

        self.assertTrue(v1.size == 5)
        self.assertTrue(v2.size == 5)
        self.assertTrue(v3.size == 3)

        a1 = v1.getDict()['__array__']; self.assertTrue(a is a1)
        a2 = v2.getDict()['__array__']; self.assertTrue(a is a2)
        a3 = v3.getDict()['__array__']; self.assertTrue(a is a2)

    def testCreateMPI(self):
        import numpy
        a = numpy.zeros(5, dtype=PETSc.ScalarType)

        v1 = PETSc.Vec().createWithArray(a, comm=PETSc.COMM_WORLD)
        v2 = PETSc.Vec().createWithArray(a, size=(5,None), comm=PETSc.COMM_WORLD)
        v3 = PETSc.Vec().createWithArray(a, size=(3,None), comm=PETSc.COMM_WORLD)

        self.assertTrue(v1.local_size == 5)
        self.assertTrue(v2.local_size == 5)
        self.assertTrue(v3.local_size == 3)

        a1 = v1.getDict()['__array__']; self.assertTrue(a is a1)
        a2 = v2.getDict()['__array__']; self.assertTrue(a is a2)
        a3 = v3.getDict()['__array__']; self.assertTrue(a is a2)

    def testSetMPIGhost(self):
        import numpy
        v = PETSc.Vec().create()
        v.setType(PETSc.Vec.Type.MPI)
        v.setSizes((5,None))
        ghosts = [i % v.size for i in range(v.owner_range[1],v.owner_range[1]+3)]
        v.setMPIGhost(ghosts)
        v.setArray(numpy.array(range(*v.owner_range),dtype=PETSc.ScalarType))
        v.ghostUpdate()
        with v.localForm() as loc:
            self.assertTrue((loc[0:v.local_size] == range(*v.owner_range)).all())
            self.assertTrue((loc[v.local_size:] == ghosts).all())

# --------------------------------------------------------------------

if __name__ == '__main__':
    unittest.main()

# --------------------------------------------------------------------