{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2.6 Stokes equation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Find $u \\in [H^1_D]^2$ and $p \\in L_2$ such that\n",
    "\n",
    "$$\n",
    "\\DeclareMathOperator{\\Div}{div}\n",
    "\\begin{array}{ccccll}\n",
    "\\int \\nabla u : \\nabla v & + & \\int \\Div v \\, p & = & \\int f v & \\forall \\, v \\\\\n",
    "\\int \\Div u \\, q &  &  & = & 0 & \\forall \\, q\n",
    "\\end{array}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define channel geometry and mesh it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ngsolve import *\n",
    "from ngsolve.webgui import Draw\n",
    "from netgen.occ import *\n",
    "\n",
    "shape = Rectangle(2,0.41).Circle(0.2,0.2,0.05).Reverse().Face()\n",
    "shape.edges.name=\"wall\"\n",
    "shape.edges.Min(X).name=\"inlet\"\n",
    "shape.edges.Max(X).name=\"outlet\"\n",
    "Draw (shape);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geo = OCCGeometry(shape, dim=2)\n",
    "mesh = Mesh(geo.GenerateMesh(maxh=0.05))\n",
    "mesh.Curve(3)\n",
    "Draw (mesh);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use Taylor Hood finite element pairing: Vector-valued continuous $P^2$ elements for velocity, and continuous $P^1$ for pressure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = VectorH1(mesh, order=2, dirichlet=\"wall|inlet|cyl\")\n",
    "Q = H1(mesh, order=1)\n",
    "X = V*Q"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setup bilinear-form for Stokes. We give names for all scalar field components. The divergence is constructed from partial derivatives of the velocity components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(u,p),(v,q) = X.TnT()\n",
    "\n",
    "stokes = InnerProduct(Grad(u), Grad(v))*dx + div(u)*q*dx + div(v)*p*dx\n",
    "a = BilinearForm(stokes).Assemble()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set inhomogeneous Dirichlet boundary condition only on inlet boundary:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gf = GridFunction(X)\n",
    "gfu, gfp = gf.components\n",
    "\n",
    "uin = CF ( (1.5*4*y*(0.41-y)/(0.41*0.41), 0) )\n",
    "gfu.Set(uin, definedon=mesh.Boundaries(\"inlet\"))\n",
    "Draw(gfu, mesh, min=0, max=2)\n",
    "SetVisualization(max=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solve equation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = -a.mat * gf.vec\n",
    "inv = a.mat.Inverse(freedofs=X.FreeDofs(), inverse=\"umfpack\")\n",
    "gf.vec.data += inv * res\n",
    "Draw(gfu, mesh);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing different velocity-pressure pairs\n",
    "Now we define a Stokes setup function to test different spaces:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def SolveStokes(X):\n",
    "    (u,p),(v,q) = X.TnT()\n",
    "\n",
    "    stokes = InnerProduct(Grad(u), Grad(v))*dx + div(u)*q*dx + div(v)*p*dx\n",
    "    a = BilinearForm(stokes).Assemble()    \n",
    "\n",
    "    gf = GridFunction(X)\n",
    "    gfu, gfp = gf.components\n",
    "\n",
    "    uin = CF ( (1.5*4*y*(0.41-y)/(0.41*0.41), 0) )\n",
    "    gfu.Set(uin, definedon=mesh.Boundaries(\"inlet\"))\n",
    "    \n",
    "    res = -a.mat * gf.vec\n",
    "    inv = a.mat.Inverse(freedofs=X.FreeDofs(), inverse=\"umfpack\")\n",
    "    gf.vec.data += inv * res\n",
    "     \n",
    "    Draw(gfu)\n",
    "    \n",
    "    return gfu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Higher order Taylor-Hood elements:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = VectorH1(mesh, order=4, dirichlet=\"wall|inlet|cyl\")\n",
    "Q = H1(mesh, order=3)\n",
    "X = V*Q\n",
    "\n",
    "gfu = SolveStokes(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With discontinuous pressure elements P2-P1 is unstable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [],
   "source": [
    "V = VectorH1(mesh, order=2, dirichlet=\"wall|inlet|cyl\")\n",
    "Q = L2(mesh, order=1)\n",
    "print (\"V.ndof =\", V.ndof, \", Q.ndof =\", Q.ndof)\n",
    "X = V*Q\n",
    "\n",
    "gfu = SolveStokes(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, P2-P1 elements work on sub-divided simplicial meshes using Alfeld splits:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh2 = Mesh(geo.GenerateMesh(maxh=0.05)).Curve(3)\n",
    "mesh2.SplitElements_Alfeld()\n",
    "V = VectorH1(mesh2, order=2, dirichlet=\"wall|inlet|cyl\")\n",
    "Q = L2(mesh2, order=1)\n",
    "print (\"V.ndof =\", V.ndof, \", Q.ndof =\", Q.ndof)\n",
    "X = V*Q\n",
    "gfu = SolveStokes(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$P^{2,+} \\times P^{1,dc}$ elements:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = VectorH1(mesh, order=2, dirichlet=\"wall|inlet|cyl\")\n",
    "V.SetOrder(TRIG,3)\n",
    "V.Update()\n",
    "print (V.ndof)\n",
    "Q = L2(mesh, order=1)\n",
    "X = V*Q\n",
    "print (\"V.ndof =\", V.ndof, \", Q.ndof =\", Q.ndof)\n",
    "\n",
    "gfu = SolveStokes(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "the mini element:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = VectorH1(mesh, order=1, dirichlet=\"wall|inlet|cyl\")\n",
    "V.SetOrder(TRIG,3)\n",
    "V.Update()\n",
    "Q = H1(mesh, order=1)\n",
    "X = V*Q\n",
    "\n",
    "gfu = SolveStokes(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stokes as a block-system\n",
    "We can now define separate bilinear-form and matrices for A and B, and combine them to a block-system:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = VectorH1(mesh, order=3, dirichlet=\"wall|inlet|cyl\")\n",
    "Q = H1(mesh, order=2)\n",
    "\n",
    "u,v = V.TnT()\n",
    "p,q = Q.TnT()\n",
    "\n",
    "a = BilinearForm(InnerProduct(Grad(u),Grad(v))*dx).Assemble()\n",
    "b = BilinearForm(div(u)*q*dx).Assemble()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Needed as preconditioner for the pressure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mp = BilinearForm(p*q*dx).Assemble()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Two right hand sides for the two spaces:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = LinearForm(V).Assemble()\n",
    "g = LinearForm(Q).Assemble();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Two `GridFunction`s for velocity and pressure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gfu = GridFunction(V, name=\"u\")\n",
    "gfp = GridFunction(Q, name=\"p\")\n",
    "uin = CoefficientFunction( (1.5*4*y*(0.41-y)/(0.41*0.41), 0) )\n",
    "gfu.Set(uin, definedon=mesh.Boundaries(\"inlet\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Combine everything to a block-system.\n",
    "`BlockMatrix` and `BlockVector` store references to the original matrices and vectors, no new large matrices are allocated. The same for the transpose matrix `b.mat.T`. It stores a wrapper for the original matrix, and replaces the call of the `Mult` function by `MultTrans`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "K = BlockMatrix( [ [a.mat, b.mat.T], [b.mat, None] ] )\n",
    "C = BlockMatrix( [ [a.mat.Inverse(V.FreeDofs()), None], [None, mp.mat.Inverse()] ] )\n",
    "\n",
    "rhs = BlockVector ( [f.vec, g.vec] )\n",
    "sol = BlockVector( [gfu.vec, gfp.vec] )\n",
    "\n",
    "solvers.MinRes (mat=K, pre=C, rhs=rhs, sol=sol, printrates=True, initialize=False);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Draw (gfu);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A stabilised lowest order formulation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = Mesh( geo.GenerateMesh(maxh=0.02))\n",
    "\n",
    "V = VectorH1(mesh, order=1, dirichlet=\"wall|inlet|cyl\")\n",
    "Q = H1(mesh, order=1)\n",
    "\n",
    "u,v = V.TnT()\n",
    "p,q = Q.TnT()\n",
    "\n",
    "a = BilinearForm(InnerProduct(Grad(u),Grad(v))*dx).Assemble()\n",
    "b = BilinearForm(div(u)*q*dx).Assemble()\n",
    "h = specialcf.mesh_size\n",
    "c = BilinearForm(-0.1*h*h*grad(p)*grad(q)*dx).Assemble()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mp = BilinearForm(p*q*dx).Assemble()\n",
    "\n",
    "f = LinearForm(V)\n",
    "f.Assemble()\n",
    "\n",
    "g = LinearForm(Q)\n",
    "g.Assemble()\n",
    "\n",
    "gfu = GridFunction(V, name=\"u\")\n",
    "gfp = GridFunction(Q, name=\"p\")\n",
    "uin = CF( (1.5*4*y*(0.41-y)/(0.41*0.41), 0) )\n",
    "gfu.Set(uin, definedon=mesh.Boundaries(\"inlet\"))\n",
    "\n",
    "K = BlockMatrix( [ [a.mat, b.mat.T], [b.mat, c.mat] ] )\n",
    "C = BlockMatrix( [ [a.mat.Inverse(V.FreeDofs()), None], [None, mp.mat.Inverse()] ] )\n",
    "\n",
    "rhs = BlockVector ( [f.vec, g.vec] )\n",
    "sol = BlockVector( [gfu.vec, gfp.vec] )\n",
    "\n",
    "solvers.MinRes (mat=K, pre=C, rhs=rhs, sol=sol, printrates=True, initialize=False, maxsteps=200);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Draw (gfu);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Draw (gfp);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## (P1-iso-P2)-P1\n",
    "\n",
    "The mesh is uniformly refined once. The velocity space is $P^1$ on the refined mesh, but the pressure is continuous $P^1$ in the coarse space. We obtain this by representing the pressure in the range of the coarse-to-fine prolongation operator. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = Mesh( geo.GenerateMesh(maxh=0.02))\n",
    "\n",
    "V = VectorH1(mesh, order=1, dirichlet=\"wall|inlet|cyl\")\n",
    "Q = H1(mesh, order=1)\n",
    "\n",
    "u,v = V.TnT()\n",
    "p,q = Q.TnT()\n",
    "\n",
    "mp = BilinearForm(p*q*dx).Assemble()\n",
    "invmp = mp.mat.Inverse(inverse=\"sparsecholesky\")\n",
    "\n",
    "mesh.ngmesh.Refine()\n",
    "V.Update()\n",
    "Q.Update()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# prolongation operator from mesh-level 0 to level 1:\n",
    "prol = Q.Prolongation().Operator(1)\n",
    "invmp2 = prol @ invmp @ prol.T\n",
    "\n",
    "a = BilinearForm(InnerProduct(Grad(u),Grad(v))*dx).Assemble()\n",
    "b = BilinearForm(div(u)*q*dx).Assemble()\n",
    "\n",
    "f = LinearForm(V)\n",
    "f.Assemble()\n",
    "\n",
    "g = LinearForm(Q)\n",
    "g.Assemble()\n",
    "\n",
    "gfu = GridFunction(V, name=\"u\")\n",
    "gfp = GridFunction(Q, name=\"p\")\n",
    "uin = CF( (1.5*4*y*(0.41-y)/(0.41*0.41), 0) )\n",
    "gfu.Set(uin, definedon=mesh.Boundaries(\"inlet\"))\n",
    "\n",
    "K = BlockMatrix( [ [a.mat, b.mat.T], [b.mat, None] ] )\n",
    "C = BlockMatrix( [ [a.mat.Inverse(V.FreeDofs()), None], [None, invmp2] ] )\n",
    "\n",
    "rhs = BlockVector ( [f.vec, g.vec] )\n",
    "sol = BlockVector( [gfu.vec, gfp.vec] )\n",
    "\n",
    "solvers.MinRes (mat=K, pre=C, rhs=rhs, sol=sol, \\\n",
    "                printrates=True, initialize=False, maxsteps=200);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Draw (gfu)\n",
    "Draw (gfp);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  },
  "nbsphinx": {
   "allow_errors": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}