[4/n][Adjoint Module] Launch a mode through a stratified stack into the simulation (and backprop)

python/adjoint/angular_spectrum.py

@@ -36,7 +36,7 @@
 """
 
 import math
-from typing import Callable, Dict, NamedTuple, Optional, Sequence, Tuple, Union
+from typing import Callable, Dict, NamedTuple, Optional, Sequence, Tuple, Union, List
 
 import jax
 import jax.numpy as jnp
@@ -1083,6 +1083,42 @@ def from_monitor(
             **kwargs,
         )
 
+    @classmethod
+    def from_volume(
+        cls,
+        simulation: mp.Simulation,
+        volume: mp.Volume,
+        stack: Stack,
+        frequencies,
+        sign: int = 1,
+        **kwargs,
+    ):
+        """Builds a propagator for a plane, without needing a monitor there.
+
+        `from_monitor` cannot be used on the injection path: a DFT monitor
+        requires the field components to be allocated, which does not happen
+        until sources are added, and the sources are what this is being built to
+        produce. This takes the frequencies directly instead.
+
+        Args:
+            simulation: supplies the grid; it need not have been run.
+            volume: the source plane.
+            stack: the layers between that plane and the mode.
+            frequencies: the frequencies to propagate.
+            sign: +1 if the outgoing direction is along +normal.
+            **kwargs: forwarded to the constructor, e.g. `pad_factor`.
+        """
+        normal, pitch, num_points = cls._plane_geometry(simulation, volume)
+        return cls(
+            stack,
+            onp.asarray(frequencies),
+            pitch,
+            num_points,
+            normal=normal,
+            sign=sign,
+            **kwargs,
+        )
+
     def fields_from_monitor(
         self, simulation: mp.Simulation, monitor
     ) -> TangentialFields:
@@ -1152,6 +1188,167 @@ def report_monitor(self, simulation, monitor) -> Dict[str, onp.ndarray]:
         fields = self.fields_from_monitor(simulation, monitor)
         return {key: onp.asarray(value) for key, value in self.report(fields).items()}
 
+    def incident_fields(self, mode: Mode, distance: float) -> TangentialFields:
+        """Tangential fields at the monitor plane for a mode arriving from afar.
+
+        The time reverse of `spectrum`. A `Mode` is specified where it is
+        physically meaningful -- at the fiber facet, say, on the far side of the
+        stack -- and this carries it back down through the layers to the monitor
+        plane, where Meep can inject it. The scalar transfer function of a
+        stratified medium is the same in both directions, so the same product of
+        transmission and propagation phases applies.
+
+        The result is purely in-going: its up-going content is zero by
+        construction, so injecting it launches a beam toward the structure and
+        not away from it.
+
+        Args:
+            mode: the target field, e.g. from `gaussian_mode`.
+            distance: how far the mode's plane is from the monitor, the same
+                convention `spectrum` uses.
+
+        Returns:
+            A `TangentialFields` on the monitor's own sample grid.
+        """
+        remaining = distance - self.stack_thickness
+        concrete = _as_concrete(remaining)
+        if concrete is not None and concrete < 0:
+            raise ValueError(
+                f"distance={distance} does not clear the stack, which is "
+                f"{_as_concrete(self.stack_thickness)} thick."
+            )
+
+        k0 = 2 * onp.pi * jnp.asarray(self.frequencies)[:, None]
+        amplitudes = mode.spectrum(self, self._kt, k0, self._indices[-1])
+
+        # Carry each polarization across the stack to the monitor plane.
+        down = {}
+        for polarization in (S_POLARIZATION, P_POLARIZATION):
+            transmission, _ = self._transmission(polarization)
+            to_interface = jnp.exp(1j * self._wavevectors[0] * self._thicknesses[0])
+            beyond = jnp.exp(1j * self._wavevectors[-1] * remaining)
+            down[polarization] = (
+                amplitudes[..., polarization] * to_interface * transmission * beyond
+            )
+
+        # Invert `decompose` for a spectrum with no up-going part. Setting
+        # up = (E - cross)/2 to zero gives E = cross = down, and the admittance
+        # relations then fix the magnetic components.
+        admittance_s, admittance_p = self._admittances[0]
+        outgoing = self.sign
+        electric_s = down[S_POLARIZATION]
+        electric_p = down[P_POLARIZATION]
+
+        azimuth = self._azimuth()
+        if azimuth is None:
+            # cross_s = -outgoing * H_p / Y_s and cross_p = outgoing * H_s / Y_p
+            magnetic_p = -outgoing * admittance_s * electric_s
+            magnetic_s = outgoing * admittance_p * electric_p
+            electric_u, electric_v = electric_s, electric_p
+            magnetic_u, magnetic_v = magnetic_s, magnetic_p
+        else:
+            # cross_s = outgoing * H_p / Y_s and cross_p = -outgoing * H_s / Y_p
+            magnetic_p = outgoing * admittance_s * electric_s
+            magnetic_s = -outgoing * admittance_p * electric_p
+            cosine, sine = azimuth
+            # the rotation into (s, p) is a reflection, hence its own inverse
+            electric_u = -electric_s * sine + electric_p * cosine
+            electric_v = electric_s * cosine + electric_p * sine
+            magnetic_u = -magnetic_s * sine + magnetic_p * cosine
+            magnetic_v = magnetic_s * cosine + magnetic_p * sine
+
+        (e_u, e_v), (h_u, h_v) = _PLANE_AXES[self.normal]
+        return TangentialFields(
+            E={
+                e_u: self._inverse_transform(electric_u),
+                e_v: self._inverse_transform(electric_v),
+            },
+            H={
+                # `_tangential_spectra` pairs E_u with H_v, so undo that here
+                h_v: self._inverse_transform(magnetic_v),
+                h_u: self._inverse_transform(magnetic_u),
+            },
+            normal=self.normal,
+            sign=self.sign,
+        )
+
+    def equivalent_sources(
+        self,
+        fields: TangentialFields,
+        time_src,
+        center,
+        size,
+        frequency: Optional[float] = None,
+        amplitude: complex = 1.0,
+        differentiable: bool = True,
+    ) -> List[mp.Source]:
+        """Equivalent surface currents that launch `fields` into the simulation.
+
+        A thin wrapper over `mp.get_equiv_sources`, which already applies the
+        equivalence principle -- an electric sheet `K = n_hat x H` and a
+        magnetic sheet `N = -n_hat x E` -- and emits ordinary sources carrying
+        `amp_data`. Both sheets are needed: one alone radiates half the field in
+        each direction instead of all of it in one.
+
+        Since `amp_data` is differentiable, the sources it returns backpropagate
+        to whatever produced `fields`, which for the angular spectrum means the
+        mode's waist, tilt and offset.
+
+        Args:
+            fields: tangential fields on the plane, e.g. from `incident_fields`.
+            time_src: the time profile, e.g. `mp.GaussianSource`.
+            center, size: the source plane, normally the same as the monitor.
+            frequency: which frequency of `fields` to inject. Defaults to the
+                only one when the propagator carries a single frequency.
+            amplitude: an overall scale.
+            differentiable: flag the sources for the adjoint solver.
+
+        Returns:
+            A list of `mp.Source`, one per non-zero current component.
+        """
+        frequencies = onp.asarray(self.frequencies)
+        if frequency is None:
+            if frequencies.size != 1:
+                raise ValueError(
+                    f"The propagator carries {frequencies.size} frequencies, so "
+                    "`frequency` has to say which one to inject; a single Meep "
+                    "source has one time profile and cannot carry independent "
+                    "amplitudes at several."
+                )
+            index = 0
+        else:
+            index = int(onp.argmin(onp.abs(frequencies - frequency)))
+
+        def slice_at(store, component):
+            value = store.get(component)
+            if value is None:
+                return onp.zeros(self.num_points)
+            return onp.asarray(value)[index] * amplitude
+
+        # get_equiv_sources wants all six components, in order, shaped for
+        # amp_data's trilinear interpolation.
+        def shaped(values):
+            values = onp.asarray(values, dtype=onp.complex128)
+            spatial = tuple(values.shape) if values.ndim else (1,)
+            return values.reshape(spatial + (1,) * (3 - len(spatial)))
+
+        field = [shaped(slice_at(fields.E, c)) for c in (mp.Ex, mp.Ey, mp.Ez)] + [
+            shaped(slice_at(fields.H, c)) for c in (mp.Hx, mp.Hy, mp.Hz)
+        ]
+
+        normal_axis = {mp.X: 0, mp.Y: 1, mp.Z: 2}[self.normal]
+        n_hat = onp.zeros(3)
+        # `sign` is the outgoing direction, so the beam travels along its
+        # negation and that is the normal the equivalence principle wants.
+        n_hat[normal_axis] = -self.sign
+
+        sources = mp.get_equiv_sources(field, n_hat, time_src, center, size)
+        if differentiable:
+            for source in sources:
+                source.differentiable = ("amp_data",)
+                source.name = f"asm_{mp.component_name(source.component)}"
+        return sources
+
     def objective_arguments(self, simulation, volume, **kwargs):
         """The `FourierFields` an objective function needs, for the adjoint path.