From 233c2a908d08a791ebf181008d94a56b70aaedb7 Mon Sep 17 00:00:00 2001 From: Jimmy Shen <14003693+jmmshn@users.noreply.github.com> Date: Thu, 23 Jan 2025 13:50:51 -0800 Subject: [PATCH] allow newer pymatgen verions (#1105) Co-authored-by: Jimmy Shen Co-authored-by: J. George --- src/atomate2/vasp/sets/core.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/atomate2/vasp/sets/core.py b/src/atomate2/vasp/sets/core.py index 618ded50e5..daa9cebcf1 100644 --- a/src/atomate2/vasp/sets/core.py +++ b/src/atomate2/vasp/sets/core.py @@ -687,13 +687,18 @@ def incar_updates(self) -> dict: @staticmethod def _get_ensemble_defaults(structure: Structure, ensemble: str) -> dict[str, Any]: """Get default params for the ensemble.""" + # Handle both old (ntypesp) and new (n_elems) pymatgen versions + n_types = getattr(structure, "n_elems", None) + if n_types is None: + n_types = structure.ntypesp + defaults = { "nve": {"MDALGO": 1, "ISIF": 2, "ANDERSEN_PROB": 0.0}, "nvt": {"MDALGO": 2, "ISIF": 2, "SMASS": 0}, "npt": { "MDALGO": 3, "ISIF": 3, - "LANGEVIN_GAMMA": [10] * structure.ntypesp, + "LANGEVIN_GAMMA": [10] * n_types, "LANGEVIN_GAMMA_L": 1, "PMASS": 10, "PSTRESS": 0,