Add --exclude-low-support and --convergence-criterion flags

LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2022-2025 Benedetto Polimeni
+Copyright (c) 2022-2026 Benedetto Polimeni
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

irescue/count.py

@@ -94,7 +94,7 @@ def parse_maps(maps_file, feature_index):
 
 
 def compute_cell_counts(
-    equivalence_classes, features_index, max_iters, tolerance, dumpEC, no_umi
+    equivalence_classes, features_index, max_iters, tolerance, dumpEC, no_umi, exclude_low_support, convergence_criterion
 ):
     """
     Calculate TE counts of a single cell, given a list of equivalence classes.
@@ -259,18 +259,42 @@ def compute_cell_counts(
         em_array = em_array.tocsr()
 
         # save an array with features > 0, as in em_array order
-        tokeep = np.flatnonzero(em_array.sum(axis=0))
+        if not exclude_low_support:
+            tokeep = np.flatnonzero(em_array.sum(axis=0))
+        else:
+            # only keep features supported by >=2 multimapping reads,
+            # or >=1 multimapping and >=1 uniquely mapping
+            tokeep1 = np.where((em_array.sum(axis=0) >= 2).A1)[0]
+            tokeep2 = np.intersect1d(
+                # (-1 because of 0-based indexing of em_array
+                #  against 1-based indexing of features)
+                np.array(list(counts.keys())) - 1,
+                np.where((em_array.sum(axis=0) == 1).A1)[0]
+            )
+            tokeep = np.union1d(tokeep1, tokeep2)
+
         # remove unmapped features from em_array
         em_array = em_array[:, tokeep]
-        # run EM
-        em_counts, em_stats = run_em(
-            em_array, cycles=max_iters, tolerance=tolerance
-        )
-        em_counts = em_counts * em_array.shape[0]
 
-        for i, c in zip(tokeep + 1, em_counts):
-            if c > 0:
-                counts[i] += c
+        # removing some features may yield empty rows
+        # (not necessary if exclude_low_support is disabled)
+        if exclude_low_support:
+            em_array = em_array[(em_array.sum(axis=1)>0).A1, :]
+
+        if em_array.shape[1] > 0:
+            # run EM
+            em_counts, em_stats = run_em(
+                em_array, cycles=max_iters, tolerance=tolerance, convergence_criterion=convergence_criterion
+            )
+            em_counts = em_counts * em_array.shape[0]
+
+            # add EM-optimized counts to uniquely mapped counts
+            # (add +1 to features to multimapped features to keep
+            #  because of 0-based indexing of em_array against
+            #  1-based indexing of features)
+            for i, c in zip(tokeep + 1, em_counts):
+                if c > 0:
+                    counts[i] += c
     return dict(counts), dump, em_stats
 
 
@@ -292,6 +316,8 @@ def run_count(
     features_index,
     tmpdir,
     no_umi,
+    exclude_low_support,
+    convergence_criterion,
     dumpEC,
     max_iters,
     tolerance,
@@ -323,11 +349,13 @@ def run_count(
                 tolerance=tolerance,
                 dumpEC=dumpEC,
                 no_umi=no_umi,
+                exclude_low_support=exclude_low_support,
+                convergence_criterion=convergence_criterion
             )
             writerr(
                 f"[{taskn}] Write cell {cellidx} ({cellbarcode.decode()}). "
                 f"EM cycles: {em_stats[0]}. Converged: {em_stats[1]}. "
-                f"Log likelihood: {em_stats[2]}. Increment: {em_stats[3]}.",
+                f"Log likelihood: {em_stats[2] if convergence_criterion=='likelihood' else 'Not computed because of convergence criterion choice'}. Increment: {em_stats[3]}.",
                 level=1,
                 send=verbose,
             )

irescue/em.py

@@ -29,7 +29,7 @@ def log_likelihood(matrix, counts):
     return log_likelihood
 
 
-def run_em(matrix, cycles=100, tolerance=1e-4):
+def run_em(matrix, cycles=100, tolerance=1e-4, convergence_criterion="likelihood"):
     """
     Run Expectation-Maximization (EM) algorithm to redistribute read counts
     across a set of features.
@@ -42,6 +42,10 @@ def run_em(matrix, cycles=100, tolerance=1e-4):
         Number of EM cycles.
     tolerance : float
         Tolerance threshold of log-likelihood difference to infer convergence.
+    convergence_criterion : str
+        Criterion to determine convergence:
+            "likelihood": log-likelihood change < tolerance.
+            "estimates": feature abundances change < tolerance.
 
     Returns
     -------
@@ -59,8 +63,11 @@ def run_em(matrix, cycles=100, tolerance=1e-4):
     nFeatures = matrix.shape[1]
     counts = np.full(shape=nFeatures, fill_value=1 / nFeatures)
 
-    # Initial log-likelihood
-    prev_loglik = log_likelihood(matrix, counts)
+    # Initial log-likelihood (or initial counts)
+    if convergence_criterion=="likelihood":
+        prev = log_likelihood(matrix, counts)
+    else:
+        prev = counts.copy()
 
     converged = False
     curr_cycle = 0
@@ -71,15 +78,21 @@ def run_em(matrix, cycles=100, tolerance=1e-4):
         e_matrix = e_step(matrix=matrix, counts=counts)
         counts = m_step(matrix=e_matrix)
 
-        # Compute the new log-likelihood
-        loglik = log_likelihood(matrix, counts)
+        # Compute the new log-likelihood (depending on convergence criterion)
+        if convergence_criterion=="likelihood":
+            curr = log_likelihood(matrix, counts)
+        else:
+            curr = counts.copy()
 
         # Check for convergence
-        loglikdiff = loglik - prev_loglik
-        if np.abs(loglikdiff) < tolerance:
+        if convergence_criterion=="likelihood":
+            diff = np.abs(curr-prev)
+        else:
+            diff = np.abs(curr-prev).sum()
+        if diff < tolerance:
             converged = True
             break
 
-        prev_loglik = loglik
+        prev = curr.copy()
 
-    return counts, (curr_cycle, converged, loglik, loglikdiff)
+    return counts, (curr_cycle, converged, curr, diff)

irescue/main.py

@@ -177,10 +177,30 @@ def parseArguments():
         metavar="FLOAT",
         default=1e-4,
         help=(
-            "Log-likelihood change below which convergence is assumed "
+            "Change between EM iterations below which convergence is assumed, "
+            "calculated on --convergence-criterion."
             "(Default: %(default)s)."
         ),
     )
+    parser.add_argument(
+        "--convergence-criterion",
+        type=str,
+        metavar="STR",
+        choices=["likelihood", "estimates"],
+        default="likelihood",
+        help=(
+            "Criterion to define convergence. "
+            "One of: %(choices)s. (Default: %(default)s)."
+        ),
+    )
+    parser.add_argument(
+        "--exclude-low-support",
+        action="store_true",
+        help=(
+            "Exclude features supported by only 1 multimapping read. "
+            "This should improve performance without significantly affecting results."
+        ),
+    )
     parser.add_argument(
         "--dump-ec",
         action="store_true",
@@ -379,6 +399,8 @@ def main():
         feature_index,
         dirs["tmp"],
         args.no_umi,
+        args.exclude_low_support,
+        args.convergence_criterion,
         args.dump_ec,
         args.max_iters,
         args.tolerance,