{ "cells": [ { "cell_type": "markdown", "id": "50483cd6", "metadata": {}, "source": [ "# Optimal Transport in linear ICA\n", "\n", "#### In this notebook we test the performance of a OT Natural Gradient ICA versus Fast ICA over dimensions." ] }, { "cell_type": "code", "execution_count": 1, "id": "fac3658e", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "import torch\n", "import matplotlib.pyplot as plt\n", "import pandas as pd\n", "import seaborn as sns\n", "from sklearn.decomposition import FastICA\n", "from tqdm.notebook import tqdm\n", "\n", "# Import your package\n", "from wasserstein_ica import WassersteinICA" ] }, { "cell_type": "code", "execution_count": 2, "id": "18192e4a", "metadata": {}, "outputs": [], "source": [ "# ==========================================\n", "# 1. Helpers & Data Generation\n", "# ==========================================\n", "\n", "def amari_error(W_est, A_true):\n", " if W_est is None or np.any(np.isnan(W_est)):\n", " return np.nan\n", " \n", " P = np.abs(W_est @ A_true)\n", " n = P.shape[0]\n", " \n", " row_sum = np.sum(P, axis=1)\n", " row_max = np.max(P, axis=1)\n", " term1 = np.sum((row_sum / row_max) - 1)\n", " \n", " col_sum = np.sum(P, axis=0)\n", " col_max = np.max(P, axis=0)\n", " term2 = np.sum((col_sum / col_max) - 1)\n", " \n", " return (term1 + term2) / (2 * n)\n", "\n", "def generate_dataset(n_dim, n_samples, seed=None):\n", " if seed is not None:\n", " np.random.seed(seed)\n", " torch.manual_seed(seed)\n", " \n", " # Standard Laplace (Super-Gaussian)\n", " sources = [np.random.laplace(0, 1, n_samples) for _ in range(n_dim)]\n", " S = np.stack(sources)\n", " \n", " # 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": 3, "id": "96f75f68", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "--- FastICA vs. OT Fixed-Point Showdown ---\n", "Dimensions: [5, 10, 15, 20, 25, 30, 35, 40, 45, 50]\n", "Samples: 5000\n", "Trials per dim: 5\n" ] } ], "source": [ "# ==========================================\n", "# 2. Experiment Setup\n", "# ==========================================\n", "\n", "# Test a range of dimensions up to 30\n", "DIMENSIONS = [n for n in range(5, 51, 5)] # Varying D\n", "N_SAMPLES = 5000\n", "N_TRIALS = 5 # Run multiple trials for error bars\n", "\n", "print(f\"--- FastICA vs. OT Fixed-Point Showdown ---\")\n", "print(f\"Dimensions: {DIMENSIONS}\")\n", "print(f\"Samples: {N_SAMPLES}\")\n", "print(f\"Trials per dim: {N_TRIALS}\")\n", "\n", "results = []" ] }, { "cell_type": "code", "execution_count": 4, "id": "bbb23f67", "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "168ee53e894c4ffdaab913dd2038b0b2", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Dimensions: 0%| | 0/10 [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", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
MethodFastICAW-ICA (Cold Start)W-ICA (Natural Gradient)
Dimension
50.04820.10032.3385
100.11830.14283.7281
150.17711.38674.5186
200.23243.10775.9845
250.29715.04716.9972
300.36536.79548.5238
350.42877.98879.5671
400.49529.970211.2877
450.561211.881712.4396
500.631012.981613.8901
\n", "
" ], "text/plain": [ "Method FastICA W-ICA (Cold Start) W-ICA (Natural Gradient)\n", "Dimension \n", "5 0.0482 0.1003 2.3385\n", "10 0.1183 0.1428 3.7281\n", "15 0.1771 1.3867 4.5186\n", "20 0.2324 3.1077 5.9845\n", "25 0.2971 5.0471 6.9972\n", "30 0.3653 6.7954 8.5238\n", "35 0.4287 7.9887 9.5671\n", "40 0.4952 9.9702 11.2877\n", "45 0.5612 11.8817 12.4396\n", "50 0.6310 12.9816 13.8901" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "df = pd.DataFrame(results)\n", "\n", "plt.figure(figsize=(10, 6))\n", "# seaborn's lineplot automatically plots the mean and a confidence interval for the trials\n", "sns.lineplot(data=df, x='Dimension', y='Amari Error', hue='Method', marker='o', linewidth=2.5)\n", "\n", "plt.title(\"Amari Error vs. Dimension\\n(FastICA vs. Wasserstein Fixed-Point)\", fontsize=14)\n", "plt.ylabel(\"Amari Error (Lower is Better)\", fontsize=12)\n", "plt.xlabel(\"Number of Dimensions\", fontsize=12)\n", "plt.grid(True, linestyle='--', alpha=0.7)\n", "#plt.yscale('log') # Log scale is often better for Amari Error to see small polish improvements\n", "plt.tight_layout()\n", "plt.show()\n", "\n", "# Display mean scores table\n", "display(df.groupby(['Dimension', 'Method'])['Amari Error'].mean().unstack().round(4))" ] }, { "cell_type": "code", "execution_count": null, "id": "03630b2b", "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 }