1*c4762a1bSJed Brownimport sys, petsc4py 2*c4762a1bSJed Brownpetsc4py.init(sys.argv) 3*c4762a1bSJed Brownfrom petsc4py import PETSc 4*c4762a1bSJed Brownimport math 5*c4762a1bSJed Brown 6*c4762a1bSJed Brownclass MyODE: 7*c4762a1bSJed Brown def __init__(self,da): 8*c4762a1bSJed Brown self.da = da 9*c4762a1bSJed Brown def getCorners(self): 10*c4762a1bSJed Brown 'get corners and ghost corners, take first element of array because this is a 1D problem' 11*c4762a1bSJed Brown xs, xm = self.da.getCorners() 12*c4762a1bSJed Brown gxs, gxm = self.da.getGhostCorners() 13*c4762a1bSJed Brown return xs[0], xm[0], gxs[0], gxm[0] 14*c4762a1bSJed Brown def function(self, ts,t,x,xdot,f): 15*c4762a1bSJed Brown mx = da.getSizes(); mx = mx[0]; hx = 1.0/mx 16*c4762a1bSJed Brown (xs,xm,gxs,gxm) = self.getCorners() 17*c4762a1bSJed Brown xx = da.createLocalVector() 18*c4762a1bSJed Brown xxdot = da.createLocalVector() 19*c4762a1bSJed Brown da.globalToLocal(x,xx) 20*c4762a1bSJed Brown da.globalToLocal(xdot,xxdot) 21*c4762a1bSJed Brown dt = ts.getTimeStep() 22*c4762a1bSJed Brown x0 = ts.getSolution() 23*c4762a1bSJed Brown lxs, lxe = xs, xs+xm 24*c4762a1bSJed Brown if xs == 0: f[0] = xx[0]/hx; lxs+=1; 25*c4762a1bSJed Brown if lxe == mx: f[mx-1] = xx[mx-1-gxs]/hx; lxe-=1 26*c4762a1bSJed Brown for i in range(lxs,lxe): 27*c4762a1bSJed Brown f[i] = xxdot[i-gxs] + (2.0*xx[i-gxs] - xx[i-1-gxs] - xx[i+1-gxs])/hx - hx*math.exp(xx[i-gxs]) 28*c4762a1bSJed Brown f.assemble() 29*c4762a1bSJed Brown def jacobian(self,ts,t,x,xdot,shift,J,P): 30*c4762a1bSJed Brown mx = da.getSizes(); mx = mx[0]; hx = 1.0/mx 31*c4762a1bSJed Brown (xs,xm,gxs,gxm) = self.getCorners() 32*c4762a1bSJed Brown xx = da.createLocalVector() 33*c4762a1bSJed Brown da.globalToLocal(x,xx) 34*c4762a1bSJed Brown x0 = ts.getSolution() 35*c4762a1bSJed Brown dt = ts.getTimeStep() 36*c4762a1bSJed Brown P.zeroEntries() 37*c4762a1bSJed Brown lxs, lxe = xs, xs+xm 38*c4762a1bSJed Brown if xs == 0: P.setValues([0],[0],1.0/hx); lxs+=1 39*c4762a1bSJed Brown if lxe == mx: P.setValues([mx-1],[mx-1],1.0/hx); lxe-=1 40*c4762a1bSJed Brown for i in range(lxs,lxe): 41*c4762a1bSJed Brown P.setValues([i],[i-1,i,i+1],[-1.0/hx,2.0/hx-hx*math.exp(xx[i-gxs])+shift,-1.0/hx]) 42*c4762a1bSJed Brown P.assemble() 43*c4762a1bSJed Brown return True # same_nz 44*c4762a1bSJed Brown 45*c4762a1bSJed BrownM = PETSc.Options().getInt('M', 9) 46*c4762a1bSJed Brownda = PETSc.DA().create([M],comm=PETSc.COMM_WORLD) 47*c4762a1bSJed Brownf = da.createGlobalVector() 48*c4762a1bSJed Brownx = f.duplicate() 49*c4762a1bSJed BrownJ = da.getMatrix(PETSc.Mat.Type.AIJ); 50*c4762a1bSJed Brown 51*c4762a1bSJed Brownts = PETSc.TS().create(PETSc.COMM_WORLD) 52*c4762a1bSJed Brownts.setProblemType(PETSc.TS.ProblemType.NONLINEAR) 53*c4762a1bSJed Brownts.setType(ts.Type.GL) 54*c4762a1bSJed Brown 55*c4762a1bSJed Brownode = MyODE(da) 56*c4762a1bSJed Brownts.setIFunction(ode.function, f) 57*c4762a1bSJed Brownts.setIJacobian(ode.jacobian, J) 58*c4762a1bSJed Brown 59*c4762a1bSJed Brownts.setTimeStep(0.1) 60*c4762a1bSJed Brownts.setDuration(10, 1.0) 61*c4762a1bSJed Brownts.setFromOptions() 62*c4762a1bSJed Brownx.set(1.0) 63*c4762a1bSJed Brownts.solve(x) 64