smiles_to_xyz.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "adfb9149",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "from tqdm.auto import tqdm\n",
    "from openbabel import openbabel as ob\n",
    "from openbabel import pybel as pb\n",
    "from pymatgen import Molecule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9ea65792",
   "metadata": {},
   "outputs": [],
   "source": [
    "train = pd.read_csv(\"Clean/train.csv\", index_col=0)\n",
    "test = pd.read_csv(\"Clean/test.csv\", index_col=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d530c1b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def xyz_from_smiles(smiles, forcefield=\"gaff\", steps=1000, folder=None):\n",
    "    \"\"\" Find xyz coordinates of molecules from smile strings\"\"\"\n",
    "    \n",
    "    mols = []\n",
    "    mols_smi = [pb.readstring('smi', mol) for mol in smiles]\n",
    "    if folder:\n",
    "        os.makedirs(folder, exist_ok=True)\n",
    "    for idx, mol in tqdm(enumerate(mols_smi), total=len(mols_smi)):\n",
    "\n",
    "        # optimize 3D geometry\n",
    "        mol.make3D(forcefield=forcefield, steps=steps)\n",
    "        mol.localopt(forcefield=forcefield, steps=steps)\n",
    "\n",
    "        # convert to pymatgen class\n",
    "        species = []\n",
    "        coords = []\n",
    "        for atom in ob.OBMolAtomIter(mol.OBMol):\n",
    "            species.append(atom.GetAtomicNum())\n",
    "            coords.append([atom.GetX(), atom.GetY(), atom.GetZ()])\n",
    "        mol = Molecule(species, coords)\n",
    "        if folder:\n",
    "            mol.to(filename=f\"{folder}/{idx}.xyz\")\n",
    "        mols.append(mol)\n",
    "    return mols"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbe7d023",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_xyz = xyz_from_smiles(train[\"Smiles\"], folder=\"structures/train\")\n",
    "test_xyz = xyz_from_smiles(test[\"Smiles\"], folder=\"structures/test\")"
   ]
  }
 ],
 "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.6"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}