diff --git a/ICs.py b/ICs.py index a50b2fa..188e6c8 100644 --- a/ICs.py +++ b/ICs.py @@ -4,12 +4,12 @@ def main_ICs(parsed_args): Main function for the initial conditions generator. """ from args_main import parse_arguments_main - from parameters_card import parse_arguments_card + from parameters_card import parse_arguments_card_for_ICs from low_level import print_starting_module, print_message, print_ending_module from os.path import isfile main_dict = parse_arguments_main(parsed_args) - card_dict, _ = parse_arguments_card(parsed_args) + card_dict = parse_arguments_card_for_ICs(parsed_args) print_starting_module("ICs", verbose=parsed_args.verbose) @@ -35,7 +35,7 @@ def main_ICs(parsed_args): case "monofonic": if card_dict["ICsMode"] == 1 or card_dict["ICsMode"] == 0: - raise ValueError("ICsMode is 1 or 2 but ICs_gen is monofonic.") + raise ValueError("ICsMode is 1 or 0 but ICs_gen is monofonic.") elif card_dict["ICsMode"] == 2: print_message("ICsMode is 2 so the initial conditions will be generated using the monofonic ICs generator.", 1, "ICs", verbose=parsed_args.verbose) ICs_monofonic(parsed_args) @@ -53,18 +53,27 @@ def main_ICs(parsed_args): def ICs_monofonic(parsed_args): from monofonic import main_monofonic - main_monofonic(parsed_args) + from os.path import isfile + from low_level import print_message + from parameters_monofonic import parse_arguments_monofonic + + monofonic_dict = parse_arguments_monofonic(parsed_args) + + if not isfile(monofonic_dict["output"]+"DM_delta.h5") or not isfile(monofonic_dict["output"]+"DM_phi2.h5") or not isfile(monofonic_dict["output"]+"DM_phi.h5") or parsed_args.force: + main_monofonic(parsed_args) + else: + print_message(f"Monofonic output files found in {monofonic_dict['output']}. Use -F to overwrite.", 1, "monofonic", verbose=parsed_args.verbose) + def ICs_sbmy(parsed_args): - from pysbmy.power import PowerSpectrum from low_level import print_starting_module, print_message from os.path import isfile - from parameters_card import parse_arguments_card + from parameters_card import parse_arguments_card_for_ICs print_starting_module("sbmy IC", verbose=parsed_args.verbose) - card_dict, _ = parse_arguments_card(parsed_args) + card_dict = parse_arguments_card_for_ICs(parsed_args) power_spectrum_file = card_dict["InputPowerSpectrum"] @@ -167,7 +176,7 @@ def create_sbmy_white_noise_field(parsed_args, card_dict, white_noise_field_file if __name__ == "__main__": from argparse import ArgumentParser from args_main import register_arguments_main - from parameters_card import register_arguments_card + from parameters_card import register_arguments_card_for_ICs from cosmo_params import register_arguments_cosmo from parameters_monofonic import register_arguments_monofonic from slurm_submission import register_arguments_slurm @@ -177,7 +186,7 @@ if __name__ == "__main__": register_arguments_main(parser) register_arguments_monofonic(parser) register_arguments_slurm(parser) - register_arguments_card(parser) + register_arguments_card_for_ICs(parser) register_arguments_cosmo(parser) parsed_args = parser.parse_args() diff --git a/args_main.py b/args_main.py index 4c60e31..a022f22 100644 --- a/args_main.py +++ b/args_main.py @@ -10,7 +10,7 @@ def register_arguments_main(parser:ArgumentParser): parser.add_argument("-wd", "--workdir", type=str, default=None, help="Directory where the work is done (default is -d/work).") parser.add_argument("-ld", "--logdir", type=str, default=None, help="Directory where the logs are saved (default is -d/logs).") parser.add_argument("-n", "--simname", type=str, default=None, help="Name of the simulation (for outputs and parameter card).") - parser.add_argument("-M", "--mode", type=str, default="tCOLA", help="Mode of the program: ICs, PM, LPT, tCOLA, sCOLA, pre_sCOLA, post_sCOLA, alltCOLA, allsCOLA.") + parser.add_argument("-M", "--mode", type=str, default="tCOLA", help="Mode of the program: ICs, PM, LPT, tCOLA, sCOLA, pre_sCOLA, post_sCOLA, allPM, alltCOLA, allsCOLA, TS.") parser.add_argument("-ic","--ICs_gen", type=str, default="ext", help="Initial conditions generator: ext, sbmy, monofonic") parser.add_argument("-v", "--verbose", type=int, default=1, help="Verbosity level (0: no output, 1: basic output, 2: subprograms output).") parser.add_argument("--seed", type=int, default=None, help="Seed for the random number generator.") @@ -49,7 +49,7 @@ def parse_arguments_main(parsed_args): if main_dict["logdir"] is None: main_dict["logdir"] = main_dict["directory"]+"logs/" if main_dict["simname"] is None: - main_dict["simname"] = main_dict["mode"] + main_dict["simname"] = "sCOLA" if "sCOLA" in main_dict["mode"] else main_dict["mode"] if main_dict["seed"] is None: import numpy as np main_dict["seed"] = np.random.randint(0, 2**32) diff --git a/low_level.py b/low_level.py index 988c145..05b5126 100644 --- a/low_level.py +++ b/low_level.py @@ -193,8 +193,64 @@ def wait_until_file_exists(filename:str, verbose:int=1, limit:int=24*3600): from os.path import isfile k=0 - while not isfile(filename): + while not isfile(filename) and k=limit: + raise TimeoutError(f"File {filename} not found after {limit} seconds.") + if k>60: + print_message(f"File {filename} exists. {k//60} minutes elapsed.", 3, "low level", verbose=verbose) + + +def get_progress_from_logfile(filename): + """ + Get the number of operations done and the total number of operations from a log file. + """ + current_operation = 0 + total_operations = 0 + from os.path import isfile + if not isfile(filename): + raise FileNotFoundError(f"Log file {filename} not found.") + with open(filename, "r") as f: + lines = f.readlines() + for line in lines: + if " operation " in line: + try: + splitted_line = line.split(" operation ")[1] + splitted_line = splitted_line.split(":")[0] + current_operation = int(splitted_line.split("/")[0]) + total_operations = int(splitted_line.split("/")[1]) + except: + pass + return current_operation, total_operations + + +def progress_bar_from_logfile(filename:str, desc:str="", verbose:int=1, **kwargs): + """ + Print a progress bar from a log file. + """ + from tqdm import tqdm + from time import sleep + k=0 + limit=3600 + update_interval=0.2 + sleep(2) # Wait for the process to be launched, and for the previous log file to be overwritten if necessary. + wait_until_file_exists(filename, verbose=verbose, limit=limit) + current_operation, total_operations = get_progress_from_logfile(filename) + previous_operation = 0 + + if current_operation == total_operations: + # print_message("Progress bar not needed, the process is already finished.", 3, "low level", verbose=verbose) + return + + with tqdm(desc=desc, total=total_operations, disable=(verbose==0), **kwargs) as pbar: + while current_operation < total_operations and k/update_interval < limit: + sleep(update_interval) + current_operation, total_operations = get_progress_from_logfile(filename) + if current_operation > previous_operation: + pbar.update(current_operation-previous_operation) + previous_operation = current_operation + k+=1 + if k/update_interval >= limit: + print_message(f"Progress bar timed out after {limit} seconds.", 3, "low level", verbose=verbose) \ No newline at end of file diff --git a/main.py b/main.py new file mode 100644 index 0000000..1913370 --- /dev/null +++ b/main.py @@ -0,0 +1,239 @@ + + +def main(parsed_args): + from low_level import print_starting_module, print_message, print_ending_module, wait_until_file_exists + from os.path import isfile + from args_main import parse_arguments_main + + print_starting_module("control", verbose=parsed_args.verbose) + main_dict = parse_arguments_main(parsed_args) + + match main_dict["mode"]: + + case "ICs" | "InitialConditions" | "InitialConditionsGenerator" | "ICsGenerator" | "ICsGen" | "ini": + print_message("Running initial conditions generator.", 1, "control", verbose=parsed_args.verbose) + from ICs import main_ICs + main_ICs(parsed_args) + print_message("Initial conditions generator finished.", 1, "control", verbose=parsed_args.verbose) + + + case "TS" | "timestepping": + print_message("Running timestepping generator.", 1, "control", verbose=parsed_args.verbose) + from timestepping import main_timestepping + main_timestepping(parsed_args) + print_message("Timestepping generator finished.", 1, "control", verbose=parsed_args.verbose) + + + case "PM" | "LPT" | "tCOLA" | "simbelmyne" | "sbmy": + print_message(f"Running Simbelmyne in mode {main_dict["mode"]}.", 1, "control", verbose=parsed_args.verbose) + from simbelmyne import main_simbelmyne + from parameters_card import parse_arguments_card + from os.path import isfile + + card_dict = parse_arguments_card(parsed_args) + + check_consistency(card_dict, main_dict["mode"]) + + output_is_required = (card_dict["WriteFinalDensity"] or card_dict["WriteFinalSnapshot"]) + if not output_is_required: + print_message("Output is not required. Skipping simbelmyne.", 1, "control", verbose=parsed_args.verbose) + else: + if (card_dict["WriteFinalDensity"] and not isfile(card_dict["OutputFinalDensity"])) or (card_dict["WriteFinalSnapshot"] and not isfile(card_dict["OutputFinalSnapshot"])) or parsed_args.force: + main_simbelmyne(parsed_args) + print_message("Simbelmyne finished.", 1, "control", verbose=parsed_args.verbose) + else: + print_message("Output files already exist. Use -F to overwrite.", 1, "control", verbose=parsed_args.verbose) + + + case "pre_sCOLA": + print_message("Running pre-sCOLA.", 1, "control", verbose=parsed_args.verbose) + from scola import main_pre_scola + from parameters_card import parse_arguments_card + + card_dict = parse_arguments_card(parsed_args) + + check_consistency(card_dict, main_dict["mode"]) + + main_pre_scola(parsed_args) + print_message("Pre-sCOLA finished.", 1, "control", verbose=parsed_args.verbose) + + + case "post_sCOLA": + print_message("Running post-sCOLA.", 1, "control", verbose=parsed_args.verbose) + from scola import main_post_scola + from parameters_card import parse_arguments_card + + card_dict = parse_arguments_card(parsed_args) + + check_consistency(card_dict, main_dict["mode"]) + + main_post_scola(parsed_args) + print_message("Post-sCOLA finished.", 1, "control", verbose=parsed_args.verbose) + + + case "sCOLA": + print_message("Running sCOLA.", 1, "control", verbose=parsed_args.verbose) + from scola import main_scola + from parameters_card import parse_arguments_card + + card_dict = parse_arguments_card(parsed_args) + + check_consistency(card_dict, main_dict["mode"]) + + main_scola(parsed_args) + print_message("sCOLA finished.", 1, "control", verbose=parsed_args.verbose) + + + case "alltCOLA" | "allPM": + print_message(f"Running ICs and Simbelmyne in mode {main_dict["mode"]}.", 1, "control", verbose=parsed_args.verbose) + from parameters_card import parse_arguments_card, main_parameter_card + from timestepping import main_timestepping + from ICs import main_ICs + from simbelmyne import main_simbelmyne + from os.path import isfile + from low_level import wait_until_file_exists + + card_dict = main_parameter_card(parsed_args) + + check_consistency(card_dict, main_dict["mode"]) + + main_timestepping(parsed_args) + + ## Check consistency of ICs_gen and ICs + if main_dict["ICs_gen"] == "monofonic": + from parameters_monofonic import parse_arguments_monofonic + monofonic_dict = parse_arguments_monofonic(parsed_args) + if monofonic_dict["output"]+"DM_delta.h5" != card_dict["ICs"]: + raise ValueError(f"ICs {card_dict['ICs']} does not match monofonic output {monofonic_dict['output']+'DM_delta.h5'}") + + print_message("Running initial conditions generator.", 1, "control", verbose=parsed_args.verbose) + main_ICs(parsed_args) + + ## Wait for the initial conditions to be generated + if main_dict["execution"] == "slurm" and card_dict["ICsMode"] == 2: + wait_until_file_exists(card_dict["ICs"]) + if main_dict["execution"] == "slurm" and card_dict["ICsMode"] == 1: + wait_until_file_exists(card_dict["InputWhiteNoise"]) + wait_until_file_exists(card_dict["InputPowerSpectrum"]) + + print_message("Initial conditions generator finished.", 1, "control", verbose=parsed_args.verbose) + + output_is_required = (card_dict["WriteFinalDensity"] or card_dict["WriteFinalSnapshot"]) + if not output_is_required: + print_message("Output is not required. Skipping simbelmyne.", 1, "control", verbose=parsed_args.verbose) + else: + if (card_dict["WriteFinalDensity"] and not isfile(card_dict["OutputFinalDensity"])) or (card_dict["WriteFinalSnapshot"] and not isfile(card_dict["OutputFinalSnapshot"])) or parsed_args.force: + main_simbelmyne(parsed_args) + print_message("Simbelmyne finished.", 1, "control", verbose=parsed_args.verbose) + else: + print_message("Output files already exist. Use -F to overwrite.", 1, "control", verbose=parsed_args.verbose) + + print_message(f"Running {main_dict["mode"]} finished.", 1, "control", verbose=parsed_args.verbose) + + + case "allsCOLA": + print_message(f"Running ICs, pre_sCOLA, sCOLA and post_sCOLA.", 1, "control", verbose=parsed_args.verbose) + from parameters_card import parse_arguments_card, main_parameter_card + from timestepping import main_timestepping + from ICs import main_ICs + from scola import main_scola, main_pre_scola, main_post_scola + + card_dict = main_parameter_card(parsed_args) + + check_consistency(card_dict, main_dict["mode"]) + + main_timestepping(parsed_args) + + ## Check consistency of ICs_gen and OutputLPTPotential + if main_dict["ICs_gen"] == "monofonic": + from parameters_monofonic import parse_arguments_monofonic + monofonic_dict = parse_arguments_monofonic(parsed_args) + if monofonic_dict["output"]+"DM_phi.h5" != card_dict["OutputLPTPotential1"]: + raise ValueError(f"OutputLPTPotential1 {card_dict['OutputLPTPotential1']} does not match monofonic output {monofonic_dict['output']+'DM_phi.h5'}") + if monofonic_dict["output"]+"DM_phi2.h5" != card_dict["OutputLPTPotential2"]: + raise ValueError(f"OutputLPTPotential2 {card_dict['OutputLPTPotential2']} does not match monofonic output {monofonic_dict['output']+'DM_phi2.h5'}") + + print_message("Running initial conditions generator.", 1, "control", verbose=parsed_args.verbose) + main_ICs(parsed_args) + + ## Wait for the initial conditions to be generated + if main_dict["execution"] == "slurm" and card_dict["ICsMode"] == 2: + wait_until_file_exists(card_dict["ICs"]) + if main_dict["execution"] == "slurm" and card_dict["ICsMode"] == 1: + wait_until_file_exists(card_dict["InputWhiteNoise"]) + wait_until_file_exists(card_dict["InputPowerSpectrum"]) + + print_message("Initial conditions generator finished.", 1, "control", verbose=parsed_args.verbose) + + ## If we don't rely on monofonic to generate the initial LPT potentials, use pre_sCOLA + if not main_dict["ICs_gen"] == "monofonic": + print_message("Running pre-sCOLA.", 1, "control", verbose=parsed_args.verbose) + main_pre_scola(parsed_args) + if parsed_args.execution == "slurm": + wait_until_file_exists(card_dict["OutputLPTPotential2"], verbose=parsed_args.verbose, limit=5*60) + print_message("Pre-sCOLA finished.", 1, "control", verbose=parsed_args.verbose) + else: + print_message("ICs_gen is monofonic. Skipping pre-sCOLA.", 1, "control", verbose=parsed_args.verbose) + + print_message("Running sCOLA.", 1, "control", verbose=parsed_args.verbose) + main_scola(parsed_args) + if parsed_args.execution == "slurm": + from tqdm import tqdm + from low_level import progress_bar_from_logfile + for b in tqdm(range(1,parsed_args.N_tiles**3+1), desc="sCOLA", unit="box", disable=(parsed_args.verbose==0)): + progress_bar_from_logfile(main_dict["logdir"]+main_dict["simname"]+".log_"+str(b), desc=f"Box {b}/{parsed_args.N_tiles**3+1}", verbose=parsed_args.verbose, leave=False) + print_message("sCOLA finished.", 1, "control", verbose=parsed_args.verbose) + + print_message("Running post-sCOLA.", 1, "control", verbose=parsed_args.verbose) + main_post_scola(parsed_args) + print_message("Post-sCOLA finished.", 1, "control", verbose=parsed_args.verbose) + + print_message("Running allsCOLA finished.", 1, "control", verbose=parsed_args.verbose) + + + print_ending_module("control", verbose=parsed_args.verbose) + + + + +def check_consistency(card_dict, mode): + ## Check consistency of EvolutionMode and ModulePMCOLA + if mode == "PM" or mode == "allPM": + if card_dict["EvolutionMode"] != 1: + raise ValueError(f"EvolutionMode is not 1: EvolutionMode={card_dict["EvolutionMode"]}") + if card_dict["ModulePMCOLA"] != 1: + raise ValueError(f"ModulePMCOLA is not 1: ModulePMCOLA={card_dict["ModulePMCOLA"]}") + elif mode == "tCOLA" or mode == "alltCOLA": + if card_dict["EvolutionMode"] != 2: + raise ValueError(f"EvolutionMode is not 2: EvolutionMode={card_dict["EvolutionMode"]}") + if card_dict["ModulePMCOLA"] != 1: + raise ValueError(f"ModulePMCOLA is not 1: ModulePMCOLA={card_dict["ModulePMCOLA"]}") + elif mode == "LPT": + if card_dict["ModulePMCOLA"] !=0: + raise ValueError(f"ModulePMCOLA is not 0: ModulePMCOLA={card_dict["ModulePMCOLA"]}") + elif mode == "pre_sCOLA" or mode == "post_sCOLA" or mode == "sCOLA" or mode=="allsCOLA": + if card_dict["EvolutionMode"] != 3: + raise ValueError(f"EvolutionMode is not 3: EvolutionMode={card_dict["EvolutionMode"]}") + if card_dict["ModulePMCOLA"] != 1: + raise ValueError(f"ModulePMCOLA is not 1: ModulePMCOLA={card_dict["ModulePMCOLA"]}") + + +if __name__ == "__main__": + from argparse import ArgumentParser + from args_main import register_arguments_main + from timestepping import register_arguments_timestepping, main_timestepping + from parameters_card import register_arguments_card, main_parameter_card + from cosmo_params import register_arguments_cosmo + from parameters_monofonic import register_arguments_monofonic + from slurm_submission import register_arguments_slurm + from low_level import wait_until_file_exists + + parser = ArgumentParser(description="Run sCOLA.") + register_arguments_main(parser) + register_arguments_timestepping(parser) + register_arguments_monofonic(parser) + register_arguments_slurm(parser) + register_arguments_card(parser) + register_arguments_cosmo(parser) + parsed_args = parser.parse_args() + main(parsed_args) \ No newline at end of file diff --git a/monofonic.py b/monofonic.py index 6c81911..07b4c6d 100644 --- a/monofonic.py +++ b/monofonic.py @@ -86,7 +86,7 @@ if __name__ == "__main__": from argparse import ArgumentParser from parameters_monofonic import register_arguments_monofonic from args_main import register_arguments_main - from parameters_card import register_arguments_card + from parameters_card import register_arguments_card_for_ICs from cosmo_params import register_arguments_cosmo from slurm_submission import register_arguments_slurm @@ -95,7 +95,7 @@ if __name__ == "__main__": register_arguments_main(parser) register_arguments_monofonic(parser) register_arguments_slurm(parser) - register_arguments_card(parser) + register_arguments_card_for_ICs(parser) register_arguments_cosmo(parser) parsed_args = parser.parse_args diff --git a/parameters_card.py b/parameters_card.py index 4ff4903..223b5ac 100644 --- a/parameters_card.py +++ b/parameters_card.py @@ -49,6 +49,26 @@ def register_arguments_card(parser:ArgumentParser): parser.add_argument("--OutputRngStateLPT", type=str, default=None, help="Output RNG state file.") +def register_arguments_card_for_ICs(parser:ArgumentParser): + parser.add_argument("-Np","--N_particles", type=int, default=128, help="Number of particles per axis.") + parser.add_argument("-Nlpt","--N_LPT_mesh", type=int, default=None, help="Number of mesh points per axis for the LPT mesh.") + parser.add_argument("-L","--L", type=float, default=100.0, help="Size of the simulation box (in Mpc/h).") + parser.add_argument("-corner","--corner", type=float, nargs=3, default=[0.0, 0.0, 0.0], help="Corner of the simulation box.") + parser.add_argument("--ICsMode", type=int, default=None, help="Initial conditions mode.") + parser.add_argument("--ICs", type=str, default=None, help="Initial conditions file.") + parser.add_argument("--InputWhiteNoise", type=str, default=None, help="Input white noise file.") + parser.add_argument("--InputPowerSpectrum", type=str, default=None, help="Input power spectrum file.") + + +def register_arguments_card_for_timestepping(parser:ArgumentParser): + parser.add_argument("-zini","--RedshiftLPT", type=float, default=19.0, help="Redshift for the LPT evolution.") + parser.add_argument("-zfinal","--RedshiftFCs", type=float, default=0.0, help="Redshift for the FCs evolution.") + parser.add_argument("--TimeSteppingFileName", type=str, default=None, help="Time stepping file.") + parser.add_argument("-lc","--GenerateLightcone", type=bool, default=False, help="Generate lightcone.") + + + + def parse_arguments_card(parsed_args): """ Parse the arguments for the parameter card. @@ -56,9 +76,9 @@ def parse_arguments_card(parsed_args): from args_main import parse_arguments_main from cosmo_params import parse_arguments_cosmo - param_file=parsed_args.paramfile cosmo_dict=parse_arguments_cosmo(parsed_args) card_dict=dict( + paramfile=parsed_args.paramfile, N_particles=parsed_args.N_particles, N_LPT_mesh=parsed_args.N_LPT_mesh, N_PM_mesh=parsed_args.N_PM_mesh, @@ -117,8 +137,8 @@ def parse_arguments_card(parsed_args): ligthcone_prefix = "lightcone_" if card_dict["GenerateLightcone"] else "" ## Now we set the parameters that are None to the default values - if param_file is None: - param_file = main_dict["paramdir"]+"parameters_"+main_dict["simname"]+".sbmy" + if card_dict["paramfile"] is None: + card_dict["paramfile"] = main_dict["paramdir"]+"parameters_"+main_dict["simname"]+".sbmy" if card_dict["N_LPT_mesh"] is None: card_dict["N_LPT_mesh"] = card_dict["N_particles"] # Default is the same as the number of particles if card_dict["N_PM_mesh"] is None: @@ -172,7 +192,7 @@ def parse_arguments_card(parsed_args): if card_dict["OutputFinalDensity"] is None: card_dict["OutputFinalDensity"] = main_dict["resultdir"]+ligthcone_prefix+"final_density_"+main_dict["simname"]+".h5" if card_dict["OutputTilesBase"] is None: - card_dict["OutputTilesBase"] = main_dict["resultdir"]+"sCOLA_tile_"+main_dict["simname"] + card_dict["OutputTilesBase"] = main_dict["workdir"]+"sCOLA_tile" if card_dict["OutputLPTPotential1"] is None: card_dict["OutputLPTPotential1"] = main_dict["workdir"]+"initial_conditions_DM_phi.h5" if card_dict["OutputLPTPotential2"] is None: @@ -184,7 +204,69 @@ def parse_arguments_card(parsed_args): if card_dict["OutputRngStateLPT"] is None: card_dict["OutputRngStateLPT"] = main_dict["workdir"]+"rng_state.h5" - return card_dict, param_file + return card_dict + + + +def parse_arguments_card_for_ICs(parsed_args): + """ + Parse the arguments for the parameter card for ICs. + """ + from args_main import parse_arguments_main + + main_dict = parse_arguments_main(parsed_args) + card_dict = dict( + N_particles=parsed_args.N_particles, + N_LPT_mesh=parsed_args.N_LPT_mesh, + L=parsed_args.L, + corner=parsed_args.corner, + ICsMode=parsed_args.ICsMode, + ICs=parsed_args.ICs, + InputWhiteNoise=parsed_args.InputWhiteNoise, + InputPowerSpectrum=parsed_args.InputPowerSpectrum, + ) + + ## Now we set the parameters that are None to the default values + if card_dict["N_LPT_mesh"] is None: + card_dict["N_LPT_mesh"] = card_dict["N_particles"] # Default is the same as the number of particles + if card_dict["InputWhiteNoise"] is None: + card_dict["InputWhiteNoise"] = main_dict["workdir"]+"white_noise.h5" + if card_dict["ICs"] is None: + card_dict["ICs"] = main_dict["workdir"]+"initial_conditions_DM_delta.h5" + if card_dict["InputPowerSpectrum"] is None: + card_dict["InputPowerSpectrum"] = main_dict["workdir"]+"power_spectrum.h5" + if card_dict["ICsMode"] is None: + match main_dict["ICs_gen"]: + case "ext" | "monofonic": + card_dict["ICsMode"] = 2 + case "sbmy": + card_dict["ICsMode"] = 1 + case _: + raise ValueError(f"ICs generator {main_dict['ICs_gen']} not recognized.") + + return card_dict + + + +def parse_arguments_card_for_timestepping(parsed_args): + """ + Parse the arguments for the parameter card for timestepping. + """ + from args_main import parse_arguments_main + + main_dict = parse_arguments_main(parsed_args) + card_dict = dict( + RedshiftLPT=parsed_args.RedshiftLPT, + RedshiftFCs=parsed_args.RedshiftFCs, + TimeSteppingFileName=parsed_args.TimeSteppingFileName, + GenerateLightcone=parsed_args.GenerateLightcone, + ) + + ## Now we set the parameters that are None to the default values + if card_dict["TimeSteppingFileName"] is None: + card_dict["TimeSteppingFileName"] = main_dict["paramdir"]+"time_stepping.h5" + + return card_dict @@ -341,13 +423,14 @@ def main_parameter_card(parsed_args): from low_level import print_message, print_starting_module, print_ending_module print_starting_module("card", verbose=parsed_args.verbose) print_message("Parsing arguments for the parameter card.", 1, "card", verbose=parsed_args.verbose) - card_dict, param_file = parse_arguments_card(parsed_args) + card_dict = parse_arguments_card(parsed_args) + paramfile = card_dict["paramfile"] parameter_card_dict = create_parameter_card_dict(**card_dict) - if isfile(param_file) and not parsed_args.force: - print_message(f"Parameter card {param_file} exists. Use --force to overwrite.", 1, "card", verbose=parsed_args.verbose) + if isfile(paramfile) and not parsed_args.force: + print_message(f"Parameter card {paramfile} exists. Use --force to overwrite.", 1, "card", verbose=parsed_args.verbose) return card_dict - write_parameter_card(parameter_card_dict, param_file, verbose=parsed_args.verbose) - print_message(f"Parameter card written to {param_file}", 2, "card", verbose=parsed_args.verbose) + write_parameter_card(parameter_card_dict, paramfile, verbose=parsed_args.verbose) + print_message(f"Parameter card written to {paramfile}", 2, "card", verbose=parsed_args.verbose) print_ending_module("card", verbose=parsed_args.verbose) return card_dict diff --git a/parameters_monofonic.py b/parameters_monofonic.py index ac934dd..5554517 100644 --- a/parameters_monofonic.py +++ b/parameters_monofonic.py @@ -17,11 +17,11 @@ def parse_arguments_monofonic(parsed_args): Parse the arguments for the monofonIC parameters. """ from args_main import parse_arguments_main - from parameters_card import parse_arguments_card + from parameters_card import parse_arguments_card_for_ICs from cosmo_params import parse_arguments_cosmo main_dict = parse_arguments_main(parsed_args) - card_dict, _ = parse_arguments_card(parsed_args) + card_dict = parse_arguments_card_for_ICs(parsed_args) cosmo_dict = parse_arguments_cosmo(parsed_args) monofonic_dict = dict( @@ -146,13 +146,13 @@ def main_parameters_monofonic(parsed_args): if __name__ == "__main__": from args_main import register_arguments_main - from parameters_card import register_arguments_card + from parameters_card import register_arguments_card_for_ICs from cosmo_params import register_arguments_cosmo parser = ArgumentParser(description="Create monofonIC configuration file.") register_arguments_main(parser) register_arguments_monofonic(parser) - register_arguments_card(parser) + register_arguments_card_for_ICs(parser) register_arguments_cosmo(parser) parsed_args = parser.parse_args() diff --git a/scola.py b/scola.py index af94f25..01e6d32 100644 --- a/scola.py +++ b/scola.py @@ -1,84 +1,115 @@ def main_scola(parsed_args): from args_main import parse_arguments_main - from low_level import print_starting_module, print_message, print_ending_module + from low_level import print_starting_module, print_message, print_ending_module, progress_bar_from_logfile, wait_until_file_exists from os.path import isfile import subprocess print_starting_module("scola", verbose=parsed_args.verbose) main_dict = parse_arguments_main(parsed_args) - param_file=parsed_args.paramfile - if param_file is None: - param_file = main_dict["paramdir"]+"parameters_"+main_dict["simname"]+".sbmy" - log_file = main_dict["logdir"]+main_dict["mode"]+".log" + paramfile=parsed_args.paramfile + if paramfile is None: + paramfile = main_dict["paramdir"]+"parameters_"+main_dict["simname"]+".sbmy" + log_file = main_dict["logdir"]+main_dict["simname"]+".log" nboxes_tot = int(parsed_args.N_tiles**3) if parsed_args.execution == "local": + from parameters_card import parse_arguments_card + from tqdm import tqdm + card_dict = parse_arguments_card(parsed_args) print_message("Running sCOLA in local mode.", 1, "scola", verbose=parsed_args.verbose) - + + tile_bar = tqdm(total=nboxes_tot, desc="sCOLA", unit="box", disable=(parsed_args.verbose!=1)) # The progress bar cannot work if there are print statements in the loop for b in range(1,nboxes_tot+1): - print_message(f"Running box {b}/{nboxes_tot}.", 2, "scola", verbose=parsed_args.verbose) - command_args = ["scola", param_file, log_file, "-b", str(b)] + if not isfile(card_dict["OutputTilesBase"]+str(b)+".h5") or parsed_args.force: + if parsed_args.verbose > 1: + print_message(f"Running box {b}/{nboxes_tot}.", 2, "scola", verbose=parsed_args.verbose) + command_args = ["scola", paramfile, log_file, "-b", str(b)] + if parsed_args.verbose < 2: + subprocess.Popen(command_args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) # Use Popen instead of run to be able to display the progress bar + progress_bar_from_logfile(log_file+f"_{b}", desc=f" Box {b}/{nboxes_tot+1}", verbose=parsed_args.verbose, leave=False) + else: + subprocess.run(command_args) - if parsed_args.verbose < 2: - from io import BytesIO - from low_level import stdout_redirector, stderr_redirector - - f = BytesIO() - g = BytesIO() - with stdout_redirector(f): - with stderr_redirector(g): - subprocess.run(command_args) - g.close() - f.close() + if parsed_args.verbose > 1: + print_message(f"Box {b}/{nboxes_tot} finished.", 3, "scola", verbose=parsed_args.verbose) else: - subprocess.run(command_args) - - print_message(f"Box {b}/{nboxes_tot} finished.", 3, "scola", verbose=parsed_args.verbose) + if parsed_args.verbose > 1: + print_message(f"Box {b}/{nboxes_tot} already exists. Use -F to overwrite.", 2, "scola", verbose=parsed_args.verbose) + tile_bar.update(1) + tile_bar.close() print_message("sCOLA finished.", 1, "scola", verbose=parsed_args.verbose) elif parsed_args.execution == "slurm": from slurm_submission import create_slurm_script, parse_arguments_slurm from args_main import parse_arguments_main - print_message("Running monofonic in slurm mode.", 1, "scola", verbose=parsed_args.verbose) + print_message("Running scola in slurm mode.", 1, "scola", verbose=parsed_args.verbose) slurm_dict=parse_arguments_slurm(parsed_args) main_dict=parse_arguments_main(parsed_args) - slurm_script = slurm_dict["scripts"]+"scola.sh" - if not isfile(slurm_script): - print_message(f"SLURM script {slurm_script} does not exist. Creating it.", 2, "scola", verbose=parsed_args.verbose) - create_slurm_script( - slurm_template=slurm_dict["scola_template"], - slurm_script=slurm_script, - job="scola", - job_config_file=param_file, - job_log=log_file, - array=(1, nboxes_tot+1), - ) - print_message(f"SLURM script written to {slurm_script}.", 3, "scola", verbose=parsed_args.verbose) + if have_no_tiles(parsed_args) or parsed_args.force: + ## Submit all boxes + print_message("Submitting all boxes.", 2, "scola", verbose=parsed_args.verbose) + slurm_script = slurm_dict["scripts"]+"scola.sh" + + if not isfile(slurm_script): + print_message(f"SLURM script {slurm_script} does not exist. Creating it.", 2, "scola", verbose=parsed_args.verbose) + create_slurm_script( + slurm_template=slurm_dict["scola_template"], + slurm_script=slurm_script, + job="scola", + job_config_file=paramfile, + job_log=log_file, + array=(1, nboxes_tot), + ) + print_message(f"SLURM script written to {slurm_script}.", 3, "scola", verbose=parsed_args.verbose) + else: + print_message(f"SLURM script {slurm_script} exists.", 2, "scola", verbose=parsed_args.verbose) + + print_message(f"Submitting scola job to SLURM.", 1, "scola", verbose=parsed_args.verbose) + command_args = ["sbatch", slurm_script] + + if parsed_args.verbose < 2: + subprocess.run(command_args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + else: + subprocess.run(command_args) + + print_message("sCOLA job submitted.", 2, "scola", verbose=parsed_args.verbose) + else: - print_message(f"SLURM script {slurm_script} exists.", 2, "scola", verbose=parsed_args.verbose) - - print_message(f"Submitting monofonic job to SLURM.", 1, "scola", verbose=parsed_args.verbose) + ## Submit missing boxes + missing_tiles_arrays = get_missing_tiles_arrays(parsed_args) + print_message(f"Submitting missing boxes: {missing_tiles_arrays}.", 2, "scola", verbose=parsed_args.verbose) - command_args = ["sbatch", slurm_script] + for missing_tiles in missing_tiles_arrays: + slurm_script = slurm_dict["scripts"]+"scola_"+str(missing_tiles[0])+"_"+str(missing_tiles[-1])+".sh" - if parsed_args.verbose < 2: - from io import BytesIO - from low_level import stdout_redirector, stderr_redirector + if not isfile(slurm_script): + print_message(f"SLURM script {slurm_script} does not exist. Creating it.", 2, "scola", verbose=parsed_args.verbose) + create_slurm_script( + slurm_template=slurm_dict["scola_template"], + slurm_script=slurm_script, + job="scola", + job_config_file=paramfile, + job_log=log_file, + array=missing_tiles, + ) + print_message(f"SLURM script written to {slurm_script}.", 3, "scola", verbose=parsed_args.verbose) + else: + print_message(f"SLURM script {slurm_script} exists.", 2, "scola", verbose=parsed_args.verbose) + + print_message(f"Submitting scola job to SLURM.", 1, "scola", verbose=parsed_args.verbose) + command_args = ["sbatch", slurm_script] - f = BytesIO() - g = BytesIO() - with stdout_redirector(f): - with stderr_redirector(g): + if parsed_args.verbose < 2: + subprocess.run(command_args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + else: subprocess.run(command_args) - g.close() - f.close() - else: - subprocess.run(command_args) - print_message("Monofonic job submitted.", 2, "scola", verbose=parsed_args.verbose) + print_message("sCOLA job submitted.", 2, "scola", verbose=parsed_args.verbose) + + else: raise ValueError(f"Execution mode {parsed_args.execution} not recognized.") @@ -86,9 +117,114 @@ def main_scola(parsed_args): print_ending_module("scola", verbose=parsed_args.verbose) -def main_pre_post_scola(parsed_args): - from simbelmyne import main_simbelmyne - main_simbelmyne(parsed_args) +def main_pre_scola(parsed_args): + """ + Pre-scola function. It calls simbelmyne to generate the LPT phi1 and phi2 potential fields. + If they already exist, it does nothing. If all tiles exist, raises error. + """ + from os.path import isfile + from low_level import print_starting_module, print_message, print_ending_module + from parameters_card import parse_arguments_card + print_starting_module("pre-scola", verbose=parsed_args.verbose) + card_dict = parse_arguments_card(parsed_args) + + if not isfile(card_dict["OutputLPTPotential1"]) or not isfile(card_dict["OutputLPTPotential2"]) or parsed_args.force: + if not have_all_tiles(parsed_args): + print_message("Running pre-scola.", 1, "pre-scola", verbose=parsed_args.verbose) + from simbelmyne import main_simbelmyne + main_simbelmyne(parsed_args) + else: + raise NotImplementedError("All tiles exists, so calling simbelmyne would generate the final output instead of the LPT potentials.") + else: + print_message("LPT potentials already exist. Use -F to overwrite.", 1, "pre-scola", verbose=parsed_args.verbose) + + print_ending_module("pre-scola", verbose=parsed_args.verbose) + + +def main_post_scola(parsed_args): + """ + Post-scola function. It calls simbelmyne to generate the final output from all tiles. + If the output already exists, it does nothing. If tiles are missing, print the missing tiles. + """ + from os.path import isfile + from low_level import print_starting_module, print_message, print_ending_module + from parameters_card import parse_arguments_card + print_starting_module("post-scola", verbose=parsed_args.verbose) + card_dict = parse_arguments_card(parsed_args) + + output_is_required = (card_dict["WriteFinalDensity"] or card_dict["WriteFinalSnapshot"]) + if not output_is_required: + print_message("Output is not required. Skipping post-scola.", 1, "post-scola", verbose=parsed_args.verbose) + print_ending_module("post-scola", verbose=parsed_args.verbose) + return + + if (card_dict["WriteFinalDensity"] and not isfile(card_dict["OutputFinalDensity"])) or (card_dict["WriteFinalSnapshot"] and not isfile(card_dict["OutputFinalSnapshot"])) or parsed_args.force: + if have_all_tiles(parsed_args): + print_message("Running post-scola.", 1, "post-scola", verbose=parsed_args.verbose) + from simbelmyne import main_simbelmyne + main_simbelmyne(parsed_args) + else: + missing_tiles_arrays = get_missing_tiles_arrays(parsed_args) + count_missing_tiles = sum([m[1]-m[0]+1 for m in missing_tiles_arrays]) + nboxes_tot = int(parsed_args.N_tiles**3) + print_message(f"All tiles are not available. Missing {count_missing_tiles} out of {nboxes_tot} ({100*count_missing_tiles/nboxes_tot:.1f}%)", 1, "post-scola", verbose=parsed_args.verbose) + print_message(f"Missing tiles: {missing_tiles_arrays}", 2, "post-scola", verbose=parsed_args.verbose) + else: + print_message("Output already exists. Use -F to overwrite.", 1, "post-scola", verbose=parsed_args.verbose) + + print_ending_module("post-scola", verbose=parsed_args.verbose) + + + +def have_all_tiles(parsed_args): + from os.path import isfile + from parameters_card import parse_arguments_card + + card_dict = parse_arguments_card(parsed_args) + nboxes_tot = int(parsed_args.N_tiles**3) + all_tiles = True + for b in range(1,nboxes_tot+1): + if not isfile(card_dict["OutputTilesBase"]+str(b)+".h5"): + all_tiles = False + return all_tiles + + +def have_no_tiles(parsed_args): + from os.path import isfile + from parameters_card import parse_arguments_card + + card_dict = parse_arguments_card(parsed_args) + nboxes_tot = int(parsed_args.N_tiles**3) + no_tiles = True + for b in range(1,nboxes_tot+1): + if isfile(card_dict["OutputTilesBase"]+str(b)+".h5"): + no_tiles = False + return no_tiles + + +def get_missing_tiles_arrays(parsed_args): + from os.path import isfile + from parameters_card import parse_arguments_card + + card_dict = parse_arguments_card(parsed_args) + nboxes_tot = int(parsed_args.N_tiles**3) + missing_tiles_arrays = [] + in_sequence_of_missing = False + for b in range(1,nboxes_tot+1): + if not isfile(card_dict["OutputTilesBase"]+str(b)+".h5"): + if not in_sequence_of_missing: + missing_tiles_arrays.append([b]) + in_sequence_of_missing = True + elif in_sequence_of_missing: + missing_tiles_arrays[-1].append(b-1) + in_sequence_of_missing = False + if in_sequence_of_missing: + missing_tiles_arrays[-1].append(nboxes_tot) + for m in missing_tiles_arrays: + assert len(m) == 2 + assert m[0] <= m[1] + return missing_tiles_arrays + if __name__ == "__main__": from argparse import ArgumentParser @@ -98,10 +234,9 @@ if __name__ == "__main__": from cosmo_params import register_arguments_cosmo from parameters_monofonic import register_arguments_monofonic from slurm_submission import register_arguments_slurm - from low_level import print_starting_module, print_message, print_ending_module, wait_until_file_exists + from low_level import wait_until_file_exists parser = ArgumentParser(description="Run sCOLA.") - # TODO: reduce the volume of arguments register_arguments_main(parser) register_arguments_timestepping(parser) register_arguments_monofonic(parser) @@ -113,11 +248,11 @@ if __name__ == "__main__": card_dict = main_parameter_card(parsed_args) nboxes_tot = int(parsed_args.N_tiles**3) main_timestepping(parsed_args) - main_pre_post_scola(parsed_args) + main_pre_scola(parsed_args) if parsed_args.execution == "slurm": - wait_until_file_exists(card_dict["OutputLPTPotential2"]) + wait_until_file_exists(card_dict["OutputLPTPotential2"], verbose=parsed_args.verbose, limit=5*60) main_scola(parsed_args) if parsed_args.execution == "slurm": for b in range(1,nboxes_tot+1): - wait_until_file_exists(f"{card_dict['OutputTilesBase']}{b}.h5") - main_pre_post_scola(parsed_args) \ No newline at end of file + wait_until_file_exists(f"{card_dict['OutputTilesBase']}{b}.h5", verbose=parsed_args.verbose, limit=5*60) + main_post_scola(parsed_args) \ No newline at end of file diff --git a/simbelmyne.py b/simbelmyne.py index 236a2f8..cb3d377 100644 --- a/simbelmyne.py +++ b/simbelmyne.py @@ -1,33 +1,25 @@ def main_simbelmyne(parsed_args): from args_main import parse_arguments_main - from low_level import print_starting_module, print_message, print_ending_module + from low_level import print_starting_module, print_message, print_ending_module, progress_bar_from_logfile from os.path import isfile import subprocess print_starting_module("simbelmyne", verbose=parsed_args.verbose) main_dict = parse_arguments_main(parsed_args) - param_file=parsed_args.paramfile - if param_file is None: - param_file = main_dict["paramdir"]+"parameters_"+main_dict["simname"]+".sbmy" - log_file = main_dict["logdir"]+main_dict["mode"]+".log" + paramfile=parsed_args.paramfile + if paramfile is None: + paramfile = main_dict["paramdir"]+"parameters_"+main_dict["simname"]+".sbmy" + log_file = main_dict["logdir"]+main_dict["simname"]+".log" if parsed_args.execution == "local": print_message("Running Simbelyne in local mode.", 1, "simbelmyne", verbose=parsed_args.verbose) - command_args = ["simbelmyne", param_file, log_file] + command_args = ["simbelmyne", paramfile, log_file] if parsed_args.verbose < 2: - from io import BytesIO - from low_level import stdout_redirector, stderr_redirector - - f = BytesIO() - g = BytesIO() - with stdout_redirector(f): - with stderr_redirector(g): - subprocess.run(command_args) - g.close() - f.close() + subprocess.Popen(command_args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) # Use Popen instead of run to be able to display the progress bar + progress_bar_from_logfile(log_file, desc=main_dict["simname"], verbose=parsed_args.verbose, leave=False) else: subprocess.run(command_args) @@ -39,7 +31,7 @@ def main_simbelmyne(parsed_args): print_message("Running simbelmyne in slurm mode.", 1, "simbelmyne", verbose=parsed_args.verbose) slurm_dict=parse_arguments_slurm(parsed_args) main_dict=parse_arguments_main(parsed_args) - slurm_script = slurm_dict["scripts"]+"simbelmyne.sh" + slurm_script = slurm_dict["scripts"]+"simbelmyne_"+main_dict["simname"]+".sh" if not isfile(slurm_script): print_message(f"SLURM script {slurm_script} does not exist. Creating it.", 2, "simbelmyne", verbose=parsed_args.verbose) @@ -47,7 +39,7 @@ def main_simbelmyne(parsed_args): slurm_template=slurm_dict["simbelmyne_template"], slurm_script=slurm_script, job="simbelmyne", - job_config_file=param_file, + job_config_file=paramfile, job_log=log_file, ) print_message(f"SLURM script written to {slurm_script}.", 3, "simbelmyne", verbose=parsed_args.verbose) @@ -59,20 +51,12 @@ def main_simbelmyne(parsed_args): command_args = ["sbatch", slurm_script] if parsed_args.verbose < 2: - from io import BytesIO - from low_level import stdout_redirector, stderr_redirector - - f = BytesIO() - g = BytesIO() - with stdout_redirector(f): - with stderr_redirector(g): - subprocess.run(command_args) - g.close() - f.close() + subprocess.run(command_args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) else: subprocess.run(command_args) - print_message("Monofonic job submitted.", 2, "simbelmyne", verbose=parsed_args.verbose) + print_message("Simbelmyne job submitted.", 2, "simbelmyne", verbose=parsed_args.verbose) + progress_bar_from_logfile(log_file, desc=main_dict["simname"], verbose=parsed_args.verbose) else: raise ValueError(f"Execution mode {parsed_args.execution} not recognized.") diff --git a/slurm_submission.py b/slurm_submission.py index 442a2e2..fab0025 100644 --- a/slurm_submission.py +++ b/slurm_submission.py @@ -108,6 +108,7 @@ def create_slurm_script(slurm_template:str, job:str, job_config_file:str, job_log:str, + array:tuple|None=None, ): """ Creates a SLURM submission script based on the provided template. @@ -116,6 +117,12 @@ def create_slurm_script(slurm_template:str, - simbelmyne - scola """ + + if array is not None and job != "scola": + raise ValueError(f"Array job range provided for job type {job}.") + if array is None and job == "scola": + raise ValueError(f"Array job range not provided for job type {job}.") + from os.path import isfile if not isfile(slurm_template): raise FileNotFoundError(f"SLURM template {slurm_template} does not exist.") @@ -127,6 +134,10 @@ def create_slurm_script(slurm_template:str, # Create the script file with open(slurm_script, "w") as f: for line in template: + if array is not None and "--array" in line: + line = line.replace("%a-%b",f"{array[0]}-{array[1]}") + if array is not None and ("--output" in line or "--error" in line): + line = line.replace("%j","%a_%A") f.write(line) # Add the job command @@ -136,7 +147,7 @@ def create_slurm_script(slurm_template:str, case "simbelmyne": f.write(f"{job} {job_config_file} {job_log}") case "scola": - f.write(f"{job} {job_config_file} {job_log} "+"${SLURM_ARRAY_TASK_ID}") + f.write(f"{job} {job_config_file} {job_log} "+"-b ${SLURM_ARRAY_TASK_ID}") case _: raise ValueError(f"Job type {job} not recognized.") diff --git a/timestepping.py b/timestepping.py index 76e2697..38b73ab 100644 --- a/timestepping.py +++ b/timestepping.py @@ -18,12 +18,10 @@ def parse_arguments_timestepping(parsed_args): """ Parse the arguments for the timestepping. """ - from args_main import parse_arguments_main - from parameters_card import parse_arguments_card + from parameters_card import parse_arguments_card_for_timestepping from cosmo_params import parse_arguments_cosmo, z2a - main_dict = parse_arguments_main(parsed_args) - card_dict, _ = parse_arguments_card(parsed_args) + card_dict = parse_arguments_card_for_timestepping(parsed_args) cosmo_dict = parse_arguments_cosmo(parsed_args) timestepping_dict = dict( @@ -113,14 +111,14 @@ def main_timestepping(parsed_args): if __name__ == "__main__": from args_main import register_arguments_main - from parameters_card import register_arguments_card + from parameters_card import register_arguments_card_for_timestepping from cosmo_params import register_arguments_cosmo parser = ArgumentParser(description="Create timestepping file.") # TODO: reduce the volume of arguments register_arguments_main(parser) register_arguments_timestepping(parser) - register_arguments_card(parser) + register_arguments_card_for_timestepping(parser) register_arguments_cosmo(parser) parsed_args = parser.parse_args() main_timestepping(parsed_args) \ No newline at end of file