{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a25c2be1",
   "metadata": {},
   "source": [
    "# Import Required Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "96d233ea",
   "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 LDMIBSS 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_Antisparse_Copula\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d0fcdf8",
   "metadata": {},
   "source": [
    "# Source Generation and Mixing Scenario"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f347c550",
   "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.679 & 0.675 & 0.664 & 0.667 & \\\\\n",
       " 0.679 & 1.000 & 0.671 & 0.669 & 0.672 & \\\\\n",
       " 0.675 & 0.671 & 1.000 & 0.668 & 0.664 & \\\\\n",
       " 0.664 & 0.669 & 0.668 & 1.000 & 0.663 & \\\\\n",
       " 0.667 & 0.672 & 0.664 & 0.663 & 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.527 & -0.300 & 2.984 & 1.358 & 1.430 & \\\\\n",
       " -1.016 & -0.553 & -0.611 & -1.812 & -1.394 & \\\\\n",
       " -0.648 & 0.151 & -1.119 & 1.254 & -0.979 & \\\\\n",
       " -0.031 & 2.301 & 0.590 & 1.426 & 0.344 & \\\\\n",
       " 0.457 & 0.813 & 1.170 & 1.047 & -1.027 & \\\\\n",
       " -0.027 & 0.497 & 1.102 & -0.285 & 0.146 & \\\\\n",
       " 0.634 & 0.994 & -1.043 & 0.192 & -0.202 & \\\\\n",
       " -0.577 & -0.443 & 1.760 & -0.536 & 0.041 & \\\\\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.984126430541846\n"
     ]
    }
   ],
   "source": [
    "N = 10000\n",
    "NumberofSources = 5\n",
    "NumberofMixtures = 8\n",
    "S = generate_correlated_copula_sources(\n",
    "    rho=0.7,\n",
    "    df=4,\n",
    "    n_sources=NumberofSources,\n",
    "    size_sources=N,\n",
    "    decreasing_correlation=False,\n",
    ")\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",
    "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": "4f161ed1",
   "metadata": {},
   "source": [
    "# Visualize Generated Sources and Mixtures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7fac953e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 172708 characters]\n",
      "text/plain": [
       "<Figure size 1094.4x648 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 282268 characters]\n",
      "text/plain": [
       "<Figure size 1080x1296 with 8 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": "08db0906",
   "metadata": {},
   "source": [
    "# Algorithm Hyperparameter Selection and Weight Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "83583247",
   "metadata": {},
   "outputs": [],
   "source": [
    "s_dim = S.shape[0]\n",
    "x_dim = X.shape[0]\n",
    "debug_iteration_point = 1000\n",
    "model = LDMIBSS(s_dim=s_dim, x_dim=x_dim, set_ground_truth=True, S=S, A=A)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e93d0f26",
   "metadata": {},
   "source": [
    "# Run LDMI Algorithm on Mixture Signals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bb6210bf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 638784 characters]\n",
      "text/plain": [
       "<Figure size 3600x2400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10000/10000 [00:35<00:00, 282.49it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 638784 characters]\n",
      "text/plain": [
       "<Figure size 3600x2400 with 4 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "model.fit_batch_nnantisparse(\n",
    "    X,\n",
    "    epsilon=1e-5,\n",
    "    mu_start=100,\n",
    "    n_iterations=10000,\n",
    "    method=\"correlation\",\n",
    "    debug_iteration_point=debug_iteration_point,\n",
    "    plot_in_jupyter=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47b79106",
   "metadata": {},
   "source": [
    "# Calculate Resulting Component SNRs and Overall SINR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "072a1a94",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Component SNR Values : [24.77427883 27.5844347  32.69580687 28.27487838 25.4177908 ]\n",
      "\n",
      "Overall SINR : 27.051022977290764\n"
     ]
    }
   ],
   "source": [
    "Wf = model.W\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",
    "SINRwsm = 10 * np.log10(CalculateSINRjit(Y_, S)[0])\n",
    "\n",
    "print(\"Overall SINR : {}\".format(SINRwsm))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99e8c11e",
   "metadata": {},
   "source": [
    "# Vizualize Extracted Signals Compared to Original Sources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a10fd9e9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 174408 characters]\n",
      "text/plain": [
       "<Figure size 1094.4x648 with 5 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 172708 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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
