HTO demultiplexing

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

# Read count tables and write a sparse matrix to CyteArc.
import numpy as np
import pandas as pd
from scipy.sparse import csr_matrix
import cytearc

# 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 two public GEO files so a changed download cannot alter the example silently.
sources = {
    "hto": (
        "https://ftp.ncbi.nlm.nih.gov/geo/samples/GSM2895nnn/GSM2895283/suppl/"
        "GSM2895283_Hashtag-HTO-count.csv.gz",
        "68ab6cd853964026a79f5c3da58a901808eab18c25684b3eecd41ebb71b9bce2",
    ),
    "rna": (
        "https://ftp.ncbi.nlm.nih.gov/geo/samples/GSM2895nnn/GSM2895282/suppl/"
        "GSM2895282_Hashtag-RNA.umi.txt.gz",
        "98a2be289ce79c2a2104eeaff3e39528ea25dfa0b3f3c1faf7246843f65ce3d9",
    ),
}
paths = {}
for assay, (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[assay] = path

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

# Read the small tag table and remove the three sequencing-summary rows.
hto_counts = pd.read_csv(paths["hto"], index_col=0).drop(
    index=["bad_struct", "no_match", "total_reads"]
).T

# Use barcodes as rows and short tag identifiers as columns.
hto_counts.index.name = "barcode"
hto_counts.columns = hto_counts.columns.str.split("-").str[0]
hto_counts.head()

# Sum RNA UMIs by barcode without loading the complete expression matrix.
with gzip.open(paths["rna"], "rt") as source:
    rna_barcodes = next(source).rstrip().split("\t")[1:]
    rna_totals = np.zeros(len(rna_barcodes), dtype=np.int64)
    for line in source:
        gene, values = line.split("\t", 1)
        gene_counts = np.fromstring(values, sep="\t", dtype=np.int64)
        if gene_counts.size != rna_totals.size:
            raise ValueError(f"Unexpected number of RNA counts for {gene}")
        rna_totals += gene_counts

# Require RNA support and a matching HTO barcode.
rna_totals = pd.Series(rna_totals, index=rna_barcodes, name="RNA_nCounts")
shared_barcodes = rna_totals.index.intersection(hto_counts.index)
selected_barcodes = shared_barcodes[rna_totals.loc[shared_barcodes] >= 200]
selected_counts = hto_counts.loc[selected_barcodes]

# Report the input population before assigning any tags.
pd.Series({
    "Barcodes in the HTO file": len(hto_counts),
    "Barcodes in both files": len(shared_barcodes),
    "Shared barcodes with at least 200 RNA UMIs": len(selected_counts),
}, name="Barcodes")

# Write the selected raw tag counts into a new HTO store.
store_path = output_dir / "cell_hashing.zarr"
cytearc.SparseToZarr(
    csr_matrix(selected_counts.to_numpy()),
    str(store_path),
    cell_ids=selected_counts.index.to_numpy(),
    feature_ids=selected_counts.columns.to_numpy(),
    assay_name="HTO",
    assay_type="HTO",
    nthreads=4,
    mem_budget="1G",
).dump()

# Preserve every RNA-supported barcode when opening the small tag panel.
ds = cytearc.DataStore(
    str(store_path),
    default_assay="HTO",
    min_features_per_cell=-1,
    nthreads=4,
    mem_budget="1G",
)

# Align the imported RNA totals to the datastore's barcodes.
ds.cells.insert("RNA_nCounts", rna_totals.loc[ds.cells.fetch_all("ids")].to_numpy())
ds

# Freeze the selected barcodes and demultiplex their raw HTO counts.
cell_selection = ds.snapshot_cell_selection("I")
identities = ds.qc.hto_demultiplexing(cell_selection, random_seed=0)

# Load assignments with their barcode index rather than relying on row order.
labels = ds.artifacts.load_values(identities).to_pandas()
labels.name = "HTO identity"

# Summarize the three classes while retaining each singlet's tag in labels.
classes = labels.where(labels.isin(["Negative", "Doublet"]), "Singlet")
class_counts = classes.value_counts().reindex(["Singlet", "Doublet", "Negative"])
pd.DataFrame({
    "Barcodes": class_counts,
    "Percent": (100 * class_counts / len(labels)).round(1),
})

# Count assignments for each tag and retain the two unresolved classes.
tag_names = ds.HTO.feats.fetch_all("ids").astype(str).tolist()
identity_order = tag_names + ["Doublet", "Negative"]
labels.value_counts().reindex(identity_order).rename("Barcodes").to_frame()

# Compare all eight tag signals across the saved assignment groups.
heatmap = ds.plots.matrixplot(
    features=tag_names,
    groups=identities,
    group_order=identity_order,
    figsize=(9, 4),
    show=False,
)

# Label the color scale and groups for the HTO assay.
heatmap.axes["matrixplot"].images[0].colorbar.set_label("Mean CLR HTO count")
heatmap.axes["matrixplot"].set_xlabel("Assigned tag or class")
plt.show()

# Inspect individual-cell distributions for two tags using the same saved identities.
distributions = ds.plots.distribution(
    ["BatchA", "BatchB"],
    grouping=identities,
    groups=identity_order,
    max_points=0,
    italicize_features=False,
    figsize=(10, 5),
    show=False,
)

# Name the plotted measurements and assignment groups.
for ax in distributions.axes.values():
    ax.set_xlabel("Assigned tag or class")
    ax.set_ylabel("CLR HTO count")
plt.show()

# Keep barcodes assigned to one of the eight tags.
singlets = ds.select_cells(identities, include=tag_names)
singlet_mask = np.asarray(ds.artifacts.load(singlets)["values"][:], dtype=bool)

# Confirm how many barcodes the selection retains and how many remain in the store.
pd.Series({
    "Selected singlets": int(singlet_mask.sum()),
    "Barcodes retained in the store": ds.cells.N,
}, name="Barcodes")