add kernel fmri experiments

Author: csinva

fmri/analyze_fmri.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -20,7 +20,7 @@
     "import h5py\n",
     "from copy import deepcopy\n",
     "from skimage.filters import gabor_kernel\n",
-    "import gabor_feats\n",
+    "# import gabor_feats\n",
     "from sklearn.linear_model import RidgeCV\n",
     "import seaborn as sns\n",
     "from scipy.io import loadmat\n",
@@ -166,17 +166,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 168,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "[\"/roi/FFAlh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/FFArh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/IPlh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/IPrh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/MTlh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/MTplh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/MTprh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/MTrh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/OBJlh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/OBJrh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/PPAlh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/PPArh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/RSCrh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/STSrh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/VOlh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/VOrh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/latocclh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/latoccrh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/v1lh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/v1rh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/v2lh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/v2rh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/v3alh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/v3arh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/v3blh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/v3brh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/v3lh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/v3rh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/v4lh (EArray(18, 64, 64), zlib(3)) ''\", \"/roi/v4rh (EArray(18, 64, 64), zlib(3)) ''\"]\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "f = tables.open_file(oj(out_dir, 'VoxelResponses_subject1.mat'))\n",
     "xs = []\n",
@@ -514,13 +506,115 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# visualize features / decompositions\n",
+    "# kernel features"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 30,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "X = np.array(loadmat(oj(out_dir, 'mot_energy_feats_st.mat'))['S_fin'])\n",
+    "X_test = np.array(loadmat(oj(out_dir, 'mot_energy_feats_sv.mat'))['S_fin'])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from jax import random\n",
+    "from jax.experimental import stax\n",
+    "from jax import random\n",
+    "from neural_tangents import stax\n",
+    "\n",
+    "# kernel function\n",
+    "init_fn, apply_fn, kernel_fn = stax.serial(\n",
+    "    stax.Dense(512), stax.Relu(),\n",
+    "    stax.Dense(512), stax.Relu(),\n",
+    "    stax.Dense(1)\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# training kernel mat\n",
+    "kernel = kernel_fn(X, X, 'ntk')\n",
+    "fname = oj(out_dir, f'mot_energy_feats_kernel_mat_ntk.pkl')\n",
+    "if not os.path.exists(fname):\n",
+    "    save_pkl(kernel, fname)\n",
+    "    \n",
+    "# training kernel mat\n",
+    "kernel_test = kernel_fn(X_test, X, 'ntk')\n",
+    "fname = oj(out_dir, f'mot_energy_feats_kernel_test_with_train_ntk.pkl')\n",
+    "if not os.path.exists(fname):\n",
+    "    save_pkl(kernel_test, fname)\n",
+    "    \n",
+    "# save out eigenvals\n",
+    "fname = oj(out_dir, f'eigenvals_eigenvecs_mot_energy_kernel_ntk.pkl')\n",
+    "if not os.path.exists(fname):\n",
+    "    kernel = load_pkl(oj(out_dir, f'mot_energy_feats_kernel_mat_ntk.pkl'))\n",
+    "    eigenvals, eigenvecs = npl.eig(kernel)\n",
+    "    save_pkl((eigenvals, eigenvecs), fname)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      " 30%|███       | 6/20 [46:02<1:56:09, 497.85s/it]"
+     ]
+    }
+   ],
+   "source": [
+    "# save kernel pinvs\n",
+    "reg_params = np.logspace(3, 6, 20).round().astype(int)\n",
+    "kernel = load_pkl(oj(out_dir, f'mot_energy_feats_kernel_mat_ntk.pkl'))\n",
+    "for reg_param in tqdm(reg_params):\n",
+    "    fname = oj(out_dir, f'pinv_mot_energy_kernel_ntk_{reg_param}.pkl')\n",
+    "    if not os.path.exists(fname):\n",
+    "        inv = npl.pinv(kernel + reg_param * np.eye(kernel.shape[0]))\n",
+    "        save_pkl(inv, fname)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# need to save kernel matrix\n",
+    "# need to save test-time kernel mat\n",
+    "\n",
+    "\n",
+    "# save eigenvalues\n",
+    "\n",
+    "# make new script\n",
+    "# need to switch to use Kernel ridge + eigenvalues"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# visualize preprocessed features\n",
     "**load and look at features**"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -579,7 +673,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.5"
+   "version": "3.8.3"
   }
  },
  "nbformat": 4,

fmri/analyze_fmri_kernel.ipynb

@@ -39,20 +39,24 @@ def load_h5(fname):
     f.close()
     return data
 
