Cell type 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
import zarr
import seaborn as sns
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

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('/path/to/Nature_CODEX_Intestine/adata.h5ad')
adata = adata[adata.obs['array'] == 'B009B'].copy()

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='CODEX',
           Health_status='Health',
           Tissue_type='Intestine',
           Disease_type='Normal',)
adata.write_h5ad('/path/to/B009B.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/B009B.h5ad'], cell_type_col="Cell Type")

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

tokenized_list = tokenizer.tokenize()

100%|██████████| 1/1 [00:19<00:00, 19.19s/it]

Save tokenized data

After tokenization, the generated protein token sequences and corresponding cell type labels are stored in Zarr format for efficient access during model fine-tuning.

The saving step only requires specifying an output directory. Two arrays are created:

• tokens: the tokenized protein sequences for each cell, stored as an integer matrix with shape [N_cells, token_dim].

• label: the encoded cell type labels corresponding to each cell, stored as a one-dimensional integer array.

Zarr is used as the storage format because it supports chunked and scalable access, which is suitable for large-scale spatial proteomics datasets.

save_dir = '/path/to/zarr/dir/'
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]

# tokens: [N, token_dim]
tokens_z = root.create_dataset(
    name='tokens',
    shape=(0, token_dim),
    chunks=(10000, token_dim),
    dtype='int64',
    maxshape=(None, token_dim)
)

# label: [N]
label_z = root.create_dataset(
    name='label',
    shape=(0,),
    chunks=(10000,),
    dtype='int64',
    maxshape=(None,)
)

for tokens, neigh in tqdm(
    zip(tokenized_list, tokenizer.cell_type_list),
    total=len(tokenized_list),
    desc="Writing Zarr (tokens + label)"
):
    # tokens: [n_cells, token_dim]
    # neigh:  [n_cells] 或 [n_cells, 1]

    tokens_z.append(tokens.astype(np.int64))
    label_z.append(neigh.reshape(-1).astype(np.int64))

print("Zarr saved.")
print("tokens shape:", root['tokens'].shape)
print("label shape:", root['label'].shape)

/data_d/WTG/mamba/envs/spaproLLM/lib/python3.9/site-packages/zarr/creation.py:190: UserWarning: ignoring keyword argument 'maxshape'
  compressor, fill_value = _kwargs_compat(compressor, fill_value, kwargs)
Writing Zarr (tokens + label): 100%|██████████| 1/1 [00:03<00:00,  3.54s/it]

Zarr saved.
tokens shape: (214837, 239)
label shape: (214837,)