Merge pull request #46 from quadbio/feat/subset_categories

Enable subsetting categories

Author: Marius1311

src/cellmapper/_docs.py

@@ -91,6 +91,13 @@
 knn_dist_metric
     Distance metric to use for nearest neighbors. See the knn algorithms documentation for details. """
 
+_subset_categories = """\
+subset_categories
+    For categorical data, optionally specify a subset of categories to include in the mapping.
+    If None (default), all categories are included. If specified, only the listed categories
+    will be mapped, and others will be ignored. For numerical data, this parameter is ignored
+    with a warning. Can be a single category string or a list of category strings."""
+
 
 d = DocstringProcessor(
     t=_t,
@@ -106,4 +113,5 @@
     n_neighbors=_n_neighbors,
     use_rep=_use_rep,
     knn_dist_metric=_knn_dist_metric,
+    subset_categories=_subset_categories,
 )

src/cellmapper/model/cellmapper.py

@@ -517,6 +517,7 @@ def map(
         symmetrize: bool | None = None,
         self_edges: bool | None = None,
         prediction_postfix: str = "pred",
+        subset_categories: None | list[str] | str = None,
     ) -> "CellMapper":
         """
         Map data from reference to query datasets.
@@ -540,6 +541,7 @@ def map(
         %(symmetrize)s
         %(self_edges)s
         %(prediction_postfix)s
+        %(subset_categories)s
         """
         if self.knn is None:
             self.compute_neighbors(
@@ -553,27 +555,23 @@ def map(
             self.compute_mapping_matrix(kernel_method=kernel_method, symmetrize=symmetrize, self_edges=self_edges)
 
         if obs_keys is not None:
-            # Handle both single key and list of keys for backward compatibility
-            if isinstance(obs_keys, str):
+            # Normalize to list for consistent handling
+            obs_keys_list = [obs_keys] if isinstance(obs_keys, str) else obs_keys
+            for obs_key in obs_keys_list:
                 self.map_obs(
-                    key=obs_keys, t=t, diffusion_method=diffusion_method, prediction_postfix=prediction_postfix
+                    key=obs_key,
+                    t=t,
+                    diffusion_method=diffusion_method,
+                    prediction_postfix=prediction_postfix,
+                    subset_categories=subset_categories,
                 )
-            else:
-                for obs_key in obs_keys:
-                    self.map_obs(
-                        key=obs_key, t=t, diffusion_method=diffusion_method, prediction_postfix=prediction_postfix
-                    )
         if obsm_keys is not None:
-            # Handle both single key and list of keys for backward compatibility
-            if isinstance(obsm_keys, str):
+            # Normalize to list for consistent handling
+            obsm_keys_list = [obsm_keys] if isinstance(obsm_keys, str) else obsm_keys
+            for obsm_key in obsm_keys_list:
                 self.map_obsm(
-                    key=obsm_keys, t=t, diffusion_method=diffusion_method, prediction_postfix=prediction_postfix
+                    key=obsm_key, t=t, diffusion_method=diffusion_method, prediction_postfix=prediction_postfix
                 )
-            else:
-                for obsm_key in obsm_keys:
-                    self.map_obsm(
-                        key=obsm_key, t=t, diffusion_method=diffusion_method, prediction_postfix=prediction_postfix
-                    )
         if layer_key is not None:
             self.map_layers(key=layer_key, t=t, diffusion_method=diffusion_method)
         if obs_keys is None and obsm_keys is None and layer_key is None:
@@ -650,6 +648,7 @@ def map_obs(
         prediction_postfix: str = "pred",
         confidence_postfix: str = "conf",
         return_probabilities: bool = False,
+        subset_categories: None | list[str] | str = None,
     ) -> np.ndarray | csr_matrix | None:
         """
         Map observation data from reference dataset to query dataset.
@@ -672,6 +671,7 @@ def map_obs(
         return_probabilities
             If True, return the probability matrix for categorical data.
             Only applicable for categorical data. The matrix is never densified.
+        %(subset_categories)s
 
         Returns
         -------
@@ -705,6 +705,43 @@ def map_obs(
             or pd.api.types.is_string_dtype(reference_data)
         )
 
+        # Handle subset_categories parameter and warnings
+        if subset_categories is not None:
+            if not is_categorical:
+                logger.warning(
+                    "subset_categories parameter specified for numerical data in key '%s'. This parameter will be ignored for numerical data.",
+                    key,
+                )
+                subset_categories = None
+            else:
+                # Convert single string to list for consistent handling
+                if isinstance(subset_categories, str):
+                    subset_categories = [subset_categories]
+
+                # Check if specified categories exist in the data
+                available_categories = set(
+                    reference_data.cat.categories if hasattr(reference_data, "cat") else reference_data.unique()
+                )
+                invalid_categories = set(subset_categories) - available_categories
+
+                if invalid_categories:
+                    logger.warning(
+                        "Some specified categories for key '%s' do not exist in the data and will be ignored: %s. Available categories: %s",
+                        key,
+                        list(invalid_categories),
+                        list(available_categories),
+                    )
+                    # Filter out invalid categories
+                    subset_categories = [cat for cat in subset_categories if cat in available_categories]
+
+                    # If no valid categories remain, set to None to use all
+                    if not subset_categories:
+                        logger.warning(
+                            "No valid categories remaining for key '%s' after filtering. Using all available categories.",
+                            key,
+                        )
+                        subset_categories = None
+
         # Log the operation being performed
         data_type = "categorical" if is_categorical else "numerical"
         if t is None:
@@ -720,7 +757,13 @@ def map_obs(
 
         if is_categorical:
             return self._map_obs_categorical(
-                key, prediction_postfix, confidence_postfix, t, diffusion_method, return_probabilities
+                key,
+                prediction_postfix,
+                confidence_postfix,
+                t,
+                diffusion_method,
+                return_probabilities,
+                subset_categories,
             )
         else:
             if return_probabilities:
@@ -736,12 +779,35 @@ def _map_obs_categorical(
         t: int | None,
         diffusion_method: Literal["iterative", "spectral"],
         return_probabilities: bool = False,
+        subset_categories: None | list[str] = None,
     ) -> np.ndarray | csr_matrix | None:
         """Map categorical observation data using one-hot encoding."""
-        onehot = OneHotEncoder(dtype=np.float32)
-        xtab = onehot.fit_transform(
-            self.reference.obs[[key]],
-        )  # shape = (n_reference_cells x n_categories), sparse csr matrix, float32
+        # Get the reference data
+        reference_data = self.reference.obs[key]
+
+        if subset_categories is not None:
+            # Create a filtered version of reference data for one-hot encoding
+            # Only include rows that have the desired categories
+            mask = reference_data.isin(subset_categories)
+
+            # Create a filtered DataFrame with only the subset categories
+            filtered_reference_data = reference_data.copy()
+            filtered_reference_data[~mask] = pd.NA  # Set non-subset categories to missing
+
+            # Create one-hot encoding only for the subset categories
+            onehot = OneHotEncoder(dtype=np.float32, handle_unknown="ignore")
+            # Create a DataFrame with only subset categories for fitting
+            subset_df = pd.DataFrame({key: pd.Categorical(subset_categories, categories=subset_categories)})
+            onehot.fit(subset_df)
+
+            # Transform the full reference data (missing values will be ignored)
+            xtab = onehot.transform(filtered_reference_data.to_frame())
+        else:
+            # Use the original approach for all categories
+            onehot = OneHotEncoder(dtype=np.float32)
+            xtab = onehot.fit_transform(self.reference.obs[[key]])
+
+        # Apply the mapping
         ytab = self.mapping_operator.apply(
             xtab, t=t, diffusion_method=diffusion_method
         )  # shape = (n_query_cells x n_categories), sparse csr matrix, float32

tests/model/test_query_to_reference_mapping.py

@@ -307,3 +307,135 @@ def test_map_obs_pseudotime_cross_mapping(self, query_reference_adata):
 
         # Verify no confidence scores for numerical data
         assert "dpt_pseudotime_conf" not in cmap.query.obs
+
+    def test_map_obs_subset_categories(self, query_reference_adata):
+        """Test mapping with subset_categories parameter for categorical data."""
+        query, reference = query_reference_adata
+
+        # Create CellMapper and compute mapping matrix
+        cmap = CellMapper(query=query, reference=reference)
+        cmap.compute_neighbors(n_neighbors=30, use_rep="X_pca", knn_method="sklearn")
+        cmap.compute_mapping_matrix(kernel_method="gauss")
+
+        # Get available leiden categories in reference
+        available_categories = list(reference.obs["leiden"].cat.categories)
+
+        # Test with subset of categories
+        subset_cats = available_categories[:2]  # Take first 2 categories
+        cmap.map_obs(key="leiden", subset_categories=subset_cats)
+
+        # Check that mapping was performed
+        assert "leiden_pred" in cmap.query.obs
+        assert "leiden_conf" in cmap.query.obs
+
+        # Check that predictions only contain subset categories (or might be missing if no assignment)
+        predicted_categories = set(cmap.query.obs["leiden_pred"].dropna().unique())
+        assert predicted_categories.issubset(set(subset_cats)), (
+            f"Predicted categories {predicted_categories} not subset of {subset_cats}"
+        )
+
+    def test_map_obs_subset_categories_single_string(self, query_reference_adata):
+        """Test mapping with subset_categories as single string."""
+        query, reference = query_reference_adata
+
+        # Create CellMapper and compute mapping matrix
+        cmap = CellMapper(query=query, reference=reference)
+        cmap.compute_neighbors(n_neighbors=30, use_rep="X_pca", knn_method="sklearn")
+        cmap.compute_mapping_matrix(kernel_method="gauss")
+
+        # Get first available category
+        first_category = reference.obs["leiden"].cat.categories[0]
+
+        # Test with single category as string
+        cmap.map_obs(key="leiden", subset_categories=first_category)
+
+        # Check that mapping was performed and only contains the specified category
+        assert "leiden_pred" in cmap.query.obs
+        predicted_categories = set(cmap.query.obs["leiden_pred"].dropna().unique())
+        assert predicted_categories.issubset({first_category}), (
+            f"Predicted categories {predicted_categories} not subset of {first_category}"
+        )
+
+    def test_map_obs_subset_categories_invalid_categories(self, query_reference_adata, caplog):
+        """Test mapping with some invalid categories in subset_categories."""
+        query, reference = query_reference_adata
+
+        # Create CellMapper and compute mapping matrix
+        cmap = CellMapper(query=query, reference=reference)
+        cmap.compute_neighbors(n_neighbors=30, use_rep="X_pca", knn_method="sklearn")
+        cmap.compute_mapping_matrix(kernel_method="gauss")
+
+        # Mix valid and invalid categories
+        valid_category = reference.obs["leiden"].cat.categories[0]
+        invalid_categories = ["nonexistent1", "nonexistent2"]
+        mixed_categories = [valid_category] + invalid_categories
+
+        # Test with mixed valid/invalid categories - should work without errors
+        cmap.map_obs(key="leiden", subset_categories=mixed_categories)
+
+        # Check that mapping still worked with valid categories
+        assert "leiden_pred" in cmap.query.obs
+        predicted_categories = set(cmap.query.obs["leiden_pred"].dropna().unique())
+        assert predicted_categories.issubset({valid_category})
+
+    def test_map_obs_subset_categories_all_invalid(self, query_reference_adata, caplog):
+        """Test mapping with all invalid categories in subset_categories."""
+        query, reference = query_reference_adata
+
+        # Create CellMapper and compute mapping matrix
+        cmap = CellMapper(query=query, reference=reference)
+        cmap.compute_neighbors(n_neighbors=30, use_rep="X_pca", knn_method="sklearn")
+        cmap.compute_mapping_matrix(kernel_method="gauss")
+
+        # Use only invalid categories
+        invalid_categories = ["nonexistent1", "nonexistent2"]
+
+        # Test with all invalid categories - should fallback to using all categories
+        cmap.map_obs(key="leiden", subset_categories=invalid_categories)
+
+        # Check that mapping still worked with all categories (fallback)
+        assert "leiden_pred" in cmap.query.obs
+        # Should have predictions from all available categories since it fell back
+        predicted_categories = set(cmap.query.obs["leiden_pred"].dropna().unique())
+        available_categories = set(reference.obs["leiden"].cat.categories)
+        # At least one category should be predicted (could be subset due to k-NN mapping)
+        assert len(predicted_categories) > 0
+        assert predicted_categories.issubset(available_categories)
+
+    def test_map_obs_subset_categories_numerical_warning(self, query_reference_adata, caplog):
+        """Test that subset_categories generates warning for numerical data."""
+        query, reference = query_reference_adata
+
+        # Create CellMapper and compute mapping matrix
+        cmap = CellMapper(query=query, reference=reference)
+        cmap.compute_neighbors(n_neighbors=30, use_rep="X_pca", knn_method="sklearn")
+        cmap.compute_mapping_matrix(kernel_method="gauss")
+
+        # Test with numerical data and subset_categories - should work and ignore the parameter
+        cmap.map_obs(key="dpt_pseudotime", subset_categories=["some_category"])
+
+        # Check that mapping still worked normally (parameter was ignored)
+        assert "dpt_pseudotime_pred" in cmap.query.obs
+        # Confidence scores should not be created for numerical data
+        assert "dpt_pseudotime_conf" not in cmap.query.obs
+
+    def test_map_method_with_subset_categories(self, query_reference_adata):
+        """Test that subset_categories parameter works through the high-level map method."""
+        query, reference = query_reference_adata
+
+        # Create CellMapper
+        cmap = CellMapper(query=query, reference=reference)
+
+        # Get available categories
+        available_categories = list(reference.obs["leiden"].cat.categories)
+        subset_cats = available_categories[:2]
+
+        # Test high-level map method with subset_categories
+        cmap.map(
+            obs_keys="leiden", n_neighbors=30, use_rep="X_pca", kernel_method="gauss", subset_categories=subset_cats
+        )
+
+        # Check that mapping was performed with subset
+        assert "leiden_pred" in cmap.query.obs
+        predicted_categories = set(cmap.query.obs["leiden_pred"].dropna().unique())
+        assert predicted_categories.issubset(set(subset_cats))