added Cython PBC

.travis.yml

@@ -27,7 +27,7 @@ env:
     - MAIN_CMD="pytest ${PYTEST_LIST}"
     - SETUP_CMD="${PYTEST_FLAGS}"
     - BUILD_CMD="pip install -v package/ && pip install testsuite/"
-    - CONDA_MIN_DEPENDENCIES="mmtf-python mock six biopython networkx cython joblib matplotlib scipy griddataformats hypothesis gsd codecov"
+    - CONDA_MIN_DEPENDENCIES="mmtf-python mock six biopython networkx cython joblib matplotlib scipy griddataformats hypothesis gsd codecov cymem"
     - CONDA_DEPENDENCIES="${CONDA_MIN_DEPENDENCIES} seaborn>=0.7.0 clustalw=2.1 netcdf4 scikit-learn"
     - CONDA_CHANNELS='biobuilds conda-forge'
     - CONDA_CHANNEL_PRIORITY=True

package/MDAnalysis/lib/_cutil.pxd

@@ -1,3 +1,4 @@
 cdef float _dot(float *, float *)
 cdef void _cross(float *, float *, float *)
-cdef float _norm(float *)
+cdef float _norm(float *)
+cdef float norm2(float *) nogil

package/MDAnalysis/lib/_cutil.pyx

@@ -305,9 +305,9 @@ cdef float _norm(float * a):
     """
     Calculates the magnitude of the vector
     """
-    cdef float result
-    cdef ssize_t n
-    result = 0.0
-    for n in range(3):
-        result += a[n]*a[n]
-    return sqrt(result)
+    return sqrt(norm2(a))
+
+
+cdef float norm2(float* a) nogil:
+    """Square of norm"""
+    return a[0] * a[0] + a[1] * a[1] + a[2] * a[2]

package/MDAnalysis/lib/new_distances.pyx

@@ -0,0 +1,45 @@
+import cython
+import numpy as np
+
+from libc.math import sqrt
+from ._cutil cimport norm2
+from pbc cimport PBC, minimum_image
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+cdef void inner_distance_array(const float[:, :] a, const float[:, :] b,
+                               PBC box, float[:, :] result) nogil:
+    """C level of distance_array"""
+    cdef int i, j
+    cdef float[3] dx
+
+    for i in range(a.shape[0]):
+        for j in range(b.shape[0]):
+            dx[0] = a[i, 0] - b[j, 0]
+            dx[1] = a[i, 1] - b[j, 1]
+            dx[2] = a[i, 2] - b[j, 2]
+
+            minimum_image(dx, box)
+
+            result[i, j] = sqrt(norm2(dx))
+
+
+def distance_array(coords1, coords2, box=None, result=None, backend=None):
+    cdef PBC cbox
+    cdef float[:, :] a, b, result_view
+
+    a = np.asarray(coords1, dtype=np.float32)
+    b = np.asarray(coords2, dtype=np.float32)
+    if not box is None:
+        box = np.asarray(box, dtype=np.float32)
+    if result is None:
+        result = np.zeros((a.shape[0], b.shape[0]), dtype=np.float32)
+
+    cbox = PBC(box)
+    result_view = result
+
+    with nogil:
+        inner_distance_array(a, b, cbox, result_view)
+
+    return result

package/MDAnalysis/lib/pbc.pxd

@@ -0,0 +1,28 @@
+from cymem.cymem cimport Pool
+
+cdef class PBC:
+    # Cymem Memory pool, takes care of heap memory allocation
+    cdef Pool mem
+    # Triclinic vector representation of box
+    cdef float** box
+    # Full box diagonal
+    cdef float* fbox_diag
+    # Half box diagonal
+    cdef float* hbox_diag
+    # Minus half box diagonal
+    cdef float* mhbox_diag
+    # Maximum cutoff squared
+    cdef float max_cutoff2
+    # What type of periodic boundaries [ORTHO, TRICLINIC, NONE]
+    cdef public int pbc_type
+    # triclinic vectors to search for minimum image
+    cdef int ntric_vec
+    cdef float **tric_vec
+
+    cdef void define_box(self, float[:] box)
+    cdef void ortho_define(self, float[:] box)
+    cdef void triclinic_define(self, float[:] box)
+    cdef void find_triclinic_shifts(self)
+    cdef float calc_max_cutoff2(self)
+
+cdef void minimum_image(float* dx, PBC pbc) nogil

package/MDAnalysis/lib/pbc.pyx

