{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "34225de3-1b63-4b1e-9f3b-661bc57bb979",
   "metadata": {},
   "source": [
    "# Train on mouse brain data\n",
    "\n",
    "This could take several hours. The trained model is provided, which can be used to try the interpretation part without training again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b2d584e4-d3db-4c08-8bf5-76ab340c225a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "sys.path.append(\"../\")\n",
    "device = \"cuda\"\n",
    "import importlib\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fa4a4014-6a35-415d-baa9-738cc509d53f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "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.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import scanpy as sc\n",
    "import squidpy as sq\n",
    "import pandas as pd\n",
    "from tqdm.notebook import tqdm\n",
    "import scipy as sp\n",
    "import numpy as np\n",
    "import multiprocessing\n",
    "import pickle as pkl\n",
    "import torch\n",
    "import gc\n",
    "import sklearn.metrics\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import matplotlib\n",
    "\n",
    "plt.rcParams['pdf.fonttype'] = 42\n",
    "plt.rcParams['mathtext.fontset'] = 'dejavuserif'\n",
    "plt.rcParams['font.family'] = 'arial'\n",
    "\n",
    "pltkw = dict(bbox_inches='tight', transparent=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d606b6a9-01be-44c4-b25a-6d001831dfb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import steamboat as sf\n",
    "import steamboat.integrated_model\n",
    "# importlib.reload(spaceformer.benchmarks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2abd2d15-e58d-4fed-820f-de6b20b75eef",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f139c7d7-44cc-4f39-8567-e602e857041a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AnnData object with n_obs × n_vars = 2616328 × 1122\n",
       "    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'\n",
       "    var: 'gene_symbol'\n",
       "    obsm: 'spatial'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if False:\n",
    "    adata = sc.read_h5ad(\"../data/Ex2_mouse_brain/Zhuang-ABCA-1-raw.h5ad\")\n",
    "    obs = pd.read_csv(\"../data/Ex2_mouse_brain/label.csv\", index_col=0)\n",
    "    \n",
    "    # obs = obs[obs['brain_section_label'].isin(obs['brain_section_label'].value_counts()[:10].index.tolist())]\n",
    "    # adata = adata[obs.index, :]\n",
    "    # adata.obs = obs\n",
    "    # gc.collect()\n",
    "    \n",
    "    adata = adata[obs.index, :]\n",
    "    gc.collect()\n",
    "    adata.obs = obs\n",
    "    gc.collect()\n",
    "    \n",
    "    adata.obsm['spatial'] = adata.obs[['x', 'y']].to_numpy()\n",
    "    \n",
    "    adata.write_h5ad(\"../data/Ex2_mouse_brain/Zhuang-ABCA-1-labeled.h5ad\")\n",
    "\n",
    "else:\n",
    "    adata = sc.read_h5ad(\"../data/Ex2_mouse_brain/Zhuang-ABCA-1-labeled.h5ad\")\n",
    "    \n",
    "adata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cefa36b9-eae5-4f0f-8d17-b05cc49380cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "if False:\n",
    "    n_genes = 100\n",
    "    np.random.seed(0)\n",
    "    chosen_gene_mask = np.zeros(adata.shape[1])\n",
    "    chosen_gene_mask[:n_genes] = 1.\n",
    "    np.random.shuffle(chosen_gene_mask)\n",
    "    chosen_gene_mask = chosen_gene_mask > 0.\n",
    "    adata = adata[:, chosen_gene_mask]\n",
    "gc.collect()\n",
    "\n",
    "# Normalizing to median total counts\n",
    "sc.pp.normalize_total(adata)\n",
    "# Logarithmize the data\n",
    "sc.pp.log1p(adata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a54f5f8-cb39-4346-999b-9e212e6a3d16",
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib\n",
    "importlib.reload(steamboat.tools)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b7d2abae-2ea1-47cd-bebc-b8cd8e60be5f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n",
      "C:\\Users\\lshh\\AppData\\Local\\Temp\\ipykernel_165432\\3785553451.py:4: ImplicitModificationWarning: Trying to modify attribute `.obs` of view, initializing view as actual.\n",
      "  adatas[-1].obs['global'] = 0\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cb52f44b73a245099285cce4c7afff1e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/129 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using ['global'] as regional annotations.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9137b3cae73d44d98b7a26d7a5f884ad",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/129 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "adatas = []\n",
    "for i in adata.obs['brain_section_label'].unique():\n",
    "    adatas.append(adata[adata.obs['brain_section_label'] == i])\n",
    "    adatas[-1].obs['global'] = 0\n",
    "\n",
    "adatas = sf.prep_adatas(adatas, norm=False, log1p=False)\n",
    "dataset = sf.make_dataset(adatas, sparse_graph=True, regional_obs=['global'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "75854034-e22e-478e-b6e0-f182c9707f95",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2025-02-04 02:31:23,369::train::INFO] Epoch 1: train_loss 153.74057\n",
      "[2025-02-04 02:49:43,469::train::INFO] Epoch 201: train_loss 81.14060\n",
      "[2025-02-04 03:08:30,928::train::INFO] Epoch 401: train_loss 79.56349\n",
      "[2025-02-04 03:26:45,193::train::INFO] Epoch 601: train_loss 79.01989\n",
      "[2025-02-04 03:44:40,598::train::INFO] Epoch 801: train_loss 78.66930\n",
      "[2025-02-04 04:02:05,800::train::INFO] Epoch 1001: train_loss 78.65283\n",
      "[2025-02-04 04:19:31,921::train::INFO] Epoch 1201: train_loss 78.60099\n",
      "[2025-02-04 04:36:46,760::train::INFO] Epoch 1401: train_loss 78.45714\n",
      "[2025-02-04 04:53:53,999::train::INFO] Epoch 1601: train_loss 78.38100\n",
      "[2025-02-04 05:11:00,665::train::INFO] Epoch 1801: train_loss 78.21965\n",
      "[2025-02-04 05:28:09,231::train::INFO] Epoch 2001: train_loss 78.08940\n",
      "[2025-02-04 05:45:18,409::train::INFO] Epoch 2201: train_loss 78.01633\n",
      "[2025-02-04 06:02:26,967::train::INFO] Epoch 2401: train_loss 78.06247\n",
      "[2025-02-04 06:16:04,476::train::INFO] Epoch 2560: train_loss 78.13670\n",
      "[2025-02-04 06:16:04,477::train::INFO] Stopping criterion met.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Steamboat(\n",
       "  (spatial_gather): BilinearAttention(\n",
       "    (bias): NonNegBias(\n",
       "      (elu): ELU(alpha=1.0)\n",
       "    )\n",
       "    (q): NonNegLinear(\n",
       "      (elu): ELU(alpha=1.0)\n",
       "    )\n",
       "    (k_local): NonNegLinear(\n",
       "      (elu): ELU(alpha=1.0)\n",
       "    )\n",
       "    (k_regionals): ModuleList(\n",
       "      (0): NonNegLinear(\n",
       "        (elu): ELU(alpha=1.0)\n",
       "      )\n",
       "    )\n",
       "    (w_ego): NonNegScale(\n",
       "      (elu): ELU(alpha=1.0)\n",
       "    )\n",
       "    (w_local): NonNegScale3(\n",
       "      (elu): ELU(alpha=1.0)\n",
       "    )\n",
       "    (w_global): NonNegScale3(\n",
       "      (elu): ELU(alpha=1.0)\n",
       "    )\n",
       "    (tanh): Tanh()\n",
       "    (v): NonNegLinear(\n",
       "      (elu): ELU(alpha=1.0)\n",
       "    )\n",
       "    (cosine_similarity): CosineSimilarity()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sf.set_random_seed(0)\n",
    "model = sf.Steamboat(adata.var_names.tolist(), n_heads=50, n_scales=3)\n",
    "model = model.to(device)\n",
    "# model.load_state_dict(torch.load('backup_mmbrain4/model_weights.pth', weights_only=True))\n",
    "\n",
    "model.fit(dataset, entry_masking_rate=0.1, feature_masking_rate=0.1,\n",
    "          max_epoch=10000, \n",
    "          loss_fun=torch.nn.MSELoss(reduction='sum'),\n",
    "          opt=torch.optim.Adam, opt_args=dict(lr=0.1), stop_eps=1e-3, report_per=200, stop_tol=200, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "6b464f55-99b0-486f-a026-a5cca20df976",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), 'saved_models/mmbrain.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4abd7cc9-8de0-4479-8b5c-de04be84fbc5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "67ce50154dcf459aa72e1c017dd54f72",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/129 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Ego, local, or global?\n",
    "import scipy as sp\n",
    "import scipy.sparse\n",
    "from tqdm.auto import tqdm\n",
    "def annotate_adatas(adatas: list[sc.AnnData], dataset, model, \n",
    "                    device='cuda', get_recon=False):\n",
    "    \"\"\"_summary_\n",
    "\n",
    "    :param adatas: _description_\n",
    "    :param dataset: _description_\n",
    "    :param model: _description_\n",
    "    :param device: _description_, defaults to 'cuda'\n",
    "    \"\"\"\n",
    "    # Safeguards\n",
    "    assert len(adatas) == len(dataset), \"mismatch in lenghths of adatas and dataset\"\n",
    "    for adata, data in zip(adatas, dataset):\n",
    "        assert adata.shape[0] == data[0].shape[0], f\"adata[{i}] has {adata.shape[0]} cells but dataset[{i}] has {data[0].shape[0]}.\"\n",
    "\n",
    "    for i, (x, adj_list, regional_xs, regional_adj_lists) in tqdm(enumerate(dataset), total=len(dataset)):\n",
    "        adj_list = adj_list.squeeze(0).to(device)\n",
    "        x = x.squeeze(0).to(device)\n",
    "        regional_adj_lists = [regional_adj_list.to(device) for regional_adj_list in regional_adj_lists]\n",
    "        regional_xs = [regional_x.to(device) for regional_x in regional_xs]\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            res, details = model(adj_list, x, x, regional_adj_lists, regional_xs, get_details=True)\n",
    "            \n",
    "            if get_recon:\n",
    "                adatas[i].obsm['X_recon'] = res.cpu().numpy()\n",
    "\n",
    "            adatas[i].obsm['q_score'] = details['embq'].cpu().numpy()\n",
    "            \n",
    "            adatas[i].obsm['k_score'] = details['embk'][0].cpu().numpy()\n",
    "            adatas[i].obsm['k_score_global'] = model.spatial_gather.k_regionals[0](x).cpu().numpy()\n",
    "            for j, emb in enumerate(details['embk'][1]):\n",
    "                adatas[i].uns[f'embk_{j}'] = emb.cpu().numpy()\n",
    "            adatas[i].obsm['attn'] = details['attn'].cpu().numpy()\n",
    "            adatas[i].obsm['ego_attn'] = details['attnm'][0].cpu().numpy()\n",
    "            adatas[i].obsm['local_attn'] = details['attnm'][1].cpu().numpy()\n",
    "            for j, matrix in enumerate(details['attnm'][2]):\n",
    "                adatas[i].obsm[f'regional_attn_{j}'] = matrix.cpu().numpy()\n",
    "\n",
    "            # local attention (as graph)\n",
    "            for j in range(model.spatial_gather.n_heads):\n",
    "                w = details['attnp'][1].cpu().numpy()[:, j, :].flatten()\n",
    "                # print(details['attnp'][1].cpu().shape)\n",
    "                # print(len(w))\n",
    "                uv = adj_list.cpu().numpy()\n",
    "                u = uv[0]\n",
    "                v = uv[1]\n",
    "                # print(len(u))\n",
    "                # print(len(v))\n",
    "                if uv.shape[0] == 3: # masked for unequal neighbors\n",
    "                    m = (uv[2] > 0)\n",
    "                    w, u, v = w[m], u[m], v[m]\n",
    "                adatas[i].obsp[f'local_attn_{j}'] = sp.sparse.csr_matrix((w, (u, v)), \n",
    "                                                                            shape=(adatas[i].shape[0], \n",
    "                                                                                adatas[i].shape[0]))\n",
    "            \n",
    "annotate_adatas(adatas, dataset, model, get_recon=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "e39a1e19-6ae0-42e3-97e5-220d324d761f",
   "metadata": {},
   "outputs": [],
   "source": [
    "v_weights = model.spatial_gather.v.weight.detach().cpu().numpy().sum(axis=0)\n",
    "adata.uns['steamboat_v_weights'] = v_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "8a201243-bcca-4cde-9b2e-22ccf90d5560",
   "metadata": {},
   "outputs": [],
   "source": [
    "## For all samples:\n",
    "def gather_obs(adata, adatas):\n",
    "    all_embq = []\n",
    "    all_embk = []\n",
    "    all_embk_glb = []\n",
    "    all_ego_attn = []\n",
    "    all_local_attn = []\n",
    "    all_global_attn = []\n",
    "    all_attn = []\n",
    "    all_recon = []\n",
    "    for i in range(len(adatas)):\n",
    "        all_embq.append(adatas[i].obsm['q_score'])\n",
    "        all_embk.append(adatas[i].obsm['k_score'])\n",
    "        all_ego_attn.append(adatas[i].obsm['ego_attn'])\n",
    "        all_local_attn.append(adatas[i].obsm['local_attn'])\n",
    "        all_global_attn.append(adatas[i].obsm['regional_attn_0'])\n",
    "        all_recon.append(adatas[i].obsm['X_recon'])\n",
    "        all_attn.append(adatas[i].obsm['attn'])\n",
    "        all_embk_glb.append(adatas[i].obsm['k_score_global'])\n",
    "    adata.obsm['q_score'] = np.vstack(all_embq)\n",
    "    adata.obsm['k_score'] = np.vstack(all_embk)\n",
    "    adata.obsm['ego_attn'] = np.vstack(all_ego_attn)\n",
    "    adata.obsm['local_attn'] = np.vstack(all_local_attn)\n",
    "    adata.obsm['global_attn'] = np.vstack(all_global_attn)\n",
    "    adata.obsm['X_recon'] = np.vstack(all_recon)\n",
    "    adata.obsm['attn'] = np.vstack(all_attn)\n",
    "    adata.obsm['k_score_global'] = np.vstack(all_embk_glb)\n",
    "gather_obs(adata, adatas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "bc2364f1-bad5-4ba8-bc1c-8ce87b460650",
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.write_h5ad(\"saved_h5ad/mmbrain.h5ad\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "5fa04e78-bd41-415d-ab68-0749c2eda6a9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ac03c39d52f34dbe88c07a4874d7d08e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/129 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i, adata in enumerate(tqdm(adatas)):\n",
    "    adata.write_h5ad(f\"saved_h5ad/mmbrain_{i}.h5ad\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}