add_default_pXC50.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "import os\n",
    "import glob\n",
    "from tqdm.auto import tqdm\n",
    "from rdkit import Chem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Adds helpful information for machine learning\n",
    "\n",
    "def compute_max_atomic_num(df):\n",
    "    smiles = df['neut-smiles'].values\n",
    "    mols = [Chem.MolFromSmiles(s) for s in smiles]\n",
    "    max_atomic_nums = []\n",
    "    max_num_atoms = []\n",
    "\n",
    "    for m in mols:\n",
    "        atom_nums = [a.GetAtomicNum() for a in m.GetAtoms()]\n",
    "        max_atomic_nums.append(np.max(atom_nums))\n",
    "        \n",
    "        num_atoms = len(m.GetAtoms())\n",
    "        max_num_atoms.append(num_atoms)\n",
    "    \n",
    "    df['Largest atomic number'] = np.max(max_atomic_nums)\n",
    "    df['# atoms'] = np.max(max_num_atoms)\n",
    "\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "paths = glob.glob('retrieved/*')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "60"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(paths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "OUT_DIR = 'retrieved_default_pXC50'\n",
    "\n",
    "for p in tqdm(paths):\n",
    "    ds_name = p.split('/')[-1]\n",
    "\n",
    "    out_dir = os.path.join(OUT_DIR, ds_name)\n",
    "    Path(out_dir).mkdir(exist_ok=True, parents=True)\n",
    "\n",
    "    if '-' not in ds_name:\n",
    "        df = pd.read_csv(os.path.join(p, 'SD.csv'))\n",
    "        df = df[~df['DR'].isna()]\n",
    "        df = compute_max_atomic_num(df)\n",
    "        df.to_csv(os.path.join(out_dir, 'DR.csv'), index=False)\n",
    "    else:    \n",
    "        df = pd.read_csv(os.path.join(p, 'DR.csv'))\n",
    "        df['DR determination'] = df['DR'].apply(lambda x: 'default' if np.isnan(x) else 'experimental')\n",
    "\n",
    "        df_dr_not_nan = df[~df['DR'].isna()]\n",
    "        dr_values = df_dr_not_nan['DR'].values\n",
    "\n",
    "        assert not np.isnan(dr_values).any()\n",
    "\n",
    "        min_dr = np.min(dr_values)\n",
    "\n",
    "        df['DR'] = df['DR'].apply(lambda x: min_dr if np.isnan(x) else x)\n",
    "\n",
    "        df = compute_max_atomic_num(df)\n",
    "        df.to_csv(os.path.join(out_dir, 'DR.csv'), index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch-geometric-2",
   "language": "python",
   "name": "torch-geometric-2"
  },
  "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.9.7"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "16d10e7bf223e4d2b7180a0577bd9603736d8148e85f5ee6f2a84aa0bc4db21d"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}