1313from gmso .core .bond import Bond
1414from gmso .core .bond_type import BondType
1515from gmso .core .box import Box
16- from gmso .core .dihedral import Dihedral
16+ from gmso .core .dihedral import Dihedral , LayeredDihedral
1717from gmso .core .dihedral_type import DihedralType
18- from gmso .core .improper import Improper
18+ from gmso .core .improper import Improper , LayeredImproper
1919from gmso .core .improper_type import ImproperType
2020from gmso .core .pairpotential_type import PairPotentialType
2121from gmso .formats .formats_registry import loads_as , saves_as
@@ -90,7 +90,9 @@ def _to_json(top, types=False, update=True):
9090 Bond : json_dict ["bonds" ],
9191 Angle : json_dict ["angles" ],
9292 Dihedral : json_dict ["dihedrals" ],
93+ LayeredDihedral : json_dict ["impropers" ],
9394 Improper : json_dict ["impropers" ],
95+ LayeredImproper : json_dict ["impropers" ],
9496 AtomType : json_dict ["atom_types" ],
9597 BondType : json_dict ["bond_types" ],
9698 AngleType : json_dict ["angle_types" ],
@@ -101,20 +103,29 @@ def _to_json(top, types=False, update=True):
101103 for connections , exclude_attr in [
102104 (top ._bonds , "bond_type" ),
103105 (top ._angles , "angle_type" ),
104- (top ._dihedrals , "dihedral_type" ),
105- (top ._impropers , "improper_type" ),
106+ (top ._dihedrals , { "dihedral_type" , "dihedral_types" } ),
107+ (top ._impropers , { "improper_type" , "improper_types" } ),
106108 ]:
107109 for connection in connections :
108110 connection_dict = connection .json_dict (
109111 exclude = {exclude_attr , "connection_members" }
112+ if not isinstance (exclude_attr , set )
113+ else exclude_attr .union ("connection_members" )
110114 )
111115 target = targets [type (connection )]
112116 connection_dict ["connection_members" ] = [
113117 top .get_index (member )
114118 for member in connection .connection_members
115119 ]
116120 target .append (connection_dict )
117- connection_type = getattr (connection , exclude_attr )
121+ if isinstance (exclude_attr , set ):
122+ for attr in exclude_attr :
123+ try :
124+ connection_type = getattr (connection , attr )
125+ except AttributeError :
126+ pass
127+ else :
128+ connection_type = getattr (connection , exclude_attr )
118129 if types and connection_type :
119130 connection_dict [exclude_attr ] = id (connection_type )
120131 if types :
0 commit comments