From 38700189af255a8770ff40c8a7e43482ae17db4d Mon Sep 17 00:00:00 2001
From: Boris Leistedt <boris.leistedt@gmail.com>
Date: Mon, 23 Jan 2017 17:35:08 +0000
Subject: [PATCH] Added support for full rank yerr, not just diagonal.

---
 george/basic.py           |  5 ++++-
 george/gp.py              | 15 ++++++++++++---
 george/testing/test_gp.py | 39 ++++++++++++++++++++++++++++++++++++++-
 3 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/george/basic.py b/george/basic.py
index 895726af..74948a87 100644
--- a/george/basic.py
+++ b/george/basic.py
@@ -63,7 +63,10 @@ def compute(self, x, yerr):
         """
         # Compute the kernel matrix.
         K = self.kernel.value(x)
-        K[np.diag_indices_from(K)] += yerr ** 2
+        if x.size == yerr.size:  # yerr is diagonal only
+            K[np.diag_indices_from(K)] += yerr ** 2
+        else:  # yerr is full square matrix
+            K += yerr ** 2
 
         # Factor the matrix and compute the log-determinant.
         self._factor = (cholesky(K, overwrite_a=True, lower=False), False)
diff --git a/george/gp.py b/george/gp.py
index 215e0a10..9e4e9727 100644
--- a/george/gp.py
+++ b/george/gp.py
@@ -158,7 +158,8 @@ def compute(self, x, yerr=TINY, sort=True, **kwargs):
         :param x: ``(nsamples,)`` or ``(nsamples, ndim)``
             The independent coordinates of the data points.
 
-        :param yerr: (optional) ``(nsamples,)`` or scalar
+        :param yerr: (optional)
+            ``(nsamples,)``, ``(nsamples, nsamples)`` or scalar
             The Gaussian uncertainties on the data points at coordinates
             ``x``. These values will be added in quadrature to the diagonal of
             the covariance matrix.
@@ -175,7 +176,14 @@ def compute(self, x, yerr=TINY, sort=True, **kwargs):
         self._x, self.inds = self.parse_samples(x, sort)
         self._x = np.ascontiguousarray(self._x, dtype=np.float64)
         try:
-            self._yerr = float(yerr) * np.ones(len(x))
+            if isinstance(yerr, np.ndarray):
+                if yerr.ndim == 1:
+                    assert yerr.size == len(x)
+                else:
+                    assert yerr.shape[0] == len(x) and yerr.shape[1] == len(x)
+                self._yerr = 1*yerr
+            else:
+                self._yerr = float(yerr) * np.ones(len(x))
         except TypeError:
             self._yerr = self._check_dimensions(yerr)[self.inds]
         self._yerr = np.ascontiguousarray(self._yerr, dtype=np.float64)
@@ -377,7 +385,8 @@ def optimize(self, x, y, yerr=TINY, sort=True, dims=None, verbose=True,
         :param y: ``(nsamples, )``
             The observations at the coordinates ``x``.
 
-        :param yerr: (optional) ``(nsamples,)`` or scalar
+        :param yerr: (optional)
+            ``(nsamples,)``, ``(nsamples, nsamples)`` or scalar
             The Gaussian uncertainties on the data points at coordinates
             ``x``. These values will be added in quadrature to the diagonal of
             the covariance matrix.
diff --git a/george/testing/test_gp.py b/george/testing/test_gp.py
index a430b920..398b9a58 100644
--- a/george/testing/test_gp.py
+++ b/george/testing/test_gp.py
@@ -3,7 +3,8 @@
 from __future__ import division, print_function
 
 __all__ = [
-    "test_gradient", "test_prediction", "test_repeated_prediction_cache"
+    "test_gradient", "test_prediction", "test_repeated_prediction_cache",
+    'test_prediction_fullrankcovmatrix'
 ]
 
 import numpy as np
@@ -111,3 +112,39 @@ def test_repeated_prediction_cache():
     assert not np.allclose(a0, a1), \
         "Different kernel parameters must give different alphas " \
         "(problem with GP cache)."
+
+
+def test_prediction_fullrankcovmatrix(solver=BasicSolver):
+    """Basic sanity checks for GP regression,
+        this time with a full rank yerr covariance matrix"""
+
+    kernel = kernels.ExpSquaredKernel(1.0)
+    gp = GP(kernel, solver=solver)
+
+    nobj = 4
+    x = np.random.randn(nobj)
+    x.sort()
+    y = np.random.randn(nobj)
+    x_pred = np.random.randn(nobj)
+    chol = np.random.uniform(size=nobj*nobj).reshape((nobj, nobj))
+    yerr_full = np.dot(chol, chol.transpose())
+    yerr_diag = np.diag(yerr_full)
+    yerr_diag_full = np.diag(yerr_diag)
+
+    gp.compute(x, yerr=yerr_diag)
+    mu_diag, cov_diag = gp.predict(y, x_pred)
+
+    gp.compute(x, yerr=yerr_diag_full)
+    mu_diag_full, cov_diag_full = gp.predict(y, x_pred)
+
+    assert np.allclose(mu_diag, mu_diag_full), \
+        "GPs with diagonal error matrix and only diagonal provided should agree"
+    assert np.allclose(cov_diag, cov_diag_full), \
+        "GPs with diagonal error matrix and only diagonal provided should agree"
+
+    gp.compute(x, yerr=yerr_full)
+    mu_full, cov_full = gp.predict(y, x_pred)
+
+    assert np.all(cov_full > -1e-15), \
+        "Covariance matrix must be nonnegative ({0}).\n{1}" \
+        .format(solver.__name__, cov_full)