# This file is part of MDTools.
# Copyright (C) 2021 The MDTools Development Team and all contributors
# listed in the file AUTHORS.rst
#
# MDTools is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation, either version 3 of the License, or (at your
# option) any later version.
#
# MDTools is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
# for more details.
#
# You should have received a copy of the GNU General Public License
# along with MDTools. If not, see <http://www.gnu.org/licenses/>.
"""
Functions to select atoms from MDAnalysis
:class:`AtomGroup <MDAnalysis.core.groups.AtomGroup>` instances.
"""
# Standard libraries
import os
import warnings
from datetime import datetime
# Third party libraries
import psutil
import numpy as np
import MDAnalysis as mda
# Local application/library specific imports
import mdtools as mdt
[docs]def universe(top, trj, verbose=True, **kwargs):
"""
Create a MDAnalysis :class:`~MDAnalysis.core.universe.Universe` from
a topology and a trajectory file.
Parameters
----------
top : str
Filename of the topology file. See |supported_topology_formats|
of MDAnalysis.
trj : str
Filename of the trajectory. See |supported_coordinate_formats|
of MDAnalysis.
verbose : bool, optional
If ``True``, print some basic information about the created
MDAnalysis :class:`~MDAnalysis.core.universe.Universe` to
standard output. In fact, if `verbose` is ``False``,
:func:`mdtools.select.universe` simply calls the constructor of
:class:`MDAnalysis.core.universe.Universe` and returns the
resulting MDAnalysis :class:`MDAnalysis.core.universe.Universe`.
kwargs : dict, optional
Additional keyword arguments to pass to the constructor of
:class:`MDAnalysis.core.universe.Universe`.
Returns
-------
universe : MDAnalysis.core.universe.Universe
The created MDAnalysis
:class:`~MDAnalysis.core.universe.Universe`.
See Also
--------
:func:`mdtools.run_time_info.ag_info_str` :
Create a string containing information about a MDAnalysis
:class:`~MDAnalysis.core.groups.AtomGroup`.
Notes
-----
This function is just a convinient wrapper around the constructor of
:class:`MDAnalysis.core.universe.Universe` that can print the
information generated by :func:`mdtools.run_time_info.ag_info_str`
about the created :class:`MDAnalysis.core.universe.Universe` to
standard output.
"""
if verbose:
print("Creating universe...")
timer = datetime.now()
proc = psutil.Process(os.getpid())
universe = mda.Universe(top, trj, **kwargs)
if universe.atoms.n_atoms == 0:
warnings.warn("The created Universe contains no atoms",
RuntimeWarning)
if verbose:
# rsplit removes last line from ag_info_str about
# UpdatingAtomGroups
print(mdt.rti.ag_info_str(ag=universe.atoms).rsplit('\n', 1)[0])
print("Total number of frames: {}"
.format(universe.trajectory.n_frames))
print("Elapsed time: {}"
.format(datetime.now()-timer))
print("Current memory usage: {:.2f} MiB"
.format(proc.memory_info().rss/2**20))
return universe
[docs]def atoms(ag, sel, warn_empty=True, verbose=True, **kwargs):
"""
Select :class:`Atoms <MDAnalysis.core.groups.Atom>` from a
MDAnalysis :class:`~MDAnalysis.core.groups.AtomGroup` instance based
on a selection string.
Parameters
----------
ag : MDAnalysis.core.groups.AtomGroup
MDAnalysis :class:`~MDAnalysis.core.groups.AtomGroup`,
:class:`~MDAnalysis.core.groups.UpdatingAtomGroup` or
:class:`~MDAnalysis.core.universe.Universe` from which the
:class:`Atoms <MDAnalysis.core.groups.Atom>` will be selected.
sel : str
Selection string for selecting
:class:`Atoms <MDAnalysis.core.groups.Atom>` out of `ag`. See
MDAnalysis' |selection_syntax| for possible choices.
warn_empty : bool, optional
If ``True`` (default), raise a :exc:`RuntimeWarning` if the
created :class:`~MDAnalysis.core.groups.AtomGroup` contains no
:class:`Atoms <MDAnalysis.core.groups.Atom>` and is not an
:class:`~MDAnalysis.core.groups.UpdatingAtomGroup`.
verbose : bool, optional
If ``True``, print some basic information about the created
MDAnalysis :class:`~MDAnalysis.core.groups.AtomGroup` to
standard output. In fact, if `verbose` is ``False``,
:func:`mdtools.select.atoms` simply calls
:meth:`~MDAnalysis.core.groups.AtomGroup.select_atoms` and
returns the resulting :class:`~MDAnalysis.core.groups.AtomGroup`.
kwargs : dict, optional
Additional keyword arguments to pass to
:meth:`MDAnalysis.core.groups.AtomGroup.select_atoms`.
Returns
-------
selection : MDAnalysis.core.groups.AtomGroup or MDAnalysis.core.groups.UpdatingAtomGroup
A MDAnalysis :class:`~MDAnalysis.core.groups.AtomGroup` or
:class:`~MDAnalysis.core.groups.UpdatingAtomGroup` (if `updating`
is ``True``) containing the selected
:class:`Atoms <MDAnalysis.core.groups.Atom>`.
See Also
--------
:meth:`MDAnalysis.core.groups.AtomGroup.select_atoms` :
Underlying function of this function
:func:`mdtools.run_time_info.ag_info_str` :
Create a string containing information about a MDAnalysis
:class:`~MDAnalysis.core.groups.AtomGroup`.
Notes
-----
This function is just a convinient wrapper around
:meth:`MDAnalysis.core.groups.AtomGroup.select_atoms` that can
print the information generated by
:func:`mdtools.run_time_info.ag_info_str` about the created
:class:`MDAnalysis.core.groups.AtomGroup` to standard output.
"""
if verbose:
print("Creating selection...")
timer = datetime.now()
proc = psutil.Process(os.getpid())
print("Selection string: '{}'".format(sel))
selection = ag.select_atoms(sel, **kwargs)
if (warn_empty and
selection.n_atoms == 0 and
not isinstance(selection, mda.core.groups.UpdatingAtomGroup)):
warnings.warn("The created AtomGroup contains no atoms",
RuntimeWarning)
if verbose:
print(mdt.rti.ag_info_str(ag=selection))
print("Elapsed time: {}"
.format(datetime.now()-timer))
print("Current memory usage: {:.2f} MiB"
.format(proc.memory_info().rss/2**20))
return selection
[docs]def atoms_around_point(
ag, point, cutoff, warn_empty=False, verbose=False, **kwargs):
"""
Select :class:`Atoms <MDAnalysis.core.groups.Atom>` from a
MDAnalysis :class:`~MDAnalysis.core.groups.AtomGroup` instance that
are within a cutoff of a point in space.
Parameters
----------
ag : MDAnalysis.core.groups.AtomGroup
MDAnalysis :class:`~MDAnalysis.core.groups.AtomGroup`,
:class:`~MDAnalysis.core.groups.UpdatingAtomGroup` or
:class:`~MDAnalysis.core.universe.Universe` from which the
:class:`Atoms <MDAnalysis.core.groups.Atom>` will be selected.
point : array_like
Array of shape ``(3,)`` containing the coordinates in Angstrom
of a point in space. All
:class:`Atoms <MDAnalysis.core.groups.Atom>` of `ag` that are
within `cutoff` of that point will be selected.
cutoff : scalar
Select all :class:`Atoms <MDAnalysis.core.groups.Atom>` that are
within `cutoff` of `point`. Must be greater than zero.
warn_empty : bool, optional
If ``True``, raise a :exc:`RuntimeWarning` if the created
:class:`~MDAnalysis.core.groups.AtomGroup` contains no
:class:`Atoms <MDAnalysis.core.groups.Atom>` and is not an
:class:`~MDAnalysis.core.groups.UpdatingAtomGroup`.
verbose : bool, optional
If ``True``, print some basic information about the created
MDAnalysis :class:`~MDAnalysis.core.groups.AtomGroup` to
standard output.
kwargs : dict, optional
Additional keyword arguments to pass to
:meth:`MDAnalysis.core.groups.AtomGroup.select_atoms`.
Returns
-------
selection : MDAnalysis.core.groups.AtomGroup
A MDAnalysis :class:`~MDAnalysis.core.groups.AtomGroup`
containing the selected
:class:`Atoms <MDAnalysis.core.groups.Atom>`.
See Also
--------
:func:`mdtools.select.atoms` :
Select
:class:`Atoms <MDAnalysis.core.groups.Atom>` from a
MDAnalysis
:class:`~MDAnalysis.core.groups.AtomGroup`
Notes
-----
This function uses the MDAnalysis selection syntax
`point x y z distance`_. It is just a convinient wrapper around
:func:`mdtools.select.atoms` that allows you to simply parse an
array like `point` and a `cutoff` without the need to assemble them
to a MDAnalysis selection string yourself.
.. _point x y z distance: https://userguide.mdanalysis.org/stable/selections.html#geometric
"""
point = np.asarray(point)
if point.shape != (3,):
raise ValueError("'point' must have shape (3,) but has shape {}"
.format(point.shape))
if cutoff <= 0:
raise ValueError("'cutoff' ({}) must be greater than zero."
.format(cutoff))
point = ' '.join(str(i) for i in point)
sel = "point {} {}".format(point, cutoff)
return atoms(ag=ag,
sel=sel,
warn_empty=warn_empty,
verbose=verbose,
**kwargs)