Average Quaternions.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[4]\n",
      "  [2]\n",
      "  [4]\n",
      "  [5]\n",
      "  [2]\n",
      "  [6]\n",
      "  [5]]\n",
      "\n",
      " [[7]\n",
      "  [7]\n",
      "  [7]\n",
      "  [6]\n",
      "  [6]\n",
      "  [0]\n",
      "  [4]]]\n",
      "[[[ 1.91431642 -2.58644125 -4.52399429]\n",
      "  [ 1.66216634  1.91478373 -3.1754777 ]\n",
      "  [-1.37756934 -3.97945781 -3.61607495]\n",
      "  [-2.76494531 -0.74811043  0.0862    ]\n",
      "  [-2.96312376 -4.93497626 -3.54723092]\n",
      "  [-0.15330327  1.72761479 -2.95415722]\n",
      "  [-3.2569135  -4.5694895  -2.04913778]]\n",
      "\n",
      " [[-0.93503135  1.48367244 -1.4673057 ]\n",
      "  [-1.92046654 -3.04796543 -2.20874811]\n",
      "  [ 0.46850553  0.17085381  0.34586542]\n",
      "  [-4.21028859 -4.26402036 -3.66641393]\n",
      "  [-4.34885968 -2.41807983  1.67236981]\n",
      "  [-2.83175343  1.93669575 -0.80385302]\n",
      "  [-1.80028721  1.90769661  0.02631152]]]\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import time \n",
    "from math import pi as Pi\n",
    "\n",
    "batch_size = 2\n",
    "MaxNAtom = 7\n",
    "MaxNeigh = 7 \n",
    "xyzs = np.random.random((batch_size,MaxNAtom,3))*7.0 - 5.0\n",
    "Zs = np.random.randint(8,size=(batch_size,MaxNAtom,1)) # some z are zero. \n",
    "print Zs\n",
    "print xyzs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.linalg.norm(xyzs[0,2] - xyzs[0,6])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def RotToQuat(axes_):\n",
    "    \"\"\"\n",
    "    axes is a ... X 3 3 tensor of axes\n",
    "    this generates a ... X 4 tensor of quaternions.\n",
    "    which are 1:1 with those axes.\n",
    "    \"\"\"\n",
    "    w = (1./2.)*tf.sqrt(tf.abs(1 + axes_[...,0, 0] + axes_[...,1, 1] + axes_[...,2, 2]))\n",
    "    x = tf.sign(axes_[...,2, 1] - axes_[...,1, 2])*tf.abs(0.5*tf.sqrt(tf.abs(1.0 + axes_[...,0, 0] - axes_[...,1, 1] - axes_[...,2, 2])))\n",
    "    y = tf.sign(axes_[...,0, 2] - axes_[...,2, 0])*tf.abs(0.5*tf.sqrt(tf.abs(1.0 - axes_[...,0, 0] + axes_[...,1, 1] - axes_[...,2, 2])))\n",
    "    z = tf.sign(axes_[...,1, 0] - axes_[...,0, 1])*tf.abs(0.5*tf.sqrt(tf.abs(1.0 - axes_[...,0, 0] - axes_[...,1, 1] + axes_[...,2, 2])))\n",
    "    return tf.stack([w,x,y,z],axis=-1)\n",
    "\n",
    "def QuatToRot(q):\n",
    "    \"\"\"\n",
    "    a_ ... X 4 tensor of quaternions\n",
    "    this generates a ... X 3 X 3 of rotation matrices.\n",
    "    \"\"\"\n",
    "    tmp=tf.stack([1 - 2.*(q[...,2]*q[...,2] + q[...,3]*q[...,3]), 2*(q[...,1]*q[...,2] - q[...,3]*q[...,0]),\n",
    "    2*(q[...,1]*q[...,3] + q[...,2]*q[...,0]),2*(q[...,1]*q[...,2] + q[...,3]*q[...,0]), 1 - 2.*(q[...,1]*q[...,1] + q[...,3]*q[...,3]),\n",
    "    2*(q[...,2]*q[...,3] - q[...,1]*q[...,0]),2*(q[...,1]*q[...,3] - q[...,2]*q[...,0]), 2*(q[...,2]*q[...,3] + q[...,1]*q[...,0]),\n",
    "    1 - 2.*(q[...,1]*q[...,1] + q[...,2]*q[...,2])],axis=-1)\n",
    "    return tf.reshape(tmp,[-1,3,3])\n",
    "\n",
    "def safe_normalize(x_):\n",
    "    nrm = tf.clip_by_value(tf.norm(x_,axis=-1,keepdims=True),1e-36,1e36)\n",
    "    nrm_ok = tf.logical_and(tf.not_equal(nrm,0.),tf.logical_not(tf.is_nan(nrm)))\n",
    "    safe_nrm = tf.where(nrm_ok,nrm,tf.ones_like(nrm))\n",
    "    return x_*tf.where(nrm_ok,1.0/safe_nrm,tf.zeros_like(nrm))\n",
    "\n",
    "def safe_inv_norm(x_):\n",
    "    nrm = tf.clip_by_value(tf.norm(x_,axis=-1,keepdims=True),1e-36,1e36)\n",
    "    nrm_ok = tf.logical_and(tf.not_equal(nrm,0.),tf.logical_not(tf.is_nan(nrm)))\n",
    "    safe_nrm = tf.where(nrm_ok,nrm,tf.ones_like(nrm))\n",
    "    return tf.where(nrm_ok,1.0/safe_nrm,tf.zeros_like(nrm))\n",
    "\n",
    "def TF_AxisAngleRotation(axes, vecs, psis):\n",
    "    \"\"\"Returns a batch of uniform rotation matrices, and the angles of each.\n",
    "    Finds random unit vector and then random angle around it.\n",
    "\n",
    "    Args:     sz_: number of rotation matrices     max_dist: maximum\n",
    "    rotation in units of 2*Pi\n",
    "    \"\"\"\n",
    "    ct = tf.cos(psis)\n",
    "    st = tf.sin(psis)\n",
    "    omct = 1.0 - ct\n",
    "    sz_ = tf.shape(axes)\n",
    "    matrices = tf.reshape(\n",
    "        tf.stack(\n",
    "            [\n",
    "                ct + axes[:, 0] * axes[:, 0] * omct,\n",
    "                axes[:, 0] * axes[:, 1] * omct - axes[:, 2] * st,\n",
    "                axes[:, 0] * axes[:, 2] * omct + axes[:, 1] * st,\n",
    "                axes[:, 1] * axes[:, 0] * omct + axes[:, 2] * st,\n",
    "                ct + axes[:, 1] * axes[:, 1] * omct,\n",
    "                axes[:, 1] * axes[:, 2] * omct - axes[:, 0] * st,\n",
    "                axes[:, 2] * axes[:, 0] * omct - axes[:, 1] * st,\n",
    "                axes[:, 2] * axes[:, 1] * omct + axes[:, 0] * st,\n",
    "                ct + axes[:, 2] * axes[:, 2] * omct\n",
    "            ],\n",
    "            axis=-1), [sz_[0], 3, 3])\n",
    "    return tf.einsum('ijk,ik->ij', matrices, vecs)\n",
    "\n",
    "def PointQuats(dxyzs,z1z2): \n",
    "    \"\"\"This version returns two sets of axes for nearest and next-nearest\n",
    "    neighbor. If the energy from both these representations is averaged the\n",
    "    result will be permutationally invariant (WRT nearest-next-nearest\n",
    "    motion) and rotationally invariant.\n",
    "\n",
    "    Args:\n",
    "        dxyz: a nMol X maxNatom X maxNeigh X 3 tensor of\n",
    "            atoms. (differenced from center of embedding\n",
    "        sparse_mask:  nMol X maxNatom X maxNeigh X 1 sparse mask\n",
    "    \"\"\"\n",
    "    # Append orthogonal axes to dxyzs\n",
    "    argshape = tf.shape(dxyzs)\n",
    "    realdata = tf.reshape(dxyzs,(argshape[0] * argshape[1], argshape[2], 3))\n",
    "    z1z2r = tf.reshape(z1z2,(argshape[0] * argshape[1], argshape[2], 2))\n",
    "    rcut = 5.0\n",
    "    axis_cutoff = rcut*rcut\n",
    "\n",
    "    weightstore = []\n",
    "    cuttore = []\n",
    "    axtore = []\n",
    "\n",
    "    # Only do upper diagonal. also sort z. \n",
    "    perms = []\n",
    "    for i in range(MaxNeigh): \n",
    "        perms = perms + [[i,j] for j in range(i+1,MaxNeigh)]\n",
    "    for perm in perms:\n",
    "        v1 = tf.reshape(dxyzs[:, :, perm[0], :],(argshape[0] * argshape[1], 3)) + tf.constant(\n",
    "                            np.array([1e-23, 0., 0.]), dtype=tf.float64)\n",
    "        v2 = tf.reshape(dxyzs[:, :, perm[1], :],(argshape[0] * argshape[1], 3)) + tf.constant(\n",
    "                            np.array([0., 1e-23, 0.]), dtype=tf.float64)\n",
    "        w1 = tf.reshape(tf.reduce_sum(dxyzs[:, :, perm[0], :] * dxyzs[:, :, perm[0], :],\n",
    "                axis=-1), (argshape[0] * argshape[1], 1))\n",
    "        w2 = tf.reshape(tf.reduce_sum(dxyzs[:, :, perm[1], :] * dxyzs[:, :, perm[1], :],\n",
    "                axis=-1), (argshape[0] * argshape[1], 1))\n",
    "        \n",
    "        mustmsk1 = tf.logical_or(tf.equal(z1z2r[:,perm[0],:1],0),tf.equal(z1z2r[:,perm[0],1:],0))\n",
    "        mustmsk2 = tf.logical_or(mustmsk1,tf.equal(z1z2r[:,perm[1],1:],0))\n",
    "                \n",
    "        v1n = safe_inv_norm(v1) * v1\n",
    "        v2n = safe_inv_norm(v2) * v2\n",
    "        w3p = tf.abs(tf.reduce_sum(v1n * v2n, axis=-1)[..., tf.newaxis])\n",
    "\n",
    "        v3 = tf.cross(v1n, v2n)\n",
    "        posz = tf.tile(tf.greater(tf.reduce_sum(v3*tf.constant([[1.,1.,1.]],dtype=tf.float64),keepdims=True,axis=-1),0.),[1,3])\n",
    "        v3 = tf.where(posz,v3,-1.*v3)\n",
    "        v3 *= safe_inv_norm(v3)\n",
    "\n",
    "        cw1 = tf.where(tf.less(w1, axis_cutoff), tf.cos(w1 / axis_cutoff * Pi / 2.0),\n",
    "            tf.zeros_like(w1))\n",
    "        cw2 = tf.where(tf.less(w2, axis_cutoff), tf.cos(w2 / axis_cutoff * Pi / 2.0),\n",
    "            tf.zeros_like(w2))\n",
    "        ca = tf.where(tf.greater(w3p, 0.92), tf.zeros_like(w3p),tf.cos(w3p / 0.92 * Pi / 2.))\n",
    "        Cuts = tf.where(mustmsk2,tf.zeros_like(cw1), (cw1 * cw2 * ca + 1e-19))\n",
    "        \n",
    "        # Compute the average of v1, v2, and their projections onto the\n",
    "        # plane.\n",
    "        v_av = (v1n + v2n) / 2.0\n",
    "        v_av *= safe_inv_norm(v_av)\n",
    "\n",
    "        # Rotate pi/4 cw and ccw to obtain v1,v2\n",
    "        first = TF_AxisAngleRotation(v3, v_av, tf.constant(Pi / 4., dtype=tf.float64))\n",
    "        second = TF_AxisAngleRotation(v3, v_av,tf.constant(-Pi / 4., dtype=tf.float64))\n",
    "        vs = tf.concat([first[:, tf.newaxis, :], second[:, tf.newaxis, :],v3[:, tf.newaxis, :]],axis=1)\n",
    "        axtore.append(vs)\n",
    "        weightstore.append(Cuts)\n",
    "\n",
    "    axs = tf.stack(axtore, axis=-3)\n",
    "    ws = tf.stack(weightstore, axis=-2) \n",
    "    dn = tf.clip_by_value(tf.reduce_sum(ws, axis=1, keepdims=True), 1e-19, 1e19)\n",
    "    wws = tf.where(tf.greater_equal(ws, 0.), ws / (dn), tf.zeros_like(ws))\n",
    "    \n",
    "    qs = RotToQuat(axs)\n",
    "    return qs,wws"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def PowerMethod(mat_,v_): \n",
    "    \"\"\"\n",
    "    finds the dominant eigenvector of a matrix by the power method. \n",
    "    also returns resdiuum.   \n",
    "    Args: \n",
    "        mat_: a matrix \n",
    "        v_: a guess eigenvector. \n",
    "    \"\"\"\n",
    "    i0 = tf.constant(0)\n",
    "    v = safe_normalize(v_ + 1e-19)\n",
    "    cond = lambda i, m: i < 35\n",
    "    body = lambda i, vx: [i+1,tf.stop_gradient(safe_normalize(tf.einsum('ijk,ik->ij',mat_,vx)))]\n",
    "    i_f, vf = tf.while_loop(\n",
    "        cond, body, loop_vars=[i0, v],\n",
    "        shape_invariants=[i0.get_shape(), v_.get_shape()])\n",
    "    diff = vf - safe_normalize(tf.einsum('ijk,ik->ij',mat_,vf))   \n",
    "    return vf, tf.reduce_sum(diff*diff,axis=-1)\n",
    "\n",
    "def AverageQuatMat(quats,weights): \n",
    "    \"\"\"\n",
    "    quats is a [...,nquat,4] tensor of quaternions. \n",
    "    weights is a [...,nquat,1] tensor of weights. \n",
    "    constructs properly averaged quat as per \n",
    "    https://ntrs.nasa.gov/archive/nasa/casi.ntrs.nasa.gov/20070017872.pdf\n",
    "    \"\"\"\n",
    "    nquat = tf.shape(quats)[1]\n",
    "    qr = tf.reshape(quats,(-1,nquat,4))\n",
    "    wr = tf.reshape(weights,(-1,nquat,1))\n",
    "    wq = qr*wr\n",
    "    todiag = tf.matmul(tf.transpose(qr,perm=(0,2,1)),wq)\n",
    "    return todiag\n",
    "\n",
    "def TF_RotationBatch(thetas, phis, psis):\n",
    "    \"\"\"Returns a batch of uniform rotation matrices, and the angles of each.\n",
    "    Finds random unit vector and then random angle around it.\n",
    "\n",
    "    Args:     sz_: number of rotation matrices     max_dist: maximum\n",
    "    rotation in units of 2*Pi\n",
    "    \"\"\"\n",
    "    Pi = 3.14159265359\n",
    "    sz_ = tf.shape(thetas)\n",
    "    axes = tf.zeros(shape=sz_ + [3])\n",
    "    axes = tf.stack(\n",
    "        [\n",
    "            tf.sin(thetas) * tf.cos(phis),\n",
    "            tf.sin(thetas) * tf.sin(phis),\n",
    "            tf.cos(thetas)\n",
    "        ],\n",
    "        axis=-1)\n",
    "    ct = tf.cos(psis)\n",
    "    st = tf.sin(psis)\n",
    "    omct = 1.0 - ct\n",
    "    matrices = tf.reshape(\n",
    "        tf.stack(\n",
    "            [\n",
    "                ct + axes[:, 0] * axes[:, 0] * omct,\n",
    "                axes[:, 0] * axes[:, 1] * omct - axes[:, 2] * st,\n",
    "                axes[:, 0] * axes[:, 2] * omct + axes[:, 1] * st,\n",
    "                axes[:, 1] * axes[:, 0] * omct + axes[:, 2] * st,\n",
    "                ct + axes[:, 1] * axes[:, 1] * omct,\n",
    "                axes[:, 1] * axes[:, 2] * omct - axes[:, 0] * st,\n",
    "                axes[:, 2] * axes[:, 0] * omct - axes[:, 1] * st,\n",
    "                axes[:, 2] * axes[:, 1] * omct + axes[:, 0] * st,\n",
    "                ct + axes[:, 2] * axes[:, 2] * omct\n",
    "            ],\n",
    "            axis=-1), [sz_[0], 3, 3])\n",
    "    return matrices\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tf.Graph().as_default():        \n",
    "    xyzs_tf = tf.Variable(xyzs)\n",
    "    ZTF = tf.Variable(Zs)\n",
    "    \n",
    "    #thetas = tf.acos(2.0 * tf.random_uniform([batch_size], dtype=tf.float64) - 1.0)\n",
    "    #phis = tf.random_uniform([batch_size], dtype=tf.float64) * 2 * Pi\n",
    "    #psis = tf.random_uniform([batch_size], dtype=tf.float64) * 2 * Pi\n",
    "    thetas = tf.acos(2.0 * tf.random_uniform([batch_size], dtype=tf.float64) - 1.0)\n",
    "    phis = tf.random_uniform([batch_size], dtype=tf.float64) * 2 * Pi\n",
    "    psis = tf.random_uniform([batch_size], dtype=tf.float64) * 2 * Pi    \n",
    "    matrices = TF_RotationBatch(thetas, phis, psis)\n",
    "    xyzs_shifted = xyzs_tf - xyzs_tf[:, 0, :][:,tf.newaxis, :]\n",
    "    tmpxyzs = tf.einsum('ijk,ikl->ijl', xyzs_shifted, matrices)\n",
    "    \n",
    "    testtf = tf.constant([[[1., 2, 3], [2, 5, 6], [3, 6, 3]]],dtype=tf.float64)\n",
    "    testtf2 = tf.constant([[1,1,1]],dtype=tf.float64)\n",
    "    \n",
    "    dxyzs = tf.expand_dims(tmpxyzs, axis=2) - tf.expand_dims(tmpxyzs, axis=1) # mol X MaxNAtom X MaxNeigh X3\n",
    "    # Compute quats. \n",
    "    z1 = tf.tile(tf.expand_dims(ZTF, axis=2),(1,1,MaxNeigh,1))\n",
    "    z2 = tf.tile(tf.expand_dims(ZTF, axis=1),(1,MaxNeigh,1,1))\n",
    "    z1z2 = tf.concat([z1,z2],axis=-1)\n",
    "    quats, weights = PointQuats(dxyzs,z1z2)\n",
    "    nquat = tf.shape(quats)[1]\n",
    "    aqm = AverageQuatMat(quats,weights)\n",
    "    #aqm = tf.eye(4,batch_shape=tf.shape(quats)[:1],dtype=tf.float64)\n",
    "    aquats, resids = PowerMethod(aqm,quats[:,0,:])\n",
    "    rots = tf.reshape(QuatToRot(aquats),(batch_size, MaxNAtom,3,3))\n",
    "    # now transform the coords. \n",
    "    cdxyzs = tf.reshape(tf.einsum('mijk,milk->mijl', dxyzs, rots),(batch_size, MaxNAtom, MaxNeigh,3))\n",
    "\n",
    "    init = tf.global_variables_initializer()\n",
    "    sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))\n",
    "    sess.run(init)\n",
    "\n",
    "    #print sess.run(PowerMethod(testtf,testtf2))\n",
    "    #print sess.run(xyzs_tf)\n",
    "    #print sess.run(quats)\n",
    "    #print sess.run(aquats)\n",
    "    Ws = sess.run(weights) \n",
    "    D = sess.run(dxyzs)\n",
    "    CD = sess.run(cdxyzs)\n",
    "    #print sess.run(aquats).shape\n",
    "    #print sess.run(cdxyzs).shape\n",
    "    #rmats = sess.run(rots)\n",
    "\n",
    "    print sess.run(tf.gradients(tf.reduce_sum(dxyzs*dxyzs,axis=2),psis))\n",
    "    print sess.run(tf.gradients(tf.reduce_sum(cdxyzs*cdxyzs,axis=2),psis))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Ws"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print Ws.shape\n",
    "print Ws[2]\n",
    "with tf.Graph().as_default(): \n",
    "    print sess.run(tf.gradients(weights,psis))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "with tf.Graph().as_default(): \n",
    "    Xqs,Xmat,Xaq = sess.run([quats[0,0,:],aqm[0],aquats])\n",
    "    print Xqs,Xmat,Xaq[0]\n",
    "    print sess.run(tf.gradients(aquats[2,2],psis))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print Zs[0]\n",
    "print Ws[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print D[0,0,0],D[0,0,1]\n",
    "print CD[0,0,0],CD[0,0,1]\n",
    "print Ws[0,0,0],Ws[0,1,0]\n",
    "print np.linalg.norm(D[0,0,0]-D[0,0,1])\n",
    "print np.linalg.norm(CD[0,0,0]-CD[0,0,1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print D[0,1,3],D[0,1,2]\n",
    "print CD[0,1,3],CD[0,1,2]\n",
    "print np.linalg.norm(D[0,1,3]-D[0,1,2])\n",
    "print np.linalg.norm(CD[0,1,3]-CD[0,1,2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "import matplotlib.pyplot as plt\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "ax.scatter(D[0,0,0], D[0,0,1], D[0,0,0])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def polykern_sr_9_11(r):\n",
    "    \"\"\"\n",
    "    Only LR cutoff.\n",
    "    \"\"\"\n",
    "    LRInner = 9.\n",
    "    LROuter = 11.\n",
    "    a = 1595.69\n",
    "    b = -668.25\n",
    "    c = 104.625\n",
    "    d = -7.25\n",
    "    e = 0.1875\n",
    "    r2 = r * r\n",
    "    r3 = r2 * r\n",
    "    r4 = r3 * r\n",
    "    kernm1 = tf.where(tf.less(r, LRInner), tf.clip_by_value(1.0/(1e-19+r),-1e19,1e19),\n",
    "                (a + b * r + c * r2 + d * r3 + e * r4) /r)\n",
    "    kern0 = tf.where(tf.less_equal(r,0.),tf.zeros_like(r),kernm1)\n",
    "    kern = tf.where(tf.greater(r, LROuter), tf.zeros_like(r),kern0)\n",
    "    return kern"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tf.Graph().as_default():        \n",
    "    rs = tf.Variable(np.linspace(0.0,30.0,300),dtype=tf.float64)\n",
    "    init = tf.global_variables_initializer()\n",
    "    sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))\n",
    "    sess.run(init)\n",
    "    print sess.run(polykern_sr_9_11(rs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.linspace(0.0,40.0,100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with tf.Graph().as_default():        \n",
    "    rs = tf.Variable(np.linspace(0.0,30.0,300),dtype=tf.float64)\n",
    "    init = tf.global_variables_initializer()\n",
    "    sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))\n",
    "    sess.run(init)\n",
    "    \n",
    "def QuatToRot(q):\n",
    "    \"\"\"\n",
    "    a_ ... X 4 tensor of quaternions\n",
    "    this generates a ... X 3 X 3 of rotation matrices.\n",
    "    \"\"\"\n",
    "    tmp=tf.stack([1 - 2.*(q[...,2]*q[...,2] + q[...,3]*q[...,3]), 2*(q[...,1]*q[...,2] - q[...,3]*q[...,0]),\n",
    "    2*(q[...,1]*q[...,3] + q[...,2]*q[...,0]),2*(q[...,1]*q[...,2] + q[...,3]*q[...,0]), 1 - 2.*(q[...,1]*q[...,1] + q[...,3]*q[...,3]),\n",
    "    2*(q[...,2]*q[...,3] - q[...,1]*q[...,0]),2*(q[...,1]*q[...,3] - q[...,2]*q[...,0]), 2*(q[...,2]*q[...,3] + q[...,1]*q[...,0]),\n",
    "    1 - 2.*(q[...,1]*q[...,1] + q[...,2]*q[...,2])],axis=-1)\n",
    "    return tf.reshape(tmp,[-1,3,3])\n",
    "\n",
    "def VectorsToOrient(v1,v2):\n",
    "    v1n = safe_normalize(v1)\n",
    "    v2n = safe_normalize(v2)\n",
    "    v3 = safe_normalize(tf.cross(v1n, v2n)+tf.constant(np.array([0., 0., 1e-19]), dtype=tf.float64))\n",
    "    # Compute the average of v1, v2, and their projections onto the\n",
    "    # plane.\n",
    "    v_av = (v1n + v2n) / 2.0\n",
    "    v_av = safe_normalize(v_av)\n",
    "    # Rotate pi/4 cw and ccw to obtain v1,v2\n",
    "    first = TF_AxisAngleRotation(v3, v_av, tf.constant(-Pi / 4., dtype=tf.float64))\n",
    "    second = TF_AxisAngleRotation(v3, v_av,tf.constant(Pi / 4., dtype=tf.float64))\n",
    "    vs = tf.concat([first[:, tf.newaxis, :], second[:, tf.newaxis, :],v3[:, tf.newaxis, :]],axis=1)\n",
    "    return vs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tf.Graph().as_default():        \n",
    "    rs = tf.Variable(np.array([[.3,.3,.3,.3]]),dtype=tf.float64)\n",
    "    init = tf.global_variables_initializer()\n",
    "    sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))\n",
    "    sess.run(init)\n",
    "    print sess.run(QuatToRot(rs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[-4.52013408 -0.22379047 -1.69580106]\n",
      " [-3.55814909  0.54129489  0.76857379]]\n",
      "[[[-0.6665386  -0.19328034 -0.71997847]\n",
      "  [-0.74178371  0.26789563  0.61480799]\n",
      "  [ 0.07404879  0.94386156 -0.32193499]]\n",
      "\n",
      " [[ 0.45647495 -0.70744694  0.53958267]\n",
      "  [-0.7966752  -0.05496639  0.60190308]\n",
      "  [-0.39615558 -0.70462581 -0.58869621]]]\n",
      "[[ 0.26430193  0.31124779 -0.75106079 -0.5188227 ]\n",
      " [ 0.45078053 -0.72459257  0.51895445 -0.04948542]]\n",
      "[[[-0.6665386  -0.19328034 -0.71997847]\n",
      "  [-0.74178371  0.26789563  0.61480799]\n",
      "  [ 0.07404879  0.94386156 -0.32193499]]\n",
      "\n",
      " [[ 0.45647495 -0.70744694  0.53958267]\n",
      "  [-0.7966752  -0.05496639  0.60190308]\n",
      "  [-0.39615558 -0.70462581 -0.58869621]]]\n"
     ]
    }
   ],
   "source": [
    "def safe_normalize(x_):\n",
    "    nrm = tf.clip_by_value(tf.norm(x_,axis=-1,keepdims=True),1e-36,1e36)\n",
    "    nrm_ok = tf.logical_and(tf.not_equal(nrm,0.),tf.logical_not(tf.is_nan(nrm)))\n",
    "    safe_nrm = tf.where(nrm_ok,nrm,tf.ones_like(nrm))\n",
    "    return x_*tf.where(nrm_ok,1.0/safe_nrm,tf.zeros_like(nrm))\n",
    "def safe_inv_norm(x_):\n",
    "    nrm = tf.clip_by_value(tf.norm(x_,axis=-1,keepdims=True),1e-36,1e36)\n",
    "    nrm_ok = tf.logical_and(tf.not_equal(nrm,0.),tf.logical_not(tf.is_nan(nrm)))\n",
    "    safe_nrm = tf.where(nrm_ok,nrm,tf.ones_like(nrm))\n",
    "    return tf.where(nrm_ok,1.0/safe_nrm,tf.zeros_like(nrm))\n",
    "def safe_norm(x_):\n",
    "    nrm = tf.clip_by_value(tf.norm(x_, axis=-1, keepdims=True), 1e-36, 1e36)\n",
    "    nrm_ok = tf.logical_and(\n",
    "        tf.not_equal(nrm, 0.), tf.logical_not(tf.is_nan(nrm)))\n",
    "    safe_nrm = tf.where(nrm_ok, nrm, tf.zeros_like(nrm))\n",
    "    return safe_nrm\n",
    "def RotToQuat(axes_):\n",
    "    \"\"\"\n",
    "    axes is a ... X 3 3 tensor of axes\n",
    "    this generates a ... X 4 tensor of quaternions.\n",
    "    which are 1:1 with those axes.\n",
    "    \"\"\"\n",
    "    w = (1./2.)*tf.sqrt(1e-19+tf.abs(1 + axes_[...,0, 0] + axes_[...,1, 1] + axes_[...,2, 2]))\n",
    "    x = tf.sign(axes_[...,2, 1] - axes_[...,1, 2])*tf.abs(0.5*tf.sqrt(1e-19+tf.abs(1.0 + axes_[...,0, 0] - axes_[...,1, 1] - axes_[...,2, 2])))\n",
    "    y = tf.sign(axes_[...,0, 2] - axes_[...,2, 0])*tf.abs(0.5*tf.sqrt(1e-19+tf.abs(1.0 - axes_[...,0, 0] + axes_[...,1, 1] - axes_[...,2, 2])))\n",
    "    z = tf.sign(axes_[...,1, 0] - axes_[...,0, 1])*tf.abs(0.5*tf.sqrt(1e-19+tf.abs(1.0 - axes_[...,0, 0] - axes_[...,1, 1] + axes_[...,2, 2])))\n",
    "    return tf.stack([w,x,y,z],axis=-1)\n",
    "def TF_AxisAngleRotation(axes, vecs, psis):\n",
    "    \"\"\"Returns a batch of uniform rotation matrices, and the angles of each.\n",
    "    Finds random unit vector and then random angle around it.\n",
    "\n",
    "    Args:     sz_: number of rotation matrices     max_dist: maximum\n",
    "    rotation in units of 2*Pi\n",
    "    \"\"\"\n",
    "    ct = tf.cos(psis)\n",
    "    st = tf.sin(psis)\n",
    "    omct = 1.0 - ct\n",
    "    sz_ = tf.shape(axes)\n",
    "    matrices = tf.reshape(\n",
    "        tf.stack(\n",
    "            [\n",
    "                ct + axes[:, 0] * axes[:, 0] * omct,\n",
    "                axes[:, 0] * axes[:, 1] * omct - axes[:, 2] * st,\n",
    "                axes[:, 0] * axes[:, 2] * omct + axes[:, 1] * st,\n",
    "                axes[:, 1] * axes[:, 0] * omct + axes[:, 2] * st,\n",
    "                ct + axes[:, 1] * axes[:, 1] * omct,\n",
    "                axes[:, 1] * axes[:, 2] * omct - axes[:, 0] * st,\n",
    "                axes[:, 2] * axes[:, 0] * omct - axes[:, 1] * st,\n",
    "                axes[:, 2] * axes[:, 1] * omct + axes[:, 0] * st,\n",
    "                ct + axes[:, 2] * axes[:, 2] * omct\n",
    "            ],\n",
    "            axis=-1), [sz_[0], 3, 3])\n",
    "    return tf.einsum('ijk,ik->ij', matrices, vecs)\n",
    "with tf.Graph().as_default():        \n",
    "    xyzs = tf.Variable(np.random.random((batch_size,MaxNAtom,3))*7.0 - 5.0)\n",
    "    init = tf.global_variables_initializer()\n",
    "    sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))\n",
    "    sess.run(init)\n",
    "    print sess.run(xyzs[0,:2])\n",
    "    print sess.run(VectorsToOrient(xyzs[:,0],xyzs[:,1]))\n",
    "    print sess.run(RotToQuat(VectorsToOrient(xyzs[:,0],xyzs[:,1])))\n",
    "    print sess.run(QuatToRot(RotToQuat(VectorsToOrient(xyzs[:,0],xyzs[:,1]))))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def slerp(v0, v1, t=0.05):\n",
    "    \"\"\"\n",
    "    Interpolate between quaternions v0 and v1.\n",
    "    \"\"\"\n",
    "    v0 = safe_normalize(v0)\n",
    "    v1 = safe_normalize(v1)\n",
    "    dot = tf.reduce_sum(v0*v1,axis=-1,keepdims=True)\n",
    "    # If the dot product is negative, slerp won't take\n",
    "    # the shorter path. Note that v1 and -v1 are equivalent when\n",
    "    # the negation is applied to all four components. Fix by\n",
    "    # reversing one quaternion.\n",
    "    signflip = tf.where(tf.less_equal(dot,0.),-1.*tf.ones_like(dot),tf.ones_like(dot))\n",
    "    v1 *= signflip\n",
    "    dot *= signflip\n",
    "    # Linear answer.\n",
    "    linq = safe_normalize(v0 + t*(v1-v0))\n",
    "    #\n",
    "    sdot = tf.clip_by_value(dot,-1.0,1.0)\n",
    "    theta_0 = tf.acos(sdot)\n",
    "    theta = theta_0*t\n",
    "    sin_theta = tf.sin(theta)\n",
    "    sin_theta_0 = tf.sin(theta_0)\n",
    "    s0 = tf.cos(theta) - dot * sin_theta / (sin_theta_0+1e-19)\n",
    "    s1 = sin_theta / (sin_theta_0+1e-19)\n",
    "    sq = safe_normalize((s0 * v0) + (s1 * v1))\n",
    "    #\n",
    "    DOT_THRESHOLD = 0.9995\n",
    "    tdot = tf.concat([dot,dot,dot,dot],axis=-1)\n",
    "    slerpd = tf.where(tf.greater(tdot,DOT_THRESHOLD),linq,sq)\n",
    "    ttiled = tf.concat([t,t,t,t],axis=-1)\n",
    "    slerpdorv1 = tf.where(tf.greater(ttiled,1.0-1e-14),v1,slerpd)\n",
    "    return tf.where(tf.less(ttiled,1e-14),v0,slerpdorv1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 0.10908506  0.92581498  0.00113723 -0.36189195]\n",
      " [ 0.16528367 -0.00838225 -0.77582942  0.60885118]]\n",
      "[[ 0.5700113   0.51692443 -0.53746492  0.34497495]\n",
      " [ 0.0237009   0.67409075 -0.21600999 -0.70596007]]\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "Can't concatenate scalars (use tf.stack instead) for 'concat_3' (op: 'ConcatV2') with input shapes: [], [], [], [], [].",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-9-5158a17d868d>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      8\u001b[0m     \u001b[0;32mprint\u001b[0m \u001b[0msess\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mq1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m     \u001b[0;32mprint\u001b[0m \u001b[0msess\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mq2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 10\u001b[0;31m     \u001b[0;32mprint\u001b[0m \u001b[0msess\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mslerp\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mq1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mq2\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m0.0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     11\u001b[0m     \u001b[0;32mprint\u001b[0m \u001b[0msess\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mslerp\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mq1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mq2\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m0.2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     12\u001b[0m     \u001b[0;32mprint\u001b[0m \u001b[0msess\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mslerp\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mq1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mq2\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m0.4\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-8-e48c3844a2a9>\u001b[0m in \u001b[0;36mslerp\u001b[0;34m(v0, v1, t)\u001b[0m\n\u001b[1;32m     28\u001b[0m     \u001b[0mtdot\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mdot\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mdot\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mdot\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mdot\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0maxis\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     29\u001b[0m     \u001b[0mslerpd\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwhere\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgreater\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtdot\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mDOT_THRESHOLD\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mlinq\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0msq\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 30\u001b[0;31m     \u001b[0mttiled\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mt\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mt\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mt\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mt\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0maxis\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     31\u001b[0m     \u001b[0mslerpdorv1\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwhere\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgreater\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mttiled\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m1.0\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1e-14\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mv1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mslerpd\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     32\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwhere\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mless\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mttiled\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m1e-14\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mv0\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mslerpdorv1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/Library/Python/2.7/site-packages/tensorflow/python/ops/array_ops.pyc\u001b[0m in \u001b[0;36mconcat\u001b[0;34m(values, axis, name)\u001b[0m\n\u001b[1;32m   1173\u001b[0m               tensor_shape.scalar())\n\u001b[1;32m   1174\u001b[0m       \u001b[0;32mreturn\u001b[0m \u001b[0midentity\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mscope\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1175\u001b[0;31m   \u001b[0;32mreturn\u001b[0m \u001b[0mgen_array_ops\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_concat_v2\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maxis\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0maxis\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1176\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1177\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/Library/Python/2.7/site-packages/tensorflow/python/ops/gen_array_ops.pyc\u001b[0m in \u001b[0;36m_concat_v2\u001b[0;34m(values, axis, name)\u001b[0m\n\u001b[1;32m    623\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0m_ctx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0min_graph_mode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    624\u001b[0m     _, _, _op = _op_def_lib._apply_op_helper(\n\u001b[0;32m--> 625\u001b[0;31m         \"ConcatV2\", values=values, axis=axis, name=name)\n\u001b[0m\u001b[1;32m    626\u001b[0m     \u001b[0m_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_op\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moutputs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    627\u001b[0m     \u001b[0m_inputs_flat\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_op\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/Library/Python/2.7/site-packages/tensorflow/python/framework/op_def_library.pyc\u001b[0m in \u001b[0;36m_apply_op_helper\u001b[0;34m(self, op_type_name, name, **keywords)\u001b[0m\n\u001b[1;32m    785\u001b[0m         op = g.create_op(op_type_name, inputs, output_types, name=scope,\n\u001b[1;32m    786\u001b[0m                          \u001b[0minput_types\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minput_types\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mattrs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mattr_protos\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 787\u001b[0;31m                          op_def=op_def)\n\u001b[0m\u001b[1;32m    788\u001b[0m       \u001b[0;32mreturn\u001b[0m \u001b[0moutput_structure\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mop_def\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_stateful\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mop\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    789\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/Library/Python/2.7/site-packages/tensorflow/python/framework/ops.pyc\u001b[0m in \u001b[0;36mcreate_op\u001b[0;34m(self, op_type, inputs, dtypes, input_types, name, attrs, op_def, compute_shapes, compute_device)\u001b[0m\n\u001b[1;32m   3271\u001b[0m         op_def=op_def)\n\u001b[1;32m   3272\u001b[0m     self._create_op_helper(ret, compute_shapes=compute_shapes,\n\u001b[0;32m-> 3273\u001b[0;31m                            compute_device=compute_device)\n\u001b[0m\u001b[1;32m   3274\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mret\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3275\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/Library/Python/2.7/site-packages/tensorflow/python/framework/ops.pyc\u001b[0m in \u001b[0;36m_create_op_helper\u001b[0;34m(self, op, compute_shapes, compute_device)\u001b[0m\n\u001b[1;32m   3311\u001b[0m     \u001b[0;31m# compute_shapes argument.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3312\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mop\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_c_op\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mcompute_shapes\u001b[0m\u001b[0;34m:\u001b[0m  \u001b[0;31m# pylint: disable=protected-access\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3313\u001b[0;31m       \u001b[0mset_shapes_for_outputs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mop\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   3314\u001b[0m     \u001b[0;31m# TODO(b/XXXX): move to Operation.__init__ once _USE_C_API flag is removed.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3315\u001b[0m     \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_add_op\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mop\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/Library/Python/2.7/site-packages/tensorflow/python/framework/ops.pyc\u001b[0m in \u001b[0;36mset_shapes_for_outputs\u001b[0;34m(op)\u001b[0m\n\u001b[1;32m   2499\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0m_set_shapes_for_outputs_c_api\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mop\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2500\u001b[0m   \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2501\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0m_set_shapes_for_outputs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mop\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2502\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2503\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/Library/Python/2.7/site-packages/tensorflow/python/framework/ops.pyc\u001b[0m in \u001b[0;36m_set_shapes_for_outputs\u001b[0;34m(op)\u001b[0m\n\u001b[1;32m   2472\u001b[0m       \u001b[0mshape_func\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_call_cpp_shape_fn_and_require_op\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2473\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2474\u001b[0;31m   \u001b[0mshapes\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mshape_func\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mop\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2475\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0mshapes\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2476\u001b[0m     raise RuntimeError(\n",
      "\u001b[0;32m/Library/Python/2.7/site-packages/tensorflow/python/framework/ops.pyc\u001b[0m in \u001b[0;36mcall_with_requiring\u001b[0;34m(op)\u001b[0m\n\u001b[1;32m   2402\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2403\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mcall_with_requiring\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mop\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2404\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mcall_cpp_shape_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mop\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrequire_shape_fn\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2405\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2406\u001b[0m   \u001b[0m_call_cpp_shape_fn_and_require_op\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcall_with_requiring\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/Library/Python/2.7/site-packages/tensorflow/python/framework/common_shapes.pyc\u001b[0m in \u001b[0;36mcall_cpp_shape_fn\u001b[0;34m(op, require_shape_fn)\u001b[0m\n\u001b[1;32m    625\u001b[0m     res = _call_cpp_shape_fn_impl(op, input_tensors_needed,\n\u001b[1;32m    626\u001b[0m                                   \u001b[0minput_tensors_as_shapes_needed\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 627\u001b[0;31m                                   require_shape_fn)\n\u001b[0m\u001b[1;32m    628\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mres\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdict\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    629\u001b[0m       \u001b[0;31m# Handles the case where _call_cpp_shape_fn_impl calls unknown_shape(op).\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/Library/Python/2.7/site-packages/tensorflow/python/framework/common_shapes.pyc\u001b[0m in \u001b[0;36m_call_cpp_shape_fn_impl\u001b[0;34m(op, input_tensors_needed, input_tensors_as_shapes_needed, require_shape_fn)\u001b[0m\n\u001b[1;32m    689\u001b[0m       \u001b[0mmissing_shape_fn\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mTrue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    690\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 691\u001b[0;31m       \u001b[0;32mraise\u001b[0m \u001b[0mValueError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0merr\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmessage\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    692\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    693\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0mmissing_shape_fn\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mValueError\u001b[0m: Can't concatenate scalars (use tf.stack instead) for 'concat_3' (op: 'ConcatV2') with input shapes: [], [], [], [], []."
     ]
    }
   ],
   "source": [
    "with tf.Graph().as_default():        \n",
    "    xyzs = tf.Variable(np.random.random((batch_size,MaxNAtom,3))*7.0 - 5.0)\n",
    "    q1=RotToQuat(VectorsToOrient(xyzs[:,0],xyzs[:,1]))\n",
    "    q2=RotToQuat(VectorsToOrient(xyzs[:,1],xyzs[:,2]))\n",
    "    init = tf.global_variables_initializer()\n",
    "    sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))\n",
    "    sess.run(init)\n",
    "    print sess.run(q1)\n",
    "    print sess.run(q2)\n",
    "    print sess.run(slerp(q1,q2,0.0))\n",
    "    print sess.run(slerp(q1,q2,0.2))\n",
    "    print sess.run(slerp(q1,q2,0.4))\n",
    "    print sess.run(slerp(q1,q2,0.6))    \n",
    "    print sess.run(slerp(q1,q2,1.0))        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}