diff --git a/docs/notebooks/crystal_bz.ipynb b/docs/notebooks/crystal_bz.ipynb
index 7955115..8ed5e63 100644
--- a/docs/notebooks/crystal_bz.ipynb
+++ b/docs/notebooks/crystal_bz.ipynb
@@ -624,30 +624,30 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# eigensolve = functools.partial(\n",
-    "#     fmm.eigensolve_isotropic_media,\n",
-    "#     wavelength=jnp.asarray(wavelengths),\n",
-    "#     in_plane_wavevector=in_plane_wavevector,\n",
-    "#     primitive_lattice_vectors=primitive_lattice_vectors,\n",
-    "#     expansion=expansion,\n",
-    "#     formulation=fmm.Formulation.FFT,\n",
-    "# )\n",
+    "eigensolve = functools.partial(\n",
+    "    fmm.eigensolve_isotropic_media,\n",
+    "    wavelength=jnp.asarray(wavelengths),\n",
+    "    in_plane_wavevector=in_plane_wavevector,\n",
+    "    primitive_lattice_vectors=primitive_lattice_vectors,\n",
+    "    expansion=expansion,\n",
+    "    formulation=fmm.Formulation.FFT,\n",
+    ")\n",
     "\n",
-    "# mask = unit_cell_pattern(pitch, diameter, resolution)\n",
-    "# permittivity_crystal = jnp.where(mask, permittivity_ambient, permittivity_slab)\n",
-    "# solve_result_crystal = eigensolve(permittivity=permittivity_crystal)\n",
-    "# solve_result_ambient = eigensolve(\n",
-    "#     permittivity=jnp.asarray(permittivity_ambient)[jnp.newaxis, jnp.newaxis]\n",
-    "# )\n",
+    "mask = unit_cell_pattern(pitch, diameter, resolution)\n",
+    "permittivity_crystal = jnp.where(mask, permittivity_ambient, permittivity_slab)\n",
+    "solve_result_crystal = eigensolve(permittivity=permittivity_crystal)\n",
+    "solve_result_ambient = eigensolve(\n",
+    "    permittivity=jnp.asarray(permittivity_ambient)[jnp.newaxis, jnp.newaxis]\n",
+    ")\n",
     "\n",
-    "# s_matrices_interior = scattering.stack_s_matrices_interior(\n",
-    "#     layer_solve_results=[\n",
-    "#         solve_result_ambient,\n",
-    "#         solve_result_crystal,\n",
-    "#         solve_result_ambient,\n",
-    "#     ],\n",
-    "#     layer_thicknesses=[thickness_ambient, thickness_slab, thickness_ambient],\n",
-    "# )"
+    "s_matrices_interior = scattering.stack_s_matrices_interior(\n",
+    "    layer_solve_results=[\n",
+    "        solve_result_ambient,\n",
+    "        solve_result_crystal,\n",
+    "        solve_result_ambient,\n",
+    "    ],\n",
+    "    layer_thicknesses=[thickness_ambient, thickness_slab, thickness_ambient],\n",
+    ")"
    ]
   },
   {
@@ -663,36 +663,36 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# def _paraxial_gaussian_field_fn(x, y, z):\n",
-    "#     # Returns the fields of a z-propagating, x-polarized Gaussian beam.\n",
-    "#     # See https://en.wikipedia.org/wiki/Gaussian_beam\n",
-    "\n",
-    "#     # Adjust array dimensions for proper batching\n",
-    "#     wavelengths_padded = wavelengths[..., jnp.newaxis, jnp.newaxis]\n",
-    "\n",
-    "#     k = 2 * jnp.pi / wavelengths_padded\n",
-    "#     z_r = jnp.pi * beam_waist**2 * jnp.sqrt(permittivity_ambient) / wavelengths_padded\n",
-    "#     w_z = beam_waist * jnp.sqrt(1 + (z / z_r) ** 2)\n",
-    "#     r = jnp.sqrt(x**2 + y**2)\n",
-    "#     ex = (\n",
-    "#         beam_waist\n",
-    "#         / w_z\n",
-    "#         * jnp.exp(-(r**2) / w_z**2)\n",
-    "#         * jnp.exp(\n",
-    "#             1j\n",
-    "#             * (\n",
-    "#                 (k * z)  # Phase\n",
-    "#                 + k * r**2 / 2 * z / (z**2 + z_r**2)  # Wavefront curvature\n",
-    "#                 - jnp.arctan(z / z_r)  # Gouy phase\n",
-    "#             )\n",
-    "#         )\n",
-    "#     )\n",
-    "#     ey = jnp.zeros_like(ex)\n",
-    "#     ez = jnp.zeros_like(ex)\n",
-    "#     hx = jnp.zeros_like(ex)\n",
-    "#     hy = ex / jnp.sqrt(permittivity_ambient)\n",
-    "#     hz = jnp.zeros_like(ex)\n",
-    "#     return (ex, ey, ez), (hx, hy, hz)"
+    "def _paraxial_gaussian_field_fn(x, y, z):\n",
+    "    # Returns the fields of a z-propagating, x-polarized Gaussian beam.\n",
+    "    # See https://en.wikipedia.org/wiki/Gaussian_beam\n",
+    "\n",
+    "    # Adjust array dimensions for proper batching\n",
+    "    wavelengths_padded = wavelengths[..., jnp.newaxis, jnp.newaxis]\n",
+    "\n",
+    "    k = 2 * jnp.pi / wavelengths_padded\n",
+    "    z_r = jnp.pi * beam_waist**2 * jnp.sqrt(permittivity_ambient) / wavelengths_padded\n",
+    "    w_z = beam_waist * jnp.sqrt(1 + (z / z_r) ** 2)\n",
+    "    r = jnp.sqrt(x**2 + y**2)\n",
+    "    ex = (\n",
+    "        beam_waist\n",
+    "        / w_z\n",
+    "        * jnp.exp(-(r**2) / w_z**2)\n",
+    "        * jnp.exp(\n",
+    "            1j\n",
+    "            * (\n",
+    "                (k * z)  # Phase\n",
+    "                + k * r**2 / 2 * z / (z**2 + z_r**2)  # Wavefront curvature\n",
+    "                - jnp.arctan(z / z_r)  # Gouy phase\n",
+    "            )\n",
+    "        )\n",
+    "    )\n",
+    "    ey = jnp.zeros_like(ex)\n",
+    "    ez = jnp.zeros_like(ex)\n",
+    "    hx = jnp.zeros_like(ex)\n",
+    "    hy = ex / jnp.sqrt(permittivity_ambient)\n",
+    "    hz = jnp.zeros_like(ex)\n",
+    "    return (ex, ey, ez), (hx, hy, hz)"
    ]
   },
   {