IMP logo
IMP Reference Guide  develop.031dafb4d2,2024/05/21
The Integrative Modeling Platform
samplers.py
1 """@namespace IMP.pmi.samplers
2  Sampling of the system.
3 """
4 
5 from __future__ import print_function
6 import IMP
7 import IMP.core
8 from IMP.pmi.tools import get_restraint_set
9 
10 
11 class _SerialReplicaExchange(object):
12  """Dummy replica exchange class used in non-MPI builds.
13  It should act similarly to IMP.mpi.ReplicaExchange
14  on a single processor.
15  """
16  def __init__(self):
17  self.__params = {}
18 
19  def get_number_of_replicas(self):
20  return 1
21 
22  def create_temperatures(self, tmin, tmax, nrep):
23  return [tmin]
24 
25  def get_my_index(self):
26  return 0
27 
28  def set_my_parameter(self, key, val):
29  self.__params[key] = val
30 
31  def get_my_parameter(self, key):
32  return self.__params[key]
33 
34  def get_friend_index(self, step):
35  return 0
36 
37  def get_friend_parameter(self, key, findex):
38  return self.get_my_parameter(key)
39 
40  def do_exchange(self, myscore, fscore, findex):
41  return False
42 
43  def set_was_used(self, was_used):
44  self.was_used = was_used
45 
46 
47 class MonteCarlo(object):
48  """Sample using Monte Carlo"""
49 
50  # check that isd is installed
51  try:
52  import IMP.isd
53  isd_available = True
54  except ImportError:
55  isd_available = False
56 
57  def __init__(self, model, objects=None, temp=1.0, filterbyname=None,
58  score_moved=False):
59  """Setup Monte Carlo sampling
60  @param model The IMP Model
61  @param objects What to sample (a list of Movers)
62  @param temp The MC temperature
63  @param filterbyname Not used
64  @param score_moved If True, attempt to speed up sampling by
65  caching scoring function terms on particles that didn't move
66  """
67  self.losp = [
68  "Rigid_Bodies",
69  "Floppy_Bodies",
70  "Nuisances",
71  "X_coord",
72  "Weights"
73  "Surfaces"]
74  self.simulated_annealing = False
75  self.selfadaptive = False
76  # that is -1 because mc has not yet run
77  self.nframe = -1
78  self.temp = temp
79  self.mvs = []
80  self.mvslabels = []
81  self.label = "None"
82  self.model = model
83  self.movers_data = {}
84 
85  self.mvs = objects
86 
87  # SerialMover
88  self.smv = IMP.core.SerialMover(self.mvs)
89 
90  self.mc = IMP.core.MonteCarlo(self.model)
91  self.mc.set_scoring_function(get_restraint_set(self.model))
92  self.mc.set_return_best(False)
93  self.mc.set_score_moved(score_moved)
94  self.mc.set_kt(self.temp)
95  self.mc.add_mover(self.smv)
96 
97  def set_kt(self, temp):
98  self.temp = temp
99  self.mc.set_kt(temp)
100 
101  def get_mc(self):
102  return self.mc
103 
104  def set_scoring_function(self, objectlist):
105  rs = IMP.RestraintSet(self.model, 1.0, 'sfo')
106  for ob in objectlist:
107  rs.add_restraint(ob.get_restraint())
109  self.mc.set_scoring_function(sf)
110 
111  def set_simulated_annealing(
112  self,
113  min_temp,
114  max_temp,
115  min_temp_time,
116  max_temp_time):
117  self.simulated_annealing = True
118  self.tempmin = min_temp
119  self.tempmax = max_temp
120  self.timemin = min_temp_time
121  self.timemax = max_temp_time
122 
123  def set_self_adaptive(self, isselfadaptive=True):
124  self.selfadaptive = isselfadaptive
125 
127  '''
128  Return a dictionary with the mover parameters for nuisance parameters
129  '''
130  output = {}
131  for i in range(self.get_number_of_movers()):
132  mv = self.smv.get_mover(i)
133  name = mv.get_name()
134  if "Nuisances" in name:
135  stepsize = IMP.core.NormalMover.get_from(mv).get_sigma()
136  output[name] = stepsize
137  return output
138 
139  def get_number_of_movers(self):
140  return len(self.smv.get_movers())
141 
142  def get_particle_types(self):
143  return self.losp
144 
145  def optimize(self, nstep):
146  self.nframe += 1
147  self.mc.optimize(nstep * self.get_number_of_movers())
148 
149  # apply simulated annealing protocol
150  if self.simulated_annealing:
151  self.temp = self.temp_simulated_annealing()
152  self.mc.set_kt(self.temp)
153 
154  # apply self adaptive protocol
155  if self.selfadaptive:
156  for i, mv in enumerate(self.mvs):
157 
158  mvacc = mv.get_number_of_accepted()
159  mvprp = mv.get_number_of_proposed()
160  if mv not in self.movers_data:
161  accept = float(mvacc) / float(mvprp)
162  self.movers_data[mv] = (mvacc, mvprp)
163  else:
164  oldmvacc, oldmvprp = self.movers_data[mv]
165  accept = float(mvacc-oldmvacc) / float(mvprp-oldmvprp)
166  self.movers_data[mv] = (mvacc, mvprp)
167  if accept < 0.05:
168  accept = 0.05
169  if accept > 1.0:
170  accept = 1.0
171 
172  if isinstance(mv, IMP.core.NormalMover):
173  stepsize = mv.get_sigma()
174  if 0.4 > accept or accept > 0.6:
175  mv.set_sigma(stepsize * 2 * accept)
176 
177  if isinstance(mv, IMP.isd.WeightMover):
178  stepsize = mv.get_radius()
179  if 0.4 > accept or accept > 0.6:
180  mv.set_radius(stepsize * 2 * accept)
181 
182  if isinstance(mv, IMP.core.RigidBodyMover):
183  mr = mv.get_maximum_rotation()
184  mt = mv.get_maximum_translation()
185  if 0.4 > accept or accept > 0.6:
186  mv.set_maximum_rotation(mr * 2 * accept)
187  mv.set_maximum_translation(mt * 2 * accept)
188 
189  if isinstance(mv, IMP.pmi.TransformMover):
190  mr = mv.get_maximum_rotation()
191  mt = mv.get_maximum_translation()
192  if 0.4 > accept or accept > 0.6:
193  mv.set_maximum_rotation(mr * 2 * accept)
194  mv.set_maximum_translation(mt * 2 * accept)
195 
196  if isinstance(mv, IMP.core.BallMover):
197  mr = mv.get_radius()
198  if 0.4 > accept or accept > 0.6:
199  mv.set_radius(mr * 2 * accept)
200 
201  def get_nuisance_movers(self, nuisances, maxstep):
202  mvs = []
203  for nuisance in nuisances:
204  print(nuisance, maxstep)
205  mvs.append(
206  IMP.core.NormalMover([nuisance],
207  IMP.FloatKeys([IMP.FloatKey("nuisance")]),
208  maxstep))
209  return mvs
210 
211  def get_rigid_body_movers(self, rbs, maxtrans, maxrot):
212  mvs = []
213  for rb in rbs:
214  mvs.append(IMP.core.RigidBodyMover(rb.get_model(), rb,
215  maxtrans, maxrot))
216  return mvs
217 
218  def get_super_rigid_body_movers(self, rbs, maxtrans, maxrot):
219  mvs = []
220  for rb in rbs:
221  if len(rb) == 2:
222  # normal Super Rigid Body
223  srbm = IMP.pmi.TransformMover(self.model, maxtrans, maxrot)
224  elif len(rb) == 3:
225  if isinstance(rb[2], tuple) and len(rb[2]) == 3 \
226  and isinstance(rb[2][0], float) \
227  and isinstance(rb[2][1], float) \
228  and isinstance(rb[2][2], float):
229  # super rigid body with 2D rotation, rb[2] is the axis
230  srbm = IMP.pmi.TransformMover(
231  self.model, IMP.algebra.Vector3D(rb[2]), maxtrans,
232  maxrot)
233  else:
234  print(
235  "Setting up a super rigid body with wrong parameters")
236  raise
237 
238  for xyz in rb[0]:
239  srbm.add_xyz_particle(xyz)
240  for rb in rb[1]:
241  srbm.add_rigid_body_particle(rb)
242  mvs.append(srbm)
243  return mvs
244 
245  def get_floppy_body_movers(self, fbs, maxtrans):
246  mvs = []
247  for fb in fbs:
248  # check is that is a rigid body member:
250  # if so force the particles to move anyway
251  floatkeys = \
253  for fk in floatkeys:
254  fb.set_is_optimized(fk, True)
255  mvs.append(
256  IMP.core.BallMover(fb.get_model(), fb,
257  IMP.FloatKeys(floatkeys),
258  maxtrans))
259  else:
260  # otherwise use the normal ball mover
261  mvs.append(IMP.core.BallMover(fb.get_model(), fb, maxtrans))
262  return mvs
263 
264  def get_X_movers(self, fbs, maxtrans):
265  mvs = []
266  Xfloatkey = IMP.core.XYZ.get_xyz_keys()[0]
267  for fb in fbs:
268  # check is that is a rigid body member:
270  raise ValueError("particle is part of a rigid body")
271  else:
272  # otherwise use the normal ball mover
273  mvs.append(IMP.core.NormalMover([fb], [Xfloatkey], maxtrans))
274  return mvs
275 
276  def get_weight_movers(self, weights, maxstep):
277  mvs = []
278  for weight in weights:
279  if weight.get_number_of_weights() > 1:
280  mvs.append(IMP.isd.WeightMover(weight, maxstep))
281  return mvs
282 
283  def get_surface_movers(self, surfaces, maxtrans, maxrot, refprob):
284  mvs = []
285  for surface in surfaces:
286  mvs.append(IMP.core.SurfaceMover(surface, maxtrans, maxrot,
287  refprob))
288  return mvs
289 
290  def temp_simulated_annealing(self):
291  if self.nframe % (self.timemin + self.timemax) < self.timemin:
292  value = 0.0
293  else:
294  value = 1.0
295  temp = self.tempmin + (self.tempmax - self.tempmin) * value
296  return temp
297 
298  def set_label(self, label):
299  self.label = label
300 
301  def get_frame_number(self):
302  return self.nframe
303 
304  def get_output(self):
305  output = {}
306  for i, mv in enumerate(self.smv.get_movers()):
307  mvname = mv.get_name()
308  mvacc = mv.get_number_of_accepted()
309  mvprp = mv.get_number_of_proposed()
310  try:
311  mvacr = float(mvacc) / float(mvprp)
312  except: # noqa: E722
313  mvacr = 0.0
314  output["MonteCarlo_Acceptance_" +
315  mvname + "_" + str(i)] = str(mvacr)
316  if "Nuisances" in mvname:
317  output["MonteCarlo_StepSize_" + mvname + "_" + str(i)] = \
318  str(IMP.core.NormalMover.get_from(mv).get_sigma())
319  if "Weights" in mvname:
320  output["MonteCarlo_StepSize_" + mvname + "_" + str(i)] = \
321  str(IMP.isd.WeightMover.get_from(mv).get_radius())
322  output["MonteCarlo_Temperature"] = str(self.mc.get_kt())
323  output["MonteCarlo_Nframe"] = str(self.nframe)
324  return output
325 
326 
327 class MolecularDynamics(object):
328  """Sample using molecular dynamics"""
329 
330  def __init__(self, model, objects, kt, gamma=0.01, maximum_time_step=1.0,
331  sf=None):
332  """Setup MD
333  @param model The IMP Model
334  @param objects What to sample. Use flat list of particles
335  @param kt Temperature
336  @param gamma Viscosity parameter
337  @param maximum_time_step MD max time step
338  """
339  self.model = model
340 
341  # check if using PMI1 objects dictionary, or just list of particles
342  try:
343  for obj in objects:
344  psamp = obj.get_particles_to_sample()
345  to_sample = psamp['Floppy_Bodies_SimplifiedModel'][0]
346  except: # noqa: E722
347  to_sample = objects
348 
350  self.model, to_sample, kt/0.0019872041, gamma)
351  self.md = IMP.atom.MolecularDynamics(self.model)
352  self.md.set_maximum_time_step(maximum_time_step)
353  if sf:
354  self.md.set_scoring_function(sf)
355  else:
356  self.md.set_scoring_function(get_restraint_set(self.model))
357  self.md.add_optimizer_state(self.ltstate)
358  self.simulated_annealing = False
359  self.nframe = -1
360 
361  def set_kt(self, kt):
362  temp = kt/0.0019872041
363  self.ltstate.set_temperature(temp)
364  self.md.assign_velocities(temp)
365 
366  def set_simulated_annealing(self, min_temp, max_temp, min_temp_time,
367  max_temp_time):
368  self.simulated_annealing = True
369  self.tempmin = min_temp
370  self.tempmax = max_temp
371  self.timemin = min_temp_time
372  self.timemax = max_temp_time
373 
374  def temp_simulated_annealing(self):
375  if self.nframe % (self.timemin + self.timemax) < self.timemin:
376  value = 0.0
377  else:
378  value = 1.0
379  temp = self.tempmin + (self.tempmax - self.tempmin) * value
380  return temp
381 
382  def set_gamma(self, gamma):
383  self.ltstate.set_gamma(gamma)
384 
385  def optimize(self, nsteps):
386  # apply simulated annealing protocol
387  self.nframe += 1
388  if self.simulated_annealing:
389  self.temp = self.temp_simulated_annealing()
390  self.set_kt(self.temp)
391  self.md.optimize(nsteps)
392 
393  def get_output(self):
394  output = {}
395  output["MolecularDynamics_KineticEnergy"] = \
396  str(self.md.get_kinetic_energy())
397  return output
398 
399 
400 class ConjugateGradients(object):
401  """Sample using conjugate gradients"""
402 
403  def __init__(self, model, objects):
404  self.model = model
405  self.nframe = -1
406  self.cg = IMP.core.ConjugateGradients(self.model)
407  self.cg.set_scoring_function(get_restraint_set(self.model))
408 
409  def set_label(self, label):
410  self.label = label
411 
412  def get_frame_number(self):
413  return self.nframe
414 
415  def optimize(self, nstep):
416  self.nframe += 1
417  self.cg.optimize(nstep)
418 
419  def set_scoring_function(self, objectlist):
420  rs = IMP.RestraintSet(self.model, 1.0, 'sfo')
421  for ob in objectlist:
422  rs.add_restraint(ob.get_restraint())
424  self.cg.set_scoring_function(sf)
425 
426  def get_output(self):
427  output = {}
428  output["ConjugatedGradients_Nframe"] = str(self.nframe)
429  return output
430 
431 
432 class ReplicaExchange(object):
433  """Sample using replica exchange"""
434 
435  def __init__(self, model, tempmin, tempmax, samplerobjects, test=True,
436  replica_exchange_object=None):
437  '''
438  samplerobjects can be a list of MonteCarlo or MolecularDynamics
439  '''
440 
441  self.model = model
442  self.samplerobjects = samplerobjects
443  # min and max temperature
444  self.TEMPMIN_ = tempmin
445  self.TEMPMAX_ = tempmax
446 
447  if replica_exchange_object is None:
448  # initialize Replica Exchange class
449  try:
450  import IMP.mpi
451  print('ReplicaExchange: MPI was found. '
452  'Using Parallel Replica Exchange')
453  self.rem = IMP.mpi.ReplicaExchange()
454  except ImportError:
455  print('ReplicaExchange: Could not find MPI. '
456  'Using Serial Replica Exchange')
457  self.rem = _SerialReplicaExchange()
458 
459  else:
460  # get the replica exchange class instance from elsewhere
461  print('got existing rex object')
462  self.rem = replica_exchange_object
463 
464  # get number of replicas
465  nproc = self.rem.get_number_of_replicas()
466 
467  if nproc % 2 != 0 and not test:
468  raise Exception(
469  "number of replicas has to be even. "
470  "set test=True to run with odd number of replicas.")
471  # create array of temperatures, in geometric progression
472  temp = self.rem.create_temperatures(
473  self.TEMPMIN_,
474  self.TEMPMAX_,
475  nproc)
476  # get replica index
477  self.temperatures = temp
478 
479  myindex = self.rem.get_my_index()
480  # set initial value of the parameter (temperature) to exchange
481  self.rem.set_my_parameter("temp", [self.temperatures[myindex]])
482  for so in self.samplerobjects:
483  so.set_kt(self.temperatures[myindex])
484  self.nattempts = 0
485  self.nmintemp = 0
486  self.nmaxtemp = 0
487  self.nsuccess = 0
488 
489  def get_temperatures(self):
490  return self.temperatures
491 
492  def get_my_temp(self):
493  return self.rem.get_my_parameter("temp")[0]
494 
495  def get_my_index(self):
496  return self.rem.get_my_index()
497 
498  def swap_temp(self, nframe, score=None):
499  if score is None:
500  score = self.model.evaluate(False)
501  # get my replica index and temperature
502  _ = self.rem.get_my_index()
503  mytemp = self.rem.get_my_parameter("temp")[0]
504 
505  if mytemp == self.TEMPMIN_:
506  self.nmintemp += 1
507 
508  if mytemp == self.TEMPMAX_:
509  self.nmaxtemp += 1
510 
511  # score divided by kbt
512  myscore = score / mytemp
513 
514  # get my friend index and temperature
515  findex = self.rem.get_friend_index(nframe)
516  ftemp = self.rem.get_friend_parameter("temp", findex)[0]
517  # score divided by kbt
518  fscore = score / ftemp
519 
520  # try exchange
521  flag = self.rem.do_exchange(myscore, fscore, findex)
522 
523  self.nattempts += 1
524  # if accepted, change temperature
525  if (flag):
526  for so in self.samplerobjects:
527  so.set_kt(ftemp)
528  self.nsuccess += 1
529 
530  def get_output(self):
531  output = {}
532  if self.nattempts != 0:
533  output["ReplicaExchange_SwapSuccessRatio"] = str(
534  float(self.nsuccess) / self.nattempts)
535  output["ReplicaExchange_MinTempFrequency"] = str(
536  float(self.nmintemp) / self.nattempts)
537  output["ReplicaExchange_MaxTempFrequency"] = str(
538  float(self.nmaxtemp) / self.nattempts)
539  else:
540  output["ReplicaExchange_SwapSuccessRatio"] = str(0)
541  output["ReplicaExchange_MinTempFrequency"] = str(0)
542  output["ReplicaExchange_MaxTempFrequency"] = str(0)
543  output["ReplicaExchange_CurrentTemp"] = str(self.get_my_temp())
544  return output
545 
546 
547 class MPI_values(object):
548  def __init__(self, replica_exchange_object=None):
549  """Query values (ie score, and others)
550  from a set of parallel jobs"""
551 
552  if replica_exchange_object is None:
553  # initialize Replica Exchange class
554  try:
555  import IMP.mpi
556  print('MPI_values: MPI was found. '
557  'Using Parallel Replica Exchange')
558  self.rem = IMP.mpi.ReplicaExchange()
559  except ImportError:
560  print('MPI_values: Could not find MPI. '
561  'Using Serial Replica Exchange')
562  self.rem = _SerialReplicaExchange()
563 
564  else:
565  # get the replica exchange class instance from elsewhere
566  print('got existing rex object')
567  self.rem = replica_exchange_object
568 
569  def set_value(self, name, value):
570  self.rem.set_my_parameter(name, [value])
571 
572  def get_values(self, name):
573  values = []
574  for i in range(self.rem.get_number_of_replicas()):
575  v = self.rem.get_friend_parameter(name, i)[0]
576  values.append(v)
577  return values
578 
579  def get_percentile(self, name):
580  value = self.rem.get_my_parameter(name)[0]
581  values = sorted(self.get_values(name))
582  ind = values.index(value)
583  percentile = float(ind)/len(values)
584  return percentile
def __init__
samplerobjects can be a list of MonteCarlo or MolecularDynamics
Definition: samplers.py:435
A Monte Carlo optimizer.
Definition: MonteCarlo.h:44
A class to implement Hamiltonian Replica Exchange.
Maintains temperature during molecular dynamics.
Sample using molecular dynamics.
Definition: samplers.py:327
def get_nuisance_movers_parameters
Return a dictionary with the mover parameters for nuisance parameters.
Definition: samplers.py:126
Miscellaneous utilities.
Definition: tools.py:1
Modify the transformation of a rigid body.
Simple conjugate gradients optimizer.
Sample using conjugate gradients.
Definition: samplers.py:400
Create a scoring function on a list of restraints.
Move continuous particle variables by perturbing them within a ball.
Modify a surface orientation.
Definition: SurfaceMover.h:37
static FloatKeys get_internal_coordinate_keys()
Definition: rigid_bodies.h:727
Object used to hold a set of restraints.
Definition: RestraintSet.h:41
Simple molecular dynamics simulator.
Code that uses the MPI parallel library.
A mover that perturbs a Weight particle.
Definition: WeightMover.h:20
Modify the transformation of a rigid body.
def __init__
Setup Monte Carlo sampling.
Definition: samplers.py:57
Modify a set of continuous variables using a normal distribution.
Definition: NormalMover.h:23
Basic functionality that is expected to be used by a wide variety of IMP users.
Sample using Monte Carlo.
Definition: samplers.py:47
The general base class for IMP exceptions.
Definition: exception.h:48
VectorD< 3 > Vector3D
Definition: VectorD.h:408
static const FloatKeys & get_xyz_keys()
Get a vector containing the keys for x,y,z.
Applies a list of movers one at a time.
Definition: SerialMover.h:26
static bool get_is_setup(const IMP::ParticleAdaptor &p)
Definition: rigid_bodies.h:770
Sample using replica exchange.
Definition: samplers.py:432
Inferential scoring building on methods developed as part of the Inferential Structure Determination ...