diff --git a/pymatgen/core/tests/test_trajectory.py b/pymatgen/core/tests/test_trajectory.py index ed924b436a5..ba5eade4c93 100644 --- a/pymatgen/core/tests/test_trajectory.py +++ b/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, diff --git a/pymatgen/io/vasp/inputs.py b/pymatgen/io/vasp/inputs.py index 63250d9451f..54da79f251e 100644 --- a/pymatgen/io/vasp/inputs.py +++ b/pymatgen/io/vasp/inputs.py @@ -94,7 +94,7 @@ 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, @@ -102,32 +102,43 @@ def __init__( 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: diff --git a/pymatgen/io/vasp/tests/test_inputs.py b/pymatgen/io/vasp/tests/test_inputs.py index 070077cf71a..11554b186e3 100644 --- a/pymatgen/io/vasp/tests/test_inputs.py +++ b/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):