Skip to content

Commit

Permalink
Removed dependency on kdtree flags, substituted capped_distance in cl…
Browse files Browse the repository at this point in the history 
…ass Distance based selection + point selection except CylindricalSelection
  • Loading branch information
@ayushsuhane
ayushsuhane committed Aug 12, 2018
1 parent ce277f0 commit 4904d4d
Show file tree
Hide file tree
Showing 2 changed files with 62 additions and 171 deletions.
172 changes: 39 additions & 133 deletions package/MDAnalysis/core/selection.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -239,11 +239,6 @@ class DistanceSelection(Selection):
- _apply_distmat
"""
def __init__(self):

if flags['use_KDTree_routines'] in (True, 'fast', 'always'):
self.apply = self._apply_KDTree
else:
self.apply = self._apply_distmat

self.periodic = flags['use_periodic_selections']

Expand Down Expand Up @@ -274,74 +269,23 @@ def __init__(self, parser, tokens):
self.cutoff = float(tokens.popleft())
self.sel = parser.parse_expression(self.precedence)

def _apply_KDTree(self, group):
"""KDTree based selection is about 7x faster than distmat
for typical problems.
"""
def apply(self, group):
indices = []
sel = self.sel.apply(group)
# All atoms in group that aren't in sel
sys = group[~np.in1d(group.indices, sel.indices)]

if not sys:
if not sys or not sel:
return sys[[]]

box = self.validate_dimensions(group.dimensions)
pairs = distances.capped_distance(sel.positions, sys.positions,
self.cutoff, box=box,
return_distances=False)
if pairs.size > 0:
indices = np.sort(pairs[:, 1])

cut = self.cutoff if box is not None else None
kdtree = PeriodicKDTree(box=box, leafsize=10)
kdtree.set_coords(sys.positions, cutoff=cut)
kdtree.search(sel.positions, self.cutoff)
unique_idx = np.asarray(kdtree.get_indices())

return sys[unique_idx.astype(np.int32)].unique

def _apply_distmat(self, group):
sel = self.sel.apply(group)
sys = group[~np.in1d(group.indices, sel.indices)]

box = self.validate_dimensions(group.dimensions)
dist = distances.distance_array(
sys.positions, sel.positions, box)

mask = (dist <= self.cutoff).any(axis=1)

return sys[mask].unique

def _apply_nsgrid(self, group):
sel = self.sel.apply(group)
# All atoms in group that aren't in sel
sys = group[~np.in1d(group.indices, sel.indices)]
box = self.validate_dimensions(group.dimensions)
# Box handling for no PBC case
if box is None:
pseudobox = np.zeros(6, dtype=np.float32)
all_coords = np.concatenate([group.positions, sel.positions])
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*self.cutoff:
# just enough box size so that NSGrid doesnot fails
sizefactor = 2.2*self.cutoff/boxsize
else:
sizefactor = 1.2
pseudobox[:3] = sizefactor*boxsize
pseudobox[3:] = 90.
shiftref, shiftconf = sys.positions.copy(), sel.positions.copy()
# Extra padding near the origin
shiftref -= lmin - 0.1*boxsize
shiftconf -= lmin - 0.1*boxsize
gridsearch = nsgrid.FastNS(self.cutoff, shiftref, box=pseudobox, pbc=False)
results = gridsearch.search(shiftconf)
else:
gridsearch = nsgrid.FastNS(self.cutoff, sys.positions, box=box)
results = gridsearch.search(sel.positions)

indices = np.sort(results.get_pairs()[:, 1])
return sys[indices].unique

return sys[np.asarray(indices, dtype=np.int64)].unique

class SphericalLayerSelection(DistanceSelection):
token = 'sphlayer'
Expand All @@ -352,39 +296,22 @@ def __init__(self, parser, tokens):
self.inRadius = float(tokens.popleft())
self.exRadius = float(tokens.popleft())
self.sel = parser.parse_expression(self.precedence)

def _apply_KDTree(self, group):
"""Selection using KDTree and PeriodicKDTree for aperiodic and
fully-periodic systems, respectively.
"""
sel = self.sel.apply(group)
box = self.validate_dimensions(group.dimensions)
periodic = box is not None
ref = sel.center_of_geometry(pbc=periodic)
kdtree = PeriodicKDTree(box=box)

cutoff = self.exRadius if box is not None else None
kdtree.set_coords(group.positions, cutoff=cutoff)
kdtree.search(ref, self.exRadius)
found_ExtIndices = kdtree.get_indices()
kdtree.search(ref, self.inRadius)
found_IntIndices = kdtree.get_indices()
found_indices = list(set(found_ExtIndices) - set(found_IntIndices))
return group[found_indices].unique

def _apply_distmat(self, group):

def apply(self, group):
indices = []
sel = self.sel.apply(group)
box = self.validate_dimensions(group.dimensions)
periodic = box is not None
ref = sel.center_of_geometry(pbc=periodic).reshape(1, 3).astype(
np.float32)
d = distances.distance_array(ref,
group.positions,
box=box)[0]
mask = d < self.exRadius
mask &= d > self.inRadius
pairs = distances.capped_distance(ref, group.positions, self.exRadius,
min_cutoff=self.inRadius,
box=box,
return_distances=False)
if pairs.size > 0:
indices = np.sort(pairs[:, 1])

return group[mask].unique
return group[np.asarray(indices, dtype=np.int64)].unique


class SphericalZoneSelection(DistanceSelection):
Expand All @@ -396,39 +323,27 @@ def __init__(self, parser, tokens):
self.cutoff = float(tokens.popleft())
self.sel = parser.parse_expression(self.precedence)

def _apply_KDTree(self, group):
"""Selection using KDTree and PeriodicKDTree for aperiodic and
fully-periodic systems, respectively.
"""
def apply(self, group):
indices = []
sel = self.sel.apply(group)
box = self.validate_dimensions(group.dimensions)
periodic = box is not None
ref = sel.center_of_geometry(pbc=periodic)

cut = self.cutoff if box is not None else None
kdtree = PeriodicKDTree(box=box)
kdtree.set_coords(group.positions, cutoff=cut)
kdtree.search(ref, self.cutoff)
found_indices = kdtree.get_indices()
return group[found_indices].unique
ref = sel.center_of_geometry(pbc=periodic).reshape(1, 3).astype(
np.float32)
pairs = distances.capped_distance(ref, group.positions, self.cutoff,
box=box,
return_distances=False)
if pairs.size > 0:
indices = np.sort(pairs[:, 1])

def _apply_distmat(self, group):
sel = self.sel.apply(group)
box = self.validate_dimensions(group.dimensions)
periodic = box is not None
ref = sel.center_of_geometry(pbc=periodic).reshape(1, 3).\
astype(np.float32)
d = distances.distance_array(ref,
group.positions,
box=box)[0]
idx = d < self.cutoff
return group[idx].unique
return group[np.asarray(indices, dtype=np.int64)].unique


class CylindricalSelection(Selection):
def __init__(self):
self.periodic = flags['use_periodic_selections']


def apply(self, group):
sel = self.sel.apply(group)

Expand Down Expand Up @@ -522,25 +437,16 @@ def __init__(self, parser, tokens):
self.ref = np.array([x, y, z], dtype=np.float32)
self.cutoff = float(tokens.popleft())

def _apply_KDTree(self, group):
box = self.validate_dimensions(group.dimensions)
kdtree = PeriodicKDTree(box=box)
cut = self.cutoff if box is not None else None
kdtree.set_coords(group.positions, cutoff=cut)
kdtree.search(self.ref, self.cutoff)
found_indices = kdtree.get_indices()

return group[found_indices].unique

def _apply_distmat(self, group):
ref_coor = self.ref[np.newaxis, ...]

ref_coor = np.asarray(ref_coor, dtype=np.float32)
def apply(self, group):
indices = []
box = self.validate_dimensions(group.dimensions)
pairs = distances.capped_distance(self.ref[None, :], group.positions, self.cutoff,
box=box,
return_distances=False)
if pairs.size > 0:
indices = np.sort(pairs[:, 1])

dist = distances.distance_array(group.positions, ref_coor, box)
mask = (dist <= self.cutoff).any(axis=1)
return group[mask].unique
return group[np.asarray(indices, dtype=np.int64)].unique


class AtomSelection(Selection):
Expand Down
61 changes: 23 additions & 38 deletions testsuite/MDAnalysisTests/core/test_atomselections.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -496,40 +496,14 @@ class BaseDistanceSelection(object):
Cylindrical methods don't use KDTree
"""

