draft: Add atom type merging for Interchange.to_gromacs()

openff/interchange/_tests/unit_tests/interop/gromacs/export/test_export.py

@@ -232,6 +232,81 @@ def test_simple_roundtrip(self, sage, smiles, reader):
             },
         )
 
+    @pytest.mark.slow
+    @pytest.mark.parametrize("reader", ["intermol", "internal"])
+    @pytest.mark.parametrize(
+        "smiles",
+        [
+            "CC",  # Two identical carbons
+            "C1CCCCC1",  # Identical carbons in the ring
+            "C1[C@@H](N)C[C@@H](N)CC1",  # Identical carbons and nitrogens in the ring
+        ],
+    )
+    def test_energies_with_merging_atom_types(self, sage, smiles, reader):
+        """
+        Tests #962
+        """
+        molecule = MoleculeWithConformer.from_smiles(smiles)
+        molecule.name = molecule.to_hill_formula()
+        topology = molecule.to_topology()
+
+        out = Interchange.from_smirnoff(force_field=sage, topology=topology)
+        out.box = [4, 4, 4]
+        out.positions = molecule.conformers[0]
+
+        get_gromacs_energies(out).compare(
+            get_gromacs_energies(out, merge_atom_types=True),
+            tolerances={
+                "Bond": 0.002 * molecule.n_bonds * unit.kilojoule / unit.mol,
+                "Electrostatics": 0.05 * unit.kilojoule / unit.mol,
+            },
+        )
+
+    @pytest.mark.slow
+    @pytest.mark.skip("from_top is not yet refactored for new Topology API")
+    @pytest.mark.parametrize("reader", ["intermol", "internal"])
+    @pytest.mark.parametrize(
+        "smiles",
+        [
+            "CC",  # Two identical carbons
+            "C1CCCCC1",  # Identical carbons in the ring
+            "C1[C@@H](N)C[C@@H](N)CC1",  # Identical carbons and nitrogens in the ring
+        ],
+    )
+    def test_simple_roundtrip_with_merging_atom_types(self, sage, smiles, reader):
+        """
+        Tests #962
+        """
+        molecule = MoleculeWithConformer.from_smiles(smiles)
+        molecule.name = molecule.to_hill_formula()
+        topology = molecule.to_topology()
+
+        out = Interchange.from_smirnoff(force_field=sage, topology=topology)
+        out.box = [4, 4, 4]
+        out.positions = molecule.conformers[0]
+
+        out.to_top("out.top")
+        out.to_top("out_merged.top", merge_atom_types=True)
+        out.to_gro("out.gro")
+
+        converted = Interchange.from_gromacs("out.top", "out.gro", reader=reader)
+        converted_merged = Interchange.from_gromacs(
+            "out_merged.top",
+            "out.gro",
+            reader=reader,
+        )
+
+        assert numpy.allclose(converted.positions, converted_merged.positions)
+        assert numpy.allclose(converted.box, converted_merged.box)
+
+        get_gromacs_energies(converted_merged).compare(
+            get_gromacs_energies(converted),
+            tolerances={
+                "Bond": 0.002 * molecule.n_bonds * unit.kilojoule / unit.mol,
+                "Electrostatics": 0.05 * unit.kilojoule / unit.mol,
+            },
+        )
+
     @skip_if_missing("parmed")
     def test_num_impropers(self, sage):
         out = Interchange.from_smirnoff(

openff/interchange/components/interchange.py

@@ -421,6 +421,7 @@ def to_gromacs(
         prefix: str,
         decimal: int = 3,
         hydrogen_mass: float = 1.007947,
+        merge_atom_types: bool = False,
     ):
         """
         Export this Interchange object to GROMACS files.
@@ -436,6 +437,9 @@ def to_gromacs(
             The mass to use for hydrogen atoms if not present in the topology. If non-trivially different
             than the default value, mass will be transferred from neighboring heavy atoms. Note that this is currently
             not applied to any waters and is unsupported when virtual sites are present.
+        merge_atom_types: bool, default = False
+            The flag to define behaviour of GROMACSWriter. If True, then similar atom types will be merged.
+            If False, each atom will have its own atom type.
 
         """
         from openff.interchange.interop.gromacs.export._export import GROMACSWriter
@@ -447,7 +451,7 @@ def to_gromacs(
             gro_file=prefix + ".gro",
         )
 
-        writer.to_top()
+        writer.to_top(merge_atom_types=merge_atom_types)
         writer.to_gro(decimal=decimal)
 
     def to_top(

openff/interchange/drivers/gromacs.py

@@ -51,6 +51,7 @@ def get_gromacs_energies(
     mdp: str = "auto",
     round_positions: int = 8,
     detailed: bool = False,
+    merge_atom_types: bool = False,
 ) -> EnergyReport:
     """
     Given an OpenFF Interchange object, return single-point energies as computed by GROMACS.
@@ -67,6 +68,8 @@ def get_gromacs_energies(
         A decimal precision for the positions in the `.gro` file.
     detailed : bool, default=False
         If True, return a detailed report containing the energies of each term.
+    merge_atom_types: bool, default=False
+        If True, energy should be computed with merging atom types.
 
     Returns
     -------
@@ -79,6 +82,7 @@ def get_gromacs_energies(
             interchange=interchange,
             mdp=mdp,
             round_positions=round_positions,
+            merge_atom_types=merge_atom_types,
         ),
         detailed=detailed,
     )
@@ -88,11 +92,16 @@ def _get_gromacs_energies(
     interchange: Interchange,
     mdp: str = "auto",
     round_positions: int = 8,
+    merge_atom_types: bool = False,
 ) -> dict[str, unit.Quantity]:
     with tempfile.TemporaryDirectory() as tmpdir:
         with temporary_cd(tmpdir):
             prefix = "_tmp"
-            interchange.to_gromacs(prefix=prefix, decimal=round_positions)
+            interchange.to_gromacs(
+                prefix=prefix,
+                decimal=round_positions,
+                merge_atom_types=merge_atom_types,
+            )
 
             if mdp == "auto":
                 mdconfig = MDConfig.from_interchange(interchange)

openff/interchange/interop/gromacs/export/_export.py

@@ -25,16 +25,20 @@ class GROMACSWriter(DefaultModel):
     top_file: pathlib.Path | str | None = None
     gro_file: pathlib.Path | str | None = None
 
-    def to_top(self):
+    def to_top(self, merge_atom_types: bool = False):
         """Write a GROMACS topology file."""
         if self.top_file is None:
             raise ValueError("No TOP file specified.")
 
         with open(self.top_file, "w") as top:
             self._write_defaults(top)
-            self._write_atomtypes(top)
+            mapping_to_reduced_atom_types = self._write_atomtypes(top, merge_atom_types)
 
-            self._write_moleculetypes(top)
+            self._write_moleculetypes(
+                top,
+                mapping_to_reduced_atom_types,
+                merge_atom_types,
+            )
 
             self._write_system(top)
             self._write_molecules(top)
@@ -60,31 +64,100 @@ def _write_defaults(self, top):
             f"{self.system.coul_14:8.6f}\n\n",
         )
 
-    def _write_atomtypes(self, top):
+    def _write_atomtypes(self, top, merge_atom_types: bool) -> dict[str, str]:
         top.write("[ atomtypes ]\n")
         top.write(
             ";type, bondingtype, atomic_number, mass, charge, ptype, sigma, epsilon\n",
         )
 
+        reduced_atom_types = []
+        mapping_to_reduced_atom_types = {}
+
+        def _is_atom_type_in_list(
+            atom_type,
+            atom_type_list,
+        ) -> bool | str:
+            """
+            Checks if the atom type is already in list.
+            """
+            from openff.units import Quantity
+
+            mass_tolerance = Quantity("1e-5 amu")
+            sigma_tolerance = Quantity("1e-5 nanometer")
+            epsilon_tolerance = Quantity("1e-5 kilojoule_per_mole")
+            for _at_name, _atom_type in atom_type_list:
+                if (
+                    atom_type.atomic_number == _atom_type.atomic_number
+                    and abs(atom_type.mass - _atom_type.mass) < mass_tolerance
+                    and abs(atom_type.sigma - _atom_type.sigma) < sigma_tolerance
+                    and abs(atom_type.epsilon - _atom_type.epsilon) < epsilon_tolerance
+                ):
+                    return _at_name
+            return False
+
+        def _get_new_entry_name(atom_type_list) -> str:
+            """
+            Entry name for atom type to be added.
+            """
+            if len(reduced_atom_types) > 0:
+                _previous_name = reduced_atom_types[-1][0]
+                _previous_idx = int(_previous_name.split("_")[1])
+                return f"AT_{_previous_idx+1}"
+            else:
+                return "AT_0"
+
         for atom_type in self.system.atom_types.values():
             if not isinstance(atom_type, LennardJonesAtomType):
                 raise NotImplementedError(
                     "Only Lennard-Jones atom types are currently supported.",
                 )
 
+            if merge_atom_types:
+                if _is_atom_type_in_list(atom_type, reduced_atom_types):
+                    mapping_to_reduced_atom_types[atom_type.name] = (
+                        _is_atom_type_in_list(
+                            atom_type,
+                            reduced_atom_types,
+                        )
+                    )
+                else:
+                    _at_name = _get_new_entry_name(reduced_atom_types)
+                    reduced_atom_types.append((_at_name, atom_type))
+                    mapping_to_reduced_atom_types[atom_type.name] = _at_name
+            else:
+                top.write(
+                    f"{atom_type.name :<11s}\t"
+                    f"{atom_type.atomic_number :6d}\t"
+                    f"{atom_type.mass.m :.16g}\t"
+                    f"{atom_type.charge.m :.16f}\t"
+                    f"{atom_type.particle_type :5s}\t"
+                    f"{atom_type.sigma.m :.16g}\t"
+                    f"{atom_type.epsilon.m :.16g}\n",
+                )
+
+        if not merge_atom_types:
+            top.write("\n")
+            return mapping_to_reduced_atom_types
+
+        for atom_type_name, atom_type in reduced_atom_types:
             top.write(
-                f"{atom_type.name :<11s}\t"
+                f"{atom_type_name :<11s}\t"
                 f"{atom_type.atomic_number :6d}\t"
                 f"{atom_type.mass.m :.16g}\t"
                 f"{atom_type.charge.m :.16f}\t"
                 f"{atom_type.particle_type :5s}\t"
                 f"{atom_type.sigma.m :.16g}\t"
                 f"{atom_type.epsilon.m :.16g}\n",
             )
-
         top.write("\n")
-
-    def _write_moleculetypes(self, top):
+        return mapping_to_reduced_atom_types
+
+    def _write_moleculetypes(
+        self,
+        top,
+        mapping_to_reduced_atom_types,
+        merge_atom_types: bool,
+    ):
         for molecule_name, molecule_type in self.system.molecule_types.items():
             top.write("[ moleculetype ]\n")
 
@@ -93,7 +166,12 @@ def _write_moleculetypes(self, top):
                 f"{molecule_type.nrexcl:10d}\n\n",
             )
 
-            self._write_atoms(top, molecule_type)
+            self._write_atoms(
+                top,
+                molecule_type,
+                mapping_to_reduced_atom_types,
+                merge_atom_types,
+            )
             self._write_pairs(top, molecule_type)
             self._write_bonds(top, molecule_type)
             self._write_angles(top, molecule_type)
@@ -104,21 +182,39 @@ def _write_moleculetypes(self, top):
 
         top.write("\n")
 
-    def _write_atoms(self, top, molecule_type):
+    def _write_atoms(
+        self,
+        top,
+        molecule_type,
+        mapping_to_reduced_atom_types,
+        merge_atom_types: bool,
+    ):
         top.write("[ atoms ]\n")
         top.write(";index, atom type, resnum, resname, name, cgnr, charge, mass\n")
 
         for atom in molecule_type.atoms:
-            top.write(
-                f"{atom.index :6d} "
-                f"{atom.atom_type :6s}"
-                f"{atom.residue_index :8d} "
-                f"{atom.residue_name :8s} "
-                f"{atom.name :6s}"
-                f"{atom.charge_group_number :6d}"
-                f"{atom.charge.m :20.12f}"
-                f"{atom.mass.m :20.12f}\n",
-            )
+            if merge_atom_types:
+                top.write(
+                    f"{atom.index :6d} "
+                    f"{mapping_to_reduced_atom_types[atom.atom_type] :6s}"
+                    f"{atom.residue_index :8d} "
+                    f"{atom.residue_name :8s} "
+                    f"{atom.name :6s}"
+                    f"{atom.charge_group_number :6d}"
+                    f"{atom.charge.m :20.12f}"
+                    f"{atom.mass.m :20.12f}\n",
+                )
+            else:
+                top.write(
+                    f"{atom.index :6d} "
+                    f"{atom.atom_type :6s}"
+                    f"{atom.residue_index :8d} "
+                    f"{atom.residue_name :8s} "
+                    f"{atom.name :6s}"
+                    f"{atom.charge_group_number :6d}"
+                    f"{atom.charge.m :20.12f}"
+                    f"{atom.mass.m :20.12f}\n",
+                )
 
         top.write("\n")