fourier_Helmholtz3D.py

r"""
Solve Helmholtz equation on (0, 2pi)x(0, 2pi)x(0, 2pi) with periodic bcs

.. math::

    \nabla^2 u + u = f,

Use Fourier basis and find :math:`u` in :math:`V^3` such that

.. math::

    (v, \nabla^2 u + u) = (v, f), \quad \forall v \in V^3

where V is the Fourier basis :math:`span{exp(1jkx)}_{k=-N/2}^{N/2-1}` and
:math:`V^3` is a tensorproductspace.

"""
import os
import numpy as np
from sympy import symbols, cos, sin
from mpi4py import MPI
from shenfun import inner, div, grad, TestFunction, TrialFunction, FunctionSpace, \
    TensorProductSpace, Array, Function, dx

comm = MPI.COMM_WORLD

# Use sympy to compute a rhs, given an analytical solution
x, y, z = symbols("x,y,z", real=True)
ue = cos(4*x) + sin(4*y) + sin(6*z)
fe = ue.diff(x, 2) + ue.diff(y, 2) + ue.diff(z, 2) + ue

# Size of discretization
N = 16

K0 = FunctionSpace(N, 'F', dtype='D')
K1 = FunctionSpace(N, 'F', dtype='D')
K2 = FunctionSpace(N, 'F', dtype='d')
T = TensorProductSpace(comm, (K0, K1, K2), slab=True)
u = TrialFunction(T)
v = TestFunction(T)

# Get f on quad points
fj = Array(T, buffer=fe)

# Compute right hand side
f_hat = Function(T)
f_hat = inner(v, fj, output_array=f_hat)

# Solve Poisson equation
A = inner(v, u+div(grad(u)))
f_hat = A.solve(f_hat)

uq = T.backward(f_hat, fast_transform=True)

uj = Array(T, buffer=ue)
print(np.sqrt(dx((uj-uq)**2)))
assert np.allclose(uj, uq)

if 'pytest' not in os.environ:
    import matplotlib.pyplot as plt
    plt.figure()
    X = T.local_mesh(True) # With broadcasting=True the shape of X is local_shape, even though the number of datapoints are still the same as in 1D
    plt.contourf(X[0][:, :, 0], X[1][:, :, 0], uq[:, :, 0])
    plt.colorbar()

    plt.figure()
    plt.contourf(X[0][:, :, 0], X[1][:, :, 0], uj[:, :, 0])
    plt.colorbar()

    plt.figure()
    plt.contourf(X[0][:, :, 0], X[1][:, :, 0], uq[:, :, 0]-uj[:, :, 0])
    plt.colorbar()
    plt.title('Error')
    #plt.show()