xref: /petsc/src/binding/petsc4py/demo/legacy/ode/fastslowsplit.py (revision d53c7dcfa6b151504f6ca0aeb7a446ad68fe4d68) 
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
10import petsc4py
11
12petsc4py.init(sys.argv)
13
14from petsc4py import PETSc
15import numpy as np
16
17
18class FSS:
19    n = 2
20    comm = PETSc.COMM_SELF
21
22    def __init__(self):
23        pass
24
25    def initialCondition(self, u):
26        u[0] = np.sqrt(2.0)
27        u[1] = np.sqrt(3.0)
28        u.assemble()
29
30    def evalRHSFunction(self, ts, t, u, f):
31        f[0] = (
32            -2.0 * (-1.0 + u[0] * u[0] - np.cos(t)) / (2.0 * u[0])
33            + 0.05 * (-2.0 + u[1] * u[1] - np.cos(5.0 * t)) / (2.0 * u[1])
34            - np.sin(t) / (2.0 * u[0])
35        )
36        f[1] = (
37            0.05 * (-1.0 + u[0] * u[0] - np.cos(t)) / (2.0 * u[0])
38            - (-2.0 + u[1] * u[1] - np.cos(5.0 * t)) / (2.0 * u[1])
39            - 5.0 * np.sin(5.0 * t) / (2.0 * u[1])
40        )
41        f.assemble()
42
43    def evalRHSFunctionSlow(self, ts, t, u, f):
44        f[0] = (
45            -2.0 * (-1.0 + u[0] * u[0] - np.cos(t)) / (2.0 * u[0])
46            + 0.05 * (-2.0 + u[1] * u[1] - np.cos(5.0 * t)) / (2.0 * u[1])
47            - np.sin(t) / (2.0 * u[0])
48        )
49        f.assemble()
50
51    def evalRHSFunctionFast(self, ts, t, u, f):
52        f[0] = 0.05 * (-1.0 + u[0] * u[0] - np.cos(t)) / (2.0 * u[0]) - 5.0 * np.sin(
53            5.0 * t
54        ) / (2.0 * u[1])
55        f.assemble()
56
57    def evalIFunctionFast(self, ts, t, u, udot, f):
58        f[0] = udot[0] + (-2.0 + u[1] * u[1] - np.cos(5.0 * t)) / (2.0 * u[1])
59        f.assemble()
60
61    def evalIJacobianFast(self, ts, t, u, udot, a, A, B):
62        A[0, 0] = a + (2.0 + np.cos(5.0 * t)) / (2.0 * u[1] * u[1]) + 0.5
63        A.assemble()
64        if A != B:
65            B.assemble()
66
67
68OptDB = PETSc.Options()
69explicitform_ = OptDB.getBool('explicitform', False)
70
71ode = FSS()
72
73Jim = PETSc.Mat().createDense([1, 1], comm=ode.comm)
74Jim.setUp()
75
76u = PETSc.Vec().createSeq(ode.n, comm=ode.comm)
77
78ts = PETSc.TS().create(comm=ode.comm)
79ts.setProblemType(ts.ProblemType.NONLINEAR)
80
81if not explicitform_:
82    iss = PETSc.IS().createGeneral([0], comm=ode.comm)
83    isf = PETSc.IS().createGeneral([1], comm=ode.comm)
84    ts.setType(ts.Type.ARKIMEX)
85    ts.setARKIMEXFastSlowSplit(True)
86    ts.setRHSSplitIS('slow', iss)
87    ts.setRHSSplitIS('fast', isf)
88    ts.setRHSSplitRHSFunction('slow', ode.evalRHSFunctionSlow, None)
89    ts.setRHSSplitRHSFunction('fast', ode.evalRHSFunctionFast, None)
90    ts.setRHSSplitIFunction('fast', ode.evalIFunctionFast, None)
91    ts.setRHSSplitIJacobian('fast', ode.evalIJacobianFast, Jim, Jim)
92else:
93    f = u.duplicate()
94    ts.setType(ts.Type.RK)
95    ts.setRHSFunction(ode.evalRHSFunction, f)
96
97ts.setTimeStep(0.01)
98ts.setMaxTime(0.3)
99ts.setExactFinalTime(PETSc.TS.ExactFinalTime.MATCHSTEP)
100
101ts.setFromOptions()
102ode.initialCondition(u)
103ts.solve(u)
104u.view()
105
106del ode, Jim, u, ts
107