{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Solving the Poisson equation\n",
    "\n",
    "\n",
    "## Solving 1D Poisson with periodic boundary conditions in Struphy\n",
    "\n",
    "This tutorial shows how to solve a manufactured 1D Poisson problem with the `Poisson` model.\n",
    "\n",
    "We use a periodic domain $x \\in [0, L)$ and choose an analytic solution\n",
    "\n",
    "$$\n",
    "\\phi_\\mathrm{exact}(x) = \\cos(kx),\n",
    "$$\n",
    "\n",
    "so the source for the stabilized Poisson problem\n",
    "\n",
    "$$\n",
    "-\\phi'' + \\varepsilon\\phi = \\rho\n",
    "$$\n",
    "\n",
    "is\n",
    "\n",
    "$$\n",
    "\\rho(x) = (k^2 + \\varepsilon)\\cos(kx).\n",
    "$$\n",
    "\n",
    "Then we run a Struphy `Simulation`, compare numerical and exact solutions, and compute the max-norm error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from struphy import Simulation, DerhamOptions, domains, grids, perturbations\n",
    "from struphy.initial.base import GenericPerturbation\n",
    "from struphy.models import Poisson\n",
    "\n",
    "Poisson.pde()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Manufactured periodic test case in 1D (aligned with verification test pattern)\n",
    "Lx = 2.0 * np.pi\n",
    "mode = 2\n",
    "k = mode * 2.0 * np.pi / Lx\n",
    "\n",
    "# Tiny stabilization makes the periodic problem uniquely solvable\n",
    "stab_eps = 1e-8\n",
    "\n",
    "# Build the source through the model variable as in test_verif_Poisson.py\n",
    "model = Poisson()\n",
    "model.propagators.poisson.options = model.propagators.poisson.Options(\n",
    "    rho=model.em_fields.source,\n",
    "    stab_eps=stab_eps,\n",
    "    )\n",
    "\n",
    "source_amp = k**2 + stab_eps\n",
    "model.em_fields.source.add_perturbation(\n",
    "    perturbations.ModesCos(ls=(mode,), amps=(source_amp,))\n",
    "    )\n",
    "\n",
    "phi_exact = lambda e1, e2, e3: np.cos(k * e1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "`Poisson.Options` controls the elliptic solve. In this example we only set:\n",
    "\n",
    "- `rho`: the manufactured source term.\n",
    "- `stab_eps`: a tiny stabilization to remove the constant null-space in the periodic case.\n",
    "\n",
    "All other option fields are left at their defaults (solver choice, preconditioner, tolerances, etc.)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1D periodic setup: periodic bcs are the default (None in each direction)\n",
    "domain = domains.Cuboid(l1=0.0, r1=Lx)\n",
    "grid = grids.TensorProductGrid(num_elements=(64, 1, 1))\n",
    "derham_opts = DerhamOptions()\n",
    "\n",
    "sim = Simulation(\n",
    "    model=model,\n",
    "    domain=domain,\n",
    "    grid=grid,\n",
    "    derham_opts=derham_opts,\n",
    "    verbose=False,\n",
    "    )\n",
    "\n",
    "# For a stationary Poisson solve, one step is enough\n",
    "sim.run(one_time_step=True, verbose=False)\n",
    "sim.pproc(verbose=False)\n",
    "sim.load_plotting_data(verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract 1D line data and compare to analytic solution\n",
    "x = sim.grids_phy[0][:, 0, 0]\n",
    "\n",
    "t_last = max(sim.spline_values.em_fields.phi_log.data.keys())\n",
    "phi_num = sim.spline_values.em_fields.phi_log.data[t_last][0][:, 0, 0]\n",
    "phi_ref = phi_exact(x, 0.0, 0.0)\n",
    "\n",
    "err = phi_num - phi_ref\n",
    "err_max = np.max(np.abs(err))\n",
    "\n",
    "plt.figure(figsize=(9, 4))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(x, phi_ref, \"k--\", label=\"analytic\")\n",
    "plt.plot(x, phi_num, \"o\", ms=3, label=\"Struphy\")\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(r\"$\\phi(x)$\")\n",
    "plt.title(\"Poisson solution\")\n",
    "plt.legend()\n",
    "plt.grid(alpha=0.3)\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.plot(x, err, \"r\", lw=1.5)\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(r\"$\\phi_h - \\phi_\\mathrm{exact}$\")\n",
    "plt.title(\"Pointwise error\")\n",
    "plt.grid(alpha=0.3)\n",
    "\n",
    "plt.tight_layout()\n",
    "print(f\"max-norm error ||phi_h - phi_exact||_inf = {err_max:.3e}\")\n",
    "\n",
    "# Electric field comparison: E = -grad(phi)\n",
    "E_num = -np.gradient(phi_num, x, edge_order=2)\n",
    "E_ref = k * np.sin(k * x)\n",
    "E_err = E_num - E_ref\n",
    "E_err_max = np.max(np.abs(E_err))\n",
    "\n",
    "plt.figure(figsize=(9, 4))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(x, E_ref, \"k--\", label=r\"analytic $-\\nabla\\phi$\")\n",
    "plt.plot(x, E_num, \"o\", ms=3, label=r\"Struphy $-\\nabla\\phi_h$\")\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(r\"$E_x(x)$\")\n",
    "plt.title(\"Electric field comparison\")\n",
    "plt.legend()\n",
    "plt.grid(alpha=0.3)\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.plot(x, E_err, \"m\", lw=1.5)\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(r\"$E_{x,h} - E_{x,\\mathrm{exact}}$\")\n",
    "plt.title(\"Electric field pointwise error\")\n",
    "plt.grid(alpha=0.3)\n",
    "\n",
    "plt.tight_layout()\n",
    "print(f\"max-norm error ||E_h - E_exact||_inf = {E_err_max:.3e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "## 2D Manufactured Test on a Rectangle\n",
    "\n",
    "We now solve a 2D manufactured Poisson problem on a rectangle with different side lengths:\n",
    "\n",
    "$$\n",
    "(x,y) \\in [0,L_x] \\times [0,L_y], \\qquad L_x \\neq L_y.\n",
    "$$\n",
    "\n",
    "Choose\n",
    "\n",
    "$$\n",
    "\\phi_\\mathrm{exact}(x,y) = \\sin\\left(\\frac{2\\pi x}{L_x}\\right)\\sin\\left(\\frac{2\\pi y}{L_y}\\right),\n",
    "$$\n",
    "\n",
    "which satisfies homogeneous Dirichlet conditions on all rectangle boundaries.\n",
    "For\n",
    "\n",
    "$$\n",
    "-\\Delta\\phi + \\varepsilon\\phi = \\rho,\n",
    "$$\n",
    "\n",
    "the manufactured source is\n",
    "\n",
    "$$\n",
    "\\rho(x,y) = \\left[\\left(\\frac{2\\pi}{L_x}\\right)^2 + \\left(\\frac{2\\pi}{L_y}\\right)^2 + \\varepsilon\\right] \\phi_\\mathrm{exact}(x,y).\n",
    "$$\n",
    "\n",
    "To stay consistent with Struphy's source handling, we build this source through a model perturbation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2D manufactured Poisson setup and solve\n",
    "Lx2 = 2.0\n",
    "Ly2 = 1.0\n",
    "eps2 = 1e-8\n",
    "\n",
    "kx2 = 2.0 * np.pi / Lx2\n",
    "ky2 = 2.0 * np.pi / Ly2\n",
    "\n",
    "phi2_exact = lambda e1, e2, e3: np.sin(kx2 * e1) * np.sin(ky2 * e2)\n",
    "\n",
    "model2 = Poisson()\n",
    "model2.propagators.poisson.options = model2.propagators.poisson.Options(\n",
    "    rho=model2.em_fields.source,\n",
    "    stab_eps=eps2,\n",
    "    )\n",
    "\n",
    "source_amp2 = kx2**2 + ky2**2 + eps2\n",
    "model2.em_fields.source.add_perturbation(\n",
    "    perturbations.ModesSinSin(ls=(1,), ms=(1,), amps=(source_amp2,))\n",
    "    )\n",
    "\n",
    "domain2 = domains.Cuboid(l1=0.0, r1=Lx2, l2=0.0, r2=Ly2)\n",
    "grid2 = grids.TensorProductGrid(num_elements=(48, 32, 1))\n",
    "derham_opts2 = DerhamOptions(\n",
    "    degree=(3, 3, 1),\n",
    "    bcs=((\"dirichlet\", \"dirichlet\"), (\"dirichlet\", \"dirichlet\"), None),\n",
    "    )\n",
    "\n",
    "sim2 = Simulation(\n",
    "    model=model2,\n",
    "    domain=domain2,\n",
    "    grid=grid2,\n",
    "    derham_opts=derham_opts2,\n",
    "    verbose=False,\n",
    "    )\n",
    "\n",
    "sim2.run(one_time_step=True, verbose=False)\n",
    "sim2.pproc(verbose=False)\n",
    "sim2.load_plotting_data(verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2D diagnostics and plots\n",
    "t2_last = max(sim2.spline_values.em_fields.phi_log.data.keys())\n",
    "X = sim2.grids_phy[0][:, :, 0]\n",
    "Y = sim2.grids_phy[1][:, :, 0]\n",
    "\n",
    "phi2_num = sim2.spline_values.em_fields.phi_log.data[t2_last][0][:, :, 0]\n",
    "phi2_ref = phi2_exact(X, Y, 0.0)\n",
    "err2 = phi2_num - phi2_ref\n",
    "err2_max = np.max(np.abs(err2))\n",
    "\n",
    "fig, axs = plt.subplots(1, 3, figsize=(15, 4.5), constrained_layout=True)\n",
    "\n",
    "im0 = axs[0].contourf(X, Y, phi2_num, levels=40, cmap=\"viridis\")\n",
    "axs[0].set_title(\"Numerical $\\\\phi_h(x,y)$\")\n",
    "axs[0].set_xlabel(\"x\")\n",
    "axs[0].set_ylabel(\"y\")\n",
    "plt.colorbar(im0, ax=axs[0])\n",
    "\n",
    "im1 = axs[1].contourf(X, Y, phi2_ref, levels=40, cmap=\"viridis\")\n",
    "axs[1].set_title(\"Analytic $\\\\phi_{\\\\mathrm{exact}}(x,y)$\")\n",
    "axs[1].set_xlabel(\"x\")\n",
    "axs[1].set_ylabel(\"y\")\n",
    "plt.colorbar(im1, ax=axs[1])\n",
    "\n",
    "im2 = axs[2].contourf(X, Y, err2, levels=40, cmap=\"coolwarm\")\n",
    "axs[2].set_title(\"Error $\\\\phi_h-\\\\phi_{\\\\mathrm{exact}}$\")\n",
    "axs[2].set_xlabel(\"x\")\n",
    "axs[2].set_ylabel(\"y\")\n",
    "plt.colorbar(im2, ax=axs[2])\n",
    "\n",
    "print(f\"2D max-norm error ||phi_h - phi_exact||_inf = {err2_max:.3e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "## 2D Manufactured Test on a Thin Annulus (HollowCylinder)\n",
    "\n",
    "As a second 2D geometry test, we solve Poisson on a disc-shaped domain with a small hole, i.e. a thin annulus.\n",
    "\n",
    "We use the physical manufactured solution\n",
    "\n",
    "$$\n",
    "\\phi_\\mathrm{exact}(r,\\theta)=\\sin\\!\\left(\\pi\\frac{r-a_1}{a_2-a_1}\\right)\\sin(m\\theta),\n",
    "$$\n",
    "\n",
    "with $r=\\sqrt{x^2+y^2}$ and $\\theta=\\mathrm{atan2}(y,x)$.\n",
    "This vanishes at $r=a_1$ and $r=a_2$, so we impose homogeneous Dirichlet conditions in the radial direction.\n",
    "\n",
    "For\n",
    "\n",
    "$$\n",
    "-\\Delta\\phi + \\varepsilon\\phi = \\rho,\n",
    "$$\n",
    "\n",
    "the source $\\rho$ is computed analytically in polar coordinates and injected as a physical-space perturbation.\n",
    "\n",
    "Below, all plots are shown in physical coordinates $(x,y)$ only."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "# HollowCylinder manufactured annulus test (physical-space definition)\n",
    "a1 = 0.8\n",
    "a2 = 1.0\n",
    "Lz3 = 1.0\n",
    "m3 = 3\n",
    "eps3 = 1e-8\n",
    "w3 = a2 - a1\n",
    "alpha3 = np.pi / w3\n",
    "\n",
    "def phi3_exact(x, y, z):\n",
    "    r = np.sqrt(x**2 + y**2)\n",
    "    theta = np.arctan2(y, x)\n",
    "    s = alpha3 * (r - a1)\n",
    "    return np.sin(s) * np.sin(m3 * theta)\n",
    "\n",
    "def rho3_exact(x, y, z):\n",
    "    r = np.sqrt(x**2 + y**2)\n",
    "    theta = np.arctan2(y, x)\n",
    "    s = alpha3 * (r - a1)\n",
    "\n",
    "    sin_s = np.sin(s)\n",
    "    cos_s = np.cos(s)\n",
    "    sin_mt = np.sin(m3 * theta)\n",
    "\n",
    "    # -Delta(phi) + eps*phi for phi(r,theta)=sin(s)sin(m theta)\n",
    "    term = (alpha3**2) * sin_s - (alpha3 / r) * cos_s + (m3**2 / r**2) * sin_s + eps3 * sin_s\n",
    "    return term * sin_mt\n",
    "\n",
    "model3 = Poisson()\n",
    "model3.propagators.poisson.options = model3.propagators.poisson.Options(\n",
    "    rho=model3.em_fields.source,\n",
    "    stab_eps=eps3,\n",
    "    )\n",
    "\n",
    "model3.em_fields.source.add_perturbation(\n",
    "    GenericPerturbation(rho3_exact, given_in_basis=\"physical\")\n",
    "    )\n",
    "\n",
    "domain3 = domains.HollowCylinder(a1=a1, a2=a2, Lz=Lz3)\n",
    "grid3 = grids.TensorProductGrid(num_elements=(36, 96, 1))\n",
    "derham_opts3 = DerhamOptions(\n",
    "    degree=(3, 3, 1),\n",
    "    bcs=((\"dirichlet\", \"dirichlet\"), None, None),\n",
    "    )\n",
    "\n",
    "sim3 = Simulation(\n",
    "    model=model3,\n",
    "    domain=domain3,\n",
    "    grid=grid3,\n",
    "    derham_opts=derham_opts3,\n",
    "    verbose=False,\n",
    "    )\n",
    "\n",
    "sim3.run(one_time_step=True, verbose=False)\n",
    "sim3.pproc(verbose=False)\n",
    "sim3.load_plotting_data(verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Annulus diagnostics and plots in physical coordinates only\n",
    "t3_last = max(sim3.spline_values.em_fields.phi_log.data.keys())\n",
    "X3 = sim3.grids_phy[0][:, :, 0]\n",
    "Y3 = sim3.grids_phy[1][:, :, 0]\n",
    "\n",
    "phi3_num = sim3.spline_values.em_fields.phi_log.data[t3_last][0][:, :, 0]\n",
    "phi3_ref = phi3_exact(X3, Y3, 0.0)\n",
    "err3 = phi3_num - phi3_ref\n",
    "err3_max = np.max(np.abs(err3))\n",
    "\n",
    "fig3, axs3 = plt.subplots(1, 3, figsize=(15, 4.8), constrained_layout=True)\n",
    "\n",
    "im30 = axs3[0].contourf(X3, Y3, phi3_num, levels=40, cmap=\"viridis\")\n",
    "axs3[0].set_title(\"Numerical $\\\\phi_h(x,y)$ on annulus\")\n",
    "axs3[0].set_xlabel(\"x\")\n",
    "axs3[0].set_ylabel(\"y\")\n",
    "axs3[0].set_aspect(\"equal\")\n",
    "plt.colorbar(im30, ax=axs3[0])\n",
    "\n",
    "im31 = axs3[1].contourf(X3, Y3, phi3_ref, levels=40, cmap=\"viridis\")\n",
    "axs3[1].set_title(\"Analytic $\\\\phi_{\\\\mathrm{exact}}(x,y)$ on annulus\")\n",
    "axs3[1].set_xlabel(\"x\")\n",
    "axs3[1].set_ylabel(\"y\")\n",
    "axs3[1].set_aspect(\"equal\")\n",
    "plt.colorbar(im31, ax=axs3[1])\n",
    "\n",
    "im32 = axs3[2].contourf(X3, Y3, err3, levels=40, cmap=\"coolwarm\")\n",
    "axs3[2].set_title(\"Error $\\\\phi_h-\\\\phi_{\\\\mathrm{exact}}$\")\n",
    "axs3[2].set_xlabel(\"x\")\n",
    "axs3[2].set_ylabel(\"y\")\n",
    "axs3[2].set_aspect(\"equal\")\n",
    "plt.colorbar(im32, ax=axs3[2])\n",
    "\n",
    "print(f\"Annulus max-norm error ||phi_h - phi_exact||_inf = {err3_max:.3e}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env (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": 5
}