def test_tric_PBC(self):
from MDAnalysis.lib.distances import apply_PBC
U = MDAnalysis.Universe(TRIC)
atoms = U.atoms.coordinates()
box1 = U.dimensions
box2 = MDAnalysis.coordinates.core.triclinic_vectors(box1)
#print box2
#print box2.shape
def numpy_PBC(coords, box):
coords -= np.array(
[box[2] * val for val in np.floor(coords[:, 2] / box[2][2])])
coords -= np.array(
[box[1] * val for val in np.floor(coords[:, 1] / box[1][1])])
coords -= np.array(
[box[0] * val for val in np.floor(coords[:, 0] / box[0][0])])
return coords
cyth1 = apply_PBC(atoms, box1)
cyth2 = apply_PBC(atoms, box2)
reference = numpy_PBC(atoms, box2)
assert_almost_equal(
cyth1,
reference,
self.prec,
err_msg="Triclinic apply_PBC failed comparison with np")
assert_almost_equal(
cyth2,
reference,
self.prec,
err_msg="Trlclinic apply_PBC failed comparison with np")
def test_tric_PBC(self):
from MDAnalysis.lib.distances import apply_PBC
U = MDAnalysis.Universe(TRIC)
atoms = U.atoms.coordinates()
box1 = U.dimensions
box2 = MDAnalysis.coordinates.core.triclinic_vectors(box1)
#print box2
#print box2.shape
def numpy_PBC(coords, box):
coords -= np.array([box[2] * val for val in np.floor(coords[:, 2] / box[2][2])])
coords -= np.array([box[1] * val for val in np.floor(coords[:, 1] / box[1][1])])
coords -= np.array([box[0] * val for val in np.floor(coords[:, 0] / box[0][0])])
return coords
cyth1 = apply_PBC(atoms, box1, backend=self.backend)
cyth2 = apply_PBC(atoms, box2, backend=self.backend)
reference = numpy_PBC(atoms, box2)
assert_almost_equal(cyth1, reference, self.prec,
err_msg="Triclinic apply_PBC failed comparison with np")
assert_almost_equal(cyth2, reference, self.prec,
err_msg="Trlclinic apply_PBC failed comparison with np")
def test_tric_PBC(self, backend):
from MDAnalysis.lib.distances import apply_PBC
U = MDAnalysis.Universe(TRIC)
atoms = U.atoms.positions
box1 = U.dimensions
box2 = MDAnalysis.coordinates.core.triclinic_vectors(box1)
def numpy_PBC(coords, box):
# move to fractional coordinates
fractional = MDAnalysis.lib.distances.transform_RtoS(coords, box)
# move fractional coordinates to central cell
fractional -= np.floor(fractional)
# move back to real coordinates
return MDAnalysis.lib.distances.transform_StoR(fractional, box)
cyth1 = apply_PBC(atoms, box1, backend=backend)
cyth2 = apply_PBC(atoms, box2, backend=backend)
reference = numpy_PBC(atoms, box2)
assert_almost_equal(
cyth1,
reference,
decimal=4,
err_msg="Triclinic apply_PBC failed comparison with np")
assert_almost_equal(
cyth2,
reference,
decimal=4,
err_msg="Triclinic apply_PBC failed comparison with np")
box = np.array([10, 7, 3, 45, 60, 90], dtype=np.float32)
r = np.array([[5.75, 0.36066014, 0.75]], dtype=np.float32)
r_in_cell = MDAnalysis.lib.distances.apply_PBC(r, box)[0]
assert_almost_equal([5.75, 7.3606596, 0.75], r_in_cell, self.prec)
def test_ortho_PBC(self):
from MDAnalysis.lib.distances import apply_PBC
U = MDAnalysis.Universe(PSF, DCD)
atoms = U.atoms.coordinates()
box1 = np.array([2.5, 2.5, 3.5], dtype=np.float32)
box2 = np.array([2.5, 2.5, 3.5, 90.0, 90.0, 90.0], dtype=np.float32)
cyth1 = apply_PBC(atoms, box1)
cyth2 = apply_PBC(atoms, box2)
reference = atoms - np.floor(atoms / box1) * box1
assert_almost_equal(cyth1, reference, self.prec, err_msg="Ortho apply_PBC #1 failed comparison with np")
assert_almost_equal(cyth2, reference, self.prec, err_msg="Ortho apply_PBC #2 failed comparison with np")
def test_ortho_PBC(self):
from MDAnalysis.lib.distances import apply_PBC
U = MDAnalysis.Universe(PSF, DCD)
atoms = U.atoms.coordinates()
box1 = np.array([2.5, 2.5, 3.5], dtype=np.float32)
box2 = np.array([2.5, 2.5, 3.5, 90., 90., 90.], dtype=np.float32)
cyth1 = apply_PBC(atoms, box1, backend=self.backend)
cyth2 = apply_PBC(atoms, box2, backend=self.backend)
reference = atoms - np.floor(atoms / box1) * box1
assert_almost_equal(cyth1, reference, self.prec,
err_msg="Ortho apply_PBC #1 failed comparison with np")
assert_almost_equal(cyth2, reference, self.prec,
err_msg="Ortho apply_PBC #2 failed comparison with np")
def search(self, centers, radius):
"""Search all points within radius of centers and its periodic images.
Wrapping of center coordinates is enforced to enable comparison to
wrapped coordinates of points in the tree.
Parameter
---------
centers: array_like (N,3)
origins around which to search for neighbors
radius: float
maximum distance around which to search for neighbors. The search
radius is half the smallest periodicity if radius exceeds this value
"""
if not self.built:
raise RuntimeError('Unbuilt tree. Run tree.set_coords(...) first')
centers = np.asarray(centers)
if centers.shape == (self.dim, ):
centers = centers.reshape((1, self.dim))
wrapped_centers = apply_PBC(centers, self.box)
self._indices = set() # clear previous search
# possible loop for parallel execution (benchmark before!)
for c in itertools.chain(wrapped_centers,
self.find_images(wrapped_centers, radius)):
self.kdt.search_center_radius(c, radius)
n = self.kdt.get_count()
if n:
new_indices = np.empty(n, int)
self.kdt.get_indices(new_indices)
self._indices.update(new_indices)
self._indices = sorted(list(self._indices))
def search_tree(self, centers, radius):
"""
Searches all the pairs within `radius` between `centers`
and ``coords``
``coords`` are the already initialized coordinates in the tree
during :meth:`set_coords`.
``centers`` are wrapped around the primary unit cell
if PBC is desired. Minimum image convention (PBC) is
activated if the `box` argument is provided during
class initialization
Parameters
----------
centers: array_like (N,3)
coordinate array to search for neighbors
radius: float
maximum distance to search for neighbors.
Returns
-------
pairs : array
all the pairs between ``coords`` and ``centers``
Note
----
This method constructs another tree from the ``centers``
and queries the previously built tree (built in
:meth:`set_coords`)
"""
if not self._built:
raise RuntimeError('Unbuilt tree. Run tree.set_coords(...)')
centers = np.asarray(centers)
if centers.shape == (self.dim, ):
centers = centers.reshape((1, self.dim))
# Sanity check
if self.pbc:
if self.cutoff < radius:
raise RuntimeError('Set cutoff greater or equal to the radius.')
# Bring all query points to the central cell
wrapped_centers = apply_PBC(centers, self.box)
other_tree = cKDTree(wrapped_centers, leafsize=self.leafsize)
pairs = other_tree.query_ball_tree(self.ckdt, radius)
pairs = np.array([[i, j] for i, lst in enumerate(pairs) for j in lst],
dtype=np.int64)
if pairs.size > 0:
pairs[:, 1] = undo_augment(pairs[:, 1],
self.mapping,
len(self.coords))
else:
other_tree = cKDTree(centers, leafsize=self.leafsize)
pairs = other_tree.query_ball_tree(self.ckdt, radius)
pairs = np.array([[i, j] for i, lst in enumerate(pairs) for j in lst],
dtype=np.int64)
if pairs.size > 0:
pairs = unique_rows(pairs)
return pairs
def set_coords(self, coords, cutoff=None):
"""Constructs KDTree from the coordinates
Wrapping of coordinates to the primary unit cell is enforced
before any distance evaluations. If periodic boundary conditions
are enabled, then duplicate particles are generated in the
vicinity of the box. An additional array `mapping` is also
generated which can be later used to trace the origin of
duplicate particle coordinates.
For non-periodic calculations, cutoff should not be provided
the parameter is only required for periodic calculations.
Parameters
----------
coords: array_like
Coordinate array of shape ``(N, 3)`` for N atoms.
cutoff: float
Specified cutoff distance to create duplicate images
Typically equivalent to the desired search radius
or the maximum of the desired cutoff radius. Relevant images
corresponding to every atom which lies
within ``cutoff`` distance from either of the box boundary
will be generated.
See Also
--------
MDAnalysis.lib.distances.augment_coordinates
"""
# If no cutoff distance is provided but PBC aware
if self.pbc and (cutoff is None):
raise RuntimeError('Provide a cutoff distance'
' with tree.set_coords(...)')
# set coords dtype to float32
# augment coordinates will work only with float32
coords = np.asarray(coords, dtype=np.float32)
if self.pbc:
self.cutoff = cutoff
# Bring the coordinates in the central cell
self.coords = apply_PBC(coords, self.box)
# generate duplicate images
self.aug, self.mapping = augment_coordinates(self.coords,
self.box,
self.cutoff)
# Images + coords
self.all_coords = np.concatenate([self.coords, self.aug])
self.ckdt = cKDTree(self.all_coords, leafsize=self.leafsize)
else:
# if cutoff distance is provided for non PBC calculations
if cutoff is not None:
raise RuntimeError('Donot provide cutoff distance for'
' non PBC aware calculations')
self.coords = coords
self.ckdt = cKDTree(self.coords, self.leafsize)
self._built = True
def set_coords(self, coords, cutoff=None):
"""Constructs KDTree from the coordinates
Wrapping of coordinates to the primary unit cell is enforced
before any distance evaluations. If periodic boundary conditions
are enabled, then duplicate particles are generated in the
vicinity of the box. An additional array `mapping` is also
generated which can be later used to trace the origin of
duplicate particle coordinates.
For non-periodic calculations, cutoff should not be provided
the parameter is only required for periodic calculations.
Parameters
----------
coords: array_like
Coordinate array of shape ``(N, 3)`` for N atoms.
cutoff: float
Specified cutoff distance to create duplicate images
Typically equivalent to the desired search radius
or the maximum of the desired cutoff radius. Relevant images
corresponding to every atom which lies
within ``cutoff`` distance from either of the box boundary
will be generated.
See Also
--------
MDAnalysis.lib.distances.augment_coordinates
"""
# If no cutoff distance is provided but PBC aware
if self.pbc and (cutoff is None):
raise RuntimeError('Provide a cutoff distance'
' with tree.set_coords(...)')
# set coords dtype to float32
# augment coordinates will work only with float32
coords = np.asarray(coords, dtype=np.float32)
if self.pbc:
self.cutoff = cutoff
# Bring the coordinates in the central cell
self.coords = apply_PBC(coords, self.box)
# generate duplicate images
self.aug, self.mapping = augment_coordinates(
self.coords, self.box, self.cutoff)
# Images + coords
self.all_coords = np.concatenate([self.coords, self.aug])
self.ckdt = cKDTree(self.all_coords, leafsize=self.leafsize)
else:
# if cutoff distance is provided for non PBC calculations
if cutoff is not None:
raise RuntimeError('Donot provide cutoff distance for'
' non PBC aware calculations')
self.coords = coords
self.ckdt = cKDTree(self.coords, self.leafsize)
self._built = True
def _srankSlow(self, pbc=False):
"""
Calculate the internal product between bond vector and z-axis.
The z-axis is defined here as (0,0,1) by default
Parameters
----------
self: object
TopologyGroup where each element correspond to the coordinate
of a tuple of atoms belonging to a bond type
pbc : bool
apply periodic boundary conditions when calculating bond vectors
this is important when connecting vectors between atoms might
require minimum image convention
Returns
-------
cosine square : ndarray
"""
if not self.btype == 'bond':
raise TypeError("TopologyGroup is not of type 'bond'")
# define z vector
z = np.array([0, 0 ,1])
if pbc:
# box dimensions
box = self._ags[0].dimensions
# get positions of atoms to primary box
a = apply_PBC(self._ags[0].positions, box)
b = apply_PBC(self._ags[1].positions, box)
# calculate bond vector for all pair of atoms
v = b - a
# calculate angle for all bonds
ang = np.array([_angle(i,z) for i in v])
# return angle
return ang
else:
v = self._ags[1].positions - self._ags[0].positions
ang = np.array([_angle(i,z) for i in v])
return ang
def test_undoaugment(b, qres):
radius = 1.5
q = transform_StoR(np.array(qres[0], dtype=np.float32), b)
if q.shape == (3, ):
q = q.reshape((1, 3))
q = apply_PBC(q, b)
aug, mapping = augment_coordinates(q, b, radius)
for idx, val in enumerate(aug):
imageid = np.asarray([len(q) + idx], dtype=np.intp)
assert_equal(mapping[idx], undo_augment(imageid, mapping, len(q))[0])
def transformation(ts):
box_half = 0.5 * ts.dimensions[:3]
for i in range(iterations):
# wrap atoms back into box
pos_center = apply_PBC(ts.positions[ix], ts.dimensions)
# get center of mass
com = np.sum(pos_center * weights, 0) / sum_weights
# com = self.ag.center_of_mass(pbc=False)
# center around the group
ts.positions = apply_PBC(ts.positions, ts.dimensions)
ts.positions -= com - box_half
ts.positions = apply_PBC(ts.positions, ts.dimensions)
# move center of atoms to origin
#ts.positions -= box_half
if center:
ts.positions -= box_half
return ts
def test_undoaugment(b, qres):
radius = 1.5
q = transform_StoR(np.array(qres[0], dtype=np.float32), b)
if q.shape == (3, ):
q = q.reshape((1, 3))
q = apply_PBC(q, b)
aug, mapping = augment_coordinates(q, b, radius)
for idx, val in enumerate(aug):
imageid = np.asarray([len(q) + idx], dtype=np.int64)
assert_equal(mapping[idx], undo_augment(imageid, mapping, len(q))[0])
def set_coords(self, coords):
"""
Add coordinates of the points. Wrapping of coordinates to the central
cell is enforced along the periodic axes.
Parameters
----------
coords: array_like
Positions of points, shape=(N, 3) for N atoms.
"""
_check_array(np.asanyarray(coords), 'coords')
wrapped_data = apply_PBC(coords, self.box)
self.kdt.set_data(wrapped_data)
self.built = 1
def test_augment(b, q, res):
radius = 1.5
q = transform_StoR(np.array(q, dtype=np.float32), b)
if q.shape == (3, ):
q = q.reshape((1, 3))
q = apply_PBC(q, b)
aug, mapping = augment_coordinates(q, b, radius)
if aug.size > 0:
aug = np.sort(aug, axis=0)
else:
aug = list()
if len(res) > 0:
cs = transform_StoR(np.array(res, dtype=np.float32), b)
cs = np.sort(cs, axis=0)
else:
cs = list()
assert_almost_equal(aug, cs, decimal=5)
def search(self, centers, radius):
"""Search all points within radius from centers and their periodic images.
All the centers coordinates are wrapped around the central cell
to enable distance evaluations from points in the tree
and their images.
Parameters
----------
centers: array_like (N,3)
coordinate array to search for neighbors
radius: float
maximum distance to search for neighbors.
"""
if not self._built:
raise RuntimeError('Unbuilt tree. Run tree.set_coords(...)')
centers = np.asarray(centers)
if centers.shape == (self.dim, ):
centers = centers.reshape((1, self.dim))
# Sanity check
if self.pbc:
if self.cutoff < radius:
raise RuntimeError('Set cutoff greater or equal to the radius.')
# Bring all query points to the central cell
wrapped_centers = apply_PBC(centers, self.box)
indices = list(self.ckdt.query_ball_point(wrapped_centers,
radius))
self._indices = np.array(list(
itertools.chain.from_iterable(indices)),
dtype=np.int64)
if self._indices.size > 0:
self._indices = undo_augment(self._indices,
self.mapping,
len(self.coords))
else:
wrapped_centers = np.asarray(centers)
indices = list(self.ckdt.query_ball_point(wrapped_centers,
radius))
self._indices = np.array(list(
itertools.chain.from_iterable(indices)),
dtype=np.int64)
self._indices = np.asarray(unique_int_1d(self._indices))
return self._indices
def test_augment(b, q, res):
radius = 1.5
q = transform_StoR(np.array(q, dtype=np.float32), b)
if q.shape == (3, ):
q = q.reshape((1, 3))
q = apply_PBC(q, b)
aug, mapping = augment_coordinates(q, b, radius)
if aug.size > 0:
aug = np.sort(aug, axis=0)
else:
aug = list()
if len(res) > 0:
cs = transform_StoR(np.array(res, dtype=np.float32), b)
cs = np.sort(cs, axis=0)
else:
cs = list()
assert_almost_equal(aug, cs, decimal=5)
def search(self, centers, radius):
"""Search all points within radius from centers and their periodic images.
All the centers coordinates are wrapped around the central cell
to enable distance evaluations from points in the tree
and their images.
Parameters
----------
centers: array_like (N,3)
coordinate array to search for neighbors
radius: float
maximum distance to search for neighbors.
"""
if not self._built:
raise RuntimeError('Unbuilt tree. Run tree.set_coords(...)')
centers = np.asarray(centers)
if centers.shape == (self.dim, ):
centers = centers.reshape((1, self.dim))
# Sanity check
if self.pbc:
if self.cutoff < radius:
raise RuntimeError(
'Set cutoff greater or equal to the radius.')
# Bring all query points to the central cell
wrapped_centers = apply_PBC(centers, self.box)
indices = list(self.ckdt.query_ball_point(wrapped_centers, radius))
self._indices = np.array(list(
itertools.chain.from_iterable(indices)),
dtype=np.int64)
if self._indices.size > 0:
self._indices = undo_augment(self._indices, self.mapping,
len(self.coords))
else:
wrapped_centers = np.asarray(centers)
indices = list(self.ckdt.query_ball_point(wrapped_centers, radius))
self._indices = np.array(list(
itertools.chain.from_iterable(indices)),
dtype=np.int64)
self._indices = np.asarray(unique_int_1d(self._indices))
return self._indices
def test_find_images(b, qcs):
"""
Test the generation of images for a given query vector and type of box.
Parameters
----------
b : list
MDAnalysis dimensions like list
qns : tuple
a query point and a list of expected neighbors.
"""
b = np.array(b, dtype=np.float32)
q = transform_StoR(np.array(qcs[0], dtype=np.float32), b)
tree = PeriodicKDTree(b)
coords = transform_StoR(f_dataset, b)
tree.set_coords(coords) # Input real space coordinates
cs = np.sort(transform_StoR(np.array(qcs[1], dtype=np.float32), b), axis=0)
q_wrapped = apply_PBC(q.reshape((1, 3)), b)
found_images = np.sort(tree.find_images(q_wrapped, radius), axis=0)
assert_almost_equal(found_images, cs, decimal=6)
def fractional_to_real(xyz, dims):
"""Convert fractional positions to real positions
Parameters
----------
xyz : np.ndarray
fractional positions
dims : dict
dict of box information
Returns
-------
xyz : np.ndarray
real space representation of fractional positions
"""
box = np.array([
dims['a'], dims['b'], dims['c'], dims['alpha'], dims['beta'],
dims['gamma']
],
dtype=np.float32)
xyz = distances.transform_StoR(xyz, box)
return distances.apply_PBC(xyz, box)
def search_tree(self, centers, radius):
"""
Searches all the pairs within `radius` between `centers`
and ``coords``
``coords`` are the already initialized coordinates in the tree
during :meth:`set_coords`.
``centers`` are wrapped around the primary unit cell
if PBC is desired. Minimum image convention (PBC) is
activated if the `box` argument is provided during
class initialization
Parameters
----------
centers: array_like (N,3)
coordinate array to search for neighbors
radius: float
maximum distance to search for neighbors.
Returns
-------
pairs : array
all the pairs between ``coords`` and ``centers``
Note
----
This method constructs another tree from the ``centers``
and queries the previously built tree (built in
:meth:`set_coords`)
"""
if not self._built:
raise RuntimeError('Unbuilt tree. Run tree.set_coords(...)')
centers = np.asarray(centers)
if centers.shape == (self.dim, ):
centers = centers.reshape((1, self.dim))
# Sanity check
if self.pbc:
if self.cutoff < radius:
raise RuntimeError(
'Set cutoff greater or equal to the radius.')
# Bring all query points to the central cell
wrapped_centers = apply_PBC(centers, self.box)
other_tree = cKDTree(wrapped_centers, leafsize=self.leafsize)
pairs = other_tree.query_ball_tree(self.ckdt, radius)
pairs = np.array([[i, j] for i, lst in enumerate(pairs)
for j in lst],
dtype=np.int64)
if pairs.size > 0:
pairs[:, 1] = undo_augment(pairs[:, 1], self.mapping,
len(self.coords))
else:
other_tree = cKDTree(centers, leafsize=self.leafsize)
pairs = other_tree.query_ball_tree(self.ckdt, radius)
pairs = np.array([[i, j] for i, lst in enumerate(pairs)
for j in lst],
dtype=np.int64)
if pairs.size > 0:
pairs = unique_rows(pairs)
return pairs
