from collections.abc import Sequence, Callable, Mapping
from typing import Literal, Tuple
import keras
import numpy as np
from bayesflow.adapters import Adapter
from bayesflow.networks import InferenceNetwork, SummaryNetwork
from bayesflow.masking import MaskName
from bayesflow.types import Tensor
from bayesflow.utils import split_arrays
from bayesflow.utils.serialization import serializable
from .basic_estimator import BasicEstimator
from .helpers import build_prior_score_fn
[docs]
@serializable("bayesflow.estimators")
class CompositionalEstimator(BasicEstimator):
"""
Defines a wrapper for estimating arbitrary compositional distributions of the form:
`p(inference_variables | [summary(summary_variables), ...], inference_conditions)`
Parameters
----------
inference_network : InferenceNetwork
The inference network used for posterior or likelihood approximation.
adapter : bayesflow.adapters.Adapter, optional
Adapter for data processing. You can use :py:meth:`build_adapter`
to create it. If ``None`` (default), an identity adapter is used
that makes a shallow copy and passes data through unchanged.
summary_network : SummaryNetwork, optional
The summary network used for data summarization (default is None).
standardize : str | Sequence[str] | None
The variables to standardize before passing to the networks. Can be either
"all" or any subset of ["inference_variables", "summary_variables", "inference_conditions"].
(default is "inference_variables").
**kwargs : dict, optional
Additional arguments passed to the :py:class:`bayesflow.estimators.Estimator` class.
"""
def __init__(
self,
*,
inference_network: InferenceNetwork,
adapter: Adapter | None = None,
summary_network: SummaryNetwork | None = None,
standardize: str | Sequence[str] | None = "inference_variables",
**kwargs,
):
super().__init__(
adapter=adapter,
inference_network=inference_network,
summary_network=summary_network,
standardize=standardize,
**kwargs,
)
[docs]
def compositional_sample(
self,
*,
num_samples: int,
conditions: dict[str, np.ndarray] | None = None,
compute_prior_score: Callable[[dict[str, Tensor], Tensor | None], dict[str, Tensor]] | None = None,
split: bool = False,
batch_size: int | None = None,
sample_shape: Literal["infer"] | Tuple[int] | int = "infer",
return_summaries: bool = False,
summary_outputs: Tensor | np.ndarray | None = None,
to_numpy: bool = True,
**kwargs,
) -> dict[str, np.ndarray | Tensor]:
"""
Generates compositional samples from the estimator given input conditions. The `conditions` dictionary is
preprocessed using the `adapter`. Samples are converted to NumPy arrays after inference.
Expected shape of each condition variable is (n_datasets, n_compositional, ...), where n_compositional >= 2.
Parameters
----------
num_samples : int
Number of samples to generate.
conditions : dict[str, np.ndarray], optional
Dictionary of conditioning variables as NumPy arrays.
compute_prior_score : Callable[[dict[str, Tensor], Tensor | None], dict[str, Tensor]], optional
A function that computes the score of the log prior distribution in the original parameter space, using
backend tensors. Optionally, the function can have a time argument, otherwise the prior score is
multiplied with (1-t), where t is diffusion time. If none provided, the unconditional score is used.
The score is mapped into the sampling space by a change of variables over the adapter, so adapters with
non-zero log determinant are supported, provided their transforms are differentiable.
split : bool, default=False
Whether to split the output arrays along the last axis and return one sample array per target variable.
batch_size : int or None, optional
If provided, the conditions are split into batches of size `batch_size`, for which samples are generated
sequentially. Can help with memory management for large sample sizes.
sample_shape : str or tuple of int, optional
Trailing structural dimensions of each generated sample, excluding the batch and target (intrinsic)
dimension. For example, use `(time,)` for time series or `(height, width)` for images.
If set to `"infer"` (default), the structural dimensions are inferred from the `inference_conditions`.
In that case, all non-vector dimensions except the last (channel) dimension are treated as structural
dimensions. For example, if the final `inference_conditions` have shape `(batch_size, time, channels)`,
then `sample_shape` is inferred as `(time,)`, and the generated samples will have shape
`(num_conditions, num_samples, time, target_dim)`.
return_summaries: bool, optional
If set to True and a summary network is present, will return the learned summary statistics for
the provided conditions.
summary_outputs : Tensor | np.ndarray | None, optional
Precomputed summary outputs to be used as conditions for sampling. If provided, these will be used instead
of the conditions. Should have shape (n_datasets, n_compositional_conditions, ...).
to_numpy: bool, optional
If True, the returned samples be converted to numpy arrays
(as opposed to returned as keras backend tensors).
**kwargs : dict
Additional keyword arguments for the sampling process.
Returns
-------
dict[str, np.ndarray | Tensor]
Dictionary containing generated samples with the same keys as `conditions`.
"""
resolved_conditions, adapted, summary_outputs = self._prepare_compositional_conditions(
conditions=conditions, batch_size=batch_size, summary_outputs=summary_outputs
)
if compute_prior_score is None:
compute_prior_score_pre = None
else:
compute_prior_score_pre = build_prior_score_fn(
compute_prior_score, adapter=self.adapter, standardizer=self.standardizer
)
kwargs = self._maybe_inject_guidance_unstandardize(kwargs)
inference_kwargs = self._prepare_inference_kwargs(kwargs, adapted=adapted)
inference_kwargs["compute_prior_score"] = compute_prior_score_pre
# NOTE: infer option of sample cannot handle the compositional dimensions
if sample_shape == "infer":
sample_shape = tuple(keras.ops.shape(resolved_conditions)[2:-1])
samples = self.sampler.sample(
inference_network=self.inference_network,
num_samples=num_samples,
conditions=resolved_conditions,
batch_size=batch_size,
sample_shape=sample_shape,
**inference_kwargs,
)
# Unstandardize and inverse-adapt samples (tree-aware for nested dict outputs)
samples = keras.tree.map_structure(
lambda s: self.standardizer.maybe_standardize(
s, key="inference_variables", stage="inference", forward=False
),
samples,
)
samples = keras.tree.map_structure(
lambda s: self.adapter({"inference_variables": s}, inverse=True, strict=False),
samples,
)
if to_numpy:
samples = keras.tree.map_structure(keras.ops.convert_to_numpy, samples)
if return_summaries and summary_outputs is not None:
samples["_summaries"] = summary_outputs
if split:
samples = split_arrays(samples, axis=-1)
return samples
def _prepare_compositional_conditions(
self,
conditions: Mapping[str, np.ndarray] | None,
batch_size: int | None = None,
summary_outputs: Tensor | np.ndarray | None = None,
**kwargs,
) -> tuple[Tensor | None, dict[str, Tensor], Tensor | None]:
if summary_outputs is not None:
num_datasets, num_items = keras.ops.shape(summary_outputs)[:2]
summary_outputs = keras.ops.reshape(
summary_outputs, (num_datasets * num_items,) + keras.ops.shape(summary_outputs)[2:]
)
flattened_conditions = None
elif conditions is not None:
original_shapes = {}
flattened_conditions = {}
target_mask_names = {
MaskName.FIXED_TARGET,
MaskName.INFER_TARGET,
MaskName.FIXED_TARGET_VALUE,
}
for key, value in conditions.items():
original_shapes[key] = value.shape
if key in target_mask_names:
flattened_conditions[key] = value
continue
num_datasets, num_items = value.shape[:2]
flattened_shape = (num_datasets * num_items,) + value.shape[2:]
flattened_conditions[key] = value.reshape(flattened_shape)
first_compositional_key = next(key for key in original_shapes if key not in target_mask_names)
num_datasets, num_items = original_shapes[first_compositional_key][:2]
else:
raise ValueError(
"At least one of 'conditions' or 'summary_outputs' must be provided for compositional sampling."
)
if num_items <= 1:
raise ValueError(
"At least two conditioning variables are required for compositional sampling, got "
f"{num_items}. Use 'sample' instead."
)
resolved_conditions, adapted, summary_outputs = self._prepare_conditions(
data=flattened_conditions,
summary_outputs=summary_outputs,
batch_size=batch_size,
**kwargs,
)
# Reshape tensors back to (num_datasets, num_items, ...)
restored_shape = (num_datasets, num_items) + keras.ops.shape(resolved_conditions)[1:]
resolved_conditions = keras.ops.reshape(resolved_conditions, restored_shape)
if summary_outputs is not None:
summary_outputs = keras.ops.reshape(summary_outputs, restored_shape)
observed_condition_mask = adapted.get(MaskName.OBSERVED_CONDITION)
if observed_condition_mask is not None:
adapted = dict(adapted)
adapted[MaskName.OBSERVED_CONDITION] = keras.ops.reshape(
observed_condition_mask,
(num_datasets, num_items) + tuple(keras.ops.shape(observed_condition_mask)[1:]),
)
return resolved_conditions, adapted, summary_outputs
