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