Skip to content

Commit 190f0ae

Browse files
mfschubertmfschubert
authored
Merge pull request invrs-io#92 from invrs-io/docs2
Add metagrating notebook
2 parents d480657 + b635f10 commit 190f0ae

File tree

3 files changed

+276
-3
lines changed

3 files changed

+276
-3
lines changed

‎docs/_toc.yml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,3 +7,4 @@ parts:
77
- caption: Challenges
88
chapters:
99
- file: notebooks/ceviche_challenge
10+
- file: notebooks/metagrating_challenge

‎docs/notebooks/ceviche_challenge.ipynb

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
"id": "84e596a3",
2020
"metadata": {},
2121
"source": [
22-
"## Simulating an existing design.\n",
22+
"## Simulating an existing design\n",
2323
"\n",
2424
"We'll begin by loading, visualizing, and simulating one design from the photonics-opt-testbed repo."
2525
]
@@ -168,15 +168,17 @@
168168
" im = ax.imshow(density, cmap=\"gray\")\n",
169169
" for c in contours:\n",
170170
" plt.plot(c[:, 1], c[:, 0], \"r\")\n",
171-
" ax.axis(False)\n",
171+
" ax.set_xticks([])\n",
172+
" ax.set_yticks([])\n",
172173
"\n",
173174
" ax = plt.subplot(133)\n",
174175
" fields = onp.real(aux[\"fields\"][2, 0, :, :])\n",
175176
" im = ax.imshow(fields, cmap=\"bwr\")\n",
176177
" im.set_clim([-onp.amax(onp.abs(fields)), onp.amax(onp.abs(fields))])\n",
177178
" for c in contours:\n",
178179
" plt.plot(c[:, 1], c[:, 0], \"k\")\n",
179-
" ax.axis(False)\n",
180+
" ax.set_xticks([])\n",
181+
" ax.set_yticks([])\n",
180182
"\n",
181183
"\n",
182184
"plot_mode_converter(component, params, response, aux)"

‎docs/notebooks/metagrating_challenge.ipynb

