visualization.py 文件源码

def plot_kde_matrix(df, w, limits=None, colorbar=True):
    """
    Plot a KDE matrix.

    Parameters
    ----------
    df: Pandas Dataframe
        The rows are the observations, the columns the variables.
    w: np.narray
        The corresponding weights.
    colorbar: bool
        Whether to plot the colorbars or not.
    limits: dictionary, optional
        Dictionary of the form ``{"name": (lower_limit, upper_limit)}``.
    """
    grid = sns.PairGrid(df, diag_sharey=False)
    if limits is None:
        limits = {}

    default = (None, None)

    def off_diagonal(x, y, **kwargs):
        df = pd.concat((x, y), axis=1)
        plot_kde_2d(df, w,
                    x.name, y.name,
                    xmin=limits.get(x.name, default)[0],
                    xmax=limits.get(x.name, default)[1],
                    ymin=limits.get(y.name, default)[0],
                    ymax=limits.get(y.name, default)[1],
                    ax=plt.gca(), title=False, colorbar=colorbar)

    def scatter(x, y, **kwargs):
        alpha = w / w.max()
        colors = np.zeros((alpha.size, 4))
        colors[:, 3] = alpha
        plt.gca().scatter(x, y, color="k")
        plt.gca().set_xlim(*limits.get(x.name, default))
        plt.gca().set_ylim(*limits.get(y.name, default))

    def diagonal(x, **kwargs):
        df = pd.concat((x,), axis=1)
        plot_kde_1d(df, w, x.name,
                    xmin=limits.get(x.name, default)[0],
                    xmax=limits.get(x.name, default)[1],
                    ax=plt.gca())

    grid.map_diag(diagonal)
    grid.map_upper(scatter)
    grid.map_lower(off_diagonal)
    return grid