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.
- witwin/maxwell/__init__.py +164 -0
- witwin/maxwell/adapters/__init__.py +1 -0
- witwin/maxwell/adapters/gds.py +203 -0
- witwin/maxwell/adapters/tidy3d.py +1669 -0
- witwin/maxwell/adjoint_inputs.py +70 -0
- witwin/maxwell/compiler/__init__.py +11 -0
- witwin/maxwell/compiler/materials.py +1355 -0
- witwin/maxwell/compiler/monitors.py +164 -0
- witwin/maxwell/compiler/sources.py +222 -0
- witwin/maxwell/fdfd/__init__.py +3 -0
- witwin/maxwell/fdfd/adjoint/__init__.py +3 -0
- witwin/maxwell/fdfd/adjoint/bridge.py +180 -0
- witwin/maxwell/fdfd/ams.py +300 -0
- witwin/maxwell/fdfd/krylov.py +280 -0
- witwin/maxwell/fdfd/plotting.py +110 -0
- witwin/maxwell/fdfd/postprocess.py +33 -0
- witwin/maxwell/fdfd/solver.py +1057 -0
- witwin/maxwell/fdtd/__init__.py +4 -0
- witwin/maxwell/fdtd/adjoint/__init__.py +21 -0
- witwin/maxwell/fdtd/adjoint/bridge.py +746 -0
- witwin/maxwell/fdtd/adjoint/capabilities.py +29 -0
- witwin/maxwell/fdtd/adjoint/core.py +3614 -0
- witwin/maxwell/fdtd/adjoint/dispatch.py +684 -0
- witwin/maxwell/fdtd/adjoint/native.py +2561 -0
- witwin/maxwell/fdtd/adjoint/profiler.py +195 -0
- witwin/maxwell/fdtd/adjoint/reverse_common.py +207 -0
- witwin/maxwell/fdtd/adjoint/seeds.py +524 -0
- witwin/maxwell/fdtd/boundary/__init__.py +37 -0
- witwin/maxwell/fdtd/boundary/bloch.py +90 -0
- witwin/maxwell/fdtd/boundary/common.py +41 -0
- witwin/maxwell/fdtd/boundary/cpml.py +527 -0
- witwin/maxwell/fdtd/boundary/runtime.py +209 -0
- witwin/maxwell/fdtd/checkpoint.py +227 -0
- witwin/maxwell/fdtd/coords.py +32 -0
- witwin/maxwell/fdtd/cuda/__init__.py +11 -0
- witwin/maxwell/fdtd/cuda/backend.py +3532 -0
- witwin/maxwell/fdtd/cuda/build.py +292 -0
- witwin/maxwell/fdtd/cuda/extension.cpp +2053 -0
- witwin/maxwell/fdtd/cuda/kernels/adjoint.cu +5290 -0
- witwin/maxwell/fdtd/cuda/kernels/aniso.cu +1704 -0
- witwin/maxwell/fdtd/cuda/kernels/boundary.cu +424 -0
- witwin/maxwell/fdtd/cuda/kernels/common.cu +74 -0
- witwin/maxwell/fdtd/cuda/kernels/common.cuh +82 -0
- witwin/maxwell/fdtd/cuda/kernels/dispersive.cu +782 -0
- witwin/maxwell/fdtd/cuda/kernels/electric.cu +3613 -0
- witwin/maxwell/fdtd/cuda/kernels/magnetic.cu +781 -0
- witwin/maxwell/fdtd/cuda/kernels/observers.cu +347 -0
- witwin/maxwell/fdtd/cuda/kernels/projection.cu +211 -0
- witwin/maxwell/fdtd/cuda/kernels/sources.cu +1482 -0
- witwin/maxwell/fdtd/cuda/kernels/spectral.cu +297 -0
- witwin/maxwell/fdtd/cuda/launch.h +23 -0
- witwin/maxwell/fdtd/cuda/prebuilt/witwin_maxwell_fdtd_cuda.pyd +0 -0
- witwin/maxwell/fdtd/cuda/runtime/__init__.py +6 -0
- witwin/maxwell/fdtd/cuda/runtime/dispatcher.py +6 -0
- witwin/maxwell/fdtd/cuda/runtime/graph.py +33 -0
- witwin/maxwell/fdtd/cuda/tensors.h +35 -0
- witwin/maxwell/fdtd/dispersion.py +48 -0
- witwin/maxwell/fdtd/evaluation.py +290 -0
- witwin/maxwell/fdtd/excitation/__init__.py +27 -0
- witwin/maxwell/fdtd/excitation/injection.py +1102 -0
- witwin/maxwell/fdtd/excitation/modes.py +1675 -0
- witwin/maxwell/fdtd/excitation/spatial.py +469 -0
- witwin/maxwell/fdtd/excitation/temporal.py +510 -0
- witwin/maxwell/fdtd/excitation/tfsf_apply.py +264 -0
- witwin/maxwell/fdtd/excitation/tfsf_common.py +421 -0
- witwin/maxwell/fdtd/excitation/tfsf_specs.py +490 -0
- witwin/maxwell/fdtd/excitation/tfsf_state.py +778 -0
- witwin/maxwell/fdtd/io.py +31 -0
- witwin/maxwell/fdtd/material_pullback.py +371 -0
- witwin/maxwell/fdtd/meshing/__init__.py +3 -0
- witwin/maxwell/fdtd/meshing/autogrid.py +335 -0
- witwin/maxwell/fdtd/observers.py +1453 -0
- witwin/maxwell/fdtd/plotting.py +121 -0
- witwin/maxwell/fdtd/postprocess.py +205 -0
- witwin/maxwell/fdtd/runtime/__init__.py +119 -0
- witwin/maxwell/fdtd/runtime/initialization.py +236 -0
- witwin/maxwell/fdtd/runtime/materials.py +1563 -0
- witwin/maxwell/fdtd/runtime/module_cache.py +31 -0
- witwin/maxwell/fdtd/runtime/spectral.py +581 -0
- witwin/maxwell/fdtd/runtime/stepping.py +1987 -0
- witwin/maxwell/fdtd/solver.py +502 -0
- witwin/maxwell/geometry/__init__.py +5 -0
- witwin/maxwell/geometry/polyslab.py +299 -0
- witwin/maxwell/materials.py +3 -0
- witwin/maxwell/media.py +1467 -0
- witwin/maxwell/monitors.py +743 -0
- witwin/maxwell/ports.py +113 -0
- witwin/maxwell/postprocess/RCS.py +139 -0
- witwin/maxwell/postprocess/__init__.py +42 -0
- witwin/maxwell/postprocess/diffraction.py +219 -0
- witwin/maxwell/postprocess/directivity.py +152 -0
- witwin/maxwell/postprocess/emission.py +67 -0
- witwin/maxwell/postprocess/modal.py +324 -0
- witwin/maxwell/postprocess/nfft.py +266 -0
- witwin/maxwell/postprocess/scattering_parameters.py +358 -0
- witwin/maxwell/postprocess/stratton_chu.py +1193 -0
- witwin/maxwell/result.py +967 -0
- witwin/maxwell/scene.py +1338 -0
- witwin/maxwell/simulation.py +659 -0
- witwin/maxwell/sources.py +917 -0
- witwin/maxwell/visualization/__init__.py +28 -0
- witwin/maxwell/visualization/interactive.py +484 -0
- witwin/maxwell/visualization/plots.py +568 -0
- witwin_maxwell-0.3.0.dist-info/METADATA +184 -0
- witwin_maxwell-0.3.0.dist-info/RECORD +107 -0
- witwin_maxwell-0.3.0.dist-info/WHEEL +4 -0
- 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
|
+
)
|