import os
import warnings
from pathlib import Path

from tqdm import TqdmWarning

warnings.filterwarnings("ignore", category=TqdmWarning)

import liana as li
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import plotnine as p9
import scanpy as sc
import squidpy as sq
from liana.method import cellphonedb
import lamindb as ln

sc.settings.verbosity = 1
sc.set_figure_params(dpi=80, dpi_save=150, frameon=False, figsize=(5, 5))
plt.rcParams["figure.dpi"] = 80
plt.rcParams["font.family"] = "DejaVu Sans"  # Arial lacks some of the glyphs LIANA+ plots use

from importlib.metadata import version

for pkg in ("scanpy", "anndata", "squidpy", "liana", "numpy", "pandas"):
    print(f"{pkg:>8}: {version(pkg)}")

  scanpy: 1.12.4
 anndata: 0.13.2
 squidpy: 1.8.3
   liana: 2.0.1.dev1+g08590dbd4
   numpy: 2.5.3
  pandas: 3.0.6

!lamin load https://lamin.ai/theislab/sc-best-practices/artifact/rYFrnTsJZLfB200S0002

→ connected lamindb: theislab/sc-best-practices
! connected in read-only mode, please use ln.DB("theislab/sc-best-practices") instead
... synchronizing ovarian_xenium_annotated_rendered.h5ad: 100.0%
→ artifact is here: /home/jovyan/.cache/lamindb/lamin-eu-central-1/VPwcjx3CDAa2/spatial/ccc/ovarian_xenium_annotated_rendered.h5ad

adata = sc.read_h5ad("/home/jovyan/.cache/lamindb/lamin-eu-central-1/VPwcjx3CDAa2/spatial/ccc/ovarian_xenium_annotated_rendered.h5ad")
adata

AnnData object with n_obs × n_vars = 400600 × 3197
    obs: 'celltype', 'sample_id', 'sliding_window', 'tissue_region', 'local_programme', 'global_programme'
    uns: 'bivariate_morans', 'celltype_colors', 'celltype_nhood_enrichment', 'cpdb_hard', 'cpdb_plain', 'cpdb_soft_bw100', 'cpdb_soft_bw30', 'cpdb_soft_bw300', 'cross_pcf', 'global_interactions', 'inflow_features', 'interscale_dim_importance', 'interscale_gene_ranks', 'interscale_net_flow_mean', 'interscale_net_flow_sd', 'interscale_net_streams', 'interscale_programme_genes', 'interscale_programmes', 'lric_ct', 'moranI', 'rendered_note'
    obsm: 'inflow_scores', 'spatial'
    layers: None (.X)

CT_KEY = "celltype"
SPATIAL_KEY = "spatial"

# The running example, threaded through Parts 1-3.
EXAMPLE_SOURCE = "Tumor Associated Fibroblasts"   # sender cell type
EXAMPLE_TARGET = "Macrophages"                    # receiver cell type
EXAMPLE_LIGAND = "THBS1"
EXAMPLE_RECEPTOR = "LRP1"
EXAMPLE_INTERACTION = f"{EXAMPLE_SOURCE}^{EXAMPLE_LIGAND}^{EXAMPLE_RECEPTOR}"
print("running example:", EXAMPLE_INTERACTION, "received by", EXAMPLE_TARGET)

print(f"{adata.n_obs:,} cells x {adata.n_vars:,} genes")
# `adata.layers` iterates a `None` key aliasing `.X` in anndata 0.13, so filter it out here
print("\nlayers:", [k for k in adata.layers if k is not None])
print("obsm:", list(adata.obsm.keys()))
print("obs:", list(adata.obs.columns))
print("sample_id levels:", adata.obs["sample_id"].unique().tolist(), "-> single section, no library_key needed")

# `.X` holds log-normalised values. The original object also carried the raw `counts` layer and
# two further layers ("norm", "log1p_norm") that were byte-identical copies of `.X`; all three
# were dropped here, because together they were 7 of the object's 9.5 GB and the analyses below
# read only `.X`. InterScale (Part 4) wants the data under the name `log1p_norm`, which we alias
# back onto `.X` there at no memory cost.
print(f"\n.X min={adata.X.min():.2f} max={adata.X.max():.2f} "
      f"integer-valued={np.allclose(adata.X[:1000].data, np.round(adata.X[:1000].data))}")
if "rendered_note" in adata.uns:
    print("\n" + adata.uns["rendered_note"])

print("\ncells per type:")
print(adata.obs[CT_KEY].value_counts().to_string())

running example: Tumor Associated Fibroblasts^THBS1^LRP1 received by Macrophages
400,600 cells x 3,197 genes

layers: []
obsm: ['inflow_scores', 'spatial']
obs: ['celltype', 'sample_id', 'sliding_window', 'tissue_region', 'local_programme', 'global_programme']
sample_id levels: ['sample'] -> single section, no library_key needed

.X min=0.00 max=5.45 integer-valued=False

Pre-rendered companion to practical 9 (p9_ccc.ipynb). .X is log1p-normalised; the raw 'counts' layer and the byte-identical 'log1p_norm'/'norm' copies of .X were dropped, taking the object from 9.5 GB to under 1 GB. Every expensive result is stored in .uns/.obsm so the notebook's guards skip recomputation - measured peaks on the full object were 43.7 GB for li.mt.cross_pcf, 35.7 GB for li.mt.inflow and over 120 GB for InterScale's get_model_output, none of which fit the 32 GB teaching server. 'moranI' spans the original 4,447-gene panel; .var is the 3,197 genes that passed the spatially-variable-gene filter.

cells per type:
celltype
Tumor Cells                     218172
Tumor Associated Fibroblasts     71984
Macrophages                      57876
Malignant Cells Lining Cyst      20367
T & NK Cells                     19814
Endothelial Cells                 7135
Pericytes                         4443
Granulosa Cells                    809

sq.pl.spatial_scatter(
    adata,
    shape=None,             
    color=CT_KEY,
    library_id='sample',
    img=False,
    size=0.12,                 # tune; the original s=0.12 is not directly equivalent
    figsize=(7.5, 8),
    title="Cell types in the ovarian adenocarcinoma section",
    legend_fontsize=8,
    legend_loc="right margin",
    frameon=True,
    axis_label=["x (µm)", "y (µm)"],
)

/app/.pixi/envs/default/lib/python3.12/site-packages/squidpy/pl/_spatial_utils.py:468: FutureWarning: The method obs_vector is deprecated and will be removed in the future. Use anndata.acc.A instead of obs_vector. E.g. `vec = adata[A.obs['foo']]` or `vec = adata[A.layers['l']['bar', :]]`
/app/.pixi/envs/default/lib/python3.12/site-packages/squidpy/pl/_spatial_utils.py:979: UserWarning: No data for colormapping provided via 'c'. Parameters 'cmap', 'norm' will be ignored

%%time
# Delaunay triangulation of the cell centroids, unpruned, to look at the edge lengths
sq.gr.spatial_neighbors_delaunay(adata, spatial_key=SPATIAL_KEY, key_added="contact")
edges = np.asarray(adata.obsp["contact_distances"].data)

# the distance to the single nearest neighbour is NOT the quantity that prunes edges
from sklearn.neighbors import NearestNeighbors

coords = np.asarray(adata.obsm[SPATIAL_KEY])
nn_dist = NearestNeighbors(n_neighbors=2).fit(coords).kneighbors(coords)[0][:, 1]

print(f"nearest-neighbour distance: median {np.median(nn_dist):6.1f} µm, 95th percentile {np.quantile(nn_dist, 0.95):6.1f} µm")
print(f"Delaunay edge length:       median {np.median(edges):6.1f} µm, 90th percentile {np.quantile(edges, 0.90):6.1f} µm")
print(f"mean degree of the unpruned graph: {np.asarray((adata.obsp['contact_connectivities'] > 0).sum(1)).mean():.2f}")

INFO     Creating graph using `None` transform and `1` libraries.                                                  
nearest-neighbour distance: median    8.2 µm, 95th percentile   12.3 µm
Delaunay edge length:       median   12.3 µm, 90th percentile   20.7 µm
mean degree of the unpruned graph: 6.00
CPU times: user 4.07 s, sys: 89.2 ms, total: 4.15 s
Wall time: 4.19 s

CONTACT_CUTOFF = 20.0  # µm, one cell diameter

fig, ax = plt.subplots(figsize=(6.5, 3.6))
ax.hist(edges, bins=np.arange(0, 80, 1), color="#999999", label="Delaunay edge lengths")
ax.hist(nn_dist, bins=np.arange(0, 80, 1), histtype="step", color="#4C72B0", lw=1.5, label="nearest-neighbour distances")
ax.axvline(CONTACT_CUTOFF, color="#C44E52", lw=2, ls="--", label=f"contact cutoff = {CONTACT_CUTOFF:.0f} µm")
ax.set_xlabel("distance (µm)")
ax.set_ylabel("count")
ax.set_title("Where to cut the contact graph")
ax.legend(frameon=False, fontsize=9)
plt.show()

print(f"{(edges <= CONTACT_CUTOFF).mean():.1%} of Delaunay edges are shorter than {CONTACT_CUTOFF:.0f} µm")

88.7% of Delaunay edges are shorter than 20 µm

