Limit numpy thread usage for Transformation classes

.travis.yml

@@ -29,7 +29,7 @@ env:
     - MAIN_CMD="pytest ${PYTEST_LIST}"
     - SETUP_CMD="${PYTEST_FLAGS}"
     - BUILD_CMD="pip install -e package/ && (cd testsuite/ && python setup.py build)"
-    - CONDA_MIN_DEPENDENCIES="mmtf-python biopython networkx cython matplotlib scipy griddataformats hypothesis gsd codecov"
+    - CONDA_MIN_DEPENDENCIES="mmtf-python biopython networkx cython matplotlib scipy griddataformats hypothesis gsd codecov threadpoolctl"
     - CONDA_DEPENDENCIES="${CONDA_MIN_DEPENDENCIES} seaborn>=0.7.0 clustalw=2.1 netcdf4 scikit-learn joblib>=0.12 chemfiles tqdm>=4.43.0 tidynamics>=1.0.0 rdkit>=2020.03.1 h5py"
     - CONDA_CHANNELS='biobuilds conda-forge'
     - CONDA_CHANNEL_PRIORITY=True

package/MDAnalysis/lib/util.py

@@ -202,6 +202,8 @@
 import functools
 from functools import wraps
 import textwrap
+from contextlib import ContextDecorator
+from threadpoolctl import threadpool_limits
 
 import mmtf
 import numpy as np
@@ -2353,3 +2355,24 @@ def check_box(box):
     if np.all(box[3:] == 90.):
         return 'ortho', box[:3]
     return 'tri_vecs', triclinic_vectors(box)
+
+
+class threadpool_limits_decorator(threadpool_limits, ContextDecorator):
+    def __init__(self, limits=None, user_api=None):
+        self._limits, self._user_api, self._prefixes = \
+            self._check_params(limits, user_api)
+
+    def __enter__(self):
+        self._original_info = self._set_threadpool_limits()
+        self.origin_num_threads = self.get_original_num_threads()
+        return self
+
+    def __exit__(self, *exc):
+        self.unregister()
+
+    def unregister(self):
+        if self._original_info is not None:
+            for module in self._original_info:
+                module.set_num_threads(
+                    self.origin_num_threads[module.user_api]
+                )

package/MDAnalysis/transformations/__init__.py

@@ -125,6 +125,7 @@ def wrapped(ts):
     method instead of being written as a function/closure.
 """
 
+from .base import TransformationBase
 from .translate import translate, center_in_box
 from .rotate import rotateby
 from .positionaveraging import PositionAverager

package/MDAnalysis/transformations/base.py

@@ -0,0 +1,42 @@
+# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
+# 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.
+#
+# 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
+#
+
+"""\
+Transformations Base Class --- :mod:`MDAnalysis.transformations.base`
+=================================================================
+
+.. autoclass:: TransformationBase
+
+"""
+from threadpoolctl import threadpool_limits
+
+
+class TransformationBase(object):
+    def __init__(self, max_threads):
+        self.max_threads = max_threads
+
+    def __call__(self, ts):
+        with threadpool_limits(self.max_threads):
+            return self._transform(ts)
+
+    def _transform(self):
+        pass

package/MDAnalysis/transformations/fit.py

@@ -37,8 +37,10 @@
 from ..analysis import align
 from ..lib.transformations import euler_from_matrix, euler_matrix
 
+from .base import TransformationBase
 
-class fit_translation(object):
+
+class fit_translation(TransformationBase):
     """Translates a given AtomGroup so that its center of geometry/mass matches
     the respective center of the given reference. A plane can be given by the
     user using the option `plane`, and will result in the removal of
@@ -85,7 +87,9 @@ class fit_translation(object):
         The transformation was changed from a function/closure to a class
         with ``__call__``.
     """
-    def __init__(self, ag, reference, plane=None, weights=None):
+    def __init__(self, ag, reference, plane=None, weights=None, max_threads=1):
+        super().__init__(max_threads)
+
         self.ag = ag
         self.reference = reference
         self.plane = plane
@@ -117,7 +121,7 @@ def __init__(self, ag, reference, plane=None, weights=None):
         self.weights = align.get_weights(self.ref.atoms, weights=self.weights)
         self.ref_com = self.ref.center(self.weights)
 
-    def __call__(self, ts):
+    def _transform(self, ts):
         mobile_com = np.asarray(self.mobile.atoms.center(self.weights),
                                 np.float32)
         vector = self.ref_com - mobile_com
@@ -128,7 +132,7 @@ def __call__(self, ts):
         return ts
 
 
-class fit_rot_trans(object):
+class fit_rot_trans(TransformationBase):
     """Perform a spatial superposition by minimizing the RMSD.
 
     Spatially align the group of atoms `ag` to `reference` by doing a RMSD
@@ -175,7 +179,9 @@ class fit_rot_trans(object):
     -------
     MDAnalysis.coordinates.base.Timestep
     """
-    def __init__(self, ag, reference, plane=None, weights=None):
+    def __init__(self, ag, reference, plane=None, weights=None, max_threads=1):
+        super().__init__(max_threads)
+
         self.ag = ag
         self.reference = reference
         self.plane = plane
@@ -207,7 +213,7 @@ def __init__(self, ag, reference, plane=None, weights=None):
         self.ref_com = self.ref.center(self.weights)
         self.ref_coordinates = self.ref.atoms.positions - self.ref_com
 
-    def __call__(self, ts):
+    def _transform(self, ts):
         mobile_com = self.mobile.atoms.center(self.weights)
         mobile_coordinates = self.mobile.atoms.positions - mobile_com
         rotation, dump = align.rotation_matrix(mobile_coordinates,

package/MDAnalysis/transformations/positionaveraging.py

@@ -35,8 +35,10 @@
 import numpy as np
 import warnings
 
+from .base import TransformationBase
 
-class PositionAverager(object):
+
+class PositionAverager(TransformationBase):
     """
     Averages the coordinates of a given timestep so that the coordinates
     of the AtomGroup correspond to the average positions of the N previous
