terra.perturb_dataset

terra.perturb_dataset#

terra.perturb_dataset(dataset, perturb_df, model_folder_path, seq_len_cell=256, n_segments=11, nproc=1, batch_size=1000, keep_in_memory=False, return_only_perturbed_cells=False, pad_gene_tokens=True, adjust_positions=False)#

Perturb a tokenized Hugging Face dataset according to a perturbation dataframe.

The perturbations described in perturb_df (gene knockouts or fold changes, applied to a cell’s own gene tokens or to its neighborhood) are applied to the gene tokens/expression of the dataset. Perturbations on all cells are applied batch-wise; perturbations targeting specific cells are matched by cell_id (and, for neighborhood targets, by the stored neighbor IDs).

Parameters:

  • dataset (Dataset) – Tokenized Hugging Face dataset to perturb. Perturbing the neighborhood of specific cells requires the dataset to have been tokenized with add_neigh_cell_ids=True (so neighbor IDs are stored).

  • perturb_df (DataFrame) – Dataframe describing the perturbations. Expected columns are perturbed_cell_id (a cell ID or "all"), perturbed_ensembl_id (an ensembl ID or "all"), perturbation_target ("cell" or "neighborhood") and perturbation_type ("knockout" or "foldchange"), plus a foldchange value for fold-change rows.

  • model_folder_path (str) – Path to the folder containing the model config and token dictionary; used to map ensembl IDs to token IDs.

  • seq_len_cell (int (default: 256)) – Number of cell gene tokens (excluding neighborhood gene tokens).

  • n_segments (int (default: 11)) – Number of segments (the index cell plus its neighbors); used when re-packing positions after perturbation.

  • nproc (int (default: 1)) – Number of processes used to map the perturbation over the dataset. Defaults to 1 (single process). Values >1 use multiprocessing, which can deadlock once torch/CUDA has been initialized in the process (the worker fork inherits a broken CUDA state) and adds dataset-serialization overhead; the perturbation itself is fast single-process.

  • batch_size (int (default: 1000)) – Number of rows per batch passed to the mapping function.

  • keep_in_memory (bool (default: False)) – Whether to keep the perturbed dataset in memory during mapping.

  • return_only_perturbed_cells (bool (default: False)) – If True, return only the cells whose tokens the perturbations actually edit (computed from the unperturbed dataset).

  • pad_gene_tokens (bool (default: True)) – If True, also zero out the gene tokens (not just the expression) of knocked-out genes.

  • adjust_positions (bool (default: False)) – If True, re-pack each segment so padded (zeroed) gene positions are moved to the end after perturbation.

Returns:

  • dataset (Dataset) – The perturbed Hugging Face Dataset, filtered to the affected cells if return_only_perturbed_cells is True.

Contents