Calculate the Brier scores of an ensemble for exceeding given thresholds.

    According to the threshold decomposition of CRPS, the resulting Brier
    scores can thus be summed along the last axis to calculate CRPS, as

    .. math::
        CRPS(F, x) = \int_z BS(F(z), H(z - x)) dz

    where $F(x) = \int_{z \leq x} p(z) dz$ is the cumulative distribution
    function (CDF) of the forecast distribution $F$, $x$ is a point estimate of
    the true observation (observational error is neglected), $BS$ denotes the
    Brier score and $H(x)$ denotes the Heaviside step function, which we define
    here as equal to 1 for x >= 0 and 0 otherwise.

    It is more efficient to calculate CRPS directly, but this threshold
    decomposition itself provides a useful summary of model quality as a
    function of measurement values.

    The Numba accelerated version of this function is much faster for
    calculating many thresholds simultaneously: it runs in time
    O(N * (E * log(E) + T)), where N is the number of observations, E is the
    ensemble size and T is the number of thresholds.

    The non-Numba accelerated version requires time and space O(N * E * T).

    Parameters
    ----------
    observations : float or array_like
        Observations float or array. Missing values (NaN) are given scores of
        NaN.
    forecasts : float or array_like
        Array of forecasts ensemble members, of the same shape as observations
        except for the extra axis corresponding to the ensemble. If forecasts
        has the same shape as observations, the forecasts are treated as
        deterministic. Missing values (NaN) are ignored.
    threshold : scalar or 1d array_like
        Threshold value(s) at which to calculate exceedence Brier scores.
    issorted : bool, optional
        Optimization flag to indicate that the elements of `ensemble` are
        already sorted along `axis`.
    axis : int, optional
        Axis in forecasts which corresponds to different ensemble members,
        along which to calculate the threshold decomposition.

    Returns
    -------
    out : np.ndarray
        Brier scores at each thresold for each ensemble forecast against the
        observations. If ``threshold`` is a scalar, the result will have the
        same shape as observations. Otherwise, it will have an additional final
        dimension corresponding to the threshold levels.

    References
    ----------
    Gneiting, T. and Ranjan, R. Comparing density forecasts using threshold-
       and quantile-weighted scoring rules. J. Bus. Econ. Stat. 29, 411-422
       (2011). http://www.stat.washington.edu/research/reports/2008/tr533.pdf

    See also
    --------
    crps_ensemble, brier_score

def threshold_brier_score(observations, forecasts, threshold, issorted=False,
                          axis=-1):
    """
    Calculate the Brier scores of an ensemble for exceeding given thresholds.

    According to the threshold decomposition of CRPS, the resulting Brier
    scores can thus be summed along the last axis to calculate CRPS, as

    .. math::
        CRPS(F, x) = \int_z BS(F(z), H(z - x)) dz

    where $F(x) = \int_{z \leq x} p(z) dz$ is the cumulative distribution
    function (CDF) of the forecast distribution $F$, $x$ is a point estimate of
    the true observation (observational error is neglected), $BS$ denotes the
    Brier score and $H(x)$ denotes the Heaviside step function, which we define
    here as equal to 1 for x >= 0 and 0 otherwise.

    It is more efficient to calculate CRPS directly, but this threshold
    decomposition itself provides a useful summary of model quality as a
    function of measurement values.

    The Numba accelerated version of this function is much faster for
    calculating many thresholds simultaneously: it runs in time
    O(N * (E * log(E) + T)), where N is the number of observations, E is the
    ensemble size and T is the number of thresholds.

    The non-Numba accelerated version requires time and space O(N * E * T).

    Parameters
    ----------
    observations : float or array_like
        Observations float or array. Missing values (NaN) are given scores of
        NaN.
    forecasts : float or array_like
        Array of forecasts ensemble members, of the same shape as observations
        except for the extra axis corresponding to the ensemble. If forecasts
        has the same shape as observations, the forecasts are treated as
        deterministic. Missing values (NaN) are ignored.
    threshold : scalar or 1d array_like
        Threshold value(s) at which to calculate exceedence Brier scores.
    issorted : bool, optional
        Optimization flag to indicate that the elements of `ensemble` are
        already sorted along `axis`.
    axis : int, optional
        Axis in forecasts which corresponds to different ensemble members,
        along which to calculate the threshold decomposition.

    Returns
    -------
    out : np.ndarray
        Brier scores at each thresold for each ensemble forecast against the
        observations. If ``threshold`` is a scalar, the result will have the
        same shape as observations. Otherwise, it will have an additional final
        dimension corresponding to the threshold levels.

    References
    ----------
    Gneiting, T. and Ranjan, R. Comparing density forecasts using threshold-
       and quantile-weighted scoring rules. J. Bus. Econ. Stat. 29, 411-422
       (2011). http://www.stat.washington.edu/research/reports/2008/tr533.pdf

    See also
    --------
    crps_ensemble, brier_score
    """
    observations = np.asarray(observations)
    threshold = np.asarray(threshold)
    forecasts = np.asarray(forecasts)

    if axis != -1:
        forecasts = move_axis_to_end(forecasts, axis)

    if forecasts.shape == observations.shape:
        forecasts = forecasts[..., np.newaxis]

    if observations.shape != forecasts.shape[:-1]:
        raise ValueError('observations and forecasts must have matching '
                         'shapes or matching shapes except along `axis=%s`'
                         % axis)

    scalar_threshold = threshold.ndim == 0

    if threshold.ndim > 1:
        raise ValueError('threshold must be scalar or 1-dimensional')
    if threshold.ndim == 1 and not (np.sort(threshold) == threshold).all():
        raise ValueError('1D thresholds must be sorted')

    threshold = threshold.reshape((1,) * observations.ndim + (-1,))

    if not issorted:
        forecasts = np.sort(forecasts, axis=-1)

    result = _threshold_brier_score_core(observations, forecasts, threshold)

    if scalar_threshold:
        result = result.squeeze(axis=-1)

    return result