sbmy_control/main.py

241 lines
No EOL
12 KiB
Python

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}", 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)
case _:
raise ValueError(f"Unknown mode: {main_dict['mode']}")
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)