1import unittest 2from petsc4py import PETSc 3from sys import getrefcount 4import numpy 5 6 7# -------------------------------------------------------------------- 8class Objective: 9 def __call__(self, tao, x): 10 return (x[0] - 1.0) ** 2 + (x[1] - 2.0) ** 2 11 12 13class Gradient: 14 def __call__(self, tao, x, g): 15 g[0] = 2.0 * (x[0] - 1.0) 16 g[1] = 2.0 * (x[1] - 2.0) 17 g.assemble() 18 19 20class MyTao: 21 def __init__(self): 22 self.log = {} 23 24 def _log(self, method): 25 self.log.setdefault(method, 0) 26 self.log[method] += 1 27 28 def create(self, tao): 29 self._log('create') 30 self.testvec = PETSc.Vec() 31 32 def destroy(self, tao): 33 self._log('destroy') 34 self.testvec.destroy() 35 36 def setFromOptions(self, tao): 37 self._log('setFromOptions') 38 39 def setUp(self, tao): 40 self._log('setUp') 41 self.testvec = tao.getSolution().duplicate() 42 43 def solve(self, tao): 44 self._log('solve') 45 46 def step(self, tao, x, g, s): 47 self._log('step') 48 tao.computeGradient(x, g) 49 g.copy(s) 50 s.scale(-1.0) 51 52 def preStep(self, tao): 53 self._log('preStep') 54 55 def postStep(self, tao): 56 self._log('postStep') 57 58 def monitor(self, tao): 59 self._log('monitor') 60 61 62class TestTaoPython(unittest.TestCase): 63 def setUp(self): 64 self.tao = PETSc.TAO() 65 self.tao.createPython(MyTao(), comm=PETSc.COMM_SELF) 66 ctx = self.tao.getPythonContext() 67 self.assertEqual(getrefcount(ctx), 3) 68 self.assertEqual(ctx.log['create'], 1) 69 self.nsolve = 0 70 71 def tearDown(self): 72 ctx = self.tao.getPythonContext() 73 self.assertEqual(getrefcount(ctx), 3) 74 self.assertTrue('destroy' not in ctx.log) 75 self.tao.destroy() 76 self.tao = None 77 PETSc.garbage_cleanup() 78 self.assertEqual(ctx.log['destroy'], 1) 79 self.assertEqual(getrefcount(ctx), 2) 80 81 def testGetType(self): 82 ctx = self.tao.getPythonContext() 83 pytype = f'{ctx.__module__}.{type(ctx).__name__}' 84 self.assertTrue(self.tao.getPythonType() == pytype) 85 86 def testSolve(self): 87 tao = self.tao 88 ctx = tao.getPythonContext() 89 x = PETSc.Vec().create(tao.getComm()) 90 x.setType('standard') 91 x.setSizes(2) 92 y1 = x.duplicate() 93 y2 = x.duplicate() 94 tao.setObjective(Objective()) 95 tao.setGradient(Gradient(), None) 96 tao.setMonitor(ctx.monitor) 97 tao.setFromOptions() 98 tao.setMaximumIterations(3) 99 tao.setSolution(x) 100 101 # Call the solve method of MyTAO 102 x.set(0.5) 103 tao.solve() 104 n = tao.getIterationNumber() 105 self.assertTrue(n == 0) 106 107 # Call the default solve method and use step of MyTAO 108 ctx.solve = None 109 x.set(0.5) 110 tao.solve() 111 n = tao.getIterationNumber() 112 self.assertTrue(n == 3) 113 x.copy(y1) 114 115 # Call the default solve method with the default step method 116 ctx.step = None 117 x.set(0.5) 118 tao.solve() 119 n = tao.getIterationNumber() 120 self.assertTrue(n == 3) 121 x.copy(y2) 122 123 self.assertTrue(y1.equal(y2)) 124 self.assertTrue(ctx.log['monitor'] == 2 * (n + 1)) 125 self.assertTrue(ctx.log['preStep'] == 2 * n) 126 self.assertTrue(ctx.log['postStep'] == 2 * n) 127 self.assertTrue(ctx.log['solve'] == 1) 128 self.assertTrue(ctx.log['setUp'] == 1) 129 self.assertTrue(ctx.log['setFromOptions'] == 1) 130 self.assertTrue(ctx.log['step'] == n) 131 tao.cancelMonitor() 132 133 134# -------------------------------------------------------------------- 135 136if numpy.iscomplexobj(PETSc.ScalarType()): 137 del TestTaoPython 138 139if __name__ == '__main__': 140 unittest.main() 141 142# -------------------------------------------------------------------- 143