Add LayeredDihedral and LayeredImproper support

gmso/__init__.py

@@ -6,11 +6,11 @@
 from .core.bond import Bond
 from .core.bond_type import BondType
 from .core.box import Box
-from .core.dihedral import Dihedral
+from .core.dihedral import Dihedral, LayeredDihedral
 from .core.dihedral_type import DihedralType
 from .core.element import Element
 from .core.forcefield import ForceField
-from .core.improper import Improper
+from .core.improper import Improper, LayeredImproper
 from .core.improper_type import ImproperType
 from .core.pairpotential_type import PairPotentialType
 from .core.subtopology import SubTopology

gmso/core/dihedral.py

@@ -1,13 +1,15 @@
-from typing import Callable, ClassVar, Optional, Tuple
+from typing import Callable, ClassVar, Iterable, Optional, Tuple
 
-from pydantic import Field
+from boltons.setutils import IndexedSet
+from pydantic import Field, ValidationError, validator
 
 from gmso.abc.abstract_connection import Connection
 from gmso.core.atom import Atom
 from gmso.core.dihedral_type import DihedralType
+from gmso.utils.misc import validate_type
 
 
-class Dihedral(Connection):
+class BaseDihedral(Connection):
     __base_doc__ = """A 4-partner connection between sites.
 
     This is a subclass of the gmso.Connection superclass.
@@ -31,19 +33,6 @@ class Dihedral(Connection):
         ..., description="The 4 atoms involved in the dihedral."
     )
 
-    dihedral_type_: Optional[DihedralType] = Field(
-        default=None, description="DihedralType of this dihedral."
-    )
-
-    @property
-    def dihedral_type(self):
-        return self.__dict__.get("dihedral_type_")
-
-    @property
-    def connection_type(self):
-        # ToDo: Deprecate this?
-        return self.__dict__.get("dihedral_type_")
-
     def equivalent_members(self):
         """Get a set of the equivalent connection member tuples
 
@@ -96,18 +85,79 @@ def _equivalent_members_hash(self):
             )
         )
 
+    def is_layered(self):
+        return hasattr(self, "dihedral_types_")
+
+    class Config:
+        fields = {
+            "connection_members_": "connection_members",
+        }
+        alias_to_fields = {
+            "connection_members": "connection_members_",
+        }
+
+
+class Dihedral(BaseDihedral):
+    dihedral_type_: Optional[DihedralType] = Field(
+        default=None, description="DihedralType of this dihedral."
+    )
+
+    @property
+    def dihedral_type(self):
+        return self.__dict__.get("dihedral_type_")
+
+    @property
+    def connection_type(self):
+        # ToDo: Deprecate this?
+        return self.__dict__.get("dihedral_type_")
+
     def __setattr__(self, key, value):
         if key == "connection_type":
-            super(Dihedral, self).__setattr__("dihedral_type", value)
+            super().__setattr__("dihedral_type", value)
         else:
-            super(Dihedral, self).__setattr__(key, value)
+            super().__setattr__(key, value)
 
     class Config:
         fields = {
             "dihedral_type_": "dihedral_type",
-            "connection_members_": "connection_members",
         }
         alias_to_fields = {
             "dihedral_type": "dihedral_type_",
-            "connection_members": "connection_members_",
+        }
+
+
+class LayeredDihedral(BaseDihedral):
+    dihedral_types_: Optional[IndexedSet] = Field(
+        default=None, description="DihedralTypes of this dihedral."
+    )
+
+    @property
+    def dihedral_types(self):
+        return self.__dict__.get("dihedral_types_")
+
+    @property
+    def connection_types(self):
+        # ToDo: Deprecate this?
+        return self.__dict__.get("dihedral_types_")
+
+    def __setattr__(self, key, value):
+        if key == "connection_types":
+            super().__setattr__("dihedral_types", value)
+        else:
+            super().__setattr__(key, value)
+
+    @validator("dihedral_types_", pre=True, always=True)
+    def validate_dihedral_types(cls, dihedral_types):
+        if not isinstance(dihedral_types, Iterable):
+            raise ValidationError("DihedralTypes should be iterable", cls)
+
+        validate_type(dihedral_types, DihedralType)
+        return IndexedSet(dihedral_types)
+
+    class Config:
+        fields = {
+            "dihedral_types_": "dihedral_types",
+        }
+        alias_to_fields = {
+            "dihedral_types": "dihedral_types_",
         }

