remove vina

app.py

@@ -77,12 +77,10 @@ class NpEncoder(json.JSONEncoder):
         return super(NpEncoder, self).default(obj)
     
 
-def evaluate(value: str, out_fn: str):
-    protein_path, ligand_path = load(value)
-
-    metrics = Metrics(protein_path, ligand_path, out_fn).evaluate()
-
-    return json.dumps(metrics, indent=4, cls=NpEncoder)
+# def evaluate(value: str, out_fn: str):
+#     protein_path, ligand_path = load(value)
+#     metrics = Metrics(protein_path, ligand_path, out_fn).evaluate()
+#     return json.dumps(metrics, indent=4, cls=NpEncoder)
 
 
 
@@ -104,9 +102,9 @@ with gr.Blocks() as demo:
     btn3 = gr.Button('visualize')
     btn3.click(show, inputs=[dropdown, candidates], outputs=[gen_complex])
 
-    metrics = gr.Textbox(label='metrics')
-    btn4 = gr.Button('evaluate')
-    btn4.click(evaluate, inputs=[dropdown, candidates], outputs=[metrics])
+    # metrics = gr.Textbox(label='metrics')
+    # btn4 = gr.Button('evaluate')
+    # btn4.click(evaluate, inputs=[dropdown, candidates], outputs=[metrics])
 
 if __name__ == '__main__':
     demo.launch(share=True)

core/datasets/__init__.py

@@ -1,6 +1,6 @@
 import torch
 from torch.utils.data import Subset
-from core.datasets.pl_pair_dataset import PocketLigandPairDataset, PocketLigandPairDatasetFeaturized, PocketLigandGeneratedPairDataset
+# from core.datasets.pl_pair_dataset import PocketLigandPairDataset, PocketLigandPairDatasetFeaturized, PocketLigandGeneratedPairDataset
 from core.datasets.pdbbind import PDBBindDataset
 
 

core/datasets/pl_pair_dataset.py

@@ -1,452 +1,452 @@
-import os
-import pickle
-import lmdb
-from torch.utils.data import Dataset
-from tqdm.auto import tqdm
-import sys
-from time import time
-
-import torch
-from torch_geometric.transforms import Compose
-
-from core.datasets.utils import PDBProtein, parse_sdf_file, ATOM_FAMILIES_ID
-from core.datasets.pl_data import ProteinLigandData, torchify_dict
-
-import core.utils.transforms as trans
-
-
-class DBReader:
-    def __init__(self, path) -> None:
-        self.path = path
-        self.db = None
-        self.keys = None
-
-    def _connect_db(self):
-        """
-            Establish read-only database connection
-        """
-        assert self.db is None, 'A connection has already been opened.'
-        self.db = lmdb.open(
-            self.path,
-            map_size=10*(1024*1024*1024),   # 10GB
-            create=False,
-            subdir=False,
-            readonly=True,
-            lock=False,
-            readahead=False,
-            meminit=False,
-        )
-        with self.db.begin() as txn:
-            self.keys = list(txn.cursor().iternext(values=False))
-
-    def _close_db(self):
-        self.db.close()
-        self.db = None
-        self.keys = None
-
-    def __del__(self):
-        if self.db is not None:
-            self._close_db()
-
-    def __len__(self):
-        if self.db is None:
-            self._connect_db()
-        return len(self.keys)
-
-    def __getitem__(self, idx):
-        if self.db is None:
-            self._connect_db()
-        key = self.keys[idx]
-        data = pickle.loads(self.db.begin().get(key))
-        data = ProteinLigandData(**data)
-        data.id = idx
-        assert data.protein_pos.size(0) > 0
-        return data
+# import os
+# import pickle
+# import lmdb
+# from torch.utils.data import Dataset
+# from tqdm.auto import tqdm
+# import sys
+# from time import time
+
+# import torch
+# from torch_geometric.transforms import Compose
+
+# from core.datasets.utils import PDBProtein, parse_sdf_file, ATOM_FAMILIES_ID
+# from core.datasets.pl_data import ProteinLigandData, torchify_dict
+
+# import core.utils.transforms as trans
+
+
+# class DBReader:
+#     def __init__(self, path) -> None:
+#         self.path = path
+#         self.db = None
+#         self.keys = None
+
+#     def _connect_db(self):
+#         """
+#             Establish read-only database connection
+#         """
+#         assert self.db is None, 'A connection has already been opened.'
+#         self.db = lmdb.open(
+#             self.path,
+#             map_size=10*(1024*1024*1024),   # 10GB
+#             create=False,
+#             subdir=False,
+#             readonly=True,
+#             lock=False,
+#             readahead=False,
+#             meminit=False,
+#         )
+#         with self.db.begin() as txn:
+#             self.keys = list(txn.cursor().iternext(values=False))
+
+#     def _close_db(self):
+#         self.db.close()
+#         self.db = None
+#         self.keys = None
+
+#     def __del__(self):
+#         if self.db is not None:
+#             self._close_db()
+
+#     def __len__(self):
+#         if self.db is None:
+#             self._connect_db()
+#         return len(self.keys)
+
+#     def __getitem__(self, idx):
+#         if self.db is None:
+#             self._connect_db()
+#         key = self.keys[idx]
+#         data = pickle.loads(self.db.begin().get(key))
+#         data = ProteinLigandData(**data)
+#         data.id = idx
+#         assert data.protein_pos.size(0) > 0
+#         return data
     
 
-class PocketLigandPairDataset(Dataset):
+# class PocketLigandPairDataset(Dataset):
 
-    def __init__(self, raw_path, transform=None, version='final'):
-        super().__init__()
-        self.raw_path = raw_path.rstrip('/')
-        self.index_path = os.path.join(self.raw_path, 'index.pkl')
-        self.processed_path = os.path.join(os.path.dirname(self.raw_path),
-                                           os.path.basename(self.raw_path) + f'_processed_{version}.lmdb')
-        self.transform = transform
-        self.reader = DBReader(self.processed_path)
+#     def __init__(self, raw_path, transform=None, version='final'):
+#         super().__init__()
+#         self.raw_path = raw_path.rstrip('/')
+#         self.index_path = os.path.join(self.raw_path, 'index.pkl')
+#         self.processed_path = os.path.join(os.path.dirname(self.raw_path),
+#                                            os.path.basename(self.raw_path) + f'_processed_{version}.lmdb')
+#         self.transform = transform
+#         self.reader = DBReader(self.processed_path)
 
-        if not os.path.exists(self.processed_path):
-            print(f'{self.processed_path} does not exist, begin processing data')
-            self._process()
+#         if not os.path.exists(self.processed_path):
+#             print(f'{self.processed_path} does not exist, begin processing data')
+#             self._process()
         
-    def _process(self):
-        db = lmdb.open(
-            self.processed_path,
-            map_size=10*(1024*1024*1024),   # 10GB
-            create=True,
-            subdir=False,
-            readonly=False,  # Writable
-        )
-        with open(self.index_path, 'rb') as f:
-            index = pickle.load(f)
-
-        num_skipped = 0
-        with db.begin(write=True, buffers=True) as txn:
-            for i, (pocket_fn, ligand_fn, *_) in enumerate(tqdm(index)):
-                if pocket_fn is None: continue
-                try:
-                    # data_prefix = '/data/work/jiaqi/binding_affinity'
-                    data_prefix = self.raw_path
-                    pocket_dict = PDBProtein(os.path.join(data_prefix, pocket_fn)).to_dict_atom()
-                    ligand_dict = parse_sdf_file(os.path.join(data_prefix, ligand_fn))
-                    data = ProteinLigandData.from_protein_ligand_dicts(
-                        protein_dict=torchify_dict(pocket_dict),
-                        ligand_dict=torchify_dict(ligand_dict),
-                    )
-                    data.protein_filename = pocket_fn
-                    data.ligand_filename = ligand_fn
-                    data = data.to_dict()  # avoid torch_geometric version issue
-                    txn.put(
-                        key=str(i).encode(),
-                        value=pickle.dumps(data)
-                    )
-                except:
-                    num_skipped += 1
-                    print('Skipping (%d) %s' % (num_skipped, ligand_fn, ))
-                    continue
-        db.close()
+#     def _process(self):
+#         db = lmdb.open(
+#             self.processed_path,
+#             map_size=10*(1024*1024*1024),   # 10GB
+#             create=True,
+#             subdir=False,
+#             readonly=False,  # Writable
+#         )
+#         with open(self.index_path, 'rb') as f:
+#             index = pickle.load(f)
+
+#         num_skipped = 0
+#         with db.begin(write=True, buffers=True) as txn:
+#             for i, (pocket_fn, ligand_fn, *_) in enumerate(tqdm(index)):
+#                 if pocket_fn is None: continue
+#                 try:
+#                     # data_prefix = '/data/work/jiaqi/binding_affinity'
+#                     data_prefix = self.raw_path
+#                     pocket_dict = PDBProtein(os.path.join(data_prefix, pocket_fn)).to_dict_atom()
+#                     ligand_dict = parse_sdf_file(os.path.join(data_prefix, ligand_fn))
+#                     data = ProteinLigandData.from_protein_ligand_dicts(
+#                         protein_dict=torchify_dict(pocket_dict),
+#                         ligand_dict=torchify_dict(ligand_dict),
+#                     )
+#                     data.protein_filename = pocket_fn
+#                     data.ligand_filename = ligand_fn
+#                     data = data.to_dict()  # avoid torch_geometric version issue
+#                     txn.put(
+#                         key=str(i).encode(),
+#                         value=pickle.dumps(data)
+#                     )
+#                 except:
+#                     num_skipped += 1
+#                     print('Skipping (%d) %s' % (num_skipped, ligand_fn, ))
+#                     continue
+#         db.close()
     
