From a1d685c6593f73b0441c89507cdda548b2aebc82 Mon Sep 17 00:00:00 2001 From: ahxbcn Date: Sun, 11 Jan 2026 00:03:19 +0800 Subject: [PATCH 01/15] Simplify arguments of abacus_dos_run tool --- src/abacusagent/modules/dos.py | 16 ++-- src/abacusagent/modules/submodules/dos.py | 91 ++++++++--------------- src/abacusagent/modules/util/comm.py | 43 +++++++---- 3 files changed, 63 insertions(+), 87 deletions(-) diff --git a/src/abacusagent/modules/dos.py b/src/abacusagent/modules/dos.py index 0555cf5..6f85da7 100644 --- a/src/abacusagent/modules/dos.py +++ b/src/abacusagent/modules/dos.py @@ -3,16 +3,15 @@ from abacusagent.init_mcp import mcp from abacusagent.modules.submodules.dos import abacus_dos_run as _abacus_dos_run + @mcp.tool() def abacus_dos_run( abacus_inputs_dir: Path, pdos_mode: Literal['species', 'species+shell', 'species+orbital'] = 'species+shell', dos_edelta_ev: float = 0.01, dos_sigma: float = 0.07, - dos_scale: float = 0.01, - dos_emin_ev: float = None, - dos_emax_ev: float = None, - dos_nche: int = None, + dos_emin_ev: float = -10.0, + dos_emax_ev: float = 10.0, ) -> Dict[str, Any]: """Run the DOS and PDOS calculation. @@ -29,11 +28,8 @@ def abacus_dos_run( - "species+orbital": Orbital-resolved PDOS will be plotted. PDOS of orbitals in the same shell of a species will be plotted in a subplot. dos_edelta_ev: Step size in writing Density of States (DOS) in eV. dos_sigma: Width of the Gaussian factor when obtaining smeared Density of States (DOS) in eV. - dos_scale: Defines the energy range of DOS output as (emax-emin)*(1+dos_scale), centered at (emax+emin)/2. - This parameter will be used when dos_emin_ev and dos_emax_ev are not set. - dos_emin_ev: Minimal range for Density of States (DOS) in eV. - dos_emax_ev: Maximal range for Density of States (DOS) in eV. - dos_nche: The order of Chebyshev expansions when using Stochastic Density Functional Theory (SDFT) to calculate DOS. + dos_emin_ev: Minimal range for Density of States (DOS) in eV. Default is -10.0. + dos_emax_ev: Maximal range for Density of States (DOS) in eV. Default is 10.0. Returns: Dict[str, Any]: A dictionary containing: @@ -47,4 +43,4 @@ def abacus_dos_run( - nscf_work_path: Path to the work directory of NSCF calculation. - nscf_normal_end: If the SCF calculation ended normally. """ - return _abacus_dos_run(abacus_inputs_dir, pdos_mode, dos_edelta_ev, dos_sigma, dos_scale, dos_emin_ev, dos_emax_ev, dos_nche) + return _abacus_dos_run(abacus_inputs_dir, pdos_mode, dos_edelta_ev, dos_sigma, dos_emin_ev, dos_emax_ev) diff --git a/src/abacusagent/modules/submodules/dos.py b/src/abacusagent/modules/submodules/dos.py index 13c197e..52e6df5 100644 --- a/src/abacusagent/modules/submodules/dos.py +++ b/src/abacusagent/modules/submodules/dos.py @@ -46,10 +46,8 @@ def abacus_dos_run( pdos_mode: Literal['species', 'species+shell', 'species+orbital'] = 'species+shell', dos_edelta_ev: float = 0.01, dos_sigma: float = 0.07, - dos_scale: float = 0.01, - dos_emin_ev: float = None, - dos_emax_ev: float = None, - dos_nche: int = None, + dos_emin_ev: float = -10.0, + dos_emax_ev: float = 10.0, ) -> Dict[str, Any]: """Run the DOS and PDOS calculation. @@ -63,14 +61,11 @@ def abacus_dos_run( pdos_mode: Mode of plotted PDOS file. - "species": Total PDOS of any species will be plotted in a picture. - "species+shell": PDOS for any shell (s, p, d, f, g,...) of any species will be plotted. PDOS of a shell of a species willbe plotted in a subplot. - - “species+orbital": Orbital-resolved PDOS will be plotted. PDOS of orbitals in the same shell of a species will be plotted in a subplot. + - "species+orbital": Orbital-resolved PDOS will be plotted. PDOS of orbitals in the same shell of a species will be plotted in a subplot. dos_edelta_ev: Step size in writing Density of States (DOS) in eV. dos_sigma: Width of the Gaussian factor when obtaining smeared Density of States (DOS) in eV. - dos_scale: Defines the energy range of DOS output as (emax-emin)*(1+dos_scale), centered at (emax+emin)/2. - This parameter will be used when dos_emin_ev and dos_emax_ev are not set. - dos_emin_ev: Minimal range for Density of States (DOS) in eV. - dos_emax_ev: Maximal range for Density of States (DOS) in eV. - dos_nche: The order of Chebyshev expansions when using Stochastic Density Functional Theory (SDFT) to calculate DOS. + dos_emin_ev: Minimal range for Density of States (DOS) in eV. Default is -10.0. + dos_emax_ev: Maximal range for Density of States (DOS) in eV. Default is 10.0. Returns: Dict[str, Any]: A dictionary containing: @@ -98,17 +93,15 @@ def abacus_dos_run( metrics_scf = abacus_dos_run_scf(abacus_inputs_dir) metrics_nscf = abacus_dos_run_nscf(metrics_scf["scf_work_path"], dos_edelta_ev=dos_edelta_ev, - dos_sigma=dos_sigma, - dos_scale=dos_scale, - dos_emin_ev=dos_emin_ev, - dos_emax_ev=dos_emax_ev, - dos_nche=dos_nche) + dos_sigma=dos_sigma) fig_paths = plot_dos_pdos(metrics_scf["scf_work_path"], metrics_nscf["nscf_work_path"], metrics_nscf["nscf_work_path"], nspin, - pdos_mode) + pdos_mode, + dos_emin_ev, + dos_emax_ev) return_dict = {"dos_fig_path": fig_paths[0]} try: @@ -167,16 +160,13 @@ def abacus_dos_run_scf(abacus_inputs_dir: Path, def abacus_dos_run_nscf(abacus_inputs_dir: Path, dos_edelta_ev: float = None, - dos_sigma: float = None, - dos_scale: float = None, - dos_emin_ev: float = None, - dos_emax_ev: float = None, - dos_nche: int = None,) -> Dict[str, Any]: + dos_sigma: float = None) -> Dict[str, Any]: work_path = generate_work_path() link_abacusjob(src=abacus_inputs_dir, dst=work_path, - copy_files=["INPUT"]) + copy_files=["INPUT", "KPT"], + exclude=["*log", "*json"]) input_param = ReadInput(os.path.join(work_path, "INPUT")) input_param["calculation"] = "nscf" @@ -189,10 +179,6 @@ def abacus_dos_run_nscf(abacus_inputs_dir: Path, for dos_param, value in { "dos_edelta_ev": dos_edelta_ev, "dos_sigma": dos_sigma, - "dos_scale": dos_scale, - "dos_emin_ev": dos_emin_ev, - "dos_emax_ev": dos_emax_ev, - "dos_nche": dos_nche }.items(): if value is not None: input_param[dos_param] = value @@ -244,23 +230,6 @@ def parse_pdos_file(file_path): return energy_values, orbitals -def parse_log_file(file_path): - """Parse Fermi energy from log file and convert to eV.""" - ry_to_ev = 13.605698066 - fermi_energy = None - - with open(file_path, 'r') as f: - for line in f: - if "Fermi energy is" in line: - match = re.search(r'Fermi energy is\s*([\d.-]+)', line) - if match: - fermi_energy = float(match.group(1)) - - if fermi_energy is None: - raise ValueError("Fermi energy not found in log file") - - return fermi_energy * ry_to_ev - def parse_basref_file(file_path): """Parse basref file to create mapping for custom labels.""" label_map = {} @@ -286,7 +255,7 @@ def parse_basref_file(file_path): return label_map -def plot_pdos(energy_values, orbitals, fermi_level, label_map, output_dir, nspin, mode, dpi=300): +def plot_pdos(energy_values, orbitals, fermi_level, label_map, output_dir, nspin, mode, dos_emin_ev, dos_emax_ev, dpi=300): """Plot PDOS data separated by atom/species with custom labels.""" # Create output directory if it doesn't exist os.makedirs(output_dir, exist_ok=True) @@ -303,17 +272,17 @@ def plot_pdos(energy_values, orbitals, fermi_level, label_map, output_dir, nspin atom_species_groups[key].append(orbital) if mode == "species": - pdos_pic_file = plot_pdos_species(shifted_energy, orbitals, output_dir, nspin, dpi) + pdos_pic_file = plot_pdos_species(shifted_energy, orbitals, output_dir, nspin, dos_emin_ev, dos_emax_ev, dpi) elif mode == "species+shell": - pdos_pic_file = plot_pdos_species_shell(shifted_energy, orbitals, output_dir, nspin, dpi) + pdos_pic_file = plot_pdos_species_shell(shifted_energy, orbitals, output_dir, nspin, dos_emin_ev, dos_emax_ev, dpi) elif mode == "species+orbital": - pdos_pic_file = plot_pdos_species_orbital(shifted_energy, orbitals, output_dir, nspin, label_map, dpi) + pdos_pic_file = plot_pdos_species_orbital(shifted_energy, orbitals, output_dir, nspin, label_map, dos_emin_ev, dos_emax_ev, dpi) else: raise ValueError(f"Not allowed mode {mode}") return pdos_pic_file -def plot_pdos_species(shifted_energy, orbitals, output_dir, nspin, dpi): +def plot_pdos_species(shifted_energy, orbitals, output_dir, nspin, dos_emin_ev, dos_emax_ev, dpi): species = {} for orbital in orbitals: species_one = orbital['species'] @@ -322,7 +291,6 @@ def plot_pdos_species(shifted_energy, orbitals, output_dir, nspin, dpi): else: species[species_one] += orbital['data'] - num_species = len(species) plt.plot(figsize=(10, 6)) for species_name, pdos_data in species.items(): if nspin == 1: @@ -334,7 +302,7 @@ def plot_pdos_species(shifted_energy, orbitals, output_dir, nspin, dpi): plt.axvline(x=0, color='black', linestyle=':', linewidth=1.0) plt.xlabel('Energy (eV)', fontsize=10) plt.ylabel(r"States ($eV^{-1}$)", fontsize=10) - plt.xlim(max(min(shifted_energy), -20), min(20, max(shifted_energy))) + plt.xlim(dos_emin_ev, dos_emax_ev) if nspin == 1: plt.ylim(bottom=0) plt.legend(fontsize=8, ncol=nspin) @@ -347,7 +315,7 @@ def plot_pdos_species(shifted_energy, orbitals, output_dir, nspin, dpi): return Path(pdos_pic_file).absolute() -def plot_pdos_species_shell(shifted_energy, orbitals, output_dir, nspin, dpi): +def plot_pdos_species_shell(shifted_energy, orbitals, output_dir, nspin, dos_emin_ev, dos_emax_ev, dpi): species_shells = {} for orbital in orbitals: species = orbital['species'] @@ -381,7 +349,7 @@ def plot_pdos_species_shell(shifted_energy, orbitals, output_dir, nspin, dpi): ax.axvline(x=0, color='black', linestyle=':', linewidth=1.0) ax.set_title(f'PDOS for {species}', fontsize=12, pad=10) ax.set_ylabel(r"States ($eV^{-1}$)", fontsize=10) - ax.set_xlim(max(min(shifted_energy), -20), min(20, max(shifted_energy))) + ax.set_xlim(dos_emin_ev, dos_emax_ev) #if nspin == 1: # ax.set_ylim(bottom=0) ax.legend(fontsize=8, ncol=nspin) @@ -398,7 +366,7 @@ def plot_pdos_species_shell(shifted_energy, orbitals, output_dir, nspin, dpi): return Path(pdos_pic_file).absolute() -def plot_pdos_species_orbital(shifted_energy, orbitals, output_dir, nspin, label_map, dpi): +def plot_pdos_species_orbital(shifted_energy, orbitals, output_dir, nspin, label_map, dos_emin_ev, dos_emax_ev, dpi): plt.rcParams["text.usetex"] = False plt.rcParams["axes.prop_cycle"] = plt.cycler("color", plt.cm.tab20.colors) @@ -455,7 +423,7 @@ def plot_pdos_species_orbital(shifted_energy, orbitals, output_dir, nspin, label ax.axvline(x=0, color='black', linestyle=':', linewidth=1.0) ax.set_title(f'PDOS for {species}-{angular_momentum}', fontsize=12, pad=10) - ax.set_xlim(max(min(shifted_energy), -20), min(20, max(shifted_energy))) + ax.set_xlim(dos_emin_ev, dos_emax_ev) if nspin == 1: ax.set_ylim(bottom=0) ax.set_ylabel(r"States ($eV^{-1}$)", fontsize=10) @@ -477,6 +445,8 @@ def plot_dos(file_path: List[Path], fermi_level: float, output_file: str = 'DOS.png', nspin: Literal[1, 2] = 1, + dos_emin_ev: float = -10.0, + dos_emax_ev: float = 10.0, dpi: int=300): """Plot total DOS from DOS1_smearing.dat and DOS2_smearing (if nspin=2) file.""" # Read first two columns from file @@ -486,9 +456,6 @@ def plot_dos(file_path: List[Path], if nspin == 2: data = np.loadtxt(file_path[1], usecols=(0, 1)) dos_dn = data[:, 1] - - # Determine energy limits based on data within x range - x_min, x_max = max(min(energy), -20), min(20, max(energy)) # Create plot plt.figure(figsize=(8, 6)) @@ -497,12 +464,12 @@ def plot_dos(file_path: List[Path], elif nspin == 2: plt.plot(energy, dos, linestyle='-', label='spin up') plt.plot(energy, -dos_dn, linestyle='--', label='spin down') - plt.axvline(x=0, color='k', linestyle='--', alpha=0.5) + plt.axvline(x=0, color='k', linestyle=':', alpha=0.5) plt.xlabel('Energy (eV)') plt.ylabel(r'States ($eV^{-1}$)') plt.title('Density of States') plt.grid(True, alpha=0.3) - plt.xlim(x_min, x_max) + plt.xlim(dos_emin_ev, dos_emax_ev) #plt.ylim(y_min, y_max) #plt.legend() @@ -518,6 +485,8 @@ def plot_dos_pdos(scf_job_path: Path, output_dir: Path, nspin: Literal[1, 2] = 1, mode: Literal['species', 'species+shell', 'species+orbital'] = 'species+shell', + dos_emin_ev: float = -10.0, + dos_emax_ev: float = 10.0, dpi=300) -> List[str]: """Plot DOS and PDOS from the NSCF job path. @@ -557,7 +526,7 @@ def plot_dos_pdos(scf_job_path: Path, fermi_level = collect_metrics(scf_job_path, ['efermi'])['efermi'] # Plot DOS and get file path - dos_plot_file = plot_dos(dos_file, fermi_level, dos_output, nspin, dpi) + dos_plot_file = plot_dos(dos_file, fermi_level, dos_output, nspin, dos_emin_ev, dos_emax_ev, dpi) all_plot_files = [dos_plot_file] print("DOS file plotted") @@ -567,7 +536,7 @@ def plot_dos_pdos(scf_job_path: Path, if basis_type != 'pw': label_map = parse_basref_file(basref_file) energy_values, orbitals = parse_pdos_file(pdos_file) - pdos_plot_file = plot_pdos(energy_values, orbitals, fermi_level, label_map, output_dir, nspin, mode, dpi) + pdos_plot_file = plot_pdos(energy_values, orbitals, fermi_level, label_map, output_dir, nspin, mode, dos_emin_ev, dos_emax_ev, dpi) # Combine file paths into a single list all_plot_files.append(pdos_plot_file) diff --git a/src/abacusagent/modules/util/comm.py b/src/abacusagent/modules/util/comm.py index 5645f7b..d52e7cf 100644 --- a/src/abacusagent/modules/util/comm.py +++ b/src/abacusagent/modules/util/comm.py @@ -7,6 +7,7 @@ import json import traceback import uuid +import shutil import glob from abacustest.lib_prepare.abacus import ReadInput @@ -233,7 +234,7 @@ def link_abacusjob(src: str, exclude (Optional[List[str]]): List of files to exclude. If None, no files are excluded. copy_files (List[str]): List of files to copy from src to dst. Default is ["INPUT", "STRU", "KPT"]. overwrite (bool): If True, existing files in the destination will be overwritten. Default is True. - exclude_directories (bool): If True, directories will be excluded from linking. Default is False. + exclude_directories (bool): If True, directories will be not copied. Default is False. Notes: - If somes files are included in both include and exclude, the file will be excluded. @@ -250,9 +251,9 @@ def link_abacusjob(src: str, if include is None: include = ["*"] - include_files = [] + include_paths = [] for pattern in include: - include_files.extend(src.glob(pattern)) + include_paths.extend(src.glob(pattern)) if exclude is None: exclude = [] @@ -262,30 +263,40 @@ def link_abacusjob(src: str, os.makedirs(dst, exist_ok=True) # Remove excluded files from included files - include_files = [f for f in include_files if f not in exclude_files] - if not include_files: + include_paths = [f for f in include_paths if f not in exclude_files] + if not include_paths: traceback.print_stack() print("No files to link after applying include and exclude patterns.\n", f"Include patterns: {include}, Exclude patterns: {exclude}, Source: {src}, Destination: {dst}\n", f"Files in source: {list(src.glob('*'))}" ) else: - for file in include_files: - if file == dst: + for path in include_paths: + if path == dst: continue - if exclude_directories and os.path.isdir(file): + if exclude_directories and os.path.isdir(path): continue - dst_file = dst / file.name - if dst_file.exists(): - if overwrite: - dst_file.unlink() + if os.path.isfile(path): + dst_file = dst / path.name + if dst_file.exists(): + if overwrite: + dst_file.unlink() + else: + continue + if str(path.name) in copy_files: + os.system(f"cp {path} {dst_file}") else: - continue - if str(file.name) in copy_files: - os.system(f"cp {file} {dst_file}") + os.symlink(path, dst_file) else: - os.symlink(file, dst_file) + dst_dir = dst / path.name + if dst_dir.exists(): + if overwrite: + shutil.rmtree(dst_dir) + shutil.copytree(path, dst_dir) + else: + shutil.copytree(path, dst_dir) + def generate_work_path(create: bool = True) -> str: """ From 8904959e0c7ceefc76e8ed4c9e9d92aa9b285d90 Mon Sep 17 00:00:00 2001 From: ahxbcn Date: Sun, 11 Jan 2026 10:30:21 +0800 Subject: [PATCH 02/15] Update integrate test for DOS --- tests/integrate_test/test_dos.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/integrate_test/test_dos.py b/tests/integrate_test/test_dos.py index 3731f08..70cbe65 100644 --- a/tests/integrate_test/test_dos.py +++ b/tests/integrate_test/test_dos.py @@ -46,9 +46,8 @@ def test_abacus_dos_run_species(self): pdos_mode='species', dos_edelta_ev = 0.01, dos_sigma = 0.07, - dos_scale = 0.01, - dos_emin_ev = -20, - dos_emax_ev = 20) + dos_emin_ev=-20, + dos_emax_ev=20) dos_fig_path = outputs['dos_fig_path'] pdos_fig_path = outputs['pdos_fig_path'] @@ -125,7 +124,6 @@ def test_abacus_dos_run_species_nspin2(self): pdos_mode='species', dos_edelta_ev = 0.01, dos_sigma = 0.07, - dos_scale = 0.01, dos_emin_ev = -20, dos_emax_ev = 20) From ef27229555427bd4fa22c0d6aeb8125e0cb96da9 Mon Sep 17 00:00:00 2001 From: ahxbcn Date: Sun, 11 Jan 2026 13:56:59 +0800 Subject: [PATCH 03/15] Dump processed DOS and PDOS to file --- src/abacusagent/modules/submodules/dos.py | 159 ++++++++++++++++++++-- 1 file changed, 146 insertions(+), 13 deletions(-) diff --git a/src/abacusagent/modules/submodules/dos.py b/src/abacusagent/modules/submodules/dos.py index 52e6df5..109074a 100644 --- a/src/abacusagent/modules/submodules/dos.py +++ b/src/abacusagent/modules/submodules/dos.py @@ -1,13 +1,14 @@ import os import re import numpy as np +from numpy.typing import NDArray import matplotlib.pyplot as plt from abacustest.lib_prepare.abacus import ReadInput, WriteInput from abacustest.lib_collectdata.collectdata import RESULT from abacustest.lib_model.comm import check_abacus_inputs from pathlib import Path -from typing import Dict, Any, List, Literal +from typing import Dict, Any, List, Literal, Optional from abacusagent.modules.util.comm import generate_work_path, link_abacusjob, run_abacus, has_chgfile, collect_metrics from abacusagent.modules.util.chemical_elements import MAX_ANGULAR_MOMENTUM_OF_ELEMENTS @@ -255,6 +256,109 @@ def parse_basref_file(file_path): return label_map +def write_pdos_data_species(energy: NDArray[np.float64], + species_pdos_data: Dict[str, NDArray[np.float64]], + nspin: Literal[1, 2] = 1, + output_dir: str = './') -> None: + """Write PDOS of species to file""" + with open(os.path.join(output_dir, "PDOS_species.dat"), 'w') as f: + if nspin == 1: + header_line = ' Energy (eV)' + for species_name in species_pdos_data.keys(): + header_line += f'{species_name:>20s}' + f.write(header_line + "\n") + for i in range(len(energy)): + line = f"{energy[i]:20.8f}" + for species_name in species_pdos_data.keys(): + line += f"{species_pdos_data[species_name][i]:20.8f}" + f.write(line + "\n") + elif nspin == 2: + header_spin_up, header_spin_down = '', '' + for species_name in species_pdos_data.keys(): + header_spin_up += f'{species_name:>17s}_up' + header_spin_down += f'{species_name:>17s}_dn' + f.write(f" Energy (eV){header_spin_up}{header_spin_down}\n") + for i in range(len(energy)): + line = f"{energy[i]:20.8f}" + line_spin_up, line_spin_down = '', '' + for species_name, pdos_data in species_pdos_data.items(): + # even index for spin up, odd index for spin down + line_spin_up += f"{pdos_data[i*2]:20.8f}" + line_spin_down += f"{pdos_data[i*2+1]:20.8f}" + line += line_spin_up + line_spin_down + f.write(line + "\n") + +def write_pdos_data_species_shell(energy: NDArray[np.float64], + species_pdos_data_dict: Dict[str, NDArray[np.float64]], + species: str, + nspin: Literal[1, 2] = 1, + output_dir: str = './') -> None: + # Write PDOS of different shells (s, p, d...) for a species to file + with open(os.path.join(output_dir, f"PDOS_{species}.dat"), "w") as f: + if nspin == 1: + header_line = ' Energy (eV)' + for l in species_pdos_data_dict.keys(): + header_line += f'{l:>20s}' + f.write(header_line + "\n") + for i in range(len(energy)): + line = f"{energy[i]:14.6f}" + for l, pdos_data in species_pdos_data_dict.items(): + line += f"{pdos_data[i]:20.8f}" + f.write(line + "\n") + elif nspin == 2: + header_spin_up, header_spin_down = '', '' + for l in species_pdos_data_dict.keys(): + header_spin_up += f'{l:>17s}_up' + header_spin_down += f'{l:>17}_dn' + f.write(f" Energy (eV){header_spin_up}{header_spin_down}\n") + for i in range(len(energy)): + line = f"{energy[i]:14.6f}" + line_spin_up, line_spin_down = '', '' + for l, pdos_data in species_pdos_data_dict.items(): + # even index for spin up, odd index for spin down + line_spin_up += f"{pdos_data[i*2]:20.8f}" + line_spin_down += f"{pdos_data[i*2+1]:20.8f}" + f.write(line + line_spin_up + line_spin_down + "\n") + +def write_pdos_data_species_orbital(energy: NDArray[np.float64], + species_pdos: Dict[str, NDArray[np.float64]], + species: str, + orbital_label: Dict[str, Dict[str, Dict[str, str]]], + nspin: Literal[1, 2] = 1, + output_dir: str = "./") -> None: + # Write PDOS data of different orbitals of a species to file + with open(os.path.join(output_dir, f"PDOS_{species}.dat"), "w") as f: + if nspin == 1: + header_line = f" Energy (eV)" + for l, species_shell_pdos in species_pdos.items(): + for m in species_shell_pdos.keys(): + orbital_name = orbital_label[species][str(angular_momentum_map.index(l))][str(m)] + header_line += f"{orbital_name.replace('$', ''):>16s}" + f.write(header_line+"\n") + for i in range(len(energy)): + line = f"{energy[i]:14.6f}" + for l, species_shell_pdos in species_pdos.items(): + for m in species_shell_pdos.keys(): + line += f"{species_pdos[l][m][i]:20.8f}" + f.write(line + "\n") + elif nspin == 2: + header_spin_up, header_spin_down = '', '' + for l, species_shell_pdos in species_pdos.items(): + for m in species_shell_pdos.keys(): + orbital_name = orbital_label[species][str(angular_momentum_map.index(l))][str(m)] + header_spin_up += f"{orbital_name.replace('$', ''):>17s}_up" + header_spin_down += f"{orbital_name.replace('$', ''):>17s}_dn" + f.write(f" Energy (eV){header_spin_up}{header_spin_down}\n") + for i in range(len(energy)): + line = f"{energy[i]:14.6f}" + line_spin_up, line_spin_down = '', '' + for l, species_shell_pdos in species_pdos.items(): + for m, species_orbital_pdos in species_shell_pdos.items(): + # even index for spin up, odd index for spin down + line_spin_up += f"{species_orbital_pdos[i*2]:20.8f}" + line_spin_down += f"{species_orbital_pdos[i*2+1]:20.8f}" + f.write(line + line_spin_up + line_spin_down + "\n") + def plot_pdos(energy_values, orbitals, fermi_level, label_map, output_dir, nspin, mode, dos_emin_ev, dos_emax_ev, dpi=300): """Plot PDOS data separated by atom/species with custom labels.""" # Create output directory if it doesn't exist @@ -291,6 +395,8 @@ def plot_pdos_species(shifted_energy, orbitals, output_dir, nspin, dos_emin_ev, else: species[species_one] += orbital['data'] + write_pdos_data_species(shifted_energy, species, nspin, output_dir) + # Plot PDOS of different species plt.plot(figsize=(10, 6)) for species_name, pdos_data in species.items(): if nspin == 1: @@ -337,6 +443,9 @@ def plot_pdos_species_shell(shifted_energy, orbitals, output_dir, nspin, dos_emi axes = [axes] for species_idx, (species, pdos_data_dict) in enumerate(species_shells.items()): + # Write PDOS data to the selected species to file + write_pdos_data_species_shell(shifted_energy, pdos_data_dict, species, nspin, output_dir) + # Plot PDOS of different shells for a species ax = axes[species_idx] for l, pdos_data in pdos_data_dict.items(): @@ -350,12 +459,10 @@ def plot_pdos_species_shell(shifted_energy, orbitals, output_dir, nspin, dos_emi ax.set_title(f'PDOS for {species}', fontsize=12, pad=10) ax.set_ylabel(r"States ($eV^{-1}$)", fontsize=10) ax.set_xlim(dos_emin_ev, dos_emax_ev) - #if nspin == 1: - # ax.set_ylim(bottom=0) + if nspin == 1: + ax.set_ylim(bottom=0) ax.legend(fontsize=8, ncol=nspin) ax.grid(alpha=0.3) - - #ax.set_ylim(bottom=0) axes[-1].set_xlabel('Energy (eV)', fontsize=10) @@ -371,6 +478,7 @@ def plot_pdos_species_orbital(shifted_energy, orbitals, output_dir, nspin, label plt.rcParams["text.usetex"] = False plt.rcParams["axes.prop_cycle"] = plt.cycler("color", plt.cm.tab20.colors) + # Initialize dict of pdos data for each orbital orbital_label = {} for (atom_index, species, l, m, z), full_label in label_map.items(): if species not in orbital_label.keys(): @@ -385,7 +493,8 @@ def plot_pdos_species_orbital(shifted_energy, orbitals, output_dir, nspin, label orbital_label[species][str(l)][str(m)] = orbital_name else: pass - + + # Sum over zeta to get the pdos of orbitals species_orbitals = {} for orbital in orbitals: species = orbital['species'] @@ -404,6 +513,7 @@ def plot_pdos_species_orbital(shifted_energy, orbitals, output_dir, nspin, label else: species_orbitals[species][angular_momentum][mag_quantum_num] += orbital['data'] + # Plot PDOS for each species and each orbital total_subplots = 0 for species, species_pdos in species_orbitals.items(): total_subplots += len(species_pdos) @@ -411,6 +521,8 @@ def plot_pdos_species_orbital(shifted_energy, orbitals, output_dir, nspin, label subplot_count = 0 for species, species_pdos in species_orbitals.items(): + # Write PDOS data of orbitals of the selected species to file + write_pdos_data_species_orbital(shifted_energy, species_pdos, species, orbital_label, nspin, output_dir) for angular_momentum, species_shell_pdos in species_pdos.items(): for m, species_orbital_pdos in species_shell_pdos.items(): ax = axes[subplot_count] @@ -441,9 +553,26 @@ def plot_pdos_species_orbital(shifted_energy, orbitals, output_dir, nspin, label return Path(pdos_pic_file).absolute() +def write_dos_data(energy: NDArray[np.float64], + dos: NDArray[np.float64], + dos_dn: Optional[NDArray[np.float64]] = None, + nspin: Literal[1, 2] = 1, + out_filename: str = 'DOS_shifted.dat') -> None: + """Write DOS data to file.""" + with open(out_filename, "w") as f: + if nspin == 1: + f.write(" Energy (eV) DOS\n") + for i in range(len(energy)): + f.write(f"{energy[i]:16.6f}{dos[i]:16.6f}\n") + elif nspin == 2: + f.write(" Energy (eV) DOS_up DOS_down\n") + for i in range(len(energy)): + f.write(f"{energy[i]:16.6f}{dos[i]:16.6f}{dos_dn[i]:16.6f}\n") + f.close() + def plot_dos(file_path: List[Path], fermi_level: float, - output_file: str = 'DOS.png', + dos_plot_file: str = 'DOS.png', nspin: Literal[1, 2] = 1, dos_emin_ev: float = -10.0, dos_emax_ev: float = 10.0, @@ -456,7 +585,13 @@ def plot_dos(file_path: List[Path], if nspin == 2: data = np.loadtxt(file_path[1], usecols=(0, 1)) dos_dn = data[:, 1] - + + # Write shifted dos data to file + dos_output_file = os.path.join(os.path.dirname(dos_plot_file), "DOS_shifted.dat") + if nspin == 1: + write_dos_data(energy, dos, nspin=nspin, out_filename=dos_output_file) + elif nspin == 2: + write_dos_data(energy, dos, dos_dn, nspin=nspin, out_filename=dos_output_file) # Create plot plt.figure(figsize=(8, 6)) if nspin == 1: @@ -470,15 +605,13 @@ def plot_dos(file_path: List[Path], plt.title('Density of States') plt.grid(True, alpha=0.3) plt.xlim(dos_emin_ev, dos_emax_ev) - #plt.ylim(y_min, y_max) - #plt.legend() # Save plot - os.makedirs(os.path.dirname(output_file), exist_ok=True) - plt.savefig(output_file, dpi=dpi, bbox_inches='tight') + os.makedirs(os.path.dirname(dos_plot_file), exist_ok=True) + plt.savefig(dos_plot_file, dpi=dpi, bbox_inches='tight') plt.close() - return Path(output_file).absolute() + return Path(dos_plot_file).absolute() def plot_dos_pdos(scf_job_path: Path, nscf_job_path: Path, From e488b00a9a6a15a16b5d0c302d1b477b70bb8bc6 Mon Sep 17 00:00:00 2001 From: ahxbcn Date: Tue, 13 Jan 2026 08:34:53 +0800 Subject: [PATCH 04/15] Add a new mode "atom" in abacus_dos_run to plot PDOS of selected atoms --- src/abacusagent/modules/dos.py | 22 ++-- src/abacusagent/modules/submodules/dos.py | 137 +++++++++++++++++++++- 2 files changed, 140 insertions(+), 19 deletions(-) diff --git a/src/abacusagent/modules/dos.py b/src/abacusagent/modules/dos.py index 6f85da7..2362093 100644 --- a/src/abacusagent/modules/dos.py +++ b/src/abacusagent/modules/dos.py @@ -1,5 +1,5 @@ from pathlib import Path -from typing import Dict, Any, List, Literal +from typing import Dict, Any, List, Literal, Optional from abacusagent.init_mcp import mcp from abacusagent.modules.submodules.dos import abacus_dos_run as _abacus_dos_run @@ -7,7 +7,8 @@ @mcp.tool() def abacus_dos_run( abacus_inputs_dir: Path, - pdos_mode: Literal['species', 'species+shell', 'species+orbital'] = 'species+shell', + pdos_mode: Literal['atoms', 'species', 'species+shell', 'species+orbital'] = 'species+shell', + pdos_atom_indices: Optional[List[int]] = None, dos_edelta_ev: float = 0.01, dos_sigma: float = 0.07, dos_emin_ev: float = -10.0, @@ -23,24 +24,15 @@ def abacus_dos_run( Args: abacus_inputs_dir: Path to the ABACUS input files, which contains the INPUT, STRU, KPT, and pseudopotential or orbital files. pdos_mode: Mode of plotted PDOS file. + - "atoms": PDOS of a list of atoms will be plotted. - "species": Total PDOS of any species will be plotted in a picture. - "species+shell": PDOS for any shell (s, p, d, f, g,...) of any species will be plotted. PDOS of a shell of a species willbe plotted in a subplot. - "species+orbital": Orbital-resolved PDOS will be plotted. PDOS of orbitals in the same shell of a species will be plotted in a subplot. + pdos_atom_indices: A list of atom indices, only used if pdos_mode is "atoms". dos_edelta_ev: Step size in writing Density of States (DOS) in eV. dos_sigma: Width of the Gaussian factor when obtaining smeared Density of States (DOS) in eV. dos_emin_ev: Minimal range for Density of States (DOS) in eV. Default is -10.0. dos_emax_ev: Maximal range for Density of States (DOS) in eV. Default is 10.0. - - Returns: - Dict[str, Any]: A dictionary containing: - - dos_fig_path: Path to the plotted DOS. - - pdos_fig_path: Path to the plotted PDOS. Only for LCAO basis. - - scf_work_path: Path to the work directory of SCF calculation. - - scf_normal_end: If the SCF calculation ended normally. - - scf_steps: Number of steps of SCF iteration. - - scf_converge: If the SCF calculation converged. - - scf_energy: The calculated energy of SCF calculation. - - nscf_work_path: Path to the work directory of NSCF calculation. - - nscf_normal_end: If the SCF calculation ended normally. + """ - return _abacus_dos_run(abacus_inputs_dir, pdos_mode, dos_edelta_ev, dos_sigma, dos_emin_ev, dos_emax_ev) + return _abacus_dos_run(abacus_inputs_dir, pdos_mode, pdos_atom_indices, dos_edelta_ev, dos_sigma, dos_emin_ev, dos_emax_ev) diff --git a/src/abacusagent/modules/submodules/dos.py b/src/abacusagent/modules/submodules/dos.py index 109074a..6da9346 100644 --- a/src/abacusagent/modules/submodules/dos.py +++ b/src/abacusagent/modules/submodules/dos.py @@ -8,7 +8,7 @@ from abacustest.lib_model.comm import check_abacus_inputs from pathlib import Path -from typing import Dict, Any, List, Literal, Optional +from typing import Dict, Any, List, Literal, Optional, Tuple from abacusagent.modules.util.comm import generate_work_path, link_abacusjob, run_abacus, has_chgfile, collect_metrics from abacusagent.modules.util.chemical_elements import MAX_ANGULAR_MOMENTUM_OF_ELEMENTS @@ -44,7 +44,8 @@ def abacus_dos_run( abacus_inputs_dir: Path, - pdos_mode: Literal['species', 'species+shell', 'species+orbital'] = 'species+shell', + pdos_mode: Literal['atoms', 'species', 'species+shell', 'species+orbital'] = 'species+shell', + pdos_atom_indices: Optional[List[int]] = None, dos_edelta_ev: float = 0.01, dos_sigma: float = 0.07, dos_emin_ev: float = -10.0, @@ -60,9 +61,11 @@ def abacus_dos_run( Args: abacus_inputs_dir: Path to the ABACUS input files, which contains the INPUT, STRU, KPT, and pseudopotential or orbital files. pdos_mode: Mode of plotted PDOS file. + - "atoms": PDOS of a list of atoms will be plotted. - "species": Total PDOS of any species will be plotted in a picture. - "species+shell": PDOS for any shell (s, p, d, f, g,...) of any species will be plotted. PDOS of a shell of a species willbe plotted in a subplot. - "species+orbital": Orbital-resolved PDOS will be plotted. PDOS of orbitals in the same shell of a species will be plotted in a subplot. + pdos_atom_indices: A list of atom indices, only used if pdos_mode is "atoms". dos_edelta_ev: Step size in writing Density of States (DOS) in eV. dos_sigma: Width of the Gaussian factor when obtaining smeared Density of States (DOS) in eV. dos_emin_ev: Minimal range for Density of States (DOS) in eV. Default is -10.0. @@ -101,6 +104,7 @@ def abacus_dos_run( metrics_nscf["nscf_work_path"], nspin, pdos_mode, + pdos_atom_indices, dos_emin_ev, dos_emax_ev) @@ -115,6 +119,8 @@ def abacus_dos_run( return return_dict except Exception as e: + import traceback + traceback.print_exc() return {"message": f"Calculating DOS and PDOS failed: {e}"} def abacus_dos_run_scf(abacus_inputs_dir: Path, @@ -320,6 +326,71 @@ def write_pdos_data_species_shell(energy: NDArray[np.float64], line_spin_down += f"{pdos_data[i*2+1]:20.8f}" f.write(line + line_spin_up + line_spin_down + "\n") +def write_pdos_data_atom(energy: NDArray[np.float64], + atom_pdos: List[Dict[str, NDArray[np.float64]]], + atom_index: int, + label_map: Dict[Tuple[int, str, int, int, int], str] = {}, + nspin: Literal[1, 2] = 1, + output_dir: str = './') -> None: + # Initialize dict of pdos data for each orbital + orbital_label = {} + for (_, species, l, m, z), full_label in label_map.items(): + if str(l) not in orbital_label.keys(): + orbital_label[str(l)] = {} + if str(m) not in orbital_label[str(l)].keys(): + orbital_name = full_label.split('(')[1].split(')')[0] + if orbital_name in orbital_rep_map.keys(): + orbital_label[str(l)][str(m)] = orbital_rep_map[orbital_name] + else: + orbital_label[str(l)][str(m)] = orbital_name + else: + pass + + # Write PDOS data of an atom to file + atom_pdos_lms = {} + orbital_names = [] + for orbital in atom_pdos: + if orbital['l'] <= angular_momentum_map.index(MAX_ANGULAR_MOMENTUM_OF_ELEMENTS[orbital['species']]): + if orbital['l'] not in atom_pdos_lms.keys(): + atom_pdos_lms[orbital['l']] = {} + atom_pdos_lms[orbital['l']][orbital['m']] = orbital['data'] + orbital_names.append(orbital_label[str(orbital['l'])][str(orbital['m'])]) + else: + if orbital['m'] not in atom_pdos_lms[orbital['l']].keys(): + atom_pdos_lms[orbital['l']][orbital['m']] = orbital['data'] + orbital_names.append(orbital_label[str(orbital['l'])][str(orbital['m'])]) + else: + atom_pdos_lms[orbital['l']][orbital['m']] += orbital['data'] + + atom_species = atom_pdos[0]['species'] + with open(os.path.join(output_dir, f"PDOS_{atom_species}{atom_index}.dat"), "w") as f: + if nspin == 1: + header_line = ' Energy (eV)' + for orbital_name in orbital_names: + header_line += f'{orbital_name.replace("$", ""):>20s}' + f.write(header_line + "\n") + for i in range(len(energy)): + line = f"{energy[i]:14.6f}" + for l, pdos_data in atom_pdos_lms.items(): + for m, pdos_data_m in pdos_data.items(): + line += f"{pdos_data_m[i]:20.8f}" + f.write(line + "\n") + elif nspin == 2: + header_spin_up, header_spin_down = '', '' + for orbital_name in orbital_names: + header_spin_up += f'{orbital_name.replace("$", ""):>17s}_up' + header_spin_down += f'{orbital_name.replace("$", ""):>17}_dn' + f.write(f" Energy (eV){header_spin_up}{header_spin_down}\n") + for i in range(len(energy)): + line = f"{energy[i]:14.6f}" + line_spin_up, line_spin_down = '', '' + for l, pdos_data in atom_pdos_lms.items(): + for m, pdos_data_m in pdos_data.items(): + # even index for spin up, odd index for spin down + line_spin_up += f"{pdos_data_m[i*2]:20.8f}" + line_spin_down += f"{pdos_data_m[i*2+1]:20.8f}" + f.write(line + line_spin_up + line_spin_down + "\n") + def write_pdos_data_species_orbital(energy: NDArray[np.float64], species_pdos: Dict[str, NDArray[np.float64]], species: str, @@ -359,7 +430,7 @@ def write_pdos_data_species_orbital(energy: NDArray[np.float64], line_spin_down += f"{species_orbital_pdos[i*2+1]:20.8f}" f.write(line + line_spin_up + line_spin_down + "\n") -def plot_pdos(energy_values, orbitals, fermi_level, label_map, output_dir, nspin, mode, dos_emin_ev, dos_emax_ev, dpi=300): +def plot_pdos(energy_values, orbitals, fermi_level, label_map, output_dir, nspin, mode, pdos_atom_indices, dos_emin_ev, dos_emax_ev, dpi=300): """Plot PDOS data separated by atom/species with custom labels.""" # Create output directory if it doesn't exist os.makedirs(output_dir, exist_ok=True) @@ -381,6 +452,8 @@ def plot_pdos(energy_values, orbitals, fermi_level, label_map, output_dir, nspin pdos_pic_file = plot_pdos_species_shell(shifted_energy, orbitals, output_dir, nspin, dos_emin_ev, dos_emax_ev, dpi) elif mode == "species+orbital": pdos_pic_file = plot_pdos_species_orbital(shifted_energy, orbitals, output_dir, nspin, label_map, dos_emin_ev, dos_emax_ev, dpi) + elif mode == "atoms": + pdos_pic_file = plot_pdos_atoms(shifted_energy, orbitals, output_dir, nspin, label_map, pdos_atom_indices, dos_emin_ev, dos_emax_ev, dpi) else: raise ValueError(f"Not allowed mode {mode}") @@ -553,6 +626,61 @@ def plot_pdos_species_orbital(shifted_energy, orbitals, output_dir, nspin, label return Path(pdos_pic_file).absolute() +def plot_pdos_atoms(shifted_energy, orbitals, output_dir, nspin, label_map, pdos_atom_indices, dos_emin_ev, dos_emax_ev, dpi): + # Plot PDOS for selected atoms + atom_pdoses = {} + for orbital in orbitals: + atom_index = orbital['atom_index'] + if atom_index in pdos_atom_indices: + if atom_index not in atom_pdoses.keys(): + atom_pdoses[atom_index] = [orbital] + else: + atom_pdoses[atom_index].append(orbital) + + fig, axes = plt.subplots(nrows=len(pdos_atom_indices), ncols=1, figsize=(8, 4*len(pdos_atom_indices))) + if len(pdos_atom_indices) == 1: + axes = [axes] + + for atom_index, atom_pdos in atom_pdoses.items(): + # Write PDOS data of all legal orbitals for each selected atom to file + write_pdos_data_atom(shifted_energy, atom_pdos, atom_index, label_map, nspin, output_dir) + # Plot PDOS of all legal shells for each selected atom + atom_pdos_ls = {} + for orbital in atom_pdos: + orbital_l_name = angular_momentum_map[orbital['l']] + if orbital_l_name not in atom_pdos_ls.keys(): + atom_pdos_ls[orbital_l_name] = orbital['data'] + else: + atom_pdos_ls[orbital_l_name] += orbital['data'] + + species = atom_pdos[0]['species'] + ax = axes[pdos_atom_indices.index(atom_index)] + for l, pdos_data in atom_pdos_ls.items(): + if angular_momentum_map.index(l) <= angular_momentum_map.index(MAX_ANGULAR_MOMENTUM_OF_ELEMENTS[species]): + if nspin == 1: + ax.plot(shifted_energy, pdos_data, label=f'{species}{atom_index}-{l}', linewidth=1.0) + elif nspin == 2: + ax.plot(shifted_energy, pdos_data[::2], color=color_map[l], label=f'{species}{atom_index}-{l} ' + r'$\uparrow$', linestyle='-', linewidth=1.0) + ax.plot(shifted_energy, -pdos_data[1::2], color=color_map[l], label=f'{species}{atom_index}-{l} ' + r'$\downarrow$', linestyle='--', linewidth=1.0) + + ax.axvline(x=0, color='black', linestyle=':', linewidth=1.0) + ax.set_title(f'PDOS for {species}{atom_index}', fontsize=12, pad=10) + ax.set_ylabel(r"States ($eV^{-1}$)", fontsize=10) + ax.set_xlim(dos_emin_ev, dos_emax_ev) + if nspin == 1: + ax.set_ylim(bottom=0) + ax.legend(fontsize=8, ncol=nspin) + ax.grid(alpha=0.3) + + axes[-1].set_xlabel('Energy (eV)', fontsize=10) + + plt.tight_layout() + pdos_pic_file = os.path.join(output_dir, 'PDOS.png') + plt.savefig(pdos_pic_file, dpi=dpi, bbox_inches='tight') + plt.close() + + return Path(pdos_pic_file).absolute() + def write_dos_data(energy: NDArray[np.float64], dos: NDArray[np.float64], dos_dn: Optional[NDArray[np.float64]] = None, @@ -618,6 +746,7 @@ def plot_dos_pdos(scf_job_path: Path, output_dir: Path, nspin: Literal[1, 2] = 1, mode: Literal['species', 'species+shell', 'species+orbital'] = 'species+shell', + pdos_atom_indices: Optional[List[int]] = None, dos_emin_ev: float = -10.0, dos_emax_ev: float = 10.0, dpi=300) -> List[str]: @@ -669,7 +798,7 @@ def plot_dos_pdos(scf_job_path: Path, if basis_type != 'pw': label_map = parse_basref_file(basref_file) energy_values, orbitals = parse_pdos_file(pdos_file) - pdos_plot_file = plot_pdos(energy_values, orbitals, fermi_level, label_map, output_dir, nspin, mode, dos_emin_ev, dos_emax_ev, dpi) + pdos_plot_file = plot_pdos(energy_values, orbitals, fermi_level, label_map, output_dir, nspin, mode, pdos_atom_indices, dos_emin_ev, dos_emax_ev, dpi) # Combine file paths into a single list all_plot_files.append(pdos_plot_file) From 9f29ad910e5268efcd195051441f9b21389004fb Mon Sep 17 00:00:00 2001 From: ahxbcn Date: Tue, 13 Jan 2026 08:39:12 +0800 Subject: [PATCH 05/15] Update wrapper for DOS --- src/abacusagent/modules/tool_wrapper.py | 27 ++++++++++------------- tests/integrate_test/test_tool_wrapper.py | 3 +-- 2 files changed, 13 insertions(+), 17 deletions(-) diff --git a/src/abacusagent/modules/tool_wrapper.py b/src/abacusagent/modules/tool_wrapper.py index 3195121..9a0387c 100644 --- a/src/abacusagent/modules/tool_wrapper.py +++ b/src/abacusagent/modules/tool_wrapper.py @@ -382,13 +382,12 @@ def abacus_dos_run( relax_precision: Literal['low', 'medium', 'high'] = 'medium', relax_method: Literal["cg", "bfgs", "bfgs_trad", "cg_bfgs", "sd", "fire"] = "cg", fixed_axes: Literal["None", "volume", "shape", "a", "b", "c", "ab", "ac", "bc"] = None, - pdos_mode: Literal['species', 'species+shell', 'species+orbital'] = 'species+shell', + pdos_mode: Literal['atoms', 'species', 'species+shell', 'species+orbital'] = 'species+shell', + pdos_atom_indices: Optional[List[int]] = None, dos_edelta_ev: float = 0.01, dos_sigma: float = 0.07, - dos_scale: float = 0.01, - dos_emin_ev: float = None, - dos_emax_ev: float = None, - dos_nche: int = None, + dos_emin_ev: float = -10.0, + dos_emax_ev: float = 10.0, ) -> Dict[str, Any]: """ Run the DOS and PDOS calculation. @@ -432,17 +431,16 @@ def abacus_dos_run( - ac: fix both a and c axes - bc: fix both b and c axes pdos_mode: Mode of plotted PDOS file. + - "atoms": PDOS of a list of atoms will be plotted. - "species": Total PDOS of any species will be plotted in a picture. - "species+shell": PDOS for any shell (s, p, d, f, g,...) of any species will be plotted. PDOS of a shell of a species willbe plotted in a subplot. - - “species+orbital": Orbital-resolved PDOS will be plotted. PDOS of orbitals in the same shell of a species will be plotted in a subplot. + - "species+orbital": Orbital-resolved PDOS will be plotted. PDOS of orbitals in the same shell of a species will be plotted in a subplot. + pdos_atom_indices: A list of atom indices, only used if pdos_mode is "atoms". dos_edelta_ev: Step size in writing Density of States (DOS) in eV. dos_sigma: Width of the Gaussian factor when obtaining smeared Density of States (DOS) in eV. - dos_scale: Defines the energy range of DOS output as (emax-emin)*(1+dos_scale), centered at (emax+emin)/2. - This parameter will be used when dos_emin_ev and dos_emax_ev are not set. - dos_emin_ev: Minimal range for Density of States (DOS) in eV. - dos_emax_ev: Maximal range for Density of States (DOS) in eV. - dos_nche: The order of Chebyshev expansions when using Stochastic Density Functional Theory (SDFT) to calculate DOS. - + dos_emin_ev: Minimal range for Density of States (DOS) in eV. Default is -10.0. + dos_emax_ev: Maximal range for Density of States (DOS) in eV. Default is 10.0. + Returns: Dict[str, Any]: A dictionary containing: - dos_fig_path: Path to the plotted DOS. @@ -474,12 +472,11 @@ def abacus_dos_run( dos_results = _abacus_dos_run(abacus_inputs_dir, pdos_mode, + pdos_atom_indices, dos_edelta_ev, dos_sigma, - dos_scale, dos_emin_ev, - dos_emax_ev, - dos_nche) + dos_emax_ev) return {'dos_fig_path': dos_results.get('dos_fig_path', None), 'pdos_fig_path': dos_results.get('pdos_fig_path', None), diff --git a/tests/integrate_test/test_tool_wrapper.py b/tests/integrate_test/test_tool_wrapper.py index c8f33cb..58ef0cf 100644 --- a/tests/integrate_test/test_tool_wrapper.py +++ b/tests/integrate_test/test_tool_wrapper.py @@ -175,8 +175,7 @@ def test_run_abacus_calculation_dos(self): fixed_axes=None, pdos_mode='species+shell', dos_edelta_ev=0.01, - dos_sigma=0.07, - dos_scale=0.01) + dos_sigma=0.07) print(outputs) dos_fig_path = outputs['dos_fig_path'] From 51101931e4caaf6c4214ec5e1f4caedcaa79a856 Mon Sep 17 00:00:00 2001 From: ahxbcn Date: Tue, 13 Jan 2026 10:55:08 +0800 Subject: [PATCH 06/15] Fix a bug in writing PDOS data to file in nspin=2 case in atom mode --- src/abacusagent/modules/submodules/dos.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/abacusagent/modules/submodules/dos.py b/src/abacusagent/modules/submodules/dos.py index 6da9346..1028be3 100644 --- a/src/abacusagent/modules/submodules/dos.py +++ b/src/abacusagent/modules/submodules/dos.py @@ -389,7 +389,7 @@ def write_pdos_data_atom(energy: NDArray[np.float64], # even index for spin up, odd index for spin down line_spin_up += f"{pdos_data_m[i*2]:20.8f}" line_spin_down += f"{pdos_data_m[i*2+1]:20.8f}" - f.write(line + line_spin_up + line_spin_down + "\n") + f.write(line + line_spin_up + line_spin_down + "\n") def write_pdos_data_species_orbital(energy: NDArray[np.float64], species_pdos: Dict[str, NDArray[np.float64]], From f2ecd7df2d76074df1f6dc319f3d2a3378e11850 Mon Sep 17 00:00:00 2001 From: ahxbcn Date: Tue, 13 Jan 2026 17:54:41 +0800 Subject: [PATCH 07/15] Update return dos and pdos data file in abacus_dos_run --- src/abacusagent/modules/dos.py | 12 +++ src/abacusagent/modules/submodules/dos.py | 90 +++++++++++------- src/abacusagent/modules/tool_wrapper.py | 4 + tests/integrate_test/data/ref_results.json | 12 +++ tests/integrate_test/test_dos.py | 102 +++++++++++++++++++++ tests/integrate_test/test_tool_wrapper.py | 6 ++ 6 files changed, 194 insertions(+), 32 deletions(-) diff --git a/src/abacusagent/modules/dos.py b/src/abacusagent/modules/dos.py index 2362093..d18a7d0 100644 --- a/src/abacusagent/modules/dos.py +++ b/src/abacusagent/modules/dos.py @@ -34,5 +34,17 @@ def abacus_dos_run( dos_emin_ev: Minimal range for Density of States (DOS) in eV. Default is -10.0. dos_emax_ev: Maximal range for Density of States (DOS) in eV. Default is 10.0. + Returns: + Dict[str, Any]: A dictionary containing: + - dos_fig_path: Path to the plotted DOS. + - pdos_fig_path: Path to the plotted PDOS. Only for LCAO basis. + - dos_data_path: Path to the data used in plotting DOS. + - pdos_data_paths: Path to the data used in plotting PDOS. Only for LCAO basis. + - scf_work_path: Path to the work directory of SCF calculation. + - scf_normal_end: If the SCF calculation ended normally. + - scf_steps: Number of steps of SCF iteration. + - scf_converge: If the SCF calculation converged. + - scf_energy: The calculated energy of SCF calculation. + - nscf_work_path: Path to the work directory of NSCF calculation """ return _abacus_dos_run(abacus_inputs_dir, pdos_mode, pdos_atom_indices, dos_edelta_ev, dos_sigma, dos_emin_ev, dos_emax_ev) diff --git a/src/abacusagent/modules/submodules/dos.py b/src/abacusagent/modules/submodules/dos.py index 1028be3..8df868e 100644 --- a/src/abacusagent/modules/submodules/dos.py +++ b/src/abacusagent/modules/submodules/dos.py @@ -75,6 +75,8 @@ def abacus_dos_run( Dict[str, Any]: A dictionary containing: - dos_fig_path: Path to the plotted DOS. - pdos_fig_path: Path to the plotted PDOS. Only for LCAO basis. + - dos_data_path: Path to the data used in plotting DOS. + - pdos_data_paths: Path to the data used in plotting PDOS. Only for LCAO basis. - scf_work_path: Path to the work directory of SCF calculation. - scf_normal_end: If the SCF calculation ended normally. - scf_steps: Number of steps of SCF iteration. @@ -99,18 +101,20 @@ def abacus_dos_run( dos_edelta_ev=dos_edelta_ev, dos_sigma=dos_sigma) - fig_paths = plot_dos_pdos(metrics_scf["scf_work_path"], - metrics_nscf["nscf_work_path"], - metrics_nscf["nscf_work_path"], - nspin, - pdos_mode, - pdos_atom_indices, - dos_emin_ev, - dos_emax_ev) - + fig_paths, dos_pdos_data_paths = plot_dos_pdos(metrics_scf["scf_work_path"], + metrics_nscf["nscf_work_path"], + metrics_nscf["nscf_work_path"], + nspin, + pdos_mode, + pdos_atom_indices, + dos_emin_ev, + dos_emax_ev) + return_dict = {"dos_fig_path": fig_paths[0]} + return_dict['dos_data_path'] = dos_pdos_data_paths[0] try: return_dict['pdos_fig_path'] = fig_paths[1] + return_dict['pdos_data_paths'] = dos_pdos_data_paths[1:] except: pass # Do nothing if PDOS file is not plotted @@ -267,7 +271,8 @@ def write_pdos_data_species(energy: NDArray[np.float64], nspin: Literal[1, 2] = 1, output_dir: str = './') -> None: """Write PDOS of species to file""" - with open(os.path.join(output_dir, "PDOS_species.dat"), 'w') as f: + pdos_data_file = Path(os.path.join(output_dir, "PDOS_species.dat")).absolute() + with open(pdos_data_file, 'w') as f: if nspin == 1: header_line = ' Energy (eV)' for species_name in species_pdos_data.keys(): @@ -293,6 +298,8 @@ def write_pdos_data_species(energy: NDArray[np.float64], line_spin_down += f"{pdos_data[i*2+1]:20.8f}" line += line_spin_up + line_spin_down f.write(line + "\n") + + return pdos_data_file.absolute() def write_pdos_data_species_shell(energy: NDArray[np.float64], species_pdos_data_dict: Dict[str, NDArray[np.float64]], @@ -300,6 +307,7 @@ def write_pdos_data_species_shell(energy: NDArray[np.float64], nspin: Literal[1, 2] = 1, output_dir: str = './') -> None: # Write PDOS of different shells (s, p, d...) for a species to file + pdos_data_file = Path(os.path.join(output_dir, f"PDOS_{species}.dat")).absolute() with open(os.path.join(output_dir, f"PDOS_{species}.dat"), "w") as f: if nspin == 1: header_line = ' Energy (eV)' @@ -325,6 +333,8 @@ def write_pdos_data_species_shell(energy: NDArray[np.float64], line_spin_up += f"{pdos_data[i*2]:20.8f}" line_spin_down += f"{pdos_data[i*2+1]:20.8f}" f.write(line + line_spin_up + line_spin_down + "\n") + + return Path(pdos_data_file).absolute() def write_pdos_data_atom(energy: NDArray[np.float64], atom_pdos: List[Dict[str, NDArray[np.float64]]], @@ -363,7 +373,8 @@ def write_pdos_data_atom(energy: NDArray[np.float64], atom_pdos_lms[orbital['l']][orbital['m']] += orbital['data'] atom_species = atom_pdos[0]['species'] - with open(os.path.join(output_dir, f"PDOS_{atom_species}{atom_index}.dat"), "w") as f: + pdos_data_file = Path(os.path.join(output_dir, f"PDOS_{atom_species}{atom_index}.dat")) + with open(pdos_data_file, "w") as f: if nspin == 1: header_line = ' Energy (eV)' for orbital_name in orbital_names: @@ -390,6 +401,8 @@ def write_pdos_data_atom(energy: NDArray[np.float64], line_spin_up += f"{pdos_data_m[i*2]:20.8f}" line_spin_down += f"{pdos_data_m[i*2+1]:20.8f}" f.write(line + line_spin_up + line_spin_down + "\n") + + return pdos_data_file def write_pdos_data_species_orbital(energy: NDArray[np.float64], species_pdos: Dict[str, NDArray[np.float64]], @@ -398,7 +411,8 @@ def write_pdos_data_species_orbital(energy: NDArray[np.float64], nspin: Literal[1, 2] = 1, output_dir: str = "./") -> None: # Write PDOS data of different orbitals of a species to file - with open(os.path.join(output_dir, f"PDOS_{species}.dat"), "w") as f: + pdos_data_file = Path(os.path.join(output_dir, f"PDOS_{species}.dat")) + with open(pdos_data_file, "w") as f: if nspin == 1: header_line = f" Energy (eV)" for l, species_shell_pdos in species_pdos.items(): @@ -430,6 +444,8 @@ def write_pdos_data_species_orbital(energy: NDArray[np.float64], line_spin_down += f"{species_orbital_pdos[i*2+1]:20.8f}" f.write(line + line_spin_up + line_spin_down + "\n") + return pdos_data_file + def plot_pdos(energy_values, orbitals, fermi_level, label_map, output_dir, nspin, mode, pdos_atom_indices, dos_emin_ev, dos_emax_ev, dpi=300): """Plot PDOS data separated by atom/species with custom labels.""" # Create output directory if it doesn't exist @@ -447,17 +463,17 @@ def plot_pdos(energy_values, orbitals, fermi_level, label_map, output_dir, nspin atom_species_groups[key].append(orbital) if mode == "species": - pdos_pic_file = plot_pdos_species(shifted_energy, orbitals, output_dir, nspin, dos_emin_ev, dos_emax_ev, dpi) + pdos_pic_file, pdos_data_files = plot_pdos_species(shifted_energy, orbitals, output_dir, nspin, dos_emin_ev, dos_emax_ev, dpi) elif mode == "species+shell": - pdos_pic_file = plot_pdos_species_shell(shifted_energy, orbitals, output_dir, nspin, dos_emin_ev, dos_emax_ev, dpi) + pdos_pic_file, pdos_data_files = plot_pdos_species_shell(shifted_energy, orbitals, output_dir, nspin, dos_emin_ev, dos_emax_ev, dpi) elif mode == "species+orbital": - pdos_pic_file = plot_pdos_species_orbital(shifted_energy, orbitals, output_dir, nspin, label_map, dos_emin_ev, dos_emax_ev, dpi) + pdos_pic_file, pdos_data_files = plot_pdos_species_orbital(shifted_energy, orbitals, output_dir, nspin, label_map, dos_emin_ev, dos_emax_ev, dpi) elif mode == "atoms": - pdos_pic_file = plot_pdos_atoms(shifted_energy, orbitals, output_dir, nspin, label_map, pdos_atom_indices, dos_emin_ev, dos_emax_ev, dpi) + pdos_pic_file, pdos_data_files = plot_pdos_atoms(shifted_energy, orbitals, output_dir, nspin, label_map, pdos_atom_indices, dos_emin_ev, dos_emax_ev, dpi) else: raise ValueError(f"Not allowed mode {mode}") - return pdos_pic_file + return pdos_pic_file, pdos_data_files def plot_pdos_species(shifted_energy, orbitals, output_dir, nspin, dos_emin_ev, dos_emax_ev, dpi): species = {} @@ -468,7 +484,7 @@ def plot_pdos_species(shifted_energy, orbitals, output_dir, nspin, dos_emin_ev, else: species[species_one] += orbital['data'] - write_pdos_data_species(shifted_energy, species, nspin, output_dir) + pdos_data_files = [write_pdos_data_species(shifted_energy, species, nspin, output_dir)] # Plot PDOS of different species plt.plot(figsize=(10, 6)) for species_name, pdos_data in species.items(): @@ -492,7 +508,7 @@ def plot_pdos_species(shifted_energy, orbitals, output_dir, nspin, dos_emin_ev, plt.savefig(pdos_pic_file, dpi=dpi) plt.close() - return Path(pdos_pic_file).absolute() + return Path(pdos_pic_file).absolute(), pdos_data_files def plot_pdos_species_shell(shifted_energy, orbitals, output_dir, nspin, dos_emin_ev, dos_emax_ev, dpi): species_shells = {} @@ -514,10 +530,12 @@ def plot_pdos_species_shell(shifted_energy, orbitals, output_dir, nspin, dos_emi fig, axes = plt.subplots(nrows=num_species, ncols=1, figsize=(8, 4*num_species)) if num_species == 1: axes = [axes] - + + pdos_data_files = [] for species_idx, (species, pdos_data_dict) in enumerate(species_shells.items()): # Write PDOS data to the selected species to file - write_pdos_data_species_shell(shifted_energy, pdos_data_dict, species, nspin, output_dir) + pdos_data_file = write_pdos_data_species_shell(shifted_energy, pdos_data_dict, species, nspin, output_dir) + pdos_data_files.append(pdos_data_file) # Plot PDOS of different shells for a species ax = axes[species_idx] @@ -544,7 +562,7 @@ def plot_pdos_species_shell(shifted_energy, orbitals, output_dir, nspin, dos_emi plt.savefig(pdos_pic_file, dpi=dpi, bbox_inches='tight') plt.close() - return Path(pdos_pic_file).absolute() + return Path(pdos_pic_file).absolute(), pdos_data_files def plot_pdos_species_orbital(shifted_energy, orbitals, output_dir, nspin, label_map, dos_emin_ev, dos_emax_ev, dpi): @@ -593,9 +611,11 @@ def plot_pdos_species_orbital(shifted_energy, orbitals, output_dir, nspin, label fig, axes = plt.subplots(nrows=total_subplots, ncols=1, figsize=(8, 4*total_subplots)) subplot_count = 0 + pdos_data_files = [] for species, species_pdos in species_orbitals.items(): # Write PDOS data of orbitals of the selected species to file - write_pdos_data_species_orbital(shifted_energy, species_pdos, species, orbital_label, nspin, output_dir) + pdos_data_file = write_pdos_data_species_orbital(shifted_energy, species_pdos, species, orbital_label, nspin, output_dir) + pdos_data_files.append(pdos_data_file) for angular_momentum, species_shell_pdos in species_pdos.items(): for m, species_orbital_pdos in species_shell_pdos.items(): ax = axes[subplot_count] @@ -624,7 +644,7 @@ def plot_pdos_species_orbital(shifted_energy, orbitals, output_dir, nspin, label plt.savefig(pdos_pic_file, dpi=dpi, bbox_inches='tight') plt.close() - return Path(pdos_pic_file).absolute() + return Path(pdos_pic_file).absolute(), pdos_data_files def plot_pdos_atoms(shifted_energy, orbitals, output_dir, nspin, label_map, pdos_atom_indices, dos_emin_ev, dos_emax_ev, dpi): # Plot PDOS for selected atoms @@ -641,9 +661,11 @@ def plot_pdos_atoms(shifted_energy, orbitals, output_dir, nspin, label_map, pdos if len(pdos_atom_indices) == 1: axes = [axes] + pdos_data_files = [] for atom_index, atom_pdos in atom_pdoses.items(): # Write PDOS data of all legal orbitals for each selected atom to file - write_pdos_data_atom(shifted_energy, atom_pdos, atom_index, label_map, nspin, output_dir) + pdos_data_file = write_pdos_data_atom(shifted_energy, atom_pdos, atom_index, label_map, nspin, output_dir) + pdos_data_files.append(pdos_data_file) # Plot PDOS of all legal shells for each selected atom atom_pdos_ls = {} for orbital in atom_pdos: @@ -679,7 +701,7 @@ def plot_pdos_atoms(shifted_energy, orbitals, output_dir, nspin, label_map, pdos plt.savefig(pdos_pic_file, dpi=dpi, bbox_inches='tight') plt.close() - return Path(pdos_pic_file).absolute() + return Path(pdos_pic_file).absolute(), pdos_data_files def write_dos_data(energy: NDArray[np.float64], dos: NDArray[np.float64], @@ -698,6 +720,8 @@ def write_dos_data(energy: NDArray[np.float64], f.write(f"{energy[i]:16.6f}{dos[i]:16.6f}{dos_dn[i]:16.6f}\n") f.close() + return Path(out_filename).absolute() + def plot_dos(file_path: List[Path], fermi_level: float, dos_plot_file: str = 'DOS.png', @@ -717,9 +741,9 @@ def plot_dos(file_path: List[Path], # Write shifted dos data to file dos_output_file = os.path.join(os.path.dirname(dos_plot_file), "DOS_shifted.dat") if nspin == 1: - write_dos_data(energy, dos, nspin=nspin, out_filename=dos_output_file) + dos_data_file = write_dos_data(energy, dos, nspin=nspin, out_filename=dos_output_file) elif nspin == 2: - write_dos_data(energy, dos, dos_dn, nspin=nspin, out_filename=dos_output_file) + dos_data_file = write_dos_data(energy, dos, dos_dn, nspin=nspin, out_filename=dos_output_file) # Create plot plt.figure(figsize=(8, 6)) if nspin == 1: @@ -739,7 +763,7 @@ def plot_dos(file_path: List[Path], plt.savefig(dos_plot_file, dpi=dpi, bbox_inches='tight') plt.close() - return Path(dos_plot_file).absolute() + return Path(dos_plot_file).absolute(), Path(dos_data_file).absolute() def plot_dos_pdos(scf_job_path: Path, nscf_job_path: Path, @@ -788,8 +812,9 @@ def plot_dos_pdos(scf_job_path: Path, fermi_level = collect_metrics(scf_job_path, ['efermi'])['efermi'] # Plot DOS and get file path - dos_plot_file = plot_dos(dos_file, fermi_level, dos_output, nspin, dos_emin_ev, dos_emax_ev, dpi) + dos_plot_file, dos_data_file = plot_dos(dos_file, fermi_level, dos_output, nspin, dos_emin_ev, dos_emax_ev, dpi) all_plot_files = [dos_plot_file] + dos_pdos_data_files = [dos_data_file] print("DOS file plotted") @@ -798,10 +823,11 @@ def plot_dos_pdos(scf_job_path: Path, if basis_type != 'pw': label_map = parse_basref_file(basref_file) energy_values, orbitals = parse_pdos_file(pdos_file) - pdos_plot_file = plot_pdos(energy_values, orbitals, fermi_level, label_map, output_dir, nspin, mode, pdos_atom_indices, dos_emin_ev, dos_emax_ev, dpi) + pdos_plot_file, pdos_data_files = plot_pdos(energy_values, orbitals, fermi_level, label_map, output_dir, nspin, mode, pdos_atom_indices, dos_emin_ev, dos_emax_ev, dpi) # Combine file paths into a single list all_plot_files.append(pdos_plot_file) + dos_pdos_data_files += pdos_data_files else: print(f"Warning: PDOS calculation not supported for PW basis type, skipping PDOS plotting") elif os.path.exists(pdos_file) and not os.path.exists(basref_file): @@ -815,4 +841,4 @@ def plot_dos_pdos(scf_job_path: Path, for file in all_plot_files: print(f"- {file}") - return all_plot_files + return all_plot_files, dos_pdos_data_files diff --git a/src/abacusagent/modules/tool_wrapper.py b/src/abacusagent/modules/tool_wrapper.py index 9a0387c..b9adbf4 100644 --- a/src/abacusagent/modules/tool_wrapper.py +++ b/src/abacusagent/modules/tool_wrapper.py @@ -445,6 +445,8 @@ def abacus_dos_run( Dict[str, Any]: A dictionary containing: - dos_fig_path: Path to the plotted DOS. - pdos_fig_path: Path to the plotted PDOS. Only for LCAO basis. + - dos_data_path: Path to the data used in plotting DOS. + - pdos_data_paths: Path to the data used in plotting PDOS. Only for LCAO basis. - scf_normal_end: If the SCF calculation ended normally. - scf_converge: If the SCF calculation converged. - scf_energy: The calculated energy of SCF calculation. @@ -480,6 +482,8 @@ def abacus_dos_run( return {'dos_fig_path': dos_results.get('dos_fig_path', None), 'pdos_fig_path': dos_results.get('pdos_fig_path', None), + 'dos_data_path': dos_results.get('dos_data_path', None), + 'pdos_data_paths': dos_results.get('pdos_data_paths', None), 'scf_normal_end': dos_results.get('scf_normal_end', None), 'scf_converge': dos_results.get('scf_converge', None), 'scf_energy': dos_results.get('scf_energy', None), diff --git a/tests/integrate_test/data/ref_results.json b/tests/integrate_test/data/ref_results.json index bd9879f..a7a5ec0 100644 --- a/tests/integrate_test/data/ref_results.json +++ b/tests/integrate_test/data/ref_results.json @@ -122,6 +122,12 @@ "scf_energy": -1688.5908919 } }, + "test_abacus_dos_run_atoms": + { + "result": { + "scf_energy": -1688.5908919 + } + }, "test_abacus_dos_run_species_nspin2": { "result": { @@ -140,6 +146,12 @@ "scf_energy": -3220.427369911483 } }, + "test_abacus_dos_run_atoms_nspin2": + { + "result": { + "scf_energy": -3220.427369911483 + } + }, "test_abacus_dos_run_pw_nspin1": { "result": { diff --git a/tests/integrate_test/test_dos.py b/tests/integrate_test/test_dos.py index 70cbe65..6b52a5c 100644 --- a/tests/integrate_test/test_dos.py +++ b/tests/integrate_test/test_dos.py @@ -48,12 +48,18 @@ def test_abacus_dos_run_species(self): dos_sigma = 0.07, dos_emin_ev=-20, dos_emax_ev=20) + print(outputs) dos_fig_path = outputs['dos_fig_path'] + dos_data_path = outputs['dos_data_path'] pdos_fig_path = outputs['pdos_fig_path'] + pdos_data_paths = outputs['pdos_data_paths'] self.assertIsInstance(dos_fig_path, get_path_type()) + self.assertIsInstance(dos_data_path, get_path_type()) self.assertIsInstance(pdos_fig_path, get_path_type()) + for pdos_data_path in pdos_data_paths: + self.assertIsInstance(pdos_data_path, get_path_type()) self.assertTrue(outputs['scf_normal_end']) self.assertTrue(outputs['scf_converge']) self.assertTrue(outputs['nscf_normal_end']) @@ -75,10 +81,15 @@ def test_abacus_dos_run_species_shell(self): pdos_mode='species+shell') dos_fig_path = outputs['dos_fig_path'] + dos_data_path = outputs['dos_data_path'] pdos_fig_path = outputs['pdos_fig_path'] + pdos_data_paths = outputs['pdos_data_paths'] self.assertIsInstance(dos_fig_path, get_path_type()) + self.assertIsInstance(dos_data_path, get_path_type()) self.assertIsInstance(pdos_fig_path, get_path_type()) + for pdos_data_path in pdos_data_paths: + self.assertIsInstance(pdos_data_path, get_path_type()) self.assertTrue(outputs['scf_normal_end']) self.assertTrue(outputs['scf_converge']) self.assertTrue(outputs['nscf_normal_end']) @@ -100,15 +111,52 @@ def test_abacus_dos_run_species_orbital(self): pdos_mode='species+orbital') dos_fig_path = outputs['dos_fig_path'] + dos_data_path = outputs['dos_data_path'] + pdos_fig_path = outputs['pdos_fig_path'] + pdos_data_paths = outputs['pdos_data_paths'] + + self.assertIsInstance(dos_fig_path, get_path_type()) + self.assertIsInstance(dos_data_path, get_path_type()) + self.assertIsInstance(pdos_fig_path, get_path_type()) + for pdos_data_path in pdos_data_paths: + self.assertIsInstance(pdos_data_path, get_path_type()) + self.assertTrue(outputs['scf_normal_end']) + self.assertTrue(outputs['scf_converge']) + self.assertTrue(outputs['nscf_normal_end']) + self.assertAlmostEqual(outputs['scf_energy'], ref_results['scf_energy']) + + def test_abacus_dos_run_atoms(self): + """ + Test the abacus_dos_run function with PDOS plotting mode set to different species and shell. + """ + test_func_name = inspect.currentframe().f_code.co_name + ref_results = load_test_ref_result(test_func_name) + + test_work_dir = self.test_path / test_func_name + shutil.copytree(self.abacus_inputs_dir_nacl_prim, test_work_dir) + shutil.copy2(self.stru_dos_nacl_prim, test_work_dir / "STRU") + shutil.copy2(self.input_dos_nacl_prim, test_work_dir / "INPUT") + + outputs = abacus_dos_run(test_work_dir, + pdos_mode='atoms', + pdos_atom_indices = [1, 2]) + + dos_fig_path = outputs['dos_fig_path'] + dos_data_path = outputs['dos_data_path'] pdos_fig_path = outputs['pdos_fig_path'] + pdos_data_paths = outputs['pdos_data_paths'] self.assertIsInstance(dos_fig_path, get_path_type()) + self.assertIsInstance(dos_data_path, get_path_type()) self.assertIsInstance(pdos_fig_path, get_path_type()) + for pdos_data_path in pdos_data_paths: + self.assertIsInstance(pdos_data_path, get_path_type()) self.assertTrue(outputs['scf_normal_end']) self.assertTrue(outputs['scf_converge']) self.assertTrue(outputs['nscf_normal_end']) self.assertAlmostEqual(outputs['scf_energy'], ref_results['scf_energy']) + def test_abacus_dos_run_species_nspin2(self): """ Test the abacus_dos_run function with nspin=2 case and PDOS plotting mode set to different species. @@ -128,10 +176,15 @@ def test_abacus_dos_run_species_nspin2(self): dos_emax_ev = 20) dos_fig_path = outputs['dos_fig_path'] + dos_data_path = outputs['dos_data_path'] pdos_fig_path = outputs['pdos_fig_path'] + pdos_data_paths = outputs['pdos_data_paths'] self.assertIsInstance(dos_fig_path, get_path_type()) + self.assertIsInstance(dos_data_path, get_path_type()) self.assertIsInstance(pdos_fig_path, get_path_type()) + for pdos_data_path in pdos_data_paths: + self.assertIsInstance(pdos_data_path, get_path_type()) self.assertTrue(outputs['scf_normal_end']) self.assertTrue(outputs['scf_converge']) self.assertTrue(outputs['nscf_normal_end']) @@ -152,10 +205,15 @@ def test_abacus_dos_run_species_shell_nspin2(self): pdos_mode='species+shell') dos_fig_path = outputs['dos_fig_path'] + dos_data_path = outputs['dos_data_path'] pdos_fig_path = outputs['pdos_fig_path'] + pdos_data_paths = outputs['pdos_data_paths'] self.assertIsInstance(dos_fig_path, get_path_type()) + self.assertIsInstance(dos_data_path, get_path_type()) self.assertIsInstance(pdos_fig_path, get_path_type()) + for pdos_data_path in pdos_data_paths: + self.assertIsInstance(pdos_data_path, get_path_type()) self.assertTrue(outputs['scf_normal_end']) self.assertTrue(outputs['scf_converge']) self.assertTrue(outputs['nscf_normal_end']) @@ -176,10 +234,49 @@ def test_abacus_dos_run_species_orbital_nspin2(self): pdos_mode='species+orbital') dos_fig_path = outputs['dos_fig_path'] + dos_data_path = outputs['dos_data_path'] + pdos_fig_path = outputs['pdos_fig_path'] + pdos_data_paths = outputs['pdos_data_paths'] + + self.assertIsInstance(dos_fig_path, get_path_type()) + self.assertIsInstance(dos_data_path, get_path_type()) + self.assertIsInstance(pdos_fig_path, get_path_type()) + for pdos_data_path in pdos_data_paths: + self.assertIsInstance(pdos_data_path, get_path_type()) + self.assertTrue(outputs['scf_normal_end']) + self.assertTrue(outputs['scf_converge']) + self.assertTrue(outputs['nscf_normal_end']) + self.assertAlmostEqual(outputs['scf_energy'], ref_results['scf_energy']) + + def test_abacus_dos_run_atoms_nspin2(self): + """ + Test the abacus_dos_run function with nspin=2 case and PDOS plotting mode set to different species. + """ + test_func_name = inspect.currentframe().f_code.co_name + ref_results = load_test_ref_result(test_func_name) + + test_work_dir = self.test_path / test_func_name + shutil.copytree(self.abacus_inputs_dir_fe_bcc_prim, test_work_dir) + shutil.copy2(self.stru_dos_fe_bcc_prim, test_work_dir / "STRU") + + outputs = abacus_dos_run(test_work_dir, + pdos_mode='atoms', + pdos_atom_indices = [1], + dos_edelta_ev = 0.01, + dos_sigma = 0.07, + dos_emin_ev = -20, + dos_emax_ev = 20) + + dos_fig_path = outputs['dos_fig_path'] + dos_data_path = outputs['dos_data_path'] pdos_fig_path = outputs['pdos_fig_path'] + pdos_data_paths = outputs['pdos_data_paths'] self.assertIsInstance(dos_fig_path, get_path_type()) + self.assertIsInstance(dos_data_path, get_path_type()) self.assertIsInstance(pdos_fig_path, get_path_type()) + for pdos_data_path in pdos_data_paths: + self.assertIsInstance(pdos_data_path, get_path_type()) self.assertTrue(outputs['scf_normal_end']) self.assertTrue(outputs['scf_converge']) self.assertTrue(outputs['nscf_normal_end']) @@ -199,10 +296,13 @@ def test_abacus_dos_run_pw_nspin1(self): shutil.copy2(self.input_dos_pw_nacl_prim, test_work_dir / "INPUT") outputs = abacus_dos_run(test_work_dir) + print(outputs) dos_fig_path = outputs['dos_fig_path'] + dos_data_path = outputs['dos_data_path'] self.assertIsInstance(dos_fig_path, get_path_type()) + self.assertIsInstance(dos_data_path, get_path_type()) self.assertTrue(outputs['scf_normal_end']) self.assertTrue(outputs['scf_converge']) self.assertTrue(outputs['nscf_normal_end']) @@ -224,8 +324,10 @@ def test_abacus_dos_run_pw_nspin2(self): print(outputs) dos_fig_path = outputs['dos_fig_path'] + dos_data_path = outputs['dos_data_path'] self.assertIsInstance(dos_fig_path, get_path_type()) + self.assertIsInstance(dos_data_path, get_path_type()) self.assertTrue(outputs['scf_normal_end']) self.assertTrue(outputs['scf_converge']) self.assertTrue(outputs['nscf_normal_end']) diff --git a/tests/integrate_test/test_tool_wrapper.py b/tests/integrate_test/test_tool_wrapper.py index 58ef0cf..0039271 100644 --- a/tests/integrate_test/test_tool_wrapper.py +++ b/tests/integrate_test/test_tool_wrapper.py @@ -180,9 +180,15 @@ def test_run_abacus_calculation_dos(self): dos_fig_path = outputs['dos_fig_path'] pdos_fig_path = outputs['pdos_fig_path'] + dos_data_path = outputs['dos_data_path'] + pdos_data_paths = outputs['pdos_data_paths'] self.assertIsInstance(dos_fig_path, get_path_type()) + self.assertIsInstance(dos_data_path, get_path_type()) self.assertIsInstance(pdos_fig_path, get_path_type()) + for pdos_data_path in pdos_data_paths: + self.assertIsInstance(pdos_data_path, get_path_type()) + self.assertTrue(outputs['scf_normal_end']) self.assertTrue(outputs['scf_converge']) self.assertTrue(outputs['nscf_normal_end']) From d3cf7895f3b26e249fae27bd221eda8cce706cfd Mon Sep 17 00:00:00 2001 From: ahxbcn Date: Tue, 13 Jan 2026 19:08:58 +0800 Subject: [PATCH 08/15] Unify process of files and directories in link_abacusjob --- src/abacusagent/modules/submodules/dos.py | 4 ++- src/abacusagent/modules/util/comm.py | 41 ++++++++++------------- 2 files changed, 20 insertions(+), 25 deletions(-) diff --git a/src/abacusagent/modules/submodules/dos.py b/src/abacusagent/modules/submodules/dos.py index 8df868e..9133ca1 100644 --- a/src/abacusagent/modules/submodules/dos.py +++ b/src/abacusagent/modules/submodules/dos.py @@ -1,5 +1,6 @@ import os import re +import glob import numpy as np from numpy.typing import NDArray import matplotlib.pyplot as plt @@ -176,12 +177,13 @@ def abacus_dos_run_nscf(abacus_inputs_dir: Path, work_path = generate_work_path() link_abacusjob(src=abacus_inputs_dir, dst=work_path, - copy_files=["INPUT", "KPT"], + copy_files=["INPUT", "KPT"] + glob.glob(os.path.join(abacus_inputs_dir, "OUT.*")), exclude=["*log", "*json"]) input_param = ReadInput(os.path.join(work_path, "INPUT")) input_param["calculation"] = "nscf" input_param["init_chg"] = "file" + input_param["out_chg"] = -1 if input_param.get("basis_type", "pw") == "lcao": input_param["out_dos"] = 2 # only for LCAO basis, and will output DOS and PDOS else: diff --git a/src/abacusagent/modules/util/comm.py b/src/abacusagent/modules/util/comm.py index d52e7cf..9935b94 100644 --- a/src/abacusagent/modules/util/comm.py +++ b/src/abacusagent/modules/util/comm.py @@ -234,7 +234,7 @@ def link_abacusjob(src: str, exclude (Optional[List[str]]): List of files to exclude. If None, no files are excluded. copy_files (List[str]): List of files to copy from src to dst. Default is ["INPUT", "STRU", "KPT"]. overwrite (bool): If True, existing files in the destination will be overwritten. Default is True. - exclude_directories (bool): If True, directories will be not copied. Default is False. + exclude_directories (bool): If True, directories will be excluded from linking or copying. Default is False. Notes: - If somes files are included in both include and exclude, the file will be excluded. @@ -257,13 +257,13 @@ def link_abacusjob(src: str, if exclude is None: exclude = [] - exclude_files = [] + exclude_paths = [] for pattern in exclude: - exclude_files.extend(src.glob(pattern)) + exclude_paths.extend(src.glob(pattern)) os.makedirs(dst, exist_ok=True) # Remove excluded files from included files - include_paths = [f for f in include_paths if f not in exclude_files] + include_paths = [f for f in include_paths if f not in exclude_paths] if not include_paths: traceback.print_stack() print("No files to link after applying include and exclude patterns.\n", @@ -274,29 +274,22 @@ def link_abacusjob(src: str, for path in include_paths: if path == dst: continue - if exclude_directories and os.path.isdir(path): + if exclude_directories and path.is_dir(): continue - - if os.path.isfile(path): - dst_file = dst / path.name - if dst_file.exists(): - if overwrite: - dst_file.unlink() - else: - continue - if str(path.name) in copy_files: - os.system(f"cp {path} {dst_file}") + + dst_path = dst / path.name + if dst_path.exists(): + if overwrite: + dst_path.unlink() else: - os.symlink(path, dst_file) - else: - dst_dir = dst / path.name - if dst_dir.exists(): - if overwrite: - shutil.rmtree(dst_dir) - shutil.copytree(path, dst_dir) + continue + if str(path.name) in copy_files or str(path) in copy_files: + if os.path.isfile(path): + shutil.copy(path, dst_path) else: - shutil.copytree(path, dst_dir) - + shutil.copytree(path, dst_path) + else: + os.symlink(path, dst_path) def generate_work_path(create: bool = True) -> str: """ From 3e4ec5d069c940ee5814b918bb18f788f8370edc Mon Sep 17 00:00:00 2001 From: ahxbcn Date: Thu, 22 Jan 2026 16:13:42 +0800 Subject: [PATCH 09/15] New PDOS plot and write functions using PDOSData in dos.py --- src/abacusagent/modules/submodules/dos.py | 1105 ++++++++------------- 1 file changed, 414 insertions(+), 691 deletions(-) diff --git a/src/abacusagent/modules/submodules/dos.py b/src/abacusagent/modules/submodules/dos.py index 9133ca1..aa6d74f 100644 --- a/src/abacusagent/modules/submodules/dos.py +++ b/src/abacusagent/modules/submodules/dos.py @@ -7,45 +7,55 @@ from abacustest.lib_prepare.abacus import ReadInput, WriteInput from abacustest.lib_collectdata.collectdata import RESULT from abacustest.lib_model.comm import check_abacus_inputs +from abacustest.lib_model.comm_dos import DOSData, PDOSData, l_map, orbital_names from pathlib import Path from typing import Dict, Any, List, Literal, Optional, Tuple -from abacusagent.modules.util.comm import generate_work_path, link_abacusjob, run_abacus, has_chgfile, collect_metrics +from abacusagent.modules.util.comm import ( + generate_work_path, + link_abacusjob, + run_abacus, + has_chgfile, + collect_metrics, +) from abacusagent.modules.util.chemical_elements import MAX_ANGULAR_MOMENTUM_OF_ELEMENTS -angular_momentum_map = ['s', 'p', 'd', 'f', 'g'] +angular_momentum_map = ["s", "p", "d", "f", "g"] color_map = { - 's': '#FF5733', - 'p': '#33FF57', - 'd': '#3357FF', - 'f': '#F033FF', - 'g': '#33FFF0' + "s": "#FF5733", + "p": "#33FF57", + "d": "#3357FF", + "f": "#F033FF", + "g": "#33FFF0", } orbital_rep_map = { - 's': 's', - 'px': r'$p_x$', - 'py': r'$p_y$', - 'pz': r'$p_z$', - 'dz^2': r'$d_{z^2}$', - 'dxz': r'$d_{xz}$', - 'dyz': r'$d_{yz}$', - 'dxy': r'$d_{xy}$', - 'dx^2-y^2': r'$d_{x^2-y^2}$', - 'fz^3': r'$f_{z^3}$', - 'fxz^2': r'$f_{xz^2}$', - 'fyz^2': r'$f_{yz^2}$', - 'fzx^2-zy^2': r'$f_{zx^2-zy^2}$', - 'fxyz': r'$f_{xyz}$', - 'fx^3-3*xy^2': r'$f_{x^3-3xy^2}$', - 'f3yx^2-y^3': r'$f_{3yx^2-y^3}$' + "s": "s", + "px": r"$p_x$", + "py": r"$p_y$", + "pz": r"$p_z$", + "dz^2": r"$d_{z^2}$", + "dxz": r"$d_{xz}$", + "dyz": r"$d_{yz}$", + "dxy": r"$d_{xy}$", + "dx^2-y^2": r"$d_{x^2-y^2}$", + "fz^3": r"$f_{z^3}$", + "fxz^2": r"$f_{xz^2}$", + "fyz^2": r"$f_{yz^2}$", + "fzx^2-zy^2": r"$f_{zx^2-zy^2}$", + "fxyz": r"$f_{xyz}$", + "fx^3-3*xy^2": r"$f_{x^3-3xy^2}$", + "f3yx^2-y^3": r"$f_{3yx^2-y^3}$", } + def abacus_dos_run( abacus_inputs_dir: Path, - pdos_mode: Literal['atoms', 'species', 'species+shell', 'species+orbital'] = 'species+shell', + pdos_mode: Literal[ + "atoms", "species", "species+shell", "species+orbital" + ] = "species+shell", pdos_atom_indices: Optional[List[int]] = None, dos_edelta_ev: float = 0.01, dos_sigma: float = 0.07, @@ -53,12 +63,12 @@ def abacus_dos_run( dos_emax_ev: float = 10.0, ) -> Dict[str, Any]: """Run the DOS and PDOS calculation. - - This function will firstly run a SCF calculation with out_chg set to 1, + + This function will firstly run a SCF calculation with out_chg set to 1, then run a NSCF calculation with init_chg set to 'file' and out_dos set to 1 or 2. If the INPUT parameter "basis_type" is "PW", then out_dos will be set to 1, and only DOS will be calculated and plotted. If the INPUT parameter "basis_type" is "LCAO", then out_dos will be set to 2, and both DOS and PDOS will be calculated and plotted. - + Args: abacus_inputs_dir: Path to the ABACUS input files, which contains the INPUT, STRU, KPT, and pseudopotential or orbital files. pdos_mode: Mode of plotted PDOS file. @@ -68,10 +78,10 @@ def abacus_dos_run( - "species+orbital": Orbital-resolved PDOS will be plotted. PDOS of orbitals in the same shell of a species will be plotted in a subplot. pdos_atom_indices: A list of atom indices, only used if pdos_mode is "atoms". dos_edelta_ev: Step size in writing Density of States (DOS) in eV. - dos_sigma: Width of the Gaussian factor when obtaining smeared Density of States (DOS) in eV. + dos_sigma: Width of the Gaussian factor when obtaining smeared Density of States (DOS) in eV. dos_emin_ev: Minimal range for Density of States (DOS) in eV. Default is -10.0. dos_emax_ev: Maximal range for Density of States (DOS) in eV. Default is 10.0. - + Returns: Dict[str, Any]: A dictionary containing: - dos_fig_path: Path to the plotted DOS. @@ -90,34 +100,41 @@ def abacus_dos_run( is_valid, msg = check_abacus_inputs(abacus_inputs_dir) if not is_valid: raise RuntimeError(f"Invalid ABACUS input files: {msg}") - + input_file = os.path.join(abacus_inputs_dir, "INPUT") input_params = ReadInput(input_file) nspin = input_params.get("nspin", 1) if nspin in [4]: - raise ValueError("Currently DOS calculation can only be plotted using for nspin=1 and nspin=2") + raise ValueError( + "Currently DOS calculation can only be plotted using for nspin=1 and nspin=2" + ) + print("Performing SCF calculation...") metrics_scf = abacus_dos_run_scf(abacus_inputs_dir) - metrics_nscf = abacus_dos_run_nscf(metrics_scf["scf_work_path"], - dos_edelta_ev=dos_edelta_ev, - dos_sigma=dos_sigma) - - fig_paths, dos_pdos_data_paths = plot_dos_pdos(metrics_scf["scf_work_path"], - metrics_nscf["nscf_work_path"], - metrics_nscf["nscf_work_path"], - nspin, - pdos_mode, - pdos_atom_indices, - dos_emin_ev, - dos_emax_ev) - + + print("Performing NSCF calculation...") + metrics_nscf = abacus_dos_run_nscf( + metrics_scf["scf_work_path"], + dos_edelta_ev=dos_edelta_ev, + dos_sigma=dos_sigma, + ) + + fig_paths, dos_pdos_data_paths = plot_write_dos_pdos( + metrics_scf["scf_work_path"], + metrics_nscf["nscf_work_path"], + pdos_mode, + pdos_atom_indices, + dos_emin_ev, + dos_emax_ev, + ) + return_dict = {"dos_fig_path": fig_paths[0]} - return_dict['dos_data_path'] = dos_pdos_data_paths[0] + return_dict["dos_data_path"] = dos_pdos_data_paths[0] try: - return_dict['pdos_fig_path'] = fig_paths[1] - return_dict['pdos_data_paths'] = dos_pdos_data_paths[1:] + return_dict["pdos_fig_path"] = fig_paths[1] + return_dict["pdos_data_paths"] = dos_pdos_data_paths[1:] except: - pass # Do nothing if PDOS file is not plotted + pass # Do nothing if PDOS file is not plotted return_dict.update(metrics_scf) return_dict.update(metrics_nscf) @@ -125,23 +142,26 @@ def abacus_dos_run( return return_dict except Exception as e: import traceback + traceback.print_exc() return {"message": f"Calculating DOS and PDOS failed: {e}"} -def abacus_dos_run_scf(abacus_inputs_dir: Path, - force_run: bool = False) -> Dict[str, Any]: + +def abacus_dos_run_scf( + abacus_inputs_dir: Path, force_run: bool = False +) -> Dict[str, Any]: """ Run the SCF calculation to generate the charge density file. If the charge file already exists, it will skip the SCF calculation. - + Args: abacus_inputs_dir: Path to the ABACUS input files, which contains the INPUT, STRU, KPT, and pseudopotential or orbital files. force_run: If True, it will run the SCF calculation even if the charge file already exists. - + Returns: Dict[str, Any]: A dictionary containing the work path, normal end status, SCF steps, convergence status, and energies. """ - + input_param = ReadInput(os.path.join(abacus_inputs_dir, "INPUT")) # check if charge file has been generated if has_chgfile(abacus_inputs_dir) and not force_run: @@ -149,9 +169,7 @@ def abacus_dos_run_scf(abacus_inputs_dir: Path, work_path = abacus_inputs_dir else: work_path = generate_work_path() - link_abacusjob(src=abacus_inputs_dir, - dst=work_path, - copy_files=["INPUT"]) + link_abacusjob(src=abacus_inputs_dir, dst=work_path, copy_files=["INPUT"]) input_param = ReadInput(os.path.join(work_path, "INPUT")) input_param["calculation"] = "scf" @@ -161,686 +179,391 @@ def abacus_dos_run_scf(abacus_inputs_dir: Path, run_abacus(work_path) rs = RESULT(path=work_path, fmt="abacus") - + return { "scf_work_path": Path(work_path).absolute(), "scf_normal_end": rs["normal_end"], "scf_steps": rs["scf_steps"], "scf_converge": rs["converge"], - "scf_energy": rs["energy"] + "scf_energy": rs["energy"], } -def abacus_dos_run_nscf(abacus_inputs_dir: Path, - dos_edelta_ev: float = None, - dos_sigma: float = None) -> Dict[str, Any]: - + +def abacus_dos_run_nscf( + abacus_inputs_dir: Path, dos_edelta_ev: float = None, dos_sigma: float = None +) -> Dict[str, Any]: work_path = generate_work_path() - link_abacusjob(src=abacus_inputs_dir, - dst=work_path, - copy_files=["INPUT", "KPT"] + glob.glob(os.path.join(abacus_inputs_dir, "OUT.*")), - exclude=["*log", "*json"]) - + link_abacusjob( + src=abacus_inputs_dir, + dst=work_path, + copy_files=["INPUT", "KPT"] + + glob.glob(os.path.join(abacus_inputs_dir, "OUT.*")), + exclude=["*log", "*json"], + ) + input_param = ReadInput(os.path.join(work_path, "INPUT")) input_param["calculation"] = "nscf" input_param["init_chg"] = "file" input_param["out_chg"] = -1 if input_param.get("basis_type", "pw") == "lcao": - input_param["out_dos"] = 2 # only for LCAO basis, and will output DOS and PDOS + input_param["out_dos"] = 2 # only for LCAO basis, and will output DOS and PDOS else: input_param["out_dos"] = 1 - + for dos_param, value in { "dos_edelta_ev": dos_edelta_ev, "dos_sigma": dos_sigma, }.items(): if value is not None: input_param[dos_param] = value - - + WriteInput(input_param, os.path.join(work_path, "INPUT")) - + run_abacus(work_path) - + rs = RESULT(path=work_path, fmt="abacus") - + return { "nscf_work_path": Path(work_path).absolute(), - "nscf_normal_end": rs["normal_end"] + "nscf_normal_end": rs["normal_end"], } -def parse_pdos_file(file_path): - """Parse the PDOS file and extract energy values and orbital data.""" - with open(file_path, 'r') as f: - content = f.read() - - energy_match = re.search(r'(.*?)', content, re.DOTALL) - if not energy_match: - raise ValueError("Energy values not found in the file.") - - energy_text = energy_match.group(1) - energy_values = np.array([float(line.strip()) for line in energy_text.strip().split()]) - - orbital_pattern = re.compile(r'(.*?)', re.DOTALL) - orbitals = [] - - for match in orbital_pattern.finditer(content): - index, atom_index, species, l, m, z, orbital_content = match.groups() - - data_match = re.search(r'(.*?)', orbital_content, re.DOTALL) - if data_match: - data_text = data_match.group(1) - data_values = np.array([float(line.strip()) for line in data_text.strip().split()]) - - orbitals.append({ - 'index': int(index), - 'atom_index': int(atom_index), - 'species': species, - 'l': int(l), - 'm': int(m), - 'z': int(z), - 'data': data_values - }) - - return energy_values, orbitals -def parse_basref_file(file_path): - """Parse basref file to create mapping for custom labels.""" - label_map = {} - - with open(file_path, 'r') as f: - for line in f: - line = line.strip() - if not line or line.startswith('#'): - continue - - parts = line.split() - if len(parts) >= 6: - # Add 1 to atom_index as per requirement - atom_index = int(parts[0]) + 1 - species = parts[1] - l = int(parts[2]) - m = int(parts[3]) - z = int(parts[4]) - symbol = parts[5] - - key = (atom_index, species, l, m, z) - label_map[key] = f'{species}{atom_index}({symbol})' - - return label_map - -def write_pdos_data_species(energy: NDArray[np.float64], - species_pdos_data: Dict[str, NDArray[np.float64]], - nspin: Literal[1, 2] = 1, - output_dir: str = './') -> None: - """Write PDOS of species to file""" - pdos_data_file = Path(os.path.join(output_dir, "PDOS_species.dat")).absolute() - with open(pdos_data_file, 'w') as f: - if nspin == 1: - header_line = ' Energy (eV)' - for species_name in species_pdos_data.keys(): - header_line += f'{species_name:>20s}' - f.write(header_line + "\n") - for i in range(len(energy)): - line = f"{energy[i]:20.8f}" - for species_name in species_pdos_data.keys(): - line += f"{species_pdos_data[species_name][i]:20.8f}" - f.write(line + "\n") - elif nspin == 2: - header_spin_up, header_spin_down = '', '' - for species_name in species_pdos_data.keys(): - header_spin_up += f'{species_name:>17s}_up' - header_spin_down += f'{species_name:>17s}_dn' - f.write(f" Energy (eV){header_spin_up}{header_spin_down}\n") - for i in range(len(energy)): - line = f"{energy[i]:20.8f}" - line_spin_up, line_spin_down = '', '' - for species_name, pdos_data in species_pdos_data.items(): - # even index for spin up, odd index for spin down - line_spin_up += f"{pdos_data[i*2]:20.8f}" - line_spin_down += f"{pdos_data[i*2+1]:20.8f}" - line += line_spin_up + line_spin_down - f.write(line + "\n") - - return pdos_data_file.absolute() - -def write_pdos_data_species_shell(energy: NDArray[np.float64], - species_pdos_data_dict: Dict[str, NDArray[np.float64]], - species: str, - nspin: Literal[1, 2] = 1, - output_dir: str = './') -> None: - # Write PDOS of different shells (s, p, d...) for a species to file - pdos_data_file = Path(os.path.join(output_dir, f"PDOS_{species}.dat")).absolute() - with open(os.path.join(output_dir, f"PDOS_{species}.dat"), "w") as f: - if nspin == 1: - header_line = ' Energy (eV)' - for l in species_pdos_data_dict.keys(): - header_line += f'{l:>20s}' - f.write(header_line + "\n") - for i in range(len(energy)): - line = f"{energy[i]:14.6f}" - for l, pdos_data in species_pdos_data_dict.items(): - line += f"{pdos_data[i]:20.8f}" - f.write(line + "\n") - elif nspin == 2: - header_spin_up, header_spin_down = '', '' - for l in species_pdos_data_dict.keys(): - header_spin_up += f'{l:>17s}_up' - header_spin_down += f'{l:>17}_dn' - f.write(f" Energy (eV){header_spin_up}{header_spin_down}\n") - for i in range(len(energy)): - line = f"{energy[i]:14.6f}" - line_spin_up, line_spin_down = '', '' - for l, pdos_data in species_pdos_data_dict.items(): - # even index for spin up, odd index for spin down - line_spin_up += f"{pdos_data[i*2]:20.8f}" - line_spin_down += f"{pdos_data[i*2+1]:20.8f}" - f.write(line + line_spin_up + line_spin_down + "\n") - - return Path(pdos_data_file).absolute() - -def write_pdos_data_atom(energy: NDArray[np.float64], - atom_pdos: List[Dict[str, NDArray[np.float64]]], - atom_index: int, - label_map: Dict[Tuple[int, str, int, int, int], str] = {}, - nspin: Literal[1, 2] = 1, - output_dir: str = './') -> None: - # Initialize dict of pdos data for each orbital - orbital_label = {} - for (_, species, l, m, z), full_label in label_map.items(): - if str(l) not in orbital_label.keys(): - orbital_label[str(l)] = {} - if str(m) not in orbital_label[str(l)].keys(): - orbital_name = full_label.split('(')[1].split(')')[0] - if orbital_name in orbital_rep_map.keys(): - orbital_label[str(l)][str(m)] = orbital_rep_map[orbital_name] - else: - orbital_label[str(l)][str(m)] = orbital_name - else: - pass - - # Write PDOS data of an atom to file - atom_pdos_lms = {} - orbital_names = [] - for orbital in atom_pdos: - if orbital['l'] <= angular_momentum_map.index(MAX_ANGULAR_MOMENTUM_OF_ELEMENTS[orbital['species']]): - if orbital['l'] not in atom_pdos_lms.keys(): - atom_pdos_lms[orbital['l']] = {} - atom_pdos_lms[orbital['l']][orbital['m']] = orbital['data'] - orbital_names.append(orbital_label[str(orbital['l'])][str(orbital['m'])]) - else: - if orbital['m'] not in atom_pdos_lms[orbital['l']].keys(): - atom_pdos_lms[orbital['l']][orbital['m']] = orbital['data'] - orbital_names.append(orbital_label[str(orbital['l'])][str(orbital['m'])]) - else: - atom_pdos_lms[orbital['l']][orbital['m']] += orbital['data'] - - atom_species = atom_pdos[0]['species'] - pdos_data_file = Path(os.path.join(output_dir, f"PDOS_{atom_species}{atom_index}.dat")) - with open(pdos_data_file, "w") as f: - if nspin == 1: - header_line = ' Energy (eV)' - for orbital_name in orbital_names: - header_line += f'{orbital_name.replace("$", ""):>20s}' - f.write(header_line + "\n") - for i in range(len(energy)): - line = f"{energy[i]:14.6f}" - for l, pdos_data in atom_pdos_lms.items(): - for m, pdos_data_m in pdos_data.items(): - line += f"{pdos_data_m[i]:20.8f}" - f.write(line + "\n") - elif nspin == 2: - header_spin_up, header_spin_down = '', '' - for orbital_name in orbital_names: - header_spin_up += f'{orbital_name.replace("$", ""):>17s}_up' - header_spin_down += f'{orbital_name.replace("$", ""):>17}_dn' - f.write(f" Energy (eV){header_spin_up}{header_spin_down}\n") - for i in range(len(energy)): - line = f"{energy[i]:14.6f}" - line_spin_up, line_spin_down = '', '' - for l, pdos_data in atom_pdos_lms.items(): - for m, pdos_data_m in pdos_data.items(): - # even index for spin up, odd index for spin down - line_spin_up += f"{pdos_data_m[i*2]:20.8f}" - line_spin_down += f"{pdos_data_m[i*2+1]:20.8f}" - f.write(line + line_spin_up + line_spin_down + "\n") - - return pdos_data_file - -def write_pdos_data_species_orbital(energy: NDArray[np.float64], - species_pdos: Dict[str, NDArray[np.float64]], - species: str, - orbital_label: Dict[str, Dict[str, Dict[str, str]]], - nspin: Literal[1, 2] = 1, - output_dir: str = "./") -> None: - # Write PDOS data of different orbitals of a species to file - pdos_data_file = Path(os.path.join(output_dir, f"PDOS_{species}.dat")) - with open(pdos_data_file, "w") as f: - if nspin == 1: - header_line = f" Energy (eV)" - for l, species_shell_pdos in species_pdos.items(): - for m in species_shell_pdos.keys(): - orbital_name = orbital_label[species][str(angular_momentum_map.index(l))][str(m)] - header_line += f"{orbital_name.replace('$', ''):>16s}" - f.write(header_line+"\n") - for i in range(len(energy)): - line = f"{energy[i]:14.6f}" - for l, species_shell_pdos in species_pdos.items(): - for m in species_shell_pdos.keys(): - line += f"{species_pdos[l][m][i]:20.8f}" - f.write(line + "\n") - elif nspin == 2: - header_spin_up, header_spin_down = '', '' - for l, species_shell_pdos in species_pdos.items(): - for m in species_shell_pdos.keys(): - orbital_name = orbital_label[species][str(angular_momentum_map.index(l))][str(m)] - header_spin_up += f"{orbital_name.replace('$', ''):>17s}_up" - header_spin_down += f"{orbital_name.replace('$', ''):>17s}_dn" - f.write(f" Energy (eV){header_spin_up}{header_spin_down}\n") - for i in range(len(energy)): - line = f"{energy[i]:14.6f}" - line_spin_up, line_spin_down = '', '' - for l, species_shell_pdos in species_pdos.items(): - for m, species_orbital_pdos in species_shell_pdos.items(): - # even index for spin up, odd index for spin down - line_spin_up += f"{species_orbital_pdos[i*2]:20.8f}" - line_spin_down += f"{species_orbital_pdos[i*2+1]:20.8f}" - f.write(line + line_spin_up + line_spin_down + "\n") - - return pdos_data_file - -def plot_pdos(energy_values, orbitals, fermi_level, label_map, output_dir, nspin, mode, pdos_atom_indices, dos_emin_ev, dos_emax_ev, dpi=300): - """Plot PDOS data separated by atom/species with custom labels.""" - # Create output directory if it doesn't exist +def plot_write_pdos_species( + pdos_data: PDOSData, + output_dir: Path, + dos_emin_ev: float, + dos_emax_ev: float, +) -> Tuple[str, List[str]]: + """Plot PDOS by species using PDOSData class.""" + os.makedirs(output_dir, exist_ok=True) + + # Get unique species + species_set = set() + for orbital in pdos_data.projected_dos: + species_set.add(orbital["species"]) + + # Sum PDOS for each species + species_pdos_data = [] + species_labels = [] + for species in species_set: + species_pdos_data.append(pdos_data.get_pdos_by_species(species)) + species_labels.append(species) + + # Use PDOSData.plot_pdos method + pdos_pic_file = os.path.join(output_dir, "PDOS.png") + PDOSData.plot_pdos( + pdosdatas=[species_pdos_data], + labels=[species_labels], + titles=["Projected density of States of different species"], + energy=pdos_data.energy, + energy_min=dos_emin_ev, + energy_max=dos_emax_ev, + pdos_fig_name=pdos_pic_file, + ) + + # Write data file using PDOSData.write_pdos method + pdos_data_file = os.path.join(output_dir, "PDOS.dat") + PDOSData.write_pdos( + pdosdatas=species_pdos_data, + energy=pdos_data.energy, + labels=species_labels, + filename=str(pdos_data_file), + ) + + return pdos_pic_file, pdos_data_file + + +def plot_write_pdos_species_shell( + pdos_data: PDOSData, + output_dir: Path, + dos_emin_ev: float, + dos_emax_ev: float, +) -> Tuple[str, List[str]]: + """Plot PDOS by species and shell using PDOSData class.""" os.makedirs(output_dir, exist_ok=True) - # Shift energy values by Fermi level - shifted_energy = energy_values - fermi_level - - # Group orbitals by atom_index and species - atom_species_groups = {} - for orbital in orbitals: - key = (orbital['atom_index'], orbital['species']) - if key not in atom_species_groups: - atom_species_groups[key] = [] - atom_species_groups[key].append(orbital) - - if mode == "species": - pdos_pic_file, pdos_data_files = plot_pdos_species(shifted_energy, orbitals, output_dir, nspin, dos_emin_ev, dos_emax_ev, dpi) - elif mode == "species+shell": - pdos_pic_file, pdos_data_files = plot_pdos_species_shell(shifted_energy, orbitals, output_dir, nspin, dos_emin_ev, dos_emax_ev, dpi) - elif mode == "species+orbital": - pdos_pic_file, pdos_data_files = plot_pdos_species_orbital(shifted_energy, orbitals, output_dir, nspin, label_map, dos_emin_ev, dos_emax_ev, dpi) - elif mode == "atoms": - pdos_pic_file, pdos_data_files = plot_pdos_atoms(shifted_energy, orbitals, output_dir, nspin, label_map, pdos_atom_indices, dos_emin_ev, dos_emax_ev, dpi) - else: - raise ValueError(f"Not allowed mode {mode}") - - return pdos_pic_file, pdos_data_files - -def plot_pdos_species(shifted_energy, orbitals, output_dir, nspin, dos_emin_ev, dos_emax_ev, dpi): - species = {} - for orbital in orbitals: - species_one = orbital['species'] - if species_one not in species.keys(): - species[species_one] = orbital['data'] - else: - species[species_one] += orbital['data'] - - pdos_data_files = [write_pdos_data_species(shifted_energy, species, nspin, output_dir)] - # Plot PDOS of different species - plt.plot(figsize=(10, 6)) - for species_name, pdos_data in species.items(): - if nspin == 1: - plt.plot(shifted_energy, pdos_data, label=species_name, linewidth=1.0) - elif nspin == 2: - plt.plot(shifted_energy, pdos_data[::2], label=f'{species_name} ' + r'$\uparrow$', linestyle='-', linewidth=1.0) - plt.plot(shifted_energy, -pdos_data[1::2], label=f'{species_name} ' + r'$\downarrow$', linestyle='--', linewidth=1.0) - - plt.axvline(x=0, color='black', linestyle=':', linewidth=1.0) - plt.xlabel('Energy (eV)', fontsize=10) - plt.ylabel(r"States ($eV^{-1}$)", fontsize=10) - plt.xlim(dos_emin_ev, dos_emax_ev) - if nspin == 1: - plt.ylim(bottom=0) - plt.legend(fontsize=8, ncol=nspin) - plt.grid(alpha=0.3) - plt.title('Projected density of States of different species') - - pdos_pic_file = os.path.join(output_dir, 'PDOS.png') - plt.savefig(pdos_pic_file, dpi=dpi) - plt.close() - - return Path(pdos_pic_file).absolute(), pdos_data_files - -def plot_pdos_species_shell(shifted_energy, orbitals, output_dir, nspin, dos_emin_ev, dos_emax_ev, dpi): - species_shells = {} - for orbital in orbitals: - species = orbital['species'] - if species not in species_shells.keys(): - species_shells[species] = {} # Initialize species kind - - angular_momentum = angular_momentum_map[orbital['l']] - # The orbital with higher angular momentum than in realistic atoms will be ignored. - if angular_momentum_map.index(angular_momentum) <= angular_momentum_map.index(MAX_ANGULAR_MOMENTUM_OF_ELEMENTS[orbital['species']]): - if angular_momentum not in species_shells[species].keys(): - species_shells[species][angular_momentum] = orbital['data'] # Initialize DOS for angular momentum of a species - else: - species_shells[species][angular_momentum] += orbital['data'] # Add DOS of a angular momentum of a species - - # Plot PDOS for each species and each shell - num_species = len(species_shells) - fig, axes = plt.subplots(nrows=num_species, ncols=1, figsize=(8, 4*num_species)) - if num_species == 1: - axes = [axes] - - pdos_data_files = [] - for species_idx, (species, pdos_data_dict) in enumerate(species_shells.items()): - # Write PDOS data to the selected species to file - pdos_data_file = write_pdos_data_species_shell(shifted_energy, pdos_data_dict, species, nspin, output_dir) - pdos_data_files.append(pdos_data_file) - # Plot PDOS of different shells for a species - ax = axes[species_idx] - - for l, pdos_data in pdos_data_dict.items(): - if nspin == 1: - ax.plot(shifted_energy, pdos_data, color=color_map[l], label=f'{species}-{l}', linewidth=1.0) - elif nspin == 2: - ax.plot(shifted_energy, pdos_data[::2], color=color_map[l], label=f'{species}-{l}' + r' $\uparrow$', linestyle='-', linewidth=1.0) - ax.plot(shifted_energy, -pdos_data[1::2], color=color_map[l], label=f'{species}-{l}' + r' $\downarrow$', linestyle='--', linewidth=1.0) - - ax.axvline(x=0, color='black', linestyle=':', linewidth=1.0) - ax.set_title(f'PDOS for {species}', fontsize=12, pad=10) - ax.set_ylabel(r"States ($eV^{-1}$)", fontsize=10) - ax.set_xlim(dos_emin_ev, dos_emax_ev) - if nspin == 1: - ax.set_ylim(bottom=0) - ax.legend(fontsize=8, ncol=nspin) - ax.grid(alpha=0.3) - - axes[-1].set_xlabel('Energy (eV)', fontsize=10) - - plt.tight_layout() - pdos_pic_file = os.path.join(output_dir, 'PDOS.png') - plt.savefig(pdos_pic_file, dpi=dpi, bbox_inches='tight') - plt.close() - - return Path(pdos_pic_file).absolute(), pdos_data_files - -def plot_pdos_species_orbital(shifted_energy, orbitals, output_dir, nspin, label_map, dos_emin_ev, dos_emax_ev, dpi): - - plt.rcParams["text.usetex"] = False - plt.rcParams["axes.prop_cycle"] = plt.cycler("color", plt.cm.tab20.colors) - - # Initialize dict of pdos data for each orbital - orbital_label = {} - for (atom_index, species, l, m, z), full_label in label_map.items(): - if species not in orbital_label.keys(): - orbital_label[species] = {} - if str(l) not in orbital_label[species].keys(): - orbital_label[species][str(l)] = {} - if str(m) not in orbital_label[species][str(l)].keys(): - orbital_name = full_label.split('(')[1].split(')')[0] - if orbital_name in orbital_rep_map.keys(): - orbital_label[species][str(l)][str(m)] = orbital_rep_map[orbital_name] - else: - orbital_label[species][str(l)][str(m)] = orbital_name - else: - pass - - # Sum over zeta to get the pdos of orbitals - species_orbitals = {} - for orbital in orbitals: - species = orbital['species'] - if species not in species_orbitals.keys(): - species_orbitals[species] = {} - - angular_momentum = angular_momentum_map[orbital['l']] - # The orbital with higher angular momentum than in realistic atoms will be ignored. - if angular_momentum_map.index(angular_momentum) <= angular_momentum_map.index(MAX_ANGULAR_MOMENTUM_OF_ELEMENTS[orbital['species']]): - if angular_momentum not in species_orbitals[species].keys(): - species_orbitals[species][angular_momentum] = {} - - mag_quantum_num = orbital['m'] - if mag_quantum_num not in species_orbitals[species][angular_momentum].keys(): - species_orbitals[species][angular_momentum][mag_quantum_num] = orbital['data'] - else: - species_orbitals[species][angular_momentum][mag_quantum_num] += orbital['data'] - - # Plot PDOS for each species and each orbital - total_subplots = 0 - for species, species_pdos in species_orbitals.items(): - total_subplots += len(species_pdos) - fig, axes = plt.subplots(nrows=total_subplots, ncols=1, figsize=(8, 4*total_subplots)) - - subplot_count = 0 - pdos_data_files = [] - for species, species_pdos in species_orbitals.items(): - # Write PDOS data of orbitals of the selected species to file - pdos_data_file = write_pdos_data_species_orbital(shifted_energy, species_pdos, species, orbital_label, nspin, output_dir) - pdos_data_files.append(pdos_data_file) - for angular_momentum, species_shell_pdos in species_pdos.items(): - for m, species_orbital_pdos in species_shell_pdos.items(): - ax = axes[subplot_count] - orbital_name = orbital_label[species][str(angular_momentum_map.index(angular_momentum))][str(m)] - if nspin == 1: - ax.plot(shifted_energy, species_orbital_pdos, label=f'{orbital_name}', linewidth=1.0) - elif nspin == 2: - ax.plot(shifted_energy, species_orbital_pdos[::2], label=f'{orbital_name} '+r'$\uparrow$', linestyle='-', linewidth=1.0) - ax.plot(shifted_energy, -species_orbital_pdos[1::2], label=f'{orbital_name} '+r'$\downarrow$', linestyle='--', linewidth=1.0) - - ax.axvline(x=0, color='black', linestyle=':', linewidth=1.0) - ax.set_title(f'PDOS for {species}-{angular_momentum}', fontsize=12, pad=10) - ax.set_xlim(dos_emin_ev, dos_emax_ev) - if nspin == 1: - ax.set_ylim(bottom=0) - ax.set_ylabel(r"States ($eV^{-1}$)", fontsize=10) - ax.legend(fontsize=8, ncol=nspin) - ax.grid(alpha=0.3) - - subplot_count += 1 - - axes[-1].set_xlabel('Energy (eV)', fontsize=10) - - plt.tight_layout() - pdos_pic_file = os.path.join(output_dir, 'PDOS.png') - plt.savefig(pdos_pic_file, dpi=dpi, bbox_inches='tight') - plt.close() - - return Path(pdos_pic_file).absolute(), pdos_data_files - -def plot_pdos_atoms(shifted_energy, orbitals, output_dir, nspin, label_map, pdos_atom_indices, dos_emin_ev, dos_emax_ev, dpi): - # Plot PDOS for selected atoms - atom_pdoses = {} - for orbital in orbitals: - atom_index = orbital['atom_index'] - if atom_index in pdos_atom_indices: - if atom_index not in atom_pdoses.keys(): - atom_pdoses[atom_index] = [orbital] - else: - atom_pdoses[atom_index].append(orbital) - - fig, axes = plt.subplots(nrows=len(pdos_atom_indices), ncols=1, figsize=(8, 4*len(pdos_atom_indices))) - if len(pdos_atom_indices) == 1: - axes = [axes] - - pdos_data_files = [] - for atom_index, atom_pdos in atom_pdoses.items(): - # Write PDOS data of all legal orbitals for each selected atom to file - pdos_data_file = write_pdos_data_atom(shifted_energy, atom_pdos, atom_index, label_map, nspin, output_dir) - pdos_data_files.append(pdos_data_file) - # Plot PDOS of all legal shells for each selected atom - atom_pdos_ls = {} - for orbital in atom_pdos: - orbital_l_name = angular_momentum_map[orbital['l']] - if orbital_l_name not in atom_pdos_ls.keys(): - atom_pdos_ls[orbital_l_name] = orbital['data'] - else: - atom_pdos_ls[orbital_l_name] += orbital['data'] + # Get unique species + species_set = set() + for orbital in pdos_data.projected_dos: + species_set.add(orbital["species"]) + + # Prepare data for PDOSData.plot_pdos + all_pdos_datas, all_labels, all_titles = [], [], [] + + for species in sorted(species_set): + species_pdos_data = [] + species_labels = [] + # Get shells for this species + species_shells = [] + for orbital in pdos_data.projected_dos: + if orbital["species"] == species: + l = orbital["l"] + if l not in species_shells: + species_shells.append(l) + species_shell_pdos = pdos_data.get_pdos_by_species_shell(species, l) + species_pdos_data.append(species_shell_pdos) + + shell_label = f"{species}-{l_map[l]}" + species_labels.append(shell_label) - species = atom_pdos[0]['species'] - ax = axes[pdos_atom_indices.index(atom_index)] - for l, pdos_data in atom_pdos_ls.items(): - if angular_momentum_map.index(l) <= angular_momentum_map.index(MAX_ANGULAR_MOMENTUM_OF_ELEMENTS[species]): - if nspin == 1: - ax.plot(shifted_energy, pdos_data, label=f'{species}{atom_index}-{l}', linewidth=1.0) - elif nspin == 2: - ax.plot(shifted_energy, pdos_data[::2], color=color_map[l], label=f'{species}{atom_index}-{l} ' + r'$\uparrow$', linestyle='-', linewidth=1.0) - ax.plot(shifted_energy, -pdos_data[1::2], color=color_map[l], label=f'{species}{atom_index}-{l} ' + r'$\downarrow$', linestyle='--', linewidth=1.0) + all_titles.append(f"PDOS of {species}") + all_pdos_datas.append(species_pdos_data) + all_labels.append(species_labels) + + # Use PDOSData.plot_pdos method + pdos_pic_file = os.path.join(output_dir, "PDOS.png") + PDOSData.plot_pdos( + pdosdatas=all_pdos_datas, + labels=all_labels, + titles=all_titles, + energy=pdos_data.energy, + energy_min=dos_emin_ev, + energy_max=dos_emax_ev, + pdos_fig_name=pdos_pic_file, + ) + + all_pdos_data_flattened = [data for species_pdos_datas in all_pdos_datas for data in species_pdos_datas] + all_labels_flattened = [label for species_labels in all_labels for label in species_labels] + pdos_data_file = os.path.join(output_dir, "PDOS.dat") + PDOSData.write_pdos( + pdosdatas=all_pdos_data_flattened, + energy=pdos_data.energy, + labels=all_labels_flattened, + filename=pdos_data_file, + ) + + return pdos_pic_file, pdos_data_file + + +def plot_write_pdos_species_orbital( + pdos_data: PDOSData, + output_dir: Path, + dos_emin_ev: float, + dos_emax_ev: float, +) -> Tuple[str, List[str]]: + """Plot PDOS by species and orbital using PDOSData class.""" + os.makedirs(output_dir, exist_ok=True) + + # Get unique species + species_shell_set = set() + for orbital in pdos_data.projected_dos: + species_shell_set.add((orbital["species"], orbital["l"])) + + # Prepare data for PDOSData.plot_pdos + all_pdos_datas, all_labels, all_titles = [], [], [] + + for (species, l) in sorted(species_shell_set): + orbital_pdos_data = [] + species_orbital_labels = [] + # Get orbitals for this shell + species_orbitals = [] + for orbital in pdos_data.projected_dos: + if orbital["species"] == species and orbital["l"] == l: + m = orbital["m"] + if m not in species_orbitals: + species_orbitals.append(m) + species_orbital_pdos = pdos_data.get_pdos_by_species_orbital(species, l, m) + orbital_pdos_data.append(species_orbital_pdos) + + orbital_label = f"{species}-{orbital_names[(l, m)]}" + species_orbital_labels.append(orbital_label) - ax.axvline(x=0, color='black', linestyle=':', linewidth=1.0) - ax.set_title(f'PDOS for {species}{atom_index}', fontsize=12, pad=10) - ax.set_ylabel(r"States ($eV^{-1}$)", fontsize=10) - ax.set_xlim(dos_emin_ev, dos_emax_ev) - if nspin == 1: - ax.set_ylim(bottom=0) - ax.legend(fontsize=8, ncol=nspin) - ax.grid(alpha=0.3) - - axes[-1].set_xlabel('Energy (eV)', fontsize=10) - - plt.tight_layout() - pdos_pic_file = os.path.join(output_dir, 'PDOS.png') - plt.savefig(pdos_pic_file, dpi=dpi, bbox_inches='tight') - plt.close() - - return Path(pdos_pic_file).absolute(), pdos_data_files - -def write_dos_data(energy: NDArray[np.float64], - dos: NDArray[np.float64], - dos_dn: Optional[NDArray[np.float64]] = None, - nspin: Literal[1, 2] = 1, - out_filename: str = 'DOS_shifted.dat') -> None: - """Write DOS data to file.""" - with open(out_filename, "w") as f: - if nspin == 1: - f.write(" Energy (eV) DOS\n") - for i in range(len(energy)): - f.write(f"{energy[i]:16.6f}{dos[i]:16.6f}\n") - elif nspin == 2: - f.write(" Energy (eV) DOS_up DOS_down\n") - for i in range(len(energy)): - f.write(f"{energy[i]:16.6f}{dos[i]:16.6f}{dos_dn[i]:16.6f}\n") - f.close() - - return Path(out_filename).absolute() - -def plot_dos(file_path: List[Path], - fermi_level: float, - dos_plot_file: str = 'DOS.png', - nspin: Literal[1, 2] = 1, - dos_emin_ev: float = -10.0, - dos_emax_ev: float = 10.0, - dpi: int=300): - """Plot total DOS from DOS1_smearing.dat and DOS2_smearing (if nspin=2) file.""" - # Read first two columns from file - data = np.loadtxt(file_path[0], usecols=(0, 1)) - energy = data[:, 0] - fermi_level # Shift by Fermi level - dos = data[:, 1] - if nspin == 2: - data = np.loadtxt(file_path[1], usecols=(0, 1)) - dos_dn = data[:, 1] - - # Write shifted dos data to file - dos_output_file = os.path.join(os.path.dirname(dos_plot_file), "DOS_shifted.dat") - if nspin == 1: - dos_data_file = write_dos_data(energy, dos, nspin=nspin, out_filename=dos_output_file) - elif nspin == 2: - dos_data_file = write_dos_data(energy, dos, dos_dn, nspin=nspin, out_filename=dos_output_file) - # Create plot - plt.figure(figsize=(8, 6)) - if nspin == 1: - plt.plot(energy, dos, linestyle='-') - elif nspin == 2: - plt.plot(energy, dos, linestyle='-', label='spin up') - plt.plot(energy, -dos_dn, linestyle='--', label='spin down') - plt.axvline(x=0, color='k', linestyle=':', alpha=0.5) - plt.xlabel('Energy (eV)') - plt.ylabel(r'States ($eV^{-1}$)') - plt.title('Density of States') - plt.grid(True, alpha=0.3) - plt.xlim(dos_emin_ev, dos_emax_ev) - - # Save plot - os.makedirs(os.path.dirname(dos_plot_file), exist_ok=True) - plt.savefig(dos_plot_file, dpi=dpi, bbox_inches='tight') - plt.close() - - return Path(dos_plot_file).absolute(), Path(dos_data_file).absolute() - -def plot_dos_pdos(scf_job_path: Path, - nscf_job_path: Path, - output_dir: Path, - nspin: Literal[1, 2] = 1, - mode: Literal['species', 'species+shell', 'species+orbital'] = 'species+shell', - pdos_atom_indices: Optional[List[int]] = None, - dos_emin_ev: float = -10.0, - dos_emax_ev: float = 10.0, - dpi=300) -> List[str]: - """Plot DOS and PDOS from the NSCF job path. - + all_titles.append(f"PDOS of {species}-{l_map[l]}") + all_pdos_datas.append(orbital_pdos_data) + all_labels.append(species_orbital_labels) + + # Use PDOSData.plot_pdos method + pdos_pic_file = os.path.join(output_dir, "PDOS.png") + PDOSData.plot_pdos( + pdosdatas=all_pdos_datas, + labels=all_labels, + titles=all_titles, + energy=pdos_data.energy, + energy_min=dos_emin_ev, + energy_max=dos_emax_ev, + pdos_fig_name=pdos_pic_file, + ) + + all_pdos_data_flattened = [data for species_pdos_datas in all_pdos_datas for data in species_pdos_datas] + all_labels_flattened = [label.replace("$", "") for species_labels in all_labels for label in species_labels] + pdos_data_file = os.path.join(output_dir, "PDOS.dat") + PDOSData.write_pdos( + pdosdatas=all_pdos_data_flattened, + energy=pdos_data.energy, + labels=all_labels_flattened, + filename=pdos_data_file, + ) + + return pdos_pic_file, pdos_data_file + + +def plot_write_pdos_atoms( + pdos_data: PDOSData, + output_dir: Path, + pdos_atom_indices: List[int], + dos_emin_ev: float, + dos_emax_ev: float, +) -> Tuple[str, List[str]]: + """Plot PDOS for selected atoms using PDOSData class.""" + os.makedirs(output_dir, exist_ok=True) + + # Prepare data for PDOSData.plot_pdos + all_pdos_datas, all_labels, all_titles = [], [], [] + + for atom_index in pdos_atom_indices: + # Obtain all shells of the selected atom + atom_shell_set = set() + for orbital in pdos_data.projected_dos: + if orbital["atom_index"] == atom_index: + species = orbital["species"] + atom_shell_set.add((orbital["l"])) + + for l in atom_shell_set: + orbital_pdos_data = [] + species_orbital_labels = [] + # Get shells for this species + species_orbitals = [] + for orbital in pdos_data.projected_dos: + if orbital["species"] == species and orbital["l"] == l: + m = orbital["m"] + if m not in species_orbitals: + species_orbitals.append(m) + species_orbital_pdos = pdos_data.get_pdos_by_species_orbital(species, l, m) + orbital_pdos_data.append(species_orbital_pdos) + + orbital_label = f"{species}{atom_index}-{orbital_names[(l, m)]}" + species_orbital_labels.append(orbital_label) + + all_titles.append(f"PDOS of {species}{atom_index}-{l_map[l]}") + all_pdos_datas.append(orbital_pdos_data) + all_labels.append(species_orbital_labels) + + # Use PDOSData.plot_pdos method + pdos_pic_file = os.path.join(output_dir, "PDOS.png") + PDOSData.plot_pdos( + pdosdatas=all_pdos_datas, + labels=all_labels, + titles=all_titles, + energy=pdos_data.energy, + energy_min=dos_emin_ev, + energy_max=dos_emax_ev, + pdos_fig_name=pdos_pic_file, + ) + + all_pdos_data_flattened = [data for species_pdos_datas in all_pdos_datas for data in species_pdos_datas] + all_labels_flattened = [label.replace("$", "") for species_labels in all_labels for label in species_labels] + pdos_data_file = os.path.join(output_dir, "PDOS.dat") + PDOSData.write_pdos( + pdosdatas=all_pdos_data_flattened, + energy=pdos_data.energy, + labels=all_labels_flattened, + filename=pdos_data_file, + ) + + return pdos_pic_file, pdos_data_file + +def plot_write_dos_pdos( + scf_job_path: Path, + nscf_job_path: Path, + mode: Literal[ + "species", "species+shell", "species+orbital", "atoms" + ] = "species+shell", + pdos_atom_indices: Optional[List[int]] = None, + dos_emin_ev: float = -10.0, + dos_emax_ev: float = 5.0, +) -> Tuple[List[str], List[str]]: + """Plot DOS and PDOS from the NSCF job path using PDOSData class. + Args: - nscf_job_path (Path): Path to the NSCF job directory containing the OUT.* files. - output_dir (Path): Directory where the output plots will be saved. - dpi (int): Dots per inch for the saved plots. - + scf_job_path (Path): Path to the SCF job directory of the DOS calculation + nscf_job_path (Path): Path to the NSCF job directory of the DOS calculation + mode (str): PDOS plotting mode ('species', 'species+shell', 'species+orbital', or 'atoms'). + pdos_atom_indices (List[int], optional): List of atom indices for atom-specific PDOS. Only valid for 'atoms' mode. + dos_emin_ev (float): Minimum energy for DOS and PDOS plots. + dos_emax_ev (float): Maximum energy for DOS and PDOS plots. + Returns: - List[str]: List of paths to the generated plot files. - + Tuple[List[str], List[str]]: Tuple containing list of plot file paths and data file paths. """ - input_param = ReadInput(os.path.join(nscf_job_path, "INPUT")) - input_dir = os.path.join(nscf_job_path, "OUT." + input_param.get("suffix","ABACUS")) - basis_type = input_param.get('basis_type', 'pw') - - # Construct file paths based on input directory - pdos_file = os.path.join(input_dir, "PDOS") - log_file = os.path.join(input_dir, "running_nscf.log") - basref_file = os.path.join(input_dir, "Orbital") - dos_file = [os.path.join(input_dir, "DOS1_smearing.dat")] - dos_output = os.path.join(output_dir, "DOS.png") - if nspin == 2: - dos_file += [os.path.join(input_dir, "DOS2_smearing.dat")] + work_path = generate_work_path() - # Validate input files exist - for file_path in [log_file, dos_file[0]]: - if not os.path.exists(file_path): - print(f"Error: File not found - {file_path}") - raise FileNotFoundError(f"Required file not found: {file_path}") - if nspin == 2: - if not os.path.exists(dos_file[1]): - print(f"Error: File not found - {dos_file[1]}") - raise FileNotFoundError(f"Required file not found: {dos_file[1]}") + input_param = ReadInput(os.path.join(nscf_job_path, "INPUT")) + basis_type = input_param.get("basis_type", "pw") + results = RESULT(fmt="abacus", path=scf_job_path) + efermi = results['efermi'] + + # Construct file paths + dos_plot_file = os.path.join(work_path, "DOS.png") + dos_data_file = os.path.join(work_path, "DOS.dat") + + dosdata = DOSData.ReadFromAbacusJob(str(nscf_job_path), efermi) + DOSData.plot_dos( + dosdata.dosdata, + dosdata.energy, + dos_emin_ev, + dos_emax_ev, + "Density of States", + dos_plot_file, + ) + DOSData.write_dos(dosdata.dosdata, dosdata.energy, dos_data_file) + + all_plot_files = [Path(dos_plot_file).absolute()] + dos_pdos_data_files = [Path(dos_data_file).absolute()] - fermi_level = collect_metrics(scf_job_path, ['efermi'])['efermi'] - - # Plot DOS and get file path - dos_plot_file, dos_data_file = plot_dos(dos_file, fermi_level, dos_output, nspin, dos_emin_ev, dos_emax_ev, dpi) - all_plot_files = [dos_plot_file] - dos_pdos_data_files = [dos_data_file] - print("DOS file plotted") - # Plot PDOS (only for LCAO basis - PW basis doesn't support PDOS in ABACUS LTSv3.10) - if os.path.exists(pdos_file) and os.path.exists(basref_file): - if basis_type != 'pw': - label_map = parse_basref_file(basref_file) - energy_values, orbitals = parse_pdos_file(pdos_file) - pdos_plot_file, pdos_data_files = plot_pdos(energy_values, orbitals, fermi_level, label_map, output_dir, nspin, mode, pdos_atom_indices, dos_emin_ev, dos_emax_ev, dpi) - + # Plot PDOS using PDOSData class (only for LCAO basis) + if basis_type != "pw": + try: + # Load PDOS data using PDOSData class + pdos_data = PDOSData.ReadFromAbacusJob(str(nscf_job_path), efermi) + + # Plot PDOS based on mode + if mode == "species": + pdos_plot_file, pdos_data_file = plot_write_pdos_species( + pdos_data, work_path, dos_emin_ev, dos_emax_ev + ) + elif mode == "species+shell": + pdos_plot_file, pdos_data_file = plot_write_pdos_species_shell( + pdos_data, work_path, dos_emin_ev, dos_emax_ev + ) + elif mode == "species+orbital": + pdos_plot_file, pdos_data_file = plot_write_pdos_species_orbital( + pdos_data, work_path, dos_emin_ev, dos_emax_ev + ) + elif mode == "atoms": + if pdos_atom_indices is None or len(pdos_atom_indices) == 0: + raise ValueError( + "For 'atoms' mode, pdos_atom_indices must be provided" + ) + pdos_plot_file, pdos_data_file = plot_write_pdos_atoms( + pdos_data, + work_path, + pdos_atom_indices, + dos_emin_ev, + dos_emax_ev, + ) + else: + raise ValueError(f"Unsupported mode: {mode}") + # Combine file paths into a single list - all_plot_files.append(pdos_plot_file) - dos_pdos_data_files += pdos_data_files - else: - print(f"Warning: PDOS calculation not supported for PW basis type, skipping PDOS plotting") - elif os.path.exists(pdos_file) and not os.path.exists(basref_file): - print(f"Warning: PDOS file exists but Orbital file not found - {basref_file}, skipping PDOS plotting") - elif not os.path.exists(pdos_file) and os.path.exists(basref_file): - print(f"Warning: Orbital file exists but PDOS file not found - {pdos_file}, skipping PDOS plotting") + all_plot_files.append(Path(pdos_plot_file).absolute()) + dos_pdos_data_files.append(Path(pdos_data_file).absolute()) + + except Exception as e: + import traceback + traceback.print_exc() + print(f"Warning: Failed to plot PDOS: {e}") + print("Skipping PDOS plotting") else: - print("Warning: Both PDOS and Orbital files not found, skipping PDOS plotting") + print( + f"Warning: PDOS calculation not supported for PW basis type, skipping PDOS plotting" + ) print("Plots generated:") for file in all_plot_files: print(f"- {file}") - + return all_plot_files, dos_pdos_data_files From 17ee6a25c288eb105dee41c2ce69aca17e85001d Mon Sep 17 00:00:00 2001 From: ahxbcn Date: Fri, 23 Jan 2026 10:56:39 +0800 Subject: [PATCH 10/15] Use method in PDOSData to plot and write DOS and PDOS --- src/abacusagent/modules/submodules/dos.py | 313 ++-------------------- 1 file changed, 20 insertions(+), 293 deletions(-) diff --git a/src/abacusagent/modules/submodules/dos.py b/src/abacusagent/modules/submodules/dos.py index aa6d74f..10ac025 100644 --- a/src/abacusagent/modules/submodules/dos.py +++ b/src/abacusagent/modules/submodules/dos.py @@ -1,9 +1,5 @@ import os -import re import glob -import numpy as np -from numpy.typing import NDArray -import matplotlib.pyplot as plt from abacustest.lib_prepare.abacus import ReadInput, WriteInput from abacustest.lib_collectdata.collectdata import RESULT from abacustest.lib_model.comm import check_abacus_inputs @@ -17,38 +13,8 @@ link_abacusjob, run_abacus, has_chgfile, - collect_metrics, ) -from abacusagent.modules.util.chemical_elements import MAX_ANGULAR_MOMENTUM_OF_ELEMENTS - - -angular_momentum_map = ["s", "p", "d", "f", "g"] -color_map = { - "s": "#FF5733", - "p": "#33FF57", - "d": "#3357FF", - "f": "#F033FF", - "g": "#33FFF0", -} - -orbital_rep_map = { - "s": "s", - "px": r"$p_x$", - "py": r"$p_y$", - "pz": r"$p_z$", - "dz^2": r"$d_{z^2}$", - "dxz": r"$d_{xz}$", - "dyz": r"$d_{yz}$", - "dxy": r"$d_{xy}$", - "dx^2-y^2": r"$d_{x^2-y^2}$", - "fz^3": r"$f_{z^3}$", - "fxz^2": r"$f_{xz^2}$", - "fyz^2": r"$f_{yz^2}$", - "fzx^2-zy^2": r"$f_{zx^2-zy^2}$", - "fxyz": r"$f_{xyz}$", - "fx^3-3*xy^2": r"$f_{x^3-3xy^2}$", - "f3yx^2-y^3": r"$f_{3yx^2-y^3}$", -} + def abacus_dos_run( @@ -228,241 +194,6 @@ def abacus_dos_run_nscf( "nscf_normal_end": rs["normal_end"], } - -def plot_write_pdos_species( - pdos_data: PDOSData, - output_dir: Path, - dos_emin_ev: float, - dos_emax_ev: float, -) -> Tuple[str, List[str]]: - """Plot PDOS by species using PDOSData class.""" - os.makedirs(output_dir, exist_ok=True) - - # Get unique species - species_set = set() - for orbital in pdos_data.projected_dos: - species_set.add(orbital["species"]) - - # Sum PDOS for each species - species_pdos_data = [] - species_labels = [] - for species in species_set: - species_pdos_data.append(pdos_data.get_pdos_by_species(species)) - species_labels.append(species) - - # Use PDOSData.plot_pdos method - pdos_pic_file = os.path.join(output_dir, "PDOS.png") - PDOSData.plot_pdos( - pdosdatas=[species_pdos_data], - labels=[species_labels], - titles=["Projected density of States of different species"], - energy=pdos_data.energy, - energy_min=dos_emin_ev, - energy_max=dos_emax_ev, - pdos_fig_name=pdos_pic_file, - ) - - # Write data file using PDOSData.write_pdos method - pdos_data_file = os.path.join(output_dir, "PDOS.dat") - PDOSData.write_pdos( - pdosdatas=species_pdos_data, - energy=pdos_data.energy, - labels=species_labels, - filename=str(pdos_data_file), - ) - - return pdos_pic_file, pdos_data_file - - -def plot_write_pdos_species_shell( - pdos_data: PDOSData, - output_dir: Path, - dos_emin_ev: float, - dos_emax_ev: float, -) -> Tuple[str, List[str]]: - """Plot PDOS by species and shell using PDOSData class.""" - os.makedirs(output_dir, exist_ok=True) - - # Get unique species - species_set = set() - for orbital in pdos_data.projected_dos: - species_set.add(orbital["species"]) - - # Prepare data for PDOSData.plot_pdos - all_pdos_datas, all_labels, all_titles = [], [], [] - - for species in sorted(species_set): - species_pdos_data = [] - species_labels = [] - # Get shells for this species - species_shells = [] - for orbital in pdos_data.projected_dos: - if orbital["species"] == species: - l = orbital["l"] - if l not in species_shells: - species_shells.append(l) - species_shell_pdos = pdos_data.get_pdos_by_species_shell(species, l) - species_pdos_data.append(species_shell_pdos) - - shell_label = f"{species}-{l_map[l]}" - species_labels.append(shell_label) - - all_titles.append(f"PDOS of {species}") - all_pdos_datas.append(species_pdos_data) - all_labels.append(species_labels) - - # Use PDOSData.plot_pdos method - pdos_pic_file = os.path.join(output_dir, "PDOS.png") - PDOSData.plot_pdos( - pdosdatas=all_pdos_datas, - labels=all_labels, - titles=all_titles, - energy=pdos_data.energy, - energy_min=dos_emin_ev, - energy_max=dos_emax_ev, - pdos_fig_name=pdos_pic_file, - ) - - all_pdos_data_flattened = [data for species_pdos_datas in all_pdos_datas for data in species_pdos_datas] - all_labels_flattened = [label for species_labels in all_labels for label in species_labels] - pdos_data_file = os.path.join(output_dir, "PDOS.dat") - PDOSData.write_pdos( - pdosdatas=all_pdos_data_flattened, - energy=pdos_data.energy, - labels=all_labels_flattened, - filename=pdos_data_file, - ) - - return pdos_pic_file, pdos_data_file - - -def plot_write_pdos_species_orbital( - pdos_data: PDOSData, - output_dir: Path, - dos_emin_ev: float, - dos_emax_ev: float, -) -> Tuple[str, List[str]]: - """Plot PDOS by species and orbital using PDOSData class.""" - os.makedirs(output_dir, exist_ok=True) - - # Get unique species - species_shell_set = set() - for orbital in pdos_data.projected_dos: - species_shell_set.add((orbital["species"], orbital["l"])) - - # Prepare data for PDOSData.plot_pdos - all_pdos_datas, all_labels, all_titles = [], [], [] - - for (species, l) in sorted(species_shell_set): - orbital_pdos_data = [] - species_orbital_labels = [] - # Get orbitals for this shell - species_orbitals = [] - for orbital in pdos_data.projected_dos: - if orbital["species"] == species and orbital["l"] == l: - m = orbital["m"] - if m not in species_orbitals: - species_orbitals.append(m) - species_orbital_pdos = pdos_data.get_pdos_by_species_orbital(species, l, m) - orbital_pdos_data.append(species_orbital_pdos) - - orbital_label = f"{species}-{orbital_names[(l, m)]}" - species_orbital_labels.append(orbital_label) - - all_titles.append(f"PDOS of {species}-{l_map[l]}") - all_pdos_datas.append(orbital_pdos_data) - all_labels.append(species_orbital_labels) - - # Use PDOSData.plot_pdos method - pdos_pic_file = os.path.join(output_dir, "PDOS.png") - PDOSData.plot_pdos( - pdosdatas=all_pdos_datas, - labels=all_labels, - titles=all_titles, - energy=pdos_data.energy, - energy_min=dos_emin_ev, - energy_max=dos_emax_ev, - pdos_fig_name=pdos_pic_file, - ) - - all_pdos_data_flattened = [data for species_pdos_datas in all_pdos_datas for data in species_pdos_datas] - all_labels_flattened = [label.replace("$", "") for species_labels in all_labels for label in species_labels] - pdos_data_file = os.path.join(output_dir, "PDOS.dat") - PDOSData.write_pdos( - pdosdatas=all_pdos_data_flattened, - energy=pdos_data.energy, - labels=all_labels_flattened, - filename=pdos_data_file, - ) - - return pdos_pic_file, pdos_data_file - - -def plot_write_pdos_atoms( - pdos_data: PDOSData, - output_dir: Path, - pdos_atom_indices: List[int], - dos_emin_ev: float, - dos_emax_ev: float, -) -> Tuple[str, List[str]]: - """Plot PDOS for selected atoms using PDOSData class.""" - os.makedirs(output_dir, exist_ok=True) - - # Prepare data for PDOSData.plot_pdos - all_pdos_datas, all_labels, all_titles = [], [], [] - - for atom_index in pdos_atom_indices: - # Obtain all shells of the selected atom - atom_shell_set = set() - for orbital in pdos_data.projected_dos: - if orbital["atom_index"] == atom_index: - species = orbital["species"] - atom_shell_set.add((orbital["l"])) - - for l in atom_shell_set: - orbital_pdos_data = [] - species_orbital_labels = [] - # Get shells for this species - species_orbitals = [] - for orbital in pdos_data.projected_dos: - if orbital["species"] == species and orbital["l"] == l: - m = orbital["m"] - if m not in species_orbitals: - species_orbitals.append(m) - species_orbital_pdos = pdos_data.get_pdos_by_species_orbital(species, l, m) - orbital_pdos_data.append(species_orbital_pdos) - - orbital_label = f"{species}{atom_index}-{orbital_names[(l, m)]}" - species_orbital_labels.append(orbital_label) - - all_titles.append(f"PDOS of {species}{atom_index}-{l_map[l]}") - all_pdos_datas.append(orbital_pdos_data) - all_labels.append(species_orbital_labels) - - # Use PDOSData.plot_pdos method - pdos_pic_file = os.path.join(output_dir, "PDOS.png") - PDOSData.plot_pdos( - pdosdatas=all_pdos_datas, - labels=all_labels, - titles=all_titles, - energy=pdos_data.energy, - energy_min=dos_emin_ev, - energy_max=dos_emax_ev, - pdos_fig_name=pdos_pic_file, - ) - - all_pdos_data_flattened = [data for species_pdos_datas in all_pdos_datas for data in species_pdos_datas] - all_labels_flattened = [label.replace("$", "") for species_labels in all_labels for label in species_labels] - pdos_data_file = os.path.join(output_dir, "PDOS.dat") - PDOSData.write_pdos( - pdosdatas=all_pdos_data_flattened, - energy=pdos_data.energy, - labels=all_labels_flattened, - filename=pdos_data_file, - ) - - return pdos_pic_file, pdos_data_file - def plot_write_dos_pdos( scf_job_path: Path, nscf_job_path: Path, @@ -473,19 +204,21 @@ def plot_write_dos_pdos( dos_emin_ev: float = -10.0, dos_emax_ev: float = 5.0, ) -> Tuple[List[str], List[str]]: - """Plot DOS and PDOS from the NSCF job path using PDOSData class. + """ + Plot DOS, PDOS and write data used in plotting to files using SCF and NSCF job directories from abacus_dos_run. Args: scf_job_path (Path): Path to the SCF job directory of the DOS calculation nscf_job_path (Path): Path to the NSCF job directory of the DOS calculation - mode (str): PDOS plotting mode ('species', 'species+shell', 'species+orbital', or 'atoms'). + mode: Mode for plotting PDOS and write PDOS data. + - "atoms": PDOS of a list of atoms will be plotted. + - "species": Total PDOS of any species will be plotted in a picture. + - "species+shell": PDOS for any shell (s, p, d, f, g,...) of any species will be plotted. PDOS of a shell of a species willbe plotted in a subplot. + - "species+orbital": Orbital-resolved PDOS will be plotted. PDOS of orbitals in the same shell of a species will be plotted in a subplot. + pdos_atom_indices: A list of atom indices, only used if pdos_mode is "atoms". pdos_atom_indices (List[int], optional): List of atom indices for atom-specific PDOS. Only valid for 'atoms' mode. dos_emin_ev (float): Minimum energy for DOS and PDOS plots. - dos_emax_ev (float): Maximum energy for DOS and PDOS plots. - - Returns: - Tuple[List[str], List[str]]: Tuple containing list of plot file paths and data file paths. - """ + dos_emax_ev (float): Maximum energy for DOS and PDOS plots. """ work_path = generate_work_path() input_param = ReadInput(os.path.join(nscf_job_path, "INPUT")) @@ -519,32 +252,26 @@ def plot_write_dos_pdos( try: # Load PDOS data using PDOSData class pdos_data = PDOSData.ReadFromAbacusJob(str(nscf_job_path), efermi) + pdos_plot_file = Path(os.path.join(work_path, "PDOS.png")).absolute() + pdos_data_file = Path(os.path.join(work_path, "PDOS.dat")).absolute() # Plot PDOS based on mode if mode == "species": - pdos_plot_file, pdos_data_file = plot_write_pdos_species( - pdos_data, work_path, dos_emin_ev, dos_emax_ev - ) + pdos_data.plot_species_pdos(dos_emin_ev, dos_emax_ev, pdos_plot_file) + pdos_data.write_species_pdos(pdos_data_file) elif mode == "species+shell": - pdos_plot_file, pdos_data_file = plot_write_pdos_species_shell( - pdos_data, work_path, dos_emin_ev, dos_emax_ev - ) + pdos_data.plot_species_shell_pdos(dos_emin_ev, dos_emax_ev, pdos_plot_file) + pdos_data.write_species_shell_pdos(pdos_data_file) elif mode == "species+orbital": - pdos_plot_file, pdos_data_file = plot_write_pdos_species_orbital( - pdos_data, work_path, dos_emin_ev, dos_emax_ev - ) + pdos_data.plot_species_orbital_pdos(dos_emin_ev, dos_emax_ev, pdos_plot_file) + pdos_data.write_species_orbital_pdos(pdos_data_file) elif mode == "atoms": if pdos_atom_indices is None or len(pdos_atom_indices) == 0: raise ValueError( "For 'atoms' mode, pdos_atom_indices must be provided" ) - pdos_plot_file, pdos_data_file = plot_write_pdos_atoms( - pdos_data, - work_path, - pdos_atom_indices, - dos_emin_ev, - dos_emax_ev, - ) + pdos_data.plot_atoms_pdos(pdos_atom_indices, dos_emin_ev, dos_emax_ev, pdos_plot_file) + pdos_data.write_atoms_pdos(pdos_atom_indices, pdos_data_file) else: raise ValueError(f"Unsupported mode: {mode}") From bca7ddea95f407d8cdb987522ed48f442c604fe9 Mon Sep 17 00:00:00 2001 From: ahxbcn Date: Fri, 23 Jan 2026 10:57:46 +0800 Subject: [PATCH 11/15] Individual DOS and PDOS postprocess tool --- src/abacusagent/modules/dos.py | 34 +++++++++++++++++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/src/abacusagent/modules/dos.py b/src/abacusagent/modules/dos.py index d18a7d0..1b3257d 100644 --- a/src/abacusagent/modules/dos.py +++ b/src/abacusagent/modules/dos.py @@ -1,8 +1,9 @@ from pathlib import Path -from typing import Dict, Any, List, Literal, Optional +from typing import Dict, Any, List, Literal, Optional, Tuple from abacusagent.init_mcp import mcp from abacusagent.modules.submodules.dos import abacus_dos_run as _abacus_dos_run +from abacusagent.modules.submodules.dos import plot_write_dos_pdos as _plot_write_dos_pdos @mcp.tool() def abacus_dos_run( @@ -48,3 +49,34 @@ def abacus_dos_run( - nscf_work_path: Path to the work directory of NSCF calculation """ return _abacus_dos_run(abacus_inputs_dir, pdos_mode, pdos_atom_indices, dos_edelta_ev, dos_sigma, dos_emin_ev, dos_emax_ev) + +def plot_write_dos_pdos( + scf_job_path: Path, + nscf_job_path: Path, + mode: Literal[ + "species", "species+shell", "species+orbital", "atoms" + ] = "species+shell", + pdos_atom_indices: Optional[List[int]] = None, + dos_emin_ev: float = -10.0, + dos_emax_ev: float = 5.0, +) -> Tuple[List[str], List[str]]: + """ + Plot DOS, PDOS and write data used in plotting to files using SCF and NSCF job directories from abacus_dos_run. + + Args: + scf_job_path (Path): Path to the SCF job directory of the DOS calculation + nscf_job_path (Path): Path to the NSCF job directory of the DOS calculation + mode: Mode for plotting PDOS and write PDOS data. + - "atoms": PDOS of a list of atoms will be plotted. + - "species": Total PDOS of any species will be plotted in a picture. + - "species+shell": PDOS for any shell (s, p, d, f, g,...) of any species will be plotted. PDOS of a shell of a species willbe plotted in a subplot. + - "species+orbital": Orbital-resolved PDOS will be plotted. PDOS of orbitals in the same shell of a species will be plotted in a subplot. + pdos_atom_indices: A list of atom indices, only used if pdos_mode is "atoms". + pdos_atom_indices (List[int], optional): List of atom indices for atom-specific PDOS. Only valid for 'atoms' mode. + dos_emin_ev (float): Minimum energy for DOS and PDOS plots. + dos_emax_ev (float): Maximum energy for DOS and PDOS plots. + + Returns: + Tuple[List[str], List[str]]: Tuple containing list of plot file paths and data file paths. + """ + return _plot_write_dos_pdos(scf_job_path, nscf_job_path, mode, pdos_atom_indices, dos_emin_ev, dos_emax_ev) From b536fc41b3fb1d879b00fad7e28eaae6b6d1a340 Mon Sep 17 00:00:00 2001 From: ahxbcn Date: Fri, 23 Jan 2026 12:03:52 +0800 Subject: [PATCH 12/15] Change method for plotting and writing DOS --- src/abacusagent/modules/submodules/dos.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/abacusagent/modules/submodules/dos.py b/src/abacusagent/modules/submodules/dos.py index 10ac025..116813a 100644 --- a/src/abacusagent/modules/submodules/dos.py +++ b/src/abacusagent/modules/submodules/dos.py @@ -232,15 +232,13 @@ def plot_write_dos_pdos( dos_data_file = os.path.join(work_path, "DOS.dat") dosdata = DOSData.ReadFromAbacusJob(str(nscf_job_path), efermi) - DOSData.plot_dos( - dosdata.dosdata, - dosdata.energy, + dosdata.plot_dos( dos_emin_ev, dos_emax_ev, "Density of States", dos_plot_file, ) - DOSData.write_dos(dosdata.dosdata, dosdata.energy, dos_data_file) + dosdata.write_dos(dos_data_file) all_plot_files = [Path(dos_plot_file).absolute()] dos_pdos_data_files = [Path(dos_data_file).absolute()] From dd193cab03115024a0943fade7aff28c5aa4780b Mon Sep 17 00:00:00 2001 From: ahxbcn Date: Fri, 23 Jan 2026 12:12:23 +0800 Subject: [PATCH 13/15] Update unittest for dos.py --- tests/test_dos.py | 40 ++++++++++++++++++++++------------------ 1 file changed, 22 insertions(+), 18 deletions(-) diff --git a/tests/test_dos.py b/tests/test_dos.py index 0d498f7..041ff71 100644 --- a/tests/test_dos.py +++ b/tests/test_dos.py @@ -1,42 +1,46 @@ import unittest -import os, sys, glob +import os, sys, glob, shutil from pathlib import Path + os.environ["ABACUSAGENT_MODEL"] = "test" -from abacusagent.modules.submodules.dos import plot_dos_pdos as mkplots +from abacusagent.modules.submodules.dos import plot_write_dos_pdos as mkplots + class TestPlotDos(unittest.TestCase): def setUp(self): self.data_dir = Path(__file__).parent / "plot_dos" - def tearDown(self): for pngfile in glob.glob(os.path.join(self.data_dir, "*.png")): os.remove(pngfile) - + def test_run_dos(self): """ Test the run_dos function with a valid input. """ # ignore the screen output - sys.stdout = open(os.devnull, 'w') + sys.stdout = open(os.devnull, "w") # Call the run_dos function - results_test = mkplots(self.data_dir, self.data_dir, self.data_dir, dpi=20) - results_ref = [Path(self.data_dir) / "DOS.png", - Path(self.data_dir) / "PDOS.png"] - - self.assertListEqual([Path(p) for p in results_test], results_ref) + results_figs, results_datas = mkplots(self.data_dir, self.data_dir, "species", dos_emin_ev=-1, dos_emax_ev=1) + + output_dir = Path(glob.glob("*plot_write_dos_pdos*")[0]).absolute() + results_figs_ref = [ + output_dir / "DOS.png", + output_dir / "PDOS.png", + ] + results_datas_ref = [ + output_dir / "DOS.dat", + output_dir / "PDOS.dat", + ] + + self.assertListEqual([Path(p) for p in results_figs], results_figs_ref) + self.assertListEqual([Path(p) for p in results_datas], results_datas_ref) if os.path.exists(self.data_dir / "metrics.json"): os.remove(self.data_dir / "metrics.json") - - - - - - - - + for dir in glob.glob("*plot_write_dos_pdos*"): + shutil.rmtree(dir) From 46fcd5b9a223218f6e3b6a3fe40068f44a4d00f3 Mon Sep 17 00:00:00 2001 From: ahxbcn Date: Fri, 23 Jan 2026 12:38:51 +0800 Subject: [PATCH 14/15] Update integrate test for dos.py --- src/abacusagent/modules/submodules/dos.py | 2 +- tests/integrate_test/test_dos.py | 40 +++++++++-------------- 2 files changed, 17 insertions(+), 25 deletions(-) diff --git a/src/abacusagent/modules/submodules/dos.py b/src/abacusagent/modules/submodules/dos.py index 116813a..d2f14db 100644 --- a/src/abacusagent/modules/submodules/dos.py +++ b/src/abacusagent/modules/submodules/dos.py @@ -98,7 +98,7 @@ def abacus_dos_run( return_dict["dos_data_path"] = dos_pdos_data_paths[0] try: return_dict["pdos_fig_path"] = fig_paths[1] - return_dict["pdos_data_paths"] = dos_pdos_data_paths[1:] + return_dict["pdos_data_path"] = dos_pdos_data_paths[1] except: pass # Do nothing if PDOS file is not plotted diff --git a/tests/integrate_test/test_dos.py b/tests/integrate_test/test_dos.py index 6b52a5c..e7b4902 100644 --- a/tests/integrate_test/test_dos.py +++ b/tests/integrate_test/test_dos.py @@ -53,13 +53,12 @@ def test_abacus_dos_run_species(self): dos_fig_path = outputs['dos_fig_path'] dos_data_path = outputs['dos_data_path'] pdos_fig_path = outputs['pdos_fig_path'] - pdos_data_paths = outputs['pdos_data_paths'] + pdos_data_path = outputs['pdos_data_path'] self.assertIsInstance(dos_fig_path, get_path_type()) self.assertIsInstance(dos_data_path, get_path_type()) self.assertIsInstance(pdos_fig_path, get_path_type()) - for pdos_data_path in pdos_data_paths: - self.assertIsInstance(pdos_data_path, get_path_type()) + self.assertIsInstance(pdos_data_path, get_path_type()) self.assertTrue(outputs['scf_normal_end']) self.assertTrue(outputs['scf_converge']) self.assertTrue(outputs['nscf_normal_end']) @@ -83,13 +82,12 @@ def test_abacus_dos_run_species_shell(self): dos_fig_path = outputs['dos_fig_path'] dos_data_path = outputs['dos_data_path'] pdos_fig_path = outputs['pdos_fig_path'] - pdos_data_paths = outputs['pdos_data_paths'] + pdos_data_path = outputs['pdos_data_path'] self.assertIsInstance(dos_fig_path, get_path_type()) self.assertIsInstance(dos_data_path, get_path_type()) self.assertIsInstance(pdos_fig_path, get_path_type()) - for pdos_data_path in pdos_data_paths: - self.assertIsInstance(pdos_data_path, get_path_type()) + self.assertIsInstance(pdos_data_path, get_path_type()) self.assertTrue(outputs['scf_normal_end']) self.assertTrue(outputs['scf_converge']) self.assertTrue(outputs['nscf_normal_end']) @@ -113,13 +111,12 @@ def test_abacus_dos_run_species_orbital(self): dos_fig_path = outputs['dos_fig_path'] dos_data_path = outputs['dos_data_path'] pdos_fig_path = outputs['pdos_fig_path'] - pdos_data_paths = outputs['pdos_data_paths'] + pdos_data_path = outputs['pdos_data_path'] self.assertIsInstance(dos_fig_path, get_path_type()) self.assertIsInstance(dos_data_path, get_path_type()) self.assertIsInstance(pdos_fig_path, get_path_type()) - for pdos_data_path in pdos_data_paths: - self.assertIsInstance(pdos_data_path, get_path_type()) + self.assertIsInstance(pdos_data_path, get_path_type()) self.assertTrue(outputs['scf_normal_end']) self.assertTrue(outputs['scf_converge']) self.assertTrue(outputs['nscf_normal_end']) @@ -144,13 +141,12 @@ def test_abacus_dos_run_atoms(self): dos_fig_path = outputs['dos_fig_path'] dos_data_path = outputs['dos_data_path'] pdos_fig_path = outputs['pdos_fig_path'] - pdos_data_paths = outputs['pdos_data_paths'] + pdos_data_path = outputs['pdos_data_path'] self.assertIsInstance(dos_fig_path, get_path_type()) self.assertIsInstance(dos_data_path, get_path_type()) self.assertIsInstance(pdos_fig_path, get_path_type()) - for pdos_data_path in pdos_data_paths: - self.assertIsInstance(pdos_data_path, get_path_type()) + self.assertIsInstance(pdos_data_path, get_path_type()) self.assertTrue(outputs['scf_normal_end']) self.assertTrue(outputs['scf_converge']) self.assertTrue(outputs['nscf_normal_end']) @@ -178,13 +174,12 @@ def test_abacus_dos_run_species_nspin2(self): dos_fig_path = outputs['dos_fig_path'] dos_data_path = outputs['dos_data_path'] pdos_fig_path = outputs['pdos_fig_path'] - pdos_data_paths = outputs['pdos_data_paths'] + pdos_data_path = outputs['pdos_data_path'] self.assertIsInstance(dos_fig_path, get_path_type()) self.assertIsInstance(dos_data_path, get_path_type()) self.assertIsInstance(pdos_fig_path, get_path_type()) - for pdos_data_path in pdos_data_paths: - self.assertIsInstance(pdos_data_path, get_path_type()) + self.assertIsInstance(pdos_data_path, get_path_type()) self.assertTrue(outputs['scf_normal_end']) self.assertTrue(outputs['scf_converge']) self.assertTrue(outputs['nscf_normal_end']) @@ -207,13 +202,12 @@ def test_abacus_dos_run_species_shell_nspin2(self): dos_fig_path = outputs['dos_fig_path'] dos_data_path = outputs['dos_data_path'] pdos_fig_path = outputs['pdos_fig_path'] - pdos_data_paths = outputs['pdos_data_paths'] + pdos_data_path = outputs['pdos_data_path'] self.assertIsInstance(dos_fig_path, get_path_type()) self.assertIsInstance(dos_data_path, get_path_type()) self.assertIsInstance(pdos_fig_path, get_path_type()) - for pdos_data_path in pdos_data_paths: - self.assertIsInstance(pdos_data_path, get_path_type()) + self.assertIsInstance(pdos_data_path, get_path_type()) self.assertTrue(outputs['scf_normal_end']) self.assertTrue(outputs['scf_converge']) self.assertTrue(outputs['nscf_normal_end']) @@ -236,13 +230,12 @@ def test_abacus_dos_run_species_orbital_nspin2(self): dos_fig_path = outputs['dos_fig_path'] dos_data_path = outputs['dos_data_path'] pdos_fig_path = outputs['pdos_fig_path'] - pdos_data_paths = outputs['pdos_data_paths'] + pdos_data_path = outputs['pdos_data_path'] self.assertIsInstance(dos_fig_path, get_path_type()) self.assertIsInstance(dos_data_path, get_path_type()) self.assertIsInstance(pdos_fig_path, get_path_type()) - for pdos_data_path in pdos_data_paths: - self.assertIsInstance(pdos_data_path, get_path_type()) + self.assertIsInstance(pdos_data_path, get_path_type()) self.assertTrue(outputs['scf_normal_end']) self.assertTrue(outputs['scf_converge']) self.assertTrue(outputs['nscf_normal_end']) @@ -270,13 +263,12 @@ def test_abacus_dos_run_atoms_nspin2(self): dos_fig_path = outputs['dos_fig_path'] dos_data_path = outputs['dos_data_path'] pdos_fig_path = outputs['pdos_fig_path'] - pdos_data_paths = outputs['pdos_data_paths'] + pdos_data_path = outputs['pdos_data_path'] self.assertIsInstance(dos_fig_path, get_path_type()) self.assertIsInstance(dos_data_path, get_path_type()) self.assertIsInstance(pdos_fig_path, get_path_type()) - for pdos_data_path in pdos_data_paths: - self.assertIsInstance(pdos_data_path, get_path_type()) + self.assertIsInstance(pdos_data_path, get_path_type()) self.assertTrue(outputs['scf_normal_end']) self.assertTrue(outputs['scf_converge']) self.assertTrue(outputs['nscf_normal_end']) From a4fe7f81bd8a10bc0ef2435eaa2e2ef411a7b9fb Mon Sep 17 00:00:00 2001 From: ahxbcn Date: Fri, 23 Jan 2026 14:58:15 +0800 Subject: [PATCH 15/15] Update version of abacustest --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 46659c0..d0c8fa4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,7 +20,7 @@ description = "ABACUS agent tools for connecting LLMs to first-principles calcul requires-python = ">=3.11" dependencies = [ - "abacustest>=0.4.51", + "abacustest>=0.4.55", "numpy<2.0", "pymatgen>=2025.5.28", "bohr-agent-sdk==0.1.121",