Merge pull request #2 from invrs-io/ceviche

Ceviche defaults

notebooks/ceviche.ipynb

@@ -0,0 +1,102 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b4b78fc5-34c7-4e23-b77f-9e34a8718cd3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from importlib import reload\n",
+    "\n",
+    "import time\n",
+    "import jax\n",
+    "import jax.numpy as jnp\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as onp\n",
+    "\n",
+    "from ceviche_challenges import units as u\n",
+    "from ceviche_challenges import beam_splitter, mode_converter, params, waveguide_bend, wdm\n",
+    "\n",
+    "from invrs_gym.challenge.ceviche import defaults"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "95f003fc-c137-4099-81ad-1e6cfae0f41b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = wdm.model.WdmModel(defaults.LIGHTWEIGHT_SIM_PARAMS, defaults.LIGHTWEIGHT_WDM_SPEC)\n",
+    "density = jnp.full(model.design_variable_shape, 1.0)\n",
+    "plt.imshow(model.density(density)[20:-20, 20:-20])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "33898107-7832-4ad1-a980-2873fb993dfb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = waveguide_bend.model.WaveguideBendModel(defaults.LIGHTWEIGHT_SIM_PARAMS, defaults.WAVEGUIDE_BEND_SPEC)\n",
+    "density = jnp.full(model.design_variable_shape, 1.0)\n",
+    "plt.imshow(model.density(density)[20:-20, 20:-20])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "23bf803f-5359-48de-a497-064fe6ef688a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = beam_splitter.model.BeamSplitterModel(defaults.LIGHTWEIGHT_SIM_PARAMS, defaults.BEAM_SPLITTER_SPEC)\n",
+    "density = jnp.full(model.design_variable_shape, 1.0)\n",
+    "plt.imshow(model.density(density)[20:-20, 20:-20])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b553c6fc-0c0f-4450-b492-fb27ba84644d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = mode_converter.model.ModeConverterModel(defaults.LIGHTWEIGHT_SIM_PARAMS, defaults.MODE_CONVERTER_SPEC)\n",
+    "density = jnp.full(model.design_variable_shape, 1.0)\n",
+    "plt.imshow(model.density(density)[20:-20, 20:-20])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "474e7001-99c1-4a7a-94bc-e1191628d313",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "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.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

src/invrs_gym/challenge/ceviche/defaults.py

