1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
|
import sys
import petsc4py
petsc4py.init(sys.argv)
from petsc4py import PETSc
class BouncingBall:
n = 2
comm = PETSc.COMM_SELF
def __init__(self):
pass
def initialCondition(self, u):
u[0] = 0.0
u[1] = 20.0
u.assemble()
def evalRHSFunction(self, ts, t, u, f):
f[0] = u[1]
f[1] = -9.8
f.assemble()
def evalRHSJacobian(self, ts, t, u, A, B):
J = A
J[0, 0] = 0.0
J[1, 0] = 0.0
J[0, 1] = 1.0
J[1, 1] = 0.0
J.assemble()
if A != B:
B.assemble()
class Monitor:
def __init__(self):
pass
def __call__(self, ts, k, t, x):
PETSc.Sys.Print(f'Position at time {t}: {x[0]}')
ode = BouncingBall()
J = PETSc.Mat().create()
J.setSizes([ode.n, ode.n])
J.setType('aij')
J.setUp()
J.assemble()
u = PETSc.Vec().createSeq(ode.n, comm=ode.comm)
f = u.duplicate()
ts = PETSc.TS().create(comm=ode.comm)
ts.setProblemType(ts.ProblemType.NONLINEAR)
ts.setType(ts.Type.BEULER)
ts.setRHSFunction(ode.evalRHSFunction, f)
ts.setRHSJacobian(ode.evalRHSJacobian, J)
# ts.setSaveTrajectory()
ts.setTime(0.0)
ts.setTimeStep(0.01)
ts.setMaxTime(15.0)
# ts.setMaxSteps(1000)
ts.setExactFinalTime(PETSc.TS.ExactFinalTime.MATCHSTEP)
# ts.setMonitor(Monitor())
direction = [-1]
terminate = [False]
def event(ts, t, X, fvalue):
fvalue[0] = X[0]
def postevent(ts, events, t, X, forward):
X[0] = 0.0
X[1] = -0.9 * X[1]
X.assemble()
ts.setEventHandler(direction, terminate, event, postevent)
ts.setEventTolerances(1e-6, vtol=[1e-9])
ts.setFromOptions()
ode.initialCondition(u)
ts.solve(u)
|
