calc utils

calc_utils/futils/rdkit2ase.py

@@ -0,0 +1,145 @@
+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