{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Cavity LDOS Bounds: Cartesian vs Polar Coordinates\n",
    "\n",
    "This notebook demonstrates how to compute limits using the Polar coordinate formulation in Dolphindes. To verify consistency, we also compare the results with the Cartesian coordinate formulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import numpy as np\n",
    "import scipy.sparse as sp\n",
    "import matplotlib.pyplot as plt\n",
    "import sys, time, os\n",
    "sys.path.append(\"../../\")\n",
    "\n",
    "from dolphindes import photonics, geometry\n",
    "from dolphindes.cvxopt import gcd\n",
    "from dolphindes.maxwell import plot_real_polar_field, plot_cplx_polar_field, expand_symmetric_field, TM_Polar_FDFD"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Problem Setup\n",
    "\n",
    "We consider a 2D TM problem. We want to maximize the LDOS at the center of a circular design region. \n",
    "The design region is an annulus $R_{in} < r < R_{out}$.\n",
    "We place a point source at the center."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "wavelength = 1.0\n",
    "omega = 2 * np.pi / wavelength\n",
    "chi = 4 + 1e-4j\n",
    "\n",
    "# Geometry parameters\n",
    "R_inner = 0.2\n",
    "R_outer = 0.6\n",
    "Npml = 10\n",
    "gpr = 50\n",
    "dl = 1/gpr"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "## 1. Cartesian Coordinates\n",
    "\n",
    "This should look pretty much identical to prior tutorials. We are just getting the cartesian value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Grid setup\n",
    "L_domain = 2 * (R_outer + 0.5)\n",
    "Nx = Ny = int(L_domain / dl) + 2 * Npml\n",
    "\n",
    "# Source at center\n",
    "ji_cart = np.zeros((Nx, Ny), dtype=complex)\n",
    "cx, cy = Nx // 2, Ny // 2\n",
    "ji_cart[cx, cy] = 1.0 / dl**2\n",
    "\n",
    "# Design mask (Annulus)\n",
    "x = (np.arange(Nx) - cx) * dl\n",
    "y = (np.arange(Ny) - cy) * dl\n",
    "X, Y = np.meshgrid(x, y, indexing='ij')\n",
    "R = np.sqrt(X**2 + Y**2)\n",
    "des_mask_cart = (R >= R_inner) & (R <= R_outer)\n",
    "ndof = np.sum(des_mask_cart)\n",
    "\n",
    "chi_bg_cart = np.zeros((Nx, Ny), dtype=complex)\n",
    "plt.imshow(des_mask_cart.T + np.real(ji_cart.T).astype(bool), origin='lower', cmap='Greys')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solver Setup\n",
    "geo_cart = geometry.CartesianFDFDGeometry(Nx, Ny, Npml, Npml, dl, dl)\n",
    "prob_cart = photonics.Photonics_TM_FDFD(\n",
    "    omega, geo_cart, chi, des_mask_cart, ji_cart, chi_background=chi_bg_cart, sparseQCQP=True\n",
    ")\n",
    "\n",
    "ei_cart = prob_cart.get_ei(ji_cart, update=True)\n",
    "vac_ldos_cart = -np.sum(0.5 * np.real(ji_cart.conj() * ei_cart) * dl**2)\n",
    "plt.imshow(ei_cart.T.real, origin='lower', cmap='RdBu')\n",
    "plt.colorbar()\n",
    "print(f\"Cartesian Vacuum LDOS: {vac_ldos_cart:.4f}\")\n",
    "\n",
    "# Objective: Maximize LDOS\n",
    "ei_des_cart = ei_cart[des_mask_cart]\n",
    "s0_cart = -0.25 * 1j * omega * ei_des_cart.conj()\n",
    "ndof = np.sum(des_mask_cart)\n",
    "A0_cart = sp.csc_array(np.zeros((ndof, ndof), dtype=complex))\n",
    "prob_cart.set_objective(A0=A0_cart, s0=s0_cart, c0=vac_ldos_cart, denseToSparse=True)\n",
    "\n",
    "# QCQP Setup\n",
    "prob_cart.setup_QCQP(Pdiags='global', verbose=1)\n",
    "res_cart = prob_cart.bound_QCQP(method='bfgs')\n",
    "print(f\"Cartesian Bound: {res_cart[0]:.4f}\")\n",
    "print(f\"Cartesian Enhancement Bound: {res_cart[0]/vac_ldos_cart:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "## 2. Polar Coordinates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "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 = 3.0 # center circle radius\n",
    "w_pml = 0.5 # surrounding pml thickness\n",
    "R_tot = R_nonpml + w_pml # total computational domain radius\n",
    "R_i = R_inner # inner radius of the cavity\n",
    "R_o = R_outer # outer radius of the cavity\n",
    "assert R_o < R_nonpml, \"Outer radius must be within non-PML region\"\n",
    "\n",
    "gpr = 20\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",
    "Nr_i = int(np.round(R_i / dr)) # inner radius grid point\n",
    "Nr_o = int(np.round(R_o / dr)) # outer radius grid point\n",
    "Nphi = 100\n",
    "n_sectors = 1\n",
    "Nphi_sector = int(Nphi/n_sectors)  # azimuthal points in one sector\n",
    "\n",
    "assert Nphi % n_sectors == 0, \"Nphi must be divisible by n_sectors\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Normalize source so integrated current in the central pixel = 1\n",
    "J_r = np.zeros(Nr)\n",
    "J_r[0] = 1.0 / (np.pi * dr**2)  # central pixel current density = 1/area\n",
    "J = np.kron(np.ones(Nphi_sector), J_r)\n",
    "J_full = expand_symmetric_field(J, n_sectors, Nr)\n",
    "phi_grid = np.linspace(0, 2*np.pi, Nphi, endpoint=False)\n",
    "r_grid = (np.arange(Nr) + 0.5) * dr\n",
    "\n",
    "# Design Mask\n",
    "des_mask_polar = np.zeros((Nr, Nphi), dtype=bool)\n",
    "for i, r in enumerate(r_grid):\n",
    "    if R_inner <= r <= R_outer:\n",
    "        des_mask_polar[i, :] = True\n",
    "ndof = np.sum(des_mask_polar)\n",
    "\n",
    "# Use radial spacing 'dr' for polar geometry\n",
    "geo_polar = geometry.PolarFDFDGeometry(Nphi, Nr, Npml, dr)\n",
    "\n",
    "# Solver Setup\n",
    "prob_polar = photonics.Photonics_TM_FDFD(\n",
    "    omega, geo_polar, chi, des_mask_polar, J, sparseQCQP=True\n",
    ")\n",
    "\n",
    "prob_polar.setup_EM_solver()\n",
    "phi_grid_sector, r_grid, phi_grid_full = prob_polar.EM_solver.get_symmetric_grids()\n",
    "areas = prob_polar.geometry.get_pixel_areas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "ei = prob_polar.get_ei()\n",
    "plot_cplx_polar_field(np.real(ei), phi_grid, r_grid)\n",
    "# Evaluate vacuum LDOS (per-sector J, multiply by n_sectors)\n",
    "vac_ldos_polar = -0.5 * np.real(J.conj().dot(areas * ei)) * n_sectors\n",
    "print(f'vac_ldos_polar: {vac_ldos_polar}')\n",
    "# print(f'vac_ldos_cart: {vac_ldos_cart}')\n",
    "\n",
    "plot_real_polar_field(prob_polar.des_mask, phi_grid_sector, r_grid, cmap='Greys')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "ei_des_polar = ei[prob_polar.des_mask.flatten(order='F')]\n",
    "s0_polar = -0.25 * 1j * omega * ei_des_polar.conj()\n",
    "ndof = np.sum(prob_polar.des_mask)\n",
    "A0_polar = sp.csc_array(np.zeros((ndof, ndof), dtype=complex))\n",
    "prob_polar.set_objective(A0=A0_polar, s0=s0_polar, c0=vac_ldos_polar, denseToSparse=True)\n",
    "\n",
    "prob_polar.setup_QCQP(Pdiags='global', verbose=1)\n",
    "res_polar = prob_polar.bound_QCQP(method='bfgs')\n",
    "print(f\"Polar Bound: {res_polar[0]:.4f}\")\n",
    "print(f\"Polar Enhancement Bound: {res_polar[0]/vac_ldos_polar:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "## 3. Comparison and GCD\n",
    "\n",
    "We expect the enhancement bounds (Purcell factors) to be similar, though differences in discretization near the singularity (source) will cause some discrepancy in the absolute vacuum LDOS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Cartesian Enhancement: {res_cart[0]/vac_ldos_cart:.4f}\")\n",
    "print(f\"Polar Enhancement:     {res_polar[0]/vac_ldos_polar:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "### Running GCD\n",
    "Now we run General Constraint Descent to tighten the bounds for both cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "gcd_params = gcd.GCDHyperparameters(\n",
    "    max_proj_cstrt_num=16,\n",
    "    max_gcd_iter_num=20,\n",
    "    gcd_iter_period=5,\n",
    "    gcd_tol=1e-3\n",
    ")\n",
    "\n",
    "print(\"Running GCD on Cartesian...\")\n",
    "t0 = time.time()\n",
    "prob_cart.QCQP.run_gcd(gcd_params)\n",
    "print(f\"Cartesian GCD Time: {time.time()-t0:.2f}s\")\n",
    "print(f\"Cartesian GCD Bound: {prob_cart.QCQP.current_dual:.4f}\")\n",
    "print(f\"Cartesian GCD Enhancement: {prob_cart.QCQP.current_dual/vac_ldos_cart:.4f}\")\n",
    "\n",
    "print(\"\\nRunning GCD on Polar...\")\n",
    "t0 = time.time()\n",
    "prob_polar.QCQP.run_gcd(gcd_params)\n",
    "print(f\"Polar GCD Time: {time.time()-t0:.2f}s\")\n",
    "print(f\"Polar GCD Bound: {prob_polar.QCQP.current_dual:.4f}\")\n",
    "print(f\"Polar GCD Enhancement: {prob_polar.QCQP.current_dual/vac_ldos_polar:.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "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
}