Tests for top N ranking post script, implements #21

Author: brianjimenez

bin/post/lgd_top.py

@@ -6,13 +6,14 @@
 import os
 import numpy as np
 from lightdock.mathutil.cython.quaternion import Quaternion
-from lightdock.util.analysis import read_lightdock_output
+from lightdock.util.analysis import read_ranking_file
 from lightdock.util.logger import LoggingManager
 from lightdock.constants import DEFAULT_NMODES_REC, DEFAULT_NMODES_LIG, DEFAULT_REC_NM_FILE, DEFAULT_LIG_NM_FILE
 from lightdock.pdbutil.PDBIO import parse_complex_from_file, write_pdb_to_file
 from lightdock.structure.complex import Complex
 from lightdock.structure.nm import read_nmodes
 from lightdock.util.parser import CommandLineParser, get_lightdock_structures
+from lightdock.prep.simulation import get_setup_from_file
 
 
 log = LoggingManager.get_logger('lightdock_top')
@@ -33,9 +34,22 @@
     # Number of structures to generate
     parser.add_argument("top", help="number of structures to generate", type=CommandLineParser.valid_integer_number,
                         metavar="top")
+    # Optional, setup file
+    parser.add_argument("--setup", "-setup", "-s", help="Simulation setup file",
+                            dest="setup_file", metavar="setup_file", type=CommandLineParser.valid_file, 
+                            default=None)
 
     args = parser.parse_args()
 
+    # Load setup configuration if provided
+    setup = get_setup_from_file(args.setup_file) if args.setup_file else None
+
+    num_anm_rec = DEFAULT_NMODES_REC
+    num_anm_lig = DEFAULT_NMODES_LIG
+    if setup and setup['use_anm']:
+        num_anm_rec = setup['anm_rec']
+        num_anm_lig = setup['anm_lig']
+
     # Receptor
     structures = []
     for structure in get_lightdock_structures(args.receptor_structures):
@@ -55,7 +69,7 @@
     ligand = Complex.from_structures(structures)
 
     # Read ranking file
-    glowworms = read_lightdock_output(args.lightdock_ranking_file, initial=1, final=args.top)
+    predictions = read_ranking_file(args.lightdock_ranking_file)
 
     # Destination path is the same as the lightdock output
     destination_path = os.path.dirname(args.lightdock_ranking_file)
@@ -72,24 +86,25 @@
     if os.path.exists(nm_lig_file):
         nmodes_lig = read_nmodes(nm_lig_file)
 
-    for i, glowworm in enumerate(glowworms):
-        receptor_pose = receptor.atom_coordinates[glowworm.receptor_id].clone()
-        ligand_pose = ligand.atom_coordinates[glowworm.ligand_id].clone()
-        # Use normal modes if provided:
-        if nmodes_rec.any():
-            for nm in range(DEFAULT_NMODES_REC):
-                rec_extent = np.array([float(x) for x in glowworm.pose[7:7 + DEFAULT_NMODES_REC]])
-                receptor_pose.coordinates += nmodes_rec[nm] * rec_extent[nm]
-        if nmodes_lig.any():
-            for nm in range(DEFAULT_NMODES_LIG):
-                lig_extent = np.array([float(x) for x in glowworm.pose[-DEFAULT_NMODES_LIG:]])
-                ligand_pose.coordinates += nmodes_lig[nm] * lig_extent[nm]
-
-        # We rotate first, ligand it's at initial position
-        rotation = Quaternion(glowworm.pose[3], glowworm.pose[4], glowworm.pose[5], glowworm.pose[6])
-        ligand_pose.rotate(rotation)
-        ligand_pose.translate([glowworm.pose[0], glowworm.pose[1], glowworm.pose[2]])
-
-        write_pdb_to_file(receptor, os.path.join(destination_path, 'top_%s.pdb' % str(i+1)), receptor_pose)
-        write_pdb_to_file(ligand,  os.path.join(destination_path, 'top_%s.pdb' % str(i+1)), ligand_pose)
+    for i, glowworm in enumerate(predictions):
+        if i < args.top:
+            receptor_pose = receptor.atom_coordinates[glowworm.receptor_id].clone()
+            ligand_pose = ligand.atom_coordinates[glowworm.ligand_id].clone()
+            # Use normal modes if provided:
+            if nmodes_rec.any():
+                for nm in range(num_anm_rec):
+                    rec_extent = np.array([float(x) for x in glowworm.pose[7:7 + num_anm_rec]])
+                    receptor_pose.coordinates += nmodes_rec[nm] * rec_extent[nm]
+            if nmodes_lig.any():
+                for nm in range(num_anm_lig):
+                    lig_extent = np.array([float(x) for x in glowworm.pose[-num_anm_lig:]])
+                    ligand_pose.coordinates += nmodes_lig[nm] * lig_extent[nm]
+
+            # We rotate first, ligand it's at initial position
+            rotation = Quaternion(glowworm.pose[3], glowworm.pose[4], glowworm.pose[5], glowworm.pose[6])
+            ligand_pose.rotate(rotation)
+            ligand_pose.translate([glowworm.pose[0], glowworm.pose[1], glowworm.pose[2]])
+
+            write_pdb_to_file(receptor, os.path.join(destination_path, 'top_%s.pdb' % str(i+1)), receptor_pose)
+            write_pdb_to_file(ligand,  os.path.join(destination_path, 'top_%s.pdb' % str(i+1)), ligand_pose)
     log.info("Generated %d conformations" % args.top)