examples/04_multimode_simulations.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "textile-panel",
   "metadata": {},
   "source": [
    "# Multimode simulations\n",
    "> SAX can handle multiple modes too!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "prescribed-plant",
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import combinations_with_replacement, product\n",
    "\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import sax"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "requested-calvin",
   "metadata": {},
   "source": [
    "## Ports and modes per port"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "grateful-programming",
   "metadata": {},
   "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": {},
   "outputs": [],
   "source": [
    "\n",
    "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": {},
   "source": [
    "If we would disregard any backreflection, this can be further simplified:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "numerical-novel",
   "metadata": {},
   "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": {},
   "source": [
    "Sometimes cross-polarization terms can also be ignored:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "encouraging-territory",
   "metadata": {},
   "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": {},
   "source": [
    "## Multimode waveguide"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "asian-noise",
   "metadata": {},
   "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": {},
   "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": "accessory-sheep",
   "metadata": {},
   "source": [
    "## Multimode MZI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "joint-semiconductor",
   "metadata": {},
   "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": {},
   "outputs": [],
   "source": [
    "mzi = sax.circuit(\n",
    "    instances={\n",
    "        \"lft\": sax.models.coupler,  # single mode models will be automatically converted to multimode models without cross polarization.\n",
    "        \"top\": sax.partial(waveguide, length=25.0),\n",
    "        \"btm\": sax.partial(waveguide, length=15.0),\n",
    "        \"rgt\": sax.models.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",
    "    modes=(\"te\", \"tm\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bridal-insider",
   "metadata": {},
   "outputs": [],
   "source": [
    "mzi()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cellular-contemporary",
   "metadata": {},
   "source": [
    "we can convert this model back to a singlemode `SDict` as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aggregate-lemon",
   "metadata": {},
   "outputs": [],
   "source": [
    "mzi_te = sax.singlemode(mzi, mode=\"te\")\n",
    "mzi_te()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sax",
   "language": "python",
   "name": "sax"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}