{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "executive-protection",
   "metadata": {},
   "source": [
    "\n",
    "# Bound constrained minimization with PETSc-VI solvers\n",
    "\n",
    "This example solve the bound constrained minimization problem\n",
    " in the domain $(x,y)$ in $\\Omega \\equiv (0,L_x)\\times(0,L_y)$\n",
    "\n",
    " $$\\min_{u\\in\\mathcal C}\\, \\mathcal{F}(u),\\qquad \\text{with}\\qquad u\\in\\mathcal{C}\\equiv\\{u\\in H_1(\\Omega):\\;0\\leq u\\leq 1,\\;u(0,y)= 0,\\;u(L_x,y) = 1\\}$$\n",
    "\n",
    " where $\\mathcal{F}(u)$ is the functional defined by the form\n",
    "\n",
    " $$\\mathcal{F}(u) =\\int_\\Omega \\left(\\frac{w(u)}{\\ell}+\\ell\\, \\nabla u\\cdot  \\nabla u\\right)\\mathrm{d}x$$\n",
    " \n",
    " The solution $u^*\\in\\mathcal{C} $ to this problem must verify the following Variational Inequality (VI):\n",
    " \n",
    "$$ \\mathcal{F}'(u^*)(v-u^*)\\geq 0,\\quad \\forall v \\in\\mathcal{C} $$\n",
    "\n",
    "We show below how to solve this problem with FEniCS using the VI solvers provided by PETSc\n",
    "\n",
    "To define the problem we need to specify the \n",
    " - The functional to be minimized $\\mathcal{F}$, and its first and second directional derivatives\n",
    " - The lower $u_L$ and upper  $u_U$ bounds, here given by $0$ and $1$ respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "residential-discussion",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import required libraries\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import dolfinx\n",
    "import dolfinx.plot\n",
    "import dolfinx.io\n",
    "import ufl\n",
    "\n",
    "from mpi4py import MPI\n",
    "from petsc4py import PETSc\n",
    "\n",
    "import pyvista\n",
    "from pyvista.utilities.xvfb import start_xvfb\n",
    "start_xvfb(wait=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "contemporary-lawsuit",
   "metadata": {},
   "source": [
    "## Mesh, function space, boundary condition\n",
    "\n",
    "We first define the mesh, the function space, and the boundary conditions, as usual\n",
    "\n",
    "We discretize the field with standard $P_1$ finite elements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "printable-library",
   "metadata": {},
   "outputs": [],
   "source": [
    "L = 1.0\n",
    "H = 0.2\n",
    "mesh = dolfinx.RectangleMesh(MPI.COMM_WORLD, [(0.0, 0.0, 0.0), (L, H, 0.0)], [100, 15])\n",
    "\n",
    "with dolfinx.io.XDMFFile(MPI.COMM_WORLD, \"mesh.xdmf\", \"w\") as f:\n",
    "    f.write_mesh(mesh)\n",
    "\n",
    "V = dolfinx.FunctionSpace(mesh, (\"CG\", 1))           \n",
    "                               \n",
    "zero = dolfinx.Function(V)\n",
    "with zero.vector.localForm() as loc:\n",
    "    loc.set(0.0)\n",
    "    \n",
    "one = dolfinx.Function(V)\n",
    "with one.vector.localForm() as loc:\n",
    "    loc.set(1.0)\n",
    "    \n",
    "def left(x):\n",
    "    is_close = np.isclose(x[0], 0.0)\n",
    "    return is_close\n",
    "\n",
    "def right(x):\n",
    "    is_close = np.isclose(x[0], L)\n",
    "    return is_close\n",
    "\n",
    "left_facets = dolfinx.mesh.locate_entities_boundary(mesh, mesh.topology.dim - 1, left)\n",
    "left_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, left_facets)\n",
    "\n",
    "right_facets = dolfinx.mesh.locate_entities_boundary(mesh, mesh.topology.dim - 1, left)\n",
    "right_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, right_facets)\n",
    "\n",
    "bcs = [dolfinx.DirichletBC(zero, left_dofs), dolfinx.DirichletBC(one, right_dofs)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "racial-oklahoma",
   "metadata": {},
   "source": [
    "## The functional to by minimized\n",
    "\n",
    "We define here the functional and its directional derivatives"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "settled-detection",
   "metadata": {},
   "outputs": [],
   "source": [
    "u = dolfinx.Function(V)\n",
    "\n",
    "ell = 0.3\n",
    "\n",
    "def w(alpha):\n",
    "    return alpha\n",
    "\n",
    "functional = (ell * ufl.inner(ufl.grad(u), ufl.grad(u)) + w(u)/ell)*ufl.dx\n",
    "\n",
    "dfunctional = ufl.derivative(functional, u, ufl.TestFunction(V))\n",
    "\n",
    "ddfunctional = ufl.derivative(dfunctional, u, ufl.TrialFunction(V))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tropical-triangle",
   "metadata": {},
   "source": [
    "## Set up the SNES VI solver\n",
    "\n",
    "To slve the problem we use the variational inequality solver `SNES` provided by `PETSc`. \n",
    "\n",
    "To use the `SNES` solver we need to provide suitable function to assemble the residual and the Jacobian, as for the Netwon solver, but with few differences. A class doing this job ggiven the the FEniCS form, the Function and the BCs is provided in the file [../python/snes_problem.py](https://gitlab.com/newfrac/newfrac-fenicsx-training/-/tree/main/notebooks/python/snes_problem.py).\n",
    "\n",
    "**Exersice \"Find the differences\":** Which are the difference between `SNESProblem` and `NonlinearPDEProblem`? Why are they necessary in your opinion? Look how the class are used. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "contrary-cambodia",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"../python/\")\n",
    "from snes_problem import SNESProblem\n",
    "\n",
    "snes_problem = SNESProblem(dfunctional, ddfunctional, u, bcs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sacred-kinase",
   "metadata": {},
   "source": [
    "We set up now the `SNES` solver though the `petsc4py` python interface\n",
    "\n",
    "Note that we assing the solvers:\n",
    "- the function to evaluate the function (residual)\n",
    "- the function to evaluate the Jacobian\n",
    "- the bounds\n",
    "\n",
    "We select the specific VI solver `vinewtonrsls`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "sacred-liberia",
   "metadata": {},
   "outputs": [],
   "source": [
    "b = dolfinx.cpp.la.create_vector(V.dofmap.index_map, V.dofmap.index_map_bs)\n",
    "J = dolfinx.fem.create_matrix(snes_problem.a)\n",
    "\n",
    "# Create Newton solver and solve\n",
    "solver_snes = PETSc.SNES().create()\n",
    "solver_snes.setType(\"vinewtonrsls\")\n",
    "solver_snes.setFunction(snes_problem.F, b)\n",
    "solver_snes.setJacobian(snes_problem.J, J)\n",
    "solver_snes.setTolerances(rtol=1.0e-9, max_it=50)\n",
    "solver_snes.getKSP().setType(\"preonly\")\n",
    "solver_snes.getKSP().setTolerances(rtol=1.0e-9)\n",
    "solver_snes.getKSP().getPC().setType(\"lu\")\n",
    "\n",
    "# We set the bound (Note: they are passed as reference and not as values)\n",
    "solver_snes.setVariableBounds(zero.vector,one.vector)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sharp-search",
   "metadata": {},
   "source": [
    "## Solve the problem\n",
    "\n",
    "We can now solve the problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "moving-monkey",
   "metadata": {},
   "outputs": [],
   "source": [
    "solver_snes.solve(None, u.vector)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "interpreted-module",
   "metadata": {},
   "source": [
    "## Prostprocessing\n",
    "\n",
    "We save to file and plot the solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "pursuant-controversy",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "Path(\"output\").mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "with dolfinx.io.XDMFFile(MPI.COMM_WORLD, \"output/u.xdmf\", \"w\") as f:\n",
    "    f.write_mesh(mesh)\n",
    "    f.write_function(u)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "chicken-knight",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dolfinx.plot\n",
    "topology, cell_types = dolfinx.plot.create_vtk_topology(mesh, mesh.topology.dim)\n",
    "import pyvista\n",
    "grid = pyvista.UnstructuredGrid(topology, cell_types, mesh.geometry.x)\n",
    "\n",
    "grid.point_arrays[\"alpha\"] = u.compute_point_values().real\n",
    "grid.set_active_scalars(\"alpha\")\n",
    "\n",
    "if pyvista.OFF_SCREEN:\n",
    "    from pyvista.utilities.xvfb import start_xvfb\n",
    "    start_xvfb(wait=0.1)\n",
    "\n",
    "plotter = pyvista.Plotter()\n",
    "plotter.add_mesh(grid, show_edges=False, show_scalar_bar=True)\n",
    "plotter.view_xy()\n",
    "if not pyvista.OFF_SCREEN:\n",
    "    plotter.show()\n",
    "figure = plotter.screenshot(\"output/u.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "external-thirty",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import evaluate_on_points\n",
    "\n",
    "points = np.zeros((3, 101))\n",
    "points[0] = np.linspace(0., 1., 101)\n",
    "points[1] = .1\n",
    "\n",
    "points_on_proc, u_values = evaluate_on_points(u, points)\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig = plt.figure()\n",
    "plt.plot(points_on_proc[:, 0], u_values, \"b--\", linewidth = 2)\n",
    "plt.grid(True)\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(\"u\")\n",
    "# If run in parallel as a python file, we save a plot per processor\n",
    "plt.savefig(f\"output/membrane_rank{MPI.COMM_WORLD.rank:d}.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "unnecessary-imaging",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "incorporate-tissue",
   "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": 5
}