Slight doc changes

Author: stefanradev93

bayesflow/amortizers.py

@@ -20,7 +20,6 @@
 
 import numpy as np
 import tensorflow as tf
-from tensorflow.keras import Model
 
 from bayesflow.exceptions import ConfigurationError, SummaryStatsError
 from bayesflow.losses import log_loss, mmd_summary_space, kl_dirichlet
@@ -59,7 +58,7 @@ def log_prob(input_dict, **kwargs):
         pass
 
 
-class AmortizedPosterior(Model, AmortizedTarget):
+class AmortizedPosterior(tf.keras.Model, AmortizedTarget):
     """A wrapper to connect an inference network for parameter estimation with an optional summary network
     as in the original BayesFlow set-up described in the paper:
 
@@ -115,7 +114,7 @@ def __init__(self, inference_net, summary_net=None, latent_dist=None, latent_is_
         any `sumamry_conditions`, i.e., `summary_conditions` should be set to None, otherwise these will be ignored.
         """
 
-        Model.__init__(self, **kwargs)
+        tf.keras.Model.__init__(self, **kwargs)
 
         self.inference_net = inference_net
         self.summary_net = summary_net
@@ -409,7 +408,7 @@ def _determine_summary_loss(self, loss_fun):
             raise NotImplementedError("Could not infer summary_loss_fun, argument should be of type (None, callable, or str)!")
 
 
-class AmortizedLikelihood(Model, AmortizedTarget):
+class AmortizedLikelihood(tf.keras.Model, AmortizedTarget):
     """An interface for a surrogate model of a simulator, or an implicit likelihood
     ``p(params | data, context).''
     """
@@ -427,7 +426,7 @@ def __init__(self, surrogate_net, latent_dist=None, **kwargs):
             a multivariate unit Gaussian.
         """
 
-        Model.__init__(self, **kwargs)
+        tf.keras.Model.__init__(self, **kwargs)
 
         self.surrogate_net = surrogate_net
         self.latent_dim = self.surrogate_net.latent_dim
@@ -616,7 +615,7 @@ def _determine_latent_dist(self, latent_dist):
             return latent_dist
 
 
-class AmortizedPosteriorLikelihood(Model, AmortizedTarget):
+class AmortizedPosteriorLikelihood(tf.keras.Model, AmortizedTarget):
     """An interface for jointly learning a surrogate model of the simulator and an approximate
     posterior given a generative model.
     """
@@ -632,7 +631,7 @@ def __init__(self, amortized_posterior, amortized_likelihood, **kwargs):
             The generative neural likelihood approximator.
         """
 
-        Model.__init__(self, **kwargs)
+        tf.keras.Model.__init__(self, **kwargs)
 
         self.amortized_posterior = amortized_posterior
         self.amortized_likelihood = amortized_likelihood

bayesflow/diagnostics.py

@@ -409,7 +409,7 @@ def plot_sbc_histograms(post_samples, prior_samples, param_names=None, fig_size=
     return f
 
 
-def plot_posterior_2d(posterior_draws, prior=None, prior_draws=None, param_names=None, height=2, 
+def plot_posterior_2d(posterior_draws, prior=None, prior_draws=None, param_names=None, height=3, 
                       legend_fontsize=14, post_color='#8f2727', prior_color='gray', post_alpha=0.9, 
                       prior_alpha=0.7):
     """Generates a bivariate pairplot given posterior draws and optional prior or prior draws.
@@ -423,7 +423,7 @@ def plot_posterior_2d(posterior_draws, prior=None, prior_draws=None, param_names
         will be used.
     param_names       : list or None, optional, default: None
         The parameter names for nice plot titles. Inferred if None
-    height            : float, optional, default: 2.
+    height            : float, optional, default: 3.
         The height of the pairplot.
     legend_fontsize   : int, optional, default: 14
         The font size of the legend text.
@@ -466,9 +466,12 @@ def plot_posterior_2d(posterior_draws, prior=None, prior_draws=None, param_names
     # Attempt to determine parameter names
     if param_names is None:
         if hasattr(prior, 'param_names'):
-            param_names = prior.param_names
+            if prior.param_names is not None:
+                param_names = prior.param_names
+            else:
+                param_names = [f'$p_{i}$' for i in range(1, n_params+1)]
         else:
-            param_names = [f'Param. {p}' for p in range(1, n_params+1)]
+            param_names = [f'$p_{i}$' for i in range(1, n_params+1)]
 
     # Pack posterior draws into a dataframe
     posterior_draws_df = pd.DataFrame(posterior_draws, columns=param_names)

bayesflow/trainers.py

@@ -85,7 +85,7 @@ class Trainer:
     """
 
     def __init__(self, amortizer, generative_model=None, configurator=None, checkpoint_path=None, 
-                 max_to_keep=3, default_lr=0.0005, skip_checks=False, memory=True, **kwargs):
+                 max_to_keep=3, default_lr=0.001, skip_checks=False, memory=True, **kwargs):
         """Creates a trainer which will use a generative model (or data simulated from it) to optimize
         a neural arhcitecture (amortizer) for amortized posterior inference, likelihood inference, or both.
 
@@ -102,7 +102,7 @@ def __init__(self, amortizer, generative_model=None, configurator=None, checkpoi
             Optional file path for storing the trained amortizer, loss history and optional memory.
         max_to_keep       : int, optional, default: 3
             Number of checkpoints and loss history snapshots to keep.
-        default_lr        : float, optional, default: 0.0005
+        default_lr        : float, optional, default: 0.001
             The default learning rate to use for default optimizers.
         skip_checks       : boolean, optional, default: False
             If True, do not perform consistency checks, i.e., simulator runs and passed through nets
@@ -271,7 +271,7 @@ def diagnose_sbc_histograms(self, inputs=None, n_samples=None, **kwargs):
 
             # Check for prior names and override keyword if available
             plot_kwargs = kwargs.pop('plot_args', {})
-            if type(self.generative_model) is GenerativeModel and plot_kwargs.get('param_names') is not None:
+            if type(self.generative_model) is GenerativeModel and plot_kwargs.get('param_names') is None:
                 plot_kwargs['param_names'] = self.generative_model.param_names
 
             return plot_sbc_histograms(post_samples, prior_samples, **plot_kwargs)
@@ -414,7 +414,9 @@ def train_offline(self, simulations_dict, epochs, batch_size, save_checkpoint=Tr
 
         # Convert to custom data set
         data_set = SimulationDataset(simulations_dict, batch_size)
+        # Prepare optimizer and initislize loss history
         self._setup_optimizer(optimizer, epochs, len(data_set.data))
+
         self.loss_history.start_new_run()
         for ep in range(1, epochs+1):