{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# A quick GNN baseline\n\nIn this notebook, I put together a very basic GNN starter. \nNone of the below in necessarily optimal (it's basically what I was able to throw together on a Sunday afternoon), but should be enough to get started for further experimentation. I wrote this on my PC, so this notebook seems to be a little bit bottlenecked by the low amount of CPU cores, but I'm sure it's possible to work around that. The below doesn't actually do very well, but most of the work was the dataloader. Maybe you can do better? \n\n## Graph neural networks \n\n[Graph neural networks](https://en.wikipedia.org/wiki/Graph_neural_network) (GNNs) are a class of neural networks that operate on graph data structures. Graphs can be used to describe many objects, such as social networks, road networks, or in this case, molecules. \n\nUnlike sequences and grids (images) which have strict definitions of adjacency, graphs have a more flexible definition of adjacency, that is defined by the edges of the graph. \n\n## RNA as graphs\n\nThere is no canonical way to represent an RNA molecule as a graph, and I believe the \"trick\" will be constructing the adjacency. In the case below, the base pairs are simply connected to the neares nodes (defined by `EDGE_DISTANCE`) on each side. If we had the 3D structure, that could be a good value to use for connections as well. From the info we do have, the pairing probablility might also work. \n\n# PyG\n\n[PyG](https://pyg.org/) (AKA pytorch geometric) is a library that contains all you need to get going with learning on graphs. \n\nBelow I've used the PyG dataset class to create a dataset, and imported a simple GNN, called EdgeCNN, to begin with. ","metadata":{"execution":{"iopub.status.busy":"2023-09-25T11:02:15.908105Z","iopub.execute_input":"2023-09-25T11:02:15.908370Z"}}},{"cell_type":"code","source":"!pip install torch_geometric","metadata":{"execution":{"iopub.status.busy":"2023-09-26T08:53:21.486203Z","iopub.execute_input":"2023-09-26T08:53:21.486719Z","iopub.status.idle":"2023-09-26T08:53:51.043260Z","shell.execute_reply.started":"2023-09-26T08:53:21.486683Z","shell.execute_reply":"2023-09-26T08:53:51.042168Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nfrom torch.utils.data import random_split\nfrom torch_geometric.data import Data, Dataset\nfrom torch_geometric.loader import DataLoader\nimport pandas as pd\nfrom pathlib import Path\nimport numpy as np\nfrom sklearn.preprocessing import OneHotEncoder\nimport polars as pl\nimport re\nfrom tqdm import tqdm","metadata":{"execution":{"iopub.status.busy":"2023-09-26T08:53:51.045861Z","iopub.execute_input":"2023-09-26T08:53:51.046242Z","iopub.status.idle":"2023-09-26T08:53:55.328187Z","shell.execute_reply.started":"2023-09-26T08:53:51.046205Z","shell.execute_reply":"2023-09-26T08:53:55.327226Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DATA_DIR = Path(\"/kaggle/input/stanford-ribonanza-rna-folding/\")\nTRAIN_CSV = DATA_DIR / \"train_data.csv\"\nTRAIN_PARQUET_FILE = \"train_data.parquet\"\nTEST_CSV = DATA_DIR / \"test_sequences.csv\"\nTEST_PARQUET_FILE = \"test_sequences.parquet\"\nPRED_CSV = \"submission.csv\"","metadata":{"execution":{"iopub.status.busy":"2023-09-26T08:53:55.329443Z","iopub.execute_input":"2023-09-26T08:53:55.330069Z","iopub.status.idle":"2023-09-26T08:53:55.336988Z","shell.execute_reply.started":"2023-09-26T08:53:55.330037Z","shell.execute_reply":"2023-09-26T08:53:55.336094Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Convert to Parquet\n\nPandas is too slow, so the below converts the training and testing data to parquet. ","metadata":{}},{"cell_type":"code","source":"def to_parquet(csv_file, parquet_file):\n    dummy_df = pl.scan_csv(csv_file)\n\n    new_schema = {}\n    for key, value in dummy_df.schema.items():\n        if key.startswith(\"reactivity\"):\n            new_schema[key] = pl.Float32\n        else:\n            new_schema[key] = value\n\n    df = pl.scan_csv(csv_file, schema=new_schema)\n    \n    df.sink_parquet(\n            parquet_file,\n            compression='uncompressed',\n            row_group_size=10,\n    )","metadata":{"execution":{"iopub.status.busy":"2023-09-26T08:53:55.339652Z","iopub.execute_input":"2023-09-26T08:53:55.340229Z","iopub.status.idle":"2023-09-26T08:53:55.355368Z","shell.execute_reply.started":"2023-09-26T08:53:55.340198Z","shell.execute_reply":"2023-09-26T08:53:55.354418Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"to_parquet(TRAIN_CSV, TRAIN_PARQUET_FILE)\nto_parquet(TEST_CSV, TEST_PARQUET_FILE)","metadata":{"execution":{"iopub.status.busy":"2023-09-26T08:53:55.356814Z","iopub.execute_input":"2023-09-26T08:53:55.357204Z","iopub.status.idle":"2023-09-26T08:56:29.852587Z","shell.execute_reply.started":"2023-09-26T08:53:55.357173Z","shell.execute_reply":"2023-09-26T08:56:29.851558Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Define adjacency\n\nThere are countless ways you can define adjacency. \nThe function below creates an edge connection array that specifies how the edges are connected. \n\nIn this case, an edge is connected to the neighbouring `n` molecules to each side of it, plus optionally itself. ","metadata":{}},{"cell_type":"code","source":"def nearest_adjacency(sequence_length, n=2, loops=True):\n    base = np.arange(sequence_length)\n    connections = []\n    for i in range(-n, n + 1):\n        if i == 0 and not loops:\n            continue\n        elif i == 0 and loops:\n            stack = np.vstack([base, base])\n            connections.append(stack)\n            continue\n\n        neighbours = base.take(range(i,sequence_length+i), mode='wrap')\n        stack = np.vstack([base, neighbours])\n        \n        if i < 0:\n            connections.append(stack[:, -i:])\n        elif i > 0:\n            connections.append(stack[:, :-i])\n\n    return np.hstack(connections)","metadata":{"execution":{"iopub.status.busy":"2023-09-26T08:56:29.854194Z","iopub.execute_input":"2023-09-26T08:56:29.854884Z","iopub.status.idle":"2023-09-26T08:56:29.863171Z","shell.execute_reply.started":"2023-09-26T08:56:29.854842Z","shell.execute_reply":"2023-09-26T08:56:29.862169Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"EDGE_DISTANCE = 4","metadata":{"execution":{"iopub.status.busy":"2023-09-26T08:56:29.864913Z","iopub.execute_input":"2023-09-26T08:56:29.865682Z","iopub.status.idle":"2023-09-26T08:56:29.875799Z","shell.execute_reply.started":"2023-09-26T08:56:29.865649Z","shell.execute_reply":"2023-09-26T08:56:29.874817Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Defining the dataloader. \n\nThe below defines a simple dataloader that parses a parquet file into node embeddings (for now just the one hot encoded bases A, G, U and C), the adjacency (using the function above), and the targets (the reactivity scores). ","metadata":{}},{"cell_type":"code","source":"class SimpleGraphDataset(Dataset):\n    def __init__(self, parquet_name, edge_distance=5, root=None, transform=None, pre_transform=None, pre_filter=None):\n        super().__init__(root, transform, pre_transform, pre_filter)\n        # Set csv name\n        self.parquet_name = parquet_name\n        # Set edge distance\n        self.edge_distance = edge_distance\n        # Initialize one hot encoder\n        self.node_encoder = OneHotEncoder(sparse_output=False, max_categories=5)\n        # For one-hot encoder to possible values\n        self.node_encoder.fit(np.array(['A', 'G', 'U', 'C']).reshape(-1, 1))\n        # Load dataframe\n        self.df = pl.read_parquet(self.parquet_name)\n        self.df = self.df.filter(pl.col(\"SN_filter\") == 1.0)\n        # Get reactivity columns names\n        reactivity_match = re.compile('(reactivity_[0-9])')\n        reactivity_names = [col for col in self.df.columns if reactivity_match.match(col)]\n        \n        self.reactivity_df = self.df.select(reactivity_names) \n\n        self.sequence_df = self.df.select(\"sequence\")\n\n    def parse_row(self, idx):\n        # Read row at idx\n        sequence_row = self.sequence_df.row(idx)  \n        reactivity_row = self.reactivity_df.row(idx)\n        # Get sequence string and convert to array\n        sequence = np.array(list(sequence_row[0])).reshape(-1, 1)\n        # Encode sequence array\n        encoded_sequence = self.node_encoder.transform(sequence)\n        # Get sequence length\n        sequence_length = len(sequence)\n        # Get edge index \n        edges_np = nearest_adjacency(sequence_length, n=self.edge_distance, loops=False)\n        # Convert to torch tensor\n        edge_index = torch.tensor(edges_np, dtype=torch.long)\n\n        # Get reactivity targets for nodes\n        reactivity = np.array(reactivity_row, dtype=np.float32)[0:sequence_length]\n     \n        # Create valid masks for nodes\n        valid_mask = np.argwhere(~np.isnan(reactivity)).reshape(-1)\n        torch_valid_mask = torch.tensor(valid_mask, dtype=torch.long)\n\n        # Replace nan values for reactivity with 0. \n        # Not actually super important as they get masked\n        reactivity = np.nan_to_num(reactivity, copy=False, nan=0.0)\n\n\n        # Define node features as one-hot encoded sequence\n        node_features = torch.Tensor(encoded_sequence)\n\n        # Targets \n        targets = torch.Tensor(reactivity)\n\n        data = Data(x=node_features, edge_index=edge_index, y=targets, valid_mask=torch_valid_mask)\n\n        return data\n\n    def len(self):\n        return len(self.df)\n\n    def get(self, idx):\n        data = self.parse_row(idx) \n        return data","metadata":{"execution":{"iopub.status.busy":"2023-09-26T08:56:29.877416Z","iopub.execute_input":"2023-09-26T08:56:29.878179Z","iopub.status.idle":"2023-09-26T08:56:29.893019Z","shell.execute_reply.started":"2023-09-26T08:56:29.878148Z","shell.execute_reply":"2023-09-26T08:56:29.892088Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Data handling\n\nDefine the train and validation datasets, and load them with dataloaders. ","metadata":{}},{"cell_type":"code","source":"full_train_dataset = SimpleGraphDataset(parquet_name=TRAIN_PARQUET_FILE, edge_distance=EDGE_DISTANCE)\ngenerator1 = torch.Generator().manual_seed(42)\ntrain_dataset, val_dataset = random_split(full_train_dataset, [0.7, 0.3], generator1)","metadata":{"execution":{"iopub.status.busy":"2023-09-26T08:56:29.894126Z","iopub.execute_input":"2023-09-26T08:56:29.895633Z","iopub.status.idle":"2023-09-26T08:56:45.733375Z","shell.execute_reply.started":"2023-09-26T08:56:29.895581Z","shell.execute_reply":"2023-09-26T08:56:45.732381Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataloader = DataLoader(train_dataset, batch_size=256, shuffle=True, num_workers=2)\nval_dataloader = DataLoader(val_dataset, batch_size=256, shuffle=False, num_workers=2)","metadata":{"execution":{"iopub.status.busy":"2023-09-26T09:13:36.238816Z","iopub.execute_input":"2023-09-26T09:13:36.239222Z","iopub.status.idle":"2023-09-26T09:13:36.245336Z","shell.execute_reply.started":"2023-09-26T09:13:36.239191Z","shell.execute_reply":"2023-09-26T09:13:36.244144Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Loss and metrics\n\nDefine the loss function and the MSE, and define the MAE as a loss. \n\nThe target values are clipped to `(0, 1)` as the competition metric is. ","metadata":{}},{"cell_type":"code","source":"import torch.nn.functional as F\n\ndef loss_fn(output, target):\n    clipped_target = torch.clip(target, min=0, max=1)\n    mses = F.mse_loss(output, clipped_target, reduction='mean')\n    return mses\n\ndef mae_fn(output, target):\n    clipped_target = torch.clip(target, min=0, max=1)\n    maes = F.l1_loss(output, clipped_target, reduction='mean')\n    return maes","metadata":{"execution":{"iopub.status.busy":"2023-09-26T09:13:37.232743Z","iopub.execute_input":"2023-09-26T09:13:37.233297Z","iopub.status.idle":"2023-09-26T09:13:37.244304Z","shell.execute_reply.started":"2023-09-26T09:13:37.233258Z","shell.execute_reply":"2023-09-26T09:13:37.243270Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Define the model\n\nBelow we define as simple EdgeCNN from PyG. \nAs a start, we give it 128 hidden channels, and 4 layers. ","metadata":{}},{"cell_type":"code","source":"from torch_geometric.nn.models import EdgeCNN\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nmodel = EdgeCNN(in_channels=full_train_dataset.num_features, hidden_channels=128,\n                num_layers=4, out_channels=1).to(device)","metadata":{"execution":{"iopub.status.busy":"2023-09-26T09:13:38.253171Z","iopub.execute_input":"2023-09-26T09:13:38.253510Z","iopub.status.idle":"2023-09-26T09:13:38.272437Z","shell.execute_reply.started":"2023-09-26T09:13:38.253484Z","shell.execute_reply":"2023-09-26T09:13:38.271562Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Make sure we are using the GPU\ndevice","metadata":{"execution":{"iopub.status.busy":"2023-09-26T09:13:39.021544Z","iopub.execute_input":"2023-09-26T09:13:39.021883Z","iopub.status.idle":"2023-09-26T09:13:39.028219Z","shell.execute_reply.started":"2023-09-26T09:13:39.021856Z","shell.execute_reply":"2023-09-26T09:13:39.027338Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training \n\nTrain the model for 10 epochs. \nIs this a good learning rate? ","metadata":{}},{"cell_type":"code","source":"n_epochs = 10\n\n\noptimizer = torch.optim.Adam(model.parameters(), lr=0.0003, weight_decay=5e-4)\n\n\nfor epoch in range(n_epochs):\n    train_losses = []\n    train_maes = []\n    model.train()\n    for batch in (pbar := tqdm(train_dataloader, position=0, leave=True)):\n        batch = batch.to(device)\n        optimizer.zero_grad()\n        out = model(batch.x, batch.edge_index)\n        out = torch.squeeze(out)\n        loss = loss_fn(out[batch.valid_mask], batch.y[batch.valid_mask])\n        mae = mae_fn(out[batch.valid_mask], batch.y[batch.valid_mask])\n        loss.backward()\n        train_losses.append(loss.detach().cpu().numpy())\n        train_maes.append(mae.detach().cpu().numpy())\n        optimizer.step()\n        pbar.set_description(f\"Train loss {loss.detach().cpu().numpy():.4f}\")       \n\n    print(f\"Epoch {epoch} train loss: \", np.mean(train_losses))    \n    print(f\"Epoch {epoch} train mae: \", np.mean(train_maes))    \n    \n    val_losses = []\n    val_maes = []\n    model.eval()\n    for batch in (pbar := tqdm(val_dataloader, position=0, leave=True)):\n        batch = batch.to(device)\n        optimizer.zero_grad()\n        out = model(batch.x, batch.edge_index)\n        out = torch.squeeze(out)\n        loss = loss_fn(out[batch.valid_mask], batch.y[batch.valid_mask])\n        mae = mae_fn(out[batch.valid_mask], batch.y[batch.valid_mask])\n        val_losses.append(loss.detach().cpu().numpy())\n        val_maes.append(mae.detach().cpu().numpy())\n        pbar.set_description(f\"Validation loss {loss.detach().cpu().numpy():.4f}\")      \n        \n    print(f\"Epoch {epoch} val loss: \", np.mean(val_losses)) \n    print(f\"Epoch {epoch} val mae: \", np.mean(val_maes))\n\n","metadata":{"execution":{"iopub.status.busy":"2023-09-26T09:13:41.178440Z","iopub.execute_input":"2023-09-26T09:13:41.179133Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch.cuda.empty_cache()","metadata":{"execution":{"iopub.status.busy":"2023-09-26T09:12:08.293755Z","iopub.execute_input":"2023-09-26T09:12:08.294128Z","iopub.status.idle":"2023-09-26T09:12:08.427688Z","shell.execute_reply.started":"2023-09-26T09:12:08.294080Z","shell.execute_reply":"2023-09-26T09:12:08.426649Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Inference on the test set\n\nHere we define a light weight dataset to handle the inference step. ","metadata":{}},{"cell_type":"code","source":"class InferenceGraphDataset(Dataset):\n    def __init__(self, parquet_name, edge_distance=2, root=None, transform=None, pre_transform=None, pre_filter=None):\n        super().__init__(root, transform, pre_transform, pre_filter)\n        # Set csv name\n        self.parquet_name = parquet_name\n        # Set edge distance\n        self.edge_distance = edge_distance\n        # Initialize one hot encoder\n        self.node_encoder = OneHotEncoder(sparse_output=False, max_categories=4)\n        # For one-hot encoder to possible values\n        self.node_encoder.fit(np.array(['A', 'G', 'U', 'C']).reshape(-1, 1))\n        # Load dataframe\n        self.df = pl.read_parquet(self.parquet_name)\n\n        self.sequence_df = self.df.select(\"sequence\")\n        self.id_min_df = self.df.select(\"id_min\")\n\n    def parse_row(self, idx):\n        # Read row at idx\n        sequence_row = self.sequence_df.row(idx)  \n        id_min = self.id_min_df.row(idx)[0]\n\n        # Get sequence string and convert to array\n        sequence = np.array(list(sequence_row[0])).reshape(-1, 1)\n        # Encode sequence array\n        encoded_sequence = self.node_encoder.transform(sequence)\n        # Get sequence length\n        sequence_length = len(sequence)\n        # Get edge index \n        edges_np = nearest_adjacency(sequence_length, n=self.edge_distance, loops=False)\n        # Convert to torch tensor\n        edge_index = torch.tensor(edges_np, dtype=torch.long)\n\n        # Define node features as one-hot encoded sequence\n        node_features = torch.Tensor(encoded_sequence)\n        ids = torch.arange(id_min, id_min+sequence_length, 1)\n\n        data = Data(x=node_features, edge_index=edge_index, ids=ids)\n\n        return data\n\n    def len(self):\n        return len(self.df)\n\n    def get(self, idx):\n        data = self.parse_row(idx) \n        return data","metadata":{"execution":{"iopub.status.busy":"2023-09-26T09:04:58.407134Z","iopub.execute_input":"2023-09-26T09:04:58.407835Z","iopub.status.idle":"2023-09-26T09:04:58.420656Z","shell.execute_reply.started":"2023-09-26T09:04:58.407803Z","shell.execute_reply":"2023-09-26T09:04:58.419603Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"infer_dataset = InferenceGraphDataset(parquet_name=TEST_PARQUET_FILE, edge_distance=EDGE_DISTANCE)\ninfer_dataloader = DataLoader(infer_dataset, batch_size=128, shuffle=False, num_workers=2)","metadata":{"execution":{"iopub.status.busy":"2023-09-26T09:12:21.067676Z","iopub.execute_input":"2023-09-26T09:12:21.068044Z","iopub.status.idle":"2023-09-26T09:12:21.765431Z","shell.execute_reply.started":"2023-09-26T09:12:21.068013Z","shell.execute_reply":"2023-09-26T09:12:21.764406Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nmodel.eval().to(device)","metadata":{"execution":{"iopub.status.busy":"2023-09-26T09:12:23.137596Z","iopub.execute_input":"2023-09-26T09:12:23.137944Z","iopub.status.idle":"2023-09-26T09:12:23.146916Z","shell.execute_reply.started":"2023-09-26T09:12:23.137915Z","shell.execute_reply":"2023-09-26T09:12:23.145749Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ids = np.empty(shape = (0,1),dtype=int)\npreds = np.empty(shape = (0,1),dtype=np.float32)\n\n\nfor batch in tqdm(infer_dataloader):\n    batch = batch.to(device)\n    out = model(batch.x, batch.edge_index).detach().cpu().numpy()\n\n    ids = np.append(ids, batch.ids.detach().cpu().numpy())\n    preds = np.append(preds, out)\n\n","metadata":{"execution":{"iopub.status.busy":"2023-09-26T09:12:26.829330Z","iopub.execute_input":"2023-09-26T09:12:26.829689Z","iopub.status.idle":"2023-09-26T09:13:26.763263Z","shell.execute_reply.started":"2023-09-26T09:12:26.829661Z","shell.execute_reply":"2023-09-26T09:13:26.762057Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Submission\n\nCreate a csv file with the submission values. \nAs you can see, I don't currently distinguish between `DMS_MaP` and `2A3_MaP`, so just write the same value to both. ","metadata":{}},{"cell_type":"code","source":"submission_df = pl.DataFrame({\"id\": ids, \"reactivity_DMS_MaP\": preds, \"reactivity_2A3_MaP\": preds})","metadata":{"execution":{"iopub.status.busy":"2023-09-26T09:11:59.624775Z","iopub.status.idle":"2023-09-26T09:11:59.625516Z","shell.execute_reply.started":"2023-09-26T09:11:59.625265Z","shell.execute_reply":"2023-09-26T09:11:59.625288Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission_df.write_csv(PRED_CSV)","metadata":{"execution":{"iopub.status.busy":"2023-09-26T09:11:59.626859Z","iopub.status.idle":"2023-09-26T09:11:59.627576Z","shell.execute_reply.started":"2023-09-26T09:11:59.627322Z","shell.execute_reply":"2023-09-26T09:11:59.627344Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Conclusion\n\nThis is a basic GNN \"starter kit\" that does the basics. \nThere are many things that can be improved, but I hope this helps people get started. ","metadata":{}}]}