[1/2] Fix adjoint gradient chunk-dependence without gathering the DFT monitors

[smartalecH]

(written by me)

This is a lighter alternative to #3264, for #2578 (but I use the same test, thanks!)

@jin-castle picked up on a few bugs we had in the existing adjoint-/forward-fields recombination routine. In particular:

• We were assuming chunks in the linked lists were always aligned. This is fine for isotropic materials (e.g. no smoothing) but breaks down when there are cross terms (due to anisotropic material tensors).
• Some of the indexing wasn't properly taking into account the extra halo/padded region we use. Particularly important again for the cross terms.
• In 3D, the indexing had some bugs in that third dimension where we were wrapping around (aliasing) into the wrong part toward the end of the chunk.

Note that this doesn't require serializing the DFT fields (e.g. broadcasting to all the processes) and the tests show notable gradient accuracy. In particular, no dependence on the number of chunks.

(cc @stevengj , @oskooi )

[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.

[smartalecH]

@stevengj I think we're ready to review (@jin-castle feel free to take a look too)

[jin-castle]

The current head (897dbe06) matches the revision I validated. The iyee_shift(c) clamp, spatial chunk matching, per-dimension bounds checks, and unchanged regression test all look good.

One point from my earlier comment may still be worth resolving: the remaining silent-zero paths. I agree that unconditional forward/adjoint chunk-count equality is not a valid invariant. The minimal 2d case with forward [Dx,Dy,Dz] chunk counts [0,0,1] showed that an entirely empty forward-component list can legitimately represent a zero contribution.

For a nonempty forward-component list, however, I would expect an adjoint chunk using that component either to find the corresponding (fc, sn, shift) key or to have a documented reason why the missing spatial key represents a legitimate zero. Currently, matching_dft_chunk() returning null is handled by an unconditional continue, so an unexpected spatial mismatch is indistinguishable from an intentional zero. Could a missing key in the nonempty-list case remain a hard diagnostic, unless there is a legitimate partial-list configuration that requires the current behavior?

Similarly, forward_dft_value() converts every lookup outside the persistent DFT box to zero. With the corrected component-aware padding, this should be legitimate at the global component-grid boundary, but a lookup that is still inside the global component grid and outside the persistent local DFT box would indicate an unexpected interior miss. Could that case also preserve a hard diagnostic?

These checks would keep the current fully distributed data path: they would require no gathering or global materialization of the DFT arrays, no additional communication, and no material increase in DFT storage. Apart from this invariant/diagnostic question, the patch matches my validation.

[stevengj]

This looks reasonable…

[stevengj]

Suggested change

	/* Clamp to this component's own grid, not to the centered grid. Component c
	runs from little_corner()+iyee_shift(c) to big_corner()+iyee_shift(c) (cf.
	LOOP_OVER_VOL), so clamping to the bare corners truncates the topmost node
	of a yee-shifted direction -- precisely the ghost node that the adjoint
	restriction stencil reaches for, and which step_boundaries() already keeps
	current. It also left `is` off the component lattice, which desynchronized
	the LOOP_OVER_IVECS counter from grid_volume::index(). */
	/* Clamp to this component's persist-padded grid, which is one cell bigger than the
	usual "owned" grid given by big_owned_corner(). */