xref: /petsc/src/binding/petsc4py/demo/legacy/ode/ce.py (revision bcee047adeeb73090d7e36cc71e39fc287cdbb97) !
1# Stiff scalar valued ODE problem with an exact solution
2
3import sys, petsc4py
4petsc4py.init(sys.argv)
5
6from petsc4py import PETSc
7from math import sin, cos, exp
8
9class CE(object):
10    n = 1
11    comm = PETSc.COMM_SELF
12    def __init__(self, lambda_=1.0):
13        self.lambda_ = lambda_
14    def evalSolution(self, t, x):
15        l = self.lambda_
16        x[0] = l/(l*l+1)*(l*cos(t)+sin(t)) - l*l/(l*l+1)*exp(-l*t)
17        x.assemble()
18    def evalFunction(self, ts, t, x, xdot, f):
19        l = self.lambda_
20        f[0] = xdot[0] + l*(x[0] - cos(t))
21        f.assemble()
22    def evalJacobian(self, ts, t, x, xdot, a, A, B):
23        J = B
24        l = self.lambda_
25        J[0,0] = a + l
26        J.assemble()
27        if A != B: A.assemble()
28        return True # same nonzero pattern
29
30OptDB = PETSc.Options()
31
32lambda_ = OptDB.getScalar('lambda', 10.0)
33ode = CE(lambda_)
34
35J = PETSc.Mat().createDense([ode.n,ode.n], comm=ode.comm)
36J.setUp()
37x = PETSc.Vec().createSeq(ode.n, comm=ode.comm)
38f = x.duplicate()
39
40ts = PETSc.TS().create(comm=ode.comm)
41ts.setProblemType(ts.ProblemType.NONLINEAR)
42ts.setType(ts.Type.GLLE)
43
44ts.setIFunction(ode.evalFunction, f)
45ts.setIJacobian(ode.evalJacobian, J)
46
47ts.setTime(0.0)
48ts.setTimeStep(0.001)
49ts.setMaxTime(10)
50ts.setMaxSteps(10000)
51ts.setExactFinalTime(PETSc.TS.ExactFinalTime.INTERPOLATE)
52
53class Monitor(object):
54    def __init__(self, ode):
55        self.ode = ode
56        self.x = PETSc.Vec().createSeq(ode.n, comm=ode.comm)
57    def __call__(self, ts, k, t, x):
58        ode.evalSolution(t, self.x)
59        self.x.axpy(-1, x)
60        e = self.x.norm()
61        h = ts.getTimeStep()
62        PETSc.Sys.Print("step %3d t=%8.2e h=%8.2e error=%8.2e" %
63                        (k, t, h, e), comm=self.ode.comm)
64
65ts.setMonitor(Monitor(ode))
66
67ts.setFromOptions()
68ode.evalSolution(0.0, x)
69ts.solve(x)
70
71del ode, J, x, f, ts
72