### Install NeuroBayes from source

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/installation.md

Use these commands to clone the repository and perform an editable installation.

```bash
git clone https://github.com/ziatdinovmax/NeuroBayes.git
cd NeuroBayes
pip install -e .
```

--------------------------------

### Clone and run NeuroBayes examples locally

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/examples/index.md

Commands to set up the local environment and launch the Jupyter notebook server for the examples directory.

```bash
git clone https://github.com/ziatdinovmax/NeuroBayes.git
cd NeuroBayes
pip install -e .
jupyter notebook examples/
```

--------------------------------

### Install GPyTorch and DeepChem

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/GPyTorch_esol.ipynb

Install the necessary libraries for Gaussian Process modeling and data processing.

```python
!pip install gpytorch
!pip install deepchem
```

--------------------------------

### Install NeuroBayes via pip

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/installation.md

Directly install the package from the GitHub repository using pip.

```bash
pip install git+https://github.com/ziatdinovmax/NeuroBayes.git
```

--------------------------------

### Install NeuroBayes via PyPI

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/README.md

Use this command to install the stable version of the library from the Python Package Index.

```bash
pip install neurobayes
```

--------------------------------

### Model Training Output

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/GPyTorch_esol.ipynb

Example output showing the training loss progression and final performance metrics.

```text
Iteration 100/1000 - Loss: 1.455
Iteration 200/1000 - Loss: 1.319
Iteration 300/1000 - Loss: 1.262
Iteration 400/1000 - Loss: 1.231
Iteration 500/1000 - Loss: 1.213
Iteration 600/1000 - Loss: 1.202
Iteration 700/1000 - Loss: 1.196
Iteration 800/1000 - Loss: 1.192
Iteration 900/1000 - Loss: 1.190
Iteration 1000/1000 - Loss: 1.189
Model Performance:
Test RMSE: 0.6887
Test R²: 0.8876
NLPD: 0.9575
Coverage probability: 0.9291
```

--------------------------------

### Fit PartialBNN Model

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/api/models/partial_bnn.md

Example usage of the fit method for training the model with specific configuration parameters.

```python
model.fit(
  X, y, num_warmup=1000, num_samples=1000,
  sgd_lr=1e-3, sgd_epochs=100,
  select_neurons_config={
- 'method' - 'variance',
- 'layer_names' - ['Dense0', 'Dense2'],
- 'num_pairs_per_layer' - 5
  }
  )
```

--------------------------------

### Apply Custom Prior to GP Model

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/api/utils/priors.md

Example of passing a custom uniform prior to a GP model during initialization and training.

```python
>>> model = gpax.ExactGP(input_dm, kernel, lengthscale_prior_dist=gpax.priors.uniform_dist(1, 3))
>>> model.fit(rng_key, X, y)
```

--------------------------------

### Get MCMC sampler

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/api/models/partial_bnn.md

Accesses the underlying MCMC sampler instance.

```python
@property
def mcmc()
```

--------------------------------

### Evaluate Model Performance

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/PartialBayesianTransformer_esol.ipynb

Computes RMSE and R^2 metrics on the test set and displays individual prediction examples.

```python
# Get predictions on test set
predictions = nn_model.predict(test_data['input_ids']).squeeze()

# Compute RMSE and R^2
test_rmse = np.sqrt(np.mean((test_data['solubility'] - predictions) ** 2))
test_r2 = r2_score(test_data['solubility'], predictions)

print("\nModel Performance:")
print(f"Test RMSE: {test_rmse:.4f}")
print(f"Test R²: {test_r2:.4f}")

# Show example predictions
print("\nExample Predictions:")
for i in range(min(5, len(X_test))):
    pred = float(nn_model.predict(test_data['input_ids'][i:i+1])[0, 0])  # Access the scalar value
    print(f"SMILES: {X_test[i]}")
    print(f"True Solubility: {y_test[i]:.2f}")
    print(f"Predicted Solubility: {pred:.2f}\n")
```

--------------------------------

### Integrate MLP with NeuroBayes Models

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/networks/mlp.md

Examples of using MLP architectures with BNN, Partial BNN, Heteroskedastic BNN, and Deep Kernel Learning models.

```python
mlp = nb.FlaxMLP(hidden_dims=[64, 32], target_dim=1)
model = nb.BNN(mlp)
```

```python
mlp = nb.FlaxMLP(hidden_dims=[64, 32], target_dim=1)
model = nb.PartialBNN(mlp, num_probabilistic_layers=1)
```

```python
mlp2head = nb.FlaxMLP2Head(hidden_dims=[64, 32], target_dim=1)
model = nb.HeteroskedasticBNN(mlp2head)
```

```python
mlp = nb.FlaxMLP(hidden_dims=[64, 32], target_dim=2)
model = nb.DKL(mlp, kernel=nb.kernels.RBFKernel)
```

--------------------------------

### Initialize and train a DeterministicNN

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/networks/detnn.md

Demonstrates creating a neural network architecture, initializing the DeterministicNN wrapper, and executing the training process.