-    def __len__(self):
-        return len(self.reader)
-
-    def __getitem__(self, idx):
-        data = self.reader[idx]
-        if self.transform is not None:
-            data = self.transform(data)
-        return data
-
-
-class PocketLigandGeneratedPairDataset(Dataset):
-
-    def __init__(self, raw_path, transform=None, version='4-decompdiff'):
-        super().__init__()
-        self.raw_path = raw_path.rstrip('/')
-        self.generated_path = os.path.join('/sharefs/share/sbdd_data/all_results', f'{version}_docked_pose_checked.pt')
-        self.processed_path = os.path.join(os.path.dirname(self.raw_path),
-                                           os.path.basename(self.raw_path) + f'_processed_{version}.lmdb')
-        self.transform = transform
-        self.reader = DBReader(self.processed_path)
-
-        if not os.path.exists(self.processed_path):
-            print(f'{self.processed_path} does not exist, begin processing data')
-            self._process()
+#     def __len__(self):
+#         return len(self.reader)
+
+#     def __getitem__(self, idx):
+#         data = self.reader[idx]
+#         if self.transform is not None:
+#             data = self.transform(data)
+#         return data
+
+
+# class PocketLigandGeneratedPairDataset(Dataset):
+
+#     def __init__(self, raw_path, transform=None, version='4-decompdiff'):
+#         super().__init__()
+#         self.raw_path = raw_path.rstrip('/')
+#         self.generated_path = os.path.join('/sharefs/share/sbdd_data/all_results', f'{version}_docked_pose_checked.pt')
+#         self.processed_path = os.path.join(os.path.dirname(self.raw_path),
+#                                            os.path.basename(self.raw_path) + f'_processed_{version}.lmdb')
+#         self.transform = transform
+#         self.reader = DBReader(self.processed_path)
+
+#         if not os.path.exists(self.processed_path):
+#             print(f'{self.processed_path} does not exist, begin processing data')
+#             self._process()
         
-    def _process(self):
-        db = lmdb.open(
-            self.processed_path,
-            map_size=10*(1024*1024*1024),   # 10GB
-            create=True,
-            subdir=False,
-            readonly=False,  # Writable
-        )
-        with open(self.generated_path, 'rb') as f:
-            results = torch.load(f)
-
-        num_skipped = 0
-        with db.begin(write=True, buffers=True) as txn:
-            idx = -1
-            for i, res in tqdm(enumerate(results), total=len(results)):
-                if isinstance(res, dict):
-                    res = [res]
-                for r in res:
-                    idx += 1
-                    mol = r["mol"]
-                    ligand_fn = r["ligand_filename"]
-                    pocket_fn = os.path.join(
-                        os.path.dirname(ligand_fn),
-                        os.path.basename(ligand_fn)[:-4] + '_pocket10.pdb'
-                    )
-
-                    if pocket_fn is None: continue
-                    try:
-                        data_prefix = self.raw_path
-                        pocket_dict = PDBProtein(os.path.join(data_prefix, pocket_fn)).to_dict_atom()
-                        ligand_dict = parse_sdf_file(mol)
-                        # ligand_dict = parse_sdf_file(os.path.join(data_prefix, ligand_fn))
-                        data = ProteinLigandData.from_protein_ligand_dicts(
-                            protein_dict=torchify_dict(pocket_dict),
-                            ligand_dict=torchify_dict(ligand_dict),
-                        )
-                        data.protein_filename = pocket_fn
-                        data.ligand_filename = ligand_fn
-                        data = data.to_dict()  # avoid torch_geometric version issue
-                        txn.put(
-                            key=str(idx).encode(),
-                            value=pickle.dumps(data)
-                        )
-                    except Exception as e:
-                        num_skipped += 1
-                        print('Skipping (%d) %s' % (num_skipped, ligand_fn, ), e)
-                        continue
-        db.close()
+#     def _process(self):
+#         db = lmdb.open(
+#             self.processed_path,
+#             map_size=10*(1024*1024*1024),   # 10GB
+#             create=True,
+#             subdir=False,
+#             readonly=False,  # Writable
+#         )
+#         with open(self.generated_path, 'rb') as f:
+#             results = torch.load(f)
+
+#         num_skipped = 0
+#         with db.begin(write=True, buffers=True) as txn:
+#             idx = -1
+#             for i, res in tqdm(enumerate(results), total=len(results)):
+#                 if isinstance(res, dict):
+#                     res = [res]
+#                 for r in res:
+#                     idx += 1
+#                     mol = r["mol"]
+#                     ligand_fn = r["ligand_filename"]
+#                     pocket_fn = os.path.join(
+#                         os.path.dirname(ligand_fn),
+#                         os.path.basename(ligand_fn)[:-4] + '_pocket10.pdb'
+#                     )
+
+#                     if pocket_fn is None: continue
+#                     try:
+#                         data_prefix = self.raw_path
+#                         pocket_dict = PDBProtein(os.path.join(data_prefix, pocket_fn)).to_dict_atom()
+#                         ligand_dict = parse_sdf_file(mol)
+#                         # ligand_dict = parse_sdf_file(os.path.join(data_prefix, ligand_fn))
+#                         data = ProteinLigandData.from_protein_ligand_dicts(
+#                             protein_dict=torchify_dict(pocket_dict),
+#                             ligand_dict=torchify_dict(ligand_dict),
+#                         )
+#                         data.protein_filename = pocket_fn
+#                         data.ligand_filename = ligand_fn
+#                         data = data.to_dict()  # avoid torch_geometric version issue
+#                         txn.put(
+#                             key=str(idx).encode(),
+#                             value=pickle.dumps(data)
+#                         )
+#                     except Exception as e:
+#                         num_skipped += 1
+#                         print('Skipping (%d) %s' % (num_skipped, ligand_fn, ), e)
+#                         continue
+#         db.close()
     
-    def __len__(self):
-        return len(self.reader)
-
-    def __getitem__(self, idx):
-        data = self.reader[idx]
-        if self.transform is not None:
-            data = self.transform(data)
-        return data
-
-
-class PocketLigandPairDatasetFromComplex(Dataset):
-    def __init__(self, raw_path, transform=None, version='final', radius=10.0):
-        super().__init__()
-        self.raw_path = raw_path.rstrip('/')
-        self.index_path = os.path.join(self.raw_path, 'index.pkl')
-        base_name = os.path.basename(self.raw_path)
-        if 'pocket' in base_name:
-            self.processed_path = os.path.join(os.path.dirname(self.raw_path),
-                                               os.path.basename(self.raw_path) + f'_processed_{version}.lmdb')
-        else:
-            self.processed_path = os.path.join(os.path.dirname(self.raw_path),
-                                            os.path.basename(self.raw_path) + f'_pocket{radius}_processed_{version}.lmdb')
-        self.transform = transform
-        self.reader = DBReader(self.processed_path)
-
-        self.radius = radius
-
-        if not os.path.exists(self.processed_path):
-            print(f'{self.processed_path} does not exist, begin processing data')
-            self._process()
-
-    def _process(self):
-        db = lmdb.open(
-            self.processed_path,
-            map_size=10*(1024*1024*1024),   # 50GB
-            create=True,
-            subdir=False,
-            readonly=False,  # Writable
-            max_readers=256,
-        )
-        with open(self.index_path, 'rb') as f:
-            index = pickle.load(f)
-
-        print('Processing data...', 'index', self.index_path, index[0])
-
-        num_skipped = 0
-        with db.begin(write=True, buffers=True) as txn:
-            for i, (pocket_fn, ligand_fn, *_) in enumerate(tqdm(index)):
-                if pocket_fn is None: continue
-                try:
-                    data_prefix = self.raw_path
-                    # clip pocket
-                    ligand_dict = parse_sdf_file(os.path.join(data_prefix, ligand_fn))
-                    protein = PDBProtein(os.path.join(data_prefix, pocket_fn))
-                    selected = protein.query_residues_ligand(ligand_dict, self.radius)
-                    pdb_block_pocket = protein.residues_to_pdb_block(selected)
-                    pocket_dict = PDBProtein(pdb_block_pocket).to_dict_atom()
-
-                    # pocket_dict = PDBProtein(os.path.join(data_prefix, pocket_fn)).to_dict_atom()
-                    # ligand_dict = parse_sdf_file(os.path.join(data_prefix, ligand_fn))
+#     def __len__(self):
+#         return len(self.reader)
+
+#     def __getitem__(self, idx):
+#         data = self.reader[idx]
+#         if self.transform is not None:
+#             data = self.transform(data)
+#         return data
+
+
+# class PocketLigandPairDatasetFromComplex(Dataset):
+#     def __init__(self, raw_path, transform=None, version='final', radius=10.0):
+#         super().__init__()
+#         self.raw_path = raw_path.rstrip('/')
+#         self.index_path = os.path.join(self.raw_path, 'index.pkl')
+#         base_name = os.path.basename(self.raw_path)
+#         if 'pocket' in base_name:
+#             self.processed_path = os.path.join(os.path.dirname(self.raw_path),
+#                                                os.path.basename(self.raw_path) + f'_processed_{version}.lmdb')
+#         else:
+#             self.processed_path = os.path.join(os.path.dirname(self.raw_path),
+#                                             os.path.basename(self.raw_path) + f'_pocket{radius}_processed_{version}.lmdb')
+#         self.transform = transform
+#         self.reader = DBReader(self.processed_path)
+
+#         self.radius = radius
+
+#         if not os.path.exists(self.processed_path):
+#             print(f'{self.processed_path} does not exist, begin processing data')
+#             self._process()
+
+#     def _process(self):
+#         db = lmdb.open(
+#             self.processed_path,
+#             map_size=10*(1024*1024*1024),   # 50GB
+#             create=True,
+#             subdir=False,
+#             readonly=False,  # Writable
+#             max_readers=256,
+#         )
+#         with open(self.index_path, 'rb') as f:
+#             index = pickle.load(f)
+
+#         print('Processing data...', 'index', self.index_path, index[0])
+
+#         num_skipped = 0
+#         with db.begin(write=True, buffers=True) as txn:
+#             for i, (pocket_fn, ligand_fn, *_) in enumerate(tqdm(index)):
+#                 if pocket_fn is None: continue
+#                 try:
+#                     data_prefix = self.raw_path
+#                     # clip pocket
+#                     ligand_dict = parse_sdf_file(os.path.join(data_prefix, ligand_fn))
+#                     protein = PDBProtein(os.path.join(data_prefix, pocket_fn))
+#                     selected = protein.query_residues_ligand(ligand_dict, self.radius)
+#                     pdb_block_pocket = protein.residues_to_pdb_block(selected)
+#                     pocket_dict = PDBProtein(pdb_block_pocket).to_dict_atom()
+
+#                     # pocket_dict = PDBProtein(os.path.join(data_prefix, pocket_fn)).to_dict_atom()
+#                     # ligand_dict = parse_sdf_file(os.path.join(data_prefix, ligand_fn))
                     
-                    data = ProteinLigandData.from_protein_ligand_dicts(
-                        protein_dict=torchify_dict(pocket_dict),
-                        ligand_dict=torchify_dict(ligand_dict),
-                    )
-                    data.protein_filename = pocket_fn
-                    data.ligand_filename = ligand_fn
-                    data = data.to_dict()  # avoid torch_geometric version issue
-                    txn.put(
-                        key=str(i).encode(),
-                        value=pickle.dumps(data)
-                    )
-                except Exception as e:
-                    num_skipped += 1
-                    print('Skipping (%d) %s' % (num_skipped, ligand_fn, ), e)
-                    with open('skipped.txt', 'a') as f:
-                        f.write('Skip %s due to %s\n' % (ligand_fn, e))
-                    continue
-        db.close()
-
-    def __len__(self):
-        return len(self.reader)
-
-    def __getitem__(self, idx):
-        data = self.reader[idx]
-        if self.transform is not None:
-            data = self.transform(data)
-        return data
+#                     data = ProteinLigandData.from_protein_ligand_dicts(
+#                         protein_dict=torchify_dict(pocket_dict),
+#                         ligand_dict=torchify_dict(ligand_dict),
+#                     )
+#                     data.protein_filename = pocket_fn
+#                     data.ligand_filename = ligand_fn
+#                     data = data.to_dict()  # avoid torch_geometric version issue
+#                     txn.put(
+#                         key=str(i).encode(),
+#                         value=pickle.dumps(data)
+#                     )
+#                 except Exception as e:
+#                     num_skipped += 1
+#                     print('Skipping (%d) %s' % (num_skipped, ligand_fn, ), e)
+#                     with open('skipped.txt', 'a') as f:
+#                         f.write('Skip %s due to %s\n' % (ligand_fn, e))
+#                     continue
+#         db.close()
+
+#     def __len__(self):
+#         return len(self.reader)
+
+#     def __getitem__(self, idx):
+#         data = self.reader[idx]
+#         if self.transform is not None:
+#             data = self.transform(data)
+#         return data
     
 
-class PocketLigandPairDatasetFeaturized(Dataset):
-    def __init__(self, raw_path, ligand_atom_mode, version='simple'):
-        """
-        in simple version, only these features are saved for better IO:
-            protein_pos, protein_atom_feature, protein_element, 
-            ligand_pos, ligand_atom_feature_full, ligand_element
-        """
-        self.raw_path = raw_path
-        self.ligand_atom_mode = ligand_atom_mode
-        self.version = version
-
-        if version == 'simple':
-            self.features_to_save = [
-                'protein_pos', 'protein_atom_feature', 'protein_element', 
-                'ligand_pos', 'ligand_atom_feature_full', 'ligand_element',
-                'protein_filename', 'ligand_filename',
-            ]
-        else:
-            raise NotImplementedError
-
-        self.transformed_path = os.path.join(
-            os.path.dirname(self.raw_path), os.path.basename(self.raw_path) + 
-            f'_{ligand_atom_mode}_transformed_{version}.pt'
-        )
-        if not os.path.exists(self.transformed_path):
-            print(f'{self.transformed_path} does not exist, begin transforming data')
-            self._transform()
-        else:
-            print(f'reading data from {self.transformed_path}...')
-            tic = time()
-            tr_data = torch.load(self.transformed_path)
-            toc = time()
-            print(f'{toc - tic} elapsed')
-            self.train_data, self.test_data = tr_data['train'], tr_data['test']
-            self.protein_atom_feature_dim = tr_data['protein_atom_feature_dim']
-            self.ligand_atom_feature_dim = tr_data['ligand_atom_feature_dim']
+# class PocketLigandPairDatasetFeaturized(Dataset):
+#     def __init__(self, raw_path, ligand_atom_mode, version='simple'):
+#         """
+#         in simple version, only these features are saved for better IO:
+#             protein_pos, protein_atom_feature, protein_element, 
+#             ligand_pos, ligand_atom_feature_full, ligand_element
+#         """
+#         self.raw_path = raw_path
+#         self.ligand_atom_mode = ligand_atom_mode
+#         self.version = version
+
+#         if version == 'simple':
+#             self.features_to_save = [
+#                 'protein_pos', 'protein_atom_feature', 'protein_element', 
+#                 'ligand_pos', 'ligand_atom_feature_full', 'ligand_element',
+#                 'protein_filename', 'ligand_filename',
+#             ]
+#         else:
+#             raise NotImplementedError
+
+#         self.transformed_path = os.path.join(
+#             os.path.dirname(self.raw_path), os.path.basename(self.raw_path) + 
+#             f'_{ligand_atom_mode}_transformed_{version}.pt'
+#         )
+#         if not os.path.exists(self.transformed_path):
+#             print(f'{self.transformed_path} does not exist, begin transforming data')
+#             self._transform()
+#         else:
+#             print(f'reading data from {self.transformed_path}...')
+#             tic = time()
+#             tr_data = torch.load(self.transformed_path)
+#             toc = time()
+#             print(f'{toc - tic} elapsed')
+#             self.train_data, self.test_data = tr_data['train'], tr_data['test']
+#             self.protein_atom_feature_dim = tr_data['protein_atom_feature_dim']
+#             self.ligand_atom_feature_dim = tr_data['ligand_atom_feature_dim']
         
