diff --git a/src/scirpy/pl/_logoplots.py b/src/scirpy/pl/_logoplots.py index 12f37b39c..9192d7d79 100644 --- a/src/scirpy/pl/_logoplots.py +++ b/src/scirpy/pl/_logoplots.py @@ -32,13 +32,14 @@ def logoplot_cdr3_motif( ax: plt.Axes | None = None, fig_kws: dict | None = None, **kwargs, -) -> Logo: +) -> Logo | list[Logo]: """ Generates logoplots of CDR3 sequences This is a user friendly wrapper function around the logomaker python package. Enables the analysis of potential amino acid motifs by displaying logo plots. - Subsetting of AnnData/MuData has to be performed manually beforehand (or while calling) and only cells with equal cdr3 sequence lengths are permitted. + Subsetting of AnnData/MuData has to be performed manually beforehand (or while calling). + If multiple CDR3 sequence lengths are present, separate logo plots are generated for each length. Parameters ---------- @@ -72,6 +73,7 @@ def logoplot_cdr3_motif( x coordinate span of each character ax Add the plot to a predefined Axes object. + Can only be used when all selected sequences have the same CDR3 length. fig_kws Parameters passed to the :func:`matplotlib.pyplot.figure` call if no `ax` is specified. @@ -81,35 +83,75 @@ def logoplot_cdr3_motif( Returns ------- - Returns a object of class logomaker.Logo (see here for more information https://logomaker.readthedocs.io/en/latest/implementation.html#matrix-functions) + Returns one object of class :class:`logomaker.Logo` if all selected sequences have the + same length. If multiple sequence lengths are present, returns a list of + :class:`logomaker.Logo` objects, one per length. """ params = DataHandler(adata, airr_mod, airr_key, chain_idx_key) if isinstance(chains, str): chains = [chains] - if ax is None: - fig_kws = {} if fig_kws is None else fig_kws - if "figsize" not in fig_kws: - fig_kws["figsize"] = (6, 2) - ax = _init_ax(fig_kws) - - # make sure that sequences are prealigned i.e. they need to have the the same length + # sequences need to be aligned for each logo plot, so we split by sequence length airr_df = get_airr(params, [cdr3_col], chains) sequence_list = [] for chain in chains: for sequence in airr_df[chain + "_" + cdr3_col]: if not pd.isnull(sequence): sequence_list.append(sequence) + sequences_by_length: dict[int, list[str]] = {} + for sequence in sequence_list: + sequences_by_length.setdefault(len(sequence), []).append(sequence) + n_lengths = len(sequences_by_length) + + def _make_logo(seq_list: list[str], ax_: plt.Axes, title: str) -> Logo: + motif = alignment_to_matrix( + seq_list, to_type=to_type, pseudocount=pseudocount, background=background, center_weights=center_weights + ) + cdr3_logo = Logo( + motif, + color_scheme=color_scheme, + vpad=vpad, + width=width, + font_name=font_name, + ax=ax_, + **kwargs, + ) + cdr3_logo.style_xticks(anchor=0, spacing=1, rotation=45) + cdr3_logo.ax.set_ylabel(f"{to_type}") + cdr3_logo.ax.grid(False) + cdr3_logo.ax.set_xlim([-1, len(motif)]) + cdr3_logo.ax.set_title(title) + return cdr3_logo + + chain_label = "/".join(chains) + if n_lengths == 1: + if ax is None: + fig_kws = {} if fig_kws is None else fig_kws + if "figsize" not in fig_kws: + fig_kws["figsize"] = (6, 2) + ax = _init_ax(fig_kws) + only_sequences = next(iter(sequences_by_length.values())) + return _make_logo(only_sequences, ax, chain_label) + + if ax is not None: + raise ValueError( + "Selected sequences contain multiple CDR3 lengths, but a single `ax` was provided. " + "Please omit `ax` to auto-generate one subplot per length." + ) - motif = alignment_to_matrix( - sequence_list, to_type=to_type, pseudocount=pseudocount, background=background, center_weights=center_weights - ) - cdr3_logo = Logo(motif, color_scheme=color_scheme, vpad=vpad, width=width, font_name=font_name, ax=ax, **kwargs) + fig_kws = {} if fig_kws is None else fig_kws + if "figsize" not in fig_kws: + fig_kws["figsize"] = (6, 2 * n_lengths) - cdr3_logo.style_xticks(anchor=0, spacing=1, rotation=45) - cdr3_logo.ax.set_ylabel(f"{to_type}") - cdr3_logo.ax.grid(False) - cdr3_logo.ax.set_xlim([-1, len(motif)]) - cdr3_logo.ax.set_title("/".join(chains)) - return cdr3_logo + _, axes = plt.subplots(n_lengths, 1, squeeze=False, **fig_kws) + logos = [] + for ax_, sequence_length in zip(axes.ravel(), sorted(sequences_by_length), strict=False): + logos.append( + _make_logo( + sequences_by_length[sequence_length], + ax_, + f"{chain_label} (CDR3 length {sequence_length})", + ) + ) + return logos diff --git a/src/scirpy/tests/test_plotting.py b/src/scirpy/tests/test_plotting.py index 93330a05c..e9b92c046 100644 --- a/src/scirpy/tests/test_plotting.py +++ b/src/scirpy/tests/test_plotting.py @@ -197,3 +197,19 @@ def test_clonotype_network_pie_legend_filters_unused(adata_clonotype_network): def test_logoplot(adata_cdr3): p = pl.logoplot_cdr3_motif(adata_cdr3, chains="VJ_1") assert isinstance(p, logomaker.Logo) + + +@pytest.mark.extra +def test_logoplot_mixed_lengths_returns_multiple_logos(adata_cdr3): + p = pl.logoplot_cdr3_motif(adata_cdr3, chains="VDJ_1") + assert isinstance(p, list) + assert len(p) == 2 + assert all(isinstance(logo, logomaker.Logo) for logo in p) + assert all("CDR3 length" in logo.ax.get_title() for logo in p) + + +@pytest.mark.extra +def test_logoplot_mixed_lengths_with_single_axis_raises(adata_cdr3): + _, ax = plt.subplots(1, 1) + with pytest.raises(ValueError, match="multiple CDR3 lengths"): + pl.logoplot_cdr3_motif(adata_cdr3, chains="VDJ_1", ax=ax)