Don't rely on jsanitize in Atoms <--> Structure object interconversion

pymatgen/io/ase.py

@@ -7,10 +7,11 @@
 from __future__ import annotations
 
 import warnings
+from collections.abc import Iterable
 from typing import TYPE_CHECKING
 
 import numpy as np
-from monty.json import jsanitize
+from monty.dev import requires
 
 from pymatgen.core.structure import Molecule, Structure
 
@@ -23,6 +24,7 @@
     from ase import Atoms
     from ase.calculators.singlepoint import SinglePointDFTCalculator
     from ase.constraints import FixAtoms
+    from ase.spacegroup import Spacegroup
 
     ase_loaded = True
 except ImportError:
@@ -38,6 +40,7 @@
 
 # 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.
+@requires(ase_loaded, "ASE needs to be installed.")
 class AseAtomsAdaptor:
     """Adaptor serves as a bridge between ASE Atoms and pymatgen objects."""
 
@@ -53,11 +56,6 @@ def get_atoms(structure: SiteCollection, **kwargs) -> Atoms:
         Returns:
             Atoms: ASE Atoms object
         """
-        if not ase_loaded:
-            raise ImportError(
-                "AseAtomsAdaptor requires the ASE package.\n"
-                "Use `pip install ase` or `conda install ase -c conda-forge`"
-            )
         if not structure.is_ordered:
             raise ValueError("ASE Atoms only supports ordered structures")
 
@@ -112,20 +110,6 @@ def get_atoms(structure: SiteCollection, **kwargs) -> Atoms:
             atoms.charge = structure.charge
             atoms.spin_multiplicity = structure.spin_multiplicity
 
-        # Set the Atoms final magnetic moments and charges if present.
-        # This uses the SinglePointDFTCalculator as the dummy calculator
-        # to store results.
-        charges = structure.site_properties.get("final_charge")
-        magmoms = structure.site_properties.get("final_magmom")
-        if magmoms or charges:
-            if magmoms and charges:
-                calc = SinglePointDFTCalculator(atoms, magmoms=magmoms, charges=charges)
-            elif magmoms:
-                calc = SinglePointDFTCalculator(atoms, magmoms=magmoms)
-            else:
-                calc = SinglePointDFTCalculator(atoms, charges=charges)
-            atoms.calc = calc
-
         # Get the oxidation states from the structure
         oxi_states: list[float | None] = [getattr(site.specie, "oxi_state", None) for site in structure]
 
@@ -136,8 +120,11 @@ def get_atoms(structure: SiteCollection, **kwargs) -> Atoms:
             fix_atoms = []
             for site in structure:
                 selective_dynamics: ArrayLike = site.properties.get("selective_dynamics")  # type: ignore[assignment]
-                if not (np.all(selective_dynamics) or not np.any(selective_dynamics)):
-                    # should be [True, True, True] or [False, False, False]
+                if (
+                    isinstance(selective_dynamics, Iterable)
+                    and True in selective_dynamics
+                    and False in selective_dynamics
+                ):
                     raise ValueError(
                         "ASE FixAtoms constraint does not support selective dynamics in only some dimensions. "
                         f"Remove the {selective_dynamics=} and try again if you do not need them."
@@ -160,11 +147,27 @@ def get_atoms(structure: SiteCollection, **kwargs) -> Atoms:
             atoms.set_array("oxi_states", np.array(oxi_states))
 
         # Atoms.info <---> Structure.properties
-        # Atoms.calc <---> Structure.calc
         if properties := getattr(structure, "properties"):  # noqa: B009
             atoms.info = properties
+
+        # Regenerate Spacegroup object from `.todict()` representation
+        if isinstance(atoms.info.get("spacegroup"), dict):
+            atoms.info["spacegroup"] = Spacegroup(
+                atoms.info["spacegroup"]["number"], setting=atoms.info["spacegroup"].get("setting", 1)
+            )
+
+        # Atoms.calc <---> Structure.calc
         if calc := getattr(structure, "calc", None):
             atoms.calc = calc
+        else:
+            # Set the Atoms final magnetic moments and charges if present.
+            # This uses the SinglePointDFTCalculator as the dummy calculator
+            # to store results.
+            charges = structure.site_properties.get("final_charge")
+            magmoms = structure.site_properties.get("final_magmom")
+            if charges or magmoms:
+                calc = SinglePointDFTCalculator(atoms, magmoms=magmoms, charges=charges)
+                atoms.calc = calc
 
         return atoms
 
@@ -212,13 +215,16 @@ def get_structure(atoms: Atoms, cls: type[Structure] = Structure, **cls_kwargs)
                 else:
                     unsupported_constraint_type = True
             if unsupported_constraint_type:
-                warnings.warn("Only FixAtoms is supported by Pymatgen. Other constraints will not be set.")
+                warnings.warn("Only FixAtoms is supported by Pymatgen. Other constraints will not be set.", UserWarning)
             sel_dyn = [[False] * 3 if atom.index in constraint_indices else [True] * 3 for atom in atoms]
         else:
             sel_dyn = None
 
         # Atoms.info <---> Structure.properties
-        properties = jsanitize(getattr(atoms, "info", {}))
+        # But first make sure `spacegroup` is JSON serializable
+        if atoms.info.get("spacegroup") and isinstance(atoms.info["spacegroup"], Spacegroup):
+            atoms.info["spacegroup"] = atoms.info["spacegroup"].todict()
+        properties = getattr(atoms, "info", {})
 
         # Return a Molecule object if that was specifically requested;
         # otherwise return a Structure object as expected