Source code for molfoundry.conformer

import os
import tempfile
from io import StringIO
from typing import Optional, Dict, Tuple
from .exceptions import *

from .graph import Graph




class Conformer:
    def __init__(self, graph: Graph) -> None:
        if not graph.isMolecule:
            raise LogicError("Trying to create conformer from non-molecule graph.")
        self._graph = graph
        self._embedded = False
        self._vertex_coords: Dict[int, Tuple[float, float, float]] = {
            v.id: (0, 0, 0)
            for v in graph.vertices
        }

    # noinspection PyShadowingBuiltins


    def getVertexCoordinate(self, id: int) -> Tuple[float, float, float]:
        return self._vertex_coords.get(id, (0, 0, 0))




    def embedRDKit(self, seed: int = 42, num_confs: int = 50) -> None:
        try:
            # noinspection PyUnresolvedReferences
            from rdkit import Chem
            # noinspection PyUnresolvedReferences
            from rdkit.Chem import AllChem
        except ImportError:
            raise LogicError("Failed to import RDKit. Please make sure to install RDKit.")

        mol = Chem.MolFromSmiles(self._graph.smilesWithIds, sanitize=False)
        if mol is None:
            raise LogicError("RDKit could not parse graph.")

        Chem.SanitizeMol(mol)
        # Ensure no conformers are present
        mol.RemoveAllConformers()

        params = AllChem.ETKDGv3()
        params.randomSeed = seed
        params.pruneRmsThresh = 0.2
        params.useRandomCoords = True
        params.useExpTorsionAnglePrefs = True
        params.useBasicKnowledge = True

        cids = AllChem.EmbedMultipleConfs(mol, numConfs=num_confs, params=params)
        if not cids:
            cid = AllChem.EmbedMolecule(mol, useRandomCoords=True, randomSeed=seed)
            if cid < 0:
                raise LogicError("Embedding failed.")
            cids = [cid]

        energies = []
        if AllChem.MMFFHasAllMoleculeParams(mol):
            props = AllChem.MMFFGetMoleculeProperties(mol, mmffVariant="MMFF94s")
            for cid in cids:
                ff = AllChem.MMFFGetMoleculeForceField(mol, props, confId=cid)
                ff.Minimize(maxIts=5000)
                energies.append((cid, ff.CalcEnergy()))
        else:
            for cid in cids:
                ff = AllChem.UFFGetMoleculeForceField(mol, confId=cid)
                ff.Minimize(maxIts=5000)
                energies.append((cid, ff.CalcEnergy()))

        best_cid, best_e = min(energies, key=lambda x: x[1])
        # Keep best
        for cid in sorted([c for c in cids if c != best_cid], reverse=True):
            mol.RemoveConformer(cid)
        best_conformer = mol.GetConformer()

        atom_id_index_map = {}
        for atom in mol.GetAtoms():
            if atom.HasProp("molAtomMapNumber"):
                aid = int(atom.GetProp("molAtomMapNumber"))
            else:
                raise LogicError(f"Atom without molAtomMapNumber: idx={atom.GetIdx()} symbol={atom.GetSymbol()}")
            if aid in atom_id_index_map:
                raise LogicError(f"Duplicate atom map number {aid}")
            atom_id_index_map[aid] = atom.GetIdx()

        n = mol.GetNumAtoms()
        expected = set(range(n))
        got = set(atom_id_index_map.keys())
        if got != expected:
            raise LogicError(f"AtomMap IDs mismatch. expected={sorted(expected)} got={sorted(got)}")

        for x in sorted(atom_id_index_map.keys()):
            idx = atom_id_index_map[x]
            coord = best_conformer.GetAtomPosition(idx)
            self._vertex_coords[x] = (coord.x, coord.y, coord.z)
        self._embedded = True




    def optimizeXTB(self, fmax=0.1, steps=50):
        """
        Local geometry optimization using BFGS

        fmax   – Energy barrier (eV/Å)
        steps  – max BFGS steps
        """
        if not self._embedded:
            raise LogicError("Trying to optimize conformer without prior embedding. Call embed*(...) first.")

        try:
            # noinspection PyUnresolvedReferences
            from ase.io import read
            # noinspection PyUnresolvedReferences
            from ase.optimize import BFGS
            # noinspection PyUnresolvedReferences
            from tblite.ase import TBLite
            from contextlib import redirect_stdout, redirect_stderr
        except ImportError:
            raise LogicError("Failed to import ASE or TBLite. Please make sure to install ASE and TBLite.")

        atoms = read(StringIO(self.xyz()))
        atoms = atoms.copy()
        atoms.calc = TBLite(method="GFN2-xTB")
        opt = BFGS(atoms, trajectory=f"{tempfile.gettempdir()}/__mf__conf__opt.traj",
                   logfile=f"{tempfile.gettempdir()}/__mf__conf__opt.log")
        try:
            # Prevent log spamming from TBLite
            with open(os.devnull, "w") as fnull, redirect_stdout(fnull), redirect_stderr(fnull):
                opt.run(fmax=fmax, steps=steps)
        except CalculationFailed as e:
            raise LogicError(f"Failed to optimize embedding: {e}")
        for i, atom in enumerate(atoms):
            x, y, z = atom.position
            self._vertex_coords[i] = (x, y, z)




    def xyz(self, comment: Optional[str] = None) -> str:
        if not self._embedded:
            print("Warning: Generating conformer xyz without prior embedding. Coordinates will all be zero.")
        result = f"{self._graph.numVertices}\n"
        result += (comment or "Generated by MolFoundry").rstrip("\n\r") + "\n"
        vertices = self._graph.vertices
        for _id in sorted(self._vertex_coords.keys()):
            coord = self._vertex_coords[_id]
            symbol = vertices[_id].atomId.symbol
            result += f"{symbol:<2} {coord[0]: .6f} {coord[1]: .6f} {coord[2]: .6f}\n"
        return result