datasetup_train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3528a7c6-cbcb-41c1-a4ac-28e4748f5e42",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Contains functions for embeding train set, training and testing a PyTorch model.\n",
    "\"\"\"\n",
    "\n",
    "import os\n",
    "import torch\n",
    "import numpy as np\n",
    "import h5py\n",
    "import pandas as pd\n",
    "from torch_geometric.data import Data, DataLoader as PyGDataLoader\n",
    "from tqdm import tqdm\n",
    "from rdkit import Chem\n",
    "import duckdb\n",
    "from multiprocessing import Pool\n",
    "from sklearn.model_selection import train_test_split\n",
    "import argparse\n",
    "from pathlib import Path\n",
    "import model_builder, engine  # Assuming these are your custom modules\n",
    "\n",
    "# Set up hyperparameters and system configuration\n",
    "PACK_NODE_DIM = 9\n",
    "PACK_EDGE_DIM = 1\n",
    "NODE_DIM = PACK_NODE_DIM * 8\n",
    "EDGE_DIM = PACK_EDGE_DIM * 8\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "NUM_WORKERS = 4\n",
    "DEFAULT_BATCH_SIZE = 128\n",
    "torch.manual_seed(42)\n",
    "\n",
    "\n",
    "# Download data from a specified path\n",
    "def download_data(path):\n",
    "    con = duckdb.connect()\n",
    "    sql_query = \"(SELECT * FROM read_csv('{}') WHERE binds = 0 ORDER BY random() LIMIT 15000000) UNION ALL (SELECT * FROM read_csv('{}') WHERE binds = 1 ORDER BY random())\".format(path, path)\n",
    "    df_list = []  # Liste pour stocker les blocs de DataFrame\n",
    "    try:\n",
    "        result = con.execute(sql_query)\n",
    "        while True:\n",
    "            df_chunk = result.fetch_df_chunk(1000)  # recup les donnees par chunks de 10 000\n",
    "            if df_chunk.empty:\n",
    "                break\n",
    "            df_list.append(df_chunk)  # Ajoute chaque bloc à la liste\n",
    "        df = pd.concat(df_list, ignore_index=True)  # Concatene tous les blocs en un seul DataFrame\n",
    "    except Exception as e:\n",
    "        print(\"An error occurred: \" + str(e))\n",
    "        raise\n",
    "    finally:\n",
    "        con.close()\n",
    "    return df\n",
    "\n",
    "# Preprocess the data'\n",
    "def preprocessing(df):\n",
    "    data_train = [smile.replace('[Dy]', 'C') for smile in df[\"molecule_smiles\"]]\n",
    "    labels = df[\"binds\"]\n",
    "    return data_train, labels\n",
    "\n",
    "# Convert SMILES strings to graph data\n",
    "def smile_to_graph(args):\n",
    "\tsmiles, label = args\n",
    "\tmol = Chem.MolFromSmiles(smiles)\n",
    "\tN = mol.GetNumAtoms()\n",
    "\tnode_feature = []\n",
    "\tedge_feature = []\n",
    "\tedge = []\n",
    "\tfor i in range(mol.GetNumAtoms()):\n",
    "\t\tatom_i = mol.GetAtomWithIdx(i)\n",
    "\t\tatom_i_features = get_atom_feature(atom_i)\n",
    "\t\tnode_feature.append(atom_i_features)\n",
    "\n",
    "\t\tfor j in range(mol.GetNumAtoms()):\n",
    "\t\t\tbond_ij = mol.GetBondBetweenAtoms(i, j)\n",
    "\t\t\tif bond_ij is not None:\n",
    "\t\t\t\tedge.append([i, j])\n",
    "\t\t\t\tbond_features_ij = get_bond_feature(bond_ij)\n",
    "\t\t\t\tedge_feature.append(bond_features_ij)\n",
    "\tnode_feature=np.stack(node_feature)\n",
    "\tedge_feature=np.stack(edge_feature)\n",
    "\tedge = np.array(edge,dtype=np.uint8)\n",
    "\treturn N,edge,node_feature,edge_feature, label\n",
    "\n",
    "######### Helper functions for feature extraction from smile\n",
    "def one_of_k_encoding(x, allowable_set, allow_unk=False):\n",
    "\tif x not in allowable_set:\n",
    "\t\tif allow_unk:\n",
    "\t\t\tx = allowable_set[-1]\n",
    "\t\telse:\n",
    "\t\t\traise Exception('input {} not in allowable set{allowable_set}!!!'.format(x))\n",
    "\treturn list(map(lambda s: x == s, allowable_set))\n",
    "\n",
    "\n",
    "##Get features of an atom (one-hot encoding:)\n",
    "'''\n",
    "\t1.atom element: 44+1 dimensions    \n",
    "\t2.the atom's hybridization: 5 dimensions\n",
    "\t3.degree of atom: 6 dimensions                        \n",
    "\t4.total number of H bound to atom: 6 dimensions\n",
    "\t5.number of implicit H bound to atom: 6 dimensions    \n",
    "\t6.whether the atom is on ring: 1 dimension\n",
    "\t7.whether the atom is aromatic: 1 dimension           \n",
    "\tTotal: 70 dimensions\n",
    "'''\n",
    "\n",
    "ATOM_SYMBOL = [\n",
    "\t'C', 'N', 'O', 'S', 'F', 'Si', 'P', 'Cl', 'Br', 'Mg',\n",
    "\t'Na', 'Ca', 'Fe', 'As', 'Al', 'I', 'B', 'V', 'K', 'Tl',\n",
    "\t'Yb', 'Sb', 'Sn', 'Ag', 'Pd', 'Co', 'Se', 'Ti', 'Zn', 'H',\n",
    "\t'Li', 'Ge', 'Cu', 'Au', 'Ni', 'Cd', 'In', 'Mn', 'Zr', 'Cr',\n",
    "\t'Pt', 'Hg', 'Pb', 'Dy',\n",
    "\t#'Unknown'\n",
    "]\n",
    "##print('ATOM_SYMBOL', len(ATOM_SYMBOL))44\n",
    "HYBRIDIZATION_TYPE = [\n",
    "\tChem.rdchem.HybridizationType.S,\n",
    "\tChem.rdchem.HybridizationType.SP,\n",
    "\tChem.rdchem.HybridizationType.SP2,\n",
    "\tChem.rdchem.HybridizationType.SP3,\n",
    "\tChem.rdchem.HybridizationType.SP3D\n",
    "]\n",
    "\n",
    "def get_atom_feature(atom):\n",
    "\tfeature = (\n",
    "\t\t one_of_k_encoding(atom.GetSymbol(), ATOM_SYMBOL)\n",
    "\t   + one_of_k_encoding(atom.GetHybridization(), HYBRIDIZATION_TYPE)\n",
    "\t   + one_of_k_encoding(atom.GetDegree(), [0, 1, 2, 3, 4, 5])\n",
    "\t   + one_of_k_encoding(atom.GetTotalNumHs(), [0, 1, 2, 3, 4, 5])\n",
    "\t   + one_of_k_encoding(atom.GetImplicitValence(), [0, 1, 2, 3, 4, 5])\n",
    "\t   + [atom.IsInRing()]\n",
    "\t   + [atom.GetIsAromatic()]\n",
    "\t)\n",
    "\t##feature = np.array(feature, dtype=np.uint8)\n",
    "\tfeature = np.packbits(feature)\n",
    "\treturn feature\n",
    "\n",
    "\n",
    "##Get features of an edge (one-hot encoding)\n",
    "'''\n",
    "\t1.single/double/triple/aromatic: 4 dimensions       \n",
    "\t2.the atom's hybridization: 1 dimensions\n",
    "\t3.whether the bond is on ring: 1 dimension          \n",
    "\tTotal: 6 dimensions\n",
    "'''\n",
    "\n",
    "def get_bond_feature(bond):\n",
    "\tbond_type = bond.GetBondType()\n",
    "\tfeature = [\n",
    "\t\tbond_type == Chem.rdchem.BondType.SINGLE,\n",
    "\t\tbond_type == Chem.rdchem.BondType.DOUBLE,\n",
    "\t\tbond_type == Chem.rdchem.BondType.TRIPLE,\n",
    "\t\tbond_type == Chem.rdchem.BondType.AROMATIC,\n",
    "\t\tbond.GetIsConjugated(),\n",
    "\t\tbond.IsInRing()\n",
    "\t]\n",
    "\t##feature = np.array(feature, dtype=np.uint8)\n",
    "\tfeature = np.packbits(feature)\n",
    "\treturn feature\n",
    "#########\n",
    "\n",
    "\n",
    "# Convert into a list of graphs \n",
    "def to_pyg_list(graph):\n",
    "    L = len(graph)\n",
    "    for i in tqdm(range(L)):\n",
    "        N, edge, node_feature, edge_feature, label = graph[i]\n",
    "        graph[i] = Data(\n",
    "            idx=i,\n",
    "            edge_index=torch.from_numpy(edge.T).int().to(device),\n",
    "            x=torch.from_numpy(node_feature).byte().to(device),\n",
    "            edge_attr=torch.from_numpy(edge_feature).byte().to(device),\n",
    "            y=torch.tensor(label).long().to(device),\n",
    "        )\n",
    "        torch.cuda.empty_cache()  # Liberer la memoire CUDA non utilisee apres chaque iteration\n",
    "    return graph\n",
    "\n",
    "\n",
    "# Convert a list of graphs to PyTorch Geometric DataLoader\n",
    "def to_pyg_loader(graphs, batch_size=32):\n",
    "    return PyGDataLoader(graphs, batch_size=batch_size)\n",
    "\n",
    "\n",
    "# Main function to get data in the required format\n",
    "def get_data_good_format(path, batch_size=32):\n",
    "    df = download_data(path)\n",
    "    smiles, labels = preprocessing(df)\n",
    "    # Transformer les Smiles en Graph \n",
    "    train_data = list(zip(smiles, labels))\n",
    "    num_train = len(train_data)\n",
    "    with Pool(NUM_WORKERS) as pool:\n",
    "        train_graphs = list(tqdm(pool.imap(smile_to_graph, train_data), total=num_train))\n",
    "    # Transformer les graph en objet Data Pytorch \n",
    "    train_graphs = to_pyg_list(train_graphs)\n",
    "    # Separe les donnees en ensembles dentrainement, de validation et de test\n",
    "    train_val_graphs, test_graphs = train_test_split(train_graphs, test_size=0.1, random_state=42)\n",
    "    train_graphs, val_graphs = train_test_split(train_val_graphs, test_size=0.1, random_state=42)\n",
    "    # Cree des DataLoader pour chaque ensemble de donnees\n",
    "    #train_loader = PyGDataLoader(train_graphs, batch_size=32, shuffle=True)\n",
    "    #val_loader = PyGDataLoader(val_graphs, batch_size=32, shuffle=False)\n",
    "    #test_loader = PyGDataLoader(test_graphs, batch_size=32, shuffle=False)\n",
    "    return train_graphs, val_graphs, test_graphs\n",
    "\n",
    "def main():\n",
    "    parser = argparse.ArgumentParser(description=\"Train a GNN model on chemical data.\")\n",
    "    parser.add_argument('--train_path', type=str, default='/user1/icmub/lg361770/Calculs/IA/leash_compet/train.csv', help='Path to the training data CSV file.')\n",
    "    parser.add_argument('--epochs', type=int, default=15, help='Number of epochs to train.')\n",
    "    args = parser.parse_args()\n",
    "\n",
    "    train_graphs, val_graphs, test_graphs = get_data_good_format(args.train_path, DEFAULT_BATCH_SIZE)\n",
    "    train_loader = PyGDataLoader(train_graphs, batch_size=DEFAULT_BATCH_SIZE, shuffle=True)\n",
    "    val_loader = PyGDataLoader(val_graphs, batch_size=DEFAULT_BATCH_SIZE, shuffle=False)\n",
    "    test_loader = PyGDataLoader(test_graphs, batch_size=DEFAULT_BATCH_SIZE, shuffle=False)\n",
    "\n",
    "    model = model_builder.Net().to(device)\n",
    "    engine.run_experiment(model=model, model_name=model.__class__.__name__, val_loader=val_loader, test_loader=test_loader, train_loader=train_loader, n_epochs=args.epochs)\n",
    "\n",
    "    MODEL_PATH = Path(\"/user1/icmub/lg361770/Calculs/IA/leash_compet/data_15M\")\n",
    "    MODEL_PATH.mkdir(parents=True, exist_ok=True)\n",
    "    MODEL_SAVE_PATH = MODEL_PATH / \"01_pytorch_GNN_15_000_000.pth\"\n",
    "    torch.save(obj=model.state_dict(), f=MODEL_SAVE_PATH)\n",
    "    print(f\"Model saved to: {MODEL_SAVE_PATH}\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "100jML",
   "language": "python",
   "name": "100jml"
  },
  "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}