xref: /petsc/src/binding/petsc4py/test/test_ts.py (revision 552edb6364df478b294b3111f33a8f37ca096b20) 
1import unittest
2from petsc4py import PETSc
3
4# --------------------------------------------------------------------
5
6
7class MyODE:
8    """
9    du/dt + u**2 = 0;
10    u0,u1,u2 = 1,2,3
11    """
12
13    def __init__(self):
14        self.rhsfunction_calls = 0
15        self.rhsjacobian_calls = 0
16        self.ifunction_calls = 0
17        self.ijacobian_calls = 0
18        self.presolve_calls = 0
19        self.update_calls = 0
20        self.postsolve_calls = 0
21        self.monitor_calls = 0
22
23    def rhsfunction(self, ts, t, u, F):
24        # print ('MyODE.rhsfunction()')
25        self.rhsfunction_calls += 1
26        f = -(u * u)
27        f.copy(F)
28
29    def rhsjacobian(self, ts, t, u, J, P):
30        # print ('MyODE.rhsjacobian()')
31        self.rhsjacobian_calls += 1
32        P.zeroEntries()
33        diag = -2 * u
34        P.setDiagonal(diag)
35        P.assemble()
36        if J != P:
37            J.assemble()
38        return True  # same_nz
39
40    def ifunction(self, ts, t, u, du, F):
41        # print ('MyODE.ifunction()')
42        self.ifunction_calls += 1
43        f = du + u * u
44        f.copy(F)
45
46    def ijacobian(self, ts, t, u, du, a, J, P):
47        # print ('MyODE.ijacobian()')
48        self.ijacobian_calls += 1
49        P.zeroEntries()
50        diag = a + 2 * u
51        P.setDiagonal(diag)
52        P.assemble()
53        if J != P:
54            J.assemble()
55        return True  # same_nz
56
57    def monitor(self, ts, s, t, u):
58        self.monitor_calls += 1
59        # dt = ts.time_step
60        # ut = ts.vec_sol.norm()
61        # prn = PETSc.Sys.Print
62        # prn('TS: step %2d, T:%f, dT:%f, u:%f' % (s,t,dt,ut))
63
64
65class BaseTestTSNonlinear:
66    TYPE = None
67
68    def setUp(self):
69        self.ts = PETSc.TS().create(PETSc.COMM_SELF)
70        eft = PETSc.TS.ExactFinalTime.STEPOVER
71        self.ts.setExactFinalTime(eft)
72        ptype = PETSc.TS.ProblemType.NONLINEAR
73        self.ts.setProblemType(ptype)
74        self.ts.setType(self.TYPE)
75        if PETSc.ScalarType().dtype.char in 'fF':
76            snes = self.ts.getSNES()
77            snes.setTolerances(rtol=1e-6)
78
79    def tearDown(self):
80        self.ts = None
81        PETSc.garbage_cleanup()
82
83
84class BaseTestTSNonlinearRHS(BaseTestTSNonlinear):
85    def testSolveRHS(self, nullsol=False):
86        ts = self.ts
87        dct = self.ts.getDict()
88        self.assertTrue(dct is not None)
89        self.assertTrue(isinstance(dct, dict))
90
91        ode = MyODE()
92        J = PETSc.Mat().create(ts.comm)
93        J.setSizes(3)
94        J.setFromOptions()
95        J.setUp()
96        u, f = J.createVecs()
97
98        ts.setAppCtx(ode)
99        ts.setRHSFunction(ode.rhsfunction, f)
100        ts.setRHSJacobian(ode.rhsjacobian, J, J)
101        ts.setMonitor(ode.monitor)
102
103        ts.snes.ksp.pc.setType('none')
104
105        T0, dT, nT = 0.00, 0.1, 10
106        T = T0 + nT * dT
107        ts.setTime(T0)
108        ts.setTimeStep(dT)
109        ts.setMaxTime(T)
110        ts.setMaxSteps(nT)
111        ts.setFromOptions()
112        u[0], u[1], u[2] = 1, 2, 3
113        if nullsol:
114            ts.setSolution(u)
115            ts.solve()
116        else:
117            ts.solve(u)
118
119        self.assertTrue(ode.rhsfunction_calls > 0)
120        self.assertTrue(ode.rhsjacobian_calls > 0)
121
122        dct = self.ts.getDict()
123        self.assertTrue('__appctx__' in dct)
124        self.assertTrue('__rhsfunction__' in dct)
125        self.assertTrue('__rhsjacobian__' in dct)
126        self.assertTrue('__monitor__' in dct)
127
128        n = ode.monitor_calls
129        ts.monitor(ts.step_number, ts.time)
130        self.assertEqual(ode.monitor_calls, n + 1)
131        n = ode.monitor_calls
132        ts.monitorCancel()
133        ts.monitor(ts.step_number, ts.time)
134        self.assertEqual(ode.monitor_calls, n)
135
136    def testFDColorRHS(self):
137        ts = self.ts
138        ode = MyODE()
139        J = PETSc.Mat().create(ts.comm)
140        J.setSizes(5)
141        J.setType('aij')
142        J.setPreallocationNNZ(nnz=1)
143        u, f = J.createVecs()
144
145        ts.setAppCtx(ode)
146        ts.setRHSFunction(ode.rhsfunction, f)
147        ts.setRHSJacobian(ode.rhsjacobian, J, J)
148        ts.setMonitor(ode.monitor)
149
150        T0, dT, nT = 0.00, 0.1, 10
151        T = T0 + nT * dT
152        ts.setTime(T0)
153        ts.setTimeStep(dT)
154        ts.setMaxTime(T)
155        ts.setMaxSteps(nT)
156        ts.setFromOptions()
157        u[0], u[1], u[2] = 1, 2, 3
158
159        ts.setSolution(u)
160        ode.rhsjacobian(ts, 0, u, J, J)
161        ts.setUp()
162        ts.snes.setUseFD(True)
163        ts.solve(u)
164
165    def testResetAndSolveRHS(self):
166        self.ts.reset()
167        self.ts.setStepNumber(0)
168        self.testSolveRHS()
169        self.ts.reset()
170        self.ts.setStepNumber(0)
171        self.testSolveRHS()
172        self.ts.reset()
173        self.ts.setStepNumber(0)
174        self.testSolveRHS(nullsol=True)
175        self.ts.reset()
176
177
178class BaseTestTSNonlinearI(BaseTestTSNonlinear):
179    def testSolveI(self):
180        ts = self.ts
181        dct = self.ts.getDict()
182        self.assertTrue(dct is not None)
183        self.assertTrue(isinstance(dct, dict))
184
185        ode = MyODE()
186        J = PETSc.Mat().create(ts.comm)
187        J.setSizes(3)
188        J.setFromOptions()
189        J.setUp()
190        u, f = J.createVecs()
191
192        ts.setAppCtx(ode)
193        ts.setIFunction(ode.ifunction, f)
194        ts.setIJacobian(ode.ijacobian, J, J)
195        ts.setMonitor(ode.monitor)
196
197        ts.snes.ksp.pc.setType('none')
198
199        T0, dT, nT = 0.00, 0.1, 10
200        T = T0 + nT * dT
201        ts.setTime(T0)
202        ts.setTimeStep(dT)
203        ts.setMaxTime(T)
204        ts.setMaxSteps(nT)
205        ts.setFromOptions()
206        u[0], u[1], u[2] = 1, 2, 3
207        ts.solve(u)
208
209        self.assertTrue(ode.ifunction_calls > 0)
210        self.assertTrue(ode.ijacobian_calls > 0)
211
212        dct = self.ts.getDict()
213        self.assertTrue('__appctx__' in dct)
214        self.assertTrue('__ifunction__' in dct)
215        self.assertTrue('__ijacobian__' in dct)
216        self.assertTrue('__monitor__' in dct)
217
218        n = ode.monitor_calls
219        ts.monitor(ts.step_number, ts.time)
220        self.assertEqual(ode.monitor_calls, n + 1)
221        n = ode.monitor_calls
222        ts.monitorCancel()
223        ts.monitor(ts.step_number, ts.time)
224        self.assertEqual(ode.monitor_calls, n)
225
226    def testFDColorI(self):
227        ts = self.ts
228        ode = MyODE()
229        J = PETSc.Mat().create(ts.comm)
230        J.setSizes(5)
231        J.setType('aij')
232        J.setPreallocationNNZ(nnz=1)
233        J.setFromOptions()
234        u, f = J.createVecs()
235
236        ts.setAppCtx(ode)
237        ts.setIFunction(ode.ifunction, f)
238        ts.setIJacobian(ode.ijacobian, J, J)
239        ts.setMonitor(ode.monitor)
240
241        T0, dT, nT = 0.00, 0.1, 10
242        T = T0 + nT * dT
243        ts.setTime(T0)
244        ts.setTimeStep(dT)
245        ts.setMaxTime(T)
246        ts.setMaxSteps(nT)
247        ts.setFromOptions()
248        u[0], u[1], u[2] = 1, 2, 3
249
250        ts.setSolution(u)
251        ode.ijacobian(ts, 0, u, 0 * u, 1, J, J)
252        ts.setUp()
253        ts.snes.setUseFD(True)
254        ts.solve(u)
255
256    def testResetAndSolveI(self):
257        self.ts.reset()
258        self.ts.setStepNumber(0)
259        self.testSolveI()
260        self.ts.reset()
261        self.ts.setStepNumber(0)
262        self.testSolveI()
263        self.ts.reset()
264
265
266class TestTSBeuler(BaseTestTSNonlinearRHS, BaseTestTSNonlinearI, unittest.TestCase):
267    TYPE = PETSc.TS.Type.BEULER
268
269
270class TestTSCN(BaseTestTSNonlinearRHS, BaseTestTSNonlinearI, unittest.TestCase):
271    TYPE = PETSc.TS.Type.CN
272
273
274class TestTSTheta(BaseTestTSNonlinearRHS, BaseTestTSNonlinearI, unittest.TestCase):
275    TYPE = PETSc.TS.Type.THETA
276
277
278class TestTSAlpha(BaseTestTSNonlinearRHS, BaseTestTSNonlinearI, unittest.TestCase):
279    TYPE = PETSc.TS.Type.ALPHA
280
281
282# --------------------------------------------------------------------
283
284if __name__ == '__main__':
285    unittest.main()
286
287# --------------------------------------------------------------------
288