Improve performance of wrap method

[hejamu]

Changes made in this Pull Request:

• Changed the wrap method to use index mapping instead of a loop.

The method was significantly slower than expected when wrapping compounds (residues, fragments, molecules, segments). See the corresponding issue on the MAICoS repo.

It just never got the same performance increases as for example unwrap. I considered using the _split_by_compound_indices approach from unwrap, but that is actually not needed for wrapping. Simply calculating the shifts and applying them in one go yields an enormous speed up. Most the performance penalty actually came from the np.where done in each loop I believe.

Here is a scaling plot comparing the speeds:

LLM / AI generated code disclosure

LLMs or other AI-powered tools (beyond simple IDE use cases) were used in this contribution: no

PR Checklist

• Issue raised/referenced?
• Tests updated/added?
• Documentation updated/added?
• package/CHANGELOG file updated?
• Is your name in package/AUTHORS? (If it is not, add it!)
• LLM/AI disclosure was updated.

Developers Certificate of Origin

I certify that I can submit this code contribution as described in the Developer Certificate of Origin, under the MDAnalysis LICENSE.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 93.82%. Comparing base (e6c1b44) to head (ddf80d9).
⚠️ Report is 3 commits behind head on develop.

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #5226      +/-   ##
===========================================
- Coverage    93.83%   93.82%   -0.01%     
===========================================
  Files          180      180              
  Lines        22475    22473       -2     
  Branches      3190     3189       -1     
===========================================
- Hits         21090    21086       -4     
- Misses         923      924       +1     
- Partials       462      463       +1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[orbeckst]

This looks good but I didn't have time to dive into it. Can you confirm that it produces the same results? I know that the tests covered it but did you look if the tests for it actually likely to find any disagreement?

[hejamu]

The tests are already quite comprehensive. I only noticed that the tests do not check whether it also works with non-contiguous compound indexes. Since the current code and the PR use fancy-indexing, I added such a test.

[PardhavMaradani]

Linking the performance gains here to the ones in an older open PR #2982 (relevant changes to groups.py) to note that this PR would address this aspect of that old PR for future reference. Thanks

[hejamu]

@orbeckst, could you take another look if we can merge this?

[orbeckst]

The fix looks good, nice work.

Can you move your benchmark files and data into a gist https://gist.github.com/ and just link it from the PR? We don't want the files in the repo but it would be good to store them properly in case we need them in the future.

EDIT: A super-nice bonus (optional!) would be to create an actual ASV benchmark test in https://github.com/MDAnalysis/mdanalysis/tree/develop/benchmarks/benchmarks that would then show up in https://www.mdanalysis.org/benchmarks/ and demonstrate the speed up publicly.

[hejamu]

Gist: https://gist.github.com/hejamu/e38f3cabe0b0fe9a5fb1d0c0fe55e3cc

I added the benchmark to asv, here the results from my machine.

| Change   | Before [6b8d4174] <develop>   | After [9b3eda02] <faster_wrap>   |   Ratio | Benchmark (Parameter)                                      |
|----------|-------------------------------|----------------------------------|---------|------------------------------------------------------------|
| -        | 69.8±0.3μs                    | 56.5±0.4μs                       |    0.81 | ag_methods.AtomGroupMethodsBench.time_wrap_compound(100)   |
| -        | 393±1μs                       | 186±1μs                          |    0.47 | ag_methods.AtomGroupMethodsBench.time_wrap_compound(1000)  |
| -        | 10.1±0.09ms                   | 907±4μs                          |    0.09 | ag_methods.AtomGroupMethodsBench.time_wrap_compound(10000) |

[hejamu]

Test failures related to: #5236

[orbeckst]

Fantastic, thank you!