################################################################################ # Copyright (C) 2011-2014 Jaakko Luttinen # # This file is licensed under the MIT License. ################################################################################ """ Module for the Gaussian distribution and similar distributions. """ import numpy as np from scipy import special import truncnorm from bayespy.utils import (random, misc, linalg) from bayespy.utils.linalg import dot, mvdot from .expfamily import (ExponentialFamily, ExponentialFamilyDistribution, useconstructor) from .wishart import (WishartMoments, WishartPriorMoments) from .gamma import (GammaMoments, GammaDistribution, GammaPriorMoments) from .deterministic import Deterministic from .node import (Moments, ensureparents) # # MOMENTS # class GaussianMoments(Moments): r""" Class for the moments of Gaussian variables. """ def __init__(self, shape): self.shape = shape self.ndim = len(shape) self.dims = (shape, 2*shape) super().__init__() def compute_fixed_moments(self, x): r""" Compute the moments for a fixed value """ x = np.asanyarray(x) x = misc.atleast_nd(x, self.ndim) return [x, linalg.outer(x, x, ndim=self.ndim)] @classmethod def from_values(cls, x, ndim): r""" Return the shape of the moments for a fixed value. """ if ndim == 0: return cls(()) else: return cls(np.shape(x)[-ndim:]) def get_instance_conversion_kwargs(self): return dict(ndim=self.ndim) def get_instance_converter(self, ndim): if ndim == self.ndim or ndim is None: return None return GaussianToGaussian(self, ndim) class GaussianToGaussian(): def __init__(self, moments_from, ndim_to): if not isinstance(moments_from, GaussianMoments): raise ValueError() if ndim_to < 0: return ValueError("ndim_to must be non-negative") self.shape_from = moments_from.shape self.ndim_from = moments_from.ndim self.ndim_to = ndim_to if self.ndim_to > self.ndim_from: raise ValueError() if self.ndim_to == 0: self.moments = GaussianMoments(()) else: self.moments = GaussianMoments(self.shape_from[-self.ndim_to:]) return def compute_moments(self, u): if self.ndim_to == self.ndim_from: return u u0 = u[0] u1 = misc.get_diag(u[1], ndim=self.ndim_from, ndim_to=self.ndim_to) return [u0, u1] def compute_message_to_parent(self, m, u_parent): # Handle broadcasting in m_child m0 = m[0] * np.ones(self.shape_from) m1 = ( misc.make_diag(m[1], ndim=self.ndim_from, ndim_from=self.ndim_to) * misc.identity(*self.shape_from) ) return [m0, m1] def compute_weights_to_parent(self, weights): diff = self.ndim_from - self.ndim_to if diff == 0: return weights return np.sum( weights * np.ones(self.shape_from[:diff]), #misc.atleast_nd(weights, diff), axis=tuple(range(-diff, 0)) ) def plates_multiplier_from_parent(self, plates_multiplier): diff = self.ndim_from - self.ndim_to return plates_multiplier + diff * (1,) def plates_from_parent(self, plates): diff = self.ndim_from - self.ndim_to if diff == 0: return plates return plates + self.shape_from[:diff] def plates_to_parent(self, plates): diff = self.ndim_from - self.ndim_to if diff == 0: return plates return plates[:-diff] class GaussianGammaMoments(Moments): r""" Class for the moments of Gaussian-gamma-ISO variables. """ def __init__(self, shape): r""" Create moments object for Gaussian-gamma isotropic variables ndim=0: scalar ndim=1: vector ndim=2: matrix ... """ self.shape = shape self.ndim = len(shape) self.dims = (shape, 2*shape, (), ()) super().__init__() def compute_fixed_moments(self, x_alpha): r""" Compute the moments for a fixed value `x` is a mean vector. `alpha` is a precision scale """ (x, alpha) = x_alpha x = np.asanyarray(x) alpha = np.asanyarray(alpha) u0 = x * misc.add_trailing_axes(alpha, self.ndim) u1 = (linalg.outer(x, x, ndim=self.ndim) * misc.add_trailing_axes(alpha, 2*self.ndim)) u2 = np.copy(alpha) u3 = np.log(alpha) u = [u0, u1, u2, u3] return u @classmethod def from_values(cls, x_alpha, ndim): r""" Return the shape of the moments for a fixed value. """ (x, alpha) = x_alpha if ndim == 0: shape = ( (), (), (), () ) else: shape = np.shape(x)[-ndim:] return cls(shape) def get_instance_conversion_kwargs(self): return dict(ndim=self.ndim) def get_instance_converter(self, ndim): # FIXME/TODO: IMPLEMENT THIS CORRECTLY! if ndim != self.ndim: raise NotImplementedError( "Conversion to different ndim in GaussianMoments not yet " "implemented." ) return None class GaussianWishartMoments(Moments): r""" Class for the moments of Gaussian-Wishart variables. """ def __init__(self, shape): self.shape = shape self.ndim = len(shape) self.dims = ( shape, (), 2*shape, () ) super().__init__() def compute_fixed_moments(self, x, Lambda): r""" Compute the moments for a fixed value `x` is a vector. `Lambda` is a precision matrix """ x = np.asanyarray(x) Lambda = np.asanyarray(Lambda) u0 = linalg.mvdot(Lambda, x, ndim=self.ndim) u1 = np.einsum( '...i,...ij,...j->...', misc.flatten_axes(x, self.ndim), misc.flatten_axes(Lambda, self.ndim, self.ndim), misc.flatten_axes(x, self.ndim) ) u2 = np.copy(Lambda) u3 = linalg.logdet_cov(Lambda, ndim=self.ndim) return [u0, u1, u2, u3] @classmethod def from_values(self, x, Lambda, ndim): r""" Return the shape of the moments for a fixed value. """ if ndim == 0: return cls(()) else: if np.ndim(x) < ndim: raise ValueError("Mean must be a vector") shape = np.shape(x)[-ndim:] if np.shape(Lambda)[-2*ndim:] != shape + shape: raise ValueError("Shapes inconsistent") return cls(shape) # # DISTRIBUTIONS # class GaussianDistribution(ExponentialFamilyDistribution): r""" Class for the VMP formulas of Gaussian variables. Currently, supports only vector variables. Notes ----- Message passing equations: .. math:: \mathbf{x} &\sim \mathcal{N}(\boldsymbol{\mu}, \mathbf{\Lambda}), .. math:: \mathbf{x},\boldsymbol{\mu} \in \mathbb{R}^{D}, \quad \mathbf{\Lambda} \in \mathbb{R}^{D \times D}, \quad \mathbf{\Lambda} \text{ symmetric positive definite} .. math:: \log\mathcal{N}( \mathbf{x} | \boldsymbol{\mu}, \mathbf{\Lambda} ) &= - \frac{1}{2} \mathbf{x}^{\mathrm{T}} \mathbf{\Lambda} \mathbf{x} + \mathbf{x}^{\mathrm{T}} \mathbf{\Lambda} \boldsymbol{\mu} - \frac{1}{2} \boldsymbol{\mu}^{\mathrm{T}} \mathbf{\Lambda} \boldsymbol{\mu} + \frac{1}{2} \log |\mathbf{\Lambda}| - \frac{D}{2} \log (2\pi) """ def __init__(self, shape): self.shape = shape self.ndim = len(shape) self.set_limits(None, None) super().__init__() def set_limits(self, minimum=None, maximum=None): self.minimum = minimum self.maximum = maximum self.has_limits = minimum is not None or maximum is not None return def compute_message_to_parent(self, parent, index, u, u_mu_Lambda): r""" Compute the message to a parent node. .. math:: \boldsymbol{\phi}_{\boldsymbol{\mu}} (\mathbf{x}, \mathbf{\Lambda}) &= \left[ \begin{matrix} \mathbf{\Lambda} \mathbf{x} \\ - \frac{1}{2} \mathbf{\Lambda} \end{matrix} \right] \\ \boldsymbol{\phi}_{\mathbf{\Lambda}} (\mathbf{x}, \boldsymbol{\mu}) &= \left[ \begin{matrix} - \frac{1}{2} \mathbf{xx}^{\mathrm{T}} + \frac{1}{2} \mathbf{x}\boldsymbol{\mu}^{\mathrm{T}} + \frac{1}{2} \boldsymbol{\mu}\mathbf{x}^{\mathrm{T}} - \frac{1}{2} \boldsymbol{\mu\mu}^{\mathrm{T}} \\ \frac{1}{2} \end{matrix} \right] """ if index == 0: x = u[0] xx = u[1] m0 = x m1 = -0.5 m2 = -0.5*xx m3 = 0.5 return [m0, m1, m2, m3] else: raise ValueError("Index out of bounds") def compute_phi_from_parents(self, u_mu_Lambda, mask=True): r""" Compute the natural parameter vector given parent moments. .. math:: \boldsymbol{\phi} (\boldsymbol{\mu}, \mathbf{\Lambda}) &= \left[ \begin{matrix} \mathbf{\Lambda} \boldsymbol{\mu} \\ - \frac{1}{2} \mathbf{\Lambda} \end{matrix} \right] """ Lambda_mu = u_mu_Lambda[0] Lambda = u_mu_Lambda[2] return [Lambda_mu, -0.5 * Lambda] def compute_moments_and_cgf(self, phi, mask=True): r""" Compute the moments and :math:`g(\phi)`. .. math:: \overline{\mathbf{u}} (\boldsymbol{\phi}) &= \left[ \begin{matrix} - \frac{1}{2} \boldsymbol{\phi}^{-1}_2 \boldsymbol{\phi}_1 \\ \frac{1}{4} \boldsymbol{\phi}^{-1}_2 \boldsymbol{\phi}_1 \boldsymbol{\phi}^{\mathrm{T}}_1 \boldsymbol{\phi}^{-1}_2 - \frac{1}{2} \boldsymbol{\phi}^{-1}_2 \end{matrix} \right] \\ g_{\boldsymbol{\phi}} (\boldsymbol{\phi}) &= \frac{1}{4} \boldsymbol{\phi}^{\mathrm{T}}_1 \boldsymbol{\phi}^{-1}_2 \boldsymbol{\phi}_1 + \frac{1}{2} \log | -2 \boldsymbol{\phi}_2 | """ # TODO: Compute -2*phi[1] and simplify the formulas L = linalg.chol(-2*phi[1], ndim=self.ndim) k = np.shape(phi[0])[-1] Cov = linalg.chol_inv(L, ndim=self.ndim) mu = linalg.chol_solve(L, phi[0], ndim=self.ndim) # G g = (-0.5 * linalg.inner(mu, phi[0], ndim=self.ndim) + 0.5 * linalg.chol_logdet(L, ndim=self.ndim)) if self.has_limits: if self.ndim != 1: raise NotImplementedError("Limits for ndim!=1 not yet supported") (p, u0, u1)= truncnorm.moments( mu, Cov, self.minimum, self.maximum, 2, ) logp = np.log(p) else: u0 = mu u1 = Cov + linalg.outer(u0, u0, ndim=self.ndim) logp = 0 u = [u0, u1] return (u, g - logp) def compute_cgf_from_parents(self, u_mu_Lambda): r""" Compute :math:`\mathrm{E}_{q(p)}[g(p)]` .. math:: g (\boldsymbol{\mu}, \mathbf{\Lambda}) &= - \frac{1}{2} \operatorname{tr}(\boldsymbol{\mu\mu}^{\mathrm{T}} \mathbf{\Lambda} ) + \frac{1}{2} \log |\mathbf{\Lambda}| """ mu_Lambda_mu = u_mu_Lambda[1] logdet_Lambda = u_mu_Lambda[3] g = -0.5*mu_Lambda_mu + 0.5*logdet_Lambda return g def compute_fixed_moments_and_f(self, x, mask=True): r""" Compute the moments and :math:`f(x)` for a fixed value. .. math:: \mathbf{u} (\mathbf{x}) &= \left[ \begin{matrix} \mathbf{x} \\ \mathbf{xx}^{\mathrm{T}} \end{matrix} \right] \\ f(\mathbf{x}) &= - \frac{D}{2} \log(2\pi) """ k = np.shape(x)[-1] u = [x, linalg.outer(x, x, ndim=self.ndim)] f = -k/2*np.log(2*np.pi) return (u, f) def compute_gradient(self, g, u, phi): r""" Compute the standard gradient with respect to the natural parameters. Gradient of the moments: .. math:: \mathrm{d}\overline{\mathbf{u}} &= \begin{bmatrix} \frac{1}{2} \phi_2^{-1} \mathrm{d}\phi_2 \phi_2^{-1} \phi_1 - \frac{1}{2} \phi_2^{-1} \mathrm{d}\phi_1 \\ - \frac{1}{4} \phi_2^{-1} \mathrm{d}\phi_2 \phi_2^{-1} \phi_1 \phi_1^{\mathrm{T}} \phi_2^{-1} - \frac{1}{4} \phi_2^{-1} \phi_1 \phi_1^{\mathrm{T}} \phi_2^{-1} \mathrm{d}\phi_2 \phi_2^{-1} + \frac{1}{2} \phi_2^{-1} \mathrm{d}\phi_2 \phi_2^{-1} + \frac{1}{4} \phi_2^{-1} \mathrm{d}\phi_1 \phi_1^{\mathrm{T}} \phi_2^{-1} + \frac{1}{4} \phi_2^{-1} \phi_1 \mathrm{d}\phi_1^{\mathrm{T}} \phi_2^{-1} \end{bmatrix} \\ &= \begin{bmatrix} 2 (\overline{u}_2 - \overline{u}_1 \overline{u}_1^{\mathrm{T}}) \mathrm{d}\phi_2 \overline{u}_1 + (\overline{u}_2 - \overline{u}_1 \overline{u}_1^{\mathrm{T}}) \mathrm{d}\phi_1 \\ u_2 d\phi_2 u_2 - 2 u_1 u_1^T d\phi_2 u_1 u_1^T + 2 (u_2 - u_1 u_1^T) d\phi_1 u_1^T \end{bmatrix} Standard gradient given the gradient with respect to the moments, that is, given the Riemannian gradient :math:`\tilde{\nabla}`: .. math:: \nabla = \begin{bmatrix} (\overline{u}_2 - \overline{u}_1 \overline{u}_1^{\mathrm{T}}) \tilde{\nabla}_1 + 2 (u_2 - u_1 u_1^T) \tilde{\nabla}_2 u_1 \\ (u_2 - u_1 u_1^T) \tilde{\nabla}_1 u_1^T + u_1 \tilde{\nabla}_1^T (u_2 - u_1 u_1^T) + 2 u_2 \tilde{\nabla}_2 u_2 - 2 u_1 u_1^T \tilde{\nabla}_2 u_1 u_1^T \end{bmatrix} """ ndim = 1 x = u[0] xx = u[1] # Some helpful variables x_x = linalg.outer(x, x, ndim=self.ndim) Cov = xx - x_x cov_g0 = linalg.mvdot(Cov, g[0], ndim=self.ndim) cov_g0_x = linalg.outer(cov_g0, x, ndim=self.ndim) g1_x = linalg.mvdot(g[1], x, ndim=self.ndim) # Compute gradient terms d0 = cov_g0 + 2 * linalg.mvdot(Cov, g1_x, ndim=self.ndim) d1 = (cov_g0_x + linalg.transpose(cov_g0_x, ndim=self.ndim) + 2 * linalg.mmdot(xx, linalg.mmdot(g[1], xx, ndim=self.ndim), ndim=self.ndim) - 2 * x_x * misc.add_trailing_axes(linalg.inner(g1_x, x, ndim=self.ndim), 2*self.ndim)) return [d0, d1] def random(self, *phi, plates=None): r""" Draw a random sample from the distribution. """ # TODO/FIXME: You shouldn't draw random values for # observed/fixed elements! # Note that phi[1] is -0.5*inv(Cov) U = linalg.chol(-2*phi[1], ndim=self.ndim) mu = linalg.chol_solve(U, phi[0], ndim=self.ndim) shape = plates + self.shape z = np.random.randn(*shape) # Denote Lambda = -2*phi[1] # Then, Cov = inv(Lambda) = inv(U'*U) = inv(U) * inv(U') # Thus, compute mu + U\z z = linalg.solve_triangular(U, z, trans='N', lower=False, ndim=self.ndim) return mu + z class GaussianARDDistribution(ExponentialFamilyDistribution): r""" ... Log probability density function: .. math:: \log p(x|\mu, \alpha) = -\frac{1}{2} x^T \mathrm{diag}(\alpha) x + x^T \mathrm{diag}(\alpha) \mu - \frac{1}{2} \mu^T \mathrm{diag}(\alpha) \mu + \frac{1}{2} \sum_i \log \alpha_i - \frac{D}{2} \log(2\pi) Parent has moments: .. math:: \begin{bmatrix} \alpha \circ \mu \\ \alpha \circ \mu \circ \mu \\ \alpha \\ \log(\alpha) \end{bmatrix} """ def __init__(self, shape): self.shape = shape self.ndim = len(shape) super().__init__() def compute_message_to_parent(self, parent, index, u, u_mu_alpha): r""" ... .. math:: m = \begin{bmatrix} x \\ [-\frac{1}{2}, \ldots, -\frac{1}{2}] \\ -\frac{1}{2} \mathrm{diag}(xx^T) \\ [\frac{1}{2}, \ldots, \frac{1}{2}] \end{bmatrix} """ if index == 0: x = u[0] x2 = misc.get_diag(u[1], ndim=self.ndim) m0 = x m1 = -0.5 * np.ones(self.shape) m2 = -0.5 * x2 m3 = 0.5 * np.ones(self.shape) return [m0, m1, m2, m3] else: raise ValueError("Invalid parent index") def compute_weights_to_parent(self, index, weights): r""" Maps the mask to the plates of a parent. """ if index != 0: raise IndexError() return misc.add_trailing_axes(weights, self.ndim) def compute_phi_from_parents(self, u_mu_alpha, mask=True): alpha_mu = u_mu_alpha[0] alpha = u_mu_alpha[2] #mu = u_mu[0] #alpha = u_alpha[0] ## if np.ndim(mu) < self.ndim_mu: ## raise ValueError("Moment of mu does not have enough dimensions") ## mu = misc.add_axes(mu, ## axis=np.ndim(mu)-self.ndim_mu, ## num=self.ndim-self.ndim_mu) phi0 = alpha_mu phi1 = -0.5 * alpha if self.ndim > 0: # Ensure that phi is not using broadcasting for variable # dimension axes ones = np.ones(self.shape) phi0 = ones * phi0 phi1 = ones * phi1 # Make a diagonal matrix phi1 = misc.diag(phi1, ndim=self.ndim) return [phi0, phi1] def compute_moments_and_cgf(self, phi, mask=True): if self.ndim == 0: # Use scalar equations u0 = -phi[0] / (2*phi[1]) u1 = u0**2 - 1 / (2*phi[1]) u = [u0, u1] g = (-0.5 * u[0] * phi[0] + 0.5 * np.log(-2*phi[1])) # TODO/FIXME: You could use these equations if phi is a scalar # in practice although ndim>0 (because the shape can be, e.g., # (1,1,1,1) for ndim=4). else: # Reshape to standard vector and matrix D = np.prod(self.shape) phi0 = np.reshape(phi[0], phi[0].shape[:-self.ndim] + (D,)) phi1 = np.reshape(phi[1], phi[1].shape[:-2*self.ndim] + (D,D)) # Compute the moments L = linalg.chol(-2*phi1) Cov = linalg.chol_inv(L) u0 = linalg.chol_solve(L, phi0) u1 = linalg.outer(u0, u0) + Cov # Compute CGF g = (- 0.5 * np.einsum('...i,...i', u0, phi0) + 0.5 * linalg.chol_logdet(L)) # Reshape to arrays u0 = np.reshape(u0, u0.shape[:-1] + self.shape) u1 = np.reshape(u1, u1.shape[:-2] + self.shape + self.shape) u = [u0, u1] return (u, g) def compute_cgf_from_parents(self, u_mu_alpha): r""" Compute the value of the cumulant generating function. """ # Compute sum(mu^2 * alpha) correctly for broadcasted shapes alpha_mu2 = u_mu_alpha[1] logdet_alpha = u_mu_alpha[3] axes = tuple(range(-self.ndim, 0)) # TODO/FIXME: You could use plate multiplier type of correction instead # of explicitly broadcasting with ones. if self.ndim > 0: alpha_mu2 = misc.sum_multiply(alpha_mu2, np.ones(self.shape), axis=axes) if self.ndim > 0: logdet_alpha = misc.sum_multiply(logdet_alpha, np.ones(self.shape), axis=axes) # Compute g g = -0.5*alpha_mu2 + 0.5*logdet_alpha return g def compute_fixed_moments_and_f(self, x, mask=True): r""" Compute u(x) and f(x) for given x. """ if self.ndim > 0 and np.shape(x)[-self.ndim:] != self.shape: raise ValueError("Invalid shape") k = np.prod(self.shape) u = [x, linalg.outer(x, x, ndim=self.ndim)] f = -k/2*np.log(2*np.pi) return (u, f) def plates_to_parent(self, index, plates): r""" Resolves the plate mapping to a parent. Given the plates of the node's moments, this method returns the plates that the message to a parent has for the parent's distribution. """ if index != 0: raise IndexError() return plates + self.shape def plates_from_parent(self, index, plates): r""" Resolve the plate mapping from a parent. Given the plates of a parent's moments, this method returns the plates that the moments has for this distribution. """ if index != 0: raise IndexError() if self.ndim == 0: return plates else: return plates[:-self.ndim] def random(self, *phi, plates=None): r""" Draw a random sample from the Gaussian distribution. """ # TODO/FIXME: You shouldn't draw random values for # observed/fixed elements! D = self.ndim if D == 0: dims = () else: dims = np.shape(phi[0])[-D:] if np.prod(dims) == 1.0: # Scalar Gaussian phi1 = phi[1] if D > 0: # Because the covariance matrix has shape (1,1,...,1,1), # that is 2*D number of ones, remove the extra half of the # shape phi1 = np.reshape(phi1, np.shape(phi1)[:-2*D] + D*(1,)) var = -0.5 / phi1 std = np.sqrt(var) mu = var * phi[0] shape = plates + dims z = np.random.randn(*shape) x = mu + std * z else: N = np.prod(dims) dims_cov = dims + dims # Reshape precision matrix plates_cov = np.shape(phi[1])[:-2*D] V = -2 * np.reshape(phi[1], plates_cov + (N,N)) # Compute Cholesky U = linalg.chol(V) # Reshape mean vector plates_phi0 = np.shape(phi[0])[:-D] phi0 = np.reshape(phi[0], plates_phi0 + (N,)) mu = linalg.chol_solve(U, phi0) # Compute mu + U\z shape = plates + (N,) z = np.random.randn(*shape) # Denote Lambda = -2*phi[1] # Then, Cov = inv(Lambda) = inv(U'*U) = inv(U) * inv(U') # Thus, compute mu + U\z x = mu + linalg.solve_triangular(U, z, trans='N', lower=False) x = np.reshape(x, plates + dims) return x def compute_gradient(self, g, u, phi): r""" Compute the standard gradient with respect to the natural parameters. Gradient of the moments: .. math:: \mathrm{d}\overline{\mathbf{u}} &= \begin{bmatrix} \frac{1}{2} \phi_2^{-1} \mathrm{d}\phi_2 \phi_2^{-1} \phi_1 - \frac{1}{2} \phi_2^{-1} \mathrm{d}\phi_1 \\ - \frac{1}{4} \phi_2^{-1} \mathrm{d}\phi_2 \phi_2^{-1} \phi_1 \phi_1^{\mathrm{T}} \phi_2^{-1} - \frac{1}{4} \phi_2^{-1} \phi_1 \phi_1^{\mathrm{T}} \phi_2^{-1} \mathrm{d}\phi_2 \phi_2^{-1} + \frac{1}{2} \phi_2^{-1} \mathrm{d}\phi_2 \phi_2^{-1} + \frac{1}{4} \phi_2^{-1} \mathrm{d}\phi_1 \phi_1^{\mathrm{T}} \phi_2^{-1} + \frac{1}{4} \phi_2^{-1} \phi_1 \mathrm{d}\phi_1^{\mathrm{T}} \phi_2^{-1} \end{bmatrix} \\ &= \begin{bmatrix} 2 (\overline{u}_2 - \overline{u}_1 \overline{u}_1^{\mathrm{T}}) \mathrm{d}\phi_2 \overline{u}_1 + (\overline{u}_2 - \overline{u}_1 \overline{u}_1^{\mathrm{T}}) \mathrm{d}\phi_1 \\ u_2 d\phi_2 u_2 - 2 u_1 u_1^T d\phi_2 u_1 u_1^T + 2 (u_2 - u_1 u_1^T) d\phi_1 u_1^T \end{bmatrix} Standard gradient given the gradient with respect to the moments, that is, given the Riemannian gradient :math:`\tilde{\nabla}`: .. math:: \nabla = \begin{bmatrix} (\overline{u}_2 - \overline{u}_1 \overline{u}_1^{\mathrm{T}}) \tilde{\nabla}_1 + 2 (u_2 - u_1 u_1^T) \tilde{\nabla}_2 u_1 \\ (u_2 - u_1 u_1^T) \tilde{\nabla}_1 u_1^T + u_1 \tilde{\nabla}_1^T (u_2 - u_1 u_1^T) + 2 u_2 \tilde{\nabla}_2 u_2 - 2 u_1 u_1^T \tilde{\nabla}_2 u_1 u_1^T \end{bmatrix} """ ndim = self.ndim x = u[0] xx = u[1] # Some helpful variables x_x = linalg.outer(x, x, ndim=ndim) Cov = xx - x_x cov_g0 = linalg.mvdot(Cov, g[0], ndim=ndim) cov_g0_x = linalg.outer(cov_g0, x, ndim=ndim) g1_x = linalg.mvdot(g[1], x, ndim=ndim) # Compute gradient terms d0 = cov_g0 + 2 * linalg.mvdot(Cov, g1_x, ndim=ndim) d1 = (cov_g0_x + linalg.transpose(cov_g0_x, ndim=ndim) + 2 * linalg.mmdot(xx, linalg.mmdot(g[1], xx, ndim=ndim), ndim=ndim) - 2 * x_x * misc.add_trailing_axes(linalg.inner(g1_x, x, ndim=ndim), 2*ndim)) return [d0, d1] class GaussianGammaDistribution(ExponentialFamilyDistribution): r""" Class for the VMP formulas of Gaussian-Gamma-ISO variables. Currently, supports only vector variables. Log pdf of the prior: .. math:: \log p(\mathbf{x}, \tau | \boldsymbol{\mu}, \mathbf{\Lambda}, a, b) =& - \frac{1}{2} \tau \mathbf{x}^T \mathbf{\Lambda} \mathbf{x} + \frac{1}{2} \tau \mathbf{x}^T \mathbf{\Lambda} \boldsymbol{\mu} + \frac{1}{2} \tau \boldsymbol{\mu}^T \mathbf{\Lambda} \mathbf{x} - \frac{1}{2} \tau \boldsymbol{\mu}^T \mathbf{\Lambda} \boldsymbol{\mu} + \frac{1}{2} \log|\mathbf{\Lambda}| + \frac{D}{2} \log\tau - \frac{D}{2} \log(2\pi) \\ & - b \tau + a \log\tau - \log\tau + a \log b - \log \Gamma(a) Log pdf of the posterior approximation: .. math:: \log q(\mathbf{x}, \tau) =& \tau \mathbf{x}^T \boldsymbol{\phi}_1 + \tau \mathbf{x}^T \mathbf{\Phi}_2 \mathbf{x} + \tau \phi_3 + \log\tau \phi_4 + g(\boldsymbol{\phi}_1, \mathbf{\Phi}_2, \phi_3, \phi_4) + f(x, \tau) """ def __init__(self, shape): self.shape = shape self.ndim = len(shape) super().__init__() def compute_message_to_parent(self, parent, index, u, u_mu_Lambda, u_a, u_b): r""" Compute the message to a parent node. - Parent :math:`(\boldsymbol{\mu}, \mathbf{\Lambda})` Moments: .. math:: \begin{bmatrix} \mathbf{\Lambda}\boldsymbol{\mu} \\ \boldsymbol{\mu}^T\mathbf{\Lambda}\boldsymbol{\mu} \\ \mathbf{\Lambda} \\ \log|\mathbf{\Lambda}| \end{bmatrix} Message: .. math:: \begin{bmatrix} \langle \tau \mathbf{x} \rangle \\ - \frac{1}{2} \langle \tau \rangle \\ - \frac{1}{2} \langle \tau \mathbf{xx}^T \rangle \\ \frac{1}{2} \end{bmatrix} - Parent :math:`a`: Moments: .. math:: \begin{bmatrix} a \\ \log \Gamma(a) \end{bmatrix} Message: .. math:: \begin{bmatrix} \langle \log\tau \rangle + \langle \log b \rangle \\ -1 \end{bmatrix} - Parent :math:`b`: Moments: .. math:: \begin{bmatrix} b \\ \log b \end{bmatrix} Message: .. math:: \begin{bmatrix} - \langle \tau \rangle \\ \langle a \rangle \end{bmatrix} """ x_tau = u[0] xx_tau = u[1] tau = u[2] logtau = u[3] if index == 0: m0 = x_tau m1 = -0.5 * tau m2 = -0.5 * xx_tau m3 = 0.5 return [m0, m1, m2, m3] elif index == 1: logb = u_b[1] m0 = logtau + logb m1 = -1 return [m0, m1] elif index == 2: a = u_a[0] m0 = -tau m1 = a return [m0, m1] else: raise ValueError("Index out of bounds") def compute_phi_from_parents(self, u_mu_Lambda, u_a, u_b, mask=True): r""" Compute the natural parameter vector given parent moments. """ Lambda_mu = u_mu_Lambda[0] mu_Lambda_mu = u_mu_Lambda[1] Lambda = u_mu_Lambda[2] a = u_a[0] b = u_b[0] phi = [Lambda_mu, -0.5*Lambda, -0.5*mu_Lambda_mu - b, a] return phi def compute_moments_and_cgf(self, phi, mask=True): r""" Compute the moments and :math:`g(\phi)`. """ # Compute helpful variables V = -2*phi[1] L_V = linalg.chol(V, ndim=self.ndim) logdet_V = linalg.chol_logdet(L_V, ndim=self.ndim) mu = linalg.chol_solve(L_V, phi[0], ndim=self.ndim) Cov = linalg.chol_inv(L_V, ndim=self.ndim) a = phi[3] b = -phi[2] - 0.5 * linalg.inner(mu, phi[0], ndim=self.ndim) log_b = np.log(b) # Compute moments u2 = a / b u3 = -log_b + special.psi(a) u0 = mu * misc.add_trailing_axes(u2, self.ndim) u1 = Cov + ( linalg.outer(mu, mu, ndim=self.ndim) * misc.add_trailing_axes(u2, 2 * self.ndim) ) u = [u0, u1, u2, u3] # Compute g g = 0.5*logdet_V + a*log_b - special.gammaln(a) return (u, g) def compute_cgf_from_parents(self, u_mu_Lambda, u_a, u_b): r""" Compute :math:`\mathrm{E}_{q(p)}[g(p)]` """ logdet_Lambda = u_mu_Lambda[3] a = u_a[0] gammaln_a = u_a[1] log_b = u_b[1] g = 0.5*logdet_Lambda + a*log_b - gammaln_a return g def compute_fixed_moments_and_f(self, x_alpha, mask=True): r""" Compute the moments and :math:`f(x)` for a fixed value. """ (x, alpha) = x_alpha logalpha = np.log(alpha) u0 = x * misc.add_trailing_axes(alpha, self.ndim) u1 = linalg.outer(x, x, ndim=self.ndim) * misc.add_trailing_axes(alpha, 2*self.ndim) u2 = alpha u3 = logalpha u = [u0, u1, u2, u3] if self.ndim > 0: D = np.prod(np.shape(x)[-self.ndim:]) else: D = 1 f = (D/2 - 1) * logalpha - D/2 * np.log(2*np.pi) return (u, f) def random(self, *phi, plates=None): r""" Draw a random sample from the distribution. """ # TODO/FIXME: This is incorrect, I think. Gamma distribution parameters # aren't directly those, because phi has some parts from the Gaussian # distribution. alpha = GammaDistribution().random( phi[2], phi[3], plates=plates ) mu = GaussianARDDistribution(self.shape).random( misc.add_trailing_axes(alpha, self.ndim) * phi[0], misc.add_trailing_axes(alpha, 2*self.ndim) * phi[1], plates=plates ) return (mu, alpha) class GaussianWishartDistribution(ExponentialFamilyDistribution): r""" Class for the VMP formulas of Gaussian-Wishart variables. Currently, supports only vector variables. .. math:: \log p(\mathbf{x}, \mathbf{\Lambda} | \boldsymbol{\mu}, \alpha, n, \mathbf{V}) =& - \frac{1}{2} \alpha \mathbf{x}^T \mathbf{\Lambda} \mathbf{x} + \frac{1}{2} \alpha \mathbf{x}^T \mathbf{\Lambda} \boldsymbol{\mu} + \frac{1}{2} \alpha \boldsymbol{\mu}^T \mathbf{\Lambda} \mathbf{x} - \frac{1}{2} \alpha \boldsymbol{\mu}^T \mathbf{\Lambda} \boldsymbol{\mu} + \frac{1}{2} \log|\mathbf{\Lambda}| + \frac{D}{2} \log\alpha - \frac{D}{2} \log(2\pi) \\ & - \frac{1}{2} \mathrm{tr}(\mathbf{V}\mathbf{\Lambda}) + \frac{n-d-1}{2} \log|\mathbf{\Lambda}| - \frac{nd}{2}\log 2 - \frac{n}{2} \log|\mathbf{V}| - \log\Gamma_d(\frac{n}{2}) Posterior approximation: .. math:: \log q(\mathbf{x}, \mathbf{\Lambda}) =& \mathbf{x}^T \mathbf{\Lambda} \boldsymbol{\phi}_1 + \phi_2 \mathbf{x}^T \mathbf{\Lambda} \mathbf{x} + \mathrm{tr}(\mathbf{\Lambda} \mathbf{\Phi}_3) + \phi_4 \log|\mathbf{\Lambda}| + g(\boldsymbol{\phi}_1, \phi_2, \mathbf{\Phi}_3, \phi_4) + f(\mathbf{x}, \mathbf{\Lambda}) """ def compute_message_to_parent(self, parent, index, u, u_mu_alpha, u_n, u_V): r""" Compute the message to a parent node. For parent :math:`q(\boldsymbol{\mu}, \alpha)`: .. math:: \alpha \boldsymbol{\mu}^T \mathbf{m}_1 \Rightarrow & \mathbf{m}_1 = \langle \mathbf{\Lambda x} \rangle \\ \alpha \boldsymbol{\mu}^T \mathbf{M}_2 \boldsymbol{\mu} \Rightarrow & \mathbf{M}_2 = - \frac{1}{2} \langle \mathbf{\Lambda} \rangle \\ \alpha m_3 \Rightarrow & m_3 = - \frac{1}{2} \langle \mathbf{x}^T \mathbf{\Lambda} \mathbf{x} \rangle \\ m_4 \log \alpha \Rightarrow & m_4 = \frac{d}{2} For parent :math:`q(\mathbf{V})`: .. math:: \mathbf{M}_1 &= \frac{\partial \langle \log p \rangle}{\partial \langle \mathbf{V} \rangle} = -\frac{1}{2} \langle \mathbf{\Lambda} \rangle \\ \mathbf{M}_2 &= \frac{\partial \langle \log p \rangle}{\partial \langle \log|\mathbf{V}| \rangle} = ... """ if index == 0: m0 m1 m2 m3 raise NotImplementedError() elif index == 1: raise NotImplementedError() elif index == 2: raise NotImplementedError() else: raise ValueError("Index out of bounds") def compute_phi_from_parents(self, u_mu_alpha, u_n, u_V, mask=True): r""" Compute the natural parameter vector given parent moments. """ alpha_mu = u_mu_alpha[0] alpha_mumu = u_mu_alpha[1] alpha = u_mu_alpha[2] V = u_V[0] n = u_n[0] phi0 = alpha_mu phi1 = -0.5 * alpha phi2 = -0.5 * (V + alpha_mumu) phi3 = 0.5 * n return [phi0, phi1, phi2, phi3] def compute_moments_and_cgf(self, phi, mask=True): r""" Compute the moments and :math:`g(\phi)`. """ # TODO/FIXME: This isn't probably correct. Phi[2:] has terms that are # related to the Gaussian also, not only Wishart. u_Lambda = WishartDistribution((D,)).compute_moments_and_cgf(phi[2:]) raise NotImplementedError() return (u, g) def compute_cgf_from_parents(self, u_mu_alpha, u_n, u_V): r""" Compute :math:`\mathrm{E}_{q(p)}[g(p)]` """ raise NotImplementedError() return g def compute_fixed_moments_and_f(self, x, Lambda, mask=True): r""" Compute the moments and :math:`f(x)` for a fixed value. """ raise NotImplementedError() return (u, f) def random(self, *params, plates=None): r""" Draw a random sample from the distribution. """ raise NotImplementedError() # # NODES # class _GaussianTemplate(ExponentialFamily): def translate(self, b, debug=False): """ Transforms the current posterior by adding a bias to the mean Parameters ---------- b : array Constant to add """ ndim = len(self.dims[0]) if ndim > 0 and np.shape(b)[-ndim:] != self.dims[0]: raise ValueError("Bias has incorrect shape") x = self.u[0] xb = linalg.outer(x, b, ndim=ndim) bx = linalg.transpose(xb, ndim=ndim) bb = linalg.outer(b, b, ndim=ndim) uh = [ self.u[0] + b, self.u[1] + xb + bx + bb ] Lambda = -2 * self.phi[1] Lambda_b = linalg.mvdot(Lambda, b, ndim=ndim) dg = -0.5 * ( linalg.inner(b, Lambda_b, ndim=ndim) + 2 * linalg.inner(x, Lambda_b, ndim=ndim) ) phih = [ self.phi[0] + Lambda_b, self.phi[1] ] self._check_shape(uh) self._check_shape(phih) self.u = uh self.phi = phih self.g = self.g + dg # TODO: This is all just debugging stuff and can be removed if debug: uh = [ui.copy() for ui in uh] gh = self.g.copy() self._update_moments_and_cgf() if any(not np.allclose(uih, ui, atol=1e-6) for (uih, ui) in zip(uh, self.u)): raise RuntimeError("BUG") if not np.allclose(self.g, gh, atol=1e-6): raise RuntimeError("BUG") return class Gaussian(_GaussianTemplate): r""" Node for Gaussian variables. The node represents a :math:`D`-dimensional vector from the Gaussian distribution: .. math:: \mathbf{x} &\sim \mathcal{N}(\boldsymbol{\mu}, \mathbf{\Lambda}), where :math:`\boldsymbol{\mu}` is the mean vector and :math:`\mathbf{\Lambda}` is the precision matrix (i.e., inverse of the covariance matrix). .. math:: \mathbf{x},\boldsymbol{\mu} \in \mathbb{R}^{D}, \quad \mathbf{\Lambda} \in \mathbb{R}^{D \times D}, \quad \mathbf{\Lambda} \text{ symmetric positive definite} Parameters ---------- mu : Gaussian-like node or GaussianGamma-like node or GaussianWishart-like node or array Mean vector Lambda : Wishart-like node or array Precision matrix See also -------- Wishart, GaussianARD, GaussianWishart, GaussianGamma """ def __init__(self, mu, Lambda, **kwargs): r""" Create Gaussian node """ super().__init__(mu, Lambda, **kwargs) @classmethod def _constructor(cls, mu, Lambda, ndim=1, **kwargs): r""" Constructs distribution and moments objects. """ mu_Lambda = WrapToGaussianWishart(mu, Lambda, ndim=ndim) shape = mu_Lambda._moments.shape moments = GaussianMoments(shape) parent_moments = (mu_Lambda._moments,) if mu_Lambda.dims != ( shape, (), shape+shape, () ): raise Exception("Parents have wrong dimensionality") distribution = GaussianDistribution(shape) parents = [mu_Lambda] return (parents, kwargs, moments.dims, cls._total_plates(kwargs.get('plates'), distribution.plates_from_parent(0, mu_Lambda.plates)), distribution, moments, parent_moments) def initialize_from_parameters(self, mu, Lambda): u = self._parent_moments[0].compute_fixed_moments(mu, Lambda) self._initialize_from_parent_moments(u) def observe_limits(self, minimum=-np.inf, maximum=np.inf): self._distribution.set_limits(minimum, maximum) self._update_mask() return def _set_mask(self, mask): self.mask = np.logical_or( mask, np.logical_or( self.observed, self._distribution.has_limits, ), ) def __str__(self): ndim = len(self.dims[0]) mu = self.u[0] Cov = self.u[1] - linalg.outer(mu, mu, ndim=ndim) return ("%s ~ Gaussian(mu, Cov)\n" " mu = \n" "%s\n" " Cov = \n" "%s\n" % (self.name, mu, Cov)) def rotate(self, R, inv=None, logdet=None, Q=None): # TODO/FIXME: Combine and refactor all these rotation transformations # into _GaussianTemplate if self._moments.ndim != 1: raise NotImplementedError("Not implemented for ndim!=1 yet") if inv is not None: invR = inv else: invR = np.linalg.inv(R) if logdet is not None: logdetR = logdet else: logdetR = np.linalg.slogdet(R)[1] # It would be more efficient and simpler, if you just rotated the # moments and didn't touch phi. However, then you would need to call # update() before lower_bound_contribution. This is more error-safe. # Rotate plates, if plate rotation matrix is given. Assume that there's # only one plate-axis if Q is not None: # Rotate moments using Q self.u[0] = np.einsum('ik,kj->ij', Q, self.u[0]) sumQ = np.sum(Q, axis=0) # Rotate natural parameters using Q self.phi[1] = np.einsum('d,dij->dij', sumQ**(-2), self.phi[1]) self.phi[0] = np.einsum('dij,dj->di', -2*self.phi[1], self.u[0]) # Transform parameters using R self.phi[0] = mvdot(invR.T, self.phi[0]) self.phi[1] = dot(invR.T, self.phi[1], invR) if Q is not None: self._update_moments_and_cgf() else: # Transform moments and g using R self.u[0] = mvdot(R, self.u[0]) self.u[1] = dot(R, self.u[1], R.T) self.g -= logdetR def rotate_matrix(self, R1, R2, inv1=None, logdet1=None, inv2=None, logdet2=None, Q=None): r""" The vector is reshaped into a matrix by stacking the row vectors. Computes R1*X*R2', which is identical to kron(R1,R2)*x (??) Note that this is slightly different from the standard Kronecker product definition because Numpy stacks row vectors instead of column vectors. Parameters ---------- R1 : ndarray A matrix from the left R2 : ndarray A matrix from the right """ if self._moments.ndim != 1: raise NotImplementedError("Not implemented for ndim!=1 yet") if Q is not None: # Rotate moments using Q self.u[0] = np.einsum('ik,kj->ij', Q, self.u[0]) sumQ = np.sum(Q, axis=0) # Rotate natural parameters using Q self.phi[1] = np.einsum('d,dij->dij', sumQ**(-2), self.phi[1]) self.phi[0] = np.einsum('dij,dj->di', -2*self.phi[1], self.u[0]) if inv1 is None: inv1 = np.linalg.inv(R1) if logdet1 is None: logdet1 = np.linalg.slogdet(R1)[1] if inv2 is None: inv2 = np.linalg.inv(R2) if logdet2 is None: logdet2 = np.linalg.slogdet(R2)[1] D1 = np.shape(R1)[0] D2 = np.shape(R2)[0] # Reshape into matrices sh0 = np.shape(self.phi[0])[:-1] + (D1,D2) sh1 = np.shape(self.phi[1])[:-2] + (D1,D2,D1,D2) phi0 = np.reshape(self.phi[0], sh0) phi1 = np.reshape(self.phi[1], sh1) # Apply rotations to phi #phi0 = dot(inv1, phi0, inv2.T) phi0 = dot(inv1.T, phi0, inv2) phi1 = np.einsum('...ia,...abcd->...ibcd', inv1.T, phi1) phi1 = np.einsum('...ic,...abcd->...abid', inv1.T, phi1) phi1 = np.einsum('...ib,...abcd->...aicd', inv2.T, phi1) phi1 = np.einsum('...id,...abcd->...abci', inv2.T, phi1) # Reshape back into vectors self.phi[0] = np.reshape(phi0, self.phi[0].shape) self.phi[1] = np.reshape(phi1, self.phi[1].shape) # It'd be better to rotate the moments too.. self._update_moments_and_cgf() class GaussianARD(_GaussianTemplate): r""" Node for Gaussian variables with ARD prior. The node represents a :math:`D`-dimensional vector from the Gaussian distribution: .. math:: \mathbf{x} &\sim \mathcal{N}(\boldsymbol{\mu}, \mathrm{diag}(\boldsymbol{\alpha})), where :math:`\boldsymbol{\mu}` is the mean vector and :math:`\mathrm{diag}(\boldsymbol{\alpha})` is the diagonal precision matrix (i.e., inverse of the covariance matrix). .. math:: \mathbf{x},\boldsymbol{\mu} \in \mathbb{R}^{D}, \quad \alpha_d > 0 \text{ for } d=0,\ldots,D-1 *Note:* The form of the posterior approximation is a Gaussian distribution with full covariance matrix instead of a diagonal matrix. Parameters ---------- mu : Gaussian-like node or GaussianGamma-like node or array Mean vector alpha : gamma-like node or array Diagonal elements of the precision matrix See also -------- Gamma, Gaussian, GaussianGamma, GaussianWishart """ def __init__(self, mu, alpha, ndim=None, shape=None, **kwargs): r""" Create GaussianARD node. """ super().__init__(mu, alpha, ndim=ndim, shape=shape, **kwargs) @classmethod def _constructor(cls, mu, alpha, ndim=None, shape=None, **kwargs): r""" Constructs distribution and moments objects. If __init__ uses useconstructor decorator, this method is called to construct distribution and moments objects. The method is given the same inputs as __init__. For some nodes, some of these can't be "static" class attributes, then the node class must overwrite this method to construct the objects manually. The point of distribution class is to move general distribution but not-node specific code. The point of moments class is to define the messaging protocols. """ mu_alpha = WrapToGaussianGamma(mu, alpha, ndim=0) if ndim is None: if shape is not None: ndim = len(shape) else: shape = () ndim = 0 else: if shape is not None: if ndim != len(shape): raise ValueError("Given shape and ndim inconsistent") else: if ndim == 0: shape = () else: if ndim > len(mu_alpha.plates): raise ValueError( "Cannot determine shape for ndim={0} because parent " "full shape has ndim={1}." .format(ndim, len(mu_alpha.plates)) ) shape = mu_alpha.plates[-ndim:] moments = GaussianMoments(shape) parent_moments = [GaussianGammaMoments(())] distribution = GaussianARDDistribution(shape) plates = cls._total_plates(kwargs.get('plates'), distribution.plates_from_parent(0, mu_alpha.plates)) parents = [mu_alpha] return (parents, kwargs, moments.dims, plates, distribution, moments, parent_moments) def initialize_from_parameters(self, mu, alpha): # Explicit broadcasting so the shapes match mu = mu * np.ones(np.shape(alpha)) alpha = alpha * np.ones(np.shape(mu)) # Compute parent moments u = self._parent_moments[0].compute_fixed_moments([mu, alpha]) # Initialize distribution self._initialize_from_parent_moments(u) def initialize_from_mean_and_covariance(self, mu, Cov): ndim = len(self._distribution.shape) u = [mu, Cov + linalg.outer(mu, mu, ndim=ndim)] mask = np.logical_not(self.observed) # TODO: You could compute the CGF but it requires Cholesky of # Cov. Do it later. self._set_moments_and_cgf(u, np.nan, mask=mask) return def __str__(self): mu = self.u[0] Cov = self.u[1] - linalg.outer(mu, mu) return ("%s ~ Gaussian(mu, Cov)\n" " mu = \n" "%s\n" " Cov = \n" "%s\n" % (self.name, mu, Cov)) def rotate(self, R, inv=None, logdet=None, axis=-1, Q=None, subset=None, debug=False): if Q is not None: raise NotImplementedError() if subset is not None: raise NotImplementedError() # TODO/FIXME: Combine and refactor all these rotation transformations # into _GaussianTemplate ndim = len(self._distribution.shape) if inv is not None: invR = inv else: invR = np.linalg.inv(R) if logdet is not None: logdetR = logdet else: logdetR = np.linalg.slogdet(R)[1] self.phi[0] = rotate_mean(self.phi[0], invR.T, axis=axis, ndim=ndim) self.phi[1] = rotate_covariance(self.phi[1], invR.T, axis=axis, ndim=ndim) self.u[0] = rotate_mean(self.u[0], R, axis=axis, ndim=ndim) self.u[1] = rotate_covariance(self.u[1], R, axis=axis, ndim=ndim) s = list(self.dims[0]) s.pop(axis) self.g -= logdetR * np.prod(s) # TODO: This is all just debugging stuff and can be removed if debug: uh = [ui.copy() for ui in self.u] gh = self.g.copy() self._update_moments_and_cgf() if any(not np.allclose(uih, ui, atol=1e-6) for (uih, ui) in zip(uh, self.u)): raise RuntimeError("BUG") if not np.allclose(self.g, gh, atol=1e-6): raise RuntimeError("BUG") return def rotate_plates(self, Q, plate_axis=-1): r""" Approximate rotation of a plate axis. Mean is rotated exactly but covariance/precision matrix is rotated approximately. """ ndim = len(self._distribution.shape) # Rotate moments using Q if not isinstance(plate_axis, int): raise ValueError("Plate axis must be integer") if plate_axis >= 0: plate_axis -= len(self.plates) if plate_axis < -len(self.plates) or plate_axis >= 0: raise ValueError("Axis out of bounds") u0 = rotate_mean(self.u[0], Q, ndim=ndim+(-plate_axis), axis=0) sumQ = misc.add_trailing_axes(np.sum(Q, axis=0), 2*ndim-plate_axis-1) phi1 = sumQ**(-2) * self.phi[1] phi0 = -2 * matrix_dot_vector(phi1, u0, ndim=ndim) self.phi[0] = phi0 self.phi[1] = phi1 self._update_moments_and_cgf() return class GaussianGamma(ExponentialFamily): r""" Node for Gaussian-gamma (isotropic) random variables. The prior: .. math:: p(x, \alpha| \mu, \Lambda, a, b) p(x|\alpha, \mu, \Lambda) = \mathcal{N}(x | \mu, \alpha Lambda) p(\alpha|a, b) = \mathcal{G}(\alpha | a, b) The posterior approximation :math:`q(x, \alpha)` has the same Gaussian-gamma form. Currently, supports only vector variables. """ @classmethod def _constructor(cls, mu, Lambda, a, b, ndim=1, **kwargs): r""" Constructs distribution and moments objects. This method is called if useconstructor decorator is used for __init__. `mu` is the mean/location vector `alpha` is the scale `V` is the scale matrix `n` is the degrees of freedom """ # Convert parent nodes mu_Lambda = WrapToGaussianWishart(mu, Lambda, ndim=ndim) a = cls._ensure_moments(a, GammaPriorMoments) b = cls._ensure_moments(b, GammaMoments) shape = mu_Lambda.dims[0] distribution = GaussianGammaDistribution(shape) moments = GaussianGammaMoments(shape) parent_moments = ( mu_Lambda._moments, a._moments, b._moments, ) # Check shapes if mu_Lambda.dims != ( shape, (), 2*shape, () ): raise ValueError("mu and Lambda have wrong shape") if a.dims != ( (), () ): raise ValueError("a has wrong shape") if b.dims != ( (), () ): raise ValueError("b has wrong shape") # List of parent nodes parents = [mu_Lambda, a, b] return (parents, kwargs, moments.dims, cls._total_plates(kwargs.get('plates'), distribution.plates_from_parent(0, mu_Lambda.plates), distribution.plates_from_parent(1, a.plates), distribution.plates_from_parent(2, b.plates)), distribution, moments, parent_moments) def translate(self, b, debug=False): if self._moments.ndim != 1: raise NotImplementedError("Only ndim=1 supported at the moment") tau = self.u[2] x = self.u[0] / tau[...,None] xb = linalg.outer(x, b, ndim=1) bx = linalg.transpose(xb, ndim=1) bb = linalg.outer(b, b, ndim=1) uh = [ self.u[0] + tau[...,None] * b, self.u[1] + tau[...,None,None] * (xb + bx + bb), self.u[2], self.u[3] ] Lambda = -2 * self.phi[1] dtau = -0.5 * ( np.einsum('...ij,...i,...j->...', Lambda, b, b) + 2 * np.einsum('...ij,...i,...j->...', Lambda, b, x) ) phih = [ self.phi[0] + np.einsum('...ij,...j->...i', Lambda, b), self.phi[1], self.phi[2] + dtau, self.phi[3] ] self._check_shape(uh) self._check_shape(phih) self.phi = phih self.u = uh # TODO: This is all just debugging stuff and can be removed if debug: uh = [ui.copy() for ui in uh] gh = self.g.copy() self._update_moments_and_cgf() if any(not np.allclose(uih, ui, atol=1e-6) for (uih, ui) in zip(uh, self.u)): raise RuntimeError("BUG") if not np.allclose(self.g, gh, atol=1e-6): raise RuntimeError("BUG") return def rotate(self, R, inv=None, logdet=None, debug=False): if self._moments.ndim != 1: raise NotImplementedError("Only ndim=1 supported at the moment") if inv is None: inv = np.linalg.inv(R) if logdet is None: logdet = np.linalg.slogdet(R)[1] uh = [ rotate_mean(self.u[0], R), rotate_covariance(self.u[1], R), self.u[2], self.u[3] ] phih = [ rotate_mean(self.phi[0], inv.T), rotate_covariance(self.phi[1], inv.T), self.phi[2], self.phi[3] ] self._check_shape(uh) self._check_shape(phih) self.phi = phih self.u = uh self.g = self.g - logdet # TODO: This is all just debugging stuff and can be removed if debug: uh = [ui.copy() for ui in uh] gh = self.g.copy() self._update_moments_and_cgf() if any(not np.allclose(uih, ui, atol=1e-6) for (uih, ui) in zip(uh, self.u)): raise RuntimeError("BUG") if not np.allclose(self.g, gh, atol=1e-6): raise RuntimeError("BUG") return def plotmatrix(self): r""" Creates a matrix of marginal plots. On diagonal, are marginal plots of each variable. Off-diagonal plot (i,j) shows the joint marginal density of x_i and x_j. """ import bayespy.plot as bpplt if self.ndim != 1: raise NotImplementedError("Only ndim=1 supported at the moment") if np.prod(self.plates) != 1: raise ValueError("Currently, does not support plates in the node.") if len(self.dims[0]) != 1: raise ValueError("Currently, supports only vector variables") # Dimensionality of the Gaussian D = self.dims[0][0] # Compute standard parameters tau = self.u[2] mu = self.u[0] mu = mu / misc.add_trailing_axes(tau, 1) Cov = self.u[1] - linalg.outer(self.u[0], mu, ndim=1) Cov = Cov / misc.add_trailing_axes(tau, 2) a = self.phi[3] b = -self.phi[2] - 0.5*linalg.inner(self.phi[0], mu, ndim=1) # Create subplots (fig, axes) = bpplt.pyplot.subplots(D+1, D+1) # Plot marginal Student t distributions for i in range(D): for j in range(i+1): if i == j: bpplt._pdf_t(*(random.gaussian_gamma_to_t(mu[i], Cov[i,i], a, b, ndim=0)), axes=axes[i,i]) else: S = Cov[np.ix_([i,j],[i,j])] (m, S, nu) = random.gaussian_gamma_to_t(mu[[i,j]], S, a, b) bpplt._contour_t(m, S, nu, axes=axes[i,j]) bpplt._contour_t(m, S, nu, axes=axes[j,i], transpose=True) # Plot Gaussian-gamma marginal distributions for k in range(D): bpplt._contour_gaussian_gamma(mu[k], Cov[k,k], a, b, axes=axes[D,k]) bpplt._contour_gaussian_gamma(mu[k], Cov[k,k], a, b, axes=axes[k,D], transpose=True) # Plot gamma marginal distribution bpplt._pdf_gamma(a, b, axes=axes[D,D]) return axes def get_gaussian_location(self): r""" Return the mean and variance of the distribution """ if self._moments.ndim != 1: raise NotImplementedError("Only ndim=1 supported at the moment") tau = self.u[2] tau_mu = self.u[0] return tau_mu / tau[...,None] def get_gaussian_mean_and_variance(self): r""" Return the mean and variance of the distribution """ if self.ndim != 1: raise NotImplementedError("Only ndim=1 supported at the moment") a = self.phi[3] nu = 2*a if np.any(nu <= 1): raise ValueError("Mean not defined for degrees of freedom <= 1") if np.any(nu <= 2): raise ValueError("Variance not defined if degrees of freedom <= 2") tau = self.u[2] tau_mu = self.u[0] mu = tau_mu / misc.add_trailing_axes(tau, 1) var = misc.get_diag(self.u[1], ndim=1) - tau_mu*mu var = var / misc.add_trailing_axes(tau, 1) var = nu / (nu-2) * var return (mu, var) def get_marginal_logpdf(self, gaussian=None, gamma=None): r""" Get the (marginal) log pdf of a subset of the variables Parameters ---------- gaussian : list or None Indices of the Gaussian variables to keep or None gamma : bool or None True if keep the gamma variable, otherwise False or None Returns ------- function A function which computes log-pdf """ if self.ndim != 1: raise NotImplementedError("Only ndim=1 supported at the moment") if gaussian is None and not gamma: raise ValueError("Must give some variables") # Compute standard parameters tau = self.u[2] mu = self.u[0] mu = mu / misc.add_trailing_axes(tau, 1) Cov = np.linalg.inv(-2*self.phi[1]) if not np.allclose(Cov, self.u[1] - linalg.outer(self.u[0], mu, ndim=1)): raise Exception("WAAAT") #Cov = Cov / misc.add_trailing_axes(tau, 2) a = self.phi[3] b = -self.phi[2] - 0.5*linalg.inner(self.phi[0], mu, ndim=1) if not gamma: # Student t distributions inds = list(gaussian) mu = mu[inds] Cov = Cov[np.ix_(inds, inds)] (mu, Cov, nu) = random.gaussian_gamma_to_t(mu, Cov, a, b, ndim=1) L = linalg.chol(Cov) logdet_Cov = linalg.chol_logdet(L) D = len(inds) def logpdf(x): y = x - mu v = linalg.chol_solve(L, y) z2 = linalg.inner(y, v, ndim=1) return random.t_logpdf(z2, logdet_Cov, nu, D) return logpdf elif gaussian is None: # Gamma distribution def logpdf(x): logx = np.log(x) return random.gamma_logpdf(b*x, logx, a*logx, a*np.log(b), special.gammaln(a)) return logpdf else: # Gaussian-gamma distribution inds = list(gaussian) mu = mu[inds] Cov = Cov[np.ix_(inds, inds)] D = len(inds) L = linalg.chol(Cov) logdet_Cov = linalg.chol_logdet(L) def logpdf(x): tau = x[...,-1] logtau = np.log(tau) x = x[...,:-1] y = x - mu v = linalg.chol_solve(L, y) * tau[...,None] z2 = linalg.inner(y, v, ndim=1) return (random.gaussian_logpdf(z2, 0, 0, logdet_Cov + D*logtau, D) + random.gamma_logpdf(b*tau, logtau, a*logtau, a*np.log(b), special.gammaln(a))) return logpdf class GaussianWishart(ExponentialFamily): r""" Node for Gaussian-Wishart random variables. The prior: .. math:: p(x, \Lambda| \mu, \alpha, V, n) p(x|\Lambda, \mu, \alpha) = \mathcal(N)(x | \mu, \alpha^{-1} Lambda^{-1}) p(\Lambda|V, n) = \mathcal(W)(\Lambda | n, V) The posterior approximation :math:`q(x, \Lambda)` has the same Gaussian-Wishart form. Currently, supports only vector variables. """ _distribution = GaussianWishartDistribution() @classmethod def _constructor(cls, mu, alpha, n, V, **kwargs): r""" Constructs distribution and moments objects. This method is called if useconstructor decorator is used for __init__. `mu` is the mean/location vector `alpha` is the scale `n` is the degrees of freedom `V` is the scale matrix """ # Convert parent nodes mu_alpha = WrapToGaussianGamma(mu, alpha, ndim=1) D = mu_alpha.dims[0][0] shape = mu_alpha._moments.shape moments = GaussianWishartMoments(shape) n = cls._ensure_moments(n, WishartPriorMoments, d=D) V = cls._ensure_moments(V, WishartMoments, ndim=1) parent_moments = ( mu_alpha._moments, n._moments, V._moments ) # Check shapes if mu_alpha.dims != ( (D,), (D,D), (), () ): raise ValueError("mu and alpha have wrong shape") if V.dims != ( (D,D), () ): raise ValueError("Precision matrix has wrong shape") if n.dims != ( (), () ): raise ValueError("Degrees of freedom has wrong shape") parents = [mu_alpha, n, V] return (parents, kwargs, moments.dims, cls._total_plates(kwargs.get('plates'), cls._distribution.plates_from_parent(0, mu_alpha.plates), cls._distribution.plates_from_parent(1, n.plates), cls._distribution.plates_from_parent(2, V.plates)), cls._distribution, moments, parent_moments) # # CONVERTERS # class GaussianToGaussianGamma(Deterministic): r""" Converter for Gaussian moments to Gaussian-gamma isotropic moments Combines the Gaussian moments with gamma moments for a fixed value 1. """ def __init__(self, X, **kwargs): r""" """ if not isinstance(X._moments, GaussianMoments): raise ValueError("Wrong moments, should be Gaussian") shape = X._moments.shape self.ndim = X._moments.ndim self._moments = GaussianGammaMoments(shape) self._parent_moments = [GaussianMoments(shape)] shape = X.dims[0] dims = ( shape, 2*shape, (), () ) super().__init__(X, dims=dims, **kwargs) def _compute_moments(self, u_X): r""" """ x = u_X[0] xx = u_X[1] u = [x, xx, 1, 0] return u def _compute_message_to_parent(self, index, m_child, u_X): r""" """ if index == 0: m = m_child[:2] return m else: raise ValueError("Invalid parent index") def _compute_function(self, x): return (x, 1) GaussianMoments.add_converter(GaussianGammaMoments, GaussianToGaussianGamma) class GaussianGammaToGaussianWishart(Deterministic): r""" """ def __init__(self, X_alpha, **kwargs): raise NotImplementedError() GaussianGammaMoments.add_converter(GaussianWishartMoments, GaussianGammaToGaussianWishart) # # WRAPPERS # # These wrappers form a single node from two nodes for messaging purposes. # class WrapToGaussianGamma(Deterministic): r""" """ def __init__(self, X, alpha, ndim=None, **kwargs): r""" """ # In case X is a numerical array, convert it to Gaussian first try: X = self._ensure_moments(X, GaussianMoments, ndim=ndim) except Moments.NoConverterError: pass try: ndim = X._moments.ndim except AttributeError as err: raise TypeError("ndim needs to be given explicitly") from err X = self._ensure_moments(X, GaussianGammaMoments, ndim=ndim) if len(X.dims[0]) != ndim: raise RuntimeError("Conversion failed ndim.") shape = X.dims[0] dims = ( shape, 2 * shape, (), () ) self.shape = shape self.ndim = len(shape) self._moments = GaussianGammaMoments(shape) self._parent_moments = [ GaussianGammaMoments(shape), GammaMoments() ] super().__init__(X, alpha, dims=dims, **kwargs) def _compute_moments(self, u_X, u_alpha): r""" """ (tau_x, tau_xx, tau, logtau) = u_X (alpha, logalpha) = u_alpha u0 = tau_x * misc.add_trailing_axes(alpha, self.ndim) u1 = tau_xx * misc.add_trailing_axes(alpha, 2 * self.ndim) u2 = tau * alpha u3 = logtau + logalpha return [u0, u1, u2, u3] def _compute_message_to_parent(self, index, m_child, u_X, u_alpha): r""" """ if index == 0: alpha = u_alpha[0] m0 = m_child[0] * misc.add_trailing_axes(alpha, self.ndim) m1 = m_child[1] * misc.add_trailing_axes(alpha, 2 * self.ndim) m2 = m_child[2] * alpha m3 = m_child[3] return [m0, m1, m2, m3] elif index == 1: (tau_x, tau_xx, tau, logtau) = u_X m0 = ( linalg.inner(m_child[0], tau_x, ndim=self.ndim) + linalg.inner(m_child[1], tau_xx, ndim=2*self.ndim) + m_child[2] * tau ) m1 = m_child[3] return [m0, m1] else: raise ValueError("Invalid parent index") class WrapToGaussianWishart(Deterministic): r""" Wraps Gaussian and Wishart nodes into a Gaussian-Wishart node. The following node combinations can be wrapped: * Gaussian and Wishart * Gaussian-gamma and Wishart * Gaussian-Wishart and gamma """ def __init__(self, X, Lambda, ndim=1, **kwargs): r""" """ # Just in case X is an array, convert it to a Gaussian node first. try: X = self._ensure_moments(X, GaussianMoments, ndim=ndim) except Moments.NoConverterError: pass try: # Try combo Gaussian-Gamma and Wishart X = self._ensure_moments(X, GaussianGammaMoments, ndim=ndim) except Moments.NoConverterError: # Have to use Gaussian-Wishart and Gamma X = self._ensure_moments(X, GaussianWishartMoments, ndim=ndim) Lambda = self._ensure_moments(Lambda, GammaMoments, ndim=ndim) shape = X.dims[0] if Lambda.dims != ((), ()): raise ValueError( "Mean and precision have inconsistent shapes: {0} and {1}" .format( X.dims, Lambda.dims ) ) self.wishart = False else: # Gaussian-Gamma and Wishart shape = X.dims[0] Lambda = self._ensure_moments(Lambda, WishartMoments, ndim=ndim) if Lambda.dims != (2 * shape, ()): raise ValueError( "Mean and precision have inconsistent shapes: {0} and {1}" .format( X.dims, Lambda.dims ) ) self.wishart = True self.ndim = len(shape) self._parent_moments = ( X._moments, Lambda._moments, ) self._moments = GaussianWishartMoments(shape) super().__init__(X, Lambda, dims=self._moments.dims, **kwargs) def _compute_moments(self, u_X_alpha, u_Lambda): r""" """ if self.wishart: alpha_x = u_X_alpha[0] alpha_xx = u_X_alpha[1] alpha = u_X_alpha[2] log_alpha = u_X_alpha[3] Lambda = u_Lambda[0] logdet_Lambda = u_Lambda[1] D = np.prod(self.dims[0]) u0 = linalg.mvdot(Lambda, alpha_x, ndim=self.ndim) u1 = linalg.inner(Lambda, alpha_xx, ndim=2*self.ndim) u2 = Lambda * misc.add_trailing_axes(alpha, 2*self.ndim) u3 = logdet_Lambda + D * log_alpha u = [u0, u1, u2, u3] return u else: raise NotImplementedError() def _compute_message_to_parent(self, index, m_child, u_X_alpha, u_Lambda): r""" ... Message from the child is :math:`[m_0, m_1, m_2, m_3]`: .. math:: \alpha m_0^T \Lambda x + m_1 \alpha x^T \Lambda x + \mathrm{tr}(\alpha m_2 \Lambda) + m_3 (\log | \alpha \Lambda |) In case of Gaussian-gamma and Wishart parents: Message to the first parent (x, alpha): .. math:: \tilde{m_0} &= \Lambda m_0 \\ \tilde{m_1} &= m_1 \Lambda \\ \tilde{m_2} &= \mathrm{tr}(m_2 \Lambda) \\ \tilde{m_3} &= m_3 \cdot D Message to the second parent (Lambda): .. math:: \tilde{m_0} &= \alpha (\frac{1}{2} m_0 x^T + \frac{1}{2} x m_0^T + m_1 xx^T + m_2) \\ \tilde{m_1} &= m_3 """ if index == 0: if self.wishart: # Message to Gaussian-gamma (isotropic) Lambda = u_Lambda[0] D = np.prod(self.dims[0]) m0 = linalg.mvdot(Lambda, m_child[0], ndim=self.ndim) m1 = Lambda * misc.add_trailing_axes(m_child[1], 2*self.ndim) m2 = linalg.inner(Lambda, m_child[2], ndim=2*self.ndim) m3 = D * m_child[3] m = [m0, m1, m2, m3] return m else: # Message to Gaussian-Wishart raise NotImplementedError() elif index == 1: if self.wishart: # Message to Wishart alpha_x = u_X_alpha[0] alpha_xx = u_X_alpha[1] alpha = u_X_alpha[2] m0 = (0.5*linalg.outer(alpha_x, m_child[0], ndim=self.ndim) + 0.5*linalg.outer(m_child[0], alpha_x, ndim=self.ndim) + alpha_xx * misc.add_trailing_axes(m_child[1], 2*self.ndim) + misc.add_trailing_axes(alpha, 2*self.ndim) * m_child[2]) m1 = m_child[3] m = [m0, m1] return m else: # Message to gamma (isotropic) raise NotImplementedError() else: raise ValueError("Invalid parent index") def reshape_gaussian_array(dims_from, dims_to, x0, x1): r""" Reshape the moments Gaussian array variable. The plates remain unaffected. """ num_dims_from = len(dims_from) num_dims_to = len(dims_to) # Reshape the first moment / mean num_plates_from = np.ndim(x0) - num_dims_from plates_from = np.shape(x0)[:num_plates_from] shape = ( plates_from + (1,)*(num_dims_to-num_dims_from) + dims_from ) x0 = np.ones(dims_to) * np.reshape(x0, shape) # Reshape the second moment / covariance / precision num_plates_from = np.ndim(x1) - 2*num_dims_from plates_from = np.shape(x1)[:num_plates_from] shape = ( plates_from + (1,)*(num_dims_to-num_dims_from) + dims_from + (1,)*(num_dims_to-num_dims_from) + dims_from ) x1 = np.ones(dims_to+dims_to) * np.reshape(x1, shape) return (x0, x1) def transpose_covariance(Cov, ndim=1): r""" Transpose the covariance array of Gaussian array variable. That is, swap the last ndim axes with the ndim axes before them. This makes transposing easy for array variables when the covariance is not a matrix but a multidimensional array. """ axes_in = [Ellipsis] + list(range(2*ndim,0,-1)) axes_out = [Ellipsis] + list(range(ndim,0,-1)) + list(range(2*ndim,ndim,-1)) return np.einsum(Cov, axes_in, axes_out) def left_rotate_covariance(Cov, R, axis=-1, ndim=1): r""" Rotate the covariance array of Gaussian array variable. ndim is the number of axes for the Gaussian variable. For vector variable, ndim=1 and covariance is a matrix. """ if not isinstance(axis, int): raise ValueError("Axis must be an integer") if axis < -ndim or axis >= ndim: raise ValueError("Axis out of range") # Force negative axis if axis >= 0: axis -= ndim # Rotation from left axes_R = [Ellipsis, ndim+abs(axis)+1, ndim+abs(axis)] axes_Cov = [Ellipsis] + list(range(ndim+abs(axis), 0, -1)) axes_out = [Ellipsis, ndim+abs(axis)+1] + list(range(ndim+abs(axis)-1, 0, -1)) Cov = np.einsum(R, axes_R, Cov, axes_Cov, axes_out) return Cov def right_rotate_covariance(Cov, R, axis=-1, ndim=1): r""" Rotate the covariance array of Gaussian array variable. ndim is the number of axes for the Gaussian variable. For vector variable, ndim=1 and covariance is a matrix. """ if not isinstance(axis, int): raise ValueError("Axis must be an integer") if axis < -ndim or axis >= ndim: raise ValueError("Axis out of range") # Force negative axis if axis >= 0: axis -= ndim # Rotation from right axes_R = [Ellipsis, abs(axis)+1, abs(axis)] axes_Cov = [Ellipsis] + list(range(abs(axis), 0, -1)) axes_out = [Ellipsis, abs(axis)+1] + list(range(abs(axis)-1, 0, -1)) Cov = np.einsum(R, axes_R, Cov, axes_Cov, axes_out) return Cov def rotate_covariance(Cov, R, axis=-1, ndim=1): r""" Rotate the covariance array of Gaussian array variable. ndim is the number of axes for the Gaussian variable. For vector variable, ndim=1 and covariance is a matrix. """ # Rotate from left and right Cov = left_rotate_covariance(Cov, R, ndim=ndim, axis=axis) Cov = right_rotate_covariance(Cov, R, ndim=ndim, axis=axis) return Cov def rotate_mean(mu, R, axis=-1, ndim=1): r""" Rotate the mean array of Gaussian array variable. ndim is the number of axes for the Gaussian variable. For vector variable, ndim=1 and mu is a vector. """ if not isinstance(axis, int): raise ValueError("Axis must be an integer") if axis < -ndim or axis >= ndim: raise ValueError("Axis out of range") # Force negative axis if axis >= 0: axis -= ndim # Rotation from right axes_R = [Ellipsis, abs(axis)+1, abs(axis)] axes_mu = [Ellipsis] + list(range(abs(axis), 0, -1)) axes_out = [Ellipsis, abs(axis)+1] + list(range(abs(axis)-1, 0, -1)) mu = np.einsum(R, axes_R, mu, axes_mu, axes_out) return mu def array_to_vector(x, ndim=1): if ndim == 0: return x shape_x = np.shape(x) D = np.prod(shape_x[-ndim:]) return np.reshape(x, shape_x[:-ndim] + (D,)) def array_to_matrix(A, ndim=1): if ndim == 0: return A shape_A = np.shape(A) D = np.prod(shape_A[-ndim:]) return np.reshape(A, shape_A[:-2*ndim] + (D,D)) def vector_to_array(x, shape): shape_x = np.shape(x) return np.reshape(x, np.shape(x)[:-1] + tuple(shape)) def matrix_dot_vector(A, x, ndim=1): if ndim < 0: raise ValueError("ndim must be non-negative integer") if ndim == 0: return A*x dims_x = np.shape(x)[-ndim:] A = array_to_matrix(A, ndim=ndim) x = array_to_vector(x, ndim=ndim) y = np.einsum('...ik,...k->...i', A, x) return vector_to_array(y, dims_x)