examples/01_quick_start.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ab894499",
   "metadata": {},
   "source": [
    "# SAX Quick Start\n",
    "> Let's go over the core functionality of SAX."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c91396d",
   "metadata": {},
   "source": [
    "## Environment variables"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79240f32",
   "metadata": {},
   "source": [
    "SAX is based on JAX... here are some useful environment variables for working with JAX:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40ce30d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# select float32 or float64 as default dtype\n",
    "%env JAX_ENABLE_X64=0\n",
    "\n",
    "# select cpu or gpu\n",
    "%env JAX_PLATFORM_NAME=cpu\n",
    "\n",
    "# set custom CUDA location for gpu:\n",
    "%env XLA_FLAGS=--xla_gpu_cuda_data_dir=/usr/lib/cuda\n",
    "\n",
    "# Using GPU?\n",
    "from jax.lib import xla_bridge\n",
    "print(xla_bridge.get_backend().platform)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4938b729-8b27-406f-b5fd-e118a9fb39f1",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0450fde",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.example_libraries.optimizers as opt\n",
    "import jax.numpy as jnp\n",
    "import matplotlib.pyplot as plt  # plotting\n",
    "import sax\n",
    "import tqdm  # progress bars"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c4d4c4e",
   "metadata": {},
   "source": [
    "## Scatter *dictionaries*\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, dictonaries allow for string indexing, which makes them much more pleasant to use in this context. Let’s for example create an `SDict` for a 50/50 coupler:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7f90205",
   "metadata": {},
   "source": [
    "```\n",
    "in1          out1\n",
    "   \\        /\n",
    "    ========\n",
    "   /        \\\n",
    "in0          out0\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89f215e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "coupling = 0.5\n",
    "kappa = coupling ** 0.5\n",
    "tau = (1 - coupling) ** 0.5\n",
    "coupler_dict = {\n",
    "    (\"in0\", \"out0\"): tau,\n",
    "    (\"out0\", \"in0\"): tau,\n",
    "    (\"in0\", \"out1\"): 1j * kappa,\n",
    "    (\"out1\", \"in0\"): 1j * kappa,\n",
    "    (\"in1\", \"out0\"): 1j * kappa,\n",
    "    (\"out0\", \"in1\"): 1j * kappa,\n",
    "    (\"in1\", \"out1\"): tau,\n",
    "    (\"out1\", \"in1\"): tau,\n",
    "}\n",
    "coupler_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb9cb2b6",
   "metadata": {},
   "source": [
    "Only the non-zero port combinations need to be specified. Any non-existent port-combination (for example `(\"in0\", \"in1\")`) is considered to be zero by SAX.\n",
    "\n",
    "Obviously, 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": "494f445d",
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "\n",
    "coupler_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67e730f6",
   "metadata": {},
   "source": [
    "## Parametrized Models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec1e1214",
   "metadata": {},
   "source": [
    "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`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33efb99f",
   "metadata": {},
   "outputs": [],
   "source": [
    "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\n",
    "\n",
    "\n",
    "coupler(coupling=0.3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ef93f65",
   "metadata": {},
   "source": [
    "We can define a waveguide in the same way:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5eac9176",
   "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",
    "    sdict = sax.reciprocal(\n",
    "        {\n",
    "            (\"in0\", \"out0\"): transmission,\n",
    "        }\n",
    "    )\n",
    "    return sdict"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67713c18",
   "metadata": {},
   "source": [
    "That's pretty straightforward. Let's now move on to parametrized circuits:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "579a9ce1",
   "metadata": {},
   "source": [
    "## Circuit Models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc3c78a0",
   "metadata": {},
   "source": [
    "Existing models can now be combined into a circuit using `sax.circuit`, which basically creates a new `Model` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6885e1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "mzi = sax.circuit(\n",
    "    instances={\n",
    "        \"lft\": coupler,\n",
    "        \"top\": waveguide,\n",
    "        \"btm\": waveguide,\n",
    "        \"rgt\": coupler,\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",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e95325c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "mzi?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2c74cff",
   "metadata": {},
   "source": [
    "The `circuit` function just creates a similar function as we created for the waveguide and the coupler, but in stead of taking parameters directly it takes parameter *dictionaries* for each of the instances in the circuit. The keys in these parameter dictionaries should correspond to the keyword arguments of each individual subcomponent. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba1202d1",
   "metadata": {},
   "source": [
    "Let's now do a simulation for the MZI we just constructed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2108f6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%time mzi()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb43b835",
   "metadata": {},
   "outputs": [],
   "source": [
    "mzi2 = jax.jit(mzi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61f2bb6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%time mzi2()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e0d41ea",
   "metadata": {},
   "source": [
    "Or in the case we want an MZI with different arm lengths:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f53bdc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "mzi(top={\"length\": 25.0}, btm={\"length\": 15.0})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a331252a",
   "metadata": {},
   "source": [
    "## Simulating the parametrized MZI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ba8626b",
   "metadata": {},
   "source": [
    "We can simulate the above mzi for multiple wavelengths as well by specifying the wavelength at the top level of the circuit call. Each setting specified at the top level of the circuit call will be propagated to all subcomponents of the circuit which have that setting:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6c8b5c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "wl = jnp.linspace(1.51, 1.59, 1000)\n",
    "%time S = mzi(wl=wl, top={\"length\": 25.0}, btm={\"length\": 15.0})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28d8cc75",
   "metadata": {},
   "source": [
    "Let's see what this gives:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7de73dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(wl * 1e3, abs(S[\"in0\", \"out0\"]) ** 2)\n",
    "plt.ylim(-0.05, 1.05)\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.ylabel(\"T\")\n",
    "plt.ylim(-0.05, 1.05)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb9f0074",
   "metadata": {},
   "source": [
    "## Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "431aee51",
   "metadata": {},
   "source": [
    "We'd like to optimize an MZI such that one of the minima is at 1550nm. To do this, we need to define a loss function for the circuit at 1550nm. This function should take the parameters that you want to optimize as positional arguments:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7dabe735",
   "metadata": {},
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def loss(delta_length):\n",
    "    S = mzi(wl=1.55, top={\"length\": 15.0 + delta_length}, btm={\"length\": 15.0})\n",
    "    return (abs(S[\"in0\", \"out0\"]) ** 2).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a60a8eab",
   "metadata": {},
   "outputs": [],
   "source": [
    "%time loss(10.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af1eb06b",
   "metadata": {},
   "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": "16369cba",
   "metadata": {},
   "outputs": [],
   "source": [
    "grad = jax.jit(\n",
    "    jax.grad(\n",
    "        loss,\n",
    "        argnums=0,  # JAX gradient function for the first positional argument, jitted\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cd2f2cf",
   "metadata": {},
   "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": "e553d69e",
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_delta_length = 10.0\n",
    "optim_init, optim_update, optim_params = opt.adam(step_size=0.1)\n",
    "optim_state = optim_init(initial_delta_length)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0a03fc2",
   "metadata": {},
   "source": [
    "Given all this, a single training step can be defined:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eea4f7b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_step(step, optim_state):\n",
    "    settings = optim_params(optim_state)\n",
    "    lossvalue = loss(settings)\n",
    "    gradvalue = grad(settings)\n",
    "    optim_state = optim_update(step, gradvalue, optim_state)\n",
    "    return lossvalue, optim_state"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba1612df",
   "metadata": {},
   "source": [
    "And we can use this step function to start the training of the MZI:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61ef4479",
   "metadata": {},
   "outputs": [],
   "source": [
    "range_ = tqdm.trange(300)\n",
    "for step in range_:\n",
    "    lossvalue, optim_state = train_step(step, optim_state)\n",
    "    range_.set_postfix(loss=f\"{lossvalue:.6f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb9a1e20",
   "metadata": {},
   "outputs": [],
   "source": [
    "delta_length = optim_params(optim_state)\n",
    "delta_length"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15d8754b",
   "metadata": {},
   "source": [
    "Let's see what we've got over a range of wavelengths:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b36d10a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "S = mzi(wl=wl, top={\"length\": 15.0 + delta_length}, btm={\"length\": 15.0})\n",
    "plt.plot(wl * 1e3, abs(S[\"in1\", \"out1\"]) ** 2)\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.ylabel(\"T\")\n",
    "plt.ylim(-0.05, 1.05)\n",
    "plt.plot([1550, 1550], [0, 1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a990da96",
   "metadata": {},
   "source": [
    "The minimum of the MZI is perfectly located at 1550nm."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d93aa85",
   "metadata": {},
   "source": [
    "## MZI Chain"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c312e13",
   "metadata": {},
   "source": [
    "Let's now create a chain of MZIs. For this, we first create a subcomponent: a directional coupler with arms:\n",
    "\n",
    "\n",
    "```\n",
    "                                 top\n",
    "                            in0 ----- out0 -> out1\n",
    "    in1 <- in1        out1                 \n",
    "              \\  dc  /                  \n",
    "               ======                  \n",
    "              /      \\                \n",
    "    in0 <- in0        out0       btm    \n",
    "                            in0 ----- out0 -> out0\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ebe80c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "dc_with_arms = sax.circuit(\n",
    "    instances={\n",
    "        \"lft\": coupler,\n",
    "        \"top\": waveguide,\n",
    "        \"btm\": waveguide,\n",
    "    },\n",
    "    connections={\n",
    "        \"lft,out0\": \"btm,in0\",\n",
    "        \"lft,out1\": \"top,in0\",\n",
    "    },\n",
    "    ports={\n",
    "        \"in0\": \"lft,in0\",\n",
    "        \"in1\": \"lft,in1\",\n",
    "        \"out0\": \"btm,out0\",\n",
    "        \"out1\": \"top,out0\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1eda468f",
   "metadata": {},
   "source": [
    "An MZI chain can now be created by cascading these directional couplers with arms:\n",
    "\n",
    "```\n",
    "      _    _    _    _             _    _  \n",
    "    \\/   \\/   \\/   \\/     ...    \\/   \\/   \n",
    "    /\\_  /\\_  /\\_  /\\_           /\\_  /\\_  \n",
    "```\n",
    "\n",
    "Let's create a *model factory* (`ModelFactory`) for this. In SAX, a *model factory* is any keyword-only function that generates a `Model`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11c4ab07",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mzi_chain(num_mzis=1) -> sax.Model:\n",
    "    chain = sax.circuit(\n",
    "        instances={f\"dc{i}\": dc_with_arms for i in range(num_mzis + 1)},\n",
    "        connections={\n",
    "            **{f\"dc{i},out0\": f\"dc{i+1},in0\" for i in range(num_mzis)},\n",
    "            **{f\"dc{i},out1\": f\"dc{i+1},in1\" for i in range(num_mzis)},\n",
    "        },\n",
    "        ports={\n",
    "            \"in0\": f\"dc0,in0\",\n",
    "            \"in1\": f\"dc0,in1\",\n",
    "            \"out0\": f\"dc{num_mzis},out0\",\n",
    "            \"out1\": f\"dc{num_mzis},out1\",\n",
    "        },\n",
    "    )\n",
    "    return chain"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03f225e9",
   "metadata": {},
   "source": [
    "Let's for example create a chain with 15 MZIs. We can also update the settings dictionary as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "576d4de8",
   "metadata": {},
   "outputs": [],
   "source": [
    "chain = mzi_chain(num_mzis=15)\n",
    "settings = sax.get_settings(chain)\n",
    "for dc in settings:\n",
    "    settings[dc][\"top\"][\"length\"] = 25.0\n",
    "    settings[dc][\"btm\"][\"length\"] = 15.0\n",
    "settings = sax.update_settings(settings, wl=jnp.linspace(1.5, 1.6, 1000))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11e3a10f",
   "metadata": {},
   "source": [
    "We can simulate this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6e1155a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%time S = chain(**settings) # time to evaluate the MZI\n",
    "func = jax.jit(chain)\n",
    "%time S = func(**settings) # time to jit the MZI\n",
    "%time S = func(**settings) # time to evaluate the MZI after jitting"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70db7179",
   "metadata": {},
   "source": [
    "Where we see that the unjitted evaluation of the MZI chain takes about a second, while the jitting of the MZI chain takes about two minutes (on a CPU). However, after the MZI chain has been jitted, the evaluation is in the order of about a millisecond!\n",
    "\n",
    "Anyway, let's see what this gives:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e4a2aab",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(1e3 * settings[\"dc0\"][\"top\"][\"wl\"], jnp.abs(S[\"in0\", \"out0\"]) ** 2)\n",
    "plt.ylim(-0.05, 1.05)\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.ylabel(\"T\")\n",
    "plt.ylim(-0.05, 1.05)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sax",
   "language": "python",
   "name": "sax"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}