```python
import neurobayes as nb
import jax.numpy as jnp

# Create a neural network architecture
mlp = nb.FlaxMLP(hidden_dims=[64, 32, 16], target_dim=1, activation='tanh')

# Create a deterministic neural network
det_nn = nb.DeterministicNN(
    architecture=mlp,
    input_shape=(10,),           # Input feature dimension
    loss='homoskedastic',        # 'homoskedastic', 'heteroskedastic', or 'classification'
    learning_rate=0.01,          # Learning rate for optimizer
    map=True,                    # Use maximum a posteriori estimation (regularization)
    sigma=1.0,                   # Prior sigma for MAP regularization
    collect_gradients=False      # Whether to collect gradients during training
)

# Train the model
det_nn.train(
    X_train,
    y_train,
    epochs=200,
    batch_size=32
)

# Make predictions
predictions = det_nn.predict(X_test)
```

--------------------------------

### Initialize and fit Gaussian Process models

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/models/gp.md

Demonstrates the basic workflow for both standard MCMC-based GP and faster Variational Inference GP models using Matern kernels.

```python
import neurobayes as nb

# MCMC-based GP
gp = nb.GP(kernel=nb.kernels.MaternKernel)
gp.fit(X, y, num_warmup=500, num_samples=1000)
mean, variance = gp.predict(X_test)

# Variational Inference GP (faster)
vi_gp = nb.VIGP(kernel=nb.kernels.MaternKernel)
vi_gp.fit(X, y, num_steps=1000, step_size=1e-3)
mean, variance = vi_gp.predict(X_test)
```

--------------------------------

### Get diagnostic results

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/api/models/partial_bnn.md

Retrieves diagnostic information if diagnostics were enabled during the fitting process.

```python
@property
def diagnostic_results()
```

--------------------------------

### Configure Stochastic Weight Averaging (SWA)

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/networks/detnn.md

Shows how to define an SWA configuration dictionary and pass it to the DeterministicNN constructor to improve generalization.

```python
# Configure SWA
swa_config = {
    'schedule': 'linear',    # Options: 'constant', 'linear', 'cyclic'
    'start_pct': 0.80,         # When to start collecting weights (fraction of total epochs)
    'swa_lr': 0.001            # Learning rate during SWA period
}

det_nn = nb.DeterministicNN(
    architecture=mlp,
    input_shape=(10,),
    learning_rate=0.01,
    swa_config=swa_config
)
```

--------------------------------

### Import dependencies

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/heteroskedastic_noisemodel.ipynb

Initializes the environment and imports necessary libraries.

```python
import sys
sys.path.append("..")

import neurobayes as nb

import numpy as np
import matplotlib.pyplot as plt

import jax.numpy as jnp
```

--------------------------------

### Initialize and fit a Bayesian neural network

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/bnn_simple_classification.ipynb

Sets up a BNN with a probabilistic architecture and performs MCMC sampling to fit the model.

```python
architecture = nb.FlaxMLP(
    hidden_dims=[8, 8, 8, 8],
    target_dim=2  # no classification flag here, since softmax will be applied at the BNN level
)

# Initialize BNN
bnn = nb.BNN(
    architecture,
    num_classes=2
)

# Fit the model
bnn.fit(
    X=X,
    y=y,
    num_warmup=1000,
    num_samples=1000,
)

fig, (ax1, ax2) = plot_bnn_predictions(X, y, bnn)
plt.show()

fig, ax = plot_bnn_realizations(X, y, bnn)
plt.show()
```

```python
architecture = nb.FlaxMLP(hidden_dims=[8, 8, 8, 8], target_dim=2)

# Initialize BNN
bnn = nb.BNN(
    architecture,
    num_classes=2
)

# Fit the model
bnn.fit(
    X=X,
    y=y,
    num_warmup=1000,
    num_samples=1000,
)

fig, (ax1, ax2) = plot_bnn_predictions(X, y, bnn)
plt.show()

fig, ax = plot_bnn_realizations(X, y, bnn)
plt.show()
```

--------------------------------

### Manual two-stage training workflow

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/models/partial_bnn.md

Use this approach to fine-tune a deterministic model before converting it to a Partial BNN for MCMC training.

```python
import neurobayes as nb

# 1. Train a deterministic model with extensive experimentation
mlp = nb.FlaxMLP(hidden_dims=[64, 32, 16, 8], target_dim=1, activation='silu')
det_model = nb.DeterministicNN(
    mlp, 
    input_shape=(X.shape[-1],),
    learning_rate=0.001,
    map=True,                # Use MAP training (with regularization)
    sigma=1.0               # Prior sigma for regularization
)

# Experiment with different training strategies
det_model.train(X, y, epochs=500, batch_size=64)

# Evaluate deterministic model performance
y_pred = det_model.predict(X_val)
print(f"Validation MSE: {(nb.utils.mse(y_pred, y_val)}")

# Try different architectures, hyperparameters, etc.
# ...

# When satisfied, get the final weights
pretrained_weights = det_model.get_params()

# 2. Create Partial BNN with these carefully tuned weights
pbn = nb.PartialBNN(
    mlp,
    deterministic_weights=pretrained_weights,
    probabilistic_layer_names=['Dense1']
)

# 3. Train with just MCMC 
# (automatically skips deterministic pre-training if pretrained weights are provided)
pbn.fit(X, y, 
        num_warmup=500, 
        num_samples=1000)
```

--------------------------------

### Get convolution and pooling operations

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/api/networks/convnet.md

Retrieves the appropriate convolution and pooling functions for 1D, 2D, or 3D input data.

```python
def get_conv_and_pool_ops(
        input_dim: int,
        kernel_size: Union[int, Tuple[int, ...]]) -> Tuple[Callable, Callable]
```

--------------------------------

### Train BNN with pre-trained priors

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/pretrained_priors.ipynb

Initializes and trains a BNN using the parameters from the deterministic model as informative priors.

```python
model = nb.BNN(
    net, pretrained_priors=detnn.state.params  # use trained weights to set priors for BNN
)
model.fit(X2, y2, num_warmup=1000, num_samples=1000, num_chains=1)
```

--------------------------------

### auto_priors(func, params_begin_with, dist_type='normal', loc=0.0, scale=1.0)

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/api/utils/priors.md

Generates a function that samples from normal or log-normal distributions for function parameters starting from a specified index.

```APIDOC
## auto_priors

### Description
Generates a function that, when invoked, samples from normal or log-normal distributions for each parameter of the given deterministic function, except the first one.

### Parameters
- **func** (Callable) - The deterministic function for which to set normal or log-normal priors.
- **params_begin_with** (int) - Parameters to account for start from this number.
- **dist_type** (str, optional) - Type of distribution ('normal' or 'log-normal'). Defaults to 'normal'.
- **loc** (float, optional) - Mean of the distribution. Defaults to 0.0.
- **scale** (float, optional) - Standard deviation of the distribution. Defaults to 1.0.

### Returns
- **Callable** - A function that returns a dictionary of sampled values from the specified distribution for each parameter of the original function.
```

--------------------------------

### Initialize and train a deterministic neural network

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/bnn_simple_classification.ipynb

Configures a FlaxMLP architecture and trains a DeterministicNN model with SWA configuration.

```python
architecture = nb.FlaxMLP(hidden_dims=[8, 8, 8, 8], target_dim=2, classification=True)
detnn = nb.DeterministicNN(
    architecture, input_shape=X.shape,
    loss='classification', learning_rate=5e-3, map=False,
    swa_config={'schedule': 'linear', 'start_pct': 0.75, 'swa_lr': 1e-4}
)

detnn.train(X, y, epochs=500)

plot_nn_predictions(X, y, detnn)
plt.show()
```

```python
architecture = nb.FlaxMLP(hidden_dims=[8, 8, 8, 8], target_dim=2, classification=True)
detnn = nb.DeterministicNN(
    architecture, input_shape=X.shape,
    loss='classification', learning_rate=5e-3, map=False,
    swa_config={'schedule': 'linear', 'start_pct': 0.75, 'swa_lr': 1e-4}
)

detnn.train(X, y, epochs=500)

plot_nn_predictions(X, y, detnn)
plt.show()
```

--------------------------------

### Initialize and Train a Gaussian Process Model

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/README.md

Configures a GP model using a Matern kernel and executes training and prediction using the standard NeuroBayes interface.

```python3
# Specify kernel
kernel = nb.kernels.MaternKernel
# Initialize GP model
model = nb.GP(kernel)
# Train the same way as BNN
model.fit(X_measured, y_measured, num_warmup=1000, num_samples=1000)
# Make a prediction the same way as with BNN
posterior_mean, posterior_var = model.predict(X_domain)
```

--------------------------------

### Configure and Train Partial BNN with Two Probabilistic Layers

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/partialBNN_simple_regression.ipynb

Initializes a Partial BNN with the last two hidden layers and the output layer as probabilistic.

```python
pbnn = nb.PartialBNN(detnn.model, detnn.get_params(), num_probabilistic_layers=2)

pbnn.fit(X, y, num_warmup=1000, num_samples=1000)
```

--------------------------------

### Initialize and Train a Deep Kernel Learning Model

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/README.md

Sets up a DKL model by combining a Flax MLP feature extractor with a Matern kernel.

```python3
# Set a number of latent dimensions
latent_dim = 2
# Initialize NN architecture for the feature extractor part of the DKL
architecture = nb.FlaxMLP(hidden_dims = [64, 32, 16, 8], target_dim=latent_dim)
# Specify kernel for the GP part of DKL
kernel = nb.kernels.MaternKernel

# Initialize DKL model
model = nb.DKL(net, kernel)
# Train and make a prediction the same way as with GP and BNN
model.fit(X_measured, y_measured, num_warmup=1000, num_samples=1000)
posterior_mean, posterior_var = model.predict(X_domain)
```

--------------------------------

### Configure and Train Partial BNN with One Probabilistic Layer

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/partialBNN_simple_regression.ipynb

Initializes a Partial BNN with the last hidden layer and output layer as probabilistic, then fits the model.

```python
pbnn = nb.PartialBNN(detnn.model, detnn.get_params(), num_probabilistic_layers=1)

pbnn.fit(X, y, num_warmup=1000, num_samples=1000)
```

--------------------------------

### Configure and Train Partial BNN with Named Probabilistic Layers

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/partialBNN_simple_regression.ipynb

