xref: /petsc/src/binding/petsc4py/test/test_nsp.py (revision b69d2765e9c8cde17308f61ee677dcd992b8a9cf) 
1import unittest
2from petsc4py import PETSc
3
4# --------------------------------------------------------------------
5
6
7def allclose(seq1, seq2):
8    for v1, v2 in zip(seq1, seq2):
9        if abs(v1 - v2) > 1e-5:
10            return False
11    return True
12
13
14class TestNullSpace(unittest.TestCase):
15    def setUp(self):
16        u1 = PETSc.Vec().createSeq(3)
17        u2 = PETSc.Vec().createSeq(3)
18        u1[0], u1[1], u1[2] = [1, 2, 0]
19        u1.normalize()
20        u2[0], u2[1], u2[2] = [2, -1, 0]
21        u2.normalize()
22        basis = [u1, u2]
23        nullsp = PETSc.NullSpace().create(False, basis, comm=PETSc.COMM_SELF)
24        self.basis = basis
25        self.nullsp = nullsp
26
27    def tearDown(self):
28        self.basis = None
29        self.nullsp = None
30        PETSc.garbage_cleanup()
31
32    def _remove(self):
33        v = PETSc.Vec().createSeq(3)
34        v[0], v[1], v[2] = [7, 8, 9]
35        w = v.copy()
36        self.nullsp.remove(w)
37        return (v, w)
38
39    def testRemove(self):
40        v, w = self._remove()
41        self.assertTrue(allclose(v.array, [7, 8, 9]))
42        self.assertTrue(allclose(w.array, [0, 0, 9]))
43        del v, w
44
45    def testRemoveInplace(self):
46        v, w = self._remove()
47        self.nullsp.remove(v)
48        self.assertTrue(v.equal(w))
49        del v, w
50
51    def testRemoveWithFunction(self):
52        def myremove(nsp, vec):
53            vec.setArray([1, 2, 3])
54
55        self.nullsp.setFunction(myremove)
56        v, w = self._remove()
57        self.assertTrue(allclose(v.array, [7, 8, 9]))
58        self.assertTrue(allclose(w.array, [1, 2, 3]))
59        self.nullsp.remove(v)
60        self.assertTrue(allclose(v.array, [1, 2, 3]))
61        self.nullsp.setFunction(None)
62        self.testRemove()
63
64    def testGetSetFunction(self):
65        def rem(nsp, vec):
66            vec.set(0)
67
68        self.nullsp.setFunction(rem)
69        dct = self.nullsp.getDict()
70        self.assertTrue(dct is not None)
71        fun, a, kw = dct['__function__']
72        self.assertTrue(fun is rem)
73        self.nullsp.setFunction(None)
74        fun = dct.get('__function__')
75        self.assertTrue(fun is None)
76
77
78# --------------------------------------------------------------------
79
80if __name__ == '__main__':
81    unittest.main()
82
83# --------------------------------------------------------------------
84