Fix shifts in plot_fancy_dataspecs

validphys2/src/validphys/covmats.py

@@ -252,23 +252,16 @@ def shifts_from_systematics(lcd_wc, theory_predictions):
         points) containing the numerical value of the systematic shifts
         due to correlated uncertainties
     """
-
-    # Separate statistical and systematic errors
-    stat_errors = lcd_wc.stat_errors.to_numpy()
-    syst_errors = lcd_wc.systematic_errors(None)
-
-    # Determine the uncorrelated part of the error
-    alpha2 = stat_errors**2
-    is_uncorr = syst_errors.columns.isin(("UNCORR", "THEORYUNCORR"))
-    alpha2 += (syst_errors.loc[:, is_uncorr].to_numpy() ** 2).sum(axis=1)
-    alpha = np.sqrt(alpha2)
+
+    # Separate the uncorrelated and correlated parts of the exp uncertainty 
+    alpha = unco_unc(lcd_wc)
+    beta  = corr_unc(lcd_wc)
 
     if alpha.all() == 0:
         shifts = np.zeros(len(alpha))
     else:
 
         # Determine the correlated part of the error
-        beta = syst_errors.loc[:, ~is_uncorr].to_numpy()
         beta = beta / alpha[:, np.newaxis]
 
         # The number of data points and the number of correlated systematics
@@ -292,9 +285,57 @@ def shifts_from_systematics(lcd_wc, theory_predictions):
         # Compute the shifts
         shifts = -np.matmul(beta * alpha[:, np.newaxis], r)
 
-    return shifts, alpha
+    return shifts
+
+def unco_unc(lcd_wc):
+    """Extract the uncorrelated part of the experimental uncertainty
+    from a :py:class:`validphys.coredata.CommonData` object
+     Parameters
+    ----------
+    loaded_commondata_with_cuts : validphys.coredata.CommonData
+        CommonData which stores information about systematic errors,
+        their treatment and description.
+    Returns
+    -------
+    alpha: np.array
+        Numpy array of dimension N_dat (where N_dat is the number of data
+        points) containing the numerical value of the uncorrelated 
+        part of the experimental uncertainty
+    """
+    # Separate statistical and systematic errors
+    stat_errors = lcd_wc.stat_errors.to_numpy()
+    syst_errors = lcd_wc.systematic_errors(None)
+
+    # Determine the uncorrelated part of the error
+    alpha2 = stat_errors**2
+    is_uncorr = syst_errors.columns.isin(("UNCORR", "THEORYUNCORR"))
+    alpha2 += (syst_errors.loc[:, is_uncorr].to_numpy() ** 2).sum(axis=1)
+    alpha = np.sqrt(alpha2)
+
+    return alpha
 
+def corr_unc(lcd_wc):
+    """Extract the correlated part of the experimental uncertainty
+    from a :py:class:`validphys.coredata.CommonData` object
+     Parameters
+    ----------
+    loaded_commondata_with_cuts : validphys.coredata.CommonData
+        CommonData which stores information about systematic errors,
+        their treatment and description.
+    Returns
+    -------
+    beta: np.array
+        Numpy array of dimension N_dat (where N_dat is the number of data
+        points) containing the numerical value of the correlated 
+        part of the experimental uncertainty
+    """
+    # Separate statistical and systematic errors
+    syst_errors = lcd_wc.systematic_errors(None)
+    is_uncorr = syst_errors.columns.isin(("UNCORR", "THEORYUNCORR"))
+    beta = syst_errors.loc[:, ~is_uncorr].to_numpy()
 
+    return beta    
+
 @check_cuts_considered
 @functools.lru_cache
 def dataset_t0_predictions(t0dataset, t0set):

validphys2/src/validphys/dataplots.py

@@ -25,9 +25,9 @@
 from validphys.checks import check_not_using_pdferr
 from validphys.commondata import loaded_commondata_with_cuts
 from validphys.core import CutsPolicy, MCStats, cut_mask, load_commondata
-from validphys.covmats import shifts_from_systematics
+from validphys.covmats import shifts_from_systematics, unco_unc
 from validphys.plotoptions.core import get_info, kitable, transform_result
-from validphys.results import chi2_stat_labels, chi2_stats
+from validphys.results import chi2_stat_labels, chi2_stats, DataResult
 from validphys.sumrules import POL_LIMS
 from validphys.utils import sane_groupby_iter, scale_from_grid, split_ranges
 
@@ -263,10 +263,6 @@ def _plot_fancy_impl(
     nkinlabels = len(table.columns)
     ndata = len(table)
 
-    # Compute shifts due to the correlated part of the exp cov matrix
-    lcd_wc = loaded_commondata_with_cuts(commondata, cutlist[0])
-    theory_predictions = results[1].central_value
-
     # This is easier than cheking every time
     if labellist is None:
         labellist = [None] * len(results)
@@ -281,36 +277,40 @@ def _plot_fancy_impl(
         # We modify the table, so we pass only the label columns
         norm_cv, _ = transform_result(cv, err, table.iloc[:, :nkinlabels], info)
 
-    cvcols = []
-
-    # Compute systematic shifts
-    # For unknown reasons, `shifts_from_systematics` may
-    # randomly fails. If a LinAlgError is raised, shifts are not included in
-    # the final plot.
-    if with_shift:
-        try:
-            shifts, alpha = shifts_from_systematics(lcd_wc, theory_predictions)
-        except np.linalg.LinAlgError:
-            log.warning(
-                "Error occurred in computing systematic shifts for "
-                f"{info.ds_metadata.name}. These will be disregarded in the plots."
-            )
-            with_shift = False
-
+    cvcols = []     
+
     for i, (result, cuts) in enumerate(zip(results, cutlist)):
-        # We modify the table, so we pass only the label columns
+
         mask = cut_mask(cuts)
         cv = np.full(ndata, np.nan)
         err = np.full(ndata, np.nan)
-        # Shift the theory when with_shift option is True
-        if i == 1 and with_shift:
+
+        shifts = None
+        alpha = None
+        do_shift = with_shift
+
+        # Compute shifts due to the correlated part on the experiemntal ucnertainty
+        if do_shift:
+            lcd_wc = loaded_commondata_with_cuts(commondata, cuts)
+            theory_predictions = result.central_value
+
+            # For unknown reasons, `shifts_from_systematics` may randomly fail.
+            # If a LinAlgError is raised, shifts are not included in the final plot. 
+            try:
+                shifts = shifts_from_systematics(lcd_wc, theory_predictions)
+            except np.linalg.LinAlgError:
+                log.warning(
+                    "Error occurred in computing systematic shifts for "
+                    f"{info.ds_metadata.name}. These will be disregarded in the plots."
+                )
+                shifts = 0.
+
             cv[mask] = result.central_value - shifts
+            err[mask] = unco_unc(lcd_wc)
+
+        # No shifts
         else:
             cv[mask] = result.central_value
-        # Retain only the uncorrelated part of the error if shifting the data
-        if i == 0 and with_shift:
-            err[mask] = alpha
-        else:
             err[mask] = result.std_error
 
         cv, err = transform_result(cv, err, table.iloc[:, :nkinlabels], info)