1# -*- coding: utf-8 -*- 2# MolMod is a collection of molecular modelling tools for python. 3# Copyright (C) 2007 - 2019 Toon Verstraelen <Toon.Verstraelen@UGent.be>, Center 4# for Molecular Modeling (CMM), Ghent University, Ghent, Belgium; all rights 5# reserved unless otherwise stated. 6# 7# This file is part of MolMod. 8# 9# MolMod is free software; you can redistribute it and/or 10# modify it under the terms of the GNU General Public License 11# as published by the Free Software Foundation; either version 3 12# of the License, or (at your option) any later version. 13# 14# MolMod is distributed in the hope that it will be useful, 15# but WITHOUT ANY WARRANTY; without even the implied warranty of 16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 17# GNU General Public License for more details. 18# 19# You should have received a copy of the GNU General Public License 20# along with this program; if not, see <http://www.gnu.org/licenses/> 21# 22# -- 23"""Efficiently compute all distances below a cutoff 24 25 Binning is a useful technique for efficiently calculating all distances 26 between a number of coordinates when you are only interested in the distances 27 below a given cutoff. The algorithm consists of two major steps: 28 29 1) Divide the given set of coordinates into bins on a regular grid 30 2) Calculate the distances (or other useful things) between coordinates in 31 neighboring bins. 32""" 33 34 35from __future__ import division 36 37from builtins import range 38import numpy as np 39 40from molmod.unit_cells import UnitCell 41 42 43__all__ = ["PairSearchIntra", "PairSearchInter"] 44 45 46class Binning(object): 47 """Division of coordinates in regular bins""" 48 def __init__(self, coordinates, cutoff, grid_cell, integer_cell=None): 49 """Initialize a Binning object 50 51 Arguments: 52 | ``coordinates`` -- a Nx3 numpy array with coordinates to be 53 triaged into bins 54 | ``cutoff`` -- The maximum distance between coordinates pairs. 55 This affects the variable self.neighbor_bins, 56 which is a set of relative integer bin coordinates 57 that lie within the cuttof of the central bin. 58 | ``grid_cell`` -- A unit cell object specifying the size and 59 shape of the bins 60 61 Optional argument: 62 | ``integer_cell`` -- the periodicity of the system in terms if 63 integer grid cells. 64 """ 65 self.grid_cell = grid_cell 66 self.integer_cell = integer_cell 67 68 # setup the bins 69 self._bins = {} 70 71 fractional = grid_cell.to_fractional(coordinates) 72 for i in range(len(coordinates)): 73 key = tuple(fractional[i].astype(int)) 74 if integer_cell is not None: 75 key = self.wrap_key(key) 76 bin = self._bins.get(key) 77 if bin is None: 78 bin = [] 79 self._bins[key] = bin 80 bin.append((i, coordinates[i])) 81 82 # compute the neigbouring bins within the cutoff 83 if self.integer_cell is None: 84 self.neighbor_indexes = grid_cell.get_radius_indexes(cutoff) 85 else: 86 max_ranges = np.diag(self.integer_cell.matrix).astype(int) 87 max_ranges[True^self.integer_cell.active] = -1 88 self.neighbor_indexes = grid_cell.get_radius_indexes(cutoff, max_ranges) 89 90 def __iter__(self): 91 """Iterate over (key,bin) pairs""" 92 return iter(self._bins.items()) 93 94 def iter_surrounding(self, center_key): 95 """Iterate over all bins surrounding the given bin""" 96 for shift in self.neighbor_indexes: 97 key = tuple(np.add(center_key, shift).astype(int)) 98 if self.integer_cell is not None: 99 key = self.wrap_key(key) 100 bin = self._bins.get(key) 101 if bin is not None: 102 yield key, bin 103 104 def wrap_key(self, key): 105 """Translate the key into the central cell 106 107 This method is only applicable in case of a periodic system. 108 """ 109 return tuple(np.round( 110 self.integer_cell.shortest_vector(key) 111 ).astype(int)) 112 113 114class PairSearchBase(object): 115 """Base class for :class:`PairSearchIntra` and :class:`PairSearchInter`""" 116 def _setup_grid(self, cutoff, unit_cell, grid): 117 """Choose a proper grid for the binning process""" 118 if grid is None: 119 # automatically choose a decent grid 120 if unit_cell is None: 121 grid = cutoff/2.9 122 else: 123 # The following would be faster, but it is not reliable 124 # enough yet. 125 #grid = unit_cell.get_optimal_subcell(cutoff/2.0) 126 divisions = np.ceil(unit_cell.spacings/cutoff) 127 divisions[divisions<1] = 1 128 grid = unit_cell/divisions 129 130 if isinstance(grid, float): 131 grid_cell = UnitCell(np.array([ 132 [grid, 0, 0], 133 [0, grid, 0], 134 [0, 0, grid] 135 ])) 136 elif isinstance(grid, UnitCell): 137 grid_cell = grid 138 else: 139 raise TypeError("Grid must be None, a float or a UnitCell instance.") 140 141 if unit_cell is not None: 142 # The columns of integer_matrix are the unit cell vectors in 143 # fractional coordinates of the grid cell. 144 integer_matrix = grid_cell.to_fractional(unit_cell.matrix.transpose()).transpose() 145 if abs((integer_matrix - np.round(integer_matrix))*self.unit_cell.active).max() > 1e-6: 146 raise ValueError("The unit cell vectors are not an integer linear combination of grid cell vectors.") 147 integer_matrix = integer_matrix.round() 148 integer_cell = UnitCell(integer_matrix, unit_cell.active) 149 else: 150 integer_cell = None 151 152 return grid_cell, integer_cell 153 154 155class PairSearchIntra(PairSearchBase): 156 """Iterator over all pairs of coordinates with a distance below a cutoff. 157 158 Example usage:: 159 160 coordinates = np.random.uniform(0,10,(10,3)) 161 for i, j, delta, distance in PairSearchIntra(coordinates, 2.5): 162 print i, j, distance 163 164 Note that for periodic systems the minimum image convention is applied. 165 """ 166 167 def __init__(self, coordinates, cutoff, unit_cell=None, grid=None): 168 """ 169 Arguments: 170 | ``coordinates`` -- A Nx3 numpy array with Cartesian coordinates 171 | ``radius`` -- The cutoff radius for the pair distances. 172 Distances larger than the cutoff will be neglected 173 in the pair search. 174 175 Optional arguments: 176 | ``unit_cell`` -- Specifies the periodic boundary conditions 177 | ``grid`` -- Specification of the grid, can be a floating point 178 number which will result in cubic bins with edge length 179 equal to the given number. Otherwise a UnitCell object 180 can be specified to construct non-cubic bins. In the 181 latter case and when a unit_cell is given, the unit cell 182 vectors must be integer linear combinations of the grid 183 cell vectors (for those directions that are active in 184 the unit cell). If this is not the case, a ValueError is 185 raised. 186 187 The default value of grid depends on other parameters: 188 189 1) When no unit cell is given, it is equal to cutoff/2.9. 190 2) When a unit cell is given, the grid cell is as close to cubic 191 as possible, with spacings below cutoff/2 that are integer 192 divisions of the unit cell spacings 193 """ 194 self.cutoff = cutoff 195 self.unit_cell = unit_cell 196 grid_cell, integer_cell = self._setup_grid(cutoff, unit_cell, grid) 197 self.bins = Binning(coordinates, cutoff, grid_cell, integer_cell) 198 199 def __iter__(self): 200 """Iterate over all pairs with a distance below the cutoff""" 201 for key0, bin0 in self.bins: 202 for key1, bin1 in self.bins.iter_surrounding(key0): 203 for i0, c0 in bin0: 204 for i1, c1 in bin1: 205 if i1 >= i0: 206 continue 207 delta = c1 - c0 208 if self.unit_cell is not None: 209 delta = self.unit_cell.shortest_vector(delta) 210 distance = np.linalg.norm(delta) 211 if distance <= self.cutoff: 212 yield i0, i1, delta, distance 213 214class PairSearchInter(PairSearchBase): 215 """Iterator over all pairs of coordinates with a distance below a cutoff. 216 217 Example usage:: 218 219 coordinates0 = np.random.uniform(0,10,(10,3)) 220 coordinates1 = np.random.uniform(0,10,(10,3)) 221 for i, j, delta, distance in PairSearchInter(coordinates0, coordinates1, 2.5): 222 print i, j, distance 223 224 Note that for periodic systems the minimum image convention is applied. 225 """ 226 227 def __init__(self, coordinates0, coordinates1, cutoff, unit_cell=None, grid=None): 228 """ 229 Arguments: 230 | ``coordinates0`` -- A Nx3 numpy array with Cartesian coordinates 231 | ``coordinates1`` -- A Nx3 numpy array with Cartesian coordinates 232 | ``radius`` -- The cutoff radius for the pair distances. 233 Distances larger than the cutoff will be neglected 234 in the pair search. 235 236 Optional arguments: 237 | ``unit_cell`` -- Specifies the periodic boundary conditions 238 | ``grid`` -- Specification of the grid, can be a floating point 239 number which will result in cubic bins with edge length 240 equal to the given number. Otherwise a UnitCell object 241 can be specified to construct non-cubic bins. In the 242 latter case and when a unit_cell is given, the unit cell 243 vectors must be integer linear combinations of the grid 244 cell vectors (for those directions that are active in 245 the unit cell). If this is not the case, a ValueError is 246 raised. 247 248 The default value of grid depends on other parameters: 249 1) When no unit cell is given, it is equal to cutoff/2.9. 250 2) When a unit cell is given, the grid cell is as close to cubic 251 as possible, with spacings below cutoff/2 that are integer 252 divisions of the unit cell spacings 253 """ 254 self.cutoff = cutoff 255 self.unit_cell = unit_cell 256 grid_cell, integer_cell = self._setup_grid(cutoff, unit_cell, grid) 257 self.bins0 = Binning(coordinates0, cutoff, grid_cell, integer_cell) 258 self.bins1 = Binning(coordinates1, cutoff, grid_cell, integer_cell) 259 260 def __iter__(self): 261 """Iterate over all pairs with a distance below the cutoff""" 262 for key0, bin0 in self.bins0: 263 for key1, bin1 in self.bins1.iter_surrounding(key0): 264 for i0, c0 in bin0: 265 for i1, c1 in bin1: 266 delta = c1 - c0 267 if self.unit_cell is not None: 268 delta = self.unit_cell.shortest_vector(delta) 269 distance = np.linalg.norm(delta) 270 if distance <= self.cutoff: 271 yield i0, i1, delta, distance 272