{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {
    "papermill": {
     "duration": 0.005932,
     "end_time": "2024-06-17T18:16:22.415701",
     "exception": false,
     "start_time": "2024-06-17T18:16:22.409769",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Forward-only Backend"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {
    "papermill": {
     "duration": 0.004257,
     "end_time": "2024-06-17T18:16:22.424876",
     "exception": false,
     "start_time": "2024-06-17T18:16:22.420619",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "> An efficient backend when the components in circuit have low back-reflection. In this case, only forward-direction matrix multiplication is calculated to reduce computational cost for complicated circuit. The improvement compared to other backends ('klu' for example) is demonstrated with an example of cascaded AMZI structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {
    "papermill": {
     "duration": 1.699982,
     "end_time": "2024-06-17T18:16:24.129137",
     "exception": false,
     "start_time": "2024-06-17T18:16:22.429155",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import sax\n",
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "> Define waveguide and directional coupler components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {
    "papermill": {
     "duration": 0.015845,
     "end_time": "2024-06-17T18:16:24.148047",
     "exception": false,
     "start_time": "2024-06-17T18:16:24.132202",
     "status": "completed"
    },
    "tags": []
   },
   "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",
    "    sdict = sax.reciprocal(\n",
    "        {\n",
    "            (\"in0\", \"out0\"): transmission,\n",
    "        }\n",
    "    )\n",
    "    return sdict\n",
    "\n",
    "\n",
    "def coupler(coupling=0.5) -> sax.SDict:\n",
    "    kappa = coupling**0.5\n",
    "    tau = (1 - coupling) ** 0.5\n",
    "    coupler_dict = sax.reciprocal(\n",
    "        {\n",
    "            (\"in0\", \"out0\"): tau,\n",
    "            (\"in0\", \"out1\"): 1j * kappa,\n",
    "            (\"in1\", \"out0\"): 1j * kappa,\n",
    "            (\"in1\", \"out1\"): tau,\n",
    "        }\n",
    "    )\n",
    "    return coupler_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "> Create a cascaded AMZI structure with n identical AMZIs, with the next AMZI is connected to the \"cross\" port of the previous one. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {
    "papermill": {
     "duration": 1.050787,
     "end_time": "2024-06-17T18:16:25.201607",
     "exception": false,
     "start_time": "2024-06-17T18:16:24.150820",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def cascaded_amzi_generator(n, backend=\"klu\"):\n",
    "    netlist = {\n",
    "        \"instances\": {},\n",
    "        \"connections\": {},\n",
    "        \"ports\": {},\n",
    "    }\n",
    "\n",
    "    models = {\n",
    "        \"coupler\": coupler,\n",
    "        \"waveguide\": waveguide,\n",
    "    }\n",
    "\n",
    "    # Build the netlist\n",
    "    for i in range(1, n + 1):\n",
    "        # Define instance names\n",
    "        left_name = f\"left_{i}\"\n",
    "        right_name = f\"right_{i}\"\n",
    "        top_name = f\"top_{i}\"\n",
    "        btm_name = f\"btm_{i}\"\n",
    "\n",
    "        # Add instances\n",
    "        netlist[\"instances\"][left_name] = \"coupler\"\n",
    "        netlist[\"instances\"][right_name] = \"coupler\"\n",
    "        netlist[\"instances\"][top_name] = \"waveguide\"\n",
    "        netlist[\"instances\"][btm_name] = \"waveguide\"\n",
    "\n",
    "        # Internal connections within AMZI i\n",
    "        netlist[\"connections\"][f\"{left_name},out0\"] = f\"{btm_name},in0\"\n",
    "        netlist[\"connections\"][f\"{btm_name},out0\"] = f\"{right_name},in0\"\n",
    "\n",
    "        netlist[\"connections\"][f\"{left_name},out1\"] = f\"{top_name},in0\"\n",
    "        netlist[\"connections\"][f\"{top_name},out0\"] = f\"{right_name},in1\"\n",
    "\n",
    "        if i > 1:\n",
    "            # Connections between AMZIs via the cross port\n",
    "            prev_right_name = f\"right_{i - 1}\"\n",
    "            netlist[\"connections\"][f\"{prev_right_name},out0\"] = f\"{left_name},in1\"\n",
    "            netlist[\"connections\"][f\"{prev_right_name},out1\"] = f\"{left_name},in0\"\n",
    "\n",
    "    # Define external ports\n",
    "    netlist[\"ports\"][\"in0\"] = \"left_1,in0\"\n",
    "    netlist[\"ports\"][\"in1\"] = \"left_1,in1\"\n",
    "    netlist[\"ports\"][\"out0\"] = f\"right_{n},out0\"\n",
    "    netlist[\"ports\"][\"out1\"] = f\"right_{n},out1\"\n",
    "\n",
    "    # Create the circuit\n",
    "    mzi_ideal, info = sax.circuit(netlist=netlist, models=models, backend=backend)\n",
    "    return mzi_ideal, info"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "> Consider a case with 20 cascaded AMZIs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "wavelengths = np.linspace(1.500, 1.600, 100_000)\n",
    "n = 10\n",
    "params = {\"wl\": wavelengths}\n",
    "for i in range(1, n + 1):\n",
    "    params[f\"left_{i}\"] = {\"coupling\": 0.5}\n",
    "    params[f\"right_{i}\"] = {\"coupling\": 0.5}\n",
    "    params[f\"top_{i}\"] = {\"length\": 50}\n",
    "    params[f\"btm_{i}\"] = {\"length\": 0}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "> Run simulation with 'klu' backend and 'forward-only' backend"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "mzi_ideal, info = cascaded_amzi_generator(n, backend=\"klu\")\n",
    "S = mzi_ideal(**params)\n",
    "transmissions_klu = 10 * jnp.log10(jnp.abs(S[\"in0\", \"out1\"]) ** 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "mzi_ideal, info = cascaded_amzi_generator(n, backend=\"forward\")\n",
    "S = mzi_ideal(**params)\n",
    "transmissions_forward = 10 * jnp.log10(jnp.abs(S[\"in0\", \"out1\"]) ** 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "> As the circuit becomes more complex, the forward-only backend experiences a more significant speed-up, and the results remain accurate as long as there is no backreflection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(figsize=(6, 4))\n",
    "\n",
    "plt.plot(wavelengths, transmissions_klu, label=\"klu\")\n",
    "plt.plot(wavelengths, transmissions_forward, \"--\", label=\"forward\")\n",
    "\n",
    "plt.xlim(1.53, 1.57)\n",
    "plt.xlabel(\"wavelength (um)\")\n",
    "plt.ylabel(\"transmission (dB)\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}