tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Quick start and overview of RNAglib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we describe the main object types of RNAglib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.1. RNADataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "RNADataset objects represent a set of RNAs, each one being represented by its 3D structure.\n",
    "\n",
    "Each item of the RNADataset is encoded by a dictionary containing (under the key \"rna\") the networkx Graph representing the RNA.\n",
    "\n",
    "It is also possible to add Representation and FeaturesComputer objects to a RNADataset.\n",
    "\n",
    "To create a default RNA Dataset, you can run the code below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Index file not found at /Users/wissam/.rnaglib/indexes/rnaglib-nr-1.0.0.json. Run rnaglib_index\n",
      "Database was found and not overwritten\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/anaconda3/envs/RNApython/lib/python3.11/site-packages/networkx/readwrite/json_graph/node_link.py:287: FutureWarning: \n",
      "The default value will be changed to `edges=\"edges\" in NetworkX 3.6.\n",
      "\n",
      "To make this warning go away, explicitly set the edges kwarg, e.g.:\n",
      "\n",
      "  nx.node_link_graph(data, edges=\"links\") to preserve current behavior, or\n",
      "  nx.node_link_graph(data, edges=\"edges\") for forward compatibility.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from rnaglib.data_loading import RNADataset\n",
    "\n",
    "dataset = RNADataset()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When calling the `__get_item__` method of a `RNADataset` object, which takes as argument the index of one RNA, the following steps happen:\n",
    "* If the dataset has `dataset.in_memory=False`, the graph of this RNA is loaded (otherwise, it has already been loaded)\n",
    "* A dictionary encoding the RNA called `rna_dict` is being built. This dictionary has 3 items: the graph of the RNA, the path of the graph and the path of the structures of the RNA\n",
    "* If some transforms to apply have been specified in `dataset.transforms`, then these transforms are being applied to the dataset.\n",
    "* The features dictionary of this RNA is being computed using the transform `dataset.features_computer` which is an attribute of the dataset and maps a dictionary of type `rna_dict` to a  dictionary of features.\n",
    "* Each representation associated with the dataset (that is to say contained in `dataset.representations`) is being applied to the considered RNA and appended to the dictionary `rna_dict`\n",
    "* The method returns the dictionary `rna_dict_` which contains the graph of the RNA (under the key `rna`), the path to the graph (under the key `graph_path`), the path to the RNA structures (under the key `cif_path`) and the RNA representations for each representation of  `dataset.representations` (under the keys corresponding to the representation names such as `graph` or `point_cloud`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'rna': <networkx.classes.digraph.DiGraph at 0x16a651350>,\n",
       " 'graph_path': PosixPath('/Users/wissam/.rnaglib/datasets/rnaglib-nr-1.0.0/graphs/1a9n.json'),\n",
       " 'cif_path': PosixPath('/Users/wissam/.rnaglib/datasets/rnaglib-nr-1.0.0/structures/1a9n.cif')}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.2. Transform\n",
    "\n",
    "The Transform class groups all the functions which map the dictionaries representing RNAs (i.e. the items of a RNADataset object) into other objects (other dictionqries or objects of a different nature).\n",
    "\n",
    "A specific tutorial gives further details about this class: https://rnaglib.org/en/latest/rnaglib.transforms.html\n",
    "\n",
    "Below are detailed some subclasses of Transform: Representation, FeaturesComputer, FilterTransform, AnnotationTransform, PartitionTransform, Compose and ComposeFilters.\n",
    "\n",
    "##### 1.2.1. Representation\n",
    "\n",
    "A Representation object is a Transform that maps a RNA dictionary (as defined above) to a mathematical representation of this RNA. In the current version of RNAGlib, 4 representations are already implemented: GraphRepresentation, PointCloudRepresentation, VoxelRepresentation and RingRepresentation\n",
    "\n",
    "GraphRepresentation converts RNA into a Leontis-Westhof graph (2.5D) where nodes are residues and edges are either base pairs or backbones.\n",
    "\n",
    "PointCloudRepresentation converts RNA into a 3D point cloud based representation.\n",
    "\n",
    "VoxelRepresentation converts RNA into a voxel/3D grid representation.\n",
    "\n",
    "RingRepresentation converts RNA into a ring-based representation.\n",
    "\n",
    "##### 1.2.2. FeaturesComputer\n",
    "\n",
    "A FeaturesComputer is a Transform that maps a RNA Dictionary to a dictionary of features and targets (both RNA-level and node-level features and targets) of this RNA.\n",
    "\n",
    "##### 1.2.3. FilterTransform\n",
    "\n",
    "A FilterTransform returns the RNAs of the RNADataset that pass a certain filter.\n",
    "\n",
    "##### 1.2.4. AnnotationTransform\n",
    "\n",
    "An AnnotationTransform is a transform computing additional node features within each RNA graph.\n",
    "\n",
    "##### 1.2.5. PartitionTransform\n",
    "\n",
    "A PartitionTransform is a transform which breaks up each RNA structure into substructures.\n",
    "\n",
    "##### 1.2.6. Compose\n",
    "\n",
    "A Compose object is a Transform which consists in the composition of a series of transforms.\n",
    "\n",
    "##### 1.2.7. ComposeFilters\n",
    "\n",
    "A ComposeFilters object is a Transform consisting in the composition of a series of filters (objects of type FilterTransform)\n",
    "\n",
    "#### 1.3. Tasks\n",
    "\n",
    "A Task is an object representing a benchmarking task to be performed on the RNA. It is associated with a specific RNADataset. Once implemented, the Task object can be called to evaluate the performance on the defined Task of various models. One particular category of Tasks is already implemented as a subclass of Tasks: ResidueClassificationTask, which groups all the tasks consisting in classifying the amino-acids of the RNA.\n",
    "\n",
    "#### 1.4. Encoders\n",
    "\n",
    "Encoders are objects that vectorize features with a specific encoding. Indeed, the features available in the RNA NetworkX graph might have different types, including text, therefore it is necessary to vectorize them to perform learning using them."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Using the tasks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, generate the necessary index files\n",
    "\n",
    "You can use the following command: \n",
    "```\n",
    "$ rnaglib_index\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.1. Instantiate the task"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Choose the task appropriate to your model. Here, we chose _RNA-Site_, a task instance called `BindingSiteDetection` for illustration.\n",
    "When instantiating the task, custom splitters or other arguments can be passed if needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating task dataset from scratch...\n",
      "Database was found and not overwritten\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/anaconda3/envs/RNApython/lib/python3.11/site-packages/networkx/readwrite/json_graph/node_link.py:287: FutureWarning: \n",
      "The default value will be changed to `edges=\"edges\" in NetworkX 3.6.\n",
      "\n",
      "To make this warning go away, explicitly set the edges kwarg, e.g.:\n",
      "\n",
      "  nx.node_link_graph(data, edges=\"links\") to preserve current behavior, or\n",
      "  nx.node_link_graph(data, edges=\"edges\") for forward compatibility.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ">>> Saving dataset.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/anaconda3/envs/RNApython/lib/python3.11/site-packages/networkx/readwrite/json_graph/node_link.py:142: FutureWarning: \n",
      "The default value will be `edges=\"edges\" in NetworkX 3.6.\n",
      "\n",
      "To make this warning go away, explicitly set the edges kwarg, e.g.:\n",
      "\n",
      "  nx.node_link_data(G, edges=\"links\") to preserve current behavior, or\n",
      "  nx.node_link_data(G, edges=\"edges\") for forward compatibility.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ">>> Done\n"
     ]
    }
   ],
   "source": [
    "from rnaglib.tasks import BindingSiteDetection\n",
    "from rnaglib.transforms import FeaturesComputer\n",
    "\n",
    "task = BindingSiteDetection(root=\"tutorial\", recompute = True) # set recompute=True to use the dataset designed for the Task, otherwise the dataset located at tutorial/dataset will be used"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.2. [In option] Customize the task dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 2.2.1. Apply already implemented transforms to the dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You might want to apply transforms or preprocessing to your dataset which is not implemented by default in the task. In this case, you can apply additional transforms to the dataset.\n",
    "\n",
    "We illustrate this below with the application of the transform PDBIDNameTransform."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rnaglib.data_loading import RNADataset\n",
    "from rnaglib.transforms import PDBIDNameTransform\n",
    "\n",
    "rnas = PDBIDNameTransform()(task.dataset)\n",
    "task.dataset = RNADataset(rnas=[r[\"rna\"] for r in rnas])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You might want to go further by implementing custom transforms to do preprocessing on the dataset. We show below how to create custom annotators and custom filters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 2.2.2. Create a custom annotator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You might want to create a custom annotator to add new features to the nodes of the graphs, for instance to perform a new task  using those new annotations. The custom annotator will typically be called in the `process` method of a `Task` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rnaglib.transforms import AnnotationTransform\n",
    "from networkx import set_node_attributes\n",
    "\n",
    "class CustomAnnotator(AnnotationTransform):\n",
    "    def forward (self, rna_dict: dict) -> dict:        \n",
    "        custom_annotation = {\n",
    "            node: self._custom_annotation(nodedata)\n",
    "            for node, nodedata in rna_dict['rna'].nodes(data=True)\n",
    "        }\n",
    "        set_node_attributes(rna_dict['rna'], custom_annotation, \"custom_annotation\")\n",
    "        return rna_dict\n",
    "    @staticmethod\n",
    "    def _has_binding_site(nodedata: dict) -> bool:\n",
    "        return ... # RNA dictionary-wise formula to compute the custom annotation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once defined, you can apply your custom annotator to the dataset using the following code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rnas = CustomAnnotator()(task.dataset)\n",
    "task.dataset = RNADataset(rnas=[r[\"rna\"] for r in rnas])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 2.2.3. Create a custom filter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Several filters are already implemented and available in `rnaglib.transforms`: `SizeFilter` which rejects RNAs which are not in the given size bounds, `RNAAttributeFilter` that rejects RNAs that lack a certain annotation at the whole RNA level, `ResidueAttributeFilter` which rejects RNAs that lack a certain annotation at the whole residue-level, `RibosomalFilter` that rejects ribsosomal RNA and `NameFilter` that filters RNA based on their names. However, you might want to create your own filter. This one could be for instance called in the `process` method of a new `Task` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rnaglib.transforms import FilterTransform\n",
    "\n",
    "class CustomFilter(FilterTransform):\n",
    "\n",
    "    def __init__(self, ..., **kwargs):\n",
    "        ...\n",
    "        super().__init__(**kwargs)\n",
    "\n",
    "    def forward(self, rna_dict: dict) -> bool:\n",
    "\n",
    "        ...\n",
    "\n",
    "        return ... # should return a Boolean (True if the RNA described by rna_dict passes the filter, False otherwise)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once defined, you can apply your custom annotator to the dataset using the following code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rnas = CustomFilter()(task.dataset)\n",
    "task.dataset = RNADataset(rnas=[r[\"rna\"] for r in rnas])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.3. [Optional] Customize the features"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 2.3.1. Add features from the graph to the dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You might want to use input features which are different from the default ones specified for this task in RNAglib. In this case, it is necessary to add it to the features computer of the RNA.\n",
    "\n",
    "The features can be chosen among the list of features available in the RNA graph: 'index', 'index_chain', 'chain_name', 'nt_resnum', 'nt_name', 'nt_code', 'nt_id', 'nt_type', 'dbn', 'summary', 'alpha', 'beta', 'gamma', 'delta', 'epsilon', 'zeta', 'epsilon_zeta', 'bb_type', 'chi', 'glyco_bond', 'C5prime_xyz', 'P_xyz', 'form', 'ssZp', 'Dp', 'splay_angle', 'splay_distance', 'splay_ratio', 'eta', 'theta', 'eta_prime', 'theta_prime', 'eta_base', 'theta_base', 'v0', 'v1', 'v2', 'v3', 'v4', 'amplitude', 'phase_angle', 'puckering', 'sugar_class', 'bin', 'cluster', 'suiteness', 'filter_rmsd', 'frame', 'sse', 'binding_protein', 'binding_ion', 'binding_small-molecule'.\n",
    "\n",
    "When adding a feature to the features computer, you have to specify a dictionary named `custom_encoders` mapping each feature to the encoder chosen to encode the feature. Canonical encoders corresponding to each feature are available in [NODE_FEATURE_MAP](https://github.com/cgoliver/rnaglib/blob/30bded91462f655c235ef57efc07e834456615a4/src/rnaglib/config/feature_encoders.py#L7)\n",
    "\n",
    "In the example below, we add the feature named `\"phase_angle\"` to the features computer of the dataset and specify that it should be encoded using the pre-implemented FloatEncoder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rnaglib.encoders import FloatEncoder\n",
    "\n",
    "task.dataset.features_computer.add_feature(feature_names=\"phase_angle\", custom_encoders={\"phase_angle\":FloatEncoder()})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 2.3.2. Create custom features"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The strategy to create custom features consists in creating a Transform object which takes as input a RNADataset and transforms it by adding the new features to the graphs representing all of the items of the RNADataset.\n",
    "\n",
    "To do so, you have to build a subclass of `Transform` and specify:\n",
    "\n",
    "* its `name`\n",
    "\n",
    "* its associated `encoder`\n",
    "\n",
    "* its `forward` method taking as input the dictionary representing one RNA and returning the updated RNA dictionary (containing its additional features)\n",
    "\n",
    "Once the custom features have been created, you still have to add them to the FeaturesComputer of the graph. To do so, you can check the documentation above (cf. section \"Adding features to the features computer of a RNADataset\").\n",
    "\n",
    "Below is the structure to write such a transform:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rnaglib.transforms import Transform\n",
    "\n",
    "class AddCustomFeature(Transform):\n",
    "    name = \"add_custom_feature\"\n",
    "    encoder = ...\n",
    "    def __init__(\n",
    "            self, **kwargs\n",
    "    ):\n",
    "        super().__init__(**kwargs)\n",
    "    def forward(self, rna_dict: Dict) -> Dict:\n",
    "\n",
    "        ... # compute and add additional features\n",
    "\n",
    "        return rna_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the transform above has been defined, there remains to apply it to the dataset as illustrated below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rnas = AddCustomFeature()(task.dataset)\n",
    "task.dataset = RNADataset(rnas=[r[\"rna\"] for r in rnas])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.4. Add a representation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 2.4.1. Add an already implemented representation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is necessary to add to the dataset a representation of the RNA structure. If the representation you want to add to perform the task is already implemented, you have to follow the code below. Already implemented representations include graphs (`GraphRepresentation` class), point clouds (`PointCloudRepresentation` class), voxels (`VoxelRepresentation` class) and rings (`RingRepresentation` class)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ">>> Adding <rnaglib.transforms.represent.graph.GraphRepresentation object at 0x3ad265390> representations.\n"
     ]
    }
   ],
   "source": [
    "from rnaglib.transforms import GraphRepresentation\n",
    "\n",
    "task.dataset.add_representation(GraphRepresentation(framework='pyg'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 2.4.2. Create a custom representation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, you might want to use a representation which doesn't belong to the aforementioned already implemented representations. In this case, you have to define your transformation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rnaglib.transforms import Representation\n",
    "\n",
    "class CustomRepresentation(Representation):\n",
    "    \"\"\"\n",
    "    Converts RNA into a custom representation.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self,\n",
    "                 **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        pass\n",
    "\n",
    "    def __call__(self, rna_graph, features_dict):\n",
    "\n",
    "        ... # computes the representation\n",
    "\n",
    "        return representation\n",
    "\n",
    "    @property\n",
    "    def name(self):\n",
    "        return \"custom_representation\" # the name of the representation\n",
    "\n",
    "    def batch(self, samples):\n",
    "\n",
    "        ... # defines the way to batch representations of different samples together\n",
    "\n",
    "        return batched_samples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the transformation has been defined, you have to add it to the dataset as in the case in which the representation has already been implemented"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "task.dataset.add_representation(CustomRepresentation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.5. Set loaders"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we changed the dataset by adding a representation (and maybe some additional features) to it, it is necessary to call `set_loaders` in order to update the train, val and test dataloaders."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "task.set_loaders()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.6. Build a model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first define the architecture of the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class LinearModel(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(LinearModel,self).__init__()\n",
    "        self.linear = torch.nn.Linear(5, 1)\n",
    "\n",
    "    def forward(self, data):  \n",
    "        x, edge_index, edge_type = data.x, data.edge_index, data.edge_attr\n",
    "        x = self.linear(x)\n",
    "        return F.sigmoid(x)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then instantiate the model and the optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.optim as optim\n",
    "\n",
    "# Define model\n",
    "learning_rate = 0.0001\n",
    "epochs = 2\n",
    "device = \"cpu\"\n",
    "model = LinearModel()\n",
    "model = model.to(device)\n",
    "optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
    "criterion = torch.nn.BCELoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.7. [Optional] Define your own `evaluate` method"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the representation you have chosen isn't the canonical representation for the task or if you want to get performace metrics which aren't implemented by default, you have to defined a custom `evaluate` method. Below is an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import accuracy_score, f1_score, roc_auc_score, matthews_corrcoef\n",
    "\n",
    "my_representation = \"graph\" # or \"point_cloud\" or \"custom_representation\"\n",
    "\n",
    "def evaluate(loader):\n",
    "    model.eval()\n",
    "    all_preds = []\n",
    "    all_labels = []\n",
    "    total_loss = 0\n",
    "\n",
    "    for batch in loader:\n",
    "        graph = batch[my_representation]\n",
    "        graph = graph.to(device)\n",
    "        out = model(graph)\n",
    "        loss = criterion(out, torch.flatten(graph.y).long())\n",
    "        total_loss += loss.item()\n",
    "        preds = out.argmax(dim=1)\n",
    "        all_preds.extend(preds.tolist())\n",
    "        all_labels.extend(graph.y.tolist())\n",
    "\n",
    "    avg_loss = total_loss / len(loader)\n",
    "\n",
    "    accuracy = accuracy_score(all_labels, all_preds)\n",
    "    f1 = f1_score(all_labels, all_preds)\n",
    "    auc = roc_auc_score(all_labels, all_preds)\n",
    "    mcc = matthews_corrcoef(all_labels, all_preds)\n",
    "\n",
    "    return accuracy, f1, auc, avg_loss, mcc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.8. Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, TrainAcc 0.7101 Val Acc: 0.7604\n",
      "Epoch 2, TrainAcc 0.7100 Val Acc: 0.7604\n"
     ]
    }
   ],
   "source": [
    "def train():\n",
    "    model.train()\n",
    "    for batch in task.train_dataloader:\n",
    "        graph = batch[\"graph\"]\n",
    "        graph = graph.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        out = model(graph)\n",
    "        loss = criterion(out, graph.y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    train()\n",
    "    train_metrics = task.evaluate(model, task.train_dataloader) # you might be using evaluate instead of task.evaluate\n",
    "    val_metrics = task.evaluate(model, task.val_dataloader) # you might be using evaluate instead of task.evaluate\n",
    "    print(\n",
    "    f\"\"\"Epoch {epoch + 1}, TrainAcc {train_metrics[\"accuracy\"]:.4f} Val Acc: {val_metrics[\"accuracy\"]:.4f}\"\"\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.9. Evaluate the model on the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Accuracy: 0.6861, Test F1 Score: 0.0005, Test AUC: 0.4964, Test MCC: -0.0068\n"
     ]
    }
   ],
   "source": [
    "test_metrics = task.evaluate(\n",
    "    model, task.test_dataloader, device\n",
    ") # you might be using evaluate instead of task.evaluate\n",
    "\n",
    "print(\n",
    "    f\"\"\"Test Accuracy: {test_metrics[\"accuracy\"]:.4f}, Test F1 Score: {test_metrics[\"f1\"]:.4f}, Test AUC: {test_metrics[\"auc\"]:.4f}, Test MCC: {test_metrics[\"mcc\"]:.4f}\"\"\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Creating custom tasks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create a custom Task\n",
    "\n",
    "In order to create a custom task, you have to define it as a subclass of a task category (for instance ResidueClassificationClass or a subclass you have created by yourself) and to specify the following:\n",
    "\n",
    "* a target variable: the variable which has to be predicted by the model\n",
    "* an input variable or a list of input variables: the inputs of the model\n",
    "* a method `get_tasks_var` specifying the FeaturesComputer to build to perform the task (in general, it will call the aforementioned target and input variables)\n",
    "* a method `process` creqting the dataset and applying some preprocessing to the dataset (especially annotation and filtering transforms) if needed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the task belongs to another task category than ResidueClassificationClass (that is to say, node-level classification task), you have to define a new Task subclass corresponding to this task category and to specify:\n",
    "* a method named `dummy_model` returning a dummy model to use to check the task is working well without any effort to define a model\n",
    "* a method named `evaluate` which, given a model, outputs a dictionary containing performace metrics of this model on the task of interest."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For instance, in the cell below, we define a toy task called AnglePrediction consisting in predicting the phase angle by using the nucleotide code.\n",
    "\n",
    "Since it is a regression and not a classification task, we first need to define a new subclass of Tasks class which we will call `ResidueRegression`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from rnaglib.tasks import Task\n",
    "from rnaglib.utils import DummyResidueModel\n",
    "from sklearn.metrics import root_mean_squared_error, mean_absolute_error\n",
    "\n",
    "class ResidueRegression(Task):\n",
    "    def __init__(self, root, splitter=None, **kwargs):\n",
    "        super().__init__(root=root, splitter=splitter, **kwargs)\n",
    "\n",
    "    @property\n",
    "    def dummy_model(self) -> torch.nn:\n",
    "        return DummyResidueModel()\n",
    "\n",
    "    def evaluate(self, model: torch.nn, device: str = \"cpu\") -> dict:\n",
    "        model.eval()\n",
    "        all_probs = []\n",
    "        all_preds = []\n",
    "        all_labels = []\n",
    "\n",
    "        with torch.no_grad():\n",
    "            for batch in self.test_dataloader:\n",
    "                graph = batch[\"graph\"]\n",
    "                graph = graph.to(device)\n",
    "                out = model(graph)\n",
    "\n",
    "                preds = out > 0.5\n",
    "                all_probs.extend(out.cpu().flatten().tolist())\n",
    "                all_preds.extend(preds.cpu().flatten().tolist())\n",
    "                all_labels.extend(graph.cpu().y.flatten().tolist())\n",
    "\n",
    "        # Compute performance metrics\n",
    "        RMSE = root_mean_squared_error(all_labels, all_preds)\n",
    "        MAE = mean_absolute_error(all_labels, preds)\n",
    "\n",
    "\n",
    "        return {\"RMSE\": RMSE, \"MAE\": MAE}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the subclass `ResidueRegression` is defined, one can define the specific task `AnglePrediction`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rnaglib.transforms import PDBIDNameTransform\n",
    "from rnaglib.encoders import BoolEncoder\n",
    "\n",
    "class AnglePrediction(ResidueRegression):\n",
    "    # Target variable\n",
    "    target_var = \"phase_angle\"\n",
    "    # Input variable\n",
    "    input_var = \"nt_code\"\n",
    "\n",
    "    def __init__(self, root, splitter=None, **kwargs):\n",
    "        super().__init__(root=root, splitter=splitter, **kwargs)\n",
    "        \n",
    "    # Creation and preprocessing of the dataset\n",
    "    def process(self) -> RNADataset:\n",
    "        rnas = RNADataset(debug=False, redundancy='all', rna_id_subset=SPLITTING_VARS['PDB_TO_CHAIN_TR60_TE18'].keys())\n",
    "        dataset = RNADataset(rnas=[r[\"rna\"] for r in rnas])\n",
    "        # TODO: remove wrong chains using  SPLITTING_VARS[\"PDB_TO_CHAIN_TR60_TE18\"]\n",
    "        rnas = PDBIDNameTransform()(rnas)\n",
    "        dataset = RNADataset(rnas=[r[\"rna\"] for r in rnas]) \n",
    "        return dataset\n",
    "    \n",
    "    # Computation of the FeaturesComputer\n",
    "    def get_task_vars(self) -> FeaturesComputer:\n",
    "        return FeaturesComputer(\n",
    "            nt_features=[self.input_var],\n",
    "            nt_targets=self.target_var,\n",
    "            custom_encoders={self.target_var: BoolEncoder()},\n",
    "        )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "RNApython",
   "language": "python",
   "name": "python3"
  },
  "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.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}