mirror of
https://github.com/gcorso/DiffDock.git
synced 2026-06-04 18:04:23 +08:00
561f70a897
* Ensure we calculate rotatable bonds on the version of the ligand with no hydrogens. Also fix spelling of rotable -> rotatable. Closes GH-220 (@Nobody-Zhang) * Vectorize SO3 calculations. Closes PR GH-218 (@tornikeo) * Pin pytorch-lightning version. Closes GH-193 (@mikael-h-christensen) * Guard against divide by zero in torus.py. Closes GH-161 (@amorehead) * Update e3nn version to 0.5.1. Closes GH-155 (@amorehead) * Add a little more info on docker container to README.md
86 lines
3.0 KiB
Python
86 lines
3.0 KiB
Python
import copy, time
|
|
import numpy as np
|
|
from collections import defaultdict
|
|
from rdkit import Chem, RDLogger
|
|
from rdkit.Chem import AllChem, rdMolTransforms
|
|
from rdkit import Geometry
|
|
import networkx as nx
|
|
from scipy.optimize import differential_evolution
|
|
|
|
RDLogger.DisableLog('rdApp.*')
|
|
|
|
"""
|
|
Conformer matching routines from Torsional Diffusion
|
|
"""
|
|
|
|
def GetDihedral(conf, atom_idx):
|
|
return rdMolTransforms.GetDihedralRad(conf, atom_idx[0], atom_idx[1], atom_idx[2], atom_idx[3])
|
|
|
|
|
|
def SetDihedral(conf, atom_idx, new_vale):
|
|
rdMolTransforms.SetDihedralRad(conf, atom_idx[0], atom_idx[1], atom_idx[2], atom_idx[3], new_vale)
|
|
|
|
|
|
def apply_changes(mol, values, rotatable_bonds, conf_id):
|
|
opt_mol = copy.copy(mol)
|
|
[SetDihedral(opt_mol.GetConformer(conf_id), rotatable_bonds[r], values[r]) for r in range(len(rotatable_bonds))]
|
|
return opt_mol
|
|
|
|
|
|
def optimize_rotatable_bonds(mol, true_mol, rotatable_bonds, probe_id=-1, ref_id=-1, seed=0, popsize=15, maxiter=500,
|
|
mutation=(0.5, 1), recombination=0.8):
|
|
opt = OptimizeConformer(mol, true_mol, rotatable_bonds, seed=seed, probe_id=probe_id, ref_id=ref_id)
|
|
max_bound = [np.pi] * len(opt.rotatable_bonds)
|
|
min_bound = [-np.pi] * len(opt.rotatable_bonds)
|
|
bounds = (min_bound, max_bound)
|
|
bounds = list(zip(bounds[0], bounds[1]))
|
|
|
|
# Optimize conformations
|
|
result = differential_evolution(opt.score_conformation, bounds,
|
|
maxiter=maxiter, popsize=popsize,
|
|
mutation=mutation, recombination=recombination, disp=False, seed=seed)
|
|
opt_mol = apply_changes(opt.mol, result['x'], opt.rotatable_bonds, conf_id=probe_id)
|
|
|
|
return opt_mol
|
|
|
|
|
|
class OptimizeConformer:
|
|
def __init__(self, mol, true_mol, rotatable_bonds, probe_id=-1, ref_id=-1, seed=None):
|
|
super(OptimizeConformer, self).__init__()
|
|
if seed:
|
|
np.random.seed(seed)
|
|
self.rotatable_bonds = rotatable_bonds
|
|
self.mol = mol
|
|
self.true_mol = true_mol
|
|
self.probe_id = probe_id
|
|
self.ref_id = ref_id
|
|
|
|
def score_conformation(self, values):
|
|
for i, r in enumerate(self.rotatable_bonds):
|
|
SetDihedral(self.mol.GetConformer(self.probe_id), r, values[i])
|
|
return AllChem.AlignMol(self.mol, self.true_mol, self.probe_id, self.ref_id)
|
|
|
|
|
|
def get_torsion_angles(mol):
|
|
torsions_list = []
|
|
G = nx.Graph()
|
|
for i, atom in enumerate(mol.GetAtoms()):
|
|
G.add_node(i)
|
|
nodes = set(G.nodes())
|
|
for bond in mol.GetBonds():
|
|
start, end = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
|
|
G.add_edge(start, end)
|
|
for e in G.edges():
|
|
G2 = copy.deepcopy(G)
|
|
G2.remove_edge(*e)
|
|
if nx.is_connected(G2): continue
|
|
l = list(sorted(nx.connected_components(G2), key=len)[0])
|
|
if len(l) < 2: continue
|
|
n0 = list(G2.neighbors(e[0]))
|
|
n1 = list(G2.neighbors(e[1]))
|
|
torsions_list.append(
|
|
(n0[0], e[0], e[1], n1[0])
|
|
)
|
|
return torsions_list
|
|
|