Add gnomad

Author: gonzalobenegas

experiments/evals/config/config.yaml

@@ -33,6 +33,7 @@ scorings:
 datasets:
   - traitgym_mendelian_promoter
   - traitgym_complex_promoter
+  - gnomad_promoter
   - sat_mut_mpra_promoter_F9
   - sat_mut_mpra_promoter_GP1BA
   - sat_mut_mpra_promoter_HBB

experiments/evals/workflow/Snakefile

@@ -2,6 +2,7 @@ configfile: "config/config.yaml"
 
 
 include: "rules/common.smk"
+include: "rules/gnomad.smk"
 include: "rules/metrics.smk"
 include: "rules/model.smk"
 include: "rules/sat_mut_mpra.smk"
@@ -32,3 +33,11 @@ rule all:
                 "absLLR.plus.score",
             ]
         ),
+        expand(
+            "results/metrics/gnomad_promoter/{metric}/{model}_{scoring}.tsv",
+            metric=["AUPRC", "AUROC"],
+            model=config["models"].keys(),
+            scoring=[
+                "LLR.minus.score",
+            ]
+        ),

experiments/evals/workflow/rules/common.smk

@@ -7,7 +7,7 @@ import gpn.model  # noqa: F401  # Registers the GPN architecture
 import numpy as np
 import pandas as pd
 from scipy.stats import spearmanr
-from sklearn.metrics import average_precision_score
+from sklearn.metrics import average_precision_score, roc_auc_score
 from transformers import AutoTokenizer, AutoModelForMaskedLM
 
 

experiments/evals/workflow/rules/gnomad.smk

@@ -0,0 +1,9 @@
+rule gnomad_promoter_dataset:
+    output:
+        "results/dataset/gnomad_promoter.parquet",
+    run:
+        V = pd.read_parquet("hf://datasets/songlab/gnomad_balanced/test.parquet")
+        V = V[V.consequence == "upstream_gene"]
+        V = V.groupby("label").sample(n=5000, random_state=42).reset_index(drop=True)
+        V = V.sort_values(COORDINATES)
+        V.to_parquet(output[0], index=False)

experiments/evals/workflow/rules/metrics.smk

@@ -13,6 +13,21 @@ rule metrics_AUPRC:
         pd.DataFrame({"AUPRC": [AUPRC]}).to_csv(output[0], sep="\t", index=False, float_format="%.3f")
 
 
+rule metrics_AUROC:
+    input:
+        "results/dataset/{dataset}.parquet",
+        "results/prediction/{dataset}/{model}.parquet",
+    output:
+        "results/metrics/{dataset}/AUROC/{model}.tsv",
+    wildcard_constraints:
+        dataset="|".join(config["datasets"]),
+    run:
+        y_true = pd.read_parquet(input[0], columns=["label"]).label
+        y_pred = pd.read_parquet(input[1], columns=["score"]).score
+        AUROC = roc_auc_score(y_true, y_pred)
+        pd.DataFrame({"AUROC": [AUROC]}).to_csv(output[0], sep="\t", index=False, float_format="%.3f")
+
+
 rule metrics_Spearman:
     input:
         "results/dataset/{dataset}.parquet",