{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Polar TM Maxwell Solver Test / Tutorial with discrete rotational symmetry\n",
    "\n",
    "Sanity checks for correctness of the polar TM Maxwell solver with discrete rotational symmetry, which also demonstrates its basic usage. High level API is essentially similar to original TM_FDFD. NOTE: because of the polar coordinates, the FD pixels are no longer of unit size. "
   ]
  },
  {
   "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 12-fold discrete rotational symmetry, which allows us to use the n_sector feature of the polar solver to restrict computation within the 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 # computational domain 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: 12-fold rotational symmetry (30° sectors)\n",
    "n_sectors = 12\n",
    "\n",
    "assert Nphi_full % n_sectors == 0, \"Nphi must be divisible by n_sectors\"\n",
    "\n",
    "Nphi_sector = Nphi_full // n_sectors  # azimuthal points in one sector\n",
    "\n",
    "# geometry\n",
    "# Important: n_sectors=1 by default so make sure to specify otherwise\n",
    "geo_sector = PolarFDFDGeometry(Nphi_sector, Nr, Npml, dr, \n",
    "                               n_sectors=n_sectors)\n",
    "\n",
    "# FDFD solver\n",
    "FDFD_sector = TM_Polar_FDFD(omega, geo_sector)\n",
    "\n",
    "# Get coordinate grids for plotting later\n",
    "phi_grid_sector, r_grid, phi_grid_full = FDFD_sector.get_symmetric_grids()\n",
    "print(phi_grid_sector.shape)\n",
    "print(phi_grid_sector)\n",
    "print(FDFD_sector.phi_grid)"
   ]
  },
  {
   "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_sector instead of Nphi_full for the symmetric case\n",
    "J_center_dipole = np.kron(np.ones(Nphi_sector), 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_sector.get_TM_field(J_center_dipole)\n",
    "plot_cplx_polar_field(E_center_dipole, phi_grid_sector, 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",
    "# expand to full circle for visualization\n",
    "E_center_dipole_full = expand_symmetric_field(E_center_dipole, n_sectors, Nr)\n",
    "print(\"Full circle visualization:\")\n",
    "plot_cplx_polar_field(E_center_dipole_full, phi_grid_full, r_grid)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "## Putting in structures\n",
    "\n",
    "Now try a non-vacuum structure and dipole in center.\n",
    "rotational unit cell in $\\phi \\in [0, \\pi / 6)$: annular pie with $\\phi \\in [\\pi/24,\\pi/8]$,  $r \\in [1.0,1.5]$,  $\\chi=3+0.1i$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "r_ring_inner = 1.0\n",
    "r_ring_outer =1.5\n",
    "\n",
    "# Create radial mask for the ring\n",
    "chi_r_grid = np.zeros(Nr)\n",
    "chi_r_grid[int(r_ring_inner/dr):int(r_ring_outer/dr)] = 1.0\n",
    "\n",
    "## setup structure grid\n",
    "chi_phi_grid = np.zeros(Nphi_sector, dtype=complex)\n",
    "phi_start = int(Nphi_sector * 0.25)  # start at 25% of sector\n",
    "phi_end = int(Nphi_sector * 0.75)    # end at 75% of sector\n",
    "chi_phi_grid[phi_start:phi_end] = 1.0\n",
    "chi_grid = np.kron(chi_phi_grid, chi_r_grid)\n",
    "plot_real_polar_field(np.real(chi_grid), phi_grid_sector, r_grid, cmap='gist_yarg') # mask of structure\n",
    "\n",
    "chi_grid_full = expand_symmetric_field(chi_grid, n_sectors, Nr)\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_sector.get_TM_field(J_center_dipole, chi_grid)\n",
    "# plot_cplx_polar_field(E, phi_grid_sector, r_grid)\n",
    "\n",
    "print(\"Full circle visualization:\")\n",
    "E_struct_full = expand_symmetric_field(E_struct, n_sectors, Nr)\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",
    "\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",
    "# Find grid indices\n",
    "ind_r = int(source_r / dr)\n",
    "ind_phi = int(source_phi_fraction * Nphi_sector)\n",
    "\n",
    "print(f\"Source at r = {r_grid[ind_r]:.3f}, phi = {phi_grid_sector[ind_phi]*180/np.pi:.1f}°\")\n",
    "\n",
    "# Setup source vector for one sector\n",
    "J_dipoles = np.zeros(Nphi_sector * Nr)\n",
    "\n",
    "# Calculate pixel area at this radius\n",
    "pixel_area = r_grid[ind_r] * dr * (2*np.pi / n_sectors / Nphi_sector)\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",
    "# Solve with vacuum Maxwell operator (no structure)\n",
    "E_dipoles = FDFD_sector.get_TM_field(J_dipoles)\n",
    "\n",
    "E_dipoles_full = expand_symmetric_field(E_dipoles, n_sectors, Nr)\n",
    "print(\"Full circle - showing 12-fold 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",
    "J_dipoles_full = expand_symmetric_field(J_dipoles, n_sectors, Nr)\n",
    "\n",
    "# Solve full problem\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}\")\n",
    "\n",
    "# Check field at source location\n",
    "source_idx_sector = ind_phi * Nr + ind_r\n",
    "source_idx_full = ind_phi * Nr + ind_r  # first sector\n",
    "print(f\"\\nField at source location:\")\n",
    "print(f\"Symmetric solver: {E_dipoles[source_idx_sector]:.6f}\")\n",
    "print(f\"Full solver: {E_dipoles_full_ref[source_idx_full]:.6f}\")\n",
    "print(f\"Difference: {abs(E_dipoles[source_idx_sector] - E_dipoles_full_ref[source_idx_full]):.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
}