0_3_0: cache prediction for the reference sequence to speed up predictions (#20)

* ensemble vep init

* backward compability of grelu, ensembling, testcases, custom fasta

* gene dataset

* gene expression prediction and sequence shifting

* fix testcase

* conflig

* cached reference prediction

* dtype reference caching safetensors

* bug fix for h5 datalaoder

---------

Co-authored-by: Muhammed Hasan Celik <[email protected]>

Author: MuhammedHasan

src/decima/cli/predict_genes.py

@@ -38,8 +38,24 @@
     is_flag=True,
     help="Save the replicates in the output parquet file. Default: False.",
 )
+@click.option(
+    "--float-precision",
+    type=str,
+    default="32",
+    help="Floating-point precision to be used in calculations. Avaliable options include: '16-true', '16-mixed', 'bf16-true', 'bf16-mixed', '32-true', '64-true', '32', '16', and 'bf16'.",
+)
 def cli_predict_genes(
-    output, genes, model, metadata, device, batch_size, num_workers, max_seq_shift, genome, save_replicates
+    output,
+    genes,
+    model,
+    metadata,
+    device,
+    batch_size,
+    num_workers,
+    max_seq_shift,
+    genome,
+    save_replicates,
+    float_precision,
 ):
     if model in ["0", "1", "2", "3"]:
         model = int(model)
@@ -63,5 +79,6 @@ def cli_predict_genes(
         max_seq_shift=max_seq_shift,
         genome=genome,
         save_replicates=save_replicates,
+        float_precision=float_precision,
     )
     ad.write_h5ad(output)

src/decima/cli/vep.py

@@ -68,6 +68,17 @@
     is_flag=True,
     help="Save the replicates in the output parquet file. Default: False.",
 )
+@click.option(
+    "--disable-reference-cache",
+    is_flag=True,
+    help="Disables the reference cache which significantly speeds up the computation by caching the reference expression predictios in the metadata.",
+)
+@click.option(
+    "--float-precision",
+    type=str,
+    default="32",
+    help="Floating-point precision to be used in calculations. Avaliable options include: '16-true', '16-mixed', 'bf16-true', 'bf16-mixed', '32-true', '64-true', '32', '16', and 'bf16'.",
+)
 def cli_predict_variant_effect(
     variants,
     output_pq,
@@ -85,6 +96,8 @@ def cli_predict_variant_effect(
     gene_col,
     genome,
     save_replicates,
+    disable_reference_cache,
+    float_precision,
 ):
     """Predict variant effect and save to parquet
 
@@ -108,6 +121,8 @@ def cli_predict_variant_effect(
 
         >>> decima vep -v "data/sample.vcf" -o "vep_results.parquet" --genome "path/to/fasta/hg38.fa"  # use custom genome build
     """
+    reference_cache = not disable_reference_cache
+
     if model in ["0", "1", "2", "3"]:  # replicate index
         model = int(model)
 
@@ -137,6 +152,8 @@ def cli_predict_variant_effect(
         gene_col=gene_col,
         genome=genome,
         save_replicates=save_replicates,
+        reference_cache=reference_cache,
+        float_precision=float_precision,
     )
 
 

src/decima/constants.py

@@ -1,3 +1,11 @@
+import os
+
+
 DECIMA_CONTEXT_SIZE = 524288
 SUPPORTED_GENOMES = {"hg38"}
 NUM_CELLS = 8856
+
+if "DECIMA_ENSEMBLE_MODELS_NAMES" in os.environ:
+    ENSEMBLE_MODELS_NAMES = os.environ["DECIMA_ENSEMBLE_MODELS_NAMES"].split(",")
+else:
+    ENSEMBLE_MODELS_NAMES = ["v1_rep0", "v1_rep1", "v1_rep2", "v1_rep3"]

src/decima/core/metadata.py

@@ -50,6 +50,7 @@ class GeneMetadata:
     gene_id: str
     pearson: float
     size_factor_pearson: float
+    ensembl_canonical_tss: Optional[bool]
 
     @classmethod
     def from_series(cls, name: str, series: pd.Series) -> "GeneMetadata":

src/decima/core/result.py

@@ -3,6 +3,7 @@
 import numpy as np
 import torch
 import pandas as pd
+
 from grelu.sequence.format import intervals_to_strings, strings_to_one_hot
 
 from decima.constants import DECIMA_CONTEXT_SIZE
@@ -143,7 +144,9 @@ def get_cell_metadata(self, cell: str) -> CellMetadata:
             raise KeyError(f"Cell {cell} not found in dataset. See avaliable cells with `result.cells`.")
         return CellMetadata.from_series(cell, self.cell_metadata.loc[cell])
 
-    def predicted_expression_matrix(self, genes: Optional[List[str]] = None) -> pd.DataFrame:
+    def predicted_expression_matrix(
+        self, genes: Optional[List[str]] = None, model_name: Optional[str] = None
+    ) -> pd.DataFrame:
         """Get predicted expression matrix for all or specific genes.
 
         Args:
@@ -152,10 +155,14 @@ def predicted_expression_matrix(self, genes: Optional[List[str]] = None) -> pd.D
         Returns:
             pd.DataFrame: Predicted expression matrix (cells x genes)
         """
+        model_name = "preds" if (model_name is None) or (model_name == "ensemble") else model_name
         if genes is None:
-            return pd.DataFrame(self.anndata.layers["preds"], index=self.cells, columns=self.genes)
+            return pd.DataFrame(self.anndata.layers[model_name], index=self.cells, columns=self.genes)
         else:
-            return pd.DataFrame(self.anndata[:, genes].layers["preds"], index=self.cells, columns=genes)
+            return pd.DataFrame(self.anndata[:, genes].layers[model_name], index=self.cells, columns=genes)
+
+    def predicted_gene_expression(self, gene, model_name):
+        return torch.from_numpy(self.anndata[:, gene].layers[model_name].ravel())
 
     def _pad_gene_metadata(self, gene_meta: pd.Series, padding: int = 0) -> pd.Series:
         """
@@ -184,6 +191,7 @@ def prepare_one_hot(self, gene: str, variants: Optional[List[Dict]] = None, padd
         Returns:
             torch.Tensor: One-hot encoding of the gene
         """
+
         assert gene in self.genes, f"{gene} is not in the anndata object"
         gene_meta = self._pad_gene_metadata(self.gene_metadata.loc[gene], padding)
 
@@ -201,6 +209,7 @@ def prepare_one_hot(self, gene: str, variants: Optional[List[Dict]] = None, padd
 
     def gene_sequence(self, gene: str, stranded: bool = True) -> str:
         """Get sequence for a gene."""
+
         try:
             assert gene in self.genes, f"{gene} is not in the anndata object"
         except AssertionError:

src/decima/data/dataset.py

@@ -1,16 +1,17 @@
 import warnings
 import torch
 import h5py
-import numpy as np
 import bioframe
+import numpy as np
+import pandas as pd
 from more_itertools import flatten
 from torch.utils.data import Dataset, default_collate
 from grelu.sequence.format import indices_to_strings
 from grelu.data.augment import Augmenter, _split_overall_idx
 from grelu.sequence.utils import reverse_complement
 
-from decima.constants import DECIMA_CONTEXT_SIZE
-from decima.data.read_hdf5 import _extract_center, index_genes, indices_to_one_hot
+from decima.constants import DECIMA_CONTEXT_SIZE, ENSEMBLE_MODELS_NAMES
+from decima.data.read_hdf5 import _extract_center
 from decima.core.result import DecimaResult
 
 from decima.model.metrics import WarningType
@@ -221,9 +222,12 @@ def __init__(
         distance_type="tss",
         min_distance=0,
         max_distance=float("inf"),
+        model_name=None,
+        reference_cache=True,
     ):
         super().__init__()
 
+        self.reference_cache = reference_cache
         self.result = DecimaResult.load(metadata_anndata)
 
         self.variants = self._overlap_genes(
@@ -253,6 +257,19 @@ def __init__(
         self.n_augmented = len(self.augmenter)
         self.padded_seq_len = DECIMA_CONTEXT_SIZE + (2 * self.max_seq_shift)
 
+        if (model_name is None) or (not reference_cache):
+            self.model_names = list()  # no reference caching
+        elif model_name == "ensemble":
+            self.model_names = ENSEMBLE_MODELS_NAMES
+        else:
+            self.model_names = [model_name]
+
+        for model_name in self.model_names:
+            assert model_name in self.result.anndata.layers.keys(), (
+                f"Model {model_name} not found in the metadata annotation. "
+                "You may not using the correct metadata file for this model."
+            )
+
     @staticmethod
     def overlap_genes(
         df_variants,
@@ -372,16 +389,25 @@ def __len__(self):
 
     def validate_allele_seq(self, gene, variant):
         seq = self.result.gene_sequence(gene)
-        vstart = variant.rel_pos
-        vend = vstart + len(variant.ref)
-        return (seq[vstart:vend] == variant.ref_tx) or (seq[vstart:vend] == variant.alt_tx)
+        pos = variant.rel_pos
+        ref_match = seq[pos : pos + len(variant.ref)] == variant.ref_tx
+        alt_match = seq[pos : pos + len(variant.alt)] == variant.alt_tx
+        return ref_match, alt_match
+
+    def predicted_expression_cache(self, gene):
+        return {model_name: self.result.predicted_gene_expression(gene, model_name) for model_name in self.model_names}
 
     def __getitem__(self, idx):
         seq_idx, augment_idx, allele_idx = _split_overall_idx(idx, (self.n_seqs, self.n_augmented, self.n_alleles))
 
         variant = self.variants.iloc[seq_idx]
         rel_pos = variant.rel_pos + self.max_seq_shift
 
+        # by default cache values are nan if matched with reference genome
+        # then it will be replaced with the predicted expression from cache.
+        pred_expr = {model_name: torch.full((self.result.shape[0],), torch.nan) for model_name in self.model_names}
+        ref_match, alt_match = self.validate_allele_seq(variant.gene, variant)
+
         warnings = list()
         if allele_idx:
             seq, mask = self.result.prepare_one_hot(
@@ -391,8 +417,11 @@ def __getitem__(self, idx):
             )
             allele = seq[:, rel_pos : rel_pos + len(variant.alt)]
             allele_tx = variant.alt_tx
+
+            if alt_match:
+                pred_expr = self.predicted_expression_cache(variant.gene)
         else:
-            if not self.validate_allele_seq(variant.gene, variant):
+            if (not ref_match) and (not alt_match):
                 warnings.append(WarningType.ALLELE_MISMATCH_WITH_REFERENCE_GENOME)
 
             seq, mask = self.result.prepare_one_hot(
@@ -403,23 +432,33 @@ def __getitem__(self, idx):
             allele = seq[:, rel_pos : rel_pos + len(variant.ref)]
             allele_tx = variant.ref_tx
 
-        if len(variant.ref_tx) == len(variant.alt_tx):  # not SNV there would be shifts
+            if ref_match:
+                pred_expr = self.predicted_expression_cache(variant.gene)
+
+        if len(variant.ref) == len(variant.alt):  # not SNV there would be shifts
             assert indices_to_strings(allele.argmax(axis=0)) == allele_tx
 
         inputs = torch.vstack([seq, mask])
-
         inputs = _extract_center(inputs, seq_len=self.padded_seq_len)
         inputs = self.augmenter(seq=inputs, idx=augment_idx)
-        return {
+
+        data = {
             "seq": inputs,
             "warning": warnings,
         }
+        if len(self.model_names) > 0:
+            data["pred_expr"] = pred_expr
+
+        return data
 
     def collate_fn(self, batch):
-        return {
+        _batch = {
             "seq": default_collate([i["seq"] for i in batch]),
             "warning": list(flatten([b["warning"] for b in batch])),
         }
+        if "pred_expr" in batch[0]:
+            _batch["pred_expr"] = default_collate([b["pred_expr"] for b in batch])
+        return _batch
 
     def __str__(self):
         return (

src/decima/hub/__init__.py

@@ -15,17 +15,6 @@ def login_wandb():
         wandb.login(host=os.environ.get("WANDB_HOST", DEFAULT_WANDB_HOST), relogin=True, anonymous="must", timeout=0)
 
 
-def get_model_name(model: Union[str, int] = 0) -> str:
-    if isinstance(model, int):
-        return f"decima_rep{model}"
-    elif isinstance(model, str):
-        return model
-    else:
-        raise ValueError(
-            f"Invalid model: {model} it need to be a string of model_name on wandb or an integer of replicate number {0, 1, 2, 3}"
-        )
-
-
 def load_decima_model(model: Union[str, int] = 0, device: Optional[str] = None):
     """Load a pre-trained Decima model from wandb or local path.
 
@@ -45,39 +34,32 @@ def load_decima_model(model: Union[str, int] = 0, device: Optional[str] = None):
     if isinstance(model, LightningModel):
         return model
     elif model == "ensemble":
-        model = EnsembleLightningModel(
+        return EnsembleLightningModel(
             [
                 load_decima_model(0, device),
                 load_decima_model(1, device),
                 load_decima_model(2, device),
                 load_decima_model(3, device),
             ]
         )
-        model.name = "ensemble"
-        return model
     elif isinstance(model, str):
-        model_name = get_model_name(model)
         if Path(model).exists():
-            model = LightningModel.load_from_checkpoint(model, map_location=device)
-            model.name = model_name
-            return model
-    elif isinstance(model, int):
-        model_name = get_model_name(model)
+            return LightningModel.load_safetensor(model, device=device)
+    elif model in {0, 1, 2, 3}:
+        model_name = f"rep{model}"
     else:
         raise ValueError(
             f"Invalid model: {model} it need to be a string of model_name on wandb "
             "or an integer of replicate number {0, 1, 2, 3}, or a path to a local model"
         )
 
     if model_name.upper() in os.environ:
-        return LightningModel.load_from_checkpoint(os.environ[model_name.upper()], map_location=device)
+        return LightningModel.load_safetensor(os.environ[model_name.upper()], device=device)
 
     art = get_artifact(model_name, project="decima")
     with TemporaryDirectory() as d:
         art.download(d)
-        model = LightningModel.load_from_checkpoint(Path(d) / "model.ckpt", map_location=device)
-        model.name = str(model_name)
-        return model
+        return LightningModel.load_safetensor(Path(d) / f"{model_name}.safetensors", device=device)
 
 
 def load_decima_metadata(path: Optional[str] = None):
@@ -95,7 +77,7 @@ def load_decima_metadata(path: Optional[str] = None):
     if "DECIMA_METADATA" in os.environ:
         return anndata.read_h5ad(os.environ["DECIMA_METADATA"])
 
-    art = get_artifact("decima_metadata", project="decima")
+    art = get_artifact("metadata", project="decima")
     with TemporaryDirectory() as d:
         art.download(d)
         return anndata.read_h5ad(Path(d) / "metadata.h5ad")

src/decima/interpret/attributions.py

@@ -5,7 +5,6 @@
 import torch
 import pyBigWig
 from pyfaidx import Faidx
-import genomepy
 from captum.attr import InputXGradient, Saliency, IntegratedGradients
 from grelu.interpret.motifs import scan_sequences
 from grelu.sequence.format import convert_input_type, strings_to_one_hot
@@ -451,6 +450,8 @@ def save_bigwig(self, bigwig_path: str):
 
         if self._chrom is not None:
             name = self.chrom
+            import genomepy
+
             sizes = genomepy.Genome("hg38").sizes
             bw.addHeader([(chrom, size) for chrom, size in sizes.items()])
         else: