{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "competitive-seating",
   "metadata": {
    "papermill": {
     "duration": 0.002829,
     "end_time": "2023-09-03T04:40:07.682429",
     "exception": false,
     "start_time": "2023-09-03T04:40:07.679600",
     "status": "completed"
    },
    "tags": []
   },
   "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": {
    "papermill": {
     "duration": 1.967603,
     "end_time": "2023-09-03T04:40:09.652605",
     "exception": false,
     "start_time": "2023-09-03T04:40:07.685002",
     "status": "completed"
    },
    "tags": []
   },
   "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": {
    "papermill": {
     "duration": 0.002693,
     "end_time": "2023-09-03T04:40:09.658132",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.655439",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## MZI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "collectible-feedback",
   "metadata": {
    "papermill": {
     "duration": 0.002765,
     "end_time": "2023-09-03T04:40:09.663551",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.660786",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Let's first see how we can define a SAX circuit from YAML:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fixed-hurricane",
   "metadata": {
    "papermill": {
     "duration": 0.019541,
     "end_time": "2023-09-03T04:40:09.685725",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.666184",
     "status": "completed"
    },
    "tags": []
   },
   "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": {
    "papermill": {
     "duration": 0.034341,
     "end_time": "2023-09-03T04:40:09.729842",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.695501",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "yaml.safe_load(netlist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e91dc2b-bcdb-4a32-ba81-d0e99ed6a562",
   "metadata": {
    "papermill": {
     "duration": 1.003594,
     "end_time": "2023-09-03T04:40:10.735629",
     "exception": false,
     "start_time": "2023-09-03T04:40:09.732035",
     "status": "completed"
    },
    "tags": []
   },
   "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": {
    "papermill": {
     "duration": 1.442552,
     "end_time": "2023-09-03T04:40:12.181254",
     "exception": false,
     "start_time": "2023-09-03T04:40:10.738702",
     "status": "completed"
    },
    "tags": []
   },
   "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": {
    "papermill": {
     "duration": 0.004263,
     "end_time": "2023-09-03T04:40:12.219071",
     "exception": false,
     "start_time": "2023-09-03T04:40:12.214808",
     "status": "completed"
    },
    "tags": []
   },
   "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": {
    "papermill": {
     "duration": 0.010415,
     "end_time": "2023-09-03T04:40:12.232816",
     "exception": false,
     "start_time": "2023-09-03T04:40:12.222401",
     "status": "completed"
    },
    "tags": []
   },
   "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": {
    "papermill": {
     "duration": 0.003564,
     "end_time": "2023-09-03T04:40:12.240320",
     "exception": false,
     "start_time": "2023-09-03T04:40:12.236756",
     "status": "completed"
    },
    "tags": []
   },
   "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": {
    "papermill": {
     "duration": 0.038857,
     "end_time": "2023-09-03T04:40:12.282793",
     "exception": false,
     "start_time": "2023-09-03T04:40:12.243936",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "mzi, _ = sax.circuit(\n",
    "    yaml.safe_load(netlist),\n",
    "    models={\"straight\": waveguide_without_dispersion, \"coupler\": sax.models.coupler},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "focal-question",
   "metadata": {
    "papermill": {
     "duration": 0.002445,
     "end_time": "2023-09-03T04:40:12.287492",
     "exception": false,
     "start_time": "2023-09-03T04:40:12.285047",
     "status": "completed"
    },
    "tags": []
   },
   "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": {
    "papermill": {
     "duration": 0.258497,
     "end_time": "2023-09-03T04:40:12.548460",
     "exception": false,
     "start_time": "2023-09-03T04:40:12.289963",
     "status": "completed"
    },
    "tags": []
   },
   "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": "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.10.12"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 6.554796,
   "end_time": "2023-09-03T04:40:12.970747",
   "environment_variables": {},
   "exception": null,
   "input_path": "./03_circuit_from_yaml.ipynb",
   "output_path": "./03_circuit_from_yaml.ipynb",
   "parameters": {},
   "start_time": "2023-09-03T04:40:06.415951",
   "version": "2.4.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}