Is it possible to also share the code for Figure 3 in the paper

[yuanzhi-zhu]

I found Figure 3 to be very helpful for understanding diffuson models in general, and it would be better if one can play with the code.

I tried something like

import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt

schedule = 'VESDE'

if schedule == 'VESDE':
    sigma_min = 0.02
    sigma_max = 100
    rho = 0
elif schedule == 'EDM':
    sigma_min = 0.002
    sigma_max = 80
    rho = 7.0

def gauss_norm(x):
    return np.exp(-x**2/2)/np.sqrt(2*np.pi)

def pred_x0_theory(x,sigma_t):
    c = np.array([
        [-1],
        [ 1],
    ])
    nominator = 0
    denominator = 0
    for i in range(c.shape[0]):
        nominator += c[i] * gauss_norm((x-c[i])/sigma_t)
        denominator += gauss_norm((x-c[i])/sigma_t)
    return nominator / (denominator + 0)

# define the ODE
def model(x, t):
    def s_t(t):
        return 1

    def s_t_p(t):
        return 0

    def sigma_t(t):
        # VESDE
        if schedule == 'VESDE':
            return sigma_min**2 * (sigma_max**2/sigma_min**2)**t
        elif schedule == 'EDM':
            rho_inv = 1.0 / rho
            sigmas = sigma_min**rho_inv + t * (
                sigma_max**rho_inv - sigma_min**rho_inv
            )
            sigmas = sigmas**rho
            return sigmas

    def sigma_t_p(t, h=1e-5):
        return (sigma_t(t + h) - sigma_t(t - h)) / (2. * h)

    first_term = s_t_p(t) / s_t(t) * x
    score = (pred_x0_theory(x/s_t(t), sigma_t(t)) - x) / sigma_t(t)**2
    second_term = s_t(t)**2 * sigma_t_p(t) * sigma_t(t) * score
    dxdt = first_term - second_term
    return dxdt

# initial condition
x0 = 1
print(x0)

# time points
t = np.linspace(0.01,0.6,100)

# solve ODE
y = odeint(model, x0, t)

# plot results
plt.plot(t, y)
plt.xlabel('time')
plt.ylabel('x(t)')
plt.show()

but it does not work out well.

thank you so much in advance for your help.

[yuanzhi-zhu]

fixed

import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt

schedule = 'EDM'

if schedule == 'VESDE':
    sigma_min = 0.02
    sigma_max = 100
    rho = 0
elif schedule == 'EDM':
    sigma_min = 0.002
    sigma_max = 80
    rho = 7.0

def gauss_norm(x):
    return np.exp(-x**2/2)/np.sqrt(2*np.pi)

def pred_x0_theory(x,sigma_t):
    c = np.array([
        [-1],
        [ 1],
    ])
    nominator = 0
    denominator = 0
    for i in range(c.shape[0]):
        nominator += c[i] * gauss_norm((x-c[i])/sigma_t)
        denominator += gauss_norm((x-c[i])/sigma_t)
    return nominator / (denominator + 0)

# define the ODE
def model(x, t):
    def s_t(t):
        return 1

    def s_t_p(t):
        return 0

    def sigma_t(t):
        if schedule == 'VESDE':
            # return self.sigma_min**2 * (self.sigma_max**2/self.sigma_min**2)**t
            return np.sqrt(t)
        elif schedule == 'EDM':
            # rho_inv = 1.0 / self.rho
            # sigmas = self.sigma_min**rho_inv + t * (
            #     self.sigma_max**rho_inv - self.sigma_min**rho_inv
            # )
            # sigmas = sigmas**self.rho
            # return sigmas
            return t

    def sigma_t_p(t, h=1e-5):
        return (sigma_t(t + h) - sigma_t(t - h)) / (2. * h)

    first_term = s_t_p(t) / s_t(t) * x
    score = (pred_x0_theory(x/s_t(t), sigma_t(t)) - x) / sigma_t(t)**2
    second_term = s_t(t)**2 * sigma_t_p(t) * sigma_t(t) * score
    dxdt = first_term - second_term
    return dxdt

# time points
t = np.linspace(0.001,25,10000)
batch_size = 30
noise_init = 0.001

# Initial conditions
x1_values = np.random.normal(1, noise_init, (batch_size,))
x1m_values = np.random.normal(-1, noise_init, (batch_size,))
x0_values = np.concatenate((x1_values, x1m_values))

# For each initial condition, solve the ODE and plot the solution
for x0 in x0_values:
    x = odeint(model, x0, t)
    plt.plot(t, x, 'orange')

plt.xlabel('t')
plt.ylabel('x')
plt.title('Trajectories of the ODE')
plt.show()