From fddfa777576f8d09c016a6f1189c46d6d75dcb50 Mon Sep 17 00:00:00 2001
From: Mark Nestor Costantini <mnc33@cam.ac.uk>
Date: Mon, 26 Jun 2023 20:38:37 +0100
Subject: [PATCH] added data_index for correct indexing of closure data

---
 validphys2/src/validphys/results.py | 33 +++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/validphys2/src/validphys/results.py b/validphys2/src/validphys/results.py
index 393891d223..629c04b988 100644
--- a/validphys2/src/validphys/results.py
+++ b/validphys2/src/validphys/results.py
@@ -26,6 +26,7 @@
 )
 from validphys.convolution import PredictionsRequireCutsError, predictions
 from validphys.core import PDF, DataGroupSpec, DataSetSpec, Stats
+from validphys.plotoptions import get_info
 
 log = logging.getLogger(__name__)
 
@@ -151,6 +152,38 @@ def from_convolution(cls, pdf, posset):
         return cls(stats)
 
 
+def data_index(data):
+    """
+    Given a core.DataGroupSpec instance, return pd.MultiIndex
+    with the following levels:
+    
+    1. experiment
+    2. datasets
+    3. datapoints indices (cuts already applied to)
+    
+
+    Parameters
+    ----------
+    data: core.DataGroupSpec
+
+    Returns
+    -------
+    pd.MultiIndex
+
+    """
+    tuples = []
+
+
+    for ds in data.datasets:
+        
+        experiment = get_info(ds).experiment
+
+        for i in ds.cuts.load():
+            tp = (experiment, ds.name, i)
+            tuples.append(tp)
+    return pd.MultiIndex.from_tuples(tuples, names=('experiment', 'dataset', 'id'))
+
+
 # TODO: finish deprecating all dependencies on this index largely in theorycovmat module
 groups_data = collect("data", ("group_dataset_inputs_by_metadata",))