Package csb :: Package statistics :: Package samplers :: Package mc :: Module multichain
[frames] | no frames]

Source Code for Module csb.statistics.samplers.mc.multichain

   1  """ 
   2  Implements several extended-ensemble Monte Carlo sampling algorithms. 
   3   
   4  Here is a short example which shows how to sample from a PDF using the replica 
   5  exchange with non-equilibrium switches (RENS) method. It draws 5000 samples from 
   6  a 1D normal distribution using the RENS algorithm working on three Markov chains 
   7  being generated by the HMC algorithm: 
   8   
   9   
  10      >>> import numpy 
  11      >>> from numpy import sqrt 
  12      >>> from csb.io.plots import Chart 
  13      >>> from csb.statistics.pdf import Normal 
  14      >>> from csb.statistics.samplers import State 
  15      >>> from csb.statistics.samplers.mc import ThermostattedMDRENSSwapParameterInfo, AlternatingAdjacentSwapScheme 
  16      >>> from csb.statistics.samplers.mc.multichain import ThermostattedMDRENS 
  17      >>> from csb.statistics.samplers.mc.singlechain import HMCSampler 
  18   
  19      >>> # Pick some initial state for the different Markov chains: 
  20      >>> initial_state = State(numpy.array([1.])) 
  21   
  22      >>> # Set standard deviations: 
  23      >>> std_devs = [1./sqrt(5), 1. / sqrt(3), 1.] 
  24   
  25      >>> # Set HMC timesteps and trajectory length: 
  26      >>> hmc_timesteps = [0.6, 0.7, 0.6] 
  27      >>> hmc_trajectory_length = 20 
  28      >>> hmc_gradients = [lambda q, t: 1 / (std_dev ** 2) * q for std_dev in std_devs] 
  29   
  30      >>> # Set parameters for the thermostatted RENS algorithm: 
  31      >>> rens_trajectory_length = 30 
  32      >>> rens_timesteps = [0.3, 0.5] 
  33   
  34      >>> # Set interpolation gradients as a function of the work parameter l: 
  35      >>> rens_gradients = [lambda q, l, i=i: (l / (std_devs[i + 1] ** 2) + (1 - l) / (std_devs[i] ** 2)) * q  
  36                            for i in range(len(std_devs)-1)] 
  37   
  38      >>> # Initialize HMC samplers: 
  39      >>> samplers = [HMCSampler(Normal(sigma=std_devs[i]), initial_state, hmc_gradients[i], hmc_timesteps[i], 
  40                      hmc_trajectory_length) for i in range(len(std_devs))] 
  41   
  42      >>> # Create swap parameter objects: 
  43      params = [ThermostattedMDRENSSwapParameterInfo(samplers[0], samplers[1], rens_timesteps[0], 
  44                rens_trajectory_length, rens_gradients[0]), 
  45                ThermostattedMDRENSSwapParameterInfo(samplers[1], samplers[2], rens_timesteps[1], 
  46                rens_trajectory_length, rens_gradients[1])] 
  47   
  48      >>> # Initialize thermostatted RENS algorithm: 
  49      >>> algorithm = ThermostattedMDRENS(samplers, params) 
  50   
  51      >>> # Initialize swapping scheme: 
  52      >>> swapper = AlternatingAdjacentSwapScheme(algorithm) 
  53   
  54      >>> # Initialize empty list which will store the samples: 
  55      >>> states = [] 
  56      >>> for i in range(5000): 
  57              if i % 5 == 0: 
  58                  swapper.swap_all() 
  59              states.append(algorithm.sample()) 
  60   
  61      >>> # Print acceptance rates: 
  62      >>> print('HMC acceptance rates:', [s.acceptance_rate for s in samplers]) 
  63      >>> print('swap acceptance rates:', algorithm.acceptance_rates) 
  64   
  65      >>> # Create and plot histogram for first sampler and numpy.random.normal reference: 
  66      >>> chart = Chart() 
  67      >>> rawstates = [state[0].position[0] for state in states] 
  68      >>> chart.plot.hist([numpy.random.normal(size=5000, scale=std_devs[0]), rawstates], bins=30, normed=True) 
  69      >>> chart.plot.legend(['numpy.random.normal', 'RENS + HMC']) 
  70      >>> chart.show() 
  71   
  72   
  73  For L{ReplicaExchangeMC} (RE), the procedure is easier because apart from the 
  74  two sampler instances the corresponding L{RESwapParameterInfo} objects take 
  75  no arguments. 
  76   
  77  Every replica exchange algorithm in this module (L{ReplicaExchangeMC}, L{MDRENS}, 
  78  L{ThermostattedMDRENS}) is used in a similar way. A simulation is always 
  79  initialized with a list of samplers (instances of classes derived from 
  80  L{AbstractSingleChainMC}) and a list of L{AbstractSwapParameterInfo} objects 
  81  suited for the algorithm under consideration. Every L{AbstractSwapParameterInfo} 
  82  object holds all the information needed to perform a swap between two samplers. 
  83  The usual scheme is to swap only adjacent replicae in a scheme:: 
  84   
  85      1 <--> 2, 3 <--> 4, ... 
  86      2 <--> 3, 4 <--> 5, ... 
  87      1 <--> 2, 3 <--> 4, ... 
  88       
  89  This swapping scheme is implemented in the L{AlternatingAdjacentSwapScheme} class, 
  90  but different schemes can be easily implemented by deriving from L{AbstractSwapScheme}. 
  91  Then the simulation is run by looping over the number of samples to be drawn 
  92  and calling the L{AbstractExchangeMC.sample} method of the algorithm. By calling 
  93  the L{AbstractSwapScheme.swap_all} method of the specific L{AbstractSwapScheme} 
  94  implementation, all swaps defined in the list of L{AbstractSwapParameterInfo} 
  95  objects are performed according to the swapping scheme. The 
  96  L{AbstractSwapScheme.swap_all} method may be called for example after sampling 
  97  intervals of a fixed length or randomly. 
  98  """ 
  99   
 100  import numpy 
 101   
 102  import csb.numeric 
 103   
 104  from abc import ABCMeta, abstractmethod 
 105   
 106  from csb.statistics.samplers import EnsembleState 
 107  from csb.statistics.samplers.mc import AbstractMC, Trajectory, MCCollection, augment_state 
 108  from csb.statistics.samplers.mc.propagators import MDPropagator, ThermostattedMDPropagator 
 109  from csb.statistics.samplers.mc.neqsteppropagator import NonequilibriumStepPropagator 
 110  from csb.statistics.samplers.mc.neqsteppropagator import Protocol, Step, ReducedHamiltonian 
 111  from csb.statistics.samplers.mc.neqsteppropagator import ReducedHamiltonianPerturbation 
 112  from csb.statistics.samplers.mc.neqsteppropagator import HMCPropagation, HMCPropagationParam 
 113  from csb.statistics.samplers.mc.neqsteppropagator import HamiltonianSysInfo, NonequilibriumTrajectory 
 114  from csb.numeric.integrators import AbstractGradient, FastLeapFrog 
