Skip to content

Commit

Permalink
Implemented compact wrapping and nojump; updated transforms.
Browse files Browse the repository at this point in the history 
Optimized AtomGroup.wrap() for compounds.

Added transformations.nojump; transformations.wrap and transformations.unwrap
now transparently expose AtomGroup.wrap() and AtomGroup.unwrap().

Added function apply_compact_PBC and added 'center' keyword to
apply_PBC.

Some PBC-related doc cleaning and clarification.
  • Loading branch information
@mnmelo
mnmelo committed Oct 13, 2020
1 parent a9543ba commit 69bc40d
Show file tree
Hide file tree
Showing 5 changed files with 356 additions and 63 deletions.
57 changes: 46 additions & 11 deletions package/MDAnalysis/core/groups.py
View file
Expand Up @@ -1345,7 +1345,7 @@ def pack_into_box(self, box=None, inplace=True):
"""
return self.wrap(box=box, inplace=inplace)

def wrap(self, compound="atoms", center="com", box=None, inplace=True):
def wrap(self, compound="atoms", center="com", compact=False, box=None, inplace=True):
r"""Shift the contents of this group back into the primary unit cell
according to periodic boundary conditions.
Expand All @@ -1354,6 +1354,28 @@ def wrap(self, compound="atoms", center="com", box=None, inplace=True):
``'atoms'``, each compound as a whole will be shifted so that its
`center` lies within the primary unit cell.
Wrapping with `compound`='atoms':
+-----------+ +-----------+
| | | |
| | 3 6 | 3 6 |
| | ! ! | ! ! |
| 1-|-2-5-8 -> |-2-5-8 1-|
| | ! ! | ! ! |
| | 4 7 | 4 7 |
| | | |
+-----------+ +-----------+
Wrapping with `compound`='fragments':
+-----------+ +-----------+
| | | |
| | 3 6 | 3 6 |
| | ! ! | ! ! |
| 1-|-2-5-8 -> 1-|-2-5-8 |
| | ! ! | ! ! |
| | 4 7 | 4 7 |
| | | |
+-----------+ +-----------+
Parameters
----------
compound : {'atoms', 'group', 'segments', 'residues', 'molecules', \
Expand All @@ -1364,6 +1386,10 @@ def wrap(self, compound="atoms", center="com", box=None, inplace=True):
center : {'com', 'cog'}
How to define the center of a given group of atoms. If `compound` is
``'atoms'``, this parameter is meaningless and therefore ignored.
compact : bool, optional
If ``True``, wrapping will be done to the space-filling volume
closest to the primary unit cell's center (only useful for
non-orthogonal periodic boxes).
box : array_like, optional
The unitcell dimensions of the system, which can be orthogonal or
triclinic and must be provided in the same format as returned by
Expand Down Expand Up @@ -1407,6 +1433,12 @@ def wrap(self, compound="atoms", center="com", box=None, inplace=True):
x_i' = x_i - \left\lfloor\frac{x_i}{L_i}\right\rfloor\,.
`compact` changes the shifting behavior: centers will be shifted by
the integer combination of box vectors that minimizes
:math:`\left\|\mathbf{r_{compact}}-\mathbf{r_{box\_center}}\right\|`,
the distance of a center's final position, :math:`\mathbf{r_{compact}}`,
to the primary unit cell's center.
When specifying a `compound`, the translation is calculated based on
each compound. The same translation is applied to all atoms
within this compound, meaning it will not be broken by the shift.
Expand All @@ -1420,6 +1452,7 @@ def wrap(self, compound="atoms", center="com", box=None, inplace=True):
This can be either ``'com'`` for center of mass, or ``'cog'`` for center
of geometry.
`box` allows a unit cell to be given for the transformation. If not
specified, the :attr:`~MDAnalysis.coordinates.base.Timestep.dimensions`
information from the current
Expand Down Expand Up @@ -1451,6 +1484,7 @@ def wrap(self, compound="atoms", center="com", box=None, inplace=True):
raise ValueError("Invalid box: Box has invalid shape or not all "
"box dimensions are positive. You can specify a "
"valid box using the 'box' argument.")
packer = distances.apply_compact_PBC if compact else distances.apply_PBC

# no matter what kind of group we have, we need to work on its (unique)
# atoms:
Expand All @@ -1472,7 +1506,7 @@ def wrap(self, compound="atoms", center="com", box=None, inplace=True):
return np.zeros((0, 3), dtype=np.float32)

if comp == "atoms" or len(atoms) == 1:
positions = distances.apply_PBC(atoms.positions, box)
positions = packer(atoms.positions, box)
else:
ctr = center.lower()
if ctr == 'com':
Expand All @@ -1496,7 +1530,7 @@ def wrap(self, compound="atoms", center="com", box=None, inplace=True):
else: # ctr == 'cog'
ctrpos = atoms.center_of_geometry(pbc=False, compound=comp)
ctrpos = ctrpos.astype(np.float32, copy=False)
target = distances.apply_PBC(ctrpos, box)
target = packer(ctrpos, box)
positions += target - ctrpos
else:
if comp == 'segments':
Expand All @@ -1510,6 +1544,7 @@ def wrap(self, compound="atoms", center="com", box=None, inplace=True):
errmsg = ("Cannot use compound='molecules', this "
"requires molnums.")
raise NoDataError(errmsg) from None

else: # comp == 'fragments'
try:
compound_indices = atoms.fragindices
Expand All @@ -1529,16 +1564,16 @@ def wrap(self, compound="atoms", center="com", box=None, inplace=True):
else: # ctr == 'cog'
ctrpos = atoms.center_of_geometry(pbc=False, compound=comp)
ctrpos = ctrpos.astype(np.float32, copy=False)
target = distances.apply_PBC(ctrpos, box)
target = packer(ctrpos, box)
shifts = target - ctrpos

# There may be gaps in the used indices. We must translate
# them to a (0, 1, 2, ..., n_indices-1) sequence
unique_indices = unique_int_1d(compound_indices)
transl_indices = np.arange(unique_indices.max()+1)
transl_indices[unique_indices] = np.arange(len(unique_indices))
# apply the shifts:
unique_compound_indices = unique_int_1d(compound_indices)
shift_idx = 0
for i in unique_compound_indices:
mask = np.where(compound_indices == i)
positions[mask] += shifts[shift_idx]
shift_idx += 1
positions += shifts[transl_indices[compound_indices]]

if inplace:
atoms.positions = positions
Expand All @@ -1559,7 +1594,7 @@ def unwrap(self, compound='fragments', reference='com', inplace=True):
| | | |
| 6 3 | | 3 | 6
| ! ! | | ! | !
|-5-8 1-2-| ==> | 1-2-|-5-8
|-5-8 1-2-| -> | 1-2-|-5-8
| ! ! | | ! | !
| 7 4 | | 4 | 7
| | | |
Expand Down
99 changes: 96 additions & 3 deletions package/MDAnalysis/lib/distances.py
View file
Expand Up @@ -61,13 +61,15 @@
.. autofunction:: calc_angles
.. autofunction:: calc_dihedrals
.. autofunction:: apply_PBC
.. autofunction:: apply_compact_PBC
.. autofunction:: transform_RtoS
.. autofunction:: transform_StoR
.. autofunction:: augment_coordinates(coordinates, box, r)
.. autofunction:: undo_augment(results, translation, nreal)
"""
import numpy as np
from numpy.lib.utils import deprecate
import itertools

from .util import check_coords, check_box
from .mdamath import triclinic_vectors
Expand Down Expand Up @@ -1133,7 +1135,7 @@ def transform_RtoS(coords, box, backend="serial"):
Returns
-------
newcoords : numpy.ndarray (``dtype=numpy.float32``, ``shape=coords.shape``)
An array containing fractional coordiantes.
An array containing fractional coordinates.
.. versionchanged:: 0.13.0
Expand Down Expand Up @@ -1482,8 +1484,8 @@ def calc_dihedrals(coords1, coords2, coords3, coords4, box=None, result=None,


@check_coords('coords')
def apply_PBC(coords, box, backend="serial"):
"""Moves coordinates into the primary unit cell.
def apply_PBC(coords, box, center=None, backend="serial"):
"""Moves coordinates into an arbitrarily centered unit cell.
Parameters
----------
Expand All @@ -1495,6 +1497,9 @@ def apply_PBC(coords, box, backend="serial"):
triclinic and must be provided in the same format as returned by
:attr:`MDAnalysis.coordinates.base.Timestep.dimensions`:
``[lx, ly, lz, alpha, beta, gamma]``.
center : numpy.ndarray
Coordinate array of shape ``(3,)`` of the center of the output unit
cell representation. Defaults to the primary unit cell's center.
backend : {'serial', 'OpenMP'}, optional
Keyword selecting the type of acceleration.
Expand All @@ -1511,13 +1516,101 @@ def apply_PBC(coords, box, backend="serial"):
.. versionchanged:: 0.19.0
Internal dtype conversion of input coordinates to ``numpy.float32``.
Now also accepts (and, likewise, returns) single coordinates.
.. versionchanged:: 2.0.0
Added *center* keyword.
"""
if len(coords) == 0:
return coords
boxtype, box = check_box(box)

