witwin-maxwell 0.3.0__py3-none-win_amd64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (107) hide show
  1. witwin/maxwell/__init__.py +164 -0
  2. witwin/maxwell/adapters/__init__.py +1 -0
  3. witwin/maxwell/adapters/gds.py +203 -0
  4. witwin/maxwell/adapters/tidy3d.py +1669 -0
  5. witwin/maxwell/adjoint_inputs.py +70 -0
  6. witwin/maxwell/compiler/__init__.py +11 -0
  7. witwin/maxwell/compiler/materials.py +1355 -0
  8. witwin/maxwell/compiler/monitors.py +164 -0
  9. witwin/maxwell/compiler/sources.py +222 -0
  10. witwin/maxwell/fdfd/__init__.py +3 -0
  11. witwin/maxwell/fdfd/adjoint/__init__.py +3 -0
  12. witwin/maxwell/fdfd/adjoint/bridge.py +180 -0
  13. witwin/maxwell/fdfd/ams.py +300 -0
  14. witwin/maxwell/fdfd/krylov.py +280 -0
  15. witwin/maxwell/fdfd/plotting.py +110 -0
  16. witwin/maxwell/fdfd/postprocess.py +33 -0
  17. witwin/maxwell/fdfd/solver.py +1057 -0
  18. witwin/maxwell/fdtd/__init__.py +4 -0
  19. witwin/maxwell/fdtd/adjoint/__init__.py +21 -0
  20. witwin/maxwell/fdtd/adjoint/bridge.py +746 -0
  21. witwin/maxwell/fdtd/adjoint/capabilities.py +29 -0
  22. witwin/maxwell/fdtd/adjoint/core.py +3614 -0
  23. witwin/maxwell/fdtd/adjoint/dispatch.py +684 -0
  24. witwin/maxwell/fdtd/adjoint/native.py +2561 -0
  25. witwin/maxwell/fdtd/adjoint/profiler.py +195 -0
  26. witwin/maxwell/fdtd/adjoint/reverse_common.py +207 -0
  27. witwin/maxwell/fdtd/adjoint/seeds.py +524 -0
  28. witwin/maxwell/fdtd/boundary/__init__.py +37 -0
  29. witwin/maxwell/fdtd/boundary/bloch.py +90 -0
  30. witwin/maxwell/fdtd/boundary/common.py +41 -0
  31. witwin/maxwell/fdtd/boundary/cpml.py +527 -0
  32. witwin/maxwell/fdtd/boundary/runtime.py +209 -0
  33. witwin/maxwell/fdtd/checkpoint.py +227 -0
  34. witwin/maxwell/fdtd/coords.py +32 -0
  35. witwin/maxwell/fdtd/cuda/__init__.py +11 -0
  36. witwin/maxwell/fdtd/cuda/backend.py +3532 -0
  37. witwin/maxwell/fdtd/cuda/build.py +292 -0
  38. witwin/maxwell/fdtd/cuda/extension.cpp +2053 -0
  39. witwin/maxwell/fdtd/cuda/kernels/adjoint.cu +5290 -0
  40. witwin/maxwell/fdtd/cuda/kernels/aniso.cu +1704 -0
  41. witwin/maxwell/fdtd/cuda/kernels/boundary.cu +424 -0
  42. witwin/maxwell/fdtd/cuda/kernels/common.cu +74 -0
  43. witwin/maxwell/fdtd/cuda/kernels/common.cuh +82 -0
  44. witwin/maxwell/fdtd/cuda/kernels/dispersive.cu +782 -0
  45. witwin/maxwell/fdtd/cuda/kernels/electric.cu +3613 -0
  46. witwin/maxwell/fdtd/cuda/kernels/magnetic.cu +781 -0
  47. witwin/maxwell/fdtd/cuda/kernels/observers.cu +347 -0
  48. witwin/maxwell/fdtd/cuda/kernels/projection.cu +211 -0
  49. witwin/maxwell/fdtd/cuda/kernels/sources.cu +1482 -0
  50. witwin/maxwell/fdtd/cuda/kernels/spectral.cu +297 -0
  51. witwin/maxwell/fdtd/cuda/launch.h +23 -0
  52. witwin/maxwell/fdtd/cuda/prebuilt/witwin_maxwell_fdtd_cuda.pyd +0 -0
  53. witwin/maxwell/fdtd/cuda/runtime/__init__.py +6 -0
  54. witwin/maxwell/fdtd/cuda/runtime/dispatcher.py +6 -0
  55. witwin/maxwell/fdtd/cuda/runtime/graph.py +33 -0
  56. witwin/maxwell/fdtd/cuda/tensors.h +35 -0
  57. witwin/maxwell/fdtd/dispersion.py +48 -0
  58. witwin/maxwell/fdtd/evaluation.py +290 -0
  59. witwin/maxwell/fdtd/excitation/__init__.py +27 -0
  60. witwin/maxwell/fdtd/excitation/injection.py +1102 -0
  61. witwin/maxwell/fdtd/excitation/modes.py +1675 -0
  62. witwin/maxwell/fdtd/excitation/spatial.py +469 -0
  63. witwin/maxwell/fdtd/excitation/temporal.py +510 -0
  64. witwin/maxwell/fdtd/excitation/tfsf_apply.py +264 -0
  65. witwin/maxwell/fdtd/excitation/tfsf_common.py +421 -0
  66. witwin/maxwell/fdtd/excitation/tfsf_specs.py +490 -0
  67. witwin/maxwell/fdtd/excitation/tfsf_state.py +778 -0
  68. witwin/maxwell/fdtd/io.py +31 -0
  69. witwin/maxwell/fdtd/material_pullback.py +371 -0
  70. witwin/maxwell/fdtd/meshing/__init__.py +3 -0
  71. witwin/maxwell/fdtd/meshing/autogrid.py +335 -0
  72. witwin/maxwell/fdtd/observers.py +1453 -0
  73. witwin/maxwell/fdtd/plotting.py +121 -0
  74. witwin/maxwell/fdtd/postprocess.py +205 -0
  75. witwin/maxwell/fdtd/runtime/__init__.py +119 -0
  76. witwin/maxwell/fdtd/runtime/initialization.py +236 -0
  77. witwin/maxwell/fdtd/runtime/materials.py +1563 -0
  78. witwin/maxwell/fdtd/runtime/module_cache.py +31 -0
  79. witwin/maxwell/fdtd/runtime/spectral.py +581 -0
  80. witwin/maxwell/fdtd/runtime/stepping.py +1987 -0
  81. witwin/maxwell/fdtd/solver.py +502 -0
  82. witwin/maxwell/geometry/__init__.py +5 -0
  83. witwin/maxwell/geometry/polyslab.py +299 -0
  84. witwin/maxwell/materials.py +3 -0
  85. witwin/maxwell/media.py +1467 -0
  86. witwin/maxwell/monitors.py +743 -0
  87. witwin/maxwell/ports.py +113 -0
  88. witwin/maxwell/postprocess/RCS.py +139 -0
  89. witwin/maxwell/postprocess/__init__.py +42 -0
  90. witwin/maxwell/postprocess/diffraction.py +219 -0
  91. witwin/maxwell/postprocess/directivity.py +152 -0
  92. witwin/maxwell/postprocess/emission.py +67 -0
  93. witwin/maxwell/postprocess/modal.py +324 -0
  94. witwin/maxwell/postprocess/nfft.py +266 -0
  95. witwin/maxwell/postprocess/scattering_parameters.py +358 -0
  96. witwin/maxwell/postprocess/stratton_chu.py +1193 -0
  97. witwin/maxwell/result.py +967 -0
  98. witwin/maxwell/scene.py +1338 -0
  99. witwin/maxwell/simulation.py +659 -0
  100. witwin/maxwell/sources.py +917 -0
  101. witwin/maxwell/visualization/__init__.py +28 -0
  102. witwin/maxwell/visualization/interactive.py +484 -0
  103. witwin/maxwell/visualization/plots.py +568 -0
  104. witwin_maxwell-0.3.0.dist-info/METADATA +184 -0
  105. witwin_maxwell-0.3.0.dist-info/RECORD +107 -0
  106. witwin_maxwell-0.3.0.dist-info/WHEEL +4 -0
  107. witwin_maxwell-0.3.0.dist-info/licenses/LICENSE +13 -0