Initializes a Partial BNN by specifying the names of the layers to be treated as probabilistic.

```python
pbnn = nb.PartialBNN(detnn.model, detnn.get_params(), probabilistic_layer_names=['Dense0', 'Dense4'])

pbnn.fit(X, y, num_warmup=1000, num_samples=1000)
```

--------------------------------

### Initialize and train a Fully Bayesian Neural Network

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/README.md

Uses the FlaxMLP architecture to create a BNN where all weights are probabilistic distributions.

```python
import neurobayes as nb

# Initialize NN architecture
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)
```

--------------------------------

### Generate and visualize data

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/partialBNN_simple_regression.ipynb

Creates the toy dataset and displays the noisy observations.

```python
# Generate data
X, y = make_data(num_samples=128)

# Visualize
plt.scatter(X, y, s=6, color = 'k')
plt.axis([-1.5, 1.5, -2, 2])
plt.show()
```

--------------------------------

### Initialize and Fit PBNN

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/pbnn_simple_classification.ipynb

Configures the PBNN with specific probabilistic layers and performs model fitting using HMC sampling.

```python
# Initialize PBNN
pbnn = nb.PartialBNN(
    detnn.model,
    detnn.state.params,
    probabilistic_layer_names=['Dense0', 'Dense4'], # first hidden and output layers
    num_classes=2
)

# Fit the model
pbnn.fit(
    X=X,
    y=y,
    num_warmup=1000,
    num_samples=1000,
)

fig, (ax1, ax2) = plot_bnn_predictions(X, y, pbnn)
plt.show()

fig, ax = plot_bnn_realizations(X, y, pbnn)
plt.show()
```

--------------------------------

### Initialize FlaxMLP

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/networks/mlp.md

Configure a standard MLP for regression or classification tasks.

```python
import neurobayes as nb

# Create an MLP for regression
mlp = nb.FlaxMLP(
    hidden_dims=[64, 32],     # Hidden layer dimensions
    target_dim=1,             # Output dimension (number of classes for classification)
    activation='tanh',        # Activation function ('tanh' or 'silu')
    hidden_dropout=0.1,       # Dropout rate for hidden layers
    classification=False      # Set to True for classification tasks
)
```

--------------------------------

### Initialize and train DKL and VIDKL models

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/models/dkl.md

Demonstrates defining a feature extractor using FlaxMLP and performing inference with both MCMC-based DKL and Variational Inference DKL.

```python
import neurobayes as nb
import jax.numpy as jnp

# Define neural network feature extractor
mlp = nb.FlaxMLP(hidden_dims=[64, 32, 16], target_dim=2, activation='tanh')

# MCMC-based DKL
dkl = nb.DKL(mlp, kernel=nb.kernels.RBFKernel)
dkl.fit(X_train, y_train, num_warmup=500, num_samples=1000)
mean, variance = dkl.predict(X_test)

# Variational Inference DKL (faster)
vi_dkl = nb.VIDKL(mlp, kernel=nb.kernels.RBFKernel)
vi_dkl.fit(X_train, y_train, num_steps=1000)
mean, variance = vi_dkl.predict(X_test)
```

--------------------------------

### Initialize and Train Deterministic Transformer

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/PartialBayesianTransformer_esol.ipynb

Sets up a FlaxTransformer architecture and trains it using the DeterministicNN wrapper.

```python
tokenizer = SMILESTokenizer()
model = nb.flax_nets.FlaxTransformer(
    vocab_size=tokenizer.vocab_size,
    d_model=16,
    nhead=4,
    num_layers=2,
    dim_feedforward=64)

nn_model = nb.DeterministicNN(
    architecture=model,
    input_shape=128,
    learning_rate=1e-3,
    map=False,
    swa_config={'schedule': 'constant', 'start_pct': 0.90},
    collect_gradients=False
)


nn_model.train(
    train_data['input_ids'],
    train_data['solubility'],
    epochs=50,
    batch_size=32
)
```

--------------------------------

### Initialize FlaxConvNet

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/networks/convnet.md

Configure a standard convolutional neural network for 2D data.

```python
import neurobayes as nb

# Create a ConvNet for 2D image data
cnn = nb.FlaxConvNet(
    input_dim=2,              # Dimensionality (1=1D, 2=2D, 3=3D)
    conv_layers=[32, 64],     # Filters in convolutional layers
    fc_layers=[128, 64],      # Hidden dimensions for fully connected layers
    target_dim=1,             # Output dimension (number of classes for classification)
    activation='tanh',        # Activation function
    kernel_size=3,            # Kernel size for convolutions
    conv_dropout=0.1,         # Dropout rate for conv layers
    hidden_dropout=0.1,       # Dropout rate for FC hidden layers
    classification=False      # Set to True for classification
)
```

--------------------------------

### Initialize and train a Partial BNN

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/pbnn_simple_classification.ipynb

Configures a FlaxMLP architecture and fits a PartialBNN model using a MAP prior for regularization.