gmso/core/improper.py

@@ -1,14 +1,16 @@
 """Support for improper style connections (4-member connection)."""
-from typing import Callable, ClassVar, Optional, Tuple
+from typing import Callable, ClassVar, Iterable, Optional, Tuple
 
-from pydantic import Field
+from boltons.setutils import IndexedSet
+from pydantic import Field, ValidationError, validator
 
 from gmso.abc.abstract_connection import Connection
 from gmso.core.atom import Atom
 from gmso.core.improper_type import ImproperType
+from gmso.utils.misc import validate_type
 
 
-class Improper(Connection):
+class BaseImproper(Connection):
     __base_doc__ = """sA 4-partner connection between sites.
 
     This is a subclass of the gmso.Connection superclass.
@@ -39,21 +41,6 @@ class Improper(Connection):
         "then the three atoms connected to the central site.",
     )
 
-    improper_type_: Optional[ImproperType] = Field(
-        default=None, description="ImproperType of this improper."
-    )
-
-    @property
-    def improper_type(self):
-        """Return Potential object for this connection if it exists."""
-        return self.__dict__.get("improper_type_")
-
-    @property
-    def connection_type(self):
-        """Return Potential object for this connection if it exists."""
-        # ToDo: Deprecate this?
-        return self.__dict__.get("improper_type_")
-
     def equivalent_members(self):
         """Get a set of the equivalent connection member tuples.
 
@@ -106,21 +93,83 @@ def _equivalent_members_hash(self):
             )
         )
 
+    class Config:
+        """Pydantic configuration to link fields to their public attribute."""
+
+        fields = {
+            "connection_members_": "connection_members",
+        }
+        alias_to_fields = {
+            "connection_members": "connection_members_",
+        }
+
+
+class Improper(BaseImproper):
+    improper_type_: Optional[ImproperType] = Field(
+        default=None, description="ImproperType of this improper."
+    )
+
+    @property
+    def improper_type(self):
+        """Return Potential object for this connection if it exists."""
+        return self.__dict__.get("improper_type_")
+
+    @property
+    def connection_type(self):
+        """Return Potential object for this connection if it exists."""
+        # ToDo: Deprecate this?
+        return self.__dict__.get("improper_type_")
+
     def __setattr__(self, key, value):
         """Set attribute override to support connection_type key."""
         if key == "connection_type":
-            super(Improper, self).__setattr__("improper_type", value)
+            super().__setattr__("improper_type", value)
         else:
-            super(Improper, self).__setattr__(key, value)
+            super().__setattr__(key, value)
 
     class Config:
         """Pydantic configuration to link fields to their public attribute."""
 
         fields = {
             "improper_type_": "improper_type",
-            "connection_members_": "connection_members",
         }
         alias_to_fields = {
             "improper_type": "improper_type_",
-            "connection_members": "connection_members_",
+        }
+
+
+class LayeredImproper(BaseImproper):
+    improper_types_: Optional[IndexedSet] = Field(
+        default=None, description="ImproperTypes of this improper."
+    )
+
+    @property
+    def improper_types(self):
+        return self.__dict__.get("improper_types_")
+
+    @property
+    def connection_types(self):
+        # ToDo: Deprecate this?
+        return self.__dict__.get("improper_types_")
+
+    def __setattr__(self, key, value):
+        if key == "connection_types":
+            super().__setattr__("improper_types_", value)
+        else:
+            super().__setattr__(key, value)
+
+    @validator("improper_types_", pre=True, always=True)
+    def validate_dihedral_types(cls, improper_types):
+        if not isinstance(improper_types, Iterable):
+            raise ValidationError("ImproperTypes should be iterable", cls)
+
+        validate_type(improper_types, ImproperType)
+        return IndexedSet(improper_types)
+
+    class Config:
+        fields = {
+            "improper_types_": "improper_types",
+        }
+        alias_to_fields = {
+            "improper_types": "improper_types_",
         }