xref: /petsc/src/binding/petsc4py/test/test_ts_py.py (revision b122ec5aa1bd4469eb4e0673542fb7de3f411254) 
1import unittest
2from petsc4py import PETSc
3from sys import getrefcount
4import gc
5
6# --------------------------------------------------------------------
7
8class MyODE:
9    """
10    du/dt + u**2 = 0;
11    u0 = 1
12    """
13
14    def __init__(self):
15        self.function_calls = 0
16        self.jacobian_calls = 0
17
18    def function(self,ts,t,u,du,F):
19        #print 'MyODE.function()'
20        self.function_calls += 1
21        f = du + u * u
22        f.copy(F)
23
24    def jacobian(self,ts,t,u,du,a,J,P):
25        #print 'MyODE.jacobian()'
26        self.jacobian_calls += 1
27        P.zeroEntries()
28        diag = a + 2 * u
29        P.setDiagonal(diag)
30        P.assemble()
31        if J != P: J.assemble()
32        return False # same_nz
33
34class MyTS:
35    def __init__(self):
36        self.log = {}
37
38    def _log(self, method, *args):
39        self.log.setdefault(method, 0)
40        self.log[method] += 1
41
42    def create(self, ts, *args):
43        self._log('create', *args)
44        self.vec_update = PETSc.Vec()
45
46    def destroy(self, ts, *args):
47        self._log('destroy', *args)
48        self.vec_update.destroy()
49
50    def setFromOptions(self, ts, *args):
51        self._log('setFromOptions', *args)
52
53    def setUp(self, ts, *args):
54        self._log('setUp', ts, *args)
55        self.vec_update = ts.getSolution().duplicate()
56
57    def reset(self, ts, *args):
58        self._log('reset', ts, *args)
59
60    def solveStep(self, ts, t, u, *args):
61        self._log('solveStep', ts, t, u, *args)
62        ts.snes.solve(None, u)
63
64    def adaptStep(self, ts, t, u, *args):
65        self._log('adaptStep', ts, t, u, *args)
66        return (ts.getTimeStep(), True)
67
68
69class TestTSPython(unittest.TestCase):
70
71    def setUp(self):
72        self.ts = PETSc.TS()
73        self.ts.createPython(MyTS(), comm=PETSc.COMM_SELF)
74        eft = PETSc.TS.ExactFinalTime.STEPOVER
75        self.ts.setExactFinalTime(eft)
76        ctx = self.ts.getPythonContext()
77        self.assertEqual(getrefcount(ctx),  3)
78        self.assertEqual(ctx.log['create'], 1)
79        self.nsolve = 0
80
81    def tearDown(self):
82        ctx = self.ts.getPythonContext()
83        self.assertEqual(getrefcount(ctx), 3)
84        self.assertTrue('destroy' not in ctx.log)
85        self.ts.destroy() # XXX
86        self.ts = None
87        self.assertEqual(ctx.log['destroy'], 1)
88        self.assertEqual(getrefcount(ctx),   2)
89
90    def testSolve(self):
91        ts = self.ts
92        ts.setProblemType(ts.ProblemType.NONLINEAR)
93        ode = MyODE()
94        J = PETSc.Mat().create(ts.comm)
95        J.setSizes(3);
96        J.setFromOptions()
97        J.setUp()
98        u, f = J.createVecs()
99
100        ts.setAppCtx(ode)
101        ts.setIFunction(ode.function, f)
102        ts.setIJacobian(ode.jacobian, J, J)
103        ts.snes.ksp.pc.setType('none')
104
105        T0, dT, nT = 0.0, 0.1, 10
106        T = T0 + nT*dT
107        ts.setTime(T0)
108        ts.setTimeStep(dT)
109        ts.setMaxTime(T)
110        ts.setMaxSteps(nT)
111        ts.setFromOptions()
112        u[0], u[1], u[2] = 1, 2, 3
113        ts.solve(u)
114        self.nsolve +=1
115
116        self.assertTrue(ode.function_calls > 0)
117        self.assertTrue(ode.jacobian_calls > 0)
118
119        ctx = self.ts.getPythonContext()
120        ncalls = self.nsolve * ts.step_number
121        self.assertTrue(ctx.log['solveStep'] == ncalls)
122        self.assertTrue(ctx.log['adaptStep'] == ncalls)
123        del ctx
124
125        dct = self.ts.getDict()
126        self.assertTrue('__appctx__'    in dct)
127        self.assertTrue('__ifunction__' in dct)
128        self.assertTrue('__ijacobian__' in dct)
129
130    def testFDColor(self):
131        #
132        ts = self.ts
133        ts.setProblemType(ts.ProblemType.NONLINEAR)
134        ode = MyODE()
135        J = PETSc.Mat().create(ts.comm)
136        J.setSizes(5); J.setType('aij');
137        J.setPreallocationNNZ(1)
138        J.setFromOptions()
139        u, f = J.createVecs()
140
141        ts.setAppCtx(ode)
142        ts.setIFunction(ode.function, f)
143        ts.setIJacobian(ode.jacobian, J, J)
144
145        T0, dT, nT = 0.00, 0.1, 10
146        T = T0 + nT*dT
147        ts.setTime(T0)
148        ts.setTimeStep(dT)
149        ts.setMaxTime(T)
150        ts.setMaxSteps(nT)
151        ts.setFromOptions()
152        u[:] = 1, 2, 3, 4, 5
153
154        ts.setSolution(u)
155        ode.jacobian(ts,0.0,u,u,1.0,J,J)
156        ts.snes.setUseFD(True)
157        ts.solve(u)
158        self.nsolve +=1
159
160    def testResetAndSolve(self):
161        self.ts.reset()
162        self.ts.setStepNumber(0)
163        self.testSolve()
164        self.ts.reset()
165        self.ts.setStepNumber(0)
166        self.testFDColor()
167        self.ts.reset()
168        self.ts.setStepNumber(0)
169        self.testSolve()
170        self.ts.reset()
171
172    def testSetAdaptLimits(self):
173        self.ts.setStepLimits(1.0, 2.0)
174        hmin, hmax = self.ts.getStepLimits()
175        self.assertEqual(1.0, hmin)
176        self.assertEqual(2.0, hmax)
177
178# --------------------------------------------------------------------
179
180if __name__ == '__main__':
181    unittest.main()
182
183# --------------------------------------------------------------------
184