{ "cells": [ { "cell_type": "markdown", "id": "00b7086c", "metadata": {}, "source": [ "# Spatial niche tokenizer" ] }, { "cell_type": "code", "execution_count": 9, "id": "cfcfecba", "metadata": {}, "outputs": [], "source": [ "import scanpy as sc\n", "import anndata as ad\n", "import pandas as pd\n", "import numpy as np\n", "import math\n", "import numba\n", "from scipy.sparse import issparse\n", "from tqdm import tqdm\n", "from typing import List, Dict, Optional, Union\n", "import pyarrow as pa\n", "import pyarrow.parquet as pq\n", "import os\n", "import sys\n", "from sklearn.neighbors import NearestNeighbors\n", "import os\n", "import zarr\n", "import numpy as np\n", "from tqdm import tqdm\n", "sys.path.append('/data_d/WTG/SPpretrain/src')\n", "from tokenizer import FineTuneTokenizer\n", "from constants import PAD_TOKEN, CLS_TOKEN, PROTEIN_TOKEN_BASE, side_maps, protein_id_map, SIDE_COLUMNS, all_proteins, DataSchema\n", "pd.set_option('display.max_columns', None)" ] }, { "cell_type": "markdown", "id": "dcb832f5", "metadata": {}, "source": [ "## Read in data and set metadata.\n", "\n", "\n", "We first load the spatial proteomics dataset stored in the AnnData format. In this example, we use the CODEX intestine dataset.\n", "\n", "The `DataSchema` object is then used to inspect and standardize the metadata fields required by Spatium. It defines the allowed categories for major metadata attributes, including technology, tissue type, health status, and disease type. Values outside the predefined vocabulary will be automatically replaced with `PAD`.\n", "\n", "After checking the available metadata options, we assign the corresponding metadata labels to the selected dataset using `schema.set()`. The processed AnnData object is finally saved for tokenization.\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "acb5035b", "metadata": {}, "outputs": [], "source": [ "adata = sc.read_h5ad('/data_d/WTG/SPpretrain/fine_tune/IMC_COAD_neigh/adata_with_clinical.h5ad')" ] }, { "cell_type": "code", "execution_count": 4, "id": "44c736c9", "metadata": {}, "outputs": [], "source": [ "adata.obsm['spatial'] = adata.obs[['AreaShape_Center_X', 'AreaShape_Center_Y']].values\n" ] }, { "cell_type": "code", "execution_count": 10, "id": "91e7b884", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Available fields:\n", "\n", "[Technology]\n", " CODEX, CyCIF, IMC\n", "\n", "[Tissue_type]\n", " Bladder, Brain, Breast, Inflammatory, Intestine, Kidney, Liver, Lung, Lymph_node, Ovarian, Pancreas, Placenta, Salivary gland, Skin, Soft tissue, Spleen, Thymus, Tonsil\n", "\n", "[Health_status]\n", " Disease, Health\n", "\n", "[Disease_type]\n", " Bladder_cancer, Breast_cancer, COAD, CRC, Covid-19, Dermal sarcoma, Diabetes, Glioblastoma, HCC, HGSC, Leiomyosarcoma, Lung_cancer, Melanoma, Myxofibrosarcoma, Normal, Renal_Cell_Carcinoma, Salivary cystadenoma, Squamous cell carcinoma\n", "\n", "Note: if a value not in the allowed options is filled, it will be automatically replaced with 'PAD'.\n" ] } ], "source": [ "schema = DataSchema()\n", "schema.describe()" ] }, { "cell_type": "code", "execution_count": null, "id": "f514b0c6", "metadata": {}, "outputs": [], "source": [ "schema.set(adata,\n", " Technology='IMC',\n", " Health_status='Disease',\n", " Tissue_type='Intestine',\n", " Disease_type='COAD',)\n", "adata.write_h5ad('/path/to/adata_with_clinical.h5ad')\n" ] }, { "cell_type": "markdown", "id": "87d9f4b7", "metadata": {}, "source": [ "## Tokenization\n", "\n", "\n", "Spatium converts spatial proteomics data into protein token sequences using the `FineTuneTokenizer`. The tokenizer requires a standardized protein vocabulary, protein ID mapping, and metadata mappings generated during preprocessing.\n", "\n", "Please note that before tokenization, the processed AnnData object with standardized metadata should be saved to disk. The tokenizer takes the path of the saved AnnData file rather than the in-memory object. Therefore, metadata annotation using `DataSchema()` should be completed before this step.\n", "\n", "\n", "\n", "- `cell_type_col`: For downstream tasks involving cell type annotation, the column name containing cell type labels in `adata.obs` should be provided through the `cell_type_col` parameter. The tokenizer will use this information to associate each spatial cell with its corresponding cell type label during data preparation." ] }, { "cell_type": "code", "execution_count": null, "id": "1eecde96", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "Collecting adata: 100%|██████████| 1/1 [00:00<00:00, 1.91it/s]\n" ] } ], "source": [ "tokenizer = FineTuneTokenizer(\n", " all_proteins=all_proteins,\n", " protein_id_map=protein_id_map,\n", " side_maps=side_maps,\n", " protein_token_base=PROTEIN_TOKEN_BASE\n", ")\n", "\n", "tokenizer.collect_data(['/path/to/adata_with_clinical.h5ad'], cell_type_col='cell_type')" ] }, { "cell_type": "code", "execution_count": 6, "id": "32ace917", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "100%|██████████| 1/1 [00:08<00:00, 8.42s/it]\n" ] } ], "source": [ "tokenized_list = tokenizer.tokenize()" ] }, { "cell_type": "markdown", "id": "4a468158", "metadata": {}, "source": [ "## Calculate neighborhood cell type composition\n", "\n", "To characterize the local spatial microenvironment of each cell, we calculate the neighborhood cell type composition within each sample.\n", "\n", "For every cell, the `k` nearest neighboring cells are identified based on spatial coordinates. The cell types of these neighbors are then summarized as a proportion vector, where each element represents the fraction of a specific cell type within the local neighborhood.\n", "\n", "The cell type labels are obtained from `tokenizer.cell_type_list[0]`, which contains the encoded cell type identity for each cell. The analysis is performed independently within each sample (`alt_identifier`) to avoid mixing cells from different spatial regions.\n", "\n", "The parameter `k` controls the size of the neighborhood. Larger values capture broader spatial patterns, while smaller values focus on more local cellular interactions." ] }, { "cell_type": "code", "execution_count": 13, "id": "53685256", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "(162337, 10)\n" ] } ], "source": [ "import numpy as np\n", "from sklearn.neighbors import NearestNeighbors\n", "\n", "cell_types = tokenizer.cell_type_list[0] # [n_cells,]\n", "coords = adata.obsm['spatial'] # [n_cells, 2]\n", "sample_ids = adata.obs['alt_identifier'].values\n", "\n", "num_types = cell_types.max() + 1\n", "k = 30\n", "\n", "neighbor_ratio = np.zeros((adata.n_obs, num_types))\n", "\n", "for sample_id in np.unique(sample_ids):\n", "\n", " idx = np.where(sample_ids == sample_id)[0]\n", "\n", " coords_sub = coords[idx]\n", " cell_types_sub = cell_types[idx]\n", "\n", " k_use = min(k, len(idx) - 1)\n", "\n", " if k_use <= 0:\n", " continue\n", "\n", " # one-hot\n", " one_hot_types = np.eye(num_types)[cell_types_sub]\n", "\n", " # kNN\n", " nbrs = NearestNeighbors(\n", " n_neighbors=k_use + 1,\n", " algorithm='auto'\n", " ).fit(coords_sub)\n", "\n", " distances, indices = nbrs.kneighbors(coords_sub)\n", "\n", " neighbor_indices = indices[:, 1:]\n", "\n", " neighbor_one_hot = one_hot_types[neighbor_indices]\n", " neighbor_sum = neighbor_one_hot.sum(axis=1)\n", "\n", " ratio_sub = neighbor_sum / neighbor_sum.sum(\n", " axis=1,\n", " keepdims=True\n", " )\n", "\n", " neighbor_ratio[idx] = ratio_sub\n", "print(neighbor_ratio.shape)\n", "neighbor_ratio_list = [neighbor_ratio]" ] }, { "cell_type": "markdown", "id": "3ea0cb9e", "metadata": {}, "source": [ "## Save tokenized data with neighborhood information\n", "\n", "After calculating the neighborhood cell type composition, the tokenized protein sequences and spatial neighborhood features are stored together in a Zarr file.\n", "\n", "The output Zarr file contains two arrays:\n", "\n", "- `tokens`: tokenized protein sequences generated by `FineTuneTokenizer`, with shape `[N_cells, token_dim]`.\n", "- `neighbor_ratio`: neighborhood cell type composition vectors, with shape `[N_cells, N_cell_types]`. Each row represents the relative proportion of different cell types among the spatial neighbors of a given cell.\n", "\n", "The output directory only needs to be specified through `save_dir`. Zarr is used to support efficient chunked storage and scalable loading for large spatial proteomics datasets." ] }, { "cell_type": "code", "execution_count": null, "id": "a70e07fd", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "Writing Zarr (tokens + neighbor_ratio): 0%| | 0/1 [00:00