cellflow.solvers.OTFlowMatching.predict

OTFlowMatching.predict(x, condition, rng=None, **kwargs)[source]

Predict the translated source x under condition condition.

This function solves the ODE learnt with the ConditionalVelocityField.

Parameters:

  • x (ndarray[tuple[Any, ...], dtype[float64]] | dict[str, ndarray[tuple[Any, ...], dtype[float64]]]) – A dictionary with keys indicating the name of the condition and values containing the input data as arrays.

  • condition (dict[str, ndarray[tuple[Any, ...], dtype[float64]]] | dict[str, dict[str, ndarray[tuple[Any, ...], dtype[float64]]]]) – A dictionary with keys indicating the name of the condition and values containing the condition of input data as arrays.

  • rng (Array | None) – Random number generator to sample from the latent distribution, only used if condition_mode='stochastic'. If None, the mean embedding is used.

  • kwargs (Any) – Keyword arguments for diffrax.diffeqsolve(). May also include guidance_scale (a float): a per-call classifier-free guidance scale that, when not 1.0, applies ClassifierFreeGuidance for this call (requires cfg_enabled), overriding the construction-time guidance. Handy for sweeping w without rebuilding the solver.

Return type:

ndarray[tuple[Any, ...], dtype[float64]] | dict[str, ndarray[tuple[Any, ...], dtype[float64]]]

Returns:

: The push-forward distribution of x under condition condition.