diff --git a/src/smefit/analyze/chi2_utils.py b/src/smefit/analyze/chi2_utils.py index 5699e71d..535df507 100644 --- a/src/smefit/analyze/chi2_utils.py +++ b/src/smefit/analyze/chi2_utils.py @@ -107,6 +107,43 @@ def compute(datasets, smeft_predictions): total_chi2_rep / datasets.Commondata.size, ) + @staticmethod + def compute_ext_chi2(external_chi2, best_fit): + r"""Compute the external likelihood for each dataset. + + Parameters + ---------- + external_chi2: dict + dictionary whose keys are the different datasets and the values + the compute_chi2 function of the external likelihood module. Obtained + by fit.external_chi2 + best_fit: pd.core.series.Series + best fit value for all the Wilson coefficients. Obtained by + the fit.best_fit property + + Returns + ------- + pd.DataFrame: + External chi2 for each dataset + """ + ext_chi2_values = [ext_chi2_func(best_fit) for ext_chi2_func in external_chi2] + + # The SM ext. likelihood can be obtained through the ext. likelihood module + # by setting all the WCs to zero + sm_coeff = pd.Series(0, index=best_fit.index, dtype=best_fit.dtype) + ext_chi2_sm_values = [ + ext_chi2_func(sm_coeff) for ext_chi2_func in external_chi2 + ] + + indices = [ext_func.__self__.__class__.__name__ for ext_func in external_chi2] + return pd.DataFrame( + { + "sm_chi2": np.array(ext_chi2_sm_values), + "ext_chi2": np.array(ext_chi2_values), + }, + index=indices, + ) + @staticmethod def add_normalized_chi2(chi2_df): r"""Add the normalized :math:`\chi^2` to the table. @@ -185,7 +222,7 @@ def _split_table_entries_sm(self, chi2_dict, chi2_dict_group): ] self.chi2_df_sm_grouped = self.chi2_df_sm_grouped.drop_duplicates() - def write(self, chi2_dict, chi2_dict_group): + def write(self, chi2_dict, chi2_dict_group, chi2_ext_dict): r"""Write the :math:`\chi^2` latex tables. Parameters @@ -194,6 +231,8 @@ def write(self, chi2_dict, chi2_dict_group): tables computed with compute() method for each fit chi2_dict_group: dict tables obtained with group_chi2_df() method for each fit + chi2_ext_dict: dict + tables computed with compute_ext_chi2() method for each fit Returns ------- @@ -206,6 +245,76 @@ def write(self, chi2_dict, chi2_dict_group): L.extend(self.write_chi2_grouped(chi2_dict, chi2_dict_group)) L.extend(["\n", "\n"]) L.extend(self.write_chi2_summary(chi2_dict_group)) + + if chi2_ext_dict != {}: + L.extend(["\n", "\n"]) + L.extend(self.write_external_chi2(chi2_ext_dict)) + + return L + + def write_external_chi2(self, ext_chi2_dict): + r"""Write the external likelihood latex tables for each dataset. + + Parameters + ---------- + ext_chi2_dict : dict + tables computed with compute_ext_chi2() method per each fit + + Returns + ------- + list(str) + list with the latex commands + """ + L = [ + r"\begin{table}[H]", + r"\centering", + r"\begin{tabular}{|l|" + "c|c|" * len(ext_chi2_dict) + "}", + r"\hline", + ] + + temp = r"" + for label in ext_chi2_dict.keys(): + temp += f"& \\multicolumn{{2}}{{c|}}{{{label}}}" + temp += r"\\ \hline" + L.append(temp) + L.append( + r"Process " + r" & SM & Best fit" * len(ext_chi2_dict) + r"\\ \hline", + ) + + # Extract unique ext. likelihood datasets + datasets = set() + for df in ext_chi2_dict.values(): + datasets.update(df.index) + datasets = list(datasets) + + total_chi2 = {group: 0 for group in ext_chi2_dict.keys()} + total_chi2_sm = {group: 0 for group in ext_chi2_dict.keys()} + for dataset in datasets: + temp = f"{dataset}" + for group, ext_chi2_df in ext_chi2_dict.items(): + if dataset in ext_chi2_df.index: + temp += f" & {ext_chi2_df.loc[dataset, 'sm_chi2']:.3f}" + temp += f" & {ext_chi2_df.loc[dataset, 'ext_chi2']:.3f}" + total_chi2[group] += ext_chi2_df.loc[dataset, "ext_chi2"] + total_chi2_sm[group] += ext_chi2_df.loc[dataset, "sm_chi2"] + else: + temp += " & &" + temp += r" \\ \hline" + L.append(temp) + + temp = r" \hline Total" + for group in ext_chi2_dict.keys(): + temp += f" & {total_chi2_sm[group]:.3f} & {total_chi2[group]:.3f}" + temp += r" \\ \hline" + + L.extend( + [ + temp, + r"\end{tabular}", + r"\caption{External likelihood table for each dataset.}", + r"\end{table}", + ] + ) return L def write_chi2_grouped(self, chi2_dict, chi2_dict_group): @@ -289,15 +398,19 @@ def write_chi2_summary(self, chi2_dict_group): r"\begin{tabular}{|l|" + "c|c|" * len(chi2_dict_group) + "}", r"\hline", ] - temp = r"" - for label in chi2_dict_group: - temp += f"& \\multicolumn{{2}}{{c|}}{{{label}}}" + temp_parts = [ + f"& \\multicolumn{{2}}{{c|}}{{{label}}}" for label in chi2_dict_group + ] + temp = "".join(temp_parts) temp += r"\\ \hline" L.append(temp) + L.append( r"Process " - + r" & $N_{\rm data}$ & $\chi^2/N_{\rm data}$" * len(chi2_dict_group) - + r"\\ \hline", + + " ".join( + [r"& $N_{\rm data}$ & $\chi^2/N_{\rm data}$" for _ in chi2_dict_group] + ) + + r"\\ \hline" ) for group in self.data_info.index.levels[0]: diff --git a/src/smefit/analyze/report.py b/src/smefit/analyze/report.py index 469398e9..930725d5 100644 --- a/src/smefit/analyze/report.py +++ b/src/smefit/analyze/report.py @@ -68,6 +68,7 @@ def __init__(self, report_path, result_path, report_config): fit.load_datasets() self.fits.append(fit) + self.fits = np.array(self.fits) # Get names of datasets for each fit @@ -159,7 +160,13 @@ def summary(self): tables=summary.fit_settings(), ) - def chi2(self, table=True, plot_experiment=None, plot_distribution=None): + def chi2( + self, + table=True, + plot_experiment=None, + plot_distribution=None, + table_external_chi2=None, + ): r""":math:`\chi^2` table and plots runner. Parameters @@ -180,6 +187,7 @@ def chi2(self, table=True, plot_experiment=None, plot_distribution=None): chi2_dict = {} chi2_dict_group = {} chi2_replica = {} + ext_chi2_dict = {} for fit in self.fits: # This computes the chi2 by taking the mean of the replicas _, chi2_total_rep = chi2_cal.compute( @@ -195,8 +203,13 @@ def chi2(self, table=True, plot_experiment=None, plot_distribution=None): chi2_dict[fit.label] = chi2_cal.add_normalized_chi2(chi2_df_best) chi2_dict_group[fit.label] = chi2_cal.group_chi2_df(chi2_df_best) + if table_external_chi2: + ext_chi2_dict[fit.label] = chi2_cal.compute_ext_chi2( + fit.external_chi2, fit.best_fit + ) + if table: - lines = chi2_cal.write(chi2_dict, chi2_dict_group) + lines = chi2_cal.write(chi2_dict, chi2_dict_group, ext_chi2_dict) compile_tex(self.report, lines, "chi2_tables") links_list = [("chi2_tables", "Tables")] diff --git a/src/smefit/external_chi2.py b/src/smefit/external_chi2.py new file mode 100644 index 00000000..f106a36b --- /dev/null +++ b/src/smefit/external_chi2.py @@ -0,0 +1,62 @@ +# -*- coding: utf-8 -*- +import importlib +import pathlib +import sys + +from smefit import log + +_logger = log.logging.getLogger(__name__) + + +def load_external_chi2(external_chi2, coefficients, rge_dict): + """ + Loads the external chi2 modules. + We assume that the external chi2 only need to know which coefficients we want to + fit and the RGE dictionary, which specifies the reference scale and the theory settings + for the RGE evolution. + + Parameters + ---------- + external_chi2: dict + dict of external chi2s. + Each key is a dictionary with the name of the chi2 class, specifying the + path to the module and the parameters to be passed to the chi2 class. + coefficients: CoefficientManager + The coefficient manager object. + rge_dict: dict + The RGE dictionary, specifying the reference scale and theory settings. + + Returns + ------- + ext_chi2_modules: list + List of external chi2 objects that can be evaluated by passing a coefficients instance + """ + # dynamical import + ext_chi2_modules = [] + + for class_name, module in external_chi2.items(): + _logger.info("Loading external chi2 module: %s", class_name) + + module_path = module["path"] + path = pathlib.Path(module_path) + base_path, stem = path.parent, path.stem + sys.path = [str(base_path)] + sys.path + try: + chi2_module = importlib.import_module(stem) + except ModuleNotFoundError: + print( + f"Module {stem} not found in {base_path}. Adjust and rerun. Exiting the code." + ) + sys.exit(1) + + my_chi2_class = getattr(chi2_module, class_name) + + extra_keys = {key: value for key, value in module.items() if key != "path"} + + chi2_ext = my_chi2_class( + coefficients=coefficients, rge_dict=rge_dict, **extra_keys + ) + + ext_chi2_modules.append(chi2_ext.compute_chi2) + + return ext_chi2_modules diff --git a/src/smefit/fit_manager.py b/src/smefit/fit_manager.py index b89ced01..da14bd6f 100644 --- a/src/smefit/fit_manager.py +++ b/src/smefit/fit_manager.py @@ -10,6 +10,7 @@ from .coefficients import CoefficientManager from .compute_theory import make_predictions +from .external_chi2 import load_external_chi2 from .loader import load_datasets @@ -57,6 +58,7 @@ def __init__(self, path, name, label=None): self.results = None self.datasets = None self.rgemat = None + self.external_chi2 = None def __repr__(self): return self.name @@ -145,6 +147,20 @@ def load_datasets(self): rgemat=self.rgemat, ) + external_chi2_dict = self.config.get("external_chi2", None) + self.external_chi2 = ( + load_external_chi2( + external_chi2_dict, self.coefficients, self.config.get("rge", None) + ) + if external_chi2_dict + else None + ) + + @property + def best_fit(self): + """Best fit value for the Wilson coefficients.""" + return self.results["best_fit_point"].iloc[0, :] + @property def smeft_predictions(self): """Compute |SMEFT| predictions for each replica. diff --git a/tests/fake_results/fake_results.yaml b/tests/fake_results/fake_results.yaml index 2fd1ff0e..3161e630 100644 --- a/tests/fake_results/fake_results.yaml +++ b/tests/fake_results/fake_results.yaml @@ -12,6 +12,9 @@ theory_path: ./tests/fake_data use_quad: False use_theory_covmat: True +external_chi2: + ExternalChi2: + path: ./tests/fake_external_chi2/test_ext_chi2.py # Datasets to include datasets: