feat: support multiple nanoAOD versions in sample database

processor/tasks/CROWNBase.py

@@ -13,6 +13,8 @@
 # import timeout_decorator
 import time
 
+from processor.tasks.helpers.NanoAODVersions import NanoAODVersions
+
 
 class ProduceBase(Task, WrapperTask):
     """
@@ -23,14 +25,20 @@ class ProduceBase(Task, WrapperTask):
     sample_list = luigi.Parameter()
     analysis = luigi.Parameter()
     config = luigi.Parameter()
-    dataset_database = luigi.Parameter(
-        default="sample_database/datasets.json",
-        significant=False,
-    )
+    nanoAOD_version = luigi.Parameter(default=NanoAODVersions.v12.value)
+    dataset_database = luigi.Parameter(default=None, significant=False)
     shifts = luigi.Parameter()
     scopes = luigi.Parameter()
     silent = False
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # Dynamically set the default value of dataset_database based on nanoAOD_version
+        if self.dataset_database is None:
+            self.dataset_database = (
+                f"sample_database/{self.nanoAOD_version}/datasets.json"
+            )
+
     def parse_samplelist(self, sample_list):
         """
         The function `parse_samplelist` takes a sample list as input and returns a list of samples, handling
@@ -114,7 +122,7 @@ def set_sample_data(self, samples):
                 console.log(
                     "Sample {} not found in {}".format(nick, self.dataset_database)
                 )
-                raise Exception("Sample not found in DB")
+                raise Exception(f"Sample not found in DB: {nick}")
             sample_data = sample_db[nick]
             data["details"][nick]["era"] = str(sample_data["era"])
             data["details"][nick]["sample_type"] = sample_data["sample_type"]

processor/tasks/ConfigureDatasets.py

@@ -1,9 +1,9 @@
 import luigi
 import os
 import json
-import yaml
 from framework import Task
 from framework import console
+from processor.tasks.helpers.NanoAODVersions import NanoAODVersions
 
 
 def ensure_dir(file_path):
@@ -18,37 +18,27 @@ class ConfigureDatasets(Task):
     """
 
     nick = luigi.Parameter()
+    nanoAOD_version = luigi.Parameter(default=NanoAODVersions.v12.value)
     era = luigi.Parameter()
     sample_type = luigi.Parameter()
     silent = luigi.BoolParameter(default=False, significant=False)
 
     def output(self):
-        target = self.remote_target("sample_database/{}.json".format(self.nick))
+        target = self.remote_target(
+            f"sample_database/{self.nanoAOD_version}/{self.nick}.json"
+        )
         return target
 
     def load_filelist_config(self):
         # first check if a json exists, if not, check for a yaml
-        sample_configfile_json = "sample_database/{era}/{type}/{nick}.json".format(
-            era=self.era, type=self.sample_type, nick=self.nick
-        )
-        sample_configfile_yaml = "sample_database/{era}/{type}/{nick}.yaml".format(
-            era=self.era, type=self.sample_type, nick=self.nick
-        )
+        sample_configfile_json = f"sample_database/{self.nanoAOD_version}/{self.era}/{self.sample_type}/{self.nick}.json"
         if os.path.exists(sample_configfile_json):
             with open(sample_configfile_json, "r") as stream:
                 try:
                     sample_data = json.load(stream)
                 except json.JSONDecodeError as exc:
                     print(exc)
                     raise Exception("Failed to load sample information")
-        elif os.path.exists(sample_configfile_yaml):
-            console.log("[DEPRECATED] Loading from YAML")
-            with open(sample_configfile_yaml, "r") as stream:
-                try:
-                    sample_data = yaml.safe_load(stream)
-                except yaml.YAMLError as exc:
-                    print(exc)
-                    raise Exception("Failed to load sample information")
         else:
             console.log("[DEPRECATED] Loading from DAS is not supported anymore")
             raise Exception("Failed to load sample information")

processor/tasks/MLTraining.py

@@ -2,7 +2,7 @@
 
 """
 Collection of tasks used to create training datasets and config files
-for the NN trainings of the NMSSM analysis 
+for the NN trainings of the NMSSM analysis
 """
 import yaml
 import os

processor/tasks/helpers/NanoAODVersions.py

@@ -0,0 +1,7 @@
+from enum import Enum
+
+
+class NanoAODVersions(Enum):
+    v9 = "nanoAOD_v9"
+    v12 = "nanoAOD_v12"
+    v15 = "nanoAOD_v15"