Check that the correct test dataset is being used for each trained re…

…ference.

src/singler/classify_integrated.py

@@ -76,6 +76,10 @@ def classify_integrated(
     """
     if isinstance(test_data, summarizedexperiment.SummarizedExperiment):
         test_data = test_data.assay(assay_type)
+
+    if test_data.shape[0] != integrated_prebuilt._test_num_features: # TODO: move to singlepp.
+        raise ValueError("number of rows in 'test_data' is not consistent with 'test_features=' used to create 'integrated_prebuilt'")
+
     ref_labs = integrated_prebuilt.reference_labels
 
     # Applying the sanity checks.

src/singler/classify_single.py

@@ -65,6 +65,10 @@ def classify_single(
     """
     if isinstance(test_data, summarizedexperiment.SummarizedExperiment):
         test_data = test_data.assay(assay_type)
+
+    if test_data.shape[0] != ref_prebuilt._test_num_features: # TODO: move to singlepp
+        raise ValueError("number of rows in 'test_data' is not consistent with 'test_features=' used to create 'ref_prebuilt'")
+
     test_ptr = mattress.initialize(test_data)
 
     best, raw_scores, delta = lib.classify_single(

src/singler/train_integrated.py

@@ -14,13 +14,14 @@ class TrainedIntegratedReferences:
     """Object containing integrated references, typically constructed by
     :py:meth:`~singler.train_integrated.train_integrated`."""
 
-    def __init__(self, ptr, ref_names, ref_labels):
+    def __init__(self, ptr: int, ref_names: Optional[Sequence], ref_labels: list, test_num_features: int):
         self._ptr = ptr
         self._names = ref_names
         self._labels = ref_labels
+        self._test_num_features = test_num_features # TODO: move to singlepp.
 
     @property
-    def reference_names(self) -> Union[Sequence[str], None]:
+    def reference_names(self) -> Union[Sequence, None]:
         """Sequence containing the names of the references. Alternatively
         ``None``, if no names were supplied."""
         return self._names
@@ -98,5 +99,6 @@ def train_integrated(
     return TrainedIntegratedReferences(
         ptr=ibuilt,
         ref_names=ref_names, 
-        ref_labels=[x.labels for x in ref_prebuilt]
+        ref_labels=[x.labels for x in ref_prebuilt],
+        test_num_features = len(test_features),
     )

src/singler/train_single.py

@@ -20,19 +20,21 @@ class TrainedSingleReference:
 
     def __init__(
         self,
-        ptr,
-        full_data,
+        ptr: int,
+        full_data: Any,
         full_label_codes: numpy.ndarray,
         labels: Sequence,
         features: Sequence,
-        markers: dict[Any, dict[Any, Sequence]]
+        markers: dict[Any, dict[Any, Sequence]],
+        test_num_features: int,
     ):
         self._ptr = ptr
         self._full_data = full_data
         self._full_label_codes = full_label_codes
         self._features = features
         self._labels = labels
         self._markers = markers
+        self._test_num_features = test_num_features # TODO: move to singlepp.
 
     def num_markers(self) -> int:
         """
@@ -233,10 +235,12 @@ def train_single(
     if test_features is None:
         test_features_idx = numpy.array(range(len(ref_features)), dtype=numpy.uint32)
         ref_features_idx = numpy.array(range(len(ref_features)), dtype=numpy.uint32)
+        test_num_features = len(ref_features)
     else:
         common_features = _stable_intersect(test_features, ref_features)
         test_features_idx = biocutils.match(common_features, test_features, dtype=numpy.uint32)
         ref_features_idx = biocutils.match(common_features, ref_features, dtype=numpy.uint32)
+        test_num_features = len(test_features)
 
     ref_ptr = mattress.initialize(ref_data)
     builder, _ = knncolle.define_builder(nn_parameters)
@@ -255,6 +259,7 @@ def train_single(
         labels = unique_labels,
         features = ref_features,
         markers = markers,
+        test_num_features = test_num_features,
     )