code review

Author: Muhammed Hasan Celik

src/decima/__init__.py

@@ -1,5 +1,7 @@
 import sys
+from decima.core.result import DecimaResult
 from decima.interpret.save_attributions import predict_save_attributions
+from decima.vep import predict_variant_effect
 
 
 if sys.version_info[:2] >= (3, 8):
@@ -16,3 +18,6 @@
     __version__ = "unknown"
 finally:
     del version, PackageNotFoundError
+
+
+__all__ = ["DecimaResult", "predict_variant_effect", "predict_save_attributions"]

src/decima/core/metadata.py

@@ -10,8 +10,8 @@ class GeneMetadata:
     Attributes:
         name: Gene name
         chrom: Chromosome where the gene is located
-        start: Start position in the chromosome
-        end: End position in the chromosome
+        start: Start position of the region around the gene to perform predictions in the chromosome
+        end: End position of the region around the gene to perform predictions in the chromosome
         strand: Strand orientation (+ or -)
         gene_type: Type of gene (e.g., protein_coding)
         frac_nan: Fraction of NaN values
@@ -91,8 +91,8 @@ class CellMetadata:
     disease: str
     study: str
     dataset: str
-    region: str
-    subregion: str
+    region: Optional[str]
+    subregion: Optional[str]
     celltype_coarse: Optional[str]
     n_cells: int
     total_counts: float

src/decima/data/dataset.py

@@ -328,9 +328,6 @@ def __getitem__(self, idx):
         variant = self.variants.iloc[seq_idx]
 
         warnings = list()
-        if not self.validate_allele_seq(variant.gene, variant):
-            warnings.append(WarningType.ALLELE_MISMATCH_WITH_REFERENCE_GENOME)
-
         if allele_idx:
             seq, mask = self.result.prepare_one_hot(
                 variant.gene,
@@ -339,6 +336,9 @@ def __getitem__(self, idx):
             allele = seq[:, variant.rel_pos : variant.rel_pos + len(variant.alt)]
             allele_tx = variant.alt_tx
         else:
+            if not self.validate_allele_seq(variant.gene, variant):
+                warnings.append(WarningType.ALLELE_MISMATCH_WITH_REFERENCE_GENOME)
+
             seq, mask = self.result.prepare_one_hot(
                 variant.gene,
                 variants=[{"chrom": variant.chrom, "pos": variant.pos, "ref": variant.alt, "alt": variant.ref}],

src/decima/interpret/attributions.py

@@ -400,15 +400,14 @@ def scan_motifs(self, motifs: str = "hocomoco_v12", window: int = 18, pthresh: f
         return df.sort_values("p-value")
 
     def plot_peaks(self, overlapping_min_dist=1000, figsize=(10, 2)):
-        """Plot attribution scores in a window around a relative location.
+        """Plot attribution scores and highlight peaks.
 
         Args:
-            relative_loc: Position relative to TSS to center plot on
-            window: Number of bases to show on each side of center
+            overlapping_min_dist: Minimum distance between peaks to consider them overlapping
             figsize: Figure size in inches (width, height)
 
         Returns:
-            matplotlib.pyplot.Figure: Attribution plot showing the window around the specified location
+            plotnine.ggplot: The plotted figure showing attribution scores with highlighted peaks
         """
         return plot_peaks(
             self.attrs,

src/decima/model/lightning.py

@@ -19,7 +19,8 @@
 
 from .decima_model import DecimaModel
 from .loss import TaskWisePoissonMultinomialLoss
-from .metrics import DiseaseLfcMSE, WarningCounter, WarningType
+from .metrics import DiseaseLfcMSE, WarningCounter
+
 
 default_train_params = {
     "lr": 4e-5,
@@ -448,11 +449,7 @@ def predict_on_dataset(
 
         expression = np.mean(expression, axis=1)  # B T
 
-        num_warnings = self.warning_counter.compute()
-        # allele mismatch is counted twice for each variant due to the two alleles
-        num_warnings[WarningType.ALLELE_MISMATCH_WITH_REFERENCE_GENOME.value] //= 2
-
-        return {"expression": expression, "warnings": num_warnings}
+        return {"expression": expression, "warnings": self.warning_counter.compute()}
 
     def get_task_idxs(
         self,

src/decima/plot/visualize.py

@@ -197,7 +197,7 @@ def plot_peaks(attrs, tss_pos, df_peaks=None, overlapping_min_dist=1000, figsize
 
     Args:
         attr: Attribution scores array
-        tss_pos: Position of TSS
+        tss_pos: Position of TSS (relative to the gene)
         df_peaks: DataFrame containing peak information
         overlapping_min_dist: Minimum distance between peaks to consider them overlapping
         figsize: Figure size

src/decima/utils/inject.py

@@ -80,22 +80,16 @@ def inject(self, variant: Dict):
         variant_start = variant["pos"]
         variant_end = variant_start + len(variant["ref"])
 
-        if variant_end <= self.start:
+        if variant_start < self.start:
             warnings.warn(
                 f"Variant position `{variant['pos']}` is upstream of the interval `[{self.start}, {self.end}]`. Skipping..."
             )
             return self
-        elif variant_start < self.start < variant_end:
-            _, right_variant = self._split_variant(variant, self.start)
-            return self.inject(right_variant)
-        elif self.end < variant_start:
+        elif self.end < variant_end:
             warnings.warn(
                 f"Variant position `{variant['pos']}` is downstream of the interval `[{self.start}, {self.end}]`. Skipping..."
             )
             return self
-        elif variant_start < self.end < variant_end:
-            left_variant, _ = self._split_variant(variant, self.end)
-            return self.inject(left_variant)
 
         if variant_start < self.anchor < variant_end:
             left_variant, right_variant = self._split_variant(variant, self.anchor)