Information Theory Fundamentals for Deep Learning: Complete Mathematical Guide 2025

import numpy as np
from scipy.stats import entropy
import matplotlib.pyplot as plt

def shannon_entropy(probabilities, base=2):
    """
    Calculate Shannon entropy of a probability distribution.
    
    Args:
        probabilities: Array of probabilities (must sum to 1)
        base: Logarithm base (2 for bits, e for nats)
    
    Returns:
        Entropy value
    """
    # Remove zero probabilities to avoid log(0)
    p = np.array(probabilities)
    p = p[p > 0]
    
    if base == 2:
        return -np.sum(p * np.log2(p))
    else:
        return -np.sum(p * np.log(p))

def differential_entropy(samples, bins=100):
    """
    Estimate differential entropy for continuous distributions.
    
    Args:
        samples: Array of continuous samples
        bins: Number of histogram bins
    
    Returns:
        Estimated entropy in bits
    """
    # Histogram-based estimation
    hist, bin_edges = np.histogram(samples, bins=bins, density=True)
    bin_width = bin_edges[1] - bin_edges[0]
    
    # Remove zero bins
    hist = hist[hist > 0]
    
    # Differential entropy
    return -np.sum(hist * np.log2(hist)) * bin_width

# Example: Entropy of different distributions
print("=== Entropy Examples ===")

# Fair coin (maximum entropy for binary)
fair_coin = [0.5, 0.5]
print(f"Fair coin entropy: {shannon_entropy(fair_coin):.4f} bits")

# Biased coin
biased_coin = [0.9, 0.1]
print(f"Biased coin (90-10) entropy: {shannon_entropy(biased_coin):.4f} bits")

# Deterministic (no uncertainty)
deterministic = [1.0, 0.0]
print(f"Deterministic entropy: {shannon_entropy(deterministic):.4f} bits")

# 6-sided fair die
fair_die = [1/6] * 6
print(f"Fair die entropy: {shannon_entropy(fair_die):.4f} bits")

=== Entropy Examples ===
Fair coin entropy: 1.0000 bits
Biased coin (90-10) entropy: 0.4690 bits
Deterministic entropy: 0.0000 bits
Fair die entropy: 2.5850 bits

def joint_entropy(joint_prob_matrix, base=2):
    """
    Calculate joint entropy H(X, Y).
    
    Args:
        joint_prob_matrix: 2D array of joint probabilities
        base: Logarithm base
    
    Returns:
        Joint entropy value
    """
    p = joint_prob_matrix.flatten()
    p = p[p > 0]
    
    if base == 2:
        return -np.sum(p * np.log2(p))
    else:
        return -np.sum(p * np.log(p))

def conditional_entropy(joint_prob_matrix, base=2):
    """
    Calculate conditional entropy H(X|Y).
    Uses chain rule: H(X|Y) = H(X,Y) - H(Y)
    
    Args:
        joint_prob_matrix: 2D array where rows are X, columns are Y
        base: Logarithm base
    
    Returns:
        Conditional entropy H(X|Y)
    """
    # H(X, Y)
    H_joint = joint_entropy(joint_prob_matrix, base)
    
    # H(Y) - marginal entropy of Y
    p_y = joint_prob_matrix.sum(axis=0)  # Sum over X
    H_y = shannon_entropy(p_y, base)
    
    return H_joint - H_y

# Example: XOR relationship
# X, Y are inputs, Z = X XOR Y
joint_prob = np.array([
    [0.25, 0.0],   # X=0, Y=0 or Y=1
    [0.0, 0.25],   # X=1, Y=0 or Y=1
    [0.0, 0.25],   # etc.
    [0.25, 0.0]
])

print(f"Joint entropy H(X,Y): {joint_entropy(joint_prob):.4f} bits")

import numpy as np
from sklearn.feature_selection import mutual_info_classif
from sklearn.neighbors import KernelDensity
import torch
import torch.nn as nn

def mutual_information_discrete(joint_prob_matrix):
    """
    Calculate mutual information I(X; Y) for discrete variables.
    
    Args:
        joint_prob_matrix: 2D array of joint probabilities p(x, y)
    
    Returns:
        Mutual information in bits
    """
    # Marginals
    p_x = joint_prob_matrix.sum(axis=1, keepdims=True)
    p_y = joint_prob_matrix.sum(axis=0, keepdims=True)
    
    # Product of marginals
    p_independent = p_x @ p_y
    
    # Mutual information
    mi = 0
    for i in range(joint_prob_matrix.shape[0]):
        for j in range(joint_prob_matrix.shape[1]):
            if joint_prob_matrix[i, j] > 0 and p_independent[i, j] > 0:
                mi += joint_prob_matrix[i, j] * np.log2(
                    joint_prob_matrix[i, j] / p_independent[i, j]
                )
    
    return mi

def kraskov_mi_estimate(X, Y, k=3):
    """
    Kraskov's KNN-based mutual information estimator.
    More accurate for continuous variables.
    
    Args:
        X: First variable samples (n_samples, n_features)
        Y: Second variable samples (n_samples, n_features)
        k: Number of nearest neighbors
    
    Returns:
        Estimated mutual information in nats
    """
    from scipy.special import digamma
    from sklearn.neighbors import NearestNeighbors
    
    n = len(X)
    
    # Reshape if 1D
    if X.ndim == 1:
        X = X.reshape(-1, 1)
    if Y.ndim == 1:
        Y = Y.reshape(-1, 1)
    
    # Joint space
    XY = np.hstack([X, Y])
    
    # Find k-th neighbor distances in joint space
    nn_xy = NearestNeighbors(n_neighbors=k+1, metric='chebyshev')
    nn_xy.fit(XY)
    distances_xy, _ = nn_xy.kneighbors(XY)
    eps = distances_xy[:, k]  # k-th neighbor distance
    
    # Count neighbors within eps in marginal spaces
    nn_x = NearestNeighbors(metric='chebyshev')
    nn_x.fit(X)
    
    nn_y = NearestNeighbors(metric='chebyshev')
    nn_y.fit(Y)
    
    n_x = np.array([len(nn_x.radius_neighbors([x], radius=e, return_distance=False)[0]) - 1 
                   for x, e in zip(X, eps)])
    n_y = np.array([len(nn_y.radius_neighbors([y], radius=e, return_distance=False)[0]) - 1 
                   for y, e in zip(Y, eps)])
    
    # Kraskov estimator
    mi = digamma(k) - np.mean(digamma(n_x + 1) + digamma(n_y + 1)) + digamma(n)
    
    return max(0, mi)  # MI is non-negative

# Example: MI between correlated variables
np.random.seed(42)
n_samples = 1000

# Highly correlated
x = np.random.randn(n_samples)
y = x + 0.1 * np.random.randn(n_samples)  # Y ≈ X
mi_high = kraskov_mi_estimate(x, y, k=5)
print(f"High correlation MI estimate: {mi_high:.4f} nats")

# Moderate correlation
y_moderate = x + np.random.randn(n_samples)
mi_moderate = kraskov_mi_estimate(x, y_moderate, k=5)
print(f"Moderate correlation MI estimate: {mi_moderate:.4f} nats")

# Independent
y_independent = np.random.randn(n_samples)
mi_independent = kraskov_mi_estimate(x, y_independent, k=5)
print(f"Independent MI estimate: {mi_independent:.4f} nats")

class InformationPlaneTracker:
    """
    Track information plane dynamics during neural network training.
    
    Computes I(X; T) and I(T; Y) for each layer T.
    """
    
    def __init__(self, model, layer_names):
        """
        Args:
            model: PyTorch neural network
            layer_names: List of layer names to track
        """
        self.model = model
        self.layer_names = layer_names
        self.activations = {}
        self.hooks = []
        
        # Register forward hooks to capture activations
        self._register_hooks()
    
    def _register_hooks(self):
        """Register forward hooks to capture layer activations."""
        for name, module in self.model.named_modules():
            if name in self.layer_names:
                hook = module.register_forward_hook(
                    lambda m, inp, out, name=name: self._save_activation(name, out)
                )
                self.hooks.append(hook)
    
    def _save_activation(self, name, output):
        """Save activation output."""
        self.activations[name] = output.detach().cpu().numpy()
    
    def compute_information(self, X, Y, n_bins=30):
        """
        Compute information plane coordinates for each layer.
        
        Args:
            X: Input data (numpy array)
            Y: Labels (numpy array)
            n_bins: Number of bins for discretization
        
        Returns:
            Dictionary with I(X;T) and I(T;Y) for each layer
        """
        # Forward pass to get activations
        self.model.eval()
        with torch.no_grad():
            _ = self.model(torch.tensor(X, dtype=torch.float32))
        
        results = {}
        
        for layer_name in self.layer_names:
            T = self.activations[layer_name]
            
            # Flatten activations if needed
            if T.ndim > 2:
                T = T.reshape(T.shape[0], -1)
            
            # Discretize activations for MI estimation
            T_discrete = self._discretize(T, n_bins)
            X_discrete = self._discretize(X, n_bins) if X.ndim > 1 else X
            
            # Estimate mutual information
            I_XT = self._estimate_mi(X_discrete, T_discrete)
            I_TY = self._estimate_mi(T_discrete, Y)
            
            results[layer_name] = {'I_XT': I_XT, 'I_TY': I_TY}
        
        return results
    
    def _discretize(self, data, n_bins):
        """Discretize continuous data into bins."""
        if data.ndim == 1:
            return np.digitize(data, np.percentile(data, np.linspace(0, 100, n_bins)))
        else:
            # For high-dimensional data, use PCA to reduce dimensions first
            from sklearn.decomposition import PCA
            pca = PCA(n_components=min(10, data.shape[1]))
            data_reduced = pca.fit_transform(data)
            return np.digitize(data_reduced[:, 0], 
                              np.percentile(data_reduced[:, 0], np.linspace(0, 100, n_bins)))
    
    def _estimate_mi(self, X, Y):
        """Estimate mutual information using histogram method."""
        # Joint histogram
        c_xy = np.histogram2d(X.flatten(), Y.flatten(), bins=30)[0]
        c_xy = c_xy / c_xy.sum()  # Normalize to joint probability
        
        return mutual_information_discrete(c_xy + 1e-10)
    
    def cleanup(self):
        """Remove all hooks."""
        for hook in self.hooks:
            hook.remove()

# Example usage
class SimpleNet(nn.Module):
    def __init__(self, input_dim, hidden_dims, output_dim):
        super().__init__()
        layers = []
        prev_dim = input_dim
        
        for i, hidden_dim in enumerate(hidden_dims):
            layers.append(nn.Linear(prev_dim, hidden_dim))
            layers.append(nn.ReLU())
            prev_dim = hidden_dim
        
        layers.append(nn.Linear(prev_dim, output_dim))
        self.layers = nn.Sequential(*layers)
    
    def forward(self, x):
        return self.layers(x)

# This demonstrates the concept - full implementation requires training loop
print("Information Plane Tracker initialized for deep learning analysis")

def gaussian_channel_capacity(signal_power, noise_variance):
    """
    Calculate capacity of Gaussian channel.
    
    Args:
        signal_power: Signal power P
        noise_variance: Noise variance σ²
    
    Returns:
        Channel capacity in bits per symbol
    """
    snr = signal_power / noise_variance
    return 0.5 * np.log2(1 + snr)

def layer_information_capacity(weights, noise_std=0.1):
    """
    Estimate information capacity of a neural network layer.
    
    Treats the layer as a Gaussian channel where:
    - Signal power ≈ variance of weighted activations
    - Noise power = activation noise variance
    
    Args:
        weights: Layer weight matrix (numpy array)
        noise_std: Standard deviation of activation noise
    
    Returns:
        Estimated capacity per input dimension
    """
    # Singular value decomposition
    u, s, vh = np.linalg.svd(weights, full_matrices=False)
    
    # Each singular value represents a "channel"
    # Water-filling solution for optimal power allocation
    signal_powers = s ** 2
    noise_power = noise_std ** 2
    
    # Total capacity (sum over all channels)
    capacities = []
    for power in signal_powers:
        if power > noise_power:  # Only count channels above noise floor
            cap = 0.5 * np.log2(1 + power / noise_power)
            capacities.append(cap)
    
    return np.sum(capacities), capacities

# Example: Analyze different layer configurations
print("=== Layer Information Capacity Analysis ===\n")

# Wide layer (more channels)
wide_weights = np.random.randn(64, 256) * 0.5
cap_wide, channels_wide = layer_information_capacity(wide_weights)
print(f"Wide layer (256 -> 64):")
print(f"  Total capacity: {cap_wide:.2f} bits")
print(f"  Active channels: {len(channels_wide)}")

# Narrow layer (bottleneck)
narrow_weights = np.random.randn(64, 32) * 0.5
cap_narrow, channels_narrow = layer_information_capacity(narrow_weights)
print(f"\nNarrow layer (32 -> 64):")
print(f"  Total capacity: {cap_narrow:.2f} bits")
print(f"  Active channels: {len(channels_narrow)}")

# The narrow layer has lower capacity - acts as bottleneck
print(f"\nCapacity ratio: {cap_narrow/cap_wide:.2%}")

def verify_data_processing_inequality(X, Y, Z, k=5):
    """
    Verify the data processing inequality for a Markov chain X -> Y -> Z.
    
    Args:
        X, Y, Z: Random variable samples (Markov chain)
        k: Number of neighbors for MI estimation
    
    Returns:
        Dictionary with MI values and verification status
    """
    I_XY = kraskov_mi_estimate(X, Y, k)
    I_XZ = kraskov_mi_estimate(X, Z, k)
    I_YZ = kraskov_mi_estimate(Y, Z, k)
    
    # DPI: I(X;Z) <= min(I(X;Y), I(Y;Z))
    bound = min(I_XY, I_YZ)
    
    return {
        'I(X;Y)': I_XY,
        'I(Y;Z)': I_YZ,
        'I(X;Z)': I_XZ,
        'DPI bound': bound,
        'DPI satisfied': I_XZ <= bound + 0.1  # Small tolerance for estimation error
    }

