implement the linear kernel with psi2 format

GPy/kern/_src/linear.py

@@ -103,197 +103,36 @@ class Linear(Kern):
     def gradients_X_diag(self, dL_dKdiag, X):
         return 2.*self.variances*dL_dKdiag[:,None]*X
 
+    def input_sensitivity(self):
+        return np.ones(self.input_dim) * self.variances
+
     #---------------------------------------#
     #             PSI statistics            #
     #---------------------------------------#
 
     def psi0(self, Z, variational_posterior):
-        if isinstance(variational_posterior, variational.SpikeAndSlabPosterior):
+        return self.psicomp.psicomputations(self.variances, Z, variational_posterior)[0]
-            return self.psicomp.psicomputations(self.variances, Z, variational_posterior)[0]
-        else:
-            return np.sum(self.variances * self._mu2S(variational_posterior), 1)
 
     def psi1(self, Z, variational_posterior):
-        if isinstance(variational_posterior, variational.SpikeAndSlabPosterior):
+        return self.psicomp.psicomputations(self.variances, Z, variational_posterior)[1]
-            return self.psicomp.psicomputations(self.variances, Z, variational_posterior)[1]
-        else:
-            return self.K(variational_posterior.mean, Z) #the variance, it does nothing
 
     @Cache_this(limit=1)
     def psi2(self, Z, variational_posterior):
-        if isinstance(variational_posterior, variational.SpikeAndSlabPosterior):
+        return self.psicomp.psicomputations(self.variances, Z, variational_posterior)[2]
-            return self.psicomp.psicomputations(self.variances, Z, variational_posterior)[2]
-        else:
-            ZA = Z * self.variances
-            ZAinner = self._ZAinner(variational_posterior, Z)
-            return np.dot(ZAinner, ZA.T)
 
     def update_gradients_expectations(self, dL_dpsi0, dL_dpsi1, dL_dpsi2, Z, variational_posterior):
-        if isinstance(variational_posterior, variational.SpikeAndSlabPosterior):
+        dL_dvar = self.psicomp.psiDerivativecomputations(dL_dpsi0, dL_dpsi1, dL_dpsi2, self.variances, Z, variational_posterior)[0]
-            dL_dvar,_,_,_,_ = self.psicomp.psiDerivativecomputations(dL_dpsi0, dL_dpsi1, dL_dpsi2, self.variances, Z, variational_posterior)
+        if self.ARD:
-            if self.ARD:
+            self.variances.gradient = dL_dvar
-                self.variances.gradient = dL_dvar
-            else:
-                self.variances.gradient = dL_dvar.sum()
         else:
-            #psi1
+            self.variances.gradient = dL_dvar.sum()
-            self.update_gradients_full(dL_dpsi1, variational_posterior.mean, Z)
-            # psi0:
-            tmp = dL_dpsi0[:, None] * self._mu2S(variational_posterior)
-            if self.ARD: self.variances.gradient += tmp.sum(0)
-            else: self.variances.gradient += tmp.sum()
-            #psi2
-            if self.ARD:
-                tmp = dL_dpsi2[:, :, :, None] * (self._ZAinner(variational_posterior, Z)[:, :, None, :] * Z[None, None, :, :])
-                self.variances.gradient += 2.*tmp.sum(0).sum(0).sum(0)
-            else:
-                self.variances.gradient += 2.*np.sum(dL_dpsi2 * self.psi2(Z, variational_posterior))/self.variances
 
     def gradients_Z_expectations(self, dL_dpsi0, dL_dpsi1, dL_dpsi2, Z, variational_posterior):
-        if isinstance(variational_posterior, variational.SpikeAndSlabPosterior):
+        return self.psicomp.psiDerivativecomputations(dL_dpsi0, dL_dpsi1, dL_dpsi2, self.variances, Z, variational_posterior)[1]
-            _,dL_dZ,_,_,_ = self.psicomp.psiDerivativecomputations(dL_dpsi0, dL_dpsi1, dL_dpsi2, self.variances, Z, variational_posterior)
-            return dL_dZ
-        else:
-            #psi1
-            grad = self.gradients_X(dL_dpsi1.T, Z, variational_posterior.mean)
-            #psi2
-            self._weave_dpsi2_dZ(dL_dpsi2, Z, variational_posterior, grad)
-            return grad
 
     def gradients_qX_expectations(self, dL_dpsi0, dL_dpsi1, dL_dpsi2, Z, variational_posterior):
