Snapshot from this documentation build. SHA-256: 7a8726a3dca15133e1e5a42c574b04caf8224d46b236707ddfb0d7cffd10ccc4
1 """Generate, run, and assess a real standalone mean-field MHD/P1 comparison.
2
3 This is a controlled weak-perturbation test with full local chemistry. A
4 supplied Ly-alpha table can add its nonlocal response. It does not certify
5 random-field nonlinear recombination or a full photon Boltzmann hierarchy.
6 """
7 from __future__ import annotations
8
9 import argparse
10 import csv
11 import hashlib
12 import json
13 import math
14 from pathlib import Path
15 import subprocess
16 import sys
17 import tomllib
18
19 import numpy as np
20
21 from . import lmhd, lya_rt, matched_units
22
23 ROOT = Path(__file__).resolve().parents[2]
24 FIELDS = {
25 "B_y": "b_y_over_B0", "gas_delta": "delta_b",
26 "gas_v_x": "v_b_z", "gas_v_y": "v_b_y",
27 "photon_v_y": "v_gamma_y",
28 "photon_baryon_slip_x": "slip_z", "photon_baryon_slip_y": "slip_y",
29 "dm_delta": "delta_dm", "dm_v_x": "v_dm_z", "x_HII": "delta_xe",
30 }
31 VELOCITIES = {name for name in FIELDS if "_v_" in name or "slip_" in name}
32
33
34 def mode_at(rows, field, z):
35 selected = sorted((-math.log1p(float(r["z"])),
36 complex(float(r["real"]),float(r["imag"])))
37 for r in rows if r["field"] == field)
38 if not selected:
39 raise ValueError(f"missing production Fourier mode: {field}")
40 x, values = np.asarray([v[0] for v in selected]), np.asarray([v[1] for v in selected])
41 target = -math.log1p(z)
42 if target < x[0]-1e-9 or target > x[-1]+1e-9:
43 raise ValueError(f"mode {field} does not span z={z}")
44 return complex(np.interp(target,x,values.real),np.interp(target,x,values.imag))
45
46
47 def physical_mode(value, field, z, handoff):
48 if field in ("B_y","B_z"):
49 return value / float(handoff["guide_field_code"])
50 if field in VELOCITIES and field != "dm_v_x":
51 value *= float(handoff["velocity_unit_cm_s_at_handoff"]) * math.sqrt(
52 (1+z)/(1+float(handoff["zhandoff"])))
53 return value
54
55
56 def write_config(path, settings):
57 # JSON scalar/string/array representations are valid for these TOML values.
58 path.write_text("\n\n".join("["+section+"]\n"+"\n".join(
59 f"{key} = {json.dumps(value,allow_nan=False)}" for key,value in values.items()
60 ) for section,values in settings.items())+"\n")
61
62
63 def case_settings(handoff, directory, args):
64 radiation = 4.15e-5 / args.hubble**2
65 settings = {
66 "run": {"n":args.resolution,"backend":args.backend,"cfl":args.cfl,
67 "reconstruction":"ppm","riemann":"hlld","max_cycles":3000000,
68 "print_every":1000,"checkpoint_final":True},
69 "cosmology": {"zstart":args.zstart,"zend":args.zend,
70 "box_ckpc_h":args.box_ckpc_h,"h":args.hubble,
71 "omega_m":matched_units.OMEGA_M,"omega_b":matched_units.OMEGA_B,
72 "omega_r":radiation,"omega_lambda":1-matched_units.OMEGA_M-radiation},
73 "magnetic": {"initial_condition":"meanfield_linear","b0_ng":args.field_pg*0.001},
74 "chemistry": {"mode":"mixing","helium":True,"expansion":"on",
75 "coupling":"strang","hydrogen_storage":"centered","smoothing":"local",
76 "ionized_fast_zmin":0,"x_hii_initial":handoff["background"]["xHII"],
77 "velocity_unit_cms":handoff["velocity_unit_cm_s_at_handoff"]},
78 "physics": {"gravity":True,"gas_gravity":True,"compton_drag":True,
79 "drag_mode":"radiation","photon_response":"moments_exponential",
80 "nonlinear_cfl":0.1,"glm_project_every":max(1,round(1.6/args.cfl)),
81 "pressure":handoff["background"]["pressure_code_mean"]},
82 "diagnostics": {"redshifts":np.linspace(args.zstart,args.zend,args.outputs).tolist(),
83 "angle_bins":1,"spectrum_bins":max(4,args.resolution//2),"energy_ledger_every":1},
84 "environment": {"MHD_INITIAL_STATE_DIR":str(directory.resolve()),
85 "MHD_MEANFIELD_THETA":args.theta,"MHD_MEANFIELD_PERT_FRAC":args.epsilon,
86 "MHD_MEANFIELD_POLARIZATION":"magnetosonic","MHD_PMF_KMODE":1,
87 "MHD_PMF_NORMALIZATION_Z":args.zstart,"MHD_CHEM_FH":matched_units.HYDROGEN_MASS_FRACTION,
88 "MHD_CHEM_FALPHA":0,"MHD_CHEM_ALLOW_LOCAL_MIXING":1,
89 "MHD_RADIATION_MAGNETIC_COMPENSATION":1,"MHD_RADIATION_FREE_STREAM_TOL":0,
90 "MHD_DRAG_HELIUM_ELECTRONS":1,"MHD_PROFILE_SKIP_CYCLES":5},
91 }
92 solver = getattr(args, "mhd_solver", "glm")
93 if solver != "glm":
94 settings["run"]["mhd_solver"] = solver
95 # CT rejects GLM B projection (its default is 0); this is the only other change.
96 del settings["physics"]["glm_project_every"]
97 if args.response_table:
98 settings["chemistry"]["nonlocal_table"] = str(args.response_table.resolve())
99 return settings
100
101
102 def assess(output, handoff, args):
103 run = output / "run"
104 metadata = tomllib.loads((run/"metadata.toml").read_text())
105 if metadata["status"] != "COMPLETED":
106 raise ValueError(f"production run is {metadata['status']}; no science PASS can be assigned")
107 with (run/"modes.tsv").open() as stream:
108 rows = list(csv.DictReader(stream,delimiter="\t"))
109 redshifts = sorted({float(r["z"]) for r in rows if r["field"]=="B_y"},reverse=True)
110 if len(redshifts)<2:
111 raise ValueError("at least two evolved epochs are required")
112 z0 = redshifts[0]
113 radiation = 4.15e-5/args.hubble**2
114 cosmo = lmhd.Cosmology(h=args.hubble,omega_m=matched_units.OMEGA_M,
115 omega_b=matched_units.OMEGA_B,omega_lambda=1-matched_units.OMEGA_M-radiation,
116 omega_r=radiation,hydrogen_mass_fraction=matched_units.HYDROGEN_MASS_FRACTION)
117 hydrogen = lmhd.PeeblesIonization(cosmo,zmax=3600,zmin=min(400,args.zend),samples=12000,recfast_v2=True)
118 temperature = lmhd.ComptonTemperature(cosmo,hydrogen,zmax=3600,zmin=min(400,args.zend),samples=12000)
119 table = lya_rt.ResponseTable.load(args.response_table) if args.response_table else None
120 k = matched_units.comoving_k_mpc(1,args.box_ckpc_h,args.hubble)
121
122 def rates(z):
123 local = lmhd.recfast_v2_local_response(cosmo,z,hydrogen(z),temperature(z))
124 if table is not None and table.redshift[0] <= z <= table.redshift[-1]:
125 delta = table.rates(z,k,differential_only=True)
126 delta[2] = table.rates(z,k)[2]
127 local = local + delta
128 return local
129
130 # Use the actual transferred Fourier state, without fitting any evolved
131 # amplitude. Runtime axes (x,y,z) map to oracle axes (z,y,x).
132 mode = lambda field: physical_mode(mode_at(rows,field,z0),field,z0,handoff)
133 velocity = handoff["velocity_unit_cm_s_at_handoff"]
134 initial = np.array([
135 mode("B_z"), mode("B_y"),mode_at(rows,"gas_delta",z0),
136 mode_at(rows,"gas_v_z",z0)*velocity,mode("gas_v_y"),mode("gas_v_x"),
137 mode_at(rows,"photon_delta",z0),mode_at(rows,"photon_v_z",z0)*velocity,
138 mode("photon_v_y"),mode_at(rows,"photon_v_x",z0)*velocity,
139 mode("dm_delta"),mode("dm_v_x"),mode("x_HII"),
140 mode_at(rows,"gas_temperature_fraction",z0)*temperature(z0),
141 ],dtype=np.complex128)
142 reference = lmhd.evolve_photon_p1(lmhd.ReferenceMode(k_mpc=k,theta=args.theta,
143 b0_gauss=args.field_pg*1e-12,perturbation_fraction=args.epsilon),cosmo,z0,
144 redshifts[-1],redshifts,hydrogen,max_dln_a=args.max_dlna,
145 include_gravity=True,include_dark_matter=True,include_photon_stress=True,
146 include_helium_electrons=True,recombination_rates=rates,
147 baryon_temperature=temperature,initial_state=initial)
148 comparisons=[]
149 for i,z in enumerate(reference.redshift):
150 for field,target in FIELDS.items():
151 candidate=physical_mode(mode_at(rows,field,float(z)),field,float(z),handoff)
152 expected=complex(reference.field(target)[i])
153 floor=1e-3 if field in VELOCITIES else 1e-6 if field=="x_HII" else 1e-7
154 status="INITIAL" if i==0 else "BELOW_FLOOR" if abs(expected)<floor else "PASS"
155 # The reduced oracle does not evolve explicit helium perturbations.
156 if field=="x_HII" and z>1700:
157 status="HELIUM_MODEL_DIAGNOSTIC"
158 amplitude=abs(abs(candidate)/abs(expected)-1) if abs(expected)>floor else math.nan
159 phase=abs(np.angle(candidate/expected)) if abs(expected)>floor else math.nan
160 if status=="PASS" and (not np.isfinite(amplitude+phase) or
161 amplitude>args.amplitude_tolerance or phase>args.phase_tolerance):
162 status="FAIL"
163 comparisons.append(dict(z=float(z),field=field,real=candidate.real,imag=candidate.imag,
164 reference_real=expected.real,reference_imag=expected.imag,
165 amplitude_error=amplitude,phase_error_rad=phase,signal_floor=floor,status=status))
166 active=[r for r in comparisons if r["status"] in ("PASS","FAIL")]
167 if not active:
168 raise ValueError("no resolved evolved observables; comparison would be vacuous")
169 passed=all(r["status"]=="PASS" for r in active)
170 gate=dict(status="PASS" if passed else "FAIL",scope="resolved mean-field full-MHD/shared photon-P1 modes",
171 source_fingerprint_sha256=metadata["source_fingerprint_sha256"],
172 reference_source_sha256={p.name:hashlib.sha256(p.read_bytes()).hexdigest()
173 for p in (Path(__file__),Path(lmhd.__file__),Path(lya_rt.__file__))},
174 initialization="actual jointly imported modes; no evolved amplitude fitting",
175 background="independent RECFAST-v2 hydrogen/temperature and Saha helium",
176 explicit_helium_perturbation="model extension; xHII above z1700 diagnostic",
177 thermal_perturbations="full-network extension; not gated against reduced oracle",
178 mhd_solver=getattr(args,"mhd_solver","glm"),
179 response_table=str(args.response_table) if args.response_table else None,
180 response_table_sha256=(hashlib.sha256(args.response_table.read_bytes()).hexdigest()
181 if args.response_table else None),
182 amplitude_tolerance=args.amplitude_tolerance,phase_tolerance_rad=args.phase_tolerance,
183 active_observables=len(active),failures=[r for r in active if r["status"]=="FAIL"],
184 max_amplitude_error=max(r["amplitude_error"] for r in active),
185 max_phase_error_rad=max(r["phase_error_rad"] for r in active),
186 mode_input_sha256=hashlib.sha256((run/"modes.tsv").read_bytes()).hexdigest())
187 with (output/"comparison.tsv").open("w") as stream:
188 writer=csv.DictWriter(stream,fieldnames=list(comparisons[0]),delimiter="\t",lineterminator="\n")
189 writer.writeheader(); writer.writerows(comparisons)
190 (output/"gate.json").write_text(json.dumps(gate,indent=2,allow_nan=False)+"\n")
191 import matplotlib
192 matplotlib.use("Agg")
193 import matplotlib.pyplot as plt
194 fig,axes=plt.subplots(2,2,figsize=(9,6),layout="constrained")
195 for ax,field in zip(axes.flat,("B_y","gas_delta","gas_v_x","dm_delta")):
196 selected=[r for r in comparisons if r["field"]==field]
197 ax.plot([r["z"] for r in selected],[abs(complex(r["real"],r["imag"])) for r in selected],label="3D MHD")
198 ax.plot([r["z"] for r in selected],[abs(complex(r["reference_real"],r["reference_imag"])) for r in selected],"--",label="P1 oracle")
199 ax.invert_xaxis(); ax.set(xlabel="Redshift",ylabel=field)
200 axes.flat[0].legend()
201 fig.suptitle(f"Standalone mean-field comparison: {gate['status']}")
202 fig.savefig(output/"comparison.pdf");fig.savefig(output/"comparison.png",dpi=160);plt.close(fig)
203 return gate
204
205
206 def main():
207 parser=argparse.ArgumentParser(description=__doc__)
208 parser.add_argument("output",type=Path)
209 parser.add_argument("--backend",choices=("cpu","metal","cuda"),default="cpu")
210 parser.add_argument("--mhd-solver",choices=("glm","ct"),default="glm")
211 parser.add_argument("--resolution",type=int,default=32)
212 parser.add_argument("--source-zstart",type=float,default=1e5)
213 parser.add_argument("--zstart",type=float,default=3200)
214 parser.add_argument("--zend",type=float,default=3100)
215 parser.add_argument("--field-pg",type=float,default=300)
216 parser.add_argument("--epsilon",type=float,default=0.025)
217 parser.add_argument("--theta",type=float,default=math.pi/4)
218 parser.add_argument("--box-ckpc-h",type=float,default=1)
219 parser.add_argument("--hubble",type=float,default=0.674)
220 parser.add_argument("--outputs",type=int,default=9)
221 parser.add_argument("--cfl",type=float,default=0.2)
222 parser.add_argument("--threads",type=int,default=4)
223 parser.add_argument("--julia",default="julia")
224 parser.add_argument("--max-dlna",type=float,default=1e-4)
225 parser.add_argument("--amplitude-tolerance",type=float,default=0.05)
226 parser.add_argument("--phase-tolerance",type=float,default=0.03)
227 parser.add_argument("--response-table",type=Path)
228 parser.add_argument("--execute",action="store_true")
229 args=parser.parse_args()
230 if not (ROOT/"scripts"/"run_pmf.jl").is_file():
231 parser.error("the native Julia comparison requires the pmf-glm source checkout; run it from an editable installation there")
232 if not args.source_zstart>args.zstart>args.zend>0 or args.outputs<2 or args.resolution<8 or args.resolution%2:
233 parser.error("require source-zstart > zstart > zend >0, even resolution >=8, outputs >=2")
234 output=args.output.resolve()
235 if output.exists() and any(output.iterdir()):
236 parser.error("use a new or empty comparison directory")
237 output.mkdir(parents=True,exist_ok=True)
238 command=[sys.executable,"-m","tools.reference.build_p1_meanfield_handoff",str(output/"handoff"),
239 "--resolution",str(args.resolution),"--zstart",str(args.source_zstart),
240 "--zhandoff",str(args.zstart),"--field-pg",str(args.field_pg),"--epsilon",str(args.epsilon),
241 "--theta",str(args.theta),"--box-ckpc-h",str(args.box_ckpc_h),"--hubble",str(args.hubble),
242 "--max-dlna",str(args.max_dlna),"--hydrogen-storage","centered"]
243 if args.response_table:
244 command += ["--response-table",str(args.response_table.resolve())]
245 with (output/"handoff.log").open("w") as stream:
246 subprocess.run(command,cwd=ROOT,check=True,stdout=stream)
247 handoff=json.loads((output/"handoff"/"metadata.json").read_text())
248 settings=case_settings(handoff,output/"handoff",args)
249 config=output/"case.toml"
250 write_config(config,settings)
251 project=ROOT if args.backend=="cpu" else ROOT/"environments"/args.backend
252 launch=[args.julia,f"--project={project}",f"--threads={args.threads}",str(ROOT/"scripts/run_pmf.jl"),str(config),f"--output={output/'run'}"]
253 print("Prepared:",config,"\nLaunch:"," ".join(launch),flush=True)
254 if not args.execute:
255 print("Add --execute in a new output directory to run and assess the complete comparison.")
256 return 0
257 subprocess.run(launch,cwd=ROOT,check=True)
258 gate=assess(output,handoff,args)
259 print(json.dumps(gate,indent=2))
260 return 0 if gate["status"]=="PASS" else 1
261
262
263 if __name__=="__main__":
264 raise SystemExit(main())