dihedral_gan_mk2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook addresses the question, \"Can we represent a molecule as a graph via a 1D column vector or a 2D matrix of fixed length, with maximum number of atoms n_rows?\" Then, can we use this representation to learn neural fingerprints? E.g., can we make an aromatic ring detector? \n",
    "\n",
    "Scheme:\n",
    "feature_matrix = X\n",
    "for each ligand:\n",
    "    choose a central atom. this can be the atom (node) that minimizes distance to furthest heavy atom in graph.\n",
    "    set first row of X to be this central atom\n",
    "    set next four rows to be the atoms bonded to that centrl atom\n",
    "        set zeros for rows where row ind > n_bonds of atom\n",
    "    for each of those atoms:\n",
    "        repeat. find their neighbors. add to matrix.\n",
    "\n",
    "algorithm: breadth-first search:\n",
    "1. create networkx graph based on molecule\n",
    "2. find \"central\" atom (different strategies)\n",
    "3. define atom matrix of size (1+4+4*3^(L-1)) x (n_features_per_atom)\n",
    "4. start atom queue q\n",
    "5. central_atom.layer = 0; central_atom.row_idx = 0;\n",
    "6. q.enqueue(central_atom)\n",
    "7. define adjacency matrix of size (1+4+4*3^(L-1)) x 4\n",
    "\n",
    "def get_row_idx(curr_layer, prev_row_idx, curr_neighbor_idx):\n",
    "    if curr_layer == 0:\n",
    "        return(0)\n",
    "    if curr_layer == 1:\n",
    "        row_idx = 1 + curr_neighbor_idx\n",
    "    if layer == 2:\n",
    "        last_max = 5\n",
    "        row_idx = last_max + (3*(prev_row_idx-last_max)) + curr_neighbor_idx\n",
    "    if layer > 2:\n",
    "        last_max = 5 + 4*3^(curr_layer-2) \n",
    "        row_idx = last_max + 3*(prev_row_idx-last_max) + curr_neighbor_idx\n",
    "    return(row_idx)\n",
    "    \n",
    "\n",
    "while q.is_not_empty():\n",
    "    a = q.dequeue()\n",
    "    a.visited = True\n",
    "    for n_idx, n in enumerate(a.neighbors()):\n",
    "        if not n.visited:\n",
    "            row_idx = c\n",
    "            n.layer = a.layer + 1\n",
    "            row_idx = get_row_idx(n.layer, a.row_idx, n_idx)\n",
    "            n.row_idx = row_idx\n",
    "            adj_matrix[a.row_idx][n_idx] = n.row_idx\n",
    "            atom_matrix[row_idx][elem_to_idx[n.elem]] = 1\n",
    "\n",
    "input_matrix = tf.concat([atom_matrix, atom_matrix[adj_matrix[:,0]], atom_matrix[adj_matrix[:,1]], atom_matrix[adj_matrix[:,2]], atom_matrix[adj_matrix[:,3]]\n",
    "\n",
    "neural net:\n",
    "h1 = relu([tf.zeros([n_features_per_atom, 4]) * input_matrix + bias))\n",
    "h1_conc = tf.concat([h1, h1[adj_matrix[:,0], ..., h1[adj_matrix[:,3])\n",
    "\n",
    "repeat h1 to get h2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.61143659, -0.99900387,  0.48252666])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "np.cos( [-0.91292151,  3.09695411,  1.06725919])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "dihedral predictor pseudocode:\n",
    "\n",
    "get bonds for molecule\n",
    "create networkx graph out of molecule (use atom indices)\n",
    "\n",
    "for each edge:\n",
    "   for neighbor_i in atom_i.neighbors():\n",
    "       if neighbor_i == atom_j: continue\n",
    "       for neighbor_j in atom_j.neighbors():\n",
    "           if neighbor_j == atom_i: continue\n",
    "           dihedrals.append((neighbor_i, atom_i, neighbor_j, atom_j))\n",
    "           check to make sure (atom_j, neighbor_j, atom_i, neighbor_i)) not already in list\n",
    "\n",
    "for dihedral in dihedrals:\n",
    "    angle =  rdMolTransforms.GetDihedralDeg(c, 0,1,2,3)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "#MolGAN\n",
    "\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "\n",
    "\n",
    "def xavier_init(size):\n",
    "    in_dim = size[0]\n",
    "    xavier_stddev = 1. / tf.sqrt(in_dim / 2.)\n",
    "    return tf.cast(tf.random_normal(shape=size, stddev=xavier_stddev), tf.float32)\n",
    "\n",
    "atom_dim = 75\n",
    "hidden_dim = 50\n",
    "\n",
    "X = tf.placeholder(tf.float32, shape=[None, 1])\n",
    "\n",
    "D_W1 = tf.Variable(xavier_init([atom_dim, hidden_dim]))\n",
    "D_b1 = tf.Variable(tf.zeros(shape=[hidden_dim]))\n",
    "\n",
    "D_W2 = tf.Variable(xavier_init([hidden_dim, hidden_dim]))\n",
    "D_b2 = tf.Variable(tf.zeros(shape=[hidden_dim]))\n",
    "\n",
    "theta_D = [D_W1, D_W2, D_b1, D_b2]\n",
    "\n",
    "\n",
    "Z = tf.placeholder(tf.float32, shape=[1, 50])\n",
    "\n",
    "G_W1 = tf.Variable(xavier_init([100, 128]))\n",
    "G_b1 = tf.Variable(tf.zeros(shape=[128]))\n",
    "\n",
    "G_W2 = tf.Variable(xavier_init([128, 1]))\n",
    "G_b2 = tf.Variable(tf.zeros(shape=[1]))\n",
    "\n",
    "theta_G = [G_W1, G_W2, G_b1, G_b2]\n",
    "\n",
    "\n",
    "def sample_Z(m, n):\n",
    "    return np.sort(np.random.uniform(-1., 1., size=[m, n]))\n",
    "\n",
    "\n",
    "def generator(z, is_training=True):\n",
    "    G_h1 = tf.nn.tanh(tf.matmul(z, G_W1) + G_b1)\n",
    "    G_h1 = tf.contrib.layers.batch_norm(G_h1, \n",
    "                                        center=True, scale=True, \n",
    "                                        is_training=is_training)\n",
    "    G_log_prob = tf.matmul(G_h1, G_W2) + G_b2\n",
    "    G_prob = tf.nn.sigmoid(G_log_prob)\n",
    "\n",
    "    return G_prob\n",
    "\n",
    "\n",
    "def discriminator(x):\n",
    "    D_h1 = tf.nn.relu(tf.matmul(x, D_W1) + D_b1)\n",
    "    D_h1 = tf.contrib.layers.batch_norm(D_h1, \n",
    "                                        center=True, scale=True, \n",
    "                                        is_training=True)\n",
    "    D_logit = tf.matmul(D_h1, D_W2) + D_b2\n",
    "    D_prob = tf.nn.sigmoid(D_logit)\n",
    "\n",
    "    return D_prob, D_logit\n",
    "\n",
    "def sampler(z, is_training=False):\n",
    "    G_h1 = tf.nn.tanh(tf.matmul(z, G_W1) + G_b1)\n",
    "    G_h1 = tf.contrib.layers.batch_norm(G_h1, \n",
    "                                        center=True, scale=True, \n",
    "                                        is_training=is_training)\n",
    "    G_log_prob = tf.matmul(G_h1, G_W2) + G_b2\n",
    "    G_prob = tf.nn.sigmoid(G_log_prob)\n",
    "\n",
    "    return G_log_prob\n",
    "\n",
    "G_pred = sampler(Z)\n",
    "\n",
    "G_sample = generator(Z)\n",
    "D_real, D_logit_real = discriminator(X)\n",
    "D_fake, D_logit_fake = discriminator(G_sample)\n",
    "\n",
    "# D_loss = -tf.reduce_mean(tf.log(D_real) + tf.log(1. - D_fake))\n",
    "# G_loss = -tf.reduce_mean(tf.log(D_fake))\n",
    "\n",
    "# Alternative losses:\n",
    "# -------------------\n",
    "D_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=D_logit_real, labels=tf.ones_like(D_logit_real)))\n",
    "D_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=D_logit_fake, labels=tf.zeros_like(D_logit_fake)))\n",
    "D_loss = D_loss_real + D_loss_fake\n",
    "G_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=D_logit_fake, labels=tf.ones_like(D_logit_fake)))\n",
    "\n",
    "D_solver = tf.train.AdamOptimizer(learning_rate=1e-4).minimize(D_loss, var_list=theta_D)\n",
    "G_solver = tf.train.AdamOptimizer(learning_rate=1e-4).minimize(G_loss, var_list=theta_G)\n",
    "\n",
    "mb_size = 128\n",
    "Z_dim = 100\n",
    "\n",
    "sess = tf.Session()\n",
    "sess.run(tf.global_variables_initializer())\n",
    "\n",
    "if not os.path.exists('out/'):\n",
    "    os.makedirs('out/')\n",
    "\n",
    "i = 0\n",
    "\n",
    "for it in range(100000):\n",
    "    if it % 1000 == 0:\n",
    "        samples = sess.run(G_sample, feed_dict={Z: sample_Z(16, Z_dim)})\n",
    "\n",
    "    X_mb = np.sort(np.random.normal(size=(mb_size,1)))\n",
    "\n",
    "    _, D_loss_curr = sess.run([D_solver, D_loss], feed_dict={X: X_mb, Z: sample_Z(mb_size, Z_dim)})\n",
    "    _, G_loss_curr = sess.run([G_solver, G_loss], feed_dict={Z: sample_Z(mb_size, Z_dim)})\n",
    "    #_, G_loss_curr = sess.run([G_solver, G_loss], feed_dict={Z: sample_Z(mb_size, Z_dim)})\n",
    "\n",
    "    if it % 1000 == 0:\n",
    "        print('Iter: {}'.format(it))\n",
    "        print('D loss: {:.4}'. format(D_loss_curr))\n",
    "        print('G_loss: {:.4}'.format(G_loss_curr))\n",
    "        print()\n",
    "    if it % 1000 == 0:\n",
    "        plt.hist(sess.run(G_pred, feed_dict={Z: sample_Z(1000, 100)}))\n",
    "        plt.show()\n",
    "\"\"\"\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import matplotlib\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from rdkit import Chem\n",
    "from rdkit.Chem import AllChem, rdMolTransforms\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_torsions_angles(mol):\n",
    "    torsion_tuples = []\n",
    "    for bond in mol.GetBonds():\n",
    "        atom_i = bond.GetBeginAtom()\n",
    "        atom_j = bond.GetEndAtom()\n",
    "        if atom_i.IsInRing() or atom_j.IsInRing():\n",
    "            continue\n",
    "        for neighbor_i in atom_i.GetNeighbors():\n",
    "            if neighbor_i.GetIdx() == atom_j.GetIdx():\n",
    "                continue\n",
    "            \n",
    "            for neighbor_j in atom_j.GetNeighbors():\n",
    "                if neighbor_j.GetIdx() == atom_i.GetIdx():\n",
    "                    continue\n",
    "                torsion_tuple = (neighbor_i.GetIdx(), atom_i.GetIdx(), atom_j.GetIdx(), neighbor_j.GetIdx())\n",
    "                reverse_torsion_tuple = (neighbor_j.GetIdx(), atom_j.GetIdx(), atom_i.GetIdx(), neighbor_i.GetIdx())\n",
    "                if torsion_tuple not in torsion_tuples and reverse_torsion_tuple not in torsion_tuples:\n",
    "                    torsion_tuples.append(torsion_tuple)\n",
    "    c = mol.GetConformer(0)\n",
    "    torsions = []\n",
    "    torsion_matrix = np.zeros((250,1))\n",
    "    torsion_indices = np.zeros((250,200,4)).astype(np.uint8)\n",
    "    for i, torsion_tuple in enumerate(torsion_tuples):\n",
    "        torsion_matrix[i] = np.abs(rdMolTransforms.GetDihedralRad(c, *torsion_tuple))\n",
    "        torsion_indices[i][torsion_tuple[0]][0] = 1\n",
    "        torsion_indices[i][torsion_tuple[1]][1] = 1\n",
    "        torsion_indices[i][torsion_tuple[2]][2] = 1\n",
    "        torsion_indices[i][torsion_tuple[3]][3] = 1\n",
    "    return((torsion_indices, csr_matrix(torsion_matrix)))\n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def featurize_mols(mol_files):\n",
    "    featurizer = AdjacencyFingerprint(max_n_atoms=200)\n",
    "    features = []\n",
    "    for mol_file in mol_files:\n",
    "        mol = Chem.MolFromMol2File(mol_file)\n",
    "        if mol is None:\n",
    "            features.append(None)\n",
    "            continue\n",
    "        torsions = get_torsions_angles(mol)\n",
    "        graph_feat = featurizer.featurize([mol])[0]\n",
    "        features.append((mol_file, torsions, graph_feat))\n",
    "    return(features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import pickle\n",
    "feature_file = \"./dihedral_features_pdbbind.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 = [f for f in find_files(pdbbind_dir, \"*ligand.mol2\")]\n",
    "    features = featurize_mols(ligand_files)\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": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "features = [f for f in features if f is not None and len(np.where(f[1][1].toarray() == 0)[0]) < 250]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "len(features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "n_layers = 2\n",
    "S = 1\n",
    "B = 200\n",
    "L_list = [50, 50, 50, 50]\n",
    "p = 75\n",
    "\n",
    "x = tf.placeholder(tf.float32, shape=[S, B, p])\n",
    "\n",
    "non_zero_inds = tf.placeholder(tf.int32, shape=[None, S*250])\n",
    "\n",
    "adj_matrix = tf.placeholder(tf.float32, shape=[S, B, B])\n",
    "dihed_indices = tf.placeholder(tf.float32, shape=[S, 250, B, 4])\n",
    "\n",
    "label_placeholder = tf.placeholder(\n",
    "    dtype='float32', shape=[S*250], name=\"label_placeholder\")\n",
    "\n",
    "phase = tf.placeholder(dtype='bool', name='phase')\n",
    "\n",
    "z = tf.placeholder(tf.float32,\n",
    "                          [None, L_list[0]], name='z')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#DISCRIMINATOR\n",
    "\n",
    "D_W_list = [None for i in range(n_layers)]\n",
    "D_b_list = [None for i in range(n_layers)]\n",
    "D_h_list = [None for i in range(n_layers)]\n",
    "\n",
    "\n",
    "\n",
    "for layer_idx in range(n_layers):\n",
    "    if layer_idx == 0:\n",
    "      L_in = p\n",
    "      L_out = L_list[0]\n",
    "      atom_matrix = x\n",
    "    else:\n",
    "      L_in = L_list[layer_idx-1]\n",
    "      L_out = L_list[layer_idx]\n",
    "      atom_matrix = D_h_list[layer_idx-1]\n",
    "\n",
    "    D_W_list[layer_idx] = tf.Variable(tf.truncated_normal([L_in, L_out], seed=2017), name=\"D_W_list%d\" %layer_idx)\n",
    "    D_b_list[layer_idx] = tf.Variable(tf.zeros([1, L_out]))\n",
    "\n",
    "    \n",
    "D_W2 = tf.Variable(tf.truncated_normal([L_list[-1]*4+1, 100]))\n",
    "D_b2 = tf.Variable(tf.zeros([1, 100]))\n",
    "\n",
    "D_W3 = tf.Variable(tf.truncated_normal([100, 1]))\n",
    "D_b3 = tf.Variable(tf.zeros([1, 1]))\n",
    "\n",
    "D_W4 = tf.Variable(tf.truncated_normal([1, 10]))\n",
    "D_b4 = tf.Variable(tf.zeros([1, 10]))\n",
    "\n",
    "D_W5 = tf.Variable(tf.truncated_normal([10, 1]))\n",
    "D_b5 = tf.Variable(tf.zeros([1, 1]))\n",
    "\n",
    "def adjacency_conv_layer(atom_matrix, W, b, L_in, L_out, layer_idx, is_training=True):\n",
    "    print(\"layer_idx: %d\" %(layer_idx))\n",
    "    h = tf.matmul(adj_matrix, atom_matrix, name=\"adj_conv1\")\n",
    "    h = tf.reshape(h, shape=(S*B, L_in))\n",
    "\n",
    "    h = tf.nn.sigmoid(tf.matmul(h, W) + b)\n",
    "    h = tf.reshape(h, (S, B, L_out))\n",
    "    h = tf.contrib.layers.batch_norm(h, \n",
    "                                        center=True, scale=True, \n",
    "                                        is_training=is_training)\n",
    "\n",
    "    return(h)\n",
    "\n",
    "def discriminator(angles, is_training=True):\n",
    "    for layer_idx in range(n_layers):\n",
    "        if layer_idx == 0:\n",
    "          L_in = p\n",
    "          L_out = L_list[0]\n",
    "          atom_matrix = x\n",
    "        else:\n",
    "          L_in = L_list[layer_idx-1]\n",
    "          L_out = L_list[layer_idx]\n",
    "          atom_matrix = D_h_list[layer_idx-1]\n",
    "        \n",
    "        D_h_list[layer_idx] = adjacency_conv_layer(atom_matrix, D_W_list[layer_idx], D_b_list[layer_idx], L_in, L_out, layer_idx)\n",
    "\n",
    "    L_final = L_list[n_layers-1]\n",
    "    h_final = tf.reshape(D_h_list[layer_idx], (S, B, L_final))\n",
    "\n",
    "    #add dihedral regressor layers\n",
    "\n",
    "    d0 = []\n",
    "    for i in range(0, S):\n",
    "        mol_tuple = []\n",
    "        for j in range(0, 4):\n",
    "          entry = h_final[i]\n",
    "          indices = dihed_indices[i][:,:,j]\n",
    "          atom_list = tf.matmul(indices, entry, name=\"disc1\")\n",
    "          atom_list = tf.reshape(atom_list, (250, L_final))\n",
    "          mol_tuple.append(atom_list)\n",
    "        mol_tuple = tf.reshape(tf.stack(mol_tuple, axis=1), (250, L_final*4))\n",
    "        d0.append(mol_tuple)\n",
    "\n",
    "    d0 = tf.concat(d0, axis=0)\n",
    "    d0 = tf.concat([d0, tf.reshape(angles, (-1,1))], axis=1)\n",
    "    d0 = tf.matmul(tf.cast(non_zero_inds, tf.float32), d0)\n",
    "\n",
    "\n",
    "    d2 = tf.nn.tanh(tf.matmul(d0, D_W2, name=\"disc2\") + D_b2)\n",
    "\n",
    "    d2 = tf.contrib.layers.batch_norm(d2, \n",
    "                                        center=True, scale=True, \n",
    "                                        is_training=is_training)\n",
    "    \n",
    "    d3 = tf.matmul(d2, D_W3, name=\"disc3\") + D_b3\n",
    "\n",
    "    d3 = tf.nn.tanh(d3)\n",
    "\n",
    "    d3 = tf.contrib.layers.batch_norm(d3, \n",
    "                                        center=True, scale=True, \n",
    "                                        is_training=is_training)\n",
    "    \n",
    "                                                                                                                                                               \n",
    "    d4 = tf.nn.tanh(tf.matmul(d3, D_W4, name=\"disc3\") + D_b4)\n",
    "    \n",
    "    D_logit = tf.nn.tanh(tf.matmul(d4, D_W5, name=\"disc4\") + D_b5)\n",
    "    \n",
    "    D_prob = tf.nn.sigmoid(D_logit)\n",
    "    return(D_prob, D_logit)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def atan2(y, x, epsilon=1.0e-12):\n",
    "  # Add a small number to all zeros, to avoid division by zero:\n",
    "  x = tf.where(tf.equal(x, 0.0), x+epsilon, x)\n",
    "  y = tf.where(tf.equal(y, 0.0), y+epsilon, y)\n",
    "\n",
    "  angle = tf.where(tf.greater(x,0.0), tf.atan(y/x), tf.zeros_like(x))\n",
    "  angle = tf.where(tf.logical_and(tf.less(x,0.0),  tf.greater_equal(y,0.0)), tf.atan(y/x) + np.pi, angle)\n",
    "  angle = tf.where(tf.logical_and(tf.less(x,0.0),  tf.less(y,0.0)), tf.atan(y/x) - np.pi, angle)\n",
    "  angle = tf.where(tf.logical_and(tf.equal(x,0.0), tf.greater(y,0.0)), 0.5*np.pi * tf.ones_like(x), angle)\n",
    "  angle = tf.where(tf.logical_and(tf.equal(x,0.0), tf.less(y,0.0)), -0.5*np.pi * tf.ones_like(x), angle)\n",
    "  angle = tf.where(tf.logical_and(tf.equal(x,0.0), tf.equal(y,0.0)), tf.zeros_like(x), angle)\n",
    "  return angle\n",
    "\n",
    "#GENERATOR\n",
    "\n",
    "G_W_list = [None for i in range(n_layers)]\n",
    "G_b_list = [None for i in range(n_layers)]\n",
    "G_h_list = [None for i in range(n_layers)]\n",
    "\n",
    "\n",
    "for layer_idx in range(n_layers):\n",
    "    if layer_idx == 0:\n",
    "      L_in = p\n",
    "      L_out = L_list[0]\n",
    "      atom_matrix = x\n",
    "    else:\n",
    "      L_in = L_list[layer_idx-1]\n",
    "      L_out = L_list[layer_idx]\n",
    "      atom_matrix = G_h_list[layer_idx-1]\n",
    "\n",
    "    G_W_list[layer_idx] = tf.Variable(tf.truncated_normal([L_in, L_out], seed=2017), name=\"G_W_list%d\" %layer_idx)\n",
    "    G_b_list[layer_idx] = tf.Variable(tf.zeros([1, L_out]))\n",
    "\n",
    "    \n",
    "G_W2 = tf.Variable(tf.truncated_normal([L_list[-1]*4, 100]))\n",
    "G_b2 = tf.Variable(tf.zeros([1, 100]))\n",
    "\n",
    "G_W3 = tf.Variable(tf.truncated_normal([100, 100]))\n",
    "G_b3 = tf.Variable(tf.zeros([1, 100]))\n",
    "\n",
    "G_W4 = tf.Variable(tf.truncated_normal([100, 1]))\n",
    "G_b4 = tf.Variable(tf.zeros([1, 1]))\n",
    "\n",
    "\n",
    "def gen_adjacency_conv_layer(atom_matrix, W, b, L_in, L_out, layer_idx, z, is_training):\n",
    "    print(\"layer_idx: %d\" %(layer_idx))\n",
    "    h = tf.matmul(adj_matrix, atom_matrix, name=\"gen0_%d\" %layer_idx)\n",
    "    h = tf.reshape(h, shape=(S*B, L_in))\n",
    "\n",
    "    h = tf.nn.tanh(tf.matmul(h, W, name=\"gen1\") + b)\n",
    "\n",
    "    h = tf.add(h, z)\n",
    "    h = tf.reshape(h, (S, B, L_out))\n",
    "    h = tf.contrib.layers.batch_norm(h, \n",
    "                                        center=True, scale=True, \n",
    "                                        is_training=is_training)\n",
    "\n",
    "    return(h)\n",
    "\n",
    "def generator(z, is_training=True):\n",
    "    for layer_idx in range(n_layers):\n",
    "        if layer_idx == 0:\n",
    "          L_in = p\n",
    "          L_out = L_list[0]\n",
    "          atom_matrix = x\n",
    "        else:\n",
    "          L_in = L_list[layer_idx-1]\n",
    "          L_out = L_list[layer_idx]\n",
    "          atom_matrix = G_h_list[layer_idx-1]\n",
    "        \n",
    "        G_h_list[layer_idx] = gen_adjacency_conv_layer(atom_matrix, G_W_list[layer_idx], G_b_list[layer_idx], L_in, L_out, layer_idx, z, is_training)\n",
    "\n",
    "    L_final = L_list[n_layers-1]\n",
    "    g_h_final = tf.reshape(G_h_list[layer_idx], (S, B, L_final))\n",
    "\n",
    "    #add dihedral regressor layers\n",
    "\n",
    "    g_d0 = []\n",
    "    for i in range(0, S):\n",
    "        mol_tuple = []\n",
    "        for j in range(0, 4):\n",
    "          entry = g_h_final[i]\n",
    "          indices = dihed_indices[i][:,:,j]\n",
    "          atom_list = tf.matmul(indices, entry, name='gen2')\n",
    "          atom_list = tf.reshape(atom_list, (250, L_final))\n",
    "          mol_tuple.append(atom_list)\n",
    "        mol_tuple = tf.reshape(tf.stack(mol_tuple, axis=1), (250, L_final*4))\n",
    "        g_d0.append(mol_tuple)\n",
    "\n",
    "    g_d0 = tf.concat(g_d0, axis=0)\n",
    "    \n",
    "    g_d2 = tf.nn.tanh(tf.matmul(g_d0, G_W2, name='gen3') + G_b2)\n",
    "\n",
    "    g_d2 = tf.contrib.layers.batch_norm(g_d2, \n",
    "                                        center=True, scale=True, \n",
    "                                        is_training=is_training)\n",
    "    \n",
    "    g_d3 = tf.matmul(g_d2, G_W3) + G_b3\n",
    "    g_d3 = tf.nn.tanh(g_d3)\n",
    "\n",
    "    g_d3 = tf.contrib.layers.batch_norm(g_d3, \n",
    "                                        center=True, scale=True, \n",
    "                                        is_training=is_training)\n",
    "                                                                                                                                                               \n",
    "    g_d4 = tf.matmul(g_d3, G_W4) + G_b4\n",
    "    \n",
    "    #output = g_d4\n",
    "    d3_cos = tf.cos(g_d4)\n",
    "    d3_sin = tf.sin(g_d4)\n",
    "    output = tf.abs(atan2(d3_sin, d3_cos))\n",
    "    \n",
    "    G_logit = tf.nn.sigmoid(output)\n",
    "\n",
    "    return(output, G_logit)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "theta_D = [D_W_list[0], D_W_list[1], D_W2, D_W3, D_W4, D_W5, D_b_list[0], D_b_list[1], D_b2, D_b3, D_b4, D_b5]\n",
    "\n",
    "theta_G = [G_W_list[0], G_W_list[1], G_W2, G_W3, G_W4, G_b_list[0], G_b_list[1], G_b2, G_b3, G_b4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def construct_feed_dict(X, start=None,\n",
    "                      stop=None, y=None,\n",
    "                      keep_prob=1.0, train=False):\n",
    "    if start is None:\n",
    "      start = 0\n",
    "      stop = len(X)\n",
    "\n",
    "    adj = [X[idx][2][0].toarray().astype(np.float32) for idx in range(start, stop)]\n",
    "    A_batch = [X[idx][2][1].toarray() for idx in range(start, stop)]\n",
    "    D_batch = [X[idx][1][0] for idx in range(start, stop)]\n",
    "    y_batch = [X[idx][1][1].toarray() for idx in range(start, stop)]\n",
    "\n",
    "    y_batch = np.squeeze(np.concatenate(y_batch))\n",
    "\n",
    "    non_zero_batch = np.where(y_batch != 0.)[0]\n",
    "\n",
    "    onehotter = OneHotEncoder(n_values = S*X[0][1][1].shape[0])\n",
    "    non_zero_onehot = onehotter.fit_transform(non_zero_batch).toarray().reshape((len(non_zero_batch),S*X[0][1][1].shape[0]))\n",
    "    \n",
    "    z_batch = np.random.uniform(-1., 1., size=(1,50))\n",
    "    #y_batch = np.random.random(size=(S*250))\n",
    "    \n",
    "    feed_dict = {x: A_batch,\n",
    "                 adj_matrix: adj,\n",
    "                 phase: train,\n",
    "                 label_placeholder: y_batch,\n",
    "                 non_zero_inds: non_zero_onehot,\n",
    "                 dihed_indices: D_batch,\n",
    "                 z:z_batch\n",
    "                }\n",
    "    return(feed_dict)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "G_pred, G_pred_logit = generator(z, is_training=False)\n",
    "G_sample, G_sample_logit = generator(z, is_training=True)\n",
    "\n",
    "D_real, D_logit_real = discriminator(label_placeholder)\n",
    "D_fake, D_logit_fake = discriminator(G_sample)\n",
    "\n",
    "# D_loss = -tf.reduce_mean(tf.log(D_real) + tf.log(1. - D_fake))\n",
    "# G_loss = -tf.reduce_mean(tf.log(D_fake))\n",
    "\n",
    "# Alternative losses:\n",
    "# -------------------\n",
    "D_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=D_logit_real, labels=tf.ones_like(D_logit_real)))\n",
    "D_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=D_logit_fake, labels=tf.zeros_like(D_logit_fake)))\n",
    "D_loss = D_loss_real + D_loss_fake\n",
    "G_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=D_logit_fake, labels=tf.ones_like(D_logit_fake)))\n",
    "\n",
    "D_solver = tf.train.AdamOptimizer(learning_rate=1e-4).minimize(D_loss, var_list=theta_D)\n",
    "G_solver = tf.train.AdamOptimizer(learning_rate=1e-4).minimize(G_loss, var_list=theta_G)\n",
    "\n",
    "mb_size = S\n",
    "\n",
    "sess = tf.Session()\n",
    "sess.run(tf.global_variables_initializer())\n",
    "\n",
    "train_dataset = features[:1]\n",
    "n_train = len(train_dataset)\n",
    "preds = []\n",
    "\n",
    "for it in range(1000):\n",
    "    if it % 100 == 0:\n",
    "        print(\"Training epoch %d\" %it)\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",
    "        feed_dict = construct_feed_dict(train_dataset, start, stop)\n",
    "\n",
    "        _, D_loss_curr = sess.run([D_solver, D_loss], feed_dict=feed_dict)\n",
    "        _, G_loss_curr = sess.run([G_solver, G_loss], feed_dict=feed_dict)\n",
    "\n",
    "    if it % 100 == 0:\n",
    "        print('Iter: {}'.format(it))\n",
    "        print('D loss: {:.4}'. format(D_loss_curr))\n",
    "        print('G_loss: {:.4}'.format(G_loss_curr))\n",
    "        print()\n",
    "\n",
    "    if it % 100 == 0:\n",
    "        samples = sess.run(G_sample, feed_dict=feed_dict)\n",
    "        print(samples[:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "batch_sched"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "for it in range(10001):\n",
    "    if it % 1000 == 0:\n",
    "        print(\"Training epoch %d\" %it)\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",
    "        feed_dict = construct_feed_dict(train_dataset, start, stop)\n",
    "\n",
    "        _, D_loss_curr = sess.run([D_solver, D_loss], feed_dict=feed_dict)\n",
    "        _, G_loss_curr = sess.run([G_solver, G_loss], feed_dict=feed_dict)\n",
    "\n",
    "    if it % 1000 == 0:\n",
    "        print('Iter: {}'.format(it))\n",
    "        print('D loss: {:.4}'. format(D_loss_curr))\n",
    "        print('G_loss: {:.4}'.format(G_loss_curr))\n",
    "        print()\n",
    "\n",
    "    if it % 1000 == 0:\n",
    "        samples = sess.run(G_sample, feed_dict=feed_dict)\n",
    "        print(samples[:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "preds = []\n",
    "for k in range(0,100):\n",
    "    feed_dict = construct_feed_dict(train_dataset, 0, 1)\n",
    "    samples = sess.run(G_sample, feed_dict=feed_dict)[2]\n",
    "    preds.append(samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "plt.hist(np.concatenate(preds))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "preds[4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "train_dataset[0][1][1].toarray()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "res = gan.predict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "plt.hist(res, bins=50)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "features[0][1][1].toarray()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "features[0][1][1].toarray()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "preds = []\n",
    "for i in range(0,100):\n",
    "    preds.append(gan.predict(features)[0][2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "np.mean(preds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "features[0][2].atom_features"
   ]
  },
  {
   "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": []
  },
  {
   "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": []
  },
  {
   "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": []
  },
  {
   "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": []
  },
  {
   "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": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}