xref: /petsc/src/binding/petsc4py/demo/legacy/ode/orego.py (revision bcee047adeeb73090d7e36cc71e39fc287cdbb97) 
1# Oregonator: stiff 3-variable oscillatory ODE system from chemical reactions,
2# problem OREGO in Hairer&Wanner volume 2
3# See also http://www.scholarpedia.org/article/Oregonator
4
5import sys, petsc4py
6petsc4py.init(sys.argv)
7
8from petsc4py import PETSc
9
10class Orego(object):
11    n = 3
12    comm = PETSc.COMM_SELF
13    def evalSolution(self, t, x):
14        assert t == 0.0, "only for t=0.0"
15        x.setArray([1, 2, 3])
16    def evalFunction(self, ts, t, x, xdot, f):
17        f.setArray([xdot[0] - 77.27*(x[1] + x[0]*(1 - 8.375e-6*x[0] - x[1])),
18                    xdot[1] - 1/77.27*(x[2] - (1 + x[0])*x[1]),
19                    xdot[2] - 0.161*(x[0] - x[2])])
20    def evalJacobian(self, ts, t, x, xdot, a, A, B):
21        B[:,:] = [[a - 77.27*((1 - 8.375e-6*x[0] - x[1]) - 8.375e-6*x[0]),   -77.27*(1 - x[0]),               0],
22                  [1/77.27*x[1],                                             a + 1/77.27*(1 + x[0]),   -1/77.27],
23                  [-0.161,                                                           0,               a + 0.161]]
24        B.assemble()
25        if A != B: A.assemble()
26        return True # same nonzero pattern
27
28OptDB = PETSc.Options()
29ode = Orego()
30
31J = PETSc.Mat().createDense([ode.n, ode.n], comm=ode.comm)
32J.setUp()
33x = PETSc.Vec().createSeq(ode.n, comm=ode.comm)
34f = x.duplicate()
35
36ts = PETSc.TS().create(comm=ode.comm)
37ts.setType(ts.Type.ROSW)        # Rosenbrock-W. ARKIMEX is a nonlinearly implicit alternative.
38
39ts.setIFunction(ode.evalFunction, f)
40ts.setIJacobian(ode.evalJacobian, J)
41
42history = []
43def monitor(ts, i, t, x):
44    xx = x[:].tolist()
45    history.append((i, t, xx))
46ts.setMonitor(monitor)
47
48ts.setTime(0.0)
49ts.setTimeStep(0.1)
50ts.setMaxTime(360)
51ts.setMaxSteps(2000)
52ts.setExactFinalTime(PETSc.TS.ExactFinalTime.INTERPOLATE)
53ts.setMaxSNESFailures(-1)       # allow an unlimited number of failures (step will be rejected and retried)
54
55# Set a different tolerance on each variable. Can use a scalar or a vector for either or both atol and rtol.
56vatol = x.duplicate(array=[1e-2, 1e-1, 1e-4])
57ts.setTolerances(atol=vatol,rtol=1e-3) # adaptive controller attempts to match this tolerance
58
59snes = ts.getSNES()             # Nonlinear solver
60snes.setTolerances(max_it=10)   # Stop nonlinear solve after 10 iterations (TS will retry with shorter step)
61ksp = snes.getKSP()             # Linear solver
62ksp.setType(ksp.Type.PREONLY)   # Just use the preconditioner without a Krylov method
63pc = ksp.getPC()                # Preconditioner
64pc.setType(pc.Type.LU)          # Use a direct solve
65
66ts.setFromOptions()             # Apply run-time options, e.g. -ts_adapt_monitor -ts_type arkimex -snes_converged_reason
67ode.evalSolution(0.0, x)
68ts.solve(x)
69print('steps %d (%d rejected, %d SNES fails), nonlinear its %d, linear its %d'
70      % (ts.getStepNumber(), ts.getStepRejections(), ts.getSNESFailures(),
71         ts.getSNESIterations(), ts.getKSPIterations()))
72
73if OptDB.getBool('plot_history', True):
74    try:
75        from matplotlib import pylab
76        from matplotlib import rc
77    except ImportError:
78        print("matplotlib not available")
79        raise SystemExit
80
81    import numpy as np
82    ii = np.asarray([v[0] for v in history])
83    tt = np.asarray([v[1] for v in history])
84    xx = np.asarray([v[2] for v in history])
85
86    rc('text', usetex=True)
87    pylab.suptitle('Oregonator: TS \\texttt{%s}' % ts.getType())
88    pylab.subplot(2,2,1)
89    pylab.subplots_adjust(wspace=0.3)
90    pylab.semilogy(ii[:-1], np.diff(tt), )
91    pylab.xlabel('step number')
92    pylab.ylabel('timestep')
93
94    for i in range(0,3):
95        pylab.subplot(2,2,i+2)
96        pylab.semilogy(tt, xx[:,i], "rgb"[i])
97        pylab.xlabel('time')
98        pylab.ylabel('$x_%d$' % i)
99
100    # pylab.savefig('orego-history.png')
101    pylab.show()
102