diff --git a/LICENSE b/LICENSE index 9618698..87748b0 100644 --- a/LICENSE +++ b/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 diff --git a/irescue/count.py b/irescue/count.py index 2f544f2..94274dc 100644 --- a/irescue/count.py +++ b/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, ) diff --git a/irescue/em.py b/irescue/em.py index 32c94e7..d42c859 100644 --- a/irescue/em.py +++ b/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) diff --git a/irescue/main.py b/irescue/main.py index 226197e..ceccf3c 100644 --- a/irescue/main.py +++ b/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,