# Simulate a Markov chain (like neural network layers)
np.random.seed(42)
n = 1000

X = np.random.randn(n)
Y = X + 0.3 * np.random.randn(n)  # First transformation (add noise)
Z = Y + 0.3 * np.random.randn(n)  # Second transformation (more noise)

result = verify_data_processing_inequality(X, Y, Z)
print("=== Data Processing Inequality Verification ===")
for key, value in result.items():
    if isinstance(value, float):
        print(f"{key}: {value:.4f}")
    else:
        print(f"{key}: {value}")

import torch
import torch.nn as nn
import torch.nn.functional as F

class RateDistortionVAE(nn.Module):
    """
    Variational Autoencoder with explicit rate-distortion tradeoff.
    
    The loss function is: L = D + β * R
    - D: Distortion (reconstruction error)
    - R: Rate (KL divergence from prior)
    - β: Lagrange multiplier controlling tradeoff
    """
    
    def __init__(self, input_dim, latent_dim, hidden_dims=[256, 128], beta=1.0):
        super().__init__()
        
        self.input_dim = input_dim
        self.latent_dim = latent_dim
        self.beta = beta
        
        # Encoder
        encoder_layers = []
        prev_dim = input_dim
        for h_dim in hidden_dims:
            encoder_layers.extend([
                nn.Linear(prev_dim, h_dim),
                nn.BatchNorm1d(h_dim),
                nn.ReLU()
            ])
            prev_dim = h_dim
        
        self.encoder = nn.Sequential(*encoder_layers)
        self.fc_mu = nn.Linear(prev_dim, latent_dim)
        self.fc_logvar = nn.Linear(prev_dim, latent_dim)
        
        # Decoder
        decoder_layers = []
        hidden_dims_reversed = hidden_dims[::-1]
        prev_dim = latent_dim
        for h_dim in hidden_dims_reversed:
            decoder_layers.extend([
                nn.Linear(prev_dim, h_dim),
                nn.BatchNorm1d(h_dim),
                nn.ReLU()
            ])
            prev_dim = h_dim
        
        decoder_layers.append(nn.Linear(prev_dim, input_dim))
        self.decoder = nn.Sequential(*decoder_layers)
    
    def encode(self, x):
        """Encode input to latent distribution parameters."""
        h = self.encoder(x)
        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)
        return mu, logvar
    
    def reparameterize(self, mu, logvar):
        """Reparameterization trick for backpropagation through sampling."""
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std
    
    def decode(self, z):
        """Decode latent representation to reconstruction."""
        return self.decoder(z)
    
    def forward(self, x):
        """Forward pass with encoding, sampling, and decoding."""
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        x_recon = self.decode(z)
        return x_recon, mu, logvar, z
    
    def loss_function(self, x, x_recon, mu, logvar):
        """
        Compute rate-distortion loss.
        
        Returns:
            total_loss: β * rate + distortion
            distortion: MSE reconstruction loss
            rate: KL divergence from N(0, I) prior
        """
        # Distortion: Reconstruction error (MSE)
        distortion = F.mse_loss(x_recon, x, reduction='sum') / x.size(0)
        
        # Rate: KL divergence KL(q(z|x) || p(z))
        # Closed form for Gaussian: -0.5 * sum(1 + log(σ²) - μ² - σ²)
        rate = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) / x.size(0)
        
        # Total loss with Lagrange multiplier β
        total_loss = distortion + self.beta * rate
        
        return total_loss, distortion, rate
    
    def compute_rate_distortion(self, dataloader):
        """
        Compute average rate and distortion over a dataset.
        
        Returns:
            avg_rate: Average bits per sample
            avg_distortion: Average MSE
        """
        self.eval()
        total_rate = 0
        total_distortion = 0
        n_samples = 0
        
        with torch.no_grad():
            for batch in dataloader:
                x = batch[0] if isinstance(batch, (list, tuple)) else batch
                x_recon, mu, logvar, _ = self(x)
                
                _, distortion, rate = self.loss_function(x, x_recon, mu, logvar)
                
                total_rate += rate.item() * x.size(0)
                total_distortion += distortion.item() * x.size(0)
                n_samples += x.size(0)
        
        # Convert rate from nats to bits
        avg_rate = (total_rate / n_samples) / np.log(2)
        avg_distortion = total_distortion / n_samples
        
        return avg_rate, avg_distortion

