def gEKPLS(xt, yt, xtest, ytest, funXLimits, ndim):
    # 'n_comp' must be an integer in [1,ndim[,  'theta0' a list of n_comp values

    t = GEKPLS(n_comp=2,
               theta0=[1e-2, 1e-2],
               xlimits=funXLimits,
               delta_x=1e-2,
               extra_points=1,
               print_prediction=False)
    t.set_training_values(xt, yt)
    # Add the gradient information
    for i in range(ndim):
        t.set_training_derivatives(xt, yt[:, 1 + i].reshape((yt.shape[0], 1)),
                                   i)
    t.train()

    # Prediction of the validation points
    y = t.predict_values(xtest)
    print('GEKPLS1,  err: ' + str(compute_rms_error(t, xtest, ytest)))
    title = ('GEKPLS model')
    return t, title, xtest, ytest

    # Prediction of the derivatives with regards to each direction space
    yd_prediction = np.zeros((ntest, ndim))
    for i in range(ndim):
        yd_prediction[:, i] = t.predict_derivatives(xtest, kx=i).T
        print('GEKPLS1, err of the ' + str(i + 1) + '-th derivative: ' +
              str(compute_rms_error(t, xtest, ydtest[:, i], kx=i)))

        if plot_status:
            plt.plot(ydtest[:, i], ydtest[:, i], '-.')
            plt.plot(ydtest[:, i], yd_prediction[:, i], '.')

        if plot_status:
            plt.show()
    def test_gekpls(self):
        import numpy as np
        from mpl_toolkits.mplot3d import Axes3D
        import matplotlib.pyplot as plt

        from smt.surrogate_models import GEKPLS
        from smt.problems import Sphere
        from smt.sampling_methods import LHS

        # Construction of the DOE
        fun = Sphere(ndim=2)
        sampling = LHS(xlimits=fun.xlimits, criterion="m")
        xt = sampling(20)
        yt = fun(xt)
        # Compute the gradient
        for i in range(2):
            yd = fun(xt, kx=i)
            yt = np.concatenate((yt, yd), axis=1)

        # Build the GEKPLS model
        n_comp = 2
        sm = GEKPLS(
            theta0=[1e-2] * n_comp,
            xlimits=fun.xlimits,
            extra_points=1,
            print_prediction=False,
            n_comp=n_comp,
        )
        sm.set_training_values(xt, yt[:, 0])
        for i in range(2):
            sm.set_training_derivatives(xt, yt[:, 1 + i].reshape((yt.shape[0], 1)), i)
        sm.train()

        # Test the model
        X = np.arange(fun.xlimits[0, 0], fun.xlimits[0, 1], 0.25)
        Y = np.arange(fun.xlimits[1, 0], fun.xlimits[1, 1], 0.25)
        X, Y = np.meshgrid(X, Y)
        Z = np.zeros((X.shape[0], X.shape[1]))

        for i in range(X.shape[0]):
            for j in range(X.shape[1]):
                Z[i, j] = sm.predict_values(
                    np.hstack((X[i, j], Y[i, j])).reshape((1, 2))
                )

        fig = plt.figure()
        ax = fig.gca(projection="3d")
        surf = ax.plot_surface(X, Y, Z)

        plt.show()
def dEKPLS(xt, yt, xtest, ytest, funXLimits, ndim):
    # 'n_comp' must be an integer in [1,ndim[,  'theta0' a list of n_comp values

    t = GEKPLS(n_comp=1,
               theta0=[1e-2],
               xlimits=funXLimits,
               delta_x=1e-2,
               extra_points=1,
               print_prediction=False)
    t.set_training_values(xt, yt)
    # Add the gradient information
    for i in range(ndim):
        t.set_training_derivatives(xt, yt[:, 1 + i].reshape((yt.shape[0], 1)),
                                   i)
    t.train()

    title = 'GEKPLS model'
    return t, title, xtest, ytest