xref: /petsc/src/binding/petsc4py/demo/legacy/ode/rober.py (revision d53c7dcfa6b151504f6ca0aeb7a446ad68fe4d68) 
1# Stiff 3-variable ODE system from chemical reactions,
2# due to Robertson (1966),
3# problem ROBER in Hairer&Wanner, ODE 2, 1996
4
5import sys
6import petsc4py
7
8petsc4py.init(sys.argv)
9
10from petsc4py import PETSc
11
12
13class Rober:
14    n = 3
15    comm = PETSc.COMM_SELF
16
17    def evalSolution(self, t, x):
18        if t != 0.0:
19            raise ValueError('Only for t=0')
20        x[:] = [1, 0, 0]
21        x.assemble()
22
23    def evalFunction(self, ts, t, x, xdot, f):
24        f[:] = [
25            xdot[0] + 0.04 * x[0] - 1e4 * x[1] * x[2],
26            xdot[1] - 0.04 * x[0] + 1e4 * x[1] * x[2] + 3e7 * x[1] ** 2,
27            xdot[2] - 3e7 * x[1] ** 2,
28        ]
29        f.assemble()
30
31    def evalJacobian(self, ts, t, x, xdot, a, A, B):
32        J = B
33        J[:, :] = [
34            [a + 0.04, -1e4 * x[2], -1e4 * x[1]],
35            [-0.04, a + 1e4 * x[2] + 3e7 * 2 * x[1], 1e4 * x[1]],
36            [0, -3e7 * 2 * x[1], a],
37        ]
38        J.assemble()
39        if A != B:
40            A.assemble()
41
42
43OptDB = PETSc.Options()
44ode = Rober()
45
46J = PETSc.Mat().createDense([ode.n, ode.n], comm=ode.comm)
47J.setUp()
48x = PETSc.Vec().createSeq(ode.n, comm=ode.comm)
49f = x.duplicate()
50
51ts = PETSc.TS().create(comm=ode.comm)
52ts.setProblemType(ts.ProblemType.NONLINEAR)
53ts.setType(ts.Type.THETA)
54
55ts.setIFunction(ode.evalFunction, f)
56ts.setIJacobian(ode.evalJacobian, J)
57
58history = []
59
60
61def monitor(ts, i, t, x):
62    xx = x[:].tolist()
63    history.append((i, t, xx))
64
65
66ts.setMonitor(monitor)
67
68ts.setTime(0.0)
69ts.setTimeStep(0.001)
70ts.setMaxTime(1e30)
71ts.setMaxSteps(100)
72ts.setExactFinalTime(PETSc.TS.ExactFinalTime.INTERPOLATE)
73
74ts.setFromOptions()
75ode.evalSolution(0.0, x)
76ts.solve(x)
77
78del ode, J, x, ts
79
80from matplotlib import pylab
81import numpy as np
82
83ii = np.asarray([v[0] for v in history])
84tt = np.asarray([v[1] for v in history])
85xx = np.asarray([v[2] for v in history])
86
87pylab.suptitle('Rober')
88pylab.subplot(2, 2, 1)
89pylab.subplots_adjust(wspace=0.3)
90pylab.semilogy(
91    ii[:-1],
92    np.diff(tt),
93)
94pylab.xlabel('step number')
95pylab.ylabel('timestep')
96
97for i in range(3):
98    pylab.subplot(2, 2, i + 2)
99    pylab.semilogx(tt, xx[:, i], 'rgb'[i])
100    pylab.xlabel('t')
101    pylab.ylabel('x%d' % i)
102
103pylab.show()
104