def plot_rate_distortion_curve(betas, input_dim, latent_dim):
    """
    Plot empirical rate-distortion curve by varying β.
    
    Higher β = lower rate (more compression), higher distortion
    Lower β = higher rate (less compression), lower distortion
    """
    rates = []
    distortions = []
    
    # Generate synthetic data
    data = torch.randn(1000, input_dim)
    
    for beta in betas:
        model = RateDistortionVAE(input_dim, latent_dim, beta=beta)
        
        # Quick training (for demonstration)
        optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
        
        for _ in range(100):  # Short training
            model.train()
            x_recon, mu, logvar, z = model(data)
            loss, _, _ = model.loss_function(data, x_recon, mu, logvar)
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        
        # Compute rate and distortion
        model.eval()
        with torch.no_grad():
            x_recon, mu, logvar, _ = model(data)
            _, distortion, rate = model.loss_function(data, x_recon, mu, logvar)
            
            rates.append(rate.item() / np.log(2))  # Convert to bits
            distortions.append(distortion.item())
    
    return rates, distortions

print("Rate-Distortion VAE model defined for neural compression analysis")

def cross_entropy(p_true, q_pred, epsilon=1e-15):
    """
    Compute cross-entropy H(p, q).
    
    Args:
        p_true: True distribution (ground truth)
        q_pred: Predicted distribution (model output)
        epsilon: Small constant for numerical stability
    
    Returns:
        Cross-entropy value
    """
    # Clip predictions to avoid log(0)
    q_pred = np.clip(q_pred, epsilon, 1 - epsilon)
    
    return -np.sum(p_true * np.log(q_pred))

def kl_divergence(p, q, epsilon=1e-15):
    """
    Compute KL divergence D_KL(p || q).
    
    Measures how much q differs from p.
    
    Args:
        p: True distribution
        q: Approximate distribution
        epsilon: Numerical stability constant
    
    Returns:
        KL divergence (non-negative)
    """
    p = np.clip(p, epsilon, 1)
    q = np.clip(q, epsilon, 1)
    
    return np.sum(p * np.log(p / q))

def cross_entropy_decomposition(p_true, q_pred):
    """
    Show the decomposition: H(p,q) = H(p) + D_KL(p||q)
    
    Returns:
        Dictionary with all components
    """
    H_p = shannon_entropy(p_true, base=np.e)  # Use natural log for consistency
    H_pq = cross_entropy(p_true, q_pred)
    D_KL = kl_divergence(p_true, q_pred)
    
    return {
        'H(p)': H_p,
        'H(p,q)': H_pq,
        'D_KL(p||q)': D_KL,
        'H(p) + D_KL': H_p + D_KL,
        'Decomposition holds': np.isclose(H_pq, H_p + D_KL)
    }

# Example: Classification loss
print("=== Cross-Entropy Decomposition ===\n")

# True distribution (one-hot for class 2)
p_true = np.array([0, 0, 1, 0, 0])

# Good prediction (confident and correct)
q_good = np.array([0.01, 0.02, 0.9, 0.05, 0.02])

# Bad prediction (incorrect)
q_bad = np.array([0.5, 0.3, 0.1, 0.05, 0.05])

# Uncertain prediction
q_uncertain = np.array([0.2, 0.2, 0.2, 0.2, 0.2])

print("True distribution (class 2):", p_true)
print()

for name, q in [('Good prediction', q_good), ('Bad prediction', q_bad), ('Uncertain', q_uncertain)]:
    print(f"{name}: {q}")
    result = cross_entropy_decomposition(p_true, q)
    print(f"  Cross-entropy: {result['H(p,q)']:.4f}")
    print(f"  KL divergence: {result['D_KL(p||q)']:.4f}")
    print()

def vae_elbo_loss(x, x_recon, mu, logvar, beta=1.0):
    """
    Evidence Lower Bound (ELBO) loss for VAE.
    
    ELBO = E_q[log p(x|z)] - β * D_KL(q(z|x) || p(z))
    
    Maximizing ELBO minimizes KL(q(z|x) || p(z|x))
    
    Args:
        x: Original input
        x_recon: Reconstructed input
        mu: Mean of approximate posterior q(z|x)
        logvar: Log variance of approximate posterior
        beta: Weight on KL term (β-VAE)
    
    Returns:
        loss: Negative ELBO
        reconstruction_term: E_q[log p(x|z)]
        kl_term: D_KL(q(z|x) || p(z))
    """
    # Reconstruction term (assuming Gaussian likelihood)
    reconstruction_loss = F.mse_loss(x_recon, x, reduction='sum')
    
    # KL divergence: KL(N(mu, sigma) || N(0, I))
    # = -0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
    kl_divergence = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    
    # Negative ELBO (minimize this)
    loss = reconstruction_loss + beta * kl_divergence
    
    return loss, reconstruction_loss.item(), kl_divergence.item()

