diff --git a/CHANGELOG.md b/CHANGELOG.md index 59232c228..cb0441d0a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,6 +27,9 @@ and this project adheres to [Semantic Versioning][]. ### Features + - Add a linear-gap Needleman-Wunsch distance metric for CDR3 amino acid sequences via + `metric="needleman_wunsch"`. The `alignment` and `fastalignment` metrics are now + deprecated. When `gap_open == gap_extend`, use `needleman_wunsch` instead. - Add support for TCRBLOSUM alpha/beta substitution matrices in the `tcrdist` distance metric via `base_matrix="tcrblosum"`, and allow configuring the substitution-to-distance cap with `distance_cap`. diff --git a/docs/api.rst b/docs/api.rst index 809d2864d..8e8ed3172 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -311,3 +311,4 @@ distance metrics ir_dist.metrics.AlignmentDistanceCalculator ir_dist.metrics.FastAlignmentDistanceCalculator ir_dist.metrics.TCRdistDistanceCalculator + ir_dist.metrics.NeedlemanWunschDistanceCalculator diff --git a/src/scirpy/ir_dist/__init__.py b/src/scirpy/ir_dist/__init__.py index 4231e843d..9f9e7811c 100644 --- a/src/scirpy/ir_dist/__init__.py +++ b/src/scirpy/ir_dist/__init__.py @@ -36,9 +36,10 @@ def IrNeighbors(*args, **kwargs): "identity", "levenshtein", "hamming", - "gpu_haming", + "gpu_hamming", "normalized_hamming", "tcrdist", + "needleman_wunsch", ] | metrics.DistanceCalculator ) @@ -55,6 +56,10 @@ def IrNeighbors(*args, **kwargs): Uses the BLOSUM62 substitution matrix by default. TCRBLOSUM alpha/beta substitution matrices (:cite:`TCRBLOSUM`) can be selected with `base_matrix="tcrblosum"`. See :class:`~scirpy.ir_dist.metrics.TCRdistDistanceCalculator`. + * `needleman_wunsch` -- Distance based on linear-gap Needleman-Wunsch global alignment. + Uses the BLOSUM62 substitution matrix by default. TCRBLOSUM alpha/beta substitution matrices + (:cite:`TCRBLOSUM`) can be selected with `base_matrix="tcrblosum"`. + See :class:`~scirpy.ir_dist.metrics.NeedlemanWunschDistanceCalculator`. * `hamming` -- Hamming distance for CDR3 sequences of equal length. See :class:`~scirpy.ir_dist.metrics.HammingDistanceCalculator`. * `gpu_hamming` -- Hamming distance for CDR3 sequences of equal length calculated with a GPU. @@ -62,10 +67,12 @@ def IrNeighbors(*args, **kwargs): * `normalized_hamming` -- Normalized Hamming distance (in percent) for CDR3 sequences of equal length. See :class:`~scirpy.ir_dist.metrics.HammingDistanceCalculator`. * `alignment` -- Distance based on pairwise sequence alignments using the - BLOSUM62 matrix. This option is incompatible with nucleotide sequences. + BLOSUM62 matrix. Deprecated; if `gap_open == gap_extend`, use `needleman_wunsch` instead. + This option is incompatible with nucleotide sequences. See :class:`~scirpy.ir_dist.metrics.FastAlignmentDistanceCalculator`. * `fastalignment` -- Distance based on pairwise sequence alignments using the BLOSUM62 matrix. Faster implementation of `alignment` with some loss. + Deprecated; if `gap_open == gap_extend`, use `needleman_wunsch` instead. This option is incompatible with nucleotide sequences. See :class:`~scirpy.ir_dist.metrics.FastAlignmentDistanceCalculator`. * any instance of :class:`~scirpy.ir_dist.metrics.DistanceCalculator`. @@ -76,7 +83,8 @@ def IrNeighbors(*args, **kwargs): All distances `> cutoff` will be replaced by `0` and eliminated from the sparse matrix. A sensible cutoff depends on the distance metric, you can find information in the corresponding docs. If set to `None`, the cutoff - will be `10` for the `alignment` and `fastalignment` metric, and `2` for `levenshtein` and `hamming`. + will be `10` for the `alignment`, `fastalignment`, and `needleman_wunsch` metric, + and `2` for `levenshtein` and `hamming`. For the identity metric, the cutoff is ignored and always set to `0`. """ @@ -119,6 +127,8 @@ def _get_distance_calculator( dist_calc = metrics.GPUHammingDistanceCalculator(**kwargs) elif metric == "tcrdist": dist_calc = metrics.TCRdistDistanceCalculator(n_jobs=n_jobs, chain_type=chain_type, **kwargs) + elif metric == "needleman_wunsch": + dist_calc = metrics.NeedlemanWunschDistanceCalculator(n_jobs=n_jobs, chain_type=chain_type, **kwargs) else: raise ValueError("Invalid distance metric.") diff --git a/src/scirpy/ir_dist/_substitution_matrices.py b/src/scirpy/ir_dist/_substitution_matrices.py new file mode 100644 index 000000000..cd47783cf --- /dev/null +++ b/src/scirpy/ir_dist/_substitution_matrices.py @@ -0,0 +1,223 @@ +from dataclasses import dataclass + +import numpy as np + + +@dataclass(frozen=True, slots=True) +class SubstitutionMatrix: + alphabet: str + matrix: np.ndarray + + +CANONICAL_AA_ALPHABET = "ARNDCQEGHILKMFPSTWYV" +AA_ALPHABET_WITH_AMBIGUOUS = f"{CANONICAL_AA_ALPHABET}BZX" +AA_ALPHABET_WITH_UNKNOWN = f"{AA_ALPHABET_WITH_AMBIGUOUS}*" + + +def _map_matrix_to_alphabet( + matrix: np.ndarray, + source_alphabet: str, + target_alphabet: str, +) -> np.ndarray: + """Map a matrix from the source alphabet to the target alphabet. + + Rows and columns for characters that occur only in `target_alphabet` are + initialized to zero. + + Parameters + ---------- + matrix: + Square matrix in the order specified by `source_alphabet`. + source_alphabet: + Alphabet describing the rows and columns of `matrix`. + target_alphabet: + Alphabet describing the rows and columns of the returned matrix. It must + contain every character from `source_alphabet`. + + Returns + ------- + mapped_matrix: + Matrix in the order specified by `target_alphabet`. Entries for target + characters absent from `source_alphabet` are initialized to zero. + """ + expected_shape = (len(source_alphabet), len(source_alphabet)) + if matrix.shape != expected_shape: + raise ValueError(f"`matrix` must have shape {expected_shape} to match `source_alphabet`.") + + if len(set(source_alphabet)) != len(source_alphabet): + raise ValueError("`source_alphabet` must not contain duplicate characters.") + if len(set(target_alphabet)) != len(target_alphabet): + raise ValueError("`target_alphabet` must not contain duplicate characters.") + + target_indices = {character: index for index, character in enumerate(target_alphabet)} + missing_characters = set(source_alphabet) - set(target_alphabet) + if missing_characters: + raise ValueError( + f"`target_alphabet` is missing characters from `source_alphabet`: {sorted(missing_characters)}" + ) + + mapped_indices = [target_indices[character] for character in source_alphabet] + mapped_matrix = np.zeros((len(target_alphabet), len(target_alphabet)), dtype=matrix.dtype) + mapped_matrix[np.ix_(mapped_indices, mapped_indices)] = matrix + return mapped_matrix + + +def _substitution_to_distance_matrix( + substitution_matrix: np.ndarray, + alphabet: str = AA_ALPHABET_WITH_UNKNOWN, + matrix_alphabet: str = CANONICAL_AA_ALPHABET, + distance_cap: int | None = 4, + distance_offset: int = 4, +) -> np.ndarray: + """Create a distance lookup matrix from an amino-acid substitution matrix. + + Parameters + ---------- + substitution_matrix: + Amino-acid substitution matrix in the order specified by `matrix_alphabet`. + alphabet: + Alphabet describing the rows and columns of the returned distance matrix. + matrix_alphabet: + Alphabet describing the rows and columns of `substitution_matrix`. + distance_cap: + Maximum distance assigned to a mismatch. If `None`, mismatch distances are uncapped. + distance_offset: + Offset from which the substitution score is subtracted. + + Returns + ------- + distance_matrix: + Distance lookup matrix in the order specified by `alphabet`. + """ + expected_shape = (len(matrix_alphabet), len(matrix_alphabet)) + if substitution_matrix.shape != expected_shape: + raise ValueError(f"`substitution_matrix` must have shape {expected_shape} to match `matrix_alphabet`.") + + distance_matrix = np.zeros(expected_shape, dtype=np.int32) + for i, aa1 in enumerate(matrix_alphabet): + for j, aa2 in enumerate(matrix_alphabet): + distance = 0 if aa1 == aa2 else distance_offset - substitution_matrix[i, j] + if distance_cap is not None: + distance = min(distance_cap, distance) + distance_matrix[i, j] = distance + return _map_matrix_to_alphabet(distance_matrix, matrix_alphabet, alphabet) + + +# fmt: off +_BLOSUM62_MATRIX = np.array( + [ + # A R N D C Q E G H I L K M F P S T W Y V + [ 4, -1, -2, -2, 0, -1, -1, 0, -2, -1, -1, -1, -1, -2, -1, 1, 0, -3, -2, 0], # A + [-1, 5, 0, -2, -3, 1, 0, -2, 0, -3, -2, 2, -1, -3, -2, -1, -1, -3, -2, -3], # R + [-2, 0, 6, 1, -3, 0, 0, 0, 1, -3, -3, 0, -2, -3, -2, 1, 0, -4, -2, -3], # N + [-2, -2, 1, 6, -3, 0, 2, -1, -1, -3, -4, -1, -3, -3, -1, 0, -1, -4, -3, -3], # D + [ 0, -3, -3, -3, 9, -3, -4, -3, -3, -1, -1, -3, -1, -2, -3, -1, -1, -2, -2, -1], # C + [-1, 1, 0, 0, -3, 5, 2, -2, 0, -3, -2, 1, 0, -3, -1, 0, -1, -2, -1, -2], # Q + [-1, 0, 0, 2, -4, 2, 5, -2, 0, -3, -3, 1, -2, -3, -1, 0, -1, -3, -2, -2], # E + [ 0, -2, 0, -1, -3, -2, -2, 6, -2, -4, -4, -2, -3, -3, -2, 0, -2, -2, -3, -3], # G + [-2, 0, 1, -1, -3, 0, 0, -2, 8, -3, -3, -1, -2, -1, -2, -1, -2, -2, 2, -3], # H + [-1, -3, -3, -3, -1, -3, -3, -4, -3, 4, 2, -3, 1, 0, -3, -2, -1, -3, -1, 3], # I + [-1, -2, -3, -4, -1, -2, -3, -4, -3, 2, 4, -2, 2, 0, -3, -2, -1, -2, -1, 1], # L + [-1, 2, 0, -1, -3, 1, 1, -2, -1, -3, -2, 5, -1, -3, -1, 0, -1, -3, -2, -2], # K + [-1, -1, -2, -3, -1, 0, -2, -3, -2, 1, 2, -1, 5, 0, -2, -1, -1, -1, -1, 1], # M + [-2, -3, -3, -3, -2, -3, -3, -3, -1, 0, 0, -3, 0, 6, -4, -2, -2, 1, 3, -1], # F + [-1, -2, -2, -1, -3, -1, -1, -2, -2, -3, -3, -1, -2, -4, 7, -1, -1, -4, -3, -2], # P + [ 1, -1, 1, 0, -1, 0, 0, 0, -1, -2, -2, 0, -1, -2, -1, 4, 1, -3, -2, -2], # S + [ 0, -1, 0, -1, -1, -1, -1, -2, -2, -1, -1, -1, -1, -2, -1, 1, 5, -2, -2, 0], # T + [-3, -3, -4, -4, -2, -2, -3, -2, -2, -3, -2, -3, -1, 1, -4, -3, -2, 11, 2, -3], # W + [-2, -2, -2, -3, -2, -1, -2, -3, 2, -1, -1, -2, -1, 3, -3, -2, -2, 2, 7, -1], # Y + [ 0, -3, -3, -3, -1, -2, -2, -3, -3, 3, 1, -2, 1, -1, -2, -2, 0, -3, -1, 4], # V + ], + dtype=np.int32, +) +_TCRBLOSUM_ALPHA_MATRIX = np.array( + [ + # A R N D C Q E G H I L K M F P S T W Y V + [ 2, -1, -1, -1, 0, 0, 0, 0, 0, -1, -1, -1, -1, -1, 0, 0, -1, 0, -1, 0], # A + [-1, 1, 0, 0, 1, 0, 0, 0, 0, 0, -1, 0, 0, 0, 0, 0, 0, 0, 0, -1], # R + [-1, 0, 1, 0, 0, 0, 0, 0, 0, -1, -2, 1, 0, 0, 0, 0, 0, 0, 0, -2], # N + [-1, 0, 0, 1, -5, 0, 0, 0, 0, -1, -2, 0, 0, 0, 0, 0, 0, 0, 0, -1], # D + [ 0, 1, 0, -5, 2, -4, -4, 0, -2, -5, 0, -5, -4, -4, -4, 0, -6, -2, -5, 0], # C + [ 0, 0, 0, 0, -4, 2, 0, 0, 0, -1, -2, 1, 0, 0, 0, 0, 0, 0, 0, -2], # Q + [ 0, 0, 0, 0, -4, 0, 1, 0, 1, -1, 0, -1, 0, 0, 0, 0, 0, 0, 0, 0], # E + [ 0, 0, 0, 0, 0, 0, 0, 1, 0, -2, -1, -1, 0, 0, 0, 0, 0, 0, 0, 0], # G + [ 0, 0, 0, 0, -2, 0, 1, 0, 2, 0, 0, -1, 0, 0, 1, 0, 0, 1, 0, 0], # H + [-1, 0, -1, -1, -5, -1, -1, -2, 0, 3, 0, -1, 0, 0, 0, 0, 1, -1, 0, 0], # I + [-1, -1, -2, -2, 0, -2, 0, -1, 0, 0, 2, -4, 0, 1, 0, -1, -1, -1, 0, 0], # L + [-1, 0, 1, 0, -5, 1, -1, -1, -1, -1, -4, 3, 0, -3, 0, -2, -1, -2, -4, -3], # K + [-1, 0, 0, 0, -4, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, -1, 0], # M + [-1, 0, 0, 0, -4, 0, 0, 0, 0, 0, 1, -3, 0, 1, 0, 0, 0, 0, 0, 0], # F + [ 0, 0, 0, 0, -4, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0], # P + [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -1, -2, 0, 0, 0, 1, 0, 0, 0, -1], # S + [-1, 0, 0, 0, -6, 0, 0, 0, 0, 1, -1, -1, 0, 0, 0, 0, 1, 0, 0, 0], # T + [ 0, 0, 0, 0, -2, 0, 0, 0, 1, -1, -1, -2, 0, 0, 0, 0, 0, 2, 0, -1], # W + [-1, 0, 0, 0, -5, 0, 0, 0, 0, 0, 0, -4, -1, 0, 0, 0, 0, 0, 1, -1], # Y + [ 0, -1, -2, -1, 0, -2, 0, 0, 0, 0, 0, -3, 0, 0, 0, -1, 0, -1, -1, 1], # V + ], + dtype=np.int32, +) +_TCRBLOSUM_BETA_MATRIX = np.array( + [ + # A R N D C Q E G H I L K M F P S T W Y V + [ 0, 0, 0, 0, -5, 0, -1, 0, 0, 0, 0, 0, 0, -1, 0, 0, 0, 0, -1, 0], # A + [ 0, 2, 0, 0, -4, -1, -1, 0, 0, 0, 0, 0, 0, -1, 0, 0, 0, 0, -1, 0], # R + [ 0, 0, 1, 1, -4, 0, 0, 0, 0, 0, -1, 0, 0, -1, 0, -1, 0, 0, 0, 0], # N + [ 0, 0, 1, 1, -4, 0, 0, 0, 0, 0, 0, 0, 0, -1, 0, -1, 0, 0, 0, 0], # D + [-5, -4, -4, -4, 2, -6, -5, 0, -3, -3, -5, -2, -1, -5, -4, 0, -5, -2, -5, -4], # C + [ 0, -1, 0, 0, -6, 2, -1, -1, -1, 0, 1, -1, 0, -2, -1, -2, -1, 0, 0, -1], # Q + [-1, -1, 0, 0, -5, -1, 2, 0, -1, 0, -1, 1, 0, -2, 0, -2, 1, 0, -1, 0], # E + [ 0, 0, 0, 0, 0, -1, 0, 0, 0, 0, 0, 0, 0, -1, 0, 0, 0, 0, -1, 0], # G + [ 0, 0, 0, 0, -3, -1, -1, 0, 2, 0, 0, -1, 0, 2, 0, -1, 0, 0, 1, 0], # H + [ 0, 0, 0, 0, -3, 0, 0, 0, 0, 2, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0], # I + [ 0, 0, -1, 0, -5, 1, -1, 0, 0, 0, 1, 0, 0, 0, 0, -1, 0, 0, 0, 0], # L + [ 0, 0, 0, 0, -2, -1, 1, 0, -1, 0, 0, 1, 0, -1, 0, 0, 0, 0, -1, 0], # K + [ 0, 0, 0, 0, -1, 0, 0, 0, 0, 2, 0, 0, 2, 0, 0, 0, 0, 0, -1, 0], # M + [-1, -1, -1, -1, -5, -2, -2, -1, 2, 0, 0, -1, 0, 2, 0, -2, 0, 0, 2, -1], # F + [ 0, 0, 0, 0, -4, -1, 0, 0, 0, 0, 0, 0, 0, 0, 1, -1, 0, 0, -1, 0], # P + [ 0, 0, -1, -1, 0, -2, -2, 0, -1, 0, -1, 0, 0, -2, -1, 1, 0, 0, -2, 0], # S + [ 0, 0, 0, 0, -5, -1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], # T + [ 0, 0, 0, 0, -2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0], # W + [-1, -1, 0, 0, -5, 0, -1, -1, 1, 0, 0, -1, -1, 2, -1, -2, 0, 0, 2, -1], # Y + [ 0, 0, 0, 0, -4, -1, 0, 0, 0, 0, 0, 0, 0, -1, 0, 0, 0, 0, -1, 0], # V + ], + dtype=np.int32, +) +# fmt: on + + +# fmt: off +_BLOSUM62_WITH_AMBIGUOUS_MATRIX = np.array( + [ + # A R N D C Q E G H I L K M F P S T W Y V B Z X + [ 4, -1, -2, -2, 0, -1, -1, 0, -2, -1, -1, -1, -1, -2, -1, 1, 0, -3, -2, 0, -2, -1, 0], # A + [-1, 5, 0, -2, -3, 1, 0, -2, 0, -3, -2, 2, -1, -3, -2, -1, -1, -3, -2, -3, -1, 0, -1], # R + [-2, 0, 6, 1, -3, 0, 0, 0, 1, -3, -3, 0, -2, -3, -2, 1, 0, -4, -2, -3, 3, 0, -1], # N + [-2, -2, 1, 6, -3, 0, 2, -1, -1, -3, -4, -1, -3, -3, -1, 0, -1, -4, -3, -3, 4, 1, -1], # D + [ 0, -3, -3, -3, 9, -3, -4, -3, -3, -1, -1, -3, -1, -2, -3, -1, -1, -2, -2, -1, -3, -3, -2], # C + [-1, 1, 0, 0, -3, 5, 2, -2, 0, -3, -2, 1, 0, -3, -1, 0, -1, -2, -1, -2, 0, 3, -1], # Q + [-1, 0, 0, 2, -4, 2, 5, -2, 0, -3, -3, 1, -2, -3, -1, 0, -1, -3, -2, -2, 1, 4, -1], # E + [ 0, -2, 0, -1, -3, -2, -2, 6, -2, -4, -4, -2, -3, -3, -2, 0, -2, -2, -3, -3, -1, -2, -1], # G + [-2, 0, 1, -1, -3, 0, 0, -2, 8, -3, -3, -1, -2, -1, -2, -1, -2, -2, 2, -3, 0, 0, -1], # H + [-1, -3, -3, -3, -1, -3, -3, -4, -3, 4, 2, -3, 1, 0, -3, -2, -1, -3, -1, 3, -3, -3, -1], # I + [-1, -2, -3, -4, -1, -2, -3, -4, -3, 2, 4, -2, 2, 0, -3, -2, -1, -2, -1, 1, -4, -3, -1], # L + [-1, 2, 0, -1, -3, 1, 1, -2, -1, -3, -2, 5, -1, -3, -1, 0, -1, -3, -2, -2, 0, 1, -1], # K + [-1, -1, -2, -3, -1, 0, -2, -3, -2, 1, 2, -1, 5, 0, -2, -1, -1, -1, -1, 1, -3, -1, -1], # M + [-2, -3, -3, -3, -2, -3, -3, -3, -1, 0, 0, -3, 0, 6, -4, -2, -2, 1, 3, -1, -3, -3, -1], # F + [-1, -2, -2, -1, -3, -1, -1, -2, -2, -3, -3, -1, -2, -4, 7, -1, -1, -4, -3, -2, -2, -1, -2], # P + [ 1, -1, 1, 0, -1, 0, 0, 0, -1, -2, -2, 0, -1, -2, -1, 4, 1, -3, -2, -2, 0, 0, 0], # S + [ 0, -1, 0, -1, -1, -1, -1, -2, -2, -1, -1, -1, -1, -2, -1, 1, 5, -2, -2, 0, -1, -1, 0], # T + [-3, -3, -4, -4, -2, -2, -3, -2, -2, -3, -2, -3, -1, 1, -4, -3, -2, 11, 2, -3, -4, -3, -2], # W + [-2, -2, -2, -3, -2, -1, -2, -3, 2, -1, -1, -2, -1, 3, -3, -2, -2, 2, 7, -1, -3, -2, -1], # Y + [ 0, -3, -3, -3, -1, -2, -2, -3, -3, 3, 1, -2, 1, -1, -2, -2, 0, -3, -1, 4, -3, -2, -1], # V + [-2, -1, 3, 4, -3, 0, 1, -1, 0, -3, -4, 0, -3, -3, -2, 0, -1, -4, -3, -3, 4, 1, -1], # B + [-1, 0, 0, 1, -3, 3, 4, -2, 0, -3, -3, 1, -1, -3, -1, 0, -1, -3, -2, -2, 1, 4, -1], # Z + [ 0, -1, -1, -1, -2, -1, -1, -1, -1, -1, -1, -1, -1, -1, -2, 0, 0, -2, -1, -1, -1, -1, -1], # X + ], + dtype=np.int32, +) +# fmt: on + +BLOSUM62 = SubstitutionMatrix(CANONICAL_AA_ALPHABET, _BLOSUM62_MATRIX) +BLOSUM62_WITH_AMBIGUOUS = SubstitutionMatrix(AA_ALPHABET_WITH_AMBIGUOUS, _BLOSUM62_WITH_AMBIGUOUS_MATRIX) +TCRBLOSUM_ALPHA = SubstitutionMatrix(CANONICAL_AA_ALPHABET, _TCRBLOSUM_ALPHA_MATRIX) +TCRBLOSUM_BETA = SubstitutionMatrix(CANONICAL_AA_ALPHABET, _TCRBLOSUM_BETA_MATRIX) diff --git a/src/scirpy/ir_dist/metrics.py b/src/scirpy/ir_dist/metrics.py index f6f7f2252..5dc8d5ddb 100644 --- a/src/scirpy/ir_dist/metrics.py +++ b/src/scirpy/ir_dist/metrics.py @@ -16,6 +16,17 @@ from scirpy.util import _doc_params, _get_usable_cpus, _parallelize_with_joblib +from ._substitution_matrices import ( + AA_ALPHABET_WITH_AMBIGUOUS, + AA_ALPHABET_WITH_UNKNOWN, + BLOSUM62, + BLOSUM62_WITH_AMBIGUOUS, + TCRBLOSUM_ALPHA, + TCRBLOSUM_BETA, + _map_matrix_to_alphabet, + _substitution_to_distance_matrix, +) + #: Deprecation of the object-level `block_size` parameter, shared by all `ParallelDistanceCalculator`s. _deprecation_block_size = Deprecation( "0.15.0", "The block size is now set in the `calc_dist_mat` function instead of the object level." @@ -345,43 +356,8 @@ def _compute_block(self, seqs1, seqs2, origin): return result -def _substitution_to_distance_matrix( - substitution_matrix: np.ndarray, - alphabet: str = "ARNDCQEGHILKMFPSTWYVBZX*", - matrix_alphabet: str = "ARNDCQEGHILKMFPSTWYV", - distance_cap: int | None = 4, - distance_offset: int = 4, -) -> np.ndarray: - """Creates a numba compatible distance matrix from a substitution matrix. - - Parameters - ---------- - substitution_matrix: - Amino-acid substitution matrix in the order specified by `matrix_alphabet`. - distance_cap: - Maximum distance assigned to a mismatch. If `None`, mismatch distances are uncapped. - distance_offset: - Offset from which the substitution score is subtracted. - - Returns - ------- - distance_matrix: - distance lookup matrix - """ - dm = np.zeros((len(alphabet), len(alphabet)), 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): - d = 0 if aa1 == aa2 else distance_offset - substitution_matrix[i, j] - if distance_cap is not None: - d = min(distance_cap, d) - dm[alphabet.index(aa1), alphabet.index(aa2)] = d - return dm - - def _seqs2mat( - seqs: Sequence[str], alphabet: str = "ARNDCQEGHILKMFPSTWYVBZX", max_len: None | int = None + seqs: Sequence[str], alphabet: str = AA_ALPHABET_WITH_AMBIGUOUS, max_len: None | int = None ) -> tuple[np.ndarray, np.ndarray]: """Convert a collection of gene sequences into a numpy matrix of integers for fast comparison. @@ -552,7 +528,7 @@ def calc_dist_mat(self, seqs: Sequence[str], seqs2: Sequence[str] | None = None) arguments = [(split_seqs[x], seqs2, is_symmetric, start_columns[x]) for x in range(self.n_blocks)] delayed_jobs = [joblib.delayed(self._calc_dist_mat_block)(*args) for args in arguments] - results = joblib.Parallel(return_as="list")(delayed_jobs) + results = list(_parallelize_with_joblib(delayed_jobs, total=len(delayed_jobs))) block_matrices_csr, block_row_mins = zip(*results, strict=False) distance_matrix_csr = scipy.sparse.vstack(block_matrices_csr) @@ -1193,14 +1169,14 @@ class TCRdistDistanceCalculator(_MetricDistanceCalculator): If True, insert gaps at a fixed position after the cysteine residue statring the CDR3 (typically position 6). If False, find the "optimal" position for inserting the gaps to make up the difference in length cutoff: - Will eleminate distances > cutoff to make efficient + Will eliminate distances > cutoff to make efficient use of sparse matrices. n_jobs: Number of numba parallel threads to use for the pairwise distance calculation n_blocks: Number of joblib delayed objects (blocks to compute) given to joblib.Parallel histogram: - Determines whether a nearest neighbor histogram should be created + Determines whether a nearest neighbor histogram should be created. Not implemented for this metric base_matrix: Amino acid substitution matrix used by TCRdist. `"blosum62"` uses the original BLOSUM62 substitution matrix, while `"tcrblosum"` uses TCRBLOSUM substitution @@ -1217,90 +1193,6 @@ class TCRdistDistanceCalculator(_MetricDistanceCalculator): is set automatically and should not be provided. """ - parasail_aa_alphabet = "ARNDCQEGHILKMFPSTWYVBZX" - parasail_aa_alphabet_with_unknown = "ARNDCQEGHILKMFPSTWYVBZX*" - # fmt: off - matrix_alphabet = "ARNDCQEGHILKMFPSTWYV" - blosum62_substitution_matrix = np.array( - [ - # A R N D C Q E G H I L K M F P S T W Y V - [ 4, -1, -2, -2, 0, -1, -1, 0, -2, -1, -1, -1, -1, -2, -1, 1, 0, -3, -2, 0], # A - [-1, 5, 0, -2, -3, 1, 0, -2, 0, -3, -2, 2, -1, -3, -2, -1, -1, -3, -2, -3], # R - [-2, 0, 6, 1, -3, 0, 0, 0, 1, -3, -3, 0, -2, -3, -2, 1, 0, -4, -2, -3], # N - [-2, -2, 1, 6, -3, 0, 2, -1, -1, -3, -4, -1, -3, -3, -1, 0, -1, -4, -3, -3], # D - [ 0, -3, -3, -3, 9, -3, -4, -3, -3, -1, -1, -3, -1, -2, -3, -1, -1, -2, -2, -1], # C - [-1, 1, 0, 0, -3, 5, 2, -2, 0, -3, -2, 1, 0, -3, -1, 0, -1, -2, -1, -2], # Q - [-1, 0, 0, 2, -4, 2, 5, -2, 0, -3, -3, 1, -2, -3, -1, 0, -1, -3, -2, -2], # E - [ 0, -2, 0, -1, -3, -2, -2, 6, -2, -4, -4, -2, -3, -3, -2, 0, -2, -2, -3, -3], # G - [-2, 0, 1, -1, -3, 0, 0, -2, 8, -3, -3, -1, -2, -1, -2, -1, -2, -2, 2, -3], # H - [-1, -3, -3, -3, -1, -3, -3, -4, -3, 4, 2, -3, 1, 0, -3, -2, -1, -3, -1, 3], # I - [-1, -2, -3, -4, -1, -2, -3, -4, -3, 2, 4, -2, 2, 0, -3, -2, -1, -2, -1, 1], # L - [-1, 2, 0, -1, -3, 1, 1, -2, -1, -3, -2, 5, -1, -3, -1, 0, -1, -3, -2, -2], # K - [-1, -1, -2, -3, -1, 0, -2, -3, -2, 1, 2, -1, 5, 0, -2, -1, -1, -1, -1, 1], # M - [-2, -3, -3, -3, -2, -3, -3, -3, -1, 0, 0, -3, 0, 6, -4, -2, -2, 1, 3, -1], # F - [-1, -2, -2, -1, -3, -1, -1, -2, -2, -3, -3, -1, -2, -4, 7, -1, -1, -4, -3, -2], # P - [ 1, -1, 1, 0, -1, 0, 0, 0, -1, -2, -2, 0, -1, -2, -1, 4, 1, -3, -2, -2], # S - [ 0, -1, 0, -1, -1, -1, -1, -2, -2, -1, -1, -1, -1, -2, -1, 1, 5, -2, -2, 0], # T - [-3, -3, -4, -4, -2, -2, -3, -2, -2, -3, -2, -3, -1, 1, -4, -3, -2, 11, 2, -3], # W - [-2, -2, -2, -3, -2, -1, -2, -3, 2, -1, -1, -2, -1, 3, -3, -2, -2, 2, 7, -1], # Y - [ 0, -3, -3, -3, -1, -2, -2, -3, -3, 3, 1, -2, 1, -1, -2, -2, 0, -3, -1, 4], # V - ], - dtype=np.int32, - ) - tcrblosum_alpha_substitution_matrix = np.array( - [ - # A R N D C Q E G H I L K M F P S T W Y V - [ 2, -1, -1, -1, 0, 0, 0, 0, 0, -1, -1, -1, -1, -1, 0, 0, -1, 0, -1, 0], # A - [-1, 1, 0, 0, 1, 0, 0, 0, 0, 0, -1, 0, 0, 0, 0, 0, 0, 0, 0, -1], # R - [-1, 0, 1, 0, 0, 0, 0, 0, 0, -1, -2, 1, 0, 0, 0, 0, 0, 0, 0, -2], # N - [-1, 0, 0, 1, -5, 0, 0, 0, 0, -1, -2, 0, 0, 0, 0, 0, 0, 0, 0, -1], # D - [ 0, 1, 0, -5, 2, -4, -4, 0, -2, -5, 0, -5, -4, -4, -4, 0, -6, -2, -5, 0], # C - [ 0, 0, 0, 0, -4, 2, 0, 0, 0, -1, -2, 1, 0, 0, 0, 0, 0, 0, 0, -2], # Q - [ 0, 0, 0, 0, -4, 0, 1, 0, 1, -1, 0, -1, 0, 0, 0, 0, 0, 0, 0, 0], # E - [ 0, 0, 0, 0, 0, 0, 0, 1, 0, -2, -1, -1, 0, 0, 0, 0, 0, 0, 0, 0], # G - [ 0, 0, 0, 0, -2, 0, 1, 0, 2, 0, 0, -1, 0, 0, 1, 0, 0, 1, 0, 0], # H - [-1, 0, -1, -1, -5, -1, -1, -2, 0, 3, 0, -1, 0, 0, 0, 0, 1, -1, 0, 0], # I - [-1, -1, -2, -2, 0, -2, 0, -1, 0, 0, 2, -4, 0, 1, 0, -1, -1, -1, 0, 0], # L - [-1, 0, 1, 0, -5, 1, -1, -1, -1, -1, -4, 3, 0, -3, 0, -2, -1, -2, -4, -3], # K - [-1, 0, 0, 0, -4, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, -1, 0], # M - [-1, 0, 0, 0, -4, 0, 0, 0, 0, 0, 1, -3, 0, 1, 0, 0, 0, 0, 0, 0], # F - [ 0, 0, 0, 0, -4, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0], # P - [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -1, -2, 0, 0, 0, 1, 0, 0, 0, -1], # S - [-1, 0, 0, 0, -6, 0, 0, 0, 0, 1, -1, -1, 0, 0, 0, 0, 1, 0, 0, 0], # T - [ 0, 0, 0, 0, -2, 0, 0, 0, 1, -1, -1, -2, 0, 0, 0, 0, 0, 2, 0, -1], # W - [-1, 0, 0, 0, -5, 0, 0, 0, 0, 0, 0, -4, -1, 0, 0, 0, 0, 0, 1, -1], # Y - [ 0, -1, -2, -1, 0, -2, 0, 0, 0, 0, 0, -3, 0, 0, 0, -1, 0, -1, -1, 1], # V - ], - dtype=np.int32, - ) - tcrblosum_beta_substitution_matrix = np.array( - [ - # A R N D C Q E G H I L K M F P S T W Y V - [ 0, 0, 0, 0, -5, 0, -1, 0, 0, 0, 0, 0, 0, -1, 0, 0, 0, 0, -1, 0], # A - [ 0, 2, 0, 0, -4, -1, -1, 0, 0, 0, 0, 0, 0, -1, 0, 0, 0, 0, -1, 0], # R - [ 0, 0, 1, 1, -4, 0, 0, 0, 0, 0, -1, 0, 0, -1, 0, -1, 0, 0, 0, 0], # N - [ 0, 0, 1, 1, -4, 0, 0, 0, 0, 0, 0, 0, 0, -1, 0, -1, 0, 0, 0, 0], # D - [-5, -4, -4, -4, 2, -6, -5, 0, -3, -3, -5, -2, -1, -5, -4, 0, -5, -2, -5, -4], # C - [ 0, -1, 0, 0, -6, 2, -1, -1, -1, 0, 1, -1, 0, -2, -1, -2, -1, 0, 0, -1], # Q - [-1, -1, 0, 0, -5, -1, 2, 0, -1, 0, -1, 1, 0, -2, 0, -2, 1, 0, -1, 0], # E - [ 0, 0, 0, 0, 0, -1, 0, 0, 0, 0, 0, 0, 0, -1, 0, 0, 0, 0, -1, 0], # G - [ 0, 0, 0, 0, -3, -1, -1, 0, 2, 0, 0, -1, 0, 2, 0, -1, 0, 0, 1, 0], # H - [ 0, 0, 0, 0, -3, 0, 0, 0, 0, 2, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0], # I - [ 0, 0, -1, 0, -5, 1, -1, 0, 0, 0, 1, 0, 0, 0, 0, -1, 0, 0, 0, 0], # L - [ 0, 0, 0, 0, -2, -1, 1, 0, -1, 0, 0, 1, 0, -1, 0, 0, 0, 0, -1, 0], # K - [ 0, 0, 0, 0, -1, 0, 0, 0, 0, 2, 0, 0, 2, 0, 0, 0, 0, 0, -1, 0], # M - [-1, -1, -1, -1, -5, -2, -2, -1, 2, 0, 0, -1, 0, 2, 0, -2, 0, 0, 2, -1], # F - [ 0, 0, 0, 0, -4, -1, 0, 0, 0, 0, 0, 0, 0, 0, 1, -1, 0, 0, -1, 0], # P - [ 0, 0, -1, -1, 0, -2, -2, 0, -1, 0, -1, 0, 0, -2, -1, 1, 0, 0, -2, 0], # S - [ 0, 0, 0, 0, -5, -1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], # T - [ 0, 0, 0, 0, -2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0], # W - [-1, -1, 0, 0, -5, 0, -1, -1, 1, 0, 0, -1, -1, 2, -1, -2, 0, 0, 2, -1], # Y - [ 0, 0, 0, 0, -4, -1, 0, 0, 0, 0, 0, 0, 0, -1, 0, 0, 0, 0, -1, 0], # V - ], - dtype=np.int32, - ) - # fmt: on - def __init__( self, cutoff: int = 20, @@ -1330,46 +1222,43 @@ def __init__( raise ValueError("`distance_cap` must be non-negative, `None`, or 'default'.") if base_matrix == "blosum62": + substitution_matrix = BLOSUM62 matrix_distance_cap = 4 if distance_cap == "default" else distance_cap - self.tcr_nb_distance_matrix = _substitution_to_distance_matrix( - self.blosum62_substitution_matrix, - self.parasail_aa_alphabet_with_unknown, - self.matrix_alphabet, - distance_cap=matrix_distance_cap, - distance_offset=4, - ) + distance_offset = 4 elif base_matrix == "tcrblosum": if chain_type == "VJ": - tcrdist_substitution_matrix = self.tcrblosum_alpha_substitution_matrix + substitution_matrix = TCRBLOSUM_ALPHA elif chain_type == "VDJ": - tcrdist_substitution_matrix = self.tcrblosum_beta_substitution_matrix + substitution_matrix = TCRBLOSUM_BETA else: raise ValueError("`chain_type` must be 'VJ' or 'VDJ' when `base_matrix='tcrblosum'`.") # Use one `distance_offset` across alpha and beta matrices so equal substitution scores # map to equal distances for both chain types. This keeps VJ and VDJ distances on a # shared scale for clonotype clustering. - off_diagonal = ~np.eye(len(self.matrix_alphabet), dtype=bool) + off_diagonal = ~np.eye(len(substitution_matrix.alphabet), dtype=bool) max_score = int( max( - np.max(self.tcrblosum_alpha_substitution_matrix[off_diagonal]), - np.max(self.tcrblosum_beta_substitution_matrix[off_diagonal]), + np.max(TCRBLOSUM_ALPHA.matrix[off_diagonal]), + np.max(TCRBLOSUM_BETA.matrix[off_diagonal]), ) ) matrix_distance_cap = None if distance_cap == "default" else distance_cap - self.tcr_nb_distance_matrix = _substitution_to_distance_matrix( - tcrdist_substitution_matrix, - self.parasail_aa_alphabet_with_unknown, - self.matrix_alphabet, - distance_cap=matrix_distance_cap, - distance_offset=max_score + 1, - ) + distance_offset = max_score + 1 else: raise ValueError(f"Unknown `base_matrix`: {base_matrix!r}") + self.tcr_nb_distance_matrix = _substitution_to_distance_matrix( + substitution_matrix.matrix, + AA_ALPHABET_WITH_UNKNOWN, + substitution_matrix.alphabet, + distance_cap=matrix_distance_cap, + distance_offset=distance_offset, + ) + super().__init__(n_jobs=n_jobs, n_blocks=n_blocks, histogram=histogram) def _tcrdist_mat( @@ -1530,10 +1419,248 @@ def _nb_tcrdist_mat(): _metric_mat = _tcrdist_mat +class NeedlemanWunschDistanceCalculator(_MetricDistanceCalculator): + """Computes pairwise global-alignment distances with linear-gap Needleman-Wunsch. + + For each sequence pair, a global alignment score is computed using the + Needleman-Wunsch dynamic programming algorithm with one linear gap penalty + for every gap position. The alignment score is converted into a distance by + subtracting it from the best possible self-alignment score of the two + sequences: + ``min(self_score(seq1), self_score(seq2)) - alignment_score(seq1, seq2)``. + Distances are therefore small for sequence pairs that can be globally + aligned with few or conservative substitutions and short gaps, and larger + for sequence pairs requiring strongly penalized substitutions or many gap + positions. + + Parameters + ---------- + gap_penalty: + Linear penalty for each gap position + cutoff: + Will eliminate distances > cutoff to make efficient use of sparse matrices + n_jobs: + Number of numba parallel threads to use for the pairwise distance calculation + n_blocks: + Number of joblib delayed objects (blocks to compute) given to joblib.Parallel + histogram: + Determines whether a nearest neighbor histogram should be created. Not implemented for this metric + base_matrix: + Amino acid substitution matrix. `"blosum62"` uses BLOSUM62, while + `"tcrblosum"` uses TCRBLOSUM alpha/beta substitution matrices depending on + `chain_type` + chain_type: + Required when `base_matrix="tcrblosum"`. `"VJ"` selects the alpha-chain matrix + and `"VDJ"` selects the beta-chain matrix. When called via `ir_dist`, this value + is set automatically and should not be provided + """ + + def __init__( + self, + cutoff: int = 10, + *, + gap_penalty: int = 4, + n_jobs: int = -1, + n_blocks: int = 1, + histogram: bool = False, + base_matrix: Literal["blosum62", "tcrblosum"] = "blosum62", + chain_type: Literal["VJ", "VDJ"] | None = None, + ): + if cutoff < 0: + raise ValueError("`cutoff` must be non-negative.") + if gap_penalty < 0: + raise ValueError("`gap_penalty` must be non-negative.") + + self.cutoff = cutoff + self.gap_penalty = gap_penalty + self.histogram = histogram + + if base_matrix == "blosum62": + substitution_matrix = BLOSUM62_WITH_AMBIGUOUS + elif base_matrix == "tcrblosum": + if chain_type == "VJ": + substitution_matrix = TCRBLOSUM_ALPHA + elif chain_type == "VDJ": + substitution_matrix = TCRBLOSUM_BETA + else: + raise ValueError("`chain_type` must be 'VJ' or 'VDJ' when `base_matrix='tcrblosum'`.") + else: + raise ValueError(f"Unknown `base_matrix`: {base_matrix!r}") + + self.nw_substitution_matrix = _map_matrix_to_alphabet( + substitution_matrix.matrix, + substitution_matrix.alphabet, + AA_ALPHABET_WITH_UNKNOWN, + ) + super().__init__(n_jobs=n_jobs, n_blocks=n_blocks, histogram=histogram) + + def _needleman_wunsch_mat( + self, + *, + seqs: Sequence[str], + seqs2: Sequence[str], + is_symmetric: bool = False, + start_column: int = 0, + ) -> tuple[list[np.ndarray], list[np.ndarray], np.ndarray, np.ndarray]: + """Computes pairwise linear-gap Needleman-Wunsch distances.""" + max_seq_len = max(len(s) for s in (*seqs, *seqs2)) + + seqs_mat1, seqs_L1 = _seqs2mat(seqs, max_len=max_seq_len) + seqs_mat2, seqs_L2 = _seqs2mat(seqs2, max_len=max_seq_len) + + cutoff = self.cutoff + gap_penalty = self.gap_penalty + substitution_matrix = self.nw_substitution_matrix + start_column *= is_symmetric + + nb.set_num_threads(_get_usable_cpus(n_jobs=self.n_jobs, use_numba=True)) + num_threads = nb.get_num_threads() + jit_parallel = num_threads > 1 + + @nb.njit + def _self_scores(seqs_mat, seqs_L): + scores = np.zeros(seqs_mat.shape[0], dtype=np.int32) + for row in range(seqs_mat.shape[0]): + score = 0 + for i in range(seqs_L[row]): + aa = seqs_mat[row, i] + score += substitution_matrix[aa, aa] + scores[row] = score + return scores + + self_scores1 = _self_scores(seqs_mat1, seqs_L1) + self_scores2 = ( + self_scores1 + if is_symmetric and seqs_mat1.shape[0] == seqs_mat2.shape[0] + else _self_scores(seqs_mat2, seqs_L2) + ) + + @nb.jit(nopython=True, parallel=jit_parallel, nogil=True) + def _nb_needleman_wunsch_mat(): + assert seqs_mat1.shape[0] == seqs_L1.shape[0] + assert seqs_mat2.shape[0] == seqs_L2.shape[0] + + num_rows = seqs_mat1.shape[0] + num_cols = seqs_mat2.shape[0] + max_len = seqs_mat1.shape[1] + + data_rows = nb.typed.List() + indices_rows = nb.typed.List() + row_element_counts = np.zeros(num_rows, dtype=np.int32) + + empty_row = np.zeros(0, dtype=np.int32) + for _ in range(0, num_rows): + data_rows.append([empty_row]) + indices_rows.append([empty_row]) + + data_row_matrix = np.empty((num_threads, num_cols), dtype=np.int32) + indices_row_matrix = np.empty((num_threads, num_cols), dtype=np.int32) + previous_rows = np.empty((num_threads, max_len + 1), dtype=np.int32) + current_rows = np.empty((num_threads, max_len + 1), dtype=np.int32) + + # Only alignment paths within a band around the diagonal can stay within the cutoff, + # because each step away from the diagonal requires one gap penalty. + band_width = max_len if gap_penalty == 0 else cutoff // gap_penalty + + for row_index in nb.prange(num_rows): + thread_id = nb.get_thread_id() + row_end_index = 0 + seq1_len = seqs_L1[row_index] + + col_start = start_column + row_index * is_symmetric + if is_symmetric: + data_row_matrix[thread_id, row_end_index] = 1 + indices_row_matrix[thread_id, row_end_index] = col_start + row_end_index += 1 + col_start += 1 + + for col_index in range(col_start, num_cols): + seq2_len = seqs_L2[col_index] + len_diff = abs(seq1_len - seq2_len) + + # Skip pairs whose length difference alone cannot stay within the cutoff. + if len_diff * gap_penalty > cutoff: + continue + + min_self_score = min(self_scores1[row_index], self_scores2[col_index]) + + band_start = 1 - band_width + band_end = 1 + band_width + + j_end = min(band_width, seq2_len) + + for j in range(j_end + 1): + previous_rows[thread_id, j] = -j * gap_penalty + + for i in range(1, seq1_len + 1): + aa1 = seqs_mat1[row_index, i - 1] + + j_start = max(1, band_start) + j_end = min(seq2_len, band_end) + + if i <= band_width: + current_rows[thread_id, 0] = -i * gap_penalty + + for j in range(j_start, j_end + 1): + aa2 = seqs_mat2[col_index, j - 1] + match_score = previous_rows[thread_id, j - 1] + substitution_matrix[aa1, aa2] + best_score = match_score + + if band_width == 0: + best_score = match_score + elif j == band_start: + delete_score = previous_rows[thread_id, j] - gap_penalty + best_score = max(match_score, delete_score) + elif j == band_end: + insert_score = current_rows[thread_id, j - 1] - gap_penalty + best_score = max(match_score, insert_score) + else: + delete_score = previous_rows[thread_id, j] - gap_penalty + insert_score = current_rows[thread_id, j - 1] - gap_penalty + best_score = max(match_score, delete_score, insert_score) + + current_rows[thread_id, j] = best_score + + copy_start = j_start * (i > band_width) + + for j in range(copy_start, j_end + 1): + previous_rows[thread_id, j] = current_rows[thread_id, j] + + band_start += 1 + band_end += 1 + + distance = max(0, min_self_score - previous_rows[thread_id, seq2_len]) + 1 + + if distance <= cutoff + 1: + data_row_matrix[thread_id, row_end_index] = distance + indices_row_matrix[thread_id, row_end_index] = col_index + row_end_index += 1 + + data_rows[row_index][0] = data_row_matrix[thread_id, 0:row_end_index].copy() + indices_rows[row_index][0] = indices_row_matrix[thread_id, 0:row_end_index].copy() + row_element_counts[row_index] = row_end_index + + data_rows_flat = [] + indices_rows_flat = [] + + for i in range(len(data_rows)): + data_rows_flat.append(data_rows[i][0]) + indices_rows_flat.append(indices_rows[i][0]) + + return data_rows_flat, indices_rows_flat, row_element_counts + + data_rows, indices_rows, row_element_counts = _nb_needleman_wunsch_mat() + return data_rows, indices_rows, row_element_counts, np.array([None]) + + _metric_mat = _needleman_wunsch_mat + + @deprecated( Deprecation( "0.15.0", - "`FastAlignmentDistanceCalculator` achieves (depending on the settings) identical results at a higher speed.", + "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.", ) ) @_doc_params(params=_doc_params_parallel_distance_calculator) @@ -1657,6 +1784,14 @@ def _self_alignment_scores(self, seqs: Sequence) -> dict: ) +@deprecated( + Deprecation( + "0.25.0", + "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.", + ) +) @_doc_params(params=_doc_params_parallel_distance_calculator) class FastAlignmentDistanceCalculator(ParallelDistanceCalculator): """\ diff --git a/src/scirpy/tests/data/needleman_wunsch_test_data/needleman_wunsch_WU3k_csr_result.npz b/src/scirpy/tests/data/needleman_wunsch_test_data/needleman_wunsch_WU3k_csr_result.npz new file mode 100644 index 000000000..439f4b818 Binary files /dev/null and b/src/scirpy/tests/data/needleman_wunsch_test_data/needleman_wunsch_WU3k_csr_result.npz differ diff --git a/src/scirpy/tests/test_deprecations.py b/src/scirpy/tests/test_deprecations.py index 7a2bfbeed..4bea4acc1 100644 --- a/src/scirpy/tests/test_deprecations.py +++ b/src/scirpy/tests/test_deprecations.py @@ -51,7 +51,7 @@ def test_no_spurious_block_size_warning(calculator): def test_alignment_distance_calculator_deprecated(): - with pytest.warns(FutureWarning, match="FastAlignmentDistanceCalculator"): + with pytest.warns(FutureWarning, match="NeedlemanWunschDistanceCalculator"): AlignmentDistanceCalculator() diff --git a/src/scirpy/tests/test_ir_dist.py b/src/scirpy/tests/test_ir_dist.py index 3f1b4ba96..41db0a55a 100644 --- a/src/scirpy/tests/test_ir_dist.py +++ b/src/scirpy/tests/test_ir_dist.py @@ -201,6 +201,49 @@ def test_ir_dist_tcrdist_tcrblosum_chain_routing(mudata): npt.assert_array_equal(res["VDJ"]["distances"].toarray(), np.array([[1, 19], [19, 1]])) +@pytest.mark.parametrize("mudata", [False, True], ids=["AnnData", "MuData"]) +def test_ir_dist_needleman_wunsch_tcrblosum_chain_routing(mudata): + # `ir_dist` should automatically route VJ to TCRBLOSUM alpha and VDJ to beta. + adata = _make_adata( + pd.DataFrame( + [ + ["cell1", "AAACAAAA", "AAACAAAA", "TRA", "TRB"], + ["cell2", "AAAHAAAA", "AAAHAAAA", "TRA", "TRB"], + ], + columns=[ + "cell_id", + "IR_VJ_1_junction_aa", + "IR_VDJ_1_junction_aa", + "IR_VJ_1_locus", + "IR_VDJ_1_locus", + ], + ).set_index("cell_id"), + mudata, + ) + + ir.pp.ir_dist( + adata, + metric="needleman_wunsch", + sequence="aa", + cutoff=20, + gap_penalty=4, + base_matrix="tcrblosum", + key_added="ir_dist_needleman_wunsch_tcrblosum", + n_jobs=1, + ) + + res = ( + adata.mod["airr"].uns["ir_dist_needleman_wunsch_tcrblosum"] + if isinstance(adata, MuData) + else adata.uns["ir_dist_needleman_wunsch_tcrblosum"] + ) + expected_seqs = np.array(["AAACAAAA", "AAAHAAAA"]) + npt.assert_array_equal(res["VJ"]["seqs"], expected_seqs) + npt.assert_array_equal(res["VDJ"]["seqs"], expected_seqs) + npt.assert_array_equal(res["VJ"]["distances"].toarray(), np.array([[1, 5], [5, 1]])) + npt.assert_array_equal(res["VDJ"]["distances"].toarray(), np.array([[1, 6], [6, 1]])) + + @pytest.mark.parametrize("with_adata2", [False, True]) @pytest.mark.parametrize("joblib_backend", ["loky", "multiprocessing", "threading"]) @pytest.mark.parametrize("n_jobs", [1, 2]) @@ -581,7 +624,7 @@ def test_compute_distances_second_anndata( npt.assert_equal(dist, expected_dist if not swap_query_reference else expected_dist.T) -@pytest.mark.parametrize("metric", ["identity", "levenshtein", "alignment", "tcrdist", "hamming"]) +@pytest.mark.parametrize("metric", ["identity", "levenshtein", "alignment", "tcrdist", "hamming", "needleman_wunsch"]) def test_ir_dist_empty_anndata(adata_cdr3, metric): adata_empty = adata_cdr3.mod["airr"].copy() if isinstance(adata_cdr3, MuData) else adata_cdr3.copy() # reset chain indices such that no chain will actually be used. diff --git a/src/scirpy/tests/test_ir_dist_metrics.py b/src/scirpy/tests/test_ir_dist_metrics.py index f970df33c..e19ab7b2c 100644 --- a/src/scirpy/tests/test_ir_dist_metrics.py +++ b/src/scirpy/tests/test_ir_dist_metrics.py @@ -6,6 +6,7 @@ import scipy.sparse import scirpy as ir +from scirpy.ir_dist._substitution_matrices import _map_matrix_to_alphabet, _substitution_to_distance_matrix from scirpy.ir_dist.metrics import ( AlignmentDistanceCalculator, DistanceCalculator, @@ -14,9 +15,9 @@ HammingDistanceCalculator, IdentityDistanceCalculator, LevenshteinDistanceCalculator, + NeedlemanWunschDistanceCalculator, ParallelDistanceCalculator, TCRdistDistanceCalculator, - _substitution_to_distance_matrix, ) from .util import _squarify @@ -69,7 +70,7 @@ def test_squarify(): ) -def test_substitution_to_distance_matrix_converts_substitution_matrix(): +def test_substitution_to_distance_matrix(): substitution_matrix = np.array( [ [4, 3, 0, -1], @@ -125,6 +126,45 @@ def test_substitution_to_distance_matrix_converts_substitution_matrix(): ) +def test_map_matrix_to_alphabet(): + matrix = np.array([[1, 2], [3, 4]], dtype=np.int32) + + mapped_matrix = _map_matrix_to_alphabet(matrix, source_alphabet="AB", target_alphabet="BCA") + + npt.assert_array_equal( + mapped_matrix, + np.array( + [ + [4, 0, 3], + [0, 0, 0], + [2, 0, 1], + ], + dtype=np.int32, + ), + strict=True, + ) + + +@pytest.mark.parametrize( + ("matrix", "source_alphabet", "target_alphabet", "match"), + [ + (np.zeros((1, 1)), "AB", "AB", "must have shape"), + (np.zeros((2, 2)), "AA", "A", "source_alphabet.*duplicate"), + (np.zeros((2, 2)), "AB", "ABBA", "target_alphabet.*duplicate"), + (np.zeros((2, 2)), "AB", "AC", "missing characters"), + ], + ids=["shape", "duplicate-source", "duplicate-target", "missing-target-character"], +) +def test_map_matrix_to_alphabet_rejects_invalid_input( + matrix: np.ndarray, + source_alphabet: str, + target_alphabet: str, + match: str, +): + with pytest.raises(ValueError, match=match): + _map_matrix_to_alphabet(matrix, source_alphabet, target_alphabet) + + def test_block_iter(): seqs1 = list("ABCDE") seqs2 = list("HIJKLM") @@ -371,13 +411,23 @@ def test_fast_alignment_dist_with_two_seq_arrays(): @pytest.mark.extra @pytest.mark.parametrize( - "metric", ["alignment", "fastalignment", "identity", "hamming", "normalized_hamming", "levenshtein", "tcrdist"] + "metric", + [ + "alignment", + "fastalignment", + "identity", + "hamming", + "normalized_hamming", + "levenshtein", + "tcrdist", + "needleman_wunsch", + ], ) @pytest.mark.parametrize("n_jobs", [-1, 1, 2]) def test_sequence_dist_all_metrics(metric, n_jobs): # Smoke test, no assertions! # Smoke test, no assertions! - metrics_with_n_blocks = ["hamming", "normalized_hamming", "tcrdist"] + metrics_with_n_blocks = ["hamming", "normalized_hamming", "tcrdist", "needleman_wunsch"] n_blocks_params = [1, 2] unique_seqs = np.array(["AAA", "ARA", "AFFFFFA", "FAFAFA", "FFF"]) @@ -820,6 +870,166 @@ def test_tcrdist(test_parameters, test_input, expected_result): assert np.array_equal(res.todense(), expected_result) +@pytest.mark.parametrize( + "test_parameters,test_input,expected_result", + [ + # test more complex strings with unequal length and set high cutoff such that cutoff is neglected + ( + {"cutoff": 200, "gap_penalty": 11, "n_jobs": 1}, + ( + np.array(["AA", "AAA", "AARA", "AHA", "AHLAA"]), + np.array(["AA", "AAA", "AARA", "AHA", "AHLAA"]), + ), + np.array( + [ + [1, 12, 23, 12, 34], + [12, 1, 12, 7, 23], + [23, 12, 1, 20, 23], + [12, 7, 20, 1, 23], + [34, 23, 23, 23, 1], + ] + ), + ), + # test cutoff filtering + ( + {"cutoff": 10, "gap_penalty": 4, "n_jobs": 1}, + ( + np.array(["AAACAAAA", "AAARAAAA", "AAAHAAAA"]), + np.array(["AAACAAAA", "AAARAAAA", "AAAHAAAA"]), + ), + np.array([[1, 9, 0], [9, 1, 6], [0, 6, 1]]), + ), + # test asymmetric sequence arrays + ( + {"cutoff": 20, "gap_penalty": 4, "n_jobs": 1}, + ( + np.array(["AAAA", "AATA", "HHHH", "WWWW"]), + np.array(["WWWW", "AAAA", "ATAA"]), + ), + np.array([[0, 1, 5], [0, 5, 10], [0, 0, 0], [1, 0, 0]]), + ), + # test ambiguous BLOSUM62 symbols + ( + {"cutoff": 20, "gap_penalty": 4, "n_jobs": 1}, + ( + np.array(["AABA", "AAZA", "AADA", "AAEA"]), + np.array(["AABA", "AAZA", "AADA", "AAEA"]), + ), + np.array([[1, 4, 1, 4], [4, 1, 4, 1], [1, 4, 1, 4], [4, 1, 4, 1]]), + ), + # test ambiguous X scores with asymmetric arrays + ( + {"cutoff": 6, "gap_penalty": 2, "n_jobs": 1}, + ( + np.array(["X", "AX", "XA", "AA"]), + np.array(["A", "XX", "AA"]), + ), + np.array([[1, 2, 2], [2, 1, 1], [2, 1, 1], [3, 1, 1]]), + ), + # test empty input arrays + ( + {"cutoff": 20, "gap_penalty": 4, "n_jobs": 1}, + (np.array([]), np.array([])), + np.empty((0, 0)), + ), + # test standard parameters with second sequences array set to None + ( + {"cutoff": 20, "gap_penalty": 4, "n_jobs": 1}, + (np.array(["AA", "AAA", "AHA"]), None), + np.array([[1, 5, 5], [5, 1, 7], [5, 7, 1]]), + ), + # test with gap_penalty set to 0 + ( + {"cutoff": 20, "gap_penalty": 0, "n_jobs": 1}, + ( + np.array(["AA", "AAA", "AHA"]), + np.array(["AA", "AAA", "AHA"]), + ), + np.array([[1, 1, 1], [1, 1, 5], [1, 5, 1]]), + ), + # test with low gap_penalty and high cutoff + ( + {"cutoff": 50, "gap_penalty": 1, "n_jobs": 1}, + ( + np.array(["AA", "AAA", "AHA", "AHLAA"]), + np.array(["AA", "AAA", "AHA", "AHLAA"]), + ), + np.array([[1, 2, 2, 4], [2, 1, 7, 3], [2, 7, 1, 3], [4, 3, 3, 1]]), + ), + # test with high gap_penalty and high cutoff + ( + {"cutoff": 50, "gap_penalty": 8, "n_jobs": 1}, + ( + np.array(["AA", "AAA", "AHA", "AHLAA"]), + np.array(["AA", "AAA", "AHA", "AHLAA"]), + ), + np.array([[1, 9, 9, 25], [9, 1, 7, 17], [9, 7, 1, 17], [25, 17, 17, 1]]), + ), + # test with high gap_penalty and tight cutoff filtering + ( + {"cutoff": 8, "gap_penalty": 8, "n_jobs": 1}, + ( + np.array(["AA", "AAA", "AHA", "AHLAA"]), + np.array(["AA", "AAA", "AHA", "AHLAA"]), + ), + np.array([[1, 9, 9, 0], [9, 1, 7, 0], [9, 7, 1, 0], [0, 0, 0, 1]]), + ), + # test asymmetric arrays with lower gap_penalty + ( + {"cutoff": 12, "gap_penalty": 2, "n_jobs": 1}, + ( + np.array(["AA", "AHA", "AHLAA"]), + np.array(["AAA", "AHAA", "WW"]), + ), + np.array([[3, 5, 0], [7, 3, 0], [5, 3, 0]]), + ), + # test with cutoff set to 0 + ( + {"cutoff": 0, "gap_penalty": 4, "n_jobs": 1}, + ( + np.array(["AA", "AAA", "AHA"]), + np.array(["AA", "AAA", "AHA"]), + ), + np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]), + ), + # test very small input sequences + ( + {"cutoff": 20, "gap_penalty": 4, "n_jobs": 1}, + (np.array(["A"]), np.array(["C"])), + np.array([[5]]), + ), + # test empty second input array + ( + {"cutoff": 20, "gap_penalty": 4, "n_jobs": 1}, + (np.array(["A", "AA"]), np.array([])), + np.empty((2, 0)), + ), + ], +) +def test_needleman_wunsch(test_parameters, test_input, expected_result): + # Check direct calculator results for small edge cases and parameter combinations. + needleman_wunsch_calculator = NeedlemanWunschDistanceCalculator(**test_parameters) + seq1, seq2 = test_input + + res = needleman_wunsch_calculator.calc_dist_mat(seq1, seq2) + + assert isinstance(res, scipy.sparse.csr_matrix) + assert res.shape == expected_result.shape + assert np.array_equal(res.todense(), expected_result) + + +def test_needleman_wunsch_clamps_negative_scores(): + # Ambiguous X scores can yield negative raw distances and should be clamped to zero. + seqs = np.array(["X", "A", "XX", "AA"]) + needleman_wunsch_calculator = NeedlemanWunschDistanceCalculator(cutoff=10, gap_penalty=4, n_jobs=1) + expected_result = np.array([[1, 1, 4, 4], [1, 1, 3, 5], [4, 3, 1, 1], [4, 5, 1, 1]]) + + res = needleman_wunsch_calculator.calc_dist_mat(seqs, seqs) + + assert isinstance(res, scipy.sparse.csr_matrix) + assert np.array_equal(res.todense(), expected_result) + + def test_sequence_dist_tcrdist_tcrblosum(): # `sequence_dist` needs an explicit `chain_type` for `tcrdist` with `tcrblosum`; # `ir_dist` handles this automatically. @@ -835,6 +1045,31 @@ def test_sequence_dist_tcrdist_tcrblosum(): npt.assert_array_equal(res.toarray(), np.array([[1, 7], [7, 1]])) +def test_sequence_dist_needleman_wunsch_tcrblosum(): + # `sequence_dist` needs an explicit `chain_type` for `needleman_wunsch` with `tcrblosum`; + # `ir_dist` handles this automatically. + seqs = np.array(["AAACAAAA", "AAAHAAAA"]) + direct_calculator = NeedlemanWunschDistanceCalculator( + cutoff=20, + gap_penalty=4, + n_jobs=1, + base_matrix="tcrblosum", + chain_type="VJ", + ) + expected_result = direct_calculator.calc_dist_mat(seqs, seqs) + + res = ir.ir_dist.sequence_dist( + seqs, + metric="needleman_wunsch", + cutoff=20, + gap_penalty=4, + n_jobs=1, + base_matrix="tcrblosum", + chain_type="VJ", + ) + npt.assert_array_equal(res.toarray(), expected_result.toarray()) + + def test_sequence_dist_tcrdist_distance_cap(): seqs = np.array(["AAACAAAA", "AAAHAAAA"]) @@ -895,6 +1130,21 @@ def test_tcrdist_base_matrix_validation(kwargs, match): TCRdistDistanceCalculator(**kwargs) +@pytest.mark.parametrize( + "kwargs,match", + [ + ({"base_matrix": "tcrblosum"}, r"`chain_type` must be 'VJ' or 'VDJ' when `base_matrix='tcrblosum'`\."), + ({"base_matrix": "foo"}, r"Unknown `base_matrix`: 'foo'"), + ({"cutoff": -1}, r"`cutoff` must be non-negative\."), + ({"gap_penalty": -1}, r"`gap_penalty` must be non-negative\."), + ], +) +def test_needleman_wunsch_base_matrix_validation(kwargs, match): + # Invalid matrix and gap settings should fail before distance computation starts. + with pytest.raises(ValueError, match=match): + NeedlemanWunschDistanceCalculator(**kwargs) + + def test_tcrdist_reference(): # test tcrdist against reference implementation from . import TESTDATA @@ -919,6 +1169,28 @@ def test_tcrdist_reference(): assert np.array_equal(res.indptr, reference_result.indptr) +def test_needleman_wunsch_reference(): + # test needleman-wunsch against a precomputed linear-gap alignment reference + from . import TESTDATA + + seqs = np.load(TESTDATA / "tcrdist_test_data/tcrdist_WU3k_seqs.npy") + reference_result = scipy.sparse.load_npz( + TESTDATA / "needleman_wunsch_test_data/needleman_wunsch_WU3k_csr_result.npz" + ) + needleman_wunsch_calculator = NeedlemanWunschDistanceCalculator( + cutoff=20, + gap_penalty=4, + n_jobs=4, + n_blocks=2, + ) + + res = needleman_wunsch_calculator.calc_dist_mat(seqs, seqs) + + assert np.array_equal(res.data, reference_result.data) + assert np.array_equal(res.indices, reference_result.indices) + assert np.array_equal(res.indptr, reference_result.indptr) + + def test_hamming_reference(): # test hamming distance against reference implementation from . import TESTDATA @@ -982,6 +1254,14 @@ def test_tcrdist_histogram_not_implemented(): _ = tcrdist_calculator.calc_dist_mat(seqs, seqs) +def test_needleman_wunsch_histogram_not_implemented(): + # Histogram mode should fail explicitly until it is implemented for needleman_wunsch. + with pytest.raises(NotImplementedError, match=None): + needleman_wunsch_calculator = NeedlemanWunschDistanceCalculator(histogram=True) + seqs = np.array(["AAAA", "AA", "AABB", "ABA"]) + _ = needleman_wunsch_calculator.calc_dist_mat(seqs, seqs) + + @pytest.mark.gpu def test_gpu_hamming_reference(): # test hamming distance against reference implementation