def classify(self, command):
self.setup()
if self.batch_file is None:
if self.gpu:
self.sbatch_header = self.sbatch_gpu_header
else:
self.sbatch_header = self.sbatch_cpu_header
else:
with open(self.batch_file, 'r') as f:
self.sbatch_header = f.read()
self.sbatch_header = self.clean_header(self.sbatch_header)
header_dict = {
"REPLACE_NAME": self.job_base_name,
"REPLACE_LOGFILE": "output.log",
"REPLACE_WALLTIME": "00:55:00",
"REPLACE_MEM": "8GB",
"APPEND": ["#SBATCH --ntasks=1", "#SBATCH --cpus-per-task=4"]
}
header_dict = merge_dict(header_dict, self.batch_replace)
self.update_header(header_dict)
setup_dict = {
"job_name": self.job_base_name,
"conda_env": self.conda_env,
"path_to_classifier": self.path_to_classifier,
"command_opts": command
}
format_dict = {
"done_file":
self.done_file,
"sbatch_header":
self.sbatch_header,
"task_setup":
self.update_setup(setup_dict, self.task_setup['nearest_neighbour'])
}
slurm_script = self.slurm.format(**format_dict)
new_hash = self.get_hash_from_string(slurm_script)
if self._check_regenerate(new_hash):
self.logger.debug("Regenerating")
shutil.rmtree(self.output_dir, ignore_errors=True)
mkdirs(self.output_dir)
slurm_output_file = self.output_dir + "/job.slurm"
with open(slurm_output_file, "w") as f:
f.write(slurm_script)
self.save_new_hash(new_hash)
self.logger.info(f"Submitting batch job {slurm_output_file}")
subprocess.run(["sbatch", slurm_output_file], cwd=self.output_dir)
else:
self.logger.info("Hash check passed, not rerunning")
self.should_be_done()
return True
def _run(self):
# Get the m0diff files for everything
for b in self.biascor_deps:
for m in b.output["m0dif_dirs"]:
self.logger.info(f"Looking at M0diff dir {m}")
sim_number = 1
if os.path.basename(m).isdigit():
sim_number = int(os.path.basename(m))
files = [
f for f in sorted(os.listdir(m))
if f.endswith(".M0DIF") or f.endswith(".M0DIF.gz")
]
for f in files:
muopt_num = int(f.split("MUOPT")[-1].split(".")[0])
fitopt_num = int(f.split("FITOPT")[-1].split("_")[0])
if muopt_num == 0:
muopt = "DEFAULT"
else:
muopt = b.output["muopts"][muopt_num -
1] # Because 0 is default
if fitopt_num == 0:
fitopt = "DEFAULT"
else:
fitopt = b.output["fitopt_index"][fitopt_num]
self.biascor_m0diffs.append(
(b.name, sim_number, muopt, muopt_num, fitopt,
fitopt_num, os.path.join(m, f)))
data_fitres_files = [
os.path.join(l.output["fitres_dirs"][0],
l.output["fitopt_map"]["DEFAULT"])
for l in self.lcfit_deps if l.output["is_data"]
]
data_fitres_output = [
d.split("/")[-4] + ".csv.gz" for d in data_fitres_files
]
sim_fitres_files = [
os.path.join(l.output["fitres_dirs"][0],
l.output["fitopt_map"]["DEFAULT"])
for l in self.lcfit_deps if not l.output["is_data"]
]
sim_fitres_output = [
d.split("/")[-4] + ".csv.gz" for d in sim_fitres_files
]
types = list(
set([
a for l in self.lcfit_deps
for a in l.sim_task.output["types_dict"]["IA"]
]))
input_yml_file = "input.yml"
output_dict = {
"COSMOMC": {
"INPUT_FILES": self.cosmomc_input_files,
"PARSED_FILES": self.cosmomc_output_files,
"PARSED_COVOPTS": self.cosmomc_covopts,
"PARAMS": self.params,
"SHIFT": self.options.get("SHIFT", False),
"PRIOR": self.options.get("PRIOR"),
"NAMES": self.names,
"CONTOUR_COVOPTS": self.covopts,
"SINGULAR_BLIND": self.singular_blind,
},
"BIASCOR": {
"WFIT_SUMMARY_INPUT": self.wsummary_files,
"WFIT_SUMMARY_OUTPUT": "all_biascor.csv",
"FITRES_INPUT": self.biascor_fitres_input_files,
"FITRES_PROB_COLS": self.biascor_prob_col_names,
"FITRES_PARSED": self.biascor_fitres_output_files,
"M0DIFF_INPUTS": self.biascor_m0diffs,
"M0DIFF_PARSED": self.biascor_m0diff_output,
"FITRES_COMBINED": self.biascor_fitres_combined,
},
"OUTPUT_NAME": self.name,
"BLIND": self.blind_params,
"LCFIT": {
"DATA_FITRES_INPUT": data_fitres_files,
"SIM_FITRES_INPUT": sim_fitres_files,
"DATA_FITRES_PARSED": data_fitres_output,
"SIM_FITRES_PARSED": sim_fitres_output,
"IA_TYPES": types,
},
}
if self.batch_file is None:
if self.gpu:
self.sbatch_header = self.sbatch_gpu_header
else:
self.sbatch_header = self.sbatch_cpu_header
else:
with open(self.batch_file, 'r') as f:
self.sbatch_header = f.read()
self.sbatch_header = self.clean_header(self.sbatch_header)
header_dict = {
"REPLACE_NAME": self.job_name,
"REPLACE_WALLTIME": "1:00:00",
"REPLACE_LOGFILE": self.logfile,
"REPLACE_MEM": "20GB",
"APPEND": ["#SBATCH --ntasks=1", "#SBATCH --cpus-per-task=1"]
}
header_dict = merge_dict(header_dict, self.batch_replace)
self.update_header(header_dict)
setup_dict = {"output_dir": self.output_dir}
format_dict = {
"sbatch_header": self.sbatch_header,
"task_setup": self.update_setup(setup_dict,
self.task_setup['analyse']),
"input_yml": input_yml_file
}
final_slurm = self.get_slurm_raw().format(**format_dict)
new_hash = self.get_hash_from_string(final_slurm +
json.dumps(output_dict))
if self._check_regenerate(new_hash):
self.logger.debug("Regenerating and launching task")
shutil.rmtree(self.output_dir, ignore_errors=True)
mkdirs(self.output_dir)
self.save_new_hash(new_hash)
for c in self.path_to_codes:
shutil.copy(c, self.output_dir)
input_yml_path = os.path.join(self.output_dir, input_yml_file)
with open(input_yml_path, "w") as f:
json.dump(output_dict, f, indent=2)
self.logger.debug(
f"Input yml file written out to {input_yml_path}")
slurm_output_file = os.path.join(self.output_dir, "slurm.job")
with open(slurm_output_file, "w") as f:
f.write(final_slurm)
self.logger.info(f"Submitting batch job for analyse chains")
subprocess.run(["sbatch", slurm_output_file], cwd=self.output_dir)
else:
self.logger.info("Hash check passed, not rerunning")
return True
def classify(self, training):
model = self.options.get("MODEL")
model_path = ""
if not training:
assert model is not None, "If TRAIN is not specified, you have to point to a model to use"
if not os.path.exists(get_output_loc(model)):
for t in self.dependencies:
if model == t.name:
self.logger.debug(
f"Found task dependency {t.name} with model file {t.output['model_filename']}"
)
model = t.output["model_filename"]
model_path = get_output_loc(model)
self.logger.debug(f"Looking for model in {model_path}")
assert os.path.exists(model_path), f"Cannot find {model_path}"
types = self.get_types()
if types is None:
types = OrderedDict({
"1": "Ia",
"0": "unknown",
"2": "SNIax",
"3": "SNIa-pec",
"20": "SNIIP",
"21": "SNIIL",
"22": "SNIIn",
"29": "SNII",
"32": "SNIb",
"33": "SNIc",
"39": "SNIbc",
"41": "SLSN-I",
"42": "SLSN-II",
"43": "SLSN-R",
"80": "AGN",
"81": "galaxy",
"98": "None",
"99": "pending",
"101": "Ia",
"120": "SNII",
"130": "SNIbc",
})
else:
has_ia = False
has_cc = False
self.logger.debug(f"Input types set to {types}")
for key, value in types.items():
if value.upper() == "IA":
has_ia = True
elif value.upper() in ["II", "IBC"]:
has_cc = True
if not has_ia:
self.logger.debug("No Ia type found, injecting type")
types[1] = "Ia"
types = dict(
sorted(types.items(),
key=lambda x: -1 if x[0] == 1 else x[0]))
self.logger.debug(f"Inject types with Ias are {types}")
if not has_cc:
self.logger.debug("No cc type found, injecting type")
types[29] = "II"
str_types = json.dumps(types)
self.logger.debug(f"Types set to {str_types}")
sim_dep = self.get_simulation_dependency()
light_curve_dir = sim_dep.output["photometry_dirs"][self.index]
self.raw_dir = light_curve_dir
fit = self.get_fit_dependency()
fit_dir = f"" if fit is None else f"--fits_dir {fit['fitres_dirs'][self.index]}"
cyclic = "--cyclic" if self.variant in ["vanilla", "variational"
] and self.cyclic else ""
batch_size = f"--batch_size {self.batch_size}"
num_layers = f"--num_layers {self.num_layers}"
hidden_dim = f"--hidden_dim {self.hidden_dim}"
variant = f"--model {self.variant}"
if self.variant == "bayesian":
variant += " --num_inference_samples 20"
clump = sim_dep.output.get("clump_file")
if clump is None:
clump_txt = ""
else:
clump_txt = f"--photo_window_files {clump}"
if self.batch_file is None:
if self.gpu:
self.sbatch_header = self.sbatch_gpu_header
else:
self.sbatch_header = self.sbatch_cpu_header
else:
with open(self.batch_file, 'r') as f:
self.sbatch_header = f.read()
self.sbatch_header = self.clean_header(self.sbatch_header)
if self.has_yml:
self.update_yml()
setup_file = "supernnova_yml"
else:
setup_file = "supernnova"
header_dict = {
"REPLACE_NAME": self.job_base_name,
"REPLACE_WALLTIME": "23:00:00",
"REPLACE_LOGFILE": "output.log",
"REPLACE_MEM": "32GB",
"APPEND": ["#SBATCH --ntasks=1", "#SBATCH --cpus-per-task=1"]
}
header_dict = merge_dict(header_dict, self.batch_replace)
self.update_header(header_dict)
setup_dict = {
"conda_env":
self.conda_env,
"dump_dir":
self.dump_dir,
"photometry_dir":
light_curve_dir,
"fit_dir":
fit_dir,
"path_to_classifier":
self.path_to_classifier,
"job_name":
self.job_base_name,
"command":
"--train_rnn" if training else "--validate_rnn",
"sntypes":
str_types,
"variant":
variant,
"cyclic":
cyclic,
"model":
"" if training else f"--model_files {model_path}",
"phot":
"",
"test_or_train":
"" if training else "--data_testing",
"redshift":
"--redshift " + self.redshift,
"norm":
"--norm " + self.norm,
"done_file":
self.done_file,
"clump":
clump_txt,
"done_file2":
self.done_file2,
"partition":
"gpu2" if self.gpu else "broadwl",
"gres":
"#SBATCH --gres=gpu:1" if self.gpu else "",
"cuda":
"--use_cuda" if self.gpu else "",
"clean_command":
f"rm -rf {self.dump_dir}/processed" if self.clean else "",
"seed":
f"--seed {self.seed}" if self.seed else "",
"batch_size":
batch_size,
"num_layers":
num_layers,
"hidden_dim":
hidden_dim,
"data_yml":
self.output_data_yml,
"classification_yml":
self.output_classification_yml,
"classification_command":
"train_rnn" if training else "validate_rnn"
}
format_dict = {
"sbatch_header":
self.sbatch_header,
"task_setup":
self.update_setup(setup_dict, self.task_setup[setup_file])
}
slurm_output_file = self.output_dir + "/job.slurm"
self.logger.info(
f"Running SuperNNova, slurm job outputting to {slurm_output_file}")
slurm_text = self.slurm.format(**format_dict)
new_hash = self.get_hash_from_string(slurm_text)
if not self._check_regenerate(new_hash):
self.should_be_done()
else:
self.logger.info("Rerunning. Cleaning output_dir")
shutil.rmtree(self.output_dir, ignore_errors=True)
mkdirs(self.output_dir)
if self.has_yml:
with open(self.output_data_yml, 'w') as f:
f.write(self.data_yml)
with open(self.output_classification_yml, 'w') as f:
f.write(self.classification_yml)
self.save_new_hash(new_hash)
with open(slurm_output_file, "w") as f:
f.write(slurm_text)
self.logger.info(
f"Submitting batch job to {'train' if training else 'predict using'} SuperNNova"
)
subprocess.run(["sbatch", slurm_output_file], cwd=self.output_dir)
return True
def _run(self):
val_p = self.options.get("PHOTFLAG_DETECT")
val_c = self.options.get("CUTWIN_SNR_NODETECT")
photflag = f"PHOTFLAG_DETECT = {val_p}" if val_p else ""
cutwin = f"CUTWIN_SNR_NODETECT = {val_c}" if val_c else ""
command_string = self.clump_command.format(
genversion=self.genversion,
data_path=self.data_path,
opt_setpkmjd=self.opt_setpkmjd,
photflag=photflag,
cutwin_snr_nodetect=cutwin,
photflag_mskrej=self.photflag_mskrej)
if self.batch_file is None:
if self.gpu:
self.sbatch_header = self.sbatch_gpu_header
else:
self.sbatch_header = self.sbatch_cpu_header
else:
with open(self.batch_file, 'r') as f:
self.sbatch_header = f.read()
self.sbatch_header = self.clean_header(self.sbatch_header)
header_dict = {
"REPLACE_NAME": self.job_name,
"REPLACE_WALLTIME": "0:20:00",
"REPLACE_LOGFILE": self.logfile,
"REPLACE_MEM": "2GB",
"APPEND": ["#SBATCH --ntasks=1", "#SBATCH --cpus-per-task=1"]
}
header_dict = merge_dict(header_dict, self.batch_replace)
self.update_header(header_dict)
setup_dict = {
"path_to_task": self.path_to_task,
"done_file": self.done_file
}
format_dict = {
"sbatch_header":
self.sbatch_header,
"task_setup":
self.update_setup(setup_dict, self.task_setup['dataprep'])
}
#format_dict = {"job_name": self.job_name, "log_file": self.logfile, "path_to_task": self.path_to_task, "done_file": self.done_file}
final_slurm = self.slurm.format(**format_dict)
new_hash = self.get_hash_from_string(command_string + final_slurm)
if self._check_regenerate(new_hash):
self.logger.debug("Regenerating and launching task")
shutil.rmtree(self.output_dir, ignore_errors=True)
mkdirs(self.output_dir)
self.save_new_hash(new_hash)
slurm_output_file = os.path.join(self.output_dir, "slurm.job")
clump_file = os.path.join(self.output_dir, "clump.nml")
with open(slurm_output_file, "w") as f:
f.write(final_slurm)
with open(clump_file, "w") as f:
f.write(command_string)
self.logger.info(f"Submitting batch job for data prep")
subprocess.run(["sbatch", slurm_output_file], cwd=self.output_dir)
else:
self.should_be_done()
self.logger.info("Hash check passed, not rerunning")
return True
def _run(self):
if self.static:
self.logger.info(
"CMB only constraints detected, copying static files")
cosmomc_static_loc = get_data_loc(self.static_path +
self.ini_prefix)
if cosmomc_static_loc is None:
self.logger.error(
"Seems like we can't find the static chains...")
return False
else:
new_hash = self.get_hash_from_string(cosmomc_static_loc)
if self._check_regenerate(new_hash):
self.logger.debug("Regenerating and copying static chains")
shutil.rmtree(self.chain_dir, ignore_errors=True)
shutil.copytree(cosmomc_static_loc, self.chain_dir)
for done_file in self.done_files:
df = os.path.join(self.output_dir, done_file)
with open(df, "w") as f:
f.write("SUCCESS")
self.save_new_hash(new_hash)
else:
self.should_be_done()
self.logger.info("Hash check passed, not rerunning")
else:
ini_filecontents = self.get_ini_file()
if ini_filecontents is None:
return False
if self.batch_file is None:
if self.gpu:
self.sbatch_header = self.sbatch_gpu_header
else:
self.sbatch_header = self.sbatch_cpu_header
else:
with open(self.batch_file, 'r') as f:
self.sbatch_header = f.read()
self.sbatch_header = self.clean_header(self.sbatch_header)
header_dict = {
"REPLACE_NAME":
self.job_name,
"REPLACE_WALLTIME":
"34:00:00",
"REPLACE_LOGFILE":
self.logfile,
"REPLACE_MEM":
"2GB",
"APPEND": [
f"#SBATCH --ntasks={self.ntasks}",
f"#SBATCH --array=1-{len(self.ini_files)}",
"#SBATCH --cpus-per-task=1"
]
}
header_dict = merge_dict(header_dict, self.batch_replace)
self.update_header(header_dict)
setup_dict = {
"done_files": " ".join(self.done_files),
"path_to_cosmomc": self.path_to_cosmomc,
"output_dir": self.output_dir,
"ini_files": " ".join(self.ini_files),
"num_jobs": len(self.ini_files),
"num_walkers": self.num_walkers,
}
format_dict = {
"sbatch_header": self.sbatch_header,
"task_setup": self.update_setup(setup_dict,
self.task_setup['cosmomc'])
}
final_slurm = self.slurm.format(**format_dict)
new_hash = self.get_hash_from_string(final_slurm +
" ".join(ini_filecontents))
if self._check_regenerate(new_hash):
self.logger.debug("Regenerating and launching task")
shutil.rmtree(self.output_dir, ignore_errors=True)
mkdirs(self.output_dir)
self.save_new_hash(new_hash)
slurm_output_file = os.path.join(self.output_dir, "slurm.job")
with open(slurm_output_file, "w") as f:
f.write(final_slurm)
for file, content in zip(self.ini_files, ini_filecontents):
filepath = os.path.join(self.output_dir, file)
with open(filepath, "w") as f:
f.write(content)
mkdirs(self.chain_dir)
needed_dirs = [
"data", "paramnames", "camb", "batch1", "batch2", "batch3"
]
for d in needed_dirs:
self.logger.debug(f"Creating symlink to {d} dir")
original_data_dir = os.path.join(self.path_to_cosmomc, d)
new_data_dir = os.path.join(self.output_dir, d)
os.symlink(original_data_dir,
new_data_dir,
target_is_directory=True)
self.logger.info(f"Submitting batch job for data prep")
subprocess.run(["sbatch", slurm_output_file], cwd=self.output_dir)
else:
self.should_be_done()
self.logger.info("Hash check passed, not rerunning")
return True
def _run(self):
sys_scale = self.calculate_input()
if self.batch_file is None:
if self.gpu:
self.sbatch_header = self.sbatch_gpu_header
else:
self.sbatch_header = self.sbatch_cpu_header
else:
with open(self.batch_file, 'r') as f:
self.sbatch_header = f.read()
self.sbatch_header = self.clean_header(self.sbatch_header)
header_dict = {
"REPLACE_NAME": self.job_name,
"REPLACE_WALLTIME": "02:30:00",
"REPLACE_LOGFILE": self.logfile,
"REPLACE_MEM": str(self.batch_mem),
"APPEND": ["#SBATCH --ntasks-per-node=1"]
}
header_dict = merge_dict(header_dict, self.batch_replace)
self.update_header(header_dict)
self.logger.debug(
f"Binned: {self.binned}, Rebinned x1: {self.rebinned_x1}, Rebinned c: {self.rebinned_c}"
)
setup_dict = {
"path_to_code": self.path_to_code,
"input_file": self.input_file,
"output_dir": self.output_dir,
"unbinned": "" if self.binned else "-u",
"nbin_x1": self.rebinned_x1,
"nbin_c": self.rebinned_c,
"subtract_vpec": "" if not self.subtract_vpec else "-s",
}
format_dict = {
"sbatch_header":
self.sbatch_header,
"task_setup":
self.update_setup(setup_dict, self.task_setup['create_cov']),
"done_file":
self.done_file,
}
final_slurm = self.slurm.format(**format_dict)
final_output_for_hash = self.get_output_string() + yaml.safe_dump(
sys_scale, width=2048) + final_slurm
new_hash = self.get_hash_from_string(final_output_for_hash)
if self._check_regenerate(new_hash):
self.logger.debug("Regenerating and launching task")
shutil.rmtree(self.output_dir, ignore_errors=True)
mkdirs(self.output_dir)
mkdirs(self.config_dir)
self.save_new_hash(new_hash)
# Write sys scales and the main input file
with open(self.sys_file_out, "w") as f:
f.write(yaml.safe_dump(sys_scale, width=2048))
with open(self.input_file, "w") as f:
f.write(self.get_output_string())
# Write out slurm job script
slurm_output_file = os.path.join(self.output_dir, "slurm.job")
with open(slurm_output_file, "w") as f:
f.write(final_slurm)
self.logger.info(f"Submitting batch job for create covariance")
subprocess.run(["sbatch", slurm_output_file], cwd=self.output_dir)
else:
self.should_be_done()
self.logger.info("Hash check passed, not rerunning")
return True
def classify(self, mode):
header_dict = {
"REPLACE_NAME": self.job_base_name,
"REPLACE_LOGFILE": "output.log",
"REPLACE_WALLTIME": "15:00:00", # TODO: scale based on number of heatmaps
"REPLACE_MEM": "8GB",
# "APPEND": ["#SBATCH --ntasks=1", "#SBATCH --cpus-per-task=8"]
}
header_dict = merge_dict(header_dict, self.batch_replace)
if self.batch_file is None:
if self.gpu:
self.sbatch_header = self.sbatch_gpu_header
else:
self.sbatch_header = self.sbatch_cpu_header
else:
with open(self.batch_file, 'r') as f:
self.sbatch_header = f.read()
self.sbatch_header = self.clean_header(self.sbatch_header)
self.update_header(header_dict)
setup_dict = {
"conda_env": self.conda_env,
"path_to_classifier": self.path_to_classifier,
"heatmaps_path": self.heatmaps_path,
"config_path": self.config_path,
"done_file": self.done_file,
}
format_dict = {
"sbatch_header": self.sbatch_header,
"task_setup": self.update_setup(setup_dict, self.task_setup['scone'])
}
slurm_output_file = self.output_dir + "/job.slurm"
self.logger.info(f"Running SCONE, slurm job outputting to {slurm_output_file}")
slurm_script = self.slurm.format(**format_dict)
new_hash = self.get_hash_from_string(slurm_script)
# check success of intermediate steps and don't redo them if successful
heatmaps_created = self._heatmap_creation_success() and self.keep_heatmaps
failed = False
if os.path.exists(self.done_file):
self.logger.debug(f"Found done file at {self.done_file}")
with open(self.done_file) as f:
if "FAILURE" in f.read().upper():
failed = True
if self._check_regenerate(new_hash) or failed:
self.logger.debug("Regenerating")
if not heatmaps_created:
shutil.rmtree(self.output_dir, ignore_errors=True)
mkdirs(self.output_dir)
else:
for f in [f.path for f in os.scandir(self.output_dir) if f.is_file()]:
os.remove(f)
sim_dep = self.get_simulation_dependency()
sim_dirs = sim_dep.output["photometry_dirs"]
lcdata_paths = self._get_lcdata_paths(sim_dirs)
metadata_paths = [path.replace("PHOT", "HEAD") for path in lcdata_paths]
self._write_config_file(metadata_paths, lcdata_paths, mode, self.config_path) # TODO: what if they don't want to train on all sims?
with open(slurm_output_file, "w") as f:
f.write(slurm_script)
self.save_new_hash(new_hash)
self.logger.info(f"Submitting batch job {slurm_output_file}")
subprocess.run(["sbatch", slurm_output_file], cwd=self.output_dir)
else:
self.logger.info("Hash check passed, not rerunning")
self.should_be_done()
return True