{
"cells": [
{
"cell_type": "markdown",
"id": "e8b64456",
"metadata": {},
"source": [
"# OT in linar ICA\n",
"\n",
"#### In this notebook, we genrate mixtures by randomly picking distributions, whether discrete or continous. We also keep at most one gaussian to confirm tractability of ICA with one gaussian. We compare how Fast ICA and OT ICA perform on 10 to 75 dima with 5 dim leap and 10,000 samples wiht max of 150 starts in phase 1 SGD step and phase 2 incorporating Stiefel manifold."
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "b95f518e",
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"import torch\n",
"import pandas as pd\n",
"import time\n",
"from sklearn.decomposition import FastICA\n",
"from joblib import Parallel, delayed\n",
"from tqdm.notebook import tqdm\n",
"import seaborn as sns\n",
"import matplotlib.pyplot as plt\n",
"from matplotlib.ticker import MultipleLocator\n",
"import warnings\n",
"from sklearn.exceptions import ConvergenceWarning\n",
"\n",
"from wasserstein_ica import WassersteinICA"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "4376d273",
"metadata": {},
"outputs": [],
"source": [
"# ==========================================\n",
"# Thesis Theme Configuration\n",
"# ==========================================\n",
"def set_thesis_theme():\n",
" thesis_colors = ['#0173B2', '#DE8F05', '#029E73', '#D55E00', '#CC78BC', '#CA9161']\n",
" plt.rcParams.update({\n",
" 'figure.figsize': (10, 6),\n",
" 'figure.dpi': 300,\n",
" 'axes.prop_cycle': plt.cycler(color=thesis_colors),\n",
" 'axes.grid': True,\n",
" 'grid.alpha': 0.3,\n",
" 'grid.linestyle': '--',\n",
" 'axes.axisbelow': True,\n",
" 'axes.spines.top': False,\n",
" 'axes.spines.right': False,\n",
" 'font.size': 12,\n",
" 'axes.titlesize': 14,\n",
" 'axes.labelsize': 12,\n",
" 'legend.frameon': True,\n",
" 'legend.fontsize': 11,\n",
" 'lines.linewidth': 2.5\n",
" })\n",
"set_thesis_theme()"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "8d3f07f9",
"metadata": {},
"outputs": [],
"source": [
"def amari_error(W, A):\n",
" \"\"\"\n",
" Computes the Amari Error to evaluate the performance of ICA.\n",
" An error close to 0 indicates perfect separation.\n",
" \n",
" W: Estimated unmixing matrix (n_components, n_features)\n",
" A: True mixing matrix (n_features, n_sources)\n",
" \"\"\"\n",
" # P is the global matrix representing the combined effect of mixing and unmixing\n",
" P = np.dot(W, A)\n",
" P_abs = np.abs(P)\n",
" \n",
" # 1. Sum of row-wise errors\n",
" row_max = np.max(P_abs, axis=1, keepdims=True)\n",
" row_term = np.sum(P_abs / row_max, axis=1) - 1.0\n",
" \n",
" # 2. Sum of column-wise errors\n",
" col_max = np.max(P_abs, axis=0, keepdims=True)\n",
" col_term = np.sum(P_abs / col_max, axis=0) - 1.0\n",
" \n",
" # Average over the number of components\n",
" n = P.shape[0]\n",
" return (np.sum(row_term) + np.sum(col_term)) / (2 * n)"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "ae808f18",
"metadata": {},
"outputs": [],
"source": [
"# ==========================================\n",
"# 1. The Generalized Hybrid Mixture Generator\n",
"# ==========================================\n",
"\n",
"def generate_general_mixture(n_dim, n_samples, seed=None):\n",
" \"\"\"\n",
" Generates a generalized ICA dataset.\n",
" Ensures exactly ONE source is Gaussian.\n",
" Randomly pulls the remaining (n_dim - 1) sources from a diverse pool \n",
" of discrete and continuous non-Gaussian distributions.\n",
" \"\"\"\n",
" if seed is not None:\n",
" np.random.seed(seed)\n",
" torch.manual_seed(seed)\n",
"\n",
" sources = []\n",
" \n",
" # 1. Inject Exactly ONE Gaussian Source\n",
" sources.append(np.random.normal(0, 1, size=n_samples))\n",
" \n",
" # 2. Define the Non-Gaussian Pool (all naturally or explicitly standardized)\n",
" def gen_laplace(): \n",
" return np.random.laplace(0, 1 / np.sqrt(2), size=n_samples)\n",
" def gen_bernoulli(): \n",
" return np.random.choice([-1.0, 1.0], size=n_samples)\n",
" def gen_uniform(): \n",
" return np.random.uniform(-np.sqrt(3), np.sqrt(3), size=n_samples)\n",
" def gen_student_t(): \n",
" s = np.random.standard_t(df=3, size=n_samples)\n",
" return s / np.std(s) # Standardize due to heavy tails\n",
" def gen_poisson():\n",
" s = np.random.poisson(lam=3, size=n_samples)\n",
" return (s - np.mean(s)) / np.std(s) # Center and scale\n",
" def gen_binomial():\n",
" s = np.random.binomial(n=10, p=0.5, size=n_samples)\n",
" return (s - np.mean(s)) / np.std(s) # Center and scale\n",
" def gen_chisquare():\n",
" s = np.random.chisquare(df=2, size=n_samples)\n",
" return (s - np.mean(s)) / np.std(s) # Center and scale\n",
" def gen_exponential():\n",
" s = np.random.exponential(scale=1.0, size=n_samples)\n",
" return (s - np.mean(s)) / np.std(s) # Center and scale\n",
" \n",
" pool = [gen_laplace, gen_bernoulli, gen_uniform, gen_student_t, \n",
" gen_poisson, gen_binomial, gen_chisquare, gen_exponential]\n",
" \n",
" # 3. Generate the remaining (n_dim - 1) sources\n",
" for _ in range(n_dim - 1):\n",
" # Select a random distribution generator from the pool\n",
" generator = np.random.choice(pool)\n",
" sources.append(generator())\n",
" \n",
" S = np.stack(sources)\n",
" \n",
" # Shuffle the rows so the Gaussian source is at a random index\n",
" np.random.shuffle(S)\n",
" \n",
" # 4. Generate a well-conditioned mixing matrix\n",
" cond_num = 1000\n",
" while cond_num > 100:\n",
" A = np.random.randn(n_dim, n_dim)\n",
" cond_num = np.linalg.cond(A)\n",
" \n",
" X = A @ S\n",
" return torch.tensor(X, dtype=torch.float32), A"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "9d2c0889",
"metadata": {},
"outputs": [],
"source": [
"# ==========================================\n",
"# 2. Parallel Worker Function (GPU Enabled)\n",
"# ==========================================\n",
"\n",
"def run_general_trial(dim, trial, n_samples, compute_level):\n",
" torch.set_num_threads(1) \n",
" trial_results = []\n",
" \n",
" device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
" X_torch, A_true = generate_general_mixture(n_dim=dim, n_samples=n_samples, seed=trial)\n",
" X_torch = X_torch.to(device)\n",
" X_np = X_torch.cpu().numpy()\n",
" \n",
" # Configure Compute Parameters\n",
" if compute_level == 'Low Compute':\n",
" fast_max_iter = 1000\n",
" w2_restarts = min(dim * 4, 150)\n",
" w2_def_iter = 150\n",
" w2_sym_iter = 300\n",
" else: # High Compute\n",
" fast_max_iter = 10000 # Test if infinite iterations help FastICA\n",
" w2_restarts = min(dim * 15, 600) \n",
" w2_def_iter = 300\n",
" w2_sym_iter = 800\n",
" \n",
" # --- FastICA ---\n",
" try:\n",
" # We allow ConvergenceWarnings to trigger naturally in the console\n",
" fast_ica = FastICA(n_components=dim, max_iter=fast_max_iter, tol=1e-4, random_state=trial)\n",
" fast_ica.fit(X_np.T)\n",
" W_fast_total = fast_ica.components_\n",
" score_fast = amari_error(W_fast_total, A_true)\n",
" except Exception:\n",
" score_fast = np.nan\n",
" \n",
" trial_results.append({\n",
" 'Dimension': dim, \n",
" 'Method': 'FastICA', \n",
" 'Compute Level': compute_level,\n",
" 'Amari Error': score_fast\n",
" })\n",
" \n",
" # --- OT-ICA (Stiefel) ---\n",
" try:\n",
" ica = WassersteinICA(X_torch)\n",
" ica.whiten()\n",
" W_white_np = ica.W_white.cpu().numpy()\n",
" \n",
" extracted_ws = []\n",
" \n",
" # 1. Deflation Phase\n",
" for _ in range(dim):\n",
" prev = torch.stack(extracted_ws) if extracted_ws else None\n",
" w, _ = ica.optimize_wasserstein2(\n",
" prev_components=prev, \n",
" max_iter=w2_def_iter, \n",
" n_restarts=w2_restarts, \n",
" dither_sigma=0.01\n",
" )\n",
" extracted_ws.append(w)\n",
" \n",
" W_deflation_init = torch.stack(extracted_ws)\n",
" \n",
" # 2. Symmetric Stiefel Phase \n",
" W_stiefel_unmixed = ica.optimize_symmetric(\n",
" n_components=dim, \n",
" max_iter=w2_sym_iter, \n",
" lr=0.25, \n",
" init_w=W_deflation_init, \n",
" optimizer='stiefel',\n",
" use_sinkhorn=False,\n",
" dither_sigma=0.01, \n",
" batch_size=1024 \n",
" )\n",
" W_wass_total = W_stiefel_unmixed.cpu().numpy() @ W_white_np\n",
" score_wass = amari_error(W_wass_total, A_true)\n",
" \n",
" except Exception as e:\n",
" score_wass = np.nan\n",
" \n",
" trial_results.append({\n",
" 'Dimension': dim, \n",
" 'Method': 'OT-ICA', \n",
" 'Compute Level': compute_level,\n",
" 'Amari Error': score_wass\n",
" })\n",
" \n",
" return trial_results"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "adb6cd78",
"metadata": {},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "76842289221b42cab32968f58dd2ee71",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
" 0%| | 0/130 [00:00"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"# ==========================================\n",
"# 4. Plotting the Results\n",
"# ==========================================\n",
"# Ensure theme is applied first\n",
"set_thesis_theme()\n",
"\n",
"# --- Side-by-Side Plotting ---\n",
"fig, axes = plt.subplots(1, 2, figsize=(14, 6), sharey=True)\n",
"palette = {'FastICA': '#0173B2', 'OT-ICA': '#D55E00'}\n",
"\n",
"for idx, comp in enumerate(COMPUTE_LEVELS):\n",
" sns.lineplot(\n",
" data=df_general[df_general['Compute Level'] == comp], \n",
" x='Dimension', y='Amari Error', hue='Method', \n",
" palette=palette, marker='o', linewidth=2.5, ax=axes[idx]\n",
" )\n",
" \n",
" # Force grid to appear\n",
" axes[idx].grid(True, linestyle='--', alpha=0.3)\n",
" \n",
" axes[idx].set_title(f\"{comp} Regime\", fontsize=14)\n",
" axes[idx].set_xlabel(\"Number of Dimensions\", fontsize=12)\n",
" axes[idx].set_xticks(DIMENSIONS)\n",
" \n",
" # The dotted gray line\n",
" axes[idx].axhline(0.3, color='gray', linestyle=':', alpha=0.8, linewidth=2)\n",
" \n",
" # Add text to the first graph (Left Plot)\n",
" if idx == 0:\n",
" axes[idx].set_ylabel(\"Amari Error (Lower is Better)\", fontsize=12)\n",
" axes[idx].legend(title='Method', loc='upper left')\n",
" \n",
" # Annotation: x is the max dimension, y is slightly above the line\n",
" axes[idx].text(\n",
" x=max(DIMENSIONS), y=0.35, \n",
" s='Good Separation', \n",
" color='gray', \n",
" fontsize=10, \n",
" fontweight='bold',\n",
" ha='right', # Horizontal alignment: Right\n",
" va='bottom' # Vertical alignment: Bottom\n",
" )\n",
" else:\n",
" axes[idx].set_ylabel(\"\")\n",
" if axes[idx].get_legend(): axes[idx].get_legend().remove()\n",
"\n",
"# Apply Y-limit and 0.5 ticks\n",
"plt.ylim(0, 4)\n",
"from matplotlib.ticker import MultipleLocator\n",
"axes[0].yaxis.set_major_locator(MultipleLocator(0.5))\n",
"\n",
"fig.suptitle(f\"General Hybrid Mixture Stress Test (N={N_SAMPLES})\", fontsize=16, y=1.05)\n",
"plt.tight_layout()\n",
"plt.savefig('hybrid_mixture_amari_compute.png', dpi=300, bbox_inches='tight')\n",
"plt.show()"
]
},