{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "competitive-seating",
   "metadata": {},
   "source": [
    "# Circuit from YAML\n",
    "> Sometimes it's useful to be able to define circuits from YAML definitions. To not re-invent the wheel, SAX uses [GDSFactory](https://gdsfactory.readthedocs.io/en/latest/yaml.html)'s YAML netlist spec to define its circuits. This makes it very easy to convert a GDSFactory layout to a SAX circuit model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "personalized-recipe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax.numpy as jnp\n",
    "import matplotlib.pyplot as plt\n",
    "import sax\n",
    "import yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "featured-liberal",
   "metadata": {},
   "source": [
    "## MZI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "collectible-feedback",
   "metadata": {},
   "source": [
    "Let's first see how we can define a SAX circuit from YAML:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fixed-hurricane",
   "metadata": {},
   "outputs": [],
   "source": [
    "netlist = \"\"\"\n",
    "instances:\n",
    "    lft:\n",
    "      component: coupler\n",
    "      settings:\n",
    "          coupling: 0.5\n",
    "    rgt:\n",
    "        component: coupler\n",
    "        settings:\n",
    "            coupling: 0.5\n",
    "    top:\n",
    "        component: straight\n",
    "        settings:\n",
    "            length: 25.0\n",
    "    btm:\n",
    "        component: straight\n",
    "        settings:\n",
    "            length: 15.0\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": "f57633a5-10e9-4ddd-a10d-b29c47e26878",
   "metadata": {},
   "outputs": [],
   "source": [
    "yaml.safe_load(netlist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e91dc2b-bcdb-4a32-ba81-d0e99ed6a562",
   "metadata": {},
   "outputs": [],
   "source": [
    "mzi, _ = sax.circuit(\n",
    "    yaml.safe_load(netlist),\n",
    "    models={\"coupler\": sax.models.coupler, \"straight\": sax.models.straight},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dynamic-commonwealth",
   "metadata": {},
   "outputs": [],
   "source": [
    "wl = jnp.linspace(1.5, 1.6, 1000)\n",
    "transmission = jnp.abs(mzi(wl=wl)[\"in0\", \"out0\"]) ** 2\n",
    "\n",
    "plt.plot(wl * 1e3, transmission)\n",
    "plt.xlabel(\"λ [nm]\")\n",
    "plt.ylabel(\"T\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "alpha-married",
   "metadata": {},
   "source": [
    "That was easy! However, during the above YAML conversion, only models available in `sax.models` were used. What if we want to map the YAML component names to custom models? Let's say we want to use a dispersionless waveguide for the above model for example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "important-association",
   "metadata": {},
   "outputs": [],
   "source": [
    "def waveguide_without_dispersion(wl=1.55, length=25.0, neff=2.34):\n",
    "    phase = 2 * jnp.pi * neff * length / wl\n",
    "    sdict = sax.reciprocal({(\"in0\", \"out0\"): jnp.exp(1j * phase)})\n",
    "    return sdict"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "valuable-candidate",
   "metadata": {},
   "source": [
    "We can regenerate the above circuit again, but this time we specify a models mapping:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "former-retro",
   "metadata": {},
   "outputs": [],
   "source": [
    "mzi, _ = sax.circuit(yaml.safe_load(netlist), models={\"straight\": waveguide_without_dispersion, 'coupler': sax.models.coupler})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "focal-question",
   "metadata": {},
   "source": [
    "> The `models=` keyword in `circuit_from_yaml` can be a dictionary **or** an imported python module (like for example `sax.models`). Or a list containing multiple of such dictionary mappings and imported modules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "expired-panel",
   "metadata": {},
   "outputs": [],
   "source": [
    "wl = jnp.linspace(1.5, 1.6, 1000)\n",
    "transmission = jnp.abs(mzi(wl=wl)[\"in0\", \"out0\"]) ** 2\n",
    "\n",
    "plt.plot(wl, transmission)\n",
    "plt.xlabel(\"Wavelength [nm]\")\n",
    "plt.ylabel(\"T\")\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sax",
   "language": "python",
   "name": "sax"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}