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