{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":67356,"databundleVersionId":8006601,"sourceType":"competition"}],"dockerImageVersionId":30746,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"I used a graph neural network to allow deep learning of structural molecular data. When futurizing the data, using one hot encoding, I emphasized important metals in biochemistry and especially in protein-ligand binding reactions like iron and zinc, while only selecting for a few metals that are especially prevalent in biochemistry to reduce the curse of dimensionality.","metadata":{}},{"cell_type":"code","source":"import numpy as np \nimport pandas as pd \nimport duckdb\nfrom rdkit import Chem\nfrom rdkit.Chem.rdmolops import GetAdjacencyMatrix\nimport torch\nfrom torch_geometric.data import Data\nfrom torch_geometric.loader import DataLoader\nfrom rdkit.Chem.rdmolops import GetAdjacencyMatrix\nimport torch.nn.functional as F\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch_geometric.nn import MessagePassing, global_mean_pool, global_max_pool\nfrom torch.nn import BCEWithLogitsLoss\nfrom sklearn.metrics import average_precision_score\nimport matplotlib.pyplot as plt","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-08-23T18:37:57.111795Z","iopub.execute_input":"2024-08-23T18:37:57.112164Z","iopub.status.idle":"2024-08-23T18:37:57.117335Z","shell.execute_reply.started":"2024-08-23T18:37:57.112136Z","shell.execute_reply":"2024-08-23T18:37:57.116136Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_path = r'train_path.parquet'\ncon = duckdb.connect()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#balance data for training by protein category\n#define the number of samples \nsamples_per_category = 100000  \n#For best results, if computationally available, \n#increase up to 100,000,000 \n\ndef get_balanced_data_for_protein(file_path, protein, samples):\n    query = f\"\"\"\n    (SELECT * FROM parquet_scan('{file_path}')\n     WHERE binds = 0 AND protein_name = '{protein}'\n     ORDER BY random()\n     LIMIT {samples})\n    UNION ALL\n    (SELECT * FROM parquet_scan('{file_path}')\n     WHERE binds = 1 AND protein_name = '{protein}'\n     ORDER BY random()\n     LIMIT {samples})\n    \"\"\"\n    return con.query(query).df()\n\n#proteins to query\nproteins = ['sEH', 'BRD4', 'HSA']\n\n#create a dictionary of dataframes with a separate dataframe for each protein\ndatasets = {}\nfor protein in proteins:\n    datasets[protein] = get_balanced_data_for_protein(train_path, protein, samples_per_category)\n\nsEH_df = datasets['sEH']\nbrd4_df = datasets['BRD4']\nhsa_df = datasets['HSA']","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#load the test datasets \n#filter for proteins without balancing them\ndef get_data_for_protein(file_path, protein):\n    return con.query(query).df()\n\ntest_path = r'test_path.parquet'\n#create a dictionary to store the filtered test datasets\ntest_datasets = {}\nfor protein in proteins:  # Using the same list of proteins: ['sEH', 'BRD4', 'HSA']\n    test_datasets[protein] = get_data_for_protein(test_path, protein)\n\nseh_test_df = test_datasets['sEH']\nbrd4_test_df = test_datasets['BRD4']\nhsa_test_df = test_datasets['HSA']\ncon.close()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#one hot encoding\ndef one_hot_encoding(x, permitted_list):\n    if x not in permitted_list:\n        x = permitted_list[-1]\n    binary_encoding = [int(boolean_value) for boolean_value in list(map(lambda s: x == s, permitted_list))]\n    return binary_encoding\n    \n#atom featurization\ndef get_atom_features(atom, use_chirality=True):\n    permitted_atom_types = ['C','N','O', 'Fe', 'Mo', 'Zn', 'Unknown']\n    atom_type = atom.GetSymbol() if atom.GetSymbol() in permitted_atom_types else 'Unknown'\n    atom_type_enc = one_hot_encoding(atom_type, permitted_atom_types)\n    atom_degree = one_hot_encoding(atom.GetDegree(), [0, 1, 2, 3, 4, 'MoreThanFour'])\n    is_in_ring = [int(atom.IsInRing())]\n    formal_charge_enc = one_hot_encoding(int(atom.GetFormalCharge()), [-3, -2, -1, 0, 1, 2, 3, \"Extreme\"])\n    hybridisation_type_enc = one_hot_encoding(str(atom.GetHybridization()), [\"S\", \"SP\", \"SP2\", \"SP3\", \"SP3D\", \"SP3D2\", \"OTHER\"])\n\n    if use_chirality == True:\n        chirality_enc = one_hot_encoding(str(atom.GetChiralTag()), [\"CHI_UNSPECIFIED\", \"CHI_TETRAHEDRAL_CW\", \"CHI_TETRAHEDRAL_CCW\", \"CHI_OTHER\"])\n        atom_features = atom_type_enc + atom_degree + is_in_ring + chirality_enc\n    else:\n        atom_features = atom_type_enc + atom_degree + is_in_ring\n        \n    return np.array(atom_features, dtype=np.float32)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#bond featurization\ndef get_bond_features(bond):\n    #list of bond types\n    permitted_bond_types = [Chem.rdchem.BondType.SINGLE, Chem.rdchem.BondType.DOUBLE, Chem.rdchem.BondType.TRIPLE, Chem.rdchem.BondType.AROMATIC, 'Unknown']\n    bond_type = bond.GetBondType() if bond.GetBondType() in permitted_bond_types else 'Unknown'\n    \n    #features\n    features = one_hot_encoding(bond_type, permitted_bond_types) \\\n               + [int(bond.IsInRing())]\n    \n    return np.array(features, dtype=np.float32)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def create_pytorch_geometric_graph_data_list_from_smiles_and_labels(x_smiles, ids, y=None):\n    data_list = []\n    \n    for index, smiles in enumerate(x_smiles):\n        mol = Chem.MolFromSmiles(smiles)\n        \n        if not mol:  #skip invalid SMILES strings\n            continue\n        \n        #node features\n        atom_features = [get_atom_features(atom) for atom in mol.GetAtoms()]\n        x = torch.tensor(atom_features, dtype=torch.float)\n        \n        #edge features\n        edge_index = []\n        edge_features = []\n        for bond in mol.GetBonds():\n            start, end = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()\n            edge_index += [(start, end), (end, start)]  # Undirected graph\n            bond_feature = get_bond_features(bond)\n            edge_features += [bond_feature, bond_feature]  # Same features in both directions\n        \n        edge_index = torch.tensor(edge_index, dtype=torch.long).t().contiguous()\n        edge_attr = torch.tensor(edge_features, dtype=torch.float)\n        \n        #create the Data object\n        data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr)\n        data.molecule_id = ids[index]\n        if y is not None:\n            data.y = torch.tensor([y[index]], dtype=torch.float)\n        \n        data_list.append(data)\n    \n    return data_list","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from tqdm import tqdm\ndef featurize_data_in_batches(smiles_list, labels_list, batch_size):\n    data_list = []\n    #define tqdm progress bar\n    pbar = tqdm(total=len(smiles_list), desc=\"Featurizing data\")\n    for i in range(0, len(smiles_list), batch_size):\n        smiles_batch = smiles_list[i:i+batch_size]\n        labels_batch = labels_list[i:i+batch_size]\n        ids_batch = ids_list[i:i+batch_size]\n        batch_data_list = create_pytorch_geometric_graph_data_list_from_smiles_and_labels(smiles_batch, ids_batch, labels_batch)\n        data_list.extend(batch_data_list)\n        pbar.update(len(smiles_batch))\n        \n    pbar.close()\n    return data_list","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#define the batch size\nbatch_size = 128\n#list of proteins and their dataframes\nproteins_data = {\n    'sEH': seh_df,\n    'BRD4': brd4_df,\n    'HSA': hsa_df\n}\n#dictionary to store the featurized data by protein\nfeaturized_data = {}\n#loop over each protein's dataframe\nfor protein_name, df in proteins_data.items():\n    print(f\"Processing {protein_name}...\")\n    smiles_list = df['molecule_smiles'].tolist()\n    ids_list = df['id'].tolist()\n    labels_list = df['binds'].tolist()\n #featurize the train data\n    featurized_data[protein_name] = featurize_data_in_batches(smiles_list, labels_list, batch_size)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#featurize the test data\nbatch_size = 128\nsmiles_list  = brd4_test_df['molecule_smiles'].tolist()\nids_list = brd4_test_df['id'].tolist()\nlabels_list = [-1]*len(smiles_list) #we dont have the actual labels, so we assign some dummy label list for the function. (don't choose 0 or 1)\nbrd4_test_data = featurize_data_in_batches(smiles_list, labels_list, batch_size)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"batch_size = 128\nsmiles_list  = seh_test_df['molecule_smiles'].tolist()\nids_list = seh_test_df['id'].tolist()\nlabels_list = [-1]*len(smiles_list)\nseh_test_data = featurize_data_in_batches(smiles_list, labels_list, batch_size)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"batch_size = 128\nsmiles_list  = hsa_test_df['molecule_smiles'].tolist()\nids_list = hsa_test_df['id'].tolist()\nlabels_list = [-1]*len(smiles_list)\nhsa_test_data = featurize_data_in_batches(smiles_list, labels_list, batch_size)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"featurized_data_train = {'sEH': seh_train_data,'BRD4': brd4_train_data, 'HSA': hsa_train_data}\n\nfeaturized_data_test = {\n    'sEH': seh_test_data,\n    'BRD4': brd4_test_data,\n    'HSA': hsa_test_data\n}","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch.save(featurized_data_train, r\"featurized_train_path\")","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch.save(featurized_data_test, r\"featurized_test_path\")","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Once data is featurized and stored (I used my hard drive), save time by running the algorithm from this point onward when tuning hyperparameters.","metadata":{}},{"cell_type":"code","source":"featurized_data_train = torch.load(r\"featurized_train_path\")\nfeaturized_data_train","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"featurized_data_test = torch.load(r\"featurized_test_path\")\nfeaturized_data_test","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CustomGNNLayer(MessagePassing):\n    def __init__(self, in_channels, out_channels):\n        super(CustomGNNLayer, self).__init__(aggr='Max')\n        self.lin = nn.Linear(in_channels + 6, out_channels)\n\n    def forward(self, x, edge_index, edge_attr):\n        # Start propagating messages\n        return MessagePassing.propagate(self, edge_index, x=x, edge_attr=edge_attr)\n\n    def message(self, x_j, edge_attr):\n        combined = torch.cat((x_j, edge_attr), dim=1)\n        return combined\n\n    def update(self, aggr_out):\n        return self.lin(aggr_out)\n\n#define the GNN Model\nclass GNNModel(nn.Module):\n    def __init__(self, input_dim, hidden_dim, num_layers, num_nodes, dropout_rate):\n        super(GNNModel, self).__init__()\n        self.num_layers = num_layers\n        self.num_nodes = num_nodes\n        self.convs = nn.ModuleList([CustomGNNLayer(input_dim if i == 0 else hidden_dim, hidden_dim) for i in range(num_layers)])\n        self.dropout = nn.Dropout(dropout_rate)\n        self.bns = nn.ModuleList([nn.BatchNorm1d(hidden_dim) for _ in range(num_layers)])\n        self.lin = nn.Linear(hidden_dim, 1)\n        \n    def forward(self, data):\n        x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr\n        for i in range(self.num_layers):\n            x = self.convs[i](x, edge_index, edge_attr)\n            x = self.bns[i](x)\n            x = F.relu(x)\n            x = self.dropout(x)\n\n\n        x = global_max_pool(x, data.batch) # Global pooling to get a graph-level representation\n        x = self.lin(x)\n        return x","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_model(train_loader, num_epochs, input_dim, hidden_dim, num_layers, dropout_rate, lr):\n    model = GNNModel(input_dim, hidden_dim, num_layers, num_nodes, dropout_rate)\n    optimizer = optim.AdamW(model.parameters(), lr=lr)\n    criterion = BCEWithLogitsLoss()\n    \n    for epoch in range(num_epochs):\n        model.train()\n        total_loss = 0\n        for batch in train_loader:\n            optimizer.zero_grad()\n            out = model(batch)\n            loss = criterion(out, batch.y.view(-1, 1).float())\n            loss.backward()\n            optimizer.step()\n            total_loss += loss.item()\n\n        print(f'Epoch {epoch+1}/{num_epochs}, Loss: {total_loss / len(train_loader)}')\n    \n    return model\n\ndef predict_with_model(model, test_loader):\n    model.eval()\n    predictions = []\n    molecule_ids = []\n\n    with torch.no_grad():\n        for data in test_loader:\n            output = torch.sigmoid(model(data))\n            predictions.extend(output.view(-1).tolist())\n            molecule_ids.extend(data.molecule_id)\n\n    return molecule_ids, predictions","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"proteins = ['sEH', 'BRD4', 'HSA']\nall_predictions = []\n\nfor protein in proteins:\n    print(f\"Training and predicting for {protein}\")\n    \n    #dataLoaders\n    train_loader = DataLoader(featurized_data_train[protein], batch_size=128, shuffle=True)\n    test_loader = DataLoader(featurized_data_test[protein], batch_size=128, shuffle=False)\n    #training hyperparamters \n    input_dim = train_loader.dataset[0].num_node_features\n    hidden_dim = 200\n    num_epochs = 33\n    num_layers = 6 #Should ideally be set so that all nodes can communicate with each other\n    num_nodes = 11\n    dropout_rate = 0.4\n    lr = 0.001\n    model = train_model(train_loader,num_epochs, input_dim, hidden_dim,num_layers, dropout_rate, lr)\n    \n    #predictions\n    molecule_ids, predictions = predict_with_model(model, test_loader)    \n    protein_predictions = pd.DataFrame({\n        'id': molecule_ids,\n        'binds': predictions,\n    })\n    all_predictions.append(protein_predictions)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#combine all predictions into one DataFrame\nfinal_predictions = pd.concat(all_predictions, ignore_index=True)\n#convert 'molecule_id' from tensors to integers in the DataFrame\nfinal_predictions['id'] = final_predictions['id'].apply(lambda x: x.item())\n\n#save\nfinal_predictions.to_csv(r\"final_prediction_path.csv\", index=False)\nfinal_predictions.head()","metadata":{},"execution_count":null,"outputs":[]}]}