diff --git a/sbmy_control/low_level.py b/sbmy_control/low_level.py index 26054bf..5a32cdd 100644 --- a/sbmy_control/low_level.py +++ b/sbmy_control/low_level.py @@ -64,10 +64,14 @@ def stdout_redirector(stream): try: # Create a temporary file and redirect stdout to it tfile = tempfile.TemporaryFile(mode="w+b") + _redirect_stdout(tfile.fileno()) + # Yield to caller, then redirect stdout back to the saved fd yield + _redirect_stdout(saved_stdout_fd) + # Copy contents of temporary file to the given stream tfile.flush() tfile.seek(0, io.SEEK_SET) @@ -237,7 +241,7 @@ def progress_bar_from_logfile(filename:str, desc:str="", verbose:int=1, **kwargs from tqdm import tqdm from time import sleep k=0 - limit=600 + limit=6000 update_interval=0.2 wait_until_file_exists(filename, verbose=verbose, limit=limit) current_operation, total_operations = get_progress_from_logfile(filename) diff --git a/sbmy_control/parameters_card.py b/sbmy_control/parameters_card.py index bbdc1a6..5384adb 100644 --- a/sbmy_control/parameters_card.py +++ b/sbmy_control/parameters_card.py @@ -47,6 +47,15 @@ def register_arguments_card(parser:ArgumentParser): parser.add_argument("--OutputFCsDensity", type=str, default=None, help="Output FCs density file.") parser.add_argument("--OutputFCsSnapshot", type=str, default=None, help="Output FCs snapshot file.") parser.add_argument("--OutputRngStateLPT", type=str, default=None, help="Output RNG state file.") + ## Tests with phiBCs and density + parser.add_argument("--WriteGravPot", type=bool, default=False, help="Write gravitational potential.") + parser.add_argument("--OutputGravitationalPotentialBase", type=str, default=None, help="Output gravitational potential base.") + parser.add_argument("--MeshGravPot", type=int, default=None, help="Mesh for gravitational potential.") + parser.add_argument("--WriteDensity", type=bool, default=False, help="Write density.") + parser.add_argument("--OutputDensityBase", type=str, default=None, help="Output density base.") + parser.add_argument("--MeshDensity", type=int, default=None, help="Mesh for density.") + parser.add_argument("--LoadPhiBCs", type=bool, default=False, help="Load phiBCs.") + parser.add_argument("--InputPhiBCsBase", type=str, default=None, help="Input phiBCs file base.") parser.add_argument("--WriteReferenceFrame", type=bool, default=False, help="Write reference frame (COCA).") parser.add_argument("--OutputMomentaBase", type=str, default=None, help="Output momenta base (COCA).") parser.add_argument("--ReadReferenceFrame", type=bool, default=False, help="Read reference frame (COCA).") @@ -126,6 +135,16 @@ def parse_arguments_card(parsed_args): OutputFCsDensity=parsed_args.OutputFCsDensity, OutputFCsSnapshot=parsed_args.OutputFCsSnapshot, OutputRngStateLPT=parsed_args.OutputRngStateLPT, + ## Tests with phiBCs and density + WriteGravPot=parsed_args.WriteGravPot, + OutputGravitationalPotentialBase=parsed_args.OutputGravitationalPotentialBase, + MeshGravPot=parsed_args.MeshGravPot, + WriteDensity=parsed_args.WriteDensity, + OutputDensityBase=parsed_args.OutputDensityBase, + MeshDensity=parsed_args.MeshDensity, + LoadPhiBCs=parsed_args.LoadPhiBCs, + InputPhiBCsBase=parsed_args.InputPhiBCsBase, + ## Cosmological parameters h=cosmo_dict["h"], Omega_m=cosmo_dict["Omega_m"], Omega_b=cosmo_dict["Omega_b"], @@ -211,6 +230,17 @@ def parse_arguments_card(parsed_args): card_dict["OutputFCsSnapshot"] = main_dict["resultdir"]+"final_particles_"+main_dict["simname"]+".gadget3" if card_dict["OutputRngStateLPT"] is None: card_dict["OutputRngStateLPT"] = main_dict["workdir"]+"rng_state.h5" + ## Tests with phiBCs and density + if card_dict["OutputGravitationalPotentialBase"] is None: + card_dict["OutputGravitationalPotentialBase"] = main_dict["workdir"]+"gravpot_"+main_dict["simname"] + if card_dict["MeshGravPot"] is None: + card_dict["MeshGravPot"] = card_dict["N_PM_mesh"] + if card_dict["OutputDensityBase"] is None: + card_dict["OutputDensityBase"] = main_dict["workdir"]+"density_"+main_dict["simname"] + if card_dict["MeshDensity"] is None: + card_dict["MeshDensity"] = card_dict["N_PM_mesh"] + if card_dict["InputPhiBCsBase"] is None: + card_dict["InputPhiBCsBase"] = main_dict["workdir"]+"gravpot_tCOLA" if card_dict["OutputMomentaBase"] is None: card_dict["OutputMomentaBase"] = main_dict["workdir"]+"momenta_"+main_dict["simname"]+"_" if card_dict["InputMomentaBase"] is None: @@ -333,6 +363,15 @@ def create_parameter_card_dict( OutputFCsDensity:str = 'fcs_density.h5', OutputFCsSnapshot:str = 'fcs_particles.gadget3', + ## Tests with phiBCs and density + WriteGravPot:bool = True, + OutputGravitationalPotentialBase:str = 'gravitational_potential.h5', + MeshGravPot:int = 128, + WriteDensity:bool = False, + OutputDensityBase:str = 'density.h5', + MeshDensity:int = 128, + LoadPhiBCs:bool = False, + InputPhiBCsBase:str = 'gravitational_potential.h5', ## COCA parameters WriteReferenceFrame:bool = False, OutputMomentaBase:str = 'momenta_', @@ -403,6 +442,17 @@ def create_parameter_card_dict( OutputLPTPotential1=OutputLPTPotential1, OutputLPTPotential2=OutputLPTPotential2, OutputTilesBase=OutputTilesBase, + + # Tests with phiBCs and density + WriteGravPot=int(WriteGravPot), + OutputGravitationalPotentialBase=OutputGravitationalPotentialBase, + MeshGravPot=MeshGravPot, + WriteDensity=int(WriteDensity), + OutputDensityBase=OutputDensityBase, + MeshDensity=MeshDensity, + LoadPhiBCs=int(LoadPhiBCs), + InputPhiBCsBase=InputPhiBCsBase, + h=h, Omega_m=Omega_m, Omega_b=Omega_b, diff --git a/sbmy_control/parameters_monofonic.py b/sbmy_control/parameters_monofonic.py index e4a2169..d9cf896 100644 --- a/sbmy_control/parameters_monofonic.py +++ b/sbmy_control/parameters_monofonic.py @@ -95,8 +95,6 @@ def get_config_from_dict(monofonic_dict): "Omega_L": monofonic_dict["Omega_q"], "H0": monofonic_dict["h"]*100., "n_s": monofonic_dict["n_s"], - "sigma_8": monofonic_dict["sigma8"], - "A_s": monofonic_dict["A_s"], "Tcmb": monofonic_dict["Tcmb"], "k_p": monofonic_dict["k_p"], "N_ur": monofonic_dict["N_ur"], @@ -110,7 +108,14 @@ def get_config_from_dict(monofonic_dict): "transfer": "CLASS", "ztarget": 2.5, "ZeroRadiation": False - } + } + + if "A_s" in monofonic_dict.keys() and monofonic_dict["A_s"] is not None: + config["cosmology"]["A_s"] = monofonic_dict["A_s"] + elif "sigma8" in monofonic_dict.keys() and monofonic_dict["sigma8"] is not None: + config["cosmology"]["sigma8"] = monofonic_dict["sigma8"] + else: + raise KeyError("Either A_s or sigma8 must be provided in the cosmology parameters.") config["random"] = { "generator": "NGENIC", diff --git a/sbmy_control/scripts/convert_snapshot_to_density.py b/sbmy_control/scripts/convert_snapshot_to_density.py index dadefed..869ee90 100644 --- a/sbmy_control/scripts/convert_snapshot_to_density.py +++ b/sbmy_control/scripts/convert_snapshot_to_density.py @@ -34,6 +34,55 @@ def convert_snapshot_to_density(snapshot_path, output_path, N=None, corner=(0.0, F.write(output_path) print("Density written to", output_path) print("Done.") + + +def convert_snapshots_to_density(snapshot_base_path, output_path, N=None, corner=(0.0, 0.0, 0.0)): + """ + Convert multiple snapshots to density fields. + + Parameters + ---------- + snapshot_base_path : str + Base path for the snapshot files. + output_path : str + Path to the output density file. + N : int + Size of the density field grid (N x N x N). + corner : tuple of float + Corner of the box (x, y, z). + """ + # Get all snapshot files with path "snapshot_base_path*" + import glob + snapshot_files = glob.glob(snapshot_base_path + "*") + if not snapshot_files: + raise FileNotFoundError(f"No snapshot files found at {snapshot_base_path}") + + A = None + F = None + + for snapshot_path in snapshot_files: + # Read the snapshot + print(f"Reading snapshot {snapshot_path}...") + snap = read_snapshot(snapshot_path) + + if N is None: + N = snap.Np0 + + # Calculate density + print("Calculating density...") + F = get_density_pm_snapshot(snap, N, N, N, corner[0], corner[1], corner[2]) + + if A is None: + A = F.data + else: + A += F.data + + F.data = A # Combine densities from all snapshots + # Write density to file + print("Writing density...") + F.write(output_path) + print("Density written to", output_path) + print("Done.") def console_main(): @@ -70,13 +119,23 @@ def console_main(): args = parser.parse_args() - convert_snapshot_to_density( - snapshot_path=args.snapshot, - output_path=args.output, - N=args.N, - corner=args.corner, - ) + if args.snapshot.endswith("h5") or args.snapshot.endswith("hdf5") or args.snapshot.endswith("gadget3"): + # If the snapshot is a single file, convert it to density + convert_snapshot_to_density( + snapshot_path=args.snapshot, + output_path=args.output, + N=args.N, + corner=args.corner, + ) + else: + # If the snapshot is a base path, convert all snapshots to density + convert_snapshots_to_density( + snapshot_base_path=args.snapshot, + output_path=args.output, + N=args.N, + corner=args.corner, + ) if __name__ == "__main__": - console_main() \ No newline at end of file + console_main()