Update pf.py

hansenlab · Oct 9, 2024 · 73fa25f · 73fa25f
1 parent 1f91c15
commit 73fa25f
Showing 1 changed file with 8 additions and 10 deletions.
diff --git a/mNSF/NSF/pf.py b/mNSF/NSF/pf.py
@@ -99,20 +99,18 @@ def __init__(self, J, L, Z, lik="poi", chol = True, X=None, psd_kernel=tfk.Mater
     self._init_misc()
     self.Kuu_chol = tf.Variable(self.eval_Kuu_chol(self.get_kernel()), dtype=dtp, trainable=False)
     if  not(chol):    
-      kernel = self.kernel 
       N = X.shape[0]
       mu_x = self.beta0+tfl.matmul(self.beta, X, transpose_b=True) #LxN
-      self.Kuf = kernel.matrix(self.Z, X) #LxMxN
-      self.alpha_x = tfl.cholesky_solve(self.Kuu_chol, self.Kuf) #LxMxN
+      Kuf = kernel.matrix(self.Z, X) #LxMxN
+      self.alpha_x = tfl.cholesky_solve(self.Kuu_chol, Kuf) #LxMxN 
       Kff_diag = kernel.apply(X, X, example_ndims=1)+self.nugget #LxN
-      mu_z = self.get_mu_z()
-      mu_tilde = mu_x + tfl.matvec(self.alpha_x, self.delta-mu_z, transpose_a=True) #LxN
+      #mu_tilde = mu_x + tfl.matvec(self.alpha_x, self.delta-mu_z, transpose_a=True) #LxN
       #compute the alpha(x_i)'(K_uu-Omega)alpha(x_i) term
-      self.a_t_Kchol = tfl.matmul(self.alpha_x, self.Kuu_chol, transpose_a=True) #LxNxM
-      aKa = tf.reduce_sum(tf.square(self.a_t_Kchol), axis=2) #LxN
-      self.a_t_Omega_tril = tfl.matmul(self.alpha_x, self.Omega_tril, transpose_a=True) #LxNxM
-      aOmega_a = tf.reduce_sum(tf.square(self.a_t_Omega_tril ), axis=2) #LxN
-      Sigma_tilde = Kff_diag - aKa + aOmega_a #LxN
+      a_t_Kchol = tfl.matmul(self.alpha_x, self.Kuu_chol, transpose_a=True) #LxNxM
+      aKa = tf.reduce_sum(tf.square(a_t_Kchol), axis=2) #LxN
+      a_t_Omega_tril = tfl.matmul(self.alpha_x, self.Omega_tril, transpose_a=True) #LxNxM
+      aOmega_a = tf.reduce_sum(tf.square(a_t_Omega_tril ), axis=2) #LxN
+      self.Sigma_tilde = Kff_diag - aKa + aOmega_a #LxN
     if self.lik=="gau" and not self.nonneg:
       self.feature_means = feature_means
     else: