From c04301abd37f2268820010260cddf5e5edd40b61 Mon Sep 17 00:00:00 2001 From: ayush Date: Mon, 25 Jun 2018 20:10:22 -0700 Subject: [PATCH 01/47] First commit for gridded method --- package/MDAnalysis/lib/c_gridsearch.pyx | 747 ++++++++++++++++++++++++ package/setup.py | 9 +- 2 files changed, 755 insertions(+), 1 deletion(-) create mode 100644 package/MDAnalysis/lib/c_gridsearch.pyx diff --git a/package/MDAnalysis/lib/c_gridsearch.pyx b/package/MDAnalysis/lib/c_gridsearch.pyx new file mode 100644 index 00000000000..82d7192544a --- /dev/null +++ b/package/MDAnalysis/lib/c_gridsearch.pyx @@ -0,0 +1,747 @@ +# -*- coding: utf-8; Mode: python; tab-width: 4; indent-tabs-mode:nil; -*- +# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 +# +# Copyright (C) 2013-2016 Sébastien Buchoux +# +# This file is part of FATSLiM. +# +# FATSLiM is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# FATSLiM is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with FATSLiM. If not, see . +#cython: cdivision=True +#cython: boundscheck=False + +# Preprocessor DEFs +DEF DIM = 3 +DEF XX = 0 +DEF YY = 1 +DEF ZZ = 2 +DEF RET_OK = 1 +DEF RET_ERROR = 0 +DEF EPSILON = 1e-5 +DEF NEIGHBORHOOD_ALLOCATION_INCREMENT = 50 +DEF GRID_ALLOCATION_INCREMENT = 50 + +DEF BOX_MARGIN=1.0010 +DEF MAX_NTRICVEC=12 + + +# Cython C imports (no Python here!) +from cython.parallel cimport prange +from libc.stdlib cimport malloc, realloc, free, abort +from libc.stdio cimport fprintf, stderr +from libc.math cimport sqrt +from libc.math cimport abs as real_abs + +cimport openmp + + +# Python imports +import numpy as np +cimport numpy as np + +# Ctypes +ctypedef np.int_t ns_int +ctypedef np.float32_t real +ctypedef real rvec[DIM] +ctypedef real matrix[DIM][DIM] + +cdef struct ns_grid: + ns_int size + ns_int[DIM] ncells + real[DIM] cellsize + ns_int *nbeads + ns_int **beadids + +cdef struct ns_neighborhood: + real cutoff + ns_int allocated_size + ns_int size + ns_int *beadids + real *beaddist + ### +cdef struct ns_neighborhood_holder: + ns_int size + ns_neighborhood **neighborhoods + +# Useful stuff + +cdef real rvec_norm2(const rvec a) nogil: + return a[XX]*a[XX]+a[YY]*a[YY]+a[ZZ]*a[ZZ] + +cdef void rvec_clear(rvec a) nogil: + a[XX]=0.0 + a[YY]=0.0 + a[ZZ]=0.0 + +cdef struct cPBCBox_t: + matrix box + rvec fbox_diag + rvec hbox_diag + rvec mhbox_diag + real max_cutoff2 + ns_int ntric_vec + ns_int[DIM] tric_shift[MAX_NTRICVEC] + real[DIM] tric_vec[MAX_NTRICVEC] + +# noinspection PyNoneFunctionAssignment +cdef class PBCBox(object): + cdef cPBCBox_t c_pbcbox + cdef rvec center + cdef rvec bbox_center + + def __init__(self, real[:,::1] box): + self.update(box) + + cdef void fast_update(self, real[:,::1] box) nogil: + cdef ns_int i, j, k, d, jc, kc, shift + cdef real d2old, d2new, d2new_c + cdef rvec trial, pos + cdef ns_int ii, jj ,kk + cdef ns_int *order = [0, -1, 1, -2, 2] + cdef bint use + cdef real min_hv2, min_ss, tmp + + rvec_clear(self.center) + # Update matrix + for i in range(DIM): + for j in range(DIM): + self.c_pbcbox.box[i][j] = box[i, j] + self.center[j] += 0.5 * box[i, j] + self.bbox_center[i] = 0.5 * box[i, i] + + # Update diagonals + for i in range(DIM): + self.c_pbcbox.fbox_diag[i] = box[i, i] + self.c_pbcbox.hbox_diag[i] = self.c_pbcbox.fbox_diag[i] * 0.5 + self.c_pbcbox.mhbox_diag[i] = - self.c_pbcbox.hbox_diag[i] + + # Update maximum cutoff + + # Physical limitation of the cut-off + # by half the length of the shortest box vector. + min_hv2 = min(0.25 * rvec_norm2(&box[XX, XX]), 0.25 * rvec_norm2(&box[YY, XX])) + min_hv2 = min(min_hv2, 0.25 * rvec_norm2(&box[ZZ, XX])) + + # Limitation to the smallest diagonal element due to optimizations: + # checking only linear combinations of single box-vectors (2 in x) + # in the grid search and pbc_dx is a lot faster + # than checking all possible combinations. + tmp = box[YY, YY] + if box[ZZ, YY] < 0: + tmp -= box[ZZ, YY] + else: + tmp += box[ZZ, YY] + + min_ss = min(box[XX, XX], min(tmp, box[ZZ, ZZ])) + + self.c_pbcbox.max_cutoff2 = min(min_hv2, min_ss * min_ss) + + # Update shift vectors + self.c_pbcbox.ntric_vec = 0 + # We will only use single shifts, but we will check a few + # more shifts to see if there is a limiting distance + # above which we can not be sure of the correct distance. + for kk in range(5): + k = order[kk] + + for jj in range(5): + j = order[jj] + + for ii in range(5): + i = order[ii] + + # A shift is only useful when it is trilinic + if j != 0 or k != 0: + d2old = 0 + d2new = 0 + + for d in range(DIM): + trial[d] = i*box[XX, d] + j*box[YY, d] + k*box[ZZ, d] + + # Choose the vector within the brick around 0,0,0 that + # will become the shortest due to shift try. + + if d == DIM: + trial[d] = 0 + pos[d] = 0 + else: + if trial[d] < 0: + pos[d] = min(self.c_pbcbox.hbox_diag[d], -trial[d]) + else: + pos[d] = max(-self.c_pbcbox.hbox_diag[d], -trial[d]) + + d2old += sqrt(pos[d]) + d2new += sqrt(pos[d] + trial[d]) + + if BOX_MARGIN*d2new < d2old: + if not (j < -1 or j > 1 or k < -1 or k > 1): + use = True + + for dd in range(DIM): + if dd == 0: + shift = i + elif dd == 1: + shift = j + else: + shift = k + + if shift: + d2new_c = 0 + + for d in range(DIM): + d2new_c += sqrt(pos[d] + trial[d] - shift*box[dd, d]) + + if d2new_c <= BOX_MARGIN*d2new: + use = False + + if use: # Accept this shift vector. + if self.c_pbcbox.ntric_vec >= MAX_NTRICVEC: + with gil: + print("\nWARNING: Found more than %d triclinic " + "correction vectors, ignoring some." + % MAX_NTRICVEC) + print(" There is probably something wrong with " + "your box.") + print(box) + else: + for d in range(DIM): + self.c_pbcbox.tric_vec[self.c_pbcbox.ntric_vec][d] = \ + trial[d] + self.c_pbcbox.tric_shift[self.c_pbcbox.ntric_vec][XX] = i + self.c_pbcbox.tric_shift[self.c_pbcbox.ntric_vec][YY] = j + self.c_pbcbox.tric_shift[self.c_pbcbox.ntric_vec][ZZ] = k + self.c_pbcbox.ntric_vec += 1 + + + def update(self, real[:,::1] box): + if box.shape[0] != DIM or box.shape[1] != DIM: + raise ValueError("Box must be a %i x %i matrix. (shape: %i x %i)" % + (DIM, DIM, box.shape[0], box.shape[1])) + if (box[XX, XX] == 0) or (box[YY, YY] == 0) or (box[ZZ, ZZ] == 0): + raise ValueError("Box does not correspond to PBC=xyz") + self.fast_update(box) + + cdef void fast_pbc_dx(self, rvec ref, rvec other, rvec dx) nogil: + cdef ns_int i, j + cdef rvec dx_start, trial + + for i in range(DIM): + dx[i] = other[i] - ref[i] + + for i in range (DIM-1, -1, -1): + while dx[i] > self.c_pbcbox.hbox_diag[i]: + for j in range (i, -1, -1): + dx[j] -= self.c_pbcbox.box[i][j] + + while dx[i] <= self.c_pbcbox.mhbox_diag[i]: + for j in range (i, -1, -1): + dx[j] += self.c_pbcbox.box[i][j] + + cdef real[:, ::1]fast_put_atoms_in_bbox(self, real[:,::1] coords) nogil: + cdef ns_int i, m, d, natoms, wd = 0 + cdef real[:,::1] bbox_coords + + natoms = coords.shape[0] + with gil: + if natoms == 0: + bbox_coords = np.empty((0, DIM)) + else: + bbox_coords = coords.copy() + + for i in range(natoms): + for m in range(DIM - 1, -1, -1): + while bbox_coords[i, m] < 0: + for d in range(m+1): + bbox_coords[i, d] += self.c_pbcbox.box[m][d] + while bbox_coords[i, m] >= self.c_pbcbox.box[m][m]: + for d in range(m+1): + bbox_coords[i, d] -= self.c_pbcbox.box[m][d] + return bbox_coords + + def put_atoms_in_bbox(self, real[:,::1] coords): + return np.asarray(self.fast_put_atoms_in_bbox(coords)) + + + +######################################################################################################################## +# +# Neighbor Search Stuff +# +######################################################################################################################## +cdef struct ns_grid: + ns_int size + ns_int[DIM] ncells + real[DIM] cellsize + ns_int *nbeads + ns_int **beadids + +cdef ns_grid initialize_nsgrid(matrix box, + float cutoff) nogil: + cdef ns_grid grid + cdef ns_int i + + for i in range(DIM): + grid.ncells[i] = (box[i][i] / cutoff) + if grid.ncells[i] == 0: + grid.ncells[i] = 1 + grid.cellsize[i] = box[i][i] / grid.ncells[i] + + grid.size = grid.ncells[XX] * grid.ncells[YY] * grid.ncells[ZZ] + return grid + +cdef ns_int populate_grid(ns_grid *grid, + real[:,::1] coords) nogil: + cdef ns_int ncoords = coords.shape[0] + cdef bint ret_val + + ret_val = populate_grid_array(grid, + &coords[0, 0], + ncoords) + + return ret_val + +cdef ns_int populate_grid_array(ns_grid *grid, + rvec *coords, + ns_int ncoords) nogil: + cdef ns_int i, cellindex = -1 + cdef ns_int grid_size = grid.size + cdef ns_int *allocated_size = NULL + + if grid_size != grid.ncells[XX] * grid.ncells[YY] * grid.ncells[ZZ]: # Grid not initialized + return RET_ERROR + + + # Allocate memory + grid.nbeads = malloc(sizeof(ns_int) * grid_size) + if grid.nbeads == NULL: + fprintf(stderr,"FATAL: Could not allocate memory for NS grid.nbeads (requested: %i bytes)\n", + sizeof(ns_int) * grid_size) + abort() + + allocated_size = malloc(sizeof(ns_int) * grid_size) + if allocated_size == NULL: + fprintf(stderr,"FATAL: Could not allocate memory for NS allocated_size (requested: %i bytes)\n", + sizeof(ns_int) * grid_size) + abort() + + for i in range(grid_size): + grid.nbeads[i] = 0 + allocated_size[i] = GRID_ALLOCATION_INCREMENT + + grid.beadids = malloc(sizeof(ns_int *) * grid_size) + if grid.beadids == NULL: + fprintf(stderr,"FATAL: Could not allocate memory for NS grid.beadids (requested: %i bytes)\n", + sizeof(ns_int *) * grid_size) + abort() + + for i in range(grid_size): + grid.beadids[i] = malloc(sizeof(ns_int) * allocated_size[i]) + if grid.beadids[i] == NULL: + fprintf(stderr,"FATAL: Could not allocate memory for NS grid.beadids[i] (requested: %i bytes)\n", + sizeof(ns_int) * allocated_size[i]) + abort() + + # Get cell indices for coords + for i in range(ncoords): + cellindex = (coords[i][ZZ] / grid.cellsize[ZZ]) * (grid.ncells[XX] * grid.ncells[YY]) +\ + (coords[i][YY] / grid.cellsize[YY]) * grid.ncells[XX] + \ + (coords[i][XX] / grid.cellsize[XX]) + + grid.beadids[cellindex][grid.nbeads[cellindex]] = i + grid.nbeads[cellindex] += 1 + + if grid.nbeads[cellindex] >= allocated_size[cellindex]: + allocated_size[cellindex] += GRID_ALLOCATION_INCREMENT + grid.beadids[cellindex] = realloc( grid.beadids[cellindex], sizeof(ns_int) * allocated_size[cellindex]) + free(allocated_size) + return RET_OK + +cdef void destroy_nsgrid(ns_grid *grid) nogil: + cdef ns_int i + if grid.nbeads != NULL: + free(grid.nbeads) + + for i in range(grid.size): + if grid.beadids[i] != NULL: + free(grid.beadids[i]) + free(grid.beadids) + + +cdef ns_neighborhood_holder *create_neighborhood_holder() nogil: + cdef ns_neighborhood_holder *holder + + holder = malloc(sizeof(ns_neighborhood_holder)) + + return holder + +cdef void free_neighborhood_holder(ns_neighborhood_holder *holder) nogil: + cdef ns_int i + + if holder == NULL: + return + + for i in range(holder.size): + if holder.neighborhoods[i].beadids != NULL: + free(holder.neighborhoods[i].beadids) + free(holder.neighborhoods[i]) + free(holder.neighborhoods) + free(holder) + +cdef ns_neighborhood *retrieve_neighborhood(rvec current_coords, real[:, ::1]neighborcoords, ns_grid *grid, PBCBox box, real cutoff2) nogil: + cdef ns_int d, m + cdef ns_int xi, yi, zi, bid + cdef real d2 + cdef rvec shifted_coord, dx, neighbor_coord, corrected_coords + + cdef bint already_checked[27] + cdef bint skip + cdef ns_int nchecked = 0, icheck + cdef ns_int cell_index + + cdef ns_neighborhood *neighborhood = malloc(sizeof(ns_neighborhood)) + if neighborhood == NULL: + abort() + + neighborhood.size = 0 + neighborhood.allocated_size = NEIGHBORHOOD_ALLOCATION_INCREMENT + neighborhood.beadids = malloc(NEIGHBORHOOD_ALLOCATION_INCREMENT * sizeof(ns_int)) + ###Modified here + neighborhood.beaddist = malloc(NEIGHBORHOOD_ALLOCATION_INCREMENT * sizeof(real)) + ### + if neighborhood.beadids == NULL: + abort() + + for zi in range(3): + for yi in range(3): + for xi in range(3): + # Calculate and/or reinitialize shifted coordinates + shifted_coord[XX] = current_coords[XX] + (xi - 1) * grid.cellsize[XX] + shifted_coord[YY] = current_coords[YY] + (yi - 1) * grid.cellsize[YY] + shifted_coord[ZZ] = current_coords[ZZ] + (zi - 1) * grid.cellsize[ZZ] + + # Make sure the shifted coordinates is inside the brick-shaped box + for m in range(DIM - 1, -1, -1): + + while shifted_coord[m] < 0: + for d in range(m+1): + shifted_coord[d] += box.c_pbcbox.box[m][d] + + + while shifted_coord[m] >= box.c_pbcbox.box[m][m]: + for d in range(m+1): + shifted_coord[d] -= box.c_pbcbox.box[m][d] + + # Get the cell index corresponding to the coord + cell_index = (shifted_coord[ZZ] / grid.cellsize[ZZ]) * (grid.ncells[XX] * grid.ncells[YY]) +\ + (shifted_coord[YY] / grid.cellsize[YY]) * grid.ncells[XX] + \ + (shifted_coord[XX] / grid.cellsize[XX]) + + # Just a safeguard + if cell_index >= grid.size: + continue + + # Check the cell index was not already selected + skip = False + for icheck in range(nchecked): + if already_checked[icheck] == cell_index: + skip = True + break + if skip: + continue + + # Search for neighbors inside this cell + for i_bead in range(grid.nbeads[cell_index]): + bid = grid.beadids[cell_index][i_bead] + + box.fast_pbc_dx(current_coords, &neighborcoords[bid, XX], dx) + + d2 = rvec_norm2(dx) + + if d2 < cutoff2: + if d2 < EPSILON: # Don't add the current bead as its own neighbor! + continue + + # Update neighbor lists + neighborhood.beadids[neighborhood.size] = bid + ### Modified here + neighborhood.beaddist[neighborhood.size] = d2 + ### + neighborhood.size += 1 + + if neighborhood.size >= neighborhood.allocated_size: + neighborhood.allocated_size += NEIGHBORHOOD_ALLOCATION_INCREMENT + neighborhood.beadids = realloc( neighborhood.beadids, neighborhood.allocated_size * sizeof(ns_int)) + ###Modified here + neighborhood.beaddist = realloc( neighborhood.beaddist, neighborhood.allocated_size * sizeof(real)) + ### + if neighborhood.beadids == NULL: + abort() + ###Modified + if neighborhood.beaddist == NULL: + abort() + ### + # Register the cell as checked + already_checked[nchecked] = cell_index + nchecked += 1 + + return neighborhood + + +cdef ns_neighborhood_holder *ns_core_parallel(real[:, ::1] refcoords, + real[:, ::1] neighborcoords, + ns_grid *grid, + PBCBox box, + real cutoff, + int nthreads=-1) nogil: + cdef ns_int coordid, i, j + cdef ns_int ncoords = refcoords.shape[0] + cdef ns_int ncoords_neighbors = neighborcoords.shape[0] + cdef real cutoff2 = cutoff * cutoff + cdef ns_neighborhood_holder *holder + + cdef ns_int *neighbor_buf + cdef ns_int buf_size, ibuf + + if nthreads < 0: + nthreads = openmp.omp_get_num_threads() + + holder = create_neighborhood_holder() + if holder == NULL: + fprintf(stderr,"FATAL: Could not allocate memory for NS holder\n", + sizeof(ns_int) * ncoords) + abort() + + holder.size = ncoords + holder.neighborhoods = malloc(sizeof(ns_neighborhood *) * ncoords) + if holder.neighborhoods == NULL: + fprintf(stderr,"FATAL: Could not allocate memory for NS holder.neighborhoods (requested: %i bytes)\n", + sizeof(ns_neighborhood) * ncoords) + abort() + + # Here starts the real core and the iteration over coordinates + for coordid in prange(ncoords, schedule='dynamic', num_threads=nthreads): + holder.neighborhoods[coordid] = retrieve_neighborhood(&refcoords[coordid, XX], + neighborcoords, + grid, + box, + cutoff2) + holder.neighborhoods[coordid].cutoff = cutoff + + return holder + +cdef ns_neighborhood_holder *ns_core(real[:, ::1] refcoords, + real[:, ::1] neighborcoords, + ns_grid *grid, + PBCBox box, + real cutoff) nogil: + cdef ns_int coordid, i, j + cdef ns_int ncoords = refcoords.shape[0] + cdef ns_int ncoords_neighbors = neighborcoords.shape[0] + cdef real cutoff2 = cutoff * cutoff + cdef ns_neighborhood_holder *holder + + cdef ns_int *neighbor_buf + cdef ns_int buf_size, ibuf + + holder = create_neighborhood_holder() + if holder == NULL: + fprintf(stderr,"FATAL: Could not allocate memory for NS holder\n", + sizeof(ns_int) * ncoords) + abort() + + holder.size = ncoords + holder.neighborhoods = malloc(sizeof(ns_neighborhood *) * ncoords) + if holder.neighborhoods == NULL: + fprintf(stderr,"FATAL: Could not allocate memory for NS holder.neighborhoods (requested: %i bytes)\n", + sizeof(ns_neighborhood) * ncoords) + abort() + + # Here starts the real core and the iteration over coordinates + for coordid in range(ncoords): + holder.neighborhoods[coordid] = retrieve_neighborhood(&refcoords[coordid, XX], + neighborcoords, + grid, + box, + cutoff2) + holder.neighborhoods[coordid].cutoff = cutoff + + return holder + + +# Python interface +cdef class FastNS(object): + cdef PBCBox box + cdef readonly int nthreads + cdef readonly real[:, ::1] coords + cdef real[:, ::1] coords_bbox + cdef readonly real cutoff + cdef bint prepared + cdef ns_grid *grid + + + def __init__(self, box): + if box.shape != (3, 3): + raise ValueError("Box must be provided as triclinic_dimensions (a 3x3 numpy.ndarray of unit cell vectors") + + self.box = PBCBox(box) + + self.nthreads = 1 + + self.coords = None + self.coords_bbox = None + + self.cutoff = -1 + + self.prepared = False + + self.grid = malloc(sizeof(ns_grid)) + + + def __dealloc__(self): + #destroy_nsgrid(self.grid) + self.grid.size = 0 + + #free(self.grid) + + def set_nthreads(self, nthreads, silent=False): + import multiprocessing + + if nthreads > multiprocessing.cpu_count(): + print("Warning: the number of threads requested if greater than the number of cores available. Performances may not be optimal!") + + if not silent: + print("Number of threads for NS adjusted to {}.".format(nthreads)) + + self.nthreads = nthreads + + + def set_coords(self, real[:, ::1] coords): + self.coords = coords + + # Make sure atoms are inside the brick-shaped box + self.coords_bbox = self.box.fast_put_atoms_in_bbox(coords) + + self.prepared = False + + + def set_cutoff(self, real cutoff): + self.cutoff = cutoff + + self.prepared = False + + + def prepare(self, force=False): + cdef ns_int i + cdef bint initialization_ok + + if self.prepared and not force: + print("NS already prepared, nothing to do!") + + if self.coords is None: + raise ValueError("Coordinates must be set before NS preparation!") + + if self.cutoff < 0: + raise ValueError("Cutoff must be set before NS preparation!") + + with nogil: + initialization_ok = False + + # Initializing grid + for i in range(DIM): + self.grid.ncells[i] = (self.box.c_pbcbox.box[i][i] / self.cutoff) + if self.grid.ncells[i] == 0: + self.grid.ncells[i] = 1 + self.grid.cellsize[i] = self.box.c_pbcbox.box[i][i] / self.grid.ncells[i] + + self.grid.size = self.grid.ncells[XX] * self.grid.ncells[YY] * self.grid.ncells[ZZ] + + # Populating grid + if populate_grid(self.grid, self.coords_bbox) == RET_OK: + initialization_ok = True + + + if initialization_ok: + self.prepared = True + else: + raise RuntimeError("Could not initialize NS grid") + + + def search(self, real[:, ::1]search_coords, return_ids=False): + cdef real[:, ::1] search_coords_bbox + cdef ns_int nid, i, j + cdef ns_neighborhood_holder *holder + cdef ns_neighborhood *neighborhood + + if not self.prepared: + self.prepare() + + + # Make sure atoms are inside the brick-shaped box + search_coords_bbox = self.box.fast_put_atoms_in_bbox(search_coords) + + + with nogil: + # Retrieve neighbors from grid + if self.nthreads == 1: + holder = ns_core(search_coords_bbox, self.coords_bbox, self.grid, self.box, self.cutoff) + else: + holder = ns_core_parallel(search_coords_bbox, self.coords_bbox, self.grid, self.box, self.cutoff, self.nthreads) + + + + neighbors = [] + ###Modify for distance + sqdist = [] + indx = [] + ### + for nid in range(holder.size): + neighborhood = holder.neighborhoods[nid] + + if return_ids: + neighborhood_py = np.empty(neighborhood.size, dtype=np.int64) + ###Modify for distance + neighborhood_dis = np.empty(neighborhood.size, dtype=np.float32) + neighborhood_indx = np.empty(neighborhood.size, dtype=np.int64) + ### + for i in range(neighborhood.size): + neighborhood_py[i] = neighborhood.beadids[i] + ###Modify for distance + neighborhood_dis[i] = neighborhood.beaddist[i] + ### + else: + neighborhood_py = np.empty((neighborhood.size, DIM), dtype=np.float32) + ###Modify for distance + neighborhood_dis = np.empty((neighborhood.size), dtype=np.float32) + neighborhood_indx = np.empty(neighborhood.size, dtype=np.int64) + ### + for i in range(neighborhood.size): + ###Modify for distance + neighborhood_dis[i] = neighborhood.beaddist[i] + neighborhood_indx[i] = neighborhood.beadids[i] + ### + + for j in range(DIM): + neighborhood_py[i,j] = self.coords[neighborhood.beadids[i], j] + + neighbors.append(neighborhood_py) + sqdist.append(neighborhood_dis) + indx.append(neighborhood_indx) + + # Free Memory + free_neighborhood_holder(holder) + + return neighbors, sqdist, indx + + + +__version__ = "26" \ No newline at end of file diff --git a/package/setup.py b/package/setup.py index 6e9a7b381c0..bd3aceec8df 100755 --- a/package/setup.py +++ b/package/setup.py @@ -380,9 +380,16 @@ def extensions(config): libraries=mathlib, define_macros=define_macros, extra_compile_args=extra_compile_args) + grid = MDAExtension('MDAnalysis.lib.grid', + ['MDAnalysis/lib/c_gridsearch' + source_suffix], + include_dirs=include_dirs, + libraries=mathlib + parallel_libraries, + define_macros=define_macros + parallel_macros, + extra_compile_args=extra_compile_args + parallel_args, + extra_link_args=parallel_args) pre_exts = [libdcd, distances, distances_omp, qcprot, transformation, libmdaxdr, util, encore_utils, - ap_clustering, spe_dimred, cutil, augment] + ap_clustering, spe_dimred, grid] cython_generated = [] if use_cython: From c3e227a863d5a505b453b0ce0167e5c8e406e0d9 Mon Sep 17 00:00:00 2001 From: ayush Date: Mon, 25 Jun 2018 22:06:44 -0700 Subject: [PATCH 02/47] removed the parallel handler to remove hard dependency on omp.h, can be added again for conditional compilation --- package/MDAnalysis/lib/c_gridsearch.pyx | 104 +----------------------- 1 file changed, 2 insertions(+), 102 deletions(-) diff --git a/package/MDAnalysis/lib/c_gridsearch.pyx b/package/MDAnalysis/lib/c_gridsearch.pyx index 82d7192544a..0a8089b9ef1 100644 --- a/package/MDAnalysis/lib/c_gridsearch.pyx +++ b/package/MDAnalysis/lib/c_gridsearch.pyx @@ -1,22 +1,3 @@ -# -*- coding: utf-8; Mode: python; tab-width: 4; indent-tabs-mode:nil; -*- -# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 -# -# Copyright (C) 2013-2016 Sébastien Buchoux -# -# This file is part of FATSLiM. -# -# FATSLiM is free software: you can redistribute it and/or modify -# it under the terms of the GNU General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# FATSLiM is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU General Public License -# along with FATSLiM. If not, see . #cython: cdivision=True #cython: boundscheck=False @@ -34,22 +15,14 @@ DEF GRID_ALLOCATION_INCREMENT = 50 DEF BOX_MARGIN=1.0010 DEF MAX_NTRICVEC=12 - -# Cython C imports (no Python here!) -from cython.parallel cimport prange from libc.stdlib cimport malloc, realloc, free, abort from libc.stdio cimport fprintf, stderr from libc.math cimport sqrt from libc.math cimport abs as real_abs -cimport openmp - - -# Python imports import numpy as np cimport numpy as np -# Ctypes ctypedef np.int_t ns_int ctypedef np.float32_t real ctypedef real rvec[DIM] @@ -93,7 +66,6 @@ cdef struct cPBCBox_t: ns_int[DIM] tric_shift[MAX_NTRICVEC] real[DIM] tric_vec[MAX_NTRICVEC] -# noinspection PyNoneFunctionAssignment cdef class PBCBox(object): cdef cPBCBox_t c_pbcbox cdef rvec center @@ -271,8 +243,6 @@ cdef class PBCBox(object): def put_atoms_in_bbox(self, real[:,::1] coords): return np.asarray(self.fast_put_atoms_in_bbox(coords)) - - ######################################################################################################################## # # Neighbor Search Stuff @@ -495,50 +465,6 @@ cdef ns_neighborhood *retrieve_neighborhood(rvec current_coords, real[:, ::1]nei nchecked += 1 return neighborhood - - -cdef ns_neighborhood_holder *ns_core_parallel(real[:, ::1] refcoords, - real[:, ::1] neighborcoords, - ns_grid *grid, - PBCBox box, - real cutoff, - int nthreads=-1) nogil: - cdef ns_int coordid, i, j - cdef ns_int ncoords = refcoords.shape[0] - cdef ns_int ncoords_neighbors = neighborcoords.shape[0] - cdef real cutoff2 = cutoff * cutoff - cdef ns_neighborhood_holder *holder - - cdef ns_int *neighbor_buf - cdef ns_int buf_size, ibuf - - if nthreads < 0: - nthreads = openmp.omp_get_num_threads() - - holder = create_neighborhood_holder() - if holder == NULL: - fprintf(stderr,"FATAL: Could not allocate memory for NS holder\n", - sizeof(ns_int) * ncoords) - abort() - - holder.size = ncoords - holder.neighborhoods = malloc(sizeof(ns_neighborhood *) * ncoords) - if holder.neighborhoods == NULL: - fprintf(stderr,"FATAL: Could not allocate memory for NS holder.neighborhoods (requested: %i bytes)\n", - sizeof(ns_neighborhood) * ncoords) - abort() - - # Here starts the real core and the iteration over coordinates - for coordid in prange(ncoords, schedule='dynamic', num_threads=nthreads): - holder.neighborhoods[coordid] = retrieve_neighborhood(&refcoords[coordid, XX], - neighborcoords, - grid, - box, - cutoff2) - holder.neighborhoods[coordid].cutoff = cutoff - - return holder - cdef ns_neighborhood_holder *ns_core(real[:, ::1] refcoords, real[:, ::1] neighborcoords, ns_grid *grid, @@ -577,7 +503,6 @@ cdef ns_neighborhood_holder *ns_core(real[:, ::1] refcoords, return holder - # Python interface cdef class FastNS(object): cdef PBCBox box @@ -610,21 +535,7 @@ cdef class FastNS(object): def __dealloc__(self): #destroy_nsgrid(self.grid) self.grid.size = 0 - - #free(self.grid) - - def set_nthreads(self, nthreads, silent=False): - import multiprocessing - - if nthreads > multiprocessing.cpu_count(): - print("Warning: the number of threads requested if greater than the number of cores available. Performances may not be optimal!") - - if not silent: - print("Number of threads for NS adjusted to {}.".format(nthreads)) - - self.nthreads = nthreads - - + def set_coords(self, real[:, ::1] coords): self.coords = coords @@ -689,15 +600,8 @@ cdef class FastNS(object): # Make sure atoms are inside the brick-shaped box search_coords_bbox = self.box.fast_put_atoms_in_bbox(search_coords) - with nogil: - # Retrieve neighbors from grid - if self.nthreads == 1: - holder = ns_core(search_coords_bbox, self.coords_bbox, self.grid, self.box, self.cutoff) - else: - holder = ns_core_parallel(search_coords_bbox, self.coords_bbox, self.grid, self.box, self.cutoff, self.nthreads) - - + holder = ns_core(search_coords_bbox, self.coords_bbox, self.grid, self.box, self.cutoff) neighbors = [] ###Modify for distance @@ -741,7 +645,3 @@ cdef class FastNS(object): free_neighborhood_holder(holder) return neighbors, sqdist, indx - - - -__version__ = "26" \ No newline at end of file From 007d47da7c410289fa65d71f5420e6a95cc4bb17 Mon Sep 17 00:00:00 2001 From: ayush Date: Sat, 30 Jun 2018 01:23:12 -0700 Subject: [PATCH 03/47] added the license header --- package/MDAnalysis/lib/c_gridsearch.pyx | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/package/MDAnalysis/lib/c_gridsearch.pyx b/package/MDAnalysis/lib/c_gridsearch.pyx index 0a8089b9ef1..a4e8745aec4 100644 --- a/package/MDAnalysis/lib/c_gridsearch.pyx +++ b/package/MDAnalysis/lib/c_gridsearch.pyx @@ -1,3 +1,22 @@ +# -*- coding: utf-8; Mode: python; tab-width: 4; indent-tabs-mode:nil; -*- +# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 +# +# Copyright (C) 2013-2016 Sébastien Buchoux +# +# This file is part of FATSLiM. +# +# FATSLiM is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# FATSLiM is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with FATSLiM. If not, see . #cython: cdivision=True #cython: boundscheck=False From fc61e5e1ef60360bada560addbe4f5782e955c7a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Buchoux?= Date: Sun, 1 Jul 2018 14:06:40 +0200 Subject: [PATCH 04/47] Added MDAnalysis header file to c_gridsearch.pyx --- package/MDAnalysis/lib/c_gridsearch.pyx | 32 +++++++++++++++---------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/package/MDAnalysis/lib/c_gridsearch.pyx b/package/MDAnalysis/lib/c_gridsearch.pyx index a4e8745aec4..ac73c97a5a2 100644 --- a/package/MDAnalysis/lib/c_gridsearch.pyx +++ b/package/MDAnalysis/lib/c_gridsearch.pyx @@ -1,22 +1,28 @@ -# -*- coding: utf-8; Mode: python; tab-width: 4; indent-tabs-mode:nil; -*- +# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; -*- # vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 # -# Copyright (C) 2013-2016 Sébastien Buchoux +# MDAnalysis --- https://www.mdanalysis.org # -# This file is part of FATSLiM. +# Copyright (C) 2013-2018 Sébastien Buchoux +# Copyright (c) 2018 The MDAnalysis Development Team and contributors +# (see the file AUTHORS for the full list of names) # -# FATSLiM is free software: you can redistribute it and/or modify -# it under the terms of the GNU General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. +# Released under the GNU Public Licence, v3 or any higher version # -# FATSLiM is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. +# Please cite your use of MDAnalysis in published work: # -# You should have received a copy of the GNU General Public License -# along with FATSLiM. If not, see . +# 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 +# +# + #cython: cdivision=True #cython: boundscheck=False From 4b273a9c50b6235713b52920537c9107a799884a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Buchoux?= Date: Sun, 1 Jul 2018 16:53:22 +0200 Subject: [PATCH 05/47] Corrected C (de)allocation in c_gridsearch.pyx. Should fix memory leak --- package/MDAnalysis/lib/c_gridsearch.pyx | 32 +++++++++++++++---------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/package/MDAnalysis/lib/c_gridsearch.pyx b/package/MDAnalysis/lib/c_gridsearch.pyx index ac73c97a5a2..2a7af4b33b7 100644 --- a/package/MDAnalysis/lib/c_gridsearch.pyx +++ b/package/MDAnalysis/lib/c_gridsearch.pyx @@ -145,16 +145,15 @@ cdef class PBCBox(object): # Update shift vectors self.c_pbcbox.ntric_vec = 0 - # We will only use single shifts, but we will check a few - # more shifts to see if there is a limiting distance - # above which we can not be sure of the correct distance. - for kk in range(5): + + # We will only use single shifts + for kk in range(3): k = order[kk] - for jj in range(5): + for jj in range(3): j = order[jj] - for ii in range(5): + for ii in range(3): i = order[ii] # A shift is only useful when it is trilinic @@ -177,8 +176,8 @@ cdef class PBCBox(object): else: pos[d] = max(-self.c_pbcbox.hbox_diag[d], -trial[d]) - d2old += sqrt(pos[d]) - d2new += sqrt(pos[d] + trial[d]) + d2old += pos[d]**2 + d2new += (pos[d] + trial[d])**2 if BOX_MARGIN*d2new < d2old: if not (j < -1 or j > 1 or k < -1 or k > 1): @@ -196,7 +195,7 @@ cdef class PBCBox(object): d2new_c = 0 for d in range(DIM): - d2new_c += sqrt(pos[d] + trial[d] - shift*box[dd, d]) + d2new_c += (pos[d] + trial[d] - shift*box[dd, d])**2 if d2new_c <= BOX_MARGIN*d2new: use = False @@ -209,7 +208,13 @@ cdef class PBCBox(object): % MAX_NTRICVEC) print(" There is probably something wrong with " "your box.") - print(box) + print(np.array(box)) + + for i in range(self.c_pbcbox.ntric_vec): + print(" -> shift #{}: [{}, {}, {}]".format(i+1, + self.c_pbcbox.tric_shift[i][XX], + self.c_pbcbox.tric_shift[i][YY], + self.c_pbcbox.tric_shift[i][ZZ])) else: for d in range(DIM): self.c_pbcbox.tric_vec[self.c_pbcbox.ntric_vec][d] = \ @@ -539,6 +544,9 @@ cdef class FastNS(object): cdef ns_grid *grid + def __cinit__(self): + self.grid = malloc(sizeof(ns_grid)) + def __init__(self, box): if box.shape != (3, 3): raise ValueError("Box must be provided as triclinic_dimensions (a 3x3 numpy.ndarray of unit cell vectors") @@ -554,11 +562,9 @@ cdef class FastNS(object): self.prepared = False - self.grid = malloc(sizeof(ns_grid)) - def __dealloc__(self): - #destroy_nsgrid(self.grid) + destroy_nsgrid(self.grid) self.grid.size = 0 def set_coords(self, real[:, ::1] coords): From 02b5819a8480a69fdb09bd708155f7d38f29eeb5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Buchoux?= Date: Sun, 1 Jul 2018 16:54:21 +0200 Subject: [PATCH 06/47] Added Tests for grid search. Added grid in MDAnalysis.lib.__init__ --- package/MDAnalysis/lib/__init__.py | 3 +- .../MDAnalysisTests/lib/test_gridsearch.py | 81 +++++++++++++++++++ 2 files changed, 83 insertions(+), 1 deletion(-) create mode 100644 testsuite/MDAnalysisTests/lib/test_gridsearch.py diff --git a/package/MDAnalysis/lib/__init__.py b/package/MDAnalysis/lib/__init__.py index 51c148dde65..a740bd8d5bf 100644 --- a/package/MDAnalysis/lib/__init__.py +++ b/package/MDAnalysis/lib/__init__.py @@ -29,7 +29,7 @@ from __future__ import absolute_import __all__ = ['log', 'transformations', 'util', 'mdamath', 'distances', - 'NeighborSearch', 'formats', 'pkdtree'] + 'NeighborSearch', 'formats', 'pkdtree', 'grid'] from . import log from . import transformations @@ -39,3 +39,4 @@ from . import NeighborSearch from . import formats from . import pkdtree +from . import grid \ No newline at end of file diff --git a/testsuite/MDAnalysisTests/lib/test_gridsearch.py b/testsuite/MDAnalysisTests/lib/test_gridsearch.py new file mode 100644 index 00000000000..cab23f3f372 --- /dev/null +++ b/testsuite/MDAnalysisTests/lib/test_gridsearch.py @@ -0,0 +1,81 @@ +# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*- +# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 fileencoding=utf-8 +# +# MDAnalysis --- https://www.mdanalysis.org +# Copyright (c) 2006-2018 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 +# + +from __future__ import print_function, absolute_import + +import pytest +from numpy.testing import assert_equal +import numpy as np + +import MDAnalysis as mda +from MDAnalysis.lib import grid +from MDAnalysis.lib.pkdtree import PeriodicKDTree +from MDAnalysis.lib.mdamath import triclinic_vectors + +from MDAnalysisTests.datafiles import GRO + +@pytest.fixture +def universe(): + u = mda.Universe(GRO) + return u + +def run_search(universe, ref_id): + cutoff = 3 + + coords = universe.atoms.positions + ref_pos = coords[ref_id] + triclinic_box = triclinic_vectors(universe.dimensions) + + # Run pkdtree search + pkdt = PeriodicKDTree(universe.atoms.dimensions, bucket_size=10) + pkdt.set_coords(coords) + pkdt.search(ref_pos, cutoff) + + results_pkdtree = pkdt.get_indices() + results_pkdtree.remove(ref_id) + results_pkdtree = np.array(results_pkdtree) + results_pkdtree.sort() + + # Run grid search + searcher = grid.FastNS(triclinic_box) + searcher.set_cutoff(cutoff) + searcher.set_coords(coords) + searcher.prepare() + + results_grid = searcher.search(np.array([ref_pos, ]), return_ids=True)[0][0] + results_grid.sort() + + return results_pkdtree, results_grid + +def test_gridsearch(universe): + """Check that pkdtree and grid search return the same results (No PBC needed)""" + + ref_id = 0 + results_pkdtree, results_grid = run_search(universe, ref_id) + assert_equal(results_pkdtree, results_grid) + +def test_gridsearch_PBC(universe): + """Check that pkdtree and grid search return the same results (PBC needed)""" + + ref_id = 13937 + results_pkdtree, results_grid = run_search(universe, ref_id) + assert_equal(results_pkdtree, results_grid) From 0f35031499d97ed8e0bbbd429080a9e5c68af8c1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Buchoux?= Date: Wed, 4 Jul 2018 16:17:11 +0200 Subject: [PATCH 07/47] Removed pointer to nsgrid structure to avoid need to free it --- package/MDAnalysis/lib/c_gridsearch.pyx | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/package/MDAnalysis/lib/c_gridsearch.pyx b/package/MDAnalysis/lib/c_gridsearch.pyx index 2a7af4b33b7..7a6de80389d 100644 --- a/package/MDAnalysis/lib/c_gridsearch.pyx +++ b/package/MDAnalysis/lib/c_gridsearch.pyx @@ -42,8 +42,7 @@ DEF MAX_NTRICVEC=12 from libc.stdlib cimport malloc, realloc, free, abort from libc.stdio cimport fprintf, stderr -from libc.math cimport sqrt -from libc.math cimport abs as real_abs + import numpy as np cimport numpy as np @@ -541,11 +540,7 @@ cdef class FastNS(object): cdef real[:, ::1] coords_bbox cdef readonly real cutoff cdef bint prepared - cdef ns_grid *grid - - - def __cinit__(self): - self.grid = malloc(sizeof(ns_grid)) + cdef ns_grid grid def __init__(self, box): if box.shape != (3, 3): @@ -562,9 +557,11 @@ cdef class FastNS(object): self.prepared = False + self.grid.size = 0 + def __dealloc__(self): - destroy_nsgrid(self.grid) + destroy_nsgrid(&self.grid) self.grid.size = 0 def set_coords(self, real[:, ::1] coords): @@ -608,7 +605,7 @@ cdef class FastNS(object): self.grid.size = self.grid.ncells[XX] * self.grid.ncells[YY] * self.grid.ncells[ZZ] # Populating grid - if populate_grid(self.grid, self.coords_bbox) == RET_OK: + if populate_grid(&self.grid, self.coords_bbox) == RET_OK: initialization_ok = True @@ -632,7 +629,7 @@ cdef class FastNS(object): search_coords_bbox = self.box.fast_put_atoms_in_bbox(search_coords) with nogil: - holder = ns_core(search_coords_bbox, self.coords_bbox, self.grid, self.box, self.cutoff) + holder = ns_core(search_coords_bbox, self.coords_bbox, &self.grid, self.box, self.cutoff) neighbors = [] ###Modify for distance From bdad5e3b3b38bb4ac96ec2b1dce58af4b067bc56 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Buchoux?= Date: Wed, 4 Jul 2018 18:13:41 +0200 Subject: [PATCH 08/47] Removed abort() from c_gridsearch.pyx --- package/MDAnalysis/lib/c_gridsearch.pyx | 131 +++++++----------- .../MDAnalysisTests/lib/test_gridsearch.py | 2 +- 2 files changed, 48 insertions(+), 85 deletions(-) diff --git a/package/MDAnalysis/lib/c_gridsearch.pyx b/package/MDAnalysis/lib/c_gridsearch.pyx index 7a6de80389d..6008f0387dd 100644 --- a/package/MDAnalysis/lib/c_gridsearch.pyx +++ b/package/MDAnalysis/lib/c_gridsearch.pyx @@ -31,18 +31,19 @@ DEF DIM = 3 DEF XX = 0 DEF YY = 1 DEF ZZ = 2 + +DEF RET_ALLOCATIONERROR = 2 DEF RET_OK = 1 DEF RET_ERROR = 0 DEF EPSILON = 1e-5 + DEF NEIGHBORHOOD_ALLOCATION_INCREMENT = 50 DEF GRID_ALLOCATION_INCREMENT = 50 DEF BOX_MARGIN=1.0010 DEF MAX_NTRICVEC=12 -from libc.stdlib cimport malloc, realloc, free, abort -from libc.stdio cimport fprintf, stderr - +from libc.stdlib cimport malloc, realloc, free import numpy as np cimport numpy as np @@ -64,8 +65,7 @@ cdef struct ns_neighborhood: ns_int allocated_size ns_int size ns_int *beadids - real *beaddist - ### + cdef struct ns_neighborhood_holder: ns_int size ns_neighborhood **neighborhoods @@ -80,6 +80,7 @@ cdef void rvec_clear(rvec a) nogil: a[YY]=0.0 a[ZZ]=0.0 + cdef struct cPBCBox_t: matrix box rvec fbox_diag @@ -90,6 +91,8 @@ cdef struct cPBCBox_t: ns_int[DIM] tric_shift[MAX_NTRICVEC] real[DIM] tric_vec[MAX_NTRICVEC] + +# Class to handle PBC calculations cdef class PBCBox(object): cdef cPBCBox_t c_pbcbox cdef rvec center @@ -232,6 +235,7 @@ cdef class PBCBox(object): raise ValueError("Box does not correspond to PBC=xyz") self.fast_update(box) + cdef void fast_pbc_dx(self, rvec ref, rvec other, rvec dx) nogil: cdef ns_int i, j cdef rvec dx_start, trial @@ -284,20 +288,6 @@ cdef struct ns_grid: ns_int *nbeads ns_int **beadids -cdef ns_grid initialize_nsgrid(matrix box, - float cutoff) nogil: - cdef ns_grid grid - cdef ns_int i - - for i in range(DIM): - grid.ncells[i] = (box[i][i] / cutoff) - if grid.ncells[i] == 0: - grid.ncells[i] = 1 - grid.cellsize[i] = box[i][i] / grid.ncells[i] - - grid.size = grid.ncells[XX] * grid.ncells[YY] * grid.ncells[ZZ] - return grid - cdef ns_int populate_grid(ns_grid *grid, real[:,::1] coords) nogil: cdef ns_int ncoords = coords.shape[0] @@ -323,15 +313,12 @@ cdef ns_int populate_grid_array(ns_grid *grid, # Allocate memory grid.nbeads = malloc(sizeof(ns_int) * grid_size) if grid.nbeads == NULL: - fprintf(stderr,"FATAL: Could not allocate memory for NS grid.nbeads (requested: %i bytes)\n", - sizeof(ns_int) * grid_size) - abort() + return RET_ALLOCATIONERROR allocated_size = malloc(sizeof(ns_int) * grid_size) if allocated_size == NULL: - fprintf(stderr,"FATAL: Could not allocate memory for NS allocated_size (requested: %i bytes)\n", - sizeof(ns_int) * grid_size) - abort() + # No need to free grid.nbeads as it will be freed by destroy_nsgrid + return RET_ALLOCATIONERROR for i in range(grid_size): grid.nbeads[i] = 0 @@ -339,16 +326,12 @@ cdef ns_int populate_grid_array(ns_grid *grid, grid.beadids = malloc(sizeof(ns_int *) * grid_size) if grid.beadids == NULL: - fprintf(stderr,"FATAL: Could not allocate memory for NS grid.beadids (requested: %i bytes)\n", - sizeof(ns_int *) * grid_size) - abort() + return RET_ALLOCATIONERROR for i in range(grid_size): grid.beadids[i] = malloc(sizeof(ns_int) * allocated_size[i]) if grid.beadids[i] == NULL: - fprintf(stderr,"FATAL: Could not allocate memory for NS grid.beadids[i] (requested: %i bytes)\n", - sizeof(ns_int) * allocated_size[i]) - abort() + return RET_ALLOCATIONERROR # Get cell indices for coords for i in range(ncoords): @@ -380,6 +363,8 @@ cdef ns_neighborhood_holder *create_neighborhood_holder() nogil: cdef ns_neighborhood_holder *holder holder = malloc(sizeof(ns_neighborhood_holder)) + holder.size = 0 + holder.neighborhoods = NULL return holder @@ -393,7 +378,9 @@ cdef void free_neighborhood_holder(ns_neighborhood_holder *holder) nogil: if holder.neighborhoods[i].beadids != NULL: free(holder.neighborhoods[i].beadids) free(holder.neighborhoods[i]) - free(holder.neighborhoods) + + if holder.neighborhoods != NULL: + free(holder.neighborhoods) free(holder) cdef ns_neighborhood *retrieve_neighborhood(rvec current_coords, real[:, ::1]neighborcoords, ns_grid *grid, PBCBox box, real cutoff2) nogil: @@ -409,16 +396,15 @@ cdef ns_neighborhood *retrieve_neighborhood(rvec current_coords, real[:, ::1]nei cdef ns_neighborhood *neighborhood = malloc(sizeof(ns_neighborhood)) if neighborhood == NULL: - abort() + return NULL neighborhood.size = 0 neighborhood.allocated_size = NEIGHBORHOOD_ALLOCATION_INCREMENT neighborhood.beadids = malloc(NEIGHBORHOOD_ALLOCATION_INCREMENT * sizeof(ns_int)) - ###Modified here - neighborhood.beaddist = malloc(NEIGHBORHOOD_ALLOCATION_INCREMENT * sizeof(real)) - ### + if neighborhood.beadids == NULL: - abort() + free(neighborhood) + return NULL for zi in range(3): for yi in range(3): @@ -472,28 +458,23 @@ cdef ns_neighborhood *retrieve_neighborhood(rvec current_coords, real[:, ::1]nei # Update neighbor lists neighborhood.beadids[neighborhood.size] = bid - ### Modified here - neighborhood.beaddist[neighborhood.size] = d2 - ### neighborhood.size += 1 if neighborhood.size >= neighborhood.allocated_size: neighborhood.allocated_size += NEIGHBORHOOD_ALLOCATION_INCREMENT neighborhood.beadids = realloc( neighborhood.beadids, neighborhood.allocated_size * sizeof(ns_int)) - ###Modified here - neighborhood.beaddist = realloc( neighborhood.beaddist, neighborhood.allocated_size * sizeof(real)) - ### + if neighborhood.beadids == NULL: - abort() - ###Modified - if neighborhood.beaddist == NULL: - abort() - ### + free(neighborhood) + return NULL + # Register the cell as checked already_checked[nchecked] = cell_index nchecked += 1 return neighborhood + + cdef ns_neighborhood_holder *ns_core(real[:, ::1] refcoords, real[:, ::1] neighborcoords, ns_grid *grid, @@ -510,16 +491,12 @@ cdef ns_neighborhood_holder *ns_core(real[:, ::1] refcoords, holder = create_neighborhood_holder() if holder == NULL: - fprintf(stderr,"FATAL: Could not allocate memory for NS holder\n", - sizeof(ns_int) * ncoords) - abort() + return NULL - holder.size = ncoords holder.neighborhoods = malloc(sizeof(ns_neighborhood *) * ncoords) if holder.neighborhoods == NULL: - fprintf(stderr,"FATAL: Could not allocate memory for NS holder.neighborhoods (requested: %i bytes)\n", - sizeof(ns_neighborhood) * ncoords) - abort() + free_neighborhood_holder(holder) + return NULL # Here starts the real core and the iteration over coordinates for coordid in range(ncoords): @@ -528,7 +505,12 @@ cdef ns_neighborhood_holder *ns_core(real[:, ::1] refcoords, grid, box, cutoff2) + if holder.neighborhoods[coordid] == NULL: + free_neighborhood_holder(holder) + return NULL + holder.neighborhoods[coordid].cutoff = cutoff + holder.size += 1 return holder @@ -580,7 +562,7 @@ cdef class FastNS(object): def prepare(self, force=False): - cdef ns_int i + cdef ns_int i, retcode cdef bint initialization_ok if self.prepared and not force: @@ -605,12 +587,11 @@ cdef class FastNS(object): self.grid.size = self.grid.ncells[XX] * self.grid.ncells[YY] * self.grid.ncells[ZZ] # Populating grid - if populate_grid(&self.grid, self.coords_bbox) == RET_OK: - initialization_ok = True - - - if initialization_ok: + retcode = populate_grid(&self.grid, self.coords_bbox) + if retcode == RET_OK: self.prepared = True + elif retcode == RET_ALLOCATIONERROR: + raise MemoryError("Could not allocate memory to initialize NS grid") else: raise RuntimeError("Could not initialize NS grid") @@ -631,45 +612,27 @@ cdef class FastNS(object): with nogil: holder = ns_core(search_coords_bbox, self.coords_bbox, &self.grid, self.box, self.cutoff) + if holder == NULL: + raise MemoryError("Could not allocate memory to run NS core") + neighbors = [] - ###Modify for distance - sqdist = [] - indx = [] - ### for nid in range(holder.size): neighborhood = holder.neighborhoods[nid] if return_ids: neighborhood_py = np.empty(neighborhood.size, dtype=np.int64) - ###Modify for distance - neighborhood_dis = np.empty(neighborhood.size, dtype=np.float32) - neighborhood_indx = np.empty(neighborhood.size, dtype=np.int64) - ### + for i in range(neighborhood.size): neighborhood_py[i] = neighborhood.beadids[i] - ###Modify for distance - neighborhood_dis[i] = neighborhood.beaddist[i] - ### else: neighborhood_py = np.empty((neighborhood.size, DIM), dtype=np.float32) - ###Modify for distance - neighborhood_dis = np.empty((neighborhood.size), dtype=np.float32) - neighborhood_indx = np.empty(neighborhood.size, dtype=np.int64) - ### for i in range(neighborhood.size): - ###Modify for distance - neighborhood_dis[i] = neighborhood.beaddist[i] - neighborhood_indx[i] = neighborhood.beadids[i] - ### - for j in range(DIM): neighborhood_py[i,j] = self.coords[neighborhood.beadids[i], j] neighbors.append(neighborhood_py) - sqdist.append(neighborhood_dis) - indx.append(neighborhood_indx) # Free Memory free_neighborhood_holder(holder) - return neighbors, sqdist, indx + return neighbors diff --git a/testsuite/MDAnalysisTests/lib/test_gridsearch.py b/testsuite/MDAnalysisTests/lib/test_gridsearch.py index cab23f3f372..c8d526e9074 100644 --- a/testsuite/MDAnalysisTests/lib/test_gridsearch.py +++ b/testsuite/MDAnalysisTests/lib/test_gridsearch.py @@ -61,7 +61,7 @@ def run_search(universe, ref_id): searcher.set_coords(coords) searcher.prepare() - results_grid = searcher.search(np.array([ref_pos, ]), return_ids=True)[0][0] + results_grid = searcher.search(np.array([ref_pos, ]), return_ids=True)[0] results_grid.sort() return results_pkdtree, results_grid From f65e32bc89b930f638f81a0d87313ece89bdc471 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Buchoux?= Date: Sat, 7 Jul 2018 10:30:10 +0200 Subject: [PATCH 09/47] grid allocation moved to FastNS method --- package/MDAnalysis/lib/c_gridsearch.pyx | 158 ++++++++++-------------- 1 file changed, 67 insertions(+), 91 deletions(-) diff --git a/package/MDAnalysis/lib/c_gridsearch.pyx b/package/MDAnalysis/lib/c_gridsearch.pyx index 6008f0387dd..98d97eded65 100644 --- a/package/MDAnalysis/lib/c_gridsearch.pyx +++ b/package/MDAnalysis/lib/c_gridsearch.pyx @@ -32,13 +32,9 @@ DEF XX = 0 DEF YY = 1 DEF ZZ = 2 -DEF RET_ALLOCATIONERROR = 2 -DEF RET_OK = 1 -DEF RET_ERROR = 0 DEF EPSILON = 1e-5 DEF NEIGHBORHOOD_ALLOCATION_INCREMENT = 50 -DEF GRID_ALLOCATION_INCREMENT = 50 DEF BOX_MARGIN=1.0010 DEF MAX_NTRICVEC=12 @@ -288,77 +284,6 @@ cdef struct ns_grid: ns_int *nbeads ns_int **beadids -cdef ns_int populate_grid(ns_grid *grid, - real[:,::1] coords) nogil: - cdef ns_int ncoords = coords.shape[0] - cdef bint ret_val - - ret_val = populate_grid_array(grid, - &coords[0, 0], - ncoords) - - return ret_val - -cdef ns_int populate_grid_array(ns_grid *grid, - rvec *coords, - ns_int ncoords) nogil: - cdef ns_int i, cellindex = -1 - cdef ns_int grid_size = grid.size - cdef ns_int *allocated_size = NULL - - if grid_size != grid.ncells[XX] * grid.ncells[YY] * grid.ncells[ZZ]: # Grid not initialized - return RET_ERROR - - - # Allocate memory - grid.nbeads = malloc(sizeof(ns_int) * grid_size) - if grid.nbeads == NULL: - return RET_ALLOCATIONERROR - - allocated_size = malloc(sizeof(ns_int) * grid_size) - if allocated_size == NULL: - # No need to free grid.nbeads as it will be freed by destroy_nsgrid - return RET_ALLOCATIONERROR - - for i in range(grid_size): - grid.nbeads[i] = 0 - allocated_size[i] = GRID_ALLOCATION_INCREMENT - - grid.beadids = malloc(sizeof(ns_int *) * grid_size) - if grid.beadids == NULL: - return RET_ALLOCATIONERROR - - for i in range(grid_size): - grid.beadids[i] = malloc(sizeof(ns_int) * allocated_size[i]) - if grid.beadids[i] == NULL: - return RET_ALLOCATIONERROR - - # Get cell indices for coords - for i in range(ncoords): - cellindex = (coords[i][ZZ] / grid.cellsize[ZZ]) * (grid.ncells[XX] * grid.ncells[YY]) +\ - (coords[i][YY] / grid.cellsize[YY]) * grid.ncells[XX] + \ - (coords[i][XX] / grid.cellsize[XX]) - - grid.beadids[cellindex][grid.nbeads[cellindex]] = i - grid.nbeads[cellindex] += 1 - - if grid.nbeads[cellindex] >= allocated_size[cellindex]: - allocated_size[cellindex] += GRID_ALLOCATION_INCREMENT - grid.beadids[cellindex] = realloc( grid.beadids[cellindex], sizeof(ns_int) * allocated_size[cellindex]) - free(allocated_size) - return RET_OK - -cdef void destroy_nsgrid(ns_grid *grid) nogil: - cdef ns_int i - if grid.nbeads != NULL: - free(grid.nbeads) - - for i in range(grid.size): - if grid.beadids[i] != NULL: - free(grid.beadids[i]) - free(grid.beadids) - - cdef ns_neighborhood_holder *create_neighborhood_holder() nogil: cdef ns_neighborhood_holder *holder @@ -517,7 +442,6 @@ cdef ns_neighborhood_holder *ns_core(real[:, ::1] refcoords, # Python interface cdef class FastNS(object): cdef PBCBox box - cdef readonly int nthreads cdef readonly real[:, ::1] coords cdef real[:, ::1] coords_bbox cdef readonly real cutoff @@ -530,8 +454,6 @@ cdef class FastNS(object): self.box = PBCBox(box) - self.nthreads = 1 - self.coords = None self.coords_bbox = None @@ -543,7 +465,15 @@ cdef class FastNS(object): def __dealloc__(self): - destroy_nsgrid(&self.grid) + cdef ns_int i + if self.grid.nbeads != NULL: + free(self.grid.nbeads) + + for i in range(self.grid.size): + if self.grid.beadids[i] != NULL: + free(self.grid.beadids[i]) + free(self.grid.beadids) + self.grid.size = 0 def set_coords(self, real[:, ::1] coords): @@ -562,8 +492,11 @@ cdef class FastNS(object): def prepare(self, force=False): - cdef ns_int i, retcode - cdef bint initialization_ok + cdef ns_int i, cellindex = -1 + cdef ns_int *allocated_size = NULL + cdef ns_int ncoords = self.coords.shape[0] + cdef ns_int allocation_guess + cdef rvec *coords = &self.coords_bbox[0, 0] if self.prepared and not force: print("NS already prepared, nothing to do!") @@ -575,7 +508,6 @@ cdef class FastNS(object): raise ValueError("Cutoff must be set before NS preparation!") with nogil: - initialization_ok = False # Initializing grid for i in range(DIM): @@ -583,17 +515,61 @@ cdef class FastNS(object): if self.grid.ncells[i] == 0: self.grid.ncells[i] = 1 self.grid.cellsize[i] = self.box.c_pbcbox.box[i][i] / self.grid.ncells[i] - self.grid.size = self.grid.ncells[XX] * self.grid.ncells[YY] * self.grid.ncells[ZZ] - # Populating grid - retcode = populate_grid(&self.grid, self.coords_bbox) - if retcode == RET_OK: - self.prepared = True - elif retcode == RET_ALLOCATIONERROR: - raise MemoryError("Could not allocate memory to initialize NS grid") - else: - raise RuntimeError("Could not initialize NS grid") + # This is just a guess on how much memory we might for each grid cell: + # we just assume an average bead density and we take four times this density just to be safe + allocation_guess = (4 * (ncoords / self.grid.size + 1)) + + # Allocate memory for the grid + self.grid.nbeads = malloc(sizeof(ns_int) * self.grid.size) + if self.grid.nbeads == NULL: + with gil: + raise MemoryError("Could not allocate memory for NS grid") + + # Allocate memory from temporary allocation counter + allocated_size = malloc(sizeof(ns_int) * self.grid.size) + if allocated_size == NULL: + # No need to free grid.nbeads as it will be freed by destroy_nsgrid called by __dealloc___ + with gil: + raise MemoryError("Could not allocate memory for allocation buffer") + + # Pre-allocate some memory for grid cells + for i in range(self.grid.size): + self.grid.nbeads[i] = 0 + allocated_size[i] = allocation_guess + + self.grid.beadids = malloc(sizeof(ns_int *) * self.grid.size) + if self.grid.beadids == NULL: + with gil: + raise MemoryError("Could not allocate memory for grid cells") + + for i in range(self.grid.size): + self.grid.beadids[i] = malloc(sizeof(ns_int) * allocated_size[i]) + if self.grid.beadids[i] == NULL: + with gil: + raise MemoryError("Could not allocate memory for grid cell") + + # Populate grid cells using the coordinates (ie do the heavy work) + for i in range(ncoords): + cellindex = (coords[i][ZZ] / self.grid.cellsize[ZZ]) * (self.grid.ncells[XX] * self.grid.ncells[YY]) +\ + (coords[i][YY] / self.grid.cellsize[YY]) * self.grid.ncells[XX] + \ + (coords[i][XX] / self.grid.cellsize[XX]) + + self.grid.beadids[cellindex][self.grid.nbeads[cellindex]] = i + self.grid.nbeads[cellindex] += 1 + + # We need to allocate more memory (simply double the amount of memory as + # 1. it should barely be needed + # 2. the size should stay fairly reasonable + if self.grid.nbeads[cellindex] >= allocated_size[cellindex]: + allocated_size[cellindex] *= 2 + self.grid.beadids[cellindex] = realloc( self.grid.beadids[cellindex], sizeof(ns_int) * allocated_size[cellindex]) + + # Now we can free the allocation buffer + free(allocated_size) + + self.prepared = True def search(self, real[:, ::1]search_coords, return_ids=False): From ff33b2f62769b32241f0860af7b66ecfdbcfda42 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Buchoux?= Date: Sun, 8 Jul 2018 10:35:17 +0200 Subject: [PATCH 10/47] NSResults object added to store results from NS. Started documentation --- package/MDAnalysis/lib/c_gridsearch.pyx | 235 ++++++++++++++++-- .../source/documentation_pages/lib/grid.rst | 2 + .../documentation_pages/lib_modules.rst | 2 + .../MDAnalysisTests/lib/test_gridsearch.py | 79 +++++- 4 files changed, 282 insertions(+), 36 deletions(-) create mode 100644 package/doc/sphinx/source/documentation_pages/lib/grid.rst diff --git a/package/MDAnalysis/lib/c_gridsearch.pyx b/package/MDAnalysis/lib/c_gridsearch.pyx index 98d97eded65..9c50a5b1777 100644 --- a/package/MDAnalysis/lib/c_gridsearch.pyx +++ b/package/MDAnalysis/lib/c_gridsearch.pyx @@ -25,6 +25,15 @@ #cython: cdivision=True #cython: boundscheck=False +#cython: initializedcheck=False + +""" +Neighbor search library --- :mod:`MDAnalysis.lib.grid` +====================================================== + +This Neighbor search library is a serialized Cython port of the NS grid search implemented in GROMACS. +""" + # Preprocessor DEFs DEF DIM = 3 @@ -40,6 +49,7 @@ DEF BOX_MARGIN=1.0010 DEF MAX_NTRICVEC=12 from libc.stdlib cimport malloc, realloc, free +from libc.math cimport sqrt import numpy as np cimport numpy as np @@ -248,16 +258,21 @@ cdef class PBCBox(object): for j in range (i, -1, -1): dx[j] += self.c_pbcbox.box[i][j] + cdef real fast_distance2(self, rvec a, rvec b) nogil: + cdef rvec dx + self.fast_pbc_dx(a, b, dx) + return rvec_norm2(dx) + + cdef real fast_distance(self, rvec a, rvec b) nogil: + return sqrt(self.fast_distance2(a,b)) + cdef real[:, ::1]fast_put_atoms_in_bbox(self, real[:,::1] coords) nogil: cdef ns_int i, m, d, natoms, wd = 0 cdef real[:,::1] bbox_coords natoms = coords.shape[0] with gil: - if natoms == 0: - bbox_coords = np.empty((0, DIM)) - else: - bbox_coords = coords.copy() + bbox_coords = coords.copy() for i in range(natoms): for m in range(DIM - 1, -1, -1): @@ -270,6 +285,8 @@ cdef class PBCBox(object): return bbox_coords def put_atoms_in_bbox(self, real[:,::1] coords): + if coords.shape[0] == 0: + return np.zeros((0, DIM), dtype=np.float32) return np.asarray(self.fast_put_atoms_in_bbox(coords)) ######################################################################################################################## @@ -439,6 +456,167 @@ cdef ns_neighborhood_holder *ns_core(real[:, ::1] refcoords, return holder +cdef class NSResults(object): + """ + Class used to store results returned by `MDAnalysis.lib.grid.FastNS.search` + """ + cdef PBCBox box + cdef readonly real cutoff + cdef real[:, ::1] grid_coords + cdef real[:, ::1] ref_coords + cdef ns_int **nids + cdef ns_int *nsizes + cdef ns_int size + cdef list indices + cdef list coordinates + cdef list distances + + def __init__(self, PBCBox box, real cutoff): + self.box = box + self.cutoff = cutoff + + self.size = 0 + self.nids = NULL + self.nsizes = NULL + + self.grid_coords = None + self.ref_coords = None + + self.indices = None + self.coordinates = None + self.distances = None + + + cdef populate(self, ns_neighborhood_holder *holder, grid_coords, ref_coords): + cdef ns_int nid, i + cdef ns_neighborhood *neighborhood + + self.grid_coords = grid_coords.copy() + self.ref_coords = ref_coords.copy() + + # Allocate memory + self.nsizes = malloc(sizeof(ns_int) * holder.size) + if self.nsizes == NULL: + raise MemoryError("Could not allocate memory for NSResults") + + self.nids = malloc(sizeof(ns_int *) * holder.size) + if self.nids == NULL: + raise MemoryError("Could not allocate memory for NSResults") + + for nid in range(holder.size): + neighborhood = holder.neighborhoods[nid] + + self.nsizes[nid] = neighborhood.size + + self.nids[nid] = malloc(sizeof(ns_int *) * neighborhood.size) + if self.nids[nid] == NULL: + raise MemoryError("Could not allocate memory for NSResults") + + with nogil: + for nid in range(holder.size): + neighborhood = holder.neighborhoods[nid] + + for i in range(neighborhood.size): + self.nids[nid][i] = neighborhood.beadids[i] + + self.size = holder.size + + def __dealloc__(self): + if self.nids != NULL: + for i in range(self.size): + if self.nids[i] != NULL: + free(self.nids[i]) + free(self.nids) + + if self.nsizes != NULL: + free(self.nsizes) + + + def get_indices(self): + """ + Return Neighbors indices. + + :return: list of indices + """ + cdef ns_int i, nid, size + + if self.indices is None: + indices = [] + + for nid in range(self.size): + size = self.nsizes[nid] + + tmp_incides = np.empty((size), dtype=np.int) + + for i in range(size): + tmp_incides[i] = self.nids[nid][i] + + indices.append(tmp_incides) + + self.indices = indices + + return self.indices + + + def get_coordinates(self): + """ + Return coordinates of neighbors. + + :return: list of coordinates + """ + cdef ns_int i, nid, size, beadid + + if self.coordinates is None: + coordinates = [] + + for nid in range(self.size): + size = self.nsizes[nid] + + tmp_values = np.empty((size, DIM), dtype=np.float32) + + for i in range(size): + beadid = self.nids[nid][i] + tmp_values[i] = self.grid_coords[beadid] + + coordinates.append(tmp_values) + + self.coordinates = coordinates + + return self.coordinates + + + def get_distances(self): + """ + Return coordinates of neighbors. + + :return: list of distances + """ + cdef ns_int i, nid, size, j, beadid + cdef rvec ref, other, dx + cdef real dist + + if self.distances is None: + distances = [] + + for nid in range(self.size): + size = self.nsizes[nid] + + tmp_values = np.empty((size), dtype=np.float32) + ref = &self.ref_coords[nid, 0] + + for i in range(size): + beadid = self.nids[nid][i] + other = &self.grid_coords[beadid, 0] + + tmp_values[i] = self.box.fast_distance(ref, other) + + distances.append(tmp_values) + + self.distances = distances + + return self.distances + + # Python interface cdef class FastNS(object): cdef PBCBox box @@ -448,7 +626,14 @@ cdef class FastNS(object): cdef bint prepared cdef ns_grid grid - def __init__(self, box): + def __init__(self, u): + import MDAnalysis as mda + from MDAnalysis.lib.mdamath import triclinic_vectors + + if not isinstance(u, mda.Universe): + raise TypeError("FastNS class must be initialized with a valid MDAnalysis.Universe instance") + box = triclinic_vectors(u.dimensions) + if box.shape != (3, 3): raise ValueError("Box must be provided as triclinic_dimensions (a 3x3 numpy.ndarray of unit cell vectors") @@ -466,6 +651,7 @@ cdef class FastNS(object): def __dealloc__(self): cdef ns_int i + # Deallocate NS grid if self.grid.nbeads != NULL: free(self.grid.nbeads) @@ -475,7 +661,8 @@ cdef class FastNS(object): free(self.grid.beadids) self.grid.size = 0 - + + def set_coords(self, real[:, ::1] coords): self.coords = coords @@ -568,12 +755,12 @@ cdef class FastNS(object): # Now we can free the allocation buffer free(allocated_size) - self.prepared = True - def search(self, real[:, ::1]search_coords, return_ids=False): + def search(self, search_coords): cdef real[:, ::1] search_coords_bbox + cdef real[:, ::1] search_coords_view cdef ns_int nid, i, j cdef ns_neighborhood_holder *holder cdef ns_neighborhood *neighborhood @@ -581,9 +768,18 @@ cdef class FastNS(object): if not self.prepared: self.prepare() + # Check the shape of search_coords as a array of 3D coords if needed + shape = search_coords.shape + if len(shape) == 1: + if not shape[0] == 3: + raise ValueError("Coordinates must be 3D") + else: + search_coords_view = np.array([search_coords,], dtype=np.float32) + else: + search_coords_view = search_coords # Make sure atoms are inside the brick-shaped box - search_coords_bbox = self.box.fast_put_atoms_in_bbox(search_coords) + search_coords_bbox = self.box.fast_put_atoms_in_bbox(search_coords_view) with nogil: holder = ns_core(search_coords_bbox, self.coords_bbox, &self.grid, self.box, self.cutoff) @@ -591,24 +787,11 @@ cdef class FastNS(object): if holder == NULL: raise MemoryError("Could not allocate memory to run NS core") - neighbors = [] - for nid in range(holder.size): - neighborhood = holder.neighborhoods[nid] - if return_ids: - neighborhood_py = np.empty(neighborhood.size, dtype=np.int64) - - for i in range(neighborhood.size): - neighborhood_py[i] = neighborhood.beadids[i] - else: - neighborhood_py = np.empty((neighborhood.size, DIM), dtype=np.float32) - for i in range(neighborhood.size): - for j in range(DIM): - neighborhood_py[i,j] = self.coords[neighborhood.beadids[i], j] - - neighbors.append(neighborhood_py) + results = NSResults(self.box, self.cutoff) + results.populate(holder, self.coords, search_coords_view) - # Free Memory + # Free memory allocated to holder free_neighborhood_holder(holder) - return neighbors + return results diff --git a/package/doc/sphinx/source/documentation_pages/lib/grid.rst b/package/doc/sphinx/source/documentation_pages/lib/grid.rst new file mode 100644 index 00000000000..a6143e9ca38 --- /dev/null +++ b/package/doc/sphinx/source/documentation_pages/lib/grid.rst @@ -0,0 +1,2 @@ +.. automodule:: MDAnalysis.lib.grid + :members: \ No newline at end of file diff --git a/package/doc/sphinx/source/documentation_pages/lib_modules.rst b/package/doc/sphinx/source/documentation_pages/lib_modules.rst index 5bbe3d041ed..83d7f81cfd9 100644 --- a/package/doc/sphinx/source/documentation_pages/lib_modules.rst +++ b/package/doc/sphinx/source/documentation_pages/lib_modules.rst @@ -36,6 +36,7 @@ functions whereas mathematical functions are to be found in :mod:`MDAnalysis.lib.NeighborSearch` contains classes to do neighbor searches with MDAnalysis objects. +:mod:`MDAnalysis.lib.grid` contains a fast implementation of grid neighbor search. List of modules --------------- @@ -50,6 +51,7 @@ List of modules ./lib/transformations ./lib/qcprot ./lib/util + ./lib/grid Low level file formats ---------------------- diff --git a/testsuite/MDAnalysisTests/lib/test_gridsearch.py b/testsuite/MDAnalysisTests/lib/test_gridsearch.py index c8d526e9074..359d32544b7 100644 --- a/testsuite/MDAnalysisTests/lib/test_gridsearch.py +++ b/testsuite/MDAnalysisTests/lib/test_gridsearch.py @@ -23,27 +23,46 @@ from __future__ import print_function, absolute_import import pytest -from numpy.testing import assert_equal +from numpy.testing import assert_equal, assert_allclose import numpy as np import MDAnalysis as mda from MDAnalysis.lib import grid from MDAnalysis.lib.pkdtree import PeriodicKDTree -from MDAnalysis.lib.mdamath import triclinic_vectors from MDAnalysisTests.datafiles import GRO + @pytest.fixture def universe(): u = mda.Universe(GRO) return u + +@pytest.fixture +def grid_results(): + u = mda.Universe(GRO) + cutoff = 2 + ref_pos = u.atoms.positions[13937] + return run_grid_search(u, ref_pos, cutoff) + + +def run_grid_search(u, ref_pos, cutoff): + coords = u.atoms.positions + + # Run grid search + searcher = grid.FastNS(u) + searcher.set_cutoff(cutoff) + searcher.set_coords(coords) + searcher.prepare() + + return searcher.search(ref_pos) + + def run_search(universe, ref_id): cutoff = 3 - coords = universe.atoms.positions ref_pos = coords[ref_id] - triclinic_box = triclinic_vectors(universe.dimensions) # Run pkdtree search pkdt = PeriodicKDTree(universe.atoms.dimensions, bucket_size=10) @@ -56,16 +75,13 @@ def run_search(universe, ref_id): results_pkdtree.sort() # Run grid search - searcher = grid.FastNS(triclinic_box) - searcher.set_cutoff(cutoff) - searcher.set_coords(coords) - searcher.prepare() - - results_grid = searcher.search(np.array([ref_pos, ]), return_ids=True)[0] + results_grid = run_grid_search(universe, ref_pos, cutoff) + results_grid = results_grid.get_indices()[0] results_grid.sort() return results_pkdtree, results_grid + def test_gridsearch(universe): """Check that pkdtree and grid search return the same results (No PBC needed)""" @@ -73,9 +89,52 @@ def test_gridsearch(universe): results_pkdtree, results_grid = run_search(universe, ref_id) assert_equal(results_pkdtree, results_grid) + def test_gridsearch_PBC(universe): """Check that pkdtree and grid search return the same results (PBC needed)""" ref_id = 13937 results_pkdtree, results_grid = run_search(universe, ref_id) assert_equal(results_pkdtree, results_grid) + + +def test_gridsearch_arraycoord(universe): + """Check the NS routine accepts a single bead coordinate as well as array of coordinates""" + cutoff = 2 + ref_pos = universe.atoms.positions[:5] + + results = [ + np.array([2, 1, 4, 3]), + np.array([2, 0, 3]), + np.array([0, 1, 3]), + np.array([ 2, 0, 1, 38341]), + np.array([ 6, 0, 5, 17]) + ] + + results_grid = run_grid_search(universe, ref_pos, cutoff).get_indices() + + assert_equal(results_grid, results) + + +def test_gridsearch_search_coordinates(grid_results): + """Check the NS routine can return coordinates instead of ids""" + + results = np.array( + [ + [40.32, 34.25, 55.9], + [0.61, 76.33, -0.56], + [0.48999998, 75.9, 0.19999999], + [-0.11, 76.19, 0.77] + ]) + + assert_allclose(grid_results.get_coordinates()[0], results) + + +def test_gridsearch_search_distances(grid_results): + """Check the NS routine can return PBC distances from neighbors""" + results = np.array([0.096, 0.096, 0.015, 0.179]) * 10 # These distances were obtained using gmx distance + results.sort() + + rounded_results = np.round(grid_results.get_distances()[0], 2) + + assert_allclose(sorted(rounded_results), results) \ No newline at end of file From 579fa0088984fb4720e9b169eecf56184c6faefd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Buchoux?= Date: Sun, 8 Jul 2018 20:10:33 +0200 Subject: [PATCH 11/47] Corrected Memory Leak --- package/MDAnalysis/lib/c_gridsearch.pyx | 36 ++++++++++++------------- 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/package/MDAnalysis/lib/c_gridsearch.pyx b/package/MDAnalysis/lib/c_gridsearch.pyx index 9c50a5b1777..1ece4c5a1d6 100644 --- a/package/MDAnalysis/lib/c_gridsearch.pyx +++ b/package/MDAnalysis/lib/c_gridsearch.pyx @@ -462,8 +462,8 @@ cdef class NSResults(object): """ cdef PBCBox box cdef readonly real cutoff - cdef real[:, ::1] grid_coords - cdef real[:, ::1] ref_coords + cdef np.ndarray grid_coords + cdef np.ndarray ref_coords cdef ns_int **nids cdef ns_int *nsizes cdef ns_int size @@ -471,6 +471,16 @@ cdef class NSResults(object): cdef list coordinates cdef list distances + def __dealloc__(self): + if self.nids != NULL: + for i in range(self.size): + if self.nids[i] != NULL: + free(self.nids[i]) + free(self.nids) + + if self.nsizes != NULL: + free(self.nsizes) + def __init__(self, PBCBox box, real cutoff): self.box = box self.cutoff = cutoff @@ -491,8 +501,8 @@ cdef class NSResults(object): cdef ns_int nid, i cdef ns_neighborhood *neighborhood - self.grid_coords = grid_coords.copy() - self.ref_coords = ref_coords.copy() + self.grid_coords = np.asarray(grid_coords) + self.ref_coords = np.asarray(ref_coords) # Allocate memory self.nsizes = malloc(sizeof(ns_int) * holder.size) @@ -521,17 +531,6 @@ cdef class NSResults(object): self.size = holder.size - def __dealloc__(self): - if self.nids != NULL: - for i in range(self.size): - if self.nids[i] != NULL: - free(self.nids[i]) - free(self.nids) - - if self.nsizes != NULL: - free(self.nsizes) - - def get_indices(self): """ Return Neighbors indices. @@ -594,6 +593,8 @@ cdef class NSResults(object): cdef ns_int i, nid, size, j, beadid cdef rvec ref, other, dx cdef real dist + cdef real[:, ::1] ref_coords = self.ref_coords + cdef real[:, ::1] grid_coords = self.grid_coords if self.distances is None: distances = [] @@ -602,11 +603,11 @@ cdef class NSResults(object): size = self.nsizes[nid] tmp_values = np.empty((size), dtype=np.float32) - ref = &self.ref_coords[nid, 0] + ref = &ref_coords[nid, 0] for i in range(size): beadid = self.nids[nid][i] - other = &self.grid_coords[beadid, 0] + other = &grid_coords[beadid, 0] tmp_values[i] = self.box.fast_distance(ref, other) @@ -787,7 +788,6 @@ cdef class FastNS(object): if holder == NULL: raise MemoryError("Could not allocate memory to run NS core") - results = NSResults(self.box, self.cutoff) results.populate(holder, self.coords, search_coords_view) From 0239d14bfdbe1cbb0f9b678b9dbb261a222c8c7e Mon Sep 17 00:00:00 2001 From: ayush Date: Mon, 25 Jun 2018 20:10:22 -0700 Subject: [PATCH 12/47] Added MDAnalysis/lib/c_search.pyx for Grid NS search --- package/MDAnalysis/lib/c_gridsearch.pyx | 797 ------------------------ 1 file changed, 797 deletions(-) diff --git a/package/MDAnalysis/lib/c_gridsearch.pyx b/package/MDAnalysis/lib/c_gridsearch.pyx index 1ece4c5a1d6..e69de29bb2d 100644 --- a/package/MDAnalysis/lib/c_gridsearch.pyx +++ b/package/MDAnalysis/lib/c_gridsearch.pyx @@ -1,797 +0,0 @@ -# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; -*- -# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 -# -# MDAnalysis --- https://www.mdanalysis.org -# -# Copyright (C) 2013-2018 Sébastien Buchoux -# Copyright (c) 2018 The MDAnalysis Development Team and contributors -# (see the file AUTHORS for the full list of names) -# -# Released under the GNU Public Licence, v3 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 -# -# - -#cython: cdivision=True -#cython: boundscheck=False -#cython: initializedcheck=False - -""" -Neighbor search library --- :mod:`MDAnalysis.lib.grid` -====================================================== - -This Neighbor search library is a serialized Cython port of the NS grid search implemented in GROMACS. -""" - - -# Preprocessor DEFs -DEF DIM = 3 -DEF XX = 0 -DEF YY = 1 -DEF ZZ = 2 - -DEF EPSILON = 1e-5 - -DEF NEIGHBORHOOD_ALLOCATION_INCREMENT = 50 - -DEF BOX_MARGIN=1.0010 -DEF MAX_NTRICVEC=12 - -from libc.stdlib cimport malloc, realloc, free -from libc.math cimport sqrt - -import numpy as np -cimport numpy as np - -ctypedef np.int_t ns_int -ctypedef np.float32_t real -ctypedef real rvec[DIM] -ctypedef real matrix[DIM][DIM] - -cdef struct ns_grid: - ns_int size - ns_int[DIM] ncells - real[DIM] cellsize - ns_int *nbeads - ns_int **beadids - -cdef struct ns_neighborhood: - real cutoff - ns_int allocated_size - ns_int size - ns_int *beadids - -cdef struct ns_neighborhood_holder: - ns_int size - ns_neighborhood **neighborhoods - -# Useful stuff - -cdef real rvec_norm2(const rvec a) nogil: - return a[XX]*a[XX]+a[YY]*a[YY]+a[ZZ]*a[ZZ] - -cdef void rvec_clear(rvec a) nogil: - a[XX]=0.0 - a[YY]=0.0 - a[ZZ]=0.0 - - -cdef struct cPBCBox_t: - matrix box - rvec fbox_diag - rvec hbox_diag - rvec mhbox_diag - real max_cutoff2 - ns_int ntric_vec - ns_int[DIM] tric_shift[MAX_NTRICVEC] - real[DIM] tric_vec[MAX_NTRICVEC] - - -# Class to handle PBC calculations -cdef class PBCBox(object): - cdef cPBCBox_t c_pbcbox - cdef rvec center - cdef rvec bbox_center - - def __init__(self, real[:,::1] box): - self.update(box) - - cdef void fast_update(self, real[:,::1] box) nogil: - cdef ns_int i, j, k, d, jc, kc, shift - cdef real d2old, d2new, d2new_c - cdef rvec trial, pos - cdef ns_int ii, jj ,kk - cdef ns_int *order = [0, -1, 1, -2, 2] - cdef bint use - cdef real min_hv2, min_ss, tmp - - rvec_clear(self.center) - # Update matrix - for i in range(DIM): - for j in range(DIM): - self.c_pbcbox.box[i][j] = box[i, j] - self.center[j] += 0.5 * box[i, j] - self.bbox_center[i] = 0.5 * box[i, i] - - # Update diagonals - for i in range(DIM): - self.c_pbcbox.fbox_diag[i] = box[i, i] - self.c_pbcbox.hbox_diag[i] = self.c_pbcbox.fbox_diag[i] * 0.5 - self.c_pbcbox.mhbox_diag[i] = - self.c_pbcbox.hbox_diag[i] - - # Update maximum cutoff - - # Physical limitation of the cut-off - # by half the length of the shortest box vector. - min_hv2 = min(0.25 * rvec_norm2(&box[XX, XX]), 0.25 * rvec_norm2(&box[YY, XX])) - min_hv2 = min(min_hv2, 0.25 * rvec_norm2(&box[ZZ, XX])) - - # Limitation to the smallest diagonal element due to optimizations: - # checking only linear combinations of single box-vectors (2 in x) - # in the grid search and pbc_dx is a lot faster - # than checking all possible combinations. - tmp = box[YY, YY] - if box[ZZ, YY] < 0: - tmp -= box[ZZ, YY] - else: - tmp += box[ZZ, YY] - - min_ss = min(box[XX, XX], min(tmp, box[ZZ, ZZ])) - - self.c_pbcbox.max_cutoff2 = min(min_hv2, min_ss * min_ss) - - # Update shift vectors - self.c_pbcbox.ntric_vec = 0 - - # We will only use single shifts - for kk in range(3): - k = order[kk] - - for jj in range(3): - j = order[jj] - - for ii in range(3): - i = order[ii] - - # A shift is only useful when it is trilinic - if j != 0 or k != 0: - d2old = 0 - d2new = 0 - - for d in range(DIM): - trial[d] = i*box[XX, d] + j*box[YY, d] + k*box[ZZ, d] - - # Choose the vector within the brick around 0,0,0 that - # will become the shortest due to shift try. - - if d == DIM: - trial[d] = 0 - pos[d] = 0 - else: - if trial[d] < 0: - pos[d] = min(self.c_pbcbox.hbox_diag[d], -trial[d]) - else: - pos[d] = max(-self.c_pbcbox.hbox_diag[d], -trial[d]) - - d2old += pos[d]**2 - d2new += (pos[d] + trial[d])**2 - - if BOX_MARGIN*d2new < d2old: - if not (j < -1 or j > 1 or k < -1 or k > 1): - use = True - - for dd in range(DIM): - if dd == 0: - shift = i - elif dd == 1: - shift = j - else: - shift = k - - if shift: - d2new_c = 0 - - for d in range(DIM): - d2new_c += (pos[d] + trial[d] - shift*box[dd, d])**2 - - if d2new_c <= BOX_MARGIN*d2new: - use = False - - if use: # Accept this shift vector. - if self.c_pbcbox.ntric_vec >= MAX_NTRICVEC: - with gil: - print("\nWARNING: Found more than %d triclinic " - "correction vectors, ignoring some." - % MAX_NTRICVEC) - print(" There is probably something wrong with " - "your box.") - print(np.array(box)) - - for i in range(self.c_pbcbox.ntric_vec): - print(" -> shift #{}: [{}, {}, {}]".format(i+1, - self.c_pbcbox.tric_shift[i][XX], - self.c_pbcbox.tric_shift[i][YY], - self.c_pbcbox.tric_shift[i][ZZ])) - else: - for d in range(DIM): - self.c_pbcbox.tric_vec[self.c_pbcbox.ntric_vec][d] = \ - trial[d] - self.c_pbcbox.tric_shift[self.c_pbcbox.ntric_vec][XX] = i - self.c_pbcbox.tric_shift[self.c_pbcbox.ntric_vec][YY] = j - self.c_pbcbox.tric_shift[self.c_pbcbox.ntric_vec][ZZ] = k - self.c_pbcbox.ntric_vec += 1 - - - def update(self, real[:,::1] box): - if box.shape[0] != DIM or box.shape[1] != DIM: - raise ValueError("Box must be a %i x %i matrix. (shape: %i x %i)" % - (DIM, DIM, box.shape[0], box.shape[1])) - if (box[XX, XX] == 0) or (box[YY, YY] == 0) or (box[ZZ, ZZ] == 0): - raise ValueError("Box does not correspond to PBC=xyz") - self.fast_update(box) - - - cdef void fast_pbc_dx(self, rvec ref, rvec other, rvec dx) nogil: - cdef ns_int i, j - cdef rvec dx_start, trial - - for i in range(DIM): - dx[i] = other[i] - ref[i] - - for i in range (DIM-1, -1, -1): - while dx[i] > self.c_pbcbox.hbox_diag[i]: - for j in range (i, -1, -1): - dx[j] -= self.c_pbcbox.box[i][j] - - while dx[i] <= self.c_pbcbox.mhbox_diag[i]: - for j in range (i, -1, -1): - dx[j] += self.c_pbcbox.box[i][j] - - cdef real fast_distance2(self, rvec a, rvec b) nogil: - cdef rvec dx - self.fast_pbc_dx(a, b, dx) - return rvec_norm2(dx) - - cdef real fast_distance(self, rvec a, rvec b) nogil: - return sqrt(self.fast_distance2(a,b)) - - cdef real[:, ::1]fast_put_atoms_in_bbox(self, real[:,::1] coords) nogil: - cdef ns_int i, m, d, natoms, wd = 0 - cdef real[:,::1] bbox_coords - - natoms = coords.shape[0] - with gil: - bbox_coords = coords.copy() - - for i in range(natoms): - for m in range(DIM - 1, -1, -1): - while bbox_coords[i, m] < 0: - for d in range(m+1): - bbox_coords[i, d] += self.c_pbcbox.box[m][d] - while bbox_coords[i, m] >= self.c_pbcbox.box[m][m]: - for d in range(m+1): - bbox_coords[i, d] -= self.c_pbcbox.box[m][d] - return bbox_coords - - def put_atoms_in_bbox(self, real[:,::1] coords): - if coords.shape[0] == 0: - return np.zeros((0, DIM), dtype=np.float32) - return np.asarray(self.fast_put_atoms_in_bbox(coords)) - -######################################################################################################################## -# -# Neighbor Search Stuff -# -######################################################################################################################## -cdef struct ns_grid: - ns_int size - ns_int[DIM] ncells - real[DIM] cellsize - ns_int *nbeads - ns_int **beadids - -cdef ns_neighborhood_holder *create_neighborhood_holder() nogil: - cdef ns_neighborhood_holder *holder - - holder = malloc(sizeof(ns_neighborhood_holder)) - holder.size = 0 - holder.neighborhoods = NULL - - return holder - -cdef void free_neighborhood_holder(ns_neighborhood_holder *holder) nogil: - cdef ns_int i - - if holder == NULL: - return - - for i in range(holder.size): - if holder.neighborhoods[i].beadids != NULL: - free(holder.neighborhoods[i].beadids) - free(holder.neighborhoods[i]) - - if holder.neighborhoods != NULL: - free(holder.neighborhoods) - free(holder) - -cdef ns_neighborhood *retrieve_neighborhood(rvec current_coords, real[:, ::1]neighborcoords, ns_grid *grid, PBCBox box, real cutoff2) nogil: - cdef ns_int d, m - cdef ns_int xi, yi, zi, bid - cdef real d2 - cdef rvec shifted_coord, dx, neighbor_coord, corrected_coords - - cdef bint already_checked[27] - cdef bint skip - cdef ns_int nchecked = 0, icheck - cdef ns_int cell_index - - cdef ns_neighborhood *neighborhood = malloc(sizeof(ns_neighborhood)) - if neighborhood == NULL: - return NULL - - neighborhood.size = 0 - neighborhood.allocated_size = NEIGHBORHOOD_ALLOCATION_INCREMENT - neighborhood.beadids = malloc(NEIGHBORHOOD_ALLOCATION_INCREMENT * sizeof(ns_int)) - - if neighborhood.beadids == NULL: - free(neighborhood) - return NULL - - for zi in range(3): - for yi in range(3): - for xi in range(3): - # Calculate and/or reinitialize shifted coordinates - shifted_coord[XX] = current_coords[XX] + (xi - 1) * grid.cellsize[XX] - shifted_coord[YY] = current_coords[YY] + (yi - 1) * grid.cellsize[YY] - shifted_coord[ZZ] = current_coords[ZZ] + (zi - 1) * grid.cellsize[ZZ] - - # Make sure the shifted coordinates is inside the brick-shaped box - for m in range(DIM - 1, -1, -1): - - while shifted_coord[m] < 0: - for d in range(m+1): - shifted_coord[d] += box.c_pbcbox.box[m][d] - - - while shifted_coord[m] >= box.c_pbcbox.box[m][m]: - for d in range(m+1): - shifted_coord[d] -= box.c_pbcbox.box[m][d] - - # Get the cell index corresponding to the coord - cell_index = (shifted_coord[ZZ] / grid.cellsize[ZZ]) * (grid.ncells[XX] * grid.ncells[YY]) +\ - (shifted_coord[YY] / grid.cellsize[YY]) * grid.ncells[XX] + \ - (shifted_coord[XX] / grid.cellsize[XX]) - - # Just a safeguard - if cell_index >= grid.size: - continue - - # Check the cell index was not already selected - skip = False - for icheck in range(nchecked): - if already_checked[icheck] == cell_index: - skip = True - break - if skip: - continue - - # Search for neighbors inside this cell - for i_bead in range(grid.nbeads[cell_index]): - bid = grid.beadids[cell_index][i_bead] - - box.fast_pbc_dx(current_coords, &neighborcoords[bid, XX], dx) - - d2 = rvec_norm2(dx) - - if d2 < cutoff2: - if d2 < EPSILON: # Don't add the current bead as its own neighbor! - continue - - # Update neighbor lists - neighborhood.beadids[neighborhood.size] = bid - neighborhood.size += 1 - - if neighborhood.size >= neighborhood.allocated_size: - neighborhood.allocated_size += NEIGHBORHOOD_ALLOCATION_INCREMENT - neighborhood.beadids = realloc( neighborhood.beadids, neighborhood.allocated_size * sizeof(ns_int)) - - if neighborhood.beadids == NULL: - free(neighborhood) - return NULL - - # Register the cell as checked - already_checked[nchecked] = cell_index - nchecked += 1 - - return neighborhood - - -cdef ns_neighborhood_holder *ns_core(real[:, ::1] refcoords, - real[:, ::1] neighborcoords, - ns_grid *grid, - PBCBox box, - real cutoff) nogil: - cdef ns_int coordid, i, j - cdef ns_int ncoords = refcoords.shape[0] - cdef ns_int ncoords_neighbors = neighborcoords.shape[0] - cdef real cutoff2 = cutoff * cutoff - cdef ns_neighborhood_holder *holder - - cdef ns_int *neighbor_buf - cdef ns_int buf_size, ibuf - - holder = create_neighborhood_holder() - if holder == NULL: - return NULL - - holder.neighborhoods = malloc(sizeof(ns_neighborhood *) * ncoords) - if holder.neighborhoods == NULL: - free_neighborhood_holder(holder) - return NULL - - # Here starts the real core and the iteration over coordinates - for coordid in range(ncoords): - holder.neighborhoods[coordid] = retrieve_neighborhood(&refcoords[coordid, XX], - neighborcoords, - grid, - box, - cutoff2) - if holder.neighborhoods[coordid] == NULL: - free_neighborhood_holder(holder) - return NULL - - holder.neighborhoods[coordid].cutoff = cutoff - holder.size += 1 - - return holder - -cdef class NSResults(object): - """ - Class used to store results returned by `MDAnalysis.lib.grid.FastNS.search` - """ - cdef PBCBox box - cdef readonly real cutoff - cdef np.ndarray grid_coords - cdef np.ndarray ref_coords - cdef ns_int **nids - cdef ns_int *nsizes - cdef ns_int size - cdef list indices - cdef list coordinates - cdef list distances - - def __dealloc__(self): - if self.nids != NULL: - for i in range(self.size): - if self.nids[i] != NULL: - free(self.nids[i]) - free(self.nids) - - if self.nsizes != NULL: - free(self.nsizes) - - def __init__(self, PBCBox box, real cutoff): - self.box = box - self.cutoff = cutoff - - self.size = 0 - self.nids = NULL - self.nsizes = NULL - - self.grid_coords = None - self.ref_coords = None - - self.indices = None - self.coordinates = None - self.distances = None - - - cdef populate(self, ns_neighborhood_holder *holder, grid_coords, ref_coords): - cdef ns_int nid, i - cdef ns_neighborhood *neighborhood - - self.grid_coords = np.asarray(grid_coords) - self.ref_coords = np.asarray(ref_coords) - - # Allocate memory - self.nsizes = malloc(sizeof(ns_int) * holder.size) - if self.nsizes == NULL: - raise MemoryError("Could not allocate memory for NSResults") - - self.nids = malloc(sizeof(ns_int *) * holder.size) - if self.nids == NULL: - raise MemoryError("Could not allocate memory for NSResults") - - for nid in range(holder.size): - neighborhood = holder.neighborhoods[nid] - - self.nsizes[nid] = neighborhood.size - - self.nids[nid] = malloc(sizeof(ns_int *) * neighborhood.size) - if self.nids[nid] == NULL: - raise MemoryError("Could not allocate memory for NSResults") - - with nogil: - for nid in range(holder.size): - neighborhood = holder.neighborhoods[nid] - - for i in range(neighborhood.size): - self.nids[nid][i] = neighborhood.beadids[i] - - self.size = holder.size - - def get_indices(self): - """ - Return Neighbors indices. - - :return: list of indices - """ - cdef ns_int i, nid, size - - if self.indices is None: - indices = [] - - for nid in range(self.size): - size = self.nsizes[nid] - - tmp_incides = np.empty((size), dtype=np.int) - - for i in range(size): - tmp_incides[i] = self.nids[nid][i] - - indices.append(tmp_incides) - - self.indices = indices - - return self.indices - - - def get_coordinates(self): - """ - Return coordinates of neighbors. - - :return: list of coordinates - """ - cdef ns_int i, nid, size, beadid - - if self.coordinates is None: - coordinates = [] - - for nid in range(self.size): - size = self.nsizes[nid] - - tmp_values = np.empty((size, DIM), dtype=np.float32) - - for i in range(size): - beadid = self.nids[nid][i] - tmp_values[i] = self.grid_coords[beadid] - - coordinates.append(tmp_values) - - self.coordinates = coordinates - - return self.coordinates - - - def get_distances(self): - """ - Return coordinates of neighbors. - - :return: list of distances - """ - cdef ns_int i, nid, size, j, beadid - cdef rvec ref, other, dx - cdef real dist - cdef real[:, ::1] ref_coords = self.ref_coords - cdef real[:, ::1] grid_coords = self.grid_coords - - if self.distances is None: - distances = [] - - for nid in range(self.size): - size = self.nsizes[nid] - - tmp_values = np.empty((size), dtype=np.float32) - ref = &ref_coords[nid, 0] - - for i in range(size): - beadid = self.nids[nid][i] - other = &grid_coords[beadid, 0] - - tmp_values[i] = self.box.fast_distance(ref, other) - - distances.append(tmp_values) - - self.distances = distances - - return self.distances - - -# Python interface -cdef class FastNS(object): - cdef PBCBox box - cdef readonly real[:, ::1] coords - cdef real[:, ::1] coords_bbox - cdef readonly real cutoff - cdef bint prepared - cdef ns_grid grid - - def __init__(self, u): - import MDAnalysis as mda - from MDAnalysis.lib.mdamath import triclinic_vectors - - if not isinstance(u, mda.Universe): - raise TypeError("FastNS class must be initialized with a valid MDAnalysis.Universe instance") - box = triclinic_vectors(u.dimensions) - - if box.shape != (3, 3): - raise ValueError("Box must be provided as triclinic_dimensions (a 3x3 numpy.ndarray of unit cell vectors") - - self.box = PBCBox(box) - - self.coords = None - self.coords_bbox = None - - self.cutoff = -1 - - self.prepared = False - - self.grid.size = 0 - - - def __dealloc__(self): - cdef ns_int i - # Deallocate NS grid - if self.grid.nbeads != NULL: - free(self.grid.nbeads) - - for i in range(self.grid.size): - if self.grid.beadids[i] != NULL: - free(self.grid.beadids[i]) - free(self.grid.beadids) - - self.grid.size = 0 - - - def set_coords(self, real[:, ::1] coords): - self.coords = coords - - # Make sure atoms are inside the brick-shaped box - self.coords_bbox = self.box.fast_put_atoms_in_bbox(coords) - - self.prepared = False - - - def set_cutoff(self, real cutoff): - self.cutoff = cutoff - - self.prepared = False - - - def prepare(self, force=False): - cdef ns_int i, cellindex = -1 - cdef ns_int *allocated_size = NULL - cdef ns_int ncoords = self.coords.shape[0] - cdef ns_int allocation_guess - cdef rvec *coords = &self.coords_bbox[0, 0] - - if self.prepared and not force: - print("NS already prepared, nothing to do!") - - if self.coords is None: - raise ValueError("Coordinates must be set before NS preparation!") - - if self.cutoff < 0: - raise ValueError("Cutoff must be set before NS preparation!") - - with nogil: - - # Initializing grid - for i in range(DIM): - self.grid.ncells[i] = (self.box.c_pbcbox.box[i][i] / self.cutoff) - if self.grid.ncells[i] == 0: - self.grid.ncells[i] = 1 - self.grid.cellsize[i] = self.box.c_pbcbox.box[i][i] / self.grid.ncells[i] - self.grid.size = self.grid.ncells[XX] * self.grid.ncells[YY] * self.grid.ncells[ZZ] - - # This is just a guess on how much memory we might for each grid cell: - # we just assume an average bead density and we take four times this density just to be safe - allocation_guess = (4 * (ncoords / self.grid.size + 1)) - - # Allocate memory for the grid - self.grid.nbeads = malloc(sizeof(ns_int) * self.grid.size) - if self.grid.nbeads == NULL: - with gil: - raise MemoryError("Could not allocate memory for NS grid") - - # Allocate memory from temporary allocation counter - allocated_size = malloc(sizeof(ns_int) * self.grid.size) - if allocated_size == NULL: - # No need to free grid.nbeads as it will be freed by destroy_nsgrid called by __dealloc___ - with gil: - raise MemoryError("Could not allocate memory for allocation buffer") - - # Pre-allocate some memory for grid cells - for i in range(self.grid.size): - self.grid.nbeads[i] = 0 - allocated_size[i] = allocation_guess - - self.grid.beadids = malloc(sizeof(ns_int *) * self.grid.size) - if self.grid.beadids == NULL: - with gil: - raise MemoryError("Could not allocate memory for grid cells") - - for i in range(self.grid.size): - self.grid.beadids[i] = malloc(sizeof(ns_int) * allocated_size[i]) - if self.grid.beadids[i] == NULL: - with gil: - raise MemoryError("Could not allocate memory for grid cell") - - # Populate grid cells using the coordinates (ie do the heavy work) - for i in range(ncoords): - cellindex = (coords[i][ZZ] / self.grid.cellsize[ZZ]) * (self.grid.ncells[XX] * self.grid.ncells[YY]) +\ - (coords[i][YY] / self.grid.cellsize[YY]) * self.grid.ncells[XX] + \ - (coords[i][XX] / self.grid.cellsize[XX]) - - self.grid.beadids[cellindex][self.grid.nbeads[cellindex]] = i - self.grid.nbeads[cellindex] += 1 - - # We need to allocate more memory (simply double the amount of memory as - # 1. it should barely be needed - # 2. the size should stay fairly reasonable - if self.grid.nbeads[cellindex] >= allocated_size[cellindex]: - allocated_size[cellindex] *= 2 - self.grid.beadids[cellindex] = realloc( self.grid.beadids[cellindex], sizeof(ns_int) * allocated_size[cellindex]) - - # Now we can free the allocation buffer - free(allocated_size) - self.prepared = True - - - def search(self, search_coords): - cdef real[:, ::1] search_coords_bbox - cdef real[:, ::1] search_coords_view - cdef ns_int nid, i, j - cdef ns_neighborhood_holder *holder - cdef ns_neighborhood *neighborhood - - if not self.prepared: - self.prepare() - - # Check the shape of search_coords as a array of 3D coords if needed - shape = search_coords.shape - if len(shape) == 1: - if not shape[0] == 3: - raise ValueError("Coordinates must be 3D") - else: - search_coords_view = np.array([search_coords,], dtype=np.float32) - else: - search_coords_view = search_coords - - # Make sure atoms are inside the brick-shaped box - search_coords_bbox = self.box.fast_put_atoms_in_bbox(search_coords_view) - - with nogil: - holder = ns_core(search_coords_bbox, self.coords_bbox, &self.grid, self.box, self.cutoff) - - if holder == NULL: - raise MemoryError("Could not allocate memory to run NS core") - - results = NSResults(self.box, self.cutoff) - results.populate(holder, self.coords, search_coords_view) - - # Free memory allocated to holder - free_neighborhood_holder(holder) - - return results From 5e9fe318f42893c2b3ab5592d44a062bf24880e3 Mon Sep 17 00:00:00 2001 From: ayush Date: Mon, 25 Jun 2018 22:06:44 -0700 Subject: [PATCH 13/47] removed the parallel handler to remove hard dependency on omp.h, can be added again for conditional compilation --- package/MDAnalysis/lib/c_gridsearch.pyx | 666 ++++++++++++++++++++++++ 1 file changed, 666 insertions(+) diff --git a/package/MDAnalysis/lib/c_gridsearch.pyx b/package/MDAnalysis/lib/c_gridsearch.pyx index e69de29bb2d..a4e8745aec4 100644 --- a/package/MDAnalysis/lib/c_gridsearch.pyx +++ b/package/MDAnalysis/lib/c_gridsearch.pyx @@ -0,0 +1,666 @@ +# -*- coding: utf-8; Mode: python; tab-width: 4; indent-tabs-mode:nil; -*- +# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 +# +# Copyright (C) 2013-2016 Sébastien Buchoux +# +# This file is part of FATSLiM. +# +# FATSLiM is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# FATSLiM is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with FATSLiM. If not, see . +#cython: cdivision=True +#cython: boundscheck=False + +# Preprocessor DEFs +DEF DIM = 3 +DEF XX = 0 +DEF YY = 1 +DEF ZZ = 2 +DEF RET_OK = 1 +DEF RET_ERROR = 0 +DEF EPSILON = 1e-5 +DEF NEIGHBORHOOD_ALLOCATION_INCREMENT = 50 +DEF GRID_ALLOCATION_INCREMENT = 50 + +DEF BOX_MARGIN=1.0010 +DEF MAX_NTRICVEC=12 + +from libc.stdlib cimport malloc, realloc, free, abort +from libc.stdio cimport fprintf, stderr +from libc.math cimport sqrt +from libc.math cimport abs as real_abs + +import numpy as np +cimport numpy as np + +ctypedef np.int_t ns_int +ctypedef np.float32_t real +ctypedef real rvec[DIM] +ctypedef real matrix[DIM][DIM] + +cdef struct ns_grid: + ns_int size + ns_int[DIM] ncells + real[DIM] cellsize + ns_int *nbeads + ns_int **beadids + +cdef struct ns_neighborhood: + real cutoff + ns_int allocated_size + ns_int size + ns_int *beadids + real *beaddist + ### +cdef struct ns_neighborhood_holder: + ns_int size + ns_neighborhood **neighborhoods + +# Useful stuff + +cdef real rvec_norm2(const rvec a) nogil: + return a[XX]*a[XX]+a[YY]*a[YY]+a[ZZ]*a[ZZ] + +cdef void rvec_clear(rvec a) nogil: + a[XX]=0.0 + a[YY]=0.0 + a[ZZ]=0.0 + +cdef struct cPBCBox_t: + matrix box + rvec fbox_diag + rvec hbox_diag + rvec mhbox_diag + real max_cutoff2 + ns_int ntric_vec + ns_int[DIM] tric_shift[MAX_NTRICVEC] + real[DIM] tric_vec[MAX_NTRICVEC] + +cdef class PBCBox(object): + cdef cPBCBox_t c_pbcbox + cdef rvec center + cdef rvec bbox_center + + def __init__(self, real[:,::1] box): + self.update(box) + + cdef void fast_update(self, real[:,::1] box) nogil: + cdef ns_int i, j, k, d, jc, kc, shift + cdef real d2old, d2new, d2new_c + cdef rvec trial, pos + cdef ns_int ii, jj ,kk + cdef ns_int *order = [0, -1, 1, -2, 2] + cdef bint use + cdef real min_hv2, min_ss, tmp + + rvec_clear(self.center) + # Update matrix + for i in range(DIM): + for j in range(DIM): + self.c_pbcbox.box[i][j] = box[i, j] + self.center[j] += 0.5 * box[i, j] + self.bbox_center[i] = 0.5 * box[i, i] + + # Update diagonals + for i in range(DIM): + self.c_pbcbox.fbox_diag[i] = box[i, i] + self.c_pbcbox.hbox_diag[i] = self.c_pbcbox.fbox_diag[i] * 0.5 + self.c_pbcbox.mhbox_diag[i] = - self.c_pbcbox.hbox_diag[i] + + # Update maximum cutoff + + # Physical limitation of the cut-off + # by half the length of the shortest box vector. + min_hv2 = min(0.25 * rvec_norm2(&box[XX, XX]), 0.25 * rvec_norm2(&box[YY, XX])) + min_hv2 = min(min_hv2, 0.25 * rvec_norm2(&box[ZZ, XX])) + + # Limitation to the smallest diagonal element due to optimizations: + # checking only linear combinations of single box-vectors (2 in x) + # in the grid search and pbc_dx is a lot faster + # than checking all possible combinations. + tmp = box[YY, YY] + if box[ZZ, YY] < 0: + tmp -= box[ZZ, YY] + else: + tmp += box[ZZ, YY] + + min_ss = min(box[XX, XX], min(tmp, box[ZZ, ZZ])) + + self.c_pbcbox.max_cutoff2 = min(min_hv2, min_ss * min_ss) + + # Update shift vectors + self.c_pbcbox.ntric_vec = 0 + # We will only use single shifts, but we will check a few + # more shifts to see if there is a limiting distance + # above which we can not be sure of the correct distance. + for kk in range(5): + k = order[kk] + + for jj in range(5): + j = order[jj] + + for ii in range(5): + i = order[ii] + + # A shift is only useful when it is trilinic + if j != 0 or k != 0: + d2old = 0 + d2new = 0 + + for d in range(DIM): + trial[d] = i*box[XX, d] + j*box[YY, d] + k*box[ZZ, d] + + # Choose the vector within the brick around 0,0,0 that + # will become the shortest due to shift try. + + if d == DIM: + trial[d] = 0 + pos[d] = 0 + else: + if trial[d] < 0: + pos[d] = min(self.c_pbcbox.hbox_diag[d], -trial[d]) + else: + pos[d] = max(-self.c_pbcbox.hbox_diag[d], -trial[d]) + + d2old += sqrt(pos[d]) + d2new += sqrt(pos[d] + trial[d]) + + if BOX_MARGIN*d2new < d2old: + if not (j < -1 or j > 1 or k < -1 or k > 1): + use = True + + for dd in range(DIM): + if dd == 0: + shift = i + elif dd == 1: + shift = j + else: + shift = k + + if shift: + d2new_c = 0 + + for d in range(DIM): + d2new_c += sqrt(pos[d] + trial[d] - shift*box[dd, d]) + + if d2new_c <= BOX_MARGIN*d2new: + use = False + + if use: # Accept this shift vector. + if self.c_pbcbox.ntric_vec >= MAX_NTRICVEC: + with gil: + print("\nWARNING: Found more than %d triclinic " + "correction vectors, ignoring some." + % MAX_NTRICVEC) + print(" There is probably something wrong with " + "your box.") + print(box) + else: + for d in range(DIM): + self.c_pbcbox.tric_vec[self.c_pbcbox.ntric_vec][d] = \ + trial[d] + self.c_pbcbox.tric_shift[self.c_pbcbox.ntric_vec][XX] = i + self.c_pbcbox.tric_shift[self.c_pbcbox.ntric_vec][YY] = j + self.c_pbcbox.tric_shift[self.c_pbcbox.ntric_vec][ZZ] = k + self.c_pbcbox.ntric_vec += 1 + + + def update(self, real[:,::1] box): + if box.shape[0] != DIM or box.shape[1] != DIM: + raise ValueError("Box must be a %i x %i matrix. (shape: %i x %i)" % + (DIM, DIM, box.shape[0], box.shape[1])) + if (box[XX, XX] == 0) or (box[YY, YY] == 0) or (box[ZZ, ZZ] == 0): + raise ValueError("Box does not correspond to PBC=xyz") + self.fast_update(box) + + cdef void fast_pbc_dx(self, rvec ref, rvec other, rvec dx) nogil: + cdef ns_int i, j + cdef rvec dx_start, trial + + for i in range(DIM): + dx[i] = other[i] - ref[i] + + for i in range (DIM-1, -1, -1): + while dx[i] > self.c_pbcbox.hbox_diag[i]: + for j in range (i, -1, -1): + dx[j] -= self.c_pbcbox.box[i][j] + + while dx[i] <= self.c_pbcbox.mhbox_diag[i]: + for j in range (i, -1, -1): + dx[j] += self.c_pbcbox.box[i][j] + + cdef real[:, ::1]fast_put_atoms_in_bbox(self, real[:,::1] coords) nogil: + cdef ns_int i, m, d, natoms, wd = 0 + cdef real[:,::1] bbox_coords + + natoms = coords.shape[0] + with gil: + if natoms == 0: + bbox_coords = np.empty((0, DIM)) + else: + bbox_coords = coords.copy() + + for i in range(natoms): + for m in range(DIM - 1, -1, -1): + while bbox_coords[i, m] < 0: + for d in range(m+1): + bbox_coords[i, d] += self.c_pbcbox.box[m][d] + while bbox_coords[i, m] >= self.c_pbcbox.box[m][m]: + for d in range(m+1): + bbox_coords[i, d] -= self.c_pbcbox.box[m][d] + return bbox_coords + + def put_atoms_in_bbox(self, real[:,::1] coords): + return np.asarray(self.fast_put_atoms_in_bbox(coords)) + +######################################################################################################################## +# +# Neighbor Search Stuff +# +######################################################################################################################## +cdef struct ns_grid: + ns_int size + ns_int[DIM] ncells + real[DIM] cellsize + ns_int *nbeads + ns_int **beadids + +cdef ns_grid initialize_nsgrid(matrix box, + float cutoff) nogil: + cdef ns_grid grid + cdef ns_int i + + for i in range(DIM): + grid.ncells[i] = (box[i][i] / cutoff) + if grid.ncells[i] == 0: + grid.ncells[i] = 1 + grid.cellsize[i] = box[i][i] / grid.ncells[i] + + grid.size = grid.ncells[XX] * grid.ncells[YY] * grid.ncells[ZZ] + return grid + +cdef ns_int populate_grid(ns_grid *grid, + real[:,::1] coords) nogil: + cdef ns_int ncoords = coords.shape[0] + cdef bint ret_val + + ret_val = populate_grid_array(grid, + &coords[0, 0], + ncoords) + + return ret_val + +cdef ns_int populate_grid_array(ns_grid *grid, + rvec *coords, + ns_int ncoords) nogil: + cdef ns_int i, cellindex = -1 + cdef ns_int grid_size = grid.size + cdef ns_int *allocated_size = NULL + + if grid_size != grid.ncells[XX] * grid.ncells[YY] * grid.ncells[ZZ]: # Grid not initialized + return RET_ERROR + + + # Allocate memory + grid.nbeads = malloc(sizeof(ns_int) * grid_size) + if grid.nbeads == NULL: + fprintf(stderr,"FATAL: Could not allocate memory for NS grid.nbeads (requested: %i bytes)\n", + sizeof(ns_int) * grid_size) + abort() + + allocated_size = malloc(sizeof(ns_int) * grid_size) + if allocated_size == NULL: + fprintf(stderr,"FATAL: Could not allocate memory for NS allocated_size (requested: %i bytes)\n", + sizeof(ns_int) * grid_size) + abort() + + for i in range(grid_size): + grid.nbeads[i] = 0 + allocated_size[i] = GRID_ALLOCATION_INCREMENT + + grid.beadids = malloc(sizeof(ns_int *) * grid_size) + if grid.beadids == NULL: + fprintf(stderr,"FATAL: Could not allocate memory for NS grid.beadids (requested: %i bytes)\n", + sizeof(ns_int *) * grid_size) + abort() + + for i in range(grid_size): + grid.beadids[i] = malloc(sizeof(ns_int) * allocated_size[i]) + if grid.beadids[i] == NULL: + fprintf(stderr,"FATAL: Could not allocate memory for NS grid.beadids[i] (requested: %i bytes)\n", + sizeof(ns_int) * allocated_size[i]) + abort() + + # Get cell indices for coords + for i in range(ncoords): + cellindex = (coords[i][ZZ] / grid.cellsize[ZZ]) * (grid.ncells[XX] * grid.ncells[YY]) +\ + (coords[i][YY] / grid.cellsize[YY]) * grid.ncells[XX] + \ + (coords[i][XX] / grid.cellsize[XX]) + + grid.beadids[cellindex][grid.nbeads[cellindex]] = i + grid.nbeads[cellindex] += 1 + + if grid.nbeads[cellindex] >= allocated_size[cellindex]: + allocated_size[cellindex] += GRID_ALLOCATION_INCREMENT + grid.beadids[cellindex] = realloc( grid.beadids[cellindex], sizeof(ns_int) * allocated_size[cellindex]) + free(allocated_size) + return RET_OK + +cdef void destroy_nsgrid(ns_grid *grid) nogil: + cdef ns_int i + if grid.nbeads != NULL: + free(grid.nbeads) + + for i in range(grid.size): + if grid.beadids[i] != NULL: + free(grid.beadids[i]) + free(grid.beadids) + + +cdef ns_neighborhood_holder *create_neighborhood_holder() nogil: + cdef ns_neighborhood_holder *holder + + holder = malloc(sizeof(ns_neighborhood_holder)) + + return holder + +cdef void free_neighborhood_holder(ns_neighborhood_holder *holder) nogil: + cdef ns_int i + + if holder == NULL: + return + + for i in range(holder.size): + if holder.neighborhoods[i].beadids != NULL: + free(holder.neighborhoods[i].beadids) + free(holder.neighborhoods[i]) + free(holder.neighborhoods) + free(holder) + +cdef ns_neighborhood *retrieve_neighborhood(rvec current_coords, real[:, ::1]neighborcoords, ns_grid *grid, PBCBox box, real cutoff2) nogil: + cdef ns_int d, m + cdef ns_int xi, yi, zi, bid + cdef real d2 + cdef rvec shifted_coord, dx, neighbor_coord, corrected_coords + + cdef bint already_checked[27] + cdef bint skip + cdef ns_int nchecked = 0, icheck + cdef ns_int cell_index + + cdef ns_neighborhood *neighborhood = malloc(sizeof(ns_neighborhood)) + if neighborhood == NULL: + abort() + + neighborhood.size = 0 + neighborhood.allocated_size = NEIGHBORHOOD_ALLOCATION_INCREMENT + neighborhood.beadids = malloc(NEIGHBORHOOD_ALLOCATION_INCREMENT * sizeof(ns_int)) + ###Modified here + neighborhood.beaddist = malloc(NEIGHBORHOOD_ALLOCATION_INCREMENT * sizeof(real)) + ### + if neighborhood.beadids == NULL: + abort() + + for zi in range(3): + for yi in range(3): + for xi in range(3): + # Calculate and/or reinitialize shifted coordinates + shifted_coord[XX] = current_coords[XX] + (xi - 1) * grid.cellsize[XX] + shifted_coord[YY] = current_coords[YY] + (yi - 1) * grid.cellsize[YY] + shifted_coord[ZZ] = current_coords[ZZ] + (zi - 1) * grid.cellsize[ZZ] + + # Make sure the shifted coordinates is inside the brick-shaped box + for m in range(DIM - 1, -1, -1): + + while shifted_coord[m] < 0: + for d in range(m+1): + shifted_coord[d] += box.c_pbcbox.box[m][d] + + + while shifted_coord[m] >= box.c_pbcbox.box[m][m]: + for d in range(m+1): + shifted_coord[d] -= box.c_pbcbox.box[m][d] + + # Get the cell index corresponding to the coord + cell_index = (shifted_coord[ZZ] / grid.cellsize[ZZ]) * (grid.ncells[XX] * grid.ncells[YY]) +\ + (shifted_coord[YY] / grid.cellsize[YY]) * grid.ncells[XX] + \ + (shifted_coord[XX] / grid.cellsize[XX]) + + # Just a safeguard + if cell_index >= grid.size: + continue + + # Check the cell index was not already selected + skip = False + for icheck in range(nchecked): + if already_checked[icheck] == cell_index: + skip = True + break + if skip: + continue + + # Search for neighbors inside this cell + for i_bead in range(grid.nbeads[cell_index]): + bid = grid.beadids[cell_index][i_bead] + + box.fast_pbc_dx(current_coords, &neighborcoords[bid, XX], dx) + + d2 = rvec_norm2(dx) + + if d2 < cutoff2: + if d2 < EPSILON: # Don't add the current bead as its own neighbor! + continue + + # Update neighbor lists + neighborhood.beadids[neighborhood.size] = bid + ### Modified here + neighborhood.beaddist[neighborhood.size] = d2 + ### + neighborhood.size += 1 + + if neighborhood.size >= neighborhood.allocated_size: + neighborhood.allocated_size += NEIGHBORHOOD_ALLOCATION_INCREMENT + neighborhood.beadids = realloc( neighborhood.beadids, neighborhood.allocated_size * sizeof(ns_int)) + ###Modified here + neighborhood.beaddist = realloc( neighborhood.beaddist, neighborhood.allocated_size * sizeof(real)) + ### + if neighborhood.beadids == NULL: + abort() + ###Modified + if neighborhood.beaddist == NULL: + abort() + ### + # Register the cell as checked + already_checked[nchecked] = cell_index + nchecked += 1 + + return neighborhood +cdef ns_neighborhood_holder *ns_core(real[:, ::1] refcoords, + real[:, ::1] neighborcoords, + ns_grid *grid, + PBCBox box, + real cutoff) nogil: + cdef ns_int coordid, i, j + cdef ns_int ncoords = refcoords.shape[0] + cdef ns_int ncoords_neighbors = neighborcoords.shape[0] + cdef real cutoff2 = cutoff * cutoff + cdef ns_neighborhood_holder *holder + + cdef ns_int *neighbor_buf + cdef ns_int buf_size, ibuf + + holder = create_neighborhood_holder() + if holder == NULL: + fprintf(stderr,"FATAL: Could not allocate memory for NS holder\n", + sizeof(ns_int) * ncoords) + abort() + + holder.size = ncoords + holder.neighborhoods = malloc(sizeof(ns_neighborhood *) * ncoords) + if holder.neighborhoods == NULL: + fprintf(stderr,"FATAL: Could not allocate memory for NS holder.neighborhoods (requested: %i bytes)\n", + sizeof(ns_neighborhood) * ncoords) + abort() + + # Here starts the real core and the iteration over coordinates + for coordid in range(ncoords): + holder.neighborhoods[coordid] = retrieve_neighborhood(&refcoords[coordid, XX], + neighborcoords, + grid, + box, + cutoff2) + holder.neighborhoods[coordid].cutoff = cutoff + + return holder + +# Python interface +cdef class FastNS(object): + cdef PBCBox box + cdef readonly int nthreads + cdef readonly real[:, ::1] coords + cdef real[:, ::1] coords_bbox + cdef readonly real cutoff + cdef bint prepared + cdef ns_grid *grid + + + def __init__(self, box): + if box.shape != (3, 3): + raise ValueError("Box must be provided as triclinic_dimensions (a 3x3 numpy.ndarray of unit cell vectors") + + self.box = PBCBox(box) + + self.nthreads = 1 + + self.coords = None + self.coords_bbox = None + + self.cutoff = -1 + + self.prepared = False + + self.grid = malloc(sizeof(ns_grid)) + + + def __dealloc__(self): + #destroy_nsgrid(self.grid) + self.grid.size = 0 + + def set_coords(self, real[:, ::1] coords): + self.coords = coords + + # Make sure atoms are inside the brick-shaped box + self.coords_bbox = self.box.fast_put_atoms_in_bbox(coords) + + self.prepared = False + + + def set_cutoff(self, real cutoff): + self.cutoff = cutoff + + self.prepared = False + + + def prepare(self, force=False): + cdef ns_int i + cdef bint initialization_ok + + if self.prepared and not force: + print("NS already prepared, nothing to do!") + + if self.coords is None: + raise ValueError("Coordinates must be set before NS preparation!") + + if self.cutoff < 0: + raise ValueError("Cutoff must be set before NS preparation!") + + with nogil: + initialization_ok = False + + # Initializing grid + for i in range(DIM): + self.grid.ncells[i] = (self.box.c_pbcbox.box[i][i] / self.cutoff) + if self.grid.ncells[i] == 0: + self.grid.ncells[i] = 1 + self.grid.cellsize[i] = self.box.c_pbcbox.box[i][i] / self.grid.ncells[i] + + self.grid.size = self.grid.ncells[XX] * self.grid.ncells[YY] * self.grid.ncells[ZZ] + + # Populating grid + if populate_grid(self.grid, self.coords_bbox) == RET_OK: + initialization_ok = True + + + if initialization_ok: + self.prepared = True + else: + raise RuntimeError("Could not initialize NS grid") + + + def search(self, real[:, ::1]search_coords, return_ids=False): + cdef real[:, ::1] search_coords_bbox + cdef ns_int nid, i, j + cdef ns_neighborhood_holder *holder + cdef ns_neighborhood *neighborhood + + if not self.prepared: + self.prepare() + + + # Make sure atoms are inside the brick-shaped box + search_coords_bbox = self.box.fast_put_atoms_in_bbox(search_coords) + + with nogil: + holder = ns_core(search_coords_bbox, self.coords_bbox, self.grid, self.box, self.cutoff) + + neighbors = [] + ###Modify for distance + sqdist = [] + indx = [] + ### + for nid in range(holder.size): + neighborhood = holder.neighborhoods[nid] + + if return_ids: + neighborhood_py = np.empty(neighborhood.size, dtype=np.int64) + ###Modify for distance + neighborhood_dis = np.empty(neighborhood.size, dtype=np.float32) + neighborhood_indx = np.empty(neighborhood.size, dtype=np.int64) + ### + for i in range(neighborhood.size): + neighborhood_py[i] = neighborhood.beadids[i] + ###Modify for distance + neighborhood_dis[i] = neighborhood.beaddist[i] + ### + else: + neighborhood_py = np.empty((neighborhood.size, DIM), dtype=np.float32) + ###Modify for distance + neighborhood_dis = np.empty((neighborhood.size), dtype=np.float32) + neighborhood_indx = np.empty(neighborhood.size, dtype=np.int64) + ### + for i in range(neighborhood.size): + ###Modify for distance + neighborhood_dis[i] = neighborhood.beaddist[i] + neighborhood_indx[i] = neighborhood.beadids[i] + ### + + for j in range(DIM): + neighborhood_py[i,j] = self.coords[neighborhood.beadids[i], j] + + neighbors.append(neighborhood_py) + sqdist.append(neighborhood_dis) + indx.append(neighborhood_indx) + + # Free Memory + free_neighborhood_holder(holder) + + return neighbors, sqdist, indx From 97d07fe763e87bb0c94d060f7685c0bf8086c03d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Buchoux?= Date: Sun, 1 Jul 2018 14:06:40 +0200 Subject: [PATCH 14/47] Added MDAnalysis header file to c_gridsearch.pyx --- package/MDAnalysis/lib/c_gridsearch.pyx | 32 +++++++++++++++---------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/package/MDAnalysis/lib/c_gridsearch.pyx b/package/MDAnalysis/lib/c_gridsearch.pyx index a4e8745aec4..ac73c97a5a2 100644 --- a/package/MDAnalysis/lib/c_gridsearch.pyx +++ b/package/MDAnalysis/lib/c_gridsearch.pyx @@ -1,22 +1,28 @@ -# -*- coding: utf-8; Mode: python; tab-width: 4; indent-tabs-mode:nil; -*- +# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; -*- # vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 # -# Copyright (C) 2013-2016 Sébastien Buchoux +# MDAnalysis --- https://www.mdanalysis.org # -# This file is part of FATSLiM. +# Copyright (C) 2013-2018 Sébastien Buchoux +# Copyright (c) 2018 The MDAnalysis Development Team and contributors +# (see the file AUTHORS for the full list of names) # -# FATSLiM is free software: you can redistribute it and/or modify -# it under the terms of the GNU General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. +# Released under the GNU Public Licence, v3 or any higher version # -# FATSLiM is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. +# Please cite your use of MDAnalysis in published work: # -# You should have received a copy of the GNU General Public License -# along with FATSLiM. If not, see . +# 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 +# +# + #cython: cdivision=True #cython: boundscheck=False From a26f437a7a73c0eec4f3f1ec3f97b56c0b5522e9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Buchoux?= Date: Sun, 1 Jul 2018 16:53:22 +0200 Subject: [PATCH 15/47] Corrected C (de)allocation in c_gridsearch.pyx. Should fix memory leak --- package/MDAnalysis/lib/c_gridsearch.pyx | 32 +++++++++++++++---------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/package/MDAnalysis/lib/c_gridsearch.pyx b/package/MDAnalysis/lib/c_gridsearch.pyx index ac73c97a5a2..2a7af4b33b7 100644 --- a/package/MDAnalysis/lib/c_gridsearch.pyx +++ b/package/MDAnalysis/lib/c_gridsearch.pyx @@ -145,16 +145,15 @@ cdef class PBCBox(object): # Update shift vectors self.c_pbcbox.ntric_vec = 0 - # We will only use single shifts, but we will check a few - # more shifts to see if there is a limiting distance - # above which we can not be sure of the correct distance. - for kk in range(5): + + # We will only use single shifts + for kk in range(3): k = order[kk] - for jj in range(5): + for jj in range(3): j = order[jj] - for ii in range(5): + for ii in range(3): i = order[ii] # A shift is only useful when it is trilinic @@ -177,8 +176,8 @@ cdef class PBCBox(object): else: pos[d] = max(-self.c_pbcbox.hbox_diag[d], -trial[d]) - d2old += sqrt(pos[d]) - d2new += sqrt(pos[d] + trial[d]) + d2old += pos[d]**2 + d2new += (pos[d] + trial[d])**2 if BOX_MARGIN*d2new < d2old: if not (j < -1 or j > 1 or k < -1 or k > 1): @@ -196,7 +195,7 @@ cdef class PBCBox(object): d2new_c = 0 for d in range(DIM): - d2new_c += sqrt(pos[d] + trial[d] - shift*box[dd, d]) + d2new_c += (pos[d] + trial[d] - shift*box[dd, d])**2 if d2new_c <= BOX_MARGIN*d2new: use = False @@ -209,7 +208,13 @@ cdef class PBCBox(object): % MAX_NTRICVEC) print(" There is probably something wrong with " "your box.") - print(box) + print(np.array(box)) + + for i in range(self.c_pbcbox.ntric_vec): + print(" -> shift #{}: [{}, {}, {}]".format(i+1, + self.c_pbcbox.tric_shift[i][XX], + self.c_pbcbox.tric_shift[i][YY], + self.c_pbcbox.tric_shift[i][ZZ])) else: for d in range(DIM): self.c_pbcbox.tric_vec[self.c_pbcbox.ntric_vec][d] = \ @@ -539,6 +544,9 @@ cdef class FastNS(object): cdef ns_grid *grid + def __cinit__(self): + self.grid = malloc(sizeof(ns_grid)) + def __init__(self, box): if box.shape != (3, 3): raise ValueError("Box must be provided as triclinic_dimensions (a 3x3 numpy.ndarray of unit cell vectors") @@ -554,11 +562,9 @@ cdef class FastNS(object): self.prepared = False - self.grid = malloc(sizeof(ns_grid)) - def __dealloc__(self): - #destroy_nsgrid(self.grid) + destroy_nsgrid(self.grid) self.grid.size = 0 def set_coords(self, real[:, ::1] coords): From 28554002b8ea389bf82dfe1a3650a24ced8aab67 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Buchoux?= Date: Sun, 1 Jul 2018 16:54:21 +0200 Subject: [PATCH 16/47] Removed pointer to nsgrid structure to avoid need to free it --- package/MDAnalysis/lib/c_gridsearch.pyx | 17 ++-- .../MDAnalysisTests/lib/test_gridsearch.py | 81 ++++--------------- 2 files changed, 22 insertions(+), 76 deletions(-) diff --git a/package/MDAnalysis/lib/c_gridsearch.pyx b/package/MDAnalysis/lib/c_gridsearch.pyx index 2a7af4b33b7..7a6de80389d 100644 --- a/package/MDAnalysis/lib/c_gridsearch.pyx +++ b/package/MDAnalysis/lib/c_gridsearch.pyx @@ -42,8 +42,7 @@ DEF MAX_NTRICVEC=12 from libc.stdlib cimport malloc, realloc, free, abort from libc.stdio cimport fprintf, stderr -from libc.math cimport sqrt -from libc.math cimport abs as real_abs + import numpy as np cimport numpy as np @@ -541,11 +540,7 @@ cdef class FastNS(object): cdef real[:, ::1] coords_bbox cdef readonly real cutoff cdef bint prepared - cdef ns_grid *grid - - - def __cinit__(self): - self.grid = malloc(sizeof(ns_grid)) + cdef ns_grid grid def __init__(self, box): if box.shape != (3, 3): @@ -562,9 +557,11 @@ cdef class FastNS(object): self.prepared = False + self.grid.size = 0 + def __dealloc__(self): - destroy_nsgrid(self.grid) + destroy_nsgrid(&self.grid) self.grid.size = 0 def set_coords(self, real[:, ::1] coords): @@ -608,7 +605,7 @@ cdef class FastNS(object): self.grid.size = self.grid.ncells[XX] * self.grid.ncells[YY] * self.grid.ncells[ZZ] # Populating grid - if populate_grid(self.grid, self.coords_bbox) == RET_OK: + if populate_grid(&self.grid, self.coords_bbox) == RET_OK: initialization_ok = True @@ -632,7 +629,7 @@ cdef class FastNS(object): search_coords_bbox = self.box.fast_put_atoms_in_bbox(search_coords) with nogil: - holder = ns_core(search_coords_bbox, self.coords_bbox, self.grid, self.box, self.cutoff) + holder = ns_core(search_coords_bbox, self.coords_bbox, &self.grid, self.box, self.cutoff) neighbors = [] ###Modify for distance diff --git a/testsuite/MDAnalysisTests/lib/test_gridsearch.py b/testsuite/MDAnalysisTests/lib/test_gridsearch.py index 359d32544b7..4991a2da71d 100644 --- a/testsuite/MDAnalysisTests/lib/test_gridsearch.py +++ b/testsuite/MDAnalysisTests/lib/test_gridsearch.py @@ -23,13 +23,16 @@ from __future__ import print_function, absolute_import import pytest -from numpy.testing import assert_equal, assert_allclose + +from numpy.testing import assert_equal import numpy as np import MDAnalysis as mda from MDAnalysis.lib import grid from MDAnalysis.lib.pkdtree import PeriodicKDTree +from MDAnalysis.lib.mdamath import triclinic_vectors + from MDAnalysisTests.datafiles import GRO @@ -39,30 +42,13 @@ def universe(): return u -@pytest.fixture -def grid_results(): - u = mda.Universe(GRO) - cutoff = 2 - ref_pos = u.atoms.positions[13937] - return run_grid_search(u, ref_pos, cutoff) - - -def run_grid_search(u, ref_pos, cutoff): - coords = u.atoms.positions - - # Run grid search - searcher = grid.FastNS(u) - searcher.set_cutoff(cutoff) - searcher.set_coords(coords) - searcher.prepare() - - return searcher.search(ref_pos) - - def run_search(universe, ref_id): cutoff = 3 + coords = universe.atoms.positions ref_pos = coords[ref_id] + triclinic_box = triclinic_vectors(universe.dimensions) + # Run pkdtree search pkdt = PeriodicKDTree(universe.atoms.dimensions, bucket_size=10) @@ -75,8 +61,13 @@ def run_search(universe, ref_id): results_pkdtree.sort() # Run grid search - results_grid = run_grid_search(universe, ref_pos, cutoff) - results_grid = results_grid.get_indices()[0] + + searcher = grid.FastNS(triclinic_box) + searcher.set_cutoff(cutoff) + searcher.set_coords(coords) + searcher.prepare() + + results_grid = searcher.search(np.array([ref_pos, ]), return_ids=True)[0][0] results_grid.sort() return results_pkdtree, results_grid @@ -95,46 +86,4 @@ def test_gridsearch_PBC(universe): ref_id = 13937 results_pkdtree, results_grid = run_search(universe, ref_id) - assert_equal(results_pkdtree, results_grid) - - -def test_gridsearch_arraycoord(universe): - """Check the NS routine accepts a single bead coordinate as well as array of coordinates""" - cutoff = 2 - ref_pos = universe.atoms.positions[:5] - - results = [ - np.array([2, 1, 4, 3]), - np.array([2, 0, 3]), - np.array([0, 1, 3]), - np.array([ 2, 0, 1, 38341]), - np.array([ 6, 0, 5, 17]) - ] - - results_grid = run_grid_search(universe, ref_pos, cutoff).get_indices() - - assert_equal(results_grid, results) - - -def test_gridsearch_search_coordinates(grid_results): - """Check the NS routine can return coordinates instead of ids""" - - results = np.array( - [ - [40.32, 34.25, 55.9], - [0.61, 76.33, -0.56], - [0.48999998, 75.9, 0.19999999], - [-0.11, 76.19, 0.77] - ]) - - assert_allclose(grid_results.get_coordinates()[0], results) - - -def test_gridsearch_search_distances(grid_results): - """Check the NS routine can return PBC distances from neighbors""" - results = np.array([0.096, 0.096, 0.015, 0.179]) * 10 # These distances were obtained using gmx distance - results.sort() - - rounded_results = np.round(grid_results.get_distances()[0], 2) - - assert_allclose(sorted(rounded_results), results) \ No newline at end of file + assert_equal(results_pkdtree, results_grid) \ No newline at end of file From 4d4a41511220599d333851f772a25b9c88acd1e6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Buchoux?= Date: Wed, 4 Jul 2018 18:13:41 +0200 Subject: [PATCH 17/47] Removed abort() from c_gridsearch.pyx and grid allocation moved to FastNS method. --- package/MDAnalysis/lib/c_gridsearch.pyx | 267 +++++++----------- .../MDAnalysisTests/lib/test_gridsearch.py | 2 +- 2 files changed, 104 insertions(+), 165 deletions(-) diff --git a/package/MDAnalysis/lib/c_gridsearch.pyx b/package/MDAnalysis/lib/c_gridsearch.pyx index 7a6de80389d..98d97eded65 100644 --- a/package/MDAnalysis/lib/c_gridsearch.pyx +++ b/package/MDAnalysis/lib/c_gridsearch.pyx @@ -31,18 +31,15 @@ DEF DIM = 3 DEF XX = 0 DEF YY = 1 DEF ZZ = 2 -DEF RET_OK = 1 -DEF RET_ERROR = 0 + DEF EPSILON = 1e-5 + DEF NEIGHBORHOOD_ALLOCATION_INCREMENT = 50 -DEF GRID_ALLOCATION_INCREMENT = 50 DEF BOX_MARGIN=1.0010 DEF MAX_NTRICVEC=12 -from libc.stdlib cimport malloc, realloc, free, abort -from libc.stdio cimport fprintf, stderr - +from libc.stdlib cimport malloc, realloc, free import numpy as np cimport numpy as np @@ -64,8 +61,7 @@ cdef struct ns_neighborhood: ns_int allocated_size ns_int size ns_int *beadids - real *beaddist - ### + cdef struct ns_neighborhood_holder: ns_int size ns_neighborhood **neighborhoods @@ -80,6 +76,7 @@ cdef void rvec_clear(rvec a) nogil: a[YY]=0.0 a[ZZ]=0.0 + cdef struct cPBCBox_t: matrix box rvec fbox_diag @@ -90,6 +87,8 @@ cdef struct cPBCBox_t: ns_int[DIM] tric_shift[MAX_NTRICVEC] real[DIM] tric_vec[MAX_NTRICVEC] + +# Class to handle PBC calculations cdef class PBCBox(object): cdef cPBCBox_t c_pbcbox cdef rvec center @@ -232,6 +231,7 @@ cdef class PBCBox(object): raise ValueError("Box does not correspond to PBC=xyz") self.fast_update(box) + cdef void fast_pbc_dx(self, rvec ref, rvec other, rvec dx) nogil: cdef ns_int i, j cdef rvec dx_start, trial @@ -284,102 +284,12 @@ cdef struct ns_grid: ns_int *nbeads ns_int **beadids -cdef ns_grid initialize_nsgrid(matrix box, - float cutoff) nogil: - cdef ns_grid grid - cdef ns_int i - - for i in range(DIM): - grid.ncells[i] = (box[i][i] / cutoff) - if grid.ncells[i] == 0: - grid.ncells[i] = 1 - grid.cellsize[i] = box[i][i] / grid.ncells[i] - - grid.size = grid.ncells[XX] * grid.ncells[YY] * grid.ncells[ZZ] - return grid - -cdef ns_int populate_grid(ns_grid *grid, - real[:,::1] coords) nogil: - cdef ns_int ncoords = coords.shape[0] - cdef bint ret_val - - ret_val = populate_grid_array(grid, - &coords[0, 0], - ncoords) - - return ret_val - -cdef ns_int populate_grid_array(ns_grid *grid, - rvec *coords, - ns_int ncoords) nogil: - cdef ns_int i, cellindex = -1 - cdef ns_int grid_size = grid.size - cdef ns_int *allocated_size = NULL - - if grid_size != grid.ncells[XX] * grid.ncells[YY] * grid.ncells[ZZ]: # Grid not initialized - return RET_ERROR - - - # Allocate memory - grid.nbeads = malloc(sizeof(ns_int) * grid_size) - if grid.nbeads == NULL: - fprintf(stderr,"FATAL: Could not allocate memory for NS grid.nbeads (requested: %i bytes)\n", - sizeof(ns_int) * grid_size) - abort() - - allocated_size = malloc(sizeof(ns_int) * grid_size) - if allocated_size == NULL: - fprintf(stderr,"FATAL: Could not allocate memory for NS allocated_size (requested: %i bytes)\n", - sizeof(ns_int) * grid_size) - abort() - - for i in range(grid_size): - grid.nbeads[i] = 0 - allocated_size[i] = GRID_ALLOCATION_INCREMENT - - grid.beadids = malloc(sizeof(ns_int *) * grid_size) - if grid.beadids == NULL: - fprintf(stderr,"FATAL: Could not allocate memory for NS grid.beadids (requested: %i bytes)\n", - sizeof(ns_int *) * grid_size) - abort() - - for i in range(grid_size): - grid.beadids[i] = malloc(sizeof(ns_int) * allocated_size[i]) - if grid.beadids[i] == NULL: - fprintf(stderr,"FATAL: Could not allocate memory for NS grid.beadids[i] (requested: %i bytes)\n", - sizeof(ns_int) * allocated_size[i]) - abort() - - # Get cell indices for coords - for i in range(ncoords): - cellindex = (coords[i][ZZ] / grid.cellsize[ZZ]) * (grid.ncells[XX] * grid.ncells[YY]) +\ - (coords[i][YY] / grid.cellsize[YY]) * grid.ncells[XX] + \ - (coords[i][XX] / grid.cellsize[XX]) - - grid.beadids[cellindex][grid.nbeads[cellindex]] = i - grid.nbeads[cellindex] += 1 - - if grid.nbeads[cellindex] >= allocated_size[cellindex]: - allocated_size[cellindex] += GRID_ALLOCATION_INCREMENT - grid.beadids[cellindex] = realloc( grid.beadids[cellindex], sizeof(ns_int) * allocated_size[cellindex]) - free(allocated_size) - return RET_OK - -cdef void destroy_nsgrid(ns_grid *grid) nogil: - cdef ns_int i - if grid.nbeads != NULL: - free(grid.nbeads) - - for i in range(grid.size): - if grid.beadids[i] != NULL: - free(grid.beadids[i]) - free(grid.beadids) - - cdef ns_neighborhood_holder *create_neighborhood_holder() nogil: cdef ns_neighborhood_holder *holder holder = malloc(sizeof(ns_neighborhood_holder)) + holder.size = 0 + holder.neighborhoods = NULL return holder @@ -393,7 +303,9 @@ cdef void free_neighborhood_holder(ns_neighborhood_holder *holder) nogil: if holder.neighborhoods[i].beadids != NULL: free(holder.neighborhoods[i].beadids) free(holder.neighborhoods[i]) - free(holder.neighborhoods) + + if holder.neighborhoods != NULL: + free(holder.neighborhoods) free(holder) cdef ns_neighborhood *retrieve_neighborhood(rvec current_coords, real[:, ::1]neighborcoords, ns_grid *grid, PBCBox box, real cutoff2) nogil: @@ -409,16 +321,15 @@ cdef ns_neighborhood *retrieve_neighborhood(rvec current_coords, real[:, ::1]nei cdef ns_neighborhood *neighborhood = malloc(sizeof(ns_neighborhood)) if neighborhood == NULL: - abort() + return NULL neighborhood.size = 0 neighborhood.allocated_size = NEIGHBORHOOD_ALLOCATION_INCREMENT neighborhood.beadids = malloc(NEIGHBORHOOD_ALLOCATION_INCREMENT * sizeof(ns_int)) - ###Modified here - neighborhood.beaddist = malloc(NEIGHBORHOOD_ALLOCATION_INCREMENT * sizeof(real)) - ### + if neighborhood.beadids == NULL: - abort() + free(neighborhood) + return NULL for zi in range(3): for yi in range(3): @@ -472,28 +383,23 @@ cdef ns_neighborhood *retrieve_neighborhood(rvec current_coords, real[:, ::1]nei # Update neighbor lists neighborhood.beadids[neighborhood.size] = bid - ### Modified here - neighborhood.beaddist[neighborhood.size] = d2 - ### neighborhood.size += 1 if neighborhood.size >= neighborhood.allocated_size: neighborhood.allocated_size += NEIGHBORHOOD_ALLOCATION_INCREMENT neighborhood.beadids = realloc( neighborhood.beadids, neighborhood.allocated_size * sizeof(ns_int)) - ###Modified here - neighborhood.beaddist = realloc( neighborhood.beaddist, neighborhood.allocated_size * sizeof(real)) - ### + if neighborhood.beadids == NULL: - abort() - ###Modified - if neighborhood.beaddist == NULL: - abort() - ### + free(neighborhood) + return NULL + # Register the cell as checked already_checked[nchecked] = cell_index nchecked += 1 return neighborhood + + cdef ns_neighborhood_holder *ns_core(real[:, ::1] refcoords, real[:, ::1] neighborcoords, ns_grid *grid, @@ -510,16 +416,12 @@ cdef ns_neighborhood_holder *ns_core(real[:, ::1] refcoords, holder = create_neighborhood_holder() if holder == NULL: - fprintf(stderr,"FATAL: Could not allocate memory for NS holder\n", - sizeof(ns_int) * ncoords) - abort() + return NULL - holder.size = ncoords holder.neighborhoods = malloc(sizeof(ns_neighborhood *) * ncoords) if holder.neighborhoods == NULL: - fprintf(stderr,"FATAL: Could not allocate memory for NS holder.neighborhoods (requested: %i bytes)\n", - sizeof(ns_neighborhood) * ncoords) - abort() + free_neighborhood_holder(holder) + return NULL # Here starts the real core and the iteration over coordinates for coordid in range(ncoords): @@ -528,14 +430,18 @@ cdef ns_neighborhood_holder *ns_core(real[:, ::1] refcoords, grid, box, cutoff2) + if holder.neighborhoods[coordid] == NULL: + free_neighborhood_holder(holder) + return NULL + holder.neighborhoods[coordid].cutoff = cutoff + holder.size += 1 return holder # Python interface cdef class FastNS(object): cdef PBCBox box - cdef readonly int nthreads cdef readonly real[:, ::1] coords cdef real[:, ::1] coords_bbox cdef readonly real cutoff @@ -548,8 +454,6 @@ cdef class FastNS(object): self.box = PBCBox(box) - self.nthreads = 1 - self.coords = None self.coords_bbox = None @@ -561,7 +465,15 @@ cdef class FastNS(object): def __dealloc__(self): - destroy_nsgrid(&self.grid) + cdef ns_int i + if self.grid.nbeads != NULL: + free(self.grid.nbeads) + + for i in range(self.grid.size): + if self.grid.beadids[i] != NULL: + free(self.grid.beadids[i]) + free(self.grid.beadids) + self.grid.size = 0 def set_coords(self, real[:, ::1] coords): @@ -580,8 +492,11 @@ cdef class FastNS(object): def prepare(self, force=False): - cdef ns_int i - cdef bint initialization_ok + cdef ns_int i, cellindex = -1 + cdef ns_int *allocated_size = NULL + cdef ns_int ncoords = self.coords.shape[0] + cdef ns_int allocation_guess + cdef rvec *coords = &self.coords_bbox[0, 0] if self.prepared and not force: print("NS already prepared, nothing to do!") @@ -593,7 +508,6 @@ cdef class FastNS(object): raise ValueError("Cutoff must be set before NS preparation!") with nogil: - initialization_ok = False # Initializing grid for i in range(DIM): @@ -601,18 +515,61 @@ cdef class FastNS(object): if self.grid.ncells[i] == 0: self.grid.ncells[i] = 1 self.grid.cellsize[i] = self.box.c_pbcbox.box[i][i] / self.grid.ncells[i] - self.grid.size = self.grid.ncells[XX] * self.grid.ncells[YY] * self.grid.ncells[ZZ] - # Populating grid - if populate_grid(&self.grid, self.coords_bbox) == RET_OK: - initialization_ok = True - - - if initialization_ok: - self.prepared = True - else: - raise RuntimeError("Could not initialize NS grid") + # This is just a guess on how much memory we might for each grid cell: + # we just assume an average bead density and we take four times this density just to be safe + allocation_guess = (4 * (ncoords / self.grid.size + 1)) + + # Allocate memory for the grid + self.grid.nbeads = malloc(sizeof(ns_int) * self.grid.size) + if self.grid.nbeads == NULL: + with gil: + raise MemoryError("Could not allocate memory for NS grid") + + # Allocate memory from temporary allocation counter + allocated_size = malloc(sizeof(ns_int) * self.grid.size) + if allocated_size == NULL: + # No need to free grid.nbeads as it will be freed by destroy_nsgrid called by __dealloc___ + with gil: + raise MemoryError("Could not allocate memory for allocation buffer") + + # Pre-allocate some memory for grid cells + for i in range(self.grid.size): + self.grid.nbeads[i] = 0 + allocated_size[i] = allocation_guess + + self.grid.beadids = malloc(sizeof(ns_int *) * self.grid.size) + if self.grid.beadids == NULL: + with gil: + raise MemoryError("Could not allocate memory for grid cells") + + for i in range(self.grid.size): + self.grid.beadids[i] = malloc(sizeof(ns_int) * allocated_size[i]) + if self.grid.beadids[i] == NULL: + with gil: + raise MemoryError("Could not allocate memory for grid cell") + + # Populate grid cells using the coordinates (ie do the heavy work) + for i in range(ncoords): + cellindex = (coords[i][ZZ] / self.grid.cellsize[ZZ]) * (self.grid.ncells[XX] * self.grid.ncells[YY]) +\ + (coords[i][YY] / self.grid.cellsize[YY]) * self.grid.ncells[XX] + \ + (coords[i][XX] / self.grid.cellsize[XX]) + + self.grid.beadids[cellindex][self.grid.nbeads[cellindex]] = i + self.grid.nbeads[cellindex] += 1 + + # We need to allocate more memory (simply double the amount of memory as + # 1. it should barely be needed + # 2. the size should stay fairly reasonable + if self.grid.nbeads[cellindex] >= allocated_size[cellindex]: + allocated_size[cellindex] *= 2 + self.grid.beadids[cellindex] = realloc( self.grid.beadids[cellindex], sizeof(ns_int) * allocated_size[cellindex]) + + # Now we can free the allocation buffer + free(allocated_size) + + self.prepared = True def search(self, real[:, ::1]search_coords, return_ids=False): @@ -631,45 +588,27 @@ cdef class FastNS(object): with nogil: holder = ns_core(search_coords_bbox, self.coords_bbox, &self.grid, self.box, self.cutoff) + if holder == NULL: + raise MemoryError("Could not allocate memory to run NS core") + neighbors = [] - ###Modify for distance - sqdist = [] - indx = [] - ### for nid in range(holder.size): neighborhood = holder.neighborhoods[nid] if return_ids: neighborhood_py = np.empty(neighborhood.size, dtype=np.int64) - ###Modify for distance - neighborhood_dis = np.empty(neighborhood.size, dtype=np.float32) - neighborhood_indx = np.empty(neighborhood.size, dtype=np.int64) - ### + for i in range(neighborhood.size): neighborhood_py[i] = neighborhood.beadids[i] - ###Modify for distance - neighborhood_dis[i] = neighborhood.beaddist[i] - ### else: neighborhood_py = np.empty((neighborhood.size, DIM), dtype=np.float32) - ###Modify for distance - neighborhood_dis = np.empty((neighborhood.size), dtype=np.float32) - neighborhood_indx = np.empty(neighborhood.size, dtype=np.int64) - ### for i in range(neighborhood.size): - ###Modify for distance - neighborhood_dis[i] = neighborhood.beaddist[i] - neighborhood_indx[i] = neighborhood.beadids[i] - ### - for j in range(DIM): neighborhood_py[i,j] = self.coords[neighborhood.beadids[i], j] neighbors.append(neighborhood_py) - sqdist.append(neighborhood_dis) - indx.append(neighborhood_indx) # Free Memory free_neighborhood_holder(holder) - return neighbors, sqdist, indx + return neighbors diff --git a/testsuite/MDAnalysisTests/lib/test_gridsearch.py b/testsuite/MDAnalysisTests/lib/test_gridsearch.py index 4991a2da71d..1c3f5e06527 100644 --- a/testsuite/MDAnalysisTests/lib/test_gridsearch.py +++ b/testsuite/MDAnalysisTests/lib/test_gridsearch.py @@ -67,7 +67,7 @@ def run_search(universe, ref_id): searcher.set_coords(coords) searcher.prepare() - results_grid = searcher.search(np.array([ref_pos, ]), return_ids=True)[0][0] + results_grid = searcher.search(np.array([ref_pos, ]), return_ids=True)[0] results_grid.sort() return results_pkdtree, results_grid From c6339c5965f90008aeb25fbd1e4aa713bc7cdb8a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Buchoux?= Date: Sun, 8 Jul 2018 10:35:17 +0200 Subject: [PATCH 18/47] NSResults object added to store results from NS. Started documentation --- package/MDAnalysis/lib/c_gridsearch.pyx | 235 ++++++++++++++++-- .../MDAnalysisTests/lib/test_gridsearch.py | 81 +++++- 2 files changed, 278 insertions(+), 38 deletions(-) diff --git a/package/MDAnalysis/lib/c_gridsearch.pyx b/package/MDAnalysis/lib/c_gridsearch.pyx index 98d97eded65..9c50a5b1777 100644 --- a/package/MDAnalysis/lib/c_gridsearch.pyx +++ b/package/MDAnalysis/lib/c_gridsearch.pyx @@ -25,6 +25,15 @@ #cython: cdivision=True #cython: boundscheck=False +#cython: initializedcheck=False + +""" +Neighbor search library --- :mod:`MDAnalysis.lib.grid` +====================================================== + +This Neighbor search library is a serialized Cython port of the NS grid search implemented in GROMACS. +""" + # Preprocessor DEFs DEF DIM = 3 @@ -40,6 +49,7 @@ DEF BOX_MARGIN=1.0010 DEF MAX_NTRICVEC=12 from libc.stdlib cimport malloc, realloc, free +from libc.math cimport sqrt import numpy as np cimport numpy as np @@ -248,16 +258,21 @@ cdef class PBCBox(object): for j in range (i, -1, -1): dx[j] += self.c_pbcbox.box[i][j] + cdef real fast_distance2(self, rvec a, rvec b) nogil: + cdef rvec dx + self.fast_pbc_dx(a, b, dx) + return rvec_norm2(dx) + + cdef real fast_distance(self, rvec a, rvec b) nogil: + return sqrt(self.fast_distance2(a,b)) + cdef real[:, ::1]fast_put_atoms_in_bbox(self, real[:,::1] coords) nogil: cdef ns_int i, m, d, natoms, wd = 0 cdef real[:,::1] bbox_coords natoms = coords.shape[0] with gil: - if natoms == 0: - bbox_coords = np.empty((0, DIM)) - else: - bbox_coords = coords.copy() + bbox_coords = coords.copy() for i in range(natoms): for m in range(DIM - 1, -1, -1): @@ -270,6 +285,8 @@ cdef class PBCBox(object): return bbox_coords def put_atoms_in_bbox(self, real[:,::1] coords): + if coords.shape[0] == 0: + return np.zeros((0, DIM), dtype=np.float32) return np.asarray(self.fast_put_atoms_in_bbox(coords)) ######################################################################################################################## @@ -439,6 +456,167 @@ cdef ns_neighborhood_holder *ns_core(real[:, ::1] refcoords, return holder +cdef class NSResults(object): + """ + Class used to store results returned by `MDAnalysis.lib.grid.FastNS.search` + """ + cdef PBCBox box + cdef readonly real cutoff + cdef real[:, ::1] grid_coords + cdef real[:, ::1] ref_coords + cdef ns_int **nids + cdef ns_int *nsizes + cdef ns_int size + cdef list indices + cdef list coordinates + cdef list distances + + def __init__(self, PBCBox box, real cutoff): + self.box = box + self.cutoff = cutoff + + self.size = 0 + self.nids = NULL + self.nsizes = NULL + + self.grid_coords = None + self.ref_coords = None + + self.indices = None + self.coordinates = None + self.distances = None + + + cdef populate(self, ns_neighborhood_holder *holder, grid_coords, ref_coords): + cdef ns_int nid, i + cdef ns_neighborhood *neighborhood + + self.grid_coords = grid_coords.copy() + self.ref_coords = ref_coords.copy() + + # Allocate memory + self.nsizes = malloc(sizeof(ns_int) * holder.size) + if self.nsizes == NULL: + raise MemoryError("Could not allocate memory for NSResults") + + self.nids = malloc(sizeof(ns_int *) * holder.size) + if self.nids == NULL: + raise MemoryError("Could not allocate memory for NSResults") + + for nid in range(holder.size): + neighborhood = holder.neighborhoods[nid] + + self.nsizes[nid] = neighborhood.size + + self.nids[nid] = malloc(sizeof(ns_int *) * neighborhood.size) + if self.nids[nid] == NULL: + raise MemoryError("Could not allocate memory for NSResults") + + with nogil: + for nid in range(holder.size): + neighborhood = holder.neighborhoods[nid] + + for i in range(neighborhood.size): + self.nids[nid][i] = neighborhood.beadids[i] + + self.size = holder.size + + def __dealloc__(self): + if self.nids != NULL: + for i in range(self.size): + if self.nids[i] != NULL: + free(self.nids[i]) + free(self.nids) + + if self.nsizes != NULL: + free(self.nsizes) + + + def get_indices(self): + """ + Return Neighbors indices. + + :return: list of indices + """ + cdef ns_int i, nid, size + + if self.indices is None: + indices = [] + + for nid in range(self.size): + size = self.nsizes[nid] + + tmp_incides = np.empty((size), dtype=np.int) + + for i in range(size): + tmp_incides[i] = self.nids[nid][i] + + indices.append(tmp_incides) + + self.indices = indices + + return self.indices + + + def get_coordinates(self): + """ + Return coordinates of neighbors. + + :return: list of coordinates + """ + cdef ns_int i, nid, size, beadid + + if self.coordinates is None: + coordinates = [] + + for nid in range(self.size): + size = self.nsizes[nid] + + tmp_values = np.empty((size, DIM), dtype=np.float32) + + for i in range(size): + beadid = self.nids[nid][i] + tmp_values[i] = self.grid_coords[beadid] + + coordinates.append(tmp_values) + + self.coordinates = coordinates + + return self.coordinates + + + def get_distances(self): + """ + Return coordinates of neighbors. + + :return: list of distances + """ + cdef ns_int i, nid, size, j, beadid + cdef rvec ref, other, dx + cdef real dist + + if self.distances is None: + distances = [] + + for nid in range(self.size): + size = self.nsizes[nid] + + tmp_values = np.empty((size), dtype=np.float32) + ref = &self.ref_coords[nid, 0] + + for i in range(size): + beadid = self.nids[nid][i] + other = &self.grid_coords[beadid, 0] + + tmp_values[i] = self.box.fast_distance(ref, other) + + distances.append(tmp_values) + + self.distances = distances + + return self.distances + + # Python interface cdef class FastNS(object): cdef PBCBox box @@ -448,7 +626,14 @@ cdef class FastNS(object): cdef bint prepared cdef ns_grid grid - def __init__(self, box): + def __init__(self, u): + import MDAnalysis as mda + from MDAnalysis.lib.mdamath import triclinic_vectors + + if not isinstance(u, mda.Universe): + raise TypeError("FastNS class must be initialized with a valid MDAnalysis.Universe instance") + box = triclinic_vectors(u.dimensions) + if box.shape != (3, 3): raise ValueError("Box must be provided as triclinic_dimensions (a 3x3 numpy.ndarray of unit cell vectors") @@ -466,6 +651,7 @@ cdef class FastNS(object): def __dealloc__(self): cdef ns_int i + # Deallocate NS grid if self.grid.nbeads != NULL: free(self.grid.nbeads) @@ -475,7 +661,8 @@ cdef class FastNS(object): free(self.grid.beadids) self.grid.size = 0 - + + def set_coords(self, real[:, ::1] coords): self.coords = coords @@ -568,12 +755,12 @@ cdef class FastNS(object): # Now we can free the allocation buffer free(allocated_size) - self.prepared = True - def search(self, real[:, ::1]search_coords, return_ids=False): + def search(self, search_coords): cdef real[:, ::1] search_coords_bbox + cdef real[:, ::1] search_coords_view cdef ns_int nid, i, j cdef ns_neighborhood_holder *holder cdef ns_neighborhood *neighborhood @@ -581,9 +768,18 @@ cdef class FastNS(object): if not self.prepared: self.prepare() + # Check the shape of search_coords as a array of 3D coords if needed + shape = search_coords.shape + if len(shape) == 1: + if not shape[0] == 3: + raise ValueError("Coordinates must be 3D") + else: + search_coords_view = np.array([search_coords,], dtype=np.float32) + else: + search_coords_view = search_coords # Make sure atoms are inside the brick-shaped box - search_coords_bbox = self.box.fast_put_atoms_in_bbox(search_coords) + search_coords_bbox = self.box.fast_put_atoms_in_bbox(search_coords_view) with nogil: holder = ns_core(search_coords_bbox, self.coords_bbox, &self.grid, self.box, self.cutoff) @@ -591,24 +787,11 @@ cdef class FastNS(object): if holder == NULL: raise MemoryError("Could not allocate memory to run NS core") - neighbors = [] - for nid in range(holder.size): - neighborhood = holder.neighborhoods[nid] - if return_ids: - neighborhood_py = np.empty(neighborhood.size, dtype=np.int64) - - for i in range(neighborhood.size): - neighborhood_py[i] = neighborhood.beadids[i] - else: - neighborhood_py = np.empty((neighborhood.size, DIM), dtype=np.float32) - for i in range(neighborhood.size): - for j in range(DIM): - neighborhood_py[i,j] = self.coords[neighborhood.beadids[i], j] - - neighbors.append(neighborhood_py) + results = NSResults(self.box, self.cutoff) + results.populate(holder, self.coords, search_coords_view) - # Free Memory + # Free memory allocated to holder free_neighborhood_holder(holder) - return neighbors + return results diff --git a/testsuite/MDAnalysisTests/lib/test_gridsearch.py b/testsuite/MDAnalysisTests/lib/test_gridsearch.py index 1c3f5e06527..30ab8813036 100644 --- a/testsuite/MDAnalysisTests/lib/test_gridsearch.py +++ b/testsuite/MDAnalysisTests/lib/test_gridsearch.py @@ -24,14 +24,13 @@ import pytest -from numpy.testing import assert_equal +from numpy.testing import assert_equal, assert_allclose import numpy as np import MDAnalysis as mda from MDAnalysis.lib import grid from MDAnalysis.lib.pkdtree import PeriodicKDTree -from MDAnalysis.lib.mdamath import triclinic_vectors from MDAnalysisTests.datafiles import GRO @@ -42,12 +41,32 @@ def universe(): return u + +@pytest.fixture +def grid_results(): + u = mda.Universe(GRO) + cutoff = 2 + ref_pos = u.atoms.positions[13937] + return run_grid_search(u, ref_pos, cutoff) + + +def run_grid_search(u, ref_pos, cutoff): + coords = u.atoms.positions + + # Run grid search + searcher = grid.FastNS(u) + searcher.set_cutoff(cutoff) + searcher.set_coords(coords) + searcher.prepare() + + return searcher.search(ref_pos) + + + def run_search(universe, ref_id): cutoff = 3 - coords = universe.atoms.positions ref_pos = coords[ref_id] - triclinic_box = triclinic_vectors(universe.dimensions) # Run pkdtree search @@ -61,13 +80,8 @@ def run_search(universe, ref_id): results_pkdtree.sort() # Run grid search - - searcher = grid.FastNS(triclinic_box) - searcher.set_cutoff(cutoff) - searcher.set_coords(coords) - searcher.prepare() - - results_grid = searcher.search(np.array([ref_pos, ]), return_ids=True)[0] + results_grid = run_grid_search(universe, ref_pos, cutoff) + results_grid = results_grid.get_indices()[0] results_grid.sort() return results_pkdtree, results_grid @@ -86,4 +100,47 @@ def test_gridsearch_PBC(universe): ref_id = 13937 results_pkdtree, results_grid = run_search(universe, ref_id) - assert_equal(results_pkdtree, results_grid) \ No newline at end of file + + assert_equal(results_pkdtree, results_grid) + + +def test_gridsearch_arraycoord(universe): + """Check the NS routine accepts a single bead coordinate as well as array of coordinates""" + cutoff = 2 + ref_pos = universe.atoms.positions[:5] + + results = [ + np.array([2, 1, 4, 3]), + np.array([2, 0, 3]), + np.array([0, 1, 3]), + np.array([ 2, 0, 1, 38341]), + np.array([ 6, 0, 5, 17]) + ] + + results_grid = run_grid_search(universe, ref_pos, cutoff).get_indices() + + assert_equal(results_grid, results) + + +def test_gridsearch_search_coordinates(grid_results): + """Check the NS routine can return coordinates instead of ids""" + + results = np.array( + [ + [40.32, 34.25, 55.9], + [0.61, 76.33, -0.56], + [0.48999998, 75.9, 0.19999999], + [-0.11, 76.19, 0.77] + ]) + + assert_allclose(grid_results.get_coordinates()[0], results) + + +def test_gridsearch_search_distances(grid_results): + """Check the NS routine can return PBC distances from neighbors""" + results = np.array([0.096, 0.096, 0.015, 0.179]) * 10 # These distances were obtained using gmx distance + results.sort() + + rounded_results = np.round(grid_results.get_distances()[0], 2) + + assert_allclose(sorted(rounded_results), results) From 7f4a176cefa1ea1f33fd6d87ded4061944cfafce Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Buchoux?= Date: Sun, 8 Jul 2018 20:10:33 +0200 Subject: [PATCH 19/47] Memory allocation changed to PyMem to enhance speed in c_gridsearch.pyx Pre-converting np.resize to PyMem stuff --- package/MDAnalysis/lib/c_gridsearch.pyx | 856 +++++++++++++++++++++--- 1 file changed, 771 insertions(+), 85 deletions(-) diff --git a/package/MDAnalysis/lib/c_gridsearch.pyx b/package/MDAnalysis/lib/c_gridsearch.pyx index 9c50a5b1777..ffdb91fed7d 100644 --- a/package/MDAnalysis/lib/c_gridsearch.pyx +++ b/package/MDAnalysis/lib/c_gridsearch.pyx @@ -49,6 +49,8 @@ DEF BOX_MARGIN=1.0010 DEF MAX_NTRICVEC=12 from libc.stdlib cimport malloc, realloc, free +from cpython.mem cimport PyMem_Malloc, PyMem_Realloc, PyMem_Free + from libc.math cimport sqrt import numpy as np @@ -57,14 +59,18 @@ cimport numpy as np ctypedef np.int_t ns_int ctypedef np.float32_t real ctypedef real rvec[DIM] +ctypedef ns_int ivec[DIM] +ctypedef ns_int ipair[2] ctypedef real matrix[DIM][DIM] cdef struct ns_grid: + ns_int beadpercell ns_int size ns_int[DIM] ncells real[DIM] cellsize - ns_int *nbeads - ns_int **beadids + ns_int *nbeads # size + ns_int *beadids # size * beadpercell + ns_int *cellids # size cdef struct ns_neighborhood: real cutoff @@ -103,6 +109,7 @@ cdef class PBCBox(object): cdef cPBCBox_t c_pbcbox cdef rvec center cdef rvec bbox_center + cdef bint is_triclinic def __init__(self, real[:,::1] box): self.update(box) @@ -118,10 +125,15 @@ cdef class PBCBox(object): rvec_clear(self.center) # Update matrix + self.is_triclinic = False for i in range(DIM): for j in range(DIM): self.c_pbcbox.box[i][j] = box[i, j] self.center[j] += 0.5 * box[i, j] + + if i != j: + if box[i, j] > EPSILON: + self.is_triclinic = True self.bbox_center[i] = 0.5 * box[i, i] # Update diagonals @@ -274,14 +286,22 @@ cdef class PBCBox(object): with gil: bbox_coords = coords.copy() - for i in range(natoms): - for m in range(DIM - 1, -1, -1): - while bbox_coords[i, m] < 0: - for d in range(m+1): - bbox_coords[i, d] += self.c_pbcbox.box[m][d] - while bbox_coords[i, m] >= self.c_pbcbox.box[m][m]: - for d in range(m+1): - bbox_coords[i, d] -= self.c_pbcbox.box[m][d] + if self.is_triclinic: + for i in range(natoms): + for m in range(DIM - 1, -1, -1): + while bbox_coords[i, m] < 0: + for d in range(m+1): + bbox_coords[i, d] += self.c_pbcbox.box[m][d] + while bbox_coords[i, m] >= self.c_pbcbox.box[m][m]: + for d in range(m+1): + bbox_coords[i, d] -= self.c_pbcbox.box[m][d] + else: + for i in range(natoms): + for m in range(DIM): + while bbox_coords[i, m] < 0: + bbox_coords[i, m] += self.c_pbcbox.box[m][m] + while bbox_coords[i, m] >= self.c_pbcbox.box[m][m]: + bbox_coords[i, m] -= self.c_pbcbox.box[m][m] return bbox_coords def put_atoms_in_bbox(self, real[:,::1] coords): @@ -294,12 +314,6 @@ cdef class PBCBox(object): # Neighbor Search Stuff # ######################################################################################################################## -cdef struct ns_grid: - ns_int size - ns_int[DIM] ncells - real[DIM] cellsize - ns_int *nbeads - ns_int **beadids cdef ns_neighborhood_holder *create_neighborhood_holder() nogil: cdef ns_neighborhood_holder *holder @@ -327,7 +341,7 @@ cdef void free_neighborhood_holder(ns_neighborhood_holder *holder) nogil: cdef ns_neighborhood *retrieve_neighborhood(rvec current_coords, real[:, ::1]neighborcoords, ns_grid *grid, PBCBox box, real cutoff2) nogil: cdef ns_int d, m - cdef ns_int xi, yi, zi, bid + cdef ns_int xi, yi, zi, bid, i_bead cdef real d2 cdef rvec shifted_coord, dx, neighbor_coord, corrected_coords @@ -335,6 +349,7 @@ cdef ns_neighborhood *retrieve_neighborhood(rvec current_coords, real[:, ::1]nei cdef bint skip cdef ns_int nchecked = 0, icheck cdef ns_int cell_index + cdef ns_int cell_offset cdef ns_neighborhood *neighborhood = malloc(sizeof(ns_neighborhood)) if neighborhood == NULL: @@ -374,21 +389,23 @@ cdef ns_neighborhood *retrieve_neighborhood(rvec current_coords, real[:, ::1]nei (shifted_coord[XX] / grid.cellsize[XX]) # Just a safeguard - if cell_index >= grid.size: - continue + #if cell_index >= grid.size: + # continue # Check the cell index was not already selected - skip = False - for icheck in range(nchecked): - if already_checked[icheck] == cell_index: - skip = True - break - if skip: - continue + #skip = False + #for icheck in range(nchecked): + # if already_checked[icheck] == cell_index: + # skip = True + # break + #if skip: + # continue # Search for neighbors inside this cell for i_bead in range(grid.nbeads[cell_index]): - bid = grid.beadids[cell_index][i_bead] + cell_offset = cell_index * grid.beadpercell + i_bead + + bid = grid.beadids[cell_offset] box.fast_pbc_dx(current_coords, &neighborcoords[bid, XX], dx) @@ -411,8 +428,8 @@ cdef ns_neighborhood *retrieve_neighborhood(rvec current_coords, real[:, ::1]nei return NULL # Register the cell as checked - already_checked[nchecked] = cell_index - nchecked += 1 + #already_checked[nchecked] = cell_index + #nchecked += 1 return neighborhood @@ -462,8 +479,8 @@ cdef class NSResults(object): """ cdef PBCBox box cdef readonly real cutoff - cdef real[:, ::1] grid_coords - cdef real[:, ::1] ref_coords + cdef np.ndarray grid_coords + cdef np.ndarray ref_coords cdef ns_int **nids cdef ns_int *nsizes cdef ns_int size @@ -471,6 +488,16 @@ cdef class NSResults(object): cdef list coordinates cdef list distances + def __dealloc__(self): + if self.nids != NULL: + for i in range(self.size): + if self.nids[i] != NULL: + free(self.nids[i]) + free(self.nids) + + if self.nsizes != NULL: + free(self.nsizes) + def __init__(self, PBCBox box, real cutoff): self.box = box self.cutoff = cutoff @@ -491,8 +518,8 @@ cdef class NSResults(object): cdef ns_int nid, i cdef ns_neighborhood *neighborhood - self.grid_coords = grid_coords.copy() - self.ref_coords = ref_coords.copy() + self.grid_coords = np.asarray(grid_coords) + self.ref_coords = np.asarray(ref_coords) # Allocate memory self.nsizes = malloc(sizeof(ns_int) * holder.size) @@ -521,17 +548,6 @@ cdef class NSResults(object): self.size = holder.size - def __dealloc__(self): - if self.nids != NULL: - for i in range(self.size): - if self.nids[i] != NULL: - free(self.nids[i]) - free(self.nids) - - if self.nsizes != NULL: - free(self.nsizes) - - def get_indices(self): """ Return Neighbors indices. @@ -594,6 +610,8 @@ cdef class NSResults(object): cdef ns_int i, nid, size, j, beadid cdef rvec ref, other, dx cdef real dist + cdef real[:, ::1] ref_coords = self.ref_coords + cdef real[:, ::1] grid_coords = self.grid_coords if self.distances is None: distances = [] @@ -602,11 +620,11 @@ cdef class NSResults(object): size = self.nsizes[nid] tmp_values = np.empty((size), dtype=np.float32) - ref = &self.ref_coords[nid, 0] + ref = &ref_coords[nid, 0] for i in range(size): beadid = self.nids[nid][i] - other = &self.grid_coords[beadid, 0] + other = &grid_coords[beadid, 0] tmp_values[i] = self.box.fast_distance(ref, other) @@ -616,6 +634,577 @@ cdef class NSResults(object): return self.distances +cdef class NSResults2(object): + cdef readonly real cutoff + cdef ns_int npairs + cdef bint debug + + cdef real[:, ::1] grid_coords # shape: size, DIM + cdef ns_int[:] search_ids + + cdef ns_int allocation_size + cdef ns_int[:, ::1] pairs# shape: pair_allocation, 2 + cdef real[:] pair_distances2 # shape: pair_allocation, 2 + + cdef list indices_buffer + cdef list coordinates_buffer + cdef list distances_buffer + cdef np.ndarray pairs_buffer + cdef np.ndarray pair_distances_buffer + + def __init__(self, real cutoff, real[:, ::1]grid_coords,ns_int[:] search_ids, debug=False): + + self.debug = debug + self.cutoff = cutoff + self.grid_coords = grid_coords + self.search_ids = search_ids + + # Preallocate memory + self.allocation_size = grid_coords.shape[0] + self.pairs = np.empty((self.allocation_size, 2), dtype=np.int) + self.pair_distances2 = np.empty(self.allocation_size, dtype=np.float32) + self.npairs = 0 + + # Buffer + self.indices_buffer = None + self.coordinates_buffer = None + self.distances_buffer = None + self.pair_distances_buffer = None + + cdef int add_neighbors(self, ns_int beadid_i, ns_int beadid_j, real distance2) nogil except 0: + # Reallocate memory if needed + if self.npairs >= self.allocation_size: + # We need to reallocate memory + with gil: + self.resize(self.allocation_size + (self.allocation_size * 0.5 + 1)) + + # Actually store pair and distance squared + if beadid_i < beadid_j: + self.pairs[self.npairs, 0] = beadid_i + self.pairs[self.npairs, 1] = beadid_j + else: + self.pairs[self.npairs, 1] = beadid_i + self.pairs[self.npairs, 0] = beadid_j + self.pair_distances2[self.npairs] = distance2 + self.npairs += 1 + + return self.npairs + + cdef resize(self, ns_int new_size): + cdef ns_int[:, ::1] pair_buffer + cdef real[:] dists_buffer + cdef ns_int i, j + + if new_size < self.npairs: + # Silently ignored the request + return + + if self.allocation_size >= new_size: + if self.debug: + print("NSresults reallocation requested but not needed ({} requested but {} already allocated)".format(new_size, self.allocation_size)) + return + + self.allocation_size = new_size + + if self.debug: + print("NSresults reallocated to {} pairs".format(self.allocation_size)) + + # Note: np.empty + update is faster than resize + # Allocating memory + pair_buffer = self.pairs + self.pairs = np.empty((self.allocation_size, 2), dtype=np.int) + + + dists_buffer = self.pair_distances2 + self.pair_distances2 = np.empty(self.allocation_size, dtype=np.float32) + + + # Update values + with nogil: + for i in range(self.npairs): + for j in range(2): + self.pairs[i, j] = pair_buffer[i, j] + + self.pair_distances2[i] = dists_buffer[i] + + def get_pairs(self): + if self.pairs_buffer is None: + self.pairs_buffer = np.array(self.pairs[:self.npairs]) + return self.pairs_buffer + + def get_pair_distances(self): + if self.pair_distances_buffer is None: + self.pair_distances_buffer = np.sqrt(self.pair_distances2[:self.npairs]) + return self.pair_distances_buffer + + cdef create_buffers(self): + cdef ns_int i, beadid_i, beadid_j + cdef real dist2 + cdef real[:] coord_i, coord_j + from collections import defaultdict + + indices_buffer = defaultdict(list) + coords_buffer = defaultdict(list) + dists_buffer = defaultdict(list) + + for i in range(self.npairs): + beadid_i = self.pairs[i, 0] + beadid_j = self.pairs[i, 1] + + dist2 = self.pair_distances2[i] + coord_i = self.grid_coords[beadid_i] + coord_j = self.grid_coords[beadid_j] + + indices_buffer[beadid_i].append(beadid_j) + indices_buffer[beadid_j].append(beadid_i) + + coords_buffer[beadid_i].append(coord_j) + coords_buffer[beadid_j].append((coord_i)) + + dists_buffer[beadid_i].append(dist2) + dists_buffer[beadid_j].append(dist2) + + self.indices_buffer = [] + self.coordinates_buffer = [] + self.distances_buffer = [] + + for elm in self.search_ids: + sorted_indices = np.argsort(indices_buffer[elm]) + self.indices_buffer.append(np.array(indices_buffer[elm])[sorted_indices]) + self.coordinates_buffer.append(np.array(coords_buffer[elm])[sorted_indices]) + self.distances_buffer.append(np.sqrt(dists_buffer[elm])[sorted_indices]) + + def get_indices(self): + if self.indices_buffer is None: + self.create_buffers() + return self.indices_buffer + + def get_distances(self): + if self.distances_buffer is None: + self.create_buffers() + return self.distances_buffer + + def get_coordinates(self): + if self.coordinates_buffer is None: + self.create_buffers() + return self.coordinates_buffer + + + +cdef class NSGrid(object): + cdef bint debug + cdef readonly real cutoff + cdef ns_int size + cdef ns_int ncoords + cdef ns_int[DIM] ncells + cdef ns_int[DIM] cell_offsets + cdef real[DIM] cellsize + cdef ns_int[:] nbeads # size + cdef ns_int[:, ::1] beadids # size * beadpercell + cdef ns_int[:] cellids # ncoords + cdef ns_int[:, ::1] cellxyz + + def __init__(self, ncoords, cutoff, PBCBox box, max_size, debug=False): + cdef ns_int i, x, y, z + cdef ns_int ncellx, ncelly, ncellz, size, nbeadspercell + cdef real bbox_vol + self.debug = debug + + self.ncoords = ncoords + + # Calculate best cutoff + self.cutoff = cutoff + bbox_vol = box.c_pbcbox.box[XX][XX] * box.c_pbcbox.box[YY][YY] * box.c_pbcbox.box[YY][YY] + size = bbox_vol/cutoff**3 + nbeadspercell = ncoords/size + while bbox_vol/self.cutoff**3 > max_size: + self.cutoff *= 1.2 + + + for i in range(DIM): + self.ncells[i] = (box.c_pbcbox.box[i][i] / self.cutoff) + if self.ncells[i] == 0: + self.ncells[i] = 1 + self.cellsize[i] = box.c_pbcbox.box[i][i] / self.ncells[i] + self.size = self.ncells[XX] * self.ncells[YY] * self.ncells[ZZ] + + if self.debug: + print("NSGrid: Requested cutoff: {:.3f} (Ncells={}, Avg # of beads per cell={}), Optimized cutoff= {:.3f} (Ncells={}, Avg # of beads per cell={})".format( + cutoff, size, nbeadspercell, + self.cutoff, self.size, (ncoords / self.size) + )) + print("NSGrid: Size={}x{}x{}={}".format(self.ncells[XX], self.ncells[YY], self.ncells[ZZ], self.size)) + + self.cell_offsets[XX] = 0 + self.cell_offsets[YY] = self.ncells[XX] + self.cell_offsets[ZZ] = self.ncells[XX] * self.ncells[YY] + + # Allocate memory + self.nbeads = np.zeros(self.size, dtype=np.int) + self.cellids = np.empty(self.ncoords, dtype=np.int) + #self.cellxyz = np.empty((self.size, DIM), dtype=np.int) + + #i = 0 + #for z in range(self.ncells[ZZ]): + # for y in range(self.ncells[YY]): + # for x in range(self.ncells[XX]): + # self.cellxyz[i, XX] = x + # self.cellxyz[i, YY] = y + # self.cellxyz[i, ZZ] = z + # i += 1 + + # Number of maximum bead per cell is not known, so wa can not allocate memory + self.beadids = None + + cdef ns_int coord2cellid(self, rvec coord) nogil: + return (coord[ZZ] / self.cellsize[ZZ]) * (self.ncells[XX] * self.ncells[YY]) +\ + (coord[YY] / self.cellsize[YY]) * self.ncells[XX] + \ + (coord[XX] / self.cellsize[XX]) + + cdef bint cellid2cellxyz(self, ns_int cellid, ivec cellxyz) nogil: + if cellid < 0: + return False + if cellid >= self.size: + return False + + cellxyz[ZZ] = (cellid / self.cell_offsets[ZZ]) + cellid -= cellxyz[ZZ] * self.cell_offsets[ZZ] + + cellxyz[YY] = (cellid / self.cell_offsets[YY]) + cellxyz[XX] = cellid - cellxyz[YY] * self.cell_offsets[YY] + + return True + + cdef fill_grid(self, real[:, ::1] coords): + cdef ns_int i, cellindex = -1, nbeads_max = 0 + cdef ns_int ncoords = coords.shape[0] + + cdef ns_int *beadcounts = NULL + + # Allocate memory + beadcounts = PyMem_Malloc(sizeof(ns_int) * self.size) + if not beadcounts: + raise MemoryError("Could not allocate memory for bead count buffer") + + + with nogil: + # Initialize buffers + for i in range(self.size): + beadcounts[i] = 0 + + # First loop: find cellindex for each bead + for i in range(ncoords): + cellindex = self.coord2cellid(&coords[i, 0]) + + self.nbeads[cellindex] += 1 + self.cellids[i] = cellindex + + if self.nbeads[cellindex] > nbeads_max: + nbeads_max = self.nbeads[cellindex] + + # Allocate memory + with gil: + self.beadids = np.empty((self.size, nbeads_max), dtype=np.int) + + # Second loop: fill grid + for i in range(ncoords): + + # Add bead to grid cell + cellindex = self.cellids[i] + self.beadids[cellindex, beadcounts[cellindex]] = i + beadcounts[cellindex] += 1 + + + # Now we can free the allocation buffer + PyMem_Free(beadcounts) + + + +# Python interface +cdef class FastNS2(object): + cdef bint debug + cdef PBCBox box + cdef readonly real[:, ::1] coords + cdef real[:, ::1] coords_bbox + cdef readonly real cutoff + cdef bint prepared + cdef NSGrid grid + + def __init__(self, u, cutoff, coords=None, prepare=True, debug=False, max_gridsize=5000): + import MDAnalysis as mda + from MDAnalysis.lib.mdamath import triclinic_vectors + + self.debug = debug + + if not isinstance(u, mda.Universe): + raise TypeError("FastNS class must be initialized with a valid MDAnalysis.Universe instance") + box = triclinic_vectors(u.dimensions) + + if box.shape != (3, 3): + raise ValueError("Box must be provided as triclinic_dimensions (a 3x3 numpy.ndarray of unit cell vectors") + + self.box = PBCBox(box) + + + if coords is None: + coords = u.atoms.positions + + self.coords = coords.copy() + + self.coords_bbox = self.box.fast_put_atoms_in_bbox(coords) + + if not prepare: + return + + if self.cutoff < 0: + raise ValueError("Cutoff must be positive!") + if self.cutoff * self.cutoff > self.box.c_pbcbox.max_cutoff2: + raise ValueError("Cutoff greater than maximum cutoff ({:.3f}) given the PBC") + self.cutoff = cutoff + + self.grid = NSGrid(self.coords_bbox.shape[0], cutoff, self.box, max_gridsize, debug=debug) + self.prepared = False + + if prepare: + self.prepare() + + + def prepare(self, force=False): + if self.prepared and not force: + return + + self.grid.fill_grid(self.coords_bbox) + + self.prepared = True + + def search(self, search_ids, min_size_increment=10): + cdef ns_int[:] search_ids_view, neighbors_view + cdef ns_int nid, i, j, size_search + cdef NSResults2 results + + cdef ns_int current_beadid + cdef rvec current_coords + + cdef ns_int cellx, celly, cellz, cellindex_adjacent + cdef ivec cellxyz, debug_cellxyz + + cdef real[:, ::1] search_coords + + #Temp stuff + cdef ns_int size = self.coords_bbox.shape[0] + cdef ns_int d, m + cdef ns_int xi, yi, zi, bid, i_bead + cdef real d2 + cdef rvec shifted_coord, dx, neighbor_coord, corrected_coords + + cdef bint already_checked[27] + cdef bint skip + cdef ns_int nchecked = 0, icheck + cdef ns_int cellindex + cdef ns_int cell_offset + + cdef ns_int[:] neighbor_ids + cdef ns_int[:] nneighbors + cdef ns_int nid_offset + + cdef real cutoff2 = self.cutoff * self.cutoff + + cdef ns_int[:] checked + + cdef ns_int npairs = 0 + cdef ns_int guessed_size = 0 + cdef bint first_ns = True + + import sys + flush = sys.stdout.flush + + if not self.prepared: + self.prepare() + + + search_ids_view = search_ids + size_search = search_ids.shape[0] + + checked = np.zeros(size, dtype=np.int) + + results = NSResults2(self.cutoff, self.coords_bbox, search_ids, self.debug) + + with nogil: + for i in range(size_search): + current_beadid = search_ids_view[i] + + cellindex = self.grid.cellids[current_beadid] + self.grid.cellid2cellxyz(cellindex, cellxyz) + + if self.debug: + + with gil: + print("Checking neighbors for bead #{} ({:.3f},{:.3f},{:.3f})->rect({:.3f},{:.3f},{:.3f}) - cell[{},{},{}]:" .format( + current_beadid, + self.coords[current_beadid, XX], self.coords[current_beadid, YY], self.coords[current_beadid, ZZ], + self.coords_bbox[current_beadid, XX], self.coords_bbox[current_beadid, YY], self.coords_bbox[current_beadid, ZZ], + cellxyz[XX], cellxyz[YY], cellxyz[ZZ])) + + + for xi in range(DIM): + if xi == 0: + if self.coords_bbox[current_beadid, XX] - self.cutoff > self.grid.cellsize[XX] * cellxyz[XX]: + if self.debug: + with gil: + print("\n#-> Bead X={:.3f}, Cell X={:.3f}, cutoff={:.3f} -> -X shift is ignored".format(self.coords_bbox[current_beadid, XX], self.grid.cellsize[XX] * cellxyz[XX], self.cutoff)) + continue + + if xi == 2: + if self.coords_bbox[current_beadid, XX] + self.cutoff < self.grid.cellsize[XX] * (cellxyz[XX] + 1): + if self.debug: + with gil: + print("\n#-> Bead X={:.3f}, Next cell X={:.3f}, cutoff={:.3f} -> +X shift is ignored".format(self.coords_bbox[current_beadid, XX], self.grid.cellsize[XX] * (cellxyz[XX] + 1), self.cutoff)) + continue + + for yi in range(DIM): + if xi == 0: + if self.coords_bbox[current_beadid, YY] - self.cutoff > self.grid.cellsize[YY] * cellxyz[YY]: + if self.debug: + with gil: + print("\n#-> Bead Y={:.3f}, Cell Y={:.3f}, cutoff={:.3f} -> -Y shift is ignored".format(self.coords_bbox[current_beadid, YY], self.grid.cellsize[YY] * cellxyz[YY], self.cutoff)) + continue + + if xi == 2: + if self.coords_bbox[current_beadid, YY] + self.cutoff < self.grid.cellsize[YY] * (cellxyz[YY] + 1): + if self.debug: + with gil: + print("\n#-> Bead Y={:.3f}, Next cell Y={:.3f}, cutoff={:.3f} -> +Y shift is ignored".format(self.coords_bbox[current_beadid, YY], self.grid.cellsize[YY] * (cellxyz[YY] + 1), self.cutoff)) + continue + + + for zi in range(DIM): + if xi == 0: + if self.coords_bbox[current_beadid, ZZ] - self.cutoff > self.grid.cellsize[ZZ] * cellxyz[ZZ]: + if self.debug: + with gil: + print("\n#-> Bead Z={:.3f}, Cell Z={:.3f}, cutoff={:.3f} -> -Z shift is ignored".format(self.coords_bbox[current_beadid, ZZ], self.grid.cellsize[ZZ] * cellxyz[ZZ], self.cutoff)) + continue + + if xi == 2: + if self.coords_bbox[current_beadid, ZZ] + self.cutoff < self.grid.cellsize[ZZ] * (cellxyz[ZZ] + 1): + if self.debug: + with gil: + print("\n#-> Bead Z={:.3f}, Next cell Z={:.3f}, cutoff={:.3f} -> +Z shift is ignored".format(self.coords_bbox[current_beadid, ZZ], self.grid.cellsize[XX] * (cellxyz[ZZ] + 1), self.cutoff)) + continue + + # Calculate and/or reinitialize shifted coordinates + shifted_coord[XX] = self.coords[current_beadid, XX] + (xi - 1) * self.grid.cellsize[XX] + shifted_coord[YY] = self.coords[current_beadid, YY] + (yi - 1) * self.grid.cellsize[YY] + shifted_coord[ZZ] = self.coords[current_beadid, ZZ] + (zi - 1) * self.grid.cellsize[ZZ] + + # Make sure the shifted coordinates is inside the brick-shaped box + for m in range(DIM - 1, -1, -1): + + while shifted_coord[m] < 0: + for d in range(m+1): + shifted_coord[d] += self.box.c_pbcbox.box[m][d] + + + while shifted_coord[m] >= self.box.c_pbcbox.box[m][m]: + for d in range(m+1): + shifted_coord[d] -= self.box.c_pbcbox.box[m][d] + + # Get the cell index corresponding to the coord + cellindex = self.grid.coord2cellid(shifted_coord) + + if self.debug: + self.grid.cellid2cellxyz(cellindex, debug_cellxyz) + with gil: + dist_shift = self.box.fast_distance(&self.coords[current_beadid, XX], shifted_coord) + grid_shift = np.array([(xi - 1) * self.grid.cellsize[XX], + (yi - 1) * self.grid.cellsize[YY], + (zi - 1) * self.grid.cellsize[ZZ]]) + print("-> Checking cell#{} ({},{},{}) for neighbors (dshift={:.3f}, grid_shift=({:.3f},{:.3f},{:.3f}->{:.3f})".format( + cellindex, + debug_cellxyz[XX], debug_cellxyz[YY], debug_cellxyz[ZZ], + dist_shift, + grid_shift[XX], grid_shift[YY], grid_shift[ZZ], + np.sqrt(np.sum(grid_shift**2)) + )) + if dist_shift > 2* self.cutoff: + print(" \_ This cell should be ignored!") + else: + print(" \_ This cell is needed") + + for j in range(self.grid.nbeads[cellindex]): + bid = self.grid.beadids[cellindex, j] + + if checked[bid] != 0: + continue + + d2 = self.box.fast_distance2(&self.coords_bbox[current_beadid, XX], &self.coords_bbox[bid, XX]) + + if self.debug and False: + self.grid.cellid2cellxyz(cellindex, debug_cellxyz) + with gil: + print( + "Beads #{} (cell[{},{},{}]-coords[{:.3f},{:.3f},{:.3f}]) and #{} (cell[{},{},{}]-coords[{:.3f},{:.3f},{:.3f}]) are tested (d2={:.3f})".format( + current_beadid, + cellxyz[XX], cellxyz[YY], cellxyz[ZZ], + self.coords_bbox[current_beadid, XX], + self.coords_bbox[current_beadid, YY], + self.coords_bbox[current_beadid, ZZ], + bid, + debug_cellxyz[XX], debug_cellxyz[YY], debug_cellxyz[ZZ], + self.coords_bbox[bid, XX], self.coords_bbox[bid, YY], + self.coords_bbox[bid, ZZ], + d2)) + + if d2 < cutoff2: + + if d2 < EPSILON: + continue + + if self.debug: + self.grid.cellid2cellxyz(cellindex, debug_cellxyz) + with gil: + self.box.fast_pbc_dx(&self.coords[current_beadid, XX], &self.coords[bid, XX], dx) + dx_py = np.array([dx[XX], dx[YY], dx[ZZ]]) + + print(" \_ Neighbor found: bead#{} (cell[{},{},{}]) -> dx={} -> d={:.3f}".format(bid, debug_cellxyz[XX], debug_cellxyz[YY], debug_cellxyz[ZZ], dx_py, np.sqrt(d2))) + results.add_neighbors(current_beadid, bid, d2) + npairs += 1 + + if first_ns: + guessed_size = ((npairs + 1) * 0.75 * size_search) + with gil: + if self.debug: + print("Neighbors of first beads= {} -> Approximated total of pairs={}".format(npairs, guessed_size)) + results.resize(guessed_size) + + first_ns = False + checked[current_beadid] = 1 + + if self.debug: + print("Total number of pairs={}".format(npairs)) + + ref_bead = 13937 + beads = np.array([4398, 4401, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, 47451]) - 1 + for bid in beads: + self.box.fast_pbc_dx(&self.coords[ref_bead, XX], &self.coords[bid, XX], dx) + dx_py = np.array([dx[XX], dx[YY], dx[ZZ]]) + self.box.fast_pbc_dx(&self.coords_bbox[ref_bead, XX], &self.coords_bbox[bid, XX], dx) + rect_dx_py = np.array([dx[XX], dx[YY], dx[ZZ]]) + self.grid.cellid2cellxyz(self.grid.coord2cellid(&self.coords_bbox[bid, XX]), cellxyz) + print("Bead #{} ({:.3f},{:.3f},{:.3f})->rect({:.3f},{:.3f},{:.3f}) - cell[{},{},{}]: dx=[{:.3f},{:.3f},{:.3f}] -> dist: {:.3f} ({}) - rect_dist: {:.3f} ({})".format( + bid, + self.coords[bid, XX], self.coords[bid, YY], self.coords[bid,ZZ], + self.coords_bbox[bid, XX], self.coords_bbox[bid, YY], self.coords_bbox[bid,ZZ], + cellxyz[XX], cellxyz[YY], cellxyz[ZZ], + dx[XX], dx[YY], dx[ZZ], + np.sqrt(np.sum(dx_py**2)), + self.box.fast_distance(&self.coords[ref_bead, XX], &self.coords[bid, XX]) <= self.cutoff, + np.sqrt(np.sum(rect_dx_py**2)), + np.sqrt(np.sum(rect_dx_py**2)) <= self.cutoff + )) + + return results + # Python interface cdef class FastNS(object): @@ -651,14 +1240,16 @@ cdef class FastNS(object): def __dealloc__(self): cdef ns_int i + # Deallocate NS grid if self.grid.nbeads != NULL: free(self.grid.nbeads) - for i in range(self.grid.size): - if self.grid.beadids[i] != NULL: - free(self.grid.beadids[i]) - free(self.grid.beadids) + if self.grid.beadids != NULL: + free(self.grid.beadids) + + if self.grid.cellids != NULL: + free(self.grid.cellids) self.grid.size = 0 @@ -680,11 +1271,11 @@ cdef class FastNS(object): def prepare(self, force=False): cdef ns_int i, cellindex = -1 - cdef ns_int *allocated_size = NULL cdef ns_int ncoords = self.coords.shape[0] - cdef ns_int allocation_guess cdef rvec *coords = &self.coords_bbox[0, 0] + cdef ns_int *beadcounts = NULL + if self.prepared and not force: print("NS already prepared, nothing to do!") @@ -704,57 +1295,55 @@ cdef class FastNS(object): self.grid.cellsize[i] = self.box.c_pbcbox.box[i][i] / self.grid.ncells[i] self.grid.size = self.grid.ncells[XX] * self.grid.ncells[YY] * self.grid.ncells[ZZ] - # This is just a guess on how much memory we might for each grid cell: - # we just assume an average bead density and we take four times this density just to be safe - allocation_guess = (4 * (ncoords / self.grid.size + 1)) + # Allocate memory for temporary buffers + self.grid.cellids = malloc(sizeof(ns_int) * ncoords) + if self.grid.cellids == NULL: + with gil: + raise MemoryError("Could not allocate memory for cell ids") + + beadcounts = malloc(sizeof(ns_int) * self.grid.size) + if beadcounts == NULL: + with gil: + raise MemoryError("Could not allocate memory for bead count buffer") - # Allocate memory for the grid self.grid.nbeads = malloc(sizeof(ns_int) * self.grid.size) if self.grid.nbeads == NULL: with gil: raise MemoryError("Could not allocate memory for NS grid") - # Allocate memory from temporary allocation counter - allocated_size = malloc(sizeof(ns_int) * self.grid.size) - if allocated_size == NULL: - # No need to free grid.nbeads as it will be freed by destroy_nsgrid called by __dealloc___ - with gil: - raise MemoryError("Could not allocate memory for allocation buffer") - - # Pre-allocate some memory for grid cells + # Initialize buffers for i in range(self.grid.size): self.grid.nbeads[i] = 0 - allocated_size[i] = allocation_guess + beadcounts[i] = 0 - self.grid.beadids = malloc(sizeof(ns_int *) * self.grid.size) - if self.grid.beadids == NULL: - with gil: - raise MemoryError("Could not allocate memory for grid cells") - for i in range(self.grid.size): - self.grid.beadids[i] = malloc(sizeof(ns_int) * allocated_size[i]) - if self.grid.beadids[i] == NULL: - with gil: - raise MemoryError("Could not allocate memory for grid cell") - - # Populate grid cells using the coordinates (ie do the heavy work) + # First loop: find cellindex for each bead for i in range(ncoords): cellindex = (coords[i][ZZ] / self.grid.cellsize[ZZ]) * (self.grid.ncells[XX] * self.grid.ncells[YY]) +\ (coords[i][YY] / self.grid.cellsize[YY]) * self.grid.ncells[XX] + \ (coords[i][XX] / self.grid.cellsize[XX]) - self.grid.beadids[cellindex][self.grid.nbeads[cellindex]] = i self.grid.nbeads[cellindex] += 1 + self.grid.cellids[i] = cellindex + + if self.grid.nbeads[cellindex] > self.grid.beadpercell: + self.grid.beadpercell = self.grid.nbeads[cellindex] + + # Allocate memory and set offsets + self.grid.beadids = malloc(sizeof(ns_int) * self.grid.size * self.grid.beadpercell) + if self.grid.beadids == NULL: + with gil: + raise MemoryError("Could not allocate memory for NS grid bead ids") - # We need to allocate more memory (simply double the amount of memory as - # 1. it should barely be needed - # 2. the size should stay fairly reasonable - if self.grid.nbeads[cellindex] >= allocated_size[cellindex]: - allocated_size[cellindex] *= 2 - self.grid.beadids[cellindex] = realloc( self.grid.beadids[cellindex], sizeof(ns_int) * allocated_size[cellindex]) + # Second loop: fill grid + for i in range(ncoords): + cellindex = self.grid.cellids[i] + self.grid.beadids[cellindex * self.grid.beadpercell + beadcounts[cellindex]] = i + beadcounts[cellindex] += 1 # Now we can free the allocation buffer - free(allocated_size) + free(beadcounts) + self.prepared = True @@ -787,7 +1376,6 @@ cdef class FastNS(object): if holder == NULL: raise MemoryError("Could not allocate memory to run NS core") - results = NSResults(self.box, self.cutoff) results.populate(holder, self.coords, search_coords_view) @@ -795,3 +1383,101 @@ cdef class FastNS(object): free_neighborhood_holder(holder) return results + + def search2(self, search_ids): + cdef ns_int[:] search_ids_view, neighbors_view + cdef ns_int nid, i, j, size_search + cdef ns_neighborhood_holder *holder + cdef ns_neighborhood *neighborhood + + cdef ns_int current_beadid + cdef rvec current_coords + + #Temp stuff + cdef ns_int size = self.coords_bbox.shape[0] + cdef ns_int d, m + cdef ns_int xi, yi, zi, bid, i_bead + cdef real d2 + cdef rvec shifted_coord, dx, neighbor_coord, corrected_coords + + cdef bint already_checked[27] + cdef bint skip + cdef ns_int nchecked = 0, icheck + cdef ns_int cell_index + cdef ns_int cell_offset + + cdef ns_int[:] neighbor_ids + cdef ns_int[:] nneighbors + cdef ns_int nid_offset + + cdef real cutoff2 = self.cutoff * self.cutoff + + + if not self.prepared: + self.prepare() + + nid_offset = self.grid.beadpercell * 27 + neighbor_ids = np.empty(size * nid_offset, dtype=np.int) + nneighbors = np.zeros(size, dtype=np.int) + + + search_ids_view = search_ids + size_search = search_ids_view.shape[0] + + with nogil: + + for current_beadid in range(size): + + for zi in range(3): + for yi in range(3): + for xi in range(3): + # Calculate and/or reinitialize shifted coordinates + shifted_coord[XX] = self.coords_bbox[current_beadid, XX] + (xi - 1) * self.grid.cellsize[XX] + shifted_coord[YY] = self.coords_bbox[current_beadid, YY] + (yi - 1) * self.grid.cellsize[YY] + shifted_coord[ZZ] = self.coords_bbox[current_beadid, ZZ] + (zi - 1) * self.grid.cellsize[ZZ] + + # Make sure the shifted coordinates is inside the brick-shaped box + for m in range(DIM - 1, -1, -1): + + while shifted_coord[m] < 0: + for d in range(m+1): + shifted_coord[d] += self.box.c_pbcbox.box[m][d] + + + while shifted_coord[m] >= self.box.c_pbcbox.box[m][m]: + for d in range(m+1): + shifted_coord[d] -= self.box.c_pbcbox.box[m][d] + + # Get the cell index corresponding to the coord + cell_index = (shifted_coord[ZZ] / self.grid.cellsize[ZZ]) * (self.grid.ncells[XX] * self.grid.ncells[YY]) +\ + (shifted_coord[YY] / self.grid.cellsize[YY]) * self.grid.ncells[XX] + \ + (shifted_coord[XX] / self.grid.cellsize[XX]) + + # Search for neighbors inside this cell + for i_bead in range(self.grid.nbeads[cell_index]): + cell_offset = cell_index * self.grid.beadpercell + i_bead + + bid = self.grid.beadids[cell_offset] + + if bid <= current_beadid: + continue + + self.box.fast_pbc_dx(&self.coords_bbox[current_beadid, XX], &self.coords_bbox[bid, XX], dx) + + d2 = rvec_norm2(dx) + + if d2 < cutoff2: + neighbor_ids[current_beadid * nid_offset + nneighbors[current_beadid]] = bid + nneighbors[current_beadid] += 1 + + neighbor_ids[bid * nid_offset + nneighbors[bid]] = current_beadid + nneighbors[bid] += 1 + + + + + results = [] + #for bid in search_ids: + # results.append(np.asarray(neighbor_ids[bid * nid_offset : bid * nid_offset + nneighbors[bid]])) + + return results From 86505097885e31040cea2a7f5360b39f0831a6bd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Buchoux?= Date: Mon, 16 Jul 2018 11:45:11 +0200 Subject: [PATCH 20/47] Unneeded code removed from c_gridsearch.pyx --- package/MDAnalysis/lib/c_gridsearch.pyx | 526 ++++++++---------- .../MDAnalysisTests/lib/test_gridsearch.py | 147 +++-- 2 files changed, 307 insertions(+), 366 deletions(-) diff --git a/package/MDAnalysis/lib/c_gridsearch.pyx b/package/MDAnalysis/lib/c_gridsearch.pyx index ffdb91fed7d..a1f519eac37 100644 --- a/package/MDAnalysis/lib/c_gridsearch.pyx +++ b/package/MDAnalysis/lib/c_gridsearch.pyx @@ -55,6 +55,7 @@ from libc.math cimport sqrt import numpy as np cimport numpy as np +cimport cython ctypedef np.int_t ns_int ctypedef np.float32_t real @@ -473,7 +474,7 @@ cdef ns_neighborhood_holder *ns_core(real[:, ::1] refcoords, return holder -cdef class NSResults(object): +cdef class NSResultsOld(object): """ Class used to store results returned by `MDAnalysis.lib.grid.FastNS.search` """ @@ -634,7 +635,7 @@ cdef class NSResults(object): return self.distances -cdef class NSResults2(object): +cdef class NSResults(object): cdef readonly real cutoff cdef ns_int npairs cdef bint debug @@ -643,8 +644,8 @@ cdef class NSResults2(object): cdef ns_int[:] search_ids cdef ns_int allocation_size - cdef ns_int[:, ::1] pairs# shape: pair_allocation, 2 - cdef real[:] pair_distances2 # shape: pair_allocation, 2 + cdef ipair *pairs # shape: pair_allocation + cdef real *pair_distances2 # shape: pair_allocation cdef list indices_buffer cdef list coordinates_buffer @@ -660,81 +661,96 @@ cdef class NSResults2(object): self.search_ids = search_ids # Preallocate memory - self.allocation_size = grid_coords.shape[0] - self.pairs = np.empty((self.allocation_size, 2), dtype=np.int) - self.pair_distances2 = np.empty(self.allocation_size, dtype=np.float32) + self.allocation_size = search_ids.shape[0] + 1 + self.pairs = PyMem_Malloc(sizeof(ipair) * self.allocation_size) + if not self.pairs: + MemoryError("Could not allocate memory for NSResults.pairs ({} bits requested)".format(sizeof(ipair) * self.allocation_size)) + self.pair_distances2 = PyMem_Malloc(sizeof(real) * self.allocation_size) + if not self.pair_distances2: + raise MemoryError("Could not allocate memory for NSResults.pair_distances2 ({} bits requested)".format(sizeof(real) * self.allocation_size)) self.npairs = 0 # Buffer self.indices_buffer = None self.coordinates_buffer = None self.distances_buffer = None + self.pairs_buffer = None self.pair_distances_buffer = None - cdef int add_neighbors(self, ns_int beadid_i, ns_int beadid_j, real distance2) nogil except 0: + def __dealloc__(self): + PyMem_Free(self.pairs) + PyMem_Free(self.pair_distances2) + + cdef int add_neighbors(self, ns_int beadid_i, ns_int beadid_j, real distance2) nogil: + # Important: If this function returns 0, it means that memory allocation failed + # Reallocate memory if needed if self.npairs >= self.allocation_size: # We need to reallocate memory - with gil: - self.resize(self.allocation_size + (self.allocation_size * 0.5 + 1)) + if self.resize(self.allocation_size + (self.allocation_size * 0.5 + 1)) == 0: + return 0 # Actually store pair and distance squared if beadid_i < beadid_j: - self.pairs[self.npairs, 0] = beadid_i - self.pairs[self.npairs, 1] = beadid_j + self.pairs[self.npairs][0] = beadid_i + self.pairs[self.npairs][1] = beadid_j else: - self.pairs[self.npairs, 1] = beadid_i - self.pairs[self.npairs, 0] = beadid_j + self.pairs[self.npairs][1] = beadid_i + self.pairs[self.npairs][0] = beadid_j self.pair_distances2[self.npairs] = distance2 self.npairs += 1 return self.npairs - cdef resize(self, ns_int new_size): - cdef ns_int[:, ::1] pair_buffer - cdef real[:] dists_buffer - cdef ns_int i, j + cdef int resize(self, ns_int new_size) nogil: + # Important: If this function returns 0, it means that memory allocation failed if new_size < self.npairs: # Silently ignored the request - return + return 1 if self.allocation_size >= new_size: if self.debug: - print("NSresults reallocation requested but not needed ({} requested but {} already allocated)".format(new_size, self.allocation_size)) - return + with gil: + print("NSresults: Reallocation requested but not needed ({} requested but {} already allocated)".format(new_size, self.allocation_size)) + return 1 self.allocation_size = new_size if self.debug: - print("NSresults reallocated to {} pairs".format(self.allocation_size)) + with gil: + print("NSresults: Reallocated to {} pairs".format(self.allocation_size)) - # Note: np.empty + update is faster than resize # Allocating memory - pair_buffer = self.pairs - self.pairs = np.empty((self.allocation_size, 2), dtype=np.int) - - - dists_buffer = self.pair_distances2 - self.pair_distances2 = np.empty(self.allocation_size, dtype=np.float32) + with gil: + self.pairs = PyMem_Realloc(self.pairs, sizeof(ipair) * self.allocation_size) + self.pair_distances2 = PyMem_Realloc(self.pair_distances2, sizeof(real) * self.allocation_size) + if not self.pairs: + return 0 - # Update values - with nogil: - for i in range(self.npairs): - for j in range(2): - self.pairs[i, j] = pair_buffer[i, j] + if not self.pair_distances2: + return 0 - self.pair_distances2[i] = dists_buffer[i] + return 1 def get_pairs(self): + cdef ns_int i + if self.pairs_buffer is None: - self.pairs_buffer = np.array(self.pairs[:self.npairs]) + self.pairs_buffer = np.empty((self.npairs, 2), dtype=np.int) + for i in range(self.npairs): + self.pairs_buffer[i, 0] = self.pairs[i][0] + self.pairs_buffer[i, 1] = self.pairs[i][1] return self.pairs_buffer def get_pair_distances(self): + cdef ns_int i if self.pair_distances_buffer is None: - self.pair_distances_buffer = np.sqrt(self.pair_distances2[:self.npairs]) + self.pair_distances_buffer = np.empty(self.npairs, dtype=np.float32) + for i in range(self.npairs): + self.pair_distances_buffer[i] = self.pair_distances2[i] + self.pair_distances_buffer = np.sqrt(self.pair_distances_buffer) return self.pair_distances_buffer cdef create_buffers(self): @@ -748,8 +764,8 @@ cdef class NSResults2(object): dists_buffer = defaultdict(list) for i in range(self.npairs): - beadid_i = self.pairs[i, 0] - beadid_j = self.pairs[i, 1] + beadid_i = self.pairs[i][0] + beadid_j = self.pairs[i][1] dist2 = self.pair_distances2[i] coord_i = self.grid_coords[beadid_i] @@ -799,10 +815,10 @@ cdef class NSGrid(object): cdef ns_int[DIM] ncells cdef ns_int[DIM] cell_offsets cdef real[DIM] cellsize - cdef ns_int[:] nbeads # size - cdef ns_int[:, ::1] beadids # size * beadpercell - cdef ns_int[:] cellids # ncoords - cdef ns_int[:, ::1] cellxyz + cdef ns_int nbeads_per_cell + cdef ns_int *nbeads # size + cdef ns_int *beadids # size * nbeads_per_cell + cdef ns_int *cellids # ncoords def __init__(self, ncoords, cutoff, PBCBox box, max_size, debug=False): cdef ns_int i, x, y, z @@ -820,7 +836,6 @@ cdef class NSGrid(object): while bbox_vol/self.cutoff**3 > max_size: self.cutoff *= 1.2 - for i in range(DIM): self.ncells[i] = (box.c_pbcbox.box[i][i] / self.cutoff) if self.ncells[i] == 0: @@ -840,21 +855,22 @@ cdef class NSGrid(object): self.cell_offsets[ZZ] = self.ncells[XX] * self.ncells[YY] # Allocate memory - self.nbeads = np.zeros(self.size, dtype=np.int) - self.cellids = np.empty(self.ncoords, dtype=np.int) - #self.cellxyz = np.empty((self.size, DIM), dtype=np.int) - - #i = 0 - #for z in range(self.ncells[ZZ]): - # for y in range(self.ncells[YY]): - # for x in range(self.ncells[XX]): - # self.cellxyz[i, XX] = x - # self.cellxyz[i, YY] = y - # self.cellxyz[i, ZZ] = z - # i += 1 - - # Number of maximum bead per cell is not known, so wa can not allocate memory - self.beadids = None + self.nbeads = PyMem_Malloc(sizeof(ns_int) * self.size) + if not self.nbeads: + raise MemoryError("Could not allocate memory from NSGrid.nbeads ({} bits requested)".format(sizeof(ns_int) * self.size)) + self.beadids = NULL + self.cellids = PyMem_Malloc(sizeof(ns_int) * self.ncoords) + if not self.cellids: + raise MemoryError("Could not allocate memory from NSGrid.cellids ({} bits requested)".format(sizeof(ns_int) * self.ncoords)) + self.nbeads_per_cell = 0 + + for i in range(self.size): + self.nbeads[i] = 0 + + def __dealloc__(self): + PyMem_Free(self.nbeads) + PyMem_Free(self.beadids) + PyMem_Free(self.cellids) cdef ns_int coord2cellid(self, rvec coord) nogil: return (coord[ZZ] / self.cellsize[ZZ]) * (self.ncells[XX] * self.ncells[YY]) +\ @@ -876,16 +892,14 @@ cdef class NSGrid(object): return True cdef fill_grid(self, real[:, ::1] coords): - cdef ns_int i, cellindex = -1, nbeads_max = 0 + cdef ns_int i, cellindex = -1 cdef ns_int ncoords = coords.shape[0] - cdef ns_int *beadcounts = NULL # Allocate memory beadcounts = PyMem_Malloc(sizeof(ns_int) * self.size) if not beadcounts: - raise MemoryError("Could not allocate memory for bead count buffer") - + raise MemoryError("Could not allocate memory for bead count buffer ({} bits requested)".format(sizeof(ns_int) * self.size)) with nogil: # Initialize buffers @@ -899,29 +913,30 @@ cdef class NSGrid(object): self.nbeads[cellindex] += 1 self.cellids[i] = cellindex - if self.nbeads[cellindex] > nbeads_max: - nbeads_max = self.nbeads[cellindex] + if self.nbeads[cellindex] > self.nbeads_per_cell: + self.nbeads_per_cell = self.nbeads[cellindex] # Allocate memory with gil: - self.beadids = np.empty((self.size, nbeads_max), dtype=np.int) + self.beadids = PyMem_Malloc(sizeof(ns_int) * self.size * self.nbeads_per_cell) #np.empty((self.size, nbeads_max), dtype=np.int) + if not self.beadids: + raise MemoryError("Could not allocate memory for NSGrid.beadids ({} bits requested)".format(sizeof(ns_int) * self.size * self.nbeads_per_cell)) # Second loop: fill grid for i in range(ncoords): # Add bead to grid cell cellindex = self.cellids[i] - self.beadids[cellindex, beadcounts[cellindex]] = i + self.beadids[cellindex * self.nbeads_per_cell + beadcounts[cellindex]] = i beadcounts[cellindex] += 1 - # Now we can free the allocation buffer PyMem_Free(beadcounts) # Python interface -cdef class FastNS2(object): +cdef class FastNS(object): cdef bint debug cdef PBCBox box cdef readonly real[:, ::1] coords @@ -977,69 +992,68 @@ cdef class FastNS2(object): self.prepared = True - def search(self, search_ids, min_size_increment=10): - cdef ns_int[:] search_ids_view, neighbors_view - cdef ns_int nid, i, j, size_search - cdef NSResults2 results + def search(self, search_ids=None, bint debug=False): + cdef ns_int i, j, size_search + cdef ns_int d, m + cdef NSResults results + cdef ns_int size = self.coords_bbox.shape[0] - cdef ns_int current_beadid + cdef ns_int current_beadid, bid cdef rvec current_coords - cdef ns_int cellx, celly, cellz, cellindex_adjacent + cdef ns_int cellindex, cellindex_adjacent, cellindex_probe cdef ivec cellxyz, debug_cellxyz cdef real[:, ::1] search_coords + cdef ns_int[:] search_ids_view - #Temp stuff - cdef ns_int size = self.coords_bbox.shape[0] - cdef ns_int d, m - cdef ns_int xi, yi, zi, bid, i_bead + cdef ns_int xi, yi, zi cdef real d2 - cdef rvec shifted_coord, dx, neighbor_coord, corrected_coords + cdef rvec shifted_coord, probe, dx - cdef bint already_checked[27] - cdef bint skip - cdef ns_int nchecked = 0, icheck - cdef ns_int cellindex - cdef ns_int cell_offset + cdef ns_int nchecked = 0 - cdef ns_int[:] neighbor_ids - cdef ns_int[:] nneighbors - cdef ns_int nid_offset - cdef real cutoff2 = self.cutoff * self.cutoff + cdef real cutoff2 = self.cutoff * self.cutoff cdef ns_int[:] checked - cdef ns_int npairs = 0 - cdef ns_int guessed_size = 0 - cdef bint first_ns = True - import sys - flush = sys.stdout.flush if not self.prepared: self.prepare() + if search_ids is None: + search_ids=np.arange(size) + elif type(search_ids) == np.int: + search_ids = np.array([search_ids,], dtype=np.int) + elif type(search_ids) != np.ndarray: + search_ids = np.array(search_ids, dtype=np.int) search_ids_view = search_ids size_search = search_ids.shape[0] checked = np.zeros(size, dtype=np.int) - results = NSResults2(self.cutoff, self.coords_bbox, search_ids, self.debug) + results = NSResults(self.cutoff, self.coords_bbox, search_ids, self.debug) + + cdef bint memory_error = False + + if self.debug and debug: + print("FastNS: Debug flag is set to True for FastNS.search()") with nogil: for i in range(size_search): + if memory_error: + break current_beadid = search_ids_view[i] cellindex = self.grid.cellids[current_beadid] self.grid.cellid2cellxyz(cellindex, cellxyz) - if self.debug: - + if self.debug and debug: with gil: - print("Checking neighbors for bead #{} ({:.3f},{:.3f},{:.3f})->rect({:.3f},{:.3f},{:.3f}) - cell[{},{},{}]:" .format( + print("FastNS: Checking neighbors for bead #{} ({:.3f},{:.3f},{:.3f})->rect({:.3f},{:.3f},{:.3f}) - cell[{},{},{}]:" .format( current_beadid, self.coords[current_beadid, XX], self.coords[current_beadid, YY], self.coords[current_beadid, ZZ], self.coords_bbox[current_beadid, XX], self.coords_bbox[current_beadid, YY], self.coords_bbox[current_beadid, ZZ], @@ -1047,56 +1061,89 @@ cdef class FastNS2(object): for xi in range(DIM): - if xi == 0: - if self.coords_bbox[current_beadid, XX] - self.cutoff > self.grid.cellsize[XX] * cellxyz[XX]: - if self.debug: - with gil: - print("\n#-> Bead X={:.3f}, Cell X={:.3f}, cutoff={:.3f} -> -X shift is ignored".format(self.coords_bbox[current_beadid, XX], self.grid.cellsize[XX] * cellxyz[XX], self.cutoff)) - continue + if memory_error: + break - if xi == 2: - if self.coords_bbox[current_beadid, XX] + self.cutoff < self.grid.cellsize[XX] * (cellxyz[XX] + 1): - if self.debug: - with gil: - print("\n#-> Bead X={:.3f}, Next cell X={:.3f}, cutoff={:.3f} -> +X shift is ignored".format(self.coords_bbox[current_beadid, XX], self.grid.cellsize[XX] * (cellxyz[XX] + 1), self.cutoff)) - continue - - for yi in range(DIM): + if not self.box.is_triclinic: + # If box is not triclinic (ie rect), when can already check if the shift can be skipped (ie cutoff inside the cell) if xi == 0: - if self.coords_bbox[current_beadid, YY] - self.cutoff > self.grid.cellsize[YY] * cellxyz[YY]: - if self.debug: + if self.coords_bbox[current_beadid, XX] - self.cutoff > self.grid.cellsize[XX] * cellxyz[XX]: + if self.debug and debug: with gil: - print("\n#-> Bead Y={:.3f}, Cell Y={:.3f}, cutoff={:.3f} -> -Y shift is ignored".format(self.coords_bbox[current_beadid, YY], self.grid.cellsize[YY] * cellxyz[YY], self.cutoff)) + print("FastNS: -> Bead X={:.3f}, Cell X={:.3f}, cutoff={:.3f} -> -X shift ignored".format( + self.coords_bbox[current_beadid, XX], + self.grid.cellsize[XX] * cellxyz[XX], + self.cutoff + )) continue if xi == 2: - if self.coords_bbox[current_beadid, YY] + self.cutoff < self.grid.cellsize[YY] * (cellxyz[YY] + 1): - if self.debug: + if self.coords_bbox[current_beadid, XX] + self.cutoff < self.grid.cellsize[XX] * (cellxyz[XX] + 1): + if self.debug and debug: with gil: - print("\n#-> Bead Y={:.3f}, Next cell Y={:.3f}, cutoff={:.3f} -> +Y shift is ignored".format(self.coords_bbox[current_beadid, YY], self.grid.cellsize[YY] * (cellxyz[YY] + 1), self.cutoff)) + print( + "FastNS: -> Bead X={:.3f}, Next cell X={:.3f}, cutoff={:.3f} -> +X shift ignored".format( + self.coords_bbox[current_beadid, XX], + self.grid.cellsize[XX] * (cellxyz[XX] + 1), + self.cutoff + )) continue + for yi in range(DIM): + if memory_error: + break - for zi in range(DIM): - if xi == 0: - if self.coords_bbox[current_beadid, ZZ] - self.cutoff > self.grid.cellsize[ZZ] * cellxyz[ZZ]: - if self.debug: + if not self.box.is_triclinic: + if yi == 0: + if self.coords_bbox[current_beadid, YY] - self.cutoff > self.grid.cellsize[YY] * cellxyz[YY]: + if self.debug and debug: with gil: - print("\n#-> Bead Z={:.3f}, Cell Z={:.3f}, cutoff={:.3f} -> -Z shift is ignored".format(self.coords_bbox[current_beadid, ZZ], self.grid.cellsize[ZZ] * cellxyz[ZZ], self.cutoff)) + print("FastNS: -> Bead Y={:.3f}, Cell Y={:.3f}, cutoff={:.3f} -> -Y shift is ignored".format( + self.coords_bbox[current_beadid, YY], + self.grid.cellsize[YY] * cellxyz[YY], + self.cutoff, + )) continue - if xi == 2: - if self.coords_bbox[current_beadid, ZZ] + self.cutoff < self.grid.cellsize[ZZ] * (cellxyz[ZZ] + 1): - if self.debug: + if yi == 2: + if self.coords_bbox[current_beadid, YY] + self.cutoff < self.grid.cellsize[YY] * (cellxyz[YY] + 1): + if self.debug and debug: with gil: - print("\n#-> Bead Z={:.3f}, Next cell Z={:.3f}, cutoff={:.3f} -> +Z shift is ignored".format(self.coords_bbox[current_beadid, ZZ], self.grid.cellsize[XX] * (cellxyz[ZZ] + 1), self.cutoff)) + print("FastNS: -> Bead Y={:.3f}, Next cell Y={:.3f}, cutoff={:.3f} -> +Y shift is ignored".format( + self.coords_bbox[current_beadid, YY], + self.grid.cellsize[YY] * (cellxyz[YY] +1), + self.cutoff, + )) continue + + for zi in range(DIM): + if not self.box.is_triclinic: + if zi == 0: + if self.coords_bbox[current_beadid, ZZ] - self.cutoff > self.grid.cellsize[ZZ] * cellxyz[ZZ]: + if self.coords_bbox[current_beadid, ZZ] - self.cutoff > 0: + if self.debug and debug: + with gil: + print("FastNS: -> Bead Z={:.3f}, Cell Z={:.3f}, cutoff={:.3f} -> -Z shift is ignored".format(self.coords_bbox[current_beadid, ZZ], self.grid.cellsize[ZZ] * cellxyz[ZZ], self.cutoff)) + continue + + if zi == 2: + if self.coords_bbox[current_beadid, ZZ] + self.cutoff < self.grid.cellsize[ZZ] * (cellxyz[ZZ] + 1): + if self.coords_bbox[current_beadid, ZZ] + self.cutoff < self.box.c_pbcbox.box[ZZ][ZZ]: + if self.debug and debug: + with gil: + print("FastNS: -> Bead Z={:.3f}, Next cell Z={:.3f}, cutoff={:.3f} -> +Z shift is ignored".format(self.coords_bbox[current_beadid, ZZ], self.grid.cellsize[XX] * (cellxyz[ZZ] + 1), self.cutoff)) + continue + # Calculate and/or reinitialize shifted coordinates shifted_coord[XX] = self.coords[current_beadid, XX] + (xi - 1) * self.grid.cellsize[XX] shifted_coord[YY] = self.coords[current_beadid, YY] + (yi - 1) * self.grid.cellsize[YY] shifted_coord[ZZ] = self.coords[current_beadid, ZZ] + (zi - 1) * self.grid.cellsize[ZZ] + probe[XX] = self.coords[current_beadid, XX] + (xi - 1) * self.cutoff + probe[YY] = self.coords[current_beadid, YY] + (yi - 1) * self.cutoff + probe[ZZ] = self.coords[current_beadid, ZZ] + (zi - 1) * self.cutoff + # Make sure the shifted coordinates is inside the brick-shaped box for m in range(DIM - 1, -1, -1): @@ -1109,105 +1156,116 @@ cdef class FastNS2(object): for d in range(m+1): shifted_coord[d] -= self.box.c_pbcbox.box[m][d] + while probe[m] < 0: + for d in range(m+1): + probe[d] += self.box.c_pbcbox.box[m][d] + + + while probe[m] >= self.box.c_pbcbox.box[m][m]: + for d in range(m+1): + probe[d] -= self.box.c_pbcbox.box[m][d] + # Get the cell index corresponding to the coord - cellindex = self.grid.coord2cellid(shifted_coord) + cellindex_adjacent = self.grid.coord2cellid(shifted_coord) + cellindex_probe = self.grid.coord2cellid(probe) + + if cellindex == cellindex_probe and xi != 1 and yi != 1 and zi != 1: + if self.debug and debug: + with gil: + print("FastNS: Grid shift [{}][{}][{}]: Cutoff is inside current cell -> This shift is ignored".format( + xi - 1, + yi -1, + zi -1 + )) + continue - if self.debug: - self.grid.cellid2cellxyz(cellindex, debug_cellxyz) + if self.debug and debug: + self.grid.cellid2cellxyz(cellindex_adjacent, debug_cellxyz) with gil: dist_shift = self.box.fast_distance(&self.coords[current_beadid, XX], shifted_coord) grid_shift = np.array([(xi - 1) * self.grid.cellsize[XX], (yi - 1) * self.grid.cellsize[YY], (zi - 1) * self.grid.cellsize[ZZ]]) - print("-> Checking cell#{} ({},{},{}) for neighbors (dshift={:.3f}, grid_shift=({:.3f},{:.3f},{:.3f}->{:.3f})".format( + print("FastNS: -> Checking cell#{} ({},{},{}) for neighbors (dshift={:.3f}, grid_shift=({:.3f},{:.3f},{:.3f}->{:.3f})".format( cellindex, debug_cellxyz[XX], debug_cellxyz[YY], debug_cellxyz[ZZ], dist_shift, grid_shift[XX], grid_shift[YY], grid_shift[ZZ], np.sqrt(np.sum(grid_shift**2)) )) - if dist_shift > 2* self.cutoff: - print(" \_ This cell should be ignored!") - else: - print(" \_ This cell is needed") - for j in range(self.grid.nbeads[cellindex]): - bid = self.grid.beadids[cellindex, j] + + for j in range(self.grid.nbeads[cellindex_adjacent]): + bid = self.grid.beadids[cellindex_adjacent * self.grid.nbeads_per_cell + j] if checked[bid] != 0: continue d2 = self.box.fast_distance2(&self.coords_bbox[current_beadid, XX], &self.coords_bbox[bid, XX]) - if self.debug and False: - self.grid.cellid2cellxyz(cellindex, debug_cellxyz) - with gil: - print( - "Beads #{} (cell[{},{},{}]-coords[{:.3f},{:.3f},{:.3f}]) and #{} (cell[{},{},{}]-coords[{:.3f},{:.3f},{:.3f}]) are tested (d2={:.3f})".format( - current_beadid, - cellxyz[XX], cellxyz[YY], cellxyz[ZZ], - self.coords_bbox[current_beadid, XX], - self.coords_bbox[current_beadid, YY], - self.coords_bbox[current_beadid, ZZ], - bid, - debug_cellxyz[XX], debug_cellxyz[YY], debug_cellxyz[ZZ], - self.coords_bbox[bid, XX], self.coords_bbox[bid, YY], - self.coords_bbox[bid, ZZ], - d2)) + # if self.debug: + # self.grid.cellid2cellxyz(cellindex, debug_cellxyz) + # with gil: + # print( + # "Beads #{} (cell[{},{},{}]-coords[{:.3f},{:.3f},{:.3f}]) and #{} (cell[{},{},{}]-coords[{:.3f},{:.3f},{:.3f}]) are tested (d2={:.3f})".format( + # current_beadid, + # cellxyz[XX], cellxyz[YY], cellxyz[ZZ], + # self.coords_bbox[current_beadid, XX], + # self.coords_bbox[current_beadid, YY], + # self.coords_bbox[current_beadid, ZZ], + # bid, + # debug_cellxyz[XX], debug_cellxyz[YY], debug_cellxyz[ZZ], + # self.coords_bbox[bid, XX], self.coords_bbox[bid, YY], + # self.coords_bbox[bid, ZZ], + # d2)) if d2 < cutoff2: if d2 < EPSILON: continue - if self.debug: + if self.debug and debug: self.grid.cellid2cellxyz(cellindex, debug_cellxyz) with gil: self.box.fast_pbc_dx(&self.coords[current_beadid, XX], &self.coords[bid, XX], dx) dx_py = np.array([dx[XX], dx[YY], dx[ZZ]]) - - print(" \_ Neighbor found: bead#{} (cell[{},{},{}]) -> dx={} -> d={:.3f}".format(bid, debug_cellxyz[XX], debug_cellxyz[YY], debug_cellxyz[ZZ], dx_py, np.sqrt(d2))) - results.add_neighbors(current_beadid, bid, d2) + print("FastNS: \_ Neighbor found: bead#{} (cell[{},{},{}]) -> dx={} -> d={:.3f}".format(bid, debug_cellxyz[XX], debug_cellxyz[YY], debug_cellxyz[ZZ], dx_py, np.sqrt(d2))) + if results.add_neighbors(current_beadid, bid, d2) == 0: + memory_error = True + break npairs += 1 + checked[current_beadid] = 1 - if first_ns: - guessed_size = ((npairs + 1) * 0.75 * size_search) - with gil: - if self.debug: - print("Neighbors of first beads= {} -> Approximated total of pairs={}".format(npairs, guessed_size)) - results.resize(guessed_size) + if memory_error: + raise MemoryError("Could not allocate memory to store NS results") - first_ns = False - checked[current_beadid] = 1 if self.debug: print("Total number of pairs={}".format(npairs)) - ref_bead = 13937 - beads = np.array([4398, 4401, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, 47451]) - 1 - for bid in beads: - self.box.fast_pbc_dx(&self.coords[ref_bead, XX], &self.coords[bid, XX], dx) - dx_py = np.array([dx[XX], dx[YY], dx[ZZ]]) - self.box.fast_pbc_dx(&self.coords_bbox[ref_bead, XX], &self.coords_bbox[bid, XX], dx) - rect_dx_py = np.array([dx[XX], dx[YY], dx[ZZ]]) - self.grid.cellid2cellxyz(self.grid.coord2cellid(&self.coords_bbox[bid, XX]), cellxyz) - print("Bead #{} ({:.3f},{:.3f},{:.3f})->rect({:.3f},{:.3f},{:.3f}) - cell[{},{},{}]: dx=[{:.3f},{:.3f},{:.3f}] -> dist: {:.3f} ({}) - rect_dist: {:.3f} ({})".format( - bid, - self.coords[bid, XX], self.coords[bid, YY], self.coords[bid,ZZ], - self.coords_bbox[bid, XX], self.coords_bbox[bid, YY], self.coords_bbox[bid,ZZ], - cellxyz[XX], cellxyz[YY], cellxyz[ZZ], - dx[XX], dx[YY], dx[ZZ], - np.sqrt(np.sum(dx_py**2)), - self.box.fast_distance(&self.coords[ref_bead, XX], &self.coords[bid, XX]) <= self.cutoff, - np.sqrt(np.sum(rect_dx_py**2)), - np.sqrt(np.sum(rect_dx_py**2)) <= self.cutoff - )) + # ref_bead = 13937 + # beads = np.array([4398, 4401, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, 47451]) - 1 + # for bid in beads: + # self.box.fast_pbc_dx(&self.coords[ref_bead, XX], &self.coords[bid, XX], dx) + # dx_py = np.array([dx[XX], dx[YY], dx[ZZ]]) + # self.box.fast_pbc_dx(&self.coords_bbox[ref_bead, XX], &self.coords_bbox[bid, XX], dx) + # rect_dx_py = np.array([dx[XX], dx[YY], dx[ZZ]]) + # self.grid.cellid2cellxyz(self.grid.coord2cellid(&self.coords_bbox[bid, XX]), cellxyz) + # print("Bead #{} ({:.3f},{:.3f},{:.3f})->rect({:.3f},{:.3f},{:.3f}) - cell[{},{},{}]: dx=[{:.3f},{:.3f},{:.3f}] -> dist: {:.3f} ({})".format( + # bid, + # self.coords[bid, XX], self.coords[bid, YY], self.coords[bid,ZZ], + # self.coords_bbox[bid, XX], self.coords_bbox[bid, YY], self.coords_bbox[bid,ZZ], + # cellxyz[XX], cellxyz[YY], cellxyz[ZZ], + # dx[XX], dx[YY], dx[ZZ], + # np.sqrt(np.sum(dx_py**2)), + # self.box.fast_distance(&self.coords[ref_bead, XX], &self.coords[bid, XX]) <= self.cutoff, + # )) return results # Python interface -cdef class FastNS(object): +cdef class FastNSOld(object): cdef PBCBox box cdef readonly real[:, ::1] coords cdef real[:, ::1] coords_bbox @@ -1376,108 +1434,10 @@ cdef class FastNS(object): if holder == NULL: raise MemoryError("Could not allocate memory to run NS core") - results = NSResults(self.box, self.cutoff) + results = NSResultsOld(self.box, self.cutoff) results.populate(holder, self.coords, search_coords_view) # Free memory allocated to holder free_neighborhood_holder(holder) return results - - def search2(self, search_ids): - cdef ns_int[:] search_ids_view, neighbors_view - cdef ns_int nid, i, j, size_search - cdef ns_neighborhood_holder *holder - cdef ns_neighborhood *neighborhood - - cdef ns_int current_beadid - cdef rvec current_coords - - #Temp stuff - cdef ns_int size = self.coords_bbox.shape[0] - cdef ns_int d, m - cdef ns_int xi, yi, zi, bid, i_bead - cdef real d2 - cdef rvec shifted_coord, dx, neighbor_coord, corrected_coords - - cdef bint already_checked[27] - cdef bint skip - cdef ns_int nchecked = 0, icheck - cdef ns_int cell_index - cdef ns_int cell_offset - - cdef ns_int[:] neighbor_ids - cdef ns_int[:] nneighbors - cdef ns_int nid_offset - - cdef real cutoff2 = self.cutoff * self.cutoff - - - if not self.prepared: - self.prepare() - - nid_offset = self.grid.beadpercell * 27 - neighbor_ids = np.empty(size * nid_offset, dtype=np.int) - nneighbors = np.zeros(size, dtype=np.int) - - - search_ids_view = search_ids - size_search = search_ids_view.shape[0] - - with nogil: - - for current_beadid in range(size): - - for zi in range(3): - for yi in range(3): - for xi in range(3): - # Calculate and/or reinitialize shifted coordinates - shifted_coord[XX] = self.coords_bbox[current_beadid, XX] + (xi - 1) * self.grid.cellsize[XX] - shifted_coord[YY] = self.coords_bbox[current_beadid, YY] + (yi - 1) * self.grid.cellsize[YY] - shifted_coord[ZZ] = self.coords_bbox[current_beadid, ZZ] + (zi - 1) * self.grid.cellsize[ZZ] - - # Make sure the shifted coordinates is inside the brick-shaped box - for m in range(DIM - 1, -1, -1): - - while shifted_coord[m] < 0: - for d in range(m+1): - shifted_coord[d] += self.box.c_pbcbox.box[m][d] - - - while shifted_coord[m] >= self.box.c_pbcbox.box[m][m]: - for d in range(m+1): - shifted_coord[d] -= self.box.c_pbcbox.box[m][d] - - # Get the cell index corresponding to the coord - cell_index = (shifted_coord[ZZ] / self.grid.cellsize[ZZ]) * (self.grid.ncells[XX] * self.grid.ncells[YY]) +\ - (shifted_coord[YY] / self.grid.cellsize[YY]) * self.grid.ncells[XX] + \ - (shifted_coord[XX] / self.grid.cellsize[XX]) - - # Search for neighbors inside this cell - for i_bead in range(self.grid.nbeads[cell_index]): - cell_offset = cell_index * self.grid.beadpercell + i_bead - - bid = self.grid.beadids[cell_offset] - - if bid <= current_beadid: - continue - - self.box.fast_pbc_dx(&self.coords_bbox[current_beadid, XX], &self.coords_bbox[bid, XX], dx) - - d2 = rvec_norm2(dx) - - if d2 < cutoff2: - neighbor_ids[current_beadid * nid_offset + nneighbors[current_beadid]] = bid - nneighbors[current_beadid] += 1 - - neighbor_ids[bid * nid_offset + nneighbors[bid]] = current_beadid - nneighbors[bid] += 1 - - - - - results = [] - #for bid in search_ids: - # results.append(np.asarray(neighbor_ids[bid * nid_offset : bid * nid_offset + nneighbors[bid]])) - - return results diff --git a/testsuite/MDAnalysisTests/lib/test_gridsearch.py b/testsuite/MDAnalysisTests/lib/test_gridsearch.py index 30ab8813036..a64143d446e 100644 --- a/testsuite/MDAnalysisTests/lib/test_gridsearch.py +++ b/testsuite/MDAnalysisTests/lib/test_gridsearch.py @@ -28,7 +28,7 @@ import numpy as np import MDAnalysis as mda -from MDAnalysis.lib import grid +from MDAnalysis.lib import nsgrid from MDAnalysis.lib.pkdtree import PeriodicKDTree @@ -42,105 +42,86 @@ def universe(): -@pytest.fixture -def grid_results(): - u = mda.Universe(GRO) - cutoff = 2 - ref_pos = u.atoms.positions[13937] - return run_grid_search(u, ref_pos, cutoff) - - -def run_grid_search(u, ref_pos, cutoff): +def run_grid_search(u, ids, cutoff=3): coords = u.atoms.positions # Run grid search - searcher = grid.FastNS(u) - searcher.set_cutoff(cutoff) - searcher.set_coords(coords) - searcher.prepare() - - return searcher.search(ref_pos) - - - -def run_search(universe, ref_id): - cutoff = 3 - coords = universe.atoms.positions - ref_pos = coords[ref_id] - - - # Run pkdtree search - pkdt = PeriodicKDTree(universe.atoms.dimensions, bucket_size=10) - pkdt.set_coords(coords) - pkdt.search(ref_pos, cutoff) - - results_pkdtree = pkdt.get_indices() - results_pkdtree.remove(ref_id) - results_pkdtree = np.array(results_pkdtree) - results_pkdtree.sort() - - # Run grid search - results_grid = run_grid_search(universe, ref_pos, cutoff) - results_grid = results_grid.get_indices()[0] - results_grid.sort() - - return results_pkdtree, results_grid + searcher = grid.FastNS(u, cutoff, coords, debug=True) + return searcher.search(ids, debug=False) def test_gridsearch(universe): """Check that pkdtree and grid search return the same results (No PBC needed)""" ref_id = 0 - results_pkdtree, results_grid = run_search(universe, ref_id) - assert_equal(results_pkdtree, results_grid) - - -def test_gridsearch_PBC(universe): - """Check that pkdtree and grid search return the same results (PBC needed)""" - - ref_id = 13937 - results_pkdtree, results_grid = run_search(universe, ref_id) - - assert_equal(results_pkdtree, results_grid) + cutoff = 3 + results = np.array([2, 3, 4, 5, 6, 7, 8, 9, 18, 19, 1211, 10862, 10865, 17582, 17585, 38342, + 38345]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! -def test_gridsearch_arraycoord(universe): - """Check the NS routine accepts a single bead coordinate as well as array of coordinates""" - cutoff = 2 - ref_pos = universe.atoms.positions[:5] - - results = [ - np.array([2, 1, 4, 3]), - np.array([2, 0, 3]), - np.array([0, 1, 3]), - np.array([ 2, 0, 1, 38341]), - np.array([ 6, 0, 5, 17]) - ] - - results_grid = run_grid_search(universe, ref_pos, cutoff).get_indices() + results_grid = run_grid_search(universe, ref_id, cutoff).get_indices()[0] - assert_equal(results_grid, results) + assert_equal(results, results_grid) -def test_gridsearch_search_coordinates(grid_results): - """Check the NS routine can return coordinates instead of ids""" +def test_gridsearch_PBC(universe): + """Check that pkdtree and grid search return the same results (No PBC needed)""" - results = np.array( - [ - [40.32, 34.25, 55.9], - [0.61, 76.33, -0.56], - [0.48999998, 75.9, 0.19999999], - [-0.11, 76.19, 0.77] - ]) + ref_id = 13937 + results = np.array([4398, 4401, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, + 47451]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! - assert_allclose(grid_results.get_coordinates()[0], results) + results_grid = run_grid_search(universe, ref_id).get_indices()[0] + assert_equal(results, results_grid) -def test_gridsearch_search_distances(grid_results): - """Check the NS routine can return PBC distances from neighbors""" - results = np.array([0.096, 0.096, 0.015, 0.179]) * 10 # These distances were obtained using gmx distance - results.sort() - rounded_results = np.round(grid_results.get_distances()[0], 2) - assert_allclose(sorted(rounded_results), results) +# def test_gridsearch_PBC(universe): +# """Check that pkdtree and grid search return the same results (PBC needed)""" +# +# ref_id = 13937 +# results_pkdtree, results_grid = run_search(universe, ref_id) +# assert_equal(results_pkdtree, results_grid) +# +# +# def test_gridsearch_arraycoord(universe): +# """Check the NS routine accepts a single bead coordinate as well as array of coordinates""" +# cutoff = 2 +# ref_pos = universe.atoms.positions[:5] +# +# results = [ +# np.array([2, 1, 4, 3]), +# np.array([2, 0, 3]), +# np.array([0, 1, 3]), +# np.array([ 2, 0, 1, 38341]), +# np.array([ 6, 0, 5, 17]) +# ] +# +# results_grid = run_grid_search(universe, ref_pos, cutoff).get_indices() +# +# assert_equal(results_grid, results) +# +# +# def test_gridsearch_search_coordinates(grid_results): +# """Check the NS routine can return coordinates instead of ids""" +# +# results = np.array( +# [ +# [40.32, 34.25, 55.9], +# [0.61, 76.33, -0.56], +# [0.48999998, 75.9, 0.19999999], +# [-0.11, 76.19, 0.77] +# ]) +# +# assert_allclose(grid_results.get_coordinates()[0], results) +# +# +# def test_gridsearch_search_distances(grid_results): +# """Check the NS routine can return PBC distances from neighbors""" +# results = np.array([0.096, 0.096, 0.015, 0.179]) * 10 # These distances were obtained using gmx distance +# results.sort() +# +# rounded_results = np.round(grid_results.get_distances()[0], 2) +# +# assert_allclose(sorted(rounded_results), results) From 1a0e1507ea32bbede89e71f6ea2e5b6cbad201b4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Buchoux?= Date: Wed, 18 Jul 2018 10:10:29 +0200 Subject: [PATCH 21/47] Module MDAnalysis/lib/c_gridsearch.pyx renamed to nsgrid.pyx --- .../MDAnalysis/lib/{c_gridsearch.pyx => nsgrid.pyx} | 7 +++---- .../lib/{test_gridsearch.py => test_nsgrid.py} | 11 +++++------ 2 files changed, 8 insertions(+), 10 deletions(-) rename package/MDAnalysis/lib/{c_gridsearch.pyx => nsgrid.pyx} (97%) rename testsuite/MDAnalysisTests/lib/{test_gridsearch.py => test_nsgrid.py} (96%) diff --git a/package/MDAnalysis/lib/c_gridsearch.pyx b/package/MDAnalysis/lib/nsgrid.pyx similarity index 97% rename from package/MDAnalysis/lib/c_gridsearch.pyx rename to package/MDAnalysis/lib/nsgrid.pyx index a1f519eac37..134356bd572 100644 --- a/package/MDAnalysis/lib/c_gridsearch.pyx +++ b/package/MDAnalysis/lib/nsgrid.pyx @@ -23,9 +23,9 @@ # # -#cython: cdivision=True -#cython: boundscheck=False -#cython: initializedcheck=False +# cython: cdivision=True +# cython: boundscheck=False +# cython: initializedcheck=False """ Neighbor search library --- :mod:`MDAnalysis.lib.grid` @@ -93,7 +93,6 @@ cdef void rvec_clear(rvec a) nogil: a[YY]=0.0 a[ZZ]=0.0 - cdef struct cPBCBox_t: matrix box rvec fbox_diag diff --git a/testsuite/MDAnalysisTests/lib/test_gridsearch.py b/testsuite/MDAnalysisTests/lib/test_nsgrid.py similarity index 96% rename from testsuite/MDAnalysisTests/lib/test_gridsearch.py rename to testsuite/MDAnalysisTests/lib/test_nsgrid.py index a64143d446e..d7be4f79e65 100644 --- a/testsuite/MDAnalysisTests/lib/test_gridsearch.py +++ b/testsuite/MDAnalysisTests/lib/test_nsgrid.py @@ -24,15 +24,12 @@ import pytest -from numpy.testing import assert_equal, assert_allclose +from numpy.testing import assert_equal import numpy as np import MDAnalysis as mda -from MDAnalysis.lib import nsgrid -from MDAnalysis.lib.pkdtree import PeriodicKDTree - - from MDAnalysisTests.datafiles import GRO +from MDAnalysis.lib import nsgrid @pytest.fixture @@ -46,7 +43,9 @@ def run_grid_search(u, ids, cutoff=3): coords = u.atoms.positions # Run grid search - searcher = grid.FastNS(u, cutoff, coords, debug=True) + searcher = nsgrid.FastNS(u, cutoff, coords, debug=True) + + return searcher.search(ids, debug=False) From a1e03dd8bdc969e3d5dc21379fa6cee476f929e0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Buchoux?= Date: Wed, 18 Jul 2018 16:44:11 +0200 Subject: [PATCH 22/47] Tests written for FastNS --- package/MDAnalysis/lib/__init__.py | 2 +- package/MDAnalysis/lib/nsgrid.pyx | 769 ++++--------------- testsuite/MDAnalysisTests/lib/test_nsgrid.py | 245 ++++-- 3 files changed, 326 insertions(+), 690 deletions(-) diff --git a/package/MDAnalysis/lib/__init__.py b/package/MDAnalysis/lib/__init__.py index a740bd8d5bf..d5a583f9aea 100644 --- a/package/MDAnalysis/lib/__init__.py +++ b/package/MDAnalysis/lib/__init__.py @@ -39,4 +39,4 @@ from . import NeighborSearch from . import formats from . import pkdtree -from . import grid \ No newline at end of file +from . import nsgrid \ No newline at end of file diff --git a/package/MDAnalysis/lib/nsgrid.pyx b/package/MDAnalysis/lib/nsgrid.pyx index 134356bd572..e4de849ffff 100644 --- a/package/MDAnalysis/lib/nsgrid.pyx +++ b/package/MDAnalysis/lib/nsgrid.pyx @@ -43,19 +43,14 @@ DEF ZZ = 2 DEF EPSILON = 1e-5 -DEF NEIGHBORHOOD_ALLOCATION_INCREMENT = 50 - DEF BOX_MARGIN=1.0010 DEF MAX_NTRICVEC=12 -from libc.stdlib cimport malloc, realloc, free +# Used to handle memory allocation from cpython.mem cimport PyMem_Malloc, PyMem_Realloc, PyMem_Free - from libc.math cimport sqrt - import numpy as np cimport numpy as np -cimport cython ctypedef np.int_t ns_int ctypedef np.float32_t real @@ -64,27 +59,8 @@ ctypedef ns_int ivec[DIM] ctypedef ns_int ipair[2] ctypedef real matrix[DIM][DIM] -cdef struct ns_grid: - ns_int beadpercell - ns_int size - ns_int[DIM] ncells - real[DIM] cellsize - ns_int *nbeads # size - ns_int *beadids # size * beadpercell - ns_int *cellids # size - -cdef struct ns_neighborhood: - real cutoff - ns_int allocated_size - ns_int size - ns_int *beadids - -cdef struct ns_neighborhood_holder: - ns_int size - ns_neighborhood **neighborhoods - -# Useful stuff +# Useful Functions cdef real rvec_norm2(const rvec a) nogil: return a[XX]*a[XX]+a[YY]*a[YY]+a[ZZ]*a[ZZ] @@ -93,6 +69,11 @@ cdef void rvec_clear(rvec a) nogil: a[YY]=0.0 a[ZZ]=0.0 +######################################################################################################################## +# +# Utility class to handle PBC +# +######################################################################################################################## cdef struct cPBCBox_t: matrix box rvec fbox_diag @@ -186,15 +167,10 @@ cdef class PBCBox(object): # Choose the vector within the brick around 0,0,0 that # will become the shortest due to shift try. - - if d == DIM: - trial[d] = 0 - pos[d] = 0 + if trial[d] < 0: + pos[d] = min(self.c_pbcbox.hbox_diag[d], -trial[d]) else: - if trial[d] < 0: - pos[d] = min(self.c_pbcbox.hbox_diag[d], -trial[d]) - else: - pos[d] = max(-self.c_pbcbox.hbox_diag[d], -trial[d]) + pos[d] = max(-self.c_pbcbox.hbox_diag[d], -trial[d]) d2old += pos[d]**2 d2new += (pos[d] + trial[d])**2 @@ -270,14 +246,34 @@ cdef class PBCBox(object): for j in range (i, -1, -1): dx[j] += self.c_pbcbox.box[i][j] + def dx(self, real[:] a, real[:] b): + cdef rvec dx + if a.shape[0] != DIM or b.shape[0] != DIM: + raise ValueError("Not 3 D coordinates") + + self.fast_pbc_dx(&a[XX], &b[XX], dx) + + return np.array([dx[XX], dx[YY], dx[ZZ]], dtype=np.float32) + + cdef real fast_distance2(self, rvec a, rvec b) nogil: cdef rvec dx self.fast_pbc_dx(a, b, dx) return rvec_norm2(dx) + def distance2(self, real[:] a, real[:] b): + if a.shape[0] != DIM or b.shape[0] != DIM: + raise ValueError("Not 3 D coordinates") + return self.fast_distance2(&a[XX], &b[XX]) + cdef real fast_distance(self, rvec a, rvec b) nogil: return sqrt(self.fast_distance2(a,b)) + def distance(self, real[:] a, real[:] b): + if a.shape[0] != DIM or b.shape[0] != DIM: + raise ValueError("Not 3 D coordinates") + return self.fast_distance(&a[XX], &b[XX]) + cdef real[:, ::1]fast_put_atoms_in_bbox(self, real[:,::1] coords) nogil: cdef ns_int i, m, d, natoms, wd = 0 cdef real[:,::1] bbox_coords @@ -305,341 +301,20 @@ cdef class PBCBox(object): return bbox_coords def put_atoms_in_bbox(self, real[:,::1] coords): - if coords.shape[0] == 0: - return np.zeros((0, DIM), dtype=np.float32) return np.asarray(self.fast_put_atoms_in_bbox(coords)) + ######################################################################################################################## # # Neighbor Search Stuff # ######################################################################################################################## - -cdef ns_neighborhood_holder *create_neighborhood_holder() nogil: - cdef ns_neighborhood_holder *holder - - holder = malloc(sizeof(ns_neighborhood_holder)) - holder.size = 0 - holder.neighborhoods = NULL - - return holder - -cdef void free_neighborhood_holder(ns_neighborhood_holder *holder) nogil: - cdef ns_int i - - if holder == NULL: - return - - for i in range(holder.size): - if holder.neighborhoods[i].beadids != NULL: - free(holder.neighborhoods[i].beadids) - free(holder.neighborhoods[i]) - - if holder.neighborhoods != NULL: - free(holder.neighborhoods) - free(holder) - -cdef ns_neighborhood *retrieve_neighborhood(rvec current_coords, real[:, ::1]neighborcoords, ns_grid *grid, PBCBox box, real cutoff2) nogil: - cdef ns_int d, m - cdef ns_int xi, yi, zi, bid, i_bead - cdef real d2 - cdef rvec shifted_coord, dx, neighbor_coord, corrected_coords - - cdef bint already_checked[27] - cdef bint skip - cdef ns_int nchecked = 0, icheck - cdef ns_int cell_index - cdef ns_int cell_offset - - cdef ns_neighborhood *neighborhood = malloc(sizeof(ns_neighborhood)) - if neighborhood == NULL: - return NULL - - neighborhood.size = 0 - neighborhood.allocated_size = NEIGHBORHOOD_ALLOCATION_INCREMENT - neighborhood.beadids = malloc(NEIGHBORHOOD_ALLOCATION_INCREMENT * sizeof(ns_int)) - - if neighborhood.beadids == NULL: - free(neighborhood) - return NULL - - for zi in range(3): - for yi in range(3): - for xi in range(3): - # Calculate and/or reinitialize shifted coordinates - shifted_coord[XX] = current_coords[XX] + (xi - 1) * grid.cellsize[XX] - shifted_coord[YY] = current_coords[YY] + (yi - 1) * grid.cellsize[YY] - shifted_coord[ZZ] = current_coords[ZZ] + (zi - 1) * grid.cellsize[ZZ] - - # Make sure the shifted coordinates is inside the brick-shaped box - for m in range(DIM - 1, -1, -1): - - while shifted_coord[m] < 0: - for d in range(m+1): - shifted_coord[d] += box.c_pbcbox.box[m][d] - - - while shifted_coord[m] >= box.c_pbcbox.box[m][m]: - for d in range(m+1): - shifted_coord[d] -= box.c_pbcbox.box[m][d] - - # Get the cell index corresponding to the coord - cell_index = (shifted_coord[ZZ] / grid.cellsize[ZZ]) * (grid.ncells[XX] * grid.ncells[YY]) +\ - (shifted_coord[YY] / grid.cellsize[YY]) * grid.ncells[XX] + \ - (shifted_coord[XX] / grid.cellsize[XX]) - - # Just a safeguard - #if cell_index >= grid.size: - # continue - - # Check the cell index was not already selected - #skip = False - #for icheck in range(nchecked): - # if already_checked[icheck] == cell_index: - # skip = True - # break - #if skip: - # continue - - # Search for neighbors inside this cell - for i_bead in range(grid.nbeads[cell_index]): - cell_offset = cell_index * grid.beadpercell + i_bead - - bid = grid.beadids[cell_offset] - - box.fast_pbc_dx(current_coords, &neighborcoords[bid, XX], dx) - - d2 = rvec_norm2(dx) - - if d2 < cutoff2: - if d2 < EPSILON: # Don't add the current bead as its own neighbor! - continue - - # Update neighbor lists - neighborhood.beadids[neighborhood.size] = bid - neighborhood.size += 1 - - if neighborhood.size >= neighborhood.allocated_size: - neighborhood.allocated_size += NEIGHBORHOOD_ALLOCATION_INCREMENT - neighborhood.beadids = realloc( neighborhood.beadids, neighborhood.allocated_size * sizeof(ns_int)) - - if neighborhood.beadids == NULL: - free(neighborhood) - return NULL - - # Register the cell as checked - #already_checked[nchecked] = cell_index - #nchecked += 1 - - return neighborhood - - -cdef ns_neighborhood_holder *ns_core(real[:, ::1] refcoords, - real[:, ::1] neighborcoords, - ns_grid *grid, - PBCBox box, - real cutoff) nogil: - cdef ns_int coordid, i, j - cdef ns_int ncoords = refcoords.shape[0] - cdef ns_int ncoords_neighbors = neighborcoords.shape[0] - cdef real cutoff2 = cutoff * cutoff - cdef ns_neighborhood_holder *holder - - cdef ns_int *neighbor_buf - cdef ns_int buf_size, ibuf - - holder = create_neighborhood_holder() - if holder == NULL: - return NULL - - holder.neighborhoods = malloc(sizeof(ns_neighborhood *) * ncoords) - if holder.neighborhoods == NULL: - free_neighborhood_holder(holder) - return NULL - - # Here starts the real core and the iteration over coordinates - for coordid in range(ncoords): - holder.neighborhoods[coordid] = retrieve_neighborhood(&refcoords[coordid, XX], - neighborcoords, - grid, - box, - cutoff2) - if holder.neighborhoods[coordid] == NULL: - free_neighborhood_holder(holder) - return NULL - - holder.neighborhoods[coordid].cutoff = cutoff - holder.size += 1 - - return holder - -cdef class NSResultsOld(object): - """ - Class used to store results returned by `MDAnalysis.lib.grid.FastNS.search` - """ - cdef PBCBox box - cdef readonly real cutoff - cdef np.ndarray grid_coords - cdef np.ndarray ref_coords - cdef ns_int **nids - cdef ns_int *nsizes - cdef ns_int size - cdef list indices - cdef list coordinates - cdef list distances - - def __dealloc__(self): - if self.nids != NULL: - for i in range(self.size): - if self.nids[i] != NULL: - free(self.nids[i]) - free(self.nids) - - if self.nsizes != NULL: - free(self.nsizes) - - def __init__(self, PBCBox box, real cutoff): - self.box = box - self.cutoff = cutoff - - self.size = 0 - self.nids = NULL - self.nsizes = NULL - - self.grid_coords = None - self.ref_coords = None - - self.indices = None - self.coordinates = None - self.distances = None - - - cdef populate(self, ns_neighborhood_holder *holder, grid_coords, ref_coords): - cdef ns_int nid, i - cdef ns_neighborhood *neighborhood - - self.grid_coords = np.asarray(grid_coords) - self.ref_coords = np.asarray(ref_coords) - - # Allocate memory - self.nsizes = malloc(sizeof(ns_int) * holder.size) - if self.nsizes == NULL: - raise MemoryError("Could not allocate memory for NSResults") - - self.nids = malloc(sizeof(ns_int *) * holder.size) - if self.nids == NULL: - raise MemoryError("Could not allocate memory for NSResults") - - for nid in range(holder.size): - neighborhood = holder.neighborhoods[nid] - - self.nsizes[nid] = neighborhood.size - - self.nids[nid] = malloc(sizeof(ns_int *) * neighborhood.size) - if self.nids[nid] == NULL: - raise MemoryError("Could not allocate memory for NSResults") - - with nogil: - for nid in range(holder.size): - neighborhood = holder.neighborhoods[nid] - - for i in range(neighborhood.size): - self.nids[nid][i] = neighborhood.beadids[i] - - self.size = holder.size - - def get_indices(self): - """ - Return Neighbors indices. - - :return: list of indices - """ - cdef ns_int i, nid, size - - if self.indices is None: - indices = [] - - for nid in range(self.size): - size = self.nsizes[nid] - - tmp_incides = np.empty((size), dtype=np.int) - - for i in range(size): - tmp_incides[i] = self.nids[nid][i] - - indices.append(tmp_incides) - - self.indices = indices - - return self.indices - - - def get_coordinates(self): - """ - Return coordinates of neighbors. - - :return: list of coordinates - """ - cdef ns_int i, nid, size, beadid - - if self.coordinates is None: - coordinates = [] - - for nid in range(self.size): - size = self.nsizes[nid] - - tmp_values = np.empty((size, DIM), dtype=np.float32) - - for i in range(size): - beadid = self.nids[nid][i] - tmp_values[i] = self.grid_coords[beadid] - - coordinates.append(tmp_values) - - self.coordinates = coordinates - - return self.coordinates - - - def get_distances(self): - """ - Return coordinates of neighbors. - - :return: list of distances - """ - cdef ns_int i, nid, size, j, beadid - cdef rvec ref, other, dx - cdef real dist - cdef real[:, ::1] ref_coords = self.ref_coords - cdef real[:, ::1] grid_coords = self.grid_coords - - if self.distances is None: - distances = [] - - for nid in range(self.size): - size = self.nsizes[nid] - - tmp_values = np.empty((size), dtype=np.float32) - ref = &ref_coords[nid, 0] - - for i in range(size): - beadid = self.nids[nid][i] - other = &grid_coords[beadid, 0] - - tmp_values[i] = self.box.fast_distance(ref, other) - - distances.append(tmp_values) - - self.distances = distances - - return self.distances - cdef class NSResults(object): cdef readonly real cutoff cdef ns_int npairs cdef bint debug - cdef real[:, ::1] grid_coords # shape: size, DIM + cdef real[:, ::1] coords # shape: size, DIM cdef ns_int[:] search_ids cdef ns_int allocation_size @@ -651,12 +326,13 @@ cdef class NSResults(object): cdef list distances_buffer cdef np.ndarray pairs_buffer cdef np.ndarray pair_distances_buffer + cdef np.ndarray pair_coordinates_buffer - def __init__(self, real cutoff, real[:, ::1]grid_coords,ns_int[:] search_ids, debug=False): + def __init__(self, real cutoff, real[:, ::1]coords, ns_int[:] search_ids, debug=False): self.debug = debug self.cutoff = cutoff - self.grid_coords = grid_coords + self.coords = coords self.search_ids = search_ids # Preallocate memory @@ -674,7 +350,7 @@ cdef class NSResults(object): self.coordinates_buffer = None self.distances_buffer = None self.pairs_buffer = None - self.pair_distances_buffer = None + self.pair_coordinates_buffer = None def __dealloc__(self): PyMem_Free(self.pairs) @@ -745,12 +421,25 @@ cdef class NSResults(object): def get_pair_distances(self): cdef ns_int i - if self.pair_distances_buffer is None: - self.pair_distances_buffer = np.empty(self.npairs, dtype=np.float32) + if self.pair_coordinates_buffer is None: + self.pair_coordinates_buffer = np.empty(self.npairs, dtype=np.float32) + for i in range(self.npairs): + self.pair_coordinates_buffer[i] = self.pair_distances2[i] + self.pair_coordinates_buffer = np.sqrt(self.pair_coordinates_buffer) + return self.pair_coordinates_buffer + + def get_pair_coordinates(self): + cdef ns_int i, j, bead_i, bead_j + if self.pair_coordinates_buffer is None: + self.pair_coordinates_buffer = np.empty((self.npairs, 2, DIM), dtype=np.float32) for i in range(self.npairs): - self.pair_distances_buffer[i] = self.pair_distances2[i] - self.pair_distances_buffer = np.sqrt(self.pair_distances_buffer) - return self.pair_distances_buffer + bead_i = self.pairs[i][0] + bead_j = self.pairs[i][1] + + for j in range(DIM): + self.pair_coordinates_buffer[i, 0, j] = self.coords[bead_i, j] + self.pair_coordinates_buffer[i, 1, j] = self.coords[bead_j, j] + return self.pair_coordinates_buffer cdef create_buffers(self): cdef ns_int i, beadid_i, beadid_j @@ -767,8 +456,8 @@ cdef class NSResults(object): beadid_j = self.pairs[i][1] dist2 = self.pair_distances2[i] - coord_i = self.grid_coords[beadid_i] - coord_j = self.grid_coords[beadid_j] + coord_i = self.coords[beadid_i] + coord_j = self.coords[beadid_j] indices_buffer[beadid_i].append(beadid_j) indices_buffer[beadid_j].append(beadid_i) @@ -933,12 +622,10 @@ cdef class NSGrid(object): PyMem_Free(beadcounts) - -# Python interface cdef class FastNS(object): cdef bint debug cdef PBCBox box - cdef readonly real[:, ::1] coords + cdef real[:, ::1] coords cdef real[:, ::1] coords_bbox cdef readonly real cutoff cdef bint prepared @@ -954,9 +641,6 @@ cdef class FastNS(object): raise TypeError("FastNS class must be initialized with a valid MDAnalysis.Universe instance") box = triclinic_vectors(u.dimensions) - if box.shape != (3, 3): - raise ValueError("Box must be provided as triclinic_dimensions (a 3x3 numpy.ndarray of unit cell vectors") - self.box = PBCBox(box) @@ -967,18 +651,14 @@ cdef class FastNS(object): self.coords_bbox = self.box.fast_put_atoms_in_bbox(coords) - if not prepare: - return - - if self.cutoff < 0: + if cutoff < 0: raise ValueError("Cutoff must be positive!") - if self.cutoff * self.cutoff > self.box.c_pbcbox.max_cutoff2: + if cutoff * cutoff > self.box.c_pbcbox.max_cutoff2: raise ValueError("Cutoff greater than maximum cutoff ({:.3f}) given the PBC") self.cutoff = cutoff self.grid = NSGrid(self.coords_bbox.shape[0], cutoff, self.box, max_gridsize, debug=debug) self.prepared = False - if prepare: self.prepare() @@ -991,7 +671,7 @@ cdef class FastNS(object): self.prepared = True - def search(self, search_ids=None, bint debug=False): + def search(self, search_ids=None): cdef ns_int i, j, size_search cdef ns_int d, m cdef NSResults results @@ -1012,12 +692,12 @@ cdef class FastNS(object): cdef ns_int nchecked = 0 - - cdef real cutoff2 = self.cutoff * self.cutoff cdef ns_int[:] checked cdef ns_int npairs = 0 + #cdef bint debug=False + if not self.prepared: self.prepare() @@ -1034,12 +714,12 @@ cdef class FastNS(object): checked = np.zeros(size, dtype=np.int) - results = NSResults(self.cutoff, self.coords_bbox, search_ids, self.debug) + results = NSResults(self.cutoff, self.coords, search_ids, self.debug) cdef bint memory_error = False - if self.debug and debug: - print("FastNS: Debug flag is set to True for FastNS.search()") + # if self.debug and debug: + # print("FastNS: Debug flag is set to True for FastNS.search()") with nogil: for i in range(size_search): @@ -1050,13 +730,13 @@ cdef class FastNS(object): cellindex = self.grid.cellids[current_beadid] self.grid.cellid2cellxyz(cellindex, cellxyz) - if self.debug and debug: - with gil: - print("FastNS: Checking neighbors for bead #{} ({:.3f},{:.3f},{:.3f})->rect({:.3f},{:.3f},{:.3f}) - cell[{},{},{}]:" .format( - current_beadid, - self.coords[current_beadid, XX], self.coords[current_beadid, YY], self.coords[current_beadid, ZZ], - self.coords_bbox[current_beadid, XX], self.coords_bbox[current_beadid, YY], self.coords_bbox[current_beadid, ZZ], - cellxyz[XX], cellxyz[YY], cellxyz[ZZ])) + # if self.debug and debug: + # with gil: + # print("FastNS: Checking neighbors for bead #{} ({:.3f},{:.3f},{:.3f})->rect({:.3f},{:.3f},{:.3f}) - cell[{},{},{}]:" .format( + # current_beadid, + # self.coords[current_beadid, XX], self.coords[current_beadid, YY], self.coords[current_beadid, ZZ], + # self.coords_bbox[current_beadid, XX], self.coords_bbox[current_beadid, YY], self.coords_bbox[current_beadid, ZZ], + # cellxyz[XX], cellxyz[YY], cellxyz[ZZ])) for xi in range(DIM): @@ -1067,25 +747,25 @@ cdef class FastNS(object): # If box is not triclinic (ie rect), when can already check if the shift can be skipped (ie cutoff inside the cell) if xi == 0: if self.coords_bbox[current_beadid, XX] - self.cutoff > self.grid.cellsize[XX] * cellxyz[XX]: - if self.debug and debug: - with gil: - print("FastNS: -> Bead X={:.3f}, Cell X={:.3f}, cutoff={:.3f} -> -X shift ignored".format( - self.coords_bbox[current_beadid, XX], - self.grid.cellsize[XX] * cellxyz[XX], - self.cutoff - )) + # if self.debug and debug: + # with gil: + # print("FastNS: -> Bead X={:.3f}, Cell X={:.3f}, cutoff={:.3f} -> -X shift ignored".format( + # self.coords_bbox[current_beadid, XX], + # self.grid.cellsize[XX] * cellxyz[XX], + # self.cutoff + # )) continue if xi == 2: if self.coords_bbox[current_beadid, XX] + self.cutoff < self.grid.cellsize[XX] * (cellxyz[XX] + 1): - if self.debug and debug: - with gil: - print( - "FastNS: -> Bead X={:.3f}, Next cell X={:.3f}, cutoff={:.3f} -> +X shift ignored".format( - self.coords_bbox[current_beadid, XX], - self.grid.cellsize[XX] * (cellxyz[XX] + 1), - self.cutoff - )) + # if self.debug and debug: + # with gil: + # print( + # "FastNS: -> Bead X={:.3f}, Next cell X={:.3f}, cutoff={:.3f} -> +X shift ignored".format( + # self.coords_bbox[current_beadid, XX], + # self.grid.cellsize[XX] * (cellxyz[XX] + 1), + # self.cutoff + # )) continue for yi in range(DIM): @@ -1095,24 +775,24 @@ cdef class FastNS(object): if not self.box.is_triclinic: if yi == 0: if self.coords_bbox[current_beadid, YY] - self.cutoff > self.grid.cellsize[YY] * cellxyz[YY]: - if self.debug and debug: - with gil: - print("FastNS: -> Bead Y={:.3f}, Cell Y={:.3f}, cutoff={:.3f} -> -Y shift is ignored".format( - self.coords_bbox[current_beadid, YY], - self.grid.cellsize[YY] * cellxyz[YY], - self.cutoff, - )) + # if self.debug and debug: + # with gil: + # print("FastNS: -> Bead Y={:.3f}, Cell Y={:.3f}, cutoff={:.3f} -> -Y shift is ignored".format( + # self.coords_bbox[current_beadid, YY], + # self.grid.cellsize[YY] * cellxyz[YY], + # self.cutoff, + # )) continue if yi == 2: if self.coords_bbox[current_beadid, YY] + self.cutoff < self.grid.cellsize[YY] * (cellxyz[YY] + 1): - if self.debug and debug: - with gil: - print("FastNS: -> Bead Y={:.3f}, Next cell Y={:.3f}, cutoff={:.3f} -> +Y shift is ignored".format( - self.coords_bbox[current_beadid, YY], - self.grid.cellsize[YY] * (cellxyz[YY] +1), - self.cutoff, - )) + # if self.debug and debug: + # with gil: + # print("FastNS: -> Bead Y={:.3f}, Next cell Y={:.3f}, cutoff={:.3f} -> +Y shift is ignored".format( + # self.coords_bbox[current_beadid, YY], + # self.grid.cellsize[YY] * (cellxyz[YY] +1), + # self.cutoff, + # )) continue @@ -1121,17 +801,17 @@ cdef class FastNS(object): if zi == 0: if self.coords_bbox[current_beadid, ZZ] - self.cutoff > self.grid.cellsize[ZZ] * cellxyz[ZZ]: if self.coords_bbox[current_beadid, ZZ] - self.cutoff > 0: - if self.debug and debug: - with gil: - print("FastNS: -> Bead Z={:.3f}, Cell Z={:.3f}, cutoff={:.3f} -> -Z shift is ignored".format(self.coords_bbox[current_beadid, ZZ], self.grid.cellsize[ZZ] * cellxyz[ZZ], self.cutoff)) + # if self.debug and debug: + # with gil: + # print("FastNS: -> Bead Z={:.3f}, Cell Z={:.3f}, cutoff={:.3f} -> -Z shift is ignored".format(self.coords_bbox[current_beadid, ZZ], self.grid.cellsize[ZZ] * cellxyz[ZZ], self.cutoff)) continue if zi == 2: if self.coords_bbox[current_beadid, ZZ] + self.cutoff < self.grid.cellsize[ZZ] * (cellxyz[ZZ] + 1): if self.coords_bbox[current_beadid, ZZ] + self.cutoff < self.box.c_pbcbox.box[ZZ][ZZ]: - if self.debug and debug: - with gil: - print("FastNS: -> Bead Z={:.3f}, Next cell Z={:.3f}, cutoff={:.3f} -> +Z shift is ignored".format(self.coords_bbox[current_beadid, ZZ], self.grid.cellsize[XX] * (cellxyz[ZZ] + 1), self.cutoff)) + # if self.debug and debug: + # with gil: + # print("FastNS: -> Bead Z={:.3f}, Next cell Z={:.3f}, cutoff={:.3f} -> +Z shift is ignored".format(self.coords_bbox[current_beadid, ZZ], self.grid.cellsize[XX] * (cellxyz[ZZ] + 1), self.cutoff)) continue # Calculate and/or reinitialize shifted coordinates @@ -1169,29 +849,29 @@ cdef class FastNS(object): cellindex_probe = self.grid.coord2cellid(probe) if cellindex == cellindex_probe and xi != 1 and yi != 1 and zi != 1: - if self.debug and debug: - with gil: - print("FastNS: Grid shift [{}][{}][{}]: Cutoff is inside current cell -> This shift is ignored".format( - xi - 1, - yi -1, - zi -1 - )) + # if self.debug and debug: + # with gil: + # print("FastNS: Grid shift [{}][{}][{}]: Cutoff is inside current cell -> This shift is ignored".format( + # xi - 1, + # yi -1, + # zi -1 + # )) continue - if self.debug and debug: - self.grid.cellid2cellxyz(cellindex_adjacent, debug_cellxyz) - with gil: - dist_shift = self.box.fast_distance(&self.coords[current_beadid, XX], shifted_coord) - grid_shift = np.array([(xi - 1) * self.grid.cellsize[XX], - (yi - 1) * self.grid.cellsize[YY], - (zi - 1) * self.grid.cellsize[ZZ]]) - print("FastNS: -> Checking cell#{} ({},{},{}) for neighbors (dshift={:.3f}, grid_shift=({:.3f},{:.3f},{:.3f}->{:.3f})".format( - cellindex, - debug_cellxyz[XX], debug_cellxyz[YY], debug_cellxyz[ZZ], - dist_shift, - grid_shift[XX], grid_shift[YY], grid_shift[ZZ], - np.sqrt(np.sum(grid_shift**2)) - )) + # if self.debug and debug: + # self.grid.cellid2cellxyz(cellindex_adjacent, debug_cellxyz) + # with gil: + # dist_shift = self.box.fast_distance(&self.coords[current_beadid, XX], shifted_coord) + # grid_shift = np.array([(xi - 1) * self.grid.cellsize[XX], + # (yi - 1) * self.grid.cellsize[YY], + # (zi - 1) * self.grid.cellsize[ZZ]]) + # print("FastNS: -> Checking cell#{} ({},{},{}) for neighbors (dshift={:.3f}, grid_shift=({:.3f},{:.3f},{:.3f}->{:.3f})".format( + # cellindex, + # debug_cellxyz[XX], debug_cellxyz[YY], debug_cellxyz[ZZ], + # dist_shift, + # grid_shift[XX], grid_shift[YY], grid_shift[ZZ], + # np.sqrt(np.sum(grid_shift**2)) + # )) for j in range(self.grid.nbeads[cellindex_adjacent]): @@ -1223,12 +903,12 @@ cdef class FastNS(object): if d2 < EPSILON: continue - if self.debug and debug: - self.grid.cellid2cellxyz(cellindex, debug_cellxyz) - with gil: - self.box.fast_pbc_dx(&self.coords[current_beadid, XX], &self.coords[bid, XX], dx) - dx_py = np.array([dx[XX], dx[YY], dx[ZZ]]) - print("FastNS: \_ Neighbor found: bead#{} (cell[{},{},{}]) -> dx={} -> d={:.3f}".format(bid, debug_cellxyz[XX], debug_cellxyz[YY], debug_cellxyz[ZZ], dx_py, np.sqrt(d2))) + # if self.debug and debug: + # self.grid.cellid2cellxyz(cellindex, debug_cellxyz) + # with gil: + # self.box.fast_pbc_dx(&self.coords[current_beadid, XX], &self.coords[bid, XX], dx) + # dx_py = np.array([dx[XX], dx[YY], dx[ZZ]]) + # print("FastNS: \_ Neighbor found: bead#{} (cell[{},{},{}]) -> dx={} -> d={:.3f}".format(bid, debug_cellxyz[XX], debug_cellxyz[YY], debug_cellxyz[ZZ], dx_py, np.sqrt(d2))) if results.add_neighbors(current_beadid, bid, d2) == 0: memory_error = True break @@ -1260,183 +940,4 @@ cdef class FastNS(object): # self.box.fast_distance(&self.coords[ref_bead, XX], &self.coords[bid, XX]) <= self.cutoff, # )) - return results - - -# Python interface -cdef class FastNSOld(object): - cdef PBCBox box - cdef readonly real[:, ::1] coords - cdef real[:, ::1] coords_bbox - cdef readonly real cutoff - cdef bint prepared - cdef ns_grid grid - - def __init__(self, u): - import MDAnalysis as mda - from MDAnalysis.lib.mdamath import triclinic_vectors - - if not isinstance(u, mda.Universe): - raise TypeError("FastNS class must be initialized with a valid MDAnalysis.Universe instance") - box = triclinic_vectors(u.dimensions) - - if box.shape != (3, 3): - raise ValueError("Box must be provided as triclinic_dimensions (a 3x3 numpy.ndarray of unit cell vectors") - - self.box = PBCBox(box) - - self.coords = None - self.coords_bbox = None - - self.cutoff = -1 - - self.prepared = False - - self.grid.size = 0 - - - def __dealloc__(self): - cdef ns_int i - - # Deallocate NS grid - if self.grid.nbeads != NULL: - free(self.grid.nbeads) - - if self.grid.beadids != NULL: - free(self.grid.beadids) - - if self.grid.cellids != NULL: - free(self.grid.cellids) - - self.grid.size = 0 - - - def set_coords(self, real[:, ::1] coords): - self.coords = coords - - # Make sure atoms are inside the brick-shaped box - self.coords_bbox = self.box.fast_put_atoms_in_bbox(coords) - - self.prepared = False - - - def set_cutoff(self, real cutoff): - self.cutoff = cutoff - - self.prepared = False - - - def prepare(self, force=False): - cdef ns_int i, cellindex = -1 - cdef ns_int ncoords = self.coords.shape[0] - cdef rvec *coords = &self.coords_bbox[0, 0] - - cdef ns_int *beadcounts = NULL - - if self.prepared and not force: - print("NS already prepared, nothing to do!") - - if self.coords is None: - raise ValueError("Coordinates must be set before NS preparation!") - - if self.cutoff < 0: - raise ValueError("Cutoff must be set before NS preparation!") - - with nogil: - - # Initializing grid - for i in range(DIM): - self.grid.ncells[i] = (self.box.c_pbcbox.box[i][i] / self.cutoff) - if self.grid.ncells[i] == 0: - self.grid.ncells[i] = 1 - self.grid.cellsize[i] = self.box.c_pbcbox.box[i][i] / self.grid.ncells[i] - self.grid.size = self.grid.ncells[XX] * self.grid.ncells[YY] * self.grid.ncells[ZZ] - - # Allocate memory for temporary buffers - self.grid.cellids = malloc(sizeof(ns_int) * ncoords) - if self.grid.cellids == NULL: - with gil: - raise MemoryError("Could not allocate memory for cell ids") - - beadcounts = malloc(sizeof(ns_int) * self.grid.size) - if beadcounts == NULL: - with gil: - raise MemoryError("Could not allocate memory for bead count buffer") - - self.grid.nbeads = malloc(sizeof(ns_int) * self.grid.size) - if self.grid.nbeads == NULL: - with gil: - raise MemoryError("Could not allocate memory for NS grid") - - # Initialize buffers - for i in range(self.grid.size): - self.grid.nbeads[i] = 0 - beadcounts[i] = 0 - - - # First loop: find cellindex for each bead - for i in range(ncoords): - cellindex = (coords[i][ZZ] / self.grid.cellsize[ZZ]) * (self.grid.ncells[XX] * self.grid.ncells[YY]) +\ - (coords[i][YY] / self.grid.cellsize[YY]) * self.grid.ncells[XX] + \ - (coords[i][XX] / self.grid.cellsize[XX]) - - self.grid.nbeads[cellindex] += 1 - self.grid.cellids[i] = cellindex - - if self.grid.nbeads[cellindex] > self.grid.beadpercell: - self.grid.beadpercell = self.grid.nbeads[cellindex] - - # Allocate memory and set offsets - self.grid.beadids = malloc(sizeof(ns_int) * self.grid.size * self.grid.beadpercell) - if self.grid.beadids == NULL: - with gil: - raise MemoryError("Could not allocate memory for NS grid bead ids") - - # Second loop: fill grid - for i in range(ncoords): - cellindex = self.grid.cellids[i] - self.grid.beadids[cellindex * self.grid.beadpercell + beadcounts[cellindex]] = i - beadcounts[cellindex] += 1 - - # Now we can free the allocation buffer - free(beadcounts) - - self.prepared = True - - - def search(self, search_coords): - cdef real[:, ::1] search_coords_bbox - cdef real[:, ::1] search_coords_view - cdef ns_int nid, i, j - cdef ns_neighborhood_holder *holder - cdef ns_neighborhood *neighborhood - - if not self.prepared: - self.prepare() - - # Check the shape of search_coords as a array of 3D coords if needed - shape = search_coords.shape - if len(shape) == 1: - if not shape[0] == 3: - raise ValueError("Coordinates must be 3D") - else: - search_coords_view = np.array([search_coords,], dtype=np.float32) - else: - search_coords_view = search_coords - - # Make sure atoms are inside the brick-shaped box - search_coords_bbox = self.box.fast_put_atoms_in_bbox(search_coords_view) - - with nogil: - holder = ns_core(search_coords_bbox, self.coords_bbox, &self.grid, self.box, self.cutoff) - - if holder == NULL: - raise MemoryError("Could not allocate memory to run NS core") - - results = NSResultsOld(self.box, self.cutoff) - results.populate(holder, self.coords, search_coords_view) - - # Free memory allocated to holder - free_neighborhood_holder(holder) - - return results + return results \ No newline at end of file diff --git a/testsuite/MDAnalysisTests/lib/test_nsgrid.py b/testsuite/MDAnalysisTests/lib/test_nsgrid.py index d7be4f79e65..387b7119578 100644 --- a/testsuite/MDAnalysisTests/lib/test_nsgrid.py +++ b/testsuite/MDAnalysisTests/lib/test_nsgrid.py @@ -24,11 +24,11 @@ import pytest -from numpy.testing import assert_equal +from numpy.testing import assert_equal, assert_allclose import numpy as np import MDAnalysis as mda -from MDAnalysisTests.datafiles import GRO +from MDAnalysisTests.datafiles import GRO, Martini_membrane_gro from MDAnalysis.lib import nsgrid @@ -45,12 +45,95 @@ def run_grid_search(u, ids, cutoff=3): # Run grid search searcher = nsgrid.FastNS(u, cutoff, coords, debug=True) + return searcher.search(ids) - return searcher.search(ids, debug=False) +def test_pbc_badbox(): + """Check that PBC box accepts only well-formated boxes""" + with pytest.raises(TypeError): + nsgrid.PBCBox([]) -def test_gridsearch(universe): - """Check that pkdtree and grid search return the same results (No PBC needed)""" + with pytest.raises(ValueError): + nsgrid.PBCBox(np.zeros((3))) # Bad shape + nsgrid.PBCBox(np.zeros((3, 3))) # Collapsed box + nsgrid.PBCBOX(np.array([[0, 0, 0], [0, 1, 0], [0, 0, 1]])) # 2D box + nsgrid.PBCBOX(np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])) # Box provided as array of integers + nsgrid.PBCBOX(np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=float)) # Box provided as array of double + + +def test_pbc_distances(): + """Check that PBC box computes distances""" + box = np.identity(3, dtype=np.float32) + bad = np.array([0.1, 0.2], dtype=np.float32) + a = np.array([0.1, 0.1, 0.1], dtype=np.float32) + b = np.array([1.1, -0.1, 0.2], dtype=np.float32) + dx = np.array([0, -0.2, 0.1], dtype=np.float32) + pbcbox = nsgrid.PBCBox(box) + + with pytest.raises(ValueError): + pbcbox.distance(a, bad) + pbcbox.distance(bad, a) + + pbcbox.distance2(a, bad) + pbcbox.distance2(bad, a) + + pbcbox.dx(bad, a) + pbcbox.dx(a, bad) + + assert_equal(pbcbox.dx(a, b), dx) + assert pbcbox.distance(a, b) == np.sqrt(np.sum(dx*dx)) + assert pbcbox.distance2(a, b) == np.sum(dx*dx) + + +def test_pbc_put_in_bbox(): + "Check that PBC put beads in brick-shaped box" + box = np.identity(3, dtype=np.float32) + coords = np.array( + [ + [0.1, 0.1, 0.1], + [-0.1, 1.1, 0.9] + ], + dtype=np.float32 + ) + results = np.array( + [ + [0.1, 0.1, 0.1], + [0.9, 0.1, 0.9] + ], + dtype=np.float32 + ) + + pbcbox = nsgrid.PBCBox(box) + + assert_allclose(pbcbox.put_atoms_in_bbox(coords), results, atol=1e-5) + + +def test_nsgrid_badinit(): + with pytest.raises(TypeError): + nsgrid.FastNS(None, 1) + +def test_nsgrid_badcutoff(universe): + with pytest.raises(ValueError): + run_grid_search(universe, 0, -4) + run_grid_search(universe, 0, 100000) + +def test_ns_grid_noneighbor(universe): + """Check that grid search returns empty lists/arrays when there is no neighbors""" + ref_id = 0 + cutoff = 0.5 + + results_grid = run_grid_search(universe, ref_id, cutoff) + + assert len(results_grid.get_coordinates()[0]) == 0 + assert len(results_grid.get_distances()[0]) == 0 + assert len(results_grid.get_indices()[0]) == 0 + assert len(results_grid.get_pairs()) == 0 + assert len(results_grid.get_pair_distances()) == 0 + assert len(results_grid.get_pair_coordinates()) == 0 + + +def test_nsgrid_noPBC(universe): + """Check that grid search works when no PBC is needed""" ref_id = 0 @@ -63,8 +146,31 @@ def test_gridsearch(universe): assert_equal(results, results_grid) -def test_gridsearch_PBC(universe): - """Check that pkdtree and grid search return the same results (No PBC needed)""" +def test_nsgrid_PBC_rect(): + """Check that nsgrid works with rect boxes and PBC""" + ref_id = 191 + results = np.array([191, 672, 682, 683, 684, 995, 996, 2060, 2808, 3300, 3791, + 3792]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! + + universe = mda.Universe(Martini_membrane_gro) + cutoff = 7 + + # FastNS is called differently to max coverage + searcher = nsgrid.FastNS(universe, cutoff, prepare=False) + + results_grid = searcher.search([ref_id,]).get_indices()[0] # pass the id as a list for test+coverage purpose + + searcher.prepare() # Does nothing, called here for coverage + results_grid2 = searcher.search().get_indices() # call without specifying any ids, should do NS for all beads + + assert_equal(results, results_grid) + assert_equal(len(universe.atoms), len(results_grid2)) + assert searcher.cutoff == 7 + assert_equal(results_grid, results_grid2[ref_id]) + + +def test_nsgrid_PBC(universe): + """Check that grid search works when PBC is needed""" ref_id = 13937 results = np.array([4398, 4401, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, @@ -76,51 +182,80 @@ def test_gridsearch_PBC(universe): -# def test_gridsearch_PBC(universe): -# """Check that pkdtree and grid search return the same results (PBC needed)""" -# -# ref_id = 13937 -# results_pkdtree, results_grid = run_search(universe, ref_id) -# assert_equal(results_pkdtree, results_grid) -# -# -# def test_gridsearch_arraycoord(universe): -# """Check the NS routine accepts a single bead coordinate as well as array of coordinates""" -# cutoff = 2 -# ref_pos = universe.atoms.positions[:5] -# -# results = [ -# np.array([2, 1, 4, 3]), -# np.array([2, 0, 3]), -# np.array([0, 1, 3]), -# np.array([ 2, 0, 1, 38341]), -# np.array([ 6, 0, 5, 17]) -# ] -# -# results_grid = run_grid_search(universe, ref_pos, cutoff).get_indices() -# -# assert_equal(results_grid, results) -# -# -# def test_gridsearch_search_coordinates(grid_results): -# """Check the NS routine can return coordinates instead of ids""" -# -# results = np.array( -# [ -# [40.32, 34.25, 55.9], -# [0.61, 76.33, -0.56], -# [0.48999998, 75.9, 0.19999999], -# [-0.11, 76.19, 0.77] -# ]) -# -# assert_allclose(grid_results.get_coordinates()[0], results) -# -# -# def test_gridsearch_search_distances(grid_results): -# """Check the NS routine can return PBC distances from neighbors""" -# results = np.array([0.096, 0.096, 0.015, 0.179]) * 10 # These distances were obtained using gmx distance -# results.sort() -# -# rounded_results = np.round(grid_results.get_distances()[0], 2) -# -# assert_allclose(sorted(rounded_results), results) +def test_nsgrid_pairs(universe): + """Check that grid search returns the proper pairs""" + + ref_id = 13937 + neighbors = np.array([4398, 4401, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, + 47451]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! + results = [] + for nid in neighbors: + if nid < ref_id: + results.append([nid, ref_id]) + else: + results.append([ref_id, nid]) + results = np.array(results) + + results_grid = run_grid_search(universe, ref_id).get_pairs() + + assert_equal(np.sort(results, axis=0), np.sort(results_grid, axis=0)) + + +def test_nsgrid_pair_distances(universe): + """Check that grid search returns the proper pair distances""" + + ref_id = 13937 + results = np.array([0.270, 0.285, 0.096, 0.096, 0.015, 0.278, 0.268, 0.179, 0.259, 0.290, + 0.270]) * 10 # These distances where obtained by gmx distance so they are in nm + + results_grid = run_grid_search(universe, ref_id).get_pair_distances() + + assert_allclose(np.sort(results), np.sort(results_grid), atol=1e-2) + + +def test_nsgrid_pair_coordinates(universe): + """Check that grid search return the proper pair coordinates""" + + ref_id = 13937 + neighbors = np.array([4398, 4401, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, + 47451]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! + coords = universe.atoms.positions + + results = [] + for nid in neighbors: + if nid < ref_id: + results.append([coords[nid], coords[ref_id]]) + else: + results.append([coords[ref_id], coords[nid]]) + results = np.array(results) + + results_grid = run_grid_search(universe, ref_id).get_pair_coordinates() + + assert_allclose(np.sort(results, axis=0), np.sort(results_grid, axis=0), atol=1e-5) + + +def test_nsgrid_distances(universe): + """Check that grid search returns the proper distances""" + + ref_id = 13937 + results = np.array([0.270, 0.285, 0.096, 0.096, 0.015, 0.278, 0.268, 0.179, 0.259, 0.290, + 0.270]) * 10 # These distances where obtained by gmx distance so they are in nm + + results_grid = run_grid_search(universe, ref_id).get_distances()[0] + + assert_allclose(np.sort(results), np.sort(results_grid), atol=1e-2) + + +def test_nsgrid_coordinates(universe): + """Check that grid search return the proper coordinates""" + + ref_id = 13937 + neighbors = np.array([4398, 4401, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, + 47451]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! + + results = universe.atoms.positions[neighbors] + + results_grid = run_grid_search(universe, ref_id).get_coordinates()[0] + + assert_allclose(np.sort(results, axis=0), np.sort(results_grid, axis=0), atol=1e-5) + From 8a75a5fe81259ca5f2baee8ee780d3266a23b4cb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Buchoux?= Date: Wed, 18 Jul 2018 17:29:27 +0200 Subject: [PATCH 23/47] Removed all references to c_gridsearch.pyx/grid module that reappeared after merge/rebase --- .../doc/sphinx/source/documentation_pages/lib_modules.rst | 4 ++-- package/setup.py | 7 ++++--- testsuite/MDAnalysisTests/lib/test_nsgrid.py | 4 ++-- 3 files changed, 8 insertions(+), 7 deletions(-) diff --git a/package/doc/sphinx/source/documentation_pages/lib_modules.rst b/package/doc/sphinx/source/documentation_pages/lib_modules.rst index 83d7f81cfd9..f918f37137e 100644 --- a/package/doc/sphinx/source/documentation_pages/lib_modules.rst +++ b/package/doc/sphinx/source/documentation_pages/lib_modules.rst @@ -36,7 +36,7 @@ functions whereas mathematical functions are to be found in :mod:`MDAnalysis.lib.NeighborSearch` contains classes to do neighbor searches with MDAnalysis objects. -:mod:`MDAnalysis.lib.grid` contains a fast implementation of grid neighbor search. +:mod:`MDAnalysis.lib.nsgrid` contains a fast implementation of grid neighbor search. List of modules --------------- @@ -51,7 +51,7 @@ List of modules ./lib/transformations ./lib/qcprot ./lib/util - ./lib/grid + ./lib/nsgrid Low level file formats ---------------------- diff --git a/package/setup.py b/package/setup.py index bd3aceec8df..52d29fecc7c 100755 --- a/package/setup.py +++ b/package/setup.py @@ -380,8 +380,8 @@ def extensions(config): libraries=mathlib, define_macros=define_macros, extra_compile_args=extra_compile_args) - grid = MDAExtension('MDAnalysis.lib.grid', - ['MDAnalysis/lib/c_gridsearch' + source_suffix], + nsgrid = MDAExtension('MDAnalysis.lib.nsgrid', + ['MDAnalysis/lib/nsgrid' + source_suffix], include_dirs=include_dirs, libraries=mathlib + parallel_libraries, define_macros=define_macros + parallel_macros, @@ -389,7 +389,8 @@ def extensions(config): extra_link_args=parallel_args) pre_exts = [libdcd, distances, distances_omp, qcprot, transformation, libmdaxdr, util, encore_utils, - ap_clustering, spe_dimred, grid] + ap_clustering, spe_dimred, cutil, nsgrid] + cython_generated = [] if use_cython: diff --git a/testsuite/MDAnalysisTests/lib/test_nsgrid.py b/testsuite/MDAnalysisTests/lib/test_nsgrid.py index 387b7119578..0fc6e1ec880 100644 --- a/testsuite/MDAnalysisTests/lib/test_nsgrid.py +++ b/testsuite/MDAnalysisTests/lib/test_nsgrid.py @@ -81,8 +81,8 @@ def test_pbc_distances(): pbcbox.dx(a, bad) assert_equal(pbcbox.dx(a, b), dx) - assert pbcbox.distance(a, b) == np.sqrt(np.sum(dx*dx)) - assert pbcbox.distance2(a, b) == np.sum(dx*dx) + assert_allclose(pbcbox.distance(a, b), np.sqrt(np.sum(dx*dx)), atol=1e-5) + assert_allclose(pbcbox.distance2(a, b), np.sum(dx*dx), atol=1e-5) def test_pbc_put_in_bbox(): From 7b0c56c99d5e99c4f4d092ba7d951b4052b89e76 Mon Sep 17 00:00:00 2001 From: Max Linke Date: Sat, 21 Jul 2018 14:30:23 +0200 Subject: [PATCH 24/47] remove unused code --- package/MDAnalysis/lib/nsgrid.pyx | 93 ++----------------------------- 1 file changed, 5 insertions(+), 88 deletions(-) diff --git a/package/MDAnalysis/lib/nsgrid.pyx b/package/MDAnalysis/lib/nsgrid.pyx index e4de849ffff..dc4017191c0 100644 --- a/package/MDAnalysis/lib/nsgrid.pyx +++ b/package/MDAnalysis/lib/nsgrid.pyx @@ -69,20 +69,15 @@ cdef void rvec_clear(rvec a) nogil: a[YY]=0.0 a[ZZ]=0.0 -######################################################################################################################## -# -# Utility class to handle PBC -# -######################################################################################################################## +############################### +# Utility class to handle PBC # +############################### cdef struct cPBCBox_t: matrix box rvec fbox_diag rvec hbox_diag rvec mhbox_diag real max_cutoff2 - ns_int ntric_vec - ns_int[DIM] tric_shift[MAX_NTRICVEC] - real[DIM] tric_vec[MAX_NTRICVEC] # Class to handle PBC calculations @@ -141,90 +136,12 @@ cdef class PBCBox(object): tmp += box[ZZ, YY] min_ss = min(box[XX, XX], min(tmp, box[ZZ, ZZ])) - self.c_pbcbox.max_cutoff2 = min(min_hv2, min_ss * min_ss) - # Update shift vectors - self.c_pbcbox.ntric_vec = 0 - - # We will only use single shifts - for kk in range(3): - k = order[kk] - - for jj in range(3): - j = order[jj] - - for ii in range(3): - i = order[ii] - - # A shift is only useful when it is trilinic - if j != 0 or k != 0: - d2old = 0 - d2new = 0 - - for d in range(DIM): - trial[d] = i*box[XX, d] + j*box[YY, d] + k*box[ZZ, d] - - # Choose the vector within the brick around 0,0,0 that - # will become the shortest due to shift try. - if trial[d] < 0: - pos[d] = min(self.c_pbcbox.hbox_diag[d], -trial[d]) - else: - pos[d] = max(-self.c_pbcbox.hbox_diag[d], -trial[d]) - - d2old += pos[d]**2 - d2new += (pos[d] + trial[d])**2 - - if BOX_MARGIN*d2new < d2old: - if not (j < -1 or j > 1 or k < -1 or k > 1): - use = True - - for dd in range(DIM): - if dd == 0: - shift = i - elif dd == 1: - shift = j - else: - shift = k - - if shift: - d2new_c = 0 - - for d in range(DIM): - d2new_c += (pos[d] + trial[d] - shift*box[dd, d])**2 - - if d2new_c <= BOX_MARGIN*d2new: - use = False - - if use: # Accept this shift vector. - if self.c_pbcbox.ntric_vec >= MAX_NTRICVEC: - with gil: - print("\nWARNING: Found more than %d triclinic " - "correction vectors, ignoring some." - % MAX_NTRICVEC) - print(" There is probably something wrong with " - "your box.") - print(np.array(box)) - - for i in range(self.c_pbcbox.ntric_vec): - print(" -> shift #{}: [{}, {}, {}]".format(i+1, - self.c_pbcbox.tric_shift[i][XX], - self.c_pbcbox.tric_shift[i][YY], - self.c_pbcbox.tric_shift[i][ZZ])) - else: - for d in range(DIM): - self.c_pbcbox.tric_vec[self.c_pbcbox.ntric_vec][d] = \ - trial[d] - self.c_pbcbox.tric_shift[self.c_pbcbox.ntric_vec][XX] = i - self.c_pbcbox.tric_shift[self.c_pbcbox.ntric_vec][YY] = j - self.c_pbcbox.tric_shift[self.c_pbcbox.ntric_vec][ZZ] = k - self.c_pbcbox.ntric_vec += 1 - - def update(self, real[:,::1] box): if box.shape[0] != DIM or box.shape[1] != DIM: - raise ValueError("Box must be a %i x %i matrix. (shape: %i x %i)" % - (DIM, DIM, box.shape[0], box.shape[1])) + raise ValueError("Box must be a {} x {} matrix. Got: {} x {})".format( + DIM, DIM, box.shape[0], box.shape[1])) if (box[XX, XX] == 0) or (box[YY, YY] == 0) or (box[ZZ, ZZ] == 0): raise ValueError("Box does not correspond to PBC=xyz") self.fast_update(box) From 37658683132363094a86776671887e88742c79cd Mon Sep 17 00:00:00 2001 From: Max Linke Date: Sat, 21 Jul 2018 14:30:43 +0200 Subject: [PATCH 25/47] make init_callable twice --- package/MDAnalysis/lib/nsgrid.pyx | 28 ++++++++++++++++------------ 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/package/MDAnalysis/lib/nsgrid.pyx b/package/MDAnalysis/lib/nsgrid.pyx index dc4017191c0..f051e7e413d 100644 --- a/package/MDAnalysis/lib/nsgrid.pyx +++ b/package/MDAnalysis/lib/nsgrid.pyx @@ -221,11 +221,9 @@ cdef class PBCBox(object): return np.asarray(self.fast_put_atoms_in_bbox(coords)) -######################################################################################################################## -# -# Neighbor Search Stuff -# -######################################################################################################################## +######################### +# Neighbor Search Stuff # +######################### cdef class NSResults(object): cdef readonly real cutoff cdef ns_int npairs @@ -246,7 +244,6 @@ cdef class NSResults(object): cdef np.ndarray pair_coordinates_buffer def __init__(self, real cutoff, real[:, ::1]coords, ns_int[:] search_ids, debug=False): - self.debug = debug self.cutoff = cutoff self.coords = coords @@ -254,12 +251,19 @@ cdef class NSResults(object): # Preallocate memory self.allocation_size = search_ids.shape[0] + 1 - self.pairs = PyMem_Malloc(sizeof(ipair) * self.allocation_size) - if not self.pairs: - MemoryError("Could not allocate memory for NSResults.pairs ({} bits requested)".format(sizeof(ipair) * self.allocation_size)) - self.pair_distances2 = PyMem_Malloc(sizeof(real) * self.allocation_size) - if not self.pair_distances2: - raise MemoryError("Could not allocate memory for NSResults.pair_distances2 ({} bits requested)".format(sizeof(real) * self.allocation_size)) + if not self.pairs and not self.pair_distances2: + self.pairs = PyMem_Malloc(sizeof(ipair) * self.allocation_size) + if not self.pairs: + MemoryError("Could not allocate memory for NSResults.pairs " + "({} bits requested)".format(sizeof(ipair) * self.allocation_size)) + self.pair_distances2 = PyMem_Malloc(sizeof(real) * self.allocation_size) + if not self.pair_distances2: + raise MemoryError("Could not allocate memory for NSResults.pair_distances2 " + "({} bits requested)".format(sizeof(real) * self.allocation_size)) + else: + if self.resize(self.allocation_size + (self.allocation_size * 0.5 + 1)) == 0: + raise MemoryError("foo") + self.npairs = 0 # Buffer From f23000868a1b3b8780b7c9afd0007ac5711d2225 Mon Sep 17 00:00:00 2001 From: Max Linke Date: Sat, 21 Jul 2018 15:34:44 +0200 Subject: [PATCH 26/47] use OK and ERROR as return codes --- package/MDAnalysis/lib/nsgrid.pyx | 32 ++++++++++++++----------------- 1 file changed, 14 insertions(+), 18 deletions(-) diff --git a/package/MDAnalysis/lib/nsgrid.pyx b/package/MDAnalysis/lib/nsgrid.pyx index f051e7e413d..01bf67e0c4c 100644 --- a/package/MDAnalysis/lib/nsgrid.pyx +++ b/package/MDAnalysis/lib/nsgrid.pyx @@ -46,6 +46,9 @@ DEF EPSILON = 1e-5 DEF BOX_MARGIN=1.0010 DEF MAX_NTRICVEC=12 +DEF OK=0 +DEF ERROR=1 + # Used to handle memory allocation from cpython.mem cimport PyMem_Malloc, PyMem_Realloc, PyMem_Free from libc.math cimport sqrt @@ -261,7 +264,7 @@ cdef class NSResults(object): raise MemoryError("Could not allocate memory for NSResults.pair_distances2 " "({} bits requested)".format(sizeof(real) * self.allocation_size)) else: - if self.resize(self.allocation_size + (self.allocation_size * 0.5 + 1)) == 0: + if self.resize(self.allocation_size) != OK: raise MemoryError("foo") self.npairs = 0 @@ -278,13 +281,13 @@ cdef class NSResults(object): PyMem_Free(self.pair_distances2) cdef int add_neighbors(self, ns_int beadid_i, ns_int beadid_j, real distance2) nogil: - # Important: If this function returns 0, it means that memory allocation failed + # Important: If this function returns ERROR, it means that memory allocation failed # Reallocate memory if needed if self.npairs >= self.allocation_size: # We need to reallocate memory - if self.resize(self.allocation_size + (self.allocation_size * 0.5 + 1)) == 0: - return 0 + if self.resize(self.allocation_size + (self.allocation_size * 0.5 + 1)) != OK: + return ERROR # Actually store pair and distance squared if beadid_i < beadid_j: @@ -296,20 +299,20 @@ cdef class NSResults(object): self.pair_distances2[self.npairs] = distance2 self.npairs += 1 - return self.npairs + return OK cdef int resize(self, ns_int new_size) nogil: # Important: If this function returns 0, it means that memory allocation failed if new_size < self.npairs: # Silently ignored the request - return 1 + return OK if self.allocation_size >= new_size: if self.debug: with gil: print("NSresults: Reallocation requested but not needed ({} requested but {} already allocated)".format(new_size, self.allocation_size)) - return 1 + return OK self.allocation_size = new_size @@ -323,12 +326,12 @@ cdef class NSResults(object): self.pair_distances2 = PyMem_Realloc(self.pair_distances2, sizeof(real) * self.allocation_size) if not self.pairs: - return 0 + return ERROR if not self.pair_distances2: - return 0 + return ERROR - return 1 + return OK def get_pairs(self): cdef ns_int i @@ -823,14 +826,7 @@ cdef class FastNS(object): if d2 < EPSILON: continue - - # if self.debug and debug: - # self.grid.cellid2cellxyz(cellindex, debug_cellxyz) - # with gil: - # self.box.fast_pbc_dx(&self.coords[current_beadid, XX], &self.coords[bid, XX], dx) - # dx_py = np.array([dx[XX], dx[YY], dx[ZZ]]) - # print("FastNS: \_ Neighbor found: bead#{} (cell[{},{},{}]) -> dx={} -> d={:.3f}".format(bid, debug_cellxyz[XX], debug_cellxyz[YY], debug_cellxyz[ZZ], dx_py, np.sqrt(d2))) - if results.add_neighbors(current_beadid, bid, d2) == 0: + elif results.add_neighbors(current_beadid, bid, d2) != OK: memory_error = True break npairs += 1 From a738fa9b8bd4492e66141aca15662d5f76e7f42e Mon Sep 17 00:00:00 2001 From: Max Linke Date: Sat, 21 Jul 2018 15:34:59 +0200 Subject: [PATCH 27/47] remove dead code --- package/MDAnalysis/lib/nsgrid.pyx | 171 +----------------------------- 1 file changed, 4 insertions(+), 167 deletions(-) diff --git a/package/MDAnalysis/lib/nsgrid.pyx b/package/MDAnalysis/lib/nsgrid.pyx index 01bf67e0c4c..df6d130ef14 100644 --- a/package/MDAnalysis/lib/nsgrid.pyx +++ b/package/MDAnalysis/lib/nsgrid.pyx @@ -615,14 +615,11 @@ cdef class FastNS(object): cdef rvec shifted_coord, probe, dx cdef ns_int nchecked = 0 + cdef ns_int[:] checked = np.zeros(size, dtype=np.int) cdef real cutoff2 = self.cutoff * self.cutoff - cdef ns_int[:] checked cdef ns_int npairs = 0 - #cdef bint debug=False - - if not self.prepared: self.prepare() @@ -636,134 +633,44 @@ cdef class FastNS(object): search_ids_view = search_ids size_search = search_ids.shape[0] - checked = np.zeros(size, dtype=np.int) - results = NSResults(self.cutoff, self.coords, search_ids, self.debug) cdef bint memory_error = False - # if self.debug and debug: - # print("FastNS: Debug flag is set to True for FastNS.search()") - with nogil: for i in range(size_search): if memory_error: break current_beadid = search_ids_view[i] - cellindex = self.grid.cellids[current_beadid] self.grid.cellid2cellxyz(cellindex, cellxyz) - - # if self.debug and debug: - # with gil: - # print("FastNS: Checking neighbors for bead #{} ({:.3f},{:.3f},{:.3f})->rect({:.3f},{:.3f},{:.3f}) - cell[{},{},{}]:" .format( - # current_beadid, - # self.coords[current_beadid, XX], self.coords[current_beadid, YY], self.coords[current_beadid, ZZ], - # self.coords_bbox[current_beadid, XX], self.coords_bbox[current_beadid, YY], self.coords_bbox[current_beadid, ZZ], - # cellxyz[XX], cellxyz[YY], cellxyz[ZZ])) - - for xi in range(DIM): if memory_error: break - - if not self.box.is_triclinic: - # If box is not triclinic (ie rect), when can already check if the shift can be skipped (ie cutoff inside the cell) - if xi == 0: - if self.coords_bbox[current_beadid, XX] - self.cutoff > self.grid.cellsize[XX] * cellxyz[XX]: - # if self.debug and debug: - # with gil: - # print("FastNS: -> Bead X={:.3f}, Cell X={:.3f}, cutoff={:.3f} -> -X shift ignored".format( - # self.coords_bbox[current_beadid, XX], - # self.grid.cellsize[XX] * cellxyz[XX], - # self.cutoff - # )) - continue - - if xi == 2: - if self.coords_bbox[current_beadid, XX] + self.cutoff < self.grid.cellsize[XX] * (cellxyz[XX] + 1): - # if self.debug and debug: - # with gil: - # print( - # "FastNS: -> Bead X={:.3f}, Next cell X={:.3f}, cutoff={:.3f} -> +X shift ignored".format( - # self.coords_bbox[current_beadid, XX], - # self.grid.cellsize[XX] * (cellxyz[XX] + 1), - # self.cutoff - # )) - continue - for yi in range(DIM): if memory_error: break - - if not self.box.is_triclinic: - if yi == 0: - if self.coords_bbox[current_beadid, YY] - self.cutoff > self.grid.cellsize[YY] * cellxyz[YY]: - # if self.debug and debug: - # with gil: - # print("FastNS: -> Bead Y={:.3f}, Cell Y={:.3f}, cutoff={:.3f} -> -Y shift is ignored".format( - # self.coords_bbox[current_beadid, YY], - # self.grid.cellsize[YY] * cellxyz[YY], - # self.cutoff, - # )) - continue - - if yi == 2: - if self.coords_bbox[current_beadid, YY] + self.cutoff < self.grid.cellsize[YY] * (cellxyz[YY] + 1): - # if self.debug and debug: - # with gil: - # print("FastNS: -> Bead Y={:.3f}, Next cell Y={:.3f}, cutoff={:.3f} -> +Y shift is ignored".format( - # self.coords_bbox[current_beadid, YY], - # self.grid.cellsize[YY] * (cellxyz[YY] +1), - # self.cutoff, - # )) - continue - - for zi in range(DIM): - if not self.box.is_triclinic: - if zi == 0: - if self.coords_bbox[current_beadid, ZZ] - self.cutoff > self.grid.cellsize[ZZ] * cellxyz[ZZ]: - if self.coords_bbox[current_beadid, ZZ] - self.cutoff > 0: - # if self.debug and debug: - # with gil: - # print("FastNS: -> Bead Z={:.3f}, Cell Z={:.3f}, cutoff={:.3f} -> -Z shift is ignored".format(self.coords_bbox[current_beadid, ZZ], self.grid.cellsize[ZZ] * cellxyz[ZZ], self.cutoff)) - continue - - if zi == 2: - if self.coords_bbox[current_beadid, ZZ] + self.cutoff < self.grid.cellsize[ZZ] * (cellxyz[ZZ] + 1): - if self.coords_bbox[current_beadid, ZZ] + self.cutoff < self.box.c_pbcbox.box[ZZ][ZZ]: - # if self.debug and debug: - # with gil: - # print("FastNS: -> Bead Z={:.3f}, Next cell Z={:.3f}, cutoff={:.3f} -> +Z shift is ignored".format(self.coords_bbox[current_beadid, ZZ], self.grid.cellsize[XX] * (cellxyz[ZZ] + 1), self.cutoff)) - continue - + if memory_error: + break # Calculate and/or reinitialize shifted coordinates shifted_coord[XX] = self.coords[current_beadid, XX] + (xi - 1) * self.grid.cellsize[XX] shifted_coord[YY] = self.coords[current_beadid, YY] + (yi - 1) * self.grid.cellsize[YY] shifted_coord[ZZ] = self.coords[current_beadid, ZZ] + (zi - 1) * self.grid.cellsize[ZZ] - probe[XX] = self.coords[current_beadid, XX] + (xi - 1) * self.cutoff probe[YY] = self.coords[current_beadid, YY] + (yi - 1) * self.cutoff probe[ZZ] = self.coords[current_beadid, ZZ] + (zi - 1) * self.cutoff - # Make sure the shifted coordinates is inside the brick-shaped box for m in range(DIM - 1, -1, -1): - while shifted_coord[m] < 0: for d in range(m+1): shifted_coord[d] += self.box.c_pbcbox.box[m][d] - - while shifted_coord[m] >= self.box.c_pbcbox.box[m][m]: for d in range(m+1): shifted_coord[d] -= self.box.c_pbcbox.box[m][d] - while probe[m] < 0: for d in range(m+1): probe[d] += self.box.c_pbcbox.box[m][d] - - while probe[m] >= self.box.c_pbcbox.box[m][m]: for d in range(m+1): probe[d] -= self.box.c_pbcbox.box[m][d] @@ -772,58 +679,12 @@ cdef class FastNS(object): cellindex_adjacent = self.grid.coord2cellid(shifted_coord) cellindex_probe = self.grid.coord2cellid(probe) - if cellindex == cellindex_probe and xi != 1 and yi != 1 and zi != 1: - # if self.debug and debug: - # with gil: - # print("FastNS: Grid shift [{}][{}][{}]: Cutoff is inside current cell -> This shift is ignored".format( - # xi - 1, - # yi -1, - # zi -1 - # )) - continue - - # if self.debug and debug: - # self.grid.cellid2cellxyz(cellindex_adjacent, debug_cellxyz) - # with gil: - # dist_shift = self.box.fast_distance(&self.coords[current_beadid, XX], shifted_coord) - # grid_shift = np.array([(xi - 1) * self.grid.cellsize[XX], - # (yi - 1) * self.grid.cellsize[YY], - # (zi - 1) * self.grid.cellsize[ZZ]]) - # print("FastNS: -> Checking cell#{} ({},{},{}) for neighbors (dshift={:.3f}, grid_shift=({:.3f},{:.3f},{:.3f}->{:.3f})".format( - # cellindex, - # debug_cellxyz[XX], debug_cellxyz[YY], debug_cellxyz[ZZ], - # dist_shift, - # grid_shift[XX], grid_shift[YY], grid_shift[ZZ], - # np.sqrt(np.sum(grid_shift**2)) - # )) - - for j in range(self.grid.nbeads[cellindex_adjacent]): bid = self.grid.beadids[cellindex_adjacent * self.grid.nbeads_per_cell + j] - if checked[bid] != 0: continue - d2 = self.box.fast_distance2(&self.coords_bbox[current_beadid, XX], &self.coords_bbox[bid, XX]) - - # if self.debug: - # self.grid.cellid2cellxyz(cellindex, debug_cellxyz) - # with gil: - # print( - # "Beads #{} (cell[{},{},{}]-coords[{:.3f},{:.3f},{:.3f}]) and #{} (cell[{},{},{}]-coords[{:.3f},{:.3f},{:.3f}]) are tested (d2={:.3f})".format( - # current_beadid, - # cellxyz[XX], cellxyz[YY], cellxyz[ZZ], - # self.coords_bbox[current_beadid, XX], - # self.coords_bbox[current_beadid, YY], - # self.coords_bbox[current_beadid, ZZ], - # bid, - # debug_cellxyz[XX], debug_cellxyz[YY], debug_cellxyz[ZZ], - # self.coords_bbox[bid, XX], self.coords_bbox[bid, YY], - # self.coords_bbox[bid, ZZ], - # d2)) - if d2 < cutoff2: - if d2 < EPSILON: continue elif results.add_neighbors(current_beadid, bid, d2) != OK: @@ -831,30 +692,6 @@ cdef class FastNS(object): break npairs += 1 checked[current_beadid] = 1 - if memory_error: raise MemoryError("Could not allocate memory to store NS results") - - - if self.debug: - print("Total number of pairs={}".format(npairs)) - - # ref_bead = 13937 - # beads = np.array([4398, 4401, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, 47451]) - 1 - # for bid in beads: - # self.box.fast_pbc_dx(&self.coords[ref_bead, XX], &self.coords[bid, XX], dx) - # dx_py = np.array([dx[XX], dx[YY], dx[ZZ]]) - # self.box.fast_pbc_dx(&self.coords_bbox[ref_bead, XX], &self.coords_bbox[bid, XX], dx) - # rect_dx_py = np.array([dx[XX], dx[YY], dx[ZZ]]) - # self.grid.cellid2cellxyz(self.grid.coord2cellid(&self.coords_bbox[bid, XX]), cellxyz) - # print("Bead #{} ({:.3f},{:.3f},{:.3f})->rect({:.3f},{:.3f},{:.3f}) - cell[{},{},{}]: dx=[{:.3f},{:.3f},{:.3f}] -> dist: {:.3f} ({})".format( - # bid, - # self.coords[bid, XX], self.coords[bid, YY], self.coords[bid,ZZ], - # self.coords_bbox[bid, XX], self.coords_bbox[bid, YY], self.coords_bbox[bid,ZZ], - # cellxyz[XX], cellxyz[YY], cellxyz[ZZ], - # dx[XX], dx[YY], dx[ZZ], - # np.sqrt(np.sum(dx_py**2)), - # self.box.fast_distance(&self.coords[ref_bead, XX], &self.coords[bid, XX]) <= self.cutoff, - # )) - - return results \ No newline at end of file + return results From 89f0b897ce5eb3efa84f9ea0964c7c493da5e69f Mon Sep 17 00:00:00 2001 From: Max Linke Date: Sat, 21 Jul 2018 20:42:25 +0200 Subject: [PATCH 28/47] remove possible memory leak --- package/MDAnalysis/lib/nsgrid.pyx | 19 ++++++------------- 1 file changed, 6 insertions(+), 13 deletions(-) diff --git a/package/MDAnalysis/lib/nsgrid.pyx b/package/MDAnalysis/lib/nsgrid.pyx index df6d130ef14..6dff29a6bab 100644 --- a/package/MDAnalysis/lib/nsgrid.pyx +++ b/package/MDAnalysis/lib/nsgrid.pyx @@ -506,12 +506,7 @@ cdef class NSGrid(object): cdef fill_grid(self, real[:, ::1] coords): cdef ns_int i, cellindex = -1 cdef ns_int ncoords = coords.shape[0] - cdef ns_int *beadcounts = NULL - - # Allocate memory - beadcounts = PyMem_Malloc(sizeof(ns_int) * self.size) - if not beadcounts: - raise MemoryError("Could not allocate memory for bead count buffer ({} bits requested)".format(sizeof(ns_int) * self.size)) + cdef ns_int[:] beadcounts = np.empty(self.size, dtype=np.int) with nogil: # Initialize buffers @@ -528,12 +523,12 @@ cdef class NSGrid(object): if self.nbeads[cellindex] > self.nbeads_per_cell: self.nbeads_per_cell = self.nbeads[cellindex] - # Allocate memory - with gil: - self.beadids = PyMem_Malloc(sizeof(ns_int) * self.size * self.nbeads_per_cell) #np.empty((self.size, nbeads_max), dtype=np.int) - if not self.beadids: - raise MemoryError("Could not allocate memory for NSGrid.beadids ({} bits requested)".format(sizeof(ns_int) * self.size * self.nbeads_per_cell)) + # Allocate memory + self.beadids = PyMem_Malloc(sizeof(ns_int) * self.size * self.nbeads_per_cell) #np.empty((self.size, nbeads_max), dtype=np.int) + if not self.beadids: + raise MemoryError("Could not allocate memory for NSGrid.beadids ({} bits requested)".format(sizeof(ns_int) * self.size * self.nbeads_per_cell)) + with nogil: # Second loop: fill grid for i in range(ncoords): @@ -542,8 +537,6 @@ cdef class NSGrid(object): self.beadids[cellindex * self.nbeads_per_cell + beadcounts[cellindex]] = i beadcounts[cellindex] += 1 - # Now we can free the allocation buffer - PyMem_Free(beadcounts) cdef class FastNS(object): From 08de29ba921f2c71b456a8b6f4eee9b6ec482c74 Mon Sep 17 00:00:00 2001 From: ayush Date: Thu, 26 Jul 2018 20:31:33 -0700 Subject: [PATCH 29/47] added the capped function, modified the nsgrid to use box, coordinates as input parameters, modified tests --- package/MDAnalysis/lib/distances.py | 58 ++++++++++++++++++- package/MDAnalysis/lib/nsgrid.pyx | 12 ++-- .../MDAnalysisTests/lib/test_distances.py | 7 ++- testsuite/MDAnalysisTests/lib/test_nsgrid.py | 4 +- 4 files changed, 70 insertions(+), 11 deletions(-) diff --git a/package/MDAnalysis/lib/distances.py b/package/MDAnalysis/lib/distances.py index 1cb86147235..4d662068a69 100644 --- a/package/MDAnalysis/lib/distances.py +++ b/package/MDAnalysis/lib/distances.py @@ -520,7 +520,8 @@ def _determine_method(reference, configuration, max_cutoff, min_cutoff=None, """ methods = {'bruteforce': _bruteforce_capped, - 'pkdtree': _pkdtree_capped} + 'pkdtree': _pkdtree_capped, + 'nsgrid': _nsgrid_capped} if method is not None: return methods[method] @@ -643,6 +644,61 @@ def _pkdtree_capped(reference, configuration, max_cutoff, distances.append(dist[num]) return pairs, distances +def _nsgrid_capped(reference, configuration, max_cutoff, min_cutoff=None, + box=None): + """Search all the pairs in *reference* and *configuration* within + a specified distance using Grid Search + + + Parameters + ----------- + reference : array + reference coordinates array with shape ``reference.shape = (3,)`` + or ``reference.shape = (len(reference), 3)``. + configuration : array + Configuration coordinate array with shape ``reference.shape = (3,)`` + or ``reference.shape = (len(reference), 3)`` + max_cutoff : float + Maximum cutoff distance between the reference and configuration + min_cutoff : (optional) float + Minimum cutoff distance between reference and configuration [None] + box : array + The dimensions, if provided, must be provided in the same + The unitcell dimesions for this system format as returned + by :attr:`MDAnalysis.coordinates.base.Timestep.dimensions`: + ``[lx,ly, lz, alpha, beta, gamma]``. Minimum image convention + is applied if the box is provided + + Note + ---- + Non Periodic Boundary conditions can-not be handled currently + + """ + from .nsgrid import FastNS + if reference.shape == (3, ): + reference = reference[None, :] + if configuration.shape == (3, ): + configuration = configuration[None, :] + all_coords = np.concatenate([configuration, reference]) + mapping = np.arange(len(reference), dtype=np.int64) + gridsearch = FastNS(box, max_cutoff, all_coords) + gridsearch.prepare() + search_ids = np.arange(len(configuration), len(all_coords)) + results = gridsearch.search(search_ids=search_ids) + pairs = results.get_pairs() + pairs[:, 1] = undo_augment(pairs[:, 1], mapping, len(configuration)) + pair_distance = results.get_pair_distances() + + if min_cutoff is not None: + idx = pair_distance > min_cutoff + pairs, pair_distance = pairs[idx], pair_distance[idx] + if pairs.size > 0: + # removing the pairs (i, j) from (reference, reference) + mask = (pairs[:, 0] < (len(configuration) - 1) ) + pairs, pair_distance = pairs[mask], pair_distance[mask] + pairs = np.sort(pairs, axis=1) + return pairs, pair_distance + def self_capped_distance(reference, max_cutoff, min_cutoff=None, box=None, method=None): diff --git a/package/MDAnalysis/lib/nsgrid.pyx b/package/MDAnalysis/lib/nsgrid.pyx index 6dff29a6bab..f38e8e143c4 100644 --- a/package/MDAnalysis/lib/nsgrid.pyx +++ b/package/MDAnalysis/lib/nsgrid.pyx @@ -548,21 +548,20 @@ cdef class FastNS(object): cdef bint prepared cdef NSGrid grid - def __init__(self, u, cutoff, coords=None, prepare=True, debug=False, max_gridsize=5000): + def __init__(self, box, cutoff, coords, prepare=True, debug=False, max_gridsize=5000): import MDAnalysis as mda from MDAnalysis.lib.mdamath import triclinic_vectors self.debug = debug - if not isinstance(u, mda.Universe): - raise TypeError("FastNS class must be initialized with a valid MDAnalysis.Universe instance") - box = triclinic_vectors(u.dimensions) + if box.shape != (3,3): + box = triclinic_vectors(box) self.box = PBCBox(box) - if coords is None: - coords = u.atoms.positions + #if coords is None: + # coords = u.atoms.positions self.coords = coords.copy() @@ -579,7 +578,6 @@ cdef class FastNS(object): if prepare: self.prepare() - def prepare(self, force=False): if self.prepared and not force: return diff --git a/testsuite/MDAnalysisTests/lib/test_distances.py b/testsuite/MDAnalysisTests/lib/test_distances.py index 38f634844e3..d83d8903e93 100644 --- a/testsuite/MDAnalysisTests/lib/test_distances.py +++ b/testsuite/MDAnalysisTests/lib/test_distances.py @@ -66,7 +66,7 @@ def test_capped_distance_noresults(): np.array([[0.1, 0.1, 0.1], [0.2, 0.1, 0.1]], dtype=np.float32)) -method_1 = ('bruteforce', 'pkdtree') +method_1 = ('bruteforce', 'pkdtree', 'nsgrid') min_cutoff_1 = (None, 0.1) @@ -77,6 +77,10 @@ def test_capped_distance_noresults(): @pytest.mark.parametrize('method', method_1) @pytest.mark.parametrize('min_cutoff', min_cutoff_1) def test_capped_distance_checkbrute(npoints, box, query, method, min_cutoff): + if method == 'nsgrid' and box is None: + pytest.skip('Not implemented yet') + + np.random.seed(90003) points = (np.random.uniform(low=0, high=1.0, size=(npoints, 3))*(boxes_1[0][:3])).astype(np.float32) @@ -89,6 +93,7 @@ def test_capped_distance_checkbrute(npoints, box, query, method, min_cutoff): min_cutoff=min_cutoff, box=box, method=method) + if pairs.shape != (0, ): found_pairs = pairs[:, 1] else: diff --git a/testsuite/MDAnalysisTests/lib/test_nsgrid.py b/testsuite/MDAnalysisTests/lib/test_nsgrid.py index 0fc6e1ec880..3e8ca02d3fe 100644 --- a/testsuite/MDAnalysisTests/lib/test_nsgrid.py +++ b/testsuite/MDAnalysisTests/lib/test_nsgrid.py @@ -43,7 +43,7 @@ def run_grid_search(u, ids, cutoff=3): coords = u.atoms.positions # Run grid search - searcher = nsgrid.FastNS(u, cutoff, coords, debug=True) + searcher = nsgrid.FastNS(u.dimensions, cutoff, coords, debug=True) return searcher.search(ids) @@ -156,7 +156,7 @@ def test_nsgrid_PBC_rect(): cutoff = 7 # FastNS is called differently to max coverage - searcher = nsgrid.FastNS(universe, cutoff, prepare=False) + searcher = nsgrid.FastNS(universe.dimensions, cutoff, universe.atoms.positions, prepare=False) results_grid = searcher.search([ref_id,]).get_indices()[0] # pass the id as a list for test+coverage purpose From 9d52ac91889628616e63a4709e8803abcd1e272a Mon Sep 17 00:00:00 2001 From: ayush Date: Thu, 26 Jul 2018 23:10:41 -0700 Subject: [PATCH 30/47] recovering accidently removed augment in setupd during rebase --- package/setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/package/setup.py b/package/setup.py index 52d29fecc7c..9df60bb149e 100755 --- a/package/setup.py +++ b/package/setup.py @@ -389,7 +389,7 @@ def extensions(config): extra_link_args=parallel_args) pre_exts = [libdcd, distances, distances_omp, qcprot, transformation, libmdaxdr, util, encore_utils, - ap_clustering, spe_dimred, cutil, nsgrid] + ap_clustering, spe_dimred, cutil, augment, nsgrid] cython_generated = [] From 8b8ea74933f9e2405ed8fb9a91d67012e596f032 Mon Sep 17 00:00:00 2001 From: ayush Date: Sat, 28 Jul 2018 12:16:50 -0700 Subject: [PATCH 31/47] corrected name in documentation --- package/doc/sphinx/source/documentation_pages/lib/grid.rst | 2 -- package/doc/sphinx/source/documentation_pages/lib/nsgrid.rst | 2 ++ 2 files changed, 2 insertions(+), 2 deletions(-) delete mode 100644 package/doc/sphinx/source/documentation_pages/lib/grid.rst create mode 100644 package/doc/sphinx/source/documentation_pages/lib/nsgrid.rst diff --git a/package/doc/sphinx/source/documentation_pages/lib/grid.rst b/package/doc/sphinx/source/documentation_pages/lib/grid.rst deleted file mode 100644 index a6143e9ca38..00000000000 --- a/package/doc/sphinx/source/documentation_pages/lib/grid.rst +++ /dev/null @@ -1,2 +0,0 @@ -.. automodule:: MDAnalysis.lib.grid - :members: \ No newline at end of file diff --git a/package/doc/sphinx/source/documentation_pages/lib/nsgrid.rst b/package/doc/sphinx/source/documentation_pages/lib/nsgrid.rst new file mode 100644 index 00000000000..f36f2480b0b --- /dev/null +++ b/package/doc/sphinx/source/documentation_pages/lib/nsgrid.rst @@ -0,0 +1,2 @@ +.. automodule:: MDAnalysis.lib.nsgrid + :members: \ No newline at end of file From e71b671af47353915ea50e849b0b8c7fd08cee04 Mon Sep 17 00:00:00 2001 From: ayush Date: Sat, 28 Jul 2018 14:32:16 -0700 Subject: [PATCH 32/47] minor corrections to remove pylint error --- testsuite/MDAnalysisTests/lib/test_nsgrid.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/testsuite/MDAnalysisTests/lib/test_nsgrid.py b/testsuite/MDAnalysisTests/lib/test_nsgrid.py index 3e8ca02d3fe..db623fedb95 100644 --- a/testsuite/MDAnalysisTests/lib/test_nsgrid.py +++ b/testsuite/MDAnalysisTests/lib/test_nsgrid.py @@ -112,11 +112,13 @@ def test_nsgrid_badinit(): with pytest.raises(TypeError): nsgrid.FastNS(None, 1) + def test_nsgrid_badcutoff(universe): with pytest.raises(ValueError): run_grid_search(universe, 0, -4) run_grid_search(universe, 0, 100000) + def test_ns_grid_noneighbor(universe): """Check that grid search returns empty lists/arrays when there is no neighbors""" ref_id = 0 @@ -158,10 +160,10 @@ def test_nsgrid_PBC_rect(): # FastNS is called differently to max coverage searcher = nsgrid.FastNS(universe.dimensions, cutoff, universe.atoms.positions, prepare=False) - results_grid = searcher.search([ref_id,]).get_indices()[0] # pass the id as a list for test+coverage purpose + results_grid = searcher.search([ref_id, ]).get_indices()[0] # pass the id as a list for test+coverage purpose searcher.prepare() # Does nothing, called here for coverage - results_grid2 = searcher.search().get_indices() # call without specifying any ids, should do NS for all beads + results_grid2 = searcher.search().get_indices() # call without specifying any ids, should do NS for all beads assert_equal(results, results_grid) assert_equal(len(universe.atoms), len(results_grid2)) @@ -181,7 +183,6 @@ def test_nsgrid_PBC(universe): assert_equal(results, results_grid) - def test_nsgrid_pairs(universe): """Check that grid search returns the proper pairs""" @@ -258,4 +259,3 @@ def test_nsgrid_coordinates(universe): results_grid = run_grid_search(universe, ref_id).get_coordinates()[0] assert_allclose(np.sort(results, axis=0), np.sort(results_grid, axis=0), atol=1e-5) - From a5d1355f90793a376642feed8b8611769d081c16 Mon Sep 17 00:00:00 2001 From: ayush Date: Tue, 31 Jul 2018 02:45:55 -0700 Subject: [PATCH 33/47] Used vectors in place of lists and dict, cleaned code and modified function to take coordinates as input, modified tests accordingly --- package/MDAnalysis/lib/nsgrid.pyx | 510 ++++++++----------- package/setup.py | 3 +- testsuite/MDAnalysisTests/lib/test_nsgrid.py | 88 ++-- 3 files changed, 256 insertions(+), 345 deletions(-) diff --git a/package/MDAnalysis/lib/nsgrid.pyx b/package/MDAnalysis/lib/nsgrid.pyx index f38e8e143c4..ffb240e2f0a 100644 --- a/package/MDAnalysis/lib/nsgrid.pyx +++ b/package/MDAnalysis/lib/nsgrid.pyx @@ -43,25 +43,26 @@ DEF ZZ = 2 DEF EPSILON = 1e-5 -DEF BOX_MARGIN=1.0010 -DEF MAX_NTRICVEC=12 - -DEF OK=0 -DEF ERROR=1 # Used to handle memory allocation from cpython.mem cimport PyMem_Malloc, PyMem_Realloc, PyMem_Free from libc.math cimport sqrt import numpy as np cimport numpy as np +from libcpp.vector cimport vector +from libcpp.map cimport map as cmap + ctypedef np.int_t ns_int ctypedef np.float32_t real ctypedef real rvec[DIM] ctypedef ns_int ivec[DIM] -ctypedef ns_int ipair[2] ctypedef real matrix[DIM][DIM] +ctypedef vector[ns_int] intvec +ctypedef cmap[ns_int, intvec] intmap +ctypedef vector[real] realvec +ctypedef cmap[ns_int, realvec] realmap # Useful Functions cdef real rvec_norm2(const rvec a) nogil: @@ -86,34 +87,25 @@ cdef struct cPBCBox_t: # Class to handle PBC calculations cdef class PBCBox(object): cdef cPBCBox_t c_pbcbox - cdef rvec center - cdef rvec bbox_center cdef bint is_triclinic def __init__(self, real[:,::1] box): self.update(box) cdef void fast_update(self, real[:,::1] box) nogil: - cdef ns_int i, j, k, d, jc, kc, shift - cdef real d2old, d2new, d2new_c - cdef rvec trial, pos - cdef ns_int ii, jj ,kk - cdef ns_int *order = [0, -1, 1, -2, 2] - cdef bint use + + cdef ns_int i, j cdef real min_hv2, min_ss, tmp - rvec_clear(self.center) # Update matrix self.is_triclinic = False for i in range(DIM): for j in range(DIM): self.c_pbcbox.box[i][j] = box[i, j] - self.center[j] += 0.5 * box[i, j] if i != j: if box[i, j] > EPSILON: self.is_triclinic = True - self.bbox_center[i] = 0.5 * box[i, i] # Update diagonals for i in range(DIM): @@ -151,8 +143,8 @@ cdef class PBCBox(object): cdef void fast_pbc_dx(self, rvec ref, rvec other, rvec dx) nogil: + cdef ns_int i, j - cdef rvec dx_start, trial for i in range(DIM): dx[i] = other[i] - ref[i] @@ -170,9 +162,7 @@ cdef class PBCBox(object): cdef rvec dx if a.shape[0] != DIM or b.shape[0] != DIM: raise ValueError("Not 3 D coordinates") - self.fast_pbc_dx(&a[XX], &b[XX], dx) - return np.array([dx[XX], dx[YY], dx[ZZ]], dtype=np.float32) @@ -193,9 +183,9 @@ cdef class PBCBox(object): if a.shape[0] != DIM or b.shape[0] != DIM: raise ValueError("Not 3 D coordinates") return self.fast_distance(&a[XX], &b[XX]) - + cdef real[:, ::1]fast_put_atoms_in_bbox(self, real[:,::1] coords) nogil: - cdef ns_int i, m, d, natoms, wd = 0 + cdef ns_int i, m, d, natoms cdef real[:,::1] bbox_coords natoms = coords.shape[0] @@ -223,220 +213,108 @@ cdef class PBCBox(object): def put_atoms_in_bbox(self, real[:,::1] coords): return np.asarray(self.fast_put_atoms_in_bbox(coords)) - ######################### # Neighbor Search Stuff # ######################### + cdef class NSResults(object): cdef readonly real cutoff cdef ns_int npairs - cdef bint debug - - cdef real[:, ::1] coords # shape: size, DIM - cdef ns_int[:] search_ids - - cdef ns_int allocation_size - cdef ipair *pairs # shape: pair_allocation - cdef real *pair_distances2 # shape: pair_allocation - cdef list indices_buffer - cdef list coordinates_buffer - cdef list distances_buffer - cdef np.ndarray pairs_buffer - cdef np.ndarray pair_distances_buffer - cdef np.ndarray pair_coordinates_buffer - - def __init__(self, real cutoff, real[:, ::1]coords, ns_int[:] search_ids, debug=False): - self.debug = debug + cdef real[:, ::1] coords # shape: size, DIM + cdef real[:, ::1] searchcoords + + cdef intmap indices_buffer + cdef realmap distances_buffer + cdef vector[ns_int] pairs_buffer + cdef vector[real] pair_distances_buffer + cdef vector[real] pair_distances2_buffer + + def __init__(self, real cutoff, real[:, ::1]coords, real[:, ::1]searchcoords): self.cutoff = cutoff self.coords = coords - self.search_ids = search_ids + self.searchcoords = searchcoords # Preallocate memory - self.allocation_size = search_ids.shape[0] + 1 - if not self.pairs and not self.pair_distances2: - self.pairs = PyMem_Malloc(sizeof(ipair) * self.allocation_size) - if not self.pairs: - MemoryError("Could not allocate memory for NSResults.pairs " - "({} bits requested)".format(sizeof(ipair) * self.allocation_size)) - self.pair_distances2 = PyMem_Malloc(sizeof(real) * self.allocation_size) - if not self.pair_distances2: - raise MemoryError("Could not allocate memory for NSResults.pair_distances2 " - "({} bits requested)".format(sizeof(real) * self.allocation_size)) - else: - if self.resize(self.allocation_size) != OK: - raise MemoryError("foo") - self.npairs = 0 - - # Buffer - self.indices_buffer = None - self.coordinates_buffer = None - self.distances_buffer = None - self.pairs_buffer = None - self.pair_coordinates_buffer = None - - def __dealloc__(self): - PyMem_Free(self.pairs) - PyMem_Free(self.pair_distances2) - - cdef int add_neighbors(self, ns_int beadid_i, ns_int beadid_j, real distance2) nogil: + + cdef void add_neighbors(self, ns_int beadid_i, ns_int beadid_j, real distance2) nogil: # Important: If this function returns ERROR, it means that memory allocation failed - # Reallocate memory if needed - if self.npairs >= self.allocation_size: - # We need to reallocate memory - if self.resize(self.allocation_size + (self.allocation_size * 0.5 + 1)) != OK: - return ERROR - - # Actually store pair and distance squared - if beadid_i < beadid_j: - self.pairs[self.npairs][0] = beadid_i - self.pairs[self.npairs][1] = beadid_j - else: - self.pairs[self.npairs][1] = beadid_i - self.pairs[self.npairs][0] = beadid_j - self.pair_distances2[self.npairs] = distance2 + self.pairs_buffer.push_back(beadid_i) + self.pairs_buffer.push_back(beadid_j) + self.pair_distances2_buffer.push_back(distance2) self.npairs += 1 - - return OK - - cdef int resize(self, ns_int new_size) nogil: - # Important: If this function returns 0, it means that memory allocation failed - - if new_size < self.npairs: - # Silently ignored the request - return OK - - if self.allocation_size >= new_size: - if self.debug: - with gil: - print("NSresults: Reallocation requested but not needed ({} requested but {} already allocated)".format(new_size, self.allocation_size)) - return OK - - self.allocation_size = new_size - - if self.debug: - with gil: - print("NSresults: Reallocated to {} pairs".format(self.allocation_size)) - - # Allocating memory - with gil: - self.pairs = PyMem_Realloc(self.pairs, sizeof(ipair) * self.allocation_size) - self.pair_distances2 = PyMem_Realloc(self.pair_distances2, sizeof(real) * self.allocation_size) - - if not self.pairs: - return ERROR - - if not self.pair_distances2: - return ERROR - - return OK - + def get_pairs(self): - cdef ns_int i - - if self.pairs_buffer is None: - self.pairs_buffer = np.empty((self.npairs, 2), dtype=np.int) - for i in range(self.npairs): - self.pairs_buffer[i, 0] = self.pairs[i][0] - self.pairs_buffer[i, 1] = self.pairs[i][1] - return self.pairs_buffer + return np.asarray(self.pairs_buffer).reshape(self.npairs, 2) def get_pair_distances(self): - cdef ns_int i - if self.pair_coordinates_buffer is None: - self.pair_coordinates_buffer = np.empty(self.npairs, dtype=np.float32) - for i in range(self.npairs): - self.pair_coordinates_buffer[i] = self.pair_distances2[i] - self.pair_coordinates_buffer = np.sqrt(self.pair_coordinates_buffer) - return self.pair_coordinates_buffer - - def get_pair_coordinates(self): - cdef ns_int i, j, bead_i, bead_j - if self.pair_coordinates_buffer is None: - self.pair_coordinates_buffer = np.empty((self.npairs, 2, DIM), dtype=np.float32) - for i in range(self.npairs): - bead_i = self.pairs[i][0] - bead_j = self.pairs[i][1] - - for j in range(DIM): - self.pair_coordinates_buffer[i, 0, j] = self.coords[bead_i, j] - self.pair_coordinates_buffer[i, 1, j] = self.coords[bead_j, j] - return self.pair_coordinates_buffer - + self.pair_distances_buffer = np.sqrt(self.pair_distances2_buffer) + return np.asarray(self.pair_distances_buffer) + cdef create_buffers(self): cdef ns_int i, beadid_i, beadid_j + cdef ns_int idx, nsearch cdef real dist2 - cdef real[:] coord_i, coord_j - from collections import defaultdict - - indices_buffer = defaultdict(list) - coords_buffer = defaultdict(list) - dists_buffer = defaultdict(list) - - for i in range(self.npairs): - beadid_i = self.pairs[i][0] - beadid_j = self.pairs[i][1] - - dist2 = self.pair_distances2[i] - coord_i = self.coords[beadid_i] - coord_j = self.coords[beadid_j] - - indices_buffer[beadid_i].append(beadid_j) - indices_buffer[beadid_j].append(beadid_i) - - coords_buffer[beadid_i].append(coord_j) - coords_buffer[beadid_j].append((coord_i)) - - dists_buffer[beadid_i].append(dist2) - dists_buffer[beadid_j].append(dist2) - - self.indices_buffer = [] - self.coordinates_buffer = [] - self.distances_buffer = [] - - for elm in self.search_ids: - sorted_indices = np.argsort(indices_buffer[elm]) - self.indices_buffer.append(np.array(indices_buffer[elm])[sorted_indices]) - self.coordinates_buffer.append(np.array(coords_buffer[elm])[sorted_indices]) - self.distances_buffer.append(np.sqrt(dists_buffer[elm])[sorted_indices]) - + + nsearch = len(self.searchcoords) + + cdef intmap indices_buffer + cdef realmap distances_buffer + + for i in range(0, 2*self.npairs, 2): + beadid_i = self.pairs_buffer[i] + beadid_j = self.pairs_buffer[i + 1] + + dist2 = self.pair_distances2_buffer[i//2] + + indices_buffer[beadid_i].push_back(beadid_j) + + dist = sqrt(dist2) + + distances_buffer[beadid_i].push_back(dist) + + self.indices_buffer.clear() + self.distances_buffer.clear() + + for idx in range(nsearch): + sorted_indices = np.argsort(indices_buffer[idx]) + self.indices_buffer.insert((idx, intvec())) + self.distances_buffer.insert((idx, realvec())) + for index in sorted_indices: + self.indices_buffer[idx].push_back(indices_buffer[idx][index]) + self.distances_buffer[idx].push_back(distances_buffer[idx][index]) + + def get_indices(self): - if self.indices_buffer is None: + if self.indices_buffer.empty(): self.create_buffers() return self.indices_buffer def get_distances(self): - if self.distances_buffer is None: + if self.distances_buffer.empty(): self.create_buffers() return self.distances_buffer - - def get_coordinates(self): - if self.coordinates_buffer is None: - self.create_buffers() - return self.coordinates_buffer - - - + cdef class NSGrid(object): - cdef bint debug - cdef readonly real cutoff - cdef ns_int size - cdef ns_int ncoords - cdef ns_int[DIM] ncells - cdef ns_int[DIM] cell_offsets - cdef real[DIM] cellsize - cdef ns_int nbeads_per_cell - cdef ns_int *nbeads # size - cdef ns_int *beadids # size * nbeads_per_cell - cdef ns_int *cellids # ncoords - - def __init__(self, ncoords, cutoff, PBCBox box, max_size, debug=False): - cdef ns_int i, x, y, z + cdef readonly real cutoff # cutoff + cdef ns_int size # total cells + cdef ns_int ncoords # number of coordinates + cdef ns_int[DIM] ncells # individual cells in every dimension + cdef ns_int[DIM] cell_offsets # Cell Multipliers + cdef real[DIM] cellsize # cell size in every dimension + cdef ns_int nbeads_per_cell # maximum beads + cdef ns_int *nbeads # size (Number of beads in every cell) + cdef ns_int *beadids # size * nbeads_per_cell (Beadids in every cell) + cdef ns_int *cellids # ncoords (Cell occupation id for every atom) + + def __init__(self, ncoords, cutoff, PBCBox box, max_size): + cdef ns_int i cdef ns_int ncellx, ncelly, ncellz, size, nbeadspercell + cdef ns_int xi, yi, zi cdef real bbox_vol - self.debug = debug + self.ncoords = ncoords @@ -449,18 +327,10 @@ cdef class NSGrid(object): self.cutoff *= 1.2 for i in range(DIM): - self.ncells[i] = (box.c_pbcbox.box[i][i] / self.cutoff) - if self.ncells[i] == 0: - self.ncells[i] = 1 + self.ncells[i] = (box.c_pbcbox.box[i][i] / self.cutoff + 1) self.cellsize[i] = box.c_pbcbox.box[i][i] / self.ncells[i] self.size = self.ncells[XX] * self.ncells[YY] * self.ncells[ZZ] - if self.debug: - print("NSGrid: Requested cutoff: {:.3f} (Ncells={}, Avg # of beads per cell={}), Optimized cutoff= {:.3f} (Ncells={}, Avg # of beads per cell={})".format( - cutoff, size, nbeadspercell, - self.cutoff, self.size, (ncoords / self.size) - )) - print("NSGrid: Size={}x{}x{}={}".format(self.ncells[XX], self.ncells[YY], self.ncells[ZZ], self.size)) self.cell_offsets[XX] = 0 self.cell_offsets[YY] = self.ncells[XX] @@ -478,15 +348,15 @@ cdef class NSGrid(object): for i in range(self.size): self.nbeads[i] = 0 - + def __dealloc__(self): PyMem_Free(self.nbeads) PyMem_Free(self.beadids) PyMem_Free(self.cellids) cdef ns_int coord2cellid(self, rvec coord) nogil: - return (coord[ZZ] / self.cellsize[ZZ]) * (self.ncells[XX] * self.ncells[YY]) +\ - (coord[YY] / self.cellsize[YY]) * self.ncells[XX] + \ + return (coord[ZZ] / self.cellsize[ZZ]) * (self.cell_offsets[ZZ]) +\ + (coord[YY] / self.cellsize[YY]) * self.cell_offsets[YY] + \ (coord[XX] / self.cellsize[XX]) cdef bint cellid2cellxyz(self, ns_int cellid, ivec cellxyz) nogil: @@ -540,125 +410,171 @@ cdef class NSGrid(object): cdef class FastNS(object): - cdef bint debug + cdef PBCBox box cdef real[:, ::1] coords cdef real[:, ::1] coords_bbox + cdef real[:, ::1] searchcoords + cdef real[:, ::1] searchcoords_bbox cdef readonly real cutoff cdef bint prepared cdef NSGrid grid + cdef NSGrid searchgrid + cdef ns_int max_gridsize + cdef intmap cell_neighbors + - def __init__(self, box, cutoff, coords, prepare=True, debug=False, max_gridsize=5000): + def __init__(self, box, cutoff, coords, prepare=True, max_gridsize=5000): import MDAnalysis as mda from MDAnalysis.lib.mdamath import triclinic_vectors - self.debug = debug if box.shape != (3,3): box = triclinic_vectors(box) self.box = PBCBox(box) - - #if coords is None: - # coords = u.atoms.positions - - self.coords = coords.copy() - - self.coords_bbox = self.box.fast_put_atoms_in_bbox(coords) - if cutoff < 0: raise ValueError("Cutoff must be positive!") if cutoff * cutoff > self.box.c_pbcbox.max_cutoff2: raise ValueError("Cutoff greater than maximum cutoff ({:.3f}) given the PBC") - self.cutoff = cutoff + + self.coords = coords.copy() - self.grid = NSGrid(self.coords_bbox.shape[0], cutoff, self.box, max_gridsize, debug=debug) - self.prepared = False - if prepare: - self.prepare() + self.coords_bbox = self.box.fast_put_atoms_in_bbox(coords) - def prepare(self, force=False): - if self.prepared and not force: - return + self.cutoff = cutoff + self.max_gridsize = max_gridsize + self.grid = NSGrid(self.coords_bbox.shape[0], self.cutoff, self.box, self.max_gridsize) + self.grid.fill_grid(self.coords_bbox) - self.prepared = True + - def search(self, search_ids=None): + def search(self, searchcoords): cdef ns_int i, j, size_search cdef ns_int d, m + cdef ns_int ncells # Total number of cells + cdef ns_int current_beadid, bid + cdef ns_int cellindex, cellindex_probe + cdef ns_int xi, yi, zi + cdef NSResults results - cdef ns_int size = self.coords_bbox.shape[0] + + cdef ivec cellxyz + + cdef real d2 + cdef rvec probe - cdef ns_int current_beadid, bid - cdef rvec current_coords + ncells = self.grid.size + cdef ns_int[:] checked = np.zeros(ncells, dtype=np.int) + + cdef real cutoff2 = self.cutoff * self.cutoff + cdef ns_int npairs = 0 + + # Generate another grid to search + self.searchcoords_bbox = self.box.fast_put_atoms_in_bbox(searchcoords) + self.searchgrid = NSGrid(self.searchcoords_bbox.shape[0], self.cutoff, self.box, self.max_gridsize) + #cdef ns_int size = self.searchcoords_bbox.shape[0] + + + + size_search = searchcoords.shape[0] + + results = NSResults(self.cutoff, self.coords, searchcoords) - cdef ns_int cellindex, cellindex_adjacent, cellindex_probe - cdef ivec cellxyz, debug_cellxyz + with nogil: + for i in range(size_search): + for m in range(ncells): + checked[m] = 0 + # Start with first search coordinate + current_beadid = i + # find the cellindex of the coordinate + cellindex = self.searchgrid.cellids[current_beadid] + # Find the actual position of cell + self.searchgrid.cellid2cellxyz(cellindex, cellxyz) + for xi in range(DIM): + for yi in range(DIM): + for zi in range(DIM): + #Probe the search coordinates in a brick shaped box + probe[XX] = self.searchcoords_bbox[current_beadid, XX] + (xi - 1) * self.cutoff + probe[YY] = self.searchcoords_bbox[current_beadid, YY] + (yi - 1) * self.cutoff + probe[ZZ] = self.searchcoords_bbox[current_beadid, ZZ] + (zi - 1) * self.cutoff + # Make sure the shifted coordinates is inside the brick-shaped box + for m in range(DIM - 1, -1, -1): + while probe[m] < 0: + for d in range(m+1): + probe[d] += self.box.c_pbcbox.box[m][d] + while probe[m] >= self.box.c_pbcbox.box[m][m]: + for d in range(m+1): + probe[d] -= self.box.c_pbcbox.box[m][d] + # Get the cell index corresponding to the probe + cellindex_probe = self.grid.coord2cellid(probe) + if checked[cellindex_probe] != 0: + continue + #for this cellindex search in grid + for j in range(self.grid.nbeads[cellindex_probe]): + bid = self.grid.beadids[cellindex_probe * self.grid.nbeads_per_cell + j] + #find distance between search coords[i] and coords[bid] + d2 = self.box.fast_distance2(&self.searchcoords_bbox[current_beadid, XX], &self.coords_bbox[bid, XX]) + if d2 < cutoff2: + if d2 < EPSILON: + continue + else: + results.add_neighbors(current_beadid, bid, d2) + npairs += 1 + checked[cellindex_probe] = 1 + return results - cdef real[:, ::1] search_coords - cdef ns_int[:] search_ids_view + def self_search(self): + cdef ns_int i, j, size_search + cdef ns_int d, m + cdef ns_int ncells # Total number of cells + cdef ns_int current_beadid, bid + cdef ns_int cellindex, cellindex_probe cdef ns_int xi, yi, zi + + cdef NSResults results + + + + cdef ivec cellxyz + cdef real d2 - cdef rvec shifted_coord, probe, dx + cdef rvec probe - cdef ns_int nchecked = 0 - cdef ns_int[:] checked = np.zeros(size, dtype=np.int) + ncells = self.grid.size + cdef ns_int[:] checked = np.zeros(ncells, dtype=np.int) cdef real cutoff2 = self.cutoff * self.cutoff cdef ns_int npairs = 0 - if not self.prepared: - self.prepare() - - if search_ids is None: - search_ids=np.arange(size) - elif type(search_ids) == np.int: - search_ids = np.array([search_ids,], dtype=np.int) - elif type(search_ids) != np.ndarray: - search_ids = np.array(search_ids, dtype=np.int) + size_search = self.coords.shape[0] - search_ids_view = search_ids - size_search = search_ids.shape[0] - - results = NSResults(self.cutoff, self.coords, search_ids, self.debug) - - cdef bint memory_error = False + results = NSResults(self.cutoff, self.coords, self.coords) with nogil: for i in range(size_search): - if memory_error: - break - current_beadid = search_ids_view[i] + for m in range(ncells): + checked[m] = 0 + # Start with first search coordinate + current_beadid = i + # find the cellindex of the coordinate cellindex = self.grid.cellids[current_beadid] + # Find the actual position of cell self.grid.cellid2cellxyz(cellindex, cellxyz) for xi in range(DIM): - if memory_error: - break - for yi in range(DIM): - if memory_error: - break + for yi in range(DIM): for zi in range(DIM): - if memory_error: - break # Calculate and/or reinitialize shifted coordinates - shifted_coord[XX] = self.coords[current_beadid, XX] + (xi - 1) * self.grid.cellsize[XX] - shifted_coord[YY] = self.coords[current_beadid, YY] + (yi - 1) * self.grid.cellsize[YY] - shifted_coord[ZZ] = self.coords[current_beadid, ZZ] + (zi - 1) * self.grid.cellsize[ZZ] - probe[XX] = self.coords[current_beadid, XX] + (xi - 1) * self.cutoff - probe[YY] = self.coords[current_beadid, YY] + (yi - 1) * self.cutoff - probe[ZZ] = self.coords[current_beadid, ZZ] + (zi - 1) * self.cutoff + #Probe the search coordinates in a brick shaped box + probe[XX] = self.coords_bbox[current_beadid, XX] + (xi - 1) * self.cutoff + probe[YY] = self.coords_bbox[current_beadid, YY] + (yi - 1) * self.cutoff + probe[ZZ] = self.coords_bbox[current_beadid, ZZ] + (zi - 1) * self.cutoff # Make sure the shifted coordinates is inside the brick-shaped box for m in range(DIM - 1, -1, -1): - while shifted_coord[m] < 0: - for d in range(m+1): - shifted_coord[d] += self.box.c_pbcbox.box[m][d] - while shifted_coord[m] >= self.box.c_pbcbox.box[m][m]: - for d in range(m+1): - shifted_coord[d] -= self.box.c_pbcbox.box[m][d] while probe[m] < 0: for d in range(m+1): probe[d] += self.box.c_pbcbox.box[m][d] @@ -666,23 +582,23 @@ cdef class FastNS(object): for d in range(m+1): probe[d] -= self.box.c_pbcbox.box[m][d] - # Get the cell index corresponding to the coord - cellindex_adjacent = self.grid.coord2cellid(shifted_coord) + # Get the cell index corresponding to the probe cellindex_probe = self.grid.coord2cellid(probe) - - for j in range(self.grid.nbeads[cellindex_adjacent]): - bid = self.grid.beadids[cellindex_adjacent * self.grid.nbeads_per_cell + j] - if checked[bid] != 0: + if checked[cellindex_probe] != 0: + continue + #for this cellindex search in grid + for j in range(self.grid.nbeads[cellindex_probe]): + bid = self.grid.beadids[cellindex_probe * self.grid.nbeads_per_cell + j] + if bid > current_beadid: continue + #find distance between search coords[i] and coords[bid] d2 = self.box.fast_distance2(&self.coords_bbox[current_beadid, XX], &self.coords_bbox[bid, XX]) if d2 < cutoff2: if d2 < EPSILON: continue - elif results.add_neighbors(current_beadid, bid, d2) != OK: - memory_error = True - break + else: + results.add_neighbors(current_beadid, bid, d2) + results.add_neighbors(bid, current_beadid, d2) npairs += 1 - checked[current_beadid] = 1 - if memory_error: - raise MemoryError("Could not allocate memory to store NS results") + checked[cellindex_probe] = 1 return results diff --git a/package/setup.py b/package/setup.py index 9df60bb149e..fe2e8bb7d1c 100755 --- a/package/setup.py +++ b/package/setup.py @@ -383,7 +383,8 @@ def extensions(config): nsgrid = MDAExtension('MDAnalysis.lib.nsgrid', ['MDAnalysis/lib/nsgrid' + source_suffix], include_dirs=include_dirs, - libraries=mathlib + parallel_libraries, + language='c++', + libraries=parallel_libraries, define_macros=define_macros + parallel_macros, extra_compile_args=extra_compile_args + parallel_args, extra_link_args=parallel_args) diff --git a/testsuite/MDAnalysisTests/lib/test_nsgrid.py b/testsuite/MDAnalysisTests/lib/test_nsgrid.py index db623fedb95..2808d1247aa 100644 --- a/testsuite/MDAnalysisTests/lib/test_nsgrid.py +++ b/testsuite/MDAnalysisTests/lib/test_nsgrid.py @@ -39,13 +39,14 @@ def universe(): -def run_grid_search(u, ids, cutoff=3): +def run_grid_search(u, searchcoords, cutoff=3): coords = u.atoms.positions - + if searchcoords.shape == (3, ): + searchcoords = searchcoords[None, :] # Run grid search - searcher = nsgrid.FastNS(u.dimensions, cutoff, coords, debug=True) + searcher = nsgrid.FastNS(u.dimensions, cutoff, coords) - return searcher.search(ids) + return searcher.search(searchcoords) def test_pbc_badbox(): @@ -115,8 +116,8 @@ def test_nsgrid_badinit(): def test_nsgrid_badcutoff(universe): with pytest.raises(ValueError): - run_grid_search(universe, 0, -4) - run_grid_search(universe, 0, 100000) + run_grid_search(universe, universe.atoms.positions[0], -4) + run_grid_search(universe, universe.atoms.positions[0], 100000) def test_ns_grid_noneighbor(universe): @@ -124,14 +125,12 @@ def test_ns_grid_noneighbor(universe): ref_id = 0 cutoff = 0.5 - results_grid = run_grid_search(universe, ref_id, cutoff) + results_grid = run_grid_search(universe, universe.atoms.positions[ref_id], cutoff) - assert len(results_grid.get_coordinates()[0]) == 0 assert len(results_grid.get_distances()[0]) == 0 assert len(results_grid.get_indices()[0]) == 0 assert len(results_grid.get_pairs()) == 0 assert len(results_grid.get_pair_distances()) == 0 - assert len(results_grid.get_pair_coordinates()) == 0 def test_nsgrid_noPBC(universe): @@ -143,7 +142,7 @@ def test_nsgrid_noPBC(universe): results = np.array([2, 3, 4, 5, 6, 7, 8, 9, 18, 19, 1211, 10862, 10865, 17582, 17585, 38342, 38345]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! - results_grid = run_grid_search(universe, ref_id, cutoff).get_indices()[0] + results_grid = run_grid_search(universe, universe.atoms.positions[ref_id], cutoff).get_indices()[0] assert_equal(results, results_grid) @@ -158,12 +157,11 @@ def test_nsgrid_PBC_rect(): cutoff = 7 # FastNS is called differently to max coverage - searcher = nsgrid.FastNS(universe.dimensions, cutoff, universe.atoms.positions, prepare=False) + searcher = nsgrid.FastNS(universe.dimensions, cutoff, universe.atoms.positions) - results_grid = searcher.search([ref_id, ]).get_indices()[0] # pass the id as a list for test+coverage purpose + results_grid = searcher.search(universe.atoms.positions[ref_id][None, :]).get_indices()[0] # pass the id as a list for test+coverage purpose - searcher.prepare() # Does nothing, called here for coverage - results_grid2 = searcher.search().get_indices() # call without specifying any ids, should do NS for all beads + results_grid2 = searcher.search(universe.atoms.positions).get_indices() # call without specifying any ids, should do NS for all beads assert_equal(results, results_grid) assert_equal(len(universe.atoms), len(results_grid2)) @@ -178,7 +176,7 @@ def test_nsgrid_PBC(universe): results = np.array([4398, 4401, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, 47451]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! - results_grid = run_grid_search(universe, ref_id).get_indices()[0] + results_grid = run_grid_search(universe, universe.atoms.positions[ref_id]).get_indices()[0] assert_equal(results, results_grid) @@ -190,16 +188,12 @@ def test_nsgrid_pairs(universe): neighbors = np.array([4398, 4401, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, 47451]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! results = [] - for nid in neighbors: - if nid < ref_id: - results.append([nid, ref_id]) - else: - results.append([ref_id, nid]) + results = np.array(results) - results_grid = run_grid_search(universe, ref_id).get_pairs() + results_grid = run_grid_search(universe, universe.atoms.positions[ref_id]).get_pairs() - assert_equal(np.sort(results, axis=0), np.sort(results_grid, axis=0)) + assert_equal(np.sort(neighbors, axis=0), np.sort(results_grid[:, 1], axis=0)) def test_nsgrid_pair_distances(universe): @@ -209,30 +203,30 @@ def test_nsgrid_pair_distances(universe): results = np.array([0.270, 0.285, 0.096, 0.096, 0.015, 0.278, 0.268, 0.179, 0.259, 0.290, 0.270]) * 10 # These distances where obtained by gmx distance so they are in nm - results_grid = run_grid_search(universe, ref_id).get_pair_distances() + results_grid = run_grid_search(universe, universe.atoms.positions[ref_id]).get_pair_distances() assert_allclose(np.sort(results), np.sort(results_grid), atol=1e-2) -def test_nsgrid_pair_coordinates(universe): - """Check that grid search return the proper pair coordinates""" +# def test_nsgrid_pair_coordinates(universe): +# """Check that grid search return the proper pair coordinates""" - ref_id = 13937 - neighbors = np.array([4398, 4401, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, - 47451]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! - coords = universe.atoms.positions +# ref_id = 13937 +# neighbors = np.array([4398, 4401, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, +# 47451]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! +# coords = universe.atoms.positions - results = [] - for nid in neighbors: - if nid < ref_id: - results.append([coords[nid], coords[ref_id]]) - else: - results.append([coords[ref_id], coords[nid]]) - results = np.array(results) +# results = [] +# for nid in neighbors: +# if nid < ref_id: +# results.append([coords[nid], coords[ref_id]]) +# else: +# results.append([coords[ref_id], coords[nid]]) +# results = np.array(results) - results_grid = run_grid_search(universe, ref_id).get_pair_coordinates() +# results_grid = run_grid_search(universe, ref_id).get_pair_coordinates() - assert_allclose(np.sort(results, axis=0), np.sort(results_grid, axis=0), atol=1e-5) +# assert_allclose(np.sort(results, axis=0), np.sort(results_grid, axis=0), atol=1e-5) def test_nsgrid_distances(universe): @@ -242,20 +236,20 @@ def test_nsgrid_distances(universe): results = np.array([0.270, 0.285, 0.096, 0.096, 0.015, 0.278, 0.268, 0.179, 0.259, 0.290, 0.270]) * 10 # These distances where obtained by gmx distance so they are in nm - results_grid = run_grid_search(universe, ref_id).get_distances()[0] + results_grid = run_grid_search(universe, universe.atoms.positions[ref_id]).get_distances()[0] assert_allclose(np.sort(results), np.sort(results_grid), atol=1e-2) -def test_nsgrid_coordinates(universe): - """Check that grid search return the proper coordinates""" +# def test_nsgrid_coordinates(universe): +# """Check that grid search return the proper coordinates""" - ref_id = 13937 - neighbors = np.array([4398, 4401, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, - 47451]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! +# ref_id = 13937 +# neighbors = np.array([4398, 4401, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, +# 47451]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! - results = universe.atoms.positions[neighbors] +# results = universe.atoms.positions[neighbors] - results_grid = run_grid_search(universe, ref_id).get_coordinates()[0] +# results_grid = run_grid_search(universe, ref_id).get_coordinates()[0] - assert_allclose(np.sort(results, axis=0), np.sort(results_grid, axis=0), atol=1e-5) +# assert_allclose(np.sort(results, axis=0), np.sort(results_grid, axis=0), atol=1e-5) From 46f9030e95d261e89d8645c9e8262ce6934d7b5e Mon Sep 17 00:00:00 2001 From: ayush Date: Tue, 31 Jul 2018 15:57:42 -0700 Subject: [PATCH 34/47] removed dict from buffers for speedup and modified tests --- package/MDAnalysis/lib/nsgrid.pyx | 131 ++++++++----------- testsuite/MDAnalysisTests/lib/test_nsgrid.py | 54 ++------ 2 files changed, 64 insertions(+), 121 deletions(-) diff --git a/package/MDAnalysis/lib/nsgrid.pyx b/package/MDAnalysis/lib/nsgrid.pyx index ffb240e2f0a..8598bec1e97 100644 --- a/package/MDAnalysis/lib/nsgrid.pyx +++ b/package/MDAnalysis/lib/nsgrid.pyx @@ -50,7 +50,6 @@ from libc.math cimport sqrt import numpy as np cimport numpy as np from libcpp.vector cimport vector -from libcpp.map cimport map as cmap ctypedef np.int_t ns_int @@ -60,9 +59,7 @@ ctypedef ns_int ivec[DIM] ctypedef real matrix[DIM][DIM] ctypedef vector[ns_int] intvec -ctypedef cmap[ns_int, intvec] intmap ctypedef vector[real] realvec -ctypedef cmap[ns_int, realvec] realmap # Useful Functions cdef real rvec_norm2(const rvec a) nogil: @@ -224,8 +221,8 @@ cdef class NSResults(object): cdef real[:, ::1] coords # shape: size, DIM cdef real[:, ::1] searchcoords - cdef intmap indices_buffer - cdef realmap distances_buffer + cdef vector[intvec] indices_buffer + cdef vector[realvec] distances_buffer cdef vector[ns_int] pairs_buffer cdef vector[real] pair_distances_buffer cdef vector[real] pair_distances2_buffer @@ -235,7 +232,6 @@ cdef class NSResults(object): self.coords = coords self.searchcoords = searchcoords - # Preallocate memory self.npairs = 0 cdef void add_neighbors(self, ns_int beadid_i, ns_int beadid_j, real distance2) nogil: @@ -260,8 +256,14 @@ cdef class NSResults(object): nsearch = len(self.searchcoords) - cdef intmap indices_buffer - cdef realmap distances_buffer + self.indices_buffer = vector[intvec]() + self.distances_buffer = vector[realvec]() + + # initialize rows corresponding to search + for i in range(nsearch): + self.indices_buffer.push_back(intvec()) + self.distances_buffer.push_back(realvec()) + for i in range(0, 2*self.npairs, 2): beadid_i = self.pairs_buffer[i] @@ -269,33 +271,27 @@ cdef class NSResults(object): dist2 = self.pair_distances2_buffer[i//2] - indices_buffer[beadid_i].push_back(beadid_j) + self.indices_buffer[beadid_i].push_back(beadid_j) dist = sqrt(dist2) - distances_buffer[beadid_i].push_back(dist) + self.distances_buffer[beadid_i].push_back(dist) - self.indices_buffer.clear() - self.distances_buffer.clear() + for i in range(self.searchcoords.shape[0]): + sorted_indices = np.argsort(self.indices_buffer[i]) + self.indices_buffer[i] = np.array(self.indices_buffer[i])[sorted_indices] + self.distances_buffer[i] = np.array(self.distances_buffer[i])[sorted_indices] - for idx in range(nsearch): - sorted_indices = np.argsort(indices_buffer[idx]) - self.indices_buffer.insert((idx, intvec())) - self.distances_buffer.insert((idx, realvec())) - for index in sorted_indices: - self.indices_buffer[idx].push_back(indices_buffer[idx][index]) - self.distances_buffer[idx].push_back(distances_buffer[idx][index]) - def get_indices(self): if self.indices_buffer.empty(): self.create_buffers() - return self.indices_buffer + return np.asarray(self.indices_buffer) def get_distances(self): if self.distances_buffer.empty(): self.create_buffers() - return self.distances_buffer + return np.asarray(self.distances_buffer) cdef class NSGrid(object): cdef readonly real cutoff # cutoff @@ -308,8 +304,9 @@ cdef class NSGrid(object): cdef ns_int *nbeads # size (Number of beads in every cell) cdef ns_int *beadids # size * nbeads_per_cell (Beadids in every cell) cdef ns_int *cellids # ncoords (Cell occupation id for every atom) + cdef bint force # To negate the effects of optimized cutoff - def __init__(self, ncoords, cutoff, PBCBox box, max_size): + def __init__(self, ncoords, cutoff, PBCBox box, max_size, force=False): cdef ns_int i cdef ns_int ncellx, ncelly, ncellz, size, nbeadspercell cdef ns_int xi, yi, zi @@ -320,14 +317,15 @@ cdef class NSGrid(object): # Calculate best cutoff self.cutoff = cutoff - bbox_vol = box.c_pbcbox.box[XX][XX] * box.c_pbcbox.box[YY][YY] * box.c_pbcbox.box[YY][YY] - size = bbox_vol/cutoff**3 - nbeadspercell = ncoords/size - while bbox_vol/self.cutoff**3 > max_size: - self.cutoff *= 1.2 + if not force: + bbox_vol = box.c_pbcbox.box[XX][XX] * box.c_pbcbox.box[YY][YY] * box.c_pbcbox.box[YY][YY] + size = bbox_vol/cutoff**3 + nbeadspercell = ncoords/size + while bbox_vol/self.cutoff**3 > max_size: + self.cutoff *= 1.2 for i in range(DIM): - self.ncells[i] = (box.c_pbcbox.box[i][i] / self.cutoff + 1) + self.ncells[i] = (box.c_pbcbox.box[i][i] / self.cutoff) self.cellsize[i] = box.c_pbcbox.box[i][i] / self.ncells[i] self.size = self.ncells[XX] * self.ncells[YY] * self.ncells[ZZ] @@ -414,17 +412,12 @@ cdef class FastNS(object): cdef PBCBox box cdef real[:, ::1] coords cdef real[:, ::1] coords_bbox - cdef real[:, ::1] searchcoords - cdef real[:, ::1] searchcoords_bbox cdef readonly real cutoff - cdef bint prepared cdef NSGrid grid - cdef NSGrid searchgrid cdef ns_int max_gridsize - cdef intmap cell_neighbors - def __init__(self, box, cutoff, coords, prepare=True, max_gridsize=5000): + def __init__(self, box, cutoff, coords, max_gridsize=5000): import MDAnalysis as mda from MDAnalysis.lib.mdamath import triclinic_vectors @@ -445,39 +438,39 @@ cdef class FastNS(object): self.cutoff = cutoff self.max_gridsize = max_gridsize - + # Note that self.cutoff might be different from self.grid.cutoff self.grid = NSGrid(self.coords_bbox.shape[0], self.cutoff, self.box, self.max_gridsize) self.grid.fill_grid(self.coords_bbox) - def search(self, searchcoords): + def search(self, search_coords): cdef ns_int i, j, size_search cdef ns_int d, m - cdef ns_int ncells # Total number of cells cdef ns_int current_beadid, bid cdef ns_int cellindex, cellindex_probe cdef ns_int xi, yi, zi cdef NSResults results - cdef ivec cellxyz - cdef real d2 cdef rvec probe - ncells = self.grid.size - cdef ns_int[:] checked = np.zeros(ncells, dtype=np.int) + cdef real[:, ::1] searchcoords + cdef real[:, ::1] searchcoords_bbox + cdef NSGrid searchgrid + ncells = self.grid.size + cdef real cutoff2 = self.cutoff * self.cutoff cdef ns_int npairs = 0 # Generate another grid to search - self.searchcoords_bbox = self.box.fast_put_atoms_in_bbox(searchcoords) - self.searchgrid = NSGrid(self.searchcoords_bbox.shape[0], self.cutoff, self.box, self.max_gridsize) - #cdef ns_int size = self.searchcoords_bbox.shape[0] - + searchcoords = np.array(search_coords, dtype=np.float32) + searchcoords_bbox = self.box.fast_put_atoms_in_bbox(searchcoords) + searchgrid = NSGrid(searchcoords_bbox.shape[0], self.grid.cutoff, self.box, self.max_gridsize, force=True) + searchgrid.fill_grid(searchcoords_bbox) size_search = searchcoords.shape[0] @@ -486,22 +479,18 @@ cdef class FastNS(object): with nogil: for i in range(size_search): - for m in range(ncells): - checked[m] = 0 # Start with first search coordinate current_beadid = i # find the cellindex of the coordinate - cellindex = self.searchgrid.cellids[current_beadid] - # Find the actual position of cell - self.searchgrid.cellid2cellxyz(cellindex, cellxyz) + cellindex = searchgrid.cellids[current_beadid] for xi in range(DIM): for yi in range(DIM): for zi in range(DIM): #Probe the search coordinates in a brick shaped box - probe[XX] = self.searchcoords_bbox[current_beadid, XX] + (xi - 1) * self.cutoff - probe[YY] = self.searchcoords_bbox[current_beadid, YY] + (yi - 1) * self.cutoff - probe[ZZ] = self.searchcoords_bbox[current_beadid, ZZ] + (zi - 1) * self.cutoff - # Make sure the shifted coordinates is inside the brick-shaped box + probe[XX] = searchcoords[current_beadid, XX] + (xi - 1) * searchgrid.cellsize[XX] + probe[YY] = searchcoords[current_beadid, YY] + (yi - 1) * searchgrid.cellsize[YY] + probe[ZZ] = searchcoords[current_beadid, ZZ] + (zi - 1) * searchgrid.cellsize[ZZ] + # Make sure the probe coordinates is inside the brick-shaped box for m in range(DIM - 1, -1, -1): while probe[m] < 0: for d in range(m+1): @@ -511,43 +500,36 @@ cdef class FastNS(object): probe[d] -= self.box.c_pbcbox.box[m][d] # Get the cell index corresponding to the probe cellindex_probe = self.grid.coord2cellid(probe) - if checked[cellindex_probe] != 0: - continue #for this cellindex search in grid for j in range(self.grid.nbeads[cellindex_probe]): bid = self.grid.beadids[cellindex_probe * self.grid.nbeads_per_cell + j] #find distance between search coords[i] and coords[bid] - d2 = self.box.fast_distance2(&self.searchcoords_bbox[current_beadid, XX], &self.coords_bbox[bid, XX]) + d2 = self.box.fast_distance2(&searchcoords_bbox[current_beadid, XX], &self.coords_bbox[bid, XX]) if d2 < cutoff2: if d2 < EPSILON: continue else: results.add_neighbors(current_beadid, bid, d2) npairs += 1 - checked[cellindex_probe] = 1 return results def self_search(self): cdef ns_int i, j, size_search cdef ns_int d, m - cdef ns_int ncells # Total number of cells cdef ns_int current_beadid, bid cdef ns_int cellindex, cellindex_probe cdef ns_int xi, yi, zi + cdef NSResults results - - - - cdef ivec cellxyz + + cdef ns_int size = self.coords_bbox.shape[0] + cdef ns_int[:] checked = np.zeros(size, dtype=np.int) cdef real d2 cdef rvec probe - ncells = self.grid.size - cdef ns_int[:] checked = np.zeros(ncells, dtype=np.int) - cdef real cutoff2 = self.cutoff * self.cutoff cdef ns_int npairs = 0 @@ -557,22 +539,18 @@ cdef class FastNS(object): with nogil: for i in range(size_search): - for m in range(ncells): - checked[m] = 0 # Start with first search coordinate current_beadid = i # find the cellindex of the coordinate cellindex = self.grid.cellids[current_beadid] - # Find the actual position of cell - self.grid.cellid2cellxyz(cellindex, cellxyz) for xi in range(DIM): for yi in range(DIM): for zi in range(DIM): # Calculate and/or reinitialize shifted coordinates #Probe the search coordinates in a brick shaped box - probe[XX] = self.coords_bbox[current_beadid, XX] + (xi - 1) * self.cutoff - probe[YY] = self.coords_bbox[current_beadid, YY] + (yi - 1) * self.cutoff - probe[ZZ] = self.coords_bbox[current_beadid, ZZ] + (zi - 1) * self.cutoff + probe[XX] = self.coords_bbox[current_beadid, XX] + (xi - 1) * self.grid.cellsize[XX] + probe[YY] = self.coords_bbox[current_beadid, YY] + (yi - 1) * self.grid.cellsize[XX] + probe[ZZ] = self.coords_bbox[current_beadid, ZZ] + (zi - 1) * self.grid.cellsize[XX] # Make sure the shifted coordinates is inside the brick-shaped box for m in range(DIM - 1, -1, -1): while probe[m] < 0: @@ -584,12 +562,10 @@ cdef class FastNS(object): # Get the cell index corresponding to the probe cellindex_probe = self.grid.coord2cellid(probe) - if checked[cellindex_probe] != 0: - continue #for this cellindex search in grid for j in range(self.grid.nbeads[cellindex_probe]): bid = self.grid.beadids[cellindex_probe * self.grid.nbeads_per_cell + j] - if bid > current_beadid: + if bid < current_beadid: continue #find distance between search coords[i] and coords[bid] d2 = self.box.fast_distance2(&self.coords_bbox[current_beadid, XX], &self.coords_bbox[bid, XX]) @@ -600,5 +576,4 @@ cdef class FastNS(object): results.add_neighbors(current_beadid, bid, d2) results.add_neighbors(bid, current_beadid, d2) npairs += 1 - checked[cellindex_probe] = 1 return results diff --git a/testsuite/MDAnalysisTests/lib/test_nsgrid.py b/testsuite/MDAnalysisTests/lib/test_nsgrid.py index 2808d1247aa..a6b94784cfe 100644 --- a/testsuite/MDAnalysisTests/lib/test_nsgrid.py +++ b/testsuite/MDAnalysisTests/lib/test_nsgrid.py @@ -39,8 +39,9 @@ def universe(): -def run_grid_search(u, searchcoords, cutoff=3): +def run_grid_search(u, ref_id, cutoff=3): coords = u.atoms.positions + searchcoords = u.atoms.positions[ref_id] if searchcoords.shape == (3, ): searchcoords = searchcoords[None, :] # Run grid search @@ -116,8 +117,8 @@ def test_nsgrid_badinit(): def test_nsgrid_badcutoff(universe): with pytest.raises(ValueError): - run_grid_search(universe, universe.atoms.positions[0], -4) - run_grid_search(universe, universe.atoms.positions[0], 100000) + run_grid_search(universe, 0, -4) + run_grid_search(universe, 0, 100000) def test_ns_grid_noneighbor(universe): @@ -125,7 +126,7 @@ def test_ns_grid_noneighbor(universe): ref_id = 0 cutoff = 0.5 - results_grid = run_grid_search(universe, universe.atoms.positions[ref_id], cutoff) + results_grid = run_grid_search(universe, ref_id, cutoff) assert len(results_grid.get_distances()[0]) == 0 assert len(results_grid.get_indices()[0]) == 0 @@ -142,7 +143,7 @@ def test_nsgrid_noPBC(universe): results = np.array([2, 3, 4, 5, 6, 7, 8, 9, 18, 19, 1211, 10862, 10865, 17582, 17585, 38342, 38345]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! - results_grid = run_grid_search(universe, universe.atoms.positions[ref_id], cutoff).get_indices()[0] + results_grid = run_grid_search(universe, ref_id, cutoff).get_indices()[0] assert_equal(results, results_grid) @@ -159,7 +160,7 @@ def test_nsgrid_PBC_rect(): # FastNS is called differently to max coverage searcher = nsgrid.FastNS(universe.dimensions, cutoff, universe.atoms.positions) - results_grid = searcher.search(universe.atoms.positions[ref_id][None, :]).get_indices()[0] # pass the id as a list for test+coverage purpose + results_grid = searcher.search(universe.atoms.positions[ref_id][None, :]).get_indices()[0] results_grid2 = searcher.search(universe.atoms.positions).get_indices() # call without specifying any ids, should do NS for all beads @@ -176,7 +177,7 @@ def test_nsgrid_PBC(universe): results = np.array([4398, 4401, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, 47451]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! - results_grid = run_grid_search(universe, universe.atoms.positions[ref_id]).get_indices()[0] + results_grid = run_grid_search(universe, ref_id).get_indices()[0] assert_equal(results, results_grid) @@ -191,7 +192,7 @@ def test_nsgrid_pairs(universe): results = np.array(results) - results_grid = run_grid_search(universe, universe.atoms.positions[ref_id]).get_pairs() + results_grid = run_grid_search(universe, ref_id).get_pairs() assert_equal(np.sort(neighbors, axis=0), np.sort(results_grid[:, 1], axis=0)) @@ -203,30 +204,11 @@ def test_nsgrid_pair_distances(universe): results = np.array([0.270, 0.285, 0.096, 0.096, 0.015, 0.278, 0.268, 0.179, 0.259, 0.290, 0.270]) * 10 # These distances where obtained by gmx distance so they are in nm - results_grid = run_grid_search(universe, universe.atoms.positions[ref_id]).get_pair_distances() + results_grid = run_grid_search(universe, ref_id).get_pair_distances() assert_allclose(np.sort(results), np.sort(results_grid), atol=1e-2) -# def test_nsgrid_pair_coordinates(universe): -# """Check that grid search return the proper pair coordinates""" - -# ref_id = 13937 -# neighbors = np.array([4398, 4401, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, -# 47451]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! -# coords = universe.atoms.positions - -# results = [] -# for nid in neighbors: -# if nid < ref_id: -# results.append([coords[nid], coords[ref_id]]) -# else: -# results.append([coords[ref_id], coords[nid]]) -# results = np.array(results) - -# results_grid = run_grid_search(universe, ref_id).get_pair_coordinates() - -# assert_allclose(np.sort(results, axis=0), np.sort(results_grid, axis=0), atol=1e-5) def test_nsgrid_distances(universe): @@ -236,20 +218,6 @@ def test_nsgrid_distances(universe): results = np.array([0.270, 0.285, 0.096, 0.096, 0.015, 0.278, 0.268, 0.179, 0.259, 0.290, 0.270]) * 10 # These distances where obtained by gmx distance so they are in nm - results_grid = run_grid_search(universe, universe.atoms.positions[ref_id]).get_distances()[0] + results_grid = run_grid_search(universe, ref_id).get_distances()[0] assert_allclose(np.sort(results), np.sort(results_grid), atol=1e-2) - - -# def test_nsgrid_coordinates(universe): -# """Check that grid search return the proper coordinates""" - -# ref_id = 13937 -# neighbors = np.array([4398, 4401, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, -# 47451]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! - -# results = universe.atoms.positions[neighbors] - -# results_grid = run_grid_search(universe, ref_id).get_coordinates()[0] - -# assert_allclose(np.sort(results, axis=0), np.sort(results_grid, axis=0), atol=1e-5) From 2c83681d155308322d2469ebe9b3c97d4edf8457 Mon Sep 17 00:00:00 2001 From: ayush Date: Tue, 31 Jul 2018 16:02:25 -0700 Subject: [PATCH 35/47] mend Removed dict from buffers from speedup, seperated search and self_search, construct a grid and search with another grid --- package/MDAnalysis/lib/distances.py | 14 ++------------ testsuite/MDAnalysisTests/lib/test_distances.py | 2 +- 2 files changed, 3 insertions(+), 13 deletions(-) diff --git a/package/MDAnalysis/lib/distances.py b/package/MDAnalysis/lib/distances.py index 4d662068a69..bec705c36e6 100644 --- a/package/MDAnalysis/lib/distances.py +++ b/package/MDAnalysis/lib/distances.py @@ -679,24 +679,14 @@ def _nsgrid_capped(reference, configuration, max_cutoff, min_cutoff=None, reference = reference[None, :] if configuration.shape == (3, ): configuration = configuration[None, :] - all_coords = np.concatenate([configuration, reference]) - mapping = np.arange(len(reference), dtype=np.int64) - gridsearch = FastNS(box, max_cutoff, all_coords) - gridsearch.prepare() - search_ids = np.arange(len(configuration), len(all_coords)) - results = gridsearch.search(search_ids=search_ids) + gridsearch = FastNS(box, max_cutoff, configuration) + results = gridsearch.search(reference) pairs = results.get_pairs() - pairs[:, 1] = undo_augment(pairs[:, 1], mapping, len(configuration)) pair_distance = results.get_pair_distances() if min_cutoff is not None: idx = pair_distance > min_cutoff pairs, pair_distance = pairs[idx], pair_distance[idx] - if pairs.size > 0: - # removing the pairs (i, j) from (reference, reference) - mask = (pairs[:, 0] < (len(configuration) - 1) ) - pairs, pair_distance = pairs[mask], pair_distance[mask] - pairs = np.sort(pairs, axis=1) return pairs, pair_distance diff --git a/testsuite/MDAnalysisTests/lib/test_distances.py b/testsuite/MDAnalysisTests/lib/test_distances.py index d83d8903e93..2bea060a778 100644 --- a/testsuite/MDAnalysisTests/lib/test_distances.py +++ b/testsuite/MDAnalysisTests/lib/test_distances.py @@ -57,7 +57,7 @@ def test_capped_distance_noresults(): npoints_1 = (1, 100) boxes_1 = (np.array([1, 2, 3, 90, 90, 90], dtype=np.float32), # ortho - np.array([1, 2, 3, 30, 45, 60], dtype=np.float32), # tri_box + np.array([1, 2, 3, 45, 60, 90], dtype=np.float32), # tri_box None, # Non Periodic ) From b05e2430b5542de435e000936598e5ddb1203bf9 Mon Sep 17 00:00:00 2001 From: ayush Date: Tue, 31 Jul 2018 19:23:52 -0700 Subject: [PATCH 36/47] reverted tests, removed sorting from get_indices, get_distances, modified compilation libs in setup --- package/MDAnalysis/lib/nsgrid.pyx | 5 ----- package/setup.py | 2 +- testsuite/MDAnalysisTests/lib/test_distances.py | 2 +- testsuite/MDAnalysisTests/lib/test_nsgrid.py | 8 ++++---- 4 files changed, 6 insertions(+), 11 deletions(-) diff --git a/package/MDAnalysis/lib/nsgrid.pyx b/package/MDAnalysis/lib/nsgrid.pyx index 8598bec1e97..7b8624a2836 100644 --- a/package/MDAnalysis/lib/nsgrid.pyx +++ b/package/MDAnalysis/lib/nsgrid.pyx @@ -277,11 +277,6 @@ cdef class NSResults(object): self.distances_buffer[beadid_i].push_back(dist) - for i in range(self.searchcoords.shape[0]): - sorted_indices = np.argsort(self.indices_buffer[i]) - self.indices_buffer[i] = np.array(self.indices_buffer[i])[sorted_indices] - self.distances_buffer[i] = np.array(self.distances_buffer[i])[sorted_indices] - def get_indices(self): if self.indices_buffer.empty(): diff --git a/package/setup.py b/package/setup.py index fe2e8bb7d1c..91a539db326 100755 --- a/package/setup.py +++ b/package/setup.py @@ -386,7 +386,7 @@ def extensions(config): language='c++', libraries=parallel_libraries, define_macros=define_macros + parallel_macros, - extra_compile_args=extra_compile_args + parallel_args, + extra_compile_args=cpp_extra_compile_args + parallel_args, extra_link_args=parallel_args) pre_exts = [libdcd, distances, distances_omp, qcprot, transformation, libmdaxdr, util, encore_utils, diff --git a/testsuite/MDAnalysisTests/lib/test_distances.py b/testsuite/MDAnalysisTests/lib/test_distances.py index 2bea060a778..d83d8903e93 100644 --- a/testsuite/MDAnalysisTests/lib/test_distances.py +++ b/testsuite/MDAnalysisTests/lib/test_distances.py @@ -57,7 +57,7 @@ def test_capped_distance_noresults(): npoints_1 = (1, 100) boxes_1 = (np.array([1, 2, 3, 90, 90, 90], dtype=np.float32), # ortho - np.array([1, 2, 3, 45, 60, 90], dtype=np.float32), # tri_box + np.array([1, 2, 3, 30, 45, 60], dtype=np.float32), # tri_box None, # Non Periodic ) diff --git a/testsuite/MDAnalysisTests/lib/test_nsgrid.py b/testsuite/MDAnalysisTests/lib/test_nsgrid.py index a6b94784cfe..41112e6bf25 100644 --- a/testsuite/MDAnalysisTests/lib/test_nsgrid.py +++ b/testsuite/MDAnalysisTests/lib/test_nsgrid.py @@ -145,7 +145,7 @@ def test_nsgrid_noPBC(universe): results_grid = run_grid_search(universe, ref_id, cutoff).get_indices()[0] - assert_equal(results, results_grid) + assert_equal(np.sort(results), np.sort(results_grid)) def test_nsgrid_PBC_rect(): @@ -164,10 +164,10 @@ def test_nsgrid_PBC_rect(): results_grid2 = searcher.search(universe.atoms.positions).get_indices() # call without specifying any ids, should do NS for all beads - assert_equal(results, results_grid) + assert_equal(np.sort(results), np.sort(results_grid)) assert_equal(len(universe.atoms), len(results_grid2)) assert searcher.cutoff == 7 - assert_equal(results_grid, results_grid2[ref_id]) + assert_equal(np.sort(results_grid), np.sort(results_grid2[ref_id])) def test_nsgrid_PBC(universe): @@ -179,7 +179,7 @@ def test_nsgrid_PBC(universe): results_grid = run_grid_search(universe, ref_id).get_indices()[0] - assert_equal(results, results_grid) + assert_equal(np.sort(results), np.sort(results_grid)) def test_nsgrid_pairs(universe): From 4f1e729c143f7a13721541e3d68ea5cb1c072e0f Mon Sep 17 00:00:00 2001 From: ayush Date: Wed, 1 Aug 2018 08:57:07 -0700 Subject: [PATCH 37/47] minor changes for failing tests --- package/MDAnalysis/lib/nsgrid.pyx | 38 +++++++++++-------------------- 1 file changed, 13 insertions(+), 25 deletions(-) diff --git a/package/MDAnalysis/lib/nsgrid.pyx b/package/MDAnalysis/lib/nsgrid.pyx index 7b8624a2836..518be1ef594 100644 --- a/package/MDAnalysis/lib/nsgrid.pyx +++ b/package/MDAnalysis/lib/nsgrid.pyx @@ -249,10 +249,10 @@ cdef class NSResults(object): self.pair_distances_buffer = np.sqrt(self.pair_distances2_buffer) return np.asarray(self.pair_distances_buffer) - cdef create_buffers(self): + cdef void create_buffers(self) nogil: cdef ns_int i, beadid_i, beadid_j cdef ns_int idx, nsearch - cdef real dist2 + cdef real dist2, dist nsearch = len(self.searchcoords) @@ -281,12 +281,12 @@ cdef class NSResults(object): def get_indices(self): if self.indices_buffer.empty(): self.create_buffers() - return np.asarray(self.indices_buffer) + return np.ascontiguousarray(self.indices_buffer) def get_distances(self): if self.distances_buffer.empty(): self.create_buffers() - return np.asarray(self.distances_buffer) + return np.ascontiguousarray(self.distances_buffer) cdef class NSGrid(object): cdef readonly real cutoff # cutoff @@ -456,13 +456,12 @@ cdef class FastNS(object): cdef real[:, ::1] searchcoords_bbox cdef NSGrid searchgrid - ncells = self.grid.size cdef real cutoff2 = self.cutoff * self.cutoff cdef ns_int npairs = 0 # Generate another grid to search - searchcoords = np.array(search_coords, dtype=np.float32) + searchcoords = np.asarray(search_coords, dtype=np.float32) searchcoords_bbox = self.box.fast_put_atoms_in_bbox(searchcoords) searchgrid = NSGrid(searchcoords_bbox.shape[0], self.grid.cutoff, self.box, self.max_gridsize, force=True) searchgrid.fill_grid(searchcoords_bbox) @@ -482,9 +481,9 @@ cdef class FastNS(object): for yi in range(DIM): for zi in range(DIM): #Probe the search coordinates in a brick shaped box - probe[XX] = searchcoords[current_beadid, XX] + (xi - 1) * searchgrid.cellsize[XX] - probe[YY] = searchcoords[current_beadid, YY] + (yi - 1) * searchgrid.cellsize[YY] - probe[ZZ] = searchcoords[current_beadid, ZZ] + (zi - 1) * searchgrid.cellsize[ZZ] + probe[XX] = searchcoords_bbox[current_beadid, XX] + (xi - 1) * searchgrid.cellsize[XX] + probe[YY] = searchcoords_bbox[current_beadid, YY] + (yi - 1) * searchgrid.cellsize[YY] + probe[ZZ] = searchcoords_bbox[current_beadid, ZZ] + (zi - 1) * searchgrid.cellsize[ZZ] # Make sure the probe coordinates is inside the brick-shaped box for m in range(DIM - 1, -1, -1): while probe[m] < 0: @@ -500,11 +499,8 @@ cdef class FastNS(object): bid = self.grid.beadids[cellindex_probe * self.grid.nbeads_per_cell + j] #find distance between search coords[i] and coords[bid] d2 = self.box.fast_distance2(&searchcoords_bbox[current_beadid, XX], &self.coords_bbox[bid, XX]) - if d2 < cutoff2: - if d2 < EPSILON: - continue - else: - results.add_neighbors(current_beadid, bid, d2) + if d2 < cutoff2 and d2 > EPSILON: + results.add_neighbors(current_beadid, bid, d2) npairs += 1 return results @@ -516,12 +512,7 @@ cdef class FastNS(object): cdef ns_int cellindex, cellindex_probe cdef ns_int xi, yi, zi - cdef NSResults results - - cdef ns_int size = self.coords_bbox.shape[0] - cdef ns_int[:] checked = np.zeros(size, dtype=np.int) - cdef real d2 cdef rvec probe @@ -564,11 +555,8 @@ cdef class FastNS(object): continue #find distance between search coords[i] and coords[bid] d2 = self.box.fast_distance2(&self.coords_bbox[current_beadid, XX], &self.coords_bbox[bid, XX]) - if d2 < cutoff2: - if d2 < EPSILON: - continue - else: - results.add_neighbors(current_beadid, bid, d2) - results.add_neighbors(bid, current_beadid, d2) + if d2 < cutoff2 and d2 > EPSILON: + results.add_neighbors(current_beadid, bid, d2) + results.add_neighbors(bid, current_beadid, d2) npairs += 1 return results From 7fd1f805192757464fdcd0fc22ef7344a9ba1d3d Mon Sep 17 00:00:00 2001 From: ayush Date: Wed, 1 Aug 2018 10:18:29 -0700 Subject: [PATCH 38/47] added hack to handle NoPBC, tests to follow --- package/MDAnalysis/lib/nsgrid.pyx | 130 ++++++++++++++++++++---------- 1 file changed, 89 insertions(+), 41 deletions(-) diff --git a/package/MDAnalysis/lib/nsgrid.pyx b/package/MDAnalysis/lib/nsgrid.pyx index 518be1ef594..3e3013b8e44 100644 --- a/package/MDAnalysis/lib/nsgrid.pyx +++ b/package/MDAnalysis/lib/nsgrid.pyx @@ -85,10 +85,13 @@ cdef struct cPBCBox_t: cdef class PBCBox(object): cdef cPBCBox_t c_pbcbox cdef bint is_triclinic + cdef bint periodic - def __init__(self, real[:,::1] box): + def __init__(self, real[:,::1] box, bint periodic): + self.periodic = periodic self.update(box) + cdef void fast_update(self, real[:,::1] box) nogil: cdef ns_int i, j @@ -146,14 +149,15 @@ cdef class PBCBox(object): for i in range(DIM): dx[i] = other[i] - ref[i] - for i in range (DIM-1, -1, -1): - while dx[i] > self.c_pbcbox.hbox_diag[i]: - for j in range (i, -1, -1): - dx[j] -= self.c_pbcbox.box[i][j] + if self.periodic: + for i in range (DIM-1, -1, -1): + while dx[i] > self.c_pbcbox.hbox_diag[i]: + for j in range (i, -1, -1): + dx[j] -= self.c_pbcbox.box[i][j] - while dx[i] <= self.c_pbcbox.mhbox_diag[i]: - for j in range (i, -1, -1): - dx[j] += self.c_pbcbox.box[i][j] + while dx[i] <= self.c_pbcbox.mhbox_diag[i]: + for j in range (i, -1, -1): + dx[j] += self.c_pbcbox.box[i][j] def dx(self, real[:] a, real[:] b): cdef rvec dx @@ -189,22 +193,23 @@ cdef class PBCBox(object): with gil: bbox_coords = coords.copy() - if self.is_triclinic: - for i in range(natoms): - for m in range(DIM - 1, -1, -1): - while bbox_coords[i, m] < 0: - for d in range(m+1): - bbox_coords[i, d] += self.c_pbcbox.box[m][d] - while bbox_coords[i, m] >= self.c_pbcbox.box[m][m]: - for d in range(m+1): - bbox_coords[i, d] -= self.c_pbcbox.box[m][d] - else: - for i in range(natoms): - for m in range(DIM): - while bbox_coords[i, m] < 0: - bbox_coords[i, m] += self.c_pbcbox.box[m][m] - while bbox_coords[i, m] >= self.c_pbcbox.box[m][m]: - bbox_coords[i, m] -= self.c_pbcbox.box[m][m] + if self.periodic: + if self.is_triclinic: + for i in range(natoms): + for m in range(DIM - 1, -1, -1): + while bbox_coords[i, m] < 0: + for d in range(m+1): + bbox_coords[i, d] += self.c_pbcbox.box[m][d] + while bbox_coords[i, m] >= self.c_pbcbox.box[m][m]: + for d in range(m+1): + bbox_coords[i, d] -= self.c_pbcbox.box[m][d] + else: + for i in range(natoms): + for m in range(DIM): + while bbox_coords[i, m] < 0: + bbox_coords[i, m] += self.c_pbcbox.box[m][m] + while bbox_coords[i, m] >= self.c_pbcbox.box[m][m]: + bbox_coords[i, m] -= self.c_pbcbox.box[m][m] return bbox_coords def put_atoms_in_bbox(self, real[:,::1] coords): @@ -410,17 +415,35 @@ cdef class FastNS(object): cdef readonly real cutoff cdef NSGrid grid cdef ns_int max_gridsize + cdef bint periodic - def __init__(self, box, cutoff, coords, max_gridsize=5000): + def __init__(self, cutoff, coords, box=None, max_gridsize=5000): import MDAnalysis as mda from MDAnalysis.lib.mdamath import triclinic_vectors + cdef real[:] pseudobox = np.zeros(6, dtype=np.float32) + cdef real[DIM] bmax, bmin + cdef ns_int i + self.periodic = True + + if (box is None) or np.allclose(box[:3], 0.): + bmax = np.max(coords, axis=0) + bmin = np.min(coords, axis=0) + for i in range(DIM): + pseudobox[i] = 1.1*(bmax - bmin) + pseudobox[DIM + i] = 90. + box = pseudobox + # shift the origin + coords -= bmin + self.periodic = False + + if box.shape != (3,3): box = triclinic_vectors(box) - self.box = PBCBox(box) + self.box = PBCBox(box, self.periodic) if cutoff < 0: raise ValueError("Cutoff must be positive!") @@ -455,6 +478,7 @@ cdef class FastNS(object): cdef real[:, ::1] searchcoords cdef real[:, ::1] searchcoords_bbox cdef NSGrid searchgrid + cdef bint check cdef real cutoff2 = self.cutoff * self.cutoff @@ -480,18 +504,30 @@ cdef class FastNS(object): for xi in range(DIM): for yi in range(DIM): for zi in range(DIM): + check = True #Probe the search coordinates in a brick shaped box probe[XX] = searchcoords_bbox[current_beadid, XX] + (xi - 1) * searchgrid.cellsize[XX] probe[YY] = searchcoords_bbox[current_beadid, YY] + (yi - 1) * searchgrid.cellsize[YY] probe[ZZ] = searchcoords_bbox[current_beadid, ZZ] + (zi - 1) * searchgrid.cellsize[ZZ] # Make sure the probe coordinates is inside the brick-shaped box - for m in range(DIM - 1, -1, -1): - while probe[m] < 0: - for d in range(m+1): - probe[d] += self.box.c_pbcbox.box[m][d] - while probe[m] >= self.box.c_pbcbox.box[m][m]: - for d in range(m+1): - probe[d] -= self.box.c_pbcbox.box[m][d] + if self.periodic: + for m in range(DIM - 1, -1, -1): + while probe[m] < 0: + for d in range(m+1): + probe[d] += self.box.c_pbcbox.box[m][d] + while probe[m] >= self.box.c_pbcbox.box[m][m]: + for d in range(m+1): + probe[d] -= self.box.c_pbcbox.box[m][d] + else: + for m in range(DIM, -1, -1, -1): + if probe[m] < 0: + check = False + break + if probe[m] > self.box.c_pbcbox.box[m][m]: + check = False + break + if not check: + continue # Get the cell index corresponding to the probe cellindex_probe = self.grid.coord2cellid(probe) #for this cellindex search in grid @@ -518,6 +554,7 @@ cdef class FastNS(object): cdef real cutoff2 = self.cutoff * self.cutoff cdef ns_int npairs = 0 + cdef bint check size_search = self.coords.shape[0] @@ -532,20 +569,31 @@ cdef class FastNS(object): for xi in range(DIM): for yi in range(DIM): for zi in range(DIM): + check = True # Calculate and/or reinitialize shifted coordinates #Probe the search coordinates in a brick shaped box probe[XX] = self.coords_bbox[current_beadid, XX] + (xi - 1) * self.grid.cellsize[XX] probe[YY] = self.coords_bbox[current_beadid, YY] + (yi - 1) * self.grid.cellsize[XX] probe[ZZ] = self.coords_bbox[current_beadid, ZZ] + (zi - 1) * self.grid.cellsize[XX] # Make sure the shifted coordinates is inside the brick-shaped box - for m in range(DIM - 1, -1, -1): - while probe[m] < 0: - for d in range(m+1): - probe[d] += self.box.c_pbcbox.box[m][d] - while probe[m] >= self.box.c_pbcbox.box[m][m]: - for d in range(m+1): - probe[d] -= self.box.c_pbcbox.box[m][d] - + if self.periodic: + for m in range(DIM - 1, -1, -1): + while probe[m] < 0: + for d in range(m+1): + probe[d] += self.box.c_pbcbox.box[m][d] + while probe[m] >= self.box.c_pbcbox.box[m][m]: + for d in range(m+1): + probe[d] -= self.box.c_pbcbox.box[m][d] + else: + for m in range(DIM, -1, -1, -1): + if probe[m] < 0: + check = False + break + elif probe[m] > self.box.c_pbcbox.box[m][m]: + check = False + break + if not check: + continue # Get the cell index corresponding to the probe cellindex_probe = self.grid.coord2cellid(probe) #for this cellindex search in grid From 57e9d55f5c4b546ee2cbdbe794578c550a77fe70 Mon Sep 17 00:00:00 2001 From: ayush Date: Wed, 1 Aug 2018 12:31:07 -0700 Subject: [PATCH 39/47] conform tests a/c to modifications in API --- package/MDAnalysis/lib/distances.py | 2 +- testsuite/MDAnalysisTests/lib/test_nsgrid.py | 18 +++++++++--------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/package/MDAnalysis/lib/distances.py b/package/MDAnalysis/lib/distances.py index bec705c36e6..729ebb1e3ba 100644 --- a/package/MDAnalysis/lib/distances.py +++ b/package/MDAnalysis/lib/distances.py @@ -679,7 +679,7 @@ def _nsgrid_capped(reference, configuration, max_cutoff, min_cutoff=None, reference = reference[None, :] if configuration.shape == (3, ): configuration = configuration[None, :] - gridsearch = FastNS(box, max_cutoff, configuration) + gridsearch = FastNS(max_cutoff, configuration, box=box) results = gridsearch.search(reference) pairs = results.get_pairs() pair_distance = results.get_pair_distances() diff --git a/testsuite/MDAnalysisTests/lib/test_nsgrid.py b/testsuite/MDAnalysisTests/lib/test_nsgrid.py index 41112e6bf25..cd9e514f3a1 100644 --- a/testsuite/MDAnalysisTests/lib/test_nsgrid.py +++ b/testsuite/MDAnalysisTests/lib/test_nsgrid.py @@ -45,7 +45,7 @@ def run_grid_search(u, ref_id, cutoff=3): if searchcoords.shape == (3, ): searchcoords = searchcoords[None, :] # Run grid search - searcher = nsgrid.FastNS(u.dimensions, cutoff, coords) + searcher = nsgrid.FastNS(cutoff, coords, box=u.dimensions) return searcher.search(searchcoords) @@ -56,11 +56,11 @@ def test_pbc_badbox(): nsgrid.PBCBox([]) with pytest.raises(ValueError): - nsgrid.PBCBox(np.zeros((3))) # Bad shape - nsgrid.PBCBox(np.zeros((3, 3))) # Collapsed box - nsgrid.PBCBOX(np.array([[0, 0, 0], [0, 1, 0], [0, 0, 1]])) # 2D box - nsgrid.PBCBOX(np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])) # Box provided as array of integers - nsgrid.PBCBOX(np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=float)) # Box provided as array of double + nsgrid.PBCBox(np.zeros((3)), True) # Bad shape + nsgrid.PBCBox(np.zeros((3, 3)), True) # Collapsed box + nsgrid.PBCBOX(np.array([[0, 0, 0], [0, 1, 0], [0, 0, 1]]), True) # 2D box + nsgrid.PBCBOX(np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]), True) # Box provided as array of integers + nsgrid.PBCBOX(np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=float), True) # Box provided as array of double def test_pbc_distances(): @@ -70,7 +70,7 @@ def test_pbc_distances(): a = np.array([0.1, 0.1, 0.1], dtype=np.float32) b = np.array([1.1, -0.1, 0.2], dtype=np.float32) dx = np.array([0, -0.2, 0.1], dtype=np.float32) - pbcbox = nsgrid.PBCBox(box) + pbcbox = nsgrid.PBCBox(box, True) with pytest.raises(ValueError): pbcbox.distance(a, bad) @@ -105,7 +105,7 @@ def test_pbc_put_in_bbox(): dtype=np.float32 ) - pbcbox = nsgrid.PBCBox(box) + pbcbox = nsgrid.PBCBox(box, True) assert_allclose(pbcbox.put_atoms_in_bbox(coords), results, atol=1e-5) @@ -158,7 +158,7 @@ def test_nsgrid_PBC_rect(): cutoff = 7 # FastNS is called differently to max coverage - searcher = nsgrid.FastNS(universe.dimensions, cutoff, universe.atoms.positions) + searcher = nsgrid.FastNS(cutoff, universe.atoms.positions, box=universe.dimensions) results_grid = searcher.search(universe.atoms.positions[ref_id][None, :]).get_indices()[0] From a7264ec6cda550acf49f53ed62b70f3565b7929c Mon Sep 17 00:00:00 2001 From: ayush Date: Thu, 2 Aug 2018 02:30:58 -0700 Subject: [PATCH 40/47] removed python functions, added no pbc handle using python min, max --- package/MDAnalysis/lib/nsgrid.pyx | 32 +---- testsuite/MDAnalysisTests/lib/test_nsgrid.py | 125 ++++++++----------- 2 files changed, 58 insertions(+), 99 deletions(-) diff --git a/package/MDAnalysis/lib/nsgrid.pyx b/package/MDAnalysis/lib/nsgrid.pyx index 3e3013b8e44..ab0df6d2e46 100644 --- a/package/MDAnalysis/lib/nsgrid.pyx +++ b/package/MDAnalysis/lib/nsgrid.pyx @@ -159,31 +159,12 @@ cdef class PBCBox(object): for j in range (i, -1, -1): dx[j] += self.c_pbcbox.box[i][j] - def dx(self, real[:] a, real[:] b): - cdef rvec dx - if a.shape[0] != DIM or b.shape[0] != DIM: - raise ValueError("Not 3 D coordinates") - self.fast_pbc_dx(&a[XX], &b[XX], dx) - return np.array([dx[XX], dx[YY], dx[ZZ]], dtype=np.float32) - cdef real fast_distance2(self, rvec a, rvec b) nogil: cdef rvec dx self.fast_pbc_dx(a, b, dx) return rvec_norm2(dx) - def distance2(self, real[:] a, real[:] b): - if a.shape[0] != DIM or b.shape[0] != DIM: - raise ValueError("Not 3 D coordinates") - return self.fast_distance2(&a[XX], &b[XX]) - - cdef real fast_distance(self, rvec a, rvec b) nogil: - return sqrt(self.fast_distance2(a,b)) - - def distance(self, real[:] a, real[:] b): - if a.shape[0] != DIM or b.shape[0] != DIM: - raise ValueError("Not 3 D coordinates") - return self.fast_distance(&a[XX], &b[XX]) cdef real[:, ::1]fast_put_atoms_in_bbox(self, real[:,::1] coords) nogil: cdef ns_int i, m, d, natoms @@ -212,9 +193,6 @@ cdef class PBCBox(object): bbox_coords[i, m] -= self.c_pbcbox.box[m][m] return bbox_coords - def put_atoms_in_bbox(self, real[:,::1] coords): - return np.asarray(self.fast_put_atoms_in_bbox(coords)) - ######################### # Neighbor Search Stuff # ######################### @@ -427,11 +405,12 @@ cdef class FastNS(object): cdef ns_int i self.periodic = True - if (box is None) or np.allclose(box[:3], 0.): + + if (box is None) or (np.allclose(box[:3], 0.) and box.shape[0] == 6): bmax = np.max(coords, axis=0) bmin = np.min(coords, axis=0) for i in range(DIM): - pseudobox[i] = 1.1*(bmax - bmin) + pseudobox[i] = 1.1*(bmax[i] - bmin[i]) pseudobox[DIM + i] = 90. box = pseudobox # shift the origin @@ -439,7 +418,6 @@ cdef class FastNS(object): self.periodic = False - if box.shape != (3,3): box = triclinic_vectors(box) @@ -519,7 +497,7 @@ cdef class FastNS(object): for d in range(m+1): probe[d] -= self.box.c_pbcbox.box[m][d] else: - for m in range(DIM, -1, -1, -1): + for m in range(DIM -1, -1, -1): if probe[m] < 0: check = False break @@ -585,7 +563,7 @@ cdef class FastNS(object): for d in range(m+1): probe[d] -= self.box.c_pbcbox.box[m][d] else: - for m in range(DIM, -1, -1, -1): + for m in range(DIM -1, -1, -1): if probe[m] < 0: check = False break diff --git a/testsuite/MDAnalysisTests/lib/test_nsgrid.py b/testsuite/MDAnalysisTests/lib/test_nsgrid.py index cd9e514f3a1..bcb335f59a8 100644 --- a/testsuite/MDAnalysisTests/lib/test_nsgrid.py +++ b/testsuite/MDAnalysisTests/lib/test_nsgrid.py @@ -24,7 +24,7 @@ import pytest -from numpy.testing import assert_equal, assert_allclose +from numpy.testing import assert_equal, assert_allclose, assert_array_equal import numpy as np import MDAnalysis as mda @@ -50,69 +50,18 @@ def run_grid_search(u, ref_id, cutoff=3): return searcher.search(searchcoords) -def test_pbc_badbox(): +def test_pbc_box(): """Check that PBC box accepts only well-formated boxes""" + pbc = True with pytest.raises(TypeError): nsgrid.PBCBox([]) with pytest.raises(ValueError): - nsgrid.PBCBox(np.zeros((3)), True) # Bad shape - nsgrid.PBCBox(np.zeros((3, 3)), True) # Collapsed box - nsgrid.PBCBOX(np.array([[0, 0, 0], [0, 1, 0], [0, 0, 1]]), True) # 2D box - nsgrid.PBCBOX(np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]), True) # Box provided as array of integers - nsgrid.PBCBOX(np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=float), True) # Box provided as array of double - - -def test_pbc_distances(): - """Check that PBC box computes distances""" - box = np.identity(3, dtype=np.float32) - bad = np.array([0.1, 0.2], dtype=np.float32) - a = np.array([0.1, 0.1, 0.1], dtype=np.float32) - b = np.array([1.1, -0.1, 0.2], dtype=np.float32) - dx = np.array([0, -0.2, 0.1], dtype=np.float32) - pbcbox = nsgrid.PBCBox(box, True) - - with pytest.raises(ValueError): - pbcbox.distance(a, bad) - pbcbox.distance(bad, a) - - pbcbox.distance2(a, bad) - pbcbox.distance2(bad, a) - - pbcbox.dx(bad, a) - pbcbox.dx(a, bad) - - assert_equal(pbcbox.dx(a, b), dx) - assert_allclose(pbcbox.distance(a, b), np.sqrt(np.sum(dx*dx)), atol=1e-5) - assert_allclose(pbcbox.distance2(a, b), np.sum(dx*dx), atol=1e-5) - - -def test_pbc_put_in_bbox(): - "Check that PBC put beads in brick-shaped box" - box = np.identity(3, dtype=np.float32) - coords = np.array( - [ - [0.1, 0.1, 0.1], - [-0.1, 1.1, 0.9] - ], - dtype=np.float32 - ) - results = np.array( - [ - [0.1, 0.1, 0.1], - [0.9, 0.1, 0.9] - ], - dtype=np.float32 - ) - - pbcbox = nsgrid.PBCBox(box, True) - - assert_allclose(pbcbox.put_atoms_in_bbox(coords), results, atol=1e-5) - - -def test_nsgrid_badinit(): - with pytest.raises(TypeError): - nsgrid.FastNS(None, 1) + nsgrid.PBCBox(np.zeros((3)), pbc) # Bad shape + nsgrid.PBCBox(np.zeros((3, 3)), pbc) # Collapsed box + nsgrid.PBCBOX(np.array([[0, 0, 0], [0, 1, 0], [0, 0, 1]]), pbc) # 2D box + nsgrid.PBCBOX(np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]), pbc) # Box provided as array of integers + nsgrid.PBCBOX(np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=float), pbc) # Box provided as array of double def test_nsgrid_badcutoff(universe): @@ -134,19 +83,6 @@ def test_ns_grid_noneighbor(universe): assert len(results_grid.get_pair_distances()) == 0 -def test_nsgrid_noPBC(universe): - """Check that grid search works when no PBC is needed""" - - ref_id = 0 - - cutoff = 3 - results = np.array([2, 3, 4, 5, 6, 7, 8, 9, 18, 19, 1211, 10862, 10865, 17582, 17585, 38342, - 38345]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! - - results_grid = run_grid_search(universe, ref_id, cutoff).get_indices()[0] - - assert_equal(np.sort(results), np.sort(results_grid)) - def test_nsgrid_PBC_rect(): """Check that nsgrid works with rect boxes and PBC""" @@ -221,3 +157,48 @@ def test_nsgrid_distances(universe): results_grid = run_grid_search(universe, ref_id).get_distances()[0] assert_allclose(np.sort(results), np.sort(results_grid), atol=1e-2) + + +@pytest.mark.parametrize('box, results', + ((None, [3, 13, 24]), + (np.array([0., 0., 0., 90., 90., 90.]), [3, 13, 24]), + (np.array([10., 10., 10., 90., 90., 90.]), [3, 13, 24, 39, 67]), + (np.array([10., 10., 10., 60., 75., 90.]), [3, 13, 24, 39, 60, 79]))) +def test_nsgrid_search(box, results): + np.random.seed(90003) + points = (np.random.uniform(low=0, high=1.0, + size=(100, 3))*(10.)).astype(np.float32) + cutoff = 2.0 + query = np.array([1., 1., 1.]).reshape((1,3)) + searcher = nsgrid.FastNS(cutoff, points, box=box) + searchresults = searcher.search(query) + indices = searchresults.get_indices()[0] + assert_equal(np.sort(indices), results) + + +@pytest.mark.parametrize('box, result', + ((None, 21), + (np.array([0., 0., 0., 90., 90., 90.]), 21), + (np.array([10., 10., 10., 90., 90., 90.]), 26), + (np.array([10., 10., 10., 60., 75., 90.]), 33))) +def test_nsgrid_selfsearch(box, result): + np.random.seed(90003) + points = (np.random.uniform(low=0, high=1.0, + size=(100, 3))*(10.)).astype(np.float32) + cutoff = 1.0 + searcher = nsgrid.FastNS(cutoff, points, box=box) + searchresults = searcher.self_search() + pairs = searchresults.get_pairs() + assert_equal(len(pairs)//2, result) + +def test_contiguous(universe): + ref_id = 13937 + cutoff = 3 + coords = universe.atoms.positions + searcher = nsgrid.FastNS(cutoff, coords, box=None) + searchresults = searcher.search(coords[ref_id][None, :]) + indices = searchresults.get_indices() + distances = searchresults.get_distances() + + assert indices.flags['C_CONTIGUOUS'] == True + assert distances.flags['C_CONTIGUOUS'] == True From 929fbf589d0a27ac7aeb6fc4bb3061f9647405c4 Mon Sep 17 00:00:00 2001 From: ayush Date: Fri, 3 Aug 2018 02:37:28 -0700 Subject: [PATCH 41/47] docs unchecked --- package/MDAnalysis/lib/nsgrid.pyx | 366 +++++++++++++++++++++++++++++- 1 file changed, 362 insertions(+), 4 deletions(-) diff --git a/package/MDAnalysis/lib/nsgrid.pyx b/package/MDAnalysis/lib/nsgrid.pyx index ab0df6d2e46..85207f9da1b 100644 --- a/package/MDAnalysis/lib/nsgrid.pyx +++ b/package/MDAnalysis/lib/nsgrid.pyx @@ -28,10 +28,41 @@ # cython: initializedcheck=False """ -Neighbor search library --- :mod:`MDAnalysis.lib.grid` -====================================================== +Neighbor search library --- :mod:`MDAnalysis.lib.nsgrid` +======================================================== + +About the code +--------------- + +This Neighbor search library is a serialized Cython version greatly +inspired by the NS grid search implemented in +`GROMACS `_ . + +GROMACS 4.x code (more precisely `nsgrid.c `_ +and `ns.c `_ ) +was used as reference to write this file. + +GROMACS 4.x code is released under the GNU Public Licence v2. + +About the algorithm +-------------------- + +The neighbor search implemented here is based on `cell lists `_ which allow +computation of pairs [#]_ with a cost of :math:`O(N)`, instead of :math:`O(N^2)`. +The basic algorithm is described in Appendix F, Page 552 of +``Understanding Molecular Dynamics: From Algorithm to Applications`` by Frenkel and Smit. + +In brief, the algorithm divides the domain into smaller subdomains called `cells` +and distributes every particle to these cells based on its position. Subsequently, +any distance based query first identifies the corresponding cell position in the +domain followed by distance evaluations within the identified cell and +neighboring cells only. Care must be taken to ensure that `cellsize` is +greater than the desired search distance, otherwise all of the neighbours might +not reflect in the results. + + +.. [#] a pair correspond to two particles that are considered as neighbors . -This Neighbor search library is a serialized Cython port of the NS grid search implemented in GROMACS. """ @@ -83,16 +114,47 @@ cdef struct cPBCBox_t: # Class to handle PBC calculations cdef class PBCBox(object): + """ + Cython implementation of `PBC-related `_ operations. + This class is used by classes :class:`FastNS` and :class:`NSGrid` to put all particles inside a brick-shaped box + and to compute PBC-aware distance. The class can also handle non-PBC aware distance evaluations through + ``periodic`` argument. + + .. warning:: + This class is not meant to be used by end users. + + .. warning:: + Even if MD triclinic boxes can be handled by this class, internal optimization is made based on the + assumption that particles are inside a brick-shaped box. When this is not the case, calculated distances are not + warranted to be exact. + """ cdef cPBCBox_t c_pbcbox cdef bint is_triclinic cdef bint periodic def __init__(self, real[:,::1] box, bint periodic): + """ + Parameters + ---------- + box : numpy.ndarray + box vectors of shape ``(3, 3)`` or + as returned by ``MDAnalysis.lib.mdamath.triclinic_vectors`` + ``dtype`` must be ``numpy.float32`` + periodic : boolean + ``True`` for PBC-aware calculations + ``False`` for non PBC aware calculations + """ self.periodic = periodic self.update(box) cdef void fast_update(self, real[:,::1] box) nogil: + """ + Updates the internal box parameters for + PBC-aware distance calculations. The internal + box parameters are used to define the brick-shaped + box which is eventually used for distance calculations. + """ cdef ns_int i, j cdef real min_hv2, min_ss, tmp @@ -134,6 +196,23 @@ cdef class PBCBox(object): self.c_pbcbox.max_cutoff2 = min(min_hv2, min_ss * min_ss) def update(self, real[:,::1] box): + """ + Updates internal MD box representation and parameters used for calculations. + + + Parameters + ---------- + box : numpy.ndarray + Describes the MD box vectors as returned by + :func:`MDAnalysis.lib.mdamath.triclinic_vectors`. + `dtype` must be :class:`numpy.float32` + + Note + ---- + Call to this method is only needed when the MD box is changed + as it always called when class is instantiated. + + """ if box.shape[0] != DIM or box.shape[1] != DIM: raise ValueError("Box must be a {} x {} matrix. Got: {} x {})".format( DIM, DIM, box.shape[0], box.shape[1])) @@ -143,6 +222,15 @@ cdef class PBCBox(object): cdef void fast_pbc_dx(self, rvec ref, rvec other, rvec dx) nogil: + """Dislacement between two points for both + PBC and non-PBC conditions + + Modifies the displacement vector between two points based + on the minimum image convention for PBC aware calculations. + + For non-PBC aware distance evaluations, calculates the + displacement vector without any modifications + """ cdef ns_int i, j @@ -161,12 +249,25 @@ cdef class PBCBox(object): cdef real fast_distance2(self, rvec a, rvec b) nogil: + """Distance calculation between two points + for both PBC and non-PBC aware calculations + + Returns the distance obeying minimum + image convention if periodic is set to ``True`` while + instantiating the ``PBCBox`` object. + """ cdef rvec dx self.fast_pbc_dx(a, b, dx) return rvec_norm2(dx) cdef real[:, ::1]fast_put_atoms_in_bbox(self, real[:,::1] coords) nogil: + """Shifts all ``coords`` to an orthogonal brick shaped box + + All the coordinates are brought into an orthogonal + box. The box vectors for the brick-shaped box + are defined in ``fast_update`` method. + """ cdef ns_int i, m, d, natoms cdef real[:,::1] bbox_coords @@ -198,6 +299,26 @@ cdef class PBCBox(object): ######################### cdef class NSResults(object): + """Class to store the results + + All the required outputs from :class:`FastNS` is stored in the + instance of this class. All the methods of :class:`FastNS` returns + an instance of this class, which can be used to generate the desired + results on demand. + + While more details can be found in individual method description, + a brief list of methods and their returns can be summarized as follows + - ``get_pairs`` : All the pairs of particles within the specified cutoff distance. + - ``get_pair_distances`` : Distance corresponding to every pair of neighbors + - ``get_indices`` : An array with arrays of neighbour corresponding to every individual particle. + - ``get_distances`` : Corresponding distance for every neighbor in ``get_indices`` + + + SeeAlso + ------- + MDAnalysis.lib.nsgrid.FastNS + + """ cdef readonly real cutoff cdef ns_int npairs @@ -211,6 +332,18 @@ cdef class NSResults(object): cdef vector[real] pair_distances2_buffer def __init__(self, real cutoff, real[:, ::1]coords, real[:, ::1]searchcoords): + """ + Parameters + ---------- + cutoff : float + Specified cutoff distance + coords : numpy.ndarray + Array with coordinates of atoms of shape ``(N, 3)`` for + ``N`` particles. ``dtype`` must be ``numpy.float32`` + searchcoords : numpy.ndarray + Array with query coordinates. Shape must be ``(M, 3)`` + for ``M`` queries. ``dtype`` must be ``numpy.float32`` + """ self.cutoff = cutoff self.coords = coords self.searchcoords = searchcoords @@ -218,21 +351,58 @@ cdef class NSResults(object): self.npairs = 0 cdef void add_neighbors(self, ns_int beadid_i, ns_int beadid_j, real distance2) nogil: - # Important: If this function returns ERROR, it means that memory allocation failed + """Internal function to add pairs and distances to buffers + The buffers populated using this method are used by other methods of this class. + This is the primary function used by :class:`FastNS` to save the pair of atoms, + which can be considered as neighbors. + """ self.pairs_buffer.push_back(beadid_i) self.pairs_buffer.push_back(beadid_j) self.pair_distances2_buffer.push_back(distance2) self.npairs += 1 def get_pairs(self): + """Returns all the pairs within the desired cutoff distance + + Returns an array of shape ``(N, 2)``, where N is the number of pairs + between ``reference`` and ``configuration`` within the specified distance. + For every pair ``(i, j)``, ``reference[i]`` and ``configuration[j]`` are + atom positions such that ``reference`` is the position of query + atoms while ``configuration`` coontains the position of group of + atoms used to search against the query atoms. + + Returns + ------- + pairs : numpy.ndarray + pairs of atom indices of neighbors from query + and initial atom coordinates of shape ``(N, 2)`` + """ return np.asarray(self.pairs_buffer).reshape(self.npairs, 2) def get_pair_distances(self): + """Returns all the distances corresponding to each pair of neighbors + + Returns an array of shape ``N`` where N is the number of pairs + among the query atoms and initial atoms within a specified distance. + Every element ``[i]`` corresponds to the distance between + ``pairs[i, 0]`` and ``pairs[i, 1]``, where pairs is the array + obtained from ``get_pairs()`` + + SeeAlso + ------- + MDAnalysis.lib.nsgrid.NSResults.get_pairs + + """ self.pair_distances_buffer = np.sqrt(self.pair_distances2_buffer) return np.asarray(self.pair_distances_buffer) cdef void create_buffers(self) nogil: + """ + Creates buffers to get individual neighbour list and distances + of the query atoms. + + """ cdef ns_int i, beadid_i, beadid_j cdef ns_int idx, nsearch cdef real dist2, dist @@ -262,16 +432,72 @@ cdef class NSResults(object): def get_indices(self): + """Individual neighbours of query atom + + For every queried atom ``i``, an array of all its neighbors + indices can be obtained from ``get_indices()[i]`` + + Returns + ------- + indices : np.ndarray + Indices of neighboring atoms. + Every element i.e. ``indices[i]`` will be an array of + shape ``m`` where m is the number of neighbours of + query atom[i]. ``dtype=numpy.int`` + + """ if self.indices_buffer.empty(): self.create_buffers() return np.ascontiguousarray(self.indices_buffer) def get_distances(self): + """Distance corresponding to individual neighbors of query atom + + For every queried atom ``i``, an array of all the distances + from its neighbors can be obtained from ``get_distances()[i]``. + Every ``distance[i, j]`` will correspond + to the distance between atoms whose indices can be obtained + from ``get_indices()[i, j]`` + + Returns + ------- + distances : np.ndarray + Every element i.e. ``distances[i]`` will be an array of + shape ``m`` where m is the number of neighbours of + query atom[i]. + + SeeAlso + ------- + MDAnalysis.lib.nsgrid.NSResults.get_indices + + """ if self.distances_buffer.empty(): self.create_buffers() return np.ascontiguousarray(self.distances_buffer) cdef class NSGrid(object): + """Constructs a uniform cuboidal grid for a brick-shaped box + + This class uses :class:`PBCBox` to define the brick shaped box. + It is essential to initialize the box with :class:`PBCBox` + inorder to form the grid. + + The domain is subdivided into number of cells based on the desired search + radius. Ideally cellsize should be equal to the search radius, but small + search radius leads to large cell-list data strucutres. + An optimization of cutoff is imposed to limit the size of data + structure such that the cellsize is always greater than or + equal to cutoff distance. + + Note + ---- + This class assumes that all the coordinates are already + inside the brick shaped box. Care must be taken to ensure + all the particles are within the brick shaped box as + defined by :class:`PBCBox`. This can be ensured by using + :func:`~MDAnalysis.lib.nsgrid.PBCBox.fast_put_atoms_in_bbox` + + """ cdef readonly real cutoff # cutoff cdef ns_int size # total cells cdef ns_int ncoords # number of coordinates @@ -285,6 +511,20 @@ cdef class NSGrid(object): cdef bint force # To negate the effects of optimized cutoff def __init__(self, ncoords, cutoff, PBCBox box, max_size, force=False): + """ + Parameters + ---------- + ncoords : int + Number of coordinates to fill inside the brick shaped box + cutoff : float + Desired cutoff radius + box : PBCBox + Instance of :class:`PBCBox` + max_size : int + Maximum total number of cells + force : boolean + Optimizes cutoff if set to ``False`` [False] + """ cdef ns_int i cdef ns_int ncellx, ncelly, ncellz, size, nbeadspercell cdef ns_int xi, yi, zi @@ -331,11 +571,21 @@ cdef class NSGrid(object): PyMem_Free(self.cellids) cdef ns_int coord2cellid(self, rvec coord) nogil: + """Finds the cell-id for the given coordinate inside the brick shaped box + + Note + ---- + Assumes the coordinate is already inside the brick shaped box. + Return wrong cell-id if this is not the case + """ return (coord[ZZ] / self.cellsize[ZZ]) * (self.cell_offsets[ZZ]) +\ (coord[YY] / self.cellsize[YY]) * self.cell_offsets[YY] + \ (coord[XX] / self.cellsize[XX]) + cdef bint cellid2cellxyz(self, ns_int cellid, ivec cellxyz) nogil: + """Finds actual cell position `(x, y, z)` from a cell-id + """ if cellid < 0: return False if cellid >= self.size: @@ -350,6 +600,17 @@ cdef class NSGrid(object): return True cdef fill_grid(self, real[:, ::1] coords): + """Sorts atoms into cells based on their position in the brick shaped box + + Every atom inside the brick shaped box is assigned a + cell-id based on its position. Another list ``beadids`` + sort the atom-ids in each cell. + + Note + ---- + The method fails if any coordinate is outside the brick shaped box. + + """ cdef ns_int i, cellindex = -1 cdef ns_int ncoords = coords.shape[0] cdef ns_int[:] beadcounts = np.empty(self.size, dtype=np.int) @@ -386,6 +647,15 @@ cdef class NSGrid(object): cdef class FastNS(object): + """Grid based search between two group of atoms + + Instantiates a class object which uses :class:`PBCBox` and + :class:`NSGrid` to construct a cuboidal + grid in an orthogonal brick shaped box. + + Minimum image convention is used for distance evaluations + if box dimensions are provided. + """ cdef PBCBox box cdef real[:, ::1] coords @@ -397,6 +667,40 @@ cdef class FastNS(object): def __init__(self, cutoff, coords, box=None, max_gridsize=5000): + """ + Initialize the grid and sort the coordinates in respective + cells by shifting the coordinates in a brick shaped box. + The brick shaped box is defined by :class:`PBCBox` + and cuboidal grid is initialize by :class:`NSGrid`. + If box is supplied, periodic shifts along box vectors are used + to contain all the coordinates inside the brick shaped box. + If box is not supplied, the range of coordinates i.e. + ``[xmax, ymax, zmax] - [xmin, ymin, zmin]`` is used to construct + a pseudo box. Subsequently, the origin is also shifted + to the ``[xmin, ymin, zmin]``. + + Parameters + ---------- + cutoff : float + Desired cutoff distance + coords : numpy.ndarray + atom coordinates of shape ``(N, 3)`` for ``N`` atoms. + ``dtype=numpy.float32`` + box : numpy.ndarray + Box dimension of shape (6, ). The dimensions must be + provided in the same format as returned + by :attr:`MDAnalysis.coordinates.base.Timestep.dimensions`: + ``[lx, ly, lz, alpha, beta, gamma]`` (dtype = numpy.float32) + [None] + max_gridsize : int + maximum number of cells in the grid + + Note + ---- + box=``None`` and ``[0., 0., 0., 90., 90., 90]`` + both are used for Non-PBC aware calculations + + """ import MDAnalysis as mda from MDAnalysis.lib.mdamath import triclinic_vectors @@ -442,6 +746,41 @@ cdef class FastNS(object): def search(self, search_coords): + """Search a group of atoms against initialized coordinates + + Creates a new grid with the query atoms and searches + against the initialized coordinates. The search is exclusive + i.e. only the pairs ``(i, j)`` such that ``atom[i]`` from query atoms + and ``atom[j]`` from the initialized set of coordinates is stored as + neighbors. + + PBC-aware/non PBC-aware calculations are automatically enabled during + the instantiation of :class:FastNS. + + Parameters + ---------- + search_coords : numpy.ndarray + Query coordinates of shape ``(N, 3)`` where + ``N`` is the number of queries + + Returns + ------- + results : NSResults object + The object from :class:NSResults + contains ``get_indices``, ``get_distances``. + ``get_pairs``, ``get_pair_distances`` + + Note + ---- + For non-PBC aware calculations, the current implementation doesn't work + if any of the query coordinates is beyond the specified range of + initialized coordinates in :func:`MDAnalysis.lib.nsgrid.FastNS`. + + SeeAlso + ------- + MDAnalysis.lib.nsgrid.NSResults + + """ cdef ns_int i, j, size_search cdef ns_int d, m cdef ns_int current_beadid, bid @@ -520,6 +859,25 @@ cdef class FastNS(object): def self_search(self): + """Searches all the pairs within the initialized coordinates + + All the pairs among the initialized coordinates are registered + in hald the time. Although the algorithm is still the same, but + the distance checks can be reduced to half in this particular case + as every pair need not be evaluated twice. + + Returns + ------- + results : NSResults object + The object from :class:NSResults + contains ``get_indices``, ``get_distances``. + ``get_pairs``, ``get_pair_distances`` + + SeeAlso + ------- + MDAnalysis.lib.nsgrid.NSResults + + """ cdef ns_int i, j, size_search cdef ns_int d, m cdef ns_int current_beadid, bid From c817b3945253b75be065c3b9fef6c3c5900e5e2b Mon Sep 17 00:00:00 2001 From: ayush Date: Fri, 3 Aug 2018 19:21:13 -0700 Subject: [PATCH 42/47] returning list for get_indices, checked Docs --- package/MDAnalysis/lib/nsgrid.pyx | 139 ++++++++++++------- testsuite/MDAnalysisTests/lib/test_nsgrid.py | 14 +- 2 files changed, 93 insertions(+), 60 deletions(-) diff --git a/package/MDAnalysis/lib/nsgrid.pyx b/package/MDAnalysis/lib/nsgrid.pyx index 85207f9da1b..05214b35a68 100644 --- a/package/MDAnalysis/lib/nsgrid.pyx +++ b/package/MDAnalysis/lib/nsgrid.pyx @@ -21,18 +21,20 @@ # MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations. # J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787 # -# + # cython: cdivision=True # cython: boundscheck=False # cython: initializedcheck=False +# cython: embedsignature=True """ Neighbor search library --- :mod:`MDAnalysis.lib.nsgrid` ======================================================== + About the code ---------------- +-------------- This Neighbor search library is a serialized Cython version greatly inspired by the NS grid search implemented in @@ -45,7 +47,7 @@ was used as reference to write this file. GROMACS 4.x code is released under the GNU Public Licence v2. About the algorithm --------------------- +------------------- The neighbor search implemented here is based on `cell lists `_ which allow computation of pairs [#]_ with a cost of :math:`O(N)`, instead of :math:`O(N^2)`. @@ -61,6 +63,7 @@ greater than the desired search distance, otherwise all of the neighbours might not reflect in the results. + .. [#] a pair correspond to two particles that are considered as neighbors . """ @@ -120,14 +123,16 @@ cdef class PBCBox(object): and to compute PBC-aware distance. The class can also handle non-PBC aware distance evaluations through ``periodic`` argument. - .. warning:: + .. warning:: This class is not meant to be used by end users. - .. warning:: + .. warning:: Even if MD triclinic boxes can be handled by this class, internal optimization is made based on the - assumption that particles are inside a brick-shaped box. When this is not the case, calculated distances are not - warranted to be exact. + assumption that particles are inside a brick-shaped box. When this is not the case, calculated distances are not + warranted to be exact. """ + + cdef cPBCBox_t c_pbcbox cdef bint is_triclinic cdef bint periodic @@ -144,6 +149,7 @@ cdef class PBCBox(object): ``True`` for PBC-aware calculations ``False`` for non PBC aware calculations """ + self.periodic = periodic self.update(box) @@ -154,8 +160,8 @@ cdef class PBCBox(object): PBC-aware distance calculations. The internal box parameters are used to define the brick-shaped box which is eventually used for distance calculations. - """ + """ cdef ns_int i, j cdef real min_hv2, min_ss, tmp @@ -212,7 +218,8 @@ cdef class PBCBox(object): Call to this method is only needed when the MD box is changed as it always called when class is instantiated. - """ + """ + if box.shape[0] != DIM or box.shape[1] != DIM: raise ValueError("Box must be a {} x {} matrix. Got: {} x {})".format( DIM, DIM, box.shape[0], box.shape[1])) @@ -255,7 +262,9 @@ cdef class PBCBox(object): Returns the distance obeying minimum image convention if periodic is set to ``True`` while instantiating the ``PBCBox`` object. + """ + cdef rvec dx self.fast_pbc_dx(a, b, dx) return rvec_norm2(dx) @@ -267,7 +276,9 @@ cdef class PBCBox(object): All the coordinates are brought into an orthogonal box. The box vectors for the brick-shaped box are defined in ``fast_update`` method. + """ + cdef ns_int i, m, d, natoms cdef real[:,::1] bbox_coords @@ -305,20 +316,8 @@ cdef class NSResults(object): instance of this class. All the methods of :class:`FastNS` returns an instance of this class, which can be used to generate the desired results on demand. - - While more details can be found in individual method description, - a brief list of methods and their returns can be summarized as follows - - ``get_pairs`` : All the pairs of particles within the specified cutoff distance. - - ``get_pair_distances`` : Distance corresponding to every pair of neighbors - - ``get_indices`` : An array with arrays of neighbour corresponding to every individual particle. - - ``get_distances`` : Corresponding distance for every neighbor in ``get_indices`` - - - SeeAlso - ------- - MDAnalysis.lib.nsgrid.FastNS - """ + cdef readonly real cutoff cdef ns_int npairs @@ -344,6 +343,7 @@ cdef class NSResults(object): Array with query coordinates. Shape must be ``(M, 3)`` for ``M`` queries. ``dtype`` must be ``numpy.float32`` """ + self.cutoff = cutoff self.coords = coords self.searchcoords = searchcoords @@ -357,6 +357,7 @@ cdef class NSResults(object): This is the primary function used by :class:`FastNS` to save the pair of atoms, which can be considered as neighbors. """ + self.pairs_buffer.push_back(beadid_i) self.pairs_buffer.push_back(beadid_j) self.pair_distances2_buffer.push_back(distance2) @@ -378,6 +379,7 @@ cdef class NSResults(object): pairs of atom indices of neighbors from query and initial atom coordinates of shape ``(N, 2)`` """ + return np.asarray(self.pairs_buffer).reshape(self.npairs, 2) def get_pair_distances(self): @@ -389,11 +391,18 @@ cdef class NSResults(object): ``pairs[i, 0]`` and ``pairs[i, 1]``, where pairs is the array obtained from ``get_pairs()`` - SeeAlso + Returns ------- + distances : numpy.ndarray + distances between pairs of query and initial + atom coordinates of shape ``N`` + + See Also + -------- MDAnalysis.lib.nsgrid.NSResults.get_pairs """ + self.pair_distances_buffer = np.sqrt(self.pair_distances2_buffer) return np.asarray(self.pair_distances_buffer) @@ -403,6 +412,7 @@ cdef class NSResults(object): of the query atoms. """ + cdef ns_int i, beadid_i, beadid_j cdef ns_int idx, nsearch cdef real dist2, dist @@ -439,25 +449,38 @@ cdef class NSResults(object): Returns ------- - indices : np.ndarray + indices : list Indices of neighboring atoms. - Every element i.e. ``indices[i]`` will be an array of - shape ``m`` where m is the number of neighbours of - query atom[i]. ``dtype=numpy.int`` + Every element i.e. ``indices[i]`` will be a list of + size ``m`` where m is the number of neighbours of + query atom[i]. + + .. code-block:: python + + results = NSResults() + indices = results.get_indices() + + ``indices[i]`` will output a list of neighboring + atoms of ``atom[i]`` from query atoms ``atom``. + ``indices[i][j]`` will give the atom-id of initial coordinates + such that ``initial_atom[indices[i][j]]`` is a neighbor of ``atom[i]`` """ + if self.indices_buffer.empty(): self.create_buffers() - return np.ascontiguousarray(self.indices_buffer) + return list(self.indices_buffer) def get_distances(self): """Distance corresponding to individual neighbors of query atom - For every queried atom ``i``, an array of all the distances - from its neighbors can be obtained from ``get_distances()[i]``. - Every ``distance[i, j]`` will correspond - to the distance between atoms whose indices can be obtained - from ``get_indices()[i, j]`` + For every queried atom ``i``, a list of all the distances + from its neighboring atoms can be obtained from ``get_distances()[i]``. + Every ``distance[i][j]`` will correspond + to the distance between atoms ``atom[i]`` from the query + atoms and ``atom[indices[j]]`` from the initialized + set of coordinates, where ``indices`` can be obtained + by ``get_indices()`` Returns ------- @@ -466,14 +489,25 @@ cdef class NSResults(object): shape ``m`` where m is the number of neighbours of query atom[i]. - SeeAlso - ------- + .. code-block:: python + + results = NSResults() + distances = results.get_distances() + + + atoms of ``atom[i]`` and query atoms ``atom``. + ``indices[i][j]`` will give the atom-id of initial coordinates + such that ``initial_atom[indices[i][j]]`` is a neighbor of ``atom[i]`` + + See Also + -------- MDAnalysis.lib.nsgrid.NSResults.get_indices """ + if self.distances_buffer.empty(): self.create_buffers() - return np.ascontiguousarray(self.distances_buffer) + return list(self.distances_buffer) cdef class NSGrid(object): """Constructs a uniform cuboidal grid for a brick-shaped box @@ -498,6 +532,7 @@ cdef class NSGrid(object): :func:`~MDAnalysis.lib.nsgrid.PBCBox.fast_put_atoms_in_bbox` """ + cdef readonly real cutoff # cutoff cdef ns_int size # total cells cdef ns_int ncoords # number of coordinates @@ -525,6 +560,7 @@ cdef class NSGrid(object): force : boolean Optimizes cutoff if set to ``False`` [False] """ + cdef ns_int i cdef ns_int ncellx, ncelly, ncellz, size, nbeadspercell cdef ns_int xi, yi, zi @@ -586,6 +622,7 @@ cdef class NSGrid(object): cdef bint cellid2cellxyz(self, ns_int cellid, ivec cellxyz) nogil: """Finds actual cell position `(x, y, z)` from a cell-id """ + if cellid < 0: return False if cellid >= self.size: @@ -611,6 +648,7 @@ cdef class NSGrid(object): The method fails if any coordinate is outside the brick shaped box. """ + cdef ns_int i, cellindex = -1 cdef ns_int ncoords = coords.shape[0] cdef ns_int[:] beadcounts = np.empty(self.size, dtype=np.int) @@ -656,7 +694,6 @@ cdef class FastNS(object): Minimum image convention is used for distance evaluations if box dimensions are provided. """ - cdef PBCBox box cdef real[:, ::1] coords cdef real[:, ::1] coords_bbox @@ -701,24 +738,30 @@ cdef class FastNS(object): both are used for Non-PBC aware calculations """ + import MDAnalysis as mda from MDAnalysis.lib.mdamath import triclinic_vectors cdef real[:] pseudobox = np.zeros(6, dtype=np.float32) cdef real[DIM] bmax, bmin - cdef ns_int i + cdef ns_int i, j, ncoords + self.periodic = True + self.coords = coords.copy() + ncoords = len(self.coords) if (box is None) or (np.allclose(box[:3], 0.) and box.shape[0] == 6): - bmax = np.max(coords, axis=0) - bmin = np.min(coords, axis=0) + bmax = np.max(self.coords, axis=0) + bmin = np.min(self.coords, axis=0) for i in range(DIM): pseudobox[i] = 1.1*(bmax[i] - bmin[i]) pseudobox[DIM + i] = 90. box = pseudobox # shift the origin - coords -= bmin + for i in range(ncoords): + for j in range(DIM): + self.coords[i][j] -= bmin[j] self.periodic = False @@ -732,9 +775,9 @@ cdef class FastNS(object): if cutoff * cutoff > self.box.c_pbcbox.max_cutoff2: raise ValueError("Cutoff greater than maximum cutoff ({:.3f}) given the PBC") - self.coords = coords.copy() + - self.coords_bbox = self.box.fast_put_atoms_in_bbox(coords) + self.coords_bbox = self.box.fast_put_atoms_in_bbox(self.coords) self.cutoff = cutoff self.max_gridsize = max_gridsize @@ -776,11 +819,12 @@ cdef class FastNS(object): if any of the query coordinates is beyond the specified range of initialized coordinates in :func:`MDAnalysis.lib.nsgrid.FastNS`. - SeeAlso - ------- + See Also + -------- MDAnalysis.lib.nsgrid.NSResults """ + cdef ns_int i, j, size_search cdef ns_int d, m cdef ns_int current_beadid, bid @@ -873,11 +917,12 @@ cdef class FastNS(object): contains ``get_indices``, ``get_distances``. ``get_pairs``, ``get_pair_distances`` - SeeAlso - ------- + See Also + -------- MDAnalysis.lib.nsgrid.NSResults """ + cdef ns_int i, j, size_search cdef ns_int d, m cdef ns_int current_beadid, bid diff --git a/testsuite/MDAnalysisTests/lib/test_nsgrid.py b/testsuite/MDAnalysisTests/lib/test_nsgrid.py index bcb335f59a8..d748808bb28 100644 --- a/testsuite/MDAnalysisTests/lib/test_nsgrid.py +++ b/testsuite/MDAnalysisTests/lib/test_nsgrid.py @@ -189,16 +189,4 @@ def test_nsgrid_selfsearch(box, result): searcher = nsgrid.FastNS(cutoff, points, box=box) searchresults = searcher.self_search() pairs = searchresults.get_pairs() - assert_equal(len(pairs)//2, result) - -def test_contiguous(universe): - ref_id = 13937 - cutoff = 3 - coords = universe.atoms.positions - searcher = nsgrid.FastNS(cutoff, coords, box=None) - searchresults = searcher.search(coords[ref_id][None, :]) - indices = searchresults.get_indices() - distances = searchresults.get_distances() - - assert indices.flags['C_CONTIGUOUS'] == True - assert distances.flags['C_CONTIGUOUS'] == True + assert_equal(len(pairs)//2, result) \ No newline at end of file From 66d94e5cc760e1789b9ce185d5308b2f92f187af Mon Sep 17 00:00:00 2001 From: ayush Date: Fri, 3 Aug 2018 23:00:27 -0700 Subject: [PATCH 43/47] removed parallel libraries --- package/setup.py | 6 ++---- testsuite/MDAnalysisTests/lib/test_pkdtree.py | 1 - 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/package/setup.py b/package/setup.py index 91a539db326..9343ac793b0 100755 --- a/package/setup.py +++ b/package/setup.py @@ -384,10 +384,8 @@ def extensions(config): ['MDAnalysis/lib/nsgrid' + source_suffix], include_dirs=include_dirs, language='c++', - libraries=parallel_libraries, - define_macros=define_macros + parallel_macros, - extra_compile_args=cpp_extra_compile_args + parallel_args, - extra_link_args=parallel_args) + define_macros=define_macros, + extra_compile_args=cpp_extra_compile_args) pre_exts = [libdcd, distances, distances_omp, qcprot, transformation, libmdaxdr, util, encore_utils, ap_clustering, spe_dimred, cutil, augment, nsgrid] diff --git a/testsuite/MDAnalysisTests/lib/test_pkdtree.py b/testsuite/MDAnalysisTests/lib/test_pkdtree.py index a073057cf77..c161949f9d4 100644 --- a/testsuite/MDAnalysisTests/lib/test_pkdtree.py +++ b/testsuite/MDAnalysisTests/lib/test_pkdtree.py @@ -119,7 +119,6 @@ def test_searchpairs(b, radius, result): indices = tree.search_pairs(radius) assert_equal(len(indices), len(result)) - @pytest.mark.parametrize('radius, result', ((0.1, []), (0.3, [[0, 2]]))) def test_ckd_searchpairs_nopbc(radius, result): From e28c70a6167123221f6993e7aec30d37c046822b Mon Sep 17 00:00:00 2001 From: ayush Date: Sun, 5 Aug 2018 18:53:07 -0700 Subject: [PATCH 44/47] solved contiguous issue, fixed tests, and added checks in API --- package/MDAnalysis/lib/nsgrid.pyx | 1974 +++++++++--------- testsuite/MDAnalysisTests/lib/test_nsgrid.py | 30 +- 2 files changed, 1002 insertions(+), 1002 deletions(-) diff --git a/package/MDAnalysis/lib/nsgrid.pyx b/package/MDAnalysis/lib/nsgrid.pyx index 05214b35a68..96ee469efde 100644 --- a/package/MDAnalysis/lib/nsgrid.pyx +++ b/package/MDAnalysis/lib/nsgrid.pyx @@ -1,991 +1,983 @@ -# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; -*- -# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 -# -# MDAnalysis --- https://www.mdanalysis.org -# -# Copyright (C) 2013-2018 Sébastien Buchoux -# Copyright (c) 2018 The MDAnalysis Development Team and contributors -# (see the file AUTHORS for the full list of names) -# -# Released under the GNU Public Licence, v3 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 -# - - -# cython: cdivision=True -# cython: boundscheck=False -# cython: initializedcheck=False -# cython: embedsignature=True - -""" -Neighbor search library --- :mod:`MDAnalysis.lib.nsgrid` -======================================================== - - -About the code --------------- - -This Neighbor search library is a serialized Cython version greatly -inspired by the NS grid search implemented in -`GROMACS `_ . - -GROMACS 4.x code (more precisely `nsgrid.c `_ -and `ns.c `_ ) -was used as reference to write this file. - -GROMACS 4.x code is released under the GNU Public Licence v2. - -About the algorithm -------------------- - -The neighbor search implemented here is based on `cell lists `_ which allow -computation of pairs [#]_ with a cost of :math:`O(N)`, instead of :math:`O(N^2)`. -The basic algorithm is described in Appendix F, Page 552 of -``Understanding Molecular Dynamics: From Algorithm to Applications`` by Frenkel and Smit. - -In brief, the algorithm divides the domain into smaller subdomains called `cells` -and distributes every particle to these cells based on its position. Subsequently, -any distance based query first identifies the corresponding cell position in the -domain followed by distance evaluations within the identified cell and -neighboring cells only. Care must be taken to ensure that `cellsize` is -greater than the desired search distance, otherwise all of the neighbours might -not reflect in the results. - - - -.. [#] a pair correspond to two particles that are considered as neighbors . - -""" - - -# Preprocessor DEFs -DEF DIM = 3 -DEF XX = 0 -DEF YY = 1 -DEF ZZ = 2 - -DEF EPSILON = 1e-5 - - -# Used to handle memory allocation -from cpython.mem cimport PyMem_Malloc, PyMem_Realloc, PyMem_Free -from libc.math cimport sqrt -import numpy as np -cimport numpy as np -from libcpp.vector cimport vector - - -ctypedef np.int_t ns_int -ctypedef np.float32_t real -ctypedef real rvec[DIM] -ctypedef ns_int ivec[DIM] -ctypedef real matrix[DIM][DIM] - -ctypedef vector[ns_int] intvec -ctypedef vector[real] realvec - -# Useful Functions -cdef real rvec_norm2(const rvec a) nogil: - return a[XX]*a[XX]+a[YY]*a[YY]+a[ZZ]*a[ZZ] - -cdef void rvec_clear(rvec a) nogil: - a[XX]=0.0 - a[YY]=0.0 - a[ZZ]=0.0 - -############################### -# Utility class to handle PBC # -############################### -cdef struct cPBCBox_t: - matrix box - rvec fbox_diag - rvec hbox_diag - rvec mhbox_diag - real max_cutoff2 - - -# Class to handle PBC calculations -cdef class PBCBox(object): - """ - Cython implementation of `PBC-related `_ operations. - This class is used by classes :class:`FastNS` and :class:`NSGrid` to put all particles inside a brick-shaped box - and to compute PBC-aware distance. The class can also handle non-PBC aware distance evaluations through - ``periodic`` argument. - - .. warning:: - This class is not meant to be used by end users. - - .. warning:: - Even if MD triclinic boxes can be handled by this class, internal optimization is made based on the - assumption that particles are inside a brick-shaped box. When this is not the case, calculated distances are not - warranted to be exact. - """ - - - cdef cPBCBox_t c_pbcbox - cdef bint is_triclinic - cdef bint periodic - - def __init__(self, real[:,::1] box, bint periodic): - """ - Parameters - ---------- - box : numpy.ndarray - box vectors of shape ``(3, 3)`` or - as returned by ``MDAnalysis.lib.mdamath.triclinic_vectors`` - ``dtype`` must be ``numpy.float32`` - periodic : boolean - ``True`` for PBC-aware calculations - ``False`` for non PBC aware calculations - """ - - self.periodic = periodic - self.update(box) - - - cdef void fast_update(self, real[:,::1] box) nogil: - """ - Updates the internal box parameters for - PBC-aware distance calculations. The internal - box parameters are used to define the brick-shaped - box which is eventually used for distance calculations. - - """ - cdef ns_int i, j - cdef real min_hv2, min_ss, tmp - - # Update matrix - self.is_triclinic = False - for i in range(DIM): - for j in range(DIM): - self.c_pbcbox.box[i][j] = box[i, j] - - if i != j: - if box[i, j] > EPSILON: - self.is_triclinic = True - - # Update diagonals - for i in range(DIM): - self.c_pbcbox.fbox_diag[i] = box[i, i] - self.c_pbcbox.hbox_diag[i] = self.c_pbcbox.fbox_diag[i] * 0.5 - self.c_pbcbox.mhbox_diag[i] = - self.c_pbcbox.hbox_diag[i] - - # Update maximum cutoff - - # Physical limitation of the cut-off - # by half the length of the shortest box vector. - min_hv2 = min(0.25 * rvec_norm2(&box[XX, XX]), 0.25 * rvec_norm2(&box[YY, XX])) - min_hv2 = min(min_hv2, 0.25 * rvec_norm2(&box[ZZ, XX])) - - # Limitation to the smallest diagonal element due to optimizations: - # checking only linear combinations of single box-vectors (2 in x) - # in the grid search and pbc_dx is a lot faster - # than checking all possible combinations. - tmp = box[YY, YY] - if box[ZZ, YY] < 0: - tmp -= box[ZZ, YY] - else: - tmp += box[ZZ, YY] - - min_ss = min(box[XX, XX], min(tmp, box[ZZ, ZZ])) - self.c_pbcbox.max_cutoff2 = min(min_hv2, min_ss * min_ss) - - def update(self, real[:,::1] box): - """ - Updates internal MD box representation and parameters used for calculations. - - - Parameters - ---------- - box : numpy.ndarray - Describes the MD box vectors as returned by - :func:`MDAnalysis.lib.mdamath.triclinic_vectors`. - `dtype` must be :class:`numpy.float32` - - Note - ---- - Call to this method is only needed when the MD box is changed - as it always called when class is instantiated. - - """ - - if box.shape[0] != DIM or box.shape[1] != DIM: - raise ValueError("Box must be a {} x {} matrix. Got: {} x {})".format( - DIM, DIM, box.shape[0], box.shape[1])) - if (box[XX, XX] == 0) or (box[YY, YY] == 0) or (box[ZZ, ZZ] == 0): - raise ValueError("Box does not correspond to PBC=xyz") - self.fast_update(box) - - - cdef void fast_pbc_dx(self, rvec ref, rvec other, rvec dx) nogil: - """Dislacement between two points for both - PBC and non-PBC conditions - - Modifies the displacement vector between two points based - on the minimum image convention for PBC aware calculations. - - For non-PBC aware distance evaluations, calculates the - displacement vector without any modifications - """ - - cdef ns_int i, j - - for i in range(DIM): - dx[i] = other[i] - ref[i] - - if self.periodic: - for i in range (DIM-1, -1, -1): - while dx[i] > self.c_pbcbox.hbox_diag[i]: - for j in range (i, -1, -1): - dx[j] -= self.c_pbcbox.box[i][j] - - while dx[i] <= self.c_pbcbox.mhbox_diag[i]: - for j in range (i, -1, -1): - dx[j] += self.c_pbcbox.box[i][j] - - - cdef real fast_distance2(self, rvec a, rvec b) nogil: - """Distance calculation between two points - for both PBC and non-PBC aware calculations - - Returns the distance obeying minimum - image convention if periodic is set to ``True`` while - instantiating the ``PBCBox`` object. - - """ - - cdef rvec dx - self.fast_pbc_dx(a, b, dx) - return rvec_norm2(dx) - - - cdef real[:, ::1]fast_put_atoms_in_bbox(self, real[:,::1] coords) nogil: - """Shifts all ``coords`` to an orthogonal brick shaped box - - All the coordinates are brought into an orthogonal - box. The box vectors for the brick-shaped box - are defined in ``fast_update`` method. - - """ - - cdef ns_int i, m, d, natoms - cdef real[:,::1] bbox_coords - - natoms = coords.shape[0] - with gil: - bbox_coords = coords.copy() - - if self.periodic: - if self.is_triclinic: - for i in range(natoms): - for m in range(DIM - 1, -1, -1): - while bbox_coords[i, m] < 0: - for d in range(m+1): - bbox_coords[i, d] += self.c_pbcbox.box[m][d] - while bbox_coords[i, m] >= self.c_pbcbox.box[m][m]: - for d in range(m+1): - bbox_coords[i, d] -= self.c_pbcbox.box[m][d] - else: - for i in range(natoms): - for m in range(DIM): - while bbox_coords[i, m] < 0: - bbox_coords[i, m] += self.c_pbcbox.box[m][m] - while bbox_coords[i, m] >= self.c_pbcbox.box[m][m]: - bbox_coords[i, m] -= self.c_pbcbox.box[m][m] - return bbox_coords - -######################### -# Neighbor Search Stuff # -######################### - -cdef class NSResults(object): - """Class to store the results - - All the required outputs from :class:`FastNS` is stored in the - instance of this class. All the methods of :class:`FastNS` returns - an instance of this class, which can be used to generate the desired - results on demand. - """ - - cdef readonly real cutoff - cdef ns_int npairs - - cdef real[:, ::1] coords # shape: size, DIM - cdef real[:, ::1] searchcoords - - cdef vector[intvec] indices_buffer - cdef vector[realvec] distances_buffer - cdef vector[ns_int] pairs_buffer - cdef vector[real] pair_distances_buffer - cdef vector[real] pair_distances2_buffer - - def __init__(self, real cutoff, real[:, ::1]coords, real[:, ::1]searchcoords): - """ - Parameters - ---------- - cutoff : float - Specified cutoff distance - coords : numpy.ndarray - Array with coordinates of atoms of shape ``(N, 3)`` for - ``N`` particles. ``dtype`` must be ``numpy.float32`` - searchcoords : numpy.ndarray - Array with query coordinates. Shape must be ``(M, 3)`` - for ``M`` queries. ``dtype`` must be ``numpy.float32`` - """ - - self.cutoff = cutoff - self.coords = coords - self.searchcoords = searchcoords - - self.npairs = 0 - - cdef void add_neighbors(self, ns_int beadid_i, ns_int beadid_j, real distance2) nogil: - """Internal function to add pairs and distances to buffers - - The buffers populated using this method are used by other methods of this class. - This is the primary function used by :class:`FastNS` to save the pair of atoms, - which can be considered as neighbors. - """ - - self.pairs_buffer.push_back(beadid_i) - self.pairs_buffer.push_back(beadid_j) - self.pair_distances2_buffer.push_back(distance2) - self.npairs += 1 - - def get_pairs(self): - """Returns all the pairs within the desired cutoff distance - - Returns an array of shape ``(N, 2)``, where N is the number of pairs - between ``reference`` and ``configuration`` within the specified distance. - For every pair ``(i, j)``, ``reference[i]`` and ``configuration[j]`` are - atom positions such that ``reference`` is the position of query - atoms while ``configuration`` coontains the position of group of - atoms used to search against the query atoms. - - Returns - ------- - pairs : numpy.ndarray - pairs of atom indices of neighbors from query - and initial atom coordinates of shape ``(N, 2)`` - """ - - return np.asarray(self.pairs_buffer).reshape(self.npairs, 2) - - def get_pair_distances(self): - """Returns all the distances corresponding to each pair of neighbors - - Returns an array of shape ``N`` where N is the number of pairs - among the query atoms and initial atoms within a specified distance. - Every element ``[i]`` corresponds to the distance between - ``pairs[i, 0]`` and ``pairs[i, 1]``, where pairs is the array - obtained from ``get_pairs()`` - - Returns - ------- - distances : numpy.ndarray - distances between pairs of query and initial - atom coordinates of shape ``N`` - - See Also - -------- - MDAnalysis.lib.nsgrid.NSResults.get_pairs - - """ - - self.pair_distances_buffer = np.sqrt(self.pair_distances2_buffer) - return np.asarray(self.pair_distances_buffer) - - cdef void create_buffers(self) nogil: - """ - Creates buffers to get individual neighbour list and distances - of the query atoms. - - """ - - cdef ns_int i, beadid_i, beadid_j - cdef ns_int idx, nsearch - cdef real dist2, dist - - nsearch = len(self.searchcoords) - - self.indices_buffer = vector[intvec]() - self.distances_buffer = vector[realvec]() - - # initialize rows corresponding to search - for i in range(nsearch): - self.indices_buffer.push_back(intvec()) - self.distances_buffer.push_back(realvec()) - - - for i in range(0, 2*self.npairs, 2): - beadid_i = self.pairs_buffer[i] - beadid_j = self.pairs_buffer[i + 1] - - dist2 = self.pair_distances2_buffer[i//2] - - self.indices_buffer[beadid_i].push_back(beadid_j) - - dist = sqrt(dist2) - - self.distances_buffer[beadid_i].push_back(dist) - - - def get_indices(self): - """Individual neighbours of query atom - - For every queried atom ``i``, an array of all its neighbors - indices can be obtained from ``get_indices()[i]`` - - Returns - ------- - indices : list - Indices of neighboring atoms. - Every element i.e. ``indices[i]`` will be a list of - size ``m`` where m is the number of neighbours of - query atom[i]. - - .. code-block:: python - - results = NSResults() - indices = results.get_indices() - - ``indices[i]`` will output a list of neighboring - atoms of ``atom[i]`` from query atoms ``atom``. - ``indices[i][j]`` will give the atom-id of initial coordinates - such that ``initial_atom[indices[i][j]]`` is a neighbor of ``atom[i]`` - - """ - - if self.indices_buffer.empty(): - self.create_buffers() - return list(self.indices_buffer) - - def get_distances(self): - """Distance corresponding to individual neighbors of query atom - - For every queried atom ``i``, a list of all the distances - from its neighboring atoms can be obtained from ``get_distances()[i]``. - Every ``distance[i][j]`` will correspond - to the distance between atoms ``atom[i]`` from the query - atoms and ``atom[indices[j]]`` from the initialized - set of coordinates, where ``indices`` can be obtained - by ``get_indices()`` - - Returns - ------- - distances : np.ndarray - Every element i.e. ``distances[i]`` will be an array of - shape ``m`` where m is the number of neighbours of - query atom[i]. - - .. code-block:: python - - results = NSResults() - distances = results.get_distances() - - - atoms of ``atom[i]`` and query atoms ``atom``. - ``indices[i][j]`` will give the atom-id of initial coordinates - such that ``initial_atom[indices[i][j]]`` is a neighbor of ``atom[i]`` - - See Also - -------- - MDAnalysis.lib.nsgrid.NSResults.get_indices - - """ - - if self.distances_buffer.empty(): - self.create_buffers() - return list(self.distances_buffer) - -cdef class NSGrid(object): - """Constructs a uniform cuboidal grid for a brick-shaped box - - This class uses :class:`PBCBox` to define the brick shaped box. - It is essential to initialize the box with :class:`PBCBox` - inorder to form the grid. - - The domain is subdivided into number of cells based on the desired search - radius. Ideally cellsize should be equal to the search radius, but small - search radius leads to large cell-list data strucutres. - An optimization of cutoff is imposed to limit the size of data - structure such that the cellsize is always greater than or - equal to cutoff distance. - - Note - ---- - This class assumes that all the coordinates are already - inside the brick shaped box. Care must be taken to ensure - all the particles are within the brick shaped box as - defined by :class:`PBCBox`. This can be ensured by using - :func:`~MDAnalysis.lib.nsgrid.PBCBox.fast_put_atoms_in_bbox` - - """ - - cdef readonly real cutoff # cutoff - cdef ns_int size # total cells - cdef ns_int ncoords # number of coordinates - cdef ns_int[DIM] ncells # individual cells in every dimension - cdef ns_int[DIM] cell_offsets # Cell Multipliers - cdef real[DIM] cellsize # cell size in every dimension - cdef ns_int nbeads_per_cell # maximum beads - cdef ns_int *nbeads # size (Number of beads in every cell) - cdef ns_int *beadids # size * nbeads_per_cell (Beadids in every cell) - cdef ns_int *cellids # ncoords (Cell occupation id for every atom) - cdef bint force # To negate the effects of optimized cutoff - - def __init__(self, ncoords, cutoff, PBCBox box, max_size, force=False): - """ - Parameters - ---------- - ncoords : int - Number of coordinates to fill inside the brick shaped box - cutoff : float - Desired cutoff radius - box : PBCBox - Instance of :class:`PBCBox` - max_size : int - Maximum total number of cells - force : boolean - Optimizes cutoff if set to ``False`` [False] - """ - - cdef ns_int i - cdef ns_int ncellx, ncelly, ncellz, size, nbeadspercell - cdef ns_int xi, yi, zi - cdef real bbox_vol - - - self.ncoords = ncoords - - # Calculate best cutoff - self.cutoff = cutoff - if not force: - bbox_vol = box.c_pbcbox.box[XX][XX] * box.c_pbcbox.box[YY][YY] * box.c_pbcbox.box[YY][YY] - size = bbox_vol/cutoff**3 - nbeadspercell = ncoords/size - while bbox_vol/self.cutoff**3 > max_size: - self.cutoff *= 1.2 - - for i in range(DIM): - self.ncells[i] = (box.c_pbcbox.box[i][i] / self.cutoff) - self.cellsize[i] = box.c_pbcbox.box[i][i] / self.ncells[i] - self.size = self.ncells[XX] * self.ncells[YY] * self.ncells[ZZ] - - - self.cell_offsets[XX] = 0 - self.cell_offsets[YY] = self.ncells[XX] - self.cell_offsets[ZZ] = self.ncells[XX] * self.ncells[YY] - - # Allocate memory - self.nbeads = PyMem_Malloc(sizeof(ns_int) * self.size) - if not self.nbeads: - raise MemoryError("Could not allocate memory from NSGrid.nbeads ({} bits requested)".format(sizeof(ns_int) * self.size)) - self.beadids = NULL - self.cellids = PyMem_Malloc(sizeof(ns_int) * self.ncoords) - if not self.cellids: - raise MemoryError("Could not allocate memory from NSGrid.cellids ({} bits requested)".format(sizeof(ns_int) * self.ncoords)) - self.nbeads_per_cell = 0 - - for i in range(self.size): - self.nbeads[i] = 0 - - def __dealloc__(self): - PyMem_Free(self.nbeads) - PyMem_Free(self.beadids) - PyMem_Free(self.cellids) - - cdef ns_int coord2cellid(self, rvec coord) nogil: - """Finds the cell-id for the given coordinate inside the brick shaped box - - Note - ---- - Assumes the coordinate is already inside the brick shaped box. - Return wrong cell-id if this is not the case - """ - return (coord[ZZ] / self.cellsize[ZZ]) * (self.cell_offsets[ZZ]) +\ - (coord[YY] / self.cellsize[YY]) * self.cell_offsets[YY] + \ - (coord[XX] / self.cellsize[XX]) - - - cdef bint cellid2cellxyz(self, ns_int cellid, ivec cellxyz) nogil: - """Finds actual cell position `(x, y, z)` from a cell-id - """ - - if cellid < 0: - return False - if cellid >= self.size: - return False - - cellxyz[ZZ] = (cellid / self.cell_offsets[ZZ]) - cellid -= cellxyz[ZZ] * self.cell_offsets[ZZ] - - cellxyz[YY] = (cellid / self.cell_offsets[YY]) - cellxyz[XX] = cellid - cellxyz[YY] * self.cell_offsets[YY] - - return True - - cdef fill_grid(self, real[:, ::1] coords): - """Sorts atoms into cells based on their position in the brick shaped box - - Every atom inside the brick shaped box is assigned a - cell-id based on its position. Another list ``beadids`` - sort the atom-ids in each cell. - - Note - ---- - The method fails if any coordinate is outside the brick shaped box. - - """ - - cdef ns_int i, cellindex = -1 - cdef ns_int ncoords = coords.shape[0] - cdef ns_int[:] beadcounts = np.empty(self.size, dtype=np.int) - - with nogil: - # Initialize buffers - for i in range(self.size): - beadcounts[i] = 0 - - # First loop: find cellindex for each bead - for i in range(ncoords): - cellindex = self.coord2cellid(&coords[i, 0]) - - self.nbeads[cellindex] += 1 - self.cellids[i] = cellindex - - if self.nbeads[cellindex] > self.nbeads_per_cell: - self.nbeads_per_cell = self.nbeads[cellindex] - - # Allocate memory - self.beadids = PyMem_Malloc(sizeof(ns_int) * self.size * self.nbeads_per_cell) #np.empty((self.size, nbeads_max), dtype=np.int) - if not self.beadids: - raise MemoryError("Could not allocate memory for NSGrid.beadids ({} bits requested)".format(sizeof(ns_int) * self.size * self.nbeads_per_cell)) - - with nogil: - # Second loop: fill grid - for i in range(ncoords): - - # Add bead to grid cell - cellindex = self.cellids[i] - self.beadids[cellindex * self.nbeads_per_cell + beadcounts[cellindex]] = i - beadcounts[cellindex] += 1 - - - -cdef class FastNS(object): - """Grid based search between two group of atoms - - Instantiates a class object which uses :class:`PBCBox` and - :class:`NSGrid` to construct a cuboidal - grid in an orthogonal brick shaped box. - - Minimum image convention is used for distance evaluations - if box dimensions are provided. - """ - cdef PBCBox box - cdef real[:, ::1] coords - cdef real[:, ::1] coords_bbox - cdef readonly real cutoff - cdef NSGrid grid - cdef ns_int max_gridsize - cdef bint periodic - - - def __init__(self, cutoff, coords, box=None, max_gridsize=5000): - """ - Initialize the grid and sort the coordinates in respective - cells by shifting the coordinates in a brick shaped box. - The brick shaped box is defined by :class:`PBCBox` - and cuboidal grid is initialize by :class:`NSGrid`. - If box is supplied, periodic shifts along box vectors are used - to contain all the coordinates inside the brick shaped box. - If box is not supplied, the range of coordinates i.e. - ``[xmax, ymax, zmax] - [xmin, ymin, zmin]`` is used to construct - a pseudo box. Subsequently, the origin is also shifted - to the ``[xmin, ymin, zmin]``. - - Parameters - ---------- - cutoff : float - Desired cutoff distance - coords : numpy.ndarray - atom coordinates of shape ``(N, 3)`` for ``N`` atoms. - ``dtype=numpy.float32`` - box : numpy.ndarray - Box dimension of shape (6, ). The dimensions must be - provided in the same format as returned - by :attr:`MDAnalysis.coordinates.base.Timestep.dimensions`: - ``[lx, ly, lz, alpha, beta, gamma]`` (dtype = numpy.float32) - [None] - max_gridsize : int - maximum number of cells in the grid - - Note - ---- - box=``None`` and ``[0., 0., 0., 90., 90., 90]`` - both are used for Non-PBC aware calculations - - """ - - import MDAnalysis as mda - from MDAnalysis.lib.mdamath import triclinic_vectors - - cdef real[:] pseudobox = np.zeros(6, dtype=np.float32) - cdef real[DIM] bmax, bmin - cdef ns_int i, j, ncoords - - self.periodic = True - self.coords = coords.copy() - ncoords = len(self.coords) - - - if (box is None) or (np.allclose(box[:3], 0.) and box.shape[0] == 6): - bmax = np.max(self.coords, axis=0) - bmin = np.min(self.coords, axis=0) - for i in range(DIM): - pseudobox[i] = 1.1*(bmax[i] - bmin[i]) - pseudobox[DIM + i] = 90. - box = pseudobox - # shift the origin - for i in range(ncoords): - for j in range(DIM): - self.coords[i][j] -= bmin[j] - self.periodic = False - - - if box.shape != (3,3): - box = triclinic_vectors(box) - - self.box = PBCBox(box, self.periodic) - - if cutoff < 0: - raise ValueError("Cutoff must be positive!") - if cutoff * cutoff > self.box.c_pbcbox.max_cutoff2: - raise ValueError("Cutoff greater than maximum cutoff ({:.3f}) given the PBC") - - - - self.coords_bbox = self.box.fast_put_atoms_in_bbox(self.coords) - - self.cutoff = cutoff - self.max_gridsize = max_gridsize - # Note that self.cutoff might be different from self.grid.cutoff - self.grid = NSGrid(self.coords_bbox.shape[0], self.cutoff, self.box, self.max_gridsize) - - self.grid.fill_grid(self.coords_bbox) - - - - def search(self, search_coords): - """Search a group of atoms against initialized coordinates - - Creates a new grid with the query atoms and searches - against the initialized coordinates. The search is exclusive - i.e. only the pairs ``(i, j)`` such that ``atom[i]`` from query atoms - and ``atom[j]`` from the initialized set of coordinates is stored as - neighbors. - - PBC-aware/non PBC-aware calculations are automatically enabled during - the instantiation of :class:FastNS. - - Parameters - ---------- - search_coords : numpy.ndarray - Query coordinates of shape ``(N, 3)`` where - ``N`` is the number of queries - - Returns - ------- - results : NSResults object - The object from :class:NSResults - contains ``get_indices``, ``get_distances``. - ``get_pairs``, ``get_pair_distances`` - - Note - ---- - For non-PBC aware calculations, the current implementation doesn't work - if any of the query coordinates is beyond the specified range of - initialized coordinates in :func:`MDAnalysis.lib.nsgrid.FastNS`. - - See Also - -------- - MDAnalysis.lib.nsgrid.NSResults - - """ - - cdef ns_int i, j, size_search - cdef ns_int d, m - cdef ns_int current_beadid, bid - cdef ns_int cellindex, cellindex_probe - cdef ns_int xi, yi, zi - - cdef NSResults results - - cdef real d2 - cdef rvec probe - - cdef real[:, ::1] searchcoords - cdef real[:, ::1] searchcoords_bbox - cdef NSGrid searchgrid - cdef bint check - - - cdef real cutoff2 = self.cutoff * self.cutoff - cdef ns_int npairs = 0 - - # Generate another grid to search - searchcoords = np.asarray(search_coords, dtype=np.float32) - searchcoords_bbox = self.box.fast_put_atoms_in_bbox(searchcoords) - searchgrid = NSGrid(searchcoords_bbox.shape[0], self.grid.cutoff, self.box, self.max_gridsize, force=True) - searchgrid.fill_grid(searchcoords_bbox) - - - size_search = searchcoords.shape[0] - - results = NSResults(self.cutoff, self.coords, searchcoords) - - with nogil: - for i in range(size_search): - # Start with first search coordinate - current_beadid = i - # find the cellindex of the coordinate - cellindex = searchgrid.cellids[current_beadid] - for xi in range(DIM): - for yi in range(DIM): - for zi in range(DIM): - check = True - #Probe the search coordinates in a brick shaped box - probe[XX] = searchcoords_bbox[current_beadid, XX] + (xi - 1) * searchgrid.cellsize[XX] - probe[YY] = searchcoords_bbox[current_beadid, YY] + (yi - 1) * searchgrid.cellsize[YY] - probe[ZZ] = searchcoords_bbox[current_beadid, ZZ] + (zi - 1) * searchgrid.cellsize[ZZ] - # Make sure the probe coordinates is inside the brick-shaped box - if self.periodic: - for m in range(DIM - 1, -1, -1): - while probe[m] < 0: - for d in range(m+1): - probe[d] += self.box.c_pbcbox.box[m][d] - while probe[m] >= self.box.c_pbcbox.box[m][m]: - for d in range(m+1): - probe[d] -= self.box.c_pbcbox.box[m][d] - else: - for m in range(DIM -1, -1, -1): - if probe[m] < 0: - check = False - break - if probe[m] > self.box.c_pbcbox.box[m][m]: - check = False - break - if not check: - continue - # Get the cell index corresponding to the probe - cellindex_probe = self.grid.coord2cellid(probe) - #for this cellindex search in grid - for j in range(self.grid.nbeads[cellindex_probe]): - bid = self.grid.beadids[cellindex_probe * self.grid.nbeads_per_cell + j] - #find distance between search coords[i] and coords[bid] - d2 = self.box.fast_distance2(&searchcoords_bbox[current_beadid, XX], &self.coords_bbox[bid, XX]) - if d2 < cutoff2 and d2 > EPSILON: - results.add_neighbors(current_beadid, bid, d2) - npairs += 1 - return results - - - def self_search(self): - """Searches all the pairs within the initialized coordinates - - All the pairs among the initialized coordinates are registered - in hald the time. Although the algorithm is still the same, but - the distance checks can be reduced to half in this particular case - as every pair need not be evaluated twice. - - Returns - ------- - results : NSResults object - The object from :class:NSResults - contains ``get_indices``, ``get_distances``. - ``get_pairs``, ``get_pair_distances`` - - See Also - -------- - MDAnalysis.lib.nsgrid.NSResults - - """ - - cdef ns_int i, j, size_search - cdef ns_int d, m - cdef ns_int current_beadid, bid - cdef ns_int cellindex, cellindex_probe - cdef ns_int xi, yi, zi - - cdef NSResults results - cdef real d2 - cdef rvec probe - - cdef real cutoff2 = self.cutoff * self.cutoff - cdef ns_int npairs = 0 - cdef bint check - - size_search = self.coords.shape[0] - - results = NSResults(self.cutoff, self.coords, self.coords) - - with nogil: - for i in range(size_search): - # Start with first search coordinate - current_beadid = i - # find the cellindex of the coordinate - cellindex = self.grid.cellids[current_beadid] - for xi in range(DIM): - for yi in range(DIM): - for zi in range(DIM): - check = True - # Calculate and/or reinitialize shifted coordinates - #Probe the search coordinates in a brick shaped box - probe[XX] = self.coords_bbox[current_beadid, XX] + (xi - 1) * self.grid.cellsize[XX] - probe[YY] = self.coords_bbox[current_beadid, YY] + (yi - 1) * self.grid.cellsize[XX] - probe[ZZ] = self.coords_bbox[current_beadid, ZZ] + (zi - 1) * self.grid.cellsize[XX] - # Make sure the shifted coordinates is inside the brick-shaped box - if self.periodic: - for m in range(DIM - 1, -1, -1): - while probe[m] < 0: - for d in range(m+1): - probe[d] += self.box.c_pbcbox.box[m][d] - while probe[m] >= self.box.c_pbcbox.box[m][m]: - for d in range(m+1): - probe[d] -= self.box.c_pbcbox.box[m][d] - else: - for m in range(DIM -1, -1, -1): - if probe[m] < 0: - check = False - break - elif probe[m] > self.box.c_pbcbox.box[m][m]: - check = False - break - if not check: - continue - # Get the cell index corresponding to the probe - cellindex_probe = self.grid.coord2cellid(probe) - #for this cellindex search in grid - for j in range(self.grid.nbeads[cellindex_probe]): - bid = self.grid.beadids[cellindex_probe * self.grid.nbeads_per_cell + j] - if bid < current_beadid: - continue - #find distance between search coords[i] and coords[bid] - d2 = self.box.fast_distance2(&self.coords_bbox[current_beadid, XX], &self.coords_bbox[bid, XX]) - if d2 < cutoff2 and d2 > EPSILON: - results.add_neighbors(current_beadid, bid, d2) - results.add_neighbors(bid, current_beadid, d2) - npairs += 1 - return results +# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; -*- +# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 +# +# MDAnalysis --- https://www.mdanalysis.org +# +# Copyright (C) 2013-2018 Sébastien Buchoux +# Copyright (c) 2018 The MDAnalysis Development Team and contributors +# (see the file AUTHORS for the full list of names) +# +# Released under the GNU Public Licence, v3 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 +# + + +# cython: cdivision=True +# cython: boundscheck=False +# cython: initializedcheck=False +# cython: embedsignature=True + +""" +Neighbor search library --- :mod:`MDAnalysis.lib.nsgrid` +======================================================== + + +About the code +-------------- + +This Neighbor search library is a serialized Cython version greatly +inspired by the NS grid search implemented in +`GROMACS `_ . + +GROMACS 4.x code (more precisely +`nsgrid.c `_ +and `ns.c `_ ) +was used as reference to write this file. + +GROMACS 4.x code is released under the GNU Public Licence v2. + +About the algorithm +------------------- + +The neighbor search implemented here is based on +`cell lists `_ which allow +computation of pairs [#]_ with a cost of :math:`O(N)`, instead +of :math:`O(N^2)`. The basic algorithm is described in +Appendix F, Page 552 of +``Understanding Molecular Dynamics: From Algorithm to Applications`` by Frenkel and Smit. + +In brief, the algorithm divides the domain into smaller subdomains called `cells` +and distributes every particle to these cells based on its position. Subsequently, +any distance based query first identifies the corresponding cell position in the +domain followed by distance evaluations within the identified cell and +neighboring cells only. Care must be taken to ensure that `cellsize` is +greater than the desired search distance, otherwise all of the neighbours might +not reflect in the results. + + +.. [#] a pair correspond to two particles that are considered as neighbors . + + +.. versionadded:: 0.19.0 +""" + +from MDAnalysis.lib.distances import _check_array +# Used to handle memory allocation +from cpython.mem cimport PyMem_Malloc, PyMem_Realloc, PyMem_Free +from libc.math cimport sqrt +import numpy as np +from libcpp.vector cimport vector +cimport numpy as np + +# Preprocessor DEFs +DEF DIM = 3 +DEF XX = 0 +DEF YY = 1 +DEF ZZ = 2 +DEF EPSILON = 1e-5 + +ctypedef np.int_t ns_int +ctypedef np.float32_t real +ctypedef real rvec[DIM] +ctypedef ns_int ivec[DIM] +ctypedef real matrix[DIM][DIM] + +ctypedef vector[ns_int] intvec +ctypedef vector[real] realvec + +# Useful Functions +cdef real rvec_norm2(const rvec a) nogil: + return a[XX]*a[XX] + a[YY]*a[YY] + a[ZZ]*a[ZZ] + +cdef void rvec_clear(rvec a) nogil: + a[XX] = 0.0 + a[YY] = 0.0 + a[ZZ] = 0.0 + +############################### +# Utility class to handle PBC # +############################### +cdef struct cPBCBox_t: + matrix box + rvec fbox_diag + rvec hbox_diag + rvec mhbox_diag + real max_cutoff2 + + +# Class to handle PBC calculations +cdef class PBCBox(object): + """ + Cython implementation of + `PBC-related `_ + operations. This class is used by classes :class:`FastNS` + and :class:`NSGrid` to put all particles inside a brick-shaped box + and to compute PBC-aware distance. The class can also handle + non-PBC aware distance evaluations through ``periodic`` argument. + + .. warning:: + This class is not meant to be used by end users. + + .. warning:: + Even if MD triclinic boxes can be handled by this class, internal optimization is made based on the + assumption that particles are inside a brick-shaped box. When this is not the case, calculated distances are not + warranted to be exact. + """ + + cdef cPBCBox_t c_pbcbox + cdef bint is_triclinic + cdef bint periodic + + def __init__(self, real[:, ::1] box, bint periodic): + """ + Parameters + ---------- + box : numpy.ndarray + box vectors of shape ``(3, 3)`` or + as returned by ``MDAnalysis.lib.mdamath.triclinic_vectors`` + ``dtype`` must be ``numpy.float32`` + periodic : boolean + ``True`` for PBC-aware calculations + ``False`` for non PBC aware calculations + """ + + self.periodic = periodic + self.update(box) + + cdef void fast_update(self, real[:, ::1] box) nogil: + """ + Updates the internal box parameters for + PBC-aware distance calculations. The internal + box parameters are used to define the brick-shaped + box which is eventually used for distance calculations. + + """ + cdef ns_int i, j + cdef real min_hv2, min_ss, tmp + + # Update matrix + self.is_triclinic = False + for i in range(DIM): + for j in range(DIM): + self.c_pbcbox.box[i][j] = box[i, j] + + if i != j: + if box[i, j] > EPSILON: + self.is_triclinic = True + + # Update diagonals + for i in range(DIM): + self.c_pbcbox.fbox_diag[i] = box[i, i] + self.c_pbcbox.hbox_diag[i] = self.c_pbcbox.fbox_diag[i] * 0.5 + self.c_pbcbox.mhbox_diag[i] = - self.c_pbcbox.hbox_diag[i] + + # Update maximum cutoff + + # Physical limitation of the cut-off + # by half the length of the shortest box vector. + min_hv2 = min(0.25 * rvec_norm2(&box[XX, XX]), 0.25 * rvec_norm2(&box[YY, XX])) + min_hv2 = min(min_hv2, 0.25 * rvec_norm2(&box[ZZ, XX])) + + # Limitation to the smallest diagonal element due to optimizations: + # checking only linear combinations of single box-vectors (2 in x) + # in the grid search and pbc_dx is a lot faster + # than checking all possible combinations. + tmp = box[YY, YY] + if box[ZZ, YY] < 0: + tmp -= box[ZZ, YY] + else: + tmp += box[ZZ, YY] + + min_ss = min(box[XX, XX], min(tmp, box[ZZ, ZZ])) + self.c_pbcbox.max_cutoff2 = min(min_hv2, min_ss * min_ss) + + def update(self, real[:, ::1] box): + """ + Updates internal MD box representation and parameters used for calculations. + + Parameters + ---------- + box : numpy.ndarray + Describes the MD box vectors as returned by + :func:`MDAnalysis.lib.mdamath.triclinic_vectors`. + `dtype` must be :class:`numpy.float32` + + Note + ---- + Call to this method is only needed when the MD box is changed + as it always called when class is instantiated. + + """ + + if box.shape[0] != DIM or box.shape[1] != DIM: + raise ValueError("Box must be a {} x {} matrix. Got: {} x {})".format( + DIM, DIM, box.shape[0], box.shape[1])) + if (box[XX, XX] == 0) or (box[YY, YY] == 0) or (box[ZZ, ZZ] == 0): + raise ValueError("Box does not correspond to PBC=xyz") + self.fast_update(box) + + cdef void fast_pbc_dx(self, rvec ref, rvec other, rvec dx) nogil: + """Dislacement between two points for both + PBC and non-PBC conditions + + Modifies the displacement vector between two points based + on the minimum image convention for PBC aware calculations. + + For non-PBC aware distance evaluations, calculates the + displacement vector without any modifications + """ + + cdef ns_int i, j + + for i in range(DIM): + dx[i] = other[i] - ref[i] + + if self.periodic: + for i in range(DIM-1, -1, -1): + while dx[i] > self.c_pbcbox.hbox_diag[i]: + for j in range(i, -1, -1): + dx[j] -= self.c_pbcbox.box[i][j] + + while dx[i] <= self.c_pbcbox.mhbox_diag[i]: + for j in range(i, -1, -1): + dx[j] += self.c_pbcbox.box[i][j] + + cdef real fast_distance2(self, rvec a, rvec b) nogil: + """Distance calculation between two points + for both PBC and non-PBC aware calculations + + Returns the distance obeying minimum + image convention if periodic is set to ``True`` while + instantiating the ``PBCBox`` object. + """ + + cdef rvec dx + self.fast_pbc_dx(a, b, dx) + return rvec_norm2(dx) + + cdef real[:, ::1]fast_put_atoms_in_bbox(self, real[:, ::1] coords) nogil: + """Shifts all ``coords`` to an orthogonal brick shaped box + + All the coordinates are brought into an orthogonal + box. The box vectors for the brick-shaped box + are defined in ``fast_update`` method. + + """ + + cdef ns_int i, m, d, natoms + cdef real[:, ::1] bbox_coords + + natoms = coords.shape[0] + with gil: + bbox_coords = coords.copy() + + if self.periodic: + if self.is_triclinic: + for i in range(natoms): + for m in range(DIM - 1, -1, -1): + while bbox_coords[i, m] < 0: + for d in range(m+1): + bbox_coords[i, d] += self.c_pbcbox.box[m][d] + while bbox_coords[i, m] >= self.c_pbcbox.box[m][m]: + for d in range(m+1): + bbox_coords[i, d] -= self.c_pbcbox.box[m][d] + else: + for i in range(natoms): + for m in range(DIM): + while bbox_coords[i, m] < 0: + bbox_coords[i, m] += self.c_pbcbox.box[m][m] + while bbox_coords[i, m] >= self.c_pbcbox.box[m][m]: + bbox_coords[i, m] -= self.c_pbcbox.box[m][m] + return bbox_coords + +######################### +# Neighbor Search Stuff # +######################### + +cdef class NSResults(object): + """Class to store the results + + All the required outputs from :class:`FastNS` is stored in the + instance of this class. All the methods of :class:`FastNS` returns + an instance of this class, which can be used to generate the desired + results on demand. + """ + + cdef readonly real cutoff + cdef ns_int npairs + + cdef real[:, ::1] coords # shape: size, DIM + cdef real[:, ::1] searchcoords + + cdef vector[intvec] indices_buffer + cdef vector[realvec] distances_buffer + cdef vector[ns_int] pairs_buffer + cdef vector[real] pair_distances_buffer + cdef vector[real] pair_distances2_buffer + + def __init__(self, real cutoff, real[:, ::1]coords, real[:, ::1]searchcoords): + """ + Parameters + ---------- + cutoff : float + Specified cutoff distance + coords : numpy.ndarray + Array with coordinates of atoms of shape ``(N, 3)`` for + ``N`` particles. ``dtype`` must be ``numpy.float32`` + searchcoords : numpy.ndarray + Array with query coordinates. Shape must be ``(M, 3)`` + for ``M`` queries. ``dtype`` must be ``numpy.float32`` + """ + + self.cutoff = cutoff + self.coords = coords + self.searchcoords = searchcoords + + self.npairs = 0 + + cdef void add_neighbors(self, ns_int beadid_i, ns_int beadid_j, real distance2) nogil: + """Internal function to add pairs and distances to buffers + + The buffers populated using this method are used by + other methods of this class. This is the + primary function used by :class:`FastNS` to save the pair of atoms, + which can be considered as neighbors. + """ + + self.pairs_buffer.push_back(beadid_i) + self.pairs_buffer.push_back(beadid_j) + self.pair_distances2_buffer.push_back(distance2) + self.npairs += 1 + + def get_pairs(self): + """Returns all the pairs within the desired cutoff distance + + Returns an array of shape ``(N, 2)``, where N is the number of pairs + between ``reference`` and ``configuration`` within the specified distance. + For every pair ``(i, j)``, ``reference[i]`` and ``configuration[j]`` are + atom positions such that ``reference`` is the position of query + atoms while ``configuration`` coontains the position of group of + atoms used to search against the query atoms. + + Returns + ------- + pairs : numpy.ndarray + pairs of atom indices of neighbors from query + and initial atom coordinates of shape ``(N, 2)`` + """ + + return np.asarray(self.pairs_buffer).reshape(self.npairs, 2) + + def get_pair_distances(self): + """Returns all the distances corresponding to each pair of neighbors + + Returns an array of shape ``N`` where N is the number of pairs + among the query atoms and initial atoms within a specified distance. + Every element ``[i]`` corresponds to the distance between + ``pairs[i, 0]`` and ``pairs[i, 1]``, where pairs is the array + obtained from ``get_pairs()`` + + Returns + ------- + distances : numpy.ndarray + distances between pairs of query and initial + atom coordinates of shape ``N`` + + See Also + -------- + MDAnalysis.lib.nsgrid.NSResults.get_pairs + + """ + + self.pair_distances_buffer = np.sqrt(self.pair_distances2_buffer) + return np.asarray(self.pair_distances_buffer) + + cdef void create_buffers(self) nogil: + """ + Creates buffers to get individual neighbour list and distances + of the query atoms. + + """ + + cdef ns_int i, beadid_i, beadid_j + cdef ns_int idx, nsearch + cdef real dist2, dist + + nsearch = len(self.searchcoords) + + self.indices_buffer = vector[intvec]() + self.distances_buffer = vector[realvec]() + + # initialize rows corresponding to search + for i in range(nsearch): + self.indices_buffer.push_back(intvec()) + self.distances_buffer.push_back(realvec()) + + for i in range(0, 2*self.npairs, 2): + beadid_i = self.pairs_buffer[i] + beadid_j = self.pairs_buffer[i + 1] + + dist2 = self.pair_distances2_buffer[i//2] + + self.indices_buffer[beadid_i].push_back(beadid_j) + + dist = sqrt(dist2) + + self.distances_buffer[beadid_i].push_back(dist) + + def get_indices(self): + """Individual neighbours of query atom + + For every queried atom ``i``, an array of all its neighbors + indices can be obtained from ``get_indices()[i]`` + + Returns + ------- + indices : list + Indices of neighboring atoms. + Every element i.e. ``indices[i]`` will be a list of + size ``m`` where m is the number of neighbours of + query atom[i]. + + .. code-block:: python + + results = NSResults() + indices = results.get_indices() + + ``indices[i]`` will output a list of neighboring + atoms of ``atom[i]`` from query atoms ``atom``. + ``indices[i][j]`` will give the atom-id of initial coordinates + such that ``initial_atom[indices[i][j]]`` is a neighbor of ``atom[i]`` + + """ + + if self.indices_buffer.empty(): + self.create_buffers() + return list(self.indices_buffer) + + def get_distances(self): + """Distance corresponding to individual neighbors of query atom + + For every queried atom ``i``, a list of all the distances + from its neighboring atoms can be obtained from ``get_distances()[i]``. + Every ``distance[i][j]`` will correspond + to the distance between atoms ``atom[i]`` from the query + atoms and ``atom[indices[j]]`` from the initialized + set of coordinates, where ``indices`` can be obtained + by ``get_indices()`` + + Returns + ------- + distances : np.ndarray + Every element i.e. ``distances[i]`` will be an array of + shape ``m`` where m is the number of neighbours of + query atom[i]. + + .. code-block:: python + + results = NSResults() + distances = results.get_distances() + + + atoms of ``atom[i]`` and query atoms ``atom``. + ``indices[i][j]`` will give the atom-id of initial coordinates + such that ``initial_atom[indices[i][j]]`` is a neighbor of ``atom[i]`` + + See Also + -------- + MDAnalysis.lib.nsgrid.NSResults.get_indices + + """ + + if self.distances_buffer.empty(): + self.create_buffers() + return list(self.distances_buffer) + +cdef class NSGrid(object): + """Constructs a uniform cuboidal grid for a brick-shaped box + + This class uses :class:`PBCBox` to define the brick shaped box + It is essential to initialize the box with :class:`PBCBox` + inorder to form the grid. + + The domain is subdivided into number of cells based on the desired search + radius. Ideally cellsize should be equal to the search radius, but small + search radius leads to large cell-list data strucutres. + An optimization of cutoff is imposed to limit the size of data + structure such that the cellsize is always greater than or + equal to cutoff distance. + + Note + ---- + This class assumes that all the coordinates are already + inside the brick shaped box. Care must be taken to ensure + all the particles are within the brick shaped box as + defined by :class:`PBCBox`. This can be ensured by using + :func:`~MDAnalysis.lib.nsgrid.PBCBox.fast_put_atoms_in_bbox` + + .. warning:: + This class is not meant to be used by end users. + + """ + + cdef readonly real cutoff # cutoff + cdef ns_int size # total cells + cdef ns_int ncoords # number of coordinates + cdef ns_int[DIM] ncells # individual cells in every dimension + cdef ns_int[DIM] cell_offsets # Cell Multipliers + cdef real[DIM] cellsize # cell size in every dimension + cdef ns_int nbeads_per_cell # maximum beads + cdef ns_int *nbeads # size (Number of beads in every cell) + cdef ns_int *beadids # size * nbeads_per_cell (Beadids in every cell) + cdef ns_int *cellids # ncoords (Cell occupation id for every atom) + cdef bint force # To negate the effects of optimized cutoff + + def __init__(self, ncoords, cutoff, PBCBox box, max_size, force=False): + """ + Parameters + ---------- + ncoords : int + Number of coordinates to fill inside the brick shaped box + cutoff : float + Desired cutoff radius + box : PBCBox + Instance of :class:`PBCBox` + max_size : int + Maximum total number of cells + force : boolean + Optimizes cutoff if set to ``False`` [False] + """ + + cdef ns_int i + cdef ns_int ncellx, ncelly, ncellz, size, nbeadspercell + cdef ns_int xi, yi, zi + cdef real bbox_vol + + self.ncoords = ncoords + + # Calculate best cutoff + self.cutoff = cutoff + if not force: + bbox_vol = box.c_pbcbox.box[XX][XX] * box.c_pbcbox.box[YY][YY] * box.c_pbcbox.box[YY][YY] + size = bbox_vol/cutoff**3 + nbeadspercell = ncoords/size + while bbox_vol/self.cutoff**3 > max_size: + self.cutoff *= 1.2 + + for i in range(DIM): + self.ncells[i] = (box.c_pbcbox.box[i][i] / self.cutoff) + self.cellsize[i] = box.c_pbcbox.box[i][i] / self.ncells[i] + self.size = self.ncells[XX] * self.ncells[YY] * self.ncells[ZZ] + + self.cell_offsets[XX] = 0 + self.cell_offsets[YY] = self.ncells[XX] + self.cell_offsets[ZZ] = self.ncells[XX] * self.ncells[YY] + + # Allocate memory + self.nbeads = PyMem_Malloc(sizeof(ns_int) * self.size) + if not self.nbeads: + raise MemoryError("Could not allocate memory from NSGrid.nbeads ({} bits requested)".format(sizeof(ns_int) * self.size)) + self.beadids = NULL + self.cellids = PyMem_Malloc(sizeof(ns_int) * self.ncoords) + if not self.cellids: + raise MemoryError("Could not allocate memory from NSGrid.cellids ({} bits requested)".format(sizeof(ns_int) * self.ncoords)) + self.nbeads_per_cell = 0 + + for i in range(self.size): + self.nbeads[i] = 0 + + def __dealloc__(self): + PyMem_Free(self.nbeads) + PyMem_Free(self.beadids) + PyMem_Free(self.cellids) + + cdef ns_int coord2cellid(self, rvec coord) nogil: + """Finds the cell-id for the given coordinate inside the brick shaped box + + Note + ---- + Assumes the coordinate is already inside the brick shaped box. + Return wrong cell-id if this is not the case + """ + return (coord[ZZ] / self.cellsize[ZZ]) * (self.cell_offsets[ZZ]) +\ + (coord[YY] / self.cellsize[YY]) * self.cell_offsets[YY] + \ + (coord[XX] / self.cellsize[XX]) + + cdef bint cellid2cellxyz(self, ns_int cellid, ivec cellxyz) nogil: + """Finds actual cell position `(x, y, z)` from a cell-id + """ + + if cellid < 0: + return False + if cellid >= self.size: + return False + + cellxyz[ZZ] = (cellid / self.cell_offsets[ZZ]) + cellid -= cellxyz[ZZ] * self.cell_offsets[ZZ] + + cellxyz[YY] = (cellid / self.cell_offsets[YY]) + cellxyz[XX] = cellid - cellxyz[YY] * self.cell_offsets[YY] + + return True + + cdef fill_grid(self, real[:, ::1] coords): + """Sorts atoms into cells based on their position in the brick shaped box + + Every atom inside the brick shaped box is assigned a + cell-id based on its position. Another list ``beadids`` + sort the atom-ids in each cell. + + Note + ---- + The method fails if any coordinate is outside the brick shaped box. + + """ + + cdef ns_int i, cellindex = -1 + cdef ns_int ncoords = coords.shape[0] + cdef ns_int[:] beadcounts = np.empty(self.size, dtype=np.int) + + with nogil: + # Initialize buffers + for i in range(self.size): + beadcounts[i] = 0 + + # First loop: find cellindex for each bead + for i in range(ncoords): + cellindex = self.coord2cellid(&coords[i, 0]) + + self.nbeads[cellindex] += 1 + self.cellids[i] = cellindex + + if self.nbeads[cellindex] > self.nbeads_per_cell: + self.nbeads_per_cell = self.nbeads[cellindex] + + # Allocate memory + self.beadids = PyMem_Malloc(sizeof(ns_int) * self.size * self.nbeads_per_cell) # np.empty((self.size, nbeads_max), dtype=np.int) + if not self.beadids: + raise MemoryError("Could not allocate memory for NSGrid.beadids ({} bits requested)".format(sizeof(ns_int) * self.size * self.nbeads_per_cell)) + + with nogil: + # Second loop: fill grid + for i in range(ncoords): + + # Add bead to grid cell + cellindex = self.cellids[i] + self.beadids[cellindex * self.nbeads_per_cell + beadcounts[cellindex]] = i + beadcounts[cellindex] += 1 + + +cdef class FastNS(object): + """Grid based search between two group of atoms + + Instantiates a class object which uses :class:`PBCBox` and + :class:`NSGrid` to construct a cuboidal + grid in an orthogonal brick shaped box. + + Minimum image convention is used for distance evaluations + if box dimensions are provided. + """ + cdef PBCBox box + cdef real[:, ::1] coords + cdef real[:, ::1] coords_bbox + cdef readonly real cutoff + cdef NSGrid grid + cdef ns_int max_gridsize + cdef bint periodic + + def __init__(self, cutoff, coords, box=None, max_gridsize=5000): + """ + Initialize the grid and sort the coordinates in respective + cells by shifting the coordinates in a brick shaped box. + The brick shaped box is defined by :class:`PBCBox` + and cuboidal grid is initialize by :class:`NSGrid`. + If box is supplied, periodic shifts along box vectors are used + to contain all the coordinates inside the brick shaped box. + If box is not supplied, the range of coordinates i.e. + ``[xmax, ymax, zmax] - [xmin, ymin, zmin]`` is used to construct + a pseudo box. Subsequently, the origin is also shifted + to the ``[xmin, ymin, zmin]``. + + Parameters + ---------- + cutoff : float + Desired cutoff distance + coords : numpy.ndarray + atom coordinates of shape ``(N, 3)`` for ``N`` atoms. + ``dtype=numpy.float32`` + box : numpy.ndarray + Box dimension of shape (6, ). The dimensions must be + provided in the same format as returned + by :attr:`MDAnalysis.coordinates.base.Timestep.dimensions`: + ``[lx, ly, lz, alpha, beta, gamma]`` (dtype = numpy.float32) + [None] + max_gridsize : int + maximum number of cells in the grid + + Note + ---- + box=``None`` and ``[0., 0., 0., 90., 90., 90]`` + both are used for Non-PBC aware calculations + + """ + + import MDAnalysis as mda + from MDAnalysis.lib.mdamath import triclinic_vectors + + + cdef real[:] pseudobox = np.zeros(6, dtype=np.float32) + cdef real[DIM] bmax, bmin + cdef ns_int i, j, ncoords + + _check_array(coords, 'coords') + + self.periodic = True + self.coords = coords.copy() + ncoords = len(self.coords) + + if (box is None) or (np.allclose(box[:3], 0.) and box.shape[0] == 6): + if len(coords) == 1: + raise ValueError("Need atleast two coordinates to search using" + " NSGrid without PBC") + + bmax = np.max(self.coords, axis=0) + bmin = np.min(self.coords, axis=0) + for i in range(DIM): + pseudobox[i] = 1.1*(bmax[i] - bmin[i]) + pseudobox[DIM + i] = 90. + box = pseudobox + # shift the origin + for i in range(ncoords): + for j in range(DIM): + self.coords[i][j] -= bmin[j] + self.periodic = False + + if box.shape != (3, 3): + box = triclinic_vectors(box) + + self.box = PBCBox(box, self.periodic) + + if cutoff < 0: + raise ValueError("Cutoff must be positive!") + if cutoff * cutoff > self.box.c_pbcbox.max_cutoff2: + raise ValueError("Cutoff greater than maximum cutoff ({:.3f}) given the PBC") + + self.coords_bbox = self.box.fast_put_atoms_in_bbox(self.coords) + + self.cutoff = cutoff + self.max_gridsize = max_gridsize + # Note that self.cutoff might be different from self.grid.cutoff + self.grid = NSGrid(self.coords_bbox.shape[0], self.cutoff, self.box, self.max_gridsize) + + self.grid.fill_grid(self.coords_bbox) + + def search(self, search_coords): + """Search a group of atoms against initialized coordinates + + Creates a new grid with the query atoms and searches + against the initialized coordinates. The search is exclusive + i.e. only the pairs ``(i, j)`` such that ``atom[i]`` from query atoms + and ``atom[j]`` from the initialized set of coordinates is stored as + neighbors. + + PBC-aware/non PBC-aware calculations are automatically enabled during + the instantiation of :class:FastNS. + + Parameters + ---------- + search_coords : numpy.ndarray + Query coordinates of shape ``(N, 3)`` where + ``N`` is the number of queries + + Returns + ------- + results : NSResults object + The object from :class:NSResults + contains ``get_indices``, ``get_distances``. + ``get_pairs``, ``get_pair_distances`` + + Note + ---- + For non-PBC aware calculations, the current implementation doesn't work + if any of the query coordinates is beyond the specified range of + initialized coordinates in :func:`MDAnalysis.lib.nsgrid.FastNS`. + + See Also + -------- + MDAnalysis.lib.nsgrid.NSResults + + """ + + cdef ns_int i, j, size_search + cdef ns_int d, m + cdef ns_int current_beadid, bid + cdef ns_int cellindex, cellindex_probe + cdef ns_int xi, yi, zi + + cdef NSResults results + + cdef real d2 + cdef rvec probe + + cdef real[:, ::1] searchcoords + cdef real[:, ::1] searchcoords_bbox + cdef NSGrid searchgrid + cdef bint check + + cdef real cutoff2 = self.cutoff * self.cutoff + cdef ns_int npairs = 0 + _check_array(search_coords, 'search_coords') + + # Generate another grid to search + searchcoords = np.ascontiguousarray(search_coords, dtype=np.float32) + searchcoords_bbox = self.box.fast_put_atoms_in_bbox(searchcoords) + searchgrid = NSGrid(searchcoords_bbox.shape[0], self.grid.cutoff, self.box, self.max_gridsize, force=True) + searchgrid.fill_grid(searchcoords_bbox) + + size_search = searchcoords.shape[0] + + results = NSResults(self.cutoff, self.coords, searchcoords) + + with nogil: + for i in range(size_search): + # Start with first search coordinate + current_beadid = i + # find the cellindex of the coordinate + cellindex = searchgrid.cellids[current_beadid] + for xi in range(DIM): + for yi in range(DIM): + for zi in range(DIM): + check = True + #Probe the search coordinates in a brick shaped box + probe[XX] = searchcoords_bbox[current_beadid, XX] + (xi - 1) * searchgrid.cellsize[XX] + probe[YY] = searchcoords_bbox[current_beadid, YY] + (yi - 1) * searchgrid.cellsize[YY] + probe[ZZ] = searchcoords_bbox[current_beadid, ZZ] + (zi - 1) * searchgrid.cellsize[ZZ] + # Make sure the probe coordinates is inside the brick-shaped box + if self.periodic: + for m in range(DIM - 1, -1, -1): + while probe[m] < 0: + for d in range(m+1): + probe[d] += self.box.c_pbcbox.box[m][d] + while probe[m] >= self.box.c_pbcbox.box[m][m]: + for d in range(m+1): + probe[d] -= self.box.c_pbcbox.box[m][d] + else: + for m in range(DIM -1, -1, -1): + if probe[m] < 0: + check = False + break + if probe[m] > self.box.c_pbcbox.box[m][m]: + check = False + break + if not check: + continue + # Get the cell index corresponding to the probe + cellindex_probe = self.grid.coord2cellid(probe) + # for this cellindex search in grid + for j in range(self.grid.nbeads[cellindex_probe]): + bid = self.grid.beadids[cellindex_probe * self.grid.nbeads_per_cell + j] + # find distance between search coords[i] and coords[bid] + d2 = self.box.fast_distance2(&searchcoords_bbox[current_beadid, XX], &self.coords_bbox[bid, XX]) + if d2 < cutoff2: + results.add_neighbors(current_beadid, bid, d2) + npairs += 1 + return results + + def self_search(self): + """Searches all the pairs within the initialized coordinates + + All the pairs among the initialized coordinates are registered + in hald the time. Although the algorithm is still the same, but + the distance checks can be reduced to half in this particular case + as every pair need not be evaluated twice. + + Returns + ------- + results : NSResults object + The object from :class:NSResults + contains ``get_indices``, ``get_distances``. + ``get_pairs``, ``get_pair_distances`` + + See Also + -------- + MDAnalysis.lib.nsgrid.NSResults + + """ + + cdef ns_int i, j, size_search + cdef ns_int d, m + cdef ns_int current_beadid, bid + cdef ns_int cellindex, cellindex_probe + cdef ns_int xi, yi, zi + + cdef NSResults results + cdef real d2 + cdef rvec probe + + cdef real cutoff2 = self.cutoff * self.cutoff + cdef ns_int npairs = 0 + cdef bint check + + size_search = self.coords.shape[0] + + results = NSResults(self.cutoff, self.coords, self.coords) + + with nogil: + for i in range(size_search): + # Start with first search coordinate + current_beadid = i + # find the cellindex of the coordinate + cellindex = self.grid.cellids[current_beadid] + for xi in range(DIM): + for yi in range(DIM): + for zi in range(DIM): + check = True + # Calculate and/or reinitialize shifted coordinates + #Probe the search coordinates in a brick shaped box + probe[XX] = self.coords_bbox[current_beadid, XX] + (xi - 1) * self.grid.cellsize[XX] + probe[YY] = self.coords_bbox[current_beadid, YY] + (yi - 1) * self.grid.cellsize[XX] + probe[ZZ] = self.coords_bbox[current_beadid, ZZ] + (zi - 1) * self.grid.cellsize[XX] + # Make sure the shifted coordinates is inside the brick-shaped box + if self.periodic: + for m in range(DIM - 1, -1, -1): + while probe[m] < 0: + for d in range(m+1): + probe[d] += self.box.c_pbcbox.box[m][d] + while probe[m] >= self.box.c_pbcbox.box[m][m]: + for d in range(m+1): + probe[d] -= self.box.c_pbcbox.box[m][d] + else: + for m in range(DIM -1, -1, -1): + if probe[m] < 0: + check = False + break + elif probe[m] > self.box.c_pbcbox.box[m][m]: + check = False + break + if not check: + continue + # Get the cell index corresponding to the probe + cellindex_probe = self.grid.coord2cellid(probe) + # for this cellindex search in grid + for j in range(self.grid.nbeads[cellindex_probe]): + bid = self.grid.beadids[cellindex_probe * self.grid.nbeads_per_cell + j] + if bid < current_beadid: + continue + # find distance between search coords[i] and coords[bid] + d2 = self.box.fast_distance2(&self.coords_bbox[current_beadid, XX], &self.coords_bbox[bid, XX]) + if d2 < cutoff2 and d2 > EPSILON: + results.add_neighbors(current_beadid, bid, d2) + results.add_neighbors(bid, current_beadid, d2) + npairs += 1 + return results diff --git a/testsuite/MDAnalysisTests/lib/test_nsgrid.py b/testsuite/MDAnalysisTests/lib/test_nsgrid.py index d748808bb28..0a5e5aa6126 100644 --- a/testsuite/MDAnalysisTests/lib/test_nsgrid.py +++ b/testsuite/MDAnalysisTests/lib/test_nsgrid.py @@ -77,17 +77,18 @@ def test_ns_grid_noneighbor(universe): results_grid = run_grid_search(universe, ref_id, cutoff) - assert len(results_grid.get_distances()[0]) == 0 - assert len(results_grid.get_indices()[0]) == 0 - assert len(results_grid.get_pairs()) == 0 - assert len(results_grid.get_pair_distances()) == 0 + # same indices will be selected as neighbour here + assert len(results_grid.get_distances()[0]) == 1 + assert len(results_grid.get_indices()[0]) == 1 + assert len(results_grid.get_pairs()) == 1 + assert len(results_grid.get_pair_distances()) == 1 def test_nsgrid_PBC_rect(): """Check that nsgrid works with rect boxes and PBC""" ref_id = 191 - results = np.array([191, 672, 682, 683, 684, 995, 996, 2060, 2808, 3300, 3791, + results = np.array([191, 192, 672, 682, 683, 684, 995, 996, 2060, 2808, 3300, 3791, 3792]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! universe = mda.Universe(Martini_membrane_gro) @@ -110,7 +111,7 @@ def test_nsgrid_PBC(universe): """Check that grid search works when PBC is needed""" ref_id = 13937 - results = np.array([4398, 4401, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, + results = np.array([4398, 4401, 13938, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, 47451]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! results_grid = run_grid_search(universe, ref_id).get_indices()[0] @@ -122,7 +123,7 @@ def test_nsgrid_pairs(universe): """Check that grid search returns the proper pairs""" ref_id = 13937 - neighbors = np.array([4398, 4401, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, + neighbors = np.array([4398, 4401, 13938, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, 47451]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! results = [] @@ -137,7 +138,7 @@ def test_nsgrid_pair_distances(universe): """Check that grid search returns the proper pair distances""" ref_id = 13937 - results = np.array([0.270, 0.285, 0.096, 0.096, 0.015, 0.278, 0.268, 0.179, 0.259, 0.290, + results = np.array([0.0, 0.270, 0.285, 0.096, 0.096, 0.015, 0.278, 0.268, 0.179, 0.259, 0.290, 0.270]) * 10 # These distances where obtained by gmx distance so they are in nm results_grid = run_grid_search(universe, ref_id).get_pair_distances() @@ -151,7 +152,7 @@ def test_nsgrid_distances(universe): """Check that grid search returns the proper distances""" ref_id = 13937 - results = np.array([0.270, 0.285, 0.096, 0.096, 0.015, 0.278, 0.268, 0.179, 0.259, 0.290, + results = np.array([0.0, 0.270, 0.285, 0.096, 0.096, 0.015, 0.278, 0.268, 0.179, 0.259, 0.290, 0.270]) * 10 # These distances where obtained by gmx distance so they are in nm results_grid = run_grid_search(universe, ref_id).get_distances()[0] @@ -169,7 +170,7 @@ def test_nsgrid_search(box, results): points = (np.random.uniform(low=0, high=1.0, size=(100, 3))*(10.)).astype(np.float32) cutoff = 2.0 - query = np.array([1., 1., 1.]).reshape((1,3)) + query = np.array([1., 1., 1.], dtype=np.float32).reshape((1,3)) searcher = nsgrid.FastNS(cutoff, points, box=box) searchresults = searcher.search(query) indices = searchresults.get_indices()[0] @@ -189,4 +190,11 @@ def test_nsgrid_selfsearch(box, result): searcher = nsgrid.FastNS(cutoff, points, box=box) searchresults = searcher.self_search() pairs = searchresults.get_pairs() - assert_equal(len(pairs)//2, result) \ No newline at end of file + assert_equal(len(pairs)//2, result) + +def test_gridfail(): + points = np.array([[1., 1., 1.]], dtype=np.float32) + cutoff = 0.3 + match = "Need atleast two coordinates to search using NSGrid without PBC" + with pytest.raises(ValueError, match=match): + searcher = nsgrid.FastNS(cutoff, points) \ No newline at end of file From 76c81cf3660309ab7ab3760c5c2a1001669d9d21 Mon Sep 17 00:00:00 2001 From: ayush Date: Mon, 6 Aug 2018 02:05:31 -0700 Subject: [PATCH 45/47] Added self_capped for nsgrid, moved no pbc box handling outside nsgrid and modified tests --- package/MDAnalysis/lib/__init__.py | 2 +- package/MDAnalysis/lib/distances.py | 124 +++++++++++++++--- package/MDAnalysis/lib/nsgrid.pyx | 107 ++++++++------- .../MDAnalysisTests/lib/test_distances.py | 27 ++-- testsuite/MDAnalysisTests/lib/test_nsgrid.py | 66 ++++++---- 5 files changed, 218 insertions(+), 108 deletions(-) diff --git a/package/MDAnalysis/lib/__init__.py b/package/MDAnalysis/lib/__init__.py index d5a583f9aea..37ebc2f811c 100644 --- a/package/MDAnalysis/lib/__init__.py +++ b/package/MDAnalysis/lib/__init__.py @@ -29,7 +29,7 @@ from __future__ import absolute_import __all__ = ['log', 'transformations', 'util', 'mdamath', 'distances', - 'NeighborSearch', 'formats', 'pkdtree', 'grid'] + 'NeighborSearch', 'formats', 'pkdtree', 'nsgrid'] from . import log from . import transformations diff --git a/package/MDAnalysis/lib/distances.py b/package/MDAnalysis/lib/distances.py index 729ebb1e3ba..71608a72fca 100644 --- a/package/MDAnalysis/lib/distances.py +++ b/package/MDAnalysis/lib/distances.py @@ -462,7 +462,8 @@ def capped_distance(reference, configuration, max_cutoff, min_cutoff=None, Currently only supports brute force and Periodic KDtree .. SeeAlso:: :func:'MDAnalysis.lib.distances.distance_array' - .. SeeAlso:: :func:'MDAnalysis.lib.pkdtree.PeriodicKDTree' + .. SeeAlso:: :func:'MDAnalysis.lib.pkdtree.PeriodicKDTree.search' + .. SeeAlso:: :class:'MDAnalysis.lib.nsgrid.FastNS.search' """ if box is not None: if box.shape[0] != 6: @@ -517,11 +518,12 @@ def _determine_method(reference, configuration, max_cutoff, min_cutoff=None, Currently implemented methods are present in the ``methods`` dictionary bruteforce : returns ``_bruteforce_capped`` PKDtree : return ``_pkdtree_capped` + NSGrid : return ``_nsgrid_capped` """ methods = {'bruteforce': _bruteforce_capped, - 'pkdtree': _pkdtree_capped, - 'nsgrid': _nsgrid_capped} + 'pkdtree': _pkdtree_capped, + 'nsgrid': _nsgrid_capped} if method is not None: return methods[method] @@ -569,7 +571,6 @@ def _bruteforce_capped(reference, configuration, max_cutoff, """ pairs, distance = [], [] - reference = np.asarray(reference, dtype=np.float32) configuration = np.asarray(configuration, dtype=np.float32) @@ -644,12 +645,13 @@ def _pkdtree_capped(reference, configuration, max_cutoff, distances.append(dist[num]) return pairs, distances + def _nsgrid_capped(reference, configuration, max_cutoff, min_cutoff=None, box=None): - """Search all the pairs in *reference* and *configuration* within + """Search all the pairs in *reference* and *configuration* within a specified distance using Grid Search - + Parameters ----------- reference : array @@ -662,16 +664,12 @@ def _nsgrid_capped(reference, configuration, max_cutoff, min_cutoff=None, Maximum cutoff distance between the reference and configuration min_cutoff : (optional) float Minimum cutoff distance between reference and configuration [None] - box : array + box : array The dimensions, if provided, must be provided in the same The unitcell dimesions for this system format as returned by :attr:`MDAnalysis.coordinates.base.Timestep.dimensions`: ``[lx,ly, lz, alpha, beta, gamma]``. Minimum image convention - is applied if the box is provided - - Note - ---- - Non Periodic Boundary conditions can-not be handled currently + is applied if the box is provided """ from .nsgrid import FastNS @@ -679,11 +677,40 @@ def _nsgrid_capped(reference, configuration, max_cutoff, min_cutoff=None, reference = reference[None, :] if configuration.shape == (3, ): configuration = configuration[None, :] - gridsearch = FastNS(max_cutoff, configuration, box=box) - results = gridsearch.search(reference) + + if box is None: + # create a pseudobox + # define the max range + # and supply the pseudobox + # along with only one set of coordinates + pseudobox = np.zeros(6, dtype=np.float32) + all_coords = np.concatenate([reference, configuration]) + lmax = all_coords.max(axis=0) + lmin = all_coords.min(axis=0) + # Using maximum dimension as the box size + boxsize = (lmax-lmin).max() + # to avoid failures of very close particles + # but with larger cutoff + if boxsize < 2*max_cutoff: + # just enough box size so that NSGrid doesnot fails + sizefactor = 2.2*max_cutoff/boxsize + else: + sizefactor = 1.2 + pseudobox[:3] = sizefactor*boxsize + pseudobox[3:] = 90. + shiftref, shiftconf = reference.copy(), configuration.copy() + # Extra padding near the origin + shiftref -= lmin - 0.1*boxsize + shiftconf -= lmin - 0.1*boxsize + gridsearch = FastNS(max_cutoff, shiftconf, box=pseudobox, pbc=False) + results = gridsearch.search(shiftref) + else: + gridsearch = FastNS(max_cutoff, configuration, box=box) + results = gridsearch.search(reference) + pairs = results.get_pairs() pair_distance = results.get_pair_distances() - + if min_cutoff is not None: idx = pair_distance > min_cutoff pairs, pair_distance = pairs[idx], pair_distance[idx] @@ -741,10 +768,11 @@ def self_capped_distance(reference, max_cutoff, min_cutoff=None, Note ----- - Currently only supports brute force and Periodic KDtree + Currently only supports brute force, Periodic KDtree and Grid Search .. SeeAlso:: :func:'MDAnalysis.lib.distances.self_distance_array' - .. SeeAlso:: :func:'MDAnalysis.lib.pkdtree.PeriodicKDTree' + .. SeeAlso:: :func:'MDAnalysis.lib.pkdtree.PeriodicKDTree.search' + .. SeeAlso:: :func:'MDAnalysis.lib.nsgrid.FastNS.self_search' """ if box is not None: if box.shape[0] != 6: @@ -796,10 +824,12 @@ def _determine_method_self(reference, max_cutoff, min_cutoff=None, Currently implemented methods are present in the ``methods`` dictionary bruteforce : returns ``_bruteforce_capped_self`` PKDtree : return ``_pkdtree_capped_self`` + NSGrid : return ``_nsgrid_capped_self`` """ methods = {'bruteforce': _bruteforce_capped_self, - 'pkdtree': _pkdtree_capped_self} + 'pkdtree': _pkdtree_capped_self, + 'nsgrid': _nsgrid_capped_self} if method is not None: return methods[method] @@ -904,6 +934,64 @@ def _pkdtree_capped_self(reference, max_cutoff, min_cutoff=None, return np.asarray(pairs), np.asarray(distance) +def _nsgrid_capped_self(reference, max_cutoff, min_cutoff=None, + box=None): + """Finds all the pairs among the *reference* coordinates within + a fixed distance using gridsearch + + Returns + ------- + pairs : array + Arrray of ``[i, j]`` pairs such that atom-index ``i`` + and ``j`` from reference array are within a fixed distance + distance: array + Distance between ``reference[i]`` and ``reference[j]`` + atom coordinate + + """ + from .nsgrid import FastNS + + reference = np.asarray(reference, dtype=np.float32) + if reference.shape == (3, ) or len(reference) == 1: + return [], [] + + if box is None: + # create a pseudobox + # define the max range + # and supply the pseudobox + # along with only one set of coordinates + pseudobox = np.zeros(6, dtype=np.float32) + lmax = reference.max(axis=0) + lmin = reference.min(axis=0) + # Using maximum dimension as the box size + boxsize = (lmax-lmin).max() + # to avoid failures of very close particles + # but with larger cutoff + if boxsize < 2*max_cutoff: + # just enough box size so that NSGrid doesnot fails + sizefactor = 2.2*max_cutoff/boxsize + else: + sizefactor = 1.2 + pseudobox[:3] = sizefactor*boxsize + pseudobox[3:] = 90. + shiftref = reference.copy() + # Extra padding near the origin + shiftref -= lmin - 0.1*boxsize + gridsearch = FastNS(max_cutoff, shiftref, box=pseudobox, pbc=False) + results = gridsearch.self_search() + else: + gridsearch = FastNS(max_cutoff, reference, box=box) + results = gridsearch.self_search() + + pairs = results.get_pairs()[::2, :] + pair_distance = results.get_pair_distances()[::2] + + if min_cutoff is not None: + idx = pair_distance > min_cutoff + pairs, pair_distance = pairs[idx], pair_distance[idx] + return pairs, pair_distance + + def transform_RtoS(inputcoords, box, backend="serial"): """Transform an array of coordinates from real space to S space (aka lambda space) diff --git a/package/MDAnalysis/lib/nsgrid.pyx b/package/MDAnalysis/lib/nsgrid.pyx index 96ee469efde..97222e2d9aa 100644 --- a/package/MDAnalysis/lib/nsgrid.pyx +++ b/package/MDAnalysis/lib/nsgrid.pyx @@ -58,7 +58,7 @@ Appendix F, Page 552 of ``Understanding Molecular Dynamics: From Algorithm to Applications`` by Frenkel and Smit. In brief, the algorithm divides the domain into smaller subdomains called `cells` -and distributes every particle to these cells based on its position. Subsequently, +and distributes every particle to these cells based on their positions. Subsequently, any distance based query first identifies the corresponding cell position in the domain followed by distance evaluations within the identified cell and neighboring cells only. Care must be taken to ensure that `cellsize` is @@ -130,8 +130,10 @@ cdef class PBCBox(object): This class is not meant to be used by end users. .. warning:: - Even if MD triclinic boxes can be handled by this class, internal optimization is made based on the - assumption that particles are inside a brick-shaped box. When this is not the case, calculated distances are not + Even if MD triclinic boxes can be handled by this class, + internal optimization is made based on the assumption that + particles are inside a brick-shaped box. When this is not + the case, calculated distances are not warranted to be exact. """ @@ -299,6 +301,7 @@ cdef class PBCBox(object): bbox_coords[i, m] += self.c_pbcbox.box[m][m] while bbox_coords[i, m] >= self.c_pbcbox.box[m][m]: bbox_coords[i, m] -= self.c_pbcbox.box[m][m] + return bbox_coords ######################### @@ -351,8 +354,9 @@ cdef class NSResults(object): The buffers populated using this method are used by other methods of this class. This is the - primary function used by :class:`FastNS` to save the pair of atoms, - which can be considered as neighbors. + primary function used by :class:`FastNS` to save all + the pair of atoms, + which are considered as neighbors. """ self.pairs_buffer.push_back(beadid_i) @@ -585,10 +589,12 @@ cdef class NSGrid(object): self.cell_offsets[ZZ] = self.ncells[XX] * self.ncells[YY] # Allocate memory + # Number of beads in every cell self.nbeads = PyMem_Malloc(sizeof(ns_int) * self.size) if not self.nbeads: raise MemoryError("Could not allocate memory from NSGrid.nbeads ({} bits requested)".format(sizeof(ns_int) * self.size)) self.beadids = NULL + # Cellindex of every bead self.cellids = PyMem_Malloc(sizeof(ns_int) * self.ncoords) if not self.cellids: raise MemoryError("Could not allocate memory from NSGrid.cellids ({} bits requested)".format(sizeof(ns_int) * self.ncoords)) @@ -686,7 +692,7 @@ cdef class FastNS(object): grid in an orthogonal brick shaped box. Minimum image convention is used for distance evaluations - if box dimensions are provided. + if pbc is set to ``True``. """ cdef PBCBox box cdef real[:, ::1] coords @@ -696,7 +702,7 @@ cdef class FastNS(object): cdef ns_int max_gridsize cdef bint periodic - def __init__(self, cutoff, coords, box=None, max_gridsize=5000): + def __init__(self, cutoff, coords, box, max_gridsize=5000, pbc=True): """ Initialize the grid and sort the coordinates in respective cells by shifting the coordinates in a brick shaped box. @@ -705,9 +711,10 @@ cdef class FastNS(object): If box is supplied, periodic shifts along box vectors are used to contain all the coordinates inside the brick shaped box. If box is not supplied, the range of coordinates i.e. - ``[xmax, ymax, zmax] - [xmin, ymin, zmin]`` is used to construct - a pseudo box. Subsequently, the origin is also shifted - to the ``[xmin, ymin, zmin]``. + ``[xmax, ymax, zmax] - [xmin, ymin, zmin]`` should be used + to construct a pseudo box. Subsequently, the origin should also be + shifted to ``[xmin, ymin, zmin]``. These arguments must be provided + to the function. Parameters ---------- @@ -715,53 +722,58 @@ cdef class FastNS(object): Desired cutoff distance coords : numpy.ndarray atom coordinates of shape ``(N, 3)`` for ``N`` atoms. - ``dtype=numpy.float32`` + ``dtype=numpy.float32``. For Non-PBC calculations, + all the coords must be within the bounding box specified + by ``box`` box : numpy.ndarray Box dimension of shape (6, ). The dimensions must be provided in the same format as returned by :attr:`MDAnalysis.coordinates.base.Timestep.dimensions`: - ``[lx, ly, lz, alpha, beta, gamma]`` (dtype = numpy.float32) - [None] + ``[lx, ly, lz, alpha, beta, gamma]``. For non-PBC + evaluations, provide an orthogonal bounding box + (dtype = numpy.float32) max_gridsize : int - maximum number of cells in the grid + maximum number of cells in the grid. This parameter + can be tuned for superior performance. + pbc : boolean + Handle to switch periodic boundary conditions on/off [True] Note ---- - box=``None`` and ``[0., 0., 0., 90., 90., 90]`` - both are used for Non-PBC aware calculations + * ``pbc=False`` Only works for orthogonal boxes. + * Care must be taken such that all particles are inside + the bounding box as defined by the box argument for non-PBC + calculations. + * In case of Non-PBC calculations, a bounding box must be provided + to encompass all the coordinates as well as the search coordinates. + The dimension should be similar to ``box`` argument but for + an orthogonal box. For instance, one valid set of argument + for ``box`` for the case of no PBC could be + ``[10, 10, 10, 90, 90, 90]`` + * Following operations are advisable for non-PBC calculations + + ..code-block:: python + + lmax = all_coords.max(axis=0) + lmin = all_coords.min(axis=0) + pseudobox[:3] = 1.1*(lmax - lmin) + pseudobox[3:] = 90. + shift = all_coords.copy() + shift -= lmin + gridsearch = FastNS(max_cutoff, shift, box=pseudobox, pbc=False) """ - import MDAnalysis as mda from MDAnalysis.lib.mdamath import triclinic_vectors - - cdef real[:] pseudobox = np.zeros(6, dtype=np.float32) - cdef real[DIM] bmax, bmin - cdef ns_int i, j, ncoords _check_array(coords, 'coords') - self.periodic = True + if np.allclose(box[:3], 0.0): + raise ValueError("Any of the box dimensions cannot be 0") + + self.periodic = pbc self.coords = coords.copy() - ncoords = len(self.coords) - - if (box is None) or (np.allclose(box[:3], 0.) and box.shape[0] == 6): - if len(coords) == 1: - raise ValueError("Need atleast two coordinates to search using" - " NSGrid without PBC") - - bmax = np.max(self.coords, axis=0) - bmin = np.min(self.coords, axis=0) - for i in range(DIM): - pseudobox[i] = 1.1*(bmax[i] - bmin[i]) - pseudobox[DIM + i] = 90. - box = pseudobox - # shift the origin - for i in range(ncoords): - for j in range(DIM): - self.coords[i][j] -= bmin[j] - self.periodic = False if box.shape != (3, 3): box = triclinic_vectors(box) @@ -778,6 +790,7 @@ cdef class FastNS(object): self.cutoff = cutoff self.max_gridsize = max_gridsize # Note that self.cutoff might be different from self.grid.cutoff + # due to optimization self.grid = NSGrid(self.coords_bbox.shape[0], self.cutoff, self.box, self.max_gridsize) self.grid.fill_grid(self.coords_bbox) @@ -810,8 +823,8 @@ cdef class FastNS(object): Note ---- For non-PBC aware calculations, the current implementation doesn't work - if any of the query coordinates is beyond the specified range of - initialized coordinates in :func:`MDAnalysis.lib.nsgrid.FastNS`. + if any of the query coordinates is beyond the range specified in + ``box`` in :func:`MDAnalysis.lib.nsgrid.FastNS`. See Also -------- @@ -944,10 +957,10 @@ cdef class FastNS(object): for zi in range(DIM): check = True # Calculate and/or reinitialize shifted coordinates - #Probe the search coordinates in a brick shaped box + # Probe the search coordinates in a brick shaped box probe[XX] = self.coords_bbox[current_beadid, XX] + (xi - 1) * self.grid.cellsize[XX] - probe[YY] = self.coords_bbox[current_beadid, YY] + (yi - 1) * self.grid.cellsize[XX] - probe[ZZ] = self.coords_bbox[current_beadid, ZZ] + (zi - 1) * self.grid.cellsize[XX] + probe[YY] = self.coords_bbox[current_beadid, YY] + (yi - 1) * self.grid.cellsize[YY] + probe[ZZ] = self.coords_bbox[current_beadid, ZZ] + (zi - 1) * self.grid.cellsize[ZZ] # Make sure the shifted coordinates is inside the brick-shaped box if self.periodic: for m in range(DIM - 1, -1, -1): @@ -962,7 +975,7 @@ cdef class FastNS(object): if probe[m] < 0: check = False break - elif probe[m] > self.box.c_pbcbox.box[m][m]: + elif probe[m] >= self.box.c_pbcbox.box[m][m]: check = False break if not check: @@ -979,5 +992,5 @@ cdef class FastNS(object): if d2 < cutoff2 and d2 > EPSILON: results.add_neighbors(current_beadid, bid, d2) results.add_neighbors(bid, current_beadid, d2) - npairs += 1 + npairs += 2 return results diff --git a/testsuite/MDAnalysisTests/lib/test_distances.py b/testsuite/MDAnalysisTests/lib/test_distances.py index d83d8903e93..38ae57384c8 100644 --- a/testsuite/MDAnalysisTests/lib/test_distances.py +++ b/testsuite/MDAnalysisTests/lib/test_distances.py @@ -25,12 +25,8 @@ import numpy as np from numpy.testing import assert_equal, assert_almost_equal -import itertools - import MDAnalysis as mda -from MDAnalysis.lib.mdamath import triclinic_vectors, triclinic_box - @pytest.mark.parametrize('coord_dtype', (np.float32, np.float64)) def test_transform_StoR_pass(coord_dtype): box = np.array([10, 7, 3, 45, 60, 90], dtype=np.float32) @@ -77,10 +73,6 @@ def test_capped_distance_noresults(): @pytest.mark.parametrize('method', method_1) @pytest.mark.parametrize('min_cutoff', min_cutoff_1) def test_capped_distance_checkbrute(npoints, box, query, method, min_cutoff): - if method == 'nsgrid' and box is None: - pytest.skip('Not implemented yet') - - np.random.seed(90003) points = (np.random.uniform(low=0, high=1.0, size=(npoints, 3))*(boxes_1[0][:3])).astype(np.float32) @@ -119,7 +111,7 @@ def test_self_capped_distance(npoints, box, method, min_cutoff): np.random.seed(90003) points = (np.random.uniform(low=0, high=1.0, size=(npoints, 3))*(boxes_1[0][:3])).astype(np.float32) - max_cutoff = 0.1 + max_cutoff = 0.2 pairs, distance = mda.lib.distances.self_capped_distance(points, max_cutoff, min_cutoff=min_cutoff, @@ -131,16 +123,17 @@ def test_self_capped_distance(npoints, box, method, min_cutoff): points[i+1:], box=box) if min_cutoff is not None: - idx = np.where((dist < max_cutoff) & (dist > min_cutoff))[0] + idx = np.where((dist < max_cutoff) & (dist > min_cutoff))[1] else: - idx = np.where((dist < max_cutoff))[0] + idx = np.where((dist < max_cutoff))[1] for other_idx in idx: j = other_idx + 1 + i found_pairs.append((i, j)) - found_distance.append(dist[other_idx]) + found_distance.append(dist[0, other_idx]) assert_equal(len(pairs), len(found_pairs)) -@pytest.mark.parametrize('box', (None, + +@pytest.mark.parametrize('box', (None, np.array([1, 1, 1, 90, 90, 90], dtype=np.float32), np.array([1, 1, 1, 60, 75, 80], dtype=np.float32))) @pytest.mark.parametrize('npoints,cutoff,meth', @@ -158,7 +151,7 @@ def test_method_selfselection(box, npoints, cutoff, meth): assert_equal(method.__name__, meth) -@pytest.mark.parametrize('box', (None, +@pytest.mark.parametrize('box', (None, np.array([1, 1, 1, 90, 90, 90], dtype=np.float32), np.array([1, 1, 1, 60, 75, 80], dtype=np.float32))) @pytest.mark.parametrize('npoints,cutoff,meth', @@ -178,9 +171,9 @@ def test_method_selection(box, npoints, cutoff, meth): # different boxlengths to shift a coordinate shifts = [ - ((0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0)), # no shifting - ((1, 0, 0), (0, 1, 1), (0, 0, 1), (1, 1, 0)), # single box lengths - ((-1, 0, 1), (0, -1, 0), (1, 0, 1), (-1, -1, -1)), # negative single + ((0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0)), # no shifting + ((1, 0, 0), (0, 1, 1), (0, 0, 1), (1, 1, 0)), # single box lengths + ((-1, 0, 1), (0, -1, 0), (1, 0, 1), (-1, -1, -1)), # negative single ((4, 3, -2), (-2, 2, 2), (-5, 2, 2), (0, 2, 2)), # multiple boxlengths ] diff --git a/testsuite/MDAnalysisTests/lib/test_nsgrid.py b/testsuite/MDAnalysisTests/lib/test_nsgrid.py index 0a5e5aa6126..03b3e24c7d3 100644 --- a/testsuite/MDAnalysisTests/lib/test_nsgrid.py +++ b/testsuite/MDAnalysisTests/lib/test_nsgrid.py @@ -24,7 +24,7 @@ import pytest -from numpy.testing import assert_equal, assert_allclose, assert_array_equal +from numpy.testing import assert_equal, assert_allclose import numpy as np import MDAnalysis as mda @@ -37,8 +37,6 @@ def universe(): u = mda.Universe(GRO) return u - - def run_grid_search(u, ref_id, cutoff=3): coords = u.atoms.positions searchcoords = u.atoms.positions[ref_id] @@ -84,7 +82,6 @@ def test_ns_grid_noneighbor(universe): assert len(results_grid.get_pair_distances()) == 1 - def test_nsgrid_PBC_rect(): """Check that nsgrid works with rect boxes and PBC""" ref_id = 191 @@ -97,9 +94,9 @@ def test_nsgrid_PBC_rect(): # FastNS is called differently to max coverage searcher = nsgrid.FastNS(cutoff, universe.atoms.positions, box=universe.dimensions) - results_grid = searcher.search(universe.atoms.positions[ref_id][None, :]).get_indices()[0] + results_grid = searcher.search(universe.atoms.positions[ref_id][None, :]).get_indices()[0] - results_grid2 = searcher.search(universe.atoms.positions).get_indices() # call without specifying any ids, should do NS for all beads + results_grid2 = searcher.search(universe.atoms.positions).get_indices() # call without specifying any ids, should do NS for all beads assert_equal(np.sort(results), np.sort(results_grid)) assert_equal(len(universe.atoms), len(results_grid2)) @@ -126,7 +123,7 @@ def test_nsgrid_pairs(universe): neighbors = np.array([4398, 4401, 13938, 13939, 13940, 13941, 17987, 23518, 23519, 23521, 23734, 47451]) - 1 # Atomid are from gmx select so there start from 1 and not 0. hence -1! results = [] - + results = np.array(results) results_grid = run_grid_search(universe, ref_id).get_pairs() @@ -146,8 +143,6 @@ def test_nsgrid_pair_distances(universe): assert_allclose(np.sort(results), np.sort(results_grid), atol=1e-2) - - def test_nsgrid_distances(universe): """Check that grid search returns the proper distances""" @@ -161,8 +156,7 @@ def test_nsgrid_distances(universe): @pytest.mark.parametrize('box, results', - ((None, [3, 13, 24]), - (np.array([0., 0., 0., 90., 90., 90.]), [3, 13, 24]), + ((None, [3, 13, 24]), (np.array([10., 10., 10., 90., 90., 90.]), [3, 13, 24, 39, 67]), (np.array([10., 10., 10., 60., 75., 90.]), [3, 13, 24, 39, 60, 79]))) def test_nsgrid_search(box, results): @@ -170,16 +164,30 @@ def test_nsgrid_search(box, results): points = (np.random.uniform(low=0, high=1.0, size=(100, 3))*(10.)).astype(np.float32) cutoff = 2.0 - query = np.array([1., 1., 1.], dtype=np.float32).reshape((1,3)) - searcher = nsgrid.FastNS(cutoff, points, box=box) - searchresults = searcher.search(query) + query = np.array([1., 1., 1.], dtype=np.float32).reshape((1, 3)) + + if box is None: + pseudobox = np.zeros(6, dtype=np.float32) + all_coords = np.concatenate([points, query]) + lmax = all_coords.max(axis=0) + lmin = all_coords.min(axis=0) + pseudobox[:3] = 1.1*(lmax - lmin) + pseudobox[3:] = 90. + shiftpoints, shiftquery = points.copy(), query.copy() + shiftpoints -= lmin + shiftquery -= lmin + searcher = nsgrid.FastNS(cutoff, shiftpoints, box=pseudobox, pbc=False) + searchresults = searcher.search(shiftquery) + else: + searcher = nsgrid.FastNS(cutoff, points, box) + searchresults = searcher.search(query) indices = searchresults.get_indices()[0] assert_equal(np.sort(indices), results) -@pytest.mark.parametrize('box, result', +@pytest.mark.parametrize('box, result', ((None, 21), - (np.array([0., 0., 0., 90., 90., 90.]), 21), + (np.array([0., 0., 0., 90., 90., 90.]), 21), (np.array([10., 10., 10., 90., 90., 90.]), 26), (np.array([10., 10., 10., 60., 75., 90.]), 33))) def test_nsgrid_selfsearch(box, result): @@ -187,14 +195,22 @@ def test_nsgrid_selfsearch(box, result): points = (np.random.uniform(low=0, high=1.0, size=(100, 3))*(10.)).astype(np.float32) cutoff = 1.0 - searcher = nsgrid.FastNS(cutoff, points, box=box) - searchresults = searcher.self_search() + if box is None or np.allclose(box[:3], 0): + # create a pseudobox + # define the max range + # and supply the pseudobox + # along with only one set of coordinates + pseudobox = np.zeros(6, dtype=np.float32) + lmax = points.max(axis=0) + lmin = points.min(axis=0) + pseudobox[:3] = 1.1*(lmax - lmin) + pseudobox[3:] = 90. + shiftref = points.copy() + shiftref -= lmin + searcher = nsgrid.FastNS(cutoff, shiftref, box=pseudobox, pbc=False) + searchresults = searcher.self_search() + else: + searcher = nsgrid.FastNS(cutoff, points, box=box) + searchresults = searcher.self_search() pairs = searchresults.get_pairs() assert_equal(len(pairs)//2, result) - -def test_gridfail(): - points = np.array([[1., 1., 1.]], dtype=np.float32) - cutoff = 0.3 - match = "Need atleast two coordinates to search using NSGrid without PBC" - with pytest.raises(ValueError, match=match): - searcher = nsgrid.FastNS(cutoff, points) \ No newline at end of file From 2a0b01f8863c24abacb24db522ad23c065a9c491 Mon Sep 17 00:00:00 2001 From: ayush Date: Mon, 6 Aug 2018 17:31:25 -0700 Subject: [PATCH 46/47] updated CHANGELOG --- package/CHANGELOG | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/package/CHANGELOG b/package/CHANGELOG index 93524527cd4..18373bae4ad 100644 --- a/package/CHANGELOG +++ b/package/CHANGELOG @@ -20,6 +20,11 @@ The rules for this file: Enhancements + * Added a wrapper of lib.nsgrid in lib.distances.self_capped_distance + and lib.distances.capped_distanceto automatically chose the fastest + method for distance based calculations. (PR #2008) + * Added Grid search functionality in lib.nsgrid for faster distance based + calculations. (PR #2008) * Modified around selections to work with KDTree and periodic boundary conditions. Should reduce memory usage (#974 PR #2022) * Modified topology.guessers.guess_bonds to automatically select the From 0cde1ded98b600d4b5c3e07ef0c8637f8c1785e7 Mon Sep 17 00:00:00 2001 From: ayush Date: Mon, 6 Aug 2018 18:30:29 -0700 Subject: [PATCH 47/47] added sebastien to authors --- package/CHANGELOG | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/package/CHANGELOG b/package/CHANGELOG index 18373bae4ad..021e9a8f484 100644 --- a/package/CHANGELOG +++ b/package/CHANGELOG @@ -14,7 +14,7 @@ The rules for this file: ------------------------------------------------------------------------------ ??/??/18 tylerjereddy, richardjgowers, palnabarun, orbeckst, kain88-de, zemanj, - VOD555, davidercruz, jbarnoud, ayushsuhane, hfmull + VOD555, davidercruz, jbarnoud, ayushsuhane, hfmull, sebastien.buchoux * 0.18.1