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/MDAnalysis/analysis/bat.py b/package/MDAnalysis/analysis/bat.py index 9b08c7dbff4..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 @@ -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 @@ -177,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 ---------- @@ -190,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) @@ -285,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, **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 @@ -315,6 +320,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.0]) # Check that the ag contains bonds if not hasattr(self._ag, "bonds"): @@ -394,8 +403,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,74 +413,91 @@ 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 + bonds = xp.asarray( + calc_bonds( + self._ag1.positions, + self._ag2.positions, + box=self._ag1.dimensions, + ), + dtype=xp.float32, ) - angles = calc_angles( - self._ag1.positions, - self._ag2.positions, - self._ag3.positions, - box=self._ag1.dimensions, + angles = xp.asarray( + calc_angles( + self._ag1.positions, + self._ag2.positions, + self._ag3.positions, + box=self._ag1.dimensions, + ), + dtype=xp.float32, ) - torsions = calc_dihedrals( - self._ag1.positions, - self._ag2.positions, - self._ag3.positions, - self._ag4.positions, - box=self._ag1.dimensions, + 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( - (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): @@ -539,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] 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