<filename>quantnn/density.py<gh_stars>0
"""
===============
quantnn.density
===============

This module provides generic functions to manipulate and derive statistics
from probabilistic predictions in the form of discretized probability density
functions as produced by DRNNs.
"""
from quantnn.common import InvalidDimensionException
from quantnn.generic import (
    get_array_module,
    numel,
    expand_dims,
    concatenate,
    trapz,
    cumtrapz,
    cumsum,
    reshape,
    pad_zeros_left,
    as_type,
    zeros,
    sample_uniform,
    argmax,
    take_along_axis
)


def _check_dimensions(n_y, n_b):
    if n_y != n_b - 1:
        raise InvalidDimensionException(
            f"Dimensions of the provided array 'y_pred' ({n_y}) do not match the"
            f" provided number of bin edges ({n_b}). Note that there should be"
            f" one more bin edge than bin values in 'y_pred'."
        )


def normalize(y_pred, bins, bin_axis=1):
    """
    Converts the raw DRNN output to a PDF.

    Args:
        y_pred: Tensor of predictions from a DRNN with the bin-probabilities
            along one of its axes.
        bins: The bin boundaries corresponding to the predictions.
        bin_axis: The index of the tensor axis which contains the predictions
            for each bin.

    Return:
        Tensor with the same shape as ``y_pred`` but with the values
        transformed to represent the PDF corresponding to the predicted
        bin probabilities in ``y_pred``.
    """
    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)

    norm = y_pred.sum(bin_axis)
    norm = expand_dims(xp, norm, bin_axis)

    dx = bins[1:] - bins[:-1]
    shape = [1] * n
    shape[bin_axis] = -1
    dx = reshape(xp, dx, shape)

    return y_pred / norm / dx


def posterior_cdf(y_pdf, bins, bin_axis=1):
    """
    Calculate CDF from predicted probability density function.

    Args:
        y_pdf: Tensor containing the predicted PDFs.
        bins: The bin-boundaries corresponding to the predictions.
        bin_axis: The index of the tensor axis which contains the predictions
            for each bin.

    Return:
        Tensor with the same shape as ``y_pdf`` but with the values
        transformed to represent the CDF corresponding to the predicted
        PDF in ``y_pdf``.
    """
    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)

    y_cdf = cumtrapz(xp, y_pdf, bins, bin_axis)

    selection = [slice(0, None)] * n
    selection[bin_axis] = slice(-1, None)
    y_cdf = y_cdf / y_cdf[tuple(selection)]
    return y_cdf


def posterior_mean(y_pdf, bins, bin_axis=1):
    """
    Calculate posterior mean from predicted PDFs.

    Args:
        y_pdf: Tensor containing the predicted PDFs.
        bins: The bin-boundaries corresponding to the predictions.
        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
        mean value of the PDF.
    """
    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)

    shape = [1] * n
    shape[bin_axis] = -1
    bins_r = reshape(xp, 0.5 * (bins[1:] + bins[:-1]), shape)

    return trapz(xp, bins_r * y_pdf, bins, bin_axis)


def posterior_median(y_pred, bins, bin_axis=1):
    """
    Calculate the posterior median from predicted PDFs.

    Args:
        y_pdf: Tensor containing the predicted PDFs.
        bins: The bin-boundaries corresponding to the predictions.
        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
        median value of the PDFs.
    """
    quantiles = posterior_quantiles(y_pred, bins, [0.5], bin_axis=bin_axis)
    n = len(y_pred.shape)
    selection = [slice(0, None)] * n
    selection[bin_axis] = 0
    return quantiles[tuple(selection)]


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 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 probability_larger_than(y_pred, bins, quantiles, bin_axis=1):
    """
    Calculate the probability of a sample being larger 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 larger than
        the given value ``y``.
    """
    return 1.0 - probability_less_than(y_pred, bins, quantiles, bin_axis=bin_axis)


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 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 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 posterior_maximum(y_pdf, bins, bin_axis=1):
    """
    Calculate maximum of the posterior distribution.

    Args:
        y_pdf: Rank-``k`` Tensor containing the normalized, discrete
            posterior pdf.
        bins: Array containing the bin boundaries corresponding to 'y_pdf'.
        bin_axis: The axis along which the bins of the posterior PDF are
            oriented.

    Return:
        Tensor of rank ``k - 1`` containing the values corresponding to the
        maxima of the given posterior.
    """
    if len(y_pdf.shape) == 1:
        bin_axis = 0
    xp = get_array_module(y_pdf)
    x = 0.5 * (bins[1:] + bins[:-1])
    indices = argmax(xp, y_pdf, axes=bin_axis)
    return x[indices]


<filename>quantnn/density.py<gh_stars>0
"""
===============
quantnn.density
===============

This module provides generic functions to manipulate and derive statistics
from probabilistic predictions in the form of discretized probability density
functions as produced by DRNNs.
"""
from quantnn.common import InvalidDimensionException
from quantnn.generic import (
    get_array_module,
    numel,
    expand_dims,
    concatenate,
    trapz,
    cumtrapz,
    cumsum,
    reshape,
    pad_zeros_left,
    as_type,
    zeros,
    sample_uniform,
    argmax,
    take_along_axis
)


def _check_dimensions(n_y, n_b):
    if n_y != n_b - 1:
        raise InvalidDimensionException(
            f"Dimensions of the provided array 'y_pred' ({n_y}) do not match the"
            f" provided number of bin edges ({n_b}). Note that there should be"
            f" one more bin edge than bin values in 'y_pred'."
        )


