Sparse GP doesn't accept covariance function as noise

JohnGoertz commented 3 years ago

Description of your problem

When defining a GP with pm.gp.Marginal, either a covariance function or an RV can be passed as the noise argument to pm.marginal_likelihood. However, a pm.gp.MarginalSparse GP throws an error if a covariance function is passed as the noise argument.

Building off the Sparse Approximations example

Setup

```python import matplotlib.pyplot as plt import numpy as np import pymc3 as pm # set the seed np.random.seed(1) n = 2000 # The number of data points X = 10 * np.sort(np.random.rand(n))[:, None] # Define the true covariance function and its parameters ℓ_true = 1.0 η_true = 3.0 cov_func = η_true ** 2 * pm.gp.cov.Matern52(1, ℓ_true) # A mean function that is zero everywhere mean_func = pm.gp.mean.Zero() # The latent function values are one sample from a multivariate normal # Note that we have to call `eval()` because PyMC3 built on top of Theano f_true = np.random.multivariate_normal( mean_func(X).eval(), cov_func(X).eval() + 1e-8 * np.eye(n), 1 ).flatten() # The observed data is the latent function plus a small amount of IID Gaussian noise # The standard deviation of the noise is `sigma` σ_true = 2.0 y = f_true + σ_true * np.random.randn(n) ## Plot the data and the unobserved latent function fig = plt.figure(figsize=(12, 5)) ax = fig.gca() ax.plot(X, f_true, "dodgerblue", lw=3, label="True f") ax.plot(X, y, "ok", ms=3, alpha=0.5, label="Data") ax.set_xlabel("X") ax.set_ylabel("The true f(x)") plt.legend(); ```

Marginal with covariance function noise

This works:

with pm.Model() as model:
    ℓ = pm.Gamma("ℓ", alpha=2, beta=1)
    η = pm.HalfCauchy("η", beta=5)

    cov = η ** 2 * pm.gp.cov.Matern52(1, ℓ)
    gp = pm.gp.Marginal(cov_func=cov)

    σ = pm.HalfCauchy("σ", beta=5)
    noise = pm.gp.cov.WhiteNoise(sigma=σ)
    y_ = gp.marginal_likelihood("y", X=X[::100], y=y[::100], noise=noise)

    mp = pm.find_MAP()

MarginalSparse with RV noise

This works:

with pm.Model() as model:
    ℓ = pm.Gamma("ℓ", alpha=2, beta=1)
    η = pm.HalfCauchy("η", beta=5)

    cov = η ** 2 * pm.gp.cov.Matern52(1, ℓ)
    gp = pm.gp.MarginalSparse(cov_func=cov, approx="FITC")

    # initialize 20 inducing points with K-means
    # gp.util
    Xu = pm.gp.util.kmeans_inducing_points(20, X)

    σ = pm.HalfCauchy("σ", beta=5)
    y_ = gp.marginal_likelihood("y", X=X, Xu=Xu, y=y, noise=σ)

    mp = pm.find_MAP()

MarginalSparse with RV noise

But this throws a TypeError: Unsupported dtype for TensorType: object :

with pm.Model() as model:
    ℓ = pm.Gamma("ℓ", alpha=2, beta=1)
    η = pm.HalfCauchy("η", beta=5)

    cov = η ** 2 * pm.gp.cov.Matern52(1, ℓ)
    gp = pm.gp.MarginalSparse(cov_func=cov, approx="FITC")

    # initialize 20 inducing points with K-means
    # gp.util
    Xu = pm.gp.util.kmeans_inducing_points(20, X)

    σ = pm.HalfCauchy("σ", beta=5)
    noise = pm.gp.cov.WhiteNoise(sigma=σ)
    y_ = gp.marginal_likelihood("y", X=X, Xu=Xu, y=y, noise=noise)

    mp = pm.find_MAP()

Complete error traceback

```python --------------------------------------------------------------------------- KeyError Traceback (most recent call last) /usr/local/lib/python3.7/dist-packages/theano/tensor/type.py in dtype_specs(self) 279 "complex64": (complex, "theano_complex64", "NPY_COMPLEX64"), --> 280 }[self.dtype] 281 except KeyError: KeyError: 'object' During handling of the above exception, another exception occurred: TypeError Traceback (most recent call last) 11 frames /usr/local/lib/python3.7/dist-packages/theano/tensor/type.py in dtype_specs(self) 281 except KeyError: 282 raise TypeError( --> 283 f"Unsupported dtype for {self.__class__.__name__}: {self.dtype}" 284 ) 285 TypeError: Unsupported dtype for TensorType: object ```

Versions and main components

PyMC3 Version: 3.11.2
Aesara/Theano Version: 1.1.2
Python Version: 3.7.11
Operating system: Google Colab
How did you install PyMC3: pip

bwengals commented 3 years ago

I think the behavior you're seeing is the desired behavior, although the error message you're seeing could certainly be more informative. The DTC, FITC and VFE approximations are all derived starting with the assumption that the noise is constant diagonal (σ^2 I). The docstring for MarginalSparse does say that it only accepts a scalar for the noise parameter.

The FITC approximation can be used to model heteroscedastic noise, see section 3.1 in https://arxiv.org/pdf/1606.04820.pdf, and VFE can be modified similarly, see Section C in http://www2.aueb.gr/users/mtitsias/papers/sparseGPv2.pdf, but this isn't implemented in PyMC3.

Thank you by the way for the really nicely formatted bug report.

JohnGoertz commented 3 years ago

Ah, gotcha, thanks for pointing that out! I agree that a NotImplementedError would be more informative in this case. Heteroscedastic noise would be useful, but I can see that it might be a bit tricky to implement. Maybe at some point I could dig into those details and give it a go myself, but not at the moment.

And glad to be a good example of a bug report!

pymc-devs / pymc