methods = [('kdtree', False),
('kdtree', True),
('distmat', True),
('distmat', False)]

around_methods = [('kdtree', False),
('kdtree', True),
('distmat', True),
('distmat', False),
('nsgrid', False),
('nsgrid', True)]

@staticmethod
def choosemeth(sel, meth, periodic):
"""hack in the desired apply method"""
if meth == 'kdtree':
sel.apply = sel._apply_KDTree
elif meth == 'distmat':
sel.apply = sel._apply_distmat
elif meth == 'nsgrid':
sel.apply = sel._apply_nsgrid

@pytest.mark.parametrize('periodic', (True, False))
def test_around(self, u, periodic):
sel = Parser.parse('around 5.0 resid 1', u.atoms)
if periodic:
sel.periodic = True
else:
sel.periodic = False

return sel

@pytest.mark.parametrize('meth, periodic', around_methods)
def test_around(self, u, meth, periodic):
sel = Parser.parse('around 5.0 resid 1', u.atoms)
sel = self.choosemeth(sel, meth, periodic)
result = sel.apply(u.atoms)

r1 = u.select_atoms('resid 1')
Expand All @@ -544,10 +518,14 @@ def test_around(self, u, meth, periodic):
ref.difference_update(set(r1.indices))
assert ref == set(result.indices)

