def calculate_quantiles(y_pred):
module = get_array_module(y_pred)
new_quantiles = to_array(module, quantiles, like=y_pred)
current_quantiles = to_array(module, self.quantiles, like=y_pred)
return qq.posterior_quantiles(
y_pred,
quantiles=current_quantiles,
new_quantiles=new_quantiles,
quantile_axis=self.quantile_axis,
)
def _post_process_prediction(self, y_pred, bins=None, key=None):
module = get_array_module(y_pred)
if bins is not None:
bins = to_array(module, bins, like=y_pred)
else:
if isinstance(self.bins, dict):
bins = to_array(module, self.bins[key], like=y_pred)
else:
bins = to_array(module, self.bins, like=y_pred)
module = get_array_module(y_pred)
y_pred = softmax(module, y_pred, axis=1)
bins = to_array(module, bins, like=y_pred)
y_pred = qd.normalize(y_pred, bins, bin_axis=self.bin_axis)
return y_pred
def calculate_quantile_function(y_pred, bins):
module = get_array_module(y_pred)
bins = to_array(module, bins, like=y_pred)
return qd.quantile_function(y_pred,
y,
bins,
bin_axis=self.bin_axis)
def calculate_samples(y_pred, bins):
module = get_array_module(y_pred)
bins = to_array(module, bins, like=y_pred)
return qd.sample_posterior(y_pred,
bins,
n_samples=n_samples,
bin_axis=self.bin_axis)
def calculate_probability(y_pred, bins):
module = get_array_module(y_pred)
bins = to_array(module, bins, like=y_pred)
return qd.probability_larger_than(y_pred,
bins,
y,
bin_axis=self.bin_axis)
def calculate_quantiles(y_pred, bins):
module = get_array_module(y_pred)
bins = to_array(module, bins, like=y_pred)
return qd.posterior_quantiles(y_pred,
bins,
quantiles,
bin_axis=self.bin_axis)
def calculate_crps(y_pred):
module = get_array_module(y_pred)
quantiles = to_array(module, self.quantiles, like=y_pred)
return qq.crps(y_pred,
y_true,
quantiles,
quantile_axis=self.quantile_axis)
def calculate_prob(y_pred):
module = get_array_module(y_pred)
quantiles = to_array(module, self.quantiles, like=y_pred)
return qq.probability_less_than(y_pred,
quantiles,
y,
quantile_axis=self.quantile_axis)
def calculate_samples(y_pred):
module = get_array_module(y_pred)
quantiles = to_array(module, self.quantiles, like=y_pred)
return qq.sample_posterior(y_pred,
quantiles,
n_samples=n_samples,
quantile_axis=self.quantile_axis)
def predict(self, x, device="cpu"):
"""
Evaluate the model.
Args:
x: The input data for which to evaluate the data.
device: The device on which to evaluate the prediction.
Returns:
The model prediction converted to numpy array.
"""
# Determine device to use
w = next(iter(self.parameters())).data
if isinstance(x, torch.Tensor):
x_torch = x
else:
x_torch = to_array(torch, x, like=w)
self.to(x_torch.device)
if x_torch.requires_grad:
y = self(x_torch)
else:
with torch.no_grad():
y = self(x_torch)
return y
def test_to_array(backend):
"""
Converts numpy array to array of given backend and ensures
that corresponding module object matches the backend.
"""
x = np.arange(10)
array = to_array(backend, x)
assert get_array_module(array) == backend
def test_take_along_axis(backend):
x = arange(backend, 0, 10.1, 1).reshape(-1, 1)
indices = to_array(backend, [[0], [1], [2], [3]])
i = take_along_axis(backend, x, indices, 0)
assert i[0] == 0
assert i[1] == 1
assert i[2] == 2
assert i[3] == 3
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 pdf(self, y_pred):
"""
Calculate PDF from predicted logits.
Args:
y_pred: Tensor containing the logit values predicted by
the neural network model.
"""
module = get_array_module(y_pred)
bins = to_array(module, self.bins, like=y_pred)
return qd.pdf(y_pred,
bins,
bin_axis=self.bin_axis)
def pdf(self, y_pred):
"""
Calculate PDF from predicted quantiles.
Args:
y_pred: Tensor containing the quantiles predicted by the NN
model.
"""
module = get_array_module(y_pred)
quantiles = to_array(module, self.quantiles, like=y_pred)
return qq.pdf(y_pred,
quantiles,
quantile_axis=self.quantile_axis)
def crps(self, y_pred, y_true):
"""
Calculate the CRPS score from predicted quantiles.
Args:
y_pred: Tensor containing the quantiles predicted by the NN
model.
y_true: Tensor containing the true values.
"""
module = get_array_module(y_pred)
quantiles = to_array(module, self.quantiles, like=y_pred)
return qq.crps(
y_pred, y_true, quantiles, quantile_axis=self.quantile_axis
)
def posterior_mean(self, y_pred):
"""
Calculate the posterior mean from predicted quantiles.
Args:
y_pred: Tensor containing the logit values predicted by
the neural network model.
"""
module = get_array_module(y_pred)
bins = to_array(module, self.bins, like=y_pred)
return qd.posterior_mean(
y_pred, bins,
bin_axis=self.bin_axis
)
def crps(self, y_pred, y_true):
"""
Calculate the CRPS score from predicted quantiles.
Args:
y_pred: Tensor containing the logit values predicted by
the neural network model.
y_true: Tensor containing the true values.
"""
module = get_array_module(y_pred)
bins = to_array(module, self.bins, like=y_pred)
return qd.crps(
y_pred, y_true, bins, bin_axis=self.bin_axis
)
def predict(self, y_pred):
"""
Apply post processing to model prediction. Converts predicted
logits to normalized probabilities.
Args:
y_pred: Tensor containing the logit values predicted by
the neural network model.
"""
module = get_array_module(y_pred)
bins = to_array(module, self.bins, like=y_pred)
y_pred = softmax(module, y_pred, axis=1)
y_pred = qd.normalize(y_pred, bins, bin_axis=self.bin_axis)
return y_pred
def sample_posterior(self, y_pred, n_samples=1):
"""
Sample retrieval values from posterior distribution.
Args:
y_pred: Tensor containing the quantiles predicted by the NN
model.
n_samples: The number of samples to produce.
"""
module = get_array_module(y_pred)
quantiles = to_array(module, self.quantiles, like=y_pred)
return qq.sample_posterior(
y_pred, quantiles, n_samples=n_samples,
quantile_axis=self.quantile_axis
)
def sample_posterior(self, y_pred, n_samples=1):
"""
Sample retrieval values from posterior distribution.
Args:
y_pred: Tensor containing the logit values predicted by
the neural network model.
n_samples: The number of samples to produce.
"""
module = get_array_module(y_pred)
bins = to_array(module, self.bins, like=y_pred)
return qd.sample_posterior(
y_pred, bins, n_samples=n_samples,
bin_axis=self.bin_axis
)
def test_scatter_add(backend):
x = zeros(backend, (3, 3))
y = ones(backend, (2, 3))
indices = to_array(backend, [0, 2])
z = scatter_add(backend, x, indices, y, 0)
assert np.isclose(z[0, 0], 1.0)
assert np.isclose(z[1, 0], 0.0)
assert np.isclose(z[2, 0], 1.0)
x = zeros(backend, (3, 3))
y = ones(backend, (3, 2))
z = scatter_add(backend, x, indices, y, 1)
assert np.isclose(z[0, 0], 1.0)
assert np.isclose(z[0, 1], 0.0)
assert np.isclose(z[0, 2], 1.0)
def probability_larger_than(self, y_pred, y):
"""
Calculate the probability that the retrieval value is larger than
a given threshold.
Args:
y_pred: Tensor containing the logit values predicted by
the neural network model.
y: The scalar threshold value.
"""
module = get_array_module(y_pred)
bins = to_array(module, self.bins, like=y_pred)
return qd.probability_larger_than(
y_pred=y_pred, bins=bins, y=y,
bin_axis=self.bin_axis
)
def probability_less_than(self, y_pred, y):
"""
Calculate the probability that the retrieval value is less than
a given threshold.
Args:
y_pred: Tensor containing the quantiles predicted by the NN
model.
y: The scalar threshold value.
"""
module = get_array_module(y_pred)
quantiles = to_array(module, self.quantiles, like=y_pred)
return qq.probability_less_than(
y_pred, quantiles, y,
quantile_axis=self.quantile_axis
)
def test_quantile_loss(xp):
"""
Tests calculation of the quantile loss function.
"""
#
# 1D predictions
#
quantiles = arange(xp, 0.1, 0.91, 0.1)
y_pred = arange(xp, 1.0, 9.1, 1.0)
y_true = to_array(xp, [5.0])
loss = quantile_loss(y_pred, quantiles, y_true)
assert np.isclose(loss.mean(), to_array(xp, [0.444444]))
#
# 2D predictions
#
quantiles = arange(xp, 0.1, 0.91, 0.1)
y_pred = eo.repeat(arange(xp, 1.0, 9.1, 1.0), 'q -> w q', w=10)
y_true = eo.repeat(to_array(xp, [5.0]), 'q -> w q', w=10)
loss = quantile_loss(y_pred, quantiles, y_true)
assert np.isclose(loss.mean(), to_array(xp, [0.444444]))
#
# 3D predictions, quantiles along last axis
#
quantiles = arange(xp, 0.1, 0.91, 0.1)
y_pred = eo.repeat(arange(xp, 1.0, 9.1, 1.0), 'q -> h w q', h=10, w=10)
y_true = eo.repeat(to_array(xp, [5.0]), 'q -> h w q', h=10, w=10)
loss = quantile_loss(y_pred, quantiles, y_true, quantile_axis=-1)
assert np.isclose(loss.mean(), to_array(xp, [0.444444]))
#
# 3D predictions, quantiles along first axis
#
quantiles = arange(xp, 0.1, 0.91, 0.1)
y_pred = eo.repeat(arange(xp, 1.0, 9.1, 1.0), 'q -> h q w', h=10, w=10)
y_true = eo.repeat(to_array(xp, [5.0]), 'q -> h q w', h=10, w=10)
loss = quantile_loss(y_pred, quantiles, y_true, quantile_axis=1)
assert np.isclose(loss.mean(), to_array(xp, [0.444444]))
def posterior_quantiles(self, y_pred, new_quantiles):
"""
Calculate quantiles of the posterior distribution.
Args:
y_pred: Tensor containing the logit values predicted by
the neural network model.
y: The scalar threshold value.
"""
module = get_array_module(y_pred)
bins = to_array(module, self.bins, like=y_pred)
return qd.posterior_quantiles(
y_pred,
bins=bins,
quantiles=new_quantiles,
bin_axis=self.bin_axis
)
def posterior_quantiles(self, y_pred, new_quantiles):
"""
Calculate new quantiles.
Args:
y_pred: Tensor containing the quantiles predicted by the NN
model.
new_quantiles: Array containing the new quantiles to compute.
"""
module = get_array_module(y_pred)
quantiles = to_array(module, self.quantiles, like=y_pred)
return qq.posterior_quantiles(
y_pred,
quantiles=quantiles,
new_quantiles=new_quantiles,
quantile_axis=self.quantile_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 test_posterior_quantiles(xp):
"""
Test interpolation of posterior quantiles.
"""
#
# 1D predictions
#
quantiles = arange(xp, 0.1, 0.91, 0.1)
new_quantiles = quantiles[:-1] + 0.05
y_pred = arange(xp, 1.0, 9.1, 1.0)
y_q = posterior_quantiles(y_pred, quantiles, quantiles)
assert np.all(np.isclose(y_pred, y_q))
y_q = posterior_quantiles(y_pred, quantiles, new_quantiles)
y_pred_i = 0.5 * (y_pred[1:] + y_pred[:-1])
assert np.all(np.isclose(y_pred_i, y_q))
new_quantiles = to_array(xp, [0.0])
y_q = posterior_quantiles(y_pred, quantiles, new_quantiles)
assert np.all(np.isclose(y_pred[0], y_q))
new_quantiles = to_array(xp, [10.0])
y_q = posterior_quantiles(y_pred, quantiles, new_quantiles)
assert np.all(np.isclose(y_pred[-1], y_q))
#
# 2D predictions
#
quantiles = arange(xp, 0.1, 0.91, 0.1)
new_quantiles = quantiles[:-1] + 0.05
y_pred = eo.repeat(arange(xp, 1.0, 9.1, 1.0), 'q -> w q', w=10)
y_q = posterior_quantiles(y_pred, quantiles, quantiles)
assert np.all(np.isclose(y_pred, y_q))
y_q = posterior_quantiles(y_pred, quantiles, new_quantiles)
y_pred_i = 0.5 * (y_pred[:, 1:] + y_pred[:, :-1])
assert np.all(np.isclose(y_pred_i, y_q))
new_quantiles = to_array(xp, [0.0])
y_q = posterior_quantiles(y_pred, quantiles, new_quantiles)
assert np.all(np.isclose(y_pred[:, 0], y_q))
new_quantiles = to_array(xp, [10.0])
y_q = posterior_quantiles(y_pred, quantiles, new_quantiles)
assert np.all(np.isclose(y_pred[:, -1], y_q))
#
# 3D predictions, quantiles along last axis
#
quantiles = arange(xp, 0.1, 0.91, 0.1)
new_quantiles = quantiles[:-1] + 0.05
y_pred = eo.repeat(arange(xp, 1.0, 9.1, 1.0), 'q -> h w q', h=10, w=10)
y_q = posterior_quantiles(y_pred, quantiles, quantiles,
quantile_axis=2)
assert np.all(np.isclose(y_pred, y_q))
y_q = posterior_quantiles(y_pred, quantiles, new_quantiles,
quantile_axis=2)
y_pred_i = 0.5 * (y_pred[:, :, 1:] + y_pred[:, :, :-1])
assert np.all(np.isclose(y_pred_i, y_q))
new_quantiles = to_array(xp, [0.0])
y_q = posterior_quantiles(y_pred, quantiles, new_quantiles,
quantile_axis=2)
assert np.all(np.isclose(y_pred[:, :, 0], y_q))
new_quantiles = to_array(xp, [10.0])
y_q = posterior_quantiles(y_pred, quantiles, new_quantiles,
quantile_axis=2)
assert np.all(np.isclose(y_pred[:, :, -1], y_q))
#
# 3D predictions, quantiles along first axis
#
quantiles = arange(xp, 0.1, 0.91, 0.1)
new_quantiles = quantiles[:-1] + 0.05
y_pred = eo.repeat(arange(xp, 1.0, 9.1, 1.0), 'q -> h q w', h=10, w=10)
y_q = posterior_quantiles(y_pred, quantiles, quantiles,
quantile_axis=1)
assert np.all(np.isclose(y_pred, y_q))
y_q = posterior_quantiles(y_pred, quantiles, new_quantiles,
quantile_axis=1)
y_pred_i = 0.5 * (y_pred[:, 1:, :] + y_pred[:, :-1, :])
assert np.all(np.isclose(y_pred_i, y_q))
new_quantiles = to_array(xp, [0.0])
y_q = posterior_quantiles(y_pred, quantiles, new_quantiles,
quantile_axis=1)
assert np.all(np.isclose(y_pred[:, 0, :], y_q))
new_quantiles = to_array(xp, [10.0])
y_q = posterior_quantiles(y_pred, quantiles, new_quantiles,
quantile_axis=1)
assert np.all(np.isclose(y_pred[:, -1, :], y_q))
def test_correct_a_priori(xp):
"""
Test correcting for a priori.
"""
r_x = to_array(xp, [-1, 4.99, 5.01, 10])
r_y = to_array(xp, [1, 1, 1, 1])
r = LookupTable(r_x, r_y)
#
# 1D predictions
#
quantiles = arange(xp, 0.1, 0.91, 0.1)
y_pred = arange(xp, 1.0, 9.1, 1.0)
y_pred_new = correct_a_priori(y_pred, quantiles, r)
assert np.isclose(y_pred_new[0], y_pred[0])
assert np.isclose(y_pred_new[-1], y_pred[-1])
#
# 2D predictions
#
quantiles = arange(xp, 0.1, 0.91, 0.1)
y_pred = eo.repeat(arange(xp, 1.0, 9.1, 1.0), 'q -> h q', h=10)
r_x = to_array(xp, [-1, 4.99, 5.01, 10])
r_y = to_array(xp, [1, 1, 1, 1])
y_pred_new = correct_a_priori(y_pred, quantiles, r)
assert np.isclose(y_pred_new[0, 0], y_pred[0, 0])
assert np.isclose(y_pred_new[-1, -1], y_pred[-1, -1])
quantiles = arange(xp, 0.1, 0.91, 0.1)
y_pred = eo.repeat(arange(xp, 1.0, 9.1, 1.0), 'q -> q h', h=10)
r_x = to_array(xp, [-1, 4.99, 5.01, 10])
r_y = to_array(xp, [1, 1, 1, 1])
y_pred_new = correct_a_priori(y_pred, quantiles, r, quantile_axis=0)
assert np.isclose(y_pred_new[0, 0], y_pred[0, 0])
assert np.isclose(y_pred_new[-1, -1], y_pred[-1, -1])
#
# 2D predictions
#
quantiles = arange(xp, 0.1, 0.91, 0.1)
y_pred = eo.repeat(arange(xp, 1.0, 9.1, 1.0), 'q -> h q w', h=10, w=10)
r_x = to_array(xp, [-1, 4.99, 5.01, 10])
r_y = to_array(xp, [1, 1, 1, 1])
y_pred_new = correct_a_priori(y_pred, quantiles, r)
assert np.isclose(y_pred_new[0, 0, 0], y_pred[0, 0, 0])
assert np.isclose(y_pred_new[-1, -1, -1], y_pred[-1, -1, -1])
quantiles = arange(xp, 0.1, 0.91, 0.1)
y_pred = eo.repeat(arange(xp, 1.0, 9.1, 1.0), 'q -> w h q', h=10, w=10)
r_x = to_array(xp, [-1, 4.99, 5.01, 10])
r_y = to_array(xp, [1, 1, 1, 1])
y_pred_new = correct_a_priori(y_pred, quantiles, r, quantile_axis=-1)
assert np.isclose(y_pred_new[0, 0, 0], y_pred[0, 0, 0])
assert np.isclose(y_pred_new[-1, -1, -1], y_pred[-1, -1, -1])