```python
architecture = nb.FlaxMLP(
    hidden_dims=[8, 8, 8, 8],
    target_dim=2  # no classification flag here, since softmax will be applied at the BNN level
)

# Initialize PBNN
pbnn = nb.PartialBNN(
    architecture,
    probabilistic_layer_names=['Dense0', 'Dense4'],
    num_classes=2
)

# Fit the model
pbnn.fit(
    X=X,
    y=y,
    num_warmup=1000,
    num_samples=1000,
    map_sigma=nb.utils.calculate_sigma(X)
)

fig, (ax1, ax2) = plot_bnn_predictions(X, y, pbnn)
plt.show()

fig, ax = plot_bnn_realizations(X, y, pbnn)
plt.show()
```

--------------------------------

### Initialize a PartialBNN model

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/models/partial_bnn.md

Create a base neural network architecture and wrap it in a PartialBNN instance to designate the final layer as probabilistic.

```python
import neurobayes as nb

# Create neural network architecture
mlp = nb.FlaxMLP(
    hidden_dims=[64, 32, 16],  # Hidden layer dimensions
    target_dim=1,              # Output dimension (1 for regression)
    activation='tanh'          # Activation function
)

# Create PartialBNN model with only the last layer being Bayesian
pbn = nb.PartialBNN(mlp, num_probabilistic_layers=1)
```

--------------------------------

### Sample Kernel Parameters

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/api/models/gp.md

Samples parameters for the GP kernel given a specific dimension.

```python
def sample_kernel_params(kernel_dim: int) -> Dict[str, jnp.ndarray]
```

--------------------------------

### Initialize FlaxConvNet2Head

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/networks/convnet.md

Configure a two-headed convolutional neural network for heteroskedastic modeling.

```python
import neurobayes as nb

# Create a two-headed ConvNet for heteroskedastic modeling of 2D data
cnn2head = nb.FlaxConvNet2Head(
    input_dim=2,              # Dimensionality (1=1D, 2=2D, 3=3D)
    conv_layers=[32, 64],     # Filters in convolutional layers
    fc_layers=[128, 64],      # Hidden dimensions for fully connected layers
    target_dim=1,             # Output dimension for each head
    activation='tanh',        # Activation function
    kernel_size=3,            # Kernel size for convolutions
    conv_dropout=0.1,         # Dropout rate for conv layers
    hidden_dropout=0.1        # Dropout rate for FC hidden layers
)
```

--------------------------------

### Infer JAX Device

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/api/utils/utils.md

Returns a JAX device based on the specified preference, defaulting to the first available if none is provided.

```python
def infer_device(device_preference: str = None)
```

--------------------------------

### Active Learning Iteration Logs

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/pbnn_example1d.ipynb

Logs showing the sampling progress and posterior statistics for steps 36 through 40.

```text
Output:
step 36
```

```text
Output:
sample: 100%|██████████| 2000/2000 [00:04<00:00, 426.13it/s, 7 steps of size 5.19e-01. acc. prob=0.94] 
```

```text
Output:

                 mean       std    median      5.0%     95.0%     n_eff     r_hat
k_length[0]      1.64      0.47      1.58      0.93      2.32    283.70      1.02
    k_scale      0.67      0.46      0.55      0.19      1.13    301.95      1.01
      noise      0.03      0.01      0.03      0.01      0.04    546.85      1.00

step 37
```

```text
Output:
sample: 100%|██████████| 2000/2000 [00:04<00:00, 464.74it/s, 11 steps of size 5.40e-01. acc. prob=0.93]
```

```text
Output:

                 mean       std    median      5.0%     95.0%     n_eff     r_hat
k_length[0]      1.62      0.44      1.55      0.95      2.31    332.99      1.01
    k_scale      0.65      0.43      0.54      0.17      1.10    367.14      1.01
      noise      0.03      0.01      0.03      0.01      0.04    707.50      1.00

step 38
```

```text
Output:
sample: 100%|██████████| 2000/2000 [00:04<00:00, 423.32it/s, 7 steps of size 5.20e-01. acc. prob=0.93]
```

```text
Output:

                 mean       std    median      5.0%     95.0%     n_eff     r_hat
k_length[0]      1.59      0.43      1.55      0.92      2.22    407.06      1.00
    k_scale      0.64      0.41      0.54      0.18      1.08    301.37      1.00
      noise      0.03      0.01      0.03      0.01      0.04    691.63      1.00

step 39
```

```text
Output:
sample: 100%|██████████| 2000/2000 [00:04<00:00, 410.53it/s, 7 steps of size 4.52e-01. acc. prob=0.95] 
```

```text
Output:

                 mean       std    median      5.0%     95.0%     n_eff     r_hat
k_length[0]      1.58      0.42      1.53      0.91      2.17    477.74      1.01
    k_scale      0.63      0.40      0.53      0.19      1.09    422.35      1.01
      noise      0.03      0.01      0.02      0.01      0.04    627.06      1.00

step 40
```

```text
Output:
sample: 100%|██████████| 2000/2000 [00:04<00:00, 491.72it/s, 7 steps of size 4.37e-01. acc. prob=0.95] 
```

```text
Output:

                 mean       std    median      5.0%     95.0%     n_eff     r_hat
k_length[0]      1.60      0.42      1.54      0.93      2.18    445.22      1.00
    k_scale      0.65      0.41      0.53      0.19      1.15    383.11      1.00
      noise      0.02      0.01      0.02      0.01      0.04    679.26      1.00
```

