How to plot regression predicted data in 3-D contour plot?

Question

I'm trying to plot the predicted mean data from Gaussian process regression into a 3-D contour. I've followed Plot 3D Contour from an Image using extent with Matplotlib and mplot3d example code: contour3d_demo3.py threads. Following is my code:

import numpy as np
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import RBF, ConstantKernel as C
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d
from matplotlib import cm

x_train = np.array([[0,0],[2,2],[3,3]])
y_train = np.array([[200,321,417]])

xvalues = np.array([0,1,2,3])
yvalues = np.array([0,1,2,3])

a,b = np.meshgrid(xvalues,yvalues)
positions = np.vstack([a.ravel(), b.ravel()])
x_test = (np.array(positions)).T

kernel = C(1.0, (1e-3, 1e3)) * RBF(10)

gp = GaussianProcessRegressor(kernel=kernel)

gp.fit(x_train, y_train)

y_pred_test = gp.predict(x_test)

fig = plt.figure()
ax = fig.add_subplot(projection = '3d')
x=y=np.arange(0,3,1)
X, Y = np.meshgrid(x,y)
Z = y_pred_test
cset = ax.contour(X, Y, Z, cmap=cm.coolwarm)
ax.clabel(cset, fontsize=9, inline=1)
plt.show()

After running the above code, I get following error on console:

I want x and y-axis as 2-D plane and the predicted values on the z-axis.The sample plot is as follows:

What is wrong with my code?

Thank you!

Answer 1

The specific error you've mentioned comes from your y_train, which might be a typo. It should be:

y_train_ : array-like, shape = (n_samples, [n_output_dims])

According to your x_train, you have 3 samples. So your y_train should have shape (3, 1) rather than (1, 3).

You also have other bugs in the plotting part:

1. add_subplot should have a position before projection = '3d'.
2. Z should have the same shape as X and Y for contour plot.
3. Because of 2, your x and y should match xvalues and yvalues.

Taken together, you might need to make the following changes:

...

y_train = np.array([200,321,417])

...

ax = fig.add_subplot(111, projection = '3d')
x=y=np.arange(0,4,1)
...
Z = y_pred_test.reshape(X.shape)

...

Just to mention two things:

1. The plot you will get after these changes won't match the figure you've shown. The figure in your question is a surface plot instead of a contour plot. You can use ax.plot_surface to get that type of plot.

2. I think you've already know this. But just in case, your plot won't be as smooth as your sample plot since your np.meshgrid is sparse.