terra.get_average_gene_embed

terra.get_average_gene_embed#

terra.get_average_gene_embed(dataset, model_folder_path, emb_layer=None, cell_gene_ensembl_id=[], neighborhood_gene_ensembl_id=[], batch_size=128, pin_memory=False, num_workers=12, include_spatial_cell_emb=True)#

Retrieve average gene embeddings for each gene per dataset.

Parameters:
  • dataset (Dataset) – Tokenized Hugging Face dataset.

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

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

  • cell_gene_ensembl_id (list (default: [])) – List with gene IDs for which cell gene embeddings will be retrieved.

  • neighborhood_gene_ensembl_id (list (default: [])) – List with gene IDs for which neighborhood gene embeddings will be retrieved.

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

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

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

  • include_spatial_cell_emb (bool (default: True)) – If True, also return average gene embeddings for the spatially contextualized cell embedding.

Return type:

dict[str, ndarray]

Returns:

-output_average_gene_embed (dict) Dictionary mapping average embedding/statistics names to numpy arrays.