Modify model comparison approximator

Author: stefanradev93

bayesflow/approximators/model_comparison_approximator.py

@@ -25,16 +25,16 @@ class ModelComparisonApproximator(Approximator):
 
     Parameters
     ----------
-    adapter: Adapter
-        Adapter for data processing.
+    adapter: bf.adapters.Adapter
+        Adapter for data pre-processing.
     num_models: int
         Number of models (simulators) that the approximator will compare
-    classifier_network: keras.Model
-        The network (e.g, an MLP) that is used for model classification.
+    classifier_network: keras.Layer
+        The network backbone (e.g, an MLP) that is used for model classification.
         The input of the classifier network is created by concatenating `classifier_variables`
         and (optional) output of the summary_network.
-    summary_network: SummaryNetwork, optional
-        The summary network used for data summarisation (default is None).
+    summary_network: bg.networks.SummaryNetwork, optional
+        The summary network used for data summarization (default is None).
         The input of the summary network is `summary_variables`.
     """
 
@@ -51,7 +51,7 @@ def __init__(
         self.classifier_network = classifier_network
         self.adapter = adapter
         self.summary_network = summary_network
-
+        self.num_models = num_models
         self.logits_projector = keras.layers.Dense(num_models)
 
     def build(self, data_shapes: Mapping[str, Shape]):
@@ -61,6 +61,7 @@ def build(self, data_shapes: Mapping[str, Shape]):
     @classmethod
     def build_adapter(
         cls,
+        num_models: int,
         classifier_conditions: Sequence[str] = None,
         summary_variables: Sequence[str] = None,
         model_index_name: str = "model_indices",
@@ -80,11 +81,9 @@ def build_adapter(
             adapter.rename(model_index_name, "model_indices")
             .keep(["classifier_conditions", "summary_variables", "model_indices"])
             .standardize(exclude="model_indices")
+            .one_hot("model_indices", num_models)
         )
 
-        # TODO: add one-hot encoding
-        # .one_hot("model_indices", self.num_models)
-
         return adapter
 
     @classmethod
@@ -239,7 +238,7 @@ def fit(
 
         if adapter == "auto":
             logging.info("Building automatic data adapter.")
-            adapter = self.build_adapter(**filter_kwargs(kwargs, self.build_adapter))
+            adapter = self.build_adapter(num_models=self.num_models, **filter_kwargs(kwargs, self.build_adapter))
 
         if simulator is not None:
             return super().fit(simulator=simulator, adapter=adapter, **kwargs)
@@ -252,15 +251,17 @@ def fit(
 
     @classmethod
     def from_config(cls, config, custom_objects=None):
-        adapter = deserialize(config["adapter"], custom_objects=custom_objects)
-        classifier_network = deserialize(config["classifier_network"], custom_objects=custom_objects)
-        summary_network = deserialize(config["summary_network"], custom_objects=custom_objects)
-        return cls(adapter=adapter, classifier_network=classifier_network, summary_network=summary_network, **config)
+        config["num_models"] = deserialize(config["num_models"], custom_objects=custom_objects)
+        config["adapter"] = deserialize(config["adapter"], custom_objects=custom_objects)
+        (config["classifier_network"],) = deserialize(config["classifier_network"], custom_objects=custom_objects)
+        config["summary_network"] = deserialize(config["summary_network"], custom_objects=custom_objects)
+        return super().from_config(config, custom_objects=custom_objects)
 
     def get_config(self):
         base_config = super().get_config()
 
         config = {
+            "num_models": serialize(self.num_models),
             "adapter": serialize(self.adapter),
             "classifier_network": serialize(self.classifier_network),
             "summary_network": serialize(self.summary_network),