-        if isinstance(variational_posterior, variational.SpikeAndSlabPosterior):
+        return self.psicomp.psiDerivativecomputations(dL_dpsi0, dL_dpsi1, dL_dpsi2, self.variances, Z, variational_posterior)[2:]
-            _,_,dL_dmu, dL_dS, dL_dgamma = self.psicomp.psiDerivativecomputations(dL_dpsi0, dL_dpsi1, dL_dpsi2, self.variances, Z, variational_posterior)
-            return dL_dmu, dL_dS, dL_dgamma
-        else:
-            grad_mu, grad_S = np.zeros(variational_posterior.mean.shape), np.zeros(variational_posterior.mean.shape)
-            # psi0
-            grad_mu += dL_dpsi0[:, None] * (2.0 * variational_posterior.mean * self.variances)
-            grad_S += dL_dpsi0[:, None] * self.variances
-            # psi1
-            grad_mu += (dL_dpsi1[:, :, None] * (Z * self.variances)).sum(1)
-            # psi2
-            self._weave_dpsi2_dmuS(dL_dpsi2, Z, variational_posterior, grad_mu, grad_S)
 
-            return grad_mu, grad_S
-
-    #--------------------------------------------------#
-    #            Helpers for psi statistics            #
-    #--------------------------------------------------#
-
-
-    def _weave_dpsi2_dmuS(self, dL_dpsi2, Z, vp, target_mu, target_S):
-        # Think N,num_inducing,num_inducing,input_dim
-        ZA = Z * self.variances
-        AZZA = ZA.T[:, None, :, None] * ZA[None, :, None, :]
-        AZZA = AZZA + AZZA.swapaxes(1, 2)
-        AZZA_2 = AZZA/2.
-        if config.getboolean('parallel', 'openmp'):
-            pragma_string = '#pragma omp parallel for private(m,mm,q,qq,factor,tmp)'
-            header_string = '#include <omp.h>'
-            weave_options = {'headers'           : ['<omp.h>'],
-                                'extra_compile_args': ['-fopenmp -O3'],
-                                'extra_link_args'   : ['-lgomp'],
-                                'libraries': ['gomp']}
-        else:
-            pragma_string = ''
-            header_string = ''
-            weave_options = {'extra_compile_args': ['-O3']}
-
-        #Using weave, we can exploit the symmetry of this problem:
-        code = """
-        int n, m, mm,q,qq;
-        double factor,tmp;
-        %s
-        for(n=0;n<N;n++){
-          for(m=0;m<num_inducing;m++){
-            for(mm=0;mm<=m;mm++){
-              //add in a factor of 2 for the off-diagonal terms (and then count them only once)
-              if(m==mm)
-                factor = dL_dpsi2(n,m,mm);
-              else
-                factor = 2.0*dL_dpsi2(n,m,mm);
-
-              for(q=0;q<input_dim;q++){
-
-                //take the dot product of mu[n,:] and AZZA[:,m,mm,q] TODO: blas!
-                tmp = 0.0;
-                for(qq=0;qq<input_dim;qq++){
-                  tmp += mu(n,qq)*AZZA(qq,m,mm,q);
-                }
-
-                target_mu(n,q) += factor*tmp;
-                target_S(n,q) += factor*AZZA_2(q,m,mm,q);
-              }
-            }
-          }
-        }
-        """ % pragma_string
-        support_code = """
-        %s
-        #include <math.h>
-        """ % header_string
-        mu = vp.mean
-        N,num_inducing,input_dim,mu = mu.shape[0],Z.shape[0],mu.shape[1],param_to_array(mu)
-        weave.inline(code, support_code=support_code,
-                     arg_names=['N','num_inducing','input_dim','mu','AZZA','AZZA_2','target_mu','target_S','dL_dpsi2'],
-                     type_converters=weave.converters.blitz,**weave_options)
-
-
-    def _weave_dpsi2_dZ(self, dL_dpsi2, Z, vp, target):
-        AZA = self.variances*self._ZAinner(vp, Z)
-
-        if config.getboolean('parallel', 'openmp'):
-            pragma_string = '#pragma omp parallel for private(n,mm,q)'
-            header_string = '#include <omp.h>'
-            weave_options =  {'headers'           : ['<omp.h>'],
-                              'extra_compile_args': ['-fopenmp -O3'],
-                              'extra_link_args'   : ['-lgomp'],
-                              'libraries': ['gomp']}
-        else:
-            pragma_string = ''
-            header_string = ''
-            weave_options = {'extra_compile_args': ['-O3']}
-
-        code="""
-        int n,m,mm,q;
-        %s
-        for(m=0;m<num_inducing;m++){
-          for(q=0;q<input_dim;q++){
-            for(mm=0;mm<num_inducing;mm++){
-              for(n=0;n<N;n++){
-                target(m,q) += 2*dL_dpsi2(n,m,mm)*AZA(n,mm,q);
-              }
-            }
-          }
-        }
-        """ % pragma_string
-        support_code = """
-        %s
-        #include <math.h>
-        """ % header_string
-
-        N,num_inducing,input_dim = vp.mean.shape[0],Z.shape[0],vp.mean.shape[1]
-        mu = param_to_array(vp.mean)
-        weave.inline(code, support_code=support_code,
-                     arg_names=['N','num_inducing','input_dim','AZA','target','dL_dpsi2'],
-                     type_converters=weave.converters.blitz,**weave_options)
-
-
-    @Cache_this(limit=1, ignore_args=(0,))
-    def _mu2S(self, vp):
-        return np.square(vp.mean) + vp.variance
-
-    @Cache_this(limit=1)
-    def _ZAinner(self, vp, Z):
-        ZA = Z*self.variances
-        inner = (vp.mean[:, None, :] * vp.mean[:, :, None])
-        diag_indices = np.diag_indices(vp.mean.shape[1], 2)
-        inner[:, diag_indices[0], diag_indices[1]] += vp.variance
-
-        return np.dot(ZA, inner).swapaxes(0, 1)  # NOTE: self.ZAinner \in [num_inducing x num_data x input_dim]!
-
-    def input_sensitivity(self):
-        return np.ones(self.input_dim) * self.variances
 
 class LinearFull(Kern):
     def __init__(self, input_dim, rank, W=None, kappa=None, active_dims=None, name='linear_full'):

