Source code for fe_utils.solvers.poisson

"""Solve a model poisson problem with Dirichlet boundary conditions
using the finite element method.

If run as a script, the result is plotted. This file can also be
imported as a module and convergence tests run on the solver.
from __future__ import division
from fe_utils import *
import numpy as np
from numpy import sin, pi
import scipy.sparse as sp
import scipy.sparse.linalg as splinalg
from argparse import ArgumentParser

[docs]def assemble(fs, f): """Assemble the finite element system for the Poisson problem given the function space in which to solve and the right hand side function.""" raise NotImplementedError
[docs]def boundary_nodes(fs): """Find the list of boundary nodes in fs. This is a unit-square-specific solution. A more elegant solution would employ the mesh topology and numbering. """ eps = 1.e-10 f = Function(fs) def on_boundary(x): """Return 1 if on the boundary, 0. otherwise.""" if x[0] < eps or x[0] > 1 - eps or x[1] < eps or x[1] > 1 - eps: return 1. else: return 0. f.interpolate(on_boundary) return np.flatnonzero(f.values)
[docs]def solve_poisson(degree, resolution, analytic=False, return_error=False): """Solve a model Poisson problem on a unit square mesh with ``resolution`` elements in each direction, using equispaced Lagrange elements of degree ``degree``.""" # Set up the mesh, finite element and function space required. mesh = UnitSquareMesh(resolution, resolution) fe = LagrangeElement(mesh.cell, degree) fs = FunctionSpace(mesh, fe) # Create a function to hold the analytic solution for comparison purposes. analytic_answer = Function(fs) analytic_answer.interpolate(lambda x: sin(4*pi*x[0])*x[1]**2*(1.-x[1])**2) # If the analytic answer has been requested then bail out now. if analytic: return analytic_answer, 0.0 # Create the right hand side function and populate it with the # correct values. f = Function(fs) f.interpolate(lambda x: (16*pi**2*(x[1] - 1)**2*x[1]**2 - 2*(x[1] - 1)**2 - 8*(x[1] - 1)*x[1] - 2*x[1]**2) * sin(4*pi*x[0])) # Assemble the finite element system. A, l = assemble(fs, f) # Create the function to hold the solution. u = Function(fs) # Cast the matrix to a sparse format and use a sparse solver for # the linear system. This is vastly faster than the dense # alternative. A = sp.csr_matrix(A) u.values[:] = splinalg.spsolve(A, l) # Compute the L^2 error in the solution for testing purposes. error = errornorm(analytic_answer, u) if return_error: u.values -= analytic_answer.values # Return the solution and the error in the solution. return u, error
if __name__ == "__main__": parser = ArgumentParser( description="""Solve a Poisson problem on the unit square.""") parser.add_argument("--analytic", action="store_true", help="Plot the analytic solution instead of solving the finite element problem.") parser.add_argument("--error", action="store_true", help="Plot the error instead of the solution.") parser.add_argument("resolution", type=int, nargs=1, help="The number of cells in each direction on the mesh.") parser.add_argument("degree", type=int, nargs=1, help="The degree of the polynomial basis for the function space.") args = parser.parse_args() resolution = args.resolution[0] degree =[0] analytic = args.analytic plot_error = args.error u, error = solve_poisson(degree, resolution, analytic, plot_error) u.plot()