Search the spatial transcriptomics atlas¶
Retrieve related spatial transcriptomics spots using natural-language and H&E image queries.
Installation¶
Create a Python 3.10+ environment, install the public HistAgent package and add the lightweight analysis dependencies used by these notebooks.
git clone https://github.com/zipging/HistAgent.git
cd HistAgent
python -m venv .venv
source .venv/bin/activate
python -m pip install -e .
python -m pip install jupyter numpy pandas matplotlib seaborn scipy scikit-learn torch huggingface-hub
If data/ is absent, the notebook downloads its data from
wli13/HistAgent-data. These examples run without a GPU. Running HistAgent on
new images requires access to the gated Prov-GigaPath model and an authenticated
Hugging Face session.
Setup¶
Import the analysis packages, locate the tutorial data and set the plotting style.
from pathlib import Path
import json
import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
configured_data = os.getenv("HISTAGENT_TUTORIAL_DATA", "").strip()
DATA_DIR = Path(configured_data).expanduser() if configured_data else Path("data")
if not DATA_DIR.exists():
from huggingface_hub import snapshot_download
snapshot = Path(
snapshot_download(
"wli13/HistAgent-data",
repo_type="dataset",
allow_patterns=["tutorials/**", "manifest.json"],
)
)
DATA_DIR = snapshot / "tutorials"
assert DATA_DIR.exists(), f"Tutorial data not found: {DATA_DIR}"
def rankridge_pack_root(groups):
configured = os.getenv("HISTAGENT_RANKRIDGE_PACKS", "").strip()
if configured:
root = Path(configured).expanduser()
if not (root / groups[0]).exists():
root = root / "rankridge_tutorial_packs"
assert (root / groups[0]).exists(), f"RankRidge pack not found: {root}"
return root
from huggingface_hub import snapshot_download
snapshot = Path(
snapshot_download(
"wli13/HistAgent-data",
repo_type="dataset",
allow_patterns=[
f"rankridge_tutorial_packs/{group}/**" for group in groups
],
)
)
return snapshot / "rankridge_tutorial_packs"
sns.set_theme(style="white", context="talk")
plt.rcParams.update({
"figure.dpi": 135,
"figure.facecolor": "white",
"axes.facecolor": "white",
"axes.spines.top": False,
"axes.spines.right": False,
"axes.edgecolor": "#000000",
"axes.labelcolor": "#000000",
"axes.linewidth": 1.25,
"axes.grid": False,
"xtick.color": "#000000",
"ytick.color": "#000000",
"text.color": "#000000",
"font.size": 10,
"axes.titlesize": 12,
"axes.labelsize": 10.5,
"xtick.labelsize": 9,
"ytick.labelsize": 9,
"legend.fontsize": 9,
"figure.titlesize": 12,
})
COLORS = {
"HistAgent": "#168267",
"Measured ST": "#6f7477",
"STPath": "#cb77a0",
"OmiCLIP": "#69aec7",
}
SURVIVAL_COLORS = {"High risk": "#ef4545", "Low risk": "#45b3c2"}
REGION_COLORS = {
"Poor": "#d85f6b",
"High": "#79a9d1",
"Negative": "#b7b7b7",
"Tumor": "#e32925",
"Stroma": "#62acec",
}
TASK_CMAPS = {
"svg": "YlGnBu",
"domain": "tab10",
"deconv": "Greens",
"deg": "Purples",
"pathway": "PuBuGn",
}
Overview¶
Query a representative spatial transcriptomics atlas with biological descriptions and H&E images.
from sklearn.preprocessing import normalize
vectors = np.load(DATA_DIR / "figure5_embedding_subset.npz")
spots = pd.read_csv(DATA_DIR / "figure5_embedding_subset.csv")
queries = json.loads((DATA_DIR / "figure5_embedding_queries.json").read_text())
X = vectors["spot_pca"]
Q = vectors["query_pca"]
print(f"Atlas spots: {len(spots):,}")
print(f"Text queries: {len(queries['query_order'])}")
spots[["species", "organ", "cell_type"]].head()
Natural-language retrieval¶
The query retrieves ranked atlas spots and reports the tissue, cell type, top genes and pathway evidence used in the paper.
Xn = normalize(X)
Qn = normalize(Q)
retrievals = {}
for query_i, query_id in enumerate(queries["query_order"]):
scores = Xn @ Qn[query_i]
top = np.argsort(scores)[::-1][:8]
result = spots.iloc[top][
["spot_id", "species", "organ", "cell_type",
"top_genes", "reactome_pathway"]
].copy()
result.insert(1, "cosine_similarity", scores[top])
retrievals[query_id] = result
print("\nQUERY:", queries["queries"][query_id]["display_query"])
result_display = result[
["cosine_similarity", "organ", "cell_type",
"top_genes", "reactome_pathway"]
].head(5).copy()
result_display = result_display.rename(
columns={"cosine_similarity": "Cosine similarity"}
)
result_display.index = np.arange(1, len(result_display) + 1)
result_display.index.name = "Rank"
display(result_display.round({"Cosine similarity": 3}))
H&E image retrieval¶
The query is a human cortical H&E field. HistAgent reads a 55 μm spot view and a 220 μm context view, then retrieves related spatial atlas locations from the gene ranking.
image_query = queries["image_query"]
local_view = plt.imread(DATA_DIR / image_query["local_image_file"])
context_view = plt.imread(DATA_DIR / image_query["context_image_file"])
fig, axes = plt.subplots(1, 2, figsize=(7.6, 3.8), constrained_layout=True)
for ax, image, title in zip(
axes,
[local_view, context_view],
["55 μm spot", "220 μm context"],
):
ax.imshow(image)
ax.set_title(title)
ax.set_axis_off()
plt.show()
print("Top image-inferred genes:")
print(", ".join(image_query["predicted_genes"][:15]))
full_atlas_matches = pd.DataFrame(image_query["example_results"])
print("\nTop matches from the deployed measured-ST atlas:")
full_atlas_display = full_atlas_matches[
["slice_id", "cell_type", "matched_genes"]
].copy()
full_atlas_display.index = np.arange(1, len(full_atlas_display) + 1)
full_atlas_display.index.name = "Rank"
display(full_atlas_display)
Retrieve related atlas locations¶
The weighted rank-overlap calculation below searches the matching organ subset and returns the closest spatial locations.
def weighted_rank_overlap(query_genes, atlas_genes):
atlas_rank = {
gene.strip().upper(): rank
for rank, gene in enumerate(str(atlas_genes).split(","))
}
score = 0.0
matched = []
for query_rank, gene in enumerate(query_genes):
gene = gene.upper()
atlas_gene_rank = atlas_rank.get(gene)
if atlas_gene_rank is None:
continue
score += (
(1 - min(query_rank, 49) / 60)
* (1 - min(atlas_gene_rank, 49) / 60)
)
if len(matched) < 8:
matched.append(gene)
return score, ", ".join(matched)
image_species = image_query["species"]
image_organ = image_query["organ"]
image_candidates = spots.query(
"species == @image_species and organ == @image_organ"
).copy()
overlap = image_candidates["top_genes"].apply(
lambda genes: weighted_rank_overlap(image_query["predicted_genes"], genes)
)
image_candidates["rank_overlap"] = overlap.str[0]
image_candidates["matched_genes"] = overlap.str[1]
image_result = image_candidates.nlargest(8, "rank_overlap").copy()
image_result_display = image_result[
["cell_type", "matched_genes", "reactome_pathway"]
].head(5).copy()
image_result_display.index = np.arange(1, len(image_result_display) + 1)
image_result_display.index.name = "Rank"
display(image_result_display)
Verify the calculations¶
The checks confirm that every atlas vector has one metadata row and that the returned similarities are sorted.
assert len(spots) == len(X) == 2_250
assert not spots["spot_id"].duplicated().any()
for result in retrievals.values():
values = result["cosine_similarity"].to_numpy()
assert np.all(values[:-1] >= values[1:])
assert (DATA_DIR / image_query["local_image_file"]).exists()
assert (DATA_DIR / image_query["context_image_file"]).exists()
assert len(image_query["predicted_genes"]) == 50
assert full_atlas_matches["cell_type"].notna().all()
assert image_result["rank_overlap"].is_monotonic_decreasing
print("Checks passed for 2,250 spots, text retrieval and H&E retrieval.")
Continue in the Agent Module¶
Open a retrieved atlas location in Spot Chat to inspect its molecular profile and ask follow-up questions.