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