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 probability_less_than(y_pred, quantiles, y, quantile_axis=1):
"""
Calculate the probability that the predicted value is less
than a given threshold value ``y`` given a tensor of predicted
quantiles ``y_pred``.
The probability :math:`P(Y > y)` is calculated by using the predicted
quantiles to estimate the CDF of the posterior distribution, which
is then interpolate to the given threshold value.
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.
y: The threshold value.
quantile_axis: The axis in y_pred along which the predicted quantiles
are found.
Returns:
A rank-(k-1) tensor containing for each set of predicted quantiles the
estimated probability of the true value being larger than the given
threshold.
"""
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(x_cdf.shape)
del output_shape[quantile_axis]
probabilities = xp.zeros(output_shape)
y_l = y_cdf[0]
x_index = [slice(0, None)] * n_dims
x_index[quantile_axis] = 0
x_l = x_cdf[tuple(x_index)]
for i in range(1, len(y_cdf)):
y_r = y_cdf[i]
x_index[quantile_axis] = i
x_r = x_cdf[tuple(x_index)]
mask = as_type(xp, (x_l < y) * (x_r >= y), x_l)
probabilities += y_l * (x_r - y) * mask
probabilities += y_r * (y - x_l) * mask
probabilities /= (mask * (x_r - x_l) + (1.0 - mask))
y_l = y_r
x_l = x_r
mask = as_type(xp, x_r < y, x_r)
probabilities += mask
return probabilities
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 = as_type(xp, sample_uniform(xp, tuple(output_shape)), y_cdf)
results = zeros(xp, samples.shape, like=y_pred)
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
def probability_less_than(y_pdf, bins, y, bin_axis=1):
"""
Calculate the probability of a sample being less than a given
value for a tensor of predicted PDFs.
Args:
y_pdf: Tensor containing the predicted PDFs.
bins: The bin-boundaries corresponding to the predictions.
y: The sample value.
bin_axis: The index of the tensor axis which contains the predictions
for each bin.
Return:
Tensor with rank reduced by one compared to ``y_pdf`` and with
the values along ``bin_axis`` of ``y_pdf`` replaced with the
probability that a sample of the distribution is smaller than
the given value ``y``.
"""
if len(y_pdf.shape) == 1:
bin_axis = 0
n_y = y_pdf.shape[bin_axis]
n_b = len(bins)
_check_dimensions(n_y, n_b)
xp = get_array_module(y_pdf)
n = len(y_pdf.shape)
x = 0.5 * (bins[1:] + bins[:-1])
mask = x < y
shape = [1] * n
shape[bin_axis] = -1
mask = as_type(xp, reshape(xp, mask, shape), y_pdf)
return trapz(xp, mask * y_pdf, bins, 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_as_type(backend):
"""
Ensures that conversion of types works.
"""
array = backend.ones((10, 10))
mask = array > 0.0
result = as_type(backend, mask, array)
assert array.dtype == result.dtype
def quantile_function(y_pdf, y_true, bins, bin_axis=1):
"""
Evaluates the quantile function at given y values.
"""
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)
n = len(y_pdf.shape)
y_cdf = posterior_cdf(y_pdf, bins, bin_axis=bin_axis)
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]
shape = [1] * n_dims
shape[bin_axis] = -1
bins_l = bins.reshape(shape)[selection_l]
bins_r = bins.reshape(shape)[selection_r]
d_bins = bins_r - bins_l
if len(y_true.shape) < len(y_pdf.shape):
y_true = expand_dims(xp, y_true, bin_axis)
mask_l = as_type(xp, bins_l <= y_true, y_true)
mask_r = as_type(xp, bins_r > y_true, y_true)
mask = mask_l * mask_r
dy = y_true - expand_dims(xp, (bins_l * mask).sum(bin_axis), bin_axis)
result = (dy * cdf_r + (d_bins - dy) * cdf_l) * mask
result = result / d_bins
result = result.sum(bin_axis)
result = result + 1.0 - as_type(xp, mask_r.sum(bin_axis) > 0, mask)
return result
def __call__(self, x, dist_axis=1):
"""
Evaluate the a priori.
Args:
x: Tensor containing the values at which to evaluate the a priori.
dist_axis: The axis along which the tensor x is sorted.
Returns;
Tensor with the same size as 'x' containing the values of the a priori
at 'x' obtained by linear interpolation.
"""
if len(x.shape) == 1:
dist_axis = 0
xp = get_array_module(x)
n_dims = len(x.shape)
n = x.shape[dist_axis]
x_index = [slice(0, None)] * n_dims
x_index[dist_axis] = 0
selection_l = [slice(0, None)] * n_dims
selection_l[dist_axis] = slice(0, -1)
selection_l = tuple(selection_l)
selection_r = [slice(0, None)] * n_dims
selection_r[dist_axis] = slice(1, None)
selection_r = tuple(selection_r)
r_shape = [1] * n_dims
r_shape[dist_axis] = -1
r_x = self.x.reshape(r_shape)
r_y = self.y.reshape(r_shape)
r_x_l = r_x[selection_l]
r_x_r = r_x[selection_r]
r_y_l = r_y[selection_l]
r_y_r = r_y[selection_r]
rs = []
for i in range(0, n):
x_index[dist_axis] = slice(i, i + 1)
index = tuple(x_index)
x_i = x[index]
mask = as_type(xp, (r_x_l < x_i) * (r_x_r >= x_i), x_i)
r = r_y_l * (r_x_r - x_i) * mask
r += r_y_r * (x_i - r_x_l) * mask
r /= mask * (r_x_r - r_x_l) + (1.0 - mask)
r = expand_dims(xp, r.sum(dist_axis), dist_axis)
rs.append(r)
r = concatenate(xp, rs, dist_axis)
return r
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 = as_type(xp, reshape(xp, x, x_shape), y_pred)
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 crps(y_pdf, y_true, bins, bin_axis=1):
r"""
Compute the Continuous Ranked Probability Score (CRPS) for a given
discrete probability density.
This function uses a piece-wise linear fit to the approximate posterior
CDF obtained from the predicted quantiles in :code:`y_pred` to
approximate the continuous ranked probability score (CRPS):
.. math::
CRPS(\mathbf{y}, x) = \int_{-\infty}^\infty (F_{x | \mathbf{y}}(x')
- \mathrm{1}_{x < x'})^2 \: dx'
Args:
y_pred: Tensor containing the predicted discrete posterior PDF
with the probabilities for different bins oriented along axis
``bin_axis`` in ``y_pred``.
y_true: Array containing the true point values.
bins: 1D array containing the bins corresponding to the probabilities
in ``y_pred``.
Returns:
Tensor of rank :math:`k - 1` containing the CRPS values for each of the
predictions in ``y_pred``.
"""
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)
n = len(y_pdf.shape)
y_cdf = posterior_cdf(y_pdf, bins, bin_axis=bin_axis)
x = bins
shape = [1] * n_dims
shape[bin_axis] = -1
x = x.reshape(shape)
if len(y_true.shape) < len(y_pdf.shape):
y_true = y_true.unsqueeze(bin_axis)
i = as_type(xp, x > y_true, y_cdf)
crps = trapz(xp, (y_cdf - i) ** 2, x, bin_axis)
return crps
def probability_less_than(y_pred, bins, y, 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)
n = len(y_pred.shape)
x = 0.5 * (bins[1:] + bins[:-1])
mask = x < y
shape = [1] * n
shape[bin_axis] = -1
mask = as_type(xp, reshape(xp, mask, shape), y_pred)
return trapz(xp, mask * y_pred, bins, bin_axis)
def quantile_loss(y_pred, quantiles, y_true, quantile_axis=1):
"""
Calculate the quantile loss for all predicted quantiles.
Args:
y_pred: A k-tensor containing the predicted quantiles along the
axis specified by ``quantile_axis``.
y_true: A tensor of rank k-1 containing the corresponding true
values.
quantiles: A vector or list containing the quantile fractions
corresponding to the predicted quantiles.
quantile_axis: The axis along which ``y_pred`` contains the
the predicted quantiles.
"""
if len(y_pred.shape) == 1:
quantile_axis = 0
xp = get_array_module(y_pred)
n_dims = len(y_pred.shape)
y_true_shape = list(y_pred.shape)
y_true_shape[quantile_axis] = 1
try:
y_true = reshape(xp, y_true, y_true_shape)
except Exception:
raise InvalidDimensionException(
"Could not reshape 'y_true' argument into expected shape "
f"{y_true_shape}."
)
quantiles = to_array(xp, quantiles)
quantiles_shape = [1] * n_dims
quantiles_shape[quantile_axis] = len(quantiles)
quantiles = reshape(xp, quantiles, quantiles_shape)
dy = y_pred - y_true
loss = zeros(xp, dy.shape, like=y_pred)
mask = as_type(xp, dy > 0.0, dy)
loss += mask * ((1.0 - quantiles) * dy)
loss += -(1.0 - mask) * (quantiles * dy)
return loss
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
def pdf_binned(y_pred, quantiles, bins, quantile_axis=1):
"""
Calculate binned representation of the posterior probability density
function (PDF).
The binned PDF is simple calculated by linearly interpolating the
piece-wise linear PDF computed using the :py:meth`pdf` method.
Args:
y_pred: Rank-k Tensor containing the predicted quantiles along the
quantile axis.
quantiles: The quantile fractions corresponding to the predicted
quantiles in y_pred.
bins: Rank-1 tensor containing the ``n_bins`` boundaries for the bins
to use to bin the PDF.
quantile_axis: The axis of y_pred along which the predicted
quantiles are located.
Returns:
Rank-k tensor with ``n_bins - 1`` elements along ``quantile_axis``
containing the probability of the result to fall between the
corresponding bin edges.
"""
if len(y_pred.shape) == 1:
quantile_axis = 0
xp = get_array_module(y_pred)
n = len(y_pred.shape)
x_cdf, y_cdf = cdf(y_pred, quantiles, quantile_axis=quantile_axis)
y_cdf_shape = [1] * n
y_cdf_shape[quantile_axis] = -1
y_cdf = reshape(xp, y_cdf, y_cdf_shape)
selection_l = [slice(0, None)] * n
selection_l[quantile_axis] = slice(0, -1)
selection_l = tuple(selection_l)
selection_r = [slice(0, None)] * n
selection_r[quantile_axis] = slice(1, None)
selection_r = tuple(selection_r)
selection_le = [slice(0, None)] * n
selection_le[quantile_axis] = 0
selection_le = tuple(selection_le)
selection_re = [slice(0, None)] * n
selection_re[quantile_axis] = -1
selection_re = tuple(selection_re)
y_pdf_binned = []
#
# Interpolate CDF for leftmost bin boundary.
#
b_l = bins[0]
mask_r = as_type(xp, (x_cdf[selection_r] >= b_l), y_cdf)
mask_l = as_type(xp, (x_cdf[selection_l] < b_l), y_cdf)
mask = mask_l * mask_r
mask_xr = as_type(xp, xp.sum(mask_r, quantile_axis) == 0.0, mask_r)
mask_xl = as_type(xp, xp.sum(mask_l, quantile_axis) == 0.0, mask_l)
x_cdf_l = xp.sum(x_cdf[selection_l] * mask, quantile_axis)
x_cdf_r = xp.sum(x_cdf[selection_r] * mask, quantile_axis)
d = (x_cdf_r - x_cdf_l) + (1.0 - xp.sum(mask, quantile_axis))
w_cdf_l = (x_cdf_r - b_l) / d
w_cdf_r = (b_l - x_cdf_l) / d
y_cdf_l = (
xp.sum(mask * y_cdf[selection_l] * mask, quantile_axis) * w_cdf_l +
xp.sum(mask * y_cdf[selection_r] * mask, quantile_axis) * w_cdf_r +
mask_xl * y_cdf[selection_le] + mask_xr * y_cdf[selection_re])
for i in range(len(bins) - 1):
b_r = bins[i + 1]
#
# Interpolate CDF for right bin boundary.
#
mask_r = as_type(xp, (x_cdf[selection_r] >= b_r), y_cdf)
mask_l = as_type(xp, (x_cdf[selection_l] < b_r), y_cdf)
mask = mask_l * mask_r
mask_xr = as_type(xp, xp.sum(mask_r, quantile_axis) == 0.0, mask_r)
mask_xl = as_type(xp, xp.sum(mask_l, quantile_axis) == 0.0, mask_l)
x_cdf_l = xp.sum(x_cdf[selection_l] * mask, quantile_axis)
x_cdf_r = xp.sum(x_cdf[selection_r] * mask, quantile_axis)
d = (x_cdf_r - x_cdf_l) + (1.0 - xp.sum(mask, quantile_axis))
w_cdf_l = (x_cdf_r - b_r) / d
w_cdf_r = (b_r - x_cdf_l) / d
y_cdf_r = (
xp.sum(mask * y_cdf[selection_l] * mask, quantile_axis) * w_cdf_l +
xp.sum(mask * y_cdf[selection_r] * mask, quantile_axis) * w_cdf_r +
mask_xl * y_cdf[selection_le] + mask_xr * y_cdf[selection_re])
dy_cdf = expand_dims(xp, y_cdf_r - y_cdf_l, quantile_axis)
y_pdf_binned.append(dy_cdf / (b_r - b_l))
y_cdf_l = y_cdf_r
b_l = b_r
return concatenate(xp, y_pdf_binned, quantile_axis)
def correct_a_priori(y_pred, quantiles, r, quantile_axis=1):
"""
Correct predicted quantiles for a priori.
Args:
y_pred: Rank-k tensor containing the predicted quantiles along
the axis given by 'quantile_axis'.
quantiles: Rank-1 tensor containing the quantile fractions that
correspond to the predicted quantiles.
r: A priori density ratio to use to correct the observations.
quantile_axis: The axis along which the quantile are oriented
in 'y_pred'.
"""
if len(y_pred.shape) == 1:
quantile_axis = 0
xp = get_array_module(y_pred)
n_dims = len(y_pred.shape)
x_pdf, y_pdf = pdf(y_pred, quantiles, quantile_axis=quantile_axis)
selection = [slice(0, None)] * len(y_pred.shape)
selection_c = copy(selection)
selection_c[quantile_axis] = 0
selection_c = tuple(selection_c)
selection_r = copy(selection)
selection_r[quantile_axis] = 1
selection_r = tuple(selection_r)
dy = y_pred[selection_r] - y_pred[selection_c]
dy /= quantiles[1] - quantiles[0]
x_cdf_l = y_pred[selection_c] - 2.0 * quantiles[0] * dy
x_cdf_l = expand_dims(xp, x_cdf_l, quantile_axis)
selection_l = copy(selection)
selection_l[quantile_axis] = -2
selection_l = tuple(selection_l)
selection_c = copy(selection)
selection_c[quantile_axis] = -1
selection_c = tuple(selection_c)
dy = y_pred[selection_c] - y_pred[selection_l]
dy /= quantiles[-1] - quantiles[-2]
x_cdf_r = y_pred[selection_c] + 2.0 * (1.0 - quantiles[-1]) * dy
x_cdf_r = expand_dims(xp, x_cdf_r, quantile_axis)
x_cdf = concatenate(xp, [x_cdf_l, y_pred, x_cdf_r], quantile_axis)
selection_l = [slice(0, None)] * n_dims
selection_l[quantile_axis] = slice(0, -1)
selection_l = tuple(selection_l)
selection_r = [slice(0, None)] * n_dims
selection_r[quantile_axis] = slice(1, None)
selection_r = tuple(selection_r)
x_index = [slice(0, None)] * n_dims
x_index[quantile_axis] = 0
y_pdf_new = r(x_pdf, dist_axis=quantile_axis) * y_pdf
selection = [slice(0, None)] * n_dims
selection[quantile_axis] = slice(1, -1)
selection = tuple(selection)
y_cdf_new = cumtrapz(xp, y_pdf_new[selection], x_cdf, quantile_axis)
selection = [slice(0, None)] * n_dims
selection[quantile_axis] = slice(-1, None)
selection = tuple(selection)
y_cdf_new = y_cdf_new / y_cdf_new[selection]
x_cdf_l = x_cdf[selection_l]
x_cdf_r = x_cdf[selection_r]
y_cdf_new_l = y_cdf_new[selection_l]
y_cdf_new_r = y_cdf_new[selection_r]
y_pred_new = []
for i in range(0, len(quantiles)):
q = quantiles[i]
mask = as_type(xp, (y_cdf_new_l < q) * (y_cdf_new_r >= q), x_cdf_l)
y_new = x_cdf_l * (y_cdf_new_r - q) * mask
y_new += x_cdf_r * (q - y_cdf_new_l) * mask
y_new /= mask * (y_cdf_new_r - y_cdf_new_l) + (1.0 - mask)
y_new = expand_dims(xp, y_new.sum(quantile_axis), quantile_axis)
y_pred_new.append(y_new)
y_pred_new = concatenate(xp, y_pred_new, quantile_axis)
return y_pred_new
def crps(y_pred, y_true, quantiles, quantile_axis=1):
r"""
Compute the Continuous Ranked Probability Score (CRPS) for given
predicted quantiles.
This function uses a piece-wise linear fit to the approximate posterior
CDF obtained from the predicted quantiles in :code:`y_pred` to
approximate the continuous ranked probability score (CRPS):
.. math::
CRPS(\mathbf{y}, x) = \int_{-\infty}^\infty (F_{x | \mathbf{y}}(x')
- \mathrm{1}_{x < x'})^2 \: dx'
Args:
y_pred: Tensor containing the predicted quantiles along the axis
specified by ``quantile_axis``.
y_true: Array containing the true point values.
quantiles: 1D array containing the quantile fractions corresponding
corresponding to the predicted quantiles.
Returns:
Tensor of rank :math:`k - 1` containing the CRPS values for each of the
predictions in ``y_pred``.
"""
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)
y_true_shape = list(x_cdf.shape)
y_true_shape[quantile_axis] = 1
y_true = to_array(xp, y_true)
y_true = reshape(xp, y_true, y_true_shape)
mask = as_type(xp, x_cdf > y_true, y_pred)
ind = ones(xp, x_cdf.shape, like=y_pred) * mask
output_shape = list(x_cdf.shape)
del output_shape[quantile_axis]
integral = zeros(xp, output_shape, like=y_pred)
x_index = [slice(0, None)] * n_dims
y_l = y_cdf[0]
x_index[quantile_axis] = 0
x_l = x_cdf[tuple(x_index)]
ind_l = ind[tuple(x_index)]
for i in range(1, len(y_cdf)):
y_r = y_cdf[i]
x_index[quantile_axis] = i
x_r = x_cdf[tuple(x_index)]
ind_r = ind[tuple(x_index)]
result = (ind_l - y_l) ** 2
result += (ind_r - y_r) ** 2
dx = x_r - x_l
result *= 0.5 * dx
integral += result
y_l = y_r
x_l = x_r
ind_l = ind_r
return integral
