Source code for stk.molecular.molecules.molecule.molecule

"""
Molecule
========

.. toctree::
    :maxdepth: 2

    Building Block <stk.molecular.molecules.building_block>
    Constructed Molecule <stk.molecular.molecules.constructed_molecule>

"""


from __future__ import annotations

import os
import typing

import numpy as np
import rdkit.Chem.AllChem as rdkit
from scipy.spatial.distance import euclidean

from stk.utilities import (
    rotation_matrix,
    rotation_matrix_arbitrary_axis,
    vector_angle,
)

from ....utilities import OneOrMany
from ...atoms import Atom
from ...bonds import Bond
from ..utilities import get_bond_atom_ids, sort_bond_atoms_by_id
from .utilities import updaters, writers

_T = typing.TypeVar("_T", bound="Molecule")


[docs]class Molecule: """ An abstract base class for molecules. Notes: You might notice that some of the methods of this abstract base class are implemented. This is purely for convenience when implementing subclasses. The implemented public methods are simply default implementations, which can be safely ignored or overridden, when implementing subclasses. Any private methods are implementation details of these default implementations. Examples: *Aligning a Molecule with a Vector* You want to rotate a molecule, such that it is aligned along with a specific direction. .. testcode:: aligning-a-molecule-with-a-vector import stk import numpy as np molecule1 = stk.BuildingBlock('CCCCC') # Align molecule1 along x-axis. molecule1 = molecule1.with_rotation_between_vectors( start=molecule1.get_direction(), target=np.array([1., 0., 0.]), origin=molecule1.get_centroid(), ) molecule2 = stk.ConstructedMolecule( topology_graph=stk.polymer.Linear( building_blocks=( stk.BuildingBlock( smiles='BrCCCBr', functional_groups=(stk.BromoFactory(), ), ), ), repeating_unit='A', num_repeating_units=15, ), ) # Align molecule2 along the [1, 4, -3] vector. molecule2 = molecule2.with_rotation_between_vectors( start=molecule2.get_direction(), target=np.array([1., 4., -3.]), origin=molecule2.get_centroid(), ) *Aligning a Molecule along a Plane* You want to place the benzene flat along the xy plane. .. testcode:: aligning-a-molecule-along-a-plane import stk import numpy as np benzene = stk.BuildingBlock('c1ccccc1') benzene = benzene.with_rotation_between_vectors( start=benzene.get_plane_normal(), target=np.array([0., 0., 1.]), origin=benzene.get_centroid(), ) """
[docs] def __init__( self, atoms: tuple[Atom, ...], bonds: tuple[Bond, ...], position_matrix: np.ndarray, ) -> None: """ Initialize a :class:`.Molecule`. Parameters: atoms: The atoms which compose the molecule. bonds: The bonds of the molecule. position_matrix: A ``(n, 3)`` matrix holding the position of every atom in the :class:`.Molecule`. """ self._atoms = atoms self._bonds = bonds # Take the transpose because it will make some matrix # multiplications faster. self._position_matrix = np.array( position_matrix.T, dtype=np.float64, )
def _with_displacement( self: _T, displacement: np.ndarray, ) -> _T: """ Modify molecule. """ self._position_matrix = ( self._position_matrix.T + displacement ).T return self
[docs] def with_displacement( self, displacement: np.ndarray, ) -> Molecule: """ Return a displaced clone. Parameters: displacement: The displacement vector to be applied. Returns: A displaced clone. """ return self.clone()._with_displacement(displacement)
def _with_rotation_about_axis( self: _T, angle: float, axis: np.ndarray, origin: np.ndarray, ) -> _T: """ Modify molecule. """ # Set the origin of the rotation to "origin". self._with_displacement(-origin) rot_mat = rotation_matrix_arbitrary_axis(angle, axis) # Apply the rotation matrix on the position matrix, to get the # new position matrix. self._position_matrix = rot_mat @ self._position_matrix # Return the centroid of the molecule to the original position. self._with_displacement(origin) return self
[docs] def with_rotation_about_axis( self, angle: float, axis: np.ndarray, origin: np.ndarray, ) -> Molecule: """ Return a rotated clone. The clone is rotated by `angle` about `axis` on the `origin`. Parameters: angle: The size of the rotation in radians. axis: The axis about which the rotation happens. Must have unit magnitude. origin: The origin about which the rotation happens. Returns: A rotated clone. """ return self.clone()._with_rotation_about_axis( angle=angle, axis=axis, origin=origin, )
def _with_rotation_between_vectors( self: _T, start: np.ndarray, target: np.ndarray, origin: np.ndarray, ) -> _T: """ Modify molecule. """ # Set the origin of the rotation to "origin". self._with_displacement(-origin) rot_mat = rotation_matrix(start, target) # Apply the rotation matrix to the atomic positions to yield # the new atomic positions. self._position_matrix = rot_mat @ self._position_matrix # Restore original position. self._with_displacement(origin) return self
[docs] def with_rotation_between_vectors( self, start: np.ndarray, target: np.ndarray, origin: np.ndarray, ) -> Molecule: """ Return a rotated clone. The rotation is equal to a rotation from `start` to `target`. Given two direction vectors, `start` and `target`, this method applies the rotation required transform `start` to `target` onto the clone. The rotation occurs about the `origin`. For example, if the `start` and `target` vectors are 45 degrees apart, a 45 degree rotation will be applied to the clone. The rotation will be along the appropriate direction. The great thing about this method is that you as long as you can associate a geometric feature of the molecule with a vector, then the clone can be rotated so that this vector is aligned with `target`. The defined vector can be virtually anything. This means that any geometric feature of the molecule can be easily aligned with any arbitrary direction. Parameters: start: A vector which is to be rotated so that it transforms into the `target` vector. target: The vector onto which `start` is rotated. origin: The point about which the rotation occurs. Returns: A rotated clone. """ return self.clone()._with_rotation_between_vectors( start=start, target=target, origin=origin, )
def _with_rotation_to_minimize_angle( self: _T, start: np.ndarray, target: np.ndarray, axis: np.ndarray, origin: np.ndarray, ) -> _T: # If the vector being rotated is not finite, exit. This is # probably due to a planar molecule. if not all(np.isfinite(x) for x in start): return self if np.allclose(target, [0, 0, 0], atol=1e-15): raise ValueError( "target has a magnitude of 0. It is therefore not " "possible to calculate an angle." ) self._with_displacement(-origin) # 1. Remove any component of the start and target vectors long # the axis. This puts them both on the same plane. # 2. Calculate the angle between them. # 3. Apply the rotation. tstart = start - np.dot(start, axis) * axis # If `tstart` is 0, it is parallel to the rotation axis, stop. if np.allclose(tstart, [0, 0, 0], 1e-8): self._with_displacement(origin) return self tend = target - np.dot(target, axis) * axis # If `tend` is 0, it is parallel to the rotation axis, stop. if np.allclose(tend, [0, 0, 0], 1e-8): self._with_displacement(origin) return self angle = vector_angle(tstart, tend) projection = tstart @ np.cross(axis, tend) if projection > 0: angle = 2 * np.pi - angle rotation_matrix = rotation_matrix_arbitrary_axis(angle, axis) self._position_matrix = rotation_matrix @ self._position_matrix self._with_displacement(origin) return self
[docs] def with_rotation_to_minimize_angle( self, start: np.ndarray, target: np.ndarray, axis: np.ndarray, origin: np.ndarray, ) -> Molecule: """ Return a rotated clone. The clone is rotated by the rotation required to minimize the angle between `start` and `target`. Note that this function will not necessarily overlay the `start` and `target` vectors. This is because the possible rotation is restricted to the `axis`. Parameters: start: The vector which is rotated. target: The vector which is stationary. axis: The vector about which the rotation happens. Must have unit magnitude. origin: The origin about which the rotation happens. Returns: A rotated clone. Raises: :class:`ValueError` If `target` has a magnitude of 0. In this case it is not possible to calculate an angle between `start` and `target`. """ return self.clone()._with_rotation_to_minimize_angle( start=start, target=target, axis=axis, origin=origin, )
def _clone(self: _T) -> _T: clone = self.__class__.__new__(self.__class__) Molecule.__init__( self=clone, atoms=self._atoms, bonds=self._bonds, position_matrix=self._position_matrix.T, ) return clone
[docs] def clone(self) -> Molecule: """ Return a clone. Returns: The clone. """ return self._clone()
[docs] def get_atomic_positions( self, atom_ids: typing.Optional[OneOrMany[int]] = None, ) -> typing.Iterator[np.ndarray]: """ Yield the positions of atoms. Parameters: atom_ids: The ids of the atoms whose positions are desired. If ``None``, then the positions of all atoms will be yielded. Can be a single :class:`int`, if the position of a single atom is desired. Yields: The x, y and z coordinates of an atom. """ if atom_ids is None: atom_ids = range(len(self._atoms)) elif isinstance(atom_ids, int): atom_ids = (atom_ids,) elif not isinstance(atom_ids, (list, tuple)): atom_ids = list(atom_ids) positions = self._position_matrix[:, atom_ids].T for atomic_position in positions: yield atomic_position
[docs] def get_atoms( self, atom_ids: typing.Optional[OneOrMany[int]] = None, ) -> typing.Iterator[Atom]: """ Yield the atoms in the molecule, ordered by id. Parameters: atom_ids: The ids of atoms to yield. Can be a single :class:`int` if a single atom is wanted, or ``None`` if all atoms are wanted. Yields: An atom in the molecule. """ if atom_ids is None: atom_ids = range(len(self._atoms)) elif isinstance(atom_ids, int): atom_ids = (atom_ids,) for atom_id in atom_ids: yield self._atoms[atom_id]
[docs] def get_num_atoms(self) -> int: """ Return the number of atoms in the molecule. Returns: The number of atoms in the molecule. """ return len(self._atoms)
[docs] def get_bonds(self) -> typing.Iterator[Bond]: """ Yield the bond in the molecule. Yields: A bond in the molecule. """ for bond in self._bonds: yield bond
[docs] def get_num_bonds(self) -> int: """ Return the number of bonds in the molecule. Returns: The number of bonds in the molecule. """ return len(self._bonds)
[docs] def get_centroid( self, atom_ids: typing.Optional[OneOrMany[int]] = None, ) -> np.ndarray: """ Return the centroid. Parameters: atom_ids: The ids of atoms which are used to calculate the centroid. Can be a single :class:`int`, if a single atom is to be used, or ``None`` if all atoms are to be used. Returns: The centroid of atoms specified by `atom_ids`. Raises: :class:`ValueError` If `atom_ids` has a length of ``0``. """ if atom_ids is None: atom_ids = range(len(self._atoms)) elif isinstance(atom_ids, int): atom_ids = (atom_ids,) elif not isinstance(atom_ids, (list, tuple)): atom_ids = list(atom_ids) if len(atom_ids) == 0: raise ValueError("atom_ids was of length 0.") return np.divide( self._position_matrix[:, atom_ids].sum(axis=1), len(atom_ids), )
[docs] def get_direction( self, atom_ids: typing.Optional[OneOrMany[int]] = None, ) -> np.ndarray: """ Return a vector of best fit through the atoms. Parameters: atom_ids: The ids of atoms which should be used to calculate the vector. Can be a single :class:`int`, if a single atom is to be used, or ``None``, if all atoms are to be used. Returns: The vector of best fit. Raises: :class:`ValueError` If `atom_ids` has a length of ``0``. """ if atom_ids is None: atom_ids = range(len(self._atoms)) elif isinstance(atom_ids, int): atom_ids = (atom_ids,) elif not isinstance(atom_ids, (list, tuple)): atom_ids = list(atom_ids) if len(atom_ids) == 0: raise ValueError("atom_ids was of length 0.") pos = self._position_matrix[:, atom_ids].T return np.around( a=np.linalg.svd(pos - pos.mean(axis=0))[-1][0], decimals=14, )
[docs] def get_maximum_diameter( self, atom_ids: typing.Optional[OneOrMany[int]] = None, ) -> float: """ Return the maximum diameter. This method does not account for the van der Waals radius of atoms. Parameters: atom_ids: The ids of atoms which are considered when looking for the maximum diameter. Can be a single :class:`int`, if a single atom is to be used, or ``None``, if all atoms are to be used. Returns: The maximum diameter in the molecule. Raises: :class:`ValueError` If `atom_ids` has a length of ``0``. """ if atom_ids is None: atom_ids = range(len(self._atoms)) elif isinstance(atom_ids, int): atom_ids = (atom_ids,) elif not isinstance(atom_ids, (list, tuple)): atom_ids = list(atom_ids) if len(atom_ids) == 0: raise ValueError("atom_ids was of length 0.") coords = self._position_matrix[:, atom_ids] return float(euclidean(coords.min(axis=1), coords.max(axis=1)))
[docs] def get_plane_normal( self, atom_ids: typing.Optional[OneOrMany[int]] = None, ) -> np.ndarray: """ Return the normal to the plane of best fit. Parameters: atom_ids: The ids of atoms which should be used to calculate the plane. Can be a single :class:`int`, if a single atom is to be used, or ``None``, if all atoms are to be used. Returns: Vector orthonormal to the plane of the molecule. Raises: :class:`ValueError` If `atom_ids` has a length of ``0``. """ if atom_ids is None: atom_ids = range(len(self._atoms)) elif isinstance(atom_ids, int): atom_ids = (atom_ids,) elif not isinstance(atom_ids, (list, tuple)): atom_ids = list(atom_ids) if len(atom_ids) == 0: raise ValueError("atom_ids was of length 0.") pos = self._position_matrix[:, atom_ids].T centroid = self.get_centroid(atom_ids) return np.around(np.linalg.svd(pos - centroid)[-1][2, :], 14)
[docs] def get_position_matrix(self) -> np.ndarray: """ Return a matrix holding the atomic positions. Returns: The array has the shape ``(n, 3)``. Each row holds the x, y and z coordinates of an atom. """ return np.array(self._position_matrix.T)
def _with_centroid( self: _T, position: np.ndarray, atom_ids: typing.Optional[OneOrMany[int]], ) -> _T: centroid = self.get_centroid(atom_ids=atom_ids) self._with_displacement(position - centroid) return self
[docs] def with_centroid( self, position: np.ndarray, atom_ids: typing.Optional[OneOrMany[int]] = None, ) -> Molecule: """ Return a clone with its centroid at `position`. Parameters: position: This array holds the position on which the centroid of the clone is going to be placed. atom_ids: The ids of atoms which should have their centroid set to `position`. Can be a single :class:`int`, if a single atom is to be used, or ``None``, if all atoms are to be used. Returns: A clone with its centroid at `position`. """ return self.clone()._with_centroid(position, atom_ids)
def _with_position_matrix( self: _T, position_matrix: np.ndarray, ) -> _T: """ Modify molecule. """ self._position_matrix = np.array(position_matrix.T) return self
[docs] def with_position_matrix( self, position_matrix: np.ndarray, ) -> Molecule: """ Return a clone with atomic positions set by `position_matrix`. Parameters: position_matrix: The position matrix of the clone. The shape of the matrix is ``(n, 3)``. Returns: The clone. """ return self.clone()._with_position_matrix(position_matrix)
[docs] def to_rdkit_mol(self) -> rdkit.Mol: """ Return an :mod:`rdkit` representation. Returns: The molecule in :mod:`rdkit` format. """ mol = rdkit.EditableMol(rdkit.Mol()) for atom in self._atoms: rdkit_atom = rdkit.Atom(atom.get_atomic_number()) rdkit_atom.SetFormalCharge(atom.get_charge()) mol.AddAtom(rdkit_atom) for bond in self._bonds: mol.AddBond( beginAtomIdx=bond.get_atom1().get_id(), endAtomIdx=bond.get_atom2().get_id(), order=( rdkit.BondType.DATIVE if bond.get_order() == 9 else rdkit.BondType(bond.get_order()) ), ) mol = mol.GetMol() rdkit_conf = rdkit.Conformer(len(self._atoms)) for atom_id, atom_coord in enumerate(self._position_matrix.T): rdkit_conf.SetAtomPosition(atom_id, atom_coord) mol.GetAtomWithIdx(atom_id).SetNoImplicit(True) mol.AddConformer(rdkit_conf) return mol
[docs] def with_structure_from_file( self, path: str, extension: typing.Optional[str] = None, ) -> Molecule: """ Return a clone, with its structure taken from a file. Multiple file types are supported, namely: #. ``.mol``, ``.sdf`` - MDL V2000 and V3000 files #. ``.xyz`` - XYZ files #. ``.mae`` - Schrodinger Maestro files #. ``.coord`` - Turbomole files #. ``.pdb`` - PDB files Parameters: path: The path to a molecular structure file holding updated coordinates for the :class:`.Molecule`. extension: If you want to treat the file as though it has a particular extension, put it here. Include the dot. Returns: A clone with atomic positions found in `path`. """ if extension is None: _, extension = os.path.splitext(path) return { ".mol": updaters._with_structure_from_mol, ".sdf": updaters._with_structure_from_mol, ".mae": updaters._with_structure_from_mae, ".xyz": updaters._with_structure_from_xyz, ".coord": updaters._with_structure_from_turbomole, ".pdb": updaters._with_structure_from_pdb, }[extension](self.clone(), path)
[docs] def with_canonical_atom_ordering(self) -> Molecule: """ Return a clone, with canonically ordered atoms. Returns: The clone. """ return self.clone()._with_canonical_atom_ordering()
def _with_canonical_atom_ordering(self: _T) -> _T: """ Modify the molecule. """ atom_map = { atom.get_id(): atom.with_id(new_id) for new_id, atom in zip( rdkit.CanonicalRankAtoms(self.to_rdkit_mol()), self._atoms, ) } self._atoms = tuple( sorted(atom_map.values(), key=lambda atom: atom.get_id()) ) self._bonds = tuple( sorted( ( sort_bond_atoms_by_id(bond.with_atoms(atom_map)) for bond in self._bonds ), key=get_bond_atom_ids, ) ) old_ids = { atom.get_id(): old_id for old_id, atom in atom_map.items() } self._position_matrix = np.array( np.array( [ self._position_matrix.T[old_ids[new_id]] for new_id in range(len(self._atoms)) ] ).T ) return self
[docs] def get_canonical_atom_ids(self) -> dict[int, int]: """ Map the id of each atom to its id under canonical ordering. Returns: Maps the id of each atom in the molecule to the id it would have under canonical ordering. """ return { old_id: new_id for old_id, new_id in enumerate( rdkit.CanonicalRankAtoms(self.to_rdkit_mol()), ) }
[docs] def write( self, path: str, atom_ids: typing.Optional[OneOrMany[int]] = None, ) -> Molecule: """ Write the structure to a file. This function will write the format based on the extension of `path`. #. ``.mol``, ``.sdf`` - MDL V3000 MOL file #. ``.xyz`` - XYZ file #. ``.pdb`` - PDB file Parameters: path: The `path` to which the molecule should be written. atom_ids: The atom ids of atoms to write. Can be a single :class:`int`, if a single atom is to be used, or ``None``, if all atoms are to be used. If you use this parameter, the atom ids in the file may not correspond to the atom ids in the molecule. Returns: The molecule. """ _, extension = os.path.splitext(path) { ".mol": writers._write_mdl_mol_file, ".sdf": writers._write_mdl_mol_file, ".xyz": writers._write_xyz_file, ".pdb": writers._write_pdb_file, }[extension](self, path, atom_ids) return self
def __str__(self) -> str: return repr(self) def __repr__(self) -> str: return f"<{self.__class__.__name__} at {id(self)}>"