Source code for MDAnalysis.analysis.encore.utils
# -*- 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.
# 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
#
from __future__ import division, absolute_import
from six.moves import range
import numbers
from multiprocessing.sharedctypes import SynchronizedArray
from multiprocessing import Process, Manager
from joblib import cpu_count
import numpy as np
import sys
import MDAnalysis as mda
from ...coordinates.memory import MemoryReader
[docs]class TriangularMatrix(object):
"""Triangular matrix class. This class is designed to provide a
memory-efficient representation of a triangular matrix that still behaves
as a square symmetric one. The class wraps a numpy.array object,
in which data are memorized in row-major order. It also has few additional
facilities to conveniently load/write a matrix from/to file. It can be
accessed using the [] and () operators, similarly to a normal numpy array.
"""
def __init__(self, size, metadata=None, loadfile=None):
"""Class constructor.
Parameters
----------
size : int / array_like
Size of the matrix (number of rows or columns). If an
array is provided instead, the size of the triangular matrix
will be calculated and the array copied as the matrix
elements. Otherwise, the matrix is just initialized to zero.
metadata : dict or None
Metadata dictionary. Used to generate the metadata attribute.
loadfile : str or None
Load the matrix from this file. All the attributes and data will
be determined by the matrix file itself (i.e. metadata will be
ignored); size has to be provided though.
"""
if isinstance(metadata, dict):
self.metadata = np.array(metadata.items(), dtype=object)
else:
self.metadata = metadata
self.size = size
if loadfile:
self.loadz(loadfile)
elif isinstance(size, numbers.Integral):
self.size = size
self._elements = np.zeros((size + 1) * size // 2, dtype=np.float64)
elif isinstance(size, SynchronizedArray):
self._elements = np.array(size.get_obj(), dtype=np.float64)
self.size = int((np.sqrt(1 + 8 * len(size)) - 1) / 2)
elif isinstance(size, np.ndarray):
self._elements = size
self.size = int((np.sqrt(1 + 8 * len(size)) - 1) / 2)
else:
raise TypeError
def __getitem__(self, args):
x, y = args
if x < y:
x, y = y, x
return self._elements[x * (x + 1) // 2 + y]
def __setitem__(self, args, val):
x, y = args
if x < y:
x, y = y, x
self._elements[x * (x + 1) // 2 + y] = val
[docs] def as_array(self):
"""Return standard numpy array equivalent"""
# pylint: disable=unsubscriptable-object
a = np.zeros((self.size, self.size))
a[np.tril_indices(self.size)] = self._elements
a[np.triu_indices(self.size)] = a.T[np.triu_indices(self.size)]
return a
[docs] def savez(self, fname):
"""Save matrix in the npz compressed numpy format. Save metadata and
data as well.
Parameters
----------
fname : str
Name of the file to be saved.
"""
np.savez(fname, elements=self._elements, metadata=self.metadata)
[docs] def loadz(self, fname):
"""Load matrix from the npz compressed numpy format.
Parameters
----------
fname : str
Name of the file to be loaded.
"""
loaded = np.load(fname, allow_pickle=True)
if loaded['metadata'].shape != ():
if loaded['metadata']['number of frames'] != self.size:
raise TypeError
self.metadata = loaded['metadata']
else:
if self.size*(self.size-1)/2+self.size != len(loaded['elements']):
raise TypeError
self._elements = loaded['elements']
def __add__(self, scalar):
"""Add scalar to matrix elements.
Parameters
----------
scalar : float
Scalar to be added.
"""
newMatrix = self.__class__(self.size)
newMatrix._elements = self._elements + scalar;
return newMatrix
def __iadd__(self, scalar):
"""Add scalar to matrix elements.
Parameters
----------
scalar : float
Scalar to be added.
"""
self._elements += scalar
return self
def __mul__(self, scalar):
"""Multiply with scalar.
Parameters
----------
scalar : float
Scalar to multiply with.
"""
newMatrix = self.__class__(self.size)
newMatrix._elements = self._elements * scalar;
return newMatrix
def __imul__(self, scalar):
"""Multiply with scalar.
Parameters
----------
scalar : float
Scalar to multiply with.
"""
self._elements *= scalar
return self
__rmul__ = __mul__
def __str__(self):
return str(self.as_array())
[docs]class ParallelCalculation(object):
"""
Generic parallel calculation class. Can use arbitrary functions,
arguments to functions and kwargs to functions.
Attributes
----------
n_jobs : int
Number of cores to be used for parallel calculation. If -1 use all
available cores.
function : callable object
Function to be run in parallel.
args : list of tuples
Each tuple contains the arguments that will be passed to
function(). This means that a call to function() is performed for
each tuple. function is called as function(\*args, \*\*kwargs). Runs
are distributed on the requested numbers of cores.
kwargs : list of dicts
Each tuple contains the named arguments that will be passed to
function, similarly as described for the args attribute.
nruns : int
Number of runs to be performed. Must be equal to len(args) and
len(kwargs).
"""
def __init__(self, n_jobs, function, args=None, kwargs=None):
"""
Parameters
----------
n_jobs : int
Number of cores to be used for parallel calculation. If -1 use all
available cores.
function : object that supports __call__, as functions
function to be run in parallel.
args : list of tuples
Arguments for function; see the ParallelCalculation class
description.
kwargs : list of dicts or None
kwargs for function; see the ParallelCalculation
class description.
"""
# args[i] should be a list of args, one for each run
self.n_jobs = n_jobs
if self.n_jobs == -1:
self.n_jobs = cpu_count()
self.functions = function
if not hasattr(self.functions, '__iter__'):
self.functions = [self.functions] * len(args)
if len(self.functions) != len(args):
self.functions = self.functions[:] * (len(args) // len(self.functions))
# Arguments should be present
if args is None:
args = []
self.args = args
# If kwargs are not present, use empty dicts
if kwargs:
self.kwargs = kwargs
else:
self.kwargs = [{} for i in self.args]
self.nruns = len(args)
[docs] def worker(self, q, results):
"""
Generic worker. Will run function with the prescribed args and kwargs.
Parameters
----------
q : multiprocessing.Manager.Queue object
work queue, from which the worker fetches arguments and
messages
results : multiprocessing.Manager.Queue object
results queue, where results are put after each calculation is
finished
"""
while True:
i = q.get()
if i == 'STOP':
return
results.put((i, self.functions[i](*self.args[i], **self.kwargs[i])))
[docs] def run(self):
"""
Run parallel calculation.
Returns
-------
results : tuple of ordered tuples (int, object)
int is the number of the calculation corresponding to a
certain argument in the args list, and object is the result of
corresponding calculation. For instance, in (3, output), output
is the return of function(\*args[3], \*\*kwargs[3]).
"""
results_list = []
if self.n_jobs == 1:
for i in range(self.nruns):
results_list.append((i, self.functions[i](*self.args[i],
**self.kwargs[i])))
else:
manager = Manager()
q = manager.Queue()
results = manager.Queue()
workers = [Process(target=self.worker, args=(q, results)) for i in
range(self.n_jobs)]
for i in range(self.nruns):
q.put(i)
for w in workers:
q.put('STOP')
for w in workers:
w.start()
for w in workers:
w.join()
results.put('STOP')
for i in iter(results.get, 'STOP'):
results_list.append(i)
return tuple(sorted(results_list, key=lambda x: x[0]))
[docs]def trm_indices(a, b):
"""
Generate (i,j) indeces of a triangular matrix, between elements a and b.
The matrix size is automatically determined from the number of elements.
For instance: trm_indices((0,0),(2,1)) yields (0,0) (1,0) (1,1) (2,0)
(2,1).
Parameters
----------
a : (int i, int j) tuple
starting matrix element.
b : (int i, int j) tuple
final matrix element.
"""
i, j = a
while i < b[0]:
if i == j:
yield (i, j)
j = 0
i += 1
else:
yield (i, j)
j += 1
while j <= b[1]:
yield (i, j)
j += 1
[docs]def trm_indices_nodiag(n):
"""generate (i,j) indeces of a triangular matrix of n rows (or columns),
without diagonal (e.g. no elements (0,0),(1,1),...,(n,n))
Parameters
----------
n : int
Matrix size
"""
for i in range(1, n):
for j in range(i):
yield (i, j)
[docs]def trm_indices_diag(n):
"""generate (i,j) indeces of a triangular matrix of n rows (or columns),
with diagonal
Parameters
----------
n : int
Matrix size
"""
for i in range(0, n):
for j in range(i + 1):
yield (i, j)
[docs]def merge_universes(universes):
"""
Merge list of universes into one
Parameters
----------
universes : list of Universe objects
Returns
----------
Universe object
"""
for universe in universes:
universe.transfer_to_memory()
return mda.Universe(
universes[0].filename,
np.concatenate(tuple([e.trajectory.timeseries(order='fac') for e in universes]),
axis=0),
format=MemoryReader)