Aluode/PerceptionLabPortable

1# Authors: The scikit-learn developers
2# SPDX-License-Identifier: BSD-3-Clause
3
4from ..utils._array_api import (
5    _find_matching_floating_dtype,
6    get_namespace_and_device,
7)
8
9
10def _weighted_percentile(array, sample_weight, percentile_rank=50, xp=None):
11    """Compute the weighted percentile with method 'inverted_cdf'.
12
13    When the percentile lies between two data points of `array`, the function returns
14    the lower value.
15
16    If `array` is a 2D array, the `values` are selected along axis 0.
17
18    `NaN` values are ignored by setting their weights to 0. If `array` is 2D, this
19    is done in a column-isolated manner: a `NaN` in the second column, does not impact
20    the percentile computed for the first column even if `sample_weight` is 1D.
21
22        .. versionchanged:: 0.24
23            Accepts 2D `array`.
24
25        .. versionchanged:: 1.7
26            Supports handling of `NaN` values.
27
28    Parameters
29    ----------
30    array : 1D or 2D array
31        Values to take the weighted percentile of.
32
33    sample_weight: 1D or 2D array
34        Weights for each value in `array`. Must be same shape as `array` or of shape
35        `(array.shape[0],)`.
36
37    percentile_rank: int or float, default=50
38        The probability level of the percentile to compute, in percent. Must be between
39        0 and 100.
40
41    xp : array_namespace, default=None
42        The standard-compatible namespace for `array`. Default: infer.
43
44    Returns
45    -------
46    percentile : scalar or 0D array if `array` 1D (or 0D), array if `array` 2D
47        Weighted percentile at the requested probability level.
48    """
49    xp, _, device = get_namespace_and_device(array)
50    # `sample_weight` should follow `array` for dtypes
51    floating_dtype = _find_matching_floating_dtype(array, xp=xp)
52    array = xp.asarray(array, dtype=floating_dtype, device=device)
53    sample_weight = xp.asarray(sample_weight, dtype=floating_dtype, device=device)
54
55    n_dim = array.ndim
56    if n_dim == 0:
57        return array
58    if array.ndim == 1:
59        array = xp.reshape(array, (-1, 1))
60    # When sample_weight 1D, repeat for each array.shape[1]
61    if array.shape != sample_weight.shape and array.shape[0] == sample_weight.shape[0]:
62        sample_weight = xp.tile(sample_weight, (array.shape[1], 1)).T
63    # Sort `array` and `sample_weight` along axis=0:
64    sorted_idx = xp.argsort(array, axis=0)
65    sorted_weights = xp.take_along_axis(sample_weight, sorted_idx, axis=0)
66
67    # Set NaN values in `sample_weight` to 0. Only perform this operation if NaN
68    # values present to avoid temporary allocations of size `(n_samples, n_features)`.
69    n_features = array.shape[1]
70    largest_value_per_column = array[
71        sorted_idx[-1, ...], xp.arange(n_features, device=device)
72    ]
73    # NaN values get sorted to end (largest value)
74    if xp.any(xp.isnan(largest_value_per_column)):
75        sorted_nan_mask = xp.take_along_axis(xp.isnan(array), sorted_idx, axis=0)
76        sorted_weights[sorted_nan_mask] = 0
77
78    # Compute the weighted cumulative distribution function (CDF) based on
79    # `sample_weight` and scale `percentile_rank` along it.
80    #
81    # Note: we call `xp.cumulative_sum` on the transposed `sorted_weights` to
82    # ensure that the result is of shape `(n_features, n_samples)` so
83    # `xp.searchsorted` calls take contiguous inputs as a result (for
84    # performance reasons).
85    weight_cdf = xp.cumulative_sum(sorted_weights.T, axis=1)
86    adjusted_percentile_rank = percentile_rank / 100 * weight_cdf[..., -1]
87
88    # Ignore leading `sample_weight=0` observations when `percentile_rank=0` (#20528)
89    mask = adjusted_percentile_rank == 0
90    adjusted_percentile_rank[mask] = xp.nextafter(
91        adjusted_percentile_rank[mask], adjusted_percentile_rank[mask] + 1
92    )
93    # For each feature with index j, find sample index i of the scalar value
94    # `adjusted_percentile_rank[j]` in 1D array `weight_cdf[j]`, such that:
95    # weight_cdf[j, i-1] < adjusted_percentile_rank[j] <= weight_cdf[j, i].
96    percentile_indices = xp.stack(
97        [
98            xp.searchsorted(
99                weight_cdf[feature_idx, ...], adjusted_percentile_rank[feature_idx]
100            )
101            for feature_idx in range(weight_cdf.shape[0])
102        ],
103    )
104    # In rare cases, `percentile_indices` equals to `sorted_idx.shape[0]`
105    max_idx = sorted_idx.shape[0] - 1
106    percentile_indices = xp.clip(percentile_indices, 0, max_idx)
107
108    col_indices = xp.arange(array.shape[1], device=device)
109    percentile_in_sorted = sorted_idx[percentile_indices, col_indices]
110
111    result = array[percentile_in_sorted, col_indices]
112
113    return result[0] if n_dim == 1 else result
114
115
116# TODO: refactor to do the symmetrisation inside _weighted_percentile to avoid
117# sorting the input array twice.
118def _averaged_weighted_percentile(array, sample_weight, percentile_rank=50, xp=None):
119    return (
120        _weighted_percentile(array, sample_weight, percentile_rank, xp=xp)
121        - _weighted_percentile(-array, sample_weight, 100 - percentile_rank, xp=xp)
122    ) / 2
123