From be6f9211d8cd27778b05e09303071207a2d49b2b Mon Sep 17 00:00:00 2001 From: Thomas Wiecki Date: Sat, 12 Nov 2016 01:06:01 +0900 Subject: [PATCH 1/7] PERF Make NUTS 1.5x faster by moving E0 calculation outside of buildtree inner loop. --- pymc3/step_methods/nuts.py | 96 +++++++++++++++++++------------------- 1 file changed, 49 insertions(+), 47 deletions(-) diff --git a/pymc3/step_methods/nuts.py b/pymc3/step_methods/nuts.py index f042373745..9152920054 100644 --- a/pymc3/step_methods/nuts.py +++ b/pymc3/step_methods/nuts.py @@ -13,7 +13,6 @@ __all__ = ['NUTS'] - class NUTS(ArrayStepShared): """ Automatically tunes step size and adjust number of steps for good performance. @@ -32,7 +31,7 @@ def __init__(self, vars=None, scaling=None, step_scale=0.25, is_cov=False, state k=0.75, t0=10, model=None, - profile=False, **kwargs): + profile=False,**kwargs): """ Parameters ---------- @@ -68,12 +67,14 @@ def __init__(self, vars=None, scaling=None, step_scale=0.25, is_cov=False, state scaling = model.test_point if isinstance(scaling, dict): - scaling = guess_scaling( - Point(scaling, model=model), model=model, vars=vars) + scaling = guess_scaling(Point(scaling, model=model), model=model, vars = vars) + + n = scaling.shape[0] - self.step_size = step_scale / n**(1 / 4.) + self.step_size = step_scale / n**(1/4.) + self.potential = quad_potential(scaling, is_cov, as_cov=False) @@ -88,22 +89,39 @@ def __init__(self, vars=None, scaling=None, step_scale=0.25, is_cov=False, state self.k = k self.Hbar = 0 - self.u = log(self.step_size * 10) + self.u = log(self.step_size*10) self.m = 1 + #self.H = Hamiltonian(model.fastlogp, model.fastdlogp(vars), self.potential) + shared = make_shared_replacements(vars, model) - self.leapfrog1_dE = leapfrog1_dE( - model.logpt, vars, shared, self.potential, profile=profile) + + dlogp = gradient(model.logpt, vars) + (logp, dlogp), q = join_nonshared_inputs([model.logpt, dlogp], vars, shared) + logp = CallableTensor(logp) + dlogp = CallableTensor(dlogp) + + + self.H = Hamiltonian(logp, dlogp, self.potential) + + p = theano.tensor.dvector('p') + p.tag.test_value = q.tag.test_value + E0 = energy(self.H, q, p) + self.E0 = theano.function([q, p], [E0]) + self.E0.trust_input = True + + self.leapfrog1_dE = leapfrog1_dE(self.H, q, profile=profile) super(NUTS, self).__init__(vars, shared, **kwargs) def astep(self, q0): - # Hamiltonian(self.logp, self.dlogp, self.potential) - H = self.leapfrog1_dE + leapfrog = self.leapfrog1_dE Emax = self.Emax e = self.step_size p0 = self.potential.random() + E0 = self.E0(q0, p0) + u = uniform() q = qn = qp = q0 p = pn = pp = p0 @@ -114,13 +132,11 @@ def astep(self, q0): v = bern(.5) * 2 - 1 if v == -1: - qn, pn, _, _, q1, n1, s1, a, na = buildtree( - H, qn, pn, u, v, j, e, Emax, q0, p0) + qn, pn, _, _, q1, n1, s1, a, na = buildtree(leapfrog, qn, pn, u, v, j, e, Emax, q0, p0, E0) else: - _, _, qp, pp, q1, n1, s1, a, na = buildtree( - H, qp, pp, u, v, j, e, Emax, q0, p0) + _, _, qp, pp, q1, n1, s1, a, na = buildtree(leapfrog, qp, pp, u, v, j, e, Emax, q0, p0, E0) - if s1 == 1 and bern(min(1, n1 * 1. / n)): + if s1 == 1 and bern(min(1, n1*1./n)): q = q1 n = n + n1 @@ -131,11 +147,10 @@ def astep(self, q0): p = -p - w = 1. / (self.m + self.t0) - self.Hbar = (1 - w) * self.Hbar + w * \ - (self.target_accept - a * 1. / na) + w = 1./(self.m+self.t0) + self.Hbar = (1 - w) * self.Hbar + w*(self.target_accept - a*1./na) - self.step_size = exp(self.u - (self.m**.5 / self.gamma) * self.Hbar) + self.step_size = exp(self.u - (self.m**.5/self.gamma)*self.Hbar) self.m += 1 return q @@ -147,26 +162,22 @@ def competence(var): return Competence.INCOMPATIBLE -def buildtree(H, q, p, u, v, j, e, Emax, q0, p0): +def buildtree(leapfrog1_dE, q, p, u, v, j, e, Emax, q0, p0, E0): if j == 0: - leapfrog1_dE = H - q1, p1, dE = leapfrog1_dE(q, p, array(v * e), q0, p0) + q1, p1, dE = leapfrog1_dE(q, p, array(v*e), E0[0]) n1 = int(log(u) + dE <= 0) s1 = int(log(u) + dE < Emax) return q1, p1, q1, p1, q1, n1, s1, min(1, exp(-dE)), 1 else: - qn, pn, qp, pp, q1, n1, s1, a1, na1 = buildtree( - H, q, p, u, v, j - 1, e, Emax, q0, p0) + qn, pn, qp, pp, q1, n1, s1, a1, na1 = buildtree(leapfrog1_dE, q, p, u, v, j - 1, e, Emax, q0, p0, E0) if s1 == 1: if v == -1: - qn, pn, _, _, q11, n11, s11, a11, na11 = buildtree( - H, qn, pn, u, v, j - 1, e, Emax, q0, p0) + qn, pn, _, _, q11, n11, s11, a11, na11 = buildtree(leapfrog1_dE, qn, pn, u, v, j - 1, e, Emax, q0, p0, E0) else: - _, _, qp, pp, q11, n11, s11, a11, na11 = buildtree( - H, qp, pp, u, v, j - 1, e, Emax, q0, p0) + _, _, qp, pp, q11, n11, s11, a11, na11 = buildtree(leapfrog1_dE, qp, pp, u, v, j - 1, e, Emax, q0, p0, E0) - if bern(n11 * 1. / (max(n1 + n11, 1))): + if bern(n11*1./(max(n1 + n11, 1))): q1 = q11 a1 = a1 + a11 @@ -179,7 +190,7 @@ def buildtree(H, q, p, u, v, j, e, Emax, q0, p0): return -def leapfrog1_dE(logp, vars, shared, pot, profile): +def leapfrog1_dE(H, q, profile): """Computes a theano function that computes one leapfrog step and the energy difference between the beginning and end of the trajectory. Parameters ---------- @@ -192,31 +203,22 @@ def leapfrog1_dE(logp, vars, shared, pot, profile): Returns ------- theano function which returns - q_new, p_new, delta_E + q_new, p_new, E """ - dlogp = gradient(logp, vars) - (logp, dlogp), q = join_nonshared_inputs([logp, dlogp], vars, shared) - logp = CallableTensor(logp) - dlogp = CallableTensor(dlogp) - - H = Hamiltonian(logp, dlogp, pot) - - p = tt.dvector('p') + p = theano.tensor.dvector('p') p.tag.test_value = q.tag.test_value - q0 = tt.dvector('q0') - q0.tag.test_value = q.tag.test_value - p0 = tt.dvector('p0') - p0.tag.test_value = p.tag.test_value - - e = tt.dscalar('e') + e = theano.tensor.dscalar('e') e.tag.test_value = 1 q1, p1 = leapfrog(H, q, p, 1, e) E = energy(H, q1, p1) - E0 = energy(H, q0, p0) + + E0 = theano.tensor.dscalar('E0') + E0.tag.test_value = 1 + dE = E - E0 - f = theano.function([q, p, e, q0, p0], [q1, p1, dE], profile=profile) + f = theano.function([q, p, e, E0], [q1, p1, dE], profile=profile) f.trust_input = True return f From 3513290bdcfd6fb351408318571906eff43600d8 Mon Sep 17 00:00:00 2001 From: Thomas Wiecki Date: Sat, 12 Nov 2016 01:19:24 +0900 Subject: [PATCH 2/7] STY Refactor creation of functions. --- pymc3/step_methods/nuts.py | 31 ++++++++++++++++++------------- 1 file changed, 18 insertions(+), 13 deletions(-) diff --git a/pymc3/step_methods/nuts.py b/pymc3/step_methods/nuts.py index 9152920054..567f18b406 100644 --- a/pymc3/step_methods/nuts.py +++ b/pymc3/step_methods/nuts.py @@ -92,35 +92,40 @@ def __init__(self, vars=None, scaling=None, step_scale=0.25, is_cov=False, state self.u = log(self.step_size*10) self.m = 1 - #self.H = Hamiltonian(model.fastlogp, model.fastdlogp(vars), self.potential) - shared = make_shared_replacements(vars, model) - dlogp = gradient(model.logpt, vars) - (logp, dlogp), q = join_nonshared_inputs([model.logpt, dlogp], vars, shared) - logp = CallableTensor(logp) - dlogp = CallableTensor(dlogp) + def create_hamiltonian(vars, shared, model): + dlogp = gradient(model.logpt, vars) + (logp, dlogp), q = join_nonshared_inputs([model.logpt, dlogp], vars, shared) + logp = CallableTensor(logp) + dlogp = CallableTensor(dlogp) + + return Hamiltonian(logp, dlogp, self.potential), q + def create_energy_func(q): + p = theano.tensor.dvector('p') + p.tag.test_value = q.tag.test_value + E0 = energy(self.H, q, p) + E0_func = theano.function([q, p], [E0]) + E0_func.trust_input = True - self.H = Hamiltonian(logp, dlogp, self.potential) + return E0_func - p = theano.tensor.dvector('p') - p.tag.test_value = q.tag.test_value - E0 = energy(self.H, q, p) - self.E0 = theano.function([q, p], [E0]) - self.E0.trust_input = True + self.H, q = create_hamiltonian(vars, shared, model) + self.compute_energy = create_energy_func(q) self.leapfrog1_dE = leapfrog1_dE(self.H, q, profile=profile) super(NUTS, self).__init__(vars, shared, **kwargs) + def astep(self, q0): leapfrog = self.leapfrog1_dE Emax = self.Emax e = self.step_size p0 = self.potential.random() - E0 = self.E0(q0, p0) + E0 = self.compute_energy(q0, p0) u = uniform() q = qn = qp = q0 From 321b05dd845a70c444da96f96809588ad5196619 Mon Sep 17 00:00:00 2001 From: Thomas Wiecki Date: Sat, 12 Nov 2016 01:20:26 +0900 Subject: [PATCH 3/7] DOC Adapt doc-string. --- pymc3/step_methods/nuts.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/pymc3/step_methods/nuts.py b/pymc3/step_methods/nuts.py index 567f18b406..b65c61232d 100644 --- a/pymc3/step_methods/nuts.py +++ b/pymc3/step_methods/nuts.py @@ -199,16 +199,14 @@ def leapfrog1_dE(H, q, profile): """Computes a theano function that computes one leapfrog step and the energy difference between the beginning and end of the trajectory. Parameters ---------- - logp : TensorVariable - vars : list of tensor variables - shared : list of shared variables not to compute leapfrog over - pot : quadpotential - porifle : Boolean + H : Hamiltonian + q : theano.tensor + proifle : Boolean Returns ------- theano function which returns - q_new, p_new, E + q_new, p_new, dE """ p = theano.tensor.dvector('p') p.tag.test_value = q.tag.test_value From 87fd45886aa09c06dde213c7be23ca1a6d62c26b Mon Sep 17 00:00:00 2001 From: Thomas Wiecki Date: Sat, 12 Nov 2016 01:21:29 +0900 Subject: [PATCH 4/7] STY Pep8. --- pymc3/step_methods/nuts.py | 45 +++++++++++++++++++++----------------- 1 file changed, 25 insertions(+), 20 deletions(-) diff --git a/pymc3/step_methods/nuts.py b/pymc3/step_methods/nuts.py index b65c61232d..9d5160de86 100644 --- a/pymc3/step_methods/nuts.py +++ b/pymc3/step_methods/nuts.py @@ -13,6 +13,7 @@ __all__ = ['NUTS'] + class NUTS(ArrayStepShared): """ Automatically tunes step size and adjust number of steps for good performance. @@ -31,7 +32,7 @@ def __init__(self, vars=None, scaling=None, step_scale=0.25, is_cov=False, state k=0.75, t0=10, model=None, - profile=False,**kwargs): + profile=False, **kwargs): """ Parameters ---------- @@ -67,14 +68,12 @@ def __init__(self, vars=None, scaling=None, step_scale=0.25, is_cov=False, state scaling = model.test_point if isinstance(scaling, dict): - scaling = guess_scaling(Point(scaling, model=model), model=model, vars = vars) - - + scaling = guess_scaling( + Point(scaling, model=model), model=model, vars=vars) n = scaling.shape[0] - self.step_size = step_scale / n**(1/4.) - + self.step_size = step_scale / n**(1 / 4.) self.potential = quad_potential(scaling, is_cov, as_cov=False) @@ -89,14 +88,15 @@ def __init__(self, vars=None, scaling=None, step_scale=0.25, is_cov=False, state self.k = k self.Hbar = 0 - self.u = log(self.step_size*10) + self.u = log(self.step_size * 10) self.m = 1 shared = make_shared_replacements(vars, model) def create_hamiltonian(vars, shared, model): dlogp = gradient(model.logpt, vars) - (logp, dlogp), q = join_nonshared_inputs([model.logpt, dlogp], vars, shared) + (logp, dlogp), q = join_nonshared_inputs( + [model.logpt, dlogp], vars, shared) logp = CallableTensor(logp) dlogp = CallableTensor(dlogp) @@ -118,7 +118,6 @@ def create_energy_func(q): super(NUTS, self).__init__(vars, shared, **kwargs) - def astep(self, q0): leapfrog = self.leapfrog1_dE Emax = self.Emax @@ -137,11 +136,13 @@ def astep(self, q0): v = bern(.5) * 2 - 1 if v == -1: - qn, pn, _, _, q1, n1, s1, a, na = buildtree(leapfrog, qn, pn, u, v, j, e, Emax, q0, p0, E0) + qn, pn, _, _, q1, n1, s1, a, na = buildtree( + leapfrog, qn, pn, u, v, j, e, Emax, q0, p0, E0) else: - _, _, qp, pp, q1, n1, s1, a, na = buildtree(leapfrog, qp, pp, u, v, j, e, Emax, q0, p0, E0) + _, _, qp, pp, q1, n1, s1, a, na = buildtree( + leapfrog, qp, pp, u, v, j, e, Emax, q0, p0, E0) - if s1 == 1 and bern(min(1, n1*1./n)): + if s1 == 1 and bern(min(1, n1 * 1. / n)): q = q1 n = n + n1 @@ -152,10 +153,11 @@ def astep(self, q0): p = -p - w = 1./(self.m+self.t0) - self.Hbar = (1 - w) * self.Hbar + w*(self.target_accept - a*1./na) + w = 1. / (self.m + self.t0) + self.Hbar = (1 - w) * self.Hbar + w * \ + (self.target_accept - a * 1. / na) - self.step_size = exp(self.u - (self.m**.5/self.gamma)*self.Hbar) + self.step_size = exp(self.u - (self.m**.5 / self.gamma) * self.Hbar) self.m += 1 return q @@ -169,20 +171,23 @@ def competence(var): def buildtree(leapfrog1_dE, q, p, u, v, j, e, Emax, q0, p0, E0): if j == 0: - q1, p1, dE = leapfrog1_dE(q, p, array(v*e), E0[0]) + q1, p1, dE = leapfrog1_dE(q, p, array(v * e), E0[0]) n1 = int(log(u) + dE <= 0) s1 = int(log(u) + dE < Emax) return q1, p1, q1, p1, q1, n1, s1, min(1, exp(-dE)), 1 else: - qn, pn, qp, pp, q1, n1, s1, a1, na1 = buildtree(leapfrog1_dE, q, p, u, v, j - 1, e, Emax, q0, p0, E0) + qn, pn, qp, pp, q1, n1, s1, a1, na1 = buildtree( + leapfrog1_dE, q, p, u, v, j - 1, e, Emax, q0, p0, E0) if s1 == 1: if v == -1: - qn, pn, _, _, q11, n11, s11, a11, na11 = buildtree(leapfrog1_dE, qn, pn, u, v, j - 1, e, Emax, q0, p0, E0) + qn, pn, _, _, q11, n11, s11, a11, na11 = buildtree( + leapfrog1_dE, qn, pn, u, v, j - 1, e, Emax, q0, p0, E0) else: - _, _, qp, pp, q11, n11, s11, a11, na11 = buildtree(leapfrog1_dE, qp, pp, u, v, j - 1, e, Emax, q0, p0, E0) + _, _, qp, pp, q11, n11, s11, a11, na11 = buildtree( + leapfrog1_dE, qp, pp, u, v, j - 1, e, Emax, q0, p0, E0) - if bern(n11*1./(max(n1 + n11, 1))): + if bern(n11 * 1. / (max(n1 + n11, 1))): q1 = q11 a1 = a1 + a11 From a0dd97850d30daf5b01a0a5a0736b493789e42ff Mon Sep 17 00:00:00 2001 From: Thomas Wiecki Date: Sat, 12 Nov 2016 01:23:59 +0900 Subject: [PATCH 5/7] STY Only return scalar --- pymc3/step_methods/nuts.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pymc3/step_methods/nuts.py b/pymc3/step_methods/nuts.py index 9d5160de86..691179762e 100644 --- a/pymc3/step_methods/nuts.py +++ b/pymc3/step_methods/nuts.py @@ -106,7 +106,7 @@ def create_energy_func(q): p = theano.tensor.dvector('p') p.tag.test_value = q.tag.test_value E0 = energy(self.H, q, p) - E0_func = theano.function([q, p], [E0]) + E0_func = theano.function([q, p], E0) E0_func.trust_input = True return E0_func @@ -171,7 +171,7 @@ def competence(var): def buildtree(leapfrog1_dE, q, p, u, v, j, e, Emax, q0, p0, E0): if j == 0: - q1, p1, dE = leapfrog1_dE(q, p, array(v * e), E0[0]) + q1, p1, dE = leapfrog1_dE(q, p, array(v * e), E0) n1 = int(log(u) + dE <= 0) s1 = int(log(u) + dE < Emax) From 830588c3ea32f3209cf6acc166a7c25c3ee27495 Mon Sep 17 00:00:00 2001 From: Thomas Wiecki Date: Sat, 12 Nov 2016 01:26:03 +0900 Subject: [PATCH 6/7] DOC Typo --- pymc3/step_methods/nuts.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc3/step_methods/nuts.py b/pymc3/step_methods/nuts.py index 691179762e..36b68d9855 100644 --- a/pymc3/step_methods/nuts.py +++ b/pymc3/step_methods/nuts.py @@ -206,7 +206,7 @@ def leapfrog1_dE(H, q, profile): ---------- H : Hamiltonian q : theano.tensor - proifle : Boolean + profile : Boolean Returns ------- From 4f01f7fea3c77faffbc627d047ab007784ba3802 Mon Sep 17 00:00:00 2001 From: Thomas Wiecki Date: Sat, 12 Nov 2016 10:53:18 +0100 Subject: [PATCH 7/7] MAINT Move theano.tensor to tt. Remove obsolete p0 and q0 args from build_tree. --- pymc3/step_methods/nuts.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/pymc3/step_methods/nuts.py b/pymc3/step_methods/nuts.py index 36b68d9855..5121681256 100644 --- a/pymc3/step_methods/nuts.py +++ b/pymc3/step_methods/nuts.py @@ -103,7 +103,7 @@ def create_hamiltonian(vars, shared, model): return Hamiltonian(logp, dlogp, self.potential), q def create_energy_func(q): - p = theano.tensor.dvector('p') + p = tt.dvector('p') p.tag.test_value = q.tag.test_value E0 = energy(self.H, q, p) E0_func = theano.function([q, p], E0) @@ -137,10 +137,10 @@ def astep(self, q0): if v == -1: qn, pn, _, _, q1, n1, s1, a, na = buildtree( - leapfrog, qn, pn, u, v, j, e, Emax, q0, p0, E0) + leapfrog, qn, pn, u, v, j, e, Emax, E0) else: _, _, qp, pp, q1, n1, s1, a, na = buildtree( - leapfrog, qp, pp, u, v, j, e, Emax, q0, p0, E0) + leapfrog, qp, pp, u, v, j, e, Emax, E0) if s1 == 1 and bern(min(1, n1 * 1. / n)): q = q1 @@ -169,7 +169,7 @@ def competence(var): return Competence.INCOMPATIBLE -def buildtree(leapfrog1_dE, q, p, u, v, j, e, Emax, q0, p0, E0): +def buildtree(leapfrog1_dE, q, p, u, v, j, e, Emax, E0): if j == 0: q1, p1, dE = leapfrog1_dE(q, p, array(v * e), E0) @@ -178,14 +178,14 @@ def buildtree(leapfrog1_dE, q, p, u, v, j, e, Emax, q0, p0, E0): return q1, p1, q1, p1, q1, n1, s1, min(1, exp(-dE)), 1 else: qn, pn, qp, pp, q1, n1, s1, a1, na1 = buildtree( - leapfrog1_dE, q, p, u, v, j - 1, e, Emax, q0, p0, E0) + leapfrog1_dE, q, p, u, v, j - 1, e, Emax, E0) if s1 == 1: if v == -1: qn, pn, _, _, q11, n11, s11, a11, na11 = buildtree( - leapfrog1_dE, qn, pn, u, v, j - 1, e, Emax, q0, p0, E0) + leapfrog1_dE, qn, pn, u, v, j - 1, e, Emax, E0) else: _, _, qp, pp, q11, n11, s11, a11, na11 = buildtree( - leapfrog1_dE, qp, pp, u, v, j - 1, e, Emax, q0, p0, E0) + leapfrog1_dE, qp, pp, u, v, j - 1, e, Emax, E0) if bern(n11 * 1. / (max(n1 + n11, 1))): q1 = q11 @@ -213,16 +213,16 @@ def leapfrog1_dE(H, q, profile): theano function which returns q_new, p_new, dE """ - p = theano.tensor.dvector('p') + p = tt.dvector('p') p.tag.test_value = q.tag.test_value - e = theano.tensor.dscalar('e') + e = tt.dscalar('e') e.tag.test_value = 1 q1, p1 = leapfrog(H, q, p, 1, e) E = energy(H, q1, p1) - E0 = theano.tensor.dscalar('E0') + E0 = tt.dscalar('E0') E0.tag.test_value = 1 dE = E - E0