%%time
# the pruned contact graph: a Delaunay edge is kept only if it is shorter than the cutoff.
# NOTE: `radius` has to be a tuple (min, max) - a list is silently ignored.
sq.gr.spatial_neighbors_delaunay(
    adata, spatial_key=SPATIAL_KEY, radius=(0.0, CONTACT_CUTOFF), key_added="contact"
)
degree = np.asarray((adata.obsp["contact_connectivities"] > 0).sum(1)).ravel()
print(f"contact graph: mean degree {degree.mean():.2f}, {(degree == 0).sum():,} isolated cells")

INFO     Creating graph using `None` transform and `1` libraries.                                                  
contact graph: mean degree 5.32, 1,318 isolated cells
CPU times: user 3.07 s, sys: 30.1 ms, total: 3.1 s
Wall time: 3.12 s

%%time
bw_plot, bw_df = li.pp.query_bandwidth(coordinates=coords, start=5, end=50, interval_n=40)
display(bw_df.head(10))
bw_plot

CPU times: user 28.7 s, sys: 22.2 ms, total: 28.7 s
Wall time: 28.9 s

%%time
BANDWIDTH = 20.0  # µm

li.pp.spatial_neighbors(
    adata,
    bandwidth=BANDWIDTH,
    cutoff=0.1,
    kernel="gaussian",
    spatial_key=SPATIAL_KEY,
    set_diag=False,
    verbose=True,
)
print("kernel graph in adata.obsp['spatial_connectivities']:", adata.obsp["spatial_connectivities"].shape)
print(f"mean number of non-zero weights per cell: {adata.obsp['spatial_connectivities'].getnnz(axis=1).mean():.1f}")

kernel graph in adata.obsp['spatial_connectivities']: (400600, 400600)
mean number of non-zero weights per cell: 101.0
CPU times: user 18 s, sys: 400 ms, total: 18.4 s
Wall time: 18.6 s

li.pl.connectivity(adata, idx=200_000, size=1.5, figure_size=(5.5, 4.5))

/tmp/ipykernel_6332/3997700466.py:1: UserWarning: This function will be deprecated in the next version. Please use scanpy or squidpy for plotting spatial connectivities.

%%time
# Moran's I for every gene; this takes a while, so it is skipped if it is already stored in adata.uns
if "moranI" not in adata.uns:
    sq.gr.spatial_autocorr(
        adata,
        mode="moran",
        genes=adata.var_names,          # by default squidpy would only test highly variable genes
        connectivity_key="spatial_connectivities",
        use_raw=False,
        n_perms=None,                   # the normal approximation is enough for a filter
        n_jobs=16,
        show_progress_bar=False,
    )
else:
    print("using the stored Moran's I (adata.uns['moranI'])")

using the stored Moran's I (adata.uns['moranI'])
CPU times: user 89 μs, sys: 30 μs, total: 119 μs
Wall time: 117 μs

moran = adata.uns["moranI"]
svg = moran.index[(moran["I"] > 0.01) & (moran["pval_norm_fdr_bh"] < 0.05)]

fig, ax = plt.subplots(figsize=(6, 3.5))
ax.plot(np.arange(1, len(moran) + 1), moran["I"].values, color="#27828E")
ax.axvline(len(svg), color="#EE9B00", ls="--", label=f"top {len(svg)}")
ax.set(xlabel="gene rank", ylabel="Moran's I", yscale="log", title="Spatial autocorrelation per gene")
ax.legend()
plt.tight_layout()
plt.show()

# `moran` spans the full panel the statistic was computed on, which is larger than the gene set
# carried by this object - the spatially variable genes were already selected when it was built.
print(f"\n{len(svg):,} of {len(moran):,} genes are spatially variable")
display(moran.head())

3,197 of 4,447 genes are spatially variable

keep = [g for g in adata.var_names if g in set(svg)]
if len(keep) < adata.n_vars:
    adata = adata[:, keep].copy()
else:
    # this object was already restricted to the spatially variable genes when it was built, so
    # the subset would copy `.X` for nothing - roughly 2.4 GB we would rather not spend
    print("already restricted to the spatially variable genes - nothing to subset")
print("adata after the SVG filter:", adata.shape)

already restricted to the spatially variable genes - nothing to subset
adata after the SVG filter: (400600, 3197)

resource = li.rs.select_resource("consensus")
genes = set(adata.var_names)
resource = resource[resource["ligand"].isin(genes) & resource["receptor"].isin(genes)].reset_index(drop=True)
print(f"consensus resource restricted to the panel and the SVG filter: {len(resource):,} ligand-receptor pairs")
display(resource.head())

consensus resource restricted to the panel and the SVG filter: 715 ligand-receptor pairs

%%time
if f"{CT_KEY}_nhood_enrichment" not in adata.uns:
    sq.gr.nhood_enrichment(
        adata, cluster_key=CT_KEY, connectivity_key="contact", n_perms=1000, seed=0, show_progress_bar=False
    )
else:
    print("using the stored neighbourhood enrichment")

CT_ORDER = np.unique(adata.obs[CT_KEY].values)
zscores = pd.DataFrame(
    adata.uns[f"{CT_KEY}_nhood_enrichment"]["zscore"], index=CT_ORDER, columns=CT_ORDER
)

fig, ax = plt.subplots(figsize=(6.2, 5.2))
sq.pl.nhood_enrichment(
    adata, cluster_key=CT_KEY, cmap="RdBu_r", vmin=-300, vmax=300, ax=ax, title="Neighbourhood enrichment (z-score)"
)
plt.show()

using the stored neighbourhood enrichment

CPU times: user 514 ms, sys: 32.7 ms, total: 547 ms
Wall time: 553 ms

nea = (
    zscores.reset_index(names="source")
    .melt(id_vars="source", var_name="target", value_name="zscore")
)
print("most co-localised cell type pairs (heterotypic only):")
display(nea.query("source != target").sort_values("zscore", ascending=False).head(8).reset_index(drop=True))
print(f"\nrunning example, {EXAMPLE_SOURCE} <-> {EXAMPLE_TARGET}: z = {zscores.loc[EXAMPLE_SOURCE, EXAMPLE_TARGET]:.1f}")
print(f"for contrast,   {EXAMPLE_SOURCE} <-> Tumor Cells: z = {zscores.loc[EXAMPLE_SOURCE, 'Tumor Cells']:.1f}")

most co-localised cell type pairs (heterotypic only):

running example, Tumor Associated Fibroblasts <-> Macrophages: z = 12.8
for contrast,   Tumor Associated Fibroblasts <-> Tumor Cells: z = -494.4

li.pl.annulus(
    adata, spatial_key=SPATIAL_KEY, radius_step=20, n_rings=8, seed=1337,
    figure_size=(7, 7), return_fig=False,
)

%%time
# The cross pair-correlation function counts neighbours in annuli out to 250 µm for every one of
# the 400,600 cells. That is by far the heaviest step of this chapter and needs well over 32 GB,
# so the result is stored in the prepared object and reused here.
if "cross_pcf" not in adata.uns:
    li.mt.cross_pcf(
        adata,
        groupby=CT_KEY,
        spatial_key=SPATIAL_KEY,
        max_radius=250,
        radius_step=20,
        min_cells=100,
        key_added="cross_pcf",
        verbose=True,
    )
else:
    print("using the stored cross pair-correlation function")

pcf = adata.uns["cross_pcf"]
print("g(r) per radius (the first annulus is extended, so the grid starts 0, 40, 60, ...):")
display(pcf.head())
print("radii:", sorted(pcf["radius"].unique()))

using the stored cross pair-correlation function
g(r) per radius (the first annulus is extended, so the grid starts 0, 40, 60, ...):

radii: [np.float64(0.0), np.float64(40.0), np.float64(60.0), np.float64(80.0), np.float64(100.0), np.float64(120.0), np.float64(140.0), np.float64(160.0), np.float64(180.0), np.float64(200.0), np.float64(220.0), np.float64(240.0), np.float64(260.0)]
CPU times: user 6.83 ms, sys: 3.61 ms, total: 10.4 ms
Wall time: 9.3 ms

auc = li.mt.get_lric_auc(adata, "cross_pcf", max_dist=50, min_bins=2)
print("summary of each curve: mean log2 g(r) up to 50 µm, and the radius at which g(r) peaks")
display(auc.sort_values("score", ascending=False).head(8).reset_index(drop=True))
display(auc.sort_values("score").head(4).reset_index(drop=True))

summary of each curve: mean log2 g(r) up to 50 µm, and the radius at which g(r) peaks

def gcurve(a, b):
    """g(r) for one unordered cell type pair (cross_pcf stores each pair once)."""
    sel = pcf[((pcf["source"] == a) & (pcf["target"] == b)) | ((pcf["source"] == b) & (pcf["target"] == a))]
    return sel.sort_values("radius")

PAIRS = [
    ("Endothelial Cells", "Pericytes", "#55A868", "aggregation"),
    (EXAMPLE_SOURCE, EXAMPLE_TARGET, "#4C72B0", "running example"),
    (EXAMPLE_SOURCE, "Tumor Cells", "#C44E52", "exclusion"),
]
fig, ax = plt.subplots(figsize=(6.5, 4))
for a, b, colour, label in PAIRS:
    g = gcurve(a, b)
    ax.plot(g["radius"], g["g"], "-o", ms=3, color=colour, label=f"{a} ^ {b} ({label})")
