bayer challenge

docs/_toc.yml

@@ -12,3 +12,4 @@ parts:
     - file: notebooks/library_challenge
     - file: notebooks/metalens_challenge
     - file: notebooks/photon_extractor_challenge
+    - file: notebooks/bayer_challenge

docs/notebooks/bayer_challenge.ipynb

@@ -0,0 +1,258 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Bayer sorter\n",
+    "\n",
+    "The bayer sorter challenge involves the design of a Si3N4 metasurface to split light in a wavelength-dependent way, so that red, green, and blue light is predominantly collected in red, green, and blue subpixels in a sensor array. The challenge is based on [Pixel-level Bayer-type colour router based on metasurfaces](https://www.nature.com/articles/s41467-022-31019-7) by Zou et al."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Simulating an existing design\n",
+    "\n",
+    "We'll begin by loading, visualizing, and simulating the design from Supplementary Figure 2 of Zou et al."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as onp\n",
+    "from skimage import measure\n",
+    "\n",
+    "design = onp.genfromtxt(\"../../reference_designs/bayer/zou.csv\", delimiter=\",\")\n",
+    "\n",
+    "plt.figure(figsize=(3, 3))\n",
+    "ax = plt.subplot(111)\n",
+    "im = plt.imshow(1 - design, cmap=\"gray\")\n",
+    "im.set_clim([-2, 1])\n",
+    "ax.set_xticks([])\n",
+    "ax.set_yticks([])\n",
+    "for c in measure.find_contours(design):\n",
+    "    plt.plot(c[:, 1], c[:, 0], 'k', lw=1)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Next, create the `bayer_sorter` challenge, which enables us to simulate and evaluate a bayer sorter design.\n",
+    "\n",
+    "The default simulation parameters are chosen to balance accuracy and simulation cost. For this notebook, we'll override these with settings that yield more accurate results."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import jax.numpy as jnp\n",
+    "from invrs_gym.challenges.bayer import challenge as bayer_challenge\n",
+    "\n",
+    "challenge = bayer_challenge.bayer_sorter(\n",
+    "    sim_params=dataclasses.replace(\n",
+    "        bayer_challenge.BAYER_SIM_PARAMS,\n",
+    "        approximate_num_terms=1200,\n",
+    "        wavelength=jnp.arange(0.405, 0.7, 0.01)\n",
+    "    )\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    " The `params` or optimization variables of the challenge include the metasurface pattern and also the metasurface thickness and metasurface-to-focal-plane separation. We'll obtain default initial parameters from the challenge, and then overwrite the metasurface pattern with the array loaded and plotted above."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import jax\n",
+    "\n",
+    "params = challenge.component.init(jax.random.PRNGKey(0))\n",
+    "assert params[\"density_metasurface\"].shape == design.shape\n",
+    "params[\"density_metasurface\"].array = design"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now, simulate the bayer sorter. Use a high-resolution wavelength grid so that its detailed wavelength-dependent response is characterized."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "response, aux = challenge.component.response(params)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The response contains the transmission for normally-incident plane wave at the specified wavelengths for both x-polarized and y-polarized fields. The transmission is reported for four sub-pixels; the first is for red, the second and third for green, and the final is for blue.\n",
+    "\n",
+    "Plot the transmission into red, green, and blue subpixels"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Average the transmission over the two different incident polarizations.\n",
+    "transmission = onp.mean(response.transmission, axis=-2)\n",
+    "\n",
+    "transmission_blue_pixel = transmission[:, 0]\n",
+    "transmission_green_pixel = transmission[:, 1] + transmission[:, 2]\n",
+    "transmission_red_pixel = transmission[:, 3]\n",
+    "\n",
+    "plt.plot(response.wavelength, transmission_blue_pixel, \"b\", lw=3)\n",
+    "plt.plot(response.wavelength, transmission_green_pixel, \"g\", lw=3)\n",
+    "plt.plot(response.wavelength, transmission_red_pixel, \"r\", lw=3)\n",
+    "plt.xlabel(\"Wavelength\")\n",
+    "plt.ylabel(\"Sub-pixel transmission\")\n",
+    "_ = plt.ylim(-0.05, 0.7)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The result is in very close agreement with Supplementary figure 7 of Zou et al.\n",
+    "\n",
+    "The transmission is computed by calculating the Poynting flux on the real-space grid at the focal plane, and summing within each sub-pixel quadrant. Since the fields are automatically computed during the course of a simulation, they are returned in `aux`.\n",
+    "\n",
+    "Let's plot the fields in the focal plane for each of the wavelengths."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": [
+     "hide-cell"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "import base64\n",
+    "import zlib\n",
+    "from matplotlib import colors\n",
+    "\n",
+    "# Code is from https://github.com/rsmith-nl/wavelength_to_rgb.\n",
+    "_CTBL = (\n",
+    "    b\"eNrV0ltr0AUAxuHn1rqoiAqSiA6EEJ3ogBUJeeFFRUdJOjOyWau1bLZpGztnM1oyy\"\n",
+    "    b\"2yYmcmMETPNlixNVpmssobKMJuYDVnNGksz0zTe5C9BXyF4v8Dz4y1RMlPpLGVlKs\"\n",
+    "    b\"pVVqh+Vu0cDVVa5mqdp61Ge63FdTrqLWuwolFno64m3U3WNutp1ttsY7O+Jpub9Df\"\n",
+    "    b\"a2migwY56O+sM1dpTY3iekblGq4zNcWC2QxWOlDtWJqVSIg/JfTJd7pRbZZpMlZtk\"\n",
+    "    b\"slwtl8skuUjOk3PkDDlFnNipcWZMjAtjUlwR18aNcXNMi9virpgRD0ZJlMZTMTuqo\"\n",
+    "    b\"ibqoyUWxCuxKJbEm/F2dEZXrI4PYn1siL7YHP3xTWyLwfg+9sRwjMT+GI/f4884Fj\"\n",
+    "    b\"mxP2S/7JVB+Vr6pEe65C1ZJC9KjTxduKcX1smF74TMhDgtzopz4/y4uGBdFlfFdXF\"\n",
+    "    b\"DTImpBe6WuD3ujnvj/ni4ID4WTxTKZ6IyqgtoXTTGC9EaL8fCgvt6dPwrXhmrCnR3\"\n",
+    "    b\"rIl18VH0xicF/fPYEl/G1hiI7UWA72KoaPBj7Iufigxj8VtR4nAcjeMnYxyXo3JYD\"\n",
+    "    b\"sq4/CqjMiLD8oPsll1Fp+0yUNTqly/kU9kkG2S9fChrpLtIuErekeWyVN6Q16Rd2m\"\n",
+    "    b\"SBzJcmqSvqVkulVMiTMkselUfkAZkh98gd/znZFLlerpEr5VK5RC6QiXK2nC4TTv7\"\n",
+    "    b\"sf7C/OcZfHOEwhzjIAcYZ4xdG+ZkR9jHMXvawmyF2sZNBdrCNAb5lK1/RzxY28xl9\"\n",
+    "    b\"bGIjH9PLenpYx1rep5v36OJdOlnJCpazjKV0sITFvEo7C2njJVqZTwtNNFBHLc9Tz\"\n",
+    "    b\"XNUMpsKyinjcUqZSQn/AJ7p9HY=\"\n",
+    ")\n",
+    "\n",
+    "\n",
+    "_CTBL = zlib.decompress(base64.b64decode(_CTBL))\n",
+    "\n",
+    "\n",
+    "def rgb(wavelength_nm):\n",
+    "    \"\"\"Converts a wavelength between 380 and 780 nm to an RGB color tuple.\n",
+    "\n",
+    "    Args:\n",
+    "        wavelength_nm: Wavelength in nanometers. It is rounded to the nearest integer.\n",
+    "\n",
+    "    Returns:\n",
+    "        A 3-tuple (red, green, blue) of integers in the range 0-255.\n",
+    "    \"\"\"\n",
+    "    wavelength_nm = int(round(wavelength_nm))\n",
+    "    if wavelength_nm < 380 or wavelength_nm > 780:\n",
+    "        raise ValueError(\"wavelength out of range\")\n",
+    "    idx = (wavelength_nm - 380) * 3\n",
+    "    color_str = _CTBL[idx : idx + 3]\n",
+    "    return onp.asarray([int(i) for i in color_str]) / 255\n",
+    "\n",
+    "\n",
+    "def cmap_for_wavelength(wavelength_um, background_color = \"k\"):\n",
+    "    \"\"\"Generates a\"\"\"\n",
+    "    color = rgb(wavelength_nm=wavelength_um * 1000)\n",
+    "    return colors.LinearSegmentedColormap.from_list(\n",
+    "        \"b\", [background_color, color], N=256\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "x, y = aux[\"coordinates_xy\"]\n",
+    "x = jnp.squeeze(x, axis=0)\n",
+    "y = jnp.squeeze(y, axis=0)\n",
+    "ex, ey, ez = aux[\"efield_xy\"]\n",
+    "intensity = jnp.abs(ex)**2 + jnp.abs(ey)**2 + jnp.abs(ez)**2\n",
+    "intensity = jnp.mean(intensity, axis=-1)  # Average over polarizations\n",
+    "\n",
+    "fig, axs = plt.subplots(ncols=6, nrows=5, figsize=(9, 9))\n",
+    "axs = axs.flatten()\n",
+    "for i, wavelength in enumerate(response.wavelength):\n",
+    "    cmap = cmap_for_wavelength(wavelength)\n",
+    "    axs[i].pcolormesh(x, y, intensity[i, :, :], cmap=cmap)\n",
+    "    axs[i].set_ylim(axs[i].get_ylim()[::-1])\n",
+    "    axs[i].axis(\"equal\")\n",
+    "    axs[i].axis(False)\n",
+    "    axs[i].plot([jnp.amin(x), jnp.amax(x)], [jnp.mean(y), jnp.mean(y)], \"w--\", lw=1)\n",
+    "    axs[i].plot([jnp.mean(x), jnp.mean(x)], [jnp.amin(y), jnp.amax(y)], \"w--\", lw=1)\n",
+    "    axs[i].set_title(f\"$\\lambda$={wavelength:.3f}$\\mu$m\", fontsize=10)\n",
+    "\n",
+    "plt.subplots_adjust(wspace=0.05, hspace=0.25)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "invrs",
+   "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"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

reference_designs/bayer/generate_design.ipynb

@@ -0,0 +1,69 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This file generates the color router from Supplementary figure 2 of [Pixel-level Bayer-type colour router based on metasurfaces](https://www.nature.com/articles/s41467-022-31019-7) by Zou et al."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as onp\n",
+    "from invrs_gym.utils import transforms\n",
+    "\n",
+    "pattern = \"\"\"0100100000000000\n",
+    "    0110001010000000\n",
+    "    0000100010000000\n",
+    "    0000000000000000\n",
+    "    1010000000000000\n",
+    "    1000100000000000\n",
+    "    0000011000000000\n",
+    "    1001000000000110\n",
+    "    0011110000000000\n",
+    "    1000000001000010\n",
+    "    0011000101000000\n",
+    "    0110100100100101\n",
+    "    0010010110000000\n",
+    "    0101110100010010\n",
+    "    0010100000000011\n",
+    "    0000001010110000\"\"\"\n",
+    "\n",
+    "design = [[float(i) for i in row.strip()] for row in pattern.split(\"\\n\")]\n",
+    "design = onp.asarray(design)[:, ::-1]  # Flip orientation compared to figure.\n",
+    "design = onp.kron(design, onp.ones((25, 25))).astype(int)\n",
+    "design = transforms.resample(design, (200, 200))\n",
+    "\n",
+    "plt.imshow(design)\n",
+    "\n",
+    "onp.savetxt(\"zou.csv\", design, delimiter=\",\", fmt=\"%i\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "invrs-cpu",
+   "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"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}