Extract the k largest elements from a on the given axis, and return them sorted from largest to smallest. If k is negative, extract the -k smallest elements instead, and return them sorted from smallest to largest. This performs best when ``k`` is much smaller than the chunk size. A
(a, k, axis=-1, split_every=None)
| 1325 | |
| 1326 | |
| 1327 | def topk(a, k, axis=-1, split_every=None): |
| 1328 | """Extract the k largest elements from a on the given axis, |
| 1329 | and return them sorted from largest to smallest. |
| 1330 | If k is negative, extract the -k smallest elements instead, |
| 1331 | and return them sorted from smallest to largest. |
| 1332 | |
| 1333 | This performs best when ``k`` is much smaller than the chunk size. All |
| 1334 | results will be returned in a single chunk along the given axis. |
| 1335 | |
| 1336 | Parameters |
| 1337 | ---------- |
| 1338 | x: Array |
| 1339 | Data being sorted |
| 1340 | k: int |
| 1341 | axis: int, optional |
| 1342 | split_every: int >=2, optional |
| 1343 | See :func:`reduce`. This parameter becomes very important when k is |
| 1344 | on the same order of magnitude of the chunk size or more, as it |
| 1345 | prevents getting the whole or a significant portion of the input array |
| 1346 | in memory all at once, with a negative impact on network transfer |
| 1347 | too when running on distributed. |
| 1348 | |
| 1349 | Returns |
| 1350 | ------- |
| 1351 | Selection of x with size abs(k) along the given axis. |
| 1352 | |
| 1353 | Examples |
| 1354 | -------- |
| 1355 | >>> import dask.array as da |
| 1356 | >>> x = np.array([5, 1, 3, 6]) |
| 1357 | >>> d = da.from_array(x, chunks=2) |
| 1358 | >>> d.topk(2).compute() |
| 1359 | array([6, 5]) |
| 1360 | >>> d.topk(-2).compute() |
| 1361 | array([1, 3]) |
| 1362 | """ |
| 1363 | axis = validate_axis(axis, a.ndim) |
| 1364 | |
| 1365 | # chunk and combine steps of the reduction, which recursively invoke |
| 1366 | # np.partition to pick the top/bottom k elements from the previous step. |
| 1367 | # The selection is not sorted internally. |
| 1368 | chunk_combine = partial(chunk.topk, k=k) |
| 1369 | # aggregate step of the reduction. Internally invokes the chunk/combine |
| 1370 | # function, then sorts the results internally. |
| 1371 | aggregate = partial(chunk.topk_aggregate, k=k) |
| 1372 | |
| 1373 | return reduction( |
| 1374 | a, |
| 1375 | chunk=chunk_combine, |
| 1376 | combine=chunk_combine, |
| 1377 | aggregate=aggregate, |
| 1378 | axis=axis, |
| 1379 | keepdims=True, |
| 1380 | dtype=a.dtype, |
| 1381 | split_every=split_every, |
| 1382 | output_size=abs(k), |
| 1383 | ) |
| 1384 |
no test coverage detected