From 12060928efe53d6f48454d8ab8a197488664682c Mon Sep 17 00:00:00 2001 From: Richard Gowers Date: Sun, 31 Jan 2016 20:41:38 +0000 Subject: [PATCH 1/2] First version --- package/MDAnalysis/lib/trans.pyx | 30 ++++++++++++++++++++++++++++++ package/setup.py | 5 ++++- 2 files changed, 34 insertions(+), 1 deletion(-) create mode 100644 package/MDAnalysis/lib/trans.pyx diff --git a/package/MDAnalysis/lib/trans.pyx b/package/MDAnalysis/lib/trans.pyx new file mode 100644 index 00000000000..61436e146bb --- /dev/null +++ b/package/MDAnalysis/lib/trans.pyx @@ -0,0 +1,30 @@ +import numpy as np +cimport cython +cimport numpy as np + + +@cython.boundscheck(False) +@cython.wraparound(False) +def residue_positions( + np.ndarray[np.float32_t, ndim=2] coords, + np.ndarray[np.int64_t, ndim=1] indices): + cdef Py_ssize_t i, j, k + + cdef int n_items = indices.shape[0] + cdef int n_res = np.unique(indices).shape[0] + + cdef np.ndarray[np.int64_t] counter = np.zeros(n_res, dtype=np.int64) + cdef np.ndarray[np.float32_t, ndim=2] output = np.zeros((n_res, 3), + dtype=np.float32) + + for i in range(n_items): + j = indices[i] + for k in range(3): + output[j, k] += coords[i, k] + counter[j] += 1 + + for i in range(n_res): + for k in range(3): + output[i, k] /= counter[i] + + return output diff --git a/package/setup.py b/package/setup.py index 605446d56e5..f6b05a0a58b 100755 --- a/package/setup.py +++ b/package/setup.py @@ -310,9 +310,12 @@ def extensions(config): util = MDAExtension('lib.formats.cython_util', sources=['MDAnalysis/lib/formats/cython_util.pyx'], include_dirs=include_dirs) + trans = MDAExtension('lib.trans', + sources=['MDAnalysis/lib/trans.pyx'], + include_dirs=include_dirs) extensions = [dcd, dcd_time, distances, distances_omp, qcprot, - transformation, xdrlib, util] + transformation, xdrlib, util, trans] if use_cython: extensions = cythonize(extensions) return extensions From 62fdf9927f3a9187df5670703b5a027f9c726338 Mon Sep 17 00:00:00 2001 From: Richard Gowers Date: Sun, 31 Jan 2016 21:07:46 +0000 Subject: [PATCH 2/2] Using memory views instead --- package/MDAnalysis/lib/trans.pyx | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/package/MDAnalysis/lib/trans.pyx b/package/MDAnalysis/lib/trans.pyx index 61436e146bb..8068c17a737 100644 --- a/package/MDAnalysis/lib/trans.pyx +++ b/package/MDAnalysis/lib/trans.pyx @@ -6,16 +6,15 @@ cimport numpy as np @cython.boundscheck(False) @cython.wraparound(False) def residue_positions( - np.ndarray[np.float32_t, ndim=2] coords, - np.ndarray[np.int64_t, ndim=1] indices): + float[:, :] coords not None, + long[:] indices not None): cdef Py_ssize_t i, j, k cdef int n_items = indices.shape[0] cdef int n_res = np.unique(indices).shape[0] - cdef np.ndarray[np.int64_t] counter = np.zeros(n_res, dtype=np.int64) - cdef np.ndarray[np.float32_t, ndim=2] output = np.zeros((n_res, 3), - dtype=np.float32) + cdef long[:] counter = np.zeros(n_res, dtype=np.int64) + cdef float[:, :] output = np.zeros((n_res, 3), dtype=np.float32) for i in range(n_items): j = indices[i] @@ -27,4 +26,4 @@ def residue_positions( for k in range(3): output[i, k] /= counter[i] - return output + return np.asarray(output)