"""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, Iterable, List, Tuple import jax import jax.numpy as jnp import numpy as onp import meep as mp from . import DesignRegion, EigenmodeCoefficient, utils _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: Iterable[onp.ndarray] ) -> (jnp.ndarray, List[List[mp.DftFields]]): """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: onp.ndarray ) -> List[List[mp.DftFields]]: """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: List[List[mp.DftFields]], adj_fields: List[List[mp.DftFields]], design_variable_shapes: List[Tuple[int, ...]], sum_freq_partials: bool = True, ) -> List[onp.ndarray]: """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