{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4044f211",
   "metadata": {},
   "source": [
    "# Import Required Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ed297f83",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"../../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",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "notebook_name = 'Nonnegative_Antisparse_Copula'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ee1f394",
   "metadata": {},
   "source": [
    "# Source Generation and Mixing Scenario"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "222a48b4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The following is the correlation matrix of sources\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{bmatrix} \n",
       " 1.000 & 0.773 & 0.772 & 0.772 & 0.772 & \\\\\n",
       " 0.773 & 1.000 & 0.773 & 0.772 & 0.773 & \\\\\n",
       " 0.772 & 0.773 & 1.000 & 0.772 & 0.773 & \\\\\n",
       " 0.772 & 0.772 & 0.772 & 1.000 & 0.773 & \\\\\n",
       " 0.772 & 0.773 & 0.773 & 0.773 & 1.000 & \\\\\n",
       "\\end{bmatrix}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The following is the mixture matrix A\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{bmatrix} \n",
       " -0.309 & 0.610 & 0.941 & 0.750 & -0.761 & \\\\\n",
       " 0.600 & -1.925 & -1.425 & 0.165 & -0.185 & \\\\\n",
       " 1.637 & -0.801 & 0.550 & -1.000 & -1.384 & \\\\\n",
       " 0.290 & 0.302 & -0.203 & -0.992 & -0.148 & \\\\\n",
       " 0.269 & -0.700 & -0.076 & 1.333 & 1.234 & \\\\\n",
       " -0.477 & -0.166 & 0.599 & 1.010 & -1.226 & \\\\\n",
       " -0.998 & 1.013 & 0.012 & 0.695 & 1.987 & \\\\\n",
       " -0.709 & 0.530 & -0.513 & 1.664 & -1.326 & \\\\\n",
       " -1.077 & -0.244 & -0.939 & -0.745 & 0.766 & \\\\\n",
       " -0.319 & -0.784 & -0.602 & 1.918 & 0.024 & \\\\\n",
       "\\end{bmatrix}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input SNR is : 30.00507557757993\n"
     ]
    }
   ],
   "source": [
    "N = 500000\n",
    "NumberofSources = 5\n",
    "NumberofMixtures = 10\n",
    "S = generate_correlated_copula_sources(rho = 0.8, df = 4, n_sources = NumberofSources, \n",
    "                                       size_sources = N , decreasing_correlation = False)\n",
    "\n",
    "print(\"The following is the correlation matrix of sources\")\n",
    "display_matrix(np.corrcoef(S))\n",
    "\n",
    "# Generate Mxr random mixing from i.i.d N(0,1)\n",
    "# A = np.random.randn(NumberofMixtures,NumberofSources)\n",
    "A = np.random.normal(0, 1, (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(np.sum(np.mean((X - NoisePart)**2, axis = 1)) / np.sum(np.mean(NoisePart**2, axis = 1)))\n",
    "print(\"The following is the mixture matrix A\")\n",
    "display_matrix(A)\n",
    "print(\"Input SNR is : {}\".format(SNRinp))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d498be6c",
   "metadata": {},
   "source": [
    "# Visualize Generated Sources and Mixtures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a496a13e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 168900 characters]\n",
      "text/plain": [
       "<Figure size 1094.4x648 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 294160 characters]\n",
      "text/plain": [
       "<Figure size 1080x1296 with 10 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "subplot_1D_signals(S[:,0:100], title = 'Original Signals', figsize = (15.2,9), colorcode = None)\n",
    "subplot_1D_signals(X[:,0:100], title = 'Mixture Signals', figsize = (15,18), colorcode = None)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e093506",
   "metadata": {},
   "source": [
    "# Algorithm Hyperparameter Selection and Weight Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "34669807",
   "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 = 10000 * np.eye(s_dim)\n",
    "\n",
    "debug_iteration_point = 25000\n",
    "model = OnlineCorInfoMax(s_dim = s_dim, x_dim = x_dim, muW = 30*1e-3, lambday = lambday,\n",
    "                        lambdae = lambdae, By = By, Be = Be, neural_OUTPUT_COMP_TOL = 1e-6,\n",
    "                        set_ground_truth = True, S = S, A = A)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d18d9462",
   "metadata": {},
   "source": [
    "# Run CorInfoMax Algorithm on Mixture Signals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "09e54b6e",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 811700 characters]\n",
      "text/plain": [
       "<Figure size 3600x2400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500000/500000 [02:34<00:00, 3241.30it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Algorithm took 154.445988 sec.\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 811700 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_nnantisparse(X = X, n_epochs = 1, neural_dynamic_iterations = 500,\n",
    "                               plot_in_jupyter = True, neural_lr_start = 0.9,\n",
    "                               neural_lr_stop = 0.001, debug_iteration_point = debug_iteration_point, \n",
    "                               shuffle = False)\n",
    "print('Algorithm took %f sec.' % (t.interval))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2034571",
   "metadata": {},
   "source": [
    "# Visualize SINR Convergence "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "dd151425",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Final SINR: 25.327796351348226\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 37964 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(model.SIR_list, xlabel = 'Number of Iterations / {}'.format(debug_iteration_point),\n",
    "                     ylabel = 'SINR (dB)', title = 'SINR Convergence Plot', colorcode = None, linewidth = 1.8)\n",
    "\n",
    "print('Final SINR: {}'.format(np.array(model.SIR_list[-1])))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "334d45cc",
   "metadata": {},
   "source": [
    "# Calculate Resulting Component SNRs and Overall SINR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8beebe39",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Component SNR Values : [24.5537566  25.70636042 25.26165618 24.23142037 25.62020405]\n",
      "\n",
      "Overall SINR : 25.327796351348226\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": "46fd45df",
   "metadata": {},
   "source": [
    "# Vizualize Extracted Signals Compared to Original Sources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6f8ca4cb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 174176 characters]\n",
      "text/plain": [
       "<Figure size 1094.4x648 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 171012 characters]\n",
      "text/plain": [
       "<Figure size 1094.4x648 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "subplot_1D_signals(Y_[:,0:100], title = 'Extracted Signals (Sign and Permutation Corrected)', figsize = (15.2,9), colorcode = None)\n",
    "subplot_1D_signals(S[:,0:100], title = 'Original Signals', figsize = (15.2,9), colorcode = None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce75389a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}
