Joint Distributions

import numpy as np
from scipy import stats

# === Joint PMF from a table ===
joint_pmf = np.array([
    [0.1, 0.2, 0.1],   # X=0: Y=0,1,2
    [0.2, 0.3, 0.1],   # X=1: Y=0,1,2
    [0.1, 0.1, 0.1]    # X=2: Y=0,1,2
])

# Marginal distributions
p_x = joint_pmf.sum(axis=1)  # Sum over Y -> P(X)
p_y = joint_pmf.sum(axis=0)  # Sum over X -> P(Y)
print(f"Marginal P(X): {p_x}")   # [0.4, 0.6, 0.3]  -- note: these don't sum to 1, fix indices
print(f"Marginal P(Y): {p_y}")   # [0.4, 0.6, 0.3]

# Conditional distribution P(Y | X=1)
x_idx = 1
p_y_given_x1 = joint_pmf[x_idx] / p_x[x_idx]
print(f"P(Y | X=1): {p_y_given_x1}")

# Verify: conditional probabilities sum to 1
print(f"Sum of conditional: {p_y_given_x1.sum():.4f}")  # Should be 1.0

# === Joint PDF for continuous variables ===
# Bivariate normal example
mean = [0, 0]
cov_matrix = [[1, 0.8], [0.8, 1]]  # Correlation = 0.8

# Sample from bivariate normal
np.random.seed(42)
n_samples = 10000
samples = np.random.multivariate_normal(mean, cov_matrix, size=n_samples)
X_samples, Y_samples = samples[:, 0], samples[:, 1]

# Empirical covariance and correlation
emp_cov = np.cov(X_samples, Y_samples)
emp_corr = np.corrcoef(X_samples, Y_samples)
print(f"Empirical covariance matrix:\n{emp_cov}")
print(f"Empirical correlation:\n{emp_corr}")

# === Independence check ===
# Two independent variables
X_ind = np.random.randn(5000)
Y_ind = np.random.randn(5000)
print(f"Correlation (independent): {np.corrcoef(X_ind, Y_ind)[0,1]:.4f}")

# Two dependent variables (Y = X^2)
X_dep = np.random.uniform(-1, 1, 5000)
Y_dep = X_dep ** 2
print(f"Correlation (dependent but uncorrelated): {np.corrcoef(X_dep, Y_dep)[0,1]:.4f}")

# === Mutual information for nonlinear dependence ===
from scipy.stats import entropy

def mutual_information_2d(x, y, bins=20):
    """Estimate mutual information using histogram-based method."""
    p_xy = np.histogram2d(x, y, bins=bins)[0]
    p_xy = p_xy / p_xy.sum()
    p_x = p_xy.sum(axis=1)
    p_y = p_xy.sum(axis=0)
    mi = entropy(p_x) + entropy(p_y) - entropy(p_xy.ravel())
    return mi

print(f"MI (independent): {mutual_information_2d(X_ind, Y_ind):.4f}")
print(f"MI (X, X^2):     {mutual_information_2d(X_dep, Y_dep):.4f}")

# === Visualize joint distribution ===
import matplotlib.pyplot as plt

fig, axes = plt.subplots(1, 3, figsize=(15, 4))

# Joint PMF heatmap
im = axes[0].imshow(joint_pmf, cmap='Blues', vmin=0, vmax=0.3)
axes[0].set_xlabel('Y')
axes[0].set_ylabel('X')
axes[0].set_title('Joint PMF')
plt.colorbar(im, ax=axes[0])

# Bivariate normal samples
axes[1].scatter(X_samples[:500], Y_samples[:500], alpha=0.3, s=10)
axes[1].set_xlabel('X')
axes[1].set_ylabel('Y')
axes[1].set_title('Bivariate Normal (ρ=0.8)')
axes[1].set_aspect('equal')

# Dependent but uncorrelated
axes[2].scatter(X_dep[:500], Y_dep[:500], alpha=0.3, s=10)
axes[2].set_xlabel('X')
axes[2].set_ylabel('Y = X²')
axes[2].set_title('Dependent but Uncorrelated')

plt.tight_layout()
plt.savefig('joint_distributions.png', dpi=150)
plt.show()