class BayesianLinear(nn.Module):
    """
    Bayesian linear layer with KL regularization.
    
    Learns posterior distribution over weights instead of point estimates.
    """
    
    def __init__(self, in_features, out_features, prior_std=1.0):
        super().__init__()
        
        self.in_features = in_features
        self.out_features = out_features
        self.prior_std = prior_std
        
        # Weight posterior parameters (mean and log variance)
        self.weight_mu = nn.Parameter(torch.zeros(out_features, in_features))
        self.weight_logvar = nn.Parameter(torch.zeros(out_features, in_features) - 5)
        
        # Bias posterior parameters
        self.bias_mu = nn.Parameter(torch.zeros(out_features))
        self.bias_logvar = nn.Parameter(torch.zeros(out_features) - 5)
        
        # Initialize
        nn.init.kaiming_normal_(self.weight_mu)
    
    def forward(self, x, sample=True):
        """Forward pass with optional weight sampling."""
        if sample and self.training:
            # Sample weights from posterior
            weight_std = torch.exp(0.5 * self.weight_logvar)
            weight = self.weight_mu + weight_std * torch.randn_like(weight_std)
            
            bias_std = torch.exp(0.5 * self.bias_logvar)
            bias = self.bias_mu + bias_std * torch.randn_like(bias_std)
        else:
            # Use mean
            weight = self.weight_mu
            bias = self.bias_mu
        
        return F.linear(x, weight, bias)
    
    def kl_divergence(self):
        """
        Compute KL divergence from prior.
        
        KL(q(w) || p(w)) where:
        - q(w) = N(mu, sigma^2) is posterior
        - p(w) = N(0, prior_std^2) is prior
        """
        # Weight KL
        prior_var = self.prior_std ** 2
        
        weight_var = torch.exp(self.weight_logvar)
        kl_weight = 0.5 * torch.sum(
            (self.weight_mu ** 2 + weight_var) / prior_var - 1 
            + np.log(prior_var) - self.weight_logvar
        )
        
        bias_var = torch.exp(self.bias_logvar)
        kl_bias = 0.5 * torch.sum(
            (self.bias_mu ** 2 + bias_var) / prior_var - 1 
            + np.log(prior_var) - self.bias_logvar
        )
        
        return kl_weight + kl_bias

print("Bayesian linear layer with KL divergence regularization defined")

from sklearn.datasets import make_classification
from sklearn.feature_selection import mutual_info_classif

def mi_feature_selection(X, y, k=10):
    """
    Select top-k features using mutual information criterion.
    
    I(X_i; Y) measures how informative feature X_i is about label Y.
    
    Args:
        X: Feature matrix (n_samples, n_features)
        y: Labels
        k: Number of features to select
    
    Returns:
        selected_features: Indices of top-k features
        mi_scores: Mutual information scores for all features
    """
    # Compute MI between each feature and target
    mi_scores = mutual_info_classif(X, y, random_state=42)
    
    # Select top-k
    selected_features = np.argsort(mi_scores)[-k:][::-1]
    
    return selected_features, mi_scores

# Example
np.random.seed(42)
X, y = make_classification(
    n_samples=1000, 
    n_features=20, 
    n_informative=5,
    n_redundant=5,
    n_classes=2,
    random_state=42
)

selected_idx, scores = mi_feature_selection(X, y, k=5)
print("=== Mutual Information Feature Selection ===")
print(f"Top 5 features by MI score: {selected_idx}")
print(f"Their MI scores: {scores[selected_idx]}")

