[2/n][Adjoint Module] Differentiable angular-spectrum propagation through stratified media - #3281
Draft
smartalecH wants to merge 8 commits into
Draft
[2/n][Adjoint Module] Differentiable angular-spectrum propagation through stratified media#3281smartalecH wants to merge 8 commits into
smartalecH wants to merge 8 commits into
Conversation
Meep's near-to-far transformation requires its surface to sit in a homogeneous
medium -- dft_near2far aborts otherwise -- so a structure radiating across a
material interface has to keep that interface inside the FDTD cell. A grating
coupler radiating up through a cladding into air, and then hundreds of microns
to a fiber, cannot afford that.
This propagates the tangential DFT fields on a planar monitor through an
arbitrary layer stack analytically instead, in JAX, so the layers above the
monitor leave the simulation and the stack itself becomes differentiable.
Physics, each piece checked against an independent oracle before the next was
written:
* the stack is solved by a scattering-matrix recursion, not a transfer matrix,
which would overflow on the first thick layer or evanescent order and
surface as NaN gradients rather than as wrong fields. Agrees with analytic
Fresnel to 3e-15, puts Brewster at 2e-16, nulls a quarter-wave
antireflection coating to 8e-17, matches an independently written transfer
matrix to 1e-15 in both polarizations, and conserves energy to 4e-16.
* up- and down-going radiation are separated using both tangential fields,
which an open near-to-far surface cannot do. For an up-going wave
H_t = Y (n_hat x E_t) in both polarizations, so the split is uniform once
written with the cross product. Verified exact: the spurious down-going
amplitude tracks the field amplitude left at the monitor edges, 1.2e-8 for a
field decayed to 1.3e-8 there.
* free-space propagation reproduces the analytic spreading of a Gaussian beam
to machine precision at 50, 200 and 500 um, given a padded window sized to
the spread beam.
* projection onto a mode uses the power inner product, so self-overlap is
exactly 1, and the closed-form tilted Gaussian recovers a launch angle
exactly rather than to grid resolution.
Every layer gets an infinitesimal loss before kz is computed. kz vanishes on the
light line, where its derivative is unbounded: jnp.sqrt(0.0) evaluates fine but
its VJP is infinite, which becomes NaN and destroys the gradient while leaving
the objective value looking correct. A grid point lands exactly there whenever
the padded monitor width is an integer number of wavelengths in the medium,
which is not a rare coincidence when cell sizes and wavelengths are round.
report() returns the three diagnostics that catch a badly placed monitor:
down-going fraction, evanescent fraction, and the amplitude left at the monitor
edges -- the transform is periodic, so a field that has not decayed there wraps
around.
2D only for now; the 3D case additionally needs the s/p rotation by azimuth with
its removable singularity at normal incidence, and raises.
Two ways in, neither requiring the other. from_monitor() post-processes an ordinary forward run and deals in NumPy, so someone who only wants a far field never meets JAX. It infers the plane normal, sample pitch and sample count from the monitor volume, and -- mirroring the check dft_near2far makes for the same reason -- rejects a monitor that does not lie in a homogeneous medium, which catches a plane accidentally clipping a waveguide or a PML. objective_arguments()/take() build and repack the FourierFields an objective function needs. Both tangential components of both polarizations are registered: the unused ones cost two extra DFT line monitors in the forward run and nothing in the adjoint, since a monitor whose cotangent is identically zero places no adjoint source, and registering them spares the user working out which polarization their source excites. Validated against Meep's own near2far in a homogeneous medium, which is an independent implementation of the same physics: the propagated field agrees to 1.6e-2 at resolution 20 and 4.0e-3 at resolution 40. Converging at second order is the part that matters -- it is what a wrong sample pitch, a wrong coordinate origin, or a missing phase ramp would not do, since those produce a fixed error. Getting that comparison to mean anything took two attempts, both caught by the diagnostics rather than by the numbers. A point dipole close to the plane never decays across it, and a full-width beam source plus a strong scatterer spread field over the whole monitor; in both cases report()['edge_amplitude'] was 0.68 and 0.055 respectively and the disagreement sat at 25% and 4% without improving with resolution. near2far tolerates an undecayed field because it integrates currents on an open surface; the transform here is periodic and wraps around. With a narrow apodized source the edge amplitude is 5e-9 and the comparison converges.
…a fiber
A two-etch silicon grating radiates upward, crosses a few hundred microns of
silica and the silica/air interface, and is collected by a fiber. The FDTD cell
stops a micron above the device layer; everything above it is an analytic layer
stack. Meshing 300 um of glass at resolution 20 would be several hundred million
pixels in 2D, and the interface puts near2far out of reach regardless, since it
requires a homogeneous medium.
Two modes, because the propagator is useful without an optimizer:
forward one simulation, then the far field, the coupling efficiency, and
sweeps over fiber tilt, working distance and superstrate thickness.
None of the sweeps re-runs the simulation -- those parameters live
entirely in the analytic stack, which is the point.
optimize topology optimization of both etch levels against the fiber overlap,
with the objective written in jax.numpy and handed straight to
OptimizationProblem.
Two stacked design regions rather than one: a single fully etched layer radiates
roughly symmetrically, and it is the second etch level that buys directionality.
The example prints the monitor diagnostics before any efficiency and says plainly
when they are bad, because the failure they catch is quiet. Getting the geometry
right took three rounds of exactly that:
* the monitor was initially only as wide as the grating, leaving 19% of the
peak field at its ends; the transform is periodic, so that wraps.
* widening the cell barely helped -- 4.0% to 3.9% for half again the width --
which is the tell that the residual is not the beam tail but near-grazing
radiation, which travels sideways instead of decaying.
* the real remaining flaw was 0.3 um between the monitor and the PML. A plane
that close picks up the absorber's residual reflection. A wavelength of
clearance brought it to 2.4%.
A few percent is intrinsic here and bounds the accuracy at a similar level;
power at those angles was never going to reach the fiber. The example says so
rather than quietly windowing the monitor, which would hide it.
Verified end to end at reduced settings: the optimizer takes the coupling from
0.0116 to 0.0273 in five iterations, which also confirms the adjoint gradient
reaches both design regions with the right sign.
19 tests, 9 seconds for all but one, since almost everything here has a closed
form to check against rather than needing a simulation.
* the scattering recursion against Fresnel swept to 89 degrees, the Brewster
angle, a nulled quarter-wave antireflection coating, an independently
written transfer matrix, and energy conservation. The p-polarization
reference carries an explicit sign, because these are the coefficients of
the tangential field and that is the negative of the form usually written
for the full electric vector -- the two agree on |r| and on Brewster, so a
mismatch here is easy to wave away as a convention and not notice.
* the up/down split, asserted against the edge amplitude rather than against a
fixed tolerance: whatever field survives at the ends of the monitor wraps
around and reappears as spurious down-going content, so the two track each
other, and that is what makes the diagnostic meaningful.
* Gaussian spreading against w0 sqrt(1 + (z/zR)^2) to six places at 50 and
200 um, with the padded window sized to the spread beam.
* self-overlap of exactly one, and exact recovery of a launch angle, which the
closed-form mode spectrum makes possible -- a sampled mode would quantize it.
* a gradient that stays finite when a wavevector lands exactly on the light
line. The window is chosen so that one does. Without the regularizer the
value is still perfect and every gradient is NaN, which is the failure worth
having a test for.
* agreement with Meep's own near2far in a homogeneous medium, 1.6e-2 at
resolution 20 and converging by better than half at 40. The convergence is
the assertion that matters, since a wrong pitch or coordinate origin gives a
fixed error; the test also asserts the edge amplitude is small first, because
otherwise the comparison means nothing.
Writing these turned up a real gap between the module and its own documentation:
layer thicknesses were coerced with float() in the constructor, so a traced
thickness raised ConcretizationTypeError and the stack was not in fact
differentiable, as both the module docstring and the example claimed. Thicknesses
are now kept as they arrive, and the concreteness checks that guard the distance
argument skip themselves rather than force a value.
Covers what it is for -- radiation crossing a material interface, which near2far cannot handle -- the forward-only and adjoint entry points, and the two constraints that produce quiet wrong answers rather than errors: a monitor whose field has not decayed by its ends, and a padded window narrower than the spread beam. Also states plainly that a few percent of edge amplitude is normal for a grating radiating into a cladding, is near-grazing radiation rather than the beam tail, and bounds the accuracy accordingly. Notes what is not supported and why: 3D needs the s/p rotation by azimuth with its removable singularity at normal incidence and raises rather than guessing, and both half-spaces are covered with two monitors rather than a closed surface, since a single plane is already complete for the half-space above it.
A monitor plane is a line in a 2D simulation and a rectangle in a 3D one, so the
transverse space is one- or two-dimensional. Everything downstream of the
transform now works on a flat list of transverse wavevectors and is indifferent
to which it was; only the transform and the polarization basis differ.
The 3D ingredient is the rotation into the s and p directions of each transverse
wavevector, whose azimuth is undefined at normal incidence and is pinned there
rather than left to produce a nan from atan2(0, 0). Verified against explicit
up-going plane waves at 200 random oblique directions, to 4e-16.
The two transverse dimensionalities use genuinely different polarization bases,
and the sign of the decomposition differs between them: with one transverse
direction the wavevector runs along v_hat, so the 3D convention would put s_hat
along -u_hat. Unifying them without noticing inverted the 2D split -- an
up-going field landed entirely in the down-going branch -- which the existing 2D
tests caught.
Tests, 28 passing. In 3D: the decomposition recovers a purely up-going field to
the level of the monitor truncation, a circular Gaussian spreads as
w0 sqrt(1 + (z/zR)^2), self-overlap is one, a matched mode couples above 0.99,
and the azimuth is finite at normal incidence. Two of those started out failing
for a reason worth recording: a spectrum that is purely s at every wavevector is
azimuthally polarized, carries a vortex at normal incidence, and does not decay
compactly, so it truncates badly on any finite monitor. A real Gaussian beam is
linearly polarized, with s and p content varying as cos and sin of the azimuth.
The adjoint gradient is checked against a finite difference in both
dimensionalities and agrees to within 1%. Getting that check to mean anything
took some care, because it first appeared to show a wrong gradient -- a ratio of
1.37 in 2D and 1.59 in 3D, constant in the finite-difference step, which rules
out truncation error and looks like a real defect.
Isolating the 2D chain, from a bare sum of |E|^2 on the monitor up to the
self-normalized mode overlap, every stage agreed to better than 2e-4 once the
run was converged, so neither the propagator nor the ratio-valued objective was
at fault. Holding the 3D geometry fixed and varying only the run length, with
both runs of the finite difference forced to cover the same interval:
length 40 adjoint 8.84e-6 finite difference 4.923e-6 ratio 1.796
length 80 adjoint 7.92e-6 finite difference 4.923e-6 ratio 1.608
length 140 adjoint 4.93e-6 finite difference 4.923e-6 ratio 1.001
The finite difference does not move; it is the adjoint that converges. The
binding constraint is the adjoint DFT, which in a lossless background rings for
considerably longer than the forward one, and an under-converged adjoint is
wrong by a factor that does not shrink with the step -- which is why it read as
a wrong gradient rather than as noise.
The tests therefore fix the run length outright, so both runs of a finite
difference cover the same interval. Left adaptive, stop_when_dft_decayed stops
the perturbed run at a different time, and that difference scales with the
perturbation, producing the same misleading signature.
smartalecH
changed the title
…trix CI killed test_angular_spectrum outright -- the OOM killer, not an assertion. propagate() summed the spectrum by building an explicit (num samples, num wavevectors) phase matrix. That is tolerable in 2D and untenable in 3D: a 64x64 monitor padded eightfold carries 262144 wavevectors, so evaluating it back on its own 4096 samples wants 17 GB. It ran here only because this machine has the memory to hide it; a 7 GB runner does not. Reconstructing on the monitor's own grid is exactly the inverse of the forward transform, so it is now an inverse FFT -- no dense matrix, and O(N log N). Arbitrary coordinates still need the direct sum, but it is chunked to a bounded number of entries per block, so asking for an inconvenient set of points costs time rather than the process. The 3D tests built the same matrix in their own helper and now use the same inverse transform. Their grids are also smaller, which is enough for what they check and takes the suite from 103 s to 77 s. test_angular_spectrum is also dropped from TESTS and left in ADJOINT_TESTS alone, matching test_adjoint_solver. It contains FDTD adjoint checks and has no business running in all four build jobs.
The non-uniform transform rebuilt its (num wavevectors, num samples) matrix of complex exponentials on every call, though it depends only on the wavevectors and the sample positions, both fixed at construction. It dominated. On a 120x120 monitor, selecting the 1517 wavevectors inside a 0.2 numerical aperture took 902 ms per overlap, against 66 ms for a padded FFT over 230400 wavevectors -- the supposedly cheap path was thirteen times slower than the one it was meant to undercut. Caching brings it to 228 ms. What remains is the O(num wavevectors x num samples) contraction itself, which is inherent. So an explicit wavevector set is worth reaching for when it is much smaller than the sample count, or when the padded grid would be unaffordable -- notably long propagation in 3D, where padding costs the square of the pad factor -- and not otherwise. The cache costs one complex matrix of that size, which is another reason it suits a small set.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.