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 6import petsc4py 7 8petsc4py.init(sys.argv) 9 10from petsc4py import PETSc 11 12 13class Orego: 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.setArray([1, 2, 3]) 21 22 def evalFunction(self, ts, t, x, xdot, f): 23 f.setArray( 24 [ 25 xdot[0] - 77.27 * (x[1] + x[0] * (1 - 8.375e-6 * x[0] - x[1])), 26 xdot[1] - 1 / 77.27 * (x[2] - (1 + x[0]) * x[1]), 27 xdot[2] - 0.161 * (x[0] - x[2]), 28 ] 29 ) 30 31 def evalJacobian(self, ts, t, x, xdot, a, A, B): 32 B[:, :] = [ 33 [ 34 a - 77.27 * ((1 - 8.375e-6 * x[0] - x[1]) - 8.375e-6 * x[0]), 35 -77.27 * (1 - x[0]), 36 0, 37 ], 38 [1 / 77.27 * x[1], a + 1 / 77.27 * (1 + x[0]), -1 / 77.27], 39 [-0.161, 0, a + 0.161], 40 ] 41 B.assemble() 42 if A != B: 43 A.assemble() 44 45 46OptDB = PETSc.Options() 47ode = Orego() 48 49J = PETSc.Mat().createDense([ode.n, ode.n], comm=ode.comm) 50J.setUp() 51x = PETSc.Vec().createSeq(ode.n, comm=ode.comm) 52f = x.duplicate() 53 54ts = PETSc.TS().create(comm=ode.comm) 55ts.setType(ts.Type.ROSW) # Rosenbrock-W. ARKIMEX is a nonlinearly implicit alternative. 56 57ts.setIFunction(ode.evalFunction, f) 58ts.setIJacobian(ode.evalJacobian, J) 59 60history = [] 61 62 63def monitor(ts, i, t, x): 64 xx = x[:].tolist() 65 history.append((i, t, xx)) 66 67 68ts.setMonitor(monitor) 69 70ts.setTime(0.0) 71ts.setTimeStep(0.1) 72ts.setMaxTime(360) 73ts.setMaxSteps(2000) 74ts.setExactFinalTime(PETSc.TS.ExactFinalTime.INTERPOLATE) 75ts.setMaxSNESFailures( 76 -1 77) # allow an unlimited number of failures (step will be rejected and retried) 78 79# Set a different tolerance on each variable. Can use a scalar or a vector for either or both atol and rtol. 80vatol = x.duplicate(array=[1e-2, 1e-1, 1e-4]) 81ts.setTolerances( 82 atol=vatol, rtol=1e-3 83) # adaptive controller attempts to match this tolerance 84 85snes = ts.getSNES() # Nonlinear solver 86snes.setTolerances( 87 max_it=10 88) # Stop nonlinear solve after 10 iterations (TS will retry with shorter step) 89ksp = snes.getKSP() # Linear solver 90ksp.setType(ksp.Type.PREONLY) # Just use the preconditioner without a Krylov method 91pc = ksp.getPC() # Preconditioner 92pc.setType(pc.Type.LU) # Use a direct solve 93 94ts.setFromOptions() # Apply run-time options, e.g. -ts_adapt_monitor -ts_type arkimex -snes_converged_reason 95ode.evalSolution(0.0, x) 96ts.solve(x) 97print( 98 'steps %d (%d rejected, %d SNES fails), nonlinear its %d, linear its %d' 99 % ( 100 ts.getStepNumber(), 101 ts.getStepRejections(), 102 ts.getSNESFailures(), 103 ts.getSNESIterations(), 104 ts.getKSPIterations(), 105 ) 106) 107 108if OptDB.getBool('plot_history', True): 109 from matplotlib import pylab 110 from matplotlib import rc 111 import numpy as np 112 113 ii = np.asarray([v[0] for v in history]) 114 tt = np.asarray([v[1] for v in history]) 115 xx = np.asarray([v[2] for v in history]) 116 117 rc('text', usetex=True) 118 pylab.suptitle('Oregonator: TS \\texttt{%s}' % ts.getType()) 119 pylab.subplot(2, 2, 1) 120 pylab.subplots_adjust(wspace=0.3) 121 pylab.semilogy( 122 ii[:-1], 123 np.diff(tt), 124 ) 125 pylab.xlabel('step number') 126 pylab.ylabel('timestep') 127 128 for i in range(3): 129 pylab.subplot(2, 2, i + 2) 130 pylab.semilogy(tt, xx[:, i], 'rgb'[i]) 131 pylab.xlabel('time') 132 pylab.ylabel('$x_%d$' % i) 133 134 # pylab.savefig('orego-history.png') 135 pylab.show() 136