Snapshot from this documentation build. SHA-256: c96ca6b98dfcc422686406b8f38d124481b78e5a4a50851e8350e4cea8465777
1 using Test
2
3 ENV["MHD_BACKEND"] = "cpu"
4 include(joinpath(@__DIR__, "..", "..", "src", "production", "runner.jl"))
5
6 @testset "incremental chemistry energy update" begin
7 total = 980.8842f0
8 internal = 418.7102f0
9 @test _incremental_total_energy(total, 1f0, internal, internal, 1f-30) === total
10 @test _incremental_total_energy(total, 2f0, internal, internal + 0.25f0, 1f-30) ==
11 total + 0.5f0
12 end
13
14 @testset "PMF checkpoint endpoint extension" begin
15 mktempdir() do dir
16 path = joinpath(dir, "checkpoint.jls")
17 stored = (N=4, zstart=3.0e6, zend=1.0e5, cfl=0.1)
18 expected = (N=4, zstart=3.0e6, zend=500.0, cfl=0.1)
19 source = Float32[1, 2, 3]
20 _write_pmf_checkpoint(path, stored, (cycle=7,), ["field" => source])
21 destination = zeros(Float32, 3)
22 runtime = _restore_pmf_checkpoint!(path, expected, ["field" => destination])
23 @test runtime.cycle == 7
24 @test destination == source
25
26 optional_destination = ones(Float32, 3)
27 runtime = _restore_pmf_checkpoint!(
28 path, expected,
29 ["field" => destination, "magnetic_residual_1" => optional_destination];
30 optional_missing_arrays=("magnetic_residual_1",))
31 @test runtime.cycle == 7
32 @test all(iszero, optional_destination)
33 @test_throws ErrorException _restore_pmf_checkpoint!(
34 path, expected,
35 ["field" => destination, "required_trailing" => zeros(Float32, 3)])
36
37 incompatible = (N=4, zstart=3.0e6, zend=500.0, cfl=0.2)
38 @test_throws ErrorException _restore_pmf_checkpoint!(
39 path, incompatible, ["field" => zeros(Float32, 3)])
40
41 destination .= 0
42 runtime = _restore_pmf_checkpoint!(
43 path, incompatible, ["field" => destination];
44 allowed_config_changes=(:cfl,))
45 @test runtime.cycle == 7
46 @test destination == source
47
48 physical_change = (N=8, zstart=3.0e6, zend=500.0, cfl=0.2)
49 @test_throws ErrorException _restore_pmf_checkpoint!(
50 path, physical_change, ["field" => zeros(Float32, 3)];
51 allowed_config_changes=(:cfl,))
52 end
53 end
54
55 @testset "PMF checkpoint mhd_solver guard" begin
56 mktempdir() do dir
57 base = (N=4, zstart=3.0e6, zend=1.0e5, cfl=0.1)
58 source = Float32[1, 2, 3]
59 restore(path, config; arrays=["field" => zeros(Float32, 3)]) =
60 _restore_pmf_checkpoint!(path, config, arrays)
61 legacy = joinpath(dir, "legacy.jls") # written before the key existed => GLM
62 glm = joinpath(dir, "glm.jls")
63 ct = joinpath(dir, "ct.jls")
64 _write_pmf_checkpoint(legacy, base, (cycle=1,), ["field" => source])
65 _write_pmf_checkpoint(glm, merge(base, (mhd_solver=:glm,)), (cycle=2,), ["field" => source])
66 _write_pmf_checkpoint(ct, merge(base, (mhd_solver=:ct,)), (cycle=3,),
67 ["field" => source, "ct_bxf" => source])
68 want(solver) = merge(base, (mhd_solver=solver,))
69 # A legacy checkpoint is GLM: it resumes as GLM (also with a changed endpoint) ...
70 @test restore(legacy, want(:glm)).cycle == 1
71 @test restore(legacy, merge(want(:glm), (zend=5.0e4,))).cycle == 1
72 @test restore(glm, want(:glm)).cycle == 2
73 # ... and must never resume as CT, nor a GLM/CT checkpoint as the other solver.
74 for path in (legacy, glm)
75 err = try; restore(path, want(:ct)); nothing; catch e; e; end
76 @test err isa ErrorException
77 @test occursin("mhd_solver mismatch", err.msg)
78 end
79 err = try; restore(ct, want(:glm)); nothing; catch e; e; end
80 @test err isa ErrorException && occursin("mhd_solver mismatch", err.msg)
81 # CT faces round-trip exactly.
82 faces = zeros(Float32, 3); field = zeros(Float32, 3)
83 @test restore(ct, want(:ct); arrays=["field" => field, "ct_bxf" => faces]).cycle == 3
84 @test faces == source
85 end
86 end
87
88 @testset "nonlocal table checkpoint migration guard" begin
89 mktempdir() do dir
90 function write_table(path, redshifts)
91 open(path, "w") do io
92 println(io, "redshift\tk_mpc")
93 for z in redshifts
94 println(io, "$z\t100")
95 end
96 end
97 end
98 old_table = joinpath(dir, "old.tsv")
99 new_table = joinpath(dir, "new.tsv")
100 write_table(old_table, (700.0, 1200.0, 1700.0))
101 write_table(new_table, (700.0, 1200.0, 1700.0))
102 stored = (N=4, zend=500.0, chem_nonlocal_table=old_table)
103 expected = (N=4, zend=500.0, chem_nonlocal_table=new_table)
104 source = Float32[1, 2, 3]
105
106 safe_path = joinpath(dir, "safe.jls")
107 _write_pmf_checkpoint(
108 safe_path, stored, (cycle=7, a=z_to_a(1e5)), ["field" => source]
109 )
110 destination = zeros(Float32, 3)
111 runtime = _restore_pmf_checkpoint!(
112 safe_path, expected, ["field" => destination];
113 allowed_config_changes=(:chem_nonlocal_table,),
114 allowed_config_validator=_checkpoint_nonlocal_table_change_is_safe,
115 )
116 @test runtime.cycle == 7
117 @test destination == source
118
119 unsafe_path = joinpath(dir, "unsafe.jls")
120 _write_pmf_checkpoint(
121 unsafe_path, stored, (cycle=8, a=z_to_a(1100.0)), ["field" => source]
122 )
123 @test_throws ErrorException _restore_pmf_checkpoint!(
124 unsafe_path, expected, ["field" => zeros(Float32, 3)];
125 allowed_config_changes=(:chem_nonlocal_table,),
126 allowed_config_validator=_checkpoint_nonlocal_table_change_is_safe,
127 )
128 end
129 end
130
131 @testset "Jedamzik cosmological MHD scaling" begin
132 a_ref = 1 / 4501
133 a = 1 / 1501
134 h = 0.67
135 vunit = 8.0e5
136 lbox = 0.713 * 3.0857e21
137
138 legacy_clock = _hydro_time_per_tau(a, h, vunit, lbox)
139 transformed_clock = _hydro_time_per_tau(a, h, vunit, lbox, a_ref)
140 @test transformed_clock / legacy_clock ≈ sqrt(a_ref / a)
141
142 dtphys = 1.25e10
143 legacy_dt = _hydro_code_dt(dtphys, 0.0, vunit, lbox, a)
144 transformed_dt = _hydro_code_dt(dtphys, 0.0, vunit, lbox, a, a_ref)
145 @test transformed_dt / legacy_dt ≈ sqrt(a_ref / a)
146 @test _cosmological_velocity_unit(vunit, a, a_ref) ≈
147 vunit * sqrt(a_ref / a)
148
149 # Integrating the extra H~/2 term exactly gives v~ ∝ a^(-1/2).
150 @test _hubble_momentum_decay(a_ref, a, a_ref) ≈ sqrt(a_ref / a)
151 @test _hubble_momentum_decay(a_ref, a, NaN) == 1.0
152
153 # At the normalization epoch the transformed and legacy code units agree.
154 @test _hydro_time_per_tau(a_ref, h, vunit, lbox, a_ref) ≈
155 _hydro_time_per_tau(a_ref, h, vunit, lbox)
156 @test _hydro_code_dt(dtphys, 0.0, vunit, lbox, a_ref, a_ref) ≈
157 _hydro_code_dt(dtphys, 0.0, vunit, lbox, a_ref)
158 end
159
160 @testset "exact cosmological diagnostic landing" begin
161 Om = 0.315
162 Or = 4.15e-5 / 0.674^2
163 OL = 1 - Om - Or
164 Ok = 0.0
165 a0 = z_to_a(3200.0)
166 for target_z in (3000.0, 1600.0, 500.0)
167 target_a = z_to_a(target_z)
168 dt = _dtau_to_scale_factor_rk2(Om, OL, Ok, Or, a0, target_a)
169 landed = _advance_a_rk2(Om, OL, Ok, Or, a0, dt, Inf)
170 @test landed ≈ target_a rtol=8eps(Float64) atol=0
171 @test a_to_z(landed) ≈ target_z rtol=8eps(Float64) atol=0
172 end
173
174 outputs = [3000.0, 1600.0, 500.0]
175 unrestricted = 10 * _dtau_to_scale_factor_rk2(
176 Om, OL, Ok, Or, a0, z_to_a(outputs[1]))
177 limited = _limit_dt_to_output_redshift(
178 unrestricted, Om, OL, Ok, Or, a0, outputs, 1)
179 @test _advance_a_rk2(Om, OL, Ok, Or, a0, limited, Inf) ≈
180 z_to_a(outputs[1]) rtol=8eps(Float64) atol=0
181 @test _limit_dt_to_output_redshift(
182 unrestricted, Om, OL, Ok, Or, a0, outputs, length(outputs) + 1) == unrestricted
183 end
184
185 @testset "particle supercomoving diagnostic units" begin
186 h = 0.674
187 a = z_to_a(3200.0)
188 box_ckpc_h = 1.0
189 h0 = 100.0 * h * 1.0e5 / 3.0856775814913673e24
190 lbox = box_ckpc_h * 3.0856775814913673e21 / h
191 @test _particle_velocity_cms_per_code(a, box_ckpc_h) ≈ h0 * lbox / a
192 end
193
194 @testset "imported PMF rescaling" begin
195 mktempdir() do dir
196 N = 4
197 dims = (N, N, N)
198 fill_value = (1.0f0, 2.0f0, 3.0f0)
199 for (name, value) in zip(("bx.f32", "by.f32", "bz.f32"), fill_value)
200 _write_f32_array_atomic(joinpath(dir, name), fill(Float32(value), dims))
201 end
202 be = MHDKernels.backend(:cpu)
203 state = MHDKernels.allocate_state(be, Float32, dims; dx=1/N, gamma=5/3,
204 riemann=:hll, recon=:ppm)
205 target = 0.25f0
206 measured = _import_initial_b_field!(be, dir, state; N=N, p0=1.0,
207 target_brms=target)
208 @test isapprox(measured, target; rtol=5e-6)
209 host = MHDKernels.fields_to_host(state)
210 @test all(iszero, host[2]) && all(iszero, host[3]) && all(iszero, host[4])
211 expected_energy = 1.0f0 / (Float32(5/3) - 1.0f0) + 0.5f0 * target^2
212 @test all(value -> isapprox(value, expected_energy; rtol=1e-5), host[5])
213 end
214 end
215
216 @testset "joint PMF handoff import" begin
217 mktempdir() do dir
218 N = 4
219 dims = (N, N, N)
220 be = MHDKernels.backend(:cpu)
221 state = MHDKernels.allocate_state(be, Float32, dims; dx=1/N, gamma=5/3,
222 riemann=:hll, recon=:ppm)
223 expected = [fill(Float32(i), dims) for i in 1:9]
224 expected[1] .= 1.25f0
225 expected[5] .= 8.0f0
226 for i in eachindex(expected)
227 _write_f32_array_atomic(joinpath(dir, "mhd_U_$i.f32"), expected[i])
228 end
229 measured = _import_initial_mhd_state!(be, dir, state; N=N)
230 @test isfinite(measured) && measured > 0
231 @test isfinite(_device_pressure_mean!(be, state, state.scratch[1]))
232 for i in eachindex(expected)
233 @test Array(state.U[i]) == vec(expected[i])
234 end
235
236 particles = ntuple(i -> zeros(Float32, N^3), 6)
237 particle_names = ("particle_px", "particle_py", "particle_pz",
238 "particle_vx", "particle_vy", "particle_vz")
239 for (i, name) in enumerate(particle_names)
240 _write_f32_array_atomic(joinpath(dir, "$name.f32"),
241 fill(Float32(i) / 10, N^3))
242 end
243 chemistry = ntuple(i -> zeros(Float32, N^3), 3)
244 for (i, name) in enumerate(("chem_eint", "chem_HII", "chem_H2I"))
245 _write_f32_array_atomic(joinpath(dir, "$name.f32"),
246 fill(Float32(i) / 100, N^3))
247 end
248 photon = MHDKernels.allocate_photon_moment_state(state)
249 photon_expected = reshape(ComplexF32.(1:length(photon.hat_batch),
250 -1:-1:-length(photon.hat_batch)),
251 size(photon.hat_batch))
252 open(joinpath(dir, "photon_hat_batch.c32"), "w") do io
253 write(io, photon_expected)
254 end
255 residual = ntuple(_ -> zeros(Float32, N^3), 3)
256 for (i, name) in enumerate(("magnetic_residual_1",
257 "magnetic_residual_2",
258 "magnetic_residual_3"))
259 _write_f32_array_atomic(joinpath(dir, "$name.f32"),
260 fill(-Float32(i) / 1000, N^3))
261 end
262
263 _import_initial_aux_state!(dir; N=N, particles=particles,
264 chemistry=chemistry, photons=photon,
265 magnetic_residual=residual)
266 for i in 1:6
267 @test particles[i] == fill(Float32(i) / 10, N^3)
268 end
269 for i in 1:3
270 @test chemistry[i] == fill(Float32(i) / 100, N^3)
271 @test residual[i] == fill(-Float32(i) / 1000, N^3)
272 end
273 _write_f32_array_atomic(joinpath(dir, "chem_HeII.f32"),
274 fill(0.04f0, N^3))
275 chemistry_helium = ntuple(i -> zeros(Float32, N^3), 4)
276 _import_initial_aux_state!(dir; N=N, particles=particles,
277 chemistry=chemistry_helium, photons=photon,
278 magnetic_residual=residual)
279 for i in 1:4
280 @test chemistry_helium[i] == fill(Float32(i) / 100, N^3)
281 end
282 @test Array(photon.hat_batch) == photon_expected
283
284 rm(joinpath(dir, "magnetic_residual_3.f32"))
285 @test_throws ErrorException _import_initial_aux_state!(
286 dir; N=N, particles=particles, chemistry=chemistry_helium,
287 photons=photon, magnetic_residual=residual)
288 end
289 end
290
291 @testset "hydrogen-storage handoff contract" begin
292 mktempdir() do dir
293 N = 4
294 chemistry = ntuple(_ -> zeros(Float32, N^3), 3)
295 for (i, name) in enumerate(("chem_eint", "chem_HII", "chem_H2I"))
296 _write_f32_array_atomic(joinpath(dir, "$name.f32"),
297 fill(Float32(i) / 100, N^3))
298 end
299
300 # Legacy handoffs have no marker and are ionized-storage by definition.
301 _import_initial_aux_state!(dir; N=N, particles=(), chemistry=chemistry,
302 photons=nothing)
303 @test_throws ErrorException _import_initial_aux_state!(
304 dir; N=N, particles=(), chemistry=chemistry, photons=nothing,
305 neutral_hydrogen_storage=true)
306
307 write(joinpath(dir, "chem_hydrogen_storage.txt"), "neutral\n")
308 _import_initial_aux_state!(dir; N=N, particles=(), chemistry=chemistry,
309 photons=nothing,
310 neutral_hydrogen_storage=true)
311 @test_throws ErrorException _import_initial_aux_state!(
312 dir; N=N, particles=(), chemistry=chemistry, photons=nothing)
313
314 write(joinpath(dir, "chem_hydrogen_storage.txt"), "centered\n")
315 @test_throws ErrorException _import_initial_aux_state!(
316 dir; N=N, particles=(), chemistry=chemistry, photons=nothing,
317 centered_hydrogen_storage=true)
318 reference = 0.9999998123456789
319 write(joinpath(dir, "chem_hydrogen_reference.txt"), "$reference\n")
320 _import_initial_aux_state!(
321 dir; N=N, particles=(), chemistry=chemistry, photons=nothing,
322 centered_hydrogen_storage=true)
323 @test _initial_hydrogen_reference(dir, :centered) == reference
324 @test_throws ErrorException _initial_hydrogen_reference(dir, :neutral)
325 end
326 end
327
328 @testset "cosmological helium thermodynamic closure" begin
329 h = 0.674
330 Ob = 0.05366655
331 fh = 0.7546
332 vunit = 8.0e5
333 @test _PMF_T_CMB == ChemistryKernels.comp2_cmb(0.0)
334 for z in (1.0e6, 3200.0, 2999.0)
335 helium = _equilibrium_helium_state(1.0, z, h, Ob, fh)
336 @test 1 <= helium.xe_total <= 1 + 2(1 - fh) / (4fh)
337 @test 0 <= helium.xheii <= (1 - fh) / (4fh)
338 @test 0 <= helium.heii_mass_fraction <= 1 - fh
339 @test _drag_electron_fraction(1.0, helium.xheii, z, h, Ob, fh, true) ≈
340 helium.xe_total rtol=2e-10
341
342 temperature = _PMF_T_CMB * (1 + z)
343 particles = fh * (1 + helium.xe_total) + (1 - fh) / 4
344 expected = particles * ChemistryKernels.KBOLTZ * temperature /
345 (ChemistryKernels.MH * vunit^2)
346 @test _cosmological_ionized_p0(z, fh, vunit, h, Ob) ≈ expected rtol=2e-14
347 end
348
349 # The device diagnostic receives dimensionless code density. Its Saha
350 # reconstruction must use the physical density unit or He III electrons
351 # are lost at high redshift and the inferred temperature is biased high.
352 N = 4
353 z = 1.0e6
354 be = MHDKernels.backend(:cpu)
355 state = MHDKernels.allocate_state(
356 be, Float32, (N, N, N); dx=1 / N, gamma=5 / 3,
357 )
358 helium = _equilibrium_helium_state(1.0, z, h, Ob, fh)
359 fill!(state.U[1], 1f0)
360 foreach(field -> fill!(field, 0f0), state.U[2:4])
361 foreach(field -> fill!(field, 0f0), state.U[6:8])
362 fill!(state.U[5], Float32(_cosmological_ionized_p0(z, fh, vunit, h, Ob) /
363 (state.γ - 1)))
364 HII = fill(Float32(fh), N^3)
365 H2I = zeros(Float32, N^3)
366 HeII = fill(Float32(helium.heii_mass_fraction), N^3)
367 stats = _device_temperature_stats!(
368 be, state, state.scratch[1], HII, H2I, HeII,
369 _density_unit_cgs(z, h, Ob), vunit^2, z, fh,
370 )
371 @test stats[1] ≈ _PMF_T_CMB * (1 + z) rtol=2e-5
372 @test stats[2] ≈ stats[1] rtol=2e-6
373 @test stats[3] ≈ stats[1] rtol=2e-6
374 end
375
376 @testset "baryon-DM common physical clocks" begin
377 a0 = z_to_a(1100.0); a1 = z_to_a(500.0)
378 h=0.674; Lh=2.75; vunit=8e5; lbox=Lh/h*_PMF_KPC_CM
379 unit0 = _cosmological_velocity_unit(vunit,a0,a0)
380 unit1 = _cosmological_velocity_unit(vunit,a1,a0)
381 @test _hubble_momentum_decay(a0,a1,a0)*unit1/unit0 ≈ a0/a1
382 @test _particle_velocity_cms_per_code(a1,Lh)/_particle_velocity_cms_per_code(a0,Lh) ≈ a0/a1
383 dtau=0.03
384 @test _gas_gravity_kick_dt(dtau,a1,h,vunit,lbox,a0)*unit1 ≈
385 dtau*_particle_velocity_cms_per_code(a1,Lh)
386 @test _particle_dt_limit(1e8,1e-20,0.01) ≈ 1e-10
387 end