"""
OpenMM topology transformations
===============================
.. moduleauthor:: Benjamin Ye <GitHub: @bbye98>
This module contains implementations of common OpenMM topology
transformations, like the generation of initial particle positions and
subsetting existing topology objects.
"""
from itertools import repeat
from typing import Any, Iterable, Union
import numpy as np
from openmm import app
from ..algorithm import topology as t
[docs]
def create_atoms(*args, **kwargs) -> Any:
"""
Generates initial particle positions.
This is an alias function. For more information, see
:func:`mdcraft.algorithm.topology.create_atoms`.
"""
return t.create_atoms(*args, **kwargs)
def _get_hierarchy_indices(
item: Union[app.Atom, app.topology.Bond, app.Residue, app.Chain],
bonds: dict[str, list],
) -> tuple[set[int], set[int], set[int], set[int]]:
"""
Get unique indices of all topology items related to the one passed
to this function.
For an atom, the indices of itself, the residue it belongs to, and
the chain that its residue belongs to are returned.
For a bond, the indices of the two atoms it connects, itself, the
residue(s) that the two atoms belong to, and the chain(s) that the
residue(s) belong to are returned.
For a residue, the indices of all underlying atoms and bonds,
itself, and the chain it belongs to are returned.
For a chain, the indices of all underlying atoms, bonds, and
residues, and itself are returned.
Parameters
----------
item : `openmm.app.Atom`, `openmm.app.topology.Bond`,
`openmm.app.Residue`, or `openmm.app.Chain`
Topology item of interest.
Returns
-------
atoms : `set`
Indices of all atoms associated with `item`.
bonds : `set`
Indices of all bonds associated with `item`.
residues : `set`
Indices of all residues associated with `item`.
chains : `set`
Indices of all chains associated with `item`.
"""
if isinstance(item, app.Atom):
return {item.index}, set(), {item.residue.index}, {item.residue.chain.index}
elif isinstance(item, app.topology.Bond):
return (
{item.atom1.index, item.atom2.index},
{bonds.index(item)},
{item.atom1.residue.index, item.atom2.residue.index},
{item.atom1.residue.chain.index, item.atom2.residue.chain.index},
)
elif isinstance(item, app.Residue):
return (
{a.index for a in item.atoms()},
{bonds.index(b) for b in item.bonds()},
{item.index},
{item.chain.index},
)
elif isinstance(item, app.Chain):
atom_indices = set()
bond_indices = set()
residue_indices = set()
for residue in item.residues():
a, b, r, _ = _get_hierarchy_indices(residue, bonds)
atom_indices |= a
bond_indices |= b
residue_indices |= r
return atom_indices, bond_indices, residue_indices, {item.index}
def _is_topology_object(obj: Any):
"""
Check if the argument is a topology item.
Parameters
----------
obj : `Any`
Any object.
Returns
-------
is_topology_object : `bool`
Boolean value indicating whether `obj` is a topology item.
"""
return isinstance(obj, (app.Atom, app.topology.Bond, app.Residue, app.Chain))
[docs]
def get_subset(
topology: app.Topology,
positions: np.ndarray[float],
*,
delete: list[Any] = None,
keep: list[Any] = None,
types: Union[str, Iterable[str]] = None,
) -> Union[app.Topology, np.ndarray[float]]:
r"""
Creates a topology subset and gets its corresponding particle positions.
Parameters
----------
topology : `openmm.app.Topology`
OpenMM topology.
positions : `numpy.ndarray`
Positions of the :math:`N` particles in the topology.
**Shape**: :math:`(N,\,3)`.
delete : array-like, keyword-only, optional
`openmm.app.Atom`, `openmm.app.Bond`, `openmm.app.Residue`,
and/or `openmm.app.Chain` objects, or the indices of those
objects. If indices are provided, their corresponding object
types (:code:`"atom"`, :code:`"residue"`, :code:`"chain"`) must
be provided in `types`. The specified items will be deleted from
the model.
.. note::
Only one of `delete` and `keep` can be specified.
keep : array-like, keyword-only, optional
`openmm.app.Atom`, `openmm.app.Bond`, `openmm.app.Residue`,
and/or `openmm.app.Chain` objects, or the indices of those
objects. If indices are provided, their corresponding object
types (:code:`"atom"`, :code:`"residue"`, :code:`"chain"`) must
be provided in `types`. The specified items will be kept in the
model.
.. note::
Only one of `delete` and `keep` can be specified.
types : `str` or array-like, keyword-only, optional
Object types corresponding to the indices provided in `delete`
or `keep`. If a `str` is provided, all items in the array are
assumed to have the same object type. Must be provided if not
all items in `delete` or `keep` are OpenMM topology objects.
Returns
-------
topology : `openmm.app.Topology`
OpenMM topology subset.
positions : `numpy.ndarray`
Positions of the remaining :math:`N_\mathrm{rem}` particles in
the topology.
**Shape**: :math:`(N_\mathrm{rem},\,3)`.
"""
# Set boolean flags for subroutines below
found = (delete is not None, keep is not None)
# Check if both delete and keep arguments were provided
if all(found):
emsg = (
"Only specify topology items to either delete or keep. "
"When both types are specified, the atoms, bonds, "
"residues, and/or chains to be removed from the topology "
"become ambiguous."
)
raise ValueError(emsg)
# Return original topology and positions if no items are specified
elif not any(found):
return topology, positions
# Ensure object type(s) are provided if not all items are topology
# objects
elif types is None and not all(
_is_topology_object(i) for i in next(a for a in [delete, keep] if a is not None)
):
emsg = (
"Object types must be specified for the topology items "
f"to be {'kept' if found[0] else 'deleted'}."
)
raise ValueError(emsg)
# Create a generator of types if a string is provided
elif isinstance(types, str):
same = True
types = repeat(types)
elif types is not None:
same = all(t == "atoms" for t in types)
# Create OpenMM Modeller
modeller = app.Modeller(topology, positions)
if types is not None:
# Create dictionary with topology subitems
model = {
"atom": list(topology.atoms()),
"bond": list(topology.bonds()),
"chain": list(topology.chains()),
"residue": list(topology.residues()),
}
# If indices and types of objects to be deleted are specified,
# create an iterable object of corresponding items from the
# dictionary
if found[0]:
delete = (
i if _is_topology_object(i) else model[t][i]
for i, t in zip(delete, types)
)
else:
# Preallocate sets to store indices of atoms, residues, and
# chains to keep
atoms = set()
bonds = set()
residues = set()
chains = set()
# Remove items to be kept from the master list of topology
# subitems to delete
for item, item_type in zip(keep, types):
if not _is_topology_object(item):
item = model[item_type][item]
a, b, r, c = _get_hierarchy_indices(item, model["bond"])
atoms |= a
bonds |= b
residues |= r
chains |= c
model["atom"] = np.delete(model["atom"], list(atoms))
model["residue"] = np.delete(model["residue"], list(residues))
model["chain"] = np.delete(model["chain"], list(chains))
if not bonds and same:
model["bond"] = []
else:
for i in sorted(bonds, reverse=True):
del model["bond"][i]
delete = [i for t in model.values() for i in t]
# Create subset by deleting objects from original topology
modeller.delete(delete)
return modeller.topology, modeller.positions