cellflow.model.CellFlow.prepare_model

CellFlow.prepare_model(condition_mode='deterministic', regularization=0.0, pooling='attention_token', pooling_kwargs=mappingproxy({}), layers_before_pool=Field(name=None, type=None, default=<dataclasses._MISSING_TYPE object>, default_factory=<function CellFlow.<lambda>>, init=True, repr=True, hash=None, compare=True, metadata=mappingproxy({}), kw_only=<dataclasses._MISSING_TYPE object>, _field_type=None), layers_after_pool=Field(name=None, type=None, default=<dataclasses._MISSING_TYPE object>, default_factory=<function CellFlow.<lambda>>, init=True, repr=True, hash=None, compare=True, metadata=mappingproxy({}), kw_only=<dataclasses._MISSING_TYPE object>, _field_type=None), condition_embedding_dim=256, cond_output_dropout=0.9, condition_encoder_kwargs=None, pool_sample_covariates=True, time_freqs=1024, time_max_period=10000, time_encoder_dims=(2048, 2048, 2048), time_encoder_dropout=0.0, hidden_dims=(2048, 2048, 2048), hidden_dropout=0.0, conditioning='concatenation', conditioning_kwargs=Field(name=None, type=None, default=<dataclasses._MISSING_TYPE object>, default_factory=<function CellFlow.<lambda>>, init=True, repr=True, hash=None, compare=True, metadata=mappingproxy({}), kw_only=<dataclasses._MISSING_TYPE object>, _field_type=None), decoder_dims=(4096, 4096, 4096), decoder_dropout=0.0, vf_act_fn=<PjitFunction of <function silu>>, vf_kwargs=None, probability_path=None, match_fn=<function match_linear>, optimizer=<optax.transforms._accumulation.MultiSteps object>, solver_kwargs=None, layer_norm_before_concatenation=False, linear_projection_before_concatenation=False, seed=0)

Prepare the model for training.

This function sets up the neural network architecture and specificities of the solver. When solver is an instance of cellflow.solvers._genot.GENOT, the following arguments have to be passed to 'condition_encoder_kwargs':

Parameters:

• condition_mode (Literal['deterministic', 'stochastic']) –

  Mode of the encoder, should be one of:

  • 'deterministic': Learns condition encoding point-wise.

  • 'stochastic': Learns a Gaussian distribution for representing conditions.

• regularization (float) –

  Regularization strength in the latent space:

  • For deterministic mode, it is the strength of the L2 regularization.

  • For stochastic mode, it is the strength of the VAE regularization.

• pooling (Literal['mean', 'attention_token', 'attention_seed']) –

  Pooling method, should be one of:

  • 'mean': Aggregates combinations of covariates by the mean of their learned embeddings.

  • 'attention_token': Aggregates combinations of covariates by an attention mechanism with a class token.

  • 'attention_seed': Aggregates combinations of covariates by seed attention.

• pooling_kwargs (dict[str, Any]) –

  Keyword arguments for the pooling method corresponding to:

  • cellflow.networks.TokenAttentionPooling if 'pooling' is 'attention_token'.

  • cellflow.networks.SeedAttentionPooling if 'pooling' is 'attention_seed'.

• layers_before_pool (dict[str, Sequence[dict[str, Any]]] | Sequence[dict[str, Any]]) –

  Layers applied to the condition embeddings before pooling. Can be of type

  • tuple with elements corresponding to dictionaries with keys:

    • 'layer_type' of type str indicating the type of the layer, can be 'mlp' or 'self_attention'.

    • Further keyword arguments for the layer type cellflow.networks.MLPBlock or cellflow.networks.SelfAttentionBlock.

  • dict with keys corresponding to perturbation covariate keys, and values correspondinng to the above mentioned tuples.

• layers_after_pool (Sequence[dict[str, Any]]) –

  Layers applied to the condition embeddings after pooling, and before applying the last layer of size 'condition_embedding_dim'. Should be of type tuple with elements corresponding to dictionaries with keys:

  • 'layer_type' of type str indicating the type of the layer, can be 'mlp' or 'self_attention'.

  • Further keys depend on the layer type, either for cellflow.networks.MLPBlock or for cellflow.networks.SelfAttentionBlock.

• condition_embedding_dim (int) – Dimensions of the condition embedding, i.e. the last layer of the cellflow.networks.ConditionEncoder.

• cond_output_dropout (float) – Dropout rate for the last layer of the cellflow.networks.ConditionEncoder.

• condition_encoder_kwargs (dict[str, Any] | None) – Keyword arguments for the cellflow.networks.ConditionEncoder.

• pool_sample_covariates (bool) – Whether to include sample covariates in the pooling.

• time_freqs (int) – Frequency of the sinusoidal time encoding (ott.neural.networks.layers.sinusoidal_time_encoder()).

• time_max_period (int | None) – Controls the frequency of the time embeddings, see cellflow.networks.utils.sinusoidal_time_encoder().

• time_encoder_dims (Sequence[int]) – Dimensions of the layers processing the time embedding in cellflow.networks.ConditionalVelocityField.time_encoder.

• time_encoder_dropout (float) – Dropout rate for the cellflow.networks.ConditionalVelocityField.time_encoder.

• hidden_dims (Sequence[int]) – Dimensions of the layers processing the input to the velocity field via cellflow.networks.ConditionalVelocityField.x_encoder.

• hidden_dropout (float) – Dropout rate for cellflow.networks.ConditionalVelocityField.x_encoder.

• conditioning (Literal['concatenation', 'film', 'resnet']) –

  Conditioning method, should be one of:

  • 'concatenation': Concatenate the time, data, and condition embeddings.

  • 'film': Use FiLM conditioning, i.e. learn FiLM weights from time and condition embedding to scale the data embeddings.

  • 'resnet': Use residual conditioning.

• conditioning_kwargs (dict[str, Any]) – Keyword arguments for the conditioning method.

• decoder_dims (Sequence[int]) – Dimensions of the output layers in cellflow.networks.ConditionalVelocityField.decoder.

• decoder_dropout (float) – Dropout rate for the output layer cellflow.networks.ConditionalVelocityField.decoder.

• vf_act_fn (Callable[[Array], Array]) – Activation function of the cellflow.networks.ConditionalVelocityField.

• vf_kwargs (dict[str, Any] | None) –

  Additional keyword arguments for the solver-specific vector field. For instance, when 'solver==genot', the following keyword argument can be passed:

  • 'genot_source_dims' of type tuple with the dimensions of the cellflow.networks.MLPBlock processing the source cell.

  • 'genot_source_dropout' of type float indicating the dropout rate for the source cell processing.

• probability_path (dict[Literal['constant_noise', 'bridge'], float] | None) –

  Probability path to use for training. Should be a dict of the form

  • '{"constant_noise": noise_val'

  • '{"bridge": noise_val}'

  If None, defaults to '{"constant_noise": 0.0}'.

• match_fn (Callable[[ndarray[tuple[Any, ...], dtype[float64]], ndarray[tuple[Any, ...], dtype[float64]]], ndarray[tuple[Any, ...], dtype[float64]]]) – Matching function between unperturbed and perturbed cells. Should take as input source and target data and return the optimal transport matrix, see e.g. cellflow.utils.match_linear().

• optimizer (GradientTransformation) – Optimizer used for training.

• solver_kwargs (dict[str, Any] | None) – Keyword arguments for the solver cellflow.solvers.OTFlowMatching or cellflow.solvers.GENOT.

• layer_norm_before_concatenation (bool) – If True, applies layer normalization before concatenating the embedded time, embedded data, and condition embeddings.

• linear_projection_before_concatenation (bool) – If True, applies a linear projection before concatenating the embedded time, embedded data, and embedded condition.

• seed – Random seed.

Return type:

None

Returns:

: Updates the following fields:

• cellflow.model.CellFlow.velocity_field - an instance of the cellflow.networks.ConditionalVelocityField.

• cellflow.model.CellFlow.solver - an instance of cellflow.solvers.OTFlowMatching or cellflow.solvers.GENOT.

• cellflow.model.CellFlow.trainer - an instance of the cellflow.training.CellFlowTrainer.