Skip to content

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

Open
jamesmkrieger wants to merge 7 commits into
prody:mainfrom
jamesmkrieger:clustenm-multistart
Open

ClustENM: multi-start (setAtoms accepts a list of AtomGroups) + optional fixer#2258
jamesmkrieger wants to merge 7 commits into
prody:mainfrom
jamesmkrieger:clustenm-multistart

Conversation

@jamesmkrieger

Copy link
Copy Markdown
Contributor

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.

…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>
jamesmkrieger and others added 6 commits July 21, 2026 13:03
…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>
@jamesmkrieger

Copy link
Copy Markdown
Contributor Author

This now also includes efficiency improvements and fixes for volume fitting

@karolamik13 karolamik13 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants