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/abc/abstract_connection.py

@@ -33,7 +33,7 @@ def name(self):
 
     @root_validator(pre=True)
     def validate_fields(cls, values):
-        connection_members = values.get("connection_members")
+        connection_members = values.get("connection_members", [])
 
         if all(isinstance(member, dict) for member in connection_members):
             connection_members = [

gmso/abc/gmso_base.py

@@ -120,3 +120,4 @@ class Config:
         extra = "forbid"
         json_encoders = GMSOJSONHandler.json_encoders
         allow_population_by_field_name = True
+        validate_assignment = True

gmso/core/dihedral.py

@@ -1,17 +1,21 @@
-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.
-    This class has strictly 4 members in its connection_members.
+    This is a subclass of the gmso.Connection superclass. This class
+    has strictly 4 members in its connection_members and used as
+    a base class to define many different forms of a Dihedral.
+
     The connection_type in this class corresponds to gmso.DihedralType.
     The connectivity of a dihedral is:
         m1–m2–m3–m4
@@ -21,7 +25,7 @@ class Dihedral(Connection):
     Notes
     -----
     Inherits some methods from Connection:
-        __eq__, __repr__, _validate methods
+        __eq__, _validate methods
 
     Additional _validate methods are presented
     """
@@ -31,19 +35,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 +87,102 @@ def _equivalent_members_hash(self):
             )
         )
 
+    def is_layered(self):
+        return hasattr(self, "dihedral_types_")
+
+    def __repr__(self):
+        return (
+            f"<{self.__class__.__name__} {self.name},\n "
+            f"connection_members: {self.connection_members},\n "
+            f"potential: {str(self.dihedral_types if self.is_layered() else self.dihedral_type)},\n "
+            f"id: {id(self)}>"
+        )
+
+    class Config:
+        fields = {
+            "connection_members_": "connection_members",
+        }
+        alias_to_fields = {
+            "connection_members": "connection_members_",
+        }
+
+
+class Dihedral(BaseDihedral):
+    __base_doc__ = """A 4-Partner connection between 2 sites with a single dihedral type association
+
+    Notes
+    -----
+    This class inherits from 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):
+    __base_doc__ = """A 4-Partner connection between 2 sites with a multiple dihedral type associations
+
+    Notes
+    -----
+    This class inherits from 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)
+    def validate_dihedral_types(cls, dihedral_types):
+        if not isinstance(dihedral_types, Iterable) or isinstance(
+            dihedral_types, str
+        ):
+            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_",
         }