{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Solver Layers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "**Layer approach:**\n",
    "\n",
    "* We create a `Preconditoner` object for the `BilinearForm` object\n",
    "* We create a `LinearSolver` object from the `BilinearForm.mat` and the `Preconditioner` object\n",
    "* We create a `NonlinearSolver` from the `BilinearForm` and the `LinearSolver` object (or LinearSolver-class+args)\n",
    "* We create a `TimeStepper` from a `BilinearForm` object and `NonLinearSolver` class+args"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "setting up a linear elasticity model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ngsolve import *\n",
    "from ngsolve.webgui import Draw\n",
    "\n",
    "from netgen.occ import *\n",
    "shape = Rectangle(1,0.1).Face()\n",
    "shape.edges.Max(X).name=\"right\"\n",
    "shape.edges.Min(X).name=\"left\"\n",
    "shape.edges.Max(Y).name=\"top\"\n",
    "shape.edges.Min(Y).name=\"bot\"\n",
    "shape.vertices.Min(X+Y).hpref=1\n",
    "shape.vertices.Min(X-Y).hpref=1\n",
    "\n",
    "mesh = Mesh(OCCGeometry(shape, dim=2).GenerateMesh(maxh=0.05))\n",
    "mesh.RefineHP(2)\n",
    "\n",
    "E, nu = 210, 0.2\n",
    "mu  = E / 2 / (1+nu)\n",
    "lam = E * nu / ((1+nu)*(1-2*nu))\n",
    "\n",
    "def C(u):\n",
    "    F = Id(2) + Grad(u)\n",
    "    return F.trans * F\n",
    "\n",
    "def NeoHooke (C):\n",
    "    return 0.5*mu*(Trace(C-Id(2)) + 2*mu/lam*Det(C)**(-lam/2/mu)-1)\n",
    "\n",
    "factor = Parameter(0.5)\n",
    "force = CoefficientFunction( (0,factor) )\n",
    "\n",
    "fes = H1(mesh, order=4, dirichlet=\"left\", dim=mesh.dim)\n",
    "u,v  = fes.TnT()\n",
    "\n",
    "a = BilinearForm(fes, symmetric=True)\n",
    "a += Variation(NeoHooke(C(u)).Compile()*dx)\n",
    "a += Variation((-InnerProduct(force,u)).Compile()*dx)\n",
    "gfu = GridFunction(fes)\n",
    "gfu.vec[:] = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "### Solver layers approach:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "a.Assemble()\n",
    "pre = preconditioners.MultiGrid(a)\n",
    "linsolve = solvers.CGSolver(a.mat, pre, maxiter=20)\n",
    "nlsolve = nonlinearsolvers.NewtonSolver(a=a, u=gfu, solver=linsolve)\n",
    "\n",
    "gfu.vec[:] = 0\n",
    "nlsolve.Solve(printing=False);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "scene = Draw (C(gfu)[0,0]-1, mesh, deformation=gfu, min=-0.1, max=0.1);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "`Variation` keeps the expression for the energy. We can use symbolic differentiation to convert it to a variational form:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = BilinearForm(NeoHooke(C(u)).Diff(u, v)*dx - force*v*dx).Assemble()\n",
    "pre = preconditioners.MultiGrid(a)\n",
    "\n",
    "nlsolve = nonlinearsolvers.NewtonSolver(a=a, u=gfu, \\\n",
    "            lin_solver_cls=solvers.CGSolver, \n",
    "            lin_solver_args={ \"pre\":pre })\n",
    "\n",
    "gfu.vec[:] = 0\n",
    "nlsolve.Solve(printing=False);\n",
    "\n",
    "Draw (C(gfu)[0,0]-1, mesh, deformation=gfu, min=-0.1, max=0.1);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "## Timesteppers\n",
    "\n",
    "* User provides the time-dependent equation as a variational form\n",
    "* Brand new: time-derivative `u.dt`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "force = CF( (0,-1))\n",
    "eq = NeoHooke(C(u)).Diff(u,v) * dx - force * v * dx + u.dt.dt * v * dx\n",
    "\n",
    "ts = timestepping.Newmark(equation=eq, dt=1e-1)\n",
    "\n",
    "gfu = GridFunction(fes)\n",
    "scene = Draw(gfu, deformation=True, center=(0,-0.3), radius=0.6)\n",
    "\n",
    "def callback(t, gfu):\n",
    "    # print(\"t = \", t, \"displacement = \", gfu(mesh(1,0.05)))\n",
    "    scene.Redraw()\n",
    "ts.Integrate(gfu, start_time=0, end_time=10, callback=callback);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "### New features to work on exptression trees"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = Mesh(unit_square.GenerateMesh(maxh=0.1))\n",
    "fes = H1(mesh, order=3, dirichlet=\".*\")\n",
    "u,v = fes.TnT()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "equ = u.dt*v*dx + 1e-3*u*v*dx + grad(u)*grad(v)*dx - v*dx\n",
    "print (equ)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "we can extract all proxies (trial- or test-functions) from the variational form:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "allproxies = equ.GetProxies(trial=True)\n",
    "print (\"proxies:\")\n",
    "for proxy in allproxies: print(proxy, \" python dtorder = \", proxy.dt_order)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "proxiesdt0 = [ prox for prox in allproxies if prox.dt_order==0 ]\n",
    "proxiesdt1 = [ prox for prox in allproxies if prox.dt_order==1 ]\n",
    "\n",
    "print (\"order 0 proxies:\")\n",
    "print (*proxiesdt0)\n",
    "\n",
    "print (\"order 1 proxies:\")\n",
    "print (*proxiesdt1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "## Building an implicit Euler timestepper\n",
    "\n",
    "The time-dependent equation is\n",
    "\n",
    "$$\n",
    "\\int \\partial_t u v + uv + \\nabla u \\nabla v - 1 v  \\; dx = 0\n",
    "$$\n",
    "\n",
    "An implicit Euler time-stepping method is to solve \n",
    "\n",
    "$$\n",
    "\\int \\frac{u - u_n}{\\tau} v + uv + \\nabla u \\nabla v - 1 v  \\; dx = 0\n",
    "$$\n",
    "\n",
    "where the unknown $u$ is the value of the new time-step $u_{n+1}$.\n",
    "\n",
    "So we have to replace the \n",
    "\n",
    "$$\n",
    "\\partial_t u  \\rightarrow \\frac{u - u_n}{\\tau}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "tau = Parameter(0.1)\n",
    "gfuold = GridFunction(fes)\n",
    "\n",
    "repl = {}\n",
    "for prox in allproxies:\n",
    "    if prox.dt_order==1:\n",
    "        repl[prox] = 1/tau*(prox.anti_dt - prox.anti_dt.ReplaceFunction(gfuold))\n",
    "\n",
    "for key,val in repl.items():\n",
    "    print (\"replace:\", key)\n",
    "    print (\"by\\n\", val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "ImplEuler_equ = equ.Replace(repl)\n",
    "print (ImplEuler_equ)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "bfIE = BilinearForm(ImplEuler_equ)\n",
    "bfIE.Assemble()\n",
    "pre = preconditioners.MultiGrid(bfIE)\n",
    "linsolve = krylovspace.CGSolver(bfIE.mat, pre, maxiter=20)\n",
    "\n",
    "gfu = GridFunction(fes)\n",
    "nlsolve = nonlinearsolvers.NewtonSolver(a=bfIE, u=gfu, solver=linsolve)\n",
    "\n",
    "scene = Draw(gfu, deformation=True, scale=10)\n",
    "gfuold.vec[:] = 0\n",
    "\n",
    "from time import sleep\n",
    "for i in range(20):\n",
    "    nlsolve.Solve()\n",
    "    gfuold.vec[:] = gfu.vec\n",
    "    scene.Redraw()\n",
    "    sleep(0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "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.13.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}