Add distopia! (#3914)

* add starrt of calc_bonds

* refactor to use distopia inplace API

* add in upcast

* actually return the distances

* distopia to ci

* add initial docs

* move to using _run backend

* add upcast back in

* fix yml CI files

* change test_pairwise_dist to use allclose rather than equals

* change another test to sue assert allclose

* improve docs

* change distance tests to use distopia backend

* split numpy compliance test

* changelog

* fix merge issue in changelog

* add note on varying precision to docs

* Update testsuite/MDAnalysisTests/lib/test_distances.py

Co-authored-by: Irfan Alibay <IAlibay@users.noreply.github.com>

* add warning for triclinic boxes

* darker

* add second importlib try block

* return callable or ndarray

* Revert "return callable or ndarray"

This reverts commit 56e5091.

* change docs on default

* start of stub

* Update package/MDAnalysis/lib/distances.py

Co-authored-by: Irfan Alibay <IAlibay@users.noreply.github.com>

* add align

* WIP on stub

* finalise distopia stub

* change calc_bonds to use distopia stub

* fix type annotation

* darker

* darker distances.py

* try to make black happier again

* fix distances.py again

* finally fix everything

* fix versionchanged

* remove try except

* Update package/MDAnalysis/lib/distances.py

Co-authored-by: Irfan Alibay <IAlibay@users.noreply.github.com>

* finalise merge

---------

Co-authored-by: Hugo Macdermott <hugomacdermott@Hugos-MacBook-Air-2.local>
Co-authored-by: Irfan Alibay <IAlibay@users.noreply.github.com>

.github/actions/setup-deps/action.yaml

@@ -53,6 +53,8 @@ inputs:
     default: 'chemfiles-python>=0.9'
   clustalw:
     default: 'clustalw=2.1'
+  distopia:
+    default: 'distopia>=0.2.0'
   h5py:
     default: 'h5py>=2.10'
   joblib:
@@ -115,6 +117,7 @@ runs:
         CONDA_OPT_DEPS: |
           ${{ inputs.chemfiles-python }}
           ${{ inputs.clustalw }}
+          ${{ inputs.distopia }}
           ${{ inputs.h5py }}
           ${{ inputs.joblib }}
           ${{ inputs.netcdf4 }}

package/CHANGELOG

@@ -26,6 +26,8 @@ Fixes
     (Issue #3336)
 
 Enhancements
+  * Add distopia distance calculation library bindings as a selectable backend
+    for `calc_bonds` in `MDA.lib.distances`. (Issue #3783, PR #3914)
   * AuxReaders are now pickle-able and copy-able (Issue #1785, PR #3887)
   * Add pickling support for Atom, Residue, Segment, ResidueGroup 
     and SegmentGroup. (PR #3953)

package/MDAnalysis/lib/_distopia.py

@@ -0,0 +1,71 @@
+# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; -*-
+# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
+#
+# MDAnalysis --- https://www.mdanalysis.org
+# Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors
+# (see the file AUTHORS for the full list of names)
+#
+# Released under the GNU Public Licence, v2 or any higher version
+#
+# Please cite your use of MDAnalysis in published work:
+#
+# R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler,
+# D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein.
+# MDAnalysis: A Python package for the rapid analysis of molecular dynamics
+# simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th
+# Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy.
+# doi: 10.25080/majora-629e541a-00e
+#
+# N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein.
+# MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
+# J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
+#
+
+"""Stub module for distopia --- :mod:`MDAnalysis.analysis.distopia`
+===================================================================
+
+This module is a stub to provide distopia distance functions to `distances.py`
+as a selectable backend.
+"""
+
+# check for distopia
+try:
+    import distopia
+except ImportError:
+    HAS_DISTOPIA = False
+else:
+    HAS_DISTOPIA = True
+
+from .c_distances import (
+    calc_bond_distance_triclinic as _calc_bond_distance_triclinic_serial,
+)
+import warnings
+import numpy as np
+
+
+def calc_bond_distance_ortho(
+    coords1, coords2: np.ndarray, box: np.ndarray, results: np.ndarray
+) -> None:
+    distopia.calc_bonds_ortho_float(
+        coords1, coords2, box[:3], results=results
+    )
+    # upcast is currently required, change for 3.0, see #3927
+
+
+def calc_bond_distance(
+    coords1: np.ndarray, coords2: np.ndarray, results: np.ndarray
+) -> None:
+    distopia.calc_bonds_no_box_float(
+        coords1, coords2, results=results
+    )
+    # upcast is currently required, change for 3.0, see #3927
+
+
+def calc_bond_distance_triclinic(
+    coords1: np.ndarray, coords2: np.ndarray, box: np.ndarray, results: np.ndarray
+) -> None:
+    # redirect to serial backend
+    warnings.warn(
+        "distopia does not support triclinic boxes, using serial backend instead."
+    )
+    _calc_bond_distance_triclinic_serial(coords1, coords2, box, results)

package/MDAnalysis/lib/distances.py

@@ -50,10 +50,57 @@
                                         with OpenMP
    ========== ========================= ======================================
 
+Use of the distopia library
+---------------------------
+
+MDAnalysis has developed a standalone library, `distopia`_ for accelerating
+the distance functions in this module using explicit SIMD vectorisation.
+This can provide many-fold speedups in calculating distances. Distopia is
+under active development and as such only a selection of functions in this
+module are covered. Consult the following table to see if the function
+you wish to use is covered by distopia. For more information see the
+`distopia documentation`_.
+
+.. Table:: Functions available using the `distopia`_ backend.
+    :align: center
+
+    +-------------------------------------+-----------------------------------+
+    | Functions                           | Notes                             |
+    +=====================================+===================================+
+    | MDAnalysis.lib.distances.calc_bonds | Doesn't support triclinic boxes   |
+    +-------------------------------------+-----------------------------------+
+
+If `distopia`_ is installed, the functions in this table will accept the key
+'distopia' for the `backend` keyword argument. If the distopia backend is
+selected the `distopia` library will be used to calculate the distances. Note
+that for functions listed in this table **distopia is not the default backend
+if and must be selected.**
+
+.. Note::
+   Distopia does not currently support triclinic simulation boxes. If you
+   specify `distopia` as the backend and your simulation box is triclinic,
+   the function will fall back to the default `serial` backend.
+
+.. Note::
+    Due to the use of Instruction Set Architecture (`ISA`_) specific SIMD
+    intrinsics in distopia via `VCL2`_, the precision of your results may
+    depend on the ISA available on your machine. However, in all tested cases
+    distopia satisfied the accuracy thresholds used to the functions in this
+    module. Please document any issues you encounter with distopia's accuracy
+    in the `relevant distopia issue`_ on the MDAnalysis GitHub repository.
+
+.. _distopia: https://github.com/MDAnalysis/distopia
+.. _distopia documentation: https://www.mdanalysis.org/distopia
+.. _ISA: https://en.wikipedia.org/wiki/Instruction_set_architecture
+.. _VCL2: https://github.com/vectorclass/version2
+.. _relevant distopia issue: https://github.com/MDAnalysis/mdanalysis/issues/3915
+
 .. versionadded:: 0.13.0
 .. versionchanged:: 2.3.0
    Distance functions can now accept an
    :class:`~MDAnalysis.core.groups.AtomGroup` or an :class:`np.ndarray`
+.. versionchanged:: 2.5.0
+   Interface to the `distopia`_ package added.
 
 Functions
 ---------
@@ -83,6 +130,7 @@
 from ._augment import augment_coordinates, undo_augment
 from .nsgrid import FastNS
 from .c_distances import _minimize_vectors_ortho, _minimize_vectors_triclinic
+from ._distopia import HAS_DISTOPIA
 
 
 # hack to select backend with backend=<backend> kwarg. Note that
@@ -97,8 +145,11 @@
                                           package="MDAnalysis.lib")
 except ImportError:
     pass
-del importlib
 
+if HAS_DISTOPIA:
+    _distances["distopia"] = importlib.import_module("._distopia",
+                             package="MDAnalysis.lib")
+del importlib
 
 def _run(funcname: str, args: Optional[tuple] = None,
          kwargs: Optional[dict] = None, backend: str = "serial") -> Callable:
@@ -1408,15 +1459,15 @@ def calc_bonds(coords1: Union[npt.NDArray, 'AtomGroup'],
         Preallocated result array of dtype ``numpy.float64`` and shape ``(n,)``
         (for ``n`` coordinate pairs). Avoids recreating the array in repeated
         function calls.
-    backend : {'serial', 'OpenMP'}, optional
-        Keyword selecting the type of acceleration.
+    backend : {'serial', 'OpenMP', 'distopia'}, optional
+        Keyword selecting the type of acceleration. Defaults to 'serial'.
 
     Returns
     -------
-    bondlengths : numpy.ndarray (``dtype=numpy.float64``, ``shape=(n,)``) or numpy.float64
-        Array containing the bond lengths between each pair of coordinates. If
-        two single coordinates were supplied, their distance is returned as a
-        single number instead of an array.
+    bondlengths : numpy.ndarray (``dtype=numpy.float64``, ``shape=(n,)``) or
+        numpy.float64 Array containing the bond lengths between each pair of
+        coordinates. If two single coordinates were supplied, their distance is
+        returned as a single number instead of an array.
 
 
     .. versionadded:: 0.8
@@ -1428,26 +1479,39 @@ def calc_bonds(coords1: Union[npt.NDArray, 'AtomGroup'],
     .. versionchanged:: 2.3.0
        Can now accept an :class:`~MDAnalysis.core.groups.AtomGroup` as an
        argument in any position and checks inputs using type hinting.
+    .. versionchanged:: 2.5.0
+       Can now optionally use the fast distance functions from distopia
     """
     numatom = coords1.shape[0]
     bondlengths = _check_result_array(result, (numatom,))
 
     if numatom > 0:
         if box is not None:
             boxtype, box = check_box(box)
-            if boxtype == 'ortho':
-                _run("calc_bond_distance_ortho",
-                     args=(coords1, coords2, box, bondlengths),
-                     backend=backend)
+            if boxtype == "ortho":
+                if backend == 'distopia':
+                    bondlengths = bondlengths.astype(np.float32)
+                _run(
+                    "calc_bond_distance_ortho",
+                    args=(coords1, coords2, box, bondlengths),
+                    backend=backend,
+                )
             else:
-                _run("calc_bond_distance_triclinic",
-                     args=(coords1, coords2, box, bondlengths),
-                     backend=backend)
+                _run(
+                    "calc_bond_distance_triclinic",
+                    args=(coords1, coords2, box, bondlengths),
+                    backend=backend,
+                )
         else:
-            _run("calc_bond_distance",
-                 args=(coords1, coords2, bondlengths),
-                 backend=backend)
-
+            if backend == 'distopia':
+                bondlengths = bondlengths.astype(np.float32)
+            _run(
+                "calc_bond_distance",
+                args=(coords1, coords2, bondlengths),
+                backend=backend,
+            )
+    if backend == 'distopia':
+        bondlengths = bondlengths.astype(np.float64)
     return bondlengths
 
 

testsuite/MDAnalysisTests/analysis/test_distances.py

@@ -148,7 +148,7 @@ def test_pairwise_dist(self, ag, ag2, expected):
         '''Ensure that pairwise distances between atoms are
         correctly calculated.'''
         actual = MDAnalysis.analysis.distances.dist(ag, ag2)[2]
-        assert_equal(actual, expected)
+        assert_allclose(actual, expected)
 
     def test_pairwise_dist_box(self, ag, ag2, expected_box, box):
         '''Ensure that pairwise distances between atoms are
@@ -161,7 +161,7 @@ def test_pairwise_dist_offset_effect(self, ag, ag2, expected):
         pairwise distance matrix.'''
         actual = MDAnalysis.analysis.distances.dist(
             ag, ag2, offset=229)[2]
-        assert_equal(actual, expected)
+        assert_allclose(actual, expected)
 
     def test_offset_calculation(self, ag, ag2):
         '''Test that offsets fed to dist() are correctly calculated.'''