grl.numerical_methods¶

ODE¶

class grl.numerical_methods.ODE(drift=None)[source]¶

Overview:

Base class for ordinary differential equations. The ODE is defined as:

\[dx = f(x, t)dt\]

where f(x, t) is the drift term.

Interfaces:

__init__

__init__(drift=None)[source]¶

SDE¶

class grl.numerical_methods.SDE(drift=None, diffusion=None)[source]¶

Overview:

Base class for stochastic differential equations. The SDE is defined as:

\[dx = f(x, t)dt + g(x, t)dW\]

where f(x, t) is the drift term, g(x, t) is the diffusion term, and dW is the Wiener process.

Interfaces:

__init__

__init__(drift=None, diffusion=None)[source]¶

DPMSolver¶

class grl.numerical_methods.DPMSolver(order, device, atol=1e-05, rtol=1e-05, steps=None, type='dpm_solver', method='singlestep', solver_type='dpm_solver', skip_type='time_uniform', denoise=False)[source]¶

Overview:: The DPM-Solver for sampling from the diffusion process.
Interface:: __init__, integrate

__init__(order, device, atol=1e-05, rtol=1e-05, steps=None, type='dpm_solver', method='singlestep', solver_type='dpm_solver', skip_type='time_uniform', denoise=False)[source]¶

Overview:: Initialize the DPM-Solver.

Parameters:

order (int) – The order of the DPM-Solver, which should be 1, 2, or 3.
device (str) – The device for the computation.
denoise (bool) – Whether to denoise at the final step.
atol (float) – The absolute tolerance for the adaptive solver.
rtol (float) – The relative tolerance for the adaptive solver.
steps (int) – The total number of function evaluations (NFE).
type (str) – The type for the DPM-Solver, which should be ‘dpm_solver’ or ‘dpm_solver++’.
method (str) – The method for the DPM-Solver, which should be ‘singlestep’, ‘multistep’, ‘singlestep_fixed’, or ‘adaptive’.
solver_type (str) – The type for the high-order solvers, which should be ‘dpm_solver’ or ‘taylor’. The type slightly impacts the performance. We recommend to use ‘dpm_solver’ type.
skip_type (str) – The type for the spacing of the time steps, which should be ‘logSNR’, ‘time_uniform’, or ‘time_quadratic’.
denoise – Whether to denoise at the final step.

integrate(diffusion_process, noise_function, data_prediction_function, x, condition=None, steps=None, save_intermediate=False)[source]¶

Overview:: Integrate the diffusion process by the DPM-Solver.

Parameters:

diffusion_process (DiffusionProcess) – The diffusion process.
noise_function (Callable) – The noise prediction model.
data_prediction_function (Callable) – The data prediction model.
x (Union[torch.Tensor, TensorDict]) – The initial value at time t_start.
condition (Union[torch.Tensor, TensorDict]) – The condition for the data prediction model.
steps (int) – The total number of function evaluations (NFE).
save_intermediate (bool) – If true, also return the intermediate model values.

Returns:

The approximated solution at time t_end.

Return type:

x_end (torch.Tensor)

ODESolver¶

class grl.numerical_methods.ODESolver(ode_solver='euler', dt=0.01, atol=1e-05, rtol=1e-05, library='torchdyn', **kwargs)[source]¶

Overview:: The ODE solver class.
Interfaces:: __init__, integrate

__init__(ode_solver='euler', dt=0.01, atol=1e-05, rtol=1e-05, library='torchdyn', **kwargs)[source]¶

Overview:: Initialize the ODE solver using torchdiffeq or torchdyn library.

Parameters:

ode_solver (str) – The ODE solver to use.
dt (float) – The time step.
atol (float) – The absolute tolerance.
rtol (float) – The relative tolerance.
library (str) – The library to use for the ODE solver. Currently, it supports ‘torchdiffeq’ and ‘torchdyn’.
**kwargs – Additional arguments for the ODE solver.

integrate(drift, x0, t_span, **kwargs)[source]¶

Overview:: Integrate the ODE.

Parameters:

drift (Union[nn.Module, Callable]) – The drift term of the ODE.
x0 (Union[torch.Tensor, TensorDict]) – The input initial state.
t_span (torch.Tensor) – The time at which to evaluate the ODE. The first element is the initial time, and the last element is the final time. For example, t = torch.tensor([0.0, 1.0]).

Returns:

The output trajectory of the ODE, which has the same data type as x0 and the shape of (len(t_span), *x0.shape).

Return type:

trajectory (Union[torch.Tensor, TensorDict])

SDESolver¶

