Adding external likelihood in the report

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]:

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")]
 

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

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.

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: