scripts/mrd_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import the modules\n",
    "import os\n",
    "import GPy\n",
    "import sys\n",
    "import numpy as np\n",
    "from matplotlib import cm\n",
    "import cPickle as pickle\n",
    "import scipy.stats as stats\n",
    "import sklearn.metrics as metrics\n",
    "import GPy.plotting.Tango as Tango\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "from matplotlib import pyplot as plt\n",
    "from sklearn.decomposition import PCA\n",
    "from GPy.plotting.matplot_dep.controllers.imshow_controller import ImshowController"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PROJECT_DIR = 'HSI2020/data'\n",
    "DATA_DIR = os.path.join(PROJECT_DIR, 'proc')\n",
    "JOINT_ANGLES_DIR = os.path.join(DATA_DIR, 'Baxter')\n",
    "WHILL_DIR = os.path.join(DATA_DIR, 'Whill')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_data(file_index):\n",
    "    JOINT_ANGLES_FILE = os.path.join(JOINT_ANGLES_DIR, 'joint_angles_%d.csv' % file_index)\n",
    "    WHILL_FILE = os.path.join(WHILL_DIR, 'whill_movement_%d.csv' % file_index)\n",
    "\n",
    "    # load the joint angles train data\n",
    "    joint_angles = np.loadtxt(JOINT_ANGLES_FILE, delimiter=',', skiprows=1)\n",
    "\n",
    "    # load the whill movement train data\n",
    "    whill_movement = np.loadtxt(WHILL_FILE, delimiter=',', skiprows=1)\n",
    "\n",
    "    # remove time stamp info from both observations\n",
    "    joint_angles = joint_angles[:, 1:]\n",
    "    whill_movement = whill_movement[:, 1:]\n",
    "    return joint_angles, whill_movement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define the index of the files used for testing and training\n",
    "train_files = [3]\n",
    "test_files = [1, 2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "joint_angles = []\n",
    "whill_movement = []\n",
    "for idx in train_files:\n",
    "    FILE_INDEX = idx\n",
    "    ja, wm = get_data(FILE_INDEX)\n",
    "    joint_angles.append(ja)\n",
    "    whill_movement.append(wm)\n",
    "\n",
    "joint_angles = np.concatenate(joint_angles, axis=0)\n",
    "whill_movement = np.concatenate(whill_movement, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "joint_angles.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "whill_movement.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "for idx in range(len(train_files)):\n",
    "    i = 200 * idx\n",
    "    j = 200 * (idx + 1)\n",
    "    plt.plot(whill_movement[i:j,0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "joint_angles.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "whill_movement.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add small bias\n",
    "bias = 1e-4 * np.random.normal(size=whill_movement.shape)\n",
    "whill_movement += bias\n",
    "whill_movement.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rep_cols = 5\n",
    "whill_movement = np.tile(whill_movement, reps=(1, rep_cols))\n",
    "whill_movement.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plotScales(scales1, scales2, options, yThresh=0.05):\n",
    "    fig = plt.figure()    \n",
    "    ax = fig.add_subplot(111)\n",
    "    \n",
    "    x = np.arange(1,scales1.shape[0] + 1)\n",
    "    c1 = Tango.colorsHex['mediumBlue']\n",
    "    c2 = Tango.colorsHex['darkGreen']\n",
    "    h1 = ax.bar(x, height=scales1, width=0.8, align='center', color=c1, linewidth=1.3)\n",
    "    h2 = ax.bar(x, height=scales2, width=0.5, align='center', color=c2, linewidth=0.7)\n",
    "    ax.plot([0.4, scales1.shape[0] + 0.6], [yThresh, yThresh], '--', linewidth=3, color=Tango.colorsHex['mediumRed'])\n",
    "    \n",
    "    # setting the bar plot parameters\n",
    "    ax.set_xlim(.4, scales1.shape[0]+.6)\n",
    "    ax.tick_params(axis='both')\n",
    "    ax.set_xticks(xrange(1,scales1.shape[0]+1))\n",
    "    ax.set_title(options['title'])\n",
    "    ax.set_ylabel(options['ylabel'])\n",
    "    ax.set_xlabel('Latent Dimensions')\n",
    "    ax.legend([h1,h2],options['labels'], loc='upper right')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plotLatent(inX,\n",
    "               title,\n",
    "               model=None,\n",
    "               which_indices=[0,1], \n",
    "               plot_inducing=False, \n",
    "               plot_variance=False, \n",
    "               max_points=[800, 300],\n",
    "               test_also=False):\n",
    "    s = 100\n",
    "    marker = 'o'    \n",
    "    resolution = 50\n",
    "    \n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111)\n",
    "    Tango.reset()\n",
    "\n",
    "    input1, input2 = which_indices\n",
    "\n",
    "    if inX[0].shape[0] > max_points[0]:\n",
    "        print(\"Warning\".format(inX[0].shape))\n",
    "        subsample = np.random.choice(inX[0].shape[0], size=max_points[0], replace=False)\n",
    "        inX[0] = inX[0][subsample]\n",
    "\n",
    "    if test_also:\n",
    "        if inX[1].shape[0] > max_points[1]:\n",
    "            print(\"Warning\".format(inX[1].shape))\n",
    "            subsample = np.random.choice(inX[1].shape[0], size=max_points[1], replace=False)\n",
    "            inX[1] = inX[1][subsample]\n",
    "        \n",
    "    xmin, ymin = inX[0][:, [input1, input2]].min(0)\n",
    "    xmax, ymax = inX[0][:, [input1, input2]].max(0)\n",
    "    x_r, y_r = xmax-xmin, ymax-ymin\n",
    "    xmin -= .1*x_r\n",
    "    xmax += .1*x_r\n",
    "    ymin -= .1*y_r\n",
    "    ymax += .1*y_r\n",
    "    print xmin, xmax, ymin, ymax\n",
    "    \n",
    "    if plot_variance:\n",
    "        def plotFunction(x):\n",
    "            Xtest_full = np.zeros((x.shape[0], qDim))\n",
    "            Xtest_full[:, [input1, input2]] = x\n",
    "            _, var = model.predict(np.atleast_2d(Xtest_full))\n",
    "            var = var[:, :1]\n",
    "            return -np.log(var)\n",
    "        qDim = model.X.mean.shape[1]\n",
    "        x, y = np.mgrid[xmin:xmax:1j*resolution, ymin:ymax:1j*resolution]\n",
    "        gridData = np.hstack((x.flatten()[:, None], y.flatten()[:, None]))\n",
    "        gridVariance = (plotFunction(gridData)).reshape((resolution, resolution))\n",
    "        varianceHandle = plt.imshow(gridVariance.T, interpolation='bilinear', origin='lower', cmap=cm.gray, extent=(xmin, xmax, ymin, ymax))\n",
    "        \n",
    "    trainH = ax.scatter(inX[0][:, input1], inX[0][:, input2], marker=marker, s=s, c=Tango.colorsHex['mediumBlue'], linewidth=.2, alpha=1.)\n",
    "    \n",
    "    if test_also:\n",
    "        testH = ax.scatter(inX[1][:, input1], inX[1][:, input2], marker=marker, s=s, c=Tango.colorsHex['mediumRed'], linewidth=.2, alpha=0.9)\n",
    "    \n",
    "    ax.grid(b=False) \n",
    "    ax.set_aspect('auto')\n",
    "    ax.tick_params(axis='both')\n",
    "\n",
    "    if test_also:\n",
    "        ax.legend([trainH,testH],['Train','Test'],loc='upper right')\n",
    "    else:\n",
    "        ax.legend([trainH],['Train'], loc='upper right')\n",
    "        \n",
    "    ax.set_xlabel('Latent Dimension %i' % (input1 + 1))\n",
    "    ax.set_ylabel('Latent Dimension %i' % (input2 + 1))\n",
    "    \n",
    "    if plot_inducing:\n",
    "        Z = model.Z\n",
    "        ax.scatter(Z[:, input1], Z[:, input2], c='w', s=25, marker=\"^\", linewidth=.3, alpha=.6)\n",
    "\n",
    "    ax.set_xlim((xmin, xmax))\n",
    "    ax.set_ylim((ymin, ymax))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set the overall dimensions for MRD\n",
    "latent_dim = 8\n",
    "\n",
    "# out of 8 dimensions, 6 dimensions are assigned to baxter's joint angle\n",
    "joint_angle_dim = 6\n",
    "\n",
    "num_inducing = 100\n",
    "\n",
    "dim_values = [range(joint_angle_dim), range(joint_angle_dim, latent_dim)]\n",
    "\n",
    "train_list = [joint_angles, whill_movement] # list of input and corresponding label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# optimization variables\n",
    "SNR0 = 100\n",
    "SNR1 = 50\n",
    "train_iters = 200\n",
    "cons_mod0_max_iters = 200\n",
    "cons_mod1_max_iters = 200\n",
    "fix_var_max_iters = 6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_samples = train_list[0].shape[0]\n",
    "\n",
    "scales = []\n",
    "input_X = np.zeros((num_samples, latent_dim))\n",
    "dim_distribution = [joint_angle_dim, latent_dim - joint_angle_dim]\n",
    "\n",
    "for dim, values, Y in zip(dim_distribution, dim_values, train_list):\n",
    "    X, var = GPy.util.initialization.initialize_latent('PCA', dim, Y)\n",
    "    scales.extend(var)\n",
    "    input_X[:, values] = X\n",
    "\n",
    "scales = np.asarray(scales)\n",
    "\n",
    "mrd_kernels = [GPy.kern.RBF(latent_dim, variance=1, lengthscale=1.0/scales, ARD=True) for _ in train_list]\n",
    "\n",
    "mrd_model = GPy.models.MRD(train_list,\n",
    "                           input_dim=latent_dim,\n",
    "                           num_inducing=num_inducing,\n",
    "                           kernel=mrd_kernels,\n",
    "                           X=input_X,\n",
    "                           name='mrd')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Phase 1: Optimizaition by fixing variance parameters\n",
    "var0 = mrd_model.Y0.Y.var()\n",
    "var1 = mrd_model.Y1.Y.var()\n",
    "\n",
    "mrd_model.Y0.rbf.variance.fix(var0)\n",
    "mrd_model.Y1.rbf.variance.fix(var1)\n",
    "\n",
    "mrd_model.Y0.Gaussian_noise.variance.fix(var0 / SNR0)\n",
    "mrd_model.Y1.Gaussian_noise.variance.fix(var1 / SNR1)\n",
    "\n",
    "mrd_model.optimize(messages=True, max_iters=fix_var_max_iters)\n",
    "\n",
    "print 'Phase 1 Optimizaition Done.'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Phase 2.1: Optimize each model individually\n",
    "# constrain space 0\n",
    "mrd_model.Y1.constrain_fixed()\n",
    "mrd_model.optimize(messages=True, max_iters=cons_mod0_max_iters)\n",
    "\n",
    "print 'Phase 2.1 Optimizaition Done.'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Phase 2.2: Optimize each model individually\n",
    "# constrain space 1\n",
    "mrd_model.Y0.constrain_fixed()\n",
    "mrd_model.Y1.unconstrain_fixed()\n",
    "mrd_model.Y1.rbf.variance.fix(var1)\n",
    "mrd_model.Y1.Gaussian_noise.variance.fix(var1 / SNR1)\n",
    "mrd_model.optimize(messages=True, max_iters=cons_mod1_max_iters)\n",
    "    \n",
    "print 'Phase 2.2 Optimizaition Done.'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Phase 3: Optimize the model without any constraints\n",
    "# training without constraints\n",
    "mrd_model.Y0.unconstrain_fixed()\n",
    "mrd_model.Y1.unconstrain_fixed()\n",
    "mrd_model.optimize(messages=True, max_iters=train_iters)\n",
    "    \n",
    "print 'Phase 3 Optimizaition Done.'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "scales1 = mrd_model.Y0.kern.input_sensitivity(summarize=False)\n",
    "scales2 = mrd_model.Y1.kern.input_sensitivity(summarize=False)\n",
    "\n",
    "scales1 /= scales1.max()\n",
    "scales2 /= scales2.max()\n",
    "\n",
    "options = {'title':'ARD Weights','ylabel':'ARD Weight','labels':['JointAngle', 'WhillMovement']}\n",
    "plotScales(scales1, scales2, options)\n",
    "\n",
    "plt.savefig('ARD_Weights_Baxter_to_Whill.pdf', format='pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pickle.dump(mrd_model, open('mrd_model.pkl','wb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function to compute reconstruction error\n",
    "def reconstructionErrorMRD(model, valData, testData, mKey, kKey, optimizeFlag=False):    \n",
    "    nSamplesVal = valData[mKey].shape[0]\n",
    "    nSamplesTest = testData[mKey].shape[0]\n",
    "\n",
    "    dims = val_data[mKey].shape[1]  \n",
    "    nDimIn = valData[kKey].shape[1]\n",
    "    nDimOut = dims\n",
    "    \n",
    "    qDim = model.X.mean.shape[1]\n",
    "    \n",
    "    # computing reconstruction error for test1, test2 with variances\n",
    "    predictVal = np.zeros((nSamplesVal,nDimOut))\n",
    "    predictTest = np.zeros((nSamplesTest,nDimOut))\n",
    "\n",
    "    latentVal = np.zeros((nSamplesVal,qDim))\n",
    "    latentTest = np.zeros((nSamplesTest,qDim))\n",
    "\n",
    "    for n in tqdm(range(nSamplesVal)):\n",
    "        yIn = valData[kKey][n,:]\n",
    "        yTrueOut = valData[mKey][n]\n",
    "    \n",
    "        [xPredict, infX] = model.Y0.infer_newX(yIn[None,:], optimize=False)\n",
    "        yOut = model.predict(xPredict.mean, Yindex=1)    \n",
    "        #sys.stdout.write('.')\n",
    "        \n",
    "        predictVal[n,:] = yOut[0]\n",
    "        latentVal[n,:] = xPredict.mean\n",
    "        \n",
    "    #sys.stdout.write('\\n\\n')\n",
    "        \n",
    "    for n in tqdm(range(nSamplesTest)):\n",
    "        yIn = testData[kKey][n,:]\n",
    "        yTrueOut = testData[mKey][n]\n",
    "    \n",
    "        [xPredict, infX] = model.Y0.infer_newX(yIn[None,:], optimize=optimizeFlag)\n",
    "        yOut = model.predict(xPredict.mean, Yindex=1)    \n",
    "        #sys.stdout.write('.')\n",
    "        \n",
    "        predictTest[n,:] = yOut[0]\n",
    "        latentTest[n,:] = xPredict.mean\n",
    "        \n",
    "    #sys.stdout.write('\\n\\n')\n",
    "    results = {}\n",
    "    valResults = {}\n",
    "    testResults = {}\n",
    "    \n",
    "    valResults['pred'] = predictVal\n",
    "    testResults['pred'] = predictTest\n",
    "    \n",
    "    valResults['latent'] = latentVal\n",
    "    testResults['latent'] = latentTest\n",
    "    \n",
    "    valErrors = np.sqrt(metrics.mean_squared_error(valData[mKey],predictVal,multioutput='raw_values'))\n",
    "    testErrors = np.sqrt(metrics.mean_squared_error(testData[mKey],predictTest,multioutput='raw_values'))\n",
    "\n",
    "    valNormErrors = np.divide(np.sqrt(metrics.mean_squared_error(valData[mKey],predictVal,multioutput='raw_values')), \n",
    "                              valData[mKey].max(axis=0) - valData[mKey].min(axis=0))\n",
    "    testNormErrors = np.divide(np.sqrt(metrics.mean_squared_error(testData[mKey],predictTest,multioutput='raw_values')), \n",
    "                               testData[mKey].max(axis=0) - testData[mKey].min(axis=0))\n",
    "\n",
    "    valCorr = np.zeros((1,nDimOut))\n",
    "    testCorr = np.zeros((1,nDimOut))\n",
    "    for d in range(dims):\n",
    "        valCorr[0,d],_ = stats.pearsonr(valData[mKey][1],predictVal[d])\n",
    "        testCorr[0,d],_ = stats.pearsonr(testData[mKey][1],predictTest[d])\n",
    "\n",
    "    valResults['rmse'] = valErrors\n",
    "    testResults['rmse'] = testErrors\n",
    "    \n",
    "    valResults['nrmse'] = valNormErrors\n",
    "    testResults['nrmse'] = testNormErrors\n",
    "    \n",
    "    valResults['corr'] = valCorr\n",
    "    testResults['corr'] = testCorr\n",
    "        \n",
    "    results['train'] = valResults\n",
    "    results['test'] = testResults\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mrd_model = pickle.load(open('mrd_model.pkl','rb'))\n",
    "train_X = mrd_model.X.mean\n",
    "mrd_X = [train_X, train_X]\n",
    "\n",
    "plotLatent(mrd_X, \n",
    "           'Latent Space', \n",
    "           model=mrd_model, \n",
    "           which_indices=[0, 1], \n",
    "           plot_variance=True, \n",
    "           max_points=[300, 300], \n",
    "           test_also=False, \n",
    "           plot_inducing=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}