Coverage for src/scores/continuous/quantile_loss_impl.py: 100%

1"""

2 Implementation of quantile loss (score)

3"""

4 from typing import Optional, Sequence

5

6 import xarray as xr

7

8 from scores.functions import apply_weights

9 from scores.typing import XarrayLike

10 from scores.utils import check_dims, gather_dimensions

11

12

13 def quantile_score(

14 fcst: XarrayLike,

15 obs: XarrayLike,

16 alpha: float,

17 reduce_dims: Optional[Sequence[str]] = None,

18 preserve_dims: Optional[Sequence[str]] = None,

19 weights: XarrayLike = None,

20) -> XarrayLike:

21 """

22 Calculates a score that targets alpha-quantiles.

23 Use with alpha = 0.5 for forecasts of the median.

24 Use with alpha = 0.9 for forecasts of the 90th percentile.

25

26 Args:

27 fcst: array of forecasts

28 obs: array of observations

29 alpha: A value between 0 and 1 (exclusive)

30 reduce_dims: Optionally specify which dimensions to reduce when

31 calculating the quantile score. All other dimensions will be preserved. As a

32 special case, 'all' will allow all dimensions to be reduced. Only one

33 of `reduce_dims` and `preserve_dims` can be supplied. The default behaviour

34 if neither are supplied is to reduce all dims.

35 preserve_dims: Optionally specify which dimensions to preserve when calculating

36 quantile score. All other dimensions will be reduced. As a special case, 'all'

37 will allow all dimensions to be preserved. In this case, the result will be in

38 the same shape/dimensionality as the forecast, and the errors will be the quantile

39 score at each point (i.e. single-value comparison against observed), and the

40 forecast and observed dimensions must match precisely. Only one of `reduce_dims`

41 and `preserve_dims` can be supplied. The default behaviour if neither are supplied

42 is to reduce all dims.

43 weights: Optionally provide an array for weighted averaging (e.g. by area, by latitude,

44 by population, custom)

45

46 Returns:

47 A DataArray with values being the mean generalised piecewise linear (GPL)

48 scoring function, with the dimensions specified in `dims`.

49 If `dims` is `None`, the returned DataArray will have only one element,

50 the overall mean GPL score.

51

52 Raises:

53 ValueError: if `alpha` is not between 0 and 1.

54

55 Notes:

56

57 .. math::

58

59 gpl(x) = \\begin{{cases}}\\alpha * (-x) & x \\leq 0\\\\

60 (1-\\alpha) x & x > 0\\end{{cases}}

61

62 where:

63

64 - :math:`\\alpha` is the targeted quantile.

65 - :math:`x` is the difference, fcst - obs

66

67 References:

68 T. Geinting, "Making and evaluating point forecasts",

69 J. Amer. Stat. Assoc., Vol. 106 No. 494 (June 2011), pp. 754--755,

70 Theorem 9

71

72 """

73 specified_dims = reduce_dims or preserve_dims

74 # check requested dims are a subset of fcst dimensions

75 if specified_dims is not None:

76 check_dims(xr_data=fcst, expected_dims=specified_dims, mode="superset")

77 # check obs dimensions are a subset of fcst dimensions

78 check_dims(xr_data=obs, expected_dims=fcst.dims, mode="subset")

79

80 # check that alpha is between 0 and 1 as required

81 if (alpha <= 0) or (alpha >= 1):

82 raise ValueError("alpha is not between 0 and 1")

83

84 # Do this operation once to save compute time

85 diff = fcst - obs

86

87 # calculate the score applicable when fcst <= obs

88 score_fcst_lte_obs = alpha * (-diff)

89

90 # calculate the score applicable when fcst > obs

91 score_fcst_ge_obs = (1 - alpha) * diff

92

93 result = xr.where(diff > 0, score_fcst_ge_obs, score_fcst_lte_obs)

94

95 reduce_dims = gather_dimensions(fcst.dims, obs.dims, reduce_dims, preserve_dims) # type: ignore[assignment]

96 results = apply_weights(result, weights)

97 score = results.mean(dim=reduce_dims)

98

99 return score