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, petsc4py 6petsc4py.init(sys.argv) 7 8from petsc4py import PETSc 9 10class Rober(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[:] = [1, 0, 0] 16 x.assemble() 17 def evalFunction(self, ts, t, x, xdot, f): 18 f[:] = [xdot[0] + 0.04*x[0] - 1e4*x[1]*x[2], 19 xdot[1] - 0.04*x[0] + 1e4*x[1]*x[2] + 3e7*x[1]**2, 20 xdot[2] - 3e7*x[1]**2] 21 f.assemble() 22 def evalJacobian(self, ts, t, x, xdot, a, A, B): 23 J = B 24 J[:,:] = [[a + 0.04, -1e4*x[2], -1e4*x[1]], 25 [-0.04, a + 1e4*x[2] + 3e7*2*x[1], 1e4*x[1]], 26 [0, -3e7*2*x[1], a]] 27 J.assemble() 28 if A != B: A.assemble() 29 return True # same nonzero pattern 30 31OptDB = PETSc.Options() 32ode = Rober() 33 34J = PETSc.Mat().createDense([ode.n, ode.n], comm=ode.comm) 35J.setUp() 36x = PETSc.Vec().createSeq(ode.n, comm=ode.comm) 37f = x.duplicate() 38 39ts = PETSc.TS().create(comm=ode.comm) 40ts.setProblemType(ts.ProblemType.NONLINEAR) 41ts.setType(ts.Type.THETA) 42 43ts.setIFunction(ode.evalFunction, f) 44ts.setIJacobian(ode.evalJacobian, J) 45 46history = [] 47def monitor(ts, i, t, x): 48 xx = x[:].tolist() 49 history.append((i, t, xx)) 50ts.setMonitor(monitor) 51 52ts.setTime(0.0) 53ts.setTimeStep(.001) 54ts.setMaxTime(1e30) 55ts.setMaxSteps(100) 56ts.setExactFinalTime(PETSc.TS.ExactFinalTime.INTERPOLATE) 57 58ts.setFromOptions() 59ode.evalSolution(0.0, x) 60ts.solve(x) 61 62del ode, J, x, ts 63 64try: 65 from matplotlib import pylab 66except ImportError: 67 print("matplotlib not available") 68 raise SystemExit 69 70import numpy as np 71ii = np.asarray([v[0] for v in history]) 72tt = np.asarray([v[1] for v in history]) 73xx = np.asarray([v[2] for v in history]) 74 75pylab.suptitle('Rober') 76pylab.subplot(2,2,1) 77pylab.subplots_adjust(wspace=0.3) 78pylab.semilogy(ii[:-1], np.diff(tt), ) 79pylab.xlabel('step number') 80pylab.ylabel('timestep') 81 82for i in range(0,3): 83 pylab.subplot(2,2,i+2) 84 pylab.semilogx(tt, xx[:,i], "rgb"[i]) 85 pylab.xlabel('t') 86 pylab.ylabel('x%d' % i) 87 88pylab.show() 89