Sampling from 1D distribution

# this program follows the pbrt Distribution1D class methods
import math, numpy
import matplotlib.pyplot as plt
numpy.random.seed(123)

# 1D random variable CDF
cdf = [0.0, 0.4, 0.45, 0.94, 1.0]
N = len(cdf)

def clamp(num, low, high):
    return max(low, min(num, high))

# Find the random variable where the CDF value
# is right below the input probability u.
def find_interval(pred):
    first = 0
    len = N
    while(len > 0):
        half = len >> 1
        mid = first + half
        if (pred(mid)):
            first = mid + 1
            len -= half + 1
        else:
            len = half
    return clamp(first - 1, 0, N - 2)

# continuous
def sample_c(u):
    offset = find_interval(lambda idx: cdf[idx] <= u)
    du = u - cdf[offset]
    if ((cdf[offset + 1] - cdf[offset]) > 0):
        du /= cdf[offset + 1] - cdf[offset]
    return (offset + du) / (N - 1)

# discrete
def sample_d(u):
    offset = find_interval(lambda idx: cdf[idx] <= u)
    return offset

num_samples = 1000
data_c = []
data_d = []

for i in range(0, num_samples):
    ep = numpy.random.random()
    res_c = sample_c(ep)
    res_c *= N-1
    # Map result from [0, 1) to [0, N - 1),
    # since the range of continuous distributions starts from CDF arr idx=0.
    # The range ends at CDF arr idx=N - 1.
    data_c.append(res_c)

    res_d = sample_d(ep)
    res_d += 1
    # Map the n-th possible value to n + 1,
    # since the first possible discrete sample value is at CDF arr idx=1.
    # The last possible value is at CDF arr idx=N - 1.
    data_d.append(res_d)

fig = plt.figure(1, figsize=(8, 10))

# plot CDF
ax = fig.add_subplot(211)
ax.tick_params(axis='both', labelsize = 6)
x = range(0, N)
plt.plot(x, cdf, 'bo--')
plt.xlabel('possible values')
plt.ylabel('cumulative probability density')

# plot samples
ax = fig.add_subplot(212)
ax.tick_params(axis='both', labelsize = 6)
n = range(0, num_samples)
plt.plot(data_c, n, 'y.', label='continuous')
plt.plot(data_d, n, 'g.', label='discrete')
plt.xlabel('sample values')
plt.ylabel('sample index')
plt.legend(loc='best', framealpha=1)
plt.show()

Result:

Imgur