enable template folder list

XENONnT · Jul 13, 2023 · 329980d · 329980d
1 parent 536d290
commit 329980d
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 10 deletions.
diff --git a/alea/blueice_extended_model.py b/alea/blueice_extended_model.py
@@ -93,8 +93,16 @@ def _build_ll_from_config(self, likelihood_config: dict) -> "LogLikelihoodSum":
         # Iterate through each likelihood term in the configuration
         for config in likelihood_config["likelihood_terms"]:
             likelihood_object = locate(config["likelihood_type"])
+            if isinstance(likelihood_config["template_folder"], str):
+                template_folder_list = [likelihood_config["template_folder"]]
+            elif isinstance(likelihood_config["template_folder"], list):
+                template_folder_list = likelihood_config["template_folder"]
+            else:
+                raise ValueError(
+                    "template_folder must be either a string or a list of strings.")
+
             blueice_config = adapt_likelihood_config_for_blueice(
-                config, likelihood_config["template_folder"])
+                config, template_folder_list)
             blueice_config["livetime_days"] = self.parameters[
                 blueice_config["livetime_parameter"]].nominal_value
             ll = likelihood_object(blueice_config)

diff --git a/alea/utils.py b/alea/utils.py
@@ -264,23 +264,34 @@ def adapt_inference_object_config(config_data, wimp_mass, cache_dir=None):
 
 
 def adapt_likelihood_config_for_blueice(likelihood_config: dict,
-                                        template_folder: str) -> dict:
+                                        template_folder_list: list) -> dict:
     """
     Adapt likelihood config to be compatible with blueice.
 
     Args:
         likelihood_config (dict): likelihood config dict
-        template_folder (str): base folder where templates are located.
-            If the folder starts with alea/, the alea folder is used as base.
+        template_folder_list (str): list of possible base folders where
+            templates are located. If a folder starts with alea/,
+            the alea folder is used as base.
+            Ordered by priority.
 
     Returns:
         dict: adapted likelihood config
     """
-    # if template folder starts with alea: get location of alea
-    if template_folder.startswith("alea/"):
-        import alea
-        alea_dir = os.path.dirname(os.path.abspath(alea.__file__))
-        template_folder = os.path.join(alea_dir, template_folder.replace("alea/", ""))
+    for template_folder in template_folder_list:
+        # if template folder starts with alea: get location of alea
+        if template_folder.startswith("alea/"):
+            import alea
+            alea_dir = os.path.dirname(os.path.abspath(alea.__file__))
+            template_folder = os.path.join(alea_dir, template_folder.replace("alea/", ""))
+            # check if template folder exists
+            if not os.path.isdir(template_folder):
+                template_folder = None
+            else:
+                break
+    # raise error if no template folder is found
+    if template_folder is None:
+        raise FileNotFoundError("No template folder found. Please provide a valid template folder.")
 
     likelihood_config["analysis_space"] = get_analysis_space(
         likelihood_config["analysis_space"])
@@ -290,7 +301,7 @@ def adapt_likelihood_config_for_blueice(likelihood_config: dict,
 
     for source in likelihood_config["sources"]:
         source["templatename"] = os.path.join(template_folder,
-                                              source["template_filename"])
+                                            source["template_filename"])
     return likelihood_config