Allow dipole correction in work function calculation

src/abacusagent/modules/submodules/work_function.py

@@ -1,47 +1,115 @@
 import os
 from pathlib import Path
-from typing import Literal, Optional, Dict, Any
+from typing import Literal, Optional, Dict, Any, List
 
-from abacustest.lib_prepare.abacus import ReadInput, WriteInput
+import numpy as np
+from abacustest.lib_prepare.abacus import ReadInput, WriteInput, AbacusStru
 from abacustest.lib_model.comm import check_abacus_inputs
 
 from abacusagent.constant import RY_TO_EV
 from abacusagent.modules.util.comm import run_abacus, generate_work_path, link_abacusjob, collect_metrics
 from abacusagent.modules.util.cube_manipulator import read_gaussian_cube, profile1d
 
+def identify_potential_plateaus(
+    averaged_potential: List,
+    length: float = 10.0,
+    threshold: float = 0.01
+):
+    """
+    Identify plateaus in the electrostatic potential using derivative of the electrostatic potential.
+    """
+    pot_derivatives = []
+    n_points = len(averaged_potential)
+    stepsize = length / (n_points - 1)
+    for i in range(n_points):
+        backward_idx = i - 1
+        forward_idx = 0 if i == n_points - 1 else i + 1
+        pot_derivative = (averaged_potential[forward_idx] - averaged_potential[backward_idx]) / (stepsize * 2.0)
+        pot_derivatives.append(pot_derivative)
+
+    is_plateau = [abs(deriv) < threshold for deriv in pot_derivatives]
+
+    plateau_ranges = []
+    in_plateau, start_idx = False, None
+    for i in range(n_points):
+        if is_plateau[i] and not in_plateau:
+            in_plateau = True
+            start_idx = i
+        elif not is_plateau[i] and in_plateau:
+            in_plateau = False
+            if not start_idx == i - 1:
+                plateau_ranges.append((start_idx, i-1))
+                start_idx = None
+        elif is_plateau[i] and i == n_points - 1:
+            if not start_idx == i - 1:
+                plateau_ranges.append((start_idx, i))
+                in_plateau, start_idx = False, None
+
+    if len(plateau_ranges) > 0:
+        if plateau_ranges[-1][1] == n_points - 1:
+            if len(plateau_ranges) > 1:
+                combined_plateau = (plateau_ranges[-1][0] - n_points, plateau_ranges[0][1])
+                plateau_ranges[0] = combined_plateau
+                plateau_ranges.pop()
+
+    return plateau_ranges
+
+def calculate_work_functions(averaged_potential: List, fermi_energy, length: float = 10.0):
+    """
+    Calculate the work function from the averaged electrostatic potential.
+    Dipole correction is suppoted and multiple plateau of electrostatic potential can be identified.
+    """
+    work_function_results = []
+    plateau_ranges = identify_potential_plateaus(averaged_potential, length=length, threshold=0.01)
+    for plateau_range in plateau_ranges:
+        plateau_start, plateau_end = plateau_range
+        plateau_averaged_potential = np.average(averaged_potential[plateau_start:plateau_end])
+        work_function_results.append({'work_function': plateau_averaged_potential - fermi_energy,
+                                      'plateau_start_fractional': plateau_start / len(averaged_potential),
+                                      'plateau_end_fractional': plateau_end / len(averaged_potential)})
+
+    return work_function_results
+
 def plot_averaged_elecstat_pot(
     averaged_elecstat_data,
     work_path: Path,
     axis: Literal['x', 'y', 'z'] = 'z',
     plot_filename: Optional[str] = "elecstat_pot_profile.png"
-) -> Dict[str, Any]:
+) -> Path:
     import matplotlib.pyplot as plt
     plt.plot(averaged_elecstat_data['data'][:, 0], averaged_elecstat_data['data'][:, 1], label='Electrostatic Potential')
     plt.xlim(0, 1)
     plt.xlabel("Fractional Coordinate along " + axis)
     plt.ylabel("Electrostatic Potential (eV)")
     plot_path = os.path.join(work_path, plot_filename)
     plt.savefig(plot_path, dpi=300)
+    plt.close()
 
     return plot_path
 
 def abacus_cal_work_function(
     abacus_inputs_dir: Path,
     vacuum_direction: Literal['x', 'y', 'z'] = 'z',
+    dipole_correction: bool = False,
 ) -> Dict[str, Any]:
     """
     Calculate the electrostatic potential and work function using ABACUS.
 
     Args:
         abacus_inputs_dir (Path): Path to the ABACUS input files, which contains the INPUT, STRU, KPT, and pseudopotential or orbital files.
         vacuum_direction (Literal['x', 'y', 'z']): The direction of the vacuum.
+        dipole_correction (bool): Whether to apply dipole correction along the vacuum direction. For polar slabs, it is recommended to enable dipole correction.
 
     Returns:
         A dictionary containing:
         - elecstat_pot_work_function_work_path (Path): Path to the ABACUS job directory calculating electrostatic potential and work function.
         - elecstat_pot_file (Path): Path to the cube file containing the electrostatic potential.
         - averaged_elecstat_pot_plot (Path): Path to the plot of the averaged electrostatic potential.
-        - work_function (float): The calculated work function in eV.
+        - work_function_results (list): A list of 1 or 2 dictionary. If dipole correction is not used, only 1 dictionaray will be returned. 
+          If dipole correction is used, there will be 2 dictionarys for calculated work function of 2 surfaces of the slab. Each dictionary contains 3 keys:
+            - 'work_function': calculated work function
+            - 'plateau_start_fractional': Fractional coordinate of start of the identified plateau in the given vacuum direction
+            - 'plateau_end_fractional': Fractional coordinate of end of the identified plateau in the given vacuum direction
     """
     try:
         is_valid, msg = check_abacus_inputs(abacus_inputs_dir)
@@ -51,11 +119,22 @@ def abacus_cal_work_function(
         work_path = Path(generate_work_path()).absolute()
         link_abacusjob(src=abacus_inputs_dir,dst=work_path,copy_files=["INPUT", "STRU"], exclude_directories=True)
         input_params = ReadInput(os.path.join(work_path, 'INPUT'))
+        stru = AbacusStru.ReadStru(os.path.join(work_path, input_params.get('stru_file', 'STRU')))
         if input_params.get('nspin', 1) not in [1, 2]:
             raise ValueError('Only non spin-polarized and collinear spin-polarized calculation are supported for calculating electrostatic potential and work function')
 
+        direction_map = {'x': 0, 'y': 1, 'z': 2}
         input_params['calculation'] = 'scf'
         input_params['out_pot'] = 2
+
+        if dipole_correction:
+            input_params['efield_flag'] = 1
+            input_params['dip_cor_flag'] = 1
+            input_params['efield_dir'] = direction_map[vacuum_direction]
+            input_params['efield_amp'] = 0.00
+            input_params['efield_pos_max'] = None
+            input_params['efield_pos_dec'] = None
+
         WriteInput(input_params, os.path.join(work_path, 'INPUT'))
 
         run_abacus(work_path)
@@ -69,15 +148,18 @@ def abacus_cal_work_function(
 
         profile_result = profile1d(pot, axis=vacuum_direction, average=True)
         profile_result['data'][:, 1] *= RY_TO_EV  # Convert from Rydberg to eV
-        v_vacuum = max(profile_result['data'][:, 1])
-        work_function = v_vacuum - metrics['efermi']
+
+        length = np.linalg.norm(stru.get_cell()[direction_map[vacuum_direction]])
+        work_function_results = calculate_work_functions(profile_result['data'][:, 1],
+                                                         fermi_energy=metrics['efermi'],
+                                                         length=length)
 
         # Plot the averaged electrostatic potential
         plot_path = plot_averaged_elecstat_pot(profile_result, work_path, axis=vacuum_direction)
 
         return {'elecstat_pot_work_function_work_path': Path(work_path).absolute(),
                 'elecstat_pot_file': Path(pot_file).absolute(),
                 'averaged_elecstat_pot_plot': Path(plot_path).absolute(),
-                'work_function': work_function}
+                'work_function_results': work_function_results}
     except Exception as e:
         return {'message': f"Calculating electrostatic potential and work function failed: {e}"}

src/abacusagent/modules/work_function.py

@@ -8,19 +8,25 @@
 def abacus_cal_work_function(
     abacus_inputs_dir: Path,
     vacuum_direction: Literal['x', 'y', 'z'] = 'z',
+    dipole_correction: bool = False,
 ) -> Dict[str, Any]:
     """
     Calculate the electrostatic potential and work function using ABACUS.
 
     Args:
         abacus_inputs_dir (Path): Path to the ABACUS input files, which contains the INPUT, STRU, KPT, and pseudopotential or orbital files.
         vacuum_direction (Literal['x', 'y', 'z']): The direction of the vacuum.
+        dipole_correction (bool): Whether to apply dipole correction along the vacuum direction. For polar slabs, it is recommended to enable dipole correction.
 
     Returns:
         A dictionary containing:
         - elecstat_pot_work_function_work_path (Path): Path to the ABACUS job directory calculating electrostatic potential and work function.
         - elecstat_pot_file (Path): Path to the cube file containing the electrostatic potential.
         - averaged_elecstat_pot_plot (Path): Path to the plot of the averaged electrostatic potential.
-        - work_function (float): The calculated work function in eV.
+        - work_function_results (list): A list of 1 or 2 dictionary. If dipole correction is not used, only 1 dictionaray will be returned. 
+          If dipole correction is used, there will be 2 dictionarys for calculated work function of 2 surfaces of the slab. Each dictionary contains 3 keys:
+            - 'work_function': calculated work function
+            - 'plateau_start_fractional': Fractional coordinate of start of the identified plateau in the given vacuum direction
+            - 'plateau_end_fractional': Fractional coordinate of end of the identified plateau in the given vacuum direction
     """
-    return _abacus_cal_work_function(abacus_inputs_dir, vacuum_direction)
+    return _abacus_cal_work_function(abacus_inputs_dir, vacuum_direction, dipole_correction)

tests/integrate_test/abacus_inputs_dirs/ZnO0001/INPUT_work_function

@@ -0,0 +1,13 @@
+INPUT_PARAMETERS
+calculation 	scf
+symmetry 	1
+ecutwfc 	100
+scf_thr 	1e-07
+scf_nmax 	100
+smearing_method 	gauss
+smearing_sigma 	0.015
+mixing_type 	broyden
+mixing_beta 	0.8
+basis_type 	lcao
+ks_solver 	genelpa
+kspacing 	0.14

tests/integrate_test/abacus_inputs_dirs/ZnO0001/O.upf

@@ -0,0 +1 @@
+../../pp_orbs/O.upf

tests/integrate_test/abacus_inputs_dirs/ZnO0001/O_gga_6au_100Ry_2s2p1d.orb

@@ -0,0 +1 @@
+../../pp_orbs/O_gga_6au_100Ry_2s2p1d.orb

tests/integrate_test/abacus_inputs_dirs/ZnO0001/STRU

@@ -0,0 +1,30 @@
+ATOMIC_SPECIES
+Zn 65.390000 Zn_ONCV_PBE_FR-1.0.upf
+O 15.999400 O.upf
+
+NUMERICAL_ORBITAL
+Zn_gga_9au_150Ry_4s2p2d1f.orb
+O_gga_6au_100Ry_2s2p1d.orb
+
+LATTICE_CONSTANT
+1.889726
+
+LATTICE_VECTORS
+    3.23735102000     0.00000000000     0.00000000000
+   -1.61867551000     2.80362822429     0.00000000000
+    0.00000000000     0.00000000000    20.88824520000
+
+ATOMIC_POSITIONS
+Cartesian
+
+Zn
+0.000000
+2
+   -0.00000001619     1.86908549220     4.23215023161 1 1 1 
+    1.61867552619     0.93454273208     1.62111958161 1 1 1 
+
+O
+0.000000
+2
+   -0.00000001619     1.86908549220     0.98991106839 1 1 1 
+    1.61867552619     0.93454273208     3.60094171839 1 1 1 

tests/integrate_test/abacus_inputs_dirs/ZnO0001/Zn_ONCV_PBE_FR-1.0.upf

@@ -0,0 +1 @@
+../../pp_orbs/Zn_ONCV_PBE_FR-1.0.upf

tests/integrate_test/abacus_inputs_dirs/ZnO0001/Zn_gga_9au_150Ry_4s2p2d1f.orb

@@ -0,0 +1 @@
+../../pp_orbs/Zn_gga_9au_150Ry_4s2p2d1f.orb

tests/integrate_test/data/ref_results.json

@@ -281,7 +281,20 @@
     "test_abacus_cal_work_function_al110":
     {
         "result": {
-            "work_function": 4.329085761761545
+            "work_function_results": [{"work_function": 4.328994681259157, 
+                                       "plateau_start_fractional": 0.5637860082304527, 
+                                       "plateau_end_fractional": 0.8353909465020576}]
+        }
+    },
+    "test_abacus_cal_work_function_zno0001_dipole_corr":
+    {
+        "result": {
+            "work_function_results": [{"work_function": 2.6146683619622704, 
+                                       "plateau_start_fractional": 0.37109375, 
+                                       "plateau_end_fractional": 0.62109375}, 
+                                      {"work_function": 7.8168926011874955, 
+                                       "plateau_start_fractional": 0.73046875, 
+                                       "plateau_end_fractional": 0.92578125}]
         }
     },
     "test_abacus_cal_vacancy_formation_energy_gamma_tial":