diff --git a/pymc3/step_methods/hmc/base_hmc.py b/pymc3/step_methods/hmc/base_hmc.py index 23c0a8d67f..017a65e248 100644 --- a/pymc3/step_methods/hmc/base_hmc.py +++ b/pymc3/step_methods/hmc/base_hmc.py @@ -1,19 +1,22 @@ +import numpy as np + from pymc3.model import modelcontext, Point from pymc3.step_methods import arraystep -from .quadpotential import quad_potential, QuadPotentialDiagAdapt from pymc3.step_methods.hmc import integration from pymc3.theanof import inputvars, floatX from pymc3.tuning import guess_scaling -import numpy as np +from .quadpotential import quad_potential, QuadPotentialDiagAdapt class BaseHMC(arraystep.GradientSharedStep): + """Superclass to implement Hamiltonian/hybrid monte carlo.""" + default_blocked = True def __init__(self, vars=None, scaling=None, step_scale=0.25, is_cov=False, model=None, blocked=True, potential=None, integrator="leapfrog", dtype=None, **theano_kwargs): - """Superclass to implement Hamiltonian/hybrid monte carlo + """Set up Hamiltonian samplers with common structures. Parameters ---------- diff --git a/pymc3/step_methods/hmc/hmc.py b/pymc3/step_methods/hmc/hmc.py index 3d46cee771..539475db1d 100644 --- a/pymc3/step_methods/hmc/hmc.py +++ b/pymc3/step_methods/hmc/hmc.py @@ -1,13 +1,9 @@ -''' -Created on Mar 7, 2011 +import numpy as np -@author: johnsalvatier -''' from ..arraystep import metrop_select, Competence from .base_hmc import BaseHMC from pymc3.vartypes import discrete_types from pymc3.theanof import floatX -import numpy as np __all__ = ['HamiltonianMC'] @@ -22,11 +18,17 @@ def unif(step_size, elow=.85, ehigh=1.15): class HamiltonianMC(BaseHMC): + R"""A sampler for continuous variables based on Hamiltonian mechanics. + + See NUTS sampler for automatically tuned stopping time and step size scaling. + """ + name = 'hmc' default_blocked = True def __init__(self, vars=None, path_length=2., step_rand=unif, **kwargs): - """ + """Set up the Hamiltonian Monte Carlo sampler. + Parameters ---------- vars : list of theano variables @@ -57,6 +59,7 @@ def __init__(self, vars=None, path_length=2., step_rand=unif, **kwargs): self.step_rand = step_rand def astep(self, q0): + """Perform a single HMC iteration.""" e = floatX(self.step_rand(self.step_size)) n_steps = int(self.path_length / e) @@ -76,6 +79,7 @@ def astep(self, q0): @staticmethod def competence(var): + """Check how appropriate this class is for sampling a random variable.""" if var.dtype in discrete_types: return Competence.INCOMPATIBLE return Competence.COMPATIBLE diff --git a/pymc3/step_methods/hmc/integration.py b/pymc3/step_methods/hmc/integration.py index 04605071b1..8ae76c5602 100644 --- a/pymc3/step_methods/hmc/integration.py +++ b/pymc3/step_methods/hmc/integration.py @@ -8,7 +8,10 @@ class CpuLeapfrogIntegrator(object): + """Optimized leapfrog integration using numpy.""" + def __init__(self, ndim, potential, logp_dlogp_func): + """Leapfrog integrator using CPU.""" self._ndim = ndim self._potential = potential self._logp_dlogp_func = logp_dlogp_func @@ -19,6 +22,7 @@ def __init__(self, ndim, potential, logp_dlogp_func): % (self._potential.dtype, self._dtype)) def compute_state(self, q, p): + """Compute Hamiltonian functions using a position and momentum.""" if q.dtype != self._dtype or p.dtype != self._dtype: raise ValueError('Invalid dtype. Must be %s' % self._dtype) logp, dlogp = self._logp_dlogp_func(q) @@ -28,6 +32,23 @@ def compute_state(self, q, p): return State(q, p, v, dlogp, energy) def step(self, epsilon, state, out=None): + """Leapfrog integrator step. + + Half a momentum update, full position update, half momentum update. + + Parameters + ---------- + epsilon: float, > 0 + step scale + state: State namedtuple, + current position data + out: (optional) State namedtuple, + preallocated arrays to write to in place + + Returns + ------- + None if `out` is provided, else a State namedtuple + """ pot = self._potential axpy = linalg.blas.get_blas_funcs('axpy', dtype=self._dtype) diff --git a/pymc3/step_methods/hmc/nuts.py b/pymc3/step_methods/hmc/nuts.py index 73af656c8a..aeab733069 100644 --- a/pymc3/step_methods/hmc/nuts.py +++ b/pymc3/step_methods/hmc/nuts.py @@ -1,17 +1,17 @@ from collections import namedtuple import warnings +import numpy as np +import numpy.random as nr +from scipy import stats, linalg +import six + from ..arraystep import Competence from pymc3.exceptions import SamplingError from .base_hmc import BaseHMC from pymc3.theanof import floatX from pymc3.vartypes import continuous_types -import numpy as np -import numpy.random as nr -from scipy import stats, linalg -import six - __all__ = ['NUTS'] @@ -28,7 +28,7 @@ class NUTS(BaseHMC): sample. A detailed description can be found at [1], "Algorithm 6: Efficient No-U-Turn Sampler with Dual Averaging". - Nuts provides a number of statistics that can be accessed with + NUTS provides a number of statistics that can be accessed with `trace.get_sampler_stats`: - `mean_tree_accept`: The mean acceptance probability for the tree @@ -70,6 +70,7 @@ class NUTS(BaseHMC): .. [1] Hoffman, Matthew D., & Gelman, Andrew. (2011). The No-U-Turn Sampler: Adaptively Setting Path Lengths in Hamiltonian Monte Carlo. """ + name = 'nuts' default_blocked = True @@ -92,7 +93,8 @@ def __init__(self, vars=None, Emax=1000, target_accept=0.8, max_treedepth=10, on_error='summary', early_max_treedepth=8, **kwargs): - R""" + R"""Set up the No-U-Turn sampler. + Parameters ---------- vars : list of Theano variables, default all continuous vars @@ -171,6 +173,7 @@ def __init__(self, vars=None, Emax=1000, target_accept=0.8, self.report = NutsReport(on_error, max_treedepth, target_accept) def astep(self, q0): + """Perform a single NUTS iteration.""" p0 = self.potential.random() start = self.integrator.compute_state(q0, p0) @@ -229,6 +232,7 @@ def astep(self, q0): @staticmethod def competence(var): + """Check how appropriate this class is for sampling a random variable.""" if var.dtype in continuous_types: return Competence.IDEAL return Competence.INCOMPATIBLE diff --git a/pymc3/step_methods/hmc/quadpotential.py b/pymc3/step_methods/hmc/quadpotential.py index 2e01ccb9b9..1041f5e168 100644 --- a/pymc3/step_methods/hmc/quadpotential.py +++ b/pymc3/step_methods/hmc/quadpotential.py @@ -13,6 +13,8 @@ def quad_potential(C, is_cov): """ + Compute a QuadPotential object from a scaling matrix. + Parameters ---------- C : arraylike, 0 <= ndim <= 2 @@ -28,10 +30,10 @@ def quad_potential(C, is_cov): if issparse(C): if not chol_available: raise ImportError("Sparse mass matrices require scikits.sparse") - if is_cov: + elif is_cov: return QuadPotentialSparse(C) else: - raise ValueError("Sparse precission matrices are not supported") + raise ValueError("Sparse precision matrices are not supported") partial_check_positive_definite(C) if C.ndim == 1: @@ -47,7 +49,7 @@ def quad_potential(C, is_cov): def partial_check_positive_definite(C): - """Simple but partial check for Positive Definiteness""" + """Make a simple but partial check for Positive Definiteness.""" if C.ndim == 1: d = C else: @@ -72,6 +74,7 @@ def __str__(self): class QuadPotential(object): def velocity(self, x, out=None): + """Compute the current velocity at a position in parameter space.""" raise NotImplementedError('Abstract method') def energy(self, x, velocity=None): @@ -93,13 +96,16 @@ def adapt(self, sample, grad): def isquadpotential(value): + """Check whether an object might be a QuadPotential object.""" return isinstance(value, QuadPotential) class QuadPotentialDiagAdapt(QuadPotential): """Adapt a diagonal mass matrix from the sample variances.""" + def __init__(self, n, initial_mean, initial_diag=None, initial_weight=0, adaptation_window=100, dtype=None): + """Set up a diagonal mass matrix.""" if initial_diag is not None and initial_diag.ndim != 1: raise ValueError('Initial diagonal must be one-dimensional.') if initial_mean.ndim != 1: @@ -130,18 +136,22 @@ def __init__(self, n, initial_mean, initial_diag=None, initial_weight=0, self.adaptation_window = adaptation_window def velocity(self, x, out=None): + """Compute the current velocity at a position in parameter space.""" return np.multiply(self._var, x, out=out) def energy(self, x, velocity=None): + """Compute kinetic energy at a position in parameter space.""" if velocity is not None: return 0.5 * x.dot(velocity) return 0.5 * x.dot(self._var * x) def velocity_energy(self, x, v_out): + """Compute velocity and return kinetic energy at a position in parameter space.""" self.velocity(x, out=v_out) return 0.5 * np.dot(x, v_out) def random(self): + """Draw random value from QuadPotential.""" vals = normal(size=self._n).astype(self.dtype) return self._inv_stds * vals @@ -152,6 +162,7 @@ def _update_from_weightvar(self, weightvar): self._var_theano.set_value(self._var) def adapt(self, sample, grad): + """Inform the potential about a new sample during tuning.""" window = self.adaptation_window self._foreground_var.add_sample(sample, weight=1) @@ -170,6 +181,7 @@ class QuadPotentialDiagAdaptGrad(QuadPotentialDiagAdapt): This is experimental, and may be removed without prior deprication. """ + def __init__(self, *args, **kwargs): super(QuadPotentialDiagAdaptGrad, self).__init__(*args, **kwargs) self._grads1 = np.zeros(self._n, dtype=self.dtype) @@ -184,6 +196,7 @@ def _update(self, var): self._var_theano.set_value(self._var) def adapt(self, sample, grad): + """Inform the potential about a new sample during tuning.""" self._grads1[:] += grad ** 2 self._grads2[:] += grad ** 2 self._ngrads1 += 1 @@ -248,7 +261,16 @@ def current_mean(self): class QuadPotentialDiag(QuadPotential): + """Quad potential using a diagonal covariance matrix.""" + def __init__(self, v, dtype=None): + """Use a vector to represent a diagonal matrix for a covariance matrix. + + Parameters + ---------- + v : vector, 0 <= ndim <= 1 + Diagonal of covariance matrix for the potential vector + """ if dtype is None: dtype = theano.config.floatX self.dtype = dtype @@ -260,55 +282,79 @@ def __init__(self, v, dtype=None): self.v = v def velocity(self, x, out=None): + """Compute the current velocity at a position in parameter space.""" if out is not None: np.multiply(x, self.v, out=out) return return self.v * x def random(self): + """Draw random value from QuadPotential.""" return floatX(normal(size=self.s.shape)) * self.inv_s def energy(self, x, velocity=None): + """Compute kinetic energy at a position in parameter space.""" if velocity is not None: return 0.5 * np.dot(x, velocity) return .5 * x.dot(self.v * x) def velocity_energy(self, x, v_out): + """Compute velocity and return kinetic energy at a position in parameter space.""" np.multiply(x, self.v, out=v_out) return 0.5 * np.dot(x, v_out) class QuadPotentialFullInv(QuadPotential): + """QuadPotential object for Hamiltonian calculations using inverse of covariance matrix.""" def __init__(self, A, dtype=None): + """Compute the lower cholesky decomposition of the potential. + + Parameters + ---------- + A : matrix, ndim = 2 + Inverse of covariance matrix for the potential vector + """ if dtype is None: dtype = theano.config.floatX self.dtype = dtype self.L = floatX(scipy.linalg.cholesky(A, lower=True)) def velocity(self, x, out=None): + """Compute the current velocity at a position in parameter space.""" vel = scipy.linalg.cho_solve((self.L, True), x) if out is None: return vel out[:] = vel def random(self): + """Draw random value from QuadPotential.""" n = floatX(normal(size=self.L.shape[0])) return np.dot(self.L, n) def energy(self, x, velocity=None): + """Compute kinetic energy at a position in parameter space.""" if velocity is None: velocity = self.velocity(x) return .5 * x.dot(velocity) def velocity_energy(self, x, v_out): + """Compute velocity and return kinetic energy at a position in parameter space.""" self.velocity(x, out=v_out) return 0.5 * np.dot(x, v_out) class QuadPotentialFull(QuadPotential): + """Basic QuadPotential object for Hamiltonian calculations.""" def __init__(self, A, dtype=None): + """Compute the lower cholesky decomposition of the potential. + + Parameters + ---------- + A : matrix, ndim = 2 + scaling matrix for the potential vector + """ if dtype is None: dtype = theano.config.floatX self.dtype = dtype @@ -316,18 +362,22 @@ def __init__(self, A, dtype=None): self.L = scipy.linalg.cholesky(A, lower=True) def velocity(self, x, out=None): + """Compute the current velocity at a position in parameter space.""" return np.dot(self.A, x, out=out) def random(self): + """Draw random value from QuadPotential.""" n = floatX(normal(size=self.L.shape[0])) return scipy.linalg.solve_triangular(self.L.T, n) def energy(self, x, velocity=None): + """Compute kinetic energy at a position in parameter space.""" if velocity is None: velocity = self.velocity(x) return .5 * x.dot(velocity) def velocity_energy(self, x, v_out): + """Compute velocity and return kinetic energy at a position in parameter space.""" self.velocity(x, out=v_out) return 0.5 * np.dot(x, v_out) @@ -347,16 +397,25 @@ def velocity_energy(self, x, v_out): class QuadPotentialSparse(QuadPotential): def __init__(self, A): + """Compute a sparse cholesky decomposition of the potential. + + Parameters + ---------- + A : matrix, ndim = 2 + scaling matrix for the potential vector + """ self.A = A self.size = A.shape[0] self.factor = factor = cholmod.cholesky(A) self.d_sqrt = np.sqrt(factor.D()) def velocity(self, x): + """Compute the current velocity at a position in parameter space.""" A = theano.sparse.as_sparse(self.A) return theano.sparse.dot(A, x) def random(self): + """Draw random value from QuadPotential.""" n = floatX(normal(size=self.size)) n /= self.d_sqrt n = self.factor.solve_Lt(n) @@ -364,4 +423,5 @@ def random(self): return n def energy(self, x): + """Compute kinetic energy at a position in parameter space.""" return 0.5 * x.T.dot(self.velocity(x)) diff --git a/pymc3/step_methods/hmc/trajectory.py b/pymc3/step_methods/hmc/trajectory.py new file mode 100644 index 0000000000..70efcb2259 --- /dev/null +++ b/pymc3/step_methods/hmc/trajectory.py @@ -0,0 +1,300 @@ +from collections import namedtuple + +import theano +import theano.tensor as tt +import numpy as np + +from pymc3.theanof import join_nonshared_inputs, gradient, CallableTensor, floatX + +Hamiltonian = namedtuple("Hamiltonian", "logp, dlogp, pot") + + +def _theano_hamiltonian(model_vars, shared, logpt, potential): + """Create a Hamiltonian with shared inputs. + + Parameters + ---------- + model_vars : array of variables to be sampled + shared : theano tensors that are already shared + logpt : model log probability + potential : hamiltonian potential + + Returns + ------- + Hamiltonian : namedtuple with log pdf, gradient of log pdf, and potential functions + q : Initial position for Hamiltonian Monte Carlo + dlogp_func: theano function that computes the gradient of a log pdf at a point + """ + dlogp = gradient(logpt, model_vars) + (logp, dlogp), q = join_nonshared_inputs([logpt, dlogp], model_vars, shared) + dlogp_func = theano.function(inputs=[q], outputs=dlogp) + dlogp_func.trust_input = True + logp = CallableTensor(logp) + dlogp = CallableTensor(dlogp) + return Hamiltonian(logp, dlogp, potential), q, dlogp_func + + +def _theano_energy_function(H, q, **theano_kwargs): + """Create a theano function for computing energy at a point in parameter space. + + Parameters + ---------- + H : Hamiltonian namedtuple + q : theano variable, starting position + theano_kwargs : passed to theano.function + + Returns + ------- + energy_function : theano function that computes the energy at a point (p, q) in phase space + p : Starting momentum variable. + """ + p = tt.vector('p') + p.tag.test_value = q.tag.test_value + total_energy = H.pot.energy(p) - H.logp(q) + energy_function = theano.function(inputs=[q, p], outputs=total_energy, **theano_kwargs) + energy_function.trust_input = True + + return energy_function, p + + +def _theano_velocity_function(H, p, **theano_kwargs): + v = H.pot.velocity(p) + velocity_function = theano.function(inputs=[p], outputs=v, **theano_kwargs) + velocity_function.trust_input = True + return velocity_function + + +def _theano_leapfrog_integrator(H, q, p, **theano_kwargs): + """Compute a theano function that computes one leapfrog step. + + Returns not only the new positiona and momentum, but also the energy. + + Parameters + ---------- + H : Hamiltonian + q : theano.tensor + p : theano.tensor + theano_kwargs : passed to theano.function + + Returns + ------- + theano function which returns + q_new, p_new, energy_new + """ + epsilon = tt.scalar('epsilon') + epsilon.tag.test_value = 1. + + n_steps = tt.iscalar('n_steps') + n_steps.tag.test_value = 2 + + q_new, p_new = leapfrog(H, q, p, epsilon, n_steps) + energy_new = energy(H, q_new, p_new) + + f = theano.function([q, p, epsilon, n_steps], [q_new, p_new, energy_new], **theano_kwargs) + f.trust_input = True + return f + + +def get_theano_hamiltonian_functions(model_vars, shared, logpt, potential, + use_single_leapfrog=False, + integrator="leapfrog", **theano_kwargs): + """Construct theano functions for the Hamiltonian, energy, and leapfrog integrator. + + Parameters + ---------- + model_vars : array of variables to be sampled + shared : theano tensors that are already shared + logpt : model log probability + potential : Hamiltonian potential + theano_kwargs : dictionary of keyword arguments to pass to theano functions + use_single_leapfrog : bool + if only 1 integration step is done at a time (as in NUTS), this + provides a ~2x speedup + integrator : str + Integration scheme to use. One of "leapfog", "two-stage", or + "three-stage". + + Returns + ------- + H : Hamiltonian namedtuple + energy_function : theano function computing energy at a point in phase space + leapfrog_integrator : theano function integrating the Hamiltonian from a point in phase space + theano_variables : dictionary of variables used in the computation graph which may be useful + """ + H, q, dlogp = _theano_hamiltonian(model_vars, shared, logpt, potential) + energy_function, p = _theano_energy_function(H, q, **theano_kwargs) + velocity_function = _theano_velocity_function(H, p, **theano_kwargs) + if use_single_leapfrog: + try: + _theano_integrator = INTEGRATORS_SINGLE[integrator] + except KeyError: + raise ValueError("Unknown integrator: %s" % integrator) + integrator = _theano_integrator(H, q, p, H.dlogp(q), **theano_kwargs) + else: + if integrator != "leapfrog": + raise ValueError("Only leapfrog is supported") + integrator = _theano_leapfrog_integrator(H, q, p, **theano_kwargs) + return H, energy_function, velocity_function, integrator, dlogp + + +def energy(H, q, p): + """Compute the total energy for the Hamiltonian at a given position/momentum.""" + return H.pot.energy(p) - H.logp(q) + + +def leapfrog(H, q, p, epsilon, n_steps): + r"""Leapfrog integrator. + + Estimates `p(t)` and `q(t)` at time :math:`t = n \cdot e`, by integrating the + Hamiltonian equations + + .. math:: + + \frac{dq_i}{dt} = \frac{\partial H}{\partial p_i} + + \frac{dp_i}{dt} = \frac{\partial H}{\partial q_i} + + with :math:`p(0) = p`, :math:`q(0) = q` + + Parameters + ---------- + H : Hamiltonian instance. + Tuple of `logp, dlogp, potential`. + q : Theano.tensor + initial position vector + p : Theano.tensor + initial momentum vector + epsilon : float, step size + n_steps : int, number of iterations + + Returns + ------- + position : Theano.tensor + position estimate at time :math:`n \cdot e`. + momentum : Theano.tensor + momentum estimate at time :math:`n \cdot e`. + """ + def full_update(p, q): + p = p + epsilon * H.dlogp(q) + q += epsilon * H.pot.velocity(p) + return p, q + # This first line can't be +=, possibly because of theano + p = p + 0.5 * epsilon * H.dlogp(q) # half momentum update + q += epsilon * H.pot.velocity(p) # full position update + if tt.gt(n_steps, 1): + (p_seq, q_seq), _ = theano.scan(full_update, outputs_info=[p, q], n_steps=n_steps - 1) + p, q = p_seq[-1], q_seq[-1] + p += 0.5 * epsilon * H.dlogp(q) # half momentum update + return q, p + + +def _theano_single_threestage(H, q, p, q_grad, **theano_kwargs): + """Perform a single step of a third order symplectic integration scheme. + + References + ---------- + Blanes, Sergio, Fernando Casas, and J. M. Sanz-Serna. "Numerical + Integrators for the Hybrid Monte Carlo Method." SIAM Journal on + Scientific Computing 36, no. 4 (January 2014): A1556-80. + doi:10.1137/130932740. + + Mannseth, Janne, Tore Selland Kleppe, and Hans J. Skaug. "On the + Application of Higher Order Symplectic Integrators in + Hamiltonian Monte Carlo." arXiv:1608.07048 [Stat], + August 25, 2016. http://arxiv.org/abs/1608.07048. + """ + epsilon = tt.scalar('epsilon') + epsilon.tag.test_value = 1. + + a = 12127897.0 / 102017882 + b = 4271554.0 / 14421423 + + # q_{a\epsilon} + p_ae = p + floatX(a) * epsilon * q_grad + q_be = q + floatX(b) * epsilon * H.pot.velocity(p_ae) + + # q_{\epsilon / 2} + p_e2 = p_ae + floatX(0.5 - a) * epsilon * H.dlogp(q_be) + + # p_{(1-b)\epsilon} + q_1be = q_be + floatX(1 - 2 * b) * epsilon * H.pot.velocity(p_e2) + p_1ae = p_e2 + floatX(0.5 - a) * epsilon * H.dlogp(q_1be) + + q_e = q_1be + floatX(b) * epsilon * H.pot.velocity(p_1ae) + grad_e = H.dlogp(q_e) + p_e = p_1ae + floatX(a) * epsilon * grad_e + v_e = H.pot.velocity(p_e) + + new_energy = energy(H, q_e, p_e) + + f = theano.function(inputs=[q, p, q_grad, epsilon], + outputs=[q_e, p_e, v_e, grad_e, new_energy], + **theano_kwargs) + f.trust_input = True + return f + + +def _theano_single_twostage(H, q, p, q_grad, **theano_kwargs): + """Perform a single step of a second order symplectic integration scheme. + + References + ---------- + Blanes, Sergio, Fernando Casas, and J. M. Sanz-Serna. "Numerical + Integrators for the Hybrid Monte Carlo Method." SIAM Journal on + Scientific Computing 36, no. 4 (January 2014): A1556-80. + doi:10.1137/130932740. + + Mannseth, Janne, Tore Selland Kleppe, and Hans J. Skaug. "On the + Application of Higher Order Symplectic Integrators in + Hamiltonian Monte Carlo." arXiv:1608.07048 [Stat], + August 25, 2016. http://arxiv.org/abs/1608.07048. + """ + epsilon = tt.scalar('epsilon') + epsilon.tag.test_value = 1. + + a = floatX((3 - np.sqrt(3)) / 6) + + p_ae = p + a * epsilon * q_grad + q_e2 = q + epsilon / 2 * H.pot.velocity(p_ae) + p_1ae = p_ae + (1 - 2 * a) * epsilon * H.dlogp(q_e2) + q_e = q_e2 + epsilon / 2 * H.pot.velocity(p_1ae) + grad_e = H.dlogp(q_e) + p_e = p_1ae + a * epsilon * grad_e + v_e = H.pot.velocity(p_e) + + new_energy = energy(H, q_e, p_e) + f = theano.function(inputs=[q, p, q_grad, epsilon], + outputs=[q_e, p_e, v_e, grad_e, new_energy], + **theano_kwargs) + f.trust_input = True + return f + + +def _theano_single_leapfrog(H, q, p, q_grad, **theano_kwargs): + """Leapfrog integrator for a single step. + + See above for documentation. This is optimized for the case where only a single step is + needed, in case of, for example, a recursive algorithm. + """ + epsilon = tt.scalar('epsilon') + epsilon.tag.test_value = 1. + + p_new = p + 0.5 * epsilon * q_grad # half momentum update + q_new = q + epsilon * H.pot.velocity(p_new) # full position update + q_new_grad = H.dlogp(q_new) + p_new += 0.5 * epsilon * q_new_grad # half momentum update + energy_new = energy(H, q_new, p_new) + v_new = H.pot.velocity(p_new) + + f = theano.function(inputs=[q, p, q_grad, epsilon], + outputs=[q_new, p_new, v_new, q_new_grad, energy_new], + **theano_kwargs) + f.trust_input = True + return f + + +INTEGRATORS_SINGLE = { + 'leapfrog': _theano_single_leapfrog, + 'two-stage': _theano_single_twostage, + 'three-stage': _theano_single_threestage, +} diff --git a/setup.cfg b/setup.cfg index 61f83f75e3..cdd217776e 100644 --- a/setup.cfg +++ b/setup.cfg @@ -3,3 +3,7 @@ testpaths = pymc3/tests [coverage:run] omit = *examples* + +[pydocstyle] +add-ignore = D100,D104 +convention = numpy