Source code for sknano.core.atoms._kdtree_atoms
# -*- coding: utf-8 -*-
"""
===============================================================================
Mixin classes for KDTree analysis (:mod:`sknano.core.atoms._kdtree_atoms`)
===============================================================================
.. currentmodule:: sknano.core.atoms._kdtree_atoms
"""
from __future__ import absolute_import, division, print_function
from __future__ import unicode_literals
__docformat__ = 'restructuredtext en'
import numbers
import warnings
warnings.filterwarnings('ignore', "Mean of empty slice.")
warnings.filterwarnings('ignore',
'invalid value encountered in double_scalars')
import numpy as np
try:
from scipy.spatial import KDTree
except ImportError:
raise ImportError('Install scipy version >= 0.13.0 to allow '
'nearest-neighbor queries between atoms.')
import sknano.core.atoms
from ._bonds import Bond, Bonds
__all__ = ['KDTreeAtomMixin', 'KDTreeAtomsMixin']
[docs]class KDTreeAtomMixin:
"""Mixin Atom class for KDTree analysis."""
@property
def NN(self):
"""Nearest-neighbor `Atoms`."""
try:
return self._NN
except AttributeError:
return None
@NN.setter
def NN(self, value):
"""Set nearest-neighbor `Atoms`."""
if not isinstance(value, sknano.core.atoms.Atoms):
raise TypeError('Expected an `Atoms` object.')
self._NN = value
@NN.deleter
def NN(self):
del self._NN
@property
def bonds(self):
"""Return atom `Bonds` instance."""
try:
return Bonds([Bond(self, nn) for nn in self.NN])
except (AttributeError, TypeError):
return Bonds()
@property
def neighbors(self):
"""Neighbor atoms."""
return self._neighbors
@neighbors.setter
def neighbors(self, value):
if not isinstance(value, sknano.core.atoms.Atoms):
raise TypeError('Expected an `Atoms` object.')
self._neighbors = value
@property
def neighbor_distances(self):
"""Neighbor atom distances."""
return self._neighbor_distances
@neighbor_distances.setter
def neighbor_distances(self, value):
self._neighbor_distances = np.asarray(value)
[docs] def todict(self):
super_dict = super().todict()
super_dict.update(dict(CN=self.CN, NN=self.NN))
return super_dict
[docs]class KDTreeAtomsMixin:
"""Mixin Atoms class for KDTree analysis."""
@property
def kNN(self):
"""Number of nearest neighbors to return when querying the kd-tree."""
return self._kNN
@kNN.setter
def kNN(self, value):
"""Set maximum number of nearest neighbors to return when querying
the kd-tree."""
if not isinstance(value, numbers.Number):
raise TypeError('Expected an integer >= 0')
self._kNN = self.kwargs['kNN'] = int(value)
@property
def NNrc(self):
"""Nearest neighbor radius cutoff."""
return self._NNrc
@NNrc.setter
def NNrc(self, value):
if not (isinstance(value, numbers.Number) and value >= 0):
raise TypeError('Expected a real number greater >= 0')
self._NNrc = self.kwargs['NNrc'] = value
@property
def atom_tree(self):
""":class:`~scipy:scipy.spatial.KDTree` of :attr:`~XAtoms.coords.`"""
try:
return KDTree(self.coords)
except ValueError:
return None
@property
def neighbors(self):
"""Return array of neighbor atoms."""
return np.asarray([atom.neighbors for atom in self])
@property
def distances(self):
"""Neighbor atoms distances."""
return self._distances
@distances.setter
def distances(self, value):
self._distances = np.asarray(value)
@property
def neighbor_distances(self):
distances = []
# [distances.extend(atom.neighbor_distances.tolist()) for atom in self]
[distances.extend(atom.neighbors.distances.tolist()) for atom in self]
return np.asarray(distances)
@property
def nearest_neighbors(self):
"""Return array of nearest-neighbor atoms for each `KDTAtom`."""
# self._update_neighbors()
return np.asarray([atom.NN for atom in self])
[docs] def query_atom_tree(self, k=16, eps=0, p=2, rc=np.inf):
"""Query atom tree for nearest neighbors distances and indices.
Parameters
----------
k : integer
The number of nearest neighbors to return.
eps : nonnegative float
Return approximate nearest neighbors; the kth returned value
is guaranteed to be no further than (1+eps) times the
distance to the real kth nearest neighbor.
p : float, 1<=p<=infinity
Which Minkowski p-norm to use.
1 is the sum-of-absolute-values "Manhattan" distance
2 is the usual Euclidean distance
infinity is the maximum-coordinate-difference distance
rc : nonnegative float
Radius cutoff. Return only neighbors within this distance.
This is used to prune tree searches, so if you are doing a series
of nearest-neighbor queries, it may help to supply the distance to
the nearest neighbor of the most recent point.
Returns
-------
d : array of floats
The distances to the nearest neighbors, sorted by distance.
i : array of integers
The locations of the neighbors in self.atom_tree.data. `i` is the
same shape as `d`.
"""
atom_tree = self.atom_tree
if atom_tree is not None:
d, i = atom_tree.query(self.coords, k=k+1, eps=eps, p=p,
distance_upper_bound=rc)
return d[:, 1:], i[:, 1:]
[docs] def query_ball_point(self, pts, r, p=2.0, eps=0):
"""Find all `Atoms` within distance `r` of point(s) `pts`.
Parameters
----------
pts : :class:`~sknano.core.math.Point`
The :class:`~sknano.core.math.Point` or
:class:`~sknano.core.math.Points` to search for neighbors of.
r : positive :class:`~python:float`
The radius of :class:`~sknano.core.atoms.KDTAtoms` to return
p : float, 1<=p<=infinity
Which Minkowski p-norm to use.
1 is the sum-of-absolute-values "Manhattan" distance
2 is the usual Euclidean distance
infinity is the maximum-coordinate-difference distance
eps : nonnegative :class:`~python:float`, optional
Approximate search.
Returns
-------
:class:`~sknano.core.atoms.KDTAtoms`
"""
atom_tree = self.atom_tree
if atom_tree is not None:
NNi = atom_tree.query_ball_point(pts, r, p=p, eps=eps)
return self.__class__(atoms=np.asarray(self)[NNi].tolist(),
**self.kwargs)
[docs] def update_attrs(self):
"""Update :class:`KDTAtom`\ s attributes."""
self.__update_neighbors()
self.__update_bonds()
[docs] def update_neighbors(self):
"""Update :attr:`KDTAtom.NN`."""
try:
NNd, NNi = self.query_atom_tree(k=self.kNN, rc=self.NNrc)
for j, atom in enumerate(self):
# atom.neighbors = self.__class__(**self.kwargs)
# atom.neighbors = NeighborAtoms()
# [atom.neighbors.append(self[NNi[j][k]])
# for k, d in enumerate(NNd[j]) if d <= self.NNrc]
# atom.neighbors = \
# NeighborAtoms([self[NNi[j][k]] for k, d in
# enumerate(NNd[j]) if d <= self.NNrc],
# casttype=False)
atom.neighbors = \
self.__class__([self[NNi[j][k]] for k, d in
enumerate(NNd[j]) if d <= self.NNrc],
casttype=False, **self.kwargs)
atom.neighbor_distances = \
[NNd[j][k] for k, d in enumerate(NNd[j]) if d <= self.NNrc]
atom.neighbors.distances = atom.neighbor_distances
atom.NN = atom.neighbors
except ValueError:
pass
__update_neighbors = update_neighbors
[docs] def update_bonds(self):
"""Update :attr:`KDTAtom.bonds`."""
self.bonds.clear()
[self.bonds.extend(atom.bonds) for atom in self]
__update_bonds = update_bonds
[docs] def neighbor_counts(self, r):
return np.asarray([KDTree([atom.r]).count_neighbors(
self.filtered(self.ids != atom.id).atom_tree,
np.asarray(r)) for atom in self])