Join observables into one list

n3fit/src/n3fit/model_trainer.py

@@ -39,22 +39,16 @@
 InputInfo = namedtuple("InputInfo", ["input", "split", "idx"])
 
 
-def _compute_observables(pdf_layers, observables):
-    """
-    Takes as input a list of PDF layers each corresponding to one observable (also given as a list)
-    Returns a list of obs(pdf).
-    """
-    return [obs(pdf) for obs, pdf, in zip(observables, pdf_layers)]
-
-
-def _compute_masked_losses(observables, losses, masks):
+def _compute_masked_losses(all_observables, losses, masks):
     """
     Takes as input a list of observables and (where neded) a mask to select the output.
     Returns a list of loss(masked_obs).
     Note that the list of masks don't need to be the same size as the list of layers/observables
     Also note that the masks here are at the level of datasets, masking out entire folds.
     Masking at the level of training/validation splits happens inside the losses.
     """
+    loss_names = [loss.name.strip('_val').strip('_exp') for loss in losses]
+    observables = [all_observables[loss_name] for loss_name in loss_names]
     return [loss(obs, mask=m) for obs, loss, m in zip_longest(observables, losses, masks)]
 
 
@@ -203,9 +197,9 @@ def __init__(
 
         # Initialize the dictionaries which contain all fitting information
         self.input_list = []
+        self.observables = []
         self.training = {
             "output": [],
-            "observables": [],
             "expdata": [],
             "ndata": 0,
             "model": None,
@@ -219,21 +213,13 @@ def __init__(
         }
         self.validation = {
             "output": [],
-            "observables": [],
             "expdata": [],
             "ndata": 0,
             "model": None,
             "folds": [],
             "posdatasets": [],
         }
-        self.experimental = {
-            "output": [],
-            "observables": [],
-            "expdata": [],
-            "ndata": 0,
-            "model": None,
-            "folds": [],
-        }
+        self.experimental = {"output": [], "expdata": [], "ndata": 0, "model": None, "folds": []}
         self.tr_masks = []
 
         self._fill_the_dictionaries()
@@ -432,6 +418,10 @@ def _model_generation(self, xinput, pdf_model, partition, partition_idx):
 
         # Now reorganize the uniques PDF so that each experiment receives its corresponding PDF
         split_pdf = [split_pdf_unique[i] for i in xinput.idx]
+
+        # Now compute ALL observables, without any masking, into a dictionary
+        all_observables = {obs.name: obs(pdf) for obs, pdf in zip(self.observables, split_pdf)}
+
         # If we are in a kfolding partition, select which datasets are out
         training_mask = validation_mask = experimental_mask = [None]
         if partition and partition["datasets"]:
@@ -445,8 +435,7 @@ def _model_generation(self, xinput, pdf_model, partition, partition_idx):
 
         # Training and validation leave out the kofld dataset
         # experiment leaves out the negation
-        observables_tr = _compute_observables(split_pdf, self.training["observables"])
-        output_tr = _compute_masked_losses(observables_tr, self.training["output"], training_mask)
+        output_tr = _compute_masked_losses(all_observables, self.training["output"], training_mask)
         training = MetaModel(full_model_input_dict, output_tr)
 
         # Validation skips integrability and the "true" chi2 skips also positivity,
@@ -461,16 +450,14 @@ def _model_generation(self, xinput, pdf_model, partition, partition_idx):
                 val_pdfs.append(partial_pdf)
 
         # We don't want to included the integrablity in the validation
-        observables_vl = _compute_observables(val_pdfs, self.validation["observables"])
         output_vl = _compute_masked_losses(
-            observables_vl, self.validation["output"], validation_mask
+            all_observables, self.validation["output"], validation_mask
         )
         validation = MetaModel(full_model_input_dict, output_vl)
 
         # Or the positivity in the total chi2
-        observables_ex = _compute_observables(exp_pdfs, self.experimental["observables"])
         output_ex = _compute_masked_losses(
-            observables_ex, self.experimental["output"], experimental_mask
+            all_observables, self.experimental["output"], experimental_mask
         )
         experimental = MetaModel(full_model_input_dict, output_ex)
 
@@ -571,9 +558,7 @@ def _generate_observables(
             self.input_list.append(exp_layer["inputs"])
 
             # Now save the observable layer
-            self.training["observables"].append(exp_layer["observables"])
-            self.validation["observables"].append(exp_layer["observables"])
-            self.experimental["observables"].append(exp_layer["observables"])
+            self.observables.append(exp_layer["observables"])
 
             # Now save the masked losses
             self.training["output"].append(exp_layer["output_tr"])
@@ -608,8 +593,7 @@ def _generate_observables(
 
             # The positivity should be on both training and validation models
             # Observables
-            self.training["observables"].append(pos_layer["observables"])
-            self.validation["observables"].append(pos_layer["observables"])
+            self.observables.append(pos_layer["observables"])
             # Masked losses
             self.training["output"].append(pos_layer["output_tr"])
             self.validation["output"].append(pos_layer["output_tr"])
@@ -637,7 +621,7 @@ def _generate_observables(
                 self.input_list.append(integ_layer["inputs"])
 
                 # The integrability all falls to the training
-                self.training["observables"].append(integ_layer["observables"])
+                self.observables.append(integ_layer["observables"])
                 self.training["output"].append(integ_layer["output_tr"])
                 self.training["integmultipliers"].append(integ_multiplier)
                 self.training["integinitials"].append(integ_initial)