Extract the indices of the k largest elements from a on the given axis, and return them sorted from largest to smallest. If k is negative, extract the indices of the -k smallest elements instead, and return them sorted from smallest to largest. This performs best when ``k`` is much
(a, k, axis=-1, split_every=None)
| 1402 | |
| 1403 | |
| 1404 | def argtopk(a, k, axis=-1, split_every=None): |
| 1405 | """Extract the indices of the k largest elements from a on the given axis, |
| 1406 | and return them sorted from largest to smallest. If k is negative, extract |
| 1407 | the indices of the -k smallest elements instead, and return them sorted |
| 1408 | from smallest to largest. |
| 1409 | |
| 1410 | This performs best when ``k`` is much smaller than the chunk size. All |
| 1411 | results will be returned in a single chunk along the given axis. |
| 1412 | |
| 1413 | Parameters |
| 1414 | ---------- |
| 1415 | x: Array |
| 1416 | Data being sorted |
| 1417 | k: int |
| 1418 | axis: int, optional |
| 1419 | split_every: int >=2, optional |
| 1420 | See :func:`topk`. The performance considerations for topk also apply |
| 1421 | here. |
| 1422 | |
| 1423 | Returns |
| 1424 | ------- |
| 1425 | Selection of np.intp indices of x with size abs(k) along the given axis. |
| 1426 | |
| 1427 | Examples |
| 1428 | -------- |
| 1429 | >>> import dask.array as da |
| 1430 | >>> x = np.array([5, 1, 3, 6]) |
| 1431 | >>> d = da.from_array(x, chunks=2) |
| 1432 | >>> d.argtopk(2).compute() |
| 1433 | array([3, 0]) |
| 1434 | >>> d.argtopk(-2).compute() |
| 1435 | array([1, 2]) |
| 1436 | """ |
| 1437 | axis = validate_axis(axis, a.ndim) |
| 1438 | |
| 1439 | # Generate nodes where every chunk is a tuple of (a, original index of a) |
| 1440 | idx = arange(a.shape[axis], chunks=(a.chunks[axis],), dtype=np.intp) |
| 1441 | idx = idx[tuple(slice(None) if i == axis else np.newaxis for i in range(a.ndim))] |
| 1442 | a_plus_idx = a.map_blocks(chunk.argtopk_preprocess, idx, dtype=object) |
| 1443 | |
| 1444 | # chunk and combine steps of the reduction. They acquire in input a tuple |
| 1445 | # of (a, original indices of a) and return another tuple containing the top |
| 1446 | # k elements of a and the matching original indices. The selection is not |
| 1447 | # sorted internally, as in np.argpartition. |
| 1448 | chunk_combine = partial(chunk.argtopk, k=k) |
| 1449 | # aggregate step of the reduction. Internally invokes the chunk/combine |
| 1450 | # function, then sorts the results internally, drops a and returns the |
| 1451 | # index only. |
| 1452 | aggregate = partial(chunk.argtopk_aggregate, k=k) |
| 1453 | |
| 1454 | if isinstance(axis, Number): |
| 1455 | naxis = 1 |
| 1456 | else: |
| 1457 | naxis = len(axis) |
| 1458 | |
| 1459 | meta = a._meta.astype(np.intp).reshape((0,) * (a.ndim - naxis + 1)) |
| 1460 | |
| 1461 | return reduction( |
no test coverage detected