{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b97ff966",
   "metadata": {},
   "source": [
    "# Import Required Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6d810cfb",
   "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 = \"Nonnegative_Sparse\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fc2bd2b",
   "metadata": {},
   "source": [
    "# Source Generation and Mixing Scenario"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "36c134d9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The following is the mixture matrix A\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{bmatrix} \n",
       " -1.732 & 0.414 & -1.329 & -0.933 & 1.609 & \\\\\n",
       " 0.161 & -1.356 & 0.838 & -1.254 & -0.852 & \\\\\n",
       " 0.720 & -0.617 & -0.013 & -0.162 & -1.035 & \\\\\n",
       " 0.450 & 0.884 & -0.536 & 0.041 & -1.058 & \\\\\n",
       " -0.132 & 0.065 & -1.500 & 0.936 & 1.871 & \\\\\n",
       " 0.828 & -0.311 & -0.696 & -1.469 & -0.493 & \\\\\n",
       " 0.832 & -0.321 & 0.888 & -1.014 & 0.232 & \\\\\n",
       " 1.083 & -0.172 & -0.757 & 0.198 & -0.350 & \\\\\n",
       " -0.085 & 2.317 & -0.375 & 0.765 & -0.018 & \\\\\n",
       " -0.539 & -2.031 & -0.655 & 0.403 & -2.392 & \\\\\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.993621849382706\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 8532 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=False,\n",
    ")\n",
    "S = 4 * S - 2\n",
    "S = ProjectRowstoL1NormBall(S.T).T\n",
    "S = S * (S >= 0)\n",
    "# print(\"The following is the correlation matrix of sources\")\n",
    "# display_matrix(np.corrcoef(S))\n",
    "plt.scatter(S[0, :], S[2, :])\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": "2a487ee9",
   "metadata": {},
   "source": [
    "# Visualize Generated Sources and Mixtures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "958718c4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 106004 characters]\n",
      "text/plain": [
       "<Figure size 1094.4x648 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 288476 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": "e496374f",
   "metadata": {},
   "source": [
    "# Algorithm Hyperparameter Selection and Weight Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c588218f",
   "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 = 5 * 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": "672c0ab4",
   "metadata": {},
   "source": [
    "# Run CorInfoMax Algorithm on Mixture Signals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7796089b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 555892 characters]\n",
      "text/plain": [
       "<Figure size 3600x2400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500000/500000 [04:40<00:00, 1780.18it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Algorithm took 281.031451 sec.\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 555892 characters]\n",
      "text/plain": [
       "<Figure size 3600x2400 with 4 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "with Timer() as t:\n",
    "    model.fit_batch_nnsparse(\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=True,\n",
    "    )\n",
    "print(\"Algorithm took %f sec.\" % (t.interval))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73569aa4",
   "metadata": {},
   "source": [
    "# Visualize SINR Convergence "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "fb62e36c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Final SIR: 30.43970710175448\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 28664 characters]\n",
      "text/plain": [
       "<Figure size 864x576 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "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": "55da49ea",
   "metadata": {},
   "source": [
    "# Calculate Resulting Component SNRs and Overall SINR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4577761b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Component SNR Values : [28.58159351 31.34289969 30.45133374 29.98303976 31.98815792]\n",
      "\n",
      "Overall SINR : 30.43970710175448\n"
     ]
    }
   ],
   "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": "538e289b",
   "metadata": {},
   "source": [
    "# Vizualize Extracted Signals Compared to Original Sources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "144f6593",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 126508 characters]\n",
      "text/plain": [
       "<Figure size 1094.4x648 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 107192 characters]\n",
      "text/plain": [
       "<Figure size 1094.4x648 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "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
}
