r"""
This multigrid solver is build from multigrid/MG.py
and implements a more general solver for an equation of the form
.. math::
\alpha \phi + \nabla\cdot { \beta \nabla \phi } + \gamma \cdot \nabla \phi = f
where alpha, beta, and gamma are defined on the same grid as phi.
These should all come in as cell-centered quantities. The solver
will put beta on edges. Note that gamma is a vector here, with
x- and y-components.
A cell-centered discretization for phi 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 GeneralMG2d(MG.CellCenterMG2d):
r"""
this is a multigrid solver that supports our general elliptic
equation.
we need to accept a coefficient ``CellCenterData2d`` object with
fields defined for ``alpha``, ``beta``, ``gamma_x``, and ``gamma_y`` on the
fine level.
We then restrict this data through the MG hierarchy (and
average beta to the edges).
we need a ``new compute_residual()`` and ``smooth()`` function, that
understands these 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",
xl_BC=None, xr_BC=None,
yl_BC=None, yr_BC=None,
nsmooth=10, nsmooth_bottom=50,
verbose=0,
coeffs=None,
true_function=None, vis=0, vis_title=""):
"""
here, coeffs is a CCData2d object
"""
# we'll keep a list of the beta coefficients averaged to the
# interfaces on each level -- note: these will already be
# scaled by 1/dx**2
self.beta_edge = []
# initialize the MG object with the auxiliary fields
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,
xl_BC=xl_BC, xr_BC=xr_BC,
yl_BC=yl_BC, yr_BC=yr_BC,
alpha=0.0, beta=0.0,
nsmooth=nsmooth, nsmooth_bottom=nsmooth_bottom,
verbose=verbose,
aux_field=["alpha", "beta", "gamma_x", "gamma_y"],
aux_bc=[coeffs.BCs["alpha"], coeffs.BCs["beta"],
coeffs.BCs["gamma_x"], coeffs.BCs["gamma_y"]],
true_function=true_function, vis=vis,
vis_title=vis_title)
# the coefficients come in a dictionary. 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.
for c in ["alpha", "beta", "gamma_x", "gamma_y"]:
v = self.grids[self.nlevels-1].get_var(c)
v.v()[:, :] = coeffs.get_var(c).v()
self.grids[self.nlevels-1].fill_BC(c)
n = self.nlevels-2
while n >= 0:
f_patch = self.grids[n+1]
c_patch = self.grids[n]
coeffs_c = c_patch.get_var(c)
coeffs_c.v()[:, :] = f_patch.restrict(c).v()
self.grids[n].fill_BC(c)
n -= 1
# put the beta coefficients on edges
beta = self.grids[self.nlevels-1].get_var("beta")
self.beta_edge.insert(0, ec.EdgeCoeffs(self.grids[self.nlevels-1].grid, beta))
n = self.nlevels-2
while n >= 0:
self.beta_edge.insert(0, self.beta_edge[0].restrict())
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")
myg = self.grids[level].grid
dx = myg.dx
dy = myg.dy
self.grids[level].fill_BC("v")
alpha = self.grids[level].get_var("alpha")
gamma_x = 0.5*self.grids[level].get_var("gamma_x")/dx
gamma_y = 0.5*self.grids[level].get_var("gamma_y")/dy
# these are already scaled by 1/dx**2 in the EdgeCoeffs
# construction
beta_x = self.beta_edge[level].x
beta_y = self.beta_edge[level].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 = (
alpha.ip_jp(ix, iy, s=2) -
beta_x.ip_jp(1+ix, iy, s=2) - beta_x.ip_jp(ix, iy, s=2) -
beta_y.ip_jp(ix, 1+iy, s=2) - beta_y.ip_jp(ix, iy, s=2))
v.ip_jp(ix, iy, s=2)[:, :] = (f.ip_jp(ix, iy, s=2) -
# (beta_{i+1/2,j} + gamma^x_{i,j}) phi_{i+1,j}
(beta_x.ip_jp(1+ix, iy, s=2) + gamma_x.ip_jp(ix, iy, s=2)) * v.ip_jp(1+ix, iy, s=2) -
# (beta_{i-1/2,j} - gamma^x_{i,j}) phi_{i-1,j}
(beta_x.ip_jp(ix, iy, s=2) - gamma_x.ip_jp(ix, iy, s=2)) * v.ip_jp(-1+ix, iy, s=2) -
# (beta_{i,j+1/2} + gamma^y_{i,j}) phi_{i,j+1}
(beta_y.ip_jp(ix, 1+iy, s=2) + gamma_y.ip_jp(ix, iy, s=2)) * v.ip_jp(ix, 1+iy, s=2) -
# (beta_{i,j-1/2} - gamma^y_{i,j}) phi_{i,j-1}
(beta_y.ip_jp(ix, iy, s=2) - gamma_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")
myg = self.grids[level].grid
dx = myg.dx
dy = myg.dy
alpha = self.grids[level].get_var("alpha")
gamma_x = 0.5*self.grids[level].get_var("gamma_x")/dx
gamma_y = 0.5*self.grids[level].get_var("gamma_y")/dy
# these already have a 1/dx**2 scaling in them
beta_x = self.beta_edge[level].x
beta_y = self.beta_edge[level].y
# compute the residual
# r = f - L_eta phi
L_eta_phi = (
# alpha_{i,j} phi_{i,j}
alpha.v()*v.v() +
# beta_{i+1/2,j} (phi_{i+1,j} - phi_{i,j})
beta_x.ip(1)*(v.ip(1) - v.v()) - \
# beta_{i-1/2,j} (phi_{i,j} - phi_{i-1,j})
beta_x.v()*(v.v() - v.ip(-1)) + \
# beta_{i,j+1/2} (phi_{i,j+1} - phi_{i,j})
beta_y.jp(1)*(v.jp(1) - v.v()) - \
# beta_{i,j-1/2} (phi_{i,j} - phi_{i,j-1})
beta_y.v()*(v.v() - v.jp(-1)) + \
# gamma^x_{i,j} (phi_{i+1,j} - phi_{i-1,j})
gamma_x.v()*(v.ip(1) - v.ip(-1)) + \
# gamma^y_{i,j} (phi_{i,j+1} - phi_{i,j-1})
gamma_y.v()*(v.jp(1) - v.jp(-1)))
r.v()[:, :] = f.v() - L_eta_phi