Animate contour and scatter plot

Question

I am trying to animate a scatter and bivariate gaussian distribution from a set of xy coordinates. I'll record the specific code that calls the scatter and distribution first and then how I measure the distribution afterwards.

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import scipy.stats as sts
import matplotlib.animation as animation

''' Below is a section of the script that generates the scatter and contour '''

fig, ax = plt.subplots(figsize = (10,4))

def plotmvs(df, xlim=None, ylim=None, fig=fig, ax=ax):

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

    PDFs = []
    for (group,gdf),color in zip(df.groupby('group'), ('red', 'blue')):

        ax.plot(*gdf[['X','Y']].values.T, '.', c=color, alpha = 0.5)

        kwargs = {
            'xlim': xlim,
            'ylim': ylim
        }
        X, Y, PDF = mvpdfs(gdf['X'].values, gdf['Y'].values, **kwargs)
        PDFs.append(PDF)

    PDF = PDFs[0] - PDFs[1]

    normPDF = PDF - PDF.min()
    normPDF = normPDF/normPDF.max()

    cfs = ax.contourf(X, Y, normPDF, levels=100, cmap='jet')

    return fig, ax

n = 10
time = [1]
d = ({      
    'A1_Y' : [10,20,15,20,25,40,50,60,61,65],                 
    'A1_X' : [15,10,15,20,25,25,30,40,60,61], 
    'A2_Y' : [10,13,17,10,20,24,29,30,33,40],                 
    'A2_X' : [10,13,15,17,18,19,20,21,26,30],
    'A3_Y' : [11,12,15,17,19,20,22,25,27,30],                 
    'A3_X' : [15,18,20,21,22,28,30,32,35,40], 
    'A4_Y' : [15,20,15,20,25,40,50,60,61,65],   
    'A4_X' : [16,20,15,30,45,30,40,10,11,15],                 
    'B1_Y' : [18,10,11,13,18,10,30,40,31,45],                 
    'B1_X' : [17,20,15,10,25,20,10,12,14,25], 
    'B2_Y' : [13,10,14,20,21,12,30,20,11,35],                 
    'B2_X' : [12,20,16,22,15,20,10,20,16,15],
    'B3_Y' : [15,20,15,20,25,10,20,10,15,25],                 
    'B3_X' : [18,15,13,20,21,10,20,10,11,15], 
    'B4_Y' : [19,12,15,18,14,19,13,12,11,18],   
    'B4_X' : [20,10,12,18,17,15,13,14,19,13],                                                                                    
     })        


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']

for time,tdf in df.groupby('time'):
    plotmvs(tdf)


'''MY ATTEMPT AT ANIMATING THE PLOT '''

def animate(i) :
    tdf.set_offsets([[tdf.iloc[0:,1][0+i][0], tdf.iloc[0:,0][0+i][0]], [tdf.iloc[0:,1][0+i][1], tdf.iloc[0:,0][0+i][1]], [tdf.iloc[0:,1][0+i][2], tdf.iloc[0:,0][0+i][2]], [tdf.iloc[0:,1][0+i][3], tdf.iloc[0:,0][0+i][3]], [tdf.iloc[0:,1][0+i][4], tdf.iloc[0:,0][0+i][4]]])
    normPDF = n[i,:,0,:].T
    cfs.set_data(X, Y, normPDF)

ani = animation.FuncAnimation(fig, animate, np.arange(0,10),# init_func = init,
                              interval = 10, blit = False)

A full working code on how the distribution is generated and plotted using a single frame

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import scipy.stats as sts
import matplotlib.animation as animation

def datalimits(*data, pad=.15):
    dmin,dmax = min(d.min() for d in data), max(d.max() for d in data)
    spad = pad*(dmax - dmin)
    return dmin - spad, dmax + spad

def rot(theta):
    theta = np.deg2rad(theta)
    return np.array([
        [np.cos(theta), -np.sin(theta)],
        [np.sin(theta), np.cos(theta)]
    ])

def getcov(radius=1, scale=1, theta=0):
    cov = np.array([
        [radius*(scale + 1), 0],
        [0, radius/(scale + 1)]
    ])

    r = rot(theta)
    return r @ cov @ r.T

def mvpdf(x, y, xlim, ylim, radius=1, velocity=0, scale=0, theta=0):

    X,Y = np.meshgrid(np.linspace(*xlim), np.linspace(*ylim))

    XY = np.stack([X, Y], 2)

    x,y = rot(theta) @ (velocity/2, 0) + (x, y)

    cov = getcov(radius=radius, scale=scale, theta=theta)

    PDF = sts.multivariate_normal([x, y], cov).pdf(XY)

    return X, Y, PDF