ax.axhline(1, ls="--", c="k", lw=1)
ax.set_xlabel("radius r (µm)")
ax.set_ylabel("g(r)")
ax.set_yscale("log")
ax.set_title("Cross pair-correlation function")
ax.legend(frameon=False, fontsize=8)
plt.show()

g_run = gcurve(EXAMPLE_SOURCE, EXAMPLE_TARGET)
print(f"{EXAMPLE_SOURCE} ^ {EXAMPLE_TARGET}: g(r) ranges from {g_run['g'].min():.2f} to {g_run['g'].max():.2f} over 0-260 µm")
print(f"Endothelial Cells ^ Pericytes: g(0-40 µm) = {gcurve('Endothelial Cells', 'Pericytes')['g'].iloc[0]:.2f}")
print(f"{EXAMPLE_SOURCE} ^ Tumor Cells: g(0-40 µm) = {gcurve(EXAMPLE_SOURCE, 'Tumor Cells')['g'].iloc[0]:.2f}")

Tumor Associated Fibroblasts ^ Macrophages: g(r) ranges from 0.94 to 0.96 over 0-260 µm
Endothelial Cells ^ Pericytes: g(0-40 µm) = 6.79
Tumor Associated Fibroblasts ^ Tumor Cells: g(0-40 µm) = 0.31

%%time
if "lric_ct" not in adata.uns:
    lric_resource = resource.query("ligand == @EXAMPLE_LIGAND and receptor == @EXAMPLE_RECEPTOR")
    li.mt.lric(
        adata,
        resource=lric_resource,
        groupby=CT_KEY,
        spatial_key=SPATIAL_KEY,
        cell_types=[EXAMPLE_SOURCE, EXAMPLE_TARGET],
        max_radius=250,
        radius_step=20,
        expr_prop=0.1,
        use_raw=False,
        key_added="lric_ct",
        verbose=True,
    )
else:
    print("using the stored LRIC result")

display(li.mt.get_lric_auc(adata, "lric_ct", max_dist=50, min_bins=2))
display(adata.uns["lric_ct"].query("source == @EXAMPLE_SOURCE").head())

using the stored LRIC result

CPU times: user 22.2 ms, sys: 6.02 ms, total: 28.2 ms
Wall time: 26 ms

li.pl.lric_lineplot(
    adata,
    "lric_ct",
    source=EXAMPLE_SOURCE,
    target=EXAMPLE_TARGET,
    interaction=f"{EXAMPLE_LIGAND}^{EXAMPLE_RECEPTOR}",
    figure_size=(6.5, 4),
)

%%time
if "cpdb_plain" not in adata.uns:
    cellphonedb(
        adata,
        groupby=CT_KEY,
        resource_name="consensus",
        expr_prop=0.05,
        use_raw=False,
        n_perms=1000,
        seed=1337,
        n_jobs=1,
        verbose=False,   # the 1,000-permutation progress bar would fill the output
        key_added="cpdb_plain",
    )
else:
    print("using the stored unconstrained CellPhoneDB result")

plain = adata.uns["cpdb_plain"]
print(f"\n{len(plain):,} rows: every sender x receiver x ligand-receptor combination that passed the expression filter")
display(plain.sort_values(["lr_means", "cellphone_pvals"], ascending=[False, True]).head(8).reset_index(drop=True))

using the stored unconstrained CellPhoneDB result

22,774 rows: every sender x receiver x ligand-receptor combination that passed the expression filter

CPU times: user 16.5 ms, sys: 3.84 ms, total: 20.3 ms
Wall time: 19.3 ms

KEYS = ["source", "target", "ligand_complex", "receptor_complex"]


def ranked(df, name):
    """Rank every row of a CellPhoneDB result by score, then p-value."""
    out = df.sort_values(["lr_means", "cellphone_pvals"], ascending=[False, True]).reset_index(drop=True)
    out[name] = np.arange(1, len(out) + 1)
    return out[KEYS + [name, "lr_means", "cellphone_pvals"]].rename(
        columns={"lr_means": f"lr_means_{name}", "cellphone_pvals": f"pval_{name}"}
    )


def running(df, name):
    """The running example's row in a ranked table."""
    sel = df.query("source == @EXAMPLE_SOURCE and target == @EXAMPLE_TARGET and ligand_complex == @EXAMPLE_LIGAND and receptor_complex == @EXAMPLE_RECEPTOR")
    return sel.iloc[0] if len(sel) else None


plain_r = ranked(plain, "plain")
row = running(plain_r, "plain")
print(f"running example {EXAMPLE_SOURCE} -> {EXAMPLE_TARGET}, {EXAMPLE_LIGAND}^{EXAMPLE_RECEPTOR}")
print(f"  unconstrained: rank {int(row['plain']):,} of {len(plain_r):,}, score {row['lr_means_plain']:.2f}, p = {row['pval_plain']:.3f}")

running example Tumor Associated Fibroblasts -> Macrophages, THBS1^LRP1
  unconstrained: rank 124 of 22,774, score 1.45, p = 0.000

%%time
groupby_pairs = nea.query("zscore > 2")[["source", "target"]].reset_index(drop=True)
print(f"the mask keeps {len(groupby_pairs)} of {len(nea)} cell type pairs")
display(groupby_pairs.head(10))

if "cpdb_hard" not in adata.uns:
    cellphonedb(
        adata,
        groupby=CT_KEY,
        groupby_pairs=groupby_pairs,
        resource_name="consensus",
        expr_prop=0.05,
        use_raw=False,
        n_perms=1000,
        seed=1337,
        n_jobs=1,
        verbose=False,
        key_added="cpdb_hard",
    )
else:
    print("using the stored hard-masked CellPhoneDB result")

hard = adata.uns["cpdb_hard"]
hard_r = ranked(hard, "hard")
row = running(hard_r, "hard")
print(f"\nhard-masked result: {len(hard):,} rows ({len(hard) / len(plain):.0%} of the unconstrained table)")
print(f"running example: rank {int(row['hard']):,} of {len(hard_r):,}")

the mask keeps 22 of 64 cell type pairs

using the stored hard-masked CellPhoneDB result

hard-masked result: 8,326 rows (37% of the unconstrained table)
running example: rank 42 of 8,326
CPU times: user 14.2 ms, sys: 3.89 ms, total: 18.1 ms
Wall time: 16.6 ms

# what the mask costs: an interaction between two cell types that avoid each other disappears
masked_out = {tuple(x) for x in groupby_pairs.to_numpy()}
print(f"('Tumor Cells', 'Macrophages') in the mask: {('Tumor Cells', 'Macrophages') in masked_out}"
      f"  (NEA z = {zscores.loc['Tumor Cells', 'Macrophages']:.1f})")
q = "source == 'Tumor Cells' and target == 'Macrophages' and ligand_complex == 'MFGE8' and receptor_complex == 'ITGB3'"
print("rows for Tumor Cells -> Macrophages MFGE8^ITGB3:  unconstrained", len(plain.query(q)), " hard-masked", len(hard.query(q)))

('Tumor Cells', 'Macrophages') in the mask: False  (NEA z = -172.1)
rows for Tumor Cells -> Macrophages MFGE8^ITGB3:  unconstrained 1  hard-masked 0

%%time
BANDWIDTHS = [30, 100, 300]
proximity = {}
for bw in BANDWIDTHS:
    proximity[bw] = li.pp.spatial_pair_proximity(
        adata, groupby=CT_KEY, spatial_key=SPATIAL_KEY, kernel="gaussian", bandwidth=bw, verbose=False
    )
display(proximity[30].head())

fig, axes = plt.subplots(1, 2, figsize=(11.5, 4.4))
for ax, bw in zip(axes, [30, 300]):
    mat = proximity[bw].pivot(index="source", columns="target", values="proximity").loc[CT_ORDER, CT_ORDER]
    im = ax.imshow(mat.to_numpy(), cmap="magma", vmin=0, vmax=1)
    ax.set_xticks(range(len(CT_ORDER)), CT_ORDER, rotation=90, fontsize=7)
    ax.set_yticks(range(len(CT_ORDER)), CT_ORDER, fontsize=7)
    ax.set_title(f"proximity weight, bandwidth = {bw} µm")
    fig.colorbar(im, ax=ax, shrink=0.8)
plt.tight_layout()
plt.show()

CPU times: user 1min 4s, sys: 578 ms, total: 1min 4s
Wall time: 25.2 s

%%time
softs = {}
for bw in BANDWIDTHS:
    key = f"cpdb_soft_bw{bw}"
    if key not in adata.uns:
        cellphonedb(
            adata,
            groupby=CT_KEY,
            resource_name="consensus",
            expr_prop=0.05,
            use_raw=False,
            n_perms=1000,
            seed=1337,
            n_jobs=1,
            verbose=False,
            key_added=key,
            spatial_key=SPATIAL_KEY,
            spatial_kwargs={"kernel": "gaussian", "bandwidth": bw},
        )
    softs[bw] = adata.uns[key]
    print(f"bandwidth {bw:>3} µm: {len(softs[bw]):,} rows")

bandwidth  30 µm: 22,774 rows
bandwidth 100 µm: 22,774 rows
bandwidth 300 µm: 22,774 rows
CPU times: user 0 ns, sys: 180 μs, total: 180 μs
Wall time: 179 μs

