Fix branch directory check in io.vasp.outputs.get_band_structure_from_vasp_multiple_branches (#4061)

* docstring tweak and support Path

* isfile -> isdir fix

* raise error if vasprun.xml missing in any branch dir

* sort pymatgen core imports

* add some unit test

* raise error if no vasprun.xml file found at all

* deprecation warning for fall back branch

* remove unused ignore tag

* add a temporary solution

* update inherit_incar docstring

* remove scratch dir

* extend deprecation deadline

* drop unsure test

---------

Co-authored-by: Matt Horton <[email protected]>

Author: DanielYang59

docs/index.md

@@ -1084,7 +1084,7 @@ def get_projections_on_elements_and_orbitals(
 
 
 @overload
-def get_reconstructed_band_structure(  # type: ignore[overload-overlap]
+def get_reconstructed_band_structure(
     list_bs: list[BandStructure],
     efermi: float | None = None,
 ) -> BandStructure:

src/pymatgen/electronic_structure/bandstructure.py

@@ -4535,58 +4535,62 @@ class VaspParseError(ParseError):
 
 
 def get_band_structure_from_vasp_multiple_branches(
-    dir_name: str,
+    dir_name: PathLike,
     efermi: float | None = None,
     projections: bool = False,
 ) -> BandStructureSymmLine | BandStructure | None:
     """Get band structure info from a VASP directory.
 
-    It takes into account that a run can be divided in several branches named
-    "branch_x". If the run has not been divided in branches the method will
-    turn to parsing vasprun.xml directly.
+    It takes into account that a run can be divided in several branches,
+    each inside a directory named "branch_x". If the run has not been
+    divided in branches the function will turn to parse vasprun.xml
+    directly from the selected directory.
 
     Args:
-        dir_name: Directory containing all bandstructure runs.
-        efermi: Efermi for bandstructure.
-        projections: True if you want to get the data on site projections if
-            any. Note that this is sometimes very large
+        dir_name (PathLike): Parent directory containing all bandstructure runs.
+        efermi (float): Fermi level for bandstructure.
+        projections (bool): True if you want to get the data on site
+            projections if any. Note that this is sometimes very large
 
     Returns:
-        A BandStructure Object.
-        None is there's a parsing error.
+        A BandStructure/BandStructureSymmLine Object.
+        None if no vasprun.xml found in given directory and branch directory.
     """
-    # TODO: Add better error handling!!!
-    if os.path.isfile(f"{dir_name}/branch_0"):
-        # Get all branch dir names
+    if os.path.isdir(f"{dir_name}/branch_0"):
+        # Get and sort all branch directories
         branch_dir_names = [os.path.abspath(d) for d in glob(f"{dir_name}/branch_*") if os.path.isdir(d)]
-
-        # Sort by the directory name (e.g, branch_10)
         sorted_branch_dir_names = sorted(branch_dir_names, key=lambda x: int(x.split("_")[-1]))
 
-        # Populate branches with Bandstructure instances
-        branches = []
-        for dname in sorted_branch_dir_names:
-            xml_file = f"{dname}/vasprun.xml"
-            if os.path.isfile(xml_file):
-                run = Vasprun(xml_file, parse_projected_eigen=projections)
-                branches.append(run.get_band_structure(efermi=efermi))
-            else:
-                # TODO: It might be better to throw an exception
-                warnings.warn(
-                    f"Skipping {dname}. Unable to find {xml_file}",
-                    stacklevel=2,
-                )
+        # Collect BandStructure from all branches
+        bs_branches: list[BandStructure | BandStructureSymmLine] = []
+        for directory in sorted_branch_dir_names:
+            vasprun_file = f"{directory}/vasprun.xml"
+            if not os.path.isfile(vasprun_file):
+                raise FileNotFoundError(f"cannot find vasprun.xml in {directory=}")
 
-        return get_reconstructed_band_structure(branches, efermi)
+            run = Vasprun(vasprun_file, parse_projected_eigen=projections)
+            bs_branches.append(run.get_band_structure(efermi=efermi))
 
-    xml_file = f"{dir_name}/vasprun.xml"
-    # Better handling of Errors
-    if os.path.isfile(xml_file):
-        return Vasprun(xml_file, parse_projected_eigen=projections).get_band_structure(
+        return get_reconstructed_band_structure(bs_branches, efermi)
+
+    # Read vasprun.xml directly if no branch head (branch_0) is found
+    # TODO: remove this branch and raise error directly after 2026-06-01
+    vasprun_file = f"{dir_name}/vasprun.xml"
+    if os.path.isfile(vasprun_file):
+        warnings.warn(
+            (
+                f"no branch dir found, reading directly from {dir_name=}\n"
+                "this fallback branch would be removed after 2026-06-01\n"
+                "please check your data dir or use Vasprun.get_band_structure directly"
+            ),
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return Vasprun(vasprun_file, parse_projected_eigen=projections).get_band_structure(
             kpoints_filename=None, efermi=efermi
         )
 
-    return None
+    raise FileNotFoundError(f"failed to find any vasprun.xml in selected {dir_name=}")
 
 
 class Xdatcar:

src/pymatgen/io/vasp/outputs.py

@@ -181,10 +181,11 @@ class VaspInputSet(InputGenerator, abc.ABC):
             Curtarolo) for monoclinic. Defaults True.
         validate_magmom (bool): Ensure that the missing magmom values are filled in with
             the VASP default value of 1.0.
-        inherit_incar (bool): Whether to inherit INCAR settings from previous
+        inherit_incar (bool | list[str]): Whether to inherit INCAR settings from previous
             calculation. This might be useful to port Custodian fixes to child jobs but
             can also be dangerous e.g. when switching from GGA to meta-GGA or relax to
             static jobs. Defaults to True.
+            Can also be a list of strings to specify which parameters are inherited.
         auto_kspacing (bool): If true, determines the value of KSPACING from the bandgap
             of a previous calculation.
         auto_ismear (bool): If true, the values for ISMEAR and SIGMA will be set

src/pymatgen/io/vasp/sets.py

@@ -19,7 +19,7 @@
 from pymatgen.core import Element
 from pymatgen.core.lattice import Lattice
 from pymatgen.core.structure import Structure
-from pymatgen.electronic_structure.bandstructure import BandStructureSymmLine
+from pymatgen.electronic_structure.bandstructure import BandStructure, BandStructureSymmLine
 from pymatgen.electronic_structure.core import Magmom, Orbital, OrbitalType, Spin
 from pymatgen.entries.compatibility import MaterialsProjectCompatibility
 from pymatgen.io.vasp.inputs import Incar, Kpoints, Poscar, Potcar
@@ -43,6 +43,7 @@
     Wavecar,
     Waveder,
     Xdatcar,
+    get_band_structure_from_vasp_multiple_branches,
 )
 from pymatgen.io.wannier90 import Unk
 from pymatgen.util.testing import FAKE_POTCAR_DIR, TEST_FILES_DIR, VASP_IN_DIR, VASP_OUT_DIR, MatSciTest
@@ -1533,13 +1534,40 @@ def test_init(self):
         assert len(oszicar.electronic_steps) == len(oszicar.ionic_steps)
         assert len(oszicar.all_energies) == 60
         assert oszicar.final_energy == approx(-526.63928)
-        assert set(oszicar.ionic_steps[-1]) == set({"F", "E0", "dE", "mag"})
+        assert set(oszicar.ionic_steps[-1]) == {"F", "E0", "dE", "mag"}
 
     def test_static(self):
         fpath = f"{TEST_DIR}/fixtures/static_silicon/OSZICAR"
         oszicar = Oszicar(fpath)
         assert oszicar.final_energy == approx(-10.645278)
-        assert set(oszicar.ionic_steps[-1]) == set({"F", "E0", "dE", "mag"})
+        assert set(oszicar.ionic_steps[-1]) == {"F", "E0", "dE", "mag"}
+
+
+class TestGetBandStructureFromVaspMultipleBranches:
+    def test_read_multi_branches(self):
+        """TODO: This functionality still needs a test."""
+
+    def test_missing_vasprun_in_branch_dir(self):
+        """Test vasprun.xml missing from branch_*."""
+        os.makedirs("no_vasp/branch_0", exist_ok=False)
+
+        with pytest.raises(FileNotFoundError, match="cannot find vasprun.xml in directory"):
+            get_band_structure_from_vasp_multiple_branches("no_vasp")
+
+    def test_no_branch_head(self):
+        """Test branch_0 is missing and read dir_name/vasprun.xml directly."""
+
+        copyfile(f"{VASP_OUT_DIR}/vasprun.force_hybrid_like_calc.xml.gz", "./vasprun.xml.gz")
+        decompress_file("./vasprun.xml.gz")
+
+        with pytest.warns(DeprecationWarning, match="no branch dir found, reading directly from"):
+            bs = get_band_structure_from_vasp_multiple_branches(".")
+        assert isinstance(bs, BandStructure)
+
+    def test_cannot_read_anything(self):
+        """Test no branch_0/, no dir_name/vasprun.xml, no vasprun.xml at all."""
+        with pytest.raises(FileNotFoundError, match="failed to find any vasprun.xml in selected"):
+            get_band_structure_from_vasp_multiple_branches(".")
 
 
 class TestLocpot(MatSciTest):