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 posterior_quantiles(y_pred, bins, quantiles, bin_axis=1):
if len(y_pred.shape) == 1:
bin_axis = 0
n_y = y_pred.shape[bin_axis]
n_b = len(bins)
_check_dimensions(n_y, n_b)
xp = get_array_module(y_pred)
y_cdf = posterior_cdf(y_pred, bins, bin_axis=bin_axis)
n = len(y_pred.shape)
dx = bins[1:] - bins[:-1]
x_shape = [1] * n
x_shape[bin_axis] = numel(bins)
dx = pad_zeros_left(xp, dx, 1, 0)
dx = reshape(xp, dx, x_shape)
y_qs = []
for q in quantiles:
mask = as_type(xp, y_cdf <= q, y_cdf)
y_q = bins[0] + xp.sum(mask * dx, bin_axis)
y_q = expand_dims(xp, y_q, bin_axis)
y_qs.append(y_q)
y_q = concatenate(xp, y_qs, bin_axis)
return y_q
def test_concatenate(backend):
"""
Ensures that concatenation of array yields tensor with the expected size.
"""
array_1 = backend.ones((10, 1))
array_2 = backend.ones((10, 2))
result = concatenate(backend, [array_1, array_2], 1)
assert numel(result) == 30
def fit_gaussian_to_quantiles(y_pred, quantiles, quantile_axis=1):
"""
Fits Gaussian distributions to predicted quantiles.
Fits mean and standard deviation values to quantiles by minimizing
the mean squared distance of the predicted quantiles and those of
the corresponding Gaussian distribution.
Args:
y_pred: A rank-k tensor containing the predicted quantiles along
the axis specified by ``quantile_axis``.
quantiles: Array of shape `(m,)` containing the quantile
fractions corresponding to the predictions in ``y_pred``.
Returns:
Tuple ``(mu, sigma)`` of tensors of rank k-1 containing the mean and
standard deviations of the Gaussian distributions corresponding to
the predictions in ``y_pred``.
"""
if len(y_pred.shape) == 1:
quantile_axis = 0
xp = get_array_module(y_pred)
x = to_array(xp, norm.ppf(quantiles))
n_dims = len(y_pred.shape)
x_shape = [
1,
] * n_dims
x_shape[quantile_axis] = -1
x_shape = tuple(x_shape)
x = reshape(xp, x, x_shape)
output_shape = list(y_pred.shape)
output_shape[quantile_axis] = 1
output_shape = tuple(output_shape)
d2e_00 = numel(x)
d2e_01 = x.sum()
d2e_10 = x.sum()
d2e_11 = (x**2).sum()
d2e_det_inv = 1.0 / (d2e_00 * d2e_11 - d2e_01 * d2e_11)
d2e_inv_00 = d2e_det_inv * d2e_11
d2e_inv_01 = -d2e_det_inv * d2e_01
d2e_inv_10 = -d2e_det_inv * d2e_10
d2e_inv_11 = d2e_det_inv * d2e_00
x = reshape(xp, x, x_shape)
de_0 = reshape(xp, -(y_pred - x).sum(axis=quantile_axis), output_shape)
de_1 = reshape(xp, -(x * (y_pred - x)).sum(axis=quantile_axis),
output_shape)
mu = -(d2e_inv_00 * de_0 + d2e_inv_01 * de_1)
sigma = 1.0 - (d2e_inv_10 * de_0 + d2e_inv_11 * de_1)
return mu, sigma
def test_numel(backend):
"""
Ensures that the numel function returns the right number of elements.
"""
array = backend.ones(10)
assert numel(array) == 10
def posterior_quantiles(y_pdf, bins, quantiles, bin_axis=1):
"""
Calculate posterior quantiles from predicted PDFs.
Args:
y_pdf: Tensor containing the predicted PDFs.
bins: The bin-boundaries corresponding to the predictions.
quantiles: List containing the quantiles fractions of the quantiles
to compute.
bin_axis: The index of the tensor axis which contains the predictions
for each bin.
Return:
Tensor with same rank as ``y_pdf`` but with the values
the values along ``bin_axis`` replaced with the quantiles
of the predicted distributions.
"""
if len(y_pdf.shape) == 1:
bin_axis = 0
n_y = y_pdf.shape[bin_axis]
n_b = len(bins)
n_dims = len(y_pdf.shape)
_check_dimensions(n_y, n_b)
xp = get_array_module(y_pdf)
y_cdf = posterior_cdf(y_pdf, bins, bin_axis=bin_axis)
n = len(y_pdf.shape)
dx = bins[1:] - bins[:-1]
x_shape = [1] * n
x_shape[bin_axis] = numel(bins)
dx = pad_zeros_left(xp, dx, 1, 0)
dx = reshape(xp, dx, x_shape)
selection = [slice(0, None)] * n_dims
selection[bin_axis] = slice(0, -1)
selection_l = tuple(selection)
selection[bin_axis] = slice(1, None)
selection_r = tuple(selection)
cdf_l = y_cdf[selection_l]
cdf_r = y_cdf[selection_r]
d_cdf = cdf_r - cdf_l
shape = [1] * n_dims
shape[bin_axis] = -1
bins_l = bins.reshape(shape)[selection_l]
bins_r = bins.reshape(shape)[selection_r]
y_qs = []
for q in quantiles:
mask_l = as_type(xp, cdf_l <= q, cdf_l)
mask_r = as_type(xp, cdf_r > q, cdf_l)
mask = mask_l * mask_r
d_q = q - expand_dims(xp, (cdf_l * mask).sum(bin_axis), bin_axis)
result = (d_q * bins_r + (d_cdf - d_q) * bins_l) * mask
result = result / (d_cdf + as_type(xp, d_cdf < 1e-6, d_cdf))
result = result.sum(bin_axis)
result = result + bins[-1] * (1.0 - as_type(xp, mask_r.sum(bin_axis) > 0, mask))
result = result + bins[0] * (1.0 - as_type(xp, mask_l.sum(bin_axis) > 0, mask))
result = expand_dims(xp, result, bin_axis)
y_qs.append(result)
y_q = concatenate(xp, y_qs, bin_axis)
return y_q