{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Polar TM Maxwell Solver Test / Tutorial with mirror symmetry and discrete rotational symmetry\n",
    "\n",
    "Sanity checks for correctness of the polar TM Maxwell solver with discrete rotational symmetry and mirror symmetry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.sparse as sp\n",
    "import scipy.sparse.linalg as spla\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.colors as colors\n",
    "\n",
    "import sys\n",
    "sys.path.append(\"../../\")\n",
    "\n",
    "from dolphindes.geometry import PolarFDFDGeometry\n",
    "from dolphindes.maxwell import TM_Polar_FDFD, plot_real_polar_field, plot_cplx_polar_field, expand_symmetric_field"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "We start by specifying a domain of interest shared between all basic tests below: a circular computational domain with radius $2\\lambda$, surrounded by pml. We will consider only structures and sources with 6-fold discrete rotational symmetry and mirrr symmetry, which allows us to use the n_sector and mirror feature of the polar solver to restrict computation within the irreducible domain. In this case the irreducible domain is a wedge of angle $30^\\circ$, which under reflection forms one 6-fold rotational unit cell. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "wvlgth = 1.0\n",
    "Qabs = np.inf # supports complex frequency, test by setting finite Qabs\n",
    "omega = 2*np.pi / wvlgth * (1 + 1j/2/Qabs)\n",
    "\n",
    "r_nonpml = 2.0 # center circle radius\n",
    "w_pml = 0.5 # surrounding pml thickness\n",
    "r_tot = r_nonpml + w_pml # total computational domain radius\n",
    "\n",
    "gpr = 30\n",
    "#gpr = 60 # high res option\n",
    "dr = 1.0/gpr # radial grid size\n",
    "Nr = int(np.round(r_tot / dr))\n",
    "Npml = int(np.round(w_pml / dr))\n",
    "\n",
    "# setting azimuthal grid size. Note the azimuthal pixel width gets larger with radius\n",
    "Nphi_full = 180 # this gives ~ 0.043 pixel width at the edge of computational domain for R_tot=2.5\n",
    "#Nphi_full = 360 # high res option\n",
    "\n",
    "# Example: 6-fold rotational symmetry and mirror symmetry (60° sectors, 30° half sectors)\n",
    "n_sectors = 6\n",
    "mirror = True\n",
    "assert Nphi_full % (2*n_sectors) == 0, \"Nphi_full must be divisible by 2*n_sectors\"\n",
    "\n",
    "Nphi_sector = Nphi_full // n_sectors\n",
    "Nphi_halfsector = Nphi_sector // 2  # azimuthal points in one 30° irreducible domain\n",
    "\n",
    "geo_halfsector = PolarFDFDGeometry(Nphi_halfsector, Nr, Npml, dr, \n",
    "                                   n_sectors=n_sectors, mirror=True)\n",
    "\n",
    "FDFD_halfsector = TM_Polar_FDFD(omega, geo_halfsector)\n",
    "\n",
    "# Get coordinate grids\n",
    "phi_grid_halfsector, r_grid, phi_grid_full = FDFD_halfsector.get_symmetric_grids()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "## Vacuum Dipole"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.special import hankel1 # for checking against analytical dipole field"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "### vacuum dipole at the center\n",
    "The simplest test: a dipole source at the origin. Because of the radial grid offset there is no gridpoint exactly at the origin, so we specify the dipole using all the points closest to the origin."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# setup dipole source at origin\n",
    "J_r = np.zeros(Nr)\n",
    "J_r[0] = 1.0 / (np.pi* dr**2)\n",
    "# note the normalization: each r gridpoint r_i covers the range [r_i - dr/2, r_i + dr/2)\n",
    "\n",
    "# Use Nphi_halfsector instead of Nphi_full for the symmetric case\n",
    "J_center_dipole = np.kron(np.ones(Nphi_halfsector), J_r)\n",
    "# note the ordering of the flattened fields: all r for one ray then move counter-clockwise\n",
    "\n",
    "E_center_dipole = FDFD_halfsector.get_TM_field(J_center_dipole) # get the center dipole field\n",
    "plot_cplx_polar_field(E_center_dipole, phi_grid_halfsector, r_grid) # use sector grid for plotting\n",
    "\n",
    "# check against analytical solution along phi=0 ray (first ray in sector)\n",
    "analytic_E_grid = (-omega/4) * hankel1(0, omega*r_grid)\n",
    "fig, (ax1,ax2) = plt.subplots(ncols=2, figsize=(10,5))\n",
    "ax1.plot(r_grid, np.real(E_center_dipole[:Nr]), label='FD real(E)')\n",
    "ax1.plot(r_grid, np.real(analytic_E_grid), '--', label='analytic real(E)')\n",
    "ax1.set_xlabel('r'); ax1.legend()\n",
    "ax2.plot(r_grid, np.imag(E_center_dipole[:Nr]), label='FD imag(E)')\n",
    "ax2.plot(r_grid, np.imag(analytic_E_grid), '--', label='analytic imag(E)')\n",
    "ax2.set_xlabel('r'); ax2.legend()\n",
    "plt.show()\n",
    "print(f'deviation starting around r={r_nonpml} due to pml')\n",
    "\n",
    "# Optional: expand to full circle for visualization\n",
    "E_center_dipole_full = expand_symmetric_field(E_center_dipole, n_sectors, Nr, mirror=True)\n",
    "print(\"Full circle visualization:\")\n",
    "plot_cplx_polar_field(E_center_dipole_full, phi_grid_full, r_grid)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "## Putting in structures\n",
    "\n",
    "Do a test with a non-vacuum structure and dipole in center.\n",
    "Structure in irreducible domain: annular pie with $\\phi \\in [0,\\pi/12]$,  $r \\in [1.0,1.3]$,  $\\chi=3+0.1i$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "## setup structure: ring at radius r_ring with thickness r_t\n",
    "r_ring = 1.0  # center radius of ring\n",
    "r_t = 0.3     # thickness of ring\n",
    "r_inner = r_ring\n",
    "r_outer = r_ring + r_t\n",
    "\n",
    "# Create radial mask for the ring\n",
    "chi_r_grid = np.zeros(Nr)\n",
    "chi_r_grid[int(r_inner/dr):int(r_outer/dr)] = 1.0\n",
    "\n",
    "## setup structure grid\n",
    "chi_phi_grid = np.zeros(Nphi_halfsector, dtype=complex)\n",
    "phi_start = 0  # start at phi = 0\n",
    "phi_end = int(np.round(Nphi_halfsector / 2))   # end at 1/4 of a sector, i.e. 1/2 of a half sector, or pi/12 in this case\n",
    "chi_phi_grid[phi_start:phi_end] = 1.0\n",
    "\n",
    "chi_grid = np.kron(chi_phi_grid, chi_r_grid)\n",
    "print(\"structure in irreducible domain:\")\n",
    "plot_real_polar_field(np.real(chi_grid), phi_grid_halfsector, r_grid, cmap='gist_yarg') # mask of structure over irreducible halfsector\n",
    "\n",
    "chi_grid_full = expand_symmetric_field(chi_grid, n_sectors, Nr, mirror=True)\n",
    "print(\"Full circle visualization:\")\n",
    "plot_real_polar_field(np.real(chi_grid_full), phi_grid_full, r_grid, cmap='gist_yarg')\n",
    "\n",
    "chi = 3.0 + 0.1j\n",
    "chi_grid *= chi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "E_struct = FDFD_halfsector.get_TM_field(J_center_dipole, chi_grid)\n",
    "plot_cplx_polar_field(E_struct, phi_grid_halfsector, r_grid)\n",
    "\n",
    "print(\"Full circle visualization:\")\n",
    "# important to specify mirror=True\n",
    "E_struct_full = expand_symmetric_field(E_struct, n_sectors, Nr, mirror=True)\n",
    "plot_cplx_polar_field(E_struct_full, phi_grid_full, r_grid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Verification with polar TM FDFD solver without symmetry\n",
    "\n",
    "# Create full-circle FDFD solver\n",
    "geo_full = PolarFDFDGeometry(Nphi_full, Nr, Npml, dr)\n",
    "FDFD_full = TM_Polar_FDFD(omega, geo_full)\n",
    "\n",
    "# Setup dipole source at center for full circle\n",
    "J_r_full = np.zeros(Nr)\n",
    "J_r_full[0] = 1.0 / (np.pi* dr**2)\n",
    "J_full = np.kron(np.ones(Nphi_full), J_r_full)\n",
    "\n",
    "# Setup structure for full circle (expand the symmetric structure)\n",
    "chi_full = expand_symmetric_field(chi_grid, n_sectors, Nr, \n",
    "                                  mirror=True)\n",
    "\n",
    "# Solve the full problem\n",
    "E_full_reference = FDFD_full.get_TM_field(J_full, chi_full)\n",
    "\n",
    "# Plot the full reference solution\n",
    "phi_grid_reference = np.linspace(0, 2*np.pi, Nphi_full, endpoint=False)\n",
    "plot_cplx_polar_field(E_full_reference, phi_grid_reference, r_grid)\n",
    "\n",
    "# Check specific points\n",
    "print(\"\\nField comparison at origin:\")\n",
    "print(f\"Full solver: {E_full_reference[0]:.6f}\")\n",
    "print(f\"Symmetric solver: {E_struct_full[0]:.6f}\")\n",
    "\n",
    "# Check a point in the middle of the domain\n",
    "mid_point = Nphi_full//2 * Nr + Nr//2\n",
    "print(f\"\\nField comparison at middle point:\")\n",
    "print(f\"Full solver: {E_full_reference[mid_point]:.6f}\")\n",
    "print(f\"Symmetric solver: {E_struct_full[mid_point]:.6f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "## Off-center dipoles\n",
    "One off-center dipole source in each rotational unit cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Source parameters\n",
    "source_r = 1.0  # radial position\n",
    "source_phi_fraction = 0.5  # fraction through the sector (0.5 = middle)\n",
    "\n",
    "ind_r = int(source_r / dr)\n",
    "ind_phi = int(source_phi_fraction * Nphi_halfsector)\n",
    "\n",
    "# Find grid indices\n",
    "ind_r = int(source_r / dr)\n",
    "ind_phi = int(source_phi_fraction * Nphi_halfsector)\n",
    "\n",
    "print(f\"Source at r = {r_grid[ind_r]:.3f}, phi = {phi_grid_halfsector[ind_phi]*180/np.pi:.1f}°\")\n",
    "\n",
    "# Setup source vector for one sector\n",
    "J_dipoles = np.zeros(Nphi_halfsector * Nr)\n",
    "\n",
    "# Calculate pixel area at this radius\n",
    "pixel_area = r_grid[ind_r] * dr * (2*np.pi / n_sectors / 2 / Nphi_halfsector)\n",
    "source_strength = 1.0 / pixel_area\n",
    "\n",
    "# Place source at the specified location in the sector\n",
    "J_dipoles[ind_phi * Nr + ind_r] = source_strength\n",
    "\n",
    "print(f\"Source strength: {source_strength:.2e}\")\n",
    "print(f\"Pixel area: {pixel_area:.6f}\")\n",
    "\n",
    "\n",
    "# Solve with vacuum Maxwell operator (no structure)\n",
    "E_dipoles = FDFD_halfsector.get_TM_field(J_dipoles)\n",
    "\n",
    "E_dipoles_full = expand_symmetric_field(E_dipoles, n_sectors, Nr, mirror=True)\n",
    "print(\"Full circle - showing D6 symmetry:\")\n",
    "plot_cplx_polar_field(E_dipoles_full, phi_grid_full, r_grid)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verification: compare with full solver\n",
    "\n",
    "J_dipoles_full = expand_symmetric_field(J_dipoles, n_sectors, Nr, mirror=True)\n",
    "E_dipoles_full_ref = FDFD_full.get_TM_field(J_dipoles_full)\n",
    "\n",
    "# Compare solutions\n",
    "print(\"Full solver reference:\")\n",
    "plot_cplx_polar_field(E_dipoles_full_ref, phi_grid_full, r_grid)\n",
    "\n",
    "# Quantitative comparison\n",
    "max_diff = np.max(np.abs(E_dipoles_full_ref - E_dipoles_full))\n",
    "rms_diff = np.sqrt(np.mean(np.abs(E_dipoles_full_ref - E_dipoles_full)**2))\n",
    "\n",
    "print(f\"\\nMaximum difference: {max_diff:.2e}\")\n",
    "print(f\"RMS difference: {rms_diff:.2e}\")\n",
    "print(f\"Relative RMS error: {rms_diff/np.sqrt(np.mean(np.abs(E_dipoles_full_ref)**2)):.2e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}