Train on mouse brain data

This could take several hours. The trained model is provided, which can be used to try the interpretation part without training again.

import os
import sys
sys.path.append("../")
device = "cuda"
import importlib

import scanpy as sc
import squidpy as sq
import pandas as pd
from tqdm.notebook import tqdm
import scipy as sp
import numpy as np
import multiprocessing
import pickle as pkl
import torch
import gc
import sklearn.metrics

import matplotlib.pyplot as plt
import seaborn as sns
import matplotlib

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

pltkw = dict(bbox_inches='tight', transparent=True)

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(

import steamboat as sf
import steamboat.integrated_model
# importlib.reload(spaceformer.benchmarks)

import torch

if False:
    adata = sc.read_h5ad("../data/Ex2_mouse_brain/Zhuang-ABCA-1-raw.h5ad")
    obs = pd.read_csv("../data/Ex2_mouse_brain/label.csv", index_col=0)

    # obs = obs[obs['brain_section_label'].isin(obs['brain_section_label'].value_counts()[:10].index.tolist())]
    # adata = adata[obs.index, :]
    # adata.obs = obs
    # gc.collect()

    adata = adata[obs.index, :]
    gc.collect()
    adata.obs = obs
    gc.collect()

    adata.obsm['spatial'] = adata.obs[['x', 'y']].to_numpy()

    adata.write_h5ad("../data/Ex2_mouse_brain/Zhuang-ABCA-1-labeled.h5ad")

else:
    adata = sc.read_h5ad("../data/Ex2_mouse_brain/Zhuang-ABCA-1-labeled.h5ad")

adata

AnnData object with n_obs × n_vars = 2616328 × 1122
    obs: 'brain_section_label', 'feature_matrix_label', 'donor_label', 'donor_genotype', 'donor_sex', 'cluster_alias', 'x', 'y', 'z', 'subclass_confidence_score', 'cluster_confidence_score', 'high_quality_transfer', 'neurotransmitter', 'class', 'subclass', 'supertype', 'cluster', 'neurotransmitter_color', 'class_color', 'subclass_color', 'supertype_color', 'cluster_color', 'x_ccf', 'y_ccf', 'z_ccf', 'parcellation_index', 'parcellation_organ', 'parcellation_category', 'parcellation_division', 'parcellation_structure', 'parcellation_substructure', 'parcellation_organ_color', 'parcellation_category_color', 'parcellation_division_color', 'parcellation_structure_color', 'parcellation_substructure_color'
    var: 'gene_symbol'
    obsm: 'spatial'

if False:
    n_genes = 100
    np.random.seed(0)
    chosen_gene_mask = np.zeros(adata.shape[1])
    chosen_gene_mask[:n_genes] = 1.
    np.random.shuffle(chosen_gene_mask)
    chosen_gene_mask = chosen_gene_mask > 0.
    adata = adata[:, chosen_gene_mask]
gc.collect()

# Normalizing to median total counts
sc.pp.normalize_total(adata)
# Logarithmize the data
sc.pp.log1p(adata)

import importlib
importlib.reload(steamboat.tools)

adatas = []
for i in adata.obs['brain_section_label'].unique():
    adatas.append(adata[adata.obs['brain_section_label'] == i])
    adatas[-1].obs['global'] = 0

adatas = sf.prep_adatas(adatas, norm=False, log1p=False)
dataset = sf.make_dataset(adatas, sparse_graph=True, regional_obs=['global'])

C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0
C:\Users\lshh\AppData\Local\Temp\ipykernel_165432\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.
  adatas[-1].obs['global'] = 0

Using ['global'] as regional annotations.

sf.set_random_seed(0)
model = sf.Steamboat(adata.var_names.tolist(), n_heads=50, n_scales=3)
model = model.to(device)
# model.load_state_dict(torch.load('backup_mmbrain4/model_weights.pth', weights_only=True))

model.fit(dataset, 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-3, report_per=200, stop_tol=200, device=device)

[2025-02-04 02:31:23,369::train::INFO] Epoch 1: train_loss 153.74057
[2025-02-04 02:49:43,469::train::INFO] Epoch 201: train_loss 81.14060
[2025-02-04 03:08:30,928::train::INFO] Epoch 401: train_loss 79.56349
[2025-02-04 03:26:45,193::train::INFO] Epoch 601: train_loss 79.01989
[2025-02-04 03:44:40,598::train::INFO] Epoch 801: train_loss 78.66930
[2025-02-04 04:02:05,800::train::INFO] Epoch 1001: train_loss 78.65283
[2025-02-04 04:19:31,921::train::INFO] Epoch 1201: train_loss 78.60099
[2025-02-04 04:36:46,760::train::INFO] Epoch 1401: train_loss 78.45714
[2025-02-04 04:53:53,999::train::INFO] Epoch 1601: train_loss 78.38100
[2025-02-04 05:11:00,665::train::INFO] Epoch 1801: train_loss 78.21965
[2025-02-04 05:28:09,231::train::INFO] Epoch 2001: train_loss 78.08940
[2025-02-04 05:45:18,409::train::INFO] Epoch 2201: train_loss 78.01633
[2025-02-04 06:02:26,967::train::INFO] Epoch 2401: train_loss 78.06247
[2025-02-04 06:16:04,476::train::INFO] Epoch 2560: train_loss 78.13670
[2025-02-04 06:16:04,477::train::INFO] Stopping criterion met.

Steamboat(
  (spatial_gather): BilinearAttention(
    (bias): NonNegBias(
      (elu): ELU(alpha=1.0)
    )
    (q): NonNegLinear(
      (elu): ELU(alpha=1.0)
    )
    (k_local): NonNegLinear(
      (elu): ELU(alpha=1.0)
    )
    (k_regionals): ModuleList(
      (0): NonNegLinear(
        (elu): ELU(alpha=1.0)
      )
    )
    (w_ego): NonNegScale(
      (elu): ELU(alpha=1.0)
    )
    (w_local): NonNegScale3(
      (elu): ELU(alpha=1.0)
    )
    (w_global): NonNegScale3(
      (elu): ELU(alpha=1.0)
    )
    (tanh): Tanh()
    (v): NonNegLinear(
      (elu): ELU(alpha=1.0)
    )
    (cosine_similarity): CosineSimilarity()
  )
)

torch.save(model.state_dict(), 'saved_models/mmbrain.pth')

# Ego, local, or global?
import scipy as sp
import scipy.sparse
from tqdm.auto import tqdm
def annotate_adatas(adatas: list[sc.AnnData], dataset, model,
                    device='cuda', get_recon=False):
    """_summary_

    :param adatas: _description_
    :param dataset: _description_
    :param model: _description_
    :param device: _description_, defaults to 'cuda'
    """
    # Safeguards
    assert len(adatas) == len(dataset), "mismatch in lenghths of adatas and dataset"
    for adata, data in zip(adatas, dataset):
        assert adata.shape[0] == data[0].shape[0], f"adata[{i}] has {adata.shape[0]} cells but dataset[{i}] has {data[0].shape[0]}."

    for i, (x, adj_list, regional_xs, regional_adj_lists) in tqdm(enumerate(dataset), total=len(dataset)):
        adj_list = adj_list.squeeze(0).to(device)
        x = x.squeeze(0).to(device)
        regional_adj_lists = [regional_adj_list.to(device) for regional_adj_list in regional_adj_lists]
        regional_xs = [regional_x.to(device) for regional_x in regional_xs]

        with torch.no_grad():
            res, details = model(adj_list, x, x, regional_adj_lists, regional_xs, get_details=True)

            if get_recon:
                adatas[i].obsm['X_recon'] = res.cpu().numpy()

            adatas[i].obsm['q_score'] = details['embq'].cpu().numpy()

            adatas[i].obsm['k_score'] = details['embk'][0].cpu().numpy()
            adatas[i].obsm['k_score_global'] = model.spatial_gather.k_regionals[0](x).cpu().numpy()
            for j, emb in enumerate(details['embk'][1]):
                adatas[i].uns[f'embk_{j}'] = emb.cpu().numpy()
            adatas[i].obsm['attn'] = details['attn'].cpu().numpy()
            adatas[i].obsm['ego_attn'] = details['attnm'][0].cpu().numpy()
            adatas[i].obsm['local_attn'] = details['attnm'][1].cpu().numpy()
            for j, matrix in enumerate(details['attnm'][2]):
                adatas[i].obsm[f'regional_attn_{j}'] = matrix.cpu().numpy()

            # local attention (as graph)
            for j in range(model.spatial_gather.n_heads):
                w = details['attnp'][1].cpu().numpy()[:, j, :].flatten()
                # print(details['attnp'][1].cpu().shape)
                # print(len(w))
                uv = adj_list.cpu().numpy()
                u = uv[0]
                v = uv[1]
                # print(len(u))
                # print(len(v))
                if uv.shape[0] == 3: # masked for unequal neighbors
                    m = (uv[2] > 0)
                    w, u, v = w[m], u[m], v[m]
                adatas[i].obsp[f'local_attn_{j}'] = sp.sparse.csr_matrix((w, (u, v)),
                                                                            shape=(adatas[i].shape[0],
                                                                                adatas[i].shape[0]))

annotate_adatas(adatas, dataset, model, get_recon=True)

v_weights = model.spatial_gather.v.weight.detach().cpu().numpy().sum(axis=0)
adata.uns['steamboat_v_weights'] = v_weights

## For all samples:
def gather_obs(adata, adatas):
    all_embq = []
    all_embk = []
    all_embk_glb = []
    all_ego_attn = []
    all_local_attn = []
    all_global_attn = []
    all_attn = []
    all_recon = []
    for i in range(len(adatas)):
        all_embq.append(adatas[i].obsm['q_score'])
        all_embk.append(adatas[i].obsm['k_score'])
        all_ego_attn.append(adatas[i].obsm['ego_attn'])
        all_local_attn.append(adatas[i].obsm['local_attn'])
        all_global_attn.append(adatas[i].obsm['regional_attn_0'])
        all_recon.append(adatas[i].obsm['X_recon'])
        all_attn.append(adatas[i].obsm['attn'])
        all_embk_glb.append(adatas[i].obsm['k_score_global'])
    adata.obsm['q_score'] = np.vstack(all_embq)
    adata.obsm['k_score'] = np.vstack(all_embk)
    adata.obsm['ego_attn'] = np.vstack(all_ego_attn)
    adata.obsm['local_attn'] = np.vstack(all_local_attn)
    adata.obsm['global_attn'] = np.vstack(all_global_attn)
    adata.obsm['X_recon'] = np.vstack(all_recon)
    adata.obsm['attn'] = np.vstack(all_attn)
    adata.obsm['k_score_global'] = np.vstack(all_embk_glb)
gather_obs(adata, adatas)

adata.write_h5ad("saved_h5ad/mmbrain.h5ad")

for i, adata in enumerate(tqdm(adatas)):
    adata.write_h5ad(f"saved_h5ad/mmbrain_{i}.h5ad")