Draft: fix adjoint gradient chunk-dependence without gathering the DFT monitors

[smartalecH]

Draft alternative to #3264, for #2578. Same test, no communication.

#3264 fixes a real bug, but it does so by gathering all six design-region DFT
components onto every rank (O(volume) memory and allreduce traffic per rank),
which is what @stevengj flagged. This PR attempts the same fix with local
changes only, to test whether the gather is necessary.

Claim: the cross-chunk data is already local

#3264 lists four root causes. Three are indexing bugs. The fourth — "the ±1-pixel
padding of persist dft chunks is clamped to the owning fields chunk and cannot
cross process boundaries" — I believe is not a defect, and everything else follows
from that.

The stencil reaches half a pixel, not a full pixel. With adjoint_c = Ex,
forward_c = Ey, iyee_shift(Ex) = (1,0,0), iyee_shift(Ey) = (0,1,0):

fwd_p   = ip + ŷ − x̂
fwd_pa  = fwd_p + 2x̂ = ip + x̂ + ŷ
fwd_pf  = fwd_p − 2ŷ = ip − x̂ − ŷ
fwd_paf =              ip + x̂ − ŷ

unit_a * 2 cancels against the yee shift. The four points are ip ± x̂ ± ŷ, i.e.
the four Ey nodes at the corners of the Yee cell around the Ex node — 1 ivec unit
away, half a pixel. The persist pad asks for 2 units, twice the reach.

Those nodes are always inside fc->gv. With io = little_corner(), chunk size
n_d pixels, and owns(p) = 0 < (p−io)_d ≤ 2n_d (src/vec.cpp:445):

direction	component	node offsets in chunk	owned	needed
a	adjoint (shifted)	1,3,…,2nₐ+1	1…2nₐ−1	—
a	forward (unshifted)	0,2,…,2nₐ	2…2nₐ	ip±1 even ∈ [0, 2nₐ] ✅
f	adjoint (unshifted)	0,2,…,2n_f	2…2n_f	—
f	forward (shifted)	1,3,…,2n_f+1	1…2n_f−1	ip±1 odd ∈ [1, 2n_f+1] ✅

The required range is exactly the chunk's forward-component node range. At a
chunk face the extra node is the chunk's ghost layer, which connect_chunks()
maintains (src/boundaries.cpp:412, :480) and which is current at DFT time:
update_dfts() is src/step.cpp:126, after step_boundaries(E_stuff) at :121.
The clamp to fc->gv therefore never removes a node the stencil needs.

Split in y at plane P, chunk A below:

A owns Ex in y: [io+2 … P]     A owns Ey in y: [io+1 … P−1]
B owns Ex in y: [P+2 … ]       B owns Ey in y: [P+1 … ]
A's Ey pad: ie = min(P−1+2, big_corner=P) = P → node P+1, A's high ghost ✅
A's topmost owned Ex node P needs Ey at P−1 (owned) and P+1 (ghost). Present.
B's Ey pad: is = max(P+1−2, little_corner=P) = P → node P+1
B's lowest owned Ex node P+2 needs Ey at P+1 and P+3. Present.

Changes

1. Pair chunks spatially (matching_dft_chunk). forward_dft_chunks[ci_forward][cur_chunk]
   indexed the forward list with the adjoint component's list position. A fields
   chunk contributes a dft_chunk for a component only where the monitor overlaps
   that component's owned grid, so the per-component lists differ in length and
   ordering near the monitor edges. Match on (fc, sn, shift) — exactly the tuple
   loop_in_chunks() iterates over. This also subsumes root cause 4 (idx_adj
   used on the forward array), which was only valid under the chunk-alignment
   assumption.

2. Per-dimension bounds checks (ivec_in_box). grid_volume::index() is an
   inner product of per-direction offsets with strides, so in 3d a point one pixel
   off the end in y still lands in [0,N), on an unrelated voxel. Guarding the flat
   index alone silently substitutes a neighboring field value. 1d/2d mostly
   degenerate to the right answer, which is consistent with this surviving — every
   adjoint test upstream is 2d or cylindrical.

3. Drop the abort on unequal adjoint/forward chunk counts. That is a
   legitimate configuration, not an error.

4. Two diagnostic counters, reported at verbosity > 2, that test the claim
   above directly: forward lookups landing inside the cell but outside the local
   chunk, and adjoint chunks with no matching forward chunk. Both should be zero
   at every chunk division. Happy to drop these before merge, or keep them as a
   guard.

Test

python/tests/test_adjoint_chunks.py is taken from #3264 unmodified, so the two
PRs are directly comparable. It forces the bug in serial via num_chunks, so it
runs in every CI job rather than only the MPI ones — a good framing, since this
was never an MPI bug.

What would falsify this

If the counters in (4) come up nonzero, or the test fails, then cross-chunk data
really is needed and #3264's gather (or a surface-scale halo exchange of the DFT
boundary planes) is required. That is the question this draft is meant to settle.

Not addressed

• Symmetry. matching_dft_chunk keys on (fc, sn, shift), which is correct for
  symmetry images, but I have not verified the rest of the path under symmetry.