--------------------------------

### Inspect Model Layer Configurations

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/PartialBayesianTransformer_esol.ipynb

Prints the architecture details of the model to identify layers for Bayesian conversion.

```python
nb.print_layer_configs(nn_model.model)
```

--------------------------------

### Initialize FlaxMLP2Head

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/networks/mlp.md

Configure a two-headed MLP for heteroskedastic modeling where heads predict mean and variance.

```python
import neurobayes as nb

# Create a two-headed MLP for heteroskedastic modeling
mlp2head = nb.FlaxMLP2Head(
    hidden_dims=[64, 32],     # Hidden layer dimensions
    target_dim=1,             # Output dimension for each head
    activation='tanh',        # Activation function
    hidden_dropout=0.1,       # Dropout rate for hidden layers
)
```

--------------------------------

### Visualize Comparative Model Metrics

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/PartialBayesianTransformer_esol.ipynb

Creates a side-by-side bar chart comparing RMSE, NLPD, and Coverage metrics across different model configurations and a Gaussian Process baseline.

```python
# GP results
test_rmse_gp = 0.6887
test_nlpd_gp = 0.9575
test_coverage_gp = 0.9291

# Labels for the bars
labels = ["Partially Bayesian\nToken Embedding",
          "Partially Bayesian\nFinal Layer",
          "Partially Bayesian\nAttention",
          "Gaussian\nProcess"]

# RMSE, NLPD, and Coverage values
rmse_values = [test_rmse_emb, test_rmse_fin, test_rmse_att, test_rmse_gp]
nlpd_values = [test_nlpd_emb, test_nlpd_fin, test_nlpd_att, test_nlpd_gp]
coverage_values = [test_coverage_emb, test_coverage_fin, test_coverage_att, test_coverage_gp]

# Colors for each model
colors = ['green', 'red', 'blue', 'purple']

# Create figure and axes with more width for the labels
plt.figure(figsize=(20, 6))
_, (ax0, ax1, ax2) = plt.subplots(1, 3, figsize=(20, 6))

# Function to create bars with error bars
def create_bar_plot(ax, values, title, ylabel, colors, labels, display_values=False):
    bars = ax.bar(np.arange(len(values)), values, color=colors)
    ax.set_xticks(np.arange(len(values)))
    ax.set_xticklabels(labels, rotation=45, ha='right')
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    if display_values:
        # Add value labels on top of bars
        for bar in bars:
            height = bar.get_height()
            ax.text(bar.get_x() + bar.get_width()/2., height,
                    f'{height:.3f}',
                    ha='center', va='bottom')
    return bars

# Create the three plots
create_bar_plot(ax0, rmse_values,
                "Root Mean Squared Error\n(lower is better)",
                "Root Mean Squared Error",
                colors, labels)

create_bar_plot(ax1, nlpd_values,
                "Negative Log Predictive Density\n(lower is better)",
                "Negative Log Predictive Density",
                colors, labels)

create_bar_plot(ax2, coverage_values,
                "Coverage Probability\n(closer to 0.95 is better)",
                "Coverage Probability",
                colors, labels)

# Set y-axis limits for NLPD and Coverage as in original
ax0.set_ylim(min(rmse_values) - 0.1, max(rmse_values) + 0.1)
ax1.set_ylim(min(nlpd_values) - 0.1, max(nlpd_values) + 0.1)
ax2.set_ylim(min(nlpd_values) - 0.1, 1.1)

# Adjust layout to prevent label overlap
plt.tight_layout()
plt.show()
```

--------------------------------

### Initialize and train a Partially Bayesian Neural Network

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/README.md

Configures a PBNN by specifying the number of probabilistic layers or explicitly naming the layers to be treated as Bayesian.

```python
# Number of probabilistic ('Bayesian') layers
num_probabilistic_layers = 2 # two last learnable layers + output layer

# 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 on unmeasured points or the full domain
posterior_mean, posterior_var = model.predict(X_domain)
```

```python
# Specify the names of probabilistic layers (output layer, 'Dense4', needs to be specified explicitly)
probabilistic_layer_names = ['Dense2', 'Dense3', 'Dense4']

# 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)
```

--------------------------------

### Train and Visualize Deterministic NN

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/pbnn_simple_classification.ipynb

Initializes a Flax-based MLP, trains it using a deterministic wrapper, and visualizes the resulting decision boundary.

```python
architecture = nb.FlaxMLP(hidden_dims=[8, 8, 8, 8], target_dim=2, classification=True)
detnn = nb.DeterministicNN(
    architecture, input_shape=X.shape,
    loss='classification', learning_rate=5e-3, map=False,
    swa_config={'schedule': 'linear', 'start_pct': 0.75, 'swa_lr': 1e-4}
)

detnn.train(X, y, epochs=500)

plot_nn_predictions(X, y, detnn)
plt.show()
```

--------------------------------

### Check GPU status

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/PartialBayesianTransformer_esol.ipynb

Verify GPU availability using the nvidia-smi command.

```bash
!nvidia-smi
```

--------------------------------

### DeterministicNN

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/api/networks/deterministic_nn.md

