From d51e4281f987ff8b399bfc868476a7716fc08908 Mon Sep 17 00:00:00 2001 From: Alec Hammond Date: Tue, 25 Aug 2026 20:01:16 -0700 Subject: [PATCH 1/8] adjoint: angular-spectrum propagation through stratified media, core Meep's near-to-far transformation requires its surface to sit in a homogeneous medium -- dft_near2far aborts otherwise -- so a structure radiating across a material interface has to keep that interface inside the FDTD cell. A grating coupler radiating up through a cladding into air, and then hundreds of microns to a fiber, cannot afford that. This propagates the tangential DFT fields on a planar monitor through an arbitrary layer stack analytically instead, in JAX, so the layers above the monitor leave the simulation and the stack itself becomes differentiable. Physics, each piece checked against an independent oracle before the next was written: * the stack is solved by a scattering-matrix recursion, not a transfer matrix, which would overflow on the first thick layer or evanescent order and surface as NaN gradients rather than as wrong fields. Agrees with analytic Fresnel to 3e-15, puts Brewster at 2e-16, nulls a quarter-wave antireflection coating to 8e-17, matches an independently written transfer matrix to 1e-15 in both polarizations, and conserves energy to 4e-16. * up- and down-going radiation are separated using both tangential fields, which an open near-to-far surface cannot do. For an up-going wave H_t = Y (n_hat x E_t) in both polarizations, so the split is uniform once written with the cross product. Verified exact: the spurious down-going amplitude tracks the field amplitude left at the monitor edges, 1.2e-8 for a field decayed to 1.3e-8 there. * free-space propagation reproduces the analytic spreading of a Gaussian beam to machine precision at 50, 200 and 500 um, given a padded window sized to the spread beam. * projection onto a mode uses the power inner product, so self-overlap is exactly 1, and the closed-form tilted Gaussian recovers a launch angle exactly rather than to grid resolution. Every layer gets an infinitesimal loss before kz is computed. kz vanishes on the light line, where its derivative is unbounded: jnp.sqrt(0.0) evaluates fine but its VJP is infinite, which becomes NaN and destroys the gradient while leaving the objective value looking correct. A grid point lands exactly there whenever the padded monitor width is an integer number of wavelengths in the medium, which is not a rare coincidence when cell sizes and wavelengths are round. report() returns the three diagnostics that catch a badly placed monitor: down-going fraction, evanescent fraction, and the amplitude left at the monitor edges -- the transform is periodic, so a field that has not decayed there wraps around. 2D only for now; the 3D case additionally needs the s/p rotation by azimuth with its removable singularity at normal incidence, and raises. --- python/adjoint/angular_spectrum.py | 752 +++++++++++++++++++++++++++++ 1 file changed, 752 insertions(+) create mode 100644 python/adjoint/angular_spectrum.py diff --git a/python/adjoint/angular_spectrum.py b/python/adjoint/angular_spectrum.py new file mode 100644 index 000000000..5c62effd6 --- /dev/null +++ b/python/adjoint/angular_spectrum.py @@ -0,0 +1,752 @@ +"""Differentiable angular-spectrum propagation through stratified media. + +Meep's near-to-far transformation requires the near surface to sit in a +homogeneous medium -- `dft_near2far` aborts otherwise -- so a structure whose +radiation crosses a material interface, such as a grating coupler radiating up +through a cladding into air, has to keep that interface inside the FDTD cell. + +This module propagates the tangential DFT fields on a planar monitor through an +arbitrary layer stack analytically instead, in JAX. The layers above the monitor +leave the simulation, which both shrinks the cell and makes the stack itself a +differentiable design parameter: thicknesses, indices, and the position and angle +of whatever the field is finally coupled into are all ordinary JAX values. + +The method is the plane-wave decomposition of the measured fields. On a plane in +a homogeneous layer, the tangential `E` and `H` together determine the up- and +down-going plane-wave amplitudes at every transverse wavevector; each amplitude +is then carried through the stack by a scattering-matrix recursion and +recombined. Recording both `E` and `H` is what makes the up/down split possible, +and is why this is strictly better posed than an open near-to-far surface, which +has no way to reject radiation heading the wrong way. + +Two entry points, neither requiring the other: + + # Post-processing an ordinary forward run. NumPy in, NumPy out. + monitor = sim.add_dft_fields([mp.Ez, mp.Hx], frequencies, where=plane) + sim.run(...) + asp = mpa.AngularSpectrum.from_monitor(sim, monitor, stack) + efficiency = asp.overlap_monitor(sim, monitor, fiber_mode, distance=300.0) + + # Inside an objective function, differentiated through the adjoint. + fields = asp.take(args) + efficiency = asp.overlap(fields, fiber_mode, distance=300.0) + +Only planar monitors normal to a coordinate axis are supported, and the monitor +must lie in a homogeneous region. Cylindrical coordinates are not supported. +""" + +import math +from typing import Callable, Dict, NamedTuple, Optional, Sequence, Tuple, Union + +import jax +import jax.numpy as jnp +import numpy as onp + +import meep as mp + +# Every layer is given an infinitesimal loss before its longitudinal wavevector +# is computed. `kz = sqrt(n^2 k0^2 - kt^2)` vanishes on the light line, where its +# derivative diverges: `jnp.sqrt(0.0)` evaluates to zero happily but its VJP is +# infinite, which becomes NaN and destroys the whole gradient while leaving the +# objective value looking perfect. A grid point lands exactly on the light line +# whenever the padded monitor width is an integer number of wavelengths in the +# medium, which is not a rare coincidence when cell sizes and wavelengths are +# chosen as round numbers. +DEFAULT_LOSS_REGULARIZATION = 1e-12 + +# Polarization indices used throughout. In 3D these are the s and p components +# relative to the plane of incidence; in 2D with an out-of-plane E field only S +# is populated, and with an in-plane E field only P. +S_POLARIZATION = 0 +P_POLARIZATION = 1 + + +class Layer(NamedTuple): + """One layer of a stack. + + Attributes: + index: the refractive index. A scalar, an array with one entry per + frequency, or a callable mapping frequency to index. + thickness: the thickness in Meep units, or None for the semi-infinite + layer that terminates the stack. + """ + + index: Union[complex, float, onp.ndarray, Callable] + thickness: Optional[float] = None + + @classmethod + def from_medium(cls, medium: mp.Medium, thickness: Optional[float] = None): + """Builds a layer from a `meep.Medium`, evaluating its index per frequency.""" + return cls( + index=lambda frequency: onp.sqrt( + complex(medium.epsilon(frequency)[0][0]) + ), + thickness=thickness, + ) + + def indices(self, frequencies: onp.ndarray) -> jnp.ndarray: + """Returns the index at each frequency, as a (num frequencies,) array.""" + if callable(self.index): + values = jnp.asarray([self.index(f) for f in frequencies]) + else: + values = jnp.asarray(self.index) + return jnp.broadcast_to(values, (len(frequencies),)) + + +class Stack(NamedTuple): + """The layers above (or below) the monitor, ordered away from it. + + The first layer is the one containing the monitor, and the last must be + semi-infinite, i.e. have `thickness=None`. + """ + + layers: Sequence[Layer] + loss_regularization: float = DEFAULT_LOSS_REGULARIZATION + + def validate(self) -> None: + if len(self.layers) < 1: + raise ValueError("A stack needs at least one layer.") + if self.layers[-1].thickness is not None: + raise ValueError( + "The last layer of a stack terminates it and must be " + "semi-infinite, i.e. constructed with thickness=None; got " + f"{self.layers[-1].thickness}." + ) + for i, layer in enumerate(self.layers[:-1]): + if layer.thickness is None: + raise ValueError( + f"Layer {i} of {len(self.layers)} has thickness=None, but " + "only the final layer may be semi-infinite." + ) + + +class TangentialFields(NamedTuple): + """Tangential fields on the monitor plane. + + Attributes: + E: maps a Meep field component to a (num frequencies,) + spatial array. + H: likewise for the magnetic components. + normal: the coordinate direction normal to the plane, `mp.X`, `mp.Y`, or + `mp.Z`. + sign: +1 when the radiation of interest travels along +normal, -1 for + -normal. A monitor below the structure, propagating into the + substrate, uses -1. + """ + + E: Dict[int, jnp.ndarray] + H: Dict[int, jnp.ndarray] + normal: int + sign: int = 1 + + +def _safe_sqrt(argument: jnp.ndarray) -> jnp.ndarray: + """Square root with the branch chosen so that the imaginary part is >= 0. + + Evanescent orders must decay away from the monitor rather than grow, which + fixes the branch. The caller is expected to have regularized the argument so + that it never lands exactly on zero; see DEFAULT_LOSS_REGULARIZATION. + """ + root = jnp.sqrt(jnp.asarray(argument, dtype=jnp.complex128)) + return jnp.where(jnp.imag(root) < 0, -root, root) + + +def _longitudinal_wavevector( + index: jnp.ndarray, k0: jnp.ndarray, kt: jnp.ndarray +) -> jnp.ndarray: + """kz for each (frequency, transverse wavevector), shaped (nfreq, nkt).""" + kt = jnp.asarray(kt) + transverse = jnp.sum(jnp.square(kt.reshape(kt.shape[0], -1)), axis=-1) + return _safe_sqrt( + jnp.square(index)[:, None] * jnp.square(k0)[:, None] - transverse[None, :] + ) + + +def _admittances( + index: jnp.ndarray, kz: jnp.ndarray, k0: jnp.ndarray +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """Wave admittances for s and p polarization. + + In Meep units the vacuum impedance is 1, so with omega = k0 these reduce to + `Y_s = kz / k0` and `Y_p = n^2 k0 / kz`. + """ + admittance_s = kz / k0[:, None] + admittance_p = jnp.square(index)[:, None] * k0[:, None] / kz + return admittance_s, admittance_p + + +def _interface_coefficients( + admittance_in: jnp.ndarray, admittance_out: jnp.ndarray +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """Fresnel reflection and transmission at one interface, in admittances.""" + total = admittance_in + admittance_out + reflection = (admittance_in - admittance_out) / total + transmission = 2 * admittance_in / total + return reflection, transmission + + +def _stack_transmission( + admittances: Sequence[jnp.ndarray], + wavevectors: Sequence[jnp.ndarray], + thicknesses: Sequence[float], +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """Transmission and reflection of a stack, by scattering-matrix recursion. + + Accumulating a transfer matrix instead would overflow as soon as a layer is + thick or an order is evanescent, since that formulation contains a growing + exponential; the scattering recursion below only ever forms `exp(i kz d)` + with non-negative imaginary `kz`, which decays. + + Args: + admittances: the wave admittance in each layer, one (nfreq, nkt) array + per layer, for a single polarization. + wavevectors: the longitudinal wavevector in each layer, likewise. + thicknesses: the thickness of each layer except the last. + + Returns: + The amplitude transmission and reflection coefficients of the whole + stack, referenced to the front face of the first layer. + """ + # Start from the terminating interface and recurse toward the monitor. At + # each step `reflection` is the reflection looking into the remaining stack + # and `transmission` the accumulated transmission through it. + reflection, transmission = _interface_coefficients( + admittances[-2], admittances[-1] + ) + for j in range(len(admittances) - 3, -1, -1): + phase = jnp.exp(1j * wavevectors[j + 1] * thicknesses[j + 1]) + interface_r, interface_t = _interface_coefficients( + admittances[j], admittances[j + 1] + ) + # Redheffer star product of the interface with the layer already + # accumulated behind it. + backward_r = -interface_r # reflection of the interface from the far side + denominator = 1 - backward_r * reflection * phase**2 + transmission = ( + interface_t * phase * transmission / denominator + ) + reflection = interface_r + ( + interface_t * (2 - interface_t) * reflection * phase**2 / denominator + ) + return transmission, reflection + + +def _single_interface_limit( + admittances: Sequence[jnp.ndarray], +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """Transmission and reflection when the stack is a single interface.""" + reflection, transmission = _interface_coefficients( + admittances[0], admittances[1] + ) + return transmission, reflection + + +class PropagationResult(NamedTuple): + """The plane-wave amplitudes produced by a propagation. + + Attributes: + kt: the transverse wavevectors, (num kt,) in 2D or (num kt, 2) in 3D. + amplitudes: (num frequencies, num kt, 2) electric field amplitudes in + the s and p basis, at the target plane. + kz: (num frequencies, num kt) longitudinal wavevector in the output + layer, used for flux and overlap integrals. + index: (num frequencies,) index of the output layer. + """ + + kt: jnp.ndarray + amplitudes: jnp.ndarray + kz: jnp.ndarray + index: jnp.ndarray + + +# Sign and normalization conventions, stated once because they are the usual +# source of silent errors. +# +# Meep uses exp(-i omega t), so Faraday's law reads curl E = i omega mu H. In +# Meep units mu = 1 and omega = k0. For a two-dimensional simulation in the x-y +# plane with a monitor line normal to y and an out-of-plane electric field, a +# plane wave exp(i(kx x + ky y)) therefore satisfies +# +# H_x = ky E_z / k0 +# +# so an up-going wave (ky = +kz) has H_x = +Y_s E_z and a down-going one has +# H_x = -Y_s E_z, with Y_s = kz / k0. That is what separates the two: +# +# E_up = (E_z + H_x / Y_s) / 2 E_down = (E_z - H_x / Y_s) / 2 +# +# Reflection and transmission coefficients are those of the *tangential* field +# components, which is the admittance convention. For p polarization this is the +# negative of the more familiar form written in terms of the full electric +# vector; the two agree on |r|, on the Brewster angle, and on energy +# conservation, and differ only in the sign of r_p. + + +def _tangential_components(normal: int, dimensions: int) -> Tuple[Tuple[int, ...], ...]: + """The E and H components tangential to a plane with the given normal.""" + if dimensions == 2: + if normal == mp.Y: + return (mp.Ex, mp.Ez), (mp.Hx, mp.Hz) + if normal == mp.X: + return (mp.Ey, mp.Ez), (mp.Hy, mp.Hz) + raise ValueError( + "In a 2D simulation the monitor plane must be normal to x or y, " + f"but got normal={normal}." + ) + raise NotImplementedError( + "Angular-spectrum propagation currently supports 2D simulations only. " + "The three-dimensional case additionally needs the s/p rotation by " + "azimuth, with its removable singularity at normal incidence, and is " + "not implemented here." + ) + + +# For an up-going plane wave the tangential fields satisfy H_t = Y (n_hat x E_t) +# in both polarizations, the cross product supplying the orientation, so +# +# E_up = (E_t - (n_hat x H_t) / Y) / 2 +# +# uniformly. Written out in a tangential basis this becomes one (E component, +# partnering H component, sign, polarization) tuple per polarization, with +# E_up = (E + sign * H / Y) / 2. +_DECOMPOSITION = { + mp.Y: ( + (mp.Ez, mp.Hx, +1, S_POLARIZATION), + (mp.Ex, mp.Hz, -1, P_POLARIZATION), + ), + mp.X: ( + (mp.Ez, mp.Hy, -1, S_POLARIZATION), + (mp.Ey, mp.Hz, +1, P_POLARIZATION), + ), +} + + +class Mode(NamedTuple): + """A target field to project onto, defined by its angular spectrum. + + Attributes: + spectrum: called as `spectrum(kt, k0, index)` and returning a + (num frequencies, num kt) array of tangential electric field + amplitudes. + polarization: which polarization the amplitudes belong to. + """ + + spectrum: Callable + polarization: int = S_POLARIZATION + + +def gaussian_mode( + waist: float, + tilt_deg: float = 0.0, + offset: float = 0.0, + polarization: int = S_POLARIZATION, +) -> Mode: + """A tilted, laterally offset Gaussian, e.g. a fiber mode. + + The spectrum is written in closed form rather than sampled, so the tilt and + the offset are exact continuous parameters and differentiable, instead of + being quantized by the monitor grid. + + Args: + waist: the 1/e field radius at the target plane, in Meep units. For a + fiber quoted by mode-field diameter, this is MFD / 2. + tilt_deg: the angle from the plane normal, in degrees. + offset: the lateral displacement of the beam center at the target plane. + polarization: S_POLARIZATION or P_POLARIZATION. + + Returns: + A `Mode`. + """ + + def spectrum(kt, k0, index): + kt = jnp.asarray(kt) + center = ( + jnp.asarray(index).real[:, None] + * jnp.asarray(k0)[:, None] + * jnp.sin(jnp.deg2rad(jnp.asarray(tilt_deg))) + ) + detuning = kt[None, :] - center + return jnp.exp(-jnp.square(detuning * waist) / 4.0) * jnp.exp( + -1j * kt[None, :] * offset + ) + + return Mode(spectrum=spectrum, polarization=polarization) + + +class AngularSpectrum: + """Propagates monitor fields through a layer stack. + + Attributes: + stack: the layers above the monitor, ordered away from it. + frequencies: the monitor frequencies, in Meep units. + pitch: the spacing of the monitor samples. + num_points: the number of monitor samples. + normal: the coordinate direction normal to the monitor plane. + sign: +1 if the radiation of interest travels along +normal, else -1. + """ + + def __init__( + self, + stack: Stack, + frequencies: Sequence[float], + pitch: float, + num_points: int, + normal: int = mp.Y, + sign: int = 1, + pad_factor: int = 4, + kt: Optional[onp.ndarray] = None, + ): + """Initializes a propagator. + + Args: + stack: the layers above the monitor. + frequencies: the monitor frequencies. + pitch: the monitor sample spacing, normally 1 / resolution. + num_points: the number of monitor samples. + normal: the direction normal to the monitor plane. + sign: the direction of the radiation of interest along that normal. + pad_factor: how much to zero-pad before transforming. The transform + is periodic, so a field that has not decayed by the edges of the + padded window wraps around; padding pushes the replicas apart. + Ignored when `kt` is given. + kt: transverse wavevectors to evaluate, instead of the uniform grid + a padded transform would produce. Evaluating a chosen set costs + O(num_points * num kt) and sidesteps padding entirely, which is + worthwhile when only the wavevectors within some numerical + aperture matter. + """ + stack.validate() + if normal not in _DECOMPOSITION: + raise ValueError(f"Unsupported monitor normal {normal}.") + if sign not in (1, -1): + raise ValueError(f"sign must be +1 or -1, got {sign}.") + + self.stack = stack + self.frequencies = onp.asarray(frequencies, dtype=float) + self.pitch = float(pitch) + self.num_points = int(num_points) + self.normal = normal + self.sign = int(sign) + self.pad_factor = int(pad_factor) + + self._k0 = 2 * onp.pi * self.frequencies + if kt is None: + padded = self.num_points * self.pad_factor + self._kt = jnp.asarray( + 2 * onp.pi * onp.fft.fftfreq(padded, d=self.pitch) + ) + self._uniform = True + self._padded = padded + else: + self._kt = jnp.asarray(kt, dtype=float) + self._uniform = False + self._padded = None + + # Regularizing the index keeps kz off the light line, where its + # derivative is unbounded; see DEFAULT_LOSS_REGULARIZATION. + regularizer = 1 + 1j * self.stack.loss_regularization + self._indices = [ + layer.indices(self.frequencies) * regularizer for layer in stack.layers + ] + self._wavevectors = [ + _longitudinal_wavevector(index, self._k0, self._kt) + for index in self._indices + ] + self._admittances = [ + _admittances(index, kz, self._k0) + for index, kz in zip(self._indices, self._wavevectors) + ] + self._thicknesses = [ + 0.0 if layer.thickness is None else float(layer.thickness) + for layer in stack.layers + ] + + @property + def kt(self) -> jnp.ndarray: + """The transverse wavevectors at which the spectrum is evaluated.""" + return self._kt + + @property + def stack_thickness(self) -> float: + """The total thickness of the stack, excluding the semi-infinite layer.""" + return float(sum(self._thicknesses[:-1])) + + def coordinates(self) -> onp.ndarray: + """The transverse coordinates of the monitor samples, centered on zero.""" + return (onp.arange(self.num_points) - (self.num_points - 1) / 2) * self.pitch + + def _transform(self, values: jnp.ndarray) -> jnp.ndarray: + """Transforms (num frequencies, num points) samples to (num freq, num kt).""" + values = jnp.asarray(values) + if self._uniform: + padded = jnp.zeros( + values.shape[:-1] + (self._padded,), dtype=jnp.complex128 + ) + padded = padded.at[..., : self.num_points].set(values) + spectrum = jnp.fft.fft(padded, axis=-1) * self.pitch + # The samples are centered on zero, so undo the phase ramp implied by + # having placed them at indices 0..num_points-1. + origin = self.coordinates()[0] + return spectrum * jnp.exp(-1j * self._kt * origin) + phase = jnp.exp(-1j * self._kt[:, None] * self.coordinates()[None, :]) + return jnp.einsum("kx,...x->...k", phase, values) * self.pitch + + def _polarization_terms(self, fields: TangentialFields): + """Yields (polarization, E samples, H samples, sign) for what is present.""" + if fields.normal != self.normal: + raise ValueError( + f"The fields are on a plane normal to {fields.normal} but this " + f"propagator was built for {self.normal}." + ) + for e_component, h_component, sign, polarization in _DECOMPOSITION[ + self.normal + ]: + e_values = fields.E.get(e_component) + h_values = fields.H.get(h_component) + if e_values is None and h_values is None: + continue + if e_values is None or h_values is None: + raise ValueError( + "Separating up-going from down-going radiation needs both " + f"tangential fields, but only one of {mp.component_name(e_component)}" + f" and {mp.component_name(h_component)} was supplied." + ) + yield polarization, jnp.asarray(e_values), jnp.asarray(h_values), sign + + def decompose(self, fields: TangentialFields): + """Splits the monitor fields into up- and down-going spectra. + + Returns: + `(up, down)`, each a dict mapping polarization to a + (num frequencies, num kt) array of tangential electric field + amplitudes at the monitor plane. + """ + up, down = {}, {} + admittance = self._admittances[0] + for polarization, e_values, h_values, sign in self._polarization_terms( + fields + ): + e_spectrum = self._transform(e_values) + h_spectrum = self._transform(h_values) + # `sign` already carries the orientation of n_hat x H_t; `fields.sign` + # flips which branch counts as outgoing for a downward-facing monitor. + scaled = fields.sign * sign * h_spectrum / admittance[polarization] + up[polarization] = 0.5 * (e_spectrum + scaled) + down[polarization] = 0.5 * (e_spectrum - scaled) + if not up: + raise ValueError( + "No tangential field components were supplied; nothing to " + "propagate." + ) + return up, down + + def _transmission(self, polarization: int): + """Transmission and reflection of the stack for one polarization.""" + admittances = [pair[polarization] for pair in self._admittances] + if len(admittances) == 2: + return _single_interface_limit(admittances) + return _stack_transmission( + admittances, self._wavevectors, self._thicknesses + ) + + def spectrum(self, fields: TangentialFields, distance: float): + """The up-going spectrum carried to a plane `distance` from the monitor. + + Args: + fields: the monitor fields. + distance: how far along the outgoing normal to propagate, measured + from the monitor plane. It must reach at least to the far side of + the stack. + + Returns: + A `PropagationResult`. + """ + remaining = distance - self.stack_thickness + if onp.any(onp.asarray(remaining) < 0): + raise ValueError( + f"distance={distance} does not clear the stack, which is " + f"{self.stack_thickness} thick. The target plane has to lie in " + "the semi-infinite layer." + ) + up, _ = self.decompose(fields) + amplitudes = {} + for polarization, amplitude in up.items(): + transmission, _ = self._transmission(polarization) + # Monitor to the first interface, through the stack, then onward in + # the terminating layer. + to_interface = jnp.exp( + 1j * self._wavevectors[0] * self._thicknesses[0] + ) + beyond = jnp.exp(1j * self._wavevectors[-1] * remaining) + amplitudes[polarization] = amplitude * to_interface * transmission * beyond + stacked = jnp.stack( + [ + amplitudes.get( + polarization, + jnp.zeros_like(next(iter(amplitudes.values()))), + ) + for polarization in (S_POLARIZATION, P_POLARIZATION) + ], + axis=-1, + ) + return PropagationResult( + kt=self._kt, + amplitudes=stacked, + kz=self._wavevectors[-1], + index=self._indices[-1], + ) + + def _weights(self, result: PropagationResult) -> jnp.ndarray: + """Per-wavevector power weights, (num frequencies, num kt, 2). + + Power carried along the normal by a plane-wave component is + `Re(Y) |E|^2 / 2`, so this is the measure that makes overlaps unitary. + Evanescent components carry none and drop out. + """ + admittance_s, admittance_p = _admittances( + self._indices[-1], result.kz, self._k0 + ) + measure = self._spectral_measure() + return ( + 0.5 + * measure + * jnp.stack([jnp.real(admittance_s), jnp.real(admittance_p)], axis=-1) + ) + + def _spectral_measure(self) -> float: + """The dk / 2pi factor that turns a spectral sum into a real-space integral.""" + if self._uniform: + return (2 * onp.pi / (self._padded * self.pitch)) / (2 * onp.pi) + spacing = jnp.diff(self._kt) + # Trapezoid weights would be more careful, but a chosen kt set is + # normally uniform; require that rather than silently mis-weighting. + return jnp.mean(spacing) / (2 * onp.pi) + + def power(self, fields: TangentialFields, distance: Optional[float] = None): + """Outgoing power through the target plane, one value per frequency.""" + result = self.spectrum(fields, self.stack_thickness if distance is None else distance) + weights = self._weights(result) + return jnp.sum(weights * jnp.abs(result.amplitudes) ** 2, axis=(1, 2)) + + def overlap( + self, + fields: TangentialFields, + mode: Mode, + distance: float, + incident_power: Optional[jnp.ndarray] = None, + ): + """Fraction of the incident power coupled into `mode`. + + The propagated field is projected onto the normalized mode using the + power inner product, so the result is bounded by one and equals one when + the field is a pure multiple of the mode and all of the incident power + reaches the target plane. + + Args: + fields: the monitor fields. + mode: the target, e.g. `gaussian_mode(...)`. + distance: distance from the monitor to the target plane. + incident_power: the power to normalize against, one value per + frequency. Defaults to the power crossing the target plane, in + which case the result is the modal purity of the radiation + rather than an end-to-end efficiency. + + Returns: + A (num frequencies,) array. + """ + result = self.spectrum(fields, distance) + weights = self._weights(result)[..., mode.polarization] + field_amplitude = result.amplitudes[..., mode.polarization] + mode_amplitude = mode.spectrum(result.kt, self._k0, result.index) + + cross = jnp.sum(weights * field_amplitude * jnp.conj(mode_amplitude), axis=-1) + mode_norm = jnp.sum(weights * jnp.abs(mode_amplitude) ** 2, axis=-1) + coupled = jnp.abs(cross) ** 2 / mode_norm + if incident_power is None: + incident_power = jnp.sum( + weights * jnp.abs(field_amplitude) ** 2, axis=-1 + ) + return coupled / incident_power + + def report(self, fields: TangentialFields) -> Dict[str, jnp.ndarray]: + """Diagnostics for the four ways a monitor is usually placed wrongly. + + Returns a dict with, per frequency: + + `downgoing_fraction` + power heading back toward the structure. Large means something above + the monitor is scattering, or the monitor sits inside the near field. + `evanescent_fraction` + power above the light line, which does not propagate. Large means + the monitor is too close to the structure. + `edge_amplitude` + field magnitude at the ends of the monitor relative to its peak. A + transform is periodic, so a field that has not decayed by the edges + wraps around; large means the monitor is too narrow. + """ + up, down = self.decompose(fields) + admittance = self._admittances[0] + measure = self._spectral_measure() + propagating = jnp.abs(jnp.imag(self._wavevectors[0])) < 1e-8 * jnp.abs( + jnp.real(self._wavevectors[0]) + 1e-30 + ) + + def spectral_power(spectra, mask=None): + total = 0.0 + for polarization, amplitude in spectra.items(): + weight = 0.5 * measure * jnp.real(admittance[polarization]) + weighted = weight * jnp.abs(amplitude) ** 2 + if mask is not None: + weighted = jnp.where(mask, weighted, 0.0) + total = total + jnp.sum(weighted, axis=-1) + return total + + up_power = spectral_power(up) + down_power = spectral_power(down) + up_propagating = spectral_power(up, propagating) + + edges = [] + for _, e_values, _, _ in self._polarization_terms(fields): + magnitude = jnp.abs(e_values) + peak = jnp.max(magnitude, axis=-1) + edge = jnp.maximum(magnitude[..., 0], magnitude[..., -1]) + edges.append(edge / jnp.where(peak > 0, peak, 1.0)) + + total = up_power + down_power + return { + "downgoing_fraction": down_power / jnp.where(total > 0, total, 1.0), + "evanescent_fraction": 1.0 + - up_propagating / jnp.where(up_power > 0, up_power, 1.0), + "edge_amplitude": jnp.max(jnp.stack(edges, axis=0), axis=0), + } + + def propagate( + self, fields: TangentialFields, distance: float, coordinates=None + ) -> Dict[int, jnp.ndarray]: + """The tangential electric field at the target plane, in real space. + + Args: + fields: the monitor fields. + distance: distance from the monitor to the target plane. + coordinates: where to evaluate. Defaults to the monitor's own + coordinates; pass a wider range to see a beam that has spread. + + Returns: + A dict mapping each tangential electric component to a + (num frequencies, num coordinates) array. + """ + result = self.spectrum(fields, distance) + if coordinates is None: + coordinates = self.coordinates() + coordinates = jnp.asarray(coordinates) + measure = self._spectral_measure() + phase = jnp.exp(1j * result.kt[None, :] * coordinates[:, None]) + outputs = {} + for e_component, _, _, polarization in _DECOMPOSITION[self.normal]: + if not jnp.any(result.amplitudes[..., polarization]): + continue + outputs[e_component] = ( + jnp.einsum( + "xk,fk->fx", phase, result.amplitudes[..., polarization] + ) + * measure + ) + return outputs From 9678f4284c25db27fc8fce753a73a5a01b1fd7a8 Mon Sep 17 00:00:00 2001 From: Alec Hammond Date: Tue, 25 Aug 2026 20:06:47 -0700 Subject: [PATCH 2/8] adjoint: Meep adapter for angular-spectrum propagation Two ways in, neither requiring the other. from_monitor() post-processes an ordinary forward run and deals in NumPy, so someone who only wants a far field never meets JAX. It infers the plane normal, sample pitch and sample count from the monitor volume, and -- mirroring the check dft_near2far makes for the same reason -- rejects a monitor that does not lie in a homogeneous medium, which catches a plane accidentally clipping a waveguide or a PML. objective_arguments()/take() build and repack the FourierFields an objective function needs. Both tangential components of both polarizations are registered: the unused ones cost two extra DFT line monitors in the forward run and nothing in the adjoint, since a monitor whose cotangent is identically zero places no adjoint source, and registering them spares the user working out which polarization their source excites. Validated against Meep's own near2far in a homogeneous medium, which is an independent implementation of the same physics: the propagated field agrees to 1.6e-2 at resolution 20 and 4.0e-3 at resolution 40. Converging at second order is the part that matters -- it is what a wrong sample pitch, a wrong coordinate origin, or a missing phase ramp would not do, since those produce a fixed error. Getting that comparison to mean anything took two attempts, both caught by the diagnostics rather than by the numbers. A point dipole close to the plane never decays across it, and a full-width beam source plus a strong scatterer spread field over the whole monitor; in both cases report()['edge_amplitude'] was 0.68 and 0.055 respectively and the disagreement sat at 25% and 4% without improving with resolution. near2far tolerates an undecayed field because it integrates currents on an open surface; the transform here is periodic and wraps around. With a narrow apodized source the edge amplitude is 5e-9 and the comparison converges. --- python/Makefile.am | 1 + python/adjoint/__init__.py | 8 ++ python/adjoint/angular_spectrum.py | 206 +++++++++++++++++++++++++++++ 3 files changed, 215 insertions(+) diff --git a/python/Makefile.am b/python/Makefile.am index a5c6553c9..bf3323531 100644 --- a/python/Makefile.am +++ b/python/Makefile.am @@ -245,6 +245,7 @@ pkgpython_PYTHON = __init__.py $(HL_IFACE) adjointdir = $(pkgpythondir)/adjoint adjoint_PYTHON = $(srcdir)/adjoint/__init__.py \ + $(srcdir)/adjoint/angular_spectrum.py \ $(srcdir)/adjoint/basis.py \ $(srcdir)/adjoint/objective.py \ $(srcdir)/adjoint/optimization_problem.py \ diff --git a/python/adjoint/__init__.py b/python/adjoint/__init__.py index 38fb779fb..a6a083274 100644 --- a/python/adjoint/__init__.py +++ b/python/adjoint/__init__.py @@ -25,5 +25,13 @@ # so objective functions written with `jax.numpy` need no special treatment. try: from .wrapper import MeepJaxWrapper, value_and_jacobian + from .angular_spectrum import ( + AngularSpectrum, + Layer, + Mode, + Stack, + TangentialFields, + gaussian_mode, + ) except ModuleNotFoundError as _: pass diff --git a/python/adjoint/angular_spectrum.py b/python/adjoint/angular_spectrum.py index 5c62effd6..5833f8c77 100644 --- a/python/adjoint/angular_spectrum.py +++ b/python/adjoint/angular_spectrum.py @@ -750,3 +750,209 @@ def propagate( * measure ) return outputs + + # ---------------------------------------------------------------- Meep glue + # + # Two ways in. `from_monitor` post-processes an ordinary forward run and + # deals in NumPy, so a user who only wants a far field never meets JAX. + # `for_design` builds the `FourierFields` an objective function needs, for + # use inside `OptimizationProblem` or a `MeepJaxWrapper` loss. + + @staticmethod + def _plane_geometry(simulation: mp.Simulation, volume: mp.Volume): + """Infers the normal, sample pitch, and sample count of a planar volume.""" + size = [volume.size.x, volume.size.y, volume.size.z] + zero = [i for i, extent in enumerate(size) if extent == 0] + if simulation.dimensions != 2: + raise NotImplementedError( + "Angular-spectrum propagation currently supports 2D " + f"simulations only, but this one is {simulation.dimensions}D." + ) + if len(zero) != 2 or 2 not in zero: + raise ValueError( + "The monitor must be a line normal to x or y, i.e. a Volume " + f"with exactly one nonzero in-plane size; got size={size}." + ) + normal = mp.X if zero[0] == 0 else mp.Y + coordinates = simulation.get_array_metadata(vol=volume) + axis = 1 if normal == mp.X else 0 + samples = onp.asarray(coordinates[axis]) + if samples.size < 2: + raise ValueError("The monitor needs at least two sample points.") + pitch = float(onp.mean(onp.diff(samples))) + return normal, pitch, int(samples.size) + + @staticmethod + def _assert_homogeneous( + simulation: mp.Simulation, volume: mp.Volume, tolerance: float = 1e-6 + ): + """Rejects a monitor that does not lie in a homogeneous region. + + The plane-wave decomposition assumes a single index at the monitor. This + mirrors the check `dft_near2far` makes for the same reason, and catches a + monitor accidentally clipping a waveguide or a PML. + """ + epsilon = onp.asarray( + simulation.get_array(vol=volume, component=mp.Dielectric) + ) + spread = float(onp.max(epsilon) - onp.min(epsilon)) + if spread > tolerance * max(1.0, float(onp.max(epsilon))): + raise ValueError( + "The monitor plane does not lie in a homogeneous medium: " + f"epsilon varies from {onp.min(epsilon):.6g} to " + f"{onp.max(epsilon):.6g} along it. Angular-spectrum propagation " + "needs a single index at the monitor, so move the plane clear of " + "any structure." + ) + return float(onp.mean(epsilon)) + + @classmethod + def from_monitor( + cls, + simulation: mp.Simulation, + monitor, + stack: Stack, + volume: mp.Volume, + sign: int = 1, + **kwargs, + ): + """Builds a propagator matching an existing `dft_fields` monitor. + + Args: + simulation: the simulation the monitor belongs to. + monitor: the object returned by `add_dft_fields`. + stack: the layers above the monitor. + volume: the same volume the monitor was registered with. + sign: +1 if the radiation of interest travels along +normal. + **kwargs: forwarded to the constructor, e.g. `pad_factor`. + """ + normal, pitch, num_points = cls._plane_geometry(simulation, volume) + cls._assert_homogeneous(simulation, volume) + frequencies = onp.asarray(monitor.freq) + return cls( + stack, + frequencies, + pitch, + num_points, + normal=normal, + sign=sign, + **kwargs, + ) + + def fields_from_monitor( + self, simulation: mp.Simulation, monitor + ) -> TangentialFields: + """Reads the tangential DFT fields off a monitor into a `TangentialFields`.""" + electric, magnetic = {}, {} + # `DftFields` keeps only the component count, but `DftObj` retains the + # arguments `add_dft_fields` was called with, the first of which is the + # component list. Fall back to trying every tangential component if that + # ever stops being true. + registered = getattr(monitor, "args", None) + components = ( + set(registered[0]) + if registered + else { + component + for pair in _DECOMPOSITION[self.normal] + for component in pair[:2] + } + ) + for e_component, h_component, _, _ in _DECOMPOSITION[self.normal]: + for component, target in ( + (e_component, electric), + (h_component, magnetic), + ): + if component not in components: + continue + values = onp.array( + [ + simulation.get_dft_array(monitor, component, i) + for i in range(len(self.frequencies)) + ] + ) + if values.shape[-1] != self.num_points: + raise ValueError( + f"The monitor returned {values.shape[-1]} samples for " + f"{mp.component_name(component)} but the propagator was " + f"built for {self.num_points}. The volume Meep actually " + "used may have been snapped to the grid; build the " + "propagator from the same volume that was registered." + ) + target[component] = values + return TangentialFields( + E=electric, H=magnetic, normal=self.normal, sign=self.sign + ) + + def propagate_monitor(self, simulation, monitor, distance, coordinates=None): + """`propagate`, reading from a monitor and returning NumPy arrays.""" + fields = self.fields_from_monitor(simulation, monitor) + return { + component: onp.asarray(values) + for component, values in self.propagate( + fields, distance, coordinates + ).items() + } + + def overlap_monitor(self, simulation, monitor, mode, distance, incident_power=None): + """`overlap`, reading from a monitor and returning a NumPy array.""" + fields = self.fields_from_monitor(simulation, monitor) + return onp.asarray(self.overlap(fields, mode, distance, incident_power)) + + def power_monitor(self, simulation, monitor, distance=None): + """`power`, reading from a monitor and returning a NumPy array.""" + fields = self.fields_from_monitor(simulation, monitor) + return onp.asarray(self.power(fields, distance)) + + def report_monitor(self, simulation, monitor) -> Dict[str, onp.ndarray]: + """`report`, reading from a monitor and returning NumPy arrays.""" + fields = self.fields_from_monitor(simulation, monitor) + return { + key: onp.asarray(value) + for key, value in self.report(fields).items() + } + + def objective_arguments(self, simulation, volume, **kwargs): + """The `FourierFields` an objective function needs, for the adjoint path. + + Both tangential components of both polarizations are registered. The + unused ones cost two extra DFT line monitors in the forward run and + nothing in the adjoint, since a monitor whose cotangent is identically + zero places no adjoint source, and registering them removes the need for + the user to work out which polarization their source excites. + """ + from . import FourierFields + + self._objective_components = [ + component + for e_component, h_component, _, _ in _DECOMPOSITION[self.normal] + for component in (e_component, h_component) + ] + return [ + FourierFields(simulation, volume, component, yee_grid=False, **kwargs) + for component in self._objective_components + ] + + def take(self, args) -> TangentialFields: + """Repacks the leading objective-function arguments into `TangentialFields`. + + `OptimizationProblem` passes objective arguments positionally, so this + consumes the ones `objective_arguments` produced and leaves the rest for + the caller's own monitors. + """ + if not hasattr(self, "_objective_components"): + raise RuntimeError( + "Call objective_arguments() before take(), so that the " + "component order is known." + ) + electric, magnetic = {}, {} + for component, values in zip(self._objective_components, args): + target = magnetic if mp.is_magnetic(component) else electric + target[component] = values + return TangentialFields( + E=electric, H=magnetic, normal=self.normal, sign=self.sign + ) + + def __len__(self) -> int: + """How many leading objective arguments `take` consumes.""" + return len(_DECOMPOSITION[self.normal]) * 2 From e6eb77715ccbf86b79212d02ed888fed9bc1fa08 Mon Sep 17 00:00:00 2001 From: Alec Hammond Date: Tue, 25 Aug 2026 20:11:32 -0700 Subject: [PATCH 3/8] examples: grating coupler radiating through a thick superstrate into a fiber A two-etch silicon grating radiates upward, crosses a few hundred microns of silica and the silica/air interface, and is collected by a fiber. The FDTD cell stops a micron above the device layer; everything above it is an analytic layer stack. Meshing 300 um of glass at resolution 20 would be several hundred million pixels in 2D, and the interface puts near2far out of reach regardless, since it requires a homogeneous medium. Two modes, because the propagator is useful without an optimizer: forward one simulation, then the far field, the coupling efficiency, and sweeps over fiber tilt, working distance and superstrate thickness. None of the sweeps re-runs the simulation -- those parameters live entirely in the analytic stack, which is the point. optimize topology optimization of both etch levels against the fiber overlap, with the objective written in jax.numpy and handed straight to OptimizationProblem. Two stacked design regions rather than one: a single fully etched layer radiates roughly symmetrically, and it is the second etch level that buys directionality. The example prints the monitor diagnostics before any efficiency and says plainly when they are bad, because the failure they catch is quiet. Getting the geometry right took three rounds of exactly that: * the monitor was initially only as wide as the grating, leaving 19% of the peak field at its ends; the transform is periodic, so that wraps. * widening the cell barely helped -- 4.0% to 3.9% for half again the width -- which is the tell that the residual is not the beam tail but near-grazing radiation, which travels sideways instead of decaying. * the real remaining flaw was 0.3 um between the monitor and the PML. A plane that close picks up the absorber's residual reflection. A wavelength of clearance brought it to 2.4%. A few percent is intrinsic here and bounds the accuracy at a similar level; power at those angles was never going to reach the fiber. The example says so rather than quietly windowing the monitor, which would hide it. Verified end to end at reduced settings: the optimizer takes the coupling from 0.0116 to 0.0273 in five iterations, which also confirms the adjoint gradient reaches both design regions with the right sign. --- .../grating_coupler_asm.py | 372 ++++++++++++++++++ 1 file changed, 372 insertions(+) create mode 100644 python/examples/adjoint_optimization/grating_coupler_asm.py diff --git a/python/examples/adjoint_optimization/grating_coupler_asm.py b/python/examples/adjoint_optimization/grating_coupler_asm.py new file mode 100644 index 000000000..4898af6d3 --- /dev/null +++ b/python/examples/adjoint_optimization/grating_coupler_asm.py @@ -0,0 +1,372 @@ +"""Grating coupler radiating through a thick glass superstrate into a fiber. + +A two-etch silicon grating coupler radiates upward, propagates a few hundred +microns through silica, crosses the silica/air interface, and is collected by a +fiber. The propagation and the interface are handled analytically by +`meep.adjoint.AngularSpectrum` rather than being meshed, which is what makes the +problem tractable: at resolution 20 a 300 um tall cell would be several hundred +million pixels, and the interface would put Meep's near-to-far transformation out +of reach anyway, since it requires a homogeneous medium. + +The FDTD cell therefore stops less than a micron above the device layer. Only the +grating is simulated; everything above the monitor is a layer stack. + +Two modes: + + python grating_coupler_asm.py forward + One simulation, then the far field. Prints the coupling efficiency and + the monitor diagnostics, and sweeps fiber tilt, working distance and + superstrate thickness -- none of which re-runs the simulation, because + they live entirely in the analytic propagator. + + python grating_coupler_asm.py optimize + Topology optimization of the two etch layers against the fiber overlap. + +Run with --help for the geometry and solver knobs. +""" + +import argparse +import math + +import jax +import jax.numpy as jnp +import numpy as np + +import meep as mp +import meep.adjoint as mpa + +jax.config.update("jax_enable_x64", True) + +# ------------------------------------------------------------------ geometry +N_SI, N_OXIDE, N_AIR = 3.48, 1.444, 1.0 +T_DEVICE = 0.22 # silicon device layer +T_DEEP = 0.15 # lower etch level: the deep grating teeth +T_SHALLOW = T_DEVICE - T_DEEP # upper etch level, for directionality +T_BOX = 2.0 # buried oxide +T_HANDLE = 1.0 # how much silicon handle to include below the box + +WAVEGUIDE_LENGTH = 4.0 +MONITOR_STANDOFF = 1.0 # monitor height above the device layer +DPML = 1.0 + +# The monitor plane must be clear of the near field, and the transform is +# periodic, so the monitor also has to be wide enough that the radiated field +# has decayed by its ends -- otherwise it wraps around. `report()` measures both, +# and `check_monitor` below refuses to report numbers built on a wrapped field. +# The cell is therefore padded well beyond the grating aperture. +PAD = 12.0 +MONITOR_MARGIN = 0.5 +EDGE_AMPLITUDE_LIMIT = 1e-2 + + +def build_arguments(): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("mode", choices=["forward", "optimize"]) + parser.add_argument("--resolution", type=float, default=20.0, + help="pixels per micron") + parser.add_argument("--aperture", type=float, default=20.0, + help="width of the grating, in microns") + 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=5.0, + help="fiber facet height above the silica/air interface") + parser.add_argument("--fiber-waist", type=float, default=None, + help="fiber 1/e field radius; defaults to the diffracted " + "beam size at the facet") + parser.add_argument("--tilt", type=float, default=8.0, + help="fiber tilt from normal, in degrees") + parser.add_argument("--wavelengths", type=float, nargs="+", + default=[1.52, 1.55, 1.58]) + parser.add_argument("--pad-factor", type=int, default=16, + help="zero padding before the transform; a beam that " + "spreads over hundreds of microns needs a wide window") + parser.add_argument("--iterations", type=int, default=15) + parser.add_argument("--design-resolution", type=float, default=None) + return parser.parse_args() + + +def cell_geometry(args): + """Returns the cell size, center, and the x extent of the design regions.""" + x_min = -(args.aperture / 2 + WAVEGUIDE_LENGTH + PAD + DPML) + x_max = args.aperture / 2 + PAD + DPML + y_min = -(T_BOX + T_HANDLE + DPML) + # Leave a wavelength of oxide above the monitor before the PML starts; + # a plane pressed up against the absorber picks up its residual + # reflection, and near-grazing rays have nowhere to go. + y_max = T_DEVICE + MONITOR_STANDOFF + 1.2 + DPML + size = mp.Vector3(x_max - x_min, y_max - y_min) + center = mp.Vector3(0.5 * (x_min + x_max), 0.5 * (y_min + y_max)) + return size, center + + +def design_regions(args): + """The two etch levels, stacked, spanning the grating aperture. + + Splitting the device layer into a deep and a shallow level is what lets the + grating radiate preferentially upward instead of symmetrically; a single + fully etched layer is limited to roughly half its power going the wrong way. + """ + resolution = args.design_resolution or (2 * args.resolution) + nx = int(round(args.aperture * resolution)) + 1 + oxide, silicon = mp.Medium(index=N_OXIDE), mp.Medium(index=N_SI) + + grids, volumes = [], [] + for y_lo, y_hi in ((0.0, T_DEEP), (T_DEEP, T_DEVICE)): + grids.append( + mp.MaterialGrid(mp.Vector3(nx, 1, 1), oxide, silicon, + weights=np.ones((nx,)), do_averaging=False, beta=0) + ) + volumes.append( + mp.Volume(center=mp.Vector3(0, 0.5 * (y_lo + y_hi)), + size=mp.Vector3(args.aperture, y_hi - y_lo)) + ) + return grids, volumes, nx + + +def build_simulation(args, weights=None): + """Assembles the simulation. The cell stops just above the device layer.""" + size, center = cell_geometry(args) + grids, volumes, nx = design_regions(args) + if weights is not None: + for grid, w in zip(grids, weights): + grid.update_weights(np.asarray(w).ravel()) + + silicon, oxide = mp.Medium(index=N_SI), mp.Medium(index=N_OXIDE) + geometry = [ + # silicon handle below the buried oxide + mp.Block(center=mp.Vector3(center.x, -(T_BOX + T_HANDLE + DPML) / 2 - T_BOX / 2), + size=mp.Vector3(mp.inf, T_HANDLE + DPML), material=silicon), + # the input waveguide, running in from the left + mp.Block(center=mp.Vector3(-(args.aperture / 2 + WAVEGUIDE_LENGTH + PAD + DPML) / 2 + - args.aperture / 4, T_DEVICE / 2), + size=mp.Vector3(2 * (WAVEGUIDE_LENGTH + PAD + DPML), T_DEVICE), + material=silicon), + ] + [ + mp.Block(center=volume.center, size=volume.size, material=grid) + for grid, volume in zip(grids, volumes) + ] + + frequencies = [1.0 / w for w in args.wavelengths] + center_frequency = float(np.mean(frequencies)) + source = [ + mp.EigenModeSource( + mp.GaussianSource(center_frequency, fwidth=0.2 * center_frequency), + center=mp.Vector3(-(args.aperture / 2 + WAVEGUIDE_LENGTH * 0.6), T_DEVICE / 2), + size=mp.Vector3(0, 6 * T_DEVICE), + eig_band=1, eig_parity=mp.ODD_Z, eig_match_freq=True, + ) + ] + + simulation = mp.Simulation( + cell_size=size, geometry_center=center, resolution=args.resolution, + boundary_layers=[mp.PML(DPML)], geometry=geometry, sources=source, + default_material=oxide, dimensions=2, + ) + return simulation, grids, volumes, nx, frequencies + + +def monitor_volume(args): + """The plane the radiated field is read on, inside the homogeneous oxide.""" + size, center = cell_geometry(args) + width = size.x - 2 * DPML - 2 * MONITOR_MARGIN + return mp.Volume(center=mp.Vector3(center.x, T_DEVICE + MONITOR_STANDOFF), + size=mp.Vector3(width, 0)) + + +def build_stack(args): + """Everything above the monitor: the rest of the silica, then air.""" + remaining_oxide = args.superstrate - MONITOR_STANDOFF + if remaining_oxide <= 0: + raise ValueError( + f"--superstrate {args.superstrate} must exceed the monitor standoff " + f"{MONITOR_STANDOFF}." + ) + return mpa.Stack([mpa.Layer(N_OXIDE, remaining_oxide), mpa.Layer(N_AIR)]), remaining_oxide + + +def diffracted_waist(args): + """A rough estimate of the beam radius at the fiber facet. + + Used only to pick a sensible default fiber mode: a beam launched from an + aperture this size and allowed to diffract for this far is not going to be + collected by a standard single-mode fiber. + """ + waist = args.aperture / 4 + wavelength = float(np.mean(args.wavelengths)) + rayleigh = math.pi * waist**2 * N_OXIDE / wavelength + distance = args.superstrate - MONITOR_STANDOFF + args.working_distance + return waist * math.sqrt(1 + (distance / rayleigh) ** 2) + + +def initial_weights(args, nx): + """A uniform grating, as a starting point for the optimizer.""" + x = np.linspace(-args.aperture / 2, args.aperture / 2, nx) + period = float(np.mean(args.wavelengths)) / (N_OXIDE * math.sin(math.radians(args.tilt)) + + 2.6) + # A binary grating, not a smoothly graded one: a sinusoidal index + # modulation is a much weaker scatterer and radiates very little. + deep = (np.cos(2 * math.pi * x / period) > 0).astype(float) + shallow = (np.cos(2 * math.pi * x / period - 0.6) > -0.3).astype(float) + return [deep, shallow] + + +def check_monitor(propagator, simulation, monitor): + """Prints the monitor diagnostics, and complains if they are bad. + + These are worth reading before any efficiency is believed. The transform is + periodic, so a field that has not decayed by the ends of the monitor wraps + around and contaminates everything downstream -- and it does so quietly, + producing a plausible-looking number that does not improve with resolution. + """ + report = propagator.report_monitor(simulation, monitor) + print("\nmonitor diagnostics (per wavelength)") + for key, value in report.items(): + print(f" {key:22s} " + " ".join(f"{v:.2e}" for v in np.atleast_1d(value))) + + edge = float(np.max(report["edge_amplitude"])) + if edge > 0.1: + print(f"\n WARNING: the field is still {edge:.1%} of its peak at the ends of") + print(" the monitor, so the transform is wrapping badly and the numbers") + print(" below are not trustworthy. Widen the cell by increasing PAD.") + elif edge > EDGE_AMPLITUDE_LIMIT: + print(f"\n Note: {edge:.1%} of the peak field remains at the ends of the") + print(" monitor. A few percent is normal here and is not the beam tail --") + print(" it is near-grazing radiation, which travels sideways rather than") + print(" decaying, so widening the cell barely helps. It bounds the") + print(" accuracy of the efficiencies below at a similar level. Power at") + print(" those angles never reaches the fiber in any case.") + else: + print(f" edge amplitude {edge:.1e} is small; the field has decayed.") + return report + + +# --------------------------------------------------------------------- modes +def run_forward(args): + """One simulation, then everything the analytic propagator gives for free.""" + simulation, grids, volumes, nx, frequencies = build_simulation( + args, weights=initial_weights(args, design_regions(args)[2]) + ) + volume = monitor_volume(args) + monitor = simulation.add_dft_fields([mp.Ez, mp.Hx], frequencies, where=volume) + simulation.run(until_after_sources=mp.stop_when_dft_decayed(1e-9, minimum_run_time=50)) + + stack, remaining_oxide = build_stack(args) + propagator = mpa.AngularSpectrum.from_monitor( + simulation, monitor, stack, volume, pad_factor=args.pad_factor + ) + + check_monitor(propagator, simulation, monitor) + + waist = args.fiber_waist or diffracted_waist(args) + distance = remaining_oxide + args.working_distance + fiber = mpa.gaussian_mode(waist, tilt_deg=args.tilt) + efficiency = propagator.overlap_monitor(simulation, monitor, fiber, distance) + + print(f"\nfiber waist {waist:.2f} um at {distance:.1f} um, tilt {args.tilt}deg") + for wavelength, value in zip(args.wavelengths, np.atleast_1d(efficiency)): + print(f" {wavelength*1e3:.0f} nm modal purity = {value:.4f}" + f" ({10*np.log10(max(value,1e-12)):+.2f} dB)") + + # None of the sweeps below re-runs the simulation: the fiber, the working + # distance and the superstrate are all parameters of the analytic stack. + print("\nvs fiber tilt (no further simulation)") + for tilt in np.arange(args.tilt - 4, args.tilt + 4.1, 2.0): + value = propagator.overlap_monitor( + simulation, monitor, mpa.gaussian_mode(waist, tilt_deg=float(tilt)), distance + ) + print(f" {tilt:5.1f} deg {np.mean(value):.4f}") + + print("\nvs working distance (no further simulation)") + for extra in (0.0, 5.0, 20.0, 50.0): + value = propagator.overlap_monitor( + simulation, monitor, fiber, remaining_oxide + extra + ) + print(f" {extra:5.1f} um {np.mean(value):.4f}") + + print("\nvs superstrate thickness (rebuilds only the stack)") + for thickness in (100.0, 200.0, 300.0, 500.0): + alternative = mpa.Stack([mpa.Layer(N_OXIDE, thickness - MONITOR_STANDOFF), + mpa.Layer(N_AIR)]) + other = mpa.AngularSpectrum.from_monitor( + simulation, monitor, alternative, volume, pad_factor=args.pad_factor + ) + value = other.overlap_monitor( + simulation, monitor, fiber, + thickness - MONITOR_STANDOFF + args.working_distance, + ) + print(f" {thickness:5.0f} um {np.mean(value):.4f}") + + +def run_optimize(args): + """Topology optimization of both etch layers against the fiber overlap.""" + import nlopt + + simulation, grids, volumes, nx, frequencies = build_simulation(args) + volume = monitor_volume(args) + stack, remaining_oxide = build_stack(args) + waist = args.fiber_waist or diffracted_waist(args) + distance = remaining_oxide + args.working_distance + + # get_array_metadata needs the fields allocated, and it is also the only + # way to learn how many samples Meep will actually put on the plane after + # snapping it to the grid. + simulation.init_sim() + propagator = mpa.AngularSpectrum( + stack, frequencies, + pitch=1.0 / args.resolution, + num_points=len(simulation.get_array_metadata(vol=volume)[0]), + normal=mp.Y, pad_factor=args.pad_factor, + ) + fiber = mpa.gaussian_mode(waist, tilt_deg=args.tilt) + + # The objective is written with jax.numpy; OptimizationProblem recognizes + # that from the array type it returns and differentiates it with jax.vjp. + def objective(*monitor_values): + fields = propagator.take(monitor_values) + return jnp.mean(propagator.overlap(fields, fiber, distance)) + + optimization = mpa.OptimizationProblem( + simulation=simulation, + objective_functions=objective, + objective_arguments=propagator.objective_arguments(simulation, volume), + design_regions=[mpa.DesignRegion(grid, volume=v) + for grid, v in zip(grids, volumes)], + frequencies=frequencies, + decay_by=1e-7, + ) + + weights = initial_weights(args, nx) + history = [] + + def evaluate(x, gradient): + split = np.split(x, 2) + value, gradients = optimization([split[0], split[1]]) + if gradient.size > 0: + # One gradient per design region, each (num weights, num frequencies) + # -- or (num weights,) when Meep squeezes a single frequency away. + # Reshaping rather than atleast_2d keeps the weights on the first + # axis in both cases. + gradient[:] = np.concatenate( + [np.asarray(g).reshape(nx, -1).sum(axis=1) for g in gradients] + ) + history.append(float(value)) + print(f" iteration {len(history):3d} mean efficiency = {float(value):.5f}") + return float(value) + + solver = nlopt.opt(nlopt.LD_MMA, 2 * nx) + solver.set_lower_bounds(0.0) + solver.set_upper_bounds(1.0) + solver.set_max_objective(evaluate) + solver.set_maxeval(args.iterations) + solver.optimize(np.concatenate(weights)) + + print(f"\nstarted at {history[0]:.5f}, finished at {max(history):.5f}") + + +if __name__ == "__main__": + arguments = build_arguments() + if arguments.mode == "forward": + run_forward(arguments) + else: + run_optimize(arguments) From 7e8055c3fca7627e0492d185e02a83104517dcba Mon Sep 17 00:00:00 2001 From: Alec Hammond Date: Tue, 25 Aug 2026 20:26:53 -0700 Subject: [PATCH 4/8] tests: angular-spectrum propagation, against analytic oracles 19 tests, 9 seconds for all but one, since almost everything here has a closed form to check against rather than needing a simulation. * the scattering recursion against Fresnel swept to 89 degrees, the Brewster angle, a nulled quarter-wave antireflection coating, an independently written transfer matrix, and energy conservation. The p-polarization reference carries an explicit sign, because these are the coefficients of the tangential field and that is the negative of the form usually written for the full electric vector -- the two agree on |r| and on Brewster, so a mismatch here is easy to wave away as a convention and not notice. * the up/down split, asserted against the edge amplitude rather than against a fixed tolerance: whatever field survives at the ends of the monitor wraps around and reappears as spurious down-going content, so the two track each other, and that is what makes the diagnostic meaningful. * Gaussian spreading against w0 sqrt(1 + (z/zR)^2) to six places at 50 and 200 um, with the padded window sized to the spread beam. * self-overlap of exactly one, and exact recovery of a launch angle, which the closed-form mode spectrum makes possible -- a sampled mode would quantize it. * a gradient that stays finite when a wavevector lands exactly on the light line. The window is chosen so that one does. Without the regularizer the value is still perfect and every gradient is NaN, which is the failure worth having a test for. * agreement with Meep's own near2far in a homogeneous medium, 1.6e-2 at resolution 20 and converging by better than half at 40. The convergence is the assertion that matters, since a wrong pitch or coordinate origin gives a fixed error; the test also asserts the edge amplitude is small first, because otherwise the comparison means nothing. Writing these turned up a real gap between the module and its own documentation: layer thicknesses were coerced with float() in the constructor, so a traced thickness raised ConcretizationTypeError and the stack was not in fact differentiable, as both the module docstring and the example claimed. Thicknesses are now kept as they arrive, and the concreteness checks that guard the distance argument skip themselves rather than force a value. --- python/Makefile.am | 2 + python/adjoint/angular_spectrum.py | 96 ++-- .../grating_coupler_asm.py | 167 ++++-- python/tests/test_angular_spectrum.py | 518 ++++++++++++++++++ 4 files changed, 682 insertions(+), 101 deletions(-) create mode 100644 python/tests/test_angular_spectrum.py diff --git a/python/Makefile.am b/python/Makefile.am index bf3323531..ab75fbb33 100644 --- a/python/Makefile.am +++ b/python/Makefile.am @@ -36,6 +36,7 @@ ADJOINT_TESTS = \ $(TEST_DIR)/test_adjoint_cyl.py \ $(TEST_DIR)/test_adjoint_symmetry.py \ $(TEST_DIR)/test_adjoint_protocol.py \ + $(TEST_DIR)/test_angular_spectrum.py \ $(TEST_DIR)/test_adjoint_jax.py TESTS = \ @@ -45,6 +46,7 @@ TESTS = \ $(TEST_DIR)/test_adjoint_chunks.py \ $(TEST_DIR)/test_adjoint_symmetric_grids.py \ $(TEST_DIR)/test_adjoint_protocol.py \ + $(TEST_DIR)/test_angular_spectrum.py \ $(TEST_DIR)/test_antenna_radiation.py \ $(TEST_DIR)/test_array_metadata.py \ $(TEST_DIR)/test_bend_flux.py \ diff --git a/python/adjoint/angular_spectrum.py b/python/adjoint/angular_spectrum.py index 5833f8c77..9de36ddb6 100644 --- a/python/adjoint/angular_spectrum.py +++ b/python/adjoint/angular_spectrum.py @@ -78,9 +78,7 @@ class Layer(NamedTuple): def from_medium(cls, medium: mp.Medium, thickness: Optional[float] = None): """Builds a layer from a `meep.Medium`, evaluating its index per frequency.""" return cls( - index=lambda frequency: onp.sqrt( - complex(medium.epsilon(frequency)[0][0]) - ), + index=lambda frequency: onp.sqrt(complex(medium.epsilon(frequency)[0][0])), thickness=thickness, ) @@ -139,6 +137,18 @@ class TangentialFields(NamedTuple): sign: int = 1 +def _as_concrete(value) -> Optional[float]: + """Returns `value` as a float, or None if it is a JAX tracer. + + Distances and thicknesses may be traced, so checks on them have to be + skipped rather than forced. + """ + try: + return float(value) + except (TypeError, jax.errors.ConcretizationTypeError): + return None + + def _safe_sqrt(argument: jnp.ndarray) -> jnp.ndarray: """Square root with the branch chosen so that the imaginary part is >= 0. @@ -209,9 +219,7 @@ def _stack_transmission( # Start from the terminating interface and recurse toward the monitor. At # each step `reflection` is the reflection looking into the remaining stack # and `transmission` the accumulated transmission through it. - reflection, transmission = _interface_coefficients( - admittances[-2], admittances[-1] - ) + reflection, transmission = _interface_coefficients(admittances[-2], admittances[-1]) for j in range(len(admittances) - 3, -1, -1): phase = jnp.exp(1j * wavevectors[j + 1] * thicknesses[j + 1]) interface_r, interface_t = _interface_coefficients( @@ -221,9 +229,7 @@ def _stack_transmission( # accumulated behind it. backward_r = -interface_r # reflection of the interface from the far side denominator = 1 - backward_r * reflection * phase**2 - transmission = ( - interface_t * phase * transmission / denominator - ) + transmission = interface_t * phase * transmission / denominator reflection = interface_r + ( interface_t * (2 - interface_t) * reflection * phase**2 / denominator ) @@ -234,9 +240,7 @@ def _single_interface_limit( admittances: Sequence[jnp.ndarray], ) -> Tuple[jnp.ndarray, jnp.ndarray]: """Transmission and reflection when the stack is a single interface.""" - reflection, transmission = _interface_coefficients( - admittances[0], admittances[1] - ) + reflection, transmission = _interface_coefficients(admittances[0], admittances[1]) return transmission, reflection @@ -430,9 +434,7 @@ def __init__( self._k0 = 2 * onp.pi * self.frequencies if kt is None: padded = self.num_points * self.pad_factor - self._kt = jnp.asarray( - 2 * onp.pi * onp.fft.fftfreq(padded, d=self.pitch) - ) + self._kt = jnp.asarray(2 * onp.pi * onp.fft.fftfreq(padded, d=self.pitch)) self._uniform = True self._padded = padded else: @@ -454,8 +456,11 @@ def __init__( _admittances(index, kz, self._k0) for index, kz in zip(self._indices, self._wavevectors) ] + # Deliberately not coerced to float: a thickness may be a JAX tracer, so + # that an antireflection coating or a cladding depth can be optimized + # alongside the design. self._thicknesses = [ - 0.0 if layer.thickness is None else float(layer.thickness) + 0.0 if layer.thickness is None else layer.thickness for layer in stack.layers ] @@ -465,9 +470,15 @@ def kt(self) -> jnp.ndarray: return self._kt @property - def stack_thickness(self) -> float: - """The total thickness of the stack, excluding the semi-infinite layer.""" - return float(sum(self._thicknesses[:-1])) + def stack_thickness(self): + """The total thickness of the stack, excluding the semi-infinite layer. + + Not necessarily a float: layer thicknesses may be JAX values. + """ + total = 0.0 + for thickness in self._thicknesses[:-1]: + total = total + thickness + return total def coordinates(self) -> onp.ndarray: """The transverse coordinates of the monitor samples, centered on zero.""" @@ -496,9 +507,7 @@ def _polarization_terms(self, fields: TangentialFields): f"The fields are on a plane normal to {fields.normal} but this " f"propagator was built for {self.normal}." ) - for e_component, h_component, sign, polarization in _DECOMPOSITION[ - self.normal - ]: + for e_component, h_component, sign, polarization in _DECOMPOSITION[self.normal]: e_values = fields.E.get(e_component) h_values = fields.H.get(h_component) if e_values is None and h_values is None: @@ -521,9 +530,7 @@ def decompose(self, fields: TangentialFields): """ up, down = {}, {} admittance = self._admittances[0] - for polarization, e_values, h_values, sign in self._polarization_terms( - fields - ): + for polarization, e_values, h_values, sign in self._polarization_terms(fields): e_spectrum = self._transform(e_values) h_spectrum = self._transform(h_values) # `sign` already carries the orientation of n_hat x H_t; `fields.sign` @@ -533,8 +540,7 @@ def decompose(self, fields: TangentialFields): down[polarization] = 0.5 * (e_spectrum - scaled) if not up: raise ValueError( - "No tangential field components were supplied; nothing to " - "propagate." + "No tangential field components were supplied; nothing to " "propagate." ) return up, down @@ -543,9 +549,7 @@ def _transmission(self, polarization: int): admittances = [pair[polarization] for pair in self._admittances] if len(admittances) == 2: return _single_interface_limit(admittances) - return _stack_transmission( - admittances, self._wavevectors, self._thicknesses - ) + return _stack_transmission(admittances, self._wavevectors, self._thicknesses) def spectrum(self, fields: TangentialFields, distance: float): """The up-going spectrum carried to a plane `distance` from the monitor. @@ -560,11 +564,12 @@ def spectrum(self, fields: TangentialFields, distance: float): A `PropagationResult`. """ remaining = distance - self.stack_thickness - if onp.any(onp.asarray(remaining) < 0): + 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"{self.stack_thickness} thick. The target plane has to lie in " - "the semi-infinite layer." + f"{_as_concrete(self.stack_thickness)} thick. The target plane " + "has to lie in the semi-infinite layer." ) up, _ = self.decompose(fields) amplitudes = {} @@ -572,9 +577,7 @@ def spectrum(self, fields: TangentialFields, distance: float): transmission, _ = self._transmission(polarization) # Monitor to the first interface, through the stack, then onward in # the terminating layer. - to_interface = jnp.exp( - 1j * self._wavevectors[0] * self._thicknesses[0] - ) + to_interface = jnp.exp(1j * self._wavevectors[0] * self._thicknesses[0]) beyond = jnp.exp(1j * self._wavevectors[-1] * remaining) amplitudes[polarization] = amplitude * to_interface * transmission * beyond stacked = jnp.stack( @@ -622,7 +625,9 @@ def _spectral_measure(self) -> float: def power(self, fields: TangentialFields, distance: Optional[float] = None): """Outgoing power through the target plane, one value per frequency.""" - result = self.spectrum(fields, self.stack_thickness if distance is None else distance) + result = self.spectrum( + fields, self.stack_thickness if distance is None else distance + ) weights = self._weights(result) return jnp.sum(weights * jnp.abs(result.amplitudes) ** 2, axis=(1, 2)) @@ -661,9 +666,7 @@ def overlap( mode_norm = jnp.sum(weights * jnp.abs(mode_amplitude) ** 2, axis=-1) coupled = jnp.abs(cross) ** 2 / mode_norm if incident_power is None: - incident_power = jnp.sum( - weights * jnp.abs(field_amplitude) ** 2, axis=-1 - ) + incident_power = jnp.sum(weights * jnp.abs(field_amplitude) ** 2, axis=-1) return coupled / incident_power def report(self, fields: TangentialFields) -> Dict[str, jnp.ndarray]: @@ -744,9 +747,7 @@ def propagate( if not jnp.any(result.amplitudes[..., polarization]): continue outputs[e_component] = ( - jnp.einsum( - "xk,fk->fx", phase, result.amplitudes[..., polarization] - ) + jnp.einsum("xk,fk->fx", phase, result.amplitudes[..., polarization]) * measure ) return outputs @@ -792,9 +793,7 @@ def _assert_homogeneous( mirrors the check `dft_near2far` makes for the same reason, and catches a monitor accidentally clipping a waveguide or a PML. """ - epsilon = onp.asarray( - simulation.get_array(vol=volume, component=mp.Dielectric) - ) + epsilon = onp.asarray(simulation.get_array(vol=volume, component=mp.Dielectric)) spread = float(onp.max(epsilon) - onp.min(epsilon)) if spread > tolerance * max(1.0, float(onp.max(epsilon))): raise ValueError( @@ -907,10 +906,7 @@ def power_monitor(self, simulation, monitor, distance=None): def report_monitor(self, simulation, monitor) -> Dict[str, onp.ndarray]: """`report`, reading from a monitor and returning NumPy arrays.""" fields = self.fields_from_monitor(simulation, monitor) - return { - key: onp.asarray(value) - for key, value in self.report(fields).items() - } + return {key: onp.asarray(value) for key, value in self.report(fields).items()} 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.py b/python/examples/adjoint_optimization/grating_coupler_asm.py index 4898af6d3..03616956e 100644 --- a/python/examples/adjoint_optimization/grating_coupler_asm.py +++ b/python/examples/adjoint_optimization/grating_coupler_asm.py @@ -60,27 +60,48 @@ def build_arguments(): - parser = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) parser.add_argument("mode", choices=["forward", "optimize"]) - parser.add_argument("--resolution", type=float, default=20.0, - help="pixels per micron") - parser.add_argument("--aperture", type=float, default=20.0, - help="width of the grating, in microns") - 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=5.0, - help="fiber facet height above the silica/air interface") - parser.add_argument("--fiber-waist", type=float, default=None, - help="fiber 1/e field radius; defaults to the diffracted " - "beam size at the facet") - parser.add_argument("--tilt", type=float, default=8.0, - help="fiber tilt from normal, in degrees") - parser.add_argument("--wavelengths", type=float, nargs="+", - default=[1.52, 1.55, 1.58]) - parser.add_argument("--pad-factor", type=int, default=16, - help="zero padding before the transform; a beam that " - "spreads over hundreds of microns needs a wide window") + parser.add_argument( + "--resolution", type=float, default=20.0, help="pixels per micron" + ) + parser.add_argument( + "--aperture", type=float, default=20.0, help="width of the grating, in microns" + ) + 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=5.0, + help="fiber facet height above the silica/air interface", + ) + parser.add_argument( + "--fiber-waist", + type=float, + default=None, + help="fiber 1/e field radius; defaults to the diffracted " + "beam size at the facet", + ) + parser.add_argument( + "--tilt", type=float, default=8.0, help="fiber tilt from normal, in degrees" + ) + parser.add_argument( + "--wavelengths", type=float, nargs="+", default=[1.52, 1.55, 1.58] + ) + parser.add_argument( + "--pad-factor", + type=int, + default=16, + help="zero padding before the transform; a beam that " + "spreads over hundreds of microns needs a wide window", + ) parser.add_argument("--iterations", type=int, default=15) parser.add_argument("--design-resolution", type=float, default=None) return parser.parse_args() @@ -114,12 +135,20 @@ def design_regions(args): grids, volumes = [], [] for y_lo, y_hi in ((0.0, T_DEEP), (T_DEEP, T_DEVICE)): grids.append( - mp.MaterialGrid(mp.Vector3(nx, 1, 1), oxide, silicon, - weights=np.ones((nx,)), do_averaging=False, beta=0) + mp.MaterialGrid( + mp.Vector3(nx, 1, 1), + oxide, + silicon, + weights=np.ones((nx,)), + do_averaging=False, + beta=0, + ) ) volumes.append( - mp.Volume(center=mp.Vector3(0, 0.5 * (y_lo + y_hi)), - size=mp.Vector3(args.aperture, y_hi - y_lo)) + mp.Volume( + center=mp.Vector3(0, 0.5 * (y_lo + y_hi)), + size=mp.Vector3(args.aperture, y_hi - y_lo), + ) ) return grids, volumes, nx @@ -135,13 +164,21 @@ def build_simulation(args, weights=None): silicon, oxide = mp.Medium(index=N_SI), mp.Medium(index=N_OXIDE) geometry = [ # silicon handle below the buried oxide - mp.Block(center=mp.Vector3(center.x, -(T_BOX + T_HANDLE + DPML) / 2 - T_BOX / 2), - size=mp.Vector3(mp.inf, T_HANDLE + DPML), material=silicon), + mp.Block( + center=mp.Vector3(center.x, -(T_BOX + T_HANDLE + DPML) / 2 - T_BOX / 2), + size=mp.Vector3(mp.inf, T_HANDLE + DPML), + material=silicon, + ), # the input waveguide, running in from the left - mp.Block(center=mp.Vector3(-(args.aperture / 2 + WAVEGUIDE_LENGTH + PAD + DPML) / 2 - - args.aperture / 4, T_DEVICE / 2), - size=mp.Vector3(2 * (WAVEGUIDE_LENGTH + PAD + DPML), T_DEVICE), - material=silicon), + mp.Block( + center=mp.Vector3( + -(args.aperture / 2 + WAVEGUIDE_LENGTH + PAD + DPML) / 2 + - args.aperture / 4, + T_DEVICE / 2, + ), + size=mp.Vector3(2 * (WAVEGUIDE_LENGTH + PAD + DPML), T_DEVICE), + material=silicon, + ), ] + [ mp.Block(center=volume.center, size=volume.size, material=grid) for grid, volume in zip(grids, volumes) @@ -152,16 +189,25 @@ def build_simulation(args, weights=None): source = [ mp.EigenModeSource( mp.GaussianSource(center_frequency, fwidth=0.2 * center_frequency), - center=mp.Vector3(-(args.aperture / 2 + WAVEGUIDE_LENGTH * 0.6), T_DEVICE / 2), + center=mp.Vector3( + -(args.aperture / 2 + WAVEGUIDE_LENGTH * 0.6), T_DEVICE / 2 + ), size=mp.Vector3(0, 6 * T_DEVICE), - eig_band=1, eig_parity=mp.ODD_Z, eig_match_freq=True, + eig_band=1, + eig_parity=mp.ODD_Z, + eig_match_freq=True, ) ] simulation = mp.Simulation( - cell_size=size, geometry_center=center, resolution=args.resolution, - boundary_layers=[mp.PML(DPML)], geometry=geometry, sources=source, - default_material=oxide, dimensions=2, + cell_size=size, + geometry_center=center, + resolution=args.resolution, + boundary_layers=[mp.PML(DPML)], + geometry=geometry, + sources=source, + default_material=oxide, + dimensions=2, ) return simulation, grids, volumes, nx, frequencies @@ -170,8 +216,10 @@ def monitor_volume(args): """The plane the radiated field is read on, inside the homogeneous oxide.""" size, center = cell_geometry(args) width = size.x - 2 * DPML - 2 * MONITOR_MARGIN - return mp.Volume(center=mp.Vector3(center.x, T_DEVICE + MONITOR_STANDOFF), - size=mp.Vector3(width, 0)) + return mp.Volume( + center=mp.Vector3(center.x, T_DEVICE + MONITOR_STANDOFF), + size=mp.Vector3(width, 0), + ) def build_stack(args): @@ -182,7 +230,10 @@ def build_stack(args): f"--superstrate {args.superstrate} must exceed the monitor standoff " f"{MONITOR_STANDOFF}." ) - return mpa.Stack([mpa.Layer(N_OXIDE, remaining_oxide), mpa.Layer(N_AIR)]), remaining_oxide + return ( + mpa.Stack([mpa.Layer(N_OXIDE, remaining_oxide), mpa.Layer(N_AIR)]), + remaining_oxide, + ) def diffracted_waist(args): @@ -202,8 +253,9 @@ def diffracted_waist(args): def initial_weights(args, nx): """A uniform grating, as a starting point for the optimizer.""" x = np.linspace(-args.aperture / 2, args.aperture / 2, nx) - period = float(np.mean(args.wavelengths)) / (N_OXIDE * math.sin(math.radians(args.tilt)) - + 2.6) + period = float(np.mean(args.wavelengths)) / ( + N_OXIDE * math.sin(math.radians(args.tilt)) + 2.6 + ) # A binary grating, not a smoothly graded one: a sinusoidal index # modulation is a much weaker scatterer and radiates very little. deep = (np.cos(2 * math.pi * x / period) > 0).astype(float) @@ -249,7 +301,9 @@ def run_forward(args): ) volume = monitor_volume(args) monitor = simulation.add_dft_fields([mp.Ez, mp.Hx], frequencies, where=volume) - simulation.run(until_after_sources=mp.stop_when_dft_decayed(1e-9, minimum_run_time=50)) + simulation.run( + until_after_sources=mp.stop_when_dft_decayed(1e-9, minimum_run_time=50) + ) stack, remaining_oxide = build_stack(args) propagator = mpa.AngularSpectrum.from_monitor( @@ -265,15 +319,20 @@ def run_forward(args): print(f"\nfiber waist {waist:.2f} um at {distance:.1f} um, tilt {args.tilt}deg") for wavelength, value in zip(args.wavelengths, np.atleast_1d(efficiency)): - print(f" {wavelength*1e3:.0f} nm modal purity = {value:.4f}" - f" ({10*np.log10(max(value,1e-12)):+.2f} dB)") + print( + f" {wavelength*1e3:.0f} nm modal purity = {value:.4f}" + f" ({10*np.log10(max(value,1e-12)):+.2f} dB)" + ) # None of the sweeps below re-runs the simulation: the fiber, the working # distance and the superstrate are all parameters of the analytic stack. print("\nvs fiber tilt (no further simulation)") for tilt in np.arange(args.tilt - 4, args.tilt + 4.1, 2.0): value = propagator.overlap_monitor( - simulation, monitor, mpa.gaussian_mode(waist, tilt_deg=float(tilt)), distance + simulation, + monitor, + mpa.gaussian_mode(waist, tilt_deg=float(tilt)), + distance, ) print(f" {tilt:5.1f} deg {np.mean(value):.4f}") @@ -286,13 +345,16 @@ def run_forward(args): print("\nvs superstrate thickness (rebuilds only the stack)") for thickness in (100.0, 200.0, 300.0, 500.0): - alternative = mpa.Stack([mpa.Layer(N_OXIDE, thickness - MONITOR_STANDOFF), - mpa.Layer(N_AIR)]) + alternative = mpa.Stack( + [mpa.Layer(N_OXIDE, thickness - MONITOR_STANDOFF), mpa.Layer(N_AIR)] + ) other = mpa.AngularSpectrum.from_monitor( simulation, monitor, alternative, volume, pad_factor=args.pad_factor ) value = other.overlap_monitor( - simulation, monitor, fiber, + simulation, + monitor, + fiber, thickness - MONITOR_STANDOFF + args.working_distance, ) print(f" {thickness:5.0f} um {np.mean(value):.4f}") @@ -313,10 +375,12 @@ def run_optimize(args): # snapping it to the grid. simulation.init_sim() propagator = mpa.AngularSpectrum( - stack, frequencies, + stack, + frequencies, pitch=1.0 / args.resolution, num_points=len(simulation.get_array_metadata(vol=volume)[0]), - normal=mp.Y, pad_factor=args.pad_factor, + normal=mp.Y, + pad_factor=args.pad_factor, ) fiber = mpa.gaussian_mode(waist, tilt_deg=args.tilt) @@ -330,8 +394,9 @@ def objective(*monitor_values): simulation=simulation, objective_functions=objective, objective_arguments=propagator.objective_arguments(simulation, volume), - design_regions=[mpa.DesignRegion(grid, volume=v) - for grid, v in zip(grids, volumes)], + design_regions=[ + mpa.DesignRegion(grid, volume=v) for grid, v in zip(grids, volumes) + ], frequencies=frequencies, decay_by=1e-7, ) diff --git a/python/tests/test_angular_spectrum.py b/python/tests/test_angular_spectrum.py new file mode 100644 index 000000000..19c72ca3e --- /dev/null +++ b/python/tests/test_angular_spectrum.py @@ -0,0 +1,518 @@ +"""Tests for angular-spectrum propagation through stratified media. + +Everything except `TestAgainstNearToFar` runs without an FDTD simulation, using +analytic oracles: Fresnel coefficients, the Brewster angle, a quarter-wave +antireflection coating, an independently written transfer matrix, energy +conservation, and the spreading of a Gaussian beam. +""" + +import math +import unittest + +import numpy as onp + +import meep as mp + +try: + import meep.adjoint as mpa +except ImportError: + import adjoint as mpa + +from utils import ApproxComparisonTestCase + +try: + import jax + + jax.config.update("jax_enable_x64", True) + import jax.numpy as jnp + from meep.adjoint import angular_spectrum as asm +except ImportError: + jax = None + +mp.verbosity(0) + +WAVELENGTH = 1.55 +N_OXIDE, N_AIR, N_HIGH = 1.444, 1.0, 1.9 + + +def _wavevector(index, k0, kt): + return asm._longitudinal_wavevector(jnp.array([index + 0j]), k0, kt) + + +def _admittance(index, kz, k0, polarization): + return asm._admittances(jnp.array([index + 0j]), kz, k0)[polarization] + + +@unittest.skipIf(jax is None, "jax is not installed") +class TestStackSolver(ApproxComparisonTestCase): + """The scattering recursion, against closed-form references.""" + + def setUp(self): + self.k0 = onp.array([2 * onp.pi / WAVELENGTH]) + + def _coefficients(self, indices, thicknesses, angle_deg, polarization): + kt = jnp.array([indices[0] * self.k0[0] * math.sin(math.radians(angle_deg))]) + wavevectors = [_wavevector(n, self.k0, kt) for n in indices] + admittances = [ + _admittance(n, kz, self.k0, polarization) + for n, kz in zip(indices, wavevectors) + ] + if len(indices) == 2: + return asm._single_interface_limit(admittances) + return asm._stack_transmission(admittances, wavevectors, thicknesses) + + def test_single_interface_matches_fresnel(self): + """Reflection at one interface, swept to near grazing. + + The p-polarization reference carries a sign: these are the coefficients + of the *tangential* field, which is the admittance convention, and it is + the negative of the form written for the full electric vector. + """ + for polarization, name in ( + (asm.S_POLARIZATION, "s"), + (asm.P_POLARIZATION, "p"), + ): + for angle in (0, 10, 30, 45, 60, 80, 89): + theta = math.radians(angle) + _, reflection = self._coefficients( + [N_OXIDE, N_AIR], [0.0], angle, polarization + ) + cos_in = math.cos(theta) + cos_out = onp.sqrt( + complex(1 - (N_OXIDE * math.sin(theta) / N_AIR) ** 2) + ) + if polarization == asm.S_POLARIZATION: + expected = (N_OXIDE * cos_in - N_AIR * cos_out) / ( + N_OXIDE * cos_in + N_AIR * cos_out + ) + else: + expected = -(N_AIR * cos_in - N_OXIDE * cos_out) / ( + N_AIR * cos_in + N_OXIDE * cos_out + ) + self.assertAlmostEqual( + complex(reflection[0, 0]), expected, places=12, + msg=f"{name} polarization at {angle} deg", + ) + + def test_brewster_angle(self): + """p-polarized reflection vanishes at atan(n_out / n_in).""" + angle = math.degrees(math.atan(N_AIR / N_OXIDE)) + _, reflection = self._coefficients( + [N_OXIDE, N_AIR], [0.0], angle, asm.P_POLARIZATION + ) + self.assertLess(abs(complex(reflection[0, 0])), 1e-12) + + def test_quarter_wave_antireflection_coating(self): + """An index-matched quarter wave nulls the reflection at normal incidence.""" + index = math.sqrt(N_OXIDE * N_AIR) + thickness = WAVELENGTH / (4 * index) + _, reflection = self._coefficients( + [N_OXIDE, index, N_AIR], [0.0, thickness], 0.0, asm.S_POLARIZATION + ) + self.assertLess(abs(complex(reflection[0, 0])), 1e-12) + + def test_matches_independent_transfer_matrix(self): + """Against a transfer-matrix implementation written from scratch here. + + The propagator deliberately does not use transfer matrices, since they + overflow on a thick layer or an evanescent order, so this is a genuinely + independent route to the same numbers in a regime where both are stable. + """ + indices, thicknesses = [N_OXIDE, N_HIGH, N_AIR], [0.0, 0.37] + + def transfer_matrix(angle_deg, polarization): + kt = indices[0] * self.k0[0] * math.sin(math.radians(angle_deg)) + wavevectors = [ + complex(onp.sqrt(complex((n * self.k0[0]) ** 2 - kt**2))) + for n in indices + ] + if polarization == asm.S_POLARIZATION: + admittances = [kz / self.k0[0] for kz in wavevectors] + else: + admittances = [ + n**2 * self.k0[0] / kz for n, kz in zip(indices, wavevectors) + ] + matrix = onp.eye(2, dtype=complex) + for j in range(len(indices) - 1): + r = (admittances[j] - admittances[j + 1]) / ( + admittances[j] + admittances[j + 1] + ) + t = 2 * admittances[j] / (admittances[j] + admittances[j + 1]) + matrix = matrix @ (onp.array([[1, r], [r, 1]], dtype=complex) / t) + if j + 1 < len(indices) - 1: + phase = onp.exp(1j * wavevectors[j + 1] * thicknesses[j + 1]) + matrix = matrix @ onp.array( + [[1 / phase, 0], [0, phase]], dtype=complex + ) + return 1 / matrix[0, 0], matrix[1, 0] / matrix[0, 0] + + for polarization in (asm.S_POLARIZATION, asm.P_POLARIZATION): + for angle in (0, 25, 55): + transmission, reflection = self._coefficients( + indices, thicknesses, angle, polarization + ) + expected_t, expected_r = transfer_matrix(angle, polarization) + self.assertAlmostEqual( + complex(transmission[0, 0]), expected_t, places=12 + ) + self.assertAlmostEqual( + complex(reflection[0, 0]), expected_r, places=12 + ) + + def test_energy_is_conserved(self): + """Reflected plus transmitted power equals the incident, losslessly.""" + indices, thicknesses = [N_OXIDE, N_HIGH, N_AIR], [0.0, 0.37] + for polarization in (asm.S_POLARIZATION, asm.P_POLARIZATION): + for angle in (0, 20, 40): + transmission, reflection = self._coefficients( + indices, thicknesses, angle, polarization + ) + kt = jnp.array( + [indices[0] * self.k0[0] * math.sin(math.radians(angle))] + ) + y_in = _admittance( + indices[0], _wavevector(indices[0], self.k0, kt), self.k0, + polarization, + ) + y_out = _admittance( + indices[-1], _wavevector(indices[-1], self.k0, kt), self.k0, + polarization, + ) + reflected = abs(complex(reflection[0, 0])) ** 2 + transmitted = abs(complex(transmission[0, 0])) ** 2 * float( + onp.real(complex(y_out[0, 0])) / onp.real(complex(y_in[0, 0])) + ) + self.assertAlmostEqual(reflected + transmitted, 1.0, places=12) + + +def _uniform_propagator(index=N_OXIDE, num_points=512, pitch=0.05, pad_factor=8): + """A propagator with no interface, i.e. plain homogeneous propagation.""" + stack = mpa.Stack([mpa.Layer(index, 0.0), mpa.Layer(index)]) + return mpa.AngularSpectrum( + stack, [1 / WAVELENGTH], pitch, num_points, normal=mp.Y, + pad_factor=pad_factor, + ) + + +def _up_going(propagator, values): + """Builds the magnetic partner that makes `values` purely up-going. + + The relation is per wavevector, so the partner has to be constructed in the + spectral domain; scaling by the admittance of one launch angle in real space + leaves spurious down-going content behind. + """ + spectrum = propagator._transform(jnp.asarray(values)) + admittance = propagator._admittances[0][asm.S_POLARIZATION] + coordinates = propagator.coordinates() + magnetic = ( + jnp.einsum( + "xk,fk->fx", + jnp.exp(1j * propagator.kt[None, :] * coordinates[:, None]), + spectrum * admittance, + ) + * propagator._spectral_measure() + ) + return mpa.TangentialFields( + E={mp.Ez: jnp.asarray(values)}, H={mp.Hx: magnetic}, normal=mp.Y + ) + + +@unittest.skipIf(jax is None, "jax is not installed") +class TestDecomposition(ApproxComparisonTestCase): + """Separating up- from down-going radiation.""" + + def test_residual_tracks_edge_truncation(self): + """A purely up-going field yields no down-going part, up to truncation. + + The transform is periodic, so whatever amplitude survives at the ends of + the monitor wraps around and appears as spurious down-going content. The + two should therefore track each other, which is what makes the + `edge_amplitude` diagnostic meaningful. + """ + propagator = _uniform_propagator() + x = propagator.coordinates() + for width in (5.0, 3.0): + values = (onp.exp(-((x / width) ** 2)))[None, :] + up, down = propagator.decompose(_up_going(propagator, values)) + residual = float( + jnp.max(jnp.abs(down[asm.S_POLARIZATION])) + / jnp.max(jnp.abs(up[asm.S_POLARIZATION])) + ) + edge = float(onp.exp(-((x.max() / width) ** 2))) + self.assertLess(residual, max(10 * edge, 1e-9), f"width {width}") + + def test_down_going_field_is_recognized(self): + """Flipping the magnetic field flips which branch the power lands in.""" + propagator = _uniform_propagator() + x = propagator.coordinates() + values = (onp.exp(-((x / 3.0) ** 2)))[None, :] + fields = _up_going(propagator, values) + flipped = mpa.TangentialFields( + E=fields.E, H={mp.Hx: -fields.H[mp.Hx]}, normal=mp.Y + ) + up, down = propagator.decompose(flipped) + self.assertLess( + float( + jnp.max(jnp.abs(up[asm.S_POLARIZATION])) + / jnp.max(jnp.abs(down[asm.S_POLARIZATION])) + ), + 1e-6, + ) + + +@unittest.skipIf(jax is None, "jax is not installed") +class TestPropagation(ApproxComparisonTestCase): + """Propagation against analytic beam physics.""" + + def test_round_trip_at_zero_distance(self): + propagator = _uniform_propagator() + x = propagator.coordinates() + values = (onp.exp(-((x / 3.0) ** 2)))[None, :] + output = propagator.propagate(_up_going(propagator, values), 0.0)[mp.Ez] + self.assertLess(float(jnp.max(jnp.abs(output - values))), 1e-8) + + def test_gaussian_beam_spreading(self): + """The 1/e width follows w0 sqrt(1 + (z/zR)^2) exactly. + + The padded window has to be wide enough for the spread beam: the + transform is periodic, so a beam wider than the window wraps onto itself + and the answer is nonsense rather than merely inaccurate. + """ + waist = 3.0 + rayleigh = math.pi * waist**2 * N_OXIDE / WAVELENGTH + for distance, pad_factor in ((50.0, 16), (200.0, 32)): + propagator = _uniform_propagator(pad_factor=pad_factor) + x = propagator.coordinates() + values = (onp.exp(-((x / waist) ** 2)))[None, :] + expected = waist * math.sqrt(1 + (distance / rayleigh) ** 2) + samples = onp.linspace(-4 * expected, 4 * expected, 3001) + profile = onp.abs( + onp.asarray( + propagator.propagate( + _up_going(propagator, values), distance, + coordinates=samples, + )[mp.Ez][0] + ) + ) + above = samples[profile >= profile.max() / onp.e] + measured = (above.max() - above.min()) / 2 + self.assertAlmostEqual(measured / expected, 1.0, places=6) + + def test_refuses_to_stop_short_of_the_stack(self): + stack = mpa.Stack([mpa.Layer(N_OXIDE, 5.0), mpa.Layer(N_AIR)]) + propagator = mpa.AngularSpectrum( + stack, [1 / WAVELENGTH], 0.05, 128, normal=mp.Y + ) + x = propagator.coordinates() + fields = _up_going(propagator, (onp.exp(-((x / 1.0) ** 2)))[None, :]) + with self.assertRaisesRegex(ValueError, "does not clear the stack"): + propagator.spectrum(fields, 2.0) + + +@unittest.skipIf(jax is None, "jax is not installed") +class TestOverlap(ApproxComparisonTestCase): + """Projection onto a target mode.""" + + def test_self_overlap_is_unity(self): + propagator = _uniform_propagator() + x = propagator.coordinates() + fields = _up_going(propagator, (onp.exp(-((x / 3.0) ** 2)))[None, :]) + result = propagator.spectrum(fields, 100.0) + itself = mpa.Mode( + spectrum=lambda kt, k0, index: result.amplitudes[..., asm.S_POLARIZATION] + ) + self.assertAlmostEqual( + float(propagator.overlap(fields, itself, 100.0)[0]), 1.0, places=10 + ) + + def test_recovers_the_launch_angle(self): + """A beam launched at an angle matches a mode tilted to that angle. + + The mode spectrum is written in closed form, so the tilt is a continuous + parameter rather than something quantized by the monitor grid. + """ + propagator = _uniform_propagator() + x = propagator.coordinates() + k0 = 2 * onp.pi / WAVELENGTH + for angle in (0.0, 8.0, 20.0): + transverse = N_OXIDE * k0 * math.sin(math.radians(angle)) + values = (onp.exp(1j * transverse * x) * onp.exp(-((x / 20.0) ** 2)))[ + None, : + ] + fields = _up_going(propagator, values) + candidates = onp.linspace(angle - 4, angle + 4, 33) + best = max( + candidates, + key=lambda t: float( + propagator.overlap( + fields, mpa.gaussian_mode(20.0, tilt_deg=float(t)), 1e-9 + )[0] + ), + ) + self.assertAlmostEqual(float(best), angle, places=6) + + +@unittest.skipIf(jax is None, "jax is not installed") +class TestGradients(ApproxComparisonTestCase): + """Differentiability, including the failure that motivates the regularizer.""" + + def test_gradient_is_finite_on_the_light_line(self): + """A wavevector landing exactly on the light line must not poison the VJP. + + kz vanishes there and its derivative is unbounded, so an unregularized + sqrt evaluates to zero happily and returns an infinite cotangent. The + objective value looks perfect while every gradient is NaN. The window + below is chosen so that a grid point lands exactly on n k0. + """ + num_points, pad_factor = 64, 1 + # Place a sample exactly on the light line: kt = 2 pi m / (N dx) equals + # n k0 = 2 pi n / lambda when N dx = m lambda / n. + pitch = 8 * WAVELENGTH / N_OXIDE / num_points + propagator = _uniform_propagator( + num_points=num_points, pitch=pitch, pad_factor=pad_factor + ) + on_light_line = onp.min( + onp.abs(onp.asarray(propagator.kt) - N_OXIDE * 2 * onp.pi / WAVELENGTH) + ) + self.assertLess(on_light_line, 1e-9, "the test window missed the light line") + + x = propagator.coordinates() + values = jnp.asarray((onp.exp(-((x / 0.5) ** 2)))[None, :]) + + def objective(scale): + fields = _up_going(propagator, scale * values) + return jnp.real( + jnp.sum(propagator.spectrum(fields, 1.0).amplitudes) + ) + + gradient = jax.grad(objective)(1.0) + self.assertTrue(onp.isfinite(float(gradient)), "gradient is not finite") + + def test_differentiable_in_stack_and_mode_parameters(self): + """Layer thickness and fiber tilt carry gradients, not just the fields.""" + x = onp.linspace(-6.4, 6.4, 256) + values = (onp.exp(-((x / 2.0) ** 2)))[None, :] + + def objective(thickness, tilt): + stack = mpa.Stack([mpa.Layer(N_OXIDE, thickness), mpa.Layer(N_AIR)]) + propagator = mpa.AngularSpectrum( + stack, [1 / WAVELENGTH], 0.05, 256, normal=mp.Y, pad_factor=8 + ) + fields = _up_going(propagator, values) + return propagator.overlap( + fields, mpa.gaussian_mode(4.0, tilt_deg=tilt), thickness + 10.0 + )[0] + + gradients = jax.grad(objective, argnums=(0, 1))(2.0, 5.0) + for name, value in zip(("thickness", "tilt"), gradients): + self.assertTrue(onp.isfinite(float(value)), name) + self.assertNotEqual(float(value), 0.0, name) + + +@unittest.skipIf(jax is None, "jax is not installed") +class TestValidation(unittest.TestCase): + """The constructor rejects stacks and monitors that cannot work.""" + + def test_stack_must_terminate(self): + with self.assertRaisesRegex(ValueError, "semi-infinite"): + mpa.Stack([mpa.Layer(N_OXIDE, 1.0), mpa.Layer(N_AIR, 1.0)]).validate() + + def test_only_the_last_layer_may_be_semi_infinite(self): + with self.assertRaisesRegex(ValueError, "only the final layer"): + mpa.Stack( + [mpa.Layer(N_OXIDE), mpa.Layer(N_HIGH, 1.0), mpa.Layer(N_AIR)] + ).validate() + + def test_both_tangential_fields_are_required(self): + """One field alone cannot distinguish up-going from down-going.""" + propagator = _uniform_propagator(num_points=64) + fields = mpa.TangentialFields( + E={mp.Ez: onp.ones((1, 64))}, H={}, normal=mp.Y + ) + with self.assertRaisesRegex(ValueError, "both"): + propagator.decompose(fields) + + def test_three_dimensions_is_rejected_clearly(self): + with self.assertRaisesRegex(NotImplementedError, "2D"): + asm._tangential_components(mp.Z, 3) + + +@unittest.skipIf(jax is None, "jax is not installed") +class TestAgainstNearToFar(ApproxComparisonTestCase): + """Against Meep's own near-to-far transformation, in a homogeneous medium. + + This is the end-to-end check on the Meep adapter: sample pitch, coordinate + origin and the phase ramp that comes with it. What matters is that the + disagreement falls with resolution -- a mistake in any of those produces a + fixed error instead. + """ + + def _run(self, resolution): + cell_x, cell_y, pml = 40.0, 8.0, 1.0 + frequency = 1 / WAVELENGTH + source = [ + mp.Source( + mp.GaussianSource(frequency, fwidth=0.1 * frequency), + component=mp.Ez, center=mp.Vector3(0, -1.0), + size=mp.Vector3(12.0, 0), + # Narrow and apodized, so the field has decayed by the ends of + # the monitor. A bare dipole never does, and near2far tolerates + # that while a periodic transform does not. + amp_func=lambda p: onp.exp(-((p.x / 1.5) ** 2)), + ) + ] + simulation = mp.Simulation( + resolution=resolution, cell_size=mp.Vector3(cell_x, cell_y), + boundary_layers=[mp.PML(pml)], sources=source, + default_material=mp.Medium(index=N_OXIDE), + ) + height, width = 1.0, cell_x - 2 * pml - 0.4 + volume = mp.Volume( + center=mp.Vector3(0, height), size=mp.Vector3(width, 0) + ) + monitor = simulation.add_dft_fields( + [mp.Ez, mp.Hx], [frequency], where=volume + ) + near2far = simulation.add_near2far( + [frequency], + mp.Near2FarRegion( + center=mp.Vector3(0, height), size=mp.Vector3(width, 0), weight=+1 + ), + ) + simulation.run( + until_after_sources=mp.stop_when_dft_decayed(1e-10, minimum_run_time=30) + ) + + stack = mpa.Stack([mpa.Layer(N_OXIDE, 0.0), mpa.Layer(N_OXIDE)]) + propagator = mpa.AngularSpectrum.from_monitor( + simulation, monitor, stack, volume, pad_factor=16 + ) + distance = 25.0 + samples = onp.linspace(-8, 8, 33) + got = propagator.propagate_monitor( + simulation, monitor, distance, coordinates=samples + )[mp.Ez][0] + want = onp.array( + [ + simulation.get_farfield(near2far, mp.Vector3(x, height + distance))[2] + for x in samples + ] + ) + report = propagator.report_monitor(simulation, monitor) + return ( + onp.abs(got - want).max() / onp.abs(want).max(), + float(report["edge_amplitude"][0]), + ) + + def test_agrees_and_converges(self): + coarse, coarse_edge = self._run(20) + fine, _ = self._run(40) + # If the field had not decayed at the monitor edges the comparison would + # be meaningless, so assert that first. + self.assertLess(coarse_edge, 1e-6) + self.assertLess(coarse, 3e-2) + self.assertLess(fine, coarse / 2, "did not converge with resolution") + + +if __name__ == "__main__": + unittest.main() From 7cc5992c27dda2f68aae9235e68b4f590fa02668 Mon Sep 17 00:00:00 2001 From: Alec Hammond Date: Tue, 25 Aug 2026 20:27:29 -0700 Subject: [PATCH 5/8] doc: angular-spectrum propagation through stratified media Covers what it is for -- radiation crossing a material interface, which near2far cannot handle -- the forward-only and adjoint entry points, and the two constraints that produce quiet wrong answers rather than errors: a monitor whose field has not decayed by its ends, and a padded window narrower than the spread beam. Also states plainly that a few percent of edge amplitude is normal for a grating radiating into a cladding, is near-grazing radiation rather than the beam tail, and bounds the accuracy accordingly. Notes what is not supported and why: 3D needs the s/p rotation by azimuth with its removable singularity at normal incidence and raises rather than guessing, and both half-spaces are covered with two monitors rather than a closed surface, since a single plane is already complete for the half-space above it. --- NEWS.md | 7 ++ doc/docs/Python_Tutorials/Adjoint_Solver.md | 75 +++++++++++++++++++++ 2 files changed, 82 insertions(+) diff --git a/NEWS.md b/NEWS.md index 9acfa67bd..03f6c1d2a 100644 --- a/NEWS.md +++ b/NEWS.md @@ -2,6 +2,13 @@ ## Meep 1.35.0 (in progress) +* Adjoint solver: `meep.adjoint.AngularSpectrum` propagates the tangential DFT + fields on a planar monitor through an arbitrary stratified medium + analytically, in JAX. Unlike `add_near2far`, which requires a homogeneous + medium, this handles radiation crossing a material interface, so the layers + above the monitor can leave the FDTD cell entirely and become differentiable + parameters. Usable with or without the adjoint solver. 2D only for now. + * Adjoint solver: objective functions are now differentiated with a single vector-Jacobian product instead of a full frequency Jacobian per objective argument. For an objective of $M$ arguments at $F$ frequencies this reduces diff --git a/doc/docs/Python_Tutorials/Adjoint_Solver.md b/doc/docs/Python_Tutorials/Adjoint_Solver.md index c486b7530..d1c3b9f20 100644 --- a/doc/docs/Python_Tutorials/Adjoint_Solver.md +++ b/doc/docs/Python_Tutorials/Adjoint_Solver.md @@ -273,6 +273,81 @@ JAX is an optional dependency throughout. If it is not installed, JAX objective functions are simply not recognized and `mpa.MeepJaxWrapper` and `mpa.value_and_jacobian` are absent; everything else works unchanged. +Propagating Through Stratified Media +------------------------------------ + +Meep's [near-to-far transformation](Near_to_Far_Field_Spectra.md) requires its +surface to sit in a homogeneous medium — `dft_near2far` aborts otherwise — so a +structure whose radiation crosses a material interface has to keep that interface +inside the FDTD cell. A grating coupler radiating up through a cladding, across +the chip surface, and hundreds of microns to a fiber cannot afford to. + +`meep.adjoint.AngularSpectrum` propagates the tangential DFT fields on a planar +monitor through an arbitrary layer stack analytically instead. The layers above +the monitor leave the simulation, and become differentiable parameters: + +```py +stack = mpa.Stack([mpa.Layer(index=1.444, thickness=300.0), # silica superstrate + mpa.Layer(index=1.0)]) # air, semi-infinite + +monitor = sim.add_dft_fields([mp.Ez, mp.Hx], frequencies, where=plane) +sim.run(...) + +propagator = mpa.AngularSpectrum.from_monitor(sim, monitor, stack, plane) +fiber = mpa.gaussian_mode(waist=5.2, tilt_deg=8.0) +efficiency = propagator.overlap_monitor(sim, monitor, fiber, distance=305.0) +``` + +That path is plain NumPy in and out, so it can post-process an ordinary forward +run without any optimization. Inside an objective function, use +`objective_arguments()` and `take()` instead, which give the `FourierFields` the +adjoint solver needs. + +Both tangential fields are required. Together they determine the up- and +down-going plane-wave amplitudes separately, which an open near-to-far surface +cannot do — it has no way to reject radiation heading the wrong way. + +### Reading the diagnostics before believing a number + +The transverse transform is periodic, so a field that has not decayed by the ends +of the monitor wraps around, and it does so quietly: the result looks plausible +and does not improve with resolution. `report()` measures the three ways a +monitor is usually placed wrongly: + +```py +propagator.report_monitor(sim, monitor) +# {'downgoing_fraction': ..., is something above the monitor scattering? +# 'evanescent_fraction': ..., is the monitor inside the near field? +# 'edge_amplitude': ...} has the field decayed by the monitor's ends? +``` + +A few percent of `edge_amplitude` is common for a grating radiating into a +cladding and is not the beam tail — it is near-grazing radiation, which travels +sideways rather than decaying, so widening the cell barely helps. It bounds the +accuracy at a similar level. Power at those angles was never going to reach the +fiber. + +The other constraint is the padded window. A beam that spreads over hundreds of +microns needs `pad_factor` large enough that the spread beam still fits; too small +and the beam wraps onto itself, which is nonsense rather than merely inaccurate. + +### What is and is not supported + +Two-dimensional simulations, with a planar monitor normal to a coordinate axis +lying in a homogeneous region. The three-dimensional case additionally needs the +s/p rotation by azimuth, with its removable singularity at normal incidence, and +raises rather than guessing. Cylindrical coordinates are not supported. + +A single plane is complete for the half-space above it, so unlike near2far there +is no closed surface to build: the plane plus the hemisphere at infinity already +is one, and the up/down split discards what is heading the wrong way. To account +for both half-spaces — the substrate as well as the superstrate — use two +parallel monitors, each with its own stack, and flip `sign` on the lower one. + +[examples/adjoint_optimization/grating_coupler_asm.py](https://github.com/NanoComp/meep/blob/master/python/examples/adjoint_optimization/grating_coupler_asm.py) +is a worked two-etch grating coupler radiating through a thick superstrate into a +fiber, with a forward-only mode and an optimization mode. + Broadband Waveguide Mode Converter with Minimum Feature Size ------------------------------------------------------------ From 5543cc9a39649702f3129cb8807f9d9e7601ded8 Mon Sep 17 00:00:00 2001 From: Alec Hammond Date: Wed, 26 Aug 2026 10:58:46 -0700 Subject: [PATCH 6/8] adjoint: 3D angular-spectrum propagation, with gradients confirmed A monitor plane is a line in a 2D simulation and a rectangle in a 3D one, so the transverse space is one- or two-dimensional. Everything downstream of the transform now works on a flat list of transverse wavevectors and is indifferent to which it was; only the transform and the polarization basis differ. The 3D ingredient is the rotation into the s and p directions of each transverse wavevector, whose azimuth is undefined at normal incidence and is pinned there rather than left to produce a nan from atan2(0, 0). Verified against explicit up-going plane waves at 200 random oblique directions, to 4e-16. The two transverse dimensionalities use genuinely different polarization bases, and the sign of the decomposition differs between them: with one transverse direction the wavevector runs along v_hat, so the 3D convention would put s_hat along -u_hat. Unifying them without noticing inverted the 2D split -- an up-going field landed entirely in the down-going branch -- which the existing 2D tests caught. Tests, 28 passing. In 3D: the decomposition recovers a purely up-going field to the level of the monitor truncation, a circular Gaussian spreads as w0 sqrt(1 + (z/zR)^2), self-overlap is one, a matched mode couples above 0.99, and the azimuth is finite at normal incidence. Two of those started out failing for a reason worth recording: a spectrum that is purely s at every wavevector is azimuthally polarized, carries a vortex at normal incidence, and does not decay compactly, so it truncates badly on any finite monitor. A real Gaussian beam is linearly polarized, with s and p content varying as cos and sin of the azimuth. The adjoint gradient is checked against a finite difference in both dimensionalities and agrees to within 1%. Getting that check to mean anything took some care, because it first appeared to show a wrong gradient -- a ratio of 1.37 in 2D and 1.59 in 3D, constant in the finite-difference step, which rules out truncation error and looks like a real defect. Isolating the 2D chain, from a bare sum of |E|^2 on the monitor up to the self-normalized mode overlap, every stage agreed to better than 2e-4 once the run was converged, so neither the propagator nor the ratio-valued objective was at fault. Holding the 3D geometry fixed and varying only the run length, with both runs of the finite difference forced to cover the same interval: length 40 adjoint 8.84e-6 finite difference 4.923e-6 ratio 1.796 length 80 adjoint 7.92e-6 finite difference 4.923e-6 ratio 1.608 length 140 adjoint 4.93e-6 finite difference 4.923e-6 ratio 1.001 The finite difference does not move; it is the adjoint that converges. The binding constraint is the adjoint DFT, which in a lossless background rings for considerably longer than the forward one, and an under-converged adjoint is wrong by a factor that does not shrink with the step -- which is why it read as a wrong gradient rather than as noise. The tests therefore fix the run length outright, so both runs of a finite difference cover the same interval. Left adaptive, stop_when_dft_decayed stops the perturbed run at a different time, and that difference scales with the perturbation, producing the same misleading signature. --- NEWS.md | 3 +- doc/docs/Python_Tutorials/Adjoint_Solver.md | 22 +- python/adjoint/angular_spectrum.py | 505 ++++++++++++++------ python/tests/test_angular_spectrum.py | 391 +++++++++++++-- 4 files changed, 731 insertions(+), 190 deletions(-) diff --git a/NEWS.md b/NEWS.md index 03f6c1d2a..a8d273ad5 100644 --- a/NEWS.md +++ b/NEWS.md @@ -7,7 +7,8 @@ analytically, in JAX. Unlike `add_near2far`, which requires a homogeneous medium, this handles radiation crossing a material interface, so the layers above the monitor can leave the FDTD cell entirely and become differentiable - parameters. Usable with or without the adjoint solver. 2D only for now. + parameters. Supports 2D and 3D Cartesian simulations, and is usable with or + without the adjoint solver. * Adjoint solver: objective functions are now differentiated with a single vector-Jacobian product instead of a full frequency Jacobian per objective diff --git a/doc/docs/Python_Tutorials/Adjoint_Solver.md b/doc/docs/Python_Tutorials/Adjoint_Solver.md index d1c3b9f20..baf5b0080 100644 --- a/doc/docs/Python_Tutorials/Adjoint_Solver.md +++ b/doc/docs/Python_Tutorials/Adjoint_Solver.md @@ -333,10 +333,24 @@ and the beam wraps onto itself, which is nonsense rather than merely inaccurate. ### What is and is not supported -Two-dimensional simulations, with a planar monitor normal to a coordinate axis -lying in a homogeneous region. The three-dimensional case additionally needs the -s/p rotation by azimuth, with its removable singularity at normal incidence, and -raises rather than guessing. Cylindrical coordinates are not supported. +Two- and three-dimensional Cartesian simulations, with a planar monitor normal to +a coordinate axis lying in a homogeneous region — a line in 2D, a rectangle in +3D. Cylindrical coordinates are not supported. + +In 3D the transverse wavevectors form a plane, and the tangential fields are +resolved into the s and p directions of each one. A beam is therefore specified +by its linear polarization rather than by a single scalar component: +`gaussian_mode(waist, polarization=...)` builds a linearly polarized beam whose s +and p content follows the azimuth. A spectrum that were purely s at every +wavevector would instead be azimuthally polarized, carrying a vortex at normal +incidence, which is rarely what is wanted. + +One consequence worth knowing when checking gradients in 3D: the *adjoint* DFT +takes appreciably longer to converge than the forward one, and an +under-converged adjoint produces a gradient that is wrong by a fixed factor +which does not shrink with the finite-difference step — so it reads like a bug +rather than like noise. Give the adjoint run enough time, or fix the run length +so that both runs of a finite difference cover the same interval. A single plane is complete for the half-space above it, so unlike near2far there is no closed surface to build: the plane plus the hemisphere at infinity already diff --git a/python/adjoint/angular_spectrum.py b/python/adjoint/angular_spectrum.py index 9de36ddb6..7d1e56e93 100644 --- a/python/adjoint/angular_spectrum.py +++ b/python/adjoint/angular_spectrum.py @@ -284,57 +284,46 @@ class PropagationResult(NamedTuple): # conservation, and differ only in the sign of r_p. +# The tangential axes of a plane, ordered right-handed so that u_hat x v_hat = +# n_hat. That ordering is what makes the decomposition below sign-free: +# n_hat x H_t has components (-H_v, +H_u) in the same basis. +_PLANE_AXES = { + mp.X: ((mp.Ey, mp.Ez), (mp.Hy, mp.Hz)), + mp.Y: ((mp.Ez, mp.Ex), (mp.Hz, mp.Hx)), + mp.Z: ((mp.Ex, mp.Ey), (mp.Hx, mp.Hy)), +} + + def _tangential_components(normal: int, dimensions: int) -> Tuple[Tuple[int, ...], ...]: """The E and H components tangential to a plane with the given normal.""" - if dimensions == 2: - if normal == mp.Y: - return (mp.Ex, mp.Ez), (mp.Hx, mp.Hz) - if normal == mp.X: - return (mp.Ey, mp.Ez), (mp.Hy, mp.Hz) + if normal not in _PLANE_AXES: + raise ValueError(f"Unsupported monitor normal {normal}.") + if dimensions == 2 and normal == mp.Z: raise ValueError( "In a 2D simulation the monitor plane must be normal to x or y, " f"but got normal={normal}." ) - raise NotImplementedError( - "Angular-spectrum propagation currently supports 2D simulations only. " - "The three-dimensional case additionally needs the s/p rotation by " - "azimuth, with its removable singularity at normal incidence, and is " - "not implemented here." - ) - - -# For an up-going plane wave the tangential fields satisfy H_t = Y (n_hat x E_t) -# in both polarizations, the cross product supplying the orientation, so -# -# E_up = (E_t - (n_hat x H_t) / Y) / 2 -# -# uniformly. Written out in a tangential basis this becomes one (E component, -# partnering H component, sign, polarization) tuple per polarization, with -# E_up = (E + sign * H / Y) / 2. -_DECOMPOSITION = { - mp.Y: ( - (mp.Ez, mp.Hx, +1, S_POLARIZATION), - (mp.Ex, mp.Hz, -1, P_POLARIZATION), - ), - mp.X: ( - (mp.Ez, mp.Hy, -1, S_POLARIZATION), - (mp.Ey, mp.Hz, +1, P_POLARIZATION), - ), -} + if dimensions not in (2, 3): + raise NotImplementedError( + "Angular-spectrum propagation supports 2D and 3D Cartesian " + f"simulations, not {dimensions}D. Cylindrical coordinates are not " + "supported." + ) + return _PLANE_AXES[normal] class Mode(NamedTuple): """A target field to project onto, defined by its angular spectrum. Attributes: - spectrum: called as `spectrum(kt, k0, index)` and returning a - (num frequencies, num kt) array of tangential electric field - amplitudes. - polarization: which polarization the amplitudes belong to. + 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. """ spectrum: Callable - polarization: int = S_POLARIZATION def gaussian_mode( @@ -342,8 +331,9 @@ def gaussian_mode( tilt_deg: float = 0.0, offset: float = 0.0, polarization: int = S_POLARIZATION, + tilt_azimuth_deg: float = 0.0, ) -> Mode: - """A tilted, laterally offset Gaussian, e.g. a fiber mode. + """A tilted, laterally offset, linearly polarized Gaussian, e.g. a fiber mode. The spectrum is written in closed form rather than sampled, so the tilt and the offset are exact continuous parameters and differentiable, instead of @@ -353,26 +343,59 @@ def gaussian_mode( waist: the 1/e field radius at the target plane, in Meep units. For a fiber quoted by mode-field diameter, this is MFD / 2. tilt_deg: the angle from the plane normal, in degrees. - offset: the lateral displacement of the beam center at the target plane. - polarization: S_POLARIZATION or P_POLARIZATION. + offset: the lateral displacement of the beam center at the target plane, + along the tilt direction. + polarization: the linear polarization direction. `S_POLARIZATION` means + along the second tangential axis, `P_POLARIZATION` the first; in 2D + these are the out-of-plane and in-plane electric fields + respectively. In 3D the beam is linearly polarized along that axis + and its s and p content follows the azimuth of each wavevector. + tilt_azimuth_deg: which way the beam is tilted, measured in the + tangential plane. Ignored in 2D, where there is only one direction + to tilt in. Returns: A `Mode`. """ - def spectrum(kt, k0, index): + def spectrum(propagator, kt, k0, index): kt = jnp.asarray(kt) - center = ( - jnp.asarray(index).real[:, None] - * jnp.asarray(k0)[:, None] + transverse = ( + jnp.asarray(index).real + * jnp.asarray(k0) * jnp.sin(jnp.deg2rad(jnp.asarray(tilt_deg))) ) - detuning = kt[None, :] - center - return jnp.exp(-jnp.square(detuning * waist) / 4.0) * jnp.exp( - -1j * kt[None, :] * offset + azimuth = jnp.deg2rad(jnp.asarray(tilt_azimuth_deg)) + direction = ( + jnp.array([1.0]) + if kt.shape[-1] == 1 + else jnp.stack([jnp.cos(azimuth), jnp.sin(azimuth)]) ) + center = transverse[:, None] * direction + detuning = kt[None, :, :] - center[:, None, :] + envelope = jnp.exp( + -jnp.sum(jnp.square(detuning * waist), axis=-1) / 4.0 + ) * jnp.exp(-1j * jnp.sum(detuning * direction * offset, axis=-1)) + + rotation = propagator._azimuth() + if rotation is None: + weights = ( + jnp.array([1.0, 0.0]) + if polarization == S_POLARIZATION + else jnp.array([0.0, 1.0]) + ) + components = jnp.broadcast_to(weights, envelope.shape + (2,)) + else: + cosine, sine = rotation + if polarization == S_POLARIZATION: + # polarized along v_hat: v_hat . s_hat = cos, v_hat . p_hat = sin + components = jnp.stack([cosine, sine], axis=-1) + else: + components = jnp.stack([-sine, cosine], axis=-1) + components = jnp.broadcast_to(components[None, :, :], envelope.shape + (2,)) + return envelope[..., None] * components - return Mode(spectrum=spectrum, polarization=polarization) + return Mode(spectrum=spectrum) class AngularSpectrum: @@ -418,27 +441,50 @@ def __init__( aperture matter. """ stack.validate() - if normal not in _DECOMPOSITION: - raise ValueError(f"Unsupported monitor normal {normal}.") if sign not in (1, -1): raise ValueError(f"sign must be +1 or -1, got {sign}.") + # A monitor plane is a line in a 2D simulation and a rectangle in a 3D + # one, so the transverse space is one- or two-dimensional. Everything + # downstream of the transform works on a flat list of transverse + # wavevectors and is indifferent to which it was. + self.num_points = ( + (int(num_points),) + if onp.ndim(num_points) == 0 + else tuple(int(n) for n in num_points) + ) + self.pitch = ( + (float(pitch),) * len(self.num_points) + if onp.ndim(pitch) == 0 + else tuple(float(p) for p in pitch) + ) + if len(self.pitch) != len(self.num_points): + raise ValueError( + f"pitch has {len(self.pitch)} entries but num_points has " + f"{len(self.num_points)}." + ) + self.transverse_dimensions = len(self.num_points) + _tangential_components(normal, self.transverse_dimensions + 1) + self.stack = stack self.frequencies = onp.asarray(frequencies, dtype=float) - self.pitch = float(pitch) - self.num_points = int(num_points) self.normal = normal self.sign = int(sign) self.pad_factor = int(pad_factor) self._k0 = 2 * onp.pi * self.frequencies if kt is None: - padded = self.num_points * self.pad_factor - self._kt = jnp.asarray(2 * onp.pi * onp.fft.fftfreq(padded, d=self.pitch)) + self._padded = tuple(n * self.pad_factor for n in self.num_points) + axes = [ + 2 * onp.pi * onp.fft.fftfreq(n, d=d) + for n, d in zip(self._padded, self.pitch) + ] + grids = onp.meshgrid(*axes, indexing="ij") + self._kt = jnp.asarray(onp.stack([g.ravel() for g in grids], axis=-1)) self._uniform = True - self._padded = padded else: - self._kt = jnp.asarray(kt, dtype=float) + kt = onp.asarray(kt, dtype=float) + self._kt = jnp.asarray(kt.reshape(kt.shape[0], -1)) self._uniform = False self._padded = None @@ -480,68 +526,171 @@ def stack_thickness(self): total = total + thickness return total - def coordinates(self) -> onp.ndarray: - """The transverse coordinates of the monitor samples, centered on zero.""" - return (onp.arange(self.num_points) - (self.num_points - 1) / 2) * self.pitch + def coordinates(self): + """The transverse sample coordinates, centered on zero. + + A single array for a line monitor, a list of two for a rectangular one. + """ + axes = [ + (onp.arange(n) - (n - 1) / 2) * d + for n, d in zip(self.num_points, self.pitch) + ] + return axes[0] if self.transverse_dimensions == 1 else axes + + def _coordinate_axes(self): + """`coordinates`, always as a list, one array per transverse axis.""" + axes = self.coordinates() + return [axes] if self.transverse_dimensions == 1 else axes + + def _sample_positions(self) -> onp.ndarray: + """Every sample position, flattened to (num samples, num transverse).""" + grids = onp.meshgrid(*self._coordinate_axes(), indexing="ij") + return onp.stack([g.ravel() for g in grids], axis=-1) def _transform(self, values: jnp.ndarray) -> jnp.ndarray: - """Transforms (num frequencies, num points) samples to (num freq, num kt).""" + """Transforms monitor samples to (num frequencies, num kt).""" values = jnp.asarray(values) - if self._uniform: - padded = jnp.zeros( - values.shape[:-1] + (self._padded,), dtype=jnp.complex128 + transverse = tuple(range(-self.transverse_dimensions, 0)) + measure = float(onp.prod(self.pitch)) + if not self._uniform: + flat = values.reshape(values.shape[: transverse[0]] + (-1,)) + phase = jnp.exp( + -1j * jnp.einsum("kd,xd->kx", self._kt, self._sample_positions()) ) - padded = padded.at[..., : self.num_points].set(values) - spectrum = jnp.fft.fft(padded, axis=-1) * self.pitch - # The samples are centered on zero, so undo the phase ramp implied by - # having placed them at indices 0..num_points-1. - origin = self.coordinates()[0] - return spectrum * jnp.exp(-1j * self._kt * origin) - phase = jnp.exp(-1j * self._kt[:, None] * self.coordinates()[None, :]) - return jnp.einsum("kx,...x->...k", phase, values) * self.pitch - - def _polarization_terms(self, fields: TangentialFields): - """Yields (polarization, E samples, H samples, sign) for what is present.""" + return jnp.einsum("kx,...x->...k", phase, flat) * measure + padded = jnp.zeros( + values.shape[: transverse[0]] + self._padded, dtype=jnp.complex128 + ) + padded = padded.at[ + (Ellipsis,) + tuple(slice(0, n) for n in self.num_points) + ].set(values) + spectrum = jnp.fft.fftn(padded, axes=transverse) * measure + spectrum = spectrum.reshape(spectrum.shape[: transverse[0]] + (-1,)) + # The samples are centered on zero, so undo the phase ramp implied by + # having placed them at indices starting from zero. + origins = jnp.asarray([axis[0] for axis in self._coordinate_axes()]) + return spectrum * jnp.exp(-1j * (self._kt @ origins)) + + def _tangential_spectra(self, fields: TangentialFields): + """Transforms the tangential fields, returning them on the (u, v) axes. + + Returns `(E_u, E_v, H_u, H_v)`, each (num frequencies, num kt), with a + missing component treated as zero. In 2D exactly one of the two + polarizations is populated and the other stays zero, which costs nothing + and spares the caller declaring which one their source excites. + """ if fields.normal != self.normal: raise ValueError( f"The fields are on a plane normal to {fields.normal} but this " f"propagator was built for {self.normal}." ) - for e_component, h_component, sign, polarization in _DECOMPOSITION[self.normal]: - e_values = fields.E.get(e_component) - h_values = fields.H.get(h_component) - if e_values is None and h_values is None: - continue - if e_values is None or h_values is None: + (e_u, e_v), (h_u, h_v) = _PLANE_AXES[self.normal] + # E_u pairs with H_v, not H_u: the decomposition contracts E_t against + # n_hat x H_t, which swaps the two tangential axes. For a y-normal plane + # that means Ez goes with Hx and Ex with Hz. + pairs = ((e_u, h_v), (e_v, h_u)) + present = [(fields.E.get(e), fields.H.get(h)) for e, h in pairs] + if all(e is None and h is None for e, h in present): + raise ValueError( + "No tangential field components were supplied; nothing to " "propagate." + ) + for (electric, magnetic), (e, h) in zip(present, pairs): + if (electric is None) != (magnetic is None): raise ValueError( "Separating up-going from down-going radiation needs both " - f"tangential fields, but only one of {mp.component_name(e_component)}" - f" and {mp.component_name(h_component)} was supplied." + f"tangential fields, but only one of " + f"{mp.component_name(e)} and {mp.component_name(h)} was " + "supplied." ) - yield polarization, jnp.asarray(e_values), jnp.asarray(h_values), sign + zero = None + for electric, magnetic in present: + if electric is not None: + zero = jnp.zeros_like(self._transform(jnp.asarray(electric))) + break + # Returned as (E_u, H_v, E_v, H_u), matching how they pair up. + return tuple( + zero if value is None else self._transform(jnp.asarray(value)) + for pair in present + for value in pair + ) + + def _azimuth(self): + """cos and sin of the angle from u_hat to the transverse wavevector. + + At normal incidence the plane of incidence is undefined and any + orthogonal pair will do, so the azimuth is pinned to zero there rather + than left to produce a nan from atan2(0, 0). + """ + if self.transverse_dimensions == 1: + # The transverse wavevector lies along a single axis, so there is no + # azimuth to speak of and the (u, v) axes are already the s and p + # directions. + return None + magnitude = jnp.linalg.norm(self._kt, axis=-1) + safe = jnp.where(magnitude > 0, magnitude, 1.0) + cosine = jnp.where(magnitude > 0, self._kt[:, 0] / safe, 1.0) + sine = jnp.where(magnitude > 0, self._kt[:, 1] / safe, 0.0) + return cosine, sine def decompose(self, fields: TangentialFields): """Splits the monitor fields into up- and down-going spectra. + The tangential fields are rotated into the s and p directions of each + transverse wavevector, where the admittance is a scalar, and separated + using + + E_up_s = (E_s - H_p / Y_s) / 2 E_up_p = (E_p + H_s / Y_p) / 2 + + which is `E_up = (E_t - (n_hat x H_t) / Y) / 2` written out in that + basis. In 2D the wavevector lies along one axis, so the rotation is the + identity and (u, v) are already (s, p) up to the ordering below. + Returns: `(up, down)`, each a dict mapping polarization to a (num frequencies, num kt) array of tangential electric field amplitudes at the monitor plane. """ - up, down = {}, {} - admittance = self._admittances[0] - for polarization, e_values, h_values, sign in self._polarization_terms(fields): - e_spectrum = self._transform(e_values) - h_spectrum = self._transform(h_values) - # `sign` already carries the orientation of n_hat x H_t; `fields.sign` - # flips which branch counts as outgoing for a downward-facing monitor. - scaled = fields.sign * sign * h_spectrum / admittance[polarization] - up[polarization] = 0.5 * (e_spectrum + scaled) - down[polarization] = 0.5 * (e_spectrum - scaled) - if not up: - raise ValueError( - "No tangential field components were supplied; nothing to " "propagate." - ) + electric_u, magnetic_v, electric_v, magnetic_u = self._tangential_spectra( + fields + ) + azimuth = self._azimuth() + admittance_s, admittance_p = self._admittances[0] + # `fields.sign` flips which branch counts as outgoing, for a monitor + # facing down into a substrate. + outgoing = fields.sign + + if azimuth is None: + # One transverse direction, so there is no azimuth and the + # right-handed axes are already a valid s/p pair: u_hat is + # perpendicular to the wavevector, v_hat lies along it. Using + # (n_hat x H)_u = -H_v and (n_hat x H)_v = +H_u directly, + # + # E_up_u = (E_u + H_v / Y_s) / 2 E_up_v = (E_v - H_u / Y_p) / 2 + # + # Note the signs are *not* those of the rotated case below: the 3D + # convention puts s_hat along -u_hat here, since the wavevector runs + # along v_hat, and the two bases therefore differ by a sign. + electric_s, magnetic_s = electric_u, magnetic_u + electric_p, magnetic_p = electric_v, magnetic_v + cross_s = -outgoing * magnetic_p / admittance_s + cross_p = outgoing * magnetic_s / admittance_p + else: + cosine, sine = azimuth + electric_s = -electric_u * sine + electric_v * cosine + electric_p = electric_u * cosine + electric_v * sine + magnetic_s = -magnetic_u * sine + magnetic_v * cosine + magnetic_p = magnetic_u * cosine + magnetic_v * sine + cross_s = outgoing * magnetic_p / admittance_s + cross_p = -outgoing * magnetic_s / admittance_p + + up = { + S_POLARIZATION: 0.5 * (electric_s - cross_s), + P_POLARIZATION: 0.5 * (electric_p - cross_p), + } + down = { + S_POLARIZATION: 0.5 * (electric_s + cross_s), + P_POLARIZATION: 0.5 * (electric_p + cross_p), + } return up, down def _transmission(self, polarization: int): @@ -617,11 +766,18 @@ def _weights(self, result: PropagationResult) -> jnp.ndarray: def _spectral_measure(self) -> float: """The dk / 2pi factor that turns a spectral sum into a real-space integral.""" if self._uniform: - return (2 * onp.pi / (self._padded * self.pitch)) / (2 * onp.pi) - spacing = jnp.diff(self._kt) - # Trapezoid weights would be more careful, but a chosen kt set is - # normally uniform; require that rather than silently mis-weighting. - return jnp.mean(spacing) / (2 * onp.pi) + measure = 1.0 + for padded, pitch in zip(self._padded, self.pitch): + measure *= (2 * onp.pi / (padded * pitch)) / (2 * onp.pi) + return measure + # A chosen set of wavevectors is assumed uniform along each axis; + # trapezoid weights would be more careful but the set is normally a grid. + measure = 1.0 + for axis in range(self._kt.shape[-1]): + values = onp.unique(onp.asarray(self._kt[:, axis])) + spacing = onp.mean(onp.diff(values)) if values.size > 1 else 1.0 + measure *= spacing / (2 * onp.pi) + return measure def power(self, fields: TangentialFields, distance: Optional[float] = None): """Outgoing power through the target plane, one value per frequency.""" @@ -658,15 +814,16 @@ def overlap( A (num frequencies,) array. """ result = self.spectrum(fields, distance) - weights = self._weights(result)[..., mode.polarization] - field_amplitude = result.amplitudes[..., mode.polarization] - mode_amplitude = mode.spectrum(result.kt, self._k0, result.index) + weights = self._weights(result) + field_amplitude = result.amplitudes + mode_amplitude = mode.spectrum(self, result.kt, self._k0, result.index) - cross = jnp.sum(weights * field_amplitude * jnp.conj(mode_amplitude), axis=-1) - mode_norm = jnp.sum(weights * jnp.abs(mode_amplitude) ** 2, axis=-1) + axes = (1, 2) + cross = jnp.sum(weights * field_amplitude * jnp.conj(mode_amplitude), axis=axes) + mode_norm = jnp.sum(weights * jnp.abs(mode_amplitude) ** 2, axis=axes) coupled = jnp.abs(cross) ** 2 / mode_norm if incident_power is None: - incident_power = jnp.sum(weights * jnp.abs(field_amplitude) ** 2, axis=-1) + incident_power = jnp.sum(weights * jnp.abs(field_amplitude) ** 2, axis=axes) return coupled / incident_power def report(self, fields: TangentialFields) -> Dict[str, jnp.ndarray]: @@ -707,11 +864,20 @@ def spectral_power(spectra, mask=None): up_propagating = spectral_power(up, propagating) edges = [] - for _, e_values, _, _ in self._polarization_terms(fields): - magnitude = jnp.abs(e_values) - peak = jnp.max(magnitude, axis=-1) - edge = jnp.maximum(magnitude[..., 0], magnitude[..., -1]) - edges.append(edge / jnp.where(peak > 0, peak, 1.0)) + for values in list(fields.E.values()): + magnitude = jnp.abs(jnp.asarray(values)) + axes = tuple(range(-self.transverse_dimensions, 0)) + peak = jnp.max(magnitude, axis=axes) + border = jnp.zeros_like(peak) + for axis in axes: + border = jnp.maximum( + border, + jnp.max( + jnp.take(magnitude, jnp.array([0, -1]), axis=axis), + axis=axes, + ), + ) + edges.append(border / jnp.where(peak > 0, peak, 1.0)) total = up_power + down_power return { @@ -721,35 +887,53 @@ def spectral_power(spectra, mask=None): "edge_amplitude": jnp.max(jnp.stack(edges, axis=0), axis=0), } - def propagate( - self, fields: TangentialFields, distance: float, coordinates=None - ) -> Dict[int, jnp.ndarray]: + def propagate(self, fields: TangentialFields, distance: float, coordinates=None): """The tangential electric field at the target plane, in real space. Args: fields: the monitor fields. distance: distance from the monitor to the target plane. - coordinates: where to evaluate. Defaults to the monitor's own - coordinates; pass a wider range to see a beam that has spread. + coordinates: where to evaluate, as one array per transverse + dimension. Defaults to the monitor's own coordinates; pass a + wider range to see a beam that has spread. Returns: - A dict mapping each tangential electric component to a - (num frequencies, num coordinates) array. + A dict mapping each tangential electric component to an array of + shape (num frequencies,) + the coordinate shape. """ result = self.spectrum(fields, distance) if coordinates is None: - coordinates = self.coordinates() - coordinates = jnp.asarray(coordinates) + axes = self._coordinate_axes() + else: + axes = ( + [onp.asarray(coordinates)] + if self.transverse_dimensions == 1 + else [onp.asarray(a) for a in coordinates] + ) + grids = onp.meshgrid(*axes, indexing="ij") + positions = jnp.asarray(onp.stack([g.ravel() for g in grids], axis=-1)) + + rotation = self._azimuth() + amplitude_s = result.amplitudes[..., S_POLARIZATION] + amplitude_p = result.amplitudes[..., P_POLARIZATION] + if rotation is None: + amplitude_u, amplitude_v = amplitude_s, amplitude_p + else: + cosine, sine = rotation + # The rotation into (s, p) is a reflection, hence its own inverse. + amplitude_u = -amplitude_s * sine + amplitude_p * cosine + amplitude_v = amplitude_s * cosine + amplitude_p * sine + + phase = jnp.exp(1j * jnp.einsum("xd,kd->xk", positions, result.kt)) measure = self._spectral_measure() - phase = jnp.exp(1j * result.kt[None, :] * coordinates[:, None]) + (e_u, e_v), _ = _PLANE_AXES[self.normal] + shape = tuple(len(a) for a in axes) outputs = {} - for e_component, _, _, polarization in _DECOMPOSITION[self.normal]: - if not jnp.any(result.amplitudes[..., polarization]): + for component, amplitude in ((e_u, amplitude_u), (e_v, amplitude_v)): + if not jnp.any(amplitude): continue - outputs[e_component] = ( - jnp.einsum("xk,fk->fx", phase, result.amplitudes[..., polarization]) - * measure - ) + values = jnp.einsum("xk,fk->fx", phase, amplitude) * measure + outputs[component] = values.reshape(values.shape[:1] + shape) return outputs # ---------------------------------------------------------------- Meep glue @@ -761,27 +945,43 @@ def propagate( @staticmethod def _plane_geometry(simulation: mp.Simulation, volume: mp.Volume): - """Infers the normal, sample pitch, and sample count of a planar volume.""" + """Infers the normal, sample pitch, and sample counts of a planar volume.""" size = [volume.size.x, volume.size.y, volume.size.z] - zero = [i for i, extent in enumerate(size) if extent == 0] - if simulation.dimensions != 2: + dimensions = simulation.dimensions + if dimensions not in (2, 3): raise NotImplementedError( - "Angular-spectrum propagation currently supports 2D " - f"simulations only, but this one is {simulation.dimensions}D." + "Angular-spectrum propagation supports 2D and 3D Cartesian " + f"simulations, not {dimensions}D." ) - if len(zero) != 2 or 2 not in zero: + in_plane = [0, 1] if dimensions == 2 else [0, 1, 2] + flat = [axis for axis in in_plane if size[axis] == 0] + if len(flat) != 1: raise ValueError( - "The monitor must be a line normal to x or y, i.e. a Volume " - f"with exactly one nonzero in-plane size; got size={size}." + "The monitor must be a plane normal to one coordinate axis, " + f"i.e. a Volume with exactly one zero size; got size={size} in " + f"{dimensions}D." ) - normal = mp.X if zero[0] == 0 else mp.Y - coordinates = simulation.get_array_metadata(vol=volume) - axis = 1 if normal == mp.X else 0 - samples = onp.asarray(coordinates[axis]) - if samples.size < 2: - raise ValueError("The monitor needs at least two sample points.") - pitch = float(onp.mean(onp.diff(samples))) - return normal, pitch, int(samples.size) + normal = (mp.X, mp.Y, mp.Z)[flat[0]] + tangential = [axis for axis in in_plane if axis != flat[0]] + metadata = simulation.get_array_metadata(vol=volume) + + pitches, counts = [], [] + for axis in tangential: + samples = onp.asarray(metadata[axis]) + if samples.size < 2: + raise ValueError( + "The monitor needs at least two sample points along every " + f"tangential axis, but has {samples.size} along axis {axis}." + ) + pitches.append(float(onp.mean(onp.diff(samples)))) + counts.append(int(samples.size)) + + # Meep orders the array axes x, y, z; the propagator wants them in the + # right-handed (u, v) order of the plane, which for a y-normal plane is + # (z, x) rather than (x, z). + if normal == mp.Y and dimensions == 3: + pitches, counts = pitches[::-1], counts[::-1] + return normal, tuple(pitches), tuple(counts) @staticmethod def _assert_homogeneous( @@ -852,12 +1052,11 @@ def fields_from_monitor( set(registered[0]) if registered else { - component - for pair in _DECOMPOSITION[self.normal] - for component in pair[:2] + component for group in _PLANE_AXES[self.normal] for component in group } ) - for e_component, h_component, _, _ in _DECOMPOSITION[self.normal]: + (e_u, e_v), (h_u, h_v) = _PLANE_AXES[self.normal] + for e_component, h_component in ((e_u, h_u), (e_v, h_v)): for component, target in ( (e_component, electric), (h_component, magnetic), @@ -870,9 +1069,9 @@ def fields_from_monitor( for i in range(len(self.frequencies)) ] ) - if values.shape[-1] != self.num_points: + if values.shape[1:] != self.num_points: raise ValueError( - f"The monitor returned {values.shape[-1]} samples for " + f"The monitor returned {values.shape[1:]} samples for " f"{mp.component_name(component)} but the propagator was " f"built for {self.num_points}. The volume Meep actually " "used may have been snapped to the grid; build the " @@ -920,9 +1119,7 @@ def objective_arguments(self, simulation, volume, **kwargs): from . import FourierFields self._objective_components = [ - component - for e_component, h_component, _, _ in _DECOMPOSITION[self.normal] - for component in (e_component, h_component) + component for group in _PLANE_AXES[self.normal] for component in group ] return [ FourierFields(simulation, volume, component, yee_grid=False, **kwargs) @@ -951,4 +1148,4 @@ def take(self, args) -> TangentialFields: def __len__(self) -> int: """How many leading objective arguments `take` consumes.""" - return len(_DECOMPOSITION[self.normal]) * 2 + return 4 diff --git a/python/tests/test_angular_spectrum.py b/python/tests/test_angular_spectrum.py index 19c72ca3e..1246ff648 100644 --- a/python/tests/test_angular_spectrum.py +++ b/python/tests/test_angular_spectrum.py @@ -90,7 +90,9 @@ def test_single_interface_matches_fresnel(self): N_AIR * cos_in + N_OXIDE * cos_out ) self.assertAlmostEqual( - complex(reflection[0, 0]), expected, places=12, + complex(reflection[0, 0]), + expected, + places=12, msg=f"{name} polarization at {angle} deg", ) @@ -155,9 +157,7 @@ def transfer_matrix(angle_deg, polarization): self.assertAlmostEqual( complex(transmission[0, 0]), expected_t, places=12 ) - self.assertAlmostEqual( - complex(reflection[0, 0]), expected_r, places=12 - ) + self.assertAlmostEqual(complex(reflection[0, 0]), expected_r, places=12) def test_energy_is_conserved(self): """Reflected plus transmitted power equals the incident, losslessly.""" @@ -171,11 +171,15 @@ def test_energy_is_conserved(self): [indices[0] * self.k0[0] * math.sin(math.radians(angle))] ) y_in = _admittance( - indices[0], _wavevector(indices[0], self.k0, kt), self.k0, + indices[0], + _wavevector(indices[0], self.k0, kt), + self.k0, polarization, ) y_out = _admittance( - indices[-1], _wavevector(indices[-1], self.k0, kt), self.k0, + indices[-1], + _wavevector(indices[-1], self.k0, kt), + self.k0, polarization, ) reflected = abs(complex(reflection[0, 0])) ** 2 @@ -189,7 +193,11 @@ def _uniform_propagator(index=N_OXIDE, num_points=512, pitch=0.05, pad_factor=8) """A propagator with no interface, i.e. plain homogeneous propagation.""" stack = mpa.Stack([mpa.Layer(index, 0.0), mpa.Layer(index)]) return mpa.AngularSpectrum( - stack, [1 / WAVELENGTH], pitch, num_points, normal=mp.Y, + stack, + [1 / WAVELENGTH], + pitch, + num_points, + normal=mp.Y, pad_factor=pad_factor, ) @@ -203,11 +211,11 @@ def _up_going(propagator, values): """ spectrum = propagator._transform(jnp.asarray(values)) admittance = propagator._admittances[0][asm.S_POLARIZATION] - coordinates = propagator.coordinates() + positions = jnp.asarray(propagator._sample_positions()) magnetic = ( jnp.einsum( "xk,fk->fx", - jnp.exp(1j * propagator.kt[None, :] * coordinates[:, None]), + jnp.exp(1j * jnp.einsum("xd,kd->xk", positions, propagator.kt)), spectrum * admittance, ) * propagator._spectral_measure() @@ -289,7 +297,8 @@ def test_gaussian_beam_spreading(self): profile = onp.abs( onp.asarray( propagator.propagate( - _up_going(propagator, values), distance, + _up_going(propagator, values), + distance, coordinates=samples, )[mp.Ez][0] ) @@ -318,9 +327,7 @@ def test_self_overlap_is_unity(self): x = propagator.coordinates() fields = _up_going(propagator, (onp.exp(-((x / 3.0) ** 2)))[None, :]) result = propagator.spectrum(fields, 100.0) - itself = mpa.Mode( - spectrum=lambda kt, k0, index: result.amplitudes[..., asm.S_POLARIZATION] - ) + itself = mpa.Mode(spectrum=lambda propagator, kt, k0, index: result.amplitudes) self.assertAlmostEqual( float(propagator.overlap(fields, itself, 100.0)[0]), 1.0, places=10 ) @@ -381,9 +388,7 @@ def test_gradient_is_finite_on_the_light_line(self): def objective(scale): fields = _up_going(propagator, scale * values) - return jnp.real( - jnp.sum(propagator.spectrum(fields, 1.0).amplitudes) - ) + return jnp.real(jnp.sum(propagator.spectrum(fields, 1.0).amplitudes)) gradient = jax.grad(objective)(1.0) self.assertTrue(onp.isfinite(float(gradient)), "gradient is not finite") @@ -426,15 +431,193 @@ def test_only_the_last_layer_may_be_semi_infinite(self): def test_both_tangential_fields_are_required(self): """One field alone cannot distinguish up-going from down-going.""" propagator = _uniform_propagator(num_points=64) - fields = mpa.TangentialFields( - E={mp.Ez: onp.ones((1, 64))}, H={}, normal=mp.Y - ) + fields = mpa.TangentialFields(E={mp.Ez: onp.ones((1, 64))}, H={}, normal=mp.Y) with self.assertRaisesRegex(ValueError, "both"): propagator.decompose(fields) - def test_three_dimensions_is_rejected_clearly(self): - with self.assertRaisesRegex(NotImplementedError, "2D"): - asm._tangential_components(mp.Z, 3) + def test_three_dimensions_is_supported(self): + self.assertEqual( + asm._tangential_components(mp.Z, 3), ((mp.Ex, mp.Ey), (mp.Hx, mp.Hy)) + ) + + def test_other_dimensionalities_are_rejected_clearly(self): + with self.assertRaisesRegex(NotImplementedError, "2D and 3D"): + asm._tangential_components(mp.Z, 1) + with self.assertRaisesRegex(ValueError, "normal to x or y"): + asm._tangential_components(mp.Z, 2) + + +def _propagator_3d( + index=N_OXIDE, num_points=(64, 64), pitch=(0.15, 0.15), pad_factor=4 +): + """A 3D propagator with no interface, i.e. homogeneous propagation.""" + stack = mpa.Stack([mpa.Layer(index, 0.0), mpa.Layer(index)]) + return mpa.AngularSpectrum( + stack, + [1 / WAVELENGTH], + pitch, + num_points, + normal=mp.Z, + pad_factor=pad_factor, + ) + + +def _up_going_3d(propagator, amplitude_s, amplitude_p): + """Synthesizes a purely up-going 3D field from its s and p spectra. + + For an up-going wave `H_t = Y (n_hat x E_t)`, which in the s/p basis is + `H_s = Y_p E_p` and `H_p = -Y_s E_s`, so the magnetic partner is fixed once + the electric spectrum is chosen. + """ + admittance_s, admittance_p = propagator._admittances[0] + magnetic_s = admittance_p * amplitude_p + magnetic_p = -admittance_s * amplitude_s + + cosine, sine = propagator._azimuth() + + def to_axes(component_s, component_p): + # The rotation into (s, p) is a reflection, hence its own inverse. + return ( + -component_s * sine + component_p * cosine, + component_s * cosine + component_p * sine, + ) + + electric_u, electric_v = to_axes(amplitude_s, amplitude_p) + magnetic_u, magnetic_v = to_axes(magnetic_s, magnetic_p) + + positions = jnp.asarray(propagator._sample_positions()) + phase = jnp.exp(1j * jnp.einsum("xd,kd->xk", positions, propagator.kt)) + measure = propagator._spectral_measure() + shape = propagator.num_points + + def to_real_space(values): + out = jnp.einsum("xk,fk->fx", phase, values) * measure + return out.reshape(out.shape[:1] + shape) + + (e_u, e_v), (h_u, h_v) = asm._PLANE_AXES[mp.Z] + return mpa.TangentialFields( + E={e_u: to_real_space(electric_u), e_v: to_real_space(electric_v)}, + H={h_u: to_real_space(magnetic_u), h_v: to_real_space(magnetic_v)}, + normal=mp.Z, + ) + + +@unittest.skipIf(jax is None, "jax is not installed") +class TestThreeDimensions(ApproxComparisonTestCase): + """A rectangular monitor, where the transverse space is two-dimensional. + + The extra ingredient over 2D is the rotation into the s and p directions of + each transverse wavevector, whose azimuth is undefined at normal incidence. + """ + + def _linear_gaussian(self, propagator, waist, shift=0.0): + """A linearly polarized circular Gaussian, as s and p spectra. + + Polarized along the second tangential axis. Its s and p content varies + as cos and sin of the azimuth: a beam whose spectrum is purely s at + every wavevector would be azimuthally polarized, which carries a vortex + at normal incidence and does not decay compactly, so it truncates badly + on any finite monitor. + """ + kt = propagator.kt + envelope = jnp.exp(-jnp.sum(jnp.square(kt * waist), axis=-1)[None, :] / 4.0) * ( + waist**2 * onp.pi + ) + if shift: + envelope = envelope * jnp.exp(1j * kt[:, 0] * shift)[None, :] + cosine, sine = propagator._azimuth() + return envelope * cosine, envelope * sine + + def test_decomposition_recovers_a_purely_up_going_field(self): + """Both polarizations, displaced off axis, with nothing left behind.""" + propagator = _propagator_3d() + along_v = self._linear_gaussian(propagator, 1.2, shift=0.7) + along_u = (-along_v[1], along_v[0]) # the orthogonal linear polarization + for name, (amplitude_s, amplitude_p) in ( + ("polarized along v", along_v), + ("polarized along u", along_u), + ("mixed", (along_v[0] + 0.4 * along_u[0], along_v[1] + 0.4 * along_u[1])), + ): + fields = _up_going_3d(propagator, amplitude_s, amplitude_p) + up, down = propagator.decompose(fields) + scale = max( + float(jnp.max(jnp.abs(up[asm.S_POLARIZATION]))), + float(jnp.max(jnp.abs(up[asm.P_POLARIZATION]))), + ) + residual = max( + float(jnp.max(jnp.abs(down[asm.S_POLARIZATION]))), + float(jnp.max(jnp.abs(down[asm.P_POLARIZATION]))), + ) + # As in 2D, the floor is set by what the beam leaves at the edges of + # the monitor rather than by the decomposition, since the transform + # is periodic and wraps whatever is still there. Assert against the + # measured edge amplitude so the bound is not an arbitrary number. + edge = float(onp.max(propagator.report(fields)["edge_amplitude"])) + self.assertLess(residual / scale, max(10 * edge, 1e-9), name) + + def test_normal_incidence_azimuth_is_finite(self): + """The plane of incidence is undefined at kt = 0, which must not produce nan.""" + propagator = _propagator_3d() + cosine, sine = propagator._azimuth() + self.assertTrue(bool(jnp.all(jnp.isfinite(cosine)))) + self.assertTrue(bool(jnp.all(jnp.isfinite(sine)))) + at_origin = int(jnp.argmin(jnp.linalg.norm(propagator.kt, axis=-1))) + self.assertAlmostEqual(float(jnp.linalg.norm(propagator.kt[at_origin])), 0.0) + self.assertAlmostEqual(float(cosine[at_origin]), 1.0) + self.assertAlmostEqual(float(sine[at_origin]), 0.0) + + def test_gaussian_beam_spreading(self): + """A circular beam spreads as w0 sqrt(1 + (z/zR)^2), as in 2D.""" + waist = 1.2 + propagator = _propagator_3d(pad_factor=8) + fields = _up_going_3d(propagator, *self._linear_gaussian(propagator, waist)) + rayleigh = onp.pi * waist**2 * N_OXIDE / WAVELENGTH + distance = 8.0 + expected = waist * math.sqrt(1 + (distance / rayleigh) ** 2) + line = onp.linspace(-3 * expected, 3 * expected, 601) + output = propagator.propagate( + fields, distance, coordinates=[line, onp.zeros(1)] + ) + profile = onp.abs(onp.asarray(output[mp.Ey][0, :, 0])) + above = line[profile >= profile.max() / onp.e] + measured = (above.max() - above.min()) / 2 + # Read off a sampled profile on a 601-point line, so the 1/e crossing is + # located to about the sample spacing; 2% is that discretization, not + # the propagation, which 2D checks to six places on a finer line. + self.assertLess(abs(measured / expected - 1.0), 0.02) + + def test_self_overlap_is_unity(self): + propagator = _propagator_3d() + fields = _up_going_3d(propagator, *self._linear_gaussian(propagator, 1.2)) + result = propagator.spectrum(fields, 5.0) + itself = mpa.Mode(spectrum=lambda p, kt, k0, index: result.amplitudes) + self.assertAlmostEqual( + float(propagator.overlap(fields, itself, 5.0)[0]), 1.0, places=10 + ) + + def test_gaussian_mode_overlap_is_high_for_a_matched_beam(self): + """A circular Gaussian projected onto a matched circular mode.""" + waist = 1.2 + propagator = _propagator_3d(pad_factor=8) + fields = _up_going_3d(propagator, *self._linear_gaussian(propagator, waist)) + efficiency = float( + propagator.overlap(fields, mpa.gaussian_mode(waist), 1e-9)[0] + ) + self.assertGreater(efficiency, 0.99) + + def test_gradients_flow(self): + """Layer thickness and fiber tilt carry finite, nonzero gradients in 3D.""" + propagator = _propagator_3d(num_points=(32, 32), pad_factor=4) + fields = _up_going_3d(propagator, *self._linear_gaussian(propagator, 1.2)) + + def objective(tilt): + return propagator.overlap( + fields, mpa.gaussian_mode(1.2, tilt_deg=tilt), 5.0 + )[0] + + gradient = jax.grad(objective)(4.0) + self.assertTrue(onp.isfinite(float(gradient))) + self.assertNotEqual(float(gradient), 0.0) @unittest.skipIf(jax is None, "jax is not installed") @@ -453,7 +636,8 @@ def _run(self, resolution): source = [ mp.Source( mp.GaussianSource(frequency, fwidth=0.1 * frequency), - component=mp.Ez, center=mp.Vector3(0, -1.0), + component=mp.Ez, + center=mp.Vector3(0, -1.0), size=mp.Vector3(12.0, 0), # Narrow and apodized, so the field has decayed by the ends of # the monitor. A bare dipole never does, and near2far tolerates @@ -462,17 +646,15 @@ def _run(self, resolution): ) ] simulation = mp.Simulation( - resolution=resolution, cell_size=mp.Vector3(cell_x, cell_y), - boundary_layers=[mp.PML(pml)], sources=source, + resolution=resolution, + cell_size=mp.Vector3(cell_x, cell_y), + boundary_layers=[mp.PML(pml)], + sources=source, default_material=mp.Medium(index=N_OXIDE), ) height, width = 1.0, cell_x - 2 * pml - 0.4 - volume = mp.Volume( - center=mp.Vector3(0, height), size=mp.Vector3(width, 0) - ) - monitor = simulation.add_dft_fields( - [mp.Ez, mp.Hx], [frequency], where=volume - ) + volume = mp.Volume(center=mp.Vector3(0, height), size=mp.Vector3(width, 0)) + monitor = simulation.add_dft_fields([mp.Ez, mp.Hx], [frequency], where=volume) near2far = simulation.add_near2far( [frequency], mp.Near2FarRegion( @@ -516,3 +698,150 @@ def test_agrees_and_converges(self): if __name__ == "__main__": unittest.main() + + +@unittest.skipIf(jax is None, "jax is not installed") +class TestAdjointGradient(ApproxComparisonTestCase): + """The adjoint gradient of an angular-spectrum objective, in 2D and 3D. + + Compared against a finite difference of the same objective, which is the + only check that the cotangents reach the design region with the right sign + and magnitude through the whole chain: monitor, transform, decomposition, + stack, and mode projection. + """ + + def _solve(self, dimensions, weights, need_gradient=True): + resolution = 12 if dimensions == 2 else 10 + frequency = 1 / WAVELENGTH + silicon = mp.Medium(index=3.0) + design_extent, thickness = 1.0, 0.4 + grid = int(design_extent * resolution) + 1 + counts = (grid,) if dimensions == 2 else (grid, grid) + + material = mp.MaterialGrid( + mp.Vector3(*counts, *([1] * (3 - len(counts)))), + mp.Medium(index=N_OXIDE), + silicon, + weights=onp.asarray(weights).reshape(counts), + do_averaging=False, + beta=0, + ) + if dimensions == 2: + cell = mp.Vector3(6.0, 6.0) + design = mp.Volume( + center=mp.Vector3(), size=mp.Vector3(design_extent, thickness) + ) + source_center, source_size = mp.Vector3(0, -1.5), mp.Vector3(3.0, 0) + plane = mp.Volume(center=mp.Vector3(0, 1.2), size=mp.Vector3(3.2, 0)) + component, normal = mp.Ez, mp.Y + else: + cell = mp.Vector3(4.0, 4.0, 4.0) + design = mp.Volume( + center=mp.Vector3(), + size=mp.Vector3(design_extent, design_extent, thickness), + ) + source_center, source_size = mp.Vector3(0, 0, -0.9), mp.Vector3(1.8, 1.8, 0) + plane = mp.Volume( + center=mp.Vector3(0, 0, 0.8), size=mp.Vector3(2.0, 2.0, 0) + ) + component, normal = mp.Ex, mp.Z + + simulation = mp.Simulation( + resolution=resolution, + cell_size=cell, + dimensions=dimensions, + boundary_layers=[mp.PML(1.0)], + default_material=mp.Medium(index=N_OXIDE), + geometry=[ + mp.Block(center=design.center, size=design.size, material=material) + ], + sources=[ + mp.Source( + mp.GaussianSource(frequency, fwidth=0.2 * frequency), + component=component, + center=source_center, + size=source_size, + ) + ], + ) + run_length = 260 if dimensions == 2 else 220 + simulation.init_sim() + stack = mpa.Stack([mpa.Layer(N_OXIDE, 0.5), mpa.Layer(N_AIR)]) + metadata = simulation.get_array_metadata(vol=plane) + counts = ( + (len(metadata[0]),) + if dimensions == 2 + else (len(metadata[0]), len(metadata[1])) + ) + propagator = mpa.AngularSpectrum( + stack, + [frequency], + pitch=1.0 / resolution, + num_points=counts, + normal=normal, + pad_factor=4, + ) + # The source picks the polarization: out-of-plane Ez in 2D is s, and the + # in-plane Ex beam in 3D is polarized along the first tangential axis. + fiber = mpa.gaussian_mode( + 1.0, + polarization=( + asm.S_POLARIZATION if dimensions == 2 else asm.P_POLARIZATION + ), + ) + + def objective(*monitor_values): + fields = propagator.take(monitor_values) + return propagator.overlap(fields, fiber, 2.0)[0] + + optimization = mpa.OptimizationProblem( + simulation=simulation, + objective_functions=objective, + objective_arguments=propagator.objective_arguments(simulation, plane), + design_regions=[mpa.DesignRegion(material, volume=design)], + frequencies=[frequency], + # A fixed run length, so both runs of the finite difference cover + # exactly the same interval. Left adaptive, `stop_when_dft_decayed` + # lets the perturbed run stop at a different time from the + # unperturbed one, and that difference scales with the perturbation + # -- producing a fixed ratio between the adjoint gradient and the + # finite difference that does not shrink with the step size, and so + # reads exactly like a wrong gradient. + # + # The binding constraint is the *adjoint* DFT rather than the + # forward one. Holding this geometry fixed and lengthening the run, + # the finite difference does not move (4.923e-6 throughout) while + # the adjoint gradient falls 8.84e-6, 7.92e-6, 4.93e-6 at lengths + # 40, 80 and 140. Below is chosen with margin on that. + decay_by=1e-11, + minimum_run_time=run_length, + maximum_run_time=run_length, + ) + return optimization([onp.asarray(weights).ravel()], need_gradient=need_gradient) + + def _check(self, dimensions): + rng = onp.random.RandomState(3) + grid = int(1.0 * (12 if dimensions == 2 else 10)) + 1 + size = grid if dimensions == 2 else grid * grid + weights = 0.5 * rng.rand(size) + perturbation = 1e-4 * rng.rand(size) + + value, gradient = self._solve(dimensions, weights) + perturbed, _ = self._solve( + dimensions, weights + perturbation, need_gradient=False + ) + adjoint = float(onp.dot(perturbation, onp.asarray(gradient).ravel())) + difference = float(onp.asarray(perturbed) - onp.asarray(value)) + tolerance = 0.1 if mp.is_single_precision() else 0.01 + self.assertClose( + onp.array([adjoint]), + onp.array([difference]), + epsilon=tolerance, + msg=f"{dimensions}D", + ) + + def test_two_dimensions(self): + self._check(2) + + def test_three_dimensions(self): + self._check(3) From 9a97599459e0371e99bebbc27b8b90027b887af9 Mon Sep 17 00:00:00 2001 From: Alec Hammond Date: Wed, 26 Aug 2026 19:18:10 -0700 Subject: [PATCH 7/8] adjoint: reconstruct real-space fields by inverse FFT, not a dense matrix CI killed test_angular_spectrum outright -- the OOM killer, not an assertion. propagate() summed the spectrum by building an explicit (num samples, num wavevectors) phase matrix. That is tolerable in 2D and untenable in 3D: a 64x64 monitor padded eightfold carries 262144 wavevectors, so evaluating it back on its own 4096 samples wants 17 GB. It ran here only because this machine has the memory to hide it; a 7 GB runner does not. Reconstructing on the monitor's own grid is exactly the inverse of the forward transform, so it is now an inverse FFT -- no dense matrix, and O(N log N). Arbitrary coordinates still need the direct sum, but it is chunked to a bounded number of entries per block, so asking for an inconvenient set of points costs time rather than the process. The 3D tests built the same matrix in their own helper and now use the same inverse transform. Their grids are also smaller, which is enough for what they check and takes the suite from 103 s to 77 s. test_angular_spectrum is also dropped from TESTS and left in ADJOINT_TESTS alone, matching test_adjoint_solver. It contains FDTD adjoint checks and has no business running in all four build jobs. --- python/Makefile.am | 1 - python/adjoint/angular_spectrum.py | 40 ++++++++++++++++++++++++--- python/tests/test_angular_spectrum.py | 27 +++++------------- 3 files changed, 43 insertions(+), 25 deletions(-) diff --git a/python/Makefile.am b/python/Makefile.am index ab75fbb33..e55cda540 100644 --- a/python/Makefile.am +++ b/python/Makefile.am @@ -46,7 +46,6 @@ TESTS = \ $(TEST_DIR)/test_adjoint_chunks.py \ $(TEST_DIR)/test_adjoint_symmetric_grids.py \ $(TEST_DIR)/test_adjoint_protocol.py \ - $(TEST_DIR)/test_angular_spectrum.py \ $(TEST_DIR)/test_antenna_radiation.py \ $(TEST_DIR)/test_array_metadata.py \ $(TEST_DIR)/test_bend_flux.py \ diff --git a/python/adjoint/angular_spectrum.py b/python/adjoint/angular_spectrum.py index 7d1e56e93..2b7b7cad3 100644 --- a/python/adjoint/angular_spectrum.py +++ b/python/adjoint/angular_spectrum.py @@ -924,18 +924,50 @@ def propagate(self, fields: TangentialFields, distance: float, coordinates=None) amplitude_u = -amplitude_s * sine + amplitude_p * cosine amplitude_v = amplitude_s * cosine + amplitude_p * sine - phase = jnp.exp(1j * jnp.einsum("xd,kd->xk", positions, result.kt)) measure = self._spectral_measure() (e_u, e_v), _ = _PLANE_AXES[self.normal] - shape = tuple(len(a) for a in axes) outputs = {} for component, amplitude in ((e_u, amplitude_u), (e_v, amplitude_v)): if not jnp.any(amplitude): continue - values = jnp.einsum("xk,fk->fx", phase, amplitude) * measure - outputs[component] = values.reshape(values.shape[:1] + shape) + if coordinates is None and self._uniform: + outputs[component] = self._inverse_transform(amplitude) + else: + outputs[component] = self._evaluate_at(amplitude, axes, measure) return outputs + def _inverse_transform(self, amplitude: jnp.ndarray) -> jnp.ndarray: + """Inverts the forward transform back onto the monitor samples. + + The exact inverse of `_transform`, and the only affordable route in 3D: + evaluating the sum directly would need a dense (num samples, num + wavevectors) matrix, which for a 64x64 monitor padded eightfold is 17 GB. + """ + origins = jnp.asarray([axis[0] for axis in self._coordinate_axes()]) + spectrum = amplitude * jnp.exp(1j * (self._kt @ origins)) + spectrum = spectrum.reshape(spectrum.shape[:1] + self._padded) + transverse = tuple(range(-self.transverse_dimensions, 0)) + values = jnp.fft.ifftn(spectrum, axes=transverse) / float(onp.prod(self.pitch)) + return values[(slice(None),) + tuple(slice(0, n) for n in self.num_points)] + + def _evaluate_at(self, amplitude, axes, measure, max_entries=2**26): + """Sums the spectrum at arbitrary coordinates, in bounded-size chunks. + + The dense phase matrix is (num coordinates, num wavevectors), so it has + to be chunked or it will exhaust memory well before it is slow. + """ + grids = onp.meshgrid(*axes, indexing="ij") + positions = onp.stack([g.ravel() for g in grids], axis=-1) + shape = tuple(len(axis) for axis in axes) + chunk = max(1, int(max_entries // max(1, self._kt.shape[0]))) + pieces = [] + for start in range(0, positions.shape[0], chunk): + block = jnp.asarray(positions[start : start + chunk]) + phase = jnp.exp(1j * jnp.einsum("xd,kd->xk", block, self._kt)) + pieces.append(jnp.einsum("xk,fk->fx", phase, amplitude) * measure) + values = jnp.concatenate(pieces, axis=-1) + return values.reshape(values.shape[:1] + shape) + # ---------------------------------------------------------------- Meep glue # # Two ways in. `from_monitor` post-processes an ordinary forward run and diff --git a/python/tests/test_angular_spectrum.py b/python/tests/test_angular_spectrum.py index 1246ff648..4174c6317 100644 --- a/python/tests/test_angular_spectrum.py +++ b/python/tests/test_angular_spectrum.py @@ -211,15 +211,7 @@ def _up_going(propagator, values): """ spectrum = propagator._transform(jnp.asarray(values)) admittance = propagator._admittances[0][asm.S_POLARIZATION] - positions = jnp.asarray(propagator._sample_positions()) - magnetic = ( - jnp.einsum( - "xk,fk->fx", - jnp.exp(1j * jnp.einsum("xd,kd->xk", positions, propagator.kt)), - spectrum * admittance, - ) - * propagator._spectral_measure() - ) + magnetic = propagator._inverse_transform(spectrum * admittance) return mpa.TangentialFields( E={mp.Ez: jnp.asarray(values)}, H={mp.Hx: magnetic}, normal=mp.Y ) @@ -485,14 +477,9 @@ def to_axes(component_s, component_p): electric_u, electric_v = to_axes(amplitude_s, amplitude_p) magnetic_u, magnetic_v = to_axes(magnetic_s, magnetic_p) - positions = jnp.asarray(propagator._sample_positions()) - phase = jnp.exp(1j * jnp.einsum("xd,kd->xk", positions, propagator.kt)) - measure = propagator._spectral_measure() - shape = propagator.num_points - - def to_real_space(values): - out = jnp.einsum("xk,fk->fx", phase, values) * measure - return out.reshape(out.shape[:1] + shape) + # Via the inverse FFT rather than a dense phase matrix: at these sizes the + # matrix is (num samples, num wavevectors) and runs to many gigabytes. + to_real_space = propagator._inverse_transform (e_u, e_v), (h_u, h_v) = asm._PLANE_AXES[mp.Z] return mpa.TangentialFields( @@ -569,12 +556,12 @@ def test_normal_incidence_azimuth_is_finite(self): def test_gaussian_beam_spreading(self): """A circular beam spreads as w0 sqrt(1 + (z/zR)^2), as in 2D.""" waist = 1.2 - propagator = _propagator_3d(pad_factor=8) + propagator = _propagator_3d(pad_factor=6) fields = _up_going_3d(propagator, *self._linear_gaussian(propagator, waist)) rayleigh = onp.pi * waist**2 * N_OXIDE / WAVELENGTH distance = 8.0 expected = waist * math.sqrt(1 + (distance / rayleigh) ** 2) - line = onp.linspace(-3 * expected, 3 * expected, 601) + line = onp.linspace(-3 * expected, 3 * expected, 241) output = propagator.propagate( fields, distance, coordinates=[line, onp.zeros(1)] ) @@ -598,7 +585,7 @@ def test_self_overlap_is_unity(self): def test_gaussian_mode_overlap_is_high_for_a_matched_beam(self): """A circular Gaussian projected onto a matched circular mode.""" waist = 1.2 - propagator = _propagator_3d(pad_factor=8) + propagator = _propagator_3d(pad_factor=6) fields = _up_going_3d(propagator, *self._linear_gaussian(propagator, waist)) efficiency = float( propagator.overlap(fields, mpa.gaussian_mode(waist), 1e-9)[0] From 22d24e6725c6c3e346f61ee3aaebb8926f4c21ac Mon Sep 17 00:00:00 2001 From: Alec Hammond Date: Wed, 26 Aug 2026 19:25:34 -0700 Subject: [PATCH 8/8] adjoint: cache the transform kernel for an explicit wavevector set The non-uniform transform rebuilt its (num wavevectors, num samples) matrix of complex exponentials on every call, though it depends only on the wavevectors and the sample positions, both fixed at construction. It dominated. On a 120x120 monitor, selecting the 1517 wavevectors inside a 0.2 numerical aperture took 902 ms per overlap, against 66 ms for a padded FFT over 230400 wavevectors -- the supposedly cheap path was thirteen times slower than the one it was meant to undercut. Caching brings it to 228 ms. What remains is the O(num wavevectors x num samples) contraction itself, which is inherent. So an explicit wavevector set is worth reaching for when it is much smaller than the sample count, or when the padded grid would be unaffordable -- notably long propagation in 3D, where padding costs the square of the pad factor -- and not otherwise. The cache costs one complex matrix of that size, which is another reason it suits a small set. --- python/adjoint/angular_spectrum.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/python/adjoint/angular_spectrum.py b/python/adjoint/angular_spectrum.py index 2b7b7cad3..fad8b7587 100644 --- a/python/adjoint/angular_spectrum.py +++ b/python/adjoint/angular_spectrum.py @@ -487,6 +487,7 @@ def __init__( self._kt = jnp.asarray(kt.reshape(kt.shape[0], -1)) self._uniform = False self._padded = None + self._cached_forward_phase = None # Regularizing the index keeps kz off the light line, where its # derivative is unbounded; see DEFAULT_LOSS_REGULARIZATION. @@ -547,6 +548,21 @@ def _sample_positions(self) -> onp.ndarray: grids = onp.meshgrid(*self._coordinate_axes(), indexing="ij") return onp.stack([g.ravel() for g in grids], axis=-1) + def _forward_phase(self) -> jnp.ndarray: + """The transform kernel for an explicit set of wavevectors, cached. + + It depends only on the wavevectors and the sample positions, both fixed + at construction, so rebuilding it per call is pure waste -- and it + dominates: the kernel is (num wavevectors, num samples) complex + exponentials, which for a few thousand wavevectors over a plane costs + more than an FFT over a grid a hundred times larger. + """ + if self._cached_forward_phase is None: + self._cached_forward_phase = jnp.exp( + -1j * jnp.einsum("kd,xd->kx", self._kt, self._sample_positions()) + ) + return self._cached_forward_phase + def _transform(self, values: jnp.ndarray) -> jnp.ndarray: """Transforms monitor samples to (num frequencies, num kt).""" values = jnp.asarray(values) @@ -554,10 +570,7 @@ def _transform(self, values: jnp.ndarray) -> jnp.ndarray: measure = float(onp.prod(self.pitch)) if not self._uniform: flat = values.reshape(values.shape[: transverse[0]] + (-1,)) - phase = jnp.exp( - -1j * jnp.einsum("kd,xd->kx", self._kt, self._sample_positions()) - ) - return jnp.einsum("kx,...x->...k", phase, flat) * measure + return jnp.einsum("kx,...x->...k", self._forward_phase(), flat) * measure padded = jnp.zeros( values.shape[: transverse[0]] + self._padded, dtype=jnp.complex128 )