{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "828187f9-e15f-4da1-b6bc-d94808ce7f2f",
   "metadata": {
    "papermill": {
     "duration": 0.009485,
     "end_time": "2023-09-03T04:40:07.685281",
     "exception": false,
     "start_time": "2023-09-03T04:40:07.675796",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Thin film optimization\n",
    "> Let's optimize a thin film..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07171d66-a595-4412-b975-681ffb558ab5",
   "metadata": {
    "papermill": {
     "duration": 0.024409,
     "end_time": "2023-09-03T04:40:07.797248",
     "exception": false,
     "start_time": "2023-09-03T04:40:07.772839",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "> This notebook was contributed by [simbilod](https://github.com/simbilod)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1f8afd7",
   "metadata": {
    "papermill": {
     "duration": 1.978445,
     "end_time": "2023-09-03T04:40:09.797090",
     "exception": false,
     "start_time": "2023-09-03T04:40:07.818645",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.example_libraries.optimizers as opt\n",
    "import jax.numpy as jnp\n",
    "import matplotlib.pyplot as plt\n",
    "import sax\n",
    "import tqdm.notebook as tqdm\n",
    "from tmm import coh_tmm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c09fa16c-05fb-4120-80ae-17d64453fcb1",
   "metadata": {
    "papermill": {
     "duration": 0.008116,
     "end_time": "2023-09-03T04:40:09.814063",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.805947",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "In this notebook, we apply SAX to thin-film optimization and show how it can be used for wavelength-dependent parameter optimization."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "508b27a5-7fee-409b-b381-988ca07f5b35",
   "metadata": {
    "papermill": {
     "duration": 0.007858,
     "end_time": "2023-09-03T04:40:09.830186",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.822328",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "The language of transfer/scatter matrices is commonly used to calculate optical properties of thin-films. Many specialized methods exist for their optimization. However, SAX can be useful to cut down on developer time by circumventing the need to manually take gradients of complicated or often-changed objective functions, and by generating efficient code from simple syntax. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77e24e88",
   "metadata": {
    "papermill": {
     "duration": 0.008462,
     "end_time": "2023-09-03T04:40:09.847085",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.838623",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Dielectric mirror Fabry-Pérot\n",
    "\n",
    "Consider a stack composed of only two materials, $n_A$ and $n_B$. Two types of transfer matrices characterize wave propagation in the system : interfaces described by Fresnel's equations, and propagation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa0f57dc",
   "metadata": {
    "papermill": {
     "duration": 0.008496,
     "end_time": "2023-09-03T04:40:09.864181",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.855685",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "For the two-material stack, this leads to 4 scatter matrices coefficients. Through reciprocity they can be constructed out of two independent ones :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0d3a8eb",
   "metadata": {
    "papermill": {
     "duration": 0.017434,
     "end_time": "2023-09-03T04:40:09.889783",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.872349",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def fresnel_mirror_ij(ni=1.0, nj=1.0):\n",
    "    \"\"\"Model a (fresnel) interface between twoo refractive indices\n",
    "\n",
    "    Args:\n",
    "        ni: refractive index of the initial medium\n",
    "        nf: refractive index of the final\n",
    "    \"\"\"\n",
    "    r_fresnel_ij = (ni - nj) / (ni + nj)  # i->j reflection\n",
    "    t_fresnel_ij = 2 * ni / (ni + nj)  # i->j transmission\n",
    "    r_fresnel_ji = -r_fresnel_ij  # j -> i reflection\n",
    "    t_fresnel_ji = (1 - r_fresnel_ij ** 2) / t_fresnel_ij  # j -> i transmission\n",
    "    sdict = {\n",
    "        (\"in\", \"in\"): r_fresnel_ij,\n",
    "        (\"in\", \"out\"): t_fresnel_ij,\n",
    "        (\"out\", \"in\"): t_fresnel_ji,\n",
    "        (\"out\", \"out\"): r_fresnel_ji,\n",
    "    }\n",
    "    return sdict\n",
    "\n",
    "\n",
    "def propagation_i(ni=1.0, di=0.5, wl=0.532):\n",
    "    \"\"\"Model the phase shift acquired as a wave propagates through medium A\n",
    "\n",
    "    Args:\n",
    "        ni: refractive index of medium (at wavelength wl)\n",
    "        di: [μm] thickness of layer\n",
    "        wl: [μm] wavelength\n",
    "    \"\"\"\n",
    "    prop_i = jnp.exp(1j * 2 * jnp.pi * ni * di / wl)\n",
    "    sdict = {\n",
    "        (\"in\", \"out\"): prop_i,\n",
    "        (\"out\", \"in\"): prop_i,\n",
    "    }\n",
    "    return sdict"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "007b1181",
   "metadata": {
    "papermill": {
     "duration": 0.008595,
     "end_time": "2023-09-03T04:40:09.906877",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.898282",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "A resonant cavity can be formed when a high index region is surrounded by low-index region :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14f86c1e",
   "metadata": {
    "papermill": {
     "duration": 0.835424,
     "end_time": "2023-09-03T04:40:10.751513",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.916089",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "dielectric_fabry_perot, _ = sax.circuit(\n",
    "    netlist={\n",
    "        \"instances\": {\n",
    "            \"air_B\": fresnel_mirror_ij,\n",
    "            \"B\": propagation_i,\n",
    "            \"B_air\": fresnel_mirror_ij,\n",
    "        },\n",
    "        \"connections\": {\n",
    "            \"air_B,out\": \"B,in\",\n",
    "            \"B,out\": \"B_air,in\",\n",
    "        },\n",
    "        \"ports\": {\n",
    "            \"in\": \"air_B,in\",\n",
    "            \"out\": \"B_air,out\",\n",
    "        },\n",
    "    },\n",
    "    backend='fg',\n",
    ")\n",
    "\n",
    "settings = sax.get_settings(dielectric_fabry_perot)\n",
    "settings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9285d92",
   "metadata": {
    "papermill": {
     "duration": 0.008388,
     "end_time": "2023-09-03T04:40:10.768755",
     "exception": false,
     "start_time": "2023-09-03T04:40:10.760367",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Let's choose $n_A = 1$, $n_B = 2$, $d_B = 1000$ nm, and compute over the visible spectrum :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ee9d40d",
   "metadata": {
    "papermill": {
     "duration": 0.052548,
     "end_time": "2023-09-03T04:40:10.830049",
     "exception": false,
     "start_time": "2023-09-03T04:40:10.777501",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "settings = sax.copy_settings(settings)\n",
    "settings[\"air_B\"][\"nj\"] = 2.0\n",
    "settings[\"B\"][\"ni\"] = 2.0\n",
    "settings[\"B_air\"][\"ni\"] = 2.0\n",
    "\n",
    "wls = jnp.linspace(0.380, 0.750, 200)\n",
    "settings = sax.update_settings(settings, wl=wls)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ace74e4",
   "metadata": {
    "papermill": {
     "duration": 0.008982,
     "end_time": "2023-09-03T04:40:10.847740",
     "exception": false,
     "start_time": "2023-09-03T04:40:10.838758",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Compute transmission and reflection, and compare to another package's results (https://github.com/sbyrnes321/tmm) :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "999d2eee",
   "metadata": {
    "papermill": {
     "duration": 0.730319,
     "end_time": "2023-09-03T04:40:11.586961",
     "exception": false,
     "start_time": "2023-09-03T04:40:10.856642",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "sdict = dielectric_fabry_perot(**settings)\n",
    "\n",
    "transmitted = sdict[\"in\", \"out\"]\n",
    "reflected = sdict[\"in\", \"in\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e0d4550",
   "metadata": {
    "papermill": {
     "duration": 1.717302,
     "end_time": "2023-09-03T04:40:13.310285",
     "exception": false,
     "start_time": "2023-09-03T04:40:11.592983",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# tmm syntax (https://github.com/sbyrnes321/tmm)\n",
    "d_list = [jnp.inf, 0.500, jnp.inf]\n",
    "n_list = [1, 2, 1]\n",
    "# initialize lists of y-values to plot\n",
    "rnorm = []\n",
    "tnorm = []\n",
    "Tnorm = []\n",
    "Rnorm = []\n",
    "for l in wls:\n",
    "    rnorm.append(coh_tmm(\"s\", n_list, d_list, 0, l)[\"r\"])\n",
    "    tnorm.append(coh_tmm(\"s\", n_list, d_list, 0, l)[\"t\"])\n",
    "    Tnorm.append(coh_tmm(\"s\", n_list, d_list, 0, l)[\"T\"])\n",
    "    Rnorm.append(coh_tmm(\"s\", n_list, d_list, 0, l)[\"R\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43b55a4d",
   "metadata": {
    "papermill": {
     "duration": 0.570107,
     "end_time": "2023-09-03T04:40:13.889434",
     "exception": false,
     "start_time": "2023-09-03T04:40:13.319327",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.scatter(wls * 1e3, jnp.real(transmitted), label=\"t SAX\")\n",
    "plt.plot(wls * 1e3, jnp.real(jnp.array(tnorm)), \"k\", label=\"t tmm\")\n",
    "plt.scatter(wls * 1e3, jnp.real(reflected), label=\"r SAX\")\n",
    "plt.plot(wls * 1e3, jnp.real(jnp.array(rnorm)), \"k--\", label=\"r tmm\")\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.ylabel(\"Transmitted and reflected amplitude\")\n",
    "plt.legend(loc=\"upper right\")\n",
    "plt.title(\"Real part\")\n",
    "plt.show()\n",
    "\n",
    "plt.scatter(wls * 1e3, jnp.imag(transmitted), label=\"t SAX\")\n",
    "plt.plot(wls * 1e3, jnp.imag(jnp.array(tnorm)), \"k\", label=\"t tmm\")\n",
    "plt.scatter(wls * 1e3, jnp.imag(reflected), label=\"r SAX\")\n",
    "plt.plot(wls * 1e3, jnp.imag(jnp.array(rnorm)), \"k--\", label=\"r tmm\")\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.ylabel(\"Transmitted and reflected amplitude\")\n",
    "plt.legend(loc=\"upper right\")\n",
    "plt.title(\"Imaginary part\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae720d84",
   "metadata": {
    "papermill": {
     "duration": 0.009551,
     "end_time": "2023-09-03T04:40:13.909669",
     "exception": false,
     "start_time": "2023-09-03T04:40:13.900118",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "In terms of powers, we get the following. Due to the reflections at the interfaces, resonant behaviour is observed, with evenly-spaced maxima/minima in wavevector space :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3b974a5",
   "metadata": {
    "papermill": {
     "duration": 0.449429,
     "end_time": "2023-09-03T04:40:14.365821",
     "exception": false,
     "start_time": "2023-09-03T04:40:13.916392",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.scatter(2 * jnp.pi / wls, jnp.abs(transmitted) ** 2, label=\"T SAX\")\n",
    "plt.plot(2 * jnp.pi / wls, Tnorm, \"k\", label=\"T tmm\")\n",
    "plt.scatter(2 * jnp.pi / wls, jnp.abs(reflected) ** 2, label=\"R SAX\")\n",
    "plt.plot(2 * jnp.pi / wls, Rnorm, \"k--\", label=\"R tmm\")\n",
    "plt.vlines(jnp.arange(3, 6) * jnp.pi / (2 * 0.5), ymin=0, ymax=1, color=\"k\", linestyle=\"--\", label=\"m$\\pi$/nd\")\n",
    "plt.xlabel(\"k = 2$\\pi$/λ [1/nm]\")\n",
    "plt.ylabel(\"Transmitted and reflected intensities\")\n",
    "plt.legend(loc=\"upper right\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b83a4463",
   "metadata": {
    "papermill": {
     "duration": 0.007267,
     "end_time": "2023-09-03T04:40:14.380690",
     "exception": false,
     "start_time": "2023-09-03T04:40:14.373423",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### Optimization test\n",
    "\n",
    "Let's attempt to minimize transmission at 500 nm by varying thickness."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d497768f",
   "metadata": {
    "papermill": {
     "duration": 0.013161,
     "end_time": "2023-09-03T04:40:14.401176",
     "exception": false,
     "start_time": "2023-09-03T04:40:14.388015",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def loss(thickness):\n",
    "    settings = sax.update_settings(sax.get_settings(dielectric_fabry_perot), wl=0.5)\n",
    "    settings[\"B\"][\"di\"] = thickness\n",
    "    settings[\"air_B\"][\"nj\"] = 2.0\n",
    "    settings[\"B\"][\"ni\"] = 2.0\n",
    "    settings[\"B_air\"][\"ni\"] = 2.0\n",
    "    sdict = dielectric_fabry_perot(**settings)\n",
    "    return jnp.abs(sdict[\"in\", \"out\"]) ** 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e8d21cc",
   "metadata": {
    "papermill": {
     "duration": 0.012367,
     "end_time": "2023-09-03T04:40:14.420683",
     "exception": false,
     "start_time": "2023-09-03T04:40:14.408316",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "grad = jax.jit(jax.grad(loss))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18af200b",
   "metadata": {
    "papermill": {
     "duration": 0.013353,
     "end_time": "2023-09-03T04:40:14.441569",
     "exception": false,
     "start_time": "2023-09-03T04:40:14.428216",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "initial_thickness = 0.5\n",
    "optim_init, optim_update, optim_params = opt.adam(step_size=0.01)\n",
    "optim_state = optim_init(initial_thickness)\n",
    "\n",
    "\n",
    "def train_step(step, optim_state):\n",
    "    thickness = optim_params(optim_state)\n",
    "    lossvalue = loss(thickness)\n",
    "    gradvalue = grad(thickness)\n",
    "    optim_state = optim_update(step, gradvalue, optim_state)\n",
    "    return lossvalue, optim_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98bcb50b",
   "metadata": {
    "papermill": {
     "duration": 2.595336,
     "end_time": "2023-09-03T04:40:17.044424",
     "exception": false,
     "start_time": "2023-09-03T04:40:14.449088",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "range_ = tqdm.trange(100)\n",
    "for step in range_:\n",
    "    lossvalue, optim_state = train_step(step, optim_state)\n",
    "    range_.set_postfix(loss=f\"{lossvalue:.6f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13d16ace",
   "metadata": {
    "papermill": {
     "duration": 0.013879,
     "end_time": "2023-09-03T04:40:17.066213",
     "exception": false,
     "start_time": "2023-09-03T04:40:17.052334",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "thickness = optim_params(optim_state)\n",
    "thickness"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28073b3a",
   "metadata": {
    "papermill": {
     "duration": 0.33158,
     "end_time": "2023-09-03T04:40:17.410200",
     "exception": false,
     "start_time": "2023-09-03T04:40:17.078620",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "settings = sax.update_settings(sax.get_settings(dielectric_fabry_perot), wl=wls)\n",
    "settings[\"B\"][\"di\"] = thickness\n",
    "settings[\"air_B\"][\"nj\"] = 2.0\n",
    "settings[\"B\"][\"ni\"] = 2.0\n",
    "settings[\"B_air\"][\"ni\"] = 2.0\n",
    "sdict = dielectric_fabry_perot(**settings)\n",
    "detected = sdict[\"in\", \"out\"]\n",
    "\n",
    "plt.plot(wls * 1e3, jnp.abs(transmitted) ** 2, label=\"Before (500 nm)\")\n",
    "plt.plot(wls * 1e3, jnp.abs(detected) ** 2, label=f\"After ({thickness*1e3:.0f} nm)\")\n",
    "plt.vlines(0.5 * 1e3, 0.6, 1, \"k\", linestyle=\"--\")\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.ylabel(\"Transmitted intensity\")\n",
    "plt.legend(loc=\"lower right\")\n",
    "plt.title(\"Thickness optimization\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa41c516",
   "metadata": {
    "papermill": {
     "duration": 0.008144,
     "end_time": "2023-09-03T04:40:17.426950",
     "exception": false,
     "start_time": "2023-09-03T04:40:17.418806",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## General Fabry-Pérot étalon\n",
    "\n",
    "We reuse the propagation matrix above, and instead of simple interface matrices, model Fabry-Pérot mirrors as general lossless reciprocal scatter matrices :\n",
    "\n",
    "$$ \\left(\\begin{array}{c} \n",
    "E_t \\\\\n",
    "E_r\n",
    "\\end{array}\\right) = E_{out} = SE_{in} = \\left(\\begin{array}{cc} \n",
    "t & r \\\\\n",
    "r & t\n",
    "\\end{array}\\right) \\left(\\begin{array}{c} \n",
    "E_0 \\\\\n",
    "0\n",
    "\\end{array}\\right) $$\n",
    "\n",
    "For lossless reciprocal systems, we further have the requirements\n",
    "\n",
    "$$ |t|^2 + |r|^2 = 1 $$\n",
    "\n",
    "and\n",
    "\n",
    "$$ \\angle t - \\angle r = \\pm \\pi/2 $$\n",
    "\n",
    "The general Fabry-Pérot cavity is analytically described by :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b0a7441",
   "metadata": {
    "papermill": {
     "duration": 0.0146,
     "end_time": "2023-09-03T04:40:17.449849",
     "exception": false,
     "start_time": "2023-09-03T04:40:17.435249",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def airy_t13(t12, t23, r21, r23, wl, d=1.0, n=1.0):\n",
    "    \"\"\"General Fabry-Pérot transmission transfer function (Airy formula)\n",
    "\n",
    "    Args:\n",
    "        t12 and r12 : S-parameters of the first mirror\n",
    "        t23 and r23 : S-parameters of the second mirror\n",
    "        wl : wavelength\n",
    "        d : gap between the two mirrors (in units of wavelength)\n",
    "        n : index of the gap between the two mirrors\n",
    "\n",
    "    Returns:\n",
    "        t13 : complex transmission amplitude of the mirror-gap-mirror system\n",
    "\n",
    "    Note:\n",
    "        Each mirror is assumed to be lossless and reciprocal : tij = tji, rij = rji\n",
    "    \"\"\"\n",
    "    phi = n * 2 * jnp.pi * d / wl\n",
    "    return t12 * t23 * jnp.exp(-1j * phi) / (1 - r21 * r23 * jnp.exp(-2j * phi))\n",
    "\n",
    "\n",
    "def airy_r13(t12, t23, r21, r23, wl, d=1.0, n=1.0):\n",
    "    \"\"\"General Fabry-Pérot reflection transfer function (Airy formula)\n",
    "\n",
    "    Args:\n",
    "        t12 and r12 : S-parameters of the first mirror\n",
    "        t23 and r23 : S-parameters of the second mirror\n",
    "        wl : wavelength\n",
    "        d : gap between the two mirrors (in units of wavelength)\n",
    "        n : index of the gap between the two mirrors\n",
    "\n",
    "    Returns:\n",
    "        r13 : complex reflection amplitude of the mirror-gap-mirror system\n",
    "\n",
    "    Note:\n",
    "        Each mirror is assumed to be lossless and reciprocal : tij = tji, rij = rji\n",
    "    \"\"\"\n",
    "    phi = n * 2 * jnp.pi * d / wl\n",
    "    return r21 + t12 * t12 * r23 * jnp.exp(-2j * phi) / (1 - r21 * r23 * jnp.exp(-2j * phi))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cc808ed",
   "metadata": {
    "papermill": {
     "duration": 0.007765,
     "end_time": "2023-09-03T04:40:17.465795",
     "exception": false,
     "start_time": "2023-09-03T04:40:17.458030",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We need to implement the relationship between $t$ and $r$ for lossless reciprocal mirrors. The design parameter will be the amplitude and phase of the tranmission coefficient. The reflection coefficient is then fully determined :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62c11bb6",
   "metadata": {
    "papermill": {
     "duration": 0.013439,
     "end_time": "2023-09-03T04:40:17.487110",
     "exception": false,
     "start_time": "2023-09-03T04:40:17.473671",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def t_complex(t_amp, t_ang):\n",
    "    return t_amp * jnp.exp(-1j * t_ang)\n",
    "\n",
    "\n",
    "def r_complex(t_amp, t_ang):\n",
    "    r_amp = jnp.sqrt((1.0 - t_amp ** 2))\n",
    "    r_ang = t_ang - jnp.pi / 2\n",
    "    return r_amp * jnp.exp(-1j * r_ang)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a865ea71",
   "metadata": {
    "papermill": {
     "duration": 0.008148,
     "end_time": "2023-09-03T04:40:17.503101",
     "exception": false,
     "start_time": "2023-09-03T04:40:17.494953",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Let's see the expected result for half-mirrors :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0a6b35d",
   "metadata": {
    "papermill": {
     "duration": 0.687002,
     "end_time": "2023-09-03T04:40:18.198214",
     "exception": false,
     "start_time": "2023-09-03T04:40:17.511212",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "t_initial = jnp.sqrt(0.5)\n",
    "d_gap = 2.0\n",
    "n_gap = 1.0\n",
    "r_initial = r_complex(t_initial, 0.0)\n",
    "\n",
    "wls = jnp.linspace(0.38, 0.78, 500)\n",
    "\n",
    "T_analytical_initial = jnp.abs(airy_t13(t_initial, t_initial, r_initial, r_initial, wls, d=d_gap, n=n_gap)) ** 2\n",
    "R_analytical_initial = jnp.abs(airy_r13(t_initial, t_initial, r_initial, r_initial, wls, d=d_gap, n=n_gap)) ** 2 \n",
    "\n",
    "plt.title(f\"t={t_initial:1.3f}, d={d_gap} nm, n={n_gap}\")\n",
    "plt.plot(2 * jnp.pi / wls, T_analytical_initial, label=\"T\")\n",
    "plt.plot(2 * jnp.pi / wls, R_analytical_initial, label=\"R\")\n",
    "plt.vlines(jnp.arange(6, 11) * jnp.pi / 2.0, ymin=0, ymax=1, color=\"k\", linestyle=\"--\", label=\"m$\\pi$/nd\")\n",
    "plt.xlabel(\"k = 2$\\pi$/$\\lambda$ [1/nm]\")\n",
    "plt.ylabel(\"Power (units of input)\")\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "plt.title(f\"t={t_initial:1.3f}, d={d_gap} nm, n={n_gap}\")\n",
    "plt.plot(wls * 1e3, T_analytical_initial, label=\"T\")\n",
    "plt.plot(wls * 1e3, R_analytical_initial, label=\"R\")\n",
    "plt.xlabel(\"$\\lambda$ (nm)\")\n",
    "plt.ylabel(\"Power (units of input)\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7c66218",
   "metadata": {
    "papermill": {
     "duration": 0.078716,
     "end_time": "2023-09-03T04:40:18.292007",
     "exception": false,
     "start_time": "2023-09-03T04:40:18.213291",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Is power conserved? (to within 0.1%)\n",
    "assert jnp.isclose(R_analytical_initial + T_analytical_initial, 1, 0.001).all()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e91cd0dd",
   "metadata": {
    "papermill": {
     "duration": 0.009808,
     "end_time": "2023-09-03T04:40:18.311545",
     "exception": false,
     "start_time": "2023-09-03T04:40:18.301737",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Now let's do the same with SAX by defining new elements :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80c784a5",
   "metadata": {
    "papermill": {
     "duration": 0.039666,
     "end_time": "2023-09-03T04:40:18.361150",
     "exception": false,
     "start_time": "2023-09-03T04:40:18.321484",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def mirror(t_amp=0.5**0.5, t_ang=0.0):\n",
    "    r_complex_val = r_complex(t_amp, t_ang)\n",
    "    t_complex_val = t_complex(t_amp, t_ang)\n",
    "    sdict = {\n",
    "        (\"in\", \"in\"): r_complex_val,\n",
    "        (\"in\", \"out\"): t_complex_val,\n",
    "        (\"out\", \"in\"): t_complex_val,  # (1 - r_complex_val**2)/t_complex_val, # t_ji\n",
    "        (\"out\", \"out\"): r_complex_val,  # -r_complex_val, # r_ji\n",
    "    }\n",
    "    return sdict\n",
    "\n",
    "\n",
    "fabry_perot_tunable, _ = sax.circuit(\n",
    "    netlist={\n",
    "        \"instances\": {\n",
    "            \"mirror1\": mirror,\n",
    "            \"gap\": propagation_i,\n",
    "            \"mirror2\": mirror,\n",
    "        },\n",
    "        \"connections\": {\n",
    "            \"mirror1,out\": \"gap,in\",\n",
    "            \"gap,out\": \"mirror2,in\",\n",
    "        },\n",
    "        \"ports\": {\n",
    "            \"in\": \"mirror1,in\",\n",
    "            \"out\": \"mirror2,out\",\n",
    "        },\n",
    "    },\n",
    "    backend='fg',\n",
    ")\n",
    "\n",
    "settings = sax.get_settings(fabry_perot_tunable)\n",
    "settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86e2f26e",
   "metadata": {
    "papermill": {
     "duration": 0.037048,
     "end_time": "2023-09-03T04:40:18.408077",
     "exception": false,
     "start_time": "2023-09-03T04:40:18.371029",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fabry_perot_tunable, _ = sax.circuit(\n",
    "    netlist={\n",
    "        \"instances\": {\n",
    "            \"mirror1\": mirror,\n",
    "            \"gap\": propagation_i,\n",
    "            \"mirror2\": mirror,\n",
    "        },\n",
    "        \"connections\": {\n",
    "            \"mirror1,out\": \"gap,in\",\n",
    "            \"gap,out\": \"mirror2,in\",\n",
    "        },\n",
    "        \"ports\": {\n",
    "            \"in\": \"mirror1,in\",\n",
    "            \"out\": \"mirror2,out\",\n",
    "        },\n",
    "    },\n",
    "    backend='fg',\n",
    ")\n",
    "\n",
    "settings = sax.get_settings(fabry_perot_tunable)\n",
    "settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3912513",
   "metadata": {
    "papermill": {
     "duration": 0.457045,
     "end_time": "2023-09-03T04:40:18.876416",
     "exception": false,
     "start_time": "2023-09-03T04:40:18.419371",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "N = 100\n",
    "wls = jnp.linspace(0.38, 0.78, N)\n",
    "settings = sax.get_settings(fabry_perot_tunable)\n",
    "settings = sax.update_settings(settings, wl=wls, t_amp=jnp.sqrt(0.5), t_ang=0.0)\n",
    "settings[\"gap\"][\"ni\"] = 1.0\n",
    "settings[\"gap\"][\"di\"] = 2.0\n",
    "transmitted_initial = fabry_perot_tunable(**settings)[\"in\", \"out\"]\n",
    "reflected_initial = fabry_perot_tunable(**settings)[\"out\", \"out\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0e93a4e",
   "metadata": {
    "papermill": {
     "duration": 0.430904,
     "end_time": "2023-09-03T04:40:19.317869",
     "exception": false,
     "start_time": "2023-09-03T04:40:18.886965",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "T_analytical_initial = jnp.abs(airy_t13(t_initial, t_initial, r_initial, r_initial, wls, d=d_gap, n=n_gap))**2\n",
    "R_analytical_initial = jnp.abs(airy_r13(t_initial, t_initial, r_initial, r_initial, wls, d=d_gap, n=n_gap))**2\n",
    "plt.title(f\"t={t_initial:1.3f}, d={d_gap} nm, n={n_gap}\")\n",
    "plt.plot(wls, T_analytical_initial, label=\"T theory\")\n",
    "plt.scatter(wls, jnp.abs(transmitted_initial) ** 2, label=\"T SAX\")\n",
    "plt.plot(wls, R_analytical_initial, label=\"R theory\")\n",
    "plt.scatter(wls, jnp.abs(reflected_initial) ** 2, label=\"R SAX\")\n",
    "plt.xlabel(\"k = 2$\\pi$/$\\lambda$ [1/nm]\")\n",
    "plt.ylabel(\"Power (units of input)\")\n",
    "plt.figlegend(framealpha=1.0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90a42e80",
   "metadata": {
    "papermill": {
     "duration": 0.016244,
     "end_time": "2023-09-03T04:40:19.351057",
     "exception": false,
     "start_time": "2023-09-03T04:40:19.334813",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Wavelength-dependent Fabry-Pérot étalon\n",
    "\n",
    "Let's repeat with a model where parameters can be wavelength-dependent. To comply with the optimizer object, we will stack all design parameters in a single array :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e15a411",
   "metadata": {
    "papermill": {
     "duration": 0.062589,
     "end_time": "2023-09-03T04:40:19.430218",
     "exception": false,
     "start_time": "2023-09-03T04:40:19.367629",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "ts_initial = jnp.zeros(2 * N)\n",
    "ts_initial = ts_initial.at[0:N].set(jnp.sqrt(0.5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbb87fa8",
   "metadata": {
    "papermill": {
     "duration": 0.015865,
     "end_time": "2023-09-03T04:40:19.462082",
     "exception": false,
     "start_time": "2023-09-03T04:40:19.446217",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We will simply loop over all wavelengths, and use different $t$ parameters at each wavelength."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7415ddb",
   "metadata": {
    "papermill": {
     "duration": 0.196918,
     "end_time": "2023-09-03T04:40:19.673328",
     "exception": false,
     "start_time": "2023-09-03T04:40:19.476410",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "wls = jnp.linspace(0.38, 0.78, N)\n",
    "transmitted = jnp.zeros_like(wls)\n",
    "reflected = jnp.zeros_like(wls)\n",
    "settings = sax.get_settings(fabry_perot_tunable)\n",
    "settings = sax.update_settings(settings, wl=wls, t_amp=ts_initial[:N], t_ang=ts_initial[N:])\n",
    "settings[\"gap\"][\"ni\"] = 1.0\n",
    "settings[\"gap\"][\"di\"] = 2.0\n",
    "# Perform computation\n",
    "sdict = fabry_perot_tunable(**settings)\n",
    "transmitted = jnp.abs(sdict[\"in\", \"out\"]) ** 2\n",
    "reflected = jnp.abs(sdict[\"in\", \"in\"]) ** 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4149451",
   "metadata": {
    "papermill": {
     "duration": 0.211274,
     "end_time": "2023-09-03T04:40:19.895220",
     "exception": false,
     "start_time": "2023-09-03T04:40:19.683946",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.plot(wls * 1e3, T_analytical_initial, label=\"T theory\")\n",
    "plt.scatter(wls * 1e3, transmitted, label=\"T SAX\")\n",
    "plt.plot(wls * 1e3, R_analytical_initial, label=\"R theory\")\n",
    "plt.scatter(wls * 1e3, reflected, label=\"R SAX\")\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.ylabel(\"Transmitted and reflected intensities\")\n",
    "plt.legend(loc=\"upper right\")\n",
    "plt.title(f\"t={t_initial:1.3f}, d={d_gap} nm, n={n_gap}\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fbc2915",
   "metadata": {
    "papermill": {
     "duration": 0.011947,
     "end_time": "2023-09-03T04:40:19.919270",
     "exception": false,
     "start_time": "2023-09-03T04:40:19.907323",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Since it seems to work, let's add a target and optimize some harmonics away :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8d6108a",
   "metadata": {
    "papermill": {
     "duration": 0.204952,
     "end_time": "2023-09-03T04:40:20.136036",
     "exception": false,
     "start_time": "2023-09-03T04:40:19.931084",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def lorentzian(l0, dl, wl, A):\n",
    "    return A / ((wl - l0) ** 2 + (0.5 * dl) ** 2)\n",
    "\n",
    "\n",
    "target = lorentzian(533.0, 20.0, wls * 1e3, 100.0)\n",
    "\n",
    "plt.scatter(wls * 1e3, transmitted, label=\"T SAX\")\n",
    "plt.scatter(wls * 1e3, reflected, label=\"R SAX\")\n",
    "plt.plot(wls * 1e3, target, \"r\", linewidth=2, label=\"target\")\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.ylabel(\"Transmitted and reflected intensities\")\n",
    "plt.legend(loc=\"upper right\")\n",
    "plt.title(f\"t={t_initial:1.3f}, d={d_gap} nm, n={n_gap}\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "490f2d79",
   "metadata": {
    "papermill": {
     "duration": 0.018361,
     "end_time": "2023-09-03T04:40:20.166752",
     "exception": false,
     "start_time": "2023-09-03T04:40:20.148391",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def loss(ts):\n",
    "    N = len(ts[::2])\n",
    "    wls = jnp.linspace(0.38, 0.78, N)\n",
    "    target = lorentzian(533.0, 20.0, wls * 1e3, 100.0)\n",
    "    settings = sax.get_settings(fabry_perot_tunable)\n",
    "    settings = sax.update_settings(settings, wl=wls, t_amp=ts[:N], t_ang=ts[N:])\n",
    "    settings[\"gap\"][\"ni\"] = 1.0\n",
    "    settings[\"gap\"][\"di\"] = 2.0\n",
    "    sdict = fabry_perot_tunable(**settings)\n",
    "    transmitted = jnp.abs(sdict[\"in\", \"out\"]) ** 2\n",
    "    return (jnp.abs(transmitted - target) ** 2).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed1b3bfe",
   "metadata": {
    "papermill": {
     "duration": 0.032696,
     "end_time": "2023-09-03T04:40:20.216751",
     "exception": false,
     "start_time": "2023-09-03T04:40:20.184055",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "grad = jax.jit(jax.grad(loss))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "721379af",
   "metadata": {
    "papermill": {
     "duration": 0.017754,
     "end_time": "2023-09-03T04:40:20.248010",
     "exception": false,
     "start_time": "2023-09-03T04:40:20.230256",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "optim_init, optim_update, optim_params = opt.adam(step_size=0.001)\n",
    "\n",
    "def train_step(step, optim_state):\n",
    "    ts = optim_params(optim_state)\n",
    "    lossvalue = loss(ts)\n",
    "    gradvalue = grad(ts)\n",
    "    optim_state = optim_update(step, gradvalue, optim_state)\n",
    "    return lossvalue, gradvalue, optim_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "744e0990",
   "metadata": {
    "papermill": {
     "duration": 2.506612,
     "end_time": "2023-09-03T04:40:22.767721",
     "exception": true,
     "start_time": "2023-09-03T04:40:20.261109",
     "status": "failed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "range_ = tqdm.trange(2000)\n",
    "\n",
    "optim_state = optim_init(ts_initial)\n",
    "for step in range_:\n",
    "    lossvalue, gradvalue, optim_state = train_step(step, optim_state)\n",
    "    range_.set_postfix(loss=f\"{lossvalue:.6f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2451287d",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "pending"
    },
    "tags": []
   },
   "source": [
    "The optimized parameters are now wavelength-dependent :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7893774",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "pending"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "ts_optimal = optim_params(optim_state)\n",
    "\n",
    "plt.scatter(wls * 1e3, ts_initial[:N], label=\"t initial\")\n",
    "plt.scatter(wls * 1e3, ts_optimal[:N], label=\"t optimal\")\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.ylabel(\"|t| $(\\lambda)$\")\n",
    "plt.legend(loc=\"best\")\n",
    "plt.title(f\"d={d_gap} nm, n={n_gap}\")\n",
    "plt.show()\n",
    "\n",
    "plt.scatter(wls * 1e3, ts_initial[N:], label=\"t initial\")\n",
    "plt.scatter(wls * 1e3, ts_optimal[N:], label=\"t optimal\")\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.ylabel(\"angle $t (\\lambda)$ (rad)\")\n",
    "plt.legend(loc=\"best\")\n",
    "plt.title(f\"d={d_gap} nm, n={n_gap}\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1b92e1e",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "pending"
    },
    "tags": []
   },
   "source": [
    "Visualizing the result :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "223bd973",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "pending"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "wls = jnp.linspace(0.38, 0.78, N)\n",
    "transmitted_optimal = jnp.zeros_like(wls)\n",
    "reflected_optimal = jnp.zeros_like(wls)\n",
    "\n",
    "settings = sax.get_settings(fabry_perot_tunable)\n",
    "settings = sax.update_settings(\n",
    "    settings, wl=wls, t_amp=ts_optimal[:N], t_ang=ts_optimal[N:]\n",
    ")\n",
    "settings[\"gap\"][\"ni\"] = 1.0\n",
    "settings[\"gap\"][\"di\"] = 2.0\n",
    "transmitted_optimal = jnp.abs(fabry_perot_tunable(**settings)[\"in\", \"out\"]) ** 2\n",
    "reflected_optimal = jnp.abs(fabry_perot_tunable(**settings)[\"in\", \"in\"]) ** 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52ef7be4",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "pending"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.scatter(wls * 1e3, transmitted_optimal, label=\"T\")\n",
    "plt.scatter(wls * 1e3, reflected_optimal, label=\"R\")\n",
    "plt.plot(wls * 1e3, lorentzian(533, 20, wls * 1e3, 100), \"r\", label=\"target\")\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.ylabel(\"Transmitted and reflected intensities\")\n",
    "plt.legend(loc=\"upper right\")\n",
    "plt.title(f\"Optimized t($\\lambda$), d={d_gap} nm, n={n_gap}\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd20054c",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "pending"
    },
    "tags": []
   },
   "source": [
    "The hard part is now to find physical stacks that physically implement $t(\\lambda)$. However, the ease with which we can modify and complexify the loss function opens opportunities for regularization and more complicated objective functions.\n",
    "\n",
    "The models above are available in `sax.models.thinfilm`, and can straightforwardly be extended to propagation at an angle, s and p polarizations, nonreciprocal systems, and systems with losses."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sax",
   "language": "python",
   "name": "sax"
  },
  "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.9.18"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 17.400725,
   "end_time": "2023-09-03T04:40:23.800437",
   "environment_variables": {},
   "exception": true,
   "input_path": "./05_thinfilm.ipynb",
   "output_path": "./05_thinfilm.ipynb",
   "parameters": {},
   "start_time": "2023-09-03T04:40:06.399712",
   "version": "2.4.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}