notebooks/sax_01_sax.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# SAX circuit simulator\n",
    "\n",
    "[SAX](https://flaport.github.io/sax/) is a circuit solver written in JAX, writing your component models in SAX enables you not only to get the function values but the gradients, this is useful for circuit optimization.\n",
    "\n",
    "This tutorial has been adapted from the SAX Quick Start and example notebooks.\n",
    "\n",
    "You can install sax with pip (read the SAX install instructions [here](https://github.com/flaport/sax#installation))\n",
    "\n",
    "```\n",
    "! pip install gplugins[sax]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import logging\n",
    "import os\n",
    "import sys\n",
    "from functools import partial\n",
    "\n",
    "import gdsfactory as gf\n",
    "import jax\n",
    "import jax.example_libraries.optimizers as opt\n",
    "import jax.numpy as jnp\n",
    "import matplotlib.pyplot as plt\n",
    "import meow as mw\n",
    "import numpy as np\n",
    "import sax\n",
    "from gdsfactory.generic_tech import get_generic_pdk\n",
    "from numpy.fft import fft2, fftfreq, fftshift, ifft2\n",
    "from rich.logging import RichHandler\n",
    "from scipy import constants\n",
    "from sklearn.linear_model import LinearRegression\n",
    "from tqdm.notebook import tqdm, trange\n",
    "\n",
    "import gplugins.sax as gs\n",
    "import gplugins.tidy3d as gt\n",
    "from gplugins.common.config import PATH\n",
    "\n",
    "gf.config.rich_output()\n",
    "PDK = get_generic_pdk()\n",
    "PDK.activate()\n",
    "\n",
    "logger = logging.getLogger()\n",
    "logger.removeHandler(sys.stderr)\n",
    "logging.basicConfig(level=\"WARNING\", datefmt=\"[%X]\", handlers=[RichHandler()])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Scatter *dictionaries*\n",
    "\n",
    "The core datastructure for specifying scatter parameters in SAX is a dictionary... more specifically a dictionary which maps a port combination (2-tuple) to a scatter parameter (or an array of scatter parameters when considering multiple wavelengths for example). Such a specific dictionary mapping is called ann `SDict` in SAX (`SDict ≈ Dict[Tuple[str,str], float]`).\n",
    "\n",
    "Dictionaries are in fact much better suited for characterizing S-parameters than, say, (jax-)numpy arrays due to the inherent sparse nature of scatter parameters. Moreover, dictionaries allow for string indexing, which makes them much more pleasant to use in this context.\n",
    "\n",
    "```\n",
    "o2            o3\n",
    "   \\        /\n",
    "    ========\n",
    "   /        \\\n",
    "o1            o4\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "coupling = 0.5\n",
    "kappa = coupling**0.5\n",
    "tau = (1 - coupling) ** 0.5\n",
    "coupler_dict = {\n",
    "    (\"o1\", \"o4\"): tau,\n",
    "    (\"o4\", \"o1\"): tau,\n",
    "    (\"o1\", \"o3\"): 1j * kappa,\n",
    "    (\"o3\", \"o1\"): 1j * kappa,\n",
    "    (\"o2\", \"o4\"): 1j * kappa,\n",
    "    (\"o4\", \"o2\"): 1j * kappa,\n",
    "    (\"o2\", \"o3\"): tau,\n",
    "    (\"o3\", \"o2\"): tau,\n",
    "}\n",
    "print(coupler_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    " it can still be tedious to specify every port in the circuit manually. SAX therefore offers the `reciprocal` function, which auto-fills the reverse connection if the forward connection exist. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "coupler_dict = sax.reciprocal(\n",
    "    {\n",
    "        (\"o1\", \"o4\"): tau,\n",
    "        (\"o1\", \"o3\"): 1j * kappa,\n",
    "        (\"o2\", \"o4\"): 1j * kappa,\n",
    "        (\"o2\", \"o3\"): tau,\n",
    "    }\n",
    ")\n",
    "\n",
    "coupler_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## Parametrized Models\n",
    "\n",
    "Constructing such an `SDict` is easy, however, usually we're more interested in having parametrized models for our components. To parametrize the coupler `SDict`, just wrap it in a function to obtain a SAX `Model`, which is a keyword-only function mapping to an `SDict`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def coupler(coupling=0.5) -> sax.SDict:\n",
    "    kappa = coupling**0.5\n",
    "    tau = (1 - coupling) ** 0.5\n",
    "    return sax.reciprocal(\n",
    "        {\n",
    "            (\"o1\", \"o4\"): tau,\n",
    "            (\"o1\", \"o3\"): 1j * kappa,\n",
    "            (\"o2\", \"o4\"): 1j * kappa,\n",
    "            (\"o2\", \"o3\"): tau,\n",
    "        }\n",
    "    )\n",
    "\n",
    "\n",
    "coupler(coupling=0.3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def waveguide(wl=1.55, wl0=1.55, neff=2.34, ng=3.4, length=10.0, loss=0.0) -> sax.SDict:\n",
    "    dwl = wl - wl0\n",
    "    dneff_dwl = (ng - neff) / wl0\n",
    "    neff = neff - dwl * dneff_dwl\n",
    "    phase = 2 * jnp.pi * neff * length / wl\n",
    "    transmission = 10 ** (-loss * length / 20) * jnp.exp(1j * phase)\n",
    "    return sax.reciprocal(\n",
    "        {\n",
    "            (\"o1\", \"o2\"): transmission,\n",
    "        }\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "### Waveguide model\n",
    "\n",
    "You can create a dispersive waveguide model in SAX."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "Lets compute the effective index `neff` and group index `ng` for a 1550nm 500nm straight waveguide"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "nm = 1e-3\n",
    "strip = gt.modes.Waveguide(\n",
    "    wavelength=1.55,\n",
    "    core_width=500 * nm,\n",
    "    core_thickness=220 * nm,\n",
    "    slab_thickness=0.0,\n",
    "    core_material=\"si\",\n",
    "    clad_material=\"sio2\",\n",
    "    group_index_step=10 * nm,\n",
    ")\n",
    "strip.plot_field(field_name=\"Ex\", mode_index=0)  # TE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "neff = strip.n_eff[0]\n",
    "print(neff)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "ng = strip.n_group[0]\n",
    "print(ng)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "straight_sc = partial(gs.models.straight, neff=neff, ng=ng)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "gs.plot_model(straight_sc)\n",
    "plt.ylim(-1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "gs.plot_model(straight_sc, phase=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "### Coupler model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "c = gf.components.coupler(length=10, gap=0.2)\n",
    "c.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "nm = 1e-3\n",
    "cp = gt.modes.WaveguideCoupler(\n",
    "    wavelength=1.55,\n",
    "    core_width=(500 * nm, 500 * nm),\n",
    "    gap=200 * nm,\n",
    "    core_thickness=220 * nm,\n",
    "    slab_thickness=0 * nm,\n",
    "    core_material=\"si\",\n",
    "    clad_material=\"sio2\",\n",
    ")\n",
    "cp.plot_field(field_name=\"Ex\", mode_index=0)  # even mode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "cp.plot_field(field_name=\"Ex\", mode_index=1)  # odd mode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "coupler = gt.modes.WaveguideCoupler(\n",
    "    wavelength=1.55,\n",
    "    core_width=(0.45, 0.45),\n",
    "    core_thickness=220 * nm,\n",
    "    core_material=\"si\",\n",
    "    clad_material=\"sio2\",\n",
    "    num_modes=4,\n",
    "    gap=100 * nm,\n",
    ")\n",
    "\n",
    "print(\"\\nCoupler:\", coupler)\n",
    "print(\"Effective indices:\", coupler.n_eff)\n",
    "print(\"Mode areas:\", coupler.mode_area)\n",
    "print(\"Coupling length:\", coupler.coupling_length())\n",
    "\n",
    "gaps = np.linspace(0.05, 0.15, 11)\n",
    "lengths = gt.modes.sweep_coupling_length(coupler, gaps)\n",
    "plt.plot(gaps, lengths)\n",
    "plt.xlabel(\"Gap (μm)\")\n",
    "plt.ylabel(\"Coupling length (μm)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "For a 200nm gap the effective index difference `dn` is `0.026`, which means that there is 100% power coupling over 29.4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "coupler_sc = partial(gs.models.coupler, dn=0.026, length=0, coupling0=0)\n",
    "gs.plot_model(coupler_sc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24",
   "metadata": {},
   "source": [
    "If we ignore the coupling from the bend `coupling0 = 0` we know that for a 3dB coupling we need half of the `lc` length, which is the length needed to coupler `100%` of power."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "coupler_sc = partial(gs.models.coupler, dn=0.026, length=29.4 / 2, coupling0=0)\n",
    "gs.plot_model(coupler_sc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26",
   "metadata": {},
   "source": [
    "### FDTD Sparameters model\n",
    "\n",
    "You can also fit a model from Sparameter FDTD simulation data from tidy3d, Lumerical or MEEP."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27",
   "metadata": {},
   "source": [
    "## Model fit\n",
    "\n",
    "You can fit a sax model to Sparameter FDTD simulation data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28",
   "metadata": {},
   "outputs": [],
   "source": [
    "filepath = PATH.test_data / \"sp\" / \"coupler_G224n_L20_S220.csv\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29",
   "metadata": {},
   "outputs": [],
   "source": [
    "coupler_fdtd = gs.read.model_from_csv(\n",
    "    filepath=filepath,\n",
    "    xkey=\"wavelength_nm\",\n",
    "    prefix=\"S\",\n",
    "    xunits=1e-3,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30",
   "metadata": {},
   "outputs": [],
   "source": [
    "gs.plot_model(coupler_fdtd)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31",
   "metadata": {},
   "source": [
    "Lets fit the coupler spectrum with a linear regression `sklearn` fit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32",
   "metadata": {},
   "outputs": [],
   "source": [
    "f = jnp.linspace(constants.c / 1.0e-6, constants.c / 2.0e-6, 500) * 1e-12  # THz\n",
    "wl = constants.c / (f * 1e12) * 1e6  # um\n",
    "\n",
    "coupler_fdtd = gs.read.model_from_csv(\n",
    "    filepath, xkey=\"wavelength_nm\", prefix=\"S\", xunits=1e-3\n",
    ")\n",
    "sd = coupler_fdtd(wl=wl)\n",
    "\n",
    "k = sd[\"o1\", \"o3\"]\n",
    "t = sd[\"o1\", \"o4\"]\n",
    "s = t + k\n",
    "a = t - k"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33",
   "metadata": {},
   "source": [
    "Lets fit the symmetric (t+k) and antisymmetric (t-k) transmission\n",
    "\n",
    "### Symmetric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(wl, jnp.abs(s))\n",
    "plt.grid(True)\n",
    "plt.xlabel(\"Frequency [THz]\")\n",
    "plt.ylabel(\"Transmission\")\n",
    "plt.title(\"symmetric (transmission + coupling)\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(wl, jnp.abs(a))\n",
    "plt.grid(True)\n",
    "plt.xlabel(\"Frequency [THz]\")\n",
    "plt.ylabel(\"Transmission\")\n",
    "plt.title(\"anti-symmetric (transmission - coupling)\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36",
   "metadata": {},
   "outputs": [],
   "source": [
    "r = LinearRegression()\n",
    "\n",
    "\n",
    "def fX(x, _order=8):\n",
    "    return (\n",
    "        x[:, None] ** (jnp.arange(_order)[None, :])\n",
    "    )  # artificially create more 'features' (wl**2, wl**3, wl**4, ...)\n",
    "\n",
    "\n",
    "X = fX(wl)\n",
    "r.fit(X, jnp.abs(s))\n",
    "asm, bsm = r.coef_, r.intercept_\n",
    "\n",
    "\n",
    "def fsm(x):\n",
    "    return fX(x) @ asm + bsm  # fit symmetric module fiir\n",
    "\n",
    "\n",
    "plt.plot(wl, jnp.abs(s), label=\"data\")\n",
    "plt.plot(wl, fsm(wl), label=\"fit\")\n",
    "plt.grid(True)\n",
    "plt.xlabel(\"Frequency [THz]\")\n",
    "plt.ylabel(\"Transmission\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37",
   "metadata": {},
   "outputs": [],
   "source": [
    "r = LinearRegression()\n",
    "r.fit(X, jnp.unwrap(jnp.angle(s)))\n",
    "asp, bsp = r.coef_, r.intercept_\n",
    "\n",
    "\n",
    "def fsp(x):\n",
    "    return fX(x) @ asp + bsp  # fit symmetric phase\n",
    "\n",
    "\n",
    "plt.plot(wl, jnp.unwrap(jnp.angle(s)), label=\"data\")\n",
    "plt.plot(wl, fsp(wl), label=\"fit\")\n",
    "plt.grid(True)\n",
    "plt.xlabel(\"Frequency [THz]\")\n",
    "plt.ylabel(\"Angle [rad]\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fs(x):\n",
    "    return fsm(x) * jnp.exp(1j * fsp(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39",
   "metadata": {},
   "source": [
    "Lets fit the symmetric (t+k) and antisymmetric (t-k) transmission\n",
    "\n",
    "### Anti-Symmetric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40",
   "metadata": {},
   "outputs": [],
   "source": [
    "r = LinearRegression()\n",
    "r.fit(X, jnp.abs(a))\n",
    "aam, bam = r.coef_, r.intercept_\n",
    "\n",
    "\n",
    "def fam(x):\n",
    "    return fX(x) @ aam + bam\n",
    "\n",
    "\n",
    "plt.plot(wl, jnp.abs(a))\n",
    "plt.plot(wl, fam(wl))\n",
    "plt.grid(True)\n",
    "plt.xlabel(\"Frequency [THz]\")\n",
    "plt.ylabel(\"Transmission\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41",
   "metadata": {},
   "outputs": [],
   "source": [
    "r = LinearRegression()\n",
    "r.fit(X, jnp.unwrap(jnp.angle(a)))\n",
    "aap, bap = r.coef_, r.intercept_\n",
    "\n",
    "\n",
    "def fap(x):\n",
    "    return fX(x) @ aap + bap\n",
    "\n",
    "\n",
    "plt.plot(wl, jnp.unwrap(jnp.angle(a)))\n",
    "plt.plot(wl, fap(wl))\n",
    "plt.grid(True)\n",
    "plt.xlabel(\"Frequency [THz]\")\n",
    "plt.ylabel(\"Angle [rad]\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fa(x):\n",
    "    return fam(x) * jnp.exp(1j * fap(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43",
   "metadata": {},
   "source": [
    "### Total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_ = 0.5 * (fs(wl) + fa(wl))\n",
    "\n",
    "plt.plot(wl, jnp.abs(t))\n",
    "plt.plot(wl, jnp.abs(t_))\n",
    "plt.xlabel(\"Frequency [THz]\")\n",
    "plt.ylabel(\"Transmission\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45",
   "metadata": {},
   "outputs": [],
   "source": [
    "k_ = 0.5 * (fs(wl) - fa(wl))\n",
    "\n",
    "plt.plot(wl, jnp.abs(k))\n",
    "plt.plot(wl, jnp.abs(k_))\n",
    "plt.xlabel(\"Frequency [THz]\")\n",
    "plt.ylabel(\"Coupling\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46",
   "metadata": {},
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def coupler(wl=1.5):\n",
    "    wl = jnp.asarray(wl)\n",
    "    wl_shape = wl.shape\n",
    "    wl = wl.ravel()\n",
    "    t = (0.5 * (fs(wl) + fa(wl))).reshape(*wl_shape)\n",
    "    k = (0.5 * (fs(wl) - fa(wl))).reshape(*wl_shape)\n",
    "    sdict = {\n",
    "        (\"o1\", \"o4\"): t,\n",
    "        (\"o1\", \"o3\"): k,\n",
    "        (\"o2\", \"o3\"): k,\n",
    "        (\"o2\", \"o4\"): t,\n",
    "    }\n",
    "    return sax.reciprocal(sdict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47",
   "metadata": {},
   "outputs": [],
   "source": [
    "f = jnp.linspace(constants.c / 1.0e-6, constants.c / 2.0e-6, 500) * 1e-12  # THz\n",
    "wl = constants.c / (f * 1e12) * 1e6  # um\n",
    "\n",
    "coupler_fdtd = gs.read.model_from_csv(\n",
    "    filepath, xkey=\"wavelength_nm\", prefix=\"S\", xunits=1e-3\n",
    ")\n",
    "sd = coupler_fdtd(wl=wl)\n",
    "sd_ = coupler(wl=wl)\n",
    "\n",
    "T = jnp.abs(sd[\"o1\", \"o4\"]) ** 2\n",
    "K = jnp.abs(sd[\"o1\", \"o3\"]) ** 2\n",
    "T_ = jnp.abs(sd_[\"o1\", \"o4\"]) ** 2\n",
    "K_ = jnp.abs(sd_[\"o1\", \"o3\"]) ** 2\n",
    "dP = jnp.unwrap(jnp.angle(sd[\"o1\", \"o3\"]) - jnp.angle(sd[\"o1\", \"o4\"]))\n",
    "dP_ = jnp.unwrap(jnp.angle(sd_[\"o1\", \"o3\"]) - jnp.angle(sd_[\"o1\", \"o4\"]))\n",
    "\n",
    "plt.figure(figsize=(12, 3))\n",
    "plt.plot(wl, T, label=\"T (fdtd)\", c=\"C0\", ls=\":\", lw=\"6\")\n",
    "plt.plot(wl, T_, label=\"T (model)\", c=\"C0\")\n",
    "\n",
    "plt.plot(wl, K, label=\"K (fdtd)\", c=\"C1\", ls=\":\", lw=\"6\")\n",
    "plt.plot(wl, K_, label=\"K (model)\", c=\"C1\")\n",
    "\n",
    "plt.ylim(-0.05, 1.05)\n",
    "plt.grid(True)\n",
    "\n",
    "plt.twinx()\n",
    "plt.plot(wl, dP, label=\"ΔΦ (fdtd)\", color=\"C2\", ls=\":\", lw=\"6\")\n",
    "plt.plot(wl, dP_, label=\"ΔΦ (model)\", color=\"C2\")\n",
    "\n",
    "plt.xlabel(\"Frequency [THz]\")\n",
    "plt.ylabel(\"Transmission\")\n",
    "plt.figlegend(bbox_to_anchor=(1.08, 0.9))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48",
   "metadata": {},
   "source": [
    "## SAX gdsfactory Compatibility\n",
    "> From Layout to Circuit Model\n",
    "\n",
    "If you define your SAX S parameter models for your components, you can directly simulate your circuits from gdsfactory. Let's first define our own MZI:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49",
   "metadata": {},
   "outputs": [],
   "source": [
    "@gf.cell\n",
    "def simple_mzi(delta_length=20.0):\n",
    "    c = gf.Component()\n",
    "\n",
    "    # components\n",
    "    mmi_in = gf.components.mmi1x2()\n",
    "    mmi_out = gf.components.mmi2x2()\n",
    "    bend = gf.components.bend_euler()\n",
    "    half_delay_straight = gf.components.straight(length=delta_length / 2)\n",
    "\n",
    "    # references\n",
    "    mmi_in = c.add_ref(mmi_in, name=\"mmi_in\")\n",
    "    mmi_out = c.add_ref(mmi_out, name=\"mmi_out\")\n",
    "    straight_top1 = c.add_ref(half_delay_straight, name=\"straight_top1\")\n",
    "    straight_top2 = c.add_ref(half_delay_straight, name=\"straight_top2\")\n",
    "    bend_top1 = c.add_ref(bend, name=\"bend_top1\")\n",
    "    bend_top2 = c.add_ref(bend, name=\"bend_top2\").dmirror()\n",
    "    bend_top3 = c.add_ref(bend, name=\"bend_top3\").dmirror()\n",
    "    bend_top4 = c.add_ref(bend, name=\"bend_top4\")\n",
    "    bend_btm1 = c.add_ref(bend, name=\"bend_btm1\").dmirror()\n",
    "    bend_btm2 = c.add_ref(bend, name=\"bend_btm2\")\n",
    "    bend_btm3 = c.add_ref(bend, name=\"bend_btm3\")\n",
    "    bend_btm4 = c.add_ref(bend, name=\"bend_btm4\").dmirror()\n",
    "\n",
    "    # connections\n",
    "    bend_top1.connect(\"o1\", mmi_in.ports[\"o2\"])\n",
    "    straight_top1.connect(\"o1\", bend_top1.ports[\"o2\"])\n",
    "    bend_top2.connect(\"o1\", straight_top1.ports[\"o2\"])\n",
    "    bend_top3.connect(\"o1\", bend_top2.ports[\"o2\"])\n",
    "    straight_top2.connect(\"o1\", bend_top3.ports[\"o2\"])\n",
    "    bend_top4.connect(\"o1\", straight_top2.ports[\"o2\"])\n",
    "\n",
    "    bend_btm1.connect(\"o1\", mmi_in.ports[\"o3\"])\n",
    "    bend_btm2.connect(\"o1\", bend_btm1.ports[\"o2\"])\n",
    "    bend_btm3.connect(\"o1\", bend_btm2.ports[\"o2\"])\n",
    "    bend_btm4.connect(\"o1\", bend_btm3.ports[\"o2\"])\n",
    "\n",
    "    mmi_out.connect(\"o1\", bend_btm4.ports[\"o2\"])\n",
    "\n",
    "    # ports\n",
    "    c.add_port(\n",
    "        \"o1\",\n",
    "        port=mmi_in.ports[\"o1\"],\n",
    "    )\n",
    "    c.add_port(\"o2\", port=mmi_out.ports[\"o3\"])\n",
    "    c.add_port(\"o3\", port=mmi_out.ports[\"o4\"])\n",
    "    return c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50",
   "metadata": {},
   "outputs": [],
   "source": [
    "mzi = simple_mzi(delta_length=20)\n",
    "mzi.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51",
   "metadata": {},
   "outputs": [],
   "source": [
    "mzi.plot_netlist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "netlist = mzi.get_netlist()\n",
    "print(netlist.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "The netlist has four different components:\n",
    "\n",
    "1. mmi1x2\n",
    "2. mmi2x2\n",
    "3. straight\n",
    "4. bend_euler\n",
    "\n",
    "You need models for each subcomponents to simulate the Component."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mmi1x2():\n",
    "    \"\"\"Assumes a perfect 1x2 splitter\"\"\"\n",
    "    return sax.reciprocal(\n",
    "        {\n",
    "            (\"o1\", \"o2\"): 0.5**0.5,\n",
    "            (\"o1\", \"o3\"): 0.5**0.5,\n",
    "        }\n",
    "    )\n",
    "\n",
    "\n",
    "def mmi2x2():\n",
    "    return sax.reciprocal(\n",
    "        {\n",
    "            (\"o1\", \"o3\"): 0.5**0.5,\n",
    "            (\"o1\", \"o4\"): 1j * 0.5**0.5,\n",
    "            (\"o2\", \"o3\"): 1j * 0.5**0.5,\n",
    "            (\"o2\", \"o4\"): 0.5**0.5,\n",
    "        }\n",
    "    )\n",
    "\n",
    "\n",
    "def straight(wl=1.5, length=10.0, neff=2.4) -> sax.SDict:\n",
    "    return sax.reciprocal(\n",
    "        {\n",
    "            (\"o1\", \"o2\"): jnp.exp(2j * jnp.pi * neff * length / wl),\n",
    "        }\n",
    "    )\n",
    "\n",
    "\n",
    "def bend_euler(wl=1.5, length=20.0):\n",
    "    \"\"\" \"Let's assume a reduced transmission for the euler bend compared to a straight\"\"\"\n",
    "    return {k: 0.99 * v for k, v in straight(wl=wl, length=length).items()}\n",
    "\n",
    "\n",
    "models = {\n",
    "    \"mmi1x2\": mmi1x2,\n",
    "    \"mmi2x2\": mmi2x2,\n",
    "    \"straight\": straight,\n",
    "    \"bend_euler\": bend_euler,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55",
   "metadata": {},
   "outputs": [],
   "source": [
    "circuit, _ = sax.circuit(netlist=netlist, models=models)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56",
   "metadata": {},
   "outputs": [],
   "source": [
    "wl = np.linspace(1.5, 1.6)\n",
    "S = circuit(wl=wl)\n",
    "\n",
    "plt.figure(figsize=(14, 4))\n",
    "plt.title(\"MZI\")\n",
    "plt.plot(1e3 * wl, jnp.abs(S[\"o1\", \"o2\"]) ** 2)\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.ylabel(\"T\")\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57",
   "metadata": {},
   "outputs": [],
   "source": [
    "mzi = simple_mzi(delta_length=40)  # Double the length, reduces FSR by 1/2\n",
    "mzi.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58",
   "metadata": {},
   "outputs": [],
   "source": [
    "circuit, _ = sax.circuit(netlist=mzi.get_netlist(), models=models)\n",
    "\n",
    "wl = np.linspace(1.5, 1.6, 256)\n",
    "S = circuit(wl=wl)\n",
    "\n",
    "plt.figure(figsize=(14, 4))\n",
    "plt.title(\"MZI\")\n",
    "plt.plot(1e3 * wl, jnp.abs(S[\"o1\", \"o2\"]) ** 2)\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.ylabel(\"T\")\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59",
   "metadata": {},
   "source": [
    "## Layout aware Monte Carlo\n",
    "\n",
    "You can model the manufacturing variations on the performance of photonics thanks to the fast SAX circuit simulator with layout information and wafer maps of waveguide width and layer thickness variations.\n",
    "\n",
    "The width and height variations can be extracted from:\n",
    "\n",
    "- Ring resonators [2017](https://opg.optica.org/oe/fulltext.cfm?uri=oe-25-9-9712&id=363202)\n",
    "- MZI interferometers [2019](https://ieeexplore.ieee.org/abstract/document/8675367)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60",
   "metadata": {
    "papermill": {
     "duration": 0.027193,
     "end_time": "2024-06-17T18:16:32.186710",
     "exception": false,
     "start_time": "2024-06-17T18:16:32.159517",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Waveguide Model\n",
    "\n",
    "Our waveguide model is not very good (it just has 100% transmission and no phase). Let's do something about the phase calculation. To do this, we need to find the effective index of the waveguide in relation to its parameters. We can use [meow](https://github.com/flaport/meow) to obtain the waveguide effective index. Let's first create a `find_waveguide_modes`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61",
   "metadata": {
    "papermill": {
     "duration": 0.044747,
     "end_time": "2024-06-17T18:16:32.257169",
     "exception": false,
     "start_time": "2024-06-17T18:16:32.212422",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def find_waveguide_modes(\n",
    "    wl: float = 1.55,\n",
    "    n_box: float = 1.4,\n",
    "    n_clad: float = 1.4,\n",
    "    n_core: float = 3.4,\n",
    "    t_slab: float = 0.0,\n",
    "    t_soi: float = 0.22,\n",
    "    w_core: float = 0.45,\n",
    "    du=0.02,\n",
    "    n_modes: int = 10,\n",
    "    cache_path: str = \"modes\",\n",
    "    replace_cached: bool = False,\n",
    "):\n",
    "    length = 10.0\n",
    "    delta = 10 * du\n",
    "    env = mw.Environment(wl=wl)\n",
    "    cache_path = os.path.abspath(cache_path)\n",
    "    os.makedirs(cache_path, exist_ok=True)\n",
    "    fn = f\"{wl=:.2f}-{n_box=:.2f}-{n_clad=:.2f}-{n_core=:.2f}-{t_slab=:.3f}-{t_soi=:.3f}-{w_core=:.3f}-{du=:.3f}-{n_modes=}.json\"\n",
    "    path = os.path.join(cache_path, fn)\n",
    "    if not replace_cached and os.path.exists(path):\n",
    "        return [mw.Mode.parse_obj(mode) for mode in json.load(open(path))]\n",
    "\n",
    "    # fmt: off\n",
    "    m_core = mw.SampledMaterial(name=\"slab\", n=np.asarray([n_core, n_core]), params={\"wl\": np.asarray([1.0, 2.0])}, meta={\"color\": (0.9, 0, 0, 0.9)})\n",
    "    m_clad = mw.SampledMaterial(name=\"clad\", n=np.asarray([n_clad, n_clad]), params={\"wl\": np.asarray([1.0, 2.0])})\n",
    "    m_box = mw.SampledMaterial(name=\"box\", n=np.asarray([n_box, n_box]), params={\"wl\": np.asarray([1.0, 2.0])})\n",
    "    box = mw.Structure(material=m_box, geometry=mw.Box(x_min=- 2 * w_core - delta, x_max= 2 * w_core + delta, y_min=- 2 * t_soi - delta, y_max=0.0, z_min=0.0, z_max=length))\n",
    "    slab = mw.Structure(material=m_core, geometry=mw.Box(x_min=-2 * w_core - delta, x_max=2 * w_core + delta, y_min=0.0, y_max=t_slab, z_min=0.0, z_max=length))\n",
    "    clad = mw.Structure(material=m_clad, geometry=mw.Box(x_min=-2 * w_core - delta, x_max=2 * w_core + delta, y_min=0, y_max=3 * t_soi + delta, z_min=0.0, z_max=length))\n",
    "    core = mw.Structure(material=m_core, geometry=mw.Box(x_min=-w_core / 2, x_max=w_core / 2, y_min=0.0, y_max=t_soi, z_min=0.0, z_max=length))\n",
    "\n",
    "    cell = mw.Cell(structures=[box, clad, slab, core], mesh=mw.Mesh2D( x=np.arange(-2*w_core, 2*w_core, du), y=np.arange(-2*t_soi, 3*t_soi, du), ), z_min=0.0, z_max=10.0)\n",
    "    cross_section = mw.CrossSection.from_cell(cell=cell, env=env)\n",
    "    modes = mw.compute_modes(cross_section, num_modes=n_modes)\n",
    "    # fmt: on\n",
    "\n",
    "    json.dump([json.loads(mode.json()) for mode in modes], open(path, \"w\"))\n",
    "\n",
    "    return modes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62",
   "metadata": {
    "papermill": {
     "duration": 0.033443,
     "end_time": "2024-06-17T18:16:32.316599",
     "exception": false,
     "start_time": "2024-06-17T18:16:32.283156",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We can also create a rudimentary model for the silicon refractive index:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63",
   "metadata": {
    "papermill": {
     "duration": 0.047724,
     "end_time": "2024-06-17T18:16:32.390273",
     "exception": false,
     "start_time": "2024-06-17T18:16:32.342549",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def silicon_index(wl):\n",
    "    \"\"\"a rudimentary silicon refractive index model\"\"\"\n",
    "    a, b = 0.2411478522088102, 3.3229394315868976\n",
    "    return a / wl + b"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64",
   "metadata": {
    "papermill": {
     "duration": 0.014059,
     "end_time": "2024-06-17T18:16:32.426233",
     "exception": false,
     "start_time": "2024-06-17T18:16:32.412174",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We can now easily calculate the modes of a strip waveguide:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65",
   "metadata": {
    "papermill": {
     "duration": 4.260045,
     "end_time": "2024-06-17T18:16:36.701343",
     "exception": false,
     "start_time": "2024-06-17T18:16:32.441298",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "modes = find_waveguide_modes(wl=1.5, n_core=silicon_index(wl=1.5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66",
   "metadata": {
    "papermill": {
     "duration": 0.025341,
     "end_time": "2024-06-17T18:16:36.744591",
     "exception": false,
     "start_time": "2024-06-17T18:16:36.719250",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "The fundamental mode is the mode with index 0:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67",
   "metadata": {
    "papermill": {
     "duration": 0.481667,
     "end_time": "2024-06-17T18:16:37.255497",
     "exception": false,
     "start_time": "2024-06-17T18:16:36.773830",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "mw.visualize(modes[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68",
   "metadata": {
    "papermill": {
     "duration": 119.282141,
     "end_time": "2024-06-17T18:18:36.554516",
     "exception": false,
     "start_time": "2024-06-17T18:16:37.272375",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "wavelengths, widths = np.mgrid[1.5:1.6:10j, 0.4:0.6:5j]\n",
    "neffs = np.zeros_like(wavelengths)\n",
    "neffs_ = neffs.ravel()\n",
    "\n",
    "for i, (wl, w) in enumerate(zip(tqdm(wavelengths.ravel()), widths.ravel())):\n",
    "    modes = find_waveguide_modes(\n",
    "        wl=wl, n_core=silicon_index(wl), w_core=w, replace_cached=False\n",
    "    )\n",
    "    neffs_[i] = np.real(modes[0].neff)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69",
   "metadata": {
    "papermill": {
     "duration": 0.006308,
     "end_time": "2024-06-17T18:18:36.567272",
     "exception": false,
     "start_time": "2024-06-17T18:18:36.560964",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "This results in the following effective indices:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70",
   "metadata": {
    "papermill": {
     "duration": 0.14437,
     "end_time": "2024-06-17T18:18:36.717545",
     "exception": false,
     "start_time": "2024-06-17T18:18:36.573175",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "_wls = np.unique(wavelengths.ravel())\n",
    "_widths = np.unique(widths.ravel())\n",
    "plt.figure(figsize=(8, 3))\n",
    "plt.plot(_wls * 1000, neffs)\n",
    "plt.ylabel(\"neff\")\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.title(\"Effective Index\")\n",
    "plt.grid(True)\n",
    "plt.figlegend(\n",
    "    [f\"{w=:.2f}um\" for w in _widths], ncol=len(widths), bbox_to_anchor=(0.95, -0.05)\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71",
   "metadata": {
    "papermill": {
     "duration": 0.582858,
     "end_time": "2024-06-17T18:18:37.319771",
     "exception": false,
     "start_time": "2024-06-17T18:18:36.736913",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "_grid = [jnp.sort(jnp.unique(wavelengths)), jnp.sort(jnp.unique(widths))]\n",
    "_data = jnp.asarray(neffs)\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def _get_coordinate(arr1d: jnp.ndarray, value: jnp.ndarray):\n",
    "    return jnp.interp(value, arr1d, jnp.arange(arr1d.shape[0]))\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def _get_coordinates(arrs1d: list[jnp.ndarray], values: jnp.ndarray):\n",
    "    # don't use vmap as arrays in arrs1d could have different shapes...\n",
    "    return jnp.array([_get_coordinate(a, v) for a, v in zip(arrs1d, values)])\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def neff(wl=1.55, width=0.5):\n",
    "    params = jnp.stack(jnp.broadcast_arrays(jnp.asarray(wl), jnp.asarray(width)), 0)\n",
    "    coords = _get_coordinates(_grid, params)\n",
    "    return jax.scipy.ndimage.map_coordinates(_data, coords, 1, mode=\"nearest\")\n",
    "\n",
    "\n",
    "neff(wl=[1.52, 1.58], width=[0.5, 0.55])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72",
   "metadata": {
    "papermill": {
     "duration": 0.449589,
     "end_time": "2024-06-17T18:18:37.776110",
     "exception": false,
     "start_time": "2024-06-17T18:18:37.326521",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "wavelengths_ = np.linspace(wavelengths.min(), wavelengths.max(), 100)\n",
    "widths_ = np.linspace(widths.min(), widths.max(), 100)\n",
    "wavelengths_, widths_ = np.meshgrid(wavelengths_, widths_)\n",
    "neffs_ = neff(wavelengths_, widths_)\n",
    "\n",
    "plt.figure(figsize=(8, 3))\n",
    "plt.pcolormesh(wavelengths_ * 1000, widths_, neffs_)\n",
    "plt.ylabel(\"neff\")\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.title(\"Effective Index\")\n",
    "plt.grid(True)\n",
    "plt.figlegend(\n",
    "    [f\"{w=:.2f}um\" for w in _widths], ncol=len(_widths), bbox_to_anchor=(0.95, -0.05)\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73",
   "metadata": {},
   "outputs": [],
   "source": [
    "def straight(wl=1.55, length=10.0, width=0.5):\n",
    "    S = {\n",
    "        (\"o1\", \"o2\"): jnp.exp(2j * np.pi * neff(wl=wl, width=width) / wl * length),\n",
    "    }\n",
    "    return sax.reciprocal(S)\n",
    "\n",
    "\n",
    "def mmi1x2():\n",
    "    \"\"\"Assumes a perfect 1x2 splitter\"\"\"\n",
    "    return sax.reciprocal(\n",
    "        {\n",
    "            (\"o1\", \"o2\"): 0.5**0.5,\n",
    "            (\"o1\", \"o3\"): 0.5**0.5,\n",
    "        }\n",
    "    )\n",
    "\n",
    "\n",
    "def mmi2x2():\n",
    "    S = {\n",
    "        (\"o1\", \"o3\"): 0.5**0.5,\n",
    "        (\"o1\", \"o4\"): 1j * 0.5**0.5,\n",
    "        (\"o2\", \"o3\"): 1j * 0.5**0.5,\n",
    "        (\"o2\", \"o4\"): 0.5**0.5,\n",
    "    }\n",
    "    return sax.reciprocal(S)\n",
    "\n",
    "\n",
    "def bend_euler(wl=1.5, length=20.0, width=0.5):\n",
    "    \"\"\"Let's assume a reduced transmission for the euler bend compared to a straight.\"\"\"\n",
    "    return {k: 0.99 * v for k, v in straight(wl=wl, length=length, width=width).items()}\n",
    "\n",
    "\n",
    "models = {\n",
    "    \"bend_euler\": bend_euler,\n",
    "    \"mmi1x2\": mmi1x2,\n",
    "    \"mmi2x2\": mmi2x2,\n",
    "    \"straight\": straight,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74",
   "metadata": {},
   "source": [
    "Even though this still is lossless transmission, we're at least modeling the phase correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75",
   "metadata": {},
   "outputs": [],
   "source": [
    "straight()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76",
   "metadata": {
    "papermill": {
     "duration": 0.632755,
     "end_time": "2024-06-17T18:16:29.642798",
     "exception": false,
     "start_time": "2024-06-17T18:16:29.010043",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "mzi = simple_mzi()\n",
    "mzi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77",
   "metadata": {
    "papermill": {
     "duration": 0.197193,
     "end_time": "2024-06-17T18:18:38.343152",
     "exception": false,
     "start_time": "2024-06-17T18:18:38.145959",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "mzi2, _ = sax.circuit(mzi.get_netlist(recursive=True), models=models)\n",
    "mzi2()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78",
   "metadata": {
    "papermill": {
     "duration": 1.170838,
     "end_time": "2024-06-17T18:18:39.521611",
     "exception": false,
     "start_time": "2024-06-17T18:18:38.350773",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "wl = jnp.linspace(1.51, 1.59, 1000)\n",
    "S = mzi2(wl=wl)\n",
    "plt.plot(wl, abs(S[\"o1\", \"o2\"]) ** 2)\n",
    "plt.ylim(-0.05, 1.05)\n",
    "plt.xlabel(\"λ [μm]\")\n",
    "plt.ylabel(\"T\")\n",
    "plt.ylim(-0.05, 1.05)\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79",
   "metadata": {
    "papermill": {
     "duration": 0.008521,
     "end_time": "2024-06-17T18:18:39.538984",
     "exception": false,
     "start_time": "2024-06-17T18:18:39.530463",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Optimize MZI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80",
   "metadata": {
    "papermill": {
     "duration": 0.008512,
     "end_time": "2024-06-17T18:18:39.556123",
     "exception": false,
     "start_time": "2024-06-17T18:18:39.547611",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We'd like to optimize an MZI such that one of the minima is at 1530nm. To do this, we need to define a loss function for the circuit at 1530nm. This function should take the parameters that you want to optimize as positional arguments:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81",
   "metadata": {
    "papermill": {
     "duration": 0.014098,
     "end_time": "2024-06-17T18:18:39.578823",
     "exception": false,
     "start_time": "2024-06-17T18:18:39.564725",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def loss_fn(delta_length):\n",
    "    S = mzi2(\n",
    "        wl=1.53,\n",
    "        straight_top1={\"length\": delta_length / 2},\n",
    "        straight_top2={\"length\": delta_length / 2},\n",
    "    )\n",
    "    return jnp.mean(jnp.abs(S[\"o1\", \"o2\"]) ** 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82",
   "metadata": {
    "papermill": {
     "duration": 0.008572,
     "end_time": "2024-06-17T18:18:39.595608",
     "exception": false,
     "start_time": "2024-06-17T18:18:39.587036",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We can use this loss function to define a grad function which works on the parameters of the loss function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83",
   "metadata": {
    "papermill": {
     "duration": 0.013557,
     "end_time": "2024-06-17T18:18:39.617789",
     "exception": false,
     "start_time": "2024-06-17T18:18:39.604232",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "grad_fn = jax.jit(\n",
    "    jax.grad(\n",
    "        loss_fn,\n",
    "        argnums=0,  # JAX gradient function for the first positional argument, jitted\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84",
   "metadata": {
    "papermill": {
     "duration": 0.008576,
     "end_time": "2024-06-17T18:18:39.634681",
     "exception": false,
     "start_time": "2024-06-17T18:18:39.626105",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Next, we need to define a JAX optimizer, which on its own is nothing more than three more functions:  an initialization function with which to initialize the optimizer state, an update function which will update the optimizer state (and with it the model parameters). The third function that's being returned will give the model parameters given the optimizer state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85",
   "metadata": {
    "papermill": {
     "duration": 0.014029,
     "end_time": "2024-06-17T18:18:39.657569",
     "exception": false,
     "start_time": "2024-06-17T18:18:39.643540",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "initial_delta_length = 21.0\n",
    "init_fn, update_fn, params_fn = opt.adam(step_size=0.1)\n",
    "state = init_fn(initial_delta_length)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86",
   "metadata": {
    "papermill": {
     "duration": 0.008507,
     "end_time": "2024-06-17T18:18:39.674468",
     "exception": false,
     "start_time": "2024-06-17T18:18:39.665961",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Given all this, a single training step can be defined:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87",
   "metadata": {
    "papermill": {
     "duration": 0.013659,
     "end_time": "2024-06-17T18:18:39.696702",
     "exception": false,
     "start_time": "2024-06-17T18:18:39.683043",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def step_fn(step, state):\n",
    "    params = params_fn(state)\n",
    "    loss = loss_fn(params)\n",
    "    grad = grad_fn(params)\n",
    "    state = update_fn(step, grad, state)\n",
    "    return loss, state"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88",
   "metadata": {
    "papermill": {
     "duration": 0.00864,
     "end_time": "2024-06-17T18:18:39.713650",
     "exception": false,
     "start_time": "2024-06-17T18:18:39.705010",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "And we can use this step function to start the training of the MZI:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89",
   "metadata": {
    "papermill": {
     "duration": 7.102135,
     "end_time": "2024-06-17T18:18:46.824568",
     "exception": false,
     "start_time": "2024-06-17T18:18:39.722433",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "for step in (\n",
    "    pb := trange(300)\n",
    "):  # the first two iterations take a while because the circuit is being jitted...\n",
    "    loss, state = step_fn(step, state)\n",
    "    pb.set_postfix(loss=f\"{loss:.6f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90",
   "metadata": {
    "papermill": {
     "duration": 0.013994,
     "end_time": "2024-06-17T18:18:46.852453",
     "exception": false,
     "start_time": "2024-06-17T18:18:46.838459",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "delta_length = params_fn(state)\n",
    "delta_length"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91",
   "metadata": {
    "papermill": {
     "duration": 0.008721,
     "end_time": "2024-06-17T18:18:46.869581",
     "exception": false,
     "start_time": "2024-06-17T18:18:46.860860",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Let's see what we've got over a range of wavelengths:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92",
   "metadata": {
    "papermill": {
     "duration": 0.173121,
     "end_time": "2024-06-17T18:18:47.051669",
     "exception": false,
     "start_time": "2024-06-17T18:18:46.878548",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "wl = jnp.linspace(1.5, 1.6, 1000)\n",
    "S = mzi2(\n",
    "    wl=wl,\n",
    "    straight_top1={\"length\": delta_length / 2},\n",
    "    straight_top2={\"length\": delta_length / 2},\n",
    ")\n",
    "plt.plot(wl * 1e3, abs(S[\"o1\", \"o2\"]) ** 2)\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.ylabel(\"T\")\n",
    "plt.plot([1530, 1530], [-1, 2], ls=\":\", color=\"black\")\n",
    "plt.ylim(-0.05, 1.05)\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93",
   "metadata": {
    "papermill": {
     "duration": 0.008063,
     "end_time": "2024-06-17T18:18:47.067940",
     "exception": false,
     "start_time": "2024-06-17T18:18:47.059877",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Note that we could've just as well optimized the waveguide width:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94",
   "metadata": {
    "papermill": {
     "duration": 3.426675,
     "end_time": "2024-06-17T18:18:50.503118",
     "exception": false,
     "start_time": "2024-06-17T18:18:47.076443",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def loss_fn(width):\n",
    "    S = mzi2(\n",
    "        wl=1.53,\n",
    "        straight_top1={\"width\": width},\n",
    "        straight_top2={\"width\": width},\n",
    "    )\n",
    "    return jnp.mean(jnp.abs(S[\"o1\", \"o2\"]) ** 2)\n",
    "\n",
    "\n",
    "grad_fn = jax.jit(\n",
    "    jax.grad(\n",
    "        loss_fn,\n",
    "        argnums=0,  # JAX gradient function for the first positional argument, jitted\n",
    "    )\n",
    ")\n",
    "initial_width = 0.5\n",
    "init_fn, update_fn, params_fn = opt.adam(step_size=0.01)\n",
    "state = init_fn(initial_width)\n",
    "for step in (\n",
    "    pb := trange(300)\n",
    "):  # the first two iterations take a while because the circuit is being jitted...\n",
    "    loss, state = step_fn(step, state)\n",
    "    pb.set_postfix(loss=f\"{loss:.6f}\")\n",
    "\n",
    "optim_width = params_fn(state)\n",
    "S = Sw = mzi2(\n",
    "    wl=wl,\n",
    "    straight_top1={\"width\": optim_width},\n",
    "    straight_top2={\"width\": optim_width},\n",
    ")\n",
    "plt.plot(wl * 1e3, abs(S[\"o1\", \"o2\"]) ** 2)\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.ylabel(\"T\")\n",
    "plt.plot([1530, 1530], [-1, 2], color=\"black\", ls=\":\")\n",
    "plt.ylim(-0.05, 1.05)\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95",
   "metadata": {
    "papermill": {
     "duration": 0.009255,
     "end_time": "2024-06-17T18:18:50.521343",
     "exception": false,
     "start_time": "2024-06-17T18:18:50.512088",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Layout-aware Monte Carlo\n",
    "\n",
    "Let's assume the waveguide width changes with a certain correlation length. We can create a 'wafermap' of width variations by randomly varying the width and low pass filtering with a spatial frequency being the inverse of the correlation length (there are probably better ways to do this, but this works for this tutorial)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96",
   "metadata": {
    "papermill": {
     "duration": 0.023728,
     "end_time": "2024-06-17T18:18:50.555103",
     "exception": false,
     "start_time": "2024-06-17T18:18:50.531375",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def create_wafermaps(\n",
    "    placements, correlation_length=1.0, num_maps=1, mean=0.0, std=1.0, seed=None\n",
    "):\n",
    "    dx = dy = correlation_length / 200\n",
    "    xs = [p[\"x\"] for p in placements.values()]\n",
    "    ys = [p[\"y\"] for p in placements.values()]\n",
    "    xmin, xmax, ymin, ymax = min(xs), max(xs), min(ys), max(ys)\n",
    "    wx, wy = xmax - xmin, ymax - ymin\n",
    "    xmin, xmax, ymin, ymax = xmin - wx, xmax + wx, ymin - wy, ymax + wy\n",
    "    x, y = np.arange(xmin, xmax + dx, dx), np.arange(ymin, ymax + dy, dy)\n",
    "    if seed is None:\n",
    "        r = np.random\n",
    "    else:\n",
    "        r = np.random.RandomState(seed=seed)\n",
    "    W0 = r.randn(num_maps, x.shape[0], y.shape[0])\n",
    "\n",
    "    fx = fftshift(fftfreq(x.shape[0], d=x[1] - x[0]))\n",
    "    fy = fftshift(fftfreq(y.shape[0], d=y[1] - y[0]))\n",
    "    fY, fX = np.meshgrid(fy, fx)\n",
    "\n",
    "    fW = fftshift(fft2(W0))\n",
    "\n",
    "    if correlation_length >= min(x.shape[0], y.shape[0]):\n",
    "        fW = np.zeros_like(fW)\n",
    "    else:\n",
    "        fW = np.where(np.sqrt(fX**2 + fY**2)[None] > 1 / correlation_length, 0, fW)\n",
    "\n",
    "    W = np.abs(fftshift(ifft2(fW))) ** 2\n",
    "\n",
    "    mean_ = W.mean(1, keepdims=True).mean(2, keepdims=True)\n",
    "\n",
    "    std_ = W.std(1, keepdims=True).std(2, keepdims=True)\n",
    "    if (std_ == 0).all():\n",
    "        std_ = 1\n",
    "\n",
    "    W = (W - mean_) / std_\n",
    "\n",
    "    W = W * std + mean\n",
    "\n",
    "    return x, y, W"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97",
   "metadata": {
    "papermill": {
     "duration": 2.890746,
     "end_time": "2024-06-17T18:18:53.460310",
     "exception": false,
     "start_time": "2024-06-17T18:18:50.569564",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "placements = mzi.get_netlist()[\"placements\"]\n",
    "xm, ym, wmaps = create_wafermaps(\n",
    "    placements,\n",
    "    correlation_length=100,\n",
    "    mean=0.5,\n",
    "    std=0.002,\n",
    "    num_maps=100,\n",
    "    seed=42,\n",
    ")\n",
    "\n",
    "for i, wmap in enumerate(wmaps):\n",
    "    if i > 1:\n",
    "        break\n",
    "    plt.imshow(wmap, cmap=\"RdBu\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98",
   "metadata": {
    "papermill": {
     "duration": 0.029908,
     "end_time": "2024-06-17T18:18:53.535052",
     "exception": false,
     "start_time": "2024-06-17T18:18:53.505144",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def widths(xw, yw, wmaps, x, y):\n",
    "    _wmap_grid = [xw, yw]\n",
    "    params = jnp.stack(jnp.broadcast_arrays(jnp.asarray(x), jnp.asarray(y)), 0)\n",
    "    coords = _get_coordinates(_wmap_grid, params)\n",
    "\n",
    "    map_coordinates = partial(\n",
    "        jax.scipy.ndimage.map_coordinates, coordinates=coords, order=1, mode=\"nearest\"\n",
    "    )\n",
    "    w = jax.vmap(map_coordinates)(wmaps)\n",
    "    return w"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99",
   "metadata": {
    "papermill": {
     "duration": 0.01214,
     "end_time": "2024-06-17T18:18:53.565772",
     "exception": false,
     "start_time": "2024-06-17T18:18:53.553632",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Let's now sample the MZI width variation on the wafer map (let's assume a single width variation per point):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "100",
   "metadata": {
    "papermill": {
     "duration": 3.28009,
     "end_time": "2024-06-17T18:18:56.856660",
     "exception": false,
     "start_time": "2024-06-17T18:18:53.576570",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "mzi_params = sax.get_settings(mzi2)\n",
    "placements = mzi.get_netlist()[\"placements\"]\n",
    "width_params = {\n",
    "    k: {\"width\": widths(xm, ym, wmaps, v[\"x\"], v[\"y\"])}\n",
    "    for k, v in placements.items()\n",
    "    if \"width\" in mzi_params[k]\n",
    "}\n",
    "\n",
    "S0 = mzi2(wl=wl)\n",
    "S = mzi2(\n",
    "    wl=wl[:, None],\n",
    "    **width_params,\n",
    ")\n",
    "ps = plt.plot(wl * 1e3, abs(S[\"o1\", \"o2\"]) ** 2, color=\"C0\", lw=1, alpha=0.1)\n",
    "nps = plt.plot(wl * 1e3, abs(S0[\"o1\", \"o2\"]) ** 2, color=\"C1\", lw=2, alpha=1)\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.ylabel(\"T\")\n",
    "plt.plot([1550, 1550], [-1, 2], color=\"black\", ls=\":\")\n",
    "plt.ylim(-0.05, 1.05)\n",
    "plt.grid(True)\n",
    "plt.figlegend([*ps[-1:], *nps], [\"MC\", \"nominal\"], bbox_to_anchor=(1.1, 0.9))\n",
    "rmse = jnp.mean(\n",
    "    jnp.abs(jnp.abs(S[\"o1\", \"o2\"]) ** 2 - jnp.abs(S0[\"o1\", \"o2\"][:, None]) ** 2) ** 2\n",
    ")\n",
    "plt.title(f\"{rmse=}\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "101",
   "metadata": {
    "papermill": {
     "duration": 0.01484,
     "end_time": "2024-06-17T18:18:56.885934",
     "exception": false,
     "start_time": "2024-06-17T18:18:56.871094",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Compact MZI\n",
    "\n",
    "Let's see if we can improve variability (i.e. the RMSE w.r.t. nominal) by making the MZI more compact:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "102",
   "metadata": {
    "papermill": {
     "duration": 0.025324,
     "end_time": "2024-06-17T18:18:56.926241",
     "exception": false,
     "start_time": "2024-06-17T18:18:56.900917",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "@gf.cell\n",
    "def compact_mzi():\n",
    "    c = gf.Component()\n",
    "\n",
    "    # instances\n",
    "    mmi_in = gf.components.mmi1x2()\n",
    "    mmi_out = gf.components.mmi2x2()\n",
    "    bend = gf.components.bend_euler()\n",
    "    half_delay_straight = gf.components.straight()\n",
    "    middle_straight = gf.components.straight(length=6.0)\n",
    "    half_middle_straight = gf.components.straight(3.0)\n",
    "\n",
    "    # references (sax convention: vars ending in underscore are references)\n",
    "    mmi_in = c.add_ref(mmi_in, name=\"mmi_in\")\n",
    "\n",
    "    bend_top1 = c.add_ref(bend, name=\"bend_top1\")\n",
    "    straight_top1 = c.add_ref(half_delay_straight, name=\"straight_top1\")\n",
    "    bend_top2 = c.add_ref(bend, name=\"bend_top2\").dmirror()\n",
    "    straight_top2 = c.add_ref(middle_straight, name=\"straight_top2\")\n",
    "    bend_top3 = c.add_ref(bend, name=\"bend_top3\").dmirror()\n",
    "    straight_top3 = c.add_ref(half_delay_straight, name=\"straight_top3\")\n",
    "    bend_top4 = c.add_ref(bend, name=\"bend_top4\")\n",
    "\n",
    "    straight_btm1 = c.add_ref(half_middle_straight, name=\"straight_btm1\")\n",
    "    bend_btm1 = c.add_ref(bend, name=\"bend_btm1\")\n",
    "    bend_btm2 = c.add_ref(bend, name=\"bend_btm2\").dmirror()\n",
    "    bend_btm3 = c.add_ref(bend, name=\"bend_btm3\").dmirror()\n",
    "    bend_btm4 = c.add_ref(bend, name=\"bend_btm4\")\n",
    "    straight_btm2 = c.add_ref(half_middle_straight, name=\"straight_btm2\")\n",
    "\n",
    "    mmi_out = c.add_ref(mmi_out, name=\"mmi_out\")\n",
    "\n",
    "    # connections\n",
    "    bend_top1.connect(\"o1\", mmi_in.ports[\"o2\"])\n",
    "    straight_top1.connect(\"o1\", bend_top1.ports[\"o2\"])\n",
    "    bend_top2.connect(\"o1\", straight_top1.ports[\"o2\"])\n",
    "    straight_top2.connect(\"o1\", bend_top2.ports[\"o2\"])\n",
    "    bend_top3.connect(\"o1\", straight_top2.ports[\"o2\"])\n",
    "    straight_top3.connect(\"o1\", bend_top3.ports[\"o2\"])\n",
    "    bend_top4.connect(\"o1\", straight_top3.ports[\"o2\"])\n",
    "\n",
    "    straight_btm1.connect(\"o1\", mmi_in.ports[\"o3\"])\n",
    "    bend_btm1.connect(\"o1\", straight_btm1.ports[\"o2\"])\n",
    "    bend_btm2.connect(\"o1\", bend_btm1.ports[\"o2\"])\n",
    "    bend_btm3.connect(\"o1\", bend_btm2.ports[\"o2\"])\n",
    "    bend_btm4.connect(\"o1\", bend_btm3.ports[\"o2\"])\n",
    "    straight_btm2.connect(\"o1\", bend_btm4.ports[\"o2\"])\n",
    "\n",
    "    mmi_out.connect(\"o1\", straight_btm2.ports[\"o2\"])\n",
    "\n",
    "    # ports\n",
    "    c.add_port(\n",
    "        \"o1\",\n",
    "        port=mmi_in.ports[\"o1\"],\n",
    "    )\n",
    "    c.add_port(\"o2\", port=mmi_out.ports[\"o3\"])\n",
    "    c.add_port(\"o3\", port=mmi_out.ports[\"o4\"])\n",
    "    return c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "103",
   "metadata": {
    "papermill": {
     "duration": 0.212473,
     "end_time": "2024-06-17T18:18:57.153655",
     "exception": false,
     "start_time": "2024-06-17T18:18:56.941182",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "compact_mzi1 = compact_mzi()\n",
    "compact_mzi1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "104",
   "metadata": {
    "papermill": {
     "duration": 0.625889,
     "end_time": "2024-06-17T18:18:57.793746",
     "exception": false,
     "start_time": "2024-06-17T18:18:57.167857",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "placements = compact_mzi1.get_netlist()[\"placements\"]\n",
    "mzi3, _ = sax.circuit(compact_mzi1.get_netlist(recursive=True), models=models)\n",
    "mzi3()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "105",
   "metadata": {
    "papermill": {
     "duration": 2.0279,
     "end_time": "2024-06-17T18:18:59.836722",
     "exception": false,
     "start_time": "2024-06-17T18:18:57.808822",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "mzi_params = sax.get_settings(mzi3)\n",
    "placements = compact_mzi1.get_netlist()[\"placements\"]\n",
    "width_params = {\n",
    "    k: {\"width\": widths(xm, ym, wmaps, v[\"x\"], v[\"y\"])}\n",
    "    for k, v in placements.items()\n",
    "    if \"width\" in mzi_params[k]\n",
    "}\n",
    "\n",
    "S0 = mzi3(wl=wl)\n",
    "S = mzi3(\n",
    "    wl=wl[:, None],\n",
    "    **width_params,\n",
    ")\n",
    "ps = plt.plot(wl * 1e3, abs(S[\"o1\", \"o2\"]) ** 2, color=\"C0\", lw=1, alpha=0.1)\n",
    "nps = plt.plot(wl * 1e3, abs(S0[\"o1\", \"o2\"]) ** 2, color=\"C1\", lw=2, alpha=1)\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.ylabel(\"T\")\n",
    "plt.plot([1550, 1550], [-1, 2], color=\"black\", ls=\":\")\n",
    "plt.ylim(-0.05, 1.05)\n",
    "plt.grid(True)\n",
    "plt.figlegend([*ps[-1:], *nps], [\"MC\", \"nominal\"], bbox_to_anchor=(1.1, 0.9))\n",
    "rmse = jnp.mean(\n",
    "    jnp.abs(jnp.abs(S[\"o1\", \"o2\"]) ** 2 - jnp.abs(S0[\"o1\", \"o2\"][:, None]) ** 2) ** 2\n",
    ")\n",
    "plt.title(f\"{rmse=}\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "106",
   "metadata": {},
   "source": [
    "## Phase shifter model\n",
    "\n",
    "You can create a phase shifter model that depends on the applied volage.\n",
    "For that you need first to figure out what's the phase shift for different voltages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "107",
   "metadata": {},
   "outputs": [],
   "source": [
    "delta_length = 10\n",
    "mzi_component = gf.components.mzi_phase_shifter(delta_length=delta_length)\n",
    "mzi_component.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "108",
   "metadata": {},
   "outputs": [],
   "source": [
    "def straight(wl=1.5, length=10.0, neff=2.4) -> sax.SDict:\n",
    "    return sax.reciprocal({(\"o1\", \"o2\"): jnp.exp(2j * jnp.pi * neff * length / wl)})\n",
    "\n",
    "\n",
    "def mmi1x2() -> sax.SDict:\n",
    "    \"\"\"Returns a perfect 1x2 splitter.\"\"\"\n",
    "    return sax.reciprocal(\n",
    "        {\n",
    "            (\"o1\", \"o2\"): 0.5**0.5,\n",
    "            (\"o1\", \"o3\"): 0.5**0.5,\n",
    "        }\n",
    "    )\n",
    "\n",
    "\n",
    "def bend_euler(wl=1.5, length=20.0) -> sax.SDict:\n",
    "    \"\"\"Returns bend Sparameters with reduced transmission compared to a straight.\"\"\"\n",
    "    return {k: 0.99 * v for k, v in straight(wl=wl, length=length).items()}\n",
    "\n",
    "\n",
    "def phase_shifter_heater(\n",
    "    wl: float = 1.55,\n",
    "    neff: float = 2.34,\n",
    "    voltage: float = 0,\n",
    "    length: float = 10,\n",
    "    loss: float = 0.0,\n",
    ") -> sax.SDict:\n",
    "    \"\"\"Returns simple phase shifter model\"\"\"\n",
    "    deltaphi = voltage * jnp.pi\n",
    "    phase = 2 * jnp.pi * neff * length / wl + deltaphi\n",
    "    amplitude = jnp.asarray(10 ** (-loss * length / 20), dtype=complex)\n",
    "    transmission = amplitude * jnp.exp(1j * phase)\n",
    "    return sax.reciprocal(\n",
    "        {\n",
    "            (\"o1\", \"o2\"): transmission,\n",
    "            (\"l_e1\", \"r_e1\"): 0.0,\n",
    "            (\"l_e2\", \"r_e2\"): 0.0,\n",
    "            (\"l_e3\", \"r_e3\"): 0.0,\n",
    "            (\"l_e4\", \"r_e4\"): 0.0,\n",
    "        }\n",
    "    )\n",
    "\n",
    "\n",
    "models = {\n",
    "    \"bend_euler\": bend_euler,\n",
    "    \"mmi1x2\": mmi1x2,\n",
    "    \"straight\": straight,\n",
    "    \"straight_heater_metal_undercut\": phase_shifter_heater,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "109",
   "metadata": {},
   "outputs": [],
   "source": [
    "mzi_component = gf.components.mzi_phase_shifter(delta_length=delta_length)\n",
    "netlist = sax.netlist(mzi_component.get_netlist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "110",
   "metadata": {},
   "outputs": [],
   "source": [
    "mzi_circuit, _ = sax.circuit(netlist=netlist, models=models, backend=\"filipsson_gunnar\")\n",
    "S = mzi_circuit(wl=1.55)\n",
    "{k: v for k, v in S.items() if abs(v) > 1e-5}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "111",
   "metadata": {},
   "outputs": [],
   "source": [
    "wl = np.linspace(1.5, 1.6, 256)\n",
    "S = mzi_circuit(wl=wl)\n",
    "\n",
    "plt.figure(figsize=(14, 4))\n",
    "plt.title(\"MZI\")\n",
    "plt.plot(1e3 * wl, jnp.abs(S[\"o1\", \"o2\"]) ** 2)\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.ylabel(\"T\")\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "112",
   "metadata": {},
   "source": [
    "Now you can tune the phase shift applied to one of the arms.\n",
    "\n",
    "How do you find out what's the name of the netlist component that you want to tune?\n",
    "\n",
    "You can backannotate the netlist and read the labels on the backannotated netlist or you can plot the netlist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "113",
   "metadata": {},
   "outputs": [],
   "source": [
    "mzi_component.plot_netlist()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "114",
   "metadata": {},
   "source": [
    "As you can see the top phase shifter instance name `sxt` is hard to see on the netlist.\n",
    "You can also reconstruct the component using the netlist and look at the labels in klayout."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "115",
   "metadata": {},
   "outputs": [],
   "source": [
    "mzi_yaml = mzi_component.get_netlist()\n",
    "mzi_component2 = gf.read.from_yaml(mzi_yaml)\n",
    "mzi_component2.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "116",
   "metadata": {},
   "source": [
    "The best way to get a deterministic name of the `instance` is naming the reference on your Pcell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "117",
   "metadata": {},
   "outputs": [],
   "source": [
    "voltages = np.linspace(-1, 1, num=5)\n",
    "voltages = [-0.5, 0, 0.5]\n",
    "\n",
    "for voltage in voltages:\n",
    "    S = mzi_circuit(\n",
    "        wl=wl,\n",
    "        sxt={\"voltage\": voltage},\n",
    "    )\n",
    "    plt.plot(wl * 1e3, abs(S[\"o1\", \"o2\"]) ** 2, label=str(voltage))\n",
    "    plt.xlabel(\"λ [nm]\")\n",
    "    plt.ylabel(\"T\")\n",
    "    plt.ylim(-0.05, 1.05)\n",
    "    plt.grid(True)\n",
    "\n",
    "plt.title(\"MZI vs voltage\")\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "118",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## Hierarchical circuits\n",
    "\n",
    "You can also simulate hierarchical circuits, such as lattice of MZI interferometers.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "119",
   "metadata": {},
   "outputs": [],
   "source": [
    "@gf.cell\n",
    "def mzis(delta_length=10):\n",
    "    c = gf.Component()\n",
    "    c1 = c << simple_mzi(delta_length=delta_length)\n",
    "    c2 = c << simple_mzi(delta_length=delta_length)\n",
    "    c2.connect(\"o1\", c1.ports[\"o2\"])\n",
    "\n",
    "    c.add_port(\"o1\", port=c1.ports[\"o1\"])\n",
    "    c.add_port(\"o2\", port=c2.ports[\"o2\"])\n",
    "    return c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "120",
   "metadata": {},
   "outputs": [],
   "source": [
    "def straight(wl=1.55, length=10.0, width=0.5):\n",
    "    S = {\n",
    "        (\"o1\", \"o2\"): jnp.exp(2j * np.pi * neff(wl=wl, width=width) / wl * length),\n",
    "    }\n",
    "    return sax.reciprocal(S)\n",
    "\n",
    "\n",
    "def mmi1x2():\n",
    "    \"\"\"Assumes a perfect 1x2 splitter\"\"\"\n",
    "    return sax.reciprocal(\n",
    "        {\n",
    "            (\"o1\", \"o2\"): 0.5**0.5,\n",
    "            (\"o1\", \"o3\"): 0.5**0.5,\n",
    "        }\n",
    "    )\n",
    "\n",
    "\n",
    "def mmi2x2():\n",
    "    S = {\n",
    "        (\"o1\", \"o3\"): 0.5**0.5,\n",
    "        (\"o1\", \"o4\"): 1j * 0.5**0.5,\n",
    "        (\"o2\", \"o3\"): 1j * 0.5**0.5,\n",
    "        (\"o2\", \"o4\"): 0.5**0.5,\n",
    "    }\n",
    "    return sax.reciprocal(S)\n",
    "\n",
    "\n",
    "def bend_euler(wl=1.5, length=20.0, width=0.5):\n",
    "    \"\"\"Let's assume a reduced transmission for the euler bend compared to a straight.\"\"\"\n",
    "    return {k: 0.99 * v for k, v in straight(wl=wl, length=length, width=width).items()}\n",
    "\n",
    "\n",
    "models = {\n",
    "    \"bend_euler\": bend_euler,\n",
    "    \"mmi1x2\": mmi1x2,\n",
    "    \"mmi2x2\": mmi2x2,\n",
    "    \"straight\": straight,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "121",
   "metadata": {},
   "outputs": [],
   "source": [
    "c2 = mzis()\n",
    "c2.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "122",
   "metadata": {},
   "outputs": [],
   "source": [
    "c2.plot_netlist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "123",
   "metadata": {},
   "outputs": [],
   "source": [
    "c1 = simple_mzi(delta_length=10)\n",
    "c1.plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "124",
   "metadata": {},
   "outputs": [],
   "source": [
    "c1.plot_netlist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "125",
   "metadata": {},
   "outputs": [],
   "source": [
    "wl = np.linspace(1.5, 1.6)\n",
    "netlist1 = c1.get_netlist(recursive=True)\n",
    "circuit1, _ = sax.circuit(netlist=netlist1, models=models)\n",
    "S1 = circuit1(wl=wl)\n",
    "\n",
    "netlist2 = c2.get_netlist(recursive=True)\n",
    "circuit2, _ = sax.circuit(netlist=netlist2, models=models)\n",
    "S2 = circuit2(wl=wl)\n",
    "\n",
    "plt.figure(figsize=(14, 4))\n",
    "plt.plot(1e3 * wl, jnp.abs(S1[\"o1\", \"o2\"]) ** 2, label=\"1 MZI\")\n",
    "plt.plot(1e3 * wl, jnp.abs(S2[\"o1\", \"o2\"]) ** 2, label=\"2 MZI\")\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.ylabel(\"T\")\n",
    "plt.grid(True)\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "custom_cell_magics": "kql"
  },
  "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}