def top_global(df, n=10):
    sel = df.sort_values(["lr_means", "cellphone_pvals"], ascending=[False, True]).head(n)
    return (sel["source"] + " -> " + sel["target"] + ": " + sel["ligand_complex"] + "^" + sel["receptor_complex"]).tolist()

display(
    pd.DataFrame(
        {"unconstrained": top_global(plain), **{f"soft {bw} µm": top_global(softs[bw]) for bw in BANDWIDTHS}},
        index=pd.RangeIndex(1, 11, name="rank"),
    )
)

comparison = plain_r.merge(hard_r, on=KEYS, how="left")
for bw in BANDWIDTHS:
    comparison = comparison.merge(ranked(softs[bw], f"soft{bw}"), on=KEYS, how="left")

rank_cols = ["plain", "hard"] + [f"soft{bw}" for bw in BANDWIDTHS]
summary = pd.DataFrame(
    {
        "rows in the table": [len(plain_r), len(hard_r)] + [len(softs[bw]) for bw in BANDWIDTHS],
        "cell type pairs kept": [df.groupby(["source", "target"], observed=True).ngroups for df in [plain, hard] + [softs[bw] for bw in BANDWIDTHS]],
        f"rank of {EXAMPLE_LIGAND}^{EXAMPLE_RECEPTOR} ({EXAMPLE_SOURCE} -> {EXAMPLE_TARGET})": [running(comparison, c)[c] for c in rank_cols],
    },
    index=["unconstrained", "hard mask (NEA z > 2)"] + [f"soft, bandwidth {bw} µm" for bw in BANDWIDTHS],
)
display(summary)

li.pl.dotplot(
    adata=adata,
    uns_key="cpdb_soft_bw100",
    colour="lr_means",
    size="cellphone_pvals",
    inverse_size=True,
    source_labels=[EXAMPLE_SOURCE, "Macrophages", "Tumor Cells"],
    target_labels=[EXAMPLE_TARGET, "Tumor Cells", "T & NK Cells"],
    top_n=15,
    orderby="lr_means",
    orderby_ascending=False,
    figure_size=(13, 9),
) + p9.theme(plot_margin=0.02)

%%time
if "bivariate_morans" not in adata.uns:
    biv = li.mt.bivariate(
        adata,
        local_name=None,
        global_name="morans",
        resource_name="consensus",
        connectivity_key="spatial_connectivities",
        n_perms=None,          # permutation p-values would dominate the runtime of this chapter
        nz_prop=0.01,
        use_raw=False,
        verbose=True,
    )
    # li.mt.bivariate returns its result instead of writing it to adata, so store it ourselves
    adata.uns["bivariate_morans"] = biv if isinstance(biv, pd.DataFrame) else biv.var.copy()
else:
    print("using the stored global bivariate Moran's I")

bivar = adata.uns["bivariate_morans"]
top_biv = bivar.sort_values("morans", ascending=False).head(10)
# n_perms=None means no permutation p-values are computed, so we do not show that column
display(top_biv[["ligand", "receptor", "ligand_props", "receptor_props", "morans"]].reset_index(drop=True))
print(f"\nrank of {EXAMPLE_LIGAND}^{EXAMPLE_RECEPTOR}: {int((bivar['morans'] > bivar.query('ligand == @EXAMPLE_LIGAND and receptor == @EXAMPLE_RECEPTOR')['morans'].iloc[0]).sum()) + 1} of {len(bivar):,}")

using the stored global bivariate Moran's I

rank of THBS1^LRP1: 1 of 710
CPU times: user 13 ms, sys: 1.12 ms, total: 14.1 ms
Wall time: 12 ms

fig, axes = plt.subplots(1, 2, figsize=(13, 6))
for ax, gene in zip(axes, [EXAMPLE_LIGAND, EXAMPLE_RECEPTOR]):
    v = np.asarray(adata[:, gene].X.todense()).ravel()
    o = np.argsort(v)
    s = ax.scatter(coords[o, 0], coords[o, 1], c=v[o], s=0.15, cmap="magma", linewidths=0, rasterized=True)
    ax.set_aspect("equal")
    ax.invert_yaxis()
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_title(f"{gene} (log-normalised)")
    fig.colorbar(s, ax=ax, shrink=0.6)
plt.tight_layout()
plt.show()

# a contrast: the same ligand with a different receptor, one that is expressed almost everywhere
CONTRAST_RECEPTOR = "CD47"
rows = []
for name, df in [("plain", plain), ("hard", hard)] + [(f"soft{bw}", softs[bw]) for bw in BANDWIDTHS]:
    r = ranked(df, name)
    for rec in (EXAMPLE_RECEPTOR, CONTRAST_RECEPTOR):
        sel = r.query("source == @EXAMPLE_SOURCE and target == @EXAMPLE_TARGET and ligand_complex == @EXAMPLE_LIGAND and receptor_complex == @rec")
        rows.append({"table": name, "receptor": rec, "rank": int(sel[name].iloc[0]) if len(sel) else np.nan})
contrast = pd.DataFrame(rows).pivot(index="receptor", columns="table", values="rank")
biv_small = bivar.query("ligand == @EXAMPLE_LIGAND and receptor in [@EXAMPLE_RECEPTOR, @CONTRAST_RECEPTOR]").set_index("receptor")
contrast["global Moran's I"] = biv_small["morans"].round(3)
contrast["receptor detected in"] = biv_small["receptor_props"].map("{:.0%}".format)
display(contrast)

%%time
# inflow scores every cell against every sender^ligand^receptor combination, so the object it
# returns is cells x (senders x LR pairs) - the second memory peak of this chapter. The prepared
# object therefore carries the per-cell scores for the two features plotted below, and we rebuild
# just that slice when they are present.
if "inflow_scores" not in adata.obsm:
    lrdata = li.mt.inflow(
        adata,
        groupby=CT_KEY,
        resource=resource,
        connectivity_key="spatial_connectivities",
        nz_prop=0.01,
        complex_sep=None,
        use_raw=False,
        verbose=True,
    )
    print("\ninflow returns a single AnnData: cells x sender^ligand^receptor")
else:
    import anndata as ad
    from scipy.sparse import csr_matrix

    stored = adata.obsm["inflow_scores"]
    lrdata = ad.AnnData(
        X=csr_matrix(stored.to_numpy(dtype=np.float32)),
        obs=adata.obs[[CT_KEY]].copy(),
        var=pd.DataFrame(index=pd.Index(stored.columns, name=None)),
    )
    lrdata.obsm[SPATIAL_KEY] = adata.obsm[SPATIAL_KEY]
    if "global_interactions" in adata.uns:
        lrdata.uns["global_interactions"] = adata.uns["global_interactions"]
    print("using the stored inflow scores for:", ", ".join(stored.columns))

print(lrdata)

using the stored inflow scores for: Tumor Associated Fibroblasts^THBS1^LRP1, Tumor Cells^MFGE8^ITGB3
AnnData object with n_obs × n_vars = 400600 × 2
    obs: 'celltype'
    uns: 'global_interactions'
    obsm: 'spatial'
    layers: None (.X)
CPU times: user 86.4 ms, sys: 16.8 ms, total: 103 ms
Wall time: 102 ms

FEATURE = EXAMPLE_INTERACTION
scores = np.asarray(lrdata[:, FEATURE].X.todense()).ravel()

fig, axes = plt.subplots(1, 2, figsize=(13, 6))
o = np.argsort(scores)
s = axes[0].scatter(coords[o, 0], coords[o, 1], c=scores[o], s=0.15, cmap="magma", linewidths=0, rasterized=True)
axes[0].set_title(f"inflow: {FEATURE}")
fig.colorbar(s, ax=axes[0], shrink=0.6)

aside = "Tumor Cells^MFGE8^ITGB3"
scores2 = np.asarray(lrdata[:, aside].X.todense()).ravel()
o2 = np.argsort(scores2)
s2 = axes[1].scatter(coords[o2, 0], coords[o2, 1], c=scores2[o2], s=0.15, cmap="magma", linewidths=0, rasterized=True)
axes[1].set_title(f"inflow: {aside} (removed by the hard mask)")
fig.colorbar(s2, ax=axes[1], shrink=0.6)

for ax in axes:
    ax.set_aspect("equal")
    ax.invert_yaxis()
    ax.set_xticks([])
    ax.set_yticks([])
plt.tight_layout()
plt.show()

print(f"{FEATURE}: score > 0 in {(scores > 0).mean():.1%} of cells, mean {scores.mean():.3f}")
print(f"{aside}: score > 0 in {(scores2 > 0).mean():.1%} of cells, mean {scores2.mean():.3f}")

Tumor Associated Fibroblasts^THBS1^LRP1: score > 0 in 30.9% of cells, mean 0.347
Tumor Cells^MFGE8^ITGB3: score > 0 in 44.8% of cells, mean 0.390

# which cell types actually receive the signal?
receivers = (
    pd.DataFrame({CT_KEY: adata.obs[CT_KEY].to_numpy(), "inflow": scores})
    .groupby(CT_KEY, observed=True)["inflow"]
    .agg(mean_score="mean", prop_nonzero=lambda s: (s > 0).mean())
    .sort_values("mean_score", ascending=False)
)
display(receivers)

fig, ax = li.pl.feature_by_group(
    lrdata,
    groupby=CT_KEY,
    feature=FEATURE,
    labels=["Tumor Cells", EXAMPLE_SOURCE, EXAMPLE_TARGET],
    figure_size=(12, 4.2),
)
plt.show()

%%time
if "global_interactions" not in lrdata.uns:
    li.mt.compute_global_specificity(lrdata, groupby=CT_KEY, n_perms=100, seed=1337, verbose=False, use_raw=False)
else:
    print("using the stored global interactions")

gi = lrdata.uns["global_interactions"]
display(
    gi.query("source == @EXAMPLE_SOURCE and target == @EXAMPLE_TARGET")
    .sort_values("lr_mean", ascending=False)
    .head(8)
    .reset_index(drop=True)
)
li.pl.dotplot(
    liana_res=gi,
    colour="lr_mean",
    size="pval",
    inverse_size=True,
    source_labels=[EXAMPLE_SOURCE, "Tumor Cells", "Macrophages"],
    target_labels=[EXAMPLE_TARGET, "Tumor Cells", "T & NK Cells"],
    top_n=15,
    orderby="lr_mean",
    orderby_ascending=False,
    figure_size=(13, 9),
) + p9.theme(plot_margin=0.02)

using the stored global interactions

CPU times: user 202 ms, sys: 8.07 ms, total: 210 ms
Wall time: 212 ms

# Parts 1-3 are done with, so we release what they left behind. Everything dropped here is still
# in the prepared object on disk, so re-running the cells above restores it.
import gc

for _name in ("plain", "hard", "softs", "plain_r", "hard_r", "comparison", "summary",
              "pcf", "bivar", "top_biv", "biv_small", "contrast", "lrdata", "scores", "scores2",
              "proximity", "auc", "g_run", "receivers", "edges", "nn_dist", "bw_df"):
    globals().pop(_name, None)

for _key in ("cpdb_plain", "cpdb_hard", "cpdb_soft_bw30", "cpdb_soft_bw100", "cpdb_soft_bw300",
             "cross_pcf", "lric_ct", "bivariate_morans", "global_interactions"):
    adata.uns.pop(_key, None)

for _key in list(adata.obsp):          # the contact and kernel graphs are not used from here on
    del adata.obsp[_key]
adata.obsm.pop("inflow_scores", None)

gc.collect()
print(f"adata is now {adata.n_obs:,} x {adata.n_vars:,}; obsp={list(adata.obsp)}")
print("InterScale results still in place:", sorted(k for k in adata.uns if k.startswith("interscale_")))

adata is now 400,600 x 3,197; obsp=[]
InterScale results still in place: ['interscale_dim_importance', 'interscale_gene_ranks', 'interscale_net_flow_mean', 'interscale_net_flow_sd', 'interscale_net_streams', 'interscale_programme_genes', 'interscale_programmes']

%%time
from pathlib import Path

import interscale
import torch
from huggingface_hub import snapshot_download
from interscale.config import load_config
from interscale.evaluation import (
    calculate_dim_importance,
    calculate_gene_ranks,
    gene_loadings,
    get_genes_dim,
    plot_all_spatial_net_streams,
    plot_global_directionality,
)
from interscale.evaluation.net_streams import compute_hierarchical_net_flow
from interscale.tl import set_full_reproducibility

HF_REPO = "theislab/InterScale"
HF_MODEL = "ovarian_xenium/heldout_split"
model_dir = Path(snapshot_download(HF_REPO, allow_patterns=f"{HF_MODEL}/*")) / HF_MODEL
print(*sorted(p.name for p in model_dir.iterdir()), sep="\n")

set_full_reproducibility()

is_cfg = load_config(model_dir / "config.yaml")
IS_GENES = pd.read_csv(model_dir / "genes.tsv", sep="\t", index_col=0).index.tolist()
n_per_window = adata.obs["sliding_window"].value_counts()

print(f"\n{len(IS_GENES):,} input genes | latent width {is_cfg.model.n_embed} | "
      f"masking: {is_cfg.dataset.mask_strategy} @ {is_cfg.dataset.mask_percentage} | "
      f"window key: {list(is_cfg.dataset.sample_key)} | max cells per window: "
      f"{is_cfg.model.global_component.parameters.max_seq_len}")
print(f"{len(n_per_window):,} windows, {n_per_window.median():.0f} cells in the median window")

Downloading bytes:           |  0.00B

Reconstructing (incomplete total...): |          |  0.00B /  0.00B

Fetching 5 files:   0%|          | 0/5 [00:00<?, ?it/s]

README.md
config.yaml
genes.tsv
metrics.json
model.pt

3,000 input genes | latent width 256 | masking: node @ 0.3 | window key: ['sliding_window'] | max cells per window: 3167
199 windows, 2218 cells in the median window
CPU times: user 6.69 s, sys: 904 ms, total: 7.59 s
Wall time: 7.9 s

# from interscale.geome_dataloader import GraphAnnDataModule
# from interscale.tl import prepare_geome_dataset, set_full_reproducibility
#
# set_full_reproducibility()
#
# # hold out whole windows: 70 / 15 / 15
# rng = np.random.default_rng(44)
# windows = np.array(sorted(n_per_window.index))
# rng.shuffle(windows)
# n_train, n_val = int(0.70 * len(windows)), int(0.15 * len(windows))
# split = pd.Series("test", index=windows)
# split[windows[:n_train]], split[windows[n_train:n_train + n_val]] = "train", "val"
# adata_is.obs["split"] = adata_is.obs["sliding_window"].map(split).astype("category")
#
# # the shipped config, with this run's paths; any other key can be overridden the same way
# train_cfg = load_config(model_dir / "config.yaml", overrides=["model.save", "my_interscale_run/"])
#
# interscale.model.CombinedModel._setup_anndata(
#     adata=adata_is, prediction_task="regression", layer_key="log1p_norm",
#     sample_key_list=["sliding_window"], split_key="split",
# )
# model = interscale.model.CombinedModel(adata_is, cfg=train_cfg)
# graphs, _ = prepare_geome_dataset(adata_is, train_cfg)          # one graph per window, per split
# datamodule = GraphAnnDataModule(
#     datas=graphs, batch_size=train_cfg.dataset.batch_size, num_workers=1, learning_type="node",
#     mask_strategy=train_cfg.dataset.mask_strategy, mask_percentage=train_cfg.dataset.mask_percentage,
# )
# os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")      # InterScale trains deterministically
# model.train(max_epochs=train_cfg.optim.n_epochs, datamodule=datamodule, early_stopping=True, devices=1)
# model.save("my_interscale_run/", overwrite=True)

%%time
# Running the model across all 199 windows is the one step of this chapter that cannot be made to
# fit: it materialises a dense reconstruction of all 3,000 genes for *both* components plus a
# [cells x 3,167] attention matrix, and `get_model_output` deliberately runs with autograd
# enabled. Measured, it needs well over 120 GB. The outputs each plot below plots were therefore
# computed once on a large-memory machine and stored in the prepared object.
#
# To reproduce them, run this on a machine with a few hundred GB of RAM:
#
# adata_is = adata[:, IS_GENES].copy()
# # the model's config reads .layers["log1p_norm"], which in this object is exactly .X;
# # assigning it is a reference, not a copy, so it costs no extra memory
# adata_is.layers["log1p_norm"] = adata_is.X
# torch.set_num_threads(os.cpu_count() or 1)
# interscale.model.CombinedModel._setup_anndata(
#     adata=adata_is, prediction_task="regression", layer_key="log1p_norm",
#     sample_key_list=["sliding_window"], split_key=None,
# )
# is_model = interscale.model.CombinedModel.load(
#     str(model_dir), adata_is, is_cfg, model_name="", local_component=True, global_component=True
# )
# # a freshly loaded torch module is in training mode, where dropout is active; switch it off, or
# # every run of get_model_output returns slightly different attention and reconstructions
# is_model.module.eval()
# is_result = is_model.get_model_output(adata_is)  # not under torch.no_grad(): it uses autograd internally

print("stored InterScale results:")
for _k in sorted(k for k in adata.uns if k.startswith("interscale_")):
    _v = adata.uns[_k]
    print(f"  adata.uns[{_k!r}]: {type(_v).__name__}"
          + (f" {_v.shape}" if hasattr(_v, "shape") else ""))
print("  adata.obs['local_programme'], adata.obs['global_programme']")

stored InterScale results:
  adata.uns['interscale_dim_importance']: dict
  adata.uns['interscale_gene_ranks']: DataFrame (3000, 4)
  adata.uns['interscale_net_flow_mean']: DataFrame (8, 8)
  adata.uns['interscale_net_flow_sd']: DataFrame (8, 8)
  adata.uns['interscale_net_streams']: dict
  adata.uns['interscale_programme_genes']: dict
  adata.uns['interscale_programmes']: dict
  adata.obs['local_programme'], adata.obs['global_programme']
CPU times: user 204 μs, sys: 37 μs, total: 241 μs
Wall time: 234 μs

%%time
# compute_hierarchical_net_flow reads the attention matrix from get_model_output, so its result
# was precomputed. With `is_result` in hand it would be:
#
# flows = compute_hierarchical_net_flow(
#     is_result, window_key="sliding_window", sample_key="tissue_region", condition_key=None,
#     cell_type_col=CT_KEY, compute_net=True,
# )
# flow_mean, flow_sd = flows["all_samples"]["mean"], flows["all_samples"]["std"]

flow_mean = adata.uns["interscale_net_flow_mean"]
flow_sd = adata.uns["interscale_net_flow_sd"]
order = [c for c in CT_ORDER if c in flow_mean.index]
flow_mean, flow_sd = flow_mean.loc[order, order], flow_sd.loc[order, order]

flow_long = (
    flow_mean.stack().rename("net_flow").rename_axis(["sender", "receiver"]).reset_index()
    .assign(sd_across_regions=flow_sd.stack().values)
    .query("sender != receiver")
    .sort_values("net_flow", ascending=False)
    .reset_index(drop=True)
)
print("strongest directed attention flows (sender -> receiver):")
display(flow_long.head(10).round(3))
ex = flow_long.query("sender == @EXAMPLE_SOURCE and receiver == @EXAMPLE_TARGET").iloc[0]
print(f"running example {EXAMPLE_SOURCE} -> {EXAMPLE_TARGET}: net flow {ex.net_flow:.3f} "
      f"(sd {ex.sd_across_regions:.3f}), rank {ex.name + 1} of {len(flow_long)} ordered pairs")

strongest directed attention flows (sender -> receiver):

running example Tumor Associated Fibroblasts -> Macrophages: net flow 0.178 (sd 0.067), rank 12 of 56 ordered pairs
CPU times: user 16.2 ms, sys: 5.1 ms, total: 21.3 ms
Wall time: 20.6 ms

plot_global_directionality(
    flow_mean, flow_sd, only_positive=True,
    title="InterScale attention flow, ovarian adenocarcinoma", figsize=(12, 8),
)
plt.tight_layout()
plt.show()

<Figure size 400x400 with 0 Axes>

%%time
# The streams are derived from the attention matrix, which is far too large to carry here.
# `plot_all_spatial_net_streams(..., return_streams=True)` returns the gridded vector fields
# instead of drawing them, and those are small, so they were precomputed:
#
# streams, X_lin, Y_lin = plot_all_spatial_net_streams(
#     is_result, fov_id=region_id, fov_key="tissue_region", window_key="sliding_window",
#     cell_type_col=CT_KEY, grid_res=60, return_streams=True,
# )
#
# What follows is the drawing half of that function, reading the stored fields. The grid is in
# raw micrometres (the scale factor is 1 without an image), so it overlays the coordinates
# directly.
from matplotlib.colors import hsv_to_rgb, rgb_to_hsv, to_hex, to_rgb
from matplotlib.lines import Line2D

ns = adata.uns["interscale_net_streams"]
region_id = ns["region_id"]
X_lin, Y_lin = np.asarray(ns["X_lin"]), np.asarray(ns["Y_lin"])

in_region = (adata.obs["tissue_region"] == region_id).to_numpy()
xy = np.asarray(adata.obsm[SPATIAL_KEY])[in_region]
cats = list(adata.obs[CT_KEY].cat.categories)
colours = dict(zip(cats, adata.uns["celltype_colors"]))
labels = adata.obs[CT_KEY].to_numpy()[in_region]

fig, ax = plt.subplots(figsize=(8, 7))
for cat in cats:
    m = labels == cat
    if m.any():
        ax.scatter(xy[m, 0], xy[m, 1], s=3, c=colours[cat], linewidths=0, rasterized=True, label=cat)

legend_elements = []
for cat in cats:
    if cat not in ns["U"]:
        continue
    U, V = np.asarray(ns["U"][cat]), np.asarray(ns["V"][cat])
    mag = np.sqrt(U**2 + V**2)
    if mag.max() == 0:
        continue
    thresh = 0.01 * mag.max()
    Un = np.divide(U, mag, out=np.full_like(U, np.nan), where=mag > thresh)
    Vn = np.divide(V, mag, out=np.full_like(V, np.nan), where=mag > thresh)
    if np.all(np.isnan(Un)):
        continue
    h = rgb_to_hsv(to_rgb(colours[cat]))
    dark = to_hex(hsv_to_rgb([h[0], h[1], h[2] * 0.7]))
    ax.streamplot(X_lin, Y_lin, Un, Vn, color=dark, linewidth=1.2, density=1.4, arrowsize=1.2)
    legend_elements.append(Line2D([0], [0], color=dark, lw=2, label=f"Flow: {cat}"))

if legend_elements:
    ax.legend(handles=legend_elements, loc="upper center", ncol=4, bbox_to_anchor=(0.5, -0.15),
              title=f"Net Flow Directions - {CT_KEY}")
ax.set_aspect("equal")
ax.invert_yaxis()
ax.set_xticks([])
ax.set_yticks([])
ax.set_title(f"Attention flow streams, tissue region {region_id}")
plt.show()

CPU times: user 4.73 s, sys: 43.4 ms, total: 4.77 s
Wall time: 4.82 s

%%time
# calculate_gene_ranks scores each gene against both dense reconstruction layers, so it needs
# `is_result`; the table was precomputed. To reproduce it:
#
# gene_rank_df = calculate_gene_ranks(is_result, layers_local_pred="_y_pred_local",
#                                     layers_global_pred="_y_pred_global")
# gene_rank_df = gene_rank_df.set_index("gene")

gene_rank_df = adata.uns["interscale_gene_ranks"]
cols = ["Local Rank", "Global Rank", "Rank Difference"]
n = len(gene_rank_df)
print(f"{n:,} genes | local wins on {(gene_rank_df['Rank Difference'] > 0).mean():.0%} | "
      f"Spearman(local rank, global rank) = {gene_rank_df['Local Rank'].corr(gene_rank_df['Global Rank'], method='spearman'):.2f}")
print("\nreconstructed much better by the LOCAL component:")
display(gene_rank_df.nlargest(8, "Rank Difference")[cols])
print("reconstructed much better by the GLOBAL component:")
display(gene_rank_df.nsmallest(8, "Rank Difference")[cols])

# the genes of this chapter's running examples
focus = [g for g in ("THBS1", "LRP1", "CD47", "MFGE8", "ITGB3", "POSTN") if g in gene_rank_df.index]
focus_df = gene_rank_df.loc[focus, cols].assign(
    scale=lambda d: np.where(d["Rank Difference"] > 0, "local", "global"),
    percentile_local=lambda d: (d["Local Rank"] / n * 100).round(0),
    percentile_global=lambda d: (d["Global Rank"] / n * 100).round(0),
)
print("running-example genes (percentile 100 = best reconstructed):")
display(focus_df)

3,000 genes | local wins on 52% | Spearman(local rank, global rank) = -0.05

reconstructed much better by the LOCAL component:

reconstructed much better by the GLOBAL component:

running-example genes (percentile 100 = best reconstructed):

CPU times: user 619 ms, sys: 54.1 ms, total: 674 ms
Wall time: 94.8 ms

%%time
# gene_loadings relates each latent dimension to the genes, calculate_dim_importance ranks the
# dimensions, and get_genes_dim returns the top genes of the most important ones. All three read
# the embeddings produced by get_model_output, so the result was precomputed. To reproduce it:
#
# gene_loadings(is_result, is_model, local_latent_key="_local_emb", global_latent_key="_global_emb",
#               layer_key="log1p_norm")
# N_DIMS = 4
# dims = {}
# for which in ("local", "global"):
#     res = calculate_dim_importance(is_result, which=which, mode="full", use_ratio=True,
#                                    cumulative_cutoff=None, spacing=2, n_top=48, store=False)
#     dims[which] = np.asarray(res["dim_sorted"])[:N_DIMS].tolist()
# programmes = {}
# for which in ("local", "global"):
#     df, _ = get_genes_dim(
#         is_result, which=which, dims=dims[which], n_top=10, X_layer="log1p_norm", min_frac=0.2,
#         rank_by="loading_x_sd", enforce_specificity=True, specificity_min=1.5,
#         plot=True, show=False,
#     )
#     programmes[which] = df.index

_stored = adata.uns["interscale_programmes"]
tables = {w: _stored["tables"][w] for w in ("local", "global")}
dims = {w: list(_stored["dims"][w]) for w in ("local", "global")}
programmes = {w: tables[w].index for w in tables}
print("most important dimensions:", dims)
for which in ("local", "global"):
    print(f"\n{which} embedding: top genes of the {len(dims[which])} most important dimensions")
    display(tables[which].head(12))

shared = programmes["local"].intersection(programmes["global"])
unique_local = programmes["local"].difference(programmes["global"])
unique_global = programmes["global"].difference(programmes["local"])
print(f"\nshared        ({len(shared)}):", ", ".join(shared))
print(f"unique local  ({len(unique_local)}):", ", ".join(unique_local))
print(f"unique global ({len(unique_global)}):", ", ".join(unique_global))

most important dimensions: {'local': [np.int64(163), np.int64(198), np.int64(227), np.int64(189)], 'global': [np.int64(189), np.int64(198), np.int64(227), np.int64(31)]}

local embedding: top genes of the 4 most important dimensions

global embedding: top genes of the 4 most important dimensions

shared        (16): CD47, RAN, PCNA, ATP5F1A, MYBL2, SNRPB, BMP7, CDC20, STAT1, OAS3, IFIT1, GBP1, IFIT3, RSAD2, PTPRF, MUC1
unique local  (14): CD99, CDH6, CXCL10, EIF4A1, FOLR1, HMGA1, HSPG2, IFITM1, IRF1, ITGB3, MX1, SLC2A1, ST14, TAP1
unique global (16): AP2M1, COMP, CTHRC1, GAS6, IFI35, IGFBP2, ITGB4, LTBP2, MYLK, NDRG1, PALLD, PPP1R1B, SCNN1A, SULF1, THBS1, THY1
CPU times: user 246 ms, sys: 18 ms, total: 264 ms
Wall time: 65.2 ms

