import warnings
from typing import Dict, List, Optional, Union
import numpy as np
from spyrmsd import constants, graph, utils
[docs]class Molecule:
def __init__(
self,
atomicnums: Union[np.ndarray, List[int]],
coordinates: Union[np.ndarray, List[List[float]]],
adjacency_matrix: Optional[Union[np.ndarray, List[List[int]]]] = None,
) -> None:
"""
Molecule initialisation.
Parameters
----------
atomicnums: Union[np.ndarray, List[int]]
Atomic numbers
coordinates: Union[np.ndarray, List[List[float]]]
Atomic coordinates
adjacency_matrix: Union[np.ndarray, List[List[int]]], optional
Molecular graph adjacency matrix
Notes
-----
A molecule is built from atomic numbers and atomic coordinates only.
Optionally, a good representation of the molecular graph (obtained with
OpenBabel or RDKit) can be stored as adjacency matrix.
"""
atomicnums = np.asarray(atomicnums, dtype=int)
coordinates = np.asarray(coordinates, dtype=float)
self.natoms: int = len(atomicnums)
assert atomicnums.shape == (self.natoms,)
assert coordinates.shape == (self.natoms, 3)
self.atomicnums = atomicnums
self.coordinates = coordinates
self.stripped: bool = bool(np.all(atomicnums != 1))
if adjacency_matrix is not None:
self.adjacency_matrix: np.ndarray = np.asarray(adjacency_matrix, dtype=int)
# Molecular graph
self.G: Dict[str, object] = {}
self.masses: Optional[List[float]] = None
[docs] @classmethod
def from_obabel(cls, obmol, adjacency: bool = True):
"""
Constructor from OpenBabel molecule.
Parameters
----------
obmol:
OpenBabel molecule
adjacency:
Flag to compute the adjacency matrix
Returns
-------
spyrmsd.molecule.Molecule
:code:`spyrmsd` Molecule
"""
# TODO: Check if OpenBabel is available?
from spyrmsd.optional import obabel as ob
return ob.to_molecule(obmol, adjacency=adjacency)
[docs] @classmethod
def from_rdkit(cls, rdmol, adjacency: bool = True):
"""
Constructor from RDKit molecule.
Parameters
----------
rdmol:
RDKit molecule
adjacency:
Flag to compute the adjacency matrix
Returns
-------
spyrmsd.molecule.Molecule
:code:`spyrmsd` Molecule
"""
# TODO: Check if RDKit is available?
from spyrmsd.optional import rdkit as rd
return rd.to_molecule(rdmol, adjacency=adjacency)
[docs] def translate(self, vector: Union[np.ndarray, List[float]]) -> None:
"""
Translate molecule.
Parameters
----------
vector: np.ndarray
Translation vector (in 3D)
"""
assert len(vector) == 3
vector = np.asarray(vector)
self.coordinates += vector
[docs] def rotate(
self, angle: float, axis: Union[np.ndarray, List[float]], units: str = "rad"
) -> None:
"""
Rotate molecule.
Parameters
----------
angle: float
Rotation angle
axis: np.ndarray
Axis of rotation (in 3D)
units: {"rad", "deg"}
Units of the angle (radians `rad` or degrees `deg`)
"""
axis = np.asarray(axis)
self.coordinates = utils.rotate(self.coordinates, angle, axis, units)
[docs] def center_of_mass(self) -> np.ndarray:
"""
Center of mass.
Returns
-------
np.ndarray
Center of mass
Notes
-----
Atomic masses are cached.
"""
# Get masses and cache them
if self.masses is None:
self.masses = [constants.anum_to_mass[anum] for anum in self.atomicnums]
return np.average(self.coordinates, axis=0, weights=self.masses)
[docs] def center_of_geometry(self) -> np.ndarray:
"""
Center of geometry.
Returns
-------
np.ndarray
Center of geometry
"""
return utils.center_of_geometry(self.coordinates)
# TODO: Change name (to stripH)
[docs] def strip(self) -> None:
"""
Strip hydrogen atoms.
"""
if not self.stripped:
idx = self.atomicnums != 1 # Non-hydrogen atoms
# Strip
self.atomicnums = self.atomicnums[idx]
self.coordinates = self.coordinates[idx, :]
# Update number of atoms
self.natoms = len(self.atomicnums)
# Update adjacency matrix
if self.adjacency_matrix is not None:
self.adjacency_matrix = self.adjacency_matrix[np.ix_(idx, idx)]
# Reset molecular graph when stripping
self.G = {}
self.stripped = True
[docs] def to_graph(self):
"""
Convert molecule to graph.
Returns
-------
Graph
Molecular graph.
Notes
-----
If the molecule does not have an associated adjacency matrix, a simple
bond perception is used.
The molecular graph is cached per backend.
"""
_current_backend = graph._current_backend
if _current_backend not in self.G.keys():
try:
self.G[_current_backend] = graph.graph_from_adjacency_matrix(
self.adjacency_matrix, self.atomicnums
)
except AttributeError:
warnings.warn(
"Molecule was not initialized with an adjacency matrix. "
+ "Using bond perception..."
)
# Automatic bond perception (with very simple rule)
self.adjacency_matrix = graph.adjacency_matrix_from_atomic_coordinates(
self.atomicnums, self.coordinates
)
self.G[_current_backend] = graph.graph_from_adjacency_matrix(
self.adjacency_matrix, self.atomicnums
)
return self.G[_current_backend]
def __len__(self) -> int:
"""
Molecule size.
Returns
-------
int
Number of atoms within the molecule
"""
return self.natoms
[docs]def coords_from_molecule(mol: Molecule, center: bool = False) -> np.ndarray:
"""
Atomic coordinates from molecule.
Parameters
----------
mol: molecule.Molecule
Molecule
center: bool
Center flag
Returns
-------
np.ndarray
Atomic coordinates (possibly centred)
Notes
-----
Atomic coordinates are centred according to the center of geometry, not the center
of mass.
"""
if center:
coords = mol.coordinates - mol.center_of_geometry()
else:
coords = mol.coordinates
return coords