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])