improvements
This commit is contained in:
parent
3dc03cec4f
commit
e488275523
11 changed files with 640 additions and 125 deletions
255
scola.py
255
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)
|
||||
wait_until_file_exists(f"{card_dict['OutputTilesBase']}{b}.h5", verbose=parsed_args.verbose, limit=5*60)
|
||||
main_post_scola(parsed_args)
|
||||
Loading…
Add table
Add a link
Reference in a new issue