python - 动画等高线和散点图

Question

我正在尝试根据一组xy 坐标 为散点图 和双变量高斯分布 设置动画。我会先记录调用散点和分布的具体代码，然后再记录我如何测量分布。

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)

有关如何使用单帧生成和绘制分布的完整工作代码

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)

我基本上想通过遍历每一行 xy 坐标来为这段代码添加动画。

Answer 1

这是对 OP 代码的非常快速和肮脏的修改，修复散点动画并添加(一种形式的)轮廓动画。

基本上，您首先要创建 artists为您的动画(在本例中为 Line2D 对象，由 plot() 返回)。随后，您创建一个 update 函数(以及可选的初始化函数)。在那个函数中，你更新现有的艺术家。我认为 example in the matplotlib docs解释一切。

在这种情况下，我修改了 OP 的 plotmvs 函数以用作更新函数(而不是 OP 提议的 animate 函数)。

contourf(即您的 cfs)返回的 QuadContourSet 本身不能用作艺术家，但您可以使用cfs.collections(支持 this SO answer )。但是，您仍然需要创建一个新的等高线图并删除旧的，而不仅仅是更新等高线数据。就个人而言，我更喜欢较低级别的方法:尝试在不调用 contourf 的情况下获取轮廓数据，然后像对散点图一样初始化和更新轮廓线。

不过，上面的方法是在下面的 OP 代码中实现的(只需复制、粘贴和运行):

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()