def atleast_list(a):
    """
    Promote an object to a list if not a list or generator.

    Parameters
    ----------
    a: object
        any object you want to at least be a list with one element

    Returns
    -------
    list or generator:
        untounched if :code:`a` was a generator or list, otherwise :code:`[a]`.

    Examples
    --------
    >>> a = 1.
    >>> atleast_list(a)
    [1.0]
    >>> a = [1.]
    >>> atleast_list(a)
    [1.0]
    """
    return a if isinstance(a, list) or isgenerator(a) else [a] def issequence(obj):
    """
    Test if an object is an iterable generator, list or tuple.

    Parameters
    ----------
    obj : object
        object to test

    Returns
    -------
    bool :
        True if :code:`obj` is a tuple, list or generator only.

    Examples
    --------
    >>> issequence([1, 2])
    True
    >>> issequence((1,))
    True
    >>> issequence((i for i in range(8)))
    True
    >>> issequence(np.array([1, 2, 3]))
    False
    """
    return inspect.isgenerator(obj) or isinstance(obj, (list, tuple)) def logsumexp(X, axis=0):
    """
    Log-sum-exp trick for matrix X for summation along a specified axis.

    This performs the following operation in a stable fashion,

    .. math::

        \log \sum^K_{k=1} \exp\{x_k\}

    Parameters
    ----------
        X: ndarray
            2D array of shape (N, D) to apply the log-sum-exp trick.
        axis: int, optional
            Axis to apply the summation along (works the same as axis in
            numpy.sum).

    Returns
    -------
        lseX: ndarray
            results of applying the log-sum-exp trick, this will be shape (D,)
            if :code:`axis=0` or shape (N,) if :code:`axis=1`.
    """

    mx = X.max(axis=axis)
    if (X.ndim > 1):
        mx = np.atleast_2d(mx).T if axis == 1 else np.atleast_2d(mx)

    return np.log(np.exp(X - mx).sum(axis=axis)) + np.ravel(mx) class LinearBasis(Basis):
    r"""
    Linear basis class, basically this just prepends a columns of ones onto X.

    .. math::

        \phi(\mathbf{X}) = [\mathbf{1}, \mathbf{X}]

    Parameters
    ----------
    onescol: bool, optional
        If true, prepend a column of ones onto X.
    regularizer : None, Parameter, optional
        The (initial) value of the regularizer/prior variance to apply to the
        regression weights of this basis function. The Parameter object must
        have a scalar value. If it is not set, it will take on a default value
        of ``Parameter(gamma(1.), Positive())``.
    """

    @slice_init
    def __init__(self, onescol=True, regularizer=None):
        """See this class's docstring."""
        self.onescol = onescol
        super(LinearBasis, self).__init__(regularizer)

    @slice_transform
    def transform(self, X):
        """
        Return this basis applied to X.

        Parameters
        ----------
        X: ndarray
            of shape (N, d) of observations where N is the number of samples,
            and d is the dimensionality of X.

        Returns
        -------
        ndarray:
            of shape (N, d+1), or (N, d) depending on onescol.
        """
        N, D = X.shape
        return np.hstack((np.ones((N, 1)), X)) if self.onescol else X

    def __repr__(self):
        """Representation."""
        return "{}(onescol={}, regularizer={})".format(
            self.__class__.__name__,
            self.onescol,
            self.regularizer
        ) def apply_grad(fun, grad):
    """
    Apply a function that takes a gradient matrix to a sequence of 2 or 3
    dimensional gradients.

    This is partucularly useful when the gradient of a basis concatenation
    object is quite complex, eg.

    >>> X = np.random.randn(100, 3)
    >>> y = np.random.randn(100)
    >>> N, d = X.shape
    >>> base = RandomRBF(Xdim=d, nbases=5) + RandomRBF(Xdim=d, nbases=5,
    ... lenscale=Parameter(np.ones(d), Positive()))
    >>> Phi = base.transform(X, 1., np.ones(d))
    >>> dffun = lambda dPhi: y.dot(Phi).dot(dPhi.T).dot(y)
    >>> df = apply_grad(dffun, base.grad(X, 1., np.ones(d)))
    >>> np.isscalar(df[0])
    True
    >>> df[1].shape
    (3,)

    Parameters
    ----------
    fun: callable
        the function too apply to the (2d) gradient.
    grad: ndarray or generator
        the gradient of the basis function (output of base.grad).

    Returns
    -------
    scalar, ndarray or sequence:
        the result of applying fun(grad) for a structured grad.
    """
    if issequence(grad):
        fgrad = [apply_grad(fun, g) for g in grad]
        return fgrad if len(fgrad) != 1 else fgrad[0]
    elif len(grad) == 0:
        return []
    elif (grad.ndim == 1) or (grad.ndim == 2):
        return fun(grad)
    elif grad.ndim == 3:
        return np.array([fun(grad[:, :, i]) for i in range(grad.shape[2])])
    else:
        raise ValueError("Only up to 3d gradients allowed!") class Gaussian(Bernoulli):
    r"""
    A univariate Gaussian likelihood for general regression tasks.

    No transformation function is needed since this is (conditionally)
    conjugate to the GLM prior.

    .. math::

        p(y_i | f_i) = \frac{1}{\sqrt{2 \pi \sigma^2}}
            \exp\left(- \frac{(y_i - f_i)^2}{2 \sigma^2} \right)

    where :math:`y_i` is a target, :math:`f_i` the value of the latent function
    corresponding to the target and :math:`\sigma` is the observation noise
    (standard deviation).

    Parameters
    ----------
    var: Parameter, optional
        A scalar Parameter describing the initial point and bounds for
        an optimiser to learn the variance parameter of this object.
    """

    def __init__(self, var=Parameter(gamma(1., scale=1), Positive())):
        """See class docstring."""
        self.params = var

    def _check_param(self, param):

        if param is None:
            return self.params.value

        if not self.params.bounds.check(param):
            raise ValueError("Input parameter is out of bounds!")

        return param

    def loglike(self, y, f, var=None):
        r"""
        Gaussian log likelihood.

        Parameters
        ----------
        y: ndarray
            array of 0, 1 valued integers of targets
        f: ndarray
            latent function from the GLM prior (:math:`\mathbf{f} =
            \boldsymbol\Phi \mathbf{w}`)
        var: float, ndarray, optional
            The variance of the distribution, if not input, the initial value
            of variance is used.

        Returns
        -------
        logp: ndarray
            the log likelihood of each y given each f under this
            likelihood.
        """
        # way faster than calling norm.logpdf
        var = self._check_param(var)
        y, f = np.broadcast_arrays(y, f)
        ll = - 0.5 * (np.log(2 * np.pi * var) + (y - f)**2 / var)
        return ll

    def Ey(self, f, var):
        r"""
        Expected value of the Gaussian likelihood.

        Parameters
        ----------
        f: ndarray
            latent function from the GLM prior (:math:`\mathbf{f} =
            \boldsymbol\Phi \mathbf{w}`)
        var: float, ndarray, optional
            The variance of the distribution, if not input, the initial value
            of variance is used.

        Returns
        -------
        Ey: ndarray
            expected value of y, :math:`\mathbb{E}[\mathbf{y}|\mathbf{f}]`.
        """
        var = self._check_param(var)
        return f

    def df(self, y, f, var):
        r"""
        Derivative of Gaussian log likelihood w.r.t.\  f.

        Parameters
        ----------
        y: ndarray
            array of 0, 1 valued integers of targets
        f: ndarray
            latent function from the GLM prior (:math:`\mathbf{f} =
            \boldsymbol\Phi \mathbf{w}`)
        var: float, ndarray, optional
            The variance of the distribution, if not input, the initial value
            of variance is used.

        Returns
        -------
        df: ndarray
            the derivative :math:`\partial \log p(y|f) / \partial f`
        """
        var = self._check_param(var)
        y, f = np.broadcast_arrays(y, f)
        return (y - f) / var

    def dp(self, y, f, var):
        r"""
        Derivative of Gaussian log likelihood w.r.t.\ the variance
        :math:`\sigma^2`.

        Parameters
        ----------
        y: ndarray
            array of 0, 1 valued integers of targets
        f: ndarray
            latent function from the GLM prior (:math:`\mathbf{f} =
            \boldsymbol\Phi \mathbf{w}`)
        var: float, ndarray, optional
            The variance of the distribution, if not input, the initial value
            of variance is used.

        Returns
        -------
        dp: float
            the derivative
            :math:`\partial \log p(y|f, \sigma^2)/ \partial \sigma^2`
            where :math:`sigma^2` is the variance.
        """
        var = self._check_param(var)
        y, f = np.broadcast_arrays(y, f)
        ivar = 1. / var
        return 0.5 * (((y - f) * ivar)**2 - ivar)

    def cdf(self, y, f, var):
        r"""
        Cumulative density function of the likelihood.

        Parameters
        ----------
        y: ndarray
            query quantiles, i.e.\  :math:`P(Y \leq y)`.
        f: ndarray
            latent function from the GLM prior (:math:`\mathbf{f} =
            \boldsymbol\Phi \mathbf{w}`)
        var: float, ndarray, optional
            The variance of the distribution, if not input, the initial value
            of variance is used.

        Returns
        -------
        cdf: ndarray
            Cumulative density function evaluated at y.
        """
        var = self._check_param(var)
        return norm.cdf(y, loc=f, scale=np.sqrt(var))

    def __repr__(self):
        """Representation."""
        return "{}(var={})".format(self.__class__.__name__, self.params) def sgd(fun, x0, data, args=(), bounds=None, batch_size=10, maxiter=5000,
        updater=None, eval_obj=False, random_state=None):
    """
    Stochastic Gradient Descent.

    Parameters
    ----------
    fun : callable
        the function to *minimize*, this must have the signature ``[obj,]``
        grad = fun(x, data, ...)`, where the ``eval_obj`` argument tells
        ``sgd`` if an objective function value is going to be returned by
        ``fun``.
    x0 : ndarray
        a sequence/1D array of initial values for the parameters to learn.
    data : ndarray
        a numpy array or sequence of data to input into ``fun``. This will
        be split along the first axis (axis=0), and then input into
        ``fun``.
    args : sequence, optional
        an optional sequence of arguments to give to fun.
    bounds : sequence, optional
        Bounds for variables, (min, max) pairs for each element in x, defining
        the bounds on that parameter.  Use None for one of min or max when
        there is no bound in that direction.
    batch_size : int, optional
        The number of observations in an SGD batch.
    maxiter : int, optional
        Number of mini-batch iterations before optimization terminates.
    updater : SGDUpdater, optional
        The type of gradient update to use, by default this is Adam.
    eval_obj : bool, optional
        This indicates whether or not ``fun`` also evaluates and returns
        the objective function value. If this is true, ``fun`` must return
        ``(obj, grad)`` and then a list of objective function values is
        also returned.
    random_state : int or RandomState, optional
        random seed

    Returns
    -------
    res : OptimizeResult
        x : narray
            the final result
        norms : list
            the list of gradient norms
        message : str
            the convergence condition ('maxiter reached' or error)
        objs : list
            the list of objective function evaluations if ``eval_obj``
            is True.
        fun : float
            the final objective function evaluation if ``eval_obj`` is
            True.
    """
    if updater is None:
        updater = Adam()

    # Make sure we aren't using a recycled updater
    updater.reset()

    N = _len_data(data)
    x = np.array(x0, copy=True, dtype=float)
    D = x.shape[0]

    # Make sure we have a valid batch size
    batch_size = min(batch_size, N)

    # Process bounds
    if bounds is not None:
        if len(bounds) != D:
            raise ValueError("The dimension of the bounds does not match x0!")

        lower, upper = zip(*map(normalize_bound, bounds))
        lower = np.array(lower)
        upper = np.array(upper)

    # Learning Records
    obj = None
    objs = []
    norms = []

    for batch in gen_batch(data, batch_size, maxiter, random_state):

        if not eval_obj:
            grad = fun(x, *chain(batch, args))
        else:
            obj, grad = fun(x, *chain(batch, args))
            objs.append(obj)

        norms.append(np.linalg.norm(grad))

        # Truncate gradients if bounded
        if bounds is not None:
            xlower = x <= lower
            grad[xlower] = np.minimum(grad[xlower], 0)
            xupper = x >= upper
            grad[xupper] = np.maximum(grad[xupper], 0)

        # perform update
        x = updater(x, grad)

        # Trucate steps if bounded
        if bounds is not None:
            x = np.clip(x, lower, upper)

    # Format results
    res = OptimizeResult(
        x=x,
        norms=norms,
        message='maxiter reached',
        fun=obj,
        objs=objs
    )

    return res class Adam(SGDUpdater):
    """
    Adam Algorithm.

    Parameters
    ----------
    alpha : float, optional
        stepsize to give the update.
    beta1 : float, optional
        smoothing/decay rate parameter for the gradient, must be [0, 1].
    beta2 : float, optional
        smoothing/decay rate parameter for the squared gradient, must be
        [0, 1].
    epsilon : float, optional
        "jitter" term to ensure continued learning (should be small).
    """

    def __init__(self, alpha=0.01, beta1=0.9, beta2=0.99, epsilon=1e-8):
        """See class docstring."""
        self.alpha = alpha
        self.beta1 = beta1
        self.beta2 = beta2
        self.epsilon = epsilon
        self.t = 0
        self.m = None
        self.v = None

    def __call__(self, x, grad):
        """
        Get a new parameter value from the Adam algorithm.

        Parameters
        ----------
        x : ndarray
            input parameters to optimise
        grad : ndarray
            gradient of x

        Returns
        -------
        x_new : ndarray
            the new value for x
        """
        self.t += 1

        if self.m is None:
            self.m = np.zeros_like(x)
            self.v = np.zeros_like(x)

        self.m = self.beta1 * self.m + (1 - self.beta1) * grad
        self.v = self.beta2 * self.v + (1 - self.beta2) * grad**2
        mbar = self.m / (1 - self.beta1**self.t)
        vbar = self.v / (1 - self.beta2**self.t)

        x_new = x - self.alpha * mbar / (np.sqrt(vbar) + self.epsilon)
        return x_new

    def reset(self):
        """Reset the state of this updater for a new optimisation problem."""
        self.__init__(self.alpha, self.beta1, self.beta2, self.epsilon)

    def __repr__(self):
        """Representation."""
        return "{}(alpha={}, beta1={}, beta2={}, epsilon={})".format(
            self.__class__.__name__,
            self.alpha,
            self.beta1,
            self.beta2,
            self.epsilon
        ) def structured_sgd(sgd):
    r"""
    Allow an SGD to accept nested sequences of Parameters to optimize.

    This decorator can intepret the :code:`Parameter` objects in `btypes.py`,
    and can accept nested sequences of *any* structure of these objects to
    optimise!

    It can also optionally evaluate *random starts* (i.e. random starting
    candidates) if the parameter objects have been initialised with
    distributions. For this, an additional parameter is exposed in the
    :code:`minimizer` interface.

    Parameters
    ----------
    fun : callable
        objective function that takes in arbitrary ndarrays, floats or nested
        sequences of these.
    parameters : (nested) sequences of Parameter objects
        Initial guess of the parameters of the objective function
    nstarts : int, optional
        The number random starting candidates for optimisation to evaluate.
        This will only happen for :code:`nstarts > 0` and if at least one
        :code:`Parameter` object is random.

    Examples
    --------
    >>> from ..optimize import sgd
    >>> from ..btypes import Parameter, Bound

    Define a cost function that returns a pair. The first element is the cost
    value and the second element is the gradient represented by a sequence.
    Even if the cost is a function of a single variable, the gradient must be a
    sequence containing one element.

    >>> def cost(w, lambda_, data):
    ...     N = len(data)
    ...     y, X = data[:, 0], data[:, 1:]
    ...     y_est = X.dot(w)
    ...     ww = w.T.dot(w)
    ...     obj = (y - y_est).sum() / N + lambda_ * ww
    ...     gradw = - 2 * X.T.dot(y - y_est) / N + 2 * lambda_ * w
    ...     gradl = ww
    ...     return obj, [gradw, gradl]

    Augment the SGD optimizer to take structured inputs

    >>> new_sgd = structured_sgd(sgd)

    Data

    >>> y = np.linspace(1, 10, 100) + np.random.randn(100) + 1
    >>> X = np.array([np.ones(100), np.linspace(1, 100, 100)]).T
    >>> data = np.hstack((y[:, np.newaxis], X))

    Constant Initial values

    >>> w_0 = Parameter(np.array([1., 1.]), Bound())
    >>> lambda_0 = Parameter(.25, Bound())

    >>> res = new_sgd(cost, [w_0, lambda_0], data, batch_size=10,
    ...               eval_obj=True)
    >>> res_w, res_lambda = res.x

    Random Initial values

    >>> from scipy.stats import norm, gamma
    >>> w_0 = Parameter(norm(), Bound(), shape=(2,))
    >>> lambda_0 = Parameter(gamma(1.), Bound())

    >>> res = new_sgd(cost, [w_0, lambda_0], data, batch_size=10,
    ...               eval_obj=True, nstarts=100)
    >>> res_w, res_lambda = res.x
    """
    @wraps(sgd)
    def new_sgd(fun,
                parameters,
                data,
                eval_obj=False,
                batch_size=10,
                args=(),
                random_state=None,
                nstarts=100,
                **sgd_kwargs):

        (array1d, fbounds), shapes = flatten(parameters,
                                             hstack=bt.hstack,
                                             shape=bt.shape,
                                             ravel=bt.ravel
                                             )

        flatten_args_dec = flatten_args(shapes)
        new_fun = flatten_args_dec(fun)

        # Find best random starting candidate if we are doing random starts
        if eval_obj and nstarts > 0:
            data_gen = gen_batch(data, batch_size, random_state=random_state)
            array1d = _random_starts(
                fun=fun,
                parameters=parameters,
                jac=True,
                args=args,
                data_gen=data_gen,
                nstarts=nstarts,
                random_state=random_state
            )

        if bool(eval_obj):
            new_fun = flatten_func_grad(new_fun)
        else:
            new_fun = flatten_grad(new_fun)

        result = sgd(new_fun, array1d, data=data, bounds=fbounds, args=args,
                     eval_obj=eval_obj, random_state=random_state,
                     **sgd_kwargs)

        result['x'] = tuple(unflatten(result['x'], shapes))
        return result

    return new_sgd def logtrick_sgd(sgd):
    r"""
    Log-Trick decorator for stochastic gradients.

    This decorator implements the "log trick" for optimizing positive bounded
    variables using SGD. It will apply this trick for any variables that
    correspond to a Positive() bound.

    Examples
    --------
    >>> from ..optimize import sgd
    >>> from ..btypes import Bound, Positive

    Here is an example where we may want to enforce a particular parameter or
    parameters to be strictly greater than zero,

    >>> def cost(w, data, lambda_):
    ...     N = len(data)
    ...     y, X = data[:, 0], data[:, 1:]
    ...     y_est = X.dot(w)
    ...     ww = w.T.dot(w)
    ...     obj = (y - y_est).sum() / N + lambda_ * ww
    ...     gradw = - 2 * X.T.dot(y - y_est) / N + 2 * lambda_ * w
    ...     return obj, gradw

    Now let's enforce that the `w` are positive,

    >>> bounds = [Positive(), Positive()]
    >>> new_sgd = logtrick_sgd(sgd)

    Data

    >>> y = np.linspace(1, 10, 100) + np.random.randn(100) + 1
    >>> X = np.array([np.ones(100), np.linspace(1, 100, 100)]).T
    >>> data = np.hstack((y[:, np.newaxis], X))

    Initial values

    >>> w_0 = np.array([1., 1.])
    >>> lambda_0 = .25

    >>> res = new_sgd(cost, w_0, data, args=(lambda_0,), bounds=bounds,
    ...               batch_size=10, eval_obj=True)
    >>> res.x >= 0
    array([ True,  True], dtype=bool)

    Note
    ----
    This decorator only works on unstructured optimizers. However, it can be
    use with structured_minimizer, so long as it is the inner wrapper.
    """
    @wraps(sgd)
    def new_sgd(fun, x0, data, bounds=None, eval_obj=False, **sgd_kwargs):

        if bounds is None:
            return sgd(fun, x0, data, bounds=bounds, eval_obj=eval_obj,
                       **sgd_kwargs)

        logx, expx, gradx, bounds = _logtrick_gen(bounds)

        if bool(eval_obj):
            def new_fun(x, *fargs, **fkwargs):
                o, g = fun(expx(x), *fargs, **fkwargs)
                return o, gradx(g, x)
        else:
            def new_fun(x, *fargs, **fkwargs):
                return gradx(fun(expx(x), *fargs, **fkwargs), x)

        # Transform the final result
        result = sgd(new_fun, logx(x0), data, bounds=bounds, eval_obj=eval_obj,
                     **sgd_kwargs)
        result['x'] = expx(result['x'])
        return result

    return new_sgd class Bound(namedtuple('Bound', ['lower', 'upper']), _BoundMixin):
    """
    Define bounds on a variable for the optimiser.

    This defaults to all real values allowed (i.e. no bounds).

    Parameters
    ----------
    lower : float
        The lower bound.
    upper : float
        The upper bound.

    Attributes
    ----------
    lower : float
        The lower bound.
    upper : float
        The upper bound.

    Examples
    --------
    >>> b = Bound(1e-10, upper=1e-5)
    >>> b
    Bound(lower=1e-10, upper=1e-05)
    >>> b.lower
    1e-10
    >>> b.upper
    1e-05
    >>> isinstance(b, tuple)
    True
    >>> tuple(b)
    (1e-10, 1e-05)
    >>> lower, upper = b
    >>> lower
    1e-10
    >>> upper
    1e-05
    >>> Bound(42, 10)
    Traceback (most recent call last):
        ...
    ValueError: lower bound cannot be greater than upper bound!
    """

    def __new__(cls, lower=None, upper=None):
        """See class docstring."""
        if lower is not None and upper is not None:
            if lower > upper:
                raise ValueError('lower bound cannot be greater than upper '
                                 'bound!')
        return super(Bound, cls).__new__(cls, lower, upper)

    def __getnewargs__(self):
        """Required for pickling."""
        return (self.lower, self.upper)

    def __repr__(self):
        """Representation."""
        return "{}(lower={}, upper={})".format(
            self.__class__.__name__,
            self.lower,
            self.upper
        ) class Positive(namedtuple('Positive', ['lower', 'upper']), _BoundMixin):
    """
    Define a positive only bound for the optimiser.

    This may induce the 'log trick' in the optimiser (when using an appropriate
    decorator), which will ignore the 'smallest' value (but will stay above 0).

    Parameters
    ---------
    upper : float
        The largest value allowed for the optimiser to evaluate (if not using
        the log trick).

    Examples
    --------
    >>> b = Positive()
    >>> b
    Positive(upper=None)

    Since ``tuple`` (and by extension its descendents) are immutable,
    the lower bound for all instances of ``Positive`` are guaranteed to
    be positive.
    """

    def __new__(cls, upper=None):
        """See class docstring."""
        lower = 1e-14
        if upper is not None:
            if lower > upper:
                raise ValueError('Upper bound must be greater than {}'
                                 .format(lower))

        return super(Positive, cls).__new__(cls, lower, upper)

    def __getnewargs__(self):
        """Required for pickling."""
        return (self.upper,)

    def __repr__(self):
        """Representation."""
        return "{}(upper={})".format(self.__class__.__name__, self.upper) class Parameter(object):
    """
    A Parameter class that associates a value with a bound.

    Attributes
    ----------
    value : None, scalar, ndarray, scipy.stats, optional
        a value or distribution to associate with this parameter. This is
        typically used as an initial value for an optimizer, and if a random
        starts optimiser is used (eg. revrand.optimize.structured_minimizer) it
        can draw randomly from the distribution.
    bound : Bound, optional
        a Bound tuple that describes the valid range for all of the elements in
        value
    shape : tuple, optional
        the shape of value, this is ignored if value is not a scipy.stats
        distribution (i.e. it is automatically determined).

    Note
    ----
    - If ``value`` is initialised as a distribution from ``scipy.stats``, then
      ``self.value`` is actually the mean of the distribution.
    - If you want to set value to a ``scipy.stats`` distribution, and also
      associate it with a ``shape`` (i.e. you want an ``ndarray`` random
      variable), then also set the ``shape`` parameters to the desired
      dimensions.

    Examples
    --------
    Null

    >>> p = Parameter()
    >>> p.value
    []
    >>> p.has_value
    False

    Scalar

    >>> p = Parameter(1.2, Bound(1, 2))
    >>> p.shape
    ()
    >>> p.value
    1.2
    >>> p.rvs()
    1.2
    >>> p.is_random
    False

    ndarray

    >>> p = Parameter(np.ones(3), Positive())
    >>> p.shape
    (3,)
    >>> p.value
    array([ 1.,  1.,  1.])
    >>> p.rvs()
    array([ 1.,  1.,  1.])

    ``scipy.stats`` scalar

    >>> from scipy.stats import gamma
    >>> p = Parameter(gamma(a=1, scale=1), Positive())
    >>> p.value == gamma(a=1, scale=1).mean()
    True
    >>> np.isscalar(p.rvs())
    True
    >>> p.is_random
    True

    ``scipy.stats`` ndarray

    >>> from scipy.stats import gamma
    >>> p = Parameter(gamma(a=1, scale=1), Positive(), shape=(3,))
    >>> all(p.value == np.ones(3) * gamma(a=1, scale=1).mean())
    True
    >>> p.rvs().shape == (3,)
    True
    """

    def __init__(self, value=None, bounds=Bound(), shape=()):
        """See class docstring."""
        value = [] if value is None else value

        if not hasattr(value, 'rvs'):
            if np.any(value) and not bounds.check(value):
                raise ValueError("Value not within bounds!")

            self.value = value
            self.dist = None
            self.shape = np.shape(value)
        else:
            self.dist = value
            self.shape = shape
            mean = bounds.clip(self.dist.mean())
            self.value = mean if shape == () else mean * np.ones(shape)

        self.bounds = bounds

    def rvs(self, random_state=None):
        r"""
        Draw a random value from this Parameter's distribution.

        If ``value`` was not initialised with a ``scipy.stats`` object, then
        the scalar/ndarray value is returned.

        Parameters
        ----------
        random_state : None, int or RandomState, optional
            random seed

        Returns
        -------
        ndarray :
            of size ``self.shape``, a random draw from the distribution, or
            ``self.value`` if not initialised with a ``scipy.stats`` object.

        Note
        ----
        Random draws are *clipped* to the bounds, and so it is up to the user
        to input a sensible sampling distribution!
        """
        # No sampling distibution
        if self.dist is None:
            return self.value

        # Unconstrained samples
        rs = check_random_state(random_state)
        samples = self.dist.rvs(size=self.shape, random_state=rs)

        # Bound the samples
        samples = self.bounds.clip(samples)

        return samples

    @property
    def has_value(self):
        """Test if this Parameter has a value, or is "null"."""
        return self.shape != (0,)

    @property
    def is_random(self):
        """Test if this Parameter was initialised with a distribution."""
        return self.dist is not None

    @property
    def is_scalar(self):
        """Test if the Parameter is for a scalar value."""
        return self.has_value and self.shape == ()

    def __repr__(self):
        """Representation."""
        return "{}(value={}, bounds={}, shape={})".format(
            self.__class__.__name__,
            self.value,
            self.bounds,
            self.shape
        )
