{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "textile-panel",
   "metadata": {
    "papermill": {
     "duration": 0.003438,
     "end_time": "2023-09-03T04:40:07.832874",
     "exception": false,
     "start_time": "2023-09-03T04:40:07.829436",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Multimode simulations\n",
    "> SAX can handle multiple modes too!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "prescribed-plant",
   "metadata": {
    "papermill": {
     "duration": 1.285171,
     "end_time": "2023-09-03T04:40:09.121177",
     "exception": false,
     "start_time": "2023-09-03T04:40:07.836006",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from itertools import combinations_with_replacement, product\n",
    "\n",
    "import jax.numpy as jnp\n",
    "import sax"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "requested-calvin",
   "metadata": {
    "papermill": {
     "duration": 0.003648,
     "end_time": "2023-09-03T04:40:09.128304",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.124656",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Ports and modes per port"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "grateful-programming",
   "metadata": {
    "papermill": {
     "duration": 0.003738,
     "end_time": "2023-09-03T04:40:09.135441",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.131703",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Let's denote a combination of a port and a mode by a string of the following format: `\"{port}@{mode}\"`. We can obtain all possible port-mode combinations with some magic itertools functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "going-entity",
   "metadata": {
    "papermill": {
     "duration": 0.015043,
     "end_time": "2023-09-03T04:40:09.155035",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.139992",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "ports = [\"in0\", \"out0\"]\n",
    "modes = [\"TE\", \"TM\"]\n",
    "portmodes = [\n",
    "    (f\"{p1}@{m1}\", f\"{p2}@{m2}\")\n",
    "    for (p1, m1), (p2, m2) in combinations_with_replacement(product(ports, modes), 2)\n",
    "]\n",
    "portmodes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "passive-selling",
   "metadata": {
    "papermill": {
     "duration": 0.003683,
     "end_time": "2023-09-03T04:40:09.162469",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.158786",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "If we would disregard any backreflection, this can be further simplified:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "numerical-novel",
   "metadata": {
    "papermill": {
     "duration": 0.012628,
     "end_time": "2023-09-03T04:40:09.206617",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.193989",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "portmodes_without_backreflection = [\n",
    "    (p1, p2) for p1, p2 in portmodes if p1.split(\"@\")[0] != p2.split(\"@\")[0]\n",
    "]\n",
    "portmodes_without_backreflection"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "hindu-application",
   "metadata": {
    "papermill": {
     "duration": 0.003869,
     "end_time": "2023-09-03T04:40:09.214671",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.210802",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Sometimes cross-polarization terms can also be ignored:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "encouraging-territory",
   "metadata": {
    "papermill": {
     "duration": 0.014551,
     "end_time": "2023-09-03T04:40:09.233025",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.218474",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "portmodes_without_crosspolarization = [\n",
    "    (p1, p2) for p1, p2 in portmodes if p1.split(\"@\")[1] == p2.split(\"@\")[1]\n",
    "]\n",
    "portmodes_without_crosspolarization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "mathematical-default",
   "metadata": {
    "papermill": {
     "duration": 0.020811,
     "end_time": "2023-09-03T04:40:09.258247",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.237436",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Multimode waveguide"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "asian-noise",
   "metadata": {
    "papermill": {
     "duration": 0.003737,
     "end_time": "2023-09-03T04:40:09.265759",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.262022",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Let's create a waveguide with two ports (`\"in\"`, `\"out\"`) and two modes (`\"te\"`, `\"tm\"`) without backreflection. Let's assume there is 5% cross-polarization and that the `\"tm\"`->`\"tm\"` transmission is 10% worse than the `\"te\"`->`\"te\"` transmission. Naturally in more realisic waveguide models these percentages will be length-dependent, but this is just a dummy model serving as an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "patient-wealth",
   "metadata": {
    "papermill": {
     "duration": 0.145676,
     "end_time": "2023-09-03T04:40:09.415317",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.269641",
     "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):\n",
    "    \"\"\"a simple straight waveguide model\n",
    "\n",
    "    Args:\n",
    "        wl: wavelength\n",
    "        neff: waveguide effective index\n",
    "        ng: waveguide group index (used for linear neff dispersion)\n",
    "        wl0: center wavelength at which neff is defined\n",
    "        length: [m] wavelength length\n",
    "        loss: [dB/m] waveguide loss\n",
    "    \"\"\"\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@TE\", \"out0@TE\"): 0.95 * transmission,  # 5% lost to cross-polarization\n",
    "            (\"in0@TE\", \"out0@TM\"): 0.05 * transmission,  # 5% cross-polarization\n",
    "            (\"in0@TM\", \"out0@TM\"): 0.85 * transmission,  # 10% worse tm->tm than te->te\n",
    "            (\"in0@TM\", \"out0@TE\"): 0.05 * transmission,  # 5% cross-polarization\n",
    "        }\n",
    "    )\n",
    "    return sdict\n",
    "\n",
    "\n",
    "waveguide()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cafacbde-e5eb-43c6-87ec-1a9e795b5cc2",
   "metadata": {
    "papermill": {
     "duration": 0.003511,
     "end_time": "2023-09-03T04:40:09.422649",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.419138",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Multimode Coupler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cf6d7a6-d054-4c3a-9053-3b09777d50f9",
   "metadata": {
    "papermill": {
     "duration": 0.022496,
     "end_time": "2023-09-03T04:40:09.448864",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.426368",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def coupler():\n",
    "    return {\n",
    "        (\"in0@TE\", \"out0@TE\"): 0.45**0.5,\n",
    "        (\"in0@TE\", \"out1@TE\"): 1j*0.45**0.5,\n",
    "        (\"in1@TE\", \"out0@TE\"): 1j*0.45**0.5,\n",
    "        (\"in1@TE\", \"out1@TE\"): 0.45**0.5,\n",
    "        (\"in0@TM\", \"out0@TM\"): 0.45**0.5,\n",
    "        (\"in0@TM\", \"out1@TM\"): 1j*0.45**0.5,\n",
    "        (\"in1@TM\", \"out0@TM\"): 1j*0.45**0.5,\n",
    "        (\"in1@TM\", \"out1@TM\"): 0.45**0.5,\n",
    "        (\"in0@TE\", \"out0@TM\"): 0.01**0.5,\n",
    "        (\"in0@TE\", \"out1@TM\"): 1j*0.01**0.5,\n",
    "        (\"in1@TE\", \"out0@TM\"): 1j*0.01**0.5,\n",
    "        (\"in1@TE\", \"out1@TM\"): 0.01**0.5,\n",
    "        (\"in0@TM\", \"out0@TE\"): 0.01**0.5,\n",
    "        (\"in0@TM\", \"out1@TE\"): 1j*0.01**0.5,\n",
    "        (\"in1@TM\", \"out0@TE\"): 1j*0.01**0.5,\n",
    "        (\"in1@TM\", \"out1@TE\"): 0.01**0.5,\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "accessory-sheep",
   "metadata": {
    "papermill": {
     "duration": 0.003872,
     "end_time": "2023-09-03T04:40:09.456496",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.452624",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Multimode MZI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "joint-semiconductor",
   "metadata": {
    "papermill": {
     "duration": 0.0042,
     "end_time": "2023-09-03T04:40:09.484952",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.480752",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We can now combine these models into a circuit in much the same way as before. We just need to add the `modes=` keyword:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "suffering-judges",
   "metadata": {
    "papermill": {
     "duration": 1.152105,
     "end_time": "2023-09-03T04:40:10.641534",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.489429",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "mzi, _ = sax.circuit(\n",
    "    netlist={\n",
    "        \"instances\": {\n",
    "            \"lft\": \"coupler\",  # single mode models will be automatically converted to multimode models without cross polarization.\n",
    "            \"top\": {\"component\": \"straight\", \"settings\": {\"length\": 25.0}},\n",
    "            \"btm\": {\"component\": \"straight\", \"settings\": {\"length\": 15.0}},\n",
    "            \"rgt\": \"coupler\",  # single mode models will be automatically converted to multimode models without cross polarization.\n",
    "        },\n",
    "        \"connections\": {\n",
    "            \"lft,out0\": \"btm,in0\",\n",
    "            \"btm,out0\": \"rgt,in0\",\n",
    "            \"lft,out1\": \"top,in0\",\n",
    "            \"top,out0\": \"rgt,in1\",\n",
    "        },\n",
    "        \"ports\": {\n",
    "            \"in0\": \"lft,in0\",\n",
    "            \"in1\": \"lft,in1\",\n",
    "            \"out0\": \"rgt,out0\",\n",
    "            \"out1\": \"rgt,out1\",\n",
    "        },\n",
    "    },\n",
    "    models={\n",
    "        'coupler': coupler,\n",
    "        'straight': waveguide,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bridal-insider",
   "metadata": {
    "papermill": {
     "duration": 0.663387,
     "end_time": "2023-09-03T04:40:11.309065",
     "exception": false,
     "start_time": "2023-09-03T04:40:10.645678",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "mzi()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cellular-contemporary",
   "metadata": {
    "papermill": {
     "duration": 0.002715,
     "end_time": "2023-09-03T04:40:11.314625",
     "exception": false,
     "start_time": "2023-09-03T04:40:11.311910",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "we can convert this model back to a singlemode `SDict` as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aggregate-lemon",
   "metadata": {
    "papermill": {
     "duration": 0.139365,
     "end_time": "2023-09-03T04:40:11.456844",
     "exception": false,
     "start_time": "2023-09-03T04:40:11.317479",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "mzi_te = sax.singlemode(mzi, mode=\"TE\")\n",
    "mzi_te()"
   ]
  }
 ],
 "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": 5.464137,
   "end_time": "2023-09-03T04:40:11.878535",
   "environment_variables": {},
   "exception": null,
   "input_path": "./04_multimode_simulations.ipynb",
   "output_path": "./04_multimode_simulations.ipynb",
   "parameters": {},
   "start_time": "2023-09-03T04:40:06.414398",
   "version": "2.4.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}