1#   Copyright 2020 The PyMC Developers
2#
3#   Licensed under the Apache License, Version 2.0 (the "License");
4#   you may not use this file except in compliance with the License.
5#   You may obtain a copy of the License at
6#
7#       http://www.apache.org/licenses/LICENSE-2.0
8#
9#   Unless required by applicable law or agreed to in writing, software
10#   distributed under the License is distributed on an "AS IS" BASIS,
11#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12#   See the License for the specific language governing permissions and
13#   limitations under the License.
14
15import numpy as np
16import numpy.random as nr
17import theano.tensor as tt
18
19from pymc3.distributions import draw_values
20from pymc3.model import modelcontext
21from pymc3.step_methods.arraystep import ArrayStep, Competence
22from pymc3.theanof import inputvars
23
24__all__ = ["EllipticalSlice"]
25
26
27def get_chol(cov, chol):
28    """Get Cholesky decomposition of the prior covariance.
29
30    Ensure that exactly one of the prior covariance or Cholesky
31    decomposition is passed. If the prior covariance was passed, then
32    return its Cholesky decomposition.
33
34    Parameters
35    ----------
36    cov: array, optional
37        Covariance matrix of the multivariate Gaussian prior.
38    chol: array, optional
39        Cholesky decomposition of the covariance matrix of the
40        multivariate Gaussian prior.
41    """
42
43    if len([i for i in [cov, chol] if i is not None]) != 1:
44        raise ValueError("Must pass exactly one of cov or chol")
45
46    if cov is not None:
47        chol = tt.slinalg.cholesky(cov)
48    return chol
49
50
51class EllipticalSlice(ArrayStep):
52    """Multivariate elliptical slice sampler step.
53
54    Elliptical slice sampling (ESS) [1]_ is a variant of slice sampling
55    that allows sampling from distributions with multivariate Gaussian
56    prior and arbitrary likelihood. It is generally about as fast as
57    regular slice sampling, mixes well even when the prior covariance
58    might otherwise induce a strong dependence between samples, and
59    does not depend on any tuning parameters.
60
61    The Gaussian prior is assumed to have zero mean.
62
63    Parameters
64    ----------
65    vars: list
66        List of variables for sampler.
67    prior_cov: array, optional
68        Covariance matrix of the multivariate Gaussian prior.
69    prior_chol: array, optional
70        Cholesky decomposition of the covariance matrix of the
71        multivariate Gaussian prior.
72    model: PyMC Model
73        Optional model for sampling step. Defaults to None (taken from
74        context).
75
76    References
77    ----------
78    .. [1] I. Murray, R. P. Adams, and D. J. C. MacKay. "Elliptical Slice
79       Sampling", The Proceedings of the 13th International Conference on
80       Artificial Intelligence and Statistics (AISTATS), JMLR W&CP
81       9:541-548, 2010.
82    """
83
84    default_blocked = True
85
86    def __init__(self, vars=None, prior_cov=None, prior_chol=None, model=None, **kwargs):
87        self.model = modelcontext(model)
88        chol = get_chol(prior_cov, prior_chol)
89        self.prior_chol = tt.as_tensor_variable(chol)
90
91        if vars is None:
92            vars = self.model.cont_vars
93        vars = inputvars(vars)
94
95        super().__init__(vars, [self.model.fastlogp], **kwargs)
96
97    def astep(self, q0, logp):
98        """q0: current state
99        logp: log probability function
100        """
101
102        # Draw from the normal prior by multiplying the Cholesky decomposition
103        # of the covariance with draws from a standard normal
104        chol = draw_values([self.prior_chol])[0]
105        nu = np.dot(chol, nr.randn(chol.shape[0]))
106        y = logp(q0) - nr.standard_exponential()
107
108        # Draw initial proposal and propose a candidate point
109        theta = nr.uniform(0, 2 * np.pi)
110        theta_max = theta
111        theta_min = theta - 2 * np.pi
112        q_new = q0 * np.cos(theta) + nu * np.sin(theta)
113
114        while logp(q_new) <= y:
115            # Shrink the bracket and propose a new point
116            if theta < 0:
117                theta_min = theta
118            else:
119                theta_max = theta
120            theta = nr.uniform(theta_min, theta_max)
121            q_new = q0 * np.cos(theta) + nu * np.sin(theta)
122
123        return q_new
124
125    @staticmethod
126    def competence(var, has_grad):
127        # Because it requires a specific type of prior, this step method
128        # should only be assigned explicitly.
129        return Competence.INCOMPATIBLE
130