Python FuncAnimation多行星椭圆轨道动画异常及颜色匹配问题

基础背景 #

已有的matplotlib多圆点圆周运动动画实现方案，无法适配当前“用小点代表行星，沿椭圆轨道绕中心原点运动”的目标需求。

初始实现问题 #

• 初始运行效果表现：仅最外层的行星沿轨道运动，内侧3个行星全部静止
• 问题根因：FuncAnimation实例被写在for j in range(1, planet_count):循环内部，最后一次循环生成的动画实例覆盖了之前所有循环的实例

planet_count = 5

df = pd.read_csv("Wincoll_prog_chall_1/data2.csv")

df["b"] = df["a"]*(1-df["Epsilon"]**2)


def position(theta, a, eps):
    r = a*(1-eps**2)/(1-eps*np.cos(theta))
    x = r*np.cos(theta)
    y = r*np.sin(theta)

    return r, x, y


num_to_color = {
    1: "orange",
    2: "pink",
    3: "green",
    4: "red",
    5: "brown",
    6: "yellow",
    7: "blue",
    8: "purple",
    9: "grey",
}



def animate(i, a, eps):
    theta = i*math.pi/180
    r = a*(1-eps**2)/(1-eps*np.cos(theta))
    x = r*np.cos(theta)
    y = r*np.sin(theta)
    planet.set_data(x, y)
    return planet,


plt.style.use('dark_background')
fig, ax = plt.subplots(figsize=(10, 10))


if planet_count > 5:
    ax.plot(0, 0, 'X', markersize=5, color="yellow")
else:
    ax.plot(0, 0, 'X', markersize=15, color="yellow")

# planet = []

for j in range(1, planet_count):
    theta = np.linspace(0, 2*math.pi, 1000)
    r, x, y = position(theta, df["a"][j], df["Epsilon"][j])
    ax.plot(x, y, color=num_to_color[j], linewidth=2, label=df["Object"][j])

    planet, = ax.plot(x[0], y[0], '.', markersize=15, color="white")
    
    # Animation
    animation = FuncAnimation(fig, animate, frames=360, interval=5, fargs=(
        df["a"][j], df["Epsilon"][j],), blit=True)

plt.xlabel("x", size=20)
plt.ylabel("y", size=20)
plt.legend(fontsize=12, loc="upper right", numpoints=1)
plt.grid(True, lw=0.3)
plt.show()

初步修正方案 #

• 修正思路：新增dataX、dataY两个数组，每次动画帧更新时计算所有行星的坐标存入数组后统一下发更新
• 修正后效果：所有行星可正常沿轨道完成动画，当前待实现的效果为每个行星点的颜色与对应轨道颜色匹配

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

AU = 149597870700
P = 5

df = pd.read_csv("Wincoll_prog_chall_1/data2.csv")

df["b"] = df["a"]*(1-df["Epsilon"]**2)


def position(theta, a, eps):
    r = a*(1-eps**2)/(1-eps*np.cos(theta))
    x = r*np.cos(theta)
    y = r*np.sin(theta)

    return r, x, y


num_to_color = {
    1: "orange",
    2: "pink",
    3: "green",
    4: "red",
    5: "brown",
    6: "yellow",
    7: "blue",
    8: "purple",
    9: "grey",
}

starting_planet = 1
planet_count = 5

if (planet_count < 6):
    anim_speed = 2.125
else:
    anim_speed = 50


def animate(i):
    dataX = []
    dataY = []

    years = round(1000*i*anim_speed/360)/1000

    for j in range(starting_planet, planet_count):
        theta = (i*anim_speed*np.pi)/(df["P"][j]*180)
        # theta = i*1/df["P"][j]*np.pi/180
        r = df["a"][j]*(1-df["Epsilon"][j]**2) / \
            (1-df["Epsilon"][j]*np.cos(theta))
        x = r*np.cos(theta)
        y = r*np.sin(theta)
        dataX.append(x)
        dataY.append(y)
        # planet.set_color(num_to_color[j])
    planet.set_data(dataX, dataY)
    title.set_text(u"Years: {}".format(years))
    return planet, title,


plt.style.use('dark_background')
fig, ax = plt.subplots(figsize=(10, 10))
# planet, = ax.plot([], [], '.')


if planet_count > 5:
    ax.plot(0, 0, 'X', markersize=5, color="yellow")
else:
    ax.plot(0, 0, 'X', markersize=15, color="yellow")


planet, = ax.plot([], [], '.', markersize=11)


title = ax.text(0.5, 0.97, "", bbox={'facecolor': 'w', 'alpha': 0.0, 'pad': 5},
                transform=ax.transAxes, ha="center")


dataX = []
dataY = []

for j in range(starting_planet, planet_count):
    theta = np.linspace(0, 2*math.pi, 1000)
    r, x, y = position(theta, df["a"][j], df["Epsilon"][j])
    ax.plot(x, y, color=num_to_color[j], linewidth=2, label=df["Object"][j])

    # planet.set_color(num_to_color[j])
    dataX.append(x[0])
    dataY.append(y[0])
animation = FuncAnimation(fig, animate, interval=1, blit=True)


planet.set_data(dataX, dataY)

plt.xlabel("x", size=20)
plt.ylabel("y", size=20)
plt.legend(fontsize=12, loc="upper right", numpoints=1)
plt.grid(True, lw=0.3)
plt.show()

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