From 8549b84b9ddb06b91c9a5317a6ec035e95341f30 Mon Sep 17 00:00:00 2001 From: Alec Hammond Date: Thu, 27 Aug 2026 16:41:04 -0700 Subject: [PATCH 1/4] adjoint: launch a mode through the stack, not just measure one incident_fields is the time reverse of spectrum: a Mode is specified where it is physically meaningful -- at the fiber facet, on the far side of the stack -- and carried back down 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 by construction, since it inverts decompose for a spectrum with no outgoing part. equivalent_sources then hands those fields to mp.get_equiv_sources, which already applies the equivalence principle and emits sources carrying amp_data. Because amp_data is differentiable, a beam injected this way backpropagates to its own waist, tilt and offset. from_volume builds a propagator without a monitor, which the injection path needs: a DFT monitor requires the field components to be allocated, and that does not happen until sources are added -- which is what is being built. Verified end to end: the injected field is 99.997% in-going, the radiated beam leaks 4e-3 in amplitude the wrong way (1.6e-5 in power), and the beam measured downstream is still 99.99% travelling the direction it was sent. --- python/adjoint/angular_spectrum.py | 199 +++++++++++- .../grating_coupler_asm_source.py | 284 ++++++++++++++++++ python/tests/test_angular_spectrum.py | 98 ++++++ 3 files changed, 580 insertions(+), 1 deletion(-) create mode 100644 python/examples/adjoint_optimization/grating_coupler_asm_source.py diff --git a/python/adjoint/angular_spectrum.py b/python/adjoint/angular_spectrum.py index 0fb562e22..e2d81dc4a 100644 --- a/python/adjoint/angular_spectrum.py +++ b/python/adjoint/angular_spectrum.py @@ -37,7 +37,7 @@ import math import warnings -from typing import Callable, Dict, NamedTuple, Optional, Sequence, Tuple, Union +from typing import Callable, Dict, List, NamedTuple, Optional, Sequence, Tuple, Union import jax import jax.numpy as jnp @@ -1163,6 +1163,42 @@ def _to_plane_axes(self, values, component): ) return values + @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: @@ -1224,6 +1260,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 848b5ff38..b7962e75d 100644 --- a/python/tests/test_angular_spectrum.py +++ b/python/tests/test_angular_spectrum.py @@ -1125,5 +1125,103 @@ def test_matches_the_permuted_z_normal_plane(self): ) +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) + + if __name__ == "__main__": unittest.main() From 3ea693de54539262fb6b002e1e1cb78a6e273fdf Mon Sep 17 00:00:00 2001 From: Alec Hammond Date: Fri, 9 Oct 2026 11:29:49 -0700 Subject: [PATCH 2/4] Fix angular-spectrum injection for x- and y-normal planes incident_fields assumed the 2D s axis is the first tangential one, which only holds for y-normal lines, and equivalent_sources laid the (u, v) arrays out as amp_data's (x, y, z) without reordering them or putting the singleton along the normal, which only holds for 2D y-normal lines and 3D z-normal planes. Add tests for an x-normal line and for the 3D layout. Co-Authored-By: Claude Opus 5.5 --- python/adjoint/angular_spectrum.py | 49 +++++++++----- python/tests/test_angular_spectrum.py | 96 +++++++++++++++++++++++++++ 2 files changed, 130 insertions(+), 15 deletions(-) diff --git a/python/adjoint/angular_spectrum.py b/python/adjoint/angular_spectrum.py index e2d81dc4a..ff8f18acd 100644 --- a/python/adjoint/angular_spectrum.py +++ b/python/adjoint/angular_spectrum.py @@ -1199,6 +1199,26 @@ def from_volume( **kwargs, ) + def _to_meep_layout(self, values) -> onp.ndarray: + """The inverse of `_to_plane_axes`, laid out as `amp_data` expects. + + `amp_data` is indexed (x, y, z) over the source volume, with length one + along any direction in which the volume has zero size, here the normal. + """ + values = onp.asarray(values, dtype=onp.complex128) + if self.transverse_dimensions == 2 and self.normal == mp.Y: + values = values.T + normal_axis = {mp.X: 0, mp.Y: 1, mp.Z: 2}[self.normal] + tangential = [ + axis + for axis in range(self.transverse_dimensions + 1) + if axis != normal_axis + ] + shape = [1, 1, 1] + for axis, n in zip(tangential, values.shape): + shape[axis] = n + return values.reshape(shape) + def fields_from_monitor( self, simulation: mp.Simulation, monitor ) -> TangentialFields: @@ -1313,11 +1333,15 @@ def incident_fields(self, mode: Mode, distance: float) -> TangentialFields: 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 + # the inverse of the 2D branch of `decompose`, for either s axis + if self._s_axis_2d == 0: + electric_u, electric_v = electric_s, electric_p + magnetic_v = -outgoing * admittance_s * electric_s + magnetic_u = outgoing * admittance_p * electric_p + else: + electric_u, electric_v = electric_p, electric_s + magnetic_u = outgoing * admittance_s * electric_s + magnetic_v = -outgoing * admittance_p * electric_p else: # cross_s = outgoing * H_p / Y_s and cross_p = -outgoing * H_s / Y_p magnetic_p = outgoing * admittance_s * electric_s @@ -1397,16 +1421,11 @@ def slice_at(store, component): 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) - ] + # get_equiv_sources wants all six components, in order, laid out for + # amp_data's trilinear interpolation over the source volume. + field = [ + self._to_meep_layout(slice_at(fields.E, c)) for c in (mp.Ex, mp.Ey, mp.Ez) + ] + [self._to_meep_layout(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) diff --git a/python/tests/test_angular_spectrum.py b/python/tests/test_angular_spectrum.py index b7962e75d..92f287478 100644 --- a/python/tests/test_angular_spectrum.py +++ b/python/tests/test_angular_spectrum.py @@ -1223,5 +1223,101 @@ def test_equivalent_sources_radiate_one_way(self): self.assertLess(float(onp.ravel(report["downgoing_fraction"])[0]), 0.01) +@unittest.skipIf(jax is None, "jax is not installed") +class TestInjectionOrientations(ApproxComparisonTestCase): + """Injection on planes other than the y-normal line `TestInjection` uses. + + In 2D an x-normal line carries s polarization on its second tangential + axis, and in 3D `amp_data` wants Meep's (x, y, z) layout, which differs from + the plane's (u, v) axes for every normal except z. + """ + + FCEN = 1 / 1.55 + + def _stack(self): + return mpa.Stack([mpa.Layer(index=1.0)]) + + def test_x_normal_line_radiates_one_way(self): + cell, line = mp.Vector3(14, 16), mp.Vector3(0, 10) + simulation = mp.Simulation( + cell_size=cell, + resolution=20, + boundary_layers=[mp.PML(2.0)], + force_complex_fields=True, + ) + simulation.init_sim() + volume = simulation._fit_volume_to_simulation( + mp.Volume(center=mp.Vector3(), size=line) + ) + propagator = mpa.AngularSpectrum.from_volume( + simulation, volume, self._stack(), [self.FCEN], pad_factor=4 + ) + self.assertEqual(propagator.normal, mp.X) + fields = propagator.incident_fields(mpa.gaussian_mode(waist=2.0), distance=8.0) + self.assertIn(mp.Ez, fields.E) + report = propagator.report(fields) + self.assertGreater(float(onp.ravel(report["downgoing_fraction"])[0]), 0.999) + + sources = propagator.equivalent_sources( + fields, + mp.GaussianSource(self.FCEN, fwidth=0.1 * self.FCEN), + center=mp.Vector3(), + size=line, + ) + simulation = mp.Simulation( + cell_size=cell, + resolution=20, + boundary_layers=[mp.PML(2.0)], + sources=sources, + force_complex_fields=True, + ) + # sign=+1: outgoing is +x, so the mode arrives from +x and travels to -x + intended = mp.Volume(center=mp.Vector3(-3.0, 0), size=line) + leaked = mp.Volume(center=mp.Vector3(3.0, 0), size=line) + toward = simulation.add_dft_fields( + [mp.Ez, mp.Hy], [self.FCEN], where=intended, yee_grid=False + ) + away = simulation.add_dft_fields( + [mp.Ez, mp.Hy], [self.FCEN], where=leaked, yee_grid=False + ) + simulation.run(until_after_sources=mp.stop_when_dft_decayed(1e-9)) + strong = onp.abs(simulation.get_dft_array(toward, mp.Ez, 0)).max() + weak = onp.abs(simulation.get_dft_array(away, mp.Ez, 0)).max() + self.assertLess(weak / strong, 0.02) + # and the beam is centered, i.e. laid out along y, not squeezed into x + profile = onp.abs(simulation.get_dft_array(toward, mp.Ez, 0)) + coordinates = onp.asarray(simulation.get_array_metadata(vol=intended)[1]) + centroid = onp.sum(coordinates * profile**2) / onp.sum(profile**2) + self.assertLess(abs(centroid), 0.1) + + def test_amp_data_layout_in_three_dimensions(self): + simulation = mp.Simulation(cell_size=mp.Vector3(3, 3, 3), resolution=8) + simulation.init_sim() + sizes = { + mp.X: mp.Vector3(0, 1.6, 1.0), + mp.Y: mp.Vector3(1.6, 0, 1.0), + mp.Z: mp.Vector3(1.6, 1.0, 0), + } + for normal, size in sizes.items(): + volume = simulation._fit_volume_to_simulation( + mp.Volume(center=mp.Vector3(), size=size) + ) + propagator = mpa.AngularSpectrum.from_volume( + simulation, volume, self._stack(), [self.FCEN] + ) + meep_shape = onp.shape( + simulation.get_array(vol=volume, component=mp.Dielectric) + ) + values = onp.arange(onp.prod(meep_shape)).reshape(meep_shape) + 0j + plane = propagator._to_plane_axes(values[None], mp.Ex)[0] + layout = propagator._to_meep_layout(plane) + normal_axis = {mp.X: 0, mp.Y: 1, mp.Z: 2}[normal] + self.assertEqual(layout.shape[normal_axis], 1, f"normal {normal}") + self.assertTrue( + onp.array_equal(layout.squeeze(axis=normal_axis), values), + f"normal {normal}", + ) + + if __name__ == "__main__": unittest.main() From 803271abd76d460199f8cfb204de1ded7d0c2825 Mon Sep 17 00:00:00 2001 From: Alec Hammond Date: Mon, 12 Oct 2026 12:23:12 -0700 Subject: [PATCH 3/4] Fix mode injection through a stack and the source example - incident_fields reused the monitor-to-far transmission, but the stack is crossed from the far side; by reciprocity that is t * Y_far / Y_monitor. Across an interface the injected power was wrong (2.1x for glass/air). - incident_fields used the mode's spectrum unchanged, which launches a tilted mode at the mirrored angle; inject its time reverse, conj(a(-kt)), so the same Mode describes the same fiber in both directions. - from_volume read the grid from an uninitialized simulation, which crashed; initialize it to read the grid and reset it afterwards. - The example rebuilt the equivalent currents with a component swap that picked the wrong (zero) components, so its optimization launched nothing and never stopped. Add equivalent_currents, the JAX version of what equivalent_sources places, and use it in both. Co-Authored-By: Claude Opus 5.5 --- python/adjoint/angular_spectrum.py | 120 ++++++++++++----- .../grating_coupler_asm_source.py | 19 +-- python/tests/test_angular_spectrum.py | 127 ++++++++++++++++++ 3 files changed, 214 insertions(+), 52 deletions(-) diff --git a/python/adjoint/angular_spectrum.py b/python/adjoint/angular_spectrum.py index ff8f18acd..2f3ee60f0 100644 --- a/python/adjoint/angular_spectrum.py +++ b/python/adjoint/angular_spectrum.py @@ -353,9 +353,11 @@ class Mode(NamedTuple): Attributes: spectrum: called as `spectrum(propagator, kt, k0, index)` and returning a (num frequencies, num kt, 2) array of tangential electric field - amplitudes in the s and p directions. Carrying both is what lets a - linearly polarized beam be represented in 3D, where its s and p - content varies with azimuth. + amplitudes in the s and p directions of the propagator's own + wavevectors (`propagator._azimuth()`), with the envelope evaluated at + `kt`. Carrying both is what lets a linearly polarized beam be + represented in 3D, where its s and p content varies with azimuth. + `incident_fields` evaluates it at `-kt` to time-reverse it. """ spectrum: Callable @@ -1181,14 +1183,23 @@ def from_volume( produce. This takes the frequencies directly instead. Args: - simulation: supplies the grid; it need not have been run. + simulation: supplies the grid. It need not have been initialized; + if it has not, it is initialized to read the grid and then reset. 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) + if simulation.fields is None: + # The plane's samples come from the grid, which needs initialized + # fields; drop them again so that the sources this builds are added + # normally when the simulation runs. + simulation.init_sim() + normal, pitch, num_points = cls._plane_geometry(simulation, volume) + simulation.reset_meep() + else: + normal, pitch, num_points = cls._plane_geometry(simulation, volume) return cls( stack, onp.asarray(frequencies), @@ -1199,13 +1210,46 @@ def from_volume( **kwargs, ) - def _to_meep_layout(self, values) -> onp.ndarray: + def equivalent_currents( + self, fields: TangentialFields, frequency_index: int = 0, amplitude=1.0 + ) -> Dict[int, jnp.ndarray]: + """The equivalent surface currents that launch `fields`, as `amp_data`. + + The electric sheet `K = n_hat x H` and magnetic sheet `N = -n_hat x E`, + keyed by the Meep component each drives and laid out as `amp_data` + expects. These are the arrays `equivalent_sources` places, computed in + JAX, so that a loss can pass them to `MeepJaxWrapper` and differentiate + through them. + """ + + def at(store, component): + value = store.get(component) + if value is None: + return self._to_meep_layout(jnp.zeros(self.num_points)) + return self._to_meep_layout(jnp.asarray(value)[frequency_index] * amplitude) + + ex, ey, ez = (at(fields.E, c) for c in (mp.Ex, mp.Ey, mp.Ez)) + hx, hy, hz = (at(fields.H, c) for c in (mp.Hx, mp.Hy, mp.Hz)) + # `sign` is the outgoing direction, so the beam travels along its + # negation and that is the normal the equivalence principle wants. + n = [0.0, 0.0, 0.0] + n[{mp.X: 0, mp.Y: 1, mp.Z: 2}[self.normal]] = -float(self.sign) + return { + mp.Ex: n[1] * hz - n[2] * hy, + mp.Ey: n[2] * hx - n[0] * hz, + mp.Ez: n[0] * hy - n[1] * hx, + mp.Hx: n[2] * ey - n[1] * ez, + mp.Hy: n[0] * ez - n[2] * ex, + mp.Hz: n[1] * ex - n[0] * ey, + } + + def _to_meep_layout(self, values) -> jnp.ndarray: """The inverse of `_to_plane_axes`, laid out as `amp_data` expects. `amp_data` is indexed (x, y, z) over the source volume, with length one along any direction in which the volume has zero size, here the normal. """ - values = onp.asarray(values, dtype=onp.complex128) + values = jnp.asarray(values, dtype=jnp.complex128) if self.transverse_dimensions == 2 and self.normal == mp.Y: values = values.T normal_axis = {mp.X: 0, mp.Y: 1, mp.Z: 2}[self.normal] @@ -1286,9 +1330,12 @@ def incident_fields(self, mode: Mode, distance: float) -> TangentialFields: 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. + plane, where Meep can inject it. The same `Mode` describes the same fiber + in both directions: the beam injected here is the time reverse of the one + `overlap` projects onto, i.e. its spectrum is `conj(a(-kt))`, which for + `gaussian_mode` is the same beam with the tilt negated. It is carried + across the stack from the far side, whose tangential-field transmission + follows from reciprocity as `t * Y_far / Y_monitor`. The result is purely in-going: its up-going content is zero by construction, so injecting it launches a beam toward the structure and @@ -1311,12 +1358,20 @@ def incident_fields(self, mode: Mode, distance: float) -> TangentialFields: ) k0 = 2 * onp.pi * jnp.asarray(self.frequencies)[:, None] - amplitudes = mode.spectrum(self, self._kt, k0, self._indices[-1]) + # time reversal: the arriving beam is conj(a(-kt)) of the outgoing one + amplitudes = jnp.conj(mode.spectrum(self, -self._kt, k0, self._indices[-1])) - # Carry each polarization across the stack to the monitor plane. + # Carry each polarization across the stack to the monitor plane. The + # stack is crossed from the far side, so reciprocity turns the + # monitor-to-far transmission t into t * Y_far / Y_monitor. down = {} for polarization in (S_POLARIZATION, P_POLARIZATION): transmission, _ = self._transmission(polarization) + transmission = ( + transmission + * self._admittances[-1][polarization] + / self._admittances[0][polarization] + ) to_interface = jnp.exp(1j * self._wavevectors[0] * self._thicknesses[0]) beyond = jnp.exp(1j * self._wavevectors[-1] * remaining) down[polarization] = ( @@ -1380,11 +1435,11 @@ def equivalent_sources( ) -> 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. + Applies the equivalence principle -- an electric sheet `K = n_hat x H` + and a magnetic sheet `N = -n_hat x E`, as `equivalent_currents` computes + them -- and emits ordinary sources carrying them as `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 @@ -1415,25 +1470,18 @@ def equivalent_sources( 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, laid out for - # amp_data's trilinear interpolation over the source volume. - field = [ - self._to_meep_layout(slice_at(fields.E, c)) for c in (mp.Ex, mp.Ey, mp.Ez) - ] + [self._to_meep_layout(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) + currents = self.equivalent_currents(fields, index, amplitude) + sources = [ + mp.Source( + time_src, + component, + center=center, + size=size, + amp_data=onp.asarray(values, dtype=onp.complex128), + ) + for component, values in currents.items() + if onp.any(onp.asarray(values)) + ] if differentiable: for source in sources: source.differentiable = ("amp_data",) diff --git a/python/examples/adjoint_optimization/grating_coupler_asm_source.py b/python/examples/adjoint_optimization/grating_coupler_asm_source.py index 5a7658edf..cae296a28 100644 --- a/python/examples/adjoint_optimization/grating_coupler_asm_source.py +++ b/python/examples/adjoint_optimization/grating_coupler_asm_source.py @@ -227,22 +227,9 @@ def currents(fiber): 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 + # the same currents `equivalent_sources` placed, computed in JAX + sheet = propagator.equivalent_currents(fields) + return [jnp.ravel(sheet[s.component]) for s in sources] def loss(params): values = wrapper( diff --git a/python/tests/test_angular_spectrum.py b/python/tests/test_angular_spectrum.py index 92f287478..f12aa5d39 100644 --- a/python/tests/test_angular_spectrum.py +++ b/python/tests/test_angular_spectrum.py @@ -1319,5 +1319,132 @@ def test_amp_data_layout_in_three_dimensions(self): ) +@unittest.skipIf(jax is None, "jax is not installed") +class TestInjectionThroughAStack(ApproxComparisonTestCase): + """`incident_fields` across an index change and for a tilted mode. + + The mode is specified in air above a glass layer, the usual fiber setup. The + injected beam must carry the mode's power times the (reciprocal) power + transmission into the glass, travel the way the fiber would emit it, and be + the exact time reverse of what `overlap` projects onto. + """ + + TILT = 8.0 + + def setUp(self): + self.propagator = mpa.AngularSpectrum( + mpa.Stack([mpa.Layer(N_OXIDE, 1.0), mpa.Layer(N_AIR)]), + [1 / WAVELENGTH], + 0.05, + 1024, + normal=mp.Y, + pad_factor=4, + ) + self.mode = mpa.gaussian_mode(waist=3.0, tilt_deg=self.TILT) + self.fields = self.propagator.incident_fields(self.mode, distance=5.0) + _, down = self.propagator.decompose(self.fields) + self.down = onp.asarray(down[asm.S_POLARIZATION])[0] + + def test_power_is_the_transmitted_mode_power(self): + p = self.propagator + k0 = onp.array([2 * onp.pi / WAVELENGTH]) + a = onp.squeeze( + onp.asarray(self.mode.spectrum(p, -p.kt, k0[:, None], p._indices[-1]))[ + ..., asm.S_POLARIZATION + ] + ) + y_glass = onp.asarray(p._admittances[0][asm.S_POLARIZATION])[0] + y_air = onp.asarray(p._admittances[-1][asm.S_POLARIZATION])[0] + propagating = onp.abs(onp.imag(onp.asarray(p._wavevectors[-1])[0])) < 1e-6 + transmitted = 1 - onp.abs((y_air - y_glass) / (y_air + y_glass)) ** 2 + got = onp.sum( + onp.real(y_glass)[propagating] * onp.abs(self.down[propagating]) ** 2 + ) + want = onp.sum((transmitted * onp.real(y_air) * onp.abs(a) ** 2)[propagating]) + self.assertAlmostEqual(got / want, 1.0, places=8) + + def test_tilted_mode_is_launched_the_way_the_fiber_emits(self): + # overlap with gaussian_mode(tilt) collects light travelling up at +kt, + # so the same fiber emits light travelling down at -kt + kt = onp.asarray(self.propagator.kt)[:, 0] + self.assertLess(kt[onp.argmax(onp.abs(self.down))], 0) + + def test_time_reverse_projects_fully_onto_the_mode(self): + # Exactly, without an interface: through one, the angle-dependent + # reflection reshapes the returning beam and the purity drops below one. + propagator = mpa.AngularSpectrum( + mpa.Stack([mpa.Layer(N_OXIDE, 1.0), mpa.Layer(N_OXIDE)]), + [1 / WAVELENGTH], + 0.05, + 1024, + normal=mp.Y, + pad_factor=4, + ) + fields = propagator.incident_fields(self.mode, distance=5.0) + reversed_fields = mpa.TangentialFields( + E={c: jnp.conj(v) for c, v in fields.E.items()}, + H={c: -jnp.conj(v) for c, v in fields.H.items()}, + normal=fields.normal, + ) + purity = float(propagator.overlap(reversed_fields, self.mode, 5.0)[0]) + self.assertAlmostEqual(purity, 1.0, places=6) + + +@unittest.skipIf(jax is None, "jax is not installed") +class TestEquivalentCurrents(ApproxComparisonTestCase): + """`equivalent_currents` is what `equivalent_sources` places, in JAX.""" + + def _propagator(self): + simulation = mp.Simulation( + cell_size=mp.Vector3(10, 8), resolution=10, boundary_layers=[mp.PML(1.0)] + ) + volume = mp.Volume(center=mp.Vector3(0, -1), size=mp.Vector3(6, 0)) + propagator = mpa.AngularSpectrum.from_volume( + simulation, + simulation._fit_volume_to_simulation(volume), + mpa.Stack([mpa.Layer(N_OXIDE, 1.0), mpa.Layer(N_AIR)]), + [1 / WAVELENGTH], + pad_factor=4, + ) + return propagator, volume + + def test_matches_the_placed_sources(self): + propagator, volume = self._propagator() + fields = propagator.incident_fields( + mpa.gaussian_mode(waist=2.0, tilt_deg=8.0), distance=5.0 + ) + currents = propagator.equivalent_currents(fields) + sources = propagator.equivalent_sources( + fields, + mp.GaussianSource(1 / WAVELENGTH, fwidth=0.1), + center=volume.center, + size=volume.size, + ) + # in 2D with an out-of-plane E field, K drives Ez and N drives Hx + self.assertEqual(sorted(s.component for s in sources), sorted([mp.Ez, mp.Hx])) + for source in sources: + self.assertTrue( + onp.array_equal( + onp.asarray(source.amp_data), + onp.asarray(currents[source.component]), + ) + ) + + def test_differentiable_in_the_mode(self): + propagator, _ = self._propagator() + + def power(waist): + fields = propagator.incident_fields( + mpa.gaussian_mode(waist=waist, tilt_deg=8.0), distance=5.0 + ) + currents = propagator.equivalent_currents(fields) + return jnp.sum(jnp.abs(currents[mp.Ez]) ** 2) + + gradient = float(jax.grad(power)(2.0)) + step = 1e-5 + reference = float((power(2.0 + step) - power(2.0 - step)) / (2 * step)) + self.assertAlmostEqual(gradient / reference, 1.0, places=5) + + if __name__ == "__main__": unittest.main() From 5e04191740fee118813c1e02d61fe0e11687ff9d Mon Sep 17 00:00:00 2001 From: Alec Hammond Date: Mon, 12 Oct 2026 13:45:08 -0700 Subject: [PATCH 4/4] Add an end-to-end check of angular-spectrum injection A small 2D grating coupler driven by a fiber mode carried through a superstrate stack outside the cell. It checks the launched field against the propagator with no fitted scale, transmit-vs-receive reciprocity across fiber tilts (with its convergence in resolution), and the gradient with respect to the fiber's waist, tilt and offset against finite differences of the full simulation, and draws a figure of all three. Co-Authored-By: Claude Opus 5.5 --- .../grating_coupler_asm_end_to_end.py | 519 ++++++++++++++++++ 1 file changed, 519 insertions(+) create mode 100644 python/examples/adjoint_optimization/grating_coupler_asm_end_to_end.py diff --git a/python/examples/adjoint_optimization/grating_coupler_asm_end_to_end.py b/python/examples/adjoint_optimization/grating_coupler_asm_end_to_end.py new file mode 100644 index 000000000..d38844843 --- /dev/null +++ b/python/examples/adjoint_optimization/grating_coupler_asm_end_to_end.py @@ -0,0 +1,519 @@ +"""End-to-end check of angular-spectrum injection, in a small 2D grating coupler. + +A fiber mode is specified in air above a silica superstrate, outside the FDTD +cell, carried down through the stack by `AngularSpectrum.incident_fields`, and +injected as equivalent surface currents just above a silicon grating. Three +independent checks, each against something the injection path does not use: + + 1. The field Meep launches matches the propagator's prediction a few microns + below the injection plane, in a uniform cell, with no fitted parameters: + only the source's own spectrum, which Meep applies to every source. + 2. Reciprocity. Coupling fiber -> waveguide (inject the fiber mode, measure + the waveguide mode) equals coupling waveguide -> fiber (launch the + waveguide mode, project the radiated field onto the fiber mode through + the same stack), at every fiber tilt. The two agree to discretization + error, which shrinks at second order with resolution. + 3. The gradient of the coupled power with respect to the fiber's waist, tilt + and offset -- through `incident_fields`, `equivalent_currents`, + `MeepJaxWrapper` and Meep's source adjoint -- matches finite differences + of the full simulation. + + python grating_coupler_asm_end_to_end.py # ~20 minutes, one core + python grating_coupler_asm_end_to_end.py --quick # coarser, ~2 minutes + python grating_coupler_asm_end_to_end.py --replot # redraw from saved results +""" + +import argparse + +import jax +import jax.numpy as jnp +import numpy as np + +import meep as mp +import meep.adjoint as mpa +from meep.adjoint import angular_spectrum as asm +from meep.adjoint import source_gradient + +jax.config.update("jax_enable_x64", True) +mp.verbosity(0) + +WAVELENGTH = 1.55 +FCEN = 1 / WAVELENGTH +N_OXIDE, N_AIR, N_SILICON = 1.444, 1.0, 3.48 + +CELL = mp.Vector3(26, 9) +PML = 1.0 +PLANE = mp.Volume(center=mp.Vector3(0, 2.5), size=mp.Vector3(12, 0)) # injection +STACK = mpa.Stack([mpa.Layer(N_OXIDE, 2.0), mpa.Layer(N_AIR)]) # above PLANE +FIBER_DISTANCE = 12.0 # from PLANE: 2 um of silica, then 10 um of air +WAIST = 2.5 + +WG_Y, WG_T = -1.5, 0.22 # silicon waveguide, fully etched grating +PERIOD, NUM_TEETH, GRATING_X0 = 0.62, 10, -3.0 +SOURCE_X, MONITOR_X = -11.0, -10.0 # far from the beam, so free-space light +WG_LINE = mp.Volume(center=mp.Vector3(MONITOR_X, WG_Y), size=mp.Vector3(0, 2.5)) + +SRC_TIME = mp.GaussianSource(FCEN, fwidth=0.2 * FCEN) +DECAY = mp.stop_when_dft_decayed(1e-9) + + +def device(): + silicon = mp.Medium(index=N_SILICON) + geometry = [ + mp.Block( + center=mp.Vector3(GRATING_X0 - 50, WG_Y), + size=mp.Vector3(100, WG_T), + material=silicon, + ) + ] + for i in range(NUM_TEETH): + geometry.append( + mp.Block( + center=mp.Vector3(GRATING_X0 + (i + 0.25) * PERIOD, WG_Y), + size=mp.Vector3(0.5 * PERIOD, WG_T), + material=silicon, + ) + ) + return geometry + + +def simulation(resolution, geometry=(), sources=()): + return mp.Simulation( + cell_size=CELL, + resolution=resolution, + boundary_layers=[mp.PML(PML)], + default_material=mp.Medium(index=N_OXIDE), + geometry=list(geometry), + sources=list(sources), + force_complex_fields=True, + ) + + +def fiber_mode(tilt, offset=0.0, waist=WAIST): + return mpa.gaussian_mode(waist, tilt_deg=tilt, offset=offset) + + +def injection(resolution, mode): + """The propagator for PLANE and the fiber's incident fields there.""" + sim = simulation(resolution) + propagator = mpa.AngularSpectrum.from_volume( + sim, sim._fit_volume_to_simulation(PLANE), STACK, [FCEN], pad_factor=8 + ) + return propagator, propagator.incident_fields(mode, distance=FIBER_DISTANCE) + + +def fiber_sources(propagator, fields): + return propagator.equivalent_sources( + fields, SRC_TIME, center=PLANE.center, size=PLANE.size + ) + + +def homogeneous(propagator, sign): + """A propagator through uniform silica on the same samples.""" + return mpa.AngularSpectrum( + mpa.Stack([mpa.Layer(N_OXIDE)]), + [FCEN], + propagator.pitch, + propagator.num_points, + normal=mp.Y, + sign=sign, + pad_factor=8, + ) + + +def waveguide_power(sim, flux, direction): + result = sim.get_eigenmode_coefficients(flux, [1], eig_parity=mp.ODD_Z) + return abs(result.alpha[0, 0, 0 if direction > 0 else 1]) ** 2 + + +# ------------------------------------------------------------ 1. the field + + +def check_field(resolution, tilt=10.0, below=3.5): + """Meep's launched field against the propagator's prediction.""" + propagator, fields = injection(resolution, fiber_mode(tilt)) + sim = simulation(resolution, sources=fiber_sources(propagator, fields)) + line = mp.Volume( + center=PLANE.center - mp.Vector3(0, below), size=mp.Vector3(PLANE.size.x, 0) + ) + region = mp.Volume(center=mp.Vector3(), size=CELL - mp.Vector3(2 * PML, 2 * PML)) + cut = sim.add_dft_fields([mp.Ez], [FCEN], where=line, yee_grid=False) + whole = sim.add_dft_fields([mp.Ez], [FCEN], where=region, yee_grid=False) + sim.run(until_after_sources=DECAY) + + x = np.asarray(sim.get_array_metadata(vol=line)[0]) + measured = np.asarray(sim.get_dft_array(cut, mp.Ez, 0)) + # Meep scales every source by the DTFT of its time profile; that is the + # only factor applied to the prediction. + spectrum = complex(source_gradient.time_profile_dtft(sim, SRC_TIME, [FCEN])[0]) + down = homogeneous(propagator, sign=-1) + outgoing = mpa.TangentialFields(E=fields.E, H=fields.H, normal=mp.Y, sign=-1) + predicted = spectrum * np.asarray( + down.propagate(outgoing, below, coordinates=x)[mp.Ez][0] + ) + meta = sim.get_array_metadata(vol=region) + return dict( + x=x, + measured=measured, + predicted=predicted, + map=np.abs(np.asarray(sim.get_dft_array(whole, mp.Ez, 0))), + map_x=np.asarray(meta[0]), + map_y=np.asarray(meta[1]), + line_y=PLANE.center.y - below, + ) + + +# ------------------------------------------------------- 2. reciprocity + + +def transmit_efficiency(resolution, tilts): + """Waveguide -> fiber, for every tilt from one simulation.""" + sim = simulation( + resolution, + device(), + [ + mp.EigenModeSource( + SRC_TIME, + center=mp.Vector3(SOURCE_X, WG_Y), + size=mp.Vector3(0, 2.5), + eig_band=1, + eig_parity=mp.ODD_Z, + direction=mp.X, + ) + ], + ) + flux_in = sim.add_flux(FCEN, 0, 1, mp.FluxRegion(volume=WG_LINE)) + monitor = sim.add_dft_fields([mp.Ez, mp.Hx], [FCEN], where=PLANE, yee_grid=False) + flux_up = sim.add_flux(FCEN, 0, 1, mp.FluxRegion(volume=PLANE)) + sim.run(until_after_sources=DECAY) + + power_in = waveguide_power(sim, flux_in, +1) + propagator = mpa.AngularSpectrum.from_monitor( + sim, monitor, STACK, PLANE, pad_factor=8 + ) + fields = propagator.fields_from_monitor(sim, monitor) + # Meep's flux per unit of the propagator's power, from the same plane + units = mp.get_fluxes(flux_up)[0] / float( + homogeneous(propagator, sign=+1).power(fields, 0.0)[0] + ) + return np.array( + [ + units + * float( + propagator.overlap( + fields, fiber_mode(t), FIBER_DISTANCE, incident_power=jnp.ones(1) + )[0] + ) + / power_in + for t in tilts + ] + ) + + +def receive_efficiency(resolution, tilt): + """Fiber -> waveguide: coupled power over the power the fiber emits.""" + mode = fiber_mode(tilt) + propagator, fields = injection(resolution, mode) + + # the power launched into the cell, measured by Meep without the device + reference = simulation(resolution, sources=fiber_sources(propagator, fields)) + below = mp.Volume( + center=PLANE.center - mp.Vector3(0, 0.5), + size=mp.Vector3(CELL.x - 2 * PML, 0), + ) + launched = reference.add_flux(FCEN, 0, 1, mp.FluxRegion(volume=below, weight=-1)) + reference.run(until_after_sources=DECAY) + + sim = simulation(resolution, device(), fiber_sources(propagator, fields)) + flux = sim.add_flux(FCEN, 0, 1, mp.FluxRegion(volume=WG_LINE)) + sim.run(until_after_sources=DECAY) + + # The fiber emits more than reaches the cell: the stack reflects some. + # The ratio comes from the propagator, in its own units. + k0 = 2 * np.pi * np.array([FCEN]) + a = np.squeeze( + np.asarray( + mode.spectrum( + propagator, propagator.kt, k0[:, None], propagator._indices[-1] + ) + ) + ) + admittance = [ + np.real(np.asarray(propagator._admittances[-1][p])[0]) + for p in (asm.S_POLARIZATION, asm.P_POLARIZATION) + ] + emitted = ( + 0.5 + * propagator._spectral_measure() + * np.sum( + admittance[0] * abs(a[..., 0]) ** 2 + admittance[1] * abs(a[..., 1]) ** 2 + ) + ) + outgoing = mpa.TangentialFields(E=fields.E, H=fields.H, normal=mp.Y, sign=-1) + arriving = float(homogeneous(propagator, sign=-1).power(outgoing, 0.0)[0]) + + return waveguide_power(sim, flux, -1) / ( + mp.get_fluxes(launched)[0] * emitted / arriving + ) + + +# ------------------------------------------------------------ 3. gradients + + +def coupled_power_loss(resolution, run_time=400.0): + """Waveguide power as a JAX function of (waist, tilt, offset).""" + sim = simulation(resolution, device()) + inert = mp.MaterialGrid( + mp.Vector3(4, 4), mp.Medium(index=N_OXIDE), mp.Medium(index=N_OXIDE) + ) + # MeepJaxWrapper needs a design region; this one changes nothing + region = mpa.DesignRegion( + inert, volume=mp.Volume(center=mp.Vector3(8, -2.5), size=mp.Vector3(0.4, 0.4)) + ) + sim.geometry = sim.geometry + [ + mp.Block(center=region.center, size=region.size, material=inert) + ] + propagator, seed = injection(resolution, fiber_mode(10.0)) + sources = fiber_sources(propagator, seed) + monitor = mpa.EigenmodeCoefficient( + sim, WG_LINE, mode=1, forward=False, eig_parity=mp.ODD_Z + ) + wrapper = mpa.MeepJaxWrapper( + sim, + sources, + [monitor], + [region], + [FCEN], + minimum_run_time=run_time, + maximum_run_time=run_time, + ) + + def loss(fiber): + mode = fiber_mode(fiber[1], offset=fiber[2], waist=fiber[0]) + fields = propagator.incident_fields(mode, distance=FIBER_DISTANCE) + sheet = propagator.equivalent_currents(fields) + values = wrapper( + [jnp.full((4, 4, 1), 0.5)], [jnp.ravel(sheet[s.component]) for s in sources] + ) + return jnp.sum(jnp.abs(jnp.asarray(values)) ** 2) + + return loss + + +def check_gradient(resolution): + x0 = jnp.array([WAIST, 10.0, 0.3]) # waist (um), tilt (deg), offset (um) + _, adjoint = jax.value_and_grad(coupled_power_loss(resolution))(x0) + reference = [] + for i, step in enumerate((0.02, 0.1, 0.02)): + e = jnp.zeros(3).at[i].set(step) + plus = float(coupled_power_loss(resolution)(x0 + e)) + minus = float(coupled_power_loss(resolution)(x0 - e)) + reference.append((plus - minus) / (2 * step)) + return np.asarray(adjoint), np.array(reference) + + +# ------------------------------------------------------------------ figure + +INK, INK_2, GRID, SURFACE = "#0b0b0b", "#52514e", "#e4e3df", "#fcfcfb" +BLUE, ORANGE = "#2a78d6", "#eb6834" + + +def figure(field, tilts, eta_tx, eta_rx, resolutions, gaps, adjoint, reference, path): + import matplotlib.pyplot as plt + from matplotlib.colors import LinearSegmentedColormap + + plt.rcParams.update( + { + "font.size": 9, + "axes.edgecolor": INK_2, + "axes.labelcolor": INK, + "axes.titlecolor": INK, + "xtick.color": INK_2, + "ytick.color": INK_2, + "axes.grid": True, + "grid.color": GRID, + "grid.linewidth": 0.6, + "axes.spines.top": False, + "axes.spines.right": False, + "figure.facecolor": SURFACE, + "axes.facecolor": SURFACE, + "legend.frameon": False, + } + ) + fig = plt.figure(figsize=(12, 7)) + grid = fig.add_gridspec(2, 3, hspace=0.45, wspace=0.32) + + # (a) the launched beam + ax = fig.add_subplot(grid[0, 0]) + blues = LinearSegmentedColormap.from_list( + "blues", [SURFACE, "#9cc1ee", BLUE, "#123a6b"] + ) + ax.pcolormesh( + field["map_x"], field["map_y"], field["map"].T, cmap=blues, shading="auto" + ) + ax.axhline(PLANE.center.y, color=INK_2, lw=1, ls="--") + ax.axhline(field["line_y"], color=ORANGE, lw=1.5) + for y, label in ( + (PLANE.center.y, "injection plane"), + (field["line_y"], "comparison line"), + ): + ax.text(6.8, y + 0.2, label, color=INK_2, fontsize=8, ha="right") + ax.set(xlabel="x (µm)", ylabel="y (µm)", xlim=(-8, 7), aspect="equal") + ax.grid(False) + ax.set_title("(a) Launched beam |Ez|, 10° fiber tilt", loc="left") + + # (b) launched field against the propagator + ax = fig.add_subplot(grid[0, 1]) + ax.plot(field["x"], field["measured"].real, color=BLUE, lw=2, label="Meep (FDTD)") + ax.plot( + field["x"][:: max(1, len(field["x"]) // 25)], + field["predicted"].real[:: max(1, len(field["x"]) // 25)], + "o", + color=ORANGE, + ms=4, + label="angular spectrum", + ) + error = np.linalg.norm(field["measured"] - field["predicted"]) / np.linalg.norm( + field["measured"] + ) + ax.text( + 0.98, + 0.06, + f"relative difference\n{error:.1%}", + transform=ax.transAxes, + color=INK_2, + ha="right", + ) + ax.set(xlabel="x (µm)", ylabel="Re Ez (a.u.)") + ax.legend(loc="upper right") + ax.set_title("(b) Field below the plane, no fitted scale", loc="left") + + # (c) gradients + ax = fig.add_subplot(grid[0, 2]) + names = ["waist", "tilt", "offset"] + position = np.arange(3) + scale = np.abs(reference) + ax.bar( + position - 0.18, adjoint / scale, 0.34, color=BLUE, label="adjoint (JAX + Meep)" + ) + ax.bar( + position + 0.18, + reference / scale, + 0.34, + color=ORANGE, + label="finite difference", + ) + for i in range(3): + rel = abs(adjoint[i] - reference[i]) / abs(reference[i]) + y = np.sign(reference[i]) * 1.08 + ax.text( + i, + y, + f"{rel:.0e}", + ha="center", + va="bottom" if y > 0 else "top", + color=INK_2, + fontsize=8, + ) + ax.axhline(0, color=INK_2, lw=0.8) + ax.set_xticks(position, names) + ax.set_ylim(-1.45, 1.9) + ax.set_ylabel("dP / dparameter, normalized") + ax.legend(loc="upper left", fontsize=8) + ax.set_title("(c) Gradient w.r.t. fiber parameters", loc="left") + + # (d) reciprocity + ax = fig.add_subplot(grid[1, 0:2]) + ax.plot(tilts, 100 * eta_tx, color=BLUE, lw=2, label="waveguide → fiber (transmit)") + ax.plot( + tilts, + 100 * eta_rx, + "o", + color=ORANGE, + ms=6, + label="fiber → waveguide (receive)", + ) + ax.set(xlabel="fiber tilt (degrees)", ylabel="coupling efficiency (%)") + ax.set_ylim(0, 1.35 * 100 * max(np.max(eta_tx), np.max(eta_rx))) + ax.legend(loc="upper center", ncol=2) + ax.set_title( + f"(d) Reciprocity: both directions, resolution {resolutions[-1]}", loc="left" + ) + + # (e) convergence of the remaining gap + ax = fig.add_subplot(grid[1, 2]) + ax.loglog(resolutions, 100 * np.asarray(gaps), "o-", color=BLUE, lw=2, ms=6) + guide = 100 * gaps[0] * (np.asarray(resolutions) / resolutions[0]) ** -2.0 + ax.loglog(resolutions, guide, "--", color=INK_2, lw=1) + ax.text(resolutions[1], guide[1] * 0.7, "second order", color=INK_2, fontsize=8) + ax.set_xticks(resolutions, [str(r) for r in resolutions]) + ticks = 100 * np.asarray(gaps) + ax.set_yticks(ticks, [f"{t:.0f}%" if t >= 10 else f"{t:.1f}%" for t in ticks]) + ax.minorticks_off() + ax.set(xlabel="resolution (pixels/µm)", ylabel="transmit vs receive gap") + ax.set_title("(e) The gap is discretization error", loc="left") + + fig.savefig(path, dpi=200, bbox_inches="tight") + print(f"wrote {path}") + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + parser.add_argument("--quick", action="store_true", help="coarser grids") + parser.add_argument("--output", default="grating_coupler_asm_end_to_end.png") + parser.add_argument( + "--replot", + action="store_true", + help="redraw the figure from the results saved by a previous run", + ) + args = parser.parse_args() + saved = args.output.rsplit(".", 1)[0] + ".npz" + + if args.replot: + r = dict(np.load(saved, allow_pickle=True)) + field = r.pop("field").item() + figure(field, **r, path=args.output) + return + + resolutions = [10, 15, 20] if args.quick else [20, 30, 40] + tilts = np.arange(4.0, 18.1, 2.0) + + field = check_field(resolutions[-1]) + error = np.linalg.norm(field["measured"] - field["predicted"]) + print(f"field: relative difference {error / np.linalg.norm(field['measured']):.2%}") + + eta_tx = transmit_efficiency(resolutions[-1], tilts) + eta_rx = np.array([receive_efficiency(resolutions[-1], t) for t in tilts]) + for t, a, b in zip(tilts, eta_tx, eta_rx): + print(f"tilt {t:4.1f} transmit {a:.4%} receive {b:.4%}") + + gaps = [] + for r in resolutions: + if r == resolutions[-1]: + i = int(np.argmin(np.abs(tilts - 10.0))) + tx, rx = eta_tx[i], eta_rx[i] + else: + tx, rx = transmit_efficiency(r, [10.0])[0], receive_efficiency(r, 10.0) + gaps.append(abs(tx - rx) / max(tx, rx)) + print(f"resolution {r}: transmit/receive gap {gaps[-1]:.1%}") + + adjoint, reference = check_gradient(resolutions[0]) + for name, a, b in zip(("waist", "tilt", "offset"), adjoint, reference): + print(f"d/d{name}: adjoint {a:+.6g} finite difference {b:+.6g}") + + if mp.am_master(): + results = dict( + tilts=tilts, + eta_tx=eta_tx, + eta_rx=eta_rx, + resolutions=np.array(resolutions), + gaps=np.array(gaps), + adjoint=adjoint, + reference=reference, + ) + np.savez(saved, field=np.array(field, dtype=object), **results) + figure(field, **results, path=args.output) + + +if __name__ == "__main__": + main()