def mvpdfs(xs, ys, xlim, ylim, radius=None, velocity=None, scale=None, theta=None):
    PDFs = []
    for i,(x,y) in enumerate(zip(xs,ys)):
        kwargs = {
            'xlim': xlim,
            'ylim': ylim
        }
        X, Y, PDF = mvpdf(x, y,**kwargs)
        PDFs.append(PDF)

    return X, Y, np.sum(PDFs, axis=0)

fig, ax = plt.subplots(figsize = (10,4))

def plotmvs(df, xlim=None, ylim=None, fig=fig, ax=ax):

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

    PDFs = []
    for (group,gdf),color in zip(df.groupby('group'), ('red', 'blue')):

        #Animate this scatter
        ax.plot(*gdf[['X','Y']].values.T, '.', c=color, alpha = 0.5)

        kwargs = {
            'xlim': xlim,
            'ylim': ylim
        }
        X, Y, PDF = mvpdfs(gdf['X'].values, gdf['Y'].values, **kwargs)
        PDFs.append(PDF)

    PDF = PDFs[0] - PDFs[1]

    normPDF = PDF - PDF.min()
    normPDF = normPDF/normPDF.max()

    #Animate this contour
    cfs = ax.contourf(X, Y, normPDF, levels=100, cmap='jet')

    return fig, ax

n = 10

time = [1]
d = ({      
    'A1_Y' : [10,20,15,20,25,40,50,60,61,65],                 
    'A1_X' : [15,10,15,20,25,25,30,40,60,61], 
    'A2_Y' : [10,13,17,10,20,24,29,30,33,40],                 
    'A2_X' : [10,13,15,17,18,19,20,21,26,30],
    'A3_Y' : [11,12,15,17,19,20,22,25,27,30],                 
    'A3_X' : [15,18,20,21,22,28,30,32,35,40], 
    'A4_Y' : [15,20,15,20,25,40,50,60,61,65],   
    'A4_X' : [16,20,15,30,45,30,40,10,11,15],                 
    'B1_Y' : [18,10,11,13,18,10,30,40,31,45],                 
    'B1_X' : [17,20,15,10,25,20,10,12,14,25], 
    'B2_Y' : [13,10,14,20,21,12,30,20,11,35],                 
    'B2_X' : [12,20,16,22,15,20,10,20,16,15],
    'B3_Y' : [15,20,15,20,25,10,20,10,15,25],                 
    'B3_X' : [18,15,13,20,21,10,20,10,11,15], 
    'B4_Y' : [19,12,15,18,14,19,13,12,11,18],   
    'B4_X' : [20,10,12,18,17,15,13,14,19,13],                                                                                   
     })        

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']

for time,tdf in df.groupby('time'):
    plotmvs(tdf)

I essentially want to animate this code by iterating over each row of xy coordinates.

Answer 1

Here's a very quick and dirty modification of the OP's code, fixing the scatter animation and adding (a form of) contour animation.

Basically, you start by creating the artists for your animation (in this case Line2D objects, as returned by plot()). Subsequently, you create an update function (and, optionally, an initialization function). In that function, you update the existing artists. I think the example in the matplotlib docs explains it all.

In this case, I modified the OP's plotmvs function to be used as the update function (instead of the OP's proposed animate function).

The QuadContourSet returned by contourf (i.e. your cfs) cannot be used as an artist in itself, but you can make it work using cfs.collections (props to this SO answer). However, you still need to create a new contour plot and remove the old one, instead of just updating the contour data. Personally I would prefer a lower level approach: try to get the contour-data without calling contourf, then initialize and update the contour lines just like you do for the scatter.

Nevertheless, the approach above is implemented in the OP's code below (just copy, paste, and run):

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

# quick and dirty override of datalimits(), to get a fixed contour-plot size
DATA_LIMITS = [0, 70]

def datalimits(*data, pad=.15):
    # dmin,dmax = min(d.min() for d in data), max(d.max() for d in data)
    # spad = pad*(dmax - dmin)
    return DATA_LIMITS  # dmin - spad, dmax + spad

def rot(theta):
    theta = np.deg2rad(theta)
    return np.array([
        [np.cos(theta), -np.sin(theta)],
        [np.sin(theta), np.cos(theta)]
    ])