class grl.numerical_methods.SDESolver(sde_solver='euler', sde_noise_type='diagonal', sde_type='ito', dt=0.001, atol=1e-05, rtol=1e-05, library='torchsde', **kwargs)[source]¶

__init__(sde_solver='euler', sde_noise_type='diagonal', sde_type='ito', dt=0.001, atol=1e-05, rtol=1e-05, library='torchsde', **kwargs)[source]¶

Overview:: Initialize the SDE solver using torchsde library.

Parameters:

sde_solver (str) – The SDE solver to use.
sde_noise_type (str) – The type of noise of the SDE. It can be ‘diagonal’, ‘general’, ‘scalar’ or ‘additive’.
sde_type (str) – The type of the SDE. It can be ‘ito’ or ‘stratonovich’.
dt (float) – The time step.
atol (float) – The absolute tolerance.
rtol (float) – The relative tolerance.
library (str) – The library to use for the ODE solver. Currently, it supports ‘torchsde’.
**kwargs – Additional arguments for the ODE solver.

integrate(drift, diffusion, x0, t_span, logqp=False, adaptive=False)[source]¶

Overview:: Integrate the SDE.

Parameters:

drift (nn.Module) – The function that defines the ODE.
diffusion (nn.Module) – The function that defines the ODE.

GaussianConditionalProbabilityPath¶

class grl.numerical_methods.GaussianConditionalProbabilityPath(config)[source]¶

Overview:

Gaussian conditional probability path.

General case:

\[p(x(t)|x(0))=\mathcal{N}(x(t);\mu(t,x(0)),\sigma^2(t,x(0))I)\]

If written in the form of SDE:

\[\mathrm{d}x=f(x,t)\mathrm{d}t+g(t)w_{t}\]

where \(f(x,t)\) is the drift term, \(g(t)\) is the diffusion term, and \(w_{t}\) is the Wiener process.

For diffusion model:

\[p(x(t)|x(0))=\mathcal{N}(x(t);s(t)x(0),\sigma^2(t)I)\]

or,

\[p(x(t)|x(0))=\mathcal{N}(x(t);s(t)x(0),s^2(t)e^{-2\lambda(t)}I)\]

If written in the form of SDE:

