Find regulons and co-regulated gene programs

# Import tools for downloading and decompressing the reference files.
import gzip
import shutil
from pathlib import Path
from urllib.request import Request, urlopen

# Import plotting and table tools for inspecting the results.
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

# Import CyteArc's datastore and regulon methods.
import cytearc
from cytearc import regulons as rg

# Keep warnings visible while hiding routine progress messages.
cytearc.configure_output(level="WARNING", progress=False)

# Show whole gene lists in tables.
pd.set_option("display.max_colwidth", 120)

# Keep the downloaded reference files in one folder.
inputs = Path("regulon_inputs")

# Create the folder on the first run.
inputs.mkdir(exist_ok=True)


# Download a file once and return its local path.
def fetch(url):

    # Name the local copy after the last part of the URL.
    path = inputs / url.rsplit("/", 1)[1]

    # Reuse the copy from an earlier run.
    if path.exists():
        return path

    # Name the client, because some hosts refuse Python's default one.
    request = Request(url, headers={"User-Agent": "CyteArc documentation"})

    # Write to a temporary name, so that an interrupted download is never reused.
    partial = path.with_name(path.name + ".part")
    with urlopen(request) as response, open(partial, "wb") as target:
        shutil.copyfileobj(response, target)

    # Give the complete file its final name.
    return partial.rename(path)


# Download the 10x Genomics PBMC multiome sample (Cell Ranger ARC 2.0, GRCh38).
matrix_h5 = fetch(
    "https://cf.10xgenomics.com/samples/cell-arc/2.0.0/pbmc_granulocyte_sorted_3k/"
    "pbmc_granulocyte_sorted_3k_filtered_feature_bc_matrix.h5"
)

# Download the human transcription factors that the motif annotations cover.
tf_list = fetch("https://resources.aertslab.org/cistarget/tf_lists/allTFs_hg38.txt")

# Download the cisTarget database, which ranks every gene by each motif's best match near it.
ranking_db = fetch(
    "https://resources.aertslab.org/cistarget/databases/homo_sapiens/hg38/refseq_r80/mc_v10_clust/"
    "gene_based/hg38_10kbp_up_10kbp_down_full_tx_v10_clust.genes_vs_motifs.rankings.feather"
)

# Download the table that assigns each cisTarget motif to transcription factors.
motif_table = fetch(
    "https://resources.aertslab.org/cistarget/motif2tf/motifs-v10nr_clust-nr.hgnc-m0.001-o0.0.tbl"
)

# Download the JASPAR 2024 vertebrate motifs.
jaspar = fetch(
    "https://jaspar.elixir.no/download/data/2024/CORE/"
    "JASPAR2024_CORE_vertebrates_non-redundant_pfms_jaspar.txt"
)

# Use the GENCODE release that the 10x peaks were called against.
gencode = "https://ftp.ebi.ac.uk/pub/databases/gencode/Gencode_human/release_32/"

# Download the genome sequence.
genome_gz = fetch(gencode + "GRCh38.primary_assembly.genome.fa.gz")

# Download the gene annotation.
gene_gtf = fetch(gencode + "gencode.v32.primary_assembly.annotation.gtf.gz")

# Download SPI1 (PU.1) CUT&RUN peaks from primary human monocytes (GEO GSM8669389).
spi1_cutrun = fetch("https://chip-atlas.dbcls.jp/data/hg38/eachData/bed05/SRX26995709.05.bed")

# Name the decompressed genome, which motif scanning reads at many positions.
genome_fa = inputs / "GRCh38.primary_assembly.genome.fa"

# Decompress the genome once.
if not genome_fa.exists():
    with gzip.open(genome_gz, "rb") as source, open(genome_fa, "wb") as target:
        shutil.copyfileobj(source, target)

# Write both assays of the 10x matrix into one store.
cytearc.CrToZarr(cytearc.CrH5Reader(str(matrix_h5)), zarr_loc="pbmc_multiome.zarr").dump()

# Open the store for writing; the pipeline below decides which cells pass quality control.
ds = cytearc.DataStore(
    "pbmc_multiome.zarr",
    default_assay="RNA",
    min_features_per_cell=-1,
    nthreads=4,
)

# Run the standard RNA recipe: cell and gene selection, graph, UMAP, and clusters.
run = ds.pipeline.run(label="baseline")

# Keep the cells that the pipeline analysed.
cells = run["analysis_cell_selection"]

# Mark the analysed cells among all cells of the store.
analysed = np.asarray(run.cells.fetch_all("I"), dtype=bool)

# Inspect the store's assays and their dimensions.
ds

# Define canonical PBMC marker panels.
panels = {
    "CD14 monocyte": ["CD14", "LYZ", "VCAN", "S100A8"],
    "CD16 monocyte": ["FCGR3A", "MS4A7", "LST1"],
    "Dendritic cell": ["FCER1A", "CD1C", "CLEC10A"],
    "Plasmacytoid DC": ["LILRA4", "IL3RA", "CLEC4C"],
    "B cell": ["MS4A1", "CD79A", "BANK1", "PAX5"],
    "CD4 T cell": ["CD3E", "IL7R", "CD4"],
    "CD8 T cell": ["CD3E", "CD8A", "CD8B"],
    "NK cell": ["NKG7", "GNLY", "KLRD1"],
}

# Average the normalized expression of every gene within each cluster.
means = ds.compare.compute_bulk(run["clusters"], feature_label="name")

# Drop gene names that several features share, so that each name selects one row.
means = means.loc[~means.index.duplicated(keep=False)]

# Scale each gene across clusters; a gene that is constant across clusters stays at zero.
scaled = means.sub(means.mean(axis=1), axis=0).div(means.std(axis=1).replace(0, 1), axis=0)

# Score every panel in every cluster as the mean scaled expression of its genes.
panel_scores = pd.DataFrame({name: scaled.loc[genes].mean() for name, genes in panels.items()})

# Name each cluster after its best-scoring panel.
cluster_names = panel_scores.idxmax(axis=1)

# Name each analysed cell after its cluster.
cell_type = cluster_names.reindex(run.cells.fetch("clusters").astype(str)).to_numpy()

# Count the analysed cells of each population.
pd.Series(cell_type).value_counts()

# Start a random generator with a fixed seed, so that the sample is reproducible.
rng = np.random.default_rng(0)

# Mark 500 analysed cells, drawn at random, for network inference.
sample = np.zeros(ds.cells.N, dtype=bool)
sample[rng.choice(np.flatnonzero(analysed), size=500, replace=False)] = True

# Save the sample as a cell column.
ds.cells.insert("grn_cells", sample, overwrite=True)

# Freeze the sample as the cells that the network models.
grn_cells = ds.snapshot_cell_selection("grn_cells")

# Read the human transcription factors that the motif annotations cover.
tfs = tf_list.read_text().split()

# Read the gene name of every RNA feature.
names = np.asarray(ds.RNA.feats.fetch_all("names")).astype(str)

# Mark the highly variable genes that the pipeline selected.
variable = np.asarray(run.features.fetch_all("highly_variable_features"), dtype=bool)

