Merge pull request #15 from dpeerlab/bugfix/threshold_convergence

Bugfix/threshold convergence

Author: Tobiaspk

src/segger/data/utils/threshold.py

@@ -0,0 +1,12 @@
+from skimage.filters import threshold_li
+
+def threshold_li_custom(arr, max_iter=100):
+    """Fallback to StopIteration if can't converge. Not implemented in threshold_li."""
+    n_iter = 0
+    def _callback(threshold):
+        nonlocal n_iter
+        n_iter += 1
+        if n_iter > max_iter:
+            raise StopIteration
+
+    return threshold_li(arr, iter_callback=_callback)

src/segger/data/writer.py

@@ -1,19 +1,17 @@
-from copyreg import pickle
-import os
+import gc
 import logging
+import numpy as np
 from lightning.pytorch.callbacks import BasePredictionWriter
-from skimage.filters import threshold_li, threshold_yen
+from skimage.filters import threshold_yen
+from .utils.threshold import threshold_li_custom
 from lightning.pytorch import Trainer, LightningModule
 from typing import Sequence, Any
 from pathlib import Path
 import polars as pl
 import torch
-
 from ..io import TrainingTranscriptFields, TrainingBoundaryFields
 from . import ISTDataModule
 
-# TODO: import datamodule, not trainer
-
 class ISTSegmentationWriter(BasePredictionWriter):
     """TODO: Description
     
@@ -32,10 +30,11 @@ def __init__(self, output_directory: Path, debug: bool = False):
         # setup debugging
         self.debug = debug
         self.path_debug = None
+        self.n_tx_predicted = 0
         if debug:
             logging.getLogger("segger").setLevel("DEBUG")
             self.path_debug = output_directory / "debug"
-            self.path_debug.mkdir(exist_ok=True)
+            self.path_debug.mkdir(exist_ok=True, parents=True)
 
     def write_on_epoch_end(
         self,
@@ -148,56 +147,77 @@ def assign_transcripts_to_cells(
         )
 
         # Per-gene thresholding (iterative to reduce memory usage)
-        logger.debug("Calculating per-gene similarity thresholds...")
-        feature_counts = (
+        logger.debug(f"Calculating per-gene similarity thresholds, using {segmentation.shape[0]/1e6:.1f}M transcripts...")
+        
+        segmentation_group = (
             segmentation
             .filter(pl.col('segger_cell_id').is_not_null())
-            .select(tx_fields.feature)
-            .to_series()
-            .value_counts()
+            .group_by(tx_fields.feature)
         )
-        thresholds = []
+
         n = 10_000_000
-        for feature, count in feature_counts.iter_rows():
-            similarities = (
-                segmentation
-                .filter(
-                    (pl.col(tx_fields.feature) == feature) &
-                    (pl.col('segger_cell_id').is_not_null())
-                )
-                .select('segger_similarity')
-            )
-            if count > n:
-                similarities = similarities.sample(n=n, seed=0)
-            similarities = similarities.to_series().to_numpy()
-            threshold_value = min(
-                threshold_li( similarities),
-                threshold_yen(similarities),
-            )
-            thresholds.append({
-                tx_fields.feature: feature,
-                'similarity_threshold': threshold_value,
-            })
-        thresholds = pl.DataFrame(thresholds)
-        
+        thresholds = []
+        failed_to_converge = []
+        n_groups = segmentation_group.len().height
+
+        for i, (feature, group) in enumerate(segmentation_group):
+
+            # log step
+            if (i + 1) % 50 == 0:
+                logger.debug(f"Processing feature {i+1}/{n_groups} (feature {feature[0]} | transcripts {group.shape[0]/1e3:.1f}K)...")
+
+            # sample if too many
+            arr = group["segger_similarity"]
+            if arr.shape[0] > n:
+                arr = arr.sample(n=n, seed=0)
+            arr = arr.to_numpy()
+
+            # threshold
+            try:
+                tye = threshold_yen(arr)
+                tli = threshold_li_custom(arr, max_iter=250)
+                threshold = min(tye, tli)
+            except StopIteration:
+                logger.debug(f"Failed to converge {feature[0]}. Will use 50% quantile of segger similarities of other genes as cutoff.")
+                failed_to_converge.append(feature[0])
+                continue
+
+            # append threshold
+            thresholds.append({tx_fields.feature: feature[0], "similarity_threshold": threshold.item(), "converged": True})
+            
+            # cleanup
+            del arr
+            gc.collect()
+
+        # backfill failed features in using the 80% quantile of thresholds
+        global_threshold = np.quantile([t["similarity_threshold"] for t in thresholds], .5)
+        for feature in failed_to_converge:
+            thresholds.append({tx_fields.feature: feature, "similarity_threshold": global_threshold, "converged": False})
+        logger.debug(f"Global Threshold: {global_threshold} | Used this to backfill {len(failed_to_converge)} features.")
+
         # Join
         logger.debug("Joining thresholds with segmentation...")
+        thresholds = pl.DataFrame(thresholds)
         segmentation = (
             segmentation
             .join(thresholds, on=tx_fields.feature, how='left')
             .drop(tx_fields.feature)
         )
+
+        logger.debug("Segmentation complete.")
         return segmentation
 
 
     # Debugging callbacks
     def on_predict_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0):
+        mask = batch['tx']['predict_mask']
+        self.n_tx_predicted += mask.sum().item()
         if not self.debug:
             return
         log_every = 50
         if batch_idx % log_every == 0:
             self.segger_logger.info(
-                f"Finished prediction batch '{batch_idx}'."
+                f"Finished prediction batch '{batch_idx}'. # TX so far {self.n_tx_predicted / 1e6:.1f}M"
             )
 
     def on_fit_start(self, trainer, pl_module):

src/segger/debug/segmentation.py

@@ -18,7 +18,7 @@ def run_segmentation_only(
     SLURMEnvironment.detect = lambda: False
 
     # Load data
-    writer = ISTSegmentationWriter(path_outputs)
+    writer = ISTSegmentationWriter(path_outputs, debug=True)
     adata = ad.read_h5ad(path_adata)
     with open(path_predictions, "rb") as f:
         predictions = pickle.load(f)