Update inference_app.py

inference_app.py

@@ -53,7 +53,7 @@ def generate_input_conformer(
     return _mol
 
 
-def optimize_coordinate(points, bound_buffer=15, dmin=6.05):
+def optimize_coordinate(points, bound_buffer=15, dmin=6.02):
     bounds = list(
         zip(
             np.average(points, axis=0) - [bound_buffer]*3,
@@ -70,7 +70,7 @@ def optimize_coordinate(points, bound_buffer=15, dmin=6.05):
     return result.x, result.fun
 
 
-def optimize_decoy_coordinate(points, bound_buffer=15, dmin=6.05, decoy_min=4.0, decoy_max=4.99):
+def optimize_decoy_coordinate(points, bound_buffer=15, dmin=6.02, decoy_min=4.0, decoy_max=4.98):
     bounds = list(
         zip(
             np.average(points, axis=0) - [bound_buffer]*3,
@@ -117,7 +117,9 @@ def predict(input_sequence, input_ligand, input_msa, input_protein):
     molwriter = Chem.SDWriter("test_docking_pose.sdf")
     molwriter.write(mol)
 
-    mol_coords = mol.GetConformer().GetPositions()    
+    # get only non hydrogen atoms
+    heavy_atom_mask = [at.GetAtomicNum() != 1 for at in mol.GetAtoms()]
+    mol_coords = mol.GetConformer().GetPositions()[heavy_atom_mask]
     # get opt coords
     new_coord, min_dist_sum = optimize_coordinate(mol_coords)
     # get mindist to protein