def composition( adata: AnnData, key: str, fontsize: Optional[str] = None, figsize: Optional[Tuple[float, float]] = None, dpi: Optional[float] = None, save: Optional[Union[str, Path]] = None, **kwargs, ) -> None: """ Plot a pie chart for categorical annotation. .. image:: https://raw.githubusercontent.com/theislab/cellrank/master/resources/images/composition.png :width: 400px :align: center Parameters ---------- %(adata)s key Key in ``adata.obs`` containing categorical observation. fontsize Font size for the pie chart labels. %(plotting)s **kwargs Keyworded arguments for :func:`matplotlib.pyplot.pie`. Returns ------- %(just_plots)s """ if key not in adata.obs: raise KeyError(f"Key `{key!r}` not found in `adata.obs`.") if not is_categorical_dtype(adata.obs[key]): raise TypeError( f"Observation `adata.obs[{key!r}]` is not categorical.") cats = adata.obs[key].cat.categories colors = adata.uns.get(f"{key}_colors", None) x = [np.sum(adata.obs[key] == cl) for cl in cats] cats_frac = x / np.sum(x) # plot these fractions in a pie plot fig, ax = plt.subplots(figsize=figsize, dpi=dpi) ax.pie( x=cats_frac, labels=cats, colors=colors, textprops={"fontsize": fontsize}, **kwargs, ) ax.set_title(f"composition by {key}") if save is not None: save_fig(fig, save) fig.show()
def plot_spectrum( self, n: Optional[int] = None, real_only: bool = False, show_eigengap: bool = True, show_all_xticks: bool = True, legend_loc: Optional[str] = None, title: Optional[str] = None, figsize: Optional[Tuple[float, float]] = (5, 5), dpi: int = 100, save: Optional[Union[str, Path]] = None, marker: str = ".", **kwargs, ) -> None: """ Plot the top eigenvalues in real or complex plane. Parameters ---------- n Number of eigenvalues to show. If `None`, show all that have been computed. real_only Whether to plot only the real part of the spectrum. show_eigengap When `real_only=True`, this determines whether to show the inferred eigengap as a dotted line. show_all_xticks When `real_only=True`, this determines whether to show the indices of all eigenvalues on the x-axis. legend_loc Location parameter for the legend. title Title of the figure. %(plotting)s marker Marker symbol used, valid options can be found in :mod:`matplotlib.markers`. **kwargs Keyword arguments for :func:`matplotlib.pyplot.scatter`. Returns ------- %(just_plots)s """ eig = getattr(self, P.EIG.s) if eig is None: raise RuntimeError( f"Compute `.{P.EIG}` first as `.{F.COMPUTE.fmt(P.EIG)}()`.") if n is None: n = len(eig["D"]) elif n <= 0: raise ValueError(f"Expected `n` to be > 0, found `{n}`.") if real_only: fig = self._plot_real_spectrum( n, show_eigengap=show_eigengap, show_all_xticks=show_all_xticks, dpi=dpi, figsize=figsize, legend_loc=legend_loc, title=title, marker=marker, **kwargs, ) else: fig = self._plot_complex_spectrum( n, dpi=dpi, figsize=figsize, legend_loc=legend_loc, title=title, marker=marker, **kwargs, ) if save: save_fig(fig, save) fig.show()
def plot_macrostate_composition( self, key: str, width: float = 0.8, title: Optional[str] = None, labelrot: float = 45, legend_loc: Optional[str] = "upper right out", figsize: Optional[Tuple[float, float]] = None, dpi: Optional[int] = None, save: Optional[Union[str, Path]] = None, show: bool = True, ) -> Optional[Axes]: """ Plot stacked histogram of macrostates over categorical annotations. Parameters ---------- %(adata)s key Key from :attr:`adata` ``.obs`` containing categorical annotations. width Bar width in `[0, 1]`. title Title of the figure. If `None`, create one automatically. labelrot Rotation of labels on x-axis. legend_loc Position of the legend. If `None`, don't show legend. %(plotting)s show If `False`, return :class:`matplotlib.pyplot.Axes`. Returns ------- :class:`matplotlib.pyplot.Axes` The axis object if ``show=False``. %(just_plots)s """ from cellrank.pl._utils import _position_legend macrostates = self._get(P.MACRO) if macrostates is None: raise RuntimeError( "Compute macrostates first as `.compute_macrostates()`.") if key not in self.adata.obs: raise KeyError(f"Key `{key}` not found in `adata.obs`.") if not is_categorical_dtype(self.adata.obs[key]): raise TypeError( f"Expected `adata.obs[{key!r}]` to be `categorical`, " f"found `{infer_dtype(self.adata.obs[key])}`.") mask = ~macrostates.isnull() df = (pd.DataFrame({ "macrostates": macrostates, key: self.adata.obs[key] })[mask].groupby([key, "macrostates"]).size()) try: cats_colors = self.adata.uns[f"{key}_colors"] except KeyError: cats_colors = _create_categorical_colors( len(self.adata.obs[key].cat.categories)) cat_color_mapper = dict( zip(self.adata.obs[key].cat.categories, cats_colors)) x_indices = np.arange(len(macrostates.cat.categories)) bottom = np.zeros_like(x_indices, dtype=np.float32) width = min(1, max(0, width)) fig, ax = plt.subplots(figsize=figsize, dpi=dpi, tight_layout=True) for cat, color in cat_color_mapper.items(): frequencies = df.loc[cat] # do not add to legend if category is missing if np.sum(frequencies) > 0: ax.bar( x_indices, frequencies, width, label=cat, color=color, bottom=bottom, ec="black", lw=0.5, ) bottom += np.array(frequencies) ax.set_xticks(x_indices) ax.set_xticklabels( # assuming at least 1 category frequencies.index, rotation=labelrot, ha="center" if labelrot in (0, 90) else "right", ) y_max = bottom.max() ax.set_ylim([0, y_max + 0.05 * y_max]) ax.set_yticks(np.linspace(0, y_max, 5)) ax.margins(0.05) ax.set_xlabel("macrostate") ax.set_ylabel("frequency") if title is None: title = f"distribution over {key}" ax.set_title(title) if legend_loc not in (None, "none"): _position_legend(ax, legend_loc=legend_loc) if save is not None: save_fig(fig, save) if not show: return ax
def plot_coarse_T( self, show_stationary_dist: bool = True, show_initial_dist: bool = False, cmap: Union[str, mcolors.ListedColormap] = "viridis", xtick_rotation: float = 45, annotate: bool = True, show_cbar: bool = True, title: Optional[str] = None, figsize: Tuple[float, float] = (8, 8), dpi: int = 80, save: Optional[Union[Path, str]] = None, text_kwargs: Mapping[str, Any] = MappingProxyType({}), **kwargs, ) -> None: """ Plot the coarse-grained transition matrix between macrostates. Parameters ---------- show_stationary_dist Whether to show the stationary distribution, if present. show_initial_dist Whether to show the initial distribution. cmap Colormap to use. xtick_rotation Rotation of ticks on the x-axis. annotate Whether to display the text on each cell. show_cbar Whether to show colorbar. title Title of the figure. %(plotting)s text_kwargs Keyword arguments for :func:`matplotlib.pyplot.text`. kwargs Keyword arguments for :func:`matplotlib.pyplot.imshow`. Returns ------- %(just_plots)s """ def stylize_dist(ax, data: np.ndarray, xticks_labels: Union[List[str], Tuple[str]] = ()): _ = ax.imshow(data, aspect="auto", cmap=cmap, norm=norm) for spine in ax.spines.values(): spine.set_visible(False) if xticks_labels is not None: ax.set_xticks(np.arange(data.shape[1])) ax.set_xticklabels(xticks_labels) plt.setp( ax.get_xticklabels(), rotation=xtick_rotation, ha="right", rotation_mode="anchor", ) else: ax.set_xticks([]) ax.tick_params(which="both", top=False, right=False, bottom=False, left=False) ax.set_yticks([]) def annotate_heatmap(im, valfmt: str = "{x:.2f}"): # modified from matplotlib's site data = im.get_array() kw = {"ha": "center", "va": "center"} kw.update(**text_kwargs) # Get the formatter in case a string is supplied if isinstance(valfmt, str): valfmt = mpl.ticker.StrMethodFormatter(valfmt) # Loop over the data and create a `Text` for each "pixel". # Change the text's color depending on the data. texts = [] for i in range(data.shape[0]): for j in range(data.shape[1]): kw.update( color=_get_black_or_white(im.norm(data[i, j]), cmap)) text = im.axes.text(j, i, valfmt(data[i, j], None), **kw) texts.append(text) def annotate_dist_ax(ax, data: np.ndarray, valfmt: str = "{x:.2f}"): if ax is None: return if isinstance(valfmt, str): valfmt = mpl.ticker.StrMethodFormatter(valfmt) kw = {"ha": "center", "va": "center"} kw.update(**text_kwargs) for i, val in enumerate(data): kw.update(color=_get_black_or_white(im.norm(val), cmap)) ax.text( i, 0, valfmt(val, None), **kw, ) coarse_T = self._get(P.COARSE_T) coarse_stat_d = self._get(P.COARSE_STAT_D) coarse_init_d = self._get(P.COARSE_INIT_D) if coarse_T is None: raise RuntimeError( "Compute coarse-grained transition matrix first as `.compute_macrostates()` with `n_states > 1`." ) if show_stationary_dist and coarse_stat_d is None: logg.warning("Coarse stationary distribution is `None`, ignoring") show_stationary_dist = False if show_initial_dist and coarse_init_d is None: logg.warning("Coarse initial distribution is `None`, ignoring") show_initial_dist = False hrs, wrs = [1], [1] if show_stationary_dist: hrs += [0.05] if show_initial_dist: hrs += [0.05] if show_cbar: wrs += [0.025] dont_show_dist = not show_initial_dist and not show_stationary_dist fig = plt.figure(constrained_layout=False, figsize=figsize, dpi=dpi) gs = plt.GridSpec( 1 + show_stationary_dist + show_initial_dist, 1 + show_cbar, height_ratios=hrs, width_ratios=wrs, wspace=0.05, hspace=0.05, ) if isinstance(cmap, str): cmap = plt.get_cmap(cmap) ax = fig.add_subplot(gs[0, 0]) cax = fig.add_subplot(gs[:1, -1]) if show_cbar else None init_ax, stat_ax = None, None labels = list(self.coarse_T.columns) tmp = coarse_T if show_initial_dist: tmp = np.c_[tmp, coarse_stat_d] if show_initial_dist: tmp = np.c_[tmp, coarse_init_d] minn, maxx = np.nanmin(tmp), np.nanmax(tmp) norm = mpl.colors.Normalize(vmin=minn, vmax=maxx) if show_stationary_dist: stat_ax = fig.add_subplot(gs[1, 0]) stylize_dist( stat_ax, np.array(coarse_stat_d).reshape(1, -1), xticks_labels=labels if not show_initial_dist else None, ) stat_ax.yaxis.set_label_position("right") stat_ax.set_ylabel("stationary dist", rotation=0, ha="left", va="center") if show_initial_dist: init_ax = fig.add_subplot(gs[show_stationary_dist + show_initial_dist, 0]) stylize_dist(init_ax, np.array(coarse_init_d).reshape(1, -1), xticks_labels=labels) init_ax.yaxis.set_label_position("right") init_ax.set_ylabel("initial dist", rotation=0, ha="left", va="center") im = ax.imshow(coarse_T, aspect="auto", cmap=cmap, norm=norm, **kwargs) ax.set_title( "coarse-grained transition matrix" if title is None else title) if cax is not None: _ = mpl.colorbar.ColorbarBase( cax, cmap=cmap, norm=norm, ticks=np.linspace(minn, maxx, 10), format="%0.3f", ) ax.set_yticks(np.arange(coarse_T.shape[0])) ax.set_yticklabels(labels) ax.tick_params( top=False, bottom=dont_show_dist, labeltop=False, labelbottom=dont_show_dist, ) for spine in ax.spines.values(): spine.set_visible(False) if dont_show_dist: ax.set_xticks(np.arange(coarse_T.shape[1])) ax.set_xticklabels(labels) plt.setp( ax.get_xticklabels(), rotation=xtick_rotation, ha="right", rotation_mode="anchor", ) else: ax.set_xticks([]) ax.set_yticks(np.arange(coarse_T.shape[0] + 1) - 0.5, minor=True) ax.tick_params(which="minor", bottom=dont_show_dist, left=False, top=False) if annotate: annotate_heatmap(im) if show_stationary_dist: annotate_dist_ax(stat_ax, coarse_stat_d.values) if show_initial_dist: annotate_dist_ax(init_ax, coarse_init_d) if save: save_fig(fig, save)
def gene_trends( adata: AnnData, model: _input_model_type, genes: Union[str, Sequence[str]], lineages: Optional[Union[str, Sequence[str]]] = None, backward: bool = False, data_key: str = "X", time_key: str = "latent_time", transpose: bool = False, time_range: Optional[Union[_time_range_type, List[_time_range_type]]] = None, callback: _callback_type = None, conf_int: Union[bool, float] = True, same_plot: bool = False, hide_cells: bool = False, perc: Optional[Union[Tuple[float, float], Sequence[Tuple[float, float]]]] = None, lineage_cmap: Optional[matplotlib.colors.ListedColormap] = None, abs_prob_cmap: matplotlib.colors.ListedColormap = cm.viridis, cell_color: Optional[str] = None, cell_alpha: float = 0.6, lineage_alpha: float = 0.2, size: float = 15, lw: float = 2, cbar: bool = True, margins: float = 0.015, sharex: Optional[Union[str, bool]] = None, sharey: Optional[Union[str, bool]] = None, gene_as_title: Optional[bool] = None, legend_loc: Optional[str] = "best", obs_legend_loc: Optional[str] = "best", ncols: int = 2, suptitle: Optional[str] = None, return_models: bool = False, n_jobs: Optional[int] = 1, backend: str = _DEFAULT_BACKEND, show_progress_bar: bool = True, figsize: Optional[Tuple[float, float]] = None, dpi: Optional[int] = None, save: Optional[Union[str, Path]] = None, plot_kwargs: Mapping = MappingProxyType({}), **kwargs, ) -> Optional[_return_model_type]: """ Plot gene expression trends along lineages. Each lineage is defined via it's lineage weights which we compute using :func:`cellrank.tl.lineages`. This function accepts any model based off :class:`cellrank.ul.models.BaseModel` to fit gene expression, where we take the lineage weights into account in the loss function. Parameters ---------- %(adata)s %(model)s %(genes)s lineages Names of the lineages to plot. If `None`, plot all lineages. %(backward)s data_key Key in ``adata.layers`` or `'X'` for ``adata.X`` where the data is stored. time_key Key in ``adata.obs`` where the pseudotime is stored. %(time_ranges)s transpose If ``same_plot=True``, group the trends by ``lineages`` instead of ``genes``. This enforces ``hide_cells=True``. If ``same_plot=False``, show ``lineages`` in rows and ``genes`` in columns. %(model_callback)s conf_int Whether to compute and show confidence interval. If the :paramref:`model` is :class:`cellrank.ul.models.GAMR`, it can also specify the confidence level, the default is `0.95`. same_plot Whether to plot all lineages for each gene in the same plot. hide_cells If `True`, hide all cells. perc Percentile for colors. Valid values are in interval `[0, 100]`. This can improve visualization. Can be specified individually for each lineage. lineage_cmap Categorical colormap to use when coloring in the lineages. If `None` and ``same_plot``, use the corresponding colors in ``adata.uns``, otherwise use `'black'`. abs_prob_cmap Continuous colormap to use when visualizing the absorption probabilities for each lineage. Only used when ``same_plot=False``. cell_color Key in :attr:`anndata.AnnData.obs` or :attr:`anndata.AnnData.var_names` used for coloring the cells. cell_alpha Alpha channel for cells. lineage_alpha Alpha channel for lineage confidence intervals. size Size of the points. lw Line width of the smoothed values. cbar Whether to show colorbar. Always shown when percentiles for lineages differ. Only used when ``same_plot=False``. margins Margins around the plot. sharex Whether to share x-axis. Valid options are `'row'`, `'col'` or `'none'`. sharey Whether to share y-axis. Valid options are `'row'`, `'col'` or `'none'`. gene_as_title Whether to show gene names as titles instead on y-axis. legend_loc Location of the legend displaying lineages. Only used when `same_plot=True`. obs_legend_loc Location of the legend when ``cell_color`` corresponds to a categorical variable. ncols Number of columns of the plot when plotting multiple genes. Only used when ``same_plot=True``. suptitle Suptitle of the figure. %(return_models)s %(parallel)s %(plotting)s plot_kwargs Keyword arguments for :meth:`cellrank.ul.models.BaseModel.plot`. kwargs Keyword arguments for :meth:`cellrank.ul.models.BaseModel.prepare`. Returns ------- %(plots_or_returns_models)s """ if isinstance(genes, str): genes = [genes] genes = _unique_order_preserving(genes) if data_key != "obs": _check_collection(adata, genes, "var_names", use_raw=kwargs.get("use_raw", False)) else: _check_collection(adata, genes, "obs", use_raw=kwargs.get("use_raw", False)) ln_key = str(AbsProbKey.BACKWARD if backward else AbsProbKey.FORWARD) if ln_key not in adata.obsm: raise KeyError(f"Lineages key `{ln_key!r}` not found in `adata.obsm`.") if lineages is None: lineages = adata.obsm[ln_key].names elif isinstance(lineages, str): lineages = [lineages] elif all(ln is None for ln in lineages): # no lineage, all the weights are 1 lineages = [None] cbar = False logg.debug("All lineages are `None`, setting the weights to `1`") lineages = _unique_order_preserving(lineages) if isinstance(time_range, (tuple, float, int, type(None))): time_range = [time_range] * len(lineages) elif len(time_range) != len(lineages): raise ValueError( f"Expected time ranges to be of length `{len(lineages)}`, found `{len(time_range)}`." ) kwargs["time_key"] = time_key kwargs["data_key"] = data_key kwargs["backward"] = backward kwargs["conf_int"] = conf_int # prepare doesnt take or need this models = _create_models(model, genes, lineages) all_models, models, genes, lineages = _fit_bulk( models, _create_callbacks(adata, callback, genes, lineages, **kwargs), genes, lineages, time_range, return_models=True, filter_all_failed=False, parallel_kwargs={ "show_progress_bar": show_progress_bar, "n_jobs": _get_n_cores(n_jobs, len(genes)), "backend": _get_backend(models, backend), }, **kwargs, ) lineages = sorted(lineages) tmp = adata.obsm[ln_key][lineages].colors if lineage_cmap is None and not transpose: lineage_cmap = tmp plot_kwargs = dict(plot_kwargs) plot_kwargs["obs_legend_loc"] = obs_legend_loc if transpose: all_models = pd.DataFrame(all_models).T.to_dict() models = pd.DataFrame(models).T.to_dict() genes, lineages = lineages, genes hide_cells = same_plot or hide_cells else: # information overload otherwise plot_kwargs["lineage_probability"] = False plot_kwargs["lineage_probability_conf_int"] = False tmp = pd.DataFrame(models).T.astype(bool) start_rows = np.argmax(tmp.values, axis=0) end_rows = tmp.shape[0] - np.argmax(tmp[::-1].values, axis=0) - 1 if same_plot: gene_as_title = True if gene_as_title is None else gene_as_title sharex = "all" if sharex is None else sharex if sharey is None: sharey = "row" if plot_kwargs.get("lineage_probability", False) else "none" ncols = len(genes) if ncols >= len(genes) else ncols nrows = int(np.ceil(len(genes) / ncols)) else: gene_as_title = False if gene_as_title is None else gene_as_title sharex = "col" if sharex is None else sharex if sharey is None: sharey = ("row" if not hide_cells or plot_kwargs.get( "lineage_probability", False) else "none") nrows = len(genes) ncols = len(lineages) plot_kwargs = dict(plot_kwargs) if plot_kwargs.get("xlabel", None) is None: plot_kwargs["xlabel"] = time_key fig, axes = plt.subplots( nrows=nrows, ncols=ncols, sharex=sharex, sharey=sharey, figsize=(6 * ncols, 4 * nrows) if figsize is None else figsize, tight_layout=True, dpi=dpi, ) axes = np.reshape(axes, (nrows, ncols)) cnt = 0 plot_kwargs["obs_legend_loc"] = None if same_plot else obs_legend_loc logg.info("Plotting trends") for row in range(len(axes)): for col in range(len(axes[row])): if cnt >= len(genes): break gene = genes[cnt] if (same_plot and plot_kwargs.get("lineage_probability", False) and transpose): lpc = adata.obsm[ln_key][gene].colors[0] else: lpc = None if same_plot: plot_kwargs["obs_legend_loc"] = (obs_legend_loc if row == 0 and col == len(axes[0]) - 1 else None) _trends_helper( models, gene=gene, lineage_names=lineages, transpose=transpose, same_plot=same_plot, hide_cells=hide_cells, perc=perc, lineage_cmap=lineage_cmap, abs_prob_cmap=abs_prob_cmap, lineage_probability_color=lpc, cell_color=cell_color, alpha=cell_alpha, lineage_alpha=lineage_alpha, size=size, lw=lw, cbar=cbar, margins=margins, sharey=sharey, gene_as_title=gene_as_title, legend_loc=legend_loc, figsize=figsize, fig=fig, axes=axes[row, col] if same_plot else axes[cnt], show_ylabel=col == 0, show_lineage=same_plot or (cnt == start_rows), show_xticks_and_label=((row + 1) * ncols + col >= len(genes)) if same_plot else (cnt == end_rows), **plot_kwargs, ) # plot legend on the 1st plot cnt += 1 if not same_plot: plot_kwargs["obs_legend_loc"] = None if same_plot and (col != ncols): for ax in np.ravel(axes)[cnt:]: ax.remove() fig.suptitle(suptitle, y=1.05) if save is not None: save_fig(fig, save) if return_models: return all_models
def graph( data: Union[AnnData, np.ndarray, spmatrix], graph_key: Optional[str] = None, ixs: Optional[np.array] = None, layout: Union[str, Dict, Callable] = "umap", keys: Sequence[KEYS] = ("incoming", ), keylocs: Union[KEYLOCS, Sequence[KEYLOCS]] = "uns", node_size: float = 400, labels: Optional[Union[Sequence[str], Sequence[Sequence[str]]]] = None, top_n_edges: Optional[Union[int, Tuple[int, bool, str]]] = None, self_loops: bool = True, self_loop_radius_frac: Optional[float] = None, filter_edges: Optional[Tuple[float, float]] = None, edge_reductions: Union[Callable, Sequence[Callable]] = np.sum, edge_weight_scale: float = 10, edge_width_limit: Optional[float] = None, edge_alpha: float = 1.0, edge_normalize: bool = False, edge_use_curved: bool = True, show_arrows: bool = True, font_size: int = 12, font_color: str = "black", color_nodes: bool = True, cat_cmap: ListedColormap = cm.Set3, cont_cmap: ListedColormap = cm.viridis, legend_loc: Optional[str] = "best", figsize: Optional[Tuple[float, float]] = None, dpi: Optional[int] = None, save: Optional[Union[str, Path]] = None, layout_kwargs: Dict = MappingProxyType({}), ) -> None: """ Plot a graph, visualizing incoming and outgoing edges or self-transitions. This is a utility function to look in more detail at the transition matrix in areas of interest, e.g. around an endpoint of development. This function is meant to visualise a small subset of nodes (~100-500) and the most likely transitions between them. Note that limiting edges visualized using ``top_n_edges`` will speed things up, as well as reduce the visual clutter. Parameters ---------- data The graph data to be plotted. graph_key Key in ``adata.obsp`` or ``adata.uns`` where the graph is stored. Only used when ``data`` is :class:`~anndata.Anndata` object. ixs Subset of indices of the graph to visualize. layout Layout to use for graph drawing. - If :class:`str`, search for embedding in ``adata.obsm['X_{layout}']``. Use ``layout_kwargs={'components': [0, 1]}`` to select components. - If :class:`dict`, keys should be values in interval ``[0, len(ixs))`` and values `(x, y)` pairs corresponding to node positions. keys Keys in ``adata.obs``, ``adata.obsm`` or ``adata.obsp`` to color the nodes. - If `'incoming'`, `'outgoing'` or `'self_loops'`, visualize reduction (see ``edge_reductions``) for each node based on incoming or outgoing edges, respectively. keylocs Locations of ``keys``. Can be any attribute of ``data`` if it's :class:`anndata.AnnData` object. node_size Size of the nodes. labels Labels of the nodes. top_n_edges Either top N outgoing edges in descending order or a tuple ``(top_n_edges, in_ascending_order, {'incoming', 'outgoing'})``. If `None`, show all edges. self_loops Whether visualize self transitions and also to consider them in ``top_n_edges``. self_loop_radius_frac Fraction of a unit circle to visualize self transitions. If `None`, use ``node_size / 1000``. filter_edges Whether to remove all edges not in `[min, max]` interval. edge_reductions Aggregation function to use when coloring nodes by edge weights. edge_weight_scale Number by which to scale the width of the edges. Useful when the weights are small. edge_width_limit Upper bound for the width of the edges. Useful when weights are unevenly distributed. edge_alpha Alpha channel value for edges and arrows. edge_normalize If `True`, normalize edges to `[0, 1]` interval prior to applying any scaling or truncation. edge_use_curved If `True`, use curved edges. This can improve visualization at a small performance cost. show_arrows Whether to show the arrows. Setting this to `False` may dramatically speed things up. font_size Font size for node labels. font_color Label color of the nodes. color_nodes Whether to color the nodes cat_cmap Categorical colormap used when ``keys`` contain categorical variables. cont_cmap Continuous colormap used when ``keys`` contain continuous variables. legend_loc Location of the legend. %(plotting)s layout_kwargs Additional kwargs for ``layout``. Returns ------- %(just_plots)s """ from anndata import AnnData as _AnnData import networkx as nx def plot_arrows(curves, G, pos, ax, edge_weight_scale): for line, (edge, val) in zip(curves, G.edges.items()): if edge[0] == edge[1]: continue mask = (~np.isnan(line)).all(axis=1) line = line[mask, :] if not len(line): # can be all NaNs continue line = line.reshape((-1, 2)) X, Y = line[:, 0], line[:, 1] node_start = pos[edge[0]] # reverse if np.where(np.isclose(node_start - line, [0, 0]).all(axis=1))[0][0]: X, Y = X[::-1], Y[::-1] mid = len(X) // 2 posA, posB = zip(X[mid:mid + 2], Y[mid:mid + 2]) # noqa arrow = FancyArrowPatch( posA=posA, posB=posB, # we clip because too small values # cause it to crash arrowstyle=ArrowStyle.CurveFilledB( head_length=np.clip( val["weight"] * edge_weight_scale * 4, _min_edge_weight, edge_width_limit, ), head_width=np.clip( val["weight"] * edge_weight_scale * 2, _min_edge_weight, edge_width_limit, ), ), color="k", zorder=float("inf"), alpha=edge_alpha, linewidth=0, ) ax.add_artist(arrow) def normalize_weights(): weights = np.array([v["weight"] for v in G.edges.values()]) minn = np.min(weights) weights = (weights - minn) / (np.max(weights) - minn) for v, w in zip(G.edges.values(), weights): v["weight"] = w def remove_top_n_edges(): if top_n_edges is None: return if isinstance(top_n_edges, (tuple, list)): to_keep, ascending, group_by = top_n_edges else: to_keep, ascending, group_by = top_n_edges, False, "out" if group_by not in ("incoming", "outgoing"): raise ValueError( "Argument `groupby` in `top_n_edges` must be either `'incoming`' or `'outgoing'`." ) source, target = zip(*G.edges) weights = [v["weight"] for v in G.edges.values()] tmp = pd.DataFrame({ "outgoing": source, "incoming": target, "w": weights }) if not self_loops: # remove self loops tmp = tmp[tmp["incoming"] != tmp["outgoing"]] to_keep = set( map( tuple, tmp.groupby(group_by).apply( lambda g: g.sort_values("w", ascending=ascending).take( range(min(to_keep, len(g)))))[["outgoing", "incoming"]].values, )) for e in list(G.edges): if e not in to_keep: G.remove_edge(*e) def remove_low_weight_edges(): if filter_edges is None or filter_edges == (None, None): return minn, maxx = filter_edges minn = minn if minn is not None else -np.inf maxx = maxx if maxx is not None else np.inf for e, attr in list(G.edges.items()): if attr["weight"] < minn or attr["weight"] > maxx: G.remove_edge(*e) _min_edge_weight = 0.00001 if edge_width_limit is None: logg.debug("Not limiting width of edges") edge_width_limit = float("inf") if self_loop_radius_frac is None: self_loop_radius_frac = (node_size / 2000 if node_size >= 200 else node_size / 1000) logg.debug( f"Setting self loop radius fraction to `{self_loop_radius_frac}`") if not isinstance(keylocs, (tuple, list)): keylocs = [keylocs] * len(keys) elif len(keylocs) == 1: keylocs = keylocs * 3 elif all(map(lambda k: k in ("incoming", "outgoing", "self_loops"), keys)): # don't care about keylocs since they are irrelevant logg.debug("Ignoring key locations") keylocs = [None] * len(keys) if not isinstance(edge_reductions, (tuple, list)): edge_reductions = [edge_reductions] * len(keys) if not all(map(callable, edge_reductions)): raise ValueError("Not all `edge_reductions` functions are callable.") if not isinstance(labels, (tuple, list)): labels = [labels] * len(keys) elif not len(labels): labels = [None] * len(keys) elif not isinstance(labels[0], (tuple, list)): labels = [labels] * len(keys) if len(keys) != len(labels): raise ValueError( f"`Keys` and `labels` must be of the same shape, found `{len(keys)}` and `{len(labels)}`." ) if isinstance(data, _AnnData): if graph_key is None: raise ValueError( "Argument `graph_key` cannot be `None` when `data` is `anndata.Anndata` object." ) gdata = _read_graph_data(data, graph_key) elif isinstance(data, (np.ndarray, spmatrix)): gdata = data else: raise TypeError( f"Expected argument `data` to be one of `anndata.AnnData`, `numpy.ndarray`, `scipy.sparse.spmatrix`, " f"found `{type(data).__name__!r}`.") is_sparse = issparse(gdata) if ixs is not None: gdata = gdata[ixs, :][:, ixs] else: ixs = list(range(gdata.shape[0])) start = logg.info("Creating graph") G = (nx.from_scipy_sparse_matrix(gdata, create_using=nx.DiGraph) if is_sparse else nx.from_numpy_array(gdata, create_using=nx.DiGraph)) remove_low_weight_edges() remove_top_n_edges() if edge_normalize: normalize_weights() logg.info(" Finish", time=start) # do NOT recreate the graph, for the edge reductions # gdata = nx.to_numpy_array(G) if figsize is None: figsize = (12, 8 * len(keys)) fig, axes = plt.subplots(nrows=len(keys), ncols=1, figsize=figsize, dpi=dpi) if not isinstance(axes, np.ndarray): axes = np.array([axes]) axes = np.ravel(axes) if isinstance(layout, str): if f"X_{layout}" not in data.obsm: raise KeyError( f"Unable to find embedding `'X_{layout}'` in `adata.obsm`.") components = layout_kwargs.get("components", [0, 1]) if len(components) != 2: raise ValueError( f"Components in `layout_kwargs` must be of length `2`, found `{len(components)}`." ) emb = data.obsm[f"X_{layout}"][:, components] pos = {i: emb[ix, :] for i, ix in enumerate(ixs)} logg.info(f"Embedding graph using `{layout!r}` layout") elif isinstance(layout, dict): rng = range(len(ixs)) for k, v in layout.items(): if k not in rng: raise ValueError( f"Key in `layout` must be in `range(len(ixs))`, found `{k}`." ) if len(v) != 2: raise ValueError( f"Value in `layout` must be a `tuple` or a `list` of length 2, found `{len(v)}`." ) pos = layout logg.debug("Using precomputed layout") elif callable(layout): start = logg.info( f"Embedding graph using `{layout.__name__!r}` layout") pos = layout(G, **layout_kwargs) logg.info(" Finish", time=start) else: raise TypeError(f"Argument `layout` must be either a `string`, " f"a `dict` or a `callable`, found `{type(layout)}`.") curves, lc = None, None if edge_use_curved: try: from ._utils import _curved_edges logg.debug("Creating curved edges") curves = _curved_edges(G, pos, self_loop_radius_frac, polarity="directed") lc = LineCollection( curves, colors="black", linewidths=np.clip( np.ravel([v["weight"] for v in G.edges.values()]) * edge_weight_scale, 0, edge_width_limit, ), alpha=edge_alpha, ) except ImportError as e: global _msg_shown if not _msg_shown: print( str(e)[:-1], "in order to use curved edges or specify `edge_use_curved=False`.", ) _msg_shown = True for ax, keyloc, key, labs, er in zip(axes, keylocs, keys, labels, edge_reductions): label_col = {} # dummy value if key in ("incoming", "outgoing", "self_loops"): if key in ("incoming", "outgoing"): vals = er(gdata, axis=int(key == "outgoing")) if issparse(vals): vals = vals.A vals = vals.flatten() else: vals = gdata.diagonal() if is_sparse else np.diag(gdata) node_v = dict(zip(pos.keys(), vals)) else: label_col = getattr(data, keyloc) if key in label_col: node_v = dict(zip(pos.keys(), label_col[key])) else: raise RuntimeError( f"Key `{key!r}` not found in `adata.{keyloc}`.") if labs is not None: if len(labs) != len(pos): raise RuntimeError( f"Number of labels ({len(labels)}) and nodes ({len(pos)}) mismatch." ) nx.draw_networkx_labels( G, pos, labels=labs if isinstance(labs, dict) else dict( zip(pos.keys(), labs)), ax=ax, font_color=font_color, font_size=font_size, ) if lc is not None and curves is not None: ax.add_collection(deepcopy(lc)) # copying necessary if show_arrows: plot_arrows(curves, G, pos, ax, edge_weight_scale) else: nx.draw_networkx_edges( G, pos, width=[ np.clip( v["weight"] * edge_weight_scale, _min_edge_weight, edge_width_limit, ) for _, v in G.edges.items() ], alpha=edge_alpha, edge_color="black", arrows=True, arrowstyle="-|>", ) if key in label_col and is_categorical_dtype(label_col[key]): values = label_col[key] if keyloc in ("obs", "obsm"): values = values[ixs] categories = values.cat.categories color_key = _colors(key) if color_key in data.uns: mapper = dict(zip(categories, data.uns[color_key])) else: mapper = dict( zip(categories, map(cat_cmap.get, range(len(categories))))) colors = [] seen = set() for v in values: colors.append(mapper[v]) seen.add(v) nodes_kwargs = dict(cmap=cat_cmap, node_color=colors) # noqa if legend_loc is not None: x, y = pos[0] for label in sorted(seen): ax.plot([x], [y], label=label, color=mapper[label]) ax.legend(loc=legend_loc) else: values = list(node_v.values()) vmin, vmax = np.min(values), np.max(values) nodes_kwargs = dict( # noqa cmap=cont_cmap, node_color=values, vmin=vmin, vmax=vmax) divider = make_axes_locatable(ax) cax = divider.append_axes("right", size="1.5%", pad=0.05) _ = mpl.colorbar.ColorbarBase(cax, cmap=cont_cmap, norm=mpl.colors.Normalize(vmin=vmin, vmax=vmax)) if color_nodes is False: nodes_kwargs = {} nx.draw_networkx_nodes(G, pos, node_size=node_size, ax=ax, **nodes_kwargs) ax.set_title(key) ax.axis("off") if save is not None: save_fig(fig, save) fig.show()
def plot( self, figsize: Tuple[float, float] = (8, 5), same_plot: bool = False, hide_cells: bool = False, perc: Tuple[float, float] = None, abs_prob_cmap: mcolors.ListedColormap = cm.viridis, cell_color: str = "black", lineage_color: str = "black", alpha: float = 0.8, lineage_alpha: float = 0.2, title: Optional[str] = None, size: int = 15, lw: float = 2, cbar: bool = True, margins: float = 0.015, xlabel: str = "pseudotime", ylabel: str = "expression", conf_int: bool = True, lineage_probability: bool = False, lineage_probability_conf_int: Union[bool, float] = False, lineage_probability_color: Optional[str] = None, dpi: int = None, fig: mpl.figure.Figure = None, ax: mpl.axes.Axes = None, return_fig: bool = False, save: Optional[str] = None, **kwargs, ) -> Optional[mpl.figure.Figure]: """ Plot the smoothed gene expression. Parameters ---------- figsize Size of the figure. same_plot Whether to plot all trends in the same plot. hide_cells Whether to hide the cells. perc Percentile by which to clip the absorption probabilities. abs_prob_cmap Colormap to use when coloring in the absorption probabilities. cell_color Color for the cells when not coloring absorption probabilities. lineage_color Color for the lineage. alpha Alpha channel for cells. lineage_alpha Alpha channel for lineage confidence intervals. title Title of the plot. size Size of the points. lw Line width for the smoothed values. cbar Whether to show colorbar. margins Margins around the plot. xlabel Label on the x-axis. ylabel Label on the y-axis. conf_int Whether to show the confidence interval. lineage_probability Whether to show smoothed lineage probability as a dashed line. Note that this will require 1 additional model fit. lineage_probability_conf_int Whether to compute and show smoothed lineage probability confidence interval. If :paramref:`self` is :class:`cellrank.ul.models.GAMR`, it can also specify the confidence level, the default is `0.95`. Only used when ``show_lineage_probability=True``. lineage_probability_color Color to use when plotting the smoothed ``lineage_probability``. If `None`, it's the same as ``lineage_color``. Only used when ``show_lineage_probability=True``. dpi Dots per inch. fig Figure to use, if `None`, create a new one. ax: :class:`matplotlib.axes.Axes` Ax to use, if `None`, create a new one. return_fig If `True`, return the figure object. save Filename where to save the plot. If `None`, just shows the plots. **kwargs Keyword arguments for :meth:`matplotlib.axes.Axes.legend`, e.g. to disable the legend, specify ``loc=None``. Only available when ``show_lineage_probability=True``. Returns ------- %(just_plots)s """ if self.y_test is None: raise RuntimeError("Run `.predict()` first.") if fig is None or ax is None: fig, ax = plt.subplots(figsize=figsize, constrained_layout=True) if dpi is not None: fig.set_dpi(dpi) conf_int = conf_int and self.conf_int is not None hide_cells = (hide_cells or self.x_all is None or self.w_all is None or self.y_all is None) lineage_probability_color = (lineage_color if lineage_probability_color is None else lineage_probability_color) scaler = kwargs.pop( "scaler", self._create_scaler( lineage_probability, show_conf_int=conf_int, ), ) if lineage_probability: if ylabel in ("expression", self._gene): ylabel = f"scaled {ylabel}" vmin, vmax = None, None if not hide_cells: vmin, vmax = _minmax(self.w_all, perc) _ = ax.scatter( self.x_all.squeeze(), scaler(self.y_all.squeeze()), c=cell_color if same_plot or np.allclose(self.w_all, 1.0) else self.w_all.squeeze(), s=size, cmap=abs_prob_cmap, vmin=vmin, vmax=vmax, alpha=alpha, ) if title is None: title = (f"{self._gene} @ {self._lineage}" if self._lineage is not None else f"{self._gene}") ax.plot(self.x_test, scaler(self.y_test), color=lineage_color, lw=lw, label=title) if title is not None: ax.set_title(title) if ylabel is not None: ax.set_ylabel(ylabel) if xlabel is not None: ax.set_xlabel(xlabel) ax.margins(margins) if conf_int: ax.fill_between( self.x_test.squeeze(), scaler(self.conf_int[:, 0]), scaler(self.conf_int[:, 1]), alpha=lineage_alpha, color=lineage_color, linestyle="--", ) if (lineage_probability and not isinstance(self, FittedModel) and not np.allclose(self.w, 1.0)): from cellrank.pl._utils import _is_any_gam_mgcv model = deepcopy(self) model._y = self._reshape_and_retype(self.w).copy() model = model.fit() if not lineage_probability_conf_int: y = model.predict() elif _is_any_gam_mgcv(model): y = model.predict( level=lineage_probability_conf_int if isinstance( lineage_probability_conf_int, float) else 0.95) else: y = model.predict() model.confidence_interval() ax.fill_between( model.x_test.squeeze(), model.conf_int[:, 0], model.conf_int[:, 1], alpha=lineage_alpha, color=lineage_probability_color, linestyle="--", ) handle = ax.plot( model.x_test, y, color=lineage_probability_color, lw=lw, linestyle="--", zorder=-1, label="probability", ) if kwargs.get("loc", "best") is not None: ax.legend(handles=handle, **kwargs) if (cbar and not hide_cells and not same_plot and not np.allclose(self.w_all, 1.0)): norm = mcolors.Normalize(vmin=vmin, vmax=vmax) divider = make_axes_locatable(ax) cax = divider.append_axes("right", size="2%", pad=0.1) _ = mpl.colorbar.ColorbarBase( cax, norm=norm, cmap=abs_prob_cmap, ticks=np.linspace(norm.vmin, norm.vmax, 5), ) if save is not None: save_fig(fig, save) if return_fig: return fig
def heatmap( adata: AnnData, model: _input_model_type, genes: Sequence[str], lineages: Optional[Union[str, Sequence[str]]] = None, backward: bool = False, mode: str = HeatmapMode.LINEAGES.s, time_key: str = "latent_time", time_range: Optional[Union[_time_range_type, List[_time_range_type]]] = None, callback: _callback_type = None, cluster_key: Optional[Union[str, Sequence[str]]] = None, show_absorption_probabilities: bool = False, cluster_genes: bool = False, keep_gene_order: bool = False, scale: bool = True, n_convolve: Optional[int] = 5, show_all_genes: bool = False, cbar: bool = True, lineage_height: float = 0.33, fontsize: Optional[float] = None, xlabel: Optional[str] = None, cmap: mcolors.ListedColormap = cm.viridis, dendrogram: bool = True, return_genes: bool = False, return_models: bool = False, n_jobs: Optional[int] = 1, backend: str = _DEFAULT_BACKEND, show_progress_bar: bool = True, figsize: Optional[Tuple[float, float]] = None, dpi: Optional[int] = None, save: Optional[Union[str, Path]] = None, **kwargs, ) -> Optional[Union[Dict[str, pd.DataFrame], Tuple[_return_model_type, Dict[ str, pd.DataFrame]]]]: """ Plot a heatmap of smoothed gene expression along specified lineages. Parameters ---------- %(adata)s %(model)s %(genes)s lineages Names of the lineages for which to plot. If `None`, plot all lineages. %(backward)s mode Valid options are: - `{m.LINEAGES.s!r}` - group by ``genes`` for each lineage in ``lineages``. - `{m.GENES.s!r}` - group by ``lineages`` for each gene in ``genes``. time_key Key in ``adata.obs`` where the pseudotime is stored. %(time_ranges)s %(model_callback)s cluster_key Key(s) in ``adata.obs`` containing categorical observations to be plotted on top of the heatmap. Only available when ``mode={m.LINEAGES.s!r}``. show_absorption_probabilities Whether to also plot absorption probabilities alongside the smoothed expression. Only available when ``mode={m.LINEAGES.s!r}``. cluster_genes Whether to cluster genes using :func:`seaborn.clustermap` when ``mode='lineages'``. keep_gene_order Whether to keep the gene order for later lineages after the first was sorted. Only available when ``cluster_genes=False`` and ``mode={m.LINEAGES.s!r}``. scale Whether to normalize the gene expression `0-1` range. n_convolve Size of the convolution window when smoothing absorption probabilities. show_all_genes Whether to show all genes on y-axis. cbar Whether to show the colorbar. lineage_height Height of a bar when ``mode={m.GENES.s!r}``. fontsize Size of the title's font. xlabel Label on the x-axis. If `None`, it is determined based on ``time_key``. cmap Colormap to use when visualizing the smoothed expression. dendrogram Whether to show dendrogram when ``cluster_genes=True``. return_genes Whether to return the sorted or clustered genes. Only available when ``mode={m.LINEAGES.s!r}``. %(return_models)s %(parallel)s %(plotting)s kwargs Keyword arguments for :meth:`cellrank.ul.models.BaseModel.prepare`. Returns ------- %(plots_or_returns_models)s :class:`pandas.DataFrame` If ``return_genes=True`` and ``mode={m.LINEAGES.s!r}``, returns :class:`pandas.DataFrame` containing the clustered or sorted genes. """ import seaborn as sns def find_indices(series: pd.Series, values) -> Tuple[Any]: def find_nearest(array: np.ndarray, value: float) -> int: ix = np.searchsorted(array, value, side="left") if ix > 0 and (ix == len(array) or fabs(value - array[ix - 1]) < fabs(value - array[ix])): return ix - 1 return ix series = series[np.argsort(series.values)] return tuple(series[[find_nearest(series.values, v) for v in values]].index) def subset_lineage(lname: str, rng: np.ndarray) -> np.ndarray: time_series = adata.obs[time_key] ixs = find_indices(time_series, rng) lin = adata[ixs, :].obsm[lineage_key][lname] lin = lin.X.copy().squeeze() if n_convolve is not None: lin = convolve(lin, np.ones(n_convolve) / n_convolve, mode="nearest") return lin def create_col_colors(lname: str, rng: np.ndarray) -> Tuple[np.ndarray, Cmap, Norm]: color = adata.obsm[lineage_key][lname].colors[0] lin = subset_lineage(lname, rng) h, _, v = mcolors.rgb_to_hsv(mcolors.to_rgb(color)) end_color = mcolors.hsv_to_rgb([h, 1, v]) lineage_cmap = mcolors.LinearSegmentedColormap.from_list( "lineage_cmap", ["#ffffff", end_color], N=len(rng)) norm = mcolors.Normalize(vmin=np.min(lin), vmax=np.max(lin)) scalar_map = cm.ScalarMappable(cmap=lineage_cmap, norm=norm) return ( np.array([mcolors.to_hex(c) for c in scalar_map.to_rgba(lin)]), lineage_cmap, norm, ) def create_col_categorical_color(cluster_key: str, rng: np.ndarray) -> np.ndarray: if not is_categorical_dtype(adata.obs[cluster_key]): raise TypeError( f"Expected `adata.obs[{cluster_key!r}]` to be categorical, " f"found `{adata.obs[cluster_key].dtype.name!r}`.") color_key = f"{cluster_key}_colors" if color_key not in adata.uns: logg.warning( f"Color key `{color_key!r}` not found in `adata.uns`. Creating new colors" ) colors = _create_categorical_colors( len(adata.obs[cluster_key].cat.categories)) adata.uns[color_key] = colors else: colors = adata.uns[color_key] time_series = adata.obs[time_key] ixs = find_indices(time_series, rng) mapper = dict(zip(adata.obs[cluster_key].cat.categories, colors)) return np.array([ mcolors.to_hex(mapper[v]) for v in adata[ixs, :].obs[cluster_key].values ]) def create_cbar( ax, x_delta: float, cmap: Cmap, norm: Norm, label: Optional[str] = None, ) -> Ax: cax = inset_axes( ax, width="1%", height="100%", loc="lower right", bbox_to_anchor=(x_delta, 0, 1, 1), bbox_transform=ax.transAxes, ) _ = mpl.colorbar.ColorbarBase( cax, cmap=cmap, norm=norm, label=label, ticks=np.linspace(norm.vmin, norm.vmax, 5), ) return cax @valuedispatch def _plot_heatmap(_mode: HeatmapMode) -> Fig: pass @_plot_heatmap.register(HeatmapMode.GENES) def _() -> Tuple[Fig, None]: def color_fill_rec(ax, xs, y1, y2, colors=None, cmap=cmap, **kwargs) -> None: colors = colors if cmap is None else cmap(colors) x = 0 for i, (color, x, y1, y2) in enumerate(zip(colors, xs, y1, y2)): dx = (xs[i + 1] - xs[i]) if i < len(x) else (xs[-1] - xs[-2]) ax.add_patch( plt.Rectangle((x, y1), dx, y2 - y1, color=color, ec=color, **kwargs)) ax.plot(x, y2, lw=0) fig, axes = plt.subplots( nrows=len(genes) + show_absorption_probabilities, figsize=(12, len(genes) + len(lineages) * lineage_height) if figsize is None else figsize, dpi=dpi, constrained_layout=True, ) if not isinstance(axes, Iterable): axes = [axes] axes = np.ravel(axes) if show_absorption_probabilities: data["absorption probability"] = data[next(iter(data.keys()))] for ax, (gene, models) in zip(axes, data.items()): if scale: vmin, vmax = 0, 1 else: c = np.array([m.y_test for m in models.values()]) vmin, vmax = np.nanmin(c), np.nanmax(c) norm = mcolors.Normalize(vmin=vmin, vmax=vmax) ix = 0 ys = [ix] if gene == "absorption probability": norm = mcolors.Normalize(vmin=0, vmax=1) for ln, x in ((ln, m.x_test) for ln, m in models.items()): y = np.ones_like(x) c = subset_lineage(ln, x.squeeze()) color_fill_rec(ax, x, y * ix, y * (ix + lineage_height), colors=norm(c)) ix += lineage_height ys.append(ix) else: for x, c in ((m.x_test, m.y_test) for m in models.values()): y = np.ones_like(x) c = _min_max_scale(c) if scale else c color_fill_rec(ax, x, y * ix, y * (ix + lineage_height), colors=norm(c)) ix += lineage_height ys.append(ix) xs = np.array([m.x_test for m in models.values()]) x_min, x_max = np.min(xs), np.max(xs) ax.set_xticks(np.linspace(x_min, x_max, _N_XTICKS)) ax.set_yticks(np.array(ys[:-1]) + lineage_height / 2) ax.spines["left"].set_position( ("data", 0) ) # move the left spine to the rectangles to get nicer yticks ax.set_yticklabels(models.keys(), ha="right") ax.set_title(gene, fontdict={"fontsize": fontsize}) ax.set_ylabel("lineage") for pos in ["top", "bottom", "left", "right"]: ax.spines[pos].set_visible(False) if cbar: cax, _ = mpl.colorbar.make_axes(ax) _ = mpl.colorbar.ColorbarBase( cax, ticks=np.linspace(vmin, vmax, 5), norm=norm, cmap=cmap, label="value" if gene == "absorption probability" else "scaled expression" if scale else "expression", ) ax.tick_params( top=False, bottom=False, left=True, right=False, labelleft=True, labelbottom=False, ) ax.xaxis.set_major_formatter(FormatStrFormatter("%.3f")) ax.tick_params( top=False, bottom=True, left=True, right=False, labelleft=True, labelbottom=True, ) ax.set_xlabel(xlabel) return fig, None @_plot_heatmap.register(HeatmapMode.LINEAGES) def _() -> Tuple[List[Fig], pd.DataFrame]: data_t = defaultdict(dict) # transpose for gene, lns in data.items(): for ln, y in lns.items(): data_t[ln][gene] = y figs = [] gene_order = None sorted_genes = pd.DataFrame() if return_genes else None for lname, models in data_t.items(): xs = np.array([m.x_test for m in models.values()]) x_min, x_max = np.nanmin(xs), np.nanmax(xs) df = pd.DataFrame([m.y_test for m in models.values()], index=models.keys()) df.index.name = "genes" if not cluster_genes: if gene_order is not None: df = df.loc[gene_order] else: max_sort = np.argsort( np.argmax(df.apply(_min_max_scale, axis=1).values, axis=1)) df = df.iloc[max_sort, :] if keep_gene_order: gene_order = df.index cat_colors = None if cluster_key is not None: cat_colors = np.stack( [ create_col_categorical_color( c, np.linspace(x_min, x_max, df.shape[1])) for c in cluster_key ], axis=0, ) if show_absorption_probabilities: col_colors, col_cmap, col_norm = create_col_colors( lname, np.linspace(x_min, x_max, df.shape[1])) if cat_colors is not None: col_colors = np.vstack([cat_colors, col_colors[None, :]]) else: col_colors, col_cmap, col_norm = cat_colors, None, None row_cluster = cluster_genes and df.shape[0] > 1 show_clust = row_cluster and dendrogram g = sns.clustermap( df, cmap=cmap, figsize=(10, min(len(genes) / 8 + 1, 10)) if figsize is None else figsize, xticklabels=False, row_cluster=row_cluster, col_colors=col_colors, colors_ratio=0, col_cluster=False, cbar_pos=None, yticklabels=show_all_genes or "auto", standard_scale=0 if scale else None, ) if cbar: cax = create_cbar( g.ax_heatmap, 0.1, cmap=cmap, norm=mcolors.Normalize( vmin=0 if scale else np.min(df.values), vmax=1 if scale else np.max(df.values), ), label="scaled expression" if scale else "expression", ) g.fig.add_axes(cax) if col_cmap is not None and col_norm is not None: cax = create_cbar( g.ax_heatmap, 0.25, cmap=col_cmap, norm=col_norm, label="absorption probability", ) g.fig.add_axes(cax) if g.ax_col_colors: main_bbox = _get_ax_bbox(g.fig, g.ax_heatmap) n_bars = show_absorption_probabilities + ( len(cluster_key) if cluster_key is not None else 0) _set_ax_height_to_cm( g.fig, g.ax_col_colors, height=min( 5, max(n_bars * main_bbox.height / len(df), 0.25 * n_bars)), ) g.ax_col_colors.set_title(lname, fontdict={"fontsize": fontsize}) else: g.ax_heatmap.set_title(lname, fontdict={"fontsize": fontsize}) g.ax_col_dendrogram.set_visible( False) # gets rid of top free space g.ax_heatmap.yaxis.tick_left() g.ax_heatmap.yaxis.set_label_position("right") g.ax_heatmap.set_xlabel(xlabel) g.ax_heatmap.set_xticks(np.linspace(0, len(df.columns), _N_XTICKS)) g.ax_heatmap.set_xticklabels( list( map(lambda n: round(n, 3), np.linspace(x_min, x_max, _N_XTICKS)))) if show_clust: # robustly show dendrogram, because gene names can be long g.ax_row_dendrogram.set_visible(True) dendro_box = g.ax_row_dendrogram.get_position() pad = 0.005 bb = g.ax_heatmap.yaxis.get_tightbbox( g.fig.canvas.get_renderer()).transformed( g.fig.transFigure.inverted()) dendro_box.x0 = bb.x0 - dendro_box.width - pad dendro_box.x1 = bb.x0 - pad g.ax_row_dendrogram.set_position(dendro_box) else: g.ax_row_dendrogram.set_visible(False) if return_genes: sorted_genes[lname] = (df.index[g.dendrogram_row.reordered_ind] if hasattr(g, "dendrogram_row") and g.dendrogram_row is not None else df.index) figs.append(g) return figs, sorted_genes mode = HeatmapMode(mode) lineage_key = str(AbsProbKey.BACKWARD if backward else AbsProbKey.FORWARD) if lineage_key not in adata.obsm: raise KeyError( f"Lineages key `{lineage_key!r}` not found in `adata.obsm`.") if lineages is None: lineages = adata.obsm[lineage_key].names elif isinstance(lineages, str): lineages = [lineages] lineages = _unique_order_preserving(lineages) _ = adata.obsm[lineage_key][lineages] if cluster_key is not None: if isinstance(cluster_key, str): cluster_key = [cluster_key] cluster_key = _unique_order_preserving(cluster_key) if isinstance(genes, str): genes = [genes] genes = _unique_order_preserving(genes) _check_collection(adata, genes, "var_names", use_raw=kwargs.get("use_raw", False)) kwargs["backward"] = backward kwargs["time_key"] = time_key models = _create_models(model, genes, lineages) all_models, data, genes, lineages = _fit_bulk( models, _create_callbacks(adata, callback, genes, lineages, **kwargs), genes, lineages, time_range, return_models=True, # always return (better error messages) filter_all_failed=True, parallel_kwargs={ "show_progress_bar": show_progress_bar, "n_jobs": _get_n_cores(n_jobs, len(genes)), "backend": _get_backend(models, backend), }, **kwargs, ) xlabel = time_key if xlabel is None else xlabel logg.debug(f"Plotting `{mode.s!r}` heatmap") fig, genes = _plot_heatmap(mode) if save is not None and fig is not None: if not isinstance(fig, Iterable): save_fig(fig, save) elif len(fig) == 1: save_fig(fig[0], save) else: for ln, f in zip(lineages, fig): save_fig(f, os.path.join(save, f"lineage_{ln}")) if return_genes and mode == HeatmapMode.LINEAGES: return (all_models, genes) if return_models else genes elif return_models: return all_models
def cluster_lineage( adata: AnnData, model: _model_type, genes: Sequence[str], lineage: str, backward: bool = False, time_range: _time_range_type = None, clusters: Optional[Sequence[str]] = None, n_points: int = 200, time_key: str = "latent_time", cluster_key: str = "clusters", norm: bool = True, recompute: bool = False, callback: _callback_type = None, ncols: int = 3, sharey: Union[str, bool] = False, key_added: Optional[str] = None, show_progress_bar: bool = True, n_jobs: Optional[int] = 1, backend: str = _DEFAULT_BACKEND, figsize: Optional[Tuple[float, float]] = None, dpi: Optional[int] = None, save: Optional[Union[str, Path]] = None, pca_kwargs: Dict = MappingProxyType({"svd_solver": "arpack"}), neighbors_kwargs: Dict = MappingProxyType({"use_rep": "X"}), louvain_kwargs: Dict = MappingProxyType({}), **kwargs, ) -> None: """ Cluster gene expression trends within a lineage and plot the clusters. This function is based on Palantir, see [Setty19]_. It can be used to discover modules of genes that drive development along a given lineage. Consider running this function on a subset of genes which are potential lineage drivers, identified e.g. by running :func:`cellrank.tl.lineage_drivers`. Parameters ---------- %(adata)s %(model)s %(genes)s lineage Name of the lineage for which to cluster the genes. %(backward)s %(time_ranges)s clusters Cluster identifiers to plot. If `None`, all clusters will be considered. Useful when plotting previously computed clusters. n_points Number of points used for prediction. time_key Key in ``adata.obs`` where the pseudotime is stored. cluster_key Key in ``adata.obs`` where the clustering is stored. norm Whether to z-normalize each trend to have zero mean, unit variance. recompute If `True`, recompute the clustering, otherwise try to find already existing one. %(model_callback)s ncols Number of columns for the plot. sharey Whether to share y-axis across multiple plots. key_added Postfix to add when saving the results to ``adata.uns``. %(parallel)s %(plotting)s pca_kwargs Keyword arguments for :func:`scanpy.pp.pca`. neighbors_kwargs Keyword arguments for :func:`scanpy.pp.neighbors`. louvain_kwargs Keyword arguments for :func:`scanpy.tl.louvain`. **kwargs: Keyword arguments for :meth:`cellrank.ul.models.BaseModel.prepare`. Returns ------- %(just_plots)s Updates ``adata.uns`` with the following key: - ``lineage_{lineage}_trend_{key_added}`` - an :class:`anndata.AnnData` object of shape ``(n_genes, n_points)`` containing the clustered genes. """ import scanpy as sc from anndata import AnnData as _AnnData lineage_key = str(AbsProbKey.BACKWARD if backward else AbsProbKey.FORWARD) if lineage_key not in adata.obsm: raise KeyError( f"Lineages key `{lineage_key!r}` not found in `adata.obsm`.") _ = adata.obsm[lineage_key][lineage] genes = _unique_order_preserving(genes) _check_collection(adata, genes, "var_names", kwargs.get("use_raw", False)) key_to_add = f"lineage_{lineage}_trend" if key_added is not None: logg.debug(f"Adding key `{key_added!r}`") key_to_add += f"_{key_added}" if recompute or key_to_add not in adata.uns: kwargs["time_key"] = time_key # kwargs for the model.prepare kwargs["n_test_points"] = n_points kwargs["backward"] = backward models = _create_models(model, genes, [lineage]) callbacks = _create_callbacks(adata, callback, genes, [lineage], **kwargs) backend = _get_backend(model, backend) n_jobs = _get_n_cores(n_jobs, len(genes)) start = logg.info(f"Computing gene trends using `{n_jobs}` core(s)") trends = parallelize( _cluster_lineages_helper, genes, as_array=True, unit="gene", n_jobs=n_jobs, backend=backend, extractor=np.vstack, show_progress_bar=show_progress_bar, )(models, callbacks, lineage, time_range, **kwargs) logg.info(" Finish", time=start) trends = trends.T if norm: logg.debug("Normalizing using `StandardScaler`") _ = StandardScaler(copy=False).fit_transform(trends) trends = _AnnData(trends.T) trends.obs_names = genes # sanity check if trends.n_obs != len(genes): raise RuntimeError( f"Expected to find `{len(genes)}` genes, found `{trends.n_obs}`." ) if n_points is not None and trends.n_vars != n_points: raise RuntimeError( f"Expected to find `{n_points}` points, found `{trends.n_vars}`." ) pca_kwargs = dict(pca_kwargs) n_comps = pca_kwargs.pop( "n_comps", min(50, kwargs.get("n_test_points"), len(genes)) - 1) # default value sc.pp.pca(trends, n_comps=n_comps, **pca_kwargs) sc.pp.neighbors(trends, **neighbors_kwargs) louvain_kwargs = dict(louvain_kwargs) louvain_kwargs["key_added"] = cluster_key sc.tl.louvain(trends, **louvain_kwargs) adata.uns[key_to_add] = trends else: logg.info(f"Loading data from `adata.uns[{key_to_add!r}]`") trends = adata.uns[key_to_add] if clusters is None: if cluster_key not in trends.obs: raise KeyError(f"Invalid cluster key `{cluster_key!r}`.") clusters = trends.obs[cluster_key].cat.categories nrows = int(np.ceil(len(clusters) / ncols)) fig, axes = plt.subplots( nrows, ncols, figsize=(ncols * 10, nrows * 10) if figsize is None else figsize, sharey=sharey, dpi=dpi, ) if not isinstance(axes, Iterable): axes = [axes] axes = np.ravel(axes) j = 0 for j, (ax, c) in enumerate(zip(axes, clusters)): # noqa data = trends[trends.obs[cluster_key] == c].X mean, sd = np.mean(data, axis=0), np.var(data, axis=0) sd = np.sqrt(sd) for i in range(data.shape[0]): ax.plot(data[i], color="gray", lw=0.5) ax.plot(mean, lw=2, color="black") ax.plot(mean - sd, lw=1.5, color="black", linestyle="--") ax.plot(mean + sd, lw=1.5, color="black", linestyle="--") ax.fill_between(range(len(mean)), mean - sd, mean + sd, color="black", alpha=0.1) ax.set_title(f"Cluster {c}") ax.set_xticks([]) if not sharey: ax.set_yticks([]) for j in range(j + 1, len(axes)): axes[j].remove() if save is not None: save_fig(fig, save)
def plot_single_flow( self, cluster: str, cluster_key: str, time_key: str, clusters: Optional[Sequence[Any]] = None, time_points: Optional[Sequence[Union[int, float]]] = None, min_flow: float = 0, remove_empty_clusters: bool = True, ascending: Optional[bool] = False, legend_loc: Optional[str] = "upper right out", alpha: Optional[float] = 0.8, xticks_step_size: Optional[int] = 1, figsize: Optional[Tuple[float, float]] = None, dpi: Optional[int] = None, save: Optional[Union[str, Path]] = None, show: bool = True, ) -> Optional[plt.Axes]: """ Visualize outgoing flow from a cluster of cells :cite:`mittnenzweig:21`. Parameters ---------- cluster Cluster for which to visualize outgoing compute_flow. cluster_key Key in :attr:`adata` ``.obs`` where clustering is stored. time_key Key in :attr:`adata` ``.obs`` where experimental time is stored. clusters Visualize flow only for these clusters. If `None`, use all clusters. time_points Visualize flow only for these time points. If `None`, use all time points. %(flow.parameters)s %(plotting)s show If `False`, return :class:`matplotlib.pyplot.Axes`. Returns ------- :class:`matplotlib.pyplot.Axes` The axis object if ``show=False``. %(just_plots)s Notes ----- This function is a Python reimplementation of the following `original R function <https://github.com/tanaylab/embflow/blob/main/scripts/generate_paper_figures/plot_vein.r>`_ with some minor stylistic differences. This function will not recreate the results from :cite:`mittnenzweig:21`, because there the Metacell model :cite:`baran:19` was used to compute the flow, whereas here the transition matrix is used. """ # noqa: E501 if self._transition_matrix is None: raise RuntimeError( "Compute transition matrix first as `.compute_transition_matrix()`." ) fp = FlowPlotter(self.adata, self.transition_matrix, cluster_key, time_key) fp = fp.prepare(cluster, clusters, time_points) ax = fp.plot( min_flow=min_flow, remove_empty_clusters=remove_empty_clusters, ascending=ascending, alpha=alpha, xticks_step_size=xticks_step_size, legend_loc=legend_loc, figsize=figsize, dpi=dpi, ) if save is not None: save_fig(ax.figure, save) if not show: return ax
def plot_random_walks( self, n_sims: int, max_iter: Union[int, float] = 0.25, seed: Optional[int] = None, successive_hits: int = 0, start_ixs: Indices_t = None, stop_ixs: Indices_t = None, basis: str = "umap", cmap: Union[str, LinearSegmentedColormap] = "gnuplot", linewidth: float = 1.0, linealpha: float = 0.3, ixs_legend_loc: Optional[str] = None, n_jobs: Optional[int] = None, backend: str = "loky", show_progress_bar: bool = True, figsize: Optional[Tuple[float, float]] = None, dpi: Optional[int] = None, save: Optional[Union[str, Path]] = None, **kwargs: Any, ) -> None: """ Plot random walks in an embedding. This method simulates random walks on the Markov chain defined though the corresponding transition matrix. The method is intended to give qualitative rather than quantitative insights into the transition matrix. Random walks are simulated by iteratively choosing the next cell based on the current cell's transition probabilities. Parameters ---------- n_sims Number of random walks to simulate. %(rw_sim.parameters)s start_ixs Cells from which to sample the starting points. If `None`, use all cells. %(rw_ixs)s For example ``{'clusters': ['Ngn3 low EP', 'Ngn3 high EP']}`` means that starting points for random walks will be samples uniformly from the these clusters. stop_ixs Cells which when hit, the random walk is terminated. If `None`, terminate after ``max_iters``. %(rw_ixs)s For example ``{'clusters': ['Alpha', 'Beta']}`` and ``succesive_hits=3`` means that the random walk will stop prematurely after cells in the above specified clusters have been visited successively 3 times in a row. basis Basis in :attr:`anndata.AnnData.obsm` to use as an embedding. cmap Colormap for the random walk lines. linewidth Width of the random walk lines. linealpha Alpha value of the random walk lines. ixs_legend_loc Legend location for the start/top indices. %(parallel)s %(plotting)s kwargs Keyword arguments for :func:`scvelo.pl.scatter`. Returns ------- %(just_plots)s For each random walk, the first/last cell is marked by the start/end colors of ``cmap``. """ def create_ixs(ixs: Indices_t, *, kind: str) -> Optional[np.ndarray]: if ixs is None: return None if isinstance(ixs, dict): # fmt: off if len(ixs) != 1: raise ValueError( f"Expected to find only 1 cluster key, found `{len(ixs)}`." ) cluster_key = next(iter(ixs.keys())) if cluster_key not in self.adata.obs: raise KeyError( f"Unable to find `adata.obs[{cluster_key!r}]`.") if not is_categorical_dtype(self.adata.obs[cluster_key]): raise TypeError( f"Expected `adata.obs[{cluster_key!r}]` to be categorical, " f"found `{infer_dtype(self.adata.obs[cluster_key])}`.") ixs = np.where( np.isin(self.adata.obs[cluster_key], ixs[cluster_key]))[0] # fmt: on elif isinstance(ixs, str): ixs = np.where(self.adata.obs_names == ixs)[0] else: ixs = np.where(np.isin(self.adata.obs_names, ixs))[0] if not len(ixs): logg.warning( f"No {kind} indices have been selected, using `None`") return None return ixs if self._transition_matrix is None: raise RuntimeError( "Compute transition matrix first as `.compute_transition_matrix()`." ) emb = _get_basis(self.adata, basis) if isinstance(cmap, str): cmap = plt.get_cmap(cmap) if not isinstance(cmap, LinearSegmentedColormap): if not hasattr(cmap, "colors"): raise AttributeError( "Unable to create a colormap, `cmap` does not have attribute `colors`." ) cmap = LinearSegmentedColormap.from_list("random_walk", colors=cmap.colors, N=max_iter) start_ixs = create_ixs(start_ixs, kind="start") stop_ixs = create_ixs(stop_ixs, kind="stop") rw = RandomWalk(self.transition_matrix, start_ixs=start_ixs, stop_ixs=stop_ixs) sims = rw.simulate_many( n_sims=n_sims, max_iter=max_iter, seed=seed, n_jobs=n_jobs, backend=backend, successive_hits=successive_hits, show_progress_bar=show_progress_bar, ) fig, ax = plt.subplots(figsize=figsize, dpi=dpi) scv.pl.scatter(self.adata, basis=basis, show=False, ax=ax, **kwargs) logg.info("Plotting random walks") for sim in sims: x = emb[sim][:, 0] y = emb[sim][:, 1] points = np.array([x, y]).T.reshape(-1, 1, 2) segments = np.concatenate([points[:-1], points[1:]], axis=1) n_seg = len(segments) lc = LineCollection( segments, linewidths=linewidth, colors=[cmap(float(i) / n_seg) for i in range(n_seg)], alpha=linealpha, zorder=2, ) ax.add_collection(lc) for ix in [0, -1]: ixs = [sim[ix] for sim in sims] plot_outline( x=emb[ixs][:, 0], y=emb[ixs][:, 1], outline_color=("black", to_hex(cmap(float(abs(ix))))), kwargs={ "s": kwargs.get("s", default_size(self.adata)) * 1.1, "alpha": 0.9, }, ax=ax, zorder=4, ) if ixs_legend_loc not in (None, "none"): from cellrank.pl._utils import _position_legend h1 = ax.scatter([], [], color=cmap(0.0), label="start") h2 = ax.scatter([], [], color=cmap(1.0), label="stop") legend = ax.get_legend() if legend is not None: ax.add_artist(legend) _position_legend(ax, legend_loc=ixs_legend_loc, handles=[h1, h2]) if save is not None: save_fig(fig, save)
def circular_projection( adata: AnnData, keys: Union[str, Sequence[str]], backward: bool = False, lineages: Optional[Union[str, Sequence[str]]] = None, early_cells: Optional[Union[Mapping[str, Sequence[str]], Sequence[str]]] = None, lineage_order: Optional[Literal["default", "optimal"]] = None, metric: Union[str, Callable, np.ndarray, pd.DataFrame] = "correlation", normalize_by_mean: bool = True, ncols: int = 4, space: float = 0.25, use_raw: bool = False, text_kwargs: Mapping[str, Any] = MappingProxyType({}), labeldistance: float = 1.25, labelrot: Union[Literal["default", "best"], float] = "best", show_edges: bool = True, key_added: Optional[str] = None, figsize: Optional[Tuple[float, float]] = None, dpi: Optional[int] = None, save: Optional[Union[str, Path]] = None, **kwargs, ): r""" Plot absorption probabilities on a circular embedding as done in [Velten17]_. Parameters ---------- %(adata)s keys Keys in :attr:`anndata.AnnData.obs` or :attr:`anndata.AnnData.var_names`. Additional keys are: - `'kl_divergence'` - as in [Velten17]_, computes KL-divergence between the fate probabilities of a cell and the average fate probabilities. See ``early_cells`` for more information. - `'entropy'` - as in [Setty19]_, computes entropy over a cells fate probabilities. %(backward)s lineages Lineages to plot. If `None`, plot all lineages. early_cells Cell ids or a mask marking early cells used to define the average fate probabilities. If `None`, use all cells. Only used when `'kl_divergence'` is in ``keys``. If a :class:`dict`, key specifies a cluster key in :attr:`anndata.AnnData.obs` and the values specify cluster labels containing early cells. lineage_order Can be one of the following: - `None` - it will determined automatically, based on the number of lineages. - `'optimal'` - order the lineages optimally by solving the Travelling salesman problem (TSP). Recommended for <= `20` lineages. - `'default'` - use the order as specified in ``lineages``. metric Metric to use when constructing pairwise distance matrix when ``lineage_order = 'optimal'``. For available options, see :func:`sklearn.metrics.pairwise_distances`. normalize_by_mean If `True`, normalize each lineage by its mean probability, as done in [Velten17]_. ncols Number of columns when plotting multiple ``keys``. space Horizontal and vertical space between for :func:`matplotlib.pyplot.subplots_adjust`. use_raw Whether to access :attr:`anndata.AnnData.raw` when there are ``keys`` in :attr:`anndata.AnnData.var_names`. text_kwargs Keyword arguments for :func:`matplotlib.pyplot.text`. labeldistance Distance at which the lineage labels will be drawn. labelrot How to rotate the labels. Valid options are: - `'best'` - rotate labels so that they are easily readable. - `'default'` - use :mod:`matplotlib`'s default. - `None` - same as `'default'`. If a :class:`float`, all labels will be rotated by this many degrees. show_edges Whether to show the edges surrounding the simplex. key_added Key in :attr:`anndata.AnnData.obsm` where to add the circular embedding. If `None`, it will be set to `'X_fate_simplex_{fwd,bwd}'`, based on ``backward``. %(plotting)s kwargs Keyword arguments for :func:`scvelo.pl.scatter`. Returns ------- %(just_plots)s Also updates ``adata`` with the following fields: - :attr:`anndata.AnnData.obsm` ``['{key_added}']``: the circular projection. - :attr:`anndata.AnnData.obs` ``['to_{initial,terminal}_states_{method}']``: the priming degree, if a method is present in ``keys``. """ if labeldistance is not None and labeldistance < 0: raise ValueError( f"Expected `delta` to be positive, found `{labeldistance}`.") if labelrot is None: labelrot = LabelRot.DEFAULT if isinstance(labelrot, str): labelrot = LabelRot(labelrot) suffix = "bwd" if backward else "fwd" if key_added is None: key_added = "X_fate_simplex_" + suffix if isinstance(keys, str): keys = (keys, ) keys = _unique_order_preserving(keys) keys_ = _check_collection( adata, keys, "obs", key_name="Observation", raise_exc=False) + _check_collection(adata, keys, "var_names", key_name="Gene", raise_exc=False, use_raw=use_raw) haystack = {s.s for s in PrimingDegree} keys = keys_ + [k for k in keys if k in haystack] keys = _unique_order_preserving(keys) if not len(keys): raise ValueError("No valid keys have been selected.") lineage_key = str(AbsProbKey.BACKWARD if backward else AbsProbKey.FORWARD) if lineage_key not in adata.obsm: raise KeyError( f"Lineages key `{lineage_key!r}` not found in `adata.obsm`.") probs = adata.obsm[lineage_key] if isinstance(lineages, str): lineages = (lineages, ) elif lineages is None: lineages = probs.names probs: Lineage = adata.obsm[lineage_key][lineages] n_lin = probs.shape[1] if n_lin <= 2: raise ValueError(f"Expected at least `3` lineages, found `{n_lin}`") X = probs.X.copy() if normalize_by_mean: X /= np.mean(X, axis=0)[None, :] X /= X.sum(1)[:, None] # this happens when cells for sel. lineages sum to 1 (or when the lineage average is 0, which is unlikely) X = np.nan_to_num(X, nan=1.0 / n_lin, copy=False) if lineage_order is None: lineage_order = LineageOrder.OPTIMAL if n_lin <= 15 else LineageOrder.DEFAULT logg.debug(f"Set ordering to `{lineage_order}`") lineage_order = LineageOrder(lineage_order) if lineage_order == LineageOrder.OPTIMAL: logg.info(f"Solving TSP for `{n_lin}` states") _, order = _get_optimal_order(X, metric=metric) else: order = np.arange(n_lin) probs = probs[:, order] X = X[:, order] angle_vec = np.linspace(0, 2 * np.pi, n_lin, endpoint=False) angle_vec_sin = np.cos(angle_vec) angle_vec_cos = np.sin(angle_vec) x = np.sum(X * angle_vec_sin, axis=1) y = np.sum(X * angle_vec_cos, axis=1) adata.obsm[key_added] = np.c_[x, y] nrows = int(np.ceil(len(keys) / ncols)) fig, ax = plt.subplots( nrows=nrows, ncols=ncols, figsize=(ncols * 5, nrows * 5) if figsize is None else figsize, dpi=dpi, ) fig.subplots_adjust(wspace=space, hspace=space) axes = np.ravel([ax]) text_kwargs = dict(text_kwargs) text_kwargs["ha"] = "center" text_kwargs["va"] = "center" _i = 0 for _i, (k, ax) in enumerate(zip(keys, axes)): set_lognorm, colorbar = False, kwargs.pop("colorbar", True) try: _ = PrimingDegree(k) logg.debug(f"Calculating priming degree using `method={k}`") val = probs.priming_degree(method=k, early_cells=early_cells) k = f"{lineage_key}_{k}" adata.obs[k] = val except ValueError: pass scv.pl.scatter( adata, basis=key_added, color=k, show=False, ax=ax, use_raw=use_raw, norm=LogNorm() if set_lognorm else None, colorbar=colorbar, **kwargs, ) if colorbar and set_lognorm: cbar = ax.collections[0].colorbar cax = cbar.locator.axis ticks = cax.minor.locator.tick_values(cbar.vmin, cbar.vmax) ticks = [ticks[0], ticks[len(ticks) // 2 + 1], ticks[-1]] cbar.set_ticks(ticks) cbar.set_ticklabels([f"{t:.2f}" for t in ticks]) cbar.update_ticks() patches, texts = ax.pie( np.ones_like(angle_vec), labeldistance=labeldistance, rotatelabels=True, labels=probs.names[::-1], startangle=-360 / len(angle_vec) / 2, counterclock=False, textprops=text_kwargs, ) for patch in patches: patch.set_visible(False) # clockwise for color, text in zip(probs.colors[::-1], texts): if isinstance(labelrot, (int, float)): text.set_rotation(labelrot) elif labelrot == LabelRot.BEST: rot = text.get_rotation() text.set_rotation(rot + 90 + (1 - rot // 180) * 180) elif labelrot != LabelRot.DEFAULT: raise NotImplementedError( f"Label rotation `{labelrot}` is not yet implemented.") text.set_color(color) if not show_edges: continue for i, color in enumerate(probs.colors): next = (i + 1) % n_lin x = 1.04 * np.linspace(angle_vec_sin[i], angle_vec_sin[next], _N) y = 1.04 * np.linspace(angle_vec_cos[i], angle_vec_cos[next], _N) points = np.array([x, y]).T.reshape(-1, 1, 2) segments = np.concatenate([points[:-1], points[1:]], axis=1) cmap = LinearSegmentedColormap.from_list( "abs_prob_cmap", [color, probs.colors[next]], N=_N) lc = LineCollection(segments, cmap=cmap, zorder=-1) lc.set_array(np.linspace(0, 1, _N)) lc.set_linewidth(2) ax.add_collection(lc) for j in range(_i + 1, len(axes)): axes[j].remove() if save is not None: save_fig(fig, save)
def plot_lineage_drivers( self, lineage: str, n_genes: int = 8, ncols: Optional[int] = None, use_raw: bool = False, title_fmt: str = "{gene} qval={qval:.4e}", figsize: Optional[Tuple[float, float]] = None, dpi: Optional[int] = None, save: Optional[Union[str, Path]] = None, **kwargs, ) -> None: """ Plot lineage drivers discovered by :meth:`compute_lineage_drivers`. Parameters ---------- lineage Lineage for which to plot the driver genes. n_genes Top most correlated genes to plot. ncols Number of columns. use_raw Whether to look in :paramref:`adata` ``.raw.var`` or :paramref:`adata` ``.var``. title_fmt Title format. Possible keywords include `{gene}`, `{qval}`, `{corr}` for gene name, q-value and correlation, respectively. %(plotting)s kwargs Keyword arguments for :func:`scvelo.pl.scatter`. Returns ------- %(just_plots)s """ def prepare_format( gene: str, *, pval: Optional[float], qval: Optional[float], corr: Optional[float], ) -> Dict[str, Any]: kwargs = {} if "{gene" in title_fmt: kwargs["gene"] = gene if "{pval" in title_fmt: kwargs["pval"] = float(pval) if pval is not None else np.nan if "{qval" in title_fmt: kwargs["qval"] = float(qval) if qval is not None else np.nan if "{corr" in title_fmt: kwargs["corr"] = float(corr) if corr is not None else np.nan return kwargs lin_drivers = self._get(P.LIN_DRIVERS) if lin_drivers is None: raise RuntimeError( f"Compute `.{P.LIN_DRIVERS}` first as `.compute_lineage_drivers()`." ) key = f"{lineage} corr" if key not in lin_drivers: raise KeyError( f"Lineage `{key!r}` not found in `{list(lin_drivers.columns)}`." ) if n_genes <= 0: raise ValueError(f"Expected `n_genes` to be positive, found `{n_genes}`.") kwargs.pop("save", None) genes = lin_drivers.sort_values(by=key, ascending=False).head(n_genes) ncols = 4 if ncols is None else ncols nrows = int(np.ceil(len(genes) / ncols)) fig, axes = plt.subplots( ncols=ncols, nrows=nrows, dpi=dpi, figsize=(ncols * 6, nrows * 4) if figsize is None else figsize, ) axes = np.ravel([axes]) _i = 0 for _i, (gene, ax) in enumerate(zip(genes.index, axes)): data = genes.loc[gene] scv.pl.scatter( self.adata, color=gene, ncols=ncols, use_raw=use_raw, ax=ax, show=False, title=title_fmt.format( **prepare_format( gene, pval=data.get(f"{lineage} pval", None), qval=data.get(f"{lineage} qval", None), corr=data.get(f"{lineage} corr", None), ) ), **kwargs, ) for j in range(_i + 1, len(axes)): axes[j].remove() if save is not None: save_fig(fig, save)
def plot( self, figsize: Tuple[float, float] = (15, 10), same_plot: bool = False, hide_cells: bool = False, perc: Tuple[float, float] = None, abs_prob_cmap: mcolors.ListedColormap = cm.viridis, cell_color: str = "black", lineage_color: str = "black", alpha: float = 0.8, lineage_alpha: float = 0.2, title: Optional[str] = None, size: int = 15, lw: float = 2, show_cbar: bool = True, margins: float = 0.015, xlabel: str = "pseudotime", ylabel: str = "expression", show_conf_int: bool = True, dpi: int = None, fig: mpl.figure.Figure = None, ax: mpl.axes.Axes = None, return_fig: bool = False, save: Optional[str] = None, ) -> Optional[mpl.figure.Figure]: """ Plot the smoothed gene expression. Parameters ---------- figsize Size of the figure. same_plot Whether to plot all trends in the same plot. hide_cells Whether to hide the cells. perc Percentile by which to clip the absorption probabilities./ abs_prob_cmap Colormap to use when coloring in the absorption probabilities. cell_color Color for the cells when not coloring absorption probabilities. lineage_color Color for the lineage. alpha Alpha channel for cells. lineage_alpha Alpha channel for lineage confidence intervals. title Title of the plot. size Size of the points. lw Line width for the smoothed values. show_cbar Whether to show colorbar. margins Margins around the plot. xlabel Label on the x-axis. ylabel Label on the y-axis. show_conf_int Whether to show the confidence interval. dpi Dots per inch. fig Figure to use, if `None`, create a new one. ax: :class:`matplotlib.axes.Axes` Ax to use, if `None`, create a new one. return_fig If `True`, return the figure object. save Filename where to save the plot. If `None`, just shows the plots. Returns ------- %(just_plots)s """ if fig is None or ax is None: fig, ax = plt.subplots(figsize=figsize, constrained_layout=True) if dpi is not None: fig.set_dpi(dpi) vmin, vmax = _minmax(self.w, perc) if not hide_cells: _ = ax.scatter( self.x_all.squeeze(), self.y_all.squeeze(), c=cell_color if same_plot or np.allclose(self.w_all, 1.0) else self.w_all.squeeze(), s=size, cmap=abs_prob_cmap, vmin=vmin, vmax=vmax, alpha=alpha, ) if title is None: title = f"{self._gene} @ {self._lineage}" _ = ax.plot(self.x_test, self.y_test, color=lineage_color, lw=lw, label=title) ax.set_title(title) ax.set_ylabel(ylabel) ax.set_xlabel(xlabel) ax.margins(margins) if show_conf_int and self.conf_int is not None: ax.fill_between( self.x_test.squeeze(), self.conf_int[:, 0], self.conf_int[:, 1], alpha=lineage_alpha, color=lineage_color, linestyle="--", ) if (show_cbar and not hide_cells and not same_plot and not np.allclose(self.w_all, 1)): norm = mcolors.Normalize(vmin=vmin, vmax=vmax) divider = make_axes_locatable(ax) cax = divider.append_axes("right", size="2.5%", pad=0.1) _ = mpl.colorbar.ColorbarBase(cax, norm=norm, cmap=abs_prob_cmap, label="absorption probability") if save is not None: save_fig(fig, save) if return_fig: return fig
def plot_schur_matrix( self, title: Optional[str] = "schur matrix", cmap: str = "viridis", figsize: Optional[Tuple[float, float]] = None, dpi: Optional[float] = 80, save: Optional[Union[str, Path]] = None, **kwargs, ): """ Plot the Schur matrix. Parameters ---------- title Title of the figure. cmap Colormap to use. %(plotting)s **kwargs Keyword arguments for :func:`seaborn.heatmap`. Returns ------- %(just_plots)s """ from seaborn import heatmap schur_matrix = getattr(self, P.SCHUR_MAT.s) if schur_matrix is None: raise RuntimeError( f"Compute Schur matrix first as `.{F.COMPUTE.fmt(P.SCHUR)}()`." ) fig, ax = plt.subplots( figsize=schur_matrix.shape if figsize is None else figsize, dpi=dpi) divider = make_axes_locatable( ax) # square=True make the colorbar a bit bigger cbar_ax = divider.append_axes("right", size="2%", pad=0.1) mask = np.zeros_like(schur_matrix, dtype=np.bool) mask[np.tril_indices_from(mask, k=-1)] = True mask[~np.isclose(schur_matrix, 0.0)] = False vmin, vmax = ( np.min(schur_matrix[~mask]), np.max(schur_matrix[~mask]), ) kwargs["fmt"] = kwargs.get("fmt", "0.2f") heatmap( schur_matrix, cmap=cmap, square=True, annot=True, vmin=vmin, vmax=vmax, cbar_ax=cbar_ax, cbar_kws={"ticks": np.linspace(vmin, vmax, 10)}, mask=mask, xticklabels=[], yticklabels=[], ax=ax, **kwargs, ) ax.set_title(title) if save is not None: save_fig(fig, path=save)
def cluster_fates( adata: AnnData, mode: str = ClusterFatesMode.PAGA_PIE.s, backward: bool = False, lineages: Optional[Union[str, Sequence[str]]] = None, cluster_key: Optional[str] = "clusters", clusters: Optional[Union[str, Sequence[str]]] = None, basis: Optional[str] = None, cbar: bool = True, ncols: Optional[int] = None, sharey: bool = False, fmt: str = "0.2f", xrot: float = 90, legend_kwargs: Mapping[str, Any] = MappingProxyType({"loc": "best"}), figsize: Optional[Tuple[float, float]] = None, dpi: Optional[int] = None, save: Optional[Union[str, Path]] = None, **kwargs, ) -> None: """ Plot aggregate lineage probabilities at a cluster level. This can be used to investigate how likely a certain cluster is to go to the %(terminal)s states,or in turn to have descended from the %(initial)s states. For mode `{m.PAGA.s!r}` and `{m.PAGA_PIE.s!r}`, we use *PAGA*, see [Wolf19]_. Parameters ---------- %(adata)s mode Type of plot to show. Valid options are: - `{m.BAR.s!r}` - barplot, one panel per cluster. - `{m.PAGA.s!r}` - scanpy's PAGA, one per %(initial_or_terminal)s state, colored in by fate. - `{m.PAGA_PIE.s!r}` - scanpy's PAGA with pie charts indicating aggregated fates. - `{m.VIOLIN.s!r}` - violin plots, one per %(initial_or_terminal)s state. - `{m.HEATMAP.s!r}` - a heatmap, showing average fates per cluster. - `{m.CLUSTERMAP.s!r}` - same as a heatmap, but with a dendrogram. %(backward)s lineages Lineages for which to visualize absorption probabilities. If `None`, use all lineages. cluster_key Key in ``adata.obs`` containing the clusters. clusters Clusters to visualize. If `None`, all clusters will be plotted. basis Basis for scatterplot to use when ``mode={m.PAGA_PIE.s!r}``. If `None`, don't show the scatterplot. cbar Whether to show colorbar when ``mode={m.PAGA_PIE.s!r}``. ncols Number of columns when ``mode={m.BAR.s!r}`` or ``mode={m.PAGA.s!r}``. sharey Whether to share y-axis when ``mode={m.BAR.s!r}``. fmt Format when using ``mode={m.HEATMAP.s!r}`` or ``mode={m.CLUSTERMAP.s!r}``. xrot Rotation of the labels on the x-axis. figsize Size of the figure. legend_kwargs Keyword arguments for :func:`matplotlib.axes.Axes.legend`, such as `'loc'` for legend position. For ``mode={m.PAGA_PIE.s!r}`` and ``basis='...'``, this controls the placement of the absorption probabilities legend. %(plotting)s **kwargs Keyword arguments for :func:`scvelo.pl.paga`, :func:`scanpy.pl.violin` or :func:`matplotlib.pyplot.bar`, depending on the value of ``mode``. Returns ------- %(just_plots)s """ from scanpy.plotting import violin from scvelo.plotting import paga from seaborn import heatmap, clustermap @valuedispatch def plot(mode: ClusterFatesMode, *_args, **_kwargs): raise NotImplementedError(mode.value) @plot.register(ClusterFatesMode.BAR) def _(): cols = 4 if ncols is None else ncols n_rows = ceil(len(clusters) / cols) fig = plt.figure(None, (3.5 * cols, 5 * n_rows) if figsize is None else figsize, dpi=dpi) fig.tight_layout() gs = plt.GridSpec(n_rows, cols, figure=fig, wspace=0.5, hspace=0.5) ax = None colors = list(adata.obsm[lk][:, lin_names].colors) for g, k in zip(gs, d.keys()): current_ax = fig.add_subplot(g, sharey=ax) current_ax.bar( x=np.arange(len(lin_names)), height=d[k][0], color=colors, yerr=d[k][1], ecolor="black", capsize=10, **kwargs, ) if sharey: ax = current_ax current_ax.set_xticks(np.arange(len(lin_names))) current_ax.set_xticklabels(lin_names, rotation=xrot) if not is_all: current_ax.set_xlabel(points) current_ax.set_ylabel("absorption probability") current_ax.set_title(k) return fig @plot.register(ClusterFatesMode.PAGA) def _(): kwargs["save"] = None kwargs["show"] = False if "cmap" not in kwargs: kwargs["cmap"] = cm.viridis cols = len(lin_names) if ncols is None else ncols nrows = ceil(len(lin_names) / cols) fig, axes = plt.subplots( nrows, cols, figsize=(7 * cols, 4 * nrows) if figsize is None else figsize, constrained_layout=True, dpi=dpi, ) # fig.tight_layout() can't use this because colorbar.make_axes fails i = 0 axes = [axes] if not isinstance(axes, np.ndarray) else np.ravel(axes) vmin, vmax = np.inf, -np.inf if basis is not None: kwargs["basis"] = basis kwargs["scatter_flag"] = True kwargs["color"] = cluster_key for i, (ax, lineage_name) in enumerate(zip(axes, lin_names)): colors = [v[0][i] for v in d.values()] kwargs["ax"] = ax kwargs["colors"] = tuple(colors) kwargs["title"] = f"{dir_prefix} {lineage_name}" vmin = np.min(colors + [vmin]) vmax = np.max(colors + [vmax]) paga(adata, **kwargs) if cbar: norm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax) cax, _ = mpl.colorbar.make_axes(ax, aspect=60) _ = mpl.colorbar.ColorbarBase( cax, ticks=np.linspace(norm.vmin, norm.vmax, 5), norm=norm, cmap=kwargs["cmap"], label="average absorption probability", ) for ax in axes[i + 1:]: # noqa ax.remove() return fig @plot.register(ClusterFatesMode.PAGA_PIE) def _(): colors = list(adata.obsm[lk][:, lin_names].colors) colors = { i: odict(zip(colors, mean)) for i, (mean, _) in enumerate(d.values()) } fig, ax = plt.subplots(figsize=figsize, dpi=dpi) fig.tight_layout() kwargs["ax"] = ax kwargs["show"] = False kwargs["colorbar"] = False # has to be disabled kwargs["show"] = False kwargs["node_colors"] = colors kwargs.pop("save", None) # we will handle saving kwargs["transitions"] = kwargs.get("transitions", "transitions_confidence") if "legend_loc" in kwargs: orig_ll = kwargs["legend_loc"] if orig_ll != "on data": kwargs["legend_loc"] = "none" # we will handle legend else: orig_ll = None kwargs["legend_loc"] = "on data" if basis is not None: kwargs["basis"] = basis kwargs["scatter_flag"] = True kwargs["color"] = cluster_key ax = paga(adata, **kwargs) ax.set_title(kwargs.get("title", cluster_key)) if basis is not None and orig_ll not in ("none", "on data", None): handles = [] for cluster_name, color in zip( adata.obs[f"{cluster_key}"].cat.categories, adata.uns[f"{cluster_key}_colors"], ): handles += [ax.scatter([], [], label=cluster_name, c=color)] first_legend = _position_legend( ax, legend_loc=orig_ll, handles=handles, **{k: v for k, v in legend_kwargs.items() if k != "loc"}, title=cluster_key, ) fig.add_artist(first_legend) if legend_kwargs.get("loc", None) not in ("none", "on data", None): # we need to use these, because scvelo can have its own handles and # they would be plotted here handles = [] for lineage_name, color in zip(lin_names, colors[0].keys()): handles += [ax.scatter([], [], label=lineage_name, c=color)] if len(colors[0].keys()) != len(adata.obsm[lk].names): handles += [ax.scatter([], [], label="Rest", c="grey")] second_legend = _position_legend( ax, legend_loc=legend_kwargs["loc"], handles=handles, **{k: v for k, v in legend_kwargs.items() if k != "loc"}, title=points, ) fig.add_artist(second_legend) return fig @plot.register(ClusterFatesMode.VIOLIN) def _(): kwargs.pop("ax", None) kwargs.pop("keys", None) kwargs.pop("save", None) # we will handle saving kwargs["show"] = False kwargs["groupby"] = cluster_key kwargs["rotation"] = xrot cols = len(lin_names) if ncols is None else ncols nrows = ceil(len(lin_names) / cols) fig, axes = plt.subplots( nrows, cols, figsize=(6 * cols, 4 * nrows) if figsize is None else figsize, sharey=sharey, dpi=dpi, ) fig.tight_layout() fig.subplots_adjust(wspace=0.2, hspace=0.3) if not isinstance(axes, np.ndarray): axes = [axes] axes = np.ravel(axes) with RandomKeys(adata, len(lin_names), where="obs") as keys: _i = 0 for _i, (name, key, ax) in enumerate(zip(lin_names, keys, axes)): adata.obs[key] = adata.obsm[lk][name].X ax.set_title(f"{dir_prefix} {name}") violin(adata, ylabel="absorption probability", keys=key, ax=ax, **kwargs) for ax in axes[_i + 1:]: # noqa ax.remove() return fig def plot_violin_no_cluster_key(): from anndata import AnnData as _AnnData kwargs.pop("ax", None) kwargs.pop("keys", None) # don't care kwargs.pop("save", None) kwargs["show"] = False kwargs["groupby"] = points kwargs["xlabel"] = None kwargs["rotation"] = xrot data = np.ravel(adata.obsm[lk].X.T)[..., np.newaxis] tmp = _AnnData(csr_matrix(data.shape, dtype=np.float32)) tmp.obs["absorption probability"] = data tmp.obs[points] = (pd.Series( np.concatenate([[f"{dir_prefix.lower()} {n}"] * adata.n_obs for n in adata.obsm[lk].names ])).astype("category").values) tmp.obs[points].cat.reorder_categories( [f"{dir_prefix.lower()} {n}" for n in adata.obsm[lk].names], inplace=True) tmp.uns[f"{points}_colors"] = adata.obsm[lk].colors fig, ax = plt.subplots(figsize=figsize if figsize is not None else (8, 6), dpi=dpi) ax.set_title(points.capitalize()) violin(tmp, keys=["absorption probability"], ax=ax, **kwargs) return fig @plot.register(ClusterFatesMode.HEATMAP) def _(): data = pd.DataFrame([mean for mean, _ in d.values()], columns=lin_names, index=clusters).T title = kwargs.pop("title", "average fate per cluster") vmin, vmax = data.values.min(), data.values.max() cbar_kws = { "label": "probability", "ticks": np.linspace(vmin, vmax, 5), "format": "%.3f", } kwargs.setdefault("cmap", "viridis") if use_clustermap: kwargs["cbar_pos"] = (0, 0.9, 0.025, 0.15) if cbar else None max_size = float(max(data.shape)) g = clustermap( data, annot=True, vmin=vmin, vmax=vmax, fmt=fmt, row_colors=adata.obsm[lk][lin_names].colors, dendrogram_ratio=( 0.15 * data.shape[0] / max_size, 0.15 * data.shape[1] / max_size, ), cbar_kws=cbar_kws, figsize=figsize, **kwargs, ) g.ax_heatmap.set_xlabel(cluster_key) g.ax_heatmap.set_ylabel("lineage") g.ax_col_dendrogram.set_title(title) fig = g.fig g = g.ax_heatmap else: fig, ax = plt.subplots(figsize=figsize, dpi=dpi) g = heatmap( data, vmin=vmin, vmax=vmax, annot=True, fmt=fmt, cbar=cbar, cbar_kws=cbar_kws, ax=ax, **kwargs, ) ax.set_title(title) ax.set_xlabel(cluster_key) ax.set_ylabel("lineage") g.set_xticklabels(g.get_xticklabels(), rotation=xrot) g.set_yticklabels(g.get_yticklabels(), rotation=0) return fig mode = ClusterFatesMode(mode) if cluster_key is not None: if cluster_key not in adata.obs: raise KeyError(f"Key `{cluster_key!r}` not found in `adata.obs`.") elif mode not in (mode.BAR, mode.VIOLIN): raise ValueError( f"Not specifying cluster key is only available for modes " f"`{ClusterFatesMode.BAR!r}` and `{ClusterFatesMode.VIOLIN!r}`, found `mode={mode!r}`." ) if backward: lk = AbsProbKey.BACKWARD.s points = TerminalStatesPlot.BACKWARD.s dir_prefix = DirPrefix.BACKWARD.s else: lk = AbsProbKey.FORWARD.s points = TerminalStatesPlot.FORWARD.s dir_prefix = DirPrefix.FORWARD.s if cluster_key is not None: is_all = False if clusters is not None: if isinstance(clusters, str): clusters = [clusters] clusters = _unique_order_preserving(clusters) if mode in (mode.PAGA, mode.PAGA_PIE): logg.debug( f"Setting `clusters` to all available ones because of `mode={mode!r}`" ) clusters = list(adata.obs[cluster_key].cat.categories) else: for cname in clusters: if cname not in adata.obs[cluster_key].cat.categories: raise KeyError( f"Cluster `{cname!r}` not found in `adata.obs[{cluster_key!r}]`." ) else: clusters = list(adata.obs[cluster_key].cat.categories) else: is_all = True clusters = [points] if lk not in adata.obsm: raise KeyError(f"Lineage key `{lk!r}` not found in `adata.obsm`.") if lineages is not None: if isinstance(lineages, str): lineages = [lineages] lin_names = _unique_order_preserving(lineages) else: # must be list for `sc.pl.violin`, else cats str lin_names = list(adata.obsm[lk].names) _ = adata.obsm[lk][lin_names] if mode == mode.VIOLIN and not is_all: adata = adata[np.isin(adata.obs[cluster_key], clusters)].copy() d = odict() for name in clusters: mask = (np.ones((adata.n_obs, ), dtype=np.bool) if is_all else (adata.obs[cluster_key] == name).values) mask = np.array(mask, dtype=np.bool) data = adata.obsm[lk][mask, lin_names].X mean = np.nanmean(data, axis=0) std = np.nanstd(data, axis=0) / np.sqrt(data.shape[0]) d[name] = [mean, std] logg.debug(f"Plotting in mode `{mode!r}`") use_clustermap = False if mode == mode.CLUSTERMAP: use_clustermap = True mode = mode.HEATMAP elif (mode in (ClusterFatesMode.PAGA, ClusterFatesMode.PAGA_PIE) and "paga" not in adata.uns): raise KeyError("Compute PAGA first as `scvelo.tl.paga()`.") fig = (plot_violin_no_cluster_key() if mode == ClusterFatesMode.VIOLIN and cluster_key is None else plot(mode)) if save is not None: save_fig(fig, save) fig.show()
def log_odds( adata: AnnData, lineage_1: str, lineage_2: Optional[str] = None, time_key: str = "exp_time", backward: bool = False, keys: Optional[Union[str, Sequence[str]]] = None, threshold: Optional[Union[float, Sequence]] = None, threshold_color: str = "red", layer: Optional[str] = None, use_raw: bool = False, size: float = 2.0, cmap: str = "viridis", alpha: Optional[float] = 0.8, ncols: Optional[int] = None, fontsize: Optional[Union[float, str]] = None, xticks_step_size: Optional[int] = 1, legend_loc: Optional[str] = "best", jitter: Union[bool, float] = True, seed: Optional[int] = None, figsize: Optional[Tuple[float, float]] = None, dpi: Optional[int] = None, save: Optional[Union[str, Path]] = None, show: bool = True, **kwargs: Any, ) -> Optional[Union[Axes, Sequence[Axes]]]: """ Plot log-odds ratio between lineages. Log-odds are plotted as a function of the experimental time. Parameters ---------- %(adata)s lineage_1 The first lineage for which to compute the log-odds. lineage_2 The second lineage for which to compute the log-odds. If `None`, use the rest of the lineages. time_key Key in :attr:`anndata.AnnData.obs` containing the experimental time. %(backward)s keys Key in :attr:`anndata.AnnData.obs` or :attr:`anndata.AnnData.var_names`. threshold Visualize whether total expression per cell is greater than ``threshold``. If a :class:`typing.Sequence`, it should be the same length as ``keys``. threshold_color Color to use when plotting thresholded expression values. layer Which layer to use to get expression values. If `None` or `'X'`, use :attr:`anndata.AnnData.X`. use_raw Whether to access :attr:`anndata.AnnData.raw`. If `True`, ``layer`` is ignored. size Size of the dots. cmap Colormap to use for continuous variables in ``keys``. alpha Alpha values for the dots. ncols Number of columns. fontsize Size of the font for the title, x- and y-label. xticks_step_size Show only every n-th ticks on x-axis. If `None`, don't show any ticks. legend_loc Position of the legend. If `None`, do not show the legend. jitter Amount of jitter to apply along x-axis. seed Seed for ``jitter`` to ensure reproducibility. %(plotting)s show If `False`, return :class:`matplotlib.pyplot.Axes` or a sequence of them. kwargs Keyword arguments for :func:`seaborn.stripplot`. Returns ------- :class:`matplotlib.pyplot.Axes` The axis object(s) if ``show=False``. %(just_plots)s """ from cellrank.tl.kernels._utils import _ensure_numeric_ordered def decorate(ax: Axes, *, title: Optional[str] = None, show_ylabel: bool = True) -> None: ax.set_xlabel(time_key, fontsize=fontsize) ax.set_title(title, fontdict={"fontsize": fontsize}) ax.set_ylabel(ylabel if show_ylabel else "", fontsize=fontsize) if xticks_step_size is None: ax.set_xticks([]) else: step = max(1, xticks_step_size) ax.set_xticks(np.arange(0, n_cats, step)) ax.set_xticklabels(df[time_key].cat.categories[::step]) def cont_palette(values: np.ndarray) -> Tuple[np.ndarray, ScalarMappable]: cm = copy(plt.get_cmap(cmap)) cm.set_bad("grey") sm = ScalarMappable(cmap=cm, norm=Normalize(vmin=np.nanmin(values), vmax=np.nanmax(values))) return np.array([to_hex(v) for v in (sm.to_rgba(values))]), sm def get_data( key: str, thresh: Optional[float] = None, ) -> Tuple[Optional[str], Optional[np.ndarray], Optional[np.ndarray], ScalarMappable]: try: _, palette = _get_categorical_colors(adata, key) df[key] = adata.obs[key].values[mask] df[key] = df[key].cat.remove_unused_categories() try: # seaborn doesn't like numeric categories df[key] = df[key].astype(float) palette = {float(k): v for k, v in palette.items()} except ValueError: pass # otherwise seaborn plots all palette = {k: palette[k] for k in df[key].unique()} hue, thresh_mask, sm = key, None, None except TypeError: palette, hue, thresh_mask, sm = ( cont_palette(adata.obs[key].values[mask])[0], None, None, None, ) except KeyError: try: # fmt: off if thresh is None: values = adata.raw.obs_vector( key) if use_raw else adata.obs_vector(key, layer=layer) palette, sm = cont_palette(values) hue, thresh_mask = None, None else: if use_raw: values = np.asarray( adata.raw[:, key].X[mask].sum(1)).squeeze() elif layer not in (None, "X"): values = np.asarray( adata[:, key].layers[layer][mask].sum(1)).squeeze() else: values = np.asarray( adata[:, key].X[mask].sum(1)).squeeze() thresh_mask = values > thresh hue, palette, sm = None, None, None # fmt: on except KeyError as e: raise e from None return hue, palette, thresh_mask, sm np.random.seed(seed) _ = kwargs.pop("orient", None) if use_raw and adata.raw is None: logg.warning("No raw attribute set. Setting `use_raw=False`") use_raw = False # define log-odds ln_key = str(AbsProbKey.BACKWARD if backward else AbsProbKey.FORWARD) if ln_key not in adata.obsm: raise KeyError(f"Lineages key `{ln_key!r}` not found in `adata.obsm`.") time = _ensure_numeric_ordered(adata, time_key) order = time.cat.categories[::-1 if backward else 1] fate1 = adata.obsm[ln_key][lineage_1].X.squeeze(-1) if lineage_2 is None: fate2 = 1 - fate1 ylabel = rf"$\log{{\frac{{{lineage_1}}}{{rest}}}}$" else: fate2 = adata.obsm[ln_key][lineage_2].X.squeeze(-1) ylabel = rf"$\log{{\frac{{{lineage_1}}}{{{lineage_2}}}}}$" # fmt: off df = pd.DataFrame({ "log_odds": np.log( np.divide(fate1, fate2, where=fate2 != 0, out=np.zeros_like(fate1)) + 1e-12), time_key: time, }) mask = (fate1 != 0) & (fate2 != 0) df = df[mask] n_cats = len(df[time_key].cat.categories) # fmt: on if keys is None: if figsize is None: figsize = np.array([n_cats, n_cats * 4 / 6]) / 2 fig, ax = plt.subplots(figsize=figsize, dpi=dpi, tight_layout=True) ax = sns.stripplot( time_key, "log_odds", data=df, order=order, jitter=jitter, color="k", size=size, ax=ax, **kwargs, ) decorate(ax) if save is not None: save_fig(fig, save) return None if show else ax if isinstance(keys, str): keys = (keys, ) if not len(keys): raise ValueError("No keys have been selected.") keys = _unique_order_preserving(keys) if not isinstance(threshold, Iterable): threshold = (threshold, ) * len(keys) if len(threshold) != len(keys): raise ValueError( f"Expected `threshold` to be of length `{len(keys)}`, found `{len(threshold)}`." ) ncols = max(len(keys) if ncols is None else ncols, 1) nrows = int(np.ceil(len(keys) / ncols)) if figsize is None: figsize = np.array([n_cats * ncols, n_cats * nrows * 4 / 6]) / 2 fig, axes = plt.subplots( nrows=nrows, ncols=ncols, figsize=figsize, dpi=dpi, constrained_layout=True, sharey="all", ) axes = np.ravel([axes]) i = 0 for i, (key, ax, thresh) in enumerate(zip(keys, axes, threshold)): hue, palette, thresh_mask, sm = get_data(key, thresh) show_ylabel = i % ncols == 0 ax = sns.stripplot( time_key, "log_odds", data=df if thresh_mask is None else df[~thresh_mask], hue=hue, order=order, jitter=jitter, color="black", palette=palette, size=size, alpha=alpha if alpha is not None else None if thresh_mask is None else 0.8, ax=ax, **kwargs, ) if thresh_mask is not None: sns.stripplot( time_key, "log_odds", data=df if thresh_mask is None else df[thresh_mask], hue=hue, order=order, jitter=jitter, color=threshold_color, palette=palette, size=size * 2, alpha=0.9, ax=ax, **kwargs, ) key = rf"${key} > {thresh}$" if sm is not None: cax = ax.inset_axes([1.02, 0, 0.025, 1], transform=ax.transAxes) fig.colorbar(sm, ax=ax, cax=cax) else: if legend_loc in (None, "none"): legend = ax.get_legend() if legend is not None: legend.remove() else: handles, labels = ax.get_legend_handles_labels() if len(handles): _position_legend(ax, legend_loc=legend_loc, handles=handles, labels=labels) decorate(ax, title=key, show_ylabel=show_ylabel) for ax in axes[i + 1:]: ax.remove() axes = axes[:i + 1] if save is not None: save_fig(fig, save) return None if show else axes[0] if len(axes) == 1 else axes
def plot_pie( self, reduction: Callable, title: Optional[str] = None, legend_loc: Optional[str] = "on data", legend_kwargs: Mapping = MappingProxyType({}), figsize: Optional[Tuple[float, float]] = None, dpi: Optional[float] = None, save: Optional[Union[Path, str]] = None, **kwargs, ) -> None: """ Plot a pie chart visualizing aggregated lineage probabilities. Parameters ---------- reduction Function that will be applied lineage-wise. title Title of the figure. legend_loc Location of the legend. If `None`, it is not shown. legend_kwargs Keyword arguments for :func:`matplotlib.axes.Axes.legend`. %(plotting)s Returns ------- %(just_plots)s """ if len(self.names) == 1: raise ValueError("Cannot plot pie chart for only 1 lineage.") fig, ax = plt.subplots(figsize=figsize, dpi=dpi) if "autopct" not in kwargs: autopct_found = False autopct = ( "{:.1f}%".format ) # we don't really care, we don't shot the pct, but the value else: autopct_found = True autopct = kwargs.pop("autopct") if title is None: title = reduction.__name__ if hasattr(reduction, "__name__") else None reduction = reduction(self, axis=int(self._is_transposed)).X.squeeze() reduction_norm = reduction / np.sum(reduction) wedges, texts, *autotexts = ax.pie( reduction_norm.squeeze(), labels=self.names if legend_loc == "on data" else None, autopct=autopct, wedgeprops={"edgecolor": "w"}, colors=self.colors, **kwargs, ) # if autopct is not None if len(autotexts): autotexts = autotexts[0] for name, at in zip(self.names, autotexts): ix = self._names_to_ixs[name] at.set_color(_get_bg_fg_colors(self.colors[ix])[1]) if not autopct_found: at.set_text(f"{reduction[ix]:.4f}") if legend_loc not in (None, "none", "on data"): ax.legend( wedges, self.names, title="lineages", loc=legend_loc, **legend_kwargs, ) ax.set_title(title) ax.set_aspect("equal") fig.show() if save is not None: save_fig(fig, save)
def cluster_lineage( adata: AnnData, model: _input_model_type, genes: Sequence[str], lineage: str, backward: bool = False, time_range: _time_range_type = None, clusters: Optional[Sequence[str]] = None, n_points: int = 200, time_key: str = "latent_time", norm: bool = True, recompute: bool = False, callback: _callback_type = None, ncols: int = 3, sharey: Union[str, bool] = False, key: Optional[str] = None, random_state: Optional[int] = None, use_leiden: bool = False, show_progress_bar: bool = True, n_jobs: Optional[int] = 1, backend: str = _DEFAULT_BACKEND, figsize: Optional[Tuple[float, float]] = None, dpi: Optional[int] = None, save: Optional[Union[str, Path]] = None, pca_kwargs: Dict = MappingProxyType({"svd_solver": "arpack"}), neighbors_kwargs: Dict = MappingProxyType({"use_rep": "X"}), clustering_kwargs: Dict = MappingProxyType({}), return_models: bool = False, **kwargs, ) -> Optional[_return_model_type]: """ Cluster gene expression trends within a lineage and plot the clusters. This function is based on Palantir, see [Setty19]_. It can be used to discover modules of genes that drive development along a given lineage. Consider running this function on a subset of genes which are potential lineage drivers, identified e.g. by running :func:`cellrank.tl.lineage_drivers`. Parameters ---------- %(adata)s %(model)s %(genes)s lineage Name of the lineage for which to cluster the genes. %(backward)s %(time_ranges)s clusters Cluster identifiers to plot. If `None`, all clusters will be considered. Useful when plotting previously computed clusters. n_points Number of points used for prediction. time_key Key in ``adata.obs`` where the pseudotime is stored. norm Whether to z-normalize each trend to have zero mean, unit variance. recompute If `True`, recompute the clustering, otherwise try to find already existing one. %(model_callback)s ncols Number of columns for the plot. sharey Whether to share y-axis across multiple plots. key Key in ``adata.uns`` where to save the results. If `None`, it will be saved as ``lineage_{lineage}_trend`` . random_state Random seed for reproducibility. use_leiden Whether to use :func:`scanpy.tl.leiden` for clustering or :func:`scanpy.tl.louvain`. %(parallel)s %(plotting)s pca_kwargs Keyword arguments for :func:`scanpy.pp.pca`. neighbors_kwargs Keyword arguments for :func:`scanpy.pp.neighbors`. clustering_kwargs Keyword arguments for :func:`scanpy.tl.louvain` or :func:`scanpy.tl.leiden`. %(return_models)s **kwargs: Keyword arguments for :meth:`cellrank.ul.models.BaseModel.prepare`. Returns ------- %(plots_or_returns_models)s Also updates ``adata.uns`` with the following: - ``key`` or ``lineage_{lineage}_trend`` - an :class:`anndata.AnnData` object of shape `(n_genes, n_points)` containing the clustered genes. """ import scanpy as sc from anndata import AnnData as _AnnData lineage_key = str(AbsProbKey.BACKWARD if backward else AbsProbKey.FORWARD) if lineage_key not in adata.obsm: raise KeyError( f"Lineages key `{lineage_key!r}` not found in `adata.obsm`.") _ = adata.obsm[lineage_key][lineage] genes = _unique_order_preserving(genes) _check_collection(adata, genes, "var_names", kwargs.get("use_raw", False)) if key is None: key = f"lineage_{lineage}_trend" if recompute or key not in adata.uns: kwargs["backward"] = backward kwargs["time_key"] = time_key kwargs["n_test_points"] = n_points models = _create_models(model, genes, [lineage]) all_models, models, genes, _ = _fit_bulk( models, _create_callbacks(adata, callback, genes, [lineage], **kwargs), genes, lineage, time_range, return_models=True, # always return (better error messages) filter_all_failed=True, parallel_kwargs={ "show_progress_bar": show_progress_bar, "n_jobs": _get_n_cores(n_jobs, len(genes)), "backend": _get_backend(models, backend), }, **kwargs, ) # `n_genes, n_test_points` trends = np.vstack( [model[lineage].y_test for model in models.values()]).T if norm: logg.debug("Normalizing trends") _ = StandardScaler(copy=False).fit_transform(trends) trends = _AnnData(trends.T) trends.obs_names = genes # sanity check if trends.n_obs != len(genes): raise RuntimeError( f"Expected to find `{len(genes)}` genes, found `{trends.n_obs}`." ) if trends.n_vars != n_points: raise RuntimeError( f"Expected to find `{n_points}` points, found `{trends.n_vars}`." ) random_state = np.random.mtrand.RandomState(random_state).randint( 2**16) pca_kwargs = dict(pca_kwargs) pca_kwargs.setdefault("n_comps", min(50, n_points, len(genes)) - 1) pca_kwargs.setdefault("random_state", random_state) sc.pp.pca(trends, **pca_kwargs) neighbors_kwargs = dict(neighbors_kwargs) neighbors_kwargs.setdefault("random_state", random_state) sc.pp.neighbors(trends, **neighbors_kwargs) clustering_kwargs = dict(clustering_kwargs) clustering_kwargs["key_added"] = "clusters" clustering_kwargs.setdefault("random_state", random_state) try: if use_leiden: sc.tl.leiden(trends, **clustering_kwargs) else: sc.tl.louvain(trends, **clustering_kwargs) except ImportError as e: logg.warning(str(e)) if use_leiden: sc.tl.louvain(trends, **clustering_kwargs) else: sc.tl.leiden(trends, **clustering_kwargs) logg.info(f"Saving data to `adata.uns[{key!r}]`") adata.uns[key] = trends else: all_models = None logg.info(f"Loading data from `adata.uns[{key!r}]`") trends = adata.uns[key] if "clusters" not in trends.obs: raise KeyError( "Unable to find the clustering in `trends.obs['clusters']`.") if clusters is None: clusters = trends.obs["clusters"].cat.categories for c in clusters: if c not in trends.obs["clusters"].cat.categories: raise ValueError( f"Invalid cluster name `{c!r}`. " f"Valid options are `{list(trends.obs['clusters'].cat.categories)}`." ) nrows = int(np.ceil(len(clusters) / ncols)) fig, axes = plt.subplots( nrows, ncols, figsize=(ncols * 10, nrows * 10) if figsize is None else figsize, sharey=sharey, dpi=dpi, ) if not isinstance(axes, Iterable): axes = [axes] axes = np.ravel(axes) j = 0 for j, (ax, c) in enumerate(zip(axes, clusters)): # noqa data = trends[trends.obs["clusters"] == c].X mean, sd = np.mean(data, axis=0), np.var(data, axis=0) sd = np.sqrt(sd) for i in range(data.shape[0]): ax.plot(data[i], color="gray", lw=0.5) ax.plot(mean, lw=2, color="black") ax.plot(mean - sd, lw=1.5, color="black", linestyle="--") ax.plot(mean + sd, lw=1.5, color="black", linestyle="--") ax.fill_between(range(len(mean)), mean - sd, mean + sd, color="black", alpha=0.1) ax.set_title(f"Cluster {c}") ax.set_xticks([]) if not sharey: ax.set_yticks([]) for j in range(j + 1, len(axes)): axes[j].remove() if save is not None: save_fig(fig, save) if return_models: return all_models
def gene_trends( adata: AnnData, model: _model_type, genes: Union[str, Sequence[str]], lineages: Optional[Union[str, Sequence[str]]] = None, backward: bool = False, data_key: str = "X", time_key: str = "latent_time", time_range: Optional[Union[_time_range_type, List[_time_range_type]]] = None, callback: _callback_type = None, conf_int: bool = True, same_plot: bool = False, hide_cells: bool = False, perc: Optional[Union[Tuple[float, float], Sequence[Tuple[float, float]]]] = None, lineage_cmap: Optional[matplotlib.colors.ListedColormap] = None, abs_prob_cmap: matplotlib.colors.ListedColormap = cm.viridis, cell_color: str = "black", cell_alpha: float = 0.6, lineage_alpha: float = 0.2, size: float = 15, lw: float = 2, show_cbar: bool = True, margins: float = 0.015, sharex: Optional[Union[str, bool]] = None, sharey: Optional[Union[str, bool]] = None, gene_as_title: Optional[bool] = None, legend_loc: Optional[str] = "best", ncols: int = 2, suptitle: Optional[str] = None, n_jobs: Optional[int] = 1, backend: str = _DEFAULT_BACKEND, show_progres_bar: bool = True, figsize: Optional[Tuple[float, float]] = None, dpi: Optional[int] = None, save: Optional[Union[str, Path]] = None, plot_kwargs: Mapping = MappingProxyType({}), **kwargs, ) -> None: """ Plot gene expression trends along lineages. Each lineage is defined via it's lineage weights which we compute using :func:`cellrank.tl.lineages`. This function accepts any model based off :class:`cellrank.ul.models.BaseModel` to fit gene expression, where we take the lineage weights into account in the loss function. Parameters ---------- %(adata)s %(model)s %(genes)s lineages Names of the lineages to plot. If `None`, plot all lineages. %(backward)s data_key Key in ``adata.layers`` or `'X'` for ``adata.X`` where the data is stored. time_key Key in ``adata.obs`` where the pseudotime is stored. %(time_ranges)s %(model_callback)s conf_int Whether to compute and show confidence intervals. same_plot Whether to plot all lineages for each gene in the same plot. hide_cells If `True`, hide all cells. perc Percentile for colors. Valid values are in interval `[0, 100]`. This can improve visualization. Can be specified individually for each lineage. lineage_cmap Colormap to use when coloring in the lineages. If `None` and ``same_plot``, use the corresponding colors in ``adata.uns``, otherwise use `'black'`. abs_prob_cmap Colormap to use when visualizing the absorption probabilities for each lineage. Only used when ``same_plot=False``. cell_color Color of the cells when not visualizing absorption probabilities. Only used when ``same_plot=True``. cell_alpha Alpha channel for cells. lineage_alpha Alpha channel for lineage confidence intervals. size Size of the points. lw Line width of the smoothed values. show_cbar Whether to show colorbar. Always shown when percentiles for lineages differ. Only used when ``same_plot=False``. margins Margins around the plot. sharex Whether to share x-axis. Valid options are `'row'`, `'col'` or `'none'`. sharey Whether to share y-axis. Valid options are `'row'`, `'col'` or `'none'`. gene_as_title Whether to show gene names as titles instead on y-axis. legend_loc Location of the legend displaying lineages. Only used when `same_plot=True`. ncols Number of columns of the plot when pl multiple genes. Only used when ``same_plot=True``. suptitle Suptitle of the figure. %(parallel)s %(plotting)s plot_kwargs Keyword arguments for :meth:`cellrank.ul.models.BaseModel.plot`. **kwargs Keyword arguments for :meth:`cellrank.ul.models.BaseModel.prepare`. Returns ------- %(just_plots)s """ if isinstance(genes, str): genes = [genes] genes = _unique_order_preserving(genes) if data_key != "obs": _check_collection(adata, genes, "var_names", use_raw=kwargs.get("use_raw", False)) else: _check_collection(adata, genes, "obs", use_raw=kwargs.get("use_raw", False)) ln_key = str(AbsProbKey.BACKWARD if backward else AbsProbKey.FORWARD) if ln_key not in adata.obsm: raise KeyError(f"Lineages key `{ln_key!r}` not found in `adata.obsm`.") if lineages is None: lineages = adata.obsm[ln_key].names elif isinstance(lineages, str): lineages = [lineages] elif all(map(lambda ln: ln is None, lineages)): # no lineage, all the weights are 1 lineages = [None] show_cbar = False logg.debug("All lineages are `None`, setting the weights to `1`") lineages = _unique_order_preserving(lineages) if same_plot: gene_as_title = True if gene_as_title is None else gene_as_title sharex = "all" if sharex is None else sharex sharey = "none" if sharey is None else sharey ncols = len(genes) if ncols >= len(genes) else ncols nrows = int(np.ceil(len(genes) / ncols)) else: gene_as_title = False if gene_as_title is None else gene_as_title sharex = "col" if sharex is None else sharex sharey = ( "none" if hide_cells else "row") if sharey is None else sharey nrows = len(genes) ncols = len(lineages) fig, axes = plt.subplots( nrows=nrows, ncols=ncols, sharex=sharex, sharey=sharey, figsize=(6 * ncols, 4 * nrows) if figsize is None else figsize, constrained_layout=True, ) axes = np.reshape(axes, (-1, ncols)) _ = adata.obsm[ln_key][[lin for lin in lineages if lin is not None]] if isinstance(time_range, (tuple, float, int, type(None))): time_range = [time_range] * len(lineages) elif len(time_range) != len(lineages): raise ValueError( f"Expected time ranges to be of length `{len(lineages)}`, found `{len(time_range)}`." ) kwargs["time_key"] = time_key kwargs["data_key"] = data_key kwargs["backward"] = backward callbacks = _create_callbacks(adata, callback, genes, lineages, **kwargs) kwargs["conf_int"] = conf_int # prepare doesnt take or need this models = _create_models(model, genes, lineages) plot_kwargs = dict(plot_kwargs) if plot_kwargs.get("xlabel", None) is None: plot_kwargs["xlabel"] = time_key n_jobs = _get_n_cores(n_jobs, len(genes)) backend = _get_backend(model, backend) start = logg.info(f"Computing trends using `{n_jobs}` core(s)") models = parallelize( _fit_gene_trends, genes, unit="gene" if data_key != "obs" else "obs", backend=backend, n_jobs=n_jobs, extractor=lambda modelss: {k: v for m in modelss for k, v in m.items()}, show_progress_bar=show_progres_bar, )(models, callbacks, lineages, time_range, **kwargs) logg.info(" Finish", time=start) logg.info("Plotting trends") cnt = 0 for row in range(len(axes)): for col in range(len(axes[row])): if cnt >= len(genes): break gene = genes[cnt] _trends_helper( adata, models, gene=gene, lineage_names=lineages, ln_key=ln_key, same_plot=same_plot, hide_cells=hide_cells, perc=perc, lineage_cmap=lineage_cmap, abs_prob_cmap=abs_prob_cmap, cell_color=cell_color, alpha=cell_alpha, lineage_alpha=lineage_alpha, size=size, lw=lw, show_cbar=show_cbar, margins=margins, sharey=sharey, gene_as_title=gene_as_title, legend_loc=legend_loc, dpi=dpi, figsize=figsize, fig=fig, axes=axes[row, col] if same_plot else axes[cnt], show_ylabel=col == 0, show_lineage=cnt == 0 or same_plot, show_xticks_and_label=((row + 1) * ncols + col >= len(genes)) if same_plot else (cnt == len(axes) - 1), **plot_kwargs, ) cnt += 1 if same_plot and (col != ncols): for ax in np.ravel(axes)[cnt:]: ax.remove() fig.suptitle(suptitle) if save is not None: save_fig(fig, save)