Add gradient to uniform distribution

[chaozg]

fixed #495 .

• I have tested the implementation with MALA, ULA and NUTS from cuqi.experimental.mcmc for two demos (see below for MWE code);
• sample goes well for all these three samplers
• warmup from NUTS is always stuck at some point and I created a separate issue to keep track of it #529

Demo 1: draw samples from Uniform(0, 1) with MALA, ULA and NUTS

import cuqi
import matplotlib.pyplot as plt
import numpy as np

uniform = cuqi.distribution.Uniform(np.array([0.0]), np.array([1.0]))
# sampler = cuqi.experimental.mcmc.MALA(uniform, initial_point=np.array([0.1]))
# sampler = cuqi.experimental.mcmc.ULA(uniform, initial_point=np.array([0.1]))
sampler = cuqi.experimental.mcmc.NUTS(uniform, initial_point=np.array([0.1]))
sampler.sample(10000)
samples = sampler.get_samples()
samples.plot_trace()

Demo 2: solve "Simplest" BIP with MALA, ULA and NUTS

import numpy as np
import matplotlib.pyplot as plt
import cuqi
from cuqi.utilities import plot_2D_density
np.random.seed(10)

# the forward model
A_matrix = np.array([[1.0, 1.0]])
A = cuqi.model.LinearModel(A_matrix)

# the prior
# x = Gaussian(np.zeros(2), 2.5)
x = cuqi.distribution.Uniform(np.array([-1.0, -1.0]), np.array([3.0, 3.0]))
print(x)

# the data distribution
b = cuqi.distribution.Gaussian(A@x, 0.1)

# the observed data
particular_x = np.array([1.5, 1.5])
b_given_particular_x = b(x=particular_x)
b_obs = b_given_particular_x.sample()
print(b_obs)

# the posterior
joint = cuqi.distribution.JointDistribution(x, b)
post = joint(b=b_obs)

# sampling with MALA, ULA and NUTS
# sampler = cuqi.experimental.mcmc.MALA(post, initial_point=np.array([2.5, 2.5]), scale=0.12)
# sampler = cuqi.experimental.mcmc.ULA(post, initial_point=np.array([2.5, 2.5]), scale=0.12)
sampler = cuqi.experimental.mcmc.NUTS(post, initial_point=np.array([2.5, 2.5]))

# sampler.warmup(1000)
sampler.sample(1000)
samples = sampler.get_samples()
samples.plot_trace()

# what b looks like with the drawn samples?
b_with_samples = cuqi.samples.Samples(samples.samples[:1,:] + samples.samples[1:,:])
plt.figure()
b_with_samples.plot_trace()

# plot exact posterior distribution and samples
# the posterior PDF
plt.figure()
plot_2D_density(post, -5, 5, -5, 5)
plt.title("Exact Posterior")

# samples
plt.figure()
samples.plot_pair()
plt.xlim(-5, 5)
plt.ylim(-5, 5)
plt.gca().set_aspect('equal')
plt.title("Posterior Samples")

The following plot shows what x1+x2 looks like with the drawn samples

[nabriis]

Let us try an example to see if NUTS will perform reasonably. How about we attempt the case from here with a Uniform prior instead? https://cuqi-dtu.github.io/CUQI-Book/chapter01/intro_example_short.html

[chaozg]

Hi @nabriis , I've tried the update with the simplest BIP, and the result seems reasonable, so this PR is ready for your further review : )

[chaozg]

Thank you, Amal. Done as suggested and issue is at https://github.com/CUQI-DTU/CUQI-Book/issues/90

[nabriis]

Nice!