Skip to content

Commit 759c727

Browse files
WIP- Support json serialization
1 parent 802185d commit 759c727

File tree

2 files changed

+17
-5
lines changed

2 files changed

+17
-5
lines changed

gmso/core/dihedral.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -186,3 +186,4 @@ class Config:
186186
alias_to_fields = {
187187
"dihedral_types": "dihedral_types_",
188188
}
189+
json_encoders = {IndexedSet: list}

gmso/formats/json.py

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,9 @@
1313
from gmso.core.bond import Bond
1414
from gmso.core.bond_type import BondType
1515
from gmso.core.box import Box
16-
from gmso.core.dihedral import Dihedral
16+
from gmso.core.dihedral import Dihedral, LayeredDihedral
1717
from gmso.core.dihedral_type import DihedralType
18-
from gmso.core.improper import Improper
18+
from gmso.core.improper import Improper, LayeredImproper
1919
from gmso.core.improper_type import ImproperType
2020
from gmso.core.pairpotential_type import PairPotentialType
2121
from 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

Comments
 (0)