gensbi.diffusion.solver.sm_ode_solver#

Score matching probability flow ODE solver.

Provides SMODESolver, a thin subclass of ODESolver used for dispatch in the score matching pipeline.

The probability flow ODE is:

\[dx = \left[ f(x,t) - \tfrac{1}{2}\, g(t)^2\, s_\theta(x, t) \right] dt\]

The pipeline wraps the score model with ScoreToODEDrift before constructing this solver, so the drift returned by get_drift() is already the PF-ODE drift.

All integration and log-probability logic is inherited from ODESolver.

Classes#

SMODESolver

Score matching probability flow ODE solver.

Module Contents#

class gensbi.diffusion.solver.sm_ode_solver.SMODESolver(velocity_model)[source]#

Bases: gensbi.core.ode_solver.ODESolver

Score matching probability flow ODE solver.

Uses the probability flow ODE formulation to sample deterministically from a score matching model. The velocity model passed to this solver should already be wrapped with ScoreToODEDrift + ModelWrapper by the pipeline.

All integration and log-probability logic is inherited from ODESolver.

Parameters:

velocity_model (gensbi.utils.model_wrapping.ModelWrapper)

get_drift(**kwargs)[source]#

Return the probability flow ODE drift (from ScoreToODEDrift adapter).

Return type:

Callable