fix: unidock_pipeline read only one ligand for big sdf (#96)

* fix:unidock_pipeline run only one ligand if use big sdf

* test:add bigsdf testcase

unidock_tools/src/unidock_tools/application/unidock_pipeline.py

@@ -147,7 +147,7 @@ def postprocessing(self, ligand_scores_list: zip,
             mol_score_list = mol_score_dict[fprefix]
             mol_score_list.sort(key=lambda x: x[1], reverse=False)
             logging.debug([item[1] for item in mol_score_list])
-            logging.info(f"docking result pose num: {len(mol_score_list)}, keep num: {topn_conf}")
+            logging.debug(f"docking result pose num: {len(mol_score_list)}, keep num: {topn_conf}")
             self.mol_group.update_mol_confs_by_file_prefix(fprefix,
                                                            [mol for mol, _ in mol_score_list[:topn_conf]])
             self.mol_group.update_property_by_file_prefix(fprefix, property_name=score_name,

unidock_tools/src/unidock_tools/modules/docking/unidock.py

@@ -203,8 +203,8 @@ def run(self):
         )
         logging.debug(f"Run Uni-Dock log: {resp.stdout}")
         if resp.returncode != 0:
-            logging.info(f"Run Uni-Dock log: {resp.stdout}")
-            logging.error(f"Run Uni-Dock error: {resp.stderr}")
+            logging.info(f"Run Uni-Dock log\n{resp.stdout}")
+            logging.error(f"Run Uni-Dock error\n{resp.stderr}")
 
         result_ligands = [f for f in self.pre_result_ligands if os.path.exists(f)]
         return result_ligands

unidock_tools/src/unidock_tools/utils/mol_group.py

@@ -91,9 +91,10 @@ def _initialize(self, ligand_files: List[Path]):
         for ligand_file in ligand_files:
             file_prefix = ligand_file.stem
             mols = read_ligand(ligand_file)
-            for mol in mols:
+            for i, mol in enumerate(mols):
                 if mol:
-                    self.mol_group.append(Mol(mol, {"file_prefix": file_prefix}))
+                    self.mol_group.append(Mol(mol, {"file_prefix": f"{file_prefix}_{i}" if len(mols) > 1 
+                                                    else file_prefix}))
 
     def update_property_by_idx(self, idx: int, property_name: str, value: Any, is_conf_prop: bool = False):
         if is_conf_prop:

unidock_tools/tests/applications/test_unidock.py

@@ -1,23 +1,47 @@
 from pathlib import Path
 import os
 import shutil
+import glob
+import json
 import subprocess
 import pytest
 
 
 @pytest.fixture
 def receptor():
-    return Path(os.path.join(os.path.dirname(os.path.dirname(__file__)), "inputs", "1bcu_protein.pdb"))
+    return Path(os.path.join(os.path.dirname(os.path.dirname(__file__)), "inputs", 
+                             "unidock_pipeline", "1bcu", "protein.pdb"))
 
 
 @pytest.fixture
 def ligand():
-    return Path(os.path.join(os.path.dirname(os.path.dirname(__file__)), "inputs", "1bcu_ligand.sdf"))
+    return Path(os.path.join(os.path.dirname(os.path.dirname(__file__)), "inputs", 
+                             "unidock_pipeline", "1bcu", "ligand.sdf"))
 
 
 @pytest.fixture
 def pocket():
-    return [5.0, 15.0, 50.0, 15, 15, 15]
+    with open(os.path.join(os.path.dirname(os.path.dirname(__file__)), "inputs", 
+                           "unidock_pipeline", "1bcu", "docking_grid.json")) as f:
+        pocket = json.load(f)
+    return pocket
+
+
+testset_dir_path = Path(__file__).parent.parent / "inputs" / "unidock_pipeline"
+testset_name_list = ["bigsdf", "1bcu"]
+
+
+def get_docking_args(testset_name):
+    receptor = os.path.join(testset_dir_path, testset_name, "protein.pdb")
+    ligand = os.path.join(testset_dir_path, testset_name, "ligand.sdf")
+    with open(os.path.join(testset_dir_path, testset_name, "docking_grid.json")) as f:
+        pocket = json.load(f)
+    testset_info = {
+        "receptor": receptor,
+        "ligand": ligand,
+        "pocket": pocket,
+    }
+    return testset_info
 
 
 def read_scores(sdf_file, score_name):
@@ -31,39 +55,21 @@ def read_scores(sdf_file, score_name):
     return score_list
 
 
-def test_unidock_pipeline_default(receptor, ligand, pocket):
-    results_dir = "unidock_results"
-    cmd = f"unidocktools unidock_pipeline -r {receptor} -l {ligand} -sd {results_dir} \
-        -cx {pocket[0]} -cy {pocket[1]} -cz {pocket[2]} -sx {pocket[3]} -sy {pocket[4]} -sz {pocket[5]} \
-        -sf vina -nm 1"
-    print(cmd)
-    resp = subprocess.run(cmd, shell=True, capture_output=True, encoding="utf-8")
-    print(resp.stdout)
-    assert resp.returncode == 0, f"run unidock pipeline app err:\n{resp.stderr}"
-
-    result_file = os.path.join(results_dir, "1bcu_ligand.sdf")
-    assert os.path.exists(result_file), f"docking result file not found"
-
-    score_list = read_scores(result_file, "docking_score")
-    score = score_list[0]
-    assert -20 <= score <= 0, f"Uni-Dock score not in range: {score}"
-    shutil.rmtree(results_dir, ignore_errors=True)
-
-
 def test_unidock_pipeline_ligand_index(receptor, ligand, pocket):
     index_file = Path("ligand_index.txt")
     with open(index_file, "w") as f:
         f.write(str(ligand))
     results_dir = "unidock_results_input_index"
     cmd = f"unidocktools unidock_pipeline -r {receptor} -i {index_file} -sd {results_dir} \
-        -cx {pocket[0]} -cy {pocket[1]} -cz {pocket[2]} -sx {pocket[3]} -sy {pocket[4]} -sz {pocket[5]} \
-        -sf vina -nm 1"
+            -cx {pocket['center_x']} -cy {pocket['center_y']} -cz {pocket['center_z']} \
+            -sx {pocket['size_x']} -sy {pocket['size_y']} -sz {pocket['size_z']} \
+            -sf vina -nm 1 --seed 181129"
     print(cmd)
     resp = subprocess.run(cmd, shell=True, capture_output=True, encoding="utf-8")
     print(resp.stdout)
     assert resp.returncode == 0, f"run unidock pipeline app err:\n{resp.stderr}"
 
-    result_file = os.path.join(results_dir, "1bcu_ligand.sdf")
+    result_file = os.path.join(results_dir, Path(ligand).name)
     assert os.path.exists(result_file), f"docking result file not found"
 
     score_list = read_scores(result_file, "docking_score")
@@ -76,14 +82,15 @@ def test_unidock_pipeline_ligand_index(receptor, ligand, pocket):
 def test_unidock_pipeline_scoring_ad4(receptor, ligand, pocket):
     results_dir = "unidock_results_ad4"
     cmd = f"unidocktools unidock_pipeline -r {receptor} -l {ligand} -sd {results_dir} \
-        -cx {pocket[0]} -cy {pocket[1]} -cz {pocket[2]} -sx {pocket[3]} -sy {pocket[4]} -sz {pocket[5]} \
-        -sf ad4 -nm 1"
+            -cx {pocket['center_x']} -cy {pocket['center_y']} -cz {pocket['center_z']} \
+            -sx {pocket['size_x']} -sy {pocket['size_y']} -sz {pocket['size_z']} \
+            -sf ad4 -nm 1 --seed 181129"
     print(cmd)
     resp = subprocess.run(cmd, shell=True, capture_output=True, encoding="utf-8")
     print(resp.stdout)
     assert resp.returncode == 0, f"run unidock pipeline app err:\n{resp.stderr}"
 
-    result_file = os.path.join(results_dir, "1bcu_ligand.sdf")
+    result_file = os.path.join(results_dir, Path(ligand).name)
     assert os.path.exists(result_file), f"docking result file not found"
 
     score_list = read_scores(result_file, "docking_score")
@@ -95,18 +102,45 @@ def test_unidock_pipeline_scoring_ad4(receptor, ligand, pocket):
 def test_unidock_pipeline_multi_pose(receptor, ligand, pocket):
     results_dir = "unidock_results_multi_pose"
     cmd = f"unidocktools unidock_pipeline -r {receptor} -l {ligand} -sd {results_dir} \
-        -cx {pocket[0]} -cy {pocket[1]} -cz {pocket[2]} -sx {pocket[3]} -sy {pocket[4]} -sz {pocket[5]} \
-        -sf vina -nm 4"
+            -cx {pocket['center_x']} -cy {pocket['center_y']} -cz {pocket['center_z']} \
+            -sx {pocket['size_x']} -sy {pocket['size_y']} -sz {pocket['size_z']} \
+            -sf vina -nm 4 --seed 181129"
     print(cmd)
     resp = subprocess.run(cmd, shell=True, capture_output=True, encoding="utf-8")
     print(resp.stdout)
     assert resp.returncode == 0, f"run unidock pipeline app err:\n{resp.stderr}"
 
-    result_file = os.path.join(results_dir, "1bcu_ligand.sdf")
+    result_file = os.path.join(results_dir, Path(ligand).name)
     assert os.path.exists(result_file), f"docking result file not found"
 
     score_list = read_scores(result_file, "docking_score")
     assert len(score_list) == 4, f"docking result pose num({len(score_list)}) not match"
     for score in score_list:
         assert -20 <= score <= 0, f"Uni-Dock score not in range: {score}"
+    shutil.rmtree(results_dir, ignore_errors=True)
+
+
+@pytest.mark.parametrize("testset_name", testset_name_list)
+def test_unidock_pipeline_default_arg(testset_name):
+    testset_info = get_docking_args(testset_name)
+    receptor, ligand, pocket = testset_info["receptor"], testset_info["ligand"], testset_info["pocket"]
+    with open(ligand) as f:
+        total_num = len([line for line in f.readlines() if line.strip() == "$$$$"])
+    results_dir = f"unidock_results_{testset_name}"
+    cmd = f"unidocktools unidock_pipeline -r {receptor} -l {ligand} -sd {results_dir} \
+        -cx {pocket['center_x']} -cy {pocket['center_y']} -cz {pocket['center_z']} \
+        -sx {pocket['size_x']} -sy {pocket['size_y']} -sz {pocket['size_z']} \
+        -sf vina -nm 1 --seed 181129"
+    print(cmd)
+    resp = subprocess.run(cmd, shell=True, capture_output=True, encoding="utf-8")
+    print(resp.stdout)
+    assert resp.returncode == 0, f"run unidock pipeline app err:\n{resp.stderr}"
+
+    result_files = glob.glob(os.path.join(results_dir, "*.sdf"))
+    assert len(result_files) == total_num, f"failed to run all ligands"
+
+    for result_file in result_files:
+        score_list = read_scores(result_file, "docking_score")
+        score = score_list[0]
+        assert score <= 0, f"Uni-Dock score is abnormal"
     shutil.rmtree(results_dir, ignore_errors=True)

unidock_tools/tests/inputs/unidock_pipeline/1bcu/docking_grid.json

@@ -0,0 +1,8 @@
+{
+    "center_x": 5.0,
+    "center_y": 15.0,
+    "center_z": 50.0,
+    "size_x": 15,
+    "size_y": 15,
+    "size_z": 15
+}

unidock_tools/tests/inputs/unidock_pipeline/bigsdf/docking_grid.json

@@ -0,0 +1,8 @@
+{
+    "center_x": -16,
+    "center_y": -17,
+    "center_z": -26,
+    "size_x": 20,
+    "size_y": 28,
+    "size_z": 22
+}