r"""Computes the quantile function of an empirical distribution Parameters ---------- qs: array-like, shape (n,) Quantiles at which the quantile function is evaluated cws: array-like, shape (m, ...) cumulative weights of the 1D empirical distribution, if batched, mus
(qs, cws, xs)
| 17 | |
| 18 | |
| 19 | def quantile_function(qs, cws, xs): |
| 20 | r"""Computes the quantile function of an empirical distribution |
| 21 | |
| 22 | Parameters |
| 23 | ---------- |
| 24 | qs: array-like, shape (n,) |
| 25 | Quantiles at which the quantile function is evaluated |
| 26 | cws: array-like, shape (m, ...) |
| 27 | cumulative weights of the 1D empirical distribution, if batched, must be similar to xs |
| 28 | xs: array-like, shape (n, ...) |
| 29 | locations of the 1D empirical distribution, batched against the `xs.ndim - 1` first dimensions |
| 30 | |
| 31 | Returns |
| 32 | ------- |
| 33 | q: array-like, shape (..., n) |
| 34 | The quantiles of the distribution |
| 35 | """ |
| 36 | nx = get_backend(qs, cws) |
| 37 | n = xs.shape[0] |
| 38 | if nx.__name__ == "torch": |
| 39 | # this is to ensure the best performance for torch searchsorted |
| 40 | # and avoid a warning related to non-contiguous arrays |
| 41 | cws = cws.movedim(0, -1).contiguous() |
| 42 | qs = qs.movedim(0, -1).contiguous() |
| 43 | else: |
| 44 | cws = cws.T |
| 45 | qs = qs.T |
| 46 | idx = nx.searchsorted(cws, qs).T |
| 47 | return nx.take_along_axis(xs, nx.clip(idx, 0, n - 1), axis=0) |
| 48 | |
| 49 | |
| 50 | def wasserstein_1d( |
no test coverage detected