xref: /petsc/src/binding/petsc4py/test/test_ts_py.py (revision ccfb0f9f40a0131988d7995ed9679700dae2a75a) 
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        self.assertEqual(getrefcount(self._getCtx()), 2)
78        self.assertEqual(self._getCtx().log['create'], 1)
79        self.nsolve = 0
80
81    def tearDown(self):
82        self.assertEqual(getrefcount(self._getCtx()), 2)
83        self.assertTrue('destroy' not in self._getCtx().log)
84        ctx = self._getCtx()
85        self.ts.destroy()
86        self.ts = None
87        PETSc.garbage_cleanup()
88        self.assertEqual(ctx.log['destroy'], 1)
89
90    def testGetType(self):
91        ctx = self.ts.getPythonContext()
92        pytype = f'{ctx.__module__}.{type(ctx).__name__}'
93        self.assertTrue(self.ts.getPythonType() == pytype)
94
95    def testSolve(self):
96        ts = self.ts
97        ts.setProblemType(ts.ProblemType.NONLINEAR)
98        ode = MyODE()
99        J = PETSc.Mat().create(ts.comm)
100        J.setSizes(3)
101        J.setFromOptions()
102        J.setUp()
103        u, f = J.createVecs()
104
105        ts.setAppCtx(ode)
106        ts.setIFunction(ode.function, f)
107        ts.setIJacobian(ode.jacobian, J, J)
108        ts.snes.ksp.pc.setType('none')
109
110        T0, dT, nT = 0.0, 0.1, 10
111        T = T0 + nT * dT
112        ts.setTime(T0)
113        ts.setTimeStep(dT)
114        ts.setMaxTime(T)
115        ts.setMaxSteps(nT)
116        ts.setFromOptions()
117        u[0], u[1], u[2] = 1, 2, 3
118        ts.solve(u)
119        self.nsolve += 1
120
121        self.assertTrue(ode.function_calls > 0)
122        self.assertTrue(ode.jacobian_calls > 0)
123
124        ctx = self.ts.getPythonContext()
125        ncalls = self.nsolve * ts.step_number
126        self.assertTrue(ctx.log['solveStep'] == ncalls)
127        self.assertTrue(ctx.log['adaptStep'] == ncalls)
128        del ctx
129
130        dct = self.ts.getDict()
131        self.assertTrue('__appctx__' in dct)
132        self.assertTrue('__ifunction__' in dct)
133        self.assertTrue('__ijacobian__' in dct)
134
135    def testFDColor(self):
136        #
137        ts = self.ts
138        ts.setProblemType(ts.ProblemType.NONLINEAR)
139        ode = MyODE()
140        J = PETSc.Mat().create(ts.comm)
141        J.setSizes(5)
142        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    def _getCtx(self):
185        return self.ts.getPythonContext()
186
187
188# --------------------------------------------------------------------
189
190if __name__ == '__main__':
191    unittest.main()
192
193# --------------------------------------------------------------------
194