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