Animate scatter and contour plot

Question

I am trying to generate an animated scatter plot within a functioning animated contour plot. I can get both working separately but not together.

The code below generates the contour from coordinates A and B in the df. I've attempted to include a separated animated scatter within the sample plot using the C coordinates. This attempt is currently commented out.

So I basically want to include another animated scatter using C_X and C_Y. I've trying by applying them to line_c.

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

''' Section 1 '''

DATA_LIMITS = [-85, 85]

def datalimits(*data):
    return DATA_LIMITS  # dmin - spad, dmax + spad

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)
    PDF = sts.multivariate_normal([x, y]).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)):
        X, Y, PDF = mvpdf(x, y, xlim, ylim)
        PDFs.append(PDF)

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

''' Animate Plot '''

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

#Animated coordinates for group A,B
line_a, = ax.plot([], [], '.', c='red', alpha = 0.5, markersize=5, animated=True)
line_b, = ax.plot([], [], '.', c='blue', alpha = 0.5, markersize=5, animated=True)

#Attempt to incorporate scatter for C 
line_c, = ax.plot([], [], '.', c='white', alpha = 0.5, markersize=2.5, animated=True)

cfs = None

def plotmvs(tdf, xlim=None, ylim=None, fig=fig, ax=ax):
    global cfs  
    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)):
        group_line.set_data(*gdf[['X','Y']].values.T)
        X, Y, PDF = mvpdfs(gdf['X'].values, gdf['Y'].values, xlim, ylim)
        PDFs.append(PDF)

    normPDF = PDF - PDF.min()
    normPDF = normPDF / normPDF.max()
    cfs = ax.contourf(X, Y, normPDF, cmap='viridis', alpha = 1, levels=np.linspace(-1,1,10))

#Create offset scatter for Group C
#    for (group, g2df), group_line in zip(df.groupby('group'), (line_c)):    
#        group_line.set_data(*g2df[['XX','YY']].values.T)

#        offset = lambda p: transforms.ScaledTranslation(p/82.,0, plt.gcf().dpi_scale_trans)
#        trans = plt.gca().transData
#        ax.scatter(line_c,transform=trans+offset(+2))


    return cfs.collections + [line_a, line_b]#, line_c] 

n = 10
time = range(n)  

d = ({
     'A1_X' :    [13.3,13.16,12.99,12.9,12.79,12.56,12.32,12.15,11.93,11.72],
     'A1_Y' :    [26.12,26.44,26.81,27.18,27.48,27.82,28.13,28.37,28.63,28.93],
     'A2_X' :    [6.97,6.96,7.03,6.98,6.86,6.76,6.55,6.26,6.09,5.9],
     'A2_Y' :    [10.92,10.83,10.71,10.52,10.22,10.02,9.86,9.7,9.54,9.37],
     'B1_X' :    [38.35,38.1,37.78,37.55,37.36,37.02,36.78,36.46,36.21,35.79],
     'B1_Y' :    [12.55,12.58,12.58,12.55,12.5,12.47,12.43,12.48,12.44,12.44],
     'B2_X' :    [14.6,14.38,14.16,13.8,13.45,13.11,12.71,12.3,12.06,11.61],
     'B2_Y' :    [4.66,4.44,4.24,4.1,4.01,3.84,3.67,3.56,3.44,3.47],
#    'C_X' :    [10,15,18,20,30,33,35,42,34,20],
#    'C_Y' :    [10,16,20,10,20,13,15,12,14,10],                 
     })

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 = 200
delay_ms = 1000
ani = animation.FuncAnimation(fig, plotmvs, frames=df.groupby('time'),
            blit=True, interval=interval_ms, repeat_delay=delay_ms)

plt.show()

Asmus · Accepted Answer

Ok, so I had to change a few things:

Most importanty, I turned off blit=True (mostly because I'm on a Mac, where this isn't fully supported, so your mileage may vary) and animated=True in the definitions of line_a and line_b.
The group "C" that you had created in your code comments was missing an integer for the tuples=[…] code to work (it expects to find an int(k.split('_')[0][1:])
The scatter plot should be a .scatter(), and since you're not using an init_func=init call within the FuncAnimation, we can simply create it "on the fly", within the animation function itself.
I've added a second for loop with an if group=='C' case — though you could solve this more elegantly — that creates the scatterplot scat. Note that I included the transform, although I'm not sure what this will achieve for you, in the end
I set the levels in ax.contourf() to the interval [0,1], but that's completely pedantic from my side ;-)
I've added zorder= in all plot items in order to control in which z-plane they are plotted (optional, but helpful), and slightly adjusted the look of the plotted lines (optional, only for emphasis)
In the end, you need to make sure you're returning an iterable list, hence the perhaps "odd looking" return cfs.collections + [scat] + [line_a,line_b]

Plot result:

Complete code:

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import scipy.stats as sts
import matplotlib.animation as animation
import matplotlib.transforms as transforms
from matplotlib.lines import Line2D

''' Section 1 '''

DATA_LIMITS = [-85, 85]

def datalimits(*data):
    return DATA_LIMITS  # dmin - spad, dmax + spad

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)
    PDF = sts.multivariate_normal([x, y]).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)):
        X, Y, PDF = mvpdf(x, y, xlim, ylim)
        PDFs.append(PDF)

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

