fix ligand name bug (#232)

Co-authored-by: zhougm <[email protected]>
deepmodeling · Jun 17, 2024 · 43f71d5 · 43f71d5
1 parent 24c1910
commit 43f71d5
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/unimol_docking_v2/interface/predictor/processor.py b/unimol_docking_v2/interface/predictor/processor.py
@@ -416,6 +416,10 @@ def single_clash_fix(self, input_content):
     def clash_fix(self, predicted_ligand, input_protein, input_ligand):
         if self.mode=='batch_one2many':
             input_protein = [input_protein] * len(input_ligand)
+        elif self.mode == 'single':
+            input_ligand = [input_ligand]
+            input_protein = [input_protein]
+            predicted_ligand = [predicted_ligand]
         input_content = zip(predicted_ligand, predicted_ligand, input_ligand, input_protein)
 
         with Pool(self.nthreads) as pool:

diff --git a/unimol_docking_v2/interface/predictor/unimol_predictor.py b/unimol_docking_v2/interface/predictor/unimol_predictor.py
@@ -37,8 +37,8 @@ def predict(self, input_protein:str,
                 batch_size:int):
         lmdb_name = self.preprocess(input_protein, input_ligand, input_docking_grid, output_ligand_name, output_ligand_dir)
 
-        pkt_data_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "example_data", "docking", "dict_pkt.txt")
-        mol_data_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "example_data", "docking", "dict_mol.txt")
+        pkt_data_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "example_data", "dict_pkt.txt")
+        mol_data_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "example_data", "dict_mol.txt")
         script_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "unimol", "infer.py")
         user_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "unimol")
         # inference