Surface plot for multivariate 5 degree polynomial

I am implementing a paper in Python, which was originally implemented in MATLAB. The paper says that a five degree polynomial was found using curve fitting from a set of sampling data points. I did not want to use their polynomial, so I started using the sample data points (given in paper) and tried to find a 5 degree polynomial using sklearn Polynomial Features and linear_model. As it is a multivariate equation f(x,y) where x and y are the length and width of a certain pond and f is the initial concentration of pollutant.

So my problem is that sklearn Polynomial Features transforms the test and train data points into n-polynomial points (I think as far as I understood it). However when I need the clf.predict function (where clf is the trained model) to take only the x and y values because when I am drawing the surface graph from Matplotlib, it requires meshgrid, so when I meshgrid my sklean-transformed test points, it's shape becomes like NxN whereas the predict function requires Nxn (where n is the number of degrees of polynomial in which it transformed the data) and N is the number of rows.

Is there any possible way to draw the mesh points for this polynomial?

Paper Link: http://www.ajer.org/papers/v5(11)/A05110105.pdf Paper title: Mathematical Model of Biological Oxygen Demand in Facultative Wastewater Stabilization Pond Based on Two-Dimensional Advection-Dispersion Model

If possible, please look at the figure 5 and 6 in the paper (link above). This is what I am trying to achieve.

Thank you enter code here

from math import exp
import numpy as np
from operator import itemgetter
from sklearn.preprocessing import PolynomialFeatures
from sklearn import linear_model
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter

fig = plt.figure()
ax = fig.gca(projection='3d')


def model_BOD (cn):
    cnp1 = []
    n = len(cn)
    # variables:
    dmx = 1e-5
    dmy = 1e-5
    u = 2.10e-4
    v = 2.10e-4
    obs_time = 100
    dt = 0.1

    for t in np.arange(0.1,obs_time,dt):
        for i in range(N):
            for j in range(N):
                d = j + (i-1)*N
                dxp1 = d  + N
                dyp1 = d + 1
                dxm1 = d - N
                dym1 = d - 1

                cnp1.append(t*(((-2*dmx/dx**2)+(-2*dmy/dy**2)+(1/t))*cn[dxp1] + (dmx/dx**2)*cn[dyp1] \
                                + (dmy/dy**2)*cn[dym1] - (u/(2*dx))*cn[dxp1] + (u/(2*dx))*cn[dxm1] \
                                - (v/(2*dy))*cn[dyp1] + (v/(2*dy))*cn[dym1]))
        cn = cnp1
        cnp1 = []
    return cn

N = 20
Length = 70
Width = 77
dx = Length/N
dy = Width/N

deg_of_poly = 5

datapoints = np.array([
    [12.5,70,81.32],[25,70,88.54],[37.5,70,67.58],[50,70,55.32],[62.5,70,56.84],[77,70,49.52],
    [0,11.5,71.32],[77,57.5,67.20],
    [0,23,58.54],[25,46,51.32],[37.5,46,49.52],
    [0,34.5,63.22],[25,34.5,48.32],[37.5,34.5,82.30],[50,34.5,56.42],[77,34.5,48.32],[37.5,23,67.32],
    [0,46,64.20],[77,11.5,41.89],[77,46,55.54],[77,23,52.22],
    [0,57.5,93.72],
    [0,70,98.20],[77,0,42.32]
    ])

X = datapoints[:,0:2]
Y = datapoints[:,-1]

predict_x = []
predict_y = []
for i in np.linspace(0,Width,N):
    for j in np.linspace(0,Length,N):
        predict_x.append([i,j])

predict = np.array(predict_x)

poly = PolynomialFeatures(degree=deg_of_poly)
X_ = poly.fit_transform(X)

predict_ = poly.fit_transform(predict)
clf = linear_model.LinearRegression()
clf.fit(X_, Y)
prediction = []

for k,i in enumerate(predict_):
    prediction.append(clf.predict(np.array([i]))[0])

prediction_ = model_BOD(prediction)
print prediction_
XX = []
XX = predict[:,0]
YY = []
YY = predict[:,-1]
XX,YY = np.meshgrid(X,Y)
Z = prediction
##R = np.sqrt(XX**2+YY**2)
##Z = np.tan(R)

surf = ax.plot_surface(XX,YY,Z)
plt.show()

If I understand correctly, the key logic here is to generate polynomial features from your meshgrid, make the prediction and plot the prediction with the original meshgrid. Hope the following gives what you want:

import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
from sklearn.preprocessing import PolynomialFeatures
from sklearn import linear_model

# The training set
datapoints = np.array([
    [12.5,70,81.32], [25,70,88.54], [37.5,70,67.58], [50,70,55.32], 
    [62.5,70,56.84], [77,70,49.52], [0,11.5,71.32], [77,57.5,67.20], 
    [0,23,58.54], [25,46,51.32], [37.5,46,49.52], [0,34.5,63.22], 
    [25,34.5,48.32], [37.5,34.5,82.30], [50,34.5,56.42], [77,34.5,48.32], 
    [37.5,23,67.32], [0,46,64.20], [77,11.5,41.89], [77,46,55.54], 
    [77,23,52.22], [0,57.5,93.72], [0,70,98.20], [77,0,42.32]
    ])
X = datapoints[:,0:2]
Y = datapoints[:,-1]
# 5 degree polynomial features
deg_of_poly = 5
poly = PolynomialFeatures(degree=deg_of_poly)
X_ = poly.fit_transform(X)
# Fit linear model
clf = linear_model.LinearRegression()
clf.fit(X_, Y)

# The test set, or plotting set
N = 20
Length = 70
predict_x0, predict_x1 = np.meshgrid(np.linspace(0, Length, N), 
                                     np.linspace(0, Length, N))
predict_x = np.concatenate((predict_x0.reshape(-1, 1), 
                            predict_x1.reshape(-1, 1)), 
                           axis=1)
predict_x_ = poly.fit_transform(predict_x)
predict_y = clf.predict(predict_x_)

# Plot
fig = plt.figure(figsize=(16, 6))
ax1 = fig.add_subplot(121, projection='3d')
surf = ax1.plot_surface(predict_x0, predict_x1, predict_y.reshape(predict_x0.shape), 
                        rstride=1, cstride=1, cmap=cm.jet, alpha=0.5)
ax1.scatter(datapoints[:, 0], datapoints[:, 1], datapoints[:, 2], c='b', marker='o')

ax1.set_xlim((70, 0))
ax1.set_ylim((0, 70))
fig.colorbar(surf, ax=ax1)
ax2 = fig.add_subplot(122)
cs = ax2.contourf(predict_x0, predict_x1, predict_y.reshape(predict_x0.shape))
ax2.contour(cs, colors='k')
fig.colorbar(cs, ax=ax2)
plt.show()