Order cell states with pseudotime

# Work with numeric arrays and cell masks.
import numpy as np

# Summarize cells and results in tables.
import pandas as pd

# Open count stores and run CyteArc analyses.
import cytearc

# Keep routine logs and progress bars out of the results.
cytearc.configure_output(level="WARNING", progress=False)

# Download the prepared example store.
dataset = cytearc.cytebase.connect("cytearc_docs").download_dataset(
    name="bastidas-ponce_4K_pancreas-d15_rnaseq",
    destination="cytearc_datasets",
    zarr=True,
)

# Open the count store for this analysis.
ds = cytearc.DataStore(f"{dataset}/data.zarr", nthreads=4)

# Open the saved pancreas analysis.
analysis_run = ds.pipeline.open(label="docs_default")

# Keep the graph from the saved analysis.
graph = analysis_run["connectivity_map"]

# Keep the full feature selection for marker testing.
all_features = analysis_run["feature_universe"]

# Inspect the opened store's cells and features.
ds

# Read the published cell types for the selected cells.
labels = ds.cells.fetch("clusters", key="I")

# Choose Ductal cells as the starting population.
source = labels == "Ductal"

# Choose the pooled endocrine endpoints.
sink = np.isin(labels, ["Alpha", "Beta", "Delta"])

# Stop if either chosen endpoint population is absent.
if not source.any() or not sink.any():
    raise ValueError("Source and sink labels must both be present")

# Start with zero weight for cells outside the endpoints.
source_sink_vector = np.zeros(len(labels), dtype=float)

# Share a total weight of minus one across the source cells.
source_sink_vector[source] = -1.0 / source.sum()

# Share a total weight of plus one across the sink cells.
source_sink_vector[sink] = 1.0 / sink.sum()

# Check the endpoint sizes and balanced total weights.
{
    "source cells": int(source.sum()),
    "sink cells": int(sink.sum()),
    "source weight": round(float(source_sink_vector[source].sum()), 6),
    "sink weight": round(float(source_sink_vector[sink].sum()), 6),
}

# Score pseudotime using the chosen source and sink weights.
pseudotime_ref = ds.trajectory.pseudotime(graph, ss_vec=source_sink_vector)

# Load pseudotime scores and their validity mask.
pseudotime = ds.trajectory.load_pseudotime(pseudotime_ref)

# Save the calculated values in the cell table.
ds.cells.insert("pseudotime", pseudotime.values, key="I", overwrite=True)

# Show how pseudotime changes across the pancreas UMAP.
ds.plots.embedding(layout=analysis_run['umap'], color_by='pseudotime', sort_values=True)

# Summarize the selected values and their spread.
pd.DataFrame(
    {
        "cell type": labels[pseudotime.valid],
        "pseudotime": pseudotime.values[pseudotime.valid],
    }
).groupby("cell type")["pseudotime"].describe()

# Test the selected genes for association with pseudotime.
marker_ref = ds.trajectory.markers(pseudotime_ref, features=all_features)

# Load the marker table and adjusted p-values.
markers = ds.trajectory.load_markers(marker_ref)

# Inspect the first few rows of the result.
markers.table.head()

# Keep genes with an adjusted p-value.
tested = markers.table.loc[
    markers.table["p_value_adjusted"].notna(),
    ["feature_name", "r_value", "p_value_adjusted"],
]

# Select the ten strongest positive correlations.
increasing = tested.loc[tested["r_value"] > 0].nlargest(10, "r_value")

# Select the ten strongest negative correlations.
decreasing = tested.loc[tested["r_value"] < 0].nsmallest(10, "r_value")

# Display increasing and decreasing genes together.
pd.concat({"increasing": increasing, "decreasing": decreasing})