Merge pull request #18 from invrs-io/broadband

Add broadband metagrating challenge

notebooks/diffract.ipynb

@@ -51,6 +51,9 @@
     "            print(f\"step {i} ({time.time() - t0:.2f}s): loss={value}\")\n",
     "            loss_values.append(value)\n",
     "            metrics_values.append(metrics)\n",
+    "        except KeyboardInterrupt:\n",
+    "            print(\"Terminating optimization\")\n",
+    "            break\n",
     "        except Exception as e:\n",
     "            if \"KeyboardInterrupt\" in str(e):\n",
     "                print(\"Terminating optimization\")\n",
@@ -78,7 +81,15 @@
     "    metagrating_challenge.metagrating(),\n",
     "    opt=optax.adam(0.02),\n",
     "    num_steps=200,\n",
-    ")"
+    ")\n",
+    "\n",
+    "plt.figure(figsize=(8, 5))\n",
+    "plt.subplot(121)\n",
+    "plt.plot(metagrating_loss_values)\n",
+    "plt.plot([m[\"distance_to_window\"] for m in metagrating_metrics_values])\n",
+    "plt.subplot(122)\n",
+    "plt.imshow(jnp.tile(metagrating_params.array, (1, 1)))\n",
+    "plt.colorbar()"
    ]
   },
   {
@@ -88,12 +99,26 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "reload(metagrating_challenge)\n",
+    "\n",
+    "(\n",
+    "    broadband_metagrating_params,\n",
+    "    broadband_metagrating_response,\n",
+    "    broadband_metagrating_aux,\n",
+    "    broadband_metagrating_loss_values,\n",
+    "    broadband_metagrating_metrics_values\n",
+    ") = optimize(\n",
+    "    metagrating_challenge.broadband_metagrating(),\n",
+    "    opt=optax.adam(0.02),\n",
+    "    num_steps=200,\n",
+    ")\n",
+    "\n",
     "plt.figure(figsize=(8, 5))\n",
     "plt.subplot(121)\n",
-    "plt.plot(metagrating_loss_values)\n",
-    "plt.plot([m[\"distance_to_window\"] for m in metagrating_metrics_values])\n",
+    "plt.plot(broadband_metagrating_loss_values)\n",
+    "plt.plot([m[\"distance_to_window\"] for m in broadband_metagrating_metrics_values])\n",
     "plt.subplot(122)\n",
-    "plt.imshow(jnp.tile(metagrating_params.array, (1, 1)))\n",
+    "plt.imshow(jnp.tile(broadband_metagrating_params.array, (1, 1)))\n",
     "plt.colorbar()"
    ]
   },
@@ -152,14 +177,14 @@
     "pec = challenge.photon_extractor()\n",
     "\n",
     "(\n",
-    "    extractor_loss_values,\n",
-    "    extractor_metrics_values,\n",
-    "    extractor_response,\n",
     "    extractor_params,\n",
+    "    extractor_response,\n",
     "    extractor_aux,\n",
+    "    extractor_loss_values,\n",
+    "    extractor_metrics_values\n",
     ") = optimize(challenge=pec, opt=optax.adam(0.02), num_steps=100)\n",
     "\n",
