{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import random\n",
    "import nibabel as nb\n",
    "import numpy as np\n",
    "from scipy.sparse.linalg import svds, eigs\n",
    "import matplotlib.pyplot as plt\n",
    "from nilearn import plotting\n",
    "from nilearn import image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "in_list = [nb.load(f) for f in glob.glob('data/cobre/fmri_*.nii.gz')]\n",
    "in_mask = nb.load('data/brain_mask.nii.gz')\n",
    "go = '3DMIGP'\n",
    "dPCA_int = 550\n",
    "dPCA_out = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random.shuffle(in_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mask = in_mask.get_fdata().ravel()\n",
    "\n",
    "def demean(x):\n",
    "    return x - np.mean(x, axis=0)\n",
    "\n",
    "for i, f in enumerate(in_list):\n",
    "    \n",
    "    print(i, end=' ')\n",
    "    \n",
    "    # read data\n",
    "    grot = f.get_fdata()\n",
    "    grot = np.reshape(grot, [-1, grot.shape[-1]])\n",
    "    grot = grot[mask!=0, :].T\n",
    "    grot = demean(grot)\n",
    "    \n",
    "    # var-norm\n",
    "    [uu, ss, vt] = svds(grot, k=30)\n",
    "    vt[np.abs(vt) < (2.3 * np.std(vt))] = 0;\n",
    "    stddevs = np.maximum(np.std(grot - (uu @ np.diag(ss) @ vt), axis=0), 0.001)\n",
    "    grot = grot/stddevs\n",
    "    \n",
    "    if i == 0:\n",
    "        W = demean(grot)\n",
    "    else:\n",
    "        # concat\n",
    "        W = np.concatenate((W, demean(grot)), axis=0)\n",
    "        \n",
    "        # reduce W to dPCA_int eigenvectors\n",
    "        if W.shape[0]-10 > dPCA_int:\n",
    "            uu = eigs(W@W.T, dPCA_int)[1]\n",
    "            uu = np.real(uu)\n",
    "            W = uu.T @ W\n",
    "        \n",
    "    f.uncache()\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grot = np.zeros([mask.shape[0], dPCA_out])\n",
    "grot[mask!=0, :] = W[:dPCA_out, :].T\n",
    "grot = np.reshape(grot, in_list[0].shape[:3] + (dPCA_out,))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nb.Nifti1Image(grot, affine=in_list[0].affine, header=in_list[0].header).to_filename('pyMIGP.nii.gz')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "melodic -i pyMIGP.nii.gz \\\n",
    "\t--mask=data/brain_mask.nii.gz \\\n",
    "\t-d 20 \\\n",
    "\t-v \\\n",
    "\t--nobet \\\n",
    "\t--disableMigp \\\n",
    "\t--varnorm \\\n",
    "\t-o pymigp_dim20.ica"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ics = nb.load('pymigp_dim20.ica/melodic_IC.nii.gz')\n",
    "\n",
    "N = ics.shape[-1]\n",
    "\n",
    "fig, ax = plt.subplots(int(np.ceil((N/2))),2, figsize=(12, N))\n",
    "\n",
    "for img, ax0 in zip(image.iter_img(ics), ax.ravel()):\n",
    "    coord = plotting.find_xyz_cut_coords(img, activation_threshold=1.5)\n",
    "    plotting.plot_stat_map(img, cut_coords=coord, vmax=5, axes=ax0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}