ClustENM: multi-start (setAtoms accepts a list of AtomGroups) + optional fixer#2258
ClustENM: multi-start (setAtoms accepts a list of AtomGroups) + optional fixer#2258jamesmkrieger wants to merge 7 commits into
Conversation
…alse 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 <noreply@anthropic.com>
…icit-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 <noreply@anthropic.com>
…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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
…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 <noreply@anthropic.com>
… 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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
|
This now also includes efficiency improvements and fixes for volume fitting |
karolamik13
left a comment
There was a problem hiding this comment.
I tested NMR structure 2KID (2 models).
Looks good. The code too.
Something else should I check here?
Code:
from prody import *
PDB_ID = '2KID'
CHAIN = 'A'
MODELS = [1, 2]
models = [parsePDB(PDB_ID, model=i).select('protein and chain ' + CHAIN) for i in MODELS]
clustenm = ClustENM(PDB_ID + '_multistart')
clustenm.setAtoms(models)
clustenm.run(n_gens=1, n_modes=3, n_confs=5, sim=False, rmsd=1.0, maxclust=5, solvent='imp', platform='CPU')
clustenm.writePDB(PDB_ID + '_multistart_relaxed', single=True)
print(clustenm)
print('Input structures:', len(models))
print('Output conformations:', clustenm.numConfs())
Tests:
@> PDB file is found in working directory (2kid.pdb.gz).
@> 2437 atoms and 1 coordinate set(s) were parsed in 0.06s.
@> PDB file is found in working directory (2kid.pdb.gz).
@> 2437 atoms and 1 coordinate set(s) were parsed in 0.06s.
@> Fixing the structure ...
@> 2361 atoms and 1 coordinate set(s) were parsed in 0.03s.
@> The structure was fixed in 0.83s.
@> 2361 atoms and 1 coordinate set(s) were parsed in 0.02s.
@> Multi-start: 2 initial structures set.
@> Kirchhoff was built in 0.01s.
@> Generation 0 ...
@> Minimization in generation 0 ...
@> Multi-start: minimising 2 initial structures (gen-0 single-conformer step skipped) ...
@> Completed in 10.26s.
@> #-------------------/-------------------#
@> Generation 1 ...
@> Sampling conformers in generation 1 ...
@> Hessian was built in 0.05s.
@> 3 modes were calculated in 0.02s.
@> Parameter: rmsd = 1.00 A
@> Parameter: n_confs = 5
@> Modes are scaled by 14.47714121815451.
@> Hessian was built in 0.08s.
@> 3 modes were calculated in 0.02s.
@> Parameter: rmsd = 1.00 A
@> Parameter: n_confs = 5
@> Modes are scaled by 20.484855031058046.
@> Clustering in generation 1 ...
@> Centroids were generated in 0.24s.
@> Minimization in generation 1 ...
@> Structures were sampled in 10.04s.
@> #-------------------/-------------------#
@> Creating an ensemble of conformers ...
@> Ensemble was created in 0.00s.
@> All completed in 20.55s.
@> PDB file saved as 2KID_multistart_relaxed.pdb
Ensemble 2KID_multistart
Input structures: 2
Output conformations: 5
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. Closes #2089
Validated with distorted structures from boltz: setAtoms([s0,s1]) + run(n_gens=0) minimises both (clash 318/340 -> 0, domain distance and %SS preserved). Single-structure path byte-unchanged.