-    def _transform(self):
-        raw_dataset = PocketLigandPairDataset(self.raw_path, None, 'final')
-
-        split_path = os.path.join(
-            os.path.dirname(self.raw_path), 'crossdocked_pocket10_pose_split.pt',
-        )
-        split = torch.load(split_path)
-        train_ids, test_ids = split['train'], split['test']
-        print(f'train_size: {len(train_ids)}, test_size: {len(test_ids)}')
-
-        protein_featurizer = trans.FeaturizeProteinAtom()
-        ligand_featurizer = trans.FeaturizeLigandAtom(self.ligand_atom_mode)
-        transform_list = [
-            protein_featurizer,
-            ligand_featurizer,
-            # trans.FeaturizeLigandBond(),
-        ]
-        transform = Compose(transform_list)
-        self.protein_atom_feature_dim = protein_featurizer.feature_dim
-        self.ligand_atom_feature_dim = ligand_featurizer.feature_dim
-
-        def _transform_subset(ids):
-            data_list = []
-
-            for idx in tqdm(ids):
-                data = raw_dataset[idx]
-                data = transform(data)
-                tr_data = {}
-                for k in self.features_to_save:
-                    tr_data[k] = getattr(data, k)
-                tr_data['id'] = idx
-                tr_data = ProteinLigandData(**tr_data)
-                data_list.append(tr_data)
-            return data_list
-
-        self.train_data = _transform_subset(train_ids)
-        print(f'train_size: {len(self.train_data)}, {sys.getsizeof(self.train_data)}')
-        self.test_data = _transform_subset(test_ids)
-        print(f'test_size: {len(self.test_data)}, {sys.getsizeof(self.test_data)}')
-        torch.save({
-            'train': self.train_data, 'test': self.test_data,
-            'protein_atom_feature_dim': self.protein_atom_feature_dim,
-            'ligand_atom_feature_dim': self.ligand_atom_feature_dim,
-        }, self.transformed_path)
-
-
-if __name__ == '__main__':
-    # original dataset
-    dataset = PocketLigandPairDataset('./data/crossdocked_v1.1_rmsd1.0_pocket10')
-    print(len(dataset), dataset[0])
-
-    ############################################################
-
-    # test DecompDiffDataset
-    # dataset = PocketLigandGeneratedPairDataset('/sharefs/share/sbdd_data/crossdocked_pocket10')
-    # print(len(dataset), dataset[0])
-
-    ############################################################
-
-    # test featurized dataset (GPU accelerated)
-    # path = '/sharefs/share/sbdd_data/crossdocked_v1.1_rmsd1.0_pocket10'
-    # ligand_atom_mode = 'add_aromatic'
-
-    # dataset = PocketLigandPairDatasetFeaturized(path, ligand_atom_mode)
-    # train_data, test_data = dataset.train_data, dataset.test_data
-    # print(f'train_size: {len(train_data)}, {sys.getsizeof(train_data)}')
-    # print(f'test_size: {len(test_data)}, {sys.getsizeof(test_data)}')
-    # print(test_data[0], sys.getsizeof(test_data[0]))
-
-    ############################################################
-
-    # test featurization
-    # find all atom types
-    # atom_types = {(1, False): 0}
-
-    # dataset = PocketLigandPairDataset(path, transform)
-    # for i in tqdm(range(len(dataset))):
-    #     data = dataset[i]
-    #     element_list = data.ligand_element
-    #     hybridization_list = data.ligand_hybridization
-    #     aromatic_list = [v[trans.AROMATIC_FEAT_MAP_IDX] for v in data.ligand_atom_feature]
-
-    #     types = [(e, a) for e, h, a in zip(element_list, hybridization_list, aromatic_list)]
-    #     for t in types:
-    #         t = (t[0].item(), bool(t[1].item()))
-    #         if t not in atom_types:
-    #             atom_types[t] = 0
-    #         atom_types[t] += 1
-
-    # idx = 0
-    # for k in sorted(atom_types.keys()):
-    #     print(f'{k}: {idx}, # {atom_types[k]}')
-    #     idx += 1
-
-    ############################################################
+#     def _transform(self):
+#         raw_dataset = PocketLigandPairDataset(self.raw_path, None, 'final')
+
+#         split_path = os.path.join(
+#             os.path.dirname(self.raw_path), 'crossdocked_pocket10_pose_split.pt',
+#         )
+#         split = torch.load(split_path)
+#         train_ids, test_ids = split['train'], split['test']
+#         print(f'train_size: {len(train_ids)}, test_size: {len(test_ids)}')
+
+#         protein_featurizer = trans.FeaturizeProteinAtom()
+#         ligand_featurizer = trans.FeaturizeLigandAtom(self.ligand_atom_mode)
+#         transform_list = [
+#             protein_featurizer,
+#             ligand_featurizer,
+#             # trans.FeaturizeLigandBond(),
+#         ]
+#         transform = Compose(transform_list)
+#         self.protein_atom_feature_dim = protein_featurizer.feature_dim
+#         self.ligand_atom_feature_dim = ligand_featurizer.feature_dim
+
+#         def _transform_subset(ids):
+#             data_list = []
+
+#             for idx in tqdm(ids):
+#                 data = raw_dataset[idx]
+#                 data = transform(data)
+#                 tr_data = {}
+#                 for k in self.features_to_save:
+#                     tr_data[k] = getattr(data, k)
+#                 tr_data['id'] = idx
+#                 tr_data = ProteinLigandData(**tr_data)
+#                 data_list.append(tr_data)
+#             return data_list
+
+#         self.train_data = _transform_subset(train_ids)
+#         print(f'train_size: {len(self.train_data)}, {sys.getsizeof(self.train_data)}')
+#         self.test_data = _transform_subset(test_ids)
+#         print(f'test_size: {len(self.test_data)}, {sys.getsizeof(self.test_data)}')
+#         torch.save({
+#             'train': self.train_data, 'test': self.test_data,
+#             'protein_atom_feature_dim': self.protein_atom_feature_dim,
+#             'ligand_atom_feature_dim': self.ligand_atom_feature_dim,
+#         }, self.transformed_path)
+
+
+# if __name__ == '__main__':
+#     # original dataset
+#     dataset = PocketLigandPairDataset('./data/crossdocked_v1.1_rmsd1.0_pocket10')
+#     print(len(dataset), dataset[0])
+
+#     ############################################################
+
+#     # test DecompDiffDataset
+#     # dataset = PocketLigandGeneratedPairDataset('/sharefs/share/sbdd_data/crossdocked_pocket10')
+#     # print(len(dataset), dataset[0])
+
+#     ############################################################
+
+#     # test featurized dataset (GPU accelerated)
+#     # path = '/sharefs/share/sbdd_data/crossdocked_v1.1_rmsd1.0_pocket10'
+#     # ligand_atom_mode = 'add_aromatic'
+
+#     # dataset = PocketLigandPairDatasetFeaturized(path, ligand_atom_mode)
+#     # train_data, test_data = dataset.train_data, dataset.test_data
+#     # print(f'train_size: {len(train_data)}, {sys.getsizeof(train_data)}')
+#     # print(f'test_size: {len(test_data)}, {sys.getsizeof(test_data)}')
+#     # print(test_data[0], sys.getsizeof(test_data[0]))
+
+#     ############################################################
+
+#     # test featurization
+#     # find all atom types
+#     # atom_types = {(1, False): 0}
+
+#     # dataset = PocketLigandPairDataset(path, transform)
+#     # for i in tqdm(range(len(dataset))):
+#     #     data = dataset[i]
+#     #     element_list = data.ligand_element
+#     #     hybridization_list = data.ligand_hybridization
+#     #     aromatic_list = [v[trans.AROMATIC_FEAT_MAP_IDX] for v in data.ligand_atom_feature]
+
+#     #     types = [(e, a) for e, h, a in zip(element_list, hybridization_list, aromatic_list)]
+#     #     for t in types:
+#     #         t = (t[0].item(), bool(t[1].item()))
+#     #         if t not in atom_types:
+#     #             atom_types[t] = 0
+#     #         atom_types[t] += 1
+
+#     # idx = 0
+#     # for k in sorted(atom_types.keys()):
+#     #     print(f'{k}: {idx}, # {atom_types[k]}')
+#     #     idx += 1
+
+#     ############################################################
     
-    # count atom types
-    # type_counter, aromatic_counter, full_counter = {}, {}, {}
-    # for i, data in enumerate(tqdm(dataset)):
-    #     element_list = data.ligand_element
-    #     hybridization_list = data.ligand_hybridization
-    #     aromatic_list = [v[trans.AROMATIC_FEAT_MAP_IDX] for v in data.ligand_atom_feature]
-    #     flag = False
-
-    #     for atom_type in element_list:
-    #         atom_type = int(atom_type.item())
-    #         if atom_type not in type_counter:
-    #             type_counter[atom_type] = 0
-    #         type_counter[atom_type] += 1
-
-    #     for e, a in zip(element_list, aromatic_list):
-    #         e = int(e.item())
-    #         a = bool(a.item())
-    #         key = (e, a)
-    #         if key not in aromatic_counter:
-    #             aromatic_counter[key] = 0
-    #         aromatic_counter[key] += 1
-
-    #         if key not in trans.MAP_ATOM_TYPE_AROMATIC_TO_INDEX:
-    #             flag = True
-
-    #     for e, h, a in zip(element_list, hybridization_list, aromatic_list):
-    #         e = int(e.item())
-    #         a = bool(a.item())
-    #         key = (e, h, a)
-    #         if key not in full_counter:
-    #             full_counter[key] = 0
-    #         full_counter[key] += 1
+#     # count atom types
+#     # type_counter, aromatic_counter, full_counter = {}, {}, {}
+#     # for i, data in enumerate(tqdm(dataset)):
+#     #     element_list = data.ligand_element
+#     #     hybridization_list = data.ligand_hybridization
+#     #     aromatic_list = [v[trans.AROMATIC_FEAT_MAP_IDX] for v in data.ligand_atom_feature]
+#     #     flag = False
+
+#     #     for atom_type in element_list:
+#     #         atom_type = int(atom_type.item())
+#     #         if atom_type not in type_counter:
+#     #             type_counter[atom_type] = 0
+#     #         type_counter[atom_type] += 1
+
+#     #     for e, a in zip(element_list, aromatic_list):
+#     #         e = int(e.item())
+#     #         a = bool(a.item())
+#     #         key = (e, a)
+#     #         if key not in aromatic_counter:
+#     #             aromatic_counter[key] = 0
+#     #         aromatic_counter[key] += 1
+
+#     #         if key not in trans.MAP_ATOM_TYPE_AROMATIC_TO_INDEX:
+#     #             flag = True
+
+#     #     for e, h, a in zip(element_list, hybridization_list, aromatic_list):
+#     #         e = int(e.item())
+#     #         a = bool(a.item())
+#     #         key = (e, h, a)
+#     #         if key not in full_counter:
+#     #             full_counter[key] = 0
+#     #         full_counter[key] += 1
         
-    # print('type_counter', type_counter)
-    # print('aromatic_counter', aromatic_counter)
-    # print('full_counter', full_counter)
+#     # print('type_counter', type_counter)
+#     # print('aromatic_counter', aromatic_counter)
+#     # print('full_counter', full_counter)
 

core/evaluation/docking_qvina.py

@@ -1,193 +1,193 @@
-import os
-import subprocess
-import random
-import string
-from easydict import EasyDict
-from rdkit import Chem
-from rdkit.Chem.rdForceFieldHelpers import UFFOptimizeMolecule
-
-
-def get_random_id(length=30):
-    letters = string.ascii_lowercase
-    return ''.join(random.choice(letters) for i in range(length))
-
-
-def load_pdb(path):
-    with open(path, 'r') as f:
-        return f.read()
-
-
-def parse_qvina_outputs(docked_sdf_path):
-    suppl = Chem.SDMolSupplier(docked_sdf_path)
-    results = []
-    for i, mol in enumerate(suppl):
-        if mol is None:
-            continue
-        line = mol.GetProp('REMARK').splitlines()[0].split()[2:]
-        results.append(EasyDict({
-            'rdmol': mol,
-            'mode_id': i,
-            'affinity': float(line[0]),
-            'rmsd_lb': float(line[1]),
-            'rmsd_ub': float(line[2]),
-        }))
-
-    return results
-
-
-class BaseDockingTask(object):
-
-    def __init__(self, pdb_block, ligand_rdmol):
-        super().__init__()
-        self.pdb_block = pdb_block
-        self.ligand_rdmol = ligand_rdmol
-
-    def run(self):
-        raise NotImplementedError()
-
-    def get_results(self):
-        raise NotImplementedError()
-
-
-class QVinaDockingTask(BaseDockingTask):
-
-    @classmethod
-    def from_generated_mol(cls, ligand_rdmol, ligand_filename, protein_root='./data/crossdocked', **kwargs):
-        # load original pdb
-        protein_fn = os.path.join(
-            os.path.dirname(ligand_filename),
-            os.path.basename(ligand_filename)[:10] + '.pdb'  # PDBId_Chain_rec.pdb
-        )
-        protein_path = os.path.join(protein_root, protein_fn)
-        with open(protein_path, 'r') as f:
-            pdb_block = f.read()
-        return cls(pdb_block, ligand_rdmol, **kwargs)
-
-    @classmethod
-    def from_original_data(cls, data, ligand_root='./data/crossdocked_pocket10', protein_root='./data/crossdocked',
-                           **kwargs):
-        protein_fn = os.path.join(
-            os.path.dirname(data.ligand_filename),
-            os.path.basename(data.ligand_filename)[:10] + '.pdb'
-        )
-        protein_path = os.path.join(protein_root, protein_fn)
-        with open(protein_path, 'r') as f:
-            pdb_block = f.read()
-
-        ligand_path = os.path.join(ligand_root, data.ligand_filename)
-        ligand_rdmol = next(iter(Chem.SDMolSupplier(ligand_path)))
-        return cls(pdb_block, ligand_rdmol, **kwargs)
-
-    def __init__(self, pdb_block, ligand_rdmol, conda_env='adt', tmp_dir='./tmp', use_uff=True, center=None,
-                 size_factor=1., pos=None):
-        super().__init__(pdb_block, ligand_rdmol)
-        self.conda_env = conda_env
-        self.tmp_dir = os.path.realpath(tmp_dir)
-        os.makedirs(tmp_dir, exist_ok=True)
-
-        self.task_id = get_random_id()
-        self.receptor_id = self.task_id + '_receptor'
-        self.ligand_id = self.task_id + '_ligand'
-
-        self.receptor_path = os.path.join(self.tmp_dir, self.receptor_id + '.pdb')
-        self.ligand_path = os.path.join(self.tmp_dir, self.ligand_id + '.sdf')
-
-        with open(self.receptor_path, 'w') as f:
-            f.write(pdb_block)
-
-        ligand_rdmol = Chem.AddHs(ligand_rdmol, addCoords=True)
-        if use_uff:
-            UFFOptimizeMolecule(ligand_rdmol)
-        # print('after uff smiles: ', Chem.MolToSmiles(ligand_rdmol))
-        sdf_writer = Chem.SDWriter(self.ligand_path)
-        sdf_writer.write(ligand_rdmol)
-        sdf_writer.close()
-        self.ligand_rdmol = ligand_rdmol
-
-        pos = ligand_rdmol.GetConformer(0).GetPositions()
-        if center is None:
-            self.center = (pos.max(0) + pos.min(0)) / 2
-        else:
-            self.center = center
-
-        if size_factor is None:
-            self.size_x, self.size_y, self.size_z = 20, 20, 20
-        else:
-            self.size_x, self.size_y, self.size_z = (pos.max(0) - pos.min(0)) * size_factor
-
-        self.proc = None
-        self.results = None
-        self.output = None
-        self.error_output = None
-        self.docked_sdf_path = None
-
-    def run(self, exhaustiveness=16):
-        commands = """
-eval "$(conda shell.bash hook)"
-conda activate {env}
-cd {tmp}
-# Prepare receptor (PDB->PDBQT)
-prepare_receptor4.py -r {receptor_id}.pdb
-# Prepare ligand
-obabel {ligand_id}.sdf -O{ligand_id}.pdbqt
-qvina2 \
-    --receptor {receptor_id}.pdbqt \
-    --ligand {ligand_id}.pdbqt \
-    --center_x {center_x:.4f} \
-    --center_y {center_y:.4f} \
-    --center_z {center_z:.4f} \
-    --size_x {size_x} --size_y {size_y} --size_z {size_z} \
-    --exhaustiveness {exhaust}
-obabel {ligand_id}_out.pdbqt -O{ligand_id}_out.sdf -h
-        """.format(
-            receptor_id=self.receptor_id,
-            ligand_id=self.ligand_id,
-            env=self.conda_env,
-            tmp=self.tmp_dir,
-            exhaust=exhaustiveness,
-            center_x=self.center[0],
-            center_y=self.center[1],
-            center_z=self.center[2],
-            size_x=self.size_x,
-            size_y=self.size_y,
-            size_z=self.size_z
-        )
-
-        self.docked_sdf_path = os.path.join(self.tmp_dir, '%s_out.sdf' % self.ligand_id)
-
-        self.proc = subprocess.Popen(
-            '/bin/bash',
-            shell=False,
-            stdin=subprocess.PIPE,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE
-        )
-
-        self.proc.stdin.write(commands.encode('utf-8'))
-        self.proc.stdin.close()
-
-        # return commands
-
-    def run_sync(self):
-        self.run()
-        while self.get_results() is None:
-            pass
-        results = self.get_results()
-        print('Best affinity:', results[0]['affinity'])
-        return results
-
-    def get_results(self):
-        if self.proc is None:  # Not started
-            return None
-        elif self.proc.poll() is None:  # In progress
-            return None
-        else:
-            if self.output is None:
-                self.output = self.proc.stdout.readlines()
-                self.error_output = self.proc.stderr.readlines()
-                try:
-                    self.results = parse_qvina_outputs(self.docked_sdf_path)
-                except Exception as e:
-                    print('[Error] Vina output error: %s' % self.docked_sdf_path, e)
-                    return []
-            return self.results
+# import os
+# import subprocess
+# import random
+# import string
+# from easydict import EasyDict
+# from rdkit import Chem
+# from rdkit.Chem.rdForceFieldHelpers import UFFOptimizeMolecule
+
+
+# def get_random_id(length=30):
+#     letters = string.ascii_lowercase
+#     return ''.join(random.choice(letters) for i in range(length))
+
+
+# def load_pdb(path):
+#     with open(path, 'r') as f:
+#         return f.read()
+
+
+# def parse_qvina_outputs(docked_sdf_path):
+#     suppl = Chem.SDMolSupplier(docked_sdf_path)
+#     results = []
+#     for i, mol in enumerate(suppl):
+#         if mol is None:
+#             continue
+#         line = mol.GetProp('REMARK').splitlines()[0].split()[2:]
+#         results.append(EasyDict({
+#             'rdmol': mol,
+#             'mode_id': i,
+#             'affinity': float(line[0]),
+#             'rmsd_lb': float(line[1]),
+#             'rmsd_ub': float(line[2]),
+#         }))
+
+#     return results
+
+
+# class BaseDockingTask(object):
+
+#     def __init__(self, pdb_block, ligand_rdmol):
+#         super().__init__()
+#         self.pdb_block = pdb_block
+#         self.ligand_rdmol = ligand_rdmol
+
+#     def run(self):
+#         raise NotImplementedError()
+
+#     def get_results(self):
+#         raise NotImplementedError()
+
+
+# class QVinaDockingTask(BaseDockingTask):
+
+#     @classmethod
+#     def from_generated_mol(cls, ligand_rdmol, ligand_filename, protein_root='./data/crossdocked', **kwargs):
+#         # load original pdb
+#         protein_fn = os.path.join(
+#             os.path.dirname(ligand_filename),
+#             os.path.basename(ligand_filename)[:10] + '.pdb'  # PDBId_Chain_rec.pdb
+#         )
+#         protein_path = os.path.join(protein_root, protein_fn)
+#         with open(protein_path, 'r') as f:
+#             pdb_block = f.read()
+#         return cls(pdb_block, ligand_rdmol, **kwargs)
+
+#     @classmethod
+#     def from_original_data(cls, data, ligand_root='./data/crossdocked_pocket10', protein_root='./data/crossdocked',
+#                            **kwargs):
+#         protein_fn = os.path.join(
+#             os.path.dirname(data.ligand_filename),
+#             os.path.basename(data.ligand_filename)[:10] + '.pdb'
+#         )
+#         protein_path = os.path.join(protein_root, protein_fn)
+#         with open(protein_path, 'r') as f:
+#             pdb_block = f.read()
+
+#         ligand_path = os.path.join(ligand_root, data.ligand_filename)
+#         ligand_rdmol = next(iter(Chem.SDMolSupplier(ligand_path)))
+#         return cls(pdb_block, ligand_rdmol, **kwargs)
+
+#     def __init__(self, pdb_block, ligand_rdmol, conda_env='adt', tmp_dir='./tmp', use_uff=True, center=None,
+#                  size_factor=1., pos=None):
+#         super().__init__(pdb_block, ligand_rdmol)
+#         self.conda_env = conda_env
+#         self.tmp_dir = os.path.realpath(tmp_dir)
+#         os.makedirs(tmp_dir, exist_ok=True)
+
+#         self.task_id = get_random_id()
+#         self.receptor_id = self.task_id + '_receptor'
+#         self.ligand_id = self.task_id + '_ligand'
+
+#         self.receptor_path = os.path.join(self.tmp_dir, self.receptor_id + '.pdb')
+#         self.ligand_path = os.path.join(self.tmp_dir, self.ligand_id + '.sdf')
+
+#         with open(self.receptor_path, 'w') as f:
+#             f.write(pdb_block)
+
+#         ligand_rdmol = Chem.AddHs(ligand_rdmol, addCoords=True)
+#         if use_uff:
+#             UFFOptimizeMolecule(ligand_rdmol)
+#         # print('after uff smiles: ', Chem.MolToSmiles(ligand_rdmol))
+#         sdf_writer = Chem.SDWriter(self.ligand_path)
+#         sdf_writer.write(ligand_rdmol)
+#         sdf_writer.close()
+#         self.ligand_rdmol = ligand_rdmol
+
+#         pos = ligand_rdmol.GetConformer(0).GetPositions()
+#         if center is None:
+#             self.center = (pos.max(0) + pos.min(0)) / 2
+#         else:
+#             self.center = center
+
+#         if size_factor is None:
+#             self.size_x, self.size_y, self.size_z = 20, 20, 20
+#         else:
+#             self.size_x, self.size_y, self.size_z = (pos.max(0) - pos.min(0)) * size_factor
+
+#         self.proc = None
+#         self.results = None
+#         self.output = None
+#         self.error_output = None
+#         self.docked_sdf_path = None
+
+#     def run(self, exhaustiveness=16):
+#         commands = """
+# eval "$(conda shell.bash hook)"
+# conda activate {env}
+# cd {tmp}
+# # Prepare receptor (PDB->PDBQT)
+# prepare_receptor4.py -r {receptor_id}.pdb
+# # Prepare ligand
+# obabel {ligand_id}.sdf -O{ligand_id}.pdbqt
+# qvina2 \
+#     --receptor {receptor_id}.pdbqt \
+#     --ligand {ligand_id}.pdbqt \
+#     --center_x {center_x:.4f} \
+#     --center_y {center_y:.4f} \
+#     --center_z {center_z:.4f} \
+#     --size_x {size_x} --size_y {size_y} --size_z {size_z} \
+#     --exhaustiveness {exhaust}
+# obabel {ligand_id}_out.pdbqt -O{ligand_id}_out.sdf -h
+#         """.format(
+#             receptor_id=self.receptor_id,
+#             ligand_id=self.ligand_id,
+#             env=self.conda_env,
+#             tmp=self.tmp_dir,
+#             exhaust=exhaustiveness,
+#             center_x=self.center[0],
+#             center_y=self.center[1],
+#             center_z=self.center[2],
+#             size_x=self.size_x,
+#             size_y=self.size_y,
+#             size_z=self.size_z
+#         )
+
+#         self.docked_sdf_path = os.path.join(self.tmp_dir, '%s_out.sdf' % self.ligand_id)
+
+#         self.proc = subprocess.Popen(
+#             '/bin/bash',
+#             shell=False,
+#             stdin=subprocess.PIPE,
+#             stdout=subprocess.PIPE,
+#             stderr=subprocess.PIPE
+#         )
+
+#         self.proc.stdin.write(commands.encode('utf-8'))
+#         self.proc.stdin.close()
+
+#         # return commands
+
+#     def run_sync(self):
+#         self.run()
+#         while self.get_results() is None:
+#             pass
+#         results = self.get_results()
+#         print('Best affinity:', results[0]['affinity'])
+#         return results
+
+#     def get_results(self):
+#         if self.proc is None:  # Not started
+#             return None
+#         elif self.proc.poll() is None:  # In progress
+#             return None
+#         else:
+#             if self.output is None:
+#                 self.output = self.proc.stdout.readlines()
+#                 self.error_output = self.proc.stderr.readlines()
+#                 try:
+#                     self.results = parse_qvina_outputs(self.docked_sdf_path)
+#                 except Exception as e:
+#                     print('[Error] Vina output error: %s' % self.docked_sdf_path, e)
+#                     return []
+#             return self.results

