flexible recommendations for ChorizoCreationError

meeko/chemtempgen.py

@@ -230,7 +230,8 @@ def get_smiles_with_atom_names(mol: Chem.Mol) -> tuple[str, list[str]]:
         smiles_atom_output_order = smiles_atom_output_order.replace(delimiter, ' ')
     smiles_output_order = (int(x) for x in smiles_atom_output_order.split())
 
-    atom_name = [mol.GetAtomWithIdx(atom_i).GetProp('atom_id') for atom_i in smiles_output_order]
+    atom_id_dict = {atom.GetIdx(): atom.GetProp('atom_id') for atom in mol.GetAtoms()}
+    atom_name = [atom_id_dict[atom_i] for atom_i in smiles_output_order]
 
     return smiles_exh, atom_name
 
@@ -537,8 +538,9 @@ def fetch_from_pdb(resname: str, max_retries = 5, backoff_factor = 2) -> str:
                 logging.info(f"Download failed for {resname}. Error: {e}. Retrying in {wait_time} seconds...")
                 time.sleep(wait_time)
             else:
-                logging.error(f"Max retries reached. Could not download CIF file for {resname}. Error: {e}")
-                return None
+                err = f"Max retries reached. Could not download CIF file for {resname}. Error: {e}"
+                logging.error(err)
+                raise ChemTempCreationError(err)
 
 # Constants for deprotonate
 acidic_proton_loc_canonical = {

meeko/linked_rdkit_chorizo.py

@@ -453,8 +453,29 @@ def _delete_residues(res_to_delete, raw_input_mols):
         raise ValueError(msg)
     return
 
+
 class ChorizoCreationError(RuntimeError):
-    pass
+
+    def __init__(self, error: str, recommendations: str = None): 
+        super().__init__(error) # main error message to pass to RuntimeError
+        self.error = error
+        self.recommendations = recommendations
+
+    def __str__(self):
+        msg = "" + os_linesep
+        msg += "Error: Creation of data structure for receptor failed." + os_linesep
+        msg += "" + os_linesep
+        msg += "Details:" + os_linesep
+        msg += self.error + os_linesep
+        msg += "" + os_linesep
+
+        if self.recommendations: 
+            msg += "Recommendations:" + os_linesep
+            msg += self.recommendations + os_linesep
+            msg += "" + os_linesep
+
+        return msg
+
 
 def handle_parsing_situations(
     unmatched_res,
@@ -495,7 +516,13 @@ def handle_parsing_situations(
         err += msg
 
     if err:
-        raise ChorizoCreationError(err)
+        recs = "1. (for batch processing) Use -a/--allow_bad_res to automatically remove residues" + os_linesep
+        recs += "that do not match templates, and --default_altloc to set" + os_linesep
+        recs += "a default altloc variant. Use these at your own risk." + os_linesep
+        recs += "" + os_linesep
+        recs += "2. (processing individual structure) Inspecting and fixing the input structure is recommended." + os_linesep
+        recs += "Use --wanted_altloc to set variants for specific residues."
+        raise ChorizoCreationError(err, recs)
     return
 
 
@@ -691,27 +718,38 @@ def __init__(
 
         # check if input assigned residue name in residue_templates
         err = ""
-        unknown_res_from_input = set(res_id for res_id in raw_input_mols if raw_input_mols[res_id][1] not in residue_templates.keys() | ambiguous.keys())
+        supported_resnames = residue_templates.keys()
+        unknown_res_from_input = {res_id: raw_input_mols[res_id][1] for res_id in raw_input_mols if raw_input_mols[res_id][1] not in supported_resnames}
         if len(unknown_res_from_input) > 0:
-            err += f"Input residue names {set(raw_input_mols[res_id][1] for res_id in unknown_res_from_input)} not in residue_templates" + os_linesep
+            err += f"Input residues {unknown_res_from_input} not in residue_templates" + os_linesep
         unknown_res_from_assign = set()
         if set_template:
-            unknown_res_from_assign = set(res_id for res_id in set_template if set_template[res_id] not in residue_templates.keys() | ambiguous.keys())
-            if len(unknown_res_from_input) > 0:
-                err += f"Assigned residue names {set(set_template[res_id] for res_id in unknown_res_from_assign)} not in residue_templates" + os_linesep
+            unknown_res_from_assign = {res_id: set_template[res_id] for res_id in set_template if set_template[res_id] not in supported_resnames}
+            if len(unknown_res_from_assign) > 0:
+                err += f"Assigned residues {unknown_res_from_assign} not in residue_templates" + os_linesep
         if err: 
             print(err)
-
-            unknown_res =  set(raw_input_mols[res_id][1] for res_id in unknown_res_from_input) | set(set_template[res_id] for res_id in unknown_res_from_assign)
-            for resname in unknown_res: 
-                cc = build_noncovalent_CC(resname)
-                fetch_template_dict = json.loads(export_chem_templates_to_json([cc]))['residue_templates'][resname]
-                residue_templates.update({resname: 
-                                          ResidueTemplate(
-                                              smiles = fetch_template_dict['smiles'],
-                                              atom_names = fetch_template_dict['atom_name'],
-                                              link_labels = fetch_template_dict['link_labels'])})
-            # raise ValueError(f"{err}")
+            print("Trying to resolve unknown residues by building chemical templates... " + os_linesep)
+
+            all_unknown_res = unknown_res_from_input.copy()
+            all_unknown_res.update(unknown_res_from_assign)
+
+            bonded_unknown_res = {res_id: all_unknown_res[res_id] for res_id in all_unknown_res if any(tup for tup in bonds if res_id in bonds)}
+            if bonded_unknown_res:
+                raise NotImplementedError(f"Unknown residues: {bonded_unknown_res} appear to be linking fragments. " + os_linesep
+                                          + "Guessing chemical templates with linker_labels are not currently supported. ")
+
+            try: 
+                for resname in all_unknown_res.values(): 
+                    cc = build_noncovalent_CC(resname)
+                    fetch_template_dict = json.loads(export_chem_templates_to_json([cc]))['residue_templates'][resname]
+                    residue_templates.update({resname: 
+                                            ResidueTemplate(
+                                                smiles = fetch_template_dict['smiles'],
+                                                atom_names = fetch_template_dict['atom_name'],
+                                                link_labels = fetch_template_dict['link_labels'])})
+            except Exception as e: 
+                raise ChorizoCreationError(str(e))
 
         self.residues, self.log = self._get_residues(
             raw_input_mols,

scripts/mk_prepare_receptor.py

@@ -520,20 +520,8 @@ def get_args():
                 wanted_altloc=wanted_altloc,
                 default_altloc=args.default_altloc,
             )
-        except ChorizoCreationError as err:
-            print()
-            print("Error: Creation of data structure for receptor failed.")
-            print()
-            print("Details: ")
-            print(err)
-            print()
-            print("Recommendations: ")
-            print("1. (for batch processing) Use -a/--allow_bad_res to automatically remove residues")
-            print("that do not match templates, and --default_altloc to set")
-            print("a default altloc variant. Use these at your own risk.")
-            print()
-            print("2. (processing individual structure) Inspecting and fixing the input structure is recommended.")
-            print("Use --wanted_altloc to set variants for specific residues.")
+        except ChorizoCreationError as e:
+            print(e)
             sys.exit(1)
 else:
     with open(args.read_pdb) as f:
@@ -550,17 +538,8 @@ def get_args():
             wanted_altloc=wanted_altloc,
             default_altloc=args.default_altloc,
         )
-    except ChorizoCreationError as err:
-        print()
-        print("Creation of data structure for receptor failed.")
-        print()
-        print(err)
-        print("Use -a/--allow_bad_res to automatically remove residues")
-        print("that do not match templates, and --default_altloc to set")
-        print("a default altloc variant. Use these at your own risk.")
-        print()
-        print("Inspecting and fixing the input structure is recommended.")
-        print("Use --wanted_altloc to set variants for specific residues.")
+    except ChorizoCreationError as e:
+        print(e)
         sys.exit(1)