Fix branch directory check in io.vasp.outputs.get_band_structure_from_vasp_multiple_branches

src/pymatgen/electronic_structure/bandstructure.py

@@ -1071,7 +1071,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/io/vasp/outputs.py

@@ -4277,55 +4277,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}")
-
-        return get_reconstructed_band_structure(branches, 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(
+        # 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=}")
+
+            run = Vasprun(vasprun_file, parse_projected_eigen=projections)
+            bs_branches.append(run.get_band_structure(efermi=efermi))
+
+        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 2025-09-14
+    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 2025-09-14\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:

tests/io/vasp/test_outputs.py

@@ -12,13 +12,15 @@
 import numpy as np
 import pytest
 from monty.io import zopen
+from monty.shutil import decompress_file
+from monty.tempfile import ScratchDir
 from numpy.testing import assert_allclose
 from pytest import approx
 
 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
@@ -39,6 +41,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, PymatgenTest
@@ -1429,13 +1432,50 @@ 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: use real multi-branch bandstructure calculation."""
+        with ScratchDir("."):
+            # Create branches
+            for idx in range(3):  # simulate 3 branches
+                branch_name = f"branch_{idx}"
+                os.makedirs(branch_name)
+                copyfile(f"{VASP_OUT_DIR}/vasprun.force_hybrid_like_calc.xml.gz", f"./{branch_name}/vasprun.xml.gz")
+                decompress_file(f"./{branch_name}/vasprun.xml.gz")
+
+            get_band_structure_from_vasp_multiple_branches(".")
+
+    def test_missing_vasprun_in_branch_dir(self):
+        """Test vasprun.xml missing from branch_*."""
+        with ScratchDir("."):
+            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."""
+        with ScratchDir("."):
+            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"), ScratchDir("."):
+            get_band_structure_from_vasp_multiple_branches(".")
 
 
 class TestLocpot(PymatgenTest):