diff --git a/slurm_submission.py b/slurm_submission.py new file mode 100644 index 0000000..f61e4ab --- /dev/null +++ b/slurm_submission.py @@ -0,0 +1,179 @@ +from argparse import ArgumentParser +from args_main import parse_arguments_main + +path_to_monofonic_binary = "/home/aubin/monofonic/build/monofonIC" + +def register_arguments_slurm(parser:ArgumentParser): + """ + Register the arguments for the SLURM parameters used in the slurm submissions scripts for the different binary calls. + Binary calls: + - monofonic + - simbelmyne + - scola + """ + parser.add_argument("-smf","--slurm_monofonic", type=str, default=None, help="Path to the monofonic SLURM submission script template.") + parser.add_argument("-ssbmy","--slurm_simbelmyne", type=str, default=None, help="Path to the simbelmyne SLURM submission script template.") + parser.add_argument("-sscola","--slurm_scola", type=str, default=None, help="Path to the scola SLURM submission script template.") + parser.add_argument("--slurm_logs", type=str, default=None, help="Path to the directory where the SLURM logs will be saved.") + parser.add_argument("--slurm_scripts", type=str, default=None, help="Path to the directory where the SLURM scripts will be saved.") + + +def parse_arguments_slurm(parsed_args): + """ + Parse the arguments for the SLURM parameters used in the slurm submissions scripts for the different binary calls. + """ + from pathlib import Path + main_dict = parse_arguments_main(parsed_args) + + slurm_dict = dict( + monofonic_template=parsed_args.slurm_monofonic, + simbelmyne_template=parsed_args.slurm_simbelmyne, + scola_template=parsed_args.slurm_scola, + logs=parsed_args.slurm_logs, + scripts=parsed_args.slurm_scripts + ) + + if slurm_dict["monofonic_template"] is None: + slurm_dict["monofonic_template"]=main_dict["paramdir"]+"slurm_monofonic.template" + if slurm_dict["simbelmyne_template"] is None: + slurm_dict["simbelmyne_template"]=main_dict["paramdir"]+"slurm_simbelmyne.template" + if slurm_dict["scola_template"] is None: + slurm_dict["scola_template"]=main_dict["paramdir"]+"slurm_scola.template" + if slurm_dict["logs"] is None: + slurm_dict["logs"]=main_dict["directory"]+"slurm_logs/" + if slurm_dict["scripts"] is None: + slurm_dict["scripts"]=main_dict["directory"]+"slurm_scripts/" + + + Path(slurm_dict["logs"]).mkdir(parents=True, exist_ok=True) + Path(slurm_dict["scripts"]).mkdir(parents=True, exist_ok=True) + + return slurm_dict + + +def create_slurm_template( + slurm_template:str, + job_name:str, + ntasks:int, + nthreads:int, + partition:str, + time:str, + mem:int, + log_out:str, + log_err:str, + array:tuple|None=None, + ): + """ + Creates a SLURM submission script template. + """ + + with open(slurm_template, "w") as f: + f.write("#!/bin/bash\n") + f.write(f"#SBATCH --job-name={job_name}\n") + f.write(f"#SBATCH --ntasks={ntasks}\n") + f.write(f"#SBATCH --cpus-per-task={nthreads}\n") + f.write(f"#SBATCH --partition={partition}\n") + f.write(f"#SBATCH --time={time}\n") + f.write(f"#SBATCH --mem={mem}G\n") + + if array is not None: + f.write(f"#SBATCH --array={array[0]}-{array[1]}\n") + f.write(f"#SBATCH --output={log_out}%x_%a_%A.out\n") + f.write(f"#SBATCH --error={log_err}%x_%a_%A.err\n") + else: + f.write(f"#SBATCH --output={log_out}%x_%j.out\n") + f.write(f"#SBATCH --error={log_err}%x_%j.err\n") + + f.write("\n") + + f.write("echo '################## SLURM VARIABLES ##################'\n") + f.write("echo SLURM_JOB_ID: $SLURM_JOB_ID\n") + f.write("echo SLURM_JOB_NAME: $SLURM_JOB_NAME\n") + f.write("echo SLURM_JOB_NODELIST: $SLURM_JOB_NODELIST\n") + f.write("echo SLURM_NNODES: $SLURM_NNODES\n") + f.write("echo SLURM_NTASKS: $SLURM_NTASKS\n") + f.write("echo SLURM_CPUS_PER_TASK: $SLURM_CPUS_PER_TASK\n") + f.write("echo SLURM_JOB_CPUS_PER_NODE: $SLURM_JOB_CPUS_PER_NODE\n") + f.write("echo SLURM_MEM_PER_CPU: $SLURM_MEM_PER_CPU\n") + f.write("echo SLURM_MEM_PER_NODE: $SLURM_MEM_PER_NODE\n") + f.write("echo '#####################################################'\n") + f.write("\n\n") + + f.write(f"export OMP_NUM_THREADS={nthreads}\n\n") + + + +def create_slurm_script(slurm_template:str, + slurm_script:str, + job:str, + job_config_file:str, + job_log:str, + ): + """ + Creates a SLURM submission script based on the provided template. + For three different kind of jobs: + - monofonic + - simbelmyne + - scola + """ + from os.path import isfile + if not isfile(slurm_template): + raise FileNotFoundError(f"SLURM template {slurm_template} does not exist.") + + # Copy template content + with open(slurm_template, "r") as f: + template = f.readlines() + + # Create the script file + with open(slurm_script, "w") as f: + for line in template: + f.write(line) + + # Add the job command + match job: + case "monofonic": + f.write(f"{path_to_monofonic_binary} {job_config_file} > {job_log}") + case "simbelmyne" | "scola": + f.write(f"{job} {job_config_file} {job_log}") + case _: + raise ValueError(f"Job type {job} not recognized.") + + + +if __name__ == "__main__": + from argparse import ArgumentParser + + parser = ArgumentParser(description="Generate slurm submission templates.") + parser.add_argument("-j","--job", type=str, default="monofonic", help="Job type: monofonic, simbelmyne, scola.") + parser.add_argument("-N","--ntasks", type=int, default=1, help="Number of tasks.") + parser.add_argument("-n","--nthreads", type=int, default=32, help="Number of threads per task.") + parser.add_argument("-p","--partition", type=str, default="comp,pscomp,compl", help="Partition to use.") + parser.add_argument("-t","--time", type=str, default="0-00:10:00", help="Time limit.") + parser.add_argument("-m","--mem", type=int, default=64, help="Memory limit.") + parser.add_argument("-d", "--directory", type=str, default="./", help="Main directory where the output will be saved (if other dir and filenames are not specified).") + parser.add_argument("-o","--log_out", type=str, default=None, help="File root for the output logs.") + parser.add_argument("-e","--log_err", type=str, default=None, help="File root for the error logs.") + parser.add_argument("-a","--array", type=int, nargs=2, default=None, help="Array job range.") + parser.add_argument("-s","--slurm_template", type=str, default=None, help="Path to the SLURM template.") + parser.add_argument("-jn","--job_name", type=str, default=None, help="Job name.") + + parsed_args = parser.parse_args() + + job_name = parsed_args.job if parsed_args.job_name is None else parsed_args.job_name + slurm_template = parsed_args.slurm_template if parsed_args.slurm_template is not None else f"{parsed_args.directory}params/slurm_{job_name}.template" + log_out = parsed_args.log_out if parsed_args.log_out is not None else f"{parsed_args.directory}slurm_logs/{job_name}_" + log_err = parsed_args.log_err if parsed_args.log_err is not None else f"{parsed_args.directory}slurm_logs/{job_name}_" + + create_slurm_template( + slurm_template=slurm_template, + job_name=job_name, + ntasks=parsed_args.ntasks, + nthreads=parsed_args.nthreads, + partition=parsed_args.partition, + time=parsed_args.time, + mem=parsed_args.mem, + log_out=log_out, + log_err=log_err, + array=parsed_args.array, + ) + diff --git a/timestepping.py b/timestepping.py new file mode 100644 index 0000000..76e2697 --- /dev/null +++ b/timestepping.py @@ -0,0 +1,126 @@ +from os.path import isfile +from pysbmy.timestepping import StandardTimeStepping +import numpy as np +from argparse import ArgumentParser + +def register_arguments_timestepping(parser:ArgumentParser): + """ + Register the arguments for the timestepping. + """ + parser.add_argument("-nt","--nsteps", type=int, default=10, help="Number of timesteps.") + parser.add_argument("--integrator", type=str, default="COLAm", help="Integrator to use.") + parser.add_argument("--TimeStepDistribution", type=str, default="a", help="Time step distribution.") + parser.add_argument("--Snapshots", type=int, nargs="*", default=None, help="Snapshots of steps to save.") + parser.add_argument("--n_LPT", type=float, default=-2.5, help="Modified discretisation parameters for COLAm.") + + +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 cosmo_params import parse_arguments_cosmo, z2a + + main_dict = parse_arguments_main(parsed_args) + card_dict, _ = parse_arguments_card(parsed_args) + cosmo_dict = parse_arguments_cosmo(parsed_args) + + timestepping_dict = dict( + ai=z2a(card_dict["RedshiftLPT"]), + af=z2a(card_dict["RedshiftFCs"]), + nsteps=parsed_args.nsteps, + n_LPT=parsed_args.n_LPT, + cosmo=cosmo_dict, + lightcone=card_dict["GenerateLightcone"], + ) + + ts_filename = card_dict["TimeSteppingFileName"] + + match parsed_args.integrator: + case "PM" | "StandardLeapfrog": + timestepping_dict["integrator"] = 0 + case "COLA": + timestepping_dict["integrator"] = 1 + case "COLAm" | "COLA_mod": + timestepping_dict["integrator"] = 2 + case "BF" | "BullFrog": + timestepping_dict["integrator"] = 3 + case "LPT": + timestepping_dict["integrator"] = 4 + case _: + raise ValueError(f"Integrator {parsed_args.integrator} not recognised.") + + match parsed_args.TimeStepDistribution: + case "a" | "lin_a" | "linear": + timestepping_dict["TimeStepDistribution"] = 0 + case "log" | "log_a" | "logarithmic": + timestepping_dict["TimeStepDistribution"] = 1 + case "exp" | "exp_a" | "exponential": + timestepping_dict["TimeStepDistribution"] = 2 + case "D" | "lin_D" | "growth": + timestepping_dict["TimeStepDistribution"] = 3 + case _: + raise ValueError(f"Time step distribution {parsed_args.TimeStepDistribution} not recognised.") + + snapshots = np.zeros(parsed_args.nsteps) + if parsed_args.Snapshots is not None: + for snap in parsed_args.Snapshots: + if snap < 0 or snap >= parsed_args.nsteps: + raise ValueError(f"Snapshot {snap} is out of range.") + snapshots[snap] = 1 + timestepping_dict["snapshots"] = snapshots + + return timestepping_dict, ts_filename + + +def create_timestepping(timestepping_dict, ts_filename:str, verbose:int=1): + """ + Main function for the timestepping. + """ + TS = StandardTimeStepping(**timestepping_dict) + if 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): + TS.write(ts_filename) + g.close() + f.close() + else: + TS.write(ts_filename) + + +def main_timestepping(parsed_args): + """ + Main function for the timestepping. + """ + from low_level import print_message, print_ending_module, print_starting_module + + print_starting_module("timestepping", verbose=parsed_args.verbose) + print_message("Parsing arguments for the timestepping file.", 1, "timestepping", verbose=parsed_args.verbose) + timestepping_dict, ts_filename = parse_arguments_timestepping(parsed_args) + if isfile(ts_filename) and not parsed_args.force: + print_message(f"Timestepping file {ts_filename} already exists. Use -F to overwrite.", 1, "timestepping", verbose=parsed_args.verbose) + return timestepping_dict + create_timestepping(timestepping_dict, ts_filename, verbose=parsed_args.verbose) + print_message(f"Timestepping file written to {ts_filename}", 2, "timestepping", verbose=parsed_args.verbose) + print_ending_module("timestepping", verbose=parsed_args.verbose) + + return timestepping_dict + +if __name__ == "__main__": + from args_main import register_arguments_main + from parameters_card import register_arguments_card + 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_cosmo(parser) + parsed_args = parser.parse_args() + main_timestepping(parsed_args) \ No newline at end of file