process 3d data for heatmap

make_heatmap_data.py

; not fast, just straitfoward
def make_heatmap_data(xs, ys, zs, func, cutoff, num_grids, min_samples):
    xmin = np.nanpercentile(xs, cutoff*100)
    xmax = np.nanpercentile(xs, (1-cutoff)*100)
    ymin = np.nanpercentile(ys, cutoff*100)
    ymax = np.nanpercentile(ys, (1-cutoff)*100)
    xm = (xmax-xmin)/num_grids
    ym = (ymax-ymin)/num_grids
    r = [[list() for _ in range(num_grids)] for _ in range(num_grids)]
    for x, y, z in zip(xs, ys, zs):
        if not np.isnan(x) and not np.isnan(y) and not np.isnan(z):
            ix = int((x-xmin)/xm)
            iy = int((ymax-y)/ym)
            if -1 < ix < num_grids and -1 < iy < num_grids:
                r[ix][iy] += [z]
    d = [[list() for _ in range(num_grids)] for _ in range(num_grids)]
    for ix in range(num_grids):
        for iy in range(num_grids):
            n = r[ix][iy]
            d[ix][iy] = func(n) if len(n) >= min_samples else np.nan
    return d


a = make_heatmap_data(
    df['oi_median_fri_C'].values, 
    df['oi_median_fri_P'].values, 
    df['!change'].values, func=np.mean, cutoff=0.01, num_grids=50, min_samples=10)
    
import seaborn as sns
sns.heatmap(a, vmin=-0.05, vmax=0.05, cmap=cm.seismic)