core/evaluation/docking_vina.py

@@ -1,269 +1,269 @@
-from openbabel import pybel
-from meeko import MoleculePreparation
-from meeko import obutils
-from vina import Vina
-import subprocess
-import rdkit.Chem as Chem
-from rdkit.Chem import AllChem
-import tempfile
-import AutoDockTools
-import os
-import contextlib
-from posecheck import PoseCheck
-
-from core.evaluation.docking_qvina import get_random_id, BaseDockingTask
-
-
-def suppress_stdout(func):
-    def wrapper(*a, **ka):
-        with open(os.devnull, 'w') as devnull:
-            with contextlib.redirect_stdout(devnull):
-                return func(*a, **ka)
-    return wrapper
-
-
-class PrepLig(object):
-    def __init__(self, input_mol, mol_format):
-        if mol_format == 'smi':
-            self.ob_mol = pybel.readstring('smi', input_mol)
-        elif mol_format == 'sdf': 
-            self.ob_mol = next(pybel.readfile(mol_format, input_mol))
-        else:
-            raise ValueError(f'mol_format {mol_format} not supported')
+# from openbabel import pybel
+# from meeko import MoleculePreparation
+# from meeko import obutils
+# from vina import Vina
+# import subprocess
+# import rdkit.Chem as Chem
+# from rdkit.Chem import AllChem
+# import tempfile
+# import AutoDockTools
+# import os
+# import contextlib
+# from posecheck import PoseCheck
+
+# from core.evaluation.docking_qvina import get_random_id, BaseDockingTask
+
+
+# def suppress_stdout(func):
+#     def wrapper(*a, **ka):
+#         with open(os.devnull, 'w') as devnull:
+#             with contextlib.redirect_stdout(devnull):
+#                 return func(*a, **ka)
+#     return wrapper
+
+
+# class PrepLig(object):
+#     def __init__(self, input_mol, mol_format):
+#         if mol_format == 'smi':
+#             self.ob_mol = pybel.readstring('smi', input_mol)
+#         elif mol_format == 'sdf': 
+#             self.ob_mol = next(pybel.readfile(mol_format, input_mol))
+#         else:
+#             raise ValueError(f'mol_format {mol_format} not supported')
         
-    def addH(self, polaronly=False, correctforph=True, PH=7): 
-        self.ob_mol.OBMol.AddHydrogens(polaronly, correctforph, PH)
-        obutils.writeMolecule(self.ob_mol.OBMol, 'tmp_h.sdf')
-
-    def gen_conf(self):
-        sdf_block = self.ob_mol.write('sdf')
-        rdkit_mol = Chem.MolFromMolBlock(sdf_block, removeHs=False)
-        AllChem.EmbedMolecule(rdkit_mol, Chem.rdDistGeom.ETKDGv3())
-        self.ob_mol = pybel.readstring('sdf', Chem.MolToMolBlock(rdkit_mol))
-        obutils.writeMolecule(self.ob_mol.OBMol, 'conf_h.sdf')
-
-    @suppress_stdout
-    def get_pdbqt(self, lig_pdbqt=None):
-        preparator = MoleculePreparation()
-        preparator.prepare(self.ob_mol.OBMol)
-        if lig_pdbqt is not None: 
-            preparator.write_pdbqt_file(lig_pdbqt)
-            return 
-        else: 
-            return preparator.write_pdbqt_string()
+#     def addH(self, polaronly=False, correctforph=True, PH=7): 
+#         self.ob_mol.OBMol.AddHydrogens(polaronly, correctforph, PH)
+#         obutils.writeMolecule(self.ob_mol.OBMol, 'tmp_h.sdf')
+
+#     def gen_conf(self):
+#         sdf_block = self.ob_mol.write('sdf')
+#         rdkit_mol = Chem.MolFromMolBlock(sdf_block, removeHs=False)
+#         AllChem.EmbedMolecule(rdkit_mol, Chem.rdDistGeom.ETKDGv3())
+#         self.ob_mol = pybel.readstring('sdf', Chem.MolToMolBlock(rdkit_mol))
+#         obutils.writeMolecule(self.ob_mol.OBMol, 'conf_h.sdf')
+
+#     @suppress_stdout
+#     def get_pdbqt(self, lig_pdbqt=None):
+#         preparator = MoleculePreparation()
+#         preparator.prepare(self.ob_mol.OBMol)
+#         if lig_pdbqt is not None: 
+#             preparator.write_pdbqt_file(lig_pdbqt)
+#             return 
+#         else: 
+#             return preparator.write_pdbqt_string()
         
 
-class PrepProt(object): 
-    def __init__(self, pdb_file): 
-        self.prot = pdb_file
+# class PrepProt(object): 
+#     def __init__(self, pdb_file): 
+#         self.prot = pdb_file
     
-    def del_water(self, dry_pdb_file): # optional
-        with open(self.prot) as f: 
-            lines = [l for l in f.readlines() if l.startswith('ATOM') or l.startswith('HETATM')] 
-            dry_lines = [l for l in lines if not 'HOH' in l]
+#     def del_water(self, dry_pdb_file): # optional
+#         with open(self.prot) as f: 
+#             lines = [l for l in f.readlines() if l.startswith('ATOM') or l.startswith('HETATM')] 
+#             dry_lines = [l for l in lines if not 'HOH' in l]
         
-        with open(dry_pdb_file, 'w') as f:
-            f.write(''.join(dry_lines))
-        self.prot = dry_pdb_file
+#         with open(dry_pdb_file, 'w') as f:
+#             f.write(''.join(dry_lines))
+#         self.prot = dry_pdb_file
         
-    def addH(self, prot_pqr):  # call pdb2pqr
-        self.prot_pqr = prot_pqr
-        subprocess.Popen(['pdb2pqr30','--ff=AMBER',self.prot, self.prot_pqr],
-                         stderr=subprocess.DEVNULL, stdout=subprocess.DEVNULL).communicate()
+#     def addH(self, prot_pqr):  # call pdb2pqr
+#         self.prot_pqr = prot_pqr
+#         subprocess.Popen(['pdb2pqr30','--ff=AMBER',self.prot, self.prot_pqr],
+#                          stderr=subprocess.DEVNULL, stdout=subprocess.DEVNULL).communicate()
 
-    def get_pdbqt(self, prot_pdbqt):
-        prepare_receptor = os.path.join(AutoDockTools.__path__[0], 'Utilities24/prepare_receptor4.py')
-        subprocess.Popen(['python3', prepare_receptor, '-r', self.prot_pqr, '-o', prot_pdbqt],
-                         stderr=subprocess.DEVNULL, stdout=subprocess.DEVNULL).communicate()
+#     def get_pdbqt(self, prot_pdbqt):
+#         prepare_receptor = os.path.join(AutoDockTools.__path__[0], 'Utilities24/prepare_receptor4.py')
+#         subprocess.Popen(['python3', prepare_receptor, '-r', self.prot_pqr, '-o', prot_pdbqt],
+#                          stderr=subprocess.DEVNULL, stdout=subprocess.DEVNULL).communicate()
 
 
-class VinaDock(object): 
-    def __init__(self, lig_pdbqt, prot_pdbqt): 
-        self.lig_pdbqt = lig_pdbqt
-        self.prot_pdbqt = prot_pdbqt
+# class VinaDock(object): 
+#     def __init__(self, lig_pdbqt, prot_pdbqt): 
+#         self.lig_pdbqt = lig_pdbqt
+#         self.prot_pdbqt = prot_pdbqt
     
