Skip to content

Add Needleman-Wunsch distance as a faster alignment metric - #725

Open
felixpetschko wants to merge 22 commits into
scverse:mainfrom
felixpetschko:performance/needleman-wunsch-cpu
Open

Add Needleman-Wunsch distance as a faster alignment metric#725
felixpetschko wants to merge 22 commits into
scverse:mainfrom
felixpetschko:performance/needleman-wunsch-cpu

Conversation

@felixpetschko

Copy link
Copy Markdown
Contributor

Summary

This PR adds a Numba-optimized Needleman-Wunsch sequence distance metric via metric="needleman_wunsch". Conceptually, this metric is very similar to the existing alignment metric, with two main differences: it is much faster for large datasets, and it uses a linear gap model, where every gap position receives the same penalty.

The existing alignment metric allows gap openings and gap extensions to be penalized differently. However, this distinction is not used with the current default parameters anyway, where gap_open == gap_extend. If the alignment metric is configured such that gap_open == gap_extend == gap_penalty, both metrics should return equal results.

The reason for using a linear gap model is mostly practical. Supporting separate gap-open and gap-extension penalties requires handling three dynamic programming matrices instead of one for each sequence pair, which adds substantial computational overhead. For CDR3 sequences, which are usually quite short, I would generally expect users to choose parameters that allow only a limited number of gap positions. For example, in my test runs I used cutoff=10 and gap_penalty=4, which allows up to two gap positions. In this setting, distinguishing between two separate single-position gaps and one adjacent two-position gap is probably not worth the additional runtime. Therefore, I do not think the restriction to a single gap penalty is a major limitation.

The problem with the current alignment metric is that it is too slow for large datasets. The fastalignment metric improves runtime, but it can also suffer from performance limitations and does not always return exact results. Besides that, the current default parameters do not allow gaps (cutoff=10, gap_open=gap_extend=11). In that default setting, only equal-length sequences can fall below the cutoff, so the metric effectively computes alignment distances between equal-length sequences without making use of meaningful gap placement.

Therefore, I implemented the Needleman-Wunsch distance in a similar style to the Numba-optimized TCRdist CPU implementation. It supports BLOSUM62 by default and can also use TCRBLOSUM alpha/beta matrices through base_matrix="tcrblosum". With the improved performance, I was able to run the 8 million-cell Omniscope COVID dataset with 64 CPU cores using cutoff=10 and gap_penalty=4 within around 12 hours.

Overall, I think this metric is a useful addition because it provides an exact alignment-based distance that can still handle larger datasets. It is more flexible than the Hamming distance, but substantially faster than the existing alignment metric. In contrast to TCRdist, it is also not specific to TCR CDR3 sequences and can therefore be used for BCR analyses as well.

The most useful default values for gap_penalty and cutoff are open for discussion, since the performance improvements make less restrictive parameter choices feasible. It could also be discussed whether the existing alignment metric is still needed if its additional gap parameterization and support for additional Parasail substitution matrices are not required.

Main changes

  • Add NeedlemanWunschDistanceCalculator and expose it via metric="needleman_wunsch" in sequence_dist and ir_dist.
  • Support BLOSUM62 and TCRBLOSUM alpha/beta substitution matrices in a similar way as the TCRdist metric.
  • Implement a Numba-based Needleman-Wunsch dynamic programming algorithm with cutoff-based optimizations.
  • Add test cases.

@felixpetschko
felixpetschko requested a review from grst July 16, 2026 13:34
@codecov

codecov Bot commented Jul 16, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 46.34146% with 88 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.15%. Comparing base (b844b22) to head (31d0724).

Files with missing lines Patch % Lines
src/scirpy/ir_dist/metrics.py 45.67% 88 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #725      +/-   ##
==========================================
- Coverage   78.29%   77.15%   -1.14%     
==========================================
  Files          51       51              
  Lines        4607     4763     +156     
==========================================
+ Hits         3607     3675      +68     
- Misses       1000     1088      +88     
Files with missing lines Coverage Δ
src/scirpy/ir_dist/__init__.py 92.07% <100.00%> (+0.16%) ⬆️
src/scirpy/ir_dist/metrics.py 53.56% <45.67%> (-2.81%) ⬇️
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@felixpetschko
felixpetschko marked this pull request as ready for review July 16, 2026 17:59
@grst

grst commented Jul 16, 2026

Copy link
Copy Markdown
Collaborator

Hi @felixpetschko,

thanks for working on this! From a technical perspective this looks all great.

