Merge pull request #125 from invrs-io/workflow

Reduce memory requirement for bayer notebook

docs/notebooks/bayer_challenge.ipynb

@@ -46,7 +46,7 @@
    "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."
+    "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: more terms in the Fourier basis, and more wavelengths."
    ]
   },
   {
@@ -67,7 +67,7 @@
     "    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",
+    "        wavelength=jnp.arange(0.405, 0.7, 0.02)\n",
     "    )\n",
     ")"
    ]
@@ -96,7 +96,7 @@
    "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."
+    "Now, simulate the bayer sorter."
    ]
   },
   {
@@ -130,9 +130,9 @@
     "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.plot(response.wavelength, transmission_blue_pixel, \"bo-\", lw=3)\n",
+    "plt.plot(response.wavelength, transmission_green_pixel, \"go-\", lw=3)\n",
+    "plt.plot(response.wavelength, transmission_red_pixel, \"ro-\", lw=3)\n",
     "plt.xlabel(\"Wavelength\")\n",
     "plt.ylabel(\"Sub-pixel transmission\")\n",
     "_ = plt.ylim(-0.05, 0.7)"
@@ -223,7 +223,7 @@
     "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",
+    "fig, axs = plt.subplots(ncols=5, nrows=3, figsize=(9, 7))\n",
     "axs = axs.flatten()\n",
     "for i, wavelength in enumerate(response.wavelength):\n",
     "    cmap = cmap_for_wavelength(wavelength)\n",