From 01ed7529785e689021033defa227f35508485298 Mon Sep 17 00:00:00 2001 From: Tyler Reddy Date: Sun, 12 Jul 2026 15:28:50 -0600 Subject: [PATCH 1/3] WIP, ENH: demo for array API support * 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. --- package/MDAnalysis/analysis/bat.py | 75 +++++++++++++++++------------- 1 file changed, 42 insertions(+), 33 deletions(-) diff --git a/package/MDAnalysis/analysis/bat.py b/package/MDAnalysis/analysis/bat.py index 9b08c7dbff4..d9a0f21c290 100644 --- a/package/MDAnalysis/analysis/bat.py +++ b/package/MDAnalysis/analysis/bat.py @@ -161,6 +161,8 @@ class to calculate dihedral angles for a given set of atoms or residues import warnings import numpy as np +from array_api_compat import array_namespace +import array_api_extra as xpx import copy import MDAnalysis as mda @@ -285,7 +287,7 @@ def get_supported_backends(cls): description="Bond-Angle-Torsions Coordinate Transformation", path="MDAnalysis.analysis.bat.BAT", ) - def __init__(self, ag, initial_atom=None, filename=None, **kwargs): + def __init__(self, ag, initial_atom=None, filename=None, backend_arr=None, **kwargs): r"""Parameters ---------- ag : AtomGroup or Universe @@ -315,6 +317,10 @@ def __init__(self, ag, initial_atom=None, filename=None, **kwargs): """ super(BAT, self).__init__(ag.universe.trajectory, **kwargs) self._ag = ag + if backend_arr is not None: + self._backend_arr = backend_arr + else: + self._backend_arr = np.array([1.]) # Check that the ag contains bonds if not hasattr(self._ag, "bonds"): @@ -394,8 +400,9 @@ def __init__(self, ag, initial_atom=None, filename=None, **kwargs): self.load(filename) def _prepare(self): - self.results.bat = np.zeros( - (self.n_frames, 3 * self._ag.n_atoms), dtype=np.float64 + xp = array_namespace(self._backend_arr) + self.results.bat = xp.zeros( + (self.n_frames, 3 * self._ag.n_atoms), dtype=xp.float32 ) def _single_frame(self): @@ -403,75 +410,77 @@ def _single_frame(self): # The rotation axis is a normalized vector pointing from atom 0 to 1 # It is described in two degrees of freedom # by the polar angle and azimuth + xp = array_namespace(self._backend_arr) if self._root.dimensions is None: (p0, p1, p2) = self._root.positions else: (p0, p1, p2) = make_whole(self._root, inplace=False) + p0, p1, p2 = map(xp.asarray, (p0, p1, p2)) v01 = p1 - p0 v21 = p1 - p2 # Internal coordinates - r01 = np.sqrt(np.einsum("i,i->", v01, v01)) + r01 = xp.sqrt(xp.einsum("i,i->", v01, v01)) # Distance between first two root atoms - r12 = np.sqrt(np.einsum("i,i->", v21, v21)) + r12 = xp.sqrt(xp.einsum("i,i->", v21, v21)) # Distance between second two root atoms # Angle between root atoms - a012 = np.arccos( + a012 = xp.arccos( max( -1.0, min( 1.0, - np.einsum("i,i->", v01, v21) - / np.sqrt( - np.einsum("i,i->", v01, v01) - * np.einsum("i,i->", v21, v21) + xp.einsum("i,i->", v01, v21) + / xp.sqrt( + xp.einsum("i,i->", v01, v01) + * xp.einsum("i,i->", v21, v21) ), ), ) ) # External coordinates e = v01 / r01 - phi = np.arctan2(e[1], e[0]) # Polar angle - theta = np.arccos(e[2]) # Azimuthal angle + phi = xp.arctan2(e[1], e[0]) # Polar angle + theta = xp.arccos(e[2]) # Azimuthal angle # Rotation to the z axis - cp = np.cos(phi) - sp = np.sin(phi) - ct = np.cos(theta) - st = np.sin(theta) - Rz = np.array( + cp = xp.cos(phi) + sp = xp.sin(phi) + ct = xp.cos(theta) + st = xp.sin(theta) + Rz = xp.asarray( [[cp * ct, ct * sp, -st], [-sp, cp, 0], [cp * st, sp * st, ct]] ) - pos2 = Rz.dot(p2 - p1) + pos2 = xp.matmul(Rz, (p2 - p1)) # Angle about the rotation axis - omega = np.arctan2(pos2[1], pos2[0]) - root_based = np.concatenate((p0, [phi, theta, omega, r01, r12, a012])) + omega = xp.arctan2(pos2[1], pos2[0]) + root_based = xp.concat((p0, xp.asarray([phi, theta, omega, r01, r12, a012]))) # Calculate internal coordinates from the torsion list - bonds = calc_bonds( - self._ag1.positions, self._ag2.positions, box=self._ag1.dimensions - ) - angles = calc_angles( + bonds = xp.asarray(calc_bonds( + self._ag1.positions, self._ag2.positions, box=self._ag1.dimensions), + dtype=xp.float32) + angles = xp.asarray(calc_angles( self._ag1.positions, self._ag2.positions, self._ag3.positions, box=self._ag1.dimensions, - ) - torsions = calc_dihedrals( + ), dtype=xp.float32) + torsions = xp.asarray(calc_dihedrals( self._ag1.positions, self._ag2.positions, self._ag3.positions, self._ag4.positions, box=self._ag1.dimensions, - ) + ), dtype=xp.float32) # When appropriate, calculate improper torsions - shift = torsions[self._primary_torsion_indices] - shift[self._unique_primary_torsion_indices] = 0.0 + shift = torsions[xp.asarray(self._primary_torsion_indices)] + shift = xpx.at(shift)[xp.asarray(self._unique_primary_torsion_indices)].set(0.0) torsions -= shift - # Wrap torsions to between -np.pi and np.pi - torsions = ((torsions + np.pi) % (2 * np.pi)) - np.pi + # Wrap torsions to between -xp.pi and xp.pi + torsions = ((torsions + xp.pi) % (2 * xp.pi)) - xp.pi - self.results.bat[self._frame_index, :] = np.concatenate( + self.results.bat = xpx.at(self.results.bat)[self._frame_index, :].set(xp.concat( (root_based, bonds, angles, torsions) - ) + )) def load(self, filename, start=None, stop=None, step=None): """Loads the bat trajectory from a file in numpy binary format From 685ab518e2c102bbe153ad9d35b68c10625f566b Mon Sep 17 00:00:00 2001 From: Oliver Beckstein Date: Sun, 19 Jul 2026 22:31:39 -0700 Subject: [PATCH 2/3] add array-api-* packages to CI - add array-api-compat - add array-api-extra For GitHub actions and Azure pipelines. Also added to ReadTheDocs environment. --- .github/actions/setup-deps/action.yaml | 6 ++++++ azure-pipelines.yml | 2 ++ maintainer/conda/environment.yml | 2 ++ package/pyproject.toml | 4 +++- package/requirements.txt | 2 ++ 5 files changed, 15 insertions(+), 1 deletion(-) diff --git a/.github/actions/setup-deps/action.yaml b/.github/actions/setup-deps/action.yaml index 7f40ec146ad..b9c2eb21349 100644 --- a/.github/actions/setup-deps/action.yaml +++ b/.github/actions/setup-deps/action.yaml @@ -17,6 +17,10 @@ inputs: description: 'use micromamba instead of conda' default: false # conda-installed min dependencies + array-api-compat: + default: 'array-api-compat' + array-api-extra: + default: 'array-api-extra' codecov: default: 'codecov' cython: @@ -118,6 +122,8 @@ runs: shell: bash -l {0} env: CONDA_MIN_DEPS: | + ${{ inputs.array-api-compat }} + ${{ inputs.array-api-extra }} ${{ inputs.codecov }} ${{ inputs.cython }} ${{ inputs.filelock }} diff --git a/azure-pipelines.yml b/azure-pipelines.yml index c5e80b9c579..b4342f207ad 100644 --- a/azure-pipelines.yml +++ b/azure-pipelines.yml @@ -82,6 +82,8 @@ jobs: h5py>=2.10 matplotlib numpy + array-api-compat + array-api-extra packaging pytest pytest-cov diff --git a/maintainer/conda/environment.yml b/maintainer/conda/environment.yml index e20ad6de080..c6bb2a0cb5c 100644 --- a/maintainer/conda/environment.yml +++ b/maintainer/conda/environment.yml @@ -2,6 +2,8 @@ name: mda-dev channels: - conda-forge dependencies: + - array-api-compat + - array-api-extra - chemfiles>=0.10 - codecov - cython diff --git a/package/pyproject.toml b/package/pyproject.toml index 283980b1332..2002ae66392 100644 --- a/package/pyproject.toml +++ b/package/pyproject.toml @@ -30,6 +30,8 @@ maintainers = [ requires-python = ">=3.11" dependencies = [ 'numpy>=1.26.0', + 'array-api-compat', + 'array-api-extra', 'GridDataFormats>=0.4.0', 'mmtf-python>=1.0.0', 'joblib>=0.12', @@ -130,7 +132,7 @@ line-length = 79 target-version = ['py311', 'py312', 'py313'] extend-exclude = ''' ( -__pycache__ +__pycache__ | MDAnalysis/coordinates/DCD\.py | MDAnalysis/coordinates/TRJ\.py | MDAnalysis/coordinates/__init__\.py diff --git a/package/requirements.txt b/package/requirements.txt index 17c3961a5ef..da64e0f1732 100644 --- a/package/requirements.txt +++ b/package/requirements.txt @@ -1,3 +1,5 @@ +array-api-compat +array-api-extra biopython>=1.80 codecov cython From 3dc97d4897f00c73f7dad0764a70640f85a47d9d Mon Sep 17 00:00:00 2001 From: Oliver Beckstein Date: Mon, 20 Jul 2026 00:22:07 -0700 Subject: [PATCH 3/3] reformat bat.py with black --- package/MDAnalysis/analysis/bat.py | 128 ++++++++++++++++------------- 1 file changed, 73 insertions(+), 55 deletions(-) diff --git a/package/MDAnalysis/analysis/bat.py b/package/MDAnalysis/analysis/bat.py index d9a0f21c290..99a087ac469 100644 --- a/package/MDAnalysis/analysis/bat.py +++ b/package/MDAnalysis/analysis/bat.py @@ -51,10 +51,10 @@ Each molecule also has six external coordinates that define its translation and rotation in space. The three Cartesian coordinates of the first atom are the molecule's translational degrees of freedom. Rotational degrees of freedom are -specified by the axis-angle convention. The rotation axis is a normalized vector -pointing from the first to second atom. It is described by the polar angle, -:math:`\phi`, and azimuthal angle, :math:`\theta`. :math:`\omega` is a third angle -that describes the rotation of the third atom about the axis. +specified by the axis-angle convention. The rotation axis is a normalized +vector pointing from the first to second atom. It is described by the polar +angle, :math:`\phi`, and azimuthal angle, :math:`\theta`. :math:`\omega` is a +third angle that describes the rotation of the third atom about the axis. This module was adapted from AlGDock :footcite:p:`Minh2020`. @@ -74,8 +74,8 @@ class to calculate dihedral angles for a given set of atoms or residues coordinates based on the topology of an atom group and interconverts between Cartesian and BAT coordinate systems. -For example, we can determine internal coordinates for residues 5-10 -of adenylate kinase (AdK). The trajectory is included within the test data files:: +For example, we can determine internal coordinates for residues 5-10 of +adenylate kinase (AdK). The trajectory is included within the test data files:: import MDAnalysis as mda from MDAnalysisTests.datafiles import PSF, DCD @@ -135,20 +135,20 @@ class to calculate dihedral angles for a given set of atoms or residues .. attribute:: results.bat - Contains the time series of the Bond-Angle-Torsion coordinates as a - (nframes, 3N) :class:`numpy.ndarray` array. Each row corresponds to - a frame in the trajectory. In each column, the first six elements - describe external degrees of freedom. The first three are the center - of mass of the initial atom. The next three specify the external angles - according to the axis-angle convention: :math:`\phi`, the polar angle, - :math:`\theta`, the azimuthal angle, and :math:`\omega`, a third angle - that describes the rotation of the third atom about the axis. The next - three degrees of freedom are internal degrees of freedom for the root - atoms: :math:`r_{01}`, the distance between atoms 0 and 1, - :math:`r_{12}`, the distance between atoms 1 and 2, - and :math:`a_{012}`, the angle between the three atoms. - The rest of the array consists of all the other bond distances, - all the other bond angles, and then all the other torsion angles. + Contains the time series of the Bond-Angle-Torsion coordinates as a + (nframes, 3N) :class:`numpy.ndarray` array. Each row corresponds to a + frame in the trajectory. In each column, the first six elements describe + external degrees of freedom. The first three are the center of mass of + the initial atom. The next three specify the external angles according + to the axis-angle convention: :math:`\phi`, the polar angle, + :math:`\theta`, the azimuthal angle, and :math:`\omega`, a third angle + that describes the rotation of the third atom about the axis. The next + three degrees of freedom are internal degrees of freedom for the root + atoms: :math:`r_{01}`, the distance between atoms 0 and 1, + :math:`r_{12}`, the distance between atoms 1 and 2, and :math:`a_{012}`, + the angle between the three atoms. The rest of the array consists of + all the other bond distances, all the other bond angles, and then all + the other torsion angles. References @@ -179,7 +179,8 @@ class to calculate dihedral angles for a given set of atoms or residues def _sort_atoms_by_mass(atoms, reverse=False): r"""Sorts a list of atoms by mass and then by index - The atom index is used as a tiebreaker so that the ordering is reproducible. + The atom index is used as a tiebreaker so that the ordering is + reproducible. Parameters ---------- @@ -192,6 +193,7 @@ def _sort_atoms_by_mass(atoms, reverse=False): ------- ag_n : list of Atoms Sorted list + """ return sorted(atoms, key=lambda a: (a.mass, a.index), reverse=reverse) @@ -287,17 +289,18 @@ def get_supported_backends(cls): description="Bond-Angle-Torsions Coordinate Transformation", path="MDAnalysis.analysis.bat.BAT", ) - def __init__(self, ag, initial_atom=None, filename=None, backend_arr=None, **kwargs): + def __init__( + self, ag, initial_atom=None, filename=None, backend_arr=None, **kwargs + ): r"""Parameters ---------- ag : AtomGroup or Universe - Group of atoms for which the BAT coordinates are calculated. - `ag` must have a bonds attribute. - If unavailable, bonds may be guessed using - :meth:`AtomGroup.guess_bonds `. - `ag` must only include one molecule. - If a trajectory is associated with the atoms, then the computation - iterates over the trajectory. + Group of atoms for which the BAT coordinates are calculated. `ag` + must have a bonds attribute. If unavailable, bonds may be guessed + using :meth:`AtomGroup.guess_bonds + `. `ag` must only + include one molecule. If a trajectory is associated with the + atoms, then the computation iterates over the trajectory. initial_atom : :class:`Atom ` The atom whose Cartesian coordinates define the translation of the molecule. If not specified, the heaviest terminal atom @@ -320,7 +323,7 @@ def __init__(self, ag, initial_atom=None, filename=None, backend_arr=None, **kwa if backend_arr is not None: self._backend_arr = backend_arr else: - self._backend_arr = np.array([1.]) + self._backend_arr = np.array([1.0]) # Check that the ag contains bonds if not hasattr(self._ag, "bonds"): @@ -452,35 +455,50 @@ def _single_frame(self): pos2 = xp.matmul(Rz, (p2 - p1)) # Angle about the rotation axis omega = xp.arctan2(pos2[1], pos2[0]) - root_based = xp.concat((p0, xp.asarray([phi, theta, omega, r01, r12, a012]))) + root_based = xp.concat( + (p0, xp.asarray([phi, theta, omega, r01, r12, a012])) + ) # Calculate internal coordinates from the torsion list - bonds = xp.asarray(calc_bonds( - self._ag1.positions, self._ag2.positions, box=self._ag1.dimensions), - dtype=xp.float32) - angles = xp.asarray(calc_angles( - self._ag1.positions, - self._ag2.positions, - self._ag3.positions, - box=self._ag1.dimensions, - ), dtype=xp.float32) - torsions = xp.asarray(calc_dihedrals( - self._ag1.positions, - self._ag2.positions, - self._ag3.positions, - self._ag4.positions, - box=self._ag1.dimensions, - ), dtype=xp.float32) + bonds = xp.asarray( + calc_bonds( + self._ag1.positions, + self._ag2.positions, + box=self._ag1.dimensions, + ), + dtype=xp.float32, + ) + angles = xp.asarray( + calc_angles( + self._ag1.positions, + self._ag2.positions, + self._ag3.positions, + box=self._ag1.dimensions, + ), + dtype=xp.float32, + ) + torsions = xp.asarray( + calc_dihedrals( + self._ag1.positions, + self._ag2.positions, + self._ag3.positions, + self._ag4.positions, + box=self._ag1.dimensions, + ), + dtype=xp.float32, + ) # When appropriate, calculate improper torsions shift = torsions[xp.asarray(self._primary_torsion_indices)] - shift = xpx.at(shift)[xp.asarray(self._unique_primary_torsion_indices)].set(0.0) + shift = xpx.at(shift)[ + xp.asarray(self._unique_primary_torsion_indices) + ].set(0.0) torsions -= shift # Wrap torsions to between -xp.pi and xp.pi torsions = ((torsions + xp.pi) % (2 * xp.pi)) - xp.pi - self.results.bat = xpx.at(self.results.bat)[self._frame_index, :].set(xp.concat( - (root_based, bonds, angles, torsions) - )) + self.results.bat = xpx.at(self.results.bat)[self._frame_index, :].set( + xp.concat((root_based, bonds, angles, torsions)) + ) def load(self, filename, start=None, stop=None, step=None): """Loads the bat trajectory from a file in numpy binary format @@ -548,10 +566,10 @@ def Cartesian(self, bat_frame): Returns ------- XYZ : numpy.ndarray - an array with dimensions (N,3) with Cartesian coordinates. The first - dimension has the same ordering as the AtomGroup used to initialize - the class. The molecule will be whole opposed to wrapped around a - periodic boundary. + an array with dimensions (N,3) with Cartesian coordinates. The + first dimension has the same ordering as the AtomGroup used to + initialize the class. The molecule will be whole opposed to wrapped + around a periodic boundary. """ # Split the bat vector into more convenient variables origin = bat_frame[:3]