cellflow.data.DataManager

class cellflow.data.DataManager(adata, sample_rep, control_key, perturbation_covariates=None, perturbation_covariate_reps=None, sample_covariates=None, sample_covariate_reps=None, split_covariates=None, max_combination_length=None, null_value=0.0, primary_group=None)[source]

Data manager for handling perturbation data.

Parameters:
  • adata (AnnData) – An AnnData object.

  • covariate_encoder – Encoder for the primary covariate.

  • categorical – Whether the primary covariate is categorical.

  • max_combination_length (int | None) – Maximum number of combinations of primary perturbation_covariates.

  • sample_rep (str | dict[str, str]) – Key in obsm where the sample representation is stored or 'X' to use X.

  • covariate_data – Dataframe with covariates. If None, obs is used.

  • condition_id_key – Key in obs that defines the condition id.

  • perturbation_covariates (dict[str, Sequence[str]] | None) – A dictionary where the keys indicate the name of the covariate group and the values are keys in uns. The corresponding columns should be either boolean (presence/abscence of the perturbation) or numeric (concentration or magnitude of the perturbation). If multiple groups are provided, the first is interpreted as the primary perturbation and the others as covariates corresponding to these perturbations, e.g. {"drug":("drugA", "drugB"), "time":("drugA_time", "drugB_time")}.

  • perturbation_covariate_reps (dict[str, str] | None) – A dictionary where the keys indicate the name of the covariate group and the values are keys in uns storing a dictionary with the representation of the covariates. E.g. {"drug":"drug_embeddings"} with adata.uns["drug_embeddings"] = {"drugA": np.array, "drugB": np.array}.

  • sample_covariates (Sequence[str] | None) – Keys in obs indicating sample covatiates to be taken into account for training and prediction, e.g. ["age", "cell_type"].

  • sample_covariate_reps (dict[str, str] | None) – A dictionary where the keys indicate the name of the covariate group and the values are keys in uns storing a dictionary with the representation of the covariates. E.g. {"cell_type": "cell_type_embeddings"} with adata.uns["cell_type_embeddings"] = {"cell_typeA": np.array, "cell_typeB": np.array}.

  • split_covariates (Sequence[str] | None) – Covariates in obs to split all control cells into different control populations. The perturbed cells are also split according to these columns, but if these covariates should also be encoded in the model, the corresponding column should also be used in perturbation_covariates or sample_covariates.

  • null_value (float) – Value to use for padding to max_combination_length.

  • primary_group (str | None) – Name of the primary perturbation group. If not provided, the first group in perturbation_covariates but perturbation_covariates should be an ordered dict.

  • control_key (str)

Methods

get_condition_data(covariate_data[, ...])

Get condition data for the model.

get_prediction_data(adata, sample_rep, ...)

Get predictions for control cells.

get_train_data(adata)

Get training data for the model.

get_validation_data(adata[, ...])

Get validation data for the model.

Attributes

adata

An AnnData object used for instantiating the DataManager.

control_key

Boolean key in obs indicating whether belongs to control group.

covar_to_idx

add description

covariate_reps

Dictionary which stores representation of covariates, i.e. the union of sample_covariate_reps and perturbation_covariate_reps.

is_categorical

Whether the primary covariate is categorical.

linked_perturb_covars

Dictionary with keys indicating the name of the primary covariate and values are dictionaries with keys indicating the name of the linked covariate group and values are the linked covariates.

max_combination_length

Maximum combination length of perturbation covariates.

null_value

Value to use for padding to max_combination_length.

perturb_covar_keys

List of all perturbation covariates.

perturbation_covariate_reps

Dictionary with keys indicating the name of the covariate group and values are keys in uns storing a dictionary with the representation of the covariates.

perturbation_covariates

Dictionary with keys indicating the name of the covariate group and values are keys in obs which together define the covariates.

primary_one_hot_encoder

One-hot encoder for the primary covariate.

sample_covariate_reps

Dictionary with keys indicating the name of the sample covariate group and values are keys in uns storing a dictionary with the representation of the sample covariates.

sample_covariates

Keys in obs indicating which sample the cell belongs to (e.g. cell line).

sample_rep

Key of the sample representation.

split_covariates

Covariates in obs to split all control cells into different control populations.