WIP, ENH: demo for array API support#5426
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
11 files changed ·
|
orbeckst
left a comment
orbeckst
left a comment
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.
orbeckst
left a comment
orbeckst
left a comment
There was a problem hiding this comment.
Let's install array-api-compat and array-api-extra in CI for this PR.
ERROR testsuite/MDAnalysisTests/analysis - ModuleNotFoundError: No module named 'array_api_compat'
(With all the other non-related failures I overlooked that these modules were missing and also made the tests fail.)
…dy/mdanalysis into treddy_array_api_demo
- add array-api-compat - add array-api-extra For GitHub actions and Azure pipelines. Also added to ReadTheDocs environment.
|
I based my local branch on the latest develop (to get some of the CI fixes) so the diff is a bit messy at the moment – sorry! @tylerjereddy if you want to do any rebasing please go ahead but I don't want to change your branch's history. |
Get in fixes that make CI pass again.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## develop #5426 +/- ##
===========================================
- Coverage 93.85% 93.84% -0.01%
===========================================
Files 182 182
Lines 22510 22518 +8
Branches 3202 3203 +1
===========================================
+ Hits 21126 21132 +6
- Misses 922 923 +1
- Partials 462 463 +1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
2 participants
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.