From dcc510bf494a29b133b768957fe2e4f19d97e620 Mon Sep 17 00:00:00 2001 From: Patrick Armstrong Date: Wed, 4 Oct 2023 21:28:54 -0500 Subject: [PATCH] Add SIMFILE_BIASCOR_FITOPS and SIMFILE_CCPRIOR_FITOPTS keys to BIASCOR stage, which allows specifying which FITOPT to use. --- README.md | 7 +++++++ pippin/biascor.py | 43 +++++++++++++++++++++++++++++++++++-------- 2 files changed, 42 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 39ffe9b7..b3231b5c 100644 --- a/README.md +++ b/README.md @@ -720,9 +720,16 @@ BIASCOR: # Input Ia bias correction simulations to be concatenated SIMFILE_BIASCOR: [DESFIT_DESBIASCOR, LOWZFIT_LOWZBIASCOR] + # Optional, specify FITOPT to use. Defaults to 0 for each SIMFILE_BIASCOR. If using this option, you must specify a FITOPT for each SIMFILE_BIASCOR + SIMFILE_BIASCOR_FITOPTS: [0, 1] # FITOPT000 and FITOPT001 + # For surveys that have contamination, add in the cc only simulation under CCPRIOR SIMFILE_CCPRIOR: DESFIT_DESSIMBIAS5YRCC + # Optional, specify FITOPT to use. Defaults to 0 for each SIMFILE_CCPRIOR. If using this option, you must specify a FITOPT for each SIMFILE_CCPRIOR + SIMFILE_CCPRIOR_FITOPTS: [0, 1] # FITOPT000 and FITOPT001 + + # Which classifier to use. Column name in FITRES will be determined from this property. # In the case of multiple classifiers this can either be # 1. A list of classifiers which map to the same probability column name (as defined by MERGE_CLASSIFIERS in the AGGREGATION stage) diff --git a/pippin/biascor.py b/pippin/biascor.py index 2278821c..4939bb2b 100644 --- a/pippin/biascor.py +++ b/pippin/biascor.py @@ -33,7 +33,26 @@ def __init__(self, name, output_dir, config, dependencies, options, global_confi self.merged_data = config.get("DATA") self.merged_iasim = config.get("SIMFILE_BIASCOR") + self.merged_iasim_fitopts = config.get("SIMFILE_BIASCOR_FITOPTS") + if self.merged_iasim is not None: + if self.merged_iasim_fitopts is None: + self.merged_iasim_fitopts = [0 for _ in self.merged_iasim] + else: + self.merged_iasim_fitopts = ensure_list(self.merged_iasim_fitopts) + assert len(self.merged_iasim_fitopts) == len(self.merged_iasim), \ + f"SIMFILE_BIASCOR_FITOPTS must be the same length as SIMFILE_BIASCOR ({len(self.merged_iasim)}), but is length {len(self.merged_iasim_fitopts)}" + self.logger.debug(f"SIMFILE_BIASCOR_FITOPTS: {self.merged_iasim_fitopts}") self.merged_ccsim = config.get("SIMFILE_CCPRIOR") + self.merged_ccsim_fitopts = config.get("SIMFILE_CCPRIOR_FITOPTS") + if self.merged_ccsim is not None: + if self.merged_ccsim_fitopts is None: + self.merged_ccsim_fitopts = [0 for _ in self.merged_ccsim] + else: + self.merged_ccsim_fitopts = ensure_list(self.merged_ccsim_fitopts) + assert len(self.merged_ccsim_fitopts) == len(self.merged_ccsim), \ + f"SIMFILE_CCPRIOR_FITOPTS must be the same length as SIMFILE_CCPRIOR ({len(self.merged_ccsim)}), but is length {len(self.merged_ccsim_fitopts)}" + self.logger.debug(f"SIMFILE_CCPRIOR_FITOPTS: {self.merged_ccsim_fitopts}") + self.classifier = config.get("CLASSIFIER") if self.classifier is not None: self.config["CLASSIFIER"] = self.classifier.name @@ -232,11 +251,11 @@ def _check_completion(self, squeue): return self.check_for_job(squeue, self.job_name) - def get_simfile_biascor(self, ia_sims): - return None if ia_sims is None else ",".join([os.path.join(m.output["fitres_dirs"][0], "FITOPT000.FITRES.gz") for m in ia_sims]) + def get_simfile_biascor(self, ia_sims, fitopt_num): + return None if ia_sims is None else ",".join([os.path.join(ia_sims[i].output["fitres_dirs"][0], f"FITOPT{n:03}.FITRES.gz") for (i, n) in enumerate(fitopt_num)]) - def get_simfile_ccprior(self, cc_sims): - return None if cc_sims is None else ",".join([os.path.join(m.output["fitres_dirs"][0], "FITOPT000.FITRES.gz") for m in cc_sims]) + def get_simfile_ccprior(self, cc_sims, fitopt_num): + return None if cc_sims is None else ",".join([os.path.join(cc_sims[i].output["fitres_dirs"][0], "FITOPT{n:03}.FITRES.gz") for (i, n) in enumerate(fitopt_num)]) def get_fitopt_map(self, datas): fitopts = {} @@ -270,15 +289,23 @@ def get_fitopt_map(self, datas): def write_input(self): if self.merged_iasim is not None: - for m in self.merged_iasim: + for (i, m) in enumerate(self.merged_iasim): if len(m.output["fitres_dirs"]) > 1: self.logger.warning(f"Your IA sim {m} has multiple versions! Using 0 index from options {m.output['fitres_dirs']}") + self.logger.debug(f"Fitres directory: {m.output['fitres_dirs'][0]}") + n = self.merged_iasim_fitopts[i] + assert os.path.exists(os.path.join(m.output["fitres_dirs"][0], f"FITOPT{n:03}.FITRES.gz")), f"FITOPT{n:03} does not exist for your IA sim {m}" if self.merged_ccsim is not None: - for m in self.merged_ccsim: + for (i, m) in enumerate(self.merged_ccsim): if len(m.output["fitres_dirs"]) > 1: self.logger.warning(f"Your CC sim {m} has multiple versions! Using 0 index from options {m.output['fitres_dirs']}") - self.bias_cor_fits = self.get_simfile_biascor(self.merged_iasim) - self.cc_prior_fits = self.get_simfile_ccprior(self.merged_ccsim) + self.logger.debug(f"Fitres directory: {m.output['fitres_dirs'][0]}") + n = self.merged_ccsim_fitopts[i] + assert os.path.exists(os.path.join(m.output["fitres_dirs"][0], f"FITOPT{n:03}.FITRES.gz")), f"FITOPT{n:03} does not exist for your CC sim {m}" + + + self.bias_cor_fits = self.get_simfile_biascor(self.merged_iasim, self.merged_iasim_fitopts) + self.cc_prior_fits = self.get_simfile_ccprior(self.merged_ccsim, self.merged_ccsim_fitopts) self.data = [m.output["lc_output_dir"] for m in self.merged_data] self.data_fitres = [m.output["fitres_file"] for m in self.merged_data] #print('MERGED DATA')