My main concern here is that we keep adding metrics without providing guidelines (and evidence for) which metric to use.

In contrast to TCRdist, it is also not specific to TCR CDR3 sequences and can therefore be used for BCR analyses as well.

In what way is TCRdist specific to TCR sequences that would prevent it from using it for BCR? Or rather in what way is the alignment distance superior for BCR? My understanding would be that the main difference is that TCRdist allows only for a single gap position, while alignment allows for multiple, but does this make a big difference in practice?

Also how does it compare to TCRdist in terms of speed? E.g. how long would TCRdist take on the omniscope dataset on the same hardware?

@felixpetschko

felixpetschko commented Jul 16, 2026

Copy link
Copy Markdown
Contributor Author

Hi @grst

My main concern here is that we keep adding metrics without providing guidelines (and evidence for) which metric to use.

My intention was rather to provide a faster implementation of the existing alignment distance metric than introducing a new metric. It is probably hard to provide real evidence for which metric to use. However, Needleman-Wunsch is a standard alignment algorithm and is also used in related sequence-analysis tools, so I think it is reasonable to use it.

In what way is TCRdist specific to TCR sequences that would prevent it from using it for BCR? Or rather in what way is the alignment distance superior for BCR? My understanding would be that the main difference is that TCRdist allows only for a single gap position, while alignment allows for multiple, but does this make a big difference in practice?

Actually, I was mainly focusing on performance, and I do not have proof for which metric is better in which case. My reasoning was that TCRdist's approach, with trimming from the N and C terminus and a single gap region, seems more targeted towards the TCR model. Needleman-Wunsch might be easier to justify for BCR CDR3 comparisons because it does not impose TCRdist's trimming and single-gap-region assumptions, and can handle length differences with a general global alignment.

Also how does it compare to TCRdist in terms of speed? E.g. how long would TCRdist take on the omniscope dataset on the same hardware?

TCRdist can run the full Omniscope COVID dataset on the same hardware in around 2.3 hours with default parameters, which makes it around 5 times faster than Needleman-Wunsch in my test run. The main reason is that, with Scirpy's default parameters, the gap position is computed by a formula and it is not necessary to try different gap positions. In contrast, Needleman-Wunsch computes the optimal global alignment that minimizes the distance.

@grst

grst commented Jul 17, 2026

Copy link
Copy Markdown
Collaborator

Do you think we could implement this without any user-facing changes?
I.e. integrate this into the AlignmentDistanceCalculator and fall back to the parasail implementation when gap_open != gap_extend?

