iNN_2D_sim.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2D test mapping from position $x$ to position $y$ and $\\delta q$ with NN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "# %load_ext tensorboard\n",
    "import time    \n",
    "\n",
    "import tensorflow as tf\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow.contrib.distributions as ds\n",
    "from IPython.core import display\n",
    "import data_generation as dg\n",
    "plt.style.use('ggplot')\n",
    "%matplotlib inline\n",
    "\n",
    "# https://github.com/moble/quaternion  install quaternion and numba\n",
    "import quaternion # [w,x,y,z] order\n",
    "\n",
    "\n",
    "import mapping_tools as tools\n",
    "\n",
    "# use conda env 'tf1'\n",
    "\n",
    "%load_ext autoreload \n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess = tf.InteractiveSession()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## prepare data X and Y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## spherical linear interpolation to generate data for straight lines\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create grid_lines for visualization\n",
    "downsampling = 10\n",
    "nbGrid = 20 * downsampling + 1\n",
    "def generate_grid(xk, xk2, offset = 1.5):\n",
    "    grid_lines=[]\n",
    "    x_grid_ = np.linspace(np.min(xk[:,0]) - offset, np.max(xk[:,0]) + offset, nbGrid)\n",
    "    y_grid_ = np.linspace(np.min(xk[:,1]) - offset, np.max(xk[:,1]) + offset, nbGrid)\n",
    "    x_grid, y_grid = np.meshgrid(x_grid_, y_grid_)\n",
    "    for ii in range(nbGrid):\n",
    "        for jj in range(nbGrid):               \n",
    "            x_grid_point = np.array([x_grid[ii,jj],y_grid[ii,jj] ])\n",
    "            grid_lines.append(x_grid_point)\n",
    "\n",
    "    grid_lines = np.asarray(grid_lines)\n",
    "    return grid_lines\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set sampling points for each object\n",
    "sp = 100 # 1,10,100 \n",
    "obj = dg.Shape\n",
    "\n",
    "# generate test objects\n",
    "p0 = obj(sample_points=sp, type=0,translation=np.array([0.7,0.9]),scale=0.1)\n",
    "p1 = obj(sample_points=sp, type=1,translation=np.array([0,1]),scale=0.15)\n",
    "p2 = obj(sample_points=sp, type=2,translation=np.array([1,0]),scale=0.15)\n",
    "p3 = obj(sample_points=sp, type=3,translation=np.array([1.2,0.5]),scale=0.15)\n",
    "p4 = obj(sample_points=sp, type=4,translation=np.array([0,0]),scale=0.5)\n",
    "\n",
    "\n",
    "x_data = np.concatenate([p0.xy,p1.xy,p2.xy,p3.xy,p4.xy],axis=0)\n",
    "x_key = [p0.key_points,p1.key_points,p2.key_points,p3.key_points,p4.key_points]\n",
    "x_data = np.concatenate([x_data, np.repeat([[1.,0,0,0]],5*sp,axis=0)],axis=1)\n",
    "fig = plt.figure(figsize=(10,10./2))\n",
    "ax = fig.add_subplot(111)\n",
    "ax.set_aspect(\"equal\")\n",
    "ax.scatter(x_data[:,0], x_data[:,1],s=1)\n",
    "xk = np.concatenate([p0.xy[0,None,:],p1.xy[0,None,:],p2.xy[0,None,:],p3.xy[0,None,:],p4.xy[0,None,:]],axis=0)\n",
    "ax.set_title('Objects on the local side')\n",
    "# print xk\n",
    "fig = plt.figure(figsize=(10,10./2))\n",
    "\n",
    "ax = fig.add_subplot( 111,sharey=ax)\n",
    "# sharey=ax\n",
    "ax.set_aspect(\"equal\")\n",
    "\n",
    "# set the rotation angle on objects\n",
    "angle = [0,-30,60,30,30.]\n",
    "p0 = obj(sample_points=sp, type=0,translation=np.array([0.7,0.7])+np.array([0.8,0.5]),theta=angle[0],scale=0.1)\n",
    "p1 = obj(sample_points=sp, type=1,translation=np.array([0,1])+np.array([0.4,0.2]),theta=angle[1],scale=0.15)\n",
    "p2 = obj(sample_points=sp, type=2,translation=np.array([1,0])+np.array([0.7,0.2]),theta=angle[2],scale=0.15)\n",
    "p3 = obj(sample_points=sp, type=3,translation=np.array([1.5,0.8])+np.array([0.5,0.3]),theta=angle[3],scale=0.15)\n",
    "p4 = obj(sample_points=sp, type=4,translation=np.array([0,0])+np.array([0,0.2]),theta=angle[4],scale=0.5)\n",
    "xk2 = np.concatenate([p0.xy[0,None,:],p1.xy[0,None,:],p2.xy[0,None,:],p3.xy[0,None,:],p4.xy[0,None,:]],axis=0)\n",
    "y_key = [p0.key_points,p1.key_points,p2.key_points,p3.key_points,p4.key_points]\n",
    "\n",
    "p = [p0.xy,p1.xy,p2.xy,p3.xy,p4.xy]\n",
    "\n",
    "q_all =np.zeros([5,4])\n",
    "for i in range(5):\n",
    "    q = quaternion.from_euler_angles([-angle[i]*np.pi/180.,0,0])\n",
    "    q_tmp = quaternion.as_float_array(q)\n",
    "    q_all[i,:] = q_tmp\n",
    "\n",
    "    p[i] = np.concatenate([p[i],np.repeat(q_tmp.reshape(1,-1), sp,axis=0)],axis=1)\n",
    "xk2 = np.concatenate([xk2, q_all],axis=1)\n",
    "y_data = np.concatenate([p[0],p[1],p[2],p[3],p[4]],axis=0)\n",
    "ax.scatter(y_data[:,0], y_data[:,1],s=1)\n",
    "ax.set_title('Objects on the remote side')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "delta_quaternion = quaternion.from_float_array(y_data[:,2:])/ quaternion.from_float_array(x_data[:,2:])\n",
    "delta_q = quaternion.as_float_array(delta_quaternion)\n",
    "delta_q_tan = tools.log_q(delta_q)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# to evaluate the jacobian\n",
    "def batch_jacobians(ys, xs):\n",
    "    \"\"\"\n",
    "    ys : [None, n_y] or [n_y]\n",
    "    xs : [None, n_x] or [n_x]\n",
    "    \"\"\"\n",
    "    if ys.shape.ndims == 2:\n",
    "        return tf.transpose(\n",
    "            tf.stack([tf.gradients(ys[:, i], xs)[0] for i in range(ys.shape[-1].value)]),\n",
    "            (1, 0, 2))\n",
    "    elif ys.shape.ndims == 1:\n",
    "        return tf.stack([tf.gradients(ys[i], xs)[0] for i in range(ys.shape[0].value)])\n",
    "    else:\n",
    "        raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen_uniform_samples(low, high, size):\n",
    "    if low.shape[0]==1:\n",
    "        output = (high - low) * np.random.random_sample((size,)) + low\n",
    "    elif low.shape[0]==2:\n",
    "        output_x = (high[0] - low[0]) * np.random.random_sample((size,)) + low[0]\n",
    "        output_y = (high[1] - low[1]) * np.random.random_sample((size,)) + low[1]\n",
    "        output = np.transpose(np.vstack([output_x, output_y]))\n",
    "    elif low.shape[0]==3:\n",
    "        output_x = (high[0] - low[0]) * np.random.random_sample((size,)) + low[0]\n",
    "        output_y = (high[1] - low[1]) * np.random.random_sample((size,)) + low[1]\n",
    "        output_z = (high[2] - low[2]) * np.random.random_sample((size,)) + low[2]\n",
    "        output = np.transpose(np.vstack([output_x, output_y, output_z]))\n",
    "    else:\n",
    "        raise Exception(\"Invalid dimension: \", low.ndim )\n",
    "    return output\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tmp_create_samples(x,num):\n",
    "    \n",
    "    # x: [N,7]\n",
    "    # a_re: [n, 7]\n",
    "    \n",
    "    center = np.mean(x[:,:2],axis=0)\n",
    "    # aa = np.random.multivariate_normal(center, 0.15* np.eye(2), (num, )) # random samples\n",
    "    offset_2 = 1.5\n",
    "    aa = gen_uniform_samples(center-offset_2, center+offset_2, num)\n",
    "    r = 0\n",
    "    a_re = []\n",
    "    weight = []\n",
    "    for i in range(num):\n",
    "        dis2 = np.sqrt( np.sum((aa[i,:] - center)**2) )\n",
    "        if dis2 >r:\n",
    "            a_re.append(aa[i,:] )\n",
    "            weight.append(dis2-r)\n",
    "    if len(weight)==0:\n",
    "        pass\n",
    "        return a_re, weight/ (sum(weight)/len(weight) )\n",
    "    else:\n",
    "        return a_re, weight/ (sum(weight)/len(weight) )\n",
    "\n",
    "samples, w_samples= tmp_create_samples(x_data[:,:2],200)\n",
    "s2 = np.asarray(samples)\n",
    "# # print s2\n",
    "\n",
    "fig = plt.figure(figsize=(10.,10.))\n",
    "ax = fig.add_subplot(111)\n",
    "\n",
    "ax.plot(x_data[:, 0], x_data[:, 1],'k-')\n",
    "ax.axis('equal')\n",
    "ax.scatter(s2[:,0], s2[:,1], s = 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_layer(inputs, in_size, out_size, name='hidden_layer', scale=1.0,reuse=False, activation_function=None): #activation_function=None\n",
    "    with tf.variable_scope(name) as scope:\n",
    "        if reuse:\n",
    "            scope.reuse_variables()\n",
    "        w_reg=tf.contrib.layers.l2_regularizer(scale)\n",
    "        Weights =  tf.get_variable(name='w', initializer=tf.random_normal([in_size,out_size]), regularizer=w_reg)#Weight\n",
    "        biases = tf.get_variable(name='b', initializer=tf.Variable(tf.zeros([1,out_size])+0.1) )#biases\n",
    "        # print Weights.name, biases.name\n",
    "        Wx_plus_b = tf.matmul(inputs, Weights)+biases #inputs*Weight+biases\n",
    "        if activation_function is None:\n",
    "            outputs = Wx_plus_b\n",
    "        else:\n",
    "            outputs = activation_function(Wx_plus_b)\n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def x2q_NN(x, mid_channel, reuse=False, name ='l', scale= 1., activation=tf.nn.tanh  ):\n",
    "    x = add_layer(x, 2, mid_channel, reuse=reuse, name= name+'1', scale=scale, activation_function=activation  )\n",
    "    x = add_layer(x, mid_channel, mid_channel,reuse=reuse, name=name+'2',scale=scale, activation_function=activation  )\n",
    "    x = add_layer(x, mid_channel, mid_channel,reuse=reuse, name=name+'3', scale=scale, activation_function=activation  )\n",
    "    # x = add_layer(x, mid_channel, mid_channel,reuse=reuse, name=name+'31',scale=scale,  activation_function=activation  )\n",
    "    x = add_layer(x, mid_channel, 3, reuse=reuse, name=name+'4',scale=scale, activation_function=activation  )\n",
    "    return x\n",
    "# generate weight for cost_pred \n",
    "def generate_pred_weight(p_num, lines_num,w1):\n",
    "    a = np.empty([0])\n",
    "    for i in range(p_num-1):          \n",
    "        a1 = np.array([w1])\n",
    "        a2 = np.ones(lines_num - 2)*1.0\n",
    "        a3 = np.concatenate([a1,a2,a1])\n",
    "        a = np.concatenate([a,a3])\n",
    "#     a4 = a4/np.linalg.norm(a4)\n",
    "    return a\n",
    "nbPts = xk.shape[0]\n",
    "w_test = generate_pred_weight(nbPts, 2, 10.)\n",
    "print w_test.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create placeholder to evaluate the NN\n",
    "x_ph = tf.placeholder(tf.float32, (None, 2), name = 'x_ph')\n",
    "y_ph = tf.placeholder(tf.float32, (None, 2), name = 'y_ph')\n",
    "q_tan_ph = tf.placeholder(tf.float32, (None, 3), name = 'q_tan_ph')\n",
    "q_ph = tools.exp_tf(q_tan_ph) # exp\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$\\text{log}(\\overrightarrow q)= \\text{log}(v,\\overrightarrow u)=\\left\\{\n",
    "\\begin{aligned}\n",
    "&\\text{arccos}^{*}(v) \\frac{\\overrightarrow u}{||\\overrightarrow u||},  &||\\overrightarrow u ||\\ne 0\\\\\n",
    "&[0, 0, 0]^T,  & \\text{else} \\\\\n",
    "\\end{aligned}\n",
    "\\right. $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from inv_net import real_nvp\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "activation = tf.nn.tanh \n",
    "mid_channel = 24\n",
    "scale = 0.3\n",
    "\n",
    "\n",
    "x_jac = tf.placeholder(tf.float32, (None, 2),name='x_jac')\n",
    "weight_jac = tf.placeholder(tf.float32, (None, ),name='weight_jac')\n",
    "xk_weight = tf.placeholder(tf.float32, (None, ),name='xk_weight')\n",
    "\n",
    "\n",
    "q_tan_pred = x2q_NN(x_ph, mid_channel, scale=scale, activation=activation) # hidden layer\n",
    "q_pred = tools.exp_tf(q_tan_pred)\n",
    "\n",
    "q_tan_jac = x2q_NN(x_jac, mid_channel, reuse=True, scale=scale, activation=activation) # hidden layer\n",
    "q_jac = tools.exp_tf(q_tan_jac)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scale_kernel= .5\n",
    "mid_channel = 24\n",
    "y_nn = real_nvp(x_ph, mid_channel, backward=False, reuse=False, name='nvp_0', activation=activation, scale_kernel=scale_kernel) \n",
    "\n",
    "# create inverse evaluation\n",
    "y_in = tf.placeholder(tf.float32, (None, 2),name='y_in')\n",
    "x_rec = real_nvp(y_in, mid_channel, backward=True, reuse=True, name='nvp_0', activation=activation, scale_kernel=scale_kernel) \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# placeholder to evaluate the jacobian\n",
    "y_jac = real_nvp(\n",
    "    x_jac, mid_channel, backward=False, name='nvp_0' ,\n",
    "    reuse=True, activation=activation, scale_kernel=scale_kernel) \n",
    "jac_xy = batch_jacobians(y_jac, x_jac)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "q_identity = tf.constant([1.0, 0, 0, 0])\n",
    "\n",
    "tmp_angle = 2* tf.reduce_sum (q_jac * q_identity, axis=1)**2  - 1 \n",
    "# tmp_angle2 = tf.clip_by_value(tmp_angle,0,1)\n",
    "# delta_dis = tf.acos(tmp_angle2)\n",
    "# cost_jac = tf.reduce_sum(delta_dis*weight_jac)\n",
    "jac_samples_num = 200\n",
    "cost_jac_xq =  jac_samples_num - tf.reduce_sum(tmp_angle * weight_jac)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# cost_pred = tf.reduce_sum(orientation_distance(q_pred, q_ph))\n",
    "\n",
    "cost_pred_xq = x_data.shape[0] - tf.reduce_sum(2 * tf.reduce_sum(q_pred * q_ph, axis=1)**2 - 1)\n",
    "scale_jac = tf.Variable(1.4)\n",
    "scale_sig = tf.sigmoid(scale_jac)*2\n",
    "\n",
    "cost_jac_xy = tf.reduce_sum( tf.reduce_sum((jac_xy -  scale_jac*tf.eye(2)) ** 2,[1,2]) * weight_jac )\n",
    "\n",
    "cost_pred_xy = tf.reduce_sum( tf.reduce_sum( (y_ph - y_nn) ** 2,axis=1)* xk_weight )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lmbda_jac = tf.placeholder(tf.float32, (4, ),name='lmbda_jac')   # [cost weight, smooth weight 0.95]\n",
    "\n",
    "condition = tf.placeholder(tf.int32, shape=[], name=\"condition\")\n",
    "cost_jac_xq_last = tf.placeholder(tf.float32, shape=[],name='cost_jac_xq_last')\n",
    "cost_jac_xy_last = tf.placeholder(tf.float32, shape=[],name='cost_jac_xy_last')\n",
    "cost_jac_smooth1 = tf.cond(condition >0, \n",
    "                           lambda: lmbda_jac[1]* cost_jac_xq_last + (1-lmbda_jac[1])*cost_jac_xq ,\n",
    "                           lambda: cost_jac_xq)\n",
    "cost_jac_smooth2 = tf.cond(condition >0, \n",
    "                           lambda: lmbda_jac[1]* cost_jac_xy_last + (1-lmbda_jac[1])*cost_jac_xy ,\n",
    "                           lambda: cost_jac_xy)\n",
    "cost_jac_xq_smooth =  tf.identity(cost_jac_smooth1, name=\"cost_jac_xq_smooth\")\n",
    "cost_jac_xy_smooth =  tf.identity(cost_jac_smooth2, name=\"cost_jac_xy_smooth\")\n",
    "\n",
    "reg_loss = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)\n",
    "\n",
    "# total cost\n",
    "cost_xq = cost_pred_xq + lmbda_jac[0] * cost_jac_xq_smooth + 0.3*tf.add_n(reg_loss)\n",
    "cost_xy = cost_pred_xy + lmbda_jac[3] * cost_jac_xy_smooth\n",
    "cost =   lmbda_jac[2]* cost_xq  + cost_xy\n",
    "# cost =    cost_xy\n",
    "\n",
    "# cost = cost_pred "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# gradient for velocity evaluation\n",
    "g_forward1 = tf.gradients(y_nn, x_ph)\n",
    "g_forward = tf.identity(g_forward1, name=\"g_forward\")\n",
    "\n",
    "g_backward1 = tf.gradients(x_rec, y_in)\n",
    "g_backward = tf.identity(g_backward1, name=\"g_backward\")\n",
    "\n",
    "def f_grad_forward(_x):\n",
    "    g = sess.run(g_forward,{x_ph: _x})\n",
    "    return g\n",
    "def f_g_backward(_x):\n",
    "    g = sess.run(g_backward,{y_in: _x})\n",
    "    return g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# evaluation from x to quaternion offset\n",
    "def f_x2q(_x):\n",
    "     return sess.run(q_pred, {x_ph: _x})\n",
    "\n",
    "def f_x2y(_x):\n",
    "     return sess.run(y_nn, {x_ph: _x})\n",
    "def f_y2x(_x):\n",
    "     return sess.run(x_rec, {y_in: _x})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rate = tf.Variable(0.1)\n",
    "optimizer = tf.train.AdamOptimizer(rate)\n",
    "\n",
    "train = optimizer.minimize(    cost, )\n",
    "# initialize all variables to train\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time = time.time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "_cost_jac_xq_smooth = 0.\n",
    "_cost_jac_xy_smooth = 0.\n",
    "xk_w = generate_pred_weight(nbPts, 100, 5.) \n",
    "xk_w = np.ones(sp*5)\n",
    "delta_t = 0; t1 = 0\n",
    "for i in range(500000):\n",
    "    s, w = tmp_create_samples(x_data[:,:2],jac_samples_num)\n",
    "    delta_t = time.time()-start_time - t1\n",
    "    t1 = time.time()-start_time\n",
    "    time_now =  int(t1)\n",
    "    try:\n",
    "        feed_dict = {\n",
    "            x_ph : x_data[:,:2], # feed the actual points\n",
    "            y_ph: y_data[:,:2],\n",
    "            xk_weight: xk_w, \n",
    "            q_ph: delta_q , # feed their delta quaternions\n",
    "            rate: 0.001, # rate of training\n",
    "            condition: i,\n",
    "            x_jac:s,\n",
    "            weight_jac: w,\n",
    "            cost_jac_xq_last: _cost_jac_xq_smooth,\n",
    "            cost_jac_xy_last: _cost_jac_xy_smooth,\n",
    "            lmbda_jac:  [1e-3, 0.0, 1, 2e-3  ],  # [xq_jac_w, smooth, xq weight, xy_jac_w ]            \n",
    "        }\n",
    "        _, _cost, _cost_pred_xq,_cost_xq, _cost_jac_xq_smooth,_cost_pred_xy,_cost_xy, _cost_jac_xy_smooth  = sess.run(\n",
    "            [train, cost, cost_pred_xq,cost_xq, cost_jac_xq_smooth,cost_pred_xy,cost_xy, cost_jac_xy_smooth],  feed_dict)\n",
    "#         xk2_pred = f_x2y(xk[:,:2])\n",
    "        y_data_pred = f_x2y(x_data[:,:2])\n",
    "        xk_pred_error = np.mean(np.sqrt(  np.sum( ((y_data_pred - y_data[:,:2])**2) , axis=1)   ))\n",
    "        q_eva = f_x2q(x_data[:,:2])\n",
    "        q_error = tools.ori_dis_np(q_eva, delta_q)\n",
    "        if xk_pred_error<0.008 and np.mean(q_error * 180/np.pi)<4.8:\n",
    "            break\n",
    "        if i%20==0:\n",
    "                display.clear_output(wait=True)\n",
    "                print(\"time_now=\", time_now,' delta_t', delta_t,'  i=',i, ' cost:',_cost, ' xk_error=',xk_pred_error)\n",
    "                print('cost_pred_xq: ', _cost_pred_xq, '  cost_xq: ',_cost_xq, ' cost_jac_xq_smooth: ',_cost_jac_xq_smooth)\n",
    "                print('cost_pred_xy: ', _cost_pred_xy, '  cost_xy: ',_cost_xy, ' q_error in degree: ',np.mean(q_error * 180/np.pi))\n",
    "    except KeyboardInterrupt:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# saver = tf.train.Saver()\n",
    "# save_path = './checkpoint_dir/MyModel_2D_iNN_sp'+str(sp)\n",
    "# saver.save(sess, save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "q_eva = f_x2q(x_data[:,:2])\n",
    "q_error = tools.ori_dis_np(q_eva, delta_q)\n",
    "print  'q_error(degree) mean and std: ',np.mean(q_error * 180/np.pi),  np.std(q_error * 180/np.pi)\n",
    "\n",
    "xk2_pred = f_x2y(xk[:,:2])\n",
    "xk_pred_error = np.sqrt(  np.sum( ((xk2_pred - xk2[:,:2])**2) , axis=1)   )\n",
    "print  'xk mean and std: ',np.mean(xk_pred_error),  np.std(xk_pred_error)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %timeit f_grad_forward(xk[0:1,:2])    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %timeit f_g_backward(xk2[0:1,:2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %timeit f_x2y(xk[0:1,:2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %timeit f_y2x(xk2[0:1,:2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %timeit q_eva = f_x2q(x_data[0:1,:2])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_2D_pos_ori(ax, data,color='r',alpha=1,length=0.01,head_width=0.007,lw=1,down_sampling=1,zorder =0):\n",
    "    # data: [n, 6]\n",
    "    arrow_len = np.array([length, 0, 0])\n",
    "    if data.ndim==1:\n",
    "        data = data.reshape(-1,1)        \n",
    "    for i in range(data.shape[0]):\n",
    "        x, y, q = data[i,0], data[i,1], quaternion.from_float_array(data[i,2:])\n",
    "        new_arrow = quaternion.rotate_vectors(q, arrow_len) \n",
    "        tmp =  length* np.array([np.cos(data[i,2]), np.sin(data[i,2])])\n",
    "#         print tmp\n",
    "        ax.arrow(x, y, new_arrow[0],new_arrow[1],width=0.004,head_width=head_width,color=color,alpha=alpha,lw=lw,zorder=zorder)     \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "red_color = np.array([228,73,26,255])/255.\n",
    "green_color = np.array([100,165,74,255])/255.\n",
    "xk_color = np.array([[38, 70, 83,255],\n",
    "                    [42, 157, 143,255],\n",
    "                    [242, 220, 166,255],\n",
    "                    [243, 155, 83,255],\n",
    "                    [197, 61, 27,255]])/255."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_list = []\n",
    "for i in range(0,nbGrid,downsampling):\n",
    "    plot_list += range(i*nbGrid, (i+1)*nbGrid,downsampling )\n",
    "\n",
    "grid_lines = generate_grid(xk, xk2)\n",
    "nbGrid = int(np.sqrt(grid_lines.shape[0]))\n",
    "\n",
    "\n",
    "grid_l = np.concatenate([grid_lines, np.repeat(np.array([[1.,0,0,0]]),nbGrid**2,axis=0 )   ],axis=1)\n",
    "grid_l2 = np.concatenate([f_x2y(grid_lines[:,:2]), f_x2q(grid_lines[:,:2])], axis=1)\n",
    "\n",
    "\n",
    "# fig, ax = plt.subplots(ncols=2, figsize=(fig_width, fig_width *4/8))\n",
    "# color_xk = np.array([[0,1,0,1],[0.5,0.5,0,1],[0,0,1,1],[1,0,0,1],[1,0.5,0.5,1]])\n",
    "# plt.rcParams[\"font.family\"] = \"Times New Roman\"\n",
    "plt.rcParams[\"font.family\"] = \"Times New Roman\"\n",
    "plt.rcParams[\"font.size\"] = \"8\"\n",
    "# plt.rcParams[\"text.usetex\"] = True\n",
    "plt.rcParams['pdf.fonttype'] = 42\n",
    "plt.rcParams['ps.fonttype'] = 42\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "# fig, ax = plt.subplots(ncols=2, figsize=(fig_width, fig_width/2 *4/5))\n",
    "# fig.subplots_adjust(left=0, bottom=0, right=1, top=1,\n",
    "#                 wspace=0.33, hspace=0)\n",
    "plot_list = []\n",
    "for i in range(0,nbGrid,downsampling):\n",
    "    plot_list += range(i*nbGrid, (i+1)*nbGrid,downsampling )\n",
    "fig_width = 85/25.4 *8/9*2 # mm to inch\n",
    "grid_lines = generate_grid(xk, xk2)\n",
    "nbGrid = int(np.sqrt(grid_lines.shape[0]))\n",
    "\n",
    "\n",
    "grid_l = np.concatenate([grid_lines, np.repeat(np.array([[1.,0,0,0]]),nbGrid**2,axis=0 )   ],axis=1)\n",
    "grid_l2 = np.concatenate([f_x2y(grid_lines[:,:2]), f_x2q(grid_lines[:,:2])], axis=1)\n",
    "grids1 = grid_l[:,:2].reshape(nbGrid,nbGrid,2)\n",
    "grids2 = grid_l2[:,:2].reshape(nbGrid,nbGrid,2)\n",
    "\n",
    "rate = 1.\n",
    "fig0 = plt.figure(figsize=(fig_width/3*rate, fig_width/3 *4/5*rate))\n",
    "ax0 = fig0.add_subplot(111)\n",
    "fig1 = plt.figure(figsize=(fig_width/3, fig_width/3 *4/5))\n",
    "\n",
    "ax1 = fig1.add_subplot(111, sharey=ax0)\n",
    "ax = [ax0,ax1]\n",
    "set_alpha = 1\n",
    "grid_color = 'lightgray'\n",
    "for i in range(0, nbGrid, downsampling):\n",
    "    ax[0].plot(grids1[i,:,0], grids1[i,:,1], c = grid_color, alpha=set_alpha,linewidth=0.5,zorder = 0)\n",
    "    ax[0].plot(grids1[:,i,0], grids1[:,i,1], c = grid_color, alpha=set_alpha,linewidth=0.5,zorder = 0)\n",
    "    ax[1].plot(grids2[i,:,0], grids2[i,:,1], c = grid_color, alpha=set_alpha,linewidth=0.5,zorder = 0)\n",
    "    ax[1].plot(grids2[:,i,0], grids2[:,i,1], c = grid_color, alpha=set_alpha,linewidth=0.5,zorder = 0)\n",
    "# Set the zorder for the artist. Artists with lower zorder values are drawn first.\n",
    "\n",
    "# ax[0].axis('equal');ax[1].axis('equal')\n",
    "# xk_color = np.array([[]])\n",
    "# ax[0].scatter(xk[:,0],xk[:,1],s = 10, c=xk_color,zorder = 10)\n",
    "# ax[1].scatter(xk2[:,0],xk2[:,1],s = 10,c=xk_color,zorder = 10)\n",
    "# for i in range(num_points):\n",
    "#     ax[0].scatter(xk[i,0],xk[i,1],s = 15, c=xk_color[i,:],zorder = 20,marker=(4,0,45))\n",
    "#     ax[1].scatter(xk2[i,0],xk2[i,1],s = 15,c=xk_color[i,:],zorder = 20,marker=(4,0,45+ right_point[i,3]-left_point[i,3]))\n",
    "# y_pred = f_x2y(x_data[:,:2])\n",
    "\n",
    "\n",
    "\n",
    "# for i in range(num_points):\n",
    "#     ax[0].text(xk[i,0],xk[i,1],str(i+1),fontsize=10, zorder=100)\n",
    "#     ax[1].text(xk2[i,0],xk2[i,1],str(i+1),fontsize=10, zorder=100 )\n",
    "\n",
    "# ax[0].scatter(x_data[:,0],x_data[:,1],s=0.6, c='k')\n",
    "# ax[1].scatter(y_pred[:,0],y_pred[:,1],s=0.6,c='k')\n",
    "num_points = xk.shape[0]\n",
    "xk_i = np.concatenate([xk[:,:2], np.repeat(np.array([[1.,0,0,0]]),num_points ,axis=0 )   ],axis=1)\n",
    "\n",
    "\n",
    "# plot_2D_pos_ori(ax[0],grid_l[plot_list,:],color='gray',alpha=set_alpha,lw=0.5)\n",
    "plot_2D_pos_ori(ax[1],grid_l2[plot_list,:],color='gray',alpha=set_alpha,lw=0.2,\n",
    "                head_width=0.02,length=0.05,zorder = 5)\n",
    "plot_2D_pos_ori(ax[0],grid_l[plot_list,:],color='gray',alpha=set_alpha,lw=0.2,\n",
    "                head_width=0.02,length=0.05,zorder = 5)\n",
    "# plot_2D_pos_ori(ax[1],grid_l2[plot_list,:],color='gray',alpha=set_alpha,lw=0.2,head_width=0.01,length=0.015,zorder = 5)\n",
    "# plot_2D_pos_ori(ax[0],grid_l[plot_list,:],color='gray',alpha=set_alpha,lw=0.2,head_width=0.01,length=0.015,zorder = 5)\n",
    "# plot_2D_pos_ori(ax[1], xk_tmp, color=red_color,length=0.05,zorder = 15,lw=0.2,head_width=0.02)  \n",
    "\n",
    "# plot_2D_pos_ori(ax[0], xk_i, color=red_color,length=0.05,zorder = 15,lw=0.2,head_width=0.02)  \n",
    "# ax[1].set_xticks([0,0.25,0.5],())\n",
    "# ax[0].set_xticks([0,0.25,0.5],())\n",
    "# ax[0].set_xlim(-0.1,0.75)\n",
    "# ax[0].set_xticks([0,0.2,0.4,0.6],())\n",
    "\n",
    "# ax[0].scatter(x_data[:,0], x_data[:,1],s=1)\n",
    "# ax[1].scatter(y_data[:,0], y_data[:,1],s=1)\n",
    "plt.rcParams[\"font.family\"] = \"Times New Roman\"\n",
    "plt.rcParams[\"font.size\"] = \"8\"\n",
    "# plt.rcParams[\"text.usetex\"] = True\n",
    "i=0\n",
    "for data in x_key:\n",
    "    ax[0].plot(data[:,0],data[:,1],c=xk_color[i,:],linewidth=1)\n",
    "    i = i+1\n",
    "i=0\n",
    "for data in y_key:\n",
    "    ax[1].plot(data[:,0],data[:,1],c=xk_color[i,:],linewidth=1)\n",
    "    i = i+1\n",
    "for i in range(2):\n",
    "    ax[i].patch.set_facecolor(\"white\")    \n",
    "    ax[i].axis('on')   \n",
    "    ax[i].spines['right'].set_visible(True)\n",
    "    ax[i].spines['right'].set_color('k')\n",
    "    ax[i].spines['left'].set_color('k')\n",
    "    ax[i].spines['top'].set_color('k')\n",
    "    ax[i].spines['bottom'].set_color('k')\n",
    "\n",
    "    ax[i].spines['top'].set_visible(True)           \n",
    "#     ax[i].axis('equal')\n",
    "    ax[i].grid(False)\n",
    "    ax[i].tick_params(axis='x', colors='k')\n",
    "    ax[i].tick_params(axis='y', colors='k')\n",
    "    ax[i].set_xticks([0,1,2],())\n",
    "    ax[i].set_yticks([0,1,2],())\n",
    "    \n",
    "x_lim = np.array([-0.5,2.5])\n",
    "y_lim = np.array([-0.5,1.75])+0.05  \n",
    "ax[0].set_xlim(x_lim)\n",
    "ax[0].set_ylim(y_lim)\n",
    "\n",
    "ax[0].set_aspect(1.)\n",
    "ax[1].set_aspect(1.)\n",
    "ax[1].set_xlim(x_lim)\n",
    "ax[1].set_ylim(y_lim)\n",
    "\n",
    "# fig0.savefig('figures/2D_local_sp100_iNN.pdf',format='pdf',bbox_inches='tight',  pad_inches=0.0)\n",
    "# fig1.savefig('figures/2D_remote_sp100_iNN.pdf',format='pdf',bbox_inches='tight',  pad_inches=0.0)\n",
    "             "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### for sp=1 plot\n",
    "%matplotlib inline\n",
    "# fig, ax = plt.subplots(ncols=2, figsize=(fig_width, fig_width/2 *4/5))\n",
    "# fig.subplots_adjust(left=0, bottom=0, right=1, top=1,\n",
    "#                 wspace=0.33, hspace=0)\n",
    "if sp ==1:\n",
    "    fig_width = 85/25.4 *8/9*2 # mm to inch\n",
    "\n",
    "    rate = 1.\n",
    "    fig0 = plt.figure(figsize=(fig_width/3*rate, fig_width/3 *4/5*rate))\n",
    "    ax0 = fig0.add_subplot(111)\n",
    "    fig1 = plt.figure(figsize=(fig_width/3, fig_width/3 *4/5))\n",
    "\n",
    "    ax1 = fig1.add_subplot(111, sharey=ax0)\n",
    "    ax = [ax0,ax1]\n",
    "    set_alpha = 1\n",
    "    grid_color = 'lightgray'\n",
    "    for i in range(0, nbGrid, downsampling):\n",
    "        ax[0].plot(grids1[i,:,0], grids1[i,:,1], c = grid_color, alpha=set_alpha,linewidth=0.5,zorder = 0)\n",
    "        ax[0].plot(grids1[:,i,0], grids1[:,i,1], c = grid_color, alpha=set_alpha,linewidth=0.5,zorder = 0)\n",
    "        ax[1].plot(grids2[i,:,0], grids2[i,:,1], c = grid_color, alpha=set_alpha,linewidth=0.5,zorder = 0)\n",
    "        ax[1].plot(grids2[:,i,0], grids2[:,i,1], c = grid_color, alpha=set_alpha,linewidth=0.5,zorder = 0)\n",
    "    # Set the zorder for the artist. Artists with lower zorder values are drawn first.\n",
    "\n",
    "  \n",
    "    for i in range(num_points):\n",
    "        ax[0].scatter(xk[i,0],xk[i,1],s = 15, c=xk_color[i,:],zorder = 20,marker=(4,0,45))\n",
    "        ax[1].scatter(xk2[i,0],xk2[i,1],s = 15,c=xk_color[i,:],zorder = 20,marker=(4,0,45- angle[i]))\n",
    "\n",
    "\n",
    "    xk_i = np.concatenate([xk[:,:2], np.repeat(np.array([[1.,0,0,0]]),num_points ,axis=0 )   ],axis=1)\n",
    "\n",
    "\n",
    "    # plot_2D_pos_ori(ax[0],grid_lines[plot_list,:],color='gray',alpha=set_alpha,lw=0.5)\n",
    "    plot_2D_pos_ori(ax[1],grid_l2[plot_list,:],color='gray',alpha=set_alpha,lw=0.2,\n",
    "                    head_width=0.02,length=0.05,zorder = 5)\n",
    "    plot_2D_pos_ori(ax[0],grid_l[plot_list,:],color='gray',alpha=set_alpha,lw=0.2,\n",
    "                    head_width=0.02,length=0.05,zorder = 5)\n",
    "\n",
    "    # xk_tmp  = np.concatenate([mapping_forward_evaluation(xk[:,:2]), xk2[:,2:] ],axis=1)\n",
    "    plot_2D_pos_ori(ax[1], xk2, color=red_color,length=0.15,zorder = 15,lw=0.5,head_width=0.03)  \n",
    "\n",
    "    plot_2D_pos_ori(ax[0], xk_i, color=red_color,length=0.15,zorder = 15,lw=0.5,head_width=0.03)  \n",
    "    # ax[1].set_xticks([0,0.25,0.5],())\n",
    "    # ax[0].set_xticks([0,0.25,0.5],())\n",
    "    # ax[0].set_xlim(-0.1,0.75)\n",
    "    # ax[0].set_xticks([0,1],())    \n",
    "    # ax[1].set_xticks([0,1,2],())\n",
    "    for i in range(2):\n",
    "        ax[i].patch.set_facecolor(\"white\")    \n",
    "        ax[i].axis('on')   \n",
    "        ax[i].spines['right'].set_visible(True)\n",
    "        ax[i].spines['right'].set_color('k')\n",
    "        ax[i].spines['left'].set_color('k')\n",
    "        ax[i].spines['top'].set_color('k')\n",
    "        ax[i].spines['bottom'].set_color('k')\n",
    "\n",
    "        ax[i].spines['top'].set_visible(True)           \n",
    "        ax[i].axis('equal')\n",
    "        ax[i].grid(False)\n",
    "        ax[i].tick_params(axis='x', colors='k')\n",
    "        ax[i].tick_params(axis='y', colors='k')\n",
    "    ax[1].set_xticks([0,1,2],())\n",
    "    ax[1].set_yticks([0,1,2],())\n",
    "    ax[0].set_xticks([0,1],())\n",
    "    ax[0].set_yticks([0,1],())\n",
    "\n",
    "# fig0.savefig('figures/2D_local_sp1_iNN.pdf',format='pdf',bbox_inches='tight',  pad_inches=0.0)\n",
    "# fig1.savefig('figures/2D_remote_sp1_iNN.pdf',format='pdf',bbox_inches='tight',  pad_inches=0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "file_extension": ".py",
  "kernelspec": {
   "display_name": "Python 2.7.18 64-bit ('py27': conda)",
   "language": "python",
   "name": "python271864bitpy27conda26beda4d3f90482094fea5792fd19e1a"
  },
  "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.18"
  },
  "mimetype": "text/x-python",
  "name": "python",
  "npconvert_exporter": "python",
  "pygments_lexer": "ipython2",
  "version": 2
 },
 "nbformat": 4,
 "nbformat_minor": 2
}