update README and TwoMoons example for discussion

README.md

@@ -9,107 +9,88 @@ It provides users with:
 
 - A user-friendly API for rapid Bayesian workflows
 - A rich collection of neural network architectures
-- Multi-Backend Support: [PyTorch](https://github.com/pytorch/pytorch), [TensorFlow](https://github.com/tensorflow/tensorflow), [JAX](https://github.com/google/jax), and [NumPy](https://github.com/numpy/numpy)
+- Multi-Backend Support via [Keras3](https://keras.io/keras_3/): You can use [PyTorch](https://github.com/pytorch/pytorch), [TensorFlow](https://github.com/tensorflow/tensorflow), or [JAX](https://github.com/google/jax)
 
 BayesFlow is designed to be a flexible and efficient tool that enables rapid statistical inference
 fueled by continuous progress in generative AI and Bayesian inference.
 
+## Conceptual Overview
+
+A cornerstone idea of amortized Bayesian inference is to employ generative
+neural networks for parameter estimation, model comparison, and model validation
+when working with intractable simulators whose behavior as a whole is too
+complex to be described analytically. The figure below presents a higher-level
+overview of neurally bootstrapped Bayesian inference.
+
+<img src="https://github.com/stefanradev93/BayesFlow/blob/master/img/high_level_framework.png?raw=true" width=80% height=80%>
+
+
+## Disclaimer
+
+This is the current dev version of BayesFlow, which constitutes a complete refactor of the library built on Keras 3. This way, you can now use any of the major deep learning libraries as backend for BayesFlow. The refactor is still work in progress with some of the advanced features not yet implemented. We are actively working on them and promise to catch up soon.
+
+If you encounter any issues, please don't hesitate to open an issue here on [Github](https://github.com/stefanradev93/BayesFlow/issues) or ask questions on our [Discourse Forums](https://discuss.bayesflow.org/).
+
 ## Install
 
 ### Backend
 
-First, install your machine learning backend of choice. Note that BayesFlow **will not run** without a backend.
+First, install one machine learning backend of choice. Note that BayesFlow **will not run** without a backend.
 
-Once installed, set the appropriate backend environment variable. For example, to use PyTorch:
+[Install JAX](https://jax.readthedocs.io/en/latest/installation.html)
+[Install PyTorch](https://pytorch.org/get-started/locally/)
+[Install TensorFlow](https://www.tensorflow.org/install)
 
-```bash
-export KERAS_BACKEND=torch
-```
+If you are new to machine learning and don't know which one to use, we recommend PyTorch to get started.
 
-If you use conda, you can instead set this individually for each environment:
+Once installed, [set the backend environment variable as required by keras](https://keras.io/getting_started/#configuring-your-backend). For example, inside your Python script write:
 
-```bash
-conda env config vars set KERAS_BACKEND=torch
+```python
+import os
+os.environ["KERAS_BACKEND"] = "torch"
+import keras
 ```
 
-### Using Conda
-
-We recommend installing with conda (or mamba).
+If you use conda, you can alternatively set this individually for each environment in your terminal. For example:
 
 ```bash
-conda install -c conda-forge bayesflow
+conda env config vars set KERAS_BACKEND=torch
 ```
 
+This way, you also don't have to manually set the backend every time you are starting Python to use BayesFlow.
+
 ### Using pip
 
+You can install the dev version with pip:
+
 ```bash
-pip install bayesflow
+pip install git+https://github.com/stefanradev93/bayesflow@dev
 ```
 
-### From Source
+### Using Conda (coming soon)
 
-Stable version:
+The dev version is not conda-installable yet.
 
-```bash
-git clone https://github.com/stefanradev93/bayesflow
-cd bayesflow
-conda env create --file environment.yaml --name bayesflow
-```
+### From Source
 
-Development version:
+If you want to contribute to BayesFlow, we recommend installing the dev branch from source:
 
 ```bash
-git clone https://github.com/stefanradev93/bayesflow
-cd bayesflow
-git checkout dev
+git clone -b dev [email protected]:stefanradev93/bayesflow.git
+cd <local-path-to-bayesflow-repository>
 conda env create --file environment.yaml --name bayesflow
 ```
 
 ## Getting Started
 
-Check out some of our walk-through notebooks:
+Check out some of our walk-through notebooks below. We are actively working on porting all notebooks to the new interface so more will be available soon!
 
 1. [Two moons toy example with flow matching](examples/TwoMoons_FlowMatching.ipynb)
-2. ...Under construction
 
 ## Documentation \& Help
 
 Documentation is available at https://bayesflow.org. Please use the [BayesFlow Forums](https://discuss.bayesflow.org/) for any BayesFlow-related questions and discussions, and [GitHub Issues](https://github.com/stefanradev93/BayesFlow/issues) for bug reports and feature requests.
 
-## Conceptual Overview
-
-A cornerstone idea of amortized Bayesian inference is to employ generative
-neural networks for parameter estimation, model comparison, and model validation
-when working with intractable simulators whose behavior as a whole is too
-complex to be described analytically. The figure below presents a higher-level
-overview of neurally bootstrapped Bayesian inference.
-
-<img src="https://github.com/stefanradev93/BayesFlow/blob/master/img/high_level_framework.png?raw=true" width=80% height=80%>
-
-### References and Further Reading
-
-- Radev S. T., D’Alessandro M., Mertens U. K., Voss A., Köthe U., & Bürkner P.
-C. (2021). Amortized Bayesian Model Comparison with Evidental Deep Learning.
-<em>IEEE Transactions on Neural Networks and Learning Systems</em>.
-doi:10.1109/TNNLS.2021.3124052 available for free at: https://arxiv.org/abs/2004.10629
-
-- Schmitt, M., Radev, S. T., & Bürkner, P. C. (2022). Meta-Uncertainty in
-Bayesian Model Comparison. In <em>International Conference on Artificial Intelligence
-and Statistics</em>, 11-29, PMLR, available for free at: https://arxiv.org/abs/2210.07278
-
-- Elsemüller, L., Schnuerch, M., Bürkner, P. C., & Radev, S. T. (2023). A Deep
-Learning Method for Comparing Bayesian Hierarchical Models. <em>ArXiv preprint</em>,
-available for free at: https://arxiv.org/abs/2301.11873
-
-- Radev, S. T., Schmitt, M., Pratz, V., Picchini, U., Köthe, U., & Bürkner, P.-C. (2023).
-JANA: Jointly amortized neural approximation of complex Bayesian models.
-*Proceedings of the Thirty-Ninth Conference on Uncertainty in Artificial Intelligence, 216*, 1695-1706.
-([arXiv](https://arxiv.org/abs/2302.09125))([PMLR](https://proceedings.mlr.press/v216/radev23a.html))
-
-## Support
-
-This project is currently managed by researchers from Rensselaer Polytechnic Institute, TU Dortmund University, and Heidelberg University. It is partially funded by the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation, Project 528702768). The project is further supported by Germany's Excellence Strategy -- EXC-2075 - 390740016 (Stuttgart Cluster of Excellence SimTech) and EXC-2181 - 390900948 (Heidelberg Cluster of Excellence STRUCTURES), as well as the Informatics for Life initiative funded by the Klaus Tschira Foundation.
-
 ## Citing BayesFlow
 
 You can cite BayesFlow along the lines of:
@@ -156,3 +137,7 @@ You can cite BayesFlow along the lines of:
   publisher = {PMLR}
 }
 ```
+
+## Acknowledgments
+
+This project is currently managed by researchers from Rensselaer Polytechnic Institute, TU Dortmund University, and Heidelberg University. It is partially funded by the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation, Project 528702768). The project is further supported by Germany's Excellence Strategy -- EXC-2075 - 390740016 (Stuttgart Cluster of Excellence SimTech) and EXC-2181 - 390900948 (Heidelberg Cluster of Excellence STRUCTURES), as well as the Informatics for Life initiative funded by the Klaus Tschira Foundation.

bayesflow/approximators/continuous_approximator.py

@@ -131,7 +131,6 @@ def sample(
         *,
         conditions: Mapping[str, Tensor],
         num_samples: int = None,
-        numpy: bool = True,
         batch_shape: Shape = None,
     ) -> dict[str, Tensor]:
         if num_samples is None and batch_shape is None:
@@ -144,11 +143,9 @@ def sample(
         conditions = {
             "inference_variables": self._sample(num_samples=num_samples, batch_shape=batch_shape, **conditions)
         }
+        conditions = keras.tree.map_structure(keras.ops.convert_to_numpy, conditions)
         conditions = self.data_adapter.deconfigure(conditions)
 
-        if numpy:
-            conditions = keras.tree.map_structure(keras.ops.convert_to_numpy, conditions)
-
         return conditions
 
     def _sample(

bayesflow/simulators/lambda_simulator.py

@@ -8,7 +8,14 @@
 
 
 class LambdaSimulator(Simulator):
-    def __init__(self, sample_fn: callable, *, is_batched: bool = False, cast_dtypes: Mapping[str, str] = "default"):
+    def __init__(
+        self,
+        sample_fn: callable,
+        *,
+        is_batched: bool = False,
+        cast_dtypes: Mapping[str, str] = "default",
+        reserved_arguments: Mapping[str, any] = "default",
+    ):
         """Implements a simulator based on a (batched or unbatched) sampling function.
         Outputs will always be in batched format.
         :param sample_fn: The sampling function.
@@ -19,6 +26,8 @@ def __init__(self, sample_fn: callable, *, is_batched: bool = False, cast_dtypes
         :param cast_dtypes: Output data types to cast arrays to.
             By default, we convert float64 (the default for numpy on x64 systems)
             to float32 (the default for deep learning on any system).
+        :param reserved_arguments: Reserved keyword arguments to pass to the sampling function.
+            By default, functions requesting an argument 'rng' will be passed the default numpy random generator.
         """
         self.sample_fn = sample_fn
         self.is_batched = is_batched
@@ -28,7 +37,15 @@ def __init__(self, sample_fn: callable, *, is_batched: bool = False, cast_dtypes
 
         self.cast_dtypes = cast_dtypes
 
+        if reserved_arguments == "default":
+            reserved_arguments = {"rng": np.random.default_rng()}
+
+        self.reserved_arguments = reserved_arguments
+
     def sample(self, batch_shape: Shape, **kwargs) -> dict[str, np.ndarray]:
+        # add reserved arguments
+        kwargs = self.reserved_arguments | kwargs
+
         # try to use only valid keyword arguments
         kwargs = filter_kwargs(kwargs, self.sample_fn)