Function plot_problem_surface

pymoo/visualization/util.py:254–321

(
    problem,
    n_samples,
    plot_type="wireframe",
    cmap="summer",
    show=True,
    return_figure=False,
)

def plot_problem_surface(
    problem,
    n_samples,
    plot_type="wireframe",
    cmap="summer",
    show=True,
    return_figure=False,
):  # noqa: ANN201
    try:
        from pymoo.visualization.matplotlib import plt  # noqa: F401
    except Exception as e:  # noqa: E722
        raise Exception(
            "Please install 'matplotlib' to use the plotting functionality."
        ) from e

    fig = plt.figure()

    if problem.n_var == 1 and problem.n_obj == 1:
        X = np.linspace(problem.xl[0], problem.xu[0], num=n_samples)[:, None]
        Y = problem.evaluate(X, return_values_of=["F"])
        ax = plt.plot(X, Y)

    elif problem.n_var == 2 and problem.n_obj == 1:
        X_range = np.linspace(problem.xl[0], problem.xu[0], num=n_samples)
        Y_range = np.linspace(problem.xl[1], problem.xu[1], num=n_samples)
        X, Y = np.meshgrid(X_range, Y_range)

        A = np.zeros((n_samples * n_samples, 2))
        counter = 0
        for i, x in enumerate(X_range):
            for j, y in enumerate(Y_range):
                A[counter, 0] = x
                A[counter, 1] = y
                counter += 1

        F = np.reshape(
            problem.evaluate(A, return_values_of=["F"]), (n_samples, n_samples)
        )

        # Plot the surface.
        if plot_type == "wireframe":
            ax = fig.add_subplot(111, projection="3d")
            ax.plot_wireframe(X, Y, F)
        elif plot_type == "contour":
            CS = plt.contour(X, Y, F)
            plt.clabel(CS, inline=1, fontsize=10)
        elif plot_type == "wireframe+contour":
            ax = fig.add_subplot(111, projection="3d")
            ax.plot_surface(X, Y, F, cmap=cmap, rstride=1, cstride=1)
            ax.contour(X, Y, F, 10, linestyles="solid", offset=-1)
            ax.set_xlabel("$x_1$")
            ax.set_ylabel("$x_2$")
            ax.set_zlabel("$f(x)$")
            ax.view_init(45, 45)
        else:
            raise Exception("Unknown plotting method.")

    else: