ott.neural.methods.flow_matching.evaluate_velocity_field

ott.neural.methods.flow_matching.evaluate_velocity_field#

ott.neural.methods.flow_matching.evaluate_velocity_field(model, x, cond=None, *, t0=0.0, t1=1.0, reverse=False, num_steps=None, solver=None, save_trajectory_kwargs=None, save_velocity_kwargs=None, **kwargs)[source]#

Solve an ODE.

Parameters:

  • model (Module) – Velocity field with a signature (t, x_t, cond) -> v_t.

  • x (Union[Array, Any]) – Initial point of shape [*dims].

  • cond (Optional[Array]) – Condition of shape [*cond_dims].

  • t0 (float) – Start time of the integration.

  • t1 (float) – End time of the integration.

  • reverse (bool) – Whether to integrate from \(t_1\) to \(t_0\).

  • num_steps (Optional[int]) – Number of steps used for solvers with a constant step size.

  • solver (Optional[AbstractSolver]) – ODE solver. If None and step_size = None, use Dopri5. Otherwise use Euler.

  • save_velocity_kwargs (Optional[Dict[str, Any]]) – Keyword arguments for SubSaveAt used to store the velocities along the integration path. The velocity will be saved in out.ys['v_t'].

  • save_trajectory_kwargs (Optional[Dict[str, Any]]) – Keyword arguments for SubSaveAt used to store the positions along the integration path. The trajectory will be saved in out.ys['x_t'].

  • kwargs (Any) – Keyword arguments for diffeqsolve().

Return type:

Solution

Returns:

The ODE solution.