Bayesian-Structure-Alignment.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementation of Bayesian structure alignment\n",
    "\n",
    "The goal of this notebook is to implement an algorithm to align two protein structures. This work is based heavily on the Pyro implementation of this problem published by Moreta et al (https://www.biorxiv.org/content/early/2019/06/04/575431.full.pdf). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Populating the interactive namespace from numpy and matplotlib\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/jgreisman/anaconda2/envs/md/lib/python3.6/site-packages/IPython/core/magics/pylab.py:160: UserWarning: pylab import has clobbered these variables: ['dist']\n",
      "`%matplotlib` prevents importing * from pylab and numpy\n",
      "  \"\\n`%matplotlib` prevents importing * from pylab and numpy\"\n"
     ]
    }
   ],
   "source": [
    "%pylab inline\n",
    "\n",
    "import torch\n",
    "import pyro\n",
    "import pyro.distributions as dist\n",
    "from pyro.infer import EmpiricalMarginal, SVI, TracePredictive\n",
    "from pyro.infer.mcmc import MCMC, NUTS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "atoms = 50\n",
    "iters = 1000\n",
    "X1 = np.random.random((atoms, 3))\n",
    "X2 = X1 + np.random.normal(0, 0.01, (atoms, 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "X1,X2 = torch.tensor(X1, dtype=torch.float32), torch.tensor(X2, dtype=torch.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 123,
   "metadata": {},
   "outputs": [],
   "source": [
    "def u_to_quat_to_R(u):\n",
    "    [t1,t2,r1,r2] = [\n",
    "        2*np.pi*u[1],\n",
    "        2*np.pi*u[2],\n",
    "        torch.sqrt(1-u[0]),\n",
    "        torch.sqrt(u[0]),\n",
    "    ]\n",
    "\n",
    "    [w,x,y,z] = [\n",
    "        r2*torch.cos(t2),\n",
    "        r1*torch.sin(t1),\n",
    "        r1*torch.cos(t1),\n",
    "        r2*torch.sin(t2),\n",
    "    ]\n",
    "\n",
    "    R = torch.tensor([\n",
    "        [w*w + x*x - y*y - z*z,     2*(x*y-w*z),     2*(x*z+w*y)],\n",
    "        [          2*(x*y+w*z), w*w-x*x+y*y-z*z,   2*(y*z - w*x)],\n",
    "        [          2*(x*z-w*y),     2*(y*z+w*x), w*w-x*x-y*y+z*z],\n",
    "    ], dtype=torch.float32)\n",
    "    return R"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model():\n",
    "    N = 50\n",
    "    # Prior over mean M, with N=number of atoms\n",
    "    M = pyro.sample(\"M\", dist.StudentT(1,0,3).expand_by([N,3]).to_event(2))\n",
    "    # Prior over variances U\n",
    "    U = pyro.sample(\"U\", dist.HalfNormal(1).expand_by([N]).to_event(1))\n",
    "    U = U.reshape(N,1).repeat(1,3).view(-1)\n",
    "    # Prior over translation T\n",
    "    T = pyro.sample(\"T\", dist.Normal(0,1).expand_by([3]).to_event(1))\n",
    "    # Prior over rotation R\n",
    "    u = pyro.sample(\"u\",dist.Uniform(0,1).expand_by([3]).to_event(1))\n",
    "    # Transformation: turn u via a unit quaternion into a rotation R\n",
    "    R = u_to_quat_to_R(u)\n",
    "    # Transformation: rotate and translate M for X2\n",
    "    M_RT = M @ R + T\n",
    "    # Likelihood\n",
    "    with pyro.plate(\"plate_students\", N*3,dim= -1):\n",
    "        pyro.sample(\"X1\", dist.StudentT(1, M.view(-1), U), obs=X1.view(-1))\n",
    "        pyro.sample(\"X2\", dist.StudentT(1, M_RT.view(-1), U), obs=X2.view(-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyro.contrib.autoguide import AutoDelta\n",
    "from pyro.infer import Trace_ELBO\n",
    "from pyro.optim import Adam\n",
    "\n",
    "guide = AutoDelta(model)\n",
    "optim = Adam({'lr': 0.03})\n",
    "svi = pyro.infer.SVI(model, guide, optim, loss=Trace_ELBO())\n",
    "\n",
    "\n",
    "losses=[]\n",
    "for i in range(10000):\n",
    "    losses.append(svi.step())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x120a5e128>]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYEAAAD8CAYAAACRkhiPAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvOIA7rQAAIABJREFUeJzt3Xl8VfWd//HXJ3sgQIKEHQwIiIjKEhHrUiuoCG3VebQz0I7S1g7a2pnazvLTOh06dpwynba2TFutrUy1M3WrWnlUFJFanbZuQZF9CYsQCBAMJEBC1s/vj3sSL+RmX2645/18PO4j537P9577PTnoO+f7Ped7zN0REZFwSop3A0REJH4UAiIiIaYQEBEJMYWAiEiIKQREREJMISAiEmIKARGREFMIiIiEmEJARCTEUuLdgNYMGjTI8/Ly4t0MEZEzxpo1aw67e25b6vb6EMjLy6OgoCDezRAROWOY2fttravuIBGREFMIiIiEmEJARCTEFAIiIiGmEBARCbFWQ8DMRpnZK2a22cw2mtlXg/KBZrbKzLYHP3OCcjOzpWZWaGbrzGxa1LYWBvW3m9nC7tstERFpi7acCdQCf+/u5wEzgTvMbBJwF7Da3ccDq4P3ANcD44PXIuABiIQGsBi4BJgBLG4IDhERiY9WQ8Ddi939nWD5GLAZGAHcADwSVHsEuDFYvgF41CPeALLNbBhwHbDK3Uvd/QiwCpjTpXsTZenq7by6raS7Ni8ikhDaNSZgZnnAVOBNYIi7F0MkKIDBQbURwN6ojxUFZc2Vd4sHX93B/ykERERa1OYQMLMs4GngTncvb6lqjDJvoTzWdy0yswIzKygp6dj/yNNSkqiuq+/QZ0VEwqJNIWBmqUQC4H/d/Zmg+GDQzUPw81BQXgSMivr4SGB/C+VNuPtD7p7v7vm5uW2a/qKJ1OQkahQCIiItasvVQQY8DGx29x9ErVoONFzhsxB4Lqr8luAqoZlAWdBdtBK41sxyggHha4OybpGWnERVrUJARKQlbZlA7jLgZmC9ma0Nyr4BLAGeNLNbgT3Ap4N1K4C5QCFQAXwewN1LzezbwNtBvXvdvbRL9iKG9JQkqhUCIiItajUE3P2PxO7PB5gVo74DdzSzrWXAsvY0sKPSFAIiIq1K2DuGNSYgItK6hA0BXR0kItK6xA2BZHUHiYi0JnFDQGMCIiKtStgQSE1Oorou5r1oIiISSNgQiFwiWhfvZoiI9GoJGwIaGBYRaV3ihoAGhkVEWpWwIZCaYtRoTEBEpEUJGwJpyck6ExARaUXihoAuERURaVVih0BdPZGpjEREJJaEDYH0lMiuaTppEZHmJWwIZKQmA3CyRvcKiIg0J4FDILJrJ2t0JiAi0pzEDYEUnQmIiLQmcUMg6A7SmICISPMSOAQauoN0JiAi0pyEDYFMDQyLiLQqYUMgvSEE1B0kItKsVkPAzJaZ2SEz2xBV9oSZrQ1eu81sbVCeZ2aVUesejPrMdDNbb2aFZrbUzJp7eH2XUHeQiEjrUtpQ55fAj4FHGwrc/a8als3s+0BZVP0d7j4lxnYeABYBbwArgDnAC+1vctvoPgERkda1eibg7q8BpbHWBX/N/yXwWEvbMLNhQH93f90j8zg8CtzY/ua2XePVQbpPQESkWZ0dE7gCOOju26PKxpjZu2b2qpldEZSNAIqi6hQFZd0mI5g24qSeLiYi0qy2dAe1ZAGnngUUA6Pd/QMzmw781szOB2L1/zc7s5uZLSLSdcTo0aM71DB1B4mItK7DZwJmlgL8BfBEQ5m7V7n7B8HyGmAHMIHIX/4joz4+Etjf3Lbd/SF3z3f3/Nzc3A61ryEEKqvVHSQi0pzOdAfNBra4e2M3j5nlmllysDwWGA/sdPdi4JiZzQzGEW4BnuvEd7cqOcnISE3iRHVtd36NiMgZrS2XiD4GvA6ca2ZFZnZrsGo+TQeErwTWmdl7wG+A2929YVD5S8AvgEIiZwjddmVQg6z0FI5XKQRERJrT6piAuy9opvxzMcqeBp5upn4BMLmd7euUrPQUTigERESalbB3DAP0TU/h+EmFgIhIcxI/BHQmICLSrIQOgaz0FA0Mi4i0IKFDoG96CieqdJ+AiEhzEjoEstKT1R0kItKChA6Bvmm6OkhEpCUJHQJZGSlUVNdRV9/sDBUiIqGW2CGQHrkNQoPDIiKxJXQI9G0IAXUJiYjEpBAQEQmxhA6BrPTITKLHdZmoiEhMCR0CfdN0JiAi0pLEDoGgO0j3CoiIxJbQIdAvIwgBTSInIhJTQodA/4xUAMpP1sS5JSIivVNCh0DDmUB5pc4ERERiSegQSElOIis9hbJKnQmIiMSS0CEA0D8jRd1BIiLNSPwQyEzVmYCISDNCEQLlCgERkZhaDQEzW2Zmh8xsQ1TZt8xsn5mtDV5zo9bdbWaFZrbVzK6LKp8TlBWa2V1dvyuxDdCZgIhIs9pyJvBLYE6M8vvdfUrwWgFgZpOA+cD5wWd+ambJZpYM/AS4HpgELAjqdrv+GToTEBFpTkprFdz9NTPLa+P2bgAed/cqYJeZFQIzgnWF7r4TwMweD+puaneL22lAZirlullMRCSmzowJfMXM1gXdRTlB2Qhgb1SdoqCsufKYzGyRmRWYWUFJSUknmgj9M1M4XlVLbV19p7YjIpKIOhoCDwDnAFOAYuD7QbnFqOstlMfk7g+5e7675+fm5nawiREDMiN3DR/T2YCISBMdCgF3P+jude5eD/ycD7t8ioBRUVVHAvtbKO92DVNHaHBYRKSpDoWAmQ2LensT0HDl0HJgvpmlm9kYYDzwFvA2MN7MxphZGpHB4+Udb3bbNZwJ6IYxEZGmWh0YNrPHgKuAQWZWBCwGrjKzKUS6dHYDtwG4+0Yze5LIgG8tcIe71wXb+QqwEkgGlrn7xi7fmxj6Z+pMQESkOW25OmhBjOKHW6h/H3BfjPIVwIp2ta4LNJ4JaBI5EZEmQnDHcCTndCYgItJUwoeAxgRERJqX8CGQmZpMSpLpTEBEJIaEDwEzi9w1rBAQEWki4UMANJ20iEhzQhMCmj9IRKSpcIRAhh4xKSISSyhCYEBmKscUAiIiTYQiBDQmICISWyhCIPJMgRrcm524VEQklEIRAv0zUqmpcypr6uLdFBGRXiUUIZDTJ3LX8NEKdQmJiEQLRwj0TQOg9ER1nFsiItK7hCIEBgYhcKRCISAiEi1UIaAzARGRU4UjBPooBEREYglFCPTPTCXJ4IhCQETkFKEIgeQkI7tPGqUaExAROUUoQgAil4mqO0hE5FShCYGBfdM4ckL3CYiIRGs1BMxsmZkdMrMNUWX/aWZbzGydmT1rZtlBeZ6ZVZrZ2uD1YNRnppvZejMrNLOlZmbds0ux5fRJ0yWiIiKnacuZwC+BOaeVrQImu/uFwDbg7qh1O9x9SvC6Par8AWARMD54nb7NbpXTJ03dQSIip2k1BNz9NaD0tLKX3L3hKS1vACNb2oaZDQP6u/vrHpnF7VHgxo41uWNy+qZxtEKTyImIROuKMYEvAC9EvR9jZu+a2atmdkVQNgIoiqpTFJTFZGaLzKzAzApKSkq6oIkwKCuN6rp6yiv1hDERkQadCgEzuweoBf43KCoGRrv7VODrwK/NrD8Qq/+/2T/J3f0hd8939/zc3NzONLHRoKx0AD44UdUl2xMRSQQpHf2gmS0EPg7MCrp4cPcqoCpYXmNmO4AJRP7yj+4yGgns7+h3d0R2MJPoEc0kKiLSqENnAmY2B/h/wCfdvSKqPNfMkoPlsUQGgHe6ezFwzMxmBlcF3QI81+nWt0NOMHWE7hoWEflQq2cCZvYYcBUwyMyKgMVErgZKB1YFV3q+EVwJdCVwr5nVAnXA7e7eMKj8JSJXGmUSGUOIHkfodo2TyOkyURGRRq2GgLsviFH8cDN1nwaebmZdATC5Xa3rQppJVESkqdDcMdw3PYV+6SkcKDsZ76aIiPQaoQkBgCEDMiguq4x3M0REeo1QhcCwARk6ExARiRKqEBiRncm+owoBEZEGoQqB4dmZHD5excmaung3RUSkVwhVCAwbkAGgLiERkUCoQmBEdiYA+zU4LCIChCwEhjeEgMYFRESAkIXA0KA7aP9RnQmIiEDIQiAjNZlBWWkKARGRQKhCACJdQvsUAiIiQBhDYEAmxbo6SEQECGMIZGey/2ilHjMpIkIoQyCDiuo6yir1cBkRkdCFwLABkctEd39Q0UpNEZHEF7oQqA+6ge5+Zn2cWyIiEn+hC4E5k4cCsLm4PM4tERGJv9CFQGpyZJenn50T55aIiMRf6EIAYNbEwZyoqo13M0RE4q5NIWBmy8zskJltiCobaGarzGx78DMnKDczW2pmhWa2zsymRX1mYVB/u5kt7PrdaZsROZnsKdXAsIhIW88EfgnMOa3sLmC1u48HVgfvAa4HxgevRcADEAkNYDFwCTADWNwQHD1teHYmFdV1OhsQkdBrUwi4+2tA6WnFNwCPBMuPADdGlT/qEW8A2WY2DLgOWOXupe5+BFhF02DpEblZ6QAcLNedwyISbp0ZExji7sUAwc/BQfkIYG9UvaKgrLnyHjcyJ3KvwN4jmkNIRMKtOwaGLUaZt1DedANmi8yswMwKSkpKurRxAHmD+gKw54MTXb5tEZEzSWdC4GDQzUPw81BQXgSMiqo3EtjfQnkT7v6Qu+e7e35ubm4nmhjb4H7ppKck6a5hEQm9zoTAcqDhCp+FwHNR5bcEVwnNBMqC7qKVwLVmlhMMCF8blPU4M6Oqtp6H/7grHl8vItJrpLSlkpk9BlwFDDKzIiJX+SwBnjSzW4E9wKeD6iuAuUAhUAF8HsDdS83s28DbQb173f30wWYREelBbQoBd1/QzKpZMeo6cEcz21kGLGtz67rRqIGZ7C2t5L29R7loVHa8myMiEhehvGMY4KMTImMNG/aXxbklIiLxE9oQuHlmHgBVNfXxbYiISByFNgQmDMliQGYq2w4ei3dTRETiJrQhYGaUVdbw+Nt7W68sIpKgQhsCAPnBdNLlJ/WoSREJp1CHwM2Xng3AUwVFcW6JiEh8hDoEpo2OnAn8+s3349wSEZH4CHUIjBrYB4AdJZpDSETCKdQhEE3PFhCRMAp9CCxdMBWA/Uc1rbSIhE/oQyApmOB67d6j8W2IiEgchD4EUpIiv4KXNx+Mc0tERHpe6EPg6omRB6IVHjoe55aIiPS80IdAWkrkVzAoeO6wiEiYhD4EGry5q1RnAyISOgqBKLN/8Gq8myAi0qMUAkDBP8+OdxNEROJCIYDGA0QkvBQCp3lxQ3G8myAi0mMUAoHPXjIagNv/5x2OVlTHuTUiIj2jwyFgZuea2dqoV7mZ3Wlm3zKzfVHlc6M+c7eZFZrZVjO7rmt2oWvce8PkxuWlqwvj2BIRkZ7T4RBw963uPsXdpwDTgQrg2WD1/Q3r3H0FgJlNAuYD5wNzgJ+aWXLnmt91kpOMP911NQDL/rRLj50UkVDoqu6gWcAOd29pYv4bgMfdvcrddwGFwIwu+v4uMSI7s3H52vtfi2NLRER6RleFwHzgsaj3XzGzdWa2zMxygrIRQPQDfYuCsl7lqdsvbVw+fLwqji0REel+nQ4BM0sDPgk8FRQ9AJwDTAGKge83VI3xcW9mm4vMrMDMCkpKSjrbxHa5OG9g43L+v73MPk0xLSIJrCvOBK4H3nH3gwDuftDd69y9Hvg5H3b5FAGjoj43Etgfa4Pu/pC757t7fm5ubhc0sX0K77u+cXn293UXsYgkrq4IgQVEdQWZ2bCodTcBG4Ll5cB8M0s3szHAeOCtLvj+LpeSnMSzX/4IAJU1dXFujYhI9+lUCJhZH+Aa4Jmo4u+a2XozWwd8DPgagLtvBJ4ENgEvAne4e6/9P+zU0TmNy4/8eXf8GiIi0o1SOvNhd68Azjqt7OYW6t8H3NeZ7+xJk0f0Z8O+chYv38jogX24ckIuyUmxhjZERM5M5h5zbLbXyM/P94KCgrh9f95dzzcpS09J4vFFMxncP+OUy0pFRHoDM1vj7vltqdupM4EwmHP+UF7ceOCUsqraem766Z8B+NdPns/h41VMHNqfeRcOi7UJEZFeS2cCbbC3tIJ7f7eJVZtafg7x7iXzeqhFIiLNa8+ZgEKgndbuPcqNP/lTs+uvnJDLA5+dRt90nWSJSHy0JwQ0i2g7TRmVze4l89i9ZB4XjBjQZP1r20o4f/FKXt3Wsze5iYh0hM4EOqm2rp539x7l0w++3mTdZy4Zzb/fdEEcWiUiYabuoDg5WVPHxG++2KR813fmYqZLS0WkZ6g7KE4yUpPZvWQev4mahA7g/7YfjlOLRERaphDoBvl5A3nzG7Ma39+y7C1q6+rj2CIRkdgUAt1kSP8MXvralY3vx93zAnl3PU/pCT26UkR6D4VAN5owpB8XjTz1CqJp315Fbx+HEZHwUAh0s+e+cjlzLxh6StmYu1fw5x0aJxCR+FMI9ICffnY6a//lGm6cMryx7DM/f5MTVbVxbJWIiEKgx2T3SeOH86eeUnb+4pVxao2ISIRCoIftXjKPJ2/Tc4xFpHdQCMTBjDGnPsd4b2lFHFsjImGmEIiTp7/0kcblK777Cif1GEsRiQOFQJxMPzvnlPdPvL03Ti0RkTBTCMTR7iXzuGbSEAAWL99Ida3uKhaRnqUQiLOf3/LhHE9fffzdOLZERMKo0yFgZrvNbL2ZrTWzgqBsoJmtMrPtwc+coNzMbKmZFZrZOjOb1tnvTwTLPhcJghc2HNDdxCLSo7rqTOBj7j4laurSu4DV7j4eWB28B7geGB+8FgEPdNH3n9Gunjikcfn1nR/EsSUiEjbd1R10A/BIsPwIcGNU+aMe8QaQbWZ6Ojvw35+/GIjcSSwi0lO6IgQceMnM1pjZoqBsiLsXAwQ/BwflI4Doy2CKgrLQu2pCbuNy3l3P8x8vbuGvf6FAEJHu1RVPQ7/M3feb2WBglZltaaFurMdrNekED8JkEcDo0aO7oIm9n5nRJy2ZiurI/QIP/GEHAI+9tYeJQ/sxdXROSx8XEemQTp8JuPv+4Och4FlgBnCwoZsn+HkoqF4EjIr6+Ehgf4xtPuTu+e6en5ube/rqhLXp3jlNyu5+Zj03/fTPHD5eRV29Bo1FpGt1KgTMrK+Z9WtYBq4FNgDLgYVBtYXAc8HycuCW4CqhmUBZQ7eRROxeMo8LT3sGAUSmlzjnGyvi0CIRSWSd7Q4aAjwbPEQ9Bfi1u79oZm8DT5rZrcAe4NNB/RXAXKAQqAA+38nvT0jP3XEZ1XX1lFXUMOPfV5+yrr7eSUrSQ+tFpGtYb78uPT8/3wsKCuLdjLjadfgEH/veH04pe+ueWcy4bzW3fXQs5ZU13Dl7Ak8V7OUrV49vrLP/aCWbi8uZdd4QRCQ8zGxN1CX7LddVCJwZfvJKIf+5cmu7PjMgM5WyyhoeXzSTjNRkauvqmTo6h2SdSYgkNIVAgsq76/ku3d4r/3AVYwb17dJtikj8tScENHfQGWT3knls+Nfr2HTvdQDcOXt8K59o2ce+9wfm/PC1rmiaiJyhuuI+AelBWemRQ7Z7yTwA7pw9oUmd0hPV1NTV89hbe5h3wTCuub/5/9FvOXCMoxXVZPdJ654Gi0ivpu6gEFq6ejs/WLXtlLIJQ7LYdvA4D3x2GledO5jDx6sYlJXOlgPlnDM4i/4ZqQD8Yeshzj6rL8MGZJCekkRwZZiI9CIaE5BWuTv3r9rG0t8Xdsn2/vG6c3ltWwnTzs7hmXeKGJnThzXvH2H3knkcLD9JSpJxVlZ6zHZs3F/O5BFN740QkY5RCEibuDtj7o7PDWjpKUk8vPBi/rD1EL/44y5++FdTuHGqppES6QoKAWmz4rJK3tt7lIrqOp4qKIrrVNZ/N2s8X7+m6RiHiLSPQkC6RU1dPSs3HmDi0H5sKj5GdmZq46R3X3tiLR+cqO70dzQMeItIx7UnBHR1kLRZanISH79wOADjBvc7Zd2ab17Tpm18Z8Vmnn5nH4ePV8Vcv2l/OZOG9wegorqW6tp6svukseeDCnL7pZOZltyJPRCR0+lMQOKqsrqOd/cc4TMtPDvhL6aN4Jl39p1S9vnL8vjcR/IoOVbF/rKTHCir5FPTR1Gwu5Tyk7Usfm4Dl40bRHpqMv/y8Uk8+24RF43M5pKxZ8X8jqv+8xVumDKCH63ezstf/yjjBmcBkXGT41W19AuujhI5E6g7SM448RykjuXCkQMYl5vFM+9GwueZL3+E/UcruXJCLjsOHedvH3uXpQumMiqnD7n9ml71JBJPCgE5I33iv/7I+n1l8W5Gh/zXgqmcO7QfHxyv5tJzYp9tiPQUhYCc0bYeOAbAbb8q4JkvX0aSwc7DJ0hPSeK+5zfzqekjeXt3KY+99eGTShfMGMXB8ipq6ur5v+2H49X0RpeMGcgTt10a72ZISCkERNrp2Mmaxik51u8rY9Kw/uwpreBnr+5k3oXDuGXZW1w27iz+VNj2S2gf/cIMrpzQ+Sfjrdp0kCsnDCI9RYPi0jYKAZEedKKqltv/Z03MM5At355D6YlqhmdnsnF/pKsryYyK6lqeeWcf08/OYUj/DCaPGEBNXT2DstJ5bVsJX3y0gOEDMhjYN4139hxt3N6ogZlkpCTz0teuBIg5bcefCg/z+o4PuPq8wZw/vH+r4VFdW09tfT0Fu4+w/dBxbr18TGd+HU188ZECPnpuLjfPPLvZOjV19Rjwxs5SLht3VrunIzlZU0dKkrHz8AnGDOpLavKpc2NetuT37Dta2ewlyO5O+claBmTGvgDA3Vn+3n6GZ2fy9JoivvMXF7S7jTV19Ww9cIy6eueiUdmN5XX1jruTktx183kqBETixN359Vt7uOfZDT32nTdNHcGz7+7jR/On8NXH1zZZP+/CYTy/rpiUJOO7n7qwSVfa6UYNzORgWRVJSXCypp6JQ/tR787nPjKGbzy7nts+OpaMlGSeX19M4aHj/NeCqY0D5cMHZPBPT69jZ8kJrp88lHVFZew7WglAcpI1Pid70rD+bCouB+DeG87nX57beEobfjR/ClNH5VBWWcMnfvzHFvf/6omD+f2WQy3WiWVAZioX5+Xw8uamnx2Zk8nHLxzOg6/uaNO25l4wlH1HKnmvqOmY1syxA3ljZ+kpZXfOHk/x0ZMM7p/O9oPHeXHjAWafN5iXNx/iJ5+ZxsV5OWT3SSMtpWPBoBAQiaPaunrG3fNCvJshCaCjN0/qeQIicZSSnKQ7n+WMoTuGRbrJru/M5d+Cq5l++PI2li6YSm2d88Tbe7loVDaFh45xcd5Axgzqyz//dgPXnj+Uuvp6Vm06xKyJgymtqKb46En6pCWzYkMxN888m/ue38wXrxjL1gPlfOKi4azfV8bk4QOYfnak62RgVhr/8OR7fGziYI5UVDNxaD/W7i1jQGYqNXX1fGr6SFasL2b2eUO44ruv8NDN03ni7b3cM+88Fi/fyMMLL2bLgXJy+qTxzd9u4NJzziI9NZm/zB/JQ6/upE96Cis3HuBXt85g8uKVXHXuYD57yWgmDe/P4WPVDM/O4M1dpYwe2Ic/FR6mqraesYP6cqK6jkFZadTWOceqath1uIKlq7fzzY9PovhoJROH9ecfnnqPFX93Bb9+630uH5fL5BH9eWH9Ac7KSqP0RDVr9x7lr2eezapNB5l7wVBe2nSQL390HPe/vI1f/nk3+WfncMX4XF7adIBv3ziZUTl9+M6KzfzNlWPZcqCcy8flclbfNN7de5TUZGPogAwu/49XuHzcIO6cPZ5xg7M4WlHDo6+/z1dnjee7K7cwuF8G108eys7Dx8lMTWHBz9/gq7PGs2rTQf553nl85hdv8ru/vZyiIxWUV9aycuMB5l4wjItGDWD2D17jqdsv5QcvbWP2pCEcOVHNj18p5JEvzGDhsrf4myvGcLKmnl+98T6PfmEGH5yo4tKxg3h+fTE//v12vv+XF/XIv9MOdweZ2SjgUWAoUA885O4/MrNvAX8DlARVv+HuK4LP3A3cCtQBf+fuK1v7HnUHiYi0T0/NHVQL/L27v2Nm/YA1ZrYqWHe/u3/vtEZNAuYD5wPDgZfNbIK713WiDSIi0gkdHhNw92J3fydYPgZsBlqaEP4G4HF3r3L3XUAhMKOj3y8iIp3XJQPDZpYHTAUaZgH7ipmtM7NlZpYTlI0Aoq9LK6Ll0BARkW7W6RAwsyzgaeBOdy8HHgDOAaYAxcD3G6rG+HjMAQkzW2RmBWZWUFJSEquKiIh0gU6FgJmlEgmA/3X3ZwDc/aC717l7PfBzPuzyKQJGRX18JLA/1nbd/SF3z3f3/Nzczt92LyIisXU4BCxyz/TDwGZ3/0FU+bCoajcBDbdOLgfmm1m6mY0BxgNvdfT7RUSk8zpzddBlwM3AejNruFf9G8ACM5tCpKtnN3AbgLtvNLMngU1Eriy6Q1cGiYjEV4dDwN3/SOx+/mafDOLu9wH3dfQ7RUSka/X6uYPMrAR4v4MfHwTEf3L5nqV9Tnxh21/QPrfX2e7epgHVXh8CnWFmBW29ay5RaJ8TX9j2F7TP3UkTyImIhJhCQEQkxBI9BB6KdwPiQPuc+MK2v6B97jYJPSYgIiItS/QzARERaUFChoCZzTGzrWZWaGZ3xbs9nWFmo8zsFTPbbGYbzeyrQflAM1tlZtuDnzlBuZnZ0mDf15nZtKhtLQzqbzezhfHap7Yws2Qze9fMfhe8H2NmbwZtf8LM0oLy9OB9YbA+L2obdwflW83suvjsSduZWbaZ/cbMtgTH+9JEPs5m9rXg3/QGM3vMzDIS8TgHE2keMrMNUWVddlzNbLqZrQ8+szSYzaHt3D2hXkAysAMYC6QB7wGT4t2uTuzPMGBasNwP2AZMAr4L3BWU3wX8R7A8F3iByI18M4E3g/KBwM7gZ06wnBPv/Wthv78O/Br4XfD+SWB+sPwg8KVg+cvAg8HyfOCJYHlScOzTgTHBv4nkeO9XK/v8CPDFYDkNyE7U40xkBuFdQGbU8f1cIh5n4EpgGrAhqqzLjiuR6XcuDT7zAnB9u9oX719QN/zCLwVWRr2/G7g73u3qwv17DrgG2AoMC8qGAVuD5Z8BC6Lqbw3WLwB+FlVYZDGVAAAC10lEQVR+Sr3e9CIyueBq4Grgd8E/7sNAyunHGFgJXBospwT17PTjHl2vN76A/sH/FO208oQ8znw4tfzA4Lj9DrguUY8zkHdaCHTJcQ3WbYkqP6VeW16J2B2UsM8tsFOf2zDE3Ysh8oAfYHBQrbn9P5N+Lz8E/onIY0sBzgKOuntt8D667Y37FawvC+qfSfsLkTPXEuC/g26wX5hZXxL0OLv7PuB7wB4iU86XAWtI/OPcoKuO64hg+fTyNkvEEGjzcwvOJNb0uQ3NVo1R5i2U9ypm9nHgkLuviS6OUdVbWXdG7G+UFCJdBg+4+1TgBJFuguac0fsd9IHfQKQLZzjQF7g+RtVEO86tae9+dnr/EzEE2vzcgjOFxXhuA3DQgmm7g5+HgvLm9v9M+b1cBnzSzHYDjxPpEvohkG1mDRMeRre9cb+C9QOAUs6c/W1QBBS5e8PT+X5DJBQS9TjPBna5e4m71wDPAB8h8Y9zg646rkXB8unlbZaIIfA2MD64yiCNyCDS8ji3qcOCkf4mz20gsk8NVwgsJDJW0FB+S3CVwUygLDjdXAlca2Y5wV9h1wZlvYq73+3uI909j8ix+727fxZ4BfhUUO30/W34PXwqqO+cYc+vcPcDwF4zOzcomkVk2vWEPM5EuoFmmlmf4N94w/4m9HGO0iXHNVh3zMxmBr/HW6K21TbxHjDppkGYuUSuotkB3BPv9nRyXy4ncnq3DlgbvOYS6Q9dDWwPfg4M6hvwk2Df1wP5Udv6AlAYvD4f731rw75fxYdXB40l8h93IfAUkB6UZwTvC4P1Y6M+f0/we9hKO6+YiNP+TgEKgmP9WyJXgSTscQb+FdhC5MFTvyJyhU/CHWfgMSLjHjVE/nK/tSuPK5Af/A53AD/mtIsLWnvpjmERkRBLxO4gERFpI4WAiEiIKQREREJMISAiEmIKARGREFMIiIiEmEJARCTEFAIiIiH2/wFJ8nwh1j16pwAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot(losses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "answer = guide()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Apply to a real protein\n",
    "\n",
    "I will grab two frames from an MD trajectory and use this method to align them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mdtraj\n",
    "\n",
    "# Input structures\n",
    "x1 = mdtraj.load(\"ref.pdb\")\n",
    "x2 = mdtraj.load(\"target.pdb\")\n",
    "\n",
    "# Placeholder trajectories for results\n",
    "mover = x1[0]\n",
    "m = x1[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1.6271, 0.6045, 2.5373],\n",
      "        [1.5632, 0.7143, 2.4642],\n",
      "        [1.6662, 0.7874, 2.3812],\n",
      "        ...,\n",
      "        [4.8180, 2.6577, 2.7567],\n",
      "        [4.9805, 2.5987, 2.9811],\n",
      "        [4.9957, 2.7432, 2.8764]])\n"
     ]
    }
   ],
   "source": [
    "X1 = torch.tensor(x1.xyz[0], dtype=torch.float32)\n",
    "X2 = torch.tensor(x2.xyz[0], dtype=torch.float32)\n",
    "print(X1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "metadata": {},
   "outputs": [],
   "source": [
    "N=1231\n",
    "def model2():\n",
    "    N=1231\n",
    "    # Prior over mean M, with N=number of atoms\n",
    "    M = pyro.sample(\"M\", dist.StudentT(1,0,3).expand_by([N,3]).to_event(2))\n",
    "    # Prior over variances U\n",
    "    U = pyro.sample(\"U\", dist.HalfNormal(1).expand_by([N]).to_event(1))\n",
    "    U = U.reshape(N,1).repeat(1,3).view(-1)\n",
    "    # Prior over translation T\n",
    "    T = pyro.sample(\"T\", dist.Normal(0,1).expand_by([3]).to_event(1))\n",
    "    # Prior over rotation R\n",
    "    u = pyro.sample(\"u\",dist.Uniform(0,1).expand_by([3]).to_event(1))\n",
    "    # Transformation: turn u via a unit quaternion into a rotation R\n",
    "    R = u_to_quat_to_R(u)\n",
    "    # Transformation: rotate and translate M for X2\n",
    "    M_RT = M @ R + T\n",
    "    \n",
    "    df = pyro.sample(\"df\", dist.LogNormal(10, 2))\n",
    "    \n",
    "    # Likelihood\n",
    "    with pyro.plate(\"plate_students\", N*3,dim= -1):\n",
    "        pyro.sample(\"X1\", dist.StudentT(df, M.view(-1), U), obs=X1.view(-1))\n",
    "        pyro.sample(\"X2\", dist.StudentT(df, M_RT.view(-1), U), obs=X2.view(-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyro.contrib.autoguide import AutoDelta\n",
    "from pyro.infer import Trace_ELBO\n",
    "from pyro.optim import Adam, AdagradRMSProp\n",
    "\n",
    "guide2 = AutoDelta(model2)\n",
    "# optim2 = Adam({'lr': 0.01})\n",
    "# optim2 = AdagradRMSProp({'lr': 1.0, 'momentum':0.1, 'eps':1e-16})\n",
    "optim2 = AdagradRMSProp({})\n",
    "svi2 = pyro.infer.SVI(model2, guide2, optim2, loss=Trace_ELBO())\n",
    "\n",
    "\n",
    "losses=[]\n",
    "for i in range(2000):\n",
    "    losses.append(svi2.step())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x6198846d8>]"
      ]
     },
     "execution_count": 110,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYcAAAD8CAYAAACcjGjIAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvOIA7rQAAIABJREFUeJzt3XlwnHed5/H3V/dlHZZkW4fP2El8xLmc2EDCDAnrOCGQzHJMsizxhlCppQIFtexAGHbIDEcVMDswZBfCZCeuOAwQ7ophEowJgQkQx5FNDp9YvmXLtmydtiTr+u4f/ZNpq1tnJLWc5/Oq6uqnf8+vn/7201J/+jl+3ebuiIiIxEtLdQEiIjL1KBxERCSBwkFERBIoHEREJIHCQUREEigcREQkgcJBREQSKBxERCSBwkFERBJkpLqAsSorK/N58+alugwRkYvG1q1bT7l7+Uj6XrThMG/ePGpqalJdhojIRcPMDo20r3YriYhIAoWDiIgkUDiIiEgChYOIiCRQOIiISAKFg4iIJFA4iIhIghGFg5kdNLPXzOxlM6sJbdPNbJOZ7Q3XJaHdzOxhM6s1s1fN7Jq45awN/fea2dq49mvD8mvDfW28n2i/rYca2VXfOlGLFxF5QxjNlsPb3P0qd18Rbj8IPOvui4Bnw22AW4FF4XI/8AjEwgR4CFgJXA881B8ooc/9cfdbM+ZnNIx3P/ICt379+YlavIjIG8Lr2a10B7A+TK8H7oxrf8JjNgPFZlYB3AJscvdGd28CNgFrwrxCd3/B3R14Im5ZIiKSAiMNBwd+aWZbzez+0DbT3esBwvWM0F4FHIm7b11oG6q9Lkm7iIikyEi/W+kt7n7MzGYAm8xs9xB9kx0v8DG0Jy44Fkz3A8yZM2foikVEZMxGtOXg7sfC9Ungp8SOGZwIu4QI1ydD9zpgdtzdq4Fjw7RXJ2lPVsej7r7C3VeUl4/oiwVFRGQMhg0HM8s3s2n908BqYDuwAeg/42gt8FSY3gDcE85aWgW0hN1OG4HVZlYSDkSvBjaGeW1mtiqcpXRP3LJERCQFRrJbaSbw03B2aQbwXXf/hZm9BPzAzO4DDgPvDf2fBm4DaoF24F4Ad280s88DL4V+n3P3xjD9YeBxIBd4JlxERCRFhg0Hd98PXJmk/TRwc5J2Bx4YZFnrgHVJ2muAZSOoV0REJoFGSIuISAKFg4iIJFA4iIhIAoWDiIgkUDiIiEgChYOIiCRQOIiISAKFg4iIJFA4iIhIAoWDiIgkUDiIiEgChYOIiCRQOIiISAKFg4iIJFA4iIhIAoWDiIgkUDiIiEgChYOIiCRQOIiISAKFg4iIJFA4iIhIAoWDiIgkUDiIiEgChYOIiCRQOIiISAKFg4iIJFA4iIhIAoWDiIgkUDiIiEgChYOIiCRQOIiISAKFg4iIJFA4iIhIAoWDiIgkUDiIiEgChYOIiCQYcTiYWbqZ/dHMfh5uzzezF81sr5l938yyQnt2uF0b5s+LW8anQ/seM7slrn1NaKs1swfH7+mJiMhYjGbL4WPArrjbXwa+5u6LgCbgvtB+H9Dk7guBr4V+mNkS4C5gKbAG+GYInHTgG8CtwBLg7tBXRERSZEThYGbVwDuAfw23DbgJ+FHosh64M0zfEW4T5t8c+t8BPOnu59z9AFALXB8ute6+3927gCdDXxERSZGRbjn8M/BJoC/cLgWa3b0n3K4DqsJ0FXAEIMxvCf3Ptw+4z2DtIiKSIsOGg5ndDpx0963xzUm6+jDzRtuerJb7zazGzGoaGhqGqFpERF6PkWw5vAV4l5kdJLbL5yZiWxLFZpYR+lQDx8J0HTAbIMwvAhrj2wfcZ7D2BO7+qLuvcPcV5eXlIyhdRETGYthwcPdPu3u1u88jdkD51+7+fuA54D2h21rgqTC9IdwmzP+1u3tovyuczTQfWARsAV4CFoWzn7LCY2wYl2cnIiJjkjF8l0F9CnjSzL4A/BF4LLQ/BnzbzGqJbTHcBeDuO8zsB8BOoAd4wN17AczsI8BGIB1Y5+47XkddIiLyOo0qHNz9N8BvwvR+YmcaDezTCbx3kPt/EfhikvangadHU4uIiEwcjZAWEZEECgcREUmgcBARkQQKBxERSaBwEBGRBAoHERFJoHAQEZEEkQ2H2KBtERFJJrLhcK6nb/hOIiIRFd1w6FY4iIgMJrrh0NOb6hJERKasyIZDp7YcREQGFdlwaO7oSnUJIiJTVmTD4VhzZ6pLEBGZsiIXDgXZsW8pP97SkeJKRESmrsiFw/T8LADqW7TlICIymMiFg1nsuq5JWw4iIoOJXDj0O9R4NtUliIhMWZELh/5vzTh0uj21hYiITGGRC4d+bZ09tHV2p7oMEZEpKZLhkJkeO/Cg01lFRJKLZDhUFucCcKxZB6VFRJKJZDhUhXA4qnAQEUkqcuHgOOXTsslMN4WDiMggIhcOAOlmzCrK0W4lEZFBRDIcILZr6agGwomIJBXZcKgsytVXaIiIDCKy4TCrKIcTrZ309em3pEVEBopcOPSPkK4ozqWnzzl15lxqCxIRmYIiFw4AGFQU5gD6dlYRkWSiGQ7EdiuBwkFEJJnIhkNFCAf96I+ISKLIhsP0/CyyMtKob9WWg4jIQJELh/4D0mZGRVEO9fryPRGRBJELBwAj9q2sswpzqNduJRGRBJEMh36VxRoIJyKSzLDhYGY5ZrbFzF4xsx1m9g+hfb6ZvWhme83s+2aWFdqzw+3aMH9e3LI+Hdr3mNktce1rQlutmT04/k8zuYqiHI63dNKrgXAiIhcYyZbDOeAmd78SuApYY2argC8DX3P3RUATcF/ofx/Q5O4Lga+FfpjZEuAuYCmwBvimmaWbWTrwDeBWYAlwd+g74TQQTkQkuWHDwWPOhJuZ4eLATcCPQvt64M4wfUe4TZh/s5lZaH/S3c+5+wGgFrg+XGrdfb+7dwFPhr4TrjKczqpvZxURudCIjjmET/gvAyeBTcA+oNnde0KXOqAqTFcBRwDC/BagNL59wH0Ga58wFjseTUVR7Ed/dNxBRORCIwoHd+9196uAamKf9Bcn6xaubZB5o21PYGb3m1mNmdU0NDQMX/gwqvRzoSIiSY3qbCV3bwZ+A6wCis0sI8yqBo6F6TpgNkCYXwQ0xrcPuM9g7cke/1F3X+HuK8rLy0dTelKFuRnkZaVzTGMdREQuMJKzlcrNrDhM5wJvB3YBzwHvCd3WAk+F6Q3hNmH+r93dQ/td4Wym+cAiYAvwErAonP2UReyg9YbxeHLD6R8Ipy0HEZELZQzfhQpgfTirKA34gbv/3Mx2Ak+a2ReAPwKPhf6PAd82s1piWwx3Abj7DjP7AbAT6AEecPdeADP7CLARSAfWufuOcXuGA7hfuMeqqiRPA+FERAYYNhzc/VXg6iTt+4kdfxjY3gm8d5BlfRH4YpL2p4GnR1DvuIg/yFFVnMPOYy2T9dAiIheFSI+QhtjPhZ4600Vnd2+qSxERmTIiHw5VJTqdVURkoMiHQ2U4nfVok447iIj0i1w4DBxAobEOIiKJIhcOA80szMEMjiocRETOi2Q4WNzpSlkZacyclqNwEBGJE8lwGKiyWAPhRETiKRyIHZRWOIiI/JnCgdjprMeaO+nTj/6IiAARDAdP8v5fVZxLV28fp87qR39ERCCC4QBgA74lvLKo/3RWDYQTEYGIhsNA/aOk65raU1yJiMjUoHAA5pbmAXDw1NkUVyIiMjUoHIC8rAwqinLYr3AQEQEiGA6e/BdImVeazwGFg4gIEMFwgAtHSPebX56v3UoiIkEkwyGZBWX5NLV303S2K9WliIiknMIhmF+WD8CB09p6EBFROAT94bC/QeEgIhK5cEg2Qhpgbmk+OZlp7KpvndyCRESmoMiFAyQ/IJ2eZiyuKGT70ZbJL0hEZIqJZDgMZmllITuPteoL+EQk8hQOcZZVFtF2rocj+hoNEYk4hUOcpZVFAGw/quMOIhJtkQuHoXYYXTqrgKyMNLYdbpq0ekREpqLIhUNMkiPSQHZGOtfMKWbz/tOTXI+IyNQS0XAY3JsWlLGzvpWW9u5UlyIikjIKhwFWLZiOO2w52JjqUkREUkbhMMBVc4rJzkjjhX3atSQi0RW5cBhshHS/7Ix0rps3nT/sOzU5BYmITEGRCwdIPkI63g2Lyth9vI2TrfpNaRGJpkiGw3BuXFQGwPN7tfUgItGkcEhi8axCSvOz+F2twkFEoknhkERamnHDojKe33tK37MkIpEUwXAY2Zv9DQvLOHXmHLuPt01wPSIiU08Ew2Gw8dEXunFROQC/q22Y2GJERKagYcPBzGab2XNmtsvMdpjZx0L7dDPbZGZ7w3VJaDcze9jMas3sVTO7Jm5Za0P/vWa2Nq79WjN7LdznYbPhzieaeLOKcrh0ZoEOSotIJI1ky6EH+IS7LwZWAQ+Y2RLgQeBZd18EPBtuA9wKLAqX+4FHIBYmwEPASuB64KH+QAl97o+735rX/9RevxsWlvPigUY6u3tTXYqIyKQaNhzcvd7dt4XpNmAXUAXcAawP3dYDd4bpO4AnPGYzUGxmFcAtwCZ3b3T3JmATsCbMK3T3F9zdgSfilpVSN15aRldPHy/pqzREJGJGdczBzOYBVwMvAjPdvR5iAQLMCN2qgCNxd6sLbUO11yVpnxDDjZCOt3L+dLLS03RKq4hEzojDwcwKgB8DH3f3oX4NJ9nxAh9De7Ia7jezGjOraWgY+4HikR7RyMvKYHFlIS8fbh7zY4mIXIxGFA5mlkksGL7j7j8JzSfCLiHC9cnQXgfMjrt7NXBsmPbqJO0J3P1Rd1/h7ivKy8tHUvrrdmV1EduPttCr8Q4iEiEjOVvJgMeAXe7+1bhZG4D+M47WAk/Ftd8TzlpaBbSE3U4bgdVmVhIORK8GNoZ5bWa2KjzWPXHLSrnl1cWc7eplf8OZVJciIjJpMkbQ5y3AB4DXzOzl0Pa3wJeAH5jZfcBh4L1h3tPAbUAt0A7cC+DujWb2eeCl0O9z7t5/pPfDwONALvBMuEwJy6tjvyv9al0Li2ZOS3E1IiKTY9hwcPffMfi4sZuT9HfggUGWtQ5Yl6S9Blg2XC3jYbQ7hy4pLyAvK51X65p597XVw99BROQNIKIjpEc+xi49zVhaWciOY0MdgxcReWOJZDiM1pKKQnbVt+pL+EQkMhQOI7C4opCzXb0caWpPdSkiIpNC4TACSyoLAdipXUsiEhGRCwcfzRDp4NKZ00hPM3bWKxxEJBoiFw4w8hHS/XIy07mkPF9bDiISGZEMh7FYUlGoLQcRiQyFwwgtriikvqWTprNdqS5FRGTCKRxGqP+g9C5tPYhIBEQuHMY6UmFxRThjSeEgIhEQuXCAkf2G9EBlBdnMKszhtaMt416PiMhUE8lwGKtr55ZQc7Ap1WWIiEw4hcMoXDO3hKPNHZxs7Ux1KSIiE0rhMApXzY59ffcrddq1JCJvbJELhzEMkD5vaWURGWlGzcHG4TuLiFzEIhcOADbaIdJBTmY6b15YxtPb68f0NRwiIheLSIbD6/HO5RUcaezQriUReUNTOIzS6qWzyEpPY8PLx1JdiojIhFE4jFJRbiY3L57Bj7YeoaWjO9XliIhMCIXDGDzwtoW0dvbw8LN7U12KiMiEiFw4jMeB5GVVRXxg1Vwe+90BflhzZByqEhGZWiIXDuPlM+9YzJsvKeVvfvQqn/vZTnr1+9Ii8gaicBijnMx01n/wev7rqjms+/0BLvnbp/lPX/0tz+0+qdNcReSil5HqAi5mmelpfOHOK1hcUchnfrqdvSfPcO/jL52ff9d1s/nQjfNZOGNaCqsUERk9hcM4eP/Kudx93Rye2X6cB7677Xz7ky8d4cmXYsckyqdls3L+dLp7+/jE6su4dKYCQ0SmrsiFw0Tt8ElLM96xvIJ3LH8HHV29/J9f7+Wbv9l3fn5D2zl+/mo9ABt3nDjfvriikA/dMJ9r55aQnmZUl+SOeQS3iMh4iVw4AEz0e29uVjqfXHM5f3PLZTS1d9Pa0c0H1r3IkcaOhL676lv5xA9fSbqciqIcrps3nRcPnGbN0lm8UtfCR29ayJLKQgqyMyjIzqCrt4+s9DROnemifFr2xD4xEYmMSIbDZDEzpudnMT0/i+c/eRO9fc7hxnYOnT7L6TNdfHvzIV4+0jzo/etbOtnwSmwk9voXDgFw3/qaIR9zzdJZvGN5BfNK8ynKzeSXO49TmJPJmxeWUj4tm+yM9PF7giLyhqVwmETpacb8snzml+UD8O5rq8/Pa27voqHtHHVNHZzt6qG8IJstBxp5dvfJIQNkoF/sOM4vdhwftp/Zhd9QW1WcS0F2Bu+8soKGtnMcb+3kiqoi2rt6KSvI5tSZc+eDbuWCUkrzs3CHtDQUOCJvQAqHKaI4L4vivCwWxR2oXrmglI/evAiAjq5e6praqSjO5UDDWYrzMmnp6Oaff/UnfrXr5Kgfb+DZtkebY7u89vyy7Xxb/LGR1+vmy2dQ19TBygXTuXZuCe5w2axptHf1sv1oC/UtndzzprlMz88iKz2NtDQddxFJpeiFw0U6BCE3K/18cFxRHfvRodnAv669bsj7Pbf7JAtnFNB4tosdx1rJyUyjOC+Ts+d6+c2eBjbvP019SwcTPYbv2d2xANtzoo0nwi6ygb71231J2/tNy8lgdkkel1dM4/0r59B4tpuC7AxmFmYztzSfmoON5Gal4w7Lq4to6eimKDcTM6Ots5vjLZ3MmJZDYW7GsAf9e3r7yEjXMCCJruiFA2BE51Pp2y6fAcDs6XlcObv4gnnvvLJyyPu2dXbzvzfuYdWCUrYeamJe2CX22ae2s6/h7ITVPHg9Peysb2VnfSs/2XZ03JabmW4U5WaxcEY+pQXZ/Hs4q6x8Wjal+VnkZKazs76Vrp4+3n1NNVkZaVw/v4S5pflsOdDI4opCDje2s+1QE9uPtnDfDfNZvXQWv689xfyyfJZVFfH83gYWziigND+bY80ddHT30tbZw6UzCyjOy7qgnq6ePo41d1CSn0VRbua4PU+R0bCLdTTvihUrvKZm6IOzyVzx0Ebeu2I2n33nkgmoSiD25vajrXUsry5iQXk+fzpxhvml+Zzt6qGsIJu2zm4y0tLA4JUjzXzpmd3srG+lrCCLq+eUsGJuCSV5WXx78yFeOzqy3814y8JSfl97eoKfWWqU5GXS1N5NRprxl5fNoLokl22Hm+jq6eMvLivn8lnTeHbXyfOnShfmZPD2JTMpys2krCCbquJcVi6YTtPZbuaW5tHV08cvdhxnflk+K+dP55Hf7mPxrELeemk5x5o7mD0974LH33qokatml5BmF/5QVk9vH1sONvLmS8omdX3I2JnZVndfMaK+Cgd5I+rrc8719GEWOxFg26Em9jWcpbWzm6tmF7P1UBNbDzVx+/IKGs92Mbc0n7mleWzef5p/+e1+jjZ3UFGUQ3qaUdeUeAqyJFpaGQuYR+LG91w3r4TFFYU88cIh7riqkqdePsaHbpjP3SvnMC07g3978TCd3b109fTx690nmVWUwxfvXMbvak/xL7/dz/HWTt51ZSV/d/sSevr6+M2eBl7Yd5q1b55HWUEW+xvO0tzRxZ1XVWFmnGzt5NW6Frp7+3ilroWPv30RG3ccZ82yWWzaeYJbl1UkhNxIvVrXzILyAgqyL94dLgqHISgcZLTcnfqWTiqLc3F3Orv7yEy3C45J9PTG3tya27t566XldHb38uNtdcwozOG/XD+HH2+t48UDsd8eL8zN4HtbDtPZ3cc7r6zkpsvL+e6Lh3ntaAud3X0XPPby6iL2HG/jXM+F7TIxMtONRTOmsbO+ddA+l8+axl9dXcXB0+18b8thygqy+OAN89ld38by6iLaOnt4bs9JPnTjAr6z+RBZGWlkpBnvuqqS6pI8ZhXmUJKfxZ7jreyqb6Olo5vczHQWlOdz2axppJkxszAH99gHnKb2Lo41d3K8pZPl1UUJW3ajoXAYwrKHNvI+hYNMce4+6k+3+xrO0NvnzJyWQ0a6kR8+4da3dFBRlEtPbx8OsUGZj23hxkvLuP2KSu59fAvZGennd1GdORc7lbq0IIsdR1v5p01/4rp5Jbx98UxaO7t5+Ugzv689TWFOBq2dPecf3wxuWFjG83tPjeeqkAFuX17Bw3ddPaYz+hQOQ1j20Eb++rrZ/N3tCgeRiXCupzfp2Jez53poau8iJzOd6XlZpKUZ//5qPZdXTKOjq5efvXKMq+eU8BeXlnO0uZ1/+NlOblk6i9L8LKpL8jja3M6h0+2cPttFcV4mPb3O9qMt3H39HGZPz2P38Vae+MMhSguy6O51frXrBDcuKuO1oy3kZcbO9ltSWcgPa45QkpfFodPt3HrFLJZWFvJvmw9zuLGdN19Syh/2DX/sKj3NUvo1/bs+t4bcrNGPLxrXcDCzdcDtwEl3XxbapgPfB+YBB4H3uXuTxT7qfB24DWgH/pu7bwv3WQv8r7DYL7j7+tB+LfA4kAs8DXzMR5BYCgcRmShdPX2kpxnp4dP5ybZOdhxt5S0Ly8jKuPAUZ3envauXvKz0C7b2zpzrIc3g0Ol2cjPTKc7LpKO7l4qiXAA6u3vZVd/KidZOag42cesVsyjMyWThjALazvWw7VATC8oK2Ha4id/XnqLP4XhrB195z5VUFeeO6XmNdzi8FTgDPBEXDl8BGt39S2b2IFDi7p8ys9uAjxILh5XA1919ZQiTGmAFsZEGW4FrQ6BsAT4GbCYWDg+7+zPDFa5wEBEZndGEw7CjfNz9P4DGAc13AOvD9Hrgzrj2JzxmM1BsZhXALcAmd2909yZgE7AmzCt09xfC1sITccsSEZEUGesQ0JnuXg8QrmeE9iog/keV60LbUO11SdqTMrP7zazGzGoaGhrGVPjFeoxFRGQyjff3AyQ7fO5jaE/K3R919xXuvqK8vHyMJSZ/UBER+bOxhsOJsEuIcN3/zW91xL7yp181cGyY9uok7SIikkJjDYcNwNowvRZ4Kq79HotZBbSE3U4bgdVmVmJmJcBqYGOY12Zmq8KZTvfELUtERFJk2HHgZvY94C+BMjOrAx4CvgT8wMzuAw4D7w3dnyZ2plItsVNZ7wVw90Yz+zzwUuj3OXfvP8j9Yf58Kusz4SIiIik0bDi4+92DzLo5SV8HHhhkOeuAdUnaa4Blw9UxXnQ4WkRkeJH8wvqJ/g1pEZGLXSTDQUREhqZwEBGRBAoHERFJELlw0ABpEZHhRS4cYGy/AiUiEiWRDAcRERmawkFERBIoHEREJEHkwmHNsllcPmtaqssQEZnShv36jDear/31VakuQURkyovcloOIiAxP4SAiIgkUDiIikkDhICIiCRQOIiKSQOEgIiIJFA4iIpJA4SAiIgnML9LvsDazBuDQGO9eBpwax3LGi+oaHdU1OqprdN6Idc119/KRdLxow+H1MLMad1+R6joGUl2jo7pGR3WNTtTr0m4lERFJoHAQEZEEUQ2HR1NdwCBU1+iortFRXaMT6boiecxBRESGFtUtBxERGUKkwsHM1pjZHjOrNbMHJ/mxZ5vZc2a2y8x2mNnHQvvfm9lRM3s5XG6Lu8+nQ617zOyWCaztoJm9Fh6/JrRNN7NNZrY3XJeEdjOzh0Ndr5rZNRNU02Vx6+RlM2s1s4+nan2Z2TozO2lm2+PaRr2OzGxt6L/XzNZOUF3/aGa7w2P/1MyKQ/s8M+uIW3ffirvPteFvoDbUbhNQ16hfu/H+nx2kru/H1XTQzF4O7ZOyvoZ4b0jt35e7R+ICpAP7gAVAFvAKsGQSH78CuCZMTwP+BCwB/h74n0n6Lwk1ZgPzQ+3pE1TbQaBsQNtXgAfD9IPAl8P0bcAzgAGrgBcn6bU7DsxN1foC3gpcA2wf6zoCpgP7w3VJmC6ZgLpWAxlh+stxdc2L7zdgOVuAN4WanwFunYC6RvXaTcT/bLK6Bsz/J+Czk7m+hnhvSOnfV5S2HK4Hat19v7t3AU8Cd0zWg7t7vbtvC9NtwC6gaoi73AE86e7n3P0AUEvsOUyWO4D1YXo9cGdc+xMesxkoNrOKCa7lZmCfuw816HFC15e7/wfQmOQxR7OObgE2uXujuzcBm4A1412Xu//S3XvCzc1A9VDLCLUVuvsLHnuXeSLuuYxbXUMY7LUb9//ZoeoKn/7fB3xvqGWM9/oa4r0hpX9fUQqHKuBI3O06hn5znjBmNg+4GngxNH0kbB6u6990ZHLrdeCXZrbVzO4PbTPdvR5if7zAjBTU1e8uLvyHTfX66jfadZSKGj9I7FNmv/lm9kcz+62Z3RjaqkItk1HXaF67yV5fNwIn3H1vXNukrq8B7w0p/fuKUjgk2yc46adqmVkB8GPg4+7eCjwCXAJcBdQT26yFya33Le5+DXAr8ICZvXWIvpO6Hs0sC3gX8MPQNBXW13AGq2Wy191ngB7gO6GpHpjj7lcD/wP4rpkVTmJdo33tJvs1vZsLP4RM6vpK8t4waNdBHn9c64pSONQBs+NuVwPHJrMAM8sk9uJ/x91/AuDuJ9y91937gP/Hn3eFTFq97n4sXJ8EfhpqONG/uyhcn5zsuoJbgW3ufiLUmPL1FWe062jSagwHI28H3h92fRB225wO01uJ7c+/NNQVv+tpQuoaw2s3mesrA/jPwPfj6p209ZXsvYEU/31FKRxeAhaZ2fzwafQuYMNkPXjYn/kYsMvdvxrXHr+//q+A/rMoNgB3mVm2mc0HFhE7CDbedeWb2bT+aWIHM7eHx+8/22Et8FRcXfeEMyZWAS39m74T5IJPc6leXwOMdh1tBFabWUnYpbI6tI0rM1sDfAp4l7u3x7WXm1l6mF5AbB3tD7W1mdmq8Hd6T9xzGc+6RvvaTeb/7NuB3e5+fnfRZK2vwd4bSPXf11iPZF+MF2JH+f9E7BPAZyb5sW8gton3KvByuNwGfBt4LbRvACri7vOZUOseXufZI0PUtYDYWSCvADv61wtQCjwL7A3X00O7Ad8Idb0GrJjAdZYHnAaK4tpSsr6IBVQ90E3sE9p9Y1lHxI4B1IbLvRNUVy2xfc/9f2ffCn3fHV7jV4BtwDvjlrOC2Jv1PuD/EgbIjnNdo37txvt/NlnBkQdgAAAATUlEQVRdof1x4L8P6Dsp64vB3xtS+velEdIiIpIgSruVRERkhBQOIiKSQOEgIiIJFA4iIpJA4SAiIgkUDiIikkDhICIiCRQOIiKS4P8D88fvNf6Z1EUAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot(losses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.4753, grad_fn=<ExpandBackward>)"
      ]
     },
     "execution_count": 111,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "answer = guide2()\n",
    "answer[\"df\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [],
   "source": [
    "moved = x1.xyz[0] @ u_to_quat_to_R(answer[\"u\"]).detach().numpy() + answer[\"T\"].detach().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [],
   "source": [
    "moved.xyz = moved\n",
    "moved.save(\"moved.pdb\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [],
   "source": [
    "m.xyz = answer[\"M\"].detach().numpy()\n",
    "m.save(\"M.pdb\")"
   ]
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}