add isrigid to molecule

mosdef-hub · Oct 16, 2023 · 904750f · 904750f
1 parent 982a04d
commit 904750f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/gmso/abc/abstract_site.py b/gmso/abc/abstract_site.py
@@ -10,12 +10,14 @@
 from gmso.exceptions import GMSOError
 
 try:
-    from pydantic.v1 import Field, StrictInt, StrictStr, validator
+    from pydantic.v1 import Field, StrictBool, StrictInt, StrictStr, validator
 except ImportError:
-    from pydantic import Field, StrictInt, StrictStr, validator
+    from pydantic import Field, StrictBool, StrictInt, StrictStr, validator
 
 PositionType = Union[Sequence[float], np.ndarray, u.unyt_array]
-MoleculeType = NamedTuple("Molecule", name=StrictStr, number=StrictInt)
+MoleculeType = NamedTuple(
+    "Molecule", name=StrictStr, number=StrictInt, isrigid=StrictBool
+)
 ResidueType = NamedTuple("Residue", name=StrictStr, number=StrictInt)
 
 SiteT = TypeVar("SiteT", bound="Site")

diff --git a/gmso/formats/top.py b/gmso/formats/top.py
@@ -174,7 +174,7 @@ def write_top(top, filename, top_vars=None):
             # Special treatment for water, may ned to consider a better way to tag rigid water
             # Built using this https://github.com/gromacs/gromacs/blob/main/share/top/oplsaa.ff/spce.itp as reference
             if "water" in tag.lower() and all(
-                "rigid" in site.label for site in unique_molecules[tag]["sites"]
+                site.molecule.isrigid for site in unique_molecules[tag]["sites"]
             ):
                 sites_list = unique_molecules[tag]["sites"]
                 water_sites = {