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