def normalize(y_pred, bins, bin_axis=1):
    """
    Converts the raw DRNN output to a PDF.

    Args:
        y_pred: Tensor of predictions from a DRNN with the bin-probabilities
            along one of its axes.
        bins: The bin boundaries corresponding to the predictions.
        bin_axis: The index of the tensor axis which contains the predictions
            for each bin.

    Return:
        Tensor with the same shape as ``y_pred`` but with the values
        transformed to represent the PDF corresponding to the predicted
        bin probabilities in ``y_pred``.
    """
    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)

    norm = y_pred.sum(bin_axis)
    norm = expand_dims(xp, norm, bin_axis)

    dx = bins[1:] - bins[:-1]
    shape = [1] * n
    shape[bin_axis] = -1
    dx = reshape(xp, dx, shape)

    return y_pred / norm / dx


def posterior_cdf(y_pdf, bins, bin_axis=1):
    """
    Calculate CDF from predicted probability density function.

    Args:
        y_pdf: Tensor containing the predicted PDFs.
        bins: The bin-boundaries corresponding to the predictions.
        bin_axis: The index of the tensor axis which contains the predictions
            for each bin.

    Return:
        Tensor with the same shape as ``y_pdf`` but with the values
        transformed to represent the CDF corresponding to the predicted
        PDF in ``y_pdf``.
    """
    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)

    y_cdf = cumtrapz(xp, y_pdf, bins, bin_axis)

    selection = [slice(0, None)] * n
    selection[bin_axis] = slice(-1, None)
    y_cdf = y_cdf / y_cdf[tuple(selection)]
    return y_cdf


def posterior_mean(y_pdf, bins, bin_axis=1):
    """
    Calculate posterior mean from predicted PDFs.

    Args:
        y_pdf: Tensor containing the predicted PDFs.
        bins: The bin-boundaries corresponding to the predictions.
        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
        mean value of the PDF.
    """
    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)

    shape = [1] * n
    shape[bin_axis] = -1
    bins_r = reshape(xp, 0.5 * (bins[1:] + bins[:-1]), shape)

    return trapz(xp, bins_r * y_pdf, bins, bin_axis)


def posterior_median(y_pred, bins, bin_axis=1):
    """
    Calculate the posterior median from predicted PDFs.

    Args:
        y_pdf: Tensor containing the predicted PDFs.
        bins: The bin-boundaries corresponding to the predictions.
        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
        median value of the PDFs.
    """
    quantiles = posterior_quantiles(y_pred, bins, [0.5], bin_axis=bin_axis)
    n = len(y_pred.shape)
    selection = [slice(0, None)] * n
    selection[bin_axis] = 0
    return quantiles[tuple(selection)]


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 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 probability_larger_than(y_pred, bins, quantiles, bin_axis=1):
    """
    Calculate the probability of a sample being larger 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 larger than
        the given value ``y``.
    """
    return 1.0 - probability_less_than(y_pred, bins, quantiles, bin_axis=bin_axis)


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 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 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 posterior_maximum(y_pdf, bins, bin_axis=1):
    """
    Calculate maximum of the posterior distribution.

    Args:
        y_pdf: Rank-``k`` Tensor containing the normalized, discrete
            posterior pdf.
        bins: Array containing the bin boundaries corresponding to 'y_pdf'.
        bin_axis: The axis along which the bins of the posterior PDF are
            oriented.

    Return:
        Tensor of rank ``k - 1`` containing the values corresponding to the
        maxima of the given posterior.
    """
    if len(y_pdf.shape) == 1:
        bin_axis = 0
    xp = get_array_module(y_pdf)
    x = 0.5 * (bins[1:] + bins[:-1])
    indices = argmax(xp, y_pdf, axes=bin_axis)
    return x[indices]


=============== quantnn.density ===============  This module provides generic functions to manipulate and derive statistics from probabilistic predictions in the form of discretized probability density functions as produced by DRNNs. Converts the raw DRNN output to a PDF.      Args:         y_pred: Tensor of predictions from a DRNN with the bin-probabilities             along one of its axes.         bins: The bin boundaries corresponding to the predictions.         bin_axis: The index of the tensor axis which contains the predictions             for each bin.      Return:         Tensor with the same shape as ``y_pred`` but with the values         transformed to represent the PDF corresponding to the predicted         bin probabilities in ``y_pred``. Calculate CDF from predicted probability density function.      Args:         y_pdf: Tensor containing the predicted PDFs.         bins: The bin-boundaries corresponding to the predictions.         bin_axis: The index of the tensor axis which contains the predictions             for each bin.      Return:         Tensor with the same shape as ``y_pdf`` but with the values         transformed to represent the CDF corresponding to the predicted         PDF in ``y_pdf``. Calculate posterior mean from predicted PDFs.      Args:         y_pdf: Tensor containing the predicted PDFs.         bins: The bin-boundaries corresponding to the predictions.         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         mean value of the PDF. Calculate the posterior median from predicted PDFs.      Args:         y_pdf: Tensor containing the predicted PDFs.         bins: The bin-boundaries corresponding to the predictions.         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         median value of the PDFs. 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. 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``. Calculate the probability of a sample being larger 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 larger than         the given value ``y``. 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. 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``. Evaluates the quantile function at given y values. Calculate maximum of the posterior distribution.      Args:         y_pdf: Rank-``k`` Tensor containing the normalized, discrete             posterior pdf.         bins: Array containing the bin boundaries corresponding to 'y_pdf'.         bin_axis: The axis along which the bins of the posterior PDF are             oriented.      Return:         Tensor of rank ``k - 1`` containing the values corresponding to the         maxima of the given posterior.