WIP, ENH: demo for array API support#5426
Open
tylerjereddy wants to merge 1 commit into
Open
Conversation
* This feature branch isn't really ever intended to be merged, but is a prototype of sorts for testing out the Python Array API standard in a small part of MDAnalysis -- specifically in `MDAnalysis/analysis/bat.py`. I'm not sure that this was the best choice, but it did have a large number of `np.*` function calls, so struck me as potentially interesting to convert. This is mostly related to showing that we've at least poked around a bit for the os4ls grant, rather than proposing something that has no chance of succeeding. * On this feature branch locally on ARM Mac, the NumPy-based tests that are relevant continue to pass via i.e.: `python -m pytest MDAnalysisTests/analysis/test_bat.py` * Some initial/crude single-trial benchmark results (in seconds) from https://github.com/tylerjereddy/mda_array_api_demo don't yet show performance improvements over NumPy CPU baseline: ```python {'numpy': 0.037562333018286154, 'torch cpu': 0.05449483299162239, 'torch mps': 0.6141375829756726, 'jax': 0.7568245000147726} ``` * However, I think there are some reasons to be optimistic here in terms of what this tells us for array API compat support for us: a) JAX arrays are immutable, and adding JAX support is challenging, but it was relatively straightforward here using `array_api_extra` with the typical `xpx.at(...)` style syntax that SciPy and other use to add support for it. b) The `torch` CPU timings aren't too far off of NumPy, and keep in mind that this branch has a LOT of array type coercions because we're leveraging our in house NumPy/CPU-based analysis functions--if the task were expanded (as the grant proposes), we'd get two wins in this regard: i) the algorithms that calculated distances and so on may be performed using the array type of interest; ii) the array type coercion performance costs may dissipate. c) Using the MacOS ARM GPU (metal performance shader or "mps") still succeeds in the benchmark, which may be our first demonstration of algorithmic success on GPU in our analysis code? Furthermore, MPS only supports float32, not float64, and our testsuite still passes here on CPU even with that single precision constraint applied to NumPy.
Documentation build overview
8 files changed ·
|
orbeckst
requested changes
Jul 15, 2026
orbeckst
left a comment
orbeckst
left a comment
Member
There was a problem hiding this comment.
Fantastic work @tylerjereddy ! It's great to see that using the array API will work. 🎉
(The current CI failures are due to #5236 and not this PR, which passes its relevant tests. ✅ )
I am going to make my review blocking so that we don't accidentally merge, given that we want to implement the array standard across the package in a consistent manner.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No LLMs/AI were used in this PR/demonstration, which is mostly aimed at the
os4lsgrant prep.This feature branch isn't really ever intended to be merged, but is a prototype of sorts for testing out the Python Array API standard in a small part of MDAnalysis -- specifically in
MDAnalysis/analysis/bat.py. I'm not sure that this was the best choice, but it did have a large number ofnp.*function calls, so struck me as potentially interesting to convert. This is mostly related to showing that we've at least poked around a bit for the os4ls grant, rather than proposing something that has no chance of succeeding.On this feature branch locally on ARM Mac, the NumPy-based tests that are relevant continue to pass (as long as you have
array-api-compatandarray-api-extrainstalled, which is standard in the ecosystem now and/or vendored by projects) via i.e.:python -m pytest MDAnalysisTests/analysis/test_bat.pySome initial/crude single-trial benchmark results (in seconds) from https://github.com/tylerjereddy/mda_array_api_demo don't yet show performance improvements over NumPy CPU baseline:
{'numpy': 0.037562333018286154, 'torch cpu': 0.05449483299162239, 'torch mps': 0.6141375829756726, 'jax': 0.7568245000147726}However, I think there are some reasons to be optimistic here in terms of what this tells us for array API compat support for us:
array_api_extrawith the typicalxpx.at(...)style syntax that SciPy and others use to add support for it.torchCPU timings aren't too far off of NumPy, and keep in mind that this branch has a LOT of array type coercions because we're leveraging our in house NumPy/CPU-based analysis functions--if the task were expanded (as the grant proposes), we'd get two wins in this regard: i) the algorithms that calculated distances and so on may be performed using the array type of interest; ii) the array type coercion performance costs may dissipate.float32, notfloat64, and our testsuite still passes here on CPU even with that single precision constraint applied to NumPy.