1import unittest 2from petsc4py import PETSc 3import numpy as N 4from sys import getrefcount 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 17 def setUp(self): 18 u1 = PETSc.Vec().createSeq(3) 19 u2 = PETSc.Vec().createSeq(3) 20 u1[0], u1[1], u1[2] = [1, 2, 0]; u1.normalize() 21 u2[0], u2[1], u2[2] = [2, -1, 0]; 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 31 def _remove(self): 32 v = PETSc.Vec().createSeq(3); 33 v[0], v[1], v[2] = [7, 8, 9] 34 w = v.copy() 35 self.nullsp.remove(w) 36 return (v, w) 37 38 def testRemove(self): 39 v, w = self._remove() 40 tols = (0, 1e-5) 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 self.nullsp.setFunction(myremove) 55 v, w = self._remove() 56 self.assertTrue(allclose(v.array, [7, 8, 9])) 57 self.assertTrue(allclose(w.array, [1, 2, 3])) 58 self.nullsp.remove(v) 59 self.assertTrue(allclose(v.array, [1, 2, 3])) 60 self.nullsp.setFunction(None) 61 self.testRemove() 62 63 def testGetSetFunction(self): 64 def rem(nsp, vec): 65 vec.set(0) 66 self.nullsp.setFunction(rem) 67 self.assertEqual(getrefcount(rem)-1, 2) 68 dct = self.nullsp.getDict() 69 self.assertTrue(dct is not None) 70 self.assertEqual(getrefcount(dct)-1, 2) 71 fun, a, kw = dct['__function__'] 72 self.assertTrue(fun is rem) 73 self.nullsp.setFunction(None) 74 fun = dct.get('__function__') 75 self.assertEqual(getrefcount(rem)-1, 1) 76 self.assertTrue(fun is None) 77 78# -------------------------------------------------------------------- 79 80if __name__ == '__main__': 81 unittest.main() 82 83# -------------------------------------------------------------------- 84