Update README.rst

docs/README.rst

@@ -1,145 +1,153 @@
 NeuroBayes
 ==========
 
-.. warning::
-   This package is actively in development, and breaking changes may occur frequently.
+.. important::
+   This package is actively in development, and breaking changes may occur frequently
 
-.. image:: _static/NB_logo2.png
-   :alt: NeuroBayes Logo
+.. image:: https://github.com/user-attachments/assets/36e1ea05-83cb-41a9-a52e-cf73335dd499
+   :alt: NB_logo2
 
-Overview
---------
+What is it for
+-------------
 
-Machine learning, at its core, is about approximating unknown functions – mapping inputs to outputs based on observed data. In scientific and engineering applications, this often means modeling complex relationships between process parameters and target properties. 
+Machine learning, at its core, is about approximating unknown functions – mapping inputs to outputs based on observed data. In scientific and engineering applications, this often means modeling complex relationships between process parameters and target properties. Traditionally, Gaussian Processes (GPs) have been favored in these domains for their ability to provide robust uncertainty estimates. However, GPs struggle with systems featuring discontinuities and non-stationarities, common in physical science problems, as well as with high dimensional data. **NeuroBayes** bridges this gap by combining the flexibility and scalability of neural networks with the rigorous uncertainty quantification of Bayesian methods. This repository enables the use of full BNNs and partial BNNs with the No-U-Turn Sampler for intermediate size datasets, making it a powerful tool for a wide range of scientific and engineering applications.
 
-**NeuroBayes** bridges the gap between traditional Gaussian Processes (GPs) and neural networks by combining the flexibility and scalability of neural networks with the rigorous uncertainty quantification of Bayesian methods. This package enables the use of full BNNs and partial BNNs with the No-U-Turn Sampler for intermediate size datasets, making it a powerful tool for scientific and engineering applications.
-
-Key Features
------------
-
-- Full Bayesian Neural Networks (BNN)
-- Partial Bayesian Neural Networks (PBNN)
-- Support for MLP and ConvNet architectures
-- Heteroskedastic noise modeling
-- Pre-trained priors integration
-
-.. image:: _static/NN_types.png
-   :alt: Neural Network Types
-
-Installation
------------
-
-.. code-block:: bash
+How to use
+----------
 
-   pip install neurobayes
+NeuroBayes provides two main approaches for implementing Bayesian Neural Networks: Fully Bayesian Neural Networks (BNN) and Partially Bayesian Neural Networks (PBNN). Both approaches currently support MLP and ConvNet architectures, with more architectures on the way. Here's how to use BNN and PBNN:
 
-Usage Guide
-----------
+.. image:: https://github.com/user-attachments/assets/9b58d9ec-cb7f-49de-aa58-e75990b08b83
+   :alt: NN_types
 
-Fully Bayesian Neural Networks
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Fully Bayesian Neural Nets
+~~~~~~~~~~~~~~~~~~~~~~~~~
 
-Fully Bayesian Neural Networks replace all constant weights in the network with probabilistic distributions:
+Fully Bayesian Neural Networks replace all constant weights in the network with probabilistic distributions. This approach provides comprehensive uncertainty quantification but may be computationally intensive for large models.
 
 .. code-block:: python
 
    import neurobayes as nb
 
    # Initialize NN architecture
-   architecture = nb.FlaxMLP(hidden_dims=[64, 32, 16, 8], target_dim=1)
+   architecture = nb.FlaxMLP(hidden_dims = [64, 32, 16, 8], target_dim=1)
 
    # Initialize Bayesian model
    model = nb.BNN(architecture)
-   
    # Train model
    model.fit(X_measured, y_measured, num_warmup=1000, num_samples=1000)
-   
    # Make a prediction on full domain
    posterior_mean, posterior_var = model.predict(X_domain)
 
