Pooled CRISPR screens

# Download and verify the published files.
import gzip
import hashlib
import shutil
import warnings
from pathlib import Path
from tempfile import TemporaryDirectory
from urllib.request import urlopen

# Read the count tables and assemble the store.
import anndata
import numpy as np
import pandas as pd
from scipy import sparse
import cytearc
from cytearc.plotting import FeatureRef
from cytearc.storage.count_matrix import CountMatrixPolicy

# Display the diagnostic figures.
import matplotlib.pyplot as plt

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

# Keep this example's downloads and store together.
work = TemporaryDirectory()
output_dir = Path(work.name)

# Pin the five GEO files so that a changed download cannot alter the example silently.
geo = "https://ftp.ncbi.nlm.nih.gov/geo"
sources = {
    "rna": (
        f"{geo}/samples/GSM4633nnn/GSM4633614/suppl/"
        "GSM4633614_ECCITE_cDNA_counts.tsv.gz",
        "2acb14a10b95f9b39536d0aec42896f30f6256332ccd336a2b01d037316b6080",
    ),
    "protein": (
        f"{geo}/samples/GSM4633nnn/GSM4633615/suppl/"
        "GSM4633615_ECCITE_ADT_counts.tsv.gz",
        "e0f6f1dd1fc36d29580a5f5cf13f0d21a98a9b0d1ebb9479ed6291a355d4cb7f",
    ),
    "hashtag": (
        f"{geo}/samples/GSM4633nnn/GSM4633616/suppl/"
        "GSM4633616_ECCITE_HTO_counts.tsv.gz",
        "269069672fc09710136ea693e7702295c02f6c38aa824933eb4529a4c75944f5",
    ),
    "guide": (
        f"{geo}/samples/GSM4633nnn/GSM4633618/suppl/"
        "GSM4633618_ECCITE_GDO_counts.tsv.gz",
        "93c67d2a92fb98fcbdfe28b9a0ac94ffee08da035b112d5452ceaeacdcb7570a",
    ),
    "metadata": (
        f"{geo}/series/GSE153nnn/GSE153056/suppl/GSE153056_ECCITE_metadata.tsv.gz",
        "7059b0cb857dc344e6bb53dbdc5ccadccb25bb069372ed6f45eff46b7e992a76",
    ),
}
paths = {}
for name, (url, expected_digest) in sources.items():
    path = output_dir / url.rsplit("/", 1)[1]

    # Copy the response in blocks without holding the file in memory.
    with urlopen(url, timeout=60) as response, path.open("wb") as destination:
        shutil.copyfileobj(response, destination)

    # Reject a file whose contents differ from the verified source.
    with path.open("rb") as source:
        actual_digest = hashlib.file_digest(source, "sha256").hexdigest()
    if actual_digest != expected_digest:
        raise ValueError(f"Checksum mismatch for {path.name}")
    paths[name] = path

# Show the size of each completed download.
pd.Series(
    {name: path.stat().st_size / 1_000_000 for name, path in paths.items()},
    name="Downloaded MB",
).round(2)

# Read the three small tables with cells as rows.
guide_counts = pd.read_csv(paths["guide"], sep="\t", index_col=0).T
hashtag_counts = pd.read_csv(paths["hashtag"], sep="\t", index_col=0).T
protein_counts = pd.read_csv(paths["protein"], sep="\t", index_col=0).T

# Read the authors' calls and quality metrics for each cell.
authors = pd.read_csv(paths["metadata"], sep="\t", index_col=0)

# Check that every table lists the same cells in the same order.
cells = guide_counts.index
for table in (hashtag_counts, protein_counts, authors):
    if not table.index.equals(cells):
        raise ValueError("The tables list different cells")

# Show the size of each table.
pd.DataFrame(
    {
        "Cells": [
            len(table) for table in (guide_counts, hashtag_counts, protein_counts)
        ],
        "Features": [
            table.shape[1] for table in (guide_counts, hashtag_counts, protein_counts)
        ],
    },
    index=["Guides", "Hashtags", "Proteins"],
)

# Measure the smallest entries, and each cell's top guide count.
guide_values = guide_counts.to_numpy()
evidence = pd.Series(
    {
        "Smallest entry": guide_values.min(),
        "Entries equal to 1 (%)": round(100 * float((guide_values == 1).mean()), 1),
        "Smallest top-guide count of a cell": guide_values.max(axis=1).min(),
    },
    name="Guide table",
    dtype=object,
)