Initializes a DeterministicNN model with a specified architecture, loss type, and SWA configuration.

```APIDOC
## class DeterministicNN(architecture, input_shape, loss='homoskedastic', learning_rate, map=False, sigma=1.0, swa_config=None)

### Description
Initializes the neural network model with the provided architecture and training configurations.

### Parameters
- **architecture** (Flax model) - Required - The underlying Flax model architecture.
- **input_shape** (tuple) - Required - The shape of the input data (n_features or *dims, n_channels).
- **loss** (str) - Optional - Type of loss: 'homoskedastic' (default) or 'heteroskedastic'.
- **learning_rate** (float) - Required - Initial learning rate.
- **map** (bool) - Optional - Whether to use maximum a posteriori approximation.
- **sigma** (float) - Optional - Standard deviation for Gaussian prior.
- **swa_config** (dict) - Optional - Configuration dictionary for Stochastic Weight Averaging.
```

--------------------------------

### Configure and Train with Final Dense Layer

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/PartialBayesianTransformer_esol.ipynb

Applies the Partially Bayesian approach to the final dense layer of the architecture.

```python
probabilistic_layers = ['FinalDense1']

prob_indices = nb.select_bayesian_components(
    nn_model,
    layer_names=probabilistic_layers,
    method='variance',
    num_pairs_per_layer=100
)

pbnn_model = nb.PartialBayesianTransformer(
    nn_model.model,
    deterministic_weights=nn_model.get_params(),
    probabilistic_layer_names=probabilistic_layers,
    probabilistic_neurons=prob_indices
)

pbnn_model.fit(
    train_data["input_ids"],
    train_data['solubility'],
    num_warmup=1000,
    num_samples=1000,
)
```

--------------------------------

### Train Step Method

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/api/networks/deterministic_nn.md

JIT-compiled training step for the model.

```python
@partial(jax.jit, static_argnums=(0, ))
def train_step(state, inputs, targets)
```

--------------------------------

### Initialize and train a BNN model

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/models/bnn.md

Defines a small MLP architecture, wraps it in a BNN model, and performs MCMC sampling for training.

```python
import neurobayes as nb
import jax.numpy as jnp

# Create neural network architecture
mlp = nb.FlaxMLP(
    hidden_dims=[32, 16],  # Keep architecture small for BNNs
    target_dim=1,
    activation='tanh'
)

# Create BNN model
bnn = nb.BNN(mlp)

# Train the model (this may take considerable time)
bnn.fit(X, y, 
        num_warmup=500,
        num_samples=1000)

# Make predictions
mean, variance = bnn.predict(X_test)
```

--------------------------------

### Fit a Gaussian Process model

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/pretrained_priors.ipynb

Initializes and fits a Gaussian Process model using the specified kernel.

```python
model = nb.GP(kernel=nb.kernels.MaternKernel)
model.fit(X2, y2, num_warmup=1000, num_samples=1000, num_chains=1)
```

--------------------------------

### train_step

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/api/networks/deterministic_nn.md

Performs a single JIT-compiled training step.

```APIDOC
## train_step(state, inputs, targets)

### Description
Executes a JIT-compiled training step for the model.

### Parameters
- **state** (object) - Required - The current training state.
- **inputs** (array) - Required - Input data batch.
- **targets** (array) - Required - Target data batch.
```

--------------------------------

### MCMCDiagnostics.run_diagnostics

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/api/utils/diagnostics.md

Analyzes and prints MCMC diagnostics summary in one go.

```APIDOC
## run_diagnostics(mcmc_samples: Dict[str, jnp.ndarray]) -> Dict

### Description
Analyze and print MCMC diagnostics summary in one go.

### Parameters
- **mcmc_samples** (Dict[str, jnp.ndarray]) - Required - A dictionary containing MCMC samples where keys are layer names and values are JAX arrays.

### Returns
- **Dict** - A dictionary containing the diagnostic results.
```

--------------------------------

### Initialize BNN Model

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/api/models/bnn.md

Define the BNN class structure for regression or classification tasks.

```python
class BNN()
```

--------------------------------

### Train Partial BNN

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/partialBNN_simple_regression.ipynb

Initializes and trains a Partial BNN with a MAP prior to regularize the deterministic component.

```python
model = nb.PartialBNN(architecture, num_probabilistic_layers=2)

model.fit(
    X, y, sgd_epochs=1000, sgd_lr=5e-3,
    map_sigma=nb.utils.calculate_sigma(X),
    num_warmup=1000, num_samples=1000)
```

--------------------------------

### infer_device

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/api/utils/utils.md

Returns a JAX device based on the specified preference.

```APIDOC
## infer_device(device_preference: str = None)

### Description
Returns a JAX device based on the specified preference. Defaults to the first available device if no preference is given, or if the specified device type is not available.

### Parameters
- **device_preference** (str) - Optional - The preferred device type ('cpu' or 'gpu').

### Returns
- A JAX device.
```

--------------------------------

### Configure XLA environment

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/heteroskedastic_noisemodel.ipynb

Sets environment variables for JAX execution.

```python
import os
os.environ['XLA_FLAGS'] = '--xla_cpu_use_thunk_runtime=false' # see https://github.com/jax-ml/jax/discussions/23822
```

