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