GPy/kern/_src/psi_comp/__init__.py

@@ -7,6 +7,7 @@ from ....core.parameterization import variational
 import rbf_psi_comp
 import ssrbf_psi_comp
 import sslinear_psi_comp
+import linear_psi_comp
 
 class PSICOMP_RBF(Pickleable):
 
@@ -33,7 +34,7 @@ class PSICOMP_Linear(Pickleable):
     @Cache_this(limit=2, ignore_args=(0,))
     def psicomputations(self, variance, Z, variational_posterior):
         if isinstance(variational_posterior, variational.NormalPosterior):
-            raise NotImplementedError
+            return linear_psi_comp.psicomputations(variance, Z, variational_posterior)
         elif isinstance(variational_posterior, variational.SpikeAndSlabPosterior):
             return sslinear_psi_comp.psicomputations(variance, Z, variational_posterior)
         else:
@@ -42,7 +43,7 @@ class PSICOMP_Linear(Pickleable):
     @Cache_this(limit=2, ignore_args=(0,1,2,3))
     def psiDerivativecomputations(self, dL_dpsi0, dL_dpsi1, dL_dpsi2, variance, Z, variational_posterior):
         if isinstance(variational_posterior, variational.NormalPosterior):
-            raise NotImplementedError
+            return linear_psi_comp.psiDerivativecomputations(dL_dpsi0, dL_dpsi1, dL_dpsi2, variance, Z, variational_posterior)
         elif isinstance(variational_posterior, variational.SpikeAndSlabPosterior):
             return sslinear_psi_comp.psiDerivativecomputations(dL_dpsi0, dL_dpsi1, dL_dpsi2, variance, Z, variational_posterior)
         else:

