Restraint support for GROMACS top writer (#685)

* add ability to write angle and dihedral restraint for gromacs top format

* fix typo in dihedral, add docs for angle restraints

* add #ifdef DIHRES for dihedral restrain section

* change var name, better handling of scaling factors for top writer

* fix tab for top writer

* fix typo

* better parsing unique molecules

* reformat top writer to work properly with new changes

* add test files for top and gro writer, make gro write out res info

* adjust unit tests

* add precision and adjust unit test

* adjustments to element parsing, spacing of writing out gro file, speed up writer compatibility check

* fix minor bugs

* reformat top writer, add simplify_check for compatability check

* fix unit test, add sanity check for top writer

* truncating site name in gro writer

* better handling use_molecule_info speedup during atomtyping

* Add test with restraints

* fix bug related to restraints section

* update test, add more docs about angle and dihedral restraints, make top writer refer to element from atomtype (work better for non-atomistic

* add more rigorous test for top and gro file

* add write gro test

* combine benzene ua and aa test

* add harmonic bond restraint for bond class and top writer

* add unit test for bond restraints (harmonic)

* update test file

* fix typo and removed unused imports

* fixing unit tests

gmso/core/angle.py

@@ -31,6 +31,16 @@ class Angle(Connection):
         default=None, description="AngleType of this angle."
     )
 
+    restraint_: Optional[dict] = Field(
+        default=None,
+        description="""
+        Restraint for this angle, must be a dict with the following keys:
+        'k' (unit of energy/mol), 'theta_eq' (unit of angle), 'n' (multiplicity, unitless).
+        Refer to https://manual.gromacs.org/current/reference-manual/topologies/topology-file-formats.html
+        for more information.
+        """,
+    )
+
     @property
     def angle_type(self):
         """Return the angle type if the angle is parametrized."""
@@ -41,6 +51,11 @@ def connection_type(self):
         """Return the angle type if the angle is parametrized."""
         return self.__dict__.get("angle_type_")
 
+    @property
+    def restraint(self):
+        """Return the restraint of this angle."""
+        return self.__dict__.get("restraint_")
+
     def equivalent_members(self):
         """Return a set of the equivalent connection member tuples.
 
@@ -74,8 +89,10 @@ class Config:
         fields = {
             "connection_members_": "connection_members",
             "angle_type_": "angle_type",
+            "restraint_": "restraint",
         }
         alias_to_fields = {
             "connection_members": "connection_members_",
             "angle_type": "angle_type_",
+            "restraint": "restraint_",
         }

gmso/core/atom_type.py

@@ -132,6 +132,7 @@ def clone(self, fast_copy=False):
         """Clone this AtomType, faster alternative to deepcopying."""
         return AtomType(
             name=str(self.name),
+            tags=self.tags,
             expression=None,
             parameters=None,
             independent_variables=None,

gmso/core/bond.py

@@ -30,6 +30,15 @@ class Bond(Connection):
     bond_type_: Optional[BondType] = Field(
         default=None, description="BondType of this bond."
     )
+    restraint_: Optional[dict] = Field(
+        default=None,
+        description="""
+        Restraint for this bond, must be a dict with the following keys:
+        'b0' (unit of length), 'kb' (unit of energy/(mol * length**2)).
+        Refer to https://manual.gromacs.org/current/reference-manual/topologies/topology-file-formats.html
+        for more information.
+        """,
+    )
 
     @property
     def bond_type(self):
@@ -42,6 +51,11 @@ def connection_type(self):
         # ToDo: Deprecate this?
         return self.__dict__.get("bond_type_")
 
+    @property
+    def restraint(self):
+        """Return the restraint of this bond."""
+        return self.__dict__.get("restraint_")
+
     def equivalent_members(self):
         """Get a set of the equivalent connection member tuples.
 
@@ -75,8 +89,10 @@ class Config:
         fields = {
             "bond_type_": "bond_type",
             "connection_members_": "connection_members",
+            "restraint_": "restraint",
         }
         alias_to_fields = {
             "bond_type": "bond_type_",
             "connection_members": "connection_members_",
+            "restraint": "restraint_",
         }

gmso/core/dihedral.py

@@ -35,6 +35,16 @@ class Dihedral(Connection):
         default=None, description="DihedralType of this dihedral."
     )
 
