--- a +++ b/unimol/utils/geom.py @@ -0,0 +1,53 @@ +import random +import numpy as np +import rdkit +import rdkit.Chem as Chem +import rdkit.Chem.AllChem as AllChem +import rdkit.Chem.Descriptors as Descriptors +import tqdm + +def change_torsion(pointset, index, angle=None): + # pointset: (N, 3) + #return the new pointset + # angle: rad + # index: (2) + if angle is None: + angle = random.uniform(-np.pi, np.pi) + pointset = np.array(pointset) + #get rotate axis + axis = pointset[index[0][-1]] - pointset[index[1][0]] + axis = axis / np.linalg.norm(axis) + #get rotate matrix + cos = np.cos(angle) + sin = np.sin(angle) + R = np.array([[cos + axis[0] ** 2 * (1 - cos), axis[0] * axis[1] * (1 - cos) - axis[2] * sin, axis[0] * axis[2] * (1 - cos) + axis[1] * sin], + [axis[1] * axis[0] * (1 - cos) + axis[2] * sin, cos + axis[1] ** 2 * (1 - cos), axis[1] * axis[2] * (1 - cos) - axis[0] * sin], + [axis[2] * axis[0] * (1 - cos) - axis[1] * sin, axis[2] * axis[1] * (1 - cos) + axis[0] * sin, cos + axis[2] ** 2 * (1 - cos)]]) + #rotate + pointset[index[1]] = np.dot(pointset[index[1]] - pointset[index[1][0]], R) + pointset[index[1][0]] + return pointset, angle + +def gen_conformation(mol, num_conf=20, num_worker=8): + try: + mol = Chem.AddHs(mol) + AllChem.EmbedMultipleConfs(mol, numConfs=num_conf, numThreads=num_worker, pruneRmsThresh=0.1, maxAttempts=5, useRandomCoords=False) + AllChem.MMFFOptimizeMoleculeConfs(mol, numThreads=num_worker) + mol = Chem.RemoveHs(mol) + return mol + except: + return None + +def RotatableBond(mol): + """Get the rotatable bond index of a molecule. + Args: + mol: rdkit.Chem.rdchem.Mol + Returns: + rotatable_bond: list of tuple + """ + rotatable_bond = [] + for bond in mol.GetBonds(): + if bond.IsInRing(): + continue + if bond.GetBondType() == Chem.rdchem.BondType.SINGLE: + rotatable_bond.append((bond.GetBeginAtomIdx(), bond.GetEndAtomIdx())) + return rotatable_bond \ No newline at end of file