IMP logo
IMP Reference Guide  2.17.0
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. Use flat list of particles
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  # check if using PMI1 or just passed a list of movers
86  gather_objects = False
87  try:
88  objects[0].get_particles_to_sample()
89  gather_objects = True
90  except: # noqa: E722
91  self.mvs = objects
92 
93  if gather_objects:
94  for ob in objects:
95  pts = ob.get_particles_to_sample()
96  for k in pts.keys():
97 
98  if "Rigid_Bodies" in k:
99  mvs = self.get_rigid_body_movers(
100  pts[k][0],
101  pts[k][1],
102  pts[k][2])
103  for mv in mvs:
104  mv.set_name(k)
105  self.mvs += mvs
106 
107  if "SR_Bodies" in k:
108  print(len(pts[k]))
109  mvs = self.get_super_rigid_body_movers(
110  pts[k][0],
111  pts[k][1],
112  pts[k][2])
113  for mv in mvs:
114  mv.set_name(k)
115  self.mvs += mvs
116 
117  if "Floppy_Bodies" in k:
118  mvs = self.get_floppy_body_movers(pts[k][0], pts[k][1])
119  for mv in mvs:
120  mv.set_name(k)
121  self.mvs += mvs
122 
123  if "X_coord" in k:
124  mvs = self.get_X_movers(pts[k][0], pts[k][1])
125  for mv in mvs:
126  mv.set_name(k)
127  self.mvs += mvs
128 
129  if "Nuisances" in k:
130  if not self.isd_available:
131  raise ValueError(
132  "isd module needed to use nuisances")
133  mvs = self.get_nuisance_movers(pts[k][0], pts[k][1])
134  for mv in mvs:
135  mv.set_name(k)
136  self.mvs += mvs
137 
138  if "Weights" in k:
139  if not self.isd_available:
140  raise ValueError(
141  "isd module needed to use weights")
142  mvs = self.get_weight_movers(pts[k][0], pts[k][1])
143  for mv in mvs:
144  mv.set_name(k)
145  self.mvs += mvs
146 
147  if "Surfaces" in k:
148  mvs = self.get_surface_movers(
149  pts[k][0],
150  pts[k][1],
151  pts[k][2],
152  pts[k][3])
153  for mv in mvs:
154  mv.set_name(k)
155  self.mvs += mvs
156 
157  # SerialMover
158  self.smv = IMP.core.SerialMover(self.mvs)
159 
160  self.mc = IMP.core.MonteCarlo(self.model)
161  self.mc.set_scoring_function(get_restraint_set(self.model))
162  self.mc.set_return_best(False)
163  self.mc.set_score_moved(score_moved)
164  self.mc.set_kt(self.temp)
165  self.mc.add_mover(self.smv)
166 
167  def set_kt(self, temp):
168  self.temp = temp
169  self.mc.set_kt(temp)
170 
171  def get_mc(self):
172  return self.mc
173 
174  def set_scoring_function(self, objectlist):
175  rs = IMP.RestraintSet(self.model, 1.0, 'sfo')
176  for ob in objectlist:
177  rs.add_restraint(ob.get_restraint())
179  self.mc.set_scoring_function(sf)
180 
181  def set_simulated_annealing(
182  self,
183  min_temp,
184  max_temp,
185  min_temp_time,
186  max_temp_time):
187  self.simulated_annealing = True
188  self.tempmin = min_temp
189  self.tempmax = max_temp
190  self.timemin = min_temp_time
191  self.timemax = max_temp_time
192 
193  def set_self_adaptive(self, isselfadaptive=True):
194  self.selfadaptive = isselfadaptive
195 
197  '''
198  Return a dictionary with the mover parameters for nuisance parameters
199  '''
200  output = {}
201  for i in range(self.get_number_of_movers()):
202  mv = self.smv.get_mover(i)
203  name = mv.get_name()
204  if "Nuisances" in name:
205  stepsize = IMP.core.NormalMover.get_from(mv).get_sigma()
206  output[name] = stepsize
207  return output
208 
209  def get_number_of_movers(self):
210  return len(self.smv.get_movers())
211 
212  def get_particle_types(self):
213  return self.losp
214 
215  def optimize(self, nstep):
216  self.nframe += 1
217  self.mc.optimize(nstep * self.get_number_of_movers())
218 
219  # apply simulated annealing protocol
220  if self.simulated_annealing:
221  self.temp = self.temp_simulated_annealing()
222  self.mc.set_kt(self.temp)
223 
224  # apply self adaptive protocol
225  if self.selfadaptive:
226  for i, mv in enumerate(self.mvs):
227 
228  mvacc = mv.get_number_of_accepted()
229  mvprp = mv.get_number_of_proposed()
230  if mv not in self.movers_data:
231  accept = float(mvacc) / float(mvprp)
232  self.movers_data[mv] = (mvacc, mvprp)
233  else:
234  oldmvacc, oldmvprp = self.movers_data[mv]
235  accept = float(mvacc-oldmvacc) / float(mvprp-oldmvprp)
236  self.movers_data[mv] = (mvacc, mvprp)
237  if accept < 0.05:
238  accept = 0.05
239  if accept > 1.0:
240  accept = 1.0
241 
242  if isinstance(mv, IMP.core.NormalMover):
243  stepsize = mv.get_sigma()
244  if 0.4 > accept or accept > 0.6:
245  mv.set_sigma(stepsize * 2 * accept)
246 
247  if isinstance(mv, IMP.isd.WeightMover):
248  stepsize = mv.get_radius()
249  if 0.4 > accept or accept > 0.6:
250  mv.set_radius(stepsize * 2 * accept)
251 
252  if isinstance(mv, IMP.core.RigidBodyMover):
253  mr = mv.get_maximum_rotation()
254  mt = mv.get_maximum_translation()
255  if 0.4 > accept or accept > 0.6:
256  mv.set_maximum_rotation(mr * 2 * accept)
257  mv.set_maximum_translation(mt * 2 * accept)
258 
259  if isinstance(mv, IMP.pmi.TransformMover):
260  mr = mv.get_maximum_rotation()
261  mt = mv.get_maximum_translation()
262  if 0.4 > accept or accept > 0.6:
263  mv.set_maximum_rotation(mr * 2 * accept)
264  mv.set_maximum_translation(mt * 2 * accept)
265 
266  if isinstance(mv, IMP.core.BallMover):
267  mr = mv.get_radius()
268  if 0.4 > accept or accept > 0.6:
269  mv.set_radius(mr * 2 * accept)
270 
271  def get_nuisance_movers(self, nuisances, maxstep):
272  mvs = []
273  for nuisance in nuisances:
274  print(nuisance, maxstep)
275  mvs.append(
276  IMP.core.NormalMover([nuisance],
277  IMP.FloatKeys([IMP.FloatKey("nuisance")]),
278  maxstep))
279  return mvs
280 
281  def get_rigid_body_movers(self, rbs, maxtrans, maxrot):
282  mvs = []
283  for rb in rbs:
284  mvs.append(IMP.core.RigidBodyMover(rb.get_model(), rb,
285  maxtrans, maxrot))
286  return mvs
287 
288  def get_super_rigid_body_movers(self, rbs, maxtrans, maxrot):
289  mvs = []
290  for rb in rbs:
291  if len(rb) == 2:
292  # normal Super Rigid Body
293  srbm = IMP.pmi.TransformMover(self.model, maxtrans, maxrot)
294  elif len(rb) == 3:
295  if isinstance(rb[2], tuple) and len(rb[2]) == 3 \
296  and isinstance(rb[2][0], float) \
297  and isinstance(rb[2][1], float) \
298  and isinstance(rb[2][2], float):
299  # super rigid body with 2D rotation, rb[2] is the axis
300  srbm = IMP.pmi.TransformMover(
301  self.model, IMP.algebra.Vector3D(rb[2]), maxtrans,
302  maxrot)
303  else:
304  print(
305  "Setting up a super rigid body with wrong parameters")
306  raise
307 
308  for xyz in rb[0]:
309  srbm.add_xyz_particle(xyz)
310  for rb in rb[1]:
311  srbm.add_rigid_body_particle(rb)
312  mvs.append(srbm)
313  return mvs
314 
315  def get_floppy_body_movers(self, fbs, maxtrans):
316  mvs = []
317  for fb in fbs:
318  # check is that is a rigid body member:
320  # if so force the particles to move anyway
321  floatkeys = \
323  for fk in floatkeys:
324  fb.set_is_optimized(fk, True)
325  mvs.append(
326  IMP.core.BallMover(fb.get_model(), fb,
327  IMP.FloatKeys(floatkeys),
328  maxtrans))
329  else:
330  # otherwise use the normal ball mover
331  mvs.append(IMP.core.BallMover(fb.get_model(), fb, maxtrans))
332  return mvs
333 
334  def get_X_movers(self, fbs, maxtrans):
335  mvs = []
336  Xfloatkey = IMP.core.XYZ.get_xyz_keys()[0]
337  for fb in fbs:
338  # check is that is a rigid body member:
340  raise ValueError("particle is part of a rigid body")
341  else:
342  # otherwise use the normal ball mover
343  mvs.append(IMP.core.NormalMover([fb], [Xfloatkey], maxtrans))
344  return mvs
345 
346  def get_weight_movers(self, weights, maxstep):
347  mvs = []
348  for weight in weights:
349  if weight.get_number_of_weights() > 1:
350  mvs.append(IMP.isd.WeightMover(weight, maxstep))
351  return mvs
352 
353  def get_surface_movers(self, surfaces, maxtrans, maxrot, refprob):
354  mvs = []
355  for surface in surfaces:
356  mvs.append(IMP.core.SurfaceMover(surface, maxtrans, maxrot,
357  refprob))
358  return mvs
359 
360  def temp_simulated_annealing(self):
361  if self.nframe % (self.timemin + self.timemax) < self.timemin:
362  value = 0.0
363  else:
364  value = 1.0
365  temp = self.tempmin + (self.tempmax - self.tempmin) * value
366  return temp
367 
368  def set_label(self, label):
369  self.label = label
370 
371  def get_frame_number(self):
372  return self.nframe
373 
374  def get_output(self):
375  output = {}
376  for i, mv in enumerate(self.smv.get_movers()):
377  mvname = mv.get_name()
378  mvacc = mv.get_number_of_accepted()
379  mvprp = mv.get_number_of_proposed()
380  try:
381  mvacr = float(mvacc) / float(mvprp)
382  except: # noqa: E722
383  mvacr = 0.0
384  output["MonteCarlo_Acceptance_" +
385  mvname + "_" + str(i)] = str(mvacr)
386  if "Nuisances" in mvname:
387  output["MonteCarlo_StepSize_" + mvname + "_" + str(i)] = \
388  str(IMP.core.NormalMover.get_from(mv).get_sigma())
389  if "Weights" in mvname:
390  output["MonteCarlo_StepSize_" + mvname + "_" + str(i)] = \
391  str(IMP.isd.WeightMover.get_from(mv).get_radius())
392  output["MonteCarlo_Temperature"] = str(self.mc.get_kt())
393  output["MonteCarlo_Nframe"] = str(self.nframe)
394  return output
395 
396 
397 class MolecularDynamics(object):
398  """Sample using molecular dynamics"""
399 
400  def __init__(self, model, objects, kt, gamma=0.01, maximum_time_step=1.0,
401  sf=None):
402  """Setup MD
403  @param model The IMP Model
404  @param objects What to sample. Use flat list of particles
405  @param kt Temperature
406  @param gamma Viscosity parameter
407  @param maximum_time_step MD max time step
408  """
409  self.model = model
410 
411  # check if using PMI1 objects dictionary, or just list of particles
412  try:
413  for obj in objects:
414  psamp = obj.get_particles_to_sample()
415  to_sample = psamp['Floppy_Bodies_SimplifiedModel'][0]
416  except: # noqa: E722
417  to_sample = objects
418 
420  self.model, to_sample, kt/0.0019872041, gamma)
421  self.md = IMP.atom.MolecularDynamics(self.model)
422  self.md.set_maximum_time_step(maximum_time_step)
423  if sf:
424  self.md.set_scoring_function(sf)
425  else:
426  self.md.set_scoring_function(get_restraint_set(self.model))
427  self.md.add_optimizer_state(self.ltstate)
428  self.simulated_annealing = False
429  self.nframe = -1
430 
431  def set_kt(self, kt):
432  temp = kt/0.0019872041
433  self.ltstate.set_temperature(temp)
434  self.md.assign_velocities(temp)
435 
436  def set_simulated_annealing(self, min_temp, max_temp, min_temp_time,
437  max_temp_time):
438  self.simulated_annealing = True
439  self.tempmin = min_temp
440  self.tempmax = max_temp
441  self.timemin = min_temp_time
442  self.timemax = max_temp_time
443 
444  def temp_simulated_annealing(self):
445  if self.nframe % (self.timemin + self.timemax) < self.timemin:
446  value = 0.0
447  else:
448  value = 1.0
449  temp = self.tempmin + (self.tempmax - self.tempmin) * value
450  return temp
451 
452  def set_gamma(self, gamma):
453  self.ltstate.set_gamma(gamma)
454 
455  def optimize(self, nsteps):
456  # apply simulated annealing protocol
457  self.nframe += 1
458  if self.simulated_annealing:
459  self.temp = self.temp_simulated_annealing()
460  self.set_kt(self.temp)
461  self.md.optimize(nsteps)
462 
463  def get_output(self):
464  output = {}
465  output["MolecularDynamics_KineticEnergy"] = \
466  str(self.md.get_kinetic_energy())
467  return output
468 
469 
470 class ConjugateGradients(object):
471  """Sample using conjugate gradients"""
472 
473  def __init__(self, model, objects):
474  self.model = model
475  self.nframe = -1
476  self.cg = IMP.core.ConjugateGradients(self.model)
477  self.cg.set_scoring_function(get_restraint_set(self.model))
478 
479  def set_label(self, label):
480  self.label = label
481 
482  def get_frame_number(self):
483  return self.nframe
484 
485  def optimize(self, nstep):
486  self.nframe += 1
487  self.cg.optimize(nstep)
488 
489  def set_scoring_function(self, objectlist):
490  rs = IMP.RestraintSet(self.model, 1.0, 'sfo')
491  for ob in objectlist:
492  rs.add_restraint(ob.get_restraint())
494  self.cg.set_scoring_function(sf)
495 
496  def get_output(self):
497  output = {}
498  output["ConjugatedGradients_Nframe"] = str(self.nframe)
499  return output
500 
501 
502 class ReplicaExchange(object):
503  """Sample using replica exchange"""
504 
505  def __init__(self, model, tempmin, tempmax, samplerobjects, test=True,
506  replica_exchange_object=None):
507  '''
508  samplerobjects can be a list of MonteCarlo or MolecularDynamics
509  '''
510 
511  self.model = model
512  self.samplerobjects = samplerobjects
513  # min and max temperature
514  self.TEMPMIN_ = tempmin
515  self.TEMPMAX_ = tempmax
516 
517  if replica_exchange_object is None:
518  # initialize Replica Exchange class
519  try:
520  import IMP.mpi
521  print('ReplicaExchange: MPI was found. '
522  'Using Parallel Replica Exchange')
523  self.rem = IMP.mpi.ReplicaExchange()
524  except ImportError:
525  print('ReplicaExchange: Could not find MPI. '
526  'Using Serial Replica Exchange')
527  self.rem = _SerialReplicaExchange()
528 
529  else:
530  # get the replica exchange class instance from elsewhere
531  print('got existing rex object')
532  self.rem = replica_exchange_object
533 
534  # get number of replicas
535  nproc = self.rem.get_number_of_replicas()
536 
537  if nproc % 2 != 0 and not test:
538  raise Exception(
539  "number of replicas has to be even. "
540  "set test=True to run with odd number of replicas.")
541  # create array of temperatures, in geometric progression
542  temp = self.rem.create_temperatures(
543  self.TEMPMIN_,
544  self.TEMPMAX_,
545  nproc)
546  # get replica index
547  self.temperatures = temp
548 
549  myindex = self.rem.get_my_index()
550  # set initial value of the parameter (temperature) to exchange
551  self.rem.set_my_parameter("temp", [self.temperatures[myindex]])
552  for so in self.samplerobjects:
553  so.set_kt(self.temperatures[myindex])
554  self.nattempts = 0
555  self.nmintemp = 0
556  self.nmaxtemp = 0
557  self.nsuccess = 0
558 
559  def get_temperatures(self):
560  return self.temperatures
561 
562  def get_my_temp(self):
563  return self.rem.get_my_parameter("temp")[0]
564 
565  def get_my_index(self):
566  return self.rem.get_my_index()
567 
568  def swap_temp(self, nframe, score=None):
569  if score is None:
570  score = self.model.evaluate(False)
571  # get my replica index and temperature
572  _ = self.rem.get_my_index()
573  mytemp = self.rem.get_my_parameter("temp")[0]
574 
575  if mytemp == self.TEMPMIN_:
576  self.nmintemp += 1
577 
578  if mytemp == self.TEMPMAX_:
579  self.nmaxtemp += 1
580 
581  # score divided by kbt
582  myscore = score / mytemp
583 
584  # get my friend index and temperature
585  findex = self.rem.get_friend_index(nframe)
586  ftemp = self.rem.get_friend_parameter("temp", findex)[0]
587  # score divided by kbt
588  fscore = score / ftemp
589 
590  # try exchange
591  flag = self.rem.do_exchange(myscore, fscore, findex)
592 
593  self.nattempts += 1
594  # if accepted, change temperature
595  if (flag):
596  for so in self.samplerobjects:
597  so.set_kt(ftemp)
598  self.nsuccess += 1
599 
600  def get_output(self):
601  output = {}
602  if self.nattempts != 0:
603  output["ReplicaExchange_SwapSuccessRatio"] = str(
604  float(self.nsuccess) / self.nattempts)
605  output["ReplicaExchange_MinTempFrequency"] = str(
606  float(self.nmintemp) / self.nattempts)
607  output["ReplicaExchange_MaxTempFrequency"] = str(
608  float(self.nmaxtemp) / self.nattempts)
609  else:
610  output["ReplicaExchange_SwapSuccessRatio"] = str(0)
611  output["ReplicaExchange_MinTempFrequency"] = str(0)
612  output["ReplicaExchange_MaxTempFrequency"] = str(0)
613  output["ReplicaExchange_CurrentTemp"] = str(self.get_my_temp())
614  return output
615 
616 
617 class MPI_values(object):
618  def __init__(self, replica_exchange_object=None):
619  """Query values (ie score, and others)
620  from a set of parallel jobs"""
621 
622  if replica_exchange_object is None:
623  # initialize Replica Exchange class
624  try:
625  import IMP.mpi
626  print('MPI_values: MPI was found. '
627  'Using Parallel Replica Exchange')
628  self.rem = IMP.mpi.ReplicaExchange()
629  except ImportError:
630  print('MPI_values: Could not find MPI. '
631  'Using Serial Replica Exchange')
632  self.rem = _SerialReplicaExchange()
633 
634  else:
635  # get the replica exchange class instance from elsewhere
636  print('got existing rex object')
637  self.rem = replica_exchange_object
638 
639  def set_value(self, name, value):
640  self.rem.set_my_parameter(name, [value])
641 
642  def get_values(self, name):
643  values = []
644  for i in range(self.rem.get_number_of_replicas()):
645  v = self.rem.get_friend_parameter(name, i)[0]
646  values.append(v)
647  return values
648 
649  def get_percentile(self, name):
650  value = self.rem.get_my_parameter(name)[0]
651  values = sorted(self.get_values(name))
652  ind = values.index(value)
653  percentile = float(ind)/len(values)
654  return percentile
def __init__
samplerobjects can be a list of MonteCarlo or MolecularDynamics
Definition: samplers.py:505
A Monte Carlo optimizer.
Definition: MonteCarlo.h:45
A class to implement Hamiltonian Replica Exchange.
Maintains temperature during molecular dynamics.
Sample using molecular dynamics.
Definition: samplers.py:397
def get_nuisance_movers_parameters
Return a dictionary with the mover parameters for nuisance parameters.
Definition: samplers.py:196
Miscellaneous utilities.
Definition: tools.py:1
Modify the transformation of a rigid body.
Simple conjugate gradients optimizer.
Sample using conjugate gradients.
Definition: samplers.py:470
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:34
static FloatKeys get_internal_coordinate_keys()
Definition: rigid_bodies.h:709
Object used to hold a set of restraints.
Definition: RestraintSet.h:37
Simple molecular dynamics optimizer.
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:20
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:421
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:23
static bool get_is_setup(const IMP::ParticleAdaptor &p)
Definition: rigid_bodies.h:752
Sample using replica exchange.
Definition: samplers.py:502
Inferential scoring building on methods developed as part of the Inferential Structure Determination ...