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.multichain import ThermostattedMDRENSSwapParameterInfo
      16      >>> from csb.statistics.samplers.mc.multichain import ThermostattedMDRENS, AlternatingAdjacentSwapScheme
      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):
     387
     388          return RESwapCommunicator(param_info, Trajectory([param_info.sampler1.state,
     389                                                            param_info.sampler1.state]),
     390                                                Trajectory([param_info.sampler2.state,
     391                                                            param_info.sampler2.state]))
     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      @abstractmethod
     510 -    def _propagator_factory(self, traj_info):
     511          """
     512          Factory method which produces the propagator object used to calculate
     513          the RENS trajectories.
     514
     515          @param traj_info: TrajectoryInfo object holding information neccessary to
     516                            calculate the rens trajectories.
     517          @type traj_info: L{RENSTrajInfo}
     518          @rtype: L{AbstractPropagator}
     519          """
     520          pass
     521
     522 -    def _run_traj_generator(self, traj_info):
     523          """
     524          Run the trajectory generator which generates a trajectory
     525          of a given length between the states of two samplers.
     526
     527          @param traj_info: TrajectoryInfo instance holding information
     528                            needed to generate a nonequilibrium trajectory
     529          @type traj_info: L{RENSTrajInfo}
     530
     531          @rtype: L{Trajectory}
     532          """
     533
     534          init_temperature = self._get_init_temperature(traj_info)
     535
     536          init_state = traj_info.init_state.clone()
     537
     538          if init_state.momentum is None:
     539              init_state = augment_state(init_state,
     540                                         init_temperature,
     541                                         traj_info.param_info.mass_matrix)
     542
     543          gen = self._propagator_factory(traj_info)
     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
     672          proposal1 = swapcom.traj21.final
     673          proposal2 = swapcom.traj12.final
     674
     675          state1 = swapcom.traj12.initial
     676          state2 = swapcom.traj21.initial
     677
     678          if swapcom.param_info.mass_matrix.is_unity_multiple:
     679              inverse_mass_matrix = 1.0 / swapcom.param_info.mass_matrix[0][0]
     680          else:
     681              inverse_mass_matrix = swapcom.param_info.mass_matrix.inverse
     682
     683          E1 = lambda x:-swapcom.sampler1._pdf.log_prob(x)
     684          E2 = lambda x:-swapcom.sampler2._pdf.log_prob(x)
     685          K = lambda x: 0.5 * numpy.dot(x.T, numpy.dot(inverse_mass_matrix, x))
     686
     687          w12 = (K(proposal2.momentum) + E2(proposal2.position)) / T2 - \
     688                (K(state1.momentum) + E1(state1.position)) / T1 - heat12
     689          w21 = (K(proposal1.momentum) + E1(proposal1.position)) / T1 - \
     690                (K(state2.momentum) + E2(state2.position)) / T2 - heat21
     691
     692          return w12, w21
     693
     694
     695 -class MDRENSSwapParameterInfo(RESwapParameterInfo):
     696      """
     697      Holds parameters for a MDRENS swap.
     698
     699      @param sampler1: First sampler
     700      @type sampler1: L{AbstractSingleChainMC}
     701
     702      @param sampler2: Second sampler
     703      @type sampler2: L{AbstractSingleChainMC}
     704
     705      @param timestep: Integration timestep
     706      @type timestep: float
     707
     708      @param traj_length: Trajectory length in number of timesteps
     709      @type traj_length: int
     710
     711      @param gradient: Gradient which determines the dynamics during a trajectory
     712      @type gradient: L{AbstractGradient}
     713
     714      @param protocol: Switching protocol determining the time dependence of the
     715                       switching parameter. It is a function M{f} taking the running
     716                       time t and the switching time tau to yield a value in M{[0, 1]}
     717                       with M{f(0, tau) = 0} and M{f(tau, tau) = 1}. Default is a linear
     718                       protocol, which is being set manually due to an epydoc bug
     719      @type protocol: callable
     720
     721      @param mass_matrix: Mass matrix
     722      @type mass_matrix: n-dimensional matrix of type L{InvertibleMatrix} with n being the dimension
     723                                 of the configuration space, that is, the dimension of
     724                                 the position / momentum vectors
     725      """
     726
     727 -    def __init__(self, sampler1, sampler2, timestep, traj_length, gradient,
     728                   protocol=None, mass_matrix=None):
     729
     730          super(MDRENSSwapParameterInfo, self).__init__(sampler1, sampler2)
     731
     732          self._mass_matrix = mass_matrix
     733          if self.mass_matrix is None:
     734              d = len(sampler1.state.position)
     735              self.mass_matrix = csb.numeric.InvertibleMatrix(numpy.eye(d), numpy.eye(d))
     736
     737          self._traj_length = traj_length
     738          self._gradient = gradient
     739          self._timestep = timestep
     740
     741          ## Can't pass the linear protocol as a default argument because of a reported bug
     742          ## in epydoc parsing which makes it fail building the docs.
     743          self._protocol = None
     744          if protocol is None:
     745              self._protocol = lambda t, tau: t / tau
     746          else:
     747              self._protocol = protocol
     748
     749      @property
     750 -    def timestep(self):
     751          """
     752          Integration timestep.
     753          """
     754          return self._timestep
     755      @timestep.setter
     756 -    def timestep(self, value):
     757          self._timestep = float(value)
     758
     759      @property
     760 -    def traj_length(self):
     761          """
     762          Trajectory length in number of integration steps.
     763          """
     764          return self._traj_length
     765      @traj_length.setter
     766 -    def traj_length(self, value):
     767          self._traj_length = int(value)
     768
     769      @property
     770 -    def gradient(self):
     771          """
     772          Gradient which governs the equations of motion.
     773          """
     774          return self._gradient
     775
     776      @property
     777 -    def mass_matrix(self):
     778          return self._mass_matrix
     779      @mass_matrix.setter
     780 -    def mass_matrix(self, value):
     781          self._mass_matrix = value
     782
     783      @property
     784 -    def protocol(self):
     785          """
     786          Switching protocol determining the time dependence
     787          of the switching parameter.
     788          """
     789          return self._protocol
     790      @protocol.setter
     791 -    def protocol(self, value):
     792          self._protocol = value
     793
     794
     795 -class ThermostattedMDRENS(MDRENS):
     796      """
     797      Replica Exchange with Nonequilibrium Switches (RENS, Ballard & Jarzynski, 2009)
     798      with Andersen-thermostatted Molecular Dynamics (MD) trajectories.
     799
     800      @param samplers: Samplers which sample their
     801                           respective equilibrium distributions
     802      @type samplers: list of L{AbstractSingleChainMC}
     803
     804      @param param_infos: ParameterInfo instance holding
     805                          information required to perform a MDRENS swap
     806      @type param_infos: list of L{ThermostattedMDRENSSwapParameterInfo}
     807
     808      @param integrator: Subclass of L{AbstractIntegrator} to be used to
     809                         calculate the non-equilibrium trajectories
     810      @type integrator: type
     811      """
     812
     813 -    def __init__(self, samplers, param_infos, integrator=csb.numeric.integrators.LeapFrog):
     814
     815          super(ThermostattedMDRENS, self).__init__(samplers, param_infos, integrator)
     816
     817 -    def _propagator_factory(self, traj_info):
     818
     819          protocol = self._setup_protocol(traj_info)
     820          tau = traj_info.param_info.traj_length * traj_info.param_info.timestep
     821          factory = InterpolationFactory(protocol, tau)
     822
     823          grad = factory.build_gradient(traj_info.param_info.gradient)
     824          temp = factory.build_temperature(traj_info.param_info.temperature)
     825
     826          gen = ThermostattedMDPropagator(grad,
     827                                          traj_info.param_info.timestep, temperature=temp,
     828                                          collision_probability=traj_info.param_info.collision_probability,
     829                                          update_interval=traj_info.param_info.collision_interval,
     830                                          mass_matrix=traj_info.param_info.mass_matrix,
     831                                          integrator=self._integrator)
     832
     833          return gen
     834
     835 -class ThermostattedMDRENSSwapParameterInfo(MDRENSSwapParameterInfo):
     836      """
     837      @param sampler1: First sampler
     838      @type sampler1: subclass instance of L{AbstractSingleChainMC}
     839
     840      @param sampler2: Second sampler
     841      @type sampler2: subclass instance of L{AbstractSingleChainMC}
     842
     843      @param timestep: Integration timestep
     844      @type timestep: float
     845
     846      @param traj_length: Trajectory length in number of timesteps
     847      @type traj_length: int
     848
     849      @param gradient: Gradient which determines the dynamics during a trajectory
     850      @type gradient: subclass instance of L{AbstractGradient}
     851
     852      @param mass_matrix: Mass matrix
     853      @type mass_matrix: n-dimensional L{InvertibleMatrix} with n being the dimension
     854                         of the configuration space, that is, the dimension of
     855                         the position / momentum vectors
     856
     857      @param protocol: Switching protocol determining the time dependence of the
     858                       switching parameter. It is a function f taking the running
     859                       time t and the switching time tau to yield a value in [0, 1]
     860                       with f(0, tau) = 0 and f(tau, tau) = 1
     861      @type protocol: callable
     862
     863      @param temperature: Temperature interpolation function.
     864      @type temperature: Real-valued function mapping from [0,1] to R.
     865          T(0) = temperature of the ensemble sampler1 samples from, T(1) = temperature
     866          of the ensemble sampler2 samples from
     867
     868      @param collision_probability: Probability for a collision with the heatbath during one timestep
     869      @type collision_probability: float
     870
     871      @param collision_interval: Interval during which collision may occur with probability
     872          collision_probability
     873      @type collision_interval: int
     874      """
     875
     876 -    def __init__(self, sampler1, sampler2, timestep, traj_length, gradient, mass_matrix=None,
     877                   protocol=None, temperature=lambda l: 1.0,
     878                   collision_probability=0.1, collision_interval=1):
     879
     880          super(ThermostattedMDRENSSwapParameterInfo, self).__init__(sampler1, sampler2, timestep,
     881                                                                     traj_length, gradient,
     882                                                                     mass_matrix=mass_matrix,
     883                                                                     protocol=protocol)
     884
     885          self._collision_probability = None
     886          self._collision_interval = None
     887          self._temperature = temperature
     888          self.collision_probability = collision_probability
     889          self.collision_interval = collision_interval
     890
     891      @property
     892 -    def collision_probability(self):
     893          """
     894          Probability for a collision with the heatbath during one timestep.
     895          """
     896          return self._collision_probability
     897      @collision_probability.setter
     898 -    def collision_probability(self, value):
     899          self._collision_probability = float(value)
     900
     901      @property
     902 -    def collision_interval(self):
     903          """
     904          Interval during which collision may occur with probability
     905          C{collision_probability}.
     906          """
     907          return self._collision_interval
     908      @collision_interval.setter
     909 -    def collision_interval(self, value):
     910          self._collision_interval = int(value)
     911
     912      @property
     913 -    def temperature(self):
     914          return self._temperature
     915
     916
     917 -class AbstractStepRENS(AbstractRENS):
     918      """
     919      Replica Exchange with Nonequilibrium Switches (RENS, Ballard & Jarzynski 2009)
     920      with stepwise trajectories as described in Nilmeier et al., "Nonequilibrium candidate
     921      Monte Carlo is an efficient tool for equilibrium simulation", PNAS 2011.
     922      The switching parameter dependence of the Hamiltonian is a linear interpolation
     923      between the PDFs of the sampler objects,
     924      M{H(S{lambda}) = H_2 * S{lambda} + (1 - S{lambda}) * H_1}.
     925      The perturbation kernel is a thermodynamic perturbation and the propagation is subclass
     926      responsibility.
     927      Note that due to the linear interpolations between the two Hamiltonians, the
     928      log-probability has to be evaluated four times per perturbation step which can be
     929      costly. In this case it is advisable to define the intermediate log probabilities
     930      in _run_traj_generator differently.
     931
     932      @param samplers: Samplers which sample their respective equilibrium distributions
     933      @type samplers: list of L{AbstractSingleChainMC}
     934
     935      @param param_infos: ParameterInfo instances holding
     936                          information required to perform a HMCStepRENS swaps
     937      @type param_infos: list of L{AbstractSwapParameterInfo}
     938      """
     939
     940      __metaclass__ = ABCMeta
     941
     942 -    def __init__(self, samplers, param_infos):
     943
     944          super(AbstractStepRENS, self).__init__(samplers, param_infos)
     945
     946          self._evaluate_im_works = True
     947
     948      @abstractmethod
     949 -    def _setup_propagations(self, im_sys_infos, param_info):
     950          """
     951          Set up the propagation steps using the information about the current system
     952          setup and parameters from the SwapParameterInfo object.
     953
     954          @param im_sys_infos: Information about the intermediate system setups
     955          @type im_sys_infos: List of L{AbstractSystemInfo}
     956
     957          @param param_info: SwapParameterInfo object containing parameters for the
     958                             propagations like timesteps, trajectory lengths etc.
     959          @type param_info: L{AbstractSwapParameterInfo}
     960          """
     961
     962          pass
     963
     964      @abstractmethod
     965 -    def _add_gradients(self, im_sys_infos, param_info, t_prot):
     966          """
     967          If needed, set im_sys_infos.hamiltonian.gradient.
     968
     969          @param im_sys_infos: Information about the intermediate system setups
     970          @type im_sys_infos: List of L{AbstractSystemInfo}
     971
     972          @param param_info: SwapParameterInfo object containing parameters for the
     973                             propagations like timesteps, trajectory lengths etc.
     974          @type param_info: L{AbstractSwapParameterInfo}
     975
     976          @param t_prot: Switching protocol defining the time dependence of the switching
     977                         parameter.
     978          @type t_prot: callable
     979          """
     980
     981          pass
     982
     983 -    def _setup_stepwise_protocol(self, traj_info):
     984          """
     985          Sets up the stepwise protocol consisting of perturbation and relaxation steps.
     986
     987          @param traj_info: TrajectoryInfo instance holding information
     988                            needed to generate a nonequilibrium trajectory
     989          @type traj_info: L{RENSTrajInfo}
     990
     991          @rtype: L{Protocol}
     992          """
     993
     994          pdf1 = traj_info.param_info.sampler1._pdf
     995          pdf2 = traj_info.param_info.sampler2._pdf
     996          T1 = traj_info.param_info.sampler1.temperature
     997          T2 = traj_info.param_info.sampler2.temperature
     998          traj_length = traj_info.param_info.intermediate_steps
     999          prot = self._setup_protocol(traj_info)
    1000          t_prot = lambda i: prot(float(i), float(traj_length))
    1001
    1002          im_log_probs = [lambda x, i=i: pdf2.log_prob(x) * t_prot(i) + \
    1003                                         (1 - t_prot(i)) * pdf1.log_prob(x)
    1004                          for i in range(traj_length + 1)]
    1005
    1006          im_temperatures = [T2 * t_prot(i) + (1 - t_prot(i)) * T1
    1007                             for i in range(traj_length + 1)]
    1008          im_reduced_hamiltonians = [ReducedHamiltonian(im_log_probs[i],
    1009                                                        temperature=im_temperatures[i])
    1010                                     for i in range(traj_length + 1)]
    1011          im_sys_infos = [HamiltonianSysInfo(im_reduced_hamiltonians[i])
    1012                          for i in range(traj_length + 1)]
    1013          perturbations = [ReducedHamiltonianPerturbation(im_sys_infos[i], im_sys_infos[i+1])
    1014                          for i in range(traj_length)]
    1015          if self._evaluate_im_works == False:
    1016              for p in perturbations:
    1017                  p.evaluate_work = False
    1018          im_sys_infos = self._add_gradients(im_sys_infos, traj_info.param_info, t_prot)
    1019          propagations = self._setup_propagations(im_sys_infos, traj_info.param_info)
    1020
    1021          steps = [Step(perturbations[i], propagations[i]) for i in range(traj_length)]
    1022
    1023          return Protocol(steps)
    1024
    1025 -    def _propagator_factory(self, traj_info):
    1026
    1027          protocol = self._setup_stepwise_protocol(traj_info)
    1028          gen = NonequilibriumStepPropagator(protocol)
    1029
    1030          return gen
    1031
    1032 -    def _run_traj_generator(self, traj_info):
    1033
    1034          init_temperature = self._get_init_temperature(traj_info)
    1035
    1036          gen = self._propagator_factory(traj_info)
    1037
    1038          traj = gen.generate(traj_info.init_state)
    1039
    1040          return NonequilibriumTrajectory([traj_info.init_state, traj.final], jacobian=1.0,
    1041                                           heat=traj.heat, work=traj.work, deltaH=traj.deltaH)
    1042
    1043
    1044 -class HMCStepRENS(AbstractStepRENS):
    1045      """
    1046      Replica Exchange with Nonequilibrium Switches (RENS, Ballard & Jarzynski 2009)
    1047      with stepwise trajectories as described in Nilmeier et al., "Nonequilibrium candidate
    1048      Monte Carlo is an efficient tool for equilibrium simulation", PNAS 2011.
    1049      The switching parameter dependence of the Hamiltonian is a linear interpolation
    1050      between the PDFs of the sampler objects,
    1051      M{H(S{lambda}) = H_2 * S{lambda} + (1 - S{lambda}) * H_1}.
    1052      The perturbation kernel is a thermodynamic perturbation and the propagation is done using HMC.
    1053
    1054      Note that due to the linear interpolations between the two Hamiltonians, the
    1055      log-probability and its gradient has to be evaluated four times per perturbation step which
    1056      can be costly. In this case it is advisable to define the intermediate log probabilities
    1057      in _run_traj_generator differently.
    1058
    1059      @param samplers: Samplers which sample their respective equilibrium distributions
    1060      @type samplers: list of L{AbstractSingleChainMC}
    1061
    1062      @param param_infos: ParameterInfo instances holding
    1063                          information required to perform a HMCStepRENS swaps
    1064      @type param_infos: list of L{HMCStepRENSSwapParameterInfo}
    1065      """
    1066
    1067 -    def __init__(self, samplers, param_infos):
    1068
    1069          super(HMCStepRENS, self).__init__(samplers, param_infos)
    1070
    1071      @staticmethod
    1072 -    def _add_gradients(im_sys_infos, param_info, t_prot):
    1073
    1074          im_gradients = [lambda x, t, i=i: param_info.gradient(x, t_prot(i))
    1075                          for i in range(param_info.intermediate_steps + 1)]
    1076
    1077          for i, s in enumerate(im_sys_infos):
    1078              s.hamiltonian.gradient = im_gradients[i]
    1079
    1080          return im_sys_infos
    1081
    1082      @staticmethod
    1083 -    def _setup_propagations(im_sys_infos, param_info):
    1084
    1085          propagation_params = [HMCPropagationParam(param_info.timestep,
    1086                                                    param_info.hmc_traj_length,
    1087                                                    im_sys_infos[i+1].hamiltonian.gradient,
    1088                                                    param_info.hmc_iterations,
    1089                                                    mass_matrix=param_info.mass_matrix,
    1090                                                    integrator=param_info.integrator)
    1091                                for i in range(param_info.intermediate_steps)]
    1092
    1093          propagations = [HMCPropagation(im_sys_infos[i+1], propagation_params[i], evaluate_heat=False)
    1094                          for i in range(param_info.intermediate_steps)]
    1095
    1096          return propagations
    1097
    1098 -    def _calc_works(self, swapcom):
    1099
    1100          return swapcom.traj12.work, swapcom.traj21.work
    1101
    1102
    1103 -class HMCStepRENSSwapParameterInfo(AbstractRENSSwapParameterInfo):
    1104      """
    1105      Holds all required information for performing HMCStepRENS swaps.
    1106
    1107      @param sampler1: First sampler
    1108      @type sampler1: subclass instance of L{AbstractSingleChainMC}
    1109
    1110      @param sampler2: Second sampler
    1111      @type sampler2: subclass instance of L{AbstractSingleChainMC}
    1112
    1113      @param timestep: integration timestep
    1114      @type timestep: float
    1115
    1116      @param hmc_traj_length: HMC trajectory length
    1117      @type hmc_traj_length: int
    1118
    1119      @param hmc_iterations: number of HMC iterations in the propagation step
    1120      @type hmc_iterations: int
    1121
    1122      @param gradient: gradient governing the equations of motion, function of
    1123                       position array and switching protocol
    1124      @type gradient: callable
    1125
    1126      @param intermediate_steps: number of steps in the protocol; this is a discrete version
    1127                                 of the switching time in "continuous" RENS implementations
    1128      @type intermediate_steps: int
    1129
    1130      @param protocol: Switching protocol determining the time dependence of the
    1131                       switching parameter. It is a function f taking the running
    1132                       time t and the switching time tau to yield a value in [0, 1]
    1133                       with f(0, tau) = 0 and f(tau, tau) = 1
    1134      @type protocol: callable
    1135
    1136      @param mass_matrix: mass matrix for kinetic energy definition
    1137      @type mass_matrix: L{InvertibleMatrix}
    1138
    1139      @param integrator: Integration scheme to be utilized
    1140      @type integrator: l{AbstractIntegrator}
    1141      """
    1142
    1143 -    def __init__(self, sampler1, sampler2, timestep, hmc_traj_length, hmc_iterations,
    1144                   gradient, intermediate_steps, parametrization=None, protocol=None,
    1145                   mass_matrix=None, integrator=FastLeapFrog):
    1146
    1147          super(HMCStepRENSSwapParameterInfo, self).__init__(sampler1, sampler2, protocol)
    1148
    1149          self._mass_matrix = None
    1150          self.mass_matrix = mass_matrix
    1151          if self.mass_matrix is None:
    1152              d = len(sampler1.state.position)
    1153              self.mass_matrix = csb.numeric.InvertibleMatrix(numpy.eye(d), numpy.eye(d))
    1154
    1155          self._hmc_traj_length = None
    1156          self.hmc_traj_length = hmc_traj_length
    1157          self._gradient = None
    1158          self.gradient = gradient
    1159          self._timestep = None
    1160          self.timestep = timestep
    1161          self._hmc_iterations = None
    1162          self.hmc_iterations = hmc_iterations
    1163          self._intermediate_steps = None
    1164          self.intermediate_steps = intermediate_steps
    1165          self._integrator = None
    1166          self.integrator = integrator
    1167
    1168      @property
    1169 -    def timestep(self):
    1170          """
    1171          Integration timestep.
    1172          """
    1173          return self._timestep
    1174      @timestep.setter
    1175 -    def timestep(self, value):
    1176          self._timestep = float(value)
    1177
    1178      @property
    1179 -    def hmc_traj_length(self):
    1180          """
    1181          HMC trajectory length in number of integration steps.
    1182          """
    1183          return self._hmc_traj_length
    1184      @hmc_traj_length.setter
    1185 -    def hmc_traj_length(self, value):
    1186          self._hmc_traj_length = int(value)
    1187
    1188      @property
    1189 -    def gradient(self):
    1190          """
    1191          Gradient which governs the equations of motion.
    1192          """
    1193          return self._gradient
    1194      @gradient.setter
    1195 -    def gradient(self, value):
    1196          self._gradient = value
    1197
    1198      @property
    1199 -    def mass_matrix(self):
    1200          return self._mass_matrix
    1201      @mass_matrix.setter
    1202 -    def mass_matrix(self, value):
    1203          self._mass_matrix = value
    1204
    1205      @property
    1206 -    def hmc_iterations(self):
    1207          return self._hmc_iterations
    1208      @hmc_iterations.setter
    1209 -    def hmc_iterations(self, value):
    1210          self._hmc_iterations = value
    1211
    1212      @property
    1213 -    def intermediate_steps(self):
    1214          return self._intermediate_steps
    1215      @intermediate_steps.setter
    1216 -    def intermediate_steps(self, value):
    1217          self._intermediate_steps = value
    1218
    1219      @property
    1220 -    def integrator(self):
    1221          return self._integrator
    1222      @integrator.setter
    1223 -    def integrator(self, value):
    1224          self._integrator = value
    1225
    1226
    1227 -class AbstractSwapScheme(object):
    1228      """
    1229      Provides the interface for classes defining schemes according to which swaps in
    1230      Replica Exchange-like simulations are performed.
    1231
    1232      @param algorithm: Exchange algorithm that performs the swaps
    1233      @type algorithm: L{AbstractExchangeMC}
    1234      """
    1235
    1236      __metaclass__ = ABCMeta
    1237
    1238 -    def __init__(self, algorithm):
    1239
    1240          self._algorithm = algorithm
    1241
    1242      @abstractmethod
    1243 -    def swap_all(self):
    1244          """
    1245          Advises the Replica Exchange-like algorithm to perform swaps according to
    1246          the schedule defined here.
    1247          """
    1248
    1249          pass
    1250
    1251
    1252 -class AlternatingAdjacentSwapScheme(AbstractSwapScheme):
    1253      """
    1254      Provides a swapping scheme in which tries exchanges between neighbours only
    1255      following the scheme 1 <-> 2, 3 <-> 4, ... and after a sampling period 2 <-> 3, 4 <-> 5, ...
    1256
    1257      @param algorithm: Exchange algorithm that performs the swaps
    1258      @type algorithm: L{AbstractExchangeMC}
    1259      """
    1260
    1261 -    def __init__(self, algorithm):
    1262
    1263          super(AlternatingAdjacentSwapScheme, self).__init__(algorithm)
    1264
    1265          self._current_swap_list = None
    1266          self._swap_list1 = []
    1267          self._swap_list2 = []
    1268          self._create_swap_lists()
    1269
    1270 -    def _create_swap_lists(self):
    1271
    1272          if len(self._algorithm.param_infos) == 1:
    1273              self._swap_list1.append(0)
    1274              self._swap_list2.append(0)
    1275          else:
    1276              i = 0
    1277              while i < len(self._algorithm.param_infos):
    1278                  self._swap_list1.append(i)
    1279                  i += 2
    1280
    1281              i = 1
    1282              while i < len(self._algorithm.param_infos):
    1283                  self._swap_list2.append(i)
    1284                  i += 2
    1285
    1286          self._current_swap_list = self._swap_list1
    1287
    1288 -    def swap_all(self):
    1289
    1290          for x in self._current_swap_list:
    1291              self._algorithm.swap(x)
    1292
    1293          if self._current_swap_list == self._swap_list1:
    1294              self._current_swap_list = self._swap_list2
    1295          else:
    1296              self._current_swap_list = self._swap_list1
    1297
    1298
    1299 -class SingleSwapStatistics(object):
    1300      """
    1301      Tracks swap statistics of a single sampler pair.
    1302
    1303      @param param_info: ParameterInfo instance holding swap parameters
    1304      @type param_info: L{AbstractSwapParameterInfo}
    1305      """
    1306
    1307 -    def __init__(self, param_info):
    1308          self._total_swaps = 0
    1309          self._accepted_swaps = 0
    1310
    1311      @property
    1312 -    def total_swaps(self):
    1313          return self._total_swaps
    1314
    1315      @property
    1316 -    def accepted_swaps(self):
    1317          return self._accepted_swaps
    1318
    1319      @property
    1320 -    def acceptance_rate(self):
    1321          """
    1322          Acceptance rate of the sampler pair.
    1323          """
    1324          if self.total_swaps > 0:
    1325              return float(self.accepted_swaps) / float(self.total_swaps)
    1326          else:
    1327              return 0.
    1328
    1329 -    def update(self, accepted):
    1330          """
    1331          Updates swap statistics.
    1332          """
    1333          self._total_swaps += 1
    1334          self._accepted_swaps += int(accepted)
    1335
    1336
    1337 -class SwapStatistics(object):
    1338      """
    1339      Tracks swap statistics for an AbstractExchangeMC subclass instance.
    1340
    1341      @param param_infos: list of ParameterInfo instances providing information
    1342                          needed for performing swaps
    1343      @type param_infos: list of L{AbstractSwapParameterInfo}
    1344      """
    1345
    1346 -    def __init__(self, param_infos):
    1347          self._stats = [SingleSwapStatistics(x) for x in param_infos]
    1348
    1349      @property
    1350 -    def stats(self):
    1351          return tuple(self._stats)
    1352
    1353      @property
    1354 -    def acceptance_rates(self):
    1355          """
    1356          Returns acceptance rates for all swaps.
    1357          """
    1358          return [x.acceptance_rate for x in self._stats]
    1359
    1360
    1361 -class InterpolationFactory(object):
    1362      """
    1363      Produces interpolations for functions changed during non-equilibrium
    1364      trajectories.
    1365
    1366      @param protocol: protocol to be used to generate non-equilibrium trajectories
    1367      @type protocol: function mapping t to [0...1] for fixed tau
    1368      @param tau: switching time
    1369      @type tau: float
    1370      """
    1371
    1372 -    def __init__(self, protocol, tau):
    1373
    1374          self._protocol = None
    1375          self._tau = None
    1376
    1377          self.protocol = protocol
    1378          self.tau = tau
    1379
    1380      @property
    1381 -    def protocol(self):
    1382          return self._protocol
    1383      @protocol.setter
    1384 -    def protocol(self, value):
    1385          if not hasattr(value, '__call__'):
    1386              raise TypeError(value)
    1387          self._protocol = value
    1388
    1389      @property
    1390 -    def tau(self):
    1391          return self._tau
    1392      @tau.setter
    1393 -    def tau(self, value):
    1394          self._tau = float(value)
    1395
    1396 -    def build_gradient(self, gradient):
    1397          """
    1398          Create a gradient instance with according to given protocol
    1399          and switching time.
    1400
    1401          @param gradient: gradient with G(0) = G_1 and G(1) = G_2
    1402          @type gradient: callable
    1403          """
    1404          return Gradient(gradient, self._protocol, self._tau)
    1405
    1406 -    def build_temperature(self, temperature):
    1407          """
    1408          Create a temperature function according to given protocol and
    1409          switching time.
    1410
    1411          @param temperature: temperature with T(0) = T_1 and T(1) = T_2
    1412          @type temperature: callable
    1413          """
    1414          return lambda t: temperature(self.protocol(t, self.tau))
    1415
    1416
    1417 -class Gradient(AbstractGradient):
    1418
    1419 -    def __init__(self, gradient, protocol, tau):
    1420
    1421          self._protocol = protocol
    1422          self._gradient = gradient
    1423          self._tau = tau
    1424
    1425 -    def evaluate(self, q, t):
    1426          return self._gradient(q, self._protocol(t, self._tau))
    1427
    1428
    1429 -class ReplicaHistory(object):
    1430      '''
    1431      Replica history object, works with both RE and RENS for
    1432      the AlternatingAdjacentSwapScheme.
    1433
    1434      @param samples: list holding ensemble states
    1435      @type samples: list
    1436
    1437      @param swap_interval: interval with which swaps were attempted, e.g.,
    1438                            5 means that every 5th regular MC step is replaced
    1439                            by a swap
    1440      @type swap_interval: int
    1441
    1442      @param first_swap: sample index of the first sample generated by a swap attempt.
    1443                         If None, the first RE sampled is assumed to have sample index
    1444                         swap_interval. If specified, it has to be greater than zero
    1445      @type first_swap: int
    1446      '''
    1447
    1448 -    def __init__(self, samples, swap_interval, first_swap=None):
    1449          self.samples = samples
    1450          self.swap_interval = swap_interval
    1451          if first_swap == None:
    1452              self.first_swap = swap_interval - 1
    1453          elif first_swap > 0:
    1454              self.first_swap = first_swap - 1
    1455          else:
    1456              raise(ValueError("Sample index of first swap has to be greater than zero!"))
    1457          self.n_replicas = len(samples[0])
    1458
    1459      @staticmethod
    1460 -    def _change_direction(x):
    1461          if x == 1:
    1462              return -1
    1463          if x == -1:
    1464              return 1
    1465
    1466 -    def calculate_history(self, start_ensemble):
    1467          '''
    1468          Calculates the replica history of the first state of ensemble #start_ensemble.
    1469
    1470          @param start_ensemble: index of the ensemble to start at, zero-indexed
    1471          @type start_ensemble: int
    1472
    1473          @return: replica history as a list of ensemble indices
    1474          @rtype: list of ints
    1475          '''
    1476
    1477          sample_counter = 0
    1478
    1479          # determine the direction (up = 1, down = -1) in the "temperature ladder" of
    1480          # the first swap attempt. Remember: first swap series is always 0 <-> 1, 2 <-> 3, ...
    1481          if start_ensemble % 2 == 0:
    1482              direction = +1
    1483          else:
    1484              direction = -1
    1485
    1486          # if number of replicas is not even and the start ensemble is the highest-temperature-
    1487          # ensemble, the first swap will be attempted "downwards"
    1488          if start_ensemble % 2 == 0 and start_ensemble == self.n_replicas - 1:
    1489              direction = -1
    1490
    1491          # will store the indices of the ensembles the state will visit in chronological order
    1492          history = []
    1493
    1494          # the ensemble the state is currently in
    1495          ens = start_ensemble
    1496
    1497          while sample_counter < len(self.samples):
    1498              if self.n_replicas == 2:
    1499                  if (sample_counter - self.first_swap - 1) % self.swap_interval == 0 and \
    1500                     sample_counter >= self.first_swap:
    1501                      ## swap attempt: determine whether it was successfull or not
    1502                      # state after swap attempt
    1503                      current_sample = self.samples[sample_counter][ens]
    1504
    1505                      # state before swap attempt
    1506                      previous_sample = self.samples[sample_counter - 1][history[-1]]
    1507
    1508                      # swap was accepted when position of the current state doesn't equal
    1509                      # the position of the state before the swap attempt, that is, the last
    1510                      # state in the history
    1511                      swap_accepted = not numpy.all(current_sample.position ==
    1512                                                    previous_sample.position)
    1513
    1514                      if swap_accepted:
    1515                          if ens == 0:
    1516                              ens = 1
    1517                          else:
    1518                              ens = 0
    1519                      history.append(ens)
    1520                  else:
    1521                      history.append(ens)
    1522
    1523              else:
    1524
    1525                  if (sample_counter - self.first_swap - 1) % self.swap_interval == 0 and \
    1526                     sample_counter >= self.first_swap:
    1527                      # state after swap attempt
    1528                      current_sample = self.samples[sample_counter][ens]
    1529
    1530                      # state before swap attempt
    1531                      previous_sample = self.samples[sample_counter - 1][ens]
    1532
    1533                      # swap was accepted when position of the current state doesn't equal
    1534                      # the position of the state before the swap attempt, that is, the last
    1535                      # state in the history
    1536                      swap_accepted = not numpy.all(current_sample.position == previous_sample.position)
    1537
    1538                      if swap_accepted:
    1539                          ens += direction
    1540                      else:
    1541                          if ens == self.n_replicas - 1:
    1542                              # if at the top of the ladder, go downwards again
    1543                              direction = -1
    1544                          elif ens == 0:
    1545                              # if at the bottom of the ladder, go upwards
    1546                              direction = +1
    1547                          else:
    1548                              # in between, reverse the direction of the trajectory
    1549                              # in temperature space
    1550                              direction = self._change_direction(direction)
    1551
    1552                  history.append(ens)
    1553
    1554              sample_counter += 1
    1555
    1556          return history
    1557
    1558 -    def calculate_projected_trajectories(self, ensemble):
    1559          '''
    1560          Calculates sequentially correlated trajectories projected on a specific ensemble.
    1561
    1562          @param ensemble: ensemble index of ensemble of interest, zero-indexed
    1563          @type ensemble: int
    1564
    1565          @return: list of Trajectory objects containg sequentially correlated trajectories
    1566          @rtype: list of L{Trajectory} objects.
    1567          '''
    1568
    1569          trajectories = []
    1570
    1571          for i in range(self.n_replicas):
    1572              history = self.calculate_history(i)
    1573              traj = [x[ensemble] for k, x in enumerate(self.samples) if history[k] == ensemble]
    1574              trajectories.append(Trajectory(traj))
    1575
    1576          return trajectories
    1577
    1578 -    def calculate_trajectories(self):
    1579          '''
    1580          Calculates sequentially correlated trajectories.
    1581
    1582          @return: list of Trajectory objects containg sequentially correlated trajectories
    1583          @rtype: list of L{Trajectory} objects.
    1584          '''
    1585
    1586          trajectories = []
    1587
    1588          for i in range(self.n_replicas):
    1589              history = self.calculate_history(i)
    1590              traj = [x[history[k]] for k, x in enumerate(self.samples)]
    1591              trajectories.append(Trajectory(traj))
    1592
    1593          return trajectories
    1594