Merge branch 'main' into write_charge_in_elementary

mosdef-hub · Jan 15, 2024 · 2b03f19 · 2b03f19
2 parents 3cd710a + aab43c0
commit 2b03f19
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 2 deletions.
diff --git a/gmso/tests/test_dihedral.py b/gmso/tests/test_dihedral.py
@@ -153,3 +153,35 @@ def test_equivalent_members_set(self):
             tuple(dihedral.connection_members)
             in dihedral_not_eq.equivalent_members()
         )
+
+    def test_sort_dihedral_types(self):
+        from gmso.utils.sorting import sort_by_classes, sort_by_types
+
+        atom1 = Atom(
+            name="atom1", position=[0, 0, 0], atom_type=AtomType(name="A")
+        )
+        atom2 = Atom(
+            name="atom2", position=[1, 0, 0], atom_type=AtomType(name="B")
+        )
+        atom3 = Atom(
+            name="atom3", position=[1, 1, 0], atom_type=AtomType(name="C")
+        )
+        atom4 = Atom(
+            name="atom4", position=[1, 1, 4], atom_type=AtomType(name="D")
+        )
+
+        consituentList = [
+            atom2.atom_type.name,
+            atom4.atom_type.name,
+            atom3.atom_type.name,
+            atom1.atom_type.name,
+        ]
+        dihtype = DihedralType(
+            member_types=consituentList, member_classes=consituentList
+        )
+
+        expected_sortingList = tuple(
+            [atom.atom_type.name for atom in [atom1, atom3, atom4, atom2]]
+        )
+        assert sort_by_classes(dihtype) == expected_sortingList
+        assert sort_by_types(dihtype) == expected_sortingList
diff --git a/gmso/tests/test_improper.py b/gmso/tests/test_improper.py
@@ -152,3 +152,35 @@ def test_equivalent_members_set(self):
             tuple(improper.connection_members)
             in improper_not_eq.equivalent_members()
         )
+
+    def test_sort_improper_types(self):
+        from gmso.utils.sorting import sort_by_classes, sort_by_types
+
+        atom1 = Atom(
+            name="atom1", position=[0, 0, 0], atom_type=AtomType(name="A")
+        )
+        atom2 = Atom(
+            name="atom2", position=[1, 0, 0], atom_type=AtomType(name="B")
+        )
+        atom3 = Atom(
+            name="atom3", position=[1, 1, 0], atom_type=AtomType(name="C")
+        )
+        atom4 = Atom(
+            name="atom4", position=[1, 1, 4], atom_type=AtomType(name="D")
+        )
+
+        consituentList = [
+            atom2.atom_type.name,
+            atom4.atom_type.name,
+            atom3.atom_type.name,
+            atom1.atom_type.name,
+        ]
+        imptype = ImproperType(
+            member_types=consituentList, member_classes=consituentList
+        )
+
+        expected_sortingList = tuple(
+            [atom.atom_type.name for atom in [atom2, atom3, atom4, atom1]]
+        )
+        assert sort_by_classes(imptype) == expected_sortingList
+        assert sort_by_types(imptype) == expected_sortingList
diff --git a/gmso/utils/sorting.py b/gmso/utils/sorting.py
@@ -122,7 +122,8 @@ def sort_by_classes(potential):
     elif isinstance(potential, ImproperType):
         return (
             potential.member_classes[0],
-            *sorted(potential.member_classes[1:]),
+            *sorted(potential.member_classes[1:3]),
+            potential.member_classes[3],
         )
     return ValueError(
         f"Potential {potential} not one of {potential_attribute_map.values()}"
@@ -174,7 +175,8 @@ def sort_by_types(potential):
     elif isinstance(potential, ImproperType):
         return (
             potential.member_types[0],
-            *sorted(potential.member_types[1:]),
+            *sorted(potential.member_types[1:3]),
+            potential.member_types[3],
         )
     return ValueError(
         f"Potential {potential} not one of {potential_attribute_map.values()}"