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