# Find the genes detected in at least 1% of analysed cells.
detected = ds.features.detected(cells, min_cells=int(analysed.sum() // 100))

# Mark those genes.
expressed = np.asarray(ds.artifacts.load(detected)["values"][:], dtype=bool)

# Model the highly variable genes and every detected transcription factor.
grn_genes = ds.features.snapshot(
    from_assay="RNA",
    mask=variable | (expressed & np.isin(names, tfs)),
)

# Score every factor-to-gene link with GRNBoost2.
grn = ds.regulons.grn(grn_cells, features=grn_genes, tfs=tfs, random_seed=777)

# Load the links with their importance and correlation.
adjacencies = ds.regulons.load_grn(grn).adjacencies

# Inspect the strongest links.
adjacencies.head()

# Build pySCENIC's candidate modules from the activating links.
modules = rg.modules_from_adjacencies(adjacencies)

# Open the motif-ranking database.
database = rg.RankingDatabase(ranking_db)

# Read the transcription factors that each motif is annotated to.
annotations = rg.read_motif_annotations(motif_table)

# Keep the module targets that an enriched motif of the module's factor recovers.
enriched = rg.prune_modules(modules, [database], annotations, nthreads=4)

# Merge every factor's surviving targets into one regulon.
regulons = rg.regulons_from_motifs(enriched)

# List the ten largest regulons.
regulons.groupby("source").size().nlargest(10)

# Score the regulons with at least ten genes in every analysed cell.
regulon_auc = ds.scores.aucell(regulons, cells, features=run["feature_universe"], tmin=10)

# Read the scores as a cells-by-regulons table.
activity = ds.scores.load(regulon_auc).to_pandas()

# Score how specific each regulon's activity is to each population.
rss = rg.regulon_specificity(activity, cell_type)

# List the three most specific regulons of each population.
rss.apply(lambda scores: ", ".join(scores.nlargest(3).index), axis=1).to_frame("most specific regulons")

# Read the UMAP coordinates of the analysed cells.
umap = run.cells.to_pandas_dataframe(["umap_1", "umap_2"])

# Pick the most specific regulon of three lineages.
shown = rss.loc[["CD14 monocyte", "B cell", "NK cell"]].idxmax(axis=1).unique()

# Create one panel per regulon.
figure, axes = plt.subplots(1, len(shown), figsize=(4 * len(shown), 3.6))

# Color the cells of each panel by one regulon's activity.
for axis, name in zip(axes, shown, strict=True):

    # Plot every analysed cell at its UMAP position.
    points = axis.scatter(umap["umap_1"], umap["umap_2"], c=activity[name], s=3)

    # Name the regulon shown in the panel.
    axis.set_title(f"{name} AUCell")

    # Label the activity scale.
    figure.colorbar(points, ax=axis, label="AUCell")

# Make room for the panel titles and color scales.
figure.tight_layout()

# Display the regulon maps.
plt.show()

# Score how specifically each pair of regulons is active in the same cells.
csi = rg.connection_specificity(activity)

# Group the regulons into five co-activity modules.
regulon_groups = rg.regulon_modules(csi, n_modules=5)

# Average the activity of each module's regulons, then within each population.
module_activity = activity.T.groupby(regulon_groups).mean().T.groupby(cell_type).mean()

# List the regulons of each module and the population where the module is most active.
pd.DataFrame({
    "regulons": regulon_groups.groupby(regulon_groups).apply(lambda group: ", ".join(group.index[:8])),
    "most active in": module_activity.idxmax(),
})

# Group the analysed cells into small, similar groups (metacells) with a fine Leiden partition.
metacells = ds.clusters.leiden(run["connectivity_map"], resolution=15)

# Sum the RNA and ATAC counts of each metacell as log counts per million.
profiles = {assay: ds.regulons.profiles(metacells, from_assay=assay) for assay in ("RNA", "ATAC")}

# Keep the variable genes that name a single feature of the profiles.
program_genes = [gene for gene in names[variable] if gene in profiles["RNA"].columns]

# Split the variable genes into programs of co-varying genes.
programs = rg.gene_programs(profiles["RNA"][program_genes])

# List the size and the eight most typical genes of each program.
programs.groupby("source", sort=False)["target"].agg(
    size="size",
    leading_genes=lambda genes: ", ".join(genes.head(8)),
)

# Draw a table of population averages as a labelled heatmap.
def heatmap(table, label, **colors):

    # Widen the figure with the number of columns.
    figure, axis = plt.subplots(figsize=(0.45 * table.shape[1] + 2, 3.6))

    # Color every value of the table.
    image = axis.imshow(table.to_numpy(), aspect="auto", **colors)

    # Label the columns.
    axis.set_xticks(range(table.shape[1]), table.columns, rotation=90)

    # Label the rows.
    axis.set_yticks(range(table.shape[0]), table.index)

    # Label the color scale.
    figure.colorbar(image, ax=axis, label=label)

    # Fit the labels into the figure.
    figure.tight_layout()

    # Display the heatmap.
    plt.show()


# Score each program in every analysed cell.
program_auc = ds.scores.aucell(programs, cells, features=run["feature_universe"], tmin=10)

# Average each program's activity within each population.
program_means = ds.scores.load(program_auc).to_pandas().groupby(cell_type).mean()

# Show the averages as a heatmap.
heatmap(program_means, "mean AUCell", cmap="magma")

# Test every program for overlap with every regulon among the variable genes.
overlaps = rg.set_overlaps(programs, regulons, program_genes)

# Keep the three most significant overlapping regulons of each program.
shared = overlaps[overlaps["q_value"] < 0.05].sort_values("p_value").groupby("query").head(3)

# List them by program.
shared.groupby("query")["reference"].agg(", ".join).to_frame("overlapping regulons")

# Find motifs that rank each program's genes highly, and correlate each nominated
# factor's expression with its program across metacells.
program_tfs = rg.program_regulators(programs, database, annotations, profiles=profiles["RNA"])

# Keep the strongest motif of each factor in each program; rows are sorted by NES.
program_tfs = program_tfs.drop_duplicates(["program", "TF"])

# Keep the factors whose expression follows their program closely.
supported = program_tfs[program_tfs["correlation"] > 0.6]

# Rank the factors of each program by correlation.
supported = supported.sort_values(["program", "correlation"], ascending=[True, False])

# List the three best-correlated factors of each program.
supported.groupby("program").head(3)[["program", "TF", "motif", "NES", "correlation"]]

# Use every peak with a fragment among the analysed cells.
peaks = ds.features.detected(cells, from_assay="ATAC", min_cells=1)

# Mark those peaks.
selected = np.asarray(ds.artifacts.load(peaks)["values"][:], dtype=bool)

# Read their ids, which give their coordinates as chrom:start-end.
peak_ids = np.asarray(ds.ATAC.feats.fetch_all("ids")).astype(str)[selected]

# Name factors whose motifs mark PBMC populations.
panel = ["SPI1", "CEBPA", "CEBPB", "IRF8", "KLF4", "PAX5", "EBF1", "POU2F2", "TBX21",
         "EOMES", "RUNX3", "GATA3", "TCF7", "LEF1", "CTCF", "IRF4", "MEF2C", "NFKB1"]

# Index the JASPAR motifs by factor name, keeping each factor's first motif.
first_motif = {}
for motif in rg.read_jaspar_motifs(jaspar):
    first_motif.setdefault(motif.name.upper(), motif)

# Scan the first motif of each panel factor that JASPAR holds.
motifs = [first_motif[name] for name in panel if name in first_motif]

# Find the peaks that hold a match of each motif on either strand.
with rg.IndexedFasta(genome_fa) as genome:
    matches = rg.match_motifs(peak_ids, genome, motifs)

# Assign each motif's peaks to the factor that the motif is named after.
cistromes = rg.cistromes_from_motifs(matches, rg.motif_annotations(motifs))

# Count the matching peaks of each factor.
cistromes["source"].value_counts().to_frame("peaks with a match").T

# Score the accessibility of each factor's motif peaks in every analysed cell.
motif_activity = ds.scores.chromvar(
    cistromes,
    cells,
    from_assay="ATAC",
    features=peaks,
    genome=genome_fa,
)

# Average each factor's z-scores within each population.
motif_means = ds.scores.load(motif_activity).to_pandas().groupby(cell_type).mean()

# Show the averages as a heatmap centred on zero.
heatmap(motif_means, "mean chromVAR z-score", cmap="RdBu_r", vmin=-5, vmax=5)

# Map the SPI1 CUT&RUN peaks onto the ATAC peaks.
binding = rg.binding_cistromes({"SPI1": spi1_cutrun}, peak_ids)

# Count the ATAC peaks that SPI1 binds.
print(f"{len(binding):,} ATAC peaks overlap an SPI1 CUT&RUN peak")

# Score the accessibility of the bound peaks in every analysed cell.
bound_activity = ds.scores.chromvar(
    binding,
    cells,
    from_assay="ATAC",
    features=peaks,
    genome=genome_fa,
)

# Read the z-scores of the bound peaks.
bound_z = ds.scores.load(bound_activity).to_pandas()["SPI1"]

# Compare the mean z-scores of the bound peaks and of the SPI1 motif in each population.
pd.DataFrame({
    "SPI1 CUT&RUN peaks": bound_z.groupby(cell_type).mean(),
    "SPI1 motif": motif_means["SPI1"],
}).round(2)

# Read the protein-coding genes of the annotation.
genes_table = rg.read_gene_annotation(gene_gtf, gene_types={"protein_coding"})

# Pair each peak with the genes whose body lies within 150 kb of it.
candidates = rg.link_regions_to_genes(
    peak_ids,
    genes_table,
    upstream=150_000,
    downstream=150_000,
    anchor="gene_body",
)

# Keep the candidate links of genes that the co-expression network models.
candidates = candidates[candidates["gene"].isin(adjacencies["target"])]

# Score every candidate link across metacells.
links = rg.region_gene_links(profiles["ATAC"], profiles["RNA"], candidates, nthreads=4)

# Build eRegulons from the motif cistromes and from the SPI1 CUT&RUN peaks.
eregulons = pd.concat([
    rg.build_eregulons(cistromes, links, adjacencies, nthreads=4),
    rg.build_eregulons(binding, links, adjacencies, label="CUTRUN", nthreads=4),
])

# Count the peaks and genes of each eRegulon.
sizes = eregulons.groupby("source").agg(peaks=("region", "nunique"), genes=("target", "nunique"))

# List the activator eRegulons.
sizes.filter(like="+/+", axis=0)

# Score the target genes of every eRegulon in every analysed cell; a gene linked by
# several peaks counts once.
eregulon_auc = ds.scores.aucell(
    eregulons[["source", "target"]].drop_duplicates(),
    cells,
    features=run["feature_universe"],
    tmin=10,
)

# Average each activator eRegulon's activity within each population.
eregulon_means = ds.scores.load(eregulon_auc).to_pandas().filter(like="+/+").groupby(cell_type).mean()

# Name the population where each activator eRegulon is most active.
pd.DataFrame({"most active in": eregulon_means.idxmax(), "mean AUCell": eregulon_means.max().round(3)})