146 lines
4.8 KiB
Python
146 lines
4.8 KiB
Python
from rdkit import Chem
|
|
from ase import atoms
|
|
from typing import Optional
|
|
import numpy as np
|
|
from rdkit.Geometry import Point3D
|
|
|
|
|
|
def MolToAtoms(mol: Chem.Mol, confID=-1) -> atoms.Atoms:
|
|
conf = mol.GetConformer(confID)
|
|
symbols = [atom.GetSymbol() for atom in mol.GetAtoms()]
|
|
positions = [
|
|
(conf.GetAtomPosition(i).x,
|
|
conf.GetAtomPosition(i).y,
|
|
conf.GetAtomPosition(i).z)
|
|
for i in range(mol.GetNumAtoms())
|
|
]
|
|
aseAtoms = atoms.Atoms(symbols=symbols, positions=positions)
|
|
charges = [atom.GetFormalCharge() for atom in mol.GetAtoms()]
|
|
aseAtoms.set_initial_charges(charges)
|
|
|
|
return aseAtoms
|
|
|
|
|
|
def AtomsToMol(
|
|
atoms: atoms.Atoms,
|
|
mol: Optional[Chem.Mol] = None,
|
|
conf_id: int = 0,
|
|
charge: Optional[int] = None,
|
|
allow_reorder: bool = False,
|
|
inplace = False
|
|
) -> Chem.Mol:
|
|
"""
|
|
Convert ASE Atoms -> RDKit Mol.
|
|
|
|
If mol is provided:
|
|
- verify natoms and element symbols match (unless allow_reorder=True)
|
|
- update (or add) conformer coordinates so mol matches atoms
|
|
|
|
If mol is None:
|
|
- create a new Mol with atoms only, set 3D coords
|
|
- Determine bonds from geometry using rdDetermineBonds.DetermineBonds()
|
|
|
|
Parameters
|
|
----------
|
|
atoms : ase.Atoms
|
|
Must have positions (Å).
|
|
mol : rdkit.Chem.Mol | None
|
|
Optional template mol.
|
|
conf_id : int
|
|
Conformer id to update/use. If mol has no conformer, one will be added.
|
|
charge : int | None
|
|
Total molecular charge used by DetermineBonds (recommended if known).
|
|
If None, will try to infer from formal charges when mol is provided;
|
|
otherwise defaults to 0 for new mol.
|
|
allow_reorder : bool
|
|
If True, will not enforce symbol-by-index equality (only checks counts).
|
|
Most of the time you want False to guarantee consistency.
|
|
inplace : bool
|
|
If True, will modify input mol instead of make a copy.
|
|
|
|
Returns
|
|
-------
|
|
rdkit.Chem.Mol
|
|
"""
|
|
positions = np.asarray(atoms.get_positions(), dtype=float)
|
|
symbols = list(atoms.get_chemical_symbols())
|
|
n = len(symbols)
|
|
|
|
# -------- case 1: update existing mol --------
|
|
if mol is not None:
|
|
if mol.GetNumAtoms() != n:
|
|
raise ValueError(f"mol has {mol.GetNumAtoms()} atoms but ASE atoms has {n}.")
|
|
|
|
mol_symbols = [a.GetSymbol() for a in mol.GetAtoms()]
|
|
if not allow_reorder:
|
|
if mol_symbols != symbols:
|
|
raise ValueError(
|
|
"Element symbols mismatch by index between mol and atoms.\n"
|
|
f"mol: {mol_symbols}\n"
|
|
f"atoms: {symbols}\n"
|
|
"If you REALLY know what you're doing, set allow_reorder=True "
|
|
"and handle mapping yourself."
|
|
)
|
|
else:
|
|
# only check multiset counts
|
|
if sorted(mol_symbols) != sorted(symbols):
|
|
raise ValueError("Element symbol counts differ between mol and atoms.")
|
|
|
|
if not inplace:
|
|
mol = Chem.Mol(mol) # copy
|
|
if mol.GetNumConformers() == 0:
|
|
conf = Chem.Conformer(n)
|
|
conf.Set3D(True)
|
|
mol.AddConformer(conf, assignId=True)
|
|
|
|
# pick conformer
|
|
try:
|
|
conf = mol.GetConformer(conf_id)
|
|
except ValueError:
|
|
# create new conformer if requested id doesn't exist
|
|
conf = Chem.Conformer(n)
|
|
conf.Set3D(True)
|
|
mol.AddConformer(conf, assignId=True)
|
|
conf = mol.GetConformer(mol.GetNumConformers() - 1)
|
|
if conf_id!=0:
|
|
print("Warning: Failed to pick conformer.")
|
|
|
|
for i in range(n):
|
|
x, y, z = map(float, positions[i])
|
|
conf.SetAtomPosition(i, Point3D(x, y, z))
|
|
|
|
# charge inference if not given
|
|
if charge is None:
|
|
charge = int(sum(a.GetFormalCharge() for a in mol.GetAtoms()))
|
|
|
|
return mol
|
|
|
|
# -------- case 2: build mol + determine bonds --------
|
|
rw = Chem.RWMol()
|
|
for sym in symbols:
|
|
rw.AddAtom(Chem.Atom(sym))
|
|
|
|
new_mol = rw.GetMol()
|
|
conf = Chem.Conformer(n)
|
|
conf.Set3D(True)
|
|
for i in range(n):
|
|
x, y, z = map(float, positions[i])
|
|
conf.SetAtomPosition(i, Point3D(x, y, z))
|
|
new_mol.AddConformer(conf, assignId=True)
|
|
|
|
# Determine bonds from geometry
|
|
if charge is None:
|
|
charge = 0
|
|
|
|
try:
|
|
from rdkit.Chem import rdDetermineBonds
|
|
rdDetermineBonds.DetermineBonds(new_mol, charge=charge)
|
|
except Exception as e:
|
|
raise RuntimeError(
|
|
"DetermineBonds failed. This can happen for metals/ions/odd geometries.\n"
|
|
"Consider providing a template mol, or implement custom distance-based bonding.\n"
|
|
f"Original error: {e}"
|
|
)
|
|
|
|
return new_mol
|