""" The multigrid module provides a framework for solving elliptic problems. A multigrid object is just a list of grids, from the finest mesh down (by factors of two) to a single interior zone (each grid has the same number of guardcells). The main multigrid class (MGfv) is setup to solve Poisson's equation with Dirichlet boundary conditions. Implementations for other equations can be created by making a new class that inherits the MGfv class, and provides new versions of the residual and smooth methods. The general usage is as follows: > a = multigrid.MGfv(nx, lBC=0, rBC=1, verbose=1) this creates the multigrid object a, with a finest grid of nx zones and a left/right boundary conditions of 0 and 1 respectively. Setting verbose = 1 causing debugging information to be output, so you can see the residual errors in each of the V-cycles. > a.initSolution(zeros(nx, Float64)) this initializes the solution vector with zeros > a.initRHS(zeros(nx, Float64)) this initializes the RHS on the finest grid to 0 (Laplace's equation). Any RHS can be set by passing through an array of nx values here. Then to solve, you just do: > a.solve(rtol = 1.e-10) where rtol is the desired tolerance (relative difference in solution from one cycle to the next). to access the final solution, use the getSolution method v = a.getSolution() For convinence, the grid information on the solution level is available as attributes to the class, a.imin and a.imax are the indices bounding the interior of the solution array (i.e. excluding the guardcells). a.x is the coordinate array a.dx is the grid spacing """ import patch1d import math from numarray import * def error(imin, imax, dx, r): # L2 norm of elements in r, multiplied by dx to # normalize return sqrt(dx*sum(r[imin:imax+1]**2)) class MGfv: def __init__(self, nx, xmin=0.0, xmax=1.0, lBC=0, rBC=0, verbose=0): self.nx = nx self.ng = 2 self.xmin = xmin self.xmax = xmax self.lBC = lBC self.rBC = rBC self.nsmooth = 3 self.maxCycles = 100 self.verbose = verbose # a small number used in computing the error, so we don't divide by 0 self.small = 1.e-16 # keep track of whether we've initialized the solution self.initializedSolution = 0 self.initializedRHS = 0 # assume that self.nx = 2^nlevels self.nlevels = int(math.log(self.nx)/math.log(2.0)) + 1 # a multigrid object will be a list of grids self.grids = [] # create the grids. Here, self.grids[0] will be the coarsest # grid and self.grids[nlevel-1] will be the finest grid # we store the solution, v, the rhs, f. i = 0 nx_t = 1 while (i < self.nlevels): self.grids.append(patch1d.patch1d(nx_t, ng=self.ng, xmin=xmin, xmax=xmax)) self.grids[i].registerVar("v") self.grids[i].registerVar("f") self.grids[i].init() if self.verbose: print self.grids[i] nx_t = nx_t*2 i += 1 # provide coordinate and indexing information for the solution mesh self.imin = self.grids[self.nlevels-1].imin self.imax = self.grids[self.nlevels-1].imax self.x = self.grids[self.nlevels-1].x self.dx = self.grids[self.nlevels-1].dx # after solving, keep track of the number of cycles taken and the # relative error from the previous cycle self.numCycles = 0 self.error = 1.e33 def getSolution(self): return self.grids[self.nlevels-1].getVar("v") def initSolution(self, data): v = self.grids[self.nlevels-1].getVar("v") v[self.grids[self.nlevels-1].imin: self.grids[self.nlevels-1].imax+1] = data self.initializedSolution = 1 def initRHS(self, data): f = self.grids[self.nlevels-1].getVar("f") f[self.grids[self.nlevels-1].imin: self.grids[self.nlevels-1].imax+1] = data self.initializedRHS = 1 def residual(self, level): """ compute the residual """ v = self.grids[level].getVar("v") f = self.grids[level].getVar("f") imin = self.grids[level].imin imax = self.grids[level].imax nx = self.grids[level].nx ng = self.grids[level].ng dx = self.grids[level].dx r = zeros(2*ng + nx, Float64) r[imin:imax+1] = f[imin:imax+1] - \ (v[imin-1:imax] - 2.0*v[imin:imax+1] + v[imin+1:imax+2])/(dx*dx) return r def smooth(self, level, nsmooth): v = self.grids[level].getVar("v") f = self.grids[level].getVar("f") imin = self.grids[level].imin imax = self.grids[level].imax dx = self.grids[level].dx # only the solution on the original grid satisfies the user # boundary conditions. The coarser grids evolve the error # where the BC is 0 if level == self.nlevels-1: lBC = self.lBC rBC = self.rBC else: lBC = 0 rBC = 0 # assume we are solving u_{i+1} - 2 u_i + u_{i-1} = dx**2 f # do red-black G-S i = 0 while (i < nsmooth): # set the guardcells to give the proper boundary condition # on the interface v[imin-1] = 2*lBC - v[imin] v[imax+1] = 2*rBC - v[imax] v[imin:imax+1:2] = 0.5*(v[imin-1:imax:2] + v[imin+1:imax+2:2] - dx*dx*f[imin:imax+1:2]) v[imin+1:imax+1:2] = 0.5*(v[imin:imax:2] + v[imin+2:imax+2:2] - dx*dx*f[imin+1:imax+1:2]) i += 1 def solve(self, rtol = 1.e-11): # start by making sure that we've initialized the solution # and the RHS if (not self.initializedSolution or not self.initializedRHS): print "ERROR: solution and RHS are not initialized" return -1 # for now, we will just do V-cycles, continuing until we achieve the L2 # norm of the relative solution difference is < rtol if self.verbose: print "source norm = ", error(self.grids[self.nlevels-1].imin, self.grids[self.nlevels-1].imax, self.grids[self.nlevels-1].dx, self.grids[self.nlevels-1].getVar("f")) oldSolution = self.grids[self.nlevels-1].getVar("v").copy() converged = 0 cycle = 1 while (not converged and cycle <= self.maxCycles): # zero out the solution on all but the finest grid level = 0 while (level < self.nlevels-1): v = self.grids[level].getVar("v") v[:] = 0.0 level += 1 # descending part if self.verbose: print "<<< beginning V-cycle (cycle %d) >>>\n" % cycle level = self.nlevels-1 while (level > 0): fP = self.grids[level] cP = self.grids[level-1] if self.verbose: print " level = %d, nx = %d" % (level, fP.nx) print " before G-S, residual L2 norm = %g" % \ (error(fP.imin, fP.imax, fP.dx, self.residual(level))) # smooth on the current level self.smooth(level, self.nsmooth) if self.verbose: print " after G-S, residual L2 norm = %g\n" % \ (error(fP.imin, fP.imax, fP.dx, self.residual(level))) # compute the residual r = self.residual(level) # restrict the residual down to the RHS of the coarser level patch1d._restrict(fP.nx, fP.ng, r, cP.nx, cP.ng, cP.getVar("f")) level -= 1 # solve the coarse problem exactly if self.verbose: print " <<< bottom solve >>>\n" bP = self.grids[0] v = bP.getVar("v") f = bP.getVar("f") v[bP.imin] = -0.5*(f[bP.imin]*bP.dx*bP.dx - v[bP.imin-1] - v[bP.imin+1]) # ascending part level = 1 while (level < self.nlevels): fP = self.grids[level] cP = self.grids[level-1] # allocate storage for the error e = zeros(2*fP.ng + fP.nx, Float64) # prolong the error up from the coarse grid patch1d._prolong(cP.nx, cP.ng, cP.getVar("v"), fP.nx, fP.ng, e) # correct the solution on the current grid v = fP.getVar("v") v += e if self.verbose: print " level = %d, nx = %d" % (level, fP.nx) print " before G-S, residual L2 norm = %g" % \ (error(fP.imin, fP.imax, fP.dx, self.residual(level))) # smooth self.smooth(level, self.nsmooth) if self.verbose: print " after G-S, residual L2 norm = %g\n" % \ (error(fP.imin, fP.imax, fP.dx, self.residual(level))) level += 1 # compute the error with respect to the previous solution solnP = self.grids[self.nlevels-1] diff = (solnP.getVar("v") - oldSolution)/(solnP.getVar("v") + self.small) err = error(solnP.imin, solnP.imax, solnP.dx, diff) oldSolution = solnP.getVar("v").copy() if (err < rtol): converged = 1 self.numCycles = cycle self.error = err if self.verbose: print "cycle %d: relative error = %g\n" % (cycle, err) cycle += 1