r"""
This multigrid solver is build from multigrid/MG.py and implements
a variable coefficient solver for an equation of the form
.. math::
\nabla\cdot { \eta \nabla \phi } = f
where :math:`\eta` is defined on the same grid as :math:`\phi`.
A cell-centered discretization is used throughout.
"""
import matplotlib.pyplot as plt
import numpy as np
import pyro.multigrid.edge_coeffs as ec
from pyro.multigrid import MG
np.set_printoptions(precision=3, linewidth=128)
[docs]
class VarCoeffCCMG2d(MG.CellCenterMG2d):
r"""
this is a multigrid solver that supports variable coefficients
we need to accept a coefficient array, coeffs, defined at each
level. We can do this at the fine level and restrict it
down the MG grids once.
we need a new ``compute_residual()`` and ``smooth()`` function, that
understands coeffs.
"""
def __init__(self, nx, ny, xmin=0.0, xmax=1.0, ymin=0.0, ymax=1.0,
xl_BC_type="dirichlet", xr_BC_type="dirichlet",
yl_BC_type="dirichlet", yr_BC_type="dirichlet",
nsmooth=10, nsmooth_bottom=50,
verbose=0,
coeffs=None, coeffs_bc=None,
true_function=None, vis=0, vis_title=""):
# we'll keep a list of the coefficients averaged to the interfaces
# on each level -- note: this will already be scaled by 1/dx**2
self.edge_coeffs = []
# initialize the MG object with the auxiliary "coeffs" field
MG.CellCenterMG2d.__init__(self, nx, ny, ng=1,
xmin=xmin, xmax=xmax, ymin=ymin, ymax=ymax,
xl_BC_type=xl_BC_type, xr_BC_type=xr_BC_type,
yl_BC_type=yl_BC_type, yr_BC_type=yr_BC_type,
alpha=0.0, beta=0.0,
nsmooth=nsmooth, nsmooth_bottom=nsmooth_bottom,
verbose=verbose,
aux_field=["coeffs"], aux_bc=[coeffs_bc],
true_function=true_function, vis=vis,
vis_title=vis_title)
# set the coefficients and restrict them down the hierarchy
# we only need to do this once. We need to hold the original
# coeffs in our grid so we can do a ghost cell fill.
c = self.grids[self.nlevels-1].get_var("coeffs")
if c.g.nx != nx or c.g.ny != ny:
raise IndexError("coefficient array not the same size as multigrid problem")
c.v()[:, :] = coeffs.v().copy()
self.grids[self.nlevels-1].fill_BC("coeffs")
# put the coefficients on edges
self.edge_coeffs.insert(0, ec.EdgeCoeffs(self.grids[self.nlevels-1].grid, c))
n = self.nlevels-2
while n >= 0:
# create the edge coefficients on level n by restricting from the
# finer grid
f_patch = self.grids[n+1]
c_patch = self.grids[n]
coeffs_c = c_patch.get_var("coeffs")
coeffs_c.v()[:, :] = f_patch.restrict("coeffs").v()
self.grids[n].fill_BC("coeffs")
# put the coefficients on edges
self.edge_coeffs.insert(0, self.edge_coeffs[0].restrict()) # _EdgeCoeffs(self.grids[n].grid, coeffs_c))
# if we are periodic, then we should force the edge coefficients
# to be periodic
# if self.grids[n].BCs["coeffs"].xlb == "periodic":
# self.edge_coeffs[0].x[self.grids[n].grid.ihi+1,:] = \
# self.edge_coeffs[0].x[self.grids[n].grid.ilo,:]
# if self.grids[n].BCs["coeffs"].ylb == "periodic":
# self.edge_coeffs[0].y[:,self.grids[n].grid.jhi+1] = \
# self.edge_coeffs[0].y[:,self.grids[n].grid.jlo]
n -= 1
[docs]
def smooth(self, level, nsmooth):
"""
Use red-black Gauss-Seidel iterations to smooth the solution
at a given level. This is used at each stage of the V-cycle
(up and down) in the MG solution, but it can also be called
directly to solve the elliptic problem (although it will take
many more iterations).
Parameters
----------
level : int
The level in the MG hierarchy to smooth the solution
nsmooth : int
The number of r-b Gauss-Seidel smoothing iterations to perform
"""
v = self.grids[level].get_var("v")
f = self.grids[level].get_var("f")
self.grids[level].fill_BC("v")
eta_x = self.edge_coeffs[level].x
eta_y = self.edge_coeffs[level].y
# print( "min/max c: {}, {}".format(np.min(c), np.max(c)))
# print( "min/max eta_x: {}, {}".format(np.min(eta_x), np.max(eta_x)))
# print( "min/max eta_y: {}, {}".format(np.min(eta_y), np.max(eta_y)))
# do red-black G-S
for _ in range(nsmooth):
# do the red black updating in four decoupled groups
#
#
# | | |
# --+-------+-------+--
# | | |
# | 4 | 3 |
# | | |
# --+-------+-------+--
# | | |
# jlo | 1 | 2 |
# | | |
# --+-------+-------+--
# | ilo | |
#
# groups 1 and 3 are done together, then we need to
# fill ghost cells, and then groups 2 and 4
for n, (ix, iy) in enumerate([(0, 0), (1, 1), (1, 0), (0, 1)]):
denom = (eta_x.ip_jp(1+ix, iy, s=2) + eta_x.ip_jp(ix, iy, s=2) +
eta_y.ip_jp(ix, 1+iy, s=2) + eta_y.ip_jp(ix, iy, s=2))
v.ip_jp(ix, iy, s=2)[:, :] = (-f.ip_jp(ix, iy, s=2) +
# eta_{i+1/2,j} phi_{i+1,j}
eta_x.ip_jp(1+ix, iy, s=2) * v.ip_jp(1+ix, iy, s=2) +
# eta_{i-1/2,j} phi_{i-1,j}
eta_x.ip_jp(ix, iy, s=2) * v.ip_jp(-1+ix, iy, s=2) +
# eta_{i,j+1/2} phi_{i,j+1}
eta_y.ip_jp(ix, 1+iy, s=2) * v.ip_jp(ix, 1+iy, s=2) +
# eta_{i,j-1/2} phi_{i,j-1}
eta_y.ip_jp(ix, iy, s=2) * v.ip_jp(ix, -1+iy, s=2)) / denom
if n in (1, 3):
self.grids[level].fill_BC("v")
if self.vis == 1:
plt.clf()
plt.subplot(221)
self._draw_solution()
plt.subplot(222)
self._draw_V()
plt.subplot(223)
self._draw_main_solution()
plt.subplot(224)
self._draw_main_error()
plt.suptitle(self.vis_title, fontsize=18)
plt.draw()
plt.savefig("mg_%4.4d.png" % (self.frame))
self.frame += 1
def _compute_residual(self, level):
""" compute the residual and store it in the r variable"""
v = self.grids[level].get_var("v")
f = self.grids[level].get_var("f")
r = self.grids[level].get_var("r")
eta_x = self.edge_coeffs[level].x
eta_y = self.edge_coeffs[level].y
# compute the residual
# r = f - L_eta phi
L_eta_phi = (
# eta_{i+1/2,j} (phi_{i+1,j} - phi_{i,j})
eta_x.ip(1)*(v.ip(1) - v.v()) - \
# eta_{i-1/2,j} (phi_{i,j} - phi_{i-1,j})
eta_x.v()*(v.v() - v.ip(-1)) + \
# eta_{i,j+1/2} (phi_{i,j+1} - phi_{i,j})
eta_y.jp(1)*(v.jp(1) - v.v()) - \
# eta_{i,j-1/2} (phi_{i,j} - phi_{i,j-1})
eta_y.v()*(v.v() - v.jp(-1)))
r.v()[:, :] = f.v() - L_eta_phi