Evaluating Spatial Perturbation Prediction on Perturb-FISH Data

• Dataset: ` <https://www.nature.com/articles/s41576-025-00857-8>`__

• Evaluate in silico perturbation prediction, by perturbing control cancer cells and observing the changes in nearby T cells.

• Training is performed in a leave-one-out fasion by removing all cancer cells with the KO to be tested.

import scanpy as sc
import squidpy as sq
import pandas as pd
import matplotlib.pyplot as plt
from tqdm.notebook import tqdm
import numba
import numpy as np
import torch

import sys
sys.path.append("../")
import steamboat as sf
import steamboat.tools

import matplotlib
plt.rcParams['pdf.fonttype'] = 42
matplotlib.rcParams['mathtext.fontset'] = 'dejavuserif'
matplotlib.rcParams['font.family'] = 'arial'

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\dask\dataframe\_pyarrow_compat.py:15: FutureWarning: Minimal version of pyarrow will soon be increased to 14.0.1. You are using 11.0.0. Please consider upgrading.
  warnings.warn(

adata = sc.read_h5ad("G:/data/perturbfish/cleaned/tumors_qc_test.h5ad")
adata

# The stored data is already transformed. This step is to find the spatial neighbors.
adatas = sf.prep_adatas([adata], norm=False, log1p=False, n_neighs=8)

test_vars = ['CHUK', 'IRAK1', 'TRAM1', 'LBP', 'IRAK4', 'PELI1', 'TAB2', 'MAP2K2', 'MAP2K6', 'IRF7', 'MYD88']

res = {}

for test_var in test_vars:
    # find all cells near ko or control tumor cells
    adata.obs['near_ko'] = (adata.obsp['spatial_connectivities'] @ (adata.obs['perturbation'] == test_var).to_numpy()[:, None]) > 0
    adata.obs['near_control'] = (adata.obsp['spatial_connectivities'] @ (adata.obs['perturbation'] == 'Control').to_numpy()[:, None]) > 0

    # find T cells near ko or control tumor cells
    adata.obs['t_near_ko'] = (adata.obs['near_ko'] & (adata.obs['celltype2'] == 'T cells'))
    adata.obs['t_near_control'] = (adata.obs['near_control'] & (adata.obs['celltype2'] == 'T cells'))

    # Find ground truth DEGs
    sub_adata = adata[adata.obs['t_near_ko'] | adata.obs['t_near_control']].copy()
    sub_adata.obs['test_group'] = sub_adata.obs['t_near_ko'].apply(lambda x: 'test' if x else 'control')
    sc.tl.rank_genes_groups(sub_adata, 'test_group', method='wilcoxon', key_added = "wilcoxon")
    # sc.pl.rank_genes_groups(sub_adata, n_genes=25, sharey=False, key="wilcoxon")
    cmp_df_gt = sc.get.rank_genes_groups_df(sub_adata, group="test", key='wilcoxon')
    # cmp_df_gt

    # keep a copy of ko cells and remove those cells (prevent leakage of data in training)
    substitute = adata.X[adata.obs['perturbation'] == test_var]
    train_adata = adata[adata.obs['perturbation'] != test_var, :].copy()
    train_adatas = sf.prep_adatas([train_adata], norm=False, log1p=False, n_neighs=8)
    train_dataset = sf.make_dataset(train_adatas, sparse_graph=True, regional_obs=[])

    # Train the model
    device = 'cuda'
    sf.set_random_seed(0)
    model = sf.Steamboat(adata.var_names.tolist(), n_heads=50, n_scales=2)
    model = model.to(device)

    model.fit(train_dataset.to(device), entry_masking_rate=0.1, feature_masking_rate=0.1,
              max_epoch=10000,
              loss_fun=torch.nn.MSELoss(reduction='sum'),
              opt=torch.optim.Adam, opt_args=dict(lr=0.1), stop_eps=1e-4, report_per=500, stop_tol=200, device=device)

    # Substitute control cells with random ko cells
    test_adata = train_adata.copy()
    test_adata.X[test_adata.obs['perturbation'] == 'Control'] = substitute[np.random.choice(substitute.shape[0],
                                                                                           (adata.obs['perturbation'] == 'Control').sum(),
                                                                                           replace=True)]

    # Reconstruct the cells
    test_dataset = sf.make_dataset([test_adata], sparse_graph=True, regional_obs=[])
    sf.tools.calc_obs([test_adata], test_dataset, model, get_recon=True)
    sf.tools.calc_obs([train_adata], train_dataset, model, get_recon=True)

    subset = test_adata.obs['t_near_control'] & (~test_adata.obs['t_near_ko'])
    cmp_adata = sc.AnnData(np.vstack([test_adata[subset].obsm['X_recon'],
                                      train_adata[subset].obsm['X_recon']]),
                           var=test_adata.var.copy())

    cmp_adata.obs['grp'] = ['KO'] * subset.sum() + ['WT'] * subset.sum()
    sc.tl.rank_genes_groups(cmp_adata, groupby='grp', method='wilcoxon')
    cmp_df_sf = sc.get.rank_genes_groups_df(cmp_adata, group="KO")
    # cmp_df_sf

    cmp_df_merge = pd.merge(cmp_df_gt, cmp_df_sf, left_on='names', right_on='names', suffixes=['_gt', '_sf'])
    # cmp_df_merge

    cmp_df_merge.to_csv(f"perturbfish-oos/{test_var}.csv")

    fig, ax = plt.subplots(figsize=(1.5, 1.5))
    rho = cmp_df_merge[['logfoldchanges_gt', 'logfoldchanges_sf']].corr(method='spearman').iloc[0, 1]
    cmp_df_merge.plot(x='logfoldchanges_gt', y='logfoldchanges_sf', kind='scatter', s=1, ax=ax)
    ax.set_title(f"{test_var}: {rho: .2f}")
    ax.set_xlabel('Ground-truth change')
    ax.set_ylabel('Predicted change')
    fig.savefig(f"perturbfish-oos/{test_var}.pdf")

    print(f"{test_var}: {rho: .2f}")
    res[test_var] = rho

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

Using [] as regional annotations.

[2025-05-29 06:56:23,215::train::INFO] Epoch 1: train_loss 87.13202
[2025-05-29 06:56:32,706::train::INFO] Epoch 501: train_loss 29.81152
[2025-05-29 06:56:42,214::train::INFO] Epoch 1001: train_loss 28.40072
[2025-05-29 06:56:51,777::train::INFO] Epoch 1501: train_loss 27.06084
[2025-05-29 06:57:01,549::train::INFO] Epoch 2001: train_loss 26.81887
[2025-05-29 06:57:11,132::train::INFO] Epoch 2501: train_loss 25.92734
[2025-05-29 06:57:11,173::train::INFO] Epoch 2503: train_loss 25.25998
[2025-05-29 06:57:11,174::train::INFO] Stopping criterion met.

Using [] as regional annotations.

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

CHUK:  0.05

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

Using [] as regional annotations.

[2025-05-29 06:57:26,496::train::INFO] Epoch 1: train_loss 87.13274
[2025-05-29 06:57:36,053::train::INFO] Epoch 501: train_loss 29.81694
[2025-05-29 06:57:45,623::train::INFO] Epoch 1001: train_loss 28.39801
[2025-05-29 06:57:55,242::train::INFO] Epoch 1501: train_loss 27.16671
[2025-05-29 06:58:05,040::train::INFO] Epoch 2001: train_loss 26.70186
[2025-05-29 06:58:14,658::train::INFO] Epoch 2501: train_loss 25.95170
[2025-05-29 06:58:19,828::train::INFO] Epoch 2770: train_loss 25.45509
[2025-05-29 06:58:19,829::train::INFO] Stopping criterion met.

Using [] as regional annotations.

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

IRAK1:  0.37

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

Using [] as regional annotations.

[2025-05-29 06:58:33,532::train::INFO] Epoch 1: train_loss 87.13361
[2025-05-29 06:58:43,107::train::INFO] Epoch 501: train_loss 29.81323
[2025-05-29 06:58:52,724::train::INFO] Epoch 1001: train_loss 28.38631
[2025-05-29 06:59:02,450::train::INFO] Epoch 1501: train_loss 27.05764
[2025-05-29 06:59:12,191::train::INFO] Epoch 2001: train_loss 26.82323
[2025-05-29 06:59:21,816::train::INFO] Epoch 2501: train_loss 25.95886
[2025-05-29 06:59:21,857::train::INFO] Epoch 2503: train_loss 25.30569
[2025-05-29 06:59:21,858::train::INFO] Stopping criterion met.

Using [] as regional annotations.

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

TRAM1:  0.71

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

Using [] as regional annotations.

[2025-05-29 06:59:35,976::train::INFO] Epoch 1: train_loss 87.13337
[2025-05-29 06:59:45,551::train::INFO] Epoch 501: train_loss 29.81038
[2025-05-29 06:59:55,188::train::INFO] Epoch 1001: train_loss 28.42270
[2025-05-29 07:00:04,907::train::INFO] Epoch 1501: train_loss 27.04279
[2025-05-29 07:00:14,622::train::INFO] Epoch 2001: train_loss 26.81478
[2025-05-29 07:00:24,259::train::INFO] Epoch 2501: train_loss 25.96813
[2025-05-29 07:00:24,300::train::INFO] Epoch 2503: train_loss 25.31587
[2025-05-29 07:00:24,300::train::INFO] Stopping criterion met.

Using [] as regional annotations.

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

LBP:  0.48

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

Using [] as regional annotations.

[2025-05-29 07:00:38,183::train::INFO] Epoch 1: train_loss 87.13642
[2025-05-29 07:00:47,758::train::INFO] Epoch 501: train_loss 29.82709
[2025-05-29 07:00:57,395::train::INFO] Epoch 1001: train_loss 28.40994
[2025-05-29 07:01:07,121::train::INFO] Epoch 1501: train_loss 27.16258
[2025-05-29 07:01:16,834::train::INFO] Epoch 2001: train_loss 26.71648
[2025-05-29 07:01:26,456::train::INFO] Epoch 2501: train_loss 25.93661
[2025-05-29 07:01:36,104::train::INFO] Epoch 3001: train_loss 25.19922
[2025-05-29 07:01:41,093::train::INFO] Epoch 3260: train_loss 25.13174
[2025-05-29 07:01:41,093::train::INFO] Stopping criterion met.

Using [] as regional annotations.

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

IRAK4:  0.22

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

Using [] as regional annotations.

[2025-05-29 07:01:55,812::train::INFO] Epoch 1: train_loss 87.13395
[2025-05-29 07:02:05,416::train::INFO] Epoch 501: train_loss 29.82120
[2025-05-29 07:02:15,252::train::INFO] Epoch 1001: train_loss 28.41534
[2025-05-29 07:02:24,877::train::INFO] Epoch 1501: train_loss 27.05967
[2025-05-29 07:02:34,502::train::INFO] Epoch 2001: train_loss 26.81615
[2025-05-29 07:02:44,141::train::INFO] Epoch 2501: train_loss 26.01119
[2025-05-29 07:02:44,181::train::INFO] Epoch 2503: train_loss 25.34945
[2025-05-29 07:02:44,182::train::INFO] Stopping criterion met.

Using [] as regional annotations.

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

PELI1:  0.70

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

Using [] as regional annotations.

[2025-05-29 07:03:00,284::train::INFO] Epoch 1: train_loss 87.13257
[2025-05-29 07:03:10,009::train::INFO] Epoch 501: train_loss 29.81552
[2025-05-29 07:03:19,583::train::INFO] Epoch 1001: train_loss 28.40267
[2025-05-29 07:03:29,167::train::INFO] Epoch 1501: train_loss 27.09154
[2025-05-29 07:03:38,751::train::INFO] Epoch 2001: train_loss 26.66153
[2025-05-29 07:03:48,338::train::INFO] Epoch 2501: train_loss 25.90419
[2025-05-29 07:03:48,377::train::INFO] Epoch 2503: train_loss 25.23844
[2025-05-29 07:03:48,378::train::INFO] Stopping criterion met.

Using [] as regional annotations.

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

TAB2:  0.34

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

Using [] as regional annotations.

[2025-05-29 07:04:02,295::train::INFO] Epoch 1: train_loss 87.13354
[2025-05-29 07:04:12,108::train::INFO] Epoch 501: train_loss 29.81968
[2025-05-29 07:04:21,656::train::INFO] Epoch 1001: train_loss 28.41932
[2025-05-29 07:04:31,214::train::INFO] Epoch 1501: train_loss 27.05893
[2025-05-29 07:04:40,772::train::INFO] Epoch 2001: train_loss 26.83900
[2025-05-29 07:04:50,348::train::INFO] Epoch 2501: train_loss 25.91199
[2025-05-29 07:04:50,387::train::INFO] Epoch 2503: train_loss 25.25336
[2025-05-29 07:04:50,388::train::INFO] Stopping criterion met.

Using [] as regional annotations.

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

MAP2K2:  0.70

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

Using [] as regional annotations.

[2025-05-29 07:05:03,466::train::INFO] Epoch 1: train_loss 87.13128
[2025-05-29 07:05:13,016::train::INFO] Epoch 501: train_loss 29.82348
[2025-05-29 07:05:22,590::train::INFO] Epoch 1001: train_loss 28.43519
[2025-05-29 07:05:32,161::train::INFO] Epoch 1501: train_loss 26.93173
[2025-05-29 07:05:41,734::train::INFO] Epoch 2001: train_loss 26.87436
[2025-05-29 07:05:51,334::train::INFO] Epoch 2501: train_loss 26.09140
[2025-05-29 07:05:51,375::train::INFO] Epoch 2503: train_loss 25.39158
[2025-05-29 07:05:51,375::train::INFO] Stopping criterion met.

Using [] as regional annotations.

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

MAP2K6:  0.29

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

Using [] as regional annotations.

[2025-05-29 07:06:05,847::train::INFO] Epoch 1: train_loss 87.13795
[2025-05-29 07:06:15,617::train::INFO] Epoch 501: train_loss 29.83226
[2025-05-29 07:06:25,284::train::INFO] Epoch 1001: train_loss 28.41406
[2025-05-29 07:06:34,905::train::INFO] Epoch 1501: train_loss 26.94826
[2025-05-29 07:06:44,526::train::INFO] Epoch 2001: train_loss 26.82843
[2025-05-29 07:06:54,328::train::INFO] Epoch 2501: train_loss 25.91732
[2025-05-29 07:06:54,368::train::INFO] Epoch 2503: train_loss 25.27829
[2025-05-29 07:06:54,369::train::INFO] Stopping criterion met.

Using [] as regional annotations.

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

IRF7:  0.62

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

Using [] as regional annotations.

[2025-05-29 07:07:09,214::train::INFO] Epoch 1: train_loss 87.13075
[2025-05-29 07:07:19,223::train::INFO] Epoch 501: train_loss 29.81181
[2025-05-29 07:07:28,911::train::INFO] Epoch 1001: train_loss 28.39847
[2025-05-29 07:07:38,693::train::INFO] Epoch 1501: train_loss 27.08238
[2025-05-29 07:07:48,353::train::INFO] Epoch 2001: train_loss 26.66256
[2025-05-29 07:07:58,166::train::INFO] Epoch 2501: train_loss 25.95692
[2025-05-29 07:07:58,206::train::INFO] Epoch 2503: train_loss 25.29110
[2025-05-29 07:07:58,207::train::INFO] Stopping criterion met.

Using [] as regional annotations.

C:\Users\lshh\miniconda3\envs\py311_torch211_cuda121\Lib\site-packages\numpy\core\fromnumeric.py:86: FutureWarning: The behavior of DataFrame.sum with axis=None is deprecated, in a future version this will reduce over both axes and return a scalar. To retain the old behavior, pass axis=0 (or do not pass axis)
  return reduction(axis=axis, out=out, **passkwargs)

MYD88:  0.30

pd.Series(res).sort_values(ascending=False).plot(kind='bar', figsize=(2, 1))
plt.ylabel('correlation with\nground truth')
plt.axhline(0.3, ls='--', c='r')

<matplotlib.lines.Line2D at 0x2845bb2b510>

pd.Series(res).to_csv("perturbfish-oos/steamboat_summary.csv")