implementing suggestions

src/smefit/analyze/fisher.py

@@ -384,7 +384,7 @@ def set_ticks(ax, yticks, xticks, latex_names, x_labels):
         ax.grid(visible=True, which="minor", alpha=0.2)
 
     @staticmethod
-    def plot_values(ax, dfs, cmap, norm):
+    def plot_values(ax, dfs, cmap, norm, labels=None):
         """
         Plot the values of the Fisher information.
 
@@ -398,6 +398,8 @@ def plot_values(ax, dfs, cmap, norm):
             colour map
         norm: matplotlib.colors.BoundaryNorm
             normalisation of colorbar
+        labels: list, optional
+            label elements for legend
         """
 
         df_1 = dfs[0]
@@ -432,7 +434,7 @@ def plot_values(ax, dfs, cmap, norm):
                             f"{elem_2:.1f}",
                             va="center",
                             ha="center",
-                            fontsize=8,
+                            fontsize=10,
                         )
 
                         # Create a triangle patch for the second element
@@ -456,7 +458,7 @@ def plot_values(ax, dfs, cmap, norm):
                         f"{elem_1:.1f}",
                         va="center",
                         ha="center",
-                        fontsize=8,
+                        fontsize=10,
                     )
                     if df_2 is not None:
 
@@ -480,14 +482,14 @@ def plot_values(ax, dfs, cmap, norm):
                                 closed=True,
                                 fc="none",
                                 edgecolor="black",
-                                label="$\\rm w/\\;RGE$",
+                                label=labels[0],
                             ),
                             mpatches.Polygon(
                                 [[0.5, -0.5], [0.5, 0.5], [0.5, 0.5]],
                                 closed=True,
                                 fc="none",
                                 edgecolor="black",
-                                label="$\\rm w/o\\;RGE$",
+                                label=labels[1],
                             ),
                         ]
                         # Add the legend to the plot
@@ -529,6 +531,7 @@ def plot_heatmap(
         other=None,
         summary_only=True,
         figsize=(11, 15),
+        labels=None,
         column_names=None,
     ):
 
@@ -543,18 +546,23 @@ def plot_heatmap(
             )
             # unify the fisher tables and fill missing values by zeros
             fisher_dfs = self.unify_fishers(fisher_df, fisher_df_other)
-            quad_fisher_dfs = self.unify_fishers(quad_fisher_df, quad_fisher_df_other)
 
             # reshuffle the tables according to the latex names ordering
             fisher_dfs = [
                 fisher[latex_names.index.get_level_values(level=1)]
                 for fisher in fisher_dfs
             ]
 
-            quad_fisher_dfs = [
-                fisher[latex_names.index.get_level_values(level=1)]
-                for fisher in quad_fisher_dfs
-            ]
+            if quad_fisher_df is not None:
+                quad_fisher_dfs = self.unify_fishers(
+                    quad_fisher_df, quad_fisher_df_other
+                )
+
+                # reshuffle the tables according to the latex names ordering
+                quad_fisher_dfs = [
+                    fisher[latex_names.index.get_level_values(level=1)]
+                    for fisher in quad_fisher_dfs
+                ]
 
         else:
             fisher_dfs = [fisher_df[latex_names.index.get_level_values(level=1)]]
@@ -587,7 +595,7 @@ def plot_heatmap(
         else:
             ax = plt.gca()
 
-        self.plot_values(ax, fisher_dfs, cmap, norm)
+        self.plot_values(ax, fisher_dfs, cmap, norm, labels)
 
         self.set_ticks(
             ax,
@@ -602,7 +610,7 @@ def plot_heatmap(
 
         if quad_fisher_df is not None:
             ax = fig.add_subplot(122)
-            self.plot_values(ax, quad_fisher_dfs, cmap, norm)
+            self.plot_values(ax, quad_fisher_dfs, cmap, norm, labels)
 
             self.set_ticks(
                 ax,

src/smefit/analyze/report.py

@@ -552,7 +552,7 @@ def fisher(
 
             if plot is not None:
                 fit_plot = copy.deepcopy(plot)
-                fit_plot.pop("together")
+                fit_plot.pop("together", None)
                 title = fit.label if fit_plot.pop("title") else None
                 fisher_cal.plot_heatmap(
                     self.coeff_info,
@@ -578,6 +578,7 @@ def fisher(
                 f"{self.report}/fisher_heatmap_both",
                 title=title,
                 other=fisher_1,
+                labels=[fit.label for fit in self.fits],
                 **fit_plot,
             )
             figs_list.append(f"fisher_heatmap_both")

src/smefit/fit_manager.py

@@ -127,15 +127,12 @@ def load_configuration(self):
     def load_datasets(self):
         """Load all datasets."""
 
-        if self.rgemat is None:
-            operators_to_keep2 = self.config["coefficients"]
-        else:
-            operators_to_keep2 = self.operators_to_keep
-
         self.datasets = load_datasets(
             self.config["data_path"],
             self.config["datasets"],
-            operators_to_keep2,
+            self.config["coefficients"]
+            if self.rgemat is None
+            else self.operators_to_keep,
             self.config["use_quad"],
             self.config["use_theory_covmat"],
             False,  # t0 is not used here because in the report we look at the experimental chi2
@@ -148,8 +145,6 @@ def load_datasets(self):
             rgemat=self.rgemat,
         )
 
-        # TODO: load also without rg for the fisher
-
     @property
     def smeft_predictions(self):
         """Compute |SMEFT| predictions for each replica.

src/smefit/loader.py

@@ -300,6 +300,7 @@ def load_theory(
         lin_dict = {}
 
         # save sm prediction at the chosen perturbative order
+
         sm = np.array(raw_th_data[order]["SM"])
 
         # split corrections into a linear and quadratic dict