python - 使用 matplotlib 设置或 Hook 双变量分布值

Question

我使用 matplotlib 动画 双变量 gaussian 分布。我通过调整 COV matrix 来计算此 distribution 以说明特定变量。我可以提供有关此过程的更多详细信息，但基本上每个 scatter 点都包含一个易于识别的特定位移。我遇到的问题是尝试设置/修复/固定发行版 未涵盖的区域。您可以看到值随着颜色的变化而波动。

问题:是否可以将这些中性区域设置或固定为特定值并因此设置颜色。具体来说，xy 值未涵盖的 coordinates 不应更改 contour 值。它们应该固定为 0.5。

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

DATA_LIMITS = [-100, 100]

def datalimits(*data):
    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)

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)
cfs = None

def plotmvs(tdf, xlim=None, ylim=None, fig=fig, ax=ax):
    global cfs  
    if cfs:
        for tp in cfs.collections:

            tp.remove()

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

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

    cfs = ax.contourf(X, Y, normPDF, levels=10, cmap='viridis', alpha = 0.8)

    return cfs.collections + [line_a, line_b]

n = 10
time = range(n) 
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']

interval_ms = 200
delay_ms = 1000
ani = FuncAnimation(fig, plotmvs, frames=df.groupby('time'),
                    blit=True, interval=interval_ms, repeat_delay=delay_ms)

plt.show()

Answer 1

我更改了您的规范化并为 contourf() 提供了明确的 levels，从而给出了您想要的结果。对代码的改动很小；我换了

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

    cfs = ax.contourf(X, Y, normPDF, levels=10, cmap='viridis', alpha = 0.8)

与

    normPDF = PDF * .5/max(PDF.max(), -PDF.min()) + .5

    cfs = ax.contourf(X, Y, normPDF, cmap='viridis', alpha = 0.8,
        levels=np.arange(0, 1, .1))

这里是结果: