/
wrapper.py
233 lines (202 loc) · 8.89 KB
/
wrapper.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
"""Wrapper for converting a Meep simulation into a differentiable JAX callable function.
Usage example:
```
import jax.numpy as jnp
import meep as mp
import meep.adjoint as mpa
sources = [
mp.EigenModeSource(...)
]
monitors = [
mpa.EigenmodeCoefficient(...),
mpa.EigenmodeCoefficient(...),
]
design_regions = [
mpa.DesignRegion(...)
]
frequencies = [1/1.55, 1/1.60, 1/1.65, ...]
simulation = mp.Simulation(...)
wrapped_meep = MeepJaxWrapper(
simulation,
sources,
monitors,
design_regions,
frequencies,
measurement_interval = 50.0,
dft_field_components = (mp.Ez,),
dft_threshold = 1e-6,
minimum_run_time = 0,
maximum_run_time = onp.inf,
until_after_sources = True
)
def loss(x):
monitor_values = wrapped_meep([x])
t = monitor_values[0,:] / monitor_values[1,:]
# Mean transmission vs wavelength
return jnp.mean(jnp.abs(t))
value, grad = jax.value_and_grad(loss)(x)
```
"""
from typing import Callable, List, Tuple
import jax
import jax.numpy as jnp
import meep as mp
import numpy as onp
from . import utils
from . import DesignRegion, EigenmodeCoefficient
_norm_fn = onp.linalg.norm
_reduce_fn = onp.max
class MeepJaxWrapper:
"""Wraps a Meep simulation object into a JAX-differentiable callable.
Attributes:
simulation: the pre-configured Meep simulation object.
sources: a list of Meep sources for the forward simulation.
monitors: a list of eigenmode coefficient monitors from the `meep.adjoint`
module.
design_regions: a list of design regions from the `meep.adjoint` module.
frequencies: the list of frequencies, in normalized Meep units.
measurement_interval: the time interval between DFT field convergence
measurements, in Meep time units. The default value is 50.
dft_field_components: a list of Meep field components, such as `mp.Ex`,
`mp.Hy`, etc, whose DFT will be monitored for convergence to stop the
simulation. The default is `mp.Ez`.
dft_threshold: the threshold for DFT field convergence. Once the norm of the
change in the fields (the maximum over all design regions and field
components) is less than this value, the simulation will be stopped. The
default value is 1e-6.
minimum_run_time: the minimum run time of the simulation, in Meep time
units. The default value is 0.
maximum_run_time: the maximum run time of the simulation, in Meep time
units. The default value is infinity.
until_after_sources: whether `maximum_run_time` should be ignored until the
sources have turned off. This parameter specifies whether `until` or
`until_after_sources` is used. See
https://meep.readthedocs.io/en/latest/Python_User_Interface/#Simulation
for more information. The default is true.
"""
_log_fn = print
def __init__(
self,
simulation: mp.Simulation,
sources: List[mp.Source],
monitors: List[EigenmodeCoefficient],
design_regions: List[DesignRegion],
frequencies: List[float],
dft_threshold: float = 1e-11,
minimum_run_time: float = 0,
maximum_run_time: float = onp.inf,
until_after_sources: bool = True,
finite_difference_step: float = utils.FD_DEFAULT
):
self.simulation = simulation
self.sources = sources
self.monitors = monitors
self.design_regions = design_regions
self.frequencies = frequencies
self.dft_threshold = dft_threshold
self.minimum_run_time = minimum_run_time
self.maximum_run_time = maximum_run_time
self.until_after_sources = until_after_sources
self.finite_difference_step = finite_difference_step
self._simulate_fn = self._initialize_callable()
def __call__(self, designs: List[jnp.ndarray]) -> jnp.ndarray:
"""Performs a Meep simulation, taking a list of designs and returning mode overlaps.
Args:
designs: a list of design variables as 1D, 2D, or 3D JAX arrays. Valid shapes for
design variables are (Nx, Ny, Nz) where Nx{y,z} match the elements of the
`grid_size` constructor argument of Meep's `MaterialGrid` used for the
corresponding design region. Singleton dimensions of the `grid_size` may be
omitted from the corresponding design variable. For example, a design variable
with a shape of either (10, 20) or (10, 20, 1) would be compatible with a
`grid_size` of (10, 20, 1). Similarly, a design variable with shapes of (25,),
(25, 1), or (25, 1, 1) would be compatible with a `grid_size` of (25, 1, 1).
Returns:
a complex-valued JAX ndarray of differentiable mode monitor overlaps values with
a shape of (num monitors, num frequencies).
"""
return self._simulate_fn(designs)
def _run_fwd_simulation(self, design_variables):
"""Runs forward simulation, returning monitor values and design region fields."""
utils.validate_and_update_design(self.design_regions, design_variables)
self.simulation.reset_meep()
self.simulation.change_sources(self.sources)
utils.register_monitors(self.monitors, self.frequencies)
fwd_design_region_monitors = utils.install_design_region_monitors(
self.simulation,
self.design_regions,
self.frequencies,
)
self.simulation.init_sim()
sim_run_args = {
'until_after_sources' if self.until_after_sources else 'until':
mp.stop_when_dft_decayed(self.dft_threshold,self.minimum_run_time,self.maximum_run_time)
}
self.simulation.run(**sim_run_args)
monitor_values = utils.gather_monitor_values(self.monitors)
return (jnp.asarray(monitor_values),
fwd_design_region_monitors)
def _run_adjoint_simulation(self, monitor_values_grad):
"""Runs adjoint simulation, returning design region fields."""
if not self.design_regions:
raise RuntimeError(
'An adjoint simulation was attempted when no design '
'regions are present.')
adjoint_sources = utils.create_adjoint_sources(self.monitors,
monitor_values_grad)
# TODO refactor with optimization_problem.py #
self.simulation.restart_fields()
self.simulation.clear_dft_monitors()
self.simulation.change_sources(adjoint_sources)
# #
adj_design_region_monitors = utils.install_design_region_monitors(
self.simulation,
self.design_regions,
self.frequencies,
)
self.simulation.init_sim()
sim_run_args = {
'until_after_sources' if self.until_after_sources else 'until':
mp.stop_when_dft_decayed(self.dft_threshold,self.minimum_run_time,self.maximum_run_time)
}
self.simulation.run(**sim_run_args)
return adj_design_region_monitors
def _calculate_vjps(
self,
fwd_fields,
adj_fields,
design_variable_shapes,
sum_freq_partials=True,
):
"""Calculates the VJP for a given set of forward and adjoint fields."""
return utils.calculate_vjps(
self.simulation,
self.design_regions,
self.frequencies,
fwd_fields,
adj_fields,
design_variable_shapes,
sum_freq_partials=sum_freq_partials,
finite_difference_step=self.finite_difference_step
)
def _initialize_callable(
self) -> Callable[[List[jnp.ndarray]], jnp.ndarray]:
"""Initializes the callable JAX function and registers its VJP."""
@jax.custom_vjp
def simulate(design_variables: List[jnp.ndarray]) -> jnp.ndarray:
monitor_values, _ = self._run_fwd_simulation(design_variables)
return monitor_values
def _simulate_fwd(design_variables):
"""Runs forward simulation, returning monitor values and fields."""
monitor_values, self.fwd_design_region_monitors = self._run_fwd_simulation(
design_variables)
design_variable_shapes = [x.shape for x in design_variables]
return monitor_values, (design_variable_shapes)
def _simulate_rev(res, monitor_values_grad):
"""Runs adjoint simulation, returning VJP of design wrt monitor values."""
design_variable_shapes = res
self.adj_design_region_monitors = self._run_adjoint_simulation(monitor_values_grad)
vjps = self._calculate_vjps(self.fwd_design_region_monitors, self.adj_design_region_monitors,
design_variable_shapes)
return ([jnp.asarray(vjp) for vjp in vjps], )
simulate.defvjp(_simulate_fwd, _simulate_rev)
return simulate