Создание трехмерной движущейся диаграммы рассеяния с помощью matplotlib

Question

Я хочу создать 3D-график в matplotlib из уже существующего 2D-графика. На 2D графике есть точки, которые беспорядочно перемещаются внутри круга. Теперь я сделал оси 3D, и, насколько я могу судить, точки тоже находятся в трехмерном положении. Но они не перемещаются, и я получаю сообщение об ошибке: «ValueError: слишком много значений для распаковки (ожидается 2)». Я не нашел, где я сделал ошибку (думаю, где-то программа все еще требует только координаты x, y). Не могли бы вы помочь мне?

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

def get_initial_coordinates():
    x_coord =[random.uniform(3, 7) for i in range(n_particles)]
    y_coord = [random.uniform(3, 7) for i in range(n_particles)]
    z_coord = [random.uniform(3, 7) for i in range(n_particles)]
    return x_coord, y_coord, z_coord

def get_initial_velocities():
    x_vel = [3 * (np.random.random() - 0.5) * box_width for i in range(n_particles)]
    y_vel = [3 * (np.random.random() - 0.5) * box_width for i in range(n_particles)]
    z_vel = [3 * (np.random.random() - 0.5) * box_width for i in range(n_particles)]
    return x_vel, y_vel, z_vel

def take_step(x_coord, y_coord, z_coord, x_vel, y_vel, z_vel):
    for i in range(n_particles):
     x_coord[i] += x_vel[i]*dt
     y_coord[i] += y_vel[i]*dt
     z_coord[i] += z_vel[i]*dt

    return x_coord, y_coord,z_coord, x_vel, y_vel,z_vel

n_particles = 40
box_width = 10
n_steps = 5000
dt = 0.001

x_coord, y_coord, z_coord = get_initial_coordinates()
x_vel, y_vel, z_vel = get_initial_velocities()
for i in range(n_steps):
    x_coord,y_coord,z_coord,x_vel,y_vel,z_vel= take_step(x_coord,y_coord,z_coord,x_vel,y_vel,z_vel)
#---------------------
fig = pyplot.figure()
ax= Axes3D(fig)
d, = ax.plot([x_coord[i] for i in range(n_particles)],
             [y_coord[i] for i in range(n_particles)],
             [z_coord[i] for i in range(n_particles)], 'ro')
x_coord, y_coord, z_coord = get_initial_coordinates()
x_vel, y_vel, z_vel = get_initial_velocities()
circle = plt.Circle((5, 5), 3, color='blue', fill=False)
ax.add_artist(circle)

def animate(i):
    take_step(x_coord, y_coord, z_coord, x_vel, y_vel, z_vel)
    d.set_data([x_coord[i]for i in range(n_particles)],
               [y_coord[i]for i in range(n_particles)],
               [z_coord[i]for i in range(n_particles)])
    return d,

anim = animation.FuncAnimation(fig, animate, frames=200, interval=20)
pyplot.show()

Answer 1

set_data() наследуется от Line2D, поэтому он может принимать только 2 списка, один для x, другой для y. Вы должны использовать .set_3d_properties() для компонента z.

def animate(i):
    take_step(x_coord, y_coord, z_coord, x_vel, y_vel, z_vel)
    d.set_data([x_coord[i]for i in range(n_particles)],
               [y_coord[i]for i in range(n_particles)],
               )
    d.set_3d_properties([z_coord[i]for i in range(n_particles)])
    return d,