0_parse-large-molecules.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parse PBD Molecules\n",
    "We also have the solvation energies of PDB molecules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jcesr_ml.utils import compute_atomization_energy\n",
    "from dlhub_sdk.models.datasets import TabularDataset\n",
    "from pymatgen.io.xyz import XYZ\n",
    "from ase.io.xyz import read_xyz\n",
    "from pybel import readstring\n",
    "from io import StringIO\n",
    "from glob import glob\n",
    "from tqdm import tqdm\n",
    "from tarfile import TarFile\n",
    "import pandas as pd\n",
    "import gzip\n",
    "import json\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_path = os.path.join('data', 'output', 'pdb_data.json')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Find the Files\n",
    "Load in the list of files from a TarFile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "tar = TarFile(fileobj=gzip.open(os.path.join('data', 'input', 'PDB_large_molecules_191.tar.gz')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 191 XYZ files\n"
     ]
    }
   ],
   "source": [
    "files = [x for x in tar.getmembers() if x.name.endswith('.xyz')]\n",
    "print(f'Found {len(files)} XYZ files')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "35\n",
      "Lattice=\"1.0 0.0 0.0 0.0 1.0 0.0 0.0 0.0 1.0\" Properties=species:S:1:pos:R:3 pbc=\"F F F\"\n",
      "O      -6.51686000      -0.61602700      -0.00039300 \n",
      "O      -5.51527900       1.39380300      -0.00016200 \n",
      "C       0.96293300      -0.52199000       0.00006700 \n",
      "C       2.22995000       0.34011200       0.00012200 \n",
      "C      -0.33295700       0.29586300       0.00020900 \n",
      "C       3.52637800      -0.47705100      -0.00008300 \n",
      "C      -1.59981700      -0.56637900       0.00016800 \n",
      "C       4.79338100       0.38463300       0.00001700 \n",
      "C      -2.89051900       0.25851500       0.00026700 \n",
      "C       6.09026600      -0.43186500      -0.00020100 \n",
      "C      -4.14756000      -0.61082800       0.00025100 \n",
      "C       7.35096100       0.43697800      -0.00012400 \n",
      "C      -5.42877900       0.19228700      -0.00012100 \n",
      "H       0.97446300      -1.18437600       0.87709600 \n",
      "H       0.97439900      -1.18418200      -0.87711100 \n",
      "H       2.21819200       1.00228000       0.87721300 \n",
      "H       2.21808700       1.00254900      -0.87676500 \n",
      "H      -0.34565900       0.95802400      -0.87662700 \n",
      "H      -0.34558600       0.95783300       0.87718900 \n",
      "H       3.53787700      -1.13927100      -0.87727100 \n",
      "H       3.53795600      -1.13960600       0.87685100 \n",
      "H      -1.58833100      -1.22884500      -0.87697400 \n",
      "H      -1.58828000      -1.22900400       0.87718700 \n",
      "H       4.78267500       1.04712000       0.87710500 \n",
      "H       4.78258900       1.04745500      -0.87681800 \n",
      "H      -2.91241300       0.92020300       0.87380800 \n",
      "H      -2.91248600       0.92033100      -0.87317400 \n",
      "H       6.10091700      -1.09344800       0.87632100 \n",
      "H       6.10081600      -1.09315600      -0.87694500 \n",
      "H      -4.17332500      -1.27696100       0.87206800 \n",
      "H      -4.17303400      -1.27738200      -0.87125000 \n",
      "H       7.38554000       1.08430200       0.88340400 \n",
      "H       7.38538600       1.08468000      -0.88337900 \n",
      "H       8.25930200      -0.17383700      -0.00033300 \n",
      "H      -7.28739400      -0.02857200      -0.00058700 \n",
      "Chemical_Formula                                   (Ha): C11H22O2\n",
      "B3LYP_energy                                         (Ha): -582.957501\n",
      "B3LYP_ZPE                                            (Ha):   0.317227\n",
      "B3LYP Enthalpy correction                            (Ha):   0.334797\n",
      "B3LYP_Zero_energy                                    (Ha): -582.640275\n",
      "B3LYP_Enthalpy                                       (Ha): -582.622704\n",
      "wB97XD_energy                                        (Ha): -582.936237\n",
      "wB97XD_B3LYP-Zero_energy                             (Ha): -582.619010\n",
      "G4MP2-0K                                             (Ha): -581.975795\n",
      "G4MP2-ZPE                                            (Ha):   0.312595\n",
      "G4MP2-Energy                                         (Ha): -581.959005\n",
      "G4MP2-Enthalpy                                       (Ha): -581.958061\n",
      "G4MP2-FreeE                                          (Ha): -582.023580\n",
      "G4MP2-Atomization Energy                       (kcal/mol): 3255.853205\n",
      "G4MP2-Standard Enthalpy                        (kcal/mol): -148.154943\n",
      "B3LYP Atomization Energy                       (kcal/mol): 3257.165329\n",
      "B3LYP Standard Enthalpy                        (kcal/mol): -149.569350\n",
      "wB97XD Atomization Energy                       (kcal/mol): 3255.904365\n",
      "wB97XD Standard Enthalpy                        (kcal/mol): -148.309014\n",
      "Solvation Energies (kcal/mol)\n",
      "Water    -3.1259\n",
      "ACN      -9.1403\n",
      "DMSO     -7.5710\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(tar.extractfile(files[0]).read().decode())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parse Data\n",
    "Get the key fields and render the data as a dataframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_solvation_energy(lines):\n",
    "    \"\"\"\n",
    "    Extract the solvation energies\n",
    "    \n",
    "    Args:\n",
    "        lines ([string]): Lines from the XYZ file\n",
    "    Returns:\n",
    "        (dict): Properites related to solubility calculations\n",
    "    \"\"\"\n",
    "    \n",
    "    # Generate list of JSCER properties\n",
    "    properties = ['sol_water', 'sol_acn', 'sol_dmso']\n",
    "    return dict((name, float(line.split()[-1])) for name, line in \n",
    "               zip(properties, lines[-3:]))\n",
    "# extract_g4mp2(content) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_clean_xyz_file(lines):\n",
    "    \"\"\"The XYZ files in our dataset don't parse with ASE or pymatgen easily\n",
    "    \n",
    "    This operation cleans up the file to make it parse easily\n",
    "    \n",
    "    Args:\n",
    "        lines ([string]): Entire file\n",
    "    Returns:\n",
    "        (string) File in a cleaner XYZ format\n",
    "    \"\"\"\n",
    "    \n",
    "    # Force pymatgen's XYZ parser to read the file\n",
    "    mol = XYZ._from_frame_string(\"\\n\".join(lines))\n",
    "    \n",
    "    # Convert the molecule back to XYZ format\n",
    "    return XYZ(mol).__str__()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_counts(xyz):\n",
    "    \"\"\"Given the XYZ coordinates of a molecule, compute its size\n",
    "    \n",
    "    Args:\n",
    "        xyz (string): XYZ-format file of the atomic structure\n",
    "    Returns:\n",
    "        (dict) Different measures of molecular size:\n",
    "            n_heavy_atom (int): Number of heavy (i.e., non-Hydrogen) atoms\n",
    "            n_electrons (int): Number of electrons in the system\n",
    "    \"\"\"\n",
    "    \n",
    "    mol = read_xyz(StringIO(xyz)).__next__()\n",
    "    Z = mol.get_atomic_numbers()\n",
    "    return {\n",
    "        'n_atoms': len(Z),\n",
    "        'n_heavy_atoms': (Z > 1).sum(),\n",
    "        'n_electrons': Z.sum(), \n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_files(file_list):\n",
    "    \"\"\"Convert XYZ files to a dataframe\n",
    "    \n",
    "    Args:\n",
    "        file_list ([str]): Path to files to be extracted\n",
    "    Returns:\n",
    "        (DataFrame) Dataset\n",
    "    \"\"\"\n",
    "    matrix = []\n",
    "    for file in tqdm(file_list):\n",
    "        with tar.extractfile(file) as fp:\n",
    "            lines = fp.read().decode().strip().split(\"\\n\")\n",
    "        \n",
    "        # Extract the numerical data\n",
    "        item = extract_solvation_energy(lines)\n",
    "        \n",
    "        # Store the file name and XYZ file\n",
    "        item['name'] = os.path.basename(file.name)\n",
    "        item['xyz'] = get_clean_xyz_file(lines)\n",
    "        \n",
    "        # Compute the smiles \n",
    "        item['smiles'] = readstring('xyz', item['xyz']).write('smiles')\n",
    "        \n",
    "        # Get some metrics of the atomic size\n",
    "        item.update(get_counts(item['xyz']))\n",
    "        \n",
    "        matrix.append(item)\n",
    "        \n",
    "    df = pd.DataFrame(matrix)\n",
    "    return df     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 191/191 [00:00<00:00, 494.00it/s]\n"
     ]
    }
   ],
   "source": [
    "dataset = parse_files(files)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset['source'] = 'pdb'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sort it into a reproducible order"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset.sort_values('name', inplace=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Mark a Test Set\n",
    "Just in case we want to do any kind of hold-out test later"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset['in_holdout'] = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset.loc[dataset.sample(frac=0.1, random_state=1).index, 'in_holdout'] = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save the Data with a Description\n",
    "Use the DLHub SDK to describe this data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset.to_json(output_path, 'records', lines=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make the description"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "metadata = TabularDataset.create_model(output_path, format='json', read_kwargs={'lines': True})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<dlhub_sdk.models.datasets.TabularDataset at 0x7f3fba115860>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "metadata.set_title('Solvation Energy of {} Large Molecules'.format(len(dataset)))\n",
    "metadata.set_name('pdb_solv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TBD: Get other authors from Rajeev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<dlhub_sdk.models.datasets.TabularDataset at 0x7f3fba115860>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "metadata.set_authors(['Dandu, Naveen', 'Assary, Rajeev', 'Narayanan, Badri', 'Curtiss, Larry'],\n",
    "                     [['Argonne National Laboratory']]*4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "for sol in dataset.columns:\n",
    "    if sol.startswith('sol_'):\n",
    "        metadata.annotate_column(sol, 'Solubility in {}'.format(sol.split(\"_\")[-1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<dlhub_sdk.models.datasets.TabularDataset at 0x7f3fba115860>"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "metadata.annotate_column('name', 'Filename, which contains the name of the molecule')\n",
    "metadata.annotate_column('source', 'Source of the dataset')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<dlhub_sdk.models.datasets.TabularDataset at 0x7f3fba115860>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "metadata.annotate_column('in_holdout', 'Whether the entry is in the test set')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<dlhub_sdk.models.datasets.TabularDataset at 0x7f3fba115860>"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "metadata.annotate_column('n_atoms', 'Number of atoms')\n",
    "metadata.annotate_column('n_electrons', 'Number of electrons')\n",
    "metadata.annotate_column('n_heavy_atoms', 'Number of non-hydrogen atoms')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<dlhub_sdk.models.datasets.TabularDataset at 0x7f3fba115860>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "metadata.annotate_column('smiles', 'SMILES string after relaxation', data_type='string')\n",
    "metadata.annotate_column('xyz', 'XYZ coordinates after relaxation', data_type='XYZ file')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<dlhub_sdk.models.datasets.TabularDataset at 0x7f3fba115860>"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "metadata.mark_inputs(['smiles', 'xyz'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(metadata.get_unannotated_columns()) == 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(os.path.dirname(output_path),\n",
    "                       '{}-description.json'.format(metadata['dlhub']['name'])), 'w') as fp:\n",
    "    json.dump(metadata.to_dict(), fp, indent=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "solv_ml",
   "language": "python",
   "name": "solv_ml"
  },
  "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}