Marco cominelli01 patch 2

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, keep_all_features, convergence_criterion
 ):
     """
     Calculate TE counts of a single cell, given a list of equivalence classes.
@@ -259,18 +259,32 @@ 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 keep_all_features:
+            tokeep = np.flatnonzero(em_array.sum(axis=0))
+        else:
+            # keep features supported by at least 2 multimapping reads
+            tokeep1 = np.where((em_array.sum(axis=0) >= 2).A1)[0]
+            # rescue features supported by only 1 multimapping read but also by (at least) 1 uniquely mapping read
+            tokeep2 = np.intersect1d(
+                np.array(list(counts.keys())),
+                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 step if keep_all_features=True)
+        if not keep_all_features:
+            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]
+
+            for i, c in zip(tokeep + 1, em_counts):
+                if c > 0:
+                    counts[i] += c
     return dict(counts), dump, em_stats
 
 
@@ -292,6 +306,8 @@ def run_count(
     features_index,
     tmpdir,
     no_umi,
+    keep_all_features,
+    convergence_criterion,
     dumpEC,
     max_iters,
     tolerance,
@@ -323,11 +339,13 @@ def run_count(
                 tolerance=tolerance,
                 dumpEC=dumpEC,
                 no_umi=no_umi,
+                keep_all_features=keep_all_features,
+                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,8 @@ 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" or "parameters"
 
     Returns
     -------
@@ -59,8 +61,8 @@ 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 parameters)
+    prev = log_likelihood(matrix, counts) if convergence_criterion=="likelihood" else counts
 
     converged = False
     curr_cycle = 0
@@ -71,15 +73,15 @@ 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 (or new parameters)
+        curr = log_likelihood(matrix, counts) if convergence_criterion=="likelihoods" else counts
 
         # Check for convergence
-        loglikdiff = loglik - prev_loglik
-        if np.abs(loglikdiff) < tolerance:
+        diff = np.abs(curr-prev) if convergence_criterion=="likelihood" else np.abs(curr-prev).sum()
+        if diff < tolerance:
             converged = True
             break
 
-        prev_loglik = loglik
+        prev = curr
 
-    return counts, (curr_cycle, converged, loglik, loglikdiff)
+    return counts, (curr_cycle, converged, curr, diff)

irescue/main.py

@@ -178,6 +178,25 @@ def parseArguments():
             "(Default: %(default)s)."
         ),
     )
+    parser.add_argument(
+        "--keep-all-features",
+        action="store_true",
+        help=(
+            "Perform EM considering features supported by only 1 multimapping read "
+            "(Default: %(default)s)."
+        ),
+    )
+    parser.add_argument(
+        "--convergence-criterion",
+        type=str,
+        metavar="STR",
+        choices=["likelihood", "parameters"],
+        default="likelihood",
+        help=(
+            "Criterion to define convergence. "
+            "One of: likelihood, parameters (Default: %(default)s)."
+        ),
+    )
     parser.add_argument(
         "--dump-ec",
         action="store_true",
@@ -376,6 +395,8 @@ def main():
         feature_index,
         dirs["tmp"],
         args.no_umi,
+        args.keep_all_features,
+        args.convergence_criterion,
         args.dump_ec,
         args.max_iters,
         args.tolerance,