Repository navigation
Sampling from a conditioned GP sometimes returns nans if the kernel argument of .condition() is not None #175
Description
Activity
I think that it's a bit of a red herring here that you're finding that the
kernelargument is the thing that fixes your result, since (I think!) the kernel that you're passing there is not the same as the one being used to compute the GP.Getting NaNs when sampling from a conditioned GP isn't all that uncommon since the matrices become very poorly conditioned. The usual approach is to add some extra (artificial) variance to the diagonal of the conditioned GP. You can do this by passing
diag=1e-6(or some other small number - you may need to experiment) togp.condition. (Side note: it might be useful to add an option like this to thesamplefunction too, but it's somewhat subtle how to do that sensibly since we currently re-use the factorization from init.)Give that a shot and see if you can get it to work! I know that this isn't super satisfying, but it really is standard practice (and most GP packages actually just add this diagonal silently!).
Thanks for the quick reply, @dfm !
After some more playing around, I think the problem might be the
quasisepCosine kernel! I've attached an example script that shows this:import jax jax.config.update("jax_enable_x64", True) jax.config.update("jax_platform_name", "cpu") import numpy as np import tinygp def build_gp(kernel: tinygp.kernels.Kernel) -> tinygp.GaussianProcess: """ Create a GP. Parameters ---------- kernel : tinygp.kernels.Kernel Covariance function. Returns ------- tinygp.GaussianProcess GP. """ return tinygp.GaussianProcess(kernel, x_train, diag=yerr**2) # define training set x_train = np.linspace(0, 10, 20) y = np.sin(2*np.pi*x_train/3) # sine wave with a period of 3 yerr = .5 y += np.random.normal(0, yerr, size=y.shape[0]) # define test points x_test = np.linspace(0, 10, 100) # get GP posterior with no added variance gp = build_gp(tinygp.kernels.quasisep.Cosine(scale=3.)) # scalable cosine kernel with period of 3 cond_gp = gp.condition(y, X_test=x_test).gp mu, std = cond_gp.mean, np.sqrt(cond_gp.variance) print(cond_gp.sample(jax.random.PRNGKey(1), shape=(1,))) # nans # get GP posterior with added variance cond_gp = gp.condition(y, X_test=x_test, diag=1).gp mu, std = cond_gp.mean, np.sqrt(cond_gp.variance) print(cond_gp.sample(jax.random.PRNGKey(1), shape=(1,))) # nans # get GP posterior with no added variance using standard cosine kernel gp = build_gp(tinygp.kernels.Cosine(scale=3.)) # non-scalable cosine kernel with period of 3 cond_gp = gp.condition(y, X_test=x_test).gp mu, std = cond_gp.mean, np.sqrt(cond_gp.variance) print(cond_gp.sample(jax.random.PRNGKey(1), shape=(1,))) # NOT nansWhen running this script, samples from the posteriors of three GPs are printed: the first sample is from a GP that uses the
quasisepCosine kernel and has no diagonal term passed tocondition(), the second sample is from the same GP but now withdiag=1passed tocondition(), and the final sample is from a GP that uses the standard (i.e., non-scalable) Cosine kernel with no diagonal term passed tocondition(). Of all three samples, the only one that doesn't return NaNs is the final one (non-scalable cosine kernel with no additional variance added).Even with large diagonal terms passed to
condition()I always get NaNs when trying to sample from the posterior of a GP that uses thequasisepCosine kernel. That said, I can sample from the prior of such a GP, and I have no issues when I use the base (i.e., non-scalable) Cosine kernel.I also tried both scalable and non-scalable implementations of the Matern32 and Exp kernels, and neither showed the same behaviour as the
quasisepCosine kernel. It seems to me that this is the culprit, and not passing akernelargument tocondition()as I originally thought. Can you verify this?
Hi, I've noticed that trying to draw samples from a conditioned GP sometimes returns
nans if the kernel argument incondition()is notNoneusing version 0.2.3. I've attached a simple example script that shows this using a GP with a Cosine + Matern-3/2 covariance function:The above script outputs:
for the sample from the GP when the kernel argument of
.condition()is unspecified, and:for the sample from the GP when the kernel argument is specified as the Cosine component of
Bizarrely, if I set the kernel argument to the Matern-3/2 component, sampling from the conditioned GP does not return
Kernel(). As a sanity check, I plot this distribution and I see no reason why sampling from it should be returningnans:nans. Any insight as to what might be going on here? Am I overlooking an obvious mistake in my code?