Helmholtz equation is given as
$$ \nabla^2 u + \alpha u = f. $$In this notebook we will solve this equation (without boundary conditions) on the surface of a torus, using curvilinear coordinates. The surface of the torus is parametrized by
$$ \begin{align} x(\theta, \phi) &= (R + r \cos \theta) \cos \phi \\ y(\theta, \phi) &= (R + r \cos \theta) \sin \phi \\ z(\theta, \phi) &= r \sin \theta \end{align} $$where $\theta, \phi$ are angles which make a full circle, so that their values start and end at the same point, $R$ is the distance from the center of the tube to the center of the torus, $r$ is the radius of the tube. Note that $\theta$ is the angle in the small circle (around its center), whereas $\phi$ is the angle of the large circle, around origo.
We start the implementation by importing necessary functionality from shenfun and sympy and then defining the coordinates of the surface of the torus. Note that rv represents the position vector spanning the surface.
In [ ]:
from shenfun import *
from shenfun.la import SolverGeneric1ND
import sympy as sp
from IPython.display import Math
N = 20
R = 3
r = 1
theta, phi = psi = sp.symbols('x,y', real=True, positive=True)
rv = ((R + r*sp.cos(theta))*sp.cos(phi), (R + r*sp.cos(theta))*sp.sin(phi), r*sp.sin(theta))
Now create necessary bases and function spaces. Due to the geometry of the problem, the solution will be periodic in both $\theta$ and $\phi$ directions, and we choose Fourier basis functions. The basis for the $\phi$-direction can be either real to complex or complex to complex, depending on the type of the solution. Here we assume a real solution
In [ ]:
B1 = FunctionSpace(N, 'F', dtype='D', domain=(0, 2*np.pi))
B2 = FunctionSpace(N, 'F', dtype='d', domain=(0, 2*np.pi))
T = TensorProductSpace(comm, (B1, B2), coordinates=(psi, rv, sp.Q.positive(r*sp.cos(theta)+R)))
V = VectorTensorProductSpace(T)
u = TrialFunction(T)
v = TestFunction(T)
T.coors.sg
Note the assumption sp.Q.positive(r*sp.cos(theta)+R)), which is there to help sympy when computing basis vectors and scaling factors. It is not completely necessary, but try to omit it and look at what happens to the expanded Helmholtz equation below.
In [ ]:
alpha = 1
du = div(grad(u))+alpha*u
g = sp.Symbol('g', real=True, positive=True) # The Jacobian**2 (T.coors.sg**2)
replace = [(sp.cos(theta)+R, sp.sqrt(g))] # to simplify the look
Math(du.tolatex(symbol_names={r: 'r', theta: '\\theta', phi: '\\phi'}, replace=replace))
We now create a manufactured solution that satisfies periodicity and compute the right hand side $f$.
In [ ]:
ue = sp.sin(theta*2)*sp.cos(4*phi)
f = (div(grad(u))+alpha*u).tosympy(basis=ue, psi=psi)
fj = Array(T, buffer=f)
f_hat = Function(T)
f_hat = inner(v, fj, output_array=f_hat)
Assemble coefficient matrix and solve problem. Note that the tensorproduct matrices along axis 0 can be non-diagonal due to the measure $\cos \theta + 3$. The matrices along the second axis will all be diagonal, so we can choose to use SolverGeneric1ND
In [ ]:
#mats = inner(v, (div(grad(u))+alpha*u))
mats = inner(grad(v), -grad(u)) + inner(v, alpha*u)
u_hat = Function(T)
sol = SolverGeneric1ND(mats)
u_hat = sol(f_hat, u_hat)
uj = u_hat.backward()
uq = Array(T, buffer=ue)
print('Error =', np.linalg.norm(uj-uq))
Finally, just plot the solution using mayavi. Wrap the periodic directions around and refine to get a nicer plot.
In [ ]:
from mayavi import mlab
u_hat2 = u_hat.refine([N*2, N*2])
ur = u_hat2.backward()
xx, yy, zz = u_hat2.function_space().local_curvilinear_mesh()
xx = np.hstack([xx, xx[:, 0][:, None]])
yy = np.hstack([yy, yy[:, 0][:, None]])
zz = np.hstack([zz, zz[:, 0][:, None]])
ur = np.hstack([ur, ur[:, 0][:, None]])
xx = np.vstack([xx, xx[0]])
yy = np.vstack([yy, yy[0]])
zz = np.vstack([zz, zz[0]])
ur = np.vstack([ur, ur[0]])
mlab.init_notebook('x3d', 400, 400)
mlab.figure(bgcolor=(1, 1, 1))
m = mlab.mesh(xx, yy, zz, scalars=ur.real, colormap='jet')
mlab.savefig('torus.tiff')
m