+def load_pkl(fname):
+    return pkl.load(open(fname, "rb" ))
+
 def save_pkl(d, fname):
     if os.path.exists(fname):
         os.remove(fname)
     with open(fname, 'wb') as f:
         pkl.dump(d, f)
 
-def get_roi_and_idx(run):
+def get_roi_and_idx(run, out_dir, sigmas):
     # select roi + i (which is the roi_idx)
     rois = ['v1lh', 'v2lh', 'v4lh', 'v1rh', 'v2rh', 'v4rh']
     roi = rois[run % len(rois)]
 
     f = tables.open_file(oj(out_dir, 'VoxelResponses_subject1.mat'), 'r')
     roi_idxs_all = f.get_node(f'/roi/{roi}')[:].flatten().nonzero()[0] # structure containing volume matrices (64x64x18) with indices corresponding to each roi in each hemisphere
-    roi_idxs = np.array([roi_idx for roi_idx in roi_idxs_all if ~np.isnan(sigmas[roi_idx])])
+    roi_idxs = np.array([roi_idx for roi_idx in roi_idxs_all
+                         if ~np.isnan(sigmas[roi_idx])])
 
     i = roi_idxs[run // len(rois)] # i is the roi idx
     return roi, i  
@@ -77,6 +81,7 @@ def get_roi_and_idx(run):
     save_dir = oj(out_dir, 'dec14_baselines_ard')
     suffix = '_feats' # _feats, '' for pixels
     norm = '_norm' # ''
+    reg_params = np.logspace(3, 6, 20).round().astype(int) # reg values to try (must match preprocess_fmri)
     print('saving to', save_dir)
 
 
@@ -97,6 +102,14 @@ def get_roi_and_idx(run):
     else:
         X_train = np.array(loadmat(oj(out_dir, 'mot_energy_feats_st.mat'))['S_fin'])
     X_test = np.array(loadmat(oj(out_dir, 'mot_energy_feats_sv.mat'))['S_fin'])
+    if use_small:
+#         (U, alphas, _) = pkl.load(open(oj(out_dir, f'decomp_mot_energy_small.pkl'), 'rb'))
+        (eigenvals, eigenvecs) = pkl.load(open(oj(out_dir, f'eigenvals_eigenvecs_mot_energy_small.pkl'), 'rb'))        
+        Y_train = Y_train[:, :720]
+    else:
+#         (U, alphas, _) = pkl.load(open(oj(out_dir, f'decomp_mot_energy.pkl'), 'rb'))
+        (eigenvals, eigenvecs) = pkl.load(open(oj(out_dir, f'eigenvals_eigenvecs_mot_energy.pkl'), 'rb'))
+    
 
     '''
     # load the raw responses
@@ -108,18 +121,12 @@ def get_roi_and_idx(run):
     Y_train = load_h5(oj(out_dir, 'rt_norm.h5')) # training responses: 73728 (voxels) x 7200 (timepoints)    
     Y_test = load_h5(oj(out_dir, 'rv_norm.h5') )
     sigmas = load_h5(oj(out_dir, f'out_rva_sigmas_norm.h5')) # stddev across repeats
-    if use_small:
-#         (U, alphas, _) = pkl.load(open(oj(out_dir, f'decomp_mot_energy_small.pkl'), 'rb'))
-        (eigenvals, eigenvecs) = pkl.load(open(oj(out_dir, f'eigenvals_eigenvecs_mot_energy_small.pkl'), 'rb'))        
-        Y_train = Y_train[:, :720]
-    else:
-#         (U, alphas, _) = pkl.load(open(oj(out_dir, f'decomp_mot_energy.pkl'), 'rb'))
-        (eigenvals, eigenvecs) = pkl.load(open(oj(out_dir, f'eigenvals_eigenvecs_mot_energy.pkl'), 'rb'))
+
 
 
     # loop over individual neurons
     for run in runs:
-        roi, i = get_roi_and_idx(run)
+        roi, i = get_roi_and_idx(run, out_dir, sigmas)
         results = {}
         os.makedirs(save_dir, exist_ok=True)
         print('fitting', roi, 'idx', i)
@@ -143,9 +150,6 @@ def get_roi_and_idx(run):
             print('\tskipping this voxel!')
             continue
 
-        # reg values to try
-        reg_params = np.logspace(3, 6, 20).round().astype(int)
-        
         # fit ard + mdl-rs
         baselines = {}
         for model_type, model_name in zip([ARDRegression], ['ard']):
@@ -245,4 +249,4 @@ def get_roi_and_idx(run):
             **baselines,
         }
         pkl.dump(results, open(oj(save_dir, f'ridge_{i}.pkl'), 'wb'))
-        print('\tdone!')
+        print(f'\tsuccesfully finished run {run}!')