Animate changing amount of lines between points - Python

Question

This question is different from other animation questions as I'm trying to animate an alternating amounts of lines between points. For example, it may be between 3 points or 50 points.

Using the data frame below, the points are labelled in Item. The first two time stamps contain 4 points but this drops to 3 points for the final two time stamps. I'm trying to find an efficient way to combine all the potential lines act each time stamp into a single call function to animate.

The issue I'm having is I'm plotting each line manually. Therefore, the line between each point is currently hard-coded, which doesn't account for a change in the amount of lines.

I need something that combines the available lines first and then passes this to the animation.

For instance, A, B, C, D are currently labelled points in the first two time points. But this drops to A, B, C for the last two time points.

This following doesn't account for alternating amounts of lines.

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

df1 = pd.DataFrame({
    'Time' : [1,1,1,1,2,2,2,2,3,3,3,4,4,4],  
    'Item' : ['A', 'B', 'C', 'D','A', 'B', 'C', 'D', 'A', 'B', 'C', 'A', 'B', 'C'],
    'GroupA_X' : [3, 4, 5, 1, 2, 5, 6, 2, 1, 6, 7, 2, 7, 8], 
    'GroupA_Y' : [2, 4, 5, 1, 2, 5, 5, 2, 2, 6, 5, 1, 5, 4], 
})

GA_X = np.array(df.groupby('Time')['GroupA_X'].apply(list).tolist())
GA_Y = np.array(df.groupby('Time')['GroupA_Y'].apply(list).tolist())

fig, ax = plt.subplots(figsize = (6,6))
ax.grid(False)
ax.set_xlim(0,10)
ax.set_ylim(0,10)

data = np.stack([GA_X, GA_Y], axis = 2)

vector1 = ax.annotate('', xy = data[0][0], 
            xytext = data[0][1], 
            arrowprops={'arrowstyle': "-", 'color': 'black'}, 
            ha='center')

vector2 = ax.annotate('', xy = data[0][0], 
            xytext = data[0][2], 
            arrowprops={'arrowstyle': "-", 'color': 'black'}, 
            ha='center')

vector3 = ax.annotate('', xy = data[0][1], 
            xytext = data[0][2], 
            arrowprops={'arrowstyle': "-", 'color': 'black'}, 
            ha='center')

def animate(i):
    start1 = np.r_[data[i, 0]]
    end1 = np.r_[data[i, 1]]

    vector1.set_position(start1)
    vector1.xy = end1    

    start2 = np.r_[data[i, 0]]
    end2 = np.r_[data[i, 2]]

    vector2.set_position(start2)
    vector2.xy = end2 

    start3 = np.r_[data[i, 1]]
    end3 = np.r_[data[i, 2]]

    vector3.set_position(start3)
    vector3.xy = end3 

    return 

ani = animation.FuncAnimation(fig, animate, interval = 100, blit = False)

Out:

data = np.stack([GA_X, GA_Y], axis = 2)

axis = normalize_axis_index(axis, result_ndim)

AxisError: axis 2 is out of bounds for array of dimension 2

Answer 1

From my understanding you want to plot at every time-step i a line for each item where the two points of each line are first the (x,y) in the time-step i and second the point (x,y) in the time step i+1.

(so if the item don't appear in the time-step i+1 we won't display the line for that item in step i)

Assuming this I suggest that:

1) use the dataframe itself instead of changing it to a np.array

2) move the creation of the lines inside the animate function

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


df1 = pd.DataFrame({
    'Time' : [1,1,1,1,2,2,2,2,3,3,3,4,4,4],  
    'Item' : ['A', 'B', 'C', 'D','A', 'B', 'C', 'D', 'A', 'B', 'C', 'A', 'B', 'C'],                  
    'GroupA_X' : [3, 4, 5, 1, 2, 5, 6, 2, 1, 6, 7, 2, 7, 8], 
    'GroupA_Y' : [2, 4, 5, 1, 2, 5, 5, 2, 2, 6, 5, 1, 5, 4],                         
        })

# generating the figure configs
frame_num = len(df1['Time'].unique()) 
fig = plt.figure()
ax1 = plt.axes(xlim=(0, 10), ylim=(0, 10))
line, = ax1.plot([], [], lw=2)
plt.xlabel('X')
plt.ylabel('Y')

plotlays, itemColors = [2], {'A':"black",
                             'B':"red",
                             'C':"blue",
                             'D':"purple"}


def animate(i):

    lines = []
    # filtering items per time step
    for (row_idx, row) in df1[df1["Time"] == i+1].iterrows():
        nextTimestep_item = df1[(df1["Time"] == i+2) & (df1["Item"] ==row["Item"])] 
        if nextTimestep_item.empty:
            # no plot if item is not on next timestep
            continue
        else:
            x = [row['GroupA_X'],nextTimestep_item['GroupA_X']]
            y = [row['GroupA_Y'],nextTimestep_item['GroupA_Y']]
            color = itemColors[row["Item"]]
            lobj = ax1.plot([],[],lw=2,color=color)[0]
            lobj.set_data(x, y)  
            lines.append(lobj)

    return lines

# call the animator.  blit=True means only re-draw the parts that have changed.
anim = animation.FuncAnimation(fig, animate, 
                               frames=frame_num, interval=500, blit=True)

plt.show()