如何在 Python 中创建动画散点图且安全为 .mp4？

Question

我在 python 中有一个数组 A，其中包含 n x n x n 矩阵，每个位置的值为 1 或 -1。

np.random.choice([1, -1], size=(n, n, n))

假设 A 包含 20 个这样的矩阵。我想创建一个 3d 散点图来可视化这些矩阵，如果某个位置的矩阵值为 1，则显示一个蓝点，如果为 -1，则显示一个红点。

A 的每个条目应在显示下一个条目之前显示固定的时间，例如 0.25 秒。由于在运行时对大 n 执行此操作的计算量非常大，我想将其渲染为 .mp4。

为了更好地理解，我希望它看起来像这样：click （来源：https://github.com/ddatsko/3d-ising-model）

感谢您的帮助！

（上下文：我想可视化 3D Ising 模型的优化）

我像这样尝试过，它可以工作，但是对于大于个位数的 n 来说速度慢得令人痛苦。

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()

Answer 1

我不确定你是否一定能提高效率，尽管可能有选择。

但是，如果您有多个核心，那么通过使用

joblib

并行化三重嵌套 for 循环来简单地加速它可能就足够了：

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FFMpegWriter
from joblib import Parallel, delayed
from itertools import product

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_scatter(matrix, i, j, k):
    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')


def plot_matrix(matrix, ax):
    ax.clear()
    ax.set_zlim(0, matrix_size)
    
    Parallel(n_jobs=-2, verbose=51)(
        delayed(plot_scatter)(matrix, i, j, k) for i, j, k  in product(
            range(matrix_size), range(matrix_size), range(matrix_size)
        )
    )

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