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