如何用Python高效生成3D散点图动画并保存为MP4

问题描述 #

我在Python中有一个数组A，包含多个n×n×n的矩阵，每个元素值为1或-1，可通过np.random.choice([1, -1], size=(n, n, n))生成。需要制作3D散点图动画：值为1的位置显示蓝色点，-1显示红色点；每个矩阵固定展示0.25秒后切换，最终保存为MP4文件。当前代码在n大于个位数时运行极慢，寻求优化方案。背景是可视化3D伊辛模型的优化过程，当前代码如下：

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

plt.rcParams['animation.ffmpeg_path']='path-to-ffmpeg.exe'

metadata = dict(title='Random-Boolean-Matrix', artist='Me')
writer = FFMpegWriter(fps=2, metadata=metadata)

fig, ax = plt.subplots(subplot_kw=dict(projection='3d'))

matrix_size = 3

# init A (array of 3d matrices)
A = []
for _ in range(10):
    A.append(np.random.choice([1, -1], size=(matrix_size, matrix_size, matrix_size)))

plt.xlim(0, matrix_size)
plt.ylim(0, matrix_size)

def plot_matrix(matrix, ax):
    ax.clear()
    ax.set_zlim(0, matrix_size)
    for i in range(matrix_size):
        for j in range(matrix_size):
            for k in range(matrix_size):
                if matrix[i, j, k] == 1:
                    ax.scatter(i, j, k, c='b', marker='o')
                else:
                    ax.scatter(i, j, k, c='r', marker='o')

with writer.saving(fig, "matrix_animation.mp4", 100):
    for i in range(10):
        plot_matrix(A[i], ax)
        writer.grab_frame()

核心优化思路 #

原代码的性能瓶颈在于逐点绘制散点+每次清空轴重绘，针对这两个问题，优化方向为：批量处理数据、复用绘图对象、减少重复计算。

1. 预计算所有帧的坐标与颜色 #

提前生成所有位置的网格坐标，再根据每个矩阵的值批量筛选颜色，避免实时嵌套循环遍历：

• 用np.meshgrid生成全量坐标数组
• 用np.where批量映射颜色，1对应蓝色，-1对应红色

2. 复用散点对象，仅更新数据 #

创建一次散点集合后，后续仅更新散点的颜色属性，无需清空轴重新绘制，大幅降低绘图开销。

3. 减少不必要的轴操作 #

初始化时一次性设置好轴的范围和比例，避免每帧重复执行布局计算。

4. 调整动画参数匹配需求 #

0.25秒每帧对应fps=4，替换原代码的fps=2，保证时间精度。

优化后的完整代码 #

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

plt.rcParams['animation.ffmpeg_path'] = 'path-to-ffmpeg.exe'

# 参数配置
matrix_size = 15  # 支持更大的n值
num_frames = 10
fps = 4  # 对应每帧展示0.25秒

# 生成所有3D矩阵
A = [np.random.choice([1, -1], size=(matrix_size, matrix_size, matrix_size)) for _ in range(num_frames)]

# 预生成全量网格坐标
x, y, z = np.meshgrid(np.arange(matrix_size), np.arange(matrix_size), np.arange(matrix_size))
all_coords = np.column_stack((x.flatten(), y.flatten(), z.flatten()))

# 预计算每帧的颜色数组
frame_colors = []
for mat in A:
    colors = np.where(mat.flatten() == 1, 'b', 'r')
    frame_colors.append(colors)

# 初始化3D绘图
fig, ax = plt.subplots(subplot_kw=dict(projection='3d'), figsize=(8, 8))
ax.set_xlim(0, matrix_size)
ax.set_ylim(0, matrix_size)
ax.set_zlim(0, matrix_size)
ax.set_box_aspect([1, 1, 1])  # 保持3D坐标轴比例一致

# 创建初始散点对象
scatter = ax.scatter(all_coords[:, 0], all_coords[:, 1], all_coords[:, 2], c=frame_colors[0], marker='o')

# 配置动画Writer
metadata = dict(title='3D伊辛模型动画', artist='Me')
writer = FFMpegWriter(fps=fps, metadata=metadata)

# 生成并保存动画
with writer.saving(fig, "ising_model_animation.mp4", dpi=100):
    writer.grab_frame()  # 保存第一帧
    for colors in frame_colors[1:]:
        scatter.set_facecolors(colors)  # 仅更新颜色
        fig.canvas.draw_idle()  # 强制更新绘图
        writer.grab_frame()

plt.close(fig)

额外优化建议 #

• 若n过大（如>20），可改用更小的标记（如'.'代替'o'）或降采样减少渲染负载
• 尝试使用matplotlib.animation.FuncAnimation替代手动循环，其内部有更高效的帧更新机制
• 极端大规模数据场景下，可考虑plotly或mayavi等专门的3D可视化库进一步提升性能

内容的提问来源于stack exchange，提问作者WhatIfBlueberry