%%time
import decoupler as dc
from scipy.stats import hypergeom
from statsmodels.stats.multitest import multipletests

go = dc.op.resource("MSigDB", organism="human")
go = go[go["collection"] == "go_biological_process"]
go = go[go["genesymbol"].isin(IS_GENES)].groupby("geneset")["genesymbol"].apply(set)
go = go[go.map(len).between(5, 300)]                 # sets with 5-300 panel genes


def ora(genes, sets=go, universe=None):
    universe = len(IS_GENES) if universe is None else universe
    genes = set(genes)
    rows = []
    for name, members in sets.items():
        k = len(genes & members)
        if k:
            rows.append((name, k, len(members), hypergeom.sf(k - 1, universe, len(members), len(genes)),
                         ", ".join(sorted(genes & members))))
    df = pd.DataFrame(rows, columns=["term", "overlap", "set_size", "pval", "genes"])
    df["padj"] = multipletests(df["pval"], method="fdr_bh")[1]
    df["term"] = df["term"].str.removeprefix("GOBP_").str.replace("_", " ").str.lower()
    return df.sort_values("pval").reset_index(drop=True)


for name, genes in (("local", unique_local), ("global", unique_global)):
    print(f"{name} programme -- top GO biological processes (background: the 3,000 model genes)")
    display(ora(genes).head(5)[["term", "overlap", "set_size", "padj", "genes"]])

local programme -- top GO biological processes (background: the 3,000 model genes)

global programme -- top GO biological processes (background: the 3,000 model genes)

CPU times: user 26 s, sys: 3.13 s, total: 29.1 s
Wall time: 51.5 s

# The programme scores are per cell, so they were precomputed alongside everything else:
#
# sc.tl.score_genes(is_result, list(unique_local), score_name="local_programme", use_raw=False)
# sc.tl.score_genes(is_result, list(unique_global), score_name="global_programme", use_raw=False)

xy_is = np.asarray(adata.obsm[SPATIAL_KEY])
fig, axes = plt.subplots(1, 2, figsize=(15, 7))
for ax, key in zip(axes, ("local_programme", "global_programme")):
    v = adata.obs[key].to_numpy()
    lim = np.quantile(np.abs(v), 0.99)
    sca = ax.scatter(xy_is[:, 0], xy_is[:, 1], c=v, s=0.1, cmap="RdBu_r", vmin=-lim, vmax=lim,
                     linewidths=0, rasterized=True)
    ax.set_aspect("equal")
    ax.invert_yaxis()
    ax.set(title=key.replace("_", " "), xlabel="x (µm)", ylabel="y (µm)")
    fig.colorbar(sca, ax=ax, shrink=0.7, label="score")
plt.tight_layout()
plt.show()

print("mean programme score per cell type:")
display(adata.obs.groupby(CT_KEY, observed=True)[["local_programme", "global_programme"]].mean()
        .loc[order].round(3))

mean programme score per cell type:

	I	var_norm
POSTN	0.575240	1.888382e-07
TNXB	0.562507	1.888382e-07
MMP11	0.555111	1.888382e-07
BCAM	0.522425	1.888382e-07
COL5A1	0.520887	1.888382e-07

	ligand	receptor
0	LGALS9	MET
1	LGALS9	CD44
2	LGALS9	LRP1
3	LGALS9	CD47
4	LGALS9	COLEC12

	ligand	ligand_complex	ligand_means	ligand_props	receptor	receptor_complex	receptor_means	receptor_props	source	target	lr_means
0	COL4A1	COL4A1	3.214948	0.939905	CD47	CD47	1.727521	0.941885	Pericytes	Tumor Cells	2.471235
1	COL4A1	COL4A1	2.990420	0.959916	CD47	CD47	1.727521	0.941885	Endothelial Cells	Tumor Cells	2.358971
2	COL4A1	COL4A1	3.214948	0.939905	CD47	CD47	1.027229	0.516687	Pericytes	Granulosa Cells	2.121089
3	COL4A1	COL4A1	3.214948	0.939905	CD44	CD44	0.976135	0.684667	Pericytes	T & NK Cells	2.095541
4	COL4A1	COL4A1	3.214948	0.939905	CD47	CD47	0.916300	0.658726	Pericytes	T & NK Cells	2.065624
5	COL4A1	COL4A1	2.990420	0.959916	CD47	CD47	1.027229	0.516687	Endothelial Cells	Granulosa Cells	2.008824
6	COL4A1	COL4A1	3.214948	0.939905	CD47	CD47	0.801665	0.616473	Pericytes	Macrophages	2.008307
7	THBS2	THBS2	2.276985	0.918732	CD47	CD47	1.727521	0.941885	Tumor Associated Fibroblasts	Tumor Cells	2.002253

	unconstrained	soft 30 µm	soft 100 µm	soft 300 µm
rank
1	Pericytes -> Tumor Cells: COL4A1^CD47	Endothelial Cells -> Tumor Cells: COL4A1^CD47	Pericytes -> Tumor Cells: COL4A1^CD47	Pericytes -> Tumor Cells: COL4A1^CD47
2	Endothelial Cells -> Tumor Cells: COL4A1^CD47	Tumor Associated Fibroblasts -> Tumor Associat...	Endothelial Cells -> Tumor Cells: COL4A1^CD47	Endothelial Cells -> Tumor Cells: COL4A1^CD47
3	Pericytes -> Granulosa Cells: COL4A1^CD47	Pericytes -> Endothelial Cells: COL4A1^CD47	Pericytes -> Macrophages: COL4A1^CD47	Pericytes -> T & NK Cells: COL4A1^CD44
4	Pericytes -> T & NK Cells: COL4A1^CD44	Pericytes -> Macrophages: COL4A1^CD47	Pericytes -> Macrophages: COL4A1^CD44	Pericytes -> T & NK Cells: COL4A1^CD47
5	Pericytes -> T & NK Cells: COL4A1^CD47	Pericytes -> Macrophages: COL4A1^CD44	Pericytes -> T & NK Cells: COL4A1^CD44	Pericytes -> Macrophages: COL4A1^CD47
6	Endothelial Cells -> Granulosa Cells: COL4A1^CD47	Pericytes -> Endothelial Cells: COL4A1^CD93	Pericytes -> T & NK Cells: COL4A1^CD47	Pericytes -> Macrophages: COL4A1^CD44
7	Pericytes -> Macrophages: COL4A1^CD47	Macrophages -> Macrophages: CD14^ITGB2	Pericytes -> Pericytes: COL4A1^CD44	Tumor Associated Fibroblasts -> Tumor Cells: T...
8	Tumor Associated Fibroblasts -> Tumor Cells: T...	Endothelial Cells -> Endothelial Cells: COL4A1...	Tumor Associated Fibroblasts -> Tumor Cells: T...	Tumor Associated Fibroblasts -> Tumor Cells: C...
9	Tumor Associated Fibroblasts -> Tumor Cells: C...	Pericytes -> Endothelial Cells: COL18A1^ITGA5	Tumor Associated Fibroblasts -> Tumor Cells: C...	Endothelial Cells -> T & NK Cells: COL4A1^CD44
10	Pericytes -> Macrophages: COL4A1^CD44	Pericytes -> Pericytes: COL4A1^CD44	Endothelial Cells -> T & NK Cells: COL4A1^CD44	Pericytes -> Malignant Cells Lining Cyst: COL4...

	rows in the table	cell type pairs kept	rank of THBS1^LRP1 (Tumor Associated Fibroblasts -> Macrophages)
unconstrained	22774	64	124.0
hard mask (NEA z > 2)	8326	22	42.0
soft, bandwidth 30 µm	22774	64	60.0
soft, bandwidth 100 µm	22774	64	85.0
soft, bandwidth 300 µm	22774	64	106.0

Cell-cell communication in spatial transcriptomics¶

Background and motivation¶

Setup¶

Loading the data¶

The dataset¶

Defining the neighborhood¶

Filtering to spatially variable genes¶

Part 1: Who is near whom, and at what distance?¶

Neighborhood enrichment analysis¶

The cross pair-correlation function¶

From cell types to a ligand-receptor pair¶

Part 2: Which ligand-receptor pairs, given proximity?¶

Spatial constraints on ligand-receptor analysis¶

Hard constraint: restricting cell-type pairs¶

Soft constraint: weighting by spatial proximity¶

Global bivariate Moran's I¶

Part 3: Where does it happen locally?¶

Inflow: local interaction scores by sender cell type¶

Part 4: Multi-scale models¶

InterScale¶

A pre-trained model¶

Training InterScale yourself¶

Loading the model and running it on every window of the tissue¶

Attention flow between cell types¶

Which genes are local, and which depend on tissue context?¶

Gene programs of the two scales¶

Further reading¶

	bandwidth	neighbours
0	5.000000	0.0
1	6.153846	0.0
2	7.307692	0.0
3	8.461538	1.0
4	9.615385	1.0
5	10.769231	2.0
6	11.923077	3.0
7	13.076923	4.0
8	14.230769	4.0
9	15.384615	5.0

	source	target	zscore
