Upload files to "/"
This commit is contained in:
parent
c188e4d8d2
commit
5c109ac66f
2 changed files with 305 additions and 0 deletions
179
slurm_submission.py
Normal file
179
slurm_submission.py
Normal file
|
|
@ -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,
|
||||
)
|
||||
|
||||
126
timestepping.py
Normal file
126
timestepping.py
Normal file
|
|
@ -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)
|
||||
Loading…
Add table
Reference in a new issue