GPy/kern/_src/psi_comp/linear_psi_comp.py

@@ -0,0 +1,74 @@
+# Copyright (c) 2012, GPy authors (see AUTHORS.txt).
+# Licensed under the BSD 3-clause license (see LICENSE.txt)
+
+"""
+The package for the Psi statistics computation of the linear kernel for Bayesian GPLVM
+"""
+
+import numpy as np
+
+def psicomputations(variance, Z, variational_posterior):
+    """
+    Compute psi-statistics for ss-linear kernel
+    """
+    # here are the "statistics" for psi0, psi1 and psi2
+    # Produced intermediate results:
+    # psi0    N
+    # psi1    NxM
+    # psi2    MxM
+    mu = variational_posterior.mean
+    S = variational_posterior.variance
+
+    psi0 = np.einsum('q,nq->n',variance,np.square(mu)+S)
+    psi1 = np.einsum('q,mq,nq->nm',variance,Z,mu)
+
+    tmp = np.einsum('q,mq,nq->nm',variance,Z,mu)
+    psi2 = np.einsum('q,mq,oq,nq->mo',np.square(variance),Z,Z,S) + np.einsum('nm,no->mo',tmp,tmp)
+
+    return psi0, psi1, psi2
+
+def psiDerivativecomputations(dL_dpsi0, dL_dpsi1, dL_dpsi2, variance, Z, variational_posterior):
+    mu = variational_posterior.mean
+    S = variational_posterior.variance
+
+    dL_dvar, dL_dmu, dL_dS, dL_dZ = _psi2computations(dL_dpsi2, variance, Z, mu, S)
+
+    # Compute for psi0 and psi1
+    mu2S = np.square(mu)+S
+    dL_dvar += np.einsum('n,nq->q',dL_dpsi0,mu2S) + np.einsum('nm,mq,nq->q',dL_dpsi1,Z,mu)
+    dL_dmu += np.einsum('n,q,nq->nq',dL_dpsi0,2.*variance,mu) + np.einsum('nm,q,mq->nq',dL_dpsi1,variance,Z)
+    dL_dS += np.einsum('n,q->nq',dL_dpsi0,variance)
+    dL_dZ +=  np.einsum('nm,q,nq->mq',dL_dpsi1, variance,mu)
+    
+    return dL_dvar, dL_dZ, dL_dmu, dL_dS
+
+def _psi2computations(dL_dpsi2, variance, Z, mu, S):
+    """
+    Z - MxQ
+    mu - NxQ
+    S - NxQ
+    gamma - NxQ
+    """
+    # here are the "statistics" for psi1 and psi2
+    # Produced intermediate results:
+    # _psi2_dvariance      Q
+    # _psi2_dZ             MxQ
+    # _psi2_dmu            NxQ
+    # _psi2_dS             NxQ
+    
+    variance2 = np.square(variance)
+    common_sum = np.einsum('q,mq,nq->nm',variance,Z,mu) # NxM
+
+    dL_dvar = np.einsum('mo,nq,q,mq,oq->q',dL_dpsi2,2.*S,variance,Z,Z)+\
+        np.einsum('mo,mq,nq,no->q',dL_dpsi2,Z,mu,common_sum)+\
+        np.einsum('mo,oq,nq,nm->q',dL_dpsi2,Z,mu,common_sum)
+            
+    dL_dmu = np.einsum('mo,q,mq,no->nq',dL_dpsi2,variance,Z,common_sum)+\
+        np.einsum('mo,q,oq,nm->nq',dL_dpsi2,variance,Z,common_sum)
+    
+    dL_dS = np.empty(S.shape)
+    dL_dS[:] = np.einsum('mo,q,mq,oq->q',dL_dpsi2,variance2,Z,Z)
+    
+    dL_dZ = 2.*(np.einsum('om,q,mq,nq->oq',dL_dpsi2,variance2,Z,S)+np.einsum('om,q,nq,nm->oq',dL_dpsi2,variance,mu,common_sum))
+
+    return dL_dvar, dL_dmu, dL_dS, dL_dZ