From bc59cb8b225597df9e2d23294498e92e9768dbaf Mon Sep 17 00:00:00 2001
From: Zhenwen Dai <z.dai@shef.ac.uk>
Date: Wed, 26 Mar 2014 17:09:01 +0000
Subject: [PATCH] [GPU] psi1 after debug

---
 .../latent_function_inference/var_dtc_gpu.py  |  28 ++--
 GPy/kern/_src/psi_comp/ssrbf_psi_gpucomp.py   | 152 +++++++++++++++---
 GPy/kern/_src/rbf.py                          |   2 +-
 GPy/models/ss_gplvm.py                        |  11 ++
 GPy/util/linalg_gpu.py                        |   4 +-
 5 files changed, 157 insertions(+), 40 deletions(-)

diff --git a/GPy/inference/latent_function_inference/var_dtc_gpu.py b/GPy/inference/latent_function_inference/var_dtc_gpu.py
index ba7ec602..75a07992 100644
--- a/GPy/inference/latent_function_inference/var_dtc_gpu.py
+++ b/GPy/inference/latent_function_inference/var_dtc_gpu.py
@@ -52,17 +52,17 @@ class VarDTC_GPU(object):
     def _initGPUCache(self, num_inducing, output_dim):
         if self.gpuCache == None:
             self.gpuCache = {# inference_likelihood
-                             'Kmm_gpu'              :gpuarray.empty((num_inducing,num_inducing),np.float64),
-                             'Lm_gpu'               :gpuarray.empty((num_inducing,num_inducing),np.float64),
-                             'ones_gpu'             :gpuarray.empty(num_inducing, np.float64),
-                             'LL_gpu'               :gpuarray.empty((num_inducing,num_inducing),np.float64),
-                             'b_gpu'                :gpuarray.empty((num_inducing,output_dim),np.float64),
-                             'v_gpu'                :gpuarray.empty((num_inducing,output_dim),np.float64),
-                             'vvt_gpu'              :gpuarray.empty((num_inducing,num_inducing),np.float64),
-                             'KmmInvPsi2LLInvT_gpu' :gpuarray.empty((num_inducing,num_inducing),np.float64),
-                             'KmmInvPsi2P_gpu'      :gpuarray.empty((num_inducing,num_inducing),np.float64),
-                             'dL_dpsi2R_gpu'        :gpuarray.empty((num_inducing,num_inducing),np.float64),
-                             'dL_dKmm_gpu'          :gpuarray.empty((num_inducing,num_inducing),np.float64),
+                             'Kmm_gpu'              :gpuarray.empty((num_inducing,num_inducing),np.float64,order='F'),
+                             'Lm_gpu'               :gpuarray.empty((num_inducing,num_inducing),np.float64,order='F'),
+                             'ones_gpu'             :gpuarray.empty(num_inducing, np.float64,order='F'),
+                             'LL_gpu'               :gpuarray.empty((num_inducing,num_inducing),np.float64,order='F'),
+                             'b_gpu'                :gpuarray.empty((num_inducing,output_dim),np.float64,order='F'),
+                             'v_gpu'                :gpuarray.empty((num_inducing,output_dim),np.float64,order='F'),
+                             'vvt_gpu'              :gpuarray.empty((num_inducing,num_inducing),np.float64,order='F'),
+                             'KmmInvPsi2LLInvT_gpu' :gpuarray.empty((num_inducing,num_inducing),np.float64,order='F'),
+                             'KmmInvPsi2P_gpu'      :gpuarray.empty((num_inducing,num_inducing),np.float64,order='F'),
+                             'dL_dpsi2R_gpu'        :gpuarray.empty((num_inducing,num_inducing),np.float64,order='F'),
+                             'dL_dKmm_gpu'          :gpuarray.empty((num_inducing,num_inducing),np.float64,order='F'),
                              # inference_minibatch
                              }
             self.gpuCache['ones_gpu'].fill(1.0)
@@ -134,11 +134,11 @@ class VarDTC_GPU(object):
             if het_noise:
                 beta_slice = beta[n_start:n_end]
                 psi0_full += (beta_slice*psi0).sum()
-                psi1Y_full += np.dot(psi1,beta_slice[:,None]*Y_slice) # DxM
+                psi1Y_full += np.dot(psi1.T,beta_slice[:,None]*Y_slice) # MxD
                 YRY_full += (beta_slice*np.square(Y_slice).sum(axis=-1)).sum()
             else:
                 psi0_full += psi0.sum()
-                psi1Y_full += np.dot(psi1,Y_slice) # DxM
+                psi1Y_full += np.dot(psi1.T,Y_slice) # MxD
                 
                 
             if uncertain_inputs:
@@ -275,7 +275,7 @@ class VarDTC_GPU(object):
         # Compute the Posterior distribution of inducing points p(u|Y)
         #======================================================================
                 
-        post = Posterior(woodbury_inv=KmmInvPsi2P_gpu.get(), woodbury_vector=v_gpu.get(), K=Kmm_gpu.get(), mean=None, cov=None, K_chol=Lm.get())
+        post = Posterior(woodbury_inv=KmmInvPsi2P_gpu.get(), woodbury_vector=v_gpu.get(), K=Kmm_gpu.get(), mean=None, cov=None, K_chol=Lm_gpu.get())
 
         return logL, dL_dKmm, post
 
diff --git a/GPy/kern/_src/psi_comp/ssrbf_psi_gpucomp.py b/GPy/kern/_src/psi_comp/ssrbf_psi_gpucomp.py
index 2acddae9..467b779d 100644
--- a/GPy/kern/_src/psi_comp/ssrbf_psi_gpucomp.py
+++ b/GPy/kern/_src/psi_comp/ssrbf_psi_gpucomp.py
@@ -13,32 +13,118 @@ try:
     import pycuda.gpuarray as gpuarray
     from scikits.cuda import cublas
     import pycuda.autoinit
-    from pycuda.reduction import ReductionKernel
-    from ...util.linalg_gpu import logDiagSum
-    
+    from pycuda.reduction import ReductionKernel    
     from pycuda.elementwise import ElementwiseKernel
     
     # The kernel form computing psi1
     comp_psi1 = ElementwiseKernel(
-        "double *psi1, double var, double *l, double *Z, double *mu, double *S, double *logGamma, double *log1Gamma, double *logpsi1denom, int N, int M, int Q",
-        "psi1[i] = comp_psi1_element(var,l, Z, mu, S, logGamma, log1Gamma, psi1denom, N, M, Q, i)",
+        "double *psi1, double var, double l, double *Z, double *mu, double *S, double *logGamma, double *log1Gamma, double *logpsi1denom, int N, int M, int Q",
+        "psi1[i] = comp_psi1_element(var,l, Z, mu, S, logGamma, log1Gamma, logpsi1denom, N, M, Q, i)",
         "comp_psi1",
         preamble="""
-        #define IDX_MQ(n,m,q) ((n*M+m)*Q+q)
-        #define IDX_Q(n,q) (n*Q+q)
+        #define IDX_NMQ(n,m,q) ((q*M+m)*N+n)
+        #define IDX_NQ(n,q) (q*N+n)
+        #define IDX_MQ(m,q) (q*M+m)
+        #define LOGEXPSUM(a,b) (a>=b?a+log(1.0+exp(b-a)):b+log(1.0+exp(a-b)))
+        
+        __device__ double comp_psi1_element(double var, double l, double *Z, double *mu, double *S, double *logGamma, double *log1Gamma, double *logpsi1denom, int N, int M, int Q, int idx)
+        {
+            int n = idx%N;
+            int m = idx/N;
+            double psi1_exp=0;
+            for(int q=0;q<Q;q++){
+                double muZ = mu[IDX_NQ(n,q)]-Z[IDX_MQ(m,q)];
+                double exp1 = logGamma[IDX_NQ(n,q)] - (logpsi1denom[IDX_NQ(n,q)] + muZ*muZ/(S[IDX_NQ(n,q)]+l) )/2.0;
+                double exp2 = log1Gamma[IDX_NQ(n,q)] - Z[IDX_MQ(m,q)]*Z[IDX_MQ(m,q)]/(l*2.0);
+                psi1_exp += LOGEXPSUM(exp1,exp2);
+            }
+            return var*exp(psi1_exp);
+        }
+        """)
+    
+    # The kernel form computing psi1 het_noise
+    comp_psi1_het = ElementwiseKernel(
+        "double *psi1, double var, double *l, double *Z, double *mu, double *S, double *logGamma, double *log1Gamma, double *logpsi1denom, int N, int M, int Q",
+        "psi1[i] = comp_psi1_element(var,l, Z, mu, S, logGamma, log1Gamma, logpsi1denom, N, M, Q, i)",
+        "comp_psi1",
+        preamble="""
+        #define IDX_NMQ(n,m,q) ((q*M+m)*N+n)
+        #define IDX_NQ(n,q) (q*N+n)
+        #define IDX_MQ(m,q) (q*M+m)
+        #define LOGEXPSUM(a,b) (a>=b?a+log(1.0+exp(b-a)):b+log(1.0+exp(a-b)))
         
         __device__ double comp_psi1_element(double var, double *l, double *Z, double *mu, double *S, double *logGamma, double *log1Gamma, double *logpsi1denom, int N, int M, int Q, int idx)
         {
-            int n = idx/M;
-            int m = idx%M;
-            double psi1=0;
+            int n = idx%N;
+            int m = idx/N;
+            double psi1_exp=0;
             for(int q=0;q<Q;q++){
-                double muZ = mu[IDX_Q(n,q)]-Z[IDX_Q(m,q)];
-                double exp1 = logGamma[IDX_Q(n,q)] - (logpsi1denom[IDX_Q(n,q)] + muZ*muZ/(S[IDX_Q(n,q)]+l[q]) )/2.0;
-                double exp2 = log1Gamma[IDX_Q(n,q)] - (Z[IDX_Q(m,q)]*Z[IDX_Q(m,q)]/l[q])/2.0;
-                psi1 += exp1>=exp2?exp1+log(1.0+exp(exp2-exp1)):exp2+log(1.0+exp(exp1-exp2));
+                double muZ = mu[IDX_NQ(n,q)]-Z[IDX_MQ(m,q)];
+                double exp1 = logGamma[IDX_NQ(n,q)] - (logpsi1denom[IDX_NQ(n,q)] + muZ*muZ/(S[IDX_NQ(n,q)]+l[q]) )/2.0;
+                double exp2 = log1Gamma[IDX_NQ(n,q)] - Z[IDX_MQ(m,q)]*Z[IDX_MQ(m,q)]/(l[q]*2.0);
+                psi1_exp += LOGEXPSUM(exp1,exp2);
             }
-            return var*exp(psi1);
+            return var*exp(psi1_exp);
+        }
+        """)
+    
+    # The kernel form computing psi2 het_noise
+    comp_psi2_het = ElementwiseKernel(
+        "double *psi2, double var, double *l, double *Z, double *mu, double *S, double *logGamma, double *log1Gamma, double *logpsi2denom, int N, int M, int Q",
+        "psi2[i] = comp_psi1_element(var,l, Z, mu, S, logGamma, log1Gamma, logpsi2denom, N, M, Q, i)",
+        "comp_psi2",
+        preamble="""
+        #define IDX_NMQ(n,m,q) ((q*M+m)*N+n)
+        #define IDX_NQ(n,q) (q*N+n)
+        #define IDX_MQ(m,q) (q*M+m)
+        #define LOGEXPSUM(a,b) (a>=b?a+log(1.0+exp(b-a)):b+log(1.0+exp(a-b)))
+        
+        __device__ double comp_psi1_element(double var, double *l, double *Z, double *mu, double *S, double *logGamma, double *log1Gamma, double *logpsi2denom, int N, int M, int Q, int idx)
+        {
+            // psi2 (n,m1,m2)
+            int m2 = idx/(M*N);
+            int m1 = (idx%(M*N))/N;
+            int n = idx%N;
+
+            double psi2_exp=0;
+            for(int q=0;q<Q;q++){ 
+                double dZ = Z[IDX_MQ(m1,q)]-Z[IDX_MQ(m2,q)];
+                double muZ = mu[IDX_NQ(n,q)] - (Z[IDX_MQ(m1,q)]+Z[IDX_MQ(m2,q)])/2.0;
+                double exp1 = logGamma[IDX_NQ(n,q)] - (logpsi2denom[IDX_NQ(n,q)])/2.0 - dZ*dZ/(l[q]*4.0) - muZ*muZ/(2*mu[IDX_NQ(n,q)]+l[q]);
+                double exp2 = log1Gamma[IDX_NQ(n,q)] - (Z[IDX_MQ(m1,q)]*Z[IDX_MQ(m1,q)]+Z[IDX_MQ(m2,q)]*Z[IDX_MQ(m2,q)])/(l[q]*2.0);
+                psi2_exp += LOGEXPSUM(exp1,exp2);
+            }
+            return var*var*exp(psi2_exp);
+        }
+        """)
+    
+    # The kernel form computing psi2 
+    comp_psi2 = ElementwiseKernel(
+        "double *psi2, double var, double l, double *Z, double *mu, double *S, double *logGamma, double *log1Gamma, double *logpsi2denom, int N, int M, int Q",
+        "psi2[i] = comp_psi1_element(var,l, Z, mu, S, logGamma, log1Gamma, logpsi2denom, N, M, Q, i)",
+        "comp_psi2",
+        preamble="""
+        #define IDX_NMQ(n,m,q) ((q*M+m)*N+n)
+        #define IDX_NQ(n,q) (q*N+n)
+        #define IDX_MQ(m,q) (q*M+m)
+        #define LOGEXPSUM(a,b) (a>=b?a+log(1.0+exp(b-a)):b+log(1.0+exp(a-b)))
+        
+        __device__ double comp_psi1_element(double var, double l, double *Z, double *mu, double *S, double *logGamma, double *log1Gamma, double *logpsi2denom, int N, int M, int Q, int idx)
+        {
+            // psi2 (n,m1,m2)
+            int m2 = idx/(M*N);
+            int m1 = (idx%(M*N))/N;
+            int n = idx%N;
+
+            double psi2_exp=0;
+            for(int q=0;q<Q;q++){ 
+                double dZ = Z[IDX_MQ(m1,q)]-Z[IDX_MQ(m2,q)];
+                double muZ = mu[IDX_NQ(n,q)] - (Z[IDX_MQ(m1,q)]+Z[IDX_MQ(m2,q)])/2.0;
+                double exp1 = logGamma[IDX_NQ(n,q)] - (logpsi2denom[IDX_NQ(n,q)])/2.0 - dZ*dZ/(l*4.0) - muZ*muZ/(2*mu[IDX_NQ(n,q)]+l);
+                double exp2 = log1Gamma[IDX_NQ(n,q)] - (Z[IDX_MQ(m1,q)]*Z[IDX_MQ(m1,q)]+Z[IDX_MQ(m2,q)]*Z[IDX_MQ(m2,q)])/(l*2.0);
+                psi2_exp += LOGEXPSUM(exp1,exp2);
+            }
+            return var*var*exp(psi2_exp);
         }
         """)
 except:
