Fix for gp.sample_gp (redone)

pymc3/gp/gp.py

@@ -15,47 +15,47 @@
 
 class GP(Continuous):
     """Gausian process
-    
+
     Parameters
     ----------
     mean_func : Mean
         Mean function of Gaussian process
     cov_func : Covariance
         Covariance function of Gaussian process
     X : array
-        Grid of points to evaluate Gaussian process over. Only required if the 
+        Grid of points to evaluate Gaussian process over. Only required if the
         GP is not an observed variable.
     sigma : scalar or array
         Observation standard deviation (defaults to zero)
     """
     def __init__(self, mean_func=None, cov_func=None, X=None, sigma=0, *args, **kwargs):
-        
+
         if mean_func is None:
             self.M = Zero()
         else:
             if not isinstance(mean_func, Mean):
                 raise ValueError('mean_func must be a subclass of Mean')
             self.M = mean_func
-            
+
         if cov_func is None:
             raise ValueError('A covariance function must be specified for GPP')
         if not isinstance(cov_func, Covariance):
             raise ValueError('cov_func must be a subclass of Covariance')
         self.K = cov_func
-        
+
         self.sigma = sigma
-        
+
         if X is not None:
             self.X = X
             self.mean = self.mode = self.M(X)
             kwargs.setdefault("shape", X.squeeze().shape)
-            
+
         super(GP, self).__init__(*args, **kwargs)
-                
+
     def random(self, point=None, size=None, **kwargs):
         X = self.X
         mu, cov = draw_values([self.M(X).squeeze(), self.K(X) + np.eye(X.shape[0])*self.sigma**2], point=point)
-        
+
         def _random(mean, cov, size=None):
             return stats.multivariate_normal.rvs(
                 mean, cov, None if size == mean.shape else size)
@@ -74,9 +74,9 @@ def logp(self, Y, X=None):
         Sigma = self.K(X) + tt.eye(X.shape[0])*self.sigma**2
 
         return MvNormal.dist(mu, Sigma).logp(Y)
-
 
-def sample_gp(trace, gp, X_values, samples=None, obs_noise=True, model=None, random_seed=None, progressbar=True):
+
+def sample_gp(trace, gp, X_values, samples=None, obs_noise=True, model=None, random_seed=None, progressbar=True, chol_const=True):
     """Generate samples from a posterior Gaussian process.
 
     Parameters
@@ -92,38 +92,41 @@ def sample_gp(trace, gp, X_values, samples=None, obs_noise=True, model=None, ran
         length of `trace`
     obs_noise : bool
         Flag for including observation noise in sample. Defaults to True.
-    model : Model 
+    model : Model
         Model used to generate `trace`. Optional if in `with` context manager.
     random_seed : integer > 0
         Random number seed for sampling.
     progressbar : bool
         Flag for showing progress bar.
-
+    chol_const : bool
+        Flag to a small diagonal to the posterior covariance
+        for numerical stability
+
     Returns
     -------
     Array of samples from posterior GP evaluated at Z.
     """
     model = modelcontext(model)
-    
+
     if samples is None:
         samples = len(trace)
-    
+
     if random_seed:
         np.random.seed(random_seed)
-    
+
     if progressbar:
         indices = tqdm(np.random.randint(0, len(trace), samples), total=samples)
     else:
         indices = np.random.randint(0, len(trace), samples)
 
-    K = gp.distribution.K 
-        
+    K = gp.distribution.K
+
     data = [v for v in model.observed_RVs if v.name==gp.name][0].data
 
     X = data['X']
     Y = data['Y']
     Z = X_values
-    
+
     S_xz = K(X, Z)
     S_zz = K(Z)
     if obs_noise:
@@ -136,8 +139,10 @@ def sample_gp(trace, gp, X_values, samples=None, obs_noise=True, model=None, ran
     # Posterior covariance
     S_post = S_zz - tt.dot(tt.dot(S_xz.T, S_inv), S_xz)
 
-    gp_post = MvNormal.dist(m_post, S_post, shape=Z.shape[0])
-
+    if chol_const:
+        n = S_post.shape[0]
+        correction = 1e-6 * tt.nlinalg.trace(S_post) * tt.eye(n)
+
+    gp_post = MvNormal.dist(m_post, S_post + correction, shape=Z.shape[0])
     samples = [gp_post.random(point=trace[idx]) for idx in indices]
-
     return np.array(samples)

pymc3/tests/test_gp.py

@@ -7,7 +7,6 @@
 import numpy.testing as npt
 import pytest
 
-
 class TestZeroMean(object):
     def test_value(self):
         X = np.linspace(0, 1, 10)[:, None]
@@ -402,7 +401,7 @@ def test_func_args(self):
 
     def test_sample(self):
         X = np.linspace(0, 1, 10)[:, None]
-        Y = np.random.randn(10, 1)
+        Y = np.random.randn(10)
         with Model() as model:
             M = gp.mean.Zero()
             l = Uniform('l', 0, 5)
@@ -411,3 +410,9 @@ def test_sample(self):
             # make a Gaussian model
             random_test = gp.GP('random_test', mean_func=M, cov_func=K, sigma=sigma, observed={'X':X, 'Y':Y})
             tr = sample(20, init=None, progressbar=False, random_seed=self.random_seed)
+
+        # test prediction
+        Z = np.linspace(0, 1, 5)[:, None]
+        with model:
+            out = gp.sample_gp(tr[-3:], gp=random_test, X_values=Z, obs_noise=False,
+                               random_seed=self.random_seed, progressbar=False, chol_const=True)