Predicting on new test inputs after conditioning

Thanks for tinygp, I find it very instructive and am planning to use it my teaching and research.

One feature I am so far missing is the ability to make predictions on a series of new test inputs after conditioning a GP. The current suggested usage seems to be to call GaussianProcess(kernel, X_train).condition(y_train, X_test) repeatedly for each X_test but this is a little wasteful if the training set is large, since it is possible to reuse the alpha coefficients.

I appreciate that sometimes one can combine all the predictions that will ever be needed into a single X_test matrix., but this is not always the case, e.g. when using a GP surrogate in an optimisation problem.

Here is a MWE showing a possible workaround:

import numpy as np
import tinygp

X_train = np.linspace(0, 1, 10)
y_train = np.sin(x)
X_test = np.linspace(0, 1, 100)

kernel = tinygp.kernels.ExpSquared()
gp = tinygp.GaussianProcess(kernel, x)
cond_gp = gp.condition(y, X_test).gp

y2 = cond_gp.mean_function.kernel.matmul(X_test, cond_gp.mean_function.X, cond_gp.mean_function.alpha)

As expected, this calculation reproduces the stored mean:

>>> np.abs(y2 - cond_gp.mean).max()
2.488494e-06

In my view it would be nice to expose this functionality a little more conveniently. Browsing the source, it looks like tinygp.means.Conditioned.__call__() should allow this, but for this example it gives a vector of length 1 rather than length 100, I think due to a problem with the vmap axes:

>>> cond_gp.mean_function(X_test)
Array(-0.0139983, dtype=float32)

I'd be happy to make a PR to change this function to work similarly to my example if you agree?

Once that is done, perhaps GaussianProcess.predict() could be modified to check if the GP has already been conditioned and if so use the stored alpha values? This would mean making the y argument optional. There might be a better way to expose the funtionality.

Start with tinygp.means.Conditioned.call() in src/tinygp/means.py, then inspect GaussianProcess.predict() and how conditioned alpha values are stored. Reproduce the reported X_test behavior first; done means supporting predictions on new inputs after conditioning without repeatedly recomputing the training coefficients, with the API behavior agreed and covered by tests.