def generate_forward_diffusion_sample(self, X_noiseless):
"""
Corrupt a training image with t steps worth of Gaussian noise, and
return the corrupted image, as well as the mean and covariance of the
posterior q(x^{t-1}|x^t, x^0).
"""
X_noiseless = X_noiseless.reshape(
(-1, self.n_colors, self.spatial_width, self.spatial_width))
n_images = X_noiseless.shape[0].astype('int16')
rng = Random().theano_rng
# choose a timestep in [1, self.trajectory_length-1].
# note the reverse process is fixed for the very
# first timestep, so we skip it.
# TODO for some reason random_integer is missing from the Blocks
# theano random number generator.
t = T.floor(rng.uniform(size=(1,1), low=1, high=self.trajectory_length,
dtype=theano.config.floatX))
t_weights = self.get_t_weights(t)
N = rng.normal(size=(n_images, self.n_colors, self.spatial_width, self.spatial_width),
dtype=theano.config.floatX)
# noise added this time step
beta_forward = self.get_beta_forward(t)
# decay in noise variance due to original signal this step
alpha_forward = 1. - beta_forward
# compute total decay in the fraction of the variance due to X_noiseless
alpha_arr = 1. - self.beta_arr
alpha_cum_forward_arr = T.extra_ops.cumprod(alpha_arr).reshape((self.trajectory_length,1))
alpha_cum_forward = T.dot(t_weights.T, alpha_cum_forward_arr)
# total fraction of the variance due to noise being mixed in
beta_cumulative = 1. - alpha_cum_forward
# total fraction of the variance due to noise being mixed in one step ago
beta_cumulative_prior_step = 1. - alpha_cum_forward/alpha_forward
# generate the corrupted training data
X_uniformnoise = X_noiseless + (rng.uniform(size=(n_images, self.n_colors, self.spatial_width, self.spatial_width),
dtype=theano.config.floatX)-T.constant(0.5,dtype=theano.config.floatX))*T.constant(self.uniform_noise,dtype=theano.config.floatX)
X_noisy = X_uniformnoise*T.sqrt(alpha_cum_forward) + N*T.sqrt(1. - alpha_cum_forward)
# compute the mean and covariance of the posterior distribution
mu1_scl = T.sqrt(alpha_cum_forward / alpha_forward)
mu2_scl = 1. / T.sqrt(alpha_forward)
cov1 = 1. - alpha_cum_forward/alpha_forward
cov2 = beta_forward / alpha_forward
lam = 1./cov1 + 1./cov2
mu = (
X_uniformnoise * mu1_scl / cov1 +
X_noisy * mu2_scl / cov2
) / lam
sigma = T.sqrt(1./lam)
sigma = sigma.reshape((1,1,1,1))
mu.name = 'mu q posterior'
sigma.name = 'sigma q posterior'
X_noisy.name = 'X_noisy'
t.name = 't'
return X_noisy, t, mu, sigma
