Added enformer predictor

Snakefile

@@ -23,6 +23,8 @@ with open("system_config.yaml", "r") as fd:
 with open("download_urls.yaml", "r") as fd:
     download_urls = yaml.safe_load(fd)
 
+# TODO upload enformer reference; download reference rule
+
 SNAKEMAKE_DIR = os.path.abspath(os.path.dirname(workflow.snakefile))
 
 CONDA_ENV_YAML_DIR=f"{SNAKEMAKE_DIR}/envs"
@@ -85,7 +87,7 @@ rule all:
         expand(
             f"{RESULTS_DIR}/predict/{{model_type}}/{{vcf_file}}.parquet",
             vcf_file=vcf_input_file_names,
-            model_type=config["predict_abexp_models"], 
+            model_type=config["predict_abexp_models"],
         ),
 
 

config.yaml

@@ -21,3 +21,5 @@ gtf_file: "example/chr22.gencode.v34.annotation.gtf.gz"
 
 # mapping of chromosome names. By default, maps them to "chr" prefix.
 chrom_alias_tsv: "{SNAKEMAKE_DIR}/resources/chromAlias.tsv"
+
+use_enformer_gpu: false

defaults.yaml

@@ -101,11 +101,35 @@ absplice:
     - Vagina
     - Whole_Blood
   use_spliceai_rocksdb: True
-  spliceai_rocksdb_path: 
+  spliceai_rocksdb_path:
     hg19: '{RESOURCES_DIR}/spliceai_rocksdb/spliceAI_hg19_chr{chromosome}.db'
     hg38: '{RESOURCES_DIR}/spliceai_rocksdb/spliceAI_hg38_chr{chromosome}.db'
   spliceai_rocksdb_chromosomes: ['1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12', '13', '14', '15', '16', '17', '18', '19', '20', '21', '22', 'X', 'Y']
 
+enformer:
+  batch_size: 2
+  isoform_aggregation_mode: canonical
+  variant_upstream_tss: 50
+  variant_downstream_tss: 200
+  # Whether to calculate the variant effect only for the canonical transcripts
+  canonical_only: true
+  # Whether to calculate the variant effect only for the protein coding transcripts
+  protein_coding_only: true
+  # the chromosomes in which to look for transcripts
+  chromosomes: ["chr1", "chr2", "chr3", "chr4", "chr5", "chr6", "chr7", "chr8", "chr9", "chr10", "chr11", "chr12",
+  "chr13", "chr14", "chr15", "chr16", "chr17", "chr18", "chr19", "chr20", "chr21", "chr22", "chrX", "chrY" ]
+  # How many central bins to average over?
+  num_agg_central_bins: 3
+  # Enformer will be run 3 times, once centered on the TSS, and +-<shift> from the TSS
+  shift: 43
+  # How many bins to save from each run. A smaller number will save space, since
+  # we only care about the central bins (close to the TSS)
+  num_output_bins: 21
+  tissue_mapper_pkl: "{SNAKEMAKE_DIR}/resources/enformer/tissue_mapper.pkl"
+  tracks_yml: "{SNAKEMAKE_DIR}/resources/enformer/tracks.yaml"
+  # Whether to download the reference or generate them locally
+  download_reference: true
+
 models:
   abexp_v1.0:
     model: "{SNAKEMAKE_DIR}/resources/models/abexp_v1.0/model.joblib"

download_urls.yaml

@@ -9,3 +9,5 @@ splicemap:
     psi3: 'https://zenodo.org/record/6408906/files/{tissue}_splicemap_psi3_method%3Dkn_event_filter%3Dmedian_cutoff.csv.gz?download=1'
     psi5: 'https://zenodo.org/record/6408906/files/{tissue}_splicemap_psi5_method%3Dkn_event_filter%3Dmedian_cutoff.csv.gz?download=1'
 
+# todo: add the correct URL
+enformer_reference_tar: "/data/nasif12/home_if12/tsi/downloads/enformer_ref.tar"

envs/abexp-enformer-gpu.yaml

@@ -0,0 +1,22 @@
+name: abexp-enformer-gpu
+channels:
+  - conda-forge
+  - bioconda
+dependencies:
+  - python>=3.11
+  - attrs<=21.4.0 # because of kipoi
+  - cyvcf2>=0.30.28
+  - numpy>=1.26.4
+  - pandas>=2.2.1
+  - pyarrow>=16.1.0
+  - pyranges>=0.0.129
+  - polars>=0.20.19
+  - scikit-learn>=1.4.1.post1
+  - lightgbm>=4.3.0
+  - zarr>=2.15.0
+  - xarray>=2024.5.0
+  - tensorflow-hub==0.16.1
+  - tensorflow==2.16.1
+  - tensorflow-gpu==2.16.1
+  - pip:
+    - "git+https://github.com/gtsitsiridis/kipoi_veff_analysis.git@aparent2#egg=kipoi-enformer&subdirectory=pkgs/kipoi_enformer"

envs/abexp-enformer.yaml

@@ -0,0 +1,21 @@
+name: abexp-enformer
+channels:
+  - conda-forge
+  - bioconda
+dependencies:
+  - python>=3.11
+  - attrs<=21.4.0 # because of kipoi
+  - cyvcf2>=0.30.28
+  - numpy>=1.26.4
+  - pandas>=2.2.1
+  - pyarrow>=16.1.0
+  - pyranges>=0.0.129
+  - polars>=0.20.19
+  - scikit-learn>=1.4.1.post1
+  - lightgbm>=4.3.0
+  - zarr>=2.15.0
+  - xarray>=2024.5.0
+  - tensorflow-hub==0.16.1
+  - tensorflow==2.16.1
+  - pip:
+    - "git+https://github.com/gtsitsiridis/kipoi_veff_analysis.git@aparent2#egg=kipoi-enformer&subdirectory=pkgs/kipoi_enformer"