Merge branch 'dev' into allow-networks

Author: LarsKue

bayesflow/adapters/adapter.py

@@ -667,6 +667,18 @@ def shift(self, keys: str | Sequence[str], by: float | np.ndarray):
         self.transforms.append(MapTransform({key: Shift(shift=by) for key in keys}))
         return self
 
+    def split(self, key: str, *, into: Sequence[str], indices_or_sections: int | Sequence[int] = None, axis: int = -1):
+        from .transforms import Split
+
+        if isinstance(into, str):
+            transform = Rename(key, into)
+        else:
+            transform = Split(key, into, indices_or_sections, axis)
+
+        self.transforms.append(transform)
+
+        return self
+
     def sqrt(self, keys: str | Sequence[str]):
         """Append an :py:class:`~transforms.Sqrt` transform to the adapter.
 

bayesflow/adapters/transforms/__init__.py

@@ -17,6 +17,7 @@
 from .scale import Scale
 from .serializable_custom_transform import SerializableCustomTransform
 from .shift import Shift
+from .split import Split
 from .sqrt import Sqrt
 from .standardize import Standardize
 from .to_array import ToArray

bayesflow/adapters/transforms/concatenate.py

@@ -23,7 +23,7 @@ class Concatenate(Transform):
     Examples
     --------
     Suppose you have a simulator that generates variables "beta" and "sigma" from priors and then observation
-    variables "x" and "y". We can then use concatonate in the following way
+    variables "x" and "y". We can then use concatenate in the following way
 
     >>> adapter = (
         bf.Adapter()

bayesflow/adapters/transforms/split.py

@@ -0,0 +1,112 @@
+from collections.abc import Sequence
+import numpy as np
+
+from keras.saving import (
+    deserialize_keras_object as deserialize,
+    register_keras_serializable as serializable,
+    serialize_keras_object as serialize,
+)
+
+from .transform import Transform
+
+
+@serializable(package="bayesflow.adapters")
+class Split(Transform):
+    """This is the effective inverse of the :py:class:`~Concatenate` Transform.
+
+    Parameters
+    ----------
+    key : str
+        The key to split in the forward transform.
+    into: Sequence[str]
+        The names of each split after the forward transform.
+    indices_or_sections : int | Sequence[int], optional, default: None
+        The number of sections or indices to split on. If not given, will split evenly into len(into) parts.
+    axis: int, optional, default: -1
+        The axis to split on.
+    """
+
+    def __init__(self, key: str, into: Sequence[str], indices_or_sections: int | Sequence[int] = None, axis: int = -1):
+        self.axis = axis
+        self.key = key
+        self.into = into
+
+        if indices_or_sections is None:
+            indices_or_sections = len(into)
+
+        self.indices_or_sections = indices_or_sections
+
+    @classmethod
+    def from_config(cls, config: dict, custom_objects=None) -> "Split":
+        return cls(
+            key=deserialize(config["key"], custom_objects),
+            into=deserialize(config["into"], custom_objects),
+            indices_or_sections=deserialize(config["indices_or_sections"], custom_objects),
+            axis=deserialize(config["axis"], custom_objects),
+        )
+
+    def get_config(self) -> dict:
+        return {
+            "key": serialize(self.key),
+            "into": serialize(self.into),
+            "indices_or_sections": serialize(self.indices_or_sections),
+            "axis": serialize(self.axis),
+        }
+
+    def forward(self, data: dict[str, np.ndarray], strict: bool = True, **kwargs) -> dict[str, np.ndarray]:
+        # avoid side effects
+        data = data.copy()
+
+        if strict and self.key not in data:
+            raise KeyError(self.key)
+        elif self.key not in data:
+            # we cannot produce a result, but also don't have to
+            return data
+
+        splits = np.split(data.pop(self.key), self.indices_or_sections, axis=self.axis)
+
+        if len(splits) != len(self.into):
+            raise ValueError(f"Requested {len(self.into)} splits, but produced {len(splits)}.")
+
+        for key, split in zip(self.into, splits):
+            data[key] = split
+
+        return data
+
+    def inverse(self, data: dict[str, np.ndarray], strict: bool = False, **kwargs) -> dict[str, np.ndarray]:
+        # avoid side effects
+        data = data.copy()
+
+        required_keys = set(self.into)
+        available_keys = set(data.keys())
+        common_keys = available_keys & required_keys
+        missing_keys = required_keys - available_keys
+
+        if strict and missing_keys:
+            # invalid call
+            raise KeyError(f"Missing keys: {missing_keys!r}")
+        elif missing_keys:
+            # we cannot produce a result, but should still remove the keys
+            for key in common_keys:
+                data.pop(key)
+
+            return data
+
+        # remove each part
+        splits = [data.pop(key) for key in self.into]
+
+        # concatenate them all
+        result = np.concatenate(splits, axis=self.axis)
+
+        # store the result
+        data[self.key] = result
+
+        return data
+
+    def extra_repr(self) -> str:
+        result = "[" + ", ".join(map(repr, self.key)) + "] -> " + repr(self.into)
+
+        if self.axis != -1:
+            result += f", axis={self.axis}"
+
+        return result

bayesflow/simulators/__init__.py

@@ -12,8 +12,16 @@
 from .simulator import Simulator
 
 from .benchmark_simulators import (
+    BernoulliGLM,
+    BernoulliGLMRaw,
+    GaussianLinear,
+    GaussianLinearUniform,
+    GaussianMixture,
+    InverseKinematics,
     LotkaVolterra,
     SIR,
+    SLCP,
+    SLCPDistractors,
     TwoMoons,
 )
 

bayesflow/simulators/benchmark_simulators/__init__.py

@@ -1,3 +1,11 @@
+from .bernoulli_glm import BernoulliGLM
+from .bernoulli_glm_raw import BernoulliGLMRaw
+from .gaussian_linear import GaussianLinear
+from .gaussian_linear_uniform import GaussianLinearUniform
+from .gaussian_mixture import GaussianMixture
+from .inverse_kinematics import InverseKinematics
 from .lotka_volterra import LotkaVolterra
 from .sir import SIR
+from .slcp import SLCP
+from .slcp_distractors import SLCPDistractors
 from .two_moons import TwoMoons

bayesflow/simulators/benchmark_simulators/gaussian_linear.py

@@ -75,5 +75,10 @@ def observation_model(self, params: np.ndarray):
         # Generate prior predictive samples, possibly a single if n_obs is None
         if self.n_obs is None:
             return self.rng.normal(loc=params, scale=self.obs_scale)
-        x = self.rng.normal(loc=params, scale=self.obs_scale, size=(self.n_obs, params.shape[0], params.shape[1]))
-        return np.transpose(x, (1, 0, 2))
+        if params.ndim == 2:
+            # batched sampling with n_obs
+            x = self.rng.normal(loc=params, scale=self.obs_scale, size=(self.n_obs, params.shape[0], params.shape[1]))
+            return np.transpose(x, (1, 0, 2))
+        elif params.ndim == 1:
+            # non-batched sampling with n_obs
+            return self.rng.normal(loc=params, scale=self.obs_scale, size=(self.n_obs, params.shape[0]))

bayesflow/simulators/benchmark_simulators/gaussian_linear_uniform.py

@@ -79,5 +79,10 @@ def observation_model(self, params: np.ndarray):
         # Generate prior predictive samples, possibly a single if n_obs is None
         if self.n_obs is None:
             return self.rng.normal(loc=params, scale=self.obs_scale)
-        x = self.rng.normal(loc=params, scale=self.obs_scale, size=(self.n_obs, params.shape[0], params.shape[1]))
-        return np.transpose(x, (1, 0, 2))
+        if params.ndim == 2:
+            # batched sampling with n_obs
+            x = self.rng.normal(loc=params, scale=self.obs_scale, size=(self.n_obs, params.shape[0], params.shape[1]))
+            return np.transpose(x, (1, 0, 2))
+        elif params.ndim == 1:
+            # non-batched sampling with n_obs
+            return self.rng.normal(loc=params, scale=self.obs_scale, size=(self.n_obs, params.shape[0]))

bayesflow/simulators/sequential_simulator.py

@@ -10,7 +10,7 @@
 class SequentialSimulator(Simulator):
     """Combines multiple simulators into one, sequentially."""
 
-    def __init__(self, simulators: Sequence[Simulator], expand_outputs: bool = True):
+    def __init__(self, simulators: Sequence[Simulator], expand_outputs: bool = True, replace_inputs: bool = True):
         """
         Initialize a SequentialSimulator.
 
@@ -22,10 +22,13 @@ def __init__(self, simulators: Sequence[Simulator], expand_outputs: bool = True)
         expand_outputs : bool, optional
             If True, 1D output arrays are expanded with an additional dimension at the end.
             Default is True.
+        replace_inputs : bool, optional
+            If True, **kwargs are auto-batched and replace simulator outputs.
         """
 
         self.simulators = simulators
         self.expand_outputs = expand_outputs
+        self.replace_inputs = replace_inputs
 
     @allow_batch_size
     def sample(self, batch_shape: Shape, **kwargs) -> dict[str, np.ndarray]:
@@ -53,6 +56,14 @@ def sample(self, batch_shape: Shape, **kwargs) -> dict[str, np.ndarray]:
         for simulator in self.simulators:
             data |= simulator.sample(batch_shape, **(kwargs | data))
 
+            if self.replace_inputs:
+                common_keys = set(data.keys()) & set(kwargs.keys())
+                for key in common_keys:
+                    value = kwargs.pop(key)
+                    if isinstance(data[key], np.ndarray):
+                        value = np.broadcast_to(value, data[key].shape)
+                    data[key] = value
+
         if self.expand_outputs:
             data = {
                 key: np.expand_dims(value, axis=-1) if np.ndim(value) == 1 else value for key, value in data.items()

bayesflow/utils/dict_utils.py

@@ -233,7 +233,7 @@ def make_variable_array(
     else:
         raise TypeError(f"Only dicts and tensors are supported as arguments, but your estimates are of type {type(x)}")
 
-    if len(variable_names) is not x.shape[-1]:
+    if len(variable_names) != x.shape[-1]:
         raise ValueError("Length of 'variable_names' should be the same as the number of variables.")
 
     if variable_keys is None: