Spatial niche tokenizer

import scanpy as sc
import anndata as ad
import pandas as pd
import numpy as np
import math
import numba
from scipy.sparse import issparse
from tqdm import tqdm
from typing import List, Dict, Optional, Union
import pyarrow as pa
import pyarrow.parquet as pq
import os
import sys
from sklearn.neighbors import NearestNeighbors
import os
import zarr
import numpy as np
from tqdm import tqdm
sys.path.append('/data_d/WTG/SPpretrain/src')
from tokenizer import FineTuneTokenizer
from constants import PAD_TOKEN, CLS_TOKEN, PROTEIN_TOKEN_BASE, side_maps, protein_id_map, SIDE_COLUMNS, all_proteins, DataSchema
pd.set_option('display.max_columns', None)

Read in data and set metadata.

We first load the spatial proteomics dataset stored in the AnnData format. In this example, we use the CODEX intestine dataset.

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.

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.

adata = sc.read_h5ad('/data_d/WTG/SPpretrain/fine_tune/IMC_COAD_neigh/adata_with_clinical.h5ad')

adata.obsm['spatial'] = adata.obs[['AreaShape_Center_X', 'AreaShape_Center_Y']].values

schema = DataSchema()
schema.describe()

Available fields:

[Technology]
  CODEX, CyCIF, IMC

[Tissue_type]
  Bladder, Brain, Breast, Inflammatory, Intestine, Kidney, Liver, Lung, Lymph_node, Ovarian, Pancreas, Placenta, Salivary gland, Skin, Soft tissue, Spleen, Thymus, Tonsil

[Health_status]
  Disease, Health

[Disease_type]
  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

Note: if a value not in the allowed options is filled, it will be automatically replaced with 'PAD'.

schema.set(adata,
           Technology='IMC',
           Health_status='Disease',
           Tissue_type='Intestine',
           Disease_type='COAD',)
adata.write_h5ad('/path/to/adata_with_clinical.h5ad')

Tokenization

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.

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.

• 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.

tokenizer = FineTuneTokenizer(
    all_proteins=all_proteins,
    protein_id_map=protein_id_map,
    side_maps=side_maps,
    protein_token_base=PROTEIN_TOKEN_BASE
)

tokenizer.collect_data(['/path/to/adata_with_clinical.h5ad'], cell_type_col='cell_type')

Collecting adata: 100%|██████████| 1/1 [00:00<00:00,  1.91it/s]

tokenized_list = tokenizer.tokenize()

100%|██████████| 1/1 [00:08<00:00,  8.42s/it]

Calculate neighborhood cell type composition

To characterize the local spatial microenvironment of each cell, we calculate the neighborhood cell type composition within each sample.

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.

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.

The parameter k controls the size of the neighborhood. Larger values capture broader spatial patterns, while smaller values focus on more local cellular interactions.

import numpy as np
from sklearn.neighbors import NearestNeighbors

cell_types = tokenizer.cell_type_list[0]  # [n_cells,]
coords = adata.obsm['spatial']            # [n_cells, 2]
sample_ids = adata.obs['alt_identifier'].values

num_types = cell_types.max() + 1
k = 30

neighbor_ratio = np.zeros((adata.n_obs, num_types))

for sample_id in np.unique(sample_ids):

    idx = np.where(sample_ids == sample_id)[0]

    coords_sub = coords[idx]
    cell_types_sub = cell_types[idx]

    k_use = min(k, len(idx) - 1)

    if k_use <= 0:
        continue

    # one-hot
    one_hot_types = np.eye(num_types)[cell_types_sub]

    # kNN
    nbrs = NearestNeighbors(
        n_neighbors=k_use + 1,
        algorithm='auto'
    ).fit(coords_sub)

    distances, indices = nbrs.kneighbors(coords_sub)

    neighbor_indices = indices[:, 1:]

    neighbor_one_hot = one_hot_types[neighbor_indices]
    neighbor_sum = neighbor_one_hot.sum(axis=1)

    ratio_sub = neighbor_sum / neighbor_sum.sum(
        axis=1,
        keepdims=True
    )

    neighbor_ratio[idx] = ratio_sub
print(neighbor_ratio.shape)
neighbor_ratio_list = [neighbor_ratio]

(162337, 10)

Save tokenized data with neighborhood information

After calculating the neighborhood cell type composition, the tokenized protein sequences and spatial neighborhood features are stored together in a Zarr file.

The output Zarr file contains two arrays:

• tokens: tokenized protein sequences generated by FineTuneTokenizer, with shape [N_cells, token_dim].

• 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.

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.

save_dir = '/path/to/zarr'
os.makedirs(save_dir, exist_ok=True)
zarr_path = save_dir
root = zarr.open(zarr_path, mode='w')
token_dim = tokenized_list[0].shape[1]
num_types = neighbor_ratio_list[0].shape[1]
tokens_z = root.create_dataset(
    name='tokens',
    shape=(0, token_dim),
    chunks=(10000, token_dim),
    dtype='int64',
    maxshape=(None, token_dim)
)

neighbor_ratio_z = root.create_dataset(
    name='neighbor_ratio',
    shape=(0, num_types),
    chunks=(10000, num_types),
    dtype='float32',
    maxshape=(None, num_types)
)
for tokens, neigh in tqdm(
    zip(tokenized_list, neighbor_ratio_list),
    total=len(tokenized_list),
    desc="Writing Zarr (tokens + neighbor_ratio)"
):
    # tokens: [n_cells, token_dim]
    # neigh:  [n_cells, num_types]
    tokens_z.append(tokens.astype(np.int64))
    neighbor_ratio_z.append(neigh.astype(np.float32))
print("Zarr saved.")
print("tokens shape:", root['tokens'].shape)
print("neighbor_ratio shape:", root['neighbor_ratio'].shape)

Writing Zarr (tokens + neighbor_ratio):   0%|          | 0/1 [00:00<?, ?it/s]Writing Zarr (tokens + neighbor_ratio): 100%|██████████| 1/1 [00:01<00:00,  1.08s/it]

Zarr saved.
tokens shape: (162337, 239)
neighbor_ratio shape: (162337, 10)