-Partially Bayesian Neural Networks
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Partially Bayesian Neural Network
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-Partially Bayesian Neural Networks apply probabilistic distributions to only a subset of layers:
+Partially Bayesian Neural Networks replace constant weights with probabilistic distributions only in a subset of the network's layers. This approach ismore computationally efficient while still providing good uncertainty estimates. By default, the deterministic part of PBNNs is trained using Maximum A Posteriori approximation, with stochastic weight averaging applied at the end of each training trajectory.
 
 .. code-block:: python
 
    # Number of probabilistic ('Bayesian') layers
-   num_probabilistic_layers = 2  # two last learnable layers + output layer
+   num_probabilistic_layers = 2 # two last learnable layers + output layer
 
-   # Initialize a PBNN model
+   # Intitalize a PBNN model
    model = nb.PartialBNN(architecture, num_probabilistic_layers=num_probabilistic_layers)
-   
    # Train
    model.fit(X_measured, y_measured, num_warmup=1000, num_samples=1000)
-   
-   # Make a prediction
+   # Make a prediction on unmeasured points or the full domain
    posterior_mean, posterior_var = model.predict(X_domain)
 
-Alternative layer specification:
+Alternatively, we can specify directly the names of the layers we want to be probabilistic
 
 .. code-block:: python
 
-   # Specify the names of probabilistic layers
+   # Specify the names of probabilistic layers (output layer, 'Dense4', needs to be specified explicitly)
    probabilistic_layer_names = ['Dense2', 'Dense3', 'Dense4']
 
-   # Initialize and train a PBNN model
+   # Intitalize and train a PBNN model
    model = nb.PartialBNN(architecture, probabilistic_layer_names=probabilistic_layer_names)
    model.fit(X_measured, y_measured, num_warmup=1000, num_samples=1000)
 
-.. image:: _static/BNN_PBNN.png
-   :alt: BNN vs PBNN Comparison
+.. image:: https://github.com/user-attachments/assets/8281b071-4f05-4432-8e23-babcaaad6b5d
+   :alt: BNN_PBNN
 
-Heteroskedastic Noise Modeling
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+The obtained posterior means and variances can be used in active learning and Bayesian optimization frameworks. See example of BNN-powered active learning `here <https://github.com/ziatdinovmax/NeuroBayes/blob/main/examples/bnn_example1d.ipynb>`_ and example of PBNN-powered active learning `here <https://github.com/ziatdinovmax/NeuroBayes/blob/main/examples/pbnn_example1d.ipynb>`_.
 
-For datasets with input-dependent noise levels:
+Heteroskedastic noise
+~~~~~~~~~~~~~~~~~~~
+
+By default, we assume constant observation noise across all inputs. However, this assumption often doesn't hold in real-world datasets which may exhibit input-dependent levels of noise. NeuroBayes offers heteroskedastic BNNs that can capture varying levels of noise in different regions of the data, allowing for more accurate uncertainty quantification.
+
+The usage of a heteroskedastic BNN is straightforward and follows the same pattern as the standard BNN models:
+
+For fully Bayesian heteroskedastic NN:
 
 .. code-block:: python
 
-   # Fully Bayesian heteroskedastic NN
+   # Initialize HeteroskedasticBNN model
    model = nb.HeteroskedasticBNN(architecture)
+   # Train
    model.fit(X_measured, y_measured, num_warmup=2000, num_samples=2000)
+   # Make a prediction
    posterior_mean, posterior_var = model.predict(X_domain)
 
-   # Partially Bayesian heteroskedastic NN
+For partially Bayesian heteroskedastic NN:
+
+.. code-block:: python
+
+   # Initialize and train
    model = nb.HeteroskedasticPartialBNN(architecture, num_probabilistic_layers=2)
-   model.fit(X_measured, y_measured, sgd_epochs=5000, sgd_lr=5e-3, 
-            num_warmup=1000, num_samples=1000)
+   model.fit(X_measured, y_measured, sgd_epochs=5000, sgd_lr=5e-3, num_warmup=1000, num_samples=1000)
+   # Make a prediction
    posterior_mean, posterior_var = model.predict(X_domain)
 
