def test_sample_uniform(backend):
"""
Ensures that array of random samples has array type
corresponding to the right backend module.
"""
samples = sample_uniform(backend, (10, ))
assert get_array_module(samples) == backend
def sample_posterior(y_pred, bins, n_samples=1, bin_axis=1):
"""
Sample the posterior distribution described by the predicted PDF.
The sampling is performed by interpolating the inverse of the cumulative
distribution function to value sampled from a uniform distribution.
Args:
y_pred: A rank-k tensor containing the predicted bin-probabilities
along the axis specified by ``quantile_axis``.
bins: The bin bounrdaries corresponding to the predicted
bin probabilities.
n_samples: How many samples to generate for each prediction.
bin_axis: The axis in y_pred along which the predicted bin
probabilities are located.
Returns:
A rank-k tensor with the values along ``bin_axis`` replaced by
samples of the posterior distribution.
"""
if len(y_pred.shape) == 1:
bin_axis = 0
xp = get_array_module(y_pred)
n_dims = len(y_pred.shape)
y_cdf = posterior_cdf(y_pred, bins, bin_axis=bin_axis)
n_bins = len(bins)
output_shape = list(y_cdf.shape)
output_shape[bin_axis] = n_samples
results = zeros(xp, output_shape, like=y_pred)
y_index = [slice(0, None)] * n_dims
y_index[bin_axis] = slice(0, 1)
y_l = y_cdf[tuple(y_index)]
b_l = bins[0]
samples = as_type(xp, sample_uniform(xp, tuple(output_shape)), y_cdf)
for i in range(1, n_bins):
y_index = [slice(0, None)] * n_dims
y_index[bin_axis] = slice(i, i + 1)
y_r = y_cdf[tuple(y_index)]
b_r = bins[i]
mask = as_type(xp, (y_l < samples) * (y_r >= samples), y_l)
results += b_l * (y_r - samples) * mask
results += b_r * (samples - y_l) * mask
results /= mask * (y_r - y_l) + (1.0 - mask)
b_l = b_r
y_l = y_r
mask = as_type(xp, y_r < samples, y_r)
results += mask * b_r
return results
def add(y_pdf_1, bins_1, y_pdf_2, bins_2, bins_out, bin_axis=1):
"""
Calculate the discretized PDF of the sum of two random variables
represented by their respective discretized PDFs.
Args:
y_pdf_1: The discretized PDF of the first random variable.
bins_1: The bin boundaries corresponding to 'y_pdf_1'.
y_pdf_2: The discretized PDF of the second random variable.
bins_2: The bin boundaries corresponding to 'y_pdf_2'.
bins_out: The bins boundaries for the resulting discretized PDF.
bin_axis: The dimension along which the probabilities are
oriented.
Return:
A tensor containing the discretized PDF corresponding to the sum of
the two given PDFs.
"""
if len(y_pdf_1.shape) == 1:
bin_axis = 0
xp = get_array_module(y_pdf_1)
bins_1_c = 0.5 * (bins_1[1:] + bins_1[:-1])
dx_1 = bins_1[1:] - bins_1[:-1]
shape_1 = [1] * len(y_pdf_1.shape)
shape_1[bin_axis] = numel(bins_1) - 1
dx_1 = dx_1.reshape(shape_1)
p_1 = y_pdf_1 * dx_1
bins_2_c = 0.5 * (bins_2[1:] + bins_2[:-1])
dx_2 = bins_2[1:] - bins_2[:-1]
shape_2 = [1] * len(y_pdf_2.shape)
shape_2[bin_axis] = numel(bins_2) - 1
dx_2 = dx_2.reshape(shape_2)
p_2 = y_pdf_2 * dx_2
out_shape = list(y_pdf_1.shape)
out_shape[bin_axis] = numel(bins_out) - 1
p_out = zeros(xp, out_shape, like=y_pdf_1)
rank = len(y_pdf_1.shape)
selection = [slice(0, None)] * rank
n_bins = numel(bins_1_c)
offsets = sample_uniform(xp, (n_bins,), like=bins_2)
for i in range(n_bins):
d_b = bins_1[i + 1] - bins_1[i]
b = bins_1[i] + offsets[i] * d_b
selection[bin_axis] = i
bins = bins_2_c + b
probs = p_1[tuple(selection)] * p_2
inds = digitize(xp, bins, bins_out) - 1
p_out = scatter_add(xp, p_out, inds, probs, bin_axis)
return normalize(p_out, bins_out, bin_axis=bin_axis)
def sample_posterior(y_pred, quantiles, n_samples=1, quantile_axis=1):
"""
Sample the posterior distribution described by the predicted quantiles.
The sampling is performed by interpolating the inverse of the cumulative
distribution function to value sampled from a uniform distribution.
Args:
y_pred: A rank-k tensor containing the predicted quantiles along the
axis specified by ``quantile_axis``.
quantiles: The quantile fractions corresponding to the predicted
quantiles.
n_samples: How many samples to generate for each prediction.
quantile_axis: The axis in y_pred along which the predicted quantiles
are found.
Returns:
A rank-k tensor with the values along ``quantile_axis`` replaced by
samples of the posterior distribution.
"""
if len(y_pred.shape) == 1:
quantile_axis = 0
xp = get_array_module(y_pred)
n_dims = len(y_pred.shape)
x_cdf, y_cdf = cdf(y_pred, quantiles, quantile_axis=quantile_axis)
output_shape = list(y_pred.shape)
output_shape[quantile_axis] = n_samples
samples = sample_uniform(xp, tuple(output_shape))
results = xp.zeros(samples.shape)
y_l = y_cdf[0]
x_index = [slice(0, None)] * n_dims
x_index[quantile_axis] = slice(0, 1)
x_l = x_cdf[tuple(x_index)]
for i in range(1, len(y_cdf)):
y_r = y_cdf[i]
x_index[quantile_axis] = slice(i, i + 1)
x_r = x_cdf[tuple(x_index)]
mask = as_type(xp, (samples > y_l) * (samples <= y_r), y_l)
results += (x_l * (y_r - samples)) * mask
results += (x_r * (samples - y_l)) * mask
results /= (mask * (y_r - y_l) + (1.0 - mask))
y_l = y_r
x_l = x_r
return results