xref: /petsc/src/binding/petsc4py/demo/legacy/ode/fastslowsplit.py (revision 6a210b70a61472cf8733db59e8a3fa68f6578c01)
1# \begin{eqnarray}
2#                  ys' = -2.0*\frac{-1.0+ys^2.0-\cos(t)}{2.0*ys}+0.05*\frac{-2.0+yf^2-\cos(5.0*t)}{2.0*yf}-\frac{\sin(t)}{2.0*ys}\\
3#                  yf' = 0.05*\frac{-1.0+ys^2-\cos(t)}{2.0*ys}-\frac{-2.0+yf^2-\cos(5.0*t)}{2.0*yf}-5.0*\frac{\sin(5.0*t)}{2.0*yf}\\
4# \end{eqnarray}
5
6# This example demonstrates how to use ARKIMEX for solving a fast-slow system. The system is partitioned additively and component-wise at the same time.
7# ys stands for the slow component and yf stands for the fast component. On the RHS for yf, only the term -\frac{-2.0+yf^2-\cos(5.0*t)}{2.0*yf} is treated implicitly while the rest is treated explicitly.
8
9import sys, petsc4py
10petsc4py.init(sys.argv)
11
12from petsc4py import PETSc
13import numpy as np
14
15class FSS(object):
16    n = 2
17    comm = PETSc.COMM_SELF
18    def __init__(self):
19        pass
20    def initialCondition(self, u):
21        u[0] = np.sqrt(2.0)
22        u[1] = np.sqrt(3.0)
23        u.assemble()
24    def evalRHSFunction(self, ts, t, u, f):
25        f[0] = -2.0 * (-1.0 + u[0] * u[0] - np.cos(t)) / (2.0 * u[0]) + 0.05 * (-2.0 + u[1] * u[1] - np.cos(5.0 * t)) / (2.0 * u[1]) - np.sin(t) / (2.0 * u[0])
26        f[1] = 0.05 * (-1.0 + u[0] * u[0] - np.cos(t)) / (2.0 * u[0]) - (-2.0 + u[1] * u[1] - np.cos(5.0 * t)) / (2.0 * u[1]) - 5.0 * np.sin(5.0 * t) / (2.0 * u[1])
27        f.assemble
28    def evalRHSFunctionSlow(self, ts, t, u, f):
29        f[0] = -2.0 * (-1.0 + u[0] * u[0] - np.cos(t)) / (2.0 * u[0]) + 0.05 * (-2.0 + u[1] * u[1] - np.cos(5.0 * t)) / (2.0 * u[1]) - np.sin(t) / (2.0 * u[0]);
30        f.assemble()
31    def evalRHSFunctionFast(self, ts, t, u, f):
32        f[0] = 0.05 * (-1.0 + u[0] * u[0] - np.cos(t)) / (2.0 * u[0]) - 5.0 * np.sin(5.0 * t) / (2.0 * u[1])
33        f.assemble()
34    def evalIFunctionFast(self, ts, t, u, udot, f):
35        f[0] = udot[0] + (-2.0 + u[1] * u[1] - np.cos(5.0 * t)) / (2.0 * u[1]);
36        f.assemble()
37    def evalIJacobianFast(self, ts, t, u, udot, a, A, B):
38        A[0,0] = a + (2.0 + np.cos(5.0 * t)) / (2.0 * u[1] * u[1]) + 0.5
39        A.assemble()
40        if A != B: B.assemble()
41        return True # same nonzero pattern
42
43OptDB = PETSc.Options()
44explicitform_ = OptDB.getBool('explicitform', False)
45
46ode = FSS()
47
48Jim = PETSc.Mat().createDense([1,1], comm=ode.comm)
49Jim.setUp()
50
51u = PETSc.Vec().createSeq(ode.n, comm=ode.comm)
52
53ts = PETSc.TS().create(comm=ode.comm)
54ts.setProblemType(ts.ProblemType.NONLINEAR)
55
56if not explicitform_:
57    iss = PETSc.IS().createGeneral([0], comm=ode.comm)
58    isf = PETSc.IS().createGeneral([1], comm=ode.comm)
59    ts.setType(ts.Type.ARKIMEX)
60    ts.setARKIMEXFastSlowSplit(True)
61    ts.setRHSSplitIS("slow",iss)
62    ts.setRHSSplitIS("fast",isf)
63    ts.setRHSSplitRHSFunction("slow", ode.evalRHSFunctionSlow, None)
64    ts.setRHSSplitRHSFunction("fast", ode.evalRHSFunctionFast, None)
65    ts.setRHSSplitIFunction("fast", ode.evalIFunctionFast, None)
66    ts.setRHSSplitIJacobian("fast", ode.evalIJacobianFast, Jim, Jim)
67else:
68    f = u.duplicate()
69    ts.setType(ts.Type.RK)
70    ts.setRHSFunction(ode.evalRHSFunction, f)
71
72ts.setTimeStep(0.01)
73ts.setMaxTime(0.3)
74ts.setExactFinalTime(PETSc.TS.ExactFinalTime.MATCHSTEP)
75
76ts.setFromOptions()
77ode.initialCondition(u)
78ts.solve(u)
79u.view()
80
81del ode, Jim, u, ts
82