[3/n][Adjoint Module] Adjoint gradients with respect to sources - #3283
Draft
smartalecH wants to merge 5 commits into
Draft
[3/n][Adjoint Module] Adjoint gradients with respect to sources#3283smartalecH wants to merge 5 commits into
smartalecH wants to merge 5 commits into
Conversation
Writing Maxwell's equations as A(rho) E = -i omega J gives dJ_obj/dJ = -i omega lambda, so the derivative with respect to a source is just the adjoint field sampled over that source's support. It needs no simulation beyond the adjoint run already being performed -- only a DFT monitor over the source region. The gathering is dft_fields::fourier_sourcegradient, written as the exact transpose of the fourier_sourcedata scatter that places adjoint sources: same loop, same points, same weights. That is what makes the array ordering automatic, since a cotangent comes back indexed exactly like the array handed in. Two conventions keep it simpler than expected -- chi1inv is a realnum, so the weights are real and no conjugation is involved, and update_dft already averages the four Yee sites that the scatter splits across, which is precisely the transpose of that split. Sources opt in with `differentiable=['currents', 'amplitude']`. The names are validated against a per-class whitelist and become the keys of the resulting gradient, so the two cannot drift apart. 'center' and 'size' are rejected with their own message: they move the grid points a source occupies rather than the amplitudes applied to them, which is outside this formulation rather than merely unimplemented. ArraySource supplies per-point amplitudes directly, indexed like get_dft_array over the same region -- the same ordering the gradient comes back in, because injection and measurement are a scatter and its transpose. Verified against central finite differences: 2.7e-7 for scalar amplitudes across frequencies, resolutions, complex amplitudes and both electric and magnetic sources, and 8.3e-8 for per-point currents. The normalization was pinned empirically rather than derived; it is adj_src_phase * dtft_forward / (dV * i omega), and _adj_src_phase is now exposed on ObjectiveQuantity so that placing an adjoint source and differentiating a source cannot drift.
MeepJaxWrapper takes a `sources` argument alongside `designs`, and its custom_vjp returns a cotangent for it. An ArraySource handed to the wrapper is differentiated automatically: the parameters that produced the currents live upstream in JAX, so there is nothing for Meep to name. This is the objective-side protocol run backwards -- there Meep returns monitor values and JAX owns the post-processing; here JAX supplies currents and Meep returns their cotangent. The cotangent follows the convention JAX and autograd both use for a real function of a complex input, dJ/d(Re a) - i dJ/d(Im a), so it chains with no adjustment. A convention error there would not raise, so the round trip through jax.value_and_grad is tested against a finite difference: 8.3e-8. test_source_gradient.py covers flag validation, the transpose (both the exact -dft identity for an aligned monitor and the dot-product identity against the scatter), ArraySource against an ordinary Source, and finite-difference checks of both amplitude and per-point currents gradients. Every finite-difference test pins the run length, since an adaptive stop makes a perturbed run end at a different time and turns a correct gradient into a constant-ratio failure that does not shrink with the step.
smartalecH
changed the title
fourier_sourcedata negates electric components and only those, since it was written to place adjoint sources. ArraySource was undoing that negation for every component, which flips the relative sign of the two sheets of an equivalent-current pair -- and that silently reverses the direction the pair radiates in rather than producing anything that looks like an error. The existing round-trip test only covered mp.Ez, so it could not see this.
Contracts the per-point cotangent onto beam_x0, beam_kdir, beam_w0 and
beam_E0 by finite-differencing Meep's own beam construction -- no FDTD runs,
only re-evaluating mp.gaussianbeam, and each directional derivative is
contracted as it is formed so no dense Jacobian is built.
This needed the transpose of the add_volume_source / src_vol_chunkloop path,
which is how beams (and amp_data, and amp_func) are actually placed --
distinct from fourier_sourcedata, which is how adjoint sources are placed.
Positions come out of the same C++ loop as the cotangent, following
material_grids_addgradient: rebuild the chunk's grid volume with
gv.subvolume(is, ie, c) and loop it, so dft[nf*idx+f] and the point it
belongs to cannot disagree.
Four things had to be right, each of which produced plausible wrong numbers:
- is_old/ie_old are only assigned when a chunk is created with persist, so
a monitor that does not set it must loop over is/ie instead;
- src_vol_chunkloop multiplies by gv.a once per *zero-width* direction, not
per dimension, so 1/dV is right for a point source and wrong by a factor
of the resolution for a line;
- IVEC_LOOP_WEIGHT has to be applied. Omitting it is invisible for a
grid-aligned electric source and gives exactly (1/2)^(zero-width
directions) for a component whose Yee points straddle the plane;
- the output must not be reduced. Collapsing a zero-extent direction merges
the two Yee planes a magnetic component straddles: the total weight
survives but the variation between them does not, which is exact for a
uniform sheet and wrong for a beam.
Verified against finite differences: 1.5e-8 for beam_w0, 2.8e-8 for
beam_x0.x, 2.7e-5 for beam_x0.y, 2.3e-8 for beam_kdir. beam_E0.z is checked
against an exact answer instead -- the objective is quadratic in it, so
dJ/dE0.z is exactly 2J -- and agrees to eight significant figures. The
single-component checks are exact to 1.0000 for electric and magnetic point
and line sources.
Meep already has two of the three pieces: Source(amp_data=...) takes an array profile, and get_equiv_sources applies the equivalence principle to build one. The only thing missing was differentiability, so add that rather than a new source class -- ArraySource is removed. amp_file_func trilinearly interpolates the user's array at each grid point, so the gradient is the transpose of that interpolation scattered back onto the array. That is a plain scatter in Python, because the hard part -- the cotangent with respect to the amplitude Meep applied at each grid point, and where that point is -- is already done in C++. MeepJaxWrapper picks up any source carrying amp_data automatically and returns the cotangent in whatever shape the caller passed, since Meep wants a 3D array but JAX may hand in any. Verified against finite differences on individual array entries: 1.9e-8. An amp_func still cannot be differentiated, and now says why: it is evaluated inside Meep, so there is no array for a cotangent to land on.
Collaborator
|
This is definitely nice to have. (Currently @lxvm has been avoiding this issue by using reciprocal simulations, i.e. effectively swapping the sources with output overlap integrals.) |
Contributor
|
@stevengj I still think the main reason for doing reciprocal simulations is that I am simulating a structure with fewer output channels than input channels. But yes, having adjoints with respect to sources would enable more end-to-end applications where the source terms depend on a design region outside of the meep simulation. |
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
3 participants
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.