diff --git a/package/MDAnalysis/lib/trans.pyx b/package/MDAnalysis/lib/trans.pyx new file mode 100644 index 00000000000..8068c17a737 --- /dev/null +++ b/package/MDAnalysis/lib/trans.pyx @@ -0,0 +1,29 @@ +import numpy as np +cimport cython +cimport numpy as np + + +@cython.boundscheck(False) +@cython.wraparound(False) +def residue_positions( + 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 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] + 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 np.asarray(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