class InformationBottleneckNet(nn.Module):
    """
    Neural network with explicit information bottleneck.
    
    Uses noise injection to control information flow.
    """
    
    def __init__(self, input_dim, bottleneck_dim, output_dim, noise_std=0.1):
        super().__init__()
        
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, bottleneck_dim * 2)  # Mean and log-variance
        )
        
        self.decoder = nn.Sequential(
            nn.Linear(bottleneck_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 128),
            nn.ReLU(),
            nn.Linear(128, output_dim)
        )
        
        self.bottleneck_dim = bottleneck_dim
        self.noise_std = noise_std
    
    def forward(self, x, add_noise=True):
        """
        Forward pass through information bottleneck.
        
        Args:
            x: Input tensor
            add_noise: Whether to add noise at bottleneck (training only)
        
        Returns:
            output: Model predictions
            bottleneck: Bottleneck representation (for analysis)
            kl: KL divergence from standard normal
        """
        # Encode to bottleneck
        h = self.encoder(x)
        mu = h[:, :self.bottleneck_dim]
        logvar = h[:, self.bottleneck_dim:]
        
        # Reparameterization trick
        if add_noise and self.training:
            std = torch.exp(0.5 * logvar)
            bottleneck = mu + std * torch.randn_like(std)
        else:
            bottleneck = mu
        
        # Decode
        output = self.decoder(bottleneck)
        
        # KL divergence from N(0,1)
        kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1)
        
        return output, bottleneck, kl.mean()
    
    def information_bottleneck_loss(self, x, y, beta=0.01):
        """
        Information bottleneck loss.
        
        L = H(Y|T) + β * I(X;T)
          ≈ CrossEntropy(y, y_pred) + β * KL(q(t|x) || p(t))
        
        Args:
            x: Input
            y: Target labels
            beta: Information constraint weight
        
        Returns:
            total_loss, classification_loss, ib_regularization
        """
        output, _, kl = self(x)
        
        # Classification loss (proxy for H(Y|T))
        class_loss = F.cross_entropy(output, y)
        
        # Information bottleneck regularization (proxy for I(X;T))
        ib_reg = kl
        
        total_loss = class_loss + beta * ib_reg
        
        return total_loss, class_loss.item(), ib_reg.item()

print("Information Bottleneck Network defined for compression research")

def estimate_model_capacity(model, input_dim, n_samples=10000):
    """
    Estimate the effective information capacity of a model.
    
    Uses random input-output pairs to estimate mutual information.
    
    Args:
        model: PyTorch model
        input_dim: Input dimensionality
        n_samples: Number of samples for estimation
    
    Returns:
        capacity_estimate: Estimated bits of information
    """
    model.eval()
    
    # Generate random inputs
    X = torch.randn(n_samples, input_dim)
    
    with torch.no_grad():
        Y = model(X)
        if isinstance(Y, tuple):
            Y = Y[0]
    
    X_np = X.numpy()
    Y_np = Y.numpy()
    
    # Estimate I(X; Y) using KDE-based method
    # This gives an upper bound on model capacity
    
    # Flatten outputs if needed
    if Y_np.ndim > 2:
        Y_np = Y_np.reshape(Y_np.shape[0], -1)
    
    # Use first principal component for high-dimensional outputs
    if Y_np.shape[1] > 1:
        from sklearn.decomposition import PCA
        pca = PCA(n_components=1)
        Y_reduced = pca.fit_transform(Y_np).flatten()
    else:
        Y_reduced = Y_np.flatten()
    
    # First PC of input
    pca_x = PCA(n_components=1)
    X_reduced = pca_x.fit_transform(X_np).flatten()
    
    # Estimate MI
    mi = kraskov_mi_estimate(X_reduced, Y_reduced, k=5)
    
    return mi / np.log(2)  # Convert to bits

# Example with a simple model
model = nn.Sequential(
    nn.Linear(10, 64),
    nn.ReLU(),
    nn.Linear(64, 32),
    nn.ReLU(),
    nn.Linear(32, 1)
)

print("Model capacity estimation framework defined")

def information_bottleneck_bound(I_TY, I_TX, beta):
    """
    Compute the Information Bottleneck objective.
    
    The optimal representation T should:
    - Maximize I(T; Y) - predictive power
    - Minimize I(T; X) - compression
    
    Args:
        I_TY: Mutual information between representation and target
        I_TX: Mutual information between representation and input
        beta: Lagrange multiplier (tradeoff parameter)
    
    Returns:
        IB objective value (lower is better compression, higher beta)
    """
    return I_TX - beta * I_TY

def deep_variational_ib_loss(x, y, model, beta=0.01):
    """
    Deep Variational Information Bottleneck loss (Alemi et al., 2016).
    
    Uses variational bounds to make IB tractable for deep learning.
    
    DVIB Loss = E[-log p(y|z)] + β * KL(q(z|x) || p(z))
    
    where:
    - q(z|x) is the encoder (stochastic)
    - p(y|z) is the decoder/classifier
    - p(z) is the prior (typically N(0, I))
    """
    # Get encoder output
    z_mean, z_logvar = model.encode(x)
    
    # Sample z using reparameterization
    z_std = torch.exp(0.5 * z_logvar)
    z = z_mean + z_std * torch.randn_like(z_std)
    
    # Classify from z
    y_pred = model.classify(z)
    
    # Classification loss (cross-entropy)
    ce_loss = F.cross_entropy(y_pred, y)
    
    # KL divergence from prior
    kl_loss = -0.5 * torch.sum(1 + z_logvar - z_mean.pow(2) - z_logvar.exp()) / x.size(0)
    
    # Total DVIB loss
    total_loss = ce_loss + beta * kl_loss
    
    return total_loss, ce_loss.item(), kl_loss.item()

print("Deep Variational Information Bottleneck framework defined")