@@ -105,19 +191,19 @@ def _psi1computations(variance, lengthscale, Z, mu, S, gamma):
     M = Z.shape[0]
     Q = mu.shape[1]
 
-    l_gpu = gpuarray.to_gpu(lengthscale2)
-    Z_gpu = gpuarray.to_gpu(Z)
-    mu_gpu = gpuarray.to_gpu(mu)
-    S_gpu = gpuarray.to_gpu(S)
+    l_gpu = gpuarray.to_gpu(np.asfortranarray(lengthscale2))
+    Z_gpu = gpuarray.to_gpu(np.asfortranarray(Z))
+    mu_gpu = gpuarray.to_gpu(np.asfortranarray(mu))
+    S_gpu = gpuarray.to_gpu(np.asfortranarray(S))
     #gamma_gpu = gpuarray.to_gpu(gamma)
-    logGamma_gpu = gpuarray.to_gpu(np.log(gamma))
-    log1Gamma_gpu = gpuarray.to_gpu(np.log(1.-gamma))
-    logpsi1denom_gpu = gpuarray.to_gpu(np.log(S/lengthscale2+1.))
+    logGamma_gpu = gpuarray.to_gpu(np.asfortranarray(np.log(gamma)))
+    log1Gamma_gpu = gpuarray.to_gpu(np.asfortranarray(np.log(1.-gamma)))
+    logpsi1denom_gpu = gpuarray.to_gpu(np.asfortranarray(np.log(S/lengthscale2+1.)))
     psi1_gpu = gpuarray.empty((mu.shape[0],Z.shape[0]),np.float64)
     
