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