if center is not None:
box_center = triclinic_vectors(box).sum(axis=0) * 0.5
center_displacement = box_center - center
coords += center_displacement

if boxtype == 'ortho':
_run("ortho_pbc", args=(coords, box), backend=backend)
else:
_run("triclinic_pbc", args=(coords, box), backend=backend)

if center is not None:
coords -= center_displacement

return coords

@check_coords('coords')
def apply_compact_PBC(coords, box, center=None, backend="serial"):
"""Moves coordinates into an arbitrarily centered, space-filling compact
volume.
Parameters
----------
coords : numpy.ndarray
Coordinate array of shape ``(3,)`` or ``(n, 3)`` (dtype is arbitrary,
will be converted to ``numpy.float32`` internally).
box : numpy.ndarray
The unitcell dimensions of the system, which can be orthogonal or
triclinic and must be provided in the same format as returned by
:attr:`MDAnalysis.coordinates.base.Timestep.dimensions`:
``[lx, ly, lz, alpha, beta, gamma]``.
center : numpy.ndarray
Coordinate array of shape ``(3,)`` of the center of the output compact
representation. Defaults to the primary unit cell's center.
backend : {'serial', 'OpenMP'}, optional
Keyword selecting the type of acceleration.
Returns
-------
newcoords : numpy.ndarray (``dtype=numpy.float32``, ``shape=coords.shape``)
Array containing coordinates that all lie within a compact space-filling
volume centered on the point given with `center` (compact in that all
that volume's points are closest to the center than to any of the
center's periodic images).
.. versionadded:: 2.0.0
"""
boxtype, box = check_box(box)
# shortcut for orthogonal boxes
if boxtype == 'ortho':
return apply_PBC(coords, box, center=center, backend=backend)

tric_vecs = triclinic_vectors(box)
pbc_img_coeffs = np.array(list(itertools.product([0,-1,1], repeat=3)))
pbc_img_vecs = np.matmul(pbc_img_coeffs, tric_vecs)

# Stuff could get sped up because this can work with squared distance;
# no accelerated function exists for it yet, though.
min_compact_norm = np.linalg.norm(pbc_img_vecs[1:], axis=1).min() * 0.5

if center is None:
center = tric_vecs.sum(axis=0) * 0.5
else:
center = np.asarray(center, dtype=np.float32)
pbc_imgs = pbc_img_vecs + center

# Also improvable as distance^2
center_dists = distance_array(center, coords)[0]
outside_ats = np.where(center_dists > min_compact_norm)[0]
# Let's make sure we're dealing with everyone in the primary cell
outside_pos = coords[outside_ats] = apply_PBC(coords[outside_ats],
box, backend)
# Potentially save cycles by re-filtering after PBC. It makes sense for GROMACS
# trajectories, that are saved as the rectangular cell.
center_dists = distance_array(center, outside_pos)[0]
outside_ats = outside_ats[center_dists > min_compact_norm]
outside_pos = coords[outside_ats]

# This works for the general case by checking the distance to all 27
# unit cell centers (primary + 1st neighbors). Will break for extremely
# skewed cases where the closest center is in a 2nd neighbor unit cell.
# The distance check to all 26 1st neighbors can also probably be optimized,
# both in general and also with boxtype-specific shortcuts.
# Also improvable as distance^2
closest_img_ndx = distance_array(outside_pos, pbc_imgs).argmin(axis=1)
to_move = closest_img_ndx.nonzero()[0]
closest_img_ndx = closest_img_ndx[to_move]
outside_ats = outside_ats[to_move]
outside_pos = coords[outside_ats] - pbc_img_vecs[closest_img_ndx]
coords[outside_ats] = outside_pos
return coords
2 changes: 1 addition & 1 deletion package/MDAnalysis/lib/mdamath.py
View file
Expand Up @@ -58,7 +58,7 @@ def norm(v):
.. math:: v = \sqrt{\mathbf{v}\cdot\mathbf{v}}
This version is faster then numpy.linalg.norm because it only works for a
This version is faster than numpy.linalg.norm because it only works for a
single vector and therefore can skip a lot of the additional fuss
linalg.norm does.
Expand Down
2 changes: 1 addition & 1 deletion package/MDAnalysis/transformations/__init__.py
View file
Expand Up @@ -129,4 +129,4 @@ def wrapped(ts):
from .rotate import rotateby
from .positionaveraging import PositionAverager
from .fit import fit_translation, fit_rot_trans
from .wrap import wrap, unwrap
from .wrap import wrap, unwrap, nojump

0 comments on commit 69bc40d

Please sign in to comment.