xref: /petsc/src/binding/petsc4py/test/test_ts_py.py (revision d5b43468fb8780a8feea140ccd6fa3e6a50411cc) 
1import unittest
2from petsc4py import PETSc
3from sys import getrefcount
4
5# --------------------------------------------------------------------
6
7class MyODE:
8    """
9    du/dt + u**2 = 0;
10    u0 = 1
11    """
12
13    def __init__(self):
14        self.function_calls = 0
15        self.jacobian_calls = 0
16
17    def function(self,ts,t,u,du,F):
18        #print 'MyODE.function()'
19        self.function_calls += 1
20        f = du + u * u
21        f.copy(F)
22
23    def jacobian(self,ts,t,u,du,a,J,P):
24        #print 'MyODE.jacobian()'
25        self.jacobian_calls += 1
26        P.zeroEntries()
27        diag = a + 2 * u
28        P.setDiagonal(diag)
29        P.assemble()
30        if J != P: J.assemble()
31        return False # same_nz
32
33class MyTS(object):
34
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        PETSc.garbage_cleanup()
88        self.assertEqual(ctx.log['destroy'], 1)
89        self.assertEqual(getrefcount(ctx),   2)
90
91    def testGetType(self):
92        ctx = self.ts.getPythonContext()
93        pytype = "{0}.{1}".format(ctx.__module__, type(ctx).__name__)
94        self.assertTrue(self.ts.getPythonType() == pytype)
95
96    def testSolve(self):
97        ts = self.ts
98        ts.setProblemType(ts.ProblemType.NONLINEAR)
99        ode = MyODE()
100        J = PETSc.Mat().create(ts.comm)
101        J.setSizes(3);
102        J.setFromOptions()
103        J.setUp()
104        u, f = J.createVecs()
105
106        ts.setAppCtx(ode)
107        ts.setIFunction(ode.function, f)
108        ts.setIJacobian(ode.jacobian, J, J)
109        ts.snes.ksp.pc.setType('none')
110
111        T0, dT, nT = 0.0, 0.1, 10
112        T = T0 + nT*dT
113        ts.setTime(T0)
114        ts.setTimeStep(dT)
115        ts.setMaxTime(T)
116        ts.setMaxSteps(nT)
117        ts.setFromOptions()
118        u[0], u[1], u[2] = 1, 2, 3
119        ts.solve(u)
120        self.nsolve +=1
121
122        self.assertTrue(ode.function_calls > 0)
123        self.assertTrue(ode.jacobian_calls > 0)
124
125        ctx = self.ts.getPythonContext()
126        ncalls = self.nsolve * ts.step_number
127        self.assertTrue(ctx.log['solveStep'] == ncalls)
128        self.assertTrue(ctx.log['adaptStep'] == ncalls)
129        del ctx
130
131        dct = self.ts.getDict()
132        self.assertTrue('__appctx__'    in dct)
133        self.assertTrue('__ifunction__' in dct)
134        self.assertTrue('__ijacobian__' in dct)
135
136    def testFDColor(self):
137        #
138        ts = self.ts
139        ts.setProblemType(ts.ProblemType.NONLINEAR)
140        ode = MyODE()
141        J = PETSc.Mat().create(ts.comm)
142        J.setSizes(5); J.setType('aij');
143        J.setPreallocationNNZ(1)
144        J.setFromOptions()
145        u, f = J.createVecs()
146
147        ts.setAppCtx(ode)
148        ts.setIFunction(ode.function, f)
149        ts.setIJacobian(ode.jacobian, J, J)
150
151        T0, dT, nT = 0.00, 0.1, 10
152        T = T0 + nT*dT
153        ts.setTime(T0)
154        ts.setTimeStep(dT)
155        ts.setMaxTime(T)
156        ts.setMaxSteps(nT)
157        ts.setFromOptions()
158        u[:] = 1, 2, 3, 4, 5
159
160        ts.setSolution(u)
161        ode.jacobian(ts,0.0,u,u,1.0,J,J)
162        ts.snes.setUseFD(True)
163        ts.solve(u)
164        self.nsolve +=1
165
166    def testResetAndSolve(self):
167        self.ts.reset()
168        self.ts.setStepNumber(0)
169        self.testSolve()
170        self.ts.reset()
171        self.ts.setStepNumber(0)
172        self.testFDColor()
173        self.ts.reset()
174        self.ts.setStepNumber(0)
175        self.testSolve()
176        self.ts.reset()
177
178    def testSetAdaptLimits(self):
179        self.ts.setStepLimits(1.0, 2.0)
180        hmin, hmax = self.ts.getStepLimits()
181        self.assertEqual(1.0, hmin)
182        self.assertEqual(2.0, hmax)
183
184# --------------------------------------------------------------------
185
186if __name__ == '__main__':
187    unittest.main()
188
189# --------------------------------------------------------------------
190