0	Pericytes	Endothelial Cells	230.418104
1	Endothelial Cells	Pericytes	230.418104
2	Macrophages	T & NK Cells	49.040450
3	T & NK Cells	Macrophages	49.040450
4	Pericytes	Malignant Cells Lining Cyst	30.749900
5	Malignant Cells Lining Cyst	Pericytes	30.749900
6	Endothelial Cells	Malignant Cells Lining Cyst	20.013160
7	Malignant Cells Lining Cyst	Endothelial Cells	20.013160

	source	target	interaction	radius	g
0	Endothelial Cells	Granulosa Cells	Endothelial Cells^Granulosa Cells	0.0	0.267444
1	Endothelial Cells	Granulosa Cells	Endothelial Cells^Granulosa Cells	40.0	0.388438
2	Endothelial Cells	Granulosa Cells	Endothelial Cells^Granulosa Cells	60.0	0.385292
3	Endothelial Cells	Granulosa Cells	Endothelial Cells^Granulosa Cells	80.0	0.447517
4	Endothelial Cells	Granulosa Cells	Endothelial Cells^Granulosa Cells	100.0	0.450199

	source	target	interaction	score	peak_radius
0	Endothelial Cells	Pericytes	Endothelial Cells^Pericytes	2.110385	0.0
1	Malignant Cells Lining Cyst	Pericytes	Malignant Cells Lining Cyst^Pericytes	1.463117	40.0
2	Endothelial Cells	Malignant Cells Lining Cyst	Endothelial Cells^Malignant Cells Lining Cyst	1.020345	40.0
3	Macrophages	T & NK Cells	Macrophages^T & NK Cells	0.285744	0.0
4	Pericytes	Tumor Associated Fibroblasts	Pericytes^Tumor Associated Fibroblasts	0.146760	40.0
5	Malignant Cells Lining Cyst	T & NK Cells	Malignant Cells Lining Cyst^T & NK Cells	0.016894	0.0
6	Macrophages	Malignant Cells Lining Cyst	Macrophages^Malignant Cells Lining Cyst	0.010477	0.0
7	Granulosa Cells	Tumor Cells	Granulosa Cells^Tumor Cells	-0.016826	40.0

	source	target	ligand_complex	receptor_complex	interaction	score	peak_radius
0	Macrophages	Tumor Associated Fibroblasts	THBS1	LRP1	THBS1^LRP1	0.417796	0.0
1	Tumor Associated Fibroblasts	Macrophages	THBS1	LRP1	THBS1^LRP1	0.079690	0.0

	source	target	mean_distance	interacting	proximity
0	Endothelial Cells	Endothelial Cells	13.209447	1	9.076117e-01
1	Endothelial Cells	Granulosa Cells	226.315482	1	4.387564e-13
2	Endothelial Cells	Macrophages	19.005599	1	8.181793e-01
3	Endothelial Cells	Malignant Cells Lining Cyst	33.911130	1	5.278890e-01
4	Endothelial Cells	Pericytes	33.791973	1	5.302602e-01

	ligand	receptor	ligand_props	receptor_props	morans
0	THBS1	LRP1	0.538415	0.383600	0.265674
1	CDH1	PTPRF	0.533899	0.659071	0.224855
2	CTHRC1	ROR2	0.362132	0.141817	0.221099
3	PLAU	LRP1	0.352601	0.383600	0.217663
4	PLAU	ITGB5	0.352601	0.447217	0.191752
5	POSTN	PTK7	0.349159	0.607282	0.186736
6	PLAT	LRP1	0.243258	0.383600	0.171124
7	CCN2	LRP1	0.222212	0.383600	0.165204
8	PLAU	PLAUR	0.352601	0.497821	0.157554
9	PLAU	ITGA5	0.352601	0.292973	0.151455

	mean_score	prop_nonzero
celltype
Tumor Associated Fibroblasts	1.428458	0.772797
Macrophages	0.366479	0.423094
Malignant Cells Lining Cyst	0.335311	0.594295
Pericytes	0.199431	0.440243
T & NK Cells	0.089203	0.145402
Endothelial Cells	0.074783	0.212474
Tumor Cells	0.022257	0.116023
Granulosa Cells	0.021961	0.051916

	sender	receiver	net_flow	sd_across_regions
0	Malignant Cells Lining Cyst	Endothelial Cells	0.380	0.062
1	Tumor Associated Fibroblasts	Endothelial Cells	0.335	0.077
2	Malignant Cells Lining Cyst	Pericytes	0.325	0.079
3	T & NK Cells	Endothelial Cells	0.325	0.134
4	Tumor Associated Fibroblasts	Pericytes	0.289	0.112
5	T & NK Cells	Pericytes	0.273	0.127
6	Granulosa Cells	Endothelial Cells	0.267	0.211
7	Tumor Associated Fibroblasts	Tumor Cells	0.225	0.042
8	Tumor Cells	Endothelial Cells	0.225	0.145
9	Granulosa Cells	Pericytes	0.208	0.204

	Local Rank	Global Rank	Rank Difference
gene
ANXA8	2969.0	3.0	2966.0
PRG4	2957.0	1.0	2956.0
CCL8	2942.0	26.0	2916.0
APLN	2894.0	27.0	2867.0
BCL6B	2901.0	46.0	2855.0
FRAS1	2841.0	15.0	2826.0
POU2AF1	2971.0	150.0	2821.0
STAR	2949.0	135.0	2814.0

	Local Rank	Global Rank	Rank Difference
gene
MX1	34.0	2977.0	-2943.0
CD47	39.0	2978.0	-2939.0
FBL	27.0	2963.0	-2936.0
CFB	31.0	2949.0	-2918.0
EEF1G	1.0	2919.0	-2918.0
LDHA	24.0	2940.0	-2916.0
TAP1	42.0	2953.0	-2911.0
PKM	4.0	2913.0	-2909.0

	Local Rank	Global Rank	Rank Difference	scale	percentile_local	percentile_global
gene
THBS1	2926.0	2947.0	-21.0	global	98.0	98.0
LRP1	2965.0	2935.0	30.0	local	99.0	98.0
CD47	39.0	2978.0	-2939.0	global	1.0	99.0
MFGE8	178.0	2719.0	-2541.0	global	6.0	91.0
ITGB3	363.0	2888.0	-2525.0	global	12.0	96.0
POSTN	2999.0	2999.0	0.0	global	100.0	100.0

	163	198	227	189
CD47	0.250946	0.064052	0.061843	-0.151916
FOLR1	0.225095	-0.063723	0.097895	-0.081138
ST14	0.203118	-0.026383	0.071885	-0.106362
CD99	-0.192145	-0.089149	-0.108125	0.092365
CDH6	0.147512	-0.043983	0.052146	-0.037930
ITGB3	0.147402	0.037966	0.040374	-0.051315
HMGA1	0.145910	0.080428	0.043126	-0.050929
RAN	0.043148	0.332891	0.107363	-0.073396
PCNA	0.132594	0.269553	-0.016913	-0.116129
ATP5F1A	0.099047	0.249352	0.055120	-0.059959
MYBL2	0.090483	0.248446	-0.086863	-0.057684
SNRPB	0.125678	0.243627	0.066017	-0.126895

	189	198	227	31
COMP	0.165916	-0.098727	-0.011677	0.096775
CD47	-0.140227	0.033082	0.086113	-0.079671
RAN	-0.058360	0.208783	0.055546	0.062465
PCNA	-0.080639	0.198901	-0.008590	0.010436
ITGB4	-0.093167	-0.192790	0.114685	0.008687
BMP7	-0.074336	0.182576	0.081009	0.043069
MYBL2	-0.046191	0.179908	-0.069660	-0.012530
SNRPB	-0.099614	0.163844	0.065552	0.027957
PPP1R1B	-0.043655	0.151052	0.090618	0.081622
CDC20	-0.029488	0.150335	-0.027776	0.003753
ATP5F1A	-0.059676	0.148794	0.045853	0.062917
STAT1	-0.080127	0.026849	-0.302979	-0.053553

	term	overlap	set_size	padj	genes
0	t cell migration	3	40	0.103186	CD99, CXCL10, ITGB3
1	response to mechanical stimulus	4	98	0.103186	CXCL10, IRF1, ITGB3, SLC2A1
2	receptor mediated endocytosis	4	100	0.103186	FOLR1, HSPG2, ITGB3, MX1
3	defense response to virus	4	101	0.103186	CXCL10, IFITM1, IRF1, MX1
4	cellular response to external stimulus	3	45	0.103186	IRF1, ITGB3, SLC2A1

	term	overlap	set_size	padj	genes
0	regulation of synaptic vesicle endocytosis	2	5	0.157477	AP2M1, MYLK
1	regulation of synaptic vesicle recycling	2	6	0.157477	AP2M1, MYLK
2	negative regulation of fibroblast growth facto...	2	8	0.189497	SULF1, THBS1
3	smooth muscle contraction	3	42	0.189497	COMP, MYLK, SULF1
4	cell substrate junction organization	3	43	0.189497	ITGB4, THBS1, THY1

	local_programme	global_programme
celltype
Endothelial Cells	0.174	0.120
Granulosa Cells	0.399	0.068
Macrophages	0.167	0.014
Malignant Cells Lining Cyst	0.098	0.301
Pericytes	0.072	0.243
T & NK Cells	0.278	-0.002
Tumor Associated Fibroblasts	0.072	0.663
Tumor Cells	0.348	0.098