{ "cells": [ { "cell_type": "markdown", "id": "745330c4", "metadata": {}, "source": [ "# Optimal Transport in linear ICA\n", "\n", "#### In this notebook we test the performance of a OT Fixed Point ICA versus Fast ICA over dimensions, the latter also uses Fixed point alg as it's contrast function is infinitely differentiable." ] }, { "cell_type": "code", "execution_count": 1, "id": "83186914", "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": "83baaabc", "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": "d5b84c2f", "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": "35963e2f", "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "dac3bea691dd448e9b1c83fff723a527", "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", "
MethodFastICAStandalone W-ICAW-ICA (Cold Start)
Dimension
50.04820.04470.1003
100.11830.10980.1428
150.17710.16661.3867
200.23240.22083.1077
250.29710.31815.0471
300.36530.52646.7954
350.42871.32937.9887
400.49522.65569.9702
450.56124.719911.8817
500.63106.622812.9816
\n", "
" ], "text/plain": [ "Method FastICA Standalone W-ICA W-ICA (Cold Start)\n", "Dimension \n", "5 0.0482 0.0447 0.1003\n", "10 0.1183 0.1098 0.1428\n", "15 0.1771 0.1666 1.3867\n", "20 0.2324 0.2208 3.1077\n", "25 0.2971 0.3181 5.0471\n", "30 0.3653 0.5264 6.7954\n", "35 0.4287 1.3293 7.9887\n", "40 0.4952 2.6556 9.9702\n", "45 0.5612 4.7199 11.8817\n", "50 0.6310 6.6228 12.9816" ] }, "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": "06a8f205", "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 }