predictions_master_script.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nltk\n",
    "import numpy as np\n",
    "from scipy.io import loadmat\n",
    "from sklearn.neural_network import MLPRegressor\n",
    "from sklearn.linear_model import Lasso, Ridge\n",
    "from sklearn.model_selection import cross_val_score\n",
    "from sklearn.metrics import r2_score\n",
    "from scipy import stats\n",
    "from sklearn.decomposition import PCA\n",
    "from IPython.core.debugger import set_trace\n",
    "import pickle\n",
    "import os\n",
    "from joblib import Parallel, delayed\n",
    "\n",
    "results_save_dir = \"predictions\"\n",
    "features_dir = \"features\"\n",
    "sub_data = \"sub_space_data/\"\n",
    "\n",
    "if not os.path.exists(results_save_dir):\n",
    "    os.mkdir(results_save_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Functions to estimate cost for each lambda, by voxel:\n",
    "from __future__ import division                                              \n",
    "\n",
    "from numpy.linalg import inv, svd\n",
    "import numpy as np\n",
    "from sklearn.model_selection import KFold\n",
    "from sklearn.linear_model import Ridge, RidgeCV\n",
    "import time \n",
    "from scipy.stats import zscore\n",
    "\n",
    "\n",
    "def corr(X,Y):\n",
    "    return np.mean(zscore(X)*zscore(Y),0)\n",
    "\n",
    "def R2(Pred,Real):\n",
    "    SSres = np.mean((Real-Pred)**2,0)\n",
    "    SStot = np.var(Real,0)\n",
    "    return np.nan_to_num(1-SSres/SStot)\n",
    "\n",
    "def R2r(Pred,Real):\n",
    "    R2rs = R2(Pred,Real)\n",
    "    ind_neg = R2rs<0\n",
    "    R2rs = np.abs(R2rs)\n",
    "    R2rs = np.sqrt(R2rs)\n",
    "    R2rs[ind_neg] *= - 1\n",
    "    return R2rs\n",
    "\n",
    "def ridge(X,Y,lmbda):\n",
    "    return np.dot(inv(X.T.dot(X)+lmbda*np.eye(X.shape[1])),X.T.dot(Y))\n",
    "\n",
    "def ridge_by_lambda(X, Y, Xval, Yval, lambdas=np.array([0.1,1,10,100,1000])):\n",
    "    error = np.zeros((lambdas.shape[0],Y.shape[1]))\n",
    "    for idx,lmbda in enumerate(lambdas):\n",
    "        weights = ridge(X,Y,lmbda)\n",
    "        error[idx] = 1 - R2(np.dot(Xval,weights),Yval)\n",
    "    return error\n",
    "\n",
    "def ridge_sk(X,Y,lmbda):\n",
    "    rd = Ridge(alpha = lmbda)\n",
    "    rd.fit(X,Y)\n",
    "    return rd.coef_.T\n",
    "\n",
    "def ridgeCV_sk(X,Y,lmbdas):\n",
    "    rd = RidgeCV(alphas = lmbdas,solver = 'svd')\n",
    "    rd.fit(X,Y)\n",
    "    return rd.coef_.T\n",
    "\n",
    "def ridge_by_lambda_sk(X, Y, Xval, Yval, lambdas=np.array([0.1,1,10,100,1000])):\n",
    "    error = np.zeros((lambdas.shape[0],Y.shape[1]))\n",
    "    for idx,lmbda in enumerate(lambdas):\n",
    "        weights = ridge_sk(X,Y,lmbda)\n",
    "        error[idx] = 1 -  R2(np.dot(Xval,weights),Yval)\n",
    "    return error\n",
    "\n",
    "def ridge_svd(X,Y,lmbda):\n",
    "    U, s, Vt = svd(X, full_matrices=False)\n",
    "    d = s / (s** 2 + lmbda)\n",
    "    return np.dot(Vt,np.diag(d).dot(U.T.dot(Y)))\n",
    "\n",
    "def ridge_by_lambda_svd(X, Y, Xval, Yval, lambdas=np.array([0.1,1,10,100,1000])):\n",
    "    error = np.zeros((lambdas.shape[0],Y.shape[1]))\n",
    "    U, s, Vt = svd(X, full_matrices=False)\n",
    "    for idx,lmbda in enumerate(lambdas):\n",
    "        d = s / (s** 2 + lmbda)\n",
    "        weights = np.dot(Vt,np.diag(d).dot(U.T.dot(Y)))\n",
    "        error[idx] = 1 - R2(np.dot(Xval,weights),Yval)\n",
    "    return error\n",
    "\n",
    "\n",
    "def kernel_ridge(X,Y,lmbda):\n",
    "    return np.dot(X.T.dot(inv(X.dot(X.T)+lmbda*np.eye(X.shape[0]))),Y)\n",
    "\n",
    "def kernel_ridge_by_lambda(X, Y, Xval, Yval, lambdas=np.array([0.1,1,10,100,1000])):\n",
    "    error = np.zeros((lambdas.shape[0],Y.shape[1]))\n",
    "    for idx,lmbda in enumerate(lambdas):\n",
    "        weights = kernel_ridge(X,Y,lmbda)\n",
    "        error[idx] = 1 - R2(np.dot(Xval,weights),Yval)\n",
    "    return error\n",
    "\n",
    "\n",
    "def kernel_ridge_svd(X,Y,lmbda):\n",
    "    U, s, Vt = svd(X.T, full_matrices=False)\n",
    "    d = s / (s** 2 + lmbda)\n",
    "    return np.dot(np.dot(U,np.diag(d).dot(Vt)),Y)\n",
    "\n",
    "def kernel_ridge_by_lambda_svd(X, Y, Xval, Yval, lambdas=np.array([0.1,1,10,100,1000])):\n",
    "    error = np.zeros((lambdas.shape[0],Y.shape[1]))\n",
    "    U, s, Vt = svd(X.T, full_matrices=False)\n",
    "    for idx,lmbda in enumerate(lambdas):\n",
    "        d = s / (s** 2 + lmbda)\n",
    "        weights = np.dot(np.dot(U,np.diag(d).dot(Vt)),Y)\n",
    "        error[idx] = 1 - R2(np.dot(Xval,weights),Yval)\n",
    "    return error\n",
    "\n",
    "\n",
    "def cross_val_ridge(train_features,train_data, n_splits = 10, \n",
    "                    lambdas = np.array([10**i for i in range(-6,10)]),\n",
    "                    method = 'plain',\n",
    "                    do_plot = False):\n",
    "    \n",
    "    ridge_1 = dict(plain = ridge_by_lambda,\n",
    "                   svd = ridge_by_lambda_svd,\n",
    "                   kernel_ridge = kernel_ridge_by_lambda,\n",
    "                   kernel_ridge_svd = kernel_ridge_by_lambda_svd,\n",
    "                   ridge_sk = ridge_by_lambda_sk)[method]\n",
    "    ridge_2 = dict(plain = ridge,\n",
    "                   svd = ridge_svd,\n",
    "                   kernel_ridge = kernel_ridge,\n",
    "                   kernel_ridge_svd = kernel_ridge_svd,\n",
    "                   ridge_sk = ridge_sk)[method]\n",
    "    \n",
    "    n_voxels = train_data.shape[1]\n",
    "    nL = lambdas.shape[0]\n",
    "    r_cv = np.zeros((nL, train_data.shape[1]))\n",
    "\n",
    "    kf = KFold(n_splits=n_splits)\n",
    "    start_t = time.time()\n",
    "    for icv, (trn, val) in enumerate(kf.split(train_data)):\n",
    "#         print('ntrain = {}'.format(train_features[trn].shape[0]))\n",
    "        cost = ridge_1(train_features[trn],train_data[trn],\n",
    "                               train_features[val],train_data[val], \n",
    "                               lambdas=lambdas)\n",
    "        if do_plot:\n",
    "            import matplotlib.pyplot as plt\n",
    "            plt.figure()\n",
    "            plt.imshow(cost,aspect = 'auto')\n",
    "        r_cv += cost\n",
    "#         if icv%3 ==0:\n",
    "#             print(icv)\n",
    "#         print('average iteration length {}'.format((time.time()-start_t)/(icv+1)))\n",
    "    if do_plot:\n",
    "        plt.figure()\n",
    "        plt.imshow(r_cv,aspect='auto',cmap = 'RdBu_r');\n",
    "\n",
    "    argmin_lambda = np.argmin(r_cv,axis = 0)\n",
    "    weights = np.zeros((train_features.shape[1],train_data.shape[1]))\n",
    "    for idx_lambda in range(lambdas.shape[0]): # this is much faster than iterating over voxels!\n",
    "        idx_vox = argmin_lambda == idx_lambda\n",
    "        weights[:,idx_vox] = ridge_2(train_features, train_data[:,idx_vox],lambdas[idx_lambda])\n",
    "    if do_plot:\n",
    "        plt.figure()\n",
    "        plt.imshow(weights,aspect='auto',cmap = 'RdBu_r',vmin = -0.5,vmax = 0.5);\n",
    "\n",
    "    return weights, np.array([lambdas[i] for i in argmin_lambda])\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "def GCV_ridge(train_features,train_data,lambdas = np.array([10**i for i in range(-6,10)])):\n",
    "    \n",
    "    n_lambdas = lambdas.shape[0]\n",
    "    n_voxels = train_data.shape[1]\n",
    "    n_time = train_data.shape[0]\n",
    "    n_p = train_features.shape[1]\n",
    "\n",
    "    CVerr = np.zeros((n_lambdas, n_voxels))\n",
    "\n",
    "    # % If we do an eigendecomp first we can quickly compute the inverse for many different values\n",
    "    # % of lambda. SVD uses X = UDV' form.\n",
    "    # % First compute K0 = (X'X + lambda*I) where lambda = 0.\n",
    "    #K0 = np.dot(train_features,train_features.T)\n",
    "    print('Running svd',)\n",
    "    start_time = time.time()\n",
    "    [U,D,Vt] = svd(train_features,full_matrices=False)\n",
    "    V = Vt.T\n",
    "    print(U.shape,D.shape,Vt.shape)\n",
    "    print('svd time: {}'.format(time.time() - start_time))\n",
    "\n",
    "    for i,regularizationParam in enumerate(lambdas):\n",
    "        regularizationParam = lambdas[i]\n",
    "        print('CVLoop: Testing regularization param: {}'.format(regularizationParam))\n",
    "\n",
    "        #Now we can obtain Kinv for any lambda doing Kinv = V * (D + lambda*I)^-1 U'\n",
    "        dlambda = D**2 + np.eye(n_p)*regularizationParam\n",
    "        dlambdaInv = np.diag(D / np.diag(dlambda))\n",
    "        KlambdaInv = V.dot(dlambdaInv).dot(U.T)\n",
    "        \n",
    "        # Compute S matrix of Hastie Trick  H = X(XT X + lambdaI)-1XT\n",
    "        S = np.dot(U, np.diag(D * np.diag(dlambdaInv))).dot(U.T)\n",
    "        denum = 1-np.trace(S)/n_time\n",
    "        \n",
    "        # Solve for weight matrix so we can compute residual\n",
    "        weightMatrix = KlambdaInv.dot(train_data);\n",
    "\n",
    "\n",
    "#         Snorm = np.tile(1 - np.diag(S) , (n_voxels, 1)).T\n",
    "        YdiffMat = (train_data - (train_features.dot(weightMatrix)));\n",
    "        YdiffMat = YdiffMat / denum;\n",
    "        CVerr[i,:] = (1/n_time)*np.sum(YdiffMat * YdiffMat,0);\n",
    "\n",
    "\n",
    "    # try using min of avg err\n",
    "    minerrIndex = np.argmin(CVerr,axis = 0);\n",
    "    r=np.zeros((n_voxels));\n",
    "\n",
    "    for nPar,regularizationParam in enumerate(lambdas):\n",
    "        ind = np.where(minerrIndex==nPar)[0];\n",
    "        if len(ind)>0:\n",
    "            r[ind] = regularizationParam;\n",
    "            print('{}% of outputs with regularization param: {}'.format(int(len(ind)/n_voxels*100),\n",
    "                                                                        regularizationParam))\n",
    "            # got good param, now obtain weights\n",
    "            dlambda = D**2 + np.eye(n_p)*regularizationParam\n",
    "            dlambdaInv = np.diag(D / np.diag(dlambda))\n",
    "            KlambdaInv = V.dot(dlambdaInv).dot(U.T)\n",
    "\n",
    "            weightMatrix[:,ind] = KlambdaInv.dot(train_data[:,ind]);\n",
    "\n",
    "\n",
    "    return weightMatrix, r\n",
    "\n",
    "def sub_to_subjectnum(sub):\n",
    "    for i in range(len(all_subjects)):\n",
    "        if all_subjects[i] == sub:\n",
    "            return i+1\n",
    "    return 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval(feature_name, save_regressed_y = False):    \n",
    "    print(feature_name)\n",
    "    np.random.seed(9)\n",
    "    kf = KFold(n_splits=4)\n",
    "    \n",
    "    # read features\n",
    "    features = np.load(os.path.join(features_dir,\"{}.npy\".format(feature_name)))\n",
    "    features = features.reshape(features.shape[0], -1)\n",
    "    features = [list(features[i]) for i in range(len(features))]\n",
    "    print(len(features))\n",
    "    \n",
    "    # align features with fMRI data and create final\n",
    "    history = 4\n",
    "    featlen = len(features[0])\n",
    "    x = []\n",
    "\n",
    "    x = features[:1303] + [[0.0]*featlen] + features[1303:2654] + [[0.0]*featlen] \\\n",
    "            + features[2654:3713] + [[0.0]*featlen] + features[3713:] + [[0.0]*featlen]\n",
    "\n",
    "    for j in range(history*4):\n",
    "        x = [[0.0]*featlen] + x\n",
    "\n",
    "    features_tr = []\n",
    "    k = history\n",
    "    while(k<len(x)/4):\n",
    "        if len(features_tr)<=(k-history):\n",
    "            features_tr.append([])\n",
    "        for o in range(history,0,-1):\n",
    "            features_tr[k-history] = features_tr[k-history] + np.sum([j for j in x[(k-o)*4:((k-o+1)*4)]],axis=0).tolist()\n",
    "        k += 1\n",
    "\n",
    "    features_tr = features_tr[:325] + features_tr[326:663] + features_tr[664:928] + features_tr[929:-1]\n",
    "    features_tr = np.array(features_tr)\n",
    "\n",
    "    print(features_tr.shape)\n",
    "\n",
    "    save_dir = os.path.join(results_save_dir, feature_name)\n",
    "\n",
    "    if not os.path.exists(save_dir):\n",
    "        os.mkdir(save_dir)\n",
    "\n",
    "    for sub in all_subjects:\n",
    "        print(\"Subject \" + sub)\n",
    "        if os.path.exists(os.path.join(save_dir, sub + \"_r2s.npy\")):\n",
    "            continue\n",
    "        N_s = sub_to_subjectnum(sub)\n",
    "        data = np.load(sub_data + sub + \".npy\")\n",
    "        print(data.shape)\n",
    "        print(features_tr.shape)\n",
    "\n",
    "        r2s = np.zeros(data.shape[1])   \n",
    "        corr = np.zeros((data.shape[1],2))\n",
    "\n",
    "        split_num = 0\n",
    "\n",
    "        all_preds = []\n",
    "        all_reals = []\n",
    "\n",
    "        for train_index, test_index in kf.split(data):\n",
    "            # remove 5 TRs which either precede or follow the TRs in the test set\n",
    "            train_index_without_overlap = train_index\n",
    "            for rem_val in range(test_index[0] - 5, test_index[0], 1):\n",
    "                train_index_without_overlap = train_index_without_overlap[train_index_without_overlap != rem_val]\n",
    "                \n",
    "            for rem_val in range(test_index[-1] + 1, test_index[-1] + 6, 1):\n",
    "                train_index_without_overlap = train_index_without_overlap[train_index_without_overlap != rem_val]\n",
    "            \n",
    "            x_train, x_test = features_tr[train_index_without_overlap], features_tr[test_index]\n",
    "            y_train, y_test = data[train_index_without_overlap], data[test_index]\n",
    "            \n",
    "            x_train = stats.zscore(x_train,axis=0)\n",
    "            x_train = np.nan_to_num(x_train)\n",
    "\n",
    "            x_test = stats.zscore(x_test,axis=0)\n",
    "            x_test = np.nan_to_num(x_test)\n",
    "\n",
    "            y_train = stats.zscore(y_train,axis=0)\n",
    "            y_train = np.nan_to_num(y_train)\n",
    "\n",
    "            y_test = stats.zscore(y_test,axis=0) \n",
    "            y_test = np.nan_to_num(y_test)\n",
    "            \n",
    "            print(x_train.shape)\n",
    "            print(y_train.shape)\n",
    "            print(x_test.shape)\n",
    "            print(y_test.shape)\n",
    "\n",
    "            weights, lbda = cross_val_ridge(x_train,y_train)        \n",
    "            y_pred = np.dot(x_test,weights)\n",
    "\n",
    "            np.save(os.path.join(save_dir, \"{}_y_pred_{}\".format(sub, split_num)),y_pred)\n",
    "            np.save(os.path.join(save_dir, \"{}_y_test_{}\".format(sub, split_num)),y_test)\n",
    "            np.save(os.path.join(save_dir, \"{}_weights_{}\".format(sub, split_num)),weights)\n",
    "            np.save(os.path.join(save_dir, \"{}_lbda_{}\".format(sub, split_num)),lbda)\n",
    "            \n",
    "            if save_regressed_y:\n",
    "                y_pred_train = np.dot(x_train,weights)\n",
    "                np.save(os.path.join(save_dir, \"{}_y_regress_test_{}\".format(sub, split_num)),y_test - y_pred)\n",
    "                np.save(os.path.join(save_dir, \"{}_y_regress_train_{}\".format(sub, split_num)),y_train - y_pred_train)\n",
    "\n",
    "            split_num += 1\n",
    "\n",
    "            all_reals.append(y_test)\n",
    "            all_preds.append(y_pred)\n",
    "\n",
    "        all_reals = np.vstack(all_reals)\n",
    "        all_preds = np.vstack(all_preds)\n",
    "\n",
    "        r2s = r2_score(all_reals,all_preds,multioutput=\"raw_values\")\n",
    "\n",
    "        for i in range(all_reals.shape[1]):\n",
    "            if np.nonzero(all_reals[:,i])[0].size > 0:\n",
    "                corr[i] = stats.pearsonr(all_reals[:,i],all_preds[:,i])\n",
    "            else:\n",
    "                r2s[i] = 0\n",
    "                corr[i][1] = 1\n",
    "\n",
    "        print(np.max(r2s))\n",
    "\n",
    "        np.save(os.path.join(save_dir, sub + \"_r2s\"),r2s)\n",
    "        np.save(os.path.join(save_dir, sub + \"_corr\"),corr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "num_contrege_sets_per_space = 5\n",
    "contrege_sets = [\"contrege_comp\", \"contrege_incomp\", \"incontrege\"]\n",
    "\n",
    "all_features = [\"punct_final\", \"node_count_punct\", \"syntactic_surprisal_punct\", \\\n",
    "                \"word_frequency_punct\", \"word_length_punct\", \"all_complexity_based_metrics_punct\", \\\n",
    "                \"pos_dep_tags_all_complexity_based_metrics\"]\n",
    "for set_name in contrege_sets:\n",
    "    for i in range(num_contrege_sets_per_space):\n",
    "        all_features.append(\"{}_set_{}_pos_dep_tags_all_complexity_metrics\".format(set_name, i))\n",
    "        if set_name == \"contrege_incomp\":\n",
    "            all_features.append(\"bert_PCA_dims_15_{}_set_{}_pos_dep_tags_all_complexity_metrics\".format(set_name, i))\n",
    "\n",
    "print(all_features)\n",
    "\n",
    "all_subjects = [\"F\",\"G\",\"H\",\"I\",\"J\",\"K\",\"L\",\"M\",\"N\"]\n",
    "\n",
    "for feat in all_features:\n",
    "    eval(feat)"
   ]
  },
  {
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}