-.. image:: _static/hsk.png
-   :alt: Heteroskedastic Noise Example
+.. image:: https://github.com/user-attachments/assets/5a619361-74c0-4d03-9b1a-4aa995f1c540
+   :alt: hsk
 
-Pre-trained Priors
+See example `here <https://github.com/ziatdinovmax/NeuroBayes/blob/main/examples/heteroskedastic.ipynb>`_.
+
+Pre-trained priors
 ~~~~~~~~~~~~~~~~
 
-Leverage existing knowledge through pre-trained models:
+.. image:: https://github.com/user-attachments/assets/b3e2d7da-7bb1-4919-8140-899795ac042d
+   :alt: Transfer_learning
+
+NeuroBayes extends the concept of partial BNNs to leverage pre-existing knowledge or simulations, allowing for more informed priors in Bayesian Neural Networks. This approach is particularly useful when you have theoretical models or simulations that can guide the learning process for experimental data.
+The process involves two main steps:
+
+Pre-training a deterministic neural network on theoretical or simulated data.
+Using the weights from this pre-trained network to center the prior distributions for a Fully Bayesian Neural Network (FBNN) or Partially Bayesian Neural Network (PBNN).
+
+Here's how to implement this approach:
+First, fit a deterministic NN to theoretical data:
 
 .. code-block:: python
 
-   # Train deterministic NN on theoretical data
    hidden_dims = [64, 32, 16, 8]
    net = nb.FlaxMLP(hidden_dims, target_dim=1)
-   detnn = nb.DeterministicNN(net, input_shape=(1,), learning_rate=5e-3, 
-                             map=True, sigma=nb.utils.calculate_sigma(X1))
+   detnn = nb.DeterministicNN(net, input_shape=(1,), learning_rate=5e-3, map=True, sigma=nb.utils.calculate_sigma(X1))
    detnn.train(X1, y1, epochs=5000, batch_size=None)
 
-   # Use pre-trained weights for BNN priors
-   model = nb.BNN(net, pretrained_priors=detnn.state.params)
+Note: In practice, you should use proper train-test-validation splits for robust model development.
+
+Next, train a BNN on experimental data, using the pre-trained weights to set theory-informed BNN priors:
+
+.. code-block:: python
+
+   model = nb.BNN(
+       net,
+       pretrained_priors=detnn.state.params  # use pre-trained weights to set priors for BNN
+   ) 
    model.fit(X2, y2, num_warmup=1000, num_samples=1000, num_chains=1)
+
+Make a prediction as ususal
+
+.. code-block:: python
+
    posterior_mean, posterior_var = model.predict(X_test)
 
-.. image:: _static/Transfer_learning.png
-   :alt: Transfer Learning
-
-.. image:: _static/pretrained_priors.png
-   :alt: Pre-trained Priors Example
+This approach allows you to incorporate domain knowledge or theoretical models into your Bayesian Neural Network, potentially leading to better generalization and more accurate uncertainty estimates, especially in cases where experimental data is limited.
 
-Examples
---------
+.. image:: https://github.com/user-attachments/assets/33f80877-4a5c-46d2-ba5d-ee540418e21b
+   :alt: pretrained_priors
 
-- `BNN-powered active learning example <https://github.com/ziatdinovmax/NeuroBayes/blob/main/examples/bnn_example1d.ipynb>`_
-- `PBNN-powered active learning example <https://github.com/ziatdinovmax/NeuroBayes/blob/main/examples/pbnn_example1d.ipynb>`_
-- `Heteroskedastic modeling example <https://github.com/ziatdinovmax/NeuroBayes/blob/main/examples/heteroskedastic.ipynb>`_
-- `Pre-trained priors example <https://github.com/ziatdinovmax/NeuroBayes/blob/main/examples/pretrained_priors.ipynb>`_
+See example `here <https://github.com/ziatdinovmax/NeuroBayes/blob/main/examples/pretrained_priors.ipynb>`_.