{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "\n",
    "logging.basicConfig(\n",
    "    level=logging.INFO,\n",
    "    format='%(asctime)s %(levelname)s %(name)-20s :: %(message)s'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Diffusion coefficient problem \n",
    "\n",
    "We implement the inverse problem to recover the diffusion coefficient given the solution to the PDE. formally that is we consider the problem \n",
    "$$\n",
    "\\nabla a \\nabla u = f \\text{ in } \\Omega\n",
    "$$\n",
    "$$\n",
    "u = g \\text{ on } \\partial \\Omega.\n",
    "$$\n",
    "\n",
    "As variational formulation this is given by the bilinear form\n",
    "\n",
    "$$     b(u,v) = \\int_\\Omega a \\nabla(u) \\nabla(v) dx $$\n",
    "\n",
    "and the linear form\n",
    "\n",
    "$$    F(v) = \\int_{\\partial \\Omega} f v ds. $$\n",
    "As the bilinear form depends linearly on the coefficient $a$ we may use the provided class `SecondOrderEllipticCoefficientPDE`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from regpy.operators.ngsolve import SecondOrderEllipticCoefficientPDE\n",
    "import ngsolve as ngs\n",
    "\n",
    "class diffusion(SecondOrderEllipticCoefficientPDE):\n",
    "    def __init__(self, domain, sol_domain,bdr_val = None,a_bdr_val=None):\n",
    "        super().__init__(domain, sol_domain, bdr_val=bdr_val,a_bdr_val=a_bdr_val)\n",
    "\n",
    "    def _bf(self,a,u,v):\n",
    "        return a*ngs.grad(u)*ngs.grad(v)*ngs.dx\n",
    "    \n",
    "    def _lf(self):\n",
    "        p = ngs.GridFunction(self.codomain.fes)\n",
    "        p.Set(-2*ngs.exp(ngs.x+ngs.y))\n",
    "        lf = ngs.LinearForm(self.codomain.fes)\n",
    "        lf += p * self.v * ngs.dx\n",
    "        return lf.Assemble()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we define a mesh and FES spaces for the coefficients and solution spaces and construct the operator by giving it the boundary values. Moreover define the exact solution and take their boundary values as input to the operator. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from netgen.geom2d import unit_square\n",
    "from regpy.vecsps.ngsolve import NgsSpace\n",
    "\n",
    "bdr = \"left|top|right|bottom\"\n",
    "mesh = ngs.Mesh(unit_square.GenerateMesh(maxh=0.2))\n",
    "fes_domain = ngs.H1(mesh, order=6, dirichlet = bdr)\n",
    "domain = NgsSpace(fes_domain,bdr = bdr)\n",
    "\n",
    "bdr = \"left|top|right|bottom\"\n",
    "fes_codomain = ngs.H1(mesh, order=6, dirichlet=bdr)\n",
    "codomain = NgsSpace(fes_codomain, bdr=bdr)\n",
    "\n",
    "bdr_coeff = ngs.sin(ngs.x*4)+2*ngs.y\n",
    "bdr_gf = ngs.GridFunction(codomain.fes)\n",
    "bdr_gf.Set(bdr_coeff,definedon=codomain.fes.mesh.Boundaries(codomain.bdr))\n",
    "bdr_val = codomain.from_ngs(bdr_gf)\n",
    "\n",
    "exact_solution_coeff = 0.5*ngs.exp(-4*(ngs.x-0.5)**2 +4*(ngs.y-0.5)**2)\n",
    "exact_solution = domain.from_ngs( exact_solution_coeff )\n",
    "p = ngs.GridFunction(domain.fes)\n",
    "p.Set(exact_solution_coeff,definedon=domain.fes.mesh.Boundaries(domain.bdr))\n",
    "a_bdr_val = domain.from_ngs( p )\n",
    "\n",
    "op = diffusion(\n",
    "    domain, codomain, bdr_val=bdr_val,a_bdr_val = a_bdr_val\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get data by perturbing the exact data constructed from a finer meshed solution by some random noise. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fine_mesh = ngs.Mesh(unit_square.GenerateMesh(maxh=0.05))\n",
    "fine_fes_domain = ngs.H1(fine_mesh, order=6, dirichlet = bdr)\n",
    "fine_domain = NgsSpace(fine_fes_domain,bdr = bdr)\n",
    "\n",
    "fine_exact_solution = fine_domain.from_ngs( exact_solution_coeff )\n",
    "p = ngs.GridFunction(fine_domain.fes)\n",
    "p.Set(exact_solution_coeff,definedon=fine_domain.fes.mesh.Boundaries(fine_domain.bdr))\n",
    "fine_a_bdr_val = fine_domain.from_ngs( p )\n",
    "\n",
    "fine_op = diffusion(\n",
    "    fine_domain, codomain, bdr_val=bdr_val,a_bdr_val = fine_a_bdr_val\n",
    ")\n",
    "\n",
    "exact_data = fine_op(fine_exact_solution)\n",
    "noise = 0.05 * codomain.randn()\n",
    "data = exact_data+noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ngsolve.webgui import Draw\n",
    "\n",
    "Draw(codomain.to_ngs(data))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the inversion\n",
    "\n",
    "Define a an initial guess by choosing the constant function on the domain and satisfying the apriori known boundary conditions. Then define a regularization setting by choosing appropriate norms on both the domain and codomain. We choose iterative regularized Gauss-Newton method as our regularization scheme. Combined with a stopping rule that is composed of a max iteration and discrepancy principle.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from regpy.solvers import RegularizationSetting\n",
    "from regpy.solvers.nonlinear.irgnm import IrgnmCG\n",
    "import regpy.stoprules as rules\n",
    "from regpy.hilbert import Hm0\n",
    "\n",
    "gfu_constant = ngs.GridFunction(domain.fes)\n",
    "gfu_constant.Set(1)\n",
    "gfu_init = ngs.GridFunction(domain.fes)\n",
    "gfu_init.vec.data = ngs.Projector(domain.fes.FreeDofs(), range=True).Project(gfu_constant.vec) + domain.to_ngs(a_bdr_val).vec\n",
    "init = domain.from_ngs(gfu_init)\n",
    "\n",
    "Draw(gfu_init)\n",
    "\n",
    "setting = RegularizationSetting(op=op, penalty=Hm0, data_fid=Hm0)\n",
    "\n",
    "\n",
    "irgnm = IrgnmCG(setting, data, regpar = 0.01,init = init)\n",
    "\n",
    "stoprule = (\n",
    "        rules.CountIterations(30) +\n",
    "        rules.Discrepancy(setting.h_codomain.norm,data,setting.h_codomain.norm(noise),tau = 1.3)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The inversion\n",
    "\n",
    "Do the inversion by calling `irgnm.run(stoprule)`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reco, reco_data = irgnm.run(stoprule)\n",
    "\n",
    "Draw(exact_solution_coeff, op.domain.fes.mesh, \"exact\")\n",
    "\n",
    "# Draw reconstructed solution\n",
    "Draw(domain.to_ngs(reco),op.domain.fes.mesh, \"reconstruction\")\n",
    "\n",
    "# Draw data space\n",
    "Draw(codomain.to_ngs(data),op.codomain.fes.mesh, \"exact data\")\n",
    "Draw(codomain.to_ngs(reco_data),op.codomain.fes.mesh, \"data of reconstruction\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "3.12.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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}