-    def _max_min_pdb(self, pdb, buffer):
-        with open(pdb, 'r') as f: 
-            lines = [l for l in f.readlines() if l.startswith('ATOM') or l.startswith('HEATATM')]
-            xs = [float(l[31:39]) for l in lines]
-            ys = [float(l[39:47]) for l in lines]
-            zs = [float(l[47:55]) for l in lines]
-            print(max(xs), min(xs))
-            print(max(ys), min(ys))
-            print(max(zs), min(zs))
-            pocket_center = [(max(xs) + min(xs))/2, (max(ys) + min(ys))/2, (max(zs) + min(zs))/2]
-            box_size = [(max(xs) - min(xs)) + buffer, (max(ys) - min(ys)) + buffer, (max(zs) - min(zs)) + buffer]
-            return pocket_center, box_size
+#     def _max_min_pdb(self, pdb, buffer):
+#         with open(pdb, 'r') as f: 
+#             lines = [l for l in f.readlines() if l.startswith('ATOM') or l.startswith('HEATATM')]
+#             xs = [float(l[31:39]) for l in lines]
+#             ys = [float(l[39:47]) for l in lines]
+#             zs = [float(l[47:55]) for l in lines]
+#             print(max(xs), min(xs))
+#             print(max(ys), min(ys))
+#             print(max(zs), min(zs))
+#             pocket_center = [(max(xs) + min(xs))/2, (max(ys) + min(ys))/2, (max(zs) + min(zs))/2]
+#             box_size = [(max(xs) - min(xs)) + buffer, (max(ys) - min(ys)) + buffer, (max(zs) - min(zs)) + buffer]
+#             return pocket_center, box_size
     
-    def get_box(self, ref=None, buffer=0):
-        '''
-        ref: reference pdb to define pocket. 
-        buffer: buffer size to add 
-
-        if ref is not None: 
-            get the max and min on x, y, z axis in ref pdb and add buffer to each dimension 
-        else: 
-            use the entire protein to define pocket 
-        '''
-        if ref is None: 
-            ref = self.prot_pdbqt
-        self.pocket_center, self.box_size = self._max_min_pdb(ref, buffer)
-        print(self.pocket_center, self.box_size)
-
-    def dock(self, score_func='vina', seed=0, mode='dock', exhaustiveness=8, save_pose=False, **kwargs):  # seed=0 mean random seed
-        v = Vina(sf_name=score_func, seed=seed, verbosity=0, **kwargs)
-        v.set_receptor(self.prot_pdbqt)
-        v.set_ligand_from_file(self.lig_pdbqt)
-        v.compute_vina_maps(center=self.pocket_center, box_size=self.box_size)
-        if mode == 'score_only': 
-            score = v.score()[0]
-        elif mode == 'minimize':
-            score = v.optimize()[0]
-        elif mode == 'dock':
-            v.dock(exhaustiveness=exhaustiveness, n_poses=1)
-            score = v.energies(n_poses=1)[0][0]
-        else:
-            raise ValueError
+#     def get_box(self, ref=None, buffer=0):
+#         '''
+#         ref: reference pdb to define pocket. 
+#         buffer: buffer size to add 
+
+#         if ref is not None: 
+#             get the max and min on x, y, z axis in ref pdb and add buffer to each dimension 
+#         else: 
+#             use the entire protein to define pocket 
+#         '''
+#         if ref is None: 
+#             ref = self.prot_pdbqt
+#         self.pocket_center, self.box_size = self._max_min_pdb(ref, buffer)
+#         print(self.pocket_center, self.box_size)
+
+#     def dock(self, score_func='vina', seed=0, mode='dock', exhaustiveness=8, save_pose=False, **kwargs):  # seed=0 mean random seed
+#         v = Vina(sf_name=score_func, seed=seed, verbosity=0, **kwargs)
+#         v.set_receptor(self.prot_pdbqt)
+#         v.set_ligand_from_file(self.lig_pdbqt)
+#         v.compute_vina_maps(center=self.pocket_center, box_size=self.box_size)
+#         if mode == 'score_only': 
+#             score = v.score()[0]
+#         elif mode == 'minimize':
+#             score = v.optimize()[0]
+#         elif mode == 'dock':
+#             v.dock(exhaustiveness=exhaustiveness, n_poses=1)
+#             score = v.energies(n_poses=1)[0][0]
+#         else:
+#             raise ValueError
         
-        if not save_pose: 
-            return score
-        else: 
-            if mode == 'score_only': 
-                pose = None 
-            elif mode == 'minimize': 
-                tmp = tempfile.NamedTemporaryFile()
-                with open(tmp.name, 'w') as f: 
-                    v.write_pose(tmp.name, overwrite=True)             
-                with open(tmp.name, 'r') as f: 
-                    pose = f.read()
+#         if not save_pose: 
+#             return score
+#         else: 
+#             if mode == 'score_only': 
+#                 pose = None 
+#             elif mode == 'minimize': 
+#                 tmp = tempfile.NamedTemporaryFile()
+#                 with open(tmp.name, 'w') as f: 
+#                     v.write_pose(tmp.name, overwrite=True)             
+#                 with open(tmp.name, 'r') as f: 
+#                     pose = f.read()
    
