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

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

@@ -32,9 +32,9 @@ def to_top(self):
 
         with open(self.top_file, "w") as top:
             self._write_defaults(top)
-            self._write_atomtypes(top)
+            to_reduced_atom_types = self._write_atomtypes(top)
 
-            self._write_moleculetypes(top)
+            self._write_moleculetypes(top, to_reduced_atom_types)
 
             self._write_system(top)
             self._write_molecules(top)
@@ -60,31 +60,71 @@ def _write_defaults(self, top):
             f"{self.system.coul_14:8.6f}\n\n",
         )
 
-    def _write_atomtypes(self, top):
+    def _write_atomtypes(self, top) -> dict:
         top.write("[ atomtypes ]\n")
         top.write(
             ";type, bondingtype, atomic_number, mass, charge, ptype, sigma, epsilon\n",
         )
 
+        reduced_atom_types = []
+        to_reduced_atom_types = {}
+
+        def _is_atom_type_in_list(
+            atom_type, atom_type_list, diff: float = 1e-5
+        ) -> bool | str:
+            """
+            Checks if the atom type is already in list.
+            """
+            for _at_name, _atom_type in atom_type_list:
+                if (
+                    atom_type.atomic_number == _atom_type.atomic_number
+                    and numpy.abs(atom_type.mass.m - _atom_type.mass.m) < diff
+                    and numpy.abs(atom_type.sigma.m - _atom_type.sigma.m) < diff
+                    and numpy.abs(atom_type.epsilon.m - _atom_type.epsilon.m) < diff
+                ):
+                    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 _is_atom_type_in_list(atom_type, reduced_atom_types):
+                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))
+                to_reduced_atom_types[atom_type.name] = _at_name
+
+        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")
+        return to_reduced_atom_types
 
-    def _write_moleculetypes(self, top):
+    def _write_moleculetypes(self, top, to_reduced_atom_types):
         for molecule_name, molecule_type in self.system.molecule_types.items():
             top.write("[ moleculetype ]\n")
 
@@ -93,7 +133,7 @@ def _write_moleculetypes(self, top):
                 f"{molecule_type.nrexcl:10d}\n\n",
             )
 
-            self._write_atoms(top, molecule_type)
+            self._write_atoms(top, molecule_type, to_reduced_atom_types)
             self._write_pairs(top, molecule_type)
             self._write_bonds(top, molecule_type)
             self._write_angles(top, molecule_type)
@@ -104,14 +144,14 @@ def _write_moleculetypes(self, top):
 
         top.write("\n")
 
-    def _write_atoms(self, top, molecule_type):
+    def _write_atoms(self, top, molecule_type, to_reduced_atom_types):
         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"{to_reduced_atom_types[atom.atom_type] :6s}"
                 f"{atom.residue_index :8d} "
                 f"{atom.residue_name :8s} "
                 f"{atom.name :6s}"