# Restore the counts by removing the added one.
guide_counts = guide_counts - 1
evidence

# Read the RNA table one gene at a time, keeping each gene's nonzero counts.
with gzip.open(paths["rna"], "rt") as source:
    rna_cells = [cell.strip('"') for cell in next(source).rstrip("\n").split("\t")[1:]]
    genes, positions, counts, offsets = [], [], [], [0]
    for line in source:
        gene, values = line.split("\t", 1)
        gene_counts = np.fromstring(values, sep="\t", dtype=np.int64)
        if gene_counts.size != len(rna_cells):
            raise ValueError(f"Unexpected number of RNA counts for {gene}")
        nonzero = np.flatnonzero(gene_counts)
        genes.append(gene.strip('"'))
        positions.append(nonzero.astype(np.int32))
        counts.append(gene_counts[nonzero].astype(np.int32))
        offsets.append(offsets[-1] + nonzero.size)

# Check that the RNA table lists the same cells as the other tables.
if rna_cells != list(cells):
    raise ValueError("The RNA table lists different cells")

# Build a sparse matrix with cells as rows.
rna_counts = sparse.csr_matrix(
    (np.concatenate(counts), np.concatenate(positions), np.asarray(offsets)),
    shape=(len(genes), len(cells)),
).T.tocsr()
del positions, counts

# Show the size of the RNA table.
pd.Series(
    {"Genes": rna_counts.shape[1], "Nonzero counts": rna_counts.nnz}, name="RNA table"
)

# Name each feature's assay, and give the antibodies IDs that no gene uses.
feature_types = (
    ["RNA"] * len(genes)
    + ["CRISPR"] * guide_counts.shape[1]
    + ["HTO"] * hashtag_counts.shape[1]
    + ["ADT"] * protein_counts.shape[1]
)
feature_names = (
    genes
    + list(guide_counts.columns)
    + list(hashtag_counts.columns)
    + list(protein_counts.columns)
)
feature_ids = feature_names[: -protein_counts.shape[1]] + [
    f"{name}-protein" for name in protein_counts.columns
]

# Combine the four tables, as 32-bit counts, into one AnnData with each cell's lane.
combined = anndata.AnnData(
    X=sparse.hstack(
        [rna_counts]
        + [
            sparse.csr_matrix(table.to_numpy(dtype=np.int32))
            for table in (guide_counts, hashtag_counts, protein_counts)
        ],
        format="csr",
    ),
    obs=pd.DataFrame({"lane": cells.str.split("_").str[0]}, index=cells),
    var=pd.DataFrame(
        {"feature_name": feature_names, "feature_types": feature_types},
        index=feature_ids,
    ),
)

# Write the AnnData, then release the tables it copied.
h5ad_path = output_dir / "papalexi_eccite.h5ad"
combined.write_h5ad(h5ad_path)
del combined, rna_counts

# Import the four assays into one store.
store_path = output_dir / "papalexi_eccite.zarr"
cytearc.H5adToZarr(
    cytearc.H5adReader(str(h5ad_path), feature_name_key="feature_name"),
    str(store_path),
    assay_split_key="feature_types",
    nthreads=4,
    mem_budget="4G",
    policy=CountMatrixPolicy(unitBytes=31_250_000, chunkBytes=3_125_000),
).dump()

# Open the store with RNA as the default assay.
ds = cytearc.DataStore(
    str(store_path), default_assay="RNA", nthreads=4, mem_budget="4G"
)
ds

# Align the authors' metadata with the store's cells.
authors = authors.loc[ds.cells.fetch_all("ids")]

# Record the hashtag that the authors called in each cell as its replicate.
ds.cells.insert("replicate", authors["MULTI_ID"].to_numpy())

# Compare the mean normalized count of each hashtag between the replicates.
hashtags = ds.HTO.feats.fetch_all("ids").tolist()
heatmap = ds.plots.matrixplot(
    features=hashtags,
    group_by="replicate",
    from_assay="HTO",
    figsize=(4, 6),
    show=False,
)

# Label the color scale for the hashtag counts.
heatmap.axes["matrixplot"].images[0].colorbar.set_label("Mean CLR hashtag count")
plt.show()

# Read each guide's target from its ID, and name PD-L1 by its gene symbol.
guides = ds.CRISPR.feats.fetch_all("ids")
targets = (
    pd.Series(guides).str.replace(r"g\d+$", "", regex=True).replace({"PDL1": "CD274"})
)

# Store the targets as the guide assay's target columns.
ds.CRISPR.feats.insert("target_gene_name", targets.to_numpy())
ds.CRISPR.feats.insert("target_gene_id", targets.to_numpy())

# Count the guides of each target.
targets.value_counts().rename("Guides").to_frame().T

# Analyze RNA over all cells: variable genes, PCA, graph, and UMAP.
rna = ds.pipeline.rna(
    label="rna",
    filtering=False,
    hvg={"top_n": 2000},
    pca={"dims": 40},
    leiden=False,
    doublets=False,
    markers=False,
    snapshot_columns=["replicate", "lane"],
)

# Color the map by replicate and by cell-cycle phase.
ds.plots.embedding(run=rna, color_by=["replicate", "cell_cycle_phase"], n_columns=2)

# Call guides over the RNA run's cells, keeping the warnings.
cells = rna["analysis_cell_selection"]
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    calls = ds.perturbation.call_guides(cells)

# Print the guides whose fit calls no cell.
for warning in caught:
    print(warning.message)

# Count the cells with no called guide, one guide, and several guides.
guide_calls = ds.perturbation.load_guide_calls(calls).cells
guide_calls["call"].where(
    guide_calls["call"].isin(["Negative", "Multiple"]), "One guide"
).value_counts()

# Plot guides per cell, cells per guide, and each cell's top-guide share.
ds.plots.guide_calls(calls)

# Show the fit of a targeting guide and of a non-targeting guide.
ds.plots.guide_thresholds(calls, guides=["IRF1g4", "NTg4"])

# Fit each guide in each lane separately; the table below counts the failed fits.
with warnings.catch_warnings(record=True):
    warnings.simplefilter("always")
    per_lane = ds.perturbation.call_guides(cells, sample_column="lane")

# Compare the fits that call no cell and the calls of the two choices.
comparison = {}
for name, ref in {
    "One fit per guide": calls,
    "One fit per guide and lane": per_lane,
}.items():
    fits = ds.perturbation.load_guide_thresholds(ref).table
    called = ds.perturbation.load_guide_calls(ref).cells["call"]
    comparison[name] = {
        "Fits": len(fits),
        "Fits that call no cell": int((fits["status"] != "fitted").sum()),
        "Cells with one guide": int((~called.isin(["Negative", "Multiple"])).sum()),
        "Cells with several guides": int((called == "Multiple").sum()),
        "Cells with no guide": int((called == "Negative").sum()),
    }
pd.DataFrame(comparison)

# Compare each single-guide call with the authors' guide.
called = guide_calls["call"]
single = ~called.isin(["Negative", "Multiple"])
agreement = called[single] == authors.loc[called.index[single], "guide_ID"]
pd.Series(
    {
        "Cells with one called guide": int(single.sum()),
        "Same guide as the authors (%)": round(100 * agreement.mean(), 2),
    },
    name="Guide calls",
    dtype=object,
)

# Label each cell by its target, with the non-targeting guides as controls.
labels = ds.perturbation.labels(calls, controls=["NT"])
labelled = ds.perturbation.load_labels(labels).cells

# Count the cells of each label.
labelled["label"].value_counts().rename("Cells").to_frame().T

# Find the cells with several guides that got one target or the control label.
merged = (labelled["guide_label"] == "Multiple") & (labelled["label"] != "Multiple")

# Compare the labels with the authors' targets, which name the control cells NT.
authors_label = authors["gene"].replace({"NT": "control"}).reindex(labelled.index)
compared = ~labelled["label"].isin(["Negative", "Multiple"])
differs = compared & (labelled["label"] != authors_label)

# Count the disagreements, and those where the authors labelled a control.
pd.Series(
    {
        "Cells with several guides and one label": int(merged.sum()),
        "Cells labelled by both": int(compared.sum()),
        "Same label as the authors (%)": round(100 * (1 - differs[compared].mean()), 1),
        "Different labels": int(differs.sum()),
        "Of these, a gene here and a control for the authors": int(
            (differs & (authors_label == "control")).sum()
        ),
    },
    name="Labels",
    dtype=object,
)