@@ -0,0 +1,179 @@
+"""Defines defaults for the ceviche challenges."""
+
+from ceviche_challenges import units as u
+from ceviche_challenges import (
+    beam_splitter,
+    mode_converter,
+    params,
+    waveguide_bend,
+    wdm,
+)
+
+
+WG_WIDTH = 400 * u.nm
+WG_MODE_PADDING = 560 * u.nm
+WG_LENGTH = 720 * u.nm
+
+PADDING = 400 * u.nm
+PORT_PML_OFFSET = 40 * u.nm
+
+CLADDING_PERMITTIVITY = 2.25
+SLAB_PERMITTIVITY = 12.25
+INPUT_MONITOR_OFFSET = 40 * u.nm
+
+PML_WIDTH_GRIDPOINTS = 20
+
+
+# -----------------------------------------------------------------------------
+# Simulation parameter defaults.
+# -----------------------------------------------------------------------------
+
+
+SIM_PARAMS = params.CevicheSimParams(
+    resolution=10 * u.nm,
+    wavelengths=u.Array([1265.0, 1270.0, 1275.0, 1285.0, 1290.0, 1295.0], u.nm),
+)
+
+LIGHTWEIGHT_SIM_PARAMS = params.CevicheSimParams(
+    resolution=40 * u.nm,
+    wavelengths=u.Array([1270.0, 1290.0], u.nm),
+)
+
+# -----------------------------------------------------------------------------
+# Defaults for the spec, which define the geometry of devices.
+# -----------------------------------------------------------------------------
+
+# Beamsplitter with a 2.0 x 3.2 um design region.
+BEAM_SPLITTER_SPEC = beam_splitter.spec.BeamSplitterSpec(
+    wg_width=WG_WIDTH,
+    wg_length=WG_LENGTH,
+    wg_separation=1120 * u.nm,
+    wg_mode_padding=WG_MODE_PADDING,
+    port_pml_offset=PORT_PML_OFFSET,
+    variable_region_size=(3200 * u.nm, 2000 * u.nm),
+    cladding_permittivity=CLADDING_PERMITTIVITY,
+    slab_permittivity=SLAB_PERMITTIVITY,
+    input_monitor_offset=INPUT_MONITOR_OFFSET,
+    design_symmetry=None,
+    pml_width=PML_WIDTH_GRIDPOINTS,
+)
+
+# Mode converter with a 1.6 x 1.6 um design region.
+MODE_CONVERTER_SPEC = mode_converter.spec.ModeConverterSpec(
+    left_wg_width=WG_WIDTH,
+    left_wg_mode_padding=WG_MODE_PADDING,
+    left_wg_mode_order=1,  # Fundamental mode.
+    right_wg_width=WG_WIDTH,
+    right_wg_mode_padding=WG_MODE_PADDING,
+    right_wg_mode_order=2,  # Second mode.
+    wg_length=WG_LENGTH,
+    padding=PADDING,
+    port_pml_offset=PORT_PML_OFFSET,
+    variable_region_size=(1600 * u.nm, 1600 * u.nm),
+    cladding_permittivity=CLADDING_PERMITTIVITY,
+    slab_permittivity=SLAB_PERMITTIVITY,
+    input_monitor_offset=INPUT_MONITOR_OFFSET,
+    pml_width=PML_WIDTH_GRIDPOINTS,
+)
+
+# Waveguide bend with a 1.6 x 1.6 um design region.
+WAVEGUIDE_BEND_SPEC = waveguide_bend.spec.WaveguideBendSpec(
+    wg_width=WG_WIDTH,
+    wg_length=WG_LENGTH,
+    wg_mode_padding=WG_MODE_PADDING,
+    padding=PADDING,
+    port_pml_offset=PORT_PML_OFFSET,
+    variable_region_size=(1600 * u.nm, 1600 * u.nm),
+    cladding_permittivity=CLADDING_PERMITTIVITY,
+    slab_permittivity=SLAB_PERMITTIVITY,
+    input_monitor_offset=INPUT_MONITOR_OFFSET,
+    pml_width=PML_WIDTH_GRIDPOINTS,
+)
+
+
+def make_wdm_spec(
+    design_extent_ij: u.Array,
+    intended_sim_resolution: u.Quantity,
+) -> wdm.spec.WdmSpec:
+    """Construct a `wdm.spec.WdmSpec` with the specified design region extent.
+
+    Since the `wdm.spec.WdmSpec` does not automatically compensate the simulation
+    extent to account for changes in the PML region thickness with changes in the
+    simulation resolution, the intended simulation resolution must be specified.
+    The actual simulation extent appropriate for that resolution (i.e. accounting
+    for the thickness of PML layers) is then computed.
+
+    Note that if the actual simulation resolution is coarser than the intended
+    resolution specified here, the PML may overlap with the source and monitors.
+    Care should be taken to ensure the correct resolution is specified here.
+
+    Args:
+        design_extent_ij: Specifies the size of the design region.
+        intended_sim_resolution: Specifies the simulation resolution to be used.
+
+    Returns:
+        The `wdm.spec.WdmSpec`.
+    """
+
+    design_extent_i, design_extent_j = design_extent_ij
+    design_extent_i_nm: int = u.resolve(design_extent_i, 1 * u.nm)
+    design_extent_j_nm: int = u.resolve(design_extent_j, 1 * u.nm)
+
+    wg_width_nm: int = u.resolve(WG_WIDTH, 1 * u.nm)
+
+    pad_width_nm: int = u.resolve(PADDING, 1 * u.nm)
+    pml_width_nm: int = u.resolve(
+        PML_WIDTH_GRIDPOINTS * intended_sim_resolution, 1 * u.nm
+    )
+
+    extent_i_nm: int = design_extent_i_nm + 2 * pad_width_nm + 2 * pml_width_nm
+    extent_j_nm: int = design_extent_j_nm + 2 * pad_width_nm + 2 * pml_width_nm
+
+    return wdm.spec.WdmSpec(
+        extent_ij=u.Array([extent_i_nm, extent_j_nm], u.nm),
+        input_wg_j=extent_j_nm / 2 * u.nm,
+        output_wgs_j=u.Array(
+            (
+                extent_j_nm / 2 - design_extent_j_nm / 2 + wg_width_nm / 2 + 320,
+                extent_j_nm / 2 + design_extent_j_nm / 2 - wg_width_nm / 2 - 320,
+            ),
+            u.nm,
+        ),
+        wg_width=WG_WIDTH,
+        wg_mode_padding=WG_MODE_PADDING,
+        input_mode_i=pml_width_nm * u.nm + PORT_PML_OFFSET,
+        output_mode_i=(extent_i_nm - pml_width_nm) * u.nm - PORT_PML_OFFSET,
+        variable_region=(
+            u.Array(
+                (
+                    extent_i_nm / 2 - design_extent_i_nm / 2,
+                    extent_j_nm / 2 - design_extent_j_nm / 2,
+                ),
+                u.nm,
+            ),
+            u.Array(
+                (
+                    extent_i_nm / 2 + design_extent_i_nm / 2,
+                    extent_j_nm / 2 + design_extent_j_nm / 2,
+                ),
+                u.nm,
+            ),
+        ),
+        cladding_permittivity=CLADDING_PERMITTIVITY,
+        slab_permittivity=SLAB_PERMITTIVITY,
+        input_monitor_offset=INPUT_MONITOR_OFFSET,
+        pml_width=PML_WIDTH_GRIDPOINTS,
+    )
+
+
+# Wavelength demultiplexer with 3.2 x 3.2 um design region.
+LIGHTWEIGHT_WDM_SPEC = make_wdm_spec(
+    design_extent_ij=u.Array([3200, 3200], u.nm),
+    intended_sim_resolution=LIGHTWEIGHT_SIM_PARAMS.resolution,
+)
+
+# Wavelength demultiplexer with 6.4 x 6.4 um design region.
+WDM_SPEC = make_wdm_spec(
+    design_extent_ij=u.Array([6400, 6400], u.nm),
+    intended_sim_resolution=SIM_PARAMS.resolution,
+)