\[\mathrm{d}x=\frac{s'(t)}{s(t)}x(t)\mathrm{d}t+s^2(t)\sqrt{2(\frac{s'(t)}{s(t)}-\lambda'(t))}e^{-\lambda(t)}\mathrm{d}w_{t}\]

or,

\[\mathrm{d}x=f(t)x(t)\mathrm{d}t+g(t)w_{t}\]

where \(s(t)\) is the scale factor, \(\sigma^2(t)I\) is the covariance matrix,: \(\sigma(t)\) is the standard deviation with the scale factor, \(e^{-2\lambda(t)}I\) is the covariance matrix without the scale factor, \(\lambda(t)\) is the half-logSNR, which is the difference between the log scale factor and the log standard deviation, \(\lambda(t)=\log(s(t))-\log(\sigma(t))\).

For VP SDE:

\[p(x(t)|x(0))=\mathcal{N}(x(t);x(0)e^{-\frac{1}{2}\int_{0}^{t}{\beta(s)\mathrm{d}s}},(1-e^{-\int_{0}^{t}{\beta(s)\mathrm{d}s}})I)\]

For Linear VP SDE:

\[p(x(t)|x(0))=\mathcal{N}(x(t);x(0)e^{-\frac{\beta(1)-\beta(0)}{4}t^2-\frac{\beta(0)}{2}t},(1-e^{-\frac{\beta(1)-\beta(0)}{2}t^2-\beta(0)t})I)\]

#TODO: add more details for Cosine VP SDE; General VE SDE; OPT-Flow;

HalfLogSNR(t)[source]¶

Overview:: Compute the half-logSNR of the Gaussian conditional probability path, which is

\[\log(s(t))-\log(\sigma(t))\]

Parameters:: t (torch.Tensor) – The input time.
Returns:: The half-logSNR.
Return type:: HalfLogSNR (torch.Tensor)

InverseHalfLogSNR(HalfLogSNR)[source]¶

Overview:

Compute the inverse function of the half-logSNR of the Gaussian conditional probability path. Since the half-logSNR is an invertible function, we can compute the time t from the half-logSNR. For linear VP SDE, the inverse function is

\[t(\lambda)=\frac{1}{\beta_1-\beta_0}(\sqrt{\beta_0^2+2(\beta_1-\beta_0)\log{(e^{-2\lambda}+1)}}-\beta_0)\]

or,

\[t(\lambda)=\frac{2(\beta_1-\beta_0)\log{(e^{-2\lambda}+1)}}{\sqrt{\beta_0^2+2(\beta_1-\beta_0)\log{(e^{-2\lambda}+1)}}+\beta_0}\]

Parameters:: HalfLogSNR (torch.Tensor) – The input half-logSNR.
Returns:: The time.
Return type:: t (torch.Tensor)

__init__(config)[source]¶

Overview:: Initialize the Gaussian conditional probability path.

Parameters:: config (EasyDict) – The configuration of the Gaussian conditional probability path.

covariance(t)[source]¶

Overview:

Compute the covariance matrix of the Gaussian conditional probability path, which is

\[\Sigma(t)\]

or

\[\sigma^2(t)I\]

or

\[s^2(t)e^{-2\lambda(t)}I\]

Parameters:: t (torch.Tensor) – The input time.
Returns:: The covariance matrix.
Return type:: covariance (torch.Tensor)

d_covariance_dt(t)[source]¶

Overview:: Compute the time derivative of the covariance matrix of the Gaussian conditional probability path, which is

\[\frac{\mathrm{d}\Sigma(t)}{\mathrm{d}t}\]

Parameters:: t (torch.Tensor) – The input time.
Returns:: The time derivative of the covariance matrix.
Return type:: d_covariance_dt (torch.Tensor)

d_log_scale_dt(t)[source]¶

Overview:: Compute the time derivative of the log scale factor of the Gaussian conditional probability path, which is

\[\log(s'(t))\]

Parameters:: t (torch.Tensor) – The input time.
Returns:: The time derivative of the log scale factor.
Return type:: d_log_scale_dt (Union[torch.Tensor, TensorDict])

d_scale_dt(t)[source]¶

Overview:: Compute the time derivative of the scale factor of the Gaussian conditional probability path, which is

\[s'(t)\]

Parameters:: t (torch.Tensor) – The input time.
Returns:: The time derivative of the scale factor.
Return type:: d_scale_dt (Union[torch.Tensor, TensorDict])

d_std_dt(t)[source]¶

Overview:: Compute the time derivative of the standard deviation of the Gaussian conditional probability path, which is

\[\frac{\mathrm{d}\sigma(t)}{\mathrm{d}t}\]

Parameters:: t (torch.Tensor) – The input time.
Return type:: Tensor

diffusion(t)[source]¶

Overview:: Return the diffusion term of the Gaussian conditional probability path. The diffusion term is given by the following:

\[g(x,t)\]

Parameters:

t (torch.Tensor) – The input time.
x (Union[torch.Tensor, TensorDict]) – The input state.

Returns:

The output diffusion term.

Return type:

diffusion (Union[torch.Tensor, TensorDict])

diffusion_squared(t)[source]¶

Overview:: Return the diffusion term of the Gaussian conditional probability path. The diffusion term is given by the following:

\[g^2(x,t)\]

Parameters:

t (torch.Tensor) – The input time.
x (Union[torch.Tensor, TensorDict]) – The input state.

Returns:

The output diffusion term.

Return type:

diffusion (Union[torch.Tensor, TensorDict])

drift(t, x=None)[source]¶

Overview:: Return the drift term of the Gaussian conditional probability path. The drift term is given by the following:

\[f(x,t)\]

Parameters:

t (torch.Tensor) – The input time.
x (Union[torch.Tensor, TensorDict]) – The input state.

Returns:

The output drift term.

Return type:

drift (Union[torch.Tensor, TensorDict])

drift_coefficient(t)[source]¶

Overview:

Return the drift coefficient term of the Gaussian conditional probability path. The drift term is given by the following:

\[f(t)\]

which satisfies the following SDE:

\[\mathrm{d}x=f(t)x(t)\mathrm{d}t+g(t)w_{t}\]

Parameters:: t (torch.Tensor) – The input time.
Returns:: The output drift term.
Return type:: drift (Union[torch.Tensor, TensorDict])

log_scale(t)[source]¶

Overview:: Compute the log scale factor of the Gaussian conditional probability path, which is

\[\log(s(t))\]

Parameters:: t (torch.Tensor) – The input time.
Returns:: The log scale factor.
Return type:: log_scale (torch.Tensor)

scale(t)[source]¶

Overview:: Compute the scale factor of the Gaussian conditional probability path, which is

\[s(t)\]

Parameters:: t (torch.Tensor) – The input time.
Returns:: The scale factor.
Return type:: scale (torch.Tensor)

std(t)[source]¶

Overview:

Compute the standard deviation of the Gaussian conditional probability path, which is

\[\sqrt{\Sigma(t)}\]

or

\[\sigma(t)\]

or

\[s(t)e^{-\lambda(t)}\]

Parameters:: t (torch.Tensor) – The input time.
Returns:: The standard deviation.
Return type:: std (torch.Tensor)