cellflow.solvers.OTFlowMatching

class cellflow.solvers.OTFlowMatching(vf, probability_path, match_fn=None, time_sampler=<function uniform_sampler>, guidance=None, **kwargs)[source]

(OT) flow matching [Lipman et al., 2022] extended to the conditional setting.

With an extension to OT-CFM [Pooladian et al., 2023, Tong et al., 2023], and its unbalanced version [Eyring et al., 2024].

Parameters:

  • vf (ConditionalVelocityField) – Vector field parameterized by a neural network.

  • probability_path (BaseFlow) – Probability path between the source and the target distributions.

  • match_fn (Callable[[Array, Array], Array] | None) – Function to match samples from the source and the target distributions. It has a (src, tgt) -> matching signature, see e.g. cellflow.utils.match_linear(). If None, no matching is performed, and pure probability_path matching [Lipman et al., 2022] is applied.

  • time_sampler (Callable[[Array, int], Array]) – Time sampler with a (rng, n_samples) -> time signature, see e.g. ott.solvers.utils.uniform_sampler().

  • guidance (Guidance | None) – Optional guidance strategy applied to the velocity field on the predict path, see e.g. ClassifierFreeGuidance. If None (the default), the plain conditional velocity field is used and prediction is unchanged.

  • kwargs (Any) – Keyword arguments for cellflow.networks.ConditionalVelocityField.create_train_state().

Methods

get_condition_embedding(condition[, ...])

Get learnt embeddings of the conditions.

predict(x, condition[, rng])

Predict the translated source x under condition condition.

step_fn(rng, batch)

Single step function of the solver.

Attributes

cfg_enabled

Whether classifier-free guidance is available at predict time.

is_trained

Whether the model is trained.