Add pymatgen.io.openff module

.github/workflows/test.yml

@@ -111,6 +111,8 @@ jobs:
           # TODO remove next line installing ase from main branch when FrechetCellFilter is released
           uv pip install --upgrade 'git+https://gitlab.com/ase/ase' --system
 
+          uv pip install --upgrade git+https://github.com/openforcefield/openff-toolkit.git --system
+
       - name: pytest split ${{ matrix.split }}
         run: |
           pytest --splits 10 --group ${{ matrix.split }} --durations-path tests/files/.pytest-split-durations tests

pymatgen/io/openff.py

@@ -0,0 +1,305 @@
+"""Utility functions for classical md subpackage."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import openff.toolkit as tk
+from openff.units import Quantity, unit
+
+import pymatgen
+from pymatgen.analysis.graphs import MoleculeGraph
+from pymatgen.analysis.local_env import OpenBabelNN, metal_edge_extender
+from pymatgen.core import Element, Molecule
+
+
+def molgraph_to_openff_mol(molgraph: MoleculeGraph) -> tk.Molecule:
+    """
+    Convert a Pymatgen MoleculeGraph to an OpenFF Molecule.
+
+    Args:
+        molgraph (MoleculeGraph): The Pymatgen MoleculeGraph to be converted.
+
+    Returns:
+        tk.Molecule: The converted OpenFF Molecule.
+    """
+    # create empty openff_mol and prepare a periodic table
+    p_table = {str(el): el.Z for el in Element}
+    openff_mol = tk.Molecule()
+
+    # set atom properties
+    partial_charges = []
+    # TODO: should assert that there is only one molecule
+    for i_node in range(len(molgraph.graph.nodes)):
+        node = molgraph.graph.nodes[i_node]
+        atomic_number = node.get("atomic_number") or p_table[molgraph.molecule[i_node].species_string]
+
+        # put formal charge on first atom if there is none present
+        formal_charge = node.get("formal_charge")
+        if formal_charge is None:
+            formal_charge = (i_node == 0) * molgraph.molecule.charge * unit.elementary_charge
+
+        # assume not aromatic if no info present
+        is_aromatic = node.get("is_aromatic") or False
+
+        openff_mol.add_atom(atomic_number, formal_charge, is_aromatic=is_aromatic)
+
+        # add to partial charge array
+        partial_charge = node.get("partial_charge")
+        if isinstance(partial_charge, Quantity):
+            partial_charge = partial_charge.magnitude
+        partial_charges.append(partial_charge)
+
+    charge_array = np.array(partial_charges)
+    if np.not_equal(charge_array, None).all():
+        openff_mol.partial_charges = charge_array * unit.elementary_charge
+
+    # set edge properties, default to single bond and assume not aromatic
+    for i_node, j, bond_data in molgraph.graph.edges(data=True):
+        bond_order = bond_data.get("bond_order") or 1
+        is_aromatic = bond_data.get("is_aromatic") or False
+        openff_mol.add_bond(i_node, j, bond_order, is_aromatic=is_aromatic)
+
+    openff_mol.add_conformer(molgraph.molecule.cart_coords * unit.angstrom)
+    return openff_mol
+
+
+def molgraph_from_openff_mol(molecule: tk.Molecule):
+    """
+    This is designed to closely mirror the graph structure generated by tk.Molecule.to_networkx
+
+    Args:
+        molecule (tk.Molecule): The OpenFF Molecule to convert.
+
+    Returns:
+        MoleculeGraph: The converted MoleculeGraph.
+    """
+    molgraph = MoleculeGraph.with_empty_graph(
+        Molecule([], []),
+        name="none",
+    )
+    p_table = {el.Z: str(el) for el in Element}
+    total_charge = 0
+    cum_atoms = 0
+
+    coords = molecule.conformers[0].magnitude if molecule.conformers is not None else np.zeros((molecule.n_atoms, 3))
+    for j, atom in enumerate(molecule.atoms):
+        molgraph.insert_node(
+            cum_atoms + j,
+            p_table[atom.atomic_number],
+            coords[j, :],
+        )
+        molgraph.graph.nodes[cum_atoms + j]["atomic_number"] = atom.atomic_number
+        molgraph.graph.nodes[cum_atoms + j]["is_aromatic"] = atom.is_aromatic
+        molgraph.graph.nodes[cum_atoms + j]["stereochemistry"] = atom.stereochemistry
+        # set partial charge as a pure float
+        partial_charge = None if atom.partial_charge is None else atom.partial_charge.magnitude
+        molgraph.graph.nodes[cum_atoms + j]["partial_charge"] = partial_charge
+        # set formal charge as a pure float
+        formal_charge = atom.formal_charge.magnitude  # type: ignore
+        molgraph.graph.nodes[cum_atoms + j]["formal_charge"] = formal_charge
+        total_charge += formal_charge
+    for bond in molecule.bonds:
+        molgraph.graph.add_edge(
+            cum_atoms + bond.atom1_index,
+            cum_atoms + bond.atom2_index,
+            bond_order=bond.bond_order,
+            is_aromatic=bond.is_aromatic,
+            stereochemistry=bond.stereochemistry,
+        )
+    # formal_charge += molecule.total_charge
+    cum_atoms += molecule.n_atoms
+    molgraph.molecule.set_charge_and_spin(charge=total_charge)
+    return molgraph
+
+
+def get_atom_map(inferred_mol: tk.Molecule, openff_mol: tk.Molecule) -> tuple[bool, dict[int, int]]:
+    """
+    Compute an atom mapping between two OpenFF Molecules.
+
+    Attempts to find an isomorphism between the molecules, considering various matching
+    criteria such as formal charges, stereochemistry, and bond orders. Returns the atom
+    mapping if an isomorphism is found, otherwise returns an empty mapping.
+
+    Args:
+        inferred_mol (tk.Molecule): The first OpenFF Molecule.
+        openff_mol (tk.Molecule): The second OpenFF Molecule.
+
+    Returns:
+        Tuple[bool, Dict[int, int]]: A tuple containing a boolean indicating if an
+            isomorphism was found and a dictionary representing the atom mapping.
+    """
+    # do not apply formal charge restrictions
+    kwargs = dict(
+        return_atom_map=True,
+        formal_charge_matching=False,
+    )
+    isomorphic, atom_map = tk.topology.Molecule.are_isomorphic(openff_mol, inferred_mol, **kwargs)
+    if isomorphic:
+        return True, atom_map
+    # relax stereochemistry restrictions
+    kwargs["atom_stereochemistry_matching"] = False
+    kwargs["bond_stereochemistry_matching"] = False
+    isomorphic, atom_map = tk.topology.Molecule.are_isomorphic(openff_mol, inferred_mol, **kwargs)
+    if isomorphic:
+        return True, atom_map
+    # relax bond order restrictions
+    kwargs["bond_order_matching"] = False
+    isomorphic, atom_map = tk.topology.Molecule.are_isomorphic(openff_mol, inferred_mol, **kwargs)
+    if isomorphic:
+        return True, atom_map
+    return False, {}
+
+
+def infer_openff_mol(
+    mol_geometry: pymatgen.core.Molecule,
+) -> tk.Molecule:
+    """
+    Infer an OpenFF Molecule from a Pymatgen Molecule.
+
+    Constructs a MoleculeGraph from the Pymatgen Molecule using the OpenBabelNN local
+    environment strategy and extends metal edges. Converts the resulting MoleculeGraph
+    to an OpenFF Molecule using molgraph_to_openff_mol.
+
+    Args:
+        mol_geometry (pymatgen.core.Molecule): The Pymatgen Molecule to infer from.
+
+    Returns:
+        tk.Molecule: The inferred OpenFF Molecule.
+    """
+    molgraph = MoleculeGraph.with_local_env_strategy(mol_geometry, OpenBabelNN())
+    molgraph = metal_edge_extender(molgraph)
+    return molgraph_to_openff_mol(molgraph)
+
+
+def add_conformer(
+    openff_mol: tk.Molecule, geometry: pymatgen.core.Molecule | None
+) -> tuple[tk.Molecule, dict[int, int]]:
+    """
+    Add conformers to an OpenFF Molecule based on the provided geometry.
+
+    If a geometry is provided, infers an OpenFF Molecule from it,
+    finds an atom mapping between the inferred molecule and the
+    input molecule, and adds the conformer coordinates to the input
+    molecule. If no geometry is provided, generates a single conformer.
+
+    Args:
+        openff_mol (tk.Molecule): The OpenFF Molecule to add conformers to.
+        geometry (Union[pymatgen.core.Molecule, None]): The geometry to use for adding
+            conformers.
+
+    Returns:
+        Tuple[tk.Molecule, Dict[int, int]]: A tuple containing the updated OpenFF
+            Molecule with added conformers and a dictionary representing the atom
+            mapping.
+    """
+    # TODO: test this
+    if geometry:
+        # for geometry in geometries:
+        inferred_mol = infer_openff_mol(geometry)
+        is_isomorphic, atom_map = get_atom_map(inferred_mol, openff_mol)
+        if not is_isomorphic:
+            raise ValueError(
+                f"An isomorphism cannot be found between smile {openff_mol.to_smiles()}"
+                f"and the provided molecule {geometry}."
+            )
+        new_mol = pymatgen.core.Molecule.from_sites([geometry.sites[i] for i in atom_map.values()])
+        openff_mol.add_conformer(new_mol.cart_coords * unit.angstrom)
+    else:
+        atom_map = {i: i for i in range(openff_mol.n_atoms)}
+        openff_mol.generate_conformers(n_conformers=1)
+    return openff_mol, atom_map
+
+
+def assign_partial_charges(
+    openff_mol: tk.Molecule,
+    atom_map: dict[int, int],
+    charge_method: str,
+    partial_charges: None | list[float],
+) -> tk.Molecule:
+    """
+    Assign partial charges to an OpenFF Molecule.
+
+    If partial charges are provided, assigns them to the molecule
+    based on the atom mapping. If the molecule has only one atom,
+    assigns the total charge as the partial charge. Otherwise,
+    assigns partial charges using the specified charge method.
+
+    Args:
+        openff_mol (tk.Molecule): The OpenFF Molecule to assign partial charges to.
+        atom_map (Dict[int, int]): A dictionary representing the atom mapping.
+        charge_method (str): The charge method to use if partial charges are
+            not provided.
+        partial_charges (Union[None, List[float]]): A list of partial charges to
+            assign or None to use the charge method.
+
+    Returns:
+        tk.Molecule: The OpenFF Molecule with assigned partial charges.
+    """
+    # TODO: test this
+    # assign partial charges
+    if partial_charges is not None:
+        partial_charges = np.array(partial_charges)
+        chargs = partial_charges[list(atom_map.values())]  # type: ignore[index, call-overload]
+        openff_mol.partial_charges = chargs * unit.elementary_charge
+    elif openff_mol.n_atoms == 1:
+        openff_mol.partial_charges = np.array([openff_mol.total_charge.magnitude]) * unit.elementary_charge
+    else:
+        openff_mol.assign_partial_charges(charge_method)
+    return openff_mol
+
+
+def create_openff_mol(
+    smile: str,
+    geometry: pymatgen.core.Molecule | str | Path | None = None,
+    charge_scaling: float = 1,
+    partial_charges: list[float] | None = None,
+    backup_charge_method: str = "am1bcc",
+) -> tk.Molecule:
+    """
+    Create an OpenFF Molecule from a SMILES string and optional geometry.
+
+    Constructs an OpenFF Molecule from the provided SMILES
+    string, adds conformers based on the provided geometry (if
+    any), assigns partial charges using the specified method
+    or provided partial charges, and applies charge scaling.
+
+    Args:
+        smile (str): The SMILES string of the molecule.
+        geometry (Union[pymatgen.core.Molecule, str, Path, None], optional): The
+            geometry to use for adding conformers. Can be a Pymatgen Molecule,
+            file path, or None.
+        charge_scaling (float, optional): The scaling factor for partial charges.
+            Default is 1.
+        partial_charges (Union[List[float], None], optional): A list of partial
+            charges to assign, or None to use the charge method.
+        backup_charge_method (str, optional): The backup charge method to use if
+            partial charges are not provided. Default is "am1bcc".
+
+    Returns:
+        tk.Molecule: The created OpenFF Molecule.
+    """
+    if isinstance(geometry, (str, Path)):
+        geometry = pymatgen.core.Molecule.from_file(str(geometry))
+
+    if partial_charges is not None:
+        if geometry is None:
+            raise ValueError("geometries must be set if partial_charges is set")
+        if len(partial_charges) != len(geometry):
+            raise ValueError("partial charges must have same length & order as geometry")
+
+    openff_mol = tk.Molecule.from_smiles(smile, allow_undefined_stereo=True)
+
+    # add conformer
+    openff_mol, atom_map = add_conformer(openff_mol, geometry)
+    # assign partial charges
+    openff_mol = assign_partial_charges(
+        openff_mol,
+        atom_map,
+        backup_charge_method,
+        partial_charges,
+    )
+    openff_mol.partial_charges *= charge_scaling
+
+    return openff_mol

tests/files/classical_md_mols/CCO.xyz

@@ -0,0 +1,11 @@
+9
+
+C      1.000000    1.000000    0.000000
+C     -0.515000    1.000000    0.000000
+O     -0.999000    1.000000    1.335000
+H      1.390000    1.001000   -1.022000
+H      1.386000    0.119000    0.523000
+H      1.385000    1.880000    0.526000
+H     -0.907000    0.118000   -0.516000
+H     -0.897000    1.894000   -0.501000
+H     -0.661000    0.198000    1.768000

tests/files/classical_md_mols/FEC-r.xyz

@@ -0,0 +1,12 @@
+10
+
+O      1.000000    1.000000    0.000000
+C     -0.219000    1.000000    0.000000
+O     -0.984000    1.000000    1.133000
+C     -2.322000    0.780000    0.720000
+C     -2.300000    1.205000   -0.711000
+H     -3.034000    0.686000   -1.332000
+F     -2.507000    2.542000   -0.809000
+O     -0.983000    0.948000   -1.128000
+H     -3.008000    1.375000    1.328000
+H     -2.544000   -0.285000    0.838000

tests/files/classical_md_mols/FEC-s.xyz

@@ -0,0 +1,12 @@
+10
+
+O      1.000000    1.000000    0.000000
+C     -0.219000    1.000000    0.000000
+O     -0.981000    1.000000    1.133000
+C     -2.323000    0.828000    0.723000
+C     -2.305000    1.254000   -0.707000
+H     -2.567000    2.305000   -0.862000
+F     -3.125000    0.469000   -1.445000
+O     -0.983000    1.001000   -1.127000
+H     -2.991000    1.447000    1.328000
+H     -2.610000   -0.222000    0.848000

tests/files/classical_md_mols/Li.xyz

@@ -0,0 +1,3 @@
+1
+
+Li    0.0    0.0    0.0

tests/files/classical_md_mols/PF6.xyz

@@ -0,0 +1,9 @@
+7
+
+P    0.0    0.0    0.0
+F    1.6    0.0    0.0
+F   -1.6    0.0    0.0
+F    0.0    1.6    0.0
+F    0.0   -1.6    0.0
+F    0.0    0.0    1.6
+F    0.0    0.0   -1.6