{
 "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": [
    "from dolfinx.fem import (Constant, dirichletbc, Function, FunctionSpace, locate_dofs_geometrical,\n",
    "                         locate_dofs_topological)\n",
    "from dolfinx.fem.petsc import LinearProblem\n",
    "from dolfinx.mesh import CellType, create_rectangle, locate_entities_boundary\n",
    "from ufl import Identity, Measure, TestFunction, TrialFunction, VectorElement, dot, dx, inner, grad, nabla_div, sym\n",
    "from mpi4py import MPI\n",
    "from petsc4py.PETSc import ScalarType\n",
    "\n",
    "import numpy as np\n",
    "import pyvista\n",
    "\n",
    "mesh = create_rectangle(MPI.COMM_WORLD, np.array([[0,0],[L, H]]), [30,30], cell_type=CellType.triangle)\n",
    "element = VectorElement(\"CG\", mesh.ufl_cell(), 1)\n",
    "V = 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_zero = np.array((0,)*mesh.geometry.dim, dtype=ScalarType)\n",
    "bc = dirichletbc(u_zero, locate_dofs_geometrical(V, clamped_boundary), V)"
   ]
  },
  {
   "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": "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_topological`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def right(x):\n",
    "    return np.logical_and(np.isclose(x[0], L), x[1] < H)\n",
    "boundary_facets = locate_entities_boundary(mesh, mesh.topology.dim-1, right)\n",
    "boundary_dofs_x = locate_dofs_topological(V.sub(0), mesh.topology.dim-1, boundary_facets)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now create our Dirichlet condition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "bcx = dirichletbc(ScalarType(0), 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": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "T = Constant(mesh, ScalarType((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": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = 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": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def epsilon(u):\n",
    "    return sym(grad(u)) \n",
    "def sigma(u):\n",
    "    return lambda_ * nabla_div(u) * Identity(u.geometric_dimension()) + 2*mu*epsilon(u)\n",
    "\n",
    "u = TrialFunction(V)\n",
    "v = TestFunction(V)\n",
    "f = Constant(mesh, ScalarType((0, -rho*g)))\n",
    "a = inner(sigma(u), epsilon(v)) * dx\n",
    "L = dot(f, v) * dx + 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": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "problem = 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": 10,
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[0m\u001b[2m2022-08-02 20:55:36.872 (   1.531s) [        BA8E8000]    vtkExtractEdges.cxx:435   INFO| \u001b[0mExecuting edge extractor: points are renumbered\u001b[0m\n",
      "\u001b[0m\u001b[2m2022-08-02 20:55:36.874 (   1.533s) [        BA8E8000]    vtkExtractEdges.cxx:551   INFO| \u001b[0mCreated 2760 edges\u001b[0m\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4f974f9aeb84442d9302ea6f1724ce3a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Renderer(camera=PerspectiveCamera(aspect=1.3333333333333333, children=(DirectionalLight(intensity=0.25, positi…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pyvista.set_jupyter_backend(\"pythreejs\")\n",
    "from dolfinx.plot import create_vtk_mesh\n",
    "\n",
    "# Create plotter and pyvista grid\n",
    "p = pyvista.Plotter()\n",
    "topology, cell_types, x = create_vtk_mesh(V)\n",
    "grid = pyvista.UnstructuredGrid(topology, cell_types, x)\n",
    "\n",
    "# Attach vector values to grid and warp grid by vector\n",
    "\n",
    "vals = np.zeros((x.shape[0], 3))\n",
    "vals[:,:len(uh)] = uh.x.array.reshape((x.shape[0], len(uh)))\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,opacity=0.8)\n",
    "p.view_xy()\n",
    "if not pyvista.OFF_SCREEN:\n",
    "    p.show()\n",
    "else:\n",
    "    pyvista.start_xvfb()\n",
    "    fig_array = p.screenshot(f\"component.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,py:light"
  },
  "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}