用matplotlib绘制成组的动画plt.step

a6b3iqyw  于 2023-06-30  发布在  其他
关注(0)|答案(1)|浏览(117)

我正在学习python,想画出跟随散点的plt.step线。散点绘制得很完美,所以我使用了几乎相同的方法来绘制plt.steps线,但它不起作用,我不清楚为什么,因为ax.set_data可以采用与ax.set_offsets相同的2D类型数组。
不幸的是,我遇到了以下问题:
a)只有一个plt.step行(应该与3个“轨道”相同)
B)在点后面一个单元格处绘制线
c)线的颜色与散射点不对应
我很感激任何建议!
电流输出:

代码示例:

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FuncAnimation

df = pd.DataFrame()
cf = 0
while cf < 3:
    df = pd.concat([df, pd.DataFrame(
        {
            "Track": f'Track {cf + 1}',
            "Timeline": np.linspace(0, 9, 10, dtype=int),
            "Position": np.random.randint(low=0+cf, high=3+cf, size=10)
        }
    )])
    cf = cf + 1

df = df.reset_index(drop=True)
print(df)
df.info()

# plot:
fig, ax = plt.subplots()

# Point coordinates:
y = df['Position']
x = df['Timeline']

# Labels with axes:
ax.set_xlabel('Timeline')
ax.set_ylabel('Position')
ax.set_ylim(-0.2, 4.2)
ax.invert_yaxis()
ax.set_xticks(list(np.unique(x)))
ax.set_yticks(list(np.unique(y)))
ax.set_xlim(df["Timeline"].min()-0.5, df["Timeline"].max()+0.5)

# Colors:
colors = {'Track 1': 'tab:red', 'Track 2': 'tab:blue', 'Track 3': 'blue'}

# Drawing points and lines according to positions:
frames = (len(df.groupby(['Timeline'])))

steps = []
scatters = []
for track, group in df.groupby("Track"):
    scatters.append(ax.scatter(group["Timeline"].iloc[0],
                               group["Position"].iloc[0],
                               s=45, c=colors[track]))
    steps = plt.step(group["Timeline"].iloc[0],
                     group["Position"].iloc[0],
                     color=colors[track])

def animate(i):
    for scatter, (_, group) in zip(scatters, df.groupby("Track")):
        scatter.set_offsets((group["Timeline"].iloc[i],
                             group["Position"].iloc[i]))
    for step, (_, group) in zip(steps, df.groupby('Track')):
        step.set_data(group['Timeline'].iloc[:i],
                       group['Position'].iloc[:i])
    print('step', i) #for some reason there are three 0 steps in the beginning

anim = FuncAnimation(fig, animate, frames=frames, interval=400)
ax.grid()

anim.save('test.gif', writer='pillow')
blmhpbnm

blmhpbnm1#

让我谈谈你的三个问题。
1.这里只有一行,因为尽管您打算将这些行添加到列表steps中,但您只是在每个循环中覆盖了该列表。因为ax.step返回一个只有一项的列表,所以需要对第一项建立索引,否则得到的将是一个列表,而不是所需的Line2D对象。
1.在python中使用slicing时,list[start:stop],slice inclusive of start index and exclusive of stop index。所以,当你做.iloc[:i]时,你没有包括索引i,所以你遗漏了你想要的点。相反,您需要执行.iloc[:i+1]
1.颜色不匹配是因为第一个问题,即它只获取循环中最后一项的结果(因为您一直在覆盖steps)。一旦你将它们添加到列表中,它们将被正确地着色。
下面是正确的代码:

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FuncAnimation

df = pd.DataFrame()
cf = 0
while cf < 3:
    df = pd.concat([df, pd.DataFrame(
        {
            "Track": f'Track {cf + 1}',
            "Timeline": np.linspace(0, 9, 10, dtype=int),
            "Position": np.random.randint(low=0+cf, high=3+cf, size=10)
        }
    )])
    cf = cf + 1

df = df.reset_index(drop=True)
print(df)
df.info()

# plot:
fig, ax = plt.subplots()

# Point coordinates:
y = df['Position']
x = df['Timeline']

# Labels with axes:
ax.set_xlabel('Timeline')
ax.set_ylabel('Position')
ax.set_ylim(-0.2, 4.2)
ax.invert_yaxis()
ax.set_xticks(list(np.unique(x)))
ax.set_yticks(list(np.unique(y)))
ax.set_xlim(df["Timeline"].min()-0.5, df["Timeline"].max()+0.5)
ax.grid()

# Colors:
colors = {'Track 1': 'tab:red', 'Track 2': 'tab:blue', 'Track 3': 'blue'}

# Drawing points and lines according to positions:
frames = (len(df.groupby(['Timeline'])))

steps = []
scatters = []
for track, group in df.groupby("Track"):
    scatters.append(ax.scatter(group["Timeline"].iloc[0],
                               group["Position"].iloc[0],
                               s=45, c=colors[track]))
    steps.append(ax.step(group["Timeline"].iloc[0],
                     group["Position"].iloc[0],
                     color=colors[track], alpha=0.5, linewidth=3)[0])

def animate(i):
    for scatter, (_, group) in zip(scatters, df.groupby("Track")):
        scatter.set_offsets((group["Timeline"].iloc[i],
                             group["Position"].iloc[i]))
    for step, (_, group) in zip(steps, df.groupby('Track')):
        step.set_data(group['Timeline'].iloc[:i+1],
                       group['Position'].iloc[:i+1])

anim = FuncAnimation(fig, animate, frames=frames, interval=400)
anim.save('test.gif', writer='pillow')

我还加厚了线条,并添加了一些不透明度,使其更清晰。

相关问题