metalens suitable for eval

invrs-io · Jul 2, 2024 · 7a70fcd · 7a70fcd
1 parent 86c7bc6
commit 7a70fcd
Show file tree

Hide file tree

Showing 39 changed files with 12,898 additions and 12,268 deletions.
diff --git a/.github/workflows/eval.yml b/.github/workflows/eval.yml
@@ -34,3 +34,4 @@ jobs:
           python scripts/eval.py diffractive_splitter reference_designs/diffractive_splitter/device1.csv
           python scripts/eval.py photon_extractor reference_designs/photon_extractor/device1.csv
           python scripts/eval.py polarization_sorter reference_designs/polarization_sorter/device1.json
+          python scripts/eval.py metalens reference_designs/metalens/device1.json
diff --git a/docs/notebooks/metalens_challenge.ipynb b/docs/notebooks/metalens_challenge.ipynb
@@ -29,16 +29,10 @@
     "from skimage import measure\n",
     "\n",
     "design = onp.genfromtxt(\n",
-    "    \"../../reference_designs/metalens/Ex/Rasmus70nm.csv\",\n",
+    "    \"../../reference_designs/metalens/Rasmus70nm.csv\",\n",
     "    delimiter=\",\",\n",
     ")\n",
     "\n",
-    "# Flip the design, so that row `0` corresponds to the top, and `-1` the bottom.\n",
-    "design = design[:, ::-1]\n",
-    "\n",
-    "# The `Rasmus70nm` design has grid spacing of 10 nm; some other designs use 20 nm.\n",
-    "grid_spacing = 0.010\n",
-    "\n",
     "plt.figure(figsize=(8, 2))\n",
     "ax = plt.subplot(111)\n",
     "im = ax.imshow(1 - design.T, cmap=\"gray\")\n",
@@ -47,7 +41,8 @@
     "for c in contours:\n",
     "    plt.plot(c[:, 1], c[:, 0], \"k\", lw=1)\n",
     "ax.set_xticks([])\n",
-    "_ = ax.set_yticks([])"
+    "ax.set_yticks([])\n",
+    "_ = ax.set_xlim(100, 700)"
    ]
   },
   {
@@ -67,18 +62,16 @@
     "from invrs_gym import challenges\n",
     "from invrs_gym.challenges.metalens import challenge as metalens_challenge\n",
     "\n",
+    "# By default the simulation does not use PML, but we include it here for visuals.\n",
+    "grid_spacing = metalens_challenge.METALENS_SPEC.grid_spacing\n",
     "spec = dataclasses.replace(\n",
     "    metalens_challenge.METALENS_SPEC,\n",
-    "    width_lens=design.shape[0] * grid_spacing,\n",
     "    width_pml=0.5,\n",
-    "    thickness_lens=design.shape[1] * grid_spacing,\n",
+    "    pml_lens_offset=2.5,\n",
+    "    pml_source_offset=2.0,\n",
     "    grid_spacing=grid_spacing,\n",
     ")\n",
     "\n",
-    "# Pad design so that it extends to the edge of the simulation domain.\n",
-    "padding = int(onp.around(spec.width / spec.grid_spacing)) - design.shape[0]\n",
-    "padded_design = onp.pad(design, ((padding // 2, padding // 2), (0, 0)), mode=\"edge\")\n",
-    "\n",
     "challenge = challenges.metalens(spec=spec)"
    ]
   },
@@ -98,7 +91,9 @@
     "import jax\n",
     "\n",
     "dummy_params = challenge.component.init(jax.random.PRNGKey(0))\n",
-    "params = dataclasses.replace(dummy_params, array=padded_design)"
+    "params = dataclasses.replace(\n",
+    "    dummy_params, array=design, fixed_solid=None, fixed_void=None\n",
+    ")"
    ]
   },
   {
@@ -198,26 +193,15 @@
     "import glob\n",
     "import imageruler\n",
     "\n",
-    "results = {\"invrsio\": [], \"rasmus\": [], \"mo\": []}\n",
+    "results = {\"rasmus\": [], \"mo\": []}\n",
+    "challenge = metalens_challenge.metalens()\n",
+    "dummy_params = challenge.component.init(jax.random.PRNGKey(0))\n",
     "\n",
-    "for fname in glob.glob(\"../../reference_designs/metalens/Ex/*.csv\"):\n",
+    "for fname in glob.glob(\"../../reference_designs/metalens/*.csv\"):\n",
     "    design = onp.genfromtxt(fname, delimiter=\",\")\n",
-    "    design = design[:, ::-1]\n",
-    "\n",
-    "    grid_spacing = 0.01 if \"Rasmus\" in fname else 0.02\n",
-    "    spec = dataclasses.replace(\n",
-    "        metalens_challenge.METALENS_SPEC,\n",
-    "        width_lens=design.shape[0] * grid_spacing,\n",
-    "        width_pml=0.5,\n",
-    "        thickness_lens=design.shape[1] * grid_spacing,\n",
-    "        grid_spacing=grid_spacing,\n",
+    "    params = dataclasses.replace(\n",
+    "        dummy_params, array=design, fixed_solid=None, fixed_void=None\n",
     "    )\n",
-    "    challenge = challenges.metalens(spec=spec)\n",
-    "    dummy_params = challenge.component.init(jax.random.PRNGKey(0))\n",
-    "\n",
-    "    padding = int(onp.around(spec.width / spec.grid_spacing)) - design.shape[0]\n",
-    "    padded_design = onp.pad(design, ((padding // 2, padding // 2), (0, 0)), mode=\"edge\")\n",
-    "    params = dataclasses.replace(dummy_params, array=padded_design)\n",
     "\n",
     "    response, aux = challenge.component.response(params)\n",
     "\n",
@@ -230,8 +214,6 @@
     "        key = \"mo\"\n",
     "    elif \"Rasmus\" in fname:\n",
     "        key = \"rasmus\"\n",
-    "    elif \"invrsio\" in fname:\n",
-    "        key = \"invrsio\"\n",
     "\n",
     "    results[key].append(\n",
     "        {\"response\": response, \"length_scale\": length_scale, \"aux\": aux}\n",
@@ -244,7 +226,7 @@
    "source": [
     "Plotting, we see results very similar to those of figure 3 of \"[Validation and characterization of algorithms and software for photonics inverse design](https://opg.optica.org/josab/abstract.cfm?URI=josab-41-2-A161)\" by Chen et al.\n",
     "\n",
-    "Differences are explained by updates to the imageruler algorithm, and the fact that we are using a different simulation algorithm to model the metalenses."
+    "Differences are explained by updates to the imageruler algorithm and the fact that we are using a different simulation algorithm to model the metalenses. Also, the Rasmus-generated designs have been resampled to have 20 nm grid spacing (originally 10 nm) as expected by the metalens challenge."
    ]
   },
   {