|
| 1 | +""" |
| 2 | +Comparing different samplers on a correlated bivariate normal distribution. |
| 3 | +
|
| 4 | +This example will sample a bivariate normal with Metropolis, NUTS and DEMetropolis |
| 5 | +at two correlations (0, 0.9) and print out the effective sample sizes, runtime and |
| 6 | +normalized effective sampling rates. |
| 7 | +""" |
| 8 | + |
| 9 | + |
| 10 | +import numpy as np |
| 11 | +import time |
| 12 | +import pandas as pd |
| 13 | +import pymc3 as pm |
| 14 | +import theano.tensor as tt |
| 15 | + |
| 16 | +# with this flag one can switch between defining the bivariate normal as |
| 17 | +# either a 2D MvNormal (USE_XY = False) split up the two dimensions into |
| 18 | +# two variables 'x' and 'y'. The latter is recommended because it highlights |
| 19 | +# different behaviour with respect to blocking. |
| 20 | +USE_XY = True |
| 21 | + |
| 22 | +def run(steppers, p): |
| 23 | + steppers = set(steppers) |
| 24 | + traces = {} |
| 25 | + effn = {} |
| 26 | + runtimes = {} |
| 27 | + |
| 28 | + with pm.Model() as model: |
| 29 | + if USE_XY: |
| 30 | + x = pm.Flat('x') |
| 31 | + y = pm.Flat('y') |
| 32 | + mu = np.array([0.,0.]) |
| 33 | + cov = np.array([[1.,p],[p,1.]]) |
| 34 | + z = pm.MvNormal.dist(mu=mu, cov=cov, shape=(2,)).logp(tt.stack([x,y])) |
| 35 | + pot = pm.Potential('logp_xy', z) |
| 36 | + start = {'x': 0, 'y': 0} |
| 37 | + else: |
| 38 | + mu = np.array([0.,0.]) |
| 39 | + cov = np.array([[1.,p],[p,1.]]) |
| 40 | + z = pm.MvNormal('z', mu=mu, cov=cov, shape=(2,)) |
| 41 | + start={'z': [0, 0]} |
| 42 | + |
| 43 | + for step_cls in steppers: |
| 44 | + name = step_cls.__name__ |
| 45 | + t_start = time.time() |
| 46 | + mt = pm.sample( |
| 47 | + draws=10000, |
| 48 | + chains=16, parallelize=False, |
| 49 | + step=step_cls(), |
| 50 | + start=start |
| 51 | + ) |
| 52 | + runtimes[name] = time.time() - t_start |
| 53 | + print('{} samples across {} chains'.format(len(mt) * mt.nchains, mt.nchains)) |
| 54 | + traces[name] = mt |
| 55 | + en = pm.diagnostics.effective_n(mt) |
| 56 | + print('effective: {}\r\n'.format(en)) |
| 57 | + if USE_XY: |
| 58 | + effn[name] = np.mean(en['x']) / len(mt) / mt.nchains |
| 59 | + else: |
| 60 | + effn[name] = np.mean(en['z']) / len(mt) / mt.nchains |
| 61 | + return traces, effn, runtimes |
| 62 | + |
| 63 | + |
| 64 | +if __name__ == '__main__': |
| 65 | + methods = [ |
| 66 | + pm.Metropolis, |
| 67 | + pm.Slice, |
| 68 | + pm.NUTS, |
| 69 | + pm.DEMetropolis |
| 70 | + ] |
| 71 | + names = [c.__name__ for c in methods] |
| 72 | + |
| 73 | + df_base = pd.DataFrame(columns=['p'] + names) |
| 74 | + df_base['p'] = [.0,.9] |
| 75 | + df_base = df_base.set_index('p') |
| 76 | + |
| 77 | + df_effectiven = df_base.copy() |
| 78 | + df_runtime = df_base.copy() |
| 79 | + df_performance = df_base.copy() |
| 80 | + |
| 81 | + for p in df_effectiven.index: |
| 82 | + trace, rate, runtime = run(methods, p) |
| 83 | + for name in names: |
| 84 | + df_effectiven.set_value(p, name, rate[name]) |
| 85 | + df_runtime.set_value(p, name, runtime[name]) |
| 86 | + df_performance.set_value(p, name, rate[name] / runtime[name]) |
| 87 | + |
| 88 | + print('\r\nEffective sample size [0...1]') |
| 89 | + print(df_effectiven.T.to_string(float_format='{:.3f}'.format)) |
| 90 | + |
| 91 | + print('\r\nRuntime [s]') |
| 92 | + print(df_runtime.T.to_string(float_format='{:.1f}'.format)) |
| 93 | + |
| 94 | + if 'NUTS' in names: |
| 95 | + print('\r\nNormalized effective sampling rate [0...1]') |
| 96 | + df_performance = df_performance.T / df_performance.loc[0]['NUTS'] |
| 97 | + else: |
| 98 | + print('\r\nNormalized effective sampling rate [1/s]') |
| 99 | + df_performance = df_performance.T |
| 100 | + print(df_performance.to_string(float_format='{:.3f}'.format)) |
0 commit comments