File: rober.py

package info (click to toggle)
petsc4py 3.23.1-1exp2
  • links: PTS, VCS
  • area: main
  • in suites: experimental
  • size: 3,448 kB
  • sloc: python: 12,503; ansic: 1,697; makefile: 343; f90: 313; sh: 14
file content (103 lines) | stat: -rw-r--r-- 2,196 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
# Stiff 3-variable ODE system from chemical reactions,
# due to Robertson (1966),
# problem ROBER in Hairer&Wanner, ODE 2, 1996

import sys
import petsc4py

petsc4py.init(sys.argv)

from petsc4py import PETSc


class Rober:
    n = 3
    comm = PETSc.COMM_SELF

    def evalSolution(self, t, x):
        if t != 0.0:
            raise ValueError('Only for t=0')
        x[:] = [1, 0, 0]
        x.assemble()

    def evalFunction(self, ts, t, x, xdot, f):
        f[:] = [
            xdot[0] + 0.04 * x[0] - 1e4 * x[1] * x[2],
            xdot[1] - 0.04 * x[0] + 1e4 * x[1] * x[2] + 3e7 * x[1] ** 2,
            xdot[2] - 3e7 * x[1] ** 2,
        ]
        f.assemble()

    def evalJacobian(self, ts, t, x, xdot, a, A, B):
        J = B
        J[:, :] = [
            [a + 0.04, -1e4 * x[2], -1e4 * x[1]],
            [-0.04, a + 1e4 * x[2] + 3e7 * 2 * x[1], 1e4 * x[1]],
            [0, -3e7 * 2 * x[1], a],
        ]
        J.assemble()
        if A != B:
            A.assemble()


OptDB = PETSc.Options()
ode = Rober()

J = PETSc.Mat().createDense([ode.n, ode.n], comm=ode.comm)
J.setUp()
x = PETSc.Vec().createSeq(ode.n, comm=ode.comm)
f = x.duplicate()

ts = PETSc.TS().create(comm=ode.comm)
ts.setProblemType(ts.ProblemType.NONLINEAR)
ts.setType(ts.Type.THETA)

ts.setIFunction(ode.evalFunction, f)
ts.setIJacobian(ode.evalJacobian, J)

history = []


def monitor(ts, i, t, x):
    xx = x[:].tolist()
    history.append((i, t, xx))


ts.setMonitor(monitor)

ts.setTime(0.0)
ts.setTimeStep(0.001)
ts.setMaxTime(1e30)
ts.setMaxSteps(100)
ts.setExactFinalTime(PETSc.TS.ExactFinalTime.INTERPOLATE)

ts.setFromOptions()
ode.evalSolution(0.0, x)
ts.solve(x)

del ode, J, x, ts

from matplotlib import pylab
import numpy as np

ii = np.asarray([v[0] for v in history])
tt = np.asarray([v[1] for v in history])
xx = np.asarray([v[2] for v in history])

pylab.suptitle('Rober')
pylab.subplot(2, 2, 1)
pylab.subplots_adjust(wspace=0.3)
pylab.semilogy(
    ii[:-1],
    np.diff(tt),
)
pylab.xlabel('step number')
pylab.ylabel('timestep')

for i in range(3):
    pylab.subplot(2, 2, i + 2)
    pylab.semilogx(tt, xx[:, i], 'rgb'[i])
    pylab.xlabel('t')
    pylab.ylabel('x%d' % i)

pylab.show()