def getcov(radius=1, scale=1, theta=0):
    cov = np.array([
        [radius*(scale + 1), 0],
        [0, radius/(scale + 1)]
    ])

    r = rot(theta)
    return r @ cov @ r.T

def mvpdf(x, y, xlim, ylim, radius=1, velocity=0, scale=0, theta=0):

    X,Y = np.meshgrid(np.linspace(*xlim), np.linspace(*ylim))

    XY = np.stack([X, Y], 2)

    x,y = rot(theta) @ (velocity/2, 0) + (x, y)

    cov = getcov(radius=radius, scale=scale, theta=theta)

    PDF = sts.multivariate_normal([x, y], cov).pdf(XY)

    return X, Y, PDF

def mvpdfs(xs, ys, xlim, ylim, radius=None, velocity=None, scale=None, theta=None):
    PDFs = []
    for i,(x,y) in enumerate(zip(xs,ys)):
        kwargs = {
            'xlim': xlim,
            'ylim': ylim
        }
        X, Y, PDF = mvpdf(x, y,**kwargs)
        PDFs.append(PDF)

    return X, Y, np.sum(PDFs, axis=0)


fig, ax = plt.subplots(figsize = (10,4))
ax.set_xlim(DATA_LIMITS)
ax.set_ylim(DATA_LIMITS)

# Initialize empty lines for the scatter (increased marker size to make them more visible)
line_a, = ax.plot([], [], '.', c='red', alpha = 0.5, markersize=20, animated=True)
line_b, = ax.plot([], [], '.', c='blue', alpha = 0.5, markersize=20, animated=True)
cfs = None

# Modify the plotmvs function so it updates the lines 
# (might as well rename the function to "update")
def plotmvs(tdf, xlim=None, ylim=None):
    global cfs  # as noted: quick and dirty...
    if cfs:
        for tp in cfs.collections:
            # Remove the existing contours
            tp.remove()

    # Get the data frame for time t
    df = tdf[1]

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

    PDFs = []

    for (group, gdf), group_line in zip(df.groupby('group'), (line_a, line_b)):

        #Animate this scatter
        #ax.plot(*gdf[['X','Y']].values.T, '.', c=color, alpha = 0.5)

        # Update the scatter line data
        group_line.set_data(*gdf[['X','Y']].values.T)

        kwargs = {
            'xlim': xlim,
            'ylim': ylim
        }
        X, Y, PDF = mvpdfs(gdf['X'].values, gdf['Y'].values, **kwargs)
        PDFs.append(PDF)


    PDF = PDFs[0] - PDFs[1]

    normPDF = PDF - PDF.min()
    normPDF = normPDF / normPDF.max()

    # Plot a new contour
    cfs = ax.contourf(X, Y, normPDF, levels=100, cmap='jet')

    # Return the artists (the trick is to return cfs.collections instead of cfs)
    return cfs.collections + [line_a, line_b]

n = 10
time = range(n)  # assuming n represents the length of the time vector...
d = ({
    'A1_Y' : [10,20,15,20,25,40,50,60,61,65],
    'A1_X' : [15,10,15,20,25,25,30,40,60,61],
    'A2_Y' : [10,13,17,10,20,24,29,30,33,40],
    'A2_X' : [10,13,15,17,18,19,20,21,26,30],
    'A3_Y' : [11,12,15,17,19,20,22,25,27,30],
    'A3_X' : [15,18,20,21,22,28,30,32,35,40],
    'A4_Y' : [15,20,15,20,25,40,50,60,61,65],
    'A4_X' : [16,20,15,30,45,30,40,10,11,15],
    'B1_Y' : [18,10,11,13,18,10,30,40,31,45],
    'B1_X' : [17,20,15,10,25,20,10,12,14,25],
    'B2_Y' : [13,10,14,20,21,12,30,20,11,35],
    'B2_X' : [12,20,16,22,15,20,10,20,16,15],
    'B3_Y' : [15,20,15,20,25,10,20,10,15,25],
    'B3_X' : [18,15,13,20,21,10,20,10,11,15],
    'B4_Y' : [19,12,15,18,14,19,13,12,11,18],
    'B4_X' : [20,10,12,18,17,15,13,14,19,13],
     })

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']

# Use the modified plotmvs as the update function, and supply the data frames
interval_ms = 200
delay_ms = 1000
ani = FuncAnimation(fig, plotmvs, frames=df.groupby('time'),
                    blit=True, interval=interval_ms, repeat_delay=delay_ms)

# Start the animation
plt.show()