Source code for stko._internal.optimizers.open_babel

import logging
from pathlib import Path

import numpy as np

try:
    from openbabel import openbabel
except ImportError:
    openbabel = None

from stko._internal.internal_types import MoleculeT
from stko._internal.optimizers.optimizers import Optimizer
from stko._internal.utilities.exceptions import (
    ForceFieldSetupError,
    WrapperNotInstalledError,
)

logger = logging.getLogger(__name__)


[docs] class OpenBabel(Optimizer): """Use OpenBabel to optimize molecules with forcefields. .. warning:: This will not work for Python >= 3.13! See https://github.com/JelfsMaterialsGroup/stko/issues/208 Parameters: forcefield: Forcefield to use. Options include `uff`, `gaff`, `ghemical`, `mmff94`. repeat_steps: Number of optimisation steps. Each optimisation step contains `sd_steps` steepest descent and then `cg_steps` conjugate gradient runs. sd_steps: Number of steepest descent steps per optimisations. cg_steps: Number of conjugate gradient steps per optimisations. Raises: :class:`WrapperNotInstalledError`: if `openbabel` not installed. .. warning:: this optimizer seems to be machine dependant, producing different energies after optimisation on Ubunut 18 vs. Ubuntu 20. See Also: * OpenBabel: https://github.com/openbabel/openbabel Examples: .. code-block:: python import stk import stko mol = stk.BuildingBlock('NCCNCCN') openbabel = stko.OpenBabel('uff') opt_mol = openbabel.optimize(mol) """ def __init__( self, forcefield: str, repeat_steps: int = 10, sd_steps: int = 50, cg_steps: int = 50, ) -> None: if openbabel is None: msg = "openbabel is not installed; see README for installation." raise WrapperNotInstalledError(msg) self._forcefield = forcefield self._repeat_steps = repeat_steps self._sd_steps = sd_steps self._cg_steps = cg_steps
[docs] def optimize(self, mol: MoleculeT) -> MoleculeT: temp_file = Path("temp.mol") mol.write(temp_file) try: ob_conversion = openbabel.OBConversion() ob_conversion.SetInFormat("mol") ob_mol = openbabel.OBMol() ob_conversion.ReadFile(ob_mol, str(temp_file)) ob_mol.PerceiveBondOrders() finally: temp_file.unlink() forcefield = openbabel.OBForceField.FindForceField(self._forcefield) outcome = forcefield.Setup(ob_mol) if not outcome: msg = f"Openbabel: {self._forcefield} could not be setup for {mol}" raise ForceFieldSetupError(msg) for _ in range(self._repeat_steps): forcefield.SteepestDescent(self._sd_steps) forcefield.GetCoordinates(ob_mol) forcefield.ConjugateGradients(self._cg_steps) forcefield.GetCoordinates(ob_mol) position_matrix = [ np.array([atom.GetX(), atom.GetY(), atom.GetZ()]) for atom in openbabel.OBMolAtomIter(ob_mol) ] return mol.with_position_matrix(np.asarray(position_matrix))