@@ -0,0 +1,239 @@
+import cython
+
+from cymem.cymem cimport Pool
+from libc.math cimport fabs, sin, cos, sqrt
+from numpy.math cimport deg2radf #, sinf as sin, cosf as cos, sqrtf as sqrt
+
+from _cutil cimport norm2
+
+DEF MAX_NTRICVEC=12
+DEF EPSILON=1e-5
+DEF BOX_MARGIN=1.0010
+
+cdef enum PBC_TYPES:
+    ORTHO = 1
+    TRICLINIC = 2
+    NOPBC = 3
+
+
+cdef class PBC:
+    """Cython class representing periodic boundaries in system
+
+    Used for minimum_image function
+    """
+    def __cinit__(self, float[:] box):
+        self.mem = None
+        self.box = NULL
+        self.fbox_diag = NULL
+        self.hbox_diag = NULL
+        self.mhbox_diag = NULL
+        self.tric_vec = NULL
+
+    def __init__(self, box):
+        self.mem = Pool()
+        self.box = <float**>self.mem.alloc(3, sizeof(float*))
+        for i in range(3):
+            self.box[i] = <float*>self.mem.alloc(3, sizeof(float))
+        self.fbox_diag = <float*>self.mem.alloc(3, sizeof(float))
+        self.hbox_diag = <float*>self.mem.alloc(3, sizeof(float))
+        self.mhbox_diag = <float*>self.mem.alloc(3, sizeof(float))
+        self.tric_vec = <float**>self.mem.alloc(MAX_NTRICVEC, sizeof(float*))
+        for i in range(MAX_NTRICVEC):
+            self.tric_vec[i] = <float*>self.mem.alloc(3, sizeof(float))
+        self.ntric_vec = 0
+
+        if box is None:
+            self.pbc_type = NOPBC
+        else:
+            if not box.shape == (6,):
+                raise ValueError('Wrong shaped box')
+            self.define_box(box)
+
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    cdef void define_box(self, float[:] box):
+        # turn box into triclinic vectors
+        cdef bint triclinic
+
+        triclinic = False
+        for i in range(3):
+            if fabs(box[i + 3] - 90.0) > EPSILON:
+                triclinic = True
+
+        if triclinic:
+            self.pbc_type = TRICLINIC
+            self.triclinic_define(box)
+        else:
+            self.pbc_type = ORTHO
+            self.ortho_define(box)
+
+        self.fbox_diag[0] = box[0]
+        self.fbox_diag[1] = box[1]
+        self.fbox_diag[2] = box[2]
+
+        self.hbox_diag[0] = 0.5 * self.fbox_diag[0]
+        self.hbox_diag[1] = 0.5 * self.fbox_diag[1]
+        self.hbox_diag[2] = 0.5 * self.fbox_diag[2]
+
+        self.mhbox_diag[0] = - self.hbox_diag[0]
+        self.mhbox_diag[1] = - self.hbox_diag[1]
+        self.mhbox_diag[2] = - self.hbox_diag[2]
+
+        self.max_cutoff2 = self.calc_max_cutoff2()
+
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    cdef void ortho_define(self, float[:] box):
+        self.box[0][0] = box[0]
+        self.box[0][1] = 0.0
+        self.box[0][2] = 0.0
+
+        self.box[1][0] = 0.0
+        self.box[1][1] = box[1]
+        self.box[1][2] = 0.0
+
+        self.box[2][0] = 0.0
+        self.box[2][1] = 0.0
+        self.box[2][2] = box[2]
+
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    @cython.cdivision(True)
+    cdef void triclinic_define(self, float[:] box):
+        cdef float alpha, beta, gamma
+
+        alpha = deg2radf(box[3])
+        beta = deg2radf(box[4])
+        gamma = deg2radf(box[5])
+
+        self.box[0][0] = box[0]
+        self.box[0][1] = 0.0
+        self.box[0][2] = 0.0
+
+        self.box[1][0] = box[1] * cos(gamma)
+        self.box[1][1] = box[1] * sin(gamma)
+        self.box[1][2] = 0.0
+
+        self.box[2][0] = box[2] * cos(beta)
+        self.box[2][1] = box[2] * (cos(alpha) - cos(beta) * cos(gamma) / sin(gamma))
+        self.box[2][2] = box[2] * sqrt(box[2] * box[2] - self.box[2][0] ** 2 - self.box[2][1] ** 2)
+
+    cdef void find_triclinic_shifts(self):
+        cdef int[3] order
+        cdef int i, j, k, dd, d, shift
+        cdef float[3] trial, pos
+        cdef float d2old, d2new, d2new_c
+        cdef bint use
+
+        order[0] = 0
+        order[1] = -1
+        order[2] = 1
+
+        for k in order:
+            for j in order:
+                for i in order:
+                    if j != 0 or k != 0:
+                        d2old = 0.0
+                        d2new = 0.0
+
+                        for d in range(3):
+                            trial[d] = i * self.box[0][d] + j * self.box[1][d] + k * self.box[2][d]
+                            if d == 3:
+                                trial[d] = 0.0
+                                pos[d] = 0.0
+                            else:
+                                if trial[d] < 0:
+                                    pos[d] = min(self.hbox_diag[d], -trial[d])
+                                else:
+                                    pos[d] = max(self.hbox_diag[d], -trial[d])
+
+                            d2old += pos[d] * pos[d]
+                            d2new += (pos[d] + trial[d]) * (pos[d] + trial[d])
+
+                        if BOX_MARGIN * d2new < d2old:
+                            use = True
+
+                            for dd in range(3):  # for d in [i, j, k]:
+                                if dd == 0:
+                                    shift = i
+                                elif dd == 1:
+                                    shift = j
+                                else:
+                                    shift = k
+
+                                if shift:
+                                    d2new_c = 0
+                                    for d in range(3):
+                                        d2new_c += (pos[d] + trial[d] - shift * self.box[dd][d])**2
+
+                                    if d2new_c <= BOX_MARGIN * d2new:
+                                        use = False
+
+                            if use: # Accept this shift vector.
+                                for d in range(3):
+                                    self.tric_vec[self.ntric_vec][d] = trial[d]
+                                    self.ntric_vec += 1
+
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    cdef float calc_max_cutoff2(self):
+        cdef float min_hv2, min_ss
+
+        min_hv2 = 0.25 * min(norm2(self.box[0]),
+                             norm2(self.box[1]),
+                             norm2(self.box[2]))
+
+        min_ss = min(self.box[0][0], fabs(self.box[1][1] - self.box[2][1]), self.box[2][2])
+
+        return min(min_hv2, min_ss * min_ss)
+
+
+cdef void minimum_image(float* dx, PBC pbc) nogil:
+    """Apply minimum image convention to a vector *dx*
+
+    Parameters
+    ----------
+    float* dx
+      a single 3d separation
+    PBC pbc
+      PBC object representing the box
+
+    Modifies dx in place
+    """
+    cdef int i, j
+    cdef float d2min, d2trial
+    cdef float[3] dx_start, trial
+
+    if pbc.pbc_type == TRICLINIC:
+        for i in range(2, -1, -1):
+            while (dx[i] > pbc.hbox_diag[i]):
+                for j in range(i, -1, -1):
+                    dx[j] -= pbc.box[i][j]
+            while (dx[i] <= pbc.mhbox_diag[i]):
+                for j in range(i, -1, -1):
+                    dx[j] += pbc.box[i][j]
+        d2min = norm2(dx)
+        if d2min > pbc.max_cutoff2:
+            for j in range(3):
+                dx_start[j] = dx[j]
+
+            i = 0
+            while (d2min > pbc.max_cutoff2) and (i < pbc.ntric_vec):
+                for j in range(3):
+                    trial[j] = dx_start[j] + pbc.tric_vec[i][j]
+                d2trial = norm2(trial)
+                if d2trial < d2min:
+                    for j in range(3):
+                        dx[j] = trial[j]
+                    d2min = d2trial
+                i += 1
+
+    elif pbc.pbc_type == ORTHO:
+        for i in range(3):
+            while (dx[i] > pbc.hbox_diag[i]):
+                dx[i] -= pbc.fbox_diag[i]
+            while (dx[i] <= pbc.mhbox_diag[i]):
+                dx[i] += pbc.fbox_diag[i]
+
+    elif pbc.pbc_type == NOPBC:
+        pass