''' Animate Plot '''

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

#Animated coordinates for group A,B
line_a, = ax.plot([], [], '-o', c='red', alpha = 0.5, markersize=5,zorder=3)
line_b, = ax.plot([], [], '-o', c='blue', alpha = 0.5, markersize=5,zorder=3)

cfs = None

def plotmvs(tdf, xlim=None, ylim=None, fig=fig, ax=ax):    
    global cfs  
    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)):
        group_line.set_data(*gdf[['X','Y']].values.T)
        X, Y, PDF = mvpdfs(gdf['X'].values, gdf['Y'].values, xlim, ylim)
        PDFs.append(PDF)

    for (group, gdf) in df.groupby('group'):
        if group=='C':
            offset = lambda p: transforms.ScaledTranslation(p/82.,0, plt.gcf().dpi_scale_trans)
            trans = plt.gca().transData
            scat=ax.scatter(gdf['X'].values, gdf['Y'].values, 
                    marker='o', c='white', alpha = 0.5,zorder=3,
                    transform=trans+offset(+2) )#markersize=2,zorder=3)


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

    cfs = ax.contourf(X, Y, normPDF, cmap='viridis', alpha = 1, levels=np.linspace(0,1,10),zorder=1)

    return  cfs.collections + [scat] + [line_a,line_b] # make sure that these are iterable!

n = 10
time = range(n)  

d = ({
    'A1_X' :    [13.3,13.16,12.99,12.9,12.79,12.56,12.32,12.15,11.93,11.72],
    'A1_Y' :    [26.12,26.44,26.81,27.18,27.48,27.82,28.13,28.37,28.63,28.93],
    'A2_X' :    [6.97,6.96,7.03,6.98,6.86,6.76,6.55,6.26,6.09,5.9],
    'A2_Y' :    [10.92,10.83,10.71,10.52,10.22,10.02,9.86,9.7,9.54,9.37],
    'B1_X' :    [38.35,38.1,37.78,37.55,37.36,37.02,36.78,36.46,36.21,35.79],
    'B1_Y' :    [12.55,12.58,12.58,12.55,12.5,12.47,12.43,12.48,12.44,12.44],
    'B2_X' :    [14.6,14.38,14.16,13.8,13.45,13.11,12.71,12.3,12.06,11.61],
    'B2_Y' :    [4.66,4.44,4.24,4.1,4.01,3.84,3.67,3.56,3.44,3.47],
    'C1_X' :    [10.,15.,18.,20.,30.,33.,35.,42.,34.,20.],## name contains an int so that tuples=... list works!
    'C1_Y' :    [10.,16.,20.,10.,20.,13.,15.,12.,14.,10.],
})

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 = 200
delay_ms = 1000
ani = animation.FuncAnimation(fig, plotmvs,  frames=df.groupby('time'), interval=interval_ms, repeat_delay=delay_ms,)

plt.show()

Update

If you'd rather have a single moving point, at the same time minimising/streamlining the lines of code, you could also initiate the plot elements with:

line_a, = ax.plot([], [], '-o', c='red', alpha = 0.5, markersize=5,zorder=3)
line_b, = ax.plot([], [], '-o', c='blue', alpha = 0.5, markersize=5,zorder=3)
lines=[line_a,line_b] ## this is iterable!

offset = lambda p: transforms.ScaledTranslation(p/82.,0, plt.gcf().dpi_scale_trans)
trans = plt.gca().transData

scat = ax.scatter([], [], s=5**2,marker='o', c='white', alpha = 0.5,zorder=3,transform=trans+offset(+2) )
scats=[scat] ## this is iterable, too!

and then inside plotmvs():

for (group, gdf), group_line in zip(df.groupby('group'), lines+scats):
        if group in ['A','B']:
            group_line.set_data(*gdf[['X','Y']].values.T)
            X, Y, PDF = mvpdfs(gdf['X'].values, gdf['Y'].values, xlim, ylim)
            PDFs.append(PDF)
        elif group in ['C']:
            x,y=(gdf['X'].values, gdf['Y'].values)
            scat.set_offsets( gdf[['X','Y']].values )

Note that updating scatter plots and line plots uses different functions and coordinate lists!

And finally update the return values:

return  cfs.collections + scats + lines # make sure that these are iterable!

Which yields the following animation:

Animate scatter and contour plot

Answers (1)

Related Questions