Explore candidate cell fates

# Import array operations and CyteArc for defining trajectory endpoints.
import numpy as np
import pandas as pd

import cytearc
from cytearc.plotting import CellField, ColorScale

# Keep routine progress messages out of the teaching output.
cytearc.configure_output(level="WARNING", progress=False)

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

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

# Reuse the saved run and its frozen cell selection.
analysis_run = ds.pipeline.open(label="docs_default")

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

# Inspect the opened assays and their dimensions.
ds

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

# Mark ductal cells as the source population.
source = annotations == "Ductal"

# Mark the pooled terminal populations.
sink = np.isin(annotations, ["Alpha", "Beta", "Delta"])

# Require at least one annotated source cell and one sink cell.
if not source.any() or not sink.any():
    raise ValueError("Source and sink annotations must both be present")

# Start every graph cell with zero source or sink mass.
source_sink_vector = np.zeros(len(annotations), dtype=float)

# Share total mass -1 among the source cells.
source_sink_vector[source] = -1.0 / source.sum()

# Share total mass +1 among the pooled sink cells.
source_sink_vector[sink] = 1.0 / sink.sum()

# Orient the graph using the chosen source and sinks.
pseudotime_ref = ds.trajectory.pseudotime(graph, ss_vec=source_sink_vector)

# Check the number of cells used to orient the graph.
pd.Series({"source cells": int(source.sum()), "sink cells": int(sink.sum())})

# Keep the computed clustering used to define terminal boundaries.
sink_labels_ref = analysis_run["clusters"]

# Read its labels in the run's cell order.
sink_labels = analysis_run.cells.fetch("clusters")

# Check that imported and computed labels cover the same cell count.
assert len(annotations) == len(sink_labels)

# Measure each published annotation's share within each cluster.
shares = pd.crosstab(
    sink_labels,
    annotations,
    rownames=["cluster"],
    colnames=["annotation"],
    normalize="index",
)

# Show the majority annotation and its share for every cluster.
sink_table = pd.DataFrame(
    {"annotation": shares.idxmax(axis=1), "majority share": shares.max(axis=1)}
)

# Inspect every candidate cluster before choosing terminal boundaries.
sink_table

# Choose the three candidate terminal identities.
candidate_fates = ["Alpha", "Beta", "Delta"]

# Keep one cluster for each requested terminal identity.
selected = sink_table.loc[sink_table["annotation"].isin(candidate_fates)]

# Count the candidate fates before checking for missing or duplicated matches.
n_fates = len(candidate_fates)

# Check that each requested fate identifies exactly one cluster.
valid_fates = len(selected) == n_fates == selected["annotation"].nunique()

# Reject a missing or ambiguous terminal-cluster choice.
if not valid_fates:
    raise ValueError("Each candidate fate must match exactly one cluster")

# Keep the computed cluster label as a column.
selected = selected.reset_index(names="sink label")

# Index the table by the candidate annotation.
selected = selected.set_index("annotation")

# Match the order used for the probability columns and plots.
selected = selected.loc[candidate_fates]

# Extract the computed cluster labels for the terminal boundaries.
terminal_labels = selected["sink label"].astype(int).tolist()

# Attach readable fate names to those cluster labels.
sink_names = dict(zip(terminal_labels, candidate_fates, strict=True))

# Inspect the selected terminal clusters and their annotation purity.
selected

# Calculate absorption probabilities for the chosen endpoints.
fate_ref = ds.trajectory.fate(
    pseudotime_ref, sink_labels_ref, sinks=terminal_labels
)

# Load the probabilities and their validity mask.
fate = ds.trajectory.load_fate(fate_ref)

# Compare valid probability rows with the complete graph-cell count.
{"valid probability rows": int(fate.valid.sum()), "graph cells": len(fate.valid)}

# Match fate names to the saved probability-column order.
fate_names = [sink_names[int(label)] for label in fate.sink_labels]

# Keep valid rows and label each probability column by fate.
probabilities = pd.DataFrame(fate.values[fate.valid], columns=fate_names)

# Measure each probability row's deviation from a total of one.
probabilities["row-sum error"] = abs(probabilities.sum(axis=1) - 1.0)

# Inspect probability ranges and the error in their row sums.
probabilities.agg(["min", "median", "max"])

# Save each fate probability in the matching cell order.
for index, label in enumerate(fate.sink_labels):

    # Save the values in cell metadata using the stated selection.
    ds.cells.insert(
        f"fate_prob_{sink_names[int(label)]}",
        fate.values[:, index],
        key="I",
        overwrite=True,
    )

# Compare the three fate probabilities on a shared zero-to-one scale.
ds.plots.embedding(
    layout=analysis_run["umap"],
    color_by=[
        CellField(
            f"fate_prob_{name}", kind="continuous", label=f"{name} probability"
        )
        for name in candidate_fates
    ],
    n_columns=3,
    color_scale=ColorScale(scope="shared", vmin=0, vmax=1),
    sort_values=True,
)