115 116 117 -class AbstractEnsembleMC(AbstractMC):
118 """ 119 Abstract class for Monte Carlo sampling algorithms simulating several ensembles. 120 121 @param samplers: samplers which sample from their respective equilibrium distributions 122 @type samplers: list of L{AbstractSingleChainMC} 123 """ 124 125 __metaclass__ = ABCMeta 126
127 - def __init__(self, samplers):
128 129 self._samplers = MCCollection(samplers) 130 state = EnsembleState([x.state for x in self._samplers]) 131 132 super(AbstractEnsembleMC, self).__init__(state)
133
134 - def sample(self):
135 """ 136 Draw an ensemble sample. 137 138 @rtype: L{EnsembleState} 139 """ 140 141 sample = EnsembleState([sampler.sample() for sampler in self._samplers]) 142 self.state = sample 143 144 return sample
145 146 @property
147 - def energy(self):
148 """ 149 Total ensemble energy. 150 """ 151 return sum([x.energy for x in self._samplers])
152
153 154 -class AbstractExchangeMC(AbstractEnsembleMC):
155 """ 156 Abstract class for Monte Carlo sampling algorithms employing some replica exchange method. 157 158 @param samplers: samplers which sample from their respective equilibrium distributions 159 @type samplers: list of L{AbstractSingleChainMC} 160 161 @param param_infos: list of ParameterInfo instances providing information needed 162 for performing swaps 163 @type param_infos: list of L{AbstractSwapParameterInfo} 164 """ 165 166 __metaclass__ = ABCMeta 167
168 - def __init__(self, samplers, param_infos):
169 super(AbstractExchangeMC, self).__init__(samplers) 170 171 self._swaplist1 = [] 172 self._swaplist2 = [] 173 self._currentswaplist = self._swaplist1 174 175 self._param_infos = param_infos 176 self._statistics = SwapStatistics(self._param_infos)
177
178 - def _checkstate(self, state):
179 if not isinstance(state, EnsembleState): 180 raise TypeError(state)
181
182 - def swap(self, index):
183 """ 184 Perform swap between sampler pair described by param_infos[index] 185 and return outcome (true = accepted, false = rejected). 186 187 @param index: index of swap pair in param_infos 188 @type index: int 189 190 @rtype: boolean 191 """ 192 param_info = self._param_infos[index] 193 swapcom = self._propose_swap(param_info) 194 swapcom = self._calc_pacc_swap(swapcom) 195 result = self._accept_swap(swapcom) 196 197 self.state = EnsembleState([x.state for x in self._samplers]) 198 199 self.statistics.stats[index].update(result) 200 201 return result
202 203 @abstractmethod
204 - def _propose_swap(self, param_info):
205 """ 206 Calculate proposal states for a swap between two samplers. 207 208 @param param_info: ParameterInfo instance holding swap parameters 209 @type param_info: L{AbstractSwapParameterInfo} 210 211 @rtype: L{AbstractSwapCommunicator} 212 """ 213 pass
214 215 @abstractmethod
216 - def _calc_pacc_swap(self, swapcom):
217 """ 218 Calculate probability to accept a swap given initial and proposal states. 219 220 @param swapcom: SwapCommunicator instance holding information to be communicated 221 between distinct swap substeps 222 @type swapcom: L{AbstractSwapCommunicator} 223 224 @rtype: L{AbstractSwapCommunicator} 225 """ 226 pass
227
228 - def _accept_swap(self, swapcom):
229 """ 230 Accept / reject an exchange between two samplers given proposal states and 231 the acceptance probability and returns the outcome (true = accepted, false = rejected). 232 233 @param swapcom: SwapCommunicator instance holding information to be communicated 234 between distinct swap substeps 235 @type swapcom: L{AbstractSwapCommunicator} 236 237 @rtype: boolean 238 """ 239 240 if numpy.random.random() < swapcom.acceptance_probability: 241 if swapcom.sampler1.state.momentum is None and swapcom.sampler2.state.momentum is None: 242 swapcom.traj12.final.momentum = None 243 swapcom.traj21.final.momentum = None 244 swapcom.sampler1.state = swapcom.traj21.final 245 swapcom.sampler2.state = swapcom.traj12.final 246 return True 247 else: 248 return False
249 250 @property
251 - def acceptance_rates(self):
252 """ 253 Return swap acceptance rates. 254 255 @rtype: list of floats 256 """ 257 return self.statistics.acceptance_rates
258 259 @property
260 - def param_infos(self):
261 """ 262 List of SwapParameterInfo instances holding all necessary parameters. 263 264 @rtype: list of L{AbstractSwapParameterInfo} 265 """ 266 return self._param_infos
267 268 @property
269 - def statistics(self):
270 return self._statistics
271
272 - def _update_statistics(self, index, accepted):
273 """ 274 Update statistics of a given swap process. 275 276 @param index: position of swap statistics to be updated 277 @type index: int 278 279 @param accepted: outcome of the swap 280 @type accepted: boolean 281 """ 282 283 self._stats[index][0] += 1 284 self._stats[index][1] += int(accepted)
285
286 287 -class AbstractSwapParameterInfo(object):
288 """ 289 Subclass instances hold all parameters necessary for performing a swap 290 between two given samplers. 291 """ 292 293 __metaclass__ = ABCMeta 294
295 - def __init__(self, sampler1, sampler2):
296 """ 297 @param sampler1: First sampler 298 @type sampler1: L{AbstractSingleChainMC} 299 300 @param sampler2: Second sampler 301 @type sampler2: L{AbstractSingleChainMC} 302 """ 303 304 self._sampler1 = sampler1 305 self._sampler2 = sampler2
306 307 @property
308 - def sampler1(self):
309 return self._sampler1
310 311 @property
312 - def sampler2(self):
313 return self._sampler2
314
315 316 -class AbstractSwapCommunicator(object):
317 """ 318 Holds all the information which needs to be communicated between 319 distinct swap substeps. 320 321 @param param_info: ParameterInfo instance holding swap parameters 322 @type param_info: L{AbstractSwapParameterInfo} 323 324 @param traj12: Forward trajectory 325 @type traj12: L{Trajectory} 326 327 @param traj21: Reverse trajectory 328 @type traj21: L{Trajectory} 329 """ 330 331 __metaclass__ = ABCMeta 332
333 - def __init__(self, param_info, traj12, traj21):
334 335 self._sampler1 = param_info.sampler1 336 self._sampler2 = param_info.sampler2 337 338 self._traj12 = traj12 339 self._traj21 = traj21 340 341 self._param_info = param_info 342 343 self._acceptance_probability = None 344 self._accepted = False
345 346 @property
347 - def sampler1(self):
348 return self._sampler1
349 350 @property
351 - def sampler2(self):
352 return self._sampler2
353 354 @property
355 - def traj12(self):
356 return self._traj12
357 358 @property
359 - def traj21(self):
360 return self._traj21
361 362 @property
363 - def acceptance_probability(self):
364 return self._acceptance_probability
365 @acceptance_probability.setter
366 - def acceptance_probability(self, value):
367 self._acceptance_probability = value
368 369 @property
370 - def accepted(self):
371 return self._accepted
372 @accepted.setter
373 - def accepted(self, value):
374 self._accepted = value
375 376 @property
377 - def param_info(self):
378 return self._param_info
379
380 381 -class ReplicaExchangeMC(AbstractExchangeMC):
382 """ 383 Replica Exchange (RE, Swendsen & Yang 1986) implementation. 384 """ 385
386 - def _propose_swap(self, param_info):
392
393 - def _calc_pacc_swap(self, swapcom):
394 395 E1 = lambda x:-swapcom.sampler1._pdf.log_prob(x) 396 E2 = lambda x:-swapcom.sampler2._pdf.log_prob(x) 397 398 T1 = swapcom.sampler1.temperature 399 T2 = swapcom.sampler2.temperature 400 401 state1 = swapcom.traj12.initial 402 state2 = swapcom.traj21.initial 403 404 proposal1 = swapcom.traj21.final 405 proposal2 = swapcom.traj12.final 406 407 swapcom.acceptance_probability = csb.numeric.exp(-E1(proposal1.position) / T1 408 + E1(state1.position) / T1 409 - E2(proposal2.position) / T2 410 + E2(state2.position) / T2) 411 412 return swapcom
413
414 415 -class RESwapParameterInfo(AbstractSwapParameterInfo):
416 """ 417 Holds parameters for a standard Replica Exchange swap. 418 """ 419 pass
420
421 422 -class RESwapCommunicator(AbstractSwapCommunicator):
423 """ 424 Holds all the information which needs to be communicated between distinct 425 RE swap substeps. 426 427 See L{AbstractSwapCommunicator} for constructor signature. 428 """ 429 pass
430
431 432 -class AbstractRENS(AbstractExchangeMC):
433 """ 434 Abstract Replica Exchange with Nonequilibrium Switches 435 (RENS, Ballard & Jarzynski 2009) class. 436 Subclasses implement various ways of generating trajectories 437 (deterministic or stochastic). 438 """ 439 440 __metaclass__ = ABCMeta 441
442 - def _propose_swap(self, param_info):
443 444 init_state1 = param_info.sampler1.state 445 init_state2 = param_info.sampler2.state 446 447 trajinfo12 = RENSTrajInfo(param_info, init_state1, direction="fw") 448 trajinfo21 = RENSTrajInfo(param_info, init_state2, direction="bw") 449 450 traj12 = self._run_traj_generator(trajinfo12) 451 traj21 = self._run_traj_generator(trajinfo21) 452 453 return RENSSwapCommunicator(param_info, traj12, traj21)
454
455 - def _setup_protocol(self, traj_info):
456 """ 457 Sets the protocol lambda(t) to either the forward or the reverse protocol. 458 459 @param traj_info: TrajectoryInfo object holding information neccessary to 460 calculate the rens trajectories. 461 @type traj_info: L{RENSTrajInfo} 462 """ 463 464 if traj_info.direction == "fw": 465 return traj_info.param_info.protocol 466 else: 467 return lambda t, tau: traj_info.param_info.protocol(tau - t, tau) 468 469 return protocol
470
471 - def _get_init_temperature(self, traj_info):
472 """ 473 Determine the initial temperature of a RENS trajectory. 474 475 @param traj_info: TrajectoryInfo object holding information neccessary to 476 calculate the RENS trajectory. 477 @type traj_info: L{RENSTrajInfo} 478 """ 479 480 if traj_info.direction == "fw": 481 return traj_info.param_info.sampler1.temperature 482 else: 483 return traj_info.param_info.sampler2.temperature
484 485 @abstractmethod
486 - def _calc_works(self, swapcom):
487 """ 488 Calculates the works expended during the nonequilibrium 489 trajectories. 490 491 @param swapcom: Swap communicator object holding all the 492 neccessary information. 493 @type swapcom: L{RENSSwapCommunicator} 494 495 @return: The expended during the forward and the backward 496 trajectory. 497 @rtype: 2-tuple of floats 498 """ 499 500 pass
501
502 - def _calc_pacc_swap(self, swapcom):
503 504 work12, work21 = self._calc_works(swapcom) 505 swapcom.acceptance_probability = csb.numeric.exp(-work12 - work21) 506 507 return swapcom
508
509 - def _propagator_factory(self, traj_info):
510 """ 511 Factory method which produces the propagator object used to calculate 512 the RENS trajectories. 513 514 @param traj_info: TrajectoryInfo object holding information neccessary to 515 calculate the rens trajectories. 516 @type traj_info: L{RENSTrajInfo} 517 @rtype: L{AbstractPropagator} 518 """ 519 pass
520
521 - def _run_traj_generator(self, traj_info):
522 """ 523 Run the trajectory generator which generates a trajectory 524 of a given length between the states of two samplers. 525 526 @param traj_info: TrajectoryInfo instance holding information 527 needed to generate a nonequilibrium trajectory 528 @type traj_info: L{RENSTrajInfo} 529 530 @rtype: L{Trajectory} 531 """ 532 533 init_temperature = self._get_init_temperature(traj_info) 534 535 init_state = traj_info.init_state.clone() 536 537 if init_state.momentum is None: 538 init_state = augment_state(init_state, 539 init_temperature, 540 traj_info.param_info.mass_matrix) 541 542 gen = self._propagator_factory(traj_info) 543 544 traj = gen.generate(init_state, int(traj_info.param_info.traj_length)) 545 546 return traj
547
548 549 -class AbstractRENSSwapParameterInfo(RESwapParameterInfo):
550 """ 551 Holds parameters for a RENS swap. 552 """ 553 554 __metaclass__ = ABCMeta 555
556 - def __init__(self, sampler1, sampler2, protocol):
557 558 super(AbstractRENSSwapParameterInfo, self).__init__(sampler1, sampler2) 559 560 ## Can't pass the linear protocol as a default argument because of a reported bug 561 ## in epydoc parsing which makes it fail building the docs. 562 self._protocol = None 563 if protocol is None: 564 self._protocol = lambda t, tau: t / tau 565 else: 566 self._protocol = protocol
567 568 @property
569 - def protocol(self):
570 """ 571 Switching protocol determining the time dependence 572 of the switching parameter. 573 """ 574 return self._protocol
575 @protocol.setter
576 - def protocol(self, value):
577 self._protocol = value
578
579 580 -class RENSSwapCommunicator(AbstractSwapCommunicator):
581 """ 582 Holds all the information which needs to be communicated between distinct 583 RENS swap substeps. 584 585 See L{AbstractSwapCommunicator} for constructor signature. 586 """ 587 588 pass
589
590 591 -class RENSTrajInfo(object):
592 """ 593 Holds information necessary for calculating a RENS trajectory. 594 595 @param param_info: ParameterInfo instance holding swap parameters 596 @type param_info: L{AbstractSwapParameterInfo} 597 598 @param init_state: state from which the trajectory is supposed to start 599 @type init_state: L{State} 600 601 @param direction: Either "fw" or "bw", indicating a forward or backward 602 trajectory. This is neccessary to pick the protocol or 603 the reversed protocol, respectively. 604 @type direction: string, either "fw" or "bw" 605 """ 606
607 - def __init__(self, param_info, init_state, direction):
608 609 self._param_info = param_info 610 self._init_state = init_state 611 self._direction = direction
612 613 @property
614 - def param_info(self):
615 return self._param_info
616 617 @property
618 - def init_state(self):
619 return self._init_state
620 621 @property
622 - def direction(self):
623 return self._direction
624
625 626 -class MDRENS(AbstractRENS):
627 """ 628 Replica Exchange with Nonequilibrium Switches (RENS, Ballard & Jarzynski 2009) 629 with Molecular Dynamics (MD) trajectories. 630 631 @param samplers: Samplers which sample their 632 respective equilibrium distributions 633 @type samplers: list of L{AbstractSingleChainMC} 634 635 @param param_infos: ParameterInfo instance holding 636 information required to perform a MDRENS swap 637 @type param_infos: list of L{MDRENSSwapParameterInfo} 638 639 @param integrator: Subclass of L{AbstractIntegrator} to be used to 640 calculate the non-equilibrium trajectories 641 @type integrator: type 642 """ 643
644 - def __init__(self, samplers, param_infos, 645 integrator=csb.numeric.integrators.FastLeapFrog):
646 647 super(MDRENS, self).__init__(samplers, param_infos) 648 649 self._integrator = integrator
650
651 - def _propagator_factory(self, traj_info):
652 653 protocol = self._setup_protocol(traj_info) 654 tau = traj_info.param_info.traj_length * traj_info.param_info.timestep 655 factory = InterpolationFactory(protocol, tau) 656 gen = MDPropagator(factory.build_gradient(traj_info.param_info.gradient), 657 traj_info.param_info.timestep, 658 mass_matrix=traj_info.param_info.mass_matrix, 659 integrator=self._integrator) 660 661 return gen
662
663 - def _calc_works(self, swapcom):
664 665 T1 = swapcom.param_info.sampler1.temperature 666 T2 = swapcom.param_info.sampler2.temperature 667 668 heat12 = swapcom.traj12.heat 669 heat21 = swapcom.traj21.heat 670 671 proposal1 = swapcom.traj21.final 672 proposal2 = swapcom.traj12.final 673 674 state1 = swapcom.traj12.initial 675 state2 = swapcom.traj21.initial 676 677 if swapcom.param_info.mass_matrix.is_unity_multiple: 678 inverse_mass_matrix = 1.0 / swapcom.param_info.mass_matrix[0][0] 679 else: 680 inverse_mass_matrix = swapcom.param_info.mass_matrix.inverse 681 682 E1 = lambda x:-swapcom.sampler1._pdf.log_prob(x) 683 E2 = lambda x:-swapcom.sampler2._pdf.log_prob(x) 684 K = lambda x: 0.5 * numpy.dot(x.T, numpy.dot(inverse_mass_matrix, x)) 685 686 w12 = (K(proposal2.momentum) + E2(proposal2.position)) / T2 - \ 687 (K(state1.momentum) + E1(state1.position)) / T1 - heat12 688 w21 = (K(proposal1.momentum) + E1(proposal1.position)) / T1 - \ 689 (K(state2.momentum) + E2(state2.position)) / T2 - heat21 690 691 return w12, w21
692
693 694 -class MDRENSSwapParameterInfo(RESwapParameterInfo):
695 """ 696 Holds parameters for a MDRENS swap. 697 698 @param sampler1: First sampler 699 @type sampler1: L{AbstractSingleChainMC} 700 701 @param sampler2: Second sampler 702 @type sampler2: L{AbstractSingleChainMC} 703 704 @param timestep: Integration timestep 705 @type timestep: float 706 707 @param traj_length: Trajectory length in number of timesteps 708 @type traj_length: int 709 710 @param gradient: Gradient which determines the dynamics during a trajectory 711 @type gradient: L{AbstractGradient} 712 713 @param protocol: Switching protocol determining the time dependence of the 714 switching parameter. It is a function M{f} taking the running 715 time t and the switching time tau to yield a value in M{[0, 1]} 716 with M{f(0, tau) = 0} and M{f(tau, tau) = 1}. Default is a linear 717 protocol, which is being set manually due to an epydoc bug 718 @type protocol: callable 719 720 @param mass_matrix: Mass matrix 721 @type mass_matrix: n-dimensional matrix of type L{InvertibleMatrix} with n being the dimension 722 of the configuration space, that is, the dimension of 723 the position / momentum vectors 724 """ 725
726 - def __init__(self, sampler1, sampler2, timestep, traj_length, gradient, 727 protocol=None, mass_matrix=None):
728 729 super(MDRENSSwapParameterInfo, self).__init__(sampler1, sampler2) 730 731 self._mass_matrix = mass_matrix 732 if self.mass_matrix is None: 733 d = len(sampler1.state.position) 734 self.mass_matrix = csb.numeric.InvertibleMatrix(numpy.eye(d), numpy.eye(d)) 735 736 self._traj_length = traj_length 737 self._gradient = gradient 738 self._timestep = timestep 739 740 ## Can't pass the linear protocol as a default argument because of a reported bug 741 ## in epydoc parsing which makes it fail building the docs. 742 self._protocol = None 743 if protocol is None: 744 self._protocol = lambda t, tau: t / tau 745 else: 746 self._protocol = protocol
747 748 @property
749 - def timestep(self):
750 """ 751 Integration timestep. 752 """ 753 return self._timestep
754 @timestep.setter
755 - def timestep(self, value):
756 self._timestep = float(value)
757 758 @property
759 - def traj_length(self):
760 """ 761 Trajectory length in number of integration steps. 762 """ 763 return self._traj_length
764 @traj_length.setter
765 - def traj_length(self, value):
766 self._traj_length = int(value)
767 768 @property
769 - def gradient(self):
770 """ 771 Gradient which governs the equations of motion. 772 """ 773 return self._gradient
774 775 @property
776 - def mass_matrix(self):
777 return self._mass_matrix
778 @mass_matrix.setter
779 - def mass_matrix(self, value):
780 self._mass_matrix = value
781 782 @property
783 - def protocol(self):
784 """ 785 Switching protocol determining the time dependence 786 of the switching parameter. 787 """ 788 return self._protocol
789 @protocol.setter
790 - def protocol(self, value):
791 self._protocol = value
792
793 794 -class ThermostattedMDRENS(MDRENS):
795 """ 796 Replica Exchange with Nonequilibrium Switches (RENS, Ballard & Jarzynski, 2009) 797 with Andersen-thermostatted Molecular Dynamics (MD) trajectories. 798 799 @param samplers: Samplers which sample their 800 respective equilibrium distributions 801 @type samplers: list of L{AbstractSingleChainMC} 802 803 @param param_infos: ParameterInfo instance holding 804 information required to perform a MDRENS swap 805 @type param_infos: list of L{ThermostattedMDRENSSwapParameterInfo} 806 807 @param integrator: Subclass of L{AbstractIntegrator} to be used to 808 calculate the non-equilibrium trajectories 809 @type integrator: type 810 """ 811
812 - def __init__(self, samplers, param_infos, integrator=csb.numeric.integrators.LeapFrog):
815
816 - def _propagator_factory(self, traj_info):
817 818 protocol = self._setup_protocol(traj_info) 819 tau = traj_info.param_info.traj_length * traj_info.param_info.timestep 820 factory = InterpolationFactory(protocol, tau) 821 822 grad = factory.build_gradient(traj_info.param_info.gradient) 823 temp = factory.build_temperature(traj_info.param_info.temperature) 824 825 gen = ThermostattedMDPropagator(grad, 826 traj_info.param_info.timestep, temperature=temp, 827 collision_probability=traj_info.param_info.collision_probability, 828 update_interval=traj_info.param_info.collision_interval, 829 mass_matrix=traj_info.param_info.mass_matrix, 830 integrator=self._integrator) 831 832 return gen
833
834 -class ThermostattedMDRENSSwapParameterInfo(MDRENSSwapParameterInfo):
835 """ 836 @param sampler1: First sampler 837 @type sampler1: subclass instance of L{AbstractSingleChainMC} 838 839 @param sampler2: Second sampler 840 @type sampler2: subclass instance of L{AbstractSingleChainMC} 841 842 @param timestep: Integration timestep 843 @type timestep: float 844 845 @param traj_length: Trajectory length in number of timesteps 846 @type traj_length: int 847 848 @param gradient: Gradient which determines the dynamics during a trajectory 849 @type gradient: subclass instance of L{AbstractGradient} 850 851 @param mass_matrix: Mass matrix 852 @type mass_matrix: n-dimensional L{InvertibleMatrix} with n being the dimension 853 of the configuration space, that is, the dimension of 854 the position / momentum vectors 855 856 @param protocol: Switching protocol determining the time dependence of the 857 switching parameter. It is a function f taking the running 858 time t and the switching time tau to yield a value in [0, 1] 859 with f(0, tau) = 0 and f(tau, tau) = 1 860 @type protocol: callable 861 862 @param temperature: Temperature interpolation function. 863 @type temperature: Real-valued function mapping from [0,1] to R. 864 T(0) = temperature of the ensemble sampler1 samples from, T(1) = temperature 865 of the ensemble sampler2 samples from 866 867 @param collision_probability: Probability for a collision with the heatbath during one timestep 868 @type collision_probability: float 869 870 @param collision_interval: Interval during which collision may occur with probability 871 collision_probability 872 @type collision_interval: int 873 """ 874
875 - def __init__(self, sampler1, sampler2, timestep, traj_length, gradient, mass_matrix=None, 876 protocol=None, temperature=lambda l: 1.0, 877 collision_probability=0.1, collision_interval=1):
878 879 super(ThermostattedMDRENSSwapParameterInfo, self).__init__(sampler1, sampler2, timestep, 880 traj_length, gradient, 881 mass_matrix=mass_matrix, 882 protocol=protocol) 883 884 self._collision_probability = None 885 self._collision_interval = None 886 self._temperature = temperature 887 self.collision_probability = collision_probability 888 self.collision_interval = collision_interval
889 890 @property
891 - def collision_probability(self):
892 """ 893 Probability for a collision with the heatbath during one timestep. 894 """ 895 return self._collision_probability
896 @collision_probability.setter
897 - def collision_probability(self, value):
898 self._collision_probability = float(value)
899 900 @property
901 - def collision_interval(self):
902 """ 903 Interval during which collision may occur with probability 904 C{collision_probability}. 905 """ 906 return self._collision_interval
907 @collision_interval.setter
908 - def collision_interval(self, value):
909 self._collision_interval = int(value)
910 911 @property
912 - def temperature(self):
913 return self._temperature
914
915 916 -class AbstractStepRENS(AbstractRENS):
917 """ 918 Replica Exchange with Nonequilibrium Switches (RENS, Ballard & Jarzynski 2009) 919 with stepwise trajectories as described in Nilmeier et al., "Nonequilibrium candidate 920 Monte Carlo is an efficient tool for equilibrium simulation", PNAS 2011. 921 The switching parameter dependence of the Hamiltonian is a linear interpolation 922 between the PDFs of the sampler objects, 923 M{H(S{lambda}) = H_2 * S{lambda} + (1 - S{lambda}) * H_1}. 924 The perturbation kernel is a thermodynamic perturbation and the propagation is subclass 925 responsibility. 926 Note that due to the linear interpolations between the two Hamiltonians, the 927 log-probability has to be evaluated four times per perturbation step which can be 928 costly. In this case it is advisable to define the intermediate log probabilities 929 in _run_traj_generator differently. 930 931 @param samplers: Samplers which sample their respective equilibrium distributions 932 @type samplers: list of L{AbstractSingleChainMC} 933 934 @param param_infos: ParameterInfo instances holding 935 information required to perform a HMCStepRENS swaps 936 @type param_infos: list of L{AbstractSwapParameterInfo} 937 """ 938 939 __metaclass__ = ABCMeta 940
941 - def __init__(self, samplers, param_infos):
942 943 super(AbstractStepRENS, self).__init__(samplers, param_infos) 944 945 self._evaluate_im_works = True
946 947 @abstractmethod
948 - def _setup_propagations(self, im_sys_infos, param_info):
949 """ 950 Set up the propagation steps using the information about the current system 951 setup and parameters from the SwapParameterInfo object. 952 953 @param im_sys_infos: Information about the intermediate system setups 954 @type im_sys_infos: List of L{AbstractSystemInfo} 955 956 @param param_info: SwapParameterInfo object containing parameters for the 957 propagations like timesteps, trajectory lengths etc. 958 @type param_info: L{AbstractSwapParameterInfo} 959 """ 960 961 pass
962 963 @abstractmethod
964 - def _add_gradients(self, im_sys_infos, param_info, t_prot):
965 """ 966 If needed, set im_sys_infos.hamiltonian.gradient. 967 968 @param im_sys_infos: Information about the intermediate system setups 969 @type im_sys_infos: List of L{AbstractSystemInfo} 970 971 @param param_info: SwapParameterInfo object containing parameters for the 972 propagations like timesteps, trajectory lengths etc. 973 @type param_info: L{AbstractSwapParameterInfo} 974 975 @param t_prot: Switching protocol defining the time dependence of the switching 976 parameter. 977 @type t_prot: callable 978 """ 979 980 pass
981
982 - def _setup_stepwise_protocol(self, traj_info):
983 """ 984 Sets up the stepwise protocol consisting of perturbation and relaxation steps. 985 986 @param traj_info: TrajectoryInfo instance holding information 987 needed to generate a nonequilibrium trajectory 988 @type traj_info: L{RENSTrajInfo} 989 990 @rtype: L{Protocol} 991 """ 992 993 pdf1 = traj_info.param_info.sampler1._pdf 994 pdf2 = traj_info.param_info.sampler2._pdf 995 T1 = traj_info.param_info.sampler1.temperature 996 T2 = traj_info.param_info.sampler2.temperature 997 traj_length = traj_info.param_info.intermediate_steps 998 prot = self._setup_protocol(traj_info) 999 t_prot = lambda i: prot(float(i), float(traj_length)) 1000 1001 im_log_probs = [lambda x, i=i: pdf2.log_prob(x) * t_prot(i) + \ 1002 (1 - t_prot(i)) * pdf1.log_prob(x) 1003 for i in range(traj_length + 1)] 1004 1005 im_temperatures = [T2 * t_prot(i) + (1 - t_prot(i)) * T1 1006 for i in range(traj_length + 1)] 1007 im_reduced_hamiltonians = [ReducedHamiltonian(im_log_probs[i], 1008 temperature=im_temperatures[i]) 1009 for i in range(traj_length + 1)] 1010 im_sys_infos = [HamiltonianSysInfo(im_reduced_hamiltonians[i]) 1011 for i in range(traj_length + 1)] 1012 perturbations = [ReducedHamiltonianPerturbation(im_sys_infos[i], im_sys_infos[i+1]) 1013 for i in range(traj_length)] 1014 if self._evaluate_im_works == False: 1015 for p in perturbations: 1016 p.evaluate_work = False 1017 im_sys_infos = self._add_gradients(im_sys_infos, traj_info.param_info, t_prot) 1018 propagations = self._setup_propagations(im_sys_infos, traj_info.param_info) 1019 1020 steps = [Step(perturbations[i], propagations[i]) for i in range(traj_length)] 1021 1022 return Protocol(steps)
1023
1024 - def _propagator_factory(self, traj_info):
1025 1026 protocol = self._setup_stepwise_protocol(traj_info) 1027 gen = NonequilibriumStepPropagator(protocol) 1028 1029 return gen
1030
1031 - def _run_traj_generator(self, traj_info):
1032 1033 init_temperature = self._get_init_temperature(traj_info) 1034 1035 gen = self._propagator_factory(traj_info) 1036 1037 traj = gen.generate(traj_info.init_state) 1038 1039 return NonequilibriumTrajectory([traj_info.init_state, traj.final], jacobian=1.0, 1040 heat=traj.heat, work=traj.work, deltaH=traj.deltaH)
1041
1042 1043 -class HMCStepRENS(AbstractStepRENS):
1044 """ 1045 Replica Exchange with Nonequilibrium Switches (RENS, Ballard & Jarzynski 2009) 1046 with stepwise trajectories as described in Nilmeier et al., "Nonequilibrium candidate 1047 Monte Carlo is an efficient tool for equilibrium simulation", PNAS 2011. 1048 The switching parameter dependence of the Hamiltonian is a linear interpolation 1049 between the PDFs of the sampler objects, 1050 M{H(S{lambda}) = H_2 * S{lambda} + (1 - S{lambda}) * H_1}. 1051 The perturbation kernel is a thermodynamic perturbation and the propagation is done using HMC. 1052 1053 Note that due to the linear interpolations between the two Hamiltonians, the 1054 log-probability and its gradient has to be evaluated four times per perturbation step which 1055 can be costly. In this case it is advisable to define the intermediate log probabilities 1056 in _run_traj_generator differently. 1057 1058 @param samplers: Samplers which sample their respective equilibrium distributions 1059 @type samplers: list of L{AbstractSingleChainMC} 1060 1061 @param param_infos: ParameterInfo instances holding 1062 information required to perform a HMCStepRENS swaps 1063 @type param_infos: list of L{HMCStepRENSSwapParameterInfo} 1064 """ 1065
1066 - def __init__(self, samplers, param_infos):
1067 1068 super(HMCStepRENS, self).__init__(samplers, param_infos)
1069
1070 - def _add_gradients(self, im_sys_infos, param_info, t_prot):
1071 1072 im_gradients = [lambda x, t, i=i: param_info.gradient(x, t_prot(i)) 1073 for i in range(param_info.intermediate_steps + 1)] 1074 1075 for i, s in enumerate(im_sys_infos): 1076 s.hamiltonian.gradient = im_gradients[i] 1077 1078 return im_sys_infos
1079
1080 - def _setup_propagations(self, im_sys_infos, param_info):
1081 1082 propagation_params = [HMCPropagationParam(param_info.timestep, 1083 param_info.hmc_traj_length, 1084 im_sys_infos[i+1].hamiltonian.gradient, 1085 param_info.hmc_iterations, 1086 mass_matrix=param_info.mass_matrix, 1087 integrator=param_info.integrator) 1088 for i in range(param_info.intermediate_steps)] 1089 1090 propagations = [HMCPropagation(im_sys_infos[i+1], propagation_params[i], evaluate_heat=False) 1091 for i in range(param_info.intermediate_steps)] 1092 1093 return propagations
1094
1095 - def _calc_works(self, swapcom):
1096 1097 return swapcom.traj12.work, swapcom.traj21.work
1098
1099 1100 -class HMCStepRENSSwapParameterInfo(AbstractRENSSwapParameterInfo):
1101 """ 1102 Holds all required information for performing HMCStepRENS swaps. 1103 1104 @param sampler1: First sampler 1105 @type sampler1: subclass instance of L{AbstractSingleChainMC} 1106 1107 @param sampler2: Second sampler 1108 @type sampler2: subclass instance of L{AbstractSingleChainMC} 1109 1110 @param timestep: integration timestep 1111 @type timestep: float 1112 1113 @param hmc_traj_length: HMC trajectory length 1114 @type hmc_traj_length: int 1115 1116 @param hmc_iterations: number of HMC iterations in the propagation step 1117 @type hmc_iterations: int 1118 1119 @param gradient: gradient governing the equations of motion, function of 1120 position array and switching protocol 1121 @type gradient: callable 1122 1123 @param intermediate_steps: number of steps in the protocol; this is a discrete version 1124 of the switching time in "continuous" RENS implementations 1125 @type intermediate_steps: int 1126 1127 @param protocol: Switching protocol determining the time dependence of the 1128 switching parameter. It is a function f taking the running 1129 time t and the switching time tau to yield a value in [0, 1] 1130 with f(0, tau) = 0 and f(tau, tau) = 1 1131 @type protocol: callable 1132 1133 @param mass_matrix: mass matrix for kinetic energy definition 1134 @type mass_matrix: L{InvertibleMatrix} 1135 1136 @param integrator: Integration scheme to be utilized 1137 @type integrator: l{AbstractIntegrator} 1138 """ 1139
1140 - def __init__(self, sampler1, sampler2, timestep, hmc_traj_length, hmc_iterations, 1141 gradient, intermediate_steps, parametrization=None, protocol=None, 1142 mass_matrix=None, integrator=FastLeapFrog):
1143 1144 super(HMCStepRENSSwapParameterInfo, self).__init__(sampler1, sampler2, protocol) 1145 1146 self._mass_matrix = None 1147 self.mass_matrix = mass_matrix 1148 if self.mass_matrix is None: 1149 d = len(sampler1.state.position) 1150 self.mass_matrix = csb.numeric.InvertibleMatrix(numpy.eye(d), numpy.eye(d)) 1151 1152 self._hmc_traj_length = None 1153 self.hmc_traj_length = hmc_traj_length 1154 self._gradient = None 1155 self.gradient = gradient 1156 self._timestep = None 1157 self.timestep = timestep 1158 self._hmc_iterations = None 1159 self.hmc_iterations = hmc_iterations 1160 self._intermediate_steps = None 1161 self.intermediate_steps = intermediate_steps 1162 self._integrator = None 1163 self.integrator = integrator
1164 1165 @property
1166 - def timestep(self):
1167 """ 1168 Integration timestep. 1169 """ 1170 return self._timestep
1171 @timestep.setter
1172 - def timestep(self, value):
1173 self._timestep = float(value)
1174 1175 @property
1176 - def hmc_traj_length(self):
1177 """ 1178 HMC trajectory length in number of integration steps. 1179 """ 1180 return self._hmc_traj_length
1181 @hmc_traj_length.setter
1182 - def hmc_traj_length(self, value):
1183 self._hmc_traj_length = int(value)
1184 1185 @property
1186 - def gradient(self):
1187 """ 1188 Gradient which governs the equations of motion. 1189 """ 1190 return self._gradient
1191 @gradient.setter
1192 - def gradient(self, value):
1193 self._gradient = value
1194 1195 @property
1196 - def mass_matrix(self):
1197 return self._mass_matrix
1198 @mass_matrix.setter
1199 - def mass_matrix(self, value):
1200 self._mass_matrix = value
1201 1202 @property
1203 - def hmc_iterations(self):
1204 return self._hmc_iterations
1205 @hmc_iterations.setter
1206 - def hmc_iterations(self, value):
1207 self._hmc_iterations = value
1208 1209 @property
1210 - def intermediate_steps(self):
1211 return self._intermediate_steps
1212 @intermediate_steps.setter
1213 - def intermediate_steps(self, value):
1214 self._intermediate_steps = value
1215 1216 @property
1217 - def integrator(self):
1218 return self._integrator
1219 @integrator.setter
1220 - def integrator(self, value):
1221 self._integrator = value
1222
1223 1224 -class AbstractSwapScheme(object):
1225 """ 1226 Provides the interface for classes defining schemes according to which swaps in 1227 Replica Exchange-like simulations are performed. 1228 1229 @param algorithm: Exchange algorithm that performs the swaps 1230 @type algorithm: L{AbstractExchangeMC} 1231 """ 1232 1233 __metaclass__ = ABCMeta 1234
1235 - def __init__(self, algorithm):
1236 1237 self._algorithm = algorithm
1238 1239 @abstractmethod
1240 - def swap_all(self):
1241 """ 1242 Advises the Replica Exchange-like algorithm to perform swaps according to 1243 the schedule defined here. 1244 """ 1245 1246 pass
1247
1248 1249 -class AlternatingAdjacentSwapScheme(AbstractSwapScheme):
1250 """ 1251 Provides a swapping scheme in which tries exchanges between neighbours only 1252 following the scheme 1 <-> 2, 3 <-> 4, ... and after a sampling period 2 <-> 3, 4 <-> 5, ... 1253 1254 @param algorithm: Exchange algorithm that performs the swaps 1255 @type algorithm: L{AbstractExchangeMC} 1256 """ 1257
1258 - def __init__(self, algorithm):
1259 1260 super(AlternatingAdjacentSwapScheme, self).__init__(algorithm) 1261 1262 self._current_swap_list = None 1263 self._swap_list1 = [] 1264 self._swap_list2 = [] 1265 self._create_swap_lists()
1266
1267 - def _create_swap_lists(self):
1268 1269 if len(self._algorithm.param_infos) == 1: 1270 self._swap_list1.append(0) 1271 self._swap_list2.append(0) 1272 else: 1273 i = 0 1274 while i < len(self._algorithm.param_infos): 1275 self._swap_list1.append(i) 1276 i += 2 1277 1278 i = 1 1279 while i < len(self._algorithm.param_infos): 1280 self._swap_list2.append(i) 1281 i += 2 1282 1283 self._current_swap_list = self._swap_list1
1284
1285 - def swap_all(self):
1286 1287 for x in self._current_swap_list: 1288 self._algorithm.swap(x) 1289 1290 if self._current_swap_list == self._swap_list1: 1291 self._current_swap_list = self._swap_list2 1292 else: 1293 self._current_swap_list = self._swap_list1
1294
1295 1296 -class SingleSwapStatistics(object):
1297 """ 1298 Tracks swap statistics of a single sampler pair. 1299 1300 @param param_info: ParameterInfo instance holding swap parameters 1301 @type param_info: L{AbstractSwapParameterInfo} 1302 """ 1303
1304 - def __init__(self, param_info):
1305 self._total_swaps = 0 1306 self._accepted_swaps = 0
1307 1308 @property
1309 - def total_swaps(self):
1310 return self._total_swaps
1311 1312 @property
1313 - def accepted_swaps(self):
1314 return self._accepted_swaps
1315 1316 @property
1317 - def acceptance_rate(self):
1318 """ 1319 Acceptance rate of the sampler pair. 1320 """ 1321 if self.total_swaps > 0: 1322 return float(self.accepted_swaps) / float(self.total_swaps) 1323 else: 1324 return 0.
1325
1326 - def update(self, accepted):
1327 """ 1328 Updates swap statistics. 1329 """ 1330 self._total_swaps += 1 1331 self._accepted_swaps += int(accepted)
1332
1333 1334 -class SwapStatistics(object):
1335 """ 1336 Tracks swap statistics for an AbstractExchangeMC subclass instance. 1337 1338 @param param_infos: list of ParameterInfo instances providing information 1339 needed for performing swaps 1340 @type param_infos: list of L{AbstractSwapParameterInfo} 1341 """ 1342
1343 - def __init__(self, param_infos):
1344 self._stats = [SingleSwapStatistics(x) for x in param_infos]
1345 1346 @property
1347 - def stats(self):
1348 return tuple(self._stats)
1349 1350 @property
1351 - def acceptance_rates(self):
1352 """ 1353 Returns acceptance rates for all swaps. 1354 """ 1355 return [x.acceptance_rate for x in self._stats]
1356
1357 1358 -class InterpolationFactory(object):
1359 """ 1360 Produces interpolations for functions changed during non-equilibrium 1361 trajectories. 1362 1363 @param protocol: protocol to be used to generate non-equilibrium trajectories 1364 @type protocol: function mapping t to [0...1] for fixed tau 1365 @param tau: switching time 1366 @type tau: float 1367 """ 1368
1369 - def __init__(self, protocol, tau):
1370 1371 self._protocol = None 1372 self._tau = None 1373 1374 self.protocol = protocol 1375 self.tau = tau
1376 1377 @property
1378 - def protocol(self):
1379 return self._protocol
1380 @protocol.setter
1381 - def protocol(self, value):
1382 if not hasattr(value, '__call__'): 1383 raise TypeError(value) 1384 self._protocol = value
1385 1386 @property
1387 - def tau(self):
1388 return self._tau
1389 @tau.setter
1390 - def tau(self, value):
1391 self._tau = float(value)
1392
1393 - def build_gradient(self, gradient):
1394 """ 1395 Create a gradient instance with according to given protocol 1396 and switching time. 1397 1398 @param gradient: gradient with G(0) = G_1 and G(1) = G_2 1399 @type gradient: callable 1400 """ 1401 return Gradient(gradient, self._protocol, self._tau)
1402
1403 - def build_temperature(self, temperature):
1404 """ 1405 Create a temperature function according to given protocol and 1406 switching time. 1407 1408 @param temperature: temperature with T(0) = T_1 and T(1) = T_2 1409 @type temperature: callable 1410 """ 1411 return lambda t: temperature(self.protocol(t, self.tau))
1412
1413 1414 -class Gradient(AbstractGradient):
1415
1416 - def __init__(self, gradient, protocol, tau):
1417 1418 self._protocol = protocol 1419 self._gradient = gradient 1420 self._tau = tau
1421
1422 - def evaluate(self, q, t):
1423 return self._gradient(q, self._protocol(t, self._tau))
1424