xref: /petsc/src/binding/petsc4py/demo/legacy/ode/rober.py (revision 55a74a43bb44613d95e937906bec3b8c3581b432)
1*55a74a43SLisandro Dalcin# Stiff 3-variable ODE system from chemical reactions,
2*55a74a43SLisandro Dalcin# due to Robertson (1966),
3*55a74a43SLisandro Dalcin# problem ROBER in Hairer&Wanner, ODE 2, 1996
4*55a74a43SLisandro Dalcin
5*55a74a43SLisandro Dalcinimport sys, petsc4py
6*55a74a43SLisandro Dalcinpetsc4py.init(sys.argv)
7*55a74a43SLisandro Dalcin
8*55a74a43SLisandro Dalcinfrom petsc4py import PETSc
9*55a74a43SLisandro Dalcin
10*55a74a43SLisandro Dalcinclass Rober(object):
11*55a74a43SLisandro Dalcin    n = 3
12*55a74a43SLisandro Dalcin    comm = PETSc.COMM_SELF
13*55a74a43SLisandro Dalcin    def evalSolution(self, t, x):
14*55a74a43SLisandro Dalcin        assert t == 0.0, "only for t=0.0"
15*55a74a43SLisandro Dalcin        x[:] = [1, 0, 0]
16*55a74a43SLisandro Dalcin        x.assemble()
17*55a74a43SLisandro Dalcin    def evalFunction(self, ts, t, x, xdot, f):
18*55a74a43SLisandro Dalcin        f[:] = [xdot[0] + 0.04*x[0] - 1e4*x[1]*x[2],
19*55a74a43SLisandro Dalcin                xdot[1] - 0.04*x[0] + 1e4*x[1]*x[2] + 3e7*x[1]**2,
20*55a74a43SLisandro Dalcin                xdot[2] - 3e7*x[1]**2]
21*55a74a43SLisandro Dalcin        f.assemble()
22*55a74a43SLisandro Dalcin    def evalJacobian(self, ts, t, x, xdot, a, A, B):
23*55a74a43SLisandro Dalcin        J = B
24*55a74a43SLisandro Dalcin        J[:,:] = [[a + 0.04, -1e4*x[2],                -1e4*x[1]],
25*55a74a43SLisandro Dalcin                  [-0.04,    a + 1e4*x[2] + 3e7*2*x[1], 1e4*x[1]],
26*55a74a43SLisandro Dalcin                  [0,        -3e7*2*x[1],               a]]
27*55a74a43SLisandro Dalcin        J.assemble()
28*55a74a43SLisandro Dalcin        if A != B: A.assemble()
29*55a74a43SLisandro Dalcin        return True # same nonzero pattern
30*55a74a43SLisandro Dalcin
31*55a74a43SLisandro DalcinOptDB = PETSc.Options()
32*55a74a43SLisandro Dalcinode = Rober()
33*55a74a43SLisandro Dalcin
34*55a74a43SLisandro DalcinJ = PETSc.Mat().createDense([ode.n, ode.n], comm=ode.comm)
35*55a74a43SLisandro DalcinJ.setUp()
36*55a74a43SLisandro Dalcinx = PETSc.Vec().createSeq(ode.n, comm=ode.comm)
37*55a74a43SLisandro Dalcinf = x.duplicate()
38*55a74a43SLisandro Dalcin
39*55a74a43SLisandro Dalcints = PETSc.TS().create(comm=ode.comm)
40*55a74a43SLisandro Dalcints.setProblemType(ts.ProblemType.NONLINEAR)
41*55a74a43SLisandro Dalcints.setType(ts.Type.THETA)
42*55a74a43SLisandro Dalcin
43*55a74a43SLisandro Dalcints.setIFunction(ode.evalFunction, f)
44*55a74a43SLisandro Dalcints.setIJacobian(ode.evalJacobian, J)
45*55a74a43SLisandro Dalcin
46*55a74a43SLisandro Dalcinhistory = []
47*55a74a43SLisandro Dalcindef monitor(ts, i, t, x):
48*55a74a43SLisandro Dalcin    xx = x[:].tolist()
49*55a74a43SLisandro Dalcin    history.append((i, t, xx))
50*55a74a43SLisandro Dalcints.setMonitor(monitor)
51*55a74a43SLisandro Dalcin
52*55a74a43SLisandro Dalcints.setTime(0.0)
53*55a74a43SLisandro Dalcints.setTimeStep(.001)
54*55a74a43SLisandro Dalcints.setMaxTime(1e30)
55*55a74a43SLisandro Dalcints.setMaxSteps(100)
56*55a74a43SLisandro Dalcints.setExactFinalTime(PETSc.TS.ExactFinalTime.INTERPOLATE)
57*55a74a43SLisandro Dalcin
58*55a74a43SLisandro Dalcints.setFromOptions()
59*55a74a43SLisandro Dalcinode.evalSolution(0.0, x)
60*55a74a43SLisandro Dalcints.solve(x)
61*55a74a43SLisandro Dalcin
62*55a74a43SLisandro Dalcindel ode, J, x, ts
63*55a74a43SLisandro Dalcin
64*55a74a43SLisandro Dalcintry:
65*55a74a43SLisandro Dalcin    from matplotlib import pylab
66*55a74a43SLisandro Dalcinexcept ImportError:
67*55a74a43SLisandro Dalcin    print("matplotlib not available")
68*55a74a43SLisandro Dalcin    raise SystemExit
69*55a74a43SLisandro Dalcin
70*55a74a43SLisandro Dalcinimport numpy as np
71*55a74a43SLisandro Dalcinii = np.asarray([v[0] for v in history])
72*55a74a43SLisandro Dalcintt = np.asarray([v[1] for v in history])
73*55a74a43SLisandro Dalcinxx = np.asarray([v[2] for v in history])
74*55a74a43SLisandro Dalcin
75*55a74a43SLisandro Dalcinpylab.suptitle('Rober')
76*55a74a43SLisandro Dalcinpylab.subplot(2,2,1)
77*55a74a43SLisandro Dalcinpylab.subplots_adjust(wspace=0.3)
78*55a74a43SLisandro Dalcinpylab.semilogy(ii[:-1], np.diff(tt), )
79*55a74a43SLisandro Dalcinpylab.xlabel('step number')
80*55a74a43SLisandro Dalcinpylab.ylabel('timestep')
81*55a74a43SLisandro Dalcin
82*55a74a43SLisandro Dalcinfor i in range(0,3):
83*55a74a43SLisandro Dalcin    pylab.subplot(2,2,i+2)
84*55a74a43SLisandro Dalcin    pylab.semilogx(tt, xx[:,i], "rgb"[i])
85*55a74a43SLisandro Dalcin    pylab.xlabel('t')
86*55a74a43SLisandro Dalcin    pylab.ylabel('x%d' % i)
87*55a74a43SLisandro Dalcin
88*55a74a43SLisandro Dalcinpylab.show()
89