Improve thermal electronic

dfttk/configuration.py

@@ -751,7 +751,9 @@ def process_thermal_electronic(
         kpoints_keys: list[str] = ["elec_dos"],
         kpoints_names: list[str] = ["KPOINTS.elec_dos"],
         contcar_name: str = "CONTCAR.elec_dos",
-        selected_volumes: list[float] = None,
+        vasprun_name: str = "vasprun.xml.elec_dos",
+        selected_volumes: np.ndarray | None = None,
+        selected_folders: list[str] | None = None,
         folder_prefix: str = "elec",
         order: int = 1,
     ):
@@ -769,7 +771,9 @@ def process_thermal_electronic(
             kpoints_keys (list[str], optional): List of KPOINTS keys for dictionary keys. Defaults to ["elec_dos"].
             kpoints_names (list[str], optional): List of KPOINTS names to read. Defaults to ["KPOINTS.elec_dos"].
             contcar_name (str, optional): Name of the CONTCAR file. Defaults to "CONTCAR.elec_dos".
-            selected_volumes (list[float], optional): List of volumes to process. Defaults to None.
+            vasprun_name (str, optional): Name of the vasprun.xml file. Defaults to ``"vasprun.xml.elec_dos"``.
+            selected_volumes (np.ndarray | None, optional): Array of volumes to process. Defaults to None.
+            selected_folders (list[str] | None, optional): List of selected folders to keep the electron DOS data. Defaults to None.
             folder_prefix (str, optional): Prefix for the folders containing electronic DOS data. Defaults to "elec".
             order (int, optional): Order of polynomial fitting. Defaults to 1.
         """
@@ -781,14 +785,17 @@ def process_thermal_electronic(
             kpoints_names=kpoints_names,
             contcar_name=contcar_name,
             selected_volumes=selected_volumes,
+            selected_folders=selected_folders,
             folder_prefix=folder_prefix,
         )
         self.thermal_electronic.get_thermal_electronic_data(
             volumes_fit=volumes_fit,
             temperatures=temperatures,
             order=order,
-            selected_volumes=selected_volumes,
             folder_prefix=folder_prefix,
+            vasprun_name=vasprun_name,
+            selected_volumes=selected_volumes,
+            selected_folders=selected_folders,
         )
 
     def process_qha(

dfttk/thermal_electronic/thermal_electronic.py

@@ -133,7 +133,8 @@ def read_total_electron_dos(
         path: str,
         folder_prefix: str = "elec",
         vasprun_name: str = "vasprun.xml.elec_dos",
-        selected_volumes: np.ndarray = None,
+        selected_volumes: np.ndarray | None = None,
+        selected_folders: list[str] | None = None,
     ) -> None:
         """
         Reads the total electron DOS data from VASP calculations for different volumes.
@@ -143,16 +144,25 @@ def read_total_electron_dos(
                 vasprun.xml files.
             folder_prefix: Prefix of the electronic folders. Defaults to ``"elec"``.
             vasprun_name: Name of the vasprun.xml file. Defaults to ``"vasprun.xml.elec_dos"``.
-            selected_volumes: List of selected volumes to keep the electron DOS data. Defaults to None.
+            selected_volumes: Array of volumes to process. Defaults to None.
+            selected_folders: List of selected folders to keep the electron DOS data. Defaults to None.
 
         Raises:
+            ValueError: If both `selected_volumes` and `selected_folders` are provided, 
+                as they are mutually exclusive.
             ValueError: If selected volumes are not found.
             ValueError: If the number of atoms is not the same for all volumes.
             ValueError: If the number of electrons is not the same for all volumes.
         """
 
         self.path = path
 
+        if selected_volumes is not None and selected_folders is not None:
+            raise ValueError(
+                "selected_volumes and selected_folders are mutually exclusive. "
+                "Provide only one of them."
+            )
+
         with warnings.catch_warnings():
             warnings.simplefilter("ignore")
 
@@ -168,13 +178,20 @@ def read_total_electron_dos(
             energies_list = []
             dos_list = []
 
-            # Iterate over electronic folders to get the relevant electronic DOS data
-            for elec_folder in elec_folders:
+            # Iterate over the requested folders; if none are provided, use all detected elec folders
+            folders_to_process = natsorted(
+                elec_folders if selected_folders is None else selected_folders
+            )
+            for elec_folder in folders_to_process:
 
                 # Get the volumes, number of atoms, and number of electrons from vasprun.xml
                 vasprun_path = os.path.join(path, elec_folder, vasprun_name)
                 vasprun = Vasprun(vasprun_path)
-                volume = round(vasprun.final_structure.volume, 6)
+                volume = round(vasprun.final_structure.volume, 2)
+
+                if selected_volumes is not None and volume not in selected_volumes:
+                    continue
+
                 volumes_list.append(volume)
                 number_of_atoms = vasprun.final_structure.num_sites
                 num_atoms_list.append(number_of_atoms)
@@ -197,6 +214,13 @@ def read_total_electron_dos(
                     vasprun_dos = vasprun.complete_dos.densities[Spin.up]
                 dos_list.append(vasprun_dos)
 
+            # Raise an error if at least one of the selected volumes is not found
+            if selected_volumes is not None:
+                missing = [v for v in selected_volumes if v not in volumes_list]
+                if missing:
+                    raise ValueError(
+                        f"The following selected volumes were not found: {missing}"
+                    )
             # Get the sorted indices based on volumes_list
             sorted_indices = np.argsort(volumes_list)
 
@@ -206,35 +230,7 @@ def read_total_electron_dos(
             nelect_list = [nelect_list[i] for i in sorted_indices]
             energies_list = [energies_list[i] for i in sorted_indices]
             dos_list = [dos_list[i] for i in sorted_indices]
-
-            # Filter values to only include selected volumes
-            if selected_volumes is not None:
-                # Check for missing volumes
-                missing = [v for v in selected_volumes if v not in volumes_list]
-                if missing:
-                    raise ValueError(
-                        f"The following selected volumes were not found: {missing}"
-                    )
-
-                filtered_volume_list = []
-                filtered_num_atoms_list = []
-                filtered_nelect_list = []
-                filtered_vasprun_energies_list = []
-                filtered_vasprun_dos_list = []
-
-                for i in range(len(volumes_list)):
-                    if volumes_list[i] in selected_volumes:
-                        filtered_volume_list.append(volumes_list[i])
-                        filtered_num_atoms_list.append(num_atoms_list[i])
-                        filtered_nelect_list.append(nelect_list[i])
-                        filtered_vasprun_energies_list.append(energies_list[i])
-                        filtered_vasprun_dos_list.append(dos_list[i])
-
-                volumes_list = filtered_volume_list
-                num_atoms_list = filtered_num_atoms_list
-                energies_list = filtered_vasprun_energies_list
-                dos_list = filtered_vasprun_dos_list
-
+
             self.number_of_atoms = np.unique(num_atoms_list)
             # If the number of atoms is not the same for all volumes, raise an error
             if len(self.number_of_atoms) > 1:
@@ -458,7 +454,7 @@ def plot_total_dos(self) -> go.Figure:
 
         return fig
 
-    def plot_vt(self, type: str, selected_temperatures: np.ndarray = None) -> go.Figure:
+    def plot_vt(self, type: str, selected_temperatures: np.ndarray | None = None) -> go.Figure:
         """
         Plots thermal electronic properties as a function of temperature or volume.
 
@@ -603,7 +599,7 @@ def calculate_chemical_potential(
         energies: np.ndarray,
         dos: np.ndarray,
         temperature: float,
-        chemical_potential_range: np.ndarray = None,
+        chemical_potential_range: np.ndarray | None = None,
         electron_tol: float = 0.05,
     ) -> float:
         """
@@ -899,7 +895,7 @@ def calculate_internal_energies(
         internal_energies_list = []
 
         # Calculate the internal energy for each temperature
-        for i, temperature in enumerate(temperatures):
+        for temperature in temperatures:
 
             # Calculate the Fermi-Dirac distribution function at the given temperature and chemical potential
             chemical_potential = self.calculate_chemical_potential(
@@ -1026,7 +1022,7 @@ def calculate_entropies(
         energies_fit_range: np.ndarray = np.array([-2, 2]),
         resolution: float = 0.0001,
         plot: bool = False,
-        plot_temperature: float = None,
+        plot_temperature: float | None = None,
     ) -> np.ndarray | tuple[np.ndarray, go.Figure]:
         """
         Calculates the thermal electronic contribution to the entropy for a given volume.
@@ -1175,7 +1171,7 @@ def calculate_heat_capacities(
         energies_fit_range: np.ndarray = np.array([-2, 2]),
         resolution: float = 0.0001,
         plot: bool = False,
-        plot_temperature: float = None,
+        plot_temperature: float | None = None,
     ) -> np.ndarray | tuple[np.ndarray, go.Figure]:
         """
         Calculates the thermal electronic contribution to the heat capacity for a given volume.
@@ -1215,7 +1211,7 @@ def calculate_heat_capacities(
         heat_capacities_list = []
 
         # Calculate the heat capacity for each temperature
-        for i, temperature in enumerate(temperatures):
+        for temperature in temperatures:
             # Calculate the chemical potential
             chemical_potential = self.calculate_chemical_potential(
                 energies, dos, temperature

dfttk/thermal_electronic/thermal_electronic_data.py

@@ -91,8 +91,10 @@ def __init__(self, path: str):
 
         self.number_of_atoms: int = None
         self.volumes: np.ndarray = None
+        self.energies_list: list = None
+        self.dos_list: list = None
+
         self.temperatures: np.ndarray = None
-
         self.helmholtz_energies: np.ndarray = None
         self.internal_energies: np.ndarray = None
         self.entropies: np.ndarray = None
@@ -131,7 +133,8 @@ def get_vasp_data(
         kpoints_keys: list[str] = ["elec_dos"],
         kpoints_names: list[str] = ["KPOINTS.elec_dos"],
         contcar_name: str = "CONTCAR.elec_dos",
-        selected_volumes: list[float] = None,
+        selected_volumes: np.ndarray | None = None,
+        selected_folders: list[str] | None = None,
         folder_prefix: str = "elec",
     ) -> None:
         """
@@ -144,14 +147,26 @@ def get_vasp_data(
             kpoints_keys: List of KPOINTS keys for dictionary keys. Defaults to ["elec_dos"].
             kpoints_names: List of KPOINTS files to read. Defaults to ["KPOINTS.elec_dos"].
             contcar_name: Name of the CONTCAR file to read. Defaults to "CONTCAR.elec_dos".
-            selected_volumes: List of selected volumes to keep the electron DOS data.
-                Defaults to None.
+            selected_volumes: NumPy array of selected volumes to keep the electron
+                DOS data, or None to keep all volumes. Defaults to None.
+            selected_folders: List of selected folders to keep the electron DOS
+                data. Defaults to None.
             folder_prefix: Prefix of the electronic folders. Defaults to ``"elec"``.
+
+        Raises:
+            ValueError: If both `selected_volumes` and `selected_folders` are provided,
+                as they are mutually exclusive.
         """
 
         # Get the list of electronic DOS folders
         elec_folders = self._get_elec_folders(folder_prefix=folder_prefix)
 
+        if selected_volumes is not None and selected_folders is not None:
+            raise ValueError(
+                "selected_volumes and selected_folders are mutually exclusive. "
+                "Provide only one of them."
+            )
+
         # Filter electronic DOS folders based on selected volumes if provided
         if selected_volumes is not None:
             volumes_set = {round(volume, 2) for volume in selected_volumes}
@@ -168,8 +183,13 @@ def get_vasp_data(
                 in volumes_set
             ]
 
+        # Iterate over the requested folders; if none are provided, use all detected elec folders
+        folders_to_process = natsorted(
+            elec_folders if selected_folders is None else selected_folders
+        )
+
         # Read the INCAR, KPOINTS, and structures for each electronic DOS folder
-        for elec_folder in elec_folders:
+        for elec_folder in folders_to_process:
             incar_data = {}
             for key, name in zip(incar_keys, incar_names):
                 incar_data[key] = Incar.from_file(
@@ -202,7 +222,8 @@ def get_thermal_electronic_data(
         order: int = 1,
         folder_prefix: str = "elec",
         vasprun_name: str = "vasprun.xml.elec_dos",
-        selected_volumes: np.ndarray = None,
+        selected_volumes: np.ndarray | None = None,
+        selected_folders: list[str] | None = None,
     ):
         """
         Calls the ThermalElectronic class to read the total electron DOS data, compute
@@ -214,7 +235,8 @@ def get_thermal_electronic_data(
             order: Order of the polynomial fit. Defaults to 1 (linear fit).
             folder_prefix: Prefix of the electronic folders. Defaults to ``"elec"``.
             vasprun_name: Name of the vasprun.xml file. Defaults to ``"vasprun.xml.elec_dos"``.
-            selected_volumes: List of selected volumes to keep the electron DOS data. Defaults to None.
+            selected_volumes: Array of volumes to process. Defaults to None.
+            selected_folders: List of selected folders to keep the electron DOS data. Defaults to None.
         """
 
         # Initialize ThermalElectronic object
@@ -226,6 +248,7 @@ def get_thermal_electronic_data(
             folder_prefix=folder_prefix,
             vasprun_name=vasprun_name,
             selected_volumes=selected_volumes,
+            selected_folders=selected_folders,
         )
         self.number_of_atoms = self.te.number_of_atoms
         self.volumes = self.te.volumes

tests/test_thermal_electronic.py

@@ -40,15 +40,13 @@ def cached_dos_data():
 
 def test_read_total_electron_dos():
     # Test all volumes
-    # This test is specifically for the read_total_electron_dos method, so keep the direct call
     thermal_electronic = ThermalElectronic()
     thermal_electronic.read_total_electron_dos(path=config_Al_path)
 
     assert thermal_electronic.path == config_Al_path
     assert thermal_electronic.number_of_atoms == 4
     assert np.all(thermal_electronic.volumes == np.array([60.0, 62.0, 64.0, 66.0, 68.0, 70.0, 72.0, 74.0]))
 
-    # Check that the energies_list and dos_list match expected values
     for i, (energies_list, dos_list) in enumerate(zip(thermal_electronic.energies_list, thermal_electronic.dos_list)):
         np.testing.assert_allclose(energies_list, expected_energies_list[i])
         np.testing.assert_allclose(dos_list, expected_dos_list[i])
@@ -61,7 +59,6 @@ def test_read_total_electron_dos():
     assert thermal_electronic.number_of_atoms == 4
     assert np.all(thermal_electronic.volumes == np.array([60.0, 66.0]))
 
-    # Check that the energies_list and dos_list match expected values
     for i, (energies_list, dos_list) in enumerate(zip(thermal_electronic.energies_list, thermal_electronic.dos_list)):
         np.testing.assert_allclose(energies_list, expected_energies_list[i * 3])  # 0 and 3 indices
         np.testing.assert_allclose(dos_list, expected_dos_list[i * 3])  # 0 and 3 indices
@@ -71,9 +68,25 @@ def test_read_total_electron_dos():
     with pytest.raises(ValueError, match="The following selected volumes were not found:"):
         thermal_electronic.read_total_electron_dos(path=config_Al_path, selected_volumes=invalid_volumes)
 
+    # Test selected folders
+    selected_folders = ["elec_4", "elec_7"]
+    thermal_electronic.read_total_electron_dos(path=config_Al_path, selected_folders=selected_folders)
+
+    assert thermal_electronic.path == config_Al_path
+    assert thermal_electronic.number_of_atoms == 4
+    assert np.all(thermal_electronic.volumes == np.array([60.0, 66.0]))
+
+    for i, (energies_list, dos_list) in enumerate(zip(thermal_electronic.energies_list, thermal_electronic.dos_list)):
+        np.testing.assert_allclose(energies_list, expected_energies_list[i * 3])  # 0 and 3 indices
+        np.testing.assert_allclose(dos_list, expected_dos_list[i * 3])  # 0 and 3 indices
+
+    # If selected_volumes and selected_folders are both provided, raise an error
+    with pytest.raises(ValueError, match="selected_volumes and selected_folders are mutually exclusive. Provide only one of them."):
+        thermal_electronic.read_total_electron_dos(path=config_Al_path, selected_volumes=selected_volumes, selected_folders=selected_folders)
+
     # TODO: Test inconsistent number of atoms
     # TODO: Test inconsistent number of electrons
-
+    
 
 def test_set_total_electron_dos(cached_dos_data):
     thermal_electronic = ThermalElectronic()