Repository: testingautomated-usi/uncertainty-wizard

Issue:
multi-quantifier quantification as_dict 
The methods `predict_quantified` on stochastic models and `predict_quantified` and `quantify_predictions` on ensembles should accept an optional flag (default `False`) allowing the following behavior:

- The results are returned as Dict[str, Dict[str, np.ndarray]], where the outer dict keys are a particular quantifiers alias and the inner dict keys are 'prediction' and 'quantification' (see the example below)
- If a quantifier was passed as an alias (str), the same alias should be used as dict key. If the quantifier was passed as an object, the first alias of the class should be used.

E.g., the return value of `predict_quantified(quantifiers=['var_ratio', PredictiveEntropy()], as_dict=True)` should be
```
{
     "var_ratio": {"prediction": <pred_array>, "quantification": <quantification_array>},
     "predictive_entropy": {"prediction": <pred_array>, "quantification": <quantification_array>}
}
```

The goal of this issue is to allow to generate outputs that are self-explanatory and do not need to remember the order of the passed quantifiers.
`stochastic_from_keras` should not use static file path
Executing `stochastic_from_keras` in different processes at the same time, using default parameters, can lead to failures, as both processes use the same temporary weight path. 

A workaround exists: Just pass custom, differing (e.g. random) `temp_weights_path` when calling `stochastic_from_keras`.

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/testingautomated-usi/uncertainty-wizard/pull/157?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=testingautomated-usi) Report
Base: **94.70**% // Head: **94.84**% // Increases project coverage by **`+0.14%`** :tada:
> Coverage data is based on head [(`2d1eaf0`)](https://codecov.io/gh/testingautomated-usi/uncertainty-wizard/pull/157?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=testingautomated-usi) compared to base [(`738ed10`)](https://codecov.io/gh/testingautomated-usi/uncertainty-wizard/commit/738ed105e65d347bac5cad9406431dc610d7056b?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=testingautomated-usi).
> Patch coverage: 100.00% of modified lines in pull request are covered.

> :exclamation: Current head 2d1eaf0 differs from pull request most recent head e10a9a3. Consider uploading reports for the commit e10a9a3 to get more accurate results

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #157      +/-   ##
==========================================
+ Coverage   94.70%   94.84%   +0.14%     
==========================================
  Files          27       27              
  Lines         925      931       +6     
==========================================
+ Hits          876      883       +7     
+ Misses         49       48       -1     
```


| [Impacted Files](https://codecov.io/gh/testingautomated-usi/uncertainty-wizard/pull/157?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=testingautomated-usi) | Coverage Δ | |
|---|---|---|
| [...\_wizard/models/\_stochastic/\_abstract\_stochastic.py](https://codecov.io/gh/testingautomated-usi/uncertainty-wizard/pull/157?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=testingautomated-usi#diff-dW5jZXJ0YWludHlfd2l6YXJkL21vZGVscy9fc3RvY2hhc3RpYy9fYWJzdHJhY3Rfc3RvY2hhc3RpYy5weQ==) | `90.65% <100.00%> (+0.26%)` | :arrow_up: |
| [...nty\_wizard/models/ensemble\_utils/\_lazy\_ensemble.py](https://codecov.io/gh/testingautomated-usi/uncertainty-wizard/pull/157?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=testingautomated-usi#diff-dW5jZXJ0YWludHlfd2l6YXJkL21vZGVscy9lbnNlbWJsZV91dGlscy9fbGF6eV9lbnNlbWJsZS5weQ==) | `92.00% <100.00%> (+0.84%)` | :arrow_up: |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=testingautomated-usi). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=testingautomated-usi)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/testingautomated-usi/uncertainty-wizard/pull/157?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=testingautomated-usi).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=testingautomated-usi).


Patch:
diff --git a/uncertainty_wizard/models/_stochastic/_abstract_stochastic.py b/uncertainty_wizard/models/_stochastic/_abstract_stochastic.py
index 7a058f6..7235dc9 100644
--- a/uncertainty_wizard/models/_stochastic/_abstract_stochastic.py
+++ b/uncertainty_wizard/models/_stochastic/_abstract_stochastic.py
@@ -262,6 +262,7 @@ class Stochastic(_UwizModel):
         steps=None,
         as_confidence: Union[None, bool] = None,
         broadcaster: Broadcaster = None,
+        return_alias_dict: bool = False,
     ):
         """
         Calculates predictions and uncertainties (or confidences) according to the passed quantifer(s).
@@ -277,7 +278,10 @@ class Stochastic(_UwizModel):
         :param as_confidence: If true, uncertainties are multiplied by (-1),
         if false, confidences are multiplied by (-1). Default: No transformations.
         :param broadcaster: Sampling Related Dependencies. If None, the DefaultBroadcaster will be used.
-        :return: A tuple (predictions, uncertainties_or_confidences) if a single quantifier was
+        :param return_alias_dict: If true, the result is returned as a dictionary with the quantifier aliases as keys.
+        :return: If `return_alias_dict=True`, a dict with all quantifier aliases as keys
+            and (predictions, uncertainties_or_confidences) as values.
+        Otherwise (default), a tuple (predictions, uncertainties_or_confidences) if a single quantifier was
         passed as string or instance, or a collection of such tuples if the passed quantifiers was an iterable.
         """
         all_q, pp_q, sample_q, return_single_tuple = self._quantifiers_as_list(
@@ -313,17 +317,25 @@ class Stochastic(_UwizModel):
             )
 
         results = self._run_quantifiers(
-            as_confidence, point_prediction_scores, all_q, stochastic_scores
+            as_confidence,
+            point_prediction_scores,
+            all_q,
+            stochastic_scores,
+            as_dict=return_alias_dict,
         )
-        if return_single_tuple:
+        if return_single_tuple and not return_alias_dict:
             return results[0]
         return results
 
     @staticmethod
     def _run_quantifiers(
-        as_confidence, point_prediction_scores, quantifiers, stochastic_scores
+        as_confidence,
+        point_prediction_scores,
+        quantifiers,
+        stochastic_scores,
+        as_dict=False,
     ):
-        results = []
+        results = dict() if as_dict else []
         for q in quantifiers:
             if q.takes_samples():
                 assert stochastic_scores is not None, (
@@ -341,7 +353,11 @@ class Stochastic(_UwizModel):
             superv_scores = q.cast_conf_or_unc(
                 as_confidence=as_confidence, superv_scores=superv_scores
             )
-            results.append((predictions, superv_scores))
+            if as_dict:
+                for alias in q.aliases():
+                    results[alias] = (predictions, superv_scores)
+            else:
+                results.append((predictions, superv_scores))
         return results
 
     @staticmethod
diff --git a/uncertainty_wizard/models/ensemble_utils/_lazy_ensemble.py b/uncertainty_wizard/models/ensemble_utils/_lazy_ensemble.py
index 919d4af..c7eee4a 100644
--- a/uncertainty_wizard/models/ensemble_utils/_lazy_ensemble.py
+++ b/uncertainty_wizard/models/ensemble_utils/_lazy_ensemble.py
@@ -469,6 +469,7 @@ class LazyEnsemble(_UwizModel):
         num_processes=None,
         context=None,
         models: Optional[Iterable[int]] = None,
+        return_alias_dict: bool = False,
     ):
         """
         Utility function to make quantified predictions on numpy arrays.
@@ -486,8 +487,11 @@ class LazyEnsemble(_UwizModel):
         :param context: A contextmanager which prepares a newly crated process for execution
             (e.g. by configuring the gpus). See class docstring for explanation of default values.
         :param models: A list of model indices to use for prediction. Default: `None`(All models).
-        :return: A tuple (predictions, uncertainties_or_confidences) if a single quantifier was passed as string
-            or instance, or a collection of such tuples if the passed quantifiers was an iterable.
+        :param return_alias_dict: If true, the result is returned as a dictionary with the quantifier aliases as keys.
+        :return: If `return_alias_dict=True`, a dict with all quantifier aliases as keys
+            and (predictions, uncertainties_or_confidences) as values.
+            Otherwise (default), a tuple (predictions, uncertainties_or_confidences) if a single quantifier was
+            passed as string or instance, or a collection of such tuples if the passed quantifiers was an iterable.
         """
         if verbose > 0:
             warnings.warn("Verbosity not yet supported in lazy ensemble models.")
@@ -501,6 +505,7 @@ class LazyEnsemble(_UwizModel):
             num_processes=num_processes,
             context=context,
             models=models,
+            return_alias_dict=return_alias_dict,
         )
 
     def quantify_predictions(
@@ -511,6 +516,7 @@ class LazyEnsemble(_UwizModel):
         num_processes: int = None,
         context: Callable[[int], EnsembleContextManager] = None,
         models: Optional[Iterable[int]] = None,
+        return_alias_dict: bool = False,
     ):
         """
         A utility function to make predictions on all atomic models and then infer overall predictions and uncertainty
@@ -525,7 +531,11 @@ class LazyEnsemble(_UwizModel):
         :param num_processes: The number of processes to use. Default: The default or value specified when creating the lazy ensemble.
         :param context: A contextmanager which prepares a newly crated process for execution (e.g. by configuring the gpus). See class docstring for explanation of default values.
         :param models: A list of model indices to use for prediction. Default: `None`(All models).
-        :return: A tuple (predictions, uncertainties_or_confidences) if a single quantifier was passed as string or instance, or a collection of such tuples if the passed quantifiers was an iterable.
+        :param return_alias_dict: If true, the result is returned as a dictionary with the quantifier aliases as keys.
+        :return: If `return_alias_dict=True`, a dict with all quantifier aliases as keys
+            and (predictions, uncertainties_or_confidences) as values.
+            Otherwise (default), a tuple (predictions, uncertainties_or_confidences) if a single quantifier was
+            passed as string or instance, or a collection of such tuples if the passed quantifiers was an iterable.
         """
         all_q, pp_q, sample_q, return_single_tuple = self._quantifiers_as_list(
             quantifier
@@ -549,14 +559,19 @@ class LazyEnsemble(_UwizModel):
                 scores = np.empty(scores_shape)
             scores[:, i] = predictions
 
-        results = []
+        results = dict() if return_alias_dict else list()
         for q in all_q:
             predictions, superv_scores = q.calculate(scores)
             superv_scores = q.cast_conf_or_unc(
                 as_confidence=as_confidence, superv_scores=superv_scores
             )
-            results.append((predictions, superv_scores))
-        if return_single_tuple:
+            # Add the predictions and superv_scores to the results
+            if return_alias_dict:
+                for alias in q.aliases():
+                    results[alias] = (predictions, superv_scores)
+            else:
+                results.append((predictions, superv_scores))
+        if return_single_tuple and not return_alias_dict:
             return results[0]
         return results
 


Test patch:
diff --git a/tests_unit/models_tests/test_lazy_ensemble.py b/tests_unit/models_tests/test_lazy_ensemble.py
index 5516312..62bd6c7 100644
--- a/tests_unit/models_tests/test_lazy_ensemble.py
+++ b/tests_unit/models_tests/test_lazy_ensemble.py
@@ -4,6 +4,7 @@ import numpy as np
 import tensorflow as tf
 
 import uncertainty_wizard as uwiz
+from uncertainty_wizard.quantifiers import StandardDeviation
 
 DUMMY_MODEL_PATH = "tmp/dummy_lazy_ensemble"
 
@@ -56,6 +57,25 @@ class LazyEnsembleTest(TestCase):
         self.assertEqual(pred.shape, (10, 1))
         self.assertEqual(std.shape, (10, 1))
 
+    def test_result_as_dict(self):
+        ensemble = uwiz.models.LazyEnsemble(
+            num_models=2, model_save_path=DUMMY_MODEL_PATH, default_num_processes=0
+        )
+        ensemble.create(create_function=create_dummy_atomic_model)
+        res = ensemble.predict_quantified(
+            x=np.ones((10, 1000)),
+            quantifier="std",
+            num_processes=0,
+            return_alias_dict=True,
+        )
+        assert isinstance(res, dict)
+        for alias in StandardDeviation().aliases():
+            assert alias in res
+            assert type(res[alias]) == tuple
+            assert len(res[alias]) == 2
+            assert res[alias][0].shape == (10, 1)
+            assert res[alias][1].shape == (10, 1)
+
     def test_dummy_main_and_one_distinct_process_are_equivalent(self):
         ensemble = uwiz.models.LazyEnsemble(
             num_models=2, model_save_path=DUMMY_MODEL_PATH
diff --git a/tests_unit/models_tests/test_sequential_stochastic.py b/tests_unit/models_tests/test_sequential_stochastic.py
index 56891e3..a148646 100644
--- a/tests_unit/models_tests/test_sequential_stochastic.py
+++ b/tests_unit/models_tests/test_sequential_stochastic.py
@@ -6,7 +6,7 @@ import tensorflow as tf
 import uncertainty_wizard as uwiz
 from uncertainty_wizard.internal_utils import UncertaintyWizardWarning
 from uncertainty_wizard.models import StochasticSequential
-from uncertainty_wizard.quantifiers import StandardDeviation
+from uncertainty_wizard.quantifiers import MaxSoftmax, StandardDeviation, VariationRatio
 
 
 class SequentialStochasticTest(TestCase):
@@ -17,6 +17,58 @@ class SequentialStochasticTest(TestCase):
         model.add(tf.keras.layers.Dropout(rate=0.5))
         return model
 
+    @staticmethod
+    def _dummy_classifier():
+        model = StochasticSequential()
+        model.add(tf.keras.layers.Input(shape=1000))
+        model.add(tf.keras.layers.Dropout(rate=0.5))
+        model.add(tf.keras.layers.Dense(10, activation="softmax"))
+        # compile the model
+        model.compile(
+            optimizer=tf.keras.optimizers.Adam(),
+            loss=tf.keras.losses.CategoricalCrossentropy(),
+            metrics=[tf.keras.metrics.CategoricalAccuracy()],
+        )
+        return model
+
+    def test_result_as_dict(self):
+        model = self._dummy_classifier()
+        x = np.ones((10, 1000))
+        res = model.predict_quantified(
+            x=x,
+            quantifier=[
+                "MaxSoftmax",
+                VariationRatio(),
+            ],
+            return_alias_dict=True,
+        )
+
+        self.assertTrue(isinstance(res, dict))
+        for key, values in res.items():
+            self.assertTrue(isinstance(key, str))
+            self.assertTrue(isinstance(values, tuple))
+            self.assertEqual(len(values), 2)
+            self.assertEqual(values[0].shape, (10,))
+            self.assertEqual(values[1].shape, (10,))
+
+        for q in [MaxSoftmax(), VariationRatio()]:
+            for a in q.aliases():
+                self.assertTrue(a in res.keys())
+
+    def test_return_type_default_multi_quant(self):
+        model = self._dummy_classifier()
+        x = np.ones((10, 1000))
+        res = model.predict_quantified(x=x, quantifier=["MaxSoftmax", VariationRatio()])
+        self.assertTrue(isinstance(res, list))
+        self.assertTrue(len(res), 2)
+
+    def test_return_type_default_single_quant(self):
+        model = self._dummy_classifier()
+        x = np.ones((10, 1000))
+        res = model.predict_quantified(x=x, quantifier="MaxSoftmax")
+        self.assertTrue(isinstance(res, tuple))
+        self.assertTrue(len(res), 2)
+
     def test_predict_is_deterministic(self):
         model = self._dummy_model()
         y = model.predict(x=np.ones((10, 1000)))