-    comp_psi1(psi1_gpu, variance, l_gpu, Z_gpu, mu_gpu, S_gpu, logGamma_gpu, log1Gamma_gpu, logpsi1denom_gpu, N, M, Q)
+    comp_psi1(psi1_gpu, variance, lengthscale2, Z_gpu, mu_gpu, S_gpu, logGamma_gpu, log1Gamma_gpu, logpsi1denom_gpu, N, M, Q)
     
-    print np.abs(psi1_gpu.get()-_psi1).max()
+    #print np.abs(psi1_gpu.get()-_psi1).max()
 
     return _psi1, _dpsi1_dvariance, _dpsi1_dgamma, _dpsi1_dmu, _dpsi1_dS, _dpsi1_dZ, _dpsi1_dlengthscale
 
@@ -167,4 +253,22 @@ def _psi2computations(variance, lengthscale, Z, mu, S, gamma):
     _dpsi2_dZ = 2.*_psi2_q * (_psi2_common*(-_psi2_Zdist*_psi2_denom+_psi2_mudist)*_psi2_exp_dist_sq - (1-gamma[:,None,None,:])*Z[:,None,:]/lengthscale2*_psi2_exp_Z) # NxMxMxQ
     _dpsi2_dlengthscale = 2.*lengthscale* _psi2_q * (_psi2_common*(S[:,None,None,:]/lengthscale2+_psi2_Zdist_sq*_psi2_denom+_psi2_mudist_sq)*_psi2_exp_dist_sq+(1-gamma[:,None,None,:])*_psi2_Z_sq_sum*0.5/lengthscale2*_psi2_exp_Z) # NxMxMxQ
 
