Update defaults

bayesflow/amortizers.py

@@ -1036,6 +1036,179 @@ def _determine_loss(self, loss_fun):
  )
 
 
+class TwoLevelAmortizedPosterior(tf.keras.Model, AmortizedTarget):
+ """An interface for estimating arbitrary two level hierarchical Bayesian models."""
+
+ def __init__(self, local_amortizer, global_amortizer, summary_net=None, **kwargs):
+ """Creates an wrapper for estimating two-level hierarchical Bayesian models.
+
+ Parameters
+ ----------
+ local_amortizer : bayesflow.amortizers.AmortizedPosterior
+ A posterior amortizer without a summary network which will estimate
+ the full conditional of the (varying numbers of) local parameter vectors.
+ global_amortizer : bayesflow.amortizers.AmortizedPosterior
+ A posterior amortizer without a summary network which will estimate the joint
+ posterior of hyperparameters and optional shared parameters given a representation
+ of an entire hierarchical data set. If both hyper- and shared parameters are present,
+ the first dimensions correspond to the hyperparameters and the remaining ones correspond
+ to the shared parameters.
+ summary_net : tf.keras.Model or None, optional, default: None
+ An optional summary network to compress non-vector data structures.
+ **kwargs : dict, optional, default: {}
+ Additional keyword arguments passed to the ``__init__`` method of a ``tf.keras.Model`` instance.
+ """
+
+ super().__init__(**kwargs)
+
+ self.local_amortizer = local_amortizer
+ self.global_amortizer = global_amortizer
+ self.summary_net = summary_net
+
+ def call(self, input_dict, **kwargs):
+ """Forward pass through the hierarchical amortized posterior."""
+
+ local_summaries, global_summaries = self._compute_condition(input_dict, **kwargs)
+ local_inputs, global_inputs = self._prepare_inputs(input_dict, local_summaries, global_summaries)
+ local_out = self.local_amortizer(local_inputs, **kwargs)
+ global_out = self.global_amortizer(global_inputs, **kwargs)
+ return local_out, global_out
+
+ def compute_loss(self, input_dict, **kwargs):
+ """Compute loss of all amortizers."""
+
+ local_summaries, global_summaries = self._compute_condition(input_dict, **kwargs)
+ local_inputs, global_inputs = self._prepare_inputs(input_dict, local_summaries, global_summaries)
+ local_loss = self.local_amortizer.compute_loss(local_inputs, **kwargs)
+ global_loss = self.global_amortizer.compute_loss(global_inputs, **kwargs)
+ return {"Local.Loss": local_loss, "Global.Loss": global_loss}
+
+ def sample(self, input_dict, n_samples, to_numpy=True, **kwargs):
+ """Obtains samples from the joint hierarchical posterior given observations.
+
+ Important: Currently works only for single hierarchical data sets!
+
+ Parameters
+ ----------
+ input_dict : dict
+ Input dictionary containing the following mandatory keys, if DEFAULT_KEYS unchanged:
+ `summary_conditions` - the hierarchical data set (to be embedded by the summary net)
+ As well as optional keys:
+ `direct_local_conditions` - (Context) variables used to condition the local posterior
+ `direct_global_conditions` - (Context) variables used to condition the global posterior
+ n_samples : int
+ The number of posterior draws (samples) to obtain from the approximate posterior
+ to_numpy : bool, optional, default: True
+ Flag indicating whether to return the samples as a `np.array` or a `tf.Tensor`
+ **kwargs : dict, optional, default: {}
+ Additional keyword arguments passed to the summary network as the amortizers
+
+ Returns:
+ --------
+ samples_dict : dict
+ A dictionary with keys `global_samples` and `local_samples`
+ Global samples will hold an array-like of shape (num_samples, num_replicas, num_local)
+ and local samples will hold an array-like of shape (num_samples, num_hyper + num_shared),
+ if optional shared patameters are present, otherwise (num_samples, num_hyper),
+ """
+
+ # Returned shapes will be
+ # local_summaries.shape = (1, num_groups, summary_dim_local)
+ # global_summaries.shape = (1, summary_dim_global)
+ local_summaries, global_summaries = self._get_local_global(input_dict, **kwargs)
+ num_groups = local_summaries.shape[1]
+
+ if local_summaries.shape[0] != 1 or global_summaries.shape[0] != 1:
+ raise NotImplementedError("Method currently supports only single hierarchical data sets!")
+
+ # Obtain samples from p(global | all_data)
+ inp_global = {DEFAULT_KEYS["direct_conditions"]: global_summaries}
+
+ # New, shape will be (n_samples, num_globals)
+ global_samples = self.global_amortizer.sample(inp_global, n_samples, **kwargs, to_numpy=False)
+
+ # Repeat local conditions for n_samples
+ # New shape -> (num_groups, n_samples, summary_dim_local)
+ local_summaries = tf.stack([tf.squeeze(local_summaries, axis=0)] * n_samples, axis=1)
+
+ # Repeat global samples for num_groups
+ # New shape -> (num_groups, n_samples, num_globals)
+ global_samples_rep = tf.stack([global_samples] * num_groups, axis=0)
+
+ # Concatenate local summaries with global samples
+ # New shape -> (num_groups, num_samples, summary_dim_local + num_globals)
+ local_summaries = tf.concat([local_summaries, global_samples_rep], axis=-1)
+
+ # Obtain samples from p(local_i | data_i, global_i)
+ inp_local = {DEFAULT_KEYS["direct_conditions"]: local_summaries}
+ local_samples = self.local_amortizer.sample(inp_local, n_samples, to_numpy=False, **kwargs)
+
+ if to_numpy:
+ global_samples = global_samples.numpy()
+ local_samples = local_samples.numpy()
+
+ return {"global_samples": global_samples, "local_samples": local_samples}
+
+ def log_prob(self, input_dict):
+ """Compute normalized log density."""
+
+ raise NotImplementedError
+
+ def _prepare_inputs(self, input_dict, local_summaries, global_summaries):
+ """Prepare input dictionaries for both amortizers."""
+
+ # Prepare inputs for local amortizer
+ local_inputs = {"direct_conditions": local_summaries, "parameters": input_dict["local_parameters"]}
+
+ # Prepare inputs for global amortizer
+ _parameters = input_dict["hyper_parameters"]
+ if input_dict.get("shared_parameters") is not None:
+ _parameters = tf.concat([_parameters, input_dict.get("shared_parameters")], axis=-1)
+ global_inputs = {"direct_conditions": global_summaries, "parameters": _parameters}
+ return local_inputs, global_inputs
+
+ def _compute_condition(self, input_dict, **kwargs):
+ """Determines conditionining variables for both amortizers."""
+
+ # Obtain needed summaries
+ local_summaries, global_summaries = self._get_local_global(input_dict, **kwargs)
+
+ # At this point, add globals as conditions
+ num_locals = local_summaries.shape[1]
+
+ # Add hyper parameters as conditions:
+ # p(local_n | data_n, hyper)
+ if input_dict.get("hyper_parameters") is not None:
+ _params = input_dict.get("hyper_parameters")
+ _conds = tf.stack([_params] * num_locals, axis=1)
+ local_summaries = tf.concat([local_summaries, _conds], axis=-1)
+ # Add shared parameters as conditions:
+ # p(local_n | data_n, hyper, shared)
+ if input_dict.get("shared_parameters") is not None:
+ _params = input_dict.get("shared_parameters")
+ _conds = tf.stack([_params] * num_locals, axis=1)
+ local_summaries = tf.concat([local_summaries, _conds], axis=-1)
+ return local_summaries, global_summaries
+
+ def _get_local_global(self, input_dict, **kwargs):
+ """Helper function to obtain local and global condition tensors."""
+
+ # Obtain summary conditions
+ if self.summary_net is not None:
+ local_summaries, global_summaries = self.summary_net(
+ input_dict["summary_conditions"], return_all=True, **kwargs
+ )
+ if input_dict.get("direct_local_conditions") is not None:
+ local_summaries = tf.concat([local_summaries, input_dict.get("direct_local_conditions")], axis=-1)
+ if input_dict.get("direct_global_conditions") is not None:
+ global_summaries = tf.concat([global_summaries, input_dict.get("direct_global_conditions")], axis=-1)
+ # If no summary net provided, assume direct conditions exist or fail
+ else:
+ local_summaries = input_dict.get("direct_local_conditions")
+ global_summaries = input_dict.get("direct_global_conditions")
+ return local_summaries, global_summaries
+
+
 class SingleModelAmortizer(AmortizedPosterior):
  """Deprecated class for amortizer posterior estimation."""
 

bayesflow/default_settings.py

@@ -176,12 +176,18 @@ def __init__(self, meta_dict: dict, mandatory_fields: list = []):
  "non_batchable_context": "non_batchable_context",
  "prior_batchable_context": "prior_batchable_context",
  "prior_non_batchable_context": "prior_non_batchable_context",
+ "prior_context": "prior_context",
+ "hyper_prior_draws": "hyper_prior_draws",
+ "shared_prior_draws": "shared_prior_draws",
+ "local_prior_draws": "local_prior_draws",
  "sim_batchable_context": "sim_batchable_context",
  "sim_non_batchable_context": "sim_non_batchable_context",
  "summary_conditions": "summary_conditions",
  "direct_conditions": "direct_conditions",
  "parameters": "parameters",
- "hyperparameters": "hyperparameters",
+ "hyper_parameters": "hyper_parameters",
+ "shared_parameters": "shared_parameters",
+ "local_parameters": "local_parameters",
  "observables": "observables",
  "targets": "targets",
  "conditions": "conditions",