tests/challenge/ceviche/test_defaults.py

@@ -0,0 +1,59 @@
+"""Tests for `challenge.ceviche.defaults`."""
+
+import unittest
+
+from ceviche_challenges import (
+    beam_splitter,
+    mode_converter,
+    waveguide_bend,
+    wdm,
+)
+from invrs_gym.challenge.ceviche import defaults
+
+
+class CreateModelTest(unittest.TestCase):
+    def test_create_beam_splitter(self):
+        model = beam_splitter.model.BeamSplitterModel(
+            defaults.SIM_PARAMS, defaults.BEAM_SPLITTER_SPEC
+        )
+        self.assertSequenceEqual(model.design_variable_shape, (320, 200))
+
+    def test_create_lightweight_beam_splitter(self):
+        model = beam_splitter.model.BeamSplitterModel(
+            defaults.LIGHTWEIGHT_SIM_PARAMS, defaults.BEAM_SPLITTER_SPEC
+        )
+        self.assertSequenceEqual(model.design_variable_shape, (80, 50))
+
+    def test_create_mode_converter(self):
+        model = mode_converter.model.ModeConverterModel(
+            defaults.SIM_PARAMS, defaults.MODE_CONVERTER_SPEC
+        )
+        self.assertSequenceEqual(model.design_variable_shape, (160, 160))
+
+    def test_create_lightweight_mode_converter(self):
+        model = mode_converter.model.ModeConverterModel(
+            defaults.LIGHTWEIGHT_SIM_PARAMS, defaults.MODE_CONVERTER_SPEC
+        )
+        self.assertSequenceEqual(model.design_variable_shape, (40, 40))
+
+    def test_create_waveguide_bend(self):
+        model = waveguide_bend.model.WaveguideBendModel(
+            defaults.SIM_PARAMS, defaults.WAVEGUIDE_BEND_SPEC
+        )
+        self.assertSequenceEqual(model.design_variable_shape, (160, 160))
+
+    def test_create_lightweight_waveguide_bend(self):
+        model = waveguide_bend.model.WaveguideBendModel(
+            defaults.LIGHTWEIGHT_SIM_PARAMS, defaults.WAVEGUIDE_BEND_SPEC
+        )
+        self.assertSequenceEqual(model.design_variable_shape, (40, 40))
+
+    def test_create_wdm(self):
+        model = wdm.model.WdmModel(defaults.SIM_PARAMS, defaults.WDM_SPEC)
+        self.assertSequenceEqual(model.design_variable_shape, (640, 640))
+
+    def test_create_lightweight_wdm(self):
+        model = wdm.model.WdmModel(
+            defaults.LIGHTWEIGHT_SIM_PARAMS, defaults.LIGHTWEIGHT_WDM_SPEC
+        )
+        self.assertSequenceEqual(model.design_variable_shape, (80, 80))