diff --git a/src/infercnvpy/pl/_chromosome_heatmap.py b/src/infercnvpy/pl/_chromosome_heatmap.py index 14c78e2..045cd62 100644 --- a/src/infercnvpy/pl/_chromosome_heatmap.py +++ b/src/infercnvpy/pl/_chromosome_heatmap.py @@ -15,14 +15,14 @@ from matplotlib.colors import BoundaryNorm, Colormap, ListedColormap, Normalize, TwoSlopeNorm from matplotlib.path import Path as MplPath from matplotlib.text import Text -from matplotlib.ticker import Formatter +from matplotlib.ticker import FixedFormatter from mpl_toolkits.axes_grid1 import make_axes_locatable from mpl_toolkits.axes_grid1.axes_size import Fixed from infercnvpy._logging import SKIP_OWN_FRAMES from infercnvpy._numba import n_kernel_chunks, numba_threads from infercnvpy._util import _as_matrix, _is_backed, _is_dask, _obs_column -from infercnvpy.pl._heatmap_kernels import binned_sums, pixel_rows, sampled_rows +from infercnvpy.pl._heatmap_kernels import binned_sums, column_ranges, gram, pixel_rows, sampled_rows if TYPE_CHECKING: from collections.abc import Callable @@ -121,9 +121,11 @@ def chromosome_heatmap( **kwargs Arguments of :func:`scanpy.pl.heatmap`, see `mode`. By default, the colour scale is a :class:`~matplotlib.colors.TwoSlopeNorm` centred at `vcenter=0`, with limits from the data. Pass the same - `vmin` and `vmax` to compare figures across modes, representations or datasets. `dendrogram=True` orders the groups by - hierarchical clustering of their mean CNV profiles (Pearson correlation, complete linkage, as - :func:`scanpy.tl.dendrogram`) and draws the dendrogram next to the heatmap. + `vmin` and `vmax` to compare figures across modes, representations or datasets. `dendrogram=True` orders the + groups as :func:`scanpy.tl.dendrogram` with its defaults clusters them (as `mode="legacy"`; with more windows + than `scanpy.settings.N_PCS`, it clusters the groups' means of a PCA of all cells) and draws the dendrogram + next to the heatmap. A dendrogram stored in `adata.uns[f"dendrogram_{groupby}"]` is used if it has the same + groups; otherwise the computed one is stored there, as `mode="legacy"` does. Returns ------- @@ -136,15 +138,16 @@ def chromosome_heatmap( if _use_legacy("chromosome_heatmap", adata, groupby, use_rep, mode, kwargs): return _legacy(_legacy_heatmap, adata, groupby, use_rep, cmap, figsize, ax, show, save, return_fig, kwargs) opts = _Options.pop(kwargs) + own = ax is None fig, ax = _figure(ax, figsize) cat = _categories(adata, groupby) palette = _group_colors(adata, groupby, len(cat.categories)) X = _as_matrix(adata.obsm[f"X_{use_rep}"], f"`adata.obsm['X_{use_rep}']`", lazy=True) - dendro = None + dendro, var_range = None, None with numba_threads(n_jobs): if opts.dendrogram: cat, palette, dendro = _order_by_dendrogram( - adata, groupby, cat, palette, lambda: _group_means(X, cat.codes.astype(np.int64), len(cat.categories)) + adata, groupby, cat, palette, lambda: _dendrogram(X, cat, groupby) ) codes = cat.codes.astype(np.int64) order, bounds, rows = pixel_rows(codes, len(cat.categories), _n_pixel_rows(fig, ax, swap=opts.swap_axes)) @@ -152,14 +155,16 @@ def chromosome_heatmap( values, lo, hi = sampled_rows(X, order[(bounds[:-1] + bounds[1:]) // 2]) else: values = binned_sums(X, order, bounds) / np.diff(bounds)[:, None] - values = opts.transform(values) + if opts.standard_scale == "var": + var_range = column_ranges(X) + values = opts.transform(values, var_range) if aggregate == "mean": lo, hi = np.nanmin(values), np.nanmax(values) center = 0 if opts.vcenter is None else opts.vcenter lo, hi = min(lo, center - 1e-9), max(hi, center + 1e-9) # TwoSlopeNorm needs vmin < vcenter < vmax norm = opts.color_norm(lo, hi) axes = _draw(fig, ax, adata, groupby, cat, palette, values, rows, dendro, use_rep, norm, cmap, opts) - return _finish(fig, axes, show=show, save=save, return_fig=return_fig) + return _finish(fig, axes, show=show, save=save, return_fig=return_fig, own=own) def chromosome_heatmap_summary( @@ -232,6 +237,7 @@ def chromosome_heatmap_summary( if _use_legacy("chromosome_heatmap_summary", adata, groupby, use_rep, mode, kwargs): return _legacy(_legacy_summary, adata, groupby, use_rep, cmap, figsize, ax, show, save, return_fig, kwargs) opts = _Options.pop(kwargs) + own = ax is None fig, ax = _figure(ax, figsize) cat = _categories(adata, groupby) palette = _group_colors(adata, groupby, len(cat.categories)) @@ -242,12 +248,17 @@ def chromosome_heatmap_summary( if opts.dendrogram: present = np.flatnonzero(np.bincount(cat.codes[cat.codes >= 0], minlength=len(cat.categories))) before = list(cat.categories[present]) - cat, palette, dendro = _order_by_dendrogram(adata, groupby, cat, palette, lambda: means) + # as `mode="legacy"`: every group mean 10 times, the groups sorted as `astype("category")` sorts them + groups = pd.Categorical(np.repeat(before, 10)) + by_name = means[[before.index(c) for c in groups.categories]] + cat, palette, dendro = _order_by_dendrogram( + adata, groupby, cat, palette, lambda: _dendrogram(np.repeat(means, 10, axis=0), groups, groupby, by_name) + ) means = means[[before.index(c) for c in cat.categories if c in before]] rows = (np.bincount(cat.codes[cat.codes >= 0], minlength=len(cat.categories)) > 0).astype(np.int64) norm = opts.color_norm(float(np.min(means)), float(np.max(means))) axes = _draw(fig, ax, adata, groupby, cat, palette, opts.transform(means), rows, dendro, use_rep, norm, cmap, opts) - return _finish(fig, axes, show=show, save=save, return_fig=return_fig) + return _finish(fig, axes, show=show, save=save, return_fig=return_fig, own=own) # --- options of mode="fast" ------------------------------------------------------------------------------------------- @@ -274,15 +285,22 @@ def pop(cls, kwargs: dict[str, Any]) -> _Options: kwargs.clear() return opts - def transform(self, values: np.ndarray) -> np.ndarray: - """`log` and `standard_scale` of the drawn rows, as :func:`scanpy.pl.heatmap` does for all cells.""" + def transform(self, values: np.ndarray, var_range: tuple[np.ndarray, np.ndarray] | None = None) -> np.ndarray: + """`log` and `standard_scale` of the drawn rows, as :func:`scanpy.pl.heatmap` does for all cells. + + `var_range` is the minimum and maximum of each window over all cells; without it, `standard_scale="var"` + scales over the drawn rows. + """ if self.log: values = np.log1p(values) if self.standard_scale is not None: axis = 0 if self.standard_scale == "var" else 1 - values = values - np.nanmin(values, axis=axis, keepdims=True) + if self.standard_scale == "var" and var_range is not None: + lo, hi = (np.log1p(r) if self.log else r for r in var_range) + else: + lo, hi = np.nanmin(values, axis=axis, keepdims=True), np.nanmax(values, axis=axis, keepdims=True) with np.errstate(invalid="ignore", divide="ignore"): - values = np.nan_to_num(values / np.nanmax(values, axis=axis, keepdims=True), nan=0) + values = np.nan_to_num((values - lo) / (hi - lo), nan=0) return values def color_norm(self, lo: float, hi: float) -> Normalize: @@ -339,7 +357,7 @@ def _use_legacy(func: str, adata: AnnData, groupby: str, use_rep: str, mode: str def _legacy( - fn: Callable[..., dict[str, Axes] | None], + fn: Callable[..., dict[str, Axes]], adata: AnnData, groupby: str, use_rep: str, @@ -353,23 +371,19 @@ def _legacy( ) -> Figure | dict[str, Axes] | None: if ax is not None: raise ValueError("`ax` needs mode='fast' (and only arguments that it supports, see `mode`).") - common = {"groupby": groupby, "use_rep": use_rep, "cmap": cmap, "figsize": figsize, "save": save} - if not return_fig: - return fn(adata, show=show, **common, **kwargs) - axes = fn(adata, show=False, **common, **kwargs) - assert axes is not None + axes = fn(adata, groupby=groupby, use_rep=use_rep, cmap=cmap, figsize=figsize, **kwargs) fig = axes["heatmap_ax"].get_figure(root=True) assert fig is not None - return fig + return _finish(fig, axes, show=show, save=save, return_fig=return_fig) -def _legacy_heatmap(adata: AnnData, **kwargs) -> dict[str, Axes] | None: +def _legacy_heatmap(adata: AnnData, **kwargs) -> dict[str, Axes]: from infercnvpy.pl._chromosome_heatmap_legacy import legacy_heatmap return legacy_heatmap(adata, **kwargs) -def _legacy_summary(adata: AnnData, **kwargs) -> dict[str, Axes] | None: +def _legacy_summary(adata: AnnData, **kwargs) -> dict[str, Axes]: from infercnvpy.pl._chromosome_heatmap_legacy import legacy_summary return legacy_summary(adata, **kwargs) @@ -426,35 +440,61 @@ def _group_means(X: Any, codes: np.ndarray, n_groups: int) -> np.ndarray: def _order_by_dendrogram( - adata: AnnData, groupby: str, cat: pd.Categorical, palette: list, means: Callable[[], np.ndarray] + adata: AnnData, groupby: str, cat: pd.Categorical, palette: list, compute: Callable[[], dict[str, Any]] ) -> tuple[pd.Categorical, list, dict[str, Any] | None]: """Reorder the categories (and their colours) as the leaves of the dendrogram of the non-empty groups. The dendrogram is `adata.uns[f"dendrogram_{groupby}"]` from :func:`scanpy.tl.dendrogram` if that has exactly - these groups, as for `mode="legacy"`, otherwise it is computed from the groups' mean CNV profiles (`means()`). - Unused categories go last. With fewer than two groups, there is no dendrogram. + these groups, otherwise `compute()`, which is stored there as `mode="legacy"` does. Unused categories go last. + With fewer than two groups, there is no dendrogram. """ present = np.flatnonzero(np.bincount(cat.codes[cat.codes >= 0], minlength=len(cat.categories))) if len(present) < 2: return cat, palette, None labels = [str(c) for c in cat.categories[present]] - stored = adata.uns.get(f"dendrogram_{groupby}", {}).get("dendrogram_info") - if stored is not None and sorted(map(str, stored["ivl"])) == sorted(labels): - dendro = {**stored, "leaves": [labels.index(str(c)) for c in stored["ivl"]]} - else: - dendro = _dendrogram(means(), labels) + key = f"dendrogram_{groupby}" + dendro = adata.uns.get(key, {}).get("dendrogram_info") + if dendro is None or sorted(map(str, dendro["ivl"])) != sorted(labels): + stored = compute() + if not adata.is_view: + adata.uns[key] = stored + dendro = stored["dendrogram_info"] + dendro = {**dendro, "leaves": [labels.index(str(c)) for c in dendro["ivl"]]} new = [*present[dendro["leaves"]], *np.setdiff1d(np.arange(len(cat.categories)), present)] return cat.reorder_categories(cat.categories[new]), [palette[i] for i in new], dendro -def _dendrogram(means: np.ndarray, labels: list[str]) -> dict[str, Any]: - """`scipy.cluster.hierarchy.dendrogram` of the groups, computed like :func:`scanpy.tl.dendrogram`.""" - import scipy.cluster.hierarchy as sch - from scipy.spatial import distance +def _dendrogram(X: Any, groups: pd.Categorical, groupby: str, means: np.ndarray | None = None) -> dict[str, Any]: + """What :func:`scanpy.tl.dendrogram` with its defaults stores for the rows of `X` grouped by `groups` (as legacy). + + `means` are the means of the non-empty groups in the order of the categories, if already known. With more + windows than `scanpy.settings.N_PCS`, scanpy clusters the groups' means of a PCA of all rows (`_pca_group_means`). + """ + if means is None: + means = _group_means(X, groups.codes.astype(np.int64), len(groups.categories)) + if X.shape[1] > sc.settings.N_PCS: + means = _pca_group_means(X, means) + present = groups.remove_unused_categories().categories + obs = pd.DataFrame({groupby: pd.Categorical(present, categories=present)}) + obs.index = obs.index.astype(str) + return sc.tl.dendrogram(AnnData(means, obs=obs), groupby, inplace=False) - corr = np.clip(pd.DataFrame(means.T).corr().to_numpy(), -1, 1) - z = sch.linkage(distance.squareform(1 - corr, checks=False), method="complete") - return sch.dendrogram(z, labels=labels, no_plot=True) + +def _pca_group_means(X: Any, means: np.ndarray) -> np.ndarray: + """The groups' means (`means`, in windows) in the cells' PCA of :func:`scanpy.pp.pca` with its defaults. + + The PCA comes from the eigenvectors of the covariance, with the signs of :func:`sklearn.utils.extmath.svd_flip`. + """ + import scipy.linalg + + n, m = X.shape + n_comps = min(n, m) - 1 if min(n, m) <= sc.settings.N_PCS else sc.settings.N_PCS + G, sums = gram(X) + mu = sums / n + _, vecs = scipy.linalg.eigh((G - n * np.outer(mu, mu)) / (n - 1)) + comps = vecs[:, ::-1][:, :n_comps].T + comps *= np.sign(comps[np.arange(n_comps), np.abs(comps).argmax(axis=1)])[:, None] + return (means - mu) @ comps.T # --- drawing -------------------------------------------------------------------------------------------------------- @@ -601,31 +641,20 @@ def thin(renderer: RendererBase) -> None: def _thin_group_labels(ax: Axes, labels: list[str], centers: np.ndarray, rows: np.ndarray, *, swap: bool) -> None: """Leave out the tick labels of groups that would overlap the label of a larger group, at drawing time.""" axis = ax.xaxis if swap else ax.yaxis - formatter = _ThinnedLabels(labels) + formatter = FixedFormatter(labels) axis.set_major_formatter(formatter) def thin(renderer: RendererBase) -> None: proto = axis.get_major_ticks()[0].label1 # font and rotation of the tick labels pos = ax.transData.transform(np.column_stack([centers, centers]))[:, 0 if swap else 1] extent = np.array([_extent(proto, label, renderer, along_x=swap) for label in labels]) - formatter.shown = _keep_labels(pos, extent, rows) + formatter.seq = [ + label if keep else "" for label, keep in zip(labels, _keep_labels(pos, extent, rows), strict=True) + ] ax.add_artist(_BeforeDraw(thin)) -class _ThinnedLabels(Formatter): - """Fixed tick labels (as `set_ticks(ticks, labels)`), of which only those in `shown` are drawn.""" - - def __init__(self, labels: list[str]): - self.labels = labels - self.shown = np.ones(len(labels), dtype=bool) - - def __call__(self, x: float, pos: int | None = None) -> str: - if pos is None or pos >= len(self.labels) or not self.shown[pos]: - return "" - return self.labels[pos] - - class _BeforeDraw(Artist): """Calls `callback(renderer)` when its axes draws, before the tick labels and texts (it draws first).""" @@ -661,26 +690,36 @@ def _keep_labels(pos: np.ndarray, extent: np.ndarray, priority: np.ndarray) -> n return keep -def _finish( - fig: Figure, axes: dict[str, Axes], *, show: bool | None, save: str | bool | None, return_fig: bool | None -) -> Figure | dict[str, Axes] | None: - if sc.settings.autosave if save is None else save: - _save(fig, save) +def _finish[T]( + fig: Figure, + axes: T, + *, + show: bool | None, + save: str | bool | None, + return_fig: bool | None, + name: str = "heatmap", + own: bool = True, +) -> Figure | T | None: + """Save, show and return as scanpy's `savefig_or_show` does. + + The figure is saved as `/.`, with `ext` from the end of `save` if that is + `.svg`, `.pdf` or `.png`. A saved figure is closed unless it is returned or not `own` (drawn into a given `ax`). + """ + saved = sc.settings.autosave if save is None else isinstance(save, str) or save + if saved: + ext = sc.settings.file_format_figs + if isinstance(save, str): + if save.endswith((".svg", ".pdf", ".png")): + save, ext = save[:-4], save[-3:] + name += save + figdir = Path(sc.settings.figdir) + figdir.mkdir(parents=True, exist_ok=True) + fig.savefig(figdir / f"{name}{sc.settings.plot_suffix}.{ext}", bbox_inches="tight") if return_fig: return fig - if sc.settings.autoshow if show is None else show: + show = sc.settings.autoshow if show is None else show + if show: plt.show() - return None - return axes - - -def _save(fig: Figure, save: str | bool | None) -> None: - """Save as `/heatmap.`, the file name `scanpy.pl.heatmap` used.""" - name, ext = "heatmap", sc.settings.file_format_figs - if isinstance(save, str): - if (suffix := Path(save).suffix) in {".svg", ".pdf", ".png"}: - save, ext = save.removesuffix(suffix), suffix[1:] - name += save - figdir = Path(sc.settings.figdir) - figdir.mkdir(parents=True, exist_ok=True) - fig.savefig(figdir / f"{name}{sc.settings.plot_suffix}.{ext}", bbox_inches="tight") + if saved and own: + plt.close(fig) + return None if show else axes diff --git a/src/infercnvpy/pl/_chromosome_heatmap_legacy.py b/src/infercnvpy/pl/_chromosome_heatmap_legacy.py index 9251c8f..9257713 100644 --- a/src/infercnvpy/pl/_chromosome_heatmap_legacy.py +++ b/src/infercnvpy/pl/_chromosome_heatmap_legacy.py @@ -12,11 +12,6 @@ from infercnvpy._util import _as_df, _as_matrix -try: - from scanpy.plotting._utils import savefig_or_show -except ImportError: # scanpy >= 1.13 - from scanpy.plotting.legacy._utils import savefig_or_show - if TYPE_CHECKING: from matplotlib.axes import Axes from matplotlib.colors import Colormap @@ -29,21 +24,20 @@ def legacy_heatmap( use_rep: str, cmap: str | Colormap, figsize: tuple[float, float], - show: bool | None, - save: str | bool | None, **kwargs, -) -> dict[str, Axes] | None: +) -> dict[str, Axes]: """Every cell as a row of :func:`scanpy.pl.heatmap`, in the original order within each group.""" X = _as_matrix(adata.obsm[f"X_{use_rep}"], f"`adata.obsm['X_{use_rep}']`") - tmp_adata = AnnData(X=X, obs=_as_df(adata.obs), uns=adata.uns) - tmp_adata = tmp_adata[_stable_group_order(tmp_adata, groupby, kwargs.get("dendrogram", False))] + obs = _as_df(adata.obs) + # an actual AnnData (not a view) that shares `uns`, so the colours scanpy picks are stored in `adata.uns` + order = _stable_group_order(AnnData(X=X, obs=obs, uns=adata.uns), groupby, kwargs.get("dendrogram", False)) + tmp_adata = AnnData(X=X[order], obs=obs.iloc[order], uns=adata.uns) # re-sort, as saving & loading anndata destroys the order chr_pos_dict = dict(sorted(adata.uns[use_rep]["chr_pos"].items(), key=lambda x: x[1])) chr_pos = list(chr_pos_dict.values()) # center color map at 0 - X = _as_matrix(tmp_adata.X) tmp_data = X if isinstance(X, np.ndarray) else X.data vmin = kwargs.pop("vmin", None) vmax = kwargs.pop("vmax", None) @@ -66,16 +60,12 @@ def legacy_heatmap( var_group_positions=var_group_positions, var_group_labels=list(chr_pos_dict.keys()), show=False, + save=False, **kwargs, ) return_ax_dic["heatmap_ax"].vlines(chr_pos[1:], lw=0.6, ymin=0, ymax=tmp_adata.shape[0]) - - savefig_or_show("heatmap", show=show, save=save) - show = sc.settings.autoshow if show is None else show - if not show: - return return_ax_dic - return None + return return_ax_dic def _stable_group_order(adata: AnnData, groupby: str, dendrogram: bool | str) -> np.ndarray: @@ -104,10 +94,8 @@ def legacy_summary( use_rep: str, cmap: str | Colormap, figsize: tuple[float, float], - show: bool | None, - save: str | bool | None, **kwargs, -) -> dict[str, Axes] | None: +) -> dict[str, Axes]: """Every group mean as 10 rows of :func:`scanpy.pl.heatmap`.""" # this dirty hack repeats each row 10 times, since scanpy # heatmap cannot really handle it if there's just one observation @@ -155,13 +143,9 @@ def _get_group_mean(group): var_group_positions=var_group_positions, var_group_labels=list(chr_pos_dict.keys()), show=False, + save=False, **kwargs, ) return_ax_dic["heatmap_ax"].vlines(chr_pos[1:], lw=0.6, ymin=-1, ymax=tmp_adata.shape[0]) - - savefig_or_show("heatmap", show=show, save=save) - show = sc.settings.autoshow if show is None else show - if not show: - return return_ax_dic - return None + return return_ax_dic diff --git a/src/infercnvpy/pl/_heatmap_kernels.py b/src/infercnvpy/pl/_heatmap_kernels.py index 5c87109..fc0c3ce 100644 --- a/src/infercnvpy/pl/_heatmap_kernels.py +++ b/src/infercnvpy/pl/_heatmap_kernels.py @@ -114,6 +114,37 @@ def add(Xb, row0, n, start): return sums +def column_ranges(X: Any) -> tuple[np.ndarray, np.ndarray]: + """Minimum and maximum of each column of `X` over all rows (with the zeros of a sparse matrix), NaNs ignored.""" + lo, hi = np.full(X.shape[1], np.inf), np.full(X.shape[1], -np.inf) + + def update(Xb, row0, n, start): + rows = Xb[row0 : row0 + n] + if sp.issparse(rows): + b_lo, b_hi = rows.nanmin(axis=0).toarray().ravel(), rows.nanmax(axis=0).toarray().ravel() + else: + b_lo, b_hi = np.fmin.reduce(rows, axis=0), np.fmax.reduce(rows, axis=0) + np.fmin(lo, b_lo, out=lo) + np.fmax(hi, b_hi, out=hi) + + _each_block(_blocks(X), update) + return lo, hi + + +def gram(X: Any) -> tuple[np.ndarray, np.ndarray]: + """`X.T @ X` and the column sums of `X`, in float64, `CHUNKSIZE` rows at a time.""" + G, sums = np.zeros((X.shape[1], X.shape[1])), np.zeros(X.shape[1]) + + def add(Xb, row0, n, start): + for s in range(row0, row0 + n, CHUNKSIZE): + rows = Xb[s : min(s + CHUNKSIZE, row0 + n)].astype(np.float64) + G[:] += (rows.T @ rows).toarray() if sp.issparse(rows) else rows.T @ rows + sums[:] += np.asarray(rows.sum(axis=0)).ravel() + + _each_block(_blocks(X), add) + return G, sums + + def sampled_rows(X: Any, picks: np.ndarray) -> tuple[np.ndarray, float, float]: """Rows `picks` of `X` as a dense float64 array, and the min and max of all values of `X`. diff --git a/src/infercnvpy/pl/_spatial.py b/src/infercnvpy/pl/_spatial.py index e38c39f..4aeb66a 100644 --- a/src/infercnvpy/pl/_spatial.py +++ b/src/infercnvpy/pl/_spatial.py @@ -6,10 +6,9 @@ import numpy as np import scanpy as sc from anndata import AnnData -from matplotlib.figure import Figure from infercnvpy._util import _as_df, _as_matrix -from infercnvpy.pl._chromosome_heatmap import CBAR_LABEL +from infercnvpy.pl._chromosome_heatmap import CBAR_LABEL, _finish _REGION = re.compile(r"^(?P[^:]+)(?::(?P[\d_,]+)-(?P[\d_,]+))?$") @@ -53,7 +52,8 @@ def spatial_cnv( Thin wrapper around :func:`scanpy.pl.embedding` with `basis="spatial"`. Each region becomes a panel coloured by the mean of `adata.obsm["X_{use_rep}"]` over its windows, with a diverging colour map centred at 0 (pass `vcenter=None` to turn that off). - `adata.obsm["X_{use_rep}"]` must be in memory. + As :func:`scanpy.pl.spatial` without an image, the panels have an equal aspect and the + y axis points down. `adata.obsm["X_{use_rep}"]` must be in memory. Parameters ---------- @@ -74,11 +74,13 @@ def spatial_cnv( cmap Colour map of the regions (and other continuous panels). **kwargs - Passed to :func:`scanpy.pl.embedding`, e.g. `size`, `ncols`, `vmin`/`vmax` or `show`. + Passed to :func:`scanpy.pl.embedding`, e.g. `size`, `ncols`, `vmin`/`vmax`, `show`, + `save` (a string is appended to the file name `basis`) or `return_fig`. Returns ------- - Whatever :func:`scanpy.pl.embedding` returns: `None`, or the axes if `show=False`. + The figure if `return_fig` is true. Otherwise, if `show` is false, the axes as + :func:`scanpy.pl.embedding` returns them. """ regions = [region] if isinstance(region, str) else list(region or []) colors = [color] if isinstance(color, str) else list(color or []) @@ -91,20 +93,18 @@ def spatial_cnv( obs[f"{use_rep} {r}"] = np.asarray(X[:, _region_columns(adata, use_rep, r)].mean(axis=1)).ravel() tmp = AnnData(obs=obs, uns=adata.uns) tmp.obsm[basis] = np.asarray(adata.obsm[basis]) - if regions: - kwargs.setdefault("vcenter", 0) - show, save = kwargs.pop("show", None), kwargs.pop("save", None) - titles = [f"{use_rep} {r}" for r in regions] - res = sc.pl.embedding(tmp, basis, color=titles + colors, cmap=cmap, show=False, **kwargs) - fig = res if isinstance(res, Figure) else np.ravel(res)[0].get_figure(root=True) - for ax in fig.axes: - if ax.get_title() in titles: - for points in ax.collections: - if points.colorbar is not None: - points.colorbar.set_label(f"mean {CBAR_LABEL}") - if isinstance(res, Figure): # `return_fig=True`: scanpy neither saves nor shows - return res - from infercnvpy.pl._chromosome_heatmap_legacy import savefig_or_show - - savefig_or_show(basis, show=show, save=save) - return None if (sc.settings.autoshow if show is None else show) else res + kwargs.setdefault("vcenter", [0] * len(regions) + [None] * len(colors)) + show, save, return_fig = (kwargs.pop(k, None) for k in ("show", "save", "return_fig")) + color = [f"{use_rep} {r}" for r in regions] + colors + res = sc.pl.embedding(tmp, basis, color=color, cmap=cmap, show=False, save=False, **kwargs) + panels = res if isinstance(res, list) else [res] + for ax in panels: + ax.set_aspect("equal") + ax.invert_yaxis() + for ax in panels[: len(regions)]: + for points in ax.collections: + if points.colorbar is not None: + points.colorbar.set_label(f"mean {CBAR_LABEL}") + fig = panels[0].get_figure(root=True) + assert fig is not None + return _finish(fig, res, show=show, save=save, return_fig=return_fig, name=basis, own=kwargs.get("ax") is None) diff --git a/tests/test_plotting.py b/tests/test_plotting.py index d37c568..f5eddb3 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -1,3 +1,5 @@ +import warnings + import dask import dask.array as da import numpy as np @@ -5,7 +7,7 @@ import pytest import scanpy as sc import scipy.sparse as sp -from anndata import AnnData +from anndata import AnnData, ImplicitModificationWarning from matplotlib import pyplot as plt from matplotlib.colors import Normalize from matplotlib.figure import Figure @@ -155,12 +157,96 @@ def test_dendrogram_from_mean_profiles(adata_small): tmp.obs["group"] = tmp.obs["group"].cat.remove_unused_categories() expected = sc.tl.dendrogram(tmp, "group", var_names=tmp.var_names, inplace=False)["categories_ordered"] assert [t.get_text() for t in axes["groupby_ax"].get_yticklabels()] == expected + # legacy clusters the group means in the order of their names, which can flip branches + tmp.obsm["X_cnv"], tmp.uns["cnv"] = tmp.X, adata_small.uns["cnv"] + legacy = cnv.pl.chromosome_heatmap_summary(tmp, groupby="group", dendrogram=True, show=False, mode="legacy") + expected = [t.get_text() for t in legacy["groupby_ax"].get_yticklabels()] + del adata_small.uns["dendrogram_group"] # stored by the call above, and used by later ones summary = cnv.pl.chromosome_heatmap_summary(adata_small, groupby="group", dendrogram=True, show=False) assert [t.get_text() for t in summary["groupby_ax"].get_yticklabels()] == expected X, g = adata_small.obsm["X_cnv"].toarray(), adata_small.obs["group"] np.testing.assert_allclose(_image(summary), [X[(g == c).to_numpy()].mean(axis=0) for c in expected], rtol=1e-6) +def _many_windows(seed: int) -> AnnData: + """400 cells in 8 groups x 300 windows: more windows than scanpy's `N_PCS`, so its dendrogram uses a PCA.""" + rng = np.random.default_rng(seed) + profiles = rng.normal(size=(8, 300)) * (rng.random((8, 300)) < 0.3) + labels = rng.integers(0, 8, 400) + X = (profiles[labels] + rng.normal(scale=0.7, size=(400, 300))).astype(np.float32) + groups = pd.Categorical([f"G{i}" for i in labels], categories=[f"G{i}" for i in (3, 0, 6, 1, 7, 4, 2, 5)]) + obs = pd.DataFrame({"g": groups}, index=[str(i) for i in range(400)]) + adata = AnnData(obs=obs, uns={"cnv": {"chr_pos": {"chr1": 0, "chr2": 150}}}) + adata.obsm["X_cnv"] = X + return adata + + +@pytest.mark.parametrize("fmt", [np.asarray, sp.csr_matrix, _dask, lambda X: _dask(sp.csr_matrix(X), 50)]) +@pytest.mark.parametrize("seed", [0, 1, 2]) +@both_plots +def test_dendrogram_like_legacy_with_many_windows(fmt, seed, plot): + legacy = plot(_many_windows(seed), groupby="g", dendrogram=True, show=False, mode="legacy") + expected = [t.get_text() for t in legacy["groupby_ax"].get_yticklabels()] + adata = _many_windows(seed) + adata.obsm["X_cnv"] = fmt(adata.obsm["X_cnv"]) + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) # not scanpy's warning that it falls back to a PCA + fast = plot(adata, groupby="g", dendrogram=True, show=False) + assert [t.get_text() for t in fast["groupby_ax"].get_yticklabels()] == expected + + +@both_plots +def test_dendrogram_is_stored_and_reused(plot, monkeypatch): + """The computed dendrogram is stored as `scanpy.tl.dendrogram` stores it, and later calls (also legacy) use it.""" + adata = _many_windows(0) + first = plot(adata, groupby="g", dendrogram=True, show=False) + order = [t.get_text() for t in first["groupby_ax"].get_yticklabels()] + stored = adata.uns["dendrogram_g"] + tmp = AnnData( + np.random.default_rng(0).normal(size=(8, 3)), + obs=pd.DataFrame({"g": adata.obs["g"].cat.categories.astype("category")}), + ) + assert set(stored) == set(sc.tl.dendrogram(tmp, "g", inplace=False)) + assert stored["groupby"] == ["g"] + assert stored["categories_ordered"] == order + + def fail(*args, **kwargs): + raise AssertionError("recomputed") + + monkeypatch.setattr("infercnvpy.pl._chromosome_heatmap._dendrogram", fail) + monkeypatch.setattr(sc.tl, "dendrogram", fail) + for mode in ("fast", "legacy"): + again = plot(adata, groupby="g", dendrogram=True, show=False, mode=mode) + assert [t.get_text() for t in again["groupby_ax"].get_yticklabels()] == order + + +@pytest.mark.parametrize("kwargs", [{}, {"log": True}, {"aggregate": "mean"}]) +def test_standard_scale_var_over_all_cells(adata_small, kwargs): + """Each window is scaled with its minimum and maximum over all cells, not only the drawn ones.""" + if kwargs.get("log"): + adata_small.obsm["X_cnv"] *= 0.2 # log1p needs values above -1 + axes = cnv.pl.chromosome_heatmap( + adata_small, groupby="group", figsize=(4, 0.2), standard_scale="var", show=False, **kwargs + ) + values = _image(axes) + assert len(values) < 59 + raw = _image(cnv.pl.chromosome_heatmap(adata_small, groupby="group", figsize=(4, 0.2), show=False, **kwargs)) + X = adata_small.obsm["X_cnv"].toarray() + X = np.log1p(X) if kwargs.get("log") else X + lo, hi = X.min(axis=0), X.max(axis=0) + np.testing.assert_allclose(values, (raw - lo) / (hi - lo), rtol=1e-6, atol=1e-7) + + +def test_legacy_stores_group_colors(adata_small): + """scanpy gets an actual AnnData that shares `uns`, so the colours it picks land in the caller's `uns`.""" + adata_small = adata_small[adata_small.obs["group"].notna()].copy() + del adata_small.uns["group_colors"] + with warnings.catch_warnings(): + warnings.simplefilter("error", ImplicitModificationWarning) + cnv.pl.chromosome_heatmap(adata_small, groupby="group", show=False, mode="legacy") + assert "group_colors" in adata_small.uns + + @both_plots def test_unsupported_argument_falls_back_to_legacy(adata_small, plot): adata_small = adata_small[adata_small.obs["group"].notna()].copy() @@ -308,6 +394,46 @@ def test_save(adata_small, tmp_path): assert sorted(p.name for p in tmp_path.iterdir()) == sorted([f"heatmap.{ext}", "heatmap_summary.png"]) +@pytest.mark.parametrize("mode", ["fast", "legacy"]) +@both_plots +def test_save_like_scanpy(adata_small, tmp_path, plot, mode): + """`save` names the file as `scanpy.pl.heatmap` does, and the figure is closed after saving unless returned.""" + adata_small = adata_small[adata_small.obs["group"].notna()].copy() + figdir, fmt = sc.settings.figdir, sc.settings.file_format_figs + sc.settings.figdir, sc.settings.file_format_figs = tmp_path, "pdf" + try: + plot(adata_small, groupby="group", show=False, save=".png", mode=mode) + plot(adata_small, groupby="group", show=False, save="", mode=mode) + assert not plt.get_fignums() + fig = plot(adata_small, groupby="group", save="_fig.svg", return_fig=True, mode=mode) + assert plt.fignum_exists(fig.number) + if mode == "fast": + _, ax = plt.subplots() + plot(adata_small, groupby="group", ax=ax, show=False, save="_ax") + assert plt.fignum_exists(ax.get_figure(root=True).number) + finally: + sc.settings.figdir, sc.settings.file_format_figs = figdir, fmt + expected = {"heatmap.png", "heatmap.pdf", "heatmap_fig.svg"} | ({"heatmap_ax.pdf"} if mode == "fast" else set()) + assert {p.name for p in tmp_path.iterdir()} == expected + + +@pytest.mark.parametrize("mode", ["fast", "legacy"]) +def test_autosave(adata_small, tmp_path, mode): + """With `scanpy.settings.autosave`, the heatmap is saved once (scanpy's own save does not leave a blank one).""" + adata_small = adata_small[adata_small.obs["group"].notna()].copy() + settings = sc.settings.figdir, sc.settings.file_format_figs, sc.settings.autosave + sc.settings.figdir, sc.settings.file_format_figs, sc.settings.autosave = tmp_path, "png", True + try: + with warnings.catch_warnings(): + warnings.simplefilter("error", FutureWarning) + cnv.pl.chromosome_heatmap(adata_small, groupby="group", show=False, mode=mode) + finally: + sc.settings.figdir, sc.settings.file_format_figs, sc.settings.autosave = settings + assert [p.name for p in tmp_path.iterdir()] == ["heatmap.png"] + assert (plt.imread(tmp_path / "heatmap.png")[..., :3] < 0.9).any() + assert not plt.get_fignums() + + def test_bad_arguments(adata_small): with pytest.raises(ValueError, match="tl.leiden"): cnv.pl.chromosome_heatmap(adata_small) diff --git a/tests/test_spatial.py b/tests/test_spatial.py index c8ac28b..b4578c0 100644 --- a/tests/test_spatial.py +++ b/tests/test_spatial.py @@ -1,10 +1,13 @@ import logging +import warnings import numpy as np import pandas as pd import pytest +import scanpy as sc import scipy.sparse as sp from anndata import AnnData +from matplotlib import pyplot as plt from scipy.spatial import cKDTree import infercnvpy as cnv @@ -233,3 +236,57 @@ def test_spatial_cnv_colorbar_label(): axes = cnv.pl.spatial_cnv(adata, region="chr1", color="label", show=False) [cbar] = [c.colorbar for ax in axes for c in ax.collections if c.colorbar is not None] assert cbar.ax.get_ylabel() == "mean CNV (relative to reference)" + + +@pytest.fixture +def adata_spatial() -> AnnData: + """30 spots on a tissue four times as wide as high, with a categorical and a continuous `obs` column.""" + rng = np.random.default_rng(0) + n = 30 + obs = pd.DataFrame( + {"label": pd.Categorical(rng.choice(["x", "y"], n)), "score": rng.uniform(100, 200, n)}, + index=[str(i) for i in range(n)], + ) + adata = AnnData(obs=obs, uns={"cnv": {"chr_pos": {"chr1": 0, "chr2": 2}}}) + adata.obsm["X_cnv"] = rng.normal(size=(n, 4)).astype(np.float32) + adata.obsm["spatial"] = np.column_stack([rng.uniform(0, 4000, n), rng.uniform(0, 1000, n)]) + return adata + + +def test_spatial_cnv_panels(adata_spatial): + """Region panels are found by position (whatever their title) and centred at 0; other panels are not.""" + axes = cnv.pl.spatial_cnv( + adata_spatial, region=["chr1", "chr2"], color=["label", "score"], title=["A", "B", "C", "D"], show=False + ) + assert [ax.get_title() for ax in axes] == ["A", "B", "C", "D"] + for ax in axes[:2]: + [points] = ax.collections + assert points.norm.vcenter == 0 + assert points.colorbar.ax.get_ylabel() == "mean CNV (relative to reference)" + [score] = axes[3].collections + score_range = adata_spatial.obs["score"].min(), adata_spatial.obs["score"].max() + assert (score.norm.vmin, score.norm.vmax) == pytest.approx(score_range) + assert score.colorbar.ax.get_ylabel() == "" + for ax in axes: # as `scanpy.pl.spatial` draws without an image + assert ax.get_aspect() == 1 + assert ax.yaxis_inverted() + + +def test_spatial_cnv_saves_once(adata_spatial, tmp_path): + plt.close("all") + settings = sc.settings.figdir, sc.settings.file_format_figs, sc.settings.autosave + sc.settings.figdir, sc.settings.file_format_figs = tmp_path, "png" + try: + sc.settings.autosave = True + cnv.pl.spatial_cnv(adata_spatial, region="chr1", show=False) + sc.settings.autosave = False + with warnings.catch_warnings(): + warnings.simplefilter("error", FutureWarning) + cnv.pl.spatial_cnv(adata_spatial, region="chr1", show=False, save="_cnv.png") + fig = cnv.pl.spatial_cnv(adata_spatial, region="chr1", return_fig=True, save="_fig.png") + finally: + sc.settings.figdir, sc.settings.file_format_figs, sc.settings.autosave = settings + assert sorted(p.name for p in tmp_path.iterdir()) == ["spatial.png", "spatial_cnv.png", "spatial_fig.png"] + for p in tmp_path.iterdir(): + assert (plt.imread(p)[..., :3] < 0.9).any(), p.name + assert plt.get_fignums() == [fig.number]