bug fix training and gene mask lightning model (#21)

Author: MuhammedHasan

src/decima/data/dataset.py

@@ -6,12 +6,12 @@
 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.sequence.format import indices_to_strings, indices_to_one_hot
 from grelu.data.augment import Augmenter, _split_overall_idx
 from grelu.sequence.utils import reverse_complement
 
 from decima.constants import DECIMA_CONTEXT_SIZE, ENSEMBLE_MODELS_NAMES
-from decima.data.read_hdf5 import _extract_center
+from decima.data.read_hdf5 import _extract_center, index_genes
 from decima.core.result import DecimaResult
 
 from decima.model.metrics import WarningType

src/decima/model/lightning.py

@@ -614,3 +614,29 @@ def predict_step(self, batch, batch_idx, dataloader_idx=0) -> Union[dict, Tensor
             return {"expression": expression, "warnings": batch["warning"]}
         else:
             return self(batch)
+
+
+class GeneMaskLightningModel(LightningModel):
+    def __init__(
+        self,
+        gene_mask_start,
+        gene_mask_end,
+        model_params: dict,
+        train_params: dict = {},
+        data_params: dict = {},
+        name: str = "",
+    ):
+        super().__init__(
+            model_params=model_params,
+            train_params=train_params,
+            data_params=data_params,
+            name=name,
+        )
+        self.gene_mask_start = gene_mask_start
+        self.gene_mask_end = gene_mask_end
+
+    def forward(self, x: Union[Tuple[Tensor, Tensor], Tensor, str, List[str]], logits: bool = False):
+        mask = torch.zeros((x.shape[0], 1, x.shape[2]), device=x.device, dtype=x.dtype)
+        mask[:, :, self.gene_mask_start : self.gene_mask_end] = 1
+        x = torch.cat([x, mask], dim=1)
+        return super().forward(x, logits)

tests/test_lightning.py

@@ -2,7 +2,7 @@
 import torch
 from decima.constants import DECIMA_CONTEXT_SIZE, NUM_CELLS
 from decima.data.dataset import VariantDataset
-from decima.model.lightning import LightningModel
+from decima.model.lightning import LightningModel, GeneMaskLightningModel
 from decima.model.metrics import WarningType
 
 from conftest import device
@@ -54,3 +54,14 @@ def test_LightningModel_predict_on_dataset_ensemble(lightning_model, df_variant)
     assert results["expression"].shape == (82, NUM_CELLS)
     assert results["warnings"]['unknown'] == 0
     assert results["warnings"]['allele_mismatch_with_reference_genome'] == 13
+
+
+@pytest.mark.long_running
+def test_GeneMaskLightningModel_forward():
+    seq = torch.randn(1, 4, DECIMA_CONTEXT_SIZE).to(device)
+    model = GeneMaskLightningModel(
+        gene_mask_start=200_000, gene_mask_end=300_000,
+        model_params={"n_tasks": NUM_CELLS, "init_borzoi": False}, name="v1_rep0"
+    ).to(device)
+    preds = model(seq)
+    assert preds.shape == (1, NUM_CELLS, 1)