Handle numpy array for selective dynamics in Structure (#4461)

* enable orjson.OPT_SERIALIZE_NUMPY

* np.array -> np.asarray

* remove tolist

* safer deprecation replacement

* rename

* better deprecation

* put structure related ops together

* add test

* str -> PathLike for filename

* clean up

* also enable OPT_SERIALIZE_NUMPY for molecule

* overwrite default grou

* ignore arg-type, fix needed from monty

* Revert "overwrite default grou"

This reverts commit 409d954.

Author: DanielYang59

src/pymatgen/core/structure.py

@@ -520,7 +520,7 @@ def is_valid(self, tol: float = DISTANCE_TOLERANCE) -> bool:
         return bool(np.min(all_dists) > tol)
 
     @abstractmethod
-    def to(self, filename: str = "", fmt: FileFormats = "") -> str | None:
+    def to(self, filename: PathLike = "", fmt: FileFormats = "") -> str | None:
         """Generate string representations (cif, json, poscar, ....) of SiteCollections (e.g.,
         molecules / structures). Should return str or None if written to a file.
         """
@@ -564,6 +564,7 @@ def add_site_property(self, property_name: str, values: Sequence | NDArray) -> S
         """
         if len(values) != len(self):
             raise ValueError(f"{len(values)=} must equal sites in structure={len(self)}")
+
         for site, val in zip(self, values, strict=True):
             site.properties[property_name] = val
 
@@ -2964,7 +2965,11 @@ def to(self, filename: PathLike = "", fmt: FileFormats = "", **kwargs) -> str:
             writer = Cssr(self)
 
         elif fmt == "json" or fnmatch(filename.lower(), "*.json*"):
-            json_str = json.dumps(self.as_dict(), **kwargs) if kwargs else orjson.dumps(self.as_dict()).decode()
+            json_str = (
+                json.dumps(self.as_dict(), **kwargs)
+                if kwargs
+                else orjson.dumps(self.as_dict(), option=orjson.OPT_SERIALIZE_NUMPY).decode()
+            )
 
             if filename:
                 with zopen(filename, mode="wt", encoding="utf-8") as file:
@@ -3989,11 +3994,11 @@ def get_centered_molecule(self) -> Self:
             properties=self.properties,
         )
 
-    def to(self, filename: str = "", fmt: str = "") -> str | None:
+    def to(self, filename: PathLike = "", fmt: str = "") -> str | None:
         """Outputs the molecule to a file or string.
 
         Args:
-            filename (str): If provided, output will be written to a file. If
+            filename (PathLike): If provided, output will be written to a file. If
                 fmt is not specified, the format is determined from the
                 filename. Defaults is None, i.e. string output.
             fmt (str): Format to output to. Defaults to JSON unless filename
@@ -4006,22 +4011,28 @@ def to(self, filename: str = "", fmt: str = "") -> str | None:
             str: String representation of molecule in given format. If a filename
                 is provided, the same string is written to the file.
         """
+        filename = str(filename)
         fmt = fmt.lower()
+
         writer: Any
         if fmt == "xyz" or fnmatch(filename.lower(), "*.xyz*"):
             from pymatgen.io.xyz import XYZ
 
             writer = XYZ(self)
+
         elif any(fmt == ext or fnmatch(filename.lower(), f"*.{ext}*") for ext in ("gjf", "g03", "g09", "com", "inp")):
             from pymatgen.io.gaussian import GaussianInput
 
             writer = GaussianInput(self)
+
         elif fmt == "json" or fnmatch(filename, "*.json*") or fnmatch(filename, "*.mson*"):
-            json_str = orjson.dumps(self.as_dict()).decode()
+            json_str = orjson.dumps(self.as_dict(), option=orjson.OPT_SERIALIZE_NUMPY).decode()
+
             if filename:
                 with zopen(filename, mode="wt", encoding="utf-8") as file:
                     file.write(json_str)  # type:ignore[arg-type]
             return json_str
+
         elif fmt in {"yaml", "yml"} or fnmatch(filename, "*.yaml*") or fnmatch(filename, "*.yml*"):
             yaml = YAML()
             str_io = io.StringIO()
@@ -4031,6 +4042,7 @@ def to(self, filename: str = "", fmt: str = "") -> str | None:
                 with zopen(filename, mode="wt", encoding="utf-8") as file:
                     file.write(yaml_str)  # type:ignore[arg-type]
             return yaml_str
+
         else:
             from pymatgen.io.babel import BabelMolAdaptor
 
@@ -4042,6 +4054,7 @@ def to(self, filename: str = "", fmt: str = "") -> str | None:
 
         if filename:
             writer.write_file(filename)
+
         return str(writer)
 
     @classmethod
@@ -4109,28 +4122,35 @@ def from_file(cls, filename: PathLike) -> IMolecule | Molecule:  # type:ignore[o
             Molecule
         """
         filename = str(filename)
+        fname = filename.lower()
 
         with zopen(filename, mode="rt", encoding="utf-8") as file:
-            contents: str = file.read()  # type:ignore[assignment]
-        fname = filename.lower()
+            contents: str = cast("str", file.read())
+
         if fnmatch(fname, "*.xyz*"):
             return cls.from_str(contents, fmt="xyz")
+
         if any(fnmatch(fname.lower(), f"*.{r}*") for r in ("gjf", "g03", "g09", "com", "inp")):
             return cls.from_str(contents, fmt="g09")
+
         if any(fnmatch(fname.lower(), f"*.{r}*") for r in ("out", "lis", "log")):
             from pymatgen.io.gaussian import GaussianOutput
 
             return GaussianOutput(filename).final_structure
+
         if fnmatch(fname, "*.json*") or fnmatch(fname, "*.mson*"):
             return cls.from_str(contents, fmt="json")
+
         if fnmatch(fname, "*.yaml*") or fnmatch(filename, "*.yml*"):
             return cls.from_str(contents, fmt="yaml")
-        from pymatgen.io.babel import BabelMolAdaptor
 
         if match := re.search(r"\.(pdb|mol|mdl|sdf|sd|ml2|sy2|mol2|cml|mrv)", filename.lower()):
+            from pymatgen.io.babel import BabelMolAdaptor
+
             new = BabelMolAdaptor.from_file(filename, match[1]).pymatgen_mol
             new.__class__ = cls
             return new
+
         raise ValueError("Cannot determine file type.")
 
 

src/pymatgen/io/abinit/abitimer.py

@@ -628,8 +628,8 @@ def as_dict(self):
         return {at: self.__dict__[at] for at in AbinitTimerSection.FIELDS}
 
     @deprecated(as_dict, deadline=(2026, 4, 4))
-    def to_dict(self):
-        return self.as_dict()
+    def to_dict(self, *args, **kwargs):
+        return self.as_dict(*args, **kwargs)
 
     def to_csvline(self, with_header=False):
         """Return a string with data in CSV format. Add header if `with_header`."""

src/pymatgen/io/aims/inputs.py

@@ -935,8 +935,8 @@ def as_dict(self):
         }
 
     @deprecated(replacement=as_dict, deadline=(2026, 4, 4))
-    def to_dict(self):
-        return self.as_dict()
+    def to_dict(self, *args, **kwargs):
+        return self.as_dict(*args, **kwargs)
 
     @classmethod
     def from_dict(cls, dct: dict[str, Any]) -> SpeciesDefaults:

src/pymatgen/io/vasp/inputs.py

@@ -25,6 +25,7 @@
 import numpy as np
 import orjson
 import scipy.constants as const
+from monty.dev import deprecated
 from monty.io import zopen
 from monty.json import MontyDecoder, MSONable
 from monty.os import cd
@@ -141,41 +142,41 @@ def __init__(
         site_properties: dict[str, Any] = {}
 
         if selective_dynamics is not None:
-            selective_dynamics = np.array(selective_dynamics)
+            selective_dynamics = np.asarray(selective_dynamics)
             if not selective_dynamics.all():
                 site_properties["selective_dynamics"] = selective_dynamics
 
         if velocities:
-            velocities = np.array(velocities)
+            velocities = np.asarray(velocities)
             if velocities.any():
                 site_properties["velocities"] = velocities
 
         if predictor_corrector:
-            predictor_corrector = np.array(predictor_corrector)
+            predictor_corrector = np.asarray(predictor_corrector)
             if predictor_corrector.any():
                 site_properties["predictor_corrector"] = predictor_corrector
 
         structure = Structure.from_sites(structure)
         self.structure = structure.copy(site_properties=site_properties)
         if sort_structure:
             self.structure = self.structure.get_sorted_structure()
-        self.true_names = true_names
-        self.comment = structure.formula if comment is None else comment
+
         if predictor_corrector_preamble:
             self.structure.properties["predictor_corrector_preamble"] = predictor_corrector_preamble
 
         if lattice_velocities and np.any(lattice_velocities):
             self.structure.properties["lattice_velocities"] = np.asarray(lattice_velocities)
 
-        self.temperature = -1.0
+        self.true_names: bool = true_names
+        self.comment: str = structure.formula if comment is None else comment
+        self.temperature: float = -1.0
 
     def __setattr__(self, name: str, value: Any) -> None:
         if name in {"selective_dynamics", "velocities"} and value is not None and len(value) > 0:
-            value = np.array(value)
+            value = np.asarray(value)
             dim = value.shape
             if dim[1] != 3 or dim[0] != len(self.structure):
                 raise ValueError(f"{name} array must be same length as the structure.")
-            value = value.tolist()
 
         super().__setattr__(name, value)
 
@@ -687,7 +688,7 @@ def get_str(
             lines.append("")
             if self.predictor_corrector_preamble:
                 lines.append(self.predictor_corrector_preamble)
-                pred = np.array(self.predictor_corrector)
+                pred = np.asarray(self.predictor_corrector)
                 for col in range(3):
                     for z in pred[:, col]:
                         lines.append(" ".join(format_str.format(i) for i in z))
@@ -700,7 +701,9 @@ def get_str(
 
         return "\n".join(lines) + "\n"
 
-    get_string = get_str
+    @deprecated(get_str)
+    def get_string(self, *args, **kwargs):
+        return self.get_str(*args, **kwargs)
 
     def write_file(self, filename: PathLike, **kwargs) -> None:
         """Write POSCAR to a file. The supported kwargs are the same as those for
@@ -716,7 +719,7 @@ def as_dict(self) -> dict:
             "@class": type(self).__name__,
             "structure": self.structure.as_dict(),
             "true_names": self.true_names,
-            "selective_dynamics": np.array(self.selective_dynamics).tolist(),
+            "selective_dynamics": np.asarray(self.selective_dynamics).tolist(),
             "velocities": self.velocities,
             "predictor_corrector": self.predictor_corrector,
             "comment": self.comment,
@@ -2432,7 +2435,7 @@ def parse_fortran_style_str(input_str: str) -> str | bool | float | int:
 
         def data_stats(data_list: Sequence) -> dict:
             """Used for hash-less and therefore less brittle POTCAR validity checking."""
-            arr = np.array(data_list)
+            arr = np.asarray(data_list)
             return {
                 "MEAN": np.mean(arr),
                 "ABSMEAN": np.mean(np.abs(arr)),
@@ -2622,7 +2625,7 @@ def compare_potcar_stats(
         if key_match:
             data_diff = [
                 abs(potcar_stats_1["stats"].get(key, {}).get(stat) - potcar_stats_2["stats"].get(key, {}).get(stat))
-                for stat in ["MEAN", "ABSMEAN", "VAR", "MIN", "MAX"]
+                for stat in ("MEAN", "ABSMEAN", "VAR", "MIN", "MAX")
                 for key in check_potcar_fields
             ]
             data_match = all(np.array(data_diff) < tolerance)

src/pymatgen/io/vasp/sets.py

@@ -488,13 +488,13 @@ def get_input_set(
         )
 
     @deprecated(get_input_set, deadline=(2026, 6, 6))
-    def get_vasp_input(self, structure: Structure | None = None) -> VaspInput:
+    def get_vasp_input(self, *args, **kwargs) -> VaspInput:
         """Get a VaspInput object.
 
         Returns:
             VaspInput.
         """
-        return self.get_input_set(structure=structure)
+        return self.get_input_set(*args, **kwargs)
 
     @property
     def incar_updates(self) -> dict:

tests/core/test_structure.py

@@ -1252,7 +1252,7 @@ def test_add_remove_site_property(self):
         with pytest.raises(AttributeError, match="attr='magmom' not found on PeriodicSite"):
             _ = struct[0].magmom
 
-    def test_propertied_structure(self):
+    def test_site_properties(self):
         # Make sure that site properties are set to None for missing values.
         self.struct.add_site_property("charge", [4.1, -5])
         self.struct.append("Li", [0.3, 0.3, 0.3])
@@ -1278,6 +1278,23 @@ def test_propertied_structure(self):
         assert struct.properties == props
         assert dct == struct.as_dict()
 
+    def test_selective_dynamics(self):
+        """Ensure selective dynamics as numpy arrays can be JSON serialized."""
+        struct = self.get_structure("Li2O")
+        struct.add_site_property(
+            "selective_dynamics", np.array([[True, True, True], [False, False, False], [True, True, True]])
+        )
+
+        orjson_str = struct.to(fmt="json")
+
+        # Also test round trip
+        orjson_struct = Structure.from_str(orjson_str, fmt="json")
+        assert struct == orjson_struct
+
+        with pytest.raises(TypeError, match="Object of type ndarray is not JSON serializable"):
+            # Use a dummy kwarg (default value) to force `json.dumps`
+            struct.to(fmt="json", ensure_ascii=True)
+
     def test_perturb(self):
         struct = self.get_structure("Li2O")
         struct_orig = struct.copy()