From 5461b0f07a2ee6f73e6e54e2a9526ec0cee0a5af Mon Sep 17 00:00:00 2001 From: James Krieger Date: Tue, 21 Jul 2026 12:10:09 +0100 Subject: [PATCH 1/7] ClustENM: multi-start (setAtoms accepts a list of AtomGroups) + fix=False option setAtoms now accepts a list/tuple of AtomGroups to seed a MULTI-START run: the topology is built from the first structure and a matching coordinate set is collected from each (via _fix_multi, same PDBFixer processing so atom order matches). They become the generation-0 initial population -- each minimised -- so a run can start from a SET of structures rather than one; equivalent to skipping the single-conformer gen-0. The generation loop is unchanged and then samples normal modes from all of them. At n_gens=0 this is a batched relax that keeps the originals (no clustering). setAtoms(..., fix=True) gains fix=False: skip PDBFixer and build the OpenMM topology directly from atoms as-is (only valid when already complete + protonated; forcefield build fails otherwise). Added _nofix for that path. Validated: setAtoms([s0,s1]) + run(n_gens=0) minimises both (clash 318/340 -> 0, NBD/%SS preserved). Single-structure path byte-unchanged. Co-Authored-By: Claude Opus 4.8 --- prody/dynamics/clustenm.py | 139 +++++++++++++++++++++++++++++++------ 1 file changed, 119 insertions(+), 20 deletions(-) diff --git a/prody/dynamics/clustenm.py b/prody/dynamics/clustenm.py index f5d5c712e..f3da12eb9 100644 --- a/prody/dynamics/clustenm.py +++ b/prody/dynamics/clustenm.py @@ -136,17 +136,48 @@ def _isBuilt(self): return self._confs is not None - def setAtoms(self, atoms, pH=7.0): + def setAtoms(self, atoms, pH=7.0, fix=True): ''' Sets atoms. - - :arg atoms: *atoms* parsed by parsePDB + + :arg atoms: a single AtomGroup parsed by parsePDB, OR a list/tuple of AtomGroups to seed a + MULTI-START run. In the multi-start case the topology is built from the first structure and a + matching coordinate set is collected from each; they become the generation-0 initial + population (each is minimised), so a run can start from a *set* of structures rather than one. + The members must be the same molecule (identical atom count after preparation). :arg pH: pH based on which to select protonation states for adding missing hydrogens, default is 7.0. :type pH: float + + :arg fix: run PDBFixer (replace nonstandard residues, add missing atoms + hydrogens) via _fix. + Default True. Set False to SKIP PDBFixer and build the OpenMM topology directly from *atoms* + as-is -- only valid when *atoms* is already complete AND protonated; the forcefield build will + fail otherwise (e.g. heavy-atom-only inputs must keep fix=True). + :type fix: bool ''' + # Multi-start: a list/tuple of AtomGroups -> set up the topology from the first, then collect one + # topology-matching coordset from each (the gen-0 initial population). + if isinstance(atoms, (list, tuple)): + ags = list(atoms) + if len(ags) == 0: + raise ValueError('setAtoms received an empty list of structures') + self.setAtoms(ags[0], pH=pH, fix=fix) # single-structure setup -> topology + self._atoms + coordsets = [self._atoms.getCoords()] + for ag in ags[1:]: + c = self._fix_multi(ag, pH, fix) + if c.shape[0] != self._n_atoms: + raise ValueError('multi-start member %r has %d atoms after preparation; expected %d ' + '(all members must be the same molecule)' + % (ag.getTitle(), c.shape[0], self._n_atoms)) + coordsets.append(c) + self._start_coordsets = coordsets + LOGGER.info('Multi-start: %d initial structures set.' % len(coordsets)) + return + + self._start_coordsets = None + atoms = atoms.select('not hetatm') self._nuc = atoms.select('nucleotide') @@ -167,12 +198,16 @@ def setAtoms(self, atoms, pH=7.0): if self._isBuilt(): super(ClustENM, self).setAtoms(atoms) else: - LOGGER.info('Fixing the structure ...') - LOGGER.timeit('_clustenm_fix') self._ph = pH - self._fix(atoms) - LOGGER.report('The structure was fixed in %.2fs.', - label='_clustenm_fix') + if fix: + LOGGER.info('Fixing the structure ...') + LOGGER.timeit('_clustenm_fix') + self._fix(atoms) + LOGGER.report('The structure was fixed in %.2fs.', + label='_clustenm_fix') + else: + LOGGER.info('Skipping PDBFixer (fix=False); building topology from atoms as-is ...') + self._nofix(atoms) if self._nuc is None: self._idx_cg = self._atoms.ca.getIndices() @@ -245,6 +280,64 @@ def _fix(self, atoms): self._topology = fixed.topology self._positions = fixed.positions + def _nofix(self, atoms): + + # Build the OpenMM topology/positions directly from atoms, SKIPPING PDBFixer. Requires atoms to be + # already complete AND protonated (the forcefield build fails otherwise). Mirrors _fix's tail only. + + try: + from openmm.app import PDBFile + except ImportError: + raise ImportError('Please install PDBFixer and OpenMM 7.6 in order to use ClustENM.') + + title = atoms.getTitle() + stream = createStringIO() + writePDBStream(stream, atoms) + stream.seek(0) + pdb = PDBFile(stream) + stream.close() + + self._atoms = atoms.copy() + self._atoms.setTitle(title) + self._topology = pdb.topology + self._positions = pdb.positions + + def _fix_multi(self, atoms, pH, fix): + + # Return coords of `atoms` after the SAME (optional) PDBFixer processing as _fix, so they match the + # topology built from the reference structure -- used to collect a multi-start population. + + atoms = atoms.select('not hetatm') + if not fix: + return atoms.getCoords() + + try: + from pdbfixer import PDBFixer + from openmm.app import PDBFile + except ImportError: + raise ImportError('Please install PDBFixer and OpenMM 7.6 in order to use ClustENM.') + + stream = createStringIO() + writePDBStream(stream, atoms) + stream.seek(0) + fixed = PDBFixer(pdbfile=stream) + stream.close() + + fixed.missingResidues = {} + fixed.findNonstandardResidues() + fixed.replaceNonstandardResidues() + fixed.removeHeterogens(False) + fixed.findMissingAtoms() + fixed.addMissingAtoms() + fixed.addMissingHydrogens(pH) + + out = createStringIO() + PDBFile.writeFile(fixed.topology, fixed.positions, out, keepIds=True) + out.seek(0) + ag = parsePDBStream(out) + out.close() + return ag.getCoords() + def _prep_sim(self, coords, external_forces=[]): try: @@ -1256,24 +1349,30 @@ def run(self, cutoff=15., n_modes=3, gamma=1., n_confs=50, rmsd=1.0, else: LOGGER.info('Minimization in generation 0 ...') LOGGER.timeit('_clustenm_min') - potential, conformer = self._min_sim(self._atoms.getCoords()) - if np.isnan(potential): - raise ValueError('Initial structure could not be minimized. Try again and/or check your structure.') + # Multi-start: skip the single-conformer gen-0 and seed directly with the N provided structures + # (each minimised) as the initial population; equivalent to gen-0 when there is only one. The + # generation loop below then samples normal modes from all of them. (Single-structure default = + # the one reference conformer.) + starts = self._start_coordsets if getattr(self, '_start_coordsets', None) else [self._atoms.getCoords()] + if len(starts) > 1: + LOGGER.info('Multi-start: minimising %d initial structures (gen-0 single-conformer step skipped) ...' + % len(starts)) + pot_conf = [self._min_sim(sc) for sc in starts] + pots0 = [pc[0] for pc in pot_conf] + if np.any(np.isnan(pots0)): + raise ValueError('An initial structure could not be minimized. Try again and/or check your structure(s).') LOGGER.report(label='_clustenm_min') LOGGER.info('#' + '-' * 19 + '/*\\' + '-' * 19 + '#') - self.setCoords(conformer) + confs0 = np.array([pc[1] for pc in pot_conf]) + self.setCoords(confs0[0]) - potentials = [potential] - sizes = [1] - new_shape = [1] - for s in conformer.shape: - new_shape.append(s) - conf = conformer.reshape(new_shape) - conformers = start_confs = conf - keys = [(0, 0)] + potentials = list(pots0) + sizes = [1] * len(pots0) + conformers = start_confs = confs0 + keys = [(0, j) for j in range(len(pots0))] for i in range(1, self._n_gens+1): self._cycle += 1 From 12213494503819339ec2a1898f834334d5b9e255 Mon Sep 17 00:00:00 2001 From: James Krieger Date: Tue, 21 Jul 2026 13:03:47 +0100 Subject: [PATCH 2/7] ClustENM: parallelise minimise/MD across conformers + build-once implicit-solvent system _min_sim_batch minimises (and runs MD, since _min_sim does both) a batch of coordsets, parallelised across conformers with a multiprocessing Pool when `parallel` is set -- mirroring the existing _generate/_sample parallelisation. Both the generation-loop minimise and the multi-start gen-0 seed now go through it, so gen-1+ sampling AND multi-start relaxes run in parallel. Order-preserved, same results as serial. Also: for implicit solvent the OpenMM system depends only on the (fixed) topology, so _min_sim caches the Simulation and reuses its Context (setPositions per conformer) instead of rebuilding createSystem each call. Explicit solvent (per-conformer solvent box) is not cached. Fork-safe: the cache is None at Pool-fork time, so each worker builds its own. NOTE: with a single GPU, parallel workers contend on it; use platform='CPU' (with enough allocated cores) for cross-process parallelism. Perf to be benchmarked on a proper cluster allocation. Co-Authored-By: Claude Opus 4.8 --- prody/dynamics/clustenm.py | 35 +++++++++++++++++++++++++++++++---- 1 file changed, 31 insertions(+), 4 deletions(-) diff --git a/prody/dynamics/clustenm.py b/prody/dynamics/clustenm.py index f3da12eb9..2692db4ce 100644 --- a/prody/dynamics/clustenm.py +++ b/prody/dynamics/clustenm.py @@ -411,12 +411,25 @@ def _min_sim(self, coords): # coords: coordset (numAtoms, 3) in Angstrom, which should be converted into nanometer try: + from openmm import Vec3 from openmm.app import StateDataReporter - from openmm.unit import kelvin, angstrom, nanometer, kilojoule_per_mole, MOLAR_GAS_CONSTANT_R + from openmm.unit import kelvin, angstrom, nanometer, kilojoule_per_mole, MOLAR_GAS_CONSTANT_R, Quantity except ImportError: raise ImportError('Please install PDBFixer and OpenMM 7.6 in order to use ClustENM.') - simulation = self._prep_sim(coords=coords) + # Build-once: an implicit-solvent system depends only on the (fixed) topology, so cache the + # Simulation and reuse its Context across conformers (setPositions per call) instead of rebuilding + # createSystem every call -- speeds up multi-conformer generations and multi-start. Explicit + # solvent adds a per-conformer solvent box, so it is NOT cached (rebuilt each call, as before). + if self._sol == 'imp' and getattr(self, '_sim_cache', None) is not None: + simulation = self._sim_cache + simulation.context.setPositions(Quantity([Vec3(*xyz) for xyz in coords], angstrom)) + if self._sim: + simulation.context.setVelocitiesToTemperature(0) # fresh start for the heating loop + else: + simulation = self._prep_sim(coords=coords) + if self._sol == 'imp': + self._sim_cache = simulation # automatic conversion into nanometer will be carried out. # simulation.context.setPositions(coords * angstrom) @@ -450,6 +463,19 @@ def _min_sim(self, coords): return np.nan, np.full_like(coords, np.nan) + def _min_sim_batch(self, coords_list): + + # Minimise (+ optional heat/sim) a batch of coordsets. Parallelised across conformers with a + # multiprocessing Pool when self._parallel is set (mirrors the _generate/_sample parallelisation), + # otherwise serial. Returns a list of (potential, coords) aligned with coords_list. + + coords_list = list(coords_list) + if self._parallel and len(coords_list) > 1: + repeats = cpu_count() if self._parallel is True else int(self._parallel) + with Pool(repeats) as p: + return p.map(self._min_sim, coords_list) + return [self._min_sim(c) for c in coords_list] + def _targeted_sim(self, coords0, coords1, tmdk=15., d_steps=100, n_max_steps=10000, ddtol=1e-3, n_conv=5): try: @@ -1348,6 +1374,7 @@ def run(self, cutoff=15., n_modes=3, gamma=1., n_confs=50, rmsd=1.0, LOGGER.info('Minimization & heating-up in generation 0 ...') else: LOGGER.info('Minimization in generation 0 ...') + self._sim_cache = None # build-once OpenMM sim cache (imp solvent); rebuilt fresh each run LOGGER.timeit('_clustenm_min') # Multi-start: skip the single-conformer gen-0 and seed directly with the N provided structures # (each minimised) as the initial population; equivalent to gen-0 when there is only one. The @@ -1357,7 +1384,7 @@ def run(self, cutoff=15., n_modes=3, gamma=1., n_confs=50, rmsd=1.0, if len(starts) > 1: LOGGER.info('Multi-start: minimising %d initial structures (gen-0 single-conformer step skipped) ...' % len(starts)) - pot_conf = [self._min_sim(sc) for sc in starts] + pot_conf = self._min_sim_batch(starts) pots0 = [pc[0] for pc in pot_conf] if np.any(np.isnan(pots0)): raise ValueError('An initial structure could not be minimized. Try again and/or check your structure(s).') @@ -1387,7 +1414,7 @@ def run(self, cutoff=15., n_modes=3, gamma=1., n_confs=50, rmsd=1.0, LOGGER.info('Minimization in generation %d ...' % i) LOGGER.timeit('_clustenm_min_sim') - pot_conf = [self._min_sim(conf) for conf in confs] + pot_conf = self._min_sim_batch(confs) LOGGER.report('Structures were sampled in %.2fs.', label='_clustenm_min_sim') From 80667183a7a8108731baec07ce42786f6e82d867 Mon Sep 17 00:00:00 2001 From: James Krieger Date: Tue, 21 Jul 2026 13:12:34 +0100 Subject: [PATCH 3/7] ClustENM: round-robin Pool workers across visible GPUs for multi-GPU parallel _worker_gpu_init pins each Pool worker to one GPU (round-robin over CUDA_VISIBLE_DEVICES) so a multi-GPU parallel run spreads across GPUs instead of all workers contending on GPU 0. No-op when <2 GPUs are visible, so CPU-parallel is unaffected. Co-Authored-By: Claude Opus 4.8 --- prody/dynamics/clustenm.py | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/prody/dynamics/clustenm.py b/prody/dynamics/clustenm.py index 2692db4ce..bbbfd79ea 100644 --- a/prody/dynamics/clustenm.py +++ b/prody/dynamics/clustenm.py @@ -463,16 +463,34 @@ def _min_sim(self, coords): return np.nan, np.full_like(coords, np.nan) + @staticmethod + def _worker_gpu_init(): + + # Pin each Pool worker to ONE visible GPU (round-robin over CUDA_VISIBLE_DEVICES) so a multi-GPU + # parallel run spreads across the GPUs instead of all contending on GPU 0. No-op when fewer than + # two GPUs are visible (e.g. the CPU platform), so it is harmless for CPU-parallel runs. + + import os + from multiprocessing import current_process + gpus = [g for g in os.environ.get('CUDA_VISIBLE_DEVICES', '').split(',') if g != ''] + if len(gpus) > 1: + try: + rank = int(current_process().name.rsplit('-', 1)[-1]) - 1 + except Exception: + rank = 0 + os.environ['CUDA_VISIBLE_DEVICES'] = gpus[rank % len(gpus)] + def _min_sim_batch(self, coords_list): # Minimise (+ optional heat/sim) a batch of coordsets. Parallelised across conformers with a # multiprocessing Pool when self._parallel is set (mirrors the _generate/_sample parallelisation), - # otherwise serial. Returns a list of (potential, coords) aligned with coords_list. + # otherwise serial. Workers round-robin across visible GPUs (multi-GPU) via _worker_gpu_init. + # Returns a list of (potential, coords) aligned with coords_list. coords_list = list(coords_list) if self._parallel and len(coords_list) > 1: repeats = cpu_count() if self._parallel is True else int(self._parallel) - with Pool(repeats) as p: + with Pool(repeats, initializer=self._worker_gpu_init) as p: return p.map(self._min_sim, coords_list) return [self._min_sim(c) for c in coords_list] From c8cad9cc78bb9c3982ef5a32436475397a2de457 Mon Sep 17 00:00:00 2001 From: James Krieger Date: Tue, 21 Jul 2026 13:52:12 +0100 Subject: [PATCH 4/7] prody clustenm app: accept comma-separated structures for multi-start The positional pdb argument now takes either one identifier/filename or several comma-separated ones. Each is parsed and selected, and the list is passed to ClustENM.setAtoms (multi-start): the extra structures seed the initial population as gen-1-style conformers. Single-structure use is unchanged. Co-Authored-By: Claude Opus 4.8 --- prody/apps/prody_apps/prody_clustenm.py | 41 ++++++++++++++++++------- 1 file changed, 30 insertions(+), 11 deletions(-) diff --git a/prody/apps/prody_apps/prody_clustenm.py b/prody/apps/prody_apps/prody_clustenm.py index 8521c8759..d64c37ae6 100644 --- a/prody/apps/prody_apps/prody_clustenm.py +++ b/prody/apps/prody_apps/prody_clustenm.py @@ -110,17 +110,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 +160,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, @@ -207,7 +218,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) @@ -350,7 +367,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__)) From 5223f34d83af9b47374577c85e44ca26a04b0008 Mon Sep 17 00:00:00 2001 From: James Krieger Date: Tue, 21 Jul 2026 14:41:27 +0100 Subject: [PATCH 5/7] ClustENM: use 'spawn' for the parallel minimise/MD pool (OpenMM+CUDA can't fork) OpenMM initialises the CUDA driver when its platform plugin is imported in the parent; a forked worker then fails Context creation with CUDA_ERROR_NOT_INITIALIZED. Switch _min_sim_batch's Pool to a 'spawn' context so each worker starts a fresh interpreter and initialises CUDA cleanly. self pickles fine (~6.6 MB). Because spawn re-imports the caller's module, the calling code must be guarded by `if __name__ == '__main__':` -- documented on run()'s parallel arg. Fixes the crash on the parallel+CUDA path. Note: for very short tasks (pure minimise, ~5 s) the per-worker spawn startup (OpenMM import + CUDA context) offsets the compute gain, so multi-GPU speedup is marginal there and expected to matter more for longer MD; single-GPU serial remains the fast path for small relax ensembles. Co-Authored-By: Claude Opus 4.8 --- prody/dynamics/clustenm.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/prody/dynamics/clustenm.py b/prody/dynamics/clustenm.py index bbbfd79ea..bee6ba6c1 100644 --- a/prody/dynamics/clustenm.py +++ b/prody/dynamics/clustenm.py @@ -25,7 +25,7 @@ __email__ = ['burak.kaynak@pitt.edu', 'doruker@pitt.edu', 'shz66@pitt.edu'] from itertools import product -from multiprocessing import cpu_count, Pool +from multiprocessing import cpu_count, get_context, Pool from collections import OrderedDict from os import chdir, mkdir from os.path import isdir @@ -490,7 +490,12 @@ def _min_sim_batch(self, coords_list): coords_list = list(coords_list) if self._parallel and len(coords_list) > 1: repeats = cpu_count() if self._parallel is True else int(self._parallel) - with Pool(repeats, initializer=self._worker_gpu_init) as p: + # 'spawn', not the default fork: OpenMM loads its CUDA platform plugin (initialising the CUDA + # driver) at import in the parent, and that driver state does not survive a fork -- a forked + # worker then fails Context creation with CUDA_ERROR_NOT_INITIALIZED. A spawned worker starts a + # fresh interpreter and initialises CUDA cleanly on its assigned GPU. self pickles fine. + ctx = get_context('spawn') + with ctx.Pool(repeats, initializer=self._worker_gpu_init) as p: return p.map(self._min_sim, coords_list) return [self._min_sim(c) for c in coords_list] @@ -1241,9 +1246,13 @@ def run(self, cutoff=15., n_modes=3, gamma=1., n_confs=50, rmsd=1.0, 'CPU' is needed for setting threads per simulation. :type platform: str - :arg parallel: If it is True (default is False), conformer generation will be parallelized. - This can also be set to a number for how many CPUs are used in parallel conformer generation. - Setting 0 or True means run as many as there are CPUs on the machine. + :arg parallel: If it is True (default is False), conformer generation AND the energy + minimisation / MD of each conformer are parallelized across worker processes. + This can also be set to a number for how many workers are used. + Setting 0 or True means run as many as there are CPUs on the machine. With a GPU + platform, workers are round-robined across the visible GPUs. The worker pool uses the + 'spawn' start method (OpenMM+CUDA cannot be forked), so the calling code MUST be guarded + by ``if __name__ == '__main__':``. :type parallel: bool :arg threads: Number of threads to use for an individual simulation using the thread setting from OpenMM From ec932977dba0bc5b05239200aa720c6dbac0ce3c Mon Sep 17 00:00:00 2001 From: James Krieger Date: Tue, 21 Jul 2026 16:27:15 +0100 Subject: [PATCH 6/7] ClustENM: separate parallel_sim + sim_devices (multi-GPU DeviceIndex) from parallel Decouple parallelisation of the per-conformer minimisation/MD from conformer generation: - `parallel` keeps its original meaning (parallel conformer generation only). - `parallel_sim` (bool/int) parallelises the minimise/MD across worker processes via _min_sim_batch's spawn Pool, independent of `parallel`. - `sim_devices` (list, or int N -> [0..N-1]) round-robins the parallel-sim workers across those GPUs. Applied through the OpenMM DeviceIndex platform property at Context creation (_prep_sim), which is read when the context is built -- so it is immune to the CUDA_VISIBLE_DEVICES-vs-import timing race that made the earlier env-var pinning ineffective. _worker_sim_init sets each worker's device from its rank; None = no pinning. Both run() and the `prody clustenm` app expose the two new options (--parallel_sim, --sim_devices). Defaults are off, so existing behaviour is unchanged. Correctness verified (parallel run yields the right conformers, app CLI wires the args and single/multi-start runs complete); multi-GPU speedup is marginal for short minimise (spawn start-up dominates) and expected to matter for longer MD / large ensembles. Co-Authored-By: Claude Opus 4.8 --- prody/apps/prody_apps/prody_clustenm.py | 21 ++++- prody/dynamics/clustenm.py | 107 +++++++++++++++++------- 2 files changed, 96 insertions(+), 32 deletions(-) diff --git a/prody/apps/prody_apps/prody_clustenm.py b/prody/apps/prody_apps/prody_clustenm.py index d64c37ae6..e9f528b80 100644 --- a/prody/apps/prody_apps/prody_clustenm.py +++ b/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') @@ -170,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) @@ -282,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']) diff --git a/prody/dynamics/clustenm.py b/prody/dynamics/clustenm.py index bee6ba6c1..a2261ef1c 100644 --- a/prody/dynamics/clustenm.py +++ b/prody/dynamics/clustenm.py @@ -54,6 +54,27 @@ __all__ = ['ClustENM', 'ClustRTB', 'ClustImANM', 'ClustExANM'] +# Physical device index this parallel-simulation worker is pinned to (None = no pinning). Set in the +# spawned worker by _worker_sim_init and read by _prep_sim when it builds the OpenMM Context, so each +# worker's Context lands on a distinct GPU. A module global (not an attribute) because it must persist +# in the worker process across tasks -- self is re-pickled per task, this is set once per worker. +_WORKER_DEVICE_INDEX = None + + +def _worker_sim_init(devices): + '''Pool initializer for parallel simulation: pin this worker to one device from *devices* + (round-robin by worker rank). *devices* is a list of physical device indices, or None/empty for + no pinning.''' + global _WORKER_DEVICE_INDEX + if devices: + from multiprocessing import current_process + try: + rank = int(current_process().name.rsplit('-', 1)[-1]) - 1 + except Exception: + rank = 0 + _WORKER_DEVICE_INDEX = devices[rank % len(devices)] + + class ClustENM(Ensemble): ''' @@ -98,8 +119,10 @@ def __init__(self, title=None): self._outlier = True self._mzscore = 3.5 self._v1 = False - self._platform = None + self._platform = None self._parallel = False + self._parallel_sim = False + self._sim_devices = None self._topology = None self._positions = None @@ -392,6 +415,11 @@ def _prep_sim(self, coords, external_forces=[]): properties = {'Precision': 'single'} elif self._platform in ['CUDA', 'OpenCL']: properties = {'Precision': 'single'} + # pin this Context to the device assigned to this parallel-sim worker (read at build time, + # so it is immune to the CUDA_VISIBLE_DEVICES-vs-import timing race) + if _WORKER_DEVICE_INDEX is not None: + key = 'DeviceIndex' if self._platform == 'CUDA' else 'OpenCLDeviceIndex' + properties[key] = str(_WORKER_DEVICE_INDEX) elif self._platform == 'CPU': if self._threads == 0: cpus = cpu_count() @@ -463,39 +491,25 @@ def _min_sim(self, coords): return np.nan, np.full_like(coords, np.nan) - @staticmethod - def _worker_gpu_init(): - - # Pin each Pool worker to ONE visible GPU (round-robin over CUDA_VISIBLE_DEVICES) so a multi-GPU - # parallel run spreads across the GPUs instead of all contending on GPU 0. No-op when fewer than - # two GPUs are visible (e.g. the CPU platform), so it is harmless for CPU-parallel runs. - - import os - from multiprocessing import current_process - gpus = [g for g in os.environ.get('CUDA_VISIBLE_DEVICES', '').split(',') if g != ''] - if len(gpus) > 1: - try: - rank = int(current_process().name.rsplit('-', 1)[-1]) - 1 - except Exception: - rank = 0 - os.environ['CUDA_VISIBLE_DEVICES'] = gpus[rank % len(gpus)] - def _min_sim_batch(self, coords_list): # Minimise (+ optional heat/sim) a batch of coordsets. Parallelised across conformers with a - # multiprocessing Pool when self._parallel is set (mirrors the _generate/_sample parallelisation), - # otherwise serial. Workers round-robin across visible GPUs (multi-GPU) via _worker_gpu_init. + # multiprocessing Pool when self._parallel_sim is set (independent of self._parallel, which + # parallelises conformer generation), otherwise serial. When self._sim_devices is set, workers + # round-robin across those GPUs via _worker_sim_init + the DeviceIndex property in _prep_sim. # Returns a list of (potential, coords) aligned with coords_list. coords_list = list(coords_list) - if self._parallel and len(coords_list) > 1: - repeats = cpu_count() if self._parallel is True else int(self._parallel) + if self._parallel_sim and len(coords_list) > 1: + repeats = cpu_count() if self._parallel_sim is True else int(self._parallel_sim) # 'spawn', not the default fork: OpenMM loads its CUDA platform plugin (initialising the CUDA # driver) at import in the parent, and that driver state does not survive a fork -- a forked # worker then fails Context creation with CUDA_ERROR_NOT_INITIALIZED. A spawned worker starts a # fresh interpreter and initialises CUDA cleanly on its assigned GPU. self pickles fine. + # (spawn re-imports the caller module, so the calling code must guard `if __name__=='__main__'`.) ctx = get_context('spawn') - with ctx.Pool(repeats, initializer=self._worker_gpu_init) as p: + with ctx.Pool(repeats, initializer=_worker_sim_init, + initargs=(self._sim_devices,)) as p: return p.map(self._min_sim, coords_list) return [self._min_sim(c) for c in coords_list] @@ -1246,15 +1260,27 @@ def run(self, cutoff=15., n_modes=3, gamma=1., n_confs=50, rmsd=1.0, 'CPU' is needed for setting threads per simulation. :type platform: str - :arg parallel: If it is True (default is False), conformer generation AND the energy - minimisation / MD of each conformer are parallelized across worker processes. - This can also be set to a number for how many workers are used. - Setting 0 or True means run as many as there are CPUs on the machine. With a GPU - platform, workers are round-robined across the visible GPUs. The worker pool uses the - 'spawn' start method (OpenMM+CUDA cannot be forked), so the calling code MUST be guarded - by ``if __name__ == '__main__':``. + :arg parallel: If it is True (default is False), conformer generation will be parallelized. + This can also be set to a number for how many CPUs are used in parallel conformer generation. + Setting 0 or True means run as many as there are CPUs on the machine. :type parallel: bool + :arg parallel_sim: Parallelize the per-conformer energy minimisation / MD across worker + processes (default False). Independent of *parallel* (which parallelizes conformer + generation): use this to speed up the simulation step, most useful for longer MD or large + ensembles. True or 0 uses as many workers as CPUs; an int sets the worker count. + The worker pool uses the 'spawn' start method (OpenMM+CUDA cannot be forked), so when + *parallel_sim* is used the calling code MUST be guarded by ``if __name__ == '__main__':``. + :type parallel_sim: bool, int + + :arg sim_devices: How the parallel simulation is distributed over GPUs (default None = no + pinning; every worker uses the platform's default device). A list of physical device + indices, or an int N meaning ``[0..N-1]``, over which the parallel-sim workers are + round-robined (applied via the OpenMM DeviceIndex property at Context creation, so it is + not affected by CUDA_VISIBLE_DEVICES import-time ordering). Only meaningful with + *parallel_sim* and a CUDA/OpenCL platform. + :type sim_devices: list, int, None + :arg threads: Number of threads to use for an individual simulation using the thread setting from OpenMM Default of 0 uses all CPUs on the machine. :type threads: int @@ -1302,6 +1328,27 @@ def run(self, cutoff=15., n_modes=3, gamma=1., n_confs=50, rmsd=1.0, # this is a deliberate choice and is documented self._parallel = True + # parallel_sim: parallelise the per-conformer energy minimisation / MD (independent of `parallel`, + # which parallelises conformer generation). bool or int worker count; True/0 -> cpu_count. + self._parallel_sim = kwargs.pop('parallel_sim', False) + if not isinstance(self._parallel_sim, (int, bool)): + raise TypeError('parallel_sim should be an int or bool') + if not isinstance(self._parallel_sim, bool) and self._parallel_sim == 1: + self._parallel_sim = False + elif not isinstance(self._parallel_sim, bool) and self._parallel_sim == 0: + self._parallel_sim = True + + # sim_devices: how the parallel simulation spreads over GPUs. A list of physical device indices, + # or an int N meaning devices [0..N-1], to round-robin the parallel-sim workers across (applied via + # the OpenMM DeviceIndex property at Context creation). None = no pinning (all workers default to + # the platform's default device). Only meaningful with parallel_sim and a CUDA/OpenCL platform. + sim_devices = kwargs.pop('sim_devices', None) + if isinstance(sim_devices, int) and not isinstance(sim_devices, bool): + sim_devices = list(range(sim_devices)) + elif sim_devices is not None: + sim_devices = [int(d) for d in sim_devices] + self._sim_devices = sim_devices + self._threads = kwargs.pop('threads', 0) self._targeted = kwargs.pop('targeted', False) self._tmdk = kwargs.pop('tmdk', 15.) From b3fb98ab370df03b31def3d495c3eba303002932 Mon Sep 17 00:00:00 2001 From: James Krieger Date: Tue, 21 Jul 2026 17:30:46 +0100 Subject: [PATCH 7/7] ClustENM: record a starting CC per seed for multi-start fitting With a fitmap, the starting cross-correlation was appended once (for the single reference), but multi-start seeds gen-0 with one conformer per input structure. That left self._cc short by (n_seeds - 1), so downstream consumers that index CC by conformer (e.g. the Scipion protocol) ran off the end. Append a starting CC for each extra seed too, keeping self._cc aligned with the gen-0 conformers. Co-Authored-By: Claude Opus 4.8 --- prody/dynamics/clustenm.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/prody/dynamics/clustenm.py b/prody/dynamics/clustenm.py index a2261ef1c..b84e76563 100644 --- a/prody/dynamics/clustenm.py +++ b/prody/dynamics/clustenm.py @@ -1475,6 +1475,19 @@ def run(self, cutoff=15., n_modes=3, gamma=1., n_confs=50, rmsd=1.0, conformers = start_confs = confs0 keys = [(0, j) for j in range(len(pots0))] + if self._fitmap is not None and len(starts) > 1: + # multi-start with fitting: the starting CC above was recorded for the first seed only, but + # gen-0 now holds one conformer per seed. Record a starting CC for each extra seed too, so + # self._cc stays aligned with the gen-0 conformers (labels (0, j)). + saved = self._atoms.getCoords() + for j in range(1, len(starts)): + self._atoms.setCoords(starts[j]) + sim_map = self._blurrer(self._atoms.toTEMPyStructure(), + self._fit_resolution, self._fitmap) + self._cc.append(float(self._scorer.CCC(self._fitmap, sim_map))) + self._atoms.setCoords(saved) + self._cc_prev = max(self._cc) + for i in range(1, self._n_gens+1): self._cycle += 1 LOGGER.info('Generation %d ...' % i)