Snapshot from this documentation build. SHA-256: c44df56f417219e957b9fc5381652d5b630c52682a058ad9556ab6d167bd770e
1 export RadiationClosure, RadiationWorkspace, PhotonMomentState,
2 allocate_radiation_workspace, allocate_photon_moment_state,
3 initialize_photon_moments!, photon_fields_to_workspace!, reconstruct_radiation_density!, photon_moment_energy,
4 RadiationEnergyLedger, RadiationStepTiming, radiation_energy_components,
5 step_radiation_godunov!, radiation_angular_response,
6 radiation_mode_coefficients, radiation_free_streaming_bound
7
8 const _RADIATION_RESPONSE_BRIDGE = Int32(0)
9 const _RADIATION_RESPONSE_ANGULAR = Int32(1)
10 const _RADIATION_RESPONSE_MOMENTS = Int32(2)
11 const _RADIATION_RESPONSE_MOMENTS_EXPONENTIAL = Int32(3)
12
13 @inline function _radiation_response_code(model::Symbol)
14 model === :bridge && return _RADIATION_RESPONSE_BRIDGE
15 model in (:angular, :transport) && return _RADIATION_RESPONSE_ANGULAR
16 model in (:moments, :photon_moments) && return _RADIATION_RESPONSE_MOMENTS
17 model in (:moments_exponential, :photon_moments_exponential) &&
18 return _RADIATION_RESPONSE_MOMENTS_EXPONENTIAL
19 error("radiation response_model must be :bridge, :angular, :moments, or " *
20 ":moments_exponential")
21 end
22
23 @inline _radiation_response_name(code::Int32) =
24 code == _RADIATION_RESPONSE_MOMENTS_EXPONENTIAL ? :moments_exponential :
25 code == _RADIATION_RESPONSE_MOMENTS ? :moments :
26 code == _RADIATION_RESPONSE_ANGULAR ? :angular : :bridge
27
28 @inline _radiation_evolves_moments(code::Int32) =
29 code == _RADIATION_RESPONSE_MOMENTS ||
30 code == _RADIATION_RESPONSE_MOMENTS_EXPONENTIAL
31
32 """
33 RadiationClosure(; inertia_ratio, mean_free_path, photon_speed,
34 gas_sound_speed, bridge_q2=5, longitudinal_viscosity=4/3)
35
36 Leading-order photon-baryon response used by `step_radiation_godunov!`.
37 Lengths and velocities use the MHD state's code units. `inertia_ratio` is
38 `R=3 rho_b c^2/(4 rho_gamma)`. The closure is asymptotically matched between
39 the tight-coupling diffusion limit and the local Compton-drag limit. The
40 transition kernel is deliberately exposed because it is a calibration surface,
41 not an exact Boltzmann transport solution.
42 """
43 struct RadiationClosure{T}
44 inertia_ratio::T
45 mean_free_path::T
46 photon_speed::T
47 gas_sound_speed::T
48 bridge_q2::T
49 longitudinal_viscosity::T
50 nyquist_guard_fraction::T
51 response_model::Int32
52 end
53
54 RadiationClosure{T}(R,l,c,cs,q2,lvisc) where {T} =
55 RadiationClosure{T}(R,l,c,cs,q2,lvisc,one(T),_RADIATION_RESPONSE_BRIDGE)
56 RadiationClosure{T}(R,l,c,cs,q2,lvisc,guard) where {T} =
57 RadiationClosure{T}(R,l,c,cs,q2,lvisc,guard,_RADIATION_RESPONSE_BRIDGE)
58
59 function RadiationClosure(; inertia_ratio::Real, mean_free_path::Real,
60 photon_speed::Real, gas_sound_speed::Real,
61 bridge_q2::Real=5, longitudinal_viscosity::Real=4/3,
62 nyquist_guard_fraction::Real=1,
63 response_model::Symbol=:bridge)
64 vals = promote(float(inertia_ratio), float(mean_free_path), float(photon_speed),
65 float(gas_sound_speed), float(bridge_q2),
66 float(longitudinal_viscosity),float(nyquist_guard_fraction))
67 R,lambda,c,cs,q2,lvisc,guard = vals
68 R > 0 || error("radiation inertia_ratio R must be positive")
69 lambda >= 0 || error("radiation mean_free_path must be non-negative")
70 c > 0 || error("radiation photon_speed must be positive")
71 cs >= 0 || error("radiation gas_sound_speed must be non-negative")
72 q2 > 0 || error("radiation bridge_q2 must be positive")
73 lvisc > 0 || error("radiation longitudinal_viscosity must be positive")
74 0 < guard <= 1 || error("radiation nyquist_guard_fraction must be in (0,1]")
75 return RadiationClosure(R,lambda,c,cs,q2,lvisc,guard,
76 _radiation_response_code(response_model))
77 end
78
79 struct RadiationWorkspace{T,A,V,H,C,Q,P,K,B,D,M,L,S,X}
80 packed_psi_correction::V
81 real_batch::B
82 fft_real::A
83 correction::NTuple{3,A}
84 density_perturbation::D
85 rho_mean::T
86 hat_batch::H
87 velocity_hat::NTuple{3,C}
88 density_hat::C
89 authoritative_density_hat::Q
90 # Copy of the density-mode increment taken before the inverse transforms and
91 # restored into `density_hat` after them for the commit. FFTW's multidimensional c2r transform destroys its complex input, so
92 # `density_hat` (a slice of `hat_batch`) is garbage once the corrections
93 # have been inverse-transformed on the CPU. Device FFTs (MPS, cuFFT) leave
94 # it intact, so there this is the authoritative array itself (a sentinel:
95 # `density_increment_hat === authoritative_density_hat` means "no copy").
96 density_increment_hat::Q
97 longitudinal_propagator::P
98 longitudinal_propagator_signature::K
99 magnetic_residual::M
100 energy_residual::D
101 density_hat_valid::L
102 density_limiter_active::L
103 density_limiter_min_theta::S
104 # CT states only (else `nothing`): device copy of the 6 face-state arrays
105 # (bxf,byf,bzf,rbx,rby,rbz; 6*ncells values) taken before each positivity-protected
106 # trial, because `ct_update!` overwrites them in place. 6/(9+9+21)=15% of the
107 # state+CT memory, 6 extra N^3 arrays.
108 ct_snapshot::X
109 end
110
111 @kernel function _radiation_initialize_density_k!(delta,@Const(rho),rho_mean)
112 c=@index(Global)
113 @inbounds delta[c]=rho[c]-rho_mean
114 end
115
116 @kernel function _radiation_reconstruct_density_k!(rho,@Const(delta),rho_mean,smallr)
117 c=@index(Global)
118 @inbounds rho[c]=max(rho_mean+delta[c],smallr)
119 end
120
121 """Reconstruct the compatibility density field from the authoritative perturbation."""
122 function reconstruct_radiation_density!(ws::RadiationWorkspace,s::MHDState)
123 _radiation_reconstruct_density_k!(s.be)(s.U[1],ws.density_perturbation,
124 ws.rho_mean,s.smallr;ndrange=ncells(s))
125 KA.synchronize(s.be)
126 s
127 end
128
129 function _radiation_refresh_density!(ws::RadiationWorkspace,s::MHDState)
130 _radiation_initialize_density_k!(s.be)(ws.density_perturbation,s.U[1],
131 ws.rho_mean;ndrange=ncells(s))
132 KA.synchronize(s.be)
133 ws.density_hat_valid[]=false
134 s
135 end
136
137 """
138 PhotonMomentState
139
140 Persistent linear photon moments in the RFFT half spectrum. `delta_hat` is the
141 fractional photon-energy perturbation and `velocity_hat` is the photon bulk
142 velocity in the same code units as the gas. The anisotropic stress is eliminated
143 with the angular closure in `_photon_stress_rates`; no additional full-grid
144 quadrupole arrays are required.
145 """
146 struct PhotonMomentState{H,C}
147 hat_batch::H
148 velocity_hat::NTuple{3,C}
149 delta_hat::C
150 dims::NTuple{3,Int}
151 end
152
153 mutable struct RadiationEnergyLedger
154 samples::Int
155 total_change::Float64
156 kinetic_change::Float64
157 internal_change::Float64
158 magnetic_change::Float64
159 photon_change::Float64
160 resolved_coupled_change::Float64
161 component_residual::Float64
162 component_residual_abs::Float64
163 component_magnitude_abs::Float64
164 end
165
166 RadiationEnergyLedger()=RadiationEnergyLedger(
167 0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0)
168
169 mutable struct RadiationStepTiming
170 predictor_s::Float64
171 godunov_s::Float64
172 corrector_s::Float64
173 ledger_s::Float64
174 spectral_steps::Int
175 fast_steps::Int
176 fallback_steps::Int
177 plm_hll_fallback_steps::Int
178 subdivision_steps::Int
179 global_fallback_steps::Int
180 density_limited_stages::Int
181 end
182
183 RadiationStepTiming()=RadiationStepTiming(0.0,0.0,0.0,0.0,0,0,0,0,0,0,0)
184
185 @inline function _time_radiation_stage!(f, timing::Nothing, field::Symbol)
186 f()
187 end
188
189 @inline function _time_radiation_stage!(f, timing::RadiationStepTiming, field::Symbol)
190 result=nothing
191 elapsed=@elapsed result=f()
192 setfield!(timing,field,getfield(timing,field)+elapsed)
193 result
194 end
195
196 function allocate_radiation_workspace(s::MHDState{T};
197 compensate_magnetic::Bool=true) where {T}
198 all(==(s.dims[1]), s.dims) ||
199 error("spectral radiation coupling currently requires a cubic grid")
200 T === Float32 || s.be isa CPU ||
201 error("device spectral radiation coupling currently supports Float32")
202 n = s.dims[1]
203 # The final four real fields carry persistent compensated-summation
204 # remainders for B and total energy. They prevent timestep-dependent
205 # freezing when weak perturbations evolve on top of strong Float32 guide
206 # fields and magnetic-energy backgrounds.
207 packed=device_zeros(s.be,T,(9*n^3,))
208 real_batch=reshape(view(packed,1:4*n^3),n,n,n,4)
209 fft_real=reshape(view(packed,1:n^3),n,n,n)
210 corr=ntuple(d -> reshape(view(packed,d*n^3+1:(d+1)*n^3),n,n,n),3)
211 density_perturbation=reshape(view(packed,4*n^3+1:5*n^3),s.dims)
212 magnetic_residual=compensate_magnetic ? ntuple(d ->
213 reshape(view(packed,(4+d)*n^3+1:(5+d)*n^3),s.dims),3) : nothing
214 energy_residual=reshape(view(packed,8*n^3+1:9*n^3),s.dims)
215 hdims = (n ÷ 2 + 1, n, n)
216 hat_batch=device_zeros(s.be,Complex{T},(hdims...,4))
217 density_hat=reshape(view(hat_batch,:,:,:,1),hdims)
218 hats=ntuple(d -> reshape(view(hat_batch,:,:,:,d+1),hdims),3)
219 authoritative_density_hat=device_zeros(s.be,Complex{T},hdims)
220 density_increment_hat=s.be isa CPU ?
221 device_zeros(s.be,Complex{T},hdims) : authoritative_density_hat
222 max_mode2=3*(n÷2-1)^2
223 longitudinal_propagator=device_zeros(s.be,T,(4,5,max_mode2+1))
224 rho_mean=T(sum(s.U[1])/length(s.U[1]))
225 _radiation_initialize_density_k!(s.be)(density_perturbation,s.U[1],rho_mean;
226 ndrange=ncells(s))
227 KA.synchronize(s.be)
228 ct_snapshot=uses_ct(s) ? device_zeros(s.be,T,(6*ncells(s),)) : nothing
229 RadiationWorkspace(packed,real_batch,fft_real,corr,density_perturbation,
230 rho_mean,hat_batch,hats,density_hat,
231 authoritative_density_hat,density_increment_hat,
232 longitudinal_propagator,
233 Ref(typemax(UInt)),magnetic_residual,energy_residual,
234 Ref(false),Ref(false),Ref(1.0),ct_snapshot)
235 end
236
237 function _radiation_reset_magnetic_residual!(ws::RadiationWorkspace)
238 ws.magnetic_residual === nothing ||
239 foreach(r -> fill!(r,zero(eltype(r))),ws.magnetic_residual)
240 fill!(ws.energy_residual,zero(eltype(ws.energy_residual)))
241 nothing
242 end
243
244 # Also reset the CT face remainders of a CT state (they are the CT analogue of
245 # `ws.magnetic_residual`).
246 function _radiation_reset_magnetic_residual!(ws::RadiationWorkspace,s::MHDState)
247 _radiation_reset_magnetic_residual!(ws)
248 _ct_reset_face_residual!(s)
249 nothing
250 end
251
252 function _radiation_trial_admissible(s::MHDState)
253 rho_min=Float64(minimum(s.U[1]))
254 isfinite(rho_min) && rho_min>=max(Float64(s.smallr),Float64(s.llf_dmin))
255 end
256
257 # The CT arrays that `ct_update!` overwrites in place and that therefore need a
258 # pre-trial copy: faces and their Kahan remainders, in snapshot order.
259 _ct_snapshot_arrays(ct)=(ct.bxf,ct.byf,ct.bzf,ct.rbx,ct.rby,ct.rbz)
260
261 function _ct_snapshot_save!(s::MHDState,ws::RadiationWorkspace)
262 snap=ws.ct_snapshot
263 snap===nothing && error("CT positivity fallback needs a workspace allocated " *
264 "from the CT state (ct_snapshot missing)")
265 nc=ncells(s)
266 for (q,a) in enumerate(_ct_snapshot_arrays(s.ct))
267 copyto!(view(snap,(q-1)*nc+1:q*nc),a)
268 end
269 nothing
270 end
271
272 function _ct_snapshot_restore!(s::MHDState,ws::RadiationWorkspace)
273 snap=ws.ct_snapshot
274 snap===nothing && error("CT trial restore needs a workspace allocated from " *
275 "the CT state (ct_snapshot missing)")
276 nc=ncells(s)
277 for (q,a) in enumerate(_ct_snapshot_arrays(s.ct))
278 copyto!(a,view(snap,(q-1)*nc+1:q*nc))
279 end
280 nothing
281 end
282
283 # Bring back the complete pre-step state. Cell fields are the swapped-out
284 # `s.scratch`. For CT the faces and their Kahan remainders were also advanced in
285 # place and come back from the snapshot, exactly (no remainder is lost); cell B
286 # (a face average) is already the pre-step one from the swap.
287 function _restore_radiation_trial!(s::MHDState,ws::RadiationWorkspace,
288 track_density::Bool)
289 s.U,s.scratch=s.scratch,s.U
290 # A rejected trial has already replaced the in-place Kahan remainders.
291 # For GLM, resetting loses at most one Float32 rounding remainder while
292 # preserving the accepted conservative state; fallback steps are exceptional.
293 _radiation_reset_magnetic_residual!(ws,s)
294 uses_ct(s) && _ct_snapshot_restore!(s,ws)
295 if track_density
296 _radiation_initialize_density_k!(s.be)(ws.density_perturbation,s.U[1],
297 ws.rho_mean;ndrange=ncells(s))
298 KA.synchronize(s.be)
299 ws.density_hat_valid[]=false
300 end
301 s
302 end
303
304 function _checkpoint_fallback_rejection!(s::MHDState,ws::RadiationWorkspace,
305 track_density::Bool;dt,debug_context,trial_rho_min,trial_wave_speed,
306 wave_limit)
307 _restore_radiation_trial!(s,ws,track_density)
308 error("RADIATION_CHECKPOINT_CFL_FALLBACK: rejected coupled " *
309 "trial dt=$dt $debug_context trial_rho_min=$trial_rho_min " *
310 "trial_wave_speed=$trial_wave_speed wave_limit=$wave_limit")
311 end
312
313 function allocate_photon_moment_state(s::MHDState{T}) where {T}
314 all(==(s.dims[1]),s.dims) ||
315 error("spectral photon moments currently require a cubic grid")
316 n=s.dims[1]; hdims=(n÷2+1,n,n)
317 hats=device_zeros(s.be,Complex{T},(hdims...,4))
318 delta=reshape(view(hats,:,:,:,1),hdims)
319 vel=ntuple(d->reshape(view(hats,:,:,:,d+1),hdims),3)
320 PhotonMomentState(hats,vel,delta,s.dims)
321 end
322
323 @kernel function _radiation_energy_components_k!(kinetic,internal,magnetic,total,
324 @Const(rho),@Const(mx),@Const(my),@Const(mz),@Const(E),
325 @Const(Bx),@Const(By),@Const(Bz),smallr)
326 c=@index(Global)
327 @inbounds begin
328 T=eltype(total); r=max(rho[c],smallr)
329 ek=T(0.5)*(mx[c]*mx[c]+my[c]*my[c]+mz[c]*mz[c])/r
330 eb=T(0.5)*(Bx[c]*Bx[c]+By[c]*By[c]+Bz[c]*Bz[c])
331 kinetic[c]=ek; magnetic[c]=eb; internal[c]=E[c]-ek-eb; total[c]=E[c]
332 end
333 end
334
335 """Return device-reduced MHD energy components without allocating full-grid temporaries."""
336 function radiation_energy_components(ws::RadiationWorkspace,s::MHDState)
337 kinetic=view(ws.real_batch,:,:,:,1); internal=view(ws.real_batch,:,:,:,2)
338 magnetic=view(ws.real_batch,:,:,:,3); total=view(ws.real_batch,:,:,:,4)
339 _radiation_energy_components_k!(s.be)(kinetic,internal,magnetic,total,
340 s.U[1],s.U[2],s.U[3],s.U[4],s.U[5],s.U[6],s.U[7],s.U[8],s.smallr;
341 ndrange=ncells(s))
342 KA.synchronize(s.be)
343 return (total=Float64(sum(total)),kinetic=Float64(sum(kinetic)),
344 internal=Float64(sum(internal)),magnetic=Float64(sum(magnetic)))
345 end
346
347 @kernel function _radiation_energy_changes_k!(dkinetic,dinternal,dmagnetic,dtotal,
348 @Const(rho),@Const(mx),@Const(my),@Const(mz),@Const(E),
349 @Const(Bx),@Const(By),@Const(Bz),
350 @Const(rho0),@Const(mx0),@Const(my0),@Const(mz0),@Const(E0),
351 @Const(Bx0),@Const(By0),@Const(Bz0),smallr)
352 c=@index(Global)
353 @inbounds begin
354 T=eltype(dtotal); r=max(rho[c],smallr); r0=max(rho0[c],smallr)
355 ek=T(0.5)*(mx[c]*mx[c]+my[c]*my[c]+mz[c]*mz[c])/r
356 ek0=T(0.5)*(mx0[c]*mx0[c]+my0[c]*my0[c]+mz0[c]*mz0[c])/r0
357 eb=T(0.5)*(Bx[c]*Bx[c]+By[c]*By[c]+Bz[c]*Bz[c])
358 eb0=T(0.5)*(Bx0[c]*Bx0[c]+By0[c]*By0[c]+Bz0[c]*Bz0[c])
359 dkinetic[c]=ek-ek0; dmagnetic[c]=eb-eb0
360 dinternal[c]=(E[c]-ek-eb)-(E0[c]-ek0-eb0)
361 dtotal[c]=E[c]-E0[c]
362 end
363 end
364
365 @kernel function _radiation_energy_residuals_k!(signed_residual,absolute_residual,
366 component_magnitude,
367 @Const(dkinetic),@Const(dinternal),@Const(dmagnetic),@Const(dtotal),
368 @Const(E),@Const(E0))
369 c=@index(Global)
370 @inbounds begin
371 dk=dkinetic[c]; di=dinternal[c]; db=dmagnetic[c]; dt=dtotal[c]
372 residual=dt-(dk+di+db)
373 signed_residual[c]=residual
374 absolute_residual[c]=abs(residual)
375 component_magnitude[c]=abs(E[c])+abs(E0[c])
376 end
377 end
378
379 function _radiation_energy_changes(ws::RadiationWorkspace,s::MHDState,old)
380 dk=view(ws.real_batch,:,:,:,1); di=view(ws.real_batch,:,:,:,2)
381 db=view(ws.real_batch,:,:,:,3); dt=view(ws.real_batch,:,:,:,4)
382 _radiation_energy_changes_k!(s.be)(dk,di,db,dt,
383 s.U[1],s.U[2],s.U[3],s.U[4],s.U[5],s.U[6],s.U[7],s.U[8],
384 old[1],old[2],old[3],old[4],old[5],old[6],old[7],old[8],s.smallr;
385 ndrange=ncells(s))
386 KA.synchronize(s.be)
387 change=(total=Float64(sum(dt)),kinetic=Float64(sum(dk)),
388 internal=Float64(sum(di)),magnetic=Float64(sum(db)))
389 # Form the identity residual per cell before reducing. On Metal, subtracting
390 # four independently reduced Float32 totals loses all useful digits for a
391 # nearly uniform state even though the pointwise decomposition is accurate.
392 _radiation_energy_residuals_k!(s.be)(dt,dk,di,dk,di,db,dt,s.U[5],old[5];
393 ndrange=ncells(s))
394 KA.synchronize(s.be)
395 merge(change,(component_residual=Float64(sum(dt)),
396 component_residual_abs=Float64(sum(dk)),
397 component_magnitude_abs=Float64(sum(di))))
398 end
399
400 function _record_radiation_energy!(ledger::RadiationEnergyLedger,change;
401 photon_change=0.0)
402 component_residual=hasproperty(change,:component_residual) ?
403 change.component_residual :
404 change.total-(change.kinetic+change.internal+change.magnetic)
405 component_residual_abs=hasproperty(change,:component_residual_abs) ?
406 change.component_residual_abs : abs(component_residual)
407 component_magnitude_abs=hasproperty(change,:component_magnitude_abs) ?
408 change.component_magnitude_abs :
409 2*(abs(change.total)+abs(change.kinetic)+abs(change.internal)+
410 abs(change.magnetic))
411 ledger.samples+=1
412 ledger.total_change+=change.total
413 ledger.kinetic_change+=change.kinetic
414 ledger.internal_change+=change.internal
415 ledger.magnetic_change+=change.magnetic
416 ledger.photon_change+=photon_change
417 ledger.resolved_coupled_change+=change.total+photon_change
418 ledger.component_residual+=component_residual
419 ledger.component_residual_abs+=component_residual_abs
420 ledger.component_magnitude_abs+=component_magnitude_abs
421 ledger
422 end
423
424 """
425 rescale_photon_velocities!(photons, factor)
426
427 Change the velocity units of the persistent photon dipole, leaving the density
428 monopole unchanged. In Jedamzik units a freely propagating physical dipole is
429 constant, so expansion from `a0` to `a1` requires `factor=sqrt(a1/a0)`.
430 """
431 function rescale_photon_velocities!(photons::PhotonMomentState, factor::Real)
432 isfinite(factor) && factor > 0 ||
433 throw(ArgumentError("photon velocity rescaling must be finite and positive"))
434 T=typeof(real(zero(eltype(photons.hat_batch))))
435 view(photons.hat_batch,:,:,:,2:4) .*= T(factor)
436 KA.synchronize(KA.get_backend(photons.hat_batch))
437 photons
438 end
439 export rescale_photon_velocities!
440
441 function _radiation_forward_batch!(hat_batch,real_batch)
442 for d in axes(real_batch,4)
443 _radiation_forward!(view(hat_batch,:,:,:,d),view(real_batch,:,:,:,d))
444 end
445 hat_batch
446 end
447
448 function _radiation_inverse_batch!(real_batch,hat_batch)
449 for d in axes(real_batch,4)
450 _radiation_inverse!(view(real_batch,:,:,:,d),view(hat_batch,:,:,:,d))
451 end
452 real_batch
453 end
454
455 """Inverse-transform photon fields into reusable workspace without changing evolved modes.
456
457 Multidimensional c2r FFTs may destroy their complex input, including FFTW's
458 out-of-place inverse. Diagnostic consumers must never transform the live state.
459 """
460 function photon_fields_to_workspace!(ws::RadiationWorkspace, photons::PhotonMomentState)
461 size(ws.hat_batch) == size(photons.hat_batch) ||
462 throw(ArgumentError("photon and workspace dimensions differ"))
463 copyto!(ws.hat_batch,photons.hat_batch)
464 _radiation_inverse_batch!(ws.real_batch,ws.hat_batch)
465 end
466
467 function _radiation_forward_triplet!(ws::RadiationWorkspace,fields)
468 for d in 1:3
469 copyto!(ws.correction[d],fields[d])
470 end
471 fill!(ws.fft_real,zero(eltype(ws.fft_real)))
472 _radiation_forward_batch!(ws.hat_batch,ws.real_batch)
473 ws.velocity_hat
474 end
475
476 function _radiation_inverse_triplet!(fields,ws::RadiationWorkspace)
477 _radiation_inverse_batch!(ws.real_batch,ws.hat_batch)
478 for d in 1:3
479 fields[d] === ws.correction[d] || copyto!(fields[d],ws.correction[d])
480 end
481 fields
482 end
483
484 const _RADIATION_FFT_CACHE = Dict{Any,Any}()
485 const _RADIATION_FFT_CACHE_LOCK = ReentrantLock()
486
487 _radiation_fft_backend_key(x) = typeof(KA.get_backend(x))
488
489 function _radiation_fft_plan_key(x, chat)
490 (_radiation_fft_backend_key(x), _radiation_fft_backend_key(chat),
491 typeof(x), typeof(chat), eltype(x), eltype(chat),
492 size(x), size(chat), strides(x), strides(chat), (1,2,3))
493 end
494
495 function _radiation_fft_plans(x::AbstractArray{T,N},
496 chat::AbstractArray{Complex{T},N}) where {T,N}
497 N in (3,4) || throw(ArgumentError("radiation FFT requires three spatial axes"))
498 expected = (size(x,1) ÷ 2 + 1, Base.tail(size(x))...)
499 size(chat) == expected || throw(DimensionMismatch(
500 "radiation FFT spectrum has size $(size(chat)); expected $expected"))
501 key = _radiation_fft_plan_key(x,chat)
502 lock(_RADIATION_FFT_CACHE_LOCK) do
503 get!(_RADIATION_FFT_CACHE, key) do
504 (plan_rfft(x,(1,2,3)), plan_irfft(chat,size(x,1),(1,2,3)))
505 end
506 end
507 end
508
509 function _radiation_forward!(chat::AbstractArray{Complex{T},3}, x::AbstractArray{T,3}) where {T}
510 fwd,_ = _radiation_fft_plans(x,chat)
511 FFTW.mul!(chat,fwd,x)
512 chat
513 end
514
515 function _radiation_inverse!(x::AbstractArray{T,3}, chat::AbstractArray{Complex{T},3}) where {T}
516 _,invp = _radiation_fft_plans(x,chat)
517 FFTW.mul!(x,invp,chat)
518 x
519 end
520
521 @kernel function _photon_initial_fields_k!(delta,vx,vy,vz,
522 @Const(rho),@Const(rho_perturbation),@Const(mx),@Const(my),@Const(mz),
523 rho_mean,smallr,
524 adiabatic::Val{ADIABATIC}) where {ADIABATIC}
525 c=@index(Global)
526 @inbounds begin
527 r=max(rho[c],smallr)
528 delta[c]=ADIABATIC ? typeof(r)(4)/typeof(r)(3)*rho_perturbation[c]/rho_mean : zero(r)
529 vx[c]=mx[c]/r; vy[c]=my[c]/r; vz[c]=mz[c]/r
530 end
531 end
532
533 """
534 initialize_photon_moments!(photons, ws, s; adiabatic_density=true,
535 comoving_velocity=true)
536
537 Initialize photons from the gas. Adiabatic initial conditions use
538 `delta_gamma = 4 delta_b / 3`; disabling them starts from a uniform photon
539 monopole. Photon and baryon velocities are initially comoving unless
540 `comoving_velocity=false`.
541 """
542 function initialize_photon_moments!(photons::PhotonMomentState,
543 ws::RadiationWorkspace,s::MHDState;
544 adiabatic_density::Bool=true,
545 comoving_velocity::Bool=true)
546 photons.dims==s.dims || error("photon and MHD grid dimensions differ")
547 rho_mean=ws.rho_mean
548 _photon_initial_fields_k!(s.be)(view(ws.real_batch,:,:,:,1),
549 view(ws.real_batch,:,:,:,2),view(ws.real_batch,:,:,:,3),
550 view(ws.real_batch,:,:,:,4),s.U[1],ws.density_perturbation,
551 s.U[2],s.U[3],s.U[4],rho_mean,s.smallr,
552 Val(adiabatic_density);ndrange=ncells(s))
553 KA.synchronize(s.be)
554 comoving_velocity || fill!(view(ws.real_batch,:,:,:,2:4),zero(eltype(s.U[1])))
555 _radiation_forward_batch!(photons.hat_batch,ws.real_batch)
556 photons
557 end
558
559 @kernel function _photon_energy_hat_k!(out,@Const(delta),@Const(vx),@Const(vy),
560 @Const(vz),n::Int32,rho_mean,R,cph)
561 p=@index(Global)
562 h=(n>>>1)+Int32(1); ii=Int32(p-1)%h
563 @inbounds begin
564 T=eltype(out)
565 w=(ii==0 || (iszero(n&Int32(1)) && ii==(n>>>1))) ? one(T) : T(2)
566 v2=abs2(vx[p])+abs2(vy[p])+abs2(vz[p])
567 # Photon inertia is rho_b/R. The monopole coefficient follows from
568 # ddelta_gamma/dt=-(4/3)div(v_gamma) and
569 # dv_gamma/dt=-(c^2/4)grad(delta_gamma).
570 out[p]=w*rho_mean/R*(T(0.5)*v2+T(3)/T(32)*cph*cph*abs2(delta[p]))
571 end
572 end
573
574 """Return the quadratic photon perturbation energy represented by the moment state."""
575 function photon_moment_energy(photons::PhotonMomentState,ws::RadiationWorkspace,
576 s::MHDState,closure::RadiationClosure)
577 n=s.dims[1]; nh=length(photons.delta_hat)
578 out=reshape(view(ws.packed_psi_correction,1:nh),size(photons.delta_hat))
579 T=eltype(s.U[1]); rho_mean=ws.rho_mean
580 _photon_energy_hat_k!(s.be)(out,photons.delta_hat,photons.velocity_hat...,
581 Int32(n),rho_mean,T(closure.inertia_ratio),T(closure.photon_speed);
582 ndrange=nh)
583 KA.synchronize(s.be)
584 Float64(sum(out))/Float64(n^3)
585 end
586
587 @inline function _radiation_acoustic_coefficients(g::T,c2::T,k2::T,dt::T) where {T}
588 halfg=T(0.5)*g
589 disc=halfg*halfg-c2*k2
590 scale=max(one(T),halfg*halfg+c2*k2)
591 if abs(disc) <= T(16)*eps(T)*scale
592 e=exp(-halfg*dt)
593 C=e
594 S=dt*e
595 elseif disc > zero(T)
596 mu=sqrt(disc)
597 # Form the slow root without subtracting nearly equal Float32 values.
598 # The direct expression can erase diffusion damping in the stiff limit.
599 omega2=c2*k2
600 rslow=-omega2/(halfg+mu)
601 rfast=-halfg-mu
602 e1=exp(rslow*dt)
603 e2=exp(rfast*dt)
604 S=(e1-e2)/(T(2)*mu)
605 B=(rslow*e1-rfast*e2)/(T(2)*mu)
606 return B,S
607 else
608 omega=sqrt(-disc)
609 e=exp(-halfg*dt)
610 C=e*cos(omega*dt)
611 S=e*sin(omega*dt)/omega
612 end
613 return C-halfg*S,S
614 end
615
616 @inline function _radiation_integral_S(B::T,S::T,g::T,omega2::T,dt::T) where {T}
617 if omega2 > T(32)*eps(T)*max(one(T),g*g)
618 return (one(T)-B-g*S)/omega2
619 end
620 x=g*dt
621 if abs(x) < T(1e-3)
622 return dt*dt*(T(0.5)-x/T(6)+x*x/T(24))
623 elseif g > zero(T)
624 return (x+expm1(-x))/(g*g)
625 end
626 T(0.5)*dt*dt
627 end
628
629 @inline function _radiation_mode_coefficients(k2::T, gamma_drag::T,
630 rc::RadiationClosure{T}) where {T}
631 q2=rc.mean_free_path*rc.mean_free_path*k2
632 if rc.response_model == _RADIATION_RESPONSE_ANGULAR ||
633 _radiation_evolves_moments(rc.response_model)
634 a0,along,atrans=_radiation_angular_response(q2)
635 stream=one(T)-atrans
636 inertia=rc.inertia_ratio/(rc.inertia_ratio+atrans)
637 gamma_t=inertia*gamma_drag*stream
638 # The transverse response is the exact static angular integral. The
639 # longitudinal viscosity factor retains the diffusion-limit 4/3
640 # coefficient and approaches local drag as photons free-stream.
641 gamma_l=gamma_t*(rc.longitudinal_viscosity*atrans+stream)
642 # along^2 gives the tight-coupling EOS and removes its restoring
643 # frequency in the optically thin limit. This is the remaining
644 # longitudinal closure approximation, exposed by response_model.
645 crad2=rc.photon_speed*rc.photon_speed*along*along/
646 (T(3)*(rc.inertia_ratio+along))
647 return stream,inertia,gamma_t,gamma_l,crad2
648 end
649 # Equation 57 of Jedamzik et al. is a diffusion-limit EOS. In the
650 # free-streaming system photon pressure is absent and scattering supplies
651 # only local drag. Use a C1 compact bridge so c^2 times a slowly decaying
652 # rational tail cannot create a fictitious optically thin sound speed.
653 x=min(q2/rc.bridge_q2,one(T))
654 stream=x*x*(T(3)-T(2)*x)
655 tight=one(T)-stream
656 inertia_tight=rc.inertia_ratio/(one(T)+rc.inertia_ratio)
657 inertia=inertia_tight+(one(T)-inertia_tight)*stream
658 gdiff=rc.photon_speed*rc.mean_free_path*k2/(T(5)*(one(T)+rc.inertia_ratio))
659 # Diffusion is an optically thick expansion and must switch off faster than
660 # k^2 grows; a second tight-coupling weight prevents a spurious constant
661 # damping tail in the free-streaming limit.
662 gamma_t=tight*tight*gdiff+stream*gamma_drag
663 gamma_l=tight*tight*rc.longitudinal_viscosity*gdiff+stream*gamma_drag
664 crad2=tight*rc.photon_speed*rc.photon_speed/
665 (T(3)*(one(T)+rc.inertia_ratio))
666 return stream,inertia,gamma_t,gamma_l,crad2
667 end
668
669 @inline function _radiation_angular_response(q2::T) where {T}
670 if q2 < T(1e-3)
671 q4=q2*q2; q6=q4*q2
672 a0=one(T)-q2/T(3)+q4/T(5)-q6/T(7)
673 along=one(T)-T(3)*q2/T(5)+T(3)*q4/T(7)-q6/T(3)
674 atrans=one(T)-q2/T(5)+T(3)*q4/T(35)-q6/T(21)
675 return a0,along,atrans
676 end
677 q=sqrt(q2)
678 a0=atan(q)/q
679 along=T(3)*(one(T)-a0)/q2
680 atrans=T(0.5)*(T(3)*a0-along)
681 return a0,along,atrans
682 end
683
684 @inline function _photon_stress_rates(k2::T,nu::T,rc::RadiationClosure{T}) where {T}
685 q2=rc.mean_free_path*rc.mean_free_path*k2
686 _,_,atrans=_radiation_angular_response(q2)
687 # Eliminating the angular stress with nu*(A_T^-1-1) reproduces the exact
688 # static transverse transport response. It tends to c*lambda*k^2/5 in
689 # diffusion and to O(c*k) phase mixing in free streaming. Longitudinal
690 # shear has the standard 4/3 enhancement only in the diffusion limit;
691 # the evolved photon monopole supplies the longitudinal pressure memory.
692 # Avoid losing the q^2/5 diffusion term when A_T rounds to one in f32.
693 dtrans=q2<T(1e-3) ? q2/T(5)-T(8)*q2*q2/T(175) :
694 max(inv(atrans)-one(T),zero(T))
695 st=nu*dtrans
696 sl=st*(rc.longitudinal_viscosity*atrans+(one(T)-atrans))
697 return st,sl
698 end
699
700 @inline function _photon_solve_transverse(rb,rg,a,G,nu,stress)
701 T=typeof(a); a11=one(T)+a*G; a22=one(T)+a*(nu+stress)
702 det=a11*a22-a*a*G*nu
703 return ((a22*rb+a*G*rg)/det,(a*nu*rb+a11*rg)/det)
704 end
705
706 @inline function _photon_solve_longitudinal(rb,ru,rg,rug,a,G,nu,stress,k,cs2,cph2)
707 T=typeof(a); ik=Complex{T}(zero(T),k)
708 a11=one(T)+a*G+a*a*k*k*cs2
709 a22=one(T)+a*(nu+stress)+a*a*k*k*cph2/T(3)
710 b1=ru-a*ik*cs2*rb
711 b2=rug-a*ik*cph2/T(4)*rg
712 det=a11*a22-a*a*G*nu
713 ub=(a22*b1+a*G*b2)/det
714 ug=(a*nu*b1+a11*b2)/det
715 return rb-a*ik*ub,ub,rg-a*T(4)/T(3)*ik*ug,ug
716 end
717
718 @inline function _photon_sdirk_transverse(ub,ug,force,dt,G,nu,stress)
719 T=typeof(dt); g=T(1)-inv(sqrt(T(2))); a=g*dt
720 y1b,y1g=_photon_solve_transverse(ub+a*force,ug,a,G,nu,stress)
721 f1b=G*(y1g-y1b)+force
722 f1g=nu*y1b-(nu+stress)*y1g
723 rhsb=ub+(one(T)-g)*dt*f1b+a*force
724 rhsg=ug+(one(T)-g)*dt*f1g
725 _photon_solve_transverse(rhsb,rhsg,a,G,nu,stress)
726 end
727
728 @inline function _photon_sdirk_longitudinal(db,ub,dg,ug,force,dt,G,nu,stress,
729 k,cs2,cph2)
730 T=typeof(dt); g=T(1)-inv(sqrt(T(2))); a=g*dt
731 ik=Complex{T}(zero(T),k)
732 d1,u1,e1,w1=_photon_solve_longitudinal(db,ub+a*force,dg,ug,a,G,nu,
733 stress,k,cs2,cph2)
734 f1d=-ik*u1
735 f1u=-ik*cs2*d1+G*(w1-u1)+force
736 f1e=-T(4)/T(3)*ik*w1
737 f1w=-ik*cph2/T(4)*e1+nu*u1-(nu+stress)*w1
738 rdb=db+(one(T)-g)*dt*f1d
739 rub=ub+(one(T)-g)*dt*f1u+a*force
740 rdg=dg+(one(T)-g)*dt*f1e
741 rug=ug+(one(T)-g)*dt*f1w
742 _photon_solve_longitudinal(rdb,rub,rdg,rug,a,G,nu,stress,k,cs2,cph2)
743 end
744
745 @inline function _photon_exact_drag_affine(ub,ug,force,dt,G,nu,stress)
746 # Exact action of the affine velocity subsystem
747 # ub' = G(ug-ub) + force
748 # ug' = nu*ub - (nu+stress)*ug.
749 # Projecting onto its two real eigenmodes avoids a matrix exponential and,
750 # unlike forming h + sqrt(h^2-det), retains the slow diffusion root in f32.
751 T=typeof(dt)
752 h=-T(0.5)*(G+nu+stress)
753 disc=T(0.25)*(nu+stress-G)*(nu+stress-G)+G*nu
754 d=sqrt(max(disc,zero(T)))
755 den=T(2)*d
756 if den <= T(16)*eps(T)*max(one(T),abs(h))
757 return ub+dt*force,ug
758 end
759 rfast=h-d
760 determinant=G*stress
761 rslow=iszero(rfast) ? h+d : determinant/rfast
762 slowb=((-G-rfast)*ub+G*ug)/den
763 slowg=(nu*ub+(-nu-stress-rfast)*ug)/den
764 fastb=((rslow+G)*ub-G*ug)/den
765 fastg=(-nu*ub+(rslow+nu+stress)*ug)/den
766
767 islow=dt*_drag_phi1(-rslow*dt)
768 ifast=dt*_drag_phi1(-rfast*dt)
769 sourceb=(((-G-rfast)*islow+(rslow+G)*ifast)/den)*force
770 sourceg=(nu*(islow-ifast)/den)*force
771 # `slow + fast == state` only algebraically. Reconstructing the state from
772 # those two components injects an O(eps(T)) change even when dt is zero;
773 # the split applies this update many times per MHD step. Form the action as
774 # identity plus timestep-scaled increments so refinement cannot accumulate
775 # a fixed Float32 residual.
776 return ub+expm1(rslow*dt)*slowb+expm1(rfast*dt)*fastb+sourceb,
777 ug+expm1(rslow*dt)*slowg+expm1(rfast*dt)*fastg+sourceg
778 end
779
780 @inline function _photon_exact_acoustic(db,ub,dt,k,continuity,pressure)
781 T=typeof(dt)
782 speed2=continuity*pressure
783 if speed2 <= zero(T)
784 ikdt=Complex{T}(zero(T),-k*dt)
785 return db+continuity*ikdt*ub,ub+pressure*ikdt*db
786 end
787 speed=sqrt(speed2)
788 theta=k*speed*dt
789 cosm1=-T(2)*sin(T(0.5)*theta)^2
790 si=sin(theta)
791 minus_i=Complex{T}(zero(T),-one(T))
792 return db+cosm1*db+minus_i*(continuity/speed)*si*ub,
793 ub+cosm1*ub+minus_i*(pressure/speed)*si*db
794 end
795
796 @inline function _photon_split_transverse(ub,ug,force,dt,G,nu,stress)
797 _photon_exact_drag_affine(ub,ug,force,dt,G,nu,stress)
798 end
799
800 @inline function _photon_pade3_exp(a::SMatrix{5,5,T,25}) where {T}
801 # Scaling and squaring with the [3/3] Pade approximant. The longitudinal
802 # system is balanced in photon-speed velocity units before reaching this
803 # routine, so scaling follows its physical acoustic/drag rates rather than
804 # the dimensional c_gamma^2 entry of the unbalanced equations.
805 nrm=max(abs(a[1,1])+abs(a[1,2])+abs(a[1,3])+abs(a[1,4])+abs(a[1,5]),
806 abs(a[2,1])+abs(a[2,2])+abs(a[2,3])+abs(a[2,4])+abs(a[2,5]),
807 abs(a[3,1])+abs(a[3,2])+abs(a[3,3])+abs(a[3,4])+abs(a[3,5]),
808 abs(a[4,1])+abs(a[4,2])+abs(a[4,3])+abs(a[4,4])+abs(a[4,5]),
809 abs(a[5,1])+abs(a[5,2])+abs(a[5,3])+abs(a[5,4])+abs(a[5,5]))
810 scale=one(T); squarings=Int32(0)
811 while nrm*scale>T(0.125) && squarings<Int32(32)
812 scale*=T(0.5)
813 squarings+=Int32(1)
814 end
815 b=scale*a
816 b2=b*b
817 b3=b2*b
818 identity=@SMatrix [one(T) zero(T) zero(T) zero(T) zero(T);
819 zero(T) one(T) zero(T) zero(T) zero(T);
820 zero(T) zero(T) one(T) zero(T) zero(T);
821 zero(T) zero(T) zero(T) one(T) zero(T);
822 zero(T) zero(T) zero(T) zero(T) one(T)]
823 even=identity+T(0.1)*b2
824 odd=T(0.5)*b+(one(T)/T(120))*b3
825 result=(even-odd)\(even+odd)
826 for _ in Int32(1):squarings
827 result=result*result
828 end
829 result
830 end
831
832 @inline function _photon_longitudinal_propagator(dt,G,nu,stress,k,cs2,cph2)
833 T=typeof(dt)
834 cph=sqrt(cph2)
835 invc=inv(cph)
836 a=dt*@SMatrix [zero(T) -k*cph zero(T) zero(T) zero(T);
837 k*cs2*invc -G zero(T) G one(T);
838 zero(T) zero(T) zero(T) -T(4)*k*cph/T(3) zero(T);
839 zero(T) nu k*cph/T(4) -(nu+stress) zero(T);
840 zero(T) zero(T) zero(T) zero(T) zero(T)]
841 _photon_pade3_exp(a)
842 end
843
844 @inline function _photon_apply_longitudinal_propagator(propagator,mode2,
845 db,ub,dg,ug,force,cph)
846 T=typeof(cph)
847 invc=inv(cph)
848 iub=Complex{T}(-imag(ub),real(ub))*invc
849 iug=Complex{T}(-imag(ug),real(ug))*invc
850 iforce=Complex{T}(-imag(force),real(force))*invc
851 q=mode2+Int32(1)
852 y1=propagator[1,1,q]*db+propagator[1,2,q]*iub+
853 propagator[1,3,q]*dg+propagator[1,4,q]*iug+
854 propagator[1,5,q]*iforce
855 y2=propagator[2,1,q]*db+propagator[2,2,q]*iub+
856 propagator[2,3,q]*dg+propagator[2,4,q]*iug+
857 propagator[2,5,q]*iforce
858 y3=propagator[3,1,q]*db+propagator[3,2,q]*iub+
859 propagator[3,3,q]*dg+propagator[3,4,q]*iug+
860 propagator[3,5,q]*iforce
861 y4=propagator[4,1,q]*db+propagator[4,2,q]*iub+
862 propagator[4,3,q]*dg+propagator[4,4,q]*iug+
863 propagator[4,5,q]*iforce
864 ubn=Complex{T}(cph*imag(y2),-cph*real(y2))
865 ugn=Complex{T}(cph*imag(y4),-cph*real(y4))
866 return y1,ubn,y3,ugn
867 end
868
869 @inline function _photon_split_longitudinal(db,ub,dg,ug,force,dt,G,nu,stress,
870 k,cs2,cph2)
871 T=typeof(dt)
872 iszero(dt) && return db,ub,dg,ug
873 # With vb=i*ub/c_gamma and vg=i*ug/c_gamma the coefficient matrix is real,
874 # balanced, and its exponential acts on complex Fourier amplitudes. The
875 # fifth constant component incorporates the affine magnetic acceleration.
876 cph=sqrt(cph2)
877 invc=inv(cph)
878 iub=Complex{T}(-imag(ub),real(ub))*invc
879 iug=Complex{T}(-imag(ug),real(ug))*invc
880 iforce=Complex{T}(-imag(force),real(force))*invc
881 y=_photon_longitudinal_propagator(dt,G,nu,stress,k,cs2,cph2)*
882 SVector{5,Complex{T}}(db,iub,dg,iug,iforce)
883 ubn=Complex{T}(cph*imag(y[2]),-cph*real(y[2]))
884 ugn=Complex{T}(cph*imag(y[4]),-cph*real(y[4]))
885 return y[1],ubn,y[3],ugn
886 end
887
888 @inline function _photon_moment_transverse(ub,ug,force,dt,G,nu,stress,method)
889 if method == _RADIATION_RESPONSE_MOMENTS_EXPONENTIAL
890 return _photon_split_transverse(ub,ug,force,dt,G,nu,stress)
891 end
892 _photon_sdirk_transverse(ub,ug,force,dt,G,nu,stress)
893 end
894
895 @inline function _photon_moment_longitudinal(db,ub,dg,ug,force,dt,G,nu,stress,
896 k,cs2,cph2,method)
897 if method == _RADIATION_RESPONSE_MOMENTS_EXPONENTIAL
898 return _photon_split_longitudinal(db,ub,dg,ug,force,dt,G,nu,stress,
899 k,cs2,cph2)
900 end
901 _photon_sdirk_longitudinal(db,ub,dg,ug,force,dt,G,nu,stress,k,cs2,cph2)
902 end
903
904 @inline function _photon_baseline_homogeneous(db,ub,dt,G,k,cs2)
905 B,S=_radiation_acoustic_coefficients(G,cs2,k*k,dt)
906 ikdb=Complex{typeof(dt)}(-imag(db),real(db))
907 ikub=Complex{typeof(dt)}(-imag(ub),real(ub))
908 return (B+G*S)*db-k*S*ikub,B*ub-k*cs2*S*ikdb
909 end
910
911 @inline function _photon_baseline_force(force,dt,G,k,cs2)
912 B,S=_radiation_acoustic_coefficients(G,cs2,k*k,dt)
913 J=_radiation_integral_S(B,S,G,cs2*k*k,dt)
914 iforce=Complex{typeof(dt)}(-imag(force),real(force))
915 return -k*J*iforce,S*force
916 end
917
918 """
919 radiation_angular_response(q)
920
921 Static angular transport response for an isotropic photon distribution.
922 `longitudinal` and `transverse` are the photon velocity responses to a baryon
923 velocity perturbation at dimensionless wavenumber `q = k lambda_mfp`.
924 """
925 function radiation_angular_response(q::Real)
926 q >= 0 || error("radiation angular response requires q >= 0")
927 T=typeof(float(q))
928 a0,along,atrans=_radiation_angular_response(T(q)*T(q))
929 (monopole=a0,longitudinal=along,transverse=atrans)
930 end
931
932 function radiation_mode_coefficients(k::Real, gamma_drag::Real,
933 rc::RadiationClosure)
934 T=typeof(rc.inertia_ratio)
935 stream,inertia,gt,gl,cr2=_radiation_mode_coefficients(T(k)^2,T(gamma_drag),rc)
936 (streaming_weight=stream,inertia_response=inertia,gamma_transverse=gt,
937 gamma_longitudinal=gl,radiation_sound_speed2=cr2,
938 response_model=_radiation_response_name(rc.response_model))
939 end
940
941 """
942 radiation_free_streaming_bound(kmin, gamma_drag, closure)
943
944 Conservative dimensionless bound for replacing the spectral gas response by
945 local Compton drag when every resolved mode has `k >= kmin`. The bound is the
946 largest fundamental-mode correction to transverse drag, longitudinal drag, or
947 radiation pressure relative to gas pressure. It is monotone toward zero in the
948 free-streaming limit for the angular and evolved-moment closures.
949 """
950 function radiation_free_streaming_bound(kmin::Real, gamma_drag::Real,
951 rc::RadiationClosure)
952 kmin > 0 || error("free-streaming bound requires kmin > 0")
953 gamma_drag >= 0 || error("free-streaming bound requires non-negative drag")
954 mode=radiation_mode_coefficients(kmin,gamma_drag,rc)
955 drag_t=gamma_drag>0 ? abs(one(gamma_drag)-mode.gamma_transverse/gamma_drag) :
956 zero(gamma_drag)
957 drag_l=gamma_drag>0 ? abs(one(gamma_drag)-mode.gamma_longitudinal/gamma_drag) :
958 zero(gamma_drag)
959 cs2=rc.gas_sound_speed*rc.gas_sound_speed
960 pressure=cs2>0 ? mode.radiation_sound_speed2/cs2 :
961 (mode.radiation_sound_speed2==0 ? zero(cs2) : oftype(cs2,Inf))
962 max(drag_t,drag_l,pressure)
963 end
964
965 @inline function _radiation_mode_index(i::Int32,n::Int32)
966 # A centered spectral derivative has no representable sign at the even-grid
967 # Nyquist mode. Zeroing that component preserves the Hermitian constraints
968 # required by c2r transforms; assigning either sign corrupts self-conjugate
969 # RFFT planes when odd gradients or longitudinal projections are applied.
970 i == (n >>> 1) ? Int32(0) : (i < (n >>> 1) ? i : i-n)
971 end
972
973 @inline function _radiation_mode_dealiased(ii::Int32,jj::Int32,kk::Int32,
974 n::Int32,guard)
975 # The Maxwell force entering the spectral correction is a real-space
976 # product. An optional Nyquist guard keeps aliased grid-axis modes from
977 # feeding back through the exponential/Godunov split.
978 aj=min(jj,n-jj); ak=min(kk,n-kk)
979 lim=guard*typeof(guard)(n>>>1)
980 typeof(guard)(ii) <= lim && typeof(guard)(aj) <= lim &&
981 typeof(guard)(ak) <= lim
982 end
983
984 @kernel function _photon_longitudinal_propagator_k!(propagator,kfac::T,dt::T,
985 G::T,rc::RadiationClosure{T}) where {T}
986 p=@index(Global)
987 mode2=Int32(p-1)
988 k2=kfac*kfac*T(mode2)
989 k=sqrt(k2)
990 nu=rc.inertia_ratio*G
991 _,stress=_photon_stress_rates(k2,nu,rc)
992 cs2=rc.gas_sound_speed*rc.gas_sound_speed
993 cph2=rc.photon_speed*rc.photon_speed
994 matrix=_photon_longitudinal_propagator(dt,G,nu,stress,k,cs2,cph2)
995 @inbounds for column in Int32(1):Int32(5),row in Int32(1):Int32(4)
996 propagator[row,column,p]=matrix[row,column]
997 end
998 end
999
1000 function _radiation_ensure_longitudinal_propagator!(ws::RadiationWorkspace,
1001 s::MHDState,dt,G,rc,kfac)
1002 rc.response_model==_RADIATION_RESPONSE_MOMENTS_EXPONENTIAL || return nothing
1003 signature=hash((dt,G,kfac,rc.inertia_ratio,rc.mean_free_path,
1004 rc.photon_speed,rc.gas_sound_speed,
1005 rc.longitudinal_viscosity))
1006 ws.longitudinal_propagator_signature[]==signature &&
1007 return ws.longitudinal_propagator
1008 _photon_longitudinal_propagator_k!(s.be)(ws.longitudinal_propagator,
1009 kfac,dt,G,rc;ndrange=size(ws.longitudinal_propagator,3))
1010 KA.synchronize(s.be)
1011 ws.longitudinal_propagator_signature[]=signature
1012 ws.longitudinal_propagator
1013 end
1014
1015 @kernel function _radiation_homogeneous_hat_k!(hx,hy,hz,drho,
1016 n::Int32,kfac::T,dt::T,gamma_drag::T,inv_rho_mean::T,
1017 rc::RadiationClosure{T},residual::Val{RES}) where {T,RES}
1018 p=@index(Global)
1019 h=(n>>>1)+Int32(1); q=Int32(p-1)
1020 ii=q%h; q=q÷h; jj=q%n; kk=q÷n
1021 @inbounds begin
1022 kx=kfac*T(_radiation_mode_index(ii,n)); ky=kfac*T(_radiation_mode_index(jj,n))
1023 kz=kfac*T(_radiation_mode_index(kk,n)); k2=kx*kx+ky*ky+kz*kz
1024 if k2 == zero(T) || !_radiation_mode_dealiased(
1025 ii,jj,kk,n,rc.nyquist_guard_fraction)
1026 hx[p]=zero(eltype(hx)); hy[p]=zero(eltype(hy)); hz[p]=zero(eltype(hz))
1027 drho[p]=zero(eltype(drho))
1028 else
1029 k=sqrt(k2); x=hx[p]; y=hy[p]; z=hz[p]
1030 udot=(kx*x+ky*y+kz*z)/k
1031 vlx=kx*udot/k; vly=ky*udot/k; vlz=kz*udot/k
1032 stream,inertia,gt,gl,cr2=_radiation_mode_coefficients(k2,gamma_drag,rc)
1033 ed=exp(-gt*dt); eb=exp(-gamma_drag*dt)
1034 cbase2=rc.gas_sound_speed*rc.gas_sound_speed
1035 Bd,Sd=_radiation_acoustic_coefficients(gl,cbase2+cr2,k2,dt)
1036 Bb,Sb=_radiation_acoustic_coefficients(gamma_drag,cbase2,k2,dt)
1037 delta=drho[p]*inv_rho_mean
1038 idelta=Complex{T}(-imag(delta),real(delta))
1039 uld=Bd*udot-k*(cbase2+cr2)*Sd*idelta
1040 ulb=Bb*udot-k*cbase2*Sb*idelta
1041 du=uld-ulb
1042 hx[p]=(ed-eb)*(x-vlx)+kx*du/k
1043 hy[p]=(ed-eb)*(y-vly)+ky*du/k
1044 hz[p]=(ed-eb)*(z-vlz)+kz*du/k
1045 if RES
1046 Ad=Bd+gl*Sd; Ab=Bb+gamma_drag*Sb
1047 iudot=Complex{T}(-imag(udot),real(udot))
1048 dexact=(Ad-Ab)*delta-k*(Sd-Sb)*iudot
1049 Bdh,Sdh=_radiation_acoustic_coefficients(gl,cbase2+cr2,k2,T(0.5)*dt)
1050 Bbh,Sbh=_radiation_acoustic_coefficients(gamma_drag,cbase2,k2,T(0.5)*dt)
1051 uldh=Bdh*udot-k*(cbase2+cr2)*Sdh*idelta
1052 ulbh=Bbh*udot-k*cbase2*Sbh*idelta
1053 iduh=Complex{T}(-imag(uldh-ulbh),real(uldh-ulbh))
1054 # The Godunov flux already advances continuity with the
1055 # corrected midpoint velocity. Add only the exact residual.
1056 drho[p]=(dexact+k*dt*iduh)/inv_rho_mean
1057 else
1058 drho[p]=zero(eltype(drho))
1059 end
1060 end
1061 end
1062 end
1063
1064 @kernel function _radiation_source_hat_k!(hx,hy,hz,drho,
1065 n::Int32,kfac::T,dt::T,gamma_drag::T,rho_mean::T,
1066 rc::RadiationClosure{T},residual::Val{RES}) where {T,RES}
1067 p=@index(Global)
1068 h=(n>>>1)+Int32(1); q=Int32(p-1)
1069 ii=q%h; q=q÷h; jj=q%n; kk=q÷n
1070 @inbounds begin
1071 kx=kfac*T(_radiation_mode_index(ii,n)); ky=kfac*T(_radiation_mode_index(jj,n))
1072 kz=kfac*T(_radiation_mode_index(kk,n)); k2=kx*kx+ky*ky+kz*kz
1073 if k2 == zero(T) || !_radiation_mode_dealiased(
1074 ii,jj,kk,n,rc.nyquist_guard_fraction)
1075 hx[p]=zero(eltype(hx)); hy[p]=zero(eltype(hy)); hz[p]=zero(eltype(hz))
1076 drho[p]=zero(eltype(drho))
1077 else
1078 k=sqrt(k2); x=hx[p]; y=hy[p]; z=hz[p]
1079 adot=(kx*x+ky*y+kz*z)/k
1080 alx=kx*adot/k; aly=ky*adot/k; alz=kz*adot/k
1081 stream,inertia,gt,gl,cr2=_radiation_mode_coefficients(k2,gamma_drag,rc)
1082 phid=dt*_drag_phi1(gt*dt); phib=dt*_drag_phi1(gamma_drag*dt)
1083 cbase2=rc.gas_sound_speed*rc.gas_sound_speed
1084 Bd,Sd=_radiation_acoustic_coefficients(gl,cbase2+cr2,k2,dt)
1085 Bb,Sb=_radiation_acoustic_coefficients(gamma_drag,cbase2,k2,dt)
1086 ct=inertia*phid-phib; cl=inertia*Sd-Sb
1087 hx[p]=ct*(x-alx)+kx*(cl*adot)/k
1088 hy[p]=ct*(y-aly)+ky*(cl*adot)/k
1089 hz[p]=ct*(z-alz)+kz*(cl*adot)/k
1090 if RES
1091 wd=(cbase2+cr2)*k2; wb=cbase2*k2
1092 Jd=_radiation_integral_S(Bd,Sd,gl,wd,dt)
1093 Jb=_radiation_integral_S(Bb,Sb,gamma_drag,wb,dt)
1094 iadot=Complex{T}(-imag(adot),real(adot))
1095 dexact=-k*(inertia*Jd-Jb)*iadot
1096 Bdh,Sdh=_radiation_acoustic_coefficients(gl,cbase2+cr2,k2,T(0.5)*dt)
1097 Bbh,Sbh=_radiation_acoustic_coefficients(gamma_drag,cbase2,k2,T(0.5)*dt)
1098 iduh=Complex{T}(-imag((inertia*Sdh-Sbh)*adot),
1099 real((inertia*Sdh-Sbh)*adot))
1100 drho[p]=rho_mean*(dexact+k*dt*iduh)
1101 else
1102 drho[p]=zero(eltype(drho))
1103 end
1104 end
1105 end
1106 end
1107
1108 @kernel function _photon_moment_homogeneous_hat_k!(hx,hy,hz,drho,
1109 @Const(propagator),pgx,pgy,pgz,pdelta,n::Int32,kfac::T,dt::T,G::T,
1110 inv_rho_mean::T,
1111 rc::RadiationClosure{T},residual::Val{RES},update::Val{UPDATE}) where {T,RES,UPDATE}
1112 p=@index(Global)
1113 h=(n>>>1)+Int32(1); q=Int32(p-1)
1114 ii=q%h; q=q÷h; jj=q%n; kk=q÷n
1115 @inbounds begin
1116 mi=_radiation_mode_index(ii,n); mj=_radiation_mode_index(jj,n)
1117 mk=_radiation_mode_index(kk,n); mode2=mi*mi+mj*mj+mk*mk
1118 kx=kfac*T(mi); ky=kfac*T(mj); kz=kfac*T(mk)
1119 k2=kfac*kfac*T(mode2)
1120 meanmode=ii==0 && jj==0 && kk==0
1121 if !_radiation_mode_dealiased(ii,jj,kk,n,rc.nyquist_guard_fraction) ||
1122 (k2==zero(T) && !meanmode)
1123 hx[p]=zero(eltype(hx)); hy[p]=zero(eltype(hy)); hz[p]=zero(eltype(hz))
1124 drho[p]=zero(eltype(drho))
1125 elseif meanmode
1126 eb=exp(-G*dt); nu=rc.inertia_ratio*G
1127 xb,xg=_photon_moment_transverse(hx[p],pgx[p],zero(eltype(hx)),dt,G,nu,
1128 zero(T),rc.response_model)
1129 yb,yg=_photon_moment_transverse(hy[p],pgy[p],zero(eltype(hy)),dt,G,nu,
1130 zero(T),rc.response_model)
1131 zb,zg=_photon_moment_transverse(hz[p],pgz[p],zero(eltype(hz)),dt,G,nu,
1132 zero(T),rc.response_model)
1133 hx[p]=xb-eb*hx[p]; hy[p]=yb-eb*hy[p]; hz[p]=zb-eb*hz[p]
1134 drho[p]=zero(eltype(drho))
1135 if UPDATE
1136 pgx[p]=xg; pgy[p]=yg; pgz[p]=zg
1137 end
1138 else
1139 k=sqrt(k2); x=hx[p]; y=hy[p]; z=hz[p]
1140 gx=pgx[p]; gy=pgy[p]; gz=pgz[p]; gd=pdelta[p]
1141 ub=(kx*x+ky*y+kz*z)/k; ug=(kx*gx+ky*gy+kz*gz)/k
1142 blx=kx*ub/k; bly=ky*ub/k; blz=kz*ub/k
1143 glx=kx*ug/k; gly=ky*ug/k; glz=kz*ug/k
1144 db=drho[p]*inv_rho_mean
1145 nu=rc.inertia_ratio*G; st,sl=_photon_stress_rates(k2,nu,rc)
1146 cs2=rc.gas_sound_speed*rc.gas_sound_speed
1147 cph2=rc.photon_speed*rc.photon_speed
1148 if rc.response_model==_RADIATION_RESPONSE_MOMENTS_EXPONENTIAL
1149 dbn,ubn,gdn,ugn=_photon_apply_longitudinal_propagator(
1150 propagator,mode2,db,ub,gd,ug,zero(eltype(ub)),
1151 rc.photon_speed)
1152 else
1153 dbn,ubn,gdn,ugn=_photon_moment_longitudinal(db,ub,gd,ug,
1154 zero(eltype(ub)),dt,G,nu,sl,k,cs2,cph2,rc.response_model)
1155 end
1156 tx,tgx=_photon_moment_transverse(x-blx,gx-glx,zero(eltype(x)),dt,G,nu,
1157 st,rc.response_model)
1158 ty,tgy=_photon_moment_transverse(y-bly,gy-gly,zero(eltype(y)),dt,G,nu,
1159 st,rc.response_model)
1160 tz,tgz=_photon_moment_transverse(z-blz,gz-glz,zero(eltype(z)),dt,G,nu,
1161 st,rc.response_model)
1162 dbb,ubb=_photon_baseline_homogeneous(db,ub,dt,G,k,cs2)
1163 ebase=exp(-G*dt)
1164 hx[p]=kx*ubn/k+tx-(kx*ubb/k+ebase*(x-blx))
1165 hy[p]=ky*ubn/k+ty-(ky*ubb/k+ebase*(y-bly))
1166 hz[p]=kz*ubn/k+tz-(kz*ubb/k+ebase*(z-blz))
1167 if RES
1168 if rc.response_model == _RADIATION_RESPONSE_MOMENTS_EXPONENTIAL
1169 drho[p]=dbn/inv_rho_mean
1170 else
1171 _,ubh,_,_=_photon_moment_longitudinal(db,ub,gd,ug,
1172 zero(eltype(ub)),T(0.5)*dt,G,nu,sl,k,cs2,cph2,
1173 rc.response_model)
1174 _,ubbh=_photon_baseline_homogeneous(db,ub,T(0.5)*dt,G,k,cs2)
1175 iduh=Complex{T}(-imag(ubh-ubbh),real(ubh-ubbh))
1176 drho[p]=(dbn-dbb+k*dt*iduh)/inv_rho_mean
1177 end
1178 else
1179 drho[p]=zero(eltype(drho))
1180 end
1181 if UPDATE
1182 pgx[p]=kx*ugn/k+tgx; pgy[p]=ky*ugn/k+tgy; pgz[p]=kz*ugn/k+tgz
1183 pdelta[p]=gdn
1184 end
1185 end
1186 end
1187 end
1188
1189 @kernel function _photon_moment_source_hat_k!(hx,hy,hz,drho,
1190 @Const(propagator),pgx,pgy,pgz,pdelta,n::Int32,kfac::T,dt::T,G::T,
1191 rho_mean::T,
1192 rc::RadiationClosure{T},residual::Val{RES},update::Val{UPDATE}) where {T,RES,UPDATE}
1193 p=@index(Global)
1194 h=(n>>>1)+Int32(1); q=Int32(p-1)
1195 ii=q%h; q=q÷h; jj=q%n; kk=q÷n
1196 @inbounds begin
1197 mi=_radiation_mode_index(ii,n); mj=_radiation_mode_index(jj,n)
1198 mk=_radiation_mode_index(kk,n); mode2=mi*mi+mj*mj+mk*mk
1199 kx=kfac*T(mi); ky=kfac*T(mj); kz=kfac*T(mk)
1200 k2=kfac*kfac*T(mode2)
1201 meanmode=ii==0 && jj==0 && kk==0
1202 if !_radiation_mode_dealiased(ii,jj,kk,n,rc.nyquist_guard_fraction) ||
1203 (k2==zero(T) && !meanmode)
1204 hx[p]=zero(eltype(hx)); hy[p]=zero(eltype(hy)); hz[p]=zero(eltype(hz))
1205 drho[p]=zero(eltype(drho))
1206 elseif meanmode
1207 nu=rc.inertia_ratio*G; phib=dt*_drag_phi1(G*dt)
1208 xb,xg=_photon_moment_transverse(zero(hx[p]),zero(pgx[p]),hx[p],dt,G,nu,
1209 zero(T),rc.response_model)
1210 yb,yg=_photon_moment_transverse(zero(hy[p]),zero(pgy[p]),hy[p],dt,G,nu,
1211 zero(T),rc.response_model)
1212 zb,zg=_photon_moment_transverse(zero(hz[p]),zero(pgz[p]),hz[p],dt,G,nu,
1213 zero(T),rc.response_model)
1214 hx[p]=xb-phib*hx[p]; hy[p]=yb-phib*hy[p]; hz[p]=zb-phib*hz[p]
1215 drho[p]=zero(eltype(drho))
1216 if UPDATE
1217 pgx[p]+=xg; pgy[p]+=yg; pgz[p]+=zg
1218 end
1219 else
1220 k=sqrt(k2); ax=hx[p]; ay=hy[p]; az=hz[p]
1221 al=(kx*ax+ky*ay+kz*az)/k
1222 alx=kx*al/k; aly=ky*al/k; alz=kz*al/k
1223 nu=rc.inertia_ratio*G; st,sl=_photon_stress_rates(k2,nu,rc)
1224 cs2=rc.gas_sound_speed*rc.gas_sound_speed
1225 cph2=rc.photon_speed*rc.photon_speed
1226 if rc.response_model==_RADIATION_RESPONSE_MOMENTS_EXPONENTIAL
1227 dbn,ubn,gdn,ugn=_photon_apply_longitudinal_propagator(
1228 propagator,mode2,zero(al),zero(al),zero(al),zero(al),al,
1229 rc.photon_speed)
1230 else
1231 dbn,ubn,gdn,ugn=_photon_moment_longitudinal(zero(al),zero(al),
1232 zero(al),zero(al),al,dt,G,nu,sl,k,cs2,cph2,
1233 rc.response_model)
1234 end
1235 tx,tgx=_photon_moment_transverse(zero(ax),zero(ax),ax-alx,dt,G,nu,st,
1236 rc.response_model)
1237 ty,tgy=_photon_moment_transverse(zero(ay),zero(ay),ay-aly,dt,G,nu,st,
1238 rc.response_model)
1239 tz,tgz=_photon_moment_transverse(zero(az),zero(az),az-alz,dt,G,nu,st,
1240 rc.response_model)
1241 dbb,ubb=_photon_baseline_force(al,dt,G,k,cs2)
1242 phib=dt*_drag_phi1(G*dt)
1243 hx[p]=kx*ubn/k+tx-(kx*ubb/k+phib*(ax-alx))
1244 hy[p]=ky*ubn/k+ty-(ky*ubb/k+phib*(ay-aly))
1245 hz[p]=kz*ubn/k+tz-(kz*ubb/k+phib*(az-alz))
1246 if RES
1247 if rc.response_model == _RADIATION_RESPONSE_MOMENTS_EXPONENTIAL
1248 drho[p]=rho_mean*dbn
1249 else
1250 dbh,ubh,_,_=_photon_moment_longitudinal(zero(al),zero(al),
1251 zero(al),zero(al),al,T(0.5)*dt,G,nu,sl,k,cs2,cph2,
1252 rc.response_model)
1253 dbbh,ubbh=_photon_baseline_force(al,T(0.5)*dt,G,k,cs2)
1254 iduh=Complex{T}(-imag(ubh-ubbh),real(ubh-ubbh))
1255 drho[p]=rho_mean*((dbn-dbb)+k*dt*iduh)
1256 end
1257 else
1258 drho[p]=zero(eltype(drho))
1259 end
1260 if UPDATE
1261 pgx[p]+=kx*ugn/k+tgx; pgy[p]+=ky*ugn/k+tgy; pgz[p]+=kz*ugn/k+tgz
1262 pdelta[p]+=gdn
1263 end
1264 end
1265 end
1266 end
1267
1268 @kernel function _radiation_prepare_velocity_k!(vx,vy,vz,
1269 @Const(rho),@Const(mx),@Const(my),@Const(mz),smallr)
1270 c=@index(Global)
1271 @inbounds begin
1272 r=max(rho[c],smallr); vx[c]=mx[c]/r; vy[c]=my[c]/r; vz[c]=mz[c]/r
1273 end
1274 end
1275
1276 @kernel function _radiation_commit_density_hat_k!(density_hat,@Const(update_hat),
1277 blend,additive::Val{ADDITIVE}) where {ADDITIVE}
1278 p=@index(Global)
1279 @inbounds density_hat[p]=ADDITIVE ?
1280 density_hat[p]+blend*update_hat[p] :
1281 density_hat[p]+blend*(update_hat[p]-density_hat[p])
1282 end
1283
1284 @kernel function _radiation_density_hat_increment_k!(update_hat,
1285 @Const(density_hat))
1286 p=@index(Global)
1287 @inbounds update_hat[p]-=density_hat[p]
1288 end
1289
1290 function _radiation_ensure_density_hat!(ws::RadiationWorkspace,s::MHDState)
1291 ws.density_hat_valid[] && return ws.authoritative_density_hat
1292 copyto!(ws.fft_real,ws.density_perturbation)
1293 _radiation_forward!(ws.authoritative_density_hat,ws.fft_real)
1294 ws.density_hat_valid[]=true
1295 ws.authoritative_density_hat
1296 end
1297
1298 function _radiation_commit_density_hat!(ws::RadiationWorkspace,s::MHDState;
1299 additive::Bool,density_blend=1)
1300 _radiation_commit_density_hat_k!(s.be)(ws.authoritative_density_hat,
1301 ws.density_hat,eltype(s.U[1])(density_blend),Val(additive);
1302 ndrange=length(ws.density_hat))
1303 KA.synchronize(s.be)
1304 ws.density_hat_valid[]=true
1305 ws.authoritative_density_hat
1306 end
1307
1308 @kernel function _radiation_force_to_accel_k!(fx,fy,fz,@Const(rho),smallr)
1309 c=@index(Global)
1310 @inbounds begin
1311 ir=inv(max(rho[c],smallr)); fx[c]*=ir; fy[c]*=ir; fz[c]*=ir
1312 end
1313 end
1314
1315 @kernel function _radiation_magnetic_accel_pair_k!(fx,fy,fz,
1316 @Const(rho0),@Const(Bx0),@Const(By0),@Const(Bz0),
1317 @Const(rho1),@Const(Bx1),@Const(By1),@Const(Bz1),
1318 n::Int,inv2dx,smallr)
1319 c=@index(Global)
1320 @inbounds begin
1321 i=(c-1)%n+1; j=((c-1)÷n)%n+1; k=(c-1)÷(n*n)+1
1322 x0,y0,z0=_tct_magnetic_force_at(Bx0,By0,Bz0,i,j,k,n,inv2dx)
1323 x1,y1,z1=_tct_magnetic_force_at(Bx1,By1,Bz1,i,j,k,n,inv2dx)
1324 h=typeof(inv2dx)(0.5)
1325 ir0=inv(max(rho0[c],smallr)); ir1=inv(max(rho1[c],smallr))
1326 fx[c]=h*(x0*ir0+x1*ir1)
1327 fy[c]=h*(y0*ir0+y1*ir1)
1328 fz[c]=h*(z0*ir0+z1*ir1)
1329 end
1330 end
1331
1332 @kernel function _radiation_add3_k!(x,y,z,@Const(ax),@Const(ay),@Const(az))
1333 c=@index(Global)
1334 @inbounds begin x[c]+=ax[c]; y[c]+=ay[c]; z[c]+=az[c] end
1335 end
1336
1337 @kernel function _radiation_speed2_k!(work,@Const(x),@Const(y),@Const(z))
1338 c=@index(Global)
1339 @inbounds work[c]=x[c]*x[c]+y[c]*y[c]+z[c]*z[c]
1340 end
1341
1342 @kernel function _radiation_apply_state_correction_k!(rho,mx,my,mz,E,
1343 @Const(Bx),@Const(By),@Const(Bz),
1344 @Const(drho),@Const(cx),@Const(cy),@Const(cz),rho_perturbation,
1345 energy_residual,
1346 rho_mean,gamma,smallr,smallE,density_blend,
1347 absolute_density::Val{ABS}) where {ABS}
1348 c=@index(Global)
1349 @inbounds begin
1350 T=eltype(E); r=max(rho[c],smallr)
1351 delta0=rho_perturbation[c]
1352 target_delta=ABS ? drho[c] : delta0+drho[c]
1353 deltan=density_blend==one(T) ? target_delta :
1354 delta0+density_blend*(target_delta-delta0)
1355 raw_rn=rho_mean+deltan
1356 rn=max(raw_rn,smallr)
1357 # `rho` is only a compatibility field. Preserve the authoritative
1358 # perturbation directly so increments below one ULP of rho_mean are not
1359 # lost when the compatibility density is reconstructed in Float32.
1360 rho_perturbation[c]=raw_rn==rn ? deltan : rn-rho_mean
1361 vx=mx[c]/r+density_blend*cx[c]
1362 vy=my[c]/r+density_blend*cy[c]
1363 vz=mz[c]/r+density_blend*cz[c]
1364 ek0=T(0.5)*(mx[c]^2+my[c]^2+mz[c]^2)/r
1365 eb=T(0.5)*(Bx[c]^2+By[c]^2+Bz[c]^2)
1366 eint=max(E[c]-ek0-eb,smallE)
1367 eintn=eint*(rn/r)^gamma
1368 nmx=rn*vx; nmy=rn*vy; nmz=rn*vz
1369 ek1=T(0.5)*(nmx*nmx+nmy*nmy+nmz*nmz)/rn
1370 rho[c]=rn; mx[c]=nmx; my[c]=nmy; mz[c]=nmz
1371 de=(eintn-eint)+(ek1-ek0)
1372 E[c]=max(_compensated_update!(energy_residual,c,E[c],de,Val(true)),
1373 smallE)
1374 end
1375 end
1376
1377 function _radiation_transform_velocity_density!(ws::RadiationWorkspace,s::MHDState,
1378 state,dt,gamma_drag,rc;
1379 density_residual=false,
1380 photons=nothing,
1381 update_photons=false,
1382 rho_mean=nothing)
1383 x,y,z=ws.correction
1384 _radiation_prepare_velocity_k!(s.be)(x,y,z,state[1],state[2],state[3],state[4],
1385 s.smallr;ndrange=ncells(s))
1386 KA.synchronize(s.be)
1387 persistent_density=rc.response_model==_RADIATION_RESPONSE_MOMENTS_EXPONENTIAL
1388 if persistent_density
1389 _radiation_ensure_density_hat!(ws,s)
1390 fill!(ws.fft_real,zero(eltype(ws.fft_real)))
1391 else
1392 copyto!(ws.fft_real,ws.density_perturbation)
1393 end
1394 _radiation_forward_batch!(ws.hat_batch,ws.real_batch)
1395 persistent_density &&
1396 copyto!(ws.density_hat,ws.authoritative_density_hat)
1397 n=s.dims[1]; T=eltype(x); kfac=T(2pi)/(T(n)*s.dx)
1398 meanrho=rho_mean === nothing ? ws.rho_mean : T(rho_mean)
1399 if _radiation_evolves_moments(rc.response_model)
1400 photons isa PhotonMomentState ||
1401 error("response_model=:moments requires a PhotonMomentState")
1402 _radiation_ensure_longitudinal_propagator!(ws,s,T(dt),T(gamma_drag),rc,
1403 kfac)
1404 _photon_moment_homogeneous_hat_k!(s.be)(ws.velocity_hat...,ws.density_hat,
1405 ws.longitudinal_propagator,photons.velocity_hat...,
1406 photons.delta_hat,Int32(n),kfac,T(dt),T(gamma_drag),inv(meanrho),rc,
1407 Val(density_residual),Val(update_photons);
1408 ndrange=length(ws.density_hat))
1409 else
1410 _radiation_homogeneous_hat_k!(s.be)(ws.velocity_hat...,ws.density_hat,
1411 Int32(n),kfac,T(dt),T(gamma_drag),inv(meanrho),rc,
1412 Val(density_residual);ndrange=length(ws.density_hat))
1413 end
1414 KA.synchronize(s.be)
1415 if persistent_density && density_residual
1416 # The exponential moment solve returns the absolute density mode. Only
1417 # inverse-transform its increment: replacing the real density with an
1418 # absolute MPS inverse every step feeds the transform's fixed Float32
1419 # gain error back into the next Godunov flux and destroys dt convergence.
1420 _radiation_density_hat_increment_k!(s.be)(ws.density_hat,
1421 ws.authoritative_density_hat;ndrange=length(ws.density_hat))
1422 KA.synchronize(s.be)
1423 # Preserve the increment: the c2r inverse below may destroy it (CPU).
1424 ws.density_increment_hat===ws.authoritative_density_hat ||
1425 copyto!(ws.density_increment_hat,ws.density_hat)
1426 end
1427 _radiation_inverse_triplet!(ws.correction,ws)
1428 # Restore the increment for `_radiation_commit_density_hat!` (CPU c2r).
1429 persistent_density && density_residual &&
1430 ws.density_increment_hat!==ws.authoritative_density_hat &&
1431 copyto!(ws.density_hat,ws.density_increment_hat)
1432 return ws.correction
1433 end
1434
1435 function _radiation_transform_force!(out,ws::RadiationWorkspace,s::MHDState,state,
1436 dt,gamma_drag,rc;density_residual=false,
1437 photons=nothing,update_photons=false,
1438 corrected_state=nothing,rho_mean=nothing)
1439 fx,fy,fz=out; n=s.dims[1]; inv2dx=eltype(fx)(1)/(eltype(fx)(2)*s.dx)
1440 if corrected_state === nothing
1441 _tct_magnetic_force_k!(s.be)(fx,fy,fz,state[6],state[7],state[8],n,inv2dx;
1442 ndrange=ncells(s))
1443 _radiation_force_to_accel_k!(s.be)(fx,fy,fz,state[1],s.smallr;
1444 ndrange=ncells(s))
1445 else
1446 _radiation_magnetic_accel_pair_k!(s.be)(fx,fy,fz,
1447 state[1],state[6],state[7],state[8],
1448 corrected_state[1],corrected_state[6],corrected_state[7],
1449 corrected_state[8],n,inv2dx,s.smallr;ndrange=ncells(s))
1450 end
1451 KA.synchronize(s.be)
1452 _radiation_forward_triplet!(ws,out)
1453 T=eltype(fx); kfac=T(2pi)/(T(n)*s.dx)
1454 meanrho=rho_mean === nothing ? ws.rho_mean : T(rho_mean)
1455 if _radiation_evolves_moments(rc.response_model)
1456 photons isa PhotonMomentState ||
1457 error("response_model=:moments requires a PhotonMomentState")
1458 _radiation_ensure_longitudinal_propagator!(ws,s,T(dt),T(gamma_drag),rc,
1459 kfac)
1460 _photon_moment_source_hat_k!(s.be)(ws.velocity_hat...,ws.density_hat,
1461 ws.longitudinal_propagator,photons.velocity_hat...,
1462 photons.delta_hat,Int32(n),kfac,T(dt),T(gamma_drag),meanrho,rc,
1463 Val(density_residual),Val(update_photons);
1464 ndrange=length(ws.velocity_hat[1]))
1465 else
1466 _radiation_source_hat_k!(s.be)(ws.velocity_hat...,ws.density_hat,
1467 Int32(n),kfac,T(dt),T(gamma_drag),meanrho,rc,
1468 Val(density_residual);ndrange=length(ws.velocity_hat[1]))
1469 end
1470 KA.synchronize(s.be)
1471 _radiation_inverse_triplet!(out,ws)
1472 out
1473 end
1474
1475 function _radiation_apply_correction!(s,ws,corr;density_absolute=false,
1476 density_blend=1)
1477 _radiation_apply_state_correction_k!(s.be)(s.U[1],s.U[2],s.U[3],s.U[4],s.U[5],
1478 s.U[6],s.U[7],s.U[8],ws.fft_real,corr...,ws.density_perturbation,
1479 ws.energy_residual,ws.rho_mean,s.γ,s.smallr,eltype(s.U[1])(1e-30),
1480 eltype(s.U[1])(density_blend),Val(density_absolute);
1481 ndrange=ncells(s))
1482 KA.synchronize(s.be)
1483 end
1484
1485 function _radiation_density_blend(ws::RadiationWorkspace,s::MHDState;
1486 absolute::Bool)
1487 gate=max(Float64(s.smallr),Float64(s.llf_dmin))
1488 rho_min=Float64(minimum(s.U[1]))
1489 delta_min=Float64(minimum(ws.fft_real))
1490 target_min=absolute ? Float64(ws.rho_mean)+delta_min : rho_min+delta_min
1491 if isfinite(rho_min) && isfinite(target_min) && target_min>=gate
1492 return 1.0
1493 end
1494 denom=rho_min-target_min
1495 denom>0 && isfinite(denom) || return 0.0
1496 clamp((rho_min-gate)/denom,0.0,1.0)
1497 end
1498
1499 @kernel function _radiation_internal_energy_k!(out,@Const(rho),@Const(mx),
1500 @Const(my),@Const(mz),@Const(E),@Const(Bx),@Const(By),@Const(Bz),smallr)
1501 c=@index(Global)
1502 @inbounds begin
1503 r=max(rho[c],smallr)
1504 ek=eltype(out)(0.5)*(mx[c]^2+my[c]^2+mz[c]^2)/r
1505 eb=eltype(out)(0.5)*(Bx[c]^2+By[c]^2+Bz[c]^2)
1506 out[c]=E[c]-ek-eb
1507 end
1508 end
1509
1510 function _radiation_state_bounds!(ws::RadiationWorkspace,s::MHDState)
1511 _radiation_internal_energy_k!(s.be)(ws.fft_real,s.U[1],s.U[2],s.U[3],
1512 s.U[4],s.U[5],s.U[6],s.U[7],s.U[8],s.smallr;ndrange=ncells(s))
1513 KA.synchronize(s.be)
1514 (rho_min=Float64(minimum(s.U[1])),rho_max=Float64(maximum(s.U[1])),
1515 eint_min=Float64(minimum(ws.fft_real)),
1516 eint_max=Float64(maximum(ws.fft_real)))
1517 end
1518
1519 # CT (:ref) trial admissibility. The :cube kernel flags a cell inadmissible when rho is
1520 # below the LLF gate, any field is non-finite, or the internal energy E-ek-eb falls below
1521 # pfl/(gamma-1); the host then sees a negative density. The :ref kernel has no such flag, so
1522 # CT evaluates the same conditions on the stored trial state (cell B = face average, i.e.
1523 # exactly what the next step reconstructs from): density gate on the host, finite
1524 # internal energy and eint >= pfl/(gamma-1) in one device pass over `ws.fft_real`. GLM
1525 # keeps the single-argument density check (its :cube kernel already encodes the rest).
1526 function _radiation_trial_admissible(s::MHDState,ws::RadiationWorkspace)
1527 _radiation_trial_admissible(s) || return false
1528 uses_ct(s) || return true
1529 _radiation_internal_energy_k!(s.be)(ws.fft_real,s.U[1],s.U[2],s.U[3],
1530 s.U[4],s.U[5],s.U[6],s.U[7],s.U[8],s.smallr;ndrange=ncells(s))
1531 KA.synchronize(s.be)
1532 eint_min=Float64(minimum(ws.fft_real))
1533 isfinite(eint_min) && eint_min>=Float64(s.pfl)/(Float64(s.γ)-1)
1534 end
1535
1536 @kernel function _radiation_fallback_mask_k!(mask,@Const(trial_rho),N)
1537 c=@index(Global)
1538 @inbounds begin
1539 i=(c-1)%N+1; j=((c-1)÷N)%N+1; k=(c-1)÷(N*N)+1
1540 im=i==1 ? N : i-1; ip=i==N ? 1 : i+1
1541 jm=j==1 ? N : j-1; jp=j==N ? 1 : j+1
1542 km=k==1 ? N : k-1; kp=k==N ? 1 : k+1
1543 bad=trial_rho[c]<0 ||
1544 trial_rho[((k-1)*N+(j-1))*N+im]<0 ||
1545 trial_rho[((k-1)*N+(j-1))*N+ip]<0 ||
1546 trial_rho[((k-1)*N+(jm-1))*N+i]<0 ||
1547 trial_rho[((k-1)*N+(jp-1))*N+i]<0 ||
1548 trial_rho[((km-1)*N+(j-1))*N+i]<0 ||
1549 trial_rho[((kp-1)*N+(j-1))*N+i]<0
1550 mask[c]=bad ? one(eltype(mask)) : zero(eltype(mask))
1551 end
1552 end
1553
1554 @kernel function _radiation_expand_fallback_mask_k!(out,@Const(mask),
1555 @Const(trial_rho),N)
1556 c=@index(Global)
1557 @inbounds begin
1558 i=(c-1)%N+1; j=((c-1)÷N)%N+1; k=(c-1)÷(N*N)+1
1559 im=i==1 ? N : i-1; ip=i==N ? 1 : i+1
1560 jm=j==1 ? N : j-1; jp=j==N ? 1 : j+1
1561 km=k==1 ? N : k-1; kp=k==N ? 1 : k+1
1562 bad=mask[c]>0 || trial_rho[c]<0 ||
1563 mask[((k-1)*N+(j-1))*N+im]>0 ||
1564 mask[((k-1)*N+(j-1))*N+ip]>0 ||
1565 mask[((k-1)*N+(jm-1))*N+i]>0 ||
1566 mask[((k-1)*N+(jp-1))*N+i]>0 ||
1567 mask[((km-1)*N+(j-1))*N+i]>0 ||
1568 mask[((kp-1)*N+(j-1))*N+i]>0 ||
1569 trial_rho[((k-1)*N+(j-1))*N+im]<0 ||
1570 trial_rho[((k-1)*N+(j-1))*N+ip]<0 ||
1571 trial_rho[((k-1)*N+(jm-1))*N+i]<0 ||
1572 trial_rho[((k-1)*N+(jp-1))*N+i]<0 ||
1573 trial_rho[((km-1)*N+(j-1))*N+i]<0 ||
1574 trial_rho[((kp-1)*N+(j-1))*N+i]<0
1575 out[c]=bad ? one(eltype(out)) : zero(eltype(out))
1576 end
1577 end
1578
1579 function _radiation_prepare_fallback_mask!(ws::RadiationWorkspace,s::MHDState)
1580 _radiation_fallback_mask_k!(s.be)(ws.density_perturbation,s.U[1],s.dims[1];
1581 ndrange=ncells(s))
1582 KA.synchronize(s.be)
1583 nothing
1584 end
1585
1586
1587 function _radiation_expand_fallback_mask!(ws::RadiationWorkspace,s::MHDState)
1588 _radiation_expand_fallback_mask_k!(s.be)(ws.fft_real,
1589 ws.density_perturbation,s.U[1],s.dims[1];ndrange=ncells(s))
1590 KA.synchronize(s.be)
1591 copyto!(ws.density_perturbation,ws.fft_real)
1592 # The first packed segment normally carries the original GLM psi state.
1593 copyto!(view(ws.packed_psi_correction,1:ncells(s)),s.scratch[9])
1594 nothing
1595 end
1596
1597 function _radiation_debug_rho(s, stage, threshold, context)
1598 threshold > 0 || return nothing
1599 rho_min = Float64(minimum(s.U[1]))
1600 rho_max = Float64(maximum(s.U[1]))
1601 e_min = Float64(minimum(s.U[5]))
1602 e_max = Float64(maximum(s.U[5]))
1603 if !(isfinite(rho_min) && isfinite(rho_max) && isfinite(e_min) &&
1604 isfinite(e_max) && rho_min >= threshold && e_min > 0)
1605 error("radiation stage=$stage $context crossed state gate: " *
1606 "rho=[$rho_min,$rho_max] E=[$e_min,$e_max] " *
1607 "rho_threshold=$threshold")
1608 end
1609 return nothing
1610 end
1611
1612 function _radiation_correction_cfl(ws,s,dt,correction)
1613 _radiation_speed2_k!(s.be)(ws.fft_real,correction...;ndrange=ncells(s))
1614 KA.synchronize(s.be)
1615 vmax=sqrt(max(0.0,Float64(maximum(ws.fft_real))))
1616 return vmax,Float64(dt)*vmax/Float64(s.dx)
1617 end
1618
1619 function _update_free_streaming_photons_velocity!(photons::PhotonMomentState,
1620 ws::RadiationWorkspace,
1621 s::MHDState{T},state,dt::T,
1622 gamma_drag::T,
1623 rc::RadiationClosure{T},
1624 rho_mean::T) where {T}
1625 x,y,z=ws.correction
1626 _radiation_prepare_velocity_k!(s.be)(x,y,z,state[1],state[2],state[3],state[4],
1627 s.smallr;ndrange=ncells(s))
1628 KA.synchronize(s.be)
1629 copyto!(ws.fft_real,reshape(state[1],s.dims))
1630 _radiation_forward_batch!(ws.hat_batch,ws.real_batch)
1631 n=s.dims[1]
1632 kfac=T(2pi)/(T(n)*s.dx)
1633 _radiation_ensure_longitudinal_propagator!(ws,s,dt,gamma_drag,rc,kfac)
1634 _photon_moment_homogeneous_hat_k!(s.be)(
1635 ws.velocity_hat...,ws.density_hat,ws.longitudinal_propagator,
1636 photons.velocity_hat...,photons.delta_hat,
1637 Int32(n),kfac,dt,gamma_drag,inv(rho_mean),rc,Val(true),Val(true);
1638 ndrange=length(ws.density_hat))
1639 KA.synchronize(s.be)
1640 photons
1641 end
1642
1643 function _update_free_streaming_photons_force!(photons::PhotonMomentState,
1644 ws::RadiationWorkspace,
1645 s::MHDState{T},state,dt::T,
1646 gamma_drag::T,
1647 rc::RadiationClosure{T},
1648 rho_mean::T;
1649 corrected_state) where {T}
1650 fx,fy,fz=ws.correction
1651 n=s.dims[1]
1652 inv2dx=T(0.5)/s.dx
1653 _radiation_magnetic_accel_pair_k!(s.be)(fx,fy,fz,
1654 state[1],state[6],state[7],state[8],
1655 corrected_state[1],corrected_state[6],corrected_state[7],
1656 corrected_state[8],n,inv2dx,s.smallr;ndrange=ncells(s))
1657 KA.synchronize(s.be)
1658 _radiation_forward_triplet!(ws,ws.correction)
1659 kfac=T(2pi)/(T(n)*s.dx)
1660 _radiation_ensure_longitudinal_propagator!(ws,s,dt,gamma_drag,rc,kfac)
1661 _photon_moment_source_hat_k!(s.be)(
1662 ws.velocity_hat...,ws.density_hat,ws.longitudinal_propagator,
1663 photons.velocity_hat...,photons.delta_hat,
1664 Int32(n),kfac,dt,gamma_drag,rho_mean,rc,Val(true),Val(true);
1665 ndrange=length(ws.density_hat))
1666 KA.synchronize(s.be)
1667 photons
1668 end
1669
1670 function _step_radiation_free_streaming!(s::MHDState{T},ws,dt::T;
1671 drag_impulse::T,rc::RadiationClosure{T},ch::T,glm_cr::Real,
1672 integrator::Symbol,energy_ledger,timing,photon_state) where {T}
1673 photon_before=0.0
1674 if energy_ledger!==nothing
1675 _time_radiation_stage!(timing,:ledger_s) do
1676 photon_before=photon_moment_energy(photon_state,ws,s,rc)
1677 end
1678 end
1679 _time_radiation_stage!(timing,:godunov_s) do
1680 fill!(ws.correction[1],zero(T)); fill!(ws.correction[2],zero(T))
1681 fill!(ws.correction[3],zero(T))
1682 copyto!(view(ws.packed_psi_correction,1:ncells(s)),s.U[9])
1683 decay=exp(-T(glm_cr)*ch*dt/s.dx)
1684 if integrator===:ref
1685 step_ref_radiation!(s,dt;ch=ch,decay=decay,
1686 drag_impulse=drag_impulse,
1687 correction=ws.packed_psi_correction,
1688 track_magnetic=ws.magnetic_residual!==nothing)
1689 elseif integrator===:cube
1690 step_cube_radiation!(s,dt;ch=ch,decay=decay,
1691 drag_impulse=drag_impulse,
1692 correction=ws.packed_psi_correction,
1693 track_magnetic=ws.magnetic_residual!==nothing)
1694 else
1695 error("unknown radiation integrator :$integrator (have :ref, :cube)")
1696 end
1697 end
1698 _time_radiation_stage!(timing,:corrector_s) do
1699 old=s.scratch
1700 gamma_drag=drag_impulse/dt
1701 rho_mean=ws.rho_mean
1702 _update_free_streaming_photons_velocity!(photon_state,ws,s,old,dt,
1703 gamma_drag,rc,rho_mean)
1704 _update_free_streaming_photons_force!(photon_state,ws,s,old,dt,
1705 gamma_drag,rc,rho_mean;
1706 corrected_state=s.U)
1707 end
1708 if energy_ledger!==nothing
1709 _time_radiation_stage!(timing,:ledger_s) do
1710 change=_radiation_energy_changes(ws,s,s.scratch)
1711 photon_change=photon_moment_energy(photon_state,ws,s,rc)-photon_before
1712 _record_radiation_energy!(energy_ledger,change;
1713 photon_change=photon_change)
1714 end
1715 end
1716 timing===nothing || (timing.fast_steps+=1)
1717 ws.density_hat_valid[]=false
1718 reconstruct_radiation_density!(ws,s)
1719 s
1720 end
1721
1722 """
1723 step_radiation_godunov!(s, ws, dt; drag_impulse, closure, ch,
1724 glm_cr=GLM_CR, integrator=:cube,
1725 photon_state=nothing)
1726
1727 Advance full GLM-MHD while embedding the nonlocal photon-baryon response in
1728 the Godunov midpoint states. The spectral solve is a correction to the exact
1729 local-drag predictor, so the free-streaming limit recovers
1730 `step_drag_godunov!`. The same response corrects the conservative trial state;
1731 gas internal and magnetic energies are retained while radiation removes or
1732 returns only the corresponding kinetic work.
1733
1734 Positivity protection (trial check, restore, minmod-PLM/HLL retry, subdivision or
1735 checkpoint error) runs for the `:cube` integrator and for CT (`mhd_solver=:ct`) states
1736 on `:ref`; GLM on `:ref` has none. The wave-speed gate is
1737 `positivity_wave_growth_limit*ch`, so CT callers must pass a positive `ch` (the signal
1738 speed; CT dynamics do not use it).
1739 """
1740 function step_radiation_godunov!(s::MHDState{T},ws,dt::Real;
1741 drag_impulse::Real,closure::RadiationClosure,ch::Real,
1742 glm_cr::Real=GLM_CR,integrator::Symbol=:cube,
1743 energy_ledger::Union{Nothing,RadiationEnergyLedger}=nothing,
1744 timing::Union{Nothing,RadiationStepTiming}=nothing,
1745 free_streaming_tolerance::Real=0,
1746 photon_state=nothing,debug_rho_min::Real=0,
1747 debug_context::AbstractString="",
1748 debug_correction_cfl::Real=Inf,
1749 positivity_fallback::Bool=true,
1750 positivity_max_fallback_passes::Int=8,
1751 positivity_wave_growth_limit::Real=4,
1752 positivity_max_subdivisions::Int=8,
1753 checkpoint_fallback_on_rejection::Bool=false,
1754 _positivity_subdivision_depth::Int=0) where {T}
1755 all_periodic(s) || error("spectral radiation coupling requires periodic boundaries")
1756 dt>0 || error("radiation-coupled dt must be positive")
1757 drag_impulse>=0 || error("radiation drag_impulse must be non-negative")
1758 positivity_max_fallback_passes>0 ||
1759 error("positivity_max_fallback_passes must be positive")
1760 positivity_wave_growth_limit>1 ||
1761 error("positivity_wave_growth_limit must exceed one")
1762 positivity_max_subdivisions>=0 ||
1763 error("positivity_max_subdivisions must be non-negative")
1764 ws.density_limiter_active[]=false
1765 rc=RadiationClosure{T}(T(closure.inertia_ratio),T(closure.mean_free_path),
1766 T(closure.photon_speed),T(closure.gas_sound_speed),T(closure.bridge_q2),
1767 T(closure.longitudinal_viscosity),T(closure.nyquist_guard_fraction),
1768 closure.response_model)
1769 kfund=T(2pi)/(T(s.dims[1])*s.dx)
1770 if rc.response_model == _RADIATION_RESPONSE_BRIDGE &&
1771 rc.mean_free_path*rc.mean_free_path*kfund*kfund >= rc.bridge_q2
1772 energy_ledger === nothing || error(
1773 "radiation energy accounting requires an allocated spectral workspace")
1774 return step_drag_godunov!(s,dt;drag_impulse=drag_impulse,ch=ch,
1775 glm_cr=glm_cr,integrator=integrator)
1776 end
1777 ws isa RadiationWorkspace ||
1778 error("radiation diffusion modes require an allocated RadiationWorkspace")
1779 if _radiation_evolves_moments(rc.response_model)
1780 photon_state isa PhotonMomentState ||
1781 error("response_model=:moments requires photon_state=allocate_photon_moment_state(s)")
1782 photon_state.dims==s.dims || error("photon and MHD grid dimensions differ")
1783 end
1784 gamma_drag=T(drag_impulse)/T(dt)
1785 free_streaming_tolerance>=0 ||
1786 error("free_streaming_tolerance must be non-negative")
1787 kfund=T(2pi)/(T(s.dims[1])*s.dx)
1788 if _radiation_evolves_moments(rc.response_model) &&
1789 free_streaming_tolerance>0 &&
1790 radiation_free_streaming_bound(kfund,gamma_drag,rc) <=
1791 T(free_streaming_tolerance)
1792 return _step_radiation_free_streaming!(s,ws,T(dt);
1793 drag_impulse=T(drag_impulse),rc=rc,ch=T(ch),glm_cr=glm_cr,
1794 integrator=integrator,energy_ledger=energy_ledger,timing=timing,
1795 photon_state=photon_state)
1796 end
1797 timing===nothing || (timing.spectral_steps+=1)
1798 rho_mean=ws.rho_mean
1799 photon_energy_before=_time_radiation_stage!(timing,:ledger_s) do
1800 energy_ledger!==nothing && _radiation_evolves_moments(rc.response_model) ?
1801 photon_moment_energy(photon_state,ws,s,rc) : 0.0
1802 end
1803 _time_radiation_stage!(timing,:predictor_s) do
1804 half=T(0.5)*T(dt)
1805 _radiation_transform_velocity_density!(ws,s,s.U,half,gamma_drag,rc;
1806 density_residual=false,
1807 photons=photon_state,
1808 rho_mean=rho_mean)
1809 tmp=(reshape(s.scratch[1],s.dims),reshape(s.scratch[2],s.dims),
1810 reshape(s.scratch[3],s.dims))
1811 for d in 1:3
1812 copyto!(tmp[d],ws.correction[d])
1813 end
1814 _radiation_transform_force!(ws.correction,ws,s,s.U,half,gamma_drag,rc;
1815 density_residual=false,photons=photon_state,
1816 rho_mean=rho_mean)
1817 _radiation_add3_k!(s.be)(ws.correction...,tmp...;ndrange=ncells(s))
1818 KA.synchronize(s.be)
1819 if isfinite(debug_correction_cfl)
1820 correction_vmax,correction_cfl=
1821 _radiation_correction_cfl(ws,s,dt,ws.correction)
1822 correction_cfl<=debug_correction_cfl || error(
1823 "radiation predictor $debug_context exceeds correction CFL: " *
1824 "dt_vcorr_dx=$correction_cfl limit=$debug_correction_cfl " *
1825 "vmax=$correction_vmax dt=$dt dx=$(s.dx)")
1826 end
1827 copyto!(view(ws.packed_psi_correction,1:ncells(s)),s.U[9])
1828 end
1829 subdivide_step=false
1830 trial_wave_speed=0.0
1831 _time_radiation_stage!(timing,:godunov_s) do
1832 decay=exp(-T(glm_cr)*T(ch)*T(dt)/s.dx)
1833 track_density=rc.response_model!=_RADIATION_RESPONSE_MOMENTS_EXPONENTIAL
1834 # GLM protects the :cube integrator only (the :ref GLM path is unchanged);
1835 # CT runs on :ref and is protected here too, with the pre-step faces
1836 # snapshotted below since ct_update! overwrites them in place.
1837 use_fallback=positivity_fallback && (integrator===:cube || uses_ct(s))
1838 use_fallback && uses_ct(s) && _ct_snapshot_save!(s,ws)
1839 if integrator===:ref
1840 step_ref_radiation!(s,dt;ch=ch,decay=decay,drag_impulse=drag_impulse,
1841 correction=ws.packed_psi_correction,
1842 track_density=track_density,
1843 track_magnetic=ws.magnetic_residual!==nothing)
1844 elseif integrator===:cube
1845 step_cube_radiation!(s,dt;ch=ch,decay=decay,drag_impulse=drag_impulse,
1846 correction=ws.packed_psi_correction,
1847 track_density=track_density,
1848 track_magnetic=ws.magnetic_residual!==nothing,
1849 check_admissibility=use_fallback)
1850 else
1851 error("unknown radiation integrator :$integrator (have :ref, :cube)")
1852 end
1853 if use_fallback
1854 trial_admissible=_radiation_trial_admissible(s,ws)
1855 wave_limit=Float64(positivity_wave_growth_limit)*
1856 max(Float64(ch),eps(Float64))
1857 if trial_admissible
1858 trial_wave_speed=Float64(max_wavespeed!(ws.fft_real,s))
1859 end
1860 reject_trial=!trial_admissible ||
1861 !(isfinite(trial_wave_speed) && trial_wave_speed<=wave_limit)
1862 if reject_trial
1863 high_trial_wave_speed=trial_wave_speed
1864 high_trial_rho_min=trial_admissible ? NaN :
1865 Float64(minimum(s.U[1]))
1866 _restore_radiation_trial!(s,ws,track_density)
1867
1868 # The second trial is a complete conservative update: every face
1869 # uses the same minmod-PLM/HLL flux, unlike a cell-local flux swap.
1870 if integrator===:ref
1871 step_ref_radiation!(s,dt;ch=ch,decay=decay,
1872 drag_impulse=drag_impulse,
1873 correction=ws.packed_psi_correction,
1874 track_density=track_density,
1875 track_magnetic=ws.magnetic_residual!==nothing,
1876 robust_fallback=true)
1877 else
1878 step_cube_radiation!(s,dt;ch=ch,decay=decay,
1879 drag_impulse=drag_impulse,
1880 correction=ws.packed_psi_correction,
1881 track_density=track_density,check_admissibility=true,
1882 track_magnetic=ws.magnetic_residual!==nothing,
1883 robust_fallback=true)
1884 end
1885 robust_admissible=_radiation_trial_admissible(s,ws)
1886 robust_wave_speed=robust_admissible ?
1887 Float64(max_wavespeed!(ws.fft_real,s)) : 0.0
1888 robust_rejected=!robust_admissible ||
1889 !(isfinite(robust_wave_speed) && robust_wave_speed<=wave_limit)
1890 robust_trial_rho_min=robust_admissible ? NaN :
1891 Float64(minimum(s.U[1]))
1892 if !robust_rejected
1893 trial_wave_speed=robust_wave_speed
1894 fallback_count=0
1895 if timing!==nothing
1896 timing.fallback_steps+=1
1897 timing.plm_hll_fallback_steps+=1
1898 fallback_count=timing.plm_hll_fallback_steps
1899 end
1900 ws.density_limiter_active[]=true
1901 if timing===nothing || fallback_count==1 || fallback_count%1000==0
1902 @warn "radiation step used conservative minmod-PLM/HLL fallback" dt debug_context fallback_count high_trial_rho_min high_trial_wave_speed robust_wave_speed wave_limit
1903 end
1904 elseif checkpoint_fallback_on_rejection
1905 _checkpoint_fallback_rejection!(s,ws,track_density;dt=dt,
1906 debug_context=debug_context,
1907 trial_rho_min=robust_trial_rho_min,
1908 trial_wave_speed=robust_wave_speed,wave_limit=wave_limit)
1909 elseif _positivity_subdivision_depth<positivity_max_subdivisions
1910 _restore_radiation_trial!(s,ws,track_density)
1911 subdivide_step=true
1912 timing===nothing || (timing.subdivision_steps+=1)
1913 if _positivity_subdivision_depth==0
1914 @warn "radiation PPM/HLL and minmod-PLM/HLL trials rejected; subdividing coupled step" dt debug_context high_trial_rho_min robust_trial_rho_min high_trial_wave_speed robust_wave_speed wave_limit
1915 end
1916 else
1917 _restore_radiation_trial!(s,ws,track_density)
1918 error("conservative minmod-PLM/HLL radiation fallback " *
1919 "remained inadmissible after " *
1920 "$positivity_max_subdivisions subdivisions: dt=$dt " *
1921 "$debug_context high_trial_rho_min=$high_trial_rho_min " *
1922 "robust_trial_rho_min=$robust_trial_rho_min " *
1923 "high_wave_speed=$high_trial_wave_speed " *
1924 "robust_wave_speed=$robust_wave_speed " *
1925 "wave_limit=$wave_limit ch=$ch")
1926 end
1927 end
1928 end
1929 end
1930 if subdivide_step
1931 half_dt=dt/2
1932 half_impulse=drag_impulse/2
1933 depth=_positivity_subdivision_depth+1
1934 step_radiation_godunov!(s,ws,half_dt;drag_impulse=half_impulse,
1935 closure=closure,ch=ch,glm_cr=glm_cr,integrator=integrator,
1936 energy_ledger=energy_ledger,timing=timing,
1937 free_streaming_tolerance=free_streaming_tolerance,
1938 photon_state=photon_state,debug_rho_min=debug_rho_min,
1939 debug_context=debug_context,
1940 debug_correction_cfl=debug_correction_cfl,
1941 positivity_fallback=positivity_fallback,
1942 positivity_max_fallback_passes=positivity_max_fallback_passes,
1943 positivity_wave_growth_limit=positivity_wave_growth_limit,
1944 positivity_max_subdivisions=positivity_max_subdivisions,
1945 checkpoint_fallback_on_rejection=checkpoint_fallback_on_rejection,
1946 _positivity_subdivision_depth=depth)
1947 sub_ch=max(Float64(ch),Float64(max_wavespeed!(ws.fft_real,s)))
1948 return step_radiation_godunov!(s,ws,half_dt;
1949 drag_impulse=half_impulse,closure=closure,ch=sub_ch,glm_cr=glm_cr,
1950 integrator=integrator,energy_ledger=energy_ledger,timing=timing,
1951 free_streaming_tolerance=free_streaming_tolerance,
1952 photon_state=photon_state,debug_rho_min=debug_rho_min,
1953 debug_context=debug_context,
1954 debug_correction_cfl=debug_correction_cfl,
1955 positivity_fallback=positivity_fallback,
1956 positivity_max_fallback_passes=positivity_max_fallback_passes,
1957 positivity_wave_growth_limit=positivity_wave_growth_limit,
1958 positivity_max_subdivisions=positivity_max_subdivisions,
1959 checkpoint_fallback_on_rejection=checkpoint_fallback_on_rejection,
1960 _positivity_subdivision_depth=depth)
1961 end
1962 _time_radiation_stage!(timing,:corrector_s) do
1963 _radiation_debug_rho(s,"mhd_trial",debug_rho_min,debug_context)
1964 old=s.scratch
1965 _radiation_transform_velocity_density!(ws,s,old,T(dt),gamma_drag,rc;
1966 density_residual=true,
1967 photons=photon_state,
1968 update_photons=_radiation_evolves_moments(
1969 rc.response_model),
1970 rho_mean=rho_mean)
1971 persistent_density=rc.response_model==
1972 _RADIATION_RESPONSE_MOMENTS_EXPONENTIAL
1973 density_blend=ws.density_limiter_active[] ?
1974 _radiation_density_blend(ws,s;absolute=false) : 1.0
1975 _radiation_apply_correction!(s,ws,ws.correction;
1976 density_blend=density_blend)
1977 persistent_density && _radiation_commit_density_hat!(ws,s;
1978 additive=true,density_blend=density_blend)
1979 if density_blend<1
1980 ws.density_limiter_min_theta[]=min(ws.density_limiter_min_theta[],
1981 density_blend)
1982 timing===nothing || (timing.density_limited_stages+=1)
1983 end
1984 _radiation_debug_rho(s,"density_correction",debug_rho_min,debug_context)
1985 _radiation_transform_force!(ws.correction,ws,s,old,T(dt),gamma_drag,rc;
1986 density_residual=true,photons=photon_state,
1987 update_photons=_radiation_evolves_moments(
1988 rc.response_model),
1989 corrected_state=s.U,rho_mean=rho_mean)
1990 density_blend=ws.density_limiter_active[] ?
1991 _radiation_density_blend(ws,s;absolute=false) : 1.0
1992 _radiation_apply_correction!(s,ws,ws.correction;
1993 density_blend=density_blend)
1994 persistent_density && _radiation_commit_density_hat!(ws,s;
1995 additive=true,density_blend=density_blend)
1996 if density_blend<1
1997 ws.density_limiter_min_theta[]=min(ws.density_limiter_min_theta[],
1998 density_blend)
1999 timing===nothing || (timing.density_limited_stages+=1)
2000 end
2001 _radiation_debug_rho(s,"force_correction",debug_rho_min,debug_context)
2002 end
2003 if energy_ledger !== nothing
2004 _time_radiation_stage!(timing,:ledger_s) do
2005 change=_radiation_energy_changes(ws,s,s.scratch)
2006 photon_change=_radiation_evolves_moments(rc.response_model) ?
2007 photon_moment_energy(photon_state,ws,s,rc)-photon_energy_before : 0.0
2008 _record_radiation_energy!(energy_ledger,change;photon_change=photon_change)
2009 end
2010 end
2011 reconstruct_radiation_density!(ws,s)
2012 return s
2013 end