Lines changed: 270 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,270 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"id": "823a99c7",
6+
"metadata": {},
7+
"source": [
8+
"# Metagrating challenge\n",
9+
"\n",
10+
"The metagrating challenge entails designing a beam deflector that couples a normally-incident plane wave into one with a polar angle of 50 degrees. This problem was studied in \"[Validation and characterization of algorithms and software for photonics inverse design](https://opg.optica.org/josab/ViewMedia.cfm?uri=josab-41-2-A161)\" by Chen et al.; the associated [photonics-opt-testbed repo](https://github.com/NanoComp/photonics-opt-testbed) contains several example designs."
11+
]
12+
},
13+
{
14+
"cell_type": "markdown",
15+
"id": "7fe52a19",
16+
"metadata": {},
17+
"source": [
18+
"## Simulating an existing design\n",
19+
"\n",
20+
"We'll begin by loading, visualizing, and simulating the two-dimensional from the photonics-opt-testbed repo."
21+
]
22+
},
23+
{
24+
"cell_type": "code",
25+
"execution_count": null,
26+
"id": "58100cea",
27+
"metadata": {},
28+
"outputs": [],
29+
"source": [
30+
"import matplotlib.pyplot as plt\n",
31+
"import numpy as onp\n",
32+
"\n",
33+
"def load_design(name):\n",
34+
" path = f\"../../reference_designs/metagrating/{name}.csv\"\n",
35+
" return onp.genfromtxt(path, delimiter=\",\")\n",
36+
"\n",
37+
"names = [\"device1\", \"device2\", \"device3\", \"device4\"]\n",
38+
"designs = [load_design(name) for name in names]\n",
39+
"\n",
40+
"plt.figure(figsize=(7, 4))\n",
41+
"for i, design in enumerate(designs):\n",
42+
" ax = plt.subplot(1, 4, i + 1)\n",
43+
" ax.imshow(design, cmap=\"gray\")\n",
44+
" _ = ax.axis(False)"
45+
]
46+
},
47+
{
48+
"cell_type": "markdown",
49+
"id": "11568008",
50+
"metadata": {},
51+
"source": [
52+
"Now, we'll create a `metagrating` challenge, which provides everything we need to simulate and optimize the metagrating."
53+
]
54+
},
55+
{
56+
"cell_type": "code",
57+
"execution_count": null,
58+
"id": "fed8a965",
59+
"metadata": {},
60+
"outputs": [],
61+
"source": [
62+
"from invrs_gym import challenges\n",
63+
"\n",
64+
"challenge = challenges.metagrating()"
65+
]
66+
},
67+
{
68+
"cell_type": "markdown",
69+
"id": "0ba9a3a4",
70+
"metadata": {},
71+
"source": [
72+
"To simulate the metagrating, we need to provide a `totypes.types.Density2DArray` object to the `challenge.component.params` method. Obtain dummy parameters using `component.init`, and then overwrite the `array` attribute with the reference design that we want to simulate."
73+
]
74+
},
75+
{
76+
"cell_type": "code",
77+
"execution_count": null,
78+
"id": "281c95f7",
79+
"metadata": {},
80+
"outputs": [],
81+
"source": [
82+
"import dataclasses\n",
83+
"import jax\n",
84+
"\n",
85+
"dummy_params = challenge.component.init(jax.random.PRNGKey(0)) # totypes.types.Density2DArray\n",
86+
"params = dataclasses.replace(dummy_params, array=load_design(\"device1\"))\n",
87+
"\n",
88+
"# Perform simulation using component response method.\n",
89+
"response, aux = challenge.component.response(params)"
90+
]
91+
},
92+
{
93+
"cell_type": "markdown",
94+
"id": "e8b6019b",
95+
"metadata": {},
96+
"source": [
97+
"The `response` contains the transmission and reflection efficiency into each diffraction order, and for TE- and TM-polarized cases. However, we only care about TM diffraction into the +1 order. Fortunately, the `challenge` has a `metrics` method that extracts this value."
98+
]
99+
},
100+
{
101+
"cell_type": "code",
102+
"execution_count": null,
103+
"id": "29b3a077",
104+
"metadata": {},
105+
"outputs": [],
106+
"source": [
107+
"metrics = challenge.metrics(response, params=params, aux=aux)\n",
108+
"print(f\"TM transmission into +1 order: {metrics['average_efficiency'] * 100:.1f}%\")"
109+
]
110+
},
111+
{
112+
"cell_type": "markdown",
113+
"id": "7f18842e",
114+
"metadata": {},
115+
"source": [
116+
"Now let's take a look at the remaining designs."
117+
]
118+
},
119+
{
120+
"cell_type": "code",
121+
"execution_count": null,
122+
"id": "58b02ecc",
123+
"metadata": {},
124+
"outputs": [],
125+
"source": [
126+
"for name in names:\n",
127+
" params = dataclasses.replace(dummy_params, array=load_design(name))\n",
128+
" response, aux = challenge.component.response(params)\n",
129+
" metrics = challenge.metrics(response, params=params, aux=aux)\n",
130+
" print(f\"{name} TM transmission into +1 order: {metrics['average_efficiency'] * 100:.1f}%\")"
131+
]
132+
},
133+
{
134+
"cell_type": "markdown",
135+
"id": "c806574c",
136+
"metadata": {},
137+
"source": [
138+
"These values are all very close to those reported in the [photonics-opt-testbed](https://github.com/NanoComp/photonics-opt-testbed/tree/main/Metagrating3D), indicating that our simulation is converged."
139+
]
140+
},
141+
{
142+
"cell_type": "markdown",
143+
"id": "14a8182d",
144+
"metadata": {},
145+
"source": [
146+
"## Metagrating optimization\n",
147+
"\n",
148+
"Now let's optimize a metagrating. Again we obtain initial random parameters and define the loss function. The loss function will also return the response and the efficiency value, which will let us see how efficiency improves as we optimize."
149+
]
150+
},
151+
{
152+
"cell_type": "code",
153+
"execution_count": null,
154+
"id": "f1a9d668",
155+
"metadata": {},
156+
"outputs": [],
157+
"source": [
158+
"params = challenge.component.init(jax.random.PRNGKey(0))\n",
159+
"\n",
160+
"\n",
161+
"def loss_fn(params):\n",
162+
" response, aux = challenge.component.response(params)\n",
163+
" loss = challenge.loss(response)\n",
164+
" metrics = challenge.metrics(response, params=params, aux=aux)\n",
165+
" efficiency = metrics[\"average_efficiency\"]\n",
166+
" return loss, (response, efficiency)"
167+
]
168+
},
169+
{
170+
"cell_type": "markdown",
171+
"id": "1ed89701",
172+
"metadata": {},
173+
"source": [
174+
"To design the metagrating we'll use the `density_lbfgsb` optimizer from the [invrs-opt](https://github.com/invrs-io/opt) package. Initialize the optimizer state, and then define the `step_fn` which is called at each optimization step, and then simply call it repeatedly to obtain an optimized design."
175+
]
176+
},
177+
{
178+
"cell_type": "code",
179+
"execution_count": null,
180+
"id": "c711308f",
181+
"metadata": {},
182+
"outputs": [],
183+
"source": [
184+
"import invrs_opt\n",
185+
"\n",
186+
"opt = invrs_opt.density_lbfgsb(beta=4)\n",
187+
"state = opt.init(params) # Initialize optimizer state using the initial parameters.\n",
188+
"\n",
189+
"\n",
190+
"@jax.jit\n",
191+
"def step_fn(state):\n",
192+
" params = opt.params(state)\n",
193+
" (value, (_, efficiency)), grad = jax.value_and_grad(loss_fn, has_aux=True)(params)\n",
194+
" state = opt.update(grad=grad, value=value, params=params, state=state)\n",
195+
" return state, (params, efficiency)\n",
196+
"\n",
197+
"\n",
198+
"# Call `step_fn` repeatedly to optimize, and store the results of each evaluation.\n",
199+
"efficiencies = []\n",
200+
"for _ in range(65):\n",
201+
" state, (params, efficiency) = step_fn(state)\n",
202+
" efficiencies.append(efficiency)"
203+
]
204+
},
205+
{
206+
"cell_type": "markdown",
207+
"id": "4adc4a5c",
208+
"metadata": {},
209+
"source": [
210+
"Now let's visualize the trajectory of efficiency, and the final design."
211+
]
212+
},
213+
{
214+
"cell_type": "code",
215+
"execution_count": null,
216+
"id": "18bab5b8",
217+
"metadata": {},
218+
"outputs": [],
219+
"source": [
220+
"from skimage import measure\n",
221+
"\n",
222+
"ax = plt.subplot(121)\n",
223+
"ax.plot(onp.asarray(efficiencies) * 100)\n",
224+
"ax.set_xlabel(\"Step\")\n",
225+
"ax.set_ylabel(\"Diffraction efficiency into +1 order (%)\")\n",
226+
"\n",
227+
"ax = plt.subplot(122)\n",
228+
"ax.imshow(params.array, cmap=\"gray\")\n",
229+
"\n",
230+
"contours = measure.find_contours(onp.asarray(params.array))\n",
231+
"for c in contours:\n",
232+
" ax.plot(c[:, 1], c[:, 0], 'r')\n",
233+
"\n",
234+
"ax.set_xticks([])\n",
235+
"ax.set_yticks([])\n",
236+
"\n",
237+
"print(f\"Final efficiency: {efficiencies[-1] * 100:.1f}%\")\n"
238+
]
239+
},
240+
{
241+
"cell_type": "markdown",
242+
"id": "bb44022e",
243+
"metadata": {},
244+
"source": [
245+
"The final efficiency is above 90%, similar to the reference designs. However, note that the design is not binary, which is a limitation of the `density_lbfgsb` optimizer: it generally does not produce binary solutions. A different optimizer would be required to obtain binary designs."
246+
]
247+
}
248+
],
249+
"metadata": {
250+
"kernelspec": {
251+
"display_name": "Python 3 (ipykernel)",
252+
"language": "python",
253+
"name": "python3"
254+
},
255+
"language_info": {
256+
"codemirror_mode": {
257+
"name": "ipython",
258+
"version": 3
259+
},
260+
"file_extension": ".py",
261+
"mimetype": "text/x-python",
262+
"name": "python",
263+
"nbconvert_exporter": "python",
264+
"pygments_lexer": "ipython3",
265+
"version": "3.10.12"
266+
}
267+
},
268+
"nbformat": 4,
269+
"nbformat_minor": 5
270+
}

0 commit comments

Comments
 (0)