xref: /petsc/src/binding/petsc4py/test/test_ts_py.py (revision a69119a591a03a9d906b29c0a4e9802e4d7c9795) 
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(object):
35
36    def __init__(self):
37        self.log = {}
38
39    def _log(self, method, *args):
40        self.log.setdefault(method, 0)
41        self.log[method] += 1
42
43    def create(self, ts, *args):
44        self._log('create', *args)
45        self.vec_update = PETSc.Vec()
46
47    def destroy(self, ts, *args):
48        self._log('destroy', *args)
49        self.vec_update.destroy()
50
51    def setFromOptions(self, ts, *args):
52        self._log('setFromOptions', *args)
53
54    def setUp(self, ts, *args):
55        self._log('setUp', ts, *args)
56        self.vec_update = ts.getSolution().duplicate()
57
58    def reset(self, ts, *args):
59        self._log('reset', ts, *args)
60
61    def solveStep(self, ts, t, u, *args):
62        self._log('solveStep', ts, t, u, *args)
63        ts.snes.solve(None, u)
64
65    def adaptStep(self, ts, t, u, *args):
66        self._log('adaptStep', ts, t, u, *args)
67        return (ts.getTimeStep(), True)
68
69
70class TestTSPython(unittest.TestCase):
71
72    def setUp(self):
73        self.ts = PETSc.TS()
74        self.ts.createPython(MyTS(), comm=PETSc.COMM_SELF)
75        eft = PETSc.TS.ExactFinalTime.STEPOVER
76        self.ts.setExactFinalTime(eft)
77        ctx = self.ts.getPythonContext()
78        self.assertEqual(getrefcount(ctx),  3)
79        self.assertEqual(ctx.log['create'], 1)
80        self.nsolve = 0
81
82    def tearDown(self):
83        ctx = self.ts.getPythonContext()
84        self.assertEqual(getrefcount(ctx), 3)
85        self.assertTrue('destroy' not in ctx.log)
86        self.ts.destroy() # XXX
87        self.ts = None
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