-            elif mode == 'dock': 
-                pose = v.poses(n_poses=1)
-            else:
-                raise ValueError
-            return score, pose
-
-
-class VinaDockingTask(BaseDockingTask):
-
-    @classmethod
-    def from_original_data(cls, data, ligand_root='./data/crossdocked_pocket10', protein_root='./data/crossdocked',
-                           **kwargs):
-        protein_fn = os.path.join(
-            os.path.dirname(data.ligand_filename),
-            os.path.basename(data.ligand_filename)[:10] + '.pdb'
-        )
-        protein_path = os.path.join(protein_root, protein_fn)
-
-        ligand_path = os.path.join(ligand_root, data.ligand_filename)
-        ligand_rdmol = next(iter(Chem.SDMolSupplier(ligand_path)))
-        return cls(protein_path, ligand_rdmol, **kwargs)
-
-    @classmethod
-    def from_generated_mol(cls, ligand_rdmol, protein_filename, **kwargs):
-        # load original pdb
-        # TODO: make it more general and compatible with sample_for_pocket
-        # protein_fn = os.path.join(
-        #     os.path.dirname(ligand_filename),
-        #     os.path.basename(ligand_filename)[:10] + '.pdb' if 'molecule' not in ligand_filename  # PDBId_Chain_rec.pdb
-        #         else os.path.basename(ligand_filename).replace('_molecule', '_protein').replace('.sdf', '.pdb')
-        # )
-        # protein_path = os.path.join(protein_root, protein_fn)
-        protein_path = protein_filename
-        return cls(protein_path, ligand_rdmol, **kwargs)
-
-    def __init__(self, protein_path, ligand_rdmol, tmp_dir='./tmp', center=None,
-                 size_factor=1., buffer=5.0, pos=None):
-        super().__init__(protein_path, ligand_rdmol)
-        # self.conda_env = conda_env
-        self.tmp_dir = os.path.realpath(tmp_dir)
-        os.makedirs(tmp_dir, exist_ok=True)
-
-        self.task_id = get_random_id()
-        self.receptor_id = self.task_id + '_receptor'
-        self.ligand_id = self.task_id + '_ligand'
-
-        self.receptor_path = protein_path
-        self.ligand_path = os.path.join(self.tmp_dir, self.ligand_id + '.sdf')
-
-        self.recon_ligand_mol = ligand_rdmol
-        ligand_rdmol = Chem.AddHs(ligand_rdmol, addCoords=True)
-
-        sdf_writer = Chem.SDWriter(self.ligand_path)
-        sdf_writer.write(ligand_rdmol)
-        sdf_writer.close()
-        self.ligand_rdmol = ligand_rdmol
-
-        pos = ligand_rdmol.GetConformer(0).GetPositions()
-        # if pos is None:
-        #    raise ValueError('pos is None')
-        if center is None:
-            self.center = (pos.max(0) + pos.min(0)) / 2
-        else:
-            self.center = center
-
-        if size_factor is None:
-            self.size_x, self.size_y, self.size_z = 20, 20, 20
-        else:
-            self.size_x, self.size_y, self.size_z = (pos.max(0) - pos.min(0)) * size_factor + buffer
-
-        self.proc = None
-        self.results = None
-        self.output = None
-        self.error_output = None
-        self.docked_sdf_path = None
-
-    def run(self, mode='dock', exhaustiveness=8, **kwargs):
-        ligand_pdbqt = self.ligand_path[:-4] + '.pdbqt'
-        protein_pqr = self.receptor_path[:-4] + '.pqr'
-        protein_pdbqt = self.receptor_path[:-4] + '.pdbqt'
-
-        lig = PrepLig(self.ligand_path, 'sdf')
-        lig.get_pdbqt(ligand_pdbqt)
-
-        prot = PrepProt(self.receptor_path)
-        if not os.path.exists(protein_pqr):
-            prot.addH(protein_pqr)
-        if not os.path.exists(protein_pdbqt):
-            prot.get_pdbqt(protein_pdbqt)
-
-        dock = VinaDock(ligand_pdbqt, protein_pdbqt)
-        dock.pocket_center, dock.box_size = self.center, [self.size_x, self.size_y, self.size_z]
-        score, pose = dock.dock(score_func='vina', mode=mode, exhaustiveness=exhaustiveness, save_pose=True, **kwargs)
-        return [{'affinity': score, 'pose': pose}]
+#             elif mode == 'dock': 
+#                 pose = v.poses(n_poses=1)
+#             else:
+#                 raise ValueError
+#             return score, pose
+
+
+# class VinaDockingTask(BaseDockingTask):
+
+#     @classmethod
+#     def from_original_data(cls, data, ligand_root='./data/crossdocked_pocket10', protein_root='./data/crossdocked',
+#                            **kwargs):
+#         protein_fn = os.path.join(
+#             os.path.dirname(data.ligand_filename),
+#             os.path.basename(data.ligand_filename)[:10] + '.pdb'
+#         )
+#         protein_path = os.path.join(protein_root, protein_fn)
+
+#         ligand_path = os.path.join(ligand_root, data.ligand_filename)
+#         ligand_rdmol = next(iter(Chem.SDMolSupplier(ligand_path)))
+#         return cls(protein_path, ligand_rdmol, **kwargs)
+
+#     @classmethod
+#     def from_generated_mol(cls, ligand_rdmol, protein_filename, **kwargs):
+#         # load original pdb
+#         # TODO: make it more general and compatible with sample_for_pocket
+#         # protein_fn = os.path.join(
+#         #     os.path.dirname(ligand_filename),
+#         #     os.path.basename(ligand_filename)[:10] + '.pdb' if 'molecule' not in ligand_filename  # PDBId_Chain_rec.pdb
+#         #         else os.path.basename(ligand_filename).replace('_molecule', '_protein').replace('.sdf', '.pdb')
+#         # )
+#         # protein_path = os.path.join(protein_root, protein_fn)
+#         protein_path = protein_filename
+#         return cls(protein_path, ligand_rdmol, **kwargs)
+
+#     def __init__(self, protein_path, ligand_rdmol, tmp_dir='./tmp', center=None,
+#                  size_factor=1., buffer=5.0, pos=None):
+#         super().__init__(protein_path, ligand_rdmol)
+#         # self.conda_env = conda_env
+#         self.tmp_dir = os.path.realpath(tmp_dir)
+#         os.makedirs(tmp_dir, exist_ok=True)
+
+#         self.task_id = get_random_id()
+#         self.receptor_id = self.task_id + '_receptor'
+#         self.ligand_id = self.task_id + '_ligand'
+
+#         self.receptor_path = protein_path
+#         self.ligand_path = os.path.join(self.tmp_dir, self.ligand_id + '.sdf')
+
+#         self.recon_ligand_mol = ligand_rdmol
+#         ligand_rdmol = Chem.AddHs(ligand_rdmol, addCoords=True)
+
+#         sdf_writer = Chem.SDWriter(self.ligand_path)
+#         sdf_writer.write(ligand_rdmol)
+#         sdf_writer.close()
+#         self.ligand_rdmol = ligand_rdmol
+
+#         pos = ligand_rdmol.GetConformer(0).GetPositions()
+#         # if pos is None:
+#         #    raise ValueError('pos is None')
+#         if center is None:
+#             self.center = (pos.max(0) + pos.min(0)) / 2
+#         else:
+#             self.center = center
+
+#         if size_factor is None:
+#             self.size_x, self.size_y, self.size_z = 20, 20, 20
+#         else:
+#             self.size_x, self.size_y, self.size_z = (pos.max(0) - pos.min(0)) * size_factor + buffer
+
+#         self.proc = None
+#         self.results = None
+#         self.output = None
+#         self.error_output = None
+#         self.docked_sdf_path = None
+
+#     def run(self, mode='dock', exhaustiveness=8, **kwargs):
+#         ligand_pdbqt = self.ligand_path[:-4] + '.pdbqt'
+#         protein_pqr = self.receptor_path[:-4] + '.pqr'
+#         protein_pdbqt = self.receptor_path[:-4] + '.pdbqt'
+
+#         lig = PrepLig(self.ligand_path, 'sdf')
+#         lig.get_pdbqt(ligand_pdbqt)
+
+#         prot = PrepProt(self.receptor_path)
+#         if not os.path.exists(protein_pqr):
+#             prot.addH(protein_pqr)
+#         if not os.path.exists(protein_pdbqt):
+#             prot.get_pdbqt(protein_pdbqt)
+
+#         dock = VinaDock(ligand_pdbqt, protein_pdbqt)
+#         dock.pocket_center, dock.box_size = self.center, [self.size_x, self.size_y, self.size_z]
+#         score, pose = dock.dock(score_func='vina', mode=mode, exhaustiveness=exhaustiveness, save_pose=True, **kwargs)
+#         return [{'affinity': score, 'pose': pose}]
     
-    @suppress_stdout
-    def run_pose_check(self):
-        pc = PoseCheck()
-        pc.load_protein_from_pdb(self.receptor_path)
-        # pc.load_ligands_from_sdf(self.ligand_path)
-        pc.load_ligands_from_mols([self.ligand_rdmol])
-        clashes = pc.calculate_clashes()
-        strain = pc.calculate_strain_energy()
-        return {'clashes': clashes[0], 'strain': strain[0]}
-
-
-# if __name__ == '__main__':
-#     lig_pdbqt = 'data/lig.pdbqt'
-#     mol_file = 'data/1a4k_ligand.sdf'
-#     a = PrepLig(mol_file, 'sdf')
-#     # mol_file = 'CC(=C)C(=O)OCCN(C)C'
-#     # a = PrepLig(mol_file, 'smi')
-#     a.addH()
-#     a.gen_conf()
-#     a.get_pdbqt(lig_pdbqt)
-#
-#     prot_file = 'data/1a4k_protein_chainAB.pdb'
-#     prot_dry = 'data/protein_dry.pdb'
-#     prot_pqr = 'data/protein.pqr'
-#     prot_pdbqt = 'data/protein.pdbqt'
-#     b = PrepProt(prot_file)
-#     b.del_water(prot_dry)
-#     b.addH(prot_pqr)
-#     b.get_pdbqt(prot_pdbqt)
-#
-#     dock = VinaDock(lig_pdbqt, prot_pdbqt)
-#     dock.get_box()
-#     dock.dock()
+#     @suppress_stdout
+#     def run_pose_check(self):
+#         pc = PoseCheck()
+#         pc.load_protein_from_pdb(self.receptor_path)
+#         # pc.load_ligands_from_sdf(self.ligand_path)
+#         pc.load_ligands_from_mols([self.ligand_rdmol])
+#         clashes = pc.calculate_clashes()
+#         strain = pc.calculate_strain_energy()
+#         return {'clashes': clashes[0], 'strain': strain[0]}
+
+
+# # if __name__ == '__main__':
+# #     lig_pdbqt = 'data/lig.pdbqt'
+# #     mol_file = 'data/1a4k_ligand.sdf'
+# #     a = PrepLig(mol_file, 'sdf')
+# #     # mol_file = 'CC(=C)C(=O)OCCN(C)C'
+# #     # a = PrepLig(mol_file, 'smi')
+# #     a.addH()
+# #     a.gen_conf()
+# #     a.get_pdbqt(lig_pdbqt)
+# #
+# #     prot_file = 'data/1a4k_protein_chainAB.pdb'
+# #     prot_dry = 'data/protein_dry.pdb'
+# #     prot_pqr = 'data/protein.pqr'
+# #     prot_pdbqt = 'data/protein.pdbqt'
+# #     b = PrepProt(prot_file)
+# #     b.del_water(prot_dry)
+# #     b.addH(prot_pqr)
+# #     b.get_pdbqt(prot_pdbqt)
+# #
+# #     dock = VinaDock(lig_pdbqt, prot_pdbqt)
+# #     dock.get_box()
+# #     dock.dock()
     
 

core/evaluation/metrics.py

@@ -8,8 +8,8 @@ from core.evaluation.utils import (
     convert_atomcloud_to_mol_smiles,
     mol2smiles,
 )
-from core.evaluation.docking_qvina import QVinaDockingTask
-from core.evaluation.docking_vina import VinaDockingTask
+# from core.evaluation.docking_qvina import QVinaDockingTask
+# from core.evaluation.docking_vina import VinaDockingTask
 from typing import List, Dict, Tuple
 from tqdm import tqdm
 import numpy as np

requirements.txt

@@ -1,5 +1,4 @@
 absl_py==2.1.0
-AutoDockTools_py3==1.5.7.post1+10.g9d37a13
 easydict==1.13
 fire==0.6.0
 gradio==4.36.1
@@ -7,14 +6,11 @@ gradio_molecule3d==0.0.5
 imageio==2.34.1
 lmdb==1.4.1
 matplotlib==3.4.3
-meeko==0.1.dev3
 numpy==1.23.1
 openbabel==3.1.1.1
 overrides==7.7.0
 Pillow==9.2.0
 Pillow==10.3.0
-posecheck==1.1
-posecheck.egg==info
 py3Dmol==2.1.0
 pytorch_lightning==2.0.8
 PyYAML==6.0.1
@@ -26,4 +22,3 @@ torch_geometric==2.1.0.post1
 torch_scatter==2.0.9
 torchdiffeq==0.2.3
 tqdm==4.64.0
-vina==1.2.2

sample_for_pocket.py

@@ -30,11 +30,11 @@ from pytorch_lightning import seed_everything
 # from absl import logging
 # import glob
 
-from core.evaluation.utils import scoring_func
-from core.evaluation.docking_vina import VinaDockingTask
-from posecheck import PoseCheck
-import numpy as np
-from rdkit import Chem
+# from core.evaluation.utils import scoring_func
+# from core.evaluation.docking_vina import VinaDockingTask
+# from posecheck import PoseCheck
+# import numpy as np
+# from rdkit import Chem
 
 
 def get_dataloader_from_pdb(cfg):