@@ -136,7 +138,8 @@ class PositionAverager(object):
 
     """
 
-    def __init__(self, avg_frames, check_reset=True):
+    def __init__(self, avg_frames, check_reset=True, max_threads=1):
+        super().__init__(max_threads)
         self.avg_frames = avg_frames
         self.check_reset = check_reset
         self.current_avg = 0
@@ -162,7 +165,7 @@ def rollposx(self, ts):
         self.coord_array = np.roll(self.coord_array, 1, axis=2)
         self.coord_array[..., 0] = ts.positions.copy()
 
-    def __call__(self, ts):
+    def _transform(self, ts):
         #  calling the same timestep will not add new data to coord_array
         #  This can prevent from getting different values when
         #  call `u.trajectory[i]` multiple times.

package/MDAnalysis/transformations/rotate.py

@@ -37,8 +37,10 @@
 from ..lib.transformations import rotation_matrix
 from ..lib.util import get_weights
 
+from .base import TransformationBase
 
-class rotateby(object):
+
+class rotateby(TransformationBase):
     '''
     Rotates the trajectory by a given angle on a given axis. The axis is defined by
     the user, combining the direction vector and a point. This point can be the center
@@ -118,7 +120,10 @@ def __init__(self,
                  point=None,
                  ag=None,
                  weights=None,
-                 wrap=False):
+                 wrap=False,
+                 max_threads=1):
+        super().__init__(max_threads)
+
         self.angle = angle
         self.direction = direction
         self.point = point
@@ -162,7 +167,7 @@ def __init__(self,
         else:
             raise ValueError('A point or an AtomGroup must be specified')
 
-    def __call__(self, ts):
+    def _transform(self, ts):
         if self.point is None:
             position = self.center_method()
         else:

package/MDAnalysis/transformations/translate.py

@@ -39,8 +39,10 @@
 import numpy as np
 from functools import partial
 
+from .base import TransformationBase
 
-class translate(object):
+
+class translate(TransformationBase):
     """
     Translates the coordinates of a given :class:`~MDAnalysis.coordinates.base.Timestep`
     instance by a given vector.
@@ -59,20 +61,22 @@ class translate(object):
     :class:`~MDAnalysis.coordinates.base.Timestep` object
 
     """
-    def __init__(self, vector):
+    def __init__(self, vector, max_threads=1):
+        super().__init__(max_threads)
+
         self.vector = vector
 
         if len(self.vector) > 2:
             self.vector = np.float32(self.vector)
         else:
             raise ValueError("{} vector is too short".format(self.vector))
 
-    def __call__(self, ts):
+    def _transform(self, ts):
         ts.positions += self.vector
         return ts
 
 
-class center_in_box(object):
+class center_in_box(TransformationBase):
     """
     Translates the coordinates of a given :class:`~MDAnalysis.coordinates.base.Timestep`
     instance so that the center of geometry/mass of the given :class:`~MDAnalysis.core.groups.AtomGroup`
@@ -112,7 +116,10 @@ class center_in_box(object):
         The transformation was changed from a function/closure to a class
         with ``__call__``.
     """
-    def __init__(self, ag, center='geometry', point=None, wrap=False):
+    def __init__(self, ag, center='geometry', point=None, wrap=False,
+                 max_threads=1):
+        super().__init__(max_threads)
+
         self.ag = ag
         self.center = center
         self.point = point
@@ -140,7 +147,7 @@ def __init__(self, ag, center='geometry', point=None, wrap=False):
                 raise ValueError(f'{self.ag} is not an AtomGroup object') \
                                  from None
 
-    def __call__(self, ts):
+    def _transform(self, ts):
         if self.point is None:
             boxcenter = np.sum(ts.triclinic_dimensions, axis=0) / 2
         else:

package/MDAnalysis/transformations/wrap.py

@@ -36,8 +36,10 @@
 
 from ..lib._cutil import make_whole
 
+from .base import TransformationBase
 
-class wrap(object):
+
+class wrap(TransformationBase):
     """
     Shift the contents of a given AtomGroup back into the unit cell. ::
 
@@ -85,16 +87,18 @@ class wrap(object):
         The transformation was changed from a function/closure to a class
         with ``__call__``.
     """
-    def __init__(self, ag, compound='atoms'):
+    def __init__(self, ag, compound='atoms', max_threads=1):
+        super().__init__(max_threads)
+
         self.ag = ag
         self.compound = compound
 
-    def __call__(self, ts):
+    def _transform(self, ts):
         self.ag.wrap(compound=self.compound)
         return ts
 
 
-class unwrap(object):
+class unwrap(TransformationBase):
     """
     Move all atoms in an AtomGroup so that bonds don't split over images
 
@@ -139,15 +143,17 @@ class unwrap(object):
         The transformation was changed from a function/closure to a class
         with ``__call__``.
     """
-    def __init__(self, ag):
+    def __init__(self, ag, max_threads=1):
+        super().__init__(max_threads)
+
         self.ag = ag
 
         try:
             self.ag.fragments
         except AttributeError:
             raise AttributeError("{} has no fragments".format(self.ag))
 
-    def __call__(self, ts):
+    def _transform(self, ts):
         for frag in self.ag.fragments:
             make_whole(frag)
         return ts