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