dihedral_xyz_gan_mk1-Copy1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import matplotlib\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem import AllChem, rdMolTransforms, rdmolops\n",
    "import os\n",
    "import fnmatch\n",
    "import numpy as np\n",
    "import deepchem as dc\n",
    "from scipy.sparse import csr_matrix\n",
    "from sklearn.preprocessing import OneHotEncoder\n",
    "from sklearn.mixture import GMM\n",
    "import random\n",
    "from tflearn.activations import leaky_relu\n",
    "from deepchem.utils.rdkit_util import get_xyz_from_mol\n",
    "from deepchem.feat.rdkit_grid_featurizer import convert_atom_to_voxel, compute_centroid, rotate_molecules\n",
    "from tensorflow.contrib.distributions import Normal\n",
    "np.set_printoptions(precision=4)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from deepchem.feat.graph_features import ConvMolFeaturizer\n",
    "from deepchem.feat.adjacency_fingerprints import AdjacencyFingerprint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1073,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def featurize_mols(mol_files, max_n_atoms, n_bonds=5, get_angle=True):\n",
    "    featurizer = AdjacencyFingerprint(max_n_atoms=max_n_atoms)\n",
    "    features = []\n",
    "    for mol_idx, mol_file in enumerate(mol_files):\n",
    "        if mol_idx % 32 == 0:\n",
    "            print(mol_idx)\n",
    "        try:\n",
    "            if \".pdb\" in mol_file:\n",
    "                mol = Chem.MolFromPDBFile(mol_file)\n",
    "            elif \"mol\" in mol_file:\n",
    "                mol = Chem.MolFromMol2File(mol_file)\n",
    "            else:\n",
    "                get_angle = False\n",
    "                mol = Chem.MolFromSmiles(mol_file)\n",
    "            if mol is None:\n",
    "                features.append(None)\n",
    "                continue\n",
    "            \n",
    "            n_atoms = mol.GetNumAtoms()\n",
    "            \n",
    "            if get_angle:\n",
    "                mol_xyz = get_xyz_from_mol(mol)\n",
    "                centroid = compute_centroid(mol_xyz)\n",
    "                mol_xyz -= centroid\n",
    "                temp = np.zeros((max_n_atoms, 3))\n",
    "                try:\n",
    "                    temp[:mol_xyz.shape[0]] = mol_xyz\n",
    "                except:\n",
    "                    features.append(None)\n",
    "                    continue\n",
    "                mol_xyz = temp\n",
    "                c = mol.GetConformer(0)\n",
    "\n",
    "            else:\n",
    "                mol_xyz = np.zeros((max_n_atoms,3))\n",
    "                AllChem.EmbedMultipleConfs(mol, 1)\n",
    "                c = mol.GetConformer(0)\n",
    "                for aid in range(mol.GetNumAtoms()):\n",
    "                    pos = c.GetAtomPosition(aid)\n",
    "                    mol_xyz[aid] = [pos.x, pos.y, pos.z]\n",
    "\n",
    "            torsions = []\n",
    "            torsion_tuples = []\n",
    "\n",
    "            torsion_matrix = np.zeros((n_bonds,1))\n",
    "            torsion_indices = np.zeros((n_bonds, max_n_atoms, 4)).astype(np.uint8)\n",
    "            subgraphs = np.zeros((n_bonds, max_n_atoms, max_n_atoms)).astype(np.uint8)\n",
    "            atom_index_of_bond = np.zeros((n_bonds, max_n_atoms, 2)).astype(np.uint8)\n",
    "            rot_vectors = np.zeros((n_bonds, 3))\n",
    "            \n",
    "            #mol_xyz = rotate_molecules([mol_xyz])[0]\n",
    "\n",
    "            idx = 0\n",
    "            for bond in mol.GetBonds():\n",
    "                if bond.IsInRing(): \n",
    "                    continue\n",
    "                def calc_torsions(atom_i, atom_j, bond_idx):        \n",
    "                    exist_dihed = False\n",
    "                    for neighbor_j in atom_j.GetNeighbors():\n",
    "                        if neighbor_j.GetIdx() == atom_i.GetIdx():\n",
    "                            continue\n",
    "\n",
    "                        dihed_idx = 0\n",
    "                        for neighbor_i in atom_i.GetNeighbors():\n",
    "                            if neighbor_i.GetIdx() == atom_j.GetIdx():\n",
    "                                continue\n",
    "\n",
    "                            exist_dihed=True\n",
    "                            torsion_tuple = (neighbor_i.GetIdx(), atom_i.GetIdx(), atom_j.GetIdx(), neighbor_j.GetIdx())\n",
    "\n",
    "                            if get_angle:\n",
    "                                torsion_matrix[bond_idx][dihed_idx] = rdMolTransforms.GetDihedralRad(c, *torsion_tuple)\n",
    "                            torsion_indices[bond_idx][torsion_tuple[0]][dihed_idx*4] = 1\n",
    "                            torsion_indices[bond_idx][torsion_tuple[1]][dihed_idx*4+1] = 1\n",
    "                            torsion_indices[bond_idx][torsion_tuple[2]][dihed_idx*4+2] = 1\n",
    "                            torsion_indices[bond_idx][torsion_tuple[3]][dihed_idx*4+3] = 1\n",
    "                            \n",
    "                            broken_mol = rdmolops.FragmentOnBonds(mol, [bond.GetIdx()], addDummies=False)\n",
    "                            molfrags = rdmolops.GetMolFrags(broken_mol)\n",
    "                            if bond.GetEndAtom().GetIdx() in molfrags[0]:\n",
    "                                atoms_to_rotate = molfrags[0]\n",
    "                            else:\n",
    "                                atoms_to_rotate = molfrags[1]\n",
    "\n",
    "\n",
    "                            subgraph = np.zeros((max_n_atoms, max_n_atoms)).astype(np.uint8)\n",
    "                            for atom_idx in atoms_to_rotate:\n",
    "                                subgraph[atom_idx,atom_idx] = 1\n",
    "                            subgraphs[bond_idx, :, :] = subgraph\n",
    "                            atom_index_of_bond[bond_idx, atom_j.GetIdx(), 0] = 1\n",
    "                            atom_index_of_bond[bond_idx, atom_i.GetIdx(), 1] = 1\n",
    "                            \n",
    "                            pos_i = c.GetAtomPosition(atom_i.GetIdx())\n",
    "                            pos_j = c.GetAtomPosition(atom_j.GetIdx())\n",
    "                            pt_i = np.array([pos_i.x, pos_i.y, pos_i.z])\n",
    "                            pt_j = np.array([pos_j.x, pos_j.y, pos_j.z])\n",
    "                            vec = pt_j - pt_i\n",
    "                            vec = vec / np.linalg.norm(vec)\n",
    "                \n",
    "                            bond_idx += 1\n",
    "                        break\n",
    "                    return(bond_idx)\n",
    "                \n",
    "                if idx >= n_bonds: \n",
    "                    break\n",
    "                idx = calc_torsions(bond.GetBeginAtom(), bond.GetEndAtom(), idx)\n",
    "            \n",
    "            if idx == 0:\n",
    "                features.append(None)\n",
    "                continue \n",
    "            for i in range(idx, n_bonds):\n",
    "                atom_index_of_bond[i] = atom_index_of_bond[idx-1]\n",
    "                \n",
    "                \n",
    "            graph_feat = featurizer.featurize([mol])[0]           \n",
    "            features.append((mol_file, mol_xyz, torsion_indices, torsion_matrix, subgraphs, atom_index_of_bond, n_atoms, graph_feat))\n",
    "        except:\n",
    "            features.append(None)\n",
    "    return(features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1276,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "32\n",
      "64\n",
      "96\n",
      "128\n",
      "160\n",
      "192\n",
      "224\n",
      "256\n",
      "288\n",
      "320\n",
      "352\n",
      "384\n",
      "416\n",
      "448\n",
      "480\n",
      "512\n",
      "544\n",
      "576\n",
      "608\n",
      "640\n",
      "672\n",
      "704\n",
      "736\n",
      "768\n",
      "800\n",
      "832\n",
      "864\n",
      "896\n",
      "928\n",
      "960\n",
      "992\n",
      "1024\n",
      "1056\n",
      "1088\n",
      "1120\n",
      "1152\n",
      "1184\n",
      "1216\n",
      "1248\n",
      "1280\n",
      "1312\n",
      "1344\n",
      "1376\n",
      "1408\n",
      "1440\n",
      "1472\n",
      "1504\n",
      "1536\n",
      "1568\n",
      "1600\n",
      "1632\n",
      "1664\n",
      "1696\n",
      "1728\n",
      "1760\n",
      "1792\n",
      "1824\n",
      "1856\n",
      "1888\n",
      "1920\n",
      "1952\n",
      "1984\n",
      "2016\n",
      "2048\n",
      "2080\n",
      "2112\n",
      "2144\n",
      "2176\n",
      "2208\n",
      "2240\n",
      "2272\n",
      "2304\n",
      "2336\n",
      "2368\n",
      "2400\n",
      "2432\n",
      "2464\n",
      "2496\n",
      "2528\n",
      "2560\n",
      "2592\n",
      "2624\n",
      "2656\n",
      "2688\n",
      "2720\n",
      "2752\n",
      "2784\n",
      "2816\n",
      "2848\n",
      "2880\n",
      "2912\n",
      "2944\n",
      "2976\n",
      "3008\n",
      "3040\n",
      "3072\n",
      "3104\n",
      "3136\n",
      "3168\n",
      "3200\n",
      "3232\n",
      "3264\n",
      "3296\n",
      "3328\n",
      "3360\n",
      "3392\n",
      "3424\n",
      "3456\n",
      "3488\n",
      "3520\n",
      "3552\n",
      "3584\n",
      "3616\n",
      "3648\n",
      "3680\n",
      "3712\n",
      "3744\n",
      "3776\n",
      "3808\n",
      "3840\n",
      "3872\n",
      "3904\n",
      "3936\n",
      "3968\n",
      "4000\n",
      "4032\n",
      "4064\n",
      "4096\n",
      "4128\n",
      "4160\n",
      "4192\n",
      "4224\n",
      "4256\n",
      "4288\n",
      "4320\n",
      "4352\n",
      "4384\n",
      "4416\n",
      "4448\n",
      "4480\n",
      "4512\n",
      "4544\n",
      "4576\n",
      "4608\n",
      "4640\n",
      "4672\n",
      "4704\n",
      "4736\n",
      "4768\n",
      "4800\n",
      "4832\n",
      "4864\n",
      "4896\n",
      "4928\n",
      "4960\n",
      "4992\n",
      "5024\n",
      "5056\n",
      "5088\n",
      "5120\n",
      "5152\n",
      "5184\n",
      "5216\n",
      "5248\n",
      "5280\n",
      "5312\n",
      "5344\n",
      "5376\n",
      "5408\n",
      "5440\n",
      "5472\n",
      "5504\n",
      "5536\n",
      "5568\n",
      "5600\n",
      "5632\n",
      "5664\n",
      "5696\n",
      "5728\n",
      "5760\n",
      "5792\n",
      "5824\n",
      "5856\n",
      "5888\n",
      "5920\n",
      "5952\n",
      "5984\n",
      "6016\n",
      "6048\n",
      "6080\n",
      "6112\n",
      "6144\n",
      "6176\n",
      "6208\n",
      "6240\n",
      "6272\n",
      "6304\n",
      "6336\n",
      "6368\n",
      "6400\n",
      "6432\n",
      "6464\n",
      "6496\n",
      "6528\n",
      "6560\n",
      "6592\n",
      "6624\n",
      "6656\n",
      "6688\n",
      "6720\n",
      "6752\n",
      "6784\n",
      "6816\n",
      "6848\n",
      "6880\n",
      "6912\n",
      "6944\n",
      "6976\n",
      "7008\n",
      "7040\n",
      "7072\n",
      "7104\n",
      "7136\n",
      "7168\n",
      "7200\n",
      "7232\n",
      "7264\n",
      "7296\n",
      "7328\n",
      "7360\n",
      "7392\n",
      "7424\n",
      "7456\n",
      "7488\n",
      "7520\n",
      "7552\n",
      "7584\n",
      "7616\n",
      "7648\n",
      "7680\n",
      "7712\n",
      "7744\n",
      "7776\n",
      "7808\n",
      "7840\n",
      "7872\n",
      "7904\n",
      "7936\n",
      "7968\n",
      "8000\n",
      "8032\n",
      "8064\n",
      "8096\n",
      "8128\n",
      "8160\n",
      "8192\n",
      "8224\n",
      "8256\n",
      "8288\n",
      "8320\n",
      "8352\n",
      "8384\n",
      "8416\n",
      "8448\n",
      "8480\n",
      "8512\n",
      "8544\n",
      "8576\n",
      "8608\n",
      "8640\n",
      "8672\n",
      "8704\n",
      "8736\n",
      "8768\n",
      "8800\n",
      "8832\n",
      "8864\n",
      "8896\n",
      "8928\n",
      "8960\n",
      "8992\n",
      "9024\n",
      "9056\n",
      "9088\n",
      "9120\n",
      "9152\n",
      "9184\n",
      "9216\n",
      "9248\n",
      "9280\n",
      "9312\n",
      "9344\n",
      "9376\n",
      "9408\n",
      "9440\n",
      "9472\n",
      "9504\n",
      "9536\n",
      "9568\n",
      "9600\n",
      "9632\n",
      "9664\n",
      "9696\n",
      "9728\n",
      "9760\n",
      "9792\n",
      "9824\n",
      "9856\n",
      "9888\n",
      "9920\n",
      "9952\n",
      "9984\n",
      "10016\n",
      "10048\n",
      "10080\n",
      "10112\n",
      "10144\n",
      "10176\n",
      "10208\n",
      "10240\n",
      "10272\n",
      "10304\n",
      "10336\n",
      "10368\n",
      "10400\n",
      "10432\n",
      "10464\n",
      "10496\n",
      "10528\n",
      "10560\n",
      "10592\n",
      "10624\n",
      "10656\n",
      "10688\n",
      "10720\n",
      "10752\n",
      "10784\n",
      "10816\n",
      "10848\n",
      "10880\n",
      "10912\n",
      "10944\n",
      "10976\n",
      "11008\n",
      "11040\n",
      "11072\n",
      "11104\n",
      "11136\n",
      "11168\n",
      "11200\n",
      "11232\n",
      "11264\n",
      "11296\n",
      "11328\n",
      "11360\n",
      "11392\n",
      "11424\n",
      "11456\n",
      "11488\n",
      "11520\n",
      "11552\n",
      "11584\n",
      "11616\n",
      "11648\n",
      "11680\n",
      "11712\n",
      "11744\n",
      "11776\n",
      "11808\n",
      "11840\n",
      "11872\n",
      "11904\n",
      "11936\n"
     ]
    }
   ],
   "source": [
    "import pickle\n",
    "max_n_atoms = 24\n",
    "D=10\n",
    "S=32\n",
    "feature_file = \"./dihed_xyz_features_pdbbind2.pkl\"\n",
    "#if not os.path.exists(feature_file):\n",
    "if 1== 1:\n",
    "    pdbbind_dir = \"/home/evan/Documents/deep_docking/datasets/v2015/\"\n",
    "    def find_files(directory, pattern):\n",
    "        for root, dirs, files in os.walk(directory):\n",
    "            for basename in files:\n",
    "                if fnmatch.fnmatch(basename, pattern):\n",
    "                    filename = os.path.join(root, basename)\n",
    "                    yield filename\n",
    "    ligand_files = []\n",
    "    for f in find_files(pdbbind_dir, \"*ligand.mol2\"):\n",
    "        ligand_files += [f]\n",
    "    ligand_files = ligand_files[:] + [\"CCCC\"]*32 #[\"/home/evan/Documents/deep_docking/alanine_dipeptide.pdb\"]*S  \n",
    "    features = featurize_mols(ligand_files, max_n_atoms, n_bonds=D)\n",
    "    with open(feature_file, \"wb\") as f:\n",
    "        pickle.dump(features, f, protocol=2)\n",
    "else:\n",
    "    with open(feature_file, \"rb\") as f:\n",
    "        features = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1277,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3662\n"
     ]
    }
   ],
   "source": [
    "features = [f for f in features if f is not None]\n",
    "print(len(features))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1278,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('/home/evan/Documents/deep_docking/datasets/v2015/4wks/4wks_ligand.mol2',\n",
       " array([[-0.0725, -0.0715, -1.7212],\n",
       "        [-0.8925,  0.3145, -0.4713],\n",
       "        [-0.2245, -0.1035,  1.0027],\n",
       "        [ 1.1895, -0.1395,  1.1898],\n",
       "        [ 0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ],\n",
       "        [ 0.    ,  0.    ,  0.    ]]),\n",
       " array([[[1, 0, 0, 0],\n",
       "         [0, 1, 0, 0],\n",
       "         [0, 0, 1, 0],\n",
       "         [0, 0, 0, 1],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0]],\n",
       " \n",
       "        [[0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0]],\n",
       " \n",
       "        [[0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0]],\n",
       " \n",
       "        [[0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0]],\n",
       " \n",
       "        [[0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0]],\n",
       " \n",
       "        [[0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0]],\n",
       " \n",
       "        [[0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0]],\n",
       " \n",
       "        [[0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0]],\n",
       " \n",
       "        [[0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0]],\n",
       " \n",
       "        [[0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0]]], dtype=uint8),\n",
       " array([[ 0.5528],\n",
       "        [ 0.    ],\n",
       "        [ 0.    ],\n",
       "        [ 0.    ],\n",
       "        [ 0.    ],\n",
       "        [ 0.    ],\n",
       "        [ 0.    ],\n",
       "        [ 0.    ],\n",
       "        [ 0.    ],\n",
       "        [ 0.    ]]),\n",
       " array([[[0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 1, ..., 0, 0, 0],\n",
       "         ..., \n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0]],\n",
       " \n",
       "        [[0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         ..., \n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0]],\n",
       " \n",
       "        [[0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         ..., \n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0]],\n",
       " \n",
       "        ..., \n",
       "        [[0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         ..., \n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0]],\n",
       " \n",
       "        [[0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         ..., \n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0]],\n",
       " \n",
       "        [[0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         ..., \n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0]]], dtype=uint8),\n",
       " array([[[0, 0],\n",
       "         [0, 1],\n",
       "         [1, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0]],\n",
       " \n",
       "        [[0, 0],\n",
       "         [0, 1],\n",
       "         [1, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0]],\n",
       " \n",
       "        [[0, 0],\n",
       "         [0, 1],\n",
       "         [1, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0]],\n",
       " \n",
       "        [[0, 0],\n",
       "         [0, 1],\n",
       "         [1, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0]],\n",
       " \n",
       "        [[0, 0],\n",
       "         [0, 1],\n",
       "         [1, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0]],\n",
       " \n",
       "        [[0, 0],\n",
       "         [0, 1],\n",
       "         [1, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0]],\n",
       " \n",
       "        [[0, 0],\n",
       "         [0, 1],\n",
       "         [1, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0]],\n",
       " \n",
       "        [[0, 0],\n",
       "         [0, 1],\n",
       "         [1, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0]],\n",
       " \n",
       "        [[0, 0],\n",
       "         [0, 1],\n",
       "         [1, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0]],\n",
       " \n",
       "        [[0, 0],\n",
       "         [0, 1],\n",
       "         [1, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0],\n",
       "         [0, 0]]], dtype=uint8),\n",
       " 4,\n",
       " (array([[1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0]], dtype=uint8), array([[1, 0, 0, ..., 0, 0, 0],\n",
       "         [1, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         ..., \n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0],\n",
       "         [0, 0, 0, ..., 0, 0, 0]], dtype=uint8)))"
      ]
     },
     "execution_count": 1278,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "features[3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1279,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "tf.reset_default_graph()\n",
    "\n",
    "alpha = 0.01\n",
    "n_layers = 1\n",
    "\n",
    "\n",
    "S = 32\n",
    "\n",
    "B = max_n_atoms\n",
    "\n",
    "D = 10\n",
    "p = 75\n",
    "\n",
    "z_dim = 32\n",
    "\n",
    "L_list = [p, 16, 64, 128, 256]\n",
    "\n",
    "dihed_per_bond = 1\n",
    "valence = 4\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1280,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#with tf.device('/gpu:0'):\n",
    "if 1==1:\n",
    "    x = tf.placeholder(tf.float32, [S, B, L_list[0]], name=\"atom_features\")\n",
    "    n_atoms = tf.placeholder(tf.int32, [S, 1], name=\"n_atoms\")\n",
    "    keep_prob = tf.placeholder(tf.float32, name=\"keep_prob\")\n",
    "    \n",
    "    adj_matrix = tf.placeholder(tf.float32, shape=[S, B, B], name=\"adj_matrix\")\n",
    "    \n",
    "    non_adj_matrix = tf.subtract(tf.ones_like(adj_matrix), adj_matrix)\n",
    "    \n",
    "    mol_xyz = tf.placeholder(tf.float32, shape=[S, B, 3], name=\"molxyz\")\n",
    "    \n",
    "    z = tf.random_uniform([S, D, z_dim], -1, 1)\n",
    "    zg = tf.random_uniform([S, D], -2, 2)\n",
    "    \n",
    "    mol_noise = tf.random_normal([S,D,1], mean=0,stddev=0.01)\n",
    "    \n",
    "    dihed_indices = tf.placeholder(tf.float32, shape=[S, D, B, 4])\n",
    "    dihed_tensor = tf.placeholder(tf.float32, shape=[S, D, 1])\n",
    "    rot_vectors = tf.placeholder(tf.float32, shape=[S, D, 3])\n",
    "    atom_index_of_bond = tf.placeholder(tf.float32, shape=[S, D, B, 2])\n",
    "    subgraph = tf.placeholder(tf.float32, shape=[S, D, B, B])\n",
    "    \n",
    "    phase = tf.placeholder(tf.bool)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1281,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#adapted from: http://stackoverflow.com/questions/37009647/compute-pairwise-distance-in-a-batch-without-replicating-tensor-in-tensorflow\n",
    "def compute_distance_matrix(A):\n",
    "    r = tf.reduce_sum(A*A, 1) # turn r into column vector \n",
    "    r = tf.reshape(r, [-1, 1]) \n",
    "    D = r - 2*tf.matmul(A, tf.transpose(A)) + tf.transpose(r)\n",
    "    return(D)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1282,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "building clades and adj\n",
      "0\n",
      "[24, 24]\n",
      "Building hidden layers\n"
     ]
    }
   ],
   "source": [
    "GW_list = [None for i in range(n_layers)]\n",
    "Gb_list = [None for i in range(n_layers)]\n",
    "\n",
    "DW_list = [None for i in range(n_layers)]\n",
    "Db_list = [None for i in range(n_layers)]\n",
    "\n",
    "\n",
    "Gh_list = [x]\n",
    "Dh_list = [x]\n",
    "\n",
    "adj_list = [adj_matrix]\n",
    "B_list = [B]\n",
    "clades_list = []\n",
    "graph_stride = 1.\n",
    "\n",
    "#with tf.device('/gpu:0'):\n",
    "if 1==1:\n",
    "\n",
    "    print(\"building clades and adj\")\n",
    "    for i in range(n_layers):\n",
    "        print(i)\n",
    "        B_list.append(int(np.ceil(B_list[i]/graph_stride)))\n",
    "        print(B_list)\n",
    "        clades_list.append(tf.stack([tf.one_hot(range(0,B_list[i],int(graph_stride)), depth=B_list[i])]*S, axis=0)) \n",
    "\n",
    "\n",
    "        adj_temp = tf.matmul(clades_list[i], adj_list[i])\n",
    "        sub_adj = tf.matmul(adj_temp, tf.transpose(adj_temp, perm=[0, 2, 1]))\n",
    "        sub_adj = tf.minimum(sub_adj, 1.)\n",
    "        adj_list.append(sub_adj)\n",
    "\n",
    "    print(\"Building hidden layers\")\n",
    "    for layer_idx in range(n_layers):\n",
    "        GW_list[layer_idx] = tf.Variable(tf.truncated_normal([L_list[layer_idx], L_list[layer_idx+1]], seed=2017), name=\"GW_list%d\" %layer_idx)\n",
    "        Gb_list[layer_idx] = tf.Variable(tf.ones([1, L_list[layer_idx+1]]))\n",
    "\n",
    "        DW_list[layer_idx] = tf.Variable(tf.truncated_normal([L_list[layer_idx], L_list[layer_idx+1]], seed=2017), name=\"DW_list%d\" %layer_idx)\n",
    "        Db_list[layer_idx] = tf.Variable(tf.ones([1, L_list[layer_idx+1]]))\n",
    "\n",
    "    f_dw1 = tf.Variable(tf.truncated_normal([L_list[n_layers], L_list[n_layers]]), name=\"f_dw1\")\n",
    "    f_db1 = tf.Variable(tf.ones([L_list[n_layers]]))\n",
    "\n",
    "    f_gw1 = tf.Variable(tf.truncated_normal([L_list[n_layers], L_list[n_layers]]))\n",
    "    f_gb1 = tf.Variable(tf.ones([L_list[n_layers]]))\n",
    "\n",
    "    g_w1_ini = tf.Variable(tf.truncated_normal([1, L_list[n_layers]*5+z_dim, 16], stddev=1.))\n",
    "    g_w1 = tf.tile(g_w1_ini, [S, 1, 1])\n",
    "    g_g1 = tf.Variable(tf.ones([1,1,16]))\n",
    "\n",
    "    g_b1 = tf.Variable(tf.ones([1, 1, 16]))\n",
    "\n",
    "    g_w2_ini = tf.Variable(tf.truncated_normal([1, 16, 16], stddev=1.))\n",
    "    g_w2 = tf.tile(g_w2_ini, [S, 1, 1])\n",
    "    g_g2 = tf.Variable(tf.ones([1,1,16]))\n",
    "\n",
    "    g_b2 = tf.Variable(tf.ones([1, 1, 16]))\n",
    "\n",
    "    g_w3_ini = tf.Variable(tf.truncated_normal([1, 16, 15], stddev=1.))\n",
    "    g_w3 = tf.tile(g_w3_ini, [S, 1, 1])\n",
    "    g_b3 = tf.Variable(tf.ones([1, 1, 15]))\n",
    "    g_g3 = tf.Variable(tf.ones([1,1,1]))\n",
    "\n",
    "    dd_w1_ini = tf.Variable(tf.truncated_normal([1, L_list[n_layers]*5+1, 32]))\n",
    "    dd_w1 = tf.tile(dd_w1_ini, [S, 1, 1])\n",
    "    dd_b1 = tf.Variable(tf.truncated_normal([1, 1, 32]))\n",
    "\n",
    "    dd_w2_ini = tf.Variable(tf.truncated_normal([1,32, 32]))\n",
    "    dd_w2 = tf.tile(dd_w2_ini, [S, 1, 1])\n",
    "    dd_b2 = tf.Variable(tf.truncated_normal([1,1,32]))\n",
    "\n",
    "    dd_w3_ini = tf.Variable(tf.truncated_normal([1,32, 1]))\n",
    "    dd_w3 = tf.tile(dd_w2_ini, [S, 1, 1])\n",
    "    dd_b3 = tf.Variable(tf.truncated_normal([1,1,1]))\n",
    "    \n",
    "    \n",
    "    d_w1_ini = tf.Variable(tf.truncated_normal([1, B*(L_list[n_layers]*2+1), 32]))\n",
    "    d_w1 = tf.tile(d_w1_ini, [S, 1, 1])\n",
    "    d_b1 = tf.Variable(tf.ones([1, 1, 32]))\n",
    "\n",
    "    d_w2_ini = tf.Variable(tf.truncated_normal([1, 32, 1]))\n",
    "    d_w2 = tf.tile(d_w2_ini, [S, 1, 1])\n",
    "    d_b2 = tf.Variable(tf.ones([1, 1, 1]))\n",
    "\n",
    "    lam1 = tf.Variable(tf.constant(-1.))\n",
    "    lam2 = tf.Variable(tf.constant(3.))\n",
    "    lam3 = tf.Variable(tf.constant(0.01))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1283,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def adjacency_conv_layer(x, W, b, adj, clades, L_in, L_out, layer_idx, S, B):\n",
    "    with tf.device('/gpu:0'):\n",
    "        print(\"layer_idx: %d\" %(layer_idx))\n",
    "        h = tf.matmul(adj, x, name=\"adj_mult_%d\" %layer_idx)\n",
    "        h = tf.reshape(h, shape=(S*B, L_in), name=\"adj_reshape_1_%d\" %layer_idx)\n",
    "\n",
    "        h = tf.matmul(h, W, name=\"adjconv_%d\" %layer_idx) + b\n",
    "    \n",
    "        h = tf.nn.tanh(h)\n",
    "        h = tf.reshape(h, (S, B, L_out), name=\"adj_reshape_2_%d\" %layer_idx)  \n",
    "\n",
    "        h = tf.matmul(clades, h)\n",
    "\n",
    "        print(\"within func h:\")\n",
    "        print(h)                                                                                                                              \n",
    "\n",
    "    return(h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1284,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#t = sess.run(tf.matmul(dihed_indices[:,:,:,0], x))\n",
    "#t = np.reshape(t, [t.shape[0]*t.shape[1],t.shape[2]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1285,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\ndef compute_dihedral_tensor(mols_batch, angle_tuples_batch):\\n    i_vec = tf.reshape(mols_batch, [S, B, 1, 3])\\n    i_vec = tf.tile(i_vec, [1, 1, B, 1])\\n\\n    j_vec = tf.reshape(mols_batch, [S, 1, B, 3])\\n    j_vec = tf.tile(j_vec, [1, B, 1, 1])\\n\\n    diff = tf.subtract(i_vec, j_vec)\\n    temp = tf.eye(B, batch_shape=[S])\\n    temp = tf.reshape(temp, [S, B, B, 1])\\n    temp = tf.tile(temp, [1, 1, 1, 3])\\n    diff = diff + temp\\n    diff = diff / tf.sqrt(tf.reduce_sum(tf.square(diff), axis=3, keep_dims=True))\\n    temp = tf.sqrt(tf.reduce_sum(tf.square(temp), axis=3, keep_dims=True))\\n    diff = tf.subtract(diff, temp)\\n\\n    print(\"diff should be [S, B, B, 3]\")\\n    print(diff)\\n\\n    ij = tf.tile(tf.reshape(diff, [S, B, 1, B, 3]), [1, 1, B, 1, 1])\\n    ik = tf.tile(tf.reshape(diff, [S, B, B, 1, 3]), [1, 1, 1, B, 1])\\n    cross = tf.cross(ij, ik)\\n    cross = cross / tf.sqrt(tf.reduce_sum(tf.square(cross), axis=4, keep_dims=True))\\n    \\n    ijk = tf.tile(tf.reshape(cross, [S, B, B, 1, B, 3]), [1m 1])\\n    \\n    dps = tf.reduce_sum(tf.multiply(ij ,ik), axis=4)\\n    #dps = tf.where(tf.is_nan(dps), tf.ones_like(dps) * 0., dps)\\n    print(\"dps should be [S, B, B, B]\")\\n    print(dps)\\n\\n    #angles_gen = tf.acos(tf.clip_by_value(dps, -1.0, 1.0))\\n    angles_computed = tf.where(tf.is_nan(dps), tf.zeros_like(dps), dps)\\n    angles_computed = tf.multiply(angles_computed, angle_tuples_batch)\\n    return(angles_computed)\\n'"
      ]
     },
     "execution_count": 1285,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def compute_angle_tensor(mols_batch, angle_tuples_batch):\n",
    "    i_vec = tf.reshape(mols_batch, [S, B, 1, 3])\n",
    "    i_vec = tf.tile(i_vec, [1, 1, B, 1])\n",
    "\n",
    "    j_vec = tf.reshape(mols_batch, [S, 1, B, 3])\n",
    "    j_vec = tf.tile(j_vec, [1, B, 1, 1])\n",
    "\n",
    "    diff = tf.subtract(i_vec, j_vec)\n",
    "    temp = tf.eye(B, batch_shape=[S])\n",
    "    temp = tf.reshape(temp, [S, B, B, 1])\n",
    "    temp = tf.tile(temp, [1, 1, 1, 3])\n",
    "    diff = diff + temp\n",
    "    diff = diff / tf.sqrt(tf.reduce_sum(tf.square(diff), axis=3, keep_dims=True))\n",
    "    temp = tf.sqrt(tf.reduce_sum(tf.square(temp), axis=3, keep_dims=True))\n",
    "    diff = tf.subtract(diff, temp)\n",
    "\n",
    "    print(\"diff should be [S, B, B, 3]\")\n",
    "    print(diff)\n",
    "\n",
    "    ij = tf.tile(tf.reshape(diff, [S, B, 1, B, 3]), [1, 1, B, 1, 1])\n",
    "    ik = tf.tile(tf.reshape(diff, [S, B, B, 1, 3]), [1, 1, 1, B, 1])\n",
    "    dps = tf.reduce_sum(tf.multiply(ij ,ik), axis=4)\n",
    "    #dps = tf.where(tf.is_nan(dps), tf.ones_like(dps) * 0., dps)\n",
    "    print(\"dps should be [S, B, B, B]\")\n",
    "    print(dps)\n",
    "\n",
    "    #angles_gen = tf.acos(tf.clip_by_value(dps, -1.0, 1.0))\n",
    "    angles_computed = tf.where(tf.is_nan(dps), tf.zeros_like(dps), dps)\n",
    "    angles_computed = tf.multiply(angles_computed, angle_tuples_batch)\n",
    "    return(angles_computed)\n",
    "\n",
    "def compute_dist_tensor(mols_batch):\n",
    "    dist_matrix = tf.concat([tf.reshape(compute_distance_matrix(tf.reshape(mol, [B, 3])), (1, B, B)) for mol in tf.split(mols_batch, S, axis=0)], axis=0)\n",
    "    #dist_matrix = tf.where(tf.is_nan(dist_matrix), tf.ones_like(dist_matrix) * 0., dist_matrix)\n",
    "    print(\"dist_matrix\")\n",
    "    print(dist_matrix)\n",
    "    #dist_matrix = compute_distance_matrix(mol)\n",
    "    return(dist_matrix)\n",
    "\n",
    "def compute_rot_matrices(t, vec):\n",
    "    t = tf.squeeze(t)\n",
    "    x, y, z = vec[:, 0], vec[:, 1], vec[:,2]\n",
    "    rot_matrix = tf.reshape(tf.transpose(tf.stack([tf.cos(t) + tf.square(x) * (1 - tf.cos(t)), \n",
    "                               x*y*(1-tf.cos(t)) - z*tf.sin(t),\n",
    "                               x*z*(1-tf.cos(t)) + y*tf.sin(t),\n",
    "                              y*x*(1-tf.cos(t)) + z*tf.sin(t),\n",
    "                               tf.cos(t) + y*y*(1-tf.cos(t)),\n",
    "                               y*z*(1-tf.cos(t)) - x*tf.sin(t),\n",
    "                              z*x*(1-tf.cos(t)) - y*tf.sin(t),\n",
    "                               z*y*(1-tf.cos(t)) + x*tf.sin(t),\n",
    "                               tf.cos(t)+ z*z*(1-tf.cos(t))])), (S, 3, 3))\n",
    "    print(\"rot_matrix\")\n",
    "    print(rot_matrix)\n",
    "\n",
    "    return(rot_matrix)\n",
    "\n",
    "def compute_rot_matrix(t, vec):\n",
    "    t = tf.squeeze(t)\n",
    "    x, y, z = vec[0, 0], vec[0, 1], vec[0, 2]\n",
    "    rot_matrix = tf.reshape(tf.transpose(tf.stack([tf.cos(t) + tf.square(x) * (1 - tf.cos(t)), \n",
    "                               x*y*(1-tf.cos(t)) - z*tf.sin(t),\n",
    "                               x*z*(1-tf.cos(t)) + y*tf.sin(t),\n",
    "                              y*x*(1-tf.cos(t)) + z*tf.sin(t),\n",
    "                               tf.cos(t) + y*y*(1-tf.cos(t)),\n",
    "                               y*z*(1-tf.cos(t)) - x*tf.sin(t),\n",
    "                              z*x*(1-tf.cos(t)) - y*tf.sin(t),\n",
    "                               z*y*(1-tf.cos(t)) + x*tf.sin(t),\n",
    "                               tf.cos(t)+ z*z*(1-tf.cos(t))])), (3, 3))\n",
    "    print(\"rot_matrix\")\n",
    "    print(rot_matrix)\n",
    "\n",
    "    return(rot_matrix)\n",
    "\n",
    "\n",
    "\"\"\"\n",
    "def compute_dihedral_tensor(mols_batch, angle_tuples_batch):\n",
    "    i_vec = tf.reshape(mols_batch, [S, B, 1, 3])\n",
    "    i_vec = tf.tile(i_vec, [1, 1, B, 1])\n",
    "\n",
    "    j_vec = tf.reshape(mols_batch, [S, 1, B, 3])\n",
    "    j_vec = tf.tile(j_vec, [1, B, 1, 1])\n",
    "\n",
    "    diff = tf.subtract(i_vec, j_vec)\n",
    "    temp = tf.eye(B, batch_shape=[S])\n",
    "    temp = tf.reshape(temp, [S, B, B, 1])\n",
    "    temp = tf.tile(temp, [1, 1, 1, 3])\n",
    "    diff = diff + temp\n",
    "    diff = diff / tf.sqrt(tf.reduce_sum(tf.square(diff), axis=3, keep_dims=True))\n",
    "    temp = tf.sqrt(tf.reduce_sum(tf.square(temp), axis=3, keep_dims=True))\n",
    "    diff = tf.subtract(diff, temp)\n",
    "\n",
    "    print(\"diff should be [S, B, B, 3]\")\n",
    "    print(diff)\n",
    "\n",
    "    ij = tf.tile(tf.reshape(diff, [S, B, 1, B, 3]), [1, 1, B, 1, 1])\n",
    "    ik = tf.tile(tf.reshape(diff, [S, B, B, 1, 3]), [1, 1, 1, B, 1])\n",
    "    cross = tf.cross(ij, ik)\n",
    "    cross = cross / tf.sqrt(tf.reduce_sum(tf.square(cross), axis=4, keep_dims=True))\n",
    "    \n",
    "    ijk = tf.tile(tf.reshape(cross, [S, B, B, 1, B, 3]), [1m 1])\n",
    "    \n",
    "    dps = tf.reduce_sum(tf.multiply(ij ,ik), axis=4)\n",
    "    #dps = tf.where(tf.is_nan(dps), tf.ones_like(dps) * 0., dps)\n",
    "    print(\"dps should be [S, B, B, B]\")\n",
    "    print(dps)\n",
    "\n",
    "    #angles_gen = tf.acos(tf.clip_by_value(dps, -1.0, 1.0))\n",
    "    angles_computed = tf.where(tf.is_nan(dps), tf.zeros_like(dps), dps)\n",
    "    angles_computed = tf.multiply(angles_computed, angle_tuples_batch)\n",
    "    return(angles_computed)\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1286,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def wnorm_layer(w, x, gamma, alpha):\n",
    "    axes = range(0, len(w.get_shape()))\n",
    "    wn = tf.nn.l2_normalize(w, -1)*gamma\n",
    "    y = tf.nn.tanh(tf.matmul(x, wn) - alpha) + alpha\n",
    "    return(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1287,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generator(W_list, b_list, h_list, L_list, n_layers, mols):\n",
    "    if 1==1:\n",
    "    #with tf.device('/gpu:0'):\n",
    "    #with tf.variable_scope(\"generator\"):\n",
    "        for layer_idx in range(n_layers):\n",
    "            h_list.append(tf.nn.dropout(adjacency_conv_layer(h_list[layer_idx], W_list[layer_idx],\n",
    "                                                         b_list[layer_idx], adj_list[layer_idx], \n",
    "                                                         clades_list[layer_idx],\n",
    "                                                         L_list[layer_idx], L_list[layer_idx+1],\n",
    "                                                         layer_idx, S, B_list[layer_idx]), keep_prob))\n",
    "        L_final = L_list[n_layers]  \n",
    "\n",
    "        print(\"h_list[-1]\")\n",
    "        print(h_list[-1])\n",
    "\n",
    "        print(\"B_list\")\n",
    "        print(B_list)\n",
    "\n",
    "        fingerprints = tf.tile(tf.reshape(tf.nn.tanh(tf.matmul(tf.reduce_sum(h_list[-1], axis=1), f_gw1) + f_gb1), [S, 1, L_final]), [1, D, 1])\n",
    "\n",
    "        #fingerprints = tf.tile(tf.reshape(tf.layers.batch_normalization(tf.nn.tanh(tf.matmul(tf.reduce_sum(h_list[-1], axis=1), f_gw1) + f_gb1), training=phase, name=\"gbn1\"), [S, 1, L_final]), [1, D, 1])\n",
    "        h_final = h_list[-1]\n",
    "\n",
    "        d0 = []\n",
    "        for l in range(0, valence*dihed_per_bond):\n",
    "            d0.append(tf.matmul(dihed_indices[:,:,:,l], h_final, name=\"dihed_mult\"))\n",
    "\n",
    "        d0 = tf.concat(d0, axis=2)\n",
    "        print(\"d0 should be shape %s\" %(str((S, D, L_list[n_layers]*4))))\n",
    "        d0 = tf.concat([d0, fingerprints, z], axis=2)\n",
    "\n",
    "        #PREDICT ANGLES:\n",
    "        g_h1 = tf.nn.dropout(tf.nn.tanh(tf.matmul(d0, g_w1) + g_b1), keep_prob)\n",
    "        #g_h1 = wnorm_layer(g_w1, d0, g_g1, g_b1)\n",
    "        #g_h1 = tf.nn.tanh(tf.matmul(d0, wnorm(g_w1, g_g1)) + g_b1)\n",
    "        #g_h1 = tf.reshape(g_h1, [S, 5, 16])\n",
    "        print(\"g_h1\")\n",
    "        print(g_h1)\n",
    "        #g_h2 = tf.nn.dropout(tf.nn.tanh(tf.matmul(g_h1, g_w2) + g_b2), 1.)\n",
    "        g_h2 = tf.nn.dropout(tf.nn.tanh(tf.matmul(g_h1, g_w2)+g_b2), keep_prob)\n",
    "        \n",
    "        print(\"g_h2\")\n",
    "        print(g_h2)\n",
    "        \n",
    "        #g_h2 = wnorm_layer(g_w2, g_h1, g_g2, g_b2)\n",
    "        \n",
    "        g_h3 = tf.reshape(tf.matmul(g_h2, g_w3) + g_b3, [S, D, 15])\n",
    "        #g_h3 = wnorm_layer(g_w3, g_h2, g_g3, g_b3)\n",
    "        \n",
    "        print(\"g_h3\")\n",
    "        print(g_h3)\n",
    "        #g_h3 = tf.atan2(tf.sin(g_h3), tf.cos(g_h3))\n",
    "\n",
    "        def sample_gen_cdf(params, noise):\n",
    "            angle = 0.\n",
    "            for i in range(0,5):\n",
    "                comp = tf.nn.tanh(tf.add(tf.multiply(params[:,:,i*3+1], noise), params[:,:,i*3+2]))\n",
    "                print(\"comp1\")\n",
    "                print(comp)\n",
    "                comp = tf.multiply(tf.clip_by_value(params[:,:,i*3], 0.0, 1.0), comp)\n",
    "                print(\"comp2\")\n",
    "                print(comp)\n",
    "                angle = tf.add(angle, comp)\n",
    "            return(angle)\n",
    "        \n",
    "        output = sample_gen_cdf(g_h3, zg)\n",
    "        output = tf.reshape(output, [S, D,1])\n",
    "\n",
    "        \"\"\"\n",
    "        new_mols = []\n",
    "        for i in range(0,S):\n",
    "            mol = mol_xyz[i]\n",
    "            for j in range(0,D):\n",
    "                rot_point = tf.matmul(tf.reshape(atom_index_of_bond[i,j,:], (1, B)), mol)\n",
    "                mol = tf.subtract(mol, rot_point)\n",
    "                rot_mat = compute_rot_matrix(tf.subtract(g_h3[i, j], dihed_tensor[i, j]), rot_vectors[i, j])\n",
    "                if i ==0 and j ==0:\n",
    "                    rot_mat_r = rot_mat\n",
    "\n",
    "                frag_i = tf.transpose(tf.matmul(subgraph[i, j, :, :], mol), perm=[1,0])\n",
    "                new_xyz = tf.transpose(tf.matmul(rot_mat, frag_i), perm=[1,0])\n",
    "\n",
    "                other_subgraph = tf.subtract(tf.eye(B), subgraph[i, j, :, :])\n",
    "                mol = tf.add(new_xyz, tf.matmul(other_subgraph, mol))\n",
    "            new_mols.append(mol)\n",
    "        rot_mols = tf.stack(new_mols, axis=0)\n",
    "\n",
    "        #ROTATE MOLECULE: \n",
    "        rot_mols = []\n",
    "        #(S, B, 3)\n",
    "        for i in range(0, S):\n",
    "            rot_mol = mols[i]\n",
    "            for j in range(0,10):\n",
    "                #atom_index_of_bond[:, j, :, 0] should be (S, 1, B)\n",
    "                #\n",
    "                rot_points = tf.matmul(tf.reshape(atom_index_of_bond[i, j, :, 0], (1, B)), rot_mol)\n",
    "                rot_mol = rot_mol - rot_points\n",
    "                rot_vecs = tf.reshape(tf.matmul(tf.reshape(atom_index_of_bond[i, j, :, 0], (1, B)), rot_mol) - tf.matmul(tf.reshape(atom_index_of_bond[i, j, :, 1], (1, B)), rot_mol), (1, 3))\n",
    "                rot_vecs = rot_vecs / tf.norm(rot_vecs)\n",
    "                rot_mat = compute_rot_matrix(tf.subtract(output[i, j], dihed_tensor[i, j]), rot_vecs)\n",
    "                if j == 0:\n",
    "                    rot_mat_r = rot_mat\n",
    "                    rot_points_r = rot_points\n",
    "                    rot_vecs_r = rot_vecs\n",
    "\n",
    "                frag_i = tf.transpose(tf.matmul(subgraph[i, j, :, :], rot_mol), perm=[1, 0])\n",
    "                frag_i_xyz = tf.transpose(tf.matmul(rot_mat, frag_i), perm=[1,0])\n",
    "                other_subgraph = tf.subtract(tf.eye(B), subgraph[i, j, :, :])\n",
    "                rot_mol = tf.add(frag_i_xyz, tf.matmul(other_subgraph, rot_mol))\n",
    "            rot_mols.append(rot_mol)\n",
    "        rot_mols = tf.stack(rot_mols, axis=0)\n",
    "        \n",
    "        \n",
    "        \"\"\"\n",
    "        rot_mols = mols\n",
    "        #(S, B, 3)\n",
    "        for j in range(0,D):\n",
    "            #atom_index_of_bond[:, j, :, 0] should be (S, 1, B)\n",
    "            #\n",
    "            rot_points = tf.matmul(tf.reshape(atom_index_of_bond[:, j, :, 0], (S, 1, B)), rot_mols)\n",
    "            rot_mols = rot_mols - rot_points\n",
    "            rot_vecs = tf.reshape(tf.matmul(tf.reshape(atom_index_of_bond[:, j, :, 0], (S, 1, B)), rot_mols) - tf.matmul(tf.reshape(atom_index_of_bond[:, j, :, 1], (S, 1, B)), rot_mols), (S, 3))\n",
    "            rot_vecs = tf.nn.l2_normalize(rot_vecs, dim=-1)\n",
    "            rot_mat = compute_rot_matrices(tf.subtract(output[:, j], dihed_tensor[:, j]), rot_vecs)\n",
    "            if 1==1:\n",
    "                rot_mat_r = rot_mat\n",
    "                rot_points_r = rot_points\n",
    "                rot_vecs_r = rot_vecs\n",
    "            \n",
    "            frag_i = tf.transpose(tf.matmul(subgraph[:, j, :, :], rot_mols), perm=[0, 2, 1])\n",
    "            frag_i_xyz = tf.transpose(tf.matmul(rot_mat, frag_i), perm=[0,2,1])\n",
    "            other_subgraph = tf.subtract(tf.tile(tf.reshape(tf.eye(B), (1, B, B)), [S, 1, 1]), subgraph[:, j, :, :])\n",
    "            rot_mols = tf.add(frag_i_xyz, tf.matmul(other_subgraph, rot_mols))\n",
    "            \n",
    "            \n",
    "        return(rot_mols, output, rot_mat_r, rot_points_r, rot_vecs_r)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1288,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def Normal(mu, sigma, a):\n",
    "    return(tf.sqrt(1/(2*3.14159*(sigma**2))) * tf.exp(-((a-mu)**2) / (2*sigma**2)))\n",
    "def optimal_rotational_quaternion(r):\n",
    "    \"\"\"Just need the largest eigenvalue of this to minimize RMSD over rotations\n",
    "    \n",
    "    References\n",
    "    ----------\n",
    "    [1] http://dx.doi.org/10.1002/jcc.20110\n",
    "    \"\"\"\n",
    "    # @formatter:off\n",
    "    return [\n",
    "        [r[0][0]+r[1][1]+r[2][2], r[1][2]-r[2][1],         r[2][0]-r[0][2],         r[0][1]-r[1][0]        ],\n",
    "        [r[1][2]-r[2][1],         r[0][0]-r[1][1]-r[2][2], r[0][1]+r[1][0],         r[0][2]+r[2][0]        ],\n",
    "        [r[2][0]-r[0][2],         r[0][1]+r[1][0],        -r[0][0]+r[1][1]-r[2][2], r[1][2]+r[2][1]        ],\n",
    "        [r[0][1]-r[1][0],         r[0][2]+r[2][0],         r[1][2]+r[2][1],        -r[0][0]-r[1][1]+r[2][2]],\n",
    "    ]\n",
    "    # @formatter:on\n",
    "\n",
    "\n",
    "def squared_deviation(frame, target):\n",
    "    \"\"\"Calculate squared deviation (n_atoms * RMSD^2) from `frame` to `target`\n",
    "    First we compute `R` which is the ordinary cross-correlation of xyz coordinates.\n",
    "    Turns out you can do a bunch of quaternion math to find an eigen-expression for finding optimal\n",
    "    rotations. There aren't quaternions in tensorflow, so we use the handy formula for turning\n",
    "    quaternions back into 4-matrices. This is the `F` matrix. We find its leading eigenvalue\n",
    "    to get the MSD after optimal rotation. Note: *finding* the optimal rotation requires the values\n",
    "    and vectors, but we don't care.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    frame, target : Tensor, shape=(n_atoms, 3)\n",
    "        Calculate the MSD between these two frames\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    sd : Tensor, shape=(0,)\n",
    "        Divide by number of atoms and take the square root for RMSD\n",
    "    \"\"\"\n",
    "    R = tf.matmul(frame, target, transpose_a=True, name='R')\n",
    "    R_parts = [tf.unstack(t) for t in tf.unstack(R)]\n",
    "    F_parts = optimal_rotational_quaternion(R_parts)\n",
    "    F = tf.stack(F_parts, name='F')\n",
    "    vals, vecs = tf.self_adjoint_eig(F, name='eig')\n",
    "    # This isn't differentiable for some godforsaken reason.\n",
    "    # vals = tf.self_adjoint_eigvals(F, name='vals')\n",
    "    lmax = tf.unstack(vals)[-1]\n",
    "    sd = tf.reduce_sum(frame ** 2 + target ** 2) - 2 * lmax\n",
    "    return sd\n",
    "\n",
    "def rmsd(frame, target, n_atoms):\n",
    "    \"\"\"Convenience function for actually returning the RMSD\n",
    "    \n",
    "    You should probably use squared_deviation when optimizing.\n",
    "    \"\"\"\n",
    "    return tf.sqrt(squared_deviation(frame, target) / n_atoms)\n",
    "\n",
    "\n",
    "def multi_sd(frames, target):\n",
    "    return tf.map_fn(lambda x: squared_deviation(x, target), frames, name='MultiMSD')\n",
    "\n",
    "\n",
    "def sum_sd(frames, target):\n",
    "    return tf.reduce_sum(multi_sd(frames, target), name='SumMSD')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1289,
   "metadata": {},
   "outputs": [],
   "source": [
    "def discriminator(W_list, b_list, h_list, L_list, n_layers, mol_real, mol_gen, dihed, reuse=False):\n",
    "    with tf.variable_scope(\"discriminator\"):\n",
    "\n",
    "        #with tf.device('/gpu:0'):\n",
    "        if 1==1:\n",
    "            for layer_idx in range(n_layers):\n",
    "                h_list.append(tf.nn.dropout(adjacency_conv_layer(h_list[layer_idx], W_list[layer_idx],\n",
    "                                                 b_list[layer_idx], adj_list[layer_idx], \n",
    "                                                 clades_list[layer_idx],\n",
    "                                                 L_list[layer_idx], L_list[layer_idx+1],\n",
    "                                                 layer_idx, S, B_list[layer_idx]), keep_prob))        \n",
    "            #DIHEDRAL LOSS:\n",
    "            fingerprints = tf.tile(tf.reshape(tf.nn.tanh(tf.matmul(tf.reduce_sum(h_list[-1], axis=1), f_dw1) + f_db1), [S, 1, L_list[n_layers]]), [1, D, 1])\n",
    "            #fingerprints = tf.tile(tf.reshape(tf.layers.batch_normalization(tf.nn.tanh(tf.matmul(tf.reduce_sum(h_list[-1], axis=1), f_dw1) + f_db1), training=phase, name=\"dbn1\", reuse=reuse), [S, 1, L_list[n_layers]]), [1, D, 1])\n",
    "            h_final = h_list[-1]\n",
    "\n",
    "            d0 = []\n",
    "            for l in range(0, valence*dihed_per_bond):\n",
    "                d0.append(tf.matmul(dihed_indices[:,:,:,l], h_final, name=\"dihed_mult\"))\n",
    "\n",
    "            d0 = tf.concat(d0, axis=2)\n",
    "            print(\"d0 should be shape %s\" %(str((S, D, L_list[n_layers]*4+1))))\n",
    "            print(d0)\n",
    "            d0 = tf.concat([d0, fingerprints, dihed], axis=2)\n",
    "            d0 = tf.reshape(d0, [S, D, L_list[n_layers]*5 + 1])\n",
    "            \n",
    "            d1 = tf.nn.dropout(tf.nn.tanh(tf.matmul(d0, dd_w1) + dd_b1), 1.)\n",
    "            #d1 = tf.layers.batch_normalization(tf.nn.tanh(tf.matmul(d0, dd_w1) + dd_b1), training=phase, name=\"dbn2\", reuse=reuse, scale=True)\n",
    "            d2 = tf.nn.tanh(tf.matmul(d1, dd_w2) + dd_b2)\n",
    "            d3 = tf.reduce_mean(tf.matmul(d2, dd_w3) + dd_b3, axis=1, keep_dims=True)\n",
    "            \"\"\"\n",
    "            \n",
    "            #prob = tf.nn.sigmoid(d2)\n",
    "            \n",
    "            print(\"d2\")\n",
    "            print(d2)\n",
    "            \n",
    "            prob = 0.\n",
    "            for i in range(0,3):\n",
    "                prob += d2[:, :, i*3] * Normal(mu=d2[:, :, i*3+1], sigma=d2[:, :, i*3+2], a=tf.reshape(dihed, (S,D)))\n",
    "            \n",
    "            return(prob)\n",
    "            \"\"\"\n",
    "            \n",
    "            #prob = tf.reduce_mean(tf.reshape(prob, [S, D]), axis=1, keep_dims=True)\n",
    "            #prob = tf.nn.tanh(tf.reduce_mean(prob, axis=1, keep_dims=True))\n",
    "\n",
    "\n",
    "            #COMPUTE CONTACT LOSS:\n",
    "            \n",
    "            rmsds = []\n",
    "            for i in range(0, S):\n",
    "                restricting_mat = tf.eye(B)[:n_atoms[i,0]]\n",
    "                mol_gen_i = tf.matmul(restricting_mat, mol_gen[i, :, :])\n",
    "                mol_real_i = tf.matmul(restricting_mat, mol_real[i, :, :])\n",
    "            \n",
    "                mol_gen_i = mol_gen_i - tf.reduce_mean(mol_gen_i, axis=0, keep_dims=True)\n",
    "                mol_real_i = mol_real_i - tf.reduce_mean(mol_real_i, axis=0, keep_dims=True)\n",
    "                rmsds.append(squared_deviation(mol_gen_i, mol_real_i)/tf.cast(n_atoms[i,0], tf.float32))\n",
    "            rmsds = tf.stack(rmsds)\n",
    "            dist_loss = tf.nn.sigmoid(lam1*(rmsds+lam2))\n",
    "\n",
    "            dist_real = compute_dist_tensor(mol_real)\n",
    "            dist_gen = compute_dist_tensor(mol_gen)\n",
    "            \n",
    "\n",
    "            #dist_loss = tf.nn.sigmoid(tf.multiply(tf.abs(dist_gen - dist_real) - lam2, lam1))\n",
    "            contact_dist = tf.multiply(dist_gen, non_adj_matrix)\n",
    "            contact_loss = tf.multiply(tf.reshape(tf.reduce_mean(tf.exp(tf.multiply(tf.multiply(tf.subtract(contact_dist, 2.), -1.), 3.)), axis=[1,2]), (S,1)), lam1)  \n",
    "\n",
    "            feat_i = tf.reshape(h_final, (S, B, 1, L_list[n_layers]))\n",
    "            feat_i = tf.tile(feat_i, [1, 1, B, 1])\n",
    "\n",
    "            feat_j = tf.reshape(h_final, (S, B, 1, L_list[n_layers]))\n",
    "            feat_j = tf.transpose(feat_j, perm=[0, 2, 1, 3])\n",
    "            feat_j = tf.tile(feat_j, [1, B, 1, 1])\n",
    "\n",
    "            d_h1 = tf.concat([tf.reshape(contact_dist, [S, B, B, 1]), feat_i, feat_j], axis=3)\n",
    "            d_h1 = tf.reshape(d_h1, [S, B, B*(L_list[n_layers]*2+1)])\n",
    "            d_h2 = tf.nn.dropout(tf.nn.tanh(tf.matmul(d_h1, d_w1) + d_b1), 1.)\n",
    "            d_h2 = tf.nn.sigmoid(tf.matmul(d_h2, d_w2) + d_b2)\n",
    "            print(\"d_h2\")\n",
    "            print(d_h2)\n",
    "            return(dist_loss, dist_real, dist_gen, rmsds, mol_gen_i)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1290,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "theta_D = DW_list[:n_layers] + Db_list[:n_layers] + [d_w1_ini, d_b1, d_w2_ini, d_b2, dd_w1_ini, dd_b1, dd_w2_ini, dd_b2, dd_w3_ini, dd_b3, f_dw1, f_db1, lam1, lam2] + [n for n in tf.trainable_variables() if \"dbn\" in str(n)]\n",
    "theta_G = GW_list[:n_layers] + Gb_list[:n_layers] + [g_w1_ini, g_b1, g_w2_ini, g_b2, g_w3_ini, g_b3, f_gw1, f_gb1, g_g1, g_g2, g_g3] + [n for n in tf.trainable_variables() if \"gbn\" in str(n)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1291,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "layer_idx: 0\n",
      "within func h:\n",
      "Tensor(\"MatMul_2:0\", shape=(32, 24, 16), dtype=float32, device=/device:GPU:0)\n",
      "h_list[-1]\n",
      "Tensor(\"dropout/mul:0\", shape=(32, 24, 16), dtype=float32)\n",
      "B_list\n",
      "[24, 24]\n",
      "d0 should be shape (32, 10, 64)\n",
      "g_h1\n",
      "Tensor(\"dropout_1/mul:0\", shape=(32, 10, 16), dtype=float32)\n",
      "g_h2\n",
      "Tensor(\"dropout_2/mul:0\", shape=(32, 10, 16), dtype=float32)\n",
      "g_h3\n",
      "Tensor(\"Reshape_1:0\", shape=(32, 10, 15), dtype=float32)\n",
      "comp1\n",
      "Tensor(\"Tanh_4:0\", shape=(32, 10), dtype=float32)\n",
      "comp2\n",
      "Tensor(\"Mul_1:0\", shape=(32, 10), dtype=float32)\n",
      "comp1\n",
      "Tensor(\"Tanh_5:0\", shape=(32, 10), dtype=float32)\n",
      "comp2\n",
      "Tensor(\"Mul_3:0\", shape=(32, 10), dtype=float32)\n",
      "comp1\n",
      "Tensor(\"Tanh_6:0\", shape=(32, 10), dtype=float32)\n",
      "comp2\n",
      "Tensor(\"Mul_5:0\", shape=(32, 10), dtype=float32)\n",
      "comp1\n",
      "Tensor(\"Tanh_7:0\", shape=(32, 10), dtype=float32)\n",
      "comp2\n",
      "Tensor(\"Mul_7:0\", shape=(32, 10), dtype=float32)\n",
      "comp1\n",
      "Tensor(\"Tanh_8:0\", shape=(32, 10), dtype=float32)\n",
      "comp2\n",
      "Tensor(\"Mul_9:0\", shape=(32, 10), dtype=float32)\n",
      "rot_matrix\n",
      "Tensor(\"Reshape_7:0\", shape=(32, 3, 3), dtype=float32)\n",
      "rot_matrix\n",
      "Tensor(\"Reshape_13:0\", shape=(32, 3, 3), dtype=float32)\n",
      "rot_matrix\n",
      "Tensor(\"Reshape_19:0\", shape=(32, 3, 3), dtype=float32)\n",
      "rot_matrix\n",
      "Tensor(\"Reshape_25:0\", shape=(32, 3, 3), dtype=float32)\n",
      "rot_matrix\n",
      "Tensor(\"Reshape_31:0\", shape=(32, 3, 3), dtype=float32)\n",
      "rot_matrix\n",
      "Tensor(\"Reshape_37:0\", shape=(32, 3, 3), dtype=float32)\n",
      "rot_matrix\n",
      "Tensor(\"Reshape_43:0\", shape=(32, 3, 3), dtype=float32)\n",
      "rot_matrix\n",
      "Tensor(\"Reshape_49:0\", shape=(32, 3, 3), dtype=float32)\n",
      "rot_matrix\n",
      "Tensor(\"Reshape_55:0\", shape=(32, 3, 3), dtype=float32)\n",
      "rot_matrix\n",
      "Tensor(\"Reshape_61:0\", shape=(32, 3, 3), dtype=float32)\n",
      "dihed, returned\n",
      "Tensor(\"Reshape_2:0\", shape=(32, 10, 1), dtype=float32)\n",
      "layer_idx: 0\n",
      "within func h:\n",
      "Tensor(\"discriminator/MatMul:0\", shape=(32, 24, 16), dtype=float32, device=/device:GPU:0)\n",
      "d0 should be shape (32, 10, 65)\n",
      "Tensor(\"discriminator/concat:0\", shape=(32, 10, 64), dtype=float32)\n",
      "dist_matrix\n",
      "Tensor(\"discriminator/concat_2:0\", shape=(32, 24, 24), dtype=float32)\n",
      "dist_matrix\n",
      "Tensor(\"discriminator/concat_3:0\", shape=(32, 24, 24), dtype=float32)\n",
      "d_h2\n",
      "Tensor(\"discriminator/Sigmoid_1:0\", shape=(32, 24, 1), dtype=float32)\n",
      "layer_idx: 0\n",
      "within func h:\n",
      "Tensor(\"discriminator_1/MatMul:0\", shape=(32, 24, 16), dtype=float32, device=/device:GPU:0)\n",
      "d0 should be shape (32, 10, 65)\n",
      "Tensor(\"discriminator_1/concat:0\", shape=(32, 10, 64), dtype=float32)\n",
      "dist_matrix\n",
      "Tensor(\"discriminator_1/concat_2:0\", shape=(32, 24, 24), dtype=float32)\n",
      "dist_matrix\n",
      "Tensor(\"discriminator_1/concat_3:0\", shape=(32, 24, 24), dtype=float32)\n",
      "d_h2\n",
      "Tensor(\"discriminator_1/Sigmoid_1:0\", shape=(32, 24, 1), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "import time \n",
    "\n",
    "#with tf.device('/gpu:0'):\n",
    "if 1==1:\n",
    "\n",
    "    G_sample, dihed, rot_mat_i, rot_points, rot_vecs = generator(GW_list, Gb_list, Gh_list, L_list, n_layers, mol_xyz)\n",
    "    print(\"dihed, returned\")\n",
    "    print(dihed)\n",
    "    #dihed = tf.clip_by_value(dihed, -3.2, 3.2)\n",
    "    \n",
    "    D_real, dist_real_real, dist_real_fake, rmsds_real, mol_gen_real = discriminator(DW_list, Db_list, Dh_list, L_list, n_layers, mol_xyz, mol_xyz, dihed_tensor + mol_noise, reuse=False)\n",
    "    D_fake, dist_fake_real, dist_fake_fake, rmsds_fake, mol_gen_fake = discriminator(DW_list, Db_list, Dh_list, L_list, n_layers, mol_xyz, G_sample, dihed, reuse=True)\n",
    "\n",
    "    #D_loss = tf.reduce_mean(D_real) - tf.reduce_mean(D_fake)\n",
    "    #G_loss = -tf.reduce_mean(D_fake)\n",
    "    \n",
    "    D_loss = tf.reduce_sum(tf.square(D_real-1) + tf.square(D_fake))/2\n",
    "    G_loss = tf.reduce_sum(tf.square(D_fake-1))/2\n",
    "    \n",
    "    D_solver = (tf.train.RMSPropOptimizer(learning_rate=5e-4)\n",
    "                .minimize(D_loss, var_list=theta_D))\n",
    "    G_solver = (tf.train.RMSPropOptimizer(learning_rate=5e-4)\n",
    "                .minimize(G_loss, var_list=theta_G))\n",
    "    \n",
    "    d_ops = [op for op in tf.get_collection(tf.GraphKeys.UPDATE_OPS) if \"dbn\" in str(op)]\n",
    "    g_ops = [op for op in tf.get_collection(tf.GraphKeys.UPDATE_OPS) if \"gbn\" in str(op)]\n",
    "\n",
    "    mb_size = S\n",
    "\n",
    "    sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "\n",
    "\n",
    "    preds = []\n",
    "\n",
    "    d_losses = []\n",
    "    g_losses = []\n",
    "    #print(sess.run(label_placeholder))\n",
    "    \n",
    "    \n",
    "    # WGAN lipschitz-penalty\n",
    "    \"\"\"\n",
    "    alpha = tf.random_uniform(\n",
    "        shape=[S,1,1], \n",
    "        minval=0.,\n",
    "        maxval=1.\n",
    "    )\n",
    "    differences = dihed_tensor - dihed\n",
    "    interpolates = dihed_tensor + (alpha*differences)\n",
    "    print(\"interpolates\")\n",
    "    print(interpolates)\n",
    "    \n",
    "    gradients = tf.gradients(discriminator(DW_list, Db_list, Dh_list, L_list, n_layers, mol_xyz, interpolates), [interpolates])[0]\n",
    "    print(\"Gradients\")\n",
    "    print(gradients)\n",
    "    slopes = tf.sqrt(tf.reduce_sum(tf.square(gradients), reduction_indices=[1,2]))\n",
    "    gradient_penalty = tf.reduce_mean((slopes-1.)**2)\n",
    "    LAMBDA = 10\n",
    "    D_loss += LAMBDA*gradient_penalty\n",
    "    \n",
    "    G_solver = tf.train.AdamOptimizer(learning_rate=1e-3, beta1=0.5, beta2=0.9).minimize(G_loss, var_list=gen_params)\n",
    "    D_solver = tf.train.AdamOptimizer(learning_rate=1e-3, beta1=0.5, beta2=0.9).minimize(D_loss, var_list=disc_params)\n",
    "    \"\"\"\n",
    "\n",
    "    \"\"\"\n",
    "    clip_D = []\n",
    "    for p in theta_D[:-1]:\n",
    "        #print(p)\n",
    "        clip_D.append(p.assign(tf.clip_by_value(p, -0.01, 0.01)))\n",
    "    clip_D.append(lam1.assign(tf.clip_by_value(lam1, 0.001, 0.01)))\n",
    "\n",
    "    \n",
    "    D_solver = (tf.train.RMSPropOptimizer(learning_rate=5e-4)\n",
    "                .minimize(-D_loss, var_list=theta_D))\n",
    "    G_solver = (tf.train.RMSPropOptimizer(learning_rate=5e-4)\n",
    "                .minimize(G_loss, var_list=theta_G))\n",
    "    \"\"\"\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1292,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 1292,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "clip_D = [lam1.assign(tf.clip_by_value(lam1, -100.0, -0.01))]\n",
    "\n",
    "tf.get_collection(tf.GraphKeys.UPDATE_OPS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1293,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def construct_feed_dict(X, start=None,\n",
    "                      stop=None, y=None,\n",
    "                      keep_prob_val=1., train=True,\n",
    "                       shuffle_inds=True):\n",
    "    a = time.time()\n",
    "    if start is None:\n",
    "      start = 0\n",
    "      stop = len(X)\n",
    "    \n",
    "    inds = range(start, stop)\n",
    "    if shuffle_inds:\n",
    "        random.shuffle(inds)\n",
    "\n",
    "    mol_xyz_batch = [X[idx][1] for idx in inds]\n",
    "    dihed_tuples_batch = [X[idx][2] for idx in inds]\n",
    "    dihed_tensor_batch = [X[idx][3] for idx in inds]\n",
    "    subgraph_batch = [X[idx][4] for idx in inds]\n",
    "    atom_index_of_bond_batch = [X[idx][5] for idx in inds]\n",
    "    n_atoms_batch = np.reshape(np.array([X[idx][6] for idx in inds]), (S,1))\n",
    "    \n",
    "    \n",
    "    atom_adj_batch = [X[idx][7][0] for idx in inds]\n",
    "    A_batch = np.array([X[idx][7][1] for idx in inds])\n",
    "\n",
    "    feed_dict = {x: A_batch,\n",
    "                 adj_matrix: atom_adj_batch,\n",
    "                 mol_xyz: mol_xyz_batch,\n",
    "                 dihed_indices: dihed_tuples_batch,\n",
    "                 dihed_tensor: dihed_tensor_batch,\n",
    "                 subgraph: subgraph_batch,\n",
    "                 atom_index_of_bond: atom_index_of_bond_batch,\n",
    "                 keep_prob: keep_prob_val,\n",
    "                 phase: train,\n",
    "                 n_atoms: n_atoms_batch\n",
    "                }\n",
    "    t = time.time()-a\n",
    "    #print(\"Construct feed dict: %f\" %(t))\n",
    "    return(feed_dict)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1294,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def generate_molecule(mol_file, new_coords, new_file):\n",
    "    print(mol_file)\n",
    "    print(new_coords[:4])\n",
    "    if \".pdb\" in mol_file:\n",
    "        mol = Chem.MolFromPDBFile(mol_file)\n",
    "    else:\n",
    "        mol = Chem.MolFromMol2File(mol_file)\n",
    "    c = mol.GetConformer(0)\n",
    "    for i in range(mol.GetNumAtoms()):\n",
    "        c.SetAtomPosition(i, new_coords[i].tolist())\n",
    "    Chem.MolToMolFile(mol, new_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8.7312040329\n",
      "Training epoch 0\n",
      "Iter: 0\n",
      "4wks:\n",
      "\n",
      "generated rotation matrix:\n",
      "[[[  9.9856e-01   5.1446e-02   1.5242e-02]\n",
      "  [ -5.1789e-02   9.9839e-01   2.3014e-02]\n",
      "  [ -1.4033e-02  -2.3770e-02   9.9962e-01]]\n",
      "\n",
      " [[  9.8323e-01  -1.7731e-01  -4.2681e-02]\n",
      "  [  1.7332e-01   9.8129e-01  -8.3851e-02]\n",
      "  [  5.6749e-02   7.5047e-02   9.9556e-01]]\n",
      "\n",
      " [[  6.2604e-01   6.8915e-01   3.6490e-01]\n",
      "  [ -7.7811e-01   5.8279e-01   2.3432e-01]\n",
      "  [ -5.1183e-02  -4.3063e-01   9.0108e-01]]\n",
      "\n",
      " [[  6.1384e-01  -7.8794e-01  -4.8444e-02]\n",
      "  [  6.9607e-01   5.6918e-01  -4.3762e-01]\n",
      "  [  3.7239e-01   2.3491e-01   8.9785e-01]]\n",
      "\n",
      " [[  8.7166e-01   4.5317e-01   1.8668e-01]\n",
      "  [ -4.8371e-01   8.5681e-01   1.7861e-01]\n",
      "  [ -7.9007e-02  -2.4598e-01   9.6605e-01]]\n",
      "\n",
      " [[ -3.3623e-02  -9.8117e-01   1.9018e-01]\n",
      "  [  7.3527e-01  -1.5317e-01  -6.6024e-01]\n",
      "  [  6.7694e-01   1.1764e-01   7.2658e-01]]\n",
      "\n",
      " [[ -3.9121e-01  -8.3135e-01   3.9472e-01]\n",
      "  [  5.0038e-01  -5.5212e-01  -6.6692e-01]\n",
      "  [  7.7238e-01  -6.3394e-02   6.3199e-01]]\n",
      "\n",
      " [[  9.4024e-01  -3.3376e-01  -6.7562e-02]\n",
      "  [  3.1954e-01   9.3332e-01  -1.6372e-01]\n",
      "  [  1.1770e-01   1.3235e-01   9.8419e-01]]\n",
      "\n",
      " [[  9.9899e-01  -4.3336e-02  -1.1832e-02]\n",
      "  [  4.3096e-02   9.9887e-01  -1.9850e-02]\n",
      "  [  1.2679e-02   1.9320e-02   9.9973e-01]]\n",
      "\n",
      " [[  6.1373e-01  -7.8803e-01  -4.8418e-02]\n",
      "  [  6.9614e-01   5.6905e-01  -4.3769e-01]\n",
      "  [  3.7247e-01   2.3492e-01   8.9782e-01]]\n",
      "\n",
      " [[ -1.8906e-01  -9.4376e-01   2.7123e-01]\n",
      "  [  6.6089e-01  -3.2658e-01  -6.7570e-01]\n",
      "  [  7.2628e-01   5.1507e-02   6.8546e-01]]\n",
      "\n",
      " [[  7.5505e-01   5.9323e-01   2.7924e-01]\n",
      "  [ -6.5151e-01   7.2672e-01   2.1776e-01]\n",
      "  [ -7.3747e-02  -3.4634e-01   9.3520e-01]]\n",
      "\n",
      " [[  7.0340e-01   6.3714e-01   3.1510e-01]\n",
      "  [ -7.0770e-01   6.6909e-01   2.2688e-01]\n",
      "  [ -6.6274e-02  -3.8258e-01   9.2154e-01]]\n",
      "\n",
      " [[  9.9163e-01   1.2320e-01   3.8731e-02]\n",
      "  [ -1.2519e-01   9.9066e-01   5.4085e-02]\n",
      "  [ -3.1706e-02  -5.8481e-02   9.9778e-01]]\n",
      "\n",
      " [[ -3.9282e-01  -8.3007e-01   3.9582e-01]\n",
      "  [  4.9872e-01  -5.5391e-01  -6.6668e-01]\n",
      "  [  7.7264e-01  -6.4482e-02   6.3156e-01]]\n",
      "\n",
      " [[  9.4468e-01   3.0812e-01   1.1245e-01]\n",
      "  [ -3.2128e-01   9.3828e-01   1.2810e-01]\n",
      "  [ -6.6038e-02  -1.5714e-01   9.8537e-01]]\n",
      "\n",
      " [[  9.8681e-01   1.5409e-01   4.9675e-02]\n",
      "  [ -1.5723e-01   9.8528e-01   6.7079e-02]\n",
      "  [ -3.8608e-02  -7.4005e-02   9.9651e-01]]\n",
      "\n",
      " [[ -2.1684e-01   6.4363e-01   7.3398e-01]\n",
      "  [ -9.3311e-01  -3.5758e-01   3.7892e-02]\n",
      "  [  2.8684e-01  -6.7667e-01   6.7812e-01]]\n",
      "\n",
      " [[  7.9533e-01   5.5248e-01   2.4942e-01]\n",
      "  [ -6.0117e-01   7.7166e-01   2.0769e-01]\n",
      "  [ -7.7728e-02  -3.1513e-01   9.4586e-01]]\n",
      "\n",
      " [[  6.1373e-01   6.9614e-01   3.7247e-01]\n",
      "  [ -7.8803e-01   5.6905e-01   2.3492e-01]\n",
      "  [ -4.8417e-02  -4.3769e-01   8.9782e-01]]\n",
      "\n",
      " [[  9.9014e-01  -1.3590e-01  -3.4069e-02]\n",
      "  [  1.3355e-01   9.8900e-01  -6.3646e-02]\n",
      "  [  4.2344e-02   5.8468e-02   9.9739e-01]]\n",
      "\n",
      " [[  9.9801e-01   6.0422e-02   1.8036e-02]\n",
      "  [ -6.0896e-02   9.9778e-01   2.6968e-02]\n",
      "  [ -1.6367e-02  -2.8012e-02   9.9947e-01]]\n",
      "\n",
      " [[ -1.8994e-01   6.6036e-01   7.2653e-01]\n",
      "  [ -9.4344e-01  -3.2757e-01   5.1082e-02]\n",
      "  [  2.7173e-01  -6.7574e-01   6.8523e-01]]\n",
      "\n",
      " [[  2.0865e-01   7.8626e-01   5.8160e-01]\n",
      "  [ -9.7452e-01   1.1713e-01   1.9128e-01]\n",
      "  [  8.2272e-02  -6.0669e-01   7.9067e-01]]\n",
      "\n",
      " [[  7.0266e-01   6.3771e-01   3.1559e-01]\n",
      "  [ -7.0844e-01   6.6827e-01   2.2699e-01]\n",
      "  [ -6.6151e-02  -3.8307e-01   9.2135e-01]]\n",
      "\n",
      " [[  8.8701e-01   4.2825e-01   1.7265e-01]\n",
      "  [ -4.5513e-01   8.7395e-01   1.7051e-01]\n",
      "  [ -7.7863e-02  -2.2982e-01   9.7011e-01]]\n",
      "\n",
      " [[  2.7350e-01  -9.6019e-01   5.6952e-02]\n",
      "  [  7.8735e-01   1.8947e-01  -5.8667e-01]\n",
      "  [  5.5252e-01   2.0529e-01   8.0782e-01]]\n",
      "\n",
      " [[  4.4036e-01  -8.9782e-01  -9.8477e-04]\n",
      "  [  7.6468e-01   3.7564e-01  -5.2360e-01]\n",
      "  [  4.7047e-01   2.2982e-01   8.5196e-01]]\n",
      "\n",
      " [[ -5.3008e-01  -6.8558e-01   4.9900e-01]\n",
      "  [  3.2157e-01  -7.0705e-01  -6.2982e-01]\n",
      "  [  7.8461e-01  -1.7339e-01   5.9525e-01]]\n",
      "\n",
      " [[  6.1373e-01  -7.8803e-01  -4.8419e-02]\n",
      "  [  6.9613e-01   5.6906e-01  -4.3769e-01]\n",
      "  [  3.7246e-01   2.3492e-01   8.9782e-01]]\n",
      "\n",
      " [[ -2.1226e-01  -9.3496e-01   2.8424e-01]\n",
      "  [  6.4657e-01  -3.5246e-01  -6.7655e-01]\n",
      "  [  7.3273e-01   4.0179e-02   6.7933e-01]]\n",
      "\n",
      " [[  9.7701e-01   2.0216e-01   6.7751e-02]\n",
      "  [ -2.0763e-01   9.7435e-01   8.6823e-02]\n",
      "  [ -4.8461e-02  -9.8894e-02   9.9392e-01]]]\n",
      "real dist:\n",
      "[[  4.7684e-07   2.3839e+00   7.4443e+00   1.0071e+01]\n",
      " [  2.3839e+00   0.0000e+00   2.7936e+00   7.2998e+00]\n",
      " [  7.4443e+00   2.7936e+00   0.0000e+00   2.0357e+00]\n",
      " [  1.0071e+01   7.2998e+00   2.0357e+00   0.0000e+00]]\n",
      "fake dist:\n",
      "[[  0.       2.3839   7.4443  13.9401]\n",
      " [  2.3839   0.       2.7936   7.2998]\n",
      " [  7.4443   2.7936   0.       2.0357]\n",
      " [ 13.9401   7.2998   2.0357   0.    ]]\n",
      "real rmsd:\n",
      "[ 0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.\n",
      "  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.]\n",
      "fake rmsd:\n",
      "[  1.3081e-02   4.2774e-01   2.7975e-01   6.3396e-02   1.6393e-01\n",
      "   3.2449e-02   4.1757e-01   1.6420e-01   1.6436e-01   1.6434e-01\n",
      "   1.6377e-01   1.6436e-01   3.6051e-01   1.6436e-01   1.2884e-02\n",
      "   5.8729e-02   1.6409e-01   1.8720e-02   1.6437e-01   1.1414e-01\n",
      "   1.3083e-02   8.1062e-06   1.3083e-02   7.5431e-03   8.1000e-03\n",
      "   9.2235e-02   1.6436e-01   1.6357e-01   1.3083e-02   1.3820e-02\n",
      "   1.0641e-01   1.0641e-01]\n",
      "[[  1.8970e+00]\n",
      " [  1.9998e+00]\n",
      " [  1.9973e+00]\n",
      " [ -4.9461e-01]\n",
      " [ -1.0000e+00]\n",
      " [  1.9999e+00]\n",
      " [  1.6658e+00]\n",
      " [  7.5893e-01]\n",
      " [  1.9445e+00]\n",
      " [ -1.0000e+00]\n",
      " [  1.3870e+00]\n",
      " [  1.6449e+00]\n",
      " [  1.9966e+00]\n",
      " [  1.3526e-01]\n",
      " [  1.0518e+00]\n",
      " [  1.9965e+00]\n",
      " [  1.9998e+00]\n",
      " [  1.0000e+00]\n",
      " [  1.6964e+00]\n",
      " [  2.7345e+00]\n",
      " [  1.9983e+00]\n",
      " [  4.6460e-02]\n",
      " [  1.9935e+00]\n",
      " [  1.7839e-02]\n",
      " [  1.0243e+00]\n",
      " [  4.1789e-04]\n",
      " [  8.3958e-01]\n",
      " [  5.1779e-04]\n",
      " [  8.3521e-01]\n",
      " [  1.3375e-01]\n",
      " [  5.7403e-01]\n",
      " [  2.8719e-01]]\n",
      "0.552773\n",
      "1.43061\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXQAAAD8CAYAAABn919SAAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAADBZJREFUeJzt3X+M5PVdx/HnS65EocSCt1IKnIsJIaL/lGwq0qZpChoE\nIzWpCU1a0dRc+geVGhNzptH+i8Y01kT/uLQoRkL/oGhJoVbAksbEXryjUH5cW9p6tuDBXW1C6z/S\npm//2K/Jsb39MfP93s3su89HstmZne/uvD/zvX1m7jszO6kqJEm7348tegBJ0jQMuiQ1YdAlqQmD\nLklNGHRJasKgS1ITBl2SmjDoktSEQZekJvaczSvbu3dvra6uns2rlKRd78iRI9+qqpXttjurQV9d\nXeXw4cNn8yoladdL8p872c5DLpLUhEGXpCYMuiQ1YdAlqQmDLklNbBv0JHclOZHk6VO+dlGSh5M8\nN3y+8MyOKUnazk7uof8tcOOGrx0AHq2qK4FHh/OSpAXaNuhV9Tng2xu+fAtw93D6buAdE88lSZrR\nvMfQL66q48PpF4GLJ5pHkjSn0a8UrapKsuk7TSfZD+wH2Ldv39irk6SFWz3w4Mzfc+zOm8/AJK82\n7z30l5JcAjB8PrHZhlV1sKrWqmptZWXbP0UgSZrTvEF/ALhtOH0b8MlpxpEkzWsnT1u8F/g34Kok\nzyd5L3An8MtJngNuGM5LkhZo22PoVfWuTS66fuJZJEkj+EpRSWrCoEtSEwZdkpow6JLUhEGXpCYM\nuiQ1YdAlqQmDLklNGHRJasKgS1ITBl2SmjDoktSEQZekJgy6JDVh0CWpCYMuSU0YdElqwqBLUhMG\nXZKaMOiS1IRBl6QmDLokNWHQJakJgy5JTRh0SWrCoEtSEwZdkpow6JLUhEGXpCYMuiQ1YdAlqQmD\nLklNjAp6kt9P8kySp5Pcm+THpxpMkjSbuYOe5FLg94C1qvoF4Bzg1qkGkyTNZuwhlz3ATyTZA5wH\n/Nf4kSRJ85g76FX1AvDnwDeA48DLVfXPUw0mSZrNmEMuFwK3AFcAbwDOT/Lu02y3P8nhJIdPnjw5\n/6SSpC2NOeRyA/AfVXWyqr4H3A9ct3GjqjpYVWtVtbaysjLi6iRJWxkT9G8A1yY5L0mA64Gj04wl\nSZrVmGPoh4D7gMeBp4afdXCiuSRJM9oz5pur6kPAhyaaRZI0gq8UlaQmDLokNWHQJakJgy5JTRh0\nSWrCoEtSEwZdkpow6JLUhEGXpCYMuiQ1YdAlqQmDLklNGHRJasKgS1ITBl2SmjDoktSEQZekJgy6\nJDVh0CWpCYMuSU0YdElqwqBLUhMGXZKaMOiS1IRBl6QmDLokNWHQJakJgy5JTRh0SWrCoEtSEwZd\nkpow6JLUhEGXpCZGBT3J65Lcl+RLSY4m+aWpBpMkzWbPyO//CPBPVfXOJOcC500wkyRpDnMHPclP\nAm8Ffhugql4BXplmLEnSrMYccrkCOAn8TZIvJPlokvMnmkuSNKMxh1z2ANcA76+qQ0k+AhwA/vjU\njZLsB/YD7Nu3b8TVSQJYPfDgTNsfu/PmMzSJls2Ye+jPA89X1aHh/H2sB/5VqupgVa1V1drKysqI\nq5MkbWXuoFfVi8A3k1w1fOl64NlJppIkzWzss1zeD9wzPMPl68DvjB9JkjSPUUGvqieAtYlmkSSN\n4CtFJakJgy5JTRh0SWrCoEtSEwZdkpow6JLUhEGXpCYMuiQ1YdAlqQmDLklNGHRJasKgS1ITBl2S\nmjDoktSEQZekJgy6JDVh0CWpCYMuSU0YdElqwqBLUhMGXZKaMOiS1IRBl6QmDLokNWHQJakJgy5J\nTRh0SWrCoEtSEwZdkpow6JLUhEGXpCYMuiQ1MTroSc5J8oUkn5piIEnSfKa4h34HcHSCnyNJGmFU\n0JNcBtwMfHSacSRJ8xp7D/0vgD8EfjDBLJKkEeYOepJfA05U1ZFtttuf5HCSwydPnpz36iRJ2xhz\nD/3NwK8nOQZ8HHh7kr/fuFFVHayqtapaW1lZGXF1kqStzB30qvqjqrqsqlaBW4F/qap3TzaZJGkm\nPg9dkprYM8UPqarHgMem+FmSpPl4D12SmjDoktSEQZekJgy6JDVh0CWpCYMuSU0YdElqwqBLUhMG\nXZKaMOiS1IRBl6QmDLokNWHQJakJgy5JTRh0SWrCoEtSE5O8wYV+dKweeHCm7Y/defMZmkTSRt5D\nl6QmDLokNWHQJakJgy5JTRh0SWrCoEtSEwZdkpow6JLUhEGXpCYMuiQ1YdAlqQmDLklNGHRJasKg\nS1ITBl2SmjDoktTE3EFPcnmSzyZ5NskzSe6YcjBJ0mzGvGPR94E/qKrHk1wAHEnycFU9O9FskqQZ\nzH0PvaqOV9Xjw+nvAkeBS6caTJI0m0neUzTJKvBG4NBpLtsP7AfYt2/f3Nfhe1mqq1n/bUubGf2g\naJLXAp8APlBV39l4eVUdrKq1qlpbWVkZe3WSpE2MCnqS17Ae83uq6v5pRpIkzWPMs1wCfAw4WlUf\nnm4kSdI8xtxDfzPwHuDtSZ4YPm6aaC5J0ozmflC0qv4VyISzSJJG8JWiktSEQZekJgy6JDVh0CWp\nCYMuSU0YdElqwqBLUhMGXZKaMOiS1IRBl6QmDLokNWHQJakJgy5JTRh0SWrCoEtSE5O8SbS0KMv4\nBsvL9gblZ/o2mme9y7jfOvAeuiQ1YdAlqQmDLklNGHRJasKgS1ITBl2SmjDoktSEQZekJgy6JDVh\n0CWpCYMuSU0YdElqwqBLUhMGXZKaMOiS1IRBl6QmRgU9yY1Jvpzkq0kOTDWUJGl2cwc9yTnAXwG/\nClwNvCvJ1VMNJkmazZh76G8CvlpVX6+qV4CPA7dMM5YkaVZjgn4p8M1Tzj8/fE2StACpqvm+MXkn\ncGNV/e5w/j3AL1bV7Ru22w/sH85eBXx5zln3At+a83uXSYd1uIbl0WEdHdYAZ3YdP1NVK9tttGfE\nFbwAXH7K+cuGr71KVR0EDo64HgCSHK6qtbE/Z9E6rMM1LI8O6+iwBliOdYw55PLvwJVJrkhyLnAr\n8MA0Y0mSZjX3PfSq+n6S24HPAOcAd1XVM5NNJkmayZhDLlTVQ8BDE82yndGHbZZEh3W4huXRYR0d\n1gBLsI65HxSVJC0XX/ovSU0sbdCT/GaSZ5L8IMmmjxwv858fSHJRkoeTPDd8vnCT7Y4leSrJE0kO\nn+05N7PdbZt1fzlc/sUk1yxizq3sYA1vS/LycNs/keRPFjHnVpLcleREkqc3uXw37Ift1rAb9sPl\nST6b5NmhTXecZpvF7ouqWsoP4OdYf976Y8DaJtucA3wN+FngXOBJ4OpFz37KfH8GHBhOHwD+dJPt\njgF7Fz3vrLctcBPwaSDAtcChRc89xxreBnxq0bNus463AtcAT29y+VLvhx2uYTfsh0uAa4bTFwBf\nWbbfiaW9h15VR6tquxchLfufH7gFuHs4fTfwjgXOMqud3La3AH9X6z4PvC7JJWd70C0s+7+PHamq\nzwHf3mKTZd8PO1nD0quq41X1+HD6u8BRfvjV8QvdF0sb9B1a9j8/cHFVHR9OvwhcvMl2BTyS5Mjw\nytplsJPbdtlv/53Od93w3+NPJ/n5szPapJZ9P+zUrtkPSVaBNwKHNly00H0x6mmLYyV5BHj9aS76\nYFV98mzPM4+t1nDqmaqqJJs9pegtVfVCkp8GHk7ypeEejc68x4F9VfU/SW4C/hG4csEz/SjaNfsh\nyWuBTwAfqKrvLHqeUy006FV1w8gfsaM/P3AmbbWGJC8luaSqjg//7Tqxyc94Yfh8Isk/sH6oYNFB\n38ltu/DbfxvbznfqL2RVPZTkr5Psrard9LdFln0/bGu37Ickr2E95vdU1f2n2WSh+2K3H3JZ9j8/\n8ABw23D6NuCH/teR5PwkF/z/aeBXgNM+E+As28lt+wDwW8Mj+9cCL59yiGkZbLuGJK9PkuH0m1j/\nnfjvsz7pOMu+H7a1G/bDMN/HgKNV9eFNNlvsvlj0I8dbPKL8G6wff/pf4CXgM8PX3wA8tOFR5a+w\n/myGDy567g1r+CngUeA54BHgoo1rYP0ZGE8OH88s0xpOd9sC7wPeN5wO629y8jXgKTZ5NtKSr+H2\n4XZ/Evg8cN2iZz7NGu4FjgPfG34n3rsL98N2a9gN++EtrD/e9UXgieHjpmXaF75SVJKa2O2HXCRJ\nA4MuSU0YdElqwqBLUhMGXZKaMOiS1IRBl6QmDLokNfF/bZEKMx4fh+UAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<matplotlib.figure.Figure at 0x7f99eb06ecd0>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/evan/Documents/deep_docking/datasets/v2015/4wks/4wks_ligand.mol2\n",
      "[[ 0.152   0.032  -2.724 ]\n",
      " [-0.668   0.418  -1.474 ]\n",
      " [ 0.      0.      0.    ]\n",
      " [ 1.2975  0.4563  0.3794]]\n",
      "4a9n:\n",
      "\n",
      "real:\n",
      "[[ 1.514   1.9077 -1.4783]\n",
      " [ 0.174   1.7917 -1.9773]\n",
      " [-0.051   1.5367 -3.3643]\n",
      " [-0.589   3.1577 -1.6003]]\n",
      "fake:\n",
      "[[ 1.6084  2.9833 -3.7863]\n",
      " [ 0.2684  2.8673 -4.2853]\n",
      " [ 0.0434  2.6123 -5.6723]\n",
      " [-0.4946  4.2333 -3.9083]]\n",
      "fake:\n",
      "[[ 0.9447  0.0914 -0.3151]\n",
      " [-0.2538 -0.4051 -0.8784]\n",
      " [-0.2079  0.9097 -0.3595]]\n",
      "real dist:\n",
      "[[  0.0000e+00   2.0581e+00   6.1439e+00   6.0000e+00]\n",
      " [  2.0581e+00   0.0000e+00   2.0394e+00   2.5903e+00]\n",
      " [  6.1439e+00   2.0394e+00   0.0000e+00   6.0288e+00]\n",
      " [  6.0000e+00   2.5903e+00   6.0288e+00  -1.9073e-06]]\n",
      "fake dist:\n",
      "[[  0.0000e+00   2.0581e+00   6.1439e+00   6.0000e+00]\n",
      " [  2.0581e+00  -3.8147e-06   2.0394e+00   2.5903e+00]\n",
      " [  6.1439e+00   2.0394e+00   0.0000e+00   6.0288e+00]\n",
      " [  6.0000e+00   2.5903e+00   6.0288e+00  -7.6294e-06]]\n",
      "real rmsd:\n",
      "[ 0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.\n",
      "  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.  0.]\n",
      "fake rmsd:\n",
      "[ 3.5072  2.515   0.781   3.9763  5.8625  3.5132  1.3755  3.2362  5.1244\n",
      "  4.2222  4.0652  1.8265  3.1932  3.3487  3.4099  6.0234  0.6497  6.2811\n",
      "  4.9466  4.5248  6.104   4.2677  2.0376  1.7002  0.754   5.7073  1.5416\n",
      "  6.0126  4.3017  5.3939  4.722   4.3398]\n",
      "-0.912921\n",
      "-0.865252\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXcAAAD8CAYAAACMwORRAAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAD3hJREFUeJzt3W+IXneZxvHvZYyrrEJeZKDZJLMjbFiwUmwZYqn7Iuuu\nbJoWwy5dqKBluy+GlAoVBKkWWmQRuggiNdIhbIuWLbpC3W5oU7TLVtq+SDXJxtg0dRmk0pSwrRUT\nQ4sSvffFHHT2cZLnzMwzeWZ+/X7gIefPPefcJzO5OPnN+ZOqQpLUlreNuwFJ0ugZ7pLUIMNdkhpk\nuEtSgwx3SWqQ4S5JDTLcJalBhrskNchwl6QGvX1cO968eXNNTU2Na/eStC4dPXr0Z1U1MaxubOE+\nNTXFkSNHxrV7SVqXkvy0T53DMpLUIMNdkhpkuEtSgwx3SWqQ4S5JDeod7kk2JPnvJI8tsi5J7ksy\nl+REkmtG26YkaSmWcuZ+B3DqIuuuB3Z0nxng/hX2JUlagV7hnmQbcAPwLxcp2Qs8VPMOA5uSbBlR\nj5KkJep75v5l4DPAby+yfivw8oL5090ySdIYDL1DNcmNwKtVdTTJrpXsLMkM88M2TE5OrmRTktao\nqTsfX1L9S/fesEqdvLX1OXP/EPDRJC8B3wQ+nORfB2peAbYvmN/WLft/qupAVU1X1fTExNBHI0iS\nlmlouFfVZ6tqW1VNATcD/1VVHx8oOwjc0l01cy1wtqrOjL5dSVIfy35wWJJ9AFU1CxwC9gBzwBvA\nrSPpTpK0LEsK96r6HvC9bnp2wfICbh9lY5Kk5fMOVUlqkOEuSQ0y3CWpQYa7JDXIcJekBhnuktQg\nw12SGmS4S1KDDHdJapDhLkkNMtwlqUGGuyQ1yHCXpAYZ7pLUIMNdkhpkuEtSgwx3SWrQ0HBP8s4k\n30/ywyQnk3x+kZpdSc4mOd597l6ddiVJffR5zd6vgA9X1fkkG4FnkzxRVYcH6p6pqhtH36IkaamG\nhnv3ftTz3ezG7lOr2ZQkaWV6jbkn2ZDkOPAq8GRVPbdI2XVJTiR5IsmVI+1SkrQkvcK9qn5TVR8A\ntgE7k7x/oOQYMFlVVwFfAR5dbDtJZpIcSXLktddeW0nfkqRLWNLVMlX1C+ApYPfA8nNVdb6bPgRs\nTLJ5ka8/UFXTVTU9MTGxgrYlSZfS52qZiSSbuul3AR8BXhyouSJJuumd3XZfH327kqQ++lwtswX4\nepINzIf2t6rqsST7AKpqFrgJuC3JBeBN4ObuF7GSpDHoc7XMCeDqRZbPLpjeD+wfbWuSpOXyDlVJ\napDhLkkNMtwlqUGGuyQ1yHCXpAYZ7pLUIMNdkhpkuEtSgwx3SWqQ4S5JDTLcJalBhrskNchwl6QG\nGe6S1CDDXZIaZLhLUoMMd0lqUJ93qL4zyfeT/DDJySSfX6QmSe5LMpfkRJJrVqddSVIffd6h+ivg\nw1V1PslG4NkkT1TV4QU11wM7us8Hgfu7PyVJYzD0zL3mne9mN3afwZdf7wUe6moPA5uSbBltq5Kk\nvvqcuZNkA3AU+DPgq1X13EDJVuDlBfOnu2VnBrYzA8wATE5OLrPl5Zm68/HLur9hXrr3hnG3IKlh\nvX6hWlW/qaoPANuAnUnev5ydVdWBqpququmJiYnlbEKS1MOSrpapql8ATwG7B1a9AmxfML+tWyZJ\nGoM+V8tMJNnUTb8L+Ajw4kDZQeCW7qqZa4GzVXUGSdJY9Blz3wJ8vRt3fxvwrap6LMk+gKqaBQ4B\ne4A54A3g1lXqV5LUw9Bwr6oTwNWLLJ9dMF3A7aNtTZK0XN6hKkkNMtwlqUGGuyQ1yHCXpAYZ7pLU\nIMNdkhpkuEtSgwx3SWqQ4S5JDTLcJalBhrskNchwl6QGGe6S1CDDXZIaZLhLUoMMd0lqkOEuSQ3q\n8w7V7UmeSvJCkpNJ7likZleSs0mOd5+7V6ddSVIffd6hegH4dFUdS/Ie4GiSJ6vqhYG6Z6rqxtG3\nKElaqqFn7lV1pqqOddO/BE4BW1e7MUnS8i1pzD3JFPMvy35ukdXXJTmR5IkkV17k62eSHEly5LXX\nXltys5KkfnqHe5J3A48An6qqcwOrjwGTVXUV8BXg0cW2UVUHqmq6qqYnJiaW27MkaYhe4Z5kI/PB\n/nBVfXtwfVWdq6rz3fQhYGOSzSPtVJLUW5+rZQI8AJyqqi9dpOaKro4kO7vtvj7KRiVJ/fW5WuZD\nwCeAHyU53i37HDAJUFWzwE3AbUkuAG8CN1dVrUK/kqQehoZ7VT0LZEjNfmD/qJqSJK2Md6hKUoMM\nd0lqkOEuSQ0y3CWpQYa7JDXIcJekBhnuktQgw12SGmS4S1KDDHdJapDhLkkNMtwlqUGGuyQ1yHCX\npAYZ7pLUIMNdkhpkuEtSg/q8Q3V7kqeSvJDkZJI7FqlJkvuSzCU5keSa1WlXktRHn3eoXgA+XVXH\nkrwHOJrkyap6YUHN9cCO7vNB4P7uT0nSGAw9c6+qM1V1rJv+JXAK2DpQthd4qOYdBjYl2TLybiVJ\nvfQ5c/+dJFPA1cBzA6u2Ai8vmD/dLTsz8PUzwAzA5OTk0jrVW87UnY+Pu4UVe+neG8bdgpZoOT93\na/H73PsXqkneDTwCfKqqzi1nZ1V1oKqmq2p6YmJiOZuQJPXQK9yTbGQ+2B+uqm8vUvIKsH3B/LZu\nmSRpDPpcLRPgAeBUVX3pImUHgVu6q2auBc5W1ZmL1EqSVlmfMfcPAZ8AfpTkeLfsc8AkQFXNAoeA\nPcAc8AZw6+hblST1NTTcq+pZIENqCrh9VE1JklbGO1QlqUGGuyQ1yHCXpAYZ7pLUIMNdkhpkuEtS\ngwx3SWqQ4S5JDTLcJalBhrskNchwl6QGGe6S1CDDXZIaZLhLUoMMd0lqkOEuSQ3q85q9B5O8muT5\ni6zfleRskuPd5+7RtylJWoo+r9n7GrAfeOgSNc9U1Y0j6UiStGJDz9yr6mng55ehF0nSiIxqzP26\nJCeSPJHkyhFtU5K0TH2GZYY5BkxW1fkke4BHgR2LFSaZAWYAJicnR7BrSdJiVnzmXlXnqup8N30I\n2Jhk80VqD1TVdFVNT0xMrHTXkqSLWHG4J7kiSbrpnd02X1/pdiVJyzd0WCbJN4BdwOYkp4F7gI0A\nVTUL3ATcluQC8CZwc1XVqnUsSRpqaLhX1ceGrN/P/KWSkqQ1wjtUJalBhrskNchwl6QGGe6S1CDD\nXZIaZLhLUoMMd0lqkOEuSQ0y3CWpQYa7JDXIcJekBhnuktQgw12SGmS4S1KDDHdJapDhLkkNMtwl\nqUFDwz3Jg0leTfL8RdYnyX1J5pKcSHLN6NuUJC1FnzP3rwG7L7H+emBH95kB7l95W5KklRga7lX1\nNPDzS5TsBR6qeYeBTUm2jKpBSdLSjWLMfSvw8oL5090ySdKYvP1y7izJDPNDN0xOTi57O1N3Pj6q\nlsamhWNYqpfuvWHcLVx2S/0+vxX/jlbb5fi3tha/z6M4c38F2L5gflu37A9U1YGqmq6q6YmJiRHs\nWpK0mFGE+0Hglu6qmWuBs1V1ZgTblSQt09BhmSTfAHYBm5OcBu4BNgJU1SxwCNgDzAFvALeuVrOS\npH6GhntVfWzI+gJuH1lHkqQV8w5VSWqQ4S5JDTLcJalBhrskNchwl6QGGe6S1CDDXZIaZLhLUoMM\nd0lqkOEuSQ0y3CWpQYa7JDXIcJekBhnuktQgw12SGmS4S1KDDHdJalCvcE+yO8mPk8wluXOR9buS\nnE1yvPvcPfpWJUl99XmH6gbgq8BHgNPAD5IcrKoXBkqfqaobV6FHSdIS9Tlz3wnMVdVPqurXwDeB\nvavbliRpJfqE+1bg5QXzp7tlg65LciLJE0muHEl3kqRlGTos09MxYLKqzifZAzwK7BgsSjIDzABM\nTk6OaNeSpEF9ztxfAbYvmN/WLfudqjpXVee76UPAxiSbBzdUVQeqarqqpicmJlbQtiTpUvqE+w+A\nHUnem+QdwM3AwYUFSa5Ikm56Z7fd10fdrCSpn6HDMlV1Icknge8AG4AHq+pkkn3d+lngJuC2JBeA\nN4Gbq6pWsW9J0iX0GnPvhloODSybXTC9H9g/2tYkScvlHaqS1CDDXZIaZLhLUoMMd0lqkOEuSQ0y\n3CWpQYa7JDXIcJekBhnuktQgw12SGmS4S1KDDHdJapDhLkkNMtwlqUGGuyQ1yHCXpAYZ7pLUoF7h\nnmR3kh8nmUty5yLrk+S+bv2JJNeMvlVJUl9Dwz3JBuCrwPXA+4CPJXnfQNn1wI7uMwPcP+I+JUlL\n0OfMfScwV1U/qapfA98E9g7U7AUeqnmHgU1Jtoy4V0lST33CfSvw8oL5092ypdZIki6Tt1/OnSWZ\nYX7YBuB8kh9fzv33sBn42bibWAVr4rjyzyPf5Jo4rlHq/o6aO67Oose1Cj8X47Ck79kKj/lP+xT1\nCfdXgO0L5rd1y5ZaQ1UdAA70aWwckhypqulx9zFqHtf64nGtP2vx2PoMy/wA2JHkvUneAdwMHByo\nOQjc0l01cy1wtqrOjLhXSVJPQ8/cq+pCkk8C3wE2AA9W1ckk+7r1s8AhYA8wB7wB3Lp6LUuShuk1\n5l5Vh5gP8IXLZhdMF3D7aFsbizU7ZLRCHtf64nGtP2vu2DKfy5Kklvj4AUlqkOE+IMk/dY9QOJ7k\nu0n+ZNw9jUKSLyZ5sTu2f0+yadw9jUKSv09yMslvk6ypqxWWY9ijPtajJA8meTXJ8+PuZZSSbE/y\nVJIXup/BO8bd00KG+x/6YlVdVVUfAB4D7h53QyPyJPD+qroK+B/gs2PuZ1SeB/4OeHrcjaxUz0d9\nrEdfA3aPu4lVcAH4dFW9D7gWuH0tfb8M9wFVdW7B7B8DTfxSoqq+W1UXutnDzN+LsO5V1amqWms3\nwy1Xn0d9rDtV9TTw83H3MWpVdaaqjnXTvwROsYbuzL+sd6iuF0m+ANwCnAX+csztrIZ/BP5t3E3o\nDyz2GI8PjqkXLUGSKeBq4LnxdvJ7b8lwT/KfwBWLrLqrqv6jqu4C7kryWeCTwD2XtcFlGnZcXc1d\nzP938uHL2dtK9DkuaVySvBt4BPjUwP/8x+otGe5V9dc9Sx9m/vr+dRHuw44ryT8ANwJ/VevoGtgl\nfL/Wu16P8dDakWQj88H+cFV9e9z9LOSY+4AkOxbM7gVeHFcvo5RkN/AZ4KNV9ca4+9Gi+jzqQ2tE\nkgAPAKeq6kvj7meQNzENSPII8OfAb4GfAvuqat2fPSWZA/4IeL1bdLiq9o2xpZFI8rfAV4AJ4BfA\n8ar6m/F2tXxJ9gBf5veP+vjCmFtasSTfAHYx/+TE/wXuqaoHxtrUCCT5C+AZ4EfM5wXA57o7+sfO\ncJekBjksI0kNMtwlqUGGuyQ1yHCXpAYZ7pLUIMNdkhpkuEtSgwx3SWrQ/wE3tTtmmOBmFgAAAABJ\nRU5ErkJggg==\n",
      "text/plain": [
       "<matplotlib.figure.Figure at 0x7f99f79a5290>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "lam1:\n",
      "-0.506502\n",
      "/home/evan/Documents/deep_docking/datasets/v2015/4a9n/4a9n_ligand.mol2\n",
      "[[ 5.4512  3.0005 -0.6985]\n",
      " [ 4.1112  2.8845 -1.1975]\n",
      " [ 3.8862  2.6295 -2.5845]\n",
      " [ 3.3482  4.2505 -0.8205]]\n",
      "Butane:\n",
      "\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAW4AAAD8CAYAAABXe05zAAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAACrNJREFUeJzt3FGIZnd5x/HfYzaiVakXGVoxbkdoCQTBWIZUsQhNa4mm\ntLRYUFBoscyNioJQUkovvMuV1AtvFrUWahVRQ0tilUgjQbCxuxolySqIpBixZEXE5EaJPr3YCWbS\nmczZMGfefXY/H3jZeWf/e+aZMztfDuec963uDgBzPG/TAwBwaYQbYBjhBhhGuAGGEW6AYYQbYBjh\nBhhGuAGGEW6AYU6tsdHrrruut7e319g0wBXp3LlzP+rurSVrVwn39vZ2zp49u8amAa5IVfU/S9c6\nVQIwjHADDCPcAMMIN8Awwg0wzKJwV9VLq+ozVfXtqjpfVa9bezAADrb0dsAPJflCd7+lqp6f5NdW\nnAmAZ3FkuKvq15O8IclfJUl3/zzJz9cdC4DDLDlV8sokF5L8U1V9o6o+UlUvWnkuAA6x5FTJqSS/\nm+Q93X1/VX0oye1J/uHpi6pqN8lukpw+ffq45+Qqt3373Ze0/pE7bltpEti8JUfcjyZ5tLvv33v+\nmVwM+T7dfaa7d7p7Z2tr0cvtAXgOjgx3d/9vku9X1Q17n/rDJA+vOhUAh1p6V8l7knxi746S7yX5\n6/VGAuDZLAp3dz+QZGflWQBYwCsnAYYRboBhhBtgGOEGGEa4AYYRboBhhBtgGOEGGEa4AYYRboBh\nhBtgGOEGGEa4AYYRboBhhBtgGOEGGEa4AYYRboBhhBtgGOEGGEa4AYYRboBhhBtgGOEGGEa4AYYR\nboBhTi1ZVFWPJHk8yS+SPNndO2sOBcDhFoV7zx90949WmwSARZwqARhmabg7yZeq6lxV7R60oKp2\nq+psVZ29cOHC8U0IwD5Lw/373X1TkjcleVdVveGZC7r7THfvdPfO1tbWsQ4JwK8sCnd3/2Dvz8eS\n3Jnk5jWHAuBwR4a7ql5UVS956uMkf5zkwbUHA+BgS+4q+Y0kd1bVU+v/tbu/sOpUABzqyHB39/eS\nvPoEZgFgAbcDAgwj3ADDCDfAMMINMIxwAwwj3ADDCDfAMMINMIxwAwwj3ADDCDfAMMINMIxwAwwj\n3ADDCDfAMMINMIxwAwwj3ADDCDfAMMINMIxwAwwj3ADDCDfAMMINMIxwAwyzONxVdU1VfaOq7lpz\nIACe3aUccb83yfm1BgFgmUXhrqrrk9yW5CPrjgPAUZYecf9jkr9N8ssVZwFggSPDXVV/kuSx7j53\nxLrdqjpbVWcvXLhwbAMCsN+SI+7XJ/nTqnokyaeS3FJV//LMRd19prt3untna2vrmMcE4ClHhru7\n/667r+/u7SRvTfKf3f321ScD4EDu4wYY5tSlLO7uLyf58iqTALCII26AYYQbYBjhBhhGuAGGEW6A\nYYQbYBjhBhhGuAGGEW6AYYQbYBjhBhhGuAGGEW6AYYQbYBjhBhhGuAGGEW6AYYQbYBjhBhhGuAGG\nEW6AYYQbYBjhBhhGuAGGEW6AYYQbYJgjw11VL6iqr1XVN6vqoar6wEkMBsDBTi1Y87Mkt3T3E1V1\nbZKvVNV/dPd/rTwbAAc4Mtzd3Ume2Ht67d6j1xwKgMMtOeJOVV2T5FyS307y4e6+/4A1u0l2k+T0\n6dPHOSOXue3b777kf/PIHbetMMnJudTv+SS+38txJtax6OJkd/+iu29Kcn2Sm6vqVQesOdPdO929\ns7W1ddxzArDnku4q6e6fJLk3ya3rjAPAUZbcVbJVVS/d+/iFSd6Y5NtrDwbAwZac435Zkn/eO8/9\nvCSf7u671h0LgMMsuavkW0lecwKzALCAV04CDCPcAMMIN8Awwg0wjHADDCPcAMMIN8Awwg0wjHAD\nDCPcAMMIN8Awwg0wjHADDCPcAMMIN8Awwg0wjHADDCPcAMMIN8Awwg0wjHADDCPcAMMIN8Awwg0w\njHADDHNkuKvqFVV1b1U9XFUPVdV7T2IwAA52asGaJ5O8v7u/XlUvSXKuqu7p7odXng2AAxx5xN3d\nP+zur+99/HiS80levvZgABzsks5xV9V2ktckuX+NYQA42pJTJUmSqnpxks8meV93//SAv99Nspsk\np0+fPrYBOXnbt989/mtc6vYfueO2lSa56CT2KVePRUfcVXVtLkb7E939uYPWdPeZ7t7p7p2tra3j\nnBGAp1lyV0kl+WiS8939wfVHAuDZLDnifn2SdyS5paoe2Hu8eeW5ADjEkee4u/srSeoEZgFgAa+c\nBBhGuAGGEW6AYYQbYBjhBhhGuAGGEW6AYYQbYBjhBhhGuAGGEW6AYYQbYBjhBhhGuAGGEW6AYYQb\nYBjhBhhGuAGGEW6AYYQbYBjhBhhGuAGGEW6AYYQbYBjhBhhGuAGGOTLcVfWxqnqsqh48iYEAeHZL\njrg/nuTWlecAYKEjw93d9yX58QnMAsACp45rQ1W1m2Q3SU6fPv2ct7N9+92XtP6RO257zl9rictt\nHtZxqT/nq9Havwsn8TNYe6aT+v0/touT3X2mu3e6e2dra+u4NgvAM7irBGAY4QYYZsntgJ9M8tUk\nN1TVo1X1zvXHAuAwR16c7O63ncQgACzjVAnAMMINMIxwAwwj3ADDCDfAMMINMIxwAwwj3ADDCDfA\nMMINMIxwAwwj3ADDCDfAMMINMIxwAwwj3ADDCDfAMMINMIxwAwwj3ADDCDfAMMINMIxwAwwj3ADD\nCDfAMIvCXVW3VtV3quq7VXX72kMBcLgjw11V1yT5cJI3Jbkxyduq6sa1BwPgYEuOuG9O8t3u/l53\n/zzJp5L82bpjAXCYJeF+eZLvP+35o3ufA2ADqruffUHVW5Lc2t1/s/f8HUl+r7vf/Yx1u0l2957e\nkOQ7xzzrdUl+dMzbnMq+2M/+2M/+2G/K/vit7t5asvDUgjU/SPKKpz2/fu9z+3T3mSRnFo33HFTV\n2e7eWWv7k9gX+9kf+9kf+12J+2PJqZL/TvI7VfXKqnp+krcm+fd1xwLgMEcecXf3k1X17iRfTHJN\nko9190OrTwbAgZacKkl3fz7J51ee5SirnYYZyL7Yz/7Yz/7Y74rbH0denATg8uIl7wDDjAl3Vf1l\nVT1UVb+sqivqCvGl8PYDv1JVH6uqx6rqwU3PcjmoqldU1b1V9fDe78p7Nz3TJlXVC6rqa1X1zb39\n8YFNz3RcxoQ7yYNJ/iLJfZseZFO8/cD/8/Ekt256iMvIk0ne3903Jnltkndd5f8/fpbklu5+dZKb\nktxaVa/d8EzHYky4u/t8dx/3i3qm8fYDT9Pd9yX58abnuFx09w+7++t7Hz+e5Hyu4lc590VP7D29\ndu9xRVzUGxNuknj7ARaqqu0kr0ly/2Yn2ayquqaqHkjyWJJ7uvuK2B+Lbgc8KVX1pSS/ecBf/X13\n/9tJzwMTVdWLk3w2yfu6+6ebnmeTuvsXSW6qqpcmubOqXtXd46+JXFbh7u4/2vQMl7lFbz/A1auq\nrs3FaH+iuz+36XkuF939k6q6NxeviYwPt1Mls3j7AQ5VVZXko0nOd/cHNz3PplXV1t6RdqrqhUne\nmOTbm53qeIwJd1X9eVU9muR1Se6uqi9ueqaT1t1PJnnq7QfOJ/n01fz2A1X1ySRfTXJDVT1aVe/c\n9Ewb9vok70hyS1U9sPd486aH2qCXJbm3qr6Viwc993T3XRue6Vh45STAMGOOuAG4SLgBhhFugGGE\nG2AY4QYYRrgBhhFugGGEG2CY/wPKUUzgMwRwDgAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<matplotlib.figure.Figure at 0x7f99eeec1750>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "every = 1000\n",
    "train_dataset = features[:(len(features)-S+10)]\n",
    "#train_dataset = features[0:8] + features[9:26] + features[29:36]\n",
    "#train_dataset = features[:96*2]\n",
    "n_train = len(train_dataset)\n",
    "\n",
    "for it in range(1000000):\n",
    "    a = time.time()\n",
    "    random.shuffle(train_dataset)\n",
    "    \n",
    "    batch_sched = list(range(0, n_train+1,S))\n",
    "    for j in range(0, len(batch_sched)-1):\n",
    "        start = batch_sched[j]\n",
    "        stop = batch_sched[j+1]\n",
    "        \n",
    "\n",
    "        feed_dict = construct_feed_dict(train_dataset, start, stop)\n",
    "\n",
    "        _, D_loss_curr = sess.run(\n",
    "            [D_solver, D_loss, clip_D] + d_ops, feed_dict=feed_dict)[0:2]\n",
    "        \n",
    "        _, G_loss_curr = sess.run([G_solver, G_loss] + g_ops, feed_dict=feed_dict)[0:2]\n",
    "                \n",
    "        #d_losses.append(D_loss_curr)\n",
    "        #g_losses.append(G_loss_curr)\n",
    "    \n",
    "    if it % every == 0:\n",
    "        \n",
    "        print(time.time()-a)\n",
    "        print(\"Training epoch %d\" %it)\n",
    "                \n",
    "        print('Iter: {}'.format(it))\n",
    "        #print('D loss: {:.4}'. format(D_loss_curr))\n",
    "        #print('G_loss: {:.4}'.format(G_loss_curr))\n",
    "        feed_dict = construct_feed_dict(features[3:4]*S, 0, S, shuffle_inds=False, keep_prob_val=1., train=False)\n",
    "        \n",
    "        print(\"4wks:\\n\")\n",
    "        \n",
    "        \"\"\"\n",
    "        print(\"real:\")\n",
    "        print(feed_dict[mol_xyz][0])\n",
    "        print(\"fake:\")\n",
    "        print(sess.run(G_sample, feed_dict=feed_dict)[0])\n",
    "        \n",
    "        \n",
    "        print(\"rot_points:\")\n",
    "        print(sess.run(rot_points, feed_dict=feed_dict))\n",
    "\n",
    "        print(\"rot_vecs:\")\n",
    "        print(sess.run(rot_vecs, feed_dict=feed_dict))\n",
    "        \n",
    "        print(\"mol_gen_real:\")\n",
    "        print(sess.run(mol_gen_real, feed_dict=feed_dict))\n",
    "        print(\"mol_gen_fake:\")\n",
    "        print(sess.run(mol_gen_fake, feed_dict=feed_dict))\n",
    "        \"\"\"\n",
    "        \n",
    "\n",
    "        print(\"generated rotation matrix:\")\n",
    "        print(sess.run(rot_mat_i, feed_dict=feed_dict))\n",
    "        \n",
    "        print(\"real dist:\")\n",
    "        print(sess.run(dist_fake_real, feed_dict=feed_dict)[0][:4, :4])\n",
    "        print(\"fake dist:\")\n",
    "        print(sess.run(dist_fake_fake, feed_dict=feed_dict)[0][:4, :4])\n",
    "        \n",
    "        print(\"real rmsd:\")\n",
    "        print(sess.run(rmsds_real, feed_dict=feed_dict))\n",
    "        print(\"fake rmsd:\")\n",
    "        print(sess.run(rmsds_fake, feed_dict=feed_dict))\n",
    "        print(sess.run(dihed, feed_dict=feed_dict)[:, 0])\n",
    "\n",
    "        print(np.mean(sess.run(dihed_tensor, feed_dict=feed_dict)[:, 0]))\n",
    "        print(np.mean(sess.run(dihed, feed_dict=feed_dict)[:, 0]))\n",
    "        \n",
    "        plt.clf()\n",
    "        plt.hist(sess.run(dihed, feed_dict=feed_dict)[:, 0], bins=25)\n",
    "        plt.show()\n",
    "        \n",
    "        generate_molecule(features[3][0], sess.run(G_sample, feed_dict=feed_dict)[0], \"/home/evan/Documents/deep_docking/dcgan_pg/DCGAN-tensorflow/generated_mols/%s_test%d.mol\" %(features[3][0].split(\".\")[0].split(\"/\")[-1], it))\n",
    "\n",
    "        feed_dict = construct_feed_dict(features[0:1]*S, 0, S, shuffle_inds=False, keep_prob_val=1., train=False)\n",
    "        print(\"4a9n:\\n\")\n",
    "        \n",
    "        print(\"real:\")\n",
    "        print(feed_dict[mol_xyz][0][:4])\n",
    "        print(\"fake:\")\n",
    "        print(sess.run(G_sample, feed_dict=feed_dict)[0][:4])\n",
    "        print(\"fake:\")\n",
    "        print(sess.run(rot_mat_i, feed_dict=feed_dict)[0])\n",
    "        \n",
    "        print(\"real dist:\")\n",
    "        print(sess.run(dist_fake_real, feed_dict=feed_dict)[0][:4, :4])\n",
    "        print(\"fake dist:\")\n",
    "        print(sess.run(dist_fake_fake, feed_dict=feed_dict)[0][:4, :4])\n",
    "        \n",
    "        print(\"real rmsd:\")\n",
    "        print(sess.run(rmsds_real, feed_dict=feed_dict))\n",
    "        print(\"fake rmsd:\")\n",
    "        print(sess.run(rmsds_fake, feed_dict=feed_dict))\n",
    "    \n",
    "        print(np.mean(sess.run(dihed_tensor, feed_dict=feed_dict)[:, 0]))\n",
    "        print(np.mean(sess.run(dihed, feed_dict=feed_dict)[:, 0]))\n",
    "        \n",
    "        plt.clf()\n",
    "        plt.hist(sess.run(dihed, feed_dict=feed_dict)[:, 0], bins=25)\n",
    "        plt.show()\n",
    "        \n",
    "        print(\"lam1:\")\n",
    "        print(sess.run(lam1, feed_dict=feed_dict))\n",
    "        \n",
    "        generate_molecule(features[0][0], sess.run(G_sample, feed_dict=feed_dict)[0], \"/home/evan/Documents/deep_docking/dcgan_pg/DCGAN-tensorflow/generated_mols/%s_test%d.mol\" %(features[0][0].split(\".\")[0].split(\"/\")[-1], it))\n",
    "        print(\"Butane:\\n\")\n",
    "        feed_dict = construct_feed_dict(features[len(features)-S:], 0, S, shuffle_inds=False, keep_prob_val=1., train=False)\n",
    "        plt.clf()\n",
    "        plt.hist(sess.run(dihed, feed_dict=feed_dict)[:, 0], bins=25)\n",
    "        plt.show()\n",
    "        #generate_molecule(features[1][0], sess.run(G_sample, feed_dict=feed_dict)[1], \"/home/evan/Documents/deep_docking/dcgan_pg/DCGAN-tensorflow/generated_mols/%s_test%d.mol\" %(features[1][0].split(\".\")[0].split(\"/\")[-1], it))\n",
    "\n",
    "        #feed_dict = construct_feed_dict(features, len(features)-S, len(features), shuffle_inds=False, keep_prob_val=1.)\n",
    "        #generate_molecule(features[-1][0], sess.run(G_sample, feed_dict=feed_dict)[0], \"/home/evan/Documents/deep_docking/dcgan_pg/DCGAN-tensorflow/generated_mols/%s_test%d.mol\" %(features[-1][0].split(\".\")[0].split(\"/\")[-1], it))\n",
    "\n",
    "        \n",
    "        \n",
    "\n",
    "        #print(\"output_i\")\n",
    "        #print(sess.run(G_sample_output_i, feed_dict=feed_dict))\n",
    "        #print(sess.run(G_sample, feed_dict=feed_dict))\n",
    "\n",
    "    #random.shuffle(train_dataset)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.where(features[27][4][3] == 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "features[27][3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "every = 1000\n",
    "#train_dataset = features[:96]\n",
    "train_dataset = features[:1]*S\n",
    "n_train = len(train_dataset)\n",
    "\n",
    "for it in range(1000000):\n",
    "    a = time.time()\n",
    "    random.shuffle(train_dataset)\n",
    "    \n",
    "    batch_sched = list(range(0, n_train+1,S))\n",
    "    for j in range(0, len(batch_sched)-1):\n",
    "        start = batch_sched[j]\n",
    "        stop = batch_sched[j+1]\n",
    "        \n",
    "        for _ in range(5):\n",
    "            feed_dict = construct_feed_dict(train_dataset, start, stop)\n",
    "\n",
    "            _, D_loss_curr, _ = sess.run(\n",
    "                [D_solver, D_loss, clip_D], feed_dict=feed_dict)\n",
    "        \n",
    "        feed_dict = construct_feed_dict(train_dataset, start, stop)\n",
    "\n",
    "        _, G_loss_curr = sess.run([G_solver, G_loss], feed_dict=feed_dict)\n",
    "        \n",
    "        d_losses.append(D_loss_curr)\n",
    "        g_losses.append(G_loss_curr)\n",
    "    \n",
    "    if it % every == 0:\n",
    "        \n",
    "        print(time.time()-a)\n",
    "        print(\"Training epoch %d\" %it)\n",
    "                \n",
    "        print('Iter: {}'.format(it))\n",
    "        print('D loss: {:.4}'. format(D_loss_curr))\n",
    "        print('G_loss: {:.4}'.format(G_loss_curr))\n",
    "        feed_dict = construct_feed_dict(features[:1]*S, 0, S, shuffle_inds=False, keep_prob_val=1.)\n",
    "        \n",
    "        print(\"4wks:\\n\")\n",
    "        \n",
    "        print(\"real:\")\n",
    "        print(feed_dict[mol_xyz][0][:4])\n",
    "        print(\"fake:\")\n",
    "        print(sess.run(G_sample, feed_dict=feed_dict)[0][:4])\n",
    "        print(\"fake:\")\n",
    "        print(sess.run(rot_mat_i, feed_dict=feed_dict))\n",
    "\n",
    "        print(np.mean(sess.run(dihed, feed_dict=feed_dict)[:, 0]))\n",
    "        plt.hist(sess.run(dihed, feed_dict=feed_dict)[:, 0])\n",
    "        plt.show()\n",
    "        \n",
    "        \n",
    "        generate_molecule(features[0][0], sess.run(G_sample, feed_dict=feed_dict)[0], \"/home/evan/Documents/deep_docking/dcgan_pg/DCGAN-tensorflow/generated_mols/%s_test%d.mol\" %(features[0][0].split(\".\")[0].split(\"/\")[-1], it))\n",
    "        \n",
    "        print(\"Butane:\\n\")\n",
    "        feed_dict = construct_feed_dict(features[len(features)-32:], 0, S, shuffle_inds=False, keep_prob_val=1.)\n",
    "        plt.hist(sess.run(dihed, feed_dict=feed_dict)[:, 0])\n",
    "        plt.show()\n",
    "        #generate_molecule(features[1][0], sess.run(G_sample, feed_dict=feed_dict)[1], \"/home/evan/Documents/deep_docking/dcgan_pg/DCGAN-tensorflow/generated_mols/%s_test%d.mol\" %(features[1][0].split(\".\")[0].split(\"/\")[-1], it))\n",
    "\n",
    "        #feed_dict = construct_feed_dict(features, len(features)-S, len(features), shuffle_inds=False, keep_prob_val=1.)\n",
    "        #generate_molecule(features[-1][0], sess.run(G_sample, feed_dict=feed_dict)[0], \"/home/evan/Documents/deep_docking/dcgan_pg/DCGAN-tensorflow/generated_mols/%s_test%d.mol\" %(features[-1][0].split(\".\")[0].split(\"/\")[-1], it))\n",
    "\n",
    "        \n",
    "        \n",
    "\n",
    "        #print(\"output_i\")\n",
    "        #print(sess.run(G_sample_output_i, feed_dict=feed_dict))\n",
    "        #print(sess.run(G_sample, feed_dict=feed_dict))\n",
    "\n",
    "    #random.shuffle(train_dataset)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 321,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXcAAAD8CAYAAACMwORRAAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAD55JREFUeJzt3V+IXnedx/H3xxhRViHrZtZmk8yOy+bGimgZYql7kXXX\npU2L2ZUuRFjrdi+GlAoKgkSFinfKgiw10iGsxZYVRVC7wSa41a3YXqSaZGNq2rqO0iUJWVsrpoYW\nJfrdizmr4+NMnjMzz/z7+X7BYc6f35zzPfMbPpz5zTnnSVUhSWrLS9a6AEnS6BnuktQgw12SGmS4\nS1KDDHdJapDhLkkNMtwlqUGGuyQ1yHCXpAa9dK0OvHXr1pqYmFirw0vShnTy5MkfV9XYsHZrFu4T\nExOcOHFirQ4vSRtSkv/p085hGUlqkOEuSQ0y3CWpQYa7JDXIcJekBvUO9ySbkvxXkq/Msy1J7k4y\nk+RMkutGW6YkaTEWc+X+XuDJBbbdBOzqpingnmXWJUlahl7hnmQHcDPwrws02QfcX7OOA1uSbBtR\njZKkRep75f4vwAeAXy2wfTtwbs7y+W6dJGkNDH1CNcktwDNVdTLJnuUcLMkUs8M2jI+PL2dXkrRu\nTBx8cFHtn/7YzStUyW/0uXJ/C/D2JE8DnwfemuTfBtpcAHbOWd7RrfstVXW4qiaranJsbOirESRJ\nSzQ03Kvqg1W1o6omgP3Af1bVPww0OwLc1t01cz1wqaoujr5cSVIfS35xWJIDAFU1DRwF9gIzwAvA\n7SOpTpK0JIsK96r6BvCNbn56zvoC7hxlYZKkpfMJVUlqkOEuSQ0y3CWpQYa7JDXIcJekBhnuktQg\nw12SGmS4S1KDDHdJapDhLkkNMtwlqUGGuyQ1yHCXpAYZ7pLUIMNdkhpkuEtSgwx3SWrQ0HBP8vIk\n30rynSRnk3x0njZ7klxKcrqb7lqZciVJffT5mL2fA2+tqstJNgOPJjlWVccH2j1SVbeMvkRJ0mIN\nDffu81Evd4ubu6lWsihJ0vL0GnNPsinJaeAZ4KGqemyeZjckOZPkWJJrR1qlJGlReoV7Vf2yqt4I\n7AB2J3n9QJNTwHhVvQH4JPDAfPtJMpXkRJITzz777HLqliRdxaLulqmqnwIPAzcOrH++qi5380eB\nzUm2zvP9h6tqsqomx8bGllG2JOlq+twtM5ZkSzf/CuBtwFMDba5Jkm5+d7ff50ZfriSpjz53y2wD\n7kuyidnQ/kJVfSXJAYCqmgZuBe5IcgV4Edjf/SNWkrQG+twtcwZ40zzrp+fMHwIOjbY0SdJS+YSq\nJDXIcJekBhnuktQgw12SGmS4S1KDDHdJapDhLkkNMtwlqUGGuyQ1yHCXpAYZ7pLUIMNdkhpkuEtS\ngwx3SWqQ4S5JDTLcJalBhrskNajPZ6i+PMm3knwnydkkH52nTZLcnWQmyZkk161MuZKkPvp8hurP\ngbdW1eUkm4FHkxyrquNz2twE7OqmNwP3dF8lSWtg6JV7zbrcLW7upsEPv94H3N+1PQ5sSbJttKVK\nkvrqc+VOkk3ASeDPgU9V1WMDTbYD5+Ysn+/WXRzYzxQwBTA+Pr7EkiVp5UwcfHCtSxiJXv9Qrapf\nVtUbgR3A7iSvX8rBqupwVU1W1eTY2NhSdiFJ6mFRd8tU1U+Bh4EbBzZdAHbOWd7RrZMkrYE+d8uM\nJdnSzb8CeBvw1ECzI8Bt3V0z1wOXquoikqQ10WfMfRtwXzfu/hLgC1X1lSQHAKpqGjgK7AVmgBeA\n21eoXklSD0PDvarOAG+aZ/30nPkC7hxtaZKkpfIJVUlqkOEuSQ0y3CWpQYa7JDXIcJekBhnuktQg\nw12SGmS4S1KDDHdJapDhLkkNMtwlqUGGuyQ1yHCXpAYZ7pLUIMNdkhpkuEtSgwx3SWpQn89Q3Znk\n4SRPJDmb5L3ztNmT5FKS091018qUK0nqo89nqF4B3l9Vp5K8CjiZ5KGqemKg3SNVdcvoS5QkLdbQ\nK/equlhVp7r5nwFPAttXujBJ0tItasw9yQSzH5b92Dybb0hyJsmxJNcu8P1TSU4kOfHss88uulhJ\nUj+9wz3JK4EvAu+rqucHNp8CxqvqDcAngQfm20dVHa6qyaqaHBsbW2rNkqQheoV7ks3MBvtnq+pL\ng9ur6vmqutzNHwU2J9k60kolSb31uVsmwKeBJ6vqEwu0uaZrR5Ld3X6fG2WhkqT++twt8xbgXcDj\nSU536z4EjANU1TRwK3BHkivAi8D+qqoVqFeS1MPQcK+qR4EMaXMIODSqoiRJy+MTqpLUIMNdkhpk\nuEtSgwx3SWqQ4S5JDTLcJalBhrskNchwl6QGGe6S1CDDXZIaZLhLUoMMd0lqkOEuSQ0y3CWpQYa7\nJDXIcJekBhnuktSgPp+hujPJw0meSHI2yXvnaZMkdyeZSXImyXUrU64kqY8+n6F6BXh/VZ1K8irg\nZJKHquqJOW1uAnZ105uBe7qvkqQ1MPTKvaouVtWpbv5nwJPA9oFm+4D7a9ZxYEuSbSOvVpLUS58r\n919LMgG8CXhsYNN24Nyc5fPduosD3z8FTAGMj48vrlJJAiYOPrio9k9/7OYVqmR96/0P1SSvBL4I\nvK+qnl/KwarqcFVNVtXk2NjYUnYhSeqhV7gn2cxssH+2qr40T5MLwM45yzu6dZKkNdDnbpkAnwae\nrKpPLNDsCHBbd9fM9cClqrq4QFtJ0grrM+b+FuBdwONJTnfrPgSMA1TVNHAU2AvMAC8At4++VElS\nX0PDvaoeBTKkTQF3jqooSdLy+ISqJDXIcJekBhnuktQgw12SGmS4S1KDDHdJapDhLkkNMtwlqUGG\nuyQ1yHCXpAYZ7pLUIMNdkhpkuEtSgwx3SWqQ4S5JDTLcJalBfT5m794kzyT57gLb9yS5lOR0N901\n+jIlSYvR52P2PgMcAu6/SptHquqWkVQkSVq2oVfuVfVN4CerUIskaURGNeZ+Q5IzSY4luXZE+5Qk\nLVGfYZlhTgHjVXU5yV7gAWDXfA2TTAFTAOPj4yM4tCRpPsu+cq+q56vqcjd/FNicZOsCbQ9X1WRV\nTY6NjS330JKkBSw73JNckyTd/O5un88td7+SpKUbOiyT5HPAHmBrkvPAR4DNAFU1DdwK3JHkCvAi\nsL+qasUqliQNNTTcq+qdQ7YfYvZWSUnSOuETqpLUIMNdkhpkuEtSgwx3SWqQ4S5JDTLcJalBhrsk\nNchwl6QGGe6S1CDDXZIaZLhLUoMMd0lqkOEuSQ0y3CWpQYa7JDXIcJekBhnuktSgoeGe5N4kzyT5\n7gLbk+TuJDNJziS5bvRlSpIWo8+V+2eAG6+y/SZgVzdNAfcsvyxJ0nIMDfeq+ibwk6s02QfcX7OO\nA1uSbBtVgZKkxRvFmPt24Nyc5fPdOknSGnnpah4syRSzQzeMj48veT8TBx9c9Pc8/bGbl3y8jWg1\nfkZLOcZi/L712VIstg+W8jNdjWOspJX+PV2vRnHlfgHYOWd5R7fud1TV4aqarKrJsbGxERxakjSf\nUYT7EeC27q6Z64FLVXVxBPuVJC3R0GGZJJ8D9gBbk5wHPgJsBqiqaeAosBeYAV4Abl+pYiVJ/QwN\n96p655DtBdw5sookScvmE6qS1CDDXZIaZLhLUoMMd0lqkOEuSQ0y3CWpQYa7JDXIcJekBhnuktQg\nw12SGmS4S1KDDHdJapDhLkkNMtwlqUGGuyQ1yHCXpAYZ7pLUoF7hnuTGJN9LMpPk4Dzb9yS5lOR0\nN901+lIlSX31+QzVTcCngLcB54FvJzlSVU8MNH2kqm5ZgRolSYvU58p9NzBTVT+sql8Anwf2rWxZ\nkqTl6BPu24Fzc5bPd+sG3ZDkTJJjSa4dSXWSpCUZOizT0ylgvKouJ9kLPADsGmyUZAqYAhgfHx/R\noSVJg/pcuV8Ads5Z3tGt+7Wqer6qLnfzR4HNSbYO7qiqDlfVZFVNjo2NLaNsSdLV9An3bwO7krw2\nycuA/cCRuQ2SXJMk3fzubr/PjbpYSVI/Q4dlqupKkvcAXwU2AfdW1dkkB7rt08CtwB1JrgAvAvur\nqlawbknSVfQac++GWo4OrJueM38IODTa0iRJS+UTqpLUIMNdkhpkuEtSgwx3SWqQ4S5JDTLcJalB\nhrskNchwl6QGGe6S1CDDXZIaZLhLUoMMd0lqkOEuSQ0y3CWpQYa7JDXIcJekBhnuktSgXuGe5MYk\n30syk+TgPNuT5O5u+5kk142+VElSX0PDPckm4FPATcDrgHcmed1As5uAXd00Bdwz4jolSYvQ58p9\nNzBTVT+sql8Anwf2DbTZB9xfs44DW5JsG3GtkqSe+oT7duDcnOXz3brFtpEkrZKXrubBkkwxO2wD\ncDnJ91bwcFuBH//62B9fwSOtvd8616Vabz+jBeoZybluECM/19Xo4yUe4/eqX/PxZZ3rn/Zp1Cfc\nLwA75yzv6NYttg1VdRg43Kew5UpyoqomV+NYa81zbZPn2qbVOtc+wzLfBnYleW2SlwH7gSMDbY4A\nt3V3zVwPXKqqiyOuVZLU09Ar96q6kuQ9wFeBTcC9VXU2yYFu+zRwFNgLzAAvALevXMmSpGF6jblX\n1VFmA3zuuuk58wXcOdrSlm1Vhn/WCc+1TZ5rm1ZnaHo2lyVJLfH1A5LUoGbCPck/J3mqe/3Bl5Ns\nWaDd00keT3I6yYnVrnMUFnGuV31txEaQ5O+TnE3yqyQL3mHQSL/2PdcW+vXVSR5K8v3u6x8u0G5D\n9ut6eGVLM+EOPAS8vqreAPw38MGrtP3LqnrjBr71aui59nxtxEbwXeAdwDd7tN3o/Tr0XBvq14PA\n16tqF/D1bnkhG6pf18srW5oJ96r6j6q60i0eZ/Ze+yb1PNc+r41Y96rqyapayYfd1o2e59pEvzJb\n833d/H3A365hLaO2Ll7Z0ky4D/gn4NgC2wr4WpKT3ROzG91C5/r79kqI1vp1Ia3062vmPAvzv8Br\nFmi3Eft1XbyyZVVfP7BcSb4GXDPPpg9X1b93bT4MXAE+u8Bu/qKqLiT5Y+ChJE9VVZ8/+VfViM51\nQ+hzrj0006+tuNq5zl2oqkqy0G17G6Jf16MNFe5V9ddX257kH4FbgL+qBe7xrKoL3ddnknyZ2T+h\n1t0vywjOtdcrIdaDYefacx9N9GsPTfRrkh8l2VZVF7vhiGcW2MeG6NcBI3tly3I0MyyT5EbgA8Db\nq+qFBdr8QZJX/f888DfM/hNrQ+lzrvR7bUQTWunXnlrp1yPAu7v5dwO/81fLBu7X9fHKlqpqYmL2\n1QfngNPdNN2t/xPgaDf/Z8B3uukss38Kr3ntK3Gu3fJeZu+m+cEGPte/Y3Y88ufAj4CvNtyvQ8+1\noX79I2bvkvk+8DXg1S3163x9BBwADnTzYfaOmh8AjwOTo67BJ1QlqUHNDMtIkn7DcJekBhnuktQg\nw12SGmS4S1KDDHdJapDhLkkNMtwlqUH/ByMOfwSd9Hf0AAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<matplotlib.figure.Figure at 0x7feedffcfd50>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "S=16\n",
    "feed_dict = construct_feed_dict(features, len(features)-S, len(features), shuffle_inds=False)\n",
    "\n",
    "plt.hist(np.squeeze(sess.run(dihed, feed_dict=feed_dict)), bins=25)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 216,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('/home/evan/Documents/deep_docking/datasets/v2015/2ohk/2ohk_ligand.mol2',\n",
       " array([[-1.2731,  0.9987, -2.2001],\n",
       "        [-0.5031,  1.0397, -1.0801],\n",
       "        [-0.0411,  2.2447, -0.6841],\n",
       "        [ 0.7269,  2.2857,  0.4109],\n",
       "        [ 1.0309,  1.1747,  1.1729],\n",
       "        [ 0.5309, -0.0663,  0.8089],\n",
       "        [ 0.7909, -1.2213,  1.5379],\n",
       "        [ 0.2719, -2.4343,  1.1299],\n",
       "        [-0.4991, -2.5073, -0.0101],\n",
       "        [-0.7681, -1.3733, -0.7401],\n",
       "        [-0.2671, -0.1413, -0.3461],\n",
       "        [ 0.    ,  0.    ,  0.    ]]),\n",
       " array([[[0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0],\n",
       "         [0, 0, 0, 0]]], dtype=uint8),\n",
       " array([[ 0.]]),\n",
       " array([[[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]], dtype=uint8),\n",
       " array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], dtype=uint8),\n",
       " array([[ 0.,  0.,  0.]]),\n",
       " (array([[1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "         [1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0],\n",
       "         [0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "         [0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0],\n",
       "         [0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0],\n",
       "         [0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 0],\n",
       "         [0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0],\n",
       "         [0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 1, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], dtype=uint8),\n",
       "  array([[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 1, 0, 0, 0, 0],\n",
       "         [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 1, 0, 0, 0, 1],\n",
       "         [0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 1, 0, 0, 0, 1],\n",
       "         [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 1, 0, 0, 0, 1],\n",
       "         [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 1, 0, 0, 0, 1],\n",
       "         [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 1, 0, 0, 0, 1],\n",
       "         [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 1, 0, 0, 0, 1],\n",
       "         [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 1, 0, 0, 0, 1],\n",
       "         [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 1, 0, 0, 0, 1],\n",
       "         [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 1, 0, 0, 0, 1],\n",
       "         [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 1, 0, 0, 0, 1],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0]], dtype=uint8)))"
      ]
     },
     "execution_count": 216,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "features[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "feed_dict = construct_feed_dict(features, len(features)-S, len(features), shuffle_inds=False)\n",
    "generate_molecule(features[-1][0], sess.run(G_sample, feed_dict=feed_dict)[0], \"/home/evan/Documents/deep_docking/dcgan_pg/DCGAN-tensorflow/generated_mols/%s_test%d.mol\" %(features[-1][0].split(\".\")[0].split(\"/\")[-1], it))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.min(results[0])\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.where(features[0][1][:,:,:,:-1] == 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "def discriminator(W_list, b_list, h_list, L_list, n_layers, mols):\n",
    "    mols = tf.add(mols, mol_noise)\n",
    "    with tf.device('/gpu:0'):\n",
    "        for layer_idx in range(n_layers):\n",
    "            h_list.append(tf.nn.dropout(adjacency_conv_layer(h_list[layer_idx], W_list[layer_idx],\n",
    "                                             b_list[layer_idx], adj_list[layer_idx], \n",
    "                                             clades_list[layer_idx],\n",
    "                                             L_list[layer_idx], L_list[layer_idx+1],\n",
    "                                             layer_idx, S, B_list[layer_idx]), keep_prob))\n",
    "        h_final = h_list[-1]\n",
    "        dist_matrix = tf.concat([tf.reshape(compute_distance_matrix(tf.reshape(mol, [B, 3])), (1, B, B)) for mol in tf.split(mols, S, axis=0)], axis=0)\n",
    "        print(\"dist_matrix\")\n",
    "        print(dist_matrix)\n",
    "        #dist_matrix = compute_distance_matrix(mol)\n",
    "        bond_dist_gen = tf.multiply(dist_matrix, adj_matrix)\n",
    "        wrong_bond_loss = tf.divide(tf.reshape(tf.multiply(tf.reduce_sum(tf.square(tf.subtract(bond_dist_mat, bond_dist_gen)), axis=[1,2]), lam1), (S,1)), n_bonds)\n",
    "        print(\"wrong_bond_loss\")\n",
    "        print(wrong_bond_loss)\n",
    "\n",
    "        angles = []\n",
    "        for i in range(0, B):\n",
    "            for j in range(0, B):\n",
    "                ij = mols[:, i,:] - mols[:,j,:]\n",
    "                ij = ij / tf.sqrt(tf.reduce_sum(tf.square(ij), 1, keep_dims=True))\n",
    "                #ij shape should be (S, 3)\n",
    "               # print(\"ij\")\n",
    "                #print(ij)\n",
    "                for k in range(0,B):\n",
    "                    ik = mols[:,i,:] - mols[:,k,:]\n",
    "                    ik = ik / tf.sqrt(tf.reduce_sum(tf.square(ik), 1, keep_dims=True))\n",
    "\n",
    "                    #ik shape should be (S, 3)\n",
    "                    #if i == 0 and j == 0 and k == 0: \n",
    "                    #   print(\"ik\")\n",
    "                     #   print(ik)\n",
    "                        \n",
    "                    dp = tf.reduce_sum(tf.multiply(ij, ik), axis=1, keep_dims=True)\n",
    "                    #dp shape should be (S,1)\n",
    "                    #if i ==0 and j ==0 and k==0:\n",
    "                    #    print(\"dp\")\n",
    "                    #    print(dp)\n",
    "                    if i == j or i == k or j == k:\n",
    "                        angle = tf.Variable(tf.constant(np.zeros([4,1]).astype(np.float32)))\n",
    "                    else:\n",
    "                        angle = tf.acos(tf.clip_by_value(tf.reduce_sum(tf.multiply(ij, ik), axis=1, keep_dims=True), -1.0, 1.0))\n",
    "                    #if i ==0 and j==0 and k==0:\n",
    "                        #print(\"angle\")\n",
    "                        #print(angle)\n",
    "                    angles.append(angle)\n",
    "        angles_gen = tf.reshape(tf.concat(angles, axis=1), [S, B, B, B])\n",
    "        angles_gen = tf.multiply(angles_gen, angle_tuples)\n",
    "        wrong_angles_loss = tf.divide(tf.reshape(tf.multiply(tf.reduce_sum(tf.square(tf.subtract(angle_tensor, angles_gen)), axis=[1,2,3]), lam2), (S,1)), n_angles)\n",
    "                \n",
    "        feat_i = tf.reshape(h_final, (S, B, 1, L_list[n_layers]))\n",
    "        feat_i = tf.tile(feat_i, [1, 1, B, 1])\n",
    "        \n",
    "        feat_j = tf.reshape(h_final, (S, B, 1, L_list[n_layers]))\n",
    "        feat_j = tf.transpose(feat_j, perm=[0, 2, 1, 3])\n",
    "        feat_j = tf.tile(feat_j, [1, B, 1, 1])\n",
    "        \n",
    "        d_h1 = tf.concat([tf.reshape(dist_matrix, [S, B, B, 1]), feat_i, feat_j], axis=3)\n",
    "        d_h1 = tf.reshape(d_h1, [S, B, B*(L_list[n_layers]*2+1)])\n",
    "        d_h2 = tf.nn.tanh(tf.matmul(d_h1, d_w1) + d_b1)\n",
    "        d_h2 = tf.divide(tf.reduce_sum(tf.matmul(d_h2, d_w2) + d_b2, axis=1), n_bonds)#, axis=1\n",
    "        print(\"d_h2\")\n",
    "        print(d_h2)\n",
    "        \n",
    "    \n",
    "    return(wrong_bond_loss + wrong_angles_loss + d_h2, bond_dist_gen, angles_gen, wrong_bond_loss, wrong_angles_loss)\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "mol = Chem.MolFromMol2File(ligand_files[0])\n",
    "c = mol.GetConformer(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "c.SetAtomPosition(0, [0., 0., 0.])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.where(features[2][1][:,:,:,:-1] == 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "results = sess.run(G_sample, feed_dict=feed_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "sess.run(pred_real, feed_dict=feed_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.max(results[0][:,:,:,:], axis=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "features[1][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "mol = Chem.MolFromMol2File(features[1][0])\n",
    "mol.GetNumAtoms()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from rdkit.Chem import AllChem\n",
    "AllChem.Compute2DCoords(mol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from rdkit.Chem import Draw\n",
    "\n",
    "Draw.MolToImage(mol)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}