{ "cells": [ { "cell_type": "markdown", "id": "e5e3bfc4", "metadata": {}, "source": [ "# OT in linear ICA\n", "\n", "#### In this notebook, I test the newest (06.03.2026) implementation of OT based ICA with a mixture of discrete distribution" ] }, { "cell_type": "code", "execution_count": 1, "id": "96da1ebb", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "import torch\n", "import pandas as pd\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", "\n", "from wasserstein_ica import WassersteinICA" ] }, { "cell_type": "code", "execution_count": 2, "id": "552ce44d", "metadata": {}, "outputs": [], "source": [ "# ==========================================\n", "# 1. The Bernoulli/Rademacher Trap Generator\n", "# ==========================================\n", "\n", "def generate_bernoulli_mixture(n_dim, n_samples, seed=None):\n", " \"\"\"\n", " Generates a mixture of independent centered Bernoulli {-1, 1} sources.\n", " This naturally has zero mean and unit variance.\n", " \"\"\"\n", " if seed is not None:\n", " np.random.seed(seed)\n", " torch.manual_seed(seed)\n", " \n", " # Generate independent discrete sources: -1 or 1\n", " S = np.random.choice([-1.0, 1.0], size=(n_dim, n_samples))\n", " \n", " # Generate a well-conditioned random mixing matrix A\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": 3, "id": "6eec31dd", "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": "de551289", "metadata": {}, "outputs": [], "source": [ "# ==========================================\n", "# 2. Parallel Worker Function (GPU Enabled)\n", "# ==========================================\n", "\n", "def run_bernoulli_trial(dim, trial, n_samples):\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_bernoulli_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", " # --- FastICA ---\n", " try:\n", " fast_ica = FastICA(n_components=dim, max_iter=2000, 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({'Dimension': dim, 'Method': 'FastICA', 'Amari Error': score_fast})\n", " \n", " # --- W-ICA (Stiefel + Gaussian Dithering) ---\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_restarts = dim * 4 if dim * 4 < 150 else 150\n", " \n", " # 1. Deflationary Phase (Standard W2 with Dithering)\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, max_iter=200, n_restarts=n_restarts, dither_sigma=0.05\n", " )\n", " extracted_ws.append(w)\n", " \n", " W_deflation_init = torch.stack(extracted_ws)\n", " \n", " # Symmetric Stiefel Phase\n", " W_stiefel_unmixed = ica.optimize_symmetric(\n", " n_components=dim, \n", " max_iter=400, # Increased to allow for longer settling\n", " lr=0.25, # Lowered for a more stable descent\n", " init_w=W_deflation_init, \n", " optimizer='stiefel',\n", " use_sinkhorn=False,\n", " dither_sigma=0.01, # Reduced blur to keep 50D corners sharp\n", " batch_size=1024 # Increased batch to handle high-dimensional geometry\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", " print(f\"W-ICA Error (dim {dim}): {e}\")\n", " score_wass = np.nan\n", " \n", " trial_results.append({'Dimension': dim, 'Method': 'W-ICA (Stiefel)', 'Amari Error': score_wass})\n", " \n", " return trial_results" ] }, { "cell_type": "code", "execution_count": 5, "id": "7a0b3d77", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "--- FastICA vs W-ICA: Bernoulli Mixture (10000 Samples) ---\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "d045c92c63274e70a05db903c41e67ee", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Running Bernoulli Trials: 0%| | 0/55 [00:00" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
MethodFastICAW-ICA (Stiefel)
Dimension
100.03330.0333
150.05590.0559
200.07330.0734
250.09170.1080
300.11250.2563
350.13220.1682
400.15380.3316
450.17451.3570
500.19424.8453
550.21498.1419
600.234910.6195
\n", "
" ], "text/plain": [ "Method FastICA W-ICA (Stiefel)\n", "Dimension \n", "10 0.0333 0.0333\n", "15 0.0559 0.0559\n", "20 0.0733 0.0734\n", "25 0.0917 0.1080\n", "30 0.1125 0.2563\n", "35 0.1322 0.1682\n", "40 0.1538 0.3316\n", "45 0.1745 1.3570\n", "50 0.1942 4.8453\n", "55 0.2149 8.1419\n", "60 0.2349 10.6195" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "plt.figure(figsize=(8, 5))\n", "sns.lineplot(data=df_bernoulli, x='Dimension', y='Amari Error', hue='Method', marker='o', linewidth=2.5)\n", "plt.title(f\"Amari Error vs. Dimension\\n(Discrete Centered Bernoulli Mixture, N={N_SAMPLES})\", fontsize=14)\n", "plt.ylabel(\"Amari Error (Lower is Better)\", fontsize=12)\n", "plt.xlabel(\"Number of Dimensions\", fontsize=12)\n", "plt.xticks(DIMENSIONS)\n", "plt.grid(True, linestyle='--', alpha=0.7)\n", "plt.tight_layout()\n", "plt.show()\n", "\n", "display(df_bernoulli.groupby(['Dimension', 'Method'])['Amari Error'].mean().unstack().round(4))" ] }, { "cell_type": "code", "execution_count": null, "id": "b4fc402d", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "ot_ica", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.4" } }, "nbformat": 4, "nbformat_minor": 5 }