@pytest.mark.parametrize('meth, periodic', methods)
def test_spherical_layer(self, u, meth, periodic):
@pytest.mark.parametrize('periodic', (True, False))
def test_spherical_layer(self, u, periodic):
sel = Parser.parse('sphlayer 2.4 6.0 resid 1', u.atoms)
sel = self.choosemeth(sel, meth, periodic)
if periodic:
sel.periodic = True
else:
sel.periodic = False

result = sel.apply(u.atoms)

r1 = u.select_atoms('resid 1')
Expand All @@ -558,10 +536,14 @@ def test_spherical_layer(self, u, meth, periodic):

assert ref == set(result.indices)

@pytest.mark.parametrize('meth, periodic', methods)
def test_spherical_zone(self, u, meth, periodic):
@pytest.mark.parametrize('periodic', (True, False))
def test_spherical_zone(self, u, periodic):
sel = Parser.parse('sphzone 5.0 resid 1', u.atoms)
sel = self.choosemeth(sel, meth, periodic)
if periodic:
sel.periodic = True
else:
sel.periodic = False

result = sel.apply(u.atoms)

r1 = u.select_atoms('resid 1')
Expand All @@ -572,10 +554,13 @@ def test_spherical_zone(self, u, meth, periodic):

assert ref == set(result.indices)

@pytest.mark.parametrize('meth, periodic', methods)
def test_point(self, u, meth, periodic):
@pytest.mark.parametrize('periodic', (True, False))
def test_point(self, u, periodic):
sel = Parser.parse('point 5.0 5.0 5.0 3.0', u.atoms)
sel = self.choosemeth(sel, meth, periodic)
if periodic:
sel.periodic = True
else:
sel.periodic = False
result = sel.apply(u.atoms)

box = u.dimensions if periodic else None
Expand Down

0 comments on commit 4904d4d

Please sign in to comment.