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

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

  1  """ 
  2  Implements several extended-ensemble Monte Carlo sampling algorithms. 
  3   
  4  Here is a short example which shows how to sample from a PDF using the replica 
  5  exchange with non-equilibrium switches (RENS) method. It draws 5000 samples from 
  6  a 1D normal distribution using the RENS algorithm working on three Markov chains 
  7  being generated by the HMC algorithm: 
  8   
  9   
 10      >>> import numpy 
 11      >>> from numpy import sqrt 
 12      >>> from csb.io.plots import Chart 
 13      >>> from csb.statistics.pdf import Normal 
 14      >>> from csb.statistics.samplers import State 
 15      >>> from csb.statistics.samplers.mc import ThermostattedMDRENSSwapParameterInfo, AlternatingAdjacentSwapScheme 
 16      >>> from csb.statistics.samplers.mc.multichain import ThermostattedMDRENS 
 17      >>> from csb.statistics.samplers.mc.singlechain import HMCSampler 
 18   
 19      >>> # Pick some initial state for the different Markov chains: 
 20      >>> initial_state = State(numpy.array([1.])) 
 21   
 22      >>> # Set standard deviations: 
 23      >>> std_devs = [1., 1. / sqrt(3), 1. / sqrt(5)] 
 24   
 25      >>> # Set HMC timesteps and trajectory length: 
 26      >>> hmc_timesteps = [0.7, 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.5, 0.3, 0.1] 
 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), 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 csb.numeric 
101   
102  from csb.statistics.samplers.mc import AbstractExchangeMC, AbstractRENS, RESwapCommunicator 
103  from csb.statistics.samplers.mc.propagators import MDPropagator, ThermostattedMDPropagator 
104  from csb.numeric.integrators import AbstractGradient 
105 106 107 -class InterpolationFactory(object):
108 """ 109 Produces interpolations for functions changed during non-equilibrium 110 trajectories. 111 112 @param protocol: protocol to be used to generate non-equilibrium trajectories 113 @type protocol: function mapping t to [0...1] for fixed tau 114 @param tau: switching time 115 @type tau: float 116 """ 117
118 - def __init__(self, protocol, tau):
119 120 self._protocol = None 121 self._tau = None 122 123 self.protocol = protocol 124 self.tau = tau
125 126 @property
127 - def protocol(self):
128 return self._protocol
129 @protocol.setter
130 - def protocol(self, value):
131 if not hasattr(value, '__call__'): 132 raise TypeError(value) 133 self._protocol = value
134 135 @property
136 - def tau(self):
137 return self._tau
138 @tau.setter
139 - def tau(self, value):
140 self._tau = float(value)
141
142 - def build_gradient(self, gradient):
143 """ 144 Create a gradient instance with according to given protocol 145 and switching time. 146 147 @param gradient: gradient with G(0) = G_1 and G(1) = G_2 148 @type gradient: callable 149 """ 150 return Gradient(gradient, self._protocol, self._tau)
151
152 - def build_temperature(self, temperature):
153 """ 154 Create a temperature function according to given protocol and 155 switching time. 156 157 @param temperature: temperature with T(0) = T_1 and T(1) = T_2 158 @type temperature: callable 159 """ 160 return lambda t: temperature(self.protocol(t, self.tau))
161
162 -class Gradient(AbstractGradient):
163
164 - def __init__(self, gradient, protocol, tau):
165 166 self._protocol = protocol 167 self._gradient = gradient 168 self._tau = tau
169
170 - def evaluate(self, q, t):
171 return self._gradient(q, self._protocol(t, self._tau))
172
173 -class ReplicaExchangeMC(AbstractExchangeMC):
174 """ 175 Replica Exchange (RE, Swendsen & Yang 1986) implementation. 176 """ 177
178 - def _propose_swap(self, param_info):
181
182 - def _calc_pacc_swap(self, swapcom):
183 184 E1 = lambda x:-swapcom.sampler1._pdf.log_prob(x) 185 E2 = lambda x:-swapcom.sampler2._pdf.log_prob(x) 186 187 state1 = swapcom.sampler1.state 188 state2 = swapcom.sampler2.state 189 190 proposal1 = swapcom.proposal1 191 proposal2 = swapcom.proposal2 192 193 swapcom.acceptance_probability = csb.numeric.exp(-E1(proposal1.position) + E1(state1.position) - \ 194 E2(proposal2.position) + E2(state2.position)) 195 196 return swapcom
197
198 -class MDRENS(AbstractRENS):
199 """ 200 Replica Exchange with Nonequilibrium Switches (RENS, Ballard & Jarzynski 2009) 201 with Molecular Dynamics (MD) trajectories. 202 203 @param samplers: Samplers which sample their 204 respective equilibrium distributions 205 @type samplers: list of L{AbstractSingleChainMC} 206 207 @param param_infos: ParameterInfo instance holding 208 information required to perform a MDRENS swap 209 @type param_infos: list of L{MDRENSSwapParameterInfo} 210 211 @param integrator: Subclass of L{AbstractIntegrator} to be used to 212 calculate the non-equilibrium trajectories 213 @type integrator: type 214 """ 215
216 - def __init__(self, samplers, param_infos, 217 integrator=csb.numeric.integrators.FastLeapFrog):
218 219 super(MDRENS, self).__init__(samplers, param_infos) 220 221 self._integrator = integrator
222
223 - def _run_traj_generator(self, traj_info):
224 225 tau = traj_info.param_info.traj_length * traj_info.param_info.timestep 226 factory = InterpolationFactory(traj_info.protocol, tau) 227 228 gen = MDPropagator(factory.build_gradient(traj_info.param_info.gradient), 229 traj_info.param_info.timestep, self._integrator) 230 231 traj = gen.generate(traj_info.init_state, int(traj_info.param_info.traj_length)) 232 return traj
233
234 -class ThermostattedMDRENS(MDRENS):
235 """ 236 Replica Exchange with Nonequilibrium Switches (RENS, Ballard & Jarzynski, 2009) 237 with Andersen-thermostatted Molecular Dynamics (MD) trajectories. 238 239 @param samplers: Samplers which sample their 240 respective equilibrium distributions 241 @type samplers: list of L{AbstractSingleChainMC} 242 243 @param param_infos: ParameterInfo instance holding 244 information required to perform a MDRENS swap 245 @type param_infos: list of L{ThermostattedMDRENSSwapParameterInfo} 246 247 @param integrator: Subclass of L{AbstractIntegrator} to be used to 248 calculate the non-equilibrium trajectories 249 @type integrator: type 250 """ 251
252 - def __init__(self, samplers, param_infos, integrator=csb.numeric.integrators.LeapFrog):
253 254 super(ThermostattedMDRENS, self).__init__(samplers, param_infos, integrator)
255
256 - def _run_traj_generator(self, traj_info):
257 258 tau = traj_info.param_info.traj_length * traj_info.param_info.timestep 259 factory = InterpolationFactory(traj_info.protocol, tau) 260 261 grad = factory.build_gradient(traj_info.param_info.gradient) 262 temp = factory.build_temperature(traj_info.param_info.temperature) 263 264 gen = ThermostattedMDPropagator(grad, 265 traj_info.param_info.timestep, temp, 266 traj_info.param_info.collision_probability, 267 traj_info.param_info.collision_interval, 268 self._integrator) 269 270 traj = gen.generate(traj_info.init_state, traj_info.param_info.traj_length) 271 272 return traj
273