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.statistics.samplers.mc import Trajectory 
105  from csb.numeric.integrators import AbstractGradient 
106 107 108 -class InterpolationFactory(object):
109 """ 110 Produces interpolations for functions changed during non-equilibrium 111 trajectories. 112 113 @param protocol: protocol to be used to generate non-equilibrium trajectories 114 @type protocol: function mapping t to [0...1] for fixed tau 115 @param tau: switching time 116 @type tau: float 117 """ 118
119 - def __init__(self, protocol, tau):
120 121 self._protocol = None 122 self._tau = None 123 124 self.protocol = protocol 125 self.tau = tau
126 127 @property
128 - def protocol(self):
129 return self._protocol
130 @protocol.setter
131 - def protocol(self, value):
132 if not hasattr(value, '__call__'): 133 raise TypeError(value) 134 self._protocol = value
135 136 @property
137 - def tau(self):
138 return self._tau
139 @tau.setter
140 - def tau(self, value):
141 self._tau = float(value)
142
143 - def build_gradient(self, gradient):
144 """ 145 Create a gradient instance with according to given protocol 146 and switching time. 147 148 @param gradient: gradient with G(0) = G_1 and G(1) = G_2 149 @type gradient: callable 150 """ 151 return Gradient(gradient, self._protocol, self._tau)
152
153 - def build_temperature(self, temperature):
154 """ 155 Create a temperature function according to given protocol and 156 switching time. 157 158 @param temperature: temperature with T(0) = T_1 and T(1) = T_2 159 @type temperature: callable 160 """ 161 return lambda t: temperature(self.protocol(t, self.tau))
162
163 -class Gradient(AbstractGradient):
164
165 - def __init__(self, gradient, protocol, tau):
166 167 self._protocol = protocol 168 self._gradient = gradient 169 self._tau = tau
170
171 - def evaluate(self, q, t):
172 return self._gradient(q, self._protocol(t, self._tau))
173
174 -class ReplicaExchangeMC(AbstractExchangeMC):
175 """ 176 Replica Exchange (RE, Swendsen & Yang 1986) implementation. 177 """ 178
179 - def _propose_swap(self, param_info):
185
186 - def _calc_pacc_swap(self, swapcom):
187 188 E1 = lambda x:-swapcom.sampler1._pdf.log_prob(x) 189 E2 = lambda x:-swapcom.sampler2._pdf.log_prob(x) 190 191 T1 = swapcom.sampler1.temperature 192 T2 = swapcom.sampler2.temperature 193 194 state1 = swapcom.traj12.initial 195 state2 = swapcom.traj21.initial 196 197 proposal1 = swapcom.traj21.final 198 proposal2 = swapcom.traj12.final 199 200 swapcom.acceptance_probability = csb.numeric.exp(-E1(proposal1.position) / T1 201 + E1(state1.position) / T1 202 - E2(proposal2.position) / T2 203 + E2(state2.position) / T2) 204 205 return swapcom
206
207 -class MDRENS(AbstractRENS):
208 """ 209 Replica Exchange with Nonequilibrium Switches (RENS, Ballard & Jarzynski 2009) 210 with Molecular Dynamics (MD) trajectories. 211 212 @param samplers: Samplers which sample their 213 respective equilibrium distributions 214 @type samplers: list of L{AbstractSingleChainMC} 215 216 @param param_infos: ParameterInfo instance holding 217 information required to perform a MDRENS swap 218 @type param_infos: list of L{MDRENSSwapParameterInfo} 219 220 @param integrator: Subclass of L{AbstractIntegrator} to be used to 221 calculate the non-equilibrium trajectories 222 @type integrator: type 223 """ 224
225 - def __init__(self, samplers, param_infos, 226 integrator=csb.numeric.integrators.FastLeapFrog):
227 228 super(MDRENS, self).__init__(samplers, param_infos) 229 230 self._integrator = integrator
231
232 - def _run_traj_generator(self, traj_info):
233 234 tau = traj_info.param_info.traj_length * traj_info.param_info.timestep 235 factory = InterpolationFactory(traj_info.protocol, tau) 236 237 gen = MDPropagator(factory.build_gradient(traj_info.param_info.gradient), 238 traj_info.param_info.timestep, 239 mass_matrix=traj_info.param_info.mass_matrix, 240 integrator=self._integrator) 241 242 traj = gen.generate(traj_info.init_state, int(traj_info.param_info.traj_length)) 243 return traj
244
245 -class ThermostattedMDRENS(MDRENS):
246 """ 247 Replica Exchange with Nonequilibrium Switches (RENS, Ballard & Jarzynski, 2009) 248 with Andersen-thermostatted Molecular Dynamics (MD) trajectories. 249 250 @param samplers: Samplers which sample their 251 respective equilibrium distributions 252 @type samplers: list of L{AbstractSingleChainMC} 253 254 @param param_infos: ParameterInfo instance holding 255 information required to perform a MDRENS swap 256 @type param_infos: list of L{ThermostattedMDRENSSwapParameterInfo} 257 258 @param integrator: Subclass of L{AbstractIntegrator} to be used to 259 calculate the non-equilibrium trajectories 260 @type integrator: type 261 """ 262
263 - def __init__(self, samplers, param_infos, integrator=csb.numeric.integrators.LeapFrog):
264 265 super(ThermostattedMDRENS, self).__init__(samplers, param_infos, integrator)
266
267 - def _run_traj_generator(self, traj_info):
268 269 tau = traj_info.param_info.traj_length * traj_info.param_info.timestep 270 factory = InterpolationFactory(traj_info.protocol, tau) 271 272 grad = factory.build_gradient(traj_info.param_info.gradient) 273 temp = factory.build_temperature(traj_info.param_info.temperature) 274 275 gen = ThermostattedMDPropagator(grad, 276 traj_info.param_info.timestep, temperature=temp, 277 collision_probability=traj_info.param_info.collision_probability, 278 update_interval=traj_info.param_info.collision_interval, 279 mass_matrix=traj_info.param_info.mass_matrix, 280 integrator=self._integrator) 281 282 283 traj = gen.generate(traj_info.init_state, traj_info.param_info.traj_length) 284 285 return traj
286