<jupyter_start><jupyter_text>Varying the set of genes used in regression This notebook assumes you have run preprocessing and the `--vary-gene` experiment```shmake updatemake preprocesspython mutrans.py --vary-gene```<jupyter_code>import torch
import matplotlib
import matplotlib.pyplot as plt

matplotlib.rcParams["figure.dpi"] = 200
matplotlib.rcParams["axes.edgecolor"] = "gray"
matplotlib.rcParams["savefig.bbox"] = "tight"
matplotlib.rcParams["savefig.pad_inches"] = 0.01
matplotlib.rcParams['font.family'] = 'sans-serif'
matplotlib.rcParams['font.sans-serif'] = ['Arial', 'Avenir', 'DejaVu Sans']
matplotlib.rcParams.update({
    # 'text.usetex': True,
    'text.latex.preamble': r'\usepackage{amsfonts}',
})
result = torch.load("results/mutrans.vary_gene.pt")
result.keys()
print("\n".join(result[()].keys()))
def plot_metric(metric, descending=False, ylabel=None, filenames=()):
    plt.figure(figsize=(8,4))
    empty = (("exclude", (("gene", ".*"),)),)
    include = {}
    exclude = {}
    for k, v in result.items():
        if k == empty:
            continue
        k = dict(k)
        if "include" in k:
            gene = k["include"][0][1]
            include[gene] = v[metric]
        if "exclude" in k:
            gene = k["exclude"][0][1]
            exclude[gene] = v[metric]
    rankby = [(include[g], g) for g in include]
    rankby.sort(reverse=descending)
    genes = [g for _, g in rankby]
    X = list(range(len(genes)))
    plt.axhline(result[empty][metric], color="darkgreen", linestyle=":", label="No genes",
                zorder=-2)
    plt.plot(X, [include[g] for g in genes], "k.", color="darkred", label="A single gene")
    plt.plot(X, [include[g] for g in genes], "wo", markersize=8, zorder=-1)
    plt.axhline(result[()][metric], color="k", lw=1, linestyle="--", label="All genes",
                zorder=-2)
    plt.legend(loc="best")
    plt.xticks(X, labels=[g[1:-1] for g in genes])
    plt.ylabel(metric if ylabel is None else ylabel)
    plt.xlabel("Gene")
    for f in filenames:
        plt.savefig(f)
plot_metric("loss", ylabel="ELBO loss",
            filenames=["paper/vary_gene_loss.png"])
plot_metric("rate_scale", ylabel=r"E[$\sigma_4$]",
            filenames=["paper/vary_gene_loss.png"])
plot_metric("KL")
plot_metric("MAE")
plot_metric("RMSE")
def plot_elbo(descending=True, filenames=()):
    plt.figure(figsize=(8,4))
    empty = (("exclude", (("gene", ".*"),)),)
    include = {}
    exclude = {}
    for k, v in result.items():
        if k == empty:
            continue
        k = dict(k)
        if "include" in k:
            gene = k["include"][0][1]
            include[gene] = -v["loss"]
        if "exclude" in k:
            gene = k["exclude"][0][1]
            exclude[gene] = -v["loss"]
    rankby = [(include[g], g) for g in include]
    rankby.sort(reverse=descending)
    genes = [g for _, g in rankby]
    X = list(range(len(genes)))
    plt.axhline(-result[()]["loss"], color="darkgreen", linestyle=":", label="All genes",
                zorder=-2)
    plt.plot(X, [include[g] for g in genes], "ko", color="darkred", label="One gene",
            markeredgecolor="white", markersize=7)
    plt.axhline(-result[empty]["loss"], color="k", linestyle="--", lw=1, label="No genes",
                zorder=-2)
    plt.legend(loc="best")
    plt.xticks(X, labels=[g[1:-1] for g in genes])
    plt.ylabel("ELBO")
    plt.xlabel("Gene")
    plt.tight_layout()
    for f in filenames:
        plt.savefig(f)
        
plot_elbo(filenames=["paper/vary_gene_elbo.png"])<jupyter_output><empty_output><jupyter_text>--vary-nsp<jupyter_code>result = torch.load("results/mutrans.vary_nsp.pt")
result.keys()
def plot_elbo(descending=True, filenames=()):
    plt.figure(figsize=(8,4))
    empty = (("exclude", (("gene", ".*"),)),)
    full = (("include", (("gene", "^ORF1[ab]:"),)),)
    include = {}
    for k, v in result.items():
        k = dict(k)
        if "include" not in k:
            continue
        k = dict(k["include"])
        if "region" not in k:
            continue
        gene, region = k["region"]
        include[gene, region] = -v["loss"]
    rankby = [(include[gr], gr[-1]) for gr in include]
    rankby.sort(reverse=descending)
    labels = [g for _, g in rankby]
    elbos = [e for e, _ in rankby]
    X = list(range(len(labels)))
    plt.axhline(-result[full]["loss"], color="darkgreen", linestyle=":", label="All of ORF1",
                zorder=-2)
    plt.plot(X, elbos, "ko", color="darkred", label="One nsp",
             markeredgecolor="white", markersize=7)
    plt.axhline(-result[empty]["loss"], color="k", linestyle="--", lw=1, label="No nsps",
                zorder=-2)
    plt.legend(loc="best")
    plt.xticks(X, labels=labels)
    plt.ylabel("ELBO")
    plt.xlabel("Nonstructural protein within ORF1")
    plt.tight_layout()
    for f in filenames:
        plt.savefig(f)
        
plot_elbo(filenames=["paper/vary_nsp_elbo.png"])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.18522246829071876, 'fraction_numerical': 0.008455808335011073, 'mean_word_length': 3.105785123966942, 'pattern_counts': {'":': 3, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '87270', 'license': 'Apache-2.0', 'n_tokens_mistral': 1818, 'n_tokens_neox': 1740, 'n_words': 364, 'path': 'notebooks/mutrans_gene.ipynb', 'repo_name': 'broadinstitute/pyro-cov', 'size': '4967'}