[ENH] Use np.argpartition for efficient top-k selection instead of np.argsort

[Kaustbh]

Reference Issues/PRs

Fixes: #2713

What does this implement/fix? Explain your changes.

This PR replaces np.argsort with np.argpartition in cases where only the top-k smallest elements are needed, improving performance on large arrays.

Does your contribution introduce a new dependency? If yes, which one?

Any other comments?

PR checklist

For all contributions

• I've added myself to the list of contributors. Alternatively, you can use the @all-contributors bot to do this for you after the PR has been merged.
• The PR title starts with either [ENH], [MNT], [DOC], [BUG], [REF], [DEP] or [GOV] indicating whether the PR topic is related to enhancement, maintenance, documentation, bugs, refactoring, deprecation or governance.

For new estimators and functions

• I've added the estimator/function to the online API documentation.
• (OPTIONAL) I've added myself as a __maintainer__ at the top of relevant files and want to be contacted regarding its maintenance. Unmaintained files may be removed. This is for the full file, and you should not add yourself if you are just making minor changes or do not want to help maintain its contents.

For developers with write access

• (OPTIONAL) I've updated aeon's CODEOWNERS to receive notifications about future changes to these files.

[aeon-actions-bot]

Thank you for contributing to aeon

I did not find any labels to add based on the title. Please add the [ENH], [MNT], [BUG], [DOC], [REF], [DEP] and/or [GOV] tags to your pull requests titles. For now you can add the labels manually.
I have added the following labels to this PR based on the changes made: [ $\color{#006b75}{\textsf{similarity search}}$ ]. Feel free to change these if they do not properly represent the PR.

The Checks tab will show the status of our automated tests. You can click on individual test runs in the tab or "Details" in the panel below to see more information if there is a failure.

If our pre-commit code quality check fails, any trivial fixes will automatically be pushed to your PR unless it is a draft.

Don't hesitate to ask questions on the aeon Slack channel if you have any.

PR CI actions

These checkboxes will add labels to enable/disable CI functionality for this PR. This may not take effect immediately, and a new commit may be required to run the new configuration.

• Run pre-commit checks for all files
• Run mypy typecheck tests
• Run all pytest tests and configurations
• Run all notebook example tests
• Run numba-disabled codecov tests
• Stop automatic pre-commit fixes (always disabled for drafts)
• Disable numba cache loading
• Push an empty commit to re-run CI checks

[baraline]

Test are failing due to an issue with the implementation, currently, sorted_index is only keeping the first k indexes of dist_profile with the smallest distance. When not allowing trivial matches, it is possible that the top-k matches are not included in the first k indexes.

The original issue mention "adjust the algorithms, notably for the successive call to argpartition in cases where trivial matches should be avoided."

The idea is to try to find these k best match in the first argpartition, and look into another k argpartition, until you find k matches or checked all values of dist profile.

Am I clear or is this a bit confusion ?

[Kaustbh]

Sorry for the earlier misunderstanding. Our main goal is to find k valid matches. If we find fewer than k, we continue searching in the next batch of k elements using argpartition, and repeat this until we either find all k matches or reach the end of dist_profile. Am I right?

[baraline]

No problem, this is quite a challenging topic. Yes I think you got it. We want to output k matches, but if for any reason (length of the distance profile or usage of trivial matches causing a shortage of admissible matches), we cannot find k, we will return less.

[baraline]

Thanks for the changes, I'll run a benchmark to compare performance before making the decision, but it looks OK at first glance.

[baraline]

Bit busy until 15th of june, i'll take a look at it after, thx 🙏

[Kaustbh]

Okay, no problem

[baraline]

Finally had a change to benchmark it, I confirm that this reduces runtime for large inputs, and the difference is negligible for small ones. Great job, thanks again !