Choosing dimensionality reductions

# Load tools for comparing partitions and displaying their layouts.
from itertools import combinations

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

import cytearc

# 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(
    "tenx_5K_pbmc_rnaseq", destination="cytearc_datasets", zarr=True
)

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

# Open the prepared baseline for the comparisons below.
baseline = ds.pipeline.open(label="docs_default")

# Keep the normalization fixed while comparing reductions.
normalized = baseline["normalized"]

# Reuse the saved UMAP coordinates.
umap = baseline["umap"]

# Inspect the opened assays and their dimensions.
ds

# Inspect the baseline clusters before changing PCA dimensions.
ds.plots.embedding(run=baseline, color_by="clusters")

# Compare three PCA dimension counts.
dimension_counts = (10, 15, 30)

# Reuse the baseline graph built from 15 PCs.
graph_15 = baseline["connectivity_map"]

# Reuse the baseline clustering for the 15-PC comparison.
cluster_refs = {15: baseline["leiden_0.5"]}

# Rebuild only the alternatives to the saved 15-PC analysis.
for dimensions in (10, 30):

    # Fit PCA with the current dimension count.
    pca = ds.reduction.pca(
        normalized, dims=dimensions, show_elbow_plot=dimensions == 30
    )

    # Build an index over these PCA coordinates.
    ann = ds.graph.ann_index(pca)

    # Query the same number of neighbors for every candidate.
    neighbors = ds.graph.neighbors(ann, k=11)

    # Build connectivity from this candidate's neighbors.
    graph = ds.graph.connectivity(neighbors)

    # Cluster this candidate graph at the fixed resolution.
    cluster_refs[dimensions] = ds.clusters.leiden(graph, resolution=0.5)

# Reuse the same initialization for the layout comparisons.
initialization_15 = baseline["embedding_initialization"]

# Load each candidate's labels in the common cell order.
cluster_values = {
    dimensions: np.asarray(
        ds.artifacts.load(cluster_refs[dimensions])["values"][:]
    )
    for dimensions in dimension_counts
}

# Count cells per cluster for each PCA dimension count.
cluster_sizes = pd.DataFrame(
    {
        dimensions: pd.Series(cluster_values[dimensions]).value_counts()
        for dimensions in dimension_counts
    }
).fillna(0).astype(int)

# Label both comparison axes before displaying the cluster sizes.
cluster_sizes.rename_axis(index="cluster", columns="PCs")

# Compare partition agreement between PCA dimension counts.
pd.Series(
    {
        f"{first} vs {second} PCs": ds.clusters.compute_concordance(
            cluster_refs[first], cluster_refs[second]
        )
        for first, second in combinations(dimension_counts, 2)
    },
    name="adjusted Rand index",
)

# Create one plotting axis for each comparison panel.
figure, axes = plt.subplots(1, 3, figsize=(12, 4))

# Draw each comparison on its own labeled axis.
for axis, dimensions in zip(axes, dimension_counts, strict=True):

    # Place each candidate clustering on the common baseline UMAP.
    ds.plots.embedding(
        layout=umap,
        color_by=cluster_refs[dimensions],
        target=axis,
        show_titles=False,
        show=False,
    )

    # Label the panel with the quantity being compared.
    axis.set_title(f"Leiden on {dimensions} PCs")

# Adjust spacing so panel labels remain readable.
figure.tight_layout()

# Display the completed figure.
figure

# Change UMAP packing while holding the graph and initialization fixed.
umap_tight = ds.embeddings.umap(graph_15, initialization_15, min_dist=0.1)

# Create one plotting axis for each comparison panel.
figure, axes = plt.subplots(1, 2, figsize=(10, 4))

# Draw each comparison on its own labeled axis.
for axis, layout, title in zip(
    axes, (umap, umap_tight), ("min_dist=1", "min_dist=0.1"), strict=True
):

    # Compare UMAP packing with the same cluster colors.
    ds.plots.embedding(
        layout=layout,
        color_by=cluster_refs[15],
        target=axis,
        show_titles=False,
        show=False,
    )

    # Label the panel with the quantity being compared.
    axis.set_title(title)

# Adjust spacing so panel labels remain readable.
figure.tight_layout()

# Display the completed figure.
figure

# Fit densMAP to the same graph and initialization.
densmap = ds.embeddings.umap(graph_15, initialization_15, use_density_map=True)

# Fit t-SNE to that graph without requesting iteration logs.
tsne = ds.embeddings.tsne(graph_15, initialization_15, verbose=False)

# Create one plotting axis for each comparison panel.
figure, axes = plt.subplots(1, 3, figsize=(12, 4))

# Pair each saved layout with its display title.
layout_comparisons = (("UMAP", umap), ("densMAP", densmap), ("t-SNE", tsne))

# Draw each comparison on its own labeled axis.
for axis, (title, layout) in zip(axes, layout_comparisons, strict=True):

    # Compare UMAP, densMAP, and t-SNE with the same cluster colors.
    ds.plots.embedding(
        layout=layout,
        color_by=cluster_refs[15],
        target=axis,
        show_titles=False,
        show=False,
    )

    # Label the panel with the quantity being compared.
    axis.set_title(title)

# Adjust spacing so panel labels remain readable.
figure.tight_layout()

# Display the completed figure.
figure