
{
"cell_type": "markdown",
"id": "662d99e6",
"metadata": {},
"source": [
"### We run an experiment investigating why FastICA fails, we increase iterations and tolerances to help it converge. This helps us verify if Fast ICa fails on principle in this general mixture setting or compute power."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c4102968",
"metadata": {},
"outputs": [],
"source": [
"\"\"\"\n",
"import numpy as np\n",
"import pandas as pd\n",
"import warnings\n",
"from sklearn.decomposition import FastICA\n",
"from sklearn.exceptions import ConvergenceWarning\n",
"from joblib import Parallel, delayed\n",
"from tqdm.notebook import tqdm\n",
"import seaborn as sns\n",
"import matplotlib.pyplot as plt\n",
"\"\"\""
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "e5784250",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"--- FastICA Limits: The Oscillation Test ---\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "a8759c3f90814f5f975540a7a9ccb613",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Running FastICA configs: 0%| | 0/210 [00:00"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" | Config | \n",
" Baseline (Iter 2k, Tol 1e-4) | \n",
" High Iter (Iter 10k, Tol 1e-4) | \n",
" Relaxed Tol (Iter 2k, Tol 1e-3) | \n",
"
\n",
" \n",
" | Dimension | \n",
" | \n",
" | \n",
" | \n",
"
\n",
" \n",
" \n",
" \n",
" | 10 | \n",
" 0.3794 | \n",
" 0.3794 | \n",
" 0.3678 | \n",
"
\n",
" \n",
" | 15 | \n",
" 0.5025 | \n",
" 0.5025 | \n",
" 0.4982 | \n",
"
\n",
" \n",
" | 20 | \n",
" 0.5304 | \n",
" 0.5238 | \n",
" 0.5270 | \n",
"
\n",
" \n",
" | 25 | \n",
" 0.6711 | \n",
" 0.6193 | \n",
" 0.6698 | \n",
"
\n",
" \n",
" | 30 | \n",
" 0.7127 | \n",
" 0.7150 | \n",
" 0.7216 | \n",
"
\n",
" \n",
" | 35 | \n",
" 0.8250 | \n",
" 0.8331 | \n",
" 0.8252 | \n",
"
\n",
" \n",
" | 40 | \n",
" 0.9284 | \n",
" 0.8900 | \n",
" 0.9014 | \n",
"
\n",
" \n",
" | 45 | \n",
" 1.2865 | \n",
" 1.2327 | \n",
" 1.2853 | \n",
"
\n",
" \n",
" | 50 | \n",
" 1.5391 | \n",
" 1.5432 | \n",
" 1.5394 | \n",
"
\n",
" \n",
" | 55 | \n",
" 1.5970 | \n",
" 1.5768 | \n",
" 1.5970 | \n",
"
\n",
" \n",
" | 60 | \n",
" 1.6028 | \n",
" 1.5507 | \n",
" 1.6028 | \n",
"
\n",
" \n",
" | 65 | \n",
" 1.8245 | \n",
" 1.7620 | \n",
" 1.8219 | \n",
"
\n",
" \n",
" | 70 | \n",
" 1.7998 | \n",
" 1.9520 | \n",
" 1.7998 | \n",
"
\n",
" \n",
" | 75 | \n",
" 1.9995 | \n",
" 2.0209 | \n",
" 2.0098 | \n",
"
\n",
" \n",
"
\n",
"
\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" | \n",
" Amari Error | \n",
" Converged | \n",
"
\n",
" \n",
" | Dimension | \n",
" Pool Type | \n",
" | \n",
" | \n",
"
\n",
" \n",
" \n",
" \n",
" | 30 | \n",
" Continuous Only | \n",
" 0.2672 | \n",
" 1.0 | \n",
"
\n",
" \n",
" | Discrete Only | \n",
" 4.5073 | \n",
" 0.0 | \n",
"
\n",
" \n",
" | Full Hybrid | \n",
" 0.7225 | \n",
" 0.6 | \n",
"
\n",
" \n",
" | Strictly Super-Gaussian | \n",
" 0.2813 | \n",
" 1.0 | \n",
"
\n",
" \n",
" | Zero Gaussian | \n",
" 0.5771 | \n",
" 0.4 | \n",
"
\n",
" \n",
" | 40 | \n",
" Continuous Only | \n",
" 0.3545 | \n",
" 1.0 | \n",
"
\n",
" \n",
" | Discrete Only | \n",
" 5.7794 | \n",
" 0.0 | \n",
"
\n",
" \n",
" | Full Hybrid | \n",
" 0.9800 | \n",
" 0.4 | \n",
"
\n",
" \n",
" | Strictly Super-Gaussian | \n",
" 0.3764 | \n",
" 1.0 | \n",
"
\n",
" \n",
" | Zero Gaussian | \n",
" 0.6935 | \n",
" 0.4 | \n",
"
\n",
" \n",
"
\n",
"
"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" | \n",
" Amari Error | \n",
" Converged | \n",
"
\n",
" \n",
" | Dimension | \n",
" Pool Type | \n",
" | \n",
" | \n",
"
\n",
" \n",
" \n",
" \n",
" | 30 | \n",
" Continuous Only | \n",
" 0.1598 | \n",
" 1.0 | \n",
"
\n",
" \n",
" | Discrete Only | \n",
" 3.6605 | \n",
" 1.0 | \n",
"
\n",
" \n",
" | Full Hybrid | \n",
" 0.2797 | \n",
" 1.0 | \n",
"
\n",
" \n",
" | Strictly Super-Gaussian | \n",
" 0.1759 | \n",
" 1.0 | \n",
"
\n",
" \n",
" | Zero Gaussian | \n",
" 0.1380 | \n",
" 1.0 | \n",
"
\n",
" \n",
" | 40 | \n",
" Continuous Only | \n",
" 0.2412 | \n",
" 1.0 | \n",
"
\n",
" \n",
" | Discrete Only | \n",
" 6.4085 | \n",
" 1.0 | \n",
"
\n",
" \n",
" | Full Hybrid | \n",
" 0.5367 | \n",
" 1.0 | \n",
"
\n",
" \n",
" | Strictly Super-Gaussian | \n",
" 0.2322 | \n",
" 1.0 | \n",
"
\n",
" \n",
" | Zero Gaussian | \n",
" 0.3007 | \n",
" 1.0 | \n",
"
\n",
" \n",
"
\n",
"
"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" | \n",
" Amari Error | \n",
" Converged | \n",
"
\n",
" \n",
" | Dimension | \n",
" Distribution | \n",
" | \n",
" | \n",
"
\n",
" \n",
" \n",
" \n",
" | 30 | \n",
" Bernoulli | \n",
" 0.1227 | \n",
" 1.0 | \n",
"
\n",
" \n",
" | Binomial | \n",
" 9.0120 | \n",
" 0.0 | \n",
"
\n",
" \n",
" | Poisson | \n",
" 7.8712 | \n",
" 0.0 | \n",
"
\n",
" \n",
" | 40 | \n",
" Bernoulli | \n",
" 0.1620 | \n",
" 1.0 | \n",
"
\n",
" \n",
" | Binomial | \n",
" 12.2742 | \n",
" 0.0 | \n",
"
\n",
" \n",
" | Poisson | \n",
" 11.2480 | \n",
" 0.0 | \n",
"
\n",
" \n",
"
\n",
"
"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" | \n",
" Amari Error | \n",
" Converged | \n",
"
\n",
" \n",
" | Dimension | \n",
" Distribution | \n",
" | \n",
" | \n",
"
\n",
" \n",
" \n",
" \n",
" | 30 | \n",
" Bernoulli | \n",
" 0.1830 | \n",
" 1.0 | \n",
"
\n",
" \n",
" | Binomial | \n",
" 9.4444 | \n",
" 1.0 | \n",
"
\n",
" \n",
" | Poisson | \n",
" 4.0633 | \n",
" 1.0 | \n",
"
\n",
" \n",
" | 40 | \n",
" Bernoulli | \n",
" 0.3047 | \n",
" 1.0 | \n",
"
\n",
" \n",
" | Binomial | \n",
" 12.4825 | \n",
" 1.0 | \n",
"
\n",
" \n",
" | Poisson | \n",
" 7.0400 | \n",
" 1.0 | \n",
"
\n",
" \n",
"
\n",
"