screen_dynamics.py 文件源码

def gRNA_scatter(s1, s2, prefix="", text=False, n_labels=30):
    # Scatter of gRNA change
    fig, axis = plt.subplots(3, 2, sharex=False, sharey=False, figsize=(8, 8))
    axis = axis.flatten()

    for i, screen in enumerate(s2.columns[::-1]):
        x = s1.join(s2)  # .fillna(0)
        x = x.iloc[np.random.permutation(len(x))]

        x = x.ix[x.index[~x.index.str.contains("Wnt")]]
        if prefix.startswith("mid_screen-"):
            b = x["gDNA_Jurkat"]
        else:
            b = x["plasmid_pool_TCR"]
        x = x.fillna(0)
        b = b.fillna(0)

        colors = pd.DataFrame()
        colors[sns.color_palette("colorblind")[0]] = x.index.str.contains("Wnt")
        colors[sns.color_palette("colorblind")[1]] = x.index.str.contains("CTRL")
        colors[sns.color_palette("colorblind")[2]] = x.index.str.contains("Tcr")
        colors[sns.color_palette("colorblind")[3]] = x.index.str.contains("Ess")
        colors = colors.apply(lambda x: x[x].index.tolist()[0], axis=1).tolist()

        axis[i].scatter(np.log2(1 + x[screen]), np.log2(1 + b), color=colors, alpha=0.5)
        if text:
            for j in x[x.index.str.contains("ETS1|GATA3|RUNX1")].index:
                axis[i].text(np.log2(1 + x[screen].ix[j]), np.log2(1 + b.ix[j]), j)

        # x = y line
        lims = [np.nanmin([np.log2(1 + x[screen]), np.log2(1 + b)]), np.nanmax([np.log2(1 + x[screen]), np.log2(1 + b)])]
        axis[i].plot((lims[0], lims[1]), (lims[0], lims[1]), linestyle='--', color='black', alpha=0.75)

        axis[i].set_title(screen)
    for i in range(0, len(axis), 2):
        axis[i].set_ylabel("gRNA frequency in plasmid (log2)")
    for ax in axis[-2:]:
        ax.set_xlabel("gRNA frequency in CROP-seq screen (log2)")
    sns.despine(fig)
    fig.savefig(os.path.join(results_dir, "gRNA_counts.norm.{}.scatter.{}svg".format(prefix, "text." if text else "")), bbox_inches="tight")
    fig.savefig(os.path.join(results_dir, "gRNA_counts.norm.{}.scatter.{}pdf".format(prefix, "text." if text else "")), bbox_inches="tight")