xref: /petsc/src/binding/petsc4py/test/test_vec.py (revision 98d129c30f3ee9fdddc40fdbc5a989b7be64f888) 
1from petsc4py import PETSc
2import unittest
3
4# --------------------------------------------------------------------
5
6
7class BaseTestVec:
8    COMM = None
9    TYPE = None
10
11    def setUp(self):
12        v = PETSc.Vec()
13        v.create(self.COMM)
14        v.setSizes(100)
15        v.setType(self.TYPE)
16        self.vec = v
17
18    def tearDown(self):
19        self.vec.destroy()
20        self.vec = None
21        PETSc.garbage_cleanup()
22
23    def testDuplicate(self):
24        self.vec.set(1)
25        vec = self.vec.duplicate()
26        self.assertFalse(self.vec.equal(vec))
27        self.assertEqual(self.vec.sizes, vec.sizes)
28        del vec
29
30    def testCopy(self):
31        self.vec.set(1)
32        vec = self.vec.duplicate()
33        self.vec.copy(vec)
34        self.assertTrue(self.vec.equal(vec))
35        del vec
36
37    def testDot(self):
38        self.vec.set(1)
39        d = self.vec.dot(self.vec)
40        self.assertAlmostEqual(abs(d), self.vec.getSize())
41        self.vec.dotBegin(self.vec)
42        d = self.vec.dotEnd(self.vec)
43        self.assertAlmostEqual(abs(d), self.vec.getSize())
44
45    def testNorm(self):
46        from math import sqrt
47
48        self.vec.set(1)
49        n1 = self.vec.norm(PETSc.NormType.NORM_1)
50        n2 = self.vec.norm(PETSc.NormType.NORM_2)
51        ni = self.vec.norm(PETSc.NormType.NORM_INFINITY)
52        self.assertAlmostEqual(n1, self.vec.getSize())
53        self.assertAlmostEqual(n2, sqrt(self.vec.getSize()))
54        self.assertAlmostEqual(n2, self.vec.norm())
55        self.assertAlmostEqual(ni, 1.0)
56        self.vec.normBegin(PETSc.NormType.NORM_1)
57        nn1 = self.vec.normEnd(PETSc.NormType.NORM_1)
58        self.assertAlmostEqual(nn1, n1)
59        self.vec.normBegin()
60        nn2 = self.vec.normEnd()
61        self.assertAlmostEqual(nn2, n2)
62        self.vec.normBegin(PETSc.NormType.NORM_INFINITY)
63        nni = self.vec.normEnd(PETSc.NormType.NORM_INFINITY)
64        self.assertAlmostEqual(nni, ni)
65
66    def testNormalize(self):
67        from math import sqrt
68
69        self.vec.set(1)
70        n2 = self.vec.normalize()
71        self.assertAlmostEqual(n2, sqrt(self.vec.getSize()))
72        self.assertAlmostEqual(1, self.vec.norm())
73
74    def testSumMinMax(self):
75        self.vec.set(1)
76        self.assertEqual(self.vec.sum(), self.vec.getSize())
77        self.vec.set(-7)
78        self.assertEqual(self.vec.min()[1], -7)
79        self.vec.set(10)
80        self.assertEqual(self.vec.max()[1], 10)
81
82    def testSwap(self):
83        v1 = self.vec
84        v2 = v1.duplicate()
85        v1.set(1)
86        v2.set(2)
87        v1.swap(v2)
88        idx, _ = self.vec.getOwnershipRange()
89        self.assertEqual(v1[idx], 2)
90        self.assertEqual(v2[idx], 1)
91
92    def testBsize(self):
93        self.vec.setBlockSize(1)
94        self.assertEqual(self.vec.getBlockSize(), 1)
95        self.vec.setBlockSize(1)
96
97    def testGetSetVals(self):
98        start, end = self.vec.getOwnershipRange()
99        self.vec[start] = -7
100        self.vec[end - 1] = -7
101        self.vec.assemble()
102        self.assertEqual(self.vec[start], -7)
103        self.assertEqual(self.vec[end - 1], -7)
104        for i in range(start, end):
105            self.vec[i] = i
106        self.vec.assemble()
107        values = [self.vec[i] for i in range(start, end)]
108        self.assertEqual(values, list(range(start, end)))
109        sz = self.vec.getSize()
110        self.assertEqual(self.vec.sum(), (sz - 1) / 2.0 * sz)
111
112    def testGetSetValsBlocked(self):
113        return
114        lsize, gsize = self.vec.getSizes()
115        start, end = self.vec.getOwnershipRange()
116        bsizes = list(range(1, lsize + 1))
117        nblocks = list(range(1, lsize + 1))
118        compat = [
119            (bs, nb)
120            for bs in bsizes
121            if not (gsize % bs or lsize % bs)
122            for nb in nblocks
123            if bs * nb <= lsize
124        ]
125        for bsize, nblock in compat:
126            self.vec.setBlockSize(bsize)
127            bindex = [start // bsize + i for i in range(nblock)]
128            bvalue = [float(i) for i in range(nblock * bsize)]
129            self.vec.setValuesBlocked(bindex, bvalue)
130            self.vec.assemble()
131            index = [start + i for i in range(nblock * bsize)]
132            value = self.vec.getValues(index)
133            self.assertEqual(bvalue, list(value))
134
135    def testGetSetArray(self):
136        self.vec.set(1)
137        arr0 = self.vec.getArray().copy()
138        self.assertEqual(arr0.sum(), self.vec.getLocalSize())
139        arr0 = self.vec.getArray().copy()
140        self.vec.setRandom()
141        arr1 = self.vec.getArray().copy()
142        self.vec.setArray(arr1)
143        arr1 = self.vec.getArray().copy()
144        arr2 = self.vec.getArray().copy()
145        self.assertTrue((arr1 == arr2).all())
146        import numpy
147
148        refs = self.vec.getRefCount()
149        arr3 = numpy.asarray(self.vec)
150        self.assertEqual(self.vec.getRefCount(), refs + 1)
151        self.assertTrue((arr1 == arr3).all())
152        arr3[:] = 0
153        self.assertAlmostEqual(abs(self.vec.sum()), 0)
154        self.assertEqual(self.vec.max()[1], 0)
155        self.assertEqual(self.vec.min()[1], 0)
156        self.vec.set(1)
157        self.assertAlmostEqual(abs(arr3.sum()), self.vec.getLocalSize())
158        self.assertEqual(arr3.min(), 1)
159        self.assertEqual(arr3.max(), 1)
160        del arr3
161        self.assertEqual(self.vec.getRefCount(), refs)
162
163    def testPlaceArray(self):
164        self.vec.set(1)
165        array = self.vec.getArray().copy()
166        self.vec.placeArray(array)
167        array[:] = 2
168        self.assertAlmostEqual(abs(self.vec.sum()), 2 * self.vec.getSize())
169        self.vec.resetArray()
170        self.assertAlmostEqual(abs(self.vec.sum()), self.vec.getSize())
171
172    def testLocalVector(self):
173        rank = self.vec.getComm().Get_rank()
174        self.vec.getArray()[:] = rank + 1
175        ln = self.vec.getLocalSize()
176        lvec = self.vec.createLocalVector()
177        self.vec.getLocalVector(lvec)
178        self.assertEqual(abs(lvec.sum()), (rank + 1) * ln)
179        self.vec.restoreLocalVector(lvec)
180        self.vec.getLocalVector(lvec, readonly=True)
181        self.assertEqual(abs(lvec.sum()), (rank + 1) * ln)
182        self.vec.restoreLocalVector(lvec, readonly=True)
183        lvec.destroy()
184
185    def testSetOption(self):
186        opt1 = PETSc.Vec.Option.IGNORE_OFF_PROC_ENTRIES
187        opt2 = PETSc.Vec.Option.IGNORE_NEGATIVE_INDICES
188        for opt in [opt1, opt2] * 2:
189            for flag in [True, False] * 2:
190                self.vec.setOption(opt, flag)
191
192    def testGetSetItem(self):
193        v = self.vec
194        w = v.duplicate()
195        #
196        v[...] = 7
197        self.assertEqual(v.max()[1], 7)
198        self.assertEqual(v.min()[1], 7)
199        #
200        v.setRandom()
201        w[...] = v
202        self.assertTrue(w.equal(v))
203        #
204        v.setRandom()
205        w[...] = v.getArray()
206        self.assertTrue(w.equal(v))
207        #
208        s, e = v.getOwnershipRange()
209        v.setRandom()
210        w[s:e] = v.getArray().copy()
211        w.assemble()
212        self.assertTrue(w.equal(v))
213        w1, v1 = w[s], v[s]
214        w2, v2 = w[e - 1], v[e - 1]
215        self.assertEqual(w1, v1)
216        self.assertEqual(w2, v2)
217
218    def testMAXPY(self):
219        y = self.vec
220        y.set(1)
221        x = [y.copy() for _ in range(3)]
222        a = [1] * len(x)
223        y.maxpy(a, x)
224        z = y.duplicate()
225        z.set(len(x) + 1)
226        self.assertTrue(y.equal(z))
227
228    def testBinOp(self):
229        x = self.vec
230        x.set(1)
231        n = x.getSize()
232        y = 2 + 2 * x + 1 - x * 3 - 1
233        self.assertEqual(y.min()[1], 1)
234        self.assertEqual(y.max()[1], 1)
235        z = (4 * x) / (2 * y)
236        self.assertEqual(z.min()[1], 2)
237        self.assertEqual(z.max()[1], 2)
238        z = z / 2
239        self.assertEqual(z.min()[1], 1)
240        self.assertEqual(z.max()[1], 1)
241        s = (+x) @ (-y)
242        self.assertEqual(s, -n)
243        #
244        M, N = n, 2 * n
245        A = PETSc.Mat().createDense((M, N), comm=self.COMM)
246        A.setUp()
247        rs, re = A.getOwnershipRange()
248        cs, ce = A.getOwnershipRangeColumn()
249        a, b = 3, 5
250        for i in range(rs, re):
251            for j in range(N):
252                A[i, j] = a * i + b * j
253        A.assemble()
254        y = x @ A
255        self.assertEqual(y.getSize(), N)
256        for i in range(cs, ce):
257            self.assertEqual(y[i], a * M * (M - 1) / 2 + b * i * M)
258        y.set(1)
259        z = A @ y
260        self.assertEqual(z.getSize(), M)
261        for i in range(rs, re):
262            self.assertEqual(z[i], b * N * (N - 1) / 2 + a * i * N)
263
264
265# --------------------------------------------------------------------
266
267
268class TestVecSeq(BaseTestVec, unittest.TestCase):
269    COMM = PETSc.COMM_SELF
270    TYPE = PETSc.Vec.Type.SEQ
271
272
273class TestVecMPI(BaseTestVec, unittest.TestCase):
274    COMM = PETSc.COMM_WORLD
275    TYPE = PETSc.Vec.Type.MPI
276
277
278class TestVecShared(BaseTestVec, unittest.TestCase):
279    if PETSc.COMM_WORLD.getSize() == 1:
280        TYPE = PETSc.Vec.Type.SHARED
281    else:
282        TYPE = PETSc.Vec.Type.MPI
283    COMM = PETSc.COMM_WORLD
284
285
286# class TestVecSieve(BaseTestVec, unittest.TestCase):
287#    CLASS = PETSc.VecSieve
288#    TARGS = ([],)
289
290# class TestVecGhost(BaseTestVec, unittest.TestCase):
291#    CLASS = PETSc.VecGhost
292#    TARGS = ([],)
293
294# --------------------------------------------------------------------
295
296
297class TestVecWithArray(unittest.TestCase):
298    def testCreateSeq(self):
299        import numpy
300
301        a = numpy.zeros(5, dtype=PETSc.ScalarType)
302
303        v1 = PETSc.Vec().createWithArray(a, comm=PETSc.COMM_SELF)
304        v2 = PETSc.Vec().createWithArray(a, size=5, comm=PETSc.COMM_SELF)
305        v3 = PETSc.Vec().createWithArray(a, size=3, comm=PETSc.COMM_SELF)
306
307        self.assertTrue(v1.size == 5)
308        self.assertTrue(v2.size == 5)
309        self.assertTrue(v3.size == 3)
310
311        a1 = v1.getDict()['__array__']
312        self.assertTrue(a is a1)
313        a2 = v2.getDict()['__array__']
314        self.assertTrue(a is a2)
315        a3 = v3.getDict()['__array__']
316        self.assertTrue(a is a3)
317
318    def testCreateMPI(self):
319        import numpy
320
321        a = numpy.zeros(5, dtype=PETSc.ScalarType)
322
323        v1 = PETSc.Vec().createWithArray(a, comm=PETSc.COMM_WORLD)
324        v2 = PETSc.Vec().createWithArray(a, size=(5, None), comm=PETSc.COMM_WORLD)
325        v3 = PETSc.Vec().createWithArray(a, size=(3, None), comm=PETSc.COMM_WORLD)
326
327        self.assertTrue(v1.local_size == 5)
328        self.assertTrue(v2.local_size == 5)
329        self.assertTrue(v3.local_size == 3)
330
331        a1 = v1.getDict()['__array__']
332        self.assertTrue(a is a1)
333        a2 = v2.getDict()['__array__']
334        self.assertTrue(a is a2)
335        a3 = v3.getDict()['__array__']
336        self.assertTrue(a is a3)
337
338    def testSetMPIGhost(self):
339        import numpy
340
341        v = PETSc.Vec().create()
342        v.setType(PETSc.Vec.Type.MPI)
343        v.setSizes((5, None))
344        ghosts = [i % v.size for i in range(v.owner_range[1], v.owner_range[1] + 3)]
345        v.setMPIGhost(ghosts)
346        v.setArray(numpy.array(range(*v.owner_range), dtype=PETSc.ScalarType))
347        v.ghostUpdate()
348        with v.localForm() as loc:
349            self.assertTrue((loc[0 : v.local_size] == range(*v.owner_range)).all())
350            self.assertTrue((loc[v.local_size :] == ghosts).all())
351
352
353# --------------------------------------------------------------------
354
355if __name__ == '__main__':
356    unittest.main()
357
358# --------------------------------------------------------------------
359