terra.harmonize_tokenize_embed_pipeline

terra.harmonize_tokenize_embed_pipeline#

terra.harmonize_tokenize_embed_pipeline(adata, sample_key, batch_key, model_folder_path, cache_directory_path, gene_mapping_dict_file_path=None, gene_occurrence_count_file_path=None, gene_occurrence_count_filter_value=10, ensembl_release=111, species='human', min_cells_per_gene=0, min_genes_per_cell=0, gene_perturb_df=None, nproc=4, processing_mode='parallel', harmonized_adata_save_path=None, save_dataset_path=None, num_shards=32, emb_layer=None, agg_excluded_genes=None, top_k=None, batch_size=128, pin_memory=False, num_workers=12, use_generator=True, add_neigh_cell_ids=False, ignore_spc_tokens=True, agg_type='avg')#

Harmonize, tokenize and embed an AnnData object.

Parameters:

  • adata (AnnData) – An unharmonized AnnData object to be tokenized.

  • sample_key (str | None) – Key in adata.obs where the sample information is stored.

  • batch_key (str) – Key in adata.obs where the batch identifier in the harmonized AnnData will be stored.

  • model_folder_path (str) – Path to the folder containing the model config, token dictionary, and normalization factors.

  • cache_directory_path (str) – Path where the cache is stored during dataset creation.

  • gene_mapping_dict_file_path (Optional[str] (default: None)) – Path to a pickled gene-name-to-ensembl-ID mapping dictionary used during harmonization. If None, defaults to the model bundle’s ensembl_dictionary.pkl when present, otherwise falls back to an Ensembl release lookup.

  • gene_occurrence_count_file_path (Optional[str] (default: None)) – Path to a pickled dictionary of per-gene pretraining occurrence counts. If None, defaults to the model bundle’s gene_count_dictionary.pkl when present; otherwise the occurrence filter is skipped.

  • gene_occurrence_count_filter_value (int (default: 10)) – Minimum number of pretraining occurrences a gene must have to be kept.

  • ensembl_release (int (default: 111)) – Ensembl release used to retrieve ensembl IDs when no mapping dictionary is provided.

  • species (str (default: 'human')) – Species passed to the Ensembl release lookup.

  • min_cells_per_gene (int (default: 0)) – Minimum number of cells a gene must be expressed in to be kept during harmonization.

  • min_genes_per_cell (int (default: 0)) – Minimum number of genes a cell must express to be kept during harmonization.

  • gene_perturb_df (Optional[DataFrame] (default: None)) –

    DataFrame with perturbation data, e.g.:

    gene_perturb_df = pd.DataFrame({
    'ensembl_id': ['ENSG00000169194', 'ENSG00000131724'],
    'target': ['neighborhood', 'cell'],
    'perturbation_type': ['foldchange', 'knockout'],
    'foldchange': [0.5, np.nan],
})

  • nproc (int (default: 4)) – Number of processes used for tokenization.

  • processing_mode (Literal['sequential', 'parallel'] (default: 'parallel')) – Mode of processing used for tokenization.

  • harmonized_adata_save_path (Union[Path, str, None] (default: None)) – If specified, harmonized AnnData is written to disk at this path.

  • save_dataset_path (Union[Path, str, None] (default: None)) – If specified, Hugging Face dataset is written to disk at this path.

  • num_shards (int (default: 32)) – Number of shards with which Hugging Face dataset is saved.

  • emb_layer (Optional[int] (default: None)) – Layer for which to retrieve the embedding.

  • agg_excluded_genes (Optional[list[int]] (default: None)) – List of tokens to be excluded from the aggregation.

  • top_k (Optional[int] (default: None)) – Include only top_k genes in aggregation.

  • batch_size (int (default: 128)) – Dataloader param.

  • pin_memory (bool (default: False)) – Dataloader param.

  • num_workers (int (default: 12)) – Number of workers used for model inference.

  • use_generator (bool (default: True)) – Whether to use generator for dataset creation.

  • add_neigh_cell_ids (bool (default: False)) – Whether neighbor cell IDs should be stored in tokenized data (used for perturbations).

  • ignore_spc_tokens (bool (default: True)) – Whether to ignore special tokens when retrieving layer embeddings.

  • agg_type (Literal['avg', 'softmax'] (default: 'avg')) – How gene embeddings are aggregated into cell and neighborhood embeddings ('avg' or 'softmax').

Returns:

  • adata (AnnData) – A harmonized AnnData object with embeddings stored in adata.obsm.

Contents