Tutorial 1: A First MI Estimate

Welcome to the NeuralMI library! This first tutorial covers the most basic use case: getting a single, quick estimate of mutual information (MI) between two variables.

Our goal is to introduce the main nmi.run function and show how to interpret its output. We’ll use a simple synthetic dataset where the true MI is known, allowing us to verify that our estimate is accurate.

Here we are estimating MI for a simple, common case: Independent and Identically Distributed (IID) data.

IID means that each data sample is independent of the others. There is no temporal order or sequence. This is a crucial concept, as it determines the correct way to split our data for model training and validation.

1. Imports

First, let’s import the necessary libraries. We’ll need numpy for data manipulation and, most importantly, our neural_mi library, which we import as nmi.

[1]:

import numpy as np
import neural_mi as nmi

2. Generating the Data

We will use the generate_correlated_gaussians function, which creates two multidimensional Gaussian variables, X and Y, with a precisely specified mutual information in bits.

The library’s ContinuousProcessor expects raw data in the format (n_channels, n_timepoints). Our data generator produces data of shape (n_samples, n_features), so we will simply transpose the matrices to match the expected format.

[2]:

# --- Dataset Parameters ---
n_samples = 5000
dim = 5
ground_truth_mi_bits = 2.0

# --- Generate Raw 2D Data ---
# This creates data of shape (n_samples, dim).
x_raw, y_raw = nmi.datasets.generate_correlated_gaussians(
    n_samples=n_samples,
    dim=dim,
    mi=ground_truth_mi_bits
)

# Transpose to the expected (n_channels, n_timepoints) format for the processor
x_raw_transposed = x_raw.T
y_raw_transposed = y_raw.T

print(f"Transposed X data shape: {x_raw_transposed.shape}")
print(f"Transposed Y data shape: {y_raw_transposed.shape}")

Transposed X data shape: torch.Size([5, 5000])
Transposed Y data shape: torch.Size([5, 5000])

3. The nmi.run Function: Modes of Analysis

The nmi.run function is the heart of the library. It can be configured to run in several ‘modes’, each designed for a different kind of analysis. For this tutorial, we will use the simplest mode, 'estimate', but it’s good to be aware of the others, which we will cover in future tutorials:

  • ``’estimate’``: A single, quick MI estimate. Perfect for getting a first look at your data.

  • ``’sweep’``: An exploratory sweep over a grid of hyperparameters (like window_size).

  • ``’rigorous’``: The full, bias-corrected MI estimation workflow for publication-ready results.

  • ``’dimensionality’``: Estimates the latent dimensionality of a single variable X.

  • ``’lag’``: Estimates MI across a range of time lags between X and Y.

4. Defining the Analysis Parameters

To run an estimation, we need to provide two sets of parameters:

  1. processor_params: These tell the data processor how to handle the raw data. Since each sample in our Gaussian dataset is independent, we use a window_size of 1. This tells the processor to treat each of the 5000 timepoints as a separate sample.

  2. base_params: These control the neural network model and the training process. We’ll define a simple model architecture and set the number of training epochs.

[3]:

# The processor will treat each of the 5000 columns as an independent sample.
processor_params = {'window_size': 1}

# Basic model and training parameters
base_params = {
    'n_epochs': 50, 'learning_rate': 5e-4, 'batch_size': 128,
    'patience': 10, 'embedding_dim': 16, 'hidden_dim': 64, 'n_layers': 2
}

5. Running the MI Estimation

Now we call the main nmi.run function. We specify mode='estimate' for a single, quick run and set random_seed=42 to ensure our result is reproducible.

[4]:

results = nmi.run(
    x_data=x_raw_transposed, y_data=y_raw_transposed,
    mode='estimate',
    processor_type='continuous',
    processor_params=processor_params,
    base_params=base_params,
    split_mode='random',  # For IID data
    output_units='bits', # Specify the output units
    random_seed=42       # For reproducibility
)

2025-10-20 00:02:18 - neural_mi - WARNING - `processor_type` is deprecated. Use `processor_type_x` and `processor_type_y` instead.
2025-10-20 00:02:18 - neural_mi - WARNING - `processor_params` is deprecated. Use `processor_params_x` and `processor_params_y` instead.
2025-10-20 00:02:18 - neural_mi - INFO - Starting parameter sweep sequentially (n_workers=1)...

2025-10-20 00:02:27 - neural_mi - INFO - Parameter sweep finished.

6. Interpreting the Results

The nmi.run function returns a special Results object that holds all the information from the analysis. For 'estimate' mode, the most important attribute is mi_estimate.

[5]:

estimated_mi_bits = results.mi_estimate

print(f"\n--- Results ---")
print(f"Ground Truth MI:  {ground_truth_mi_bits:.3f} bits")
print(f"Estimated MI:     {estimated_mi_bits:.3f} bits")
print(f"Estimation Error: {abs(estimated_mi_bits - ground_truth_mi_bits):.3f} bits")


--- Results ---
Ground Truth MI:  2.000 bits
Estimated MI:     1.944 bits
Estimation Error: 0.056 bits

7. Conclusion

Success! The estimated MI is very close to the ground truth value. You have successfully used the nmi.run function to get your first MI estimate.

In the next tutorial, we’ll dive deeper into how the library handles more complex, realistic neural data formats.