-    "extractor_aux = pec.component.response(extractor_params, compute_fields=True)"
+    "_, extractor_aux = pec.component.response(extractor_params, compute_fields=True)"
    ]
   },
   {
@@ -191,7 +216,122 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e5514da1-5ae2-4d91-a4dc-997bf21aadd8",
+   "id": "99ed0be6-0cc3-49ce-8d98-957168ce6fba",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(extractor_aux[0])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "32dc5a87-61c9-4785-b19b-81e3b4824fe8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Jones direct, parallelogramic, 81\n",
+    "\n",
+    "# path = \"../../../../Desktop/broadband_metagrating.png\"\n",
+    "\n",
+    "# im = onp.asarray(Image.open(path), dtype=float)\n",
+    "# im = onp.mean(im, axis=-1)\n",
+    "\n",
+    "\n",
+    "# # Trim the \"obviously border\" pixels.\n",
+    "# design = im\n",
+    "# design = design[113:553, 144:1056]\n",
+    "\n",
+    "# # Remove the text and normalize.\n",
+    "# design -= onp.percentile(design, 5)\n",
+    "# design[:100, 500:800] = 0.0\n",
+    "# design /= onp.percentile(design, 95)\n",
+    "# design[185:250, 200:455] = 1.0\n",
+    "# design = onp.clip(design, 0, 1)\n",
+    "\n",
+    "# # Remove tick marks.\n",
+    "# design[:20, 90:120] = 1.0\n",
+    "# design[-20:, 90:120] = 1.0\n",
+    "# design[210:225, :10] = 1.0\n",
+    "# design[210:225, -10:] = 1.0\n",
+    "# design[:20, 200:] = 0.0\n",
+    "# design[-20:, 200:] = 0.0\n",
+    "# design[:150, :20] = 0.0\n",
+    "# design[:150, -20:] = 0.0\n",
+    "# design[-150:, :20] = 0.0\n",
+    "# design[-150:, -20:] = 0.0\n",
+    "\n",
+    "# print(design.shape[1] / design.shape[0])\n",
+    "# print(411 / 199)\n",
+    "\n",
+    "# # Zoom to achieve the specified output shape.\n",
+    "# design_shape = (199 * 3, 411 * 3)\n",
+    "# zoom = (design_shape[0] / design.shape[0], design_shape[1] / design.shape[1])\n",
+    "# design = ndimage.zoom(design, zoom, order=3)\n",
+    "\n",
+    "# design_shape = (199, 411)\n",
+    "# zoom = (design_shape[0] / design.shape[0], design_shape[1] / design.shape[1])\n",
+    "# design = ndimage.zoom(design, zoom, order=1)\n",
+    "\n",
+    "# design = design * 1.05 - 0.025\n",
+    "# design = onp.clip(design, 0, 1)\n",
+    "# density_array = design.T\n",
+    "\n",
+    "##########\n",
+    "\n",
+    "comp = metagrating_challenge.broadband_metagrating(\n",
+    "    sim_params=dataclasses.replace(\n",
+    "        metagrating_challenge.BROADBAND_METAGRATING_SIM_PARAMS,\n",
+    "        approximate_num_terms=81,\n",
+    "        formulation=fmm.Formulation.JONES_DIRECT,\n",
+    "        truncation=basis.Truncation.PARALLELOGRAMIC,\n",
+    "    ),\n",
+    "    spec=dataclasses.replace(\n",
+    "        metagrating_challenge.BROADBAND_METAGRATING_SPEC,\n",
+    "        thickness_grating=0.528,\n",
+    "    ),\n",
+    ").component\n",
+    "density = comp.init(jax.random.PRNGKey(0))\n",
+    "print(comp.expansion.num_terms)\n",
+    "\n",
+    "density_array = onp.genfromtxt(\"../tests/challenge/diffract/broadband_metagrating_designs/device1.csv\", delimiter=\",\")\n",
+    "assert density.shape == density_array.shape\n",
+    "density = dataclasses.replace(density, array=density_array)\n",
+    "\n",
+    "response, _ = comp.response(density, wavelength=jnp.linspace(1.425, 1.675, 100))\n",
+    "\n",
+    "plt.figure(figsize=(7, 4))\n",
+    "efficiency = metagrating_challenge._value_for_order(response.transmission_efficiency, response.expansion, (1, 0))\n",
+    "ax = plt.subplot(121)\n",
+    "plt.plot(response.wavelength, efficiency)\n",
+    "plt.plot(reference_data[:, 0] / 1000, reference_data[:, 1], 'r')\n",
+    "ax.set_ylim([0, 1])\n",
+    "ax.set_xlim([1.425, 1.675])\n",
+    "ax.set_yticks(onp.arange(0, 1.1, 0.1))\n",
+    "ax.set_xlabel(\"Wavelength (um)\")\n",
+    "ax.set_ylabel(\"Transmission coefficient to +1 order\")\n",
+    "ax.grid(True)\n",
+    "ax = plt.subplot(122)\n",
+    "plt.imshow(density_array, cmap=\"gray\")\n",
+    "ax.axis(\"off\")\n",
+    "\n",
+    "plt.tight_layout()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d4d1f78a-d8c6-41fb-8a03-6010721df514",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.plot(comp.expansion.basis_coefficients[:, 1], comp.expansion.basis_coefficients[:, 0], 'o')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "adbb8d00-d98f-429a-a1c2-fdcaf1e7622f",
    "metadata": {},
    "outputs": [],
    "source": []

pyproject.toml

@@ -62,3 +62,6 @@ addopts = [
     # Allow test files to have the same name in different directories.
     "--import-mode=importlib",
 ]
+
+[tool.mypy]
+follow_imports = "normal"

src/invrs_gym/challenge/__init__.py

@@ -1,13 +1,14 @@
 __all__ = [
-    "beam_splitter_challenge",
     "lightweight_beam_splitter_challenge",
     "lightweight_mode_converter_challenge",
     "lightweight_waveguide_bend_challenge",
     "lightweight_wdm_challenge",
+    "beam_splitter_challenge",
     "mode_converter_challenge",
     "waveguide_bend_challenge",
     "wdm_challenge",
     "metagrating",
+    "broadband_metagrating",
     "diffractive_splitter",
     "photon_extractor",
 ]
@@ -22,6 +23,9 @@
     waveguide_bend_challenge,
     wdm_challenge,
 )
