Joint reanalysis of CyteBase datasets

# Import the catalog, composite creator, and notebook display tools.
from pathlib import Path

import numpy as np
import pandas as pd
from IPython.display import display

import cytearc
from cytearc.composite import create_composite
from cytearc.cytebase import Catalog

# Choose resources for source reads, composite creation, and the pipeline.
memory_budget = "6G"
workers = 4

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

# Open the October 2026 release anonymously; counts remain remote.
catalog = Catalog(version="2610", token=False)

# Pin the four CyteBase datasets under names used in the joint analysis.
dataset_ids = {
    "blood_wilk": "wilk_2020_single_cell_atlas_peripheral_sars_cov_2_infection_456e8b9b",
    "blood_lee": "lee_2020_immunophenotyping_covid_19_development_severe_covid_19_de2c780c",
    "blood_arunachalam": "jin_2021_individual_single_cell_rna_seq_pbmc_data_arunachalam_et_al_59b69042",
    "blood_schulte_schrepping": "jin_2021_individual_single_cell_rna_seq_data_schulte_schrepping_et_al_5e717147",
}

# Open each source once and retain it for composite creation.
sources = {
    name: catalog.open_datastore(
        dataset_id,
        mem_budget=memory_budget,
        nthreads=workers,
    )
    for name, dataset_id in dataset_ids.items()
}

# Select primary COVID-19 and normal cells using the original source metadata.
rows = {}
for name, source in sources.items():
    primary = source.cells.fetch_all("is_primary_data")
    disease = source.cells.fetch_all("disease")
    selected = np.flatnonzero(primary & np.isin(disease, ["COVID-19", "normal"]))

    # Require every chosen study to contribute cells.
    if selected.size == 0:
        raise ValueError(f"No primary COVID-19 or normal cells selected for {name!r}")
    rows[name] = selected

# Show the complete selected size for each study before QC.
display(
    pd.DataFrame(
        {
            "source_cells": {name: source.cells.N for name, source in sources.items()},
            "selected_cells": {name: len(selected) for name, selected in rows.items()},
        }
    )
)

# Choose a new local destination outside every source store.
joint_path = Path("joint_cytebase.zarr")

# Create an ordinary writable DataStore with counts backed by the selected sources.
joint = create_composite(
    sources,
    at=joint_path,
    assay="RNA",
    features="intersection",
    rows=rows,
    mem_budget=memory_budget,
    nthreads=workers,
)

# Confirm the combined cells and shared genes.
print(f"{joint.cells.N:,} cells and {joint.RNA.feats.N:,} shared genes")

# Freeze the same design columns, with the Harmony column first, in both runs.
snapshot_columns = ["dataset", "orig_disease", "orig_donor_id"]
design = joint.cells.to_pandas_dataframe(snapshot_columns, key=None)

# Check condition coverage and donor representation within each study.
display(
    design.groupby(["dataset", "orig_disease"], observed=True).agg(
        cells=("orig_donor_id", "size"),
        donors=("orig_donor_id", "nunique"),
    )
)

# Check whether any within-study donor appears in both conditions.
donor_conditions = design.groupby(["dataset", "orig_donor_id"], observed=True)[
    "orig_disease"
].nunique()
print(
    f"Within-study donors represented in both conditions: {(donor_conditions > 1).sum()}"
)

# Define a human gene-symbol blacklist, including current histone names.
hvg_blacklist = "|".join(
    [
        r"^MT-",
        r"^RP[SL]\d+[AXY]?\d*$|^RPLP[0-2]$|^RPSA$|^MRP[SL]\d+[A-C]?$",
        r"^HLA-",
        r"^HIST[1-4]H|^H1-[1-6]$|^H1F[1-6]$|^H2AC\d+$|^H2BC\d+$"
        r"|^H3C\d+$|^H4C\d+$|^H2AW$|^H3-4$|^H4-16$",
        r"^XIST$|^DDX3Y$|^USP9Y$|^EIF1AY$|^KDM5D$|^SRY$|^ZFY$"
        r"|^UTY$|^TMSB4Y$|^NLGN4Y$",
    ]
)

# Keep the reviewed filtering and feature settings shared between comparison runs.
filtering = {"method": "mad", "sample_column": "dataset", "n_mads": 5}
pipeline_params = {
    "species": "homo_sapiens",
    "hvg": {"blacklist": hvg_blacklist},
}

# Run joint analysis without batch correction as the comparison baseline.
baseline = joint.pipeline.run(
    assay="RNA",
    label="joint_baseline",
    filtering=filtering,
    snapshot_columns=snapshot_columns,
    params=pipeline_params,
)

# Audit QC retention against all selected cells in each study and condition.
retention = design[["dataset", "orig_disease"]].copy()
retention["retained"] = baseline.cells.fetch_all("I")
display(
    retention.groupby(["dataset", "orig_disease"], observed=True)["retained"]
    .agg(included="size", retained="sum", retained_fraction="mean")
    .round(3)
)

# Show the uncorrected embedding by source study and by cluster.
joint.plots.embedding(run=baseline, color_by=["dataset", "clusters"], n_columns=2)

# Save the corrected analysis as a separate run with matched upstream settings.
run = joint.pipeline.run(
    assay="RNA",
    label="joint_harmony",
    filtering=filtering,
    harmony_batch_columns=["dataset"],
    snapshot_columns=snapshot_columns,
    params=pipeline_params,
)

# Confirm that both comparisons start from the same frozen cells and PCA.
assert baseline["analysis_cell_selection"] == run["analysis_cell_selection"]
assert baseline["pca"] == run["pca"]

# Show the corrected embedding using the same coloring as the baseline.
joint.plots.embedding(run=run, color_by=["dataset", "clusters"], n_columns=2)

# Summarize the top five saved marker genes per corrected cluster.
markers = joint.markers.load(marker=run["markers"])
display(
    markers.groupby("group_id", sort=False)
    .head(5)
    .groupby("group_id", sort=False)["feature_name"]
    .agg(", ".join)
    .rename("top_markers")
)

# Read study membership and clustering from the run's frozen cell view.
assignments = run.cells.to_pandas_dataframe(["dataset", "clusters"])
display(pd.crosstab(assignments["dataset"], assignments["clusters"]))

# Keep the corrected run label for the reopen call below.
run_label = run.label

# Supply anonymous access for these public CyteBase sources on each reopen.
source_storage_options = {name: {"token": False} for name in dataset_ids}

# Reopen the composite through the ordinary DataStore constructor.
reopened = cytearc.DataStore(
    str(joint_path),
    min_features_per_cell=-1,
    mem_budget=memory_budget,
    nthreads=workers,
    source_storage_options=source_storage_options,
)

# Load the saved corrected run and confirm that its clustering is preserved.
reopened_run = reopened.pipeline.open(label=run_label)
assert reopened_run["clusters"] == run["clusters"]
print(f"Reopened {reopened_run.label}: {reopened_run.status}")