+    N = mu.shape[0]
+    M = Z.shape[0]
+    Q = mu.shape[1]
+
+#    l_gpu = gpuarray.to_gpu(np.asfortranarray(lengthscale2))
+    Z_gpu = gpuarray.to_gpu(np.asfortranarray(Z))
+    mu_gpu = gpuarray.to_gpu(np.asfortranarray(mu))
+    S_gpu = gpuarray.to_gpu(np.asfortranarray(S))
+    #gamma_gpu = gpuarray.to_gpu(gamma)
+    logGamma_gpu = gpuarray.to_gpu(np.asfortranarray(np.log(gamma)))
+    log1Gamma_gpu = gpuarray.to_gpu(np.asfortranarray(np.log(1.-gamma)))
+    logpsi2denom_gpu = gpuarray.to_gpu(np.asfortranarray(np.log(2.*S/lengthscale2+1.)))
+    psi2_gpu = gpuarray.empty((mu.shape[0],Z.shape[0],Z.shape[0]),np.float64)
+    
+    comp_psi2(psi2_gpu, variance, lengthscale2, Z_gpu, mu_gpu, S_gpu, logGamma_gpu, log1Gamma_gpu, logpsi2denom_gpu, N, M, Q)
+    
+    print np.abs(psi2_gpu.get()-_psi2).max()
+
     return _psi2, _dpsi2_dvariance, _dpsi2_dgamma, _dpsi2_dmu, _dpsi2_dS, _dpsi2_dZ, _dpsi2_dlengthscale
diff --git a/GPy/kern/_src/rbf.py b/GPy/kern/_src/rbf.py
index 32d5e1a5..0cf8b8de 100644
--- a/GPy/kern/_src/rbf.py
+++ b/GPy/kern/_src/rbf.py
@@ -8,7 +8,7 @@ from ...util.misc import param_to_array
 from stationary import Stationary
 from GPy.util.caching import Cache_this
 from ...core.parameterization import variational
-from psi_comp import ssrbf_psi_gpucomp
+from psi_comp import ssrbf_psi_gpucomp as ssrbf_psi_comp
 
 class RBF(Stationary):
     """
diff --git a/GPy/models/ss_gplvm.py b/GPy/models/ss_gplvm.py
index e32745c7..ec0f032a 100644
--- a/GPy/models/ss_gplvm.py
+++ b/GPy/models/ss_gplvm.py
@@ -11,6 +11,9 @@ from ..likelihoods import Gaussian
 from ..inference.optimization import SCG
 from ..util import linalg
 from ..core.parameterization.variational import SpikeAndSlabPrior, SpikeAndSlabPosterior
+from ..inference.latent_function_inference.var_dtc_parallel import update_gradients
+from ..inference.latent_function_inference.var_dtc_gpu import VarDTC_GPU
+
 
 class SSGPLVM(SparseGP):
     """
@@ -64,8 +67,16 @@ class SSGPLVM(SparseGP):
         SparseGP.__init__(self, X, Y, Z, kernel, likelihood, inference_method, name, **kwargs)
         self.add_parameter(self.X, index=0)
         self.add_parameter(self.variational_prior)
+        
+    def set_X_gradients(self, X, X_grad):
+        """Set the gradients of the posterior distribution of X in its specific form."""
+        X.mean.gradient, X.variance.gradient, X.binary_prob.gradient = X_grad
 
     def parameters_changed(self):
+        if isinstance(self.inference_method, VarDTC_GPU):
+            update_gradients(self)
+            return
+        
         super(SSGPLVM, self).parameters_changed()
         self._log_marginal_likelihood -= self.variational_prior.KL_divergence(self.X)
 
diff --git a/GPy/util/linalg_gpu.py b/GPy/util/linalg_gpu.py
index e066bc04..12d5a823 100644
--- a/GPy/util/linalg_gpu.py
+++ b/GPy/util/linalg_gpu.py
@@ -8,6 +8,8 @@
 import numpy as np
 
 try:
+    import pycuda.autoinit
     from pycuda.reduction import ReductionKernel
     logDiagSum = ReductionKernel(np.float64, neutral="0", reduce_expr="a+b", map_expr="i%step==0?log(x[i]):0", arguments="double *x, int step")
-except:
\ No newline at end of file
+except:
+    pass