{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9a9473df",
   "metadata": {},
   "source": [
    "# Import Required Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e2836749",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "sys.path.insert(0, \"../../src\")\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython import display\n",
    "import pylab as pl\n",
    "\n",
    "from CorInfoMaxBSS import *\n",
    "from general_utils import *\n",
    "from visualization_utils import *\n",
    "\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "notebook_name = \"Sparse\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6012a28a",
   "metadata": {},
   "source": [
    "# Source Generation and Mixing Scenario"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8e5ed26f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The following is the mixture matrix A\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{bmatrix} \n",
       " 0.027 & 0.328 & 2.410 & -1.619 & 0.136 & \\\\\n",
       " -0.690 & -1.047 & 0.753 & 0.602 & -1.008 & \\\\\n",
       " -1.755 & -0.302 & -0.378 & -0.133 & -0.408 & \\\\\n",
       " 0.433 & 0.360 & 0.686 & -0.041 & -1.018 & \\\\\n",
       " 1.878 & -0.016 & -2.512 & 0.571 & -0.216 & \\\\\n",
       " 0.137 & 0.445 & 0.056 & -0.781 & 0.911 & \\\\\n",
       " 0.849 & 2.524 & 0.956 & 0.075 & 0.783 & \\\\\n",
       " -0.544 & -1.443 & 0.634 & -1.249 & -0.857 & \\\\\n",
       " -0.655 & 1.399 & -0.859 & 0.144 & 0.428 & \\\\\n",
       " -0.249 & -2.637 & -1.049 & -0.798 & -0.302 & \\\\\n",
       "\\end{bmatrix}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input SNR is : 29.993846839695646\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 9228 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "N = 500000\n",
    "NumberofSources = 5\n",
    "NumberofMixtures = 10\n",
    "S = generate_correlated_copula_sources(\n",
    "    rho=0.0,\n",
    "    df=4,\n",
    "    n_sources=NumberofSources,\n",
    "    size_sources=N,\n",
    "    decreasing_correlation=True,\n",
    ")\n",
    "S = 4 * S - 2\n",
    "S = ProjectRowstoL1NormBall(S.T).T\n",
    "\n",
    "# print(\"The following is the correlation matrix of sources\")\n",
    "# display_matrix(np.corrcoef(S))\n",
    "\n",
    "plt.scatter(S[0, :], S[2, :])\n",
    "\n",
    "\n",
    "# Generate Mxr random mixing from i.i.d N(0,1)\n",
    "A = np.random.randn(NumberofMixtures, NumberofSources)\n",
    "X = np.dot(A, S)\n",
    "\n",
    "SNR = 30\n",
    "X, NoisePart = addWGN(X, SNR, return_noise=True)\n",
    "\n",
    "SNRinp = 10 * np.log10(\n",
    "    np.sum(np.mean((X - NoisePart) ** 2, axis=1))\n",
    "    / np.sum(np.mean(NoisePart**2, axis=1))\n",
    ")\n",
    "print(\"The following is the mixture matrix A\")\n",
    "display_matrix(A)\n",
    "print(\"Input SNR is : {}\".format(SNRinp))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51ee1563",
   "metadata": {},
   "source": [
    "# Visualize Generated Sources and Mixtures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "381fb2fe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 134812 characters]\n",
      "text/plain": [
       "<Figure size 1094.4x648 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 316596 characters]\n",
      "text/plain": [
       "<Figure size 1080x1296 with 10 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "subplot_1D_signals(\n",
    "    S[:, 0:100], title=\"Original Signals\", figsize=(15.2, 9), colorcode=None\n",
    ")\n",
    "subplot_1D_signals(\n",
    "    X[:, 0:100], title=\"Mixture Signals\", figsize=(15, 18), colorcode=None\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "532b84e0",
   "metadata": {},
   "source": [
    "# Algorithm Hyperparameter Selection and Weight Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "19ba2f5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "lambday = 1 - 1e-1 / 10\n",
    "lambdae = 1 - 1e-1 / 10\n",
    "s_dim = S.shape[0]\n",
    "x_dim = X.shape[0]\n",
    "\n",
    "# Inverse output covariance\n",
    "By = 1 * np.eye(s_dim)\n",
    "# Inverse error covariance\n",
    "Be = 1000 * np.eye(s_dim)\n",
    "\n",
    "debug_iteration_point = 10000\n",
    "model = OnlineCorInfoMax(\n",
    "    s_dim=s_dim,\n",
    "    x_dim=x_dim,\n",
    "    muW=30 * 1e-3,\n",
    "    lambday=lambday,\n",
    "    lambdae=lambdae,\n",
    "    By=By,\n",
    "    Be=Be,\n",
    "    neural_OUTPUT_COMP_TOL=1e-6,\n",
    "    set_ground_truth=True,\n",
    "    S=S,\n",
    "    A=A,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3be8178",
   "metadata": {},
   "source": [
    "# Run CorInfoMax Algorithm on Mixture Signals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edcfe1fb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 609864 characters]\n",
      "text/plain": [
       "<Figure size 3600x2400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 36%|███▌      | 179960/500000 [02:46<03:24, 1563.46it/s]"
     ]
    }
   ],
   "source": [
    "model.fit_batch_sparse(\n",
    "    X=X,\n",
    "    n_epochs=1,\n",
    "    neural_dynamic_iterations=500,\n",
    "    plot_in_jupyter=True,\n",
    "    neural_lr_start=0.1,\n",
    "    neural_lr_stop=0.001,\n",
    "    debug_iteration_point=debug_iteration_point,\n",
    "    shuffle=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45062dfa",
   "metadata": {},
   "source": [
    "# Visualize SINR Convergence "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59ebb722",
   "metadata": {},
   "outputs": [],
   "source": [
    "mpl.rcParams[\"xtick.labelsize\"] = 18\n",
    "mpl.rcParams[\"ytick.labelsize\"] = 18\n",
    "plot_convergence_plot(\n",
    "    model.SIR_list,\n",
    "    xlabel=\"Number of Iterations / {}\".format(debug_iteration_point),\n",
    "    ylabel=\"SINR (dB)\",\n",
    "    title=\"SINR Convergence Plot\",\n",
    "    colorcode=None,\n",
    "    linewidth=1.8,\n",
    ")\n",
    "\n",
    "print(\"Final SIR: {}\".format(np.array(model.SIR_list[-1])))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c84d1f44",
   "metadata": {},
   "source": [
    "# Calculate Resulting Component SNRs and Overall SINR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a6e03b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "Wf = model.compute_overall_mapping(return_mapping=True)\n",
    "Y = Wf @ X\n",
    "Y_ = signed_and_permutation_corrected_sources(S, Y)\n",
    "coef_ = ((Y_ * S).sum(axis=1) / (Y_ * Y_).sum(axis=1)).reshape(-1, 1)\n",
    "Y_ = coef_ * Y_\n",
    "\n",
    "print(\"Component SNR Values : {}\\n\".format(snr_jit(S, Y_)))\n",
    "\n",
    "SINR = 10 * np.log10(CalculateSINRjit(Y_, S)[0])\n",
    "\n",
    "print(\"Overall SINR : {}\".format(SINR))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d57151ec",
   "metadata": {},
   "source": [
    "# Vizualize Extracted Signals Compared to Original Sources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6079c4e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "subplot_1D_signals(\n",
    "    Y_[:, 0:100],\n",
    "    title=\"Extracted Signals (Sign and Permutation Corrected)\",\n",
    "    figsize=(15.2, 9),\n",
    "    colorcode=None,\n",
    ")\n",
    "subplot_1D_signals(\n",
    "    S[:, 0:100], title=\"Original Signals\", figsize=(15.2, 9), colorcode=None\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
