add: use_3d opt for ligand

PaddlePaddle · Aug 20, 2024 · d7bd9ee · d7bd9ee
1 parent 0720311
commit d7bd9ee
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 8 deletions.
diff --git a/apps/protein_folding/helixfold3/data/demo_3fap_protein_sm.json b/apps/protein_folding/helixfold3/data/demo_3fap_protein_sm.json
@@ -12,7 +12,8 @@
         },
         {
             "type": "ligand",
-            "smiles": "Cc1ccc(s1)[C@@H]\\2C[C@@H]3CC[C@H]([C@@](O3)(C(=O)C(=O)N4CCCC[C@H]4C(=O)O[C@@H](CC(=O)[C@@H](/C=C(/[C@H]([C@H](C(=O)[C@@H](C[C@@H](/C=C/C=C/C=C2\\C)C)C)OC)O)\\C)C)[C@H](C)C[C@@H]5CC[C@H]([C@@H](C5)OC)O)O)C",
+            "sdf": "/mnt/data/yinying/tests/helixfold/ligands/ARD_ideal.sdf",
+            "use_3d": false,
             "count": 1
         }
     ]

diff --git a/apps/protein_folding/helixfold3/helixfold/infer_scripts/preprocess.py b/apps/protein_folding/helixfold3/helixfold/infer_scripts/preprocess.py
@@ -164,14 +164,14 @@ def _get_supported_formats(self) -> Tuple[str]:
 
         return tuple(formats)
 
-    def _perform_conversion(self, input_type: str, input_value: str) -> Chem.Mol:
+    def _perform_conversion(self, input_type: str, input_value: str, generate_3d: bool=True) -> Chem.Mol:
         with tempfile.NamedTemporaryFile(suffix=".mol2") as temp_file, utils.timing(f'converting {input_type} to mol2: {input_value}'):
             if input_type == 'smiles':
-                obabel_cmd = f"{self.obabel_bin} -:'{input_value}' -omol2 -O{temp_file.name} --gen3d"
+                obabel_cmd = f"{self.obabel_bin} -:'{input_value}' -omol2 -O{temp_file.name} {'--gen3d' if generate_3d else ''}"
                 if len(input_value)>60:
                     logging.warning(f'This takes a while ...')
             else:
-                obabel_cmd = f"{self.obabel_bin} -i {input_type} {input_value} -omol2 -O{temp_file.name} --gen3d"
+                obabel_cmd = f"{self.obabel_bin} -i {input_type} {input_value} -omol2 -O{temp_file.name} {'--gen3d' if generate_3d else ''}"
             logging.debug(f'Launching command: `{obabel_cmd}`')
             ret = subprocess.run(obabel_cmd, shell=True, capture_output=True, text=True)
             mol = Chem.MolFromMol2File(temp_file.name, sanitize=False)
@@ -181,16 +181,16 @@ def _perform_conversion(self, input_type: str, input_value: str) -> Chem.Mol:
 
             return optimal_mol_wo_H
 
-    def _convert_to_mol(self, input_type: str, input_value: str) -> Chem.Mol:
+    def _convert_to_mol(self, input_type: str, input_value: str, generate_3d: bool=True) -> Chem.Mol:
         if input_type not in self.supported_formats:
             raise ValueError(f'Unsupported small molecule input: {input_type}. \nSupported formats: \n{self.supported_formats}\n')
 
         if input_type != 'smiles' and not os.path.isfile(input_value):
             raise FileNotFoundError(f'Cannot find the {input_type.upper()} file at {input_value}.')
 
-        return self._perform_conversion(input_type, input_value)
+        return self._perform_conversion(input_type, input_value, generate_3d)
 
-    __call__: Callable[[str, str], Chem.Mol] = _convert_to_mol
+    __call__: Callable[[str, str, bool], Chem.Mol] = _convert_to_mol
 def polymer_convert(items):
     """
         "type": "protein",                          
@@ -251,7 +251,7 @@ def ligand_convert(items: Mapping[str, Union[int, str]]):
         _ccd_seqs.append(f"({ligand_name})")
         # mol_wo_h = smiles_to_ETKDGMol(items['smiles'])
 
-        mol_wo_h = converter(k, items[k])
+        mol_wo_h = converter(k, items[k], items.get('use_3d', True))
         _extra_mol_infos = make_basic_info_fromMol(mol_wo_h)
         ccd_to_extra_mol_infos = {
             ligand_name: _extra_mol_infos