ClustENM: multi-start (setAtoms accepts a list of AtomGroups) + optional fixer

prody/apps/prody_apps/prody_clustenm.py

@@ -31,6 +31,10 @@
     ('threshold', 'RMSD threshold to apply when forming clusters, can be tuple of floats', 'None'),
     ('no_sim', 'whether a short MD simulation is not performed after energy minimization (otherwise it is)', False),
     ('parallel', 'whether conformer generation will be parallelized', False),
+    ('parallel_sim', 'number of worker processes to parallelize per-conformer minimisation/MD across '
+                     '(0 means all CPUs; absent/False means off); independent of --parallel', False),
+    ('sim_devices', 'comma-separated physical GPU indices to round-robin the parallel-sim workers '
+                    'across, e.g. 0,1 (default: none, no GPU pinning)', None),
     ('v1', 'whether to use original sampling method with complete enumeration of ANM modes', False),
     ('no_outlier', 'whether to not exclude outliers in each generation when using implicit solvent (always False for explicit)', False),
     ('mzscore', 'modified z-score threshold to label conformers as outliers', 3.5),
@@ -101,6 +105,10 @@ def prody_clustenm(pdb, **kwargs):
     sim = not kwargs.pop('no_sim')
     temp = kwargs.pop('temp')
     parallel = kwargs.pop('parallel')
+    parallel_sim = kwargs.pop('parallel_sim')
+    sim_devices = kwargs.pop('sim_devices')
+    if isinstance(sim_devices, str):
+        sim_devices = [int(x) for x in sim_devices.split(',') if x.strip() != '']
     write_params = kwargs.pop("write_params")
 
     solvent = kwargs.pop('solvent')
@@ -110,17 +118,28 @@ def prody_clustenm(pdb, **kwargs):
     outlier = not kwargs.pop('no_outlier')
     mzscore = kwargs.pop('mzscore')
 
-    pdb = prody.parsePDB(pdb, model=model, altloc=altloc)
+    # a single PDB identifier/filename, or several comma-separated ones -> multi-start ClustENM.
+    # The extra structures seed the initial population as gen-1-style conformers; after selection
+    # they must share the topology of the first. Mirrors setAtoms() accepting a list of AtomGroups.
+    pdbs = [p.strip() for p in pdb.split(',') if p.strip()]
+    ags = [prody.parsePDB(p, model=model, altloc=altloc) for p in pdbs]
+    pdb = ags[0]
     if prefix == '_clustenm':
         prefix = pdb.getTitle() + '_clustenm'
 
-    select = pdb.select(selstr)
-    if select is None:
-        LOGGER.warn('Selection {0} did not match any atoms.'
-                    .format(repr(selstr)))
-        return
-    LOGGER.info('{0} atoms will be used for ClustENM calculations.'
-                .format(len(select)))
+    selects = []
+    for ag in ags:
+        sel = ag.select(selstr)
+        if sel is None:
+            LOGGER.warn('Selection {0} did not match any atoms in {1}.'
+                        .format(repr(selstr), ag.getTitle()))
+            return
+        selects.append(sel)
+    select = selects[0]
+    LOGGER.info('{0} atoms will be used for ClustENM calculations{1}.'
+                .format(len(select),
+                        ' (%d starting structures)' % len(selects)
+                        if len(selects) > 1 else ''))
 
     try:
         gamma = float(kwargs.pop('gamma'))
@@ -149,7 +168,7 @@ def prody_clustenm(pdb, **kwargs):
             raise TypeError("Please provide cutoff as a float or equation using math")
 
     ens = prody.ClustENM(pdb.getTitle())
-    ens.setAtoms(select)
+    ens.setAtoms(selects if len(selects) > 1 else select)
     ens.run(n_gens=ngens, n_modes=nmodes,
             n_confs=nconfs, rmsd=eval(rmsd),
             cutoff=cutoff, gamma=gamma,
@@ -159,7 +178,8 @@ def prody_clustenm(pdb, **kwargs):
             t_steps_g=eval(t_steps_g),
             outlier=outlier, mzscore=mzscore,
             sparse=sparse, kdtree=kdtree, turbo=turbo,
-            parallel=parallel, fitmap=fitmap,
+            parallel=parallel, parallel_sim=parallel_sim,
+            sim_devices=sim_devices, fitmap=fitmap,
             fit_resolution=fit_resolution,
             nproc=nproc, **kwargs)
 
@@ -207,7 +227,13 @@ def addCommand(commands):
 carbon alpha atoms with residue numbers less than 70, and save all of the
 graphical output files:
 
-  $ prody clustenm 1aar -s "calpha and chain A and resnum < 70" -A""",
+  $ prody clustenm 1aar -s "calpha and chain A and resnum < 70" -A
+
+Start ClustENM(D) from multiple structures (comma-separated identifiers or
+filenames sharing the same topology after selection), seeding the initial
+population with all of them:
+
+  $ prody clustenm model1.pdb,model2.pdb,model3.pdb""",
   test_examples=[0, 1])
 
     group = addNMAParameters(subparser, include_nproc=True)
@@ -265,7 +291,15 @@ def addCommand(commands):
     group.add_argument('-B', '--parallel', dest='parallel',
         action='store_true',
         default=DEFAULTS['parallel'], help=HELPTEXT['parallel'])
-
+
+    group.add_argument('--parallel_sim', dest='parallel_sim', type=int,
+        default=DEFAULTS['parallel_sim'], metavar='INT',
+        help=HELPTEXT['parallel_sim'] + ' (default: %(default)s)')
+
+    group.add_argument('--sim_devices', dest='sim_devices', type=str,
+        default=DEFAULTS['sim_devices'], metavar='STR',
+        help=HELPTEXT['sim_devices'] + ' (default: %(default)s)')
+
     group.add_argument('-E', '--write_params', dest='write_params',
         action='store_true',
         default=DEFAULTS['write_params'], help=HELPTEXT['write_params'])
@@ -350,7 +384,9 @@ def addCommand(commands):
         default=DEFAULTS['platform'], metavar='STR',
         help=HELPTEXT['platform'] + ' (default: %(default)s)')
 
-    subparser.add_argument('pdb', help='PDB identifier or filename')
+    subparser.add_argument('pdb', help='PDB identifier or filename; or several comma-separated '
+        'identifiers/filenames (sharing topology after selection) to start ClustENM from multiple '
+        'structures')
 
     subparser.set_defaults(func=lambda ns: prody_clustenm(ns.__dict__.pop('pdb'),
                                                           **ns.__dict__))