{"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"},{"sourceId":8023033,"sourceType":"datasetVersion","datasetId":4727900}],"dockerImageVersionId":30673,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-04-04T15:05:19.986424Z","iopub.execute_input":"2024-04-04T15:05:19.986850Z","iopub.status.idle":"2024-04-04T15:05:21.186172Z","shell.execute_reply.started":"2024-04-04T15:05:19.986817Z","shell.execute_reply":"2024-04-04T15:05:21.184809Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Introduction\n\nIn this tutorial, we will go through another of the several representations of molecules for ML and AI: Molecular Graphs. Using PyTorch Geometric, we are able to convert molecular data from their standard SMILES strings representation into actual graphs with nodes and edges representing atoms and bonds.","metadata":{}},{"cell_type":"markdown","source":"We first start by loading all the necessary packages:\n1. DuckDB is used to load in the datasets from the parquet files.\n2. RDKit is a cheminformatics library which will help in the conversion process from SMILES to graphs.\n3. We install Torch for all our neural network needs.","metadata":{}},{"cell_type":"code","source":"!pip install duckdb","metadata":{"execution":{"iopub.status.busy":"2024-04-04T15:05:32.986474Z","iopub.execute_input":"2024-04-04T15:05:32.987235Z","iopub.status.idle":"2024-04-04T15:05:58.695239Z","shell.execute_reply.started":"2024-04-04T15:05:32.987189Z","shell.execute_reply":"2024-04-04T15:05:58.693860Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install rdkit","metadata":{"execution":{"iopub.status.busy":"2024-04-04T15:05:58.698504Z","iopub.execute_input":"2024-04-04T15:05:58.698985Z","iopub.status.idle":"2024-04-04T15:06:16.979387Z","shell.execute_reply.started":"2024-04-04T15:05:58.698940Z","shell.execute_reply":"2024-04-04T15:06:16.978344Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install torch_geometric","metadata":{"execution":{"iopub.status.busy":"2024-04-04T15:06:16.980745Z","iopub.execute_input":"2024-04-04T15:06:16.981098Z","iopub.status.idle":"2024-04-04T15:06:33.719026Z","shell.execute_reply.started":"2024-04-04T15:06:16.981068Z","shell.execute_reply":"2024-04-04T15:06:33.717741Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import duckdb\nimport pandas as pd\n\ntrain_path = '/kaggle/input/leash-predict-chemical-bindings/train.parquet'\n\ncon = duckdb.connect()\n\n# Define the number of samples you want per category (binders and non-binders) per protein\nsamples_per_category = 1000  # Adjust this number as needed, this is just an example for the tutorial. \n#Your real values should be higher than this for proper training.\n\ndef get_balanced_data_for_protein(file_path, protein, samples):\n    \"\"\"\n    Fetches a balanced dataset for a specific protein.\n    \n    Parameters:\n    - file_path: Path to the dataset file.\n    - protein: The name of the protein.\n    - samples: Number of samples per binder category (1 or 0).\n    \n    Returns:\n    - A pandas DataFrame containing the balanced dataset for the protein.\n    \"\"\"\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# List of proteins to query for\nproteins = ['sEH', 'BRD4', 'HSA']\n\n# Creating a dictionary of dataframes, one for each protein\ndatasets = {}\nfor protein in proteins:\n    datasets[protein] = get_balanced_data_for_protein(train_path, protein, samples_per_category)\n\n# At this point, `datasets` contains separate dataframes for sEH, BRD4, and HSA\n# For example, to access the dataset for sEH:\nseh_df = datasets['sEH']\n\n# And similarly for BRD4 and HSA\nbrd4_df = datasets['BRD4']\nhsa_df = datasets['HSA']\n\n\n'''We now have balanced data for training for each protein, \nwhich we will featurize for our GNN model'''","metadata":{"execution":{"iopub.status.busy":"2024-04-04T15:06:33.722129Z","iopub.execute_input":"2024-04-04T15:06:33.723556Z","iopub.status.idle":"2024-04-04T15:08:25.769896Z","shell.execute_reply.started":"2024-04-04T15:06:33.723520Z","shell.execute_reply":"2024-04-04T15:08:25.768816Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"'''We do a similar process to load in the test datasets, \nhowever this time we only filter for the proteins, and\nwe don't balance them'''\ndef get_data_for_protein(file_path, protein):\n    \"\"\"\n    Fetches the dataset for a specific protein from the test dataset.\n    \n    Parameters:\n    - file_path: Path to the dataset file.\n    - protein: The name of the protein.\n    \n    Returns:\n    - A pandas DataFrame containing the dataset for the protein.\n    \"\"\"\n    query = f\"\"\"\n    SELECT * FROM parquet_scan('{file_path}')\n    WHERE protein_name = '{protein}'\n    \"\"\"\n    return con.query(query).df()\n\ntest_path = '/kaggle/input/leash-predict-chemical-bindings/test.parquet'\n\n# Assuming the connection `con` to DuckDB is still open from the previous operations\n\n# Creating 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\n# At this point, `test_datasets` contains separate dataframes for sEH, BRD4, and HSA without balancing\n# For example, to access the test dataset for sEH:\nseh_test_df = test_datasets['sEH']\n\n# And similarly for BRD4 and HSA\nbrd4_test_df = test_datasets['BRD4']\nhsa_test_df = test_datasets['HSA']\ncon.close()","metadata":{"execution":{"iopub.status.busy":"2024-04-04T15:08:25.771079Z","iopub.execute_input":"2024-04-04T15:08:25.771387Z","iopub.status.idle":"2024-04-04T15:08:28.571975Z","shell.execute_reply.started":"2024-04-04T15:08:25.771354Z","shell.execute_reply":"2024-04-04T15:08:28.570696Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"seh_df.head()","metadata":{"execution":{"iopub.status.busy":"2024-04-04T17:04:52.506245Z","iopub.execute_input":"2024-04-04T17:04:52.506783Z","iopub.status.idle":"2024-04-04T17:04:52.524549Z","shell.execute_reply.started":"2024-04-04T17:04:52.506740Z","shell.execute_reply":"2024-04-04T17:04:52.522881Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"seh_test_df.head() #note we don't have a binds column here","metadata":{"execution":{"iopub.status.busy":"2024-04-04T17:05:01.127356Z","iopub.execute_input":"2024-04-04T17:05:01.127840Z","iopub.status.idle":"2024-04-04T17:05:01.143712Z","shell.execute_reply.started":"2024-04-04T17:05:01.127803Z","shell.execute_reply":"2024-04-04T17:05:01.142461Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Featurizing","metadata":{}},{"cell_type":"code","source":"from rdkit import Chem\nfrom rdkit.Chem.rdmolops import GetAdjacencyMatrix\nimport torch\nfrom torch_geometric.data import Data\nfrom torch_geometric.loader import DataLoader","metadata":{"execution":{"iopub.status.busy":"2024-04-04T15:08:28.573255Z","iopub.execute_input":"2024-04-04T15:08:28.573552Z","iopub.status.idle":"2024-04-04T15:08:37.711463Z","shell.execute_reply.started":"2024-04-04T15:08:28.573527Z","shell.execute_reply":"2024-04-04T15:08:37.710465Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Atom Featurisation\n## Auxiliary function for one-hot enconding transformation based on list of\n##permitted values\n\ndef one_hot_encoding(x, permitted_list):\n    \"\"\"\n    Maps input elements x which are not in the permitted list to the last element\n    of the permitted list.\n    \"\"\"\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    \n# Main atom feat. func\n\ndef get_atom_features(atom, use_chirality=True):\n    # Define a simplified list of atom types\n    permitted_atom_types = ['C', 'N', 'O', 'S', 'P', 'F', 'Cl', 'Br', 'I','Dy', '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    \n    # Consider only the most impactful features: atom degree and whether the atom is in a ring\n    atom_degree = one_hot_encoding(atom.GetDegree(), [0, 1, 2, 3, 4, 'MoreThanFour'])\n    is_in_ring = [int(atom.IsInRing())]\n    \n    # Optionally include chirality\n    if use_chirality:\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":{"iopub.status.busy":"2024-04-04T15:08:37.713374Z","iopub.execute_input":"2024-04-04T15:08:37.714076Z","iopub.status.idle":"2024-04-04T15:08:37.729840Z","shell.execute_reply.started":"2024-04-04T15:08:37.714033Z","shell.execute_reply":"2024-04-04T15:08:37.728407Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Bond featurization\n\ndef get_bond_features(bond):\n    # Simplified 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: Bond type, Is in a ring\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":{"iopub.status.busy":"2024-04-04T15:08:37.732245Z","iopub.execute_input":"2024-04-04T15:08:37.732609Z","iopub.status.idle":"2024-04-04T15:08:37.759843Z","shell.execute_reply.started":"2024-04-04T15:08:37.732579Z","shell.execute_reply":"2024-04-04T15:08:37.758545Z"},"trusted":true},"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        # Creating 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":{"iopub.status.busy":"2024-04-04T15:08:37.761468Z","iopub.execute_input":"2024-04-04T15:08:37.761924Z","iopub.status.idle":"2024-04-04T15:08:37.776458Z","shell.execute_reply.started":"2024-04-04T15:08:37.761846Z","shell.execute_reply":"2024-04-04T15:08:37.774278Z"},"trusted":true},"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":{"iopub.status.busy":"2024-04-04T15:08:37.779387Z","iopub.execute_input":"2024-04-04T15:08:37.779908Z","iopub.status.idle":"2024-04-04T15:08:37.788710Z","shell.execute_reply.started":"2024-04-04T15:08:37.779875Z","shell.execute_reply":"2024-04-04T15:08:37.787475Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Define the batch size for featurization\nbatch_size = 2**8\n# List of proteins and their corresponding dataframes\nproteins_data = {\n    'sEH': seh_df,\n    'BRD4': brd4_df,\n    'HSA': hsa_df\n}\n# Dictionary to store the featurized data for each protein\nfeaturized_data = {}\n# Loop over each protein and its 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 data\n    featurized_data[protein_name] = featurize_data_in_batches(smiles_list, labels_list, batch_size)\n    \n\nseh_train_data = featurized_data['sEH']\nbrd4_train_data = featurized_data['BRD4']\nhsa_train_data = featurized_data['HSA']","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"At this stage, you can use torch.save to save your featurized datasets","metadata":{}},{"cell_type":"markdown","source":"What follows is the featurizing process for the test data. This can also be done in a loop rather than breaking it up into many cells as done here.","metadata":{}},{"cell_type":"code","source":"batch_size = 2**8\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":{"scrolled":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"batch_size = 2**8\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":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"batch_size = 2**8\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":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"seh_test_data[0] #example of what a pytorch object looks like","metadata":{"execution":{"iopub.status.busy":"2024-04-04T17:03:59.400535Z","iopub.execute_input":"2024-04-04T17:03:59.401278Z","iopub.status.idle":"2024-04-04T17:03:59.408595Z","shell.execute_reply.started":"2024-04-04T17:03:59.401244Z","shell.execute_reply":"2024-04-04T17:03:59.407179Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training and Test","metadata":{}},{"cell_type":"code","source":"featurized_data_train = {\n    'sEH': seh_train_data,\n    'BRD4': brd4_train_data,\n    'HSA': hsa_train_data\n}\n\nfeaturized_data_test = {\n    'sEH': seh_test_data,\n    'BRD4': brd4_test_data,\n    'HSA': hsa_test_data\n}","metadata":{"execution":{"iopub.status.busy":"2024-04-04T16:49:40.343140Z","iopub.execute_input":"2024-04-04T16:49:40.343600Z","iopub.status.idle":"2024-04-04T16:49:40.350618Z","shell.execute_reply.started":"2024-04-04T16:49:40.343570Z","shell.execute_reply":"2024-04-04T16:49:40.349258Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch.nn.functional as F\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch_geometric.loader import DataLoader\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\nimport numpy as np\n\n#Define custom GNN layer\nclass 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 GNN Model\nclass GNNModel(nn.Module):\n    def __init__(self, input_dim, hidden_dim, num_layers, dropout_rate):\n        super(GNNModel, self).__init__()\n        self.num_layers = num_layers\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\n","metadata":{"execution":{"iopub.status.busy":"2024-04-04T16:49:40.352599Z","iopub.execute_input":"2024-04-04T16:49:40.353142Z","iopub.status.idle":"2024-04-04T16:49:41.367550Z","shell.execute_reply.started":"2024-04-04T16:49:40.353098Z","shell.execute_reply":"2024-04-04T16:49:41.366439Z"},"trusted":true},"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, 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\n","metadata":{"execution":{"iopub.status.busy":"2024-04-04T16:49:41.368876Z","iopub.execute_input":"2024-04-04T16:49:41.369197Z","iopub.status.idle":"2024-04-04T16:49:41.381336Z","shell.execute_reply.started":"2024-04-04T16:49:41.369161Z","shell.execute_reply":"2024-04-04T16:49:41.379976Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torch_geometric.loader import DataLoader\n\n# Other setup remains the same as before\nproteins = ['sEH', 'BRD4', 'HSA']\nall_predictions = []\n\nfor protein in proteins:\n    print(f\"Training and predicting for {protein}\")\n    \n    # Create DataLoaders for the current protein\n    train_loader = DataLoader(featurized_data_train[protein], batch_size=32, shuffle=True)\n    test_loader = DataLoader(featurized_data_test[protein], batch_size=32, shuffle=False)\n    \n    # Train model\n    input_dim = train_loader.dataset[0].num_node_features\n    hidden_dim = 64\n    num_epochs = 11\n    num_layers = 4 #Should ideally be set so that all nodes can communicate with each other\n    dropout_rate = 0.3\n    lr = 0.001\n    #These are just example values, feel free to play around with them.\n    model = train_model(train_loader,num_epochs, input_dim, hidden_dim,num_layers, dropout_rate, lr)\n    \n    # Predict\n    molecule_ids, predictions = predict_with_model(model, test_loader)\n    \n    # Collect predictions\n    protein_predictions = pd.DataFrame({\n        'id': molecule_ids,\n        'binds': predictions,\n    })\n    all_predictions.append(protein_predictions)","metadata":{"execution":{"iopub.status.busy":"2024-04-04T16:53:20.923646Z","iopub.execute_input":"2024-04-04T16:53:20.924222Z","iopub.status.idle":"2024-04-04T17:00:42.413608Z","shell.execute_reply.started":"2024-04-04T16:53:20.924181Z","shell.execute_reply":"2024-04-04T17:00:42.412362Z"},"trusted":true},"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 directly within the DataFrame\nfinal_predictions['id'] = final_predictions['id'].apply(lambda x: x.item())\n\n# Now save the modified DataFrame to a CSV file\nfinal_predictions.to_csv('final_predictions.csv', index=False)","metadata":{"execution":{"iopub.status.busy":"2024-04-04T17:00:42.449662Z","iopub.execute_input":"2024-04-04T17:00:42.450643Z","iopub.status.idle":"2024-04-04T17:00:51.249295Z","shell.execute_reply.started":"2024-04-04T17:00:42.450587Z","shell.execute_reply":"2024-04-04T17:00:51.248206Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The final predictions csv will be our final output for the model.","metadata":{}},{"cell_type":"code","source":"final_predictions.head()","metadata":{"execution":{"iopub.status.busy":"2024-04-04T17:02:30.149440Z","iopub.execute_input":"2024-04-04T17:02:30.150801Z","iopub.status.idle":"2024-04-04T17:02:30.170506Z","shell.execute_reply.started":"2024-04-04T17:02:30.150727Z","shell.execute_reply":"2024-04-04T17:02:30.169354Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}