• The kernel-normalization and eps_averaging fixes in Fix 3D MaterialGrid subpixel smoothing: kernel normalization (#3258) and gradient consistency (#3259) #3263 are independent.

cc @jin-castle @stevengj @lxvm — @jin-castle, this is meant as input to #3264,
not a replacement; if the counters fire, your analysis stands and mine is wrong.

🤖 Generated with Claude Code

[smartalecH]

I'm not sure this is true. Even with different adjoint sources, the forward and adjoint fields in the design region should be the same, and thus the chunks should be the same. Can't we check that in the code that allocates the adjoint run from the forward run?

[jin-castle]

Thanks for putting this together. I ran the two diagnostic counters, and the
result supports the local/distributed approach, but it falsifies one premise in
the current patch: the persist box is intended to include the required ghost
node, but for Yee-shifted components its high-side clamp truncates that node.

Using this PR's head (39eb51dc) with only the counter print gates forced on,
the serial case from test_adjoint_chunks.py gives:

num_chunks	n_interior_miss	n_missing_fwd_chunk	relative gradient change
1	0	0	—
3	120	0	4.425561e-01

The objective changes by only 6.7e-16. The gradient difference reproduces
the CI failure (0.442556054604398), so the failure is entirely in gradient
assembly.

The split between the counters is useful:

• n_missing_fwd_chunk = 0: matching on (fields_chunk, sn, shift) works for
  this problem. I agree that pairing the forward list with the adjoint list
  index was wrong, and matching_dft_chunk fixes that without gathering.
• n_interior_miss = 120: the required point is inside the cell but outside
  the allocated local DFT box when a chunk boundary crosses the design region.

The stencil derivation in the PR is correct: it reaches one ivec unit, and
the necessary field value is in the fields chunk's ghost layer. The issue is
the clamp used while constructing a persistent DFT chunk:

is = max(is - one_ivec(fc->gv.dim) * 2, fc->gv.little_corner());
ie = min(ie + one_ivec(fc->gv.dim) * 2, fc->gv.big_corner());

little_corner() and big_corner() do not include the component's Yee shift.
For a component shifted in direction d, the component-node range extends to
big_corner() + iyee_shift(c). Capping ie at the unshifted corner leaves it
on the wrong parity; the subsequent (ie-is)/2 sizing therefore omits the
high-side ghost node. ivec_in_box then correctly rejects the lookup instead
of silently aliasing another voxel.

Adding the component shift to the clamp fixes that off-by-one node:

is = max(is - one_ivec(fc->gv.dim) * 2,
         fc->gv.little_corner() + fc->gv.iyee_shift(c));
ie = min(ie + one_ivec(fc->gv.dim) * 2,
         fc->gv.big_corner() + fc->gv.iyee_shift(c));

With those two changed bounds on top of this PR, unchanged otherwise:

num_chunks	n_interior_miss	n_missing_fwd_chunk	relative gradient change
1	0	0	—
3	0	0	3.674299e-16

I also checked the combined patch beyond this one split:

• test_adjoint_chunks.py passes for MPI process counts 1 through 8,
  including odd counts.
• With Meep choosing the process-dependent chunk divisions, the gradient
  agrees with the one-process result to at worst 5.1e-12 through 8 processes;
  the objective is unchanged.
• A separate single-precision build passes the test at 1, 2, 3, 4, and 8
  processes.
• The full double-precision make check passes 62/62.
• Stress cases with the design region flush against the cell faces do not hit
  gv.subvolume() assertions, and a symmetry-induced change in chunk layout
  remains chunk-independent. This last check is only about the clamp and chunk
  division, not a claim about the correctness of the broader symmetry-adjoint
  path.

I also looked more closely at your inline note about the removed chunk-count
abort. I do not think unconditional same-component count equality is an
invariant in the current allocation path. add_dft_chunkloop skips a component
when fc->f[c][0] has not been allocated, while change_sources() may allocate
additional components before the adjoint monitors are installed and does not
retroactively add chunks to the retained forward monitors.

A minimal 2d example using the same Simulation, design volume, and chunk
layout gives [Dx,Dy,Dz] forward-monitor chunk counts of [0,0,1] for an Ez
forward source, then [1,1,1] after changing to an Ex adjoint source. The
missing forward polarization is identically zero in that decoupled problem, so
an unconditional count abort would reject a valid zero contribution.

I think the useful invariant is spatial rather than numerical: if a forward
component has any chunks, every adjoint chunk that uses it should find the
corresponding (fc, sn, shift) key; failure to find that key should abort (or at
least remain a hard diagnostic). If the entire forward-component list is empty,
zero is legitimate. This is also stronger than comparing counts, since equal
counts do not prove that the spatial keys match. Forcing all three design-field
components to be allocated before the forward run would make the counts equal,
but would allocate fields that are otherwise unnecessary.

My proposed minimal shape is therefore:

1. the iyee_shift(c) clamp in src/dft.cpp;
2. this PR's spatial chunk matching and per-dimension bounds check;
3. a spatial-key invariant for nonempty forward-component lists rather than an
   unconditional chunk-count equality check;
4. test_adjoint_chunks.py unchanged.

That keeps the DFT data distributed and adds no communication. It also means
the global gather in #3264 is unnecessary. I am happy for this fix to land via
#3274 with the clamp folded in; I can reduce #3264 to the same shape rather than
defend the gather.

[stevengj]

Thanks for working on this, everyone.