import unittest import jax import jax.numpy as jnp import meep.adjoint as mpa import numpy as onp import parameterized from utils import ApproxComparisonTestCase import meep as mp # The calculation of finite-difference gradients # requires that JAX be operated with double precision jax.config.update("jax_enable_x64", True) # The step size for the finite-difference # gradient calculation _FD_STEP = 1e-4 # The tolerance for the adjoint and finite-difference # gradient comparison _TOL = 0.1 if mp.is_single_precision() else 0.025 # We expect 3 design region monitor pointers # (one for each field component) _NUM_DES_REG_MON = 3 mp.verbosity(0) def build_straight_wg_simulation( wg_width=0.5, wg_padding=1.0, wg_length=1.0, pml_width=1.0, source_to_pml=0.5, source_to_monitor=0.1, frequencies=[1 / 1.55], gaussian_rel_width=0.2, sim_resolution=20, design_region_resolution=20, ): """Builds a simulation of a straight waveguide with a design region segment.""" design_region_shape = (1.0, wg_width) # Simulation domain size sx = 2 * pml_width + 2 * wg_length + design_region_shape[0] sy = ( 2 * pml_width + 2 * wg_padding + max( wg_width, design_region_shape[1], ) ) # Mean / center frequency fmean = onp.mean(frequencies) si = mp.Medium(index=3.4) sio2 = mp.Medium(index=1.44) sources = [ mp.EigenModeSource( mp.GaussianSource(frequency=fmean, fwidth=fmean * gaussian_rel_width), eig_band=1, direction=mp.NO_DIRECTION, eig_kpoint=mp.Vector3(1, 0, 0), size=mp.Vector3(0, wg_width + 2 * wg_padding, 0), center=[-sx / 2 + pml_width + source_to_pml, 0, 0], ), mp.EigenModeSource( mp.GaussianSource(frequency=fmean, fwidth=fmean * gaussian_rel_width), eig_band=1, direction=mp.NO_DIRECTION, eig_kpoint=mp.Vector3(-1, 0, 0), size=mp.Vector3(0, wg_width + 2 * wg_padding, 0), center=[sx / 2 - pml_width - source_to_pml, 0, 0], ), ] nx, ny = int(design_region_shape[0] * design_region_resolution), int( design_region_shape[1] * design_region_resolution ) mat_grid = mp.MaterialGrid( mp.Vector3(nx, ny), sio2, si, grid_type="U_DEFAULT", ) design_regions = [ mpa.DesignRegion( mat_grid, volume=mp.Volume( center=mp.Vector3(), size=mp.Vector3( design_region_shape[0], design_region_shape[1], 0, ), ), ) ] geometry = [ mp.Block( center=mp.Vector3( x=-design_region_shape[0] / 2 - wg_length / 2 - pml_width / 2 ), material=si, size=mp.Vector3(wg_length + pml_width, wg_width, 0), ), # left wg mp.Block( center=mp.Vector3( x=+design_region_shape[0] / 2 + wg_length / 2 + pml_width / 2 ), material=si, size=mp.Vector3(wg_length + pml_width, wg_width, 0), ), # right wg mp.Block( center=design_regions[0].center, size=design_regions[0].size, material=mat_grid, ), # design region ] simulation = mp.Simulation( cell_size=mp.Vector3(sx, sy), boundary_layers=[mp.PML(pml_width)], geometry=geometry, sources=sources, resolution=sim_resolution, ) monitor_centers = [ mp.Vector3(-sx / 2 + pml_width + source_to_pml + source_to_monitor), mp.Vector3(sx / 2 - pml_width - source_to_pml - source_to_monitor), ] monitor_size = mp.Vector3(y=wg_width + 2 * wg_padding) monitors = [ mpa.EigenmodeCoefficient( simulation, mp.Volume(center=center, size=monitor_size), mode=1, forward=forward, ) for center in monitor_centers for forward in [True, False] ] return simulation, sources, monitors, design_regions, frequencies class UtilsTest(unittest.TestCase): def setUp(self): super().setUp() ( self.simulation, self.sources, self.monitors, self.design_regions, self.frequencies, ) = build_straight_wg_simulation() def test_mode_monitor_helpers(self): mpa.utils.register_monitors(self.monitors, self.frequencies) self.simulation.run(until=100) monitor_values = mpa.utils.gather_monitor_values(self.monitors) self.assertEqual(monitor_values.dtype, onp.complex128) self.assertEqual( monitor_values.shape, (len(self.monitors), len(self.frequencies)) ) def test_dist_dft_pointers(self): fwd_design_region_monitors = mpa.utils.install_design_region_monitors( self.simulation, self.design_regions, self.frequencies, ) self.assertEqual(len(fwd_design_region_monitors[0]), _NUM_DES_REG_MON) class WrapperTest(ApproxComparisonTestCase): @parameterized.parameterized.expand( [ ( "1500_1550bw_01relative_gaussian_port1", onp.linspace(1 / 1.50, 1 / 1.55, 3).tolist(), 0.1, 0.5, 0, ), ( "1550_1600bw_02relative_gaussian_port1", onp.linspace(1 / 1.55, 1 / 1.60, 3).tolist(), 0.2, 0.5, 0, ), ( "1500_1600bw_03relative_gaussian_port1", onp.linspace(1 / 1.50, 1 / 1.60, 4).tolist(), 0.3, 0.5, 0, ), ( "1500_1550bw_01relative_gaussian_port2", onp.linspace(1 / 1.50, 1 / 1.55, 3).tolist(), 0.1, 0.5, 1, ), ( "1550_1600bw_02relative_gaussian_port2", onp.linspace(1 / 1.55, 1 / 1.60, 3).tolist(), 0.2, 0.5, 1, ), ( "1500_1600bw_03relative_gaussian_port2", onp.linspace(1 / 1.50, 1 / 1.60, 4).tolist(), 0.3, 0.5, 1, ), ] ) def test_wrapper_gradients( self, _, frequencies, gaussian_rel_width, design_variable_fill_value, excite_port_idx, ): """Tests gradient from the JAX-Meep wrapper against finite differences.""" ( simulation, sources, monitors, design_regions, frequencies, ) = build_straight_wg_simulation( frequencies=frequencies, gaussian_rel_width=gaussian_rel_width ) design_shape = tuple( int(i) for i in design_regions[0].design_parameters.grid_size )[:2] x = onp.ones(design_shape) * design_variable_fill_value # Define a loss function def loss_fn(x, excite_port_idx=0): wrapped_meep = mpa.MeepJaxWrapper( simulation, [sources[excite_port_idx]], monitors, design_regions, frequencies, ) monitor_values = wrapped_meep([x]) s1p, s1m, s2p, s2m = monitor_values t = s2p / s1p if excite_port_idx == 0 else s1m / s2m return jnp.mean(jnp.square(jnp.abs(t))) value, adjoint_grad = jax.value_and_grad(loss_fn)( x, excite_port_idx=excite_port_idx ) projection = [] fd_projection = [] # Project along 5 random directions in the design parameter space. for seed in range(5): # Create dp random_perturbation_vector = _FD_STEP * jax.random.normal( jax.random.PRNGKey(seed), x.shape, ) # Calculate p + dp x_perturbed = x + random_perturbation_vector # Calculate T(p + dp) value_perturbed = loss_fn(x_perturbed, excite_port_idx=excite_port_idx) projection.append( onp.dot( random_perturbation_vector.ravel(), adjoint_grad.ravel(), ) ) fd_projection.append(value_perturbed - value) projection = onp.stack(projection) fd_projection = onp.stack(fd_projection) # Check that dp . ∇T ~ T(p + dp) - T(p) self.assertClose( projection, fd_projection, epsilon=_TOL, ) if __name__ == "__main__": unittest.main()