Added support for listing available dataset data

README.md

@@ -24,12 +24,16 @@ The code for data curation can be found in the Jupyter notebook [`data_curation.
 
 ## 🚀 Installation
 
-To install the package, open your terminal and run the following command:
+To install the package, open your terminal and run the following commands:
 
 ```bash
+git clone https://github.com/ribesstefano/PROTAC-Degradation-Predictor.git
+cd PROTAC-Degradation-Predictor
 pip install .
 ```
 
+The package has been developed on a Linux machine with Python 3.10.8. It is recommended to use a virtual environment to avoid conflicts with other packages.
+
 ## 🎯 Usage
 
 After installing the package, you can use it as follows:
@@ -56,10 +60,30 @@ print(f'The given PROTAC is: {"active" if active_protac else "inactive"}')
 
 This example demonstrates how to predict the activity of a PROTAC molecule. The `is_protac_active` function takes the SMILES string of the PROTAC, the E3 ligase, the UniProt ID of the target protein, and the cell line as inputs. It returns whether the PROTAC is active or not.
 
+The function supports batch computation by passing lists of SMILES strings, E3 ligases, UniProt IDs, and cell lines. In this case, it returns a list of booleans indicating the activity of each PROTAC.
+
+
+
+
 ## 📈 Training
 
 The code for training the model can be found in the file [`run_experiments.py`](src/run_experiments.py).
 
+## 📄 Citation
+
+If you use this tool in your research, please cite the following paper:
+
+```
+@misc{ribes2024modeling,
+    title={Modeling PROTAC Degradation Activity with Machine Learning},
+    author={Stefano Ribes and Eva Nittinger and Christian Tyrchan and Rocío Mercado},
+    year={2024},
+    eprint={2406.02637},
+    archivePrefix={arXiv},
+    primaryClass={q-bio.QM}
+}
+```
+
 ## 📜 License
 
 This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.

protac_degradation_predictor/__init__.py

@@ -3,6 +3,10 @@ from .data_utils import (
     load_cell2embedding,
     get_fingerprint,
     is_active,
+    load_curated_dataset,
+    avail_cell_lines,
+    avail_e3_ligases,
+    avail_uniprots,
 )
 from .protac_dataset import (
     PROTAC_Dataset,

protac_degradation_predictor/data_utils.py

@@ -1,7 +1,7 @@
 import os
 import pkg_resources
 import pickle
-from typing import Dict, Optional
+from typing import Dict, Optional, List
 
 from .config import config
 
@@ -61,6 +61,33 @@ def load_cell2embedding(
     return cell2embedding
 
 
+def avail_e3_ligases() -> List[str]:
+    """ Get the available E3 ligases.
+    
+    Returns:
+        List[str]: The available E3 ligases.
+    """
+    return list(config.e3_ligase2uniprot.keys())
+
+
+def avail_cell_lines() -> List[str]:
+    """ Get the available cell lines.
+    
+    Returns:
+        List[str]: The available cell lines.
+    """
+    return list(load_cell2embedding().keys())
+
+
+def avail_uniprots() -> List[str]:
+    """ Get the available Uniprot IDs.
+    
+    Returns:
+        List[str]: The available Uniprot IDs.
+    """
+    return list(load_protein2embedding().keys())
+
+
 def get_fingerprint(smiles: str, morgan_fpgen = None) -> np.ndarray:
     """ Get the Morgan fingerprint of a molecule.
     
@@ -113,4 +140,15 @@ def is_active(
         if pd.notnull(pDC50) and pd.notnull(Dmax):
             return True if pDC50 >= pDC50_threshold and Dmax >= Dmax_threshold else False
         else:
-            return np.nan
+            return np.nan
+
+
+def load_curated_dataset() -> pd.DataFrame:
+    """ Load the curated PROTAC dataset as described in the paper: https://arxiv.org/abs/2406.02637
+
+    Returns:
+        pd.DataFrame: The curated PROTAC dataset.
+    """
+    with pkg_resources.resource_stream(__name__, 'data/PROTAC-Degradation-DB.csv') as f:
+        protac_df = pd.read_csv(f)
+    return protac_df

protac_degradation_predictor/protac_degradation_predictor.py

@@ -93,12 +93,18 @@ def get_protac_active_proba(
             prescaled_embeddings=False, # Normalization performed by the model
         )
         preds[ckpt_path] = sigmoid(pred).detach().cpu().numpy().flatten()
+
     # NOTE: The predictions array has shape: (n_models, batch_size)
     preds = np.array(list(preds.values()))
+    mean_preds = np.mean(preds, axis=0)
+    # Return a single value if not list as input
+    preds = preds if isinstance(protac_smiles, list) else preds[0]
+    means_preds = mean_preds if isinstance(protac_smiles, list) else mean_preds[0]
+    
     return {
         'preds': preds,
-        'mean': np.mean(preds, axis=0),
-        'majority_vote': np.mean(preds, axis=0) > 0.5,
+        'mean': mean_preds,
+        'majority_vote': mean_preds > 0.5,
     }