Binding pocket selection and prediction works

.gitignore

@@ -177,3 +177,4 @@ src/server/experiments/proteins/*
 src/server/experiments/conformations/*
 !/src/server/experiments/conformations/85e6d812-777c-4a79-97c6-054b7aa517f2
 src/server/experiments/drug_discovery/25befab3-cb27-4f04-a454-b966ebedbc0c/targets/
+src/ai/custom_models/drug_target/p2rank_2.4.1

src/ai/custom_models/drug_target/utils.py

@@ -6,6 +6,12 @@
 
 dirname = os.path.dirname
 
+import ssl
+
+# Create an unverified SSL context
+ssl._create_default_https_context = ssl._create_unverified_context
+
+
 def install_p2rank():
     # URL of the p2rank tar.gz file
     p2rank_url = "https://github.com/rdk/p2rank/releases/download/2.4.1/p2rank_2.4.1.tar.gz"

src/ai/model.py

@@ -1,4 +1,5 @@
 import shutil
+import subprocess
 
 import torch
 from torch import Tensor
@@ -11,6 +12,7 @@
 import requests
 from tqdm import tqdm
 import pickle
+from pathlib import Path
 
 from src.server import settings
 from src.server.services.progress import ProgressTracker
@@ -332,19 +334,41 @@ def load_model(self):
     # Method to get raw model outputs
     def _raw_inference(self, protein_file_path: str, save_dir: str):
         protein_filename = os.path.split(protein_file_path)[1]
-        destination_protein_file = os.path.join(save_dir, protein_filename)
-        shutil.copyfile(protein_file_path, destination_protein_file)
         ds = f"{save_dir}/protein_list.ds"
         with open(ds, "w") as out:
             out.write(f"{protein_filename}\n")
-        p2rank_exec = dirname(os.path.abspath(__file__)) + "/custom_models/drug_target/p2rank_2.4.1/prank"
-        p2rank = f"bash {p2rank_exec}"
-        cmd = f"{p2rank} predict {ds} -o {save_dir}/p2rank -threads 1"
-        os.system(cmd)
+        p2rank_exec = os.path.join(os.path.dirname(os.path.abspath(__file__)), "custom_models/drug_target/p2rank_2.4.1/prank")
+        cmd = ["bash", p2rank_exec, "predict", ds, "-o", f"{save_dir}/p2rank", "-threads", "1"]
 
+        # Run the command and wait for it to complete
+        subprocess.run(cmd, check=True)
+
+        p2rankFile = os.path.join(save_dir, 'p2rank', f"{Path(protein_filename).stem}.pdb_predictions.csv")
+        pocket = pd.read_csv(p2rankFile, skipinitialspace=True)
+
+        # Check if the dataframe is empty
+        if pocket.empty:
+            return np.array([])
+
+        print("POCKET DF: ", pocket.columns)
+
+        # Extract and process the 'residue_ids' column
+        # Splits the value by "_" and return only the integer on the right
+        residue_ids = pocket['residue_ids'].str.split()
+        all_ids = [int(item.split('_')[1]) for sublist in residue_ids for item in sublist]
+
+        pocket_ids = np.sort(np.array(all_ids))
+
+        np.save(os.path.join(save_dir, "pocket.npy"), pocket_ids)
+
+        pocket_ids = pocket_ids.tolist()
+
+        return pocket_ids
+
+
     # Method to return raw outputs in the desired format
     def predict(self, protein_file_path: str, save_dir: str):
-        self._raw_inference(protein_file_path, save_dir)
+        return self._raw_inference(protein_file_path, save_dir)
 
 
 
@@ -453,9 +477,9 @@ def load_model(self):
 
         pass
 
-    def _raw_inference(self, protein_ids, ligands, ligands_names, experiment_folder, target_positions, num_recycles):
+    def _raw_inference(self, protein_ids, ligands, ligands_names, experiment_folder, binding_pockets, num_recycles):
         experiment_progress_tracker = ProgressTracker(experiment_folder, protein_ids)
-        for ID in protein_ids:
+        for ID, binding_pocket in zip(protein_ids, binding_pockets):
             for LIGAND, LIGAND_NAME in zip(ligands, ligands_names):
                 protein_folder = os.path.join(experiment_folder, ID)
                 result_folder = os.path.join(protein_folder, 'result')
@@ -475,7 +499,7 @@ def _raw_inference(self, protein_ids, ligands, ligands_names, experiment_folder,
 
                 ID = ID.split('/')[-1]
                 # Predict
-                predict(config.CONFIG, MSA_FEATS, LIGAND_FEATS, ID, target_positions, PARAMS, num_recycles, outdir=result_folder)
+                predict(config.CONFIG, MSA_FEATS, LIGAND_FEATS, ID, binding_pocket, PARAMS, num_recycles, outdir=result_folder)
 
                 # Process the prediction
                 RAW_PDB = os.path.join(result_folder, f'{ID}_pred_raw.pdb')
@@ -507,19 +531,17 @@ def _raw_inference(self, protein_ids, ligands, ligands_names, experiment_folder,
             experiment_progress_tracker.update_progress(ID)
 
 
-    def predict(self, ligand_files_paths, protein_files, experiment_folder: str):
+    def predict(self, ligand_files_paths, protein_files, binding_pockets, experiment_folder: str):
         logger.info("Making dti predictions...")
         ligands_names, ligands_smiles = read_sdf_files(ligand_files_paths)
         protein_file_paths = [os.path.splitext(x)[0] for x in protein_files]
 
-        # TODO: ADD p2rank to identify target positions
-        target_array = np.asarray([])
         protein_names = [os.path.splitext(protein_file_path)[-2] for protein_file_path in protein_files]
         self.prepare_data(ligands_names, ligands_smiles, protein_file_paths, protein_files, experiment_folder)
         self._raw_inference(
             protein_ids=protein_names,
             ligands=ligands_smiles,
             ligands_names=ligands_names,
             experiment_folder=experiment_folder,
-            target_positions=target_array,
+            binding_pockets=binding_pockets,
             num_recycles=3)

src/server/api_handlers/drug_target.py

@@ -1,6 +1,7 @@
 import time
 
 from flask import Request, jsonify
+import numpy as np
 
 from src.server.services.sdf_pdb_combine import combine_sdf_pdb
 from src.server import settings
@@ -46,36 +47,42 @@ def load_targets(self, request):
         return {'id': experiment_id, 'name': experiment_name, 'targets': targets}
 
     def set_binding_pocket(self, request):
-        experiment_id = request.args.get('id')
-        protein_id = request.args.get('proteinId')
-        is_pocket_manual = request.args.get('isManualPocket')
+        data = request.json  # Access data sent in the request body
+        experiment_id = data.get('id')
+        protein_id = data.get('proteinId')
+        selected_residues = data.get('selectedResidues', [])
+
+        # Convert to NumPy array
+        selected_residues_array = np.array(selected_residues)
 
-        if is_pocket_manual:
-            pocket_sequence_ids = request.args.get('pocketIds')
-            self.experiments_loader.save_pocket(experiment_id,
-                                                protein_id,
-                                                is_pocket_manual,
-                                                pocket_sequence_ids)
-        else:
-            self.experiments_loader.save_pocket(experiment_id,
-                                                protein_id,
-                                                is_pocket_manual)
+        self.experiments_loader.set_binding_pocket(experiment_id, 
+                                                   protein_id,
+                                                   selected_residues_array)
+
+        return {}
 
 
     def get_binding_pocket(self, request):
         experiment_id = request.args.get('id')
         protein_id = request.args.get('proteinId')
-        is_pocket_manual = request.args.get('isManualPocket')
-
-        pocket_ids = []
 
-        pocket_ids = self.experiments_loader.load_pocket(experiment_id,
-                                                protein_id,
-                                                is_pocket_manual)
+        pocket_ids = self.experiments_loader.load_binding_pocket(experiment_id,
+                                                protein_id)
 
         return {"experimentId": experiment_id,
                 "proteinId": protein_id,
                 "pocketIds": pocket_ids}
+
+    def predict_binding_pocket(self, request):
+        experiment_id = request.args.get('id')
+        protein_id = request.args.get('proteinId')
+
+        pocket = self.drug_discovery.predict_pocket(
+                                                experiment_id=experiment_id,
+                                                protein_id=protein_id
+                                                )
+
+        return pocket
 
 
 

src/server/drug_target_blueprint.py

@@ -25,6 +25,10 @@ def store_binding_pocket():
     @drug_target_bp.route('/get-binding-pocket', methods=['GET'])
     def get_binding_pocket():
         return api_handler.get_binding_pocket(request)
+
+    @drug_target_bp.route('/predict-binding-pocket', methods=['GET'])
+    def predict_binding_pocket():
+        return api_handler.predict_binding_pocket(request)
 
     @drug_target_bp.route('/experiments')
     def get_experiments():

src/server/experiments/drug_discovery/25befab3-cb27-4f04-a454-b966ebedbc0c/metadata.json

@@ -1,5 +1,5 @@
 {
     "id": "25befab3-cb27-4f04-a454-b966ebedbc0c",
     "name": "brinzolamide - CARBONIC ANHYDRASE",
-    "date": "2023-12-25T19:57:58.717714"
+    "date": "2023-12-26T21:00:48.870574"
 }

src/server/frontend/src/components/AminoAcidLab/TargetProteinViewer.vue

@@ -1,6 +1,6 @@
 <script>
 export default {
-    props: ['target'],
+    props: ['api', 'experiment', 'target'],
     data() {
         const views = {
             default: { key: 'default', title: 'Default representation' },
@@ -26,7 +26,8 @@ export default {
             ligandComponent: null,
             selectedResidues: new Set(),
             sequence: [],
-            selectionMode: false
+            selectionMode: false,
+            showNoPocketModal: false,
         }
     },
     methods: {
@@ -93,8 +94,28 @@ export default {
         },
         sendSelectedResidues() {
             const selectedResidueArray = Array.from(this.selectedResidues);
-            console.log("Sending selected residues to backend:", selectedResidueArray);
-            // Replace with your API call to send data to the backend
+            this.api.setBindingPocket(this.experiment, this.target.metadata.id, selectedResidueArray);
+        },
+        async displayBindingPocket() {
+            const response = await this.api.loadBindingPocket(this.experiment, this.target.metadata.id);
+            if (response.data.pocketIds && response.data.pocketIds.length > 0) {
+                this.selectedResidues = new Set(response.data.pocketIds);
+                this.highlightSelectedResidues();
+            } else {
+                // Show modal if no binding pocket data is found
+                this.showNoPocketModal = true;
+            }
+        },
+        async predictBindingPocket() {
+            this.isLoading = true; // Set loading state
+            const response = await this.api.predictBindingPocket(this.experiment, this.target.metadata.id);
+            this.isLoading = false; // Reset loading state
+            this.showNoPocketModal = false; // Close the modal
+
+            if (response.data && response.data.length > 0) {
+                this.selectedResidues = new Set(response.data);
+                this.highlightSelectedResidues();
+            }
         },
         extractSequenceFromComponent(component) {
             // Extracting amino acid sequence from the PDB component
@@ -172,6 +193,41 @@ export default {
                 <button v-if="this.selectionMode" type="button" @click="sendSelectedResidues" class="btn btn-success" style="margin-top: 10px;">
                     Confirm selection
                 </button>
+                <button v-if="!this.selectionMode" type="button" @click="displayBindingPocket" class="btn btn-info" style="margin-top: 10px;">
+                    Display Binding Pocket
+                </button>
+            </div>
+        </div>
+        <div v-if="isLoading" class="loading-indicator">
+            <div class="spinner-border" role="status">
+                <span class="visually-hidden">Predicting...</span>
+            </div>
+            <span>Predicting...</span>
+        </div>
+        <div v-if="showNoPocketModal" class="modal" tabindex="-1" role="dialog">
+            <div class="modal-dialog" role="document">
+                <div class="modal-content">
+                    <div class="modal-header">
+                        <h5 class="modal-title">Binding Pocket Not Found</h5>
+                        <button type="button" class="close" data-dismiss="modal" aria-label="Close" @click="showNoPocketModal = false">
+                            <span aria-hidden="true">&times;</span>
+                        </button>
+                    </div>
+                    <div v-if="isLoading" class="loading-indicator">
+                        <div class="spinner-border" role="status">
+                            <span class="visually-hidden">Predicting...</span>
+                        </div>
+                        <span>Predicting...</span>
+                    </div>
+                    <div class="modal-body" style="color: black;">
+                        <p>You don't have a binding pocket for this protein.</p>
+                        <p>Do you wish it to be predicted?</p>
+                    </div>
+                    <div class="modal-footer">
+                        <button type="button" class="btn btn-primary" @click="predictBindingPocket">Predict</button>
+                        <button type="button" class="btn btn-secondary" data-dismiss="modal" @click="showNoPocketModal = false">Close</button>
+                    </div>
+                </div>
             </div>
         </div>
         <div class="row">
@@ -203,4 +259,11 @@ export default {
     background-color: blue;
     color: white;
 }
+
+.loading-indicator {
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    padding: 10px;
+}
 </style>

src/server/frontend/src/components/AminoAcidLab/TargetsList.vue

@@ -110,7 +110,7 @@ export default {
 
         <!-- Modal -->
         <div v-if="isModalOpen" class="modal">
-            <TargetProteinViewer :target="selectedTarget" @close="closeModal" />
+            <TargetProteinViewer :api="this.api" :experiment="this.experiment" :target="selectedTarget" @close="closeModal" />
         </div>
     </div>
 </template>

src/server/frontend/src/storage.js

@@ -102,6 +102,9 @@ const store = createStore({
             //state.drugTarget.experiment = experiment;
             state.drugTarget.experiment.targets = data.targets;
         },
+        drugTarget_setBindingPocket (state, { data }) {
+            const experimentsArray = [];
+        },
         drugTarget_inference(state, experiment) {
             state.drugTarget.experiment = experiment;
         },
@@ -262,6 +265,17 @@ const store = createStore({
             const response = await axios.get(apiConstants.drugTarget.loadTargets.path, { params: { id: experiment.metaData.id, name: experiment.metaData.name } });
             commit(apiConstants.drugTarget.loadTargets.mutation, { data: response.data });
         },
+        async drugTarget_setBindingPocket({ commit }, { experiment, proteinId, selectedResidues }){
+            const payload = {
+                id: experiment.metaData.id,
+                name: experiment.metaData.name,
+                proteinId: proteinId,
+                selectedResidues: Array.from(selectedResidues),
+                isPocketManual: true
+            };
+            const response = await axios.post(apiConstants.drugTarget.setBindingPocket.path, payload);
+            commit(apiConstants.drugTarget.loadTargets.mutation, { data: response.data });
+        },
         async drugTarget_inference({ commit }, { payload }) {
             const { form, experiment } = payload;
             const formData = new FormData(form);
@@ -412,6 +426,15 @@ export const api = {
         loadTargets: async (experiment) => {
             return await store.dispatch(apiConstants.drugTarget.loadTargets.action, { experiment });
         },
+        setBindingPocket: async (experiment, proteinId, selectedResidues) => {
+            return await store.dispatch(apiConstants.drugTarget.setBindingPocket.action, { experiment, proteinId, selectedResidues });
+        },
+        loadBindingPocket: async (experiment, proteinId) => {
+            return await axios.get(apiConstants.drugTarget.loadBindingPocket.path, { params: { id: experiment.metaData.id, proteinId: proteinId } });
+        },
+        predictBindingPocket: async (experiment, proteinId) => {
+            return await axios.get(apiConstants.drugTarget.predictBindingPocket.path, { params: { id: experiment.metaData.id, proteinId: proteinId } });
+        },
         inference: async (payload) => {
             return await store.dispatch(apiConstants.drugTarget.inference.action, { payload });
         },

src/server/frontend/src/storageConstants.js

@@ -79,6 +79,17 @@ export const apiConstants = {
             mutation: 'drugTarget_loadTargets',
             action: 'drugTarget_loadTargets'
         },
+        setBindingPocket: {
+            path: baseUrl + '/api/drug-target/set-binding-pocket',
+            mutation: 'drugTarget_setBindingPocket',
+            action: 'drugTarget_setBindingPocket'
+        },
+        loadBindingPocket: {
+            path: baseUrl + '/api/drug-target/get-binding-pocket'
+        },
+        predictBindingPocket: {
+            path: baseUrl + '/api/drug-target/predict-binding-pocket'
+        },
         inference: { // request: {name: 'exp1', sdf: 'sdf', pdb: 'pdb'} response: {exp, data: []}
             path: baseUrl + '/api/drug-target/inference',
             mutation: 'drugTarget_inference',