Stop mckp from silently swallowing the requested device

cellbender/remove_background/estimation.py

@@ -428,6 +428,7 @@ def estimate_noise(
         verbose: bool = False,
         n_chunks: Optional[int] = None,
         use_multiple_processes: bool = False,
+        device: str = "cpu",
         **kwargs,
     ) -> sp.csr_matrix:
         """Given the full probabilistic posterior, compute noise counts
@@ -442,6 +443,9 @@ def estimate_noise(
                 If None, targets about 5000 genes per chunk.
             use_multiple_processes: True to use multiprocessing. Seems faster
                 without using it, not entirely clear why
+            device: The backend the caller selected. Accepted so that requesting
+                an accelerator is not silently swallowed, but this estimator
+                computes its MAP step on the CPU regardless, to bound memory use.
 
         Returns:
             noise_count_csr: Estimated noise count matrix.
@@ -457,6 +461,11 @@ def estimate_noise(
 
         t0 = time.time()
 
+        if device != "cpu":
+            logger.debug(
+                f"Estimator 'mckp' computes its MAP step on the CPU rather than on {device}, to bound memory use."
+            )
+
         if use_multiple_processes:
             logger.info("Dividing dataset into chunks of genes")
             chunk_logic_list = self._gene_chunk_iterator(
@@ -579,6 +588,9 @@ def _chunk_estimate_noise(
 
         # First we need to compute the MAP to find out which direction to go.
         t = time.time()
+        # Deliberately on the CPU even when an accelerator was requested: MCKP
+        # is the memory-hungry estimator (issue #396) and this MAP runs over a
+        # whole chunk at once. estimate_noise() tells the user about it.
         map_dict = apply_function_dense_chunks(
             noise_log_prob_coo=noise_log_prob_coo, fun=MAP.torch_argmax, device="cpu"
         )

tests/test_estimation.py

@@ -408,3 +408,24 @@ def test_estimation_array_to_csr():
     truth_csr = coo.tocsr()
 
     assert sparse_matrix_equal(output_csr, truth_csr)
+
+
+def test_mckp_accepts_a_device_without_silently_ignoring_it(log_prob_coo, caplog):
+    """--cuda / --mps must not be swallowed without a word.
+
+    MCKP computes its MAP step on the CPU deliberately, to bound memory use
+    (issue #396). That is fine, but it used to absorb the caller's device into
+    **kwargs and say nothing, so the flag looked effective and was not.
+    """
+    import logging
+
+    converter = IndexConverter(total_n_cells=3, total_n_genes=5)
+    estimator = MultipleChoiceKnapsack(index_converter=converter)
+    with caplog.at_level(logging.DEBUG, logger="cellbender"):
+        estimator.estimate_noise(
+            noise_log_prob_coo=log_prob_coo["coo"],
+            noise_offsets=log_prob_coo["offsets"],
+            noise_targets_per_gene=np.array([1, 0, 0, 0, 0]),
+            device="cuda",
+        )
+    assert "mckp" in caplog.text and "CPU" in caplog.text