{
 "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": [
    "target_xyzs = sorted(glob(\"targets/*.xyz\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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": 5,
   "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": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cutoff_func(R_ij, central_cutoff=4.8, central_decay=0.03):\n",
    "    if R_ij <= (central_cutoff - central_decay):\n",
    "       # print('1')\n",
    "        func = 1.\n",
    "    elif ((central_cutoff - central_decay) < R_ij) and (R_ij <= (central_cutoff + central_decay)):\n",
    "       # print('function')\n",
    "        func = 0.5 * (1. + np.cos((np.pi * R_ij - central_cutoff + central_decay)))\n",
    "    else:\n",
    "       # print('zero')\n",
    "        func = 0.\n",
    "    return func"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_atomic_CM(ncharges, coords, max_natoms, central_cutoff=4.8, central_decay=0.03):\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(i):\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",
    "        # concat upper triangular and diagonal\n",
    "        upper_triang = np.triu(M)\n",
    "        non_zero_i, non_zero_j = np.nonzero(upper_triang)\n",
    "        unpadded_rep = upper_triang[non_zero_i, non_zero_j]\n",
    "        # pad to full size\n",
    "        n_zeros = size - len(unpadded_rep)\n",
    "        zeros = np.zeros(n_zeros)\n",
    "        rep[k] = np.concatenate((unpadded_rep, zeros))\n",
    "        \n",
    "    return rep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['targets/qm9.sdf', 'targets/vitc.sdf', 'targets/vitd.sdf']"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "target_files = sorted(glob(\"targets/*.sdf\"))\n",
    "target_files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "target_sdfs = [read_sdf(x) for x in target_files]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "conf_data = [get_ncharges_coords(x) for x in target_sdfs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "ncharges_list, coords_list = zip(*conf_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[9, 12, 28]"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sizes = [len(x) for x in ncharges_list]\n",
    "sizes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/puck/anaconda3/envs/rdkit/lib/python3.7/site-packages/ipykernel_launcher.py:4: 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",
      "  after removing the cwd from sys.path.\n"
     ]
    }
   ],
   "source": [
    "target_reps = np.array(\n",
    "[np.array(get_atomic_CM(np.array(ncharges_list[i]), np.array(coords_list[i]),\n",
    "                                max_natoms=sizes[i]))\n",
    "for i in range(len(ncharges_list))])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[39.89268481, 20.19979599, 16.0445292 , 10.27020053, 13.46303348,\n",
       "        16.70306767, 23.89446837, 17.57023229,  9.7738549 , 11.22347027,\n",
       "        14.10327203, 32.57835982, 15.64275949, 17.59643875, 24.93677728,\n",
       "        30.00133532, 17.36551167, 17.48814725, 17.13482138, 13.22143102,\n",
       "        32.63942034,  0.        ,  0.        ,  0.        ,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ],\n",
       "       [39.89268481, 20.19979599, 16.0445292 , 10.27020053, 13.46303348,\n",
       "        16.70306767, 23.89446837, 17.57023229,  9.7738549 , 10.24376593,\n",
       "        11.5318783 , 11.22347027, 14.10327203, 32.57835982, 15.64275949,\n",
       "        14.05740332, 18.07046347, 17.59643875, 24.93677728, 30.00133532,\n",
       "        24.37881483, 23.29821822, 17.36551167, 17.48814725, 40.31759536,\n",
       "        34.11172854, 17.13482138, 13.22143102, 28.78715844, 16.01650232,\n",
       "        12.2823109 , 40.74474394, 21.22940297, 32.63942034,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ],\n",
       "       [39.89268481, 20.19979599, 16.0445292 , 10.27020053, 13.46303348,\n",
       "        16.70306767, 23.89446837, 17.57023229,  9.7738549 , 10.24376593,\n",
       "        11.5318783 , 11.22347027, 14.10327203, 32.57835982, 15.64275949,\n",
       "        14.05740332, 18.07046347, 17.59643875, 24.93677728, 30.00133532,\n",
       "        24.37881483, 23.29821822, 17.36551167, 17.48814725, 40.31759536,\n",
       "        34.11172854, 17.13482138, 13.22143102, 28.78715844, 16.01650232,\n",
       "        12.2823109 , 40.74474394, 21.22940297, 32.63942034,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ],\n",
       "       [39.89268481, 20.19979599, 16.0445292 , 10.27020053, 13.46303348,\n",
       "        16.70306767, 23.89446837, 17.57023229,  9.7738549 , 10.24376593,\n",
       "        11.5318783 , 11.22347027, 14.10327203, 32.57835982, 15.64275949,\n",
       "        14.05740332, 18.07046347, 17.59643875, 24.93677728, 30.00133532,\n",
       "        24.37881483, 23.29821822, 17.36551167, 17.48814725, 40.31759536,\n",
       "        34.11172854, 17.13482138, 13.22143102, 28.78715844, 16.01650232,\n",
       "        12.2823109 , 40.74474394, 21.22940297, 32.63942034,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ],\n",
       "       [39.89268481, 20.19979599, 16.0445292 , 10.27020053, 13.46303348,\n",
       "        16.70306767, 23.89446837, 17.57023229,  9.7738549 , 10.24376593,\n",
       "        11.5318783 , 11.22347027, 14.10327203, 32.57835982, 15.64275949,\n",
       "        14.05740332, 18.07046347, 17.59643875, 24.93677728, 30.00133532,\n",
       "        24.37881483, 23.29821822, 17.36551167, 17.48814725, 40.31759536,\n",
       "        34.11172854, 17.13482138, 13.22143102, 28.78715844, 16.01650232,\n",
       "        12.2823109 , 40.74474394, 21.22940297, 32.63942034,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ],\n",
       "       [23.89446837, 17.57023229,  9.7738549 , 10.24376593, 11.5318783 ,\n",
       "        11.22347027, 14.10327203, 32.57835982, 15.64275949, 14.05740332,\n",
       "        18.07046347, 17.59643875, 24.93677728, 30.00133532, 24.37881483,\n",
       "        23.29821822, 17.36551167, 17.48814725, 40.31759536, 34.11172854,\n",
       "        17.13482138, 13.22143102, 28.78715844, 16.01650232, 12.2823109 ,\n",
       "        40.74474394, 21.22940297, 32.63942034,  0.        ,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ],\n",
       "       [23.89446837, 17.57023229,  9.7738549 , 10.24376593, 11.5318783 ,\n",
       "        11.22347027, 14.10327203, 32.57835982, 15.64275949, 14.05740332,\n",
       "        18.07046347, 17.59643875, 24.93677728, 30.00133532, 24.37881483,\n",
       "        23.29821822, 17.36551167, 17.48814725, 40.31759536, 34.11172854,\n",
       "        17.13482138, 13.22143102, 28.78715844, 16.01650232, 12.2823109 ,\n",
       "        40.74474394, 21.22940297, 32.63942034,  0.        ,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ],\n",
       "       [39.89268481, 20.19979599, 16.0445292 , 10.27020053, 13.46303348,\n",
       "        16.70306767, 23.89446837, 17.57023229,  9.7738549 , 10.24376593,\n",
       "        11.5318783 , 11.22347027, 14.10327203, 32.57835982, 15.64275949,\n",
       "        14.05740332, 18.07046347, 17.59643875, 24.93677728, 30.00133532,\n",
       "        24.37881483, 23.29821822, 17.36551167, 17.48814725, 40.31759536,\n",
       "        34.11172854, 17.13482138, 13.22143102, 28.78715844, 16.01650232,\n",
       "        12.2823109 , 40.74474394, 21.22940297, 32.63942034,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ],\n",
       "       [39.89268481, 20.19979599, 16.0445292 , 10.27020053, 13.46303348,\n",
       "        16.70306767, 23.89446837, 17.57023229,  9.7738549 , 10.24376593,\n",
       "        11.5318783 , 11.22347027, 14.10327203, 32.57835982, 15.64275949,\n",
       "        14.05740332, 18.07046347, 17.59643875, 24.93677728, 30.00133532,\n",
       "        24.37881483, 23.29821822, 17.36551167, 17.48814725, 40.31759536,\n",
       "        34.11172854, 17.13482138, 13.22143102, 28.78715844, 16.01650232,\n",
       "        12.2823109 , 40.74474394, 21.22940297, 32.63942034,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ,\n",
       "         0.        ,  0.        ,  0.        ,  0.        ,  0.        ]])"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "target_reps[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "target_labels = [t.split(\"/\")[-1].split(\".xyz\")[0] for t in target_sdfs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "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(\"target_aCM_data.npz\", \n",
    "         target_labels=target_labels, \n",
    "         target_reps=target_reps, \n",
    "         target_ncharges=ncharges_list,)"
   ]
  },
  {
   "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
}