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

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

  1  """ 
  2  Implements several extended-ensemble Monte Carlo sampling algorithms. 
  3   
  4  Here is a short example which shows how to sample from a PDF using the replica 
  5  exchange with non-equilibrium switches (RENS) method. It draws 5000 samples from 
  6  a 1D normal distribution using the RENS algorithm working on three Markov chains 
  7  being generated by the HMC algorithm: 
  8   
  9   
 10      >>> import numpy 
 11      >>> from numpy import sqrt 
 12      >>> from csb.io.plots import Chart 
 13      >>> from csb.statistics.pdf import Normal 
 14      >>> from csb.statistics.samplers import State 
 15      >>> from csb.statistics.samplers.mc import ThermostattedMDRENSSwapParameterInfo, AlternatingAdjacentSwapScheme 
 16      >>> from csb.statistics.samplers.mc.multichain import ThermostattedMDRENS 
 17      >>> from csb.statistics.samplers.mc.singlechain import HMCSampler 
 18   
 19      >>> # Pick some initial state for the different Markov chains: 
 20      >>> initial_state = State(numpy.array([1.])) 
 21   
 22      >>> # Set standard deviations: 
 23      >>> std_devs = [1./sqrt(5), 1. / sqrt(3), 1.] 
 24   
 25      >>> # Set HMC timesteps and trajectory length: 
 26      >>> hmc_timesteps = [0.6, 0.7, 0.6] 
 27      >>> hmc_trajectory_length = 20 
 28      >>> hmc_gradients = [lambda q, t: 1 / (std_dev ** 2) * q for std_dev in std_devs] 
 29   
 30      >>> # Set parameters for the thermostatted RENS algorithm: 
 31      >>> rens_trajectory_length = 30 
 32      >>> rens_timesteps = [0.3, 0.5] 
 33   
 34      >>> # Set interpolation gradients as a function of the work parameter l: 
 35      >>> rens_gradients = [lambda q, l, i=i: (l / (std_devs[i + 1] ** 2) + (1 - l) / (std_devs[i] ** 2)) * q  
 36                            for i in range(len(std_devs)-1)] 
 37   
 38      >>> # Initialize HMC samplers: 
 39      >>> samplers = [HMCSampler(Normal(sigma=std_devs[i]), initial_state, hmc_gradients[i], hmc_timesteps[i], 
 40                      hmc_trajectory_length) for i in range(len(std_devs))] 
 41   
 42      >>> # Create swap parameter objects: 
 43      params = [ThermostattedMDRENSSwapParameterInfo(samplers[0], samplers[1], rens_timesteps[0], 
 44                rens_trajectory_length, rens_gradients[0]), 
 45                ThermostattedMDRENSSwapParameterInfo(samplers[1], samplers[2], rens_timesteps[1], 
 46                rens_trajectory_length, rens_gradients[1])] 
 47   
 48      >>> # Initialize thermostatted RENS algorithm: 
 49      >>> algorithm = ThermostattedMDRENS(samplers, params) 
 50   
 51      >>> # Initialize swapping scheme: 
 52      >>> swapper = AlternatingAdjacentSwapScheme(algorithm) 
 53   
 54      >>> # Initialize empty list which will store the samples: 
 55      >>> states = [] 
 56      >>> for i in range(5000): 
 57              if i % 5 == 0: 
 58                  swapper.swap_all() 
 59              states.append(algorithm.sample()) 
 60   
 61      >>> # Print acceptance rates: 
 62      >>> print('HMC acceptance rates:', [s.acceptance_rate for s in samplers]) 
 63      >>> print('swap acceptance rates:', algorithm.acceptance_rates) 
 64   
 65      >>> # Create and plot histogram for first sampler and numpy.random.normal reference: 
 66      >>> chart = Chart() 
 67      >>> rawstates = [state[0].position[0] for state in states] 
 68      >>> chart.plot.hist([numpy.random.normal(size=5000, scale=std_devs[0]), rawstates], bins=30, normed=True) 
 69      >>> chart.plot.legend(['numpy.random.normal', 'RENS + HMC']) 
 70      >>> chart.show() 
 71   
 72   
 73  For L{ReplicaExchangeMC} (RE), the procedure is easier because apart from the 
 74  two sampler instances the corresponding L{RESwapParameterInfo} objects take 
 75  no arguments. 
 76   
 77  Every replica exchange algorithm in this module (L{ReplicaExchangeMC}, L{MDRENS}, 
 78  L{ThermostattedMDRENS}) is used in a similar way. A simulation is always 
 79  initialized with a list of samplers (instances of classes derived from 
 80  L{AbstractSingleChainMC}) and a list of L{AbstractSwapParameterInfo} objects 
 81  suited for the algorithm under consideration. Every L{AbstractSwapParameterInfo} 
 82  object holds all the information needed to perform a swap between two samplers. 
 83  The usual scheme is to swap only adjacent replicae in a scheme:: 
 84   
 85      1 <--> 2, 3 <--> 4, ... 
 86      2 <--> 3, 4 <--> 5, ... 
 87      1 <--> 2, 3 <--> 4, ... 
 88       
 89  This swapping scheme is implemented in the L{AlternatingAdjacentSwapScheme} class, 
 90  but different schemes can be easily implemented by deriving from L{AbstractSwapScheme}. 
 91  Then the simulation is run by looping over the number of samples to be drawn 
 92  and calling the L{AbstractExchangeMC.sample} method of the algorithm. By calling 
 93  the L{AbstractSwapScheme.swap_all} method of the specific L{AbstractSwapScheme} 
 94  implementation, all swaps defined in the list of L{AbstractSwapParameterInfo} 
 95  objects are performed according to the swapping scheme. The 
 96  L{AbstractSwapScheme.swap_all} method may be called for example after sampling 
 97  intervals of a fixed length or randomly. 
 98  """ 
 99   