# Color the map by four targets and the control cells.
ds.plots.embedding(
    layout=rna["umap"],
    color_by=labels,
    groups=["control", "IFNGR2", "JAK2", "IRF1", "SMAD4"],
)

# Measure each target gene in its perturbed cells against the control cells.
efficiency = ds.perturbation.efficiency(labels)
measured = ds.perturbation.load_efficiency(efficiency).table

# Show the tested targets, strongest knockdown first.
columns = [
    "target",
    "n_cells",
    "control_mean_umi",
    "remaining_fraction",
    "adjusted_p_value",
]
measured.query("status == 'tested'").sort_values("remaining_fraction")[columns].round(3)

# Plot each target's log2 fold change against the control cells.
ds.plots.efficiency(efficiency)

# Measure each target's E-distance from the control cells in the RNA run's PCA.
edistance = ds.perturbation.edistance(labels, rna["pca"])
distances = ds.perturbation.load_edistance(edistance).table

# Plot the E-distances, largest first.
ds.plots.edistance(edistance)

# Compute each cell's perturbation signature within its replicate, over 40 PCs.
signature = ds.perturbation.signature(
    labels, rna["pca"], dims=40, sample_column="replicate"
)

# Classify each cell of each target as knocked out or not perturbed.
classes = ds.perturbation.mixscape(signature)
mixscape = ds.perturbation.load_mixscape(classes)

# Show each target's cells, KO cells and the genes that separate them.
targets_table = mixscape.targets.set_index("target")[
    ["status", "n_cells", "n_ko", "n_de_genes"]
]
targets_table["ko_share"] = (targets_table["n_ko"] / targets_table["n_cells"]).round(2)
targets_table.sort_values("ko_share", ascending=False)

# Plot each target's shares of KO and NP cells.
ds.plots.mixscape(classes)

# Join each cell's class with its called guide, and keep the cells of one IRF1 guide.
by_guide = mixscape.cells.join(guide_calls["call"])
irf1 = by_guide[by_guide["call"].str.startswith("IRF1g")]

# Count each IRF1 guide's KO and NP cells.
shares = pd.crosstab(irf1["call"], irf1["global_class"])
shares["KO share"] = (shares["KO"] / shares.sum(axis=1)).round(2)
shares

# Plot PD-L1 protein of the KO and NP cells of eight targets and of the control cells.
regulators = ["IFNGR1", "IFNGR2", "JAK2", "STAT1", "IRF1", "CUL3", "BRD4", "MYC"]
groups = ["control"] + [
    f"{target} {kind}" for target in regulators for kind in ("KO", "NP")
]
violins = ds.plots.distribution(
    FeatureRef("PDL1", assay="ADT", label="PD-L1 protein"),
    grouping=classes,
    groups=groups,
    max_points=0,
    figsize=(10, 4),
    show=False,
)

# Name the plotted measurement and the groups.
for ax in violins.axes.values():
    ax.set_xlabel("Mixscape class")
    ax.set_ylabel("PD-L1 protein (CLR)")
plt.show()

# Read each cell's normalized PD-L1 protein.
pdl1 = pd.Series(
    ds.get_cell_vals(from_assay="ADT", cell_key="I", k="PDL1"),
    index=ds.cells.fetch_all("ids"),
)

# Compare each class's median with the median of the control cells.
medians = pdl1.groupby(mixscape.cells["mixscape_class"]).median()
shift = (medians - medians["control"]).rename("PD-L1 median minus control")
shift.drop("control").reindex(groups[1:]).round(2).to_frame().T

# Keep the classes of at least 25 cells.
sizes = mixscape.cells["mixscape_class"].value_counts()
classes_shift = shift.drop("control")
classes_shift = classes_shift[sizes.reindex(classes_shift.index) >= 25]

# Show the eight classes with the lowest PD-L1 against the control cells.
classes_shift.sort_values().head(8).round(2).to_frame().T

# Test every target against every variable gene, with the replicate as a covariate.
tests = ds.perturbation.associations(
    labels, features=rna["highly_variable_features"], sample_column="replicate"
)
result = ds.perturbation.load_associations(tests)

# Summarize the calibration check and the discoveries.
pd.Series(
    {
        "Calibration grade": result.calibration_grade,
        "False discoveries in the calibration pairs": result.n_false_discoveries,
        "Calibration pairs": len(result.calibration),
        "Discovery pairs tested": int(result.discovery["pass_qc"].sum()),
        "Discoveries at FDR 0.1": int(result.discovery["significant"].sum()),
    },
    name="Association tests",
)