def compute_fisher_information(model, data_loader, criterion):
    """
    Compute diagonal of Fisher Information Matrix.
    
    Fisher Information measures how sensitive the model's predictions
    are to changes in each parameter.
    
    Args:
        model: PyTorch model
        data_loader: Data loader for estimation
        criterion: Loss function
    
    Returns:
        fisher_diagonal: Dictionary mapping parameter names to Fisher values
    """
    model.eval()
    fisher = {n: torch.zeros_like(p) for n, p in model.named_parameters() if p.requires_grad}
    
    n_samples = 0
    
    for x, y in data_loader:
        model.zero_grad()
        
        output = model(x)
        loss = criterion(output, y)
        loss.backward()
        
        # Accumulate squared gradients (diagonal Fisher approximation)
        for n, p in model.named_parameters():
            if p.requires_grad and p.grad is not None:
                fisher[n] += p.grad.pow(2) * x.size(0)
        
        n_samples += x.size(0)
    
    # Average
    for n in fisher:
        fisher[n] /= n_samples
    
    return fisher

def elastic_weight_consolidation_loss(model, fisher, old_params, lambda_ewc=1000):
    """
    EWC loss to prevent catastrophic forgetting.
    
    L_EWC = λ/2 * Σ_i F_i * (θ_i - θ*_i)²
    
    Penalizes changes to parameters that are important for previous tasks.
    
    Args:
        model: Current model
        fisher: Fisher Information from previous task
        old_params: Parameters after learning previous task
        lambda_ewc: Regularization strength
    
    Returns:
        EWC regularization loss
    """
    ewc_loss = 0
    
    for name, param in model.named_parameters():
        if name in fisher:
            # Weighted L2 penalty based on Fisher importance
            ewc_loss += torch.sum(fisher[name] * (param - old_params[name]).pow(2))
    
    return lambda_ewc * ewc_loss / 2

print("Fisher Information and EWC framework defined for continual learning")

class InformationFlowMonitor:
    """
    Monitor information flow through neural network during training.
    
    Tracks:
    - Layer-wise entropy
    - Mutual information between layers
    - Gradient information
    """
    
    def __init__(self, model, layer_names=None):
        self.model = model
        self.activations = {}
        self.gradients = {}
        self.hooks = []
        
        if layer_names is None:
            # Monitor all ReLU layers
            layer_names = [n for n, m in model.named_modules() 
                          if isinstance(m, (nn.ReLU, nn.Linear))]
        
        self.layer_names = layer_names
        self._register_hooks()
    
    def _register_hooks(self):
        """Register forward and backward hooks."""
        for name, module in self.model.named_modules():
            if name in self.layer_names:
                # Forward hook
                fwd_hook = module.register_forward_hook(
                    lambda m, inp, out, n=name: self._save_activation(n, out)
                )
                self.hooks.append(fwd_hook)
                
                # Backward hook
                bwd_hook = module.register_full_backward_hook(
                    lambda m, grad_in, grad_out, n=name: self._save_gradient(n, grad_out)
                )
                self.hooks.append(bwd_hook)
    
    def _save_activation(self, name, output):
        if isinstance(output, torch.Tensor):
            self.activations[name] = output.detach().cpu().numpy()
    
    def _save_gradient(self, name, grad_output):
        if grad_output[0] is not None:
            self.gradients[name] = grad_output[0].detach().cpu().numpy()
    
    def compute_layer_entropy(self, n_bins=30):
        """Compute entropy of activations at each layer."""
        entropies = {}
        
        for name, act in self.activations.items():
            # Flatten and discretize
            act_flat = act.flatten()
            hist, _ = np.histogram(act_flat, bins=n_bins, density=True)
            hist = hist[hist > 0]
            
            # Entropy (scaled by bin width)
            entropy = -np.sum(hist * np.log2(hist + 1e-10)) * (act_flat.max() - act_flat.min()) / n_bins
            entropies[name] = entropy
        
        return entropies
    
    def compute_activation_statistics(self):
        """Compute statistics useful for information analysis."""
        stats = {}
        
        for name, act in self.activations.items():
            stats[name] = {
                'mean': np.mean(act),
                'std': np.std(act),
                'sparsity': np.mean(act == 0),
                'max': np.max(act),
                'shape': act.shape
            }
        
        return stats
    
    def cleanup(self):
        """Remove hooks."""
        for hook in self.hooks:
            hook.remove()

# Example usage
print("=== Information Flow Monitoring ===")
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
)

# This demonstrates the monitoring capability
monitor = InformationFlowMonitor(model, layer_names=['0', '2', '4'])
print("Information flow monitor initialized")
print(f"Monitoring layers: {monitor.layer_names}")