{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {
    "papermill": {
     "duration": 0.030769,
     "end_time": "2024-06-17T18:16:22.569120",
     "exception": false,
     "start_time": "2024-06-17T18:16:22.538351",
     "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": "1",
   "metadata": {
    "papermill": {
     "duration": 2.166678,
     "end_time": "2024-06-17T18:16:24.746814",
     "exception": false,
     "start_time": "2024-06-17T18:16:22.580136",
     "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": "2",
   "metadata": {
    "papermill": {
     "duration": 0.006354,
     "end_time": "2024-06-17T18:16:24.759099",
     "exception": false,
     "start_time": "2024-06-17T18:16:24.752745",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## MZI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {
    "papermill": {
     "duration": 0.004957,
     "end_time": "2024-06-17T18:16:24.769033",
     "exception": false,
     "start_time": "2024-06-17T18:16:24.764076",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Let's first see how we can define a SAX circuit from YAML:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {
    "papermill": {
     "duration": 0.018457,
     "end_time": "2024-06-17T18:16:24.792543",
     "exception": false,
     "start_time": "2024-06-17T18:16:24.774086",
     "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": "5",
   "metadata": {
    "papermill": {
     "duration": 0.030699,
     "end_time": "2024-06-17T18:16:24.828197",
     "exception": false,
     "start_time": "2024-06-17T18:16:24.797498",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "yaml.safe_load(netlist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {
    "papermill": {
     "duration": 1.266813,
     "end_time": "2024-06-17T18:16:26.100210",
     "exception": false,
     "start_time": "2024-06-17T18:16:24.833397",
     "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": "7",
   "metadata": {
    "papermill": {
     "duration": 1.678883,
     "end_time": "2024-06-17T18:16:27.785324",
     "exception": false,
     "start_time": "2024-06-17T18:16:26.106441",
     "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": "8",
   "metadata": {
    "papermill": {
     "duration": 0.007452,
     "end_time": "2024-06-17T18:16:27.800305",
     "exception": false,
     "start_time": "2024-06-17T18:16:27.792853",
     "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": "9",
   "metadata": {
    "papermill": {
     "duration": 0.057597,
     "end_time": "2024-06-17T18:16:27.864448",
     "exception": false,
     "start_time": "2024-06-17T18:16:27.806851",
     "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": "10",
   "metadata": {
    "papermill": {
     "duration": 0.004915,
     "end_time": "2024-06-17T18:16:27.874435",
     "exception": false,
     "start_time": "2024-06-17T18:16:27.869520",
     "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": "11",
   "metadata": {
    "papermill": {
     "duration": 0.077847,
     "end_time": "2024-06-17T18:16:27.956972",
     "exception": false,
     "start_time": "2024-06-17T18:16:27.879125",
     "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": "12",
   "metadata": {
    "papermill": {
     "duration": 0.007553,
     "end_time": "2024-06-17T18:16:27.972385",
     "exception": false,
     "start_time": "2024-06-17T18:16:27.964832",
     "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": "13",
   "metadata": {
    "papermill": {
     "duration": 0.300855,
     "end_time": "2024-06-17T18:16:28.280833",
     "exception": false,
     "start_time": "2024-06-17T18:16:27.979978",
     "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}