1717from datasets .constants import aa_short2long , atom_order , three_to_one
1818from datasets .parse_chi import get_chi_angles , get_coords , aa_idx2aa_short , get_onehot_sequence
1919from utils .torsion import get_transformation_mask
20+ from utils .logging_utils import get_logger
2021
2122
2223periodic_table = GetPeriodicTable ()
@@ -305,11 +306,11 @@ def generate_conformer(mol):
305306 failures , id = 0 , - 1
306307 while failures < 3 and id == - 1 :
307308 if failures > 0 :
308- print (f'rdkit coords could not be generated. trying again { failures } .' )
309+ get_logger (). debug (f'rdkit coords could not be generated. trying again { failures } .' )
309310 id = AllChem .EmbedMolecule (mol , ps )
310311 failures += 1
311312 if id == - 1 :
312- print ('rdkit coords could not be generated without using random coords. using random coords now.' )
313+ get_logger (). info ('rdkit coords could not be generated without using random coords. using random coords now.' )
313314 ps .useRandomCoords = True
314315 AllChem .EmbedMolecule (mol , ps )
315316 AllChem .MMFFOptimizeMolecule (mol , confId = 0 )
@@ -417,6 +418,7 @@ def write_mol_with_coords(mol, new_coords, path):
417418 w .write (mol )
418419 w .close ()
419420
421+
420422def read_molecule (molecule_file , sanitize = False , calc_charges = False , remove_hs = False ):
421423 if molecule_file .endswith ('.mol2' ):
422424 mol = Chem .MolFromMol2File (molecule_file , sanitize = False , removeHs = False )
@@ -433,8 +435,8 @@ def read_molecule(molecule_file, sanitize=False, calc_charges=False, remove_hs=F
433435 elif molecule_file .endswith ('.pdb' ):
434436 mol = Chem .MolFromPDBFile (molecule_file , sanitize = False , removeHs = False )
435437 else :
436- return ValueError ('Expect the format of the molecule_file to be '
437- 'one of .mol2, .sdf, .pdbqt and .pdb, got {}' .format (molecule_file ))
438+ raise ValueError ('Expect the format of the molecule_file to be '
439+ 'one of .mol2, .sdf, .pdbqt and .pdb, got {}' .format (molecule_file ))
438440
439441 try :
440442 if sanitize or calc_charges :
@@ -449,7 +451,12 @@ def read_molecule(molecule_file, sanitize=False, calc_charges=False, remove_hs=F
449451
450452 if remove_hs :
451453 mol = Chem .RemoveHs (mol , sanitize = sanitize )
452- except :
454+
455+ except Exception as e :
456+ # Print stacktrace
457+ import traceback
458+ msg = traceback .format_exc ()
459+ get_logger ().warning (f"Failed to process molecule: { molecule_file } \n { msg } " )
453460 return None
454461
455462 return mol
0 commit comments