@@ -0,0 +1,2561 @@
1
+ """Native CUDA reverse-step runners for the FDTD adjoint.
2
+
3
+ These runners orchestrate the fused reverse-math CUDA kernels registered in
4
+ ``cuda/backend.py`` into a full per-step reverse pass. Python only sequences the
5
+ kernel launches; every
6
+ per-cell reverse computation (electric-adjoint -> mid-H adjoint, magnetic-adjoint
7
+ -> pre-step E adjoint + eps gradient, and the magnetic decay pullback) runs
8
+ inside the compiled kernels.
9
+
10
+ The mid-step magnetic field is still reconstructed with the shared Torch replay
11
+ helper (``_forward_magnetic_fields`` / ``_forward_magnetic_fields_complex`` for the
12
+ complex split-field Bloch runner); nativizing that replay is tracked separately
13
+ under the native-replay work. The reverse *math* itself is fully native here.
14
+ """
15
+
16
+ from __future__ import annotations
17
+
18
+ from ..cuda import backend as _cuda_backend
19
+ from .dispatch import (
20
+ _NATIVE_REVERSE_LABELS,
21
+ _ReverseBackend,
22
+ register_native_reverse_backend,
23
+ )
24
+
25
+
26
+ def _scene_is_cuda(forward_state) -> bool:
27
+ reference = forward_state.get("Ex")
28
+ return bool(reference is not None and reference.is_cuda)
29
+
30
+
31
+ def _cuda_scene_native_qualifies(solver, forward_state) -> bool:
32
+ """Native reverse runners only run on a CUDA scene with the compiled extension."""
33
+ if not _scene_is_cuda(forward_state) or not _cuda_backend.is_available():
34
+ return False
35
+ try:
36
+ extension = _cuda_backend.get_compiled_extension()
37
+ except (ImportError, OSError, RuntimeError):
38
+ return False
39
+ solver._fdtd_cuda_extension = extension
40
+ return True
41
+
42
+
43
+ def _dispersive_native_qualifies(solver, forward_state) -> bool:
44
+ """Gate the native ADE runner to CUDA scenes with the compiled extension."""
45
+ return _cuda_scene_native_qualifies(solver, forward_state)
46
+
47
+
48
+ def _reverse_step_standard_native_core(
49
+ solver,
50
+ forward_state,
51
+ adjoint_state,
52
+ *,
53
+ time_value,
54
+ eps_ex,
55
+ eps_ey,
56
+ eps_ez,
57
+ resolved_source_terms,
58
+ magnetic_fields=None,
59
+ ):
60
+ """Fused native standard reverse math, without the source-term eps gradient.
61
+
62
+ Returns the ``_ReverseStepResult`` the analytic standard reference produces
63
+ *before* ``_accumulate_source_term_gradients`` runs. The public runner adds
64
+ that accumulation; the dispersive runner reuses this core as its base reverse
65
+ and defers the single source-term accumulation to the end of the full step.
66
+ """
67
+ import torch
68
+
69
+ from . import core as _adjoint
70
+ from .reverse_common import dynamic_electric_curls
71
+
72
+ # Mid-step H the forward electric update consumed (shared Torch replay).
73
+ if magnetic_fields is None:
74
+ magnetic_fields = _adjoint._forward_magnetic_fields(
75
+ solver,
76
+ forward_state,
77
+ time_value=time_value,
78
+ resolved_source_terms=resolved_source_terms,
79
+ )
80
+ hx_mid = magnetic_fields["Hx"]
81
+ hy_mid = magnetic_fields["Hy"]
82
+ hz_mid = magnetic_fields["Hz"]
83
+
84
+ # Dynamic electric curl coefficients (cast base curl by the eps leaf).
85
+ ex_curl, ey_curl, ez_curl = dynamic_electric_curls(
86
+ solver, eps_ex=eps_ex, eps_ey=eps_ey, eps_ez=eps_ez
87
+ )
88
+
89
+ # Phase 1: electric adjoint -> mid-H adjoint (the magnetic_output_adjoint).
90
+ adj_hx_mid = torch.empty_like(forward_state["Hx"])
91
+ adj_hy_mid = torch.empty_like(forward_state["Hy"])
92
+ adj_hz_mid = torch.empty_like(forward_state["Hz"])
93
+ _cuda_backend._reverse_electric_hx_standard(
94
+ AdjHxMid=adj_hx_mid,
95
+ AdjHxPost=adjoint_state["Hx"],
96
+ AdjEyPost=adjoint_state["Ey"],
97
+ AdjEzPost=adjoint_state["Ez"],
98
+ EyCurl=ey_curl,
99
+ EzCurl=ez_curl,
100
+ invDy=solver.inv_dy_e,
101
+ invDz=solver.inv_dz_e,
102
+ )
103
+ _cuda_backend._reverse_electric_hy_standard(
104
+ AdjHyMid=adj_hy_mid,
105
+ AdjHyPost=adjoint_state["Hy"],
106
+ AdjExPost=adjoint_state["Ex"],
107
+ AdjEzPost=adjoint_state["Ez"],
108
+ ExCurl=ex_curl,
109
+ EzCurl=ez_curl,
110
+ invDx=solver.inv_dx_e,
111
+ invDz=solver.inv_dz_e,
112
+ )
113
+ _cuda_backend._reverse_electric_hz_standard(
114
+ AdjHzMid=adj_hz_mid,
115
+ AdjHzPost=adjoint_state["Hz"],
116
+ AdjExPost=adjoint_state["Ex"],
117
+ AdjEyPost=adjoint_state["Ey"],
118
+ ExCurl=ex_curl,
119
+ EyCurl=ey_curl,
120
+ invDx=solver.inv_dx_e,
121
+ invDy=solver.inv_dy_e,
122
+ )
123
+
124
+ # Phase 2: magnetic adjoint -> pre-step E adjoint + eps gradient. Each kernel
125
+ # reconstructs its own curl(H) from the mid-H fields and assigns the complete
126
+ # pre-step E adjoint (electric decay pullback + magnetic forward-diff).
127
+ adj_ex_prev = torch.empty_like(forward_state["Ex"])
128
+ adj_ey_prev = torch.empty_like(forward_state["Ey"])
129
+ adj_ez_prev = torch.empty_like(forward_state["Ez"])
130
+ grad_eps_ex = torch.empty_like(eps_ex)
131
+ grad_eps_ey = torch.empty_like(eps_ey)
132
+ grad_eps_ez = torch.empty_like(eps_ez)
133
+ _cuda_backend._reverse_magnetic_ex_standard(
134
+ AdjExPrev=adj_ex_prev,
135
+ GradEpsEx=grad_eps_ex,
136
+ AdjExPost=adjoint_state["Ex"],
137
+ AdjHyMid=adj_hy_mid,
138
+ AdjHzMid=adj_hz_mid,
139
+ ExDecay=solver.cex_decay,
140
+ ExCurl=ex_curl,
141
+ EpsEx=eps_ex,
142
+ HyMid=hy_mid,
143
+ HzMid=hz_mid,
144
+ HyCurl=solver.chy_curl,
145
+ HzCurl=solver.chz_curl,
146
+ invDyE=solver.inv_dy_e,
147
+ invDzE=solver.inv_dz_e,
148
+ invDyH=solver.inv_dy_h,
149
+ invDzH=solver.inv_dz_h,
150
+ yLowBoundaryMode=solver.boundary_y_low_code,
151
+ yHighBoundaryMode=solver.boundary_y_high_code,
152
+ zLowBoundaryMode=solver.boundary_z_low_code,
153
+ zHighBoundaryMode=solver.boundary_z_high_code,
154
+ )
155
+ _cuda_backend._reverse_magnetic_ey_standard(
156
+ AdjEyPrev=adj_ey_prev,
157
+ GradEpsEy=grad_eps_ey,
158
+ AdjEyPost=adjoint_state["Ey"],
159
+ AdjHxMid=adj_hx_mid,
160
+ AdjHzMid=adj_hz_mid,
161
+ EyDecay=solver.cey_decay,
162
+ EyCurl=ey_curl,
163
+ EpsEy=eps_ey,
164
+ HxMid=hx_mid,
165
+ HzMid=hz_mid,
166
+ HxCurl=solver.chx_curl,
167
+ HzCurl=solver.chz_curl,
168
+ invDxE=solver.inv_dx_e,
169
+ invDzE=solver.inv_dz_e,
170
+ invDxH=solver.inv_dx_h,
171
+ invDzH=solver.inv_dz_h,
172
+ xLowBoundaryMode=solver.boundary_x_low_code,
173
+ xHighBoundaryMode=solver.boundary_x_high_code,
174
+ zLowBoundaryMode=solver.boundary_z_low_code,
175
+ zHighBoundaryMode=solver.boundary_z_high_code,
176
+ )
177
+ _cuda_backend._reverse_magnetic_ez_standard(
178
+ AdjEzPrev=adj_ez_prev,
179
+ GradEpsEz=grad_eps_ez,
180
+ AdjEzPost=adjoint_state["Ez"],
181
+ AdjHxMid=adj_hx_mid,
182
+ AdjHyMid=adj_hy_mid,
183
+ EzDecay=solver.cez_decay,
184
+ EzCurl=ez_curl,
185
+ EpsEz=eps_ez,
186
+ HxMid=hx_mid,
187
+ HyMid=hy_mid,
188
+ HxCurl=solver.chx_curl,
189
+ HyCurl=solver.chy_curl,
190
+ invDxE=solver.inv_dx_e,
191
+ invDyE=solver.inv_dy_e,
192
+ invDxH=solver.inv_dx_h,
193
+ invDyH=solver.inv_dy_h,
194
+ xLowBoundaryMode=solver.boundary_x_low_code,
195
+ xHighBoundaryMode=solver.boundary_x_high_code,
196
+ yLowBoundaryMode=solver.boundary_y_low_code,
197
+ yHighBoundaryMode=solver.boundary_y_high_code,
198
+ )
199
+
200
+ # Phase 3: magnetic decay pullback -> pre-step H adjoint.
201
+ adj_hx_prev = torch.empty_like(forward_state["Hx"])
202
+ adj_hy_prev = torch.empty_like(forward_state["Hy"])
203
+ adj_hz_prev = torch.empty_like(forward_state["Hz"])
204
+ _cuda_backend._reverse_magnetic_hx_decay(
205
+ AdjHxPrev=adj_hx_prev, AdjHxMid=adj_hx_mid, HxDecay=solver.chx_decay
206
+ )
207
+ _cuda_backend._reverse_magnetic_hy_decay(
208
+ AdjHyPrev=adj_hy_prev, AdjHyMid=adj_hy_mid, HyDecay=solver.chy_decay
209
+ )
210
+ _cuda_backend._reverse_magnetic_hz_decay(
211
+ AdjHzPrev=adj_hz_prev, AdjHzMid=adj_hz_mid, HzDecay=solver.chz_decay
212
+ )
213
+
214
+ step_result = _adjoint._ReverseStepResult(
215
+ pre_step_adjoint={
216
+ "Ex": adj_ex_prev,
217
+ "Ey": adj_ey_prev,
218
+ "Ez": adj_ez_prev,
219
+ "Hx": adj_hx_prev,
220
+ "Hy": adj_hy_prev,
221
+ "Hz": adj_hz_prev,
222
+ },
223
+ grad_eps_ex=grad_eps_ex,
224
+ grad_eps_ey=grad_eps_ey,
225
+ grad_eps_ez=grad_eps_ez,
226
+ backend=_NATIVE_REVERSE_LABELS[_ReverseBackend.STANDARD],
227
+ magnetic_output_adjoint={"Hx": adj_hx_mid, "Hy": adj_hy_mid, "Hz": adj_hz_mid},
228
+ )
229
+ return step_result
230
+
231
+
232
+ def _reverse_step_standard_native(
233
+ solver,
234
+ forward_state,
235
+ adjoint_state,
236
+ *,
237
+ time_value,
238
+ eps_ex,
239
+ eps_ey,
240
+ eps_ez,
241
+ resolved_source_terms,
242
+ profiler,
243
+ ):
244
+ """Fully native reverse step for the standard (open-boundary) configuration.
245
+
246
+ Mirrors ``reverse_step_standard_native`` exactly, launching the
247
+ fused reverse kernels in a fixed order. The fused kernels *assign* (not
248
+ accumulate) into their outputs, so the launch order is load-bearing: the
249
+ electric->H kernels must fully populate the mid-H adjoints before the
250
+ magnetic->E kernels (which read them) and before the decay pullback.
251
+ """
252
+ from . import core as _adjoint
253
+
254
+ step_result = _reverse_step_standard_native_core(
255
+ solver,
256
+ forward_state,
257
+ adjoint_state,
258
+ time_value=time_value,
259
+ eps_ex=eps_ex,
260
+ eps_ey=eps_ey,
261
+ eps_ez=eps_ez,
262
+ resolved_source_terms=resolved_source_terms,
263
+ )
264
+ # The native runner owns the full reverse-step contract, including the
265
+ # analytic source-term eps-gradient accumulation the reference path applies.
266
+ return _adjoint._accumulate_source_term_gradients(
267
+ step_result,
268
+ solver=solver,
269
+ adjoint_state=adjoint_state,
270
+ time_value=time_value,
271
+ eps_ex=eps_ex,
272
+ eps_ey=eps_ey,
273
+ eps_ez=eps_ez,
274
+ resolved_source_terms=resolved_source_terms,
275
+ )
276
+
277
+
278
+ def _reverse_step_cpml_native_core(
279
+ solver,
280
+ forward_state,
281
+ adjoint_state,
282
+ *,
283
+ time_value,
284
+ eps_ex,
285
+ eps_ey,
286
+ eps_ez,
287
+ resolved_source_terms,
288
+ magnetic_fields=None,
289
+ ):
290
+ """Fused native CPML reverse math, without the source-term eps gradient.
291
+
292
+ Returns the ``_ReverseStepResult`` the analytic CPML reference produces
293
+ *before* ``_accumulate_source_term_gradients`` runs. The public runner adds
294
+ that accumulation; the dispersive runner reuses this core as its CPML base
295
+ reverse and defers the single source-term accumulation to the end of the step.
296
+
297
+ Mirrors ``reverse_step_cpml_native`` exactly, launching the fused
298
+ CPML reverse kernels in a fixed order. Every per-cell reverse computation
299
+ (electric-adjoint -> pre-step E/psi_e adjoint + eps gradient + curl(H)
300
+ derivative adjoint, and magnetic-decay/psi_h pullback -> pre-step H/psi_h
301
+ adjoint + curl(E) derivative adjoint) runs inside the compiled kernels; Python
302
+ only sequences the launches.
303
+
304
+ Launch order is load-bearing. The fused CPML electric/magnetic kernels
305
+ *assign* into their pre-step / eps-gradient / psi outputs, and the transposed
306
+ difference accumulators *add* into the mid-step H adjoint and the pre-step E
307
+ adjoint. So all three electric kernels (and their backward-difference folds)
308
+ must fully populate the mid-step H adjoint before any magnetic kernel reads
309
+ it, and the electric kernels must assign the pre-step E adjoint before the
310
+ magnetic kernels add their forward-difference contributions into it.
311
+
312
+ The coefficient argument lists are long and per-component permuted (the
313
+ positive/negative CPML axis, its ``b``/``c``/``inv_kappa`` stretch vectors,
314
+ and the pre-step psi state all rotate with the component); the wiring here is
315
+ a line-for-line transcription of the Torch reference and is pinned against it
316
+ by the step-level parity test.
317
+ """
318
+ from . import core as _adjoint
319
+ from .reverse_common import allocate_cpml_reverse_context
320
+
321
+ # Mid-step H the forward electric update consumed (shared Torch replay).
322
+ if magnetic_fields is None:
323
+ magnetic_fields = _adjoint._forward_magnetic_fields(
324
+ solver,
325
+ forward_state,
326
+ time_value=time_value,
327
+ resolved_source_terms=resolved_source_terms,
328
+ )
329
+ hx_mid = magnetic_fields["Hx"]
330
+ hy_mid = magnetic_fields["Hy"]
331
+ hz_mid = magnetic_fields["Hz"]
332
+
333
+ ctx = allocate_cpml_reverse_context(
334
+ solver,
335
+ forward_state,
336
+ adjoint_state,
337
+ eps_ex=eps_ex,
338
+ eps_ey=eps_ey,
339
+ eps_ez=eps_ez,
340
+ )
341
+ pre = ctx.pre_step_adjoint
342
+ adj_h_mid = ctx.magnetic_output_adjoint
343
+
344
+ # Phase 1: electric adjoint -> pre-step E/psi_e adjoint + eps gradient, and
345
+ # fold each curl(H)-derivative adjoint into the mid-step H adjoint.
346
+ _cuda_backend._reverse_electric_cpml_ex(
347
+ AdjExPrev=pre["Ex"],
348
+ GradEpsEx=ctx.grad_eps_ex,
349
+ AdjPsiPosPrev=pre["psi_ex_y"],
350
+ AdjPsiNegPrev=pre["psi_ex_z"],
351
+ AdjDPos=ctx.adj_d_hz_dy,
352
+ AdjDNeg=ctx.adj_d_hy_dz,
353
+ AdjExPost=adjoint_state["Ex"],
354
+ AdjPsiPosPost=adjoint_state["psi_ex_y"],
355
+ AdjPsiNegPost=adjoint_state["psi_ex_z"],
356
+ ExDecay=solver.cex_decay,
357
+ ExCurl=ctx.ex_curl,
358
+ EpsEx=eps_ex,
359
+ PsiPos=forward_state["psi_ex_y"],
360
+ PsiNeg=forward_state["psi_ex_z"],
361
+ BPos=solver.cpml_b_e_y,
362
+ CPos=solver.cpml_c_e_y,
363
+ InvKappaPos=solver.cpml_inv_kappa_e_y,
364
+ BNeg=solver.cpml_b_e_z,
365
+ CNeg=solver.cpml_c_e_z,
366
+ InvKappaNeg=solver.cpml_inv_kappa_e_z,
367
+ HyMid=hy_mid,
368
+ HzMid=hz_mid,
369
+ invDy=solver.inv_dy_e,
370
+ invDz=solver.inv_dz_e,
371
+ yLowBoundaryMode=solver.boundary_y_low_code,
372
+ yHighBoundaryMode=solver.boundary_y_high_code,
373
+ zLowBoundaryMode=solver.boundary_z_low_code,
374
+ zHighBoundaryMode=solver.boundary_z_high_code,
375
+ )
376
+ _cuda_backend._accumulate_backward_diff_y(FieldGrad=adj_h_mid["Hz"], DiffGrad=ctx.adj_d_hz_dy, invDy=solver.inv_dy_e)
377
+ _cuda_backend._accumulate_backward_diff_z(FieldGrad=adj_h_mid["Hy"], DiffGrad=ctx.adj_d_hy_dz, invDz=solver.inv_dz_e)
378
+
379
+ _cuda_backend._reverse_electric_cpml_ey(
380
+ AdjEyPrev=pre["Ey"],
381
+ GradEpsEy=ctx.grad_eps_ey,
382
+ AdjPsiPosPrev=pre["psi_ey_z"],
383
+ AdjPsiNegPrev=pre["psi_ey_x"],
384
+ AdjDPos=ctx.adj_d_hx_dz,
385
+ AdjDNeg=ctx.adj_d_hz_dx,
386
+ AdjEyPost=adjoint_state["Ey"],
387
+ AdjPsiPosPost=adjoint_state["psi_ey_x"],
388
+ AdjPsiNegPost=adjoint_state["psi_ey_z"],
389
+ EyDecay=solver.cey_decay,
390
+ EyCurl=ctx.ey_curl,
391
+ EpsEy=eps_ey,
392
+ PsiPos=forward_state["psi_ey_z"],
393
+ PsiNeg=forward_state["psi_ey_x"],
394
+ BPos=solver.cpml_b_e_z,
395
+ CPos=solver.cpml_c_e_z,
396
+ InvKappaPos=solver.cpml_inv_kappa_e_z,
397
+ BNeg=solver.cpml_b_e_x,
398
+ CNeg=solver.cpml_c_e_x,
399
+ InvKappaNeg=solver.cpml_inv_kappa_e_x,
400
+ HxMid=hx_mid,
401
+ HzMid=hz_mid,
402
+ invDx=solver.inv_dx_e,
403
+ invDz=solver.inv_dz_e,
404
+ xLowBoundaryMode=solver.boundary_x_low_code,
405
+ xHighBoundaryMode=solver.boundary_x_high_code,
406
+ zLowBoundaryMode=solver.boundary_z_low_code,
407
+ zHighBoundaryMode=solver.boundary_z_high_code,
408
+ )
409
+ _cuda_backend._accumulate_backward_diff_z(FieldGrad=adj_h_mid["Hx"], DiffGrad=ctx.adj_d_hx_dz, invDz=solver.inv_dz_e)
410
+ _cuda_backend._accumulate_backward_diff_x(FieldGrad=adj_h_mid["Hz"], DiffGrad=ctx.adj_d_hz_dx, invDx=solver.inv_dx_e)
411
+
412
+ _cuda_backend._reverse_electric_cpml_ez(
413
+ AdjEzPrev=pre["Ez"],
414
+ GradEpsEz=ctx.grad_eps_ez,
415
+ AdjPsiPosPrev=pre["psi_ez_x"],
416
+ AdjPsiNegPrev=pre["psi_ez_y"],
417
+ AdjDPos=ctx.adj_d_hy_dx,
418
+ AdjDNeg=ctx.adj_d_hx_dy,
419
+ AdjEzPost=adjoint_state["Ez"],
420
+ AdjPsiPosPost=adjoint_state["psi_ez_x"],
421
+ AdjPsiNegPost=adjoint_state["psi_ez_y"],
422
+ EzDecay=solver.cez_decay,
423
+ EzCurl=ctx.ez_curl,
424
+ EpsEz=eps_ez,
425
+ PsiPos=forward_state["psi_ez_x"],
426
+ PsiNeg=forward_state["psi_ez_y"],
427
+ BPos=solver.cpml_b_e_x,
428
+ CPos=solver.cpml_c_e_x,
429
+ InvKappaPos=solver.cpml_inv_kappa_e_x,
430
+ BNeg=solver.cpml_b_e_y,
431
+ CNeg=solver.cpml_c_e_y,
432
+ InvKappaNeg=solver.cpml_inv_kappa_e_y,
433
+ HxMid=hx_mid,
434
+ HyMid=hy_mid,
435
+ invDx=solver.inv_dx_e,
436
+ invDy=solver.inv_dy_e,
437
+ xLowBoundaryMode=solver.boundary_x_low_code,
438
+ xHighBoundaryMode=solver.boundary_x_high_code,
439
+ yLowBoundaryMode=solver.boundary_y_low_code,
440
+ yHighBoundaryMode=solver.boundary_y_high_code,
441
+ )
442
+ _cuda_backend._accumulate_backward_diff_x(FieldGrad=adj_h_mid["Hy"], DiffGrad=ctx.adj_d_hy_dx, invDx=solver.inv_dx_e)
443
+ _cuda_backend._accumulate_backward_diff_y(FieldGrad=adj_h_mid["Hx"], DiffGrad=ctx.adj_d_hx_dy, invDy=solver.inv_dy_e)
444
+
445
+ # Phase 2: magnetic-decay + psi_h pullback -> pre-step H/psi_h adjoint,
446
+ # folding each curl(E)-derivative adjoint into the pre-step E adjoint.
447
+ _cuda_backend._reverse_magnetic_cpml_hx(
448
+ AdjHxPrev=pre["Hx"],
449
+ AdjPsiPosPrev=pre["psi_hx_y"],
450
+ AdjPsiNegPrev=pre["psi_hx_z"],
451
+ AdjDPos=ctx.adj_d_ez_dy,
452
+ AdjDNeg=ctx.adj_d_ey_dz,
453
+ AdjHxPost=adj_h_mid["Hx"],
454
+ AdjPsiPosPost=adjoint_state["psi_hx_y"],
455
+ AdjPsiNegPost=adjoint_state["psi_hx_z"],
456
+ HxDecay=solver.chx_decay,
457
+ HxCurl=solver.chx_curl,
458
+ BPos=solver.cpml_b_h_y,
459
+ CPos=solver.cpml_c_h_y,
460
+ InvKappaPos=solver.cpml_inv_kappa_h_y,
461
+ BNeg=solver.cpml_b_h_z,
462
+ CNeg=solver.cpml_c_h_z,
463
+ InvKappaNeg=solver.cpml_inv_kappa_h_z,
464
+ )
465
+ _cuda_backend._accumulate_forward_diff_y(FieldGrad=pre["Ez"], DiffGrad=ctx.adj_d_ez_dy, invDy=solver.inv_dy_h)
466
+ _cuda_backend._accumulate_forward_diff_z(FieldGrad=pre["Ey"], DiffGrad=ctx.adj_d_ey_dz, invDz=solver.inv_dz_h)
467
+
468
+ _cuda_backend._reverse_magnetic_cpml_hy(
469
+ AdjHyPrev=pre["Hy"],
470
+ AdjPsiPosPrev=pre["psi_hy_z"],
471
+ AdjPsiNegPrev=pre["psi_hy_x"],
472
+ AdjDPos=ctx.adj_d_ex_dz,
473
+ AdjDNeg=ctx.adj_d_ez_dx,
474
+ AdjHyPost=adj_h_mid["Hy"],
475
+ AdjPsiPosPost=adjoint_state["psi_hy_x"],
476
+ AdjPsiNegPost=adjoint_state["psi_hy_z"],
477
+ HyDecay=solver.chy_decay,
478
+ HyCurl=solver.chy_curl,
479
+ BPos=solver.cpml_b_h_z,
480
+ CPos=solver.cpml_c_h_z,
481
+ InvKappaPos=solver.cpml_inv_kappa_h_z,
482
+ BNeg=solver.cpml_b_h_x,
483
+ CNeg=solver.cpml_c_h_x,
484
+ InvKappaNeg=solver.cpml_inv_kappa_h_x,
485
+ )
486
+ _cuda_backend._accumulate_forward_diff_z(FieldGrad=pre["Ex"], DiffGrad=ctx.adj_d_ex_dz, invDz=solver.inv_dz_h)
487
+ _cuda_backend._accumulate_forward_diff_x(FieldGrad=pre["Ez"], DiffGrad=ctx.adj_d_ez_dx, invDx=solver.inv_dx_h)
488
+
489
+ _cuda_backend._reverse_magnetic_cpml_hz(
490
+ AdjHzPrev=pre["Hz"],
491
+ AdjPsiPosPrev=pre["psi_hz_x"],
492
+ AdjPsiNegPrev=pre["psi_hz_y"],
493
+ AdjDPos=ctx.adj_d_ey_dx,
494
+ AdjDNeg=ctx.adj_d_ex_dy,
495
+ AdjHzPost=adj_h_mid["Hz"],
496
+ AdjPsiPosPost=adjoint_state["psi_hz_x"],
497
+ AdjPsiNegPost=adjoint_state["psi_hz_y"],
498
+ HzDecay=solver.chz_decay,
499
+ HzCurl=solver.chz_curl,
500
+ BPos=solver.cpml_b_h_x,
501
+ CPos=solver.cpml_c_h_x,
502
+ InvKappaPos=solver.cpml_inv_kappa_h_x,
503
+ BNeg=solver.cpml_b_h_y,
504
+ CNeg=solver.cpml_c_h_y,
505
+ InvKappaNeg=solver.cpml_inv_kappa_h_y,
506
+ )
507
+ _cuda_backend._accumulate_forward_diff_x(FieldGrad=pre["Ey"], DiffGrad=ctx.adj_d_ey_dx, invDx=solver.inv_dx_h)
508
+ _cuda_backend._accumulate_forward_diff_y(FieldGrad=pre["Ex"], DiffGrad=ctx.adj_d_ex_dy, invDy=solver.inv_dy_h)
509
+
510
+ step_result = _adjoint._ReverseStepResult(
511
+ pre_step_adjoint=pre,
512
+ grad_eps_ex=ctx.grad_eps_ex,
513
+ grad_eps_ey=ctx.grad_eps_ey,
514
+ grad_eps_ez=ctx.grad_eps_ez,
515
+ backend=_NATIVE_REVERSE_LABELS[_ReverseBackend.CPML],
516
+ magnetic_output_adjoint=adj_h_mid,
517
+ )
518
+ return step_result
519
+
520
+
521
+ def _reverse_step_cpml_native(
522
+ solver,
523
+ forward_state,
524
+ adjoint_state,
525
+ *,
526
+ time_value,
527
+ eps_ex,
528
+ eps_ey,
529
+ eps_ez,
530
+ resolved_source_terms,
531
+ profiler,
532
+ ):
533
+ """Fully native reverse step for the CPML (absorbing-boundary) configuration."""
534
+ from . import core as _adjoint
535
+
536
+ step_result = _reverse_step_cpml_native_core(
537
+ solver,
538
+ forward_state,
539
+ adjoint_state,
540
+ time_value=time_value,
541
+ eps_ex=eps_ex,
542
+ eps_ey=eps_ey,
543
+ eps_ez=eps_ez,
544
+ resolved_source_terms=resolved_source_terms,
545
+ )
546
+ # The native runner owns the full reverse-step contract, including the
547
+ # analytic source-term eps-gradient accumulation the reference path applies.
548
+ return _adjoint._accumulate_source_term_gradients(
549
+ step_result,
550
+ solver=solver,
551
+ adjoint_state=adjoint_state,
552
+ time_value=time_value,
553
+ eps_ex=eps_ex,
554
+ eps_ey=eps_ey,
555
+ eps_ez=eps_ez,
556
+ resolved_source_terms=resolved_source_terms,
557
+ )
558
+
559
+
560
+ def _reverse_step_conductive_native(
561
+ solver,
562
+ forward_state,
563
+ adjoint_state,
564
+ *,
565
+ time_value,
566
+ eps_ex,
567
+ eps_ey,
568
+ eps_ez,
569
+ resolved_source_terms,
570
+ profiler,
571
+ ):
572
+ """Fully native reverse step for a static-conductive (sigma_e) CPML medium.
573
+
574
+ Mirrors :func:`reverse_step_conductive_cpml_native`. The magnetic
575
+ reverse, the curl(H)/curl(E) difference folds, and the pre-step field/psi
576
+ accumulation are the linear CPML machinery unchanged; the only conduction
577
+ specialization is the electric reverse kernel, which recomputes the eps
578
+ sensitivity of the semi-implicit ``decay``/``curl`` pair analytically instead
579
+ of applying the linear rule. Launch order is load-bearing: the conductive
580
+ electric kernels *assign* the pre-step E/psi + eps gradient and the
581
+ difference folds *accumulate* into the mid-step H adjoint, so all three
582
+ electric kernels (and their folds) must complete before any magnetic kernel
583
+ reads the mid-step H adjoint.
584
+ """
585
+ from . import core as _adjoint
586
+ from .reverse_common import allocate_cpml_reverse_context
587
+
588
+ magnetic_fields = _adjoint._forward_magnetic_fields(
589
+ solver,
590
+ forward_state,
591
+ time_value=time_value,
592
+ resolved_source_terms=resolved_source_terms,
593
+ )
594
+ hx_mid = magnetic_fields["Hx"]
595
+ hy_mid = magnetic_fields["Hy"]
596
+ hz_mid = magnetic_fields["Hz"]
597
+
598
+ coeffs = _adjoint._conductive_reverse_coefficients(
599
+ solver, eps_ex=eps_ex, eps_ey=eps_ey, eps_ez=eps_ez
600
+ )
601
+ decay_ex, curl_ex, half_ex = coeffs["Ex"]
602
+ decay_ey, curl_ey, half_ey = coeffs["Ey"]
603
+ decay_ez, curl_ez, half_ez = coeffs["Ez"]
604
+ dt = float(solver.dt)
605
+
606
+ ctx = allocate_cpml_reverse_context(
607
+ solver,
608
+ forward_state,
609
+ adjoint_state,
610
+ eps_ex=eps_ex,
611
+ eps_ey=eps_ey,
612
+ eps_ez=eps_ez,
613
+ )
614
+ pre = ctx.pre_step_adjoint
615
+ adj_h_mid = ctx.magnetic_output_adjoint
616
+
617
+ # Phase 1: conductive electric adjoint -> pre-step E/psi_e adjoint + eps
618
+ # gradient, folding each curl(H)-derivative adjoint into the mid-step H.
619
+ _cuda_backend._reverse_electric_cpml_conductive_ex(
620
+ AdjExPrev=pre["Ex"],
621
+ GradEpsEx=ctx.grad_eps_ex,
622
+ AdjPsiPosPrev=pre["psi_ex_y"],
623
+ AdjPsiNegPrev=pre["psi_ex_z"],
624
+ AdjDPos=ctx.adj_d_hz_dy,
625
+ AdjDNeg=ctx.adj_d_hy_dz,
626
+ AdjExPost=adjoint_state["Ex"],
627
+ AdjPsiPosPost=adjoint_state["psi_ex_y"],
628
+ AdjPsiNegPost=adjoint_state["psi_ex_z"],
629
+ ExDecay=decay_ex,
630
+ ExCurl=curl_ex,
631
+ ExHalf=half_ex,
632
+ ExPrev=forward_state["Ex"],
633
+ EpsEx=eps_ex,
634
+ Dt=dt,
635
+ PsiPos=forward_state["psi_ex_y"],
636
+ PsiNeg=forward_state["psi_ex_z"],
637
+ BPos=solver.cpml_b_e_y,
638
+ CPos=solver.cpml_c_e_y,
639
+ InvKappaPos=solver.cpml_inv_kappa_e_y,
640
+ BNeg=solver.cpml_b_e_z,
641
+ CNeg=solver.cpml_c_e_z,
642
+ InvKappaNeg=solver.cpml_inv_kappa_e_z,
643
+ HyMid=hy_mid,
644
+ HzMid=hz_mid,
645
+ invDy=solver.inv_dy_e,
646
+ invDz=solver.inv_dz_e,
647
+ yLowBoundaryMode=solver.boundary_y_low_code,
648
+ yHighBoundaryMode=solver.boundary_y_high_code,
649
+ zLowBoundaryMode=solver.boundary_z_low_code,
650
+ zHighBoundaryMode=solver.boundary_z_high_code,
651
+ )
652
+ _cuda_backend._accumulate_backward_diff_y(FieldGrad=adj_h_mid["Hz"], DiffGrad=ctx.adj_d_hz_dy, invDy=solver.inv_dy_e)
653
+ _cuda_backend._accumulate_backward_diff_z(FieldGrad=adj_h_mid["Hy"], DiffGrad=ctx.adj_d_hy_dz, invDz=solver.inv_dz_e)
654
+
655
+ _cuda_backend._reverse_electric_cpml_conductive_ey(
656
+ AdjEyPrev=pre["Ey"],
657
+ GradEpsEy=ctx.grad_eps_ey,
658
+ AdjPsiPosPrev=pre["psi_ey_z"],
659
+ AdjPsiNegPrev=pre["psi_ey_x"],
660
+ AdjDPos=ctx.adj_d_hx_dz,
661
+ AdjDNeg=ctx.adj_d_hz_dx,
662
+ AdjEyPost=adjoint_state["Ey"],
663
+ AdjPsiPosPost=adjoint_state["psi_ey_x"],
664
+ AdjPsiNegPost=adjoint_state["psi_ey_z"],
665
+ EyDecay=decay_ey,
666
+ EyCurl=curl_ey,
667
+ EyHalf=half_ey,
668
+ EyPrev=forward_state["Ey"],
669
+ EpsEy=eps_ey,
670
+ Dt=dt,
671
+ PsiPos=forward_state["psi_ey_z"],
672
+ PsiNeg=forward_state["psi_ey_x"],
673
+ BPos=solver.cpml_b_e_z,
674
+ CPos=solver.cpml_c_e_z,
675
+ InvKappaPos=solver.cpml_inv_kappa_e_z,
676
+ BNeg=solver.cpml_b_e_x,
677
+ CNeg=solver.cpml_c_e_x,
678
+ InvKappaNeg=solver.cpml_inv_kappa_e_x,
679
+ HxMid=hx_mid,
680
+ HzMid=hz_mid,
681
+ invDx=solver.inv_dx_e,
682
+ invDz=solver.inv_dz_e,
683
+ xLowBoundaryMode=solver.boundary_x_low_code,
684
+ xHighBoundaryMode=solver.boundary_x_high_code,
685
+ zLowBoundaryMode=solver.boundary_z_low_code,
686
+ zHighBoundaryMode=solver.boundary_z_high_code,
687
+ )
688
+ _cuda_backend._accumulate_backward_diff_z(FieldGrad=adj_h_mid["Hx"], DiffGrad=ctx.adj_d_hx_dz, invDz=solver.inv_dz_e)
689
+ _cuda_backend._accumulate_backward_diff_x(FieldGrad=adj_h_mid["Hz"], DiffGrad=ctx.adj_d_hz_dx, invDx=solver.inv_dx_e)
690
+
691
+ _cuda_backend._reverse_electric_cpml_conductive_ez(
692
+ AdjEzPrev=pre["Ez"],
693
+ GradEpsEz=ctx.grad_eps_ez,
694
+ AdjPsiPosPrev=pre["psi_ez_x"],
695
+ AdjPsiNegPrev=pre["psi_ez_y"],
696
+ AdjDPos=ctx.adj_d_hy_dx,
697
+ AdjDNeg=ctx.adj_d_hx_dy,
698
+ AdjEzPost=adjoint_state["Ez"],
699
+ AdjPsiPosPost=adjoint_state["psi_ez_x"],
700
+ AdjPsiNegPost=adjoint_state["psi_ez_y"],
701
+ EzDecay=decay_ez,
702
+ EzCurl=curl_ez,
703
+ EzHalf=half_ez,
704
+ EzPrev=forward_state["Ez"],
705
+ EpsEz=eps_ez,
706
+ Dt=dt,
707
+ PsiPos=forward_state["psi_ez_x"],
708
+ PsiNeg=forward_state["psi_ez_y"],
709
+ BPos=solver.cpml_b_e_x,
710
+ CPos=solver.cpml_c_e_x,
711
+ InvKappaPos=solver.cpml_inv_kappa_e_x,
712
+ BNeg=solver.cpml_b_e_y,
713
+ CNeg=solver.cpml_c_e_y,
714
+ InvKappaNeg=solver.cpml_inv_kappa_e_y,
715
+ HxMid=hx_mid,
716
+ HyMid=hy_mid,
717
+ invDx=solver.inv_dx_e,
718
+ invDy=solver.inv_dy_e,
719
+ xLowBoundaryMode=solver.boundary_x_low_code,
720
+ xHighBoundaryMode=solver.boundary_x_high_code,
721
+ yLowBoundaryMode=solver.boundary_y_low_code,
722
+ yHighBoundaryMode=solver.boundary_y_high_code,
723
+ )
724
+ _cuda_backend._accumulate_backward_diff_x(FieldGrad=adj_h_mid["Hy"], DiffGrad=ctx.adj_d_hy_dx, invDx=solver.inv_dx_e)
725
+ _cuda_backend._accumulate_backward_diff_y(FieldGrad=adj_h_mid["Hx"], DiffGrad=ctx.adj_d_hx_dy, invDy=solver.inv_dy_e)
726
+
727
+ # Phase 2: magnetic-decay + psi_h pullback -> pre-step H/psi_h adjoint,
728
+ # folding each curl(E)-derivative adjoint into the pre-step E adjoint. This is
729
+ # the linear CPML magnetic reverse unchanged (conduction is electric-only).
730
+ _cuda_backend._reverse_magnetic_cpml_hx(
731
+ AdjHxPrev=pre["Hx"],
732
+ AdjPsiPosPrev=pre["psi_hx_y"],
733
+ AdjPsiNegPrev=pre["psi_hx_z"],
734
+ AdjDPos=ctx.adj_d_ez_dy,
735
+ AdjDNeg=ctx.adj_d_ey_dz,
736
+ AdjHxPost=adj_h_mid["Hx"],
737
+ AdjPsiPosPost=adjoint_state["psi_hx_y"],
738
+ AdjPsiNegPost=adjoint_state["psi_hx_z"],
739
+ HxDecay=solver.chx_decay,
740
+ HxCurl=solver.chx_curl,
741
+ BPos=solver.cpml_b_h_y,
742
+ CPos=solver.cpml_c_h_y,
743
+ InvKappaPos=solver.cpml_inv_kappa_h_y,
744
+ BNeg=solver.cpml_b_h_z,
745
+ CNeg=solver.cpml_c_h_z,
746
+ InvKappaNeg=solver.cpml_inv_kappa_h_z,
747
+ )
748
+ _cuda_backend._accumulate_forward_diff_y(FieldGrad=pre["Ez"], DiffGrad=ctx.adj_d_ez_dy, invDy=solver.inv_dy_h)
749
+ _cuda_backend._accumulate_forward_diff_z(FieldGrad=pre["Ey"], DiffGrad=ctx.adj_d_ey_dz, invDz=solver.inv_dz_h)
750
+
751
+ _cuda_backend._reverse_magnetic_cpml_hy(
752
+ AdjHyPrev=pre["Hy"],
753
+ AdjPsiPosPrev=pre["psi_hy_z"],
754
+ AdjPsiNegPrev=pre["psi_hy_x"],
755
+ AdjDPos=ctx.adj_d_ex_dz,
756
+ AdjDNeg=ctx.adj_d_ez_dx,
757
+ AdjHyPost=adj_h_mid["Hy"],
758
+ AdjPsiPosPost=adjoint_state["psi_hy_x"],
759
+ AdjPsiNegPost=adjoint_state["psi_hy_z"],
760
+ HyDecay=solver.chy_decay,
761
+ HyCurl=solver.chy_curl,
762
+ BPos=solver.cpml_b_h_z,
763
+ CPos=solver.cpml_c_h_z,
764
+ InvKappaPos=solver.cpml_inv_kappa_h_z,
765
+ BNeg=solver.cpml_b_h_x,
766
+ CNeg=solver.cpml_c_h_x,
767
+ InvKappaNeg=solver.cpml_inv_kappa_h_x,
768
+ )
769
+ _cuda_backend._accumulate_forward_diff_z(FieldGrad=pre["Ex"], DiffGrad=ctx.adj_d_ex_dz, invDz=solver.inv_dz_h)
770
+ _cuda_backend._accumulate_forward_diff_x(FieldGrad=pre["Ez"], DiffGrad=ctx.adj_d_ez_dx, invDx=solver.inv_dx_h)
771
+
772
+ _cuda_backend._reverse_magnetic_cpml_hz(
773
+ AdjHzPrev=pre["Hz"],
774
+ AdjPsiPosPrev=pre["psi_hz_x"],
775
+ AdjPsiNegPrev=pre["psi_hz_y"],
776
+ AdjDPos=ctx.adj_d_ey_dx,
777
+ AdjDNeg=ctx.adj_d_ex_dy,
778
+ AdjHzPost=adj_h_mid["Hz"],
779
+ AdjPsiPosPost=adjoint_state["psi_hz_x"],
780
+ AdjPsiNegPost=adjoint_state["psi_hz_y"],
781
+ HzDecay=solver.chz_decay,
782
+ HzCurl=solver.chz_curl,
783
+ BPos=solver.cpml_b_h_x,
784
+ CPos=solver.cpml_c_h_x,
785
+ InvKappaPos=solver.cpml_inv_kappa_h_x,
786
+ BNeg=solver.cpml_b_h_y,
787
+ CNeg=solver.cpml_c_h_y,
788
+ InvKappaNeg=solver.cpml_inv_kappa_h_y,
789
+ )
790
+ _cuda_backend._accumulate_forward_diff_x(FieldGrad=pre["Ey"], DiffGrad=ctx.adj_d_ey_dx, invDx=solver.inv_dx_h)
791
+ _cuda_backend._accumulate_forward_diff_y(FieldGrad=pre["Ex"], DiffGrad=ctx.adj_d_ex_dy, invDy=solver.inv_dy_h)
792
+
793
+ step_result = _adjoint._ReverseStepResult(
794
+ pre_step_adjoint=pre,
795
+ grad_eps_ex=ctx.grad_eps_ex,
796
+ grad_eps_ey=ctx.grad_eps_ey,
797
+ grad_eps_ez=ctx.grad_eps_ez,
798
+ backend=_NATIVE_REVERSE_LABELS[_ReverseBackend.CONDUCTIVE],
799
+ magnetic_output_adjoint=adj_h_mid,
800
+ )
801
+ return _adjoint._accumulate_source_term_gradients(
802
+ step_result,
803
+ solver=solver,
804
+ adjoint_state=adjoint_state,
805
+ time_value=time_value,
806
+ eps_ex=eps_ex,
807
+ eps_ey=eps_ey,
808
+ eps_ez=eps_ez,
809
+ resolved_source_terms=resolved_source_terms,
810
+ )
811
+
812
+
813
+ def _reverse_cpml_magnetic_phase_native(solver, ctx, adjoint_state, *, suffix=""):
814
+ """Run the shared linear magnetic CPML pullback after an electric variant."""
815
+ pre = ctx.pre_step_adjoint
816
+ adj_h_mid = ctx.magnetic_output_adjoint
817
+ for component, pos_name, neg_name, d_pos, d_neg, curl, b_pos, c_pos, k_pos, b_neg, c_neg, k_neg in (
818
+ ("Hx", "psi_hx_y", "psi_hx_z", ctx.adj_d_ez_dy, ctx.adj_d_ey_dz, solver.chx_curl,
819
+ solver.cpml_b_h_y, solver.cpml_c_h_y, solver.cpml_inv_kappa_h_y,
820
+ solver.cpml_b_h_z, solver.cpml_c_h_z, solver.cpml_inv_kappa_h_z),
821
+ ("Hy", "psi_hy_z", "psi_hy_x", ctx.adj_d_ex_dz, ctx.adj_d_ez_dx, solver.chy_curl,
822
+ solver.cpml_b_h_z, solver.cpml_c_h_z, solver.cpml_inv_kappa_h_z,
823
+ solver.cpml_b_h_x, solver.cpml_c_h_x, solver.cpml_inv_kappa_h_x),
824
+ ("Hz", "psi_hz_x", "psi_hz_y", ctx.adj_d_ey_dx, ctx.adj_d_ex_dy, solver.chz_curl,
825
+ solver.cpml_b_h_x, solver.cpml_c_h_x, solver.cpml_inv_kappa_h_x,
826
+ solver.cpml_b_h_y, solver.cpml_c_h_y, solver.cpml_inv_kappa_h_y),
827
+ ):
828
+ field_key = component + suffix
829
+ pos_key, neg_key = pos_name + suffix, neg_name + suffix
830
+ getattr(_cuda_backend, f"_reverse_magnetic_cpml_{component.lower()}")(
831
+ **{f"Adj{component}Prev": pre[field_key]},
832
+ AdjPsiPosPrev=pre[pos_key], AdjPsiNegPrev=pre[neg_key],
833
+ AdjDPos=d_pos, AdjDNeg=d_neg,
834
+ **{f"Adj{component}Post": adj_h_mid[field_key]},
835
+ AdjPsiPosPost=adjoint_state[pos_key], AdjPsiNegPost=adjoint_state[neg_key],
836
+ **{f"{component}Decay": getattr(solver, f"c{component.lower()}_decay")},
837
+ **{f"{component}Curl": curl},
838
+ BPos=b_pos, CPos=c_pos, InvKappaPos=k_pos,
839
+ BNeg=b_neg, CNeg=c_neg, InvKappaNeg=k_neg,
840
+ )
841
+ _cuda_backend._accumulate_forward_diff_y(FieldGrad=pre["Ez" + suffix], DiffGrad=ctx.adj_d_ez_dy, invDy=solver.inv_dy_h)
842
+ _cuda_backend._accumulate_forward_diff_z(FieldGrad=pre["Ey" + suffix], DiffGrad=ctx.adj_d_ey_dz, invDz=solver.inv_dz_h)
843
+ _cuda_backend._accumulate_forward_diff_z(FieldGrad=pre["Ex" + suffix], DiffGrad=ctx.adj_d_ex_dz, invDz=solver.inv_dz_h)
844
+ _cuda_backend._accumulate_forward_diff_x(FieldGrad=pre["Ez" + suffix], DiffGrad=ctx.adj_d_ez_dx, invDx=solver.inv_dx_h)
845
+ _cuda_backend._accumulate_forward_diff_x(FieldGrad=pre["Ey" + suffix], DiffGrad=ctx.adj_d_ey_dx, invDx=solver.inv_dx_h)
846
+ _cuda_backend._accumulate_forward_diff_y(FieldGrad=pre["Ex" + suffix], DiffGrad=ctx.adj_d_ex_dy, invDy=solver.inv_dy_h)
847
+
848
+
849
+ def _reverse_step_kerr_native(
850
+ solver,
851
+ forward_state,
852
+ adjoint_state,
853
+ *,
854
+ time_value,
855
+ eps_ex,
856
+ eps_ey,
857
+ eps_ez,
858
+ resolved_source_terms,
859
+ profiler,
860
+ ):
861
+ """Fully native reverse step for an instantaneous Kerr (chi3) CPML medium.
862
+
863
+ Mirrors :func:`reverse_step_kerr_cpml_native`. The Kerr forward
864
+ rewrites only the electric ``curl`` coefficient each step from the pre-update
865
+ fields (``curl = (dt / eff) * decay`` with ``eff = eps + eps0 * chi3 * |E|^2``),
866
+ so the reverse math splits into: (1) a native forward collocation of the frozen
867
+ fields to the per-edge ``|E|^2`` (``fsq``); (2) the fused Kerr electric reverse
868
+ kernels, which assign the pre-step E/psi adjoint (decay pullback) plus the
869
+ grad_eps / grad_chi3 coefficient sensitivities and emit the per-edge ``|E|^2``
870
+ cotangent ``g_fsq``, while folding curl(H) into the mid-step H adjoint; (3) the
871
+ linear CPML magnetic reverse (the Kerr term is electric-only); and (4) the
872
+ shared ``collocation_transpose`` scattering ``g_fsq`` back onto the pre-step
873
+ fields. Launch order is load-bearing: the electric kernels *assign* the pre-step
874
+ E adjoint before the magnetic forward-diff folds and the collocation transpose
875
+ *accumulate* onto it, and all three electric kernels (and their curl(H) folds)
876
+ must complete before any magnetic kernel reads the mid-step H adjoint.
877
+ """
878
+ import torch
879
+
880
+ from . import core as _adjoint
881
+ from .reverse_common import allocate_cpml_reverse_context
882
+
883
+ magnetic_fields = _adjoint._forward_magnetic_fields(
884
+ solver,
885
+ forward_state,
886
+ time_value=time_value,
887
+ resolved_source_terms=resolved_source_terms,
888
+ )
889
+ hx_mid = magnetic_fields["Hx"]
890
+ hy_mid = magnetic_fields["Hy"]
891
+ hz_mid = magnetic_fields["Hz"]
892
+
893
+ eps0 = float(solver.eps0)
894
+ dt = float(solver.dt)
895
+
896
+ # Per-edge |E|^2 collocation of the frozen checkpoint fields (native forward).
897
+ fsq_ex = torch.empty_like(eps_ex)
898
+ fsq_ey = torch.empty_like(eps_ey)
899
+ fsq_ez = torch.empty_like(eps_ez)
900
+ _cuda_backend._collocate_field_square(
901
+ FsqEx=fsq_ex,
902
+ FsqEy=fsq_ey,
903
+ FsqEz=fsq_ez,
904
+ Ex=forward_state["Ex"],
905
+ Ey=forward_state["Ey"],
906
+ Ez=forward_state["Ez"],
907
+ )
908
+
909
+ ctx = allocate_cpml_reverse_context(
910
+ solver,
911
+ forward_state,
912
+ adjoint_state,
913
+ eps_ex=eps_ex,
914
+ eps_ey=eps_ey,
915
+ eps_ez=eps_ez,
916
+ )
917
+ pre = ctx.pre_step_adjoint
918
+ adj_h_mid = ctx.magnetic_output_adjoint
919
+
920
+ grad_chi3_ex = torch.empty_like(eps_ex)
921
+ grad_chi3_ey = torch.empty_like(eps_ey)
922
+ grad_chi3_ez = torch.empty_like(eps_ez)
923
+ g_fsq_ex = torch.empty_like(eps_ex)
924
+ g_fsq_ey = torch.empty_like(eps_ey)
925
+ g_fsq_ez = torch.empty_like(eps_ez)
926
+
927
+ # Phase 1: Kerr electric adjoint -> pre-step E/psi adjoint + grad_eps/grad_chi3
928
+ # + |E|^2 cotangent, folding each curl(H) derivative into the mid-step H.
929
+ _cuda_backend._reverse_electric_cpml_kerr_ex(
930
+ AdjExPrev=pre["Ex"],
931
+ GradEpsEx=ctx.grad_eps_ex,
932
+ GradChi3Ex=grad_chi3_ex,
933
+ GFsqEx=g_fsq_ex,
934
+ AdjPsiPosPrev=pre["psi_ex_y"],
935
+ AdjPsiNegPrev=pre["psi_ex_z"],
936
+ AdjDPos=ctx.adj_d_hz_dy,
937
+ AdjDNeg=ctx.adj_d_hy_dz,
938
+ AdjExPost=adjoint_state["Ex"],
939
+ AdjPsiPosPost=adjoint_state["psi_ex_y"],
940
+ AdjPsiNegPost=adjoint_state["psi_ex_z"],
941
+ ExDecay=solver.cex_decay,
942
+ EpsEx=eps_ex,
943
+ Chi3Ex=solver.kerr_chi3_Ex,
944
+ FsqEx=fsq_ex,
945
+ Dt=dt,
946
+ Eps0=eps0,
947
+ PsiPos=forward_state["psi_ex_y"],
948
+ PsiNeg=forward_state["psi_ex_z"],
949
+ BPos=solver.cpml_b_e_y,
950
+ CPos=solver.cpml_c_e_y,
951
+ InvKappaPos=solver.cpml_inv_kappa_e_y,
952
+ BNeg=solver.cpml_b_e_z,
953
+ CNeg=solver.cpml_c_e_z,
954
+ InvKappaNeg=solver.cpml_inv_kappa_e_z,
955
+ HyMid=hy_mid,
956
+ HzMid=hz_mid,
957
+ invDy=solver.inv_dy_e,
958
+ invDz=solver.inv_dz_e,
959
+ yLowBoundaryMode=solver.boundary_y_low_code,
960
+ yHighBoundaryMode=solver.boundary_y_high_code,
961
+ zLowBoundaryMode=solver.boundary_z_low_code,
962
+ zHighBoundaryMode=solver.boundary_z_high_code,
963
+ )
964
+ _cuda_backend._accumulate_backward_diff_y(FieldGrad=adj_h_mid["Hz"], DiffGrad=ctx.adj_d_hz_dy, invDy=solver.inv_dy_e)
965
+ _cuda_backend._accumulate_backward_diff_z(FieldGrad=adj_h_mid["Hy"], DiffGrad=ctx.adj_d_hy_dz, invDz=solver.inv_dz_e)
966
+
967
+ _cuda_backend._reverse_electric_cpml_kerr_ey(
968
+ AdjEyPrev=pre["Ey"],
969
+ GradEpsEy=ctx.grad_eps_ey,
970
+ GradChi3Ey=grad_chi3_ey,
971
+ GFsqEy=g_fsq_ey,
972
+ AdjPsiPosPrev=pre["psi_ey_z"],
973
+ AdjPsiNegPrev=pre["psi_ey_x"],
974
+ AdjDPos=ctx.adj_d_hx_dz,
975
+ AdjDNeg=ctx.adj_d_hz_dx,
976
+ AdjEyPost=adjoint_state["Ey"],
977
+ AdjPsiPosPost=adjoint_state["psi_ey_x"],
978
+ AdjPsiNegPost=adjoint_state["psi_ey_z"],
979
+ EyDecay=solver.cey_decay,
980
+ EpsEy=eps_ey,
981
+ Chi3Ey=solver.kerr_chi3_Ey,
982
+ FsqEy=fsq_ey,
983
+ Dt=dt,
984
+ Eps0=eps0,
985
+ PsiPos=forward_state["psi_ey_z"],
986
+ PsiNeg=forward_state["psi_ey_x"],
987
+ BPos=solver.cpml_b_e_z,
988
+ CPos=solver.cpml_c_e_z,
989
+ InvKappaPos=solver.cpml_inv_kappa_e_z,
990
+ BNeg=solver.cpml_b_e_x,
991
+ CNeg=solver.cpml_c_e_x,
992
+ InvKappaNeg=solver.cpml_inv_kappa_e_x,
993
+ HxMid=hx_mid,
994
+ HzMid=hz_mid,
995
+ invDx=solver.inv_dx_e,
996
+ invDz=solver.inv_dz_e,
997
+ xLowBoundaryMode=solver.boundary_x_low_code,
998
+ xHighBoundaryMode=solver.boundary_x_high_code,
999
+ zLowBoundaryMode=solver.boundary_z_low_code,
1000
+ zHighBoundaryMode=solver.boundary_z_high_code,
1001
+ )
1002
+ _cuda_backend._accumulate_backward_diff_z(FieldGrad=adj_h_mid["Hx"], DiffGrad=ctx.adj_d_hx_dz, invDz=solver.inv_dz_e)
1003
+ _cuda_backend._accumulate_backward_diff_x(FieldGrad=adj_h_mid["Hz"], DiffGrad=ctx.adj_d_hz_dx, invDx=solver.inv_dx_e)
1004
+
1005
+ _cuda_backend._reverse_electric_cpml_kerr_ez(
1006
+ AdjEzPrev=pre["Ez"],
1007
+ GradEpsEz=ctx.grad_eps_ez,
1008
+ GradChi3Ez=grad_chi3_ez,
1009
+ GFsqEz=g_fsq_ez,
1010
+ AdjPsiPosPrev=pre["psi_ez_x"],
1011
+ AdjPsiNegPrev=pre["psi_ez_y"],
1012
+ AdjDPos=ctx.adj_d_hy_dx,
1013
+ AdjDNeg=ctx.adj_d_hx_dy,
1014
+ AdjEzPost=adjoint_state["Ez"],
1015
+ AdjPsiPosPost=adjoint_state["psi_ez_x"],
1016
+ AdjPsiNegPost=adjoint_state["psi_ez_y"],
1017
+ EzDecay=solver.cez_decay,
1018
+ EpsEz=eps_ez,
1019
+ Chi3Ez=solver.kerr_chi3_Ez,
1020
+ FsqEz=fsq_ez,
1021
+ Dt=dt,
1022
+ Eps0=eps0,
1023
+ PsiPos=forward_state["psi_ez_x"],
1024
+ PsiNeg=forward_state["psi_ez_y"],
1025
+ BPos=solver.cpml_b_e_x,
1026
+ CPos=solver.cpml_c_e_x,
1027
+ InvKappaPos=solver.cpml_inv_kappa_e_x,
1028
+ BNeg=solver.cpml_b_e_y,
1029
+ CNeg=solver.cpml_c_e_y,
1030
+ InvKappaNeg=solver.cpml_inv_kappa_e_y,
1031
+ HxMid=hx_mid,
1032
+ HyMid=hy_mid,
1033
+ invDx=solver.inv_dx_e,
1034
+ invDy=solver.inv_dy_e,
1035
+ xLowBoundaryMode=solver.boundary_x_low_code,
1036
+ xHighBoundaryMode=solver.boundary_x_high_code,
1037
+ yLowBoundaryMode=solver.boundary_y_low_code,
1038
+ yHighBoundaryMode=solver.boundary_y_high_code,
1039
+ )
1040
+ _cuda_backend._accumulate_backward_diff_x(FieldGrad=adj_h_mid["Hy"], DiffGrad=ctx.adj_d_hy_dx, invDx=solver.inv_dx_e)
1041
+ _cuda_backend._accumulate_backward_diff_y(FieldGrad=adj_h_mid["Hx"], DiffGrad=ctx.adj_d_hx_dy, invDy=solver.inv_dy_e)
1042
+
1043
+ # Phase 2: linear CPML magnetic reverse (the Kerr term is electric-only).
1044
+ _reverse_cpml_magnetic_phase_native(solver, ctx, adjoint_state)
1045
+
1046
+ # Phase 3: scatter the |E|^2 cotangent back onto the pre-step fields (the exact
1047
+ # transpose of the collocation, own-axis 2*E plus the 4-point off-axis average).
1048
+ # It accumulates onto the pre-step E adjoint the electric kernels already
1049
+ # assigned and the magnetic folds added, so it runs last.
1050
+ _cuda_backend._collocation_transpose(
1051
+ AdjEx=pre["Ex"],
1052
+ AdjEy=pre["Ey"],
1053
+ AdjEz=pre["Ez"],
1054
+ GEx=g_fsq_ex,
1055
+ GEy=g_fsq_ey,
1056
+ GEz=g_fsq_ez,
1057
+ Ex=forward_state["Ex"],
1058
+ Ey=forward_state["Ey"],
1059
+ Ez=forward_state["Ez"],
1060
+ )
1061
+
1062
+ step_result = _adjoint._ReverseStepResult(
1063
+ pre_step_adjoint=pre,
1064
+ grad_eps_ex=ctx.grad_eps_ex,
1065
+ grad_eps_ey=ctx.grad_eps_ey,
1066
+ grad_eps_ez=ctx.grad_eps_ez,
1067
+ backend=_NATIVE_REVERSE_LABELS[_ReverseBackend.KERR],
1068
+ magnetic_output_adjoint=adj_h_mid,
1069
+ grad_chi3_ex=grad_chi3_ex,
1070
+ grad_chi3_ey=grad_chi3_ey,
1071
+ grad_chi3_ez=grad_chi3_ez,
1072
+ )
1073
+ return _adjoint._accumulate_source_term_gradients(
1074
+ step_result,
1075
+ solver=solver,
1076
+ adjoint_state=adjoint_state,
1077
+ time_value=time_value,
1078
+ eps_ex=eps_ex,
1079
+ eps_ey=eps_ey,
1080
+ eps_ez=eps_ez,
1081
+ resolved_source_terms=resolved_source_terms,
1082
+ )
1083
+
1084
+
1085
+ def _reverse_step_nonlinear_native(
1086
+ solver, forward_state, adjoint_state, *, time_value, eps_ex, eps_ey, eps_ez,
1087
+ resolved_source_terms, profiler,
1088
+ ):
1089
+ """Native reverse for chi2/chi3/two-photon absorption on standard or CPML grids."""
1090
+ import torch
1091
+
1092
+ from . import core as _adjoint
1093
+ from .reverse_common import allocate_cpml_reverse_context
1094
+
1095
+ magnetic_fields = _adjoint._forward_magnetic_fields(
1096
+ solver, forward_state, time_value=time_value,
1097
+ resolved_source_terms=resolved_source_terms,
1098
+ )
1099
+ original_state_names = tuple(forward_state)
1100
+ if not getattr(solver, "uses_cpml", False):
1101
+ psi_components = {
1102
+ "psi_ex_y": "Ex", "psi_ex_z": "Ex",
1103
+ "psi_ey_x": "Ey", "psi_ey_z": "Ey",
1104
+ "psi_ez_x": "Ez", "psi_ez_y": "Ez",
1105
+ "psi_hx_y": "Hx", "psi_hx_z": "Hx",
1106
+ "psi_hy_x": "Hy", "psi_hy_z": "Hy",
1107
+ "psi_hz_x": "Hz", "psi_hz_y": "Hz",
1108
+ }
1109
+ forward_state = dict(forward_state)
1110
+ adjoint_state = dict(adjoint_state)
1111
+ for psi_name, component_name in psi_components.items():
1112
+ forward_state[psi_name] = torch.zeros_like(forward_state[component_name])
1113
+ adjoint_state[psi_name] = torch.zeros_like(adjoint_state[component_name])
1114
+ for stagger in ("e", "h"):
1115
+ for axis in ("x", "y", "z"):
1116
+ vector = getattr(solver, f"inv_d{axis}_{stagger}")
1117
+ setattr(solver, f"cpml_b_{stagger}_{axis}", torch.ones_like(vector))
1118
+ setattr(solver, f"cpml_c_{stagger}_{axis}", torch.zeros_like(vector))
1119
+ setattr(solver, f"cpml_inv_kappa_{stagger}_{axis}", torch.ones_like(vector))
1120
+ dispersive_state = _adjoint._advance_dispersive_state(solver, forward_state)
1121
+ fsq = {name: torch.empty_like(eps) for name, eps in (
1122
+ ("Ex", eps_ex), ("Ey", eps_ey), ("Ez", eps_ez)
1123
+ )}
1124
+ _cuda_backend._collocate_field_square(
1125
+ FsqEx=fsq["Ex"], FsqEy=fsq["Ey"], FsqEz=fsq["Ez"],
1126
+ Ex=forward_state["Ex"], Ey=forward_state["Ey"], Ez=forward_state["Ez"],
1127
+ )
1128
+ ctx = allocate_cpml_reverse_context(
1129
+ solver, forward_state, adjoint_state,
1130
+ eps_ex=eps_ex, eps_ey=eps_ey, eps_ez=eps_ez,
1131
+ )
1132
+ pre, adj_h = ctx.pre_step_adjoint, ctx.magnetic_output_adjoint
1133
+ grad_chi2 = {name: torch.empty_like(value) for name, value in fsq.items()}
1134
+ grad_chi3 = {name: torch.empty_like(value) for name, value in fsq.items()}
1135
+ grad_tpa = {name: torch.empty_like(value) for name, value in fsq.items()}
1136
+ g_fsq = {name: torch.empty_like(value) for name, value in fsq.items()}
1137
+ configs = (
1138
+ (0, "Ex", eps_ex, "psi_ex_y", "psi_ex_z", ctx.adj_d_hz_dy, ctx.adj_d_hy_dz,
1139
+ magnetic_fields["Hz"], magnetic_fields["Hy"], solver.inv_dy_e, solver.inv_dz_e,
1140
+ solver.cpml_b_e_y, solver.cpml_c_e_y, solver.cpml_inv_kappa_e_y,
1141
+ solver.cpml_b_e_z, solver.cpml_c_e_z, solver.cpml_inv_kappa_e_z,
1142
+ solver.boundary_y_low_code, solver.boundary_y_high_code,
1143
+ solver.boundary_z_low_code, solver.boundary_z_high_code),
1144
+ (1, "Ey", eps_ey, "psi_ey_z", "psi_ey_x", ctx.adj_d_hx_dz, ctx.adj_d_hz_dx,
1145
+ magnetic_fields["Hx"], magnetic_fields["Hz"], solver.inv_dz_e, solver.inv_dx_e,
1146
+ solver.cpml_b_e_z, solver.cpml_c_e_z, solver.cpml_inv_kappa_e_z,
1147
+ solver.cpml_b_e_x, solver.cpml_c_e_x, solver.cpml_inv_kappa_e_x,
1148
+ solver.boundary_x_low_code, solver.boundary_x_high_code,
1149
+ solver.boundary_z_low_code, solver.boundary_z_high_code),
1150
+ (2, "Ez", eps_ez, "psi_ez_x", "psi_ez_y", ctx.adj_d_hy_dx, ctx.adj_d_hx_dy,
1151
+ magnetic_fields["Hy"], magnetic_fields["Hx"], solver.inv_dx_e, solver.inv_dy_e,
1152
+ solver.cpml_b_e_x, solver.cpml_c_e_x, solver.cpml_inv_kappa_e_x,
1153
+ solver.cpml_b_e_y, solver.cpml_c_e_y, solver.cpml_inv_kappa_e_y,
1154
+ solver.boundary_x_low_code, solver.boundary_x_high_code,
1155
+ solver.boundary_y_low_code, solver.boundary_y_high_code),
1156
+ )
1157
+ grad_eps = {"Ex": ctx.grad_eps_ex, "Ey": ctx.grad_eps_ey, "Ez": ctx.grad_eps_ez}
1158
+ for (component, name, eps, psi_pos, psi_neg, adj_d_pos, adj_d_neg,
1159
+ h_pos, h_neg, inv_pos, inv_neg, b_pos, c_pos, k_pos, b_neg, c_neg, k_neg,
1160
+ low_a, high_a, low_b, high_b) in configs:
1161
+ # Ey's reference wiring stores the two post-psi cotangents in the
1162
+ # historical x/z order; retain it for exact forward/checkpoint parity.
1163
+ post_pos, post_neg = (psi_neg, psi_pos) if name == "Ey" else (psi_pos, psi_neg)
1164
+ suffix = name[-1]
1165
+ _cuda_backend._reverse_electric_cpml_nonlinear(
1166
+ component=component, AdjPrev=pre[name], GradEps=grad_eps[name],
1167
+ GradChi2=grad_chi2[name], GradChi3=grad_chi3[name], GradTpa=grad_tpa[name],
1168
+ GFsq=g_fsq[name], AdjPsiPosPrev=pre[psi_pos], AdjPsiNegPrev=pre[psi_neg],
1169
+ AdjDPos=adj_d_pos, AdjDNeg=adj_d_neg, AdjPost=adjoint_state[name],
1170
+ AdjPsiPosPost=adjoint_state[post_pos], AdjPsiNegPost=adjoint_state[post_neg],
1171
+ EPrev=forward_state[name], ExternalDecay=getattr(solver, f"ce{suffix.lower()}_decay_external"),
1172
+ Eps=eps, Chi2=getattr(solver, f"nonlinear_chi2_{name}"),
1173
+ Chi3=getattr(solver, f"kerr_chi3_{name}"), Tpa=getattr(solver, f"tpa_sigma_{name}"),
1174
+ SigmaStatic=getattr(solver, f"sigma_e_{name}"), Fsq=fsq[name],
1175
+ Dt=float(solver.dt), Eps0=float(solver.eps0),
1176
+ PsiPos=forward_state[psi_pos], PsiNeg=forward_state[psi_neg],
1177
+ BPos=b_pos, CPos=c_pos, InvKappaPos=k_pos,
1178
+ BNeg=b_neg, CNeg=c_neg, InvKappaNeg=k_neg,
1179
+ HPosMid=h_pos, HNegMid=h_neg, InvPos=inv_pos, InvNeg=inv_neg,
1180
+ LowModeA=low_a, HighModeA=high_a, LowModeB=low_b, HighModeB=high_b,
1181
+ )
1182
+ _cuda_backend._accumulate_backward_diff_y(FieldGrad=adj_h["Hz"], DiffGrad=ctx.adj_d_hz_dy, invDy=solver.inv_dy_e)
1183
+ _cuda_backend._accumulate_backward_diff_z(FieldGrad=adj_h["Hy"], DiffGrad=ctx.adj_d_hy_dz, invDz=solver.inv_dz_e)
1184
+ _cuda_backend._accumulate_backward_diff_z(FieldGrad=adj_h["Hx"], DiffGrad=ctx.adj_d_hx_dz, invDz=solver.inv_dz_e)
1185
+ _cuda_backend._accumulate_backward_diff_x(FieldGrad=adj_h["Hz"], DiffGrad=ctx.adj_d_hz_dx, invDx=solver.inv_dx_e)
1186
+ _cuda_backend._accumulate_backward_diff_x(FieldGrad=adj_h["Hy"], DiffGrad=ctx.adj_d_hy_dx, invDx=solver.inv_dx_e)
1187
+ _cuda_backend._accumulate_backward_diff_y(FieldGrad=adj_h["Hx"], DiffGrad=ctx.adj_d_hx_dy, invDy=solver.inv_dy_e)
1188
+ _reverse_cpml_magnetic_phase_native(solver, ctx, adjoint_state)
1189
+ _cuda_backend._collocation_transpose(
1190
+ AdjEx=pre["Ex"], AdjEy=pre["Ey"], AdjEz=pre["Ez"],
1191
+ GEx=g_fsq["Ex"], GEy=g_fsq["Ey"], GEz=g_fsq["Ez"],
1192
+ Ex=forward_state["Ex"], Ey=forward_state["Ey"], Ez=forward_state["Ez"],
1193
+ )
1194
+ if dispersive_state:
1195
+ eps_by_field = {"Ex": eps_ex, "Ey": eps_ey, "Ez": eps_ez}
1196
+ grad_eps = {"Ex": ctx.grad_eps_ex, "Ey": ctx.grad_eps_ey, "Ez": ctx.grad_eps_ez}
1197
+ dt = float(solver.dt)
1198
+ for component, model, index, _tensor_names, entry in (
1199
+ _adjoint.iter_dispersive_state_specs(solver) or ()
1200
+ ):
1201
+ current = _adjoint.dispersive_state_name(component, model, index, "current")
1202
+ adj_current = torch.empty_like(forward_state[current])
1203
+ _cuda_backend._reverse_dispersive_correction(
1204
+ AdjCurrentCorrected=adj_current, GradEps=grad_eps[component],
1205
+ AdjCurrentPost=adjoint_state[current], AdjElectricPost=adjoint_state[component],
1206
+ Current=dispersive_state[current], Eps=eps_by_field[component], dt=dt,
1207
+ )
1208
+ if model == "debye":
1209
+ polarization = _adjoint.dispersive_state_name(component, model, index, "polarization")
1210
+ _cuda_backend._reverse_debye_current(
1211
+ AdjElectricPrev=pre[component], AdjPolarizationPrev=pre[polarization],
1212
+ AdjPolarizationPost=adjoint_state[polarization], AdjCurrentPost=adj_current,
1213
+ DebyeDrive=entry["drive"], decay=float(entry["decay"]), dt=dt,
1214
+ )
1215
+ elif model == "drude":
1216
+ _cuda_backend._reverse_drude_current(
1217
+ AdjElectricPrev=pre[component], AdjCurrentPrev=pre[current],
1218
+ AdjCurrentPost=adj_current, DrudeDrive=entry["drive"],
1219
+ decay=float(entry["decay"]),
1220
+ )
1221
+ else:
1222
+ polarization = _adjoint.dispersive_state_name(component, model, index, "polarization")
1223
+ _cuda_backend._reverse_lorentz_current(
1224
+ AdjElectricPrev=pre[component], AdjPolarizationPrev=pre[polarization],
1225
+ AdjCurrentPrev=pre[current], AdjPolarizationPost=adjoint_state[polarization],
1226
+ AdjCurrentPost=adj_current, LorentzDrive=entry["drive"],
1227
+ decay=float(entry["decay"]), restoring=float(entry["restoring"]), dt=dt,
1228
+ )
1229
+ result = _adjoint._ReverseStepResult(
1230
+ pre_step_adjoint={name: pre[name] for name in original_state_names}, grad_eps_ex=ctx.grad_eps_ex,
1231
+ grad_eps_ey=ctx.grad_eps_ey, grad_eps_ez=ctx.grad_eps_ez,
1232
+ backend=_NATIVE_REVERSE_LABELS[_ReverseBackend.GENERAL_NONLINEAR],
1233
+ magnetic_output_adjoint=adj_h,
1234
+ grad_chi2_ex=grad_chi2["Ex"], grad_chi2_ey=grad_chi2["Ey"], grad_chi2_ez=grad_chi2["Ez"],
1235
+ grad_chi3_ex=grad_chi3["Ex"], grad_chi3_ey=grad_chi3["Ey"], grad_chi3_ez=grad_chi3["Ez"],
1236
+ grad_tpa_ex=grad_tpa["Ex"], grad_tpa_ey=grad_tpa["Ey"], grad_tpa_ez=grad_tpa["Ez"],
1237
+ )
1238
+ return _adjoint._accumulate_source_term_gradients(
1239
+ result, solver=solver, adjoint_state=adjoint_state, time_value=time_value,
1240
+ eps_ex=eps_ex, eps_ey=eps_ey, eps_ez=eps_ez,
1241
+ resolved_source_terms=resolved_source_terms,
1242
+ )
1243
+
1244
+
1245
+ def _reverse_step_bloch_native_core(
1246
+ solver,
1247
+ forward_state,
1248
+ adjoint_state,
1249
+ *,
1250
+ time_value,
1251
+ eps_ex,
1252
+ eps_ey,
1253
+ eps_ez,
1254
+ resolved_source_terms,
1255
+ magnetic_fields=None,
1256
+ magnetic_cpml=False,
1257
+ magnetic_adjoint_additions=None,
1258
+ ):
1259
+ """Fused native complex-Bloch reverse math, without the source-term eps gradient.
1260
+
1261
+ Returns the ``_ReverseStepResult`` the analytic Bloch reference produces
1262
+ *before* ``_accumulate_source_term_gradients`` runs, with ``pre_step_adjoint``
1263
+ holding the twelve real/imag field adjoints. The public runner adds that
1264
+ accumulation; the Bloch+dispersive runner reuses this core as its complex base
1265
+ reverse and folds the electric ADE VJP on top before the single accumulation.
1266
+
1267
+ Mirrors ``reverse_step_bloch_native`` exactly, launching the fused
1268
+ complex Bloch reverse kernels in a fixed order. Every per-cell reverse
1269
+ computation (electric-adjoint -> mid-H adjoint with the boundary wrap phase,
1270
+ magnetic-adjoint -> pre-step E adjoint + eps gradient, and the real/imag
1271
+ magnetic decay pullback) runs inside the compiled kernels; Python only
1272
+ sequences the launches and carries the three per-axis Bloch phase pairs.
1273
+
1274
+ The fused kernels *assign* (not accumulate) into their outputs, so the launch
1275
+ order is load-bearing: the electric->H kernels must fully populate both real
1276
+ and imag mid-H adjoints before the magnetic->E kernels (which read them for the
1277
+ forward-diff fold) and before the decay pullback (which reads them for the
1278
+ pre-step H adjoint).
1279
+ """
1280
+ import torch
1281
+
1282
+ from . import core as _adjoint
1283
+ from .reverse_common import dynamic_electric_curls
1284
+
1285
+ # Mid-step complex H the forward electric update consumed (shared Torch replay).
1286
+ if magnetic_fields is None:
1287
+ magnetic_fields = _adjoint._forward_magnetic_fields_complex(
1288
+ solver,
1289
+ forward_state,
1290
+ time_value=time_value,
1291
+ resolved_source_terms=resolved_source_terms,
1292
+ )
1293
+
1294
+ # Dynamic electric curl coefficients (cast base curl by the eps leaf).
1295
+ ex_curl, ey_curl, ez_curl = dynamic_electric_curls(
1296
+ solver, eps_ex=eps_ex, eps_ey=eps_ey, eps_ez=eps_ez
1297
+ )
1298
+
1299
+ cos = solver.boundary_phase_cos
1300
+ sin = solver.boundary_phase_sin
1301
+ phase_cos_x, phase_sin_x = float(cos[0]), float(sin[0])
1302
+ phase_cos_y, phase_sin_y = float(cos[1]), float(sin[1])
1303
+ phase_cos_z, phase_sin_z = float(cos[2]), float(sin[2])
1304
+
1305
+ # Phase 1: electric adjoint -> mid-H adjoint (complex), with the boundary wrap
1306
+ # phase on the transposed backward differences.
1307
+ adj_hx_mid_r = torch.empty_like(forward_state["Hx"])
1308
+ adj_hx_mid_i = torch.empty_like(forward_state["Hx"])
1309
+ adj_hy_mid_r = torch.empty_like(forward_state["Hy"])
1310
+ adj_hy_mid_i = torch.empty_like(forward_state["Hy"])
1311
+ adj_hz_mid_r = torch.empty_like(forward_state["Hz"])
1312
+ adj_hz_mid_i = torch.empty_like(forward_state["Hz"])
1313
+ _cuda_backend._reverse_electric_hx_bloch(
1314
+ AdjHxMidReal=adj_hx_mid_r,
1315
+ AdjHxMidImag=adj_hx_mid_i,
1316
+ AdjHxPostReal=adjoint_state["Hx"],
1317
+ AdjHxPostImag=adjoint_state["Hx_imag"],
1318
+ AdjEyPostReal=adjoint_state["Ey"],
1319
+ AdjEyPostImag=adjoint_state["Ey_imag"],
1320
+ AdjEzPostReal=adjoint_state["Ez"],
1321
+ AdjEzPostImag=adjoint_state["Ez_imag"],
1322
+ EyCurl=ey_curl,
1323
+ EzCurl=ez_curl,
1324
+ phaseCosY=phase_cos_y,
1325
+ phaseSinY=phase_sin_y,
1326
+ phaseCosZ=phase_cos_z,
1327
+ phaseSinZ=phase_sin_z,
1328
+ invDy=solver.inv_dy_e,
1329
+ invDz=solver.inv_dz_e,
1330
+ )
1331
+ _cuda_backend._reverse_electric_hy_bloch(
1332
+ AdjHyMidReal=adj_hy_mid_r,
1333
+ AdjHyMidImag=adj_hy_mid_i,
1334
+ AdjHyPostReal=adjoint_state["Hy"],
1335
+ AdjHyPostImag=adjoint_state["Hy_imag"],
1336
+ AdjExPostReal=adjoint_state["Ex"],
1337
+ AdjExPostImag=adjoint_state["Ex_imag"],
1338
+ AdjEzPostReal=adjoint_state["Ez"],
1339
+ AdjEzPostImag=adjoint_state["Ez_imag"],
1340
+ ExCurl=ex_curl,
1341
+ EzCurl=ez_curl,
1342
+ phaseCosX=phase_cos_x,
1343
+ phaseSinX=phase_sin_x,
1344
+ phaseCosZ=phase_cos_z,
1345
+ phaseSinZ=phase_sin_z,
1346
+ invDx=solver.inv_dx_e,
1347
+ invDz=solver.inv_dz_e,
1348
+ )
1349
+ _cuda_backend._reverse_electric_hz_bloch(
1350
+ AdjHzMidReal=adj_hz_mid_r,
1351
+ AdjHzMidImag=adj_hz_mid_i,
1352
+ AdjHzPostReal=adjoint_state["Hz"],
1353
+ AdjHzPostImag=adjoint_state["Hz_imag"],
1354
+ AdjExPostReal=adjoint_state["Ex"],
1355
+ AdjExPostImag=adjoint_state["Ex_imag"],
1356
+ AdjEyPostReal=adjoint_state["Ey"],
1357
+ AdjEyPostImag=adjoint_state["Ey_imag"],
1358
+ ExCurl=ex_curl,
1359
+ EyCurl=ey_curl,
1360
+ phaseCosX=phase_cos_x,
1361
+ phaseSinX=phase_sin_x,
1362
+ phaseCosY=phase_cos_y,
1363
+ phaseSinY=phase_sin_y,
1364
+ invDx=solver.inv_dx_e,
1365
+ invDy=solver.inv_dy_e,
1366
+ )
1367
+ if magnetic_adjoint_additions:
1368
+ for name, target in (
1369
+ ("Hx", adj_hx_mid_r), ("Hy", adj_hy_mid_r), ("Hz", adj_hz_mid_r),
1370
+ ("Hx_imag", adj_hx_mid_i), ("Hy_imag", adj_hy_mid_i), ("Hz_imag", adj_hz_mid_i),
1371
+ ):
1372
+ addition = magnetic_adjoint_additions.get(name)
1373
+ if addition is not None:
1374
+ _cuda_backend._accumulate_in_place(dst=target, src=addition)
1375
+ if magnetic_cpml:
1376
+ zero_hx_r, zero_hx_i = torch.zeros_like(adj_hx_mid_r), torch.zeros_like(adj_hx_mid_i)
1377
+ zero_hy_r, zero_hy_i = torch.zeros_like(adj_hy_mid_r), torch.zeros_like(adj_hy_mid_i)
1378
+ zero_hz_r, zero_hz_i = torch.zeros_like(adj_hz_mid_r), torch.zeros_like(adj_hz_mid_i)
1379
+ else:
1380
+ zero_hx_r, zero_hx_i = adj_hx_mid_r, adj_hx_mid_i
1381
+ zero_hy_r, zero_hy_i = adj_hy_mid_r, adj_hy_mid_i
1382
+ zero_hz_r, zero_hz_i = adj_hz_mid_r, adj_hz_mid_i
1383
+
1384
+ # Phase 2: magnetic adjoint -> pre-step E adjoint (complex) + eps gradient. Each
1385
+ # kernel reconstructs its own curl(H) from the mid-H fields (for the eps
1386
+ # gradient) and assigns the complete pre-step E adjoint (electric decay pullback
1387
+ # + magnetic forward-diff fold of the mid-H adjoints).
1388
+ adj_ex_prev_r = torch.empty_like(forward_state["Ex"])
1389
+ adj_ex_prev_i = torch.empty_like(forward_state["Ex"])
1390
+ adj_ey_prev_r = torch.empty_like(forward_state["Ey"])
1391
+ adj_ey_prev_i = torch.empty_like(forward_state["Ey"])
1392
+ adj_ez_prev_r = torch.empty_like(forward_state["Ez"])
1393
+ adj_ez_prev_i = torch.empty_like(forward_state["Ez"])
1394
+ grad_eps_ex = torch.empty_like(eps_ex)
1395
+ grad_eps_ey = torch.empty_like(eps_ey)
1396
+ grad_eps_ez = torch.empty_like(eps_ez)
1397
+ _cuda_backend._reverse_magnetic_ex_bloch(
1398
+ AdjExPrevReal=adj_ex_prev_r,
1399
+ AdjExPrevImag=adj_ex_prev_i,
1400
+ GradEpsEx=grad_eps_ex,
1401
+ AdjExPostReal=adjoint_state["Ex"],
1402
+ AdjExPostImag=adjoint_state["Ex_imag"],
1403
+ AdjHyMidReal=zero_hy_r,
1404
+ AdjHyMidImag=zero_hy_i,
1405
+ AdjHzMidReal=zero_hz_r,
1406
+ AdjHzMidImag=zero_hz_i,
1407
+ ExDecay=solver.cex_decay,
1408
+ ExCurl=ex_curl,
1409
+ EpsEx=eps_ex,
1410
+ HyMidReal=magnetic_fields["Hy"],
1411
+ HyMidImag=magnetic_fields["Hy_imag"],
1412
+ HzMidReal=magnetic_fields["Hz"],
1413
+ HzMidImag=magnetic_fields["Hz_imag"],
1414
+ HyCurl=solver.chy_curl,
1415
+ HzCurl=solver.chz_curl,
1416
+ phaseCosY=phase_cos_y,
1417
+ phaseSinY=phase_sin_y,
1418
+ phaseCosZ=phase_cos_z,
1419
+ phaseSinZ=phase_sin_z,
1420
+ invDyE=solver.inv_dy_e,
1421
+ invDzE=solver.inv_dz_e,
1422
+ invDyH=solver.inv_dy_h,
1423
+ invDzH=solver.inv_dz_h,
1424
+ )
1425
+ _cuda_backend._reverse_magnetic_ey_bloch(
1426
+ AdjEyPrevReal=adj_ey_prev_r,
1427
+ AdjEyPrevImag=adj_ey_prev_i,
1428
+ GradEpsEy=grad_eps_ey,
1429
+ AdjEyPostReal=adjoint_state["Ey"],
1430
+ AdjEyPostImag=adjoint_state["Ey_imag"],
1431
+ AdjHxMidReal=zero_hx_r,
1432
+ AdjHxMidImag=zero_hx_i,
1433
+ AdjHzMidReal=zero_hz_r,
1434
+ AdjHzMidImag=zero_hz_i,
1435
+ EyDecay=solver.cey_decay,
1436
+ EyCurl=ey_curl,
1437
+ EpsEy=eps_ey,
1438
+ HxMidReal=magnetic_fields["Hx"],
1439
+ HxMidImag=magnetic_fields["Hx_imag"],
1440
+ HzMidReal=magnetic_fields["Hz"],
1441
+ HzMidImag=magnetic_fields["Hz_imag"],
1442
+ HxCurl=solver.chx_curl,
1443
+ HzCurl=solver.chz_curl,
1444
+ phaseCosX=phase_cos_x,
1445
+ phaseSinX=phase_sin_x,
1446
+ phaseCosZ=phase_cos_z,
1447
+ phaseSinZ=phase_sin_z,
1448
+ invDxE=solver.inv_dx_e,
1449
+ invDzE=solver.inv_dz_e,
1450
+ invDxH=solver.inv_dx_h,
1451
+ invDzH=solver.inv_dz_h,
1452
+ )
1453
+ _cuda_backend._reverse_magnetic_ez_bloch(
1454
+ AdjEzPrevReal=adj_ez_prev_r,
1455
+ AdjEzPrevImag=adj_ez_prev_i,
1456
+ GradEpsEz=grad_eps_ez,
1457
+ AdjEzPostReal=adjoint_state["Ez"],
1458
+ AdjEzPostImag=adjoint_state["Ez_imag"],
1459
+ AdjHxMidReal=zero_hx_r,
1460
+ AdjHxMidImag=zero_hx_i,
1461
+ AdjHyMidReal=zero_hy_r,
1462
+ AdjHyMidImag=zero_hy_i,
1463
+ EzDecay=solver.cez_decay,
1464
+ EzCurl=ez_curl,
1465
+ EpsEz=eps_ez,
1466
+ HxMidReal=magnetic_fields["Hx"],
1467
+ HxMidImag=magnetic_fields["Hx_imag"],
1468
+ HyMidReal=magnetic_fields["Hy"],
1469
+ HyMidImag=magnetic_fields["Hy_imag"],
1470
+ HxCurl=solver.chx_curl,
1471
+ HyCurl=solver.chy_curl,
1472
+ phaseCosX=phase_cos_x,
1473
+ phaseSinX=phase_sin_x,
1474
+ phaseCosY=phase_cos_y,
1475
+ phaseSinY=phase_sin_y,
1476
+ invDxE=solver.inv_dx_e,
1477
+ invDyE=solver.inv_dy_e,
1478
+ invDxH=solver.inv_dx_h,
1479
+ invDyH=solver.inv_dy_h,
1480
+ )
1481
+
1482
+ pre_by_name = {
1483
+ "Ex": adj_ex_prev_r,
1484
+ "Ey": adj_ey_prev_r,
1485
+ "Ez": adj_ez_prev_r,
1486
+ "Ex_imag": adj_ex_prev_i,
1487
+ "Ey_imag": adj_ey_prev_i,
1488
+ "Ez_imag": adj_ez_prev_i,
1489
+ }
1490
+ if magnetic_cpml:
1491
+ from types import SimpleNamespace
1492
+ for name in ("Hx", "Hy", "Hz", "Hx_imag", "Hy_imag", "Hz_imag"):
1493
+ pre_by_name[name] = torch.empty_like(forward_state[name])
1494
+ for name in _adjoint.checkpoint_schema(solver).cpml_state_names:
1495
+ pre_by_name[name] = torch.zeros_like(forward_state[name])
1496
+ def make_ctx(suffix, adj_h_values):
1497
+ return SimpleNamespace(
1498
+ pre_step_adjoint=pre_by_name,
1499
+ magnetic_output_adjoint=adj_h_values,
1500
+ adj_d_ez_dy=torch.empty_like(forward_state["Hx"]),
1501
+ adj_d_ey_dz=torch.empty_like(forward_state["Hx"]),
1502
+ adj_d_ex_dz=torch.empty_like(forward_state["Hy"]),
1503
+ adj_d_ez_dx=torch.empty_like(forward_state["Hy"]),
1504
+ adj_d_ey_dx=torch.empty_like(forward_state["Hz"]),
1505
+ adj_d_ex_dy=torch.empty_like(forward_state["Hz"]),
1506
+ )
1507
+ real_ctx = make_ctx("", {"Hx": adj_hx_mid_r, "Hy": adj_hy_mid_r, "Hz": adj_hz_mid_r})
1508
+ imag_ctx = make_ctx("_imag", {
1509
+ "Hx_imag": adj_hx_mid_i, "Hy_imag": adj_hy_mid_i, "Hz_imag": adj_hz_mid_i
1510
+ })
1511
+ _reverse_cpml_magnetic_phase_native(solver, real_ctx, adjoint_state)
1512
+ _reverse_cpml_magnetic_phase_native(solver, imag_ctx, adjoint_state, suffix="_imag")
1513
+ else:
1514
+ for component, adj_r, adj_i in (
1515
+ ("Hx", adj_hx_mid_r, adj_hx_mid_i),
1516
+ ("Hy", adj_hy_mid_r, adj_hy_mid_i),
1517
+ ("Hz", adj_hz_mid_r, adj_hz_mid_i),
1518
+ ):
1519
+ prev_r, prev_i = torch.empty_like(adj_r), torch.empty_like(adj_i)
1520
+ getattr(_cuda_backend, f"_reverse_magnetic_{component.lower()}_decay")(
1521
+ **{f"Adj{component}Prev": prev_r, f"Adj{component}Mid": adj_r,
1522
+ f"{component}Decay": getattr(solver, f"c{component.lower()}_decay")}
1523
+ )
1524
+ getattr(_cuda_backend, f"_reverse_magnetic_{component.lower()}_decay")(
1525
+ **{f"Adj{component}Prev": prev_i, f"Adj{component}Mid": adj_i,
1526
+ f"{component}Decay": getattr(solver, f"c{component.lower()}_decay")}
1527
+ )
1528
+ pre_by_name[component] = prev_r
1529
+ pre_by_name[component + "_imag"] = prev_i
1530
+ pre_by_name = {name: pre_by_name[name] for name in forward_state if name in pre_by_name}
1531
+ return _adjoint._ReverseStepResult(
1532
+ pre_step_adjoint=pre_by_name,
1533
+ grad_eps_ex=grad_eps_ex,
1534
+ grad_eps_ey=grad_eps_ey,
1535
+ grad_eps_ez=grad_eps_ez,
1536
+ backend=_NATIVE_REVERSE_LABELS[_ReverseBackend.BLOCH],
1537
+ magnetic_output_adjoint={
1538
+ "Hx": adj_hx_mid_r, "Hy": adj_hy_mid_r, "Hz": adj_hz_mid_r,
1539
+ "Hx_imag": adj_hx_mid_i, "Hy_imag": adj_hy_mid_i, "Hz_imag": adj_hz_mid_i,
1540
+ },
1541
+ )
1542
+
1543
+
1544
+ def _reverse_step_mixed_bloch_cpml_native(
1545
+ solver, forward_state, adjoint_state, *, time_value, eps_ex, eps_ey, eps_ez,
1546
+ resolved_source_terms, profiler,
1547
+ ):
1548
+ """Native reverse of one-PML-axis/two-Bloch-axis complex CPML."""
1549
+ import dataclasses
1550
+ import torch
1551
+
1552
+ from . import core as _adjoint
1553
+ from .reverse_common import dynamic_electric_curls
1554
+ from ..runtime.stepping import _bloch_cpml_pml_axis
1555
+
1556
+ pml_axis_name = _bloch_cpml_pml_axis(solver)
1557
+ if pml_axis_name is None:
1558
+ raise RuntimeError("Native mixed Bloch/CPML reverse requires one PML axis and two Bloch axes.")
1559
+ captured = []
1560
+ _adjoint._step_state(
1561
+ solver, forward_state, time_value=time_value,
1562
+ eps_ex=eps_ex, eps_ey=eps_ey, eps_ez=eps_ez,
1563
+ capture_magnetic=captured,
1564
+ )
1565
+ magnetic_fields = {name: tensor.contiguous() for name, tensor in captured[0].items()}
1566
+ dispersive_state = _adjoint._advance_dispersive_state(solver, forward_state)
1567
+ magnetic_dispersive_state = _adjoint._advance_magnetic_dispersive_state(
1568
+ solver, forward_state
1569
+ )
1570
+ curls = dict(zip(("Ex", "Ey", "Ez"), dynamic_electric_curls(
1571
+ solver, eps_ex=eps_ex, eps_ey=eps_ey, eps_ez=eps_ez
1572
+ )))
1573
+ additions = {
1574
+ name: torch.zeros_like(forward_state[name])
1575
+ for name in ("Hx", "Hy", "Hz", "Hx_imag", "Hy_imag", "Hz_imag")
1576
+ }
1577
+ electric_psi_names = (
1578
+ "psi_ex_y", "psi_ex_z", "psi_ey_x", "psi_ey_z", "psi_ez_x", "psi_ez_y"
1579
+ )
1580
+ psi_pre = {
1581
+ name + suffix: adjoint_state[name + suffix].detach().clone()
1582
+ for suffix in ("", "_imag") for name in electric_psi_names
1583
+ }
1584
+ axis = {"x": 0, "y": 1, "z": 2}[pml_axis_name]
1585
+ # component, psi, derivative target H, sign, tangent axis
1586
+ correction_specs = {
1587
+ "z": (("Ex", "psi_ex_z", "Hy", -1.0, 1), ("Ey", "psi_ey_z", "Hx", 1.0, 0)),
1588
+ "y": (("Ex", "psi_ex_y", "Hz", 1.0, 2), ("Ez", "psi_ez_y", "Hx", -1.0, 0)),
1589
+ "x": (("Ey", "psi_ey_x", "Hz", -1.0, 2), ("Ez", "psi_ez_x", "Hy", 1.0, 1)),
1590
+ }[pml_axis_name]
1591
+ inv_delta = (solver.inv_dx_e, solver.inv_dy_e, solver.inv_dz_e)[axis]
1592
+ b = (solver.cpml_b_e_x, solver.cpml_b_e_y, solver.cpml_b_e_z)[axis]
1593
+ c = (solver.cpml_c_e_x, solver.cpml_c_e_y, solver.cpml_c_e_z)[axis]
1594
+ inv_kappa = (
1595
+ solver.cpml_inv_kappa_e_x, solver.cpml_inv_kappa_e_y, solver.cpml_inv_kappa_e_z
1596
+ )[axis]
1597
+ for suffix in ("", "_imag"):
1598
+ for component, psi_name, h_name, sign, tangent_axis in correction_specs:
1599
+ field_key, psi_key, h_key = component + suffix, psi_name + suffix, h_name + suffix
1600
+ adj_derivative = torch.empty_like(forward_state[field_key])
1601
+ _cuda_backend._reverse_cpml_correction(
1602
+ AdjPsiPrev=psi_pre[psi_key], AdjDerivative=adj_derivative,
1603
+ AdjField=adjoint_state[field_key], AdjPsiPost=adjoint_state[psi_key],
1604
+ Curl=curls[component], B=b, C=c, InvKappa=inv_kappa,
1605
+ NormalAxis=axis, TangentAxis=tangent_axis,
1606
+ TangentLowMode=getattr(solver, f"boundary_{'xyz'[tangent_axis]}_low_code"),
1607
+ TangentHighMode=getattr(solver, f"boundary_{'xyz'[tangent_axis]}_high_code"),
1608
+ Sign=sign,
1609
+ )
1610
+ accumulate = (
1611
+ _cuda_backend._accumulate_backward_diff_x,
1612
+ _cuda_backend._accumulate_backward_diff_y,
1613
+ _cuda_backend._accumulate_backward_diff_z,
1614
+ )[axis]
1615
+ accumulate(**{
1616
+ "FieldGrad": additions[h_key], "DiffGrad": adj_derivative,
1617
+ ("invDx", "invDy", "invDz")[axis]: inv_delta,
1618
+ })
1619
+ result = _reverse_step_bloch_native_core(
1620
+ solver, forward_state, adjoint_state, time_value=time_value,
1621
+ eps_ex=eps_ex, eps_ey=eps_ey, eps_ez=eps_ez,
1622
+ resolved_source_terms=resolved_source_terms,
1623
+ magnetic_fields=magnetic_fields, magnetic_cpml=True,
1624
+ magnetic_adjoint_additions=additions,
1625
+ )
1626
+ pre = dict(result.pre_step_adjoint)
1627
+ pre.update(psi_pre)
1628
+ for name in _adjoint.checkpoint_schema(solver).dispersive_state_names:
1629
+ pre[name] = torch.zeros_like(forward_state[name])
1630
+ for name in _adjoint.checkpoint_schema(solver).magnetic_dispersive_state_names:
1631
+ pre[name] = torch.zeros_like(forward_state[name])
1632
+ if dispersive_state:
1633
+ grad_eps = {"Ex": result.grad_eps_ex, "Ey": result.grad_eps_ey, "Ez": result.grad_eps_ez}
1634
+ eps_by_field = {"Ex": eps_ex, "Ey": eps_ey, "Ez": eps_ez}
1635
+ dt = float(solver.dt)
1636
+ for suffix in ("", "_imag"):
1637
+ for component, model, index, _tensor_names, entry in (
1638
+ _adjoint.iter_dispersive_state_specs(solver) or ()
1639
+ ):
1640
+ current = _adjoint.dispersive_state_name(component, model, index, "current") + suffix
1641
+ adj_current = torch.empty_like(forward_state[current])
1642
+ _cuda_backend._reverse_dispersive_correction(
1643
+ AdjCurrentCorrected=adj_current, GradEps=grad_eps[component],
1644
+ AdjCurrentPost=adjoint_state[current],
1645
+ AdjElectricPost=adjoint_state[component + suffix],
1646
+ Current=dispersive_state[current], Eps=eps_by_field[component], dt=dt,
1647
+ )
1648
+ if model == "debye":
1649
+ polarization = _adjoint.dispersive_state_name(
1650
+ component, model, index, "polarization"
1651
+ ) + suffix
1652
+ _cuda_backend._reverse_debye_current(
1653
+ AdjElectricPrev=pre[component + suffix],
1654
+ AdjPolarizationPrev=pre[polarization],
1655
+ AdjPolarizationPost=adjoint_state[polarization],
1656
+ AdjCurrentPost=adj_current, DebyeDrive=entry["drive"],
1657
+ decay=float(entry["decay"]), dt=dt,
1658
+ )
1659
+ elif model == "drude":
1660
+ _cuda_backend._reverse_drude_current(
1661
+ AdjElectricPrev=pre[component + suffix], AdjCurrentPrev=pre[current],
1662
+ AdjCurrentPost=adj_current, DrudeDrive=entry["drive"],
1663
+ decay=float(entry["decay"]),
1664
+ )
1665
+ else:
1666
+ polarization = _adjoint.dispersive_state_name(
1667
+ component, model, index, "polarization"
1668
+ ) + suffix
1669
+ _cuda_backend._reverse_lorentz_current(
1670
+ AdjElectricPrev=pre[component + suffix],
1671
+ AdjPolarizationPrev=pre[polarization], AdjCurrentPrev=pre[current],
1672
+ AdjPolarizationPost=adjoint_state[polarization], AdjCurrentPost=adj_current,
1673
+ LorentzDrive=entry["drive"], decay=float(entry["decay"]),
1674
+ restoring=float(entry["restoring"]), dt=dt,
1675
+ )
1676
+ _apply_magnetic_ade_reverse_native(
1677
+ solver, forward_state, adjoint_state, magnetic_dispersive_state,
1678
+ result.magnetic_output_adjoint, pre,
1679
+ )
1680
+ pre = {name: pre[name] for name in forward_state}
1681
+ backend = (
1682
+ _ReverseBackend.GRATING_TFSF if getattr(solver, "tfsf_enabled", False)
1683
+ else _ReverseBackend.MIXED_BLOCH_CPML
1684
+ )
1685
+ result = dataclasses.replace(
1686
+ result, pre_step_adjoint=pre, backend=_NATIVE_REVERSE_LABELS[backend]
1687
+ )
1688
+ return _adjoint._accumulate_source_term_gradients(
1689
+ result, solver=solver, adjoint_state=adjoint_state, time_value=time_value,
1690
+ eps_ex=eps_ex, eps_ey=eps_ey, eps_ez=eps_ez,
1691
+ resolved_source_terms=resolved_source_terms,
1692
+ )
1693
+
1694
+
1695
+ def _reverse_step_bloch_native(
1696
+ solver,
1697
+ forward_state,
1698
+ adjoint_state,
1699
+ *,
1700
+ time_value,
1701
+ eps_ex,
1702
+ eps_ey,
1703
+ eps_ez,
1704
+ resolved_source_terms,
1705
+ profiler,
1706
+ ):
1707
+ """Fully native reverse step for the complex split-field Bloch configuration."""
1708
+ from . import core as _adjoint
1709
+
1710
+ step_result = _reverse_step_bloch_native_core(
1711
+ solver,
1712
+ forward_state,
1713
+ adjoint_state,
1714
+ time_value=time_value,
1715
+ eps_ex=eps_ex,
1716
+ eps_ey=eps_ey,
1717
+ eps_ez=eps_ez,
1718
+ resolved_source_terms=resolved_source_terms,
1719
+ )
1720
+ # The native runner owns the full reverse-step contract, including the analytic
1721
+ # source-term eps-gradient accumulation the reference path applies via ``finish``.
1722
+ return _adjoint._accumulate_source_term_gradients(
1723
+ step_result,
1724
+ solver=solver,
1725
+ adjoint_state=adjoint_state,
1726
+ time_value=time_value,
1727
+ eps_ex=eps_ex,
1728
+ eps_ey=eps_ey,
1729
+ eps_ez=eps_ez,
1730
+ resolved_source_terms=resolved_source_terms,
1731
+ )
1732
+
1733
+
1734
+ def _apply_magnetic_ade_reverse_native(
1735
+ solver, forward_state, adjoint_state, magnetic_state, magnetic_output_adjoint,
1736
+ pre_step_adjoint,
1737
+ ):
1738
+ """Apply native magnetic-current correction and ADE-state pullbacks."""
1739
+ import torch
1740
+
1741
+ from . import core as _adjoint
1742
+
1743
+ if not magnetic_state:
1744
+ return
1745
+ mu = {"Hx": solver.mu_Hx, "Hy": solver.mu_Hy, "Hz": solver.mu_Hz}
1746
+ suffixes = ("", "_imag") if getattr(solver, "complex_fields_enabled", False) else ("",)
1747
+ for suffix in suffixes:
1748
+ for component, model, index, _tensor_names, entry in (
1749
+ _adjoint.iter_magnetic_dispersive_state_specs(solver) or ()
1750
+ ):
1751
+ current = _adjoint.dispersive_state_name(component, model, index, "current") + suffix
1752
+ adj_current = torch.empty_like(forward_state[current])
1753
+ discarded_grad_mu = torch.zeros_like(mu[component])
1754
+ _cuda_backend._reverse_dispersive_correction(
1755
+ AdjCurrentCorrected=adj_current, GradEps=discarded_grad_mu,
1756
+ AdjCurrentPost=adjoint_state[current],
1757
+ AdjElectricPost=magnetic_output_adjoint[component + suffix],
1758
+ Current=magnetic_state[current], Eps=mu[component], dt=float(solver.dt),
1759
+ )
1760
+ if model == "debye":
1761
+ polarization = _adjoint.dispersive_state_name(
1762
+ component, model, index, "polarization"
1763
+ ) + suffix
1764
+ _cuda_backend._reverse_debye_current(
1765
+ AdjElectricPrev=pre_step_adjoint[component + suffix],
1766
+ AdjPolarizationPrev=pre_step_adjoint[polarization],
1767
+ AdjPolarizationPost=adjoint_state[polarization], AdjCurrentPost=adj_current,
1768
+ DebyeDrive=entry["drive"], decay=float(entry["decay"]), dt=float(solver.dt),
1769
+ )
1770
+ elif model == "drude":
1771
+ _cuda_backend._reverse_drude_current(
1772
+ AdjElectricPrev=pre_step_adjoint[component + suffix],
1773
+ AdjCurrentPrev=pre_step_adjoint[current], AdjCurrentPost=adj_current,
1774
+ DrudeDrive=entry["drive"], decay=float(entry["decay"]),
1775
+ )
1776
+ else:
1777
+ polarization = _adjoint.dispersive_state_name(
1778
+ component, model, index, "polarization"
1779
+ ) + suffix
1780
+ _cuda_backend._reverse_lorentz_current(
1781
+ AdjElectricPrev=pre_step_adjoint[component + suffix],
1782
+ AdjPolarizationPrev=pre_step_adjoint[polarization],
1783
+ AdjCurrentPrev=pre_step_adjoint[current],
1784
+ AdjPolarizationPost=adjoint_state[polarization], AdjCurrentPost=adj_current,
1785
+ LorentzDrive=entry["drive"], decay=float(entry["decay"]),
1786
+ restoring=float(entry["restoring"]), dt=float(solver.dt),
1787
+ )
1788
+
1789
+
1790
+ def _reverse_step_bloch_dispersive_native(
1791
+ solver,
1792
+ forward_state,
1793
+ adjoint_state,
1794
+ *,
1795
+ time_value,
1796
+ eps_ex,
1797
+ eps_ey,
1798
+ eps_ez,
1799
+ resolved_source_terms,
1800
+ profiler,
1801
+ ):
1802
+ """Fully native reverse step for a Bloch (complex-field) + electric-dispersive medium.
1803
+
1804
+ Mirrors :func:`reverse_step_bloch_dispersive_native`. The complex
1805
+ Bloch base reverse produces the twelve real/imag field adjoints and the base
1806
+ eps gradient; on top of it the electric ADE VJP runs on the real and imaginary
1807
+ halves independently, reusing the same fused dispersive-correction and
1808
+ ADE-current kernels the pure-real dispersive runner uses. Both halves share the
1809
+ real eps, so both correction launches *accumulate* into the same eps gradient;
1810
+ both ADE-current launches *accumulate* the field-drive term onto their half's
1811
+ pre-step E adjoint and *assign* their half's pre-step polarization/current
1812
+ adjoint. Launch order is load-bearing: the Bloch base assigns the pre-step E
1813
+ adjoint first, then per half the correction produces the corrected current
1814
+ adjoint the ADE-current kernels read, and those accumulate onto the base E
1815
+ adjoint. The single source-term accumulation runs once at the end.
1816
+ """
1817
+ import torch
1818
+
1819
+ from . import core as _adjoint
1820
+
1821
+ # Replay the post-update electric dispersive currents (Torch), real + imag.
1822
+ dispersive_state = _adjoint._advance_dispersive_state(solver, forward_state)
1823
+ magnetic_dispersive_state = _adjoint._advance_magnetic_dispersive_state(
1824
+ solver, forward_state
1825
+ )
1826
+ magnetic_fields = None
1827
+ if magnetic_dispersive_state:
1828
+ magnetic_fields = _adjoint._apply_magnetic_dispersive_corrections(
1829
+ solver,
1830
+ _adjoint._forward_magnetic_fields_complex(
1831
+ solver, forward_state, time_value=time_value,
1832
+ resolved_source_terms=resolved_source_terms,
1833
+ ),
1834
+ magnetic_dispersive_state,
1835
+ )
1836
+
1837
+ # Native complex Bloch base reverse (no source-term grads yet).
1838
+ base_result = _reverse_step_bloch_native_core(
1839
+ solver,
1840
+ forward_state,
1841
+ adjoint_state,
1842
+ time_value=time_value,
1843
+ eps_ex=eps_ex,
1844
+ eps_ey=eps_ey,
1845
+ eps_ez=eps_ez,
1846
+ resolved_source_terms=resolved_source_terms,
1847
+ magnetic_fields=magnetic_fields,
1848
+ )
1849
+
1850
+ pre_step_adjoint = dict(base_result.pre_step_adjoint)
1851
+ for name in _adjoint.checkpoint_schema(solver).dispersive_state_names:
1852
+ if name not in pre_step_adjoint:
1853
+ pre_step_adjoint[name] = torch.zeros_like(forward_state[name])
1854
+ for name in _adjoint.checkpoint_schema(solver).magnetic_dispersive_state_names:
1855
+ pre_step_adjoint[name] = torch.zeros_like(forward_state[name])
1856
+
1857
+ grad_eps_by_field = {
1858
+ "Ex": base_result.grad_eps_ex,
1859
+ "Ey": base_result.grad_eps_ey,
1860
+ "Ez": base_result.grad_eps_ez,
1861
+ }
1862
+ eps_by_field = {"Ex": eps_ex, "Ey": eps_ey, "Ez": eps_ez}
1863
+ dt = float(solver.dt)
1864
+
1865
+ # Per half (real, imaginary), per pole: correction VJP then ADE-state VJP. The
1866
+ # coefficients are identical across halves; only the field/state suffix differs.
1867
+ for suffix in ("", "_imag"):
1868
+ for component_name, model_name, index, _tensor_names, entry in _adjoint.iter_dispersive_state_specs(solver) or ():
1869
+ current_name = _adjoint.dispersive_state_name(component_name, model_name, index, "current") + suffix
1870
+ adj_current_corrected = torch.empty_like(forward_state[current_name])
1871
+ _cuda_backend._reverse_dispersive_correction(
1872
+ AdjCurrentCorrected=adj_current_corrected,
1873
+ GradEps=grad_eps_by_field[component_name],
1874
+ AdjCurrentPost=adjoint_state[current_name],
1875
+ AdjElectricPost=adjoint_state[component_name + suffix],
1876
+ Current=dispersive_state[current_name],
1877
+ Eps=eps_by_field[component_name],
1878
+ dt=dt,
1879
+ )
1880
+ if model_name == "debye":
1881
+ polarization_name = _adjoint.dispersive_state_name(component_name, model_name, index, "polarization") + suffix
1882
+ _cuda_backend._reverse_debye_current(
1883
+ AdjElectricPrev=pre_step_adjoint[component_name + suffix],
1884
+ AdjPolarizationPrev=pre_step_adjoint[polarization_name],
1885
+ AdjPolarizationPost=adjoint_state[polarization_name],
1886
+ AdjCurrentPost=adj_current_corrected,
1887
+ DebyeDrive=entry["drive"],
1888
+ decay=float(entry["decay"]),
1889
+ dt=dt,
1890
+ )
1891
+ elif model_name == "drude":
1892
+ _cuda_backend._reverse_drude_current(
1893
+ AdjElectricPrev=pre_step_adjoint[component_name + suffix],
1894
+ AdjCurrentPrev=pre_step_adjoint[current_name],
1895
+ AdjCurrentPost=adj_current_corrected,
1896
+ DrudeDrive=entry["drive"],
1897
+ decay=float(entry["decay"]),
1898
+ )
1899
+ else:
1900
+ polarization_name = _adjoint.dispersive_state_name(component_name, model_name, index, "polarization") + suffix
1901
+ _cuda_backend._reverse_lorentz_current(
1902
+ AdjElectricPrev=pre_step_adjoint[component_name + suffix],
1903
+ AdjPolarizationPrev=pre_step_adjoint[polarization_name],
1904
+ AdjCurrentPrev=pre_step_adjoint[current_name],
1905
+ AdjPolarizationPost=adjoint_state[polarization_name],
1906
+ AdjCurrentPost=adj_current_corrected,
1907
+ LorentzDrive=entry["drive"],
1908
+ decay=float(entry["decay"]),
1909
+ restoring=float(entry["restoring"]),
1910
+ dt=dt,
1911
+ )
1912
+
1913
+ _apply_magnetic_ade_reverse_native(
1914
+ solver, forward_state, adjoint_state, magnetic_dispersive_state,
1915
+ base_result.magnetic_output_adjoint, pre_step_adjoint,
1916
+ )
1917
+ step_result = _adjoint._ReverseStepResult(
1918
+ pre_step_adjoint=pre_step_adjoint,
1919
+ grad_eps_ex=grad_eps_by_field["Ex"],
1920
+ grad_eps_ey=grad_eps_by_field["Ey"],
1921
+ grad_eps_ez=grad_eps_by_field["Ez"],
1922
+ backend=_NATIVE_REVERSE_LABELS[_ReverseBackend.BLOCH_DISPERSIVE],
1923
+ source_adjoint_state=dict(adjoint_state),
1924
+ )
1925
+ return _adjoint._accumulate_source_term_gradients(
1926
+ step_result,
1927
+ solver=solver,
1928
+ adjoint_state=adjoint_state,
1929
+ time_value=time_value,
1930
+ eps_ex=eps_ex,
1931
+ eps_ey=eps_ey,
1932
+ eps_ez=eps_ez,
1933
+ resolved_source_terms=resolved_source_terms,
1934
+ )
1935
+
1936
+
1937
+ def _reverse_step_dispersive_native(
1938
+ solver,
1939
+ forward_state,
1940
+ adjoint_state,
1941
+ *,
1942
+ time_value,
1943
+ eps_ex,
1944
+ eps_ey,
1945
+ eps_ez,
1946
+ resolved_source_terms,
1947
+ profiler,
1948
+ ):
1949
+ """Fully native reverse step for the electric-dispersive (ADE) configuration.
1950
+
1951
+ Mirrors ``reverse_step_dispersive_native`` for the electric-only
1952
+ dispersive case. The reverse *math* runs entirely in native kernels:
1953
+
1954
+ 1. The mid-step post-update dispersive currents are reconstructed with the
1955
+ shared Torch ADE replay (``_advance_dispersive_state``); this is the
1956
+ once-per-step replay tracked under the native-replay work, same bar as the
1957
+ standard/CPML runners' mid-H replay.
1958
+ 2. The native standard or CPML *core* reverse produces the pre-step E/H (and
1959
+ psi) adjoint plus the base eps gradient, consuming the post-step field
1960
+ adjoints directly (the reference's ``electric_source_adjoint`` is an
1961
+ identity clone of those, so no separate copy is needed).
1962
+ 3. Per electric pole the fused ``reverseDispersiveCorrection3D`` kernel folds
1963
+ the 1/eps coupling: it assigns the corrected post-step current adjoint and
1964
+ accumulates the ``dt * J * adj_E / eps^2`` eps gradient onto the base eps
1965
+ gradient in place.
1966
+ 4. The fused ``reverseDebye/Drude/LorentzCurrent3D`` kernels pull the corrected
1967
+ current (and raw polarization) adjoint back through the ADE update, adding
1968
+ the field-drive term into the pre-step E adjoint and assigning the pre-step
1969
+ polarization/current adjoints.
1970
+
1971
+ Launch order is load-bearing: the base reverse assigns the pre-step E adjoint
1972
+ and base eps gradient first; the correction kernel then accumulates the eps
1973
+ gradient and produces the corrected current adjoint the ADE-state kernels read;
1974
+ the ADE-state kernels accumulate the field-drive term on top of the base
1975
+ pre-step E adjoint. The single source-term eps-gradient accumulation runs once
1976
+ at the end, matching the reference ``finish`` step.
1977
+ """
1978
+ import torch
1979
+
1980
+ from . import core as _adjoint
1981
+
1982
+ # Step 1: replay the post-update electric and magnetic ADE currents. Replay
1983
+ # is orchestration; all pullbacks below are compiled CUDA kernels.
1984
+ dispersive_state = _adjoint._advance_dispersive_state(solver, forward_state)
1985
+ magnetic_dispersive_state = _adjoint._advance_magnetic_dispersive_state(
1986
+ solver, forward_state
1987
+ )
1988
+ magnetic_fields = None
1989
+ if magnetic_dispersive_state:
1990
+ magnetic_fields = _adjoint._apply_magnetic_dispersive_corrections(
1991
+ solver,
1992
+ _adjoint._forward_magnetic_fields(
1993
+ solver,
1994
+ forward_state,
1995
+ time_value=time_value,
1996
+ resolved_source_terms=resolved_source_terms,
1997
+ ),
1998
+ magnetic_dispersive_state,
1999
+ )
2000
+
2001
+ # Step 2: native base reverse (standard or CPML core), no source-term grads.
2002
+ if getattr(solver, "uses_cpml", False):
2003
+ base_result = _reverse_step_cpml_native_core(
2004
+ solver,
2005
+ forward_state,
2006
+ adjoint_state,
2007
+ time_value=time_value,
2008
+ eps_ex=eps_ex,
2009
+ eps_ey=eps_ey,
2010
+ eps_ez=eps_ez,
2011
+ resolved_source_terms=resolved_source_terms,
2012
+ magnetic_fields=magnetic_fields,
2013
+ )
2014
+ else:
2015
+ base_result = _reverse_step_standard_native_core(
2016
+ solver,
2017
+ forward_state,
2018
+ adjoint_state,
2019
+ time_value=time_value,
2020
+ eps_ex=eps_ex,
2021
+ eps_ey=eps_ey,
2022
+ eps_ez=eps_ez,
2023
+ resolved_source_terms=resolved_source_terms,
2024
+ magnetic_fields=magnetic_fields,
2025
+ )
2026
+
2027
+ pre_step_adjoint = dict(base_result.pre_step_adjoint)
2028
+ for name in _adjoint.checkpoint_schema(solver).dispersive_state_names:
2029
+ if name not in pre_step_adjoint:
2030
+ pre_step_adjoint[name] = torch.zeros_like(forward_state[name])
2031
+ for name in _adjoint.checkpoint_schema(solver).magnetic_dispersive_state_names:
2032
+ pre_step_adjoint[name] = torch.zeros_like(forward_state[name])
2033
+
2034
+ grad_eps_by_field = {
2035
+ "Ex": base_result.grad_eps_ex,
2036
+ "Ey": base_result.grad_eps_ey,
2037
+ "Ez": base_result.grad_eps_ez,
2038
+ }
2039
+ eps_by_field = {"Ex": eps_ex, "Ey": eps_ey, "Ez": eps_ez}
2040
+ dt = float(solver.dt)
2041
+
2042
+ # Steps 3 + 4: per pole, correction VJP then ADE-state VJP.
2043
+ for component_name, model_name, index, _tensor_names, entry in _adjoint.iter_dispersive_state_specs(solver) or ():
2044
+ current_name = _adjoint.dispersive_state_name(component_name, model_name, index, "current")
2045
+ adj_current_corrected = torch.empty_like(forward_state[current_name])
2046
+ _cuda_backend._reverse_dispersive_correction(
2047
+ AdjCurrentCorrected=adj_current_corrected,
2048
+ GradEps=grad_eps_by_field[component_name],
2049
+ AdjCurrentPost=adjoint_state[current_name],
2050
+ AdjElectricPost=adjoint_state[component_name],
2051
+ Current=dispersive_state[current_name],
2052
+ Eps=eps_by_field[component_name],
2053
+ dt=dt,
2054
+ )
2055
+ if model_name == "debye":
2056
+ polarization_name = _adjoint.dispersive_state_name(component_name, model_name, index, "polarization")
2057
+ _cuda_backend._reverse_debye_current(
2058
+ AdjElectricPrev=pre_step_adjoint[component_name],
2059
+ AdjPolarizationPrev=pre_step_adjoint[polarization_name],
2060
+ AdjPolarizationPost=adjoint_state[polarization_name],
2061
+ AdjCurrentPost=adj_current_corrected,
2062
+ DebyeDrive=entry["drive"],
2063
+ decay=float(entry["decay"]),
2064
+ dt=dt,
2065
+ )
2066
+ elif model_name == "drude":
2067
+ _cuda_backend._reverse_drude_current(
2068
+ AdjElectricPrev=pre_step_adjoint[component_name],
2069
+ AdjCurrentPrev=pre_step_adjoint[current_name],
2070
+ AdjCurrentPost=adj_current_corrected,
2071
+ DrudeDrive=entry["drive"],
2072
+ decay=float(entry["decay"]),
2073
+ )
2074
+ else:
2075
+ polarization_name = _adjoint.dispersive_state_name(
2076
+ component_name, model_name, index, "polarization"
2077
+ )
2078
+ _cuda_backend._reverse_lorentz_current(
2079
+ AdjElectricPrev=pre_step_adjoint[component_name],
2080
+ AdjPolarizationPrev=pre_step_adjoint[polarization_name],
2081
+ AdjCurrentPrev=pre_step_adjoint[current_name],
2082
+ AdjPolarizationPost=adjoint_state[polarization_name],
2083
+ AdjCurrentPost=adj_current_corrected,
2084
+ LorentzDrive=entry["drive"],
2085
+ decay=float(entry["decay"]),
2086
+ restoring=float(entry["restoring"]),
2087
+ dt=dt,
2088
+ )
2089
+
2090
+ # Magnetic ADE is the same local state recurrence on H. The correction
2091
+ # kernel is reused with mu as its denominator; its accumulated mu gradient is
2092
+ # intentionally not exposed by the current public material compiler, matching
2093
+ # the existing static-mu adjoint contract.
2094
+ if magnetic_dispersive_state:
2095
+ mu_by_field = {
2096
+ "Hx": solver.mu_Hx,
2097
+ "Hy": solver.mu_Hy,
2098
+ "Hz": solver.mu_Hz,
2099
+ }
2100
+ magnetic_output_adjoint = base_result.magnetic_output_adjoint
2101
+ for component_name, model_name, index, _tensor_names, entry in (
2102
+ _adjoint.iter_magnetic_dispersive_state_specs(solver) or ()
2103
+ ):
2104
+ current_name = _adjoint.dispersive_state_name(
2105
+ component_name, model_name, index, "current"
2106
+ )
2107
+ adj_current_corrected = torch.empty_like(forward_state[current_name])
2108
+ discarded_grad_mu = torch.zeros_like(mu_by_field[component_name])
2109
+ _cuda_backend._reverse_dispersive_correction(
2110
+ AdjCurrentCorrected=adj_current_corrected,
2111
+ GradEps=discarded_grad_mu,
2112
+ AdjCurrentPost=adjoint_state[current_name],
2113
+ AdjElectricPost=magnetic_output_adjoint[component_name],
2114
+ Current=magnetic_dispersive_state[current_name],
2115
+ Eps=mu_by_field[component_name],
2116
+ dt=dt,
2117
+ )
2118
+ if model_name == "debye":
2119
+ polarization_name = _adjoint.dispersive_state_name(
2120
+ component_name, model_name, index, "polarization"
2121
+ )
2122
+ _cuda_backend._reverse_debye_current(
2123
+ AdjElectricPrev=pre_step_adjoint[component_name],
2124
+ AdjPolarizationPrev=pre_step_adjoint[polarization_name],
2125
+ AdjPolarizationPost=adjoint_state[polarization_name],
2126
+ AdjCurrentPost=adj_current_corrected,
2127
+ DebyeDrive=entry["drive"],
2128
+ decay=float(entry["decay"]),
2129
+ dt=dt,
2130
+ )
2131
+ elif model_name == "drude":
2132
+ _cuda_backend._reverse_drude_current(
2133
+ AdjElectricPrev=pre_step_adjoint[component_name],
2134
+ AdjCurrentPrev=pre_step_adjoint[current_name],
2135
+ AdjCurrentPost=adj_current_corrected,
2136
+ DrudeDrive=entry["drive"],
2137
+ decay=float(entry["decay"]),
2138
+ )
2139
+ else:
2140
+ polarization_name = _adjoint.dispersive_state_name(
2141
+ component_name, model_name, index, "polarization"
2142
+ )
2143
+ _cuda_backend._reverse_lorentz_current(
2144
+ AdjElectricPrev=pre_step_adjoint[component_name],
2145
+ AdjPolarizationPrev=pre_step_adjoint[polarization_name],
2146
+ AdjCurrentPrev=pre_step_adjoint[current_name],
2147
+ AdjPolarizationPost=adjoint_state[polarization_name],
2148
+ AdjCurrentPost=adj_current_corrected,
2149
+ LorentzDrive=entry["drive"],
2150
+ decay=float(entry["decay"]),
2151
+ restoring=float(entry["restoring"]),
2152
+ dt=dt,
2153
+ )
2154
+ step_result = _adjoint._ReverseStepResult(
2155
+ pre_step_adjoint=pre_step_adjoint,
2156
+ grad_eps_ex=grad_eps_by_field["Ex"],
2157
+ grad_eps_ey=grad_eps_by_field["Ey"],
2158
+ grad_eps_ez=grad_eps_by_field["Ez"],
2159
+ backend=_NATIVE_REVERSE_LABELS[_ReverseBackend.DISPERSIVE],
2160
+ source_adjoint_state=dict(adjoint_state),
2161
+ )
2162
+ # The reference dispersive path discards the base magnetic mid-step adjoint
2163
+ # (it returns no ``magnetic_output_adjoint``), so the single source-term
2164
+ # accumulation sees the same empty magnetic seed here.
2165
+ return _adjoint._accumulate_source_term_gradients(
2166
+ step_result,
2167
+ solver=solver,
2168
+ adjoint_state=adjoint_state,
2169
+ time_value=time_value,
2170
+ eps_ex=eps_ex,
2171
+ eps_ey=eps_ey,
2172
+ eps_ez=eps_ez,
2173
+ resolved_source_terms=resolved_source_terms,
2174
+ )
2175
+
2176
+
2177
+ def _accumulate_tfsf_sample_adjoint_native(adj_aux, terms, adjoint_fields, *, origin, ds):
2178
+ """Transpose the per-step TFSF incident injection into ``adj_aux`` via kernels.
2179
+
2180
+ Iterates the injection terms exactly like the forward per-term apply, slicing
2181
+ the adjoint of the field the incident patch was written into and launching the
2182
+ matching native sample-adjoint kernel (scalar / line / interpolated). The
2183
+ ``adj_field_patch`` is materialized contiguous so the kernels can index it
2184
+ linearly against the term coefficient patch.
2185
+ """
2186
+ for term in terms:
2187
+ offset_i, offset_j, offset_k = (int(offset) for offset in term["offsets"])
2188
+ coeff_patch = term["coeff_patch"]
2189
+ shape_i, shape_j, shape_k = (int(length) for length in coeff_patch.shape)
2190
+ adj_field_patch = adjoint_fields[term["field_name"]][
2191
+ offset_i : offset_i + shape_i,
2192
+ offset_j : offset_j + shape_j,
2193
+ offset_k : offset_k + shape_k,
2194
+ ].contiguous()
2195
+ component_scale = float(term["component_scale"])
2196
+ if "sample_positions" in term:
2197
+ _cuda_backend._accumulate_tfsf_interpolated_sample_adjoint(
2198
+ AdjAuxField=adj_aux,
2199
+ AdjFieldPatch=adj_field_patch,
2200
+ CoeffPatch=coeff_patch,
2201
+ SamplePositions=term["sample_positions"],
2202
+ origin=origin,
2203
+ ds=ds,
2204
+ componentScale=component_scale,
2205
+ )
2206
+ continue
2207
+ if term["scalar_sample_index"] is not None:
2208
+ _cuda_backend._accumulate_tfsf_scalar_sample_adjoint(
2209
+ AdjAuxField=adj_aux,
2210
+ AdjFieldPatch=adj_field_patch,
2211
+ CoeffPatch=coeff_patch,
2212
+ sampleIndex=int(term["scalar_sample_index"]),
2213
+ componentScale=component_scale,
2214
+ )
2215
+ continue
2216
+ from ..excitation.tfsf_specs import reference_sample_axis_code
2217
+
2218
+ _cuda_backend._accumulate_tfsf_line_sample_adjoint(
2219
+ AdjAuxField=adj_aux,
2220
+ AdjFieldPatch=adj_field_patch,
2221
+ CoeffPatch=coeff_patch,
2222
+ SampleIndices=term["sample_indices"],
2223
+ sampleAxisCode=int(reference_sample_axis_code(term)),
2224
+ componentScale=component_scale,
2225
+ )
2226
+
2227
+
2228
+ def _reverse_tfsf_auxiliary_state_native(solver, forward_state, adjoint_state, *, magnetic_output_adjoint):
2229
+ """Native CUDA transpose of the per-step TFSF auxiliary update.
2230
+
2231
+ Mirrors ``_reverse_tfsf_auxiliary_state_native`` one-for-one, driving
2232
+ the fused 1D auxiliary reverse kernels and the sample-adjoint kernels in the
2233
+ load-bearing order: the electric-side incident injection and the electric
2234
+ auxiliary advance both accumulate into the advanced-magnetic adjoint before the
2235
+ magnetic auxiliary advance reads it; the magnetic incident injection reads the
2236
+ base reverse's full mid-step H adjoint. Returns the aux pre-step adjoints, or an
2237
+ empty dict when the active provider carries no 1D auxiliary state.
2238
+ """
2239
+ from . import core as _adjoint
2240
+
2241
+ auxiliary_state = _adjoint._extract_tfsf_auxiliary_state(forward_state)
2242
+ if auxiliary_state is None:
2243
+ return {}
2244
+ tfsf_state = getattr(solver, "_tfsf_state", None)
2245
+ if tfsf_state is None:
2246
+ return {}
2247
+ auxiliary_grid = tfsf_state.get("auxiliary_grid")
2248
+ if auxiliary_grid is None:
2249
+ return {}
2250
+
2251
+ import torch
2252
+
2253
+ adj_electric_prev = torch.zeros_like(auxiliary_state["electric"])
2254
+ adj_magnetic_prev = torch.empty_like(auxiliary_state["magnetic"])
2255
+ adj_magnetic_after = adjoint_state["tfsf_aux_magnetic"].detach().clone()
2256
+
2257
+ ds = float(auxiliary_grid.ds)
2258
+ origin_electric = float(auxiliary_grid.s_min)
2259
+ origin_magnetic = float(auxiliary_grid.s_min + 0.5 * ds)
2260
+
2261
+ # (1) Electric incident injection samples the advanced magnetic grid.
2262
+ _accumulate_tfsf_sample_adjoint_native(
2263
+ adj_magnetic_after,
2264
+ tfsf_state.get("electric_terms", ()),
2265
+ adjoint_state,
2266
+ origin=origin_magnetic,
2267
+ ds=ds,
2268
+ )
2269
+ # (2) Electric auxiliary advance reverse.
2270
+ _cuda_backend._reverse_tfsf_auxiliary_electric(
2271
+ AdjElectricPrev=adj_electric_prev,
2272
+ AdjMagneticAfter=adj_magnetic_after,
2273
+ AdjElectricPost=adjoint_state["tfsf_aux_electric"],
2274
+ ElectricDecay=auxiliary_grid.electric_decay,
2275
+ ElectricCurl=auxiliary_grid.electric_curl,
2276
+ sourceIndex=int(auxiliary_grid.source_index),
2277
+ )
2278
+ # (3) Magnetic auxiliary advance reverse.
2279
+ _cuda_backend._reverse_tfsf_auxiliary_magnetic(
2280
+ AdjElectricPrev=adj_electric_prev,
2281
+ AdjMagneticPrev=adj_magnetic_prev,
2282
+ AdjMagneticAfter=adj_magnetic_after,
2283
+ MagneticDecay=auxiliary_grid.magnetic_decay,
2284
+ MagneticCurl=auxiliary_grid.magnetic_curl,
2285
+ )
2286
+ # (4) Magnetic incident injection samples the pre-step electric grid; it lands
2287
+ # on the mid-step H, so read the full mid-H adjoint from the base reverse.
2288
+ _accumulate_tfsf_sample_adjoint_native(
2289
+ adj_electric_prev,
2290
+ tfsf_state.get("magnetic_terms", ()),
2291
+ magnetic_output_adjoint,
2292
+ origin=origin_electric,
2293
+ ds=ds,
2294
+ )
2295
+ return {
2296
+ "tfsf_aux_electric": adj_electric_prev,
2297
+ "tfsf_aux_magnetic": adj_magnetic_prev,
2298
+ }
2299
+
2300
+
2301
+ def _reverse_step_tfsf_native(
2302
+ solver,
2303
+ forward_state,
2304
+ adjoint_state,
2305
+ *,
2306
+ time_value,
2307
+ eps_ex,
2308
+ eps_ey,
2309
+ eps_ez,
2310
+ resolved_source_terms,
2311
+ profiler,
2312
+ ):
2313
+ """Fully native reverse step for the TFSF (total-field/scattered-field) case.
2314
+
2315
+ The native standard or CPML core reverse produces the pre-step field/psi
2316
+ adjoint (its mid-step H replay consumes
2317
+ the materialized magnetic incident source terms, so the injected incident field
2318
+ is folded into the mid-H exactly as the forward step did), then the native TFSF
2319
+ auxiliary reverse kernels add the 1D auxiliary electric/magnetic pre-step
2320
+ adjoints. TFSF incident terms are literal patches, so this runner needs no
2321
+ source-term permittivity-gradient accumulation.
2322
+ """
2323
+ from . import core as _adjoint
2324
+
2325
+ tfsf_source_terms = _adjoint._tfsf_magnetic_source_terms(
2326
+ solver,
2327
+ forward_state,
2328
+ time_value=time_value,
2329
+ resolved_source_terms=resolved_source_terms,
2330
+ )
2331
+
2332
+ if getattr(solver, "uses_cpml", False):
2333
+ base_result = _reverse_step_cpml_native_core(
2334
+ solver,
2335
+ forward_state,
2336
+ adjoint_state,
2337
+ time_value=time_value,
2338
+ eps_ex=eps_ex,
2339
+ eps_ey=eps_ey,
2340
+ eps_ez=eps_ez,
2341
+ resolved_source_terms=tfsf_source_terms,
2342
+ )
2343
+ else:
2344
+ base_result = _reverse_step_standard_native_core(
2345
+ solver,
2346
+ forward_state,
2347
+ adjoint_state,
2348
+ time_value=time_value,
2349
+ eps_ex=eps_ex,
2350
+ eps_ey=eps_ey,
2351
+ eps_ez=eps_ez,
2352
+ resolved_source_terms=tfsf_source_terms,
2353
+ )
2354
+
2355
+ pre_step_adjoint = dict(base_result.pre_step_adjoint)
2356
+ auxiliary_grads = _reverse_tfsf_auxiliary_state_native(
2357
+ solver,
2358
+ forward_state,
2359
+ adjoint_state,
2360
+ magnetic_output_adjoint=base_result.magnetic_output_adjoint,
2361
+ )
2362
+ pre_step_adjoint.update(auxiliary_grads)
2363
+
2364
+ return _adjoint._ReverseStepResult(
2365
+ pre_step_adjoint=pre_step_adjoint,
2366
+ grad_eps_ex=base_result.grad_eps_ex,
2367
+ grad_eps_ey=base_result.grad_eps_ey,
2368
+ grad_eps_ez=base_result.grad_eps_ez,
2369
+ backend=_NATIVE_REVERSE_LABELS[_ReverseBackend.TFSF],
2370
+ magnetic_output_adjoint=base_result.magnetic_output_adjoint,
2371
+ )
2372
+
2373
+
2374
+ def _full_aniso_neg_inv_deltas(solver):
2375
+ """Cache the negated electric inverse spacings the off-diagonal fold needs.
2376
+
2377
+ The off-diagonal reverse folds each ``adj_curl`` into two mid-step H adjoints
2378
+ with opposite signs (``curl = D_a - D_b``). The shared backward-difference
2379
+ transpose only accumulates ``+D^T``, so the subtractive fold passes a negated
2380
+ inverse-spacing vector (the transpose is linear in it). These 1D vectors are
2381
+ static for the whole reverse pass, so they are negated once and cached on the
2382
+ solver rather than per step, keeping the per-cell reverse math fully native.
2383
+ """
2384
+ cached = getattr(solver, "_full_aniso_neg_inv_deltas", None)
2385
+ if cached is None:
2386
+ cached = {
2387
+ "x": -solver.inv_dx_e,
2388
+ "y": -solver.inv_dy_e,
2389
+ "z": -solver.inv_dz_e,
2390
+ }
2391
+ try:
2392
+ solver._full_aniso_neg_inv_deltas = cached
2393
+ except (AttributeError, TypeError):
2394
+ pass
2395
+ return cached
2396
+
2397
+
2398
+ def _reverse_step_full_aniso_native(
2399
+ solver,
2400
+ forward_state,
2401
+ adjoint_state,
2402
+ *,
2403
+ time_value,
2404
+ eps_ex,
2405
+ eps_ey,
2406
+ eps_ez,
2407
+ resolved_source_terms,
2408
+ profiler,
2409
+ ):
2410
+ """Fully native reverse step for a full (off-diagonal) anisotropic CPML medium.
2411
+
2412
+ The off-diagonal coupling only feeds the mid-step magnetic adjoint the diagonal
2413
+ CPML reverse already consumes, so this runner (1) forms the off-diagonal
2414
+ curl(H) cotangents with the native ``full_aniso_curl_adjoint`` op, (2) folds
2415
+ them into a clone of the post-step H adjoint with the shared native
2416
+ backward-difference transpose, and (3) delegates to the linear CPML native core
2417
+ on that augmented adjoint state. Every per-cell reverse computation runs inside
2418
+ the compiled kernels; Python only sequences the launches (matching
2419
+ ``reverse_step_full_aniso_cpml_native`` exactly).
2420
+ """
2421
+ import dataclasses
2422
+
2423
+ import torch
2424
+
2425
+ from . import core as _adjoint
2426
+
2427
+ adj_ex = adjoint_state["Ex"]
2428
+ adj_ey = adjoint_state["Ey"]
2429
+ adj_ez = adjoint_state["Ez"]
2430
+
2431
+ # (1) Off-diagonal curl(H) cotangents on their own electric-edge grids (assign).
2432
+ adj_curl_x = torch.empty_like(adj_ex)
2433
+ adj_curl_y = torch.empty_like(adj_ey)
2434
+ adj_curl_z = torch.empty_like(adj_ez)
2435
+ _cuda_backend._full_aniso_curl_adjoint(
2436
+ AdjCurlX=adj_curl_x,
2437
+ AdjCurlY=adj_curl_y,
2438
+ AdjCurlZ=adj_curl_z,
2439
+ AdjEx=adj_ex,
2440
+ AdjEy=adj_ey,
2441
+ AdjEz=adj_ez,
2442
+ CoeffExY=solver.cex_aniso_y,
2443
+ CoeffExZ=solver.cex_aniso_z,
2444
+ CoeffEyX=solver.cey_aniso_x,
2445
+ CoeffEyZ=solver.cey_aniso_z,
2446
+ CoeffEzX=solver.cez_aniso_x,
2447
+ CoeffEzY=solver.cez_aniso_y,
2448
+ )
2449
+
2450
+ # (2) Fold into a clone of the post-step H adjoint (native accumulate). The
2451
+ # +/- sign of each curl derivative matches the diagonal curl fold; the minus
2452
+ # folds pass a negated inverse spacing (the transpose is linear in it).
2453
+ neg = _full_aniso_neg_inv_deltas(solver)
2454
+ aug_hx = adjoint_state["Hx"].detach().clone()
2455
+ aug_hy = adjoint_state["Hy"].detach().clone()
2456
+ aug_hz = adjoint_state["Hz"].detach().clone()
2457
+ # curl_x = dHz/dy - dHy/dz -> Hz += Dy^T(curl_x); Hy -= Dz^T(curl_x)
2458
+ _cuda_backend._accumulate_backward_diff_y(FieldGrad=aug_hz, DiffGrad=adj_curl_x, invDy=solver.inv_dy_e)
2459
+ _cuda_backend._accumulate_backward_diff_z(FieldGrad=aug_hy, DiffGrad=adj_curl_x, invDz=neg["z"])
2460
+ # curl_y = dHx/dz - dHz/dx -> Hx += Dz^T(curl_y); Hz -= Dx^T(curl_y)
2461
+ _cuda_backend._accumulate_backward_diff_z(FieldGrad=aug_hx, DiffGrad=adj_curl_y, invDz=solver.inv_dz_e)
2462
+ _cuda_backend._accumulate_backward_diff_x(FieldGrad=aug_hz, DiffGrad=adj_curl_y, invDx=neg["x"])
2463
+ # curl_z = dHy/dx - dHx/dy -> Hy += Dx^T(curl_z); Hx -= Dy^T(curl_z)
2464
+ _cuda_backend._accumulate_backward_diff_x(FieldGrad=aug_hy, DiffGrad=adj_curl_z, invDx=solver.inv_dx_e)
2465
+ _cuda_backend._accumulate_backward_diff_y(FieldGrad=aug_hx, DiffGrad=adj_curl_z, invDy=neg["y"])
2466
+
2467
+ adjoint_state_aug = dict(adjoint_state)
2468
+ adjoint_state_aug["Hx"] = aug_hx
2469
+ adjoint_state_aug["Hy"] = aug_hy
2470
+ adjoint_state_aug["Hz"] = aug_hz
2471
+
2472
+ # (3) Diagonal CPML reverse on the augmented adjoint state.
2473
+ step_result = _reverse_step_cpml_native_core(
2474
+ solver,
2475
+ forward_state,
2476
+ adjoint_state_aug,
2477
+ time_value=time_value,
2478
+ eps_ex=eps_ex,
2479
+ eps_ey=eps_ey,
2480
+ eps_ez=eps_ez,
2481
+ resolved_source_terms=resolved_source_terms,
2482
+ )
2483
+ step_result = dataclasses.replace(
2484
+ step_result, backend=_NATIVE_REVERSE_LABELS[_ReverseBackend.FULL_ANISO]
2485
+ )
2486
+ # Source-term eps-gradient accumulation reads the electric adjoint only, so it
2487
+ # runs on the original (unaugmented) adjoint state.
2488
+ return _adjoint._accumulate_source_term_gradients(
2489
+ step_result,
2490
+ solver=solver,
2491
+ adjoint_state=adjoint_state,
2492
+ time_value=time_value,
2493
+ eps_ex=eps_ex,
2494
+ eps_ey=eps_ey,
2495
+ eps_ez=eps_ez,
2496
+ resolved_source_terms=resolved_source_terms,
2497
+ )
2498
+
2499
+
2500
+ def register_native_reverse_backends() -> None:
2501
+ """Register every available native CUDA reverse-step runner."""
2502
+ register_native_reverse_backend(
2503
+ _ReverseBackend.GENERAL_NONLINEAR,
2504
+ _reverse_step_nonlinear_native,
2505
+ qualifier=_cuda_scene_native_qualifies,
2506
+ )
2507
+ register_native_reverse_backend(
2508
+ _ReverseBackend.MIXED_BLOCH_CPML,
2509
+ _reverse_step_mixed_bloch_cpml_native,
2510
+ qualifier=_cuda_scene_native_qualifies,
2511
+ )
2512
+ register_native_reverse_backend(
2513
+ _ReverseBackend.GRATING_TFSF,
2514
+ _reverse_step_mixed_bloch_cpml_native,
2515
+ qualifier=_cuda_scene_native_qualifies,
2516
+ )
2517
+ register_native_reverse_backend(
2518
+ _ReverseBackend.STANDARD,
2519
+ _reverse_step_standard_native,
2520
+ qualifier=_cuda_scene_native_qualifies,
2521
+ )
2522
+ register_native_reverse_backend(
2523
+ _ReverseBackend.CPML,
2524
+ _reverse_step_cpml_native,
2525
+ qualifier=_cuda_scene_native_qualifies,
2526
+ )
2527
+ register_native_reverse_backend(
2528
+ _ReverseBackend.CONDUCTIVE,
2529
+ _reverse_step_conductive_native,
2530
+ qualifier=_cuda_scene_native_qualifies,
2531
+ )
2532
+ register_native_reverse_backend(
2533
+ _ReverseBackend.KERR,
2534
+ _reverse_step_kerr_native,
2535
+ qualifier=_cuda_scene_native_qualifies,
2536
+ )
2537
+ register_native_reverse_backend(
2538
+ _ReverseBackend.FULL_ANISO,
2539
+ _reverse_step_full_aniso_native,
2540
+ qualifier=_cuda_scene_native_qualifies,
2541
+ )
2542
+ register_native_reverse_backend(
2543
+ _ReverseBackend.BLOCH,
2544
+ _reverse_step_bloch_native,
2545
+ qualifier=_cuda_scene_native_qualifies,
2546
+ )
2547
+ register_native_reverse_backend(
2548
+ _ReverseBackend.BLOCH_DISPERSIVE,
2549
+ _reverse_step_bloch_dispersive_native,
2550
+ qualifier=_cuda_scene_native_qualifies,
2551
+ )
2552
+ register_native_reverse_backend(
2553
+ _ReverseBackend.DISPERSIVE,
2554
+ _reverse_step_dispersive_native,
2555
+ qualifier=_dispersive_native_qualifies,
2556
+ )
2557
+ register_native_reverse_backend(
2558
+ _ReverseBackend.TFSF,
2559
+ _reverse_step_tfsf_native,
2560
+ qualifier=_cuda_scene_native_qualifies,
2561
+ )