added checks for variant intervals exceeding chromosome ends

Author: avantikalal

src/grelu/data/dataset.py

@@ -16,6 +16,7 @@
 from torch.utils.data import Dataset
 
 from grelu.data.augment import Augmenter, _split_overall_idx
+from grelu.data.preprocess import filter_chrom_ends
 from grelu.data.utils import _check_multiclass, _create_task_data
 from grelu.sequence.format import (
     INDEX_TO_BASE_HASH,
@@ -152,6 +153,7 @@ def _load_seqs(self, seqs: Union[str, Sequence, pd.DataFrame, np.ndarray]) -> No
         seqs = resize(seqs, seq_len=self.padded_seq_len, end=self.end)
 
         if get_input_type(seqs) == "intervals":
+            seqs = filter_chrom_ends(seqs, genome=self.genome)
             self.intervals = seqs
             self.chroms = list(set(self.intervals.chrom))
         else:
@@ -603,7 +605,8 @@ def _load_seqs(self, variants: pd.DataFrame) -> None:
         from grelu.variant import variants_to_intervals
 
         self.padded_seq_len = self.seq_len + (2 * self.max_seq_shift)
-        self.intervals = variants_to_intervals(variants, seq_len=self.padded_seq_len)
+        intervals = variants_to_intervals(variants, seq_len=self.padded_seq_len)
+        self.intervals = filter_chrom_ends(intervals, genome=self.genome)
         self.seqs = convert_input_type(self.intervals, "indices", genome=self.genome)
 
     def __len__(self) -> int:
@@ -710,7 +713,8 @@ def _load_seqs(self, variants: pd.DataFrame) -> None:
         from grelu.variant import variants_to_intervals
 
         self.padded_seq_len = self.seq_len + (2 * self.max_seq_shift)
-        self.intervals = variants_to_intervals(variants, seq_len=self.padded_seq_len)
+        intervals = variants_to_intervals(variants, seq_len=self.padded_seq_len)
+        self.intervals = filter_chrom_ends(intervals, genome=self.genome)
         self.seqs = convert_input_type(self.intervals, "indices", genome=self.genome)
         self.n_seqs = self.seqs.shape[0]
 

src/grelu/data/preprocess.py

@@ -33,14 +33,15 @@ def filter_intervals(
     Returns:
         Filtered intervals in the same format (if inplace = False)
     """
-    print("Keeping {} intervals".format(sum(keep)))
-    if isinstance(data, pd.DataFrame):
-        return data.drop(index=data.index[~keep], inplace=inplace)
-    elif isinstance(data, AnnData):
-        if inplace:
-            data._inplace_subset_var(index=data.var_names[keep])
-        else:
-            return data[:, keep]
+    if sum(keep) < data.shape[0]:
+        print("Keeping {} of {} intervals".format(sum(keep), data.shape[0]))
+        if isinstance(data, pd.DataFrame):
+            return data.drop(index=data.index[~keep], inplace=inplace)
+        elif isinstance(data, AnnData):
+            if inplace:
+                data._inplace_subset_var(index=data.var_names[keep])
+            else:
+                return data[:, keep]
 
 
 def filter_obs(