diff --git a/NEWS.md b/NEWS.md index 9acfa67bd..a8d273ad5 100644 --- a/NEWS.md +++ b/NEWS.md @@ -2,6 +2,14 @@ ## 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. 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 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..baf5b0080 100644 --- a/doc/docs/Python_Tutorials/Adjoint_Solver.md +++ b/doc/docs/Python_Tutorials/Adjoint_Solver.md @@ -273,6 +273,95 @@ 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- 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 +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 ------------------------------------------------------------ diff --git a/python/Makefile.am b/python/Makefile.am index a5c6553c9..e55cda540 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 = \ @@ -245,6 +246,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 new file mode 100644 index 000000000..fad8b7587 --- /dev/null +++ b/python/adjoint/angular_spectrum.py @@ -0,0 +1,1196 @@ +"""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 _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. + + 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. + + +# 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 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}." + ) + 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(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 + + +def gaussian_mode( + waist: float, + tilt_deg: float = 0.0, + offset: float = 0.0, + polarization: int = S_POLARIZATION, + tilt_azimuth_deg: float = 0.0, +) -> 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 + 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, + 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(propagator, kt, k0, index): + kt = jnp.asarray(kt) + transverse = ( + jnp.asarray(index).real + * jnp.asarray(k0) + * jnp.sin(jnp.deg2rad(jnp.asarray(tilt_deg))) + ) + 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) + + +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 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.normal = normal + self.sign = int(sign) + self.pad_factor = int(pad_factor) + + self._k0 = 2 * onp.pi * self.frequencies + if kt is None: + 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 + else: + kt = onp.asarray(kt, dtype=float) + 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. + 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) + ] + # 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 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): + """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): + """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 _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) + 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,)) + return jnp.einsum("kx,...x->...k", self._forward_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}." + ) + (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 " + f"{mp.component_name(e)} and {mp.component_name(h)} was " + "supplied." + ) + 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. + """ + 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): + """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 + concrete = _as_concrete(remaining) + if concrete is not None and concrete < 0: + raise ValueError( + f"distance={distance} does not clear the stack, which is " + f"{_as_concrete(self.stack_thickness)} thick. 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: + 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.""" + 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) + field_amplitude = result.amplitudes + mode_amplitude = mode.spectrum(self, result.kt, self._k0, result.index) + + 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=axes) + 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 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 { + "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): + """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, 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 an array of + shape (num frequencies,) + the coordinate shape. + """ + result = self.spectrum(fields, distance) + if coordinates is None: + 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 + + measure = self._spectral_measure() + (e_u, e_v), _ = _PLANE_AXES[self.normal] + outputs = {} + for component, amplitude in ((e_u, amplitude_u), (e_v, amplitude_v)): + if not jnp.any(amplitude): + continue + 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 + # 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 counts of a planar volume.""" + size = [volume.size.x, volume.size.y, volume.size.z] + dimensions = simulation.dimensions + if dimensions not in (2, 3): + raise NotImplementedError( + "Angular-spectrum propagation supports 2D and 3D Cartesian " + f"simulations, not {dimensions}D." + ) + 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 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, 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( + 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 group in _PLANE_AXES[self.normal] for component in group + } + ) + (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), + ): + 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 group in _PLANE_AXES[self.normal] for component in group + ] + 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 4 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..03616956e --- /dev/null +++ b/python/examples/adjoint_optimization/grating_coupler_asm.py @@ -0,0 +1,437 @@ +"""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) diff --git a/python/tests/test_angular_spectrum.py b/python/tests/test_angular_spectrum.py new file mode 100644 index 000000000..4174c6317 --- /dev/null +++ b/python/tests/test_angular_spectrum.py @@ -0,0 +1,834 @@ +"""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] + magnetic = propagator._inverse_transform(spectrum * admittance) + 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 propagator, kt, k0, index: result.amplitudes) + 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_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) + + # 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( + 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=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, 241) + 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=6) + 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") +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() + + +@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)