<jupyter_start><jupyter_text>Demo - Stokes equations  **<NAME>** (email: `<EMAIL>`), Department of Mathematics, University of Oslo.Date: **January 23, 2019****Summary.** The Stokes equations describe the flow of highly viscous fluids.This is a demonstration of how the Python module [shenfun](https://github.com/spectralDNS/shenfun) can be used to solve Stokesequations using a  mixed (coupled) basis in a 3D tensor product domain.We assume homogeneous Dirichlet boundary conditions in one directionand periodicity in the remaining two. The solver described runs with MPIwithout any further considerations required from the user.The solver assembles a block matrix with sparsity pattern as shown belowfor the Legendre basis. -->Coupled block matrix for Stokes equations. Stokes' equationsStokes' equations are given in strong form as $$\begin{align*}\nabla^2 \boldsymbol{u} - \nabla p &= \boldsymbol{f} \quad \text{in }  \Omega, \\ \nabla \cdot \boldsymbol{u} &= h \quad \text{in } \Omega  \\ \int_{\Omega} p dx &= 0\end{align*}$$ where $\boldsymbol{u}$ and $p$ are, respectively, thefluid velocity vector and pressure, and the domain$\Omega = [0, 2\pi]^2 \times [-1, 1]$. The flow is assumed periodicin $x$ and $y$-directions, whereas there is a no-slip homogeneous Dirichletboundary condition on $\boldsymbol{u}$ on the boundaries of the $z$-direction, i.e.,$\boldsymbol{u}(x, y, \pm 1) = (0, 0, 0)$. (Note that we can configure shenfun withnon-periodicity in any of the three directions. However, since we are tosolve linear algebraic systems in the non-periodic direction, there is a speedbenefit from having the nonperiodic direction last. This has to do with Numpyusing a C-style row-major storage of arrays by default.)The right hand side vector $\boldsymbol{f}(\boldsymbol{x})$ is an external applied body force.The right hand side $h$ is usually zero in the regular Stokes equations. Herewe include it because it will be nonzero in the verification, which is using themethod of manufactured solutions. Note that the final $\int_{\Omega} p dx = 0$is there because there is no Dirichlet boundary condition on the pressureand the system of equations would otherwise be ill conditioned.To solve Stokes' equations with the Galerkin method we need basisfunctions for both velocity and pressure. ADirichlet basis will be used for velocity, whereas there is no boundary restrictionon the pressure basis. For both three-dimensional bases we will use one basisfunction for the $x$-direction,$\mathcal{X}(x)$, one for the $y$-direction, $\mathcal{Y}(y)$, and one for the$z$-direction, $\mathcal{Z}(z)$. Andthen we create three-dimensional basis functions like $$\begin{equation}v(x, y, z) = \mathcal{X}(x) \mathcal{Y}(y) \mathcal{Z} (z).\label{_auto1} \tag{1}\end{equation}$$ The basis functions $\mathcal{X}(x)$ and $\mathcal{Y}(y)$ are chosen as Fourierexponentials, since these functions are periodic: $$\begin{equation}\mathcal{X}_l(x) = e^{\imath l x}, \forall \, l \in \boldsymbol{l}^{N_0}, \label{_auto2} \tag{2}\end{equation}$$ $$\begin{equation}  \mathcal{Y}_m(y) =  e^{\imath m y}, \forall \, m \in \boldsymbol{m}^{N_1},\label{_auto3} \tag{3}\end{equation}$$ where $\boldsymbol{l}^{N_0} = (-N_0/2, -N_0/2+1, \ldots, N_0/2-1)$ and$\boldsymbol{m}^{N_1} = (-N_1/2, -N_1/2+1, \ldots, N_1/2-1)$.The size of the discretized problem in real physical space is$\boldsymbol{N} = (N_0, N_1, N_2)$, i.e., there are $N_0 \cdot N_1 \cdot N_2$ quadrature pointsin total.The basis functions for $\mathcal{Z}(z)$ remain to be decided.For the velocity we need homogeneous Dirichlet boundary conditions, and for thiswe use composite Legendre or Chebyshev polynomials $$\begin{equation}\mathcal{Z}^0_n(z) = \phi_n(z) - \phi_{n+2}(z), \forall \, n \in \boldsymbol{n}^{N_2-2},\label{_auto4} \tag{4}\end{equation}$$ where $\phi_n$ is the n'th Legendre or Chebyshev polynomial of the first kind.$\boldsymbol{n}^{N_2-2} = (0, 1, \ldots, N_2-3)$, and the zero on $\mathcal{Z}^0$is there to indicate the zero value on the boundary.The pressure basis that comes with no restrictions for the boundary is alittle trickier. The reason for this has to do withinf-sup stability. The obvious choice of basis is the regular Legendre orChebyshev basis, which is denoted as $$\begin{equation}\mathcal{Z}_n(z) = \phi_n(z),  \forall \, n \in \boldsymbol{n}^{N_2}. \label{eq:Zn} \tag{5}\end{equation}$$ The problem is that for the natural choice of $n \in (0, 1, \ldots, N_2-1)$there is a nullspace and one degree of freedom remains unresolved. It turns outthat the proper choice for the pressure basis is simply ([5](eq:Zn)) for$n \in \boldsymbol{n}^{N_2-2}$. (Also remember that we have to fix $\int_{\Omega} p dx = 0$.)With given basis functions we obtain the spaces $$\begin{equation}V^{N_0} = \text{span}\{ \mathcal{X}_l \}_{l\in\boldsymbol{l}^{N_0}}, \label{_auto5} \tag{6}\end{equation}$$ $$\begin{equation}  V^{N_1} = \text{span}\{ \mathcal{Y}_m \}_{m\in\boldsymbol{m}^{N_1}}, \label{_auto6} \tag{7}\end{equation}$$ $$\begin{equation}  V^{N_2} = \text{span}\{ \mathcal{Z}_n \}_{n\in\boldsymbol{n}^{N_2-2}}, \label{_auto7} \tag{8}\end{equation}$$ $$\begin{equation}  V_0^{N_2} = \text{span}\{ \mathcal{Z}^0_n \}_{n\in\boldsymbol{n}^{N_2-2}},\label{_auto8} \tag{9}\end{equation}$$ and from these we create two different tensor product spaces $$\begin{equation}W_0^{\boldsymbol{N}}(\boldsymbol{x}) = V^{N_0}(x) \otimes V^{N_1}(y) \otimes V_0^{N_2}(z), \label{_auto9} \tag{10}\end{equation}$$ $$\begin{equation}  W^{\boldsymbol{N}}(\boldsymbol{x}) = V^{N_0}(x) \otimes V^{N_1}(y) \otimes V^{N_2}(z).\label{_auto10} \tag{11}\end{equation}$$ The velocity vector is using a mixed basis, such that we will look forsolutions $\boldsymbol{u} \in [W_0^{\boldsymbol{N}}]^3 \, (=W_0^{\boldsymbol{N}} \times W_0^{\boldsymbol{N}} \times W_0^{\boldsymbol{N}})$,whereas we look for the pressure$p \in W^{\boldsymbol{N}}$. We now formulate a variational problem using the Galerkin method: Find$\boldsymbol{u} \in [W_0^{\boldsymbol{N}}]^3$ and $p \in W^{\boldsymbol{N}}$ such that $$\begin{equation}\int_{\Omega} (\nabla^2 \boldsymbol{u} - \nabla p ) \cdot \overline{\boldsymbol{v}} \, dx_w = \int_{\Omega} \boldsymbol{f} \cdot \overline{\boldsymbol{v}}\, dx_w \quad\forall \boldsymbol{v} \, \in \, [W_0^{\boldsymbol{N}}]^3, \label{eq:varform} \tag{12} \end{equation}$$ $$\begin{equation}  \int_{\Omega} \nabla \cdot \boldsymbol{u} \, \overline{q} \, dx_w = \int_{\Omega} h \overline{q} \, dx_w \quad\forall q \, \in \, W^{\boldsymbol{N}}.\label{_auto11} \tag{13}\end{equation}$$ Here $dx_w=w_xdxw_ydyw_zdz$ represents a weighted measure, with weights $w_x(x), w_y(y), w_z(z)$.Note that it is only Chebyshev polynomials thatmake use of a non-constant weight $w_x=1/\sqrt{1-x^2}$. The Fourier weights are $w_y=w_z=1/(2\pi)$and the Legendre weight is $w_x=1$.The overline in $\boldsymbol{\overline{v}}$ and $\overline{q}$ represents a complex conjugate, which is needed here becausethe Fourier exponentials are complex functions. Mixed variational formSince we are to solve for $\boldsymbol{u}$ and $p$ at the same time, we formulate amixed (coupled) problem: find $(\boldsymbol{u}, p) \in [W_0^{\boldsymbol{N}}]^3 \times W^{\boldsymbol{N}}$such that $$\begin{equation}a((\boldsymbol{u}, p), (\boldsymbol{v}, q)) = L((\boldsymbol{v}, q)) \quad \forall (\boldsymbol{v}, q) \in [W_0^{\boldsymbol{N}}]^3 \times W^{\boldsymbol{N}},\label{_auto12} \tag{14}\end{equation}$$ where bilinear ($a$) and linear ($L$) forms are given as $$\begin{equation}    a((\boldsymbol{u}, p), (\boldsymbol{v}, q)) = \int_{\Omega} (\nabla^2 \boldsymbol{u} - \nabla p) \cdot \overline{\boldsymbol{v}} \, dx_w + \int_{\Omega} \nabla \cdot \boldsymbol{u} \, \overline{q} \, dx_w, \label{_auto13} \tag{15}\end{equation}$$ $$\begin{equation}      L((\boldsymbol{v}, q)) = \int_{\Omega} \boldsymbol{f} \cdot \overline{\boldsymbol{v}}\, dx_w + \int_{\Omega} h \overline{q} \, dx_w.\label{_auto14} \tag{16}\end{equation}$$ Note that the bilinear form will assemble to block matrices, whereas the right hand sidelinear form will assemble to block vectors. Implementation PreambleWe will solve the Stokes equations using the [shenfun](https://github.com/spectralDNS/shenfun) Python module. The first thing neededis then to import some of this module's functionalityplus some other helper modules, like [Numpy](https://numpy.org) and [Sympy](https://sympy.org):<jupyter_code>import os
import sys
import numpy as np
from sympy import symbols, sin, cos
from shenfun import *<jupyter_output><empty_output><jupyter_text>We use `Sympy` for the manufactured solution and `Numpy` for testing. Manufactured solutionThe exact solutions $\boldsymbol{u}_e(\boldsymbol{x})$ and $p(\boldsymbol{x})$ are chosen to satisfy boundaryconditions, and the right hand sides $\boldsymbol{f}(\boldsymbol{x})$ and $h(\boldsymbol{x})$ are thencomputed exactly using `Sympy`. These exact right hand sides will then be used tocompute a numerical solution that can be verified against the manufacturedsolution. The chosen solution with computed right hand sides are:<jupyter_code>x, y, z = symbols('x,y,z')
uex = sin(2*y)*(1-z**2)
uey = sin(2*x)*(1-z**2)
uez = sin(2*z)*(1-z**2)
pe = -0.1*sin(2*x)*cos(4*y)
fx = uex.diff(x, 2) + uex.diff(y, 2) + uex.diff(z, 2) - pe.diff(x, 1)
fy = uey.diff(x, 2) + uey.diff(y, 2) + uey.diff(z, 2) - pe.diff(y, 1)
fz = uez.diff(x, 2) + uez.diff(y, 2) + uez.diff(z, 2) - pe.diff(z, 1)
h = uex.diff(x, 1) + uey.diff(y, 1) + uez.diff(z, 1)<jupyter_output><empty_output><jupyter_text>Tensor product spacesOne-dimensional spaces are created using the [FunctionSpace()](https://shenfun.readthedocs.io/en/latest/shenfun.forms.htmlshenfun.forms.arguments.FunctionSpace) function. A choice ofpolynomials between Legendre or Chebyshev can be made, and the sizeof the domain is given<jupyter_code>N = (20, 20, 20)
family = 'Legendre'
K0 = FunctionSpace(N[0], 'Fourier', dtype='D', domain=(0, 2*np.pi))
K1 = FunctionSpace(N[1], 'Fourier', dtype='d', domain=(0, 2*np.pi))
SD = FunctionSpace(N[2], family, bc=(0, 0))
ST = FunctionSpace(N[2], family)<jupyter_output><empty_output><jupyter_text>Next the one-dimensional spaces are used to create two tensor product spaces Q = $W^{\boldsymbol{N}}$and TD = $W_0^{\boldsymbol{N}}$, one vector V = $[W_0^{\boldsymbol{N}}]^3$ and one mixedspace  VQ = V $\times$ Q.<jupyter_code>TD = TensorProductSpace(comm, (K0, K1, SD), axes=(2, 0, 1))
Q = TensorProductSpace(comm, (K0, K1, ST), axes=(2, 0, 1))
V = VectorSpace(TD)
VQ = CompositeSpace([V, Q])<jupyter_output><empty_output><jupyter_text>Note that we choose to transform axes in the order $1, 0, 2$. This is to ensurethat the fully transformed arrays are aligned in the non-periodic direction 2.And we need the arrays aligned in this direction, because this is the onlydirection where there are tensor product matrices that are non-diagonal. AllFourier matrices are, naturally, diagonal.Test- and trialfunctions are created much like in a regular, non-mixed,formulation. However, one has to create one test- and trialfunction forthe mixed space, and then split them up afterwards<jupyter_code>up = TrialFunction(VQ)
vq = TestFunction(VQ)
u, p = up
v, q = vq<jupyter_output><empty_output><jupyter_text>With the basisfunctions in place we may assemble the different blocks of thefinal coefficient matrix. Since Legendre is using a constant weight function,the equations may also be integrated by parts to obtain a symmetric system:<jupyter_code>if family.lower() == 'chebyshev':
    A = inner(v, div(grad(u)))
    G = inner(v, -grad(p))
else:
    A = inner(grad(v), -grad(u))
    G = inner(div(v), p)
D = inner(q, div(u))<jupyter_output><empty_output><jupyter_text>The assembled subsystems `A, G` and `D` are lists containg the different blocks ofthe complete, coupled matrix. `A` actually contains 6tensor product matrices of type [TPMatrix](https://shenfun.readthedocs.io/en/latest/shenfun.htmlshenfun.matrixbase.TPMatrix). The first twomatrices are for vector component zero of the test function `v[0]` andtrial function `u[0]`, thematrices 2 and 3 are for components 1 and the last two are for components2. The first two matrices are as such for ```Python        A[0:2] = inner(v[0], div(grad(u[0])))``` Breaking it down this inner product is mathematically $$\begin{equation}\label{eq:partialeq1} \tag{17}\int_{\Omega} \boldsymbol{v}[0] \left(\frac{\partial^2 \boldsymbol{u}[0]}{\partial x^2} + \frac{\partial^2 \boldsymbol{u}[0]}{\partial y^2} + \frac{\partial^2 \boldsymbol{u}[0]}{\partial z^2}\right) w_x dx w_y dy w_z dz.\end{equation}$$ If we now use test function $\boldsymbol{v}[0]$ $$\begin{equation}\boldsymbol{v}[0]_{lmn} = \mathcal{X}_l \mathcal{Y}_m \mathcal{Z}_n,\label{_auto15} \tag{18}\end{equation}$$ and trialfunction $$\begin{equation}\boldsymbol{u}[0]_{pqr} = \sum_{p} \sum_{q} \sum_{r} \hat{\boldsymbol{u}}[0]_{pqr} \mathcal{X}_p \mathcal{Y}_q \mathcal{Z}_r,\label{_auto16} \tag{19}\end{equation}$$ where $\hat{\boldsymbol{u}}$ are the unknown degrees of freedom, and then insert these functionsinto ([17](eq:partialeq1)), then we obtain afterperforming some exact evaluations over the periodic directions $$\begin{equation} \Big( \underbrace{-\left(l^2 \delta_{lp} + m^2 \delta_{mq} \right) \int_{-1}^{1} \mathcal{Z}_r(z) \mathcal{Z}_n(z) w_z dz}_{A[0]} + \underbrace{\delta_{lp} \delta_{mq} \int_{-1}^{1} \frac{\partial^2 \mathcal{Z}_r(z)}{\partial z^2} \mathcal{Z}_n(z) w_z dz}_{A[1]} \Big) \hat{\boldsymbol{u}}[0]_{pqr},\label{_auto17} \tag{20}\end{equation}$$ Similarly for components 1 and 2 of the test and trial vectors, leading to 6 tensorproduct matrices in total for `A`. Similarly, we get three components of `G`and  three of `D`.Eliminating the Fourier diagonal matrices, we are left with block matrices like $$H(l, m) =  \begin{bmatrix}      A[0]+A[1] & 0 & 0 & G[0] \\       0 & A[2]+A[3] & 0 & G[1] \\       0 & 0 &  A[4]+A[5] & G[2] \\       D[0] & D[1] & D[2] & 0  \end{bmatrix}$$ Note that there will be one large block matrix $H(l, m)$ for each Fourierwavenumber combination $(l, m)$. To solve the problem in the end we will need toloop over these wavenumbers and solve the assembled linear systems one by one.An example of the block matrix, for $l=m=5$ and $\boldsymbol{N}=(20, 20, 20)$ is givenin Fig. [fig:BlockMat](fig:BlockMat).In the end we create a block matrix through<jupyter_code>M = BlockMatrix(A+G+D)<jupyter_output><empty_output><jupyter_text>The right hand side can easily be assembled since we have alreadydefined the functions $\boldsymbol{f}$ and $h$, see Sec. [Manufactured solution](sec:mansol)<jupyter_code># Get mesh (quadrature points)
X = TD.local_mesh(True)

# Get f and h on quad points
fh = Array(VQ, buffer=(fx, fy, fz, h))
f_, h_ = fh

# Compute inner products
fh_hat = Function(VQ)
f_hat, h_hat = fh_hat
f_hat = inner(v, f_, output_array=f_hat)
h_hat = inner(q, h_, output_array=h_hat)<jupyter_output><empty_output><jupyter_text>In the end all that is left is to solve and compare withthe exact solution.<jupyter_code># Solve problem
up_hat = M.solve(fh_hat, constraints=((3, 0, 0), (3, N[2]-1, 0)))
up = up_hat.backward()
u_, p_ = up

# Exact solution
ux, uy, uz = Array(V, buffer=(uex, uey, uez))
pe = Array(Q, buffer=pe)

error = [comm.reduce(np.linalg.norm(ux-u_[0])),
         comm.reduce(np.linalg.norm(uy-u_[1])),
         comm.reduce(np.linalg.norm(uz-u_[2])),
         comm.reduce(np.linalg.norm(pe-p_))]
print(error)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1601218247796786, 'fraction_numerical': 0.019699326075686883, 'mean_word_length': 5.124206349206349, 'pattern_counts': {'":': 0, '<': 39, '<?xml version=': 0, '>': 40, 'https://': 6, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '526153', 'license': 'BSD-2-Clause', 'n_tokens_mistral': 5308, 'n_tokens_neox': 5058, 'n_words': 1691, 'path': 'binder/stokes.ipynb', 'repo_name': 'spectralDNS/shenfun', 'size': '15453'}