(would be interesting to know if in the history of scirpy anyone has ever changed these default parameters. I'd guess not).

@felixpetschko

Copy link
Copy Markdown
Contributor Author

Do you think we could implement this without any user-facing changes?

Yes, I will do that!

(would be interesting to know if in the history of scirpy anyone has ever changed these default parameters. I'd guess not).

However, I think we should definitely change the default params and set them in a way such that gaps are allowed. Otherwise there is no alignment done at all. I would allow at least 2 gaps such that it's worth to even run the dynamic programming alignment algorithm.

@grst

grst commented Jul 17, 2026

Copy link
Copy Markdown
Collaborator

Just changing defaults is also not very good practice... so we'd at least have to warn about it.
But maybe it is better to keep needleman_wunsch as a separate alignment metric and deprecate "alignment" and "fastalignment" distances.

We could complement it with a "metrics" guide in the documentation that explains the pros/cons and usecases of the different metrics.

@felixpetschko

Copy link
Copy Markdown
Contributor Author

Alright, then let's do it like that 👍

@felixpetschko

Copy link
Copy Markdown
Contributor Author

Now I marked alignment and fastalignment as deprecated in favor of needleman_wunsch and added a guide for choosing a sequence distance metric via #728.

@grst grst left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have a few more comments, mostly on improving the organization of the metrics submodule.

Comment on lines +554 to 557
results_iter = joblib.Parallel(return_as="generator")(delayed_jobs)
results_iter = tqdm(results_iter, total=len(delayed_jobs), desc="Computing distance blocks")
results = list(results_iter)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This would fail on joblib backends that do not support return_as="generator", e.g. dask.

See scirpy.utils._parallelize_with_joblib for a helper that addresses this.

_metric_mat = _gpu_hamming_mat


PARASAIL_AA_ALPHABET = "ARNDCQEGHILKMFPSTWYVBZX"

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd move all these definitions to the top of the file, or maybe even better, a separate submodule within the ir_dist package.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, maybe worth renaming this simply to AA_ALPHABET? Or is this still specific to parasail in any way?

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

... or maybe make them dataclasses, that match together alphabet and substitution matrix

@dataclass
class SubstitutionMatrix:
    alphabet: str
    matrix: np.ndarray

BLOSUM62 = SubstitutionMatrix(alphabet = "ARN...", matrix = np.array([...]))
TCRBLOSUM_ALPHA = SubstitutionMatrix(...)

Comment on lines +1339 to +1344
parasail_aa_alphabet = PARASAIL_AA_ALPHABET
parasail_aa_alphabet_with_unknown = PARASAIL_AA_ALPHABET_WITH_UNKNOWN
matrix_alphabet = CANONICAL_AA_ALPHABET
blosum62_substitution_matrix = BLOSUM62_SUBSTITUTION_MATRIX
tcrblosum_alpha_substitution_matrix = TCRBLOSUM_ALPHA_SUBSTITUTION_MATRIX
tcrblosum_beta_substitution_matrix = TCRBLOSUM_BETA_SUBSTITUTION_MATRIX

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there any reason for defining these as class variables instead of directly referencing the constants?

Comment on lines +1611 to +1617
parasail_aa_alphabet = PARASAIL_AA_ALPHABET
parasail_aa_alphabet_with_unknown = PARASAIL_AA_ALPHABET_WITH_UNKNOWN
tcrblosum_matrix_alphabet = CANONICAL_AA_ALPHABET
blosum62_substitution_matrix = BLOSUM62_SUBSTITUTION_MATRIX
blosum62_with_ambiguous_substitution_matrix = BLOSUM62_WITH_AMBIGUOUS_SUBSTITUTION_MATRIX
tcrblosum_alpha_substitution_matrix = TCRBLOSUM_ALPHA_SUBSTITUTION_MATRIX
tcrblosum_beta_substitution_matrix = TCRBLOSUM_BETA_SUBSTITUTION_MATRIX

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Again, wouldn't it be easier to read if the constants were directly used everywhere?

Comment on lines +1656 to +1668
def _make_numba_substitution_matrix(self, substitution_matrix: np.ndarray, matrix_alphabet: str) -> np.ndarray:
score_matrix = np.zeros(
(len(self.parasail_aa_alphabet_with_unknown), len(self.parasail_aa_alphabet_with_unknown)),
dtype=np.int32,
)
if substitution_matrix.shape != (len(matrix_alphabet), len(matrix_alphabet)):
raise ValueError("`substitution_matrix` must be square and match `matrix_alphabet`.")
for i, aa1 in enumerate(matrix_alphabet):
for j, aa2 in enumerate(matrix_alphabet):
score_matrix[self.parasail_aa_alphabet.index(aa1), self.parasail_aa_alphabet.index(aa2)] = (
substitution_matrix[i, j]
)
return score_matrix

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could this become a metric-agnostic helper function? I think we already have similar code in other metrics...

"""\
FastAlignmentDistanceCalculator achieves (depending on the settings) identical results
at a higher speed.
If `gap_open == gap_extend`, use NeedlemanWunschDistanceCalculator instead.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
If `gap_open == gap_extend`, use NeedlemanWunschDistanceCalculator instead.
If `gap_open == gap_extend` (the default), use NeedlemanWunschDistanceCalculator instead, which provides identical results while being much faster. If you actually have a use-case for affine gap penalties, please let us know by opening an issue on GitHub.


@deprecated(
"""\
If `gap_open == gap_extend`, use NeedlemanWunschDistanceCalculator instead.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
If `gap_open == gap_extend`, use NeedlemanWunschDistanceCalculator instead.
If `gap_open == gap_extend` (the default), use NeedlemanWunschDistanceCalculator instead, which provides identical results while being much faster. If you actually have a use-case for affine gap penalties, please let us know by opening an issue on GitHub.



def test_needleman_wunsch_reference():
# test needleman-wunsch against a precomputed linear-gap alignment reference

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How has this been derived? Parasail?

@grst grst moved this to On Hold in scirpy-dev Jul 27, 2026
@grst grst added this to scirpy-dev Jul 27, 2026
@grst grst moved this from On Hold to In progress in scirpy-dev Jul 27, 2026
@grst

grst commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator

regarding deprecations, take a look at #735 please that switches to the decorators provided by scverse-misc.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

Status: In progress

Development

Successfully merging this pull request may close these issues.

2 participants