Snapshot from this documentation build. SHA-256: 4ad1f7b69ef27a605099c4cc9c3fdda738a0eaf6061e55cf1f903c998fe9465f
1 # ── Primordial magnetic field initial conditions ─────────────────────────────
2 # Non-helical Gaussian PMFs with a Batchelor magnetic-energy spectrum. We build
3 # a random vector potential A(k), take B(k) = i s(k) × A(k), then inverse-transform
4 # B. Here s(k)=sin(k dx)/dx is the central-difference derivative symbol, so the
5 # diagnostic divergence uses the same operator and s·B vanishes to roundoff.
6
7 export PMFSpectrumSpec, pmf_batchelor_band, pmf_batchelor_field,
8 pmf_multiband_manifest, pmf_resolved_kmax,
9 init_pmf_batchelor!, init_pmf_batchelor_tiled!
10
11 """A resolution-independent definition of a physical Batchelor PMF spectrum."""
12 Base.@kwdef struct PMFSpectrumSpec
13 box::Float64
14 reference_brms::Float64
15 normalization_kind::Symbol = :reference_band_rms
16 normalization_redshift::Float64 = NaN
17 seed::Int = 1
18 kcut::Float64
19 reference_kmax::Float64
20 reference_kcut::Float64 = kcut
21 energy_index::Float64 = 4.0
22 cutoff_order::Float64 = 2.0
23 minimum_cells_per_wavelength::Float64 = 16.0
24 mode_lock_n::Int
25 curl_symbol::Symbol = :centered
26 end
27
28 function _pmf_normalization_kind(kind, preserve_low::Bool)
29 value = Symbol(lowercase(String(kind)))
30 value = value in (:total, :total_band) ? :total_band_rms : value
31 value = value in (:reference, :reference_band, :low_band) ?
32 :reference_band_rms : value
33 value === :auto && return preserve_low ? :reference_band_rms : :total_band_rms
34 value in (:total_band_rms, :reference_band_rms) || error(
35 "normalization_kind must be :total_band_rms or :reference_band_rms")
36 value === :reference_band_rms && !preserve_low && error(
37 "reference-band normalization requires preserve_low_kmax")
38 return value
39 end
40
41 """Largest physical wavenumber resolved at the requested cells per wavelength."""
42 function pmf_resolved_kmax(dims::NTuple{3,Int}, box::Real,
43 minimum_cells_per_wavelength::Real)
44 cells = Float64(minimum_cells_per_wavelength)
45 cells > 0 || error("minimum_cells_per_wavelength must be positive")
46 return 2π * minimum(dims) / (Float64(box) * cells)
47 end
48
49 """
50 pmf_multiband_manifest(spec, edges)
51
52 Describe disjoint physical wavenumber bands and the smallest even cubic grid
53 that resolves each upper edge. Every band uses the same reference low modes
54 and normalization from `spec`.
55 """
56 function pmf_multiband_manifest(spec::PMFSpectrumSpec, edges)
57 ks = Float64.(collect(edges))
58 length(ks) >= 2 || error("multiband edges require at least two values")
59 first(ks) >= 0 || error("multiband edges must be nonnegative")
60 all(diff(ks) .> 0) || error("multiband edges must be strictly increasing")
61 spec.reference_kmax <= ks[2] ||
62 error("the first band must contain reference_kmax=$(spec.reference_kmax)")
63 return [begin
64 raw_n = ceil(Int, khi * spec.box * spec.minimum_cells_per_wavelength / (2π))
65 required_n = iseven(raw_n) ? raw_n : raw_n + 1
66 (band=i, kmin=klo, kmax=khi, required_n=required_n,
67 cells_per_wavelength=2π * required_n / (spec.box * khi),
68 box=spec.box, reference_brms=spec.reference_brms,
69 normalization_kind=spec.normalization_kind,
70 normalization_redshift=spec.normalization_redshift,
71 reference_kmax=spec.reference_kmax, seed=spec.seed,
72 mode_lock_n=spec.mode_lock_n)
73 end for (i, (klo, khi)) in enumerate(zip(ks[1:end-1], ks[2:end]))]
74 end
75
76 """Generate one band from a resolution-independent PMF spectrum definition."""
77 function pmf_batchelor_band(spec::PMFSpectrumSpec, dims::NTuple{3,Int};
78 kmin::Real=0, kmax::Real=spec.reference_kmax)
79 minimum(dims) >= spec.mode_lock_n ||
80 error("grid $dims is smaller than mode_lock_n=$(spec.mode_lock_n)")
81 return pmf_batchelor_field(dims; box=spec.box, brms=spec.reference_brms,
82 seed=spec.seed, kcut=spec.kcut, kmin=kmin, kmax=kmax,
83 energy_index=spec.energy_index, cutoff_order=spec.cutoff_order,
84 mode_lock_n=spec.mode_lock_n, preserve_low_kmax=spec.reference_kmax,
85 preserve_low_kcut=spec.reference_kcut, curl_symbol=spec.curl_symbol,
86 normalization_kind=spec.normalization_kind,
87 minimum_cells_per_wavelength=spec.minimum_cells_per_wavelength)
88 end
89
90 @inline _fftfreq_i(i::Int, n::Int, L::Float64) =
91 (i - 1 <= n ÷ 2 ? i - 1 : i - 1 - n) * (2π / L)
92
93 @inline _rfftfreq_i(i::Int, n::Int, L::Float64) = (i - 1) * (2π / L)
94
95 @inline _pmf_mode_i(i::Int, n::Int) = i - 1 <= n ÷ 2 ? i - 1 : i - 1 - n
96
97 function _random_vector_potential_k(dims::NTuple{3,Int}, seed::Integer)
98 nx, ny, nz = dims
99 rng = MersenneTwister(seed)
100 a = randn(rng, Float32, nx, ny, nz)
101 ax = FFTW.rfft(a, (1,2,3))
102 randn!(rng, a)
103 ay = FFTW.rfft(a, (1,2,3))
104 randn!(rng, a)
105 az = FFTW.rfft(a, (1,2,3))
106 return ax, ay, az
107 end
108
109 @inline function _pmf_mix64(x::UInt64)
110 z = x + 0x9e3779b97f4a7c15
111 z = xor(z, z >> 30) * 0xbf58476d1ce4e5b9
112 z = xor(z, z >> 27) * 0x94d049bb133111eb
113 return xor(z, z >> 31)
114 end
115
116 @inline function _pmf_hash_word(seed::Integer, qx::Int, qy::Int, qz::Int,
117 component::Int, lane::Int)
118 h = _pmf_mix64(reinterpret(UInt64, Int64(seed)))
119 for value in (qx, qy, qz, component, lane)
120 h = _pmf_mix64(xor(h, reinterpret(UInt64, Int64(value))))
121 end
122 return h
123 end
124
125 @inline function _pmf_hash_uniform(seed::Integer, qx::Int, qy::Int, qz::Int,
126 component::Int, lane::Int)
127 word = _pmf_hash_word(seed, qx, qy, qz, component, lane)
128 return (Float64(word >> 11) + 0.5) * 0x1.0p-53
129 end
130
131 @inline function _pmf_hash_normal(seed::Integer, qx::Int, qy::Int, qz::Int,
132 component::Int; real_only::Bool=false)
133 u1 = _pmf_hash_uniform(seed, qx, qy, qz, component, 1)
134 u2 = _pmf_hash_uniform(seed, qx, qy, qz, component, 2)
135 radius = sqrt(-2 * log(u1))
136 if real_only
137 return ComplexF32(Float32(radius * cospi(2u2)), 0)
138 end
139 phase = 2π * u2
140 return ComplexF32(Float32(radius * cos(phase)), Float32(radius * sin(phase)))
141 end
142
143 """
144 Build a vector-potential realization whose random coefficients are keyed by
145 physical Fourier indices rather than array traversal. Boxes whose lengths are
146 integer subdivisions of `reference_box` therefore receive identical phase and
147 polarization on every shared, non-Nyquist physical mode.
148 """
149 function _physical_hash_vector_potential_k(dims::NTuple{3,Int}, seed::Integer,
150 box::Float64, reference_box::Float64;
151 zero_phase::Bool=false)
152 nx, ny, nz = dims
153 ratio = reference_box / box
154 qscale = round(Int, ratio)
155 qscale > 0 && isapprox(ratio, qscale; rtol=0,
156 atol=32eps(Float64) * max(1, abs(ratio))) || error(
157 "physical-hash reference_box/box must be a positive integer; got " *
158 "reference_box=$reference_box box=$box")
159 nxh = nx ÷ 2 + 1
160 fields = ntuple(_ -> Array{ComplexF32}(undef, nxh, ny, nz), 3)
161 @inbounds for kk in 1:nz, jj in 1:ny, ii in 1:nxh
162 mx = ii - 1
163 my = _pmf_mode_i(jj, ny)
164 mz = _pmf_mode_i(kk, nz)
165 boundary_x = ii == 1 || ii == nxh
166 conjugate_mode = false
167 real_only = false
168 key_my = my
169 key_mz = mz
170 if boundary_x
171 pair_j = mod(-my, ny) + 1
172 pair_k = mod(-mz, nz) + 1
173 if (pair_j, pair_k) < (jj, kk)
174 conjugate_mode = true
175 key_my = _pmf_mode_i(pair_j, ny)
176 key_mz = _pmf_mode_i(pair_k, nz)
177 elseif pair_j == jj && pair_k == kk
178 real_only = true
179 end
180 end
181 qx = qscale * mx
182 qy = qscale * key_my
183 qz = qscale * key_mz
184 for component in 1:3
185 value = _pmf_hash_normal(seed, qx, qy, qz, component;
186 real_only=zero_phase || real_only)
187 fields[component][ii,jj,kk] = conjugate_mode ? conj(value) : value
188 end
189 end
190 return fields
191 end
192
193 function _vector_potential_k(dims::NTuple{3,Int}, seed::Integer, box::Float64,
194 mode_realization::Symbol, realization_box::Float64)
195 mode_realization === :legacy && return _random_vector_potential_k(dims, seed)
196 mode_realization === :physical_hash && return _physical_hash_vector_potential_k(
197 dims, seed, box, realization_box)
198 mode_realization === :physical_hash_zero_phase &&
199 return _physical_hash_vector_potential_k(
200 dims, seed, box, realization_box; zero_phase=true)
201 error("mode_realization must be :legacy, :physical_hash, or " *
202 ":physical_hash_zero_phase")
203 end
204
205 function _rfft_zero_pad_complex(ak::Array{ComplexF32,3}, source_dims::NTuple{3,Int},
206 target_dims::NTuple{3,Int})
207 ox, oy, oz = source_dims
208 nx, ny, nz = target_dims
209 out = zeros(ComplexF32, nx ÷ 2 + 1, ny, nz)
210 scale = Float32((nx * ny * nz) / (ox * oy * oz))
211 @inbounds for kk in 1:oz
212 mz = _pmf_mode_i(kk, oz)
213 tk = mz >= 0 ? mz + 1 : nz + mz + 1
214 for jj in 1:oy
215 my = _pmf_mode_i(jj, oy)
216 tj = my >= 0 ? my + 1 : ny + my + 1
217 for ii in 1:(ox ÷ 2 + 1)
218 mx = ii - 1
219 out[mx + 1, tj, tk] = scale * ak[ii,jj,kk]
220 end
221 end
222 end
223 return out
224 end
225
226 function _periodic_tile_field(a::Array{Float32,3}, dims::NTuple{3,Int})
227 sx, sy, sz = size(a)
228 nx, ny, nz = dims
229 nx % sx == 0 || error("tile x-size $sx must divide target x-size $nx")
230 ny % sy == 0 || error("tile y-size $sy must divide target y-size $ny")
231 nz % sz == 0 || error("tile z-size $sz must divide target z-size $nz")
232 out = Array{Float32}(undef, dims)
233 @inbounds for k in 1:nz, j in 1:ny, i in 1:nx
234 out[i,j,k] = a[mod1(i, sx), mod1(j, sy), mod1(k, sz)]
235 end
236 return out
237 end
238
239 function _install_pmf_field!(s::MHDState{T}, Bx::Array{Float32,3}, By::Array{Float32,3},
240 Bz::Array{Float32,3}; rho0::Real = 1,
241 p0::Real = 1) where {T}
242 nc = ncells(s)
243 ρ = T(rho0); p = T(p0)
244 Ehost = Vector{T}(undef, nc)
245 Bxv = vec(Bx); Byv = vec(By); Bzv = vec(Bz)
246 Threads.@threads for idx in 1:nc
247 @inbounds begin
248 bx = T(Bxv[idx]); by = T(Byv[idx]); bz = T(Bzv[idx])
249 Ehost[idx] = p/(s.γ - one(T)) + T(0.5) * (bx*bx + by*by + bz*bz)
250 end
251 end
252 fill!(s.U[1], ρ)
253 fill!(s.U[2], zero(T)); fill!(s.U[3], zero(T)); fill!(s.U[4], zero(T))
254 copyto!(s.U[5], Ehost)
255 copyto!(s.U[6], Bxv); copyto!(s.U[7], Byv); copyto!(s.U[8], Bzv)
256 fill!(s.U[9], zero(T))
257 KA.synchronize(s.be)
258 # CT states need a discretely divergence-free FACE field: derive it from the cell
259 # field just installed (spectral half-cell shift + discrete projection).
260 return uses_ct(s) ? ct_faces_from_cells!(s) : nothing
261 end
262
263 function _pmf_fft_threads()
264 env = get(ENV, "MHD_PMF_FFT_THREADS", get(ENV, "CIC_FFT_THREADS", ""))
265 n = isempty(strip(env)) ? min(Sys.CPU_THREADS, 4) : parse(Int, env)
266 return max(1, n)
267 end
268
269 function _pmf_kspace_threads()
270 get(ENV, "MHD_PMF_KSPACE_THREADS", "0") in ("1", "true", "TRUE", "yes", "on")
271 end
272
273 """
274 pmf_batchelor_field(dims; box=1, brms=1, seed=1, kcut=nothing, kmax=nothing,
275 energy_index=4, cutoff_order=2, normalization_kind=:auto,
276 fixed_mode_power=false, mode_realization=:legacy)
277
278 Generate a non-helical, approximately Batchelor-spectrum magnetic field on a
279 periodic grid. `energy_index=4` means shell magnetic energy `E_B(k) ∝ k^4`
280 below the exponential `kcut` taper. Set `kmax` to impose an additional strict
281 physical band limit; every Fourier coefficient with `|k| > kmax` is exactly
282 zero before the inverse transform. `normalization_kind=:total_band_rms` makes
283 `brms` the target `sqrt(<B^2>)` over the generated band.
284 `normalization_kind=:reference_band_rms` instead normalizes the preserved
285 `0 < k <= preserve_low_kmax` reference band, so adding ultraviolet modes raises
286 the total rms. `:auto` retains the legacy choice: reference-band normalization
287 when `preserve_low_kmax` is present and total-band normalization otherwise.
288 The returned tuple is `(Bx, By, Bz, stats)`.
289
290 With `fixed_mode_power=true`, each nonzero Fourier mode is rescaled to the
291 analytic envelope while retaining its random phase and transverse direction.
292 This removes finite-volume amplitude scatter from matched-box spectrum tests.
293 With `mode_realization=:physical_hash`, phase and polarization are keyed by
294 physical mode number relative to `realization_box`; this is intended for
295 matched-box and fixed-mode convergence studies.
296 `mode_realization=:physical_hash_zero_phase` uses the same physical keys but
297 real Fourier coefficients, matching the zero-phase non-helical construction
298 used by Jedamzik, Abel, and Ali-Haimoud for their recombination calculations.
299 """
300 function pmf_batchelor_field(dims::NTuple{3,Int}; box::Real = 1,
301 brms::Real = 1, seed::Integer = 1,
302 kcut = nothing, kmin::Real = 0,
303 kmax = nothing, energy_index::Real = 4,
304 cutoff_order::Real = 2, mode_lock_n = nothing,
305 high_mode_lock_n = nothing,
306 preserve_low_kmax = nothing,
307 preserve_low_kcut = nothing,
308 normalization_kind = :auto,
309 curl_symbol = :centered,
310 curl_reference_n = nothing,
311 fixed_mode_power::Bool = false,
312 mode_realization = :legacy,
313 realization_box = nothing,
314 minimum_cells_per_wavelength = nothing,
315 invariant_axis::Integer = 0)
316 nx, ny, nz = dims
317 invariant_axis in 0:3 || error("invariant_axis must be 0, 1, 2, or 3")
318 iseven(nx) || error("pmf_batchelor_field currently requires even nx for rfft packing")
319 L = Float64(box)
320 target = Float64(brms)
321 target > 0 || error("pmf_batchelor_field requires brms > 0")
322 realization_mode = Symbol(lowercase(String(mode_realization)))
323 realization_mode in (:legacy, :physical_hash, :physical_hash_zero_phase) ||
324 error("mode_realization must be :legacy, :physical_hash, or " *
325 ":physical_hash_zero_phase")
326 realization_L = realization_box === nothing ? L : Float64(realization_box)
327 realization_L > 0 || error("realization_box must be positive")
328 realization_L + 32eps(Float64) * max(1, realization_L) >= L || error(
329 "realization_box cannot be smaller than box")
330 kc = kcut === nothing ? (π * minimum(dims) / L) : Float64(kcut)
331 kc > 0 || error("kcut must be positive")
332 kmin_f = Float64(kmin)
333 kmin_f >= 0 || error("kmin must be nonnegative")
334 km = kmax === nothing ? Inf : Float64(kmax)
335 km > 0 || error("kmax must be positive")
336 kmin_f < km || error("kmin must be smaller than kmax")
337 resolved_kmax = if minimum_cells_per_wavelength === nothing
338 Inf
339 else
340 isfinite(km) || error("kmax is required when enforcing minimum cells per wavelength")
341 resolved = pmf_resolved_kmax(dims, L, minimum_cells_per_wavelength)
342 km <= resolved * (1 + 16eps(Float64)) || error(
343 "kmax=$km is under-resolved on grid $dims: limit is $resolved " *
344 "for $(minimum_cells_per_wavelength) cells per wavelength")
345 resolved
346 end
347 low_km = preserve_low_kmax === nothing ? nothing : Float64(preserve_low_kmax)
348 preserve_low = low_km !== nothing
349 if preserve_low
350 low_km > 0 || error("preserve_low_kmax must be positive")
351 low_km <= km || error("preserve_low_kmax cannot exceed kmax")
352 end
353 low_kc = preserve_low_kcut === nothing ? kc : Float64(preserve_low_kcut)
354 low_kc > 0 || error("preserve_low_kcut must be positive")
355 norm_kind = _pmf_normalization_kind(normalization_kind, preserve_low)
356 source_lock_n = high_mode_lock_n === nothing ?
357 (!preserve_low ? mode_lock_n : nothing) : high_mode_lock_n
358 source_dims = dims
359 if source_lock_n !== nothing && Int(source_lock_n) != minimum(dims)
360 lock_n = Int(source_lock_n)
361 lock_n > 0 || error("mode_lock_n must be positive")
362 lock_n <= minimum(dims) ||
363 error("high_mode_lock_n=$lock_n cannot exceed target dims $dims")
364 km <= π * lock_n / L ||
365 error("PMF band exceeds the high-mode-lock Nyquist frequency")
366 source_dims = (lock_n, lock_n, lock_n)
367 end
368 amp_exp = 0.5 * (Float64(energy_index) - 4.0)
369 curl_mode = Symbol(curl_symbol)
370 curl_mode in (:centered, :continuum, :spectral) ||
371 error("curl_symbol must be :centered or :continuum")
372 curl_mode === :spectral && (curl_mode = :continuum)
373 curl_ref_n = curl_reference_n === nothing ? 0 : Int(curl_reference_n)
374 if curl_ref_n != 0
375 curl_mode === :centered ||
376 error("curl_reference_n is only supported with centered curl")
377 0 < curl_ref_n <= minimum(dims) ||
378 error("curl_reference_n must be positive and no larger than target dims")
379 end
380 FFTW.set_num_threads(_pmf_fft_threads())
381 nxh = nx ÷ 2 + 1
382 axk, ayk, azk = begin
383 ax, ay, az = _vector_potential_k(source_dims, seed, L, realization_mode,
384 realization_L)
385 source_dims == dims ? (ax, ay, az) :
386 (_rfft_zero_pad_complex(ax, source_dims, dims),
387 _rfft_zero_pad_complex(ay, source_dims, dims),
388 _rfft_zero_pad_complex(az, source_dims, dims))
389 end
390 low_source_dims = dims
391 low_axk, low_ayk, low_azk = if preserve_low
392 if mode_lock_n !== nothing && Int(mode_lock_n) != minimum(dims)
393 lock_n = Int(mode_lock_n)
394 lock_n > 0 || error("mode_lock_n must be positive")
395 lock_n <= minimum(dims) ||
396 error("mode_lock_n=$lock_n cannot exceed target dims $dims")
397 low_km <= π * lock_n / L ||
398 error("preserved low-k band exceeds the mode-lock Nyquist frequency")
399 low_source_dims = (lock_n, lock_n, lock_n)
400 end
401 ax, ay, az = _vector_potential_k(low_source_dims, seed, L,
402 realization_mode, realization_L)
403 low_source_dims == dims ? (ax, ay, az) :
404 (_rfft_zero_pad_complex(ax, low_source_dims, dims),
405 _rfft_zero_pad_complex(ay, low_source_dims, dims),
406 _rfft_zero_pad_complex(az, low_source_dims, dims))
407 else
408 (axk, ayk, azk)
409 end
410 # Zero-padding multiplies each source grid by Vtarget/Vsource. Match the
411 # high-source white-noise variance to the embedded low-source realization
412 # before joining the two spectral bands.
413 high_source_scale = preserve_low ?
414 Float32(sqrt(prod(source_dims) / prod(low_source_dims))) : 1.0f0
415 bxk = similar(axk)
416 byk = similar(axk)
417 bzk = similar(axk)
418 low_bxk = preserve_low ? zeros(eltype(bxk), size(bxk)) : bxk
419 low_byk = preserve_low ? zeros(eltype(byk), size(byk)) : byk
420 low_bzk = preserve_low ? zeros(eltype(bzk), size(bzk)) : bzk
421 dx = L / nx
422 im = ComplexF32(0, 1)
423 kxv = [_rfftfreq_i(ii, nx, L) for ii in 1:nxh]
424 kyv = [_fftfreq_i(jj, ny, L) for jj in 1:ny]
425 kzv = [_fftfreq_i(kk, nz, L) for kk in 1:nz]
426 sxv = Float32[curl_mode === :centered ? sin(kx * dx) / dx : kx for kx in kxv]
427 syv = Float32[curl_mode === :centered ? sin(ky * dx) / dx : ky for ky in kyv]
428 szv = Float32[curl_mode === :centered ? sin(kz * dx) / dx : kz for kz in kzv]
429 reference_dx = curl_ref_n == 0 ? dx : L / curl_ref_n
430 rxv = Float32[curl_mode === :centered ? sin(kx * reference_dx) / reference_dx : kx
431 for kx in kxv]
432 ryv = Float32[curl_mode === :centered ? sin(ky * reference_dx) / reference_dx : ky
433 for ky in kyv]
434 rzv = Float32[curl_mode === :centered ? sin(kz * reference_dx) / reference_dx : kz
435 for kz in kzv]
436 project_reference_curl = curl_ref_n != 0 && curl_ref_n != minimum(dims)
437 function fill_bk_slice!(kk)
438 @inbounds for jj in 1:ny, ii in 1:nxh
439 kx = kxv[ii]
440 ky = kyv[jj]
441 kz = kzv[kk]
442 k2 = kx*kx + ky*ky + kz*kz
443 if k2 == 0 || k2 > km * km ||
444 (invariant_axis != 0 && (kx,ky,kz)[invariant_axis] != 0)
445 bxk[ii,jj,kk] = 0
446 byk[ii,jj,kk] = 0
447 bzk[ii,jj,kk] = 0
448 else
449 kmag = sqrt(k2)
450 # k-space coordinates are Float32. Treat a cutoff within a few
451 # Float32 ulps of a lattice shell as inclusive, as documented.
452 use_low = preserve_low &&
453 k2 <= Float32(low_km)^2 * (1f0 + 8f0 * eps(Float32))
454 local_kc = use_low ? low_kc : kc
455 amp = (kmag / local_kc)^amp_exp *
456 exp(-0.5 * (kmag / local_kc)^Float64(cutoff_order))
457 src_ax = use_low ? low_axk[ii,jj,kk] : axk[ii,jj,kk]
458 src_ay = use_low ? low_ayk[ii,jj,kk] : ayk[ii,jj,kk]
459 src_az = use_low ? low_azk[ii,jj,kk] : azk[ii,jj,kk]
460 source_scale = use_low ? 1.0f0 : high_source_scale
461 weighted_ax = ComplexF32(amp * source_scale) * src_ax
462 weighted_ay = ComplexF32(amp * source_scale) * src_ay
463 weighted_az = ComplexF32(amp * source_scale) * src_az
464 sx = sxv[ii]
465 sy = syv[jj]
466 sz = szv[kk]
467 rx = rxv[ii]
468 ry = ryv[jj]
469 rz = rzv[kk]
470 bx = im * (ry * weighted_az - rz * weighted_ay)
471 by = im * (rz * weighted_ax - rx * weighted_az)
472 bz = im * (rx * weighted_ay - ry * weighted_ax)
473 if project_reference_curl
474 s2 = sx*sx + sy*sy + sz*sz
475 if s2 > 0
476 parallel = (sx*bx + sy*by + sz*bz) / s2
477 bx -= sx * parallel
478 by -= sy * parallel
479 bz -= sz * parallel
480 end
481 end
482 bxk[ii,jj,kk] = bx
483 byk[ii,jj,kk] = by
484 bzk[ii,jj,kk] = bz
485 if use_low
486 low_bxk[ii,jj,kk] = bx
487 low_byk[ii,jj,kk] = by
488 low_bzk[ii,jj,kk] = bz
489 end
490 if kmin_f > 0 && k2 <= kmin_f * kmin_f
491 bxk[ii,jj,kk] = 0
492 byk[ii,jj,kk] = 0
493 bzk[ii,jj,kk] = 0
494 end
495 end
496 end
497 end
498 if _pmf_kspace_threads()
499 Threads.@threads for kk in 1:nz
500 fill_bk_slice!(kk)
501 end
502 else
503 for kk in 1:nz
504 fill_bk_slice!(kk)
505 end
506 end
507 if fixed_mode_power
508 function fix_mode_power_slice!(kk)
509 @inbounds for jj in 1:ny, ii in 1:nxh
510 kx = kxv[ii]
511 ky = kyv[jj]
512 kz = kzv[kk]
513 k2 = kx*kx + ky*ky + kz*kz
514 (k2 == 0 || k2 > km * km || k2 <= kmin_f * kmin_f) && continue
515 raw = abs2(bxk[ii,jj,kk]) + abs2(byk[ii,jj,kk]) +
516 abs2(bzk[ii,jj,kk])
517 raw > 0 || continue
518 kmag = sqrt(k2)
519 use_low = preserve_low &&
520 k2 <= Float32(low_km)^2 * (1f0 + 8f0 * eps(Float32))
521 local_kc = use_low ? low_kc : kc
522 target_power = (kmag / local_kc)^(Float64(energy_index) - 2.0) *
523 exp(-(kmag / local_kc)^Float64(cutoff_order))
524 scale_mode = Float32(sqrt(target_power / raw))
525 bxk[ii,jj,kk] *= scale_mode
526 byk[ii,jj,kk] *= scale_mode
527 bzk[ii,jj,kk] *= scale_mode
528 if use_low
529 low_bxk[ii,jj,kk] *= scale_mode
530 low_byk[ii,jj,kk] *= scale_mode
531 low_bzk[ii,jj,kk] *= scale_mode
532 end
533 end
534 end
535 if _pmf_kspace_threads()
536 Threads.@threads for kk in 1:nz
537 fix_mode_power_slice!(kk)
538 end
539 else
540 for kk in 1:nz
541 fix_mode_power_slice!(kk)
542 end
543 end
544 end
545 Bx = Array{Float32}(FFTW.irfft(bxk, nx, (1,2,3)))
546 By = Array{Float32}(FFTW.irfft(byk, nx, (1,2,3)))
547 Bz = Array{Float32}(FFTW.irfft(bzk, nx, (1,2,3)))
548 rms0 = sqrt((sum(abs2, Bx) + sum(abs2, By) + sum(abs2, Bz)) / length(Bx))
549 rms0 > 0 || error("generated zero PMF field; check kcut, kmax, and dims")
550 reference_rms0 = if preserve_low
551 lowsum = sum(abs2, FFTW.irfft(low_bxk, nx, (1,2,3))) +
552 sum(abs2, FFTW.irfft(low_byk, nx, (1,2,3))) +
553 sum(abs2, FFTW.irfft(low_bzk, nx, (1,2,3)))
554 low_rms0 = sqrt(lowsum / length(Bx))
555 low_rms0 > 0 || error("preserved low-k PMF field is zero")
556 low_rms0
557 else
558 rms0
559 end
560 norm_rms0 = norm_kind === :reference_band_rms ? reference_rms0 : rms0
561 scale = Float32(target / norm_rms0)
562 Bx .*= scale; By .*= scale; Bz .*= scale
563 total_brms = sqrt((sum(abs2, Bx) + sum(abs2, By) + sum(abs2, Bz)) / length(Bx))
564 reference_brms = Float64(scale) * reference_rms0
565 normalization_brms = norm_kind === :reference_band_rms ? reference_brms : total_brms
566 stats = (brms = target,
567 brms_measured = total_brms,
568 total_brms_measured = total_brms,
569 reference_band_brms_measured = reference_brms,
570 normalization_brms_measured = normalization_brms,
571 normalization_kind = String(norm_kind),
572 box = L, kcut = kc, energy_index = Float64(energy_index), seed = Int(seed),
573 kmin = kmin_f, kmax = km,
574 minimum_cells_per_wavelength = minimum_cells_per_wavelength === nothing ?
575 NaN : Float64(minimum_cells_per_wavelength),
576 actual_cells_per_wavelength = isfinite(km) ? 2π * minimum(dims) / (L * km) : NaN,
577 resolved_kmax = resolved_kmax,
578 mode_lock_n = mode_lock_n === nothing ? 0 : Int(mode_lock_n),
579 high_mode_lock_n = high_mode_lock_n === nothing ? 0 : Int(high_mode_lock_n),
580 curl_reference_n = curl_ref_n,
581 curl_symbol = String(curl_mode),
582 fixed_mode_power = fixed_mode_power,
583 invariant_axis = Int(invariant_axis),
584 mode_realization = String(realization_mode),
585 realization_box = realization_L,
586 preserve_low_kmax = preserve_low ? low_km : NaN,
587 preserve_low_kcut = preserve_low ? low_kc : NaN,
588 preserve_high_source_scale = preserve_low ? Float64(high_source_scale) : 1.0,
589 preserve_low_norm_rms0 = preserve_low ? reference_rms0 : NaN)
590 return Bx, By, Bz, stats
591 end
592
593 """
594 init_pmf_batchelor!(s; brms=nothing, b0_pG=nothing, Bunit_pG=1,
595 rho0=1, p0=1, box=Nx*dx, seed=1, kcut=nothing,
596 kmax=nothing, energy_index=4)
597
598 Initialize `s` with uniform gas and a non-helical Batchelor PMF. Use `brms`
599 directly in code units, or provide `b0_pG` with `Bunit_pG` to map present-day
600 comoving pico-Gauss into code units. Velocities start at zero and `ψ=0`.
601 """
602 function init_pmf_batchelor!(s::MHDState{T}; brms = nothing, b0_pG = nothing,
603 Bunit_pG::Real = 1, rho0::Real = 1, p0::Real = 1,
604 box::Real = s.dims[1] * s.dx, seed::Integer = 1,
605 kcut = nothing, kmin::Real = 0,
606 kmax = nothing, energy_index::Real = 4,
607 cutoff_order::Real = 2, mode_lock_n = nothing,
608 high_mode_lock_n = nothing,
609 preserve_low_kmax = nothing,
610 preserve_low_kcut = nothing,
611 normalization_kind = :auto,
612 curl_symbol = :centered,
613 curl_reference_n = nothing,
614 fixed_mode_power::Bool = false,
615 mode_realization = :legacy,
616 realization_box = nothing,
617 minimum_cells_per_wavelength = nothing,
618 invariant_axis::Integer = 0) where {T}
619 br = brms === nothing ? begin
620 b0_pG === nothing && error("provide either brms code units or b0_pG")
621 Float64(b0_pG) / Float64(Bunit_pG)
622 end : Float64(brms)
623 Bx, By, Bz, stats = pmf_batchelor_field(s.dims; box=box, brms=br, seed=seed,
624 kcut=kcut, kmin=kmin, kmax=kmax,
625 energy_index=energy_index,
626 cutoff_order=cutoff_order,
627 mode_lock_n=mode_lock_n,
628 high_mode_lock_n=high_mode_lock_n,
629 preserve_low_kmax=preserve_low_kmax,
630 preserve_low_kcut=preserve_low_kcut,
631 normalization_kind=normalization_kind,
632 curl_symbol=curl_symbol,
633 curl_reference_n=curl_reference_n,
634 fixed_mode_power=fixed_mode_power,
635 mode_realization=mode_realization,
636 realization_box=realization_box,
637 minimum_cells_per_wavelength=minimum_cells_per_wavelength,
638 invariant_axis=invariant_axis)
639 ctinfo = _install_pmf_field!(s, Bx, By, Bz; rho0=rho0, p0=p0)
640 out = merge(stats, (rho0=Float64(rho0), p0=Float64(p0), Bunit_pG=Float64(Bunit_pG)))
641 return ctinfo === nothing ? out : merge(out, (ct_faces=ctinfo,))
642 end
643
644 """
645 init_pmf_batchelor_tiled!(s; tile_box, tile_n=nothing, ...)
646
647 Initialize a Batchelor PMF on a smaller physical periodic tile and repeat that
648 tile through the full domain. This is useful for box-size controls where a
649 large box must contain the exact same physical magnetic pattern as a smaller
650 box, with no additional wavelengths larger than `tile_box`.
651 """
652 function init_pmf_batchelor_tiled!(s::MHDState{T}; tile_box::Real,
653 tile_n = nothing,
654 brms = nothing, b0_pG = nothing,
655 Bunit_pG::Real = 1, rho0::Real = 1,
656 p0::Real = 1, box::Real = s.dims[1] * s.dx,
657 seed::Integer = 1, kcut = nothing, kmin::Real = 0,
658 kmax = nothing, energy_index::Real = 4,
659 cutoff_order::Real = 2, mode_lock_n = nothing,
660 high_mode_lock_n = nothing,
661 preserve_low_kmax = nothing,
662 preserve_low_kcut = nothing,
663 normalization_kind = :auto,
664 curl_symbol = :centered,
665 curl_reference_n = nothing,
666 fixed_mode_power::Bool = false,
667 mode_realization = :legacy,
668 realization_box = nothing,
669 minimum_cells_per_wavelength = nothing,
670 invariant_axis::Integer = 0) where {T}
671 br = brms === nothing ? begin
672 b0_pG === nothing && error("provide either brms code units or b0_pG")
673 Float64(b0_pG) / Float64(Bunit_pG)
674 end : Float64(brms)
675 L = Float64(box)
676 Lt = Float64(tile_box)
677 Lt > 0 || error("tile_box must be positive")
678 Lt <= L || error("tile_box cannot exceed full box")
679 n = s.dims[1]
680 nt = tile_n === nothing ? round(Int, n * Lt / L) : Int(tile_n)
681 nt > 0 || error("tile_n must be positive")
682 nt <= n || error("tile_n cannot exceed target grid size")
683 n % nt == 0 || error("tile_n=$nt must divide target grid size $n")
684 isapprox(Lt / L, nt / n; rtol=0, atol=4eps(Float64) * max(1, Lt / L)) ||
685 error("tile_box/full box must match tile_n/target_n")
686 Bxt, Byt, Bzt, stats = pmf_batchelor_field((nt, nt, nt); box=Lt, brms=br,
687 seed=seed, kcut=kcut, kmin=kmin, kmax=kmax,
688 energy_index=energy_index,
689 cutoff_order=cutoff_order,
690 mode_lock_n=mode_lock_n,
691 high_mode_lock_n=high_mode_lock_n,
692 preserve_low_kmax=preserve_low_kmax,
693 preserve_low_kcut=preserve_low_kcut,
694 normalization_kind=normalization_kind,
695 curl_symbol=curl_symbol,
696 curl_reference_n=curl_reference_n,
697 fixed_mode_power=fixed_mode_power,
698 mode_realization=mode_realization,
699 realization_box=realization_box,
700 minimum_cells_per_wavelength=minimum_cells_per_wavelength,
701 invariant_axis=invariant_axis)
702 Bx = nt == n ? Bxt : _periodic_tile_field(Bxt, s.dims)
703 By = nt == n ? Byt : _periodic_tile_field(Byt, s.dims)
704 Bz = nt == n ? Bzt : _periodic_tile_field(Bzt, s.dims)
705 ctinfo = _install_pmf_field!(s, Bx, By, Bz; rho0=rho0, p0=p0)
706 out = merge(stats, (rho0=Float64(rho0), p0=Float64(p0),
707 Bunit_pG=Float64(Bunit_pG), full_box=L,
708 tile_box=Lt, tile_n=nt))
709 return ctinfo === nothing ? out : merge(out, (ct_faces=ctinfo,))
710 end