Plotting the Himmelblau function

plotHimmelblau.py

import numpy as np
import matplotlib.pyplot as plt
# ################################################################################
def plot_function():
    # Make data.
    X = np.arange(-5, 5, 0.25)
    Y = np.arange(-5, 5, 0.25)
    X, Y = np.meshgrid(X, Y)
    Z = (X ** 2 + Y - 11) ** 2 + (X + Y ** 2 - 7) ** 2

    # Plot the surface.
    fig = plt.figure()
    ax = fig.gca(projection='3d')
    surf = ax.plot_surface(X, Y, Z, cmap=cm.coolwarm,
                           linewidth=0, antialiased=False)
    ax.set_xlabel('X Label')
    ax.set_ylabel('Y Label')
    ax.set_zlabel('Z Label')
    return X, Y, Z
# ################################################################################
X, Y, Z = plot_function()

• change on public gist

Thank you for the public code !
I think the gist may lack an import for the cm.coolwarm part.
You may want to use matplotlib.colormaps["coolwarm"] or an equivalent matplotlib.cm.coolwarm.

For 3D plots there is this line that works well in many cases:

fig, ax = plt.subplots(subplot_kw={'projection': '3d'})

Hello @floffy-f ! Yes, thanks a lot! I think it
fig, ax = plt.subplots(subplot_kw={'projection': '3d'})
may minimizes this code

new one

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D # It is important for 3d plot
# update by Alex
# 05.03.2024
# ################################################################################
def plot_function():
    # Make data.
    X = np.arange(-5, 5, 0.25)
    Y = np.arange(-5, 5, 0.25)
    X, Y = np.meshgrid(X, Y)
    Z = (X ** 2 + Y - 11) ** 2 + (X + Y ** 2 - 7) ** 2

    # Plot the surface.
    #     # fig = plt.figure()
    #     # ax = plt.axes(projection='3d')
    #  Eqvivalent code from @floffy-f
    fig, ax = plt.subplots(subplot_kw={'projection': '3d'})

    # Syntax for plotting
    # cmap = 'viridis'
    # cmap='coolwarm'
    ax.plot_surface(X, Y, Z, cmap='viridis', \
                    edgecolor='green', linewidth=0, antialiased=False)

    ax.set_xlabel('X Label')
    ax.set_ylabel('Y Label')
    ax.set_zlabel('Z Label')
    plt.show()
    return X, Y, Z
# ################################################################################
X, Y, Z = plot_function()```

Perfect ! Thank you for your response and cool code.