fragmenstein/walton/_movements.py
from typing import (Tuple, Union, Dict)
import numpy as np
from rdkit import Chem
from rdkit import Geometry
from rdkit.Chem import rdFMCS, AllChem, rdMolAlign
from ._base import WaltonBase
class WaltonMove(WaltonBase):
# get coordinates
def get_point(self, atom_idx: Union[int, Geometry.Point3D], mol_idx: int = -1) -> Geometry.Point3D:
"""
Given atom index and mol index return a point.
If a point itself is given, return that.
"""
if isinstance(atom_idx, Geometry.Point3D):
return atom_idx
elif mol_idx == -1:
raise ValueError('If providing an atom index as a number please provide a mol_idx')
elif isinstance(atom_idx, int):
return self.get_mol(mol_idx).GetConformer(0).GetAtomPosition(atom_idx)
else:
raise TypeError(atom_idx)
# ----------- superpose ------------------------------------------------------------------
def superpose_by_map(self, maps: Dict[Tuple[int, int], Dict[int, int]]):
"""
The map is a dict of a tuple of two indices: the moving one and the fixed mol.
The items are dictionary of atom indices from the former and the correspondance in the latter.
"""
for (moved_idx, fixed_idx), atom_map in maps.items():
Chem.rdMolAlign.AlignMol(self.mols[moved_idx],
self.mols[fixed_idx],
atomMap=list(atom_map.items()))
self.superposeed = True
def superpose_by_mcs(self, fixed_mol_idx=-1, **mcs_settings):
"""
Aligns by MCS. The fix molecule is the last molecule unless specified.
"""
fixed = self.get_mol(fixed_mol_idx)
commons = []
for moved in self.mols[:1]:
res: rdFMCS.MCSResult = rdFMCS.FindMCS([moved, fixed], **mcs_settings)
common = Chem.MolFromSmarts(res.smartsString)
Chem.SanitizeMol(common)
mcs_map = list(zip(moved.GetSubstructMatch(common),
fixed.GetSubstructMatch(common)
)
)
Chem.rdMolAlign.AlignMol(moved, fixed, atomMap=mcs_map)
commons.append(common)
self.superposeed = True
return commons
# ----- transforms ----------------------------------------
def transform(self, mol_idx: int, transform_matrix: np.array):
"""
Applies an affine transform 4x4 matrix to mol_idx.
Use ``translate`` or ``rotate``.
"""
AllChem.TransformConformer(self.get_mol(mol_idx).GetConformer(0),
transform_matrix.astype(np.double)
)
def translate(self, mol_idx: int, x: float, y: float, z: float):
"""
Move the molecule at index ``mol_idx`` of ``.mols``.
"""
translation = np.array([[1, 0, 0, x],
[0, 1, 0, y],
[0, 0, 1, z],
[0, 0, 0, 1]], dtype=np.double)
self.transform(mol_idx, translation)
def rotate(self, mol_idx: int, theta: float, axis: str, degrees: bool = True):
"""
Rotate a molecule around a give angle on an axis.
This is coded terribly.
To pivot around a given atom has to translate the pivot point to the origin, rotate and translate back again.
If this is gibberish, please look up affine transform 4x4 matrices.
This is not done by Euler angle because its a one off thing.
"""
correct_unit = lambda d: np.deg2rad(d) if degrees else d # noqa No lambda PEP8, ha!
theta_rad = correct_unit(theta)
# row, column
ori_rotation = np.zeros((4, 4)) + np.diag(np.ones(4), 0)
if axis == 'x': # roll
first_row = 1
second_row = 2
elif axis == 'y': # pitch
first_row = 0
second_row = 2
elif axis == 'z': # yaw
first_row = 0
second_row = 1
else:
raise ValueError(f'Axis {axis} is not x,y or z')
# for sanity's sake
first_col = first_row
second_col = second_row
ori_rotation[first_row, first_col] = np.cos(theta_rad)
ori_rotation[first_row, second_col] = -np.sin(theta_rad)
ori_rotation[second_row, first_col] = np.sin(theta_rad)
ori_rotation[second_row, second_col] = np.cos(theta_rad)
self.transform(mol_idx, ori_rotation)
def experimental_rotate(self, mol_idx: int, angle_x: float, angle_y: float, angle_z: float, degrees: bool = True):
"""
THIS HAS A BUG.
Rotate a molecule around a give set of angles,
in degrees pivoting on the origin.
These are Tait-Bryan angles.
There is a bug in here. Use rotate instead.
To pivot around a given atom has to translate the pivot point to the origin, rotate and translate back again.
If this is gibberish, please look up affine transform 4x4 matrices.
"""
correct_unit = lambda d: np.deg2rad(d) if degrees else d # noqa
alpha = correct_unit(angle_x)
beta = correct_unit(angle_y)
gamma = correct_unit(angle_z)
ori_rotation = np.array([[np.cos(beta) * np.cos(gamma),
np.sin(alpha) * np.sin(beta) * np.cos(gamma) - np.cos(alpha) * np.sin(gamma),
np.cos(alpha) * np.sin(beta) * np.cos(gamma) + np.sin(alpha) * np.sin(gamma),
0],
[np.cos(beta) * np.sin(gamma),
np.sin(alpha) * np.sin(beta) * np.sin(gamma) + np.cos(alpha) * np.cos(gamma),
np.cos(alpha) * np.sin(beta) * np.sin(gamma) - np.sin(alpha) * np.cos(gamma),
0],
[-np.sin(beta),
np.sin(alpha) * np.cos(beta),
np.cos(alpha) * np.cos(beta),
0],
[0, 0, 0, 1]], dtype=np.double)
self.transform(mol_idx, ori_rotation)