100  import 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 T1 = swapcom.sampler1.temperature 188 T2 = swapcom.sampler2.temperature 189 190 state1 = swapcom.sampler1.state 191 state2 = swapcom.sampler2.state 192 193 proposal1 = swapcom.proposal1 194 proposal2 = swapcom.proposal2 195 196 swapcom.acceptance_probability = csb.numeric.exp(-E1(proposal1.position) / T1 + E1(state1.position) / T1 - \ 197 E2(proposal2.position) / T2 + E2(state2.position) / T2) 198 199 return swapcom
200
201 -class MDRENS(AbstractRENS):
202 """ 203 Replica Exchange with Nonequilibrium Switches (RENS, Ballard & Jarzynski 2009) 204 with Molecular Dynamics (MD) trajectories. 205 206 @param samplers: Samplers which sample their 207 respective equilibrium distributions 208 @type samplers: list of L{AbstractSingleChainMC} 209 210 @param param_infos: ParameterInfo instance holding 211 information required to perform a MDRENS swap 212 @type param_infos: list of L{MDRENSSwapParameterInfo} 213 214 @param integrator: Subclass of L{AbstractIntegrator} to be used to 215 calculate the non-equilibrium trajectories 216 @type integrator: type 217 """ 218
219 - def __init__(self, samplers, param_infos, 220 integrator=csb.numeric.integrators.FastLeapFrog):
221 222 super(MDRENS, self).__init__(samplers, param_infos) 223 224 self._integrator = integrator
225
226 - def _run_traj_generator(self, traj_info):
227 228 tau = traj_info.param_info.traj_length * traj_info.param_info.timestep 229 factory = InterpolationFactory(traj_info.protocol, tau) 230 231 gen = MDPropagator(factory.build_gradient(traj_info.param_info.gradient), 232 traj_info.param_info.timestep, self._integrator) 233 234 traj = gen.generate(traj_info.init_state, int(traj_info.param_info.traj_length)) 235 return traj
236
237 -class ThermostattedMDRENS(MDRENS):
238 """ 239 Replica Exchange with Nonequilibrium Switches (RENS, Ballard & Jarzynski, 2009) 240 with Andersen-thermostatted Molecular Dynamics (MD) trajectories. 241 242 @param samplers: Samplers which sample their 243 respective equilibrium distributions 244 @type samplers: list of L{AbstractSingleChainMC} 245 246 @param param_infos: ParameterInfo instance holding 247 information required to perform a MDRENS swap 248 @type param_infos: list of L{ThermostattedMDRENSSwapParameterInfo} 249 250 @param integrator: Subclass of L{AbstractIntegrator} to be used to 251 calculate the non-equilibrium trajectories 252 @type integrator: type 253 """ 254
255 - def __init__(self, samplers, param_infos, integrator=csb.numeric.integrators.LeapFrog):
256 257 super(ThermostattedMDRENS, self).__init__(samplers, param_infos, integrator)
258
259 - def _run_traj_generator(self, traj_info):
260 261 tau = traj_info.param_info.traj_length * traj_info.param_info.timestep 262 factory = InterpolationFactory(traj_info.protocol, tau) 263 264 grad = factory.build_gradient(traj_info.param_info.gradient) 265 temp = factory.build_temperature(traj_info.param_info.temperature) 266 267 gen = ThermostattedMDPropagator(grad, 268 traj_info.param_info.timestep, temp, 269 traj_info.param_info.collision_probability, 270 traj_info.param_info.collision_interval, 271 self._integrator) 272 273 traj = gen.generate(traj_info.init_state, traj_info.param_info.traj_length) 274 275 return traj
276