Bobasheto
Bobasheto

Reputation: 27

Making fading trails in python FuncAnimation

I have the following code that uses 4th order Runge Kutta to solve the equations of motion of a double pendulum. I am attempting to animate the positions of the two pendulums using FuncAnimation. I tried removing the 0th element of each coordinate list after every 10th new element is plotted. I thought this would create a fading trail effect. Instead it just plotted one point. (note x1 and y1 are the coordinates of the first pendulum, and x2 and y2 are the coordinates of the second coupled pendulum) I feel like I'm over complicating things. when I remove the if i % 10 ==1: loop the code animates fine, but it does not delete any previous points. I feel this should be fairly easy to do but I could not find any straight forward examples, and the documentation was not extremely helpful.

In short I have 4 lists of coordinates that describe the position of two objects. I want to animate the current position of these two objects and not any previous positions.

import matplotlib.pyplot as plt
import numpy as np
import matplotlib.animation as animation
from matplotlib.animation import FuncAnimation 
from itertools import count
from IPython import display

h=.001
t=0
l1=1
l2=1
m1=1
m2=1
theta1= 0
theta2 = np.pi/2
thetadot1=0
thetadot2=0
g=9.8
x1=[]
y1=[]
x2=[]
y2=[]
E=[]
T=[]
Theta_dot1=[]
Theta_dot2=[]
kinetic_energy=[]
potential_energy=[]



    
while(t<4):

    KE = .5*m1*((thetadot1*l1*np.cos(theta1))**2+(thetadot1*l1*np.sin(theta1))**2) + .5*m2*((thetadot1*l1*np.cos(theta1)+thetadot2*l2*np.cos(theta2))**2+(thetadot1*l1*np.sin(theta1)+thetadot2*l2*np.sin(theta2))**2)
    PE = (m1*g*(l1-l1*np.cos(theta1))+m2*g*(l1+l2-l1*np.cos(theta1)-l2*np.cos(theta2))) 
    E.append(KE+PE)
    kinetic_energy.append(KE)
    potential_energy.append(PE)
    Theta_dot1.append(thetadot1)
    Theta_dot2.append(thetadot2)
        
    T.append(t)
    x1.append(l1*np.sin(theta1))
    y1.append(-l1*np.cos(theta1))
    x2.append(l1*np.sin(theta1)+l2*np.sin(theta2))
    y2.append(-l1*np.cos(theta1)-l2*np.cos(theta2))

    P1 = (-g*(2*m1+m2)*np.sin(theta1)-m2*g*np.sin(theta1-2*theta2)-2*np.sin(theta1-theta2)*m2*(l2*thetadot2**2+l1*thetadot1**2*(np.cos(theta1-theta2))))/(l1*(2*m1+m2*(1-np.cos(2*(theta1-theta2)))))
    L1 = (2*np.sin(theta1-theta2)*(l1*thetadot1**2*(m1+m2)+np.cos(theta1)*g*(m1+m2)+l2*m2*thetadot2**2*np.cos(theta1-theta2)) )/(l2*(2*m1+m2*(1-np.cos(2*(theta1-theta2)))))
    G1 = thetadot1
    H1 = thetadot2
    
    

    thetadot1_1     = thetadot1     + (h/2)*P1
    thetadot2_1     = thetadot2     + (h/2)*L1
    theta1_1        = theta1        + (h/2)*G1
    theta2_1        = theta2        + (h/2)*H1
 
    
    
    P2 = (-g*(2*m1+m2)*np.sin(theta1_1)-m2*g*np.sin(theta1_1-2*theta2_1)-2*np.sin(theta1_1-theta2_1)*m2*(l2*thetadot2_1**2+l1*thetadot1_1**2*(np.cos(theta1_1-theta2_1))))/(l1*(2*m1+m2*(1-np.cos(2*(theta1_1-theta2_1)))))
    L2 = (2*np.sin(theta1_1-theta2_1)*(l1*thetadot1_1**2*(m1+m2)+np.cos(theta1_1)*g*(m1+m2)+l2*m2*thetadot2_1**2*np.cos(theta1_1-theta2_1)) )/(l2*(2*m1+m2*(1-np.cos(2*(theta1_1-theta2_1)))))
    G2 = thetadot1_1
    H2 = thetadot2_1
    
    thetadot1_2     = thetadot1_1     + (h/2)*P2
    thetadot2_2     = thetadot2_1     + (h/2)*L2
    theta1_2        = theta1_1        + (h/2)*G2
    theta2_2        = theta2_1        + (h/2)*H2
    
    P3 =  (-g*(2*m1+m2)*np.sin(theta1_2)-m2*g*np.sin(theta1_2-2*theta2_2)-2*np.sin(theta1_2-theta2_2)*m2*(l2*thetadot2_2**2+l1*thetadot1_2**2*(np.cos(theta1_2-theta2_2))))/(l1*(2*m1+m2*(1-np.cos(2*(theta1_2-theta2_2)))))
    L3 = (2*np.sin(theta1_2-theta2_2)*(l1*thetadot1_2**2*(m1+m2)+np.cos(theta1_2)*g*(m1+m2)+l2*m2*thetadot2_2**2*np.cos(theta1_2-theta2_2)) )/(l2*(2*m1+m2*(1-np.cos(2*(theta1_2-theta2_2)))))
    G3 = thetadot1_2
    H3 = thetadot2_2
    
    thetadot1_3     = thetadot1_2     + (h/2)*P3
    thetadot2_3     = thetadot2_2     + (h/2)*L3
    theta1_3        = theta1_2        + (h/2)*G3
    theta2_3        = theta2_2        + (h/2)*H3
    
    
    P4 = (-g*(2*m1+m2)*np.sin(theta1_3)-m2*g*np.sin(theta1_3-2*theta2_3)-2*np.sin(theta1_3-theta2_3)*m2*(l2*thetadot2_3**2+l1*thetadot1_3**2*(np.cos(theta1_3-theta2_3))))/(l1*(2*m1+m2*(1-np.cos(2*(theta1_3-theta2_3)))))
    L4 = (2*np.sin(theta1_3-theta2_3)*(l1*thetadot1_3**2*(m1+m2)+np.cos(theta1_3)*g*(m1+m2)+l2*m2*thetadot2_3**2*np.cos(theta1_3-theta2_3)) )/(l2*(2*m1+m2*(1-np.cos(2*(theta1_3-theta2_3))))) 
    G4 = thetadot1_3
    H4 = thetadot2_3
    
    
    
 
    thetadot1   = thetadot1  + (h/6.0) * (P1+(2.*P2)+(2.0*P3) + P4)
    thetadot2   = thetadot2  + (h/6.0) * (L1+(2.*L2)+(2.0*L3) + L4)
    theta1      = theta1     + (h/6.0) * (G1+(2.*G2)+(2.0*G3) + G4)
    theta2      = theta2     + (h/6.0) * (H1+(2.*H2)+(2.0*H3) + H4)
    t=t+h 
 
    