+    restraint_: Optional[dict] = Field(
+        default=None,
+        description="""
+        Restraint for this dihedral, must be a dict with the following keys:
+        'k' (unit of energy/(mol * angle**2)), 'phi_eq' (unit of angle), 'delta_phi' (unit of angle).
+        Refer to https://manual.gromacs.org/current/reference-manual/topologies/topology-file-formats.html
+        for more information.
+        """,
+    )
+
     @property
     def dihedral_type(self):
         return self.__dict__.get("dihedral_type_")
@@ -44,6 +54,11 @@ def connection_type(self):
         # ToDo: Deprecate this?
         return self.__dict__.get("dihedral_type_")
 
+    @property
+    def restraint(self):
+        """Return the restraint of this dihedral."""
+        return self.__dict__.get("restraint_")
+
     def equivalent_members(self):
         """Get a set of the equivalent connection member tuples
 
@@ -74,8 +89,10 @@ class Config:
         fields = {
             "dihedral_type_": "dihedral_type",
             "connection_members_": "connection_members",
+            "restraint_": "restraint",
         }
         alias_to_fields = {
             "dihedral_type": "dihedral_type_",
             "connection_members": "connection_members_",
+            "restraint": "restraint_",
         }

gmso/external/convert_foyer_xml.py

@@ -529,16 +529,6 @@ def _create_sub_element(root_el, name, attrib_dict=None):
 
 
 def _validate_foyer(xml_path):
-    import warnings
+    from foyer.validator import Validator
 
-    from gmso.utils.io import has_foyer
-
-    if not has_foyer:
-        warnings.warn(
-            "Cannot validate the xml using foyer, since foyer is not installed."
-            "Please install foyer using conda install -c conda-forge foyer."
-        )
-    else:
-        from foyer.validator import Validator
-
-        Validator(xml_path)
+    Validator(xml_path)

gmso/external/convert_mbuild.py

@@ -250,11 +250,7 @@ def _parse_particle(particle_map, site):
 def _parse_site(site_map, particle, search_method):
     """Parse information for a gmso.Site from a mBuild.Compound adn add it to the site map."""
     pos = particle.xyz[0] * u.nm
-    ele = (
-        search_method(particle.element.symbol)
-        if particle.element
-        else search_method(particle.name)
-    )
+    ele = search_method(particle.element.symbol) if particle.element else None
     charge = particle.charge * u.elementary_charge if particle.charge else None
     mass = particle.mass * u.amu if particle.mass else None
 

gmso/external/convert_parmed.py

@@ -256,9 +256,14 @@ def _atom_types_from_pmd(structure):
     unique_atom_types = list(unique_atom_types)
     pmd_top_atomtypes = {}
     for atom_type in unique_atom_types:
+        if atom_type.atomic_number:
+            element = element_by_atomic_number(atom_type.atomic_number).symbol
+        else:
+            element = atom_type.name
         top_atomtype = gmso.AtomType(
             name=atom_type.name,
             charge=atom_type.charge * u.elementary_charge,
+            tags={"element": element},
             expression="4*epsilon*((sigma/r)**12 - (sigma/r)**6)",
             parameters={
                 "sigma": atom_type.sigma * u.angstrom,

gmso/formats/gro.py

@@ -1,5 +1,6 @@
 """Read and write Gromos87 (.GRO) file format."""
 import datetime
+import re
 import warnings
 
 import numpy as np
@@ -10,7 +11,6 @@
 from gmso.core.atom import Atom
 from gmso.core.box import Box
 from gmso.core.topology import Topology
-from gmso.exceptions import NotYetImplementedWarning
 from gmso.formats.formats_registry import loads_as, saves_as
 
 
@@ -57,25 +57,31 @@ def read_gro(filename):
         coords = u.nm * np.zeros(shape=(n_atoms, 3))
         for row, _ in enumerate(coords):
             line = gro_file.readline()
+            content = line.split()
             if not line:
                 msg = (
                     "Incorrect number of lines in .gro file. Based on the "
                     "number in the second line of the file, {} rows of"
                     "atoms were expected, but at least one fewer was found."
                 )
                 raise ValueError(msg.format(n_atoms))
-            resid = int(line[:5])
-            res_name = line[5:10]
-            atom_name = line[10:15]
-            atom_id = int(line[15:20])
+
+            res = content[0]
+            atom_name = content[1]
+            atom_id = content[2]
             coords[row] = u.nm * np.array(
                 [
-                    float(line[20:28]),
-                    float(line[28:36]),
-                    float(line[36:44]),
+                    float(content[3]),
+                    float(content[4]),
+                    float(content[5]),
                 ]
             )
             site = Atom(name=atom_name, position=coords[row])
+
+            r = re.compile("([0-9]+)([a-zA-Z]+)")
+            m = r.match(res)
+            site.molecule = (m.group(2), int(m.group(1)) - 1)
+            site.residue = (m.group(2), int(m.group(1)) - 1)
             top.add_site(site, update_types=False)
         top.update_topology()
 
@@ -97,7 +103,7 @@ def read_gro(filename):
 
 
 @saves_as(".gro")
-def write_gro(top, filename):
+def write_gro(top, filename, precision=3):
     """Write a topology to a gro file.
 
     The Gromos87 (gro) format is a common plain text structure file used
