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