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