Add function to get all sources names from all likelihoods

alea/models/blueice_extended_model.py

@@ -2,6 +2,7 @@
 from typing import List, Dict, Callable, Optional, Union, cast
 from copy import deepcopy
 from pydoc import locate
+import itertools
 
 import numpy as np
 import scipy.stats as stats
@@ -138,6 +139,20 @@ def get_source_name_list(self, likelihood_name: str) -> list:
         ll_index = self.likelihood_names.index(likelihood_name)
         return self.likelihood_list[ll_index].source_name_list
 
+    def get_all_source_names(self) -> set:
+        """Return a set of possible source names from all likelihood terms.
+
+        Args:
+            likelihood_name (str): Name of the likelihood.
+        Returns:
+            set: set of source names.
+
+        """
+        source_names = set(
+            itertools.chain.from_iterable([ll.source_name_list for ll in self.likelihood_list[:-1]])
+        )
+        return source_names
+
     @property
     def likelihood_list(self) -> List:
         """Return a list of likelihood terms."""

tests/test_blueice_extended_model.py

@@ -70,6 +70,11 @@ def get_expectation_values(self):
             expectation_values.append(this_expectation_dict)
         return expectation_values
 
+    def test_get_all_source_names(self):
+        """Test of the get_all_source_names method."""
+        for model in self.models:
+            model.get_all_source_names()
+
     def test_expectation_values(self):
         """Test of the expectation_values method."""