88from theano import tensor as TT
99import theano
1010
11- sharedX = lambda X , name : \
12- shared (numpy .asarray (X , dtype = theano .config .floatX ), name = name )
11+ sharedX = ( lambda X , name :
12+ shared (numpy .asarray (X , dtype = theano .config .floatX ), name = name ) )
1313
1414
1515def kinetic_energy (vel ):
@@ -145,13 +145,14 @@ def leapfrog(pos, vel, step):
145145
146146 # perform leapfrog updates: the scan op is used to repeatedly compute
147147 # vel(t + (m-1/2)*stepsize) and pos(t + m*stepsize) for m in [2,n_steps].
148- (all_pos , all_vel ), scan_updates = theano .scan (leapfrog ,
149- outputs_info = [
150- dict (initial = pos_full_step ),
151- dict (initial = vel_half_step ),
152- ],
153- non_sequences = [stepsize ],
154- n_steps = n_steps - 1 )
148+ (all_pos , all_vel ), scan_updates = theano .scan (
149+ leapfrog ,
150+ outputs_info = [
151+ dict (initial = pos_full_step ),
152+ dict (initial = vel_half_step ),
153+ ],
154+ non_sequences = [stepsize ],
155+ n_steps = n_steps - 1 )
155156 final_pos = all_pos [- 1 ]
156157 final_vel = all_vel [- 1 ]
157158 # NOTE: Scan always returns an updates dictionary, in case the
@@ -171,6 +172,7 @@ def leapfrog(pos, vel, step):
171172 # return new proposal state
172173 return final_pos , final_vel
173174
175+
174176# start-snippet-1
175177def hmc_move (s_rng , positions , energy_fn , stepsize , n_steps ):
176178 """
@@ -224,10 +226,11 @@ def hmc_move(s_rng, positions, energy_fn, stepsize, n_steps):
224226 # end-snippet-4
225227 return accept , final_pos
226228
229+
227230# start-snippet-5
228231def hmc_updates (positions , stepsize , avg_acceptance_rate , final_pos , accept ,
229- target_acceptance_rate , stepsize_inc , stepsize_dec ,
230- stepsize_min , stepsize_max , avg_acceptance_slowness ):
232+ target_acceptance_rate , stepsize_inc , stepsize_dec ,
233+ stepsize_min , stepsize_max , avg_acceptance_slowness ):
231234 """This function is executed after `n_steps` of HMC sampling
232235 (`hmc_move` function). It creates the updates dictionary used by
233236 the `simulate` function. It takes care of updating: the position
@@ -293,14 +296,15 @@ def hmc_updates(positions, stepsize, avg_acceptance_rate, final_pos, accept,
293296 # perform exponential moving average
294297 mean_dtype = theano .scalar .upcast (accept .dtype , avg_acceptance_rate .dtype )
295298 new_acceptance_rate = TT .add (
296- avg_acceptance_slowness * avg_acceptance_rate ,
297- (1.0 - avg_acceptance_slowness ) * accept .mean (dtype = mean_dtype ))
299+ avg_acceptance_slowness * avg_acceptance_rate ,
300+ (1.0 - avg_acceptance_slowness ) * accept .mean (dtype = mean_dtype ))
298301 # end-snippet-6 start-snippet-8
299302 return [(positions , new_positions ),
300303 (stepsize , new_stepsize ),
301304 (avg_acceptance_rate , new_acceptance_rate )]
302305 # end-snippet-8
303306
307+
304308class HMC_sampler (object ):
305309 """
306310 Convenience wrapper for performing Hybrid Monte Carlo (HMC). It creates the
@@ -322,11 +326,11 @@ def __init__(self, **kwargs):
322326
323327 @classmethod
324328 def new_from_shared_positions (
325- cls ,
326- shared_positions ,
329+ cls ,
330+ shared_positions ,
327331 energy_fn ,
328- initial_stepsize = 0.01 ,
329- target_acceptance_rate = .9 ,
332+ initial_stepsize = 0.01 ,
333+ target_acceptance_rate = .9 ,
330334 n_steps = 20 ,
331335 stepsize_dec = 0.98 ,
332336 stepsize_min = 0.001 ,
@@ -350,8 +354,6 @@ def new_from_shared_positions(
350354 sampling to work.
351355
352356 """
353- batchsize = shared_positions .shape [0 ]
354-
355357 # allocate shared variables
356358 stepsize = sharedX (initial_stepsize , 'hmc_stepsize' )
357359 avg_acceptance_rate = sharedX (target_acceptance_rate ,
@@ -360,40 +362,40 @@ def new_from_shared_positions(
360362
361363 # define graph for an `n_steps` HMC simulation
362364 accept , final_pos = hmc_move (
363- s_rng ,
364- shared_positions ,
365- energy_fn ,
366- stepsize ,
367- n_steps )
365+ s_rng ,
366+ shared_positions ,
367+ energy_fn ,
368+ stepsize ,
369+ n_steps )
368370
369371 # define the dictionary of updates, to apply on every `simulate` call
370372 simulate_updates = hmc_updates (
371- shared_positions ,
372- stepsize ,
373- avg_acceptance_rate ,
374- final_pos = final_pos ,
375- accept = accept ,
376- stepsize_min = stepsize_min ,
377- stepsize_max = stepsize_max ,
378- stepsize_inc = stepsize_inc ,
379- stepsize_dec = stepsize_dec ,
380- target_acceptance_rate = target_acceptance_rate ,
381- avg_acceptance_slowness = avg_acceptance_slowness )
373+ shared_positions ,
374+ stepsize ,
375+ avg_acceptance_rate ,
376+ final_pos = final_pos ,
377+ accept = accept ,
378+ stepsize_min = stepsize_min ,
379+ stepsize_max = stepsize_max ,
380+ stepsize_inc = stepsize_inc ,
381+ stepsize_dec = stepsize_dec ,
382+ target_acceptance_rate = target_acceptance_rate ,
383+ avg_acceptance_slowness = avg_acceptance_slowness )
382384
383385 # compile theano function
384386 simulate = function ([], [], updates = simulate_updates )
385387
386388 # create HMC_sampler object with the following attributes ...
387389 return cls (
388- positions = shared_positions ,
389- stepsize = stepsize ,
390- stepsize_min = stepsize_min ,
391- stepsize_max = stepsize_max ,
392- avg_acceptance_rate = avg_acceptance_rate ,
393- target_acceptance_rate = target_acceptance_rate ,
394- s_rng = s_rng ,
395- _updates = simulate_updates ,
396- simulate = simulate )
390+ positions = shared_positions ,
391+ stepsize = stepsize ,
392+ stepsize_min = stepsize_min ,
393+ stepsize_max = stepsize_max ,
394+ avg_acceptance_rate = avg_acceptance_rate ,
395+ target_acceptance_rate = target_acceptance_rate ,
396+ s_rng = s_rng ,
397+ _updates = simulate_updates ,
398+ simulate = simulate )
397399
398400 def draw (self , ** kwargs ):
399401 """
0 commit comments