Enforce data-input ordering in nnfit_theory_covmat

Author: achiefa

validphys2/src/validphys/config.py

@@ -897,6 +897,19 @@ def produce_dataset_inputs_sampling_covmat(
                 else:
                     return covmats.dataset_inputs_exp_covmat
 
+    def produce_fitting_covmat_name(self, fit):
+        """Produce the name of the covmat to be used in fitting,
+        according to how it was generated by vp-setupfit.
+        """
+        runcard = fit.as_input()
+        use_thcovmat = runcard.get("theorycovmatconfig", {}).get("use_thcovmat_in_fitting", False)
+        if use_thcovmat:
+            covmat_name = "datacuts_theory_theorycovmatconfig_fitting_covmat_table.csv"
+        else:
+            covmat_name = "datacuts_theory_fitting_covmat_table.csv"
+        path = fit.path / "tables" / covmat_name
+        return path
+
     def produce_loaded_theory_covmat(
         self,
         output_path,
@@ -1363,9 +1376,9 @@ def produce_nnfit_theory_covmat(
                 f = total_theory_covmat_fitting
             else:
                 # Only scalevar uncertainties
-                from validphys.theorycovariance.construction import theory_covmat_custom
+                from validphys.theorycovariance.construction import theory_covmat_custom_fitting
 
-                f = theory_covmat_custom
+                f = theory_covmat_custom_fitting
         elif user_covmat_path is not None:
             # Only user uncertainties
             from validphys.theorycovariance.construction import user_covmat_fitting

validphys2/src/validphys/n3fit_data.py

@@ -416,13 +416,19 @@ def setupfit_fitting_covmat(dataset_inputs_fitting_covmat, diagonal_basis=True):
         log.info("working in diagonal basis.")
 
         # convert covmat to correlation
-        diag_inv_sqrt = 1 / np.sqrt(np.diag(covmat))
-        cormat = np.einsum("i, ij, j -> ij", diag_inv_sqrt, covmat, diag_inv_sqrt)
+        sigma = np.sqrt(np.diag(covmat))
+        sigma_inv = 1 / sigma
+        cormat = np.einsum("i, ij, j -> ij", sigma_inv, covmat, sigma_inv)
 
         # diagonalise the correlation matrix
-        eig_vals, uT = np.linalg.eigh(cormat)
-        uT = np.einsum("i, ik -> ik", diag_inv_sqrt, uT)
-        diagonal_rotation = uT.T
+        eig_vals, v = np.linalg.eigh(cormat)  # cormat = V @ diag(eig_vals) @ V.T
+        u = np.einsum("i, ij -> ij", sigma, v)
+        rec_covmat = u @ np.diag(eig_vals) @ u.T
+        assert np.allclose(
+            covmat, rec_covmat
+        ), "Diagonalisation failed to reproduce original covmat"
+
+        diagonal_rotation = np.einsum("ij, j -> ij", v.T, sigma_inv)
 
     return covmat, diagonal_rotation, eig_vals
 
@@ -435,9 +441,7 @@ def fitting_data_dict(
     _inv_covmat_prepared,
     kfold_masks,
     fittable_datasets_masked,
-    output_path,
     threshold=0.0,
-    diagonal_basis=True,
 ):
     """
     Provider which takes  the information from validphys ``data``.

validphys2/src/validphys/pseudodata.py

@@ -31,7 +31,7 @@ class ReplicaGenerationError(Exception):
     pass
 
 
-def fit_diagonal_basis_rotation(fit):
+def fit_diagonal_basis_rotation(fitting_covmat_name, fit):
     """Rotation matrix taking pseudodata from the original to the diagonal basis,
     or ``None`` if ``fit`` was not run in diagonal basis.
 
@@ -40,17 +40,14 @@ def fit_diagonal_basis_rotation(fit):
     applied here is bit-identical to the one used at generation time.
     """
     runcard = fit.as_input()
+
     if not runcard.get("diagonal_basis", True):
         return None
-    use_thcovmat = runcard.get("theorycovmatconfig", {}).get("use_thcovmat_in_fitting", False)
-    fname = (
-        "datacuts_theory_theorycovmatconfig_fitting_covmat_table.csv"
-        if use_thcovmat
-        else "datacuts_theory_fitting_covmat_table.csv"
-    )
+
     eigensystem = pd.read_csv(
-        fit.path / "tables" / fname, index_col=[0], header=[0], sep="\t|,", engine="python"
+        fitting_covmat_name, index_col=[0], header=[0], sep="\t|,", engine="python"
     )
+
     return eigensystem.iloc[:, 1:].values
 
 

validphys2/src/validphys/theorycovariance/construction.py

@@ -456,13 +456,10 @@ def total_theory_covmat(theory_covmat_custom, user_covmat):
     return theory_covmat_custom + user_covmat
 
 
-def theory_covmat_custom_fitting(theory_covmat_custom_per_prescription, procs_index_matched):
-    """theory_covmat_custom_per_prescription but reindexed so the order of the datasets matches
-    those in the experiment covmat so they are aligned when fitting."""
-    df = theory_covmat_custom_per_prescription.reindex(procs_index_matched).T.reindex(
-        procs_index_matched
-    )
-    return df
+def _reindex_covmat_to_fitting_order(covmat, index):
+    """Reindex a covmat DataFrame so the dataset ordering matches the experiment
+    covmat (the grouped/runcard order) used when fitting."""
+    return covmat.reindex(index).T.reindex(index)
 
 
 theory_covmats_fitting = collect(theory_covmat_custom_per_prescription, ("point_prescriptions",))
@@ -474,16 +471,22 @@ def theory_covmat_custom(theory_covmats_fitting):
     return sum(theory_covmats_fitting)
 
 
+def theory_covmat_custom_fitting(theory_covmat_custom, procs_index_matched):
+    """theory_covmat_custom (summed over all point prescriptions) reindexed so the
+    dataset ordering matches the experiment covmat for alignment when fitting."""
+    return _reindex_covmat_to_fitting_order(theory_covmat_custom, procs_index_matched)
+
+
 def total_theory_covmat_fitting(total_theory_covmat, procs_index_matched):
     """total_theory_covmat but reindexed so the order of the datasets matches
     those in the experiment covmat so they are aligned when fitting."""
-    return theory_covmat_custom_fitting(total_theory_covmat, procs_index_matched)
+    return _reindex_covmat_to_fitting_order(total_theory_covmat, procs_index_matched)
 
 
 def user_covmat_fitting(user_covmat, procs_index_matched):
     """user_covmat but reindexed so the order of the datasets matches
     those in the experiment covmat so they are aligned when fitting."""
-    return theory_covmat_custom_fitting(user_covmat, procs_index_matched)
+    return _reindex_covmat_to_fitting_order(user_covmat, procs_index_matched)
 
 
 def procs_index_matched(groups_index, procs_index):