diff --git a/.gitignore b/.gitignore index db79aeb..dd3d502 100644 --- a/.gitignore +++ b/.gitignore @@ -110,3 +110,4 @@ ENV/ # IDE settings .vscode/ +.idea/ diff --git a/docs/genetools.rst b/docs/genetools.rst index 42d02d4..a7feedf 100644 --- a/docs/genetools.rst +++ b/docs/genetools.rst @@ -36,6 +36,14 @@ genetools.stats module :undoc-members: :show-inheritance: +genetools.trajectory module +--------------------------- + +.. automodule:: genetools.trajectory + :members: + :undoc-members: + :show-inheritance: + Module contents --------------- diff --git a/genetools/__init__.py b/genetools/__init__.py index 62b817a..127cfd2 100644 --- a/genetools/__init__.py +++ b/genetools/__init__.py @@ -5,6 +5,6 @@ __version__ = "0.4.0" # We want genetools.[submodule].[function] to be accessible simply by importing genetools, without having to import genetools.[submodule] -from . import helpers, plots, scanpy_helpers, stats +from . import helpers, plots, scanpy_helpers, stats, trajectory -__all__ = ["helpers", "plots", "scanpy_helpers", "stats"] +__all__ = ["helpers", "plots", "scanpy_helpers", "stats", "trajectory"] diff --git a/genetools/helpers.py b/genetools/helpers.py index fe0675d..9c9c219 100644 --- a/genetools/helpers.py +++ b/genetools/helpers.py @@ -60,6 +60,28 @@ def merge_into_left(left, right, **kwargs): return df +def sample_cells_from_clusters(obs_df, n_cells, cluster_key, cluster_names): + """[summary] + + :param obs_df: [description] + :type obs_df: [type] + :param n_cells: [description] + :type n_cells: [type] + :param cluster_key: [description] + :type cluster_key: [type] + :param cluster_names: [description] + :type cluster_names: [type] + :return: [description] + :rtype: [type] + """ + return ( + obs_df[obs_df[cluster_key].isin(cluster_names)] + .index.to_series() + .sample(n_cells) + .values + ) + + def horizontal_concat(df_left, df_right): """Concatenate df_right horizontally to df_left, with no checks for whether the indexes match, but confirming final shape. diff --git a/genetools/plots.py b/genetools/plots.py index 0c091a6..bcb1a62 100644 --- a/genetools/plots.py +++ b/genetools/plots.py @@ -29,6 +29,11 @@ def umap_scatter( label_color="k", label_alpha=0.5, label_size=20, + highlight_cell_names=None, + highlight_marker_size=30, + highlight_marker_color="r", + highlight_zorder=5, + highlight_marker_style="s", ): """Simple umap scatter plot, with legend outside figure. @@ -68,6 +73,16 @@ def umap_scatter( :type label_alpha: float, optional :param label_size: size of cluster labels, defaults to 20 :type label_size: int, optional + :param highlight_cell_names: [description] + :type highlight_cell_names: [type] + :param highlight_marker_size: [description], defaults to 30 + :type highlight_marker_size: int, optional + :param highlight_marker_color: [description], defaults to 'r' + :type highlight_marker_color: str, optional + :param highlight_zorder: [description], defaults to 5 + :type highlight_zorder: int, optional + :param highlight_marker_style: [description], defaults to 's' + :type highlight_marker_style: str, optional :return: matplotlib figure and axes :rtype: (matplotlib.Figure, matplotlib.Axes) """ @@ -125,6 +140,17 @@ def umap_scatter( zorder=label_z_order, ) + # overlay: highlight cells + if highlight_cell_names is not None: + plt.scatter( + data.loc[highlight_cell_names][umap_1_key], + data.loc[highlight_cell_names][umap_2_key], + s=highlight_marker_size, + c=highlight_marker_color, + zorder=highlight_zorder, + marker=highlight_marker_style, + ) + sns.despine(ax=ax) plt.tight_layout() @@ -260,112 +286,189 @@ def horizontal_stacked_bar_plot( #### -# def stacked_density_plot( -# data, -# row_var, -# hue_var, -# value_var, -# col_var=None, -# overlap=False, -# suptitle=None, -# figsize=None, -# hue_order=None, -# row_order=None, -# palette=None, -# ): -# """ -# Multiple density plot. -# Adapted from old work at https://github.com/hammerlab/infino/blob/develop/analyze_cut.py#L912 - -# For row_order, consider row_order=reversed(list(range(data.ylevel.values.max()+1))) -# """ - -# with sns.plotting_context("notebook"): -# with sns.axes_style("white", rc={"axes.facecolor": (0, 0, 0, 0)}): -# g = sns.FacetGrid( -# data, -# row=row_var, -# hue=hue_var, -# col=col_var, -# row_order=row_order, -# hue_order=hue_order, -# aspect=15, -# height=0.5, -# palette=palette, -# sharey=False, # important -- they don't share y ranges. -# ) - -# ## Draw the densities in a few steps -# # this is the shaded area -# g.map(sns.kdeplot, value_var, clip_on=False, shade=True, alpha=0.8, lw=2) - -# # this is the dividing horizontal line -# g.map(plt.axhline, y=0, lw=2, clip_on=False, ls="dashed") - -# ### Add label for each facet. - -# def label(**kwargs): -# """ -# kwargs is e.g.: {'color': (0.4918017777777778, 0.25275644444444445, 0.3333333333333333), 'label': 'Name of the row'} -# """ -# color = kwargs["color"] -# label = kwargs["label"] -# ax = plt.gca() # map() changes current axis repeatedly -# # x=1 if plot_on_right else 0; ha="right" if plot_on_right else "left", -# ax.text( -# 1.25, -# 0.5, -# label, -# # fontweight="bold", -# color=color, -# # ha="right", -# ha="left", -# va="center", -# transform=ax.transAxes, -# fontsize="x-small", -# # fontsize='x-large', #15, -# # bbox=dict(facecolor='yellow', alpha=0.3) -# ) - -# g.map(label) - -# ## Beautify the plot. -# g.set(xlim=(-0.01, 1.01)) -# # seems to do the trick along with sharey=False -# g.set(ylim=(0, None)) - -# # Some `subplots_adjust` line is necessary. without this, nothing appears -# if not overlap: -# g.fig.subplots_adjust(hspace=0) - -# # Remove axes details that don't play will with overlap -# g.set_titles("") -# # g.set_titles(col_template="{col_name}", row_template="") -# g.set(yticks=[], ylabel="") -# g.despine(bottom=True, left=True) - -# # fix x axis -# g.set_xlabels("Pseudotime") - -# # resize -# if figsize: -# g.fig.set_size_inches(figsize[0], figsize[1]) -# else: -# cur_size = g.fig.get_size_inches() -# increase_vertical = 3 # 7 #4 # 3 -# g.fig.set_size_inches(cur_size[0], cur_size[1] + increase_vertical) - -# if suptitle is not None: -# g.fig.suptitle(suptitle, fontsize="medium") - -# # tighten -# g.fig.tight_layout() - -# # then reoverlap -# if overlap: -# g.fig.subplots_adjust(hspace=-0.1) - -# return g, g.fig +def stacked_density_plot( + data, + cluster_label_key, + value_key, + xlabel=None, + suptitle=None, + figsize=(6, 8.5), + palette=None, + overlap=False, + **kwargs +): + """Plot probability densities by class. + + :param data: A dataframe where each row represents one observation belonging to a particular class. + :type data: pandas.DataFrame + :param cluster_label_key: Column name identifying the class of each observation. + :type cluster_label_key: str + :param value_key: Column name identifying the observation value. Values should range from 0 to 1. + :type value_key: str + :param xlabel: X-axis label, defaults to None + :type xlabel: str, optional + :param suptitle: Figure title above stacked density plot grid, defaults to None + :type suptitle: str, optional + :param figsize: Figure size, defaults to (6, 8.5) + :type figsize: tuple, optional + :param palette: Color palette for each class, defaults to None (in which case default palette used) + :type palette: matplotlib palette name, list of colors, or dict mapping class values to colors, optional + :param overlap: Whether to overlap the stacked density curves, defaults to False + :type overlap: bool, optional + :return: Figure + :rtype: matplotlib.Figure + """ + ## 1. remove without one-cell clusters because we can't draw a density for those + + clusters_to_keep = data[cluster_label_key].value_counts() + clusters_to_keep = clusters_to_keep[clusters_to_keep > 1].index.tolist() + plot_df = data[data[cluster_label_key].isin(clusters_to_keep)].copy() + + # TODO: bring back debug logging: + # print("Removed 1-cell clusters") + # print(data.shape, "-->", plot_df.shape) + + ## 2. Set a new ylevel that only contains the remaining clusters + + # problem with using cluster_label_key as ylevel: + # it's possibly a Categorical so it still contains the empty levels, so now those will have empty rows attached + # just convert back to strings from categorical + plot_df[cluster_label_key] = plot_df[cluster_label_key].astype("str") + + ## 3. Set row order + + # sort rows by median value + # so we have nicely increasing densities as we go down the plot + ylevel_order = ( + plot_df.groupby(cluster_label_key)[value_key] + .median() + .sort_values() + .index.tolist() + ) + + ## 4. Plot. + + return _stacked_density_facetgrid( + plot_df, + row_var=cluster_label_key, + hue_var=cluster_label_key, + value_var=value_key, + xlabel=xlabel, + overlap=overlap, + suptitle=suptitle, + figsize=figsize, + row_order=ylevel_order, + palette=palette, + **kwargs + ) + + +def _stacked_density_facetgrid( + data, + row_var, + hue_var, + value_var, + xlabel=None, + col_var=None, + overlap=False, + suptitle=None, + figsize=None, + hue_order=None, + row_order=None, + palette=None, +): + """ + Customizable multiple rows+columns density plot. + + Adapted from old work at https://github.com/hammerlab/infino/blob/develop/analyze_cut.py#L912 + + For row_order, consider row_order=reversed(list(range(data.ylevel.values.max()+1))) + """ + + with sns.plotting_context("notebook"): + with sns.axes_style("white", rc={"axes.facecolor": (0, 0, 0, 0)}): + grid = sns.FacetGrid( + data, + row=row_var, + hue=hue_var, + col=col_var, + row_order=row_order, + hue_order=hue_order, + aspect=15, + height=0.5, + palette=palette, + sharey=False, # important -- they don't share y ranges. + ) + + ## Draw the densities in a few steps + # this is the shaded area + grid.map(sns.kdeplot, value_var, clip_on=False, shade=True, alpha=0.8, lw=2) + + # this is the dividing horizontal line + grid.map(plt.axhline, y=0, lw=2, clip_on=False, ls="dashed") + + ### Add label for each facet. + + def label(**kwargs): + """ + default kwargs are color (RGB tuple) and label (name of the row) + other text options to consider passing in as kwargs: + - fontweight="bold" + - ha="right" + - bbox=dict(facecolor='yellow', alpha=0.3) + """ + color = kwargs.pop("color") + label = kwargs.pop("label") + ax = plt.gca() # map() changes current axis repeatedly + # x=1 if plot_on_right else 0; ha="right" if plot_on_right else "left", + ax.text( + 1.25, + 0.5, + label, + color=color, + ha="left", + va="center", + transform=ax.transAxes, + fontsize="x-small", + **kwargs + ) + + grid.map(label) + + ## Beautify the plot. + grid.set(xlim=(-0.01, 1.01)) + # seems to do the trick along with sharey=False + grid.set(ylim=(0, None)) + + # Some `subplots_adjust` line is necessary. without this, nothing appears + if not overlap: + grid.fig.subplots_adjust(hspace=0) + + # Remove axes details that don't play will with overlap + grid.set_titles("") + # g.set_titles(col_template="{col_name}", row_template="") + grid.set(yticks=[], ylabel="") + grid.despine(bottom=True, left=True) + + # fix x axis + if xlabel is not None: + grid.set_xlabels(xlabel) + + # resize + if figsize: + grid.fig.set_size_inches(figsize[0], figsize[1]) + + if suptitle is not None: + grid.fig.suptitle(suptitle, fontsize="medium") + + # tighten + grid.fig.tight_layout() + + # overlap + if overlap: + grid.fig.subplots_adjust(hspace=-0.1) + + return grid.fig # TODO: density umap plot diff --git a/genetools/trajectory.py b/genetools/trajectory.py new file mode 100644 index 0000000..65f773e --- /dev/null +++ b/genetools/trajectory.py @@ -0,0 +1,491 @@ +# run on all +import numpy as np +import pandas as pd +from genetools import plots, stats, helpers +import matplotlib.pyplot as plt +import seaborn as sns + + +def monte_carlo_pseudotime(adata_input, root_cells): + """[summary] + + TODO: document what must be done to adata before running dpt. i.e. PCA, neighbors, and diffmap. + + :param adata_input: [description] + :type adata_input: [type] + :param root_cells: [description] + :type root_cells: [type] + :return: [description] + :rtype: [type] + """ + import scanpy as sc + + adata = adata_input.copy() + + # each entry is one trajectory from one root cell: + # # a Series, values are the pseudotime values per-cell, and name is the root cell + pseudotime_vals = [] + + for root_cell in root_cells: + # Set root cell ID + adata.uns["iroot"] = np.where(adata.obs.index == root_cell)[0][0] + + # Run DPT. Scanpy stores output in obs 'dpt_pseudotime' + sc.tl.dpt(adata, n_dcs=min(adata.obsm["X_diffmap"].shape[1], 10), copy=False) + + # Save this trajectory + pseudotime_vals.append(adata.obs["dpt_pseudotime"].rename(root_cell)) + + # percentile normalize each pseudotime trajectory + pseudotime_vals = pd.concat(pseudotime_vals, axis=1) + pseudotime_vals_percentile_normalized = pseudotime_vals.apply( + stats.percentile_normalize, axis=0 + ) + + # Create dataframe where each row is a cell's value in a trajectory from a specific root cell. + # i.e. convert to format: one value per cell, times # experiments + df = pd.melt( + pseudotime_vals_percentile_normalized.rename_axis("cell_barcode").reset_index(), + id_vars="cell_barcode", + value_vars=root_cells, + var_name="root_cell", + value_name="pseudotime", + ) + + return df + + +def get_cells_at_percentile(trajectory_df, percentile): + """get points at percentile along trajectories + + :param trajectory_df: [description] + :type trajectory_df: [type] + :param percentile: [description] (should be between 0 and 1) + :type percentile: [type] + :return: [description] + :rtype: [type] + """ + if not (0.0 <= percentile <= 1.0): + raise ValueError("percentile must be between 0 and 1") + df = trajectory_df.copy() + df["diff_pseudotime"] = (df["pseudotime"] - percentile).abs() + return df.loc[df.groupby("root_cell", sort=False)["diff_pseudotime"].idxmin()] + + +def plot_stochasticity( + trajectory_df, + xlabel="Pseudotime", + ylabel="Number of unique barcodes in bin", + title="How stochastic are different parts of the trajectory?", +): + """[summary] + + :param trajectory_df: [description] + :type trajectory_df: [type] + :param xlabel: [description], defaults to "Pseudotime" + :type xlabel: str, optional + :param ylabel: [description], defaults to "Number of unique barcodes in bin" + :type ylabel: str, optional + :param title: [description], defaults to "How stochastic are different parts of the trajectory?" + :type title: str, optional + :return: [description] + :rtype: [type] + + # Note: + # expects percentile normalized pseudotime values, + # so that we have a roughly uniform number of values in each bin + """ + df = trajectory_df.copy() + + # make bins + pt_bins = np.histogram_bin_edges( + df["pseudotime"], bins=50, range=(0, 1) + ) # try bins='auto' + + # cut data into bins + df["bin_id"] = pd.cut(df["pseudotime"], pt_bins) # alternative: np.digitize + + # change bin ID to be left-offset of bin instead + df["bin_id"] = df["bin_id"].apply(lambda interval: interval.left) + + # group by bin ID, count unique barcodes in that bin + cells_by_bin = df.groupby("bin_id", sort=False)["cell_barcode"].nunique() + + fig = plt.figure() + plt.scatter(cells_by_bin.index, cells_by_bin.values) + plt.xlim(0, 1) + if xlabel is not None: + plt.xlabel(xlabel) + if ylabel is not None: + plt.ylabel(ylabel) + if title is not None: + plt.title(title) + return fig + + +def plot_dispersions( + adata, + mean_pseudotime_per_cell, + std_pseudotime_per_cell, + hue_key=None, + palette=None, + figsize=(8, 6), + xlabel="Mean pseudotime per cell", + ylabel="Standard deviation of pseudotime per cell", + **kwargs +): + """For each cell, scatter-plot the standard deviation against the mean of its pseudotimes across trajectories. + + :param adata: [description] + :type adata: [type] + :param mean_pseudotime_per_cell: [description] + :type mean_pseudotime_per_cell: [type] + :param std_pseudotime_per_cell: [description] + :type std_pseudotime_per_cell: [type] + :param hue_key: [description], defaults to None + :type hue_key: [type], optional + :param palette: [description], defaults to None + :type palette: [type], optional + :param figsize: [description], defaults to (8,6) + :type figsize: tuple, optional + :raises ValueError: [description] + :return: [description] + :rtype: [type] + """ + pseudotime_mean_key, pseudotime_std_key = "pseudotime_mean", "pseudotime_std" + if hue_key is not None and hue_key in [pseudotime_mean_key, pseudotime_std_key]: + raise ValueError( + "Cannot specify %s or %s as hue key, because they will be overwritten" + % (pseudotime_mean_key, pseudotime_std_key) + ) + + # align by common index + plot_data = helpers.merge_into_left( + adata.obs, mean_pseudotime_per_cell.rename(pseudotime_mean_key) + ) + plot_data = helpers.merge_into_left( + plot_data, std_pseudotime_per_cell.rename(pseudotime_std_key) + ) + + # drop missing data + cols_keep = [pseudotime_mean_key, pseudotime_std_key] + if hue_key is not None: + cols_keep.append(hue_key) + plot_data = plot_data[cols_keep].dropna(how="any") + + # plot + fig, ax = plt.subplots(figsize=figsize) + sns.scatterplot( + data=plot_data, + x=pseudotime_mean_key, + y=pseudotime_std_key, + hue=hue_key, + hue_order=( + sorted(plot_data[hue_key].unique()) if hue_key is not None else None + ), + palette=( + plots._verify_or_create_palette(palette, plot_data, hue_key) + if hue_key is not None + else None + ), + ax=ax, + legend="full", + **kwargs + ) + + plt.xlabel(xlabel) + plt.ylabel(ylabel) + + sns.despine(ax=ax) + plt.tight_layout() + plots._pull_legend_out_of_figure() + + return fig, ax + + +def standard_deviation_per_cell(trajectory_df, barcode_key="cell_barcode"): + """[summary] + + :param trajectory_df: [description] + :type trajectory_df: [type] + :param barcode_key: [description], defaults to "cell_barcode" + :type barcode_key: str, optional + :return: [description] + :rtype: [type] + """ + return ( + trajectory_df.groupby([barcode_key], sort=False)["pseudotime"] + .std() + .rename("stdev_across_trajectories") + ) + + +def mean_order(trajectory_df, barcode_key="cell_barcode"): + """[summary] + ## Ensemble: Mean trajectory + # Take each cell's mean across all trajectories + # "project cells onto this ensemble medoid trajectory" + + trajectories have different linear scales: + + - so percentile normalize each trajectory first: replace start value by 0, end by 1, median by 0.5. now they have same scale, so averaging sensible (already done) + - make sure nans→0s are being treated carefully, i.e. they are not being ranked ordinally. the pre-existing 0s (start cells) should be normalized to 0. + - actually there are no nans above. + + In tests: + # make sure concattable to obs + assert all(mean_trajectory.index == adata.obs.index) + + # adata.obs["mean_trajectory"] = ensemble_pt_trajectories( + # adata, col_names, pseudotime_vals, roots + # ) + + :param trajectory_df: [description] + :type trajectory_df: [type] + :param barcode_key: [description], defaults to "cell_barcode" + :type barcode_key: str, optional + :raises ValueError: [description] + :return: [description] + :rtype: [type] + """ + if any(trajectory_df["pseudotime"].isna()): + raise ValueError("Some cells are unreachable (NaN pseudotimes)") + + # Take each cell's mean across all trajectories + mean_trajectory = trajectory_df.groupby([barcode_key], sort=False)[ + "pseudotime" + ].mean() + + # renormalize + mean_trajectory = pd.Series( + stats.percentile_normalize(mean_trajectory), + index=mean_trajectory.index, + name="mean_order", + ) + + return mean_trajectory + + +def spectral_order( + trajectory_df, + root_cell_key="root_cell", + barcode_key="cell_barcode", + n_trajectories_sample=None, + n_cells_sample=None, +): + """[summary] + + Use a kNN graph from scanpy instead of computing a dense similarity matrix: + + 1. Make P, the cells-x-cells binary comparisons matrix of 1s and -1s based on cells beating each other in pseudotime trajectory, averaged over pseudotime trajectories. + - Subsample cells at random. + - Subsample matches at random. + - At a point the basic order will pop out, no need to be perfect. + 2. Normalize each row to have same Euclidean norm + 3. Create Anndata from P. (View P as the usual cells x features matrix, where each cell is featurized in terms of the others.) + 4. Run scanpy `sc.pp.neighbors`. Now you have a kNN graph, like a compressed version of the full similarity graph above. + 5. Laplacian on `adata.uns['neighbors']['connectivities']`, the adjacency matrix of the kNN graph. + 6. 2nd smallest eigenvector + + Previously, we were creating dense similarity matrix, where the similarity measure between two cells `i, j` is the sum over all other cells `k` of $1 - abs(P_{i,j} - P_{j,k})$. The intuition is that this gives you the ranking similarity between any two cells by comparing their probabilities of being ranked higher than other reference cells. Now are are simplifying how to create this similarity matrix. + + + # Note: expects percentile normalized input + + :param trajectory_df: [description] + :type trajectory_df: [type] + :param root_cell_key: [description], defaults to "root_cell" + :type root_cell_key: str, optional + :param barcode_key: [description], defaults to "cell_barcode" + :type barcode_key: str, optional + :param n_trajectories_sample: [description], defaults to None + :type n_trajectories_sample: [type], optional + :param n_cells_sample: [description], defaults to None + :type n_cells_sample: [type], optional + :raises ValueError: [description] + :return: [description] + :rtype: [type] + """ + + import scipy.sparse + from scipy.sparse.csgraph import laplacian + from scipy.sparse.linalg import eigsh + from sklearn import preprocessing + import anndata + import scanpy as sc + + if any(trajectory_df["pseudotime"].isna()): + raise ValueError("Some cells are unreachable (NaN pseudotimes)") + + if not n_cells_sample: + # default to 20% of cells + n_cells_sample = int(trajectory_df[barcode_key].nunique() * 0.2) + if not n_trajectories_sample: + # default to 10% of trajectories + n_trajectories_sample = int(trajectory_df[root_cell_key].nunique() * 0.1) + + cells_sample = np.random.choice( + trajectory_df[barcode_key].unique(), size=n_cells_sample, replace=False + ) + matches_sample = np.random.choice( + trajectory_df[root_cell_key].unique(), size=n_trajectories_sample, replace=False + ) + # n_experiments x n_cells, percentile normalized + pseudotime_ranks_subsample = trajectory_df[ + (trajectory_df[root_cell_key].isin(matches_sample)) + & (trajectory_df[barcode_key].isin(cells_sample)) + ].pivot(index=root_cell_key, columns=barcode_key, values="pseudotime") + + # make P, the transition matrix + matches = scipy.sparse.csr_matrix( + (pseudotime_ranks_subsample.shape[1], pseudotime_ranks_subsample.shape[1]) + ) + + # create lower triangular ones + template_mat = scipy.sparse.tril(np.ones(matches.shape), format="csr") + + # for every experiment, get C comparisons matrix + for _, row in pseudotime_ranks_subsample.iterrows(): + + # reorder template matrix to match real order + # basically, we want to move column #index to be column #argsort + # and same for rows + # where #index is an index into the lower-triangular matrix created above + # and #argsort is the argsort of the real order + + # e.g. suppose the values are [4,1,2,3] + # then the argsort is [2,3,4,1] + # the index (as always) into the lower-triangular matrix is [1,2,3,4] to start + # (in actuality the index and argsort would be zero-indexed. not the raw values) + # so you would move column 1 to be column 2, column 2 to be column 3, + # column 3 to be column 4, and column 4 to be column 1; and same for rows. + # i.e. column #index --> column #argsort + + # we achieve this by creating a series from the argsort, so we will have + # 1 2 + # 2 3 + # 3 4 + # 4 1 + + # then we sort the values + # that rearranges the index into 4,1,2,3 + # which is the order we put rows and columns of the lower-triangular matrix into. + + reorder_template = pd.Series(np.argsort(row.values)).sort_values().index.values + # note: this is a copy not a view unfortunately. + matches += template_mat[reorder_template, :][:, reorder_template] + + # compute sample probabilities + Q = matches + + # normalize each row to have same euclidean norm + Q_normalized = preprocessing.normalize(Q, norm="l2") + + # kNN graph + Q_adata = anndata.AnnData(Q_normalized) + sc.pp.neighbors(Q_adata, use_rep="X") + + # laplacian of the adjacency matrix of the kNN graph + laplacian_mat = laplacian(Q_adata.uns["neighbors"]["connectivities"], normed=False) + + # take second smallest eigenvalue + _, eigenvecs = eigsh(laplacian_mat, 2, which="SM", tol=1e-8) + + # order along second smallest eigenvalue + labels = pseudotime_ranks_subsample.columns + order = [label for (_, label) in sorted(zip(eigenvecs[:, 1], labels))] + + # assign orders 0 -> 1, normalize / num_cells + spectral_order_series = pd.Series( + np.arange(len(order)) / len(order), index=order, name="spectral_order" + ) + + # sometimes eigenvector signs will be flipped, giving the reverse ordering. + # in these cases we should just reverse the ordering. + + # to determine whether this is the case: + # choose a cell we know to be closer to the start (or end), and ensure it is closer to the start (or end). + # for instance, run the mean ensembling, check whether the start or end point is ahead in the spectral ordering. + # flip order as needed to match. + + # another way to do this: + # look at mean spectral order of the root cells of all trajectories + # flip order if needed. + # (note: we subsampled the cells before running spectral ordering. + # so instead of using 0-position, use min-position cells from the trajectories.) + filtered_trajectories = trajectory_df[ + (trajectory_df[root_cell_key].isin(matches_sample)) + & (trajectory_df[barcode_key].isin(cells_sample)) + ] + if ( + spectral_order_series[ + filtered_trajectories.loc[ + filtered_trajectories.groupby(root_cell_key)["pseudotime"].idxmin() + ][barcode_key] + ].median() + > 0.5 + ): + spectral_order_series = 1 - spectral_order_series + + return spectral_order_series + + +def compare_orders( + trajectory_1, + trajectory_2, + trajectory_1_label="Trajectory 1", + trajectory_2_label="Trajectory 2", + title=None, +): + """[summary] + + # per-cell difference in ordering, hist + + # adata.obs["mean_trajectory"], adata.obs["spectral_order"] + # xlabel="Mean over percentile normalized trajectories" + # ylabel="Spectral order" + # title="Spectral vs mean pseudotime ensemble" + + :param trajectory_1: [description] + :type trajectory_1: [type] + :param trajectory_2: [description] + :type trajectory_2: [type] + :param trajectory_1_label: [description], defaults to 'Trajectory 1' + :type trajectory_1_label: str, optional + :param trajectory_2_label: [description], defaults to 'Trajectory 2' + :type trajectory_2_label: str, optional + :param title: [description], defaults to None + :type title: [type], optional + :return: [description] + :rtype: [type] + """ + + # merge indexes + compare_aggregates = pd.merge( + trajectory_1, trajectory_2, left_index=True, right_index=True, how="inner" + ) + trajectory_1, trajectory_2 = ( + compare_aggregates[trajectory_1.name], + compare_aggregates[trajectory_2.name], + ) + # print(compare_aggregates.shape) + + compare_aggregates["diff"] = trajectory_1 - trajectory_2 + compare_aggregates["abs_diff"] = compare_aggregates["diff"].abs() + # compare_aggregates.head() + + fig, ax = plt.subplots(figsize=(8, 8)) + plt.scatter(trajectory_1, trajectory_2, s=1, zorder=10) + + # plot y=x + plt.plot([0, 1], [0, 1], "k", alpha=0.75, zorder=0) + + ax.set_aspect("equal") + + plt.xlabel(trajectory_1_label) + plt.ylabel(trajectory_2_label) + if title is not None: + plt.title(title) + sns.despine() + + return fig, compare_aggregates["abs_diff"].median() diff --git a/tests/baseline/test_mean_vs_spectral_orderings.png b/tests/baseline/test_mean_vs_spectral_orderings.png new file mode 100644 index 0000000..70af8c0 Binary files /dev/null and b/tests/baseline/test_mean_vs_spectral_orderings.png differ diff --git a/tests/baseline/test_plot_dispersions.png b/tests/baseline/test_plot_dispersions.png new file mode 100644 index 0000000..74eab11 Binary files /dev/null and b/tests/baseline/test_plot_dispersions.png differ diff --git a/tests/baseline/test_plot_mean_trajectory_on_umap.png b/tests/baseline/test_plot_mean_trajectory_on_umap.png new file mode 100644 index 0000000..b501204 Binary files /dev/null and b/tests/baseline/test_plot_mean_trajectory_on_umap.png differ diff --git a/tests/baseline/test_plot_roots.png b/tests/baseline/test_plot_roots.png new file mode 100644 index 0000000..58cf1c2 Binary files /dev/null and b/tests/baseline/test_plot_roots.png differ diff --git a/tests/baseline/test_plot_spectral_trajectory_on_umap.png b/tests/baseline/test_plot_spectral_trajectory_on_umap.png new file mode 100644 index 0000000..14a6d1b Binary files /dev/null and b/tests/baseline/test_plot_spectral_trajectory_on_umap.png differ diff --git a/tests/baseline/test_plot_stdev_trajectory_on_umap.png b/tests/baseline/test_plot_stdev_trajectory_on_umap.png new file mode 100644 index 0000000..167d3c1 Binary files /dev/null and b/tests/baseline/test_plot_stdev_trajectory_on_umap.png differ diff --git a/tests/baseline/test_stacked_density_plot.png b/tests/baseline/test_stacked_density_plot.png new file mode 100644 index 0000000..397f2d7 Binary files /dev/null and b/tests/baseline/test_stacked_density_plot.png differ diff --git a/tests/baseline/test_stacked_density_plot_overlap_no_labels.png b/tests/baseline/test_stacked_density_plot_overlap_no_labels.png new file mode 100644 index 0000000..3589c12 Binary files /dev/null and b/tests/baseline/test_stacked_density_plot_overlap_no_labels.png differ diff --git a/tests/baseline/test_stochasticity_plot.png b/tests/baseline/test_stochasticity_plot.png new file mode 100644 index 0000000..d53798d Binary files /dev/null and b/tests/baseline/test_stochasticity_plot.png differ diff --git a/tests/test_trajectory.py b/tests/test_trajectory.py new file mode 100644 index 0000000..3978b8f --- /dev/null +++ b/tests/test_trajectory.py @@ -0,0 +1,239 @@ +#!/usr/bin/env python + +"""Tests for trajectory analysis — tutorial order: + +- roots choose, plot, get value counts by cluster +- run trajectories +- end points get value counts by cluster, and plot +- stochasticity plot +- mean order +- spectral order: + n_trajectories_sample = 100 # number of trajectories to look at + n_cells_sample = 4400 # number of cells to look at + root_cell_key = 'root_cell' + barcode_key = 'full_barcode' + (After running: left merge into the obs?) +- compare_trajectories +- left merge mean_trajectory back +- plot trajectory on umap +- plot stacked density + +""" + + +import pytest +import numpy as np +import random +import matplotlib +import seaborn as sns + +matplotlib.use("Agg") + +from genetools import plots, trajectory, helpers + +random_seed = 12345 +np.random.seed(random_seed) +random.seed(random_seed) + +n_roots = 100 +# TODO: multiple root clusters +root_cluster = "B" + + +@pytest.fixture(scope="module") +def roots(adata): + # Choose roots + return helpers.sample_cells_from_clusters( + adata.obs, n_roots, "louvain", [root_cluster] + ) + + +def test_roots(adata, roots): + # Confirm roots come from the right clusters + assert len(roots) == n_roots + assert all(adata[roots].obs["louvain"] == root_cluster) + + +@pytest.mark.mpl_image_compare(savefig_kwargs={"bbox_inches": "tight"}) +def test_plot_roots(adata, roots): + # Plot root points over umap + fig, _ = plots.umap_scatter( + data=adata.obs, + umap_1_key="umap_1", + umap_2_key="umap_2", + hue_key="louvain", + label_key="louvain", + highlight_cell_names=roots, + ) + return fig + + +@pytest.fixture(scope="module") +def trajectories(adata, roots): + # run pseudotime + return trajectory.monte_carlo_pseudotime(adata, roots) + + +def test_trajectory_shape(trajectories, adata, roots): + # Should have one value per cell, times the number of experiments + assert trajectories.shape[0] == adata.obs.shape[0] * len(roots) + + +def test_trajectory_ranges(trajectories): + assert all(trajectories["pseudotime"] >= 0.0) + assert all(trajectories["pseudotime"] <= 1.0) + assert not any(trajectories["pseudotime"].isna()) + + +def test_trajectory_colnames(trajectories): + expected_cols = ["cell_barcode", "root_cell", "pseudotime"] + for col in expected_cols: + assert col in trajectories.columns + + +def test_get_end_points(trajectories, roots): + df = trajectory.get_cells_at_percentile(trajectories, 1.0) + assert df.shape[0] == len(roots) + assert all(df["pseudotime"]) == 1.0 + + +@pytest.mark.xfail(raises=ValueError) +def test_get_cells_at_percentile_bounds(trajectories): + # reject if not a percentile + trajectory.get_cells_at_percentile(trajectories, 50) + + +@pytest.mark.mpl_image_compare(savefig_kwargs={"bbox_inches": "tight"}) +def test_stochasticity_plot(trajectories): + return trajectory.plot_stochasticity(trajectories) + + +@pytest.fixture(scope="module") +def stdev_by_cell(trajectories): + return trajectory.standard_deviation_per_cell(trajectories, "cell_barcode") + + +@pytest.fixture(scope="module") +def mean_ordering(trajectories): + return trajectory.mean_order(trajectories, "cell_barcode") + + +@pytest.fixture(scope="module") +def spectral_ordering(trajectories): + return trajectory.spectral_order(trajectories, "root_cell", "cell_barcode") + + +@pytest.mark.xfail(raises=ValueError) +def test_mean_ordering_rejects_unreachable_cells(trajectories): + df = trajectories.copy() + df["pseudotime"] = np.nan + trajectory.mean_order(df, "cell_barcode") + + +@pytest.mark.xfail(raises=ValueError) +def test_spectral_ordering_rejects_unreachable_cells(trajectories): + df = trajectories.copy() + df["pseudotime"] = np.nan + trajectory.spectral_order(df, "root_cell", "cell_barcode") + + +@pytest.mark.mpl_image_compare(tolerance=10) +def test_mean_vs_spectral_orderings( + adata, trajectories, mean_ordering, spectral_ordering +): + # confirm basics about both series (use ", series.name" to emit series name in errors) + for series in [mean_ordering, spectral_ordering]: + # check range + assert all(series >= 0.0), series.name + assert all(series <= 1.0), series.name + + # no NaNs + assert not any(series.isna()), series.name + + # every cell barcode included (exception: spectral order takes a subset) + assert series.shape[0] <= trajectories["cell_barcode"].nunique(), series.name + + # no duplicate barcodes, i.e. only one entry per cell barcode + # we want no repeats in the order + assert not any(series.index.duplicated()), series.name + + # test merge back into adata obs, to confirm mergeability and validate 1:1 + helpers.merge_into_left(adata.obs, series) + + # compare orderings + fig, median_abs_diff = trajectory.compare_orders( + mean_ordering, + spectral_ordering, + trajectory_1_label="Mean order", + trajectory_2_label="Spectral order", + title="Mean vs spectral", + ) + assert median_abs_diff <= 0.2 + return fig + + +# Note we can't parametrize with fixtures yet: https://docs.pytest.org/en/latest/proposals/parametrize_with_fixtures.html +# So we split the parametrization with mean_ordering, spectral_ordering, and standard deviation across three tests +@pytest.mark.mpl_image_compare(savefig_kwargs={"bbox_inches": "tight"}) +def test_plot_mean_trajectory_on_umap(adata, mean_ordering): + return _umap_trajectory(adata, mean_ordering) + + +@pytest.mark.mpl_image_compare(savefig_kwargs={"bbox_inches": "tight"}, tolerance=10) +def test_plot_stdev_trajectory_on_umap(adata, spectral_ordering): + return _umap_trajectory(adata, spectral_ordering) + + +# high tolerance +# this test is flaky because of randomness of choosing which cells are included in the spectral ordering? +@pytest.mark.mpl_image_compare(savefig_kwargs={"bbox_inches": "tight"}, tolerance=20) +def test_plot_spectral_trajectory_on_umap(adata, stdev_by_cell): + return _umap_trajectory(adata, stdev_by_cell) + + +def _umap_trajectory(adata, computed_trajectory): + plot_data = helpers.merge_into_left( + adata.obs, computed_trajectory.rename("computed_trajectory") + ) + fig, ax = plots.umap_scatter( + data=plot_data, + umap_1_key="umap_1", + umap_2_key="umap_2", + hue_key="computed_trajectory", + continuous_hue=True, + label_key="louvain", + ) + ax.set_title("Computed pseudotime trajectory") + return fig + + +@pytest.mark.mpl_image_compare(savefig_kwargs={"bbox_inches": "tight"}) +def test_plot_dispersions(adata, mean_ordering, stdev_by_cell): + fig, _ = trajectory.plot_dispersions( + adata, mean_ordering, stdev_by_cell, hue_key="louvain" + ) + return fig + + +@pytest.mark.mpl_image_compare(savefig_kwargs={"bbox_inches": "tight"}) +def test_stacked_density_plot(adata, mean_ordering): + plot_data = helpers.merge_into_left(adata.obs, mean_ordering) + return plots.stacked_density_plot( + data=plot_data, + cluster_label_key="louvain", + value_key="mean_order", + xlabel="Pseudotime", + suptitle="Mean pseudotime trajectory", + ) + + +@pytest.mark.mpl_image_compare(savefig_kwargs={"bbox_inches": "tight"}) +def test_stacked_density_plot_overlap_no_labels(adata, mean_ordering): + plot_data = helpers.merge_into_left(adata.obs, mean_ordering) + return plots.stacked_density_plot( + data=plot_data, + cluster_label_key="louvain", + value_key="mean_order", + palette=sns.color_palette("Set2"), + overlap=True, + )