# Source code for fe_utils.solvers.helmholtz

"""Solve a model helmholtz problem 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 fe_utils import *
import numpy as np
from numpy import cos, 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 Helmholtz problem given
the function space in which to solve and the right hand side
function."""

raise NotImplementedError

# Create an appropriate (complete) quadrature rule.

# Tabulate the basis functions and their gradients at the quadrature points.

# Create the left hand side matrix and right hand side vector.
# This creates a sparse matrix because creating a dense one may
# well run your machine out of memory!
A = sp.lil_matrix((fs.node_count, fs.node_count))
l = np.zeros(fs.node_count)

# Now loop over all the cells and assemble A and l

return A, l

[docs]def solve_helmholtz(degree, resolution, analytic=False, return_error=False):
"""Solve a model Helmholtz 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.

# If the analytic answer has been requested then bail out now.
if analytic:

# Create the right hand side function and populate it with the
# correct values.
f = Function(fs)
f.interpolate(lambda x: ((16*pi**2 + 1)*(x[1] - 1)**2*x[1]**2 - 12*x[1]**2 + 12*x[1] - 2) *
cos(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.

if return_error:

# Return the solution and the error in the solution.
return u, error

if __name__ == "__main__":

parser = ArgumentParser(
description="""Solve a Helmholtz problem on the unit square.""")
help="Plot the analytic solution instead of solving the finite element problem.")
help="Plot the error instead of the solution.")