-
Notifications
You must be signed in to change notification settings - Fork 78
/
rdkit.py
149 lines (113 loc) · 4.82 KB
/
rdkit.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
"""
Calls the RDKit package.
"""
from typing import TYPE_CHECKING, Dict
from qcelemental.models import AtomicResult, Provenance
from qcelemental.util import safe_version, which_import
from ..exceptions import InputError
from ..units import ureg
from .model import ProgramHarness
if TYPE_CHECKING:
from qcelemental.models import AtomicInput
from ..config import TaskConfig
class RDKitHarness(ProgramHarness):
_defaults = {
"name": "RDKit",
"scratch": False,
"thread_safe": True,
"thread_parallel": False,
"node_parallel": False,
"managed_memory": False,
}
version_cache: Dict[str, str] = {}
class Config(ProgramHarness.Config):
pass
@staticmethod
def _process_molecule_rdkit(jmol):
from rdkit import Chem
# Handle errors
if abs(jmol.molecular_charge) > 1.0e-6:
raise InputError("RDKit does not currently support charged molecules.")
if not jmol.connectivity: # Check for empty list
raise InputError("RDKit requires molecules to have a connectivity graph.")
# Build out the base molecule
base_mol = Chem.Mol()
rw_mol = Chem.RWMol(base_mol)
for sym in jmol.symbols:
rw_mol.AddAtom(Chem.Atom(sym.title()))
# Add in connectivity
bond_types = {1: Chem.BondType.SINGLE, 2: Chem.BondType.DOUBLE, 3: Chem.BondType.TRIPLE}
for atom1, atom2, bo in jmol.connectivity:
rw_mol.AddBond(atom1, atom2, bond_types[bo])
mol = rw_mol.GetMol()
# Write out the conformer
natom = len(jmol.symbols)
conf = Chem.Conformer(natom)
bohr2ang = ureg.conversion_factor("bohr", "angstrom")
for line in range(natom):
conf.SetAtomPosition(
line,
(
bohr2ang * jmol.geometry[line, 0],
bohr2ang * jmol.geometry[line, 1],
bohr2ang * jmol.geometry[line, 2],
),
)
mol.AddConformer(conf)
Chem.rdmolops.SanitizeMol(mol)
return mol
@staticmethod
def found(raise_error: bool = False) -> bool:
return which_import(
"rdkit",
return_bool=True,
raise_error=raise_error,
raise_msg="Please install via `conda install rdkit -c conda-forge`.",
)
def get_version(self) -> str:
"""Return the currently used version of RDKit."""
self.found(raise_error=True)
which_prog = which_import("rdkit")
if which_prog not in self.version_cache:
import rdkit
self.version_cache[which_prog] = safe_version(rdkit.__version__)
return self.version_cache[which_prog]
def compute(self, input_data: "AtomicInput", config: "TaskConfig") -> "AtomicResult":
"""
Runs RDKit in FF typing
"""
self.found(raise_error=True)
import rdkit
from rdkit.Chem import AllChem
# Failure flag
ret_data = {"success": False}
# Build the Molecule
jmol = input_data.molecule
mol = self._process_molecule_rdkit(jmol)
if input_data.model.method.lower() == "uff":
ff = AllChem.UFFGetMoleculeForceField(mol)
all_params = AllChem.UFFHasAllMoleculeParams(mol)
elif input_data.model.method.lower() in ["mmff94", "mmff94s"]:
props = AllChem.MMFFGetMoleculeProperties(mol, mmffVariant=input_data.model.method)
ff = AllChem.MMFFGetMoleculeForceField(mol, props)
all_params = AllChem.MMFFHasAllMoleculeParams(mol)
else:
raise InputError("RDKit only supports the UFF, MMFF94, and MMFF94s methods currently.")
if all_params is False:
raise InputError("RDKit parameters not found for all atom types in molecule.")
ff.Initialize()
ret_data["properties"] = {"return_energy": ff.CalcEnergy() * ureg.conversion_factor("kJ / mol", "hartree")}
if input_data.driver == "energy":
ret_data["return_result"] = ret_data["properties"]["return_energy"]
elif input_data.driver == "gradient":
coef = ureg.conversion_factor("kJ / mol", "hartree") * ureg.conversion_factor("angstrom", "bohr")
ret_data["return_result"] = [x * coef for x in ff.CalcGrad()]
else:
raise InputError(f"Driver {input_model.driver} not implemented for RDKit.")
ret_data["provenance"] = Provenance(
creator="rdkit", version=rdkit.__version__, routine="rdkit.Chem.AllChem.UFFGetMoleculeForceField"
)
ret_data["schema_name"] = "qcschema_output"
ret_data["success"] = True
# Form up a dict first, then sent to BaseModel to avoid repeat kwargs which don't override each other
return AtomicResult(**{**input_data.dict(), **ret_data})