Ensure MSONAtoms is indeed MSONable when Atoms.info is loaded w…

…ith goodies (#3670) * Ensure `MSONAtoms` is indeed `MSONable` * pre-commit auto-fixes * Update ase.py Signed-off-by: Andrew S. Rosen <[email protected]> * Update test_ase.py Signed-off-by: Andrew S. Rosen <[email protected]> * pre-commit auto-fixes * Update ase.py Signed-off-by: Andrew S. Rosen <[email protected]> * Update test_ase.py Signed-off-by: Andrew S. Rosen <[email protected]> * fix * pre-commit auto-fixes * Update ase.py Signed-off-by: Andrew S. Rosen <[email protected]> --------- Signed-off-by: Andrew S. Rosen <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
materialsproject · Mar 4, 2024 · b96d9e1 · b96d9e1
1 parent 442d740
commit b96d9e1
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 8 deletions.
diff --git a/pymatgen/io/ase.py b/pymatgen/io/ase.py
@@ -11,7 +11,7 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
-from monty.json import MSONable
+from monty.json import MontyDecoder, MSONable, jsanitize
 
 from pymatgen.core.structure import Molecule, Structure
 
@@ -49,19 +49,29 @@ def __init__(self, *args, **kwargs):
 class MSONAtoms(Atoms, MSONable):
     """A custom subclass of ASE Atoms that is MSONable, including `.as_dict()` and `.from_dict()` methods."""
 
-    def as_dict(s: Atoms) -> dict[str, Any]:
+    def as_dict(atoms: Atoms) -> dict[str, Any]:
         # Normally, we would want to this to be a wrapper around atoms.todict() with @module and
         # @class key-value pairs inserted. However, atoms.todict()/atoms.fromdict() is not meant
         # to be used in a round-trip fashion and does not work properly with constraints.
         # See ASE issue #1387.
-        return {"@module": "pymatgen.io.ase", "@class": "MSONAtoms", "atoms_json": encode(s)}
+        atoms_no_info = atoms.copy()
+        atoms_no_info.info = {}
+        return {
+            "@module": "pymatgen.io.ase",
+            "@class": "MSONAtoms",
+            "atoms_json": encode(atoms_no_info),
+            "atoms_info": jsanitize(atoms.info, strict=True),
+        }
 
     def from_dict(dct: dict[str, Any]) -> MSONAtoms:
         # Normally, we would want to this to be a wrapper around atoms.fromdict() with @module and
         # @class key-value pairs inserted. However, atoms.todict()/atoms.fromdict() is not meant
         # to be used in a round-trip fashion and does not work properly with constraints.
         # See ASE issue #1387.
-        return MSONAtoms(decode(dct["atoms_json"]))
+        mson_atoms = MSONAtoms(decode(dct["atoms_json"]))
+        atoms_info = MontyDecoder().process_decoded(dct["atoms_info"])
+        mson_atoms.info = atoms_info
+        return mson_atoms
 
 
 # NOTE: If making notable changes to this class, please ping @Andrew-S-Rosen on GitHub.

diff --git a/tests/io/test_ase.py b/tests/io/test_ase.py
@@ -316,15 +316,30 @@ def test_back_forth_v4():
 
 @skip_if_no_ase
 def test_msonable_atoms():
+    structure = Structure.from_file(f"{TEST_FILES_DIR}/POSCAR")
+
     atoms = ase.io.read(f"{TEST_FILES_DIR}/OUTCAR")
+    atoms_info = {"test": "hi", "structure": structure}
+    atoms.info = atoms_info
     assert not isinstance(atoms, MSONAtoms)
-    ref = {"@module": "pymatgen.io.ase", "@class": "MSONAtoms", "atoms_json": ase.io.jsonio.encode(atoms)}
+
+    ref_atoms = atoms.copy()
+    ref_atoms.info = {}
+
     msonable_atoms = MSONAtoms(atoms)
     assert atoms == msonable_atoms
-    assert msonable_atoms.as_dict() == ref
-    assert MSONAtoms.from_dict(ref) == atoms
 
-    structure = Structure.from_file(f"{TEST_FILES_DIR}/POSCAR")
+    msonable_atoms_dict = msonable_atoms.as_dict()
+    assert msonable_atoms_dict == {
+        "@module": "pymatgen.io.ase",
+        "@class": "MSONAtoms",
+        "atoms_json": ase.io.jsonio.encode(ref_atoms),
+        "atoms_info": jsanitize(atoms_info, strict=True),
+    }
+
+    atoms_back = MSONAtoms.from_dict(msonable_atoms_dict)
+    assert atoms_back == atoms
+    assert atoms_back.info == atoms.info
 
     atoms = AseAtomsAdaptor.get_atoms(structure, msonable=True)
     assert callable(atoms.as_dict)