# This file is part of MDTools.
# Copyright (C) 2021, 2022 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 warnings
from datetime import datetime
# Third-party libraries
import MDAnalysis as mda
import numpy as np
import psutil
# First-party libraries
import mdtools as mdt
[docs]
def universe(top, trj, verbose=True, **kwargs):
"""
Create an 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 parse to the constructor of
:class:`MDAnalysis.core.universe.Universe`.
Returns
-------
universe : MDAnalysis.core.universe.Universe
The created MDAnalysis
:class:`~MDAnalysis.core.universe.Universe`.
Raises
------
RuntimeWarning :
If the created :class:`~MDAnalysis.core.universe.Universe`
contains no atoms.
See Also
--------
:func:`mdtools.run_time_info.ag_info_str` :
Create a string containing information about an MDAnalysis
:class:`~MDAnalysis.core.groups.AtomGroup`.
Notes
-----
This function is just a convenient 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()
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(mdt.rti.mem_usage(proc))
)
return universe
[docs]
def atoms(ag, sel, warn_empty=True, verbose=True, **kwargs):
"""
Select :class:`Atoms <MDAnalysis.core.groups.Atom>` from an
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 parse to
:meth:`MDAnalysis.core.groups.AtomGroup.select_atoms`.
Returns
-------
selection : MDAnalysis.core.groups.AtomGroup or \
MDAnalysis.core.groups.UpdatingAtomGroup
An 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>`.
Raises
------
RuntimeWarning :
`warn_empty` is ``True`` and the created
:class:`~MDAnalysis.core.groups.AtomGroup` contains no atoms and
is not an :class:`~MDAnalysis.core.groups.UpdatingAtomGroup`.
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 an MDAnalysis
:class:`~MDAnalysis.core.groups.AtomGroup`.
Notes
-----
This function is just a convenient 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()
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(mdt.rti.mem_usage(proc))
)
return selection
[docs]
def atoms_around_point(ag, point, cutoff, **kwargs):
"""
Select :class:`Atoms <MDAnalysis.core.groups.Atom>` from an
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.
kwargs : dict, optional
Additional keyword arguments to parse to
:meth:`MDAnalysis.core.groups.AtomGroup.select_atoms`. See
there for possible choices. By default, `warn_empty` and
`verbose` are set to ``False``.
Returns
-------
selection : MDAnalysis.core.groups.AtomGroup
An 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 an
MDAnalysis :class:`~MDAnalysis.core.groups.AtomGroup`
Notes
-----
This function uses the MDAnalysis selection syntax
`point x y z distance`_. It is just a convenient 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 an 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)
kwargs.setdefault("warn_empty", False)
kwargs.setdefault("verbose", False)
return atoms(ag=ag, sel=sel, **kwargs)