examples/funcalign/rsrm_synthetic_reconstruction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## This example replicates the synthetic experiment in the paper\n",
    "## \"Capturing Shared and Individual Information in fMRI Data\", J. Turek, C. Ellis, L. Skalaban, N. Turk-Browne, T. Willke"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import some libraries that we will need"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import scipy.io\n",
    "from scipy.stats import stats\n",
    "import numpy as np\n",
    "import brainiak.funcalign.srm as srm\n",
    "import brainiak.funcalign.rsrm as rsrm\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Experiment setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "voxels = 100\n",
    "samples = 200\n",
    "subjects = 10\n",
    "features = 3\n",
    "snr = 20 # in dB\n",
    "amplitude = 8.0 # approximate amplitude used for the values in Si\n",
    "p = 0.05 # probability of being non-zero in Si"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we create some synthetic data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a Shared response R with K = 3\n",
    "theta = np.linspace(-4 * np.pi, 4 * np.pi, samples)\n",
    "z = np.linspace(-2, 2, samples)\n",
    "r = z ** 2 + 1\n",
    "x = r * np.sin(theta)\n",
    "y = r * np.cos(theta)\n",
    "curve = np.vstack((x, y, z))\n",
    "print('Curve max, min values:', np.max(curve), np.min(curve))\n",
    "\n",
    "# Create the subjects' data\n",
    "data = [None] * subjects\n",
    "W = [None] * subjects\n",
    "noise_level = 0.0\n",
    "for s in range(subjects):\n",
    "    R = curve\n",
    "    W[s], _ = np.linalg.qr(np.random.randn(voxels, 3))\n",
    "    data[s] = W[s].dot(R)\n",
    "    noise_level += np.sum(np.abs(data[s])**2)\n",
    "\n",
    "# Compute noise_sigma from desired SNR\n",
    "noise_level = noise_level / (10 ** (snr / 10))\n",
    "noise_level = np.sqrt(noise_level  / subjects / voxels / samples)\n",
    "\n",
    "for s in range(subjects):\n",
    "    n = noise_level * np.random.randn(voxels, samples)\n",
    "    data[s] += n\n",
    "    \n",
    "S = [None] * subjects\n",
    "for s in range(subjects):\n",
    "    S[s] = (np.random.rand(data[s].shape[0], samples) < p) * ((np.random.rand(data[s].shape[0], samples) * amplitude - amplitude/2) )\n",
    "    data[s] += S[s]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we fit the algorithms, SRM and RSRM, to the synthetic data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "algo_srm = srm.SRM(features=3, n_iter=20)\n",
    "algo_srm.fit(data)\n",
    "\n",
    "\n",
    "algo_rsrm = rsrm.RSRM(features=3, gamma=0.35, n_iter=20)\n",
    "algo_rsrm.fit(data)\n",
    "\n",
    "print('Done')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following function finds the orthogonal transform to align the shared response to the original curve."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_orthogonal_transform(shared_response, curve):\n",
    "    u,_,vt = np.linalg.svd(shared_response.dot(curve.T))\n",
    "    q = u.dot(vt)\n",
    "    aligned_curve = q.T.dot(shared_response)\n",
    "    return aligned_curve"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.gca(projection='3d')\n",
    "ax.plot(curve[0, :], curve[1, :], curve[2, :], '-g', label='original', lineWidth=5)\n",
    "\n",
    "proj = find_orthogonal_transform(algo_srm.s_, curve)\n",
    "ax.plot(proj[0, :], proj[1, :], proj[2, :], '-b', label='SRM', lineWidth=3)\n",
    "proj = find_orthogonal_transform(algo_rsrm.r_, curve)\n",
    "ax.plot(proj[0, :], proj[1, :], proj[2, :], '-r', label='RSRM', lineWidth=3)\n",
    "plt.legend()"
   ]
  }
 ],
 "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.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}