Fix scikit-learn API compliance and MaskedDataset logic

Author: Pranav Choudhary

pyaptamer/benchmarking/_base.py

@@ -73,12 +73,13 @@ class Benchmarking:
     >>> summary = bench.run()  # doctest: +SKIP
     """
 
-    def __init__(self, estimators, metrics, X, y, cv=None):
+    def __init__(self, estimators, metrics, X, y, cv=None, labels=None):
         self.estimators = estimators if isinstance(estimators, list) else [estimators]
         self.metrics = metrics if isinstance(metrics, list) else [metrics]
         self.X = X
         self.y = y
         self.cv = cv
+        self.labels = labels
         self.results = None
 
     def _to_scorers(self, metrics):
@@ -128,8 +129,27 @@ def run(self):
         self.scorers_ = self._to_scorers(self.metrics)
         results = {}
 
-        for estimator in self.estimators:
-            est_name = estimator.__class__.__name__
+        if self.labels is not None:
+            if len(self.labels) != len(self.estimators):
+                raise ValueError("Length of labels must match length of estimators.")
+            names = self.labels
+        else:
+            counts = {}
+            for est in self.estimators:
+                name = est.__class__.__name__
+                counts[name] = counts.get(name, 0) + 1
+            
+            names = []
+            seen = {}
+            for est in self.estimators:
+                name = est.__class__.__name__
+                if counts[name] > 1:
+                    seen[name] = seen.get(name, 0) + 1
+                    names.append(f"{name}_{seen[name]}")
+                else:
+                    names.append(name)
+
+        for estimator, est_name in zip(self.estimators, names):
 
             cv_results = cross_validate(
                 estimator,

pyaptamer/datasets/dataclasses/_masked.py

@@ -85,7 +85,7 @@ def __init__(
         self.box = np.array(list(range(max_len)))
         self.len = len(self.x)
 
-    def _mask_rna(self, x_masked: Tensor, mask_positions: list[int]) -> Tensor:
+    def _mask_rna(self, x_masked: Tensor, mask_positions: list[int]) -> tuple[Tensor, list[int]]:
         """Mask adjacent nucleotides for RNA sequences.
 
         Parameters
@@ -97,20 +97,20 @@ def _mask_rna(self, x_masked: Tensor, mask_positions: list[int]) -> Tensor:
 
         Returns
         -------
-        Tensor
-            The tensor with adjacent nucleotides masked.
+        tuple[Tensor, list[int]]
+            The tensor with adjacent nucleotides masked and the list of adjacent positions.
         """
         adjacent_positions = []
         for pos in mask_positions:
             # mask position + 1 (if within bounds)
-            if pos < self.max_len - 1:
+            if pos < self.max_len - 1 and x_masked[pos + 1] > 0:
                 adjacent_positions.append(pos + 1)
             # mask position - 1 (if within bounds)
-            if pos > 0:
+            if pos > 0 and x_masked[pos - 1] > 0:
                 adjacent_positions.append(pos - 1)
         x_masked[adjacent_positions] = self.mask_idx
 
-        return x_masked
+        return x_masked, adjacent_positions
 
     def __len__(self) -> int:
         """
@@ -151,7 +151,7 @@ def __getitem__(self, index: int) -> tuple[Tensor, Tensor, Tensor, Tensor]:
         y = torch.tensor(self.y[index], dtype=torch.int64)
 
         x_masked = x.clone().detach()
-        y_masked = x.clone().detach()
+        y_masked = y.clone().detach()
 
         # non-padding positions (0 is padding)
         seq_len = torch.sum(x_masked > 0)
@@ -173,7 +173,8 @@ def __getitem__(self, index: int) -> tuple[Tensor, Tensor, Tensor, Tensor]:
 
         # for RNA, also mask adjacent nucleotides for base pairing
         if self.is_rna:
-            x_masked = self._mask_rna(x_masked, actual_mask_positions)
+            x_masked, adjacent_positions = self._mask_rna(x_masked, actual_mask_positions)
+            no_mask_positions = [pos for pos in no_mask_positions if pos not in adjacent_positions]
 
         # zero out non-masked positions in target
         y_masked[no_mask_positions] = 0

pyaptamer/trafos/encode/_greedy.py

@@ -132,7 +132,8 @@ def _transform(self, X):
 
         return result_df
 
-    def get_test_params(self):
+    @classmethod
+    def get_test_params(cls):
         """Get test parameters for GreedyEncoder.
 
         Returns