import numpy as np
import logging
from itertools import chain
from multiprocessing import Pool
from scipy.stats.distributions import gamma, norm
from scipy.optimize import brentq
from sklearn.base import BaseEstimator, RegressorMixin
from sklearn.utils.validation import check_is_fitted, check_X_y, check_array
from sklearn.utils import check_random_state
from .utils import atleast_list, issequence
from .mathfun.special import logsumexp
from .basis_functions import LinearBasis, apply_grad
from .likelihoods import Gaussian
from .optimize import sgd, structured_sgd, logtrick_sgd, Adam
from .btypes import Bound, Positive, Parameter

    This provides a scikit learn compatible interface for the glm module.

    Parameters
    ----------
    likelihood : Object
        A likelihood object, see the likelihoods module.
    basis : Basis
        A basis object, see the basis_functions module.
    K : int, optional
        Number of diagonal Gaussian components to use to approximate the
        posterior distribution.
    maxiter : int, optional
        Maximum number of iterations of stochastic gradients to run.
    batch_size : int, optional
        number of observations to use per SGD batch.
    updater : SGDUpdater, optional
        The SGD learning rate updating algorithm to use, by default this is
        Adam. See revrand.optimize.sgd for different options.
    nsamples : int, optional
        Number of samples for sampling the expected likelihood and expected
        likelihood gradients
    nstarts : int, optional
        if there are any parameters with distributions as initial values, this
        determines how many random candidate starts shoulds be evaluated before
        commencing optimisation at the best candidate.
    random_state : None, int or RandomState, optional
        random seed

    Notes
    -----
    This approximates the posterior distribution over the weights with a
    mixture of Gaussians:

    .. math ::

        \mathbf{w} \sim \frac{1}{K} \sum^K_{k=1}
            \mathcal{N}(\mathbf{m_k}, \boldsymbol{\Psi}_k)

    where,

    .. math ::

        \boldsymbol{\Psi}_k = \text{diag}([\Psi_{k,1}, \ldots,
            \Psi_{k,D}]).

    This is so arbitrary likelihoods can be used with this algorithm, while
    still mainting flexible and tractable non-Gaussian posteriors. Additionaly
    this has the benefit that we have a reduced number of parameters to
    optimise (compared with full covariance Gaussians).

    The main differences between this implementation and the GLM in [1]_ are:
        - We use diagonal mixtures, as opposed to isotropic.
        - We use auto encoding variational Bayes (AEVB) inference [2]_ with
          stochastic gradients.

    This uses the python logging module for displaying learning status. To view
    these messages have something like,

    .. code ::

        import logging
        logging.basicConfig(level=logging.INFO)
        log = logging.getLogger(__name__)

    in your calling code.
    """

    def __init__(self,
                 likelihood=Gaussian(),
                 basis=LinearBasis(),
                 K=10,
                 maxiter=3000,
                 batch_size=10,
                 updater=None,
                 nsamples=50,
                 nstarts=500,
                 random_state=None
                 ):
        """See class docstring."""
        self.likelihood = likelihood
        self.basis = basis
        self.K = K
        self.maxiter = maxiter
        self.batch_size = batch_size
        self.updater = updater
        self.nsamples = nsamples
        self.nstarts = nstarts
        self.random_state = random_state  # For clone compatibility
        self.random_ = check_random_state(self.random_state)

    def fit(self, X, y, likelihood_args=()):
        r"""
        Learn the parameters of a Bayesian generalized linear model (GLM).

        Parameters
        ----------
        X : ndarray
            (N, d) array input dataset (N samples, d dimensions).
        y : ndarray
            (N,) array targets (N samples)
        likelihood : Object
            A likelihood object, see the likelihoods module.
        likelihood_args : sequence, optional
            sequence of arguments to pass to the likelihood function. These are
            non-learnable parameters. They can be scalars or arrays of length
            N.
        """
        X, y = check_X_y(X, y)

        # Batch magnification factor
        N, _ = X.shape
        self.B_ = X.shape[0] / self.batch_size
        self.D_ = self.basis.get_dim(X)

        # Pack data
        likelihood_args = _reshape_likelihood_args(likelihood_args, N)
        data = (X, y) + likelihood_args

        # Pack params

        params = [Parameter(WGTRND, Bound(), shape=(self.D_, self.K)),