diff --git a/python/adjoint/angular_spectrum.py b/python/adjoint/angular_spectrum.py index fad8b7587..247858d27 100644 --- a/python/adjoint/angular_spectrum.py +++ b/python/adjoint/angular_spectrum.py @@ -36,7 +36,7 @@ """ import math -from typing import Callable, Dict, NamedTuple, Optional, Sequence, Tuple, Union +from typing import Callable, Dict, NamedTuple, Optional, Sequence, Tuple, Union, List import jax import jax.numpy as jnp @@ -1083,6 +1083,42 @@ def from_monitor( **kwargs, ) + @classmethod + def from_volume( + cls, + simulation: mp.Simulation, + volume: mp.Volume, + stack: Stack, + frequencies, + sign: int = 1, + **kwargs, + ): + """Builds a propagator for a plane, without needing a monitor there. + + `from_monitor` cannot be used on the injection path: a DFT monitor + requires the field components to be allocated, which does not happen + until sources are added, and the sources are what this is being built to + produce. This takes the frequencies directly instead. + + Args: + simulation: supplies the grid; it need not have been run. + volume: the source plane. + stack: the layers between that plane and the mode. + frequencies: the frequencies to propagate. + sign: +1 if the outgoing direction is along +normal. + **kwargs: forwarded to the constructor, e.g. `pad_factor`. + """ + normal, pitch, num_points = cls._plane_geometry(simulation, volume) + return cls( + stack, + onp.asarray(frequencies), + pitch, + num_points, + normal=normal, + sign=sign, + **kwargs, + ) + def fields_from_monitor( self, simulation: mp.Simulation, monitor ) -> TangentialFields: @@ -1152,6 +1188,167 @@ def report_monitor(self, simulation, monitor) -> Dict[str, onp.ndarray]: fields = self.fields_from_monitor(simulation, monitor) return {key: onp.asarray(value) for key, value in self.report(fields).items()} + def incident_fields(self, mode: Mode, distance: float) -> TangentialFields: + """Tangential fields at the monitor plane for a mode arriving from afar. + + The time reverse of `spectrum`. A `Mode` is specified where it is + physically meaningful -- at the fiber facet, say, on the far side of the + stack -- and this carries it back down through the layers to the monitor + plane, where Meep can inject it. The scalar transfer function of a + stratified medium is the same in both directions, so the same product of + transmission and propagation phases applies. + + The result is purely in-going: its up-going content is zero by + construction, so injecting it launches a beam toward the structure and + not away from it. + + Args: + mode: the target field, e.g. from `gaussian_mode`. + distance: how far the mode's plane is from the monitor, the same + convention `spectrum` uses. + + Returns: + A `TangentialFields` on the monitor's own sample grid. + """ + remaining = distance - self.stack_thickness + concrete = _as_concrete(remaining) + if concrete is not None and concrete < 0: + raise ValueError( + f"distance={distance} does not clear the stack, which is " + f"{_as_concrete(self.stack_thickness)} thick." + ) + + k0 = 2 * onp.pi * jnp.asarray(self.frequencies)[:, None] + amplitudes = mode.spectrum(self, self._kt, k0, self._indices[-1]) + + # Carry each polarization across the stack to the monitor plane. + down = {} + for polarization in (S_POLARIZATION, P_POLARIZATION): + transmission, _ = self._transmission(polarization) + to_interface = jnp.exp(1j * self._wavevectors[0] * self._thicknesses[0]) + beyond = jnp.exp(1j * self._wavevectors[-1] * remaining) + down[polarization] = ( + amplitudes[..., polarization] * to_interface * transmission * beyond + ) + + # Invert `decompose` for a spectrum with no up-going part. Setting + # up = (E - cross)/2 to zero gives E = cross = down, and the admittance + # relations then fix the magnetic components. + admittance_s, admittance_p = self._admittances[0] + outgoing = self.sign + electric_s = down[S_POLARIZATION] + electric_p = down[P_POLARIZATION] + + azimuth = self._azimuth() + if azimuth is None: + # cross_s = -outgoing * H_p / Y_s and cross_p = outgoing * H_s / Y_p + magnetic_p = -outgoing * admittance_s * electric_s + magnetic_s = outgoing * admittance_p * electric_p + electric_u, electric_v = electric_s, electric_p + magnetic_u, magnetic_v = magnetic_s, magnetic_p + else: + # cross_s = outgoing * H_p / Y_s and cross_p = -outgoing * H_s / Y_p + magnetic_p = outgoing * admittance_s * electric_s + magnetic_s = -outgoing * admittance_p * electric_p + cosine, sine = azimuth + # the rotation into (s, p) is a reflection, hence its own inverse + electric_u = -electric_s * sine + electric_p * cosine + electric_v = electric_s * cosine + electric_p * sine + magnetic_u = -magnetic_s * sine + magnetic_p * cosine + magnetic_v = magnetic_s * cosine + magnetic_p * sine + + (e_u, e_v), (h_u, h_v) = _PLANE_AXES[self.normal] + return TangentialFields( + E={ + e_u: self._inverse_transform(electric_u), + e_v: self._inverse_transform(electric_v), + }, + H={ + # `_tangential_spectra` pairs E_u with H_v, so undo that here + h_v: self._inverse_transform(magnetic_v), + h_u: self._inverse_transform(magnetic_u), + }, + normal=self.normal, + sign=self.sign, + ) + + def equivalent_sources( + self, + fields: TangentialFields, + time_src, + center, + size, + frequency: Optional[float] = None, + amplitude: complex = 1.0, + differentiable: bool = True, + ) -> List[mp.Source]: + """Equivalent surface currents that launch `fields` into the simulation. + + A thin wrapper over `mp.get_equiv_sources`, which already applies the + equivalence principle -- an electric sheet `K = n_hat x H` and a + magnetic sheet `N = -n_hat x E` -- and emits ordinary sources carrying + `amp_data`. Both sheets are needed: one alone radiates half the field in + each direction instead of all of it in one. + + Since `amp_data` is differentiable, the sources it returns backpropagate + to whatever produced `fields`, which for the angular spectrum means the + mode's waist, tilt and offset. + + Args: + fields: tangential fields on the plane, e.g. from `incident_fields`. + time_src: the time profile, e.g. `mp.GaussianSource`. + center, size: the source plane, normally the same as the monitor. + frequency: which frequency of `fields` to inject. Defaults to the + only one when the propagator carries a single frequency. + amplitude: an overall scale. + differentiable: flag the sources for the adjoint solver. + + Returns: + A list of `mp.Source`, one per non-zero current component. + """ + frequencies = onp.asarray(self.frequencies) + if frequency is None: + if frequencies.size != 1: + raise ValueError( + f"The propagator carries {frequencies.size} frequencies, so " + "`frequency` has to say which one to inject; a single Meep " + "source has one time profile and cannot carry independent " + "amplitudes at several." + ) + index = 0 + else: + index = int(onp.argmin(onp.abs(frequencies - frequency))) + + def slice_at(store, component): + value = store.get(component) + if value is None: + return onp.zeros(self.num_points) + return onp.asarray(value)[index] * amplitude + + # get_equiv_sources wants all six components, in order, shaped for + # amp_data's trilinear interpolation. + def shaped(values): + values = onp.asarray(values, dtype=onp.complex128) + spatial = tuple(values.shape) if values.ndim else (1,) + return values.reshape(spatial + (1,) * (3 - len(spatial))) + + field = [shaped(slice_at(fields.E, c)) for c in (mp.Ex, mp.Ey, mp.Ez)] + [ + shaped(slice_at(fields.H, c)) for c in (mp.Hx, mp.Hy, mp.Hz) + ] + + normal_axis = {mp.X: 0, mp.Y: 1, mp.Z: 2}[self.normal] + n_hat = onp.zeros(3) + # `sign` is the outgoing direction, so the beam travels along its + # negation and that is the normal the equivalence principle wants. + n_hat[normal_axis] = -self.sign + + sources = mp.get_equiv_sources(field, n_hat, time_src, center, size) + if differentiable: + for source in sources: + source.differentiable = ("amp_data",) + source.name = f"asm_{mp.component_name(source.component)}" + return sources + def objective_arguments(self, simulation, volume, **kwargs): """The `FourierFields` an objective function needs, for the adjoint path. diff --git a/python/examples/adjoint_optimization/grating_coupler_asm_source.py b/python/examples/adjoint_optimization/grating_coupler_asm_source.py new file mode 100644 index 000000000..5a7658edf --- /dev/null +++ b/python/examples/adjoint_optimization/grating_coupler_asm_source.py @@ -0,0 +1,284 @@ +"""Grating coupler driven *from the fiber*, with the fiber outside the cell. + +The companion example, `grating_coupler_asm.py`, runs the coupler in the +transmit direction: light comes in along the waveguide, radiates up, and the +angular-spectrum propagator carries the radiated field hundreds of microns to a +fiber. This one runs it in the receive direction. A fiber mode is specified +where it is physically meaningful -- at the facet, far above the chip and +outside the FDTD cell -- carried *down* through the silica and across the +silica/air interface analytically, and injected at a plane just above the +device layer as equivalent surface currents. + +Nothing about the cell changes: the superstrate and the working distance stay +out of the simulation in both directions. What changes is which end of the +chain is the source. + +The whole path is differentiable. The fiber's waist, tilt and lateral offset +are ordinary JAX values that reach Meep only through the injected currents, and +Meep returns the cotangent with respect to those currents; JAX carries it back +the rest of the way. So the grating and the fiber alignment can be optimized +together, which is the thing the transmit-direction example cannot do. + + # what the fiber launches, and where it is pointing + python grating_coupler_asm_source.py forward + + # co-optimize the grating and the fiber alignment + python grating_coupler_asm_source.py optimize --iterations 20 + +The geometry helpers are imported from the transmit-direction example rather +than duplicated; the two describe the same device. +""" + +import argparse +import numpy as np + +import meep as mp +import meep.adjoint as mpa + +import grating_coupler_asm as tx + +WAVEGUIDE_MODE_STANDOFF = 0.6 # where the waveguide mode is measured, as a +# fraction of the waveguide run-in + + +def build_arguments(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("mode", choices=("forward", "optimize")) + parser.add_argument("--resolution", type=int, default=30) + parser.add_argument("--aperture", type=float, default=12.0) + parser.add_argument("--wavelengths", type=float, nargs="+", default=[1.55]) + parser.add_argument( + "--superstrate", + type=float, + default=300.0, + help="silica thickness above the device layer, in microns", + ) + parser.add_argument( + "--working-distance", + type=float, + default=50.0, + help="air gap between the silica surface and the fiber facet", + ) + parser.add_argument("--waist", type=float, default=5.2, help="fiber MFD / 2") + parser.add_argument("--tilt", type=float, default=8.0, help="fiber tilt, degrees") + parser.add_argument("--offset", type=float, default=0.0) + parser.add_argument("--iterations", type=int, default=20) + parser.add_argument("--learning-rate", type=float, default=0.02) + # consumed by the geometry helpers imported from the transmit example + parser.add_argument("--design-resolution", type=float, default=None) + parser.add_argument("--pad-factor", type=int, default=16) + parser.add_argument("--fiber-waist", type=float, default=None) + return parser.parse_args() + + +def source_plane(args): + """The plane the fiber's field is injected on. + + The same plane the transmit example reads the radiated field from: inside + the homogeneous oxide, clear of the near field, and wide enough that the + beam has decayed by its ends. + """ + return tx.monitor_volume(args) + + +def injection_propagator(simulation, args, frequency): + """A propagator for the source plane, facing up toward the fiber. + + `from_volume` rather than `from_monitor`: a DFT monitor needs the field + components to be allocated, which does not happen until sources are added, + and the sources are exactly what this is being built to produce. + """ + stack, _ = tx.build_stack(args) + return mpa.AngularSpectrum.from_volume( + simulation, + simulation._fit_volume_to_simulation(source_plane(args)), + stack, + [frequency], + sign=+1, # the fiber is in the +y direction + pad_factor=4, + ) + + +def fiber_distance(args): + """From the source plane to the fiber facet.""" + return (args.superstrate - tx.MONITOR_STANDOFF) + args.working_distance + + +def waveguide_monitor(args): + """Where the coupled power is measured, in the input waveguide.""" + return mp.Volume( + center=mp.Vector3( + -(args.aperture / 2 + tx.WAVEGUIDE_LENGTH * WAVEGUIDE_MODE_STANDOFF), + tx.T_DEVICE / 2, + ), + size=mp.Vector3(0, 6 * tx.T_DEVICE), + ) + + +def build_simulation(args, weights=None): + """The coupler with no sources; the fiber's currents are added separately.""" + simulation, grids, volumes, nx, frequencies = tx.build_simulation(args, weights) + # The transmit example drives the waveguide. Here the fiber drives the + # grating instead, so that source goes away. + simulation.sources = [] + return simulation, grids, volumes, nx, frequencies + + +def fiber_sources(simulation, args, frequency, mode=None): + """Equivalent surface currents for the fiber mode, at the source plane.""" + propagator = injection_propagator(simulation, args, frequency) + if mode is None: + mode = mpa.gaussian_mode( + waist=args.waist, tilt_deg=args.tilt, offset=args.offset + ) + fields = propagator.incident_fields(mode, distance=fiber_distance(args)) + plane = source_plane(args) + return ( + propagator, + fields, + propagator.equivalent_sources( + fields, + mp.GaussianSource(frequency, fwidth=0.2 * frequency), + center=plane.center, + size=plane.size, + ), + ) + + +def run_forward(args): + """Inject the fiber mode and report how much reaches the waveguide.""" + simulation, _, _, _, frequencies = build_simulation( + args, weights=tx.initial_weights(args, tx.design_regions(args)[2]) + ) + frequency = float(np.mean(frequencies)) + propagator, fields, sources = fiber_sources(simulation, args, frequency) + + report = propagator.report(fields) + print("\nthe field the fiber puts on the source plane") + print( + f" travelling toward the chip {float(np.ravel(report['downgoing_fraction'])[0]):.4f}" + ) + print( + f" evanescent {float(np.ravel(report['evanescent_fraction'])[0]):.3e}" + ) + print( + f" amplitude at the plane edges {float(np.ravel(report['edge_amplitude'])[0]):.3e}" + ) + if float(np.ravel(report["edge_amplitude"])[0]) > tx.EDGE_AMPLITUDE_LIMIT: + print( + " ^ the beam has not decayed by the ends of the plane, so the " + "injected field wraps around. Widen the cell or narrow the beam." + ) + + simulation.change_sources(sources) + monitor = simulation.add_mode_monitor( + [frequency], + mp.ModeRegion(volume=waveguide_monitor(args)), + ) + simulation.run(until_after_sources=mp.stop_when_dft_decayed(1e-9)) + + coefficients = simulation.get_eigenmode_coefficients( + monitor, [1], eig_parity=mp.ODD_Z + ) + # the backward-going coefficient: the waveguide runs off to the left + coupled = abs(coefficients.alpha[0, 0, 1]) ** 2 + print(f"\npower coupled into the waveguide mode: {coupled:.6e}") + return coupled + + +def run_optimize(args): + """Co-optimize the grating and the fiber alignment.""" + import jax + + jax.config.update("jax_enable_x64", True) + import jax.numpy as jnp + import optax + + simulation, grids, volumes, nx, frequencies = build_simulation(args) + frequency = float(np.mean(frequencies)) + plane = source_plane(args) + propagator = injection_propagator(simulation, args, frequency) + distance = fiber_distance(args) + + design_regions = [ + mpa.DesignRegion(grid, volume=volume) for grid, volume in zip(grids, volumes) + ] + monitor = mpa.EigenmodeCoefficient( + simulation, waveguide_monitor(args), mode=1, forward=False + ) + + # A placeholder source, replaced on every iteration by the currents JAX + # computes. Its amplitudes are what the wrapper differentiates with respect + # to; the fiber parameters live upstream and Meep never sees them. + _, _, sources = fiber_sources(simulation, args, frequency) + + wrapper = mpa.MeepJaxWrapper( + simulation, + sources, + [monitor], + design_regions, + [frequency], + minimum_run_time=200.0, + ) + + def currents(fiber): + """The amp_data array for each equivalent-current sheet, in JAX.""" + mode = mpa.gaussian_mode( + waist=fiber["waist"], tilt_deg=fiber["tilt"], offset=fiber["offset"] + ) + fields = propagator.incident_fields(mode, distance=distance) + # get_equiv_sources puts n x H on the electric components and -n x E on + # the magnetic ones, so the two are swapped relative to each other. + swap = { + mp.Ex: mp.Hx, + mp.Ey: mp.Hy, + mp.Ez: mp.Hz, + mp.Hx: mp.Ex, + mp.Hy: mp.Ey, + mp.Hz: mp.Ez, + } + out = [] + for s in sources: + paired = swap[s.component] + store = fields.H if mp.is_electric(s.component) else fields.E + out.append(jnp.ravel(store[paired][0])) + return out + + def loss(params): + values = wrapper( + [jnp.asarray(w) for w in params["weights"]], currents(params["fiber"]) + ) + return -jnp.sum(jnp.abs(jnp.asarray(values)) ** 2) + + params = { + "weights": [jnp.asarray(w) for w in tx.initial_weights(args, nx)], + "fiber": { + "waist": jnp.asarray(args.waist), + "tilt": jnp.asarray(args.tilt), + "offset": jnp.asarray(args.offset), + }, + } + optimizer = optax.adam(args.learning_rate) + state = optimizer.init(params) + + for iteration in range(args.iterations): + value, gradient = jax.value_and_grad(loss)(params) + updates, state = optimizer.update(gradient, state, params) + params = optax.apply_updates(params, updates) + params["weights"] = [jnp.clip(w, 0.0, 1.0) for w in params["weights"]] + fiber = params["fiber"] + print( + f"{iteration:3d} coupling {-float(value):.6e} " + f"waist {float(fiber['waist']):.3f} " + f"tilt {float(fiber['tilt']):.3f} " + f"offset {float(fiber['offset']):+.3f}" + ) + return params + + +if __name__ == "__main__": + arguments = build_arguments() + if arguments.mode == "forward": + run_forward(arguments) + else: + run_optimize(arguments) diff --git a/python/tests/test_angular_spectrum.py b/python/tests/test_angular_spectrum.py index 4174c6317..f0f224748 100644 --- a/python/tests/test_angular_spectrum.py +++ b/python/tests/test_angular_spectrum.py @@ -832,3 +832,101 @@ def test_two_dimensions(self): def test_three_dimensions(self): self._check(3) + + +class TestInjection(ApproxComparisonTestCase): + """Launching a mode into a simulation through the stack. + + The time reverse of the measurement path: `incident_fields` carries a mode + back down to the monitor plane and `equivalent_sources` turns it into + currents Meep can apply. + """ + + FCEN = 1 / 1.55 + RESOLUTION = 20 + WAIST = 2.0 + DISTANCE = 8.0 + CELL = mp.Vector3(16, 14) + LINE = mp.Vector3(10, 0) + PML = 2.0 + + def _stack(self): + # one semi-infinite layer, so the check isolates the injection physics + return mpa.Stack([mpa.Layer(index=1.0, thickness=0.0), mpa.Layer(index=1.0)]) + + def _propagator(self, y=0.0, sign=1): + simulation = mp.Simulation( + cell_size=self.CELL, + resolution=self.RESOLUTION, + boundary_layers=[mp.PML(self.PML)], + force_complex_fields=True, + ) + simulation.init_sim() + volume = simulation._fit_volume_to_simulation( + mp.Volume(center=mp.Vector3(0, y), size=self.LINE) + ) + propagator = mpa.AngularSpectrum.from_volume( + simulation, volume, self._stack(), [self.FCEN], sign=sign, pad_factor=4 + ) + return propagator, volume + + def test_incident_fields_are_purely_ingoing(self): + # `incident_fields` inverts `decompose` for a spectrum with no outgoing + # part, so feeding the result back through `decompose` must recover + # that: all of the power heading toward the structure, none away. + propagator, _ = self._propagator() + fields = propagator.incident_fields( + mpa.gaussian_mode(waist=self.WAIST), distance=self.DISTANCE + ) + report = propagator.report(fields) + self.assertGreater(float(onp.ravel(report["downgoing_fraction"])[0]), 0.999) + + def test_equivalent_sources_radiate_one_way(self): + # Both sheets of the equivalent-current pair are needed, with the right + # relative sign. With one sheet, or with the sign of either flipped, the + # beam appears in the other half-space instead -- at full strength, not + # as a small error. + propagator, _ = self._propagator() + fields = propagator.incident_fields( + mpa.gaussian_mode(waist=self.WAIST), distance=self.DISTANCE + ) + sources = propagator.equivalent_sources( + fields, + mp.GaussianSource(self.FCEN, fwidth=0.1 * self.FCEN), + center=mp.Vector3(0, 0), + size=self.LINE, + ) + self.assertTrue(sources) + + simulation = mp.Simulation( + cell_size=self.CELL, + resolution=self.RESOLUTION, + boundary_layers=[mp.PML(self.PML)], + sources=sources, + force_complex_fields=True, + ) + # sign=+1 means the outgoing direction is +y, so the mode arrives from + # +y and the beam should end up at negative y. + intended = mp.Volume(center=mp.Vector3(0, -3.0), size=self.LINE) + leaked = mp.Volume(center=mp.Vector3(0, 3.0), size=self.LINE) + toward = simulation.add_dft_fields( + [mp.Ez, mp.Hx], [self.FCEN], where=intended, yee_grid=False + ) + away = simulation.add_dft_fields( + [mp.Ez, mp.Hx], [self.FCEN], where=leaked, yee_grid=False + ) + simulation.run(until_after_sources=mp.stop_when_dft_decayed(1e-9)) + + strong = onp.max( + onp.abs(onp.asarray(simulation.get_dft_array(toward, mp.Ez, 0))) + ) + weak = onp.max(onp.abs(onp.asarray(simulation.get_dft_array(away, mp.Ez, 0)))) + self.assertLess(weak / strong, 0.02) + + # and the beam that arrives is still travelling the way it was sent + downward = mpa.AngularSpectrum.from_monitor( + simulation, toward, self._stack(), intended, sign=-1, pad_factor=4 + ) + measured = downward.fields_from_monitor(simulation, toward) + report = downward.report(measured) + self.assertLess(float(onp.ravel(report["downgoing_fraction"])[0]), 0.01)