-from invrs_gym.challenge.diffract.metagrating_challenge import metagrating
+from invrs_gym.challenge.diffract.metagrating_challenge import (
+    broadband_metagrating,
+    metagrating,
+)
 from invrs_gym.challenge.diffract.splitter_challenge import diffractive_splitter
 from invrs_gym.challenge.extractor.challenge import photon_extractor

src/invrs_gym/challenge/ceviche/challenge.py

@@ -244,7 +244,7 @@ def metrics(
 
 
 def _wavelength_bound(
-    band_bound: jnp.ndarray, transmission_shape: Tuple[int, int, int]
+    band_bound: jnp.ndarray, transmission_shape: Tuple[int, ...]
 ) -> jnp.ndarray:
     """Obtain per-wavelength bound compatible with `transmission_shape`."""
     assert len(transmission_shape) == 3

src/invrs_gym/challenge/diffract/metagrating_challenge.py

@@ -180,27 +180,119 @@ def _value_for_order(
 
 SYMMETRIES = (symmetry.REFLECTION_E_W,)
 
-# Minimum width and spacing are approximately 80 nm for the default dimensions
-# of 1.371 x 0.525 um and grid shape of (118, 45).
-MINIMUM_WIDTH = 7
-MINIMUM_SPACING = 7
-
 # Objective is to diffract light into the +1 transmitted order, with efficiency better
 # than 95 percent.
 TRANSMISSION_ORDER = (1, 0)
 TRANSMISSION_LOWER_BOUND = 0.95
 
 
 def metagrating(
-    minimum_width: int = MINIMUM_WIDTH,
-    minimum_spacing: int = MINIMUM_SPACING,
+    minimum_width: int = 7,
+    minimum_spacing: int = 7,
     density_initializer: DensityInitializer = common.identity_initializer,
     transmission_order: Tuple[int, int] = TRANSMISSION_ORDER,
     transmission_lower_bound: float = TRANSMISSION_LOWER_BOUND,
     spec: common.GratingSpec = METAGRATING_SPEC,
     sim_params: common.GratingSimParams = METAGRATING_SIM_PARAMS,
 ) -> MetagratingChallenge:
-    """Metagrating with 1.371 x 0.525 um design region."""
+    """Metagrating challenge with 1.371 x 0.525 um design region.
+
+    The metagrating challenge is based on the metagrating example in "Validation and
+    characterization of algorithms for photonics inverse design" by Chen et al.
+    (in preparation).
+
+    It involves maximizing diffraction of light transmitted from a silicon oxide
+    substrate into the ambient using a patterned silicon metastructure. The excitation
+    is TM-polarized plane wave with 1.05 micron wavelength.
+
+    Args:
+        minimum_width: The minimum width target for the challenge, in pixels. The
+            physical minimum width is approximately 80 nm.
+        minimum_spacing: The minimum spacing target for the challenge, in pixels.
+        density_initializer: Callable which returns the initial density, given a
+            key and seed density.
+        transmission_order: The diffraction order to be maximized.
+        transmission_lower_bound: The lower bound for transmission. When the lower
+            bound is exceeded, the challenge is considered to be solved.
+        spec: Defines the physical specification of the metagrating.
+        sim_params: Defines the simulation settings of the metagrating.
+
+    Returns:
+        The `MetagratingChallenge`.
+    """
+    return MetagratingChallenge(
+        component=MetagratingComponent(
+            spec=spec,
+            sim_params=sim_params,
+            density_initializer=density_initializer,
+            minimum_width=minimum_width,
+            minimum_spacing=minimum_spacing,
+            symmetries=SYMMETRIES,
+        ),
+        transmission_order=transmission_order,
+        transmission_lower_bound=transmission_lower_bound,
+    )
+
+
+# -----------------------------------------------------------------------------
+# Broadband metagrating with 1.807 x 0.874 um design region.
+# -----------------------------------------------------------------------------
+
+
+BROADBAND_METAGRATING_SPEC = common.GratingSpec(
+    permittivity_ambient=(1.0 + 0.0j) ** 2,
+    permittivity_grating=(3.5 + 0.0j) ** 2,
+    permittivity_encapsulation=(1.0 + 0.0j) ** 2,
+    permittivity_substrate=(1.45 + 0.0j) ** 2,
+    thickness_grating=0.528,
+    period_x=1.807,
+    period_y=0.874,
+)
+
+BROADBAND_METAGRATING_SIM_PARAMS = common.GratingSimParams(
+    grid_shape=(411, 199),
+    wavelength=jnp.asarray((1.530, 1.538, 1.546, 1.554, 1.562, 1.570)),
+    polarization=common.TM,
+    formulation=fmm.Formulation.JONES_DIRECT,
+    approximate_num_terms=400,
+    truncation=basis.Truncation.CIRCULAR,
+)
+
+
+def broadband_metagrating(
+    minimum_width: int = 41,
+    minimum_spacing: int = 41,
+    density_initializer: DensityInitializer = common.identity_initializer,
+    transmission_order: Tuple[int, int] = TRANSMISSION_ORDER,
+    transmission_lower_bound: float = TRANSMISSION_LOWER_BOUND,
+    spec: common.GratingSpec = BROADBAND_METAGRATING_SPEC,
+    sim_params: common.GratingSimParams = BROADBAND_METAGRATING_SIM_PARAMS,
+) -> MetagratingChallenge:
+    """Broadband metagrating challenge with 1.807 x 0.874 um design region.
+
+    The broadband metagrating challenge is based on "General-purpose algorithm for
+    two-material minimum feature size enforcement of nanophotonic devices" by Jenkins
+    et al. (https://pubs.acs.org/doi/abs/10.1021/acsphotonics.2c01166).
+
+    It involves maximizing diffraction of light transmitted from a silicon oxide
+    substrate into the ambient using a patterned silicon metastructure. The excitation
+    consists of TM-polarized plane waves with wavelengths near 1.550 microns.
+
+    Args:
+        minimum_width: The minimum width target for the challenge, in pixels. The
+            physical minimum width is approximately 180 nm.
+        minimum_spacing: The minimum spacing target for the challenge, in pixels.
+        density_initializer: Callable which returns the initial density, given a
+            key and seed density.
+        transmission_order: The diffraction order to be maximized.
+        transmission_lower_bound: The lower bound for transmission. When the lower
+            bound is exceeded, the challenge is considered to be solved.
+        spec: Defines the physical specification of the metagrating.
+        sim_params: Defines the simulation settings of the metagrating.
+
+    Returns:
+        The `MetagratingChallenge`.
+    """
     return MetagratingChallenge(
         component=MetagratingComponent(
             spec=spec,