syrte
6/12/2017 - 5:31 PM

compare_violin.py

def compare_violin(x, y, bins, **kwargs):
    import numpy as np
    import seaborn as sns
    import matplotlib.pyplot as plt
    
    nmax = int(kwargs.pop('nmax', 1e4))
    nmin = int(kwargs.pop('nmax', 1))
    kwargs.setdefault('showmedians', True)
    
    bins_mid = (bins[:-1] + bins[1:]) * 0.5
    
    ix = (np.isfinite(x) & np.isfinite(y)).nonzero()
    x, y = x[ix], y[ix]
    
    idx = bins.searchsorted(x) - 1
    dat, pos = [], []
    for i in range(len(bins) - 1):
        a, b = y[idx == i], bins_mid[i]
        if a.size > nmax:
            a = np.random.choice(a, int(nmax), replace=False)
        if a.size >= nmin:
            dat.append(a)
            pos.append(b)
    plt.violinplot(dataset=dat, positions=pos, **kwargs)


def compare_violin2(x, y, bins=10, **kwargs):
    import numpy as np
    import seaborn as sns
    ix = np.isfinite(x) & np.isfinite(y)
    x, y = x[ix], y[ix]
    if np.isscalar(bins):
        nbin = bins
        xmax, xmin = x.max(), x.min()
    else:
        xmin, xmax, nbin = bins
    step = (xmax - xmin) / nbin
    x = ((x - xmin) // step + 0.5) * step + xmin
    sns.violinplot(x=x, y=y, **kwargs)