{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from glob import glob\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rdkit import Chem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_sdf(sdf):\n",
    "    with open(sdf, \"r\") as f:\n",
    "        txt = f.read().rstrip()\n",
    "    return txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_ncharges_coords(sdf):\n",
    "    mol = Chem.MolFromMolBlock(sdf)\n",
    "   #mol = Chem.AddHs(mol)\n",
    "    # rdkit molobj\n",
    "    ncharges = [atom.GetAtomicNum() for atom in mol.GetAtoms()]\n",
    "    conf = mol.GetConformer()\n",
    "    coords = np.asarray(conf.GetPositions())\n",
    "    return ncharges, coords"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cutoff_func(R_ij, central_cutoff=4.8, central_decay=1):\n",
    "    if R_ij <= (central_cutoff - central_decay):\n",
    "        func = 1.\n",
    "    elif ((central_cutoff - central_decay) < R_ij) and (R_ij <= (central_cutoff + central_decay)):\n",
    "        func = 0.5 * (1. + np.cos((np.pi * R_ij - central_cutoff + central_decay)/central_decay))\n",
    "    else:\n",
    "        func = 0.\n",
    "    return func"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_atomic_CM(ncharges, coords, max_natoms, central_cutoff=4.8, central_decay=1):\n",
    "    size = int((max_natoms + 1)*max_natoms / 2)\n",
    "    rep = np.zeros((len(ncharges), size))\n",
    "    \n",
    "    # central atom loop\n",
    "    for k in range(len(ncharges)):\n",
    "        M = np.zeros((len(ncharges), len(ncharges)))\n",
    "        for i in range(len(ncharges)):\n",
    "            R_ik = np.linalg.norm(coords[i]-coords[k])\n",
    "           # print('R_ik', R_ik)\n",
    "            f_ik = cutoff_func(R_ik, central_cutoff=central_cutoff,\n",
    "                              central_decay=central_decay)\n",
    "            for j in range(len(ncharges)):\n",
    "                if i <=j:\n",
    "                    if i == j:\n",
    "                        M[i,j] = 0.5 * ncharges[i]**2.4 * f_ik**2\n",
    "                        M[j,i] = M[i,j]\n",
    "\n",
    "                    else:\n",
    "                        R_jk = np.linalg.norm(coords[j]-coords[k])\n",
    "                      #  print('R_jk', R_jk)\n",
    "                        f_jk = cutoff_func(R_jk, central_cutoff=central_cutoff,\n",
    "                                          central_decay=central_decay)\n",
    "                        R_ij = np.linalg.norm(coords[i]-coords[j])\n",
    "                      #  print('R_ij', R_ij)\n",
    "                        f_ij = cutoff_func(R_ij, central_cutoff=central_cutoff,\n",
    "                                          central_decay=central_decay)\n",
    "                        M[i,j] = (ncharges[i]*ncharges[j]/R_ij)*f_ik*f_jk*f_ij\n",
    "                        M[j,i] = M[i,j]\n",
    "\n",
    "\n",
    "        # concat upper triangular and diagonal\n",
    "        upper_triang = M[np.triu_indices(len(M))]\n",
    "        \n",
    "        # pad to full size\n",
    "        n_zeros = size - len(upper_triang)\n",
    "        zeros = np.zeros(n_zeros)\n",
    "        rep[k] = np.concatenate((upper_triang, zeros))\n",
    "\n",
    "    return rep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['targets/qm9.sdf', 'targets/vitc.sdf', 'targets/vitd.sdf']"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "target_sdfs = sorted(glob(\"targets/*.sdf\"))\n",
    "target_sdfs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "qm9_amons_files = sorted(glob(\"amons-qm9/*.sdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "qm9_amons_sdfs = [read_sdf(x) for x in qm9_amons_files]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "conf_data = [get_ncharges_coords(x) for x in qm9_amons_sdfs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "ncharges_list, coords_list = zip(*conf_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "qm9_ncharges = ncharges_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "qm9_reps = [np.array(get_atomic_CM(np.array(ncharges_list[i]),\n",
    "                                                                        np.array(coords_list[i]), \n",
    "                                                                        max_natoms=9))\n",
    "            for i in range(len(ncharges_list))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/puck/anaconda3/envs/rdkit/lib/python3.7/site-packages/ipykernel_launcher.py:1: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
      "  \"\"\"Entry point for launching an IPython kernel.\n"
     ]
    }
   ],
   "source": [
    "qm9_reps = np.array(qm9_reps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 45)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "qm9_reps[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "qm9_amons_labels = [t.split(\"/\")[-1].split(\".sdf\")[0] for t in qm9_amons_files]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "vitc_amons_files = sorted(glob(\"amons-vitc/*.sdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "vitc_amons_sdfs = [read_sdf(x) for x in vitc_amons_files]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "conf_data = [get_ncharges_coords(x) for x in vitc_amons_sdfs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "ncharges_list, coords_list = zip(*conf_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "vitc_ncharges = ncharges_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "vitc_reps = [np.array(get_atomic_CM(np.array(ncharges_list[i]), np.array(coords_list[i]), \n",
    "                                                         max_natoms=12)) for i in \n",
    "            range(len(ncharges_list))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/puck/anaconda3/envs/rdkit/lib/python3.7/site-packages/ipykernel_launcher.py:1: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
      "  \"\"\"Entry point for launching an IPython kernel.\n"
     ]
    }
   ],
   "source": [
    "vitc_reps = np.array(vitc_reps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "vitc_amons_labels = [t.split(\"/\")[-1].split(\".sdf\")[0] for t in vitc_amons_files]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "vitd_amons_files = sorted(glob(\"amons-vitd/*.sdf\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "vitd_amons_sdfs = [read_sdf(x) for x in vitd_amons_files]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "conf_data = [get_ncharges_coords(x) for x in vitd_amons_sdfs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "ncharges_list, coords_list = zip(*conf_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "vitd_ncharges = ncharges_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "vitd_reps = [np.array(get_atomic_CM(np.array(ncharges_list[i]), np.array(coords_list[i]),\n",
    "                                                                         max_natoms=28))\n",
    "    for i in range(len(ncharges_list))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/puck/anaconda3/envs/rdkit/lib/python3.7/site-packages/ipykernel_launcher.py:1: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
      "  \"\"\"Entry point for launching an IPython kernel.\n"
     ]
    }
   ],
   "source": [
    "vitd_reps = np.array(vitd_reps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "vitd_amons_labels = [t.split(\"/\")[-1].split(\".sdf\")[0] for t in vitd_amons_files]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "# np save "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/puck/anaconda3/envs/rdkit/lib/python3.7/site-packages/numpy/core/_asarray.py:136: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
      "  return array(a, dtype, copy=False, order=order, subok=True)\n"
     ]
    }
   ],
   "source": [
    "np.savez(\"amons_aCM_data.npz\", \n",
    "         vitd_amons_labels=vitd_amons_labels,\n",
    "         vitc_amons_labels=vitc_amons_labels,\n",
    "         qm9_amons_labels=qm9_amons_labels,\n",
    "         vitd_amons_ncharges=vitd_ncharges,\n",
    "         vitc_amons_ncharges=vitc_ncharges,\n",
    "         qm9_amons_ncharges=qm9_ncharges,\n",
    "         vitd_amons_reps=vitd_reps,\n",
    "         vitc_amons_reps=vitc_reps,\n",
    "         qm9_amons_reps=qm9_reps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 406)"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vitd_reps[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}