Scan dihedral angles using Open Babel (here using random angles)

scan-dihedrals.py

#!/usr/bin/env python

from __future__ import print_function

import sys
import os
import math
import random

import pybel
from pybel import ob

#  optimize dihedral angle vectors
#    using MMFF94 energies and Bayesian optimization


# global variable for the MMFF94 force field
ff = pybel._forcefields["mmff94"]

def energy():
    return ff.Energy(False)  # false = don't compute gradients


if __name__ == "__main__":
    # Syntax  dihedrals.py  filename.mol2
    # read molecule
    if (sys.argv) != 2:
        exit

    filename = sys.argv[1]
    extension = os.path.splitext(filename)[1]
    # read a single molecule from the supplied file
    # turn this into a while loop if you want all molecules in a file
    mol = next(pybel.readfile(extension[1:], filename))
    if ff.Setup(mol.OBMol) is False:
        print("Cannot set up force field, exiting")
        exit()

    print("Molecule read. Number of atoms: {}".format(len(mol.atoms)))

    # set up the pool of rotatable bonds
    rl = ob.OBRotorList()
    rl.Setup(mol.OBMol)
    print("Number of rotatable bonds: {}".format(rl.Size()))

    # iterate through them to get the current dihedrals
    currentCoords = mol.OBMol.GetCoordinates()
    rotIterator = rl.BeginRotors()
    rotor = rl.BeginRotor(rotIterator)  # first rotatable bond
    rotors = []
    angles = []  # in radians
    while rotor is not None:
        angle = rotor.CalcTorsion(currentCoords)
        angles.append(angle)
        print("dihedral: {0:3f}".format(math.degrees(angle)))
        atoms = rotor.GetDihedralAtoms()
        
        print("GetDihedralAtoms:", rotor.GetDihedralAtoms())
        rotors.append(rotor)
        rotor = rl.NextRotor(rotIterator)
    print("Current energy", round(energy(), 4))

    # each angle can obviously go from 0 .. 2*pi
    # pick a random set of angles
    # this would obviously be driven by the Bayesian optimization
    for i in range(len(angles)):
        angle = random.uniform(0, 2.0 * math.pi)
        print("dihedral: {0:.3f}".format(math.degrees(angle)))
        angles[i] = angle
        rotors[i].SetToAngle(currentCoords, angle)
    # update the coordinates and re-compute the energy
    mol.OBMol.SetCoordinates(currentCoords)
    ff.SetCoordinates(mol.OBMol)
    print("New energy: {0:.4f}".format(energy(), 4))

    exit()

    # write the current coordinates to a file
    newFile = os.path.splitext(filename)[0] + '-opt.sdf'
    mol.write("sdf", newFile)