@@ -113,7 +119,8 @@ def write_gro(top, filename):
         The `topology` to write out to the gro file.
     filename : str or file object
         The location and name of file to save to disk.
-
+    precision : int, optional, default=3
+        The number of sig fig to write out the position in.
 
     Notes
     -----
@@ -135,7 +142,7 @@ def write_gro(top, filename):
             )
         )
         out_file.write("{:d}\n".format(top.n_sites))
-        out_file.write(_prepare_atoms(top, pos_array))
+        out_file.write(_prepare_atoms(top, pos_array, precision))
         out_file.write(_prepare_box(top))
 
 
@@ -154,30 +161,44 @@ def _validate_positions(pos_array):
     return pos_array
 
 
-def _prepare_atoms(top, updated_positions):
+def _prepare_atoms(top, updated_positions, precision):
     out_str = str()
+    warnings.warn(
+        "Residue information is parsed from site.molecule,"
+        "or site.residue if site.molecule does not exist."
+        "Note that the residue idx will be bump by 1 since GROMACS utilize 1-index."
+    )
     for idx, (site, pos) in enumerate(zip(top.sites, updated_positions)):
-        warnings.warn(
-            "Residue information is not currently "
-            "stored or written to GRO files.",
-            NotYetImplementedWarning,
-        )
-        # TODO: assign residues
-        res_id = 1
-        res_name = "X"
-        atom_name = site.name
+        if site.molecule:
+            res_id = site.molecule.number + 1
+            res_name = site.molecule.name
+        elif site.residue:
+            res_id = site.residue.number + 1
+            res_name = site.molecule.name[:3]
+        else:
+            res_id = 1
+            res_name = "MOL"
+        if len(res_name) > 3:
+            res_name = res_name[:3]
+
+        atom_name = site.name if len(site.name) <= 3 else site.name[:3]
         atom_id = idx + 1
-        out_str = (
-            out_str
-            + "{0:5d}{1:5s}{2:5s}{3:5d}{4:8.3f}{5:8.3f}{6:8.3f}\n".format(
-                res_id,
-                res_name,
-                atom_name,
-                atom_id,
-                pos[0].in_units(u.nm).value,
-                pos[1].in_units(u.nm).value,
-                pos[2].in_units(u.nm).value,
-            )
+
+        varwidth = 5 + precision
+        crdfmt = f"{{:{varwidth}.{precision}f}}"
+
+        # preformat pos str
+        crt_x = crdfmt.format(pos[0].in_units(u.nm).value)[:varwidth]
+        crt_y = crdfmt.format(pos[1].in_units(u.nm).value)[:varwidth]
+        crt_z = crdfmt.format(pos[2].in_units(u.nm).value)[:varwidth]
+        out_str = out_str + "{0:5d}{1:5s}{2:5s}{3:5d}{4}{5}{6}\n".format(
+            res_id,
+            res_name,
+            atom_name,
+            atom_id,
+            crt_x,
+            crt_y,
+            crt_z,
         )
     return out_str
 
@@ -190,7 +211,7 @@ def _prepare_box(top):
         rtol=1e-5,
         atol=0.1 * u.degree,
     ):
-        out_str = out_str + " {:0.5f} {:0.5f} {:0.5f} \n".format(
+        out_str = out_str + " {:0.5f} {:0.5f} {:0.5f}\n".format(
             top.box.lengths[0].in_units(u.nm).value.round(6),
             top.box.lengths[1].in_units(u.nm).value.round(6),
             top.box.lengths[2].in_units(u.nm).value.round(6),