package/setup.py

@@ -305,6 +305,18 @@ def extensions(config):
                           include_dirs=include_dirs + ['MDAnalysis/lib/formats/include'],
                           define_macros=define_macros,
                           extra_compile_args=extra_compile_args)
+    pbc = MDAExtension('MDAnalysis.lib.pbc',
+                       ['MDAnalysis/lib/pbc' + source_suffix],
+                       include_dirs=include_dirs,
+                       libraries=mathlib,
+                       define_macros=define_macros,
+                       extra_compile_args=extra_compile_args)
+    new_distances = MDAExtension('MDAnalysis.lib.new_distances',
+                                 ['MDAnalysis/lib/new_distances' + source_suffix],
+                                 include_dirs=include_dirs,
+                                 libraries=mathlib,
+                                 define_macros=define_macros,
+                                 extra_compile_args=extra_compile_args)
     distances = MDAExtension('MDAnalysis.lib.c_distances',
                              ['MDAnalysis/lib/c_distances' + source_suffix],
                              include_dirs=include_dirs + ['MDAnalysis/lib/include'],
@@ -380,7 +392,7 @@ def extensions(config):
                               libraries=mathlib,
                               define_macros=define_macros,
                               extra_compile_args=extra_compile_args)
-    pre_exts = [libdcd, distances, distances_omp, qcprot,
+    pre_exts = [libdcd, pbc, new_distances, distances, distances_omp, qcprot,
                 transformation, libmdaxdr, util, encore_utils,
                 ap_clustering, spe_dimred, cutil, augment]