From ff1c19c30b1941155b66be8d1801a22acb146f47 Mon Sep 17 00:00:00 2001 From: Andrew Rosen Date: Mon, 4 Mar 2024 10:29:05 -0800 Subject: [PATCH 01/10] Ensure `MSONAtoms` is indeed `MSONable` --- pymatgen/io/ase.py | 15 +++++++++------ tests/io/test_ase.py | 14 ++++++++++++-- 2 files changed, 21 insertions(+), 8 deletions(-) diff --git a/pymatgen/io/ase.py b/pymatgen/io/ase.py index 2d947353606..b5d7d547477 100644 --- a/pymatgen/io/ase.py +++ b/pymatgen/io/ase.py @@ -4,14 +4,13 @@ """ from __future__ import annotations - import warnings from collections.abc import Iterable from importlib.metadata import PackageNotFoundError from typing import TYPE_CHECKING import numpy as np -from monty.json import MSONable +from monty.json import MSONable, jsanitize, MontyDecoder from pymatgen.core.structure import Molecule, Structure @@ -49,20 +48,24 @@ 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_info = atoms.info.copy() + atoms.info = {} + return {"@module": "pymatgen.io.ase", "@class": "MSONAtoms", "atoms_json": encode(atoms), "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. # There are some subtleties in here, particularly related to spins/charges. diff --git a/tests/io/test_ase.py b/tests/io/test_ase.py index 7faf309b3b3..11fba7d13e8 100644 --- a/tests/io/test_ase.py +++ b/tests/io/test_ase.py @@ -317,12 +317,22 @@ def test_back_forth_v4(): @skip_if_no_ase def test_msonable_atoms(): atoms = ase.io.read(f"{TEST_FILES_DIR}/OUTCAR") + atoms_info = {"test": "hi"} + 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 + + ref = {"@module": "pymatgen.io.ase", "@class": "MSONAtoms", "atoms_json": ase.io.jsonio.encode(ref_atoms), "atoms_info": atoms_info} assert msonable_atoms.as_dict() == ref - assert MSONAtoms.from_dict(ref) == atoms + + atoms_back = MSONAtoms.from_dict(ref) + assert atoms_back == atoms + assert atoms_back.info == atoms_info structure = Structure.from_file(f"{TEST_FILES_DIR}/POSCAR") From 912c10160298eb8bdff3e0d328b1c728fdd2bd12 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 4 Mar 2024 18:31:34 +0000 Subject: [PATCH 02/10] pre-commit auto-fixes --- pymatgen/io/ase.py | 11 +++++++++-- tests/io/test_ase.py | 9 +++++++-- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/pymatgen/io/ase.py b/pymatgen/io/ase.py index b5d7d547477..a34f1d4e639 100644 --- a/pymatgen/io/ase.py +++ b/pymatgen/io/ase.py @@ -4,13 +4,14 @@ """ from __future__ import annotations + import warnings from collections.abc import Iterable from importlib.metadata import PackageNotFoundError from typing import TYPE_CHECKING import numpy as np -from monty.json import MSONable, jsanitize, MontyDecoder +from monty.json import MontyDecoder, MSONable, jsanitize from pymatgen.core.structure import Molecule, Structure @@ -55,7 +56,12 @@ def as_dict(atoms: Atoms) -> dict[str, Any]: # See ASE issue #1387. atoms_info = atoms.info.copy() atoms.info = {} - return {"@module": "pymatgen.io.ase", "@class": "MSONAtoms", "atoms_json": encode(atoms), "atoms_info": jsanitize(atoms_info, strict=True)} + return { + "@module": "pymatgen.io.ase", + "@class": "MSONAtoms", + "atoms_json": encode(atoms), + "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 @@ -67,6 +73,7 @@ def from_dict(dct: dict[str, Any]) -> MSONAtoms: mson_atoms.info = atoms_info return mson_atoms + # NOTE: If making notable changes to this class, please ping @Andrew-S-Rosen on GitHub. # There are some subtleties in here, particularly related to spins/charges. class AseAtomsAdaptor: diff --git a/tests/io/test_ase.py b/tests/io/test_ase.py index 11fba7d13e8..8a8232bca81 100644 --- a/tests/io/test_ase.py +++ b/tests/io/test_ase.py @@ -323,11 +323,16 @@ def test_msonable_atoms(): ref_atoms = atoms.copy() ref_atoms.info = {} - + msonable_atoms = MSONAtoms(atoms) assert atoms == msonable_atoms - ref = {"@module": "pymatgen.io.ase", "@class": "MSONAtoms", "atoms_json": ase.io.jsonio.encode(ref_atoms), "atoms_info": atoms_info} + ref = { + "@module": "pymatgen.io.ase", + "@class": "MSONAtoms", + "atoms_json": ase.io.jsonio.encode(ref_atoms), + "atoms_info": atoms_info, + } assert msonable_atoms.as_dict() == ref atoms_back = MSONAtoms.from_dict(ref) From 213ddd9c15c6ae5688ae07a72bb8851418686290 Mon Sep 17 00:00:00 2001 From: "Andrew S. Rosen" Date: Mon, 4 Mar 2024 10:39:03 -0800 Subject: [PATCH 03/10] Update ase.py Signed-off-by: Andrew S. Rosen --- pymatgen/io/ase.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymatgen/io/ase.py b/pymatgen/io/ase.py index a34f1d4e639..1c454f84879 100644 --- a/pymatgen/io/ase.py +++ b/pymatgen/io/ase.py @@ -55,7 +55,7 @@ def as_dict(atoms: Atoms) -> dict[str, Any]: # to be used in a round-trip fashion and does not work properly with constraints. # See ASE issue #1387. atoms_info = atoms.info.copy() - atoms.info = {} + atoms.info: dict = {} return { "@module": "pymatgen.io.ase", "@class": "MSONAtoms", From c2005c6587ce02dcf443e320e5baba41922a61e9 Mon Sep 17 00:00:00 2001 From: "Andrew S. Rosen" Date: Mon, 4 Mar 2024 11:36:42 -0800 Subject: [PATCH 04/10] Update test_ase.py Signed-off-by: Andrew S. Rosen --- tests/io/test_ase.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/io/test_ase.py b/tests/io/test_ase.py index 8a8232bca81..47f302d42c5 100644 --- a/tests/io/test_ase.py +++ b/tests/io/test_ase.py @@ -316,8 +316,9 @@ def test_back_forth_v4(): @skip_if_no_ase def test_msonable_atoms(): + atoms = ase.io.read(f"{TEST_FILES_DIR}/OUTCAR") - atoms_info = {"test": "hi"} + atoms_info = {"test": "hi", "structure": structure} atoms.info = atoms_info assert not isinstance(atoms, MSONAtoms) From 95af7e2c8492fc47daa8874526a4ce42c8f93bb1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 4 Mar 2024 19:37:47 +0000 Subject: [PATCH 05/10] pre-commit auto-fixes --- tests/io/test_ase.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/io/test_ase.py b/tests/io/test_ase.py index 47f302d42c5..1ec92f81a1f 100644 --- a/tests/io/test_ase.py +++ b/tests/io/test_ase.py @@ -316,7 +316,6 @@ def test_back_forth_v4(): @skip_if_no_ase def test_msonable_atoms(): - atoms = ase.io.read(f"{TEST_FILES_DIR}/OUTCAR") atoms_info = {"test": "hi", "structure": structure} atoms.info = atoms_info From 4266eda0cc99fb0c0a3880f49343b5ea3c756877 Mon Sep 17 00:00:00 2001 From: "Andrew S. Rosen" Date: Mon, 4 Mar 2024 11:41:19 -0800 Subject: [PATCH 06/10] Update ase.py Signed-off-by: Andrew S. Rosen --- pymatgen/io/ase.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pymatgen/io/ase.py b/pymatgen/io/ase.py index 1c454f84879..6090f96123b 100644 --- a/pymatgen/io/ase.py +++ b/pymatgen/io/ase.py @@ -54,13 +54,13 @@ def as_dict(atoms: Atoms) -> dict[str, Any]: # @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. - atoms_info = atoms.info.copy() - atoms.info: dict = {} + atoms_no_info = atoms.copy() + atoms_no_info.info: dict = {} return { "@module": "pymatgen.io.ase", "@class": "MSONAtoms", - "atoms_json": encode(atoms), - "atoms_info": jsanitize(atoms_info, strict=True), + "atoms_json": encode(atoms_no_info), + "atoms_info": jsanitize(atoms.info, strict=True), } def from_dict(dct: dict[str, Any]) -> MSONAtoms: From a63d982b941d331402a57871e96636c03bc7add9 Mon Sep 17 00:00:00 2001 From: "Andrew S. Rosen" Date: Mon, 4 Mar 2024 11:42:37 -0800 Subject: [PATCH 07/10] Update test_ase.py Signed-off-by: Andrew S. Rosen --- tests/io/test_ase.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/io/test_ase.py b/tests/io/test_ase.py index 1ec92f81a1f..894d343721a 100644 --- a/tests/io/test_ase.py +++ b/tests/io/test_ase.py @@ -337,7 +337,7 @@ def test_msonable_atoms(): atoms_back = MSONAtoms.from_dict(ref) assert atoms_back == atoms - assert atoms_back.info == atoms_info + assert atoms_back.info == atoms.info structure = Structure.from_file(f"{TEST_FILES_DIR}/POSCAR") From 47d938456ff80d532090a10f7e89875010e7a86e Mon Sep 17 00:00:00 2001 From: Andrew Rosen Date: Mon, 4 Mar 2024 11:56:39 -0800 Subject: [PATCH 08/10] fix --- tests/io/test_ase.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/tests/io/test_ase.py b/tests/io/test_ase.py index 894d343721a..5274194e16e 100644 --- a/tests/io/test_ase.py +++ b/tests/io/test_ase.py @@ -316,6 +316,8 @@ 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 @@ -327,20 +329,13 @@ def test_msonable_atoms(): msonable_atoms = MSONAtoms(atoms) assert atoms == msonable_atoms - ref = { - "@module": "pymatgen.io.ase", - "@class": "MSONAtoms", - "atoms_json": ase.io.jsonio.encode(ref_atoms), - "atoms_info": atoms_info, - } - assert msonable_atoms.as_dict() == ref + 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(ref) + atoms_back = MSONAtoms.from_dict(msonable_atoms_dict) assert atoms_back == atoms assert atoms_back.info == atoms.info - structure = Structure.from_file(f"{TEST_FILES_DIR}/POSCAR") - atoms = AseAtomsAdaptor.get_atoms(structure, msonable=True) assert callable(atoms.as_dict) assert callable(atoms.from_dict) From 24372a789bb0b9e4aa42cc9504d13fb467e26ccf Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 4 Mar 2024 19:57:24 +0000 Subject: [PATCH 09/10] pre-commit auto-fixes --- tests/io/test_ase.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/io/test_ase.py b/tests/io/test_ase.py index 5274194e16e..ef23ae1ca86 100644 --- a/tests/io/test_ase.py +++ b/tests/io/test_ase.py @@ -330,7 +330,12 @@ def test_msonable_atoms(): assert atoms == msonable_atoms 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)} + 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 From 058bdb9e4fddb0ba9a00d5a8a93483d91486325a Mon Sep 17 00:00:00 2001 From: "Andrew S. Rosen" Date: Mon, 4 Mar 2024 12:02:08 -0800 Subject: [PATCH 10/10] Update ase.py Signed-off-by: Andrew S. Rosen --- pymatgen/io/ase.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymatgen/io/ase.py b/pymatgen/io/ase.py index 6090f96123b..7d959179bf5 100644 --- a/pymatgen/io/ase.py +++ b/pymatgen/io/ase.py @@ -55,7 +55,7 @@ def as_dict(atoms: Atoms) -> dict[str, Any]: # to be used in a round-trip fashion and does not work properly with constraints. # See ASE issue #1387. atoms_no_info = atoms.copy() - atoms_no_info.info: dict = {} + atoms_no_info.info = {} return { "@module": "pymatgen.io.ase", "@class": "MSONAtoms",