Extending CyteArc with custom analyses

# Calculate custom statistics from graph weights and count blocks.
import numpy as np

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 saved analysis and retain its exact results.
run = ds.pipeline.open(label="docs_default")

# Keep the reference to the saved connectivity graph.
graph_ref = run["connectivity_map"]

# Inspect the opened assays and their dimensions.
ds

# Load the symmetric connectivity graph as a sparse matrix.
graph = ds.graph.load(graph=graph_ref, symmetric=True, upper_only=False)

# Sum the edge weights connected to each cell.
graph_strength = np.asarray(graph.sum(axis=1)).ravel()

# Save the values in cell metadata using the stated selection.
ds.cells.insert(
    column_name="customGraphStrength", values=graph_strength, key="I", overwrite=True
)

# Summarize the graph size and the new per-cell connectivity statistic.
{
    "cells": int(graph.shape[0]),
    "edges": int(graph.nnz),
    "customGraphStrength mean": float(graph_strength.mean()),
}

# Locate high and low graph connectivity on the saved embedding.
ds.plots.embedding(layout=run["umap"], color_by="customGraphStrength")

# Locate the run's cells on the stored count-matrix axis.
cell_index = np.flatnonzero(run.cells.fetch_all("I"))

# Read the run's highly variable gene selection.
hvg_values = np.asarray(run.features.fetch_all("highly_variable_features"), dtype=bool)

# Locate those genes on the stored feature axis.
feature_index = np.flatnonzero(hvg_values)

# Create a lazy view of the selected cells and genes.
selected_counts = ds.RNA.rawData[:, feature_index][cell_index, :]

# Check the selected matrix dimensions before streaming counts.
{"selected cells": len(cell_index), "selected genes": len(feature_index)}

# Collect one small vector of detection counts per block.
detected_blocks = []

# Process a bounded count block without loading the full matrix.
for count_block in selected_counts.stream_blocks(
    nthreads=4, msg="Calculating custom detection statistic"
):

    # Keep this result in its original processing order.
    detected_blocks.append(np.count_nonzero(count_block, axis=1))

# Check how many cell results were collected across the blocks.
sum(len(block) for block in detected_blocks)

# Join the block results in their original cell order.
detected_hvgs = np.concatenate(detected_blocks)

# Save the values in cell metadata using the stated selection.
ds.cells.insert(
    column_name="customDetectedHVGs", values=detected_hvgs, key="I", overwrite=True
)

# Summarize the per-cell detected-gene counts.
{
    "cells": int(detected_hvgs.size),
    "customDetectedHVGs mean": float(detected_hvgs.mean()),
    "customDetectedHVGs max": int(detected_hvgs.max()),
}

# Keep cells above the lowest quarter of graph strength.
well_connected = graph_strength >= np.quantile(graph_strength, 0.25)

# Save the values in cell metadata using the stated selection.
ds.cells.insert(
    column_name="wellConnected",
    values=well_connected,
    fill_value=False,
    key="I",
    overwrite=True,
)

# Check how many active cells pass the custom selection.
{"selected cells": int(well_connected.sum()), "active cells": len(well_connected)}

# Compare detected-gene counts and the custom cell selection.
ds.plots.embedding(
    layout=run["umap"], color_by=["customDetectedHVGs", "wellConnected"], n_columns=2
)

# Choose a small marker panel for the comparison.
panel_genes = ["CD3D", "MS4A1", "CD14", "LYZ", "NKG7", "GNLY"]

# Materialize only the selected cells and marker panel.
adata = ds.to_anndata(cell_key="wellConnected", feature_names=panel_genes)

# Check the exported dimensions and feature identifiers.
adata.shape, adata.var_names.tolist()

# Create a temporary destination separate from the source store.
from pathlib import Path
from tempfile import TemporaryDirectory

export_directory = TemporaryDirectory()
subset_path = Path(export_directory.name) / "well_connected.zarr"

# Copy the chosen cells with all RNA genes and their existing filter flags.
writer = cytearc.SubsetZarr(
    zarr_loc=str(subset_path),
    assays=[ds.RNA],
    cell_key="wellConnected",
    reset_cell_filter=False,
)
writer.dump()

# Confirm the cell and gene counts in the independent subset.
subset = cytearc.DataStore(str(subset_path))
{"exported cells": subset.cells.N, "retained genes": subset.RNA.feats.N}