# Plot the calibration of the tests on the negative and positive controls.
ds.plots.association_calibration(tests)

# Show each target's test against its own gene, strongest decrease first.
columns = ["target", "n_cells", "z", "log2_fold_change", "p_value", "significant"]
result.positive.dropna(subset=["z"]).sort_values("z")[columns].round(3)

# Plot the fold change and p-value of every tested pair.
ds.plots.association_volcano(tests)

# Count each target's discoveries.
found = result.discovery.query("significant").groupby("target").size()
found.sort_values(ascending=False).rename("Discoveries").to_frame().T

# Profile each target over the variable genes, against each replicate's controls.
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    profiles = ds.perturbation.profiles(
        labels, rna["highly_variable_features"], sample_column="replicate"
    )

# Print the targets with too few cells to profile.
for warning in caught:
    print(warning.message)

# Cluster the targets with a significant E-distance.
strong = sorted(distances.query("adjusted_p_value < 0.05")["target"])
clusters = ds.perturbation.cluster(profiles, perturbations=strong, min_cluster_size=3)
grouped = ds.perturbation.load_clusters(clusters)

# Show each target's cluster and membership probability, by cluster.
order = grouped.table.sort_values(["cluster", "probability"], ascending=[True, False])
order[["cluster", "probability", "n_cells"]].round(2)

# Plot the correlations between the profiles, ordered by similarity.
ds.plots.perturbation_clusters(clusters)

# Join one row per target from each step's table.
effects = measured.set_index("target")[["n_cells", "remaining_fraction"]]
effects["edistance"] = distances.set_index("target")["edistance"]
effects["ko_share"] = targets_table["ko_share"]
effects["pdl1_ko_shift"] = [
    shift.get(f"{target} KO", np.nan) for target in effects.index
]
effects["discoveries"] = found.reindex(effects.index, fill_value=0)
effects["cluster"] = grouped.table["cluster"].reindex(effects.index).astype("Int64")

# Show the targets with at least 25 cells, by E-distance.
effects.query("n_cells >= 25").sort_values("edistance", ascending=False).round(2)

# Run the screen's analysis as one recipe; the run's report keeps its warnings.
with warnings.catch_warnings(record=True):
    warnings.simplefilter("always")
    screen = ds.pipeline.perturbation(
        cells=rna,
        labels={"controls": ["NT"]},
        hvg={"top_n": 2000},
        pca={"dims": 40},
        signature={"dims": 40, "sample_column": "replicate"},
        mixscape=True,
        associations={"sample_column": "replicate"},
        profiles={"sample_column": "replicate", "cells": "labelled"},
        clusters={"perturbations": strong, "min_cluster_size": 3},
        label="screen",
    )

# Show each stage's status and how many of its outputs it reused.
report = screen.report()
pd.DataFrame(
    [
        {
            "Stage": stage["stage"],
            "Status": stage["status"],
            "Outputs": len(stage["outputs"]),
            "Reused": sum(output["reused"] for output in stage["outputs"]),
        }
        for stage in report["stages"]
    ]
).set_index("Stage")

# Show the first cells' guide call, label, control flag and Mixscape class.
fields = ["guide_call", "perturbation_label", "is_control", "mixscape_class"]
screen.cells.to_pandas_dataframe(["ids", *fields]).head()

# Show the summaries of the stages that make the screen's decisions.
summaries = {stage["stage"]: stage["summary"] for stage in report["stages"]}
for name in (
    "guide_calling",
    "perturbation_labels",
    "efficiency",
    "edistance",
    "mixscape",
    "associations",
    "profiles",
    "perturbation_clusters",
):
    print(name, summaries[name])

# Print the warnings that the run recorded, with the stage that raised each one.
for warning in report["summary"]["warnings"]:
    print(warning["stage"], warning["message"])

# Sum each gene's counts per target and replicate.
pseudobulk = ds.compare.compute_bulk(
    labels,
    from_assay="RNA",
    secondary_groups="replicate",
    aggr_type="sum",
    feature_label="name",
    null_vals=["Negative", "Multiple"],
)

# Show a corner of the table: genes as rows, a target and replicate in each column.
pseudobulk.iloc[:5, :6]