--------------------------------

### Configure Bayesian layers by name

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/models/partial_bnn.md

Explicitly define which layers should be treated as Bayesian by providing their names.

```python
# Make specific layers Bayesian
pbn = nb.PartialBNN(mlp, probabilistic_layer_names=['Dense1', 'Dense3'])
```

--------------------------------

### Import required libraries

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/GPyTorch_esol.ipynb

Imports GPyTorch, PyTorch, RDKit, and standard data science libraries for modeling.

```python
import gpytorch

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import r2_score
from sklearn.preprocessing import StandardScaler
from scipy.stats import norm


import torch

from rdkit import Chem
from rdkit.Chem import Descriptors, Crippen
```

--------------------------------

### Visualize Model Comparison

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/bnn_example1d.ipynb

Generates a side-by-side bar chart comparing MSE and NLPD metrics for GP and BNN models.

```python
# Labels for the bars
labels = ["GP", "BNN"]
# MSE and NLPD values
mse_values = [mse_gp, mse_bnn]
nlpd_values = [nlpd_gp, nlpd_bnn]
# Colors for each model
colors = ['green', 'red']
# Plot
_, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
ax1.bar(labels, mse_values, color=colors)
ax1.set_ylabel("MSE")
ax2.bar(labels, nlpd_values, color=colors)
ax2.set_ylabel("NLPD")
plt.show()
```

--------------------------------

### Calculate performance metrics

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/pretrained_priors.ipynb

Computes MSE and NLPD for the BNN predictions.

```python
mse1 = nb.utils.mse(posterior_mean, y2_true)
nlpd1 = nb.utils.nlpd(y2_true, posterior_mean, posterior_var)
```

--------------------------------

### put_on_device

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/api/utils/utils.md

Places multiple data items on the specified device.

```APIDOC
## put_on_device(device=None, *data_items)

### Description
Places multiple data items on the specified device.

### Parameters
- **device** (str) - Optional - The target device as a string (e.g., 'cpu', 'gpu'). If None, the default device is used.
- **data_items** (Any) - Required - Variable number of data items (such as JAX array or dictionary) to be placed on the device.

### Returns
- A tuple of the data items placed on the specified device.
```

--------------------------------

### Import dependencies

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/PartialBayesianTransformer_esol.ipynb

Import necessary libraries for data manipulation, modeling, and evaluation.

```python
import neurobayes as nb

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import r2_score
```

--------------------------------

### Implement Pre-trained Priors

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/README.md

Leverage theoretical or simulated data by pre-training a deterministic model and using its weights to center the priors of a Bayesian Neural Network.

```python3
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.train(X1, y1, epochs=5000, batch_size=None)
```

```python3
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)
```

```python3
posterior_mean, posterior_var = model.predict(X_test)
```

--------------------------------

### Load ESOL dataset

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/GPyTorch_esol.ipynb

Reads the downloaded CSV file into a pandas DataFrame.

```python
esol_df = pd.read_csv('delaney-processed.csv')
esol_df
```

--------------------------------

### Initialize and Train Partially Bayesian Transformer

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/PartialBayesianTransformer_esol.ipynb

Configures a model with probabilistic token embedding layers and trains it using MCMC sampling.

```python
# Specify probabilistic layers
probabilistic_layers = ['TokenEmbed']

# Select a subset of weights for the selected probabilistic layers
prob_indices = nb.select_bayesian_components(
    nn_model,
    layer_names=probabilistic_layers,
    method='variance',
    num_pairs_per_layer=100
)

# Initialize a partially Bayesian Transformer
pbnn_model = nb.PartialBayesianTransformer(
    nn_model.model,
    deterministic_weights=nn_model.get_params(),
    probabilistic_layer_names=probabilistic_layers,
    probabilistic_neurons=prob_indices
)

# Train
pbnn_model.fit(
    train_data["input_ids"],
    train_data['solubility'],
    num_warmup=1000,
    num_samples=1000,
)
```

--------------------------------

### Download ESOL dataset

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/examples/PartialBayesianTransformer_esol.ipynb

Fetch the ESOL dataset CSV file from the DeepChem repository.

```bash
!wget https://raw.githubusercontent.com/deepchem/deepchem/master/datasets/delaney-processed.csv
```

--------------------------------

### Integrate ConvNet with Partial BNN

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/networks/convnet.md

Wrap a standard ConvNet within a Partial BNN model.

```python
cnn = nb.FlaxConvNet(input_dim=2, conv_layers=[8, 32, 64], fc_layers=[128, 16], target_dim=1)
model = nb.PartialBNN(cnn, num_probabilistic_layers=1)
```

--------------------------------

### Place Data on Device

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/api/utils/utils.md

Moves multiple data items to a specified device, preserving their original structure.

```python
def put_on_device(device=None, *data_items)
```

--------------------------------

### Run Full Diagnostics

Source: https://github.com/ziatdinovmax/neurobayes/blob/main/docs/api/utils/diagnostics.md

Executes both analysis and summary printing in a single operation.

```python
def run_diagnostics(mcmc_samples: Dict[str, jnp.ndarray]) -> Dict
```