Friday, 24 May 2019

Plot multiple subplots as animations

I have two separate subplots that I'm hoping to display as animations. For the subplots below, ax1 displays an animated scatter plot, while ax2 is a scatter now, I'm hoping to alter this to a line plot.

Please note: I've simplified the question to only display relevant info. However I'm hoping to keep the code similar to what it is now.

Below is my attempt:

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

DATA_LIMITS = [0, 15]

def datalimits(*data):
    return DATA_LIMITS 

fig = plt.figure(figsize=(10,18))
grid = plt.GridSpec(1, 3, wspace=0.4, hspace=0.3)

gridsize = (3, 2)
ax1 = plt.subplot2grid(gridsize, (0, 0), colspan=2, rowspan=2)
ax2 = plt.subplot2grid(gridsize, (2, 0), colspan=2, rowspan=2)
ax1.grid(False)
ax2.grid(False)

ax1.set_xlim(DATA_LIMITS)
ax1.set_ylim(DATA_LIMITS)

line_a, = ax1.plot([], [], 'o', c='red', alpha = 0.5, markersize=5,zorder=3)
line_b, = ax1.plot([], [], 'o', c='blue', alpha = 0.5, markersize=5,zorder=3)
lines=[line_a,line_b] 

scat = ax1.scatter([], [], s=20, marker='o', c='white', alpha = 1,zorder=3)
scats=[scat] 

line_d = ax2.plot([], [], 'o', c = 'k')

ax2.set_ylim(-6,6) 
ax2.set_xlim(0,15) 

def plots(tdf, xlim=None, ylim=None, fig=fig, ax=ax1):

    df = tdf[1]

    if xlim is None: xlim = datalimits(df['X'])
    if ylim is None: ylim = datalimits(df['Y'])

    for (group, gdf), group_line in zip(df.groupby('group'), lines+scats+line_d):
        if group in ['A','B','D']:
            group_line.set_data(*gdf[['X','Y']].values.T)
        elif group in ['C']:
            gdf['X'].values, gdf['Y'].values
            scat.set_offsets(gdf[['X','Y']].values)

    return [scat] + [line_a,line_b] + [line_d]          

n = 9
time = range(n)  

d = ({
     'A1_X' :    [13,14,12,13,11,12,13,12,11,10],
     'A1_Y' :    [6,6,7,7,7,8,8,8,9,10],
     'A2_X' :    [7,6,5,7,6,3,4,5,6,6],
     'A2_Y' :    [11,12,11,10,11,12,10,11,10,9],
     'B1_X' :    [8,9,8,7,6,7,5,6,7,6],
     'B1_Y' :    [3,4,3,2,3,4,2,1,2,3],
     'B2_X' :    [13,14,14,14,13,13,13,12,12,12],
     'B2_Y' :    [5,4,3,2,4,5,4,6,3,3],
     'C1_X' :   [5,6,7,5,6,5,6,5,6,5],
     'C1_Y' :   [10,11,10,11,12,11,10,8,7,6],
     'D1_X' :   [0,1,2,3,4,5,6,7,8,9],           
     'D1_Y' :   [0,1,2,3,4,3,2,1,0,-1],                
    })

tuples = [((t, k.split('_')[0][0], int(k.split('_')[0][1:]), k.split('_')[1]), v[i])
    for k,v in d.items() for i,t in enumerate(time) ]

df = pd.Series(dict(tuples)).unstack(-1)
df.index.names = ['time', 'group', 'id']

interval_ms = 1000
delay_ms = 2000
ani = animation.FuncAnimation(fig, plots, frames=df.groupby('time'), interval=interval_ms, repeat_delay=delay_ms,)

plt.show()



from Plot multiple subplots as animations

No comments:

Post a Comment