Snapshot from this documentation build. SHA-256: 52aabef93f2db11bf94957c9661524dafbf41a9064221ac2b7a46cddd1954e04
1 using MHDKernels
2 using PoissonKernels
3 using KernelAbstractions
4 using ChemistryKernels
5 using Printf
6 using Dates
7 using Serialization
8
9 const KA = KernelAbstractions
10 const BACKEND_NAME = Symbol(lowercase(get(ENV, "MHD_BACKEND", "metal")))
11 const T = let precision = lowercase(get(ENV, "MHD_PRECISION", "f32"))
12 precision in ("f32", "float32") ? Float32 :
13 precision in ("f64", "float64") ? Float64 :
14 error("MHD_PRECISION must be f32 or f64")
15 end
16 BACKEND_NAME !== :cpu && T !== Float32 &&
17 error("MHD_PRECISION=f64 is supported only by the CPU reference backend")
18 const _PMF_ARAD = 7.5657e-15
19 const _PMF_SIGT = 6.6524e-25
20 const _PMF_CL = 2.99792458e10
21 # Keep pressure initialization, drag, photon closure, and chemistry on the same
22 # Fixsen/Grackle convention used by ChemistryKernels.comp2_cmb.
23 const _PMF_T_CMB = 2.725
24 const _PMF_KPC_CM = 3.0856775814913673e21
25 const _PMF_MPC_CM = 1000 * _PMF_KPC_CM
26 const _PMF_CHECKPOINT_MAGIC = "MHDKernels PMF checkpoint v1"
27
28 function _write_pmf_checkpoint(path::AbstractString, config, runtime, arrays)
29 isempty(path) && return nothing
30 parent = dirname(path)
31 isempty(parent) || mkpath(parent)
32 tmp = path * ".tmp"
33 open(tmp, "w") do io
34 serialize(io, _PMF_CHECKPOINT_MAGIC)
35 serialize(io, config)
36 serialize(io, runtime)
37 serialize(io, length(arrays))
38 for (name, device_array) in arrays
39 serialize(io, String(name))
40 # Stream one field at a time so checkpointing does not retain a
41 # second full simulation state in host memory.
42 serialize(io, Array(device_array))
43 end
44 flush(io)
45 end
46 mv(tmp, path; force=true)
47 return filesize(path)
48 end
49
50 function _restore_pmf_checkpoint!(path::AbstractString, expected_config, arrays;
51 allowed_config_changes=(),
52 allowed_config_validator=nothing,
53 optional_missing_arrays=())
54 open(path, "r") do io
55 deserialize(io) == _PMF_CHECKPOINT_MAGIC ||
56 error("unsupported PMF checkpoint format: $path")
57 stored_config = deserialize(io)
58 runtime = deserialize(io)
59 if haskey(expected_config, :mhd_solver)
60 # A checkpoint written before the solver option existed is a GLM checkpoint.
61 # Never let the legacy-compatibility path resume it (or a GLM one) as CT, or
62 # vice versa: CT carries face fields that a GLM checkpoint cannot provide.
63 stored_solver = haskey(stored_config, :mhd_solver) ?
64 Symbol(stored_config.mhd_solver) : :glm
65 expected_solver = Symbol(expected_config.mhd_solver)
66 stored_solver == expected_solver || error(
67 "PMF checkpoint mhd_solver mismatch: checkpoint was written by " *
68 "mhd_solver=$(stored_solver)" *
69 (haskey(stored_config, :mhd_solver) ? "" : " (no mhd_solver key: legacy GLM)") *
70 " but this run requests mhd_solver=$(expected_solver); resume with the same " *
71 "solver or start a new run")
72 haskey(stored_config, :mhd_solver) ||
73 (stored_config = merge(stored_config, (mhd_solver=:glm,)))
74 end
75 if !isequal(stored_config, expected_config)
76 stored_keys = keys(stored_config)
77 compatible_legacy = all(k -> haskey(expected_config, k) &&
78 isequal(getproperty(stored_config, k),
79 getproperty(expected_config, k)),
80 stored_keys)
81 # The requested stopping redshift does not enter any update. Permit
82 # extending or shortening a run while keeping every evolution
83 # parameter in the checkpoint fingerprint strict.
84 compatible_endpoint = all(k -> k === :zend ||
85 (haskey(expected_config, k) &&
86 isequal(getproperty(stored_config, k),
87 getproperty(expected_config, k))),
88 stored_keys) &&
89 all(k -> k === :zend || haskey(stored_config, k),
90 keys(expected_config))
91 allowed = Set(Symbol.(allowed_config_changes))
92 structurally_compatible_override = !isempty(allowed) &&
93 all(k -> k === :zend || k in allowed ||
94 (haskey(expected_config, k) &&
95 isequal(getproperty(stored_config, k),
96 getproperty(expected_config, k))), stored_keys) &&
97 all(k -> k === :zend || k in allowed || haskey(stored_config, k),
98 keys(expected_config)) &&
99 all(k -> haskey(stored_config, k) && haskey(expected_config, k), allowed)
100 changed = structurally_compatible_override ?
101 filter(k -> !isequal(getproperty(stored_config, k),
102 getproperty(expected_config, k)),
103 collect(allowed)) : Symbol[]
104 compatible_override = structurally_compatible_override &&
105 (allowed_config_validator === nothing ||
106 allowed_config_validator(stored_config, expected_config,
107 runtime, changed))
108 (compatible_legacy || compatible_endpoint || compatible_override) || error(
109 "PMF checkpoint configuration mismatch\n" *
110 "stored: $(stored_config)\nexpected: $(expected_config)")
111 if compatible_legacy
112 haskey(expected_config,:velocity_transport_version) &&
113 !haskey(stored_config,:velocity_transport_version) && error(
114 "checkpoint predates corrected photon velocity evolution; " *
115 "restart this trajectory from a consistent initial state")
116 @warn "restoring legacy PMF checkpoint without newer fingerprint fields" missing=filter(
117 k -> !(k in stored_keys), collect(keys(expected_config)))
118 elseif compatible_endpoint
119 @info "restoring PMF checkpoint with changed stopping redshift" stored_zend=stored_config.zend expected_zend=expected_config.zend
120 else
121 @warn "restoring PMF checkpoint with explicitly allowed configuration changes" changed
122 end
123 end
124 narrays = deserialize(io)
125 narrays <= length(arrays) || error(
126 "PMF checkpoint array-count mismatch: stored=$narrays expected=$(length(arrays))")
127 optional = Set(String.(optional_missing_arrays))
128 for (index,(expected_name, device_array)) in enumerate(arrays)
129 if index > narrays
130 String(expected_name) in optional || error(
131 "PMF checkpoint is missing required trailing field $expected_name")
132 fill!(device_array,zero(eltype(device_array)))
133 continue
134 end
135 stored_name = deserialize(io)
136 stored_name == String(expected_name) || error(
137 "PMF checkpoint field mismatch: stored=$stored_name expected=$expected_name")
138 host_array = deserialize(io)
139 size(host_array) == size(device_array) || error(
140 "PMF checkpoint size mismatch for $stored_name: " *
141 "stored=$(size(host_array)) expected=$(size(device_array))")
142 eltype(host_array) == eltype(device_array) || error(
143 "PMF checkpoint type mismatch for $stored_name: " *
144 "stored=$(eltype(host_array)) expected=$(eltype(device_array))")
145 copyto!(device_array, host_array)
146 end
147 return runtime
148 end
149 end
150
151 function _nonlocal_table_redshift_support(path::AbstractString)
152 lines = readlines(path)
153 isempty(lines) && error("nonlocal response table is empty: $path")
154 header = split(strip(lines[1]), '\t')
155 column = findfirst(==("redshift"), header)
156 column === nothing && error("nonlocal response table has no redshift column: $path")
157 redshifts = Float64[]
158 for line in lines[2:end]
159 isempty(strip(line)) && continue
160 fields = split(strip(line), '\t')
161 push!(redshifts, parse(Float64, fields[column]))
162 end
163 isempty(redshifts) && error("nonlocal response table has no data: $path")
164 extrema(redshifts)
165 end
166
167 function _checkpoint_nonlocal_table_change_is_safe(stored_config, expected_config,
168 runtime, changed)
169 :chem_nonlocal_table in changed || return true
170 checkpoint_redshift = a_to_z(runtime.a)
171 stored_support = _nonlocal_table_redshift_support(stored_config.chem_nonlocal_table)
172 expected_support = _nonlocal_table_redshift_support(expected_config.chem_nonlocal_table)
173 outside(support) = checkpoint_redshift < support[1] || checkpoint_redshift > support[2]
174 outside(stored_support) && outside(expected_support) || error(
175 "cannot change the nonlocal response table at checkpoint z=$(checkpoint_redshift): " *
176 "stored support=$(stored_support), expected support=$(expected_support)")
177 @info "validated nonlocal-table checkpoint migration outside both supports" checkpoint_redshift stored_support expected_support
178 true
179 end
180
181 z_to_a(z) = 1.0 / (1.0 + z)
182 a_to_z(a) = 1.0 / a - 1.0
183 dadtau(Om, OL, Ok, Or, a) = sqrt(a^3 * (Om + OL*a^3 + Ok*a + Or/a))
184 dtau_for_dlna(Om, OL, Ok, Or, a, dlna) = a * dlna / dadtau(Om, OL, Ok, Or, a)
185
186 function _particle_velocity_cms_per_code(a, box_ckpc_h)
187 a > 0 || throw(ArgumentError("particle velocity conversion requires a > 0"))
188 box_ckpc_h > 0 ||
189 throw(ArgumentError("particle velocity conversion requires a positive box"))
190 # H0 * (box_ckpc_h / h) = 10^4 * box_ckpc_h cm/s.
191 return 1.0e4 * Float64(box_ckpc_h) / Float64(a)
192 end
193
194 if BACKEND_NAME === :metal
195 using Metal
196 Metal.functional() || error("Metal backend is not functional")
197 MHDKernels.has_backend(:metal) || error("MHDKernelsMetalExt did not load")
198 PoissonKernels.has_backend(:metal) || error("PoissonKernelsMetalExt did not load")
199 ChemistryKernels.has_backend(:metal) || error("ChemistryKernelsMetalExt did not load")
200 elseif BACKEND_NAME === :cuda
201 using CUDA
202 end
203
204 using KernelAbstractions: @Const, @index, @kernel
205
206 @kernel function _init_lattice_k!(px, py, pz, vx, vy, vz, N::Int)
207 p = @index(Global)
208 @inbounds begin
209 T = eltype(px)
210 q = p - 1
211 i = q % N
212 j = (q ÷ N) % N
213 k = q ÷ (N * N)
214 invN = one(T) / T(N)
215 px[p] = (T(i) + T(0.5)) * invN
216 py[p] = (T(j) + T(0.5)) * invN
217 pz[p] = (T(k) + T(0.5)) * invN
218 vx[p] = zero(T)
219 vy[p] = zero(T)
220 vz[p] = zero(T)
221 end
222 end
223
224 @kernel function _particle_speed2_k!(work, @Const(vx), @Const(vy), @Const(vz))
225 p = @index(Global)
226 @inbounds begin
227 T = eltype(work)
228 x = T(vx[p])
229 y = T(vy[p])
230 z = T(vz[p])
231 work[p] = x*x + y*y + z*z
232 end
233 end
234
235 @kernel function _grid_accel2_k!(work, @Const(φ), N::Int)
236 c = @index(Global)
237 @inbounds begin
238 q = c - 1
239 i = q % N
240 j = (q ÷ N) % N
241 k = q ÷ (N * N)
242 ip = (i + 1) % N
243 im = (i + N - 1) % N
244 jp = (j + 1) % N
245 jm = (j + N - 1) % N
246 kp = (k + 1) % N
247 km = (k + N - 1) % N
248 lin(ii, jj, kk) = ii + N * jj + N * N * kk + 1
249 hc = Float32(0.5) * Float32(N)
250 gx = -hc * (Float32(φ[lin(ip, j, k)]) - Float32(φ[lin(im, j, k)]))
251 gy = -hc * (Float32(φ[lin(i, jp, k)]) - Float32(φ[lin(i, jm, k)]))
252 gz = -hc * (Float32(φ[lin(i, j, kp)]) - Float32(φ[lin(i, j, km)]))
253 work[c] = gx*gx + gy*gy + gz*gz
254 end
255 end
256
257 function _particle_dt_limit(vmax::Float64, gmax::Float64, maxdisp::Float64)
258 if gmax > 0
259 return 2maxdisp / (sqrt(vmax*vmax + 2gmax*maxdisp) + vmax)
260 elseif vmax > 0
261 return maxdisp / vmax
262 else
263 return Inf
264 end
265 end
266
267 @kernel function _push_kdk_global_phi_mix_k!(px, py, pz, vx, vy, vz,
268 @Const(φ0), @Const(φ1),
269 theta, dcoef, ts, driftcoef,
270 wrap, dowrap::Int,
271 invcell, c05, c1,
272 n1::Int, n2::Int, n3::Int, hc)
273 p = @index(Global)
274 @inbounds begin
275 xq = px[p] + dcoef * vx[p]
276 yq = py[p] + dcoef * vy[p]
277 zq = pz[p] + dcoef * vz[p]
278 xpos = xq * invcell
279 ypos = yq * invcell
280 zpos = zq * invcell
281 i1 = unsafe_trunc(Int, xpos + c05)
282 j1 = unsafe_trunc(Int, ypos + c05)
283 k1 = unsafe_trunc(Int, zpos + c05)
284 dxr = oftype(xpos, i1) + c05 - xpos
285 dyr = oftype(ypos, j1) + c05 - ypos
286 dzr = oftype(zpos, k1) + c05 - zpos
287 ex = c1 - dxr
288 ey = c1 - dyr
289 ez = c1 - dzr
290 ax0, ay0, az0 = PoissonKernels._global_phi_force(φ0, i1, j1, k1,
291 dxr, dyr, dzr, ex, ey, ez,
292 n1, n2, n3, hc)
293 ax1, ay1, az1 = PoissonKernels._global_phi_force(φ1, i1, j1, k1,
294 dxr, dyr, dzr, ex, ey, ez,
295 n1, n2, n3, hc)
296 omθ = c1 - theta
297 ax = omθ * ax0 + theta * ax1
298 ay = omθ * ay0 + theta * ay1
299 az = omθ * az0 + theta * az1
300
301 Tv = eltype(vx)
302 nvx = Tv(vx[p] + ax * ts)
303 nvy = Tv(vy[p] + ay * ts)
304 nvz = Tv(vz[p] + az * ts)
305
306 x = px[p] + driftcoef * nvx
307 y = py[p] + driftcoef * nvy
308 z = pz[p] + driftcoef * nvz
309 if dowrap == 1
310 x = mod(x, wrap)
311 y = mod(y, wrap)
312 z = mod(z, wrap)
313 end
314 px[p] = x
315 py[p] = y
316 pz[p] = z
317 vx[p] = Tv(nvx + ax * ts)
318 vy[p] = Tv(nvy + ay * ts)
319 vz[p] = Tv(nvz + az * ts)
320 end
321 end
322
323 function _push_particles_fused_global_phi_mix!(px::AbstractVector{T}, py, pz, vx, vy, vz,
324 φ0, φ1; theta::Real, dtau::Real, nc,
325 wrap::Real=1) where {T}
326 be = KA.get_backend(px)
327 n1, n2, n3 = Int(nc[1]), Int(nc[2]), Int(nc[3])
328 half = T(0.5) * T(dtau)
329 dowrap = wrap > 0 ? 1 : 0
330 _push_kdk_global_phi_mix_k!(be)(px, py, pz, vx, vy, vz, φ0, φ1,
331 T(theta), half, half, T(dtau),
332 T(wrap), dowrap, T(n1), T(0.5), T(1),
333 n1, n2, n3, T(0.5) * T(n1);
334 ndrange=length(px))
335 return nothing
336 end
337
338 @kernel function _push_kdk_global_lattice_k!(dxp, dyp, dzp, vx, vy, vz, @Const(phi),
339 dcoef, ts, driftcoef, dowrap::Int,
340 c05, c1, N::Int, hc)
341 p = @index(Global)
342 @inbounds begin
343 Tv = eltype(vx)
344 q = p - 1
345 i = q % N
346 j = (q ÷ N) % N
347 k = q ÷ (N * N)
348 nT = Tv(N)
349 # Keep the fractional offset separate from the integer lattice cell.
350 # Forming i + N*dxp in Float32 discards the high-z displacement.
351 ux = c05 + nT * (Tv(dxp[p]) + dcoef * Tv(vx[p]))
352 uy = c05 + nT * (Tv(dyp[p]) + dcoef * Tv(vy[p]))
353 uz = c05 + nT * (Tv(dzp[p]) + dcoef * Tv(vz[p]))
354 ox = unsafe_trunc(Int, floor(ux + c05))
355 oy = unsafe_trunc(Int, floor(uy + c05))
356 oz = unsafe_trunc(Int, floor(uz + c05))
357 i1 = i + ox
358 j1 = j + oy
359 k1 = k + oz
360 dxr = Tv(ox) + c05 - ux
361 dyr = Tv(oy) + c05 - uy
362 dzr = Tv(oz) + c05 - uz
363 # Resolve both complementary weights directly. Computing 1-dxr would
364 # round a sub-ULP interpolation weight to zero when dxr is near one.
365 ex = ux - Tv(ox) + c05
366 ey = uy - Tv(oy) + c05
367 ez = uz - Tv(oz) + c05
368 ax, ay, az = PoissonKernels._global_phi_force(phi, i1, j1, k1,
369 dxr, dyr, dzr, ex, ey, ez,
370 N, N, N, hc)
371
372 nvx = Tv(vx[p] + ax * ts)
373 nvy = Tv(vy[p] + ay * ts)
374 nvz = Tv(vz[p] + az * ts)
375 x = Tv(dxp[p]) + driftcoef * nvx
376 y = Tv(dyp[p]) + driftcoef * nvy
377 z = Tv(dzp[p]) + driftcoef * nvz
378 if dowrap == 1
379 x -= floor(x + c05)
380 y -= floor(y + c05)
381 z -= floor(z + c05)
382 end
383 dxp[p] = x; dyp[p] = y; dzp[p] = z
384 vx[p] = Tv(nvx + ax * ts)
385 vy[p] = Tv(nvy + ay * ts)
386 vz[p] = Tv(nvz + az * ts)
387 end
388 end
389
390 function _push_particles_fused_global_lattice!(dxp::AbstractVector{T}, dyp, dzp,
391 vx, vy, vz, phi;
392 dtau::Real, N::Int,
393 wrap::Bool=true) where {T}
394 be = KA.get_backend(dxp)
395 half = T(0.5) * T(dtau)
396 _push_kdk_global_lattice_k!(be)(dxp, dyp, dzp, vx, vy, vz, phi,
397 half, half, T(dtau), wrap ? 1 : 0,
398 T(0.5), T(1), N, T(0.5) * T(N);
399 ndrange=length(dxp))
400 return nothing
401 end
402
403 function _deposit_particles!(rho, px, py, pz, vx, vy, vz;
404 N::Int, disp::Real, lattice_displacements::Bool)
405 if lattice_displacements
406 MHDKernels.deposit_lattice_displacement_delta!(
407 rho, px, py, pz, vx, vy, vz; N=N, disp=disp, shift=-0.5,
408 )
409 else
410 PoissonKernels.cic_deposit!(rho, px, py, pz, vx, vy, vz, one(eltype(rho));
411 N=N, disp=disp, shift=-0.5)
412 end
413 return rho
414 end
415
416 function _push_particles!(px, py, pz, vx, vy, vz, phi;
417 dtau::Real, N::Int, wrap::Bool,
418 lattice_displacements::Bool)
419 if lattice_displacements
420 _push_particles_fused_global_lattice!(px, py, pz, vx, vy, vz, phi;
421 dtau=dtau, N=N, wrap=wrap)
422 else
423 PoissonKernels.push_particles_fused_global!(px, py, pz, vx, vy, vz, phi;
424 dtau=dtau, nc=(N,N,N),
425 wrap=wrap ? 1 : 0)
426 end
427 return nothing
428 end
429
430 @kernel function _assemble_total_perturbation_k!(ρtot, @Const(gasρ), @Const(dmρ),
431 fb, fdm, gas_mean, dm_mean,
432 N::Int)
433 c = @index(Global)
434 @inbounds ρtot[c] = fb * (gasρ[c] - gas_mean) +
435 fdm * (dmρ[c] - dm_mean)
436 end
437
438 @kernel function _delta_from_density_k!(dst, @Const(src), N::Int)
439 c = @index(Global)
440 @inbounds dst[c] = src[c] - one(eltype(dst))
441 end
442
443 @kernel function _density_from_delta_k!(dst, @Const(src), N::Int)
444 c = @index(Global)
445 @inbounds dst[c] = src[c] + one(eltype(dst))
446 end
447
448 @kernel function _init_single_b_k!(rho, mx, my, mz, E, Bx, By, Bz, psi,
449 N::Int, γ, p0, b0, mode::Int, kmode::Int)
450 c = @index(Global)
451 @inbounds begin
452 T = eltype(rho)
453 i = (c - 1) % N + 1
454 x = (T(i) - T(0.5)) / T(N)
455 phase = T(2π) * T(kmode) * x
456 bx = zero(T)
457 by = mode == 1 ? b0 * sin(phase) : b0 * sin(phase)
458 bz = mode == 1 ? b0 * cos(phase) : zero(T)
459 rho[c] = one(T)
460 mx[c] = zero(T); my[c] = zero(T); mz[c] = zero(T)
461 Bx[c] = bx; By[c] = by; Bz[c] = bz; psi[c] = zero(T)
462 E[c] = p0 / (γ - one(T)) + T(0.5) * (bx * bx + by * by + bz * bz)
463 end
464 end
465
466 @kernel function _init_alfven_linear_k!(rho, mx, my, mz, E, Bx, By, Bz, psi,
467 N::Int, γ, p0, bguide, bpert, kmode::Int)
468 c = @index(Global)
469 @inbounds begin
470 T = eltype(rho)
471 i = (c - 1) % N + 1
472 x = (T(i) - T(0.5)) / T(N)
473 phase = T(2π) * T(kmode) * x
474 bx = bguide
475 by = bpert * sin(phase)
476 bz = zero(T)
477 rho[c] = one(T)
478 mx[c] = zero(T); my[c] = zero(T); mz[c] = zero(T)
479 Bx[c] = bx; By[c] = by; Bz[c] = bz; psi[c] = zero(T)
480 E[c] = p0 / (γ - one(T)) + T(0.5) * (bx * bx + by * by + bz * bz)
481 end
482 end
483
484 @kernel function _init_meanfield_linear_k!(rho, mx, my, mz, E, Bx, By, Bz, psi,
485 N::Int, γ, p0, bguide, bpert, theta,
486 polarization::Int, kmode::Int)
487 c = @index(Global)
488 @inbounds begin
489 T = eltype(rho)
490 i = (c - 1) % N + 1
491 x = (T(i) - T(0.5)) / T(N)
492 phase = T(2π) * T(kmode) * x
493 bx = bguide * cos(theta)
494 by = bguide * sin(theta) + (polarization == 1 ? bpert * cos(phase) : zero(T))
495 bz = polarization == 2 ? bpert * cos(phase) : zero(T)
496 rho[c] = one(T)
497 mx[c] = zero(T); my[c] = zero(T); mz[c] = zero(T)
498 Bx[c] = bx; By[c] = by; Bz[c] = bz; psi[c] = zero(T)
499 E[c] = p0 / (γ - one(T)) + T(0.5) * (bx * bx + by * by + bz * bz)
500 end
501 end
502
503 function _single_b_mode(env::AbstractString)
504 s = Symbol(lowercase(env))
505 s in (:forcefree, :circular) && return 1
506 s in (:sine, :single, :single_sine, :nonforcefree,
507 :magpressure, :compressive, :pressure, :density) && return 2
508 error("MHD_PMF_INIT must be batchelor, forcefree, sine/magpressure, or alfven_linear; got '$env'")
509 end
510
511 @kernel function _init_reduced_chem_k!(HII, H2I, @Const(rho), xhii, xh2, fh,
512 hydrogen_reference, N::Int,
513 ::Val{NH}, ::Val{CH}) where {NH,CH}
514 c = @index(Global)
515 @inbounds begin
516 r = rho[c]
517 HII[c] = CH ? xhii - hydrogen_reference :
518 (NH ? max(one(r) - xhii - xh2, zero(r)) : xhii) * fh * r
519 H2I[c] = xh2 * fh * r
520 end
521 end
522
523 @kernel function _recenter_hydrogen_delta_k!(hydrogen_delta, shift)
524 c = @index(Global)
525 @inbounds hydrogen_delta[c] -= shift
526 end
527
528 @kernel function _init_reduced_helium_k!(HeII, @Const(rho), heii_mass_fraction, N::Int)
529 c = @index(Global)
530 @inbounds HeII[c] = heii_mass_fraction * rho[c]
531 end
532
533 @kernel function _mhd_to_eint_k!(eint, eint_before, @Const(rho), @Const(mx),
534 @Const(my), @Const(mz), @Const(E), @Const(Bx),
535 @Const(By), @Const(Bz), small)
536 c = @index(Global)
537 @inbounds begin
538 T = eltype(eint)
539 r = max(rho[c], small)
540 ek = T(0.5) * (mx[c]*mx[c] + my[c]*my[c] + mz[c]*mz[c]) / r
541 eb = T(0.5) * (Bx[c]*Bx[c] + By[c]*By[c] + Bz[c]*Bz[c])
542 value = max((E[c] - ek - eb) / r, T(1e-30))
543 eint[c] = value
544 eint_before[c] = value
545 end
546 end
547
548 # Rebuilding E from kinetic, magnetic, and internal energy can round even a
549 # chemistry no-op. Applying only the changed internal component is bitwise
550 # identity when the network leaves it unchanged and also reads fewer arrays.
551 @inline function _incremental_total_energy(E, rho, eint_before, eint_after, small)
552 max(E + rho * (eint_after - eint_before), small)
553 end
554
555 @inline function _incremental_total_energy_compensated!(residual,c,E,rho,
556 eint_before,eint_after,
557 small,tracked::Val)
558 MHDKernels._compensated_floor_update!(
559 residual,c,E,rho*(eint_after-eint_before),small,tracked)
560 end
561
562 @kernel function _eint_delta_to_mhd_k!(E, @Const(eint), @Const(eint_before),
563 @Const(rho), energy_residual,small,
564 track_energy::Val{TRACK}) where {TRACK}
565 c = @index(Global)
566 @inbounds begin
567 T = eltype(E)
568 r = max(rho[c], small)
569 E[c] = _incremental_total_energy_compensated!(energy_residual,c,E[c],r,
570 max(eint_before[c],T(1e-30)),max(eint[c],T(1e-30)),T(1e-30),
571 track_energy)
572 end
573 end
574
575 @kernel function _chem_analytic_mixing_mhd_k!(rho, @Const(mx), @Const(my), @Const(mz),
576 E, @Const(Bx), @Const(By), @Const(Bz),
577 HII, H2I, @Const(nsm),energy_residual,
578 nsm_scalar, du, vu2, tu, dt, z,
579 f_alpha, Xe_mean, fudge, gauss,
580 hubble, Om, OL, fh,
581 hubble_expansion, adot_over_a,
582 rate_tables, dtfrac, itcap, small,
583 track_energy::Val{TRACK},
584 ::Val{NS}, ::Val{SN}) where {TRACK,NS,SN}
585 c = @index(Global)
586 @inbounds begin
587 T = eltype(E)
588 r_code = max(rho[c], small)
589 ek = T(0.5) * (mx[c]*mx[c] + my[c]*my[c] + mz[c]*mz[c]) / r_code
590 eb = T(0.5) * (Bx[c]*Bx[c] + By[c]*By[c] + Bz[c]*Bz[c])
591 e_code = max((E[c] - ek - eb) / r_code, T(1e-30))
592 nsm_code = NS ? T(nsm_scalar) : nsm[c]
593 nsm_h = nsm_code * du * (SN ? one(T) : T(fh)) / T(ChemistryKernels.MH)
594 en, hii, h2, _ = ChemistryKernels.evolve_cell_analytic_mixing(
595 r_code * du, e_code * vu2, HII[c] * du, H2I[c] * du,
596 nsm_h, dt * tu, z;
597 f_alpha=T(f_alpha), Xe_mean=T(Xe_mean),
598 smoothed_is_neutral=Val(SN), fudge=T(fudge), gauss=T(gauss),
599 hubble=T(hubble), Om=T(Om), OL=T(OL), fh=T(fh),
600 hubble_expansion=hubble_expansion, adot_over_a=T(adot_over_a),
601 rate_tables=rate_tables, cool_tables=nothing,
602 dtfrac=T(dtfrac), itcap=itcap)
603 cap = max(T(fh) * r_code, zero(T))
604 h2_code = min(max(h2 / du, zero(T)), cap)
605 hii_code = min(max(hii / du, zero(T)), max(cap - h2_code, zero(T)))
606 H2I[c] = h2_code
607 HII[c] = hii_code
608 e_new = max(en / vu2, T(1e-30))
609 E[c] = _incremental_total_energy_compensated!(energy_residual,c,E[c],
610 r_code,e_code,e_new,T(1e-30),track_energy)
611 end
612 end
613
614 @kernel function _chem_mixing_mhd_k!(rho, @Const(mx), @Const(my), @Const(mz),
615 E, @Const(Bx), @Const(By), @Const(Bz),
616 HII, H2I, @Const(nsm),energy_residual,
617 nsm_scalar, du, vu2, dt, z,
618 f_alpha, Xe_mean, fudge, gauss,
619 hubble, Om, OL, fh, hubble_expansion,
620 rate_tables, dtfrac, itcap, small,
621 hydrogen_references,
622 ::Val{NH}, ::Val{CH},
623 track_energy::Val{TRACK},
624 ::Val{NS}, ::Val{SN}) where {NH,CH,TRACK,NS,SN}
625 c = @index(Global)
626 @inbounds begin
627 T = eltype(E)
628 hydrogen_reference = real(hydrogen_references)
629 hydrogen_complement_reference = imag(hydrogen_references)
630 r_code = max(rho[c], small)
631 ek = T(0.5) * (mx[c]*mx[c] + my[c]*my[c] + mz[c]*mz[c]) / r_code
632 eb = T(0.5) * (Bx[c]*Bx[c] + By[c]*By[c] + Bz[c]*Bz[c])
633 e_code = max((E[c] - ek - eb) / r_code, T(1e-30))
634 hcap = max(T(fh) * r_code, eps(T))
635 h2_code_in = min(max(H2I[c], zero(T)), hcap)
636 xh2 = h2_code_in / hcap
637 hydrogen_code = if CH
638 fraction = NH ? T(hydrogen_complement_reference) - HII[c] - xh2 :
639 T(hydrogen_reference) + HII[c]
640 min(max(fraction, zero(T)), max(one(T) - xh2, zero(T))) * hcap
641 else
642 HII[c]
643 end
644 nsm_code = NS ? T(nsm_scalar) : nsm[c]
645 nsm_h = nsm_code * du * (SN ? one(T) : T(fh)) / T(ChemistryKernels.MH)
646 en, hii, h2, _, _ = ChemistryKernels.evolve_cell_mixing(
647 r_code * du, e_code * vu2, hydrogen_code * du, H2I[c] * du, zero(T),
648 nsm_h, dt, z;
649 f_alpha=T(f_alpha), Xe_mean=T(Xe_mean),
650 smoothed_is_neutral=Val(SN), fudge=T(fudge), gauss=T(gauss),
651 hubble=T(hubble), Om=T(Om), OL=T(OL), fh=T(fh),
652 deuterium=false, helium=false, hubble_expansion=hubble_expansion,
653 rate_tables=rate_tables, dtfrac=T(dtfrac), itcap=itcap,
654 neutral_hydrogen_storage=Val(NH))
655 cap = max(T(fh) * r_code, zero(T))
656 h2_code = min(max(h2 / du, zero(T)), cap)
657 hii_code = min(max(hii / du, zero(T)), max(cap - h2_code, zero(T)))
658 H2I[c] = h2_code
659 HII[c] = if CH
660 NH ? T(hydrogen_complement_reference) -
661 (hii_code + h2_code) / max(cap, eps(T)) :
662 hii_code / max(cap, eps(T)) - T(hydrogen_reference)
663 else
664 hii_code
665 end
666 e_new = max(en / vu2, T(1e-30))
667 E[c] = _incremental_total_energy_compensated!(energy_residual,c,E[c],
668 r_code,e_code,e_new,T(1e-30),track_energy)
669 end
670 end
671
672 @kernel function _chem_mixing_helium_mhd_k!(rho, @Const(mx), @Const(my), @Const(mz),
673 E, @Const(Bx), @Const(By), @Const(Bz),
674 HII, H2I, HeII, @Const(nsm),energy_residual,
675 nsm_scalar, du, vu2, dt, z,
676 f_alpha, Xe_mean, fudge, gauss,
677 hubble, Om, OL, fh, hubble_expansion,
678 rate_tables, dtfrac, itcap, small,
679 hydrogen_references,
680 ::Val{NH}, ::Val{CH},
681 track_energy::Val{TRACK},
682 ::Val{NS}, ::Val{SN}) where {NH,CH,TRACK,NS,SN}
683 c = @index(Global)
684 @inbounds begin
685 T = eltype(E)
686 hydrogen_reference = real(hydrogen_references)
687 hydrogen_complement_reference = imag(hydrogen_references)
688 r_code = max(rho[c], small)
689 ek = T(0.5) * (mx[c]*mx[c] + my[c]*my[c] + mz[c]*mz[c]) / r_code
690 eb = T(0.5) * (Bx[c]*Bx[c] + By[c]*By[c] + Bz[c]*Bz[c])
691 e_code = max((E[c] - ek - eb) / r_code, T(1e-30))
692 hcap = max(T(fh) * r_code, eps(T))
693 h2_code_in = min(max(H2I[c], zero(T)), hcap)
694 xh2 = h2_code_in / hcap
695 hydrogen_code = if CH
696 fraction = NH ? T(hydrogen_complement_reference) - HII[c] - xh2 :
697 T(hydrogen_reference) + HII[c]
698 min(max(fraction, zero(T)), max(one(T) - xh2, zero(T))) * hcap
699 else
700 HII[c]
701 end
702 nsm_code = NS ? T(nsm_scalar) : nsm[c]
703 nsm_h = nsm_code * du * (SN ? one(T) : T(fh)) / T(ChemistryKernels.MH)
704 en, hii, h2, _, heii = ChemistryKernels.evolve_cell_mixing(
705 r_code * du, e_code * vu2, hydrogen_code * du, H2I[c] * du, zero(T),
706 nsm_h, dt, z;
707 f_alpha=T(f_alpha), Xe_mean=T(Xe_mean),
708 smoothed_is_neutral=Val(SN), fudge=T(fudge), gauss=T(gauss),
709 hubble=T(hubble), Om=T(Om), OL=T(OL), fh=T(fh),
710 deuterium=false, helium=true, HeII_m=HeII[c] * du,
711 hubble_expansion=hubble_expansion,
712 rate_tables=rate_tables, dtfrac=T(dtfrac), itcap=itcap,
713 neutral_hydrogen_storage=Val(NH))
714 hcap = max(T(fh) * r_code, zero(T))
715 h2_code = min(max(h2 / du, zero(T)), hcap)
716 hii_code = min(max(hii / du, zero(T)), max(hcap - h2_code, zero(T)))
717 hecap = max((one(T) - T(fh)) * r_code, zero(T))
718 H2I[c] = h2_code
719 HII[c] = if CH
720 NH ? T(hydrogen_complement_reference) -
721 (hii_code + h2_code) / max(hcap, eps(T)) :
722 hii_code / max(hcap, eps(T)) - T(hydrogen_reference)
723 else
724 hii_code
725 end
726 HeII[c] = min(max(heii / du, zero(T)), hecap)
727 e_new = max(en / vu2, T(1e-30))
728 E[c] = _incremental_total_energy_compensated!(energy_residual,c,E[c],
729 r_code,e_code,e_new,T(1e-30),track_energy)
730 end
731 end
732
733 # Metal exposes at most 31 indirect buffer resources to one compute pipeline.
734 # The fully fused helium kernel crosses that limit once compensated total-energy
735 # storage is enabled. Keep chemistry on device, but operate on the already
736 # unpacked specific internal energy and apply its total-energy delta separately.
737 @kernel function _chem_mixing_helium_eint_k!(eint, HII, H2I, HeII,
738 @Const(rho), @Const(nsm),
739 nsm_scalar, du, vu2, dt, z,
740 f_alpha, Xe_mean, fudge, gauss,
741 hubble, Om, OL, fh, hubble_expansion,
742 rate_tables, dtfrac, itcap, small,
743 hydrogen_references,
744 ::Val{NH}, ::Val{CH},
745 ::Val{NS}, ::Val{SN}) where {NH,CH,NS,SN}
746 c = @index(Global)
747 @inbounds begin
748 T = eltype(eint)
749 hydrogen_reference = real(hydrogen_references)
750 hydrogen_complement_reference = imag(hydrogen_references)
751 r_code = max(rho[c], small)
752 e_code = max(eint[c], T(1e-30))
753 hcap = max(T(fh) * r_code, eps(T))
754 h2_code_in = min(max(H2I[c], zero(T)), hcap)
755 xh2 = h2_code_in / hcap
756 hydrogen_code = if CH
757 fraction = NH ? T(hydrogen_complement_reference) - HII[c] - xh2 :
758 T(hydrogen_reference) + HII[c]
759 min(max(fraction, zero(T)), max(one(T) - xh2, zero(T))) * hcap
760 else
761 HII[c]
762 end
763 nsm_code = NS ? T(nsm_scalar) : nsm[c]
764 nsm_h = nsm_code * du * (SN ? one(T) : T(fh)) / T(ChemistryKernels.MH)
765 en, hii, h2, _, heii = ChemistryKernels.evolve_cell_mixing(
766 r_code * du, e_code * vu2, hydrogen_code * du, H2I[c] * du, zero(T),
767 nsm_h, dt, z;
768 f_alpha=T(f_alpha), Xe_mean=T(Xe_mean),
769 smoothed_is_neutral=Val(SN), fudge=T(fudge), gauss=T(gauss),
770 hubble=T(hubble), Om=T(Om), OL=T(OL), fh=T(fh),
771 deuterium=false, helium=true, HeII_m=HeII[c] * du,
772 hubble_expansion=hubble_expansion,
773 rate_tables=rate_tables, dtfrac=T(dtfrac), itcap=itcap,
774 neutral_hydrogen_storage=Val(NH))
775 hcap_out = max(T(fh) * r_code, zero(T))
776 h2_code = min(max(h2 / du, zero(T)), hcap_out)
777 hii_code = min(max(hii / du, zero(T)), max(hcap_out - h2_code, zero(T)))
778 hecap = max((one(T) - T(fh)) * r_code, zero(T))
779 H2I[c] = h2_code
780 HII[c] = if CH
781 NH ? T(hydrogen_complement_reference) -
782 (hii_code + h2_code) / max(hcap_out, eps(T)) :
783 hii_code / max(hcap_out, eps(T)) - T(hydrogen_reference)
784 else
785 hii_code
786 end
787 HeII[c] = min(max(heii / du, zero(T)), hecap)
788 eint[c] = max(en / vu2, T(1e-30))
789 end
790 end
791
792 function _apply_compton_drag_mhd!(be, s, f::Real, vx0::Real, vy0::Real, vz0::Real)
793 MHDKernels.apply_linear_drag!(s, clamp(f, 0.0, 1.0);
794 target_velocity=(vx0, vy0, vz0))
795 return nothing
796 end
797
798 @inline function _lin3(i, j, k, N)
799 ii = mod(i - 1, N) + 1
800 jj = mod(j - 1, N) + 1
801 kk = mod(k - 1, N) + 1
802 return ((kk - 1) * N + (jj - 1)) * N + ii
803 end
804
805 @inline function _pressure_at(rho, mx, my, mz, E, Bx, By, Bz, c, gamma_gas, smallr, smallE)
806 r = max(rho[c], smallr)
807 ek = (mx[c] * mx[c] + my[c] * my[c] + mz[c] * mz[c]) / (2 * r)
808 eb = (Bx[c] * Bx[c] + By[c] * By[c] + Bz[c] * Bz[c]) / 2
809 eint = max(E[c] - ek - eb, smallE)
810 return (gamma_gas - one(gamma_gas)) * eint
811 end
812
813 @kernel function _gas_force_balance_diag_k!(pax, pay, paz, hax, hay, haz,
814 @Const(rho), @Const(mx), @Const(my), @Const(mz), @Const(E),
815 @Const(Bx), @Const(By), @Const(Bz), N::Int, inv2dx,
816 gamma_gas, smallr, smallE)
817 c = @index(Global)
818 @inbounds begin
819 i = (c - 1) % N + 1
820 j = ((c - 1) ÷ N) % N + 1
821 k = (c - 1) ÷ (N * N) + 1
822 cxp = _lin3(i + 1, j, k, N); cxm = _lin3(i - 1, j, k, N)
823 cyp = _lin3(i, j + 1, k, N); cym = _lin3(i, j - 1, k, N)
824 czp = _lin3(i, j, k + 1, N); czm = _lin3(i, j, k - 1, N)
825 px = -(_pressure_at(rho, mx, my, mz, E, Bx, By, Bz, cxp,
826 gamma_gas, smallr, smallE) -
827 _pressure_at(rho, mx, my, mz, E, Bx, By, Bz, cxm,
828 gamma_gas, smallr, smallE)) * inv2dx
829 py = -(_pressure_at(rho, mx, my, mz, E, Bx, By, Bz, cyp,
830 gamma_gas, smallr, smallE) -
831 _pressure_at(rho, mx, my, mz, E, Bx, By, Bz, cym,
832 gamma_gas, smallr, smallE)) * inv2dx
833 pz = -(_pressure_at(rho, mx, my, mz, E, Bx, By, Bz, czp,
834 gamma_gas, smallr, smallE) -
835 _pressure_at(rho, mx, my, mz, E, Bx, By, Bz, czm,
836 gamma_gas, smallr, smallE)) * inv2dx
837 lx, ly, lz = _magnetic_force_at(Bx, By, Bz, i, j, k, N, inv2dx)
838 irho = inv(max(rho[c], smallr))
839 pax[c] = px * irho; pay[c] = py * irho; paz[c] = pz * irho
840 hax[c] = (px + lx) * irho
841 hay[c] = (py + ly) * irho
842 haz[c] = (pz + lz) * irho
843 end
844 end
845
846 @kernel function _gas_total_accel_diag_k!(tx, ty, tz,
847 @Const(hx), @Const(hy), @Const(hz),
848 @Const(sx), @Const(sy), @Const(sz), gamma_drag)
849 c = @index(Global)
850 @inbounds begin
851 tx[c] = hx[c] + gamma_drag * sx[c]
852 ty[c] = hy[c] + gamma_drag * sy[c]
853 tz[c] = hz[c] + gamma_drag * sz[c]
854 end
855 end
856
857 @kernel function _gas_velocity_diag_k!(vx, vy, vz,
858 @Const(rho), @Const(mx), @Const(my), @Const(mz), smallr)
859 c = @index(Global)
860 @inbounds begin
861 irho = inv(max(rho[c], smallr))
862 vx[c] = mx[c] * irho
863 vy[c] = my[c] * irho
864 vz[c] = mz[c] * irho
865 end
866 end
867
868 @kernel function _mass_flux_divergence_diag_k!(divm,
869 @Const(mx), @Const(my), @Const(mz), N::Int, inv2dx)
870 c = @index(Global)
871 @inbounds begin
872 i = (c - 1) % N + 1
873 j = ((c - 1) ÷ N) % N + 1
874 k = (c - 1) ÷ (N * N) + 1
875 divm[c] = ((mx[_lin3(i + 1, j, k, N)] - mx[_lin3(i - 1, j, k, N)]) +
876 (my[_lin3(i, j + 1, k, N)] - my[_lin3(i, j - 1, k, N)]) +
877 (mz[_lin3(i, j, k + 1, N)] - mz[_lin3(i, j, k - 1, N)])) * inv2dx
878 end
879 end
880
881 @inline function _terminal_v_at(rho, mx, my, mz, E, Bx, By, Bz, i, j, k, N,
882 inv2dx, gamma_code, gamma_gas, smallr, smallE, vcap)
883 if gamma_code <= eps(gamma_code)
884 return zero(gamma_code), zero(gamma_code), zero(gamma_code)
885 end
886 T = typeof(gamma_code)
887 c = _lin3(i, j, k, N)
888 cxp = _lin3(i + 1, j, k, N); cxm = _lin3(i - 1, j, k, N)
889 cyp = _lin3(i, j + 1, k, N); cym = _lin3(i, j - 1, k, N)
890 czp = _lin3(i, j, k + 1, N); czm = _lin3(i, j, k - 1, N)
891 # Use conservative Maxwell-stress differences instead of pointwise J x B.
892 # The continuum forms are equivalent for div(B)=0, but the stress form avoids
893 # severe cell-centered aliasing of magnetic-pressure modes near marginal
894 # resolution.
895 bxp = Bx[cxp]; byp = By[cxp]; bzp = Bz[cxp]
896 bxm = Bx[cxm]; bym = By[cxm]; bzm = Bz[cxm]
897 b2p = bxp*bxp + byp*byp + bzp*bzp
898 b2m = bxm*bxm + bym*bym + bzm*bzm
899 txxp = bxp*bxp - T(0.5)*b2p
900 txxm = bxm*bxm - T(0.5)*b2m
901 tyxp = byp*bxp
902 tyxm = bym*bxm
903 tzxp = bzp*bxp
904 tzxm = bzm*bxm
905
906 bxp = Bx[cyp]; byp = By[cyp]; bzp = Bz[cyp]
907 bxm = Bx[cym]; bym = By[cym]; bzm = Bz[cym]
908 b2p = bxp*bxp + byp*byp + bzp*bzp
909 b2m = bxm*bxm + bym*bym + bzm*bzm
910 txyp = bxp*byp
911 txym = bxm*bym
912 tyyp = byp*byp - T(0.5)*b2p
913 tyym = bym*bym - T(0.5)*b2m
914 tzyp = bzp*byp
915 tzym = bzm*bym
916
917 bxp = Bx[czp]; byp = By[czp]; bzp = Bz[czp]
918 bxm = Bx[czm]; bym = By[czm]; bzm = Bz[czm]
919 b2p = bxp*bxp + byp*byp + bzp*bzp
920 b2m = bxm*bxm + bym*bym + bzm*bzm
921 txzp = bxp*bzp
922 txzm = bxm*bzm
923 tyzp = byp*bzp
924 tyzm = bym*bzm
925 tzzp = bzp*bzp - T(0.5)*b2p
926 tzzm = bzm*bzm - T(0.5)*b2m
927
928 fx = ((txxp - txxm) + (txyp - txym) + (txzp - txzm)) * inv2dx
929 fy = ((tyxp - tyxm) + (tyyp - tyym) + (tyzp - tyzm)) * inv2dx
930 fz = ((tzxp - tzxm) + (tzyp - tzym) + (tzzp - tzzm)) * inv2dx
931 pxp = _pressure_at(rho, mx, my, mz, E, Bx, By, Bz, cxp, gamma_gas, smallr, smallE)
932 pxm = _pressure_at(rho, mx, my, mz, E, Bx, By, Bz, cxm, gamma_gas, smallr, smallE)
933 pyp = _pressure_at(rho, mx, my, mz, E, Bx, By, Bz, cyp, gamma_gas, smallr, smallE)
934 pym = _pressure_at(rho, mx, my, mz, E, Bx, By, Bz, cym, gamma_gas, smallr, smallE)
935 pzp = _pressure_at(rho, mx, my, mz, E, Bx, By, Bz, czp, gamma_gas, smallr, smallE)
936 pzm = _pressure_at(rho, mx, my, mz, E, Bx, By, Bz, czm, gamma_gas, smallr, smallE)
937 fx -= (pxp - pxm) * inv2dx
938 fy -= (pyp - pym) * inv2dx
939 fz -= (pzp - pzm) * inv2dx
940 invrg = one(gamma_code) / (max(rho[c], smallr) * gamma_code)
941 vx = fx * invrg
942 vy = fy * invrg
943 vz = fz * invrg
944 v2 = vx * vx + vy * vy + vz * vz
945 if v2 > vcap * vcap
946 fac = vcap / sqrt(v2)
947 vx *= fac; vy *= fac; vz *= fac
948 end
949 return vx, vy, vz
950 end
951
952 @inline function _vxb_at(rho, mx, my, mz, E, Bx, By, Bz, i, j, k, N,
953 inv2dx, gamma_code, gamma_gas, smallr, smallE, vcap)
954 c = _lin3(i, j, k, N)
955 vx, vy, vz = _terminal_v_at(rho, mx, my, mz, E, Bx, By, Bz, i, j, k, N,
956 inv2dx, gamma_code, gamma_gas, smallr, smallE, vcap)
957 bx = Bx[c]; by = By[c]; bz = Bz[c]
958 return vy * bz - vz * by, vz * bx - vx * bz, vx * by - vy * bx
959 end
960
961 @inline function _vxb_cached_at(vx, vy, vz, Bx, By, Bz, i, j, k, N)
962 c = _lin3(i, j, k, N)
963 vxc = vx[c]; vyc = vy[c]; vzc = vz[c]
964 bx = Bx[c]; by = By[c]; bz = Bz[c]
965 return vyc * bz - vzc * by, vzc * bx - vxc * bz, vxc * by - vyc * bx
966 end
967
968 @inline function _terminal_cached_v_at(tvx, tvy, tvz, i, j, k, N, vsmooth)
969 c = _lin3(i, j, k, N)
970 vx = tvx[c]; vy = tvy[c]; vz = tvz[c]
971 if vsmooth <= zero(vsmooth)
972 return vx, vy, vz
973 end
974 cxp = _lin3(i + 1, j, k, N); cxm = _lin3(i - 1, j, k, N)
975 cyp = _lin3(i, j + 1, k, N); cym = _lin3(i, j - 1, k, N)
976 czp = _lin3(i, j, k + 1, N); czm = _lin3(i, j, k - 1, N)
977 a = min(max(vsmooth, zero(vsmooth)), one(vsmooth))
978 cc = one(a) - a
979 cn = a / typeof(a)(6)
980 sx = tvx[cxp] + tvx[cxm] + tvx[cyp] + tvx[cym] + tvx[czp] + tvx[czm]
981 sy = tvy[cxp] + tvy[cxm] + tvy[cyp] + tvy[cym] + tvy[czp] + tvy[czm]
982 sz = tvz[cxp] + tvz[cxm] + tvz[cyp] + tvz[cym] + tvz[czp] + tvz[czm]
983 return cc * vx + cn * sx, cc * vy + cn * sy, cc * vz + cn * sz
984 end
985
986 @inline function _vxb_cached_smooth_at(tvx, tvy, tvz, Bx, By, Bz, i, j, k, N, vsmooth)
987 c = _lin3(i, j, k, N)
988 vxc, vyc, vzc = _terminal_cached_v_at(tvx, tvy, tvz, i, j, k, N, vsmooth)
989 bx = Bx[c]; by = By[c]; bz = Bz[c]
990 return vyc * bz - vzc * by, vzc * bx - vxc * bz, vxc * by - vyc * bx
991 end
992
993 @inline function _magnetic_force_at(Bx, By, Bz, i, j, k, N, inv2dx)
994 T = typeof(inv2dx)
995 c = _lin3(i, j, k, N)
996 cxp = _lin3(i + 1, j, k, N); cxm = _lin3(i - 1, j, k, N)
997 cyp = _lin3(i, j + 1, k, N); cym = _lin3(i, j - 1, k, N)
998 czp = _lin3(i, j, k + 1, N); czm = _lin3(i, j, k - 1, N)
999
1000 bxp = Bx[cxp]; byp = By[cxp]; bzp = Bz[cxp]
1001 bxm = Bx[cxm]; bym = By[cxm]; bzm = Bz[cxm]
1002 b2p = bxp*bxp + byp*byp + bzp*bzp
1003 b2m = bxm*bxm + bym*bym + bzm*bzm
1004 txxp = bxp*bxp - T(0.5)*b2p
1005 txxm = bxm*bxm - T(0.5)*b2m
1006 tyxp = byp*bxp
1007 tyxm = bym*bxm
1008 tzxp = bzp*bxp
1009 tzxm = bzm*bxm
1010
1011 bxp = Bx[cyp]; byp = By[cyp]; bzp = Bz[cyp]
1012 bxm = Bx[cym]; bym = By[cym]; bzm = Bz[cym]
1013 b2p = bxp*bxp + byp*byp + bzp*bzp
1014 b2m = bxm*bxm + bym*bym + bzm*bzm
1015 txyp = bxp*byp
1016 txym = bxm*bym
1017 tyyp = byp*byp - T(0.5)*b2p
1018 tyym = bym*bym - T(0.5)*b2m
1019 tzyp = bzp*byp
1020 tzym = bzm*bym
1021
1022 bxp = Bx[czp]; byp = By[czp]; bzp = Bz[czp]
1023 bxm = Bx[czm]; bym = By[czm]; bzm = Bz[czm]
1024 b2p = bxp*bxp + byp*byp + bzp*bzp
1025 b2m = bxm*bxm + bym*bym + bzm*bzm
1026 txzp = bxp*bzp
1027 txzm = bxm*bzm
1028 tyzp = byp*bzp
1029 tyzm = bym*bzm
1030 tzzp = bzp*bzp - T(0.5)*b2p
1031 tzzm = bzm*bzm - T(0.5)*b2m
1032
1033 fx = ((txxp - txxm) + (txyp - txym) + (txzp - txzm)) * inv2dx
1034 fy = ((tyxp - tyxm) + (tyyp - tyym) + (tyzp - tyzm)) * inv2dx
1035 fz = ((tzxp - tzxm) + (tzyp - tzym) + (tzzp - tzzm)) * inv2dx
1036 return fx, fy, fz
1037 end
1038
1039 @inline function _pressure_implicit_v_at(rho_new, fxm, fym, fzm,
1040 i, j, k, N, inv2dx, gamma_code,
1041 pressure_coeff, smallr, vcap)
1042 c = _lin3(i, j, k, N)
1043 cxp = _lin3(i + 1, j, k, N); cxm = _lin3(i - 1, j, k, N)
1044 cyp = _lin3(i, j + 1, k, N); cym = _lin3(i, j - 1, k, N)
1045 czp = _lin3(i, j, k + 1, N); czm = _lin3(i, j, k - 1, N)
1046 r = max(rho_new[c], smallr)
1047 gx = pressure_coeff * (rho_new[cxp] - rho_new[cxm]) * inv2dx
1048 gy = pressure_coeff * (rho_new[cyp] - rho_new[cym]) * inv2dx
1049 gz = pressure_coeff * (rho_new[czp] - rho_new[czm]) * inv2dx
1050 invrg = one(gamma_code) / (r * max(gamma_code, eps(gamma_code)))
1051 vx = (fxm[c] - gx) * invrg
1052 vy = (fym[c] - gy) * invrg
1053 vz = (fzm[c] - gz) * invrg
1054 v2 = vx * vx + vy * vy + vz * vz
1055 if v2 > vcap * vcap
1056 fac = vcap / sqrt(v2)
1057 vx *= fac; vy *= fac; vz *= fac
1058 end
1059 return vx, vy, vz
1060 end
1061
1062 @inline function _magnetic_capped_v_at(rho, fxm, fym, fzm, c, gamma_code, smallr, vcap)
1063 r = max(rho[c], smallr)
1064 invrg = one(gamma_code) / (r * max(gamma_code, eps(gamma_code)))
1065 vx = fxm[c] * invrg
1066 vy = fym[c] * invrg
1067 vz = fzm[c] * invrg
1068 v2 = vx * vx + vy * vy + vz * vz
1069 if v2 > vcap * vcap
1070 fac = vcap / sqrt(v2)
1071 vx *= fac; vy *= fac; vz *= fac
1072 end
1073 return vx, vy, vz
1074 end
1075
1076 @inline function _vxb_pressure_implicit_at(rho_new, fxm, fym, fzm, Bx, By, Bz,
1077 i, j, k, N, inv2dx, gamma_code,
1078 pressure_coeff, smallr, vcap)
1079 c = _lin3(i, j, k, N)
1080 vx, vy, vz = _pressure_implicit_v_at(rho_new, fxm, fym, fzm, i, j, k, N,
1081 inv2dx, gamma_code, pressure_coeff,
1082 smallr, vcap)
1083 bx = Bx[c]; by = By[c]; bz = Bz[c]
1084 return vy * bz - vz * by, vz * bx - vx * bz, vx * by - vy * bx
1085 end
1086
1087 @kernel function _magnetic_force_field_k!(fx, fy, fz, @Const(Bx), @Const(By), @Const(Bz),
1088 N::Int, inv2dx)
1089 c = @index(Global)
1090 @inbounds begin
1091 i = (c - 1) % N + 1
1092 j = ((c - 1) ÷ N) % N + 1
1093 k = (c - 1) ÷ (N * N) + 1
1094 fxi, fyi, fzi = _magnetic_force_at(Bx, By, Bz, i, j, k, N, inv2dx)
1095 fx[c] = fxi
1096 fy[c] = fyi
1097 fz[c] = fzi
1098 end
1099 end
1100
1101 @kernel function _pressure_implicit_source_k!(source, @Const(rho), @Const(fx), @Const(fy), @Const(fz),
1102 N::Int, inv2dx, gamma_code, smallr, vcap)
1103 c = @index(Global)
1104 @inbounds begin
1105 T = eltype(source)
1106 i = (c - 1) % N + 1
1107 j = ((c - 1) ÷ N) % N + 1
1108 k = (c - 1) ÷ (N * N) + 1
1109 cxp = _lin3(i + 1, j, k, N); cxm = _lin3(i - 1, j, k, N)
1110 cyp = _lin3(i, j + 1, k, N); cym = _lin3(i, j - 1, k, N)
1111 czp = _lin3(i, j, k + 1, N); czm = _lin3(i, j, k - 1, N)
1112 vx, vy, vz = _magnetic_capped_v_at(rho, fx, fy, fz, c, gamma_code, smallr, vcap)
1113 vxp, _, _ = _magnetic_capped_v_at(rho, fx, fy, fz, cxp, gamma_code, smallr, vcap)
1114 vxm, _, _ = _magnetic_capped_v_at(rho, fx, fy, fz, cxm, gamma_code, smallr, vcap)
1115 _, vyp, _ = _magnetic_capped_v_at(rho, fx, fy, fz, cyp, gamma_code, smallr, vcap)
1116 _, vym, _ = _magnetic_capped_v_at(rho, fx, fy, fz, cym, gamma_code, smallr, vcap)
1117 _, _, vzp = _magnetic_capped_v_at(rho, fx, fy, fz, czp, gamma_code, smallr, vcap)
1118 _, _, vzm = _magnetic_capped_v_at(rho, fx, fy, fz, czm, gamma_code, smallr, vcap)
1119 vxf_p = T(0.5) * (vx + vxp)
1120 vxf_m = T(0.5) * (vxm + vx)
1121 vyf_p = T(0.5) * (vy + vyp)
1122 vyf_m = T(0.5) * (vym + vy)
1123 vzf_p = T(0.5) * (vz + vzp)
1124 vzf_m = T(0.5) * (vzm + vz)
1125 fx_p = vxf_p * max(ifelse(vxf_p >= zero(T), rho[c], rho[cxp]), smallr)
1126 fx_m = vxf_m * max(ifelse(vxf_m >= zero(T), rho[cxm], rho[c]), smallr)
1127 fy_p = vyf_p * max(ifelse(vyf_p >= zero(T), rho[c], rho[cyp]), smallr)
1128 fy_m = vyf_m * max(ifelse(vyf_m >= zero(T), rho[cym], rho[c]), smallr)
1129 fz_p = vzf_p * max(ifelse(vzf_p >= zero(T), rho[c], rho[czp]), smallr)
1130 fz_m = vzf_m * max(ifelse(vzf_m >= zero(T), rho[czm], rho[c]), smallr)
1131 div_flux = ((fx_p - fx_m) + (fy_p - fy_m) + (fz_p - fz_m)) * (T(2) * inv2dx)
1132 source[c] = -div_flux
1133 end
1134 end
1135
1136 @kernel function _terminal_velocity_field_k!(tvx, tvy, tvz,
1137 @Const(rho), @Const(mx), @Const(my), @Const(mz), @Const(E),
1138 @Const(Bx), @Const(By), @Const(Bz),
1139 N::Int, inv2dx, gamma_code, gamma_gas, smallr, smallE, vcap)
1140 c = @index(Global)
1141 @inbounds begin
1142 i = (c - 1) % N + 1
1143 j = ((c - 1) ÷ N) % N + 1
1144 k = (c - 1) ÷ (N * N) + 1
1145 vx, vy, vz = _terminal_v_at(rho, mx, my, mz, E, Bx, By, Bz,
1146 i, j, k, N, inv2dx, gamma_code, gamma_gas,
1147 smallr, smallE, vcap)
1148 tvx[c] = vx
1149 tvy[c] = vy
1150 tvz[c] = vz
1151 end
1152 end
1153
1154 @kernel function _terminal_velocity_overlap_diag_k!(delta2, terminal2, full2,
1155 @Const(rho), @Const(mx), @Const(my), @Const(mz),
1156 @Const(tvx), @Const(tvy), @Const(tvz), smallr)
1157 c = @index(Global)
1158 @inbounds begin
1159 r = max(rho[c], smallr)
1160 vx = mx[c] / r
1161 vy = my[c] / r
1162 vz = mz[c] / r
1163 tx = tvx[c]
1164 ty = tvy[c]
1165 tz = tvz[c]
1166 dx = vx - tx
1167 dy = vy - ty
1168 dz = vz - tz
1169 delta2[c] = r * (dx * dx + dy * dy + dz * dz)
1170 terminal2[c] = r * (tx * tx + ty * ty + tz * tz)
1171 full2[c] = r * (vx * vx + vy * vy + vz * vz)
1172 end
1173 end
1174
1175 @kernel function _terminal_velocity_update_cached_k!(orho, omx, omy, omz, oE, oBx, oBy, oBz, opsi,
1176 @Const(rho), @Const(mx), @Const(my), @Const(mz), @Const(E),
1177 @Const(Bx), @Const(By), @Const(Bz), @Const(psi),
1178 @Const(tvx), @Const(tvy), @Const(tvz),
1179 N::Int, dt, inv2dx, gamma_gas, smallr, smallE, vsmooth)
1180 c = @index(Global)
1181 @inbounds begin
1182 T = eltype(E)
1183 i = (c - 1) % N + 1
1184 j = ((c - 1) ÷ N) % N + 1
1185 k = (c - 1) ÷ (N * N) + 1
1186
1187 r = max(rho[c], smallr)
1188 ek0 = T(0.5) * (mx[c] * mx[c] + my[c] * my[c] + mz[c] * mz[c]) / r
1189 eb0 = T(0.5) * (Bx[c] * Bx[c] + By[c] * By[c] + Bz[c] * Bz[c])
1190 eint = max(E[c] - ek0 - eb0, smallE)
1191
1192 vx, vy, vz = _terminal_cached_v_at(tvx, tvy, tvz, i, j, k, N, vsmooth)
1193 cxp = _lin3(i + 1, j, k, N); cxm = _lin3(i - 1, j, k, N)
1194 cyp = _lin3(i, j + 1, k, N); cym = _lin3(i, j - 1, k, N)
1195 czp = _lin3(i, j, k + 1, N); czm = _lin3(i, j, k - 1, N)
1196 vxp_x, _, _ = _terminal_cached_v_at(tvx, tvy, tvz, i + 1, j, k, N, vsmooth)
1197 vxm_x, _, _ = _terminal_cached_v_at(tvx, tvy, tvz, i - 1, j, k, N, vsmooth)
1198 _, vyp_y, _ = _terminal_cached_v_at(tvx, tvy, tvz, i, j + 1, k, N, vsmooth)
1199 _, vym_y, _ = _terminal_cached_v_at(tvx, tvy, tvz, i, j - 1, k, N, vsmooth)
1200 _, _, vzp_z = _terminal_cached_v_at(tvx, tvy, tvz, i, j, k + 1, N, vsmooth)
1201 _, _, vzm_z = _terminal_cached_v_at(tvx, tvy, tvz, i, j, k - 1, N, vsmooth)
1202 vxf_p = T(0.5) * (vx + vxp_x)
1203 vxf_m = T(0.5) * (vxm_x + vx)
1204 vyf_p = T(0.5) * (vy + vyp_y)
1205 vyf_m = T(0.5) * (vym_y + vy)
1206 vzf_p = T(0.5) * (vz + vzp_z)
1207 vzf_m = T(0.5) * (vzm_z + vz)
1208 fx_p = vxf_p * max(ifelse(vxf_p >= zero(T), rho[c], rho[cxp]), smallr)
1209 fx_m = vxf_m * max(ifelse(vxf_m >= zero(T), rho[cxm], rho[c]), smallr)
1210 fy_p = vyf_p * max(ifelse(vyf_p >= zero(T), rho[c], rho[cyp]), smallr)
1211 fy_m = vyf_m * max(ifelse(vyf_m >= zero(T), rho[cym], rho[c]), smallr)
1212 fz_p = vzf_p * max(ifelse(vzf_p >= zero(T), rho[c], rho[czp]), smallr)
1213 fz_m = vzf_m * max(ifelse(vzf_m >= zero(T), rho[czm], rho[c]), smallr)
1214 div_rhov = ((fx_p - fx_m) + (fy_p - fy_m) + (fz_p - fz_m)) * (T(2) * inv2dx)
1215 nrho = max(rho[c] - dt * div_rhov, smallr)
1216 cyp_x, cyp_y, cyp_z = _vxb_cached_smooth_at(tvx, tvy, tvz, Bx, By, Bz, i, j + 1, k, N, vsmooth)
1217 cym_x, cym_y, cym_z = _vxb_cached_smooth_at(tvx, tvy, tvz, Bx, By, Bz, i, j - 1, k, N, vsmooth)
1218 czp_x, czp_y, czp_z = _vxb_cached_smooth_at(tvx, tvy, tvz, Bx, By, Bz, i, j, k + 1, N, vsmooth)
1219 czm_x, czm_y, czm_z = _vxb_cached_smooth_at(tvx, tvy, tvz, Bx, By, Bz, i, j, k - 1, N, vsmooth)
1220 cxp_x, cxp_y, cxp_z = _vxb_cached_smooth_at(tvx, tvy, tvz, Bx, By, Bz, i + 1, j, k, N, vsmooth)
1221 cxm_x, cxm_y, cxm_z = _vxb_cached_smooth_at(tvx, tvy, tvz, Bx, By, Bz, i - 1, j, k, N, vsmooth)
1222
1223 nbx = Bx[c] + dt * ((cyp_z - cym_z) - (czp_y - czm_y)) * inv2dx
1224 nby = By[c] + dt * ((czp_x - czm_x) - (cxp_z - cxm_z)) * inv2dx
1225 nbz = Bz[c] + dt * ((cxp_y - cxm_y) - (cyp_x - cym_x)) * inv2dx
1226 nmx = nrho * vx
1227 nmy = nrho * vy
1228 nmz = nrho * vz
1229 ek1 = T(0.5) * (nmx * nmx + nmy * nmy + nmz * nmz) / nrho
1230 eb1 = T(0.5) * (nbx * nbx + nby * nby + nbz * nbz)
1231 orho[c] = nrho
1232 omx[c] = nmx; omy[c] = nmy; omz[c] = nmz
1233 oE[c] = max(eint * (nrho / r) + ek1 + eb1, smallE)
1234 oBx[c] = nbx; oBy[c] = nby; oBz[c] = nbz
1235 opsi[c] = psi[c]
1236 end
1237 end
1238
1239 @kernel function _terminal_velocity_step_k!(orho, omx, omy, omz, oE, oBx, oBy, oBz, opsi,
1240 @Const(rho), @Const(mx), @Const(my), @Const(mz), @Const(E),
1241 @Const(Bx), @Const(By), @Const(Bz), @Const(psi),
1242 N::Int, dt, inv2dx, gamma_code, gamma_gas, smallr, smallE, vcap)
1243 c = @index(Global)
1244 @inbounds begin
1245 T = eltype(E)
1246 i = (c - 1) % N + 1
1247 j = ((c - 1) ÷ N) % N + 1
1248 k = (c - 1) ÷ (N * N) + 1
1249
1250 r = max(rho[c], smallr)
1251 ek0 = T(0.5) * (mx[c] * mx[c] + my[c] * my[c] + mz[c] * mz[c]) / r
1252 eb0 = T(0.5) * (Bx[c] * Bx[c] + By[c] * By[c] + Bz[c] * Bz[c])
1253 eint = max(E[c] - ek0 - eb0, smallE)
1254
1255 vx, vy, vz = _terminal_v_at(rho, mx, my, mz, E, Bx, By, Bz, i, j, k, N, inv2dx, gamma_code, gamma_gas, smallr, smallE, vcap)
1256 vxp_x, _, _ = _terminal_v_at(rho, mx, my, mz, E, Bx, By, Bz, i + 1, j, k, N, inv2dx, gamma_code, gamma_gas, smallr, smallE, vcap)
1257 vxm_x, _, _ = _terminal_v_at(rho, mx, my, mz, E, Bx, By, Bz, i - 1, j, k, N, inv2dx, gamma_code, gamma_gas, smallr, smallE, vcap)
1258 _, vyp_y, _ = _terminal_v_at(rho, mx, my, mz, E, Bx, By, Bz, i, j + 1, k, N, inv2dx, gamma_code, gamma_gas, smallr, smallE, vcap)
1259 _, vym_y, _ = _terminal_v_at(rho, mx, my, mz, E, Bx, By, Bz, i, j - 1, k, N, inv2dx, gamma_code, gamma_gas, smallr, smallE, vcap)
1260 _, _, vzp_z = _terminal_v_at(rho, mx, my, mz, E, Bx, By, Bz, i, j, k + 1, N, inv2dx, gamma_code, gamma_gas, smallr, smallE, vcap)
1261 _, _, vzm_z = _terminal_v_at(rho, mx, my, mz, E, Bx, By, Bz, i, j, k - 1, N, inv2dx, gamma_code, gamma_gas, smallr, smallE, vcap)
1262 cxp = _lin3(i + 1, j, k, N); cxm = _lin3(i - 1, j, k, N)
1263 cyp = _lin3(i, j + 1, k, N); cym = _lin3(i, j - 1, k, N)
1264 czp = _lin3(i, j, k + 1, N); czm = _lin3(i, j, k - 1, N)
1265 vxf_p = T(0.5) * (vx + vxp_x)
1266 vxf_m = T(0.5) * (vxm_x + vx)
1267 vyf_p = T(0.5) * (vy + vyp_y)
1268 vyf_m = T(0.5) * (vym_y + vy)
1269 vzf_p = T(0.5) * (vz + vzp_z)
1270 vzf_m = T(0.5) * (vzm_z + vz)
1271 fx_p = vxf_p * max(ifelse(vxf_p >= zero(T), rho[c], rho[cxp]), smallr)
1272 fx_m = vxf_m * max(ifelse(vxf_m >= zero(T), rho[cxm], rho[c]), smallr)
1273 fy_p = vyf_p * max(ifelse(vyf_p >= zero(T), rho[c], rho[cyp]), smallr)
1274 fy_m = vyf_m * max(ifelse(vyf_m >= zero(T), rho[cym], rho[c]), smallr)
1275 fz_p = vzf_p * max(ifelse(vzf_p >= zero(T), rho[c], rho[czp]), smallr)
1276 fz_m = vzf_m * max(ifelse(vzf_m >= zero(T), rho[czm], rho[c]), smallr)
1277 div_rhov = ((fx_p - fx_m) + (fy_p - fy_m) + (fz_p - fz_m)) * (T(2) * inv2dx)
1278 nrho = max(rho[c] - dt * div_rhov, smallr)
1279 cyp_x, cyp_y, cyp_z = _vxb_at(rho, mx, my, mz, E, Bx, By, Bz, i, j + 1, k, N, inv2dx, gamma_code, gamma_gas, smallr, smallE, vcap)
1280 cym_x, cym_y, cym_z = _vxb_at(rho, mx, my, mz, E, Bx, By, Bz, i, j - 1, k, N, inv2dx, gamma_code, gamma_gas, smallr, smallE, vcap)
1281 czp_x, czp_y, czp_z = _vxb_at(rho, mx, my, mz, E, Bx, By, Bz, i, j, k + 1, N, inv2dx, gamma_code, gamma_gas, smallr, smallE, vcap)
1282 czm_x, czm_y, czm_z = _vxb_at(rho, mx, my, mz, E, Bx, By, Bz, i, j, k - 1, N, inv2dx, gamma_code, gamma_gas, smallr, smallE, vcap)
1283 cxp_x, cxp_y, cxp_z = _vxb_at(rho, mx, my, mz, E, Bx, By, Bz, i + 1, j, k, N, inv2dx, gamma_code, gamma_gas, smallr, smallE, vcap)
1284 cxm_x, cxm_y, cxm_z = _vxb_at(rho, mx, my, mz, E, Bx, By, Bz, i - 1, j, k, N, inv2dx, gamma_code, gamma_gas, smallr, smallE, vcap)
1285
1286 nbx = Bx[c] + dt * ((cyp_z - cym_z) - (czp_y - czm_y)) * inv2dx
1287 nby = By[c] + dt * ((czp_x - czm_x) - (cxp_z - cxm_z)) * inv2dx
1288 nbz = Bz[c] + dt * ((cxp_y - cxm_y) - (cyp_x - cym_x)) * inv2dx
1289 nmx = nrho * vx
1290 nmy = nrho * vy
1291 nmz = nrho * vz
1292 ek1 = T(0.5) * (nmx * nmx + nmy * nmy + nmz * nmz) / nrho
1293 eb1 = T(0.5) * (nbx * nbx + nby * nby + nbz * nbz)
1294 orho[c] = nrho
1295 omx[c] = nmx; omy[c] = nmy; omz[c] = nmz
1296 oE[c] = max(eint * (nrho / r) + ek1 + eb1, smallE)
1297 oBx[c] = nbx; oBy[c] = nby; oBz[c] = nbz
1298 opsi[c] = psi[c]
1299 end
1300 end
1301
1302 # Centred pressure-implicit finalize. Every array argument is distinct (no output
1303 # aliases an input): this kernel reads rho_new, fxm/fym/fzm and Bx/By/Bz at
1304 # neighbouring cells (pressure gradient and the centred curl of v x B), so any
1305 # output sharing storage with them would be a cross-work-item data race
1306 # (thread/GPU-scheduling dependent results). `eint` is the thermal energy
1307 # stashed by `_pressure_implicit_stash_eint_k!`; `rho` is read at own cell only.
1308 # psi is not written (the caller keeps U[9] across the step).
1309 @kernel function _pressure_implicit_finalize_k!(orho, omx, omy, omz, oE, oBx, oBy, oBz,
1310 @Const(rho), @Const(eint),
1311 @Const(Bx), @Const(By), @Const(Bz),
1312 @Const(rho_new), @Const(fxm), @Const(fym), @Const(fzm),
1313 N::Int, dt, inv2dx, gamma_code, pressure_coeff,
1314 gamma_gas, smallr, smallE, vcap)
1315 c = @index(Global)
1316 @inbounds begin
1317 T = eltype(eint)
1318 i = (c - 1) % N + 1
1319 j = ((c - 1) ÷ N) % N + 1
1320 k = (c - 1) ÷ (N * N) + 1
1321
1322 r0 = max(rho[c], smallr)
1323
1324 nrho = max(rho_new[c], smallr)
1325 vx, vy, vz = _pressure_implicit_v_at(rho_new, fxm, fym, fzm, i, j, k, N,
1326 inv2dx, gamma_code, pressure_coeff,
1327 smallr, vcap)
1328 cyp_x, cyp_y, cyp_z = _vxb_pressure_implicit_at(rho_new, fxm, fym, fzm, Bx, By, Bz,
1329 i, j + 1, k, N, inv2dx, gamma_code,
1330 pressure_coeff, smallr, vcap)
1331 cym_x, cym_y, cym_z = _vxb_pressure_implicit_at(rho_new, fxm, fym, fzm, Bx, By, Bz,
1332 i, j - 1, k, N, inv2dx, gamma_code,
1333 pressure_coeff, smallr, vcap)
1334 czp_x, czp_y, czp_z = _vxb_pressure_implicit_at(rho_new, fxm, fym, fzm, Bx, By, Bz,
1335 i, j, k + 1, N, inv2dx, gamma_code,
1336 pressure_coeff, smallr, vcap)
1337 czm_x, czm_y, czm_z = _vxb_pressure_implicit_at(rho_new, fxm, fym, fzm, Bx, By, Bz,
1338 i, j, k - 1, N, inv2dx, gamma_code,
1339 pressure_coeff, smallr, vcap)
1340 cxp_x, cxp_y, cxp_z = _vxb_pressure_implicit_at(rho_new, fxm, fym, fzm, Bx, By, Bz,
1341 i + 1, j, k, N, inv2dx, gamma_code,
1342 pressure_coeff, smallr, vcap)
1343 cxm_x, cxm_y, cxm_z = _vxb_pressure_implicit_at(rho_new, fxm, fym, fzm, Bx, By, Bz,
1344 i - 1, j, k, N, inv2dx, gamma_code,
1345 pressure_coeff, smallr, vcap)
1346
1347 nbx = Bx[c] + dt * ((cyp_z - cym_z) - (czp_y - czm_y)) * inv2dx
1348 nby = By[c] + dt * ((czp_x - czm_x) - (cxp_z - cxm_z)) * inv2dx
1349 nbz = Bz[c] + dt * ((cxp_y - cxm_y) - (cyp_x - cym_x)) * inv2dx
1350 nmx = nrho * vx
1351 nmy = nrho * vy
1352 nmz = nrho * vz
1353 ek1 = T(0.5) * (nmx * nmx + nmy * nmy + nmz * nmz) / nrho
1354 eb1 = T(0.5) * (nbx * nbx + nby * nby + nbz * nbz)
1355 orho[c] = nrho
1356 omx[c] = nmx; omy[c] = nmy; omz[c] = nmz
1357 oE[c] = max(eint[c] * (nrho / r0)^gamma_gas + ek1 + eb1, smallE)
1358 oBx[c] = nbx; oBy[c] = nby; oBz[c] = nbz
1359 end
1360 end
1361
1362 # In-place own-cell stash of the thermal energy into E (single array, one work-item per
1363 # index, reads precede the write; mx/my/mz/rho/B are separate arrays).
1364 @kernel function _pressure_implicit_stash_eint_k!(E, @Const(rho), @Const(mx), @Const(my), @Const(mz),
1365 @Const(Bx), @Const(By), @Const(Bz), smallr, smallE)
1366 c = @index(Global)
1367 @inbounds begin
1368 T = eltype(E)
1369 r0 = max(rho[c], smallr)
1370 ek0 = T(0.5) * (mx[c] * mx[c] + my[c] * my[c] + mz[c] * mz[c]) / r0
1371 eb0 = T(0.5) * (Bx[c] * Bx[c] + By[c] * By[c] + Bz[c] * Bz[c])
1372 E[c] = max(E[c] - ek0 - eb0, smallE)
1373 end
1374 end
1375
1376 @kernel function _stash_eint_and_force_k!(eint, mx, my, mz,
1377 @Const(rho), @Const(E),
1378 @Const(Bx), @Const(By), @Const(Bz),
1379 @Const(fx), @Const(fy), @Const(fz),
1380 smallr, smallE)
1381 c = @index(Global)
1382 @inbounds begin
1383 T = eltype(E)
1384 r = max(rho[c], smallr)
1385 ek = T(0.5) * (mx[c] * mx[c] + my[c] * my[c] + mz[c] * mz[c]) / r
1386 eb = T(0.5) * (Bx[c] * Bx[c] + By[c] * By[c] + Bz[c] * Bz[c])
1387 eint[c] = max(E[c] - ek - eb, smallE)
1388 mx[c] = fx[c]; my[c] = fy[c]; mz[c] = fz[c]
1389 end
1390 end
1391
1392 @kernel function _pressure_implicit_hydro_finalize_k!(orho, omx, omy, omz, oE, oBx, oBy, oBz, opsi,
1393 @Const(rho),
1394 @Const(Bx), @Const(By), @Const(Bz), @Const(psi),
1395 @Const(eint), @Const(rho_new),
1396 @Const(fxm), @Const(fym), @Const(fzm),
1397 N::Int, inv2dx, gamma_code, pressure_coeff,
1398 gamma_gas, smallr, smallE, vcap)
1399 c = @index(Global)
1400 @inbounds begin
1401 T = eltype(eint)
1402 i = (c - 1) % N + 1
1403 j = ((c - 1) ÷ N) % N + 1
1404 k = (c - 1) ÷ (N * N) + 1
1405 r0 = max(rho[c], smallr)
1406 eb = T(0.5) * (Bx[c] * Bx[c] + By[c] * By[c] + Bz[c] * Bz[c])
1407 nr = max(rho_new[c], smallr)
1408 vx, vy, vz = _pressure_implicit_v_at(rho_new, fxm, fym, fzm, i, j, k, N,
1409 inv2dx, gamma_code, pressure_coeff,
1410 smallr, vcap)
1411 nmx = nr * vx
1412 nmy = nr * vy
1413 nmz = nr * vz
1414 ek = T(0.5) * (nmx * nmx + nmy * nmy + nmz * nmz) / nr
1415 orho[c] = nr
1416 omx[c] = nmx; omy[c] = nmy; omz[c] = nmz
1417 oE[c] = max(eint[c] * (nr / r0)^gamma_gas + ek + eb, smallE)
1418 oBx[c] = Bx[c]; oBy[c] = By[c]; oBz[c] = Bz[c]
1419 opsi[c] = psi[c]
1420 end
1421 end
1422
1423 @inline function _minmod3(a, b, c)
1424 if a > zero(a) && b > zero(b) && c > zero(c)
1425 return min(a, min(b, c))
1426 elseif a < zero(a) && b < zero(b) && c < zero(c)
1427 return max(a, max(b, c))
1428 end
1429 return zero(a)
1430 end
1431
1432 @inline function _plm_face(qll, ql, qr, qrr)
1433 T = typeof(ql)
1434 sl = _minmod3(T(1.5) * (ql - qll), T(0.5) * (qr - qll), T(1.5) * (qr - ql))
1435 sr = _minmod3(T(1.5) * (qr - ql), T(0.5) * (qrr - ql), T(1.5) * (qrr - qr))
1436 return ql + T(0.5) * sl, qr - T(0.5) * sr
1437 end
1438
1439 @inline function _induction_flux_x_at(rho, fx, fy, fz, Bx, By, Bz,
1440 i, j, k, N, inv2dx, gamma_code,
1441 pressure_coeff, smallr, vcap)
1442 ll = _lin3(i - 1, j, k, N)
1443 l = _lin3(i, j, k, N)
1444 r = _lin3(i + 1, j, k, N)
1445 rr = _lin3(i + 2, j, k, N)
1446 vlx, vly, vlz = _pressure_implicit_v_at(rho, fx, fy, fz, i, j, k, N,
1447 inv2dx, gamma_code, pressure_coeff, smallr, vcap)
1448 vrx, vry, vrz = _pressure_implicit_v_at(rho, fx, fy, fz, i + 1, j, k, N,
1449 inv2dx, gamma_code, pressure_coeff, smallr, vcap)
1450 a = max(abs(vlx), abs(vrx))
1451 z = zero(a)
1452 bxl, bxr = _plm_face(Bx[ll], Bx[l], Bx[r], Bx[rr])
1453 byl, byr = _plm_face(By[ll], By[l], By[r], By[rr])
1454 bzl, bzr = _plm_face(Bz[ll], Bz[l], Bz[r], Bz[rr])
1455 fby_l = vlx * byl - vly * bxl
1456 fby_r = vrx * byr - vry * bxr
1457 fbz_l = vlx * bzl - vlz * bxl
1458 fbz_r = vrx * bzr - vrz * bxr
1459 half = one(a) / typeof(a)(2)
1460 return z, (fby_l + fby_r - a * (byr - byl)) * half,
1461 (fbz_l + fbz_r - a * (bzr - bzl)) * half
1462 end
1463
1464 @inline function _induction_flux_y_at(rho, fx, fy, fz, Bx, By, Bz,
1465 i, j, k, N, inv2dx, gamma_code,
1466 pressure_coeff, smallr, vcap)
1467 ll = _lin3(i, j - 1, k, N)
1468 l = _lin3(i, j, k, N)
1469 r = _lin3(i, j + 1, k, N)
1470 rr = _lin3(i, j + 2, k, N)
1471 vlx, vly, vlz = _pressure_implicit_v_at(rho, fx, fy, fz, i, j, k, N,
1472 inv2dx, gamma_code, pressure_coeff, smallr, vcap)
1473 vrx, vry, vrz = _pressure_implicit_v_at(rho, fx, fy, fz, i, j + 1, k, N,
1474 inv2dx, gamma_code, pressure_coeff, smallr, vcap)
1475 a = max(abs(vly), abs(vry))
1476 z = zero(a)
1477 bxl, bxr = _plm_face(Bx[ll], Bx[l], Bx[r], Bx[rr])
1478 byl, byr = _plm_face(By[ll], By[l], By[r], By[rr])
1479 bzl, bzr = _plm_face(Bz[ll], Bz[l], Bz[r], Bz[rr])
1480 fbx_l = vly * bxl - vlx * byl
1481 fbx_r = vry * bxr - vrx * byr
1482 fbz_l = vly * bzl - vlz * byl
1483 fbz_r = vry * bzr - vrz * byr
1484 half = one(a) / typeof(a)(2)
1485 return (fbx_l + fbx_r - a * (bxr - bxl)) * half, z,
1486 (fbz_l + fbz_r - a * (bzr - bzl)) * half
1487 end
1488
1489 @inline function _induction_flux_z_at(rho, fx, fy, fz, Bx, By, Bz,
1490 i, j, k, N, inv2dx, gamma_code,
1491 pressure_coeff, smallr, vcap)
1492 ll = _lin3(i, j, k - 1, N)
1493 l = _lin3(i, j, k, N)
1494 r = _lin3(i, j, k + 1, N)
1495 rr = _lin3(i, j, k + 2, N)
1496 vlx, vly, vlz = _pressure_implicit_v_at(rho, fx, fy, fz, i, j, k, N,
1497 inv2dx, gamma_code, pressure_coeff, smallr, vcap)
1498 vrx, vry, vrz = _pressure_implicit_v_at(rho, fx, fy, fz, i, j, k + 1, N,
1499 inv2dx, gamma_code, pressure_coeff, smallr, vcap)
1500 a = max(abs(vlz), abs(vrz))
1501 z = zero(a)
1502 bxl, bxr = _plm_face(Bx[ll], Bx[l], Bx[r], Bx[rr])
1503 byl, byr = _plm_face(By[ll], By[l], By[r], By[rr])
1504 bzl, bzr = _plm_face(Bz[ll], Bz[l], Bz[r], Bz[rr])
1505 fbx_l = vlz * bxl - vlx * bzl
1506 fbx_r = vrz * bxr - vrx * bzr
1507 fby_l = vlz * byl - vly * bzl
1508 fby_r = vrz * byr - vry * bzr
1509 half = one(a) / typeof(a)(2)
1510 return (fbx_l + fbx_r - a * (bxr - bxl)) * half,
1511 (fby_l + fby_r - a * (byr - byl)) * half, z
1512 end
1513
1514 @inline function _induction_rhs_at(rho, fx, fy, fz, Bx, By, Bz,
1515 i, j, k, N, inv2dx, gamma_code,
1516 pressure_coeff, smallr, vcap)
1517 fxp = _induction_flux_x_at(rho, fx, fy, fz, Bx, By, Bz, i, j, k, N,
1518 inv2dx, gamma_code, pressure_coeff, smallr, vcap)
1519 fxm = _induction_flux_x_at(rho, fx, fy, fz, Bx, By, Bz, i - 1, j, k, N,
1520 inv2dx, gamma_code, pressure_coeff, smallr, vcap)
1521 fyp = _induction_flux_y_at(rho, fx, fy, fz, Bx, By, Bz, i, j, k, N,
1522 inv2dx, gamma_code, pressure_coeff, smallr, vcap)
1523 fym = _induction_flux_y_at(rho, fx, fy, fz, Bx, By, Bz, i, j - 1, k, N,
1524 inv2dx, gamma_code, pressure_coeff, smallr, vcap)
1525 fzp = _induction_flux_z_at(rho, fx, fy, fz, Bx, By, Bz, i, j, k, N,
1526 inv2dx, gamma_code, pressure_coeff, smallr, vcap)
1527 fzm = _induction_flux_z_at(rho, fx, fy, fz, Bx, By, Bz, i, j, k - 1, N,
1528 inv2dx, gamma_code, pressure_coeff, smallr, vcap)
1529 invdx = typeof(inv2dx)(2) * inv2dx
1530 return -((fxp[1] - fxm[1]) + (fyp[1] - fym[1]) + (fzp[1] - fzm[1])) * invdx,
1531 -((fxp[2] - fxm[2]) + (fyp[2] - fym[2]) + (fzp[2] - fzm[2])) * invdx,
1532 -((fxp[3] - fxm[3]) + (fyp[3] - fym[3]) + (fzp[3] - fzm[3])) * invdx
1533 end
1534
1535 @kernel function _induction_predict_k!(pBx, pBy, pBz,
1536 @Const(rho), @Const(fx), @Const(fy), @Const(fz),
1537 @Const(Bx), @Const(By), @Const(Bz),
1538 N::Int, dt, inv2dx, gamma_code, pressure_coeff,
1539 smallr, vcap)
1540 c = @index(Global)
1541 @inbounds begin
1542 i = (c - 1) % N + 1
1543 j = ((c - 1) ÷ N) % N + 1
1544 k = (c - 1) ÷ (N * N) + 1
1545 dbx, dby, dbz = _induction_rhs_at(rho, fx, fy, fz, Bx, By, Bz, i, j, k, N,
1546 inv2dx, gamma_code, pressure_coeff, smallr, vcap)
1547 pBx[c] = Bx[c] + dt * dbx
1548 pBy[c] = By[c] + dt * dby
1549 pBz[c] = Bz[c] + dt * dbz
1550 end
1551 end
1552
1553 @kernel function _induction_correct_k!(Bx, By, Bz, E,
1554 @Const(rho), @Const(fx), @Const(fy), @Const(fz),
1555 @Const(pBx), @Const(pBy), @Const(pBz),
1556 N::Int, dt, inv2dx, gamma_code, pressure_coeff,
1557 smallr, smallE, vcap)
1558 c = @index(Global)
1559 @inbounds begin
1560 T = eltype(E)
1561 i = (c - 1) % N + 1
1562 j = ((c - 1) ÷ N) % N + 1
1563 k = (c - 1) ÷ (N * N) + 1
1564 dbx, dby, dbz = _induction_rhs_at(rho, fx, fy, fz, pBx, pBy, pBz, i, j, k, N,
1565 inv2dx, gamma_code, pressure_coeff, smallr, vcap)
1566 obx = Bx[c]; oby = By[c]; obz = Bz[c]
1567 nbx = T(0.5) * (obx + pBx[c] + dt * dbx)
1568 nby = T(0.5) * (oby + pBy[c] + dt * dby)
1569 nbz = T(0.5) * (obz + pBz[c] + dt * dbz)
1570 eb0 = T(0.5) * (obx * obx + oby * oby + obz * obz)
1571 eb1 = T(0.5) * (nbx * nbx + nby * nby + nbz * nbz)
1572 Bx[c] = nbx; By[c] = nby; Bz[c] = nbz
1573 E[c] = max(E[c] + eb1 - eb0, smallE)
1574 end
1575 end
1576
1577 @kernel function _divb_centered_k!(divb, @Const(Bx), @Const(By), @Const(Bz), N::Int, inv2dx)
1578 c = @index(Global)
1579 @inbounds begin
1580 i = (c - 1) % N + 1
1581 j = ((c - 1) ÷ N) % N + 1
1582 k = (c - 1) ÷ (N * N) + 1
1583 divb[c] = ((Bx[_lin3(i + 1, j, k, N)] - Bx[_lin3(i - 1, j, k, N)]) +
1584 (By[_lin3(i, j + 1, k, N)] - By[_lin3(i, j - 1, k, N)]) +
1585 (Bz[_lin3(i, j, k + 1, N)] - Bz[_lin3(i, j, k - 1, N)])) * inv2dx
1586 end
1587 end
1588
1589 @kernel function _project_b_centered_k!(Bx, By, Bz, E, @Const(phi),
1590 energy_residual,N::Int,inv2dx,smallE,
1591 track_energy::Val{TRACK}) where {TRACK}
1592 c = @index(Global)
1593 @inbounds begin
1594 T = eltype(E)
1595 i = (c - 1) % N + 1
1596 j = ((c - 1) ÷ N) % N + 1
1597 k = (c - 1) ÷ (N * N) + 1
1598 obx = Bx[c]; oby = By[c]; obz = Bz[c]
1599 nbx = obx - (phi[_lin3(i + 1, j, k, N)] - phi[_lin3(i - 1, j, k, N)]) * inv2dx
1600 nby = oby - (phi[_lin3(i, j + 1, k, N)] - phi[_lin3(i, j - 1, k, N)]) * inv2dx
1601 nbz = obz - (phi[_lin3(i, j, k + 1, N)] - phi[_lin3(i, j, k - 1, N)]) * inv2dx
1602 eb0 = T(0.5) * (obx * obx + oby * oby + obz * obz)
1603 eb1 = T(0.5) * (nbx * nbx + nby * nby + nbz * nbz)
1604 Bx[c] = nbx; By[c] = nby; Bz[c] = nbz
1605 E[c] = MHDKernels._compensated_floor_update!(
1606 energy_residual,c,E[c],eb1-eb0,smallE,track_energy)
1607 end
1608 end
1609
1610 @kernel function _pressure_velocity_speed_k!(speed, @Const(rho), @Const(fx), @Const(fy), @Const(fz),
1611 N::Int, inv2dx, gamma_code, pressure_coeff,
1612 smallr, vcap)
1613 c = @index(Global)
1614 @inbounds begin
1615 i = (c - 1) % N + 1
1616 j = ((c - 1) ÷ N) % N + 1
1617 k = (c - 1) ÷ (N * N) + 1
1618 vx, vy, vz = _pressure_implicit_v_at(rho, fx, fy, fz, i, j, k, N,
1619 inv2dx, gamma_code, pressure_coeff,
1620 smallr, vcap)
1621 speed[c] = sqrt(vx * vx + vy * vy + vz * vz)
1622 end
1623 end
1624
1625 @kernel function _terminal_nonlinear_limit_k!(rho, mx, my, mz, E, Bx, By, Bz,
1626 @Const(rho0), @Const(Bx0), @Const(By0), @Const(Bz0),
1627 smallr, smallE, rho_rel_limit, b_rel_limit)
1628 c = @index(Global)
1629 @inbounds begin
1630 T = eltype(E)
1631 r_before = max(rho[c], smallr)
1632 r_after = r_before
1633 if isfinite(rho_rel_limit)
1634 rref = max(rho0[c], smallr)
1635 rlo = rref / (one(T) + rho_rel_limit)
1636 rhi = rref * (one(T) + rho_rel_limit)
1637 r_after = min(max(r_before, max(rlo, smallr)), rhi)
1638 end
1639
1640 bx_before = Bx[c]
1641 by_before = By[c]
1642 bz_before = Bz[c]
1643 bx_after = bx_before
1644 by_after = by_before
1645 bz_after = bz_before
1646 if isfinite(b_rel_limit)
1647 dbx = bx_before - Bx0[c]
1648 dby = by_before - By0[c]
1649 dbz = bz_before - Bz0[c]
1650 bmag = sqrt(Bx0[c] * Bx0[c] + By0[c] * By0[c] + Bz0[c] * Bz0[c])
1651 dbmag = sqrt(dbx * dbx + dby * dby + dbz * dbz)
1652 maxdb = b_rel_limit * max(bmag, smallE)
1653 if dbmag > maxdb
1654 fac = maxdb / max(dbmag, smallE)
1655 bx_after = Bx0[c] + fac * dbx
1656 by_after = By0[c] + fac * dby
1657 bz_after = Bz0[c] + fac * dbz
1658 end
1659 end
1660
1661 if r_after != r_before || bx_after != bx_before || by_after != by_before || bz_after != bz_before
1662 mom_scale = r_after / max(r_before, smallr)
1663 ek_before = T(0.5) * (mx[c] * mx[c] + my[c] * my[c] + mz[c] * mz[c]) / max(r_before, smallr)
1664 eb_before = T(0.5) * (bx_before * bx_before + by_before * by_before + bz_before * bz_before)
1665 eint = max(E[c] - ek_before - eb_before, smallE)
1666 mx[c] *= mom_scale
1667 my[c] *= mom_scale
1668 mz[c] *= mom_scale
1669 ek_after = T(0.5) * (mx[c] * mx[c] + my[c] * my[c] + mz[c] * mz[c]) / max(r_after, smallr)
1670 eb_after = T(0.5) * (bx_after * bx_after + by_after * by_after + bz_after * bz_after)
1671 rho[c] = r_after
1672 Bx[c] = bx_after
1673 By[c] = by_after
1674 Bz[c] = bz_after
1675 E[c] = max(eint * (r_after / max(r_before, smallr)) + ek_after + eb_after, smallE)
1676 end
1677 end
1678 end
1679
1680 # CT: the first-order velocity-step paths below also write cell-centred B (centred curl of
1681 # v x B). For mhd_solver=ct the cell v x B is computed from the same velocity into the CT
1682 # workspace before the cell kernel, the faces then advance by the curl of the edge-averaged
1683 # EMF, and cell B / total energy are re-derived from the faces after the kernel.
1684 @kernel function _terminal_ct_cell_emf_k!(ecx, ecy, ecz,
1685 @Const(rho), @Const(mx), @Const(my), @Const(mz), @Const(E),
1686 @Const(Bx), @Const(By), @Const(Bz),
1687 N::Int, inv2dx, gamma_code, gamma_gas, smallr, smallE, vcap)
1688 c = @index(Global)
1689 @inbounds begin
1690 i = (c - 1) % N + 1
1691 j = ((c - 1) ÷ N) % N + 1
1692 k = (c - 1) ÷ (N * N) + 1
1693 x, y, z = _vxb_at(rho, mx, my, mz, E, Bx, By, Bz, i, j, k, N, inv2dx,
1694 gamma_code, gamma_gas, smallr, smallE, vcap)
1695 ecx[c] = x; ecy[c] = y; ecz[c] = z
1696 end
1697 end
1698
1699 @kernel function _terminal_ct_cell_emf_cached_k!(ecx, ecy, ecz, @Const(tvx), @Const(tvy), @Const(tvz),
1700 @Const(Bx), @Const(By), @Const(Bz),
1701 N::Int, vsmooth)
1702 c = @index(Global)
1703 @inbounds begin
1704 i = (c - 1) % N + 1
1705 j = ((c - 1) ÷ N) % N + 1
1706 k = (c - 1) ÷ (N * N) + 1
1707 x, y, z = _vxb_cached_smooth_at(tvx, tvy, tvz, Bx, By, Bz, i, j, k, N, vsmooth)
1708 ecx[c] = x; ecy[c] = y; ecz[c] = z
1709 end
1710 end
1711
1712 function _terminal_velocity_step!(be, s, dt::Real, gamma_code::Real, vcap::Real)
1713 if MHDKernels.uses_ct(s)
1714 ct = s.ct
1715 _terminal_ct_cell_emf_k!(be)(ct.ecx, ct.ecy, ct.ecz, s.U[1], s.U[2], s.U[3], s.U[4],
1716 s.U[5], s.U[6], s.U[7], s.U[8],
1717 Int(s.dims[1]), T(1) / (T(2) * s.dx),
1718 T(gamma_code), T(s.γ), T(s.smallr), T(1e-30), T(vcap);
1719 ndrange=length(s.U[1]))
1720 KA.synchronize(be)
1721 end
1722 _terminal_velocity_step_k!(be)(s.scratch..., s.U...,
1723 Int(s.dims[1]), T(dt), T(1) / (T(2) * s.dx),
1724 T(gamma_code), T(s.γ), T(s.smallr), T(1e-30), T(vcap);
1725 ndrange=length(s.U[1]))
1726 KA.synchronize(be)
1727 s.U, s.scratch = s.scratch, s.U
1728 MHDKernels.uses_ct(s) && MHDKernels.terminal_ct_euler_from_cell_emf!(s, dt)
1729 return nothing
1730 end
1731
1732 function _terminal_pressure_exponential_solve!(dst, state, source, coeff_cells2, dt)
1733 PoissonKernels.rfft_exponential_source!(dst, state, source;
1734 coeff_cells2=coeff_cells2, dt=dt)
1735 return dst
1736 end
1737
1738 function _project_mhd_b!(be, s; energy_compensation=nothing)
1739 N = Int(s.dims[1])
1740 inv2dx = T(1) / (T(2) * s.dx)
1741 divb = s.scratch[1]
1742 phi = s.scratch[5]
1743 _divb_centered_k!(be)(divb, s.U[6], s.U[7], s.U[8], N, inv2dx;
1744 ndrange=length(s.U[1]))
1745 PoissonKernels.fft_poisson_rfft!(reshape(phi, N, N, N),
1746 reshape(divb, N, N, N);
1747 boxsize=Float64(N) * Float64(s.dx),
1748 greens=:centered)
1749 residual=energy_compensation===nothing ? s.U[5] : energy_compensation
1750 _project_b_centered_k!(be)(s.U[6],s.U[7],s.U[8],s.U[5],phi,residual,
1751 N,inv2dx,T(1e-30),Val(energy_compensation!==nothing);
1752 ndrange=length(s.U[1]))
1753 fill!(s.U[9], zero(T))
1754 KA.synchronize(be)
1755 return nothing
1756 end
1757
1758 function _terminal_pressure_implicit_step!(be, s, dt::Real, gamma_code::Real,
1759 vcap::Real, pressure_coeff::Real,
1760 rho_rel_limit::Real, b_rel_limit::Real;
1761 induction::Symbol=:centered,
1762 induction_cfl::Real=0.4,
1763 induction_dissipation::Real=0.05,
1764 induction_dissipation_order::Int=4,
1765 project_b::Bool=true,
1766 max_induction_subcycles::Int=256)
1767 force_s = 0.0
1768 pressure_fft_s = 0.0
1769 finalize_s = 0.0
1770 induction_s = 0.0
1771 projection_s = 0.0
1772 if !(isfinite(gamma_code) && gamma_code > 0)
1773 finalize_s = @elapsed _terminal_velocity_step!(be, s, dt, gamma_code, vcap)
1774 return (nsub=1, force_s, pressure_fft_s, finalize_s, induction_s, projection_s)
1775 end
1776 if MHDKernels.uses_ct(s) && induction !== :ct
1777 error("MHD_SOLVER=ct with MHD_TERMINAL_PRESSURE_IMPLICIT requires " *
1778 "MHD_TERMINAL_INDUCTION=ct (the face-staggered SSPRK2 induction); " *
1779 "terminal_induction=$induction advances cell-centred B with a centred curl " *
1780 "(and ssprk2 relies on a GLM-style B projection), which cannot keep CT faces consistent")
1781 end
1782 if induction === :ct
1783 result = MHDKernels.terminal_pressure_ct_step!(_terminal_pressure_exponential_solve!, s, dt;
1784 gamma_drag=gamma_code, pressure_coeff=pressure_coeff,
1785 induction_cfl=induction_cfl,
1786 induction_dissipation=induction_dissipation,
1787 induction_dissipation_order=induction_dissipation_order,
1788 vcap=vcap, max_induction_subcycles=max_induction_subcycles)
1789 # CT: the face update keeps div b at round-off, so the cell-centred B projection
1790 # (which would overwrite the face average) is skipped.
1791 if project_b && !MHDKernels.uses_ct(s)
1792 projection_s = @elapsed _project_mhd_b!(be, s)
1793 return merge(result, (; projection_s))
1794 end
1795 return result
1796 end
1797 N = Int(s.dims[1])
1798 inv2dx = T(1) / (T(2) * s.dx)
1799 fx = s.scratch[2]
1800 fy = s.scratch[3]
1801 fz = s.scratch[4]
1802 source = s.scratch[1]
1803 sol = s.scratch[5]
1804 force_s = @elapsed begin
1805 _magnetic_force_field_k!(be)(fx, fy, fz, s.U[6], s.U[7], s.U[8],
1806 N, inv2dx; ndrange=length(s.U[1]))
1807 _pressure_implicit_source_k!(be)(source, s.U[1], fx, fy, fz,
1808 N, inv2dx, T(gamma_code), T(s.smallr), T(vcap);
1809 ndrange=length(s.U[1]))
1810 KA.synchronize(be)
1811 end
1812 coeff_cells2 = Float64(dt) * max(Float64(pressure_coeff), 0.0) /
1813 max(Float64(gamma_code), eps(Float64)) /
1814 max(Float64(s.dx)^2, eps(Float64))
1815 pressure_fft_s = @elapsed begin
1816 PoissonKernels.rfft_exponential_source!(reshape(sol, N, N, N),
1817 reshape(s.U[1], N, N, N),
1818 reshape(source, N, N, N);
1819 coeff_cells2=coeff_cells2, dt=dt)
1820 copyto!(source, sol)
1821 KA.synchronize(be)
1822 end
1823 if induction in (:ssprk2, :ct)
1824 # Preserve thermal energy before reusing the old momentum arrays for
1825 # non-aliased Lorentz-force storage during the hydro finalize.
1826 finalize_s = @elapsed begin
1827 _stash_eint_and_force_k!(be)(sol, s.U[2], s.U[3], s.U[4],
1828 s.U[1], s.U[5], s.U[6], s.U[7], s.U[8],
1829 fx, fy, fz, T(s.smallr), T(1e-30);
1830 ndrange=length(s.U[1]))
1831 # Aliasing audit: outputs are s.scratch[1:9], of which scratch[1] (= `source` = rho_new)
1832 # is read at neighbour cells by `_pressure_implicit_v_at`. The only write to it is
1833 # orho[c] = max(rho_new[c], smallr), i.e. the value already stored, so the overlap is
1834 # benign unless rho_new < smallr (a floored cell). scratch[5] (= eint) is read at own
1835 # cell only, before the same work-item's write. Force fields are in U[2:4], distinct.
1836 _pressure_implicit_hydro_finalize_k!(be)(s.scratch..., s.U[1],
1837 s.U[6], s.U[7], s.U[8], s.U[9],
1838 sol, source, s.U[2], s.U[3], s.U[4],
1839 N, inv2dx, T(gamma_code), T(pressure_coeff),
1840 T(s.γ), T(s.smallr), T(1e-30), T(vcap);
1841 ndrange=length(s.U[1]))
1842 KA.synchronize(be)
1843 s.U, s.scratch = s.scratch, s.U
1844 end
1845
1846 # The induction transport has its own CFL. The pressure response stays
1847 # at macro cadence; only the nonlinear v x B coupling is subcycled.
1848 fx = s.scratch[2]; fy = s.scratch[3]; fz = s.scratch[4]
1849 speed = s.scratch[1]
1850 _magnetic_force_field_k!(be)(fx, fy, fz, s.U[6], s.U[7], s.U[8],
1851 N, inv2dx; ndrange=length(s.U[1]))
1852 _pressure_velocity_speed_k!(be)(speed, s.U[1], fx, fy, fz,
1853 N, inv2dx, T(gamma_code), T(pressure_coeff),
1854 T(s.smallr), T(vcap); ndrange=length(s.U[1]))
1855 KA.synchronize(be)
1856 vmax = Float64(maximum(speed))
1857 nind = max(1, ceil(Int, Float64(dt) * vmax /
1858 max(Float64(induction_cfl) * Float64(s.dx), eps(Float64))))
1859 nind <= max_induction_subcycles ||
1860 error("terminal induction needs $nind subcycles, above MHD_TERMINAL_INDUCTION_MAX_SUBCYCLES=$max_induction_subcycles")
1861 dti = T(dt / nind)
1862 induction_s = @elapsed begin
1863 for _ in 1:nind
1864 _magnetic_force_field_k!(be)(fx, fy, fz, s.U[6], s.U[7], s.U[8],
1865 N, inv2dx; ndrange=length(s.U[1]))
1866 _induction_predict_k!(be)(s.scratch[6], s.scratch[7], s.scratch[8],
1867 s.U[1], fx, fy, fz, s.U[6], s.U[7], s.U[8],
1868 N, dti, inv2dx, T(gamma_code), T(pressure_coeff),
1869 T(s.smallr), T(vcap); ndrange=length(s.U[1]))
1870 _magnetic_force_field_k!(be)(fx, fy, fz,
1871 s.scratch[6], s.scratch[7], s.scratch[8],
1872 N, inv2dx; ndrange=length(s.U[1]))
1873 _induction_correct_k!(be)(s.U[6], s.U[7], s.U[8], s.U[5],
1874 s.U[1], fx, fy, fz,
1875 s.scratch[6], s.scratch[7], s.scratch[8],
1876 N, dti, inv2dx, T(gamma_code), T(pressure_coeff),
1877 T(s.smallr), T(1e-30), T(vcap);
1878 ndrange=length(s.U[1]))
1879 end
1880 KA.synchronize(be)
1881 end
1882 if project_b
1883 projection_s = @elapsed begin
1884 divb = s.scratch[1]
1885 phi = s.scratch[5]
1886 _divb_centered_k!(be)(divb, s.U[6], s.U[7], s.U[8], N, inv2dx;
1887 ndrange=length(s.U[1]))
1888 PoissonKernels.fft_poisson_rfft!(reshape(phi, N, N, N),
1889 reshape(divb, N, N, N);
1890 boxsize=Float64(N) * Float64(s.dx),
1891 greens=:centered)
1892 _project_b_centered_k!(be)(s.U[6],s.U[7],s.U[8],s.U[5],phi,
1893 s.U[5],N,inv2dx,T(1e-30),Val(false);
1894 ndrange=length(s.U[1]))
1895 KA.synchronize(be)
1896 end
1897 end
1898 return (nsub=nind, force_s, pressure_fft_s, finalize_s, induction_s, projection_s)
1899 elseif induction !== :centered
1900 error("unknown terminal induction mode $induction")
1901 end
1902 # Race-free buffer plan (all arrays handed to one kernel are distinct):
1903 # sc = old scratch, u = old U. sc[1] holds rho_new (`source`), sc[2:4] the force
1904 # fields; both are read at neighbour cells. U[5] is overwritten in place with the
1905 # thermal energy (own cell). Outputs: rho,mx,my,mz,E -> sc[5:9]; Bx,By,Bz -> u[2:4]
1906 # (the old momenta, dead after the stash). u[1] (old rho) and u[6:8] (old B) stay
1907 # intact as the nonlinear-limiter reference. No extra memory is allocated.
1908 # New U = (sc5,sc6,sc7,sc8,sc9,u2,u3,u4,u9); new scratch = (sc1..4,u1,u5,u6,u7,u8).
1909 finalize_s = @elapsed begin
1910 sc = s.scratch
1911 u = s.U
1912 _pressure_implicit_stash_eint_k!(be)(u[5], u[1], u[2], u[3], u[4], u[6], u[7], u[8],
1913 T(s.smallr), T(1e-30);
1914 ndrange=length(s.U[1]))
1915 KA.synchronize(be)
1916 _pressure_implicit_finalize_k!(be)(sc[5], sc[6], sc[7], sc[8], sc[9], u[2], u[3], u[4],
1917 u[1], u[5], u[6], u[7], u[8],
1918 source, fx, fy, fz,
1919 N, T(dt), inv2dx, T(gamma_code), T(pressure_coeff),
1920 T(s.γ), T(s.smallr), T(1e-30), T(vcap);
1921 ndrange=length(s.U[1]))
1922 KA.synchronize(be)
1923 if isfinite(rho_rel_limit) || isfinite(b_rel_limit)
1924 # In place on the new fields (own-cell); references are the distinct old rho, B.
1925 _terminal_nonlinear_limit_k!(be)(sc[5], sc[6], sc[7], sc[8], sc[9], u[2], u[3], u[4],
1926 u[1], u[6], u[7], u[8],
1927 T(s.smallr), T(1e-30), T(rho_rel_limit), T(b_rel_limit);
1928 ndrange=length(s.U[1]))
1929 KA.synchronize(be)
1930 end
1931 s.U = (sc[5], sc[6], sc[7], sc[8], sc[9], u[2], u[3], u[4], u[9])
1932 s.scratch = (sc[1], sc[2], sc[3], sc[4], u[1], u[5], u[6], u[7], u[8])
1933 end
1934 return (nsub=1, force_s, pressure_fft_s, finalize_s, induction_s, projection_s)
1935 end
1936
1937 function _terminal_velocity_step_cached!(be, s, dt::Real, gamma_code::Real, vcap::Real,
1938 tvx, tvy, tvz, vsmooth::Real)
1939 _terminal_velocity_field_k!(be)(tvx, tvy, tvz, s.U[1], s.U[2], s.U[3], s.U[4],
1940 s.U[5], s.U[6], s.U[7], s.U[8],
1941 Int(s.dims[1]), T(1) / (T(2) * s.dx),
1942 T(gamma_code), T(s.γ), T(s.smallr), T(1e-30), T(vcap);
1943 ndrange=length(s.U[1]))
1944 KA.synchronize(be)
1945 if MHDKernels.uses_ct(s)
1946 ct = s.ct
1947 _terminal_ct_cell_emf_cached_k!(be)(ct.ecx, ct.ecy, ct.ecz, tvx, tvy, tvz,
1948 s.U[6], s.U[7], s.U[8], Int(s.dims[1]), T(vsmooth);
1949 ndrange=length(s.U[1]))
1950 KA.synchronize(be)
1951 end
1952 _terminal_velocity_update_cached_k!(be)(s.scratch..., s.U..., tvx, tvy, tvz,
1953 Int(s.dims[1]), T(dt), T(1) / (T(2) * s.dx),
1954 T(s.γ), T(s.smallr), T(1e-30), T(vsmooth);
1955 ndrange=length(s.U[1]))
1956 KA.synchronize(be)
1957 s.U, s.scratch = s.scratch, s.U
1958 MHDKernels.uses_ct(s) && MHDKernels.terminal_ct_euler_from_cell_emf!(s, dt)
1959 return nothing
1960 end
1961
1962 @kernel function _terminal_handoff_momenta_cached_k!(mx, my, mz, E,
1963 @Const(rho), @Const(Bx), @Const(By), @Const(Bz),
1964 @Const(tvx), @Const(tvy), @Const(tvz),
1965 smallr, smallE, alpha)
1966 c = @index(Global)
1967 @inbounds begin
1968 R = eltype(E)
1969 r = max(rho[c], smallr)
1970 ox = mx[c]
1971 oy = my[c]
1972 oz = mz[c]
1973 ek0 = R(0.5) * (ox * ox + oy * oy + oz * oz) / r
1974 eb = R(0.5) * (Bx[c] * Bx[c] + By[c] * By[c] + Bz[c] * Bz[c])
1975 eint = max(E[c] - ek0 - eb, smallE)
1976 a = min(max(alpha, zero(alpha)), one(alpha))
1977 nx = ox + a * (r * tvx[c] - ox)
1978 ny = oy + a * (r * tvy[c] - oy)
1979 nz = oz + a * (r * tvz[c] - oz)
1980 ek1 = R(0.5) * (nx * nx + ny * ny + nz * nz) / r
1981 mx[c] = nx
1982 my[c] = ny
1983 mz[c] = nz
1984 E[c] = max(eint + ek1 + eb, smallE)
1985 end
1986 end
1987
1988 function _terminal_handoff_momenta!(be, s, gamma_code::Real, vcap::Real, alpha::Real;
1989 measure::Bool=false)
1990 tvx = s.scratch[2]
1991 tvy = s.scratch[3]
1992 tvz = s.scratch[4]
1993 _terminal_velocity_field_k!(be)(tvx, tvy, tvz, s.U[1], s.U[2], s.U[3], s.U[4],
1994 s.U[5], s.U[6], s.U[7], s.U[8],
1995 Int(s.dims[1]), T(1) / (T(2) * s.dx),
1996 T(gamma_code), T(s.γ), T(s.smallr), T(1e-30), T(vcap);
1997 ndrange=length(s.U[1]))
1998 KA.synchronize(be)
1999 overlap = (relative=NaN, delta_v_rms=NaN, terminal_v_rms=NaN, full_v_rms=NaN)
2000 if measure
2001 delta2 = s.scratch[1]
2002 terminal2 = s.scratch[5]
2003 full2 = s.scratch[9]
2004 _terminal_velocity_overlap_diag_k!(be)(delta2, terminal2, full2,
2005 s.U[1], s.U[2], s.U[3], s.U[4],
2006 tvx, tvy, tvz, T(s.smallr);
2007 ndrange=length(s.U[1]))
2008 KA.synchronize(be)
2009 n = Float64(length(s.U[1]))
2010 d2 = max(Float64(sum(delta2)), 0.0)
2011 t2 = max(Float64(sum(terminal2)), 0.0)
2012 f2 = max(Float64(sum(full2)), 0.0)
2013 scale2 = 0.5 * (t2 + f2)
2014 relative = scale2 > eps(Float64) ? sqrt(d2 / scale2) : 0.0
2015 overlap = (relative=relative,
2016 delta_v_rms=sqrt(d2 / n),
2017 terminal_v_rms=sqrt(t2 / n),
2018 full_v_rms=sqrt(f2 / n))
2019 end
2020 _terminal_handoff_momenta_cached_k!(be)(s.U[2], s.U[3], s.U[4], s.U[5],
2021 s.U[1], s.U[6], s.U[7], s.U[8],
2022 tvx, tvy, tvz,
2023 T(s.smallr), T(1e-30), T(alpha);
2024 ndrange=length(s.U[1]))
2025 KA.synchronize(be)
2026 return overlap
2027 end
2028
2029 function _hybrid_handoff_alpha(alpha_max::Real, remaining::Int, total::Int)
2030 total <= 1 && return Float64(alpha_max)
2031 progress = clamp(Float64(total - remaining) / Float64(total - 1), 0.0, 1.0)
2032 return Float64(alpha_max) * 0.5 * (1.0 + cos(pi * progress))
2033 end
2034
2035 function _limit_source_dt(dt::Real, dt_explicit::Real, max_factor::Real)
2036 if isfinite(max_factor) && dt > max_factor * dt_explicit
2037 return (dt=Float64(max_factor * dt_explicit), limited=true)
2038 end
2039 return (dt=Float64(dt), limited=false)
2040 end
2041
2042 function _limit_hybrid_aware_dt(dt::Real, dt_explicit::Real, max_factor::Real)
2043 limited_dt = min(Float64(dt), Float64(max_factor * dt_explicit))
2044 nsub = max(1, ceil(Int, limited_dt / max(Float64(dt_explicit), eps(Float64))))
2045 return (dt=limited_dt, nsub=nsub, limited=limited_dt < dt)
2046 end
2047
2048 function _terminal_implicit_diffuse_B!(s, coeff_cells2::Real)
2049 coeff = Float64(coeff_cells2)
2050 if !(isfinite(coeff) && coeff > 0)
2051 return nothing
2052 end
2053 N = Int(s.dims[1])
2054 if MHDKernels.uses_ct(s)
2055 # The implicit diffusion is a translation-invariant Fourier multiplier, so applying it
2056 # to every FACE component commutes with the face divergence (div b stays at
2057 # round-off). Cell B is re-derived as the face average; as for GLM, E is untouched.
2058 ct = s.ct
2059 # The faces are views into one packed buffer; the device FFT wrappers are not
2060 # guaranteed to honour view offsets (Metal gave a corrupted field), so stage each
2061 # face through two full scratch arrays.
2062 for face in (ct.bxf, ct.byf, ct.bzf)
2063 copyto!(s.scratch[6], face)
2064 PoissonKernels.rfft_implicit_diffuse!(reshape(s.scratch[7], N, N, N),
2065 reshape(s.scratch[6], N, N, N); coeff_cells2=coeff)
2066 copyto!(face, s.scratch[7])
2067 end
2068 MHDKernels._ct_bcc_k!(s.be, 256)(s.U[6], s.U[7], s.U[8], ct.bxf, ct.byf, ct.bzf,
2069 N, N, N; ndrange=length(s.U[1]))
2070 KA.synchronize(s.be)
2071 return nothing
2072 end
2073 for v in 6:8
2074 dst = reshape(s.scratch[v], N, N, N)
2075 src = reshape(s.U[v], N, N, N)
2076 PoissonKernels.rfft_implicit_diffuse!(dst, src; coeff_cells2=coeff)
2077 copyto!(s.U[v], s.scratch[v])
2078 end
2079 KA.synchronize(s.be)
2080 return nothing
2081 end
2082
2083 @kernel function _clamp_reduced_chem_k!(HII, H2I, @Const(rho), fh,
2084 hydrogen_reference,
2085 hydrogen_complement_reference,
2086 N::Int, ::Val{CH}) where {CH}
2087 c = @index(Global)
2088 @inbounds begin
2089 T = eltype(HII)
2090 cap = max(fh * rho[c], zero(T))
2091 h2 = min(max(H2I[c], zero(T)), cap)
2092 hii = if CH
2093 xh2 = h2 / max(cap, eps(T))
2094 min(max(HII[c], -hydrogen_reference),
2095 max(hydrogen_complement_reference - xh2,
2096 -hydrogen_reference))
2097 else
2098 min(max(HII[c], zero(T)), max(cap - h2, zero(T)))
2099 end
2100 H2I[c] = h2
2101 HII[c] = hii
2102 end
2103 end
2104
2105 @kernel function _clamp_reduced_helium_k!(HeII, @Const(rho), fh, N::Int)
2106 c = @index(Global)
2107 @inbounds begin
2108 T = eltype(HeII)
2109 cap = max((one(T) - T(fh)) * rho[c], zero(T))
2110 HeII[c] = min(max(HeII[c], zero(T)), cap)
2111 end
2112 end
2113
2114 @kernel function _neutral_h_mass_k!(dst, @Const(rho), @Const(HII), @Const(H2I),
2115 fh, hydrogen_complement_reference, N::Int,
2116 ::Val{NH}, ::Val{CH}) where {NH,CH}
2117 c = @index(Global)
2118 @inbounds begin
2119 T = eltype(dst)
2120 dst[c] = CH ? max(fh * rho[c] *
2121 (hydrogen_complement_reference - HII[c]) - H2I[c],
2122 zero(T)) :
2123 NH ? max(HII[c], zero(T)) :
2124 max(fh * rho[c] - HII[c] - H2I[c], zero(T))
2125 end
2126 end
2127
2128 @kernel function _b2_diag_k!(dst, @Const(Bx), @Const(By), @Const(Bz))
2129 c = @index(Global, Linear)
2130 @inbounds begin
2131 bx = Float32(Bx[c])
2132 by = Float32(By[c])
2133 bz = Float32(Bz[c])
2134 dst[c] = bx*bx + by*by + bz*bz
2135 end
2136 end
2137
2138 @kernel function _divb_abs_periodic_diag_k!(dst, @Const(Bx), @Const(By), @Const(Bz), N::Int, inv2dx)
2139 c = @index(Global, Linear)
2140 @inbounds begin
2141 q = c - 1
2142 i = q % N
2143 j = (q ÷ N) % N
2144 k = q ÷ (N * N)
2145 ip = (i + 1) % N
2146 im = (i + N - 1) % N
2147 jp = (j + 1) % N
2148 jm = (j + N - 1) % N
2149 kp = (k + 1) % N
2150 km = (k + N - 1) % N
2151 lin(ii, jj, kk) = ii + N * jj + N * N * kk + 1
2152 db = (Float32(Bx[lin(ip, j, k)]) - Float32(Bx[lin(im, j, k)]) +
2153 Float32(By[lin(i, jp, k)]) - Float32(By[lin(i, jm, k)]) +
2154 Float32(Bz[lin(i, j, kp)]) - Float32(Bz[lin(i, j, km)])) * Float32(inv2dx)
2155 dst[c] = abs(db)
2156 end
2157 end
2158
2159 @kernel function _component2_diag_k!(dst, @Const(src))
2160 c = @index(Global, Linear)
2161 @inbounds begin
2162 x = Float32(src[c])
2163 dst[c] = x*x
2164 end
2165 end
2166
2167 @kernel function _photon_slip_diag_k!(sx,sy,sz,@Const(gx),@Const(gy),@Const(gz),
2168 @Const(rho),@Const(mx),@Const(my),@Const(mz),smallr)
2169 c=@index(Global)
2170 @inbounds begin
2171 r=max(rho[c],smallr)
2172 sx[c]=gx[c]-mx[c]/r
2173 sy[c]=gy[c]-my[c]/r
2174 sz[c]=gz[c]-mz[c]/r
2175 end
2176 end
2177
2178 @kernel function _delta2_diag_k!(dst, @Const(src), mean)
2179 c = @index(Global, Linear)
2180 @inbounds begin
2181 d = Float32(src[c]) - Float32(mean)
2182 dst[c] = d*d
2183 end
2184 end
2185
2186 @kernel function _density_cos_mode_diag_k!(dst,@Const(rho),N::Int,kmode::Int,
2187 harmonic::Int,mean)
2188 c = @index(Global, Linear)
2189 @inbounds begin
2190 i = (c - 1) % N + 1
2191 x = (Float32(i) - 0.5f0) / Float32(N)
2192 w = cos(2f0 * Float32(pi) * Float32(harmonic * kmode) * x)
2193 dst[c] = (Float32(rho[c]) - Float32(mean)) * w
2194 end
2195 end
2196
2197 @kernel function _scalar_fourier_mode_diag_k!(dst, @Const(src), N::Int, kmode::Int,
2198 quadrature::Int, mean)
2199 c = @index(Global, Linear)
2200 @inbounds begin
2201 T = eltype(dst)
2202 i = (c - 1) % N + 1
2203 x = (T(i) - T(0.5)) / T(N)
2204 phase = T(2) * T(pi) * T(kmode) * x
2205 w = quadrature == 1 ? cos(phase) : sin(phase)
2206 dst[c] = (T(src[c]) - T(mean)) * w
2207 end
2208 end
2209
2210 @kernel function _velocity_fourier_mode_diag_k!(dst, @Const(momentum), @Const(rho),
2211 N::Int, kmode::Int,
2212 quadrature::Int, smallr)
2213 c = @index(Global, Linear)
2214 @inbounds begin
2215 T = eltype(dst)
2216 i = (c - 1) % N + 1
2217 x = (T(i) - T(0.5)) / T(N)
2218 phase = T(2) * T(pi) * T(kmode) * x
2219 w = quadrature == 1 ? cos(phase) : sin(phase)
2220 dst[c] = T(momentum[c]) / max(T(rho[c]), T(smallr)) * w
2221 end
2222 end
2223
2224 @kernel function _slip_fourier_mode_diag_k!(dst, @Const(photon_velocity),
2225 @Const(momentum), @Const(rho),
2226 N::Int, kmode::Int,
2227 quadrature::Int, smallr)
2228 c = @index(Global, Linear)
2229 @inbounds begin
2230 T = eltype(dst)
2231 i = (c - 1) % N + 1
2232 x = (T(i) - T(0.5)) / T(N)
2233 phase = T(2) * T(pi) * T(kmode) * x
2234 w = quadrature == 1 ? cos(phase) : sin(phase)
2235 vb = T(momentum[c]) / max(T(rho[c]), T(smallr))
2236 dst[c] = (T(photon_velocity[c]) - vb) * w
2237 end
2238 end
2239
2240 @kernel function _species_fourier_mode_diag_k!(dst, @Const(species), @Const(rho),
2241 N::Int, kmode::Int,
2242 quadrature::Int, fh, mean)
2243 c = @index(Global, Linear)
2244 @inbounds begin
2245 T = eltype(dst)
2246 i = (c - 1) % N + 1
2247 x = (T(i) - T(0.5)) / T(N)
2248 phase = T(2) * T(pi) * T(kmode) * x
2249 w = quadrature == 1 ? cos(phase) : sin(phase)
2250 fraction = T(species[c]) / max(T(fh) * T(rho[c]), eps(T))
2251 dst[c] = (fraction - T(mean)) * w
2252 end
2253 end
2254
2255 @kernel function _species_fraction_diag_k!(dst, @Const(rho), @Const(sp), fh)
2256 c = @index(Global, Linear)
2257 @inbounds begin
2258 den = max(Float32(fh) * Float32(rho[c]), eps(Float32))
2259 dst[c] = max(Float32(sp[c]) / den, 0f0)
2260 end
2261 end
2262
2263 @kernel function _hydrogen_complement_fraction_diag_k!(dst, @Const(rho),
2264 @Const(hydrogen),
2265 @Const(H2I), fh,
2266 ::Val{NH}) where {NH}
2267 c = @index(Global, Linear)
2268 @inbounds begin
2269 den = max(Float32(fh) * Float32(rho[c]), eps(Float32))
2270 if NH
2271 dst[c] = max((Float32(hydrogen[c]) + Float32(H2I[c])) / den, 0f0)
2272 else
2273 dst[c] = max(1f0 - Float32(hydrogen[c]) / den, 0f0)
2274 end
2275 end
2276 end
2277
2278 @kernel function _centered_hii_fraction_diag_k!(dst, @Const(rho),
2279 @Const(hydrogen_delta),
2280 @Const(H2I), fh,
2281 hydrogen_reference)
2282 c = @index(Global, Linear)
2283 @inbounds begin
2284 R = eltype(dst)
2285 hden = max(R(fh) * R(rho[c]), eps(R))
2286 xh2 = max(R(H2I[c]) / hden, zero(R))
2287 dst[c] = min(max(R(hydrogen_reference) + R(hydrogen_delta[c]), zero(R)),
2288 max(one(R) - xh2, zero(R)))
2289 end
2290 end
2291
2292 @kernel function _total_electron_fraction_diag_k!(dst, @Const(rho), @Const(HII),
2293 @Const(H2I), @Const(HeII),
2294 density_unit, z, fh,
2295 hydrogen_reference,
2296 ::Val{NH}, ::Val{CH}) where {NH,CH}
2297 c = @index(Global, Linear)
2298 @inbounds begin
2299 R = eltype(dst)
2300 r = max(R(rho[c]), eps(R))
2301 hden = max(R(fh) * r, eps(R))
2302 xhii = CH ? max(R(hydrogen_reference) + R(HII[c]), zero(R)) :
2303 NH ? max(one(R) - (R(HII[c]) + R(H2I[c])) / hden, zero(R)) :
2304 max(R(HII[c]) / hden, zero(R))
2305 xheii = max(R(HeII[c]) / max(R(4) * R(fh) * r, eps(R)), zero(R))
2306 nH = max(R(density_unit) * r * R(fh) / R(ChemistryKernels.MH), R(1e-30))
2307 dst[c] = ChemistryKernels.total_electron_fraction(
2308 xhii, xheii, nH, R(_PMF_T_CMB) * (one(R) + R(z)); fh=R(fh))
2309 end
2310 end
2311
2312 @kernel function _thermal_diag_k!(dst, @Const(rho), @Const(mx), @Const(my), @Const(mz),
2313 @Const(E), @Const(Bx), @Const(By), @Const(Bz),
2314 gamma, smallr, smallE, mode::Int)
2315 c = @index(Global, Linear)
2316 @inbounds begin
2317 r = max(Float32(rho[c]), Float32(smallr))
2318 ek = 0.5f0 * (Float32(mx[c]) * Float32(mx[c]) +
2319 Float32(my[c]) * Float32(my[c]) +
2320 Float32(mz[c]) * Float32(mz[c])) / r
2321 eb = 0.5f0 * (Float32(Bx[c]) * Float32(Bx[c]) +
2322 Float32(By[c]) * Float32(By[c]) +
2323 Float32(Bz[c]) * Float32(Bz[c]))
2324 eint = max(Float32(E[c]) - ek - eb, Float32(smallE))
2325 p = (Float32(gamma) - 1f0) * eint
2326 dst[c] = mode == 1 ? p : eint / r
2327 end
2328 end
2329
2330 @kernel function _gas_temperature_diag_k!(dst, @Const(rho), @Const(mx), @Const(my),
2331 @Const(mz), @Const(E), @Const(Bx), @Const(By), @Const(Bz), @Const(HII),
2332 @Const(H2I), velocity_unit2, fh, gamma, smallr,
2333 hydrogen_reference, ::Val{NH}, ::Val{CH}) where {NH,CH}
2334 c = @index(Global, Linear)
2335 @inbounds begin
2336 T = eltype(dst)
2337 r = max(T(rho[c]), T(smallr))
2338 kinetic = T(0.5) * (T(mx[c])^2 + T(my[c])^2 + T(mz[c])^2) / r
2339 magnetic = T(0.5) * (T(Bx[c])^2 + T(By[c])^2 + T(Bz[c])^2)
2340 specific_e = max((T(E[c]) - kinetic - magnetic) / r, T(1e-30))
2341 h2_fraction = max(T(H2I[c]) / r, zero(T))
2342 hii_fraction = CH ? T(fh) * max(T(hydrogen_reference) + T(HII[c]), zero(T)) :
2343 NH ? max(T(fh) - T(HII[c]) / r - h2_fraction, zero(T)) :
2344 max(T(HII[c]) / r, zero(T))
2345 particles_per_mass_h = max(
2346 T(fh) + hii_fraction - T(0.5) * h2_fraction +
2347 (one(T) - T(fh)) / T(4),
2348 eps(T),
2349 )
2350 dst[c] = (T(gamma) - one(T)) * specific_e * T(velocity_unit2) *
2351 T(1.6735575e-24) /
2352 (T(1.380649e-16) * particles_per_mass_h)
2353 end
2354 end
2355
2356 @kernel function _gas_temperature_helium_diag_k!(dst, @Const(rho), @Const(mx),
2357 @Const(my), @Const(mz), @Const(E), @Const(Bx), @Const(By), @Const(Bz),
2358 @Const(HII), @Const(H2I), @Const(HeII), density_unit, velocity_unit2, z,
2359 fh, gamma, smallr, hydrogen_reference,
2360 ::Val{NH}, ::Val{CH}) where {NH,CH}
2361 c = @index(Global, Linear)
2362 @inbounds begin
2363 R = eltype(dst)
2364 r = max(R(rho[c]), R(smallr))
2365 kinetic = R(0.5) * (R(mx[c])^2 + R(my[c])^2 + R(mz[c])^2) / r
2366 magnetic = R(0.5) * (R(Bx[c])^2 + R(By[c])^2 + R(Bz[c])^2)
2367 specific_e = max((R(E[c]) - kinetic - magnetic) / r, R(1e-30))
2368 du = R(density_unit)
2369 hii_code = CH ? max(R(fh) * r *
2370 (R(hydrogen_reference) + R(HII[c])), zero(R)) :
2371 NH ? max(R(fh) * r - R(HII[c]) - R(H2I[c]), zero(R)) :
2372 R(HII[c])
2373 dst[c] = ChemistryKernels.temperature_from_reduced_helium(
2374 r * du, specific_e * R(velocity_unit2), hii_code * du,
2375 R(H2I[c]) * du, R(HeII[c]) * du, R(z);
2376 fh=R(fh), gamma=R(gamma))
2377 end
2378 end
2379
2380 @kernel function _logrho_hist_k!(counts, @Const(rho), nbins::Int, loglo, invw)
2381 c = @index(Global, Linear)
2382 @inbounds begin
2383 lr = log10(max(Float32(rho[c]), 1.0f-30))
2384 b = unsafe_trunc(Int, floor((lr - Float32(loglo)) * Float32(invw))) + 1
2385 if 1 <= b <= nbins
2386 KA.@atomic counts[b] += Int32(1)
2387 end
2388 end
2389 end
2390
2391 @inline function _peebles_recomb_cell(rho, e, hii_m, h2_m, dt, z, hz, fh, gamma,
2392 f_alpha, xe_mean, nsm_mass, nsm_neutral, fudge, gauss, dtfrac, itcap)
2393 R = typeof(e)
2394 mh = R(ChemistryKernels.MH)
2395 kb = R(ChemistryKernels.KBOLTZ)
2396 tiny = R(1e-30)
2397 fhd = max(R(fh) * rho, tiny)
2398 # Species are mass densities here. A fixed 1e-30 g/cm^3 floor creates a
2399 # spurious H2 fraction as the cosmological density falls; zero is safe
2400 # because number-density denominators below retain their own unitful floor.
2401 h2 = min(max(h2_m, zero(R)), fhd)
2402 hii = min(max(hii_m, zero(R)), max(fhd - h2, zero(R)))
2403 nHeI = max((one(R) - R(fh)) * rho / (R(4) * mh), tiny)
2404 tleft = dt
2405 iter = 0
2406 while tleft > zero(R) && iter < itcap
2407 iter += 1
2408 nHII = max(hii / mh, tiny)
2409 nH2I = max(h2 / mh, tiny)
2410 nHI = max((fhd - hii - h2) / mh, tiny)
2411 ne = nHII
2412 ntot = max(nHI + nHII + ne + nHeI + R(0.5) * nH2I, tiny)
2413 Tgas = max((R(gamma) - one(R)) * rho * e / (kb * ntot), one(R))
2414 Trad = R(2.725) * (one(R) + R(z))
2415 nsm_h = nsm_neutral == 1 ? max(nsm_mass / mh, tiny) :
2416 max(nsm_mass * R(fh) / mh, tiny) * max(one(R) - R(xe_mean), zero(R))
2417 nHIeff = (one(R) - R(f_alpha)) * nHI + R(f_alpha) * nsm_h
2418 k2 = ChemistryKernels.peebles_k2_mixing(
2419 Tgas, nHI, nHIeff, hz; fudge=fudge, gauss=gauss, Trad=Trad)
2420 beta = ChemistryKernels.beta1s_freq(Trad) * k2 / (ChemistryKernels.recfast_alpha(Tgas) * R(1.0e6))
2421 kion = ChemistryKernels.k1(Tgas)
2422 nHtot = max((fhd - h2) / mh, tiny)
2423 ioncoef = beta + kion * ne
2424 dn = -k2 * nHII * ne + ioncoef * nHI
2425 dtit = min(tleft, max(R(dtfrac) * nHII / max(abs(dn), tiny), R(1e-30)))
2426
2427 # Scalar backward-Euler solve for y=nHII with nHI=nHtot-y:
2428 # dy/dt = -k2*y^2 + beta*(N-y) + kion*y*(N-y)
2429 # Rates/T are frozen within this substep, but electron density is not.
2430 A = dtit * (k2 + kion)
2431 B = one(R) + dtit * beta - dtit * kion * nHtot
2432 C = -(nHII + dtit * beta * nHtot)
2433 disc = max(B * B - R(4) * A * C, zero(R))
2434 sqd = sqrt(disc)
2435 nHII_new = A > tiny ? (-R(2) * C) / max(B + sqd, tiny) :
2436 (nHII + beta * nHtot * dtit) / (one(R) + beta * dtit)
2437 nHII_new = min(max(nHII_new, tiny), nHtot)
2438 hii = nHII_new * mh
2439
2440 # Stiff Compton energy relaxation at frozen ionization over this substep.
2441 c1 = ChemistryKernels.comp1_cmb(R(z))
2442 Tnow = max((R(gamma) - one(R)) * rho * e / (kb * ntot), one(R))
2443 edot_c = -c1 * (Tnow - Trad) * nHII_new
2444 Kc = c1 * nHII_new * (Tnow / max(e, tiny)) / rho
2445 if Kc * dtit > R(1e-4)
2446 B = c1 * nHII_new * Trad / rho
2447 e = (e + B * dtit) / (one(R) + Kc * dtit)
2448 else
2449 e = e + (edot_c / rho) * dtit
2450 end
2451 e = max(e, tiny)
2452 tleft -= dtit
2453 end
2454 return e, hii, h2
2455 end
2456
2457 @kernel function _chem_peebles_k!(rho, eint, HII, H2I, @Const(nsm),
2458 nsm_scalar, nsm_is_scalar::Int, nsm_neutral::Int,
2459 du, vu2, tu, dt, z, hz, fh, gamma,
2460 f_alpha, xe_mean, fudge, gauss, dtfrac, itcap::Int)
2461 c = @index(Global)
2462 @inbounds begin
2463 R = eltype(eint)
2464 nsm_mass = nsm_is_scalar == 1 ? nsm_scalar : nsm[c]
2465 e, h, h2 = _peebles_recomb_cell(rho[c] * du, eint[c] * vu2,
2466 HII[c] * du, H2I[c] * du, dt * tu,
2467 z, hz, fh, gamma, f_alpha, xe_mean,
2468 nsm_mass * du, nsm_neutral, fudge, gauss, dtfrac, itcap)
2469 eint[c] = e / vu2
2470 HII[c] = h / du
2471 H2I[c] = h2 / du
2472 end
2473 end
2474
2475 @kernel function _chem_ionized_compton_k!(rho, eint, HII, H2I,
2476 du, vu2, dt, z, fh, gamma,
2477 hydrogen_reference, hydrogen_complement_reference,
2478 ::Val{NH}, ::Val{CH}) where {NH,CH}
2479 c = @index(Global)
2480 @inbounds begin
2481 R = eltype(eint)
2482 mh = R(ChemistryKernels.MH)
2483 kb = R(ChemistryKernels.KBOLTZ)
2484 tiny = R(1e-30)
2485 r_code = max(rho[c], tiny)
2486 r = r_code * du
2487 fhd = max(R(fh) * r, tiny)
2488 h2 = min(max(H2I[c] * du, tiny), fhd)
2489 Trad = R(2.725) * (one(R) + R(z))
2490 nH = fhd / mh
2491 hi = min(ChemistryKernels.hydrogen_saha_neutral_fraction(Trad, nH) * fhd,
2492 max(fhd - h2, zero(R)))
2493 hii = max(fhd - h2 - hi, tiny)
2494 nHII = hii / mh
2495 nH2I = h2 / mh
2496 nHI = max(hi / mh, tiny)
2497 nHeI = max((one(R) - R(fh)) * r / (R(4) * mh), tiny)
2498 ne = nHII
2499 ntot = max(nHI + nHII + ne + nHeI + R(0.5) * nH2I, tiny)
2500 e = max(eint[c] * vu2, tiny)
2501 Tnow = max((R(gamma) - one(R)) * r * e / (kb * ntot), one(R))
2502 c1 = ChemistryKernels.comp1_cmb(R(z))
2503 Kc = c1 * nHII * (Tnow / max(e, tiny)) / r
2504 if Kc * dt > R(1e-4)
2505 B = c1 * nHII * Trad / r
2506 e = (e + B * dt) / (one(R) + Kc * dt)
2507 else
2508 edot_c = -c1 * (Tnow - Trad) * nHII
2509 e = e + (edot_c / r) * dt
2510 end
2511 eint[c] = max(e, tiny) / vu2
2512 HII[c] = CH ? hydrogen_complement_reference - (h2 + hi) / fhd :
2513 (NH ? hi : hii) / du
2514 H2I[c] = h2 / du
2515 end
2516 end
2517
2518 @kernel function _chem_ionized_compton_helium_k!(rho, eint, HII, H2I, HeII,
2519 du, vu2, dt, z, fh, gamma,
2520 xe_total, heii_mass_fraction,
2521 hydrogen_reference, hydrogen_complement_reference,
2522 ::Val{NH}, ::Val{CH}) where {NH,CH}
2523 c = @index(Global)
2524 @inbounds begin
2525 R = eltype(eint)
2526 mh = R(ChemistryKernels.MH)
2527 kb = R(ChemistryKernels.KBOLTZ)
2528 tiny = R(1e-30)
2529 r_code = max(rho[c], tiny)
2530 r = r_code * du
2531 fhd = max(R(fh) * r, tiny)
2532 h2 = min(max(H2I[c] * du, tiny), fhd)
2533 nH = fhd / mh
2534 Trad = R(2.725) * (one(R) + R(z))
2535 hi = min(ChemistryKernels.hydrogen_saha_neutral_fraction(Trad, nH) * fhd,
2536 max(fhd - h2, zero(R)))
2537 hii = max(fhd - h2 - hi, tiny)
2538 nH2I = h2 / mh
2539 nHII = hii / mh
2540 nHI = max(hi / mh, tiny)
2541 nHe = max((one(R) - R(fh)) * r / (R(4) * mh), tiny)
2542 ne = max(R(xe_total) * nH, tiny)
2543 ntot = max(nHI + nHII + ne + nHe + R(0.5) * nH2I, tiny)
2544 e = max(eint[c] * vu2, tiny)
2545 Tnow = max((R(gamma) - one(R)) * r * e / (kb * ntot), one(R))
2546 c1 = ChemistryKernels.comp1_cmb(R(z))
2547 Kc = c1 * ne * (Tnow / max(e, tiny)) / r
2548 if Kc * dt > R(1e-4)
2549 B = c1 * ne * Trad / r
2550 e = (e + B * dt) / (one(R) + Kc * dt)
2551 else
2552 edot_c = -c1 * (Tnow - Trad) * ne
2553 e = e + (edot_c / r) * dt
2554 end
2555 eint[c] = max(e, tiny) / vu2
2556 HII[c] = CH ? hydrogen_complement_reference - (h2 + hi) / fhd :
2557 (NH ? hi : hii) / du
2558 H2I[c] = h2 / du
2559 HeII[c] = R(heii_mass_fraction) * r_code
2560 end
2561 end
2562
2563 function _parse_list(name, default)
2564 parse.(Float64, split(get(ENV, name, default), ","))
2565 end
2566
2567 function _parse_optional_list(name)
2568 raw = strip(get(ENV, name, ""))
2569 isempty(raw) && return Float64[]
2570 return parse.(Float64, split(raw, ","))
2571 end
2572
2573 function _interp_table(x::Float64, xs::Vector{Float64}, ys::Vector{Float64}, default::Float64)
2574 isempty(xs) && return default
2575 length(xs) == length(ys) || error("table length mismatch: $(length(xs)) z values vs $(length(ys)) values")
2576 p = sortperm(xs)
2577 xss = xs[p]; yss = ys[p]
2578 x <= xss[1] && return yss[1]
2579 x >= xss[end] && return yss[end]
2580 i = searchsortedfirst(xss, x)
2581 t = (x - xss[i-1]) / (xss[i] - xss[i-1])
2582 return yss[i-1] * (1 - t) + yss[i] * t
2583 end
2584
2585 function _density_cos_mode_norm(N::Int, kmode::Int, harmonic::Int)
2586 k = 2π * harmonic * kmode
2587 norm1 = 0.0
2588 @inbounds for i in 1:N
2589 x = (Float64(i) - 0.5) / Float64(N)
2590 w = cos(k * x)
2591 norm1 += w * w
2592 end
2593 return Float64(N) * Float64(N) * norm1
2594 end
2595
2596 function _fourier_mode_norm(N::Int, kmode::Int, quadrature::Int)
2597 k = 2π * kmode
2598 norm1 = 0.0
2599 @inbounds for i in 1:N
2600 x = (Float64(i) - 0.5) / Float64(N)
2601 w = quadrature == 1 ? cos(k * x) : sin(k * x)
2602 norm1 += w * w
2603 end
2604 return Float64(N) * Float64(N) * norm1
2605 end
2606
2607 function _device_complex_mode!(be, scratch, src, N::Int, kmode::Int;
2608 mean=zero(eltype(src)))
2609 n = length(src)
2610 mean_device = eltype(scratch)(mean)
2611 values = ntuple(2) do quadrature
2612 _scalar_fourier_mode_diag_k!(be)(scratch, src, Int(N), Int(kmode),
2613 quadrature, mean_device; ndrange=n)
2614 KA.synchronize(be)
2615 norm = _fourier_mode_norm(N, kmode, quadrature)
2616 norm > eps(Float64) ? Float64(sum(scratch)) / norm : 0.0
2617 end
2618 return complex(values[1], -values[2])
2619 end
2620
2621 function _device_complex_velocity_mode!(be, scratch, momentum, rho,
2622 N::Int, kmode::Int, smallr)
2623 n = length(rho)
2624 values = ntuple(2) do quadrature
2625 _velocity_fourier_mode_diag_k!(be)(scratch, momentum, rho, Int(N),
2626 Int(kmode), quadrature, smallr;
2627 ndrange=n)
2628 KA.synchronize(be)
2629 norm = _fourier_mode_norm(N, kmode, quadrature)
2630 norm > eps(Float64) ? Float64(sum(scratch)) / norm : 0.0
2631 end
2632 return complex(values[1], -values[2])
2633 end
2634
2635 function _device_complex_slip_mode!(be, scratch, photon_velocity, momentum, rho,
2636 N::Int, kmode::Int, smallr)
2637 n = length(rho)
2638 values = ntuple(2) do quadrature
2639 _slip_fourier_mode_diag_k!(be)(scratch, photon_velocity, momentum, rho,
2640 Int(N), Int(kmode), quadrature, smallr;
2641 ndrange=n)
2642 KA.synchronize(be)
2643 norm = _fourier_mode_norm(N, kmode, quadrature)
2644 norm > eps(Float64) ? Float64(sum(scratch)) / norm : 0.0
2645 end
2646 return complex(values[1], -values[2])
2647 end
2648
2649 function _device_complex_species_mode!(be, scratch, species, rho,
2650 N::Int, kmode::Int, fh)
2651 n = length(rho)
2652 mean_fraction = _device_species_mean!(be, scratch, rho, species, fh)
2653 values = ntuple(2) do quadrature
2654 _species_fourier_mode_diag_k!(be)(scratch, species, rho, Int(N),
2655 Int(kmode), quadrature,
2656 eltype(rho)(fh),
2657 eltype(rho)(mean_fraction);
2658 ndrange=n)
2659 KA.synchronize(be)
2660 norm = _fourier_mode_norm(N, kmode, quadrature)
2661 norm > eps(Float64) ? Float64(sum(scratch)) / norm : 0.0
2662 end
2663 return complex(values[1], -values[2])
2664 end
2665
2666 function _device_complex_hii_mode!(be, scratch, fraction_scratch, hydrogen, H2I,
2667 rho, N::Int, kmode::Int, fh,
2668 neutral_hydrogen_storage::Bool,
2669 centered_hydrogen_storage::Bool=false,
2670 hydrogen_reference::Real=1)
2671 centered_hydrogen_storage && return _device_complex_mode!(
2672 be, scratch, hydrogen, N, kmode; mean=zero(eltype(hydrogen)))
2673 neutral_hydrogen_storage || return _device_complex_species_mode!(
2674 be, scratch, hydrogen, rho, N, kmode, fh)
2675 _hydrogen_complement_fraction_diag_k!(be)(
2676 fraction_scratch, rho, hydrogen, H2I, eltype(rho)(fh), Val(true);
2677 ndrange=length(rho))
2678 KA.synchronize(be)
2679 mean_complement = Float64(sum(fraction_scratch)) / Float64(length(rho))
2680 # x_HII = 1 - (x_HI + x_H2), so every nonzero Fourier mode is the
2681 # negative complement mode. This avoids ever forming 1-x in Float32.
2682 -_device_complex_mode!(be, scratch, fraction_scratch, N, kmode;
2683 mean=mean_complement)
2684 end
2685
2686 function _device_complex_temperature_fraction_mode!(be, scratch, temperature_scratch,
2687 s, HII, H2I, HeII, N::Int, kmode::Int, velocity_unit2, z, fh,
2688 hubble, omega_b, neutral_hydrogen_storage::Bool=false,
2689 centered_hydrogen_storage::Bool=false,
2690 hydrogen_reference::Real=1)
2691 n = length(s.U[1])
2692 if HeII === nothing
2693 _gas_temperature_diag_k!(be)(
2694 temperature_scratch, s.U[1], s.U[2], s.U[3], s.U[4], s.U[5],
2695 s.U[6], s.U[7], s.U[8], HII, H2I, eltype(s.U[1])(velocity_unit2),
2696 eltype(s.U[1])(fh), s.γ, s.smallr,
2697 eltype(s.U[1])(hydrogen_reference), Val(neutral_hydrogen_storage),
2698 Val(centered_hydrogen_storage); ndrange=n,
2699 )
2700 else
2701 density_unit = _density_unit_cgs(z, hubble, omega_b)
2702 _gas_temperature_helium_diag_k!(be)(
2703 temperature_scratch, s.U[1], s.U[2], s.U[3], s.U[4], s.U[5],
2704 s.U[6], s.U[7], s.U[8], HII, H2I, HeII,
2705 eltype(s.U[1])(density_unit), eltype(s.U[1])(velocity_unit2),
2706 eltype(s.U[1])(z),
2707 eltype(s.U[1])(fh), s.γ, s.smallr,
2708 eltype(s.U[1])(hydrogen_reference), Val(neutral_hydrogen_storage),
2709 Val(centered_hydrogen_storage); ndrange=n,
2710 )
2711 end
2712 KA.synchronize(be)
2713 mean_temperature = Float64(sum(temperature_scratch)) / Float64(n)
2714 value = _device_complex_mode!(
2715 be, scratch, temperature_scratch, N, kmode; mean=mean_temperature
2716 )
2717 value / max(mean_temperature, eps(Float64))
2718 end
2719
2720 function _device_complex_thermal_fraction_mode!(be, scratch, thermal_scratch,
2721 s, N::Int, kmode::Int, mode::Int)
2722 n = length(s.U[1])
2723 _thermal_diag_k!(be)(
2724 thermal_scratch, s.U[1], s.U[2], s.U[3], s.U[4], s.U[5],
2725 s.U[6], s.U[7], s.U[8], eltype(s.U[1])(s.γ),
2726 eltype(s.U[1])(s.smallr), eltype(s.U[1])(1e-30), mode;
2727 ndrange=n,
2728 )
2729 KA.synchronize(be)
2730 mean_value = Float64(sum(thermal_scratch)) / Float64(n)
2731 value = _device_complex_mode!(
2732 be, scratch, thermal_scratch, N, kmode; mean=mean_value
2733 )
2734 value / max(abs(mean_value), eps(Float64))
2735 end
2736
2737 function _device_brms!(be, s, scratch)
2738 n = length(s.U[1])
2739 _b2_diag_k!(be)(scratch, s.U[6], s.U[7], s.U[8]; ndrange=n)
2740 KA.synchronize(be)
2741 return sqrt(Float64(sum(scratch)) / Float64(n))
2742 end
2743
2744 function _device_divb_norm!(be, s, scratch, br::Float64)
2745 n = length(s.U[1])
2746 N = Int(s.dims[1])
2747 _divb_abs_periodic_diag_k!(be)(scratch, s.U[6], s.U[7], s.U[8], N, T(0.5) / T(s.dx);
2748 ndrange=n)
2749 KA.synchronize(be)
2750 return Float64(maximum(scratch)) * Float64(s.dx) / max(br, eps(Float64))
2751 end
2752
2753 "Face-divergence diagnostic max|div b_face| dx / Brms for CT states; NaN for GLM states."
2754 function _device_divb_face_norm(s, br::Float64)
2755 MHDKernels.uses_ct(s) || return NaN
2756 return MHDKernels.max_divb_face(s) * Float64(s.dx) / max(br, eps(Float64))
2757 end
2758
2759 function _device_component_rms!(be, scratch, src)
2760 n = length(src)
2761 _component2_diag_k!(be)(scratch, src; ndrange=n)
2762 KA.synchronize(be)
2763 return sqrt(Float64(sum(scratch)) / Float64(n))
2764 end
2765
2766 function _device_delta_rms!(be, scratch, src; mean=one(T))
2767 n = length(src)
2768 _delta2_diag_k!(be)(scratch, src, mean; ndrange=n)
2769 KA.synchronize(be)
2770 return sqrt(Float64(sum(scratch)) / Float64(n))
2771 end
2772
2773 function _device_delta_rms!(be, scratch, src, mean)
2774 n = length(src)
2775 _delta2_diag_k!(be)(scratch, src, mean; ndrange=n)
2776 KA.synchronize(be)
2777 return sqrt(Float64(sum(scratch)) / Float64(n))
2778 end
2779
2780 function _device_density_cos_mode_amp!(be,scratch,rho,N::Int,kmode::Int,
2781 harmonic::Int;mean=one(eltype(rho)))
2782 n = length(rho)
2783 _density_cos_mode_diag_k!(be)(scratch,rho,Int(N),Int(kmode),Int(harmonic),mean;
2784 ndrange=n)
2785 KA.synchronize(be)
2786 norm = _density_cos_mode_norm(N, kmode, harmonic)
2787 return norm > 0 ? Float64(sum(scratch)) / norm : NaN
2788 end
2789
2790 function _device_species_stats!(be, scratch, rho, HII, H2I, HeII, fh;
2791 z=NaN, h=0.674, Ob=0.049,
2792 neutral_hydrogen_storage::Bool=false,
2793 centered_hydrogen_storage::Bool=false,
2794 hydrogen_reference::Real=1)
2795 if HII === nothing
2796 return ntuple(_ -> NaN, 12)
2797 end
2798 n = length(rho)
2799 if centered_hydrogen_storage
2800 _centered_hii_fraction_diag_k!(be)(
2801 scratch, rho, HII, H2I, T(fh), T(hydrogen_reference); ndrange=n)
2802 KA.synchronize(be)
2803 xhmean = Float64(sum(scratch)) / Float64(n)
2804 xhmin = Float64(minimum(scratch))
2805 xhmax = Float64(maximum(scratch))
2806 elseif neutral_hydrogen_storage
2807 _hydrogen_complement_fraction_diag_k!(be)(
2808 scratch, rho, HII, H2I, T(fh), Val(true); ndrange=n)
2809 KA.synchronize(be)
2810 complement_mean = Float64(sum(scratch)) / Float64(n)
2811 complement_min = Float64(minimum(scratch))
2812 complement_max = Float64(maximum(scratch))
2813 xhmean = 1.0 - complement_mean
2814 xhmin = 1.0 - complement_max
2815 xhmax = 1.0 - complement_min
2816 else
2817 _species_fraction_diag_k!(be)(scratch, rho, HII, T(fh); ndrange=n)
2818 KA.synchronize(be)
2819 xhmean = Float64(sum(scratch)) / Float64(n)
2820 xhmin = Float64(minimum(scratch))
2821 xhmax = Float64(maximum(scratch))
2822 end
2823
2824 _species_fraction_diag_k!(be)(scratch, rho, H2I, T(fh); ndrange=n)
2825 KA.synchronize(be)
2826 x2mean = Float64(sum(scratch)) / Float64(n)
2827 x2min = Float64(minimum(scratch))
2828 x2max = Float64(maximum(scratch))
2829 if HeII === nothing || !isfinite(z)
2830 return (xhmean, xhmin, xhmax, x2mean, x2min, x2max,
2831 NaN, NaN, NaN, NaN, NaN, NaN)
2832 end
2833 _species_fraction_diag_k!(be)(scratch, rho, HeII, T(4) * T(fh);
2834 ndrange=n)
2835 KA.synchronize(be)
2836 xhemean = Float64(sum(scratch)) / Float64(n)
2837 xhemin = Float64(minimum(scratch))
2838 xhemax = Float64(maximum(scratch))
2839
2840 density_unit = _density_unit_cgs(z, h, Ob)
2841 _total_electron_fraction_diag_k!(be)(
2842 scratch, rho, HII, H2I, HeII, T(density_unit), T(z), T(fh),
2843 T(hydrogen_reference), Val(neutral_hydrogen_storage),
2844 Val(centered_hydrogen_storage); ndrange=n)
2845 KA.synchronize(be)
2846 xtmean = Float64(sum(scratch)) / Float64(n)
2847 xtmin = Float64(minimum(scratch))
2848 xtmax = Float64(maximum(scratch))
2849 return (xhmean, xhmin, xhmax, x2mean, x2min, x2max,
2850 xhemean, xhemin, xhemax, xtmean, xtmin, xtmax)
2851 end
2852
2853 function _device_hii_mean!(be, scratch, rho, hydrogen, H2I, fh,
2854 neutral_hydrogen_storage::Bool,
2855 centered_hydrogen_storage::Bool=false,
2856 hydrogen_reference::Real=1)
2857 centered_hydrogen_storage && return hydrogen_reference +
2858 Float64(sum(hydrogen)) / Float64(length(hydrogen))
2859 neutral_hydrogen_storage ||
2860 return _device_species_mean!(be, scratch, rho, hydrogen, fh)
2861 _hydrogen_complement_fraction_diag_k!(be)(
2862 scratch, rho, hydrogen, H2I, T(fh), Val(true); ndrange=length(rho))
2863 KA.synchronize(be)
2864 1.0 - Float64(sum(scratch)) / Float64(length(rho))
2865 end
2866
2867 function _device_species_mean!(be, scratch, rho, species, fh)
2868 species === nothing && return NaN
2869 _species_fraction_diag_k!(be)(scratch, rho, species, T(fh); ndrange=length(rho))
2870 KA.synchronize(be)
2871 return Float64(sum(scratch)) / Float64(length(rho))
2872 end
2873
2874 function _device_thermal_stats!(be, s, scratch)
2875 n = length(s.U[1])
2876 _thermal_diag_k!(be)(scratch, s.U[1], s.U[2], s.U[3], s.U[4], s.U[5],
2877 s.U[6], s.U[7], s.U[8], T(s.γ), T(s.smallr), T(1e-30), 1;
2878 ndrange=n)
2879 KA.synchronize(be)
2880 pmean = Float64(sum(scratch)) / Float64(n)
2881 pmin = Float64(minimum(scratch))
2882 pmax = Float64(maximum(scratch))
2883
2884 _thermal_diag_k!(be)(scratch, s.U[1], s.U[2], s.U[3], s.U[4], s.U[5],
2885 s.U[6], s.U[7], s.U[8], T(s.γ), T(s.smallr), T(1e-30), 2;
2886 ndrange=n)
2887 KA.synchronize(be)
2888 emean = Float64(sum(scratch)) / Float64(n)
2889 emin = Float64(minimum(scratch))
2890 emax = Float64(maximum(scratch))
2891 return (pmean, pmin, pmax, emean, emin, emax)
2892 end
2893
2894 function _device_temperature_stats!(be, s, scratch, HII, H2I, HeII,
2895 density_unit, velocity_unit2, z, fh,
2896 neutral_hydrogen_storage::Bool=false,
2897 centered_hydrogen_storage::Bool=false,
2898 hydrogen_reference::Real=1)
2899 HII === nothing && return (NaN, NaN, NaN)
2900 n = length(s.U[1])
2901 if HeII === nothing
2902 _gas_temperature_diag_k!(be)(
2903 scratch, s.U[1], s.U[2], s.U[3], s.U[4], s.U[5],
2904 s.U[6], s.U[7], s.U[8], HII, H2I, T(velocity_unit2), T(fh),
2905 s.γ, s.smallr, T(hydrogen_reference), Val(neutral_hydrogen_storage),
2906 Val(centered_hydrogen_storage); ndrange=n)
2907 else
2908 _gas_temperature_helium_diag_k!(be)(
2909 scratch, s.U[1], s.U[2], s.U[3], s.U[4], s.U[5],
2910 s.U[6], s.U[7], s.U[8], HII, H2I, HeII, T(density_unit),
2911 T(velocity_unit2), T(z), T(fh), s.γ, s.smallr,
2912 T(hydrogen_reference), Val(neutral_hydrogen_storage),
2913 Val(centered_hydrogen_storage); ndrange=n)
2914 end
2915 KA.synchronize(be)
2916 (Float64(sum(scratch)) / Float64(n), Float64(minimum(scratch)),
2917 Float64(maximum(scratch)))
2918 end
2919
2920 function _device_pressure_mean!(be, s, scratch)
2921 n = length(s.U[1])
2922 _thermal_diag_k!(be)(scratch, s.U[1], s.U[2], s.U[3], s.U[4], s.U[5],
2923 s.U[6], s.U[7], s.U[8], T(s.γ), T(s.smallr), T(1e-30), 1;
2924 ndrange=n)
2925 KA.synchronize(be)
2926 return Float64(sum(scratch)) / Float64(n)
2927 end
2928
2929 function _final_device_stats!(be,s,ρdm,scratch,HII,H2I,HeII,fh;
2930 density_perturbation=nothing, z=NaN,
2931 h=0.674, Ob=0.049,
2932 dm_mean=one(eltype(ρdm)),
2933 neutral_hydrogen_storage::Bool=false,
2934 centered_hydrogen_storage::Bool=false,
2935 hydrogen_reference::Real=1)
2936 br = _device_brms!(be, s, scratch)
2937 divn = _device_divb_norm!(be, s, scratch, br)
2938 δb=density_perturbation === nothing ?
2939 _device_delta_rms!(be,scratch,s.U[1],one(T)) :
2940 _device_delta_rms!(be,scratch,density_perturbation,zero(T))
2941 δdm = _device_delta_rms!(be, scratch, ρdm, dm_mean)
2942 chemstats = _device_species_stats!(be, scratch, s.U[1], HII, H2I, HeII, fh;
2943 z=z, h=h, Ob=Ob,
2944 neutral_hydrogen_storage=neutral_hydrogen_storage,
2945 centered_hydrogen_storage=centered_hydrogen_storage,
2946 hydrogen_reference=hydrogen_reference)
2947 return br, divn, δb, δdm, chemstats
2948 end
2949
2950 function _adaptive_diag_watch!(be,s,scratch,HII,H2I,HeII,fh;
2951 density_perturbation=nothing, z=NaN,
2952 h=0.674, Ob=0.049,
2953 neutral_hydrogen_storage::Bool=false,
2954 centered_hydrogen_storage::Bool=false,
2955 hydrogen_reference::Real=1)
2956 δb=density_perturbation === nothing ?
2957 _device_delta_rms!(be,scratch,s.U[1],one(T)) :
2958 _device_delta_rms!(be,scratch,density_perturbation,zero(T))
2959 chemstats = _device_species_stats!(be, scratch, s.U[1], HII, H2I, HeII, fh;
2960 z=z, h=h, Ob=Ob,
2961 neutral_hydrogen_storage=neutral_hydrogen_storage,
2962 centered_hydrogen_storage=centered_hydrogen_storage,
2963 hydrogen_reference=hydrogen_reference)
2964 return δb, chemstats[1]
2965 end
2966
2967 function _rel_changed(now::Float64, old::Float64, frac::Float64)
2968 isfinite(now) || return false
2969 isfinite(old) || return true
2970 scale = max(abs(old), eps(Float64))
2971 return abs(now - old) / scale >= frac
2972 end
2973
2974 function _state_summary(s)
2975 _b2_diag_k!(s.be)(s.scratch[1], s.U[6], s.U[7], s.U[8]; ndrange=length(s.U[1]))
2976 KA.synchronize(s.be)
2977 ρmin = minimum(s.U[1])
2978 ρmax = maximum(s.U[1])
2979 Emin = minimum(s.U[5])
2980 Emax = maximum(s.U[5])
2981 B2max = maximum(s.scratch[1])
2982 finite = isfinite(Float64(ρmin)) && isfinite(Float64(ρmax)) &&
2983 isfinite(Float64(Emin)) && isfinite(Float64(Emax)) &&
2984 isfinite(Float64(B2max))
2985 return (finite = finite,
2986 ρmin = ρmin, ρmax = ρmax,
2987 Emin = Emin, Emax = Emax,
2988 Bmax = isfinite(Float64(B2max)) && Float64(B2max) >= 0 ?
2989 sqrt(Float64(B2max)) : Float64(B2max))
2990 end
2991
2992 function _assert_finite_step(s, label, cyc, a, dt, smax)
2993 if !(isfinite(dt) && dt > 0 && isfinite(Float64(smax)) && Float64(smax) > 0)
2994 ss = _state_summary(s)
2995 error(@sprintf("%s cycle=%d invalid timestep dt=%.6e smax=%.6e a=%.8e z=%.6f rho=[%.6e,%.6e] E=[%.6e,%.6e] Bmax=%.6e finite=%s",
2996 label, cyc, dt, Float64(smax), a, isfinite(a) ? a_to_z(a) : NaN,
2997 Float64(ss.ρmin), Float64(ss.ρmax), Float64(ss.Emin), Float64(ss.Emax),
2998 ss.Bmax, string(ss.finite)))
2999 end
3000 end
3001
3002 function _assert_finite_state(s, label, cyc, a, dt, smax)
3003 ss = _state_summary(s)
3004 if !(ss.finite && isfinite(Float64(ss.ρmin)) && Float64(ss.ρmin) > 0 &&
3005 isfinite(Float64(ss.Emin)) && Float64(ss.Emin) > 0)
3006 error(@sprintf("%s cycle=%d invalid state dt=%.6e smax=%.6e a=%.8e z=%.6f rho=[%.6e,%.6e] E=[%.6e,%.6e] Bmax=%.6e finite=%s",
3007 label, cyc, dt, Float64(smax), a, isfinite(a) ? a_to_z(a) : NaN,
3008 Float64(ss.ρmin), Float64(ss.ρmax), Float64(ss.Emin), Float64(ss.Emax),
3009 ss.Bmax, string(ss.finite)))
3010 end
3011 return ss
3012 end
3013
3014 function _time_phase(f, be1, be2)
3015 t = @elapsed begin
3016 f()
3017 KA.synchronize(be1)
3018 be2 === be1 || KA.synchronize(be2)
3019 end
3020 return t
3021 end
3022
3023 function _log_line(msg::AbstractString)
3024 @printf("[%s] %s\n", Dates.format(now(), dateformat"yyyy-mm-ddTHH:MM:SS"), msg)
3025 flush(stdout)
3026 end
3027
3028 function _device_memory_gib()
3029 if BACKEND_NAME === :metal
3030 Metal.synchronize()
3031 stats = Metal.alloc_stats
3032 live = max(0, stats.alloc_bytes - stats.free_bytes) / 2.0^30
3033 return (live=live, allocated=stats.alloc_bytes / 2.0^30,
3034 freed=stats.free_bytes / 2.0^30)
3035 elseif BACKEND_NAME === :cuda
3036 CUDA.synchronize()
3037 live = (CUDA.total_memory() - CUDA.available_memory()) / 2.0^30
3038 return (live=live, allocated=NaN, freed=NaN)
3039 end
3040 return (live=NaN, allocated=NaN, freed=NaN)
3041 end
3042
3043 function _log_device_memory(label::AbstractString)
3044 get(ENV, "MHD_MEMPROBE", "0") == "1" || return nothing
3045 stats = _device_memory_gib()
3046 @printf("[%s] MEMORY %-28s live_device=%.3f GiB allocated=%.3f GiB freed=%.3f GiB\n",
3047 Dates.format(now(), dateformat"yyyy-mm-ddTHH:MM:SS"), label,
3048 stats.live, stats.allocated, stats.freed)
3049 flush(stdout)
3050 return nothing
3051 end
3052
3053 function _stage(label::AbstractString, f)
3054 _log_line("START " * label)
3055 result = nothing
3056 t = @elapsed begin
3057 result = f()
3058 end
3059 _log_line(@sprintf("DONE %s %.3f s", label, t))
3060 _log_device_memory(label)
3061 return result
3062 end
3063
3064 _stage(f, label::AbstractString) = _stage(label, f)
3065
3066 function _fft_mode(name::Symbol, env::AbstractString)
3067 s = Symbol(lowercase(env))
3068 if s === :auto
3069 # :mps denotes the native real-FFT interface in this driver: MPSGraph
3070 # for Metal, cuFFT for CUDA, and FFTW for CPU arrays.
3071 return :mps
3072 elseif s in (:mps, :native, :rfft)
3073 return :mps
3074 elseif s in (:ka, :ka_rfft)
3075 return :ka
3076 elseif s in (:c2c, :ka_c2c)
3077 return :c2c
3078 else
3079 error("MHD_FFT must be auto, mps/native/rfft, ka, or c2c; got '$env'")
3080 end
3081 end
3082
3083 function _chem_mode(env::AbstractString)
3084 s = Symbol(lowercase(env))
3085 s in (:off, :none, :false, :0) && return :off
3086 s in (:peebles, :recfast, :full, :on, :true, :1) && return :peebles
3087 s in (:dense, :stiff, :network) && return :dense
3088 s in (:mixing, :lya, :lyalpha) && return :mixing
3089 s in (:analytic_mixing, :mixing_analytic, :fast_mixing, :analytic, :analytic_h2, :fast) &&
3090 return :analytic_mixing
3091 error("MHD_CHEM must be off, peebles/recfast, dense, mixing, or analytic_mixing; got '$env'")
3092 end
3093
3094 function _chem_smooth_mode(env::AbstractString)
3095 s = Symbol(lowercase(env))
3096 s in (:local, :none, :cell) && return :local
3097 s in (:global, :mean, :box) && return :global
3098 s in (:gaussian, :fft, :fft_gaussian, :rfft, :rfft_gaussian) && return :gaussian
3099 error("MHD_CHEM_SMOOTH must be local, global, or gaussian; got '$env'")
3100 end
3101
3102 function _chem_coupling_mode(env::AbstractString)
3103 s = Symbol(lowercase(env))
3104 s in (:lie, :split, :post, :operator_split) && return :lie
3105 s in (:strang, :midpoint, :half, :halfstep, :symmetric) && return :strang
3106 error("MHD_CHEM_COUPLING must be lie/split or strang/midpoint; got '$env'")
3107 end
3108
3109 function _chem_expansion_enabled(env::AbstractString, zrun::Bool, chem_mode::Symbol)
3110 s = lowercase(strip(env))
3111 s in ("auto", "default") && return zrun && chem_mode !== :off
3112 s in ("on", "true", "yes", "1") && return true
3113 s in ("off", "false", "no", "0") && return false
3114 error("MHD_CHEM_EXPANSION must be auto, on/1, or off/0; got '$env'")
3115 end
3116
3117 function _run_pmf_chemistry_step!(be_p, be_mhd, s, eint, HII, H2I, HeII, ρtot, φ;
3118 N::Int, ncells::Int, zrun::Bool,
3119 a0::Float64, a1::Float64, dt_code::Float64,
3120 hub::Float64, Om::Float64, OL::Float64, Or::Float64,
3121 Ob::Float64, boxsize::Float64,
3122 chem_mode::Symbol, chem_smooth::Symbol,
3123 chem_density_fixed::AbstractString,
3124 chem_vunit::Float64, chem_time_unit::Float64,
3125 chem_ionized_fast_zmin::Float64,
3126 chem_fused_analytic::Bool,
3127 chem_fused_mixing::Bool,
3128 chem_helium::Bool,
3129 chem_neutral_hydrogen_storage::Bool,
3130 chem_centered_hydrogen_storage::Bool,
3131 chem_expansion::Bool,
3132 chem_profile_enabled::Bool, measuring::Bool,
3133 refresh_mean::Bool,
3134 chem_phase, chem_falpha::Float64,
3135 chem_falpha_zs, chem_falpha_vals,
3136 chem_nsmean::Float64, chem_smooth_sigma::Float64,
3137 chem_xe_mean::Float64, evolving_xe_mean::Float64,
3138 evolving_xheii_mean::Float64,
3139 chem_recfast_hswitch::Bool, chem_fudge::Float64,
3140 chem_rate_tables, chem_dtfrac::Float64,
3141 chem_itcap::Int, fh,
3142 chem_nonlocal_table=nothing,
3143 chem_nonlocal_scale::Float64=1.0,
3144 chem_nonlocal_before::Bool=false,
3145 radiation_workspace=nothing,
3146 terminal_lbox_cm::Float64=NaN,
3147 cosmological_mhd_aref::Float64=NaN)
3148 chem_mode === :off && return 0.0, evolving_xe_mean, evolving_xheii_mean
3149 T = eltype(s.U[1])
3150 energy_residual = radiation_workspace === nothing ? s.U[5] :
3151 radiation_workspace.energy_residual
3152 track_energy = Val(radiation_workspace !== nothing)
3153 xe_mean = evolving_xe_mean
3154 xheii_mean = evolving_xheii_mean
3155 hydrogen_reference = evolving_xe_mean
3156 hydrogen_complement_reference = 1.0 - hydrogen_reference
3157 internal_neutral_hydrogen = chem_neutral_hydrogen_storage ||
3158 (chem_centered_hydrogen_storage && hydrogen_reference >= 0.5)
3159 tc = _time_phase(be_p, be_mhd) do
3160 profile_chem = chem_profile_enabled && measuring
3161 zchem = zrun ? a_to_z(0.5 * (a0 + a1)) : 0.0
3162 dtphys = zrun ? _cosmic_dt_seconds(a0, a1, hub, Om, OL, Or) :
3163 dt_code * chem_time_unit
3164 amid = 0.5 * (a0 + a1)
3165 velocity_unit2 = chem_vunit^2 *
3166 (isfinite(cosmological_mhd_aref) ? cosmological_mhd_aref / amid : 1.0)
3167 dunit = isempty(chem_density_fixed) ? _density_unit_cgs(zchem, hub, Ob) :
3168 parse(Float64, chem_density_fixed)
3169 chem_hubble_expansion = chem_expansion && zrun
3170 chem_adot_over_a = chem_hubble_expansion ? _hubble_s(zchem, hub, Om, OL, Or) : NaN
3171 function run_nonlocal!()
3172 chem_nonlocal_table === nothing && return nothing
3173 chem_nonlocal_scale > 0 || return nothing
3174 radiation_workspace === nothing && error(
3175 "nonlocal chemistry requires the reusable radiation FFT workspace")
3176 isfinite(terminal_lbox_cm) && terminal_lbox_cm > 0 || error(
3177 "nonlocal chemistry requires a positive comoving MHD_TERMINAL_LBOX_CM")
3178 hz = _hubble_s(zchem, hub, Om, OL, Or)
3179 theta_conversion = sqrt(velocity_unit2) * boxsize /
3180 (terminal_lbox_cm * amid * hz)
3181 nonlocal_block = () -> MHDKernels.apply_nonlocal_recombination!(
3182 HII, H2I, s, radiation_workspace, chem_nonlocal_table;
3183 redshift=zchem, xe_mean=xe_mean,
3184 dt_seconds=dtphys * chem_nonlocal_scale,
3185 fundamental_k_mpc=2pi * _PMF_MPC_CM / terminal_lbox_cm,
3186 theta_over_aH_conversion=theta_conversion,
3187 hydrogen_mass_fraction=fh,
3188 neutral_hydrogen_storage=internal_neutral_hydrogen,
3189 centered_hydrogen_storage=chem_centered_hydrogen_storage,
3190 hydrogen_reference=hydrogen_reference,
3191 hydrogen_complement_reference=hydrogen_complement_reference,
3192 velocity_unit2=velocity_unit2,
3193 temperature_reference=_PMF_T_CMB * (1 + zchem))
3194 if profile_chem
3195 chem_phase[:nonlocal] += _time_phase(nonlocal_block, be_p, be_mhd)
3196 else
3197 nonlocal_block()
3198 end
3199 nothing
3200 end
3201 chem_nonlocal_before && run_nonlocal!()
3202 ionized_fast = (chem_mode === :peebles || chem_mode === :mixing ||
3203 chem_mode === :analytic_mixing) &&
3204 chem_ionized_fast_zmin > 0 && zchem >= chem_ionized_fast_zmin
3205 use_fused_analytic = chem_fused_analytic && chem_mode === :analytic_mixing &&
3206 !ionized_fast
3207 use_fused_mixing = chem_fused_mixing && chem_mode === :mixing && !ionized_fast
3208 chem_helium && chem_mode !== :mixing && error(
3209 "MHD_CHEM_HELIUM=1 currently requires MHD_CHEM=mixing")
3210 chem_helium && !ionized_fast && !use_fused_mixing && error(
3211 "MHD_CHEM_HELIUM=1 requires the fused full mixing kernel")
3212 use_fused_chem = use_fused_analytic || use_fused_mixing
3213 split_helium_energy = BACKEND_NAME === :metal && chem_helium &&
3214 use_fused_mixing
3215 chemistry_updates_total_energy = use_fused_chem && !split_helium_energy
3216 if profile_chem && !chemistry_updates_total_energy
3217 chem_phase[:unpack] += _time_phase(be_p, be_mhd) do
3218 _mhd_to_eint_k!(be_p)(eint, s.scratch[9], s.U[1], s.U[2],
3219 s.U[3], s.U[4], s.U[5], s.U[6], s.U[7],
3220 s.U[8], T(s.smallr); ndrange=ncells)
3221 end
3222 elseif !chemistry_updates_total_energy
3223 _mhd_to_eint_k!(be_p)(eint, s.scratch[9], s.U[1], s.U[2],
3224 s.U[3], s.U[4], s.U[5], s.U[6], s.U[7],
3225 s.U[8], T(s.smallr); ndrange=ncells)
3226 KA.synchronize(be_p)
3227 end
3228 if ionized_fast
3229 helium_state = chem_helium ?
3230 _equilibrium_helium_state(1.0, zchem, hub, Ob, Float64(fh)) : nothing
3231 if profile_chem
3232 chem_phase[:solve] += _time_phase(be_p, be_mhd) do
3233 if chem_helium
3234 _chem_ionized_compton_helium_k!(be_p)(
3235 s.U[1], eint, HII, H2I, HeII,
3236 T(dunit), T(velocity_unit2), T(dtphys), T(zchem), fh,
3237 T(s.γ), T(helium_state.xe_total),
3238 T(helium_state.heii_mass_fraction),
3239 T(hydrogen_reference), T(hydrogen_complement_reference),
3240 Val(internal_neutral_hydrogen),
3241 Val(chem_centered_hydrogen_storage); ndrange=ncells)
3242 else
3243 _chem_ionized_compton_k!(be_p)(s.U[1], eint, HII, H2I,
3244 T(dunit), T(velocity_unit2),
3245 T(dtphys), T(zchem), fh, T(s.γ),
3246 T(hydrogen_reference),
3247 T(hydrogen_complement_reference),
3248 Val(internal_neutral_hydrogen),
3249 Val(chem_centered_hydrogen_storage);
3250 ndrange=ncells)
3251 end
3252 end
3253 else
3254 if chem_helium
3255 _chem_ionized_compton_helium_k!(be_p)(
3256 s.U[1], eint, HII, H2I, HeII,
3257 T(dunit), T(velocity_unit2), T(dtphys), T(zchem), fh,
3258 T(s.γ), T(helium_state.xe_total),
3259 T(helium_state.heii_mass_fraction),
3260 T(hydrogen_reference), T(hydrogen_complement_reference),
3261 Val(internal_neutral_hydrogen),
3262 Val(chem_centered_hydrogen_storage); ndrange=ncells)
3263 else
3264 _chem_ionized_compton_k!(be_p)(s.U[1], eint, HII, H2I,
3265 T(dunit), T(velocity_unit2),
3266 T(dtphys), T(zchem), fh, T(s.γ),
3267 T(hydrogen_reference),
3268 T(hydrogen_complement_reference),
3269 Val(internal_neutral_hydrogen),
3270 Val(chem_centered_hydrogen_storage);
3271 ndrange=ncells)
3272 end
3273 end
3274 elseif chem_mode === :peebles || chem_mode === :mixing ||
3275 chem_mode === :analytic_mixing
3276 f_alpha_step = clamp(_interp_table(Float64(zchem), chem_falpha_zs,
3277 chem_falpha_vals, chem_falpha),
3278 0.0, 1.0)
3279 nsm_arg = s.U[1]
3280 nsm_scalar = T(0)
3281 nsm_is_scalar = 0
3282 nsm_neutral = 0
3283 if f_alpha_step != 0
3284 if chem_smooth === :global
3285 nsm_scalar = T(chem_nsmean)
3286 nsm_is_scalar = 1
3287 elseif chem_smooth === :gaussian
3288 if profile_chem
3289 chem_phase[:smooth] += _time_phase(be_p, be_mhd) do
3290 _neutral_h_mass_k!(be_p)(ρtot, s.U[1], HII, H2I, fh,
3291 T(hydrogen_complement_reference), Int(N),
3292 Val(internal_neutral_hydrogen),
3293 Val(chem_centered_hydrogen_storage);
3294 ndrange=ncells)
3295 PoissonKernels.rfft_smooth_gaussian!(φ, ρtot;
3296 sigma_cells=chem_smooth_sigma, boxsize=boxsize)
3297 end
3298 else
3299 _neutral_h_mass_k!(be_p)(ρtot, s.U[1], HII, H2I, fh,
3300 T(hydrogen_complement_reference), Int(N),
3301 Val(internal_neutral_hydrogen),
3302 Val(chem_centered_hydrogen_storage);
3303 ndrange=ncells)
3304 PoissonKernels.rfft_smooth_gaussian!(φ, ρtot;
3305 sigma_cells=chem_smooth_sigma, boxsize=boxsize)
3306 end
3307 nsm_arg = φ
3308 nsm_neutral = 1
3309 else
3310 nsm_arg = s.U[1]
3311 end
3312 end
3313 gauss = chem_recfast_hswitch ? ChemistryKernels.recfast_gauss_factor(zchem) : 1.0
3314 xmean = f_alpha_step == 0 ? 0.0 :
3315 (isfinite(chem_xe_mean) ? chem_xe_mean : xe_mean)
3316 solve_block = () -> begin
3317 if chem_mode === :mixing
3318 if use_fused_mixing
3319 if chem_helium
3320 if split_helium_energy
3321 _chem_mixing_helium_eint_k!(be_p)(
3322 eint, HII, H2I, HeII, s.U[1], nsm_arg,
3323 nsm_scalar, T(dunit), T(velocity_unit2),
3324 T(dtphys), T(zchem), T(f_alpha_step), T(xmean),
3325 T(chem_fudge), T(gauss), T(100hub), T(Om), T(OL), fh,
3326 chem_hubble_expansion, chem_rate_tables,
3327 T(chem_dtfrac), chem_itcap, T(s.smallr),
3328 Complex{T}(T(hydrogen_reference),
3329 T(hydrogen_complement_reference)),
3330 Val(internal_neutral_hydrogen),
3331 Val(chem_centered_hydrogen_storage),
3332 Val(nsm_is_scalar == 1),
3333 Val(nsm_neutral == 1); ndrange=ncells)
3334 else
3335 _chem_mixing_helium_mhd_k!(be_p)(
3336 s.U[1], s.U[2], s.U[3], s.U[4], s.U[5],
3337 s.U[6], s.U[7], s.U[8], HII, H2I, HeII, nsm_arg,
3338 energy_residual,
3339 nsm_scalar, T(dunit), T(velocity_unit2),
3340 T(dtphys), T(zchem), T(f_alpha_step), T(xmean),
3341 T(chem_fudge), T(gauss), T(100hub), T(Om), T(OL), fh,
3342 chem_hubble_expansion, chem_rate_tables,
3343 T(chem_dtfrac), chem_itcap, T(s.smallr),
3344 Complex{T}(T(hydrogen_reference),
3345 T(hydrogen_complement_reference)),
3346 Val(internal_neutral_hydrogen),
3347 Val(chem_centered_hydrogen_storage),
3348 track_energy,
3349 Val(nsm_is_scalar == 1),
3350 Val(nsm_neutral == 1); ndrange=ncells)
3351 end
3352 else
3353 _chem_mixing_mhd_k!(be_p)(
3354 s.U[1], s.U[2], s.U[3], s.U[4], s.U[5],
3355 s.U[6], s.U[7], s.U[8], HII, H2I, nsm_arg,
3356 energy_residual,
3357 nsm_scalar, T(dunit), T(velocity_unit2),
3358 T(dtphys), T(zchem), T(f_alpha_step), T(xmean),
3359 T(chem_fudge), T(gauss), T(100hub), T(Om), T(OL), fh,
3360 chem_hubble_expansion, chem_rate_tables,
3361 T(chem_dtfrac), chem_itcap, T(s.smallr),
3362 Complex{T}(T(hydrogen_reference),
3363 T(hydrogen_complement_reference)),
3364 Val(internal_neutral_hydrogen),
3365 Val(chem_centered_hydrogen_storage),
3366 track_energy,
3367 Val(nsm_is_scalar == 1),
3368 Val(nsm_neutral == 1); ndrange=ncells)
3369 end
3370 else
3371 nsm_source = nsm_is_scalar == 1 ? nsm_scalar : nsm_arg
3372 ChemistryKernels.solve_chem_mixing_device!(
3373 s.U[1], eint, HII, H2I, nsm_source;
3374 a_value=1/(1+zchem), dt=dtphys,
3375 density_units=dunit, length_units=sqrt(velocity_unit2), time_units=1,
3376 f_alpha=f_alpha_step, Xe_mean=xmean,
3377 smoothed_is_neutral=nsm_neutral == 1,
3378 recfast_fudge=chem_fudge,
3379 recfast_hswitch=chem_recfast_hswitch,
3380 hubble_expansion=chem_hubble_expansion,
3381 hubble=100hub, Om=Om, OL=OL, fh=Float64(fh),
3382 rate_tables=chem_rate_tables, dtfrac=chem_dtfrac,
3383 itcap=chem_itcap, backend=BACKEND_NAME, precision=T)
3384 end
3385 elseif chem_mode === :analytic_mixing
3386 if use_fused_analytic
3387 _chem_analytic_mixing_mhd_k!(be_p)(
3388 s.U[1], s.U[2], s.U[3], s.U[4], s.U[5],
3389 s.U[6], s.U[7], s.U[8], HII, H2I, nsm_arg,
3390 energy_residual,
3391 nsm_scalar, T(dunit), T(velocity_unit2), T(1.0),
3392 T(dtphys), T(zchem), T(f_alpha_step), T(xmean),
3393 T(chem_fudge), T(gauss), T(100hub), T(Om), T(OL),
3394 fh, chem_hubble_expansion, T(chem_adot_over_a),
3395 chem_rate_tables, T(chem_dtfrac), chem_itcap, T(s.smallr),
3396 track_energy,
3397 Val(nsm_is_scalar == 1),
3398 Val(nsm_neutral == 1); ndrange=ncells)
3399 else
3400 nsm_source = nsm_is_scalar == 1 ? nsm_scalar : nsm_arg
3401 ChemistryKernels.solve_chem_analytic_mixing_device!(
3402 s.U[1], eint, HII, H2I, nsm_source;
3403 a_value=1/(1+zchem), dt=dtphys,
3404 density_units=dunit, length_units=sqrt(velocity_unit2), time_units=1,
3405 f_alpha=f_alpha_step, Xe_mean=xmean,
3406 smoothed_is_neutral=nsm_neutral == 1,
3407 recfast_fudge=chem_fudge,
3408 recfast_hswitch=chem_recfast_hswitch,
3409 hubble_expansion=chem_hubble_expansion,
3410 adot_over_a=chem_adot_over_a,
3411 hubble=100hub, Om=Om, OL=OL, fh=Float64(fh),
3412 rate_tables=chem_rate_tables, dtfrac=chem_dtfrac,
3413 itcap=chem_itcap, backend=BACKEND_NAME, precision=T)
3414 end
3415 else
3416 hz = _hubble_s(zchem, hub, Om, OL, Or)
3417 _chem_peebles_k!(be_p)(s.U[1], eint, HII, H2I, nsm_arg,
3418 nsm_scalar, nsm_is_scalar, nsm_neutral,
3419 T(dunit), T(velocity_unit2), T(1.0),
3420 T(dtphys), T(zchem), T(hz), fh, T(s.γ),
3421 T(f_alpha_step), T(xmean), T(chem_fudge),
3422 T(gauss), T(chem_dtfrac), chem_itcap;
3423 ndrange=ncells)
3424 end
3425 end
3426 if profile_chem
3427 chem_phase[:solve] += _time_phase(solve_block, be_p, be_mhd)
3428 else
3429 solve_block()
3430 end
3431 else
3432 if profile_chem
3433 chem_phase[:solve] += _time_phase(be_p, be_mhd) do
3434 ChemistryKernels.solve_chem_device!(s.U[1], eint, HII, H2I;
3435 a_value=1/(1+zchem), dt=dtphys, density_units=dunit,
3436 length_units=sqrt(velocity_unit2), time_units=1,
3437 hubble=100hub, Om=Om, OL=OL, fh=Float64(fh), deuterium=false,
3438 hubble_expansion=chem_hubble_expansion,
3439 adot_over_a=chem_adot_over_a,
3440 dtfrac=chem_dtfrac, itcap=chem_itcap, backend=BACKEND_NAME,
3441 precision=T)
3442 end
3443 else
3444 ChemistryKernels.solve_chem_device!(s.U[1], eint, HII, H2I;
3445 a_value=1/(1+zchem), dt=dtphys, density_units=dunit,
3446 length_units=sqrt(velocity_unit2), time_units=1,
3447 hubble=100hub, Om=Om, OL=OL, fh=Float64(fh), deuterium=false,
3448 hubble_expansion=chem_hubble_expansion,
3449 adot_over_a=chem_adot_over_a,
3450 dtfrac=chem_dtfrac, itcap=chem_itcap, backend=BACKEND_NAME,
3451 precision=T)
3452 end
3453 end
3454 chem_nonlocal_before || run_nonlocal!()
3455 if profile_chem
3456 if !use_fused_chem
3457 chem_phase[:clamp] += _time_phase(be_p, be_mhd) do
3458 _clamp_reduced_chem_k!(be_p)(
3459 HII, H2I, s.U[1], fh, T(hydrogen_reference),
3460 T(hydrogen_complement_reference), Int(N),
3461 Val(chem_centered_hydrogen_storage);
3462 ndrange=ncells)
3463 chem_helium && _clamp_reduced_helium_k!(be_p)(
3464 HeII, s.U[1], fh, Int(N); ndrange=ncells)
3465 end
3466 end
3467 if refresh_mean
3468 chem_phase[:mean] += @elapsed begin
3469 xe_mean = _device_hii_mean!(
3470 be_p, s.scratch[1], s.U[1], HII, H2I, fh,
3471 internal_neutral_hydrogen,
3472 chem_centered_hydrogen_storage, hydrogen_reference)
3473 if chem_centered_hydrogen_storage
3474 _recenter_hydrogen_delta_k!(be_p)(
3475 HII, T(xe_mean - hydrogen_reference); ndrange=ncells)
3476 KA.synchronize(be_p)
3477 end
3478 chem_helium && (xheii_mean = _device_species_mean!(
3479 be_p, s.scratch[1], s.U[1], HeII, 4T(fh)))
3480 end
3481 end
3482 if !chemistry_updates_total_energy
3483 chem_phase[:repack] += _time_phase(be_p, be_mhd) do
3484 _eint_delta_to_mhd_k!(be_p)(
3485 s.U[5],eint,s.scratch[9],s.U[1],energy_residual,
3486 T(s.smallr),track_energy;
3487 ndrange=ncells)
3488 end
3489 end
3490 else
3491 if !use_fused_chem
3492 _clamp_reduced_chem_k!(be_p)(
3493 HII, H2I, s.U[1], fh, T(hydrogen_reference),
3494 T(hydrogen_complement_reference), Int(N),
3495 Val(chem_centered_hydrogen_storage);
3496 ndrange=ncells)
3497 chem_helium && _clamp_reduced_helium_k!(be_p)(
3498 HeII, s.U[1], fh, Int(N); ndrange=ncells)
3499 KA.synchronize(be_p)
3500 end
3501 if refresh_mean
3502 xe_mean = _device_hii_mean!(
3503 be_p, s.scratch[1], s.U[1], HII, H2I, fh,
3504 internal_neutral_hydrogen,
3505 chem_centered_hydrogen_storage, hydrogen_reference)
3506 if chem_centered_hydrogen_storage
3507 _recenter_hydrogen_delta_k!(be_p)(
3508 HII, T(xe_mean - hydrogen_reference); ndrange=ncells)
3509 KA.synchronize(be_p)
3510 end
3511 chem_helium && (xheii_mean = _device_species_mean!(
3512 be_p, s.scratch[1], s.U[1], HeII, 4T(fh)))
3513 end
3514 if !chemistry_updates_total_energy
3515 _eint_delta_to_mhd_k!(be_p)(
3516 s.U[5],eint,s.scratch[9],s.U[1],energy_residual,
3517 T(s.smallr),track_energy;
3518 ndrange=ncells)
3519 end
3520 end
3521 end
3522 return tc, xe_mean, xheii_mean
3523 end
3524
3525 function _hubble_s(z, h, Om, OL, Or)
3526 H0 = 100.0 * h * 1.0e5 / _PMF_MPC_CM
3527 Ok = 1.0 - Om - OL - Or
3528 H0 * sqrt(Om * (1 + z)^3 + Ok * (1 + z)^2 + Or * (1 + z)^4 + OL)
3529 end
3530
3531 function _advance_a_rk2(Om, OL, Ok, Or, a, dt, a_end)
3532 k1 = dadtau(Om, OL, Ok, Or, a)
3533 amid = a + 0.5 * k1 * dt
3534 k2 = dadtau(Om, OL, Ok, Or, amid)
3535 return min(a + k2 * dt, a_end)
3536 end
3537
3538 function _dtau_to_scale_factor_rk2(Om, OL, Ok, Or, a, target_a)
3539 target_a > a || return 0.0
3540 estimate = (target_a - a) / max(dadtau(Om, OL, Ok, Or, a), eps(Float64))
3541 lo = 0.0
3542 hi = max(estimate, eps(Float64))
3543 while _advance_a_rk2(Om, OL, Ok, Or, a, hi, Inf) < target_a
3544 hi *= 2
3545 end
3546 for _ in 1:56
3547 mid = 0.5 * (lo + hi)
3548 if _advance_a_rk2(Om, OL, Ok, Or, a, mid, Inf) < target_a
3549 lo = mid
3550 else
3551 hi = mid
3552 end
3553 end
3554 return hi
3555 end
3556
3557 function _limit_dt_to_output_redshift(dt, Om, OL, Ok, Or, a, output_zs, next_index)
3558 next_index <= length(output_zs) || return Float64(dt)
3559 target_a = z_to_a(output_zs[next_index])
3560 tolerance = 64eps(Float64) * max(1.0, abs(a), abs(target_a))
3561 target_a > a + tolerance || return Float64(dt)
3562 estimate = dtau_for_dlna(Om, OL, Ok, Or, a, log(target_a / a))
3563 estimate <= Float64(dt) * (1 + 1e-8) || return Float64(dt)
3564 return min(Float64(dt),
3565 _dtau_to_scale_factor_rk2(Om, OL, Ok, Or, a, target_a))
3566 end
3567
3568 function _cosmic_dt_seconds(a0, a1, h, Om, OL, Or)
3569 amid = 0.5 * (a0 + a1)
3570 zmid = a_to_z(amid)
3571 max(a1 - a0, 0.0) / (amid * _hubble_s(zmid, h, Om, OL, Or))
3572 end
3573
3574 function _compton_drag_rate_s(z, xe, fh)
3575 (4 / 3) * _PMF_ARAD * (_PMF_T_CMB * (1 + z))^4 *
3576 max(xe, 0.0) * fh * _PMF_SIGT / (_PMF_CL * ChemistryKernels.MH)
3577 end
3578
3579 function _drag_electron_fraction(xhii, z, h, Ob, fh, include_helium)
3580 include_helium || return max(Float64(xhii), 0.0)
3581 nH = _density_unit_cgs(z, h, Ob) * fh / ChemistryKernels.MH
3582 return ChemistryKernels.total_electron_fraction(
3583 Float64(xhii), nH, _PMF_T_CMB * (1 + z); fh=Float64(fh))
3584 end
3585
3586 function _drag_electron_fraction(xhii, xheii, z, h, Ob, fh, include_helium)
3587 include_helium || return max(Float64(xhii), 0.0)
3588 nH = _density_unit_cgs(z, h, Ob) * fh / ChemistryKernels.MH
3589 return ChemistryKernels.total_electron_fraction(
3590 Float64(xhii), Float64(xheii), nH, _PMF_T_CMB * (1 + z);
3591 fh=Float64(fh))
3592 end
3593
3594 function _density_unit_cgs(z, h, Ob)
3595 Ob * 1.8788e-29 * h^2 * (1 + z)^3
3596 end
3597
3598 function _equilibrium_helium_state(xhii, z, h, Ob, fh)
3599 nH = _density_unit_cgs(z, h, Ob) * Float64(fh) / ChemistryKernels.MH
3600 Trad = _PMF_T_CMB * (1 + Float64(z))
3601 xe = ChemistryKernels.total_electron_fraction(
3602 Float64(xhii), nH, Trad; fh=Float64(fh))
3603 s1, s2 = ChemistryKernels.helium_saha_pair(Trad)
3604 ne = max(xe * nH, 1e-300)
3605 r1 = s1 / ne
3606 r2 = s2 / ne
3607 denominator = 1 + r1 + r1 * r2
3608 fHe = (1 - Float64(fh)) / (4Float64(fh))
3609 xheii = fHe * r1 / denominator
3610 return (
3611 xe_total=xe,
3612 xheii=xheii,
3613 heii_mass_fraction=4Float64(fh) * xheii,
3614 )
3615 end
3616
3617 function _radiation_closure(z,xe,h,Ob,fh,vunit,lbox_comoving_cm,
3618 gas_sound_speed;bridge_q2=5,
3619 longitudinal_viscosity=4/3,
3620 nyquist_guard_fraction=1,
3621 response_model=:angular)
3622 isfinite(lbox_comoving_cm) && lbox_comoving_cm>0 ||
3623 error("spectral radiation coupling needs MHD_TERMINAL_LBOX_CKPC or MHD_TERMINAL_LBOX_CM")
3624 rho_b=_density_unit_cgs(z,h,Ob)
3625 Tcmb=_PMF_T_CMB*(1+z)
3626 rho_gamma_energy=_PMF_ARAD*Tcmb^4
3627 R=3*rho_b*_PMF_CL^2/(4*rho_gamma_energy)
3628 ne=rho_b*max(xe,eps(Float64))*Float64(fh)/ChemistryKernels.MH
3629 lambda_phys=1/(ne*_PMF_SIGT)
3630 lphys=lbox_comoving_cm/(1+z)
3631 return MHDKernels.RadiationClosure(
3632 inertia_ratio=R,
3633 mean_free_path=lambda_phys/lphys,
3634 photon_speed=_PMF_CL/vunit,
3635 gas_sound_speed=gas_sound_speed,
3636 bridge_q2=bridge_q2,
3637 longitudinal_viscosity=longitudinal_viscosity,
3638 nyquist_guard_fraction=nyquist_guard_fraction,
3639 response_model=Symbol(response_model))
3640 end
3641
3642 @inline function _photon_baryon_inertia_ratio(z,h,Ob)
3643 return 3*_density_unit_cgs(z,h,Ob)*_PMF_CL^2 /
3644 (4*_PMF_ARAD*(_PMF_T_CMB*(1+z))^4)
3645 end
3646
3647 function _radiation_scale_diagnostics(rc, n)
3648 kfund = 2pi
3649 knyquist = pi * n
3650 fund = MHDKernels.radiation_mode_coefficients(kfund, 0, rc)
3651 nyquist = MHDKernels.radiation_mode_coefficients(knyquist, 0, rc)
3652 return (qfund=rc.mean_free_path * kfund,
3653 qnyquist=rc.mean_free_path * knyquist,
3654 stream_fund=fund.streaming_weight,
3655 stream_nyquist=nyquist.streaming_weight,
3656 crad_fund=sqrt(max(fund.radiation_sound_speed2, 0)),
3657 crad_nyquist=sqrt(max(nyquist.radiation_sound_speed2, 0)))
3658 end
3659
3660 function _cosmological_ionized_p0(z, fh, vunit, h, Ob)
3661 z > -1 || error("automatic MHD_P0 requires a valid start redshift")
3662 vunit > 0 || error("automatic MHD_P0 requires positive MHD_CHEM_VUNIT")
3663 nH = _density_unit_cgs(z, h, Ob) * Float64(fh) / ChemistryKernels.MH
3664 temperature = _PMF_T_CMB * (1 + Float64(z))
3665 xe = ChemistryKernels.total_electron_fraction(
3666 1.0, nH, temperature; fh=Float64(fh))
3667 # Hydrogen nuclei + helium nuclei + all free electrons, per baryon mass.
3668 particles_per_h_mass = Float64(fh) * (1 + xe) +
3669 (1 - Float64(fh)) / 4
3670 return particles_per_h_mass * ChemistryKernels.KBOLTZ * temperature /
3671 (ChemistryKernels.MH * Float64(vunit)^2)
3672 end
3673
3674 function _pmf_code_brms_from_comoving_ng(b0_ng, zref, h, Ob, vunit)
3675 zref > -1 || error("MHD_PMF_B0_NG requires a valid redshift reference")
3676 vunit > 0 || error("MHD_CHEM_VUNIT must be positive for MHD_PMF_B0_NG")
3677 ρ = _density_unit_cgs(zref, h, Ob)
3678 bphys = Float64(b0_ng) * 1e-9 * (1 + zref)^2
3679 bunit = Float64(vunit) * sqrt(4π * ρ)
3680 return bphys / bunit
3681 end
3682
3683 function _terminal_gamma_code(gamma_drag, dtphys, dtcode, vunit, lbox_comoving_cm, amid,
3684 magnetic_a_ref=NaN)
3685 if !(isfinite(lbox_comoving_cm) && lbox_comoving_cm > 0)
3686 return gamma_drag * dtphys / max(dtcode, eps(Float64))
3687 end
3688 lphys = lbox_comoving_cm * amid
3689 gamma_code = gamma_drag * lphys / max(vunit, eps(Float64))
3690 # B code units are fixed when a comoving physical amplitude is mapped at
3691 # a_ref. The terminal force clock otherwise acquires a spurious a/a_ref.
3692 return isfinite(magnetic_a_ref) ? gamma_code * amid / magnetic_a_ref : gamma_code
3693 end
3694
3695 function _hydro_time_per_tau(a, h, vunit, lbox_comoving_cm,
3696 cosmological_mhd_aref=NaN)
3697 if !(isfinite(lbox_comoving_cm) && lbox_comoving_cm > 0)
3698 return 1.0
3699 end
3700 H0 = 100.0 * h * 1.0e5 / _PMF_MPC_CM
3701 scale = isfinite(cosmological_mhd_aref) ?
3702 sqrt(a * cosmological_mhd_aref) : a
3703 return scale * vunit / (H0 * lbox_comoving_cm)
3704 end
3705
3706 function _hydro_code_dt(dtphys, dtau, vunit, lbox_comoving_cm, amid,
3707 cosmological_mhd_aref=NaN)
3708 if !(isfinite(lbox_comoving_cm) && lbox_comoving_cm > 0)
3709 return Float64(dtau)
3710 end
3711 if isfinite(cosmological_mhd_aref)
3712 return Float64(dtphys) * vunit * sqrt(cosmological_mhd_aref) /
3713 (lbox_comoving_cm * amid^(3 / 2))
3714 end
3715 return Float64(dtphys) * vunit / (lbox_comoving_cm * amid)
3716 end
3717
3718 @inline function _cosmological_velocity_unit(vunit, amid, cosmological_mhd_aref)
3719 return isfinite(cosmological_mhd_aref) ?
3720 vunit * sqrt(cosmological_mhd_aref / amid) : vunit
3721 end
3722
3723 @inline function _hubble_momentum_decay(a0, a1, cosmological_mhd_aref)
3724 return isfinite(cosmological_mhd_aref) ? sqrt(a0 / a1) : 1.0
3725 end
3726
3727 function _source_drag_impulse_cap(mode, user_cap, gas_gravity,
3728 gravity_cap, slip_factor=1.0)
3729 base=mode in (:aware,:semiimplicit) ? Float64(user_cap) : Inf
3730 return gas_gravity ? min(base,Float64(gravity_cap)/Float64(slip_factor)) : base
3731 end
3732
3733 @inline function _hydro_subinterval(dt_hydro, dtau, sub_dtau)
3734 return Float64(dt_hydro) * Float64(sub_dtau) /
3735 max(Float64(dtau), eps(Float64))
3736 end
3737
3738 function _gas_gravity_kick_dt(dtau, a, h, vunit, lbox_comoving_cm,
3739 cosmological_mhd_aref=NaN)
3740 return Float64(dtau) /
3741 _hydro_time_per_tau(a, h, vunit, lbox_comoving_cm,
3742 cosmological_mhd_aref)
3743 end
3744
3745 @inline function _drag_impulse(gamma_drag, dtphys, dtcode, rate_code)
3746 return isfinite(rate_code) ? Float64(rate_code) * Float64(dtcode) :
3747 Float64(gamma_drag) * Float64(dtphys)
3748 end
3749
3750 function _parse_z_outputs(env::AbstractString)
3751 isempty(strip(env)) && return Float64[]
3752 return sort(parse.(Float64, split(env, ",")); rev=true)
3753 end
3754
3755 @inline function _reached_output_z(znow::Real, ztarget::Real)
3756 tol = 64eps(Float64) * max(1.0, abs(Float64(znow)), abs(Float64(ztarget)))
3757 return Float64(znow) <= Float64(ztarget) + tol
3758 end
3759
3760 """Write a central density slice and line-of-sight maximum projection.
3761
3762 Frames are deliberately reduced on the host after one reusable device-to-host
3763 copy. This keeps science-movie output to O(N^2) storage per redshift rather
3764 than retaining full O(N^3) density states.
3765 """
3766 function _write_density_frame!(frame_dir::AbstractString, hostrho, hostslice, hostproj,
3767 rho, N::Int, z::Real, cycle::Int; density_offset=0)
3768 isempty(frame_dir) && return
3769 mkpath(frame_dir)
3770 copyto!(hostrho, rho)
3771 density_offset == 0 || (hostrho .+= Float32(density_offset))
3772 mid = fld(N, 2) + 1
3773 @views copyto!(hostslice, hostrho[:, :, mid])
3774 @inbounds for j in 1:N, i in 1:N
3775 vmax = hostrho[i, j, 1]
3776 for k in 2:N
3777 vmax = max(vmax, hostrho[i, j, k])
3778 end
3779 hostproj[i, j] = vmax
3780 end
3781 tag = @sprintf("cycle%09d_z%.12g", cycle, Float64(z))
3782 slice_name = "density_slice_$(tag).f32"
3783 proj_name = "density_maxproj_$(tag).f32"
3784 open(joinpath(frame_dir, slice_name), "w") do io
3785 write(io, hostslice)
3786 end
3787 open(joinpath(frame_dir, proj_name), "w") do io
3788 write(io, hostproj)
3789 end
3790 manifest = joinpath(frame_dir, "frames.tsv")
3791 exists = isfile(manifest)
3792 open(manifest, "a") do io
3793 exists || println(io, "cycle\tz\tN\tslice_f32\tmaxproj_f32\tlayout")
3794 println(io, join((cycle, Float64(z), N, slice_name, proj_name, "julia-column-major"), '\t'))
3795 end
3796 return nothing
3797 end
3798
3799 function _write_run_density_frames!(frame_dir, hostrho, hostslice, hostproj,
3800 s, rho_dm, positions, velocities, N, z, cycle;
3801 field, lattice_displacements)
3802 if field in (:gas,:both)
3803 _write_density_frame!(frame_dir,hostrho,hostslice,hostproj,s.U[1],N,z,cycle)
3804 end
3805 if field in (:dm,:both)
3806 _deposit_particles!(rho_dm,positions...,velocities...;
3807 N=N,disp=0,lattice_displacements=lattice_displacements)
3808 KA.synchronize(s.be)
3809 _write_density_frame!(joinpath(frame_dir,"dm"),hostrho,hostslice,hostproj,
3810 rho_dm,N,z,cycle;density_offset=lattice_displacements ? 1 : 0)
3811 end
3812 end
3813
3814 function _write_f32_array_atomic(path::AbstractString, array)
3815 host = Array{Float32}(undef, size(array))
3816 copyto!(host, array)
3817 tmp = path * ".tmp"
3818 open(tmp, "w") do io
3819 write(io, host)
3820 end
3821 mv(tmp, path; force=true)
3822 host = nothing
3823 GC.gc()
3824 return nothing
3825 end
3826
3827 function _write_f64_array_atomic(path::AbstractString, array)
3828 host = Array{Float64}(undef, size(array))
3829 copyto!(host, array)
3830 tmp = path * ".tmp"
3831 open(tmp, "w") do io
3832 write(io, host)
3833 end
3834 mv(tmp, path; force=true)
3835 return nothing
3836 end
3837
3838 function _read_f32_array_into!(array, path::AbstractString, dims)
3839 expected = prod(dims)
3840 filesize(path) == expected * sizeof(Float32) ||
3841 error("bad raw Float32 size for $path: got $(filesize(path)) bytes, expected $(expected * sizeof(Float32))")
3842 host = Array{Float32}(undef, dims)
3843 open(path, "r") do io
3844 read!(io, host)
3845 eof(io) || error("extra trailing bytes in $path")
3846 end
3847 copyto!(array, host)
3848 host = nothing
3849 GC.gc()
3850 return nothing
3851 end
3852
3853 function _read_c32_array_into!(array, path::AbstractString, dims)
3854 expected = prod(dims)
3855 filesize(path) == expected * sizeof(ComplexF32) ||
3856 error("bad raw ComplexF32 size for $path: got $(filesize(path)) bytes, expected $(expected * sizeof(ComplexF32))")
3857 host = Array{ComplexF32}(undef, dims)
3858 open(path, "r") do io
3859 read!(io, host)
3860 eof(io) || error("extra trailing bytes in $path")
3861 end
3862 copyto!(array, host)
3863 host = nothing
3864 GC.gc()
3865 return nothing
3866 end
3867
3868 function _import_initial_mhd_state!(be, dir::AbstractString, s; N::Int)
3869 isempty(dir) && return nothing
3870 for i in eachindex(s.U)
3871 _read_f32_array_into!(s.U[i], joinpath(dir, "mhd_U_$i.f32"), (N, N, N))
3872 end
3873 rho_min = Float64(minimum(s.U[1]))
3874 energy_min = Float64(minimum(s.U[5]))
3875 isfinite(rho_min) && rho_min > 0 ||
3876 error("imported MHD state has invalid minimum density $rho_min")
3877 isfinite(energy_min) && energy_min > 0 ||
3878 error("imported MHD state has invalid minimum energy $energy_min")
3879 brms = _device_brms!(be, s, s.scratch[1])
3880 @printf("INITIAL_STATE_MHD_IMPORT path=%s n=%d brms=%.8e rho_min=%.8e E_min=%.8e\n",
3881 dir, N, brms, rho_min, energy_min)
3882 flush(stdout)
3883 return Float32(brms)
3884 end
3885
3886 function _initial_hydrogen_storage(dir::AbstractString)
3887 storage_path = joinpath(dir, "chem_hydrogen_storage.txt")
3888 storage = isfile(storage_path) ?
3889 Symbol(lowercase(strip(read(storage_path, String)))) : :ionized
3890 storage in (:ionized, :neutral, :centered) || error(
3891 "invalid initial-state hydrogen storage marker $storage in $storage_path")
3892 storage
3893 end
3894
3895 function _initial_hydrogen_reference(dir::AbstractString, requested_storage::Symbol)
3896 isempty(dir) && return nothing
3897 isfile(joinpath(dir, "chem_HII.f32")) || return nothing
3898 imported_storage = _initial_hydrogen_storage(dir)
3899 imported_storage === requested_storage || error(
3900 "initial-state chemistry uses $imported_storage hydrogen storage, " *
3901 "but this run requests $requested_storage; regenerate or migrate the handoff")
3902 requested_storage === :centered || return nothing
3903 reference_path = joinpath(dir, "chem_hydrogen_reference.txt")
3904 isfile(reference_path) || error(
3905 "centered initial-state chemistry requires $reference_path")
3906 reference = tryparse(Float64, strip(read(reference_path, String)))
3907 reference === nothing && error("invalid centered hydrogen reference in $reference_path")
3908 isfinite(reference) && 0 <= reference <= 1 || error(
3909 "centered hydrogen reference must lie in [0,1]; got $reference")
3910 reference
3911 end
3912
3913 function _import_initial_aux_state!(dir::AbstractString; N::Int,
3914 particles, chemistry, photons,
3915 density_perturbation=nothing,
3916 magnetic_residual=nothing,
3917 energy_residual=nothing,
3918 neutral_hydrogen_storage::Bool=false,
3919 centered_hydrogen_storage::Bool=false)
3920 isempty(dir) && return nothing
3921 neutral_hydrogen_storage && centered_hydrogen_storage && error(
3922 "initial-state hydrogen storage cannot be both neutral and centered")
3923 particle_names = ("particle_px", "particle_py", "particle_pz",
3924 "particle_vx", "particle_vy", "particle_vz")
3925 particle_present = map(name -> isfile(joinpath(dir, "$name.f32")), particle_names)
3926 all(particle_present) || !any(particle_present) ||
3927 error("initial-state particle import must provide all six particle arrays")
3928 if all(particle_present)
3929 for (array, name) in zip(particles, particle_names)
3930 _read_f32_array_into!(array, joinpath(dir, "$name.f32"), (N^3,))
3931 end
3932 end
3933 if chemistry !== nothing
3934 chem_names = length(chemistry) == 4 ?
3935 ("chem_eint", "chem_HII", "chem_H2I", "chem_HeII") :
3936 ("chem_eint", "chem_HII", "chem_H2I")
3937 length(chemistry) in (3, 4) ||
3938 error("initial-state chemistry import expects three H fields or four H+He fields")
3939 chem_present = map(name -> isfile(joinpath(dir, "$name.f32")), chem_names)
3940 all(chem_present) || !any(chem_present) ||
3941 error("initial-state chemistry import must provide all requested chemistry fields")
3942 if all(chem_present)
3943 imported_storage = _initial_hydrogen_storage(dir)
3944 requested_storage = centered_hydrogen_storage ? :centered :
3945 neutral_hydrogen_storage ? :neutral : :ionized
3946 imported_storage === requested_storage || error(
3947 "initial-state chemistry uses $imported_storage hydrogen storage, " *
3948 "but this run requests $requested_storage; regenerate or migrate the handoff")
3949 requested_storage === :centered &&
3950 _initial_hydrogen_reference(dir, requested_storage)
3951 for (array, name) in zip(chemistry, chem_names)
3952 _read_f32_array_into!(array, joinpath(dir, "$name.f32"), (N^3,))
3953 end
3954 end
3955 end
3956 photon_path = joinpath(dir, "photon_hat_batch.c32")
3957 if isfile(photon_path)
3958 photons === nothing && error("initial state contains photon moments but response_model is not moments")
3959 _read_c32_array_into!(photons.hat_batch, photon_path, (N ÷ 2 + 1, N, N, 4))
3960 elseif photons !== nothing
3961 error("response_model=moments initial-state import requires photon_hat_batch.c32")
3962 end
3963 density_path = joinpath(dir, "gas_delta.f32")
3964 if isfile(density_path)
3965 density_perturbation === nothing && error(
3966 "initial state contains gas_delta.f32 but radiation density tracking is inactive")
3967 _read_f32_array_into!(density_perturbation, density_path, (N^3,))
3968 end
3969 residual_names = ("magnetic_residual_1", "magnetic_residual_2",
3970 "magnetic_residual_3")
3971 residual_present = map(
3972 name -> isfile(joinpath(dir, "$name.f32")), residual_names)
3973 all(residual_present) || !any(residual_present) || error(
3974 "initial-state magnetic residual import must provide all three fields")
3975 if all(residual_present)
3976 magnetic_residual === nothing && error(
3977 "initial state contains magnetic residuals but compensation is inactive")
3978 length(magnetic_residual) == 3 || error(
3979 "initial-state magnetic residual import expects three fields")
3980 for (array, name) in zip(magnetic_residual, residual_names)
3981 _read_f32_array_into!(array, joinpath(dir, "$name.f32"), (N^3,))
3982 end
3983 end
3984 energy_residual_path=joinpath(dir,"energy_residual.f32")
3985 if isfile(energy_residual_path)
3986 energy_residual === nothing && error(
3987 "initial state contains an energy residual but compensation is inactive")
3988 _read_f32_array_into!(energy_residual,energy_residual_path,(N^3,))
3989 end
3990 @printf("INITIAL_STATE_AUX_IMPORT path=%s particles=%s chemistry=%s photons=%s gas_delta=%s magnetic_residual=%s energy_residual=%s\n",
3991 dir, string(all(particle_present)), string(chemistry !== nothing &&
3992 isfile(joinpath(dir, "chem_eint.f32"))), string(isfile(photon_path)),
3993 string(isfile(density_path)), string(all(residual_present)),
3994 string(isfile(energy_residual_path)))
3995 flush(stdout)
3996 return nothing
3997 end
3998
3999 @kernel function _reset_imported_pmf_state_k!(rho, mx, my, mz, E,
4000 @Const(Bx), @Const(By), @Const(Bz),
4001 gamma, p0)
4002 c = @index(Global)
4003 @inbounds begin
4004 T = eltype(E)
4005 rho[c] = one(T)
4006 mx[c] = zero(T)
4007 my[c] = zero(T)
4008 mz[c] = zero(T)
4009 E[c] = T(p0) / (T(gamma) - one(T)) +
4010 T(0.5) * (Bx[c] * Bx[c] + By[c] * By[c] + Bz[c] * Bz[c])
4011 end
4012 end
4013
4014 @kernel function _scale_imported_b_k!(Bx, By, Bz, factor)
4015 c = @index(Global)
4016 @inbounds begin
4017 Bx[c] *= factor
4018 By[c] *= factor
4019 Bz[c] *= factor
4020 end
4021 end
4022
4023 function _import_initial_b_field!(be, dir::AbstractString, s; N::Int, p0,
4024 target_brms=nothing)
4025 isempty(dir) && return nothing
4026 _read_f32_array_into!(s.U[6], joinpath(dir, "bx.f32"), (N, N, N))
4027 _read_f32_array_into!(s.U[7], joinpath(dir, "by.f32"), (N, N, N))
4028 _read_f32_array_into!(s.U[8], joinpath(dir, "bz.f32"), (N, N, N))
4029 T = eltype(s.U[5])
4030 imported_brms = _device_brms!(be, s, s.scratch[1])
4031 if target_brms !== nothing
4032 target_brms > 0 || error("MHD imported target Brms must be positive")
4033 imported_brms > 0 || error("cannot rescale an imported zero magnetic field")
4034 factor = T(target_brms / imported_brms)
4035 _scale_imported_b_k!(be)(s.U[6], s.U[7], s.U[8], factor; ndrange=N^3)
4036 KA.synchronize(be)
4037 end
4038 _reset_imported_pmf_state_k!(be)(s.U[1], s.U[2], s.U[3], s.U[4], s.U[5],
4039 s.U[6], s.U[7], s.U[8], T(s.γ), T(p0);
4040 ndrange=N^3)
4041 KA.synchronize(be)
4042 brms = _device_brms!(be, s, s.scratch[1])
4043 @printf("INITIAL_B_IMPORT path=%s n=%d raw_brms=%.8e brms=%.8e rescaled=%s\n",
4044 dir, N, imported_brms, brms, string(target_brms !== nothing))
4045 flush(stdout)
4046 return Float32(brms)
4047 end
4048
4049 function _write_initial_b_field!(dir::AbstractString, s, st; N::Int, boxsize,
4050 zstart, zend, seed, kcut, kmax,
4051 mode_lock_n, high_mode_lock_n, brms,
4052 normalization_z, target_comoving_ng,
4053 reference_comoving_ng, total_comoving_ng)
4054 isempty(dir) && return nothing
4055 mkpath(dir)
4056 _write_f32_array_atomic(joinpath(dir, "bx.f32"), s.U[6])
4057 _write_f32_array_atomic(joinpath(dir, "by.f32"), s.U[7])
4058 _write_f32_array_atomic(joinpath(dir, "bz.f32"), s.U[8])
4059 open(joinpath(dir, "metadata.csv.tmp"), "w") do io
4060 println(io, "n,boxsize,zstart,zend,seed,kcut,kmax,mode_lock_n,high_mode_lock_n,normalization_kind,normalization_z,target_comoving_ng,reference_comoving_ng,total_comoving_ng,brms,brms_measured")
4061 println(io, join((N, boxsize, zstart, zend, seed, kcut,
4062 kmax === nothing ? "" : kmax,
4063 mode_lock_n === nothing ? "" : mode_lock_n,
4064 high_mode_lock_n === nothing ? "" : high_mode_lock_n,
4065 st.normalization_kind, normalization_z,
4066 target_comoving_ng, reference_comoving_ng,
4067 total_comoving_ng,
4068 brms, st.brms_measured), ','))
4069 end
4070 mv(joinpath(dir, "metadata.csv.tmp"), joinpath(dir, "metadata.csv"); force=true)
4071 @printf("INITIAL_B_EXPORT path=%s n=%d brms=%.8e\n", dir, N, st.brms_measured)
4072 flush(stdout)
4073 return nothing
4074 end
4075
4076 function _write_final_b_field!(dir::AbstractString, s; N::Int, z, cycle, brms)
4077 isempty(dir) && return nothing
4078 mkpath(dir)
4079 suffix = eltype(s.U[6]) === Float64 ? "f64" : "f32"
4080 writer = eltype(s.U[6]) === Float64 ? _write_f64_array_atomic : _write_f32_array_atomic
4081 writer(joinpath(dir, "bx.$suffix"), s.U[6])
4082 writer(joinpath(dir, "by.$suffix"), s.U[7])
4083 writer(joinpath(dir, "bz.$suffix"), s.U[8])
4084 open(joinpath(dir, "metadata.csv.tmp"), "w") do io
4085 println(io, "n,z,cycle,brms,precision")
4086 println(io, join((N, z, cycle, brms, suffix), ','))
4087 end
4088 mv(joinpath(dir, "metadata.csv.tmp"), joinpath(dir, "metadata.csv"); force=true)
4089 @printf("FINAL_B_EXPORT path=%s n=%d z=%.8e brms=%.8e\n",
4090 dir, N, z, brms)
4091 flush(stdout)
4092 return nothing
4093 end
4094
4095 function _write_history!(path, cycle, measured, a, br, divn, δb, δdm;
4096 btrans=NaN, btrans0=NaN, delta2k=NaN,
4097 thermalstats=(NaN, NaN, NaN, NaN, NaN, NaN),
4098 temperaturestats=(NaN, NaN, NaN), divb_face=NaN)
4099 isempty(path) && return
4100 exists = isfile(path)
4101 mkpath(dirname(path))
4102 open(path, "a") do io
4103 exists || println(io, join(("cycle","measured","z","a","brms","divBdx_over_brms",
4104 "delta_b_rms","delta_dm_rms",
4105 "xHII_mean","xHII_min","xHII_max",
4106 "fH2_mean","fH2_min","fH2_max",
4107 "xHeII_mean","xHeII_min","xHeII_max",
4108 "xe_total_mean","xe_total_min","xe_total_max",
4109 "btrans_rms","btrans_ratio",
4110 "delta2k_amp",
4111 "p_mean","p_min","p_max",
4112 "eint_mean","eint_min","eint_max",
4113 "temperature_mean","temperature_min",
4114 "temperature_max","divb_face_dx_over_brms"), '\t'))
4115 println(io, join((cycle, measured, a_to_z(a), a, br, divn, δb, δdm,
4116 ntuple(_ -> NaN, 12)...,
4117 btrans, btrans / btrans0, delta2k,
4118 thermalstats..., temperaturestats..., divb_face), '\t'))
4119 end
4120 end
4121
4122 function _write_history!(path, cycle, measured, a, br, divn, δb, δdm, chemstats;
4123 btrans=NaN, btrans0=NaN, delta2k=NaN,
4124 thermalstats=(NaN, NaN, NaN, NaN, NaN, NaN),
4125 temperaturestats=(NaN, NaN, NaN), divb_face=NaN)
4126 isempty(path) && return
4127 exists = isfile(path)
4128 mkpath(dirname(path))
4129 open(path, "a") do io
4130 exists || println(io, join(("cycle","measured","z","a","brms","divBdx_over_brms",
4131 "delta_b_rms","delta_dm_rms",
4132 "xHII_mean","xHII_min","xHII_max",
4133 "fH2_mean","fH2_min","fH2_max",
4134 "xHeII_mean","xHeII_min","xHeII_max",
4135 "xe_total_mean","xe_total_min","xe_total_max",
4136 "btrans_rms","btrans_ratio",
4137 "delta2k_amp",
4138 "p_mean","p_min","p_max",
4139 "eint_mean","eint_min","eint_max",
4140 "temperature_mean","temperature_min",
4141 "temperature_max","divb_face_dx_over_brms"), '\t'))
4142 println(io, join((cycle, measured, a_to_z(a), a, br, divn, δb, δdm,
4143 chemstats..., btrans, btrans / btrans0, delta2k,
4144 thermalstats..., temperaturestats..., divb_face), '\t'))
4145 end
4146 end
4147
4148 function _write_pkmu!(path, label, result; cycle, measured, a, nmu)
4149 isempty(path) && return
4150 exists = isfile(path)
4151 mkpath(dirname(path))
4152 open(path, "a") do io
4153 exists || println(io, join(("field","cycle","measured","z","a","kbin","mubin",
4154 "k","mu_lo","mu_hi","P","Nmodes"), '\t'))
4155 for b in eachindex(result.k), m in 1:nmu
4156 μlo = (m - 1) / nmu
4157 μhi = m / nmu
4158 println(io, join((label, cycle, measured, a_to_z(a), a, b, m, result.k[b],
4159 μlo, μhi, result.P[b,m], result.Nmodes[b,m]), '\t'))
4160 end
4161 end
4162 end
4163
4164 function _write_density_pdf!(path, label, counts_dev, rho; be, cycle, measured, a,
4165 loglo::Real, loghi::Real)
4166 isempty(path) && return
4167 nbins = length(counts_dev)
4168 fill!(counts_dev, Int32(0))
4169 width = (Float64(loghi) - Float64(loglo)) / Float64(nbins)
4170 invw = 1.0 / width
4171 _logrho_hist_k!(be)(counts_dev, rho, Int(nbins), T(loglo), T(invw); ndrange=length(rho))
4172 KA.synchronize(be)
4173 counts = Vector{Int32}(undef, nbins)
4174 copyto!(counts, counts_dev)
4175
4176 exists = isfile(path)
4177 mkpath(dirname(path))
4178 norm = Float64(length(rho)) * width
4179 open(path, "a") do io
4180 exists || println(io, join(("field","cycle","measured","z","a","bin",
4181 "logrho_lo","logrho_hi","rho_mid","count","pdf"), '\t'))
4182 for b in 1:nbins
4183 lo = Float64(loglo) + (b - 1) * width
4184 hi = lo + width
4185 mid = 10.0 ^ (0.5 * (lo + hi))
4186 pdf = Float64(counts[b]) / norm
4187 println(io, join((label, cycle, measured, a_to_z(a), a, b,
4188 lo, hi, mid, counts[b], pdf), '\t'))
4189 end
4190 end
4191 end
4192
4193 function _write_linear_mode!(path; be, s, rho_dm, scratch, N, gravity,
4194 cycle, measured, a, kmode, boxsize, hubble,
4195 omega_b=0.049,
4196 rho_dm_mean=one(eltype(rho_dm)),
4197 particle_displacement_x=nothing,
4198 particle_velocity_x=nothing, particle_velocities=nothing, gas_velocity_unit=NaN,
4199 particle_positions=nothing, lattice_displacements=true,
4200 particle_velocity_cms_factor=NaN,
4201 physical_box_ckpc_h=NaN,
4202 gas_delta=nothing, HII=nothing, H2I=nothing,
4203 HeII=nothing,
4204 fh=0.76f0, photon_real_batch=nothing,
4205 temperature_velocity_unit2=NaN,
4206 neutral_hydrogen_storage=false,
4207 centered_hydrogen_storage=false,
4208 hydrogen_reference=1.0)
4209 isempty(path) && return
4210 0 < kmode < N ÷ 2 ||
4211 error("complex mode diagnostics require 0 < MHD_PMF_KMODE < N/2")
4212 rows = Pair{String,ComplexF64}[]
4213 for (name, src) in (("B_x", s.U[6]), ("B_y", s.U[7]), ("B_z", s.U[8]))
4214 mean = Float64(sum(src)) / Float64(length(src))
4215 push!(rows, name => _device_complex_mode!(be, scratch, src, N, kmode;
4216 mean=mean))
4217 end
4218 gas_source = gas_delta === nothing ? s.U[1] : gas_delta
4219 gas_mean = gas_delta === nothing ? one(eltype(gas_source)) : zero(eltype(gas_source))
4220 push!(rows, "gas_delta" => _device_complex_mode!(be, scratch, gas_source, N,
4221 kmode; mean=gas_mean))
4222 for (name, momentum) in (("gas_v_x", s.U[2]), ("gas_v_y", s.U[3]),
4223 ("gas_v_z", s.U[4]))
4224 mode = _device_complex_velocity_mode!(be, scratch, momentum,s.U[1],N,kmode,s.smallr)
4225 push!(rows,name=>mode)
4226 isfinite(gas_velocity_unit) && push!(rows,name*"_cms"=>gas_velocity_unit*mode)
4227 end
4228 MHDKernels.terminal_magnetic_force!(s, s.scratch[1], s.scratch[2],
4229 s.scratch[3])
4230 KA.synchronize(s.be)
4231 for d in 1:3
4232 axis = ("x", "y", "z")[d]
4233 push!(rows, "lorentz_accel_$axis" =>
4234 _device_complex_velocity_mode!(be, scratch, s.scratch[d],
4235 s.U[1], N, kmode,
4236 s.smallr))
4237 end
4238 if gravity
4239 push!(rows, "dm_delta" => _device_complex_mode!(be, scratch, rho_dm, N,
4240 kmode; mean=rho_dm_mean))
4241 if particle_displacement_x !== nothing
4242 push!(rows, "dm_displacement_x" =>
4243 _device_complex_mode!(be, scratch,
4244 particle_displacement_x, N,
4245 kmode))
4246 end
4247 velocities = particle_velocities === nothing ?
4248 (particle_velocity_x === nothing ? () : (particle_velocity_x,)) :
4249 particle_velocities
4250 if !isempty(velocities) && isfinite(particle_velocity_cms_factor) &&
4251 particle_positions !== nothing
4252 # Lorentz modes above have consumed scratch[1:3]; no evolution or
4253 # FFT is in flight here. Reuse two integrator targets, distinct from
4254 # the reduction scratch and all live state/particle/photon fields.
4255 velocity_grid, density_grid = s.scratch[1], s.scratch[2]
4256 for d in eachindex(velocities)
4257 MHDKernels.deposit_particle_velocity!(velocity_grid,density_grid,
4258 particle_positions...,velocities[d];
4259 N=N,lattice_displacements=lattice_displacements)
4260 mode = _device_complex_mode!(be,scratch,velocity_grid,N,kmode)
4261 push!(rows,"dm_v_"*("x","y","z")[d]=>particle_velocity_cms_factor*mode)
4262 end
4263 empty_fraction=Float64(mapreduce(iszero,+,density_grid))/length(density_grid)
4264 push!(rows,"dm_velocity_empty_fraction"=>complex(empty_fraction))
4265 end
4266 end
4267 if photon_real_batch !== nothing
4268 photon_delta = reshape(view(photon_real_batch, :, :, :, 1), :)
4269 push!(rows, "photon_delta" => _device_complex_mode!(be, scratch,
4270 photon_delta, N,
4271 kmode))
4272 for d in 1:3
4273 photon_velocity = reshape(view(photon_real_batch, :, :, :, d + 1), :)
4274 axis = ("x", "y", "z")[d]
4275 push!(rows, "photon_v_$axis" => _device_complex_mode!(be, scratch,
4276 photon_velocity,
4277 N, kmode))
4278 push!(rows, "photon_baryon_slip_$axis" =>
4279 _device_complex_slip_mode!(be, scratch, photon_velocity,
4280 s.U[d + 1], s.U[1], N,
4281 kmode, s.smallr))
4282 end
4283 end
4284 HII === nothing || push!(rows, "x_HII" =>
4285 _device_complex_hii_mode!(be, scratch, s.scratch[4], HII, H2I,
4286 s.U[1], N, kmode, fh,
4287 neutral_hydrogen_storage,
4288 centered_hydrogen_storage,
4289 hydrogen_reference))
4290 H2I === nothing || push!(rows, "x_H2" =>
4291 _device_complex_species_mode!(be, scratch, H2I, s.U[1], N, kmode, fh))
4292 HeII === nothing || push!(rows, "x_HeII" =>
4293 _device_complex_species_mode!(be, scratch, HeII, s.U[1], N, kmode,
4294 4 * fh))
4295 push!(rows, "gas_pressure_fraction" =>
4296 _device_complex_thermal_fraction_mode!(
4297 be, scratch, s.scratch[4], s, N, kmode, 1,
4298 ))
4299 push!(rows, "gas_specific_internal_energy_fraction" =>
4300 _device_complex_thermal_fraction_mode!(
4301 be, scratch, s.scratch[4], s, N, kmode, 2,
4302 ))
4303 if HII !== nothing && H2I !== nothing && isfinite(temperature_velocity_unit2)
4304 push!(rows, "gas_temperature_fraction" =>
4305 _device_complex_temperature_fraction_mode!(
4306 be, scratch, s.scratch[4], s, HII, H2I, HeII, N, kmode,
4307 temperature_velocity_unit2, a_to_z(a), fh, hubble, omega_b,
4308 neutral_hydrogen_storage,
4309 centered_hydrogen_storage,
4310 hydrogen_reference,
4311 ))
4312 end
4313
4314 k_box = isfinite(physical_box_ckpc_h) ?
4315 2π * Float64(kmode) / physical_box_ckpc_h : NaN
4316 k_mpc = 1000.0 * Float64(hubble) * k_box
4317 exists = isfile(path)
4318 mkpath(dirname(path))
4319 open(path, "a") do io
4320 exists || println(io, join(("field", "cycle", "measured", "z", "a",
4321 "kmode", "k_h_ckpc", "k_comoving_mpc_inv",
4322 "real", "imag", "amplitude", "phase"), '\t'))
4323 for (field, value) in rows
4324 println(io, join((field, cycle, measured, a_to_z(a), a, kmode,
4325 k_box, k_mpc, real(value), imag(value), abs(value),
4326 angle(value)), '\t'))
4327 end
4328 end
4329 return nothing
4330 end
4331
4332 "Write DM-mass-weighted streaming statistics at actual particle positions."
4333 function _write_streaming!(path, s, positions, velocities, scratch;
4334 a, cycle, measured, gas_unit, dm_unit,
4335 lattice_displacements)
4336 isempty(path) && return
4337 out = s.scratch[1:3]
4338 MHDKernels.sample_baryon_dm_velocity!(out, s, positions, velocities;
4339 lattice_displacements=lattice_displacements, gas_unit=gas_unit, dm_unit=dm_unit)
4340 KA.synchronize(s.be)
4341 means = ntuple(d->Float64(sum(out[d]))/length(out[d]),3)
4342 work = reshape(scratch,:)
4343 second = ntuple(3) do d
4344 @. work = out[d]^2
4345 Float64(sum(work))/length(work)
4346 end
4347 variance = ntuple(3) do d
4348 mean_d = eltype(work)(means[d])
4349 @. work = (out[d]-mean_d)^2
4350 Float64(sum(work))/length(work)
4351 end
4352 dm_means = ntuple(d->dm_unit*Float64(sum(velocities[d]))/length(velocities[d]),3)
4353 exists = isfile(path)
4354 mkpath(dirname(path))
4355 open(path,"a") do io
4356 exists || println(io,"cycle\tmeasured\tz\ta\tvbc_x_cms\tvbc_y_cms\tvbc_z_cms\tvbc_rms_cms\tvbc_dispersion_cms\tvc_x_cms\tvc_y_cms\tvc_z_cms\tvb_sampled_x_cms\tvb_sampled_y_cms\tvb_sampled_z_cms")
4357 println(io,join((cycle,measured,a_to_z(a),a,means...,sqrt(sum(second)),
4358 sqrt(sum(variance)),dm_means...,
4359 (means .+ dm_means)...),'\t'))
4360 end
4361 end
4362
4363 function _emit_diagnostics!(; be_p, s, ρdm, scratch, px, py, pz, vx, vy, vz,
4364 N, ncells, gravity, cycle, measured, a, pk_path,
4365 history_path, pdf_path, pdf_counts, pdf_loglo, pdf_loghi,
4366 nmu, nbins, boxsize, axis, kmax=nothing,
4367 lattice_displacements=false,
4368 HII=nothing, H2I=nothing, HeII=nothing, fh=0.76f0,
4369 btrans0=NaN, density_kmode=1,
4370 mode_history_path="", mode_kmode=density_kmode,
4371 hubble=0.674,
4372 omega_b=0.049,
4373 radiation_workspace=nothing,
4374 radiation_photon_state=nothing,
4375 radiation_drag_rate=NaN,
4376 temperature_velocity_unit2=NaN,
4377 physical_box_ckpc_h=NaN, gas_velocity_unit=NaN, streaming_path="",
4378 neutral_hydrogen_storage=false,
4379 centered_hydrogen_storage=false,
4380 hydrogen_reference=1.0)
4381 T = eltype(s.U[1])
4382 if gravity
4383 _deposit_particles!(ρdm, px, py, pz, vx, vy, vz;
4384 N=N, disp=0, lattice_displacements=lattice_displacements)
4385 KA.synchronize(be_p)
4386 end
4387 if gravity && !isempty(streaming_path)
4388 _write_streaming!(streaming_path,s,(px,py,pz),(vx,vy,vz),scratch;
4389 a=a,cycle=cycle,measured=measured,gas_unit=gas_velocity_unit,
4390 dm_unit=_particle_velocity_cms_per_code(a,physical_box_ckpc_h),
4391 lattice_displacements=lattice_displacements)
4392 end
4393 br = _device_brms!(be_p, s, scratch)
4394 divn = _device_divb_norm!(be_p, s, scratch, br)
4395 divb_face = _device_divb_face_norm(s, br)
4396 gas_delta=radiation_workspace === nothing ? nothing :
4397 radiation_workspace.density_perturbation
4398 dm_mean = lattice_displacements ? zero(T) : one(T)
4399 δb = gas_delta === nothing ? _device_delta_rms!(be_p,scratch,s.U[1],one(T)) :
4400 _device_delta_rms!(be_p,scratch,gas_delta,zero(T))
4401 δdm = gravity ? _device_delta_rms!(be_p, scratch, ρdm, dm_mean) : NaN
4402 chemstats = _device_species_stats!(
4403 be_p, scratch, s.U[1], HII, H2I, HeII, fh;
4404 z=a_to_z(a), h=hubble, Ob=omega_b,
4405 neutral_hydrogen_storage=neutral_hydrogen_storage,
4406 centered_hydrogen_storage=centered_hydrogen_storage,
4407 hydrogen_reference=hydrogen_reference)
4408 thermalstats = _device_thermal_stats!(be_p, s, scratch)
4409 temperaturestats = isfinite(temperature_velocity_unit2) ?
4410 _device_temperature_stats!(be_p, s, scratch, HII, H2I, HeII,
4411 _density_unit_cgs(a_to_z(a), hubble, omega_b),
4412 temperature_velocity_unit2, a_to_z(a), fh,
4413 neutral_hydrogen_storage,
4414 centered_hydrogen_storage,
4415 hydrogen_reference) :
4416 (NaN, NaN, NaN)
4417 btrans = _device_component_rms!(be_p, scratch, s.U[7])
4418 delta2k = gas_delta === nothing ?
4419 _device_density_cos_mode_amp!(be_p,scratch,s.U[1],N,density_kmode,2) :
4420 _device_density_cos_mode_amp!(be_p,scratch,gas_delta,N,density_kmode,2;
4421 mean=zero(T))
4422 _write_history!(history_path, cycle, measured, a, br, divn, δb, δdm, chemstats;
4423 btrans=btrans, btrans0=btrans0, delta2k=delta2k,
4424 thermalstats=thermalstats,
4425 temperaturestats=temperaturestats, divb_face=divb_face)
4426
4427 photon_real_ready = false
4428 if !isempty(mode_history_path)
4429 if radiation_workspace !== nothing && radiation_photon_state !== nothing
4430 MHDKernels.photon_fields_to_workspace!(radiation_workspace,
4431 radiation_photon_state)
4432 photon_real_ready = true
4433 end
4434 _write_linear_mode!(mode_history_path; be=be_p, s=s, rho_dm=ρdm,
4435 scratch=scratch, N=N, gravity=gravity, cycle=cycle,
4436 measured=measured, a=a, kmode=mode_kmode,
4437 boxsize=boxsize, hubble=hubble, omega_b=omega_b,
4438 rho_dm_mean=dm_mean,
4439 particle_displacement_x=lattice_displacements ? px : nothing,
4440 particle_velocities=(vx,vy,vz),
4441 particle_positions=(px,py,pz),
4442 lattice_displacements=lattice_displacements,
4443 gas_velocity_unit=gas_velocity_unit,
4444 particle_velocity_cms_factor=
4445 isfinite(physical_box_ckpc_h) ?
4446 _particle_velocity_cms_per_code(a, physical_box_ckpc_h) : NaN,
4447 physical_box_ckpc_h=physical_box_ckpc_h,
4448 gas_delta=gas_delta,
4449 HII=HII, H2I=H2I, HeII=HeII, fh=fh,
4450 temperature_velocity_unit2=temperature_velocity_unit2,
4451 neutral_hydrogen_storage=neutral_hydrogen_storage,
4452 centered_hydrogen_storage=centered_hydrogen_storage,
4453 hydrogen_reference=hydrogen_reference,
4454 photon_real_batch=photon_real_ready ?
4455 radiation_workspace.real_batch : nothing)
4456 end
4457
4458 if pdf_counts !== nothing && !isempty(pdf_path)
4459 _write_density_pdf!(pdf_path, "gas_rho", pdf_counts, s.U[1];
4460 be=be_p, cycle=cycle, measured=measured, a=a,
4461 loglo=pdf_loglo, loghi=pdf_loghi)
4462 if gravity
4463 dm_pdf_source = ρdm
4464 if lattice_displacements
4465 _density_from_delta_k!(be_p)(scratch, ρdm, Int(N);
4466 ndrange=ncells)
4467 KA.synchronize(be_p)
4468 dm_pdf_source = scratch
4469 end
4470 _write_density_pdf!(pdf_path, "dm_rho", pdf_counts, dm_pdf_source;
4471 be=be_p, cycle=cycle, measured=measured, a=a,
4472 loglo=pdf_loglo, loghi=pdf_loghi)
4473 end
4474 end
4475
4476 isempty(pk_path) && return
4477 isotropic_nmu = 1
4478 if gas_delta === nothing
4479 _delta_from_density_k!(be_p)(scratch,s.U[1],Int(N);ndrange=ncells)
4480 KA.synchronize(be_p)
4481 else
4482 copyto!(scratch,gas_delta)
4483 end
4484 _write_pkmu!(pk_path, "gas_delta",
4485 PoissonKernels.power_spectrum_aniso_gpu(scratch; boxsize=boxsize,
4486 nmu=isotropic_nmu,
4487 nbins=nbins, axis=axis,
4488 kmax=kmax);
4489 cycle=cycle, measured=measured, a=a, nmu=isotropic_nmu)
4490 if gravity
4491 if lattice_displacements
4492 copyto!(scratch, ρdm)
4493 else
4494 _delta_from_density_k!(be_p)(scratch, ρdm, Int(N); ndrange=ncells)
4495 KA.synchronize(be_p)
4496 end
4497 _write_pkmu!(pk_path, "dm_delta",
4498 PoissonKernels.power_spectrum_aniso_gpu(scratch; boxsize=boxsize,
4499 nmu=isotropic_nmu,
4500 nbins=nbins, axis=axis,
4501 kmax=kmax);
4502 cycle=cycle, measured=measured, a=a, nmu=isotropic_nmu)
4503 end
4504 Bx = reshape(s.U[6], N, N, N)
4505 By = reshape(s.U[7], N, N, N)
4506 Bz = reshape(s.U[8], N, N, N)
4507 # Magnetic damping is isotropic in these campaigns. One mu bin preserves
4508 # the shared P(k,mu) schema while avoiding unnecessary angular output.
4509 _write_pkmu!(pk_path, "B_vector",
4510 PoissonKernels.power_spectrum_aniso_gpu((Bx, By, Bz);
4511 boxsize=boxsize,
4512 nmu=isotropic_nmu,
4513 nbins=nbins, axis=axis,
4514 kmax=kmax);
4515 cycle=cycle, measured=measured, a=a, nmu=isotropic_nmu)
4516 pressure_force = ntuple(d -> reshape(s.scratch[d], N, N, N), 3)
4517 tension_force = ntuple(d -> reshape(s.scratch[d + 3], N, N, N), 3)
4518 lorentz_force = ntuple(d -> reshape(s.scratch[d + 6], N, N, N), 3)
4519 MHDKernels.terminal_magnetic_force_split!(
4520 s, pressure_force..., tension_force...)
4521 MHDKernels.terminal_magnetic_force!(s, lorentz_force...)
4522 KA.synchronize(s.be)
4523 for (label, force) in (("B_pressure_force", pressure_force),
4524 ("B_tension_force", tension_force),
4525 ("B_lorentz_force", lorentz_force))
4526 _write_pkmu!(pk_path, label,
4527 PoissonKernels.power_spectrum_aniso_gpu(force;
4528 boxsize=boxsize, nmu=isotropic_nmu, nbins=nbins,
4529 axis=axis, kmax=kmax);
4530 cycle=cycle, measured=measured, a=a, nmu=isotropic_nmu)
4531 end
4532 if radiation_workspace !== nothing && radiation_photon_state !== nothing
4533 photon_real_ready || MHDKernels.photon_fields_to_workspace!(
4534 radiation_workspace, radiation_photon_state)
4535 photon_delta=reshape(view(radiation_workspace.real_batch,:,:,:,1),:)
4536 copyto!(scratch,photon_delta)
4537 _write_pkmu!(pk_path,"photon_delta",
4538 PoissonKernels.power_spectrum_aniso_gpu(scratch;boxsize=boxsize,
4539 nmu=isotropic_nmu,nbins=nbins,axis=axis,kmax=kmax);
4540 cycle=cycle,measured=measured,a=a,nmu=isotropic_nmu)
4541
4542 photon_velocity=ntuple(d->reshape(s.scratch[d],N,N,N),3)
4543 for d in 1:3
4544 copyto!(photon_velocity[d],view(radiation_workspace.real_batch,:,:,:,d+1))
4545 end
4546 _write_pkmu!(pk_path,"photon_velocity",
4547 PoissonKernels.power_spectrum_aniso_gpu(photon_velocity;boxsize=boxsize,
4548 nmu=nmu,nbins=nbins,axis=axis,kmax=kmax);
4549 cycle=cycle,measured=measured,a=a,nmu=nmu)
4550
4551 _photon_slip_diag_k!(s.be)(s.scratch[1],s.scratch[2],s.scratch[3],
4552 s.scratch[1],s.scratch[2],s.scratch[3],
4553 s.U[1],s.U[2],s.U[3],s.U[4],s.smallr;ndrange=ncells)
4554 KA.synchronize(s.be)
4555 slip=ntuple(d->reshape(s.scratch[d],N,N,N),3)
4556 _write_pkmu!(pk_path,"photon_baryon_slip",
4557 PoissonKernels.power_spectrum_aniso_gpu(slip;boxsize=boxsize,
4558 nmu=nmu,nbins=nbins,axis=axis,kmax=kmax);
4559 cycle=cycle,measured=measured,a=a,nmu=nmu)
4560
4561 pressure_accel=ntuple(d->reshape(s.scratch[d+3],N,N,N),3)
4562 hydro_accel=ntuple(d->reshape(s.scratch[d+6],N,N,N),3)
4563 _gas_force_balance_diag_k!(s.be)(pressure_accel...,hydro_accel...,
4564 s.U[1],s.U[2],s.U[3],s.U[4],s.U[5],s.U[6],s.U[7],s.U[8],
4565 Int(N),T(0.5)/s.dx,s.γ,s.smallr,T(1e-30);ndrange=ncells)
4566 KA.synchronize(s.be)
4567 _write_pkmu!(pk_path,"gas_pressure_acceleration",
4568 PoissonKernels.power_spectrum_aniso_gpu(pressure_accel;boxsize=boxsize,
4569 nmu=isotropic_nmu,nbins=nbins,axis=axis,kmax=kmax);
4570 cycle=cycle,measured=measured,a=a,nmu=isotropic_nmu)
4571 _write_pkmu!(pk_path,"gas_hydro_acceleration",
4572 PoissonKernels.power_spectrum_aniso_gpu(hydro_accel;boxsize=boxsize,
4573 nmu=isotropic_nmu,nbins=nbins,axis=axis,kmax=kmax);
4574 cycle=cycle,measured=measured,a=a,nmu=isotropic_nmu)
4575 if isfinite(radiation_drag_rate)
4576 _gas_total_accel_diag_k!(s.be)(pressure_accel...,
4577 hydro_accel...,slip...,T(radiation_drag_rate);ndrange=ncells)
4578 KA.synchronize(s.be)
4579 _write_pkmu!(pk_path,"gas_total_acceleration",
4580 PoissonKernels.power_spectrum_aniso_gpu(pressure_accel;boxsize=boxsize,
4581 nmu=isotropic_nmu,nbins=nbins,axis=axis,kmax=kmax);
4582 cycle=cycle,measured=measured,a=a,nmu=isotropic_nmu)
4583 end
4584
4585 gas_velocity=ntuple(d->reshape(s.scratch[d],N,N,N),3)
4586 _gas_velocity_diag_k!(s.be)(gas_velocity...,
4587 s.U[1],s.U[2],s.U[3],s.U[4],s.smallr;ndrange=ncells)
4588 KA.synchronize(s.be)
4589 _write_pkmu!(pk_path,"gas_velocity",
4590 PoissonKernels.power_spectrum_aniso_gpu(gas_velocity;boxsize=boxsize,
4591 nmu=isotropic_nmu,nbins=nbins,axis=axis,kmax=kmax);
4592 cycle=cycle,measured=measured,a=a,nmu=isotropic_nmu)
4593 _mass_flux_divergence_diag_k!(s.be)(s.scratch[4],
4594 s.U[2],s.U[3],s.U[4],Int(N),T(0.5)/s.dx;ndrange=ncells)
4595 KA.synchronize(s.be)
4596 _write_pkmu!(pk_path,"gas_mass_flux_divergence",
4597 PoissonKernels.power_spectrum_aniso_gpu(reshape(s.scratch[4],N,N,N);
4598 boxsize=boxsize,nmu=isotropic_nmu,nbins=nbins,axis=axis,kmax=kmax);
4599 cycle=cycle,measured=measured,a=a,nmu=isotropic_nmu)
4600 end
4601 end
4602
4603 function main()
4604 N = parse(Int, get(ENV, "MHD_PMF_N", "128"))
4605 steps = parse(Int, get(ENV, "MHD_STEPS", "8"))
4606 tfinal_env = get(ENV, "MHD_TFINAL", "")
4607 tfinal = isempty(tfinal_env) ? nothing : parse(Float64, tfinal_env)
4608 zstart_env = get(ENV, "MHD_ZSTART", "")
4609 zrun = !isempty(zstart_env)
4610 # A benchmark warmup advances every state array but intentionally does not
4611 # advance cosmological time. It is therefore valid only for fixed-time
4612 # throughput tests, never as the default for a physical redshift trajectory.
4613 warmup = parse(Int, get(ENV, "MHD_WARMUP", zrun ? "0" : "1"))
4614 profile_skip_cycles = parse(Int, get(ENV, "MHD_PROFILE_SKIP_CYCLES", "0"))
4615 profile_skip_cycles >= 0 || error("MHD_PROFILE_SKIP_CYCLES must be non-negative")
4616 zstart = zrun ? parse(Float64, zstart_env) : NaN
4617 zend = zrun ? parse(Float64, get(ENV, "MHD_ZEND", "20")) : NaN
4618 hub = parse(Float64, get(ENV, "MHD_H", "0.674"))
4619 Om = parse(Float64, get(ENV, "MHD_OM", "0.315"))
4620 OL = parse(Float64, get(ENV, "MHD_OL", "0.685"))
4621 Or = parse(Float64, get(ENV, "MHD_OR", string(4.15e-5 / hub^2)))
4622 Ok = 1.0 - Om - OL - Or
4623 Ob = parse(Float64, get(ENV, "MHD_OB", string(0.17037 * Om)))
4624 maxexp = parse(Float64, get(ENV, "MHD_MAXEXP", "0.01"))
4625 # Source factor 1 is the corrected-clock science reference. Larger factors
4626 # are explicit performance experiments and must pass the cadence gate.
4627 source_max_dt_factor = parse(Float64, get(ENV, "MHD_SOURCE_MAX_DT_FACTOR", "1.0"))
4628 source_max_dt_factor >= 1 || error("MHD_SOURCE_MAX_DT_FACTOR must be >= 1")
4629 maxcycles = parse(Int, get(ENV, "MHD_MAX_CYCLES", "100000000"))
4630 debug_stop_cycle = parse(Int, get(ENV, "MHD_DEBUG_STOP_CYCLE", "0"))
4631 print_every = parse(Int, get(ENV, "MHD_PRINT_EVERY", "100"))
4632 cycle_print_every = parse(Int, get(ENV, "MHD_CYCLE_PRINT_EVERY", "1"))
4633 check_every = parse(Int, get(ENV, "MHD_CHECK_EVERY", zrun ? "25" : "0"))
4634 cfl = parse(Float64, get(ENV, "MHD_CFL", "0.4"))
4635 p0_env = lowercase(strip(get(ENV, "MHD_P0", "1.0")))
4636 p0_auto = p0_env in ("auto", "cosmology", "cmb")
4637 p0 = p0_auto ? NaN : parse(Float64, p0_env)
4638 pfloor = parse(Float64, get(ENV, "MHD_PFLOOR", "1e-7"))
4639 llf_dmin = parse(Float64, get(ENV, "MHD_LLF_DMIN", "1e-4"))
4640 llf_pmin = parse(Float64, get(ENV, "MHD_LLF_PMIN", "0"))
4641 drag_enabled = get(ENV, "MHD_COMPTON_DRAG", "0") in ("1", "true", "TRUE", "yes", "on")
4642 # Source-aware full MHD is the science default. The terminal closure did not
4643 # pass the density-response validation and must remain an explicit legacy
4644 # comparison mode rather than an automatically selected production regime.
4645 default_drag_dt_mode = drag_enabled && zrun ? "radiation" :
4646 (drag_enabled ? "godunov" : "explicit")
4647 drag_dt_mode = Symbol(lowercase(get(ENV, "MHD_DRAG_DT_MODE", default_drag_dt_mode)))
4648 drag_dt_mode in (:explicit, :boost, :subcycle, :semiimplicit, :aware,
4649 :godunov, :radiation, :terminal, :hybrid) ||
4650 error("MHD_DRAG_DT_MODE must be explicit, boost, subcycle, semiimplicit, aware, godunov, radiation, terminal, or hybrid; got $(drag_dt_mode)")
4651 drag_dt_mode in (:godunov, :radiation, :terminal, :hybrid) && !drag_enabled &&
4652 error("MHD_DRAG_DT_MODE=$(drag_dt_mode) requires MHD_COMPTON_DRAG=1")
4653 radiation_bridge_q2=parse(Float64,get(ENV,"MHD_RADIATION_BRIDGE_Q2","5.0"))
4654 radiation_long_viscosity=parse(Float64,get(ENV,"MHD_RADIATION_LONG_VISCOSITY",string(4/3)))
4655 radiation_max_dlna=parse(Float64,get(ENV,"MHD_RADIATION_MAX_DLNA","2.5e-4"))
4656 radiation_highz_max_dlna=parse(Float64,get(ENV,"MHD_RADIATION_HIGHZ_MAX_DLNA",
4657 string(radiation_max_dlna)))
4658 radiation_highz_zmin=parse(Float64,get(ENV,"MHD_RADIATION_HIGHZ_ZMIN","Inf"))
4659 radiation_nyquist_guard=parse(Float64,get(ENV,"MHD_RADIATION_NYQUIST_GUARD","1"))
4660 radiation_response_model=Symbol(lowercase(get(ENV,"MHD_RADIATION_RESPONSE_MODEL","moments")))
4661 radiation_nonlinear_default=radiation_response_model in (:angular,:transport) ? "0.4" : "Inf"
4662 radiation_nonlinear_cfl=parse(Float64,get(ENV,"MHD_RADIATION_NONLINEAR_CFL",
4663 radiation_nonlinear_default))
4664 radiation_energy_ledger_every=parse(Int,get(ENV,"MHD_RADIATION_ENERGY_LEDGER_EVERY","0"))
4665 radiation_free_stream_tolerance=parse(Float64,
4666 get(ENV,"MHD_RADIATION_FREE_STREAM_TOL","5e-3"))
4667 radiation_positivity_fallback=get(ENV,"MHD_RADIATION_POSITIVITY_FALLBACK","1") in
4668 ("1","true","TRUE","yes","on")
4669 radiation_positivity_max_fallback_passes=parse(Int,
4670 get(ENV,"MHD_RADIATION_POSITIVITY_MAX_FALLBACK_PASSES","8"))
4671 radiation_positivity_wave_growth_limit=parse(Float64,
4672 get(ENV,"MHD_RADIATION_POSITIVITY_WAVE_GROWTH_LIMIT","4"))
4673 radiation_positivity_max_subdivisions=parse(Int,
4674 get(ENV,"MHD_RADIATION_POSITIVITY_MAX_SUBDIVISIONS","8"))
4675 radiation_checkpoint_fallback=get(ENV,
4676 "MHD_RADIATION_CHECKPOINT_FALLBACK","0") in
4677 ("1","true","TRUE","yes","on")
4678 radiation_bridge_q2>0 || error("MHD_RADIATION_BRIDGE_Q2 must be positive")
4679 radiation_long_viscosity>0 || error("MHD_RADIATION_LONG_VISCOSITY must be positive")
4680 radiation_max_dlna>0 || error("MHD_RADIATION_MAX_DLNA must be positive")
4681 radiation_highz_max_dlna>0 || error("MHD_RADIATION_HIGHZ_MAX_DLNA must be positive")
4682 radiation_highz_zmin>=0 || error("MHD_RADIATION_HIGHZ_ZMIN must be non-negative")
4683 radiation_nonlinear_cfl>0 || error("MHD_RADIATION_NONLINEAR_CFL must be positive")
4684 0 < radiation_nyquist_guard <= 1 || error("MHD_RADIATION_NYQUIST_GUARD must be in (0,1]")
4685 radiation_response_model in (:bridge,:angular,:transport,:moments,:photon_moments,
4686 :moments_exponential,
4687 :photon_moments_exponential) ||
4688 error("MHD_RADIATION_RESPONSE_MODEL must be bridge, angular, moments, or " *
4689 "moments_exponential")
4690 radiation_energy_ledger_every>=0 ||
4691 error("MHD_RADIATION_ENERGY_LEDGER_EVERY must be non-negative")
4692 radiation_free_stream_tolerance>=0 ||
4693 error("MHD_RADIATION_FREE_STREAM_TOL must be non-negative")
4694 radiation_positivity_max_fallback_passes>0 ||
4695 error("MHD_RADIATION_POSITIVITY_MAX_FALLBACK_PASSES must be positive")
4696 radiation_positivity_wave_growth_limit>1 ||
4697 error("MHD_RADIATION_POSITIVITY_WAVE_GROWTH_LIMIT must exceed one")
4698 radiation_positivity_max_subdivisions>=0 ||
4699 error("MHD_RADIATION_POSITIVITY_MAX_SUBDIVISIONS must be non-negative")
4700 drag_dt_mode===:radiation && !zrun &&
4701 error("MHD_DRAG_DT_MODE=radiation requires a cosmological redshift run")
4702 allow_uncertified_terminal = get(ENV, "MHD_ALLOW_UNCERTIFIED_TERMINAL", "0") in
4703 ("1", "true", "TRUE", "yes", "on")
4704 drag_dt_mode in (:terminal, :hybrid) && !allow_uncertified_terminal &&
4705 error("MHD_DRAG_DT_MODE=$(drag_dt_mode) uses the uncertified terminal closure; " *
4706 "production runs must use godunov. Set MHD_ALLOW_UNCERTIFIED_TERMINAL=1 " *
4707 "only for legacy comparison experiments.")
4708 drag_dt_boost = parse(Float64, get(ENV, "MHD_DRAG_DT_BOOST", "1.0"))
4709 drag_dt_boost >= 1 || error("MHD_DRAG_DT_BOOST must be >= 1")
4710 drag_max_impulse = parse(Float64, get(ENV, "MHD_DRAG_MAX_IMPULSE", "Inf"))
4711 drag_max_impulse > 0 || error("MHD_DRAG_MAX_IMPULSE must be positive")
4712 # Gravity is a separate end-of-step kick. Bound its splitting error even
4713 # when the local drag propagator itself can take arbitrarily stiff steps.
4714 gravity_drag_max_impulse = parse(Float64,
4715 get(ENV,"MHD_GRAVITY_DRAG_MAX_IMPULSE","0.01"))
4716 isfinite(gravity_drag_max_impulse) && gravity_drag_max_impulse > 0 ||
4717 error("MHD_GRAVITY_DRAG_MAX_IMPULSE must be finite and positive")
4718 drag_xe_env = lowercase(get(ENV, "MHD_DRAG_XE", "auto"))
4719 drag_xe_fixed = drag_xe_env == "auto" ? NaN : parse(Float64, drag_xe_env)
4720 drag_helium_electrons = get(ENV, "MHD_DRAG_HELIUM_ELECTRONS", "1") in
4721 ("1", "true", "TRUE", "yes", "on")
4722 drag_rate_code = parse(Float64, get(ENV, "MHD_DRAG_RATE_CODE", "NaN"))
4723 (isnan(drag_rate_code) || drag_rate_code >= 0) ||
4724 error("MHD_DRAG_RATE_CODE must be non-negative or NaN")
4725 init_terminal_momenta = get(ENV, "MHD_INIT_TERMINAL_MOMENTA", "0") in
4726 ("1", "true", "TRUE", "yes", "on")
4727 init_terminal_momenta && !drag_enabled &&
4728 error("MHD_INIT_TERMINAL_MOMENTA=1 requires MHD_COMPTON_DRAG=1")
4729 init_terminal_momenta && !zrun && !isfinite(drag_rate_code) &&
4730 error("automatic terminal-momentum initialization requires MHD_ZSTART")
4731 drag_vx0 = parse(Float64, get(ENV, "MHD_DRAG_VX0", "0.0"))
4732 drag_vy0 = parse(Float64, get(ENV, "MHD_DRAG_VY0", "0.0"))
4733 drag_vz0 = parse(Float64, get(ENV, "MHD_DRAG_VZ0", "0.0"))
4734 terminal_vcap = parse(Float64, get(ENV, "MHD_TERMINAL_VCAP", "Inf"))
4735 terminal_gamma_scale = parse(Float64, get(ENV, "MHD_TERMINAL_GAMMA_SCALE", "1.0"))
4736 terminal_gamma_code_override = parse(Float64, get(ENV, "MHD_TERMINAL_GAMMA_CODE", "NaN"))
4737 terminal_split = get(ENV, "MHD_TERMINAL_SPLIT", "0") != "0"
4738 terminal_vsmooth = parse(Float64, get(ENV, "MHD_TERMINAL_VSMOOTH", "0.0"))
4739 0 <= terminal_vsmooth <= 1 || error("MHD_TERMINAL_VSMOOTH must be in [0,1]")
4740 terminal_dedicated_scratch = get(ENV, "MHD_TERMINAL_DEDICATED_SCRATCH", "0") in
4741 ("1", "true", "TRUE", "yes", "on")
4742 terminal_lbox_cm_env = get(ENV, "MHD_TERMINAL_LBOX_CM", "")
4743 terminal_lbox_ckpc_env = get(ENV, "MHD_TERMINAL_LBOX_CKPC", "")
4744 terminal_lbox_cm = !isempty(strip(terminal_lbox_cm_env)) ? parse(Float64, terminal_lbox_cm_env) :
4745 (!isempty(strip(terminal_lbox_ckpc_env)) ? parse(Float64, terminal_lbox_ckpc_env) * _PMF_KPC_CM : NaN)
4746 terminal_max_disp_cells = parse(Float64, get(ENV, "MHD_TERMINAL_MAX_DISP_CELLS", "0.1"))
4747 terminal_diff_cfl = parse(Float64, get(ENV, "MHD_TERMINAL_DIFF_CFL", "Inf"))
4748 terminal_diff_cfl > 0 || error("MHD_TERMINAL_DIFF_CFL must be positive")
4749 terminal_accuracy_cfl = parse(Float64, get(ENV, "MHD_TERMINAL_ACCURACY_CFL", "Inf"))
4750 terminal_accuracy_cfl > 0 || error("MHD_TERMINAL_ACCURACY_CFL must be positive")
4751 terminal_accuracy_pressure_env = lowercase(strip(
4752 get(ENV, "MHD_TERMINAL_ACCURACY_PRESSURE", "auto")))
4753 terminal_accuracy_pressure_auto = terminal_accuracy_pressure_env in
4754 ("auto", "initial", "p0")
4755 terminal_accuracy_pressure = terminal_accuracy_pressure_auto ? NaN :
4756 parse(Float64, terminal_accuracy_pressure_env)
4757 terminal_implicit_diffuse = get(ENV, "MHD_TERMINAL_IMPLICIT_DIFFUSE", "0") in
4758 ("1", "true", "TRUE", "yes", "on")
4759 terminal_implicit_diffuse_fac = parse(Float64, get(ENV, "MHD_TERMINAL_IMPLICIT_DIFFUSE_FAC", "1.0"))
4760 terminal_implicit_diffuse_fac >= 0 || error("MHD_TERMINAL_IMPLICIT_DIFFUSE_FAC must be non-negative")
4761 terminal_pressure_implicit = get(ENV, "MHD_TERMINAL_PRESSURE_IMPLICIT", "0") in
4762 ("1", "true", "TRUE", "yes", "on")
4763 terminal_pressure_coeff_env = lowercase(strip(get(ENV, "MHD_TERMINAL_PRESSURE_COEFF", "auto")))
4764 terminal_pressure_coeff_auto = terminal_pressure_coeff_env == "auto"
4765 terminal_pressure_coeff = terminal_pressure_coeff_auto ? NaN :
4766 parse(Float64, terminal_pressure_coeff_env)
4767 (terminal_pressure_coeff_auto || terminal_pressure_coeff >= 0) ||
4768 error("MHD_TERMINAL_PRESSURE_COEFF must be auto or non-negative")
4769 terminal_rho_rel_limit = parse(Float64, get(ENV, "MHD_TERMINAL_RHO_REL_LIMIT", "Inf"))
4770 terminal_rho_rel_limit > 0 || error("MHD_TERMINAL_RHO_REL_LIMIT must be positive")
4771 terminal_b_rel_limit = parse(Float64, get(ENV, "MHD_TERMINAL_B_REL_LIMIT", "Inf"))
4772 terminal_b_rel_limit > 0 || error("MHD_TERMINAL_B_REL_LIMIT must be positive")
4773 terminal_induction = Symbol(lowercase(get(ENV, "MHD_TERMINAL_INDUCTION", "centered")))
4774 terminal_induction in (:centered, :ssprk2, :ct) ||
4775 error("MHD_TERMINAL_INDUCTION must be centered, ssprk2, or ct")
4776 terminal_induction_cfl = parse(Float64, get(ENV, "MHD_TERMINAL_INDUCTION_CFL", "0.4"))
4777 0 < terminal_induction_cfl <= 1 || error("MHD_TERMINAL_INDUCTION_CFL must be in (0,1]")
4778 terminal_induction_dissipation = parse(Float64, get(ENV, "MHD_TERMINAL_INDUCTION_DISSIPATION", "0.05"))
4779 0 <= terminal_induction_dissipation <= 1 ||
4780 error("MHD_TERMINAL_INDUCTION_DISSIPATION must be in [0,1]")
4781 terminal_induction_dissipation_order = parse(Int, get(ENV, "MHD_TERMINAL_INDUCTION_DISSIPATION_ORDER", "4"))
4782 terminal_induction_dissipation_order in (2, 4) ||
4783 error("MHD_TERMINAL_INDUCTION_DISSIPATION_ORDER must be 2 or 4")
4784 terminal_project_b = get(ENV, "MHD_TERMINAL_PROJECT_B", terminal_induction === :ssprk2 ? "1" : "0") in
4785 ("1", "true", "TRUE", "yes", "on")
4786 terminal_induction_max_subcycles = parse(Int, get(ENV, "MHD_TERMINAL_INDUCTION_MAX_SUBCYCLES", "256"))
4787 terminal_induction_max_subcycles >= 1 ||
4788 error("MHD_TERMINAL_INDUCTION_MAX_SUBCYCLES must be positive")
4789 hybrid_min_va_cs = parse(Float64, get(ENV, "MHD_HYBRID_MIN_VA_CS", "0.0"))
4790 hybrid_min_va_cs >= 0 || error("MHD_HYBRID_MIN_VA_CS must be non-negative")
4791 hybrid_exit_drag_omega = parse(Float64,
4792 get(ENV, "MHD_HYBRID_EXIT_DRAG_OVER_OMEGA",
4793 get(ENV, "MHD_HYBRID_MIN_DRAG_OVER_OMEGA", "1.0")))
4794 hybrid_enter_drag_omega = parse(Float64,
4795 get(ENV, "MHD_HYBRID_ENTER_DRAG_OVER_OMEGA", string(2 * hybrid_exit_drag_omega)))
4796 0 < hybrid_exit_drag_omega < hybrid_enter_drag_omega ||
4797 error("hybrid drag hysteresis requires 0 < exit < enter")
4798 hybrid_check_every = parse(Int, get(ENV, "MHD_HYBRID_CHECK_EVERY", "1"))
4799 hybrid_check_every >= 1 || error("MHD_HYBRID_CHECK_EVERY must be positive")
4800 hybrid_confirm_checks = parse(Int, get(ENV, "MHD_HYBRID_CONFIRM_CHECKS", "2"))
4801 hybrid_confirm_checks >= 1 || error("MHD_HYBRID_CONFIRM_CHECKS must be positive")
4802 hybrid_handoff_steps = parse(Int, get(ENV, "MHD_HYBRID_HANDOFF_STEPS", "0"))
4803 hybrid_handoff_steps >= 0 || error("MHD_HYBRID_HANDOFF_STEPS must be non-negative")
4804 hybrid_handoff_alpha = parse(Float64, get(ENV, "MHD_HYBRID_HANDOFF_ALPHA", "0.0"))
4805 0 <= hybrid_handoff_alpha <= 1 || error("MHD_HYBRID_HANDOFF_ALPHA must be in [0,1]")
4806 hybrid_handoff_error_tol = parse(Float64, get(ENV, "MHD_HYBRID_HANDOFF_ERROR_TOL", "Inf"))
4807 hybrid_handoff_error_tol > 0 || error("MHD_HYBRID_HANDOFF_ERROR_TOL must be positive")
4808 hybrid_handoff_displacement_tol = parse(Float64,
4809 get(ENV, "MHD_HYBRID_HANDOFF_DISPLACEMENT_TOL", "0.02"))
4810 hybrid_handoff_displacement_tol > 0 ||
4811 error("MHD_HYBRID_HANDOFF_DISPLACEMENT_TOL must be positive")
4812 hybrid_handoff_diag_every = parse(Int, get(ENV, "MHD_HYBRID_HANDOFF_DIAG_EVERY", "8"))
4813 hybrid_handoff_diag_every >= 1 || error("MHD_HYBRID_HANDOFF_DIAG_EVERY must be positive")
4814 hybrid_reenter_terminal = get(ENV, "MHD_HYBRID_REENTER_TERMINAL", "0") in
4815 ("1", "true", "TRUE", "yes", "on")
4816 hybrid_aware_dt_factor = parse(Float64, get(ENV, "MHD_HYBRID_AWARE_DT_FACTOR", "1.0"))
4817 hybrid_aware_dt_factor > 0 || error("MHD_HYBRID_AWARE_DT_FACTOR must be positive")
4818 brms = parse(Float64, get(ENV, "MHD_PMF_BRMS", "1e-3"))
4819 pmf_b0_ng_env = get(ENV, "MHD_PMF_B0_NG", "")
4820 initial_state_dir = strip(get(ENV, "MHD_INITIAL_STATE_DIR", ""))
4821 pmf_normalization_z_env = strip(get(ENV, "MHD_PMF_NORMALIZATION_Z", ""))
4822 pmf_normalization_z = isempty(pmf_normalization_z_env) ? zstart :
4823 parse(Float64, pmf_normalization_z_env)
4824 if zrun && isempty(initial_state_dir) &&
4825 !isapprox(pmf_normalization_z, zstart; rtol=0,
4826 atol=8eps(Float64) * max(1, zstart))
4827 error("MHD_PMF_NORMALIZATION_Z must equal MHD_ZSTART unless MHD_INITIAL_STATE_DIR supplies a pre-evolved state")
4828 end
4829 pmf_init = Symbol(lowercase(get(ENV, "MHD_PMF_INIT", "batchelor")))
4830 radiation_magnetic_compensation = get(
4831 ENV,"MHD_RADIATION_MAGNETIC_COMPENSATION",
4832 pmf_init === :meanfield_linear ? "1" : "0") in
4833 ("1","true","TRUE","yes","on")
4834 pmf_kmode = parse(Int, get(ENV, "MHD_PMF_KMODE", "1"))
4835 meanfield_theta = parse(Float64, get(ENV, "MHD_MEANFIELD_THETA",
4836 string(π / 4)))
4837 0 <= meanfield_theta <= π / 2 ||
4838 error("MHD_MEANFIELD_THETA must be in [0,pi/2]")
4839 meanfield_polarization = Symbol(lowercase(
4840 get(ENV, "MHD_MEANFIELD_POLARIZATION", "magnetosonic")))
4841 meanfield_polarization_code =
4842 meanfield_polarization in (:magnetosonic, :compressive, :inplane) ? 1 :
4843 meanfield_polarization in (:alfven, :alfvenic, :outofplane) ? 2 :
4844 error("MHD_MEANFIELD_POLARIZATION must be magnetosonic or alfven")
4845 meanfield_pert_frac = parse(Float64,
4846 get(ENV, "MHD_MEANFIELD_PERT_FRAC", "1.0e-4"))
4847 meanfield_pert_frac > 0 || error("MHD_MEANFIELD_PERT_FRAC must be positive")
4848 pmf_single_mode_normalization =
4849 Symbol(lowercase(get(ENV, "MHD_PMF_SINGLE_MODE_NORMALIZATION", "amplitude")))
4850 pmf_single_mode_normalization in (:amplitude, :rms) ||
4851 error("MHD_PMF_SINGLE_MODE_NORMALIZATION must be amplitude or rms")
4852 boxsize = parse(Float64, get(ENV, "MHD_BOXSIZE", "1.0"))
4853 boxsize > 0 || error("MHD_BOXSIZE must be positive")
4854 pmf_tile_box_env = strip(get(ENV, "MHD_PMF_TILE_BOXSIZE", ""))
4855 pmf_tile_box = isempty(pmf_tile_box_env) ? nothing : parse(Float64, pmf_tile_box_env)
4856 pmf_tile_box === nothing || (pmf_tile_box > 0 && pmf_tile_box <= boxsize) ||
4857 error("MHD_PMF_TILE_BOXSIZE must be positive and no larger than MHD_BOXSIZE")
4858 pmf_tile_n_env = strip(get(ENV, "MHD_PMF_TILE_N", ""))
4859 pmf_tile_n = isempty(pmf_tile_n_env) ? nothing : parse(Int, pmf_tile_n_env)
4860 pmf_tile_n === nothing || pmf_tile_n > 0 || error("MHD_PMF_TILE_N must be positive")
4861 kcut = parse(Float64, get(ENV, "MHD_PMF_KCUT", string(0.625π*N / boxsize)))
4862 kmax_env = strip(get(ENV, "MHD_PMF_KMAX", ""))
4863 pmf_kmax = isempty(kmax_env) ? nothing : parse(Float64, kmax_env)
4864 pmf_kmax === nothing || pmf_kmax > 0 || error("MHD_PMF_KMAX must be positive")
4865 pmf_kmin = parse(Float64, get(ENV, "MHD_PMF_KMIN", "0"))
4866 pmf_kmin >= 0 || error("MHD_PMF_KMIN must be non-negative")
4867 pmf_kmax === nothing || pmf_kmin < pmf_kmax ||
4868 error("MHD_PMF_KMIN must be smaller than MHD_PMF_KMAX")
4869 pmf_required_cells_env = strip(get(ENV, "MHD_PMF_MIN_CELLS_PER_WAVELENGTH", ""))
4870 pmf_required_cells = isempty(pmf_required_cells_env) ? nothing :
4871 parse(Float64, pmf_required_cells_env)
4872 pmf_required_cells === nothing || pmf_required_cells > 0 ||
4873 error("MHD_PMF_MIN_CELLS_PER_WAVELENGTH must be positive")
4874 pmf_required_cells === nothing || pmf_kmax !== nothing ||
4875 error("MHD_PMF_KMAX is required when enforcing PMF wavelength resolution")
4876 preserve_low_kmax_env = strip(get(ENV, "MHD_PMF_PRESERVE_LOW_KMAX", ""))
4877 pmf_preserve_low_kmax = isempty(preserve_low_kmax_env) ? nothing : parse(Float64, preserve_low_kmax_env)
4878 pmf_preserve_low_kmax === nothing || pmf_preserve_low_kmax > 0 ||
4879 error("MHD_PMF_PRESERVE_LOW_KMAX must be positive")
4880 preserve_low_kcut_env = strip(get(ENV, "MHD_PMF_PRESERVE_LOW_KCUT", ""))
4881 pmf_preserve_low_kcut = isempty(preserve_low_kcut_env) ? nothing : parse(Float64, preserve_low_kcut_env)
4882 pmf_preserve_low_kcut === nothing || pmf_preserve_low_kcut > 0 ||
4883 error("MHD_PMF_PRESERVE_LOW_KCUT must be positive")
4884 pmf_normalization_kind = Symbol(lowercase(get(ENV, "MHD_PMF_NORMALIZATION_KIND", "auto")))
4885 pmf_normalization_kind = pmf_normalization_kind in (:total, :total_band) ?
4886 :total_band_rms : pmf_normalization_kind
4887 pmf_normalization_kind = pmf_normalization_kind in (:reference, :reference_band, :low_band) ?
4888 :reference_band_rms : pmf_normalization_kind
4889 pmf_normalization_kind in (:auto, :total_band_rms, :reference_band_rms) ||
4890 error("MHD_PMF_NORMALIZATION_KIND must be auto, total_band_rms, or reference_band_rms")
4891 pmf_normalization_kind === :reference_band_rms && pmf_preserve_low_kmax === nothing &&
4892 error("reference_band_rms requires MHD_PMF_PRESERVE_LOW_KMAX")
4893 pmf_normalization_kind_resolved = pmf_normalization_kind === :auto ?
4894 (pmf_preserve_low_kmax === nothing ? :total_band_rms : :reference_band_rms) :
4895 pmf_normalization_kind
4896 pmf_fixed_mode_power = get(ENV, "MHD_PMF_FIXED_MODE_POWER", "0") in
4897 ("1", "true", "TRUE", "yes", "on")
4898 pmf_mode_realization = Symbol(lowercase(get(ENV, "MHD_PMF_MODE_REALIZATION",
4899 "legacy")))
4900 pmf_mode_realization in (:legacy, :physical_hash, :physical_hash_zero_phase) ||
4901 error("MHD_PMF_MODE_REALIZATION must be legacy, physical_hash, or " *
4902 "physical_hash_zero_phase")
4903 realization_box_env = strip(get(ENV, "MHD_PMF_REALIZATION_BOX", ""))
4904 pmf_realization_box = isempty(realization_box_env) ? boxsize :
4905 parse(Float64, realization_box_env)
4906 pmf_realization_box > 0 || error("MHD_PMF_REALIZATION_BOX must be positive")
4907 pmf_curl_symbol = Symbol(lowercase(get(ENV, "MHD_PMF_CURL_SYMBOL", "centered")))
4908 pmf_curl_symbol in (:centered, :continuum, :spectral) ||
4909 error("MHD_PMF_CURL_SYMBOL must be centered or continuum")
4910 pmf_curl_symbol === :spectral && (pmf_curl_symbol = :continuum)
4911 curl_reference_env = strip(get(ENV, "MHD_PMF_CURL_REFERENCE_N", ""))
4912 pmf_curl_reference_n = isempty(curl_reference_env) ? nothing :
4913 parse(Int, curl_reference_env)
4914 if pmf_curl_reference_n !== nothing
4915 pmf_curl_symbol === :centered ||
4916 error("MHD_PMF_CURL_REFERENCE_N requires centered curl")
4917 0 < pmf_curl_reference_n <= N ||
4918 error("MHD_PMF_CURL_REFERENCE_N must be positive and no larger than MHD_PMF_N")
4919 end
4920 mode_lock_env = strip(get(ENV, "MHD_PMF_MODE_LOCK_N", ""))
4921 pmf_mode_lock_n = isempty(mode_lock_env) ? nothing : parse(Int, mode_lock_env)
4922 high_mode_lock_env = strip(get(ENV, "MHD_PMF_HIGH_MODE_LOCK_N", ""))
4923 pmf_high_mode_lock_n = isempty(high_mode_lock_env) ? nothing :
4924 parse(Int, high_mode_lock_env)
4925 pmf_shortest_k = pmf_kmax === nothing ? kcut : pmf_kmax
4926 pmf_min_cells_per_wavelength = 2π * N / (pmf_shortest_k * boxsize)
4927 if pmf_mode_lock_n !== nothing
4928 pmf_mode_lock_n <= N || error("MHD_PMF_MODE_LOCK_N cannot exceed MHD_PMF_N")
4929 lock_limit = π * pmf_mode_lock_n / boxsize
4930 if pmf_preserve_low_kmax === nothing
4931 pmf_shortest_k <= lock_limit ||
4932 error("fixed-band PMF exceeds the mode-lock Nyquist frequency")
4933 else
4934 pmf_preserve_low_kmax <= lock_limit ||
4935 error("preserved low-k PMF band exceeds the mode-lock Nyquist frequency")
4936 end
4937 end
4938 if pmf_high_mode_lock_n !== nothing
4939 pmf_high_mode_lock_n <= N ||
4940 error("MHD_PMF_HIGH_MODE_LOCK_N cannot exceed MHD_PMF_N")
4941 pmf_shortest_k <= π * pmf_high_mode_lock_n / boxsize ||
4942 error("PMF band exceeds the high-mode-lock Nyquist frequency")
4943 end
4944 seed = parse(Int, get(ENV, "MHD_PMF_SEED", "42"))
4945 recon = Symbol(lowercase(get(ENV, "MHD_RECON", "ppm")))
4946 debug_recon_override = strip(get(ENV, "MHD_DEBUG_RECON_OVERRIDE", ""))
4947 isempty(debug_recon_override) ||
4948 (recon = Symbol(lowercase(debug_recon_override)))
4949 riemann = Symbol(lowercase(get(ENV, "MHD_RIEMANN", "hlld")))
4950 debug_riemann_override = strip(get(ENV, "MHD_DEBUG_RIEMANN_OVERRIDE", ""))
4951 isempty(debug_riemann_override) ||
4952 (riemann = Symbol(lowercase(debug_riemann_override)))
4953 glm_ch_fac = parse(Float64, get(ENV, "MHD_GLM_CH_FAC", "2.0"))
4954 glm_cr = parse(Float64, get(ENV, "MHD_GLM_CR", "1.0"))
4955 mhd_solver = Symbol(lowercase(strip(get(ENV, "MHD_SOLVER", "glm"))))
4956 mhd_solver in (:glm, :ct) ||
4957 error("MHD_SOLVER must be glm or ct; got $mhd_solver")
4958 use_ct = mhd_solver === :ct
4959 if use_ct
4960 # Constrained transport keeps B on faces. Every path below that writes
4961 # cell-centred B during the evolution is unreachable for CT or rejected here.
4962 # Terminal/hybrid: every terminal B writer has a face-staggered CT counterpart; the
4963 # genuinely cell-wise options are rejected here.
4964 if drag_dt_mode in (:terminal, :hybrid)
4965 terminal_pressure_implicit && terminal_induction !== :ct && error(
4966 "MHD_SOLVER=ct with MHD_TERMINAL_PRESSURE_IMPLICIT=1 requires " *
4967 "MHD_TERMINAL_INDUCTION=ct (centered/ssprk2 advance cell-centred B)")
4968 terminal_induction === :ssprk2 && error(
4969 "MHD_SOLVER=ct does not support MHD_TERMINAL_INDUCTION=ssprk2 (cell-centred " *
4970 "SSPRK2 with a GLM-style B projection); use ct")
4971 end
4972 haskey(ENV, "MHD_GLM_PROJECT_EVERY") &&
4973 parse(Int, ENV["MHD_GLM_PROJECT_EVERY"]) > 0 && error(
4974 "MHD_SOLVER=ct does not use GLM B projection: MHD_GLM_PROJECT_EVERY=" *
4975 "$(ENV["MHD_GLM_PROJECT_EVERY"]) would overwrite cell-centred B; leave it unset " *
4976 "(CT default 0, including the radiation mode where GLM projects every cycle)")
4977 end
4978 glm_project_every = parse(Int, get(ENV, "MHD_GLM_PROJECT_EVERY",
4979 use_ct ? "0" : (drag_dt_mode === :radiation ? "1" : "0")))
4980 glm_project_every >= 0 || error("MHD_GLM_PROJECT_EVERY must be non-negative")
4981 use_ct && glm_project_every != 0 &&
4982 error("MHD_SOLVER=ct requires MHD_GLM_PROJECT_EVERY=0")
4983 fb = T(parse(Float64, get(ENV, "MHD_FB", "0.17037")))
4984 fdm = one(T) - fb
4985 gravity = parse(Bool, get(ENV, "MHD_GRAVITY", "true"))
4986 gravity_gas_kick = get(ENV, "MHD_GRAVITY_GAS_KICK", "1") in
4987 ("1", "true", "TRUE", "yes", "on")
4988 gravity_subcycles_fixed = parse(Int, get(ENV, "MHD_GRAVITY_SUBCYCLES", "1"))
4989 gravity_subcycles_fixed >= 1 || error("MHD_GRAVITY_SUBCYCLES must be >= 1")
4990 gravity_dt_boost = parse(Float64, get(ENV, "MHD_GRAVITY_DT_BOOST", "0.0"))
4991 (gravity_dt_boost == 0 || gravity_dt_boost >= 1) ||
4992 error("MHD_GRAVITY_DT_BOOST must be 0 or >= 1")
4993 gravity_source = Symbol(lowercase(get(ENV, "MHD_GRAVITY_SOURCE", "total")))
4994 gravity_source in (:total, :gas, :dm) ||
4995 error("MHD_GRAVITY_SOURCE must be total, gas, or dm")
4996 grav_gas_weight = gravity_source === :dm ? zero(T) : fb
4997 grav_dm_weight = gravity_source === :gas ? zero(T) : fdm
4998 particle_max_disp_cells = parse(Float64, get(ENV, "MHD_PARTICLE_MAX_DISP_CELLS", "Inf"))
4999 particle_max_disp_cells > 0 || error("MHD_PARTICLE_MAX_DISP_CELLS must be positive")
5000 particle_wrap = get(ENV, "MHD_PARTICLE_WRAP", "1") in ("1", "true", "TRUE", "yes", "on")
5001 lattice_displacements = get(ENV, "MHD_PARTICLE_LATTICE_DISPLACEMENTS", "1") in
5002 ("1", "true", "TRUE", "yes", "on")
5003 dm_density_mean = lattice_displacements ? zero(T) : one(T)
5004 gravity_phi_interp = get(ENV, "MHD_GRAVITY_PHI_INTERP", "0") in ("1", "true", "TRUE", "yes", "on")
5005 gravity_midpoint_source = get(ENV, "MHD_GRAVITY_MIDPOINT_SOURCE", "1") in
5006 ("1", "true", "TRUE", "yes", "on")
5007 gravity_phi_interp && gravity_midpoint_source &&
5008 error("choose either MHD_GRAVITY_PHI_INTERP=1 or MHD_GRAVITY_MIDPOINT_SOURCE=1")
5009 gravity_phi_interp && lattice_displacements &&
5010 error("potential interpolation is not implemented for lattice-displacement particles")
5011 fft_mode = _fft_mode(BACKEND_NAME, get(ENV, "MHD_FFT", "auto"))
5012 if BACKEND_NAME === :metal && fft_mode !== :mps &&
5013 !(get(ENV, "MHD_ALLOW_NON_MPS_FFT", "0") in ("1", "true", "TRUE", "yes", "on"))
5014 error("Metal PMF runs use the MPS rfft Poisson path by default; set MHD_ALLOW_NON_MPS_FFT=1 to benchmark MHD_FFT=$(fft_mode)")
5015 end
5016 gravity_1buf = get(ENV, "MHD_GRAVITY_1BUF",
5017 (BACKEND_NAME === :metal && fft_mode === :mps) ? "1" : "0") in
5018 ("1", "true", "TRUE", "yes", "on")
5019 # Cosmological PMF production uses the fused full mixing network.
5020 # Non-cosmological hydro/MHD benchmarks retain chemistry-off.
5021 chem_mode = _chem_mode(get(ENV, "MHD_CHEM", zrun ? "mixing" : "off"))
5022 chem_helium = get(ENV, "MHD_CHEM_HELIUM",
5023 chem_mode === :mixing ? "1" : "0") in
5024 ("1", "true", "TRUE", "yes", "on")
5025 chem_helium && chem_mode !== :mixing && error(
5026 "MHD_CHEM_HELIUM=1 currently requires the full MHD_CHEM=mixing network")
5027 chem_advect_species = chem_mode !== :off &&
5028 get(ENV, "MHD_CHEM_ADVECT_SPECIES", "1") in
5029 ("1", "true", "TRUE", "yes", "on")
5030 chem_coupling_default = zrun && chem_mode !== :off ? "strang" : "lie"
5031 chem_coupling = _chem_coupling_mode(get(ENV, "MHD_CHEM_COUPLING",
5032 chem_coupling_default))
5033 chem_expansion = _chem_expansion_enabled(get(ENV, "MHD_CHEM_EXPANSION", "auto"),
5034 zrun, chem_mode)
5035 fh = T(parse(Float64, get(ENV, "MHD_CHEM_FH", "0.76")))
5036 xhii_default = zrun && zstart > 1500 ? "1.0" : "2e-4"
5037 xhii0 = T(parse(Float64, get(ENV, "MHD_XHII_INIT", xhii_default)))
5038 xh20 = T(parse(Float64, get(ENV, "MHD_XH2_INIT", "1e-12")))
5039 chem_len_unit = parse(Float64, get(ENV, "MHD_CHEM_LENGTH_UNITS", "1.0"))
5040 chem_time_unit = parse(Float64, get(ENV, "MHD_CHEM_TIME_UNITS", "1.0"))
5041 chem_vunit_default = zrun ? "8e5" : string(chem_len_unit / chem_time_unit)
5042 chem_vunit = parse(Float64, get(ENV, "MHD_CHEM_VUNIT", chem_vunit_default))
5043 chem_time_unit = chem_len_unit / chem_vunit
5044 if p0_auto
5045 zrun || error("MHD_P0=auto requires MHD_ZSTART")
5046 p0 = _cosmological_ionized_p0(zstart, fh, chem_vunit, hub, Ob)
5047 end
5048 isfinite(p0) && p0 > 0 || error("MHD_P0 must be positive and finite")
5049 terminal_accuracy_pressure_auto && (terminal_accuracy_pressure = p0)
5050 terminal_accuracy_pressure > 0 ||
5051 error("MHD_TERMINAL_ACCURACY_PRESSURE must be positive")
5052 pmf_b0_ng = isempty(strip(pmf_b0_ng_env)) ? NaN : parse(Float64, pmf_b0_ng_env)
5053 pmf_comoving_clock = get(ENV, "MHD_PMF_COMOVING_CLOCK", "1") in
5054 ("1", "true", "TRUE", "yes", "on")
5055 pmf_b_zref = zrun ? pmf_normalization_z :
5056 parse(Float64, get(ENV, "MHD_PMF_B0_ZREF", "3000"))
5057 pmf_b_aref = isfinite(pmf_b0_ng) && pmf_comoving_clock ? z_to_a(pmf_b_zref) : NaN
5058 cosmological_mhd_scaling = Symbol(lowercase(strip(
5059 get(ENV, "MHD_COSMOLOGICAL_MHD_SCALING", "legacy"))))
5060 cosmological_mhd_scaling in (:legacy, :jedamzik) ||
5061 error("MHD_COSMOLOGICAL_MHD_SCALING must be legacy or jedamzik")
5062 cosmological_mhd_scaling === :jedamzik && !zrun &&
5063 error("MHD_COSMOLOGICAL_MHD_SCALING=jedamzik requires a redshift run")
5064 cosmological_mhd_scaling === :jedamzik && !isfinite(pmf_b_aref) &&
5065 error("MHD_COSMOLOGICAL_MHD_SCALING=jedamzik requires MHD_PMF_B0_NG and the comoving clock")
5066 cosmological_mhd_aref = cosmological_mhd_scaling === :jedamzik ? pmf_b_aref : NaN
5067 physical_box_ckpc_h = terminal_lbox_cm / _PMF_KPC_CM * hub
5068 streaming_velocity = ntuple(d->parse(Float64,get(ENV,
5069 ("MHD_STREAMING_VX_KMS","MHD_STREAMING_VY_KMS","MHD_STREAMING_VZ_KMS")[d],"0")),3)
5070 all(isfinite,streaming_velocity) || error("streaming velocities must be finite")
5071 streaming_active = any(!iszero,streaming_velocity)
5072 streaming_active && !(zrun && gravity && gravity_gas_kick &&
5073 cosmological_mhd_scaling === :jedamzik) && error(
5074 "coherent streaming requires cosmology, gravity, gas gravity, and jedamzik scaling")
5075 streaming_active && !isempty(initial_state_dir) && error(
5076 "streaming boost cannot be combined with an imported correlated initial state")
5077 invariant_axis = parse(Int,get(ENV,"MHD_PMF_INVARIANT_AXIS","0"))
5078 invariant_axis in 0:3 || error("MHD_PMF_INVARIANT_AXIS must be 0, 1, 2, or 3")
5079 invariant_axis != 0 && (pmf_init !== :batchelor || !isempty(initial_state_dir)) &&
5080 error("planar initialization requires built-in Batchelor ICs")
5081
5082 if isfinite(pmf_b0_ng)
5083 brms = _pmf_code_brms_from_comoving_ng(pmf_b0_ng, pmf_b_zref, hub, Ob, chem_vunit)
5084 end
5085 meanfield_bguide = parse(Float64,
5086 get(ENV, "MHD_MEANFIELD_BGUIDE", string(brms)))
5087 meanfield_bpert = parse(Float64, get(ENV, "MHD_MEANFIELD_BPERT",
5088 string(max(abs(meanfield_bguide), eps(Float64)) * meanfield_pert_frac)))
5089 meanfield_bguide > 0 || error("MHD_MEANFIELD_BGUIDE must be positive")
5090 meanfield_bpert > 0 || error("MHD_MEANFIELD_BPERT must be positive")
5091 alfven_bguide = parse(Float64, get(ENV, "MHD_ALFVEN_BGUIDE", string(brms)))
5092 alfven_bpert_frac = parse(Float64, get(ENV, "MHD_ALFVEN_BPERT_FRAC", "1.0e-3"))
5093 alfven_bpert = parse(Float64, get(ENV, "MHD_ALFVEN_BPERT",
5094 string(max(abs(alfven_bguide), eps(Float64)) * alfven_bpert_frac)))
5095 chem_density_fixed = get(ENV, "MHD_CHEM_DENSITY_UNITS", "")
5096 mixing_chem = chem_mode === :mixing || chem_mode === :analytic_mixing
5097 chem_hydrogen_storage = Symbol(lowercase(strip(get(
5098 ENV, "MHD_CHEM_H_STORAGE", zrun && chem_mode === :mixing ? "centered" : "ionized"))))
5099 chem_hydrogen_storage in (:ionized, :neutral, :centered) ||
5100 error("MHD_CHEM_H_STORAGE must be ionized, neutral, or centered")
5101 chem_neutral_hydrogen_storage = chem_hydrogen_storage === :neutral
5102 chem_centered_hydrogen_storage = chem_hydrogen_storage === :centered
5103 (chem_neutral_hydrogen_storage || chem_centered_hydrogen_storage) &&
5104 chem_mode !== :mixing && error(
5105 "MHD_CHEM_H_STORAGE=neutral/centered currently requires MHD_CHEM=mixing")
5106 # A full-state handoff may carry chemistry payloads even when an isolation
5107 # run deliberately disables chemistry. In that case the payload is not
5108 # imported, so its storage convention must not constrain the run.
5109 imported_hydrogen_reference = chem_mode === :off ? nothing :
5110 _initial_hydrogen_reference(initial_state_dir, chem_hydrogen_storage)
5111 imported_hydrogen_reference === nothing ||
5112 (xhii0 = T(imported_hydrogen_reference))
5113 chem_dtfrac = parse(Float64, get(ENV, "MHD_CHEM_DTFRAC",
5114 zrun && mixing_chem ? "1.0" : "0.1"))
5115 # Boosted terminal steps can span the entire stiff recombination interval.
5116 # The loop exits early after consuming dt, so a high ceiling is cheap once
5117 # converged; lower ceilings silently left ionized cells partially evolved.
5118 chem_itcap_default = chem_mode === :mixing ? "1024" :
5119 chem_mode === :analytic_mixing ? "64" :
5120 chem_mode === :peebles ? "64" : "64"
5121 chem_itcap = parse(Int, get(ENV, "MHD_CHEM_ITCAP", chem_itcap_default))
5122 chem_falpha = parse(Float64, get(ENV, "MHD_CHEM_FALPHA",
5123 zrun && mixing_chem ? "1.0" : "0.0"))
5124 chem_falpha_zs = _parse_optional_list("MHD_CHEM_FALPHA_ZS")
5125 chem_falpha_vals = _parse_optional_list("MHD_CHEM_FALPHA_VALS")
5126 chem_smooth = _chem_smooth_mode(get(ENV, "MHD_CHEM_SMOOTH",
5127 zrun && mixing_chem ? "gaussian" : "local"))
5128 chem_allow_local_mixing = get(ENV, "MHD_CHEM_ALLOW_LOCAL_MIXING", "0") in
5129 ("1", "true", "TRUE", "yes", "on")
5130 chem_nsmean = parse(Float64, get(ENV, "MHD_CHEM_NSMEAN", "1.0"))
5131 chem_smooth_sigma = parse(Float64, get(ENV, "MHD_CHEM_SMOOTH_SIGMA_CELLS", "16.0"))
5132 chem_nonlocal_table_path = strip(get(ENV, "MHD_CHEM_NONLOCAL_TABLE", ""))
5133 chem_nonlocal_scale = parse(Float64, get(ENV, "MHD_CHEM_NONLOCAL_SCALE", "1.0"))
5134 chem_nonlocal_scale >= 0 || error("MHD_CHEM_NONLOCAL_SCALE must be non-negative")
5135 !isempty(chem_nonlocal_table_path) && chem_nonlocal_scale == 0 && error(
5136 "remove MHD_CHEM_NONLOCAL_TABLE rather than configuring a zero response scale")
5137 chem_nonlocal_enabled = !isempty(chem_nonlocal_table_path) && chem_nonlocal_scale > 0
5138 chem_nonlocal_enabled && !zrun && error(
5139 "MHD_CHEM_NONLOCAL_TABLE is supported only for cosmological redshift runs")
5140 chem_nonlocal_enabled && chem_mode === :off && error(
5141 "MHD_CHEM_NONLOCAL_TABLE requires MHD_CHEM to be enabled")
5142 chem_nonlocal_enabled && chem_falpha != 0 && error(
5143 "nonlocal response tables require MHD_CHEM_FALPHA=0 to avoid double-counting mixing")
5144 chem_nonlocal_enabled && chem_smooth !== :local && error(
5145 "nonlocal response tables require MHD_CHEM_SMOOTH=local")
5146 chem_xe_mean = parse(Float64, get(ENV, "MHD_CHEM_XE_MEAN", "NaN"))
5147 chem_ionized_fast_zmin = parse(Float64, get(ENV, "MHD_CHEM_IONIZED_FAST_ZMIN",
5148 zrun && chem_helium ? "3000" : "0"))
5149 chem_helium && 0 < chem_ionized_fast_zmin < 3000 && error(
5150 "helium freeze-out must be evolved explicitly: set " *
5151 "MHD_CHEM_IONIZED_FAST_ZMIN=0 or >=3000")
5152 chem_recfast_hswitch = get(ENV, "MHD_CHEM_RECFAST_HSWITCH", "1") == "1"
5153 chem_rate_tables_enabled = get(ENV, "MHD_CHEM_RATE_TABLES", "1") in
5154 ("1", "true", "TRUE", "yes", "on")
5155 chem_profile_enabled = get(ENV, "MHD_CHEM_PROFILE", "0") in
5156 ("1", "true", "TRUE", "yes", "on")
5157 chem_fused_analytic = get(ENV, "MHD_CHEM_FUSED_ANALYTIC",
5158 chem_mode === :analytic_mixing ? "1" : "0") in
5159 ("1", "true", "TRUE", "yes", "on")
5160 chem_fused_mixing = get(ENV, "MHD_CHEM_FUSED_MIXING",
5161 chem_mode === :mixing ? "1" : "0") in
5162 ("1", "true", "TRUE", "yes", "on")
5163 (chem_neutral_hydrogen_storage || chem_centered_hydrogen_storage) &&
5164 !chem_fused_mixing && error(
5165 "MHD_CHEM_H_STORAGE=neutral/centered requires MHD_CHEM_FUSED_MIXING=1")
5166 chem_helium && !chem_fused_mixing && error(
5167 "MHD_CHEM_HELIUM=1 requires MHD_CHEM_FUSED_MIXING=1")
5168 chem_mean_every = parse(Int, get(ENV, "MHD_CHEM_MEAN_EVERY", "1"))
5169 chem_mean_every >= 1 || error("MHD_CHEM_MEAN_EVERY must be positive")
5170 chem_centered_hydrogen_storage && chem_mean_every != 1 && error(
5171 "MHD_CHEM_H_STORAGE=centered requires MHD_CHEM_MEAN_EVERY=1")
5172 chem_fudge = parse(Float64, get(ENV, "MHD_CHEM_RECFAST_FUDGE",
5173 chem_recfast_hswitch ? "1.125" : "1.0"))
5174 final_diag = get(ENV, "MHD_FINAL_DIAG", "1") != "0"
5175 history_path = get(ENV, "MHD_HISTORY_TSV", "")
5176 initial_history = get(ENV, "MHD_INITIAL_HISTORY", "1") in
5177 ("1", "true", "TRUE", "yes", "on")
5178 initial_pk = get(ENV, "MHD_INITIAL_PK", "0") in
5179 ("1", "true", "TRUE", "yes", "on")
5180 pk_path = get(ENV, "MHD_PK_TSV", "")
5181 mode_history_path = get(ENV, "MHD_MODE_HISTORY_TSV", "")
5182 if pmf_init in (:meanfield, :meanfield_linear, :sv25_mode) ||
5183 !isempty(mode_history_path)
5184 0 < pmf_kmode < N ÷ 2 ||
5185 error("matched linear modes require 0 < MHD_PMF_KMODE < N/2")
5186 end
5187 streaming_path = get(ENV,"MHD_STREAMING_TSV","")
5188 !isempty(streaming_path) && !(zrun && gravity && isfinite(physical_box_ckpc_h)) &&
5189 error("streaming diagnostics require a cosmological physical box and gravity")
5190 pdf_path = get(ENV, "MHD_PDF_TSV", "")
5191 adaptive_diag_rel = parse(Float64, get(ENV, "MHD_ADAPTIVE_DIAG_REL", "0"))
5192 adaptive_diag_rel >= 0 || error("MHD_ADAPTIVE_DIAG_REL must be non-negative")
5193 adaptive_diag_zmin = parse(Float64, get(ENV, "MHD_ADAPTIVE_DIAG_ZMIN", "-Inf"))
5194 adaptive_diag_zmax = parse(Float64, get(ENV, "MHD_ADAPTIVE_DIAG_ZMAX", "Inf"))
5195 adaptive_diag_zmax >= adaptive_diag_zmin ||
5196 error("MHD_ADAPTIVE_DIAG_ZMAX must be >= MHD_ADAPTIVE_DIAG_ZMIN")
5197 pdf_nbins = parse(Int, get(ENV, "MHD_PDF_NBINS", "128"))
5198 pdf_nbins > 0 || error("MHD_PDF_NBINS must be positive")
5199 pdf_loglo = parse(Float64, get(ENV, "MHD_PDF_LOGRHO_MIN", "-4"))
5200 pdf_loghi = parse(Float64, get(ENV, "MHD_PDF_LOGRHO_MAX", "4"))
5201 pdf_loghi > pdf_loglo || error("MHD_PDF_LOGRHO_MAX must exceed MHD_PDF_LOGRHO_MIN")
5202 pk_every = parse(Int, get(ENV, "MHD_PK_EVERY", "0"))
5203 pk_zs = _parse_z_outputs(get(ENV, "MHD_PK_ZS", ""))
5204 pk_next = Ref(1)
5205 diag_clear_fft = get(ENV, "MHD_DIAG_CLEAR_FFT", "0") in
5206 ("1", "true", "TRUE", "yes", "on")
5207 frame_dir = get(ENV, "MHD_DENSITY_FRAME_DIR", "")
5208 frame_zs = _parse_z_outputs(get(ENV, "MHD_DENSITY_FRAME_ZS", ""))
5209 !isempty(frame_zs) && isempty(frame_dir) &&
5210 error("MHD_DENSITY_FRAME_ZS requires MHD_DENSITY_FRAME_DIR")
5211 frame_field = Symbol(get(ENV,"MHD_DENSITY_FRAME_FIELD","gas"))
5212 frame_field in (:gas,:dm,:both) || error("frame field must be gas, dm, or both")
5213 frame_field != :gas && !gravity && error("DM frames require gravity")
5214 frame_next = Ref(1)
5215 pk_nmu = parse(Int, get(ENV, "MHD_PK_NMU", "4"))
5216 pk_nbins = parse(Int, get(ENV, "MHD_PK_NBINS", string(max(1, N ÷ 2))))
5217 pk_axis = parse(Int, get(ENV, "MHD_PK_AXIS", "1"))
5218 pk_kmax_env = strip(get(ENV, "MHD_PK_KMAX", ""))
5219 pk_kmax = isempty(pk_kmax_env) ? nothing : parse(Float64, pk_kmax_env)
5220 pk_kmax === nothing || pk_kmax > 0 || error("MHD_PK_KMAX must be positive")
5221 debug_div_every = parse(Int, get(ENV, "MHD_DEBUG_DIV_EVERY", "0"))
5222 debug_div_every >= 0 || error("MHD_DEBUG_DIV_EVERY must be non-negative")
5223 debug_phase_finite = get(ENV, "MHD_DEBUG_PHASE_FINITE", "0") in
5224 ("1", "true", "TRUE", "yes", "on")
5225 debug_phase_cycle_min = parse(Int, get(ENV, "MHD_DEBUG_PHASE_CYCLE_MIN", "0"))
5226 debug_phase_cycle_min >= 0 || error("MHD_DEBUG_PHASE_CYCLE_MIN must be non-negative")
5227 debug_phase_zmax = parse(Float64, get(ENV, "MHD_DEBUG_PHASE_ZMAX", "Inf"))
5228 debug_phase_rho_min = parse(Float64, get(ENV, "MHD_DEBUG_PHASE_RHO_MIN", "0"))
5229 debug_phase_rho_min >= 0 || error("MHD_DEBUG_PHASE_RHO_MIN must be non-negative")
5230 debug_radiation_rho_min = parse(Float64,
5231 get(ENV, "MHD_DEBUG_RADIATION_RHO_MIN", "0"))
5232 debug_radiation_cycle_min = parse(Int,
5233 get(ENV, "MHD_DEBUG_RADIATION_CYCLE_MIN", "0"))
5234 debug_radiation_correction_cfl = parse(Float64,
5235 get(ENV, "MHD_DEBUG_RADIATION_CORRECTION_CFL", "Inf"))
5236 debug_radiation_rho_min >= 0 ||
5237 error("MHD_DEBUG_RADIATION_RHO_MIN must be non-negative")
5238 debug_radiation_cycle_min >= 0 ||
5239 error("MHD_DEBUG_RADIATION_CYCLE_MIN must be non-negative")
5240 checkpoint_path = strip(get(ENV, "MHD_CHECKPOINT_PATH", ""))
5241 checkpoint_read_path = strip(get(ENV, "MHD_CHECKPOINT_READ_PATH", checkpoint_path))
5242 checkpoint_write_path = strip(get(ENV, "MHD_CHECKPOINT_WRITE_PATH", checkpoint_path))
5243 checkpoint_every = parse(Int, get(ENV, "MHD_CHECKPOINT_EVERY", "0"))
5244 checkpoint_at_cycle = parse(Int, get(ENV, "MHD_CHECKPOINT_AT_CYCLE", "0"))
5245 checkpoint_every >= 0 || error("MHD_CHECKPOINT_EVERY must be non-negative")
5246 checkpoint_at_cycle >= 0 || error("MHD_CHECKPOINT_AT_CYCLE must be non-negative")
5247 checkpoint_resume = get(ENV, "MHD_CHECKPOINT_RESUME", "1") in
5248 ("1", "true", "TRUE", "yes", "on")
5249 checkpoint_final = get(ENV, "MHD_CHECKPOINT_FINAL", "1") in
5250 ("1", "true", "TRUE", "yes", "on")
5251 checkpoint_disable_writes = get(ENV, "MHD_CHECKPOINT_DISABLE_WRITES", "0") in
5252 ("1", "true", "TRUE", "yes", "on")
5253 checkpoint_every > 0 && isempty(checkpoint_path) &&
5254 error("MHD_CHECKPOINT_EVERY requires MHD_CHECKPOINT_PATH")
5255 adaptive_last_delta_b = Ref(NaN)
5256 adaptive_last_xhii = Ref(NaN)
5257 if (chem_mode === :mixing || chem_mode === :analytic_mixing) &&
5258 chem_falpha <= 0 && !chem_nonlocal_enabled && !chem_allow_local_mixing
5259 chem_falpha = 1.0
5260 end
5261 if (chem_mode === :mixing || chem_mode === :analytic_mixing) &&
5262 chem_smooth === :local && !chem_nonlocal_enabled && !chem_allow_local_mixing
5263 error("MHD_CHEM=$(chem_mode) needs MHD_CHEM_SMOOTH=global or gaussian")
5264 end
5265
5266 be_mhd = MHDKernels.backend(BACKEND_NAME)
5267 be_p = PoissonKernels.backend(BACKEND_NAME)
5268 pdf_counts = isempty(pdf_path) ? nothing : PoissonKernels.device_zeros(be_p, Int32, (pdf_nbins,))
5269 frame_hostrho = isempty(frame_dir) ? nothing : Array{Float32}(undef, N, N, N)
5270 frame_hostslice = isempty(frame_dir) ? nothing : Array{Float32}(undef, N, N)
5271 frame_hostproj = isempty(frame_dir) ? nothing : Array{Float32}(undef, N, N)
5272 default_mhd_integrator = BACKEND_NAME === :cpu || use_ct ? "ref" : "cube"
5273 mhd_integrator = Symbol(lowercase(get(ENV, "MHD_INTEGRATOR",
5274 default_mhd_integrator)))
5275 mhd_integrator in (:ref, :cube) ||
5276 error("MHD_INTEGRATOR must be ref or cube; got $mhd_integrator")
5277 use_ct && mhd_integrator === :cube && error(
5278 "MHD_SOLVER=ct requires MHD_INTEGRATOR=ref (the default for CT on every backend); " *
5279 "the cube integrators are GLM-only")
5280 if BACKEND_NAME === :metal && mhd_integrator === :cube
5281 f16max = Float64(floatmax(Float16))
5282 p0 < f16max || error(
5283 "MHD_P0=$(p0) overflows the Metal cube's Float16 primitive staging; " *
5284 "choose physical code units such as MHD_CHEM_VUNIT=1e8")
5285 abs(brms) < f16max / 8 || error(
5286 "MHD Brms=$(brms) leaves insufficient Float16 staging headroom; " *
5287 "increase MHD_CHEM_VUNIT")
5288 end
5289 # CT runs ch = smax (so dt is not GLM-limited), GLM runs ch = glm_ch_fac*smax. The radiation
5290 # positivity wave-speed limit is `growth * ch`; for CT scale the growth by glm_ch_fac so the
5291 # ABSOLUTE limit (growth * glm_ch_fac * smax) is the same one GLM applies.
5292 radiation_positivity_wave_growth_eff = radiation_positivity_wave_growth_limit *
5293 (use_ct ? glm_ch_fac : 1.0)
5294 radiation_positivity_wave_growth_eff > 1 || error(
5295 "MHD_RADIATION_POSITIVITY_WAVE_GROWTH_LIMIT * MHD_GLM_CH_FAC must exceed one " *
5296 "for MHD_SOLVER=ct (got $(radiation_positivity_wave_growth_eff))")
5297 drag_dt_mode in (:godunov,:radiation) && mhd_integrator === :cube &&
5298 any(x -> !iszero(x), (drag_vx0, drag_vy0, drag_vz0)) &&
5299 error("Metal/CUDA source-aware Godunov drag is specialized for the CMB frame; set MHD_DRAG_VX0/VY0/VZ0=0")
5300 ncells = N^3
5301 use_ct && _log_line(@sprintf("CT_SOLVER mhd_solver=ct integrator=ref glm_ch_limited_dt=false glm_project_every=0 note=\"CT step is not GLM-limited: dt is the plain CFL dt (GLM scales it by smax/max(smax,ch_fac*smax)); expected difference vs GLM at equal CFL\""))
5302 _log_line(@sprintf("CONFIG mhd_solver=%s backend=%s N=%d ncells=%.3fM zrun=%s cosmology=(h=%.6g Om=%.6g OL=%.6g Or=%.6g) chem=%s chem_coupling=%s chem_expansion=%s rate_tables=%s gravity=%s gas_kick=%s source=%s phi_interp=%s midpoint_source=%s grav1buf=%s fft=%s particle_wrap=%s lattice_displacements=%s",
5303 String(mhd_solver), String(BACKEND_NAME), N, ncells/1e6,
5304 zrun ? @sprintf("%.1f->%.1f", zstart, zend) : "off",
5305 hub, Om, OL, Or, String(chem_mode), String(chem_coupling),
5306 string(chem_expansion),
5307 string(chem_rate_tables_enabled),
5308 string(gravity), string(gravity_gas_kick), String(gravity_source),
5309 string(gravity_phi_interp), string(gravity_midpoint_source),
5310 string(gravity_1buf), String(fft_mode),
5311 string(particle_wrap), string(lattice_displacements)))
5312 _log_line(@sprintf("PMF_MODES boxsize=%.9g kcut=%.9g kmin=%.9g kmax=%s hard_cutoff=%s mode_lock_n=%s high_mode_lock_n=%s curl_symbol=%s curl_reference_n=%s fixed_mode_power=%s mode_realization=%s realization_box=%.9g tile_box=%s tile_n=%s preserve_low_kmax=%s preserve_low_kcut=%s minimum_cells_per_wavelength=%.6g required_minimum_cells=%s",
5313 boxsize,
5314 kcut, pmf_kmin,
5315 pmf_kmax === nothing ? "none" : @sprintf("%.9g", pmf_kmax),
5316 string(pmf_kmax !== nothing),
5317 pmf_mode_lock_n === nothing ? "none" : string(pmf_mode_lock_n),
5318 pmf_high_mode_lock_n === nothing ? "none" : string(pmf_high_mode_lock_n),
5319 String(pmf_curl_symbol),
5320 pmf_curl_reference_n === nothing ? "none" : string(pmf_curl_reference_n),
5321 string(pmf_fixed_mode_power),
5322 String(pmf_mode_realization), pmf_realization_box,
5323 pmf_tile_box === nothing ? "none" : @sprintf("%.9g", pmf_tile_box),
5324 pmf_tile_n === nothing ? "auto" : string(pmf_tile_n),
5325 pmf_preserve_low_kmax === nothing ? "none" : @sprintf("%.9g", pmf_preserve_low_kmax),
5326 pmf_preserve_low_kcut === nothing ? "none" : @sprintf("%.9g", pmf_preserve_low_kcut),
5327 pmf_min_cells_per_wavelength,
5328 pmf_required_cells === nothing ? "none" : @sprintf("%.6g", pmf_required_cells)))
5329
5330 s = _stage("allocate_mhd_state") do
5331 MHDKernels.allocate_state(be_mhd, T, (N,N,N);
5332 dx=1/N, gamma=5/3, riemann=riemann, recon=recon,
5333 pfl=pfloor, llf_dmin=llf_dmin,
5334 llf_pmin=llf_pmin, mhd_solver=mhd_solver)
5335 end
5336 st = _stage("init_pmf_" * String(pmf_init)) do
5337 if pmf_init === :batchelor
5338 base = if pmf_tile_box === nothing
5339 MHDKernels.init_pmf_batchelor!(s; brms=brms, rho0=1, p0=p0,
5340 box=boxsize, seed=seed, kcut=kcut,
5341 kmin=pmf_kmin, kmax=pmf_kmax,
5342 mode_lock_n=pmf_mode_lock_n,
5343 high_mode_lock_n=pmf_high_mode_lock_n,
5344 preserve_low_kmax=pmf_preserve_low_kmax,
5345 preserve_low_kcut=pmf_preserve_low_kcut,
5346 normalization_kind=pmf_normalization_kind_resolved,
5347 curl_symbol=pmf_curl_symbol,
5348 curl_reference_n=pmf_curl_reference_n,
5349 fixed_mode_power=pmf_fixed_mode_power,
5350 mode_realization=pmf_mode_realization,
5351 realization_box=pmf_realization_box,
5352 minimum_cells_per_wavelength=pmf_required_cells,
5353 invariant_axis=invariant_axis)
5354 else
5355 MHDKernels.init_pmf_batchelor_tiled!(s; brms=brms, rho0=1, p0=p0,
5356 box=boxsize, tile_box=pmf_tile_box,
5357 tile_n=pmf_tile_n,
5358 seed=seed, kcut=kcut,
5359 kmin=pmf_kmin, kmax=pmf_kmax,
5360 mode_lock_n=pmf_mode_lock_n,
5361 high_mode_lock_n=pmf_high_mode_lock_n,
5362 preserve_low_kmax=pmf_preserve_low_kmax,
5363 preserve_low_kcut=pmf_preserve_low_kcut,
5364 normalization_kind=pmf_normalization_kind_resolved,
5365 curl_symbol=pmf_curl_symbol,
5366 curl_reference_n=pmf_curl_reference_n,
5367 fixed_mode_power=pmf_fixed_mode_power,
5368 mode_realization=pmf_mode_realization,
5369 realization_box=pmf_realization_box,
5370 minimum_cells_per_wavelength=pmf_required_cells,
5371 invariant_axis=invariant_axis)
5372 end
5373 merge(base, (btrans0=Float32(NaN), alfven_bguide=Float32(NaN),
5374 alfven_bpert=Float32(NaN)))
5375 elseif pmf_init in (:meanfield, :meanfield_linear, :sv25_mode)
5376 _init_meanfield_linear_k!(be_mhd)(s.U..., Int(N), T(s.γ), T(p0),
5377 T(meanfield_bguide), T(meanfield_bpert), T(meanfield_theta),
5378 meanfield_polarization_code, pmf_kmode; ndrange=ncells)
5379 KA.synchronize(be_mhd)
5380 measured_brms = _device_brms!(be_mhd, s, s.scratch[1])
5381 btrans0 = abs(meanfield_bpert) / sqrt(2)
5382 (brms_measured=Float32(measured_brms),
5383 total_brms_measured=Float32(measured_brms),
5384 reference_band_brms_measured=Float32(measured_brms),
5385 normalization_brms_measured=Float32(meanfield_bguide),
5386 normalization_kind="mean_field_amplitude", seed=seed, kcut=kcut,
5387 btrans0=Float32(btrans0), alfven_bguide=Float32(NaN),
5388 alfven_bpert=Float32(NaN))
5389 elseif pmf_init in (:alfven, :alfven_linear, :linear_alfven, :guide_alfven)
5390 _init_alfven_linear_k!(be_mhd)(s.U..., Int(N), T(s.γ), T(p0),
5391 T(alfven_bguide), T(alfven_bpert), pmf_kmode;
5392 ndrange=ncells)
5393 KA.synchronize(be_mhd)
5394 measured_brms = _device_brms!(be_mhd, s, s.scratch[1])
5395 btrans0 = _device_component_rms!(be_mhd, s.scratch[1], s.U[7])
5396 (brms_measured=Float32(measured_brms),
5397 total_brms_measured=Float32(measured_brms),
5398 reference_band_brms_measured=Float32(measured_brms),
5399 normalization_brms_measured=Float32(measured_brms),
5400 normalization_kind="total_band_rms", seed=seed, kcut=kcut,
5401 btrans0=Float32(btrans0), alfven_bguide=Float32(alfven_bguide),
5402 alfven_bpert=Float32(alfven_bpert))
5403 else
5404 mode_code = _single_b_mode(String(pmf_init))
5405 single_b_amplitude = mode_code == 2 && pmf_single_mode_normalization === :rms ?
5406 sqrt(2) * brms : brms
5407 _init_single_b_k!(be_mhd)(s.U..., Int(N), T(s.γ), T(p0),
5408 T(single_b_amplitude), mode_code, pmf_kmode;
5409 ndrange=ncells)
5410 KA.synchronize(be_mhd)
5411 measured_brms = _device_brms!(be_mhd, s, s.scratch[1])
5412 (brms_measured=Float32(measured_brms),
5413 total_brms_measured=Float32(measured_brms),
5414 reference_band_brms_measured=Float32(measured_brms),
5415 normalization_brms_measured=Float32(measured_brms),
5416 normalization_kind="total_band_rms", seed=seed, kcut=kcut,
5417 btrans0=Float32(NaN), alfven_bguide=Float32(NaN),
5418 alfven_bpert=Float32(NaN))
5419 end
5420 end
5421 import_rescale_to_pmf = get(ENV, "MHD_INITIAL_B_RESCALE_TO_PMF", "0") in
5422 ("1", "true", "TRUE", "yes", "on")
5423 initial_b_import_dir = strip(get(ENV, "MHD_INITIAL_B_IMPORT_DIR", ""))
5424 isempty(initial_state_dir) || isempty(initial_b_import_dir) ||
5425 error("MHD_INITIAL_STATE_DIR and MHD_INITIAL_B_IMPORT_DIR are mutually exclusive")
5426 imported_brms = _stage("import_initial_b") do
5427 _import_initial_b_field!(be_mhd, initial_b_import_dir, s;
5428 N=N, p0=p0,
5429 target_brms=import_rescale_to_pmf ? brms : nothing)
5430 end
5431 imported_brms === nothing || (st = merge(st, (
5432 brms_measured=imported_brms, total_brms_measured=imported_brms,
5433 reference_band_brms_measured=Float32(NaN),
5434 normalization_brms_measured=imported_brms,
5435 normalization_kind="imported_total_rms", btrans0=Float32(NaN))))
5436 imported_state_brms = _stage("import_initial_mhd_state") do
5437 _import_initial_mhd_state!(be_mhd, initial_state_dir, s; N=N)
5438 end
5439 imported_state_brms === nothing || (st = merge(st, (
5440 brms_measured=imported_state_brms, total_brms_measured=imported_state_brms,
5441 reference_band_brms_measured=Float32(NaN),
5442 normalization_brms_measured=imported_state_brms,
5443 normalization_kind="imported_joint_state", btrans0=Float32(NaN))))
5444 # Constrained transport evolves B on faces. After EVERY initial-condition and import stage
5445 # (and before any checkpoint restore) build the divergence-free face field once from the
5446 # final cell field. Thermal pressure is preserved by the kernel (E is adjusted), and the
5447 # normalization statistics reported by the run are re-measured from the resulting cell B.
5448 ct_init_change = NaN
5449 ct_init_energy_ratio = NaN
5450 ct_init_divb_face = NaN
5451 if use_ct
5452 ct_first = hasproperty(st, :ct_faces) ? st.ct_faces : nothing
5453 ct_again = _stage("ct_initial_faces") do
5454 MHDKernels.ct_faces_from_cells!(s)
5455 end
5456 # The PMF initializers already projected the generated field (stats in st.ct_faces);
5457 # their change is the one relative to the generated PMF. This second call covers the
5458 # other initial-condition and import stages and is idempotent for PMF fields.
5459 ct_report = ct_first === nothing ? ct_again : ct_first
5460 ct_init_change = Float64(ct_report.cell_b_rel_l2_change)
5461 ct_init_energy_ratio = Float64(ct_report.magnetic_energy_ratio)
5462 ct_init_divb_face = Float64(ct_again.max_divb_face_dx_over_brms)
5463 @printf("CT_INITIAL_FACES cell_b_rel_l2_change=%.6e magnetic_energy_ratio=%.9f divb_face_dx_over_brms=%.6e source=%s\n",
5464 ct_init_change, ct_init_energy_ratio, ct_init_divb_face,
5465 ct_first === nothing ? "runner" : "pmf_install")
5466 @printf("CT_INITIAL_FACES_REPROJECT cell_b_rel_l2_change=%.6e magnetic_energy_ratio=%.9f divb_face_dx_over_brms=%.6e (second call on the final cell field; ~0 change = idempotent)\n",
5467 ct_again.cell_b_rel_l2_change, ct_again.magnetic_energy_ratio,
5468 ct_again.max_divb_face_dx_over_brms)
5469 ct_brms_final = _device_brms!(be_mhd, s, s.scratch[1])
5470 @printf("CT_INITIAL_NORMALIZATION brms_before_faces=%.8e brms_after_faces=%.8e (brms_measured/total_brms_measured now report the post-projection cell field; reference_band and normalization_brms stay the pre-projection values)\n",
5471 Float64(st.total_brms_measured), ct_brms_final)
5472 flush(stdout)
5473 st = merge(st, (brms_measured_preprojection=st.brms_measured,
5474 total_brms_measured_preprojection=st.total_brms_measured,
5475 brms_measured=Float32(ct_brms_final),
5476 total_brms_measured=Float32(ct_brms_final)))
5477 end
5478 if imported_state_brms !== nothing
5479 p0 = _stage("measure_imported_pressure") do
5480 _device_pressure_mean!(be_mhd, s, s.scratch[1])
5481 end
5482 terminal_accuracy_pressure_auto && (terminal_accuracy_pressure = p0)
5483 @printf("INITIAL_STATE_PRESSURE p_mean=%.8e source=imported_energy\n", p0)
5484 flush(stdout)
5485 end
5486 pmf_ng_per_code = isfinite(pmf_b0_ng) ? pmf_b0_ng / brms : NaN
5487 pmf_initial_reference_comoving_ng =
5488 Float64(st.reference_band_brms_measured) * pmf_ng_per_code
5489 pmf_initial_total_comoving_ng = Float64(st.total_brms_measured) * pmf_ng_per_code
5490 _stage("export_initial_b") do
5491 _write_initial_b_field!(get(ENV, "MHD_INITIAL_B_DIR", ""), s, st;
5492 N=N, boxsize=boxsize, zstart=zstart, zend=zend,
5493 seed=seed, kcut=kcut, kmax=pmf_kmax,
5494 mode_lock_n=pmf_mode_lock_n,
5495 high_mode_lock_n=pmf_high_mode_lock_n, brms=brms,
5496 normalization_z=pmf_normalization_z,
5497 target_comoving_ng=pmf_b0_ng,
5498 reference_comoving_ng=pmf_initial_reference_comoving_ng,
5499 total_comoving_ng=pmf_initial_total_comoving_ng)
5500 end
5501
5502 px, py, pz, vx, vy, vz, ρdm, ρtot, φ, φprev = _stage("allocate_particle_gravity_arrays") do
5503 px = PoissonKernels.device_zeros(be_p, T, (ncells,))
5504 py = similar(px); pz = similar(px)
5505 vx = similar(px); vy = similar(px); vz = similar(px)
5506 ρdm = PoissonKernels.device_zeros(be_p, T, (ncells,))
5507 ρtot = PoissonKernels.device_zeros(be_p, T, (N,N,N))
5508 φ = gravity_1buf ? ρtot : similar(ρtot)
5509 φprev = gravity && gravity_source === :gas && gravity_phi_interp ? similar(ρtot) : nothing
5510 (px, py, pz, vx, vy, vz, ρdm, ρtot, φ, φprev)
5511 end
5512 eint, HII, H2I, HeII = _stage("allocate_chemistry_arrays") do
5513 eint = chem_mode === :off ? nothing : PoissonKernels.device_zeros(be_p, T, (ncells,))
5514 HII = chem_mode === :off ? nothing : PoissonKernels.device_zeros(be_p, T, (ncells,))
5515 H2I = chem_mode === :off ? nothing : PoissonKernels.device_zeros(be_p, T, (ncells,))
5516 HeII = chem_helium ? PoissonKernels.device_zeros(be_p, T, (ncells,)) : nothing
5517 (eint, HII, H2I, HeII)
5518 end
5519 chem_rate_tables = _stage("build_chemistry_rate_tables") do
5520 (chem_mode === :mixing || chem_mode === :analytic_mixing) && chem_rate_tables_enabled ?
5521 ChemistryKernels.build_rate_tables(; backend=BACKEND_NAME, precision=T) : nothing
5522 end
5523 chem_nonlocal_table = _stage("load_nonlocal_chemistry_table") do
5524 chem_nonlocal_enabled ? MHDKernels.load_nonlocal_recombination_table(
5525 chem_nonlocal_table_path, be_mhd; precision=T) : nothing
5526 end
5527 terminal_tvx, terminal_tvy, terminal_tvz = _stage("allocate_terminal_scratch") do
5528 if terminal_dedicated_scratch
5529 (PoissonKernels.device_zeros(be_p, T, (ncells,)),
5530 PoissonKernels.device_zeros(be_p, T, (ncells,)),
5531 PoissonKernels.device_zeros(be_p, T, (ncells,)))
5532 else
5533 (nothing, nothing, nothing)
5534 end
5535 end
5536
5537 initial_helium_state = chem_helium ?
5538 _equilibrium_helium_state(Float64(xhii0), zstart, hub, Ob, Float64(fh)) :
5539 (xe_total=Float64(xhii0), xheii=0.0, heii_mass_fraction=0.0)
5540 _stage("init_particles_and_chemistry") do
5541 if lattice_displacements
5542 MHDKernels.initialize_lattice_displacements!(px, py, pz, vx, vy, vz)
5543 else
5544 _init_lattice_k!(be_p)(px, py, pz, vx, vy, vz, Int(N); ndrange=ncells)
5545 end
5546 if chem_mode !== :off
5547 _init_reduced_chem_k!(be_p)(HII, H2I, s.U[1], xhii0, xh20, fh,
5548 T(xhii0), Int(N),
5549 Val(chem_neutral_hydrogen_storage),
5550 Val(chem_centered_hydrogen_storage);
5551 ndrange=ncells)
5552 chem_helium && _init_reduced_helium_k!(be_p)(
5553 HeII, s.U[1], T(initial_helium_state.heii_mass_fraction), Int(N);
5554 ndrange=ncells)
5555 end
5556 KA.synchronize(be_p)
5557 nothing
5558 end
5559 if init_terminal_momenta && isempty(initial_state_dir)
5560 _stage("init_terminal_momenta") do
5561 gamma_init = if isfinite(drag_rate_code)
5562 drag_rate_code
5563 else
5564 drag_xe_init = isfinite(drag_xe_fixed) ? drag_xe_fixed :
5565 _drag_electron_fraction(
5566 Float64(xhii0), initial_helium_state.xheii,
5567 zstart, hub, Ob, Float64(fh),
5568 drag_helium_electrons)
5569 gamma_drag_init = _compton_drag_rate_s(zstart, drag_xe_init, Float64(fh))
5570 terminal_gamma_scale *
5571 _terminal_gamma_code(gamma_drag_init, 1.0, 1.0, chem_vunit,
5572 terminal_lbox_cm, z_to_a(zstart), pmf_b_aref)
5573 end
5574 isfinite(terminal_gamma_code_override) &&
5575 (gamma_init = terminal_gamma_code_override)
5576 gamma_init > 0 || error("terminal-momentum initialization needs positive drag")
5577 _terminal_handoff_momenta!(be_mhd, s, gamma_init, terminal_vcap, 1.0)
5578 @printf("INITIAL_TERMINAL_MOMENTA gamma_code=%.8e mode=%s\n", gamma_init,
5579 isfinite(drag_rate_code) ? "fixed" : "physical")
5580 flush(stdout)
5581 nothing
5582 end
5583 end
5584
5585 terminal_rho_mean = Float64(sum(s.U[1])) / Float64(ncells)
5586 terminal_pressure_coeff_step = terminal_pressure_coeff_auto ?
5587 Float64(s.γ) * p0 / max(terminal_rho_mean, eps(Float64)) : terminal_pressure_coeff
5588 radiation_initial_xe = isfinite(drag_xe_fixed) ? drag_xe_fixed :
5589 _drag_electron_fraction(xhii0, initial_helium_state.xheii,
5590 zstart, hub, Ob, Float64(fh),
5591 drag_helium_electrons)
5592 radiation_initial_closure = drag_dt_mode === :radiation ?
5593 _radiation_closure(zstart, radiation_initial_xe, hub, Ob, fh,
5594 chem_vunit, terminal_lbox_cm,
5595 sqrt(max(terminal_pressure_coeff_step, 0.0));
5596 bridge_q2=radiation_bridge_q2,
5597 longitudinal_viscosity=radiation_long_viscosity,
5598 nyquist_guard_fraction=radiation_nyquist_guard,
5599 response_model=radiation_response_model) : nothing
5600 radiation_needs_fft = chem_nonlocal_enabled ||
5601 (drag_dt_mode === :radiation &&
5602 (radiation_response_model in (:moments, :photon_moments,
5603 :moments_exponential,
5604 :photon_moments_exponential) ||
5605 radiation_response_model !== :bridge ||
5606 radiation_initial_closure.mean_free_path^2 * (2pi)^2 <
5607 radiation_initial_closure.bridge_q2))
5608 radiation_workspace = radiation_needs_fft ?
5609 MHDKernels.allocate_radiation_workspace(s;
5610 compensate_magnetic=radiation_magnetic_compensation) : nothing
5611 radiation_workspace_bytes = radiation_needs_fft ?
5612 9*ncells*sizeof(T)+
5613 5*(N÷2+1)*N*N*sizeof(ComplexF32)+
5614 20*(3*(N÷2-1)^2+1)*sizeof(T) : 0
5615 radiation_photon_state = drag_dt_mode === :radiation && radiation_response_model in (
5616 :moments,:photon_moments,:moments_exponential,
5617 :photon_moments_exponential) ?
5618 MHDKernels.allocate_photon_moment_state(s) : nothing
5619 if streaming_active
5620 initial_velocity_unit = _cosmological_velocity_unit(
5621 chem_vunit,z_to_a(zstart),cosmological_mhd_aref)
5622 MHDKernels.boost_gas_velocity!(s,streaming_velocity .* (1e5/initial_velocity_unit))
5623 KA.synchronize(be_mhd)
5624 @printf("STREAMING_IC frame=DM_rest sign=vb-vc z=%.9g vbc_kms=(%.9g,%.9g,%.9g)\n",
5625 zstart,streaming_velocity...)
5626 end
5627 if radiation_photon_state !== nothing
5628 MHDKernels.initialize_photon_moments!(radiation_photon_state,
5629 radiation_workspace,s;
5630 adiabatic_density=parse(Bool,get(ENV,"MHD_RADIATION_PHOTON_ADIABATIC","true")),
5631 comoving_velocity=parse(Bool,get(ENV,"MHD_RADIATION_PHOTON_COMOVING","true")))
5632 end
5633 _stage("import_initial_aux_state") do
5634 _import_initial_aux_state!(initial_state_dir; N=N,
5635 particles=(px, py, pz, vx, vy, vz),
5636 chemistry=chem_mode === :off ? nothing :
5637 (chem_helium ? (eint, HII, H2I, HeII) : (eint, HII, H2I)),
5638 photons=radiation_photon_state,
5639 density_perturbation=radiation_workspace === nothing ? nothing :
5640 radiation_workspace.density_perturbation,
5641 magnetic_residual=radiation_workspace === nothing ? nothing :
5642 radiation_workspace.magnetic_residual,
5643 energy_residual=radiation_workspace === nothing ? nothing :
5644 radiation_workspace.energy_residual,
5645 neutral_hydrogen_storage=chem_neutral_hydrogen_storage,
5646 centered_hydrogen_storage=chem_centered_hydrogen_storage)
5647 radiation_workspace === nothing ||
5648 MHDKernels.reconstruct_radiation_density!(radiation_workspace,s)
5649 end
5650 radiation_photon_bytes = radiation_photon_state === nothing ? 0 :
5651 4*(N÷2+1)*N*N*sizeof(ComplexF32)
5652 radiation_energy_ledger = drag_dt_mode === :radiation &&
5653 radiation_energy_ledger_every > 0 ?
5654 MHDKernels.RadiationEnergyLedger() : nothing
5655
5656 phase = Dict(:cfl => 0.0, :mhd => 0.0, :radiation => 0.0,
5657 :glm_projection => 0.0,
5658 :drag => 0.0, :deposit => 0.0, :assemble => 0.0,
5659 :fft => 0.0, :gas_gravity => 0.0, :push => 0.0, :chem => 0.0)
5660 radiation_step_timing = MHDKernels.RadiationStepTiming()
5661 radiation_phase = Dict(:predictor => 0.0, :corrector => 0.0, :ledger => 0.0)
5662 chem_phase = Dict(:unpack => 0.0, :smooth => 0.0, :solve => 0.0,
5663 :nonlocal => 0.0,
5664 :clamp => 0.0, :mean => 0.0, :repack => 0.0)
5665 glm_projection_steps = 0
5666 terminal_phase = Dict(:force => 0.0, :pressure_fft => 0.0, :finalize => 0.0,
5667 :induction => 0.0, :projection => 0.0)
5668 terminal_measured = 0
5669 measured = 0
5670 measured_subcycles = 0
5671 last_dt = 0.0
5672 last_dt_explicit = 0.0
5673 last_drag_gammaH = NaN
5674 last_drag_f = NaN
5675 last_drag_xe = isfinite(drag_xe_fixed) ? drag_xe_fixed : Float64(xhii0)
5676 last_radiation_closure = radiation_initial_closure
5677 evolving_xe_mean = Float64(xhii0)
5678 evolving_xheii_mean = initial_helium_state.xheii
5679 last_hybrid_va_cs = NaN
5680 last_hybrid_drag_omega = NaN
5681 last_hybrid_regime = "n/a"
5682 was_terminal = false
5683 hybrid_terminal_active = false
5684 hybrid_controller = MHDKernels.DragRegimeController(
5685 exit_drag_over_omega=hybrid_exit_drag_omega,
5686 enter_drag_over_omega=hybrid_enter_drag_omega,
5687 min_va_over_cs=hybrid_min_va_cs,
5688 confirm_checks=hybrid_confirm_checks,
5689 reenter_terminal=hybrid_reenter_terminal)
5690 handoff_remaining = 0
5691 handoff_events = 0
5692 handoff_steps_done = 0
5693 last_handoff_remaining = 0
5694 last_handoff_alpha = 0.0
5695 last_handoff_error = NaN
5696 max_handoff_error = NaN
5697 max_handoff_settled_error = NaN
5698 last_handoff_delta_v_rms = NaN
5699 last_handoff_terminal_v_rms = NaN
5700 last_handoff_full_v_rms = NaN
5701 last_handoff_displacement = NaN
5702 max_handoff_displacement = 0.0
5703 cumulative_handoff_displacement = 0.0
5704 source_limited_steps = 0
5705 drag_impulse_limited_steps = 0
5706 max_gravity_fast_impulse = 0.0
5707 last_terminal_diffuse_coeff = 0.0
5708 last_terminal_vmax = NaN
5709 last_induction_nsub = 0
5710 induction_subcycles_total = 0
5711 last_nsub = 1
5712 subcycles_total = 0
5713 last_gravity_nsub = gravity ? 1 : 0
5714 gravity_subcycles_total = 0
5715 last_particle_vmax = NaN
5716 last_particle_gmax = NaN
5717 particle_accel_max_dt = Inf
5718 boosted_steps = 0
5719 elapsed_time = 0.0
5720 a = zrun ? z_to_a(zstart) : NaN
5721 a_end = zrun ? z_to_a(zend) : NaN
5722
5723 checkpoint_config = (
5724 velocity_transport_version=2,
5725 gravity_drag_max_impulse=gravity_drag_max_impulse,
5726 backend=BACKEND_NAME, precision=string(T), N=N, zstart=zstart, zend=zend,
5727 boxsize=boxsize, pmf_init=pmf_init, pmf_b0_ng=pmf_b0_ng,
5728 streaming_velocity=streaming_velocity, invariant_axis=invariant_axis,
5729 pmf_kmode=pmf_kmode,
5730 meanfield_theta=meanfield_theta,
5731 meanfield_polarization=meanfield_polarization,
5732 meanfield_bguide=meanfield_bguide,
5733 meanfield_bpert=meanfield_bpert,
5734 pmf_seed=seed, pmf_kcut=kcut, pmf_kmin=pmf_kmin, pmf_kmax=pmf_kmax,
5735 pmf_mode_lock_n=pmf_mode_lock_n,
5736 pmf_high_mode_lock_n=pmf_high_mode_lock_n, cfl=cfl,
5737 pmf_normalization_z=pmf_normalization_z,
5738 cosmological_mhd_scaling=cosmological_mhd_scaling,
5739 initial_state_imported=!isempty(initial_state_dir),
5740 pmf_curl_reference_n=pmf_curl_reference_n,
5741 recon=recon, riemann=riemann, llf_dmin=llf_dmin, llf_pmin=llf_pmin,
5742 mhd_integrator=mhd_integrator, drag_dt_mode=drag_dt_mode,
5743 radiation_response_model=radiation_response_model,
5744 radiation_nonlinear_cfl=radiation_nonlinear_cfl,
5745 radiation_free_stream_tolerance=radiation_free_stream_tolerance,
5746 radiation_max_dlna=radiation_max_dlna,
5747 radiation_highz_max_dlna=radiation_highz_max_dlna,
5748 radiation_magnetic_compensation=radiation_magnetic_compensation,
5749 chem_mode=chem_mode, chem_coupling=chem_coupling,
5750 chem_helium=chem_helium,
5751 chem_hydrogen_storage=chem_hydrogen_storage,
5752 chem_expansion=chem_expansion,
5753 chem_nonlocal_table=chem_nonlocal_table_path,
5754 chem_nonlocal_scale=chem_nonlocal_scale, gravity=gravity,
5755 gravity_source=gravity_source, gravity_gas_kick=gravity_gas_kick,
5756 gravity_midpoint_source=gravity_midpoint_source,
5757 lattice_displacements=lattice_displacements, fft_mode=fft_mode,
5758 mhd_solver=mhd_solver,
5759 )
5760 function _checkpoint_arrays()
5761 arrays = Pair{String,Any}[]
5762 # The integrator swaps U and scratch after every accepted step, so this
5763 # tuple must be resolved at checkpoint time rather than captured once.
5764 for i in eachindex(s.U)
5765 push!(arrays, "mhd_U_$i" => s.U[i])
5766 end
5767 for (name, array) in (("particle_px", px), ("particle_py", py),
5768 ("particle_pz", pz), ("particle_vx", vx),
5769 ("particle_vy", vy), ("particle_vz", vz))
5770 push!(arrays, name => array)
5771 end
5772 if chem_mode !== :off
5773 push!(arrays, "chem_eint" => eint)
5774 push!(arrays, "chem_HII" => HII)
5775 push!(arrays, "chem_H2I" => H2I)
5776 chem_helium && push!(arrays, "chem_HeII" => HeII)
5777 end
5778 if radiation_photon_state !== nothing
5779 push!(arrays, "photon_hat_batch" => radiation_photon_state.hat_batch)
5780 end
5781 if radiation_workspace !== nothing
5782 push!(arrays, "gas_delta" => radiation_workspace.density_perturbation)
5783 if radiation_workspace.magnetic_residual !== nothing
5784 for d in 1:3
5785 push!(arrays, "magnetic_residual_$d" =>
5786 radiation_workspace.magnetic_residual[d])
5787 end
5788 end
5789 push!(arrays,"energy_residual"=>radiation_workspace.energy_residual)
5790 end
5791 if use_ct
5792 # Faces are the evolved CT state; restore them exactly (never rebuild from cells).
5793 push!(arrays, "ct_bxf" => s.ct.bxf)
5794 push!(arrays, "ct_byf" => s.ct.byf)
5795 push!(arrays, "ct_bzf" => s.ct.bzf)
5796 # Kahan remainders of the compensated face updates: active exactly when the
5797 # radiation workspace compensates (GLM's magnetic_residual_$d analogue).
5798 if radiation_workspace !== nothing &&
5799 radiation_workspace.magnetic_residual !== nothing
5800 push!(arrays, "ct_residual_1" => s.ct.rbx)
5801 push!(arrays, "ct_residual_2" => s.ct.rby)
5802 push!(arrays, "ct_residual_3" => s.ct.rbz)
5803 end
5804 end
5805 return arrays
5806 end
5807
5808 restart_cycle = 0
5809 resumed_from_checkpoint = checkpoint_resume && !isempty(checkpoint_read_path) &&
5810 isfile(checkpoint_read_path)
5811 if resumed_from_checkpoint
5812 _log_line("START restore_checkpoint path=$checkpoint_read_path")
5813 restore_seconds = @elapsed begin
5814 checkpoint_allowed_changes = Symbol[]
5815 get(ENV, "MHD_CHECKPOINT_ALLOW_CFL_CHANGE", "0") in
5816 ("1", "true", "TRUE", "yes", "on") &&
5817 append!(checkpoint_allowed_changes,
5818 (:cfl, :radiation_nonlinear_cfl))
5819 get(ENV, "MHD_CHECKPOINT_ALLOW_NONLOCAL_TABLE_CHANGE", "0") in
5820 ("1", "true", "TRUE", "yes", "on") &&
5821 push!(checkpoint_allowed_changes, :chem_nonlocal_table)
5822 runtime = _restore_pmf_checkpoint!(
5823 checkpoint_read_path, checkpoint_config, _checkpoint_arrays();
5824 allowed_config_changes=Tuple(checkpoint_allowed_changes),
5825 allowed_config_validator=
5826 _checkpoint_nonlocal_table_change_is_safe,
5827 optional_missing_arrays=("magnetic_residual_1",
5828 "magnetic_residual_2",
5829 "magnetic_residual_3",
5830 "energy_residual",
5831 "ct_residual_1", "ct_residual_2",
5832 "ct_residual_3"))
5833 restart_cycle = runtime.cycle
5834 measured = runtime.measured
5835 measured_subcycles = runtime.measured_subcycles
5836 elapsed_time = runtime.elapsed_time
5837 a = runtime.a
5838 evolving_xe_mean = runtime.evolving_xe_mean
5839 evolving_xheii_mean = runtime.evolving_xheii_mean
5840 merge!(phase, runtime.phase)
5841 merge!(chem_phase, runtime.chem_phase)
5842 glm_projection_steps = runtime.glm_projection_steps
5843 source_limited_steps = runtime.source_limited_steps
5844 boosted_steps = runtime.boosted_steps
5845 subcycles_total = runtime.subcycles_total
5846 gravity_subcycles_total = runtime.gravity_subcycles_total
5847 induction_subcycles_total = runtime.induction_subcycles_total
5848 terminal_measured = runtime.terminal_measured
5849 merge!(terminal_phase, runtime.terminal_phase)
5850 pk_next[] = runtime.pk_next
5851 frame_next[] = runtime.frame_next
5852 adaptive_last_delta_b[] = runtime.adaptive_last_delta_b
5853 adaptive_last_xhii[] = runtime.adaptive_last_xhii
5854 last_radiation_closure = runtime.last_radiation_closure
5855 reset_energy_ledger = get(
5856 ENV, "MHD_CHECKPOINT_RESET_ENERGY_LEDGER", "0") in
5857 ("1", "true", "TRUE", "yes", "on")
5858 if radiation_energy_ledger !== nothing && runtime.radiation_energy !== nothing &&
5859 !reset_energy_ledger
5860 energy = runtime.radiation_energy
5861 radiation_energy_ledger.samples = energy.samples
5862 radiation_energy_ledger.total_change = energy.total_change
5863 radiation_energy_ledger.kinetic_change = energy.kinetic_change
5864 radiation_energy_ledger.internal_change = energy.internal_change
5865 radiation_energy_ledger.magnetic_change = energy.magnetic_change
5866 radiation_energy_ledger.photon_change = energy.photon_change
5867 radiation_energy_ledger.resolved_coupled_change = energy.resolved_coupled_change
5868 radiation_energy_ledger.component_residual = energy.component_residual
5869 radiation_energy_ledger.component_residual_abs = energy.component_residual_abs
5870 radiation_energy_ledger.component_magnitude_abs =
5871 hasproperty(energy, :component_magnitude_abs) ?
5872 energy.component_magnitude_abs : NaN
5873 elseif radiation_energy_ledger !== nothing && reset_energy_ledger
5874 @info "resetting radiation energy ledger after checkpoint restore"
5875 end
5876 KA.synchronize(be_mhd)
5877 KA.synchronize(be_p)
5878 end
5879 _log_line(@sprintf("DONE restore_checkpoint %.3f s cycle=%d z=%.6f",
5880 restore_seconds, restart_cycle,
5881 zrun ? a_to_z(a) : NaN))
5882 end
5883
5884 if φprev !== nothing
5885 _stage("init_previous_gas_potential") do
5886 if radiation_workspace === nothing
5887 _assemble_total_perturbation_k!(be_p)(
5888 ρtot, s.U[1], ρdm, grav_gas_weight, grav_dm_weight,
5889 one(T), dm_density_mean, Int(N); ndrange=ncells,
5890 )
5891 else
5892 _assemble_total_perturbation_k!(be_p)(
5893 ρtot, radiation_workspace.density_perturbation, ρdm,
5894 grav_gas_weight, grav_dm_weight, zero(T), dm_density_mean,
5895 Int(N); ndrange=ncells,
5896 )
5897 end
5898 gravcoef0 = zrun ? 1.5 * Om * a : 1.0
5899 if fft_mode === :mps
5900 PoissonKernels.fft_poisson_rfft!(φprev, ρtot; G=gravcoef0, a=1, boxsize=1)
5901 elseif fft_mode === :ka
5902 PoissonKernels.fft_poisson_rfft_ka!(φprev, ρtot; G=gravcoef0, a=1, boxsize=1)
5903 else
5904 PoissonKernels.fft_poisson_root_gpu!(φprev, ρtot; G=gravcoef0, a=1, boxsize=1)
5905 end
5906 KA.synchronize(be_p)
5907 nothing
5908 end
5909 end
5910
5911 cs_floor = sqrt(Float64(s.γ) * max(p0, pfloor))
5912 va_rms = Float64(st.brms_measured)
5913 terminal_keff = parse(Float64, get(ENV, "MHD_TERMINAL_KEFF",
5914 pmf_init === :batchelor ? string(max(1.0, kcut * boxsize / (2π))) :
5915 string(max(1, pmf_kmode))))
5916 hybrid_va_rms = va_rms
5917 if isfinite(pmf_b0_ng)
5918 @printf("PMF physical normalization: kind=%s epoch=initial_handoff z=%.3f target=%.6g comoving_nG reference_band=%.6g comoving_nG total_band=%.6g comoving_nG brms_code=%.6e vunit=%.6e cm/s\n",
5919 st.normalization_kind, pmf_normalization_z, pmf_b0_ng,
5920 pmf_initial_reference_comoving_ng, pmf_initial_total_comoving_ng,
5921 brms, chem_vunit)
5922 @printf("PMF comoving terminal clock: enabled=%s a_ref=%.9e\n",
5923 string(pmf_comoving_clock), z_to_a(pmf_b_zref))
5924 @printf("PMF cosmological MHD scaling: mode=%s velocity_aref=%.9e hubble_momentum_drag=%s\n",
5925 string(cosmological_mhd_scaling), cosmological_mhd_aref,
5926 string(cosmological_mhd_scaling === :jedamzik))
5927 flush(stdout)
5928 end
5929 if pmf_init in (:meanfield, :meanfield_linear, :sv25_mode)
5930 @printf("PMF matched mean-field mode: theta=%.9e polarization=%s guide_code=%.9e perturbation_code=%.9e fraction=%.9e\n",
5931 meanfield_theta, String(meanfield_polarization), meanfield_bguide,
5932 meanfield_bpert, meanfield_bpert / meanfield_bguide)
5933 flush(stdout)
5934 end
5935 p0_auto && @printf("PMF cosmological pressure: Tcmb(zstart)=%.6e K -> p0_code=%.9e\n",
5936 _PMF_T_CMB * (1 + zstart), p0)
5937 @printf("PMF+DM lattice bench: backend=%s N=%d steps=%d warmup=%d profile_skip=%d tfinal=%s zrun=%s source_max_dt_factor=%.3g recon=%s riemann=%s chfac=%.2f cr=%.2f init=%s kmode=%d brms=%.3e p0=%.3e pfloor=%.3e va_rms/cs_floor=%.3e kcut=%.3f kmax=%s mode_lock_n=%s gravity=%s gas_kick=%s grav_source=%s phi_interp=%s midpoint_source=%s grav1buf=%s grav_sub=%d grav_dt_boost=%.2f particle_max_disp=%.3g particle_wrap=%s lattice_displacements=%s fft=%s chem=%s chem_coupling=%s chem_expansion=%s smooth=%s falpha=%.3g xHII0=%.3e drag=%s drag_dt=%s boost=%.2f drag_xe=%s terminal_split=%s terminal_vsmooth=%.3f terminal_gamma_scale=%.3e terminal_vcap=%.3e terminal_lbox_ckpc=%.3e terminal_diff_cfl=%.3e terminal_accuracy_cfl=%.3e terminal_accuracy_pressure=%.3e terminal_implicit_diffuse=%s terminal_implicit_fac=%.3e terminal_pressure_implicit=%s terminal_pressure_coeff=%.3e terminal_pressure_coeff_mode=%s terminal_rho_rel_limit=%.3e terminal_b_rel_limit=%.3e hybrid_handoff_steps=%d hybrid_handoff_alpha_max=%.3f hybrid_handoff_ramp=cosine hybrid_handoff_error_tol=%.3e hybrid_handoff_displacement_tol=%.3e hybrid_reenter=%s hybrid_aware_dt_factor=%.3f\n",
5938 String(BACKEND_NAME), N, steps, warmup, profile_skip_cycles, tfinal === nothing ? "cycle-count" : string(tfinal),
5939 zrun ? @sprintf("%.1f->%.1f", zstart, zend) : "off", source_max_dt_factor,
5940 String(recon), String(riemann), glm_ch_fac, glm_cr, String(pmf_init), pmf_kmode, brms, p0, pfloor,
5941 va_rms / max(cs_floor, eps(Float64)), kcut,
5942 pmf_kmax === nothing ? "none" : @sprintf("%.6g", pmf_kmax),
5943 pmf_mode_lock_n === nothing ? "none" : string(pmf_mode_lock_n),
5944 string(gravity), string(gravity_gas_kick), String(gravity_source),
5945 string(gravity_phi_interp), string(gravity_midpoint_source),
5946 string(gravity_1buf), gravity_subcycles_fixed,
5947 gravity_dt_boost, particle_max_disp_cells, string(particle_wrap),
5948 string(lattice_displacements), String(fft_mode),
5949 String(chem_mode), String(chem_coupling), string(chem_expansion),
5950 String(chem_smooth), chem_falpha,
5951 Float64(xhii0), string(drag_enabled),
5952 String(drag_dt_mode), drag_dt_boost, drag_xe_env, string(terminal_split), terminal_vsmooth,
5953 terminal_gamma_scale, terminal_vcap,
5954 isfinite(terminal_lbox_cm) ? terminal_lbox_cm / _PMF_KPC_CM : NaN,
5955 terminal_diff_cfl, terminal_accuracy_cfl, terminal_accuracy_pressure,
5956 string(terminal_implicit_diffuse), terminal_implicit_diffuse_fac,
5957 string(terminal_pressure_implicit), terminal_pressure_coeff_step,
5958 terminal_pressure_coeff_auto ? "auto" : "fixed",
5959 terminal_rho_rel_limit, terminal_b_rel_limit,
5960 hybrid_handoff_steps, hybrid_handoff_alpha, hybrid_handoff_error_tol,
5961 hybrid_handoff_displacement_tol,
5962 string(hybrid_reenter_terminal),
5963 hybrid_aware_dt_factor)
5964 @printf(" mhd_solver=%s integrator=%s glm_ch_fac=%.6g glm_cr=%.6g glm_project_every=%d%s\n",
5965 String(mhd_solver), String(mhd_integrator), glm_ch_fac, glm_cr, glm_project_every,
5966 use_ct ? " (CT: ch_fac/cr unused, dt not GLM-limited, no B projection)" : "")
5967 @printf(" terminal induction=%s cfl=%.3f dissipation=%.3f/order%d project_b=%s max_subcycles=%d\n",
5968 string(terminal_induction), terminal_induction_cfl, terminal_induction_dissipation,
5969 terminal_induction_dissipation_order, string(terminal_project_b),
5970 terminal_induction_max_subcycles)
5971 @printf(" chemistry coupling=%s fused_analytic=%s fused_mixing=%s helium=%s hydrogen_storage=%s ionized_fast_zmin=%.6g mean_every=%d profile=%s rate_tables=%s itcap=%d dtfrac=%.3g\n",
5972 String(chem_coupling), string(chem_fused_analytic), string(chem_fused_mixing),
5973 string(chem_helium), String(chem_hydrogen_storage),
5974 chem_ionized_fast_zmin, chem_mean_every,
5975 string(chem_profile_enabled), string(chem_rate_tables_enabled),
5976 chem_itcap, chem_dtfrac)
5977 @printf(" chemistry species transport=%s scheme=bounded_fraction_backtrace scratch=reused_mhd\n",
5978 string(chem_advect_species))
5979 drag_dt_mode === :hybrid &&
5980 @printf(" hybrid hysteresis exit=%.3f enter=%.3f check_every=%d confirm=%d min_va/cs=%.3f\n",
5981 hybrid_exit_drag_omega, hybrid_enter_drag_omega, hybrid_check_every,
5982 hybrid_confirm_checks, hybrid_min_va_cs)
5983 drag_dt_mode === :radiation &&
5984 @printf(" radiation full-MHD response=%s bridge_q2=%.6g longitudinal_viscosity=%.6g max_dlna=%.6g highz_max_dlna=%.6g highz_zmin=%.6g nonlinear_cfl=%.6g free_stream_tol=%.6g magnetic_compensation=%s fft_active=%s workspace_bytes=%d photon_state_bytes=%d\n",
5985 string(radiation_response_model),
5986 radiation_bridge_q2,radiation_long_viscosity,radiation_max_dlna,
5987 radiation_highz_max_dlna,radiation_highz_zmin,
5988 radiation_nonlinear_cfl,radiation_free_stream_tolerance,
5989 string(radiation_magnetic_compensation),
5990 string(radiation_needs_fft),radiation_workspace_bytes,
5991 radiation_photon_bytes)
5992 chem_nonlocal_enabled &&
5993 @printf(" nonlocal Lyman-alpha table=%s scale=%.6g workspace=reused-radiation\n",
5994 chem_nonlocal_table_path, chem_nonlocal_scale)
5995 isfinite(drag_rate_code) &&
5996 @printf(" fixed dimensionless drag rate Gamma_code=%.6e\n", drag_rate_code)
5997 flush(stdout)
5998
5999 function _debug_div(label, cyc_now, a_now)
6000 debug_div_every == 0 && return nothing
6001 if cyc_now != 0 && cyc_now % debug_div_every != 0
6002 return nothing
6003 end
6004 br_dbg = _device_brms!(be_mhd, s, s.scratch[1])
6005 div_dbg = _device_divb_norm!(be_mhd, s, s.scratch[1], br_dbg)
6006 @printf(" divdbg cycle=%d phase=%s z=%.6f brms=%.8e divBdx_over_brms=%.8e\n",
6007 cyc_now, label, zrun ? a_to_z(a_now) : NaN, br_dbg, div_dbg)
6008 flush(stdout)
6009 return nothing
6010 end
6011 function _debug_phase_state(label, cyc_now, a_now, dt_now, smax_now)
6012 debug_phase_finite || return nothing
6013 cyc_now < debug_phase_cycle_min && return nothing
6014 z_now = zrun ? a_to_z(a_now) : NaN
6015 zrun && z_now > debug_phase_zmax && return nothing
6016 ss = _assert_finite_state(s, label, cyc_now, a_now, dt_now, smax_now)
6017 if debug_phase_rho_min > 0 && Float64(ss.ρmin) < debug_phase_rho_min
6018 error(@sprintf("%s cycle=%d density threshold crossed z=%.6f rho_min=%.8e threshold=%.8e rho_max=%.8e E=[%.8e,%.8e] Bmax=%.8e",
6019 label, cyc_now, z_now, Float64(ss.ρmin),
6020 debug_phase_rho_min, Float64(ss.ρmax),
6021 Float64(ss.Emin), Float64(ss.Emax), ss.Bmax))
6022 end
6023 return nothing
6024 end
6025 _debug_div("init", 0, zrun ? a : 1.0)
6026
6027 function _chem_step_interval!(a0_step, a1_step, dt_code_step, refresh_mean,
6028 measuring_now, xe_mean_in, xheii_mean_in,
6029 nonlocal_before::Bool=false)
6030 return _run_pmf_chemistry_step!(
6031 be_p, be_mhd, s, eint, HII, H2I, HeII, ρtot, φ;
6032 N=N, ncells=ncells, zrun=zrun,
6033 a0=Float64(a0_step), a1=Float64(a1_step), dt_code=Float64(dt_code_step),
6034 hub=hub, Om=Om, OL=OL, Or=Or, Ob=Ob, boxsize=boxsize,
6035 chem_mode=chem_mode, chem_smooth=chem_smooth,
6036 chem_density_fixed=chem_density_fixed,
6037 chem_vunit=chem_vunit, chem_time_unit=chem_time_unit,
6038 chem_ionized_fast_zmin=chem_ionized_fast_zmin,
6039 chem_fused_analytic=chem_fused_analytic,
6040 chem_fused_mixing=chem_fused_mixing,
6041 chem_helium=chem_helium,
6042 chem_neutral_hydrogen_storage=chem_neutral_hydrogen_storage,
6043 chem_centered_hydrogen_storage=chem_centered_hydrogen_storage,
6044 chem_expansion=chem_expansion,
6045 chem_profile_enabled=chem_profile_enabled, measuring=measuring_now,
6046 refresh_mean=refresh_mean, chem_phase=chem_phase,
6047 chem_falpha=chem_falpha, chem_falpha_zs=chem_falpha_zs,
6048 chem_falpha_vals=chem_falpha_vals, chem_nsmean=chem_nsmean,
6049 chem_smooth_sigma=chem_smooth_sigma, chem_xe_mean=chem_xe_mean,
6050 evolving_xe_mean=xe_mean_in,
6051 evolving_xheii_mean=xheii_mean_in,
6052 chem_recfast_hswitch=chem_recfast_hswitch, chem_fudge=chem_fudge,
6053 chem_rate_tables=chem_rate_tables, chem_dtfrac=chem_dtfrac,
6054 chem_itcap=chem_itcap, fh=fh,
6055 chem_nonlocal_table=chem_nonlocal_table,
6056 chem_nonlocal_scale=chem_nonlocal_scale,
6057 chem_nonlocal_before=nonlocal_before,
6058 radiation_workspace=radiation_workspace,
6059 terminal_lbox_cm=terminal_lbox_cm,
6060 cosmological_mhd_aref=cosmological_mhd_aref)
6061 end
6062
6063 function _save_checkpoint!(cycle_now)
6064 (isempty(checkpoint_write_path) || checkpoint_disable_writes) && return nothing
6065 energy = radiation_energy_ledger === nothing ? nothing : (
6066 samples=radiation_energy_ledger.samples,
6067 total_change=radiation_energy_ledger.total_change,
6068 kinetic_change=radiation_energy_ledger.kinetic_change,
6069 internal_change=radiation_energy_ledger.internal_change,
6070 magnetic_change=radiation_energy_ledger.magnetic_change,
6071 photon_change=radiation_energy_ledger.photon_change,
6072 resolved_coupled_change=radiation_energy_ledger.resolved_coupled_change,
6073 component_residual=radiation_energy_ledger.component_residual,
6074 component_residual_abs=radiation_energy_ledger.component_residual_abs,
6075 component_magnitude_abs=radiation_energy_ledger.component_magnitude_abs,
6076 )
6077 runtime = (
6078 cycle=cycle_now, measured=measured,
6079 measured_subcycles=measured_subcycles,
6080 elapsed_time=elapsed_time, a=a,
6081 evolving_xe_mean=evolving_xe_mean,
6082 evolving_xheii_mean=evolving_xheii_mean,
6083 phase=copy(phase), chem_phase=copy(chem_phase),
6084 glm_projection_steps=glm_projection_steps,
6085 source_limited_steps=source_limited_steps,
6086 boosted_steps=boosted_steps,
6087 subcycles_total=subcycles_total,
6088 gravity_subcycles_total=gravity_subcycles_total,
6089 induction_subcycles_total=induction_subcycles_total,
6090 terminal_measured=terminal_measured,
6091 terminal_phase=copy(terminal_phase),
6092 pk_next=pk_next[], frame_next=frame_next[],
6093 adaptive_last_delta_b=adaptive_last_delta_b[],
6094 adaptive_last_xhii=adaptive_last_xhii[],
6095 last_radiation_closure=last_radiation_closure,
6096 radiation_energy=energy,
6097 )
6098 KA.synchronize(be_mhd)
6099 KA.synchronize(be_p)
6100 bytes = _write_pmf_checkpoint(checkpoint_write_path, checkpoint_config,
6101 runtime, _checkpoint_arrays())
6102 @printf(" checkpoint cycle=%d z=%.6f bytes=%d path=%s\n",
6103 cycle_now, zrun ? a_to_z(a) : NaN, bytes, checkpoint_write_path)
6104 flush(stdout)
6105 return nothing
6106 end
6107
6108 # Establish the actual initialized state before any MHD/drag source step.
6109 # Batchelor ICs deliberately start with uniform gas density; recording this
6110 # row prevents the first post-step diagnostic from being mistaken for an IC.
6111 if !resumed_from_checkpoint &&
6112 ((initial_history && (!isempty(history_path) || !isempty(mode_history_path))) ||
6113 (initial_pk && !isempty(pk_path)))
6114 _emit_diagnostics!(be_p=be_p, s=s, ρdm=ρdm, scratch=ρtot,
6115 px=px, py=py, pz=pz, vx=vx, vy=vy, vz=vz,
6116 N=N, ncells=ncells, gravity=gravity, cycle=0,
6117 measured=0, a=zrun ? a : 1.0,
6118 pk_path=initial_pk ? pk_path : "",
6119 history_path=initial_history ? history_path : "", pdf_path="",
6120 pdf_counts=nothing, pdf_loglo=pdf_loglo, pdf_loghi=pdf_loghi,
6121 nmu=pk_nmu, nbins=pk_nbins, boxsize=boxsize, axis=pk_axis,
6122 kmax=pk_kmax,
6123 lattice_displacements=lattice_displacements,
6124 HII=HII, H2I=H2I, HeII=HeII, fh=fh,
6125 btrans0=Float64(st.btrans0), density_kmode=pmf_kmode,
6126 mode_history_path=initial_history ? mode_history_path : "",
6127 mode_kmode=pmf_kmode, hubble=hub, omega_b=Ob,
6128 physical_box_ckpc_h=physical_box_ckpc_h, streaming_path=streaming_path,
6129 gas_velocity_unit=_cosmological_velocity_unit(chem_vunit,a,cosmological_mhd_aref),
6130 radiation_workspace=radiation_workspace,
6131 radiation_photon_state=radiation_photon_state,
6132 neutral_hydrogen_storage=chem_neutral_hydrogen_storage,
6133 centered_hydrogen_storage=chem_centered_hydrogen_storage,
6134 hydrogen_reference=evolving_xe_mean,
6135 temperature_velocity_unit2=chem_mode === :off ? NaN :
6136 chem_vunit^2 * (isfinite(cosmological_mhd_aref) ?
6137 cosmological_mhd_aref / a : 1.0))
6138 @printf("INITIAL_DIAG emitted at z=%.6f before the first MHD/drag step (history=%s pk=%s)\n",
6139 zrun ? a_to_z(a) : NaN, string(initial_history), string(initial_pk))
6140 flush(stdout)
6141 end
6142
6143 if !resumed_from_checkpoint && zrun && frame_next[] <= length(frame_zs) &&
6144 _reached_output_z(a_to_z(a), frame_zs[frame_next[]])
6145 _write_run_density_frames!(frame_dir, frame_hostrho, frame_hostslice, frame_hostproj,
6146 s,ρdm,(px,py,pz),(vx,vy,vz),Int(N),a_to_z(a),0;
6147 field=frame_field,lattice_displacements=lattice_displacements)
6148 while frame_next[] <= length(frame_zs) &&
6149 _reached_output_z(a_to_z(a), frame_zs[frame_next[]])
6150 frame_next[] += 1
6151 end
6152 @printf("INITIAL_FRAME emitted at z=%.6f\n", a_to_z(a))
6153 flush(stdout)
6154 end
6155
6156 cyc = restart_cycle
6157 while true
6158 cyc += 1
6159 if debug_stop_cycle > 0 && cyc > debug_stop_cycle
6160 error("DEBUG stop cycle=$debug_stop_cycle")
6161 end
6162 dt0_ref = Ref(zero(T))
6163 smax_ref = Ref(zero(T))
6164 tcfl = _time_phase(be_mhd, be_p) do
6165 dt_tmp, smax_tmp = MHDKernels.compute_dt!(s.scratch[1], s; cfl=cfl)
6166 dt0_ref[] = dt_tmp
6167 smax_ref[] = smax_tmp
6168 nothing
6169 end
6170 dt0 = dt0_ref[]
6171 smax = smax_ref[]
6172 # CT has no GLM wave speed: ch is the plain signal speed (kept finite and nonzero for the
6173 # step interfaces that accept it), so dt_hydro_explicit == dt0 and nothing is GLM-limited.
6174 ch = use_ct ? T(smax) : T(glm_ch_fac) * T(smax)
6175 dt_hydro_explicit = Float64(dt0) * Float64(smax) / Float64(max(T(smax), ch))
6176 hydro_per_tau = zrun ? _hydro_time_per_tau(a, hub, chem_vunit,
6177 terminal_lbox_cm,
6178 cosmological_mhd_aref) : 1.0
6179 dt_explicit = dt_hydro_explicit / hydro_per_tau
6180 dt = dt_explicit
6181 if drag_enabled && drag_dt_mode in (:boost, :subcycle, :semiimplicit, :aware,
6182 :godunov, :radiation, :terminal, :hybrid) && drag_dt_boost > 1
6183 dt *= drag_dt_boost
6184 boosted_steps += 1
6185 end
6186 if zrun
6187 if !(isfinite(a) && a > 0 && a <= a_end)
6188 error(@sprintf("cycle=%d invalid scale factor a=%.8e z=%.6f", cyc, a, isfinite(a) ? a_to_z(a) : NaN))
6189 end
6190 if isapprox(a, a_end; rtol=0, atol=8eps(Float64) * max(1.0, a_end))
6191 a = a_end
6192 break
6193 end
6194 radiation_tight_now = drag_dt_mode === :radiation &&
6195 last_radiation_closure !== nothing &&
6196 last_radiation_closure.mean_free_path^2 * (2pi)^2 <
6197 last_radiation_closure.bridge_q2
6198 radiation_dlna_now = a_to_z(a) > radiation_highz_zmin ?
6199 radiation_highz_max_dlna : radiation_max_dlna
6200 maxexp_step = radiation_tight_now ? min(maxexp, radiation_dlna_now) : maxexp
6201 dt = min(dt, dtau_for_dlna(Om, OL, Ok, Or, a, maxexp_step))
6202 if radiation_tight_now && isfinite(radiation_nonlinear_cfl)
6203 # Linear photon-acoustic propagation is integrated by the
6204 # spectral propagator. Limit only the finite-volume MHD/GLM
6205 # displacement; using c_rad here restores the linear CFL and
6206 # makes the exact propagator unusable across recombination.
6207 nonlinear_signal = max(Float64(smax), Float64(ch))
6208 dt_radiation = radiation_nonlinear_cfl * Float64(s.dx) /
6209 (max(nonlinear_signal, eps(Float64)) * hydro_per_tau)
6210 dt = min(dt, dt_radiation)
6211 end
6212 # Land on diagnostic redshifts with the same RK2 map used by the
6213 # evolution. Post-crossing snapshots can otherwise differ in z by
6214 # O(dt), contaminating convergence tests during recombination.
6215 dt = _limit_dt_to_output_redshift(
6216 dt, Om, OL, Ok, Or, a, pk_zs, pk_next[])
6217 dt = _limit_dt_to_output_redshift(
6218 dt, Om, OL, Ok, Or, a, frame_zs, frame_next[])
6219 dt_end = _dtau_to_scale_factor_rk2(Om, OL, Ok, Or, a, a_end)
6220 if dt_end <= 8eps(Float64) * max(1.0, abs(dt))
6221 a = a_end
6222 break
6223 end
6224 dt = min(dt, dt_end)
6225 end
6226 if tfinal !== nothing && cyc > warmup && elapsed_time + dt > tfinal
6227 dt = tfinal - elapsed_time
6228 end
6229 source_dt = _limit_source_dt(dt, dt_explicit, source_max_dt_factor)
6230 dt = source_dt.dt
6231 source_limited_steps += source_dt.limited
6232 _assert_finite_step(s, "pre-step", cyc, zrun ? a : NaN, dt, smax)
6233 last_dt = dt
6234 last_dt_explicit = dt_explicit
6235 measuring = cyc > warmup + profile_skip_cycles
6236 nsub = 1
6237 if drag_enabled && drag_dt_mode in (:subcycle, :semiimplicit, :aware,
6238 :godunov, :radiation, :hybrid) && dt > dt_explicit
6239 nsub = max(1, ceil(Int, dt / max(dt_explicit, eps(Float64))))
6240 end
6241 last_nsub = nsub
6242
6243 a_before = zrun ? a : 1.0
6244 a_after_step = zrun ? _advance_a_rk2(Om, OL, Ok, Or, a, dt, a_end) : 1.0
6245 zsource = zrun ? a_to_z(0.5 * (a_before + a_after_step)) : 0.0
6246 dtphys_source = zrun ? _cosmic_dt_seconds(a_before, a_after_step, hub, Om, OL, Or) :
6247 dt * chem_time_unit
6248 dt_hydro = zrun ? _hydro_code_dt(dtphys_source, dt, chem_vunit,
6249 terminal_lbox_cm,
6250 0.5 * (a_before + a_after_step),
6251 cosmological_mhd_aref) : dt
6252 tdrag = 0.0
6253 tm = 0.0
6254 trad = 0.0
6255 tgp = 0.0
6256 radiation_timing_before = (
6257 predictor=radiation_step_timing.predictor_s,
6258 corrector=radiation_step_timing.corrector_s,
6259 ledger=radiation_step_timing.ledger_s,
6260 )
6261 terminal_timed = false
6262 terminal_step_timing = (force_s=0.0, pressure_fft_s=0.0, finalize_s=0.0,
6263 induction_s=0.0, projection_s=0.0)
6264 gamma_drag = 0.0
6265 drag_impulse = 0.0
6266 if drag_enabled
6267 drag_xe = isfinite(drag_xe_fixed) ? drag_xe_fixed :
6268 _drag_electron_fraction(evolving_xe_mean, evolving_xheii_mean,
6269 zsource, hub, Ob,
6270 Float64(fh), drag_helium_electrons)
6271 last_drag_xe = drag_xe
6272 gamma_drag = _compton_drag_rate_s(zsource, drag_xe, Float64(fh))
6273 last_drag_gammaH = isfinite(drag_rate_code) ? NaN :
6274 zrun ? gamma_drag / _hubble_s(zsource, hub, Om, OL, Or) : NaN
6275 drag_impulse = _drag_impulse(gamma_drag, dtphys_source, dt, drag_rate_code)
6276 last_drag_f = exp(-drag_impulse)
6277 end
6278 # For evolved photons the fast slip eigenvalue is (1+R)*Gamma, not
6279 # Gamma alone. Use that stiffness in the split-gravity accuracy cap.
6280 gravity_slip_factor = radiation_photon_state !== nothing ?
6281 1 + _photon_baryon_inertia_ratio(a_to_z(a_after_step),hub,Ob) : 1.0
6282 source_impulse_cap = _source_drag_impulse_cap(drag_dt_mode,
6283 drag_max_impulse,gravity && gravity_gas_kick,
6284 gravity_drag_max_impulse,gravity_slip_factor)
6285 if drag_enabled && isfinite(source_impulse_cap) &&
6286 drag_impulse > source_impulse_cap
6287 drag_impulse_limited_steps += 1
6288 for _ in 1:12
6289 drag_impulse <= source_impulse_cap * (1 + 1e-6) && break
6290 dt *= 0.99 * source_impulse_cap / drag_impulse
6291 last_dt = dt
6292 a_after_step = zrun ? _advance_a_rk2(Om, OL, Ok, Or, a, dt, a_end) : 1.0
6293 zsource = zrun ? a_to_z(0.5 * (a_before + a_after_step)) : 0.0
6294 dtphys_source = zrun ?
6295 _cosmic_dt_seconds(a_before, a_after_step, hub, Om, OL, Or) :
6296 dt * chem_time_unit
6297 dt_hydro = zrun ?
6298 _hydro_code_dt(dtphys_source, dt, chem_vunit, terminal_lbox_cm,
6299 0.5 * (a_before + a_after_step),
6300 cosmological_mhd_aref) : dt
6301 drag_xe = isfinite(drag_xe_fixed) ? drag_xe_fixed :
6302 _drag_electron_fraction(evolving_xe_mean, evolving_xheii_mean,
6303 zsource, hub, Ob,
6304 Float64(fh), drag_helium_electrons)
6305 last_drag_xe = drag_xe
6306 gamma_drag = _compton_drag_rate_s(zsource, drag_xe, Float64(fh))
6307 last_drag_gammaH = isfinite(drag_rate_code) ? NaN :
6308 zrun ? gamma_drag / _hubble_s(zsource, hub, Om, OL, Or) : NaN
6309 drag_impulse = _drag_impulse(gamma_drag, dtphys_source, dt, drag_rate_code)
6310 last_drag_f = exp(-drag_impulse)
6311 end
6312 drag_impulse <= source_impulse_cap * (1 + 1e-6) || error(
6313 "failed to resolve the gravity/drag source impulse")
6314 nsub = max(1, ceil(Int, dt / max(dt_explicit, eps(Float64))))
6315 last_nsub = nsub
6316 end
6317 gamma_code = drag_enabled ?
6318 (isfinite(drag_rate_code) ? drag_rate_code : terminal_gamma_scale *
6319 _terminal_gamma_code(gamma_drag, dtphys_source, dt, chem_vunit,
6320 terminal_lbox_cm, 0.5 * (a_before + a_after_step),
6321 pmf_b_aref)) : 0.0
6322 isfinite(terminal_gamma_code_override) && (gamma_code = terminal_gamma_code_override)
6323 use_terminal = drag_dt_mode === :terminal
6324 if drag_dt_mode === :hybrid
6325 hybrid_check = cyc == 1 || ((cyc - 1) % hybrid_check_every == 0)
6326 if hybrid_check
6327 hybrid_va_rms = _device_brms!(be_mhd, s, s.scratch[1])
6328 end
6329 last_hybrid_va_cs = hybrid_va_rms / max(cs_floor, eps(Float64))
6330 omega_A = 2π * terminal_keff * max(hybrid_va_rms, eps(Float64))
6331 last_hybrid_drag_omega = gamma_code / max(omega_A, eps(Float64))
6332 if hybrid_check
6333 decision = MHDKernels.update_drag_regime!(hybrid_controller,
6334 last_hybrid_drag_omega,
6335 last_hybrid_va_cs)
6336 hybrid_terminal_active = decision.terminal
6337 end
6338 use_terminal = hybrid_terminal_active
6339 last_hybrid_regime = use_terminal ? "terminal" : "aware"
6340 if use_terminal
6341 nsub = 1
6342 last_nsub = nsub
6343 else
6344 aware_step = _limit_hybrid_aware_dt(dt, dt_explicit,
6345 hybrid_aware_dt_factor)
6346 nsub = aware_step.nsub
6347 last_nsub = nsub
6348 end
6349 if !use_terminal && aware_step.limited
6350 dt = aware_step.dt
6351 last_dt = dt
6352 a_after_step = zrun ? _advance_a_rk2(Om, OL, Ok, Or, a, dt, a_end) : 1.0
6353 zsource = zrun ? a_to_z(0.5 * (a_before + a_after_step)) : 0.0
6354 dtphys_source = zrun ? _cosmic_dt_seconds(a_before, a_after_step, hub, Om, OL, Or) :
6355 dt * chem_time_unit
6356 dt_hydro = zrun ? _hydro_code_dt(dtphys_source, dt, chem_vunit,
6357 terminal_lbox_cm,
6358 0.5 * (a_before + a_after_step),
6359 cosmological_mhd_aref) : dt
6360 if drag_enabled
6361 drag_xe = isfinite(drag_xe_fixed) ? drag_xe_fixed :
6362 _drag_electron_fraction(evolving_xe_mean,
6363 evolving_xheii_mean,
6364 zsource, hub, Ob,
6365 Float64(fh), drag_helium_electrons)
6366 last_drag_xe = drag_xe
6367 gamma_drag = _compton_drag_rate_s(zsource, drag_xe, Float64(fh))
6368 last_drag_gammaH = isfinite(drag_rate_code) ? NaN :
6369 zrun ? gamma_drag / _hubble_s(zsource, hub, Om, OL, Or) : NaN
6370 drag_impulse = _drag_impulse(gamma_drag, dtphys_source, dt, drag_rate_code)
6371 last_drag_f = exp(-drag_impulse)
6372 gamma_code = isfinite(drag_rate_code) ? drag_rate_code : terminal_gamma_scale *
6373 _terminal_gamma_code(gamma_drag, dtphys_source, dt, chem_vunit,
6374 terminal_lbox_cm, 0.5 * (a_before + a_after_step),
6375 pmf_b_aref)
6376 isfinite(terminal_gamma_code_override) && (gamma_code = terminal_gamma_code_override)
6377 last_hybrid_drag_omega = gamma_code / max(omega_A, eps(Float64))
6378 end
6379 end
6380 elseif drag_dt_mode === :terminal
6381 last_hybrid_regime = "terminal"
6382 elseif drag_dt_mode in (:aware, :subcycle, :semiimplicit)
6383 last_hybrid_regime = "aware"
6384 elseif drag_dt_mode === :godunov
6385 last_hybrid_regime = "godunov"
6386 elseif drag_dt_mode === :radiation
6387 last_hybrid_regime = "radiation"
6388 elseif drag_dt_mode === :boost
6389 last_hybrid_regime = "boost"
6390 else
6391 last_hybrid_regime = "explicit"
6392 end
6393 if (use_terminal && terminal_pressure_coeff_auto) || drag_dt_mode===:radiation
6394 pmean = _device_pressure_mean!(be_mhd, s, s.scratch[1])
6395 terminal_pressure_coeff_step = Float64(s.γ) * pmean /
6396 max(terminal_rho_mean, eps(Float64))
6397 end
6398 if drag_dt_mode === :radiation
6399 radiation_velocity_unit = _cosmological_velocity_unit(
6400 chem_vunit, 0.5 * (a_before + a_after_step),
6401 cosmological_mhd_aref)
6402 last_radiation_closure = _radiation_closure(
6403 zsource, last_drag_xe, hub, Ob, fh, radiation_velocity_unit,
6404 terminal_lbox_cm, sqrt(max(terminal_pressure_coeff_step, 0.0));
6405 bridge_q2=radiation_bridge_q2,
6406 longitudinal_viscosity=radiation_long_viscosity,
6407 nyquist_guard_fraction=radiation_nyquist_guard,
6408 response_model=radiation_response_model)
6409 end
6410 if drag_dt_mode === :hybrid && was_terminal && !use_terminal
6411 handoff_events += 1
6412 cumulative_handoff_displacement = 0.0
6413 @printf(" hybrid handoff start cycle=%d z=%.6f G/omA=%.3e va/cs=%.3e steps=%d alpha_max=%.3f ramp=cosine error_tol=%.3e displacement_tol=%.3e\n",
6414 cyc, zrun ? a_to_z(a_before) : NaN, last_hybrid_drag_omega,
6415 last_hybrid_va_cs, hybrid_handoff_steps, hybrid_handoff_alpha,
6416 hybrid_handoff_error_tol, hybrid_handoff_displacement_tol)
6417 flush(stdout)
6418 if hybrid_handoff_steps > 0
6419 handoff_remaining = max(handoff_remaining, hybrid_handoff_steps)
6420 end
6421 end
6422 if drag_dt_mode === :hybrid && !was_terminal && use_terminal && cyc > 1
6423 @printf(" hybrid terminal re-entry cycle=%d z=%.6f G/omA=%.3e va/cs=%.3e\n",
6424 cyc, zrun ? a_to_z(a_before) : NaN, last_hybrid_drag_omega,
6425 last_hybrid_va_cs)
6426 flush(stdout)
6427 end
6428 if use_terminal && (isfinite(terminal_max_disp_cells) ||
6429 isfinite(terminal_diff_cfl) || isfinite(terminal_accuracy_cfl))
6430 for _limiter_iter in 1:2
6431 dt_terminal_limit = Inf
6432 if isfinite(terminal_max_disp_cells) && terminal_max_disp_cells > 0
6433 terminal_vmax = MHDKernels.terminal_velocity_max!(s;
6434 gamma_drag=gamma_code, pressure_coeff=terminal_pressure_coeff_step,
6435 vcap=terminal_vcap)
6436 last_terminal_vmax = terminal_vmax
6437 dt_disp = terminal_max_disp_cells * Float64(s.dx) /
6438 max(terminal_vmax, eps(Float64))
6439 dt_terminal_limit = min(dt_terminal_limit, dt_disp)
6440 end
6441 if isfinite(terminal_diff_cfl)
6442 dt_diff = terminal_diff_cfl * max(gamma_code, eps(Float64)) *
6443 Float64(s.dx)^2 / max(Float64(smax)^2, eps(Float64))
6444 dt_terminal_limit = min(dt_terminal_limit, dt_diff)
6445 end
6446 if isfinite(terminal_accuracy_cfl)
6447 kacc = 2π * max(terminal_keff, eps(Float64))
6448 dt_acc = terminal_accuracy_cfl * max(gamma_code, eps(Float64)) /
6449 (4 * terminal_accuracy_pressure * max(kacc * kacc, eps(Float64)))
6450 dt_terminal_limit = min(dt_terminal_limit, dt_acc)
6451 end
6452 if dt_hydro <= dt_terminal_limit * (1 + 1e-6)
6453 break
6454 end
6455 dt *= dt_terminal_limit / dt_hydro
6456 dt = max(dt, 8eps(Float64) * max(1.0, abs(dt)))
6457 last_dt = dt
6458 nsub = 1
6459 last_nsub = nsub
6460 a_after_step = zrun ? _advance_a_rk2(Om, OL, Ok, Or, a, dt, a_end) : 1.0
6461 zsource = zrun ? a_to_z(0.5 * (a_before + a_after_step)) : 0.0
6462 dtphys_source = zrun ? _cosmic_dt_seconds(a_before, a_after_step, hub, Om, OL, Or) :
6463 dt * chem_time_unit
6464 dt_hydro = zrun ? _hydro_code_dt(dtphys_source, dt, chem_vunit,
6465 terminal_lbox_cm,
6466 0.5 * (a_before + a_after_step),
6467 cosmological_mhd_aref) : dt
6468 if drag_enabled
6469 drag_xe = isfinite(drag_xe_fixed) ? drag_xe_fixed :
6470 _drag_electron_fraction(evolving_xe_mean,
6471 evolving_xheii_mean,
6472 zsource, hub, Ob,
6473 Float64(fh), drag_helium_electrons)
6474 last_drag_xe = drag_xe
6475 gamma_drag = _compton_drag_rate_s(zsource, drag_xe, Float64(fh))
6476 last_drag_gammaH = isfinite(drag_rate_code) ? NaN :
6477 zrun ? gamma_drag / _hubble_s(zsource, hub, Om, OL, Or) : NaN
6478 drag_impulse = _drag_impulse(gamma_drag, dtphys_source, dt, drag_rate_code)
6479 last_drag_f = exp(-drag_impulse)
6480 gamma_code = isfinite(drag_rate_code) ? drag_rate_code : terminal_gamma_scale *
6481 _terminal_gamma_code(gamma_drag, dtphys_source, dt, chem_vunit,
6482 terminal_lbox_cm, 0.5 * (a_before + a_after_step),
6483 pmf_b_aref)
6484 isfinite(terminal_gamma_code_override) && (gamma_code = terminal_gamma_code_override)
6485 if drag_dt_mode === :hybrid
6486 omega_A = 2π * terminal_keff * max(hybrid_va_rms, eps(Float64))
6487 last_hybrid_drag_omega = gamma_code / max(omega_A, eps(Float64))
6488 end
6489 end
6490 end
6491 end
6492 was_terminal = use_terminal
6493 last_handoff_remaining = handoff_remaining
6494 last_handoff_alpha = 0.0
6495 subcycles_total += nsub
6496 tc = 0.0
6497 chem_refresh_mean = cyc <= warmup || chem_mean_every == 1 ||
6498 (cyc % chem_mean_every == 0)
6499 a_chem_mid = zrun ? _advance_a_rk2(Om, OL, Ok, Or, a_before, 0.5 * dt,
6500 a_after_step) : 1.0
6501 if chem_mode !== :off && chem_coupling === :strang
6502 tchem, evolving_xe_mean, evolving_xheii_mean =
6503 _chem_step_interval!(a_before, a_chem_mid, 0.5 * dt,
6504 chem_refresh_mean, measuring, evolving_xe_mean,
6505 evolving_xheii_mean,
6506 false)
6507 tc += tchem
6508 if drag_enabled
6509 drag_xe = isfinite(drag_xe_fixed) ? drag_xe_fixed :
6510 _drag_electron_fraction(evolving_xe_mean, evolving_xheii_mean,
6511 zsource, hub, Ob,
6512 Float64(fh), drag_helium_electrons)
6513 last_drag_xe = drag_xe
6514 gamma_drag = _compton_drag_rate_s(zsource, drag_xe, Float64(fh))
6515 last_drag_gammaH = isfinite(drag_rate_code) ? NaN :
6516 zrun ? gamma_drag / _hubble_s(zsource, hub, Om, OL, Or) : NaN
6517 drag_impulse = _drag_impulse(gamma_drag, dtphys_source, dt, drag_rate_code)
6518 last_drag_f = exp(-drag_impulse)
6519 gamma_code = isfinite(drag_rate_code) ? drag_rate_code : terminal_gamma_scale *
6520 _terminal_gamma_code(gamma_drag, dtphys_source, dt, chem_vunit,
6521 terminal_lbox_cm,
6522 0.5 * (a_before + a_after_step), pmf_b_aref)
6523 isfinite(terminal_gamma_code_override) && (gamma_code = terminal_gamma_code_override)
6524 if drag_dt_mode === :hybrid
6525 omega_A = 2π * terminal_keff * max(hybrid_va_rms, eps(Float64))
6526 last_hybrid_drag_omega = gamma_code / max(omega_A, eps(Float64))
6527 end
6528 end
6529 _debug_div("after_chem_half_pre", cyc, zrun ? a_before : 1.0)
6530 _debug_phase_state("after_chem_half_pre", cyc, a_before, dt, smax)
6531 end
6532 if drag_enabled && gravity && gravity_gas_kick
6533 fast_impulse=drag_impulse*gravity_slip_factor
6534 fast_impulse <= gravity_drag_max_impulse*(1+1e-6) || error(
6535 "chemistry increased the gravity/drag slip impulse beyond its " *
6536 "accuracy cap; reduce MHD_MAXEXP or MHD_SOURCE_MAX_DT_FACTOR")
6537 max_gravity_fast_impulse=max(max_gravity_fast_impulse,fast_impulse)
6538 end
6539 if use_terminal
6540 vcap_step = terminal_vcap
6541 if chem_advect_species
6542 tm += _time_phase(be_mhd, be_p) do
6543 if chem_centered_hydrogen_storage
6544 MHDKernels.conserved_species_to_fraction!(
6545 H2I, s.U[1]; maximum_total=fh)
6546 else
6547 MHDKernels.conserved_species_to_fractions!(
6548 HII, H2I, s.U[1]; maximum_total=fh)
6549 end
6550 chem_helium && MHDKernels.conserved_species_to_fraction!(
6551 HeII, s.U[1]; maximum_total=one(T) - fh)
6552 end
6553 end
6554 tm += _time_phase(be_mhd, be_p) do
6555 if terminal_pressure_implicit
6556 result = _terminal_pressure_implicit_step!(be_mhd, s, dt_hydro, gamma_code,
6557 vcap_step, terminal_pressure_coeff_step,
6558 terminal_rho_rel_limit,
6559 terminal_b_rel_limit;
6560 induction=terminal_induction,
6561 induction_cfl=terminal_induction_cfl,
6562 induction_dissipation=terminal_induction_dissipation,
6563 induction_dissipation_order=terminal_induction_dissipation_order,
6564 project_b=terminal_project_b,
6565 max_induction_subcycles=terminal_induction_max_subcycles)
6566 last_induction_nsub = result.nsub
6567 induction_subcycles_total += result.nsub
6568 terminal_step_timing = result
6569 terminal_timed = true
6570 elseif terminal_split && (terminal_dedicated_scratch || !(φ === ρtot && eint === nothing))
6571 tvx = terminal_dedicated_scratch ? terminal_tvx : ρdm
6572 tvy = terminal_dedicated_scratch ? terminal_tvy : ρtot
6573 tvz = terminal_dedicated_scratch ? terminal_tvz :
6574 (φ === ρtot && eint !== nothing) ? eint : φ
6575 _terminal_velocity_step_cached!(be_mhd, s, dt_hydro, gamma_code, vcap_step,
6576 tvx, tvy, tvz, terminal_vsmooth)
6577 else
6578 _terminal_velocity_step!(be_mhd, s, dt_hydro, gamma_code, vcap_step)
6579 end
6580 end
6581 if terminal_implicit_diffuse
6582 tm += _time_phase(be_mhd, be_p) do
6583 br_step = _device_brms!(be_mhd, s, s.scratch[1])
6584 coeff = terminal_implicit_diffuse_fac * dt_hydro * br_step^2 /
6585 max(gamma_code, eps(Float64)) / max(Float64(s.dx)^2, eps(Float64))
6586 last_terminal_diffuse_coeff = coeff
6587 _terminal_implicit_diffuse_B!(s, coeff)
6588 end
6589 else
6590 last_terminal_diffuse_coeff = 0.0
6591 end
6592 if chem_advect_species
6593 tm += _time_phase(be_mhd, be_p) do
6594 if chem_centered_hydrogen_storage
6595 MHDKernels.advect_passive_scalar!(
6596 HII, s.scratch[5], s, dt_hydro)
6597 MHDKernels.advect_species_fraction!(
6598 H2I, s.scratch[9], s, dt_hydro; maximum_total=fh)
6599 else
6600 MHDKernels.advect_species_fractions!(
6601 HII, H2I, s.scratch[5], s.scratch[9], s, dt_hydro;
6602 maximum_total=fh)
6603 end
6604 chem_helium && MHDKernels.advect_species_fraction!(
6605 HeII, eint, s, dt_hydro; maximum_total=one(T) - fh)
6606 end
6607 end
6608 else
6609 handoff_active_this_step = handoff_remaining > 0
6610 handoff_alpha_step = 0.0
6611 if handoff_remaining > 0
6612 vcap_step = terminal_vcap
6613 handoff_alpha_step = _hybrid_handoff_alpha(hybrid_handoff_alpha,
6614 handoff_remaining,
6615 hybrid_handoff_steps)
6616 last_handoff_alpha = handoff_alpha_step
6617 tdrag += _time_phase(be_mhd, be_p) do
6618 _terminal_handoff_momenta!(be_mhd, s, gamma_code, vcap_step,
6619 handoff_alpha_step)
6620 end
6621 end
6622 hubble_half_decay = sqrt(_hubble_momentum_decay(
6623 a_before, a_after_step, cosmological_mhd_aref)^(1 / nsub))
6624 for _ in 1:nsub
6625 dtsub = dt_hydro / nsub
6626 source_aware_drag = drag_enabled && drag_dt_mode === :godunov
6627 radiation_drag = drag_enabled && drag_dt_mode === :radiation
6628 exponential_drag = drag_enabled &&
6629 drag_dt_mode in (:aware, :semiimplicit, :hybrid)
6630 if hubble_half_decay < 1
6631 tdrag += _time_phase(be_mhd, be_p) do
6632 MHDKernels.apply_linear_drag!(s,hubble_half_decay;
6633 energy_compensation=radiation_drag ?
6634 radiation_workspace.energy_residual : nothing)
6635 radiation_photon_state === nothing ||
6636 MHDKernels.rescale_photon_velocities!(
6637 radiation_photon_state,inv(hubble_half_decay))
6638 end
6639 end
6640 if drag_enabled && !exponential_drag && !source_aware_drag && !radiation_drag
6641 fhalf = exp(-0.5 * drag_impulse / nsub)
6642 tdrag += _time_phase(be_mhd, be_p) do
6643 _apply_compton_drag_mhd!(be_mhd, s, fhalf, drag_vx0, drag_vy0, drag_vz0)
6644 end
6645 end
6646 godunov_before = radiation_step_timing.godunov_s
6647 if chem_advect_species
6648 tm += _time_phase(be_mhd, be_p) do
6649 if chem_centered_hydrogen_storage
6650 MHDKernels.conserved_species_to_fraction!(
6651 H2I, s.U[1]; maximum_total=fh)
6652 else
6653 MHDKernels.conserved_species_to_fractions!(
6654 HII, H2I, s.U[1]; maximum_total=fh)
6655 end
6656 chem_helium && MHDKernels.conserved_species_to_fraction!(
6657 HeII, s.U[1]; maximum_total=one(T) - fh)
6658 end
6659 end
6660 step_elapsed = _time_phase(be_mhd, be_p) do
6661 if radiation_drag
6662 debug_radiation_now = debug_radiation_rho_min > 0 &&
6663 cyc >= debug_radiation_cycle_min
6664 MHDKernels.step_radiation_godunov!(
6665 s,radiation_workspace,dtsub;
6666 drag_impulse=drag_impulse/nsub,
6667 closure=last_radiation_closure,
6668 ch=ch,glm_cr=glm_cr,integrator=mhd_integrator,
6669 photon_state=radiation_photon_state,
6670 debug_rho_min=debug_radiation_now ?
6671 debug_radiation_rho_min : 0.0,
6672 debug_context=debug_radiation_now ?
6673 "cycle=$cyc z=$(a_to_z(a_before))" : "",
6674 debug_correction_cfl=cyc >= debug_radiation_cycle_min ?
6675 debug_radiation_correction_cfl : Inf,
6676 timing=radiation_step_timing,
6677 free_streaming_tolerance=radiation_free_stream_tolerance,
6678 positivity_fallback=radiation_positivity_fallback,
6679 positivity_max_fallback_passes=
6680 radiation_positivity_max_fallback_passes,
6681 positivity_wave_growth_limit=
6682 radiation_positivity_wave_growth_eff,
6683 positivity_max_subdivisions=
6684 radiation_positivity_max_subdivisions,
6685 checkpoint_fallback_on_rejection=
6686 radiation_checkpoint_fallback,
6687 energy_ledger=radiation_energy_ledger !== nothing &&
6688 cyc % radiation_energy_ledger_every == 0 ?
6689 radiation_energy_ledger : nothing)
6690 elseif source_aware_drag
6691 MHDKernels.step_drag_godunov!(
6692 s, dtsub; drag_impulse=drag_impulse / nsub,
6693 ch=ch, glm_cr=glm_cr, integrator=mhd_integrator,
6694 target_velocity=(drag_vx0, drag_vy0, drag_vz0))
6695 else
6696 MHDKernels.step!(s, dtsub; ch=ch, glm_cr=glm_cr,
6697 integrator=mhd_integrator)
6698 end
6699 end
6700 if radiation_drag
6701 godunov_delta = radiation_step_timing.godunov_s-godunov_before
6702 tm += godunov_delta
6703 trad += max(step_elapsed-godunov_delta,0.0)
6704 else
6705 tm += step_elapsed
6706 end
6707 if exponential_drag
6708 tdrag += _time_phase(be_mhd, be_p) do
6709 MHDKernels.apply_exponential_drag_increment!(
6710 s, s.scratch, drag_impulse / nsub;
6711 target_velocity=(drag_vx0, drag_vy0, drag_vz0))
6712 end
6713 elseif drag_enabled && !source_aware_drag && !radiation_drag
6714 fhalf = exp(-0.5 * drag_impulse / nsub)
6715 tdrag += _time_phase(be_mhd, be_p) do
6716 _apply_compton_drag_mhd!(be_mhd, s, fhalf, drag_vx0, drag_vy0, drag_vz0)
6717 end
6718 end
6719 if hubble_half_decay < 1
6720 tdrag += _time_phase(be_mhd, be_p) do
6721 MHDKernels.apply_linear_drag!(s,hubble_half_decay;
6722 energy_compensation=radiation_drag ?
6723 radiation_workspace.energy_residual : nothing)
6724 radiation_photon_state === nothing ||
6725 MHDKernels.rescale_photon_velocities!(
6726 radiation_photon_state,inv(hubble_half_decay))
6727 end
6728 end
6729 if chem_advect_species
6730 tm += _time_phase(be_mhd, be_p) do
6731 if chem_centered_hydrogen_storage
6732 MHDKernels.advect_passive_scalar!(
6733 HII, s.scratch[5], s, dtsub)
6734 MHDKernels.advect_species_fraction!(
6735 H2I, s.scratch[9], s, dtsub; maximum_total=fh)
6736 else
6737 MHDKernels.advect_species_fractions!(
6738 HII, H2I, s.scratch[5], s.scratch[9], s, dtsub;
6739 maximum_total=fh)
6740 end
6741 chem_helium && MHDKernels.advect_species_fraction!(
6742 HeII, eint, s, dtsub; maximum_total=one(T) - fh)
6743 end
6744 end
6745 end
6746 if glm_project_every > 0 && cyc % glm_project_every == 0
6747 tgp += _time_phase(be_mhd, be_p) do
6748 _project_mhd_b!(be_mhd,s;
6749 energy_compensation=drag_enabled && drag_dt_mode===:radiation ?
6750 radiation_workspace.energy_residual : nothing)
6751 end
6752 glm_projection_steps += 1
6753 end
6754 if handoff_active_this_step
6755 overlap = nothing
6756 tdrag += _time_phase(be_mhd, be_p) do
6757 overlap = _terminal_handoff_momenta!(be_mhd, s, gamma_code, terminal_vcap,
6758 0.0; measure=true)
6759 end
6760 last_handoff_error = overlap.relative
6761 max_handoff_error = isnan(max_handoff_error) ? overlap.relative :
6762 max(max_handoff_error, overlap.relative)
6763 max_handoff_settled_error = isnan(max_handoff_settled_error) ? overlap.relative :
6764 max(max_handoff_settled_error, overlap.relative)
6765 last_handoff_delta_v_rms = overlap.delta_v_rms
6766 last_handoff_terminal_v_rms = overlap.terminal_v_rms
6767 last_handoff_full_v_rms = overlap.full_v_rms
6768 handoff_index = handoff_steps_done + 1
6769 raw_handoff_displacement = overlap.delta_v_rms * dt_hydro / Float64(s.dx)
6770 last_handoff_displacement = handoff_alpha_step * raw_handoff_displacement
6771 max_handoff_displacement = max(max_handoff_displacement,
6772 last_handoff_displacement)
6773 cumulative_handoff_displacement += last_handoff_displacement
6774 if handoff_index == 1 || handoff_remaining == 1 ||
6775 handoff_index % hybrid_handoff_diag_every == 0
6776 @printf(" hybrid overlap step=%d remaining=%d alpha=%.6f rel=%.6e dv_rms=%.6e vt_rms=%.6e vf_rms=%.6e raw_dcell=%.6e correction_dcell=%.6e sum_correction_dcell=%.6e\n",
6777 handoff_index, handoff_remaining, handoff_alpha_step, overlap.relative,
6778 overlap.delta_v_rms, overlap.terminal_v_rms, overlap.full_v_rms,
6779 raw_handoff_displacement, last_handoff_displacement,
6780 cumulative_handoff_displacement)
6781 flush(stdout)
6782 end
6783 overlap.relative <= hybrid_handoff_error_tol ||
6784 error("hybrid handoff velocity mismatch $(overlap.relative) exceeds " *
6785 "MHD_HYBRID_HANDOFF_ERROR_TOL=$(hybrid_handoff_error_tol)")
6786 cumulative_handoff_displacement <= hybrid_handoff_displacement_tol ||
6787 error("hybrid handoff accumulated correction displacement " *
6788 "$(cumulative_handoff_displacement) cells exceeds " *
6789 "MHD_HYBRID_HANDOFF_DISPLACEMENT_TOL=" *
6790 "$(hybrid_handoff_displacement_tol)")
6791 handoff_remaining -= 1
6792 handoff_steps_done += 1
6793 end
6794 end
6795 _debug_div("after_mhd", cyc, zrun ? a_before : 1.0)
6796 _debug_phase_state("after_mhd", cyc, a_before, dt, smax)
6797 if chem_mode !== :off
6798 if chem_coupling === :strang
6799 tchem, evolving_xe_mean, evolving_xheii_mean =
6800 _chem_step_interval!(a_chem_mid, a_after_step, 0.5 * dt,
6801 chem_refresh_mean, measuring, evolving_xe_mean,
6802 evolving_xheii_mean,
6803 true)
6804 tc += tchem
6805 else
6806 tchem, evolving_xe_mean, evolving_xheii_mean =
6807 _chem_step_interval!(a_before, a_after_step, dt,
6808 chem_refresh_mean, measuring, evolving_xe_mean,
6809 evolving_xheii_mean,
6810 false)
6811 tc += tchem
6812 end
6813 end
6814 _debug_div("after_chem", cyc, zrun ? a_before : 1.0)
6815 _debug_phase_state("after_chem", cyc, a_before, dt, smax)
6816 td = ta = tf = tg = tp = 0.0
6817 if gravity
6818 ngrav = gravity_subcycles_fixed
6819 if gravity_dt_boost > 0
6820 grav_max_dt = max(gravity_dt_boost * dt_explicit, eps(Float64))
6821 ngrav = max(ngrav, ceil(Int, dt / grav_max_dt))
6822 end
6823 if isfinite(particle_max_disp_cells)
6824 # Reuse the previous potential's acceleration bound rather than
6825 # adding a preflight deposit/FFT to the total-source path.
6826 if isfinite(particle_accel_max_dt) && particle_accel_max_dt > 0
6827 ngrav = max(ngrav, ceil(Int, dt / particle_accel_max_dt))
6828 end
6829 _particle_speed2_k!(be_p)(ρdm, vx, vy, vz; ndrange=ncells)
6830 KA.synchronize(be_p)
6831 last_particle_vmax = sqrt(max(0.0, Float64(maximum(ρdm))))
6832 if last_particle_vmax > 0
6833 particle_max_dt = particle_max_disp_cells / (Float64(N) * last_particle_vmax)
6834 ngrav = max(ngrav, ceil(Int, dt / max(particle_max_dt, eps(Float64))))
6835 end
6836 end
6837 last_gravity_nsub = max(1, ngrav)
6838 dtdm = dt / last_gravity_nsub
6839 if grav_dm_weight == 0
6840 ta += _time_phase(be_p, be_mhd) do
6841 if gravity_midpoint_source
6842 if radiation_workspace === nothing
6843 MHDKernels.predict_density_backward!(ρtot,s,0.5*dt_hydro)
6844 _assemble_total_perturbation_k!(be_p)(
6845 ρtot, ρtot, ρdm, grav_gas_weight,
6846 grav_dm_weight, one(T), dm_density_mean, Int(N);
6847 ndrange=ncells,
6848 )
6849 else
6850 MHDKernels.predict_density_perturbation_backward!(ρtot,
6851 radiation_workspace.density_perturbation,
6852 s.U[2],s.U[3],s.U[4],s.dims;
6853 dx=s.dx,dt_back=0.5*dt_hydro)
6854 _assemble_total_perturbation_k!(be_p)(
6855 ρtot, ρtot, ρdm, grav_gas_weight,
6856 grav_dm_weight, zero(T), dm_density_mean, Int(N);
6857 ndrange=ncells,
6858 )
6859 end
6860 else
6861 if radiation_workspace === nothing
6862 _assemble_total_perturbation_k!(be_p)(
6863 ρtot, s.U[1], ρdm, grav_gas_weight,
6864 grav_dm_weight, one(T), dm_density_mean, Int(N);
6865 ndrange=ncells,
6866 )
6867 else
6868 _assemble_total_perturbation_k!(be_p)(
6869 ρtot, radiation_workspace.density_perturbation,
6870 ρdm, grav_gas_weight, grav_dm_weight, zero(T),
6871 dm_density_mean, Int(N); ndrange=ncells,
6872 )
6873 end
6874 end
6875 end
6876 old_single_step = !gravity_midpoint_source && last_gravity_nsub == 1 &&
6877 gravity_subcycles_fixed == 1 && gravity_dt_boost == 0
6878 agrav = old_single_step ? a : (zrun ? 0.5 * (a_before + a_after_step) : 1.0)
6879 tf += _time_phase(be_p, be_mhd) do
6880 gravcoef = zrun ? 1.5 * Om * agrav : 1.0
6881 if fft_mode === :mps
6882 PoissonKernels.fft_poisson_rfft!(φ, ρtot; G=gravcoef, a=1, boxsize=1)
6883 elseif fft_mode === :ka
6884 PoissonKernels.fft_poisson_rfft_ka!(φ, ρtot; G=gravcoef, a=1, boxsize=1)
6885 else
6886 PoissonKernels.fft_poisson_root_gpu!(φ, ρtot; G=gravcoef, a=1, boxsize=1)
6887 end
6888 end
6889 if gravity_gas_kick
6890 dtgas = zrun ? _gas_gravity_kick_dt(
6891 dt, agrav, hub, chem_vunit, terminal_lbox_cm,
6892 cosmological_mhd_aref) : dt
6893 tg += _time_phase(be_mhd, be_p) do
6894 MHDKernels.apply_cell_center_gravity!(s,φ,dtgas;
6895 energy_compensation=drag_enabled && drag_dt_mode===:radiation ?
6896 radiation_workspace.energy_residual : nothing)
6897 end
6898 end
6899 if isfinite(particle_max_disp_cells)
6900 _grid_accel2_k!(be_p)(ρdm, φ, Int(N); ndrange=ncells)
6901 KA.synchronize(be_p)
6902 last_particle_gmax = sqrt(max(0.0, Float64(maximum(ρdm))))
6903 maxdisp = particle_max_disp_cells / Float64(N)
6904 particle_max_dt = _particle_dt_limit(last_particle_vmax, last_particle_gmax, maxdisp)
6905 if isfinite(particle_max_dt) && particle_max_dt > 0
6906 last_gravity_nsub = max(last_gravity_nsub,
6907 ceil(Int, dt / max(particle_max_dt, eps(Float64))))
6908 dtdm = dt / last_gravity_nsub
6909 end
6910 end
6911 for ig in 1:last_gravity_nsub
6912 tp += _time_phase(be_p, be_mhd) do
6913 if φprev === nothing
6914 _push_particles!(px, py, pz, vx, vy, vz, φ;
6915 dtau=dtdm, N=N, wrap=particle_wrap,
6916 lattice_displacements=lattice_displacements)
6917 else
6918 θ = (ig - 0.5) / last_gravity_nsub
6919 _push_particles_fused_global_phi_mix!(px, py, pz, vx, vy, vz, φprev, φ;
6920 theta=θ, dtau=dtdm, nc=(N,N,N), wrap=1)
6921 end
6922 end
6923 end
6924 if φprev !== nothing
6925 tf += _time_phase(be_p, be_mhd) do
6926 copyto!(φprev, φ)
6927 end
6928 end
6929 else
6930 for ig in 1:last_gravity_nsub
6931 td += _time_phase(be_p, be_mhd) do
6932 midpoint_disp = gravity_midpoint_source ? 0.5 * dtdm : 0.0
6933 _deposit_particles!(ρdm, px, py, pz, vx, vy, vz;
6934 N=N, disp=midpoint_disp,
6935 lattice_displacements=lattice_displacements)
6936 end
6937 ta += _time_phase(be_p, be_mhd) do
6938 if gravity_midpoint_source && grav_gas_weight != 0
6939 dt_back = dt - (ig - 0.5) * dtdm
6940 dt_back_hydro = _hydro_subinterval(dt_hydro, dt, dt_back)
6941 if radiation_workspace === nothing
6942 MHDKernels.predict_density_backward!(ρtot,s,dt_back_hydro)
6943 _assemble_total_perturbation_k!(be_p)(
6944 ρtot, ρtot, ρdm, grav_gas_weight,
6945 grav_dm_weight, one(T), dm_density_mean,
6946 Int(N); ndrange=ncells,
6947 )
6948 else
6949 MHDKernels.predict_density_perturbation_backward!(ρtot,
6950 radiation_workspace.density_perturbation,
6951 s.U[2],s.U[3],s.U[4],s.dims;
6952 dx=s.dx,dt_back=dt_back_hydro)
6953 _assemble_total_perturbation_k!(be_p)(
6954 ρtot, ρtot, ρdm, grav_gas_weight,
6955 grav_dm_weight, zero(T), dm_density_mean,
6956 Int(N); ndrange=ncells,
6957 )
6958 end
6959 else
6960 if radiation_workspace === nothing
6961 _assemble_total_perturbation_k!(be_p)(
6962 ρtot, s.U[1], ρdm, grav_gas_weight,
6963 grav_dm_weight, one(T), dm_density_mean,
6964 Int(N); ndrange=ncells,
6965 )
6966 else
6967 _assemble_total_perturbation_k!(be_p)(
6968 ρtot,
6969 radiation_workspace.density_perturbation,
6970 ρdm, grav_gas_weight, grav_dm_weight,
6971 zero(T), dm_density_mean, Int(N);
6972 ndrange=ncells,
6973 )
6974 end
6975 end
6976 end
6977 old_single_step = last_gravity_nsub == 1 &&
6978 gravity_subcycles_fixed == 1 &&
6979 gravity_dt_boost == 0 &&
6980 !gravity_midpoint_source
6981 agrav = old_single_step ? a :
6982 (zrun ? a_before + (a_after_step - a_before) * ((ig - 0.5) / last_gravity_nsub) : 1.0)
6983 tf += _time_phase(be_p, be_mhd) do
6984 gravcoef = zrun ? 1.5 * Om * agrav : 1.0
6985 if fft_mode === :mps
6986 PoissonKernels.fft_poisson_rfft!(φ, ρtot; G=gravcoef, a=1, boxsize=1)
6987 elseif fft_mode === :ka
6988 PoissonKernels.fft_poisson_rfft_ka!(φ, ρtot; G=gravcoef, a=1, boxsize=1)
6989 else
6990 PoissonKernels.fft_poisson_root_gpu!(φ, ρtot; G=gravcoef, a=1, boxsize=1)
6991 end
6992 end
6993 if isfinite(particle_max_disp_cells) && ig == 1
6994 _grid_accel2_k!(be_p)(ρdm, φ, Int(N); ndrange=ncells)
6995 KA.synchronize(be_p)
6996 last_particle_gmax = sqrt(max(0.0, Float64(maximum(ρdm))))
6997 maxdisp = particle_max_disp_cells / Float64(N)
6998 particle_accel_max_dt = _particle_dt_limit(last_particle_vmax,
6999 last_particle_gmax, maxdisp)
7000 end
7001 if gravity_gas_kick
7002 dtgas = zrun ? _gas_gravity_kick_dt(
7003 dtdm, agrav, hub, chem_vunit, terminal_lbox_cm,
7004 cosmological_mhd_aref) : dtdm
7005 tg += _time_phase(be_mhd, be_p) do
7006 MHDKernels.apply_cell_center_gravity!(s,φ,dtgas;
7007 energy_compensation=drag_enabled && drag_dt_mode===:radiation ?
7008 radiation_workspace.energy_residual : nothing)
7009 end
7010 end
7011 tp += _time_phase(be_p, be_mhd) do
7012 _push_particles!(px, py, pz, vx, vy, vz, φ;
7013 dtau=dtdm, N=N, wrap=particle_wrap,
7014 lattice_displacements=lattice_displacements)
7015 end
7016 end
7017 end
7018 else
7019 last_gravity_nsub = 0
7020 end
7021 _debug_div("after_gravity", cyc, zrun ? a_before : 1.0)
7022 _debug_phase_state("after_gravity", cyc, a_before, dt, smax)
7023
7024 if measuring
7025 measured += 1
7026 measured_subcycles += nsub
7027 gravity_subcycles_total += last_gravity_nsub
7028 elapsed_time += dt
7029 phase[:cfl] += tcfl
7030 phase[:mhd] += tm
7031 phase[:radiation] += trad
7032 phase[:glm_projection] += tgp
7033 phase[:drag] += tdrag
7034 phase[:deposit] += td
7035 phase[:assemble] += ta
7036 phase[:fft] += tf
7037 phase[:gas_gravity] += tg
7038 phase[:push] += tp
7039 phase[:chem] += tc
7040 radiation_phase[:predictor] +=
7041 radiation_step_timing.predictor_s-radiation_timing_before.predictor
7042 radiation_phase[:corrector] +=
7043 radiation_step_timing.corrector_s-radiation_timing_before.corrector
7044 radiation_phase[:ledger] +=
7045 radiation_step_timing.ledger_s-radiation_timing_before.ledger
7046 if terminal_timed
7047 terminal_measured += 1
7048 terminal_phase[:force] += terminal_step_timing.force_s
7049 terminal_phase[:pressure_fft] += terminal_step_timing.pressure_fft_s
7050 terminal_phase[:finalize] += terminal_step_timing.finalize_s
7051 terminal_phase[:induction] += terminal_step_timing.induction_s
7052 terminal_phase[:projection] += terminal_step_timing.projection_s
7053 end
7054 end
7055 if cycle_print_every > 0 && (cyc <= warmup || cyc % cycle_print_every == 0)
7056 @printf("cycle %3d%s dt=%.3e hdt=%.3e pdt_s=%.3e expdt=%.3e nsub=%d gsub=%d pvmax=%.3e pgmax=%.3e regime=%s handoff=%d herr=%.3e va/cs=%.3e G/omA=%.3e idiff=%.3e drag %.4f f=%.6f G/H=%.3e cfl %.4f mhd %.4f rad %.4f chem %.4f dep %.4f asm %.4f fft %.4f ggas %.4f push %.4f\n",
7057 cyc, measuring ? "" : " warm", dt, dt_hydro, dtphys_source,
7058 dt_explicit, nsub, last_gravity_nsub,
7059 last_particle_vmax, last_particle_gmax, last_hybrid_regime,
7060 last_handoff_remaining, last_handoff_error,
7061 last_hybrid_va_cs, last_hybrid_drag_omega,
7062 last_terminal_diffuse_coeff, tdrag, last_drag_f, last_drag_gammaH,
7063 tcfl, tm, trad, tc, td, ta, tf, tg, tp)
7064 flush(stdout)
7065 if drag_dt_mode === :radiation
7066 rd = _radiation_scale_diagnostics(last_radiation_closure, N)
7067 free_bound=MHDKernels.radiation_free_streaming_bound(
7068 2pi,drag_impulse/max(dt_hydro,eps(Float64)),
7069 last_radiation_closure)
7070 @printf(" radiation R=%.6e lambda/L=%.6e q=[%.3e,%.3e] stream=[%.3e,%.3e] crad/cs=[%.3e,%.3e] free_bound=%.3e fast=%d spectral=%d\n",
7071 last_radiation_closure.inertia_ratio,
7072 last_radiation_closure.mean_free_path,
7073 rd.qfund, rd.qnyquist, rd.stream_fund,
7074 rd.stream_nyquist,
7075 rd.crad_fund / max(last_radiation_closure.gas_sound_speed,
7076 eps(Float64)),
7077 rd.crad_nyquist / max(last_radiation_closure.gas_sound_speed,
7078 eps(Float64)),
7079 free_bound,radiation_step_timing.fast_steps,
7080 radiation_step_timing.spectral_steps)
7081 flush(stdout)
7082 end
7083 end
7084 if check_every > 0 && (cyc % check_every == 0)
7085 ss = _assert_finite_state(s, "post-step", cyc, zrun ? a : NaN, dt, smax)
7086 @printf(" check rho=[%.6e, %.6e] E=[%.6e, %.6e] Bmax=%.6e\n",
7087 Float64(ss.ρmin), Float64(ss.ρmax), Float64(ss.Emin), Float64(ss.Emax), ss.Bmax)
7088 flush(stdout)
7089 end
7090 if zrun && cyc > warmup
7091 a = a_after_step
7092 if measuring && print_every > 0 &&
7093 (measured % print_every == 0 || a >= a_end)
7094 @printf(" z=%.3f a=%.6e measured=%d\n", a_to_z(a), a, measured)
7095 flush(stdout)
7096 end
7097 if measuring && frame_next[] <= length(frame_zs) &&
7098 _reached_output_z(a_to_z(a), frame_zs[frame_next[]])
7099 _write_run_density_frames!(frame_dir, frame_hostrho, frame_hostslice, frame_hostproj,
7100 s,ρdm,(px,py,pz),(vx,vy,vz),Int(N),a_to_z(a),cyc;
7101 field=frame_field,lattice_displacements=lattice_displacements)
7102 while frame_next[] <= length(frame_zs) &&
7103 _reached_output_z(a_to_z(a), frame_zs[frame_next[]])
7104 frame_next[] += 1
7105 end
7106 @printf(" frame z=%.3f wrote density slice and maximum projection\n", a_to_z(a))
7107 flush(stdout)
7108 end
7109 end
7110 if measuring && (!isempty(history_path) || !isempty(mode_history_path) ||
7111 !isempty(pk_path) || !isempty(pdf_path))
7112 znow = zrun ? a_to_z(a) : NaN
7113 due_every = pk_every > 0 && measured % pk_every == 0
7114 due_z = zrun && pk_next[] <= length(pk_zs) &&
7115 _reached_output_z(znow, pk_zs[pk_next[]])
7116 due_adaptive = false
7117 adaptive_delta_b = NaN
7118 adaptive_xhii = NaN
7119 in_adaptive_window = !zrun ||
7120 (isfinite(znow) && adaptive_diag_zmin <= znow <= adaptive_diag_zmax)
7121 if adaptive_diag_rel > 0 && in_adaptive_window && !isempty(history_path)
7122 adaptive_delta_b,adaptive_xhii=_adaptive_diag_watch!(
7123 be_p,s,ρtot,HII,H2I,HeII,fh;
7124 z=znow,h=hub,Ob=Ob,
7125 density_perturbation=radiation_workspace === nothing ? nothing :
7126 radiation_workspace.density_perturbation,
7127 neutral_hydrogen_storage=chem_neutral_hydrogen_storage,
7128 centered_hydrogen_storage=chem_centered_hydrogen_storage,
7129 hydrogen_reference=evolving_xe_mean)
7130 due_adaptive = _rel_changed(adaptive_delta_b, adaptive_last_delta_b[], adaptive_diag_rel) ||
7131 _rel_changed(adaptive_xhii, adaptive_last_xhii[], adaptive_diag_rel)
7132 end
7133 if due_every || due_z || due_adaptive
7134 full_diag = due_every || due_z
7135 tdiag = @elapsed begin
7136 _emit_diagnostics!(be_p=be_p, s=s, ρdm=ρdm, scratch=ρtot,
7137 px=px, py=py, pz=pz, vx=vx, vy=vy, vz=vz,
7138 N=N, ncells=ncells, gravity=gravity, cycle=cyc,
7139 measured=measured, a=zrun ? a : 1.0,
7140 pk_path=full_diag ? pk_path : "", history_path=history_path,
7141 pdf_path=full_diag ? pdf_path : "", pdf_counts=pdf_counts,
7142 pdf_loglo=pdf_loglo, pdf_loghi=pdf_loghi,
7143 nmu=pk_nmu, nbins=pk_nbins, boxsize=boxsize,
7144 axis=pk_axis, kmax=pk_kmax,
7145 lattice_displacements=lattice_displacements,
7146 HII=HII, H2I=H2I, HeII=HeII, fh=fh,
7147 btrans0=Float64(st.btrans0),
7148 density_kmode=pmf_kmode,
7149 mode_history_path=mode_history_path,
7150 mode_kmode=pmf_kmode, hubble=hub, omega_b=Ob,
7151 physical_box_ckpc_h=physical_box_ckpc_h, streaming_path=streaming_path,
7152 gas_velocity_unit=_cosmological_velocity_unit(chem_vunit,a,cosmological_mhd_aref),
7153 radiation_workspace=radiation_workspace,
7154 radiation_photon_state=radiation_photon_state,
7155 radiation_drag_rate=drag_impulse /
7156 max(dt_hydro,eps(Float64)),
7157 neutral_hydrogen_storage=chem_neutral_hydrogen_storage,
7158 centered_hydrogen_storage=chem_centered_hydrogen_storage,
7159 hydrogen_reference=evolving_xe_mean,
7160 temperature_velocity_unit2=chem_mode === :off ? NaN :
7161 chem_vunit^2 *
7162 (isfinite(cosmological_mhd_aref) ?
7163 cosmological_mhd_aref / a : 1.0))
7164 KA.synchronize(be_mhd)
7165 KA.synchronize(be_p)
7166 # MPSGraph diagnostics create independent rFFT plans. This
7167 # opt-in Metal barrier isolates them from the next terminal
7168 # pressure/chemistry FFT step.
7169 if full_diag && diag_clear_fft
7170 PoissonKernels.clear_fft_scratch!(be_p)
7171 end
7172 end
7173 if adaptive_diag_rel > 0
7174 adaptive_last_delta_b[] = isfinite(adaptive_delta_b) ? adaptive_delta_b :
7175 (radiation_workspace === nothing ?
7176 _device_delta_rms!(be_p,ρtot,s.U[1],one(T)) :
7177 _device_delta_rms!(be_p,ρtot,
7178 radiation_workspace.density_perturbation,
7179 zero(T)))
7180 adaptive_last_xhii[] = isfinite(adaptive_xhii) ? adaptive_xhii :
7181 _device_species_stats!(
7182 be_p, ρtot, s.U[1], HII, H2I,
7183 HeII, fh; z=znow,h=hub,Ob=Ob,
7184 neutral_hydrogen_storage=
7185 chem_neutral_hydrogen_storage,
7186 centered_hydrogen_storage=
7187 chem_centered_hydrogen_storage,
7188 hydrogen_reference=evolving_xe_mean)[1]
7189 end
7190 if full_diag && !isempty(pk_path)
7191 reason = due_z ? "z" : due_every ? "every" : due_adaptive ? "adaptive" : "manual"
7192 @printf(" diag z=%.3f wrote P(k,mu) in %.3f s (%s)\n",
7193 zrun ? znow : NaN, tdiag, reason)
7194 flush(stdout)
7195 end
7196 _debug_div("after_diag", cyc, zrun ? a : 1.0)
7197 while zrun && pk_next[] <= length(pk_zs) &&
7198 _reached_output_z(znow, pk_zs[pk_next[]])
7199 pk_next[] += 1
7200 end
7201 end
7202 end
7203 if (checkpoint_every > 0 && measuring && measured % checkpoint_every == 0) ||
7204 (checkpoint_at_cycle > 0 && cyc == checkpoint_at_cycle)
7205 _save_checkpoint!(cyc)
7206 end
7207 if tfinal === nothing
7208 if zrun
7209 (a >= a_end || measured >= maxcycles) && break
7210 else
7211 cyc >= warmup + steps && break
7212 end
7213 else
7214 elapsed_time >= tfinal * (1 - 1e-9) && break
7215 end
7216 end
7217
7218 if checkpoint_final && !isempty(checkpoint_write_path) && !checkpoint_disable_writes
7219 _save_checkpoint!(cyc)
7220 end
7221
7222 if gravity && final_diag
7223 _deposit_particles!(ρdm, px, py, pz, vx, vy, vz;
7224 N=N, disp=0,
7225 lattice_displacements=lattice_displacements)
7226 KA.synchronize(be_p)
7227 end
7228 # Loop diagnostics are evaluated before advancing a. Emit a final target
7229 # after the last accepted step so z=zend P(k,mu) and PDF products exist.
7230 final_pk_due = zrun && pk_next[] <= length(pk_zs) &&
7231 _reached_output_z(a_to_z(a), pk_zs[pk_next[]])
7232 if final_diag && final_pk_due
7233 _emit_diagnostics!(be_p=be_p, s=s, ρdm=ρdm, scratch=ρtot,
7234 px=px, py=py, pz=pz, vx=vx, vy=vy, vz=vz,
7235 N=N, ncells=ncells, gravity=gravity, cycle=cyc,
7236 measured=measured, a=a,
7237 pk_path=pk_path, history_path=history_path, pdf_path=pdf_path,
7238 pdf_counts=pdf_counts, pdf_loglo=pdf_loglo, pdf_loghi=pdf_loghi,
7239 nmu=pk_nmu, nbins=pk_nbins, boxsize=boxsize, axis=pk_axis,
7240 kmax=pk_kmax,
7241 lattice_displacements=lattice_displacements,
7242 HII=HII, H2I=H2I, HeII=HeII, fh=fh,
7243 btrans0=Float64(st.btrans0), density_kmode=pmf_kmode,
7244 mode_history_path=mode_history_path,
7245 mode_kmode=pmf_kmode, hubble=hub, omega_b=Ob,
7246 physical_box_ckpc_h=physical_box_ckpc_h, streaming_path=streaming_path,
7247 gas_velocity_unit=_cosmological_velocity_unit(chem_vunit,a,cosmological_mhd_aref),
7248 radiation_workspace=radiation_workspace,
7249 radiation_photon_state=radiation_photon_state,
7250 radiation_drag_rate=drag_impulse /
7251 max(dt_hydro,eps(Float64)),
7252 neutral_hydrogen_storage=chem_neutral_hydrogen_storage,
7253 centered_hydrogen_storage=chem_centered_hydrogen_storage,
7254 hydrogen_reference=evolving_xe_mean,
7255 temperature_velocity_unit2=chem_mode === :off ? NaN :
7256 chem_vunit^2 * (isfinite(cosmological_mhd_aref) ?
7257 cosmological_mhd_aref / a : 1.0))
7258 KA.synchronize(be_mhd)
7259 KA.synchronize(be_p)
7260 @printf("FINAL_DIAG emitted at z=%.6f\n", a_to_z(a))
7261 flush(stdout)
7262 end
7263 _log_device_memory("final")
7264 memory_stats = get(ENV, "MHD_MEMPROBE", "0") == "1" ? _device_memory_gib() :
7265 (live=NaN, allocated=NaN, freed=NaN)
7266 total = sum(values(phase))
7267 mcells = ncells * measured / total / 1e6
7268 substep_mcells = ncells * measured_subcycles / total / 1e6
7269 br, divn, δb, δdm, chemstats =
7270 final_diag ? _final_device_stats!(be_p,s,ρdm,ρtot,HII,H2I,HeII,fh;
7271 density_perturbation=radiation_workspace === nothing ? nothing :
7272 radiation_workspace.density_perturbation,
7273 z=zrun ? a_to_z(a) : NaN,h=hub,Ob=Ob,
7274 dm_mean=dm_density_mean,
7275 neutral_hydrogen_storage=chem_neutral_hydrogen_storage,
7276 centered_hydrogen_storage=chem_centered_hydrogen_storage,
7277 hydrogen_reference=evolving_xe_mean) :
7278 (NaN, NaN, NaN, NaN, ntuple(_ -> NaN, 12))
7279 gravity || (δdm = NaN)
7280 divb_face_final = final_diag ? _device_divb_face_norm(s, br) : NaN
7281 _write_final_b_field!(get(ENV, "MHD_FINAL_B_DIR", ""), s;
7282 N=N, z=zrun ? a_to_z(a) : NaN,
7283 cycle=measured, brms=br)
7284 delta2k_final = final_diag && pmf_init === :magpressure ?
7285 (radiation_workspace === nothing ?
7286 _device_density_cos_mode_amp!(be_p,ρtot,s.U[1],N,pmf_kmode,2) :
7287 _device_density_cos_mode_amp!(be_p,ρtot,
7288 radiation_workspace.density_perturbation,N,pmf_kmode,2;
7289 mean=zero(T))) : NaN
7290
7291 @printf("\nMEASURED cycles=%d subcycles=%d cells=%.3fM total=%.4f s throughput=%.1f macro-Mcell/s substep=%.1f Mcell/s last_dt=%.3e\n",
7292 measured, measured_subcycles, ncells/1e6, total, mcells, substep_mcells, last_dt)
7293 @printf("EVOLVED measured_time=%.6e\n", elapsed_time)
7294 zrun && @printf("COSMO final_z=%.6f final_a=%.8e cycles=%d\n", a_to_z(a), a, measured)
7295 @printf("PHASE s/cycle: cfl %.4f | mhd %.4f | radiation %.4f | glm_project %.4f | drag %.4f | chem %.4f | deposit %.4f | assemble %.4f | fft %.4f | gas_gravity %.4f | push %.4f | total %.4f\n",
7296 phase[:cfl]/measured, phase[:mhd]/measured, phase[:radiation]/measured,
7297 phase[:glm_projection]/measured,
7298 phase[:drag]/measured, phase[:chem]/measured, phase[:deposit]/measured,
7299 phase[:assemble]/measured, phase[:fft]/measured, phase[:gas_gravity]/measured,
7300 phase[:push]/measured, total/measured)
7301 if chem_profile_enabled && measured > 0
7302 chem_profile_sum = sum(values(chem_phase))
7303 @printf("CHEM_PROFILE s/cycle: unpack %.5f | smooth %.5f | solve %.5f | nonlocal %.5f | clamp %.5f | mean %.5f | repack %.5f | sum %.5f\n",
7304 chem_phase[:unpack]/measured, chem_phase[:smooth]/measured,
7305 chem_phase[:solve]/measured, chem_phase[:nonlocal]/measured,
7306 chem_phase[:clamp]/measured,
7307 chem_phase[:mean]/measured, chem_phase[:repack]/measured,
7308 chem_profile_sum/measured)
7309 end
7310 @printf("RADIATION_PROFILE s/cycle: predictor %.5f | corrector %.5f | ledger %.5f | sum %.5f | fast_steps %d | spectral_steps %d | fallback_steps %d | plm_hll_fallback_steps %d | subdivision_steps %d | global_fallback_steps %d | density_limited_stages %d | density_min_theta %.6f | free_tol %.3e\n",
7311 radiation_phase[:predictor]/measured,
7312 radiation_phase[:corrector]/measured,
7313 radiation_phase[:ledger]/measured,
7314 phase[:radiation]/measured,radiation_step_timing.fast_steps,
7315 radiation_step_timing.spectral_steps,radiation_step_timing.fallback_steps,
7316 radiation_step_timing.plm_hll_fallback_steps,
7317 radiation_step_timing.subdivision_steps,
7318 radiation_step_timing.global_fallback_steps,
7319 radiation_step_timing.density_limited_stages,
7320 radiation_workspace===nothing ? 1.0 :
7321 radiation_workspace.density_limiter_min_theta[],
7322 radiation_free_stream_tolerance)
7323 @printf("SHARES: cfl %.0f%% | mhd %.0f%% | radiation %.0f%% | glm_project %.0f%% | drag %.0f%% | chem %.0f%% | deposit %.0f%% | assemble %.0f%% | fft %.0f%% | gas_gravity %.0f%% | push %.0f%%\n",
7324 100*phase[:cfl]/total, 100*phase[:mhd]/total, 100*phase[:radiation]/total,
7325 100*phase[:glm_projection]/total,
7326 100*phase[:drag]/total, 100*phase[:chem]/total, 100*phase[:deposit]/total,
7327 100*phase[:assemble]/total, 100*phase[:fft]/total,
7328 100*phase[:gas_gravity]/total, 100*phase[:push]/total)
7329 @printf("GLM: ch_fac=%.6g cr=%.6g project_every=%d projections=%d\n",
7330 glm_ch_fac, glm_cr, glm_project_every, glm_projection_steps)
7331 use_ct && @printf("MHD_SOLVER: ct (GLM ch_fac/cr unused; dt not GLM-limited; no B projection)\n")
7332 @printf("SOURCE_CADENCE: max_dt_factor=%.6g limited_steps=%d\n",
7333 source_max_dt_factor, source_limited_steps)
7334 terminal_denom = max(terminal_measured, 1)
7335 @printf("TERMINAL measured_cycles=%d s/terminal-cycle: force %.6f | pressure_fft %.6f | finalize %.6f | induction %.6f | projection %.6f\n",
7336 terminal_measured, terminal_phase[:force]/terminal_denom,
7337 terminal_phase[:pressure_fft]/terminal_denom, terminal_phase[:finalize]/terminal_denom,
7338 terminal_phase[:induction]/terminal_denom, terminal_phase[:projection]/terminal_denom)
7339 @printf("MEMORY device_live_gib=%.3f allocated_gib=%.3f freed_gib=%.3f (Metal allocator; use vmmap physical footprint for MPS/driver memory)\n",
7340 memory_stats.live, memory_stats.allocated, memory_stats.freed)
7341 @printf("DIAG: Brms %.3e -> %.3e | divB*dx/Brms %.3e | delta_b_rms %.3e | delta_dm_rms %.3e\n",
7342 st.brms_measured, br, divn, δb, δdm)
7343 use_ct && @printf("DIAG_CT: mhd_solver=ct divb_face*dx/Brms %.3e (exact face divergence; the centred divB above is a different diagnostic) | initial faces: cell_b_rel_l2_change=%.3e magnetic_energy_ratio=%.9f divb_face_dx_over_brms=%.3e\n",
7344 divb_face_final, ct_init_change, ct_init_energy_ratio, ct_init_divb_face)
7345 radiation_component_scale = radiation_energy_ledger === nothing ? NaN :
7346 abs(radiation_energy_ledger.kinetic_change) +
7347 abs(radiation_energy_ledger.internal_change) +
7348 abs(radiation_energy_ledger.magnetic_change)
7349 radiation_component_residual_norm = radiation_energy_ledger === nothing ? NaN :
7350 abs(radiation_energy_ledger.component_residual) /
7351 max(radiation_component_scale, eps(Float64))
7352 radiation_component_residual_abs_norm = radiation_energy_ledger === nothing ? NaN :
7353 radiation_energy_ledger.component_residual_abs /
7354 max(radiation_energy_ledger.component_magnitude_abs, eps(Float64))
7355 pmf_init === :magpressure && @printf("MODE: delta2k_amp %.8e\n", delta2k_final)
7356 @printf("GRAVITY_SOURCE_ACCURACY max_fast_slip_impulse=%.6g observed_max=%.9e (applies when gas gravity and drag are active)\n", gravity_drag_max_impulse,max_gravity_fast_impulse)
7357 drag_enabled && @printf("DRAG: mode=%s regime=%s boost=%.3f boosted_steps=%d subcycles_total=%d last_nsub=%d max_impulse=%.6e impulse_limited_steps=%d last_explicit_dt=%.3e last_dt=%.3e last_f=%.8f last_Gamma_over_H=%.6e rate_code=%.6e xe_mode=%s last_xe=%.6e terminal_gamma_scale=%.3e terminal_vcap=%.3e terminal_vmax=%.3e terminal_pressure_coeff=%.6e terminal_pressure_coeff_mode=%s induction_subcycles=%d last_induction_nsub=%d terminal_idiff_coeff=%.3e hybrid_va_cs=%.3e hybrid_drag_over_omega=%.3e handoff_events=%d handoff_steps=%d handoff_remaining=%d handoff_alpha_max=%.3f handoff_ramp=cosine handoff_alpha_last=%.3f handoff_error_last=%.6e handoff_error_max=%.6e handoff_error_settled_max=%.6e handoff_dv_rms=%.6e handoff_vt_rms=%.6e handoff_vf_rms=%.6e handoff_error_tol=%.6e handoff_correction_dcell_last=%.6e handoff_correction_dcell_max=%.6e handoff_correction_dcell_sum=%.6e handoff_correction_dcell_tol=%.6e\n",
7358 String(drag_dt_mode), last_hybrid_regime, drag_dt_boost, boosted_steps, subcycles_total, last_nsub,
7359 drag_max_impulse, drag_impulse_limited_steps,
7360 last_dt_explicit, last_dt, last_drag_f, last_drag_gammaH, drag_rate_code,
7361 drag_xe_env, last_drag_xe,
7362 terminal_gamma_scale, terminal_vcap, last_terminal_vmax,
7363 terminal_pressure_coeff_step,
7364 terminal_pressure_coeff_auto ? "auto" : "fixed",
7365 induction_subcycles_total, last_induction_nsub, last_terminal_diffuse_coeff,
7366 last_hybrid_va_cs, last_hybrid_drag_omega,
7367 handoff_events, handoff_steps_done, handoff_remaining,
7368 hybrid_handoff_alpha, last_handoff_alpha,
7369 last_handoff_error, max_handoff_error, max_handoff_settled_error,
7370 last_handoff_delta_v_rms,
7371 last_handoff_terminal_v_rms, last_handoff_full_v_rms,
7372 hybrid_handoff_error_tol, last_handoff_displacement,
7373 max_handoff_displacement, cumulative_handoff_displacement,
7374 hybrid_handoff_displacement_tol)
7375 if drag_dt_mode === :radiation
7376 rd = _radiation_scale_diagnostics(last_radiation_closure, N)
7377 @printf("RADIATION: response_model=%s R=%.8e lambda_over_box=%.8e q_fund=%.8e q_nyquist=%.8e stream_fund=%.8e stream_nyquist=%.8e crad_fund=%.8e crad_nyquist=%.8e gas_cs=%.8e bridge_q2=%.8e longitudinal_viscosity=%.8e max_dlna=%.8e highz_max_dlna=%.8e highz_zmin=%.8e nonlinear_cfl=%.8e magnetic_compensation=%s workspace_bytes=%d\n",
7378 String(radiation_response_model),
7379 last_radiation_closure.inertia_ratio,
7380 last_radiation_closure.mean_free_path,
7381 rd.qfund, rd.qnyquist, rd.stream_fund,
7382 rd.stream_nyquist, rd.crad_fund, rd.crad_nyquist,
7383 last_radiation_closure.gas_sound_speed,
7384 last_radiation_closure.bridge_q2,
7385 last_radiation_closure.longitudinal_viscosity,
7386 radiation_max_dlna,radiation_highz_max_dlna,
7387 radiation_highz_zmin,radiation_nonlinear_cfl,
7388 string(radiation_magnetic_compensation),
7389 radiation_workspace_bytes)
7390 if radiation_energy_ledger !== nothing
7391 @printf("RADIATION_ENERGY: samples=%d gas_total_change=%.8e kinetic_change=%.8e internal_change=%.8e magnetic_change=%.8e photon_change=%.8e resolved_coupled_change=%.8e component_residual=%.8e component_residual_abs=%.8e component_magnitude_abs=%.8e component_residual_norm=%.8e component_residual_abs_norm=%.8e photon_closure=%s\n",
7392 radiation_energy_ledger.samples,
7393 radiation_energy_ledger.total_change,
7394 radiation_energy_ledger.kinetic_change,
7395 radiation_energy_ledger.internal_change,
7396 radiation_energy_ledger.magnetic_change,
7397 radiation_energy_ledger.photon_change,
7398 radiation_energy_ledger.resolved_coupled_change,
7399 radiation_energy_ledger.component_residual,
7400 radiation_energy_ledger.component_residual_abs,
7401 radiation_energy_ledger.component_magnitude_abs,
7402 radiation_component_residual_norm,
7403 radiation_component_residual_abs_norm,
7404 radiation_response_model in (
7405 :moments,:photon_moments,:moments_exponential,
7406 :photon_moments_exponential) ?
7407 "moments" : "untracked")
7408 end
7409 end
7410 gravity && @printf("GRAVITY: source=%s gas_kick=%s weights=(gas %.6g, dm %.6g) midpoint_source=%s lattice_displacements=%s fixed_subcycles=%d dt_boost=%.3f particle_max_disp=%.6g last_particle_vmax=%.6e last_particle_gmax=%.6e subcycles_total=%d last_nsub=%d\n",
7411 String(gravity_source), string(gravity_gas_kick),
7412 Float64(grav_gas_weight), Float64(grav_dm_weight),
7413 string(gravity_midpoint_source), string(lattice_displacements),
7414 gravity_subcycles_fixed, gravity_dt_boost, particle_max_disp_cells,
7415 last_particle_vmax, last_particle_gmax, gravity_subcycles_total, last_gravity_nsub)
7416 chem_mode !== :off && @printf("CHEM: xHII mean/min/max %.6e %.6e %.6e | fH2 mean/min/max %.6e %.6e %.6e | xHeII mean/min/max %.6e %.6e %.6e | xe_total mean/min/max %.6e %.6e %.6e\n",
7417 chemstats...)
7418 flush(stdout)
7419
7420 out = get(ENV, "MHD_OUT_TSV", "")
7421 if !isempty(out)
7422 exists = isfile(out)
7423 open(out, "a") do io
7424 if !exists
7425 println(io, join(("timestamp","backend","N","recon","riemann","glm_ch_fac","glm_cr",
7426 "chem","chem_coupling","chem_expansion","chem_helium",
7427 "chem_rate_tables","chem_itcap","chem_dtfrac",
7428 "initial_p0","initial_p0_mode",
7429 "source_max_dt_factor","source_limited_steps",
7430 "hubble_h","omega_m","omega_l","omega_r",
7431 "pmf_kcut","pmf_kmax","pmf_hard_cutoff",
7432 "pmf_normalization_kind","pmf_normalization_z",
7433 "cosmological_mhd_scaling",
7434 "pmf_target_comoving_ng","pmf_initial_reference_comoving_ng",
7435 "pmf_initial_total_comoving_ng",
7436 "pmf_preserve_low_kmax","pmf_preserve_low_kcut",
7437 "pmf_mode_lock_n","pmf_high_mode_lock_n",
7438 "pmf_curl_symbol","pmf_curl_reference_n","pmf_fixed_mode_power",
7439 "pmf_mode_realization","pmf_realization_box",
7440 "pmf_seed","pmf_min_cells_per_wavelength",
7441 "boxsize","pk_nbins","pk_kmax",
7442 "glm_project_every","glm_projection_steps",
7443 "drag","drag_dt_mode","drag_dt_boost","drag_boosted_steps",
7444 "drag_subcycles_total","drag_last_nsub","drag_max_impulse","drag_impulse_limited_steps",
7445 "drag_last_f","drag_last_gamma_over_H","drag_xe",
7446 "drag_helium_electrons","drag_last_xe","drag_rate_code",
7447 "radiation_response_model","radiation_free_stream_tolerance",
7448 "radiation_fast_steps","radiation_spectral_steps",
7449 "radiation_fallback_steps","radiation_plm_hll_fallback_steps",
7450 "radiation_subdivision_steps","radiation_global_fallback_steps",
7451 "radiation_density_limited_stages","radiation_density_min_theta",
7452 "radiation_energy_samples",
7453 "radiation_total_change","radiation_kinetic_change",
7454 "radiation_internal_change","radiation_magnetic_change",
7455 "radiation_photon_change","radiation_resolved_coupled_change",
7456 "radiation_component_residual",
7457 "radiation_component_residual_abs",
7458 "radiation_component_magnitude_abs",
7459 "radiation_component_residual_norm",
7460 "radiation_component_residual_abs_norm",
7461 "gravity","gravity_gas_kick","gravity_source",
7462 "gravity_phi_interp","gravity_midpoint_source",
7463 "particle_lattice_displacements",
7464 "gravity_gas_weight","gravity_dm_weight",
7465 "gravity_subcycles_fixed","gravity_dt_boost","particle_max_disp_cells",
7466 "particle_last_vmax","particle_last_gmax","gravity_subcycles_total","gravity_last_nsub",
7467 "terminal_gamma_scale","terminal_vcap","terminal_diff_cfl",
7468 "terminal_accuracy_cfl","terminal_accuracy_pressure",
7469 "terminal_implicit_diffuse","terminal_implicit_diffuse_fac","terminal_idiff_coeff",
7470 "terminal_pressure_implicit","terminal_pressure_coeff",
7471 "terminal_pressure_coeff_mode",
7472 "terminal_rho_rel_limit","terminal_b_rel_limit",
7473 "terminal_max_disp_cells","terminal_last_vmax",
7474 "terminal_induction","terminal_induction_cfl","terminal_induction_dissipation",
7475 "terminal_induction_dissipation_order","terminal_project_b",
7476 "terminal_induction_max_subcycles","terminal_induction_subcycles_total",
7477 "terminal_induction_last_nsub",
7478 "hybrid_regime","hybrid_va_cs","hybrid_drag_over_omega",
7479 "hybrid_exit_drag_over_omega","hybrid_enter_drag_over_omega",
7480 "hybrid_check_every","hybrid_confirm_checks","hybrid_reenter_terminal",
7481 "hybrid_handoff_events","hybrid_handoff_steps_done","hybrid_handoff_remaining",
7482 "hybrid_handoff_alpha_max","hybrid_handoff_ramp","hybrid_handoff_alpha_last",
7483 "hybrid_handoff_error_last","hybrid_handoff_error_max",
7484 "hybrid_handoff_error_settled_max",
7485 "hybrid_handoff_delta_v_rms","hybrid_handoff_terminal_v_rms",
7486 "hybrid_handoff_full_v_rms","hybrid_handoff_error_tol",
7487 "hybrid_handoff_correction_dcell_last","hybrid_handoff_correction_dcell_max",
7488 "hybrid_handoff_correction_dcell_sum","hybrid_handoff_correction_dcell_tol",
7489 "brms0","brms","divBdx_over_brms","delta_b_rms","delta_dm_rms","delta2k_amp",
7490 "xHII_mean","xHII_min","xHII_max","fH2_mean","fH2_min","fH2_max",
7491 "xHeII_mean","xHeII_min","xHeII_max",
7492 "xe_total_mean","xe_total_min","xe_total_max",
7493 "cycles","measured_subcycles","measured_time","zstart","zend","final_z","total_s",
7494 "throughput_macro_mcells_s","throughput_substep_mcells_s",
7495 "last_dt","last_explicit_dt",
7496 "cfl_s_cyc","mhd_s_cyc","radiation_s_cyc",
7497 "glm_projection_s_cyc",
7498 "drag_s_cyc","chem_s_cyc","deposit_s_cyc","assemble_s_cyc",
7499 "fft_s_cyc","gas_gravity_s_cyc","push_s_cyc","terminal_measured_cycles",
7500 "terminal_force_s_cyc","terminal_pressure_fft_s_cyc","terminal_finalize_s_cyc",
7501 "terminal_induction_s_cyc","terminal_projection_s_cyc",
7502 "device_live_gib","device_allocated_gib","device_freed_gib",
7503 "mhd_solver","divb_face_dx_over_brms",
7504 "ct_initial_cell_b_rel_l2_change",
7505 "ct_initial_magnetic_energy_ratio",
7506 "ct_initial_divb_face_dx_over_brms"),
7507 '\t'))
7508 end
7509 println(io, join((Dates.format(now(), dateformat"yyyy-mm-ddTHH:MM:SS"),
7510 String(BACKEND_NAME), N, String(recon), String(riemann),
7511 glm_ch_fac, glm_cr, String(chem_mode), String(chem_coupling),
7512 chem_expansion, chem_helium,
7513 chem_rate_tables_enabled, chem_itcap,
7514 chem_dtfrac, p0,
7515 !isempty(initial_state_dir) ? "imported" :
7516 (p0_auto ? "cmb" : "fixed"),
7517 source_max_dt_factor,
7518 source_limited_steps, hub, Om, OL, Or,
7519 kcut, pmf_kmax === nothing ? "" : pmf_kmax, pmf_kmax !== nothing,
7520 st.normalization_kind, pmf_normalization_z,
7521 cosmological_mhd_scaling,
7522 pmf_b0_ng, pmf_initial_reference_comoving_ng,
7523 pmf_initial_total_comoving_ng,
7524 pmf_preserve_low_kmax === nothing ? "" : pmf_preserve_low_kmax,
7525 pmf_preserve_low_kcut === nothing ? "" : pmf_preserve_low_kcut,
7526 pmf_mode_lock_n === nothing ? 0 : pmf_mode_lock_n,
7527 pmf_high_mode_lock_n === nothing ? 0 : pmf_high_mode_lock_n,
7528 String(pmf_curl_symbol),
7529 pmf_curl_reference_n === nothing ? 0 : pmf_curl_reference_n,
7530 pmf_fixed_mode_power,
7531 String(pmf_mode_realization), pmf_realization_box,
7532 seed,
7533 pmf_min_cells_per_wavelength,
7534 boxsize, pk_nbins, pk_kmax === nothing ? "" : pk_kmax,
7535 glm_project_every, glm_projection_steps,
7536 drag_enabled, String(drag_dt_mode),
7537 drag_dt_boost, boosted_steps, subcycles_total, last_nsub,
7538 drag_max_impulse, drag_impulse_limited_steps,
7539 last_drag_f, last_drag_gammaH, drag_xe_env, drag_helium_electrons,
7540 last_drag_xe, drag_rate_code,
7541 String(radiation_response_model),
7542 radiation_free_stream_tolerance,
7543 radiation_step_timing.fast_steps,
7544 radiation_step_timing.spectral_steps,
7545 radiation_step_timing.fallback_steps,
7546 radiation_step_timing.plm_hll_fallback_steps,
7547 radiation_step_timing.subdivision_steps,
7548 radiation_step_timing.global_fallback_steps,
7549 radiation_step_timing.density_limited_stages,
7550 radiation_workspace === nothing ? 1.0 :
7551 radiation_workspace.density_limiter_min_theta[],
7552 radiation_energy_ledger === nothing ? 0 : radiation_energy_ledger.samples,
7553 radiation_energy_ledger === nothing ? NaN : radiation_energy_ledger.total_change,
7554 radiation_energy_ledger === nothing ? NaN : radiation_energy_ledger.kinetic_change,
7555 radiation_energy_ledger === nothing ? NaN : radiation_energy_ledger.internal_change,
7556 radiation_energy_ledger === nothing ? NaN : radiation_energy_ledger.magnetic_change,
7557 radiation_energy_ledger === nothing ? NaN : radiation_energy_ledger.photon_change,
7558 radiation_energy_ledger === nothing ? NaN : radiation_energy_ledger.resolved_coupled_change,
7559 radiation_energy_ledger === nothing ? NaN : radiation_energy_ledger.component_residual,
7560 radiation_energy_ledger === nothing ? NaN : radiation_energy_ledger.component_residual_abs,
7561 radiation_energy_ledger === nothing ? NaN : radiation_energy_ledger.component_magnitude_abs,
7562 radiation_component_residual_norm,
7563 radiation_component_residual_abs_norm,
7564 gravity, gravity_gas_kick, String(gravity_source),
7565 gravity_phi_interp, gravity_midpoint_source, lattice_displacements,
7566 grav_gas_weight, grav_dm_weight, gravity_subcycles_fixed,
7567 gravity_dt_boost, particle_max_disp_cells, last_particle_vmax,
7568 last_particle_gmax, gravity_subcycles_total, last_gravity_nsub,
7569 terminal_gamma_scale, terminal_vcap, terminal_diff_cfl,
7570 terminal_accuracy_cfl, terminal_accuracy_pressure,
7571 terminal_implicit_diffuse, terminal_implicit_diffuse_fac, last_terminal_diffuse_coeff,
7572 terminal_pressure_implicit, terminal_pressure_coeff_step,
7573 terminal_pressure_coeff_auto ? "auto" : "fixed",
7574 terminal_rho_rel_limit, terminal_b_rel_limit,
7575 terminal_max_disp_cells, last_terminal_vmax,
7576 String(terminal_induction), terminal_induction_cfl,
7577 terminal_induction_dissipation, terminal_induction_dissipation_order,
7578 terminal_project_b, terminal_induction_max_subcycles,
7579 induction_subcycles_total, last_induction_nsub,
7580 last_hybrid_regime, last_hybrid_va_cs, last_hybrid_drag_omega,
7581 hybrid_exit_drag_omega, hybrid_enter_drag_omega,
7582 hybrid_check_every, hybrid_confirm_checks, hybrid_reenter_terminal,
7583 handoff_events, handoff_steps_done, handoff_remaining,
7584 hybrid_handoff_alpha, "cosine", last_handoff_alpha,
7585 last_handoff_error, max_handoff_error, max_handoff_settled_error,
7586 last_handoff_delta_v_rms,
7587 last_handoff_terminal_v_rms, last_handoff_full_v_rms,
7588 hybrid_handoff_error_tol, last_handoff_displacement,
7589 max_handoff_displacement, cumulative_handoff_displacement,
7590 hybrid_handoff_displacement_tol,
7591 st.brms_measured, br, divn, δb, δdm, delta2k_final,
7592 chemstats...,
7593 measured, measured_subcycles, elapsed_time, zstart, zend, zrun ? a_to_z(a) : NaN,
7594 total, mcells, substep_mcells,
7595 last_dt, last_dt_explicit,
7596 phase[:cfl]/measured, phase[:mhd]/measured,
7597 phase[:radiation]/measured,
7598 phase[:glm_projection]/measured, phase[:drag]/measured, phase[:chem]/measured,
7599 phase[:deposit]/measured, phase[:assemble]/measured, phase[:fft]/measured,
7600 phase[:gas_gravity]/measured, phase[:push]/measured, terminal_measured,
7601 terminal_phase[:force]/terminal_denom,
7602 terminal_phase[:pressure_fft]/terminal_denom,
7603 terminal_phase[:finalize]/terminal_denom,
7604 terminal_phase[:induction]/terminal_denom,
7605 terminal_phase[:projection]/terminal_denom,
7606 memory_stats.live, memory_stats.allocated, memory_stats.freed,
7607 String(mhd_solver), divb_face_final, ct_init_change,
7608 ct_init_energy_ratio, ct_init_divb_face), '\t'))
7609 end
7610 end
7611 end
7612
7613 if abspath(PROGRAM_FILE) == @__FILE__
7614 main()
7615 end