matplotlibで等高線図を書く

programming python matplotlib numpy

schwefel.png:

  • ax.contour(x, y, z, zdir=, offset=<平面のどの位置に描画するか>)
  • contour関数のオプション(一部)
    • zdir: どの平面に投射するか(x, y, zのいずれか)
    • offset: 平面のどの位置に描画するか
    • cmap: カラーマップ
#!/usr/bin/env python
#coding:utf-8

from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import matplotlib.pyplot as plt
import numpy as np


def schwefel(x, y):
    return 418.9829 * 2 - (x * np.sin(np.sqrt(abs(x))) + \
                           y * np.sin(np.sqrt(abs(y))))


def work():
    N = 100
    MIN_X = -500
    MAX_X =  500

    x = np.linspace(MIN_X, MAX_X, N)
    y = np.linspace(MIN_X, MAX_X, N)
    x, y = np.meshgrid(x, y)
    z = schwefel(x, y)
    
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1, projection="3d")

    ax.contour(x, y, z, zdir='z', offset=0, cmap=cm.coolwarm)
    ax.contour(x, y, z, zdir='x', offset=MIN_X, cmap=cm.coolwarm)
    ax.contour(x, y, z, zdir='y', offset=MAX_X, cmap=cm.coolwarm)

    ax.set_title("Schwefel")

    fig.savefig("schwefel.png", bbox_inches="tight")
    

if __name__ == "__main__":
    work()