{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Component-wise Dirichlet BC\n",
    "Author: Jørgen S. Dokken\n",
    "\n",
    "In this section, we will learn how to prescribe Dirichlet boundary conditions on a component of your unknown $u_h$.\n",
    "We will illustrate the problem using a `VectorElement`. However, the method generalizes to any `MixedElement`.\n",
    "\n",
    "We will use a slightly modified version of [the linear elasticity demo](./../chapter2/linearelasticity_code), namely\n",
    "$$\n",
    "-\\nabla \\cdot \\sigma (u) = f\\quad \\text{in } \\Omega,\n",
    "$$\n",
    "\n",
    "$$\n",
    "\\sigma \\cdot n = 0 \\quad \\text{on } \\partial \\Omega_N,\n",
    "$$\n",
    "\n",
    "$$\n",
    "u= 0\\quad \\text{at } \\partial\\Omega_{D},\n",
    "$$\n",
    "\n",
    "$$\n",
    "u_x=0 \\quad \\text{at } \\partial\\Omega_{Dx},\n",
    "$$\n",
    "\n",
    "$$\n",
    "\\sigma(u)= \\lambda \\mathrm{tr}(\\epsilon(u))I + 2 \\mu \\epsilon(u), \\qquad \\epsilon(u) = \\frac{1}{2}\\left(\\nabla u + (\\nabla u )^T\\right).\n",
    "$$\n",
    "We will consider a two dimensional box spanning $[0,L]\\times[0,H]$, where\n",
    "$\\partial\\Omega_N$ is the left and right side of the beam, $\\partial\\Omega_D$ the bottom of the  beam, while $\\partial\\Omega_{Dx}$ is the right side of the beam.\n",
    "We will prescribe a displacement $u_x=0$ on the right side of the beam, while the beam is being deformed under its own weight. The sides of the box is traction free."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "L = 1\n",
    "H = 1.3\n",
    "lambda_ = 1.25 \n",
    "mu = 1\n",
    "rho = 1\n",
    "g = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As in the previous demos, we define our mesh and function space. We will create a `ufl.VectorElement` to create a two dimensional vector space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dolfinx\n",
    "import ufl\n",
    "import numpy as np\n",
    "from mpi4py import MPI\n",
    "from dolfinx.cpp.mesh import CellType\n",
    "\n",
    "mesh = dolfinx.RectangleMesh(MPI.COMM_WORLD, np.array([[0,0,0],[L, H,0]]), [30,30], cell_type=CellType.triangle)\n",
    "element = ufl.VectorElement(\"CG\", mesh.ufl_cell(), 1)\n",
    "V = dolfinx.FunctionSpace(mesh, element)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Boundary conditions\n",
    "As we would like to clamp the boundary at $x=0$, we do this by using a marker function, we use `dolfinx.fem.locate_dofs_geometrical` to identify the relevant degrees of freedom."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def clamped_boundary(x):\n",
    "    return np.isclose(x[1], 0)\n",
    "\n",
    "u_D = dolfinx.Function(V)\n",
    "with u_D.vector.localForm() as loc:\n",
    "    loc.set(0)\n",
    "bc = dolfinx.DirichletBC(u_D, dolfinx.fem.locate_dofs_geometrical(V, clamped_boundary))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we would like to constrain the $x$-component of our solution at $x=L$ to $0$. We start by creating the sub space only containing the $x$\n",
    "-component."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "V0 = V.sub(0).collapse()\n",
    "uDx = dolfinx.Function(V0)\n",
    "with uDx.vector.localForm() as uDx_loc:\n",
    "    uDx_loc.set(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we locate the degrees of freedom on the top boundary. However, as the boundary condition is in a sub space of our solution, we need to supply both the parent space $V$ and the sub space $V_0$ to `dolfinx.locate_dofs_geometrical`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def right(x):\n",
    "    return np.logical_and(np.isclose(x[0], L), x[1] < H)\n",
    "boundary_dofs_x = dolfinx.fem.locate_dofs_geometrical((V.sub(0), V0), right)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now create our Dirichlet condition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "bcx = dolfinx.DirichletBC(uDx, boundary_dofs_x, V.sub(0))\n",
    "bcs = [bc, bcx]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we want the traction $T$ over the remaining boundary to be $0$, we create a `dolfinx.Constant`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "T = dolfinx.Constant(mesh, (0, 0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also want to specify the integration measure $\\mathrm{d}s$, which should be the integral over the boundary of our domain. We do this by using `ufl`, and its built in integration measures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ufl\n",
    "ds = ufl.Measure(\"ds\", domain=mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Variational formulation\n",
    "We are now ready to create our variational formulation in close to mathematical syntax, as for the previous problems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def epsilon(u):\n",
    "    return ufl.sym(ufl.grad(u)) \n",
    "def sigma(u):\n",
    "    return lambda_ * ufl.nabla_div(u) * ufl.Identity(u.geometric_dimension()) + 2*mu*epsilon(u)\n",
    "\n",
    "u = ufl.TrialFunction(V)\n",
    "v = ufl.TestFunction(V)\n",
    "f = dolfinx.Constant(mesh, (0, -rho*g))\n",
    "a = ufl.inner(sigma(u), epsilon(v)) * ufl.dx\n",
    "L = ufl.dot(f, v) * ufl.dx + ufl.dot(T, v) * ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solve the linear variational problem\n",
    "As in the previous demos, we assemble the matrix and right hand side vector and use PETSc to solve our variational problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "problem = dolfinx.fem.LinearProblem(a, L, bcs=bcs, petsc_options={\"ksp_type\": \"preonly\", \"pc_type\": \"lu\"})\n",
    "uh = problem.solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualization\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "90df8d40913b4d9992330379b79a776e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "ViewInteractiveWidget(height=800, layout=Layout(height='auto', width='100%'), width=800)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pyvista\n",
    "import dolfinx.plot\n",
    "# Start virtual framebuffer\n",
    "pyvista.start_xvfb(wait=0.05)\n",
    "\n",
    "# Create plotter and pyvista grid\n",
    "p = pyvista.Plotter(title=\"Deflection\", window_size=[800, 800])\n",
    "topology, cell_types = dolfinx.plot.create_vtk_topology(mesh, mesh.topology.dim)\n",
    "grid = pyvista.UnstructuredGrid(topology, cell_types, mesh.geometry.x)\n",
    "\n",
    "# Attach vector values to grid and warp grid by vector\n",
    "vals_2D = uh.compute_point_values().real \n",
    "vals = np.zeros((vals_2D.shape[0], 3))\n",
    "vals[:,:2] = vals_2D\n",
    "grid[\"u\"] = vals\n",
    "actor_0 = p.add_mesh(grid, style=\"wireframe\", color=\"k\")\n",
    "warped = grid.warp_by_vector(\"u\", factor=1.5)\n",
    "actor_1 = p.add_mesh(warped, show_edges=True)\n",
    "p.view_xy()\n",
    "if not pyvista.OFF_SCREEN:\n",
    "   p.show()\n",
    "else:\n",
    "    fig_array = p.screenshot(f\"component.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}