无法一次填充 matplotlib 动画帧

Question

我目前正在尝试构建 N 体模拟，但在按我想要的方式绘制结果时遇到了一点麻烦。

在下面的代码中（包含轨道中几个点的一些示例数据），我正在导入位置和时间数据并将其组织到 pandas 数据框中。为了创建 3D 动画，我使用了 matplotlib 的动画类，效果很好。

但是，设置动画的常用方法是有限的，因为您不能单独自定义每个帧中的点（如果我在这里错了，请告诉我 :p）。由于我的动画显示的是轨道物体，我想改变它们的大小和颜色。为此，我基本上为每个主体创建了一个图形并设置它的颜色等。当它到达 update_graph 函数时，我遍历 n 主体，检索它们各自的 (x,y,z) 坐标，并更新他们的图表。

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d.axes3d import get_test_data
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation as animation
import pandas as pd

nbodies = 2

x = np.array([[1.50000000e-10, 0.00000000e+00, 0.00000000e+00],
              [9.99950000e-01, 1.00000000e-02, 0.00000000e+00],
              [4.28093585e-06, 3.22964816e-06, 0.00000000e+00],
              [-4.16142210e-01, 9.09335149e-01, 0.00000000e+00],
              [5.10376489e-06, 1.42204430e-05, 0.00000000e+00],
              [-6.53770813e-01, -7.56722445e-01, 0.00000000e+00]])
t = np.array([0.01, 0.01, 2.0, 2.0, 4.0, 4.0])
tt = np.array([0.01, 2.0, 4.0])

x = x.reshape((len(tt), nbodies, 3))

x_coords = x[:, :, 0].flatten()
y_coords = x[:, :, 1].flatten()
z_coords = x[:, :, 2].flatten()

df = pd.DataFrame({"time": t[:] ,"x" : x_coords, "y" : y_coords, "z" : z_coords})
print(df)

def update_graph(num):
    data=df[df['time']==tt[num]] # x,y,z of all bodies at current time
    for n in range(nbodies): # update graphs
        data_n = data[data['x']==x_coords[int(num * nbodies) + n]] # x,y,z of body n
        graph = graphs[n]
        graph.set_data(data_n.x, data_n.y)
        graph.set_3d_properties(data_n.z)
        graphs[n] = graph
    return graphs

plt.style.use('dark_background')
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.set_xlabel('x (AU)')
ax.set_ylabel('y (AU)')
ax.set_zlabel('z (AU)')

plt.xlim(-1.5,1.5)
plt.ylim(-1.5,1.5)

# initialize
data=df[df['time']==0]
ms_list = [5, 1]
c_list = ['yellow', 'blue']
graphs = []
for n in range(nbodies):
    graphs.append(ax.plot([], [], [], linestyle="", marker=".", 
                            markersize=ms_list[n], color=c_list[n])[0])

ani = animation.FuncAnimation(fig, update_graph, len(tt), 
                               interval=400, blit=True, repeat=True)

plt.show()

但是，这样做会给我以下错误：

Traceback (most recent call last):
  File "/home/kris/anaconda3/lib/python3.7/site-packages/matplotlib/backend_bases.py", line 1194, in _on_timer
    ret = func(*args, **kwargs)
  File "/home/kris/anaconda3/lib/python3.7/site-packages/matplotlib/animation.py", line 1447, in _step
    still_going = Animation._step(self, *args)
  File "/home/kris/anaconda3/lib/python3.7/site-packages/matplotlib/animation.py", line 1173, in _step
    self._draw_next_frame(framedata, self._blit)
  File "/home/kris/anaconda3/lib/python3.7/site-packages/matplotlib/animation.py", line 1193, in _draw_next_frame
    self._post_draw(framedata, blit)
  File "/home/kris/anaconda3/lib/python3.7/site-packages/matplotlib/animation.py", line 1216, in _post_draw
    self._blit_draw(self._drawn_artists, self._blit_cache)
  File "/home/kris/anaconda3/lib/python3.7/site-packages/matplotlib/animation.py", line 1231, in _blit_draw
    a.axes.draw_artist(a)
  File "/home/kris/anaconda3/lib/python3.7/site-packages/matplotlib/axes/_base.py", line 2661, in draw_artist
    a.draw(self.figure._cachedRenderer)
  File "/home/kris/anaconda3/lib/python3.7/site-packages/matplotlib/artist.py", line 38, in draw_wrapper
    return draw(artist, renderer, *args, **kwargs)
  File "/home/kris/anaconda3/lib/python3.7/site-packages/mpl_toolkits/mplot3d/art3d.py", line 202, in draw
    xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, renderer.M)
  File "/home/kris/anaconda3/lib/python3.7/site-packages/mpl_toolkits/mplot3d/proj3d.py", line 201, in proj_transform
    vec = _vec_pad_ones(xs, ys, zs)
  File "/home/kris/anaconda3/lib/python3.7/site-packages/mpl_toolkits/mplot3d/proj3d.py", line 189, in _vec_pad_ones
    return np.array([xs, ys, zs, np.ones_like(xs)])
  File "/home/kris/anaconda3/lib/python3.7/site-packages/pandas/core/series.py", line 871, in __getitem__
    result = self.index.get_value(self, key)
  File "/home/kris/anaconda3/lib/python3.7/site-packages/pandas/core/indexes/base.py", line 4405, in get_value
    return self._engine.get_value(s, k, tz=getattr(series.dtype, "tz", None))
  File "pandas/_libs/index.pyx", line 80, in pandas._libs.index.IndexEngine.get_value
  File "pandas/_libs/index.pyx", line 90, in pandas._libs.index.IndexEngine.get_value
  File "pandas/_libs/index.pyx", line 138, in pandas._libs.index.IndexEngine.get_loc
  File "pandas/_libs/hashtable_class_helper.pxi", line 997, in pandas._libs.hashtable.Int64HashTable.get_item
  File "pandas/_libs/hashtable_class_helper.pxi", line 1004, in pandas._libs.hashtable.Int64HashTable.get_item
KeyError: 0
Aborted (core dumped)

我不确定这究竟意味着什么，但我知道问题在于仅使用一行坐标而不是全部三行来更新图形。因为如果我有

def update_graph(num):
    data=df[df['time']==tt[num]] # x,y,z of all bodies at current time
    for n in range(nbodies): # update graphs
        #data_n = data[data['x']==x_coords[int(num * nbodies) + n]] # x,y,z of body n
        graph = graphs[n]
        graph.set_data(data.x, data.y)  # using data rather than data_n here now
        graph.set_3d_properties(data.z)
        graphs[n] = graph
    return graphs

它确实有效，并按照您的预期绘制了三个具有不同颜色和大小的身体副本。

任何帮助将不胜感激。谢谢！

Answer 1

我不明白您为什么要通过 pandas DataFrame，而您的 numpy 数组中似乎已经拥有所需的所有数据。我无法重现最初的问题，我提出了这个使用纯 numpy 数组的解决方案，这可能会解决问题：

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d.axes3d import get_test_data
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation as animation
import pandas as pd

nbodies = 2

x = np.array([[1.50000000e-10, 0.00000000e+00, 0.00000000e+00],
              [9.99950000e-01, 1.00000000e-02, 0.00000000e+00],
              [4.28093585e-06, 3.22964816e-06, 0.00000000e+00],
              [-4.16142210e-01, 9.09335149e-01, 0.00000000e+00],
              [5.10376489e-06, 1.42204430e-05, 0.00000000e+00],
              [-6.53770813e-01, -7.56722445e-01, 0.00000000e+00]])
t = np.array([0.01, 0.01, 2.0, 2.0, 4.0, 4.0])
tt = np.array([0.01, 2.0, 4.0])
x = x.reshape((len(tt), nbodies, 3))


def update_graph(i):
    data = x[i, :, :]  # x,y,z of all bodies at current time
    for body, graph in zip(data, graphs):  # update graphs
        graph.set_data(body[0], body[1])
        graph.set_3d_properties(body[2])
    return graphs


plt.style.use('dark_background')
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.set_xlabel('x (AU)')
ax.set_ylabel('y (AU)')
ax.set_zlabel('z (AU)')

plt.xlim(-1.5, 1.5)
plt.ylim(-1.5, 1.5)

# initialize
ms_list = [50, 10]
c_list = ['yellow', 'blue']
graphs = []
for n in range(nbodies):
    graphs.append(ax.plot([], [], [], linestyle="", marker=".",
                          markersize=ms_list[n], color=c_list[n])[0])

ani = animation.FuncAnimation(fig, func=update_graph, frames=len(tt),
                              interval=400, blit=True, repeat=True)

plt.show()