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
     545          traj = gen.generate(init_state, int(traj_info.param_info.traj_length))
     546
     547          return traj
     548
     549
     550 -class AbstractRENSSwapParameterInfo(RESwapParameterInfo):
     551      """
     552      Holds parameters for a RENS swap.
     553      """
     554
     555      __metaclass__ = ABCMeta
     556
     557 -    def __init__(self, sampler1, sampler2, protocol):
     558
     559          super(AbstractRENSSwapParameterInfo, self).__init__(sampler1, sampler2)
     560
     561          ## Can't pass the linear protocol as a default argument because of a reported bug
     562          ## in epydoc parsing which makes it fail building the docs.
     563          self._protocol = None
     564          if protocol is None:
     565              self._protocol = lambda t, tau: t / tau
     566          else:
     567              self._protocol = protocol
     568
     569      @property
     570 -    def protocol(self):
     571          """
     572          Switching protocol determining the time dependence
     573          of the switching parameter.
     574          """
     575          return self._protocol
     576      @protocol.setter
     577 -    def protocol(self, value):
     578          self._protocol = value
     579
     580
     581 -class RENSSwapCommunicator(AbstractSwapCommunicator):
     582      """
     583      Holds all the information which needs to be communicated between distinct
     584      RENS swap substeps.
     585
     586      See L{AbstractSwapCommunicator} for constructor signature.
     587      """
     588
     589      pass
     590
     591
     592 -class RENSTrajInfo(object):
     593      """
     594      Holds information necessary for calculating a RENS trajectory.
     595
     596      @param param_info: ParameterInfo instance holding swap parameters
     597      @type param_info: L{AbstractSwapParameterInfo}
     598
     599      @param init_state: state from which the trajectory is supposed to start
     600      @type init_state: L{State}
     601
     602      @param direction: Either "fw" or "bw", indicating a forward or backward
     603                        trajectory. This is neccessary to pick the protocol or
     604                        the reversed protocol, respectively.
     605      @type direction: string, either "fw" or "bw"
     606      """
     607
     608 -    def __init__(self, param_info, init_state, direction):
     609
     610          self._param_info = param_info
     611          self._init_state = init_state
     612          self._direction = direction
     613
     614      @property
     615 -    def param_info(self):
     616          return self._param_info
     617
     618      @property
     619 -    def init_state(self):
     620          return self._init_state
     621
     622      @property
     623 -    def direction(self):
     624          return self._direction
     625
     626
     627 -class MDRENS(AbstractRENS):
     628      """
     629      Replica Exchange with Nonequilibrium Switches (RENS, Ballard & Jarzynski 2009)
     630      with Molecular Dynamics (MD) trajectories.
     631
     632      @param samplers: Samplers which sample their
     633                           respective equilibrium distributions
     634      @type samplers: list of L{AbstractSingleChainMC}
     635
     636      @param param_infos: ParameterInfo instance holding
     637                          information required to perform a MDRENS swap
     638      @type param_infos: list of L{MDRENSSwapParameterInfo}
     639
     640      @param integrator: Subclass of L{AbstractIntegrator} to be used to
     641                         calculate the non-equilibrium trajectories
     642      @type integrator: type
     643      """
     644
     645 -    def __init__(self, samplers, param_infos,
     646                   integrator=csb.numeric.integrators.FastLeapFrog):
     647
     648          super(MDRENS, self).__init__(samplers, param_infos)
     649
     650          self._integrator = integrator
     651
     652 -    def _propagator_factory(self, traj_info):
     653
     654          protocol = self._setup_protocol(traj_info)
     655          tau = traj_info.param_info.traj_length * traj_info.param_info.timestep
     656          factory = InterpolationFactory(protocol, tau)
     657          gen = MDPropagator(factory.build_gradient(traj_info.param_info.gradient),
     658                             traj_info.param_info.timestep,
     659                             mass_matrix=traj_info.param_info.mass_matrix,
     660                             integrator=self._integrator)
     661
     662          return gen
     663
     664 -    def _calc_works(self, swapcom):
     665
     666          T1 = swapcom.param_info.sampler1.temperature
     667          T2 = swapcom.param_info.sampler2.temperature
     668
     669          heat12 = swapcom.traj12.heat
     670          heat21 = swapcom.traj21.heat
     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 -    def _add_gradients(self, im_sys_infos, param_info, t_prot):
    1072
    1073          im_gradients = [lambda x, t, i=i: param_info.gradient(x, t_prot(i))
    1074                          for i in range(param_info.intermediate_steps + 1)]
    1075
    1076          for i, s in enumerate(im_sys_infos):
    1077              s.hamiltonian.gradient = im_gradients[i]
    1078
    1079          return im_sys_infos
    1080
    1081 -    def _setup_propagations(self, im_sys_infos, param_info):
    1082
    1083          propagation_params = [HMCPropagationParam(param_info.timestep,
    1084                                                    param_info.hmc_traj_length,
    1085                                                    im_sys_infos[i+1].hamiltonian.gradient,
    1086                                                    param_info.hmc_iterations,
    1087                                                    mass_matrix=param_info.mass_matrix,
    1088                                                    integrator=param_info.integrator)
    1089                                for i in range(param_info.intermediate_steps)]
    1090
    1091          propagations = [HMCPropagation(im_sys_infos[i+1], propagation_params[i], evaluate_heat=False)
    1092                          for i in range(param_info.intermediate_steps)]
    1093
    1094          return propagations
    1095
    1096 -    def _calc_works(self, swapcom):
    1097
    1098          return swapcom.traj12.work, swapcom.traj21.work
    1099
    1100
    1101 -class HMCStepRENSSwapParameterInfo(AbstractRENSSwapParameterInfo):
    1102      """
    1103      Holds all required information for performing HMCStepRENS swaps.
    1104
    1105      @param sampler1: First sampler
    1106      @type sampler1: subclass instance of L{AbstractSingleChainMC}
    1107
    1108      @param sampler2: Second sampler
    1109      @type sampler2: subclass instance of L{AbstractSingleChainMC}
    1110
    1111      @param timestep: integration timestep
    1112      @type timestep: float
    1113
    1114      @param hmc_traj_length: HMC trajectory length
    1115      @type hmc_traj_length: int
    1116
    1117      @param hmc_iterations: number of HMC iterations in the propagation step
    1118      @type hmc_iterations: int
    1119
    1120      @param gradient: gradient governing the equations of motion, function of
    1121                       position array and switching protocol
    1122      @type gradient: callable
    1123
    1124      @param intermediate_steps: number of steps in the protocol; this is a discrete version
    1125                                 of the switching time in "continuous" RENS implementations
    1126      @type intermediate_steps: int
    1127
    1128      @param protocol: Switching protocol determining the time dependence of the
    1129                       switching parameter. It is a function f taking the running
    1130                       time t and the switching time tau to yield a value in [0, 1]
    1131                       with f(0, tau) = 0 and f(tau, tau) = 1
    1132      @type protocol: callable
    1133
    1134      @param mass_matrix: mass matrix for kinetic energy definition
    1135      @type mass_matrix: L{InvertibleMatrix}
    1136
    1137      @param integrator: Integration scheme to be utilized
    1138      @type integrator: l{AbstractIntegrator}
    1139      """
    1140
    1141 -    def __init__(self, sampler1, sampler2, timestep, hmc_traj_length, hmc_iterations,
    1142                   gradient, intermediate_steps, parametrization=None, protocol=None,
    1143                   mass_matrix=None, integrator=FastLeapFrog):
    1144
    1145          super(HMCStepRENSSwapParameterInfo, self).__init__(sampler1, sampler2, protocol)
    1146
    1147          self._mass_matrix = None
    1148          self.mass_matrix = mass_matrix
    1149          if self.mass_matrix is None:
    1150              d = len(sampler1.state.position)
    1151              self.mass_matrix = csb.numeric.InvertibleMatrix(numpy.eye(d), numpy.eye(d))
    1152
    1153          self._hmc_traj_length = None
    1154          self.hmc_traj_length = hmc_traj_length
    1155          self._gradient = None
    1156          self.gradient = gradient
    1157          self._timestep = None
    1158          self.timestep = timestep
    1159          self._hmc_iterations = None
    1160          self.hmc_iterations = hmc_iterations
    1161          self._intermediate_steps = None
    1162          self.intermediate_steps = intermediate_steps
    1163          self._integrator = None
    1164          self.integrator = integrator
    1165
    1166      @property
    1167 -    def timestep(self):
    1168          """
    1169          Integration timestep.
    1170          """
    1171          return self._timestep
    1172      @timestep.setter
    1173 -    def timestep(self, value):
    1174          self._timestep = float(value)
    1175
    1176      @property
    1177 -    def hmc_traj_length(self):
    1178          """
    1179          HMC trajectory length in number of integration steps.
    1180          """
    1181          return self._hmc_traj_length
    1182      @hmc_traj_length.setter
    1183 -    def hmc_traj_length(self, value):
    1184          self._hmc_traj_length = int(value)
    1185
    1186      @property
    1187 -    def gradient(self):
    1188          """
    1189          Gradient which governs the equations of motion.
    1190          """
    1191          return self._gradient
    1192      @gradient.setter
    1193 -    def gradient(self, value):
    1194          self._gradient = value
    1195
    1196      @property
    1197 -    def mass_matrix(self):
    1198          return self._mass_matrix
    1199      @mass_matrix.setter
    1200 -    def mass_matrix(self, value):
    1201          self._mass_matrix = value
    1202
    1203      @property
    1204 -    def hmc_iterations(self):
    1205          return self._hmc_iterations
    1206      @hmc_iterations.setter
    1207 -    def hmc_iterations(self, value):
    1208          self._hmc_iterations = value
    1209
    1210      @property
    1211 -    def intermediate_steps(self):
    1212          return self._intermediate_steps
    1213      @intermediate_steps.setter
    1214 -    def intermediate_steps(self, value):
    1215          self._intermediate_steps = value
    1216
    1217      @property
    1218 -    def integrator(self):
    1219          return self._integrator
    1220      @integrator.setter
    1221 -    def integrator(self, value):
    1222          self._integrator = value
    1223
    1224
    1225 -class AbstractSwapScheme(object):
    1226      """
    1227      Provides the interface for classes defining schemes according to which swaps in
    1228      Replica Exchange-like simulations are performed.
    1229
    1230      @param algorithm: Exchange algorithm that performs the swaps
    1231      @type algorithm: L{AbstractExchangeMC}
    1232      """
    1233
    1234      __metaclass__ = ABCMeta
    1235
    1236 -    def __init__(self, algorithm):
    1237
    1238          self._algorithm = algorithm
    1239
    1240      @abstractmethod
    1241 -    def swap_all(self):
    1242          """
    1243          Advises the Replica Exchange-like algorithm to perform swaps according to
    1244          the schedule defined here.
    1245          """
    1246
    1247          pass
    1248
    1249
    1250 -class AlternatingAdjacentSwapScheme(AbstractSwapScheme):
    1251      """
    1252      Provides a swapping scheme in which tries exchanges between neighbours only
    1253      following the scheme 1 <-> 2, 3 <-> 4, ... and after a sampling period 2 <-> 3, 4 <-> 5, ...
    1254
    1255      @param algorithm: Exchange algorithm that performs the swaps
    1256      @type algorithm: L{AbstractExchangeMC}
    1257      """
    1258
    1259 -    def __init__(self, algorithm):
    1260
    1261          super(AlternatingAdjacentSwapScheme, self).__init__(algorithm)
    1262
    1263          self._current_swap_list = None
    1264          self._swap_list1 = []
    1265          self._swap_list2 = []
    1266          self._create_swap_lists()
    1267
    1268 -    def _create_swap_lists(self):
    1269
    1270          if len(self._algorithm.param_infos) == 1:
    1271              self._swap_list1.append(0)
    1272              self._swap_list2.append(0)
    1273          else:
    1274              i = 0
    1275              while i < len(self._algorithm.param_infos):
    1276                  self._swap_list1.append(i)
    1277                  i += 2
    1278
    1279              i = 1
    1280              while i < len(self._algorithm.param_infos):
    1281                  self._swap_list2.append(i)
    1282                  i += 2
    1283
    1284          self._current_swap_list = self._swap_list1
    1285
    1286 -    def swap_all(self):
    1287
    1288          for x in self._current_swap_list:
    1289              self._algorithm.swap(x)
    1290
    1291          if self._current_swap_list == self._swap_list1:
    1292              self._current_swap_list = self._swap_list2
    1293          else:
    1294              self._current_swap_list = self._swap_list1
    1295
    1296
    1297 -class SingleSwapStatistics(object):
    1298      """
    1299      Tracks swap statistics of a single sampler pair.
    1300
    1301      @param param_info: ParameterInfo instance holding swap parameters
    1302      @type param_info: L{AbstractSwapParameterInfo}
    1303      """
    1304
    1305 -    def __init__(self, param_info):
    1306          self._total_swaps = 0
    1307          self._accepted_swaps = 0
    1308
    1309      @property
    1310 -    def total_swaps(self):
    1311          return self._total_swaps
    1312
    1313      @property
    1314 -    def accepted_swaps(self):
    1315          return self._accepted_swaps
    1316
    1317      @property
    1318 -    def acceptance_rate(self):
    1319          """
    1320          Acceptance rate of the sampler pair.
    1321          """
    1322          if self.total_swaps > 0:
    1323              return float(self.accepted_swaps) / float(self.total_swaps)
    1324          else:
    1325              return 0.
    1326
    1327 -    def update(self, accepted):
    1328          """
    1329          Updates swap statistics.
    1330          """
    1331          self._total_swaps += 1
    1332          self._accepted_swaps += int(accepted)
    1333
    1334
    1335 -class SwapStatistics(object):
    1336      """
    1337      Tracks swap statistics for an AbstractExchangeMC subclass instance.
    1338
    1339      @param param_infos: list of ParameterInfo instances providing information
    1340                          needed for performing swaps
    1341      @type param_infos: list of L{AbstractSwapParameterInfo}
    1342      """
    1343
    1344 -    def __init__(self, param_infos):
    1345          self._stats = [SingleSwapStatistics(x) for x in param_infos]
    1346
    1347      @property
    1348 -    def stats(self):
    1349          return tuple(self._stats)
    1350
    1351      @property
    1352 -    def acceptance_rates(self):
    1353          """
    1354          Returns acceptance rates for all swaps.
    1355          """
    1356          return [x.acceptance_rate for x in self._stats]
    1357
    1358
    1359 -class InterpolationFactory(object):
    1360      """
    1361      Produces interpolations for functions changed during non-equilibrium
    1362      trajectories.
    1363
    1364      @param protocol: protocol to be used to generate non-equilibrium trajectories
    1365      @type protocol: function mapping t to [0...1] for fixed tau
    1366      @param tau: switching time
    1367      @type tau: float
    1368      """
    1369
    1370 -    def __init__(self, protocol, tau):
    1371
    1372          self._protocol = None
    1373          self._tau = None
    1374
    1375          self.protocol = protocol
    1376          self.tau = tau
    1377
    1378      @property
    1379 -    def protocol(self):
    1380          return self._protocol
    1381      @protocol.setter
    1382 -    def protocol(self, value):
    1383          if not hasattr(value, '__call__'):
    1384              raise TypeError(value)
    1385          self._protocol = value
    1386
    1387      @property
    1388 -    def tau(self):
    1389          return self._tau
    1390      @tau.setter
    1391 -    def tau(self, value):
    1392          self._tau = float(value)
    1393
    1394 -    def build_gradient(self, gradient):
    1395          """
    1396          Create a gradient instance with according to given protocol
    1397          and switching time.
    1398
    1399          @param gradient: gradient with G(0) = G_1 and G(1) = G_2
    1400          @type gradient: callable
    1401          """
    1402          return Gradient(gradient, self._protocol, self._tau)
    1403
    1404 -    def build_temperature(self, temperature):
    1405          """
    1406          Create a temperature function according to given protocol and
    1407          switching time.
    1408
    1409          @param temperature: temperature with T(0) = T_1 and T(1) = T_2
    1410          @type temperature: callable
    1411          """
    1412          return lambda t: temperature(self.protocol(t, self.tau))
    1413
    1414
    1415 -class Gradient(AbstractGradient):
    1416
    1417 -    def __init__(self, gradient, protocol, tau):
    1418
    1419          self._protocol = protocol
    1420          self._gradient = gradient
    1421          self._tau = tau
    1422
    1423 -    def evaluate(self, q, t):
    1424          return self._gradient(q, self._protocol(t, self._tau))
    1425
    1426
    1427 -class ReplicaHistory(object):
    1428      '''
    1429      Replica history object, works with both RE and RENS for
    1430      the AlternatingAdjacentSwapScheme.
    1431
    1432      @param samples: list holding ensemble states
    1433      @type samples: list
    1434
    1435      @param swap_interval: interval with which swaps were attempted, e.g.,
    1436                            5 means that every 5th regular MC step is replaced
    1437                            by a swap
    1438      @type swap_interval: int
    1439
    1440      @param first_swap: sample index of the first sample generated by a swap attempt.
    1441                         If None, the first RE sampled is assumed to have sample index
    1442                         swap_interval. If specified, it has to be greater than zero
    1443      @type first_swap: int
    1444      '''
    1445
    1446 -    def __init__(self, samples, swap_interval, first_swap=None):
    1447          self.samples = samples
    1448          self.swap_interval = swap_interval
    1449          if first_swap == None:
    1450              self.first_swap = swap_interval - 1
    1451          elif first_swap > 0:
    1452              self.first_swap = first_swap - 1
    1453          else:
    1454              raise(ValueError("Sample index of first swap has to be greater than zero!"))
    1455          self.n_replicas = len(samples[0])
    1456
    1457      @staticmethod
    1458 -    def _change_direction(x):
    1459          if x == 1:
    1460              return -1
    1461          if x == -1:
    1462              return 1
    1463
    1464 -    def calculate_history(self, start_ensemble):
    1465          '''
    1466          Calculates the replica history of the first state of ensemble #start_ensemble.
    1467
    1468          @param start_ensemble: index of the ensemble to start at, zero-indexed
    1469          @type start_ensemble: int
    1470
    1471          @return: replica history as a list of ensemble indices
    1472          @rtype: list of ints
    1473          '''
    1474
    1475          sample_counter = 0
    1476
    1477          # determine the direction (up = 1, down = -1) in the "temperature ladder" of
    1478          # the first swap attempt. Remember: first swap series is always 0 <-> 1, 2 <-> 3, ...
    1479          if start_ensemble % 2 == 0:
    1480              direction = +1
    1481          else:
    1482              direction = -1
    1483
    1484          # if number of replicas is not even and the start ensemble is the highest-temperature-
    1485          # ensemble, the first swap will be attempted "downwards"
    1486          if start_ensemble % 2 == 0 and start_ensemble == self.n_replicas - 1:
    1487              direction = -1
    1488
    1489          # will store the indices of the ensembles the state will visit in chronological order
    1490          history = []
    1491
    1492          # the ensemble the state is currently in
    1493          ens = start_ensemble
    1494
    1495          while sample_counter < len(self.samples):
    1496              if self.n_replicas == 2:
    1497                  if (sample_counter - self.first_swap - 1) % self.swap_interval == 0 and \
    1498                     sample_counter >= self.first_swap:
    1499                      ## swap attempt: determine whether it was successfull or not
    1500                      # state after swap attempt
    1501                      current_sample = self.samples[sample_counter][ens]
    1502
    1503                      # state before swap attempt
    1504                      previous_sample = self.samples[sample_counter - 1][history[-1]]
    1505
    1506                      # swap was accepted when position of the current state doesn't equal
    1507                      # the position of the state before the swap attempt, that is, the last
    1508                      # state in the history
    1509                      swap_accepted = not numpy.all(current_sample.position ==
    1510                                                    previous_sample.position)
    1511
    1512                      if swap_accepted:
    1513                          if ens == 0:
    1514                              ens = 1
    1515                          else:
    1516                              ens = 0
    1517                      history.append(ens)
    1518                  else:
    1519                      history.append(ens)
    1520
    1521              else:
    1522
    1523                  if (sample_counter - self.first_swap - 1) % self.swap_interval == 0 and \
    1524                     sample_counter >= self.first_swap:
    1525                      # state after swap attempt
    1526                      current_sample = self.samples[sample_counter][ens]
    1527
    1528                      # state before swap attempt
    1529                      previous_sample = self.samples[sample_counter - 1][ens]
    1530
    1531                      # swap was accepted when position of the current state doesn't equal
    1532                      # the position of the state before the swap attempt, that is, the last
    1533                      # state in the history
    1534                      swap_accepted = not numpy.all(current_sample.position == previous_sample.position)
    1535
    1536                      if swap_accepted:
    1537                          ens += direction
    1538                      else:
    1539                          if ens == self.n_replicas - 1:
    1540                              # if at the top of the ladder, go downwards again
    1541                              direction = -1
    1542                          elif ens == 0:
    1543                              # if at the bottom of the ladder, go upwards
    1544                              direction = +1
    1545                          else:
    1546                              # in between, reverse the direction of the trajectory
    1547                              # in temperature space
    1548                              direction = self._change_direction(direction)
    1549
    1550                  history.append(ens)
    1551
    1552              sample_counter += 1
    1553
    1554          return history
    1555
    1556 -    def calculate_projected_trajectories(self, ensemble):
    1557          '''
    1558          Calculates sequentially correlated trajectories projected on a specific ensemble.
    1559
    1560          @param ensemble: ensemble index of ensemble of interest, zero-indexed
    1561          @type ensemble: int
    1562
    1563          @return: list of Trajectory objects containg sequentially correlated trajectories
    1564          @rtype: list of L{Trajectory} objects.
    1565          '''
    1566
    1567          trajectories = []
    1568
    1569          for i in range(self.n_replicas):
    1570              history = self.calculate_history(i)
    1571              traj = [x[ensemble] for k, x in enumerate(self.samples) if history[k] == ensemble]
    1572              trajectories.append(Trajectory(traj))
    1573
    1574          return trajectories
    1575
    1576 -    def calculate_trajectories(self):
    1577          '''
    1578          Calculates sequentially correlated trajectories.
    1579
    1580          @return: list of Trajectory objects containg sequentially correlated trajectories
    1581          @rtype: list of L{Trajectory} objects.
    1582          '''
    1583
    1584          trajectories = []
    1585
    1586          for i in range(self.n_replicas):
    1587              history = self.calculate_history(i)
    1588              traj = [x[history[k]] for k, x in enumerate(self.samples)]
    1589              trajectories.append(Trajectory(traj))
    1590
    1591          return trajectories
    1592