Fix ValueError when passing selective_dynamics to Poscar (#2951)

* Fixed selective_dynamics error for poscar in vasp/inputs.py

* add unit test

* Revert "add unit test"

This reverts commit 4bc5cf3.

* add unit test for selective dynamics in `io/vasp/tests`

* pre-commit auto-fixes

* correct typing

* pre-commit auto-fixes

* clean up

* doc str to google format

* PoscarTest convert selective_dynamics to list[list[bool]] before compare

---------

Co-authored-by: Janosh Riebesell <[email protected]>

pymatgen/core/tests/test_trajectory.py

@@ -164,15 +164,9 @@ def test_site_properties(self):
         species, coords, charge, spin = self._get_species_and_coords()
 
         props = [
-            {
-                "test": [[True, True, True], [False, False, False]],
-            },
-            {
-                "test": [[False, False, False], [False, False, False]],
-            },
-            {
-                "test": [[True, True, True], [False, False, False]],
-            },
+            {"test": [[True, True, True], [False, False, False]]},
+            {"test": [[False, False, False], [False, False, False]]},
+            {"test": [[True, True, True], [False, False, False]]},
         ]
         traj = Trajectory(
             species=species,

pymatgen/io/vasp/inputs.py

@@ -94,40 +94,51 @@ def __init__(
         self,
         structure: Structure,
         comment: str | None = None,
-        selective_dynamics=None,
+        selective_dynamics: ArrayLike | None = None,
         true_names: bool = True,
         velocities: ArrayLike | None = None,
         predictor_corrector: ArrayLike | None = None,
         predictor_corrector_preamble: str | None = None,
         sort_structure: bool = False,
     ):
         """
-        :param structure: Structure object.
-        :param comment: Optional comment line for POSCAR. Defaults to unit
-            cell formula of structure. Defaults to None.
-        :param selective_dynamics: bool values for selective dynamics,
-            where N is number of sites. Defaults to None.
-        :param true_names: Set to False if the names in the POSCAR are not
-            well-defined and ambiguous. This situation arises commonly in
-            vasp < 5 where the POSCAR sometimes does not contain element
-            symbols. Defaults to True.
-        :param velocities: Velocities for the POSCAR. Typically parsed
-            in MD runs or can be used to initialize velocities.
-        :param predictor_corrector: Predictor corrector for the POSCAR.
-            Typically parsed in MD runs.
-        :param predictor_corrector_preamble: Preamble to the predictor
-            corrector.
-        :param sort_structure: Whether to sort structure. Useful if species
-            are not grouped properly together.
+        Args:
+            structure (Structure): Structure object.
+            comment (str | None, optional): Optional comment line for POSCAR. Defaults to unit
+                cell formula of structure. Defaults to None.
+            selective_dynamics (ArrayLike | None, optional): Bool values for selective dynamics,
+                where N is the number of sites. Defaults to None.
+            true_names (bool, optional): Set to False if the names in the POSCAR are not
+                well-defined and ambiguous. This situation arises commonly in
+                VASP < 5 where the POSCAR sometimes does not contain element
+                symbols. Defaults to True.
+            velocities (ArrayLike | None, optional): Velocities for the POSCAR. Typically parsed
+                in MD runs or can be used to initialize velocities. Defaults to None.
+            predictor_corrector (ArrayLike | None, optional): Predictor corrector for the POSCAR.
+                Typically parsed in MD runs. Defaults to None.
+            predictor_corrector_preamble (str | None, optional): Preamble to the predictor
+                corrector. Defaults to None.
+            sort_structure (bool, optional): Whether to sort the structure. Useful if species
+                are not grouped properly together. Defaults to False.
         """
         if structure.is_ordered:
             site_properties = {}
-            if selective_dynamics:
-                site_properties["selective_dynamics"] = selective_dynamics
-            if velocities:
-                site_properties["velocities"] = velocities
-            if predictor_corrector:
-                site_properties["predictor_corrector"] = predictor_corrector
+
+            if selective_dynamics is not None:
+                selective_dynamics = np.array(selective_dynamics)
+                if not selective_dynamics.all():
+                    site_properties["selective_dynamics"] = selective_dynamics
+
+            if velocities is not None:
+                velocities = np.array(velocities)
+                if velocities.any():
+                    site_properties["velocities"] = velocities
+
+            if predictor_corrector is not None:
+                predictor_corrector = np.array(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:

pymatgen/io/vasp/tests/test_inputs.py

@@ -84,7 +84,9 @@ def test_init(self):
 0.750000 0.500000 0.750000 F F F O
 """
         poscar = Poscar.from_string(poscar_string)
-        assert poscar.selective_dynamics == [[True, True, True], [False, False, False]]
+        selective_dynamics = [list(x) for x in poscar.selective_dynamics]
+
+        assert selective_dynamics == [[True, True, True], [False, False, False]]
         self.selective_poscar = poscar
 
     def test_from_file(self):
@@ -370,6 +372,34 @@ def test_write(self):
         self.assertArrayAlmostEqual(poscar.structure.lattice.abc, p.structure.lattice.abc, 5)
         tempfname.unlink()
 
+    def test_selective_dynamics(self):
+        filepath = PymatgenTest.TEST_FILES_DIR / "POSCAR.Fe3O4"
+        poscar = Poscar.from_file(filepath)
+        structure = poscar.structure
+
+        # Fix bottom half
+        fixed_indices = structure.frac_coords[:, 2] >= 0.5
+
+        poscar = Poscar(structure, selective_dynamics=np.tile(fixed_indices.reshape(-1, 1), [1, 3]))
+        selective_dynamics = [list(x) for x in poscar.selective_dynamics]
+
+        assert selective_dynamics == [
+            [True, True, True],
+            [False, False, False],
+            [False, False, False],
+            [True, True, True],
+            [False, False, False],
+            [True, True, True],
+            [True, True, True],
+            [False, False, False],
+            [False, False, False],
+            [True, True, True],
+            [True, True, True],
+            [False, False, False],
+            [True, True, True],
+            [False, False, False],
+        ]
+
 
 class IncarTest(PymatgenTest):
     def setUp(self):