fig, axes = plt.subplots(nrows = 1, ncols = 1, figsize = (15,5))
axes.set_ylim(-2.5, 2.5)
axes.set_xlim(-2.5, 2.5)
plt.style.use("ggplot")


x_1,y_1,x_2,y_2 = [], [], [], []
def animate(i):
    
    
    
    if i % 10 ==1:     #this is the stuff that is not working
        x_1.remove(0) 
        y_1.remove(0)
        x_2.remove(0)
        y_2.remove(0)
   
    x_1.append(x1[i*10])
    y_1.append(y1[i*10])
    x_2.append(x2[i*10])
    y_2.append(y2[i*10])
    

    
    axes.plot(x_1,y_1,'.', color="red")
    axes.plot(x_2,y_2,'.', color="gray", linewidth=0.5)
    
    
    
anim = FuncAnimation(fig, animate, interval=.1)

Upvotes: 0

Views: 388

Answers (1)

Davide_sd
Davide_sd

Reputation: 13185

If you want fading lines you have to deal with the extra complexity of LineCollection, which is not easy to master.

import matplotlib.colors as colors
from matplotlib.collections import LineCollection
from matplotlib.lines import Line2D

n_points_to_render = 500
# opacity of the segments
alphas = np.linspace(0, 1, n_points_to_render)

# create "solid" color maps with a varying opacity
red = colors.to_rgb("red") + (0.0,)
redfade = colors.LinearSegmentedColormap.from_list('my', [red, "red"])
green = colors.to_rgb("green") + (0.0,)
greenfade = colors.LinearSegmentedColormap.from_list('my', [green, "green"])

def get_segments(i):
    # LineCollection requires segments
    _x1 = x1[i:i+n_points_to_render]
    _y1 = y1[i:i+n_points_to_render]
    _x2 = x2[i:i+n_points_to_render]
    _y2 = y2[i:i+n_points_to_render]
    points1 = np.vstack((_x1, _y1)).T.reshape(-1, 1, 2)
    segments1 = np.hstack((points1[:-1], points1[1:]))
    points2 = np.vstack((_x2, _y2)).T.reshape(-1, 1, 2)
    segments2 = np.hstack((points2[:-1], points2[1:]))
    return segments1, segments2

fig, ax = plt.subplots(nrows = 1, ncols = 1)
ax.set_ylim(-2.5, 2.5)
ax.set_xlim(-2.5, 2.5)
plt.style.use("ggplot")

# create and add two LineCollections
segments1, segments2 = get_segments(0)
lc1 = LineCollection(segments1, array=alphas, cmap=redfade, lw=2)
lc2 = LineCollection(segments2, array=alphas, cmap=greenfade, lw=2)
line1 = ax.add_collection(lc1)
line2 = ax.add_collection(lc2)

def animate(i):
    segments1, segments2 = get_segments(i)
    line1.set_segments(segments1)
    line2.set_segments(segments2)

# create a legend as LineCollection doesn't have any by default
l1_legend = Line2D([0, 1], [0, 1], color="r", linewidth=2)
l2_legend = Line2D([0, 1], [0, 1], color="g", linewidth=2)
ax.legend([l1_legend, l2_legend], ['Line 1', 'Line 2'])

anim = FuncAnimation(fig, animate, frames=len(x1) - n_points_to_render, interval=30)
plt.show()

Upvotes: 1

Related Questions