Matplotlib绘制三维图实现图形旋转、颜色变换和视频及动画输出

1. Matplotlib绘制三维图

通过Matplotlib.pyplot.axes.plot_surface来绘制一个三维曲面，代码如下：

# -*_ coding: utf-8 -*-

import numpy as np
import matplotlib.pyplot as plt

x = np.linspace(-6, 6, 240)
y = np.linspace(-6, 6, 240)
X, Y = np.meshgrid(x, y)
Z = np.sin(np.sqrt(X**2+Y**2))*2

fig = plt.figure(figsize=(13, 10))
ax3d = plt.axes(projection='3d')

ax3d.set_xlabel("X")
ax3d.set_ylabel("Y")
ax3d.set_zlabel("Z")
ax3d.set_zlim(-3.01, 3.01)
ax3d.set_xlim(-6.01, 6.01)
ax3d.set_ylim(-6.01, 6.01)

plt.tick_params(labelsize=10)

surf = ax3d.plot_surface(X, Y, Z, cmap = plt.cm.YlGnBu_r)
ax3d.set_title(r"Surface Z = 3sin($\sqrt{X^2 + y^2})$", fontsize=18)
fig.colorbar(surf, shrink=0.5, aspect=20)
plt.tight_layout()                          #减少窗口空白
plt.show()

2. 实现三维图形的旋转

Matplotlib.pyplot.axes.view_init()可以设置三维图形的视角，所有可以通过视角的变化来实现三维图形的旋转，代码如下：

for angle in range(0, 360):
    ax3d.view_init(30, angle)
    plt.draw()
    plt.pause(0.1)

其中，view_init()的第一个参数是上下视角旋转角度，第二个参数是x,y平面视角旋转角度。

3.实现三维图形颜色的变换

 Matplotlib.colors.LinearSegmentedColormap可以使用颜色数列自定义cmap颜色，使用surf.set_cmap()和plt.setp()可以更改三维图形曲面和网格的颜色。这样通过对颜色数列的移位，就可实现颜色的变换了，代码如下：

from matplotlib.colors import LinearSegmentedColormap

clist = [(26/255, 49/255, 139/255), (73/255, 108/255, 206/255), (130/255, 170/255, 231/255),
         (185/255, 210/255, 243/255), (230/255, 240/255, 254/255), (249/255, 219/255, 229/255),
         (247/255, 166/255, 191/255), (228/255, 107/255, 144/255), (192/255, 63/255, 103/255),
         (154/255, 19/255, 61/255)] # 颜色数值需在0和1之间
mycmp = LinearSegmentedColormap.from_list("No001", clist)

for angle in range(0, 360):
    ax3d.view_init(30, angle)
    clist = clist[-1:] + clist[:-1]     #颜色列表向左移一位
    mycmp = LinearSegmentedColormap.from_list("No001", clist)
    plt.setp(surf, color=clist[-1])
    surf.set_cmap(mycmp)
    plt.draw()
    plt.pause(0.1)

4. 三维效果的视频输出

Matplotlib.animation.FFMpegWriter利用FFMpeg实现视频的输出，首先需要安装FFMpeg程序，如果程序找不到ffmpeg.exe文件，会报错“找不到系统文件”，需要使用plt.rcParams来指定文件位置。最后，调用cmd，使用ffmpeg给得到的视频添加上音乐。

from matplotlib.animation import FFMpegWriter
import os
plt.rcParams['animation.ffmpeg_path'] ='C:\\ffmpeg\\bin\\ffmpeg.exe'
metadata = dict(title='Rotating', artist='Matplotlib', comment='Rotating Video')
writer = FFMpegWriter(fps=10, metadata=metadata, extra_args=['-vcodec', 'libx264'])
file_dir = 'D:\\works\\Python_Qt\\matplotlib_figs\\video\\'

with writer.saving(fig, file_dir+'out1_ffmpeg.mp4', 120):
    for angle in range(0, 360):
        ax3d.view_init(30, angle)
        clist = clist[-1:] + clist[:-1]     #颜色列表向左移一位
        mycmp = LinearSegmentedColormap.from_list("No001", clist)
        plt.setp(surf, color=clist[-1])
        surf.set_cmap(mycmp)
        plt.draw()
        writer.grab_frame()
        plt.pause(0.1)
cmd = "C:\\ffmpeg\\bin\\ffmpeg.exe -i %s -i %s %s" % (file_dir+'music.mp3', file_dir+'out1_ffmpeg.mp4', file_dir+'out1_ffmpeg_music.mp4')
os.system(cmd)

其中，grab_frame()用来抓取窗口内容

5. 三维效果的动画输出

Matplotlib.animation.PillowWriter利用Pillow实现gif动画的输出，代码和视频输出类似：

from matplotlib.animation import PillowWriter
metadata = dict(title='Rotating', artist='Yisl04', comment='Rotating Video')
writer = PillowWriter(fps=10, metadata=metadata)
        
with writer.saving(fig, file_dir+'out1_pillow.gif', 120):
    for angle in range(0, 360):
        ax3d.view_init(30, angle)
        clist = clist[-1:] + clist[:-1]     #颜色列表向左移一位
        mycmp = LinearSegmentedColormap.from_list("No001", clist)
        plt.setp(surf, color=clist[-1])
        surf.set_cmap(mycmp)
        plt.draw()
        writer.grab_frame()
        plt.pause(0.1) 

最后，整合上述代码：

# -*_ coding: utf-8 -*-

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FFMpegWriter
from matplotlib.colors import LinearSegmentedColormap

clist = [(26/255, 49/255, 139/255), (73/255, 108/255, 206/255), (130/255, 170/255, 231/255),
         (185/255, 210/255, 243/255), (230/255, 240/255, 254/255), (249/255, 219/255, 229/255),
         (247/255, 166/255, 191/255), (228/255, 107/255, 144/255), (192/255, 63/255, 103/255),
         (154/255, 19/255, 61/255)] # 颜色数值需在0和1之间
mycmp = LinearSegmentedColormap.from_list("No001", clist)

x = np.linspace(-6, 6, 240)
y = np.linspace(-6, 6, 240)
X, Y = np.meshgrid(x, y)
Z = np.sin(np.sqrt(X**2+Y**2))*2

fig = plt.figure(figsize=(13, 10))
ax3d = plt.axes(projection='3d')

ax3d.set_xlabel("X")
ax3d.set_ylabel("Y")
ax3d.set_zlabel("Z")
ax3d.set_zlim(-3.01, 3.01)
ax3d.set_xlim(-6.01, 6.01)
ax3d.set_ylim(-6.01, 6.01)

plt.tick_params(labelsize=10)

surf = ax3d.plot_surface(X, Y, Z, cmap = mycmp)
ax3d.set_title(r"Surface Z = 3sin($\sqrt{X^2 + y^2})$", fontsize=18)
#shrink表示图例缩放比例，aspect表示长宽比
fig.colorbar(surf, shrink=0.5, aspect=20)
plt.tight_layout()                          #减少窗口空白

#定义视频和动画输出函数
def video_output(video_format, clist):
    file_dir = 'D:\\works\\Python_Qt\\matplotlib_figs\\video\\'
    if video_format == 'ffmpeg':
        
        plt.rcParams['animation.ffmpeg_path'] ='C:\\ffmpeg\\bin\\ffmpeg.exe'
        metadata = dict(title='Rotating', artist='matplotlib', comment='Rotating Video')
        writer = FFMpegWriter(fps=10, metadata=metadata, extra_args=['-vcodec', 'libx264'])

        with writer.saving(fig, file_dir+'out1_ffmpeg.mp4', 120):
            for angle in range(0, 360):
                ax3d.view_init(30, angle)
                clist = clist[-1:] + clist[:-1]     #颜色列表向左移一位
                mycmp = LinearSegmentedColormap.from_list("No001", clist)
                plt.setp(surf, color=clist[-1])
                surf.set_cmap(mycmp)
                plt.draw()
                writer.grab_frame()
                plt.pause(0.1)
                
        cmd = "C:\\ffmpeg\\bin\\ffmpeg.exe -i %s -i %s %s" % (file_dir+'music.mp3', file_dir+'out1_ffmpeg.mp4', file_dir+'out1_ffmpeg_music.mp4')
        os.system(cmd)
                
    elif video_format == 'pillow':
        metadata = dict(title='Rotating', artist='matplotlib', comment='Rotating Video')
        writer = PillowWriter(fps=10, metadata=metadata)
        
        with writer.saving(fig, file_dir+'out1_pillow.gif', 120):
            for angle in range(0, 360):
                ax3d.view_init(30, angle)
                clist = clist[-1:] + clist[:-1]     #颜色列表向左移一位
                mycmp = LinearSegmentedColormap.from_list("No001", clist)
                plt.setp(surf, color=clist[-1])
                surf.set_cmap(mycmp)
                plt.draw()
                writer.grab_frame()
                plt.pause(0.1) 
    else:
        ax3d.set_title(r"视频输出格式错误！", fontsize=18, color='red')
        plt.show()

#输出视频
video_output('ffmpeg', clist)
#输出动画
video_output('pillow', clist)

得到的视频和动画如下所示：

Matplotlib绘制三维图实现图形旋转、颜色变换

---