diff --git a/ChangeLog b/ChangeLog index 6819c58d..ff770ff4 100644 --- a/ChangeLog +++ b/ChangeLog @@ -2,6 +2,14 @@ -Date -Name -changes +-------------- +November 19, 2025 +Name: Qihao Cheng +Changes: utils/pdos/, utils/atom/ +1. Added XC functional fallback in PDOS: unsupported SPARC functionals now default to GGA_PBE for atomic orbital generation. +2. Extended atomic DFT module to support EXX and RPA (OEP-based implementation). +3. Added thread-parallel RPA frequency evaluation and improved standalone atomic code performance. + -------------- November 07, 2025 Name: Qihao Cheng diff --git a/src/initialization.c b/src/initialization.c index 82d860d5..0c20420b 100644 --- a/src/initialization.c +++ b/src/initialization.c @@ -3722,7 +3722,7 @@ void write_output_init(SPARC_OBJ *pSPARC) { } fprintf(output_fp,"***************************************************************************\n"); - fprintf(output_fp,"* SPARC (version November 07, 2025) *\n"); + fprintf(output_fp,"* SPARC (version November 19, 2025) *\n"); fprintf(output_fp,"* Copyright (c) 2020 Material Physics & Mechanics Group, Georgia Tech *\n"); fprintf(output_fp,"* Distributed under GNU General Public License 3 (GPL) *\n"); fprintf(output_fp,"* Start time: %s *\n",c_time_str); diff --git a/utils/atom/__init__.py b/utils/atom/__init__.py index 477bb0bb..ce96af08 100644 --- a/utils/atom/__init__.py +++ b/utils/atom/__init__.py @@ -15,6 +15,59 @@ __version__ = "0.1.0" + +import sys +import os + +# --------------------------------------------------------------------------- +# Record whether NumPy was already imported BEFORE this package was imported. +# This is the key information for determining BLAS thread safety later. +# --------------------------------------------------------------------------- +_NUMPY_IMPORTED_BEFORE_ATOMIC = ("numpy" in sys.modules) + + +# --------------------------------------------------------------------------- +# Record whether BLAS environment variables already indicate single-threaded +# execution. This allows us to avoid warnings even when NumPy was imported +# earlier, because the user may have correctly configured BLAS beforehand. +# --------------------------------------------------------------------------- + +def _env_says_single_thread() -> bool: + for var in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", + "OPENBLAS_NUM_THREADS", "NUMEXPR_NUM_THREADS"): + val = os.environ.get(var) + if val is None: + return False + if val.strip() not in ("1", "1.0"): + return False + return True + +_BLAS_ENV_SINGLE_THREADED = _env_says_single_thread() + +# --------------------------------------------------------------------------- +# Record whether threadpoolctl is installed. +# --------------------------------------------------------------------------- +def _threadpoolctl_installed() -> bool: + try: + import threadpoolctl + return True + except ImportError: + return False + +_THREADPOOLCTL_INSTALLED = _threadpoolctl_installed() + +# Force MKL / OpenBLAS into single-thread mode without any extra dependency. +if not _BLAS_ENV_SINGLE_THREADED: + os.environ["OMP_NUM_THREADS"] = "1" + os.environ["MKL_NUM_THREADS"] = "1" + os.environ["OPENBLAS_NUM_THREADS"] = "1" + os.environ["NUMEXPR_NUM_THREADS"] = "1" +else: + # Do not override user settings when NumPy is already loaded. + pass + + + # Main solver class from .solver import AtomicDFTSolver diff --git a/utils/atom/mesh/builder.py b/utils/atom/mesh/builder.py index b7519298..174e3037 100644 --- a/utils/atom/mesh/builder.py +++ b/utils/atom/mesh/builder.py @@ -131,11 +131,11 @@ def gauss_legendre(n: int) -> Tuple[np.ndarray, np.ndarray]: assert n > 0, \ QUADRATURE_POINT_NUMBER_NOT_GREATER_THAN_0_ERROR.format(n) - beta = 0.5/np.sqrt(1 - 1 / (2*(np.arange(1,n)))**(2)) - T = np.diag(beta,k=1) + np.diag(beta,k=-1) - nodes, eigvecs = np.linalg.eigh(T,UPLO='L') + beta = 0.5 / np.sqrt(1.0 - 1.0 / (2.0 * (np.arange(1, n)))**2) + T = np.diag(beta, k=1) + np.diag(beta, k=-1) + nodes, eigvecs = np.linalg.eigh(T, UPLO='L') index = np.arange(nodes.size) - legendre_weights = 2*(eigvecs[0,index])**2 + legendre_weights = 2*(eigvecs[0, index])**2 return nodes, legendre_weights diff --git a/utils/atom/pseudo/local.py b/utils/atom/pseudo/local.py index 12c7a6ad..3901ac1b 100644 --- a/utils/atom/pseudo/local.py +++ b/utils/atom/pseudo/local.py @@ -185,7 +185,7 @@ def get_rho_core_correction(self, r_nodes: np.ndarray) -> np.ndarray: self.density_nlcc, k=3 ) - rho_nlcc[indices_below_cutoff[0]] = rho_nlcc_interpolator(r_below_cutoff) + rho_nlcc[indices_below_cutoff[0]] = rho_nlcc_interpolator(r_below_cutoff).reshape(-1) return rho_nlcc @@ -318,8 +318,8 @@ def print_info(self): # Basic atomic information - print(f"Valence Charge (z_valence) : {self.z_valence}") - print(f"Nuclear Charge (z_nuclear) : {self.z_nuclear}") + print(f" Valence Charge (z_valence) : {self.z_valence}") + print(f" Nuclear Charge (z_nuclear) : {self.z_nuclear}") print() if self.load_psp: diff --git a/utils/atom/scf/convergence.py b/utils/atom/scf/convergence.py index 68e80dc1..536a0f5a 100644 --- a/utils/atom/scf/convergence.py +++ b/utils/atom/scf/convergence.py @@ -10,6 +10,8 @@ # Convergence Checker Error Messages LOOP_TYPE_ERROR_MESSAGE = \ "parameter loop_type must be 'Inner' or 'Outer', get {} instead" +FOOTER_BLANK_LINE_TYPE_ERROR_MESSAGE = \ + "parameter footer_blank_line must be a boolean, get type {} instead" @dataclass class ConvergenceHistory: @@ -36,11 +38,12 @@ class ConvergenceChecker: def __init__( self, - tolerance: float = 1e-6, - n_consecutive: int = 1, - norm_type: int = 2, - loop_type: str = "Inner" - ): + tolerance : float = 1e-6, + n_consecutive : int = 1, + norm_type : int = 2, + loop_type : str = "Inner", + footer_blank_line: bool = False, + ): """ Parameters ---------- @@ -52,17 +55,28 @@ def __init__( Norm type for residual calculation (1, 2, or np.inf) loop_type : str Type of loop for display purposes ("Inner" or "Outer") + footer_blank_line: bool + If True, print a blank line after the footer """ assert loop_type in ["Inner", "Outer"], \ LOOP_TYPE_ERROR_MESSAGE.format(loop_type) - self.tolerance = tolerance - self.n_consecutive = n_consecutive - self.norm_type = norm_type - self.loop_type = loop_type - + self.tolerance = tolerance + self.n_consecutive = n_consecutive + self.norm_type = norm_type + self.loop_type = loop_type + self._footer_blank_line = footer_blank_line + self._consecutive_count = 0 self._history = ConvergenceHistory(iterations=[], residuals=[]) + + def set_footer_blank_line(self, footer_blank_line: bool): + """ + Set the footer blank line flag + """ + assert isinstance(footer_blank_line, bool), \ + FOOTER_BLANK_LINE_TYPE_ERROR_MESSAGE.format(type(footer_blank_line)) + self._footer_blank_line = footer_blank_line def check( @@ -170,8 +184,10 @@ def print_footer(self, converged: bool, n_iterations: int, prefix: str = ""): """Print table footer with final status""" status_msg = "converged" if converged else "not converged" if self.loop_type == "Outer": - print(f"[Outer] Converged after {n_iterations} iteration(s)") + print(f"{prefix} Outer SCF iteration converged after {n_iterations} iteration(s)") else: print(f"{prefix}\t SCF Iteration {status_msg} after {n_iterations} iteration(s)") + + if self._footer_blank_line: print() diff --git a/utils/atom/scf/driver.py b/utils/atom/scf/driver.py index 4891bbe3..9c9a499d 100644 --- a/utils/atom/scf/driver.py +++ b/utils/atom/scf/driver.py @@ -1,18 +1,45 @@ from __future__ import annotations +from matplotlib import use import numpy as np -from typing import Dict, Any, Optional, Union, List +from typing import Dict, Any, Optional, Union, List, Tuple, Iterable from dataclasses import dataclass, field -from .hamiltonian import HamiltonianBuilder, SwitchesFlags +from .hamiltonian import HamiltonianBuilder from .density import DensityCalculator, DensityData from .poisson import PoissonSolver from .convergence import ConvergenceChecker from .eigensolver import EigenSolver from .mixer import Mixer -from ..utils.occupation_states import OccupationInfo -from ..xc.evaluator import XCEvaluator, create_xc_evaluator, XCPotentialData +from ..xc.evaluator import XCEvaluator, create_xc_evaluator from ..xc.functional_requirements import get_functional_requirements, FunctionalRequirements from ..xc.hybrid import HartreeFockExchange +from ..xc.oep import OEPCalculator +from ..mesh.operators import RadialOperatorsBuilder +from ..utils.occupation_states import OccupationInfo + + +# Switches Flags Error Messages +XC_FUNCTIONAL_NOT_STRING_ERROR = \ + "parameter xc_functional must be a string, get type {} instead" +HYBRID_MIXING_PARAMETER_NOT_A_FLOAT_ERROR = \ + "parameter hybrid_mixing_parameter must be a float, get type {} instead" +USE_XC_FUNCTIONAL_TYPE_ERROR = \ + "parameter use_xc_functional must be a boolean, get type {} instead" +USE_HF_EXCHANGE_TYPE_ERROR = \ + "parameter use_hf_exchange must be a boolean, get type {} instead" +USE_OEP_EXCHANGE_TYPE_ERROR = \ + "parameter use_oep_exchange must be a boolean, get type {} instead" +USE_OEP_CORRELATION_TYPE_ERROR = \ + "parameter use_oep_correlation must be a boolean, get type {} instead" + +USE_OEP_NOT_BOOL_ERROR = \ + "parameter use_oep must be a boolean, get type {} instead" +USE_OEP_SHOULD_NOT_BE_NONE_FOR_PBE0_ERROR = \ + "parameter use_oep should not be None for PBE0 functional, get None instead" +USE_OEP_NOT_TRUE_FOR_OEP_FUNCTIONAL_ERROR = \ + "parameter use_oep must be True for OEP functional `{}`, get {} instead" +USE_OEP_NOT_FALSE_FOR_NON_OEP_FUNCTIONAL_ERROR = \ + "parameter use_oep must be False for non-OEP functional `{}`, get {} instead" # SCF Settings Error Messages @@ -34,12 +61,23 @@ "parameter eigenvalues must be a numpy array, get type {} instead" EIGENVECTORS_TYPE_ERROR_MESSAGE = \ "parameter eigenvectors must be a numpy array, get type {} instead" +DENSITY_DATA_TYPE_ERROR_MESSAGE = \ + "parameter density_data must be a DensityData instance, get type {} instead" +FULL_EIGENVALUES_TYPE_ERROR_MESSAGE = \ + "parameter full_eigenvalues must be a numpy array, get type {} instead" +FULL_EIGENVECTORS_TYPE_ERROR_MESSAGE = \ + "parameter full_eigenvectors must be a numpy array, get type {} instead" +FULL_L_TERMS_TYPE_ERROR_MESSAGE = \ + "parameter full_l_terms must be a numpy array, get type {} instead" + RHO_TYPE_ERROR_MESSAGE = \ "parameter rho must be a numpy array, get type {} instead" CONVERGED_TYPE_ERROR_MESSAGE = \ "parameter converged must be a boolean, get type {} instead" ITERATIONS_TYPE_ERROR_MESSAGE = \ "parameter iterations must be an integer, get type {} instead" +RHO_RESIDUAL_TYPE_ERROR_MESSAGE = \ + "parameter rho_residual must be a float, get type {} instead" RESIDUAL_TYPE_ERROR_MESSAGE = \ "parameter residual must be a float, get type {} instead" TOTAL_ENERGY_TYPE_ERROR_MESSAGE = \ @@ -50,7 +88,8 @@ "parameter outer_iterations must be an integer, get type {} instead" OUTER_CONVERGED_TYPE_ERROR_MESSAGE = \ "parameter outer_converged must be a boolean, get type {} instead" - +ENABLE_PARALLELIZATION_NOT_BOOL_ERROR = \ + "parameter enable_parallelization must be a boolean, get type {} instead" # SCF Driver Error Messages HAMILTONIAN_BUILDER_TYPE_ERROR_MESSAGE = \ @@ -65,16 +104,87 @@ "parameter mixer must be a Mixer, get type {} instead" OCCUPATION_INFO_TYPE_ERROR_MESSAGE = \ "parameter occupation_info must be a OccupationInfo, get type {} instead" +HYBRID_MIXING_PARAMETER_TYPE_ERROR_MESSAGE = \ + "parameter hybrid_mixing_parameter must be a float, get type {} instead" +OPS_BUILDER_OEP_TYPE_ERROR_MESSAGE = \ + "parameter ops_builder_oep must be a RadialOperatorsBuilder, get type {} instead" +OEP_MIXING_PARAMETER_TYPE_ERROR_MESSAGE = \ + "parameter oep_mixing_parameter must be a float, get type {} instead" +FREQUENCY_QUADRATURE_POINT_NUMBER_TYPE_ERROR_MESSAGE = \ + "parameter frequency_quadrature_point_number must be an integer, get type {} instead" +ANGULAR_MOMENTUM_CUTOFF_TYPE_ERROR_MESSAGE = \ + "parameter angular_momentum_cutoff must be an integer, get type {} instead" +ANGULAR_MOMENTUM_CUTOFF_NOT_NONE_ERROR_MESSAGE = \ + "parameter angular_momentum_cutoff must be not None, get None instead" RHO_INITIAL_TYPE_ERROR_MESSAGE = \ "parameter rho_initial must be a numpy array, get type {} instead" SETTINGS_TYPE_ERROR_MESSAGE = \ "parameter settings must be a SCFSettings or a dictionary, get type {} instead" +V_X_OEP_TYPE_ERROR_MESSAGE = \ + "parameter v_x_oep must be a numpy array, get type {} instead" +V_C_OEP_TYPE_ERROR_MESSAGE = \ + "parameter v_c_oep must be a numpy array, get type {} instead" H_HF_EXCHANGE_DICT_BY_L_TYPE_ERROR_MESSAGE = \ "provided parameter H_hf_exchange_dict_by_l must be a dictionary, get type {} instead" ORBITALS_INITIAL_TYPE_ERROR_MESSAGE = \ "parameter orbitals_initial must be a numpy array, get type {} instead" +XC_FUNCTIONAL_TYPE_ERROR_MESSAGE = \ + "parameter xc_functional must be a string, get type {} instead" +SWITCHES_TYPE_ERROR_MESSAGE = \ + "parameter switches must be a SwitchesFlags, get type {} instead" +XC_REQUIREMENTS_TYPE_ERROR_MESSAGE = \ + "parameter xc_requirements must be a FunctionalRequirements, get type {} instead" +XC_CALCULATOR_TYPE_ERROR_MESSAGE = \ + "parameter xc_calculator must be a XCEvaluator, get type {} instead" + + +OCC_EIGENVALUES_LIST_NOT_LIST_ERROR_MESSAGE = \ + "parameter occ_eigenvalues_list must be a list, get type {} instead" +OCC_EIGENVECTORS_LIST_NOT_LIST_ERROR_MESSAGE = \ + "parameter occ_eigenvectors_list must be a list, get type {} instead" +UNOCC_EIGENVALUES_LIST_NOT_LIST_ERROR_MESSAGE = \ + "parameter unocc_eigenvalues_list must be a list, get type {} instead" +UNOCC_EIGENVECTORS_LIST_NOT_LIST_ERROR_MESSAGE = \ + "parameter unocc_eigenvectors_list must be a list, get type {} instead" +OCC_EIGENVALUES_LIST_LENGTH_ERROR_MESSAGE = \ + "Number of occupied eigenvalues in 'occ_eigenvalues_list' must match number of unique l values, get {} instead" +OCC_EIGENVECTORS_LIST_LENGTH_ERROR_MESSAGE = \ + "Number of occupied eigenvectors in 'occ_eigenvectors_list' must match number of unique l values, get {} instead" +UNOCC_EIGENVALUES_LIST_LENGTH_ERROR_MESSAGE = \ + "Number of unoccupied eigenvalues in 'unocc_eigenvalues_list' must match number of unique l values, get {} instead" +UNOCC_EIGENVECTORS_LIST_LENGTH_ERROR_MESSAGE = \ + "Number of unoccupied eigenvectors in 'unocc_eigenvectors_list' must match number of unique l values, get {} instead" +OCC_EIGENVALUES_LIST_NDIM_ERROR_MESSAGE = \ + "Occupied eigenvalues in 'occ_eigenvalues_list' must be 1D arrays, get dimension {} instead" +OCC_EIGENVECTORS_LIST_NDIM_ERROR_MESSAGE = \ + "Occupied eigenvectors in 'occ_eigenvectors_list' must be 2D arrays, get dimension {} instead" +UNOCC_EIGENVALUES_LIST_NDIM_ERROR_MESSAGE = \ + "Unoccupied eigenvalues in 'unocc_eigenvalues_list' must be 1D arrays, get dimension {} instead" +UNOCC_EIGENVECTORS_LIST_NDIM_ERROR_MESSAGE = \ + "Unoccupied eigenvectors in 'unocc_eigenvectors_list' must be 2D arrays, get dimension {} instead" +OCC_EIGENVALUES_AND_EIGENVECTORS_LIST_SHAPE_MISMATCH_ERROR_MESSAGE = \ + "Occupied eigenvalues and eigenvectors in 'occ_eigenvalues_list' and 'occ_eigenvectors_list' must have the same number of rows, get {} and {} instead" +UNOCC_EIGENVALUES_AND_EIGENVECTORS_LIST_SHAPE_MISMATCH_ERROR_MESSAGE = \ + "Unoccupied eigenvalues and eigenvectors in 'unocc_eigenvalues_list' and 'unocc_eigenvectors_list' must have the same number of rows, get {} and {} instead" +OCC_EIGENVECTORS_LIST_SHAPE_ERROR_MESSAGE = \ + "Occupied eigenvectors in 'occ_eigenvectors_list' must have the same number of rows as the number of interior nodes, get {} and {} instead" +UNOCC_EIGENVECTORS_LIST_SHAPE_ERROR_MESSAGE = \ + "Unoccupied eigenvectors in 'unocc_eigenvectors_list' must have the same number of rows as the number of interior nodes, get {} and {} instead" + + +# Switches Flags Warning Messages +HYBRID_MIXING_PARAMETER_NOT_PROVIDED_WARNING = \ + "WARNING: hybrid_mixing_parameter not provided for {xc_functional} functional, using default value {default_value}" +HYBRID_MIXING_PARAMETER_NOT_1_0_WARNING = \ + "WARNING: hybrid_mixing_parameter for {xc_functional} should be 1.0, got {hybrid_mixing_parameter}" +HYBRID_MIXING_PARAMETER_NOT_0_0_WARNING = \ + "WARNING: hybrid_mixing_parameter for {xc_functional} should be 0.0, got {hybrid_mixing_parameter}" +HYBRID_MIXING_PARAMETER_NOT_0_25_WARNING = \ + "WARNING: hybrid_mixing_parameter for {xc_functional} should be 0.25, got {hybrid_mixing_parameter}" + + # SCF Driver Warning Messages INNER_SCF_DID_NOT_CONVERGE_WARNING = \ @@ -83,6 +193,190 @@ "WARNING: Hartree-Fock calculator is not available, please initialize the HF calculator first" + +class SwitchesFlags: + """ + Internal helper class for HamiltonianBuilder to manage functional switches. + Determines which Hamiltonian components to include based on the XC functional. + + Attributes: + use_xc_functional (bool) : Whether to use XC functional + use_hf_exchange (bool) : Whether to include Hartree-Fock exchange mixing + use_oep_exchange (bool) : Whether to use OEP exchange potential + use_oep_correlation (bool) : Whether to use OEP correlation potential, for RPA + use_metagga (bool) : Whether the functional requires meta-GGA terms + """ + def __init__( + self, + xc_functional : str, + use_oep : Optional[bool] = None, + hybrid_mixing_parameter: Optional[float] = None, + ): + """ + Initialize switches based on XC functional and hybrid mixing parameter. + + Parameters + ---------- + xc_functional : str + Name of XC functional + use_oep : bool, optional + Whether to use OEP exchange potential, useful only for PBE0 functional, + otherwise it has to agree with use_oep_exchange flag + hybrid_mixing_parameter : float, optional + Mixing parameter for hybrid functionals (0-1). Required only for hybrid functionals. + - For PBE0: should be 0.25 + - For HF : should be 1.0 + """ + # Type checking + assert isinstance(xc_functional, str), \ + XC_FUNCTIONAL_NOT_STRING_ERROR.format(type(xc_functional)) + + if use_oep is not None: + assert isinstance(use_oep, bool), \ + USE_OEP_NOT_BOOL_ERROR.format(type(use_oep)) + if hybrid_mixing_parameter is not None: + assert isinstance(hybrid_mixing_parameter, (float, int)), \ + HYBRID_MIXING_PARAMETER_NOT_A_FLOAT_ERROR.format(type(hybrid_mixing_parameter)) + hybrid_mixing_parameter = float(hybrid_mixing_parameter) + + self.check_xc_functional_and_use_oep_consistency(use_oep, xc_functional) + + + self.xc_functional = xc_functional + self.hybrid_mixing_parameter = hybrid_mixing_parameter + + # Initialize all flags to False, except use_xc_functional which is True by default + self.use_xc_functional = True + self.use_hf_exchange = False + self.use_oep_exchange = False + self.use_oep_correlation = False + self.use_metagga = False + + # Set flags based on functional + if xc_functional == 'None': + self.use_xc_functional = False + + # Exchange-only functionals + elif xc_functional == 'EXX': + self.use_oep_exchange = True + self.use_oep_correlation = False # EXX does not use correlation potential + self.use_xc_functional = False + if hybrid_mixing_parameter is not None: + print(HYBRID_MIXING_PARAMETER_NOT_PROVIDED_WARNING.format(xc_functional='EXX', default_value=1.0)) + self.hybrid_mixing_parameter = 1.0 + elif not np.isclose(hybrid_mixing_parameter, 1.0): + print(HYBRID_MIXING_PARAMETER_NOT_1_0_WARNING.format(xc_functional='EXX', hybrid_mixing_parameter=hybrid_mixing_parameter)) + + # Exchange-correlation functionals, for RPA + elif xc_functional == 'RPA': + self.use_oep_exchange = True + self.use_oep_correlation = True + self.use_xc_functional = False + if hybrid_mixing_parameter is not None: + print(HYBRID_MIXING_PARAMETER_NOT_PROVIDED_WARNING.format(xc_functional='RPA', default_value=1.0)) + self.hybrid_mixing_parameter = 1.0 + elif not np.isclose(hybrid_mixing_parameter, 1.0): + print(HYBRID_MIXING_PARAMETER_NOT_1_0_WARNING.format(xc_functional='RPA', hybrid_mixing_parameter=hybrid_mixing_parameter)) + + # Hartree-Fock exchange functionals + elif xc_functional == 'HF': + self.use_hf_exchange = True + self.use_xc_functional = False + # Check if hybrid_mixing_parameter is provided for hybrid functionals + if hybrid_mixing_parameter is None: + print(HYBRID_MIXING_PARAMETER_NOT_PROVIDED_WARNING.format(xc_functional='HF', default_value=1.0)) + self.hybrid_mixing_parameter = 1.0 + elif not np.isclose(hybrid_mixing_parameter, 1.0): + print(HYBRID_MIXING_PARAMETER_NOT_1_0_WARNING.format(xc_functional='HF', hybrid_mixing_parameter=hybrid_mixing_parameter)) + + # Hybrid GGA functionals + elif xc_functional == 'PBE0': + # Check if hybrid_mixing_parameter is provided for hybrid functionals + if hybrid_mixing_parameter is None: + print(HYBRID_MIXING_PARAMETER_NOT_PROVIDED_WARNING.format(xc_functional='PBE0', default_value=0.25)) + self.hybrid_mixing_parameter = 0.25 + elif not np.isclose(hybrid_mixing_parameter, 0.25): + print(HYBRID_MIXING_PARAMETER_NOT_0_25_WARNING.format(xc_functional='PBE0', hybrid_mixing_parameter=hybrid_mixing_parameter)) + + if use_oep is False: + self.use_hf_exchange = True + self.use_oep_exchange = False + else: + self.use_oep_exchange = True + self.use_hf_exchange = False + + # Set meta-GGA flag for functionals that require de_xc_dtau + elif xc_functional in ['SCAN', 'RSCAN', 'R2SCAN']: + self.use_metagga = True + + # LDA/GGA functionals: no special flags (default False) + elif xc_functional in ['LDA_PZ', 'LDA_PW', 'GGA_PBE']: + pass + + # Invalid XC functional + else: + raise ValueError(f"Invalid XC functional: {xc_functional}") + + + + @property + def use_oep(self) -> bool: + return self.use_oep_exchange or self.use_oep_correlation + + + @staticmethod + def check_xc_functional_and_use_oep_consistency(use_oep: bool, xc_functional: str): + """ + Check if the consistency between use_oep and xc_functional is correct + """ + # PBE0 functional must be used with OEP flag + if xc_functional == 'PBE0': + assert use_oep is not None, \ + USE_OEP_SHOULD_NOT_BE_NONE_FOR_PBE0_ERROR.format(xc_functional) + return + + # For other functionals, use_oep can be None, and no extra check is needed + if use_oep is None: + return + + + if xc_functional in ['EXX', 'RPA']: + assert use_oep is True, \ + USE_OEP_NOT_TRUE_FOR_OEP_FUNCTIONAL_ERROR.format(xc_functional) + elif xc_functional in ['LDA_PZ', 'LDA_PW', 'GGA_PBE', 'SCAN', 'RSCAN', 'R2SCAN', 'HF']: + assert use_oep is False, \ + USE_OEP_NOT_FALSE_FOR_NON_OEP_FUNCTIONAL_ERROR.format(xc_functional) + else: + raise ValueError(f"Invalid XC functional: {xc_functional}") + + + + def print_info(self): + """ + Print functional switch information summary. + """ + print("=" * 60) + print("\t\t FUNCTIONAL SWITCHES") + print("=" * 60) + print(f"\t XC Functional : {self.xc_functional}") + if self.hybrid_mixing_parameter is not None: + print(f"\t Hybrid Mixing Parameter : {self.hybrid_mixing_parameter}") + else: + print(f"\t Hybrid Mixing Parameter : None (not applicable)") + print() + print("\t FUNCTIONAL FLAGS:") + print(f"\t use_xc_functional : {self.use_xc_functional}") + print(f"\t use_hf_exchange : {self.use_hf_exchange}") + print(f"\t use_oep_exchange : {self.use_oep_exchange}") + print(f"\t use_oep_correlation : {self.use_oep_correlation}") + print(f"\t use_metagga : {self.use_metagga}") + print() + + + + + + @dataclass class SCFSettings: """ @@ -105,16 +399,16 @@ class SCFSettings: """ # Inner loop settings - inner_max_iter: int = 200 - rho_tol: float = 1e-6 - n_consecutive: int = 1 + inner_max_iter : int = 200 + rho_tol : float = 1e-6 + n_consecutive : int = 1 # Outer loop settings - outer_max_iter: int = 1 - outer_rho_tol: float = 1e-5 + outer_max_iter : int = 1 + outer_rho_tol : float = 1e-5 # Output settings - print_debug: bool = False + print_debug : bool = False def __post_init__(self): @@ -149,12 +443,12 @@ def from_dict(cls, settings_dict: dict) -> SCFSettings: Settings object """ return cls( - inner_max_iter=settings_dict.get('inner_max_iter', 200), - rho_tol=settings_dict.get('rho_tol', 1e-6), - n_consecutive=settings_dict.get('n_consecutive', 1), - outer_max_iter=settings_dict.get('outer_max_iter', 1), - outer_rho_tol=settings_dict.get('outer_rho_tol', 1e-5), - print_debug=settings_dict.get('print_debug', False) + inner_max_iter = settings_dict.get('inner_max_iter' , 200), + rho_tol = settings_dict.get('rho_tol' , 1e-6), + n_consecutive = settings_dict.get('n_consecutive' , 1), + outer_max_iter = settings_dict.get('outer_max_iter' , 1), + outer_rho_tol = settings_dict.get('outer_rho_tol' , 1e-5), + print_debug = settings_dict.get('print_debug' , False) ) @@ -186,61 +480,95 @@ class SCFResult: ---------- eigen_energies : np.ndarray Kohn-Sham eigenvalues (orbital energies) for all states, shape (n_states,) + orbitals : np.ndarray Converged Kohn-Sham orbitals (radial wavefunctions R_nl(r)) Shape: (n_states, n_quad_points) + density_data : DensityData Converged electron density and related quantities (rho, grad_rho, tau) + + full_eigen_energies : np.ndarray, optional + Optional storage of the complete KS spectrum (occupied + unoccupied). + Shape: (n_full_states,). + + full_orbitals : np.ndarray, optional + Optional storage of all radial orbitals associated with full_eigen_energies. + Shape: (n_full_states, n_quad_points). + + full_l_terms : np.ndarray, optional + Optional storage of the angular momentum index for each entry in full_eigen_energies. + Shape: (n_full_states,). + converged : bool Whether inner SCF loop converged + iterations : int Number of inner SCF iterations performed + rho_residual : float Final density residual of inner loop (L2 norm) + outer_iterations : int, optional Number of outer SCF iterations (for HF/OEP/RPA methods) + outer_converged : bool, optional Whether outer SCF loop converged + total_energy : float, optional Total energy of the system + energy_components : dict, optional Breakdown of energy components (kinetic, Hartree, XC, etc.) """ # Core results - eigen_energies: np.ndarray - orbitals: np.ndarray - density_data: DensityData + eigen_energies : np.ndarray + orbitals : np.ndarray + density_data : DensityData # Inner loop convergence info - converged: bool - iterations: int - rho_residual: float + converged : bool + iterations : int + rho_residual : float # Outer loop info (optional) - outer_iterations : Optional[int] = None - outer_converged : Optional[bool] = None + full_eigen_energies : Optional[np.ndarray] = None + full_orbitals : Optional[np.ndarray] = None + full_l_terms : Optional[np.ndarray] = None + + outer_iterations : Optional[int] = None + outer_converged : Optional[bool] = None # Energy info (optional) - total_energy : Optional[float] = None - energy_components : Optional[dict] = field(default=None) + total_energy : Optional[float] = None + energy_components : Optional[dict] = field(default=None) def __post_init__(self): # type check for required fields assert isinstance(self.eigen_energies, np.ndarray), \ - "eigen_energies must be a numpy array, got {}".format(type(self.eigen_energies)) + EIGENVALUES_TYPE_ERROR_MESSAGE.format(type(self.eigen_energies)) assert isinstance(self.orbitals, np.ndarray), \ - "orbitals must be a numpy array, got {}".format(type(self.orbitals)) + EIGENVECTORS_TYPE_ERROR_MESSAGE.format(type(self.orbitals)) assert isinstance(self.density_data, DensityData), \ - "density_data must be a DensityData instance, got {}".format(type(self.density_data)) + DENSITY_DATA_TYPE_ERROR_MESSAGE.format(type(self.density_data)) assert isinstance(self.converged, bool), \ CONVERGED_TYPE_ERROR_MESSAGE.format(type(self.converged)) assert isinstance(self.iterations, int), \ ITERATIONS_TYPE_ERROR_MESSAGE.format(type(self.iterations)) assert isinstance(self.rho_residual, (float, np.floating)), \ - "rho_residual must be a float, got {}".format(type(self.rho_residual)) + RHO_RESIDUAL_TYPE_ERROR_MESSAGE.format(type(self.rho_residual)) # type check for optional fields (only if not None) + if self.full_eigen_energies is not None: + assert isinstance(self.full_eigen_energies, np.ndarray), \ + FULL_EIGENVALUES_TYPE_ERROR_MESSAGE.format(type(self.full_eigen_energies)) + if self.full_orbitals is not None: + assert isinstance(self.full_orbitals, np.ndarray), \ + FULL_EIGENVECTORS_TYPE_ERROR_MESSAGE.format(type(self.full_orbitals)) + if self.full_l_terms is not None: + assert isinstance(self.full_l_terms, np.ndarray), \ + FULL_L_TERMS_TYPE_ERROR_MESSAGE.format(type(self.full_l_terms)) if self.outer_iterations is not None: assert isinstance(self.outer_iterations, int), \ OUTER_ITERATIONS_TYPE_ERROR_MESSAGE.format(type(self.outer_iterations)) @@ -271,16 +599,19 @@ def from_dict(cls, result_dict: dict) -> SCFResult: Result object """ return cls( - eigenvalues=result_dict['eigenvalues'], - eigenvectors=result_dict['eigenvectors'], - rho=result_dict['rho'], - converged=result_dict['converged'], - iterations=result_dict['iterations'], - residual=result_dict['residual'], - outer_iterations=result_dict.get('outer_iterations'), - outer_converged=result_dict.get('outer_converged'), - total_energy=result_dict.get('total_energy'), - energy_components=result_dict.get('energy_components') + eigenvalues = result_dict['eigenvalues'], + eigenvectors = result_dict['eigenvectors'], + rho = result_dict['rho'], + converged = result_dict['converged'], + iterations = result_dict['iterations'], + residual = result_dict['residual'], + full_eigen_energies = result_dict.get('full_eigen_energies'), + full_orbitals = result_dict.get('full_orbitals'), + full_l_terms = result_dict.get('full_l_terms'), + outer_iterations = result_dict.get('outer_iterations'), + outer_converged = result_dict.get('outer_converged'), + total_energy = result_dict.get('total_energy'), + energy_components = result_dict.get('energy_components') ) def to_dict(self) -> dict: @@ -293,15 +624,21 @@ def to_dict(self) -> dict: Dictionary representation of results """ result = { - 'eigenvalues': self.eigenvalues, + 'eigenvalues' : self.eigenvalues, 'eigenvectors': self.eigenvectors, - 'rho': self.rho, - 'converged': self.converged, - 'iterations': self.iterations, - 'residual': self.residual + 'rho' : self.rho, + 'converged' : self.converged, + 'iterations' : self.iterations, + 'residual' : self.residual } # Add optional fields if present + if self.full_eigen_energies is not None: + result['full_eigen_energies'] = self.full_eigen_energies + if self.full_orbitals is not None: + result['full_orbitals'] = self.full_orbitals + if self.full_l_terms is not None: + result['full_l_terms'] = self.full_l_terms if self.outer_iterations is not None: result['outer_iterations'] = self.outer_iterations if self.outer_converged is not None: @@ -366,14 +703,20 @@ class SCFDriver: def __init__( self, - hamiltonian_builder : HamiltonianBuilder, - density_calculator : DensityCalculator, - poisson_solver : PoissonSolver, - eigensolver : EigenSolver, - mixer : Mixer, - occupation_info : OccupationInfo, - xc_functional : str, - hybrid_mixing_parameter : Optional[float] = None + hamiltonian_builder : HamiltonianBuilder, + density_calculator : DensityCalculator, + poisson_solver : PoissonSolver, + eigensolver : EigenSolver, + mixer : Mixer, + occupation_info : OccupationInfo, + xc_functional : str, + use_oep : Optional[bool] = None, + hybrid_mixing_parameter : Optional[float] = None, + ops_builder_oep : Optional[RadialOperatorsBuilder] = None, + oep_mixing_parameter : Optional[float] = None, + frequency_quadrature_point_number : Optional[int] = None, # parameter for RPA correlation potential + angular_momentum_cutoff : Optional[int] = None, # parameter for RPA functional only + enable_parallelization : Optional[bool] = None, # parameter for RPA calculations only ): """ Parameters @@ -394,6 +737,9 @@ def __init__( Name of XC functional (e.g., 'LDA_PZ', 'GGA_PBE', 'SCAN') Used to determine what density-related quantities to compute Also used to initialize the XC calculator internally + use_oep : bool, optional + Whether to use OEP exchange potential, useful only for PBE0 functional, + otherwise it has to agree with use_oep_exchange flag hybrid_mixing_parameter : float, optional Mixing parameter for hybrid functionals (HF exchange fraction) Required only for hybrid functionals (PBE0, HF) @@ -401,39 +747,67 @@ def __init__( - For HF: 1.0 For non-hybrid functionals, this parameter is ignored This parameter is designed to be autodiff-compatible for delta learning + ops_builder_oep : RadialOperatorsBuilder, optional + Dedicated operators builder for OEP basis/projectors. It must be provided when OEP is enabled, otherwise it will be ignored. + oep_mixing_parameter : float, optional + Scaling parameter (λ) applied to OEP exchange/correlation potentials. + frequency_quadrature_point_number : int, optional + Number of frequency quadrature points for RPA correlation potential. + angular_momentum_cutoff : int, optional + Maximum angular momentum quantum number to include for RPA functional. """ - self.hamiltonian_builder = hamiltonian_builder - self.density_calculator = density_calculator - self.poisson_solver = poisson_solver - self.eigensolver = eigensolver - self.mixer = mixer - self.occupation_info = occupation_info - self.xc_functional = xc_functional - self.hybrid_mixing_parameter = hybrid_mixing_parameter - + + self.hamiltonian_builder = hamiltonian_builder + self.density_calculator = density_calculator + self.poisson_solver = poisson_solver + self.eigensolver = eigensolver + self.mixer = mixer + self.occupation_info = occupation_info + self.xc_functional = xc_functional + self.hybrid_mixing_parameter = hybrid_mixing_parameter + self.use_oep = use_oep + self.ops_builder_oep = ops_builder_oep + self.oep_mixing_parameter = oep_mixing_parameter + self.frequency_quadrature_point_number = frequency_quadrature_point_number + self.angular_momentum_cutoff = angular_momentum_cutoff + self.enable_parallelization = enable_parallelization + self._check_initialization() + + # Create SwitchesFlags instance (handles validation internally) - self.switches = SwitchesFlags(xc_functional, hybrid_mixing_parameter) + self.switches = SwitchesFlags( + xc_functional = xc_functional, + use_oep = use_oep, + hybrid_mixing_parameter = hybrid_mixing_parameter + ) # Get functional requirements (what to compute for this functional) self.xc_requirements : FunctionalRequirements = get_functional_requirements(xc_functional) # Initialize XC calculator internally based on functional self.xc_calculator : Optional[XCEvaluator] = self._initialize_xc_calculator( - derivative_matrix=density_calculator.derivative_matrix, - r_quad=density_calculator.quadrature_nodes + derivative_matrix = density_calculator.derivative_matrix, + r_quad = density_calculator.quadrature_nodes ) # Initialize HF exchange calculator for hybrid functionals self.hf_calculator : Optional[HartreeFockExchange] = self._initialize_hf_calculator() - + + # Initialize OEP calculator for OEP calculations + self.oep_calculator : Optional[OEPCalculator] = self._initialize_oep_calculator() + # Extract NLCC density from pseudopotential (if using pseudopotential) - self.rho_nlcc : np.ndarray = self._initialize_nlcc_density() + self.rho_nlcc : Optional[np.ndarray] = self._initialize_nlcc_density() # Initialize convergence checkers (will be configured in run method) self.inner_convergence_checker : Optional[ConvergenceChecker] = None self.outer_convergence_checker : Optional[ConvergenceChecker] = None - # type check for required parameters + + def _check_initialization(self): + """ + Check the initialization of the SCFDriver for required parameters. + """ assert isinstance(self.hamiltonian_builder, HamiltonianBuilder), \ HAMILTONIAN_BUILDER_TYPE_ERROR_MESSAGE.format(type(self.hamiltonian_builder)) assert isinstance(self.density_calculator, DensityCalculator), \ @@ -446,12 +820,35 @@ def __init__( MIXER_TYPE_ERROR_MESSAGE.format(type(self.mixer)) assert isinstance(self.occupation_info, OccupationInfo), \ OCCUPATION_INFO_TYPE_ERROR_MESSAGE.format(type(self.occupation_info)) - - + assert isinstance(self.xc_functional, str), \ + XC_FUNCTIONAL_TYPE_ERROR_MESSAGE.format(type(self.xc_functional)) + if self.hybrid_mixing_parameter is not None: + assert isinstance(self.hybrid_mixing_parameter, (float, np.floating)), \ + HYBRID_MIXING_PARAMETER_TYPE_ERROR_MESSAGE.format(type(self.hybrid_mixing_parameter)) + if self.use_oep is not None: + assert isinstance(self.use_oep, bool), \ + USE_OEP_NOT_BOOL_ERROR.format(type(self.use_oep)) + if self.ops_builder_oep is not None: + assert isinstance(self.ops_builder_oep, RadialOperatorsBuilder), \ + OPS_BUILDER_OEP_TYPE_ERROR_MESSAGE.format(type(self.ops_builder_oep)) + if self.oep_mixing_parameter is not None: + assert isinstance(self.oep_mixing_parameter, (float, np.floating)), \ + OEP_MIXING_PARAMETER_TYPE_ERROR_MESSAGE.format(type(self.oep_mixing_parameter)) + if self.frequency_quadrature_point_number is not None: + assert isinstance(self.frequency_quadrature_point_number, int), \ + FREQUENCY_QUADRATURE_POINT_NUMBER_TYPE_ERROR_MESSAGE.format(type(self.frequency_quadrature_point_number)) + if self.angular_momentum_cutoff is not None: + assert isinstance(self.angular_momentum_cutoff, int), \ + ANGULAR_MOMENTUM_CUTOFF_TYPE_ERROR_MESSAGE.format(type(self.angular_momentum_cutoff)) + if self.enable_parallelization is not None: + assert isinstance(self.enable_parallelization, bool), \ + ENABLE_PARALLELIZATION_NOT_BOOL_ERROR.format(type(self.enable_parallelization)) + + def _initialize_xc_calculator( self, - derivative_matrix: np.ndarray, - r_quad: np.ndarray + derivative_matrix : np.ndarray, + r_quad : np.ndarray ) -> Optional[XCEvaluator]: """ Initialize XC calculator based on the functional name. @@ -479,15 +876,13 @@ def _initialize_xc_calculator( ValueError If the specified functional is not implemented """ - if self.xc_functional in ['None', 'HF']: - return None - else: + if self.switches.use_xc_functional: return create_xc_evaluator( - functional_name=self.xc_functional, - derivative_matrix=derivative_matrix, - r_quad=r_quad + functional_name = self.xc_functional, + derivative_matrix = derivative_matrix, + r_quad = r_quad ) - + def _initialize_hf_calculator(self) -> Optional[HartreeFockExchange]: """ @@ -504,13 +899,42 @@ def _initialize_hf_calculator(self) -> Optional[HartreeFockExchange]: # Create HF exchange calculator with ops_builder and occupation_info hf_calculator = HartreeFockExchange( - ops_builder=self.hamiltonian_builder.ops_builder, - occupation_info=self.occupation_info + ops_builder = self.hamiltonian_builder.ops_builder, + occupation_info = self.occupation_info ) return hf_calculator + def _initialize_oep_calculator(self) -> Optional[OEPCalculator]: + """ + Initialize OEP calculator for OEP calculations. + + Returns + ------- + Optional[OEPCalculator] + OEP calculator if functional requires it, None otherwise + """ + # Only create OEP calculator for OEP calculations + if not self.switches.use_oep: + return None + else: + assert self.ops_builder_oep is not None, \ + "ops_builder_oep must be provided when OEP is enabled" + + # Create OEP calculator with ops_builder and occupation_info + oep_calculator = OEPCalculator( + ops_builder = self.hamiltonian_builder.ops_builder, + ops_builder_oep = self.ops_builder_oep, + occupation_info = self.occupation_info, + use_rpa_correlation = self.switches.use_oep_correlation, + frequency_quadrature_point_number = self.frequency_quadrature_point_number, + angular_momentum_cutoff = self.angular_momentum_cutoff, + ) + + return oep_calculator + + def _initialize_nlcc_density(self) -> np.ndarray: """ Initialize non-linear core correction (NLCC) density from pseudopotential. @@ -614,8 +1038,8 @@ def _compute_hf_exchange_matrices_dict(self, orbitals: np.ndarray) -> Dict[int, def run( self, - rho_initial : np.ndarray, - settings : Union[SCFSettings, Dict[str, Any]], + rho_initial : np.ndarray, + settings : Union[SCFSettings, Dict[str, Any]], orbitals_initial : Optional[np.ndarray] = None ) -> SCFResult: """ @@ -656,10 +1080,14 @@ def run( # Initialize outer convergence checker if outer loop is needed if needs_outer_loop: self.outer_convergence_checker = ConvergenceChecker( - tolerance = settings.outer_rho_tol, - n_consecutive = settings.n_consecutive, # Outer loop typically converges in 1 iteration - loop_type = "Outer" + tolerance = settings.outer_rho_tol, + n_consecutive = settings.n_consecutive, # Outer loop typically converges in 1 iteration + loop_type = "Outer", + footer_blank_line = True, ) + else: + self.inner_convergence_checker.set_footer_blank_line(True) + if hasattr(settings, 'print_debug') and settings.print_debug: print("="*60) @@ -736,6 +1164,8 @@ def _inner_loop( rho_initial : np.ndarray, settings : SCFSettings, orbitals_initial : Optional[np.ndarray] = None, + v_x_oep : Optional[np.ndarray] = None, + v_c_oep : Optional[np.ndarray] = None, H_hf_exchange_dict_by_l : Optional[Dict[int, np.ndarray]] = None, ) -> SCFResult: """ @@ -754,6 +1184,10 @@ def _inner_loop( Initial orbitals guess for debugging Shape: (n_grid, n_orbitals) If provided, will be used as initial orbitals instead of solving eigenvalue problem + v_x_oep : np.ndarray, optional + OEP exchange potential + v_c_oep : np.ndarray, optional + OEP correlation potential H_hf_exchange_dict_by_l : dict, optional Hartree-Fock exchange matrices dictionary (from outer loop) @@ -770,19 +1204,22 @@ def _inner_loop( SETTINGS_TYPE_ERROR_MESSAGE.format(type(settings)) if orbitals_initial is not None: assert isinstance(orbitals_initial, np.ndarray), \ - ORBITALS_INITIAL_TYPE_ERROR_MESSAGE.format(type(orbitals_initial)) + ORBITALS_INITIAL_TYPE_ERROR_MESSAGE.format(type(orbitals_initial)) + if v_x_oep is not None: + assert isinstance(v_x_oep, np.ndarray), \ + V_X_OEP_TYPE_ERROR_MESSAGE.format(type(v_x_oep)) + if v_c_oep is not None: + assert isinstance(v_c_oep, np.ndarray), \ + V_C_OEP_TYPE_ERROR_MESSAGE.format(type(v_c_oep)) if H_hf_exchange_dict_by_l is not None: assert isinstance(H_hf_exchange_dict_by_l, dict), \ - H_HF_EXCHANGE_DICT_BY_L_TYPE_ERROR_MESSAGE.format(type(H_hf_exchange_dict_by_l)) + H_HF_EXCHANGE_DICT_BY_L_TYPE_ERROR_MESSAGE.format(type(H_hf_exchange_dict_by_l)) # initialize variables max_iter = settings.inner_max_iter print_debug = settings.print_debug - # rho = rho_initial.copy() - - rho = self.density_calculator.normalize_density(rho_initial.copy()) - + rho = self.density_calculator.normalize_density(rho_initial.copy()) density_data = self.density_calculator.create_density_data_from_mixed( rho_mixed = rho, orbitals = orbitals_initial, @@ -815,6 +1252,7 @@ def _inner_loop( # For all-electron: rho_nlcc is zero, so rho_total = rho_valence v_x = np.zeros_like(rho) # Default: no exchange potential v_c = np.zeros_like(rho) # Default: no correlation potential + de_xc_dtau = None # Default: no meta-GGA derivative term if self.xc_calculator is not None: # Compute XC using new interface: DensityData → XCPotentialData @@ -824,50 +1262,59 @@ def _inner_loop( de_xc_dtau = xc_potential_data.de_xc_dtau # ===== Step 2: Build and solve for each l channel ===== - eigenvalues_all = [] - eigenvectors_all = [] - + occ_eigenvalues_list : List[np.ndarray] = [] + occ_eigenvectors_list : List[np.ndarray] = [] + unocc_eigenvalues_list : List[np.ndarray] = [] # needed only for OEP + unocc_eigenvectors_list : List[np.ndarray] = [] # needed only for OEP + for l in self.occupation_info.unique_l_values: # Build Hamiltonian for this l H_l = self.hamiltonian_builder.build_for_l_channel( - l = l, - v_hartree = v_hartree, - v_x = v_x, - v_c = v_c, - switches = self.switches, - de_xc_dtau = de_xc_dtau, - symmetrize = False + l = l, + v_hartree = v_hartree, + v_x = v_x, + v_c = v_c, + switches = self.switches, + v_x_oep = v_x_oep, + v_c_oep = v_c_oep, + de_xc_dtau = de_xc_dtau, + symmetrize = True, + exclude_boundary = True, ) - - S_inv_sqrt = self.hamiltonian_builder.ops_builder.get_S_inv_sqrt() - H_l = S_inv_sqrt[1:-1,1:-1] @ H_l[1:-1,1:-1] @ S_inv_sqrt[1:-1,1:-1] - H_l = 0.5 * (H_l + H_l.T) - # Number of states to solve for this l n_states = self.occupation_info.n_states_for_l(l) - # Solve eigenvalue problem - # eigvals, eigvecs = self.eigensolver.solve_lowest(H_l[1:-1,1:-1], n_states) - eigvals, eigvecs = self.eigensolver.solve_lowest(H_l, n_states) - - eigenvalues_all.append(eigvals) - eigenvectors_all.append(eigvecs) + if self.switches.use_oep: + full_eigvals_l, full_eigvecs_l = self.eigensolver.solve_full(H_l) + occ_eigvals_l = full_eigvals_l[:n_states] + occ_eigvecs_l = full_eigvecs_l[:, :n_states] + + # Store unoccupied eigenvalues and eigenvectors, used only for OEP and at final readout + unocc_eigenvalues_list.append(full_eigvals_l[n_states:]) + unocc_eigenvectors_list.append(full_eigvecs_l[:, n_states:]) + else: + # Solve eigenvalue problem + occ_eigvals_l, occ_eigvecs_l = self.eigensolver.solve_lowest(H_l, n_states) + + # Store occupied eigenvalues and eigenvectors + occ_eigenvalues_list.append(occ_eigvals_l) + occ_eigenvectors_list.append(occ_eigvecs_l) # Reorder eigenstates to match occupation list order - eigenvalues, eigenvectors = self._reorder_eigenstates_by_occupation( - eigenvalues_all, eigenvectors_all + occ_eigenvalues, occ_eigenvectors = self._reorder_eigenstates_by_occupation( + occ_eigenvalues_list, occ_eigenvectors_list ) # Interpolate eigenvectors to quadrature points, also symmetrize the eigenvectors - orbitals = self.hamiltonian_builder.interpolate_eigenvectors_to_quadrature( - eigenvectors = eigenvectors, + occ_orbitals = self.hamiltonian_builder.interpolate_eigenvectors_to_quadrature( + eigenvectors = occ_eigenvectors, symmetrize = True, pad_width = 1, ) # ===== Step 3: Compute new density ===== # Compute new density from orbitals - rho_new = self.density_calculator.compute_density(orbitals, normalize=True) + rho_new = self.density_calculator.compute_density(occ_orbitals, normalize=True) # ===== Step 4: Check convergence ===== converged, residual = self.inner_convergence_checker.check( @@ -885,7 +1332,7 @@ def _inner_loop( # Update density_data for next iteration using mixed density density_data = self.density_calculator.create_density_data_from_mixed( rho_mixed = rho, - orbitals = orbitals, + orbitals = occ_orbitals, compute_gradient = self.xc_requirements.needs_gradient, compute_tau = self.xc_requirements.needs_tau, rho_nlcc = self.rho_nlcc @@ -900,23 +1347,47 @@ def _inner_loop( # Create final density_data from converged orbitals, do not include NLCC - final_density_data : DensityData = self.density_calculator.create_density_data_from_orbitals( - orbitals = orbitals, - compute_gradient = self.xc_requirements.needs_gradient, - compute_tau = self.xc_requirements.needs_tau, - rho_nlcc = None - ) + final_density_data : DensityData = \ + self.density_calculator.create_density_data_from_orbitals( + orbitals = occ_orbitals, + compute_gradient = self.xc_requirements.needs_gradient, + compute_tau = self.xc_requirements.needs_tau, + rho_nlcc = None + ) + + # Construct full eigenvalues, orbitals and l terms for OEP, only needed for OEP and at final readout + if self.switches.use_oep: + full_eigen_energies, full_orbitals, full_l_terms = \ + self._compute_full_orbitals_and_eigenvalues( + rho_initial = rho_initial, + orbitals_initial = orbitals_initial, + v_x_oep = v_x_oep, + v_c_oep = v_c_oep, + switches = self.switches, + xc_requirements = self.xc_requirements, + xc_calculator = self.xc_calculator + ) + else: + full_eigen_energies = None + full_orbitals = None + full_l_terms = None - # Return final state - return SCFResult( - eigen_energies = eigenvalues, - orbitals = orbitals, - density_data = final_density_data, - converged = converged, - iterations = iteration + 1, - rho_residual = residual + + result = SCFResult( + eigen_energies = occ_eigenvalues, + orbitals = occ_orbitals, + density_data = final_density_data, + converged = converged, + iterations = iteration + 1, + rho_residual = residual, + full_eigen_energies = full_eigen_energies, + full_orbitals = full_orbitals, + full_l_terms = full_l_terms, ) + return result + + def _outer_loop( self, @@ -954,35 +1425,59 @@ def _outer_loop( SETTINGS_TYPE_ERROR_MESSAGE.format(type(settings)) if orbitals_initial is not None: assert isinstance(orbitals_initial, np.ndarray), \ - ORBITALS_INITIAL_TYPE_ERROR_MESSAGE.format(type(orbitals_initial)) + ORBITALS_INITIAL_TYPE_ERROR_MESSAGE.format(type(orbitals_initial)) + # initialize variables max_outer_iter = settings.outer_max_iter print_debug = settings.print_debug - rho = rho_initial.copy() - orbitals = orbitals_initial + rho = rho_initial.copy() + orbitals = orbitals_initial - # Compute HF exchange matrices from initial orbitals if provided - if orbitals_initial is not None: + # Compute HF exchange matrices from initial orbitals if needed + if self.switches.use_hf_exchange and (orbitals_initial is not None): H_hf_exchange_dict_by_l = self._compute_hf_exchange_matrices_dict(orbitals_initial) else: H_hf_exchange_dict_by_l = self._get_zero_hf_exchange_matrices_dict() - - + # Reset outer convergence checker self.outer_convergence_checker.reset() + + # Compute default full orbitals and eigenvalues if use_oep is True + if self.switches.use_oep: + full_eigen_energies, full_orbitals, full_l_terms = \ + self._compute_default_full_orbitals_and_eigenvalues( + rho_initial = rho_initial, + orbitals_initial = orbitals_initial, + xc_functional = 'GGA_PBE' + ) + for outer_iter in range(max_outer_iter): if print_debug: print(f"Outer iteration {outer_iter + 1}") + + # Compute OEP potentials from initial orbitals if use_oep is True + if self.switches.use_oep: + v_x_oep, v_c_oep = self.oep_calculator.compute_oep_potentials( + full_eigen_energies = full_eigen_energies, + full_orbitals = full_orbitals, + full_l_terms = full_l_terms, + enable_parallelization = self.enable_parallelization, + ) + else: + v_x_oep = None + v_c_oep = None # Run inner SCF with fixed HF exchange inner_result : SCFResult = self._inner_loop( rho_initial = rho, settings = settings, - H_hf_exchange_dict_by_l = H_hf_exchange_dict_by_l + H_hf_exchange_dict_by_l = H_hf_exchange_dict_by_l, + v_x_oep = v_x_oep, + v_c_oep = v_c_oep, ) # update rho and orbitals @@ -990,13 +1485,13 @@ def _outer_loop( orbitals = inner_result.orbitals # update HF exchange dictionary - H_hf_exchange_dict_by_l = self._compute_hf_exchange_matrices_dict(orbitals) + if self.switches.use_hf_exchange: + H_hf_exchange_dict_by_l = self._compute_hf_exchange_matrices_dict(orbitals) - # Check outer loop convergence outer_converged, outer_residual = self.outer_convergence_checker.check( rho, rho_new, outer_iter + 1, - print_status=print_debug + print_status = print_debug ) if outer_converged: @@ -1004,22 +1499,329 @@ def _outer_loop( # Update for next outer iteration rho = rho_new + if self.switches.use_oep: + # if OEP is used, also update full orbitals and eigenvalues from inner result + full_eigen_energies = inner_result.full_eigen_energies + full_orbitals = inner_result.full_orbitals + full_l_terms = inner_result.full_l_terms # Update outer loop info in result outer_iterations = outer_iter + 1 - outer_converged = (outer_iter < max_outer_iter - 1) + outer_converged = (outer_iter < max_outer_iter - 1) # Print outer loop footer if debug enabled if print_debug: self.outer_convergence_checker.print_footer(outer_converged, outer_iterations) - return SCFResult( - eigen_energies = inner_result.eigen_energies, - orbitals = inner_result.orbitals, - density_data = inner_result.density_data, - converged = outer_converged, - iterations = outer_iterations, - rho_residual = outer_residual, + outer_result = SCFResult( + eigen_energies = inner_result.eigen_energies, + orbitals = inner_result.orbitals, + density_data = inner_result.density_data, + converged = outer_converged, + iterations = outer_iterations, + rho_residual = outer_residual, + full_eigen_energies = inner_result.full_eigen_energies, + full_orbitals = inner_result.full_orbitals, + full_l_terms = inner_result.full_l_terms, ) + + return outer_result + + + + + def _compute_full_orbitals_and_eigenvalues( + self, + rho_initial : np.ndarray, + switches : SwitchesFlags, + xc_requirements : FunctionalRequirements, + xc_calculator : Optional[XCEvaluator] = None, + orbitals_initial : Optional[np.ndarray] = None, + v_x_oep : Optional[np.ndarray] = None, + v_c_oep : Optional[np.ndarray] = None, + ) -> SCFResult: + """ + Compute full orbitals and eigenvalues from initial density and orbitals + Note: No inner SCF loop needed for default full orbitals and eigenvalues, just solve eigenvalue problem for each l channel. + + Fixed: external potential, HF exchange (if any) + Iterate: rho → V_H, V_xc → H → solve → orbitals → rho' + + Parameters + ---------- + rho_initial : np.ndarray + Initial density guess + switches : SwitchesFlags + Switches flags + xc_requirements : FunctionalRequirements + Functional requirements + xc_calculator : Optional[XCEvaluator] + XC calculator, optional + orbitals_initial : np.ndarray, optional + Initial orbitals guess for debugging + Shape: (n_grid, n_orbitals) + If provided, will be used as initial orbitals instead of solving eigenvalue problem + v_x_oep : np.ndarray, optional + OEP exchange potential + v_c_oep : np.ndarray, optional + OEP correlation potential + Returns + ------- + full_eigen_energies : np.ndarray + Full eigenvalues (occupied followed by virtual) + full_orbitals : np.ndarray + Full orbitals interpolated to quadrature grid + full_l_terms : np.ndarray + Angular momentum index for each entry in full_eigen_energies + """ + # type check for required fields + assert isinstance(rho_initial, np.ndarray), \ + RHO_INITIAL_TYPE_ERROR_MESSAGE.format(type(rho_initial)) + assert isinstance(switches, SwitchesFlags), \ + SWITCHES_TYPE_ERROR_MESSAGE.format(type(switches)) + assert isinstance(xc_requirements, FunctionalRequirements), \ + XC_REQUIREMENTS_TYPE_ERROR_MESSAGE.format(type(xc_requirements)) + + # type check for optional fields + if xc_calculator is not None: + assert isinstance(xc_calculator, XCEvaluator), \ + XC_CALCULATOR_TYPE_ERROR_MESSAGE.format(type(xc_calculator)) + if orbitals_initial is not None: + assert isinstance(orbitals_initial, np.ndarray), \ + ORBITALS_INITIAL_TYPE_ERROR_MESSAGE.format(type(orbitals_initial)) + if v_x_oep is not None: + assert isinstance(v_x_oep, np.ndarray), \ + V_X_OEP_TYPE_ERROR_MESSAGE.format(type(v_x_oep)) + if v_c_oep is not None: + assert isinstance(v_c_oep, np.ndarray), \ + V_C_OEP_TYPE_ERROR_MESSAGE.format(type(v_c_oep)) + + # initialize variables + rho = self.density_calculator.normalize_density(rho_initial.copy()) + density_data = self.density_calculator.create_density_data_from_mixed( + rho_mixed = rho, + orbitals = orbitals_initial, + compute_gradient = xc_requirements.needs_gradient, + compute_tau = xc_requirements.needs_tau and (orbitals_initial is not None), + rho_nlcc = self.rho_nlcc + ) + + # ===== Step 1: Compute potentials ===== + # Hartree potential + v_hartree = self.poisson_solver.solve_hartree(rho) + + # XC potential + # For pseudopotentials: use rho_total = rho_valence + rho_nlcc + # For all-electron: rho_nlcc is zero, so rho_total = rho_valence + v_x = np.zeros_like(rho) # Default: no exchange potential + v_c = np.zeros_like(rho) # Default: no correlation potential + de_xc_dtau = None # Default: no meta-GGA derivative term + + if xc_calculator is not None: + # Compute XC using new interface: DensityData → XCPotentialData + xc_potential_data = xc_calculator.compute_xc(density_data) + v_x = xc_potential_data.v_x + v_c = xc_potential_data.v_c + de_xc_dtau = xc_potential_data.de_xc_dtau + + # ===== Step 2: Build and solve for each l channel ===== + occ_eigenvalues_list : List[np.ndarray] = [] + occ_eigenvectors_list : List[np.ndarray] = [] + unocc_eigenvalues_list : List[np.ndarray] = [] # needed only for OEP + unocc_eigenvectors_list : List[np.ndarray] = [] # needed only for OEP + + + if self.angular_momentum_cutoff is not None: + unique_l_values = list(range(self.angular_momentum_cutoff + 1)) + else: + unique_l_values = self.occupation_info.unique_l_values + + + for l in unique_l_values: + # Build Hamiltonian for this l + H_l = self.hamiltonian_builder.build_for_l_channel( + l = l, + v_hartree = v_hartree, + v_x = v_x, + v_c = v_c, + switches = switches, + v_x_oep = v_x_oep, + v_c_oep = v_c_oep, + de_xc_dtau = de_xc_dtau, + symmetrize = True, + exclude_boundary = True, + ) + + + # Number of states to solve for this l + n_occ_states = self.occupation_info.n_states_for_l(l) + + # Solve eigenvalue problem + full_eigenvalues_l, full_eigenvectors_l = self.eigensolver.solve_full(H_l) + + # Append eigenvalues and eigenvectors to lists + occ_eigenvalues_list.append(full_eigenvalues_l[:n_occ_states]) + occ_eigenvectors_list.append(full_eigenvectors_l[:, :n_occ_states]) + unocc_eigenvalues_list.append(full_eigenvalues_l[n_occ_states:]) + unocc_eigenvectors_list.append(full_eigenvectors_l[:, n_occ_states:]) + + + return self._construct_full_eigenvalues_and_orbitals_and_l_terms( + occ_eigenvalues_list = occ_eigenvalues_list, + occ_eigenvectors_list = occ_eigenvectors_list, + unocc_eigenvalues_list = unocc_eigenvalues_list, + unocc_eigenvectors_list = unocc_eigenvectors_list, + angular_momentum_cutoff = self.angular_momentum_cutoff + ) + + + def _compute_default_full_orbitals_and_eigenvalues( + self, + rho_initial : np.ndarray, + orbitals_initial : np.ndarray, + xc_functional : str = 'GGA_PBE' + ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Compute full orbitals and eigenvalues from initial density and orbitals + Note: No inner SCF loop needed for default full orbitals and eigenvalues, just solve eigenvalue problem for each l channel. + + Parameters + ---------- + rho_initial : np.ndarray + Initial density + orbitals_initial : np.ndarray + Initial orbitals + xc_functional : str + XC functional, default: 'GGA_PBE' + + Returns + ------- + full_eigen_energies : np.ndarray + Full eigenvalues (occupied followed by virtual) + full_orbitals : np.ndarray + Full orbitals interpolated to quadrature grid + full_l_terms : np.ndarray + Angular momentum index for each entry in full_eigen_energies + """ + # type check for required fields + assert isinstance(rho_initial, np.ndarray), \ + RHO_INITIAL_TYPE_ERROR_MESSAGE.format(type(rho_initial)) + assert isinstance(orbitals_initial, np.ndarray), \ + ORBITALS_INITIAL_TYPE_ERROR_MESSAGE.format(type(orbitals_initial)) + assert isinstance(xc_functional, str), \ + XC_FUNCTIONAL_TYPE_ERROR_MESSAGE.format(type(xc_functional)) + + _switches : SwitchesFlags = SwitchesFlags(xc_functional) + _xc_requirements : FunctionalRequirements = get_functional_requirements(xc_functional) + _xc_calculator : XCEvaluator = \ + create_xc_evaluator( + functional_name = xc_functional, + derivative_matrix = self.density_calculator.derivative_matrix, + r_quad = self.density_calculator.quadrature_nodes + ) + + return self._compute_full_orbitals_and_eigenvalues( + rho_initial = rho_initial, + orbitals_initial = orbitals_initial, + switches = _switches, + xc_requirements = _xc_requirements, + xc_calculator = _xc_calculator + ) + + + def _construct_full_eigenvalues_and_orbitals_and_l_terms( + self, + occ_eigenvalues_list : List[np.ndarray], + occ_eigenvectors_list : List[np.ndarray], + unocc_eigenvalues_list : List[np.ndarray], + unocc_eigenvectors_list : List[np.ndarray], + angular_momentum_cutoff : Optional[int] = None, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Construct full eigenvalues, orbitals and l terms for OEP, only needed for OEP and at final readout + """ + # type check for required fields + assert isinstance(occ_eigenvalues_list, list), \ + OCC_EIGENVALUES_LIST_NOT_LIST_ERROR_MESSAGE.format(type(occ_eigenvalues_list)) + assert isinstance(occ_eigenvectors_list, list), \ + OCC_EIGENVECTORS_LIST_NOT_LIST_ERROR_MESSAGE.format(type(occ_eigenvectors_list)) + assert isinstance(unocc_eigenvalues_list, list), \ + UNOCC_EIGENVALUES_LIST_NOT_LIST_ERROR_MESSAGE.format(type(unocc_eigenvalues_list)) + assert isinstance(unocc_eigenvectors_list, list), \ + UNOCC_EIGENVECTORS_LIST_NOT_LIST_ERROR_MESSAGE.format(type(unocc_eigenvectors_list)) + if angular_momentum_cutoff is not None: + assert isinstance(angular_momentum_cutoff, int), \ + ANGULAR_MOMENTUM_CUTOFF_TYPE_ERROR_MESSAGE.format(type(angular_momentum_cutoff)) + + # length checks + unique_l_values_number = len(self.occupation_info.unique_l_values) if angular_momentum_cutoff is None else angular_momentum_cutoff + 1 + assert len(occ_eigenvalues_list) == unique_l_values_number, \ + OCC_EIGENVALUES_LIST_LENGTH_ERROR_MESSAGE.format(len(occ_eigenvalues_list)) + assert len(occ_eigenvectors_list) == unique_l_values_number, \ + OCC_EIGENVECTORS_LIST_LENGTH_ERROR_MESSAGE.format(len(occ_eigenvectors_list)) + assert len(unocc_eigenvalues_list) == unique_l_values_number, \ + UNOCC_EIGENVALUES_LIST_LENGTH_ERROR_MESSAGE.format(len(unocc_eigenvalues_list)) + assert len(unocc_eigenvectors_list) == unique_l_values_number, \ + UNOCC_EIGENVECTORS_LIST_LENGTH_ERROR_MESSAGE.format(len(unocc_eigenvectors_list)) + + # Dimention and shape checks + n_physical_nodes = self.hamiltonian_builder.ops_builder.physical_nodes.shape[0] + n_interior_nodes = n_physical_nodes - 2 + + for i in range(unique_l_values_number): + occ_eigenvalues = occ_eigenvalues_list[i] + occ_eigenvectors = occ_eigenvectors_list[i] + unocc_eigenvalues = unocc_eigenvalues_list[i] + unocc_eigenvectors = unocc_eigenvectors_list[i] + + # check dimention + assert occ_eigenvalues.ndim == 1, \ + OCC_EIGENVALUES_LIST_NDIM_ERROR_MESSAGE.format(occ_eigenvalues.ndim) + assert occ_eigenvectors.ndim == 2, \ + OCC_EIGENVECTORS_LIST_NDIM_ERROR_MESSAGE.format(occ_eigenvectors.ndim) + assert unocc_eigenvalues.ndim == 1, \ + UNOCC_EIGENVALUES_LIST_NDIM_ERROR_MESSAGE.format(unocc_eigenvalues.ndim) + assert unocc_eigenvectors.ndim == 2, \ + UNOCC_EIGENVECTORS_LIST_NDIM_ERROR_MESSAGE.format(unocc_eigenvectors.ndim) + + # occupied shape check + assert occ_eigenvalues.shape[0] == occ_eigenvectors.shape[1], \ + OCC_EIGENVALUES_AND_EIGENVECTORS_LIST_SHAPE_MISMATCH_ERROR_MESSAGE.format(occ_eigenvalues.shape[0], occ_eigenvectors.shape[1]) + assert occ_eigenvectors.shape[0] == n_interior_nodes, \ + OCC_EIGENVECTORS_LIST_SHAPE_ERROR_MESSAGE.format(occ_eigenvectors.shape[0], n_interior_nodes) + + # unoccupied shape check + assert unocc_eigenvalues.shape[0] == unocc_eigenvectors.shape[1], \ + UNOCC_EIGENVALUES_AND_EIGENVECTORS_LIST_SHAPE_MISMATCH_ERROR_MESSAGE.format(unocc_eigenvalues.shape[0], unocc_eigenvectors.shape[1]) + assert unocc_eigenvectors.shape[0] == n_interior_nodes, \ + UNOCC_EIGENVECTORS_LIST_SHAPE_ERROR_MESSAGE.format(unocc_eigenvectors.shape[0], n_interior_nodes) + + + # Reorder eigenstates to match occupation list order + occ_eigenvalues, occ_eigenvectors = self._reorder_eigenstates_by_occupation( + occ_eigenvalues_list, occ_eigenvectors_list + ) + + full_eigenvalues = np.concatenate([occ_eigenvalues, *unocc_eigenvalues_list], axis = 0) + full_eigenvectors = np.concatenate([occ_eigenvectors, *unocc_eigenvectors_list], axis = 1) + + # Interpolate eigenvectors to quadrature points, also symmetrize the eigenvectors + full_orbitals = self.hamiltonian_builder.interpolate_eigenvectors_to_quadrature( + eigenvectors = full_eigenvectors, + symmetrize = True, + pad_width = 1, + ) + + # Compute angular momentum index for each entry in full_eigen_energies + unocc_l_terms = np.concatenate([ + np.full(vals.shape[0], l, dtype=int) + for l, vals in zip(range(unique_l_values_number), unocc_eigenvalues_list) + if vals.size > 0 + ]) if len(unocc_eigenvalues_list) > 0 else np.empty(0, dtype=int) + + full_l_terms = np.concatenate([self.occupation_info.occ_l, unocc_l_terms]) + + return full_eigenvalues, full_orbitals, full_l_terms + diff --git a/utils/atom/scf/eigensolver.py b/utils/atom/scf/eigensolver.py index daed5991..b2681a35 100644 --- a/utils/atom/scf/eigensolver.py +++ b/utils/atom/scf/eigensolver.py @@ -24,6 +24,15 @@ def __init__(self, xc_functional: Optional[str] = None): """ self.xc_functional = xc_functional + + @staticmethod + def solve_full(H: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """ + Solve for all eigenvalues and eigenvectors. + """ + return np.linalg.eigh(H, UPLO='L') + + def solve_lowest(self, H: np.ndarray, k: int) -> tuple[np.ndarray, np.ndarray]: """ Solve for k lowest eigenvalues and eigenvectors. diff --git a/utils/atom/scf/energy.py b/utils/atom/scf/energy.py index 8667592b..61a546d1 100644 --- a/utils/atom/scf/energy.py +++ b/utils/atom/scf/energy.py @@ -5,14 +5,16 @@ from __future__ import annotations import numpy as np -from typing import Optional, TYPE_CHECKING +from typing import Optional, Tuple, TYPE_CHECKING from dataclasses import dataclass +from ..scf.driver import SwitchesFlags from ..mesh.operators import GridData from ..utils.occupation_states import OccupationInfo from ..scf.poisson import PoissonSolver from ..xc.evaluator import XCEvaluator, XCPotentialData from ..xc.hybrid import HartreeFockExchange +from ..xc.oep import OEPCalculator from ..mesh.operators import RadialOperatorsBuilder from ..pseudo.local import LocalPseudopotential from ..pseudo.non_local import NonLocalPseudopotential @@ -38,6 +40,36 @@ "parameter mixing_parameter must be in [0, 1], got {} instead" +FULL_EIGEN_ENERGIES_NOT_NONE_ERROR_MESSAGE = \ + "parameter `full_eigen_energies` must be not None, get None instead" +FULL_ORBITALS_NOT_NONE_ERROR_MESSAGE = \ + "parameter `full_orbitals` must be not None, get None instead" +FULL_L_TERMS_NOT_NONE_ERROR_MESSAGE = \ + "parameter `full_l_terms` must be not None, get None instead" + + +SWITCHES_NOT_A_SWITCHESFLAGS_ERROR_MESSAGE = \ + "parameter `switches` must be a SwitchesFlags instance, get {} instead" +GRID_DATA_NOT_A_GRIDDATA_ERROR_MESSAGE = \ + "parameter `grid_data` must be a GridData instance, get {} instead" +OCCUPATION_INFO_NOT_A_OCCUPATIONINFO_ERROR_MESSAGE = \ + "parameter `occupation_info` must be a OccupationInfo instance, get {} instead" +OPS_BUILDER_NOT_A_RADIALOPERATORSBUILDER_ERROR_MESSAGE = \ + "parameter `ops_builder` must be a RadialOperatorsBuilder instance, get {} instead" +POISSON_SOLVER_NOT_A_POISSONSOLVER_ERROR_MESSAGE = \ + "parameter `poisson_solver` must be a PoissonSolver instance, get {} instead" +PSEUDO_NOT_A_LOCALPSEUDOPOTENTIAL_ERROR_MESSAGE = \ + "parameter `pseudo` must be a LocalPseudopotential instance, get {} instead" +XC_CALCULATOR_NOT_A_XCEVALUATOR_ERROR_MESSAGE = \ + "parameter `xc_calculator` must be a XCEvaluator instance, get {} instead" +HF_CALCULATOR_NOT_A_HARTREEFOCKEXCHANGE_ERROR_MESSAGE = \ + "parameter `hf_calculator` must be a HartreeFockExchange instance, get {} instead" +OEP_CALCULATOR_NOT_A_OEPCALCULATOR_ERROR_MESSAGE = \ + "parameter `oep_calculator` must be a OEPCalculator instance, get {} instead" +DERIVATIVE_MATRIX_NOT_A_NDARRAY_ERROR_MESSAGE = \ + "parameter `derivative_matrix` must be a np.ndarray instance, get {} instead" + + @dataclass class EnergyComponents: """Container for all energy components""" @@ -54,6 +86,7 @@ class EnergyComponents: # Optional: advanced functionals and corrections nonlocal_psp : float = 0.0 # E_nl (non-local pseudopotential) hf_exchange : float = 0.0 # E_HF (Hartree-Fock exchange) + oep_exchange : float = 0.0 # E_OEP (OEP exchange) rpa_correlation : float = 0.0 # E_RPA (RPA correlation) @property @@ -66,13 +99,14 @@ def total_potential(self) -> float: """Total potential energy""" return (self.external + self.nonlocal_psp + self.hartree + self.exchange + self.correlation + - self.hf_exchange + self.rpa_correlation) + self.hf_exchange + self.oep_exchange + self.rpa_correlation) @property def total(self) -> float: """Total energy""" return self.total_kinetic + self.total_potential + def print_info(self, title: str = "Energy Components"): """Print formatted energy information""" # Title @@ -96,6 +130,8 @@ def print_info(self, title: str = "Energy Components"): print(f"\t Nonlocal PSP : {self.nonlocal_psp:16.8f} Ha") if abs(self.hf_exchange) > 1e-12: print(f"\t HF Exchange : {self.hf_exchange:16.8f} Ha") + if abs(self.oep_exchange) > 1e-12: + print(f"\t OEP Exchange : {self.oep_exchange:16.8f} Ha") if abs(self.rpa_correlation) > 1e-12: print(f"\t RPA Correlation : {self.rpa_correlation:16.8f} Ha") print(f"\t Total Potential : {self.total_potential:16.8f} Ha") @@ -111,7 +147,7 @@ class EnergyCalculator: Computes total energy from converged Kohn-Sham solution Total energy: - E = T_s + E_ext + E_H + E_xc + E_HF + E_RPA + E = T_s + E_ext + E_H + E_xc + E_HF + E_OEP + E_RPA where: T_s = kinetic energy of non-interacting electrons @@ -119,11 +155,13 @@ class EnergyCalculator: E_H = Hartree energy (classical electron-electron repulsion) E_xc = exchange-correlation energy E_HF = Hartree-Fock exchange (optional) + E_OEP = OEP exchange (optional) E_RPA = RPA correlation (optional) """ def __init__( self, + switches : SwitchesFlags, grid_data : GridData, occupation_info : OccupationInfo, ops_builder : RadialOperatorsBuilder, @@ -131,11 +169,14 @@ def __init__( pseudo : LocalPseudopotential, xc_calculator : Optional[XCEvaluator] = None, hf_calculator : Optional[HartreeFockExchange] = None, + oep_calculator : Optional[OEPCalculator] = None, derivative_matrix : Optional[np.ndarray] = None, - ): + ): """ Parameters ---------- + switches : SwitchesFlags + Switches flags for the energy calculation grid_data : GridData Grid information (quadrature points, weights) from standard grid occupation_info : OccupationInfo @@ -152,11 +193,18 @@ def __init__( hf_calculator : HartreeFockExchange, optional HF exchange calculator for hybrid functionals Should be the same instance as used in SCFDriver + oep_calculator : OEPCalculator, optional + OEP exchange calculator for OEP exchange functional + Should be the same instance as used in SCFDriver derivative_matrix : np.ndarray, optional Derivative matrix for kinetic energy computation Note: If None, uses ops_builder.derivative_matrix. If provided, should be from dense grid for more accurate results. """ + assert isinstance(grid_data, GridData), \ + GRID_DATA_NOT_A_GRIDDATA_ERROR_MESSAGE.format(type(grid_data)) + + self.switches = switches self.quadrature_nodes = grid_data.quadrature_nodes self.quadrature_weights = grid_data.quadrature_weights self.occupation_info = occupation_info @@ -165,7 +213,10 @@ def __init__( self.pseudo = pseudo self.xc_calculator = xc_calculator self.hf_calculator = hf_calculator - + self.oep_calculator = oep_calculator + self._check_initialization() + + # Use provided derivative_matrix or fall back to ops_builder's matrix if derivative_matrix is not None: self.derivative_matrix = derivative_matrix @@ -190,18 +241,49 @@ def __init__( # Non-local pseudopotential calculator if not self.all_electron: self.nonlocal_calculator = NonLocalPseudopotential( - pseudo=self.pseudo, - ops_builder=self.ops_builder + pseudo = self.pseudo, + ops_builder = self.ops_builder ) else: self.nonlocal_calculator = None - + + def _check_initialization(self): + """ + Check if the initialization is correct + """ + # type check for required fields + assert isinstance(self.switches, SwitchesFlags), \ + SWITCHES_NOT_A_SWITCHESFLAGS_ERROR_MESSAGE.format(type(self.switches)) + assert isinstance(self.occupation_info, OccupationInfo), \ + OCCUPATION_INFO_NOT_A_OCCUPATIONINFO_ERROR_MESSAGE.format(type(self.occupation_info)) + assert isinstance(self.ops_builder, RadialOperatorsBuilder), \ + OPS_BUILDER_NOT_A_RADIALOPERATORSBUILDER_ERROR_MESSAGE.format(type(self.ops_builder)) + assert isinstance(self.poisson_solver, PoissonSolver), \ + POISSON_SOLVER_NOT_A_POISSONSOLVER_ERROR_MESSAGE.format(type(self.poisson_solver)) + assert isinstance(self.pseudo, LocalPseudopotential), \ + PSEUDO_NOT_A_LOCALPSEUDOPOTENTIAL_ERROR_MESSAGE.format(type(self.pseudo)) + + # type check for optional fields + if self.xc_calculator is not None: + assert isinstance(self.xc_calculator, XCEvaluator), \ + XC_CALCULATOR_NOT_A_XCEVALUATOR_ERROR_MESSAGE.format(type(self.xc_calculator)) + if self.hf_calculator is not None: + assert isinstance(self.hf_calculator, HartreeFockExchange), \ + HF_CALCULATOR_NOT_A_HARTREEFOCKEXCHANGE_ERROR_MESSAGE.format(type(self.hf_calculator)) + if self.oep_calculator is not None: + assert isinstance(self.oep_calculator, OEPCalculator), \ + OEP_CALCULATOR_NOT_A_OEPCALCULATOR_ERROR_MESSAGE.format(type(self.oep_calculator)) + def compute_energy( self, - orbitals: np.ndarray, - density_data: 'DensityData', - mixing_parameter: Optional[float] = None, + orbitals : np.ndarray, + density_data : 'DensityData', + mixing_parameter : Optional[float] = None, + full_eigen_energies : Optional[np.ndarray] = None, + full_orbitals : Optional[np.ndarray] = None, + full_l_terms : Optional[np.ndarray] = None, + enable_parallelization : Optional[bool] = None, ) -> EnergyComponents: """ Compute total energy from converged SCF solution @@ -214,7 +296,17 @@ def compute_energy( density_data : DensityData Electron density and related quantities (rho, grad_rho, tau) This density data should not include NLCC - + mixing_parameter : Optional[float] + Mixing parameter for the exchange-correlation energy + full_eigen_energies : Optional[np.ndarray] + Full eigenvalues of the Kohn-Sham orbitals + full_orbitals : Optional[np.ndarray] + Full orbitals of the Kohn-Sham orbitals + full_l_terms : Optional[np.ndarray] + Full l terms of the Kohn-Sham orbitals + enable_parallelization: Optional[bool] + Flag for parallelization of RPA calculations + Returns ------- energy : EnergyComponents @@ -250,24 +342,34 @@ def compute_energy( # 6. Hartree-Fock exchange energy (if HF calculator is available) E_hf_exchange = self._compute_hf_exchange_energy(orbitals) - - # 7. Advanced energy terms (placeholders for future implementation) - E_rpa_correlation = 0.0 + + # 7. OEP exchange energy (if OEP calculator is available) + E_oep_exchange = self._compute_oep_exchange_energy(orbitals) + + # 8. Advanced energy terms (placeholders for future implementation) + E_rpa_correlation = self._compute_rpa_correlation_energy( + full_eigen_energies = full_eigen_energies, + full_orbitals = full_orbitals, + full_l_terms = full_l_terms, + enable_parallelization = enable_parallelization, + ) if mixing_parameter is not None: E_x *= (1.0 - mixing_parameter) E_hf_exchange *= mixing_parameter - + E_oep_exchange *= mixing_parameter + return EnergyComponents( - kinetic_radial=T_radial, - kinetic_angular=T_angular, - external=E_ext, - hartree=E_hartree, - exchange=E_x, - correlation=E_c, - nonlocal_psp=E_nonlocal_psp, - hf_exchange=E_hf_exchange, - rpa_correlation=E_rpa_correlation + kinetic_radial = T_radial, + kinetic_angular = T_angular, + external = E_ext, + hartree = E_hartree, + exchange = E_x, + correlation = E_c, + nonlocal_psp = E_nonlocal_psp, + hf_exchange = E_hf_exchange, + oep_exchange = E_oep_exchange, + rpa_correlation = E_rpa_correlation ) @@ -275,7 +377,7 @@ def _compute_kinetic_energy( self, orbitals: np.ndarray, derivative_matrix: np.ndarray - ) -> tuple[float, float]: + ) -> tuple[float, float]: """ Compute kinetic energy: T = T_radial + T_angular @@ -352,9 +454,13 @@ def _compute_xc_energy(self, density_data: 'DensityData') -> tuple[float, float] E_x, E_c : float, float Exchange and correlation energy components """ - if self.xc_calculator is None: + if not self.switches.use_xc_functional: return 0.0, 0.0 + assert isinstance(self.xc_calculator, XCEvaluator), \ + XC_CALCULATOR_NOT_A_XCEVALUATOR_ERROR_MESSAGE.format(type(self.xc_calculator)) + + # Get total density including NLCC for pseudopotential calculations density_data_total = self.get_total_density_data_for_xc(density_data) @@ -403,50 +509,84 @@ def get_total_density_data_for_xc(self, density_data: 'DensityData') -> 'Density # Use static method to add NLCC correction # This properly handles recomputing grad_rho if needed density_data_total = DensityCalculator.add_nlcc_to_density_data( - density_data_valence=density_data, - rho_nlcc=self.rho_nlcc if not self.all_electron else None, - quadrature_nodes=self.quadrature_nodes, - derivative_matrix=self.derivative_matrix + density_data_valence = density_data, + rho_nlcc = self.rho_nlcc if not self.all_electron else None, + quadrature_nodes = self.quadrature_nodes, + derivative_matrix = self.derivative_matrix ) return density_data_total - def compute_xc_potential(self, density_data: 'DensityData'): + def compute_local_xc_potential( + self, + density_data : 'DensityData', + full_eigen_energies : Optional[np.ndarray] = None, + full_orbitals : Optional[np.ndarray] = None, + full_l_terms : Optional[np.ndarray] = None, + enable_parallelization : Optional[bool] = None, + ) -> Tuple[np.ndarray, np.ndarray]: """ Compute exchange-correlation potential for given density data. - - This method handles NLCC correction for pseudopotential calculations. - The input density_data contains valence density only, and NLCC is added - for XC calculations. - Parameters ---------- + density_data : DensityData - Valence electron density data (does not include NLCC) - + Electron density and related quantities (rho, grad_rho, tau) + This density data should not include NLCC + full_eigen_energies : Optional[np.ndarray] + Full eigenvalues of the Kohn-Sham orbitals + full_orbitals : Optional[np.ndarray] + Full orbitals of the Kohn-Sham orbitals + full_l_terms : Optional[np.ndarray] + Full l terms of the Kohn-Sham orbitals + Returns ------- - xc_potential_data : XCPotentialData - Exchange-correlation potential and energy density data + v_x_local, v_c_local : np.ndarray, np.ndarray + Exchange and correlation potentials """ - if self.xc_calculator is None: - # Return zero XC potential data when no XC calculator is available - n_grid = len(self.quadrature_nodes) - return XCPotentialData( - v_x=np.zeros(n_grid), - v_c=np.zeros(n_grid), - e_x=np.zeros(n_grid), - e_c=np.zeros(n_grid) - ) - - # Get total density data (including NLCC for pseudopotential calculations) - density_data_total = self.get_total_density_data_for_xc(density_data) + + n_grid = len(self.quadrature_nodes) + + v_x = np.zeros(n_grid) + v_c = np.zeros(n_grid) + v_x_oep = np.zeros(n_grid) + v_c_oep = np.zeros(n_grid) + + # Compute XC potential using XC functional + if self.switches.use_xc_functional: + + assert isinstance(self.xc_calculator, XCEvaluator), \ + XC_CALCULATOR_NOT_A_XCEVALUATOR_ERROR_MESSAGE.format(type(self.xc_calculator)) + + density_data_total = self.get_total_density_data_for_xc(density_data) + + # Compute XC potential using total density + xc_potential_data = self.xc_calculator.compute_xc(density_data_total) + + v_x = xc_potential_data.v_x + v_c = xc_potential_data.v_c - # Compute XC potential using total density - xc_potential_data = self.xc_calculator.compute_xc(density_data_total) + + if self.switches.use_oep: + v_x_oep, v_c_oep = self.oep_calculator.compute_oep_potentials( + full_eigen_energies = full_eigen_energies, + full_orbitals = full_orbitals, + full_l_terms = full_l_terms, + enable_parallelization = enable_parallelization, + ) - return xc_potential_data + # Mix the XC potential using the hybrid mixing parameter + if self.switches.hybrid_mixing_parameter is not None: + v_x_local = v_x * (1.0 - self.switches.hybrid_mixing_parameter) + v_x_oep * self.switches.hybrid_mixing_parameter + v_c_local = v_c * (1.0 - self.switches.hybrid_mixing_parameter) + v_c_oep * self.switches.hybrid_mixing_parameter + else: + v_x_local = v_x + v_x_oep + v_c_local = v_c + v_c_oep + + return v_x_local, v_c_local + def _compute_nonlocal_psp_energy(self, orbitals: np.ndarray) -> float: @@ -498,10 +638,57 @@ def _compute_hf_exchange_energy( float HF exchange energy (0.0 if no HF calculator available) """ - if self.hf_calculator is None: + if not self.switches.use_hf_exchange: return 0.0 + + assert isinstance(self.hf_calculator, HartreeFockExchange), \ + HF_CALCULATOR_NOT_A_HARTREEFOCKEXCHANGE_ERROR_MESSAGE.format(type(self.hf_calculator)) # Delegate computation to HF calculator return self.hf_calculator.compute_exchange_energy(orbitals) + + def _compute_oep_exchange_energy(self, orbitals: np.ndarray) -> float: + """ + Compute OEP exchange energy. + """ + if not self.switches.use_oep_exchange: + return 0.0 + + assert isinstance(self.oep_calculator, OEPCalculator), \ + OEP_CALCULATOR_NOT_A_OEPCALCULATOR_ERROR_MESSAGE.format(type(self.oep_calculator)) + + return self.oep_calculator.compute_exchange_energy(orbitals) + + + + def _compute_rpa_correlation_energy( + self, + full_eigen_energies : np.ndarray, + full_orbitals : np.ndarray, + full_l_terms : np.ndarray, + enable_parallelization: Optional[bool] = False, + ) -> float: + """ + Compute RPA correlation energy. + """ + if not self.switches.use_oep_correlation: + return 0.0 + + assert isinstance(self.oep_calculator, OEPCalculator), \ + OEP_CALCULATOR_NOT_A_OEPCALCULATOR_ERROR_MESSAGE.format(type(self.oep_calculator)) + + assert full_eigen_energies is not None, \ + FULL_EIGEN_ENERGIES_NOT_NONE_ERROR_MESSAGE + assert full_orbitals is not None, \ + FULL_ORBITALS_NOT_NONE_ERROR_MESSAGE + assert full_l_terms is not None, \ + FULL_L_TERMS_NOT_NONE_ERROR_MESSAGE + + return self.oep_calculator.compute_correlation_energy( + full_eigen_energies = full_eigen_energies, + full_orbitals = full_orbitals, + full_l_terms = full_l_terms, + enable_parallelization = enable_parallelization, + ) \ No newline at end of file diff --git a/utils/atom/scf/hamiltonian.py b/utils/atom/scf/hamiltonian.py index 57700f32..b85f2e12 100644 --- a/utils/atom/scf/hamiltonian.py +++ b/utils/atom/scf/hamiltonian.py @@ -5,15 +5,17 @@ from __future__ import annotations import numpy as np -from typing import Dict, Optional, Any -from dataclasses import dataclass +from typing import Dict, Optional, TYPE_CHECKING from ..mesh.operators import RadialOperatorsBuilder from ..pseudo.local import LocalPseudopotential from ..utils.occupation_states import OccupationInfo +# Type checking for SwitchesFlags, avoid circular import +if TYPE_CHECKING: + from .driver import SwitchesFlags -# Error messages +# Hamiltonian Builder Error messages OPS_BUILDER_MUST_BE_A_RADIAL_OPERATORS_BUILDER_ERROR = \ "parameter ops_builder must be a RadialOperatorsBuilder, get type {} instead" LOCAL_PSEUDOPOTENTIAL_MUST_BE_A_LOCAL_PSEUDOPOTENTIAL_ERROR = \ @@ -22,141 +24,27 @@ "parameter occupation_info must be an OccupationInfo, get type {} instead" ALL_ELECTRON_MUST_BE_A_BOOLEAN_ERROR = \ "parameter all_electron must be a boolean, get type {} instead" - EIGENVECTORS_MUST_BE_A_NUMPY_ARRAY_ERROR = \ "parameter eigenvectors must be a numpy array, get type {} instead" EIGENVECTORS_MUST_BE_A_2D_ARRAY_ERROR = \ "parameter eigenvectors must be a 2D array, get dimension {} instead" HARTREE_FOCK_EXCHANGE_MATRIX_FOR_L_CHANNEL_NOT_AVAILABLE_ERROR = \ "Hartree-Fock exchange matrix for l={l} is not available, please set the HF exchange matrices first" - MIXING_PARAMETER_NOT_A_FLOAT_ERROR = \ "parameter mixing_parameter must be a float, get type {} instead" MIXING_PARAMETER_NOT_IN_ZERO_ONE_ERROR = \ "parameter mixing_parameter must be in [0, 1], got {} instead" - DE_XC_DTAU_NOT_AVAILABLE_ERROR = \ "parameter de_xc_dtau is not available for l={l}, please set the de_xc_dtau first" -# Warning messages +# Hamiltonian Builder Warning messages HARTREE_FOCK_EXCHANGE_MATRIX_FOR_L_CHANNEL_NOT_AVAILABLE_WARNING = \ "WARNING: Hartree-Fock exchange matrix for l={l} is not available, please set the HF exchange matrices first" HF_EXCHANGE_MATRIX_NOT_AVAILABLE_WARNING = \ "WARNING: Hartree-Fock exchange matrix is not available, please set the HF exchange matrices first" - - -class SwitchesFlags: - """ - Internal helper class for HamiltonianBuilder to manage functional switches. - Determines which Hamiltonian components to include based on the XC functional. - - Attributes: - use_hf_exchange (bool): Whether to use Hartree-Fock exchange - use_oep_exchange (bool): Whether to use Optimized Effective Potential (OEP) exchange - use_oep_correlation (bool): Whether to use OEP correlation - use_rpa_correlation (bool): Whether to use Random Phase Approximation (RPA) correlation - use_metagga (bool): Whether to use meta-GGA functional (requires de_xc_dtau in Hamiltonian) - OEP (bool): Whether to use OEP methods - inside_OEP (bool): Whether currently inside OEP iteration - uscrpa (bool): Whether to use RPA (Random Phase Approximation) - "use RPA" flag - """ - def __init__(self, xc_functional: str, hybrid_mixing_parameter: Optional[float] = None): - """ - Initialize switches based on XC functional and hybrid mixing parameter. - - Parameters - ---------- - xc_functional : str - Name of XC functional - hybrid_mixing_parameter : float, optional - Mixing parameter for hybrid functionals (0-1). Required only for hybrid functionals. - - For PBE0: should be 0.25 - - For HF: should be 1.0 - """ - # Type checking - assert isinstance(xc_functional, str), \ - "parameter xc_functional must be a string, get type {} instead".format(type(xc_functional)) - if hybrid_mixing_parameter is not None: - assert isinstance(hybrid_mixing_parameter, (float, int)), \ - "parameter hybrid_mixing_parameter must be a float, get type {} instead".format(type(hybrid_mixing_parameter)) - hybrid_mixing_parameter = float(hybrid_mixing_parameter) - - self.xc_functional = xc_functional - self.hybrid_mixing_parameter = hybrid_mixing_parameter - - # Initialize all flags to False - self.OEP = False - self.inside_OEP = False - self.uscrpa = False - self.use_hf_exchange = False - self.use_oep_exchange = False - self.use_oep_correlation = False - self.use_rpa_correlation = False - self.use_metagga = False - - # Set flags based on functional - if xc_functional == 'RPA': - self.OEP = True - self.uscrpa = True - self.use_oep_exchange = True - self.use_oep_correlation = True - self.use_rpa_correlation = True - elif xc_functional == 'OEPx': - self.OEP = True - self.uscrpa = True - self.use_oep_exchange = True - self.use_oep_correlation = True - elif xc_functional == 'HF': - self.use_hf_exchange = True - # Check if hybrid_mixing_parameter is provided for hybrid functionals - if hybrid_mixing_parameter is None: - print("WARNING: hybrid_mixing_parameter not provided for HF functional, using default value 1.0") - self.hybrid_mixing_parameter = 1.0 - elif not np.isclose(hybrid_mixing_parameter, 1.0): - print("WARNING: hybrid_mixing_parameter for HF should be 1.0, got {}".format(hybrid_mixing_parameter)) - elif xc_functional == 'PBE0': - self.use_hf_exchange = True - # Check if hybrid_mixing_parameter is provided for hybrid functionals - if hybrid_mixing_parameter is None: - print("WARNING: hybrid_mixing_parameter not provided for PBE0 functional, using default value 0.25") - self.hybrid_mixing_parameter = 0.25 - elif not np.isclose(hybrid_mixing_parameter, 0.25): - print("WARNING: hybrid_mixing_parameter for PBE0 should be 0.25, got {}".format(hybrid_mixing_parameter)) - - # Set meta-GGA flag for functionals that require de_xc_dtau - if xc_functional in ['SCAN', 'RSCAN', 'R2SCAN']: - self.use_metagga = True - - # LDA/GGA functionals: no special flags (default False) - - def print_info(self): - """ - Print functional switch information summary. - """ - print("=" * 60) - print("\t\t FUNCTIONAL SWITCHES") - print("=" * 60) - print(f"\t XC Functional : {self.xc_functional}") - if self.hybrid_mixing_parameter is not None: - print(f"\t Hybrid Mixing Parameter : {self.hybrid_mixing_parameter}") - else: - print(f"\t Hybrid Mixing Parameter : None (not applicable)") - print() - print("\t FUNCTIONAL FLAGS:") - print(f"\t use_hf_exchange : {self.use_hf_exchange}") - print(f"\t use_oep_exchange : {self.use_oep_exchange}") - print(f"\t use_oep_correlation : {self.use_oep_correlation}") - print(f"\t use_rpa_correlation : {self.use_rpa_correlation}") - print(f"\t use_metagga : {self.use_metagga}") - print(f"\t OEP : {self.OEP}") - print(f"\t inside_OEP : {self.inside_OEP}") - print(f"\t use_RPA (uscrpa) : {self.uscrpa}") - print() - - class HamiltonianBuilder: """ Constructs Kohn-Sham Hamiltonian matrices for each angular momentum channel @@ -251,12 +139,12 @@ def _compute_fixed_matrices(self): from ..pseudo.non_local import NonLocalPseudopotential nonlocal_calculator = NonLocalPseudopotential( - pseudo=self.pseudo, - ops_builder=self.ops_builder + pseudo = self.pseudo, + ops_builder = self.ops_builder ) self.H_nonlocal = nonlocal_calculator.compute_all_nonlocal_matrices( - l_channels=self.occupation_info.unique_l_values + l_channels = self.occupation_info.unique_l_values ) else: self.H_nonlocal = {} @@ -264,15 +152,16 @@ def _compute_fixed_matrices(self): def build_for_l_channel( self, - l: int, - v_hartree: np.ndarray, - v_x: np.ndarray, - v_c: np.ndarray, - switches: SwitchesFlags, - v_x_oep: Optional[np.ndarray] = None, - v_c_oep: Optional[np.ndarray] = None, - de_xc_dtau: Optional[np.ndarray] = None, - symmetrize: bool = False + l : int, + v_hartree : np.ndarray, + v_x : np.ndarray, + v_c : np.ndarray, + switches : 'SwitchesFlags', + v_x_oep : Optional[np.ndarray] = None, + v_c_oep : Optional[np.ndarray] = None, + de_xc_dtau : Optional[np.ndarray] = None, + symmetrize : bool = False, + exclude_boundary : bool = False, ) -> np.ndarray: """ Build total Hamiltonian for angular momentum channel l @@ -287,7 +176,7 @@ def build_for_l_channel( Exchange potential V_x(r) at quadrature points v_c : np.ndarray Correlation potential V_c(r) at quadrature points - switches : SwitchesFlags + switches : 'SwitchesFlags' Functional switches determining which components to include v_x_oep : np.ndarray, optional OEP exchange potential (if used) @@ -298,6 +187,8 @@ def build_for_l_channel( symmetrize : bool, optional Whether to apply symmetrization using S^(-1/2) (default: False) Transforms: H → S^(-1/2) @ H @ S^(-1/2), then H → (H+H^T)/2 + exclude_boundary : bool, optional + Whether to exclude boundary nodes from the Hamiltonian matrix (default: False) Returns ------- @@ -318,7 +209,7 @@ def build_for_l_channel( # Determine whether to mix the exchange functionals of HF and GGA_PBE hybrid_mixing_parameter = 0.0 - if switches.use_hf_exchange: + if switches.use_hf_exchange or switches.use_oep_exchange: assert isinstance(switches.hybrid_mixing_parameter, float), \ MIXING_PARAMETER_NOT_A_FLOAT_ERROR.format(type(switches.hybrid_mixing_parameter)) assert 0.0 <= switches.hybrid_mixing_parameter <= 1.0, \ @@ -347,22 +238,13 @@ def build_for_l_channel( v_xc_total = (1.0 - hybrid_mixing_parameter) * v_x v_xc_total += v_c if v_x_oep is not None: - v_xc_total += v_x_oep + v_xc_total += v_x_oep * hybrid_mixing_parameter if v_c_oep is not None: v_xc_total += v_c_oep H_xc = self.ops_builder.build_potential_matrix(v_xc_total) H += H_xc - - # np.savetxt("H_part.txt", self.H_kinetic + self.H_ext) - # np.savetxt("H_hartree.txt", H_hartree) - # np.savetxt("H_xc.txt", H_xc) - # np.savetxt("v_hartree.txt", v_hartree) - # np.savetxt("v_x.txt", v_x) - # np.savetxt("v_c.txt", v_c) - # raise RuntimeError("Stop here") - # Add meta-GGA kinetic density term (radial component) if switches.use_metagga: H_metagga_tau = self.ops_builder.build_metagga_kinetic_density_matrix(de_xc_dtau) @@ -387,11 +269,18 @@ def build_for_l_channel( assert l in self.H_hf_exchange_dict, \ HARTREE_FOCK_EXCHANGE_MATRIX_FOR_L_CHANNEL_NOT_AVAILABLE_ERROR.format(l) H += self.H_hf_exchange_dict[l] * hybrid_mixing_parameter - + + + # Exclude boundary nodes from the Hamiltonian matrix + if exclude_boundary: + H = H[1:-1,1:-1] + # Apply symmetrization transformation (if requested) if symmetrize: # Get S^(-1/2) from ops_builder S_inv_sqrt = self.ops_builder.get_S_inv_sqrt() + if exclude_boundary: + S_inv_sqrt = S_inv_sqrt[1:-1,1:-1] # Transform: H → S^(-1/2) @ H @ S^(-1/2) H = S_inv_sqrt @ H @ S_inv_sqrt diff --git a/utils/atom/solver.py b/utils/atom/solver.py index b0fe83c9..a6b7a860 100644 --- a/utils/atom/solver.py +++ b/utils/atom/solver.py @@ -24,9 +24,15 @@ from __future__ import annotations import os +os.environ["OMP_NUM_THREADS"] = "1" +os.environ["MKL_NUM_THREADS"] = "1" +os.environ["OPENBLAS_NUM_THREADS"] = "1" +os.environ["NUMEXPR_NUM_THREADS"] = "1" + import sys from pathlib import Path + # Fix the relative import issue when running as a script try: __package__ @@ -40,11 +46,15 @@ if str(parent_dir) not in sys.path: sys.path.insert(0, str(parent_dir)) + import numpy as np +np.set_printoptions(precision=20) +np.set_printoptions(threshold=sys.maxsize) + from typing import Optional, Dict, Any, Tuple # Mesh & operators -from .mesh.builder import Quadrature1D, Mesh1D, RPAFrequencyGrid +from .mesh.builder import Quadrature1D, Mesh1D from .mesh.operators import GridData, RadialOperatorsBuilder # Typing imports @@ -52,7 +62,7 @@ from .pseudo.non_local import NonLocalPseudopotential from .utils.occupation_states import OccupationInfo from .scf.energy import EnergyComponents -from .scf.driver import SCFResult +from .scf.driver import SCFResult, SwitchesFlags from .xc.evaluator import XCPotentialData from .xc.functional_requirements import get_functional_requirements @@ -65,7 +75,6 @@ EnergyCalculator, EigenSolver, Mixer, - ConvergenceChecker ) # XC and snapshot @@ -75,13 +84,27 @@ # Valid XC Functional -VALID_XC_FUNCTIONAL_LIST = ['GGA_PBE', 'RPA', 'OEPx', 'LDA_PZ', 'LDA_PW', 'SCAN', 'RSCAN', 'R2SCAN', 'PBE0', 'HF'] +VALID_XC_FUNCTIONAL_LIST = [ + 'LDA_PZ' , # LDA Perdew-Zunger + 'LDA_PW' , # LDA Perdew-Wang + 'GGA_PBE', # GGA Perdew-Burke-Ernzerhof + 'SCAN' , # SCAN functional, meta-GGA + 'RSCAN' , # RSCAN functional, meta-GGA + 'R2SCAN' , # R2SCAN functional, meta-GGA + 'HF' , # Hartree-Fock + 'PBE0' , # PBE0 Perdew-Burke-Ernzerhof, hybrid functional + 'EXX' , # Exact Exchange, using OEP method + 'RPA' , # Random Phase Approximation, with exact exchange +] + +# Valid XC Functional for OEP +VALID_XC_FUNCTIONAL_FOR_OEP_LIST = ['EXX', 'RPA', 'PBE0'] + # Valid Mesh Type VALID_MESH_TYPE_LIST = ['exponential', 'polynomial', 'uniform'] - # Type Check Error Messages ATOMIC_NUMBER_NOT_INTEGER_ERROR = \ "parameter atomic_number must be an integer, get {} instead" @@ -103,22 +126,22 @@ "parameter scf_tolerance must be a float, get {} instead" ALL_ELECTRON_FLAG_NOT_BOOL_ERROR = \ "parameter all_electron_flag must be a boolean, get {} instead" -ENABLE_DOMAIN_SIZE_TEST_NOT_BOOL_ERROR = \ - "parameter enable_domain_size_test must be a boolean, get {} instead" +USE_OEP_NOT_BOOL_ERROR = \ + "parameter use_oep must be a boolean, get {} instead" +USE_OEP_NOT_TRUE_FOR_OEP_FUNCTIONAL_ERROR = \ + "parameter use_oep must be True for OEP functional `{}`, get {} instead" +USE_OEP_NOT_FALSE_FOR_NON_OEP_FUNCTIONAL_ERROR = \ + "parameter use_oep must be False for non-OEP functional `{}`, get {} instead" PSP_DIR_PATH_NOT_STRING_ERROR = \ "parameter psp_dir_path must be a string, get {} instead" PSP_FILE_NAME_NOT_STRING_ERROR = \ "parameter psp_file_name must be a string, get {} instead" -FREQUENCY_INTEGRATION_POINT_NUMBER_NOT_INTEGER_ERROR = \ - "parameter frequency_integration_point_number must be an integer, get {} instead" -EIGENTOLERANCE_X_NOT_FLOAT_ERROR = \ - "parameter eigentolerance_x must be a float, get {} instead" -EIGENTOLERANCE_C_NOT_FLOAT_ERROR = \ - "parameter eigentolerance_c must be a float, get {} instead" -L_MAX_QUANTUM_NUMBER_NOT_INTEGER_ERROR = \ - "parameter l_max_quantum_number must be an integer, get {} instead" -SMOOTHENING_CUTOFF_FREQUENCY_NOT_FLOAT_ERROR = \ - "parameter smoothening_cutoff_frequency must be a float, get {} instead" +FREQUENCY_QUADRATURE_POINT_NUMBER_NOT_INTEGER_ERROR = \ + "parameter frequency_quadrature_point_number must be an integer, get {} instead" +ANGULAR_MOMENTUM_CUTOFF_NOT_INTEGER_ERROR = \ + "parameter angular_momentum_cutoff must be an integer, get {} instead" +ENABLE_PARALLELIZATION_NOT_BOOL_ERROR = \ + "parameter enable_parallelization must be a boolean, get {} instead" DOUBLE_HYBRID_FLAG_NOT_BOOL_ERROR = \ "parameter double_hybrid_flag must be a boolean, get {} instead" HYBRID_MIXING_PARAMETER_NOT_FLOAT_ERROR = \ @@ -127,6 +150,15 @@ "parameter mesh_spacing must be a float, get {} instead" PRINT_DEBUG_NOT_BOOL_ERROR = \ "parameter print_debug must be a boolean, get {} instead" +OEP_BASIS_NUMBER_NOT_INTEGER_ERROR = \ + "parameter oep_basis_number must be an integer, get {} instead" +OEP_BASIS_NUMBER_NOT_GREATER_THAN_0_ERROR = \ + "parameter oep_basis_number must be greater than 0, get {} instead" +OEP_MIXING_PARAMETER_NOT_FLOAT_ERROR = \ + "parameter oep_mixing_parameter must be a float, get {} instead" +OEP_MIXING_PARAMETER_NOT_IN_ZERO_ONE_ERROR = \ + "parameter oep_mixing_parameter must be in [0, 1], get {} instead" + # Value Check Error Messages ATOMIC_NUMBER_NOT_GREATER_THAN_0_ERROR = \ @@ -142,10 +174,10 @@ QUADRATURE_POINT_NUMBER_NOT_GREATER_THAN_0_ERROR = \ "parameter quadrature_point_number must be greater than 0, get {} instead" -FREQUENCY_INTEGRATION_POINT_NUMBER_NOT_GREATER_THAN_0_ERROR = \ - "parameter frequency_integration_point_number must be greater than 0, get {} instead" -L_MAX_QUANTUM_NUMBER_NEGATIVE_ERROR = \ - "parameter l_max_quantum_number must be non-negative, get {} instead" +FREQUENCY_QUADRATURE_POINT_NUMBER_NOT_GREATER_THAN_0_ERROR = \ + "parameter frequency_quadrature_point_number must be greater than 0, get {} instead" +ANGULAR_MOMENTUM_CUTOFF_NEGATIVE_ERROR = \ + "parameter angular_momentum_cutoff must be non-negative, get {} instead" XC_FUNCTIONAL_TYPE_ERROR_MESSAGE = \ "parameter xc_functional must be a string, get type {} instead" XC_FUNCTIONAL_NOT_IN_VALID_LIST_ERROR = \ @@ -160,16 +192,13 @@ "parameter default psp directory path {} does not exist, please provide a valid psp directory path" PSP_FILE_NAME_NOT_EXISTS_ERROR = \ "parameter psp file name `{}` does not exist in the psp file path `{}`, please provide a valid psp file name" -EIGENTOLERANCE_X_NOT_GREATER_THAN_0_ERROR = \ - "parameter eigentolerance_x must be greater than 0, get {} instead" -EIGENTOLERANCE_C_NOT_GREATER_THAN_0_ERROR = \ - "parameter eigentolerance_c must be greater than 0, get {} instead" -SMOOTHENING_CUTOFF_FREQUENCY_NEGATIVE_ERROR = \ - "parameter smoothening_cutoff_frequency must be non-negative, get {} instead" HYBRID_MIXING_PARAMETER_NOT_IN_ZERO_ONE_ERROR = \ "parameter hybrid_mixing_parameter must be in [0, 1], get {} instead" HYBRID_MIXING_PARAMETER_NOT_ONE_FOR_NON_HYBRID_FUNCTIONAL_ERROR = \ "parameter hybrid_mixing_parameter must be 1.0 for non-hybrid functional, get {} instead" +HYBRID_MIXING_PARAMETER_NOT_ONE_ERROR = \ + "parameter hybrid_mixing_parameter must be 1.0 for functional {}, get {} instead" + DENSITY_MIXING_PARAMETER_NOT_FLOAT_ERROR = \ "parameter density_mixing_parameter must be a float, get {} instead" DENSITY_MIXING_PARAMETER_NOT_IN_ZERO_ONE_ERROR = \ @@ -184,14 +213,10 @@ "WARNING: parameter psp_dir_path is not None for all-electron calculation, so it will be ignored" PSP_FILE_NAME_NOT_NONE_FOR_ALL_ELECTRON_CALCULATION_WARNING = \ "WARNING: parameter psp_file_name is not None for all-electron calculation, so it will be ignored" -FREQUENCY_INTEGRATION_POINT_NUMBER_NOT_NONE_FOR_OEPX_AND_NONE_XC_FUNCTIONAL_WARNING = \ - "WARNING: parameter frequency_integration_point_number is not None for XC functional `{}`, so it will be ignored" -EIGENTOLERANCE_X_NOT_NONE_FOR_XC_FUNCTIONAL_OTHER_THAN_RPA_WARNING = \ - "WARNING: parameter eigentolerance_x is not None for XC functional `{}`, so it will be ignored" -EIGENTOLERANCE_C_NOT_NONE_FOR_XC_FUNCTIONAL_OTHER_THAN_RPA_WARNING = \ - "WARNING: parameter eigentolerance_c is not None for XC functional `{}`, so it will be ignored" -SMOOTHENING_CUTOFF_FREQUENCY_NOT_NONE_FOR_XC_FUNCTIONAL_OTHER_THAN_RPA_WARNING = \ - "WARNING: parameter smoothening_cutoff_frequency is not None for XC functional `{}`, so it will be ignored" +FREQUENCY_QUADRATURE_POINT_NUMBER_NOT_NONE_FOR_OEPX_AND_NONE_XC_FUNCTIONAL_WARNING = \ + "WARNING: parameter frequency_quadrature_point_number is not None for XC functional `{}`, so it will be ignored" +ANGULAR_MOMENTUM_CUTOFF_NOT_NONE_FOR_XC_FUNCTIONAL_OTHER_THAN_RPA_WARNING = \ + "WARNING: parameter angular_momentum_cutoff is not None for XC functional `{}`, so it will be ignored" NO_HYBRID_MIXING_PARAMETER_PROVIDED_FOR_HYBRID_FUNCTIONAL_WARNING = \ "WARNING: hybrid_mixing_parameter not provided for {} functional, using default value {}" HYBRID_MIXING_PARAMETER_NOT_IN_ZERO_ONE_WARNING = \ @@ -202,10 +227,59 @@ "WARNING: hybrid_mixing_parameter for {} must be a float, got {}" HYBRID_MIXING_PARAMETER_NOT_IN_ZERO_ONE_WARNING = \ "WARNING: hybrid_mixing_parameter for {} must be in [0, 1], got {}" -HYBRID_MIXING_PARAMETER_NOT_ONE_FOR_HF_ERROR = \ - "WARNING: hybrid_mixing_parameter for {} must be 1.0 for HF functional, got {}" WARM_START_NOT_CONVERGED_WARNING = \ "WARNING: warm start calculation for {} did not converge, using intermediate result" +ENABLE_PARALLELIZATION_NOT_NONE_FOR_XC_FUNCTIONAL_OTHER_THAN_RPA_WARNING = \ + "WARNING: parameter enable_parallelization is not None for XC functional `{}`, so it will be ignored" + +NUMPY_IMPORTED_BEFORE_ATOMIC_WARNING = \ + """ + [ATOM WARNING] NumPy was imported before the 'atom' package. + + This prevents the package from forcing BLAS libraries (MKL / OpenBLAS / NumExpr) + into single-thread mode. When parallel RPA calculations attempt to use + multiple Python threads or processes, each NumPy/SciPy linear algebra call may + internally spawn many BLAS threads. + + This can lead to: + • Severe CPU oversubscription (N_threads x BLAS_threads >> CPU cores), + • Very slow execution or apparent freezing/hanging, + • Poor scaling in parallel regions. + + To safely enable parallel execution, please choose ONE of the following: + + 1) Import 'atom' BEFORE importing NumPy/SciPy, e.g.: + + import atom + import numpy as np + + 2) Configure BLAS to single-thread mode *before importing NumPy*, e.g.: + + # In your shell (recommended): + export OMP_NUM_THREADS=1 + export MKL_NUM_THREADS=1 + export OPENBLAS_NUM_THREADS=1 + export NUMEXPR_NUM_THREADS=1 + + or in Python BEFORE NumPy is imported: + + import os + os.environ["OMP_NUM_THREADS"] = "1" + os.environ["MKL_NUM_THREADS"] = "1" + os.environ["OPENBLAS_NUM_THREADS"] = "1" + os.environ["NUMEXPR_NUM_THREADS"] = "1" + + 3) Install 'threadpoolctl' in your environment, which allows the package to + dynamically limit BLAS threads even if NumPy was already imported: + + pip install threadpoolctl + + After installation, 'atom' will automatically detect it and + apply safe single-thread limits for parallel execution. + + Parallel mode is disabled for this run to avoid deadlocks or CPU thrashing. + """ + class AtomicDFTSolver: """ @@ -217,95 +291,94 @@ class AtomicDFTSolver: """ # Basic physical parameters - atomic_number : int # Atomic number of the element to calculate (e.g., 13 for Aluminum) - domain_size : float # Radial computational domain size in atomic units (typically 10-30) - number_of_finite_elements : int # Number of finite elements in the computational domain - polynomial_order : int # Polynomial order of basis functions within each finite element - quadrature_point_number : int # Number of quadrature points for numerical integration (recommended: 3-4x polynomial_order) + atomic_number : int # Atomic number of the element to calculate (e.g., 13 for Aluminum) + domain_size : float # Radial computational domain size in atomic units (typically 10-30) + number_of_finite_elements : int # Number of finite elements in the computational domain + polynomial_order : int # Polynomial order of basis functions within each finite element + quadrature_point_number : int # Number of quadrature points for numerical integration (recommended: 3-4x polynomial_order) # Exchange-correlation functional parameters - xc_functional : str # XC functional type: 'GGA_PBE', 'RPA', 'OEPx', 'LDA_PZ', 'LDA_PW', 'SCAN', 'RSCAN', 'R2SCAN' - mesh_type : str # Mesh distribution type: 'exponential' (higher density near nucleus), 'polynomial', or 'uniform' - mesh_concentration : float # Mesh concentration parameter (controls point density distribution) + xc_functional : str # XC functional type: 'GGA_PBE', 'RPA', 'EXX', 'LDA_PZ', 'LDA_PW', 'SCAN', 'RSCAN', 'R2SCAN' + mesh_type : str # Mesh distribution type: 'exponential' (higher density near nucleus), 'polynomial', or 'uniform' + mesh_concentration : float # Mesh concentration parameter (controls point density distribution) # Self-consistent field (SCF) convergence parameters - scf_tolerance : float # SCF convergence tolerance (typically 1e-8) - all_electron_flag : bool # True for all-electron calculation, False for pseudopotential calculation - enable_domain_size_test : bool # Flag for domain size convergence testing + scf_tolerance : float # SCF convergence tolerance (typically 1e-8) + all_electron_flag : bool # True for all-electron calculation, False for pseudopotential calculation + use_oep : bool # Enable optimized effective potential workflow in SCF # Pseudopotential parameters - psp_dir_path : str # Path to pseudopotential files directory (required when all_electron_flag=False) - psp_file_name : str # Name of the pseudopotential file (required when all_electron_flag=False) + psp_dir_path : str # Path to pseudopotential files directory (required when all_electron_flag=False) + psp_file_name : str # Name of the pseudopotential file (required when all_electron_flag=False) - # Advanced functional parameters (for OEPx, RPA, etc.) - frequency_integration_point_number : int # Number of frequency integration points for RPA calculations - eigentolerance_x : float # Eigenvalue convergence tolerance for exchange term - eigentolerance_c : float # Eigenvalue convergence tolerance for correlation term - l_max_quantum_number : int # Maximum angular momentum quantum number to include - smoothening_cutoff_frequency : float # Smoothing cutoff frequency for numerical stability - double_hybrid_flag : bool # Flag for double-hybrid functional methods - hybrid_mixing_parameter : float # Mixing parameter for hybrid/double-hybrid functionals (e.g., 0.25 for PBE0) + # Advanced functional parameters (for EXX, RPA, etc.) + frequency_quadrature_point_number : int # Number of frequency quadrature points for RPA calculations + angular_momentum_cutoff : int # Maximum angular momentum quantum number to include + double_hybrid_flag : bool # Flag for double-hybrid functional methods + hybrid_mixing_parameter : float # Mixing parameter for hybrid/double-hybrid functionals (e.g., 0.25 for PBE0) + oep_basis_number : int # Basis size used in OEP calculations when enabled + enable_parallelization : bool # Flag for parallelization of RPA calculations # Grid and computational parameters - mesh_spacing : float # Minimum mesh spacing (should match output file spacing) - density_mixing_parameter : float # Density mixing parameter for SCF convergence (alpha in linear mixing) - print_debug : bool # Flag for printing debug information during calculation + mesh_spacing : float # Minimum mesh spacing (should match output file spacing) + density_mixing_parameter : float # Density mixing parameter for SCF convergence (alpha in linear mixing) + oep_mixing_parameter : float # Scaling parameter (λ) for OEP exchange/correlation potentials + print_debug : bool # Flag for printing debug information during calculation def __init__(self, - atomic_number : int, # Only atomic_number is required, all other parameters have default values - domain_size : Optional[float] = None, # 20.0 by default - number_of_finite_elements : Optional[int] = None, # 17 by default - polynomial_order : Optional[int] = None, # 31 by default - quadrature_point_number : Optional[int] = None, # 95 by default - xc_functional : Optional[str] = None, # 'GGA_PBE' by default - mesh_type : Optional[str] = None, # 'exponential' by default - mesh_concentration : Optional[float] = None, # 61.0 by default - scf_tolerance : Optional[float] = None, # 1e-8 by default - all_electron_flag : Optional[bool] = None, # False by default - enable_domain_size_test : Optional[bool] = None, # False by default - psp_dir_path : Optional[str] = None, # ../psps by default - psp_file_name : Optional[str] = None, # {atomic_number}.psp8 by default - frequency_integration_point_number : Optional[int] = None, # if xc_functional is 'RPA', 1200 by default - eigentolerance_x : Optional[float] = None, # for RPA, 1e-9 by default; for OEPx, 1e-11 by default, otherwise not needed - eigentolerance_c : Optional[float] = None, # for RPA, 1e-9 by default; for OEPx, 1e-09 by default, otherwise not needed - l_max_quantum_number : Optional[int] = None, # for RPA, 4 by default; for OEPx, 8 by default, otherwise 0 by default - smoothening_cutoff_frequency : Optional[float] = None, # for xc other than RPA, not needed; otherwise 0.0 for most atoms - double_hybrid_flag : Optional[bool] = None, # False by default - hybrid_mixing_parameter : Optional[float] = None, # 1.0 by default (0.25 for PBE0, variable for RPA) - mesh_spacing : Optional[float] = None, # 0.1 by default - density_mixing_parameter : Optional[float] = None, # 0.5 by default (alpha in linear mixing) - print_debug : Optional[bool] = None, # False by default + atomic_number : int, # Only atomic_number is required, all other parameters have default values + domain_size : Optional[float] = None, # 20.0 by default + number_of_finite_elements : Optional[int] = None, # 17 by default + polynomial_order : Optional[int] = None, # 31 by default + quadrature_point_number : Optional[int] = None, # 95 by default + oep_basis_number : Optional[int] = None, # not needed by default, if needed, int(polynomial_order * 0.25) by default + xc_functional : Optional[str] = None, # 'GGA_PBE' by default + mesh_type : Optional[str] = None, # 'exponential' by default + mesh_concentration : Optional[float] = None, # 61.0 by default + scf_tolerance : Optional[float] = None, # 1e-8 by default + all_electron_flag : Optional[bool] = None, # False by default + use_oep : Optional[bool] = None, # False by default + psp_dir_path : Optional[str] = None, # ../psps by default + psp_file_name : Optional[str] = None, # {atomic_number}.psp8 by default + frequency_quadrature_point_number : Optional[int] = None, # for RPA, 25 by default, otherwise not needed + angular_momentum_cutoff : Optional[int] = None, # for RPA, 4 by default, otherwise not needed + enable_parallelization : Optional[bool] = None, # for RPA, False by default, otherwise not needed + double_hybrid_flag : Optional[bool] = None, # False by default # TODO: implement double hybrid functional, should be simple + hybrid_mixing_parameter : Optional[float] = None, # 1.0 by default (0.25 for PBE0, variable for RPA) + mesh_spacing : Optional[float] = None, # 0.1 by default + density_mixing_parameter : Optional[float] = None, # 0.5 by default (alpha in linear mixing) + oep_mixing_parameter : Optional[float] = None, # 1.0 by default (scales OEP potentials), used for double hybrid functional only + print_debug : Optional[bool] = None, # False by default ): """ Initialize the AtomicDFTSolver with computational parameters. Args: - atomic_number (int) : Atomic number of the element (e.g., 13 for Aluminum) - domain_size (float): Radial domain size in atomic units (typically 10-30) - number_of_finite_elements (int) : Number of finite elements in the domain - polynomial_order (int) : Polynomial order of basis functions (typically 20-40) - quadrature_point_number (int) : Quadrature points for integration (3-4x polynomial_order) - xc_functional (str) : Exchange-correlation functional ('GGA_PBE', 'RPA', 'OEPx', etc.) - mesh_type (str) : Mesh type ('exponential', 'polynomial', 'uniform') - mesh_concentration (float): Mesh concentration parameter (controls point density) - scf_tolerance (float): SCF convergence tolerance (typically 1e-8) - all_electron_flag (bool) : True for all-electron, False for pseudopotential - enable_domain_size_test (bool) : Enable domain size convergence testing - psp_dir_path (str) : Path to pseudopotential directory (required if all_electron_flag=False) - psp_file_name (str) : Name of pseudopotential file (required if all_electron_flag=False) - frequency_integration_point_number (int) : Frequency points for RPA calculations - eigentolerance_x (float): Exchange eigenvalue convergence tolerance - eigentolerance_c (float): Correlation eigenvalue convergence tolerance - l_max_quantum_number (int) : Maximum angular momentum quantum number - smoothening_cutoff_frequency (float): Smoothing cutoff for numerical stability - double_hybrid_flag (bool) : Enable double-hybrid functional methods - hybrid_mixing_parameter (float): Mixing parameter for hybrid functionals (e.g., 0.25 for PBE0) - mesh_spacing (float): Minimum mesh spacing (should match output file) - density_mixing_parameter (float): Density mixing parameter for SCF (alpha in linear mixing) - print_debug (bool) : Enable debug output + atomic_number (int) : Atomic number of the element (e.g., 13 for Aluminum) + domain_size (float) : Radial domain size in atomic units (typically 10-30) + number_of_finite_elements (int) : Number of finite elements in the domain + polynomial_order (int) : Polynomial order of basis functions (typically 20-40) + quadrature_point_number (int) : Quadrature points for integration (3-4x polynomial_order) + xc_functional (str) : Exchange-correlation functional ('GGA_PBE', 'RPA', 'EXX', etc.) + mesh_type (str) : Mesh type ('exponential', 'polynomial', 'uniform') + mesh_concentration (float) : Mesh concentration parameter (controls point density) + scf_tolerance (float) : SCF convergence tolerance (typically 1e-8) + all_electron_flag (bool) : True for all-electron, False for pseudopotential + use_oep (bool) : Enable optimized effective potential calculations (default: False) + oep_basis_number (int) : Size of OEP auxiliary basis; defaults to int(polynomial_order * 0.25) when use_oep=True + psp_dir_path (str) : Path to pseudopotential directory (required if all_electron_flag=False) + psp_file_name (str) : Name of pseudopotential file (required if all_electron_flag=False) + frequency_quadrature_point_number (int) : Frequency quadrature points for RPA calculations, used for RPA functional only + angular_momentum_cutoff (int) : Maximum angular momentum quantum number, used for RPA functional only + double_hybrid_flag (bool) : Enable double-hybrid functional methods + hybrid_mixing_parameter (float) : Mixing parameter for hybrid functionals (e.g., 0.25 for PBE0) + mesh_spacing (float) : Mesh spacing for the uniform grid, used to set the output mesh spacing + density_mixing_parameter (float) : Density mixing parameter for SCF (alpha in linear mixing) + oep_mixing_parameter (float) : Mixing parameter for OEP functionals (lambda in OEP) + print_debug (bool) : Enable debug output Raises: ValueError: If any parameter has incorrect type @@ -316,29 +389,29 @@ def __init__(self, """ # Initialize the class attributes - self.atomic_number = atomic_number - self.domain_size = domain_size - self.number_of_finite_elements = number_of_finite_elements - self.polynomial_order = polynomial_order - self.quadrature_point_number = quadrature_point_number - self.xc_functional = xc_functional - self.mesh_type = mesh_type - self.mesh_concentration = mesh_concentration - self.scf_tolerance = scf_tolerance - self.all_electron_flag = all_electron_flag - self.enable_domain_size_test = enable_domain_size_test - self.psp_dir_path = psp_dir_path - self.psp_file_name = psp_file_name - self.frequency_integration_point_number = frequency_integration_point_number - self.eigentolerance_x = eigentolerance_x - self.eigentolerance_c = eigentolerance_c - self.l_max_quantum_number = l_max_quantum_number - self.smoothening_cutoff_frequency = smoothening_cutoff_frequency - self.double_hybrid_flag = double_hybrid_flag - self.hybrid_mixing_parameter = hybrid_mixing_parameter - self.mesh_spacing = mesh_spacing - self.density_mixing_parameter = density_mixing_parameter - self.print_debug = print_debug + self.atomic_number = atomic_number + self.domain_size = domain_size + self.number_of_finite_elements = number_of_finite_elements + self.polynomial_order = polynomial_order + self.quadrature_point_number = quadrature_point_number + self.xc_functional = xc_functional + self.mesh_type = mesh_type + self.mesh_concentration = mesh_concentration + self.scf_tolerance = scf_tolerance + self.all_electron_flag = all_electron_flag + self.use_oep = use_oep + self.oep_basis_number = oep_basis_number + self.psp_dir_path = psp_dir_path + self.psp_file_name = psp_file_name + self.frequency_quadrature_point_number = frequency_quadrature_point_number + self.angular_momentum_cutoff = angular_momentum_cutoff + self.enable_parallelization = enable_parallelization + self.double_hybrid_flag = double_hybrid_flag + self.hybrid_mixing_parameter = hybrid_mixing_parameter + self.mesh_spacing = mesh_spacing + self.density_mixing_parameter = density_mixing_parameter + self.oep_mixing_parameter = oep_mixing_parameter + self.print_debug = print_debug # set the default parameters, if not provided @@ -357,31 +430,36 @@ def __init__(self, all_electron_flag = self.all_electron_flag) - # XC evaluator for delta-learning usage (to be initialized when needed) - self.xc_evaluator = None # TODO: Initialize when delta-learning is implemented - # Grid data and operators (initialized in __init__) self.grid_data_standard : Optional[GridData] = None self.grid_data_dense : Optional[GridData] = None + self.grid_data_oep : Optional[GridData] = None self.ops_builder_standard : Optional[RadialOperatorsBuilder] = None self.ops_builder_dense : Optional[RadialOperatorsBuilder] = None + self.ops_builder_oep : Optional[RadialOperatorsBuilder] = None # SCF components (initialized in __init__) - self.hamiltonian_builder : Optional[HamiltonianBuilder] = None - self.density_calculator : Optional[DensityCalculator] = None - self.poisson_solver : Optional[PoissonSolver] = None - self.energy_calculator : Optional[EnergyCalculator] = None - self.scf_driver : Optional[SCFDriver] = None + self.hamiltonian_builder : Optional[HamiltonianBuilder] = None + self.density_calculator : Optional[DensityCalculator] = None + self.poisson_solver : Optional[PoissonSolver] = None + self.energy_calculator : Optional[EnergyCalculator] = None + self.scf_driver : Optional[SCFDriver] = None # Initialize grids and operators - self.grid_data_standard, self.grid_data_dense = self._initialize_grids() + self.grid_data_standard, self.grid_data_dense, self.grid_data_oep = self._initialize_grids() self.ops_builder_standard = RadialOperatorsBuilder.from_grid_data( self.grid_data_standard, verbose=self.print_debug ) self.ops_builder_dense = RadialOperatorsBuilder.from_grid_data( self.grid_data_dense, verbose=self.print_debug ) - + + if self.use_oep: + # Initialize OEP operators builder + self.ops_builder_oep = RadialOperatorsBuilder.from_grid_data( + self.grid_data_oep, verbose=self.print_debug + ) + # Initialize SCF components self._initialize_scf_components( ops_builder_standard = self.ops_builder_standard, @@ -451,6 +529,7 @@ def set_and_check_initial_parameters(self): assert self.xc_functional in VALID_XC_FUNCTIONAL_LIST, \ XC_FUNCTIONAL_NOT_IN_VALID_LIST_ERROR.format(VALID_XC_FUNCTIONAL_LIST, self.xc_functional) + # mesh type if self.mesh_type is None: self.mesh_type = 'exponential' @@ -462,9 +541,9 @@ def set_and_check_initial_parameters(self): # mesh concentration if self.mesh_concentration is None: # default value if self.mesh_type == 'exponential': - self.mesh_concentration = 61.0 + self.mesh_concentration = 100.0 elif self.mesh_type == 'polynomial': - self.mesh_concentration = 3.0 + self.mesh_concentration = 2.0 elif self.mesh_type == 'uniform': self.mesh_concentration = None if self.mesh_type in ['exponential', 'polynomial']: # type check @@ -487,9 +566,9 @@ def set_and_check_initial_parameters(self): if self.scf_tolerance is None: # For most functionals, the default tolerance is 1e-8 self.scf_tolerance = 1e-8 - if self.xc_functional in ['SCAN']: - # SCAN functional suffers from convergence issues, so we use a higher tolerance - self.scf_tolerance = 1e-4 + if self.xc_functional in ['SCAN', 'RSCAN', 'R2SCAN']: + # SCAN, RSCAN, R2SCAN functionals suffer from convergence issues, so we use a higher tolerance + self.scf_tolerance = 1e-6 try: self.scf_tolerance = float(self.scf_tolerance) except: @@ -507,13 +586,61 @@ def set_and_check_initial_parameters(self): assert isinstance(self.all_electron_flag, bool), \ ALL_ELECTRON_FLAG_NOT_BOOL_ERROR.format(type(self.all_electron_flag)) - # enable domain size test - if self.enable_domain_size_test is None: - self.enable_domain_size_test = False - if self.enable_domain_size_test in [0, 1]: - self.enable_domain_size_test = False if self.enable_domain_size_test == 0 else True - assert isinstance(self.enable_domain_size_test, bool), \ - ENABLE_DOMAIN_SIZE_TEST_NOT_BOOL_ERROR.format(type(self.enable_domain_size_test)) + # use OEP flag + if self.use_oep in [0, 1]: + self.use_oep = False if self.use_oep == 0 else True + if self.xc_functional in VALID_XC_FUNCTIONAL_FOR_OEP_LIST: + # OEP functional must be used with OEP flag + if self.xc_functional == 'PBE0': + if self.use_oep is None: + self.use_oep = False + else: + if self.use_oep is None: + self.use_oep = True + assert self.use_oep is True, \ + USE_OEP_NOT_TRUE_FOR_OEP_FUNCTIONAL_ERROR.format(self.xc_functional, self.use_oep) + else: + # Other functionals must not be used with OEP flag, otherwise raise error + if self.use_oep is None: + self.use_oep = False + assert self.use_oep is False, \ + USE_OEP_NOT_FALSE_FOR_NON_OEP_FUNCTIONAL_ERROR.format(self.xc_functional, self.use_oep) + + assert isinstance(self.use_oep, bool), \ + USE_OEP_NOT_BOOL_ERROR.format(type(self.use_oep)) + + + # OEP auxiliary basis size + if self.oep_basis_number is None: + if self.use_oep: + default_oep_basis = max(1, int(self.polynomial_order * 0.25)) + self.oep_basis_number = default_oep_basis + else: + if not isinstance(self.oep_basis_number, int): + try: + self.oep_basis_number = int(self.oep_basis_number) + except Exception: + raise ValueError(OEP_BASIS_NUMBER_NOT_INTEGER_ERROR.format(type(self.oep_basis_number))) + assert isinstance(self.oep_basis_number, int), \ + OEP_BASIS_NUMBER_NOT_INTEGER_ERROR.format(type(self.oep_basis_number)) + assert self.oep_basis_number > 0, \ + OEP_BASIS_NUMBER_NOT_GREATER_THAN_0_ERROR.format(self.oep_basis_number) + + + # OEP mixing parameter (λ scaling for OEP potentials) + if self.oep_mixing_parameter is None: + if self.use_oep: + self.oep_mixing_parameter = 1.0 + else: + if not isinstance(self.oep_mixing_parameter, float): + try: + self.oep_mixing_parameter = float(self.oep_mixing_parameter) + except: + raise ValueError(OEP_MIXING_PARAMETER_NOT_FLOAT_ERROR.format(type(self.oep_mixing_parameter))) + assert isinstance(self.oep_mixing_parameter, float), \ + OEP_MIXING_PARAMETER_NOT_FLOAT_ERROR.format(type(self.oep_mixing_parameter)) + assert self.oep_mixing_parameter > 0.0 and self.oep_mixing_parameter <= 1.0, \ + OEP_MIXING_PARAMETER_NOT_IN_ZERO_ONE_ERROR.format(self.oep_mixing_parameter) # psp directory path if self.all_electron_flag == False: @@ -555,78 +682,45 @@ def set_and_check_initial_parameters(self): # frequency integration point number if self.xc_functional in ['RPA', ]: - if self.frequency_integration_point_number is None: - self.frequency_integration_point_number = 1200 - assert isinstance(self.frequency_integration_point_number, int), \ - FREQUENCY_INTEGRATION_POINT_NUMBER_NOT_INTEGER_ERROR.format(type(self.frequency_integration_point_number)) - assert self.frequency_integration_point_number > 0, \ - FREQUENCY_INTEGRATION_POINT_NUMBER_NOT_GREATER_THAN_0_ERROR.format(self.frequency_integration_point_number) + if self.frequency_quadrature_point_number is None: + self.frequency_quadrature_point_number = 25 + assert isinstance(self.frequency_quadrature_point_number, int), \ + FREQUENCY_QUADRATURE_POINT_NUMBER_NOT_INTEGER_ERROR.format(type(self.frequency_quadrature_point_number)) + assert self.frequency_quadrature_point_number > 0, \ + FREQUENCY_QUADRATURE_POINT_NUMBER_NOT_GREATER_THAN_0_ERROR.format(self.frequency_quadrature_point_number) else: - if self.frequency_integration_point_number is not None: - print(FREQUENCY_INTEGRATION_POINT_NUMBER_NOT_NONE_FOR_OEPX_AND_NONE_XC_FUNCTIONAL_WARNING.format(self.xc_functional)) - self.frequency_integration_point_number = None - - # eigentolerance x - if self.xc_functional in ['OEPx', 'RPA']: # Default value: 1e-11 for OEPx, 1e-9 for RPA - if self.eigentolerance_x is None: - self.eigentolerance_x = 1e-11 if self.xc_functional == 'OEPx' else 1e-9 - assert isinstance(self.eigentolerance_x, float), \ - EIGENTOLERANCE_X_NOT_FLOAT_ERROR.format(type(self.eigentolerance_x)) - assert self.eigentolerance_x > 0., \ - EIGENTOLERANCE_X_NOT_GREATER_THAN_0_ERROR.format(self.eigentolerance_x) + if self.frequency_quadrature_point_number is not None: + print(FREQUENCY_QUADRATURE_POINT_NUMBER_NOT_NONE_FOR_OEPX_AND_NONE_XC_FUNCTIONAL_WARNING.format(self.xc_functional)) + self.frequency_quadrature_point_number = None + + # angular_momentum_cutoff, used for RPA functional only + if self.xc_functional in ['RPA']: + if self.angular_momentum_cutoff is None: + self.angular_momentum_cutoff = 4 + assert isinstance(self.angular_momentum_cutoff, int), \ + ANGULAR_MOMENTUM_CUTOFF_NOT_INTEGER_ERROR.format(type(self.angular_momentum_cutoff)) + assert self.angular_momentum_cutoff >= 0., \ + ANGULAR_MOMENTUM_CUTOFF_NEGATIVE_ERROR.format(self.angular_momentum_cutoff) else: - if self.eigentolerance_x is not None: - print(EIGENTOLERANCE_X_NOT_NONE_FOR_XC_FUNCTIONAL_OTHER_THAN_RPA_WARNING.format(self.xc_functional)) - self.eigentolerance_x = None - - # eigentolerance c - if self.xc_functional in ['OEPx', 'RPA']: # Default value: 1e-9 for OEPx, 1e-9 for RPA - if self.eigentolerance_c is None: - self.eigentolerance_c = 1e-9 - assert isinstance(self.eigentolerance_c, float), \ - EIGENTOLERANCE_C_NOT_FLOAT_ERROR.format(type(self.eigentolerance_c)) - assert self.eigentolerance_c > 0., \ - EIGENTOLERANCE_C_NOT_GREATER_THAN_0_ERROR.format(self.eigentolerance_c) - else: - if self.eigentolerance_c is not None: - print(EIGENTOLERANCE_C_NOT_NONE_FOR_XC_FUNCTIONAL_OTHER_THAN_RPA_WARNING.format(self.xc_functional)) - self.eigentolerance_c = None - - # l_max_quantum_number - if self.l_max_quantum_number is None: - if self.xc_functional == 'OEPx': - self.l_max_quantum_number = 8 - elif self.xc_functional == 'RPA': - self.l_max_quantum_number = 4 - else: - self.l_max_quantum_number = 0 - assert isinstance(self.l_max_quantum_number, int), \ - L_MAX_QUANTUM_NUMBER_NOT_INTEGER_ERROR.format(type(self.l_max_quantum_number)) - assert self.l_max_quantum_number >= 0., \ - L_MAX_QUANTUM_NUMBER_NEGATIVE_ERROR.format(self.l_max_quantum_number) + if self.angular_momentum_cutoff is not None: + print(ANGULAR_MOMENTUM_CUTOFF_NOT_NONE_FOR_XC_FUNCTIONAL_OTHER_THAN_RPA_WARNING.format(self.xc_functional)) + self.angular_momentum_cutoff = None + + # enable parallelization flag + if self.xc_functional in ['RPA']: + if self.enable_parallelization is None: + self.enable_parallelization = False + assert isinstance(self.enable_parallelization, bool), \ + ENABLE_PARALLELIZATION_NOT_BOOL_ERROR.format(type(self.enable_parallelization)) + if self.enable_parallelization: + from . import _NUMPY_IMPORTED_BEFORE_ATOMIC, _BLAS_ENV_SINGLE_THREADED, _THREADPOOLCTL_INSTALLED + if _NUMPY_IMPORTED_BEFORE_ATOMIC and not _BLAS_ENV_SINGLE_THREADED and not _THREADPOOLCTL_INSTALLED: + print(NUMPY_IMPORTED_BEFORE_ATOMIC_WARNING) + self.enable_parallelization = False - # smoothening cutoff frequency - if self.xc_functional in ['RPA', ]: - if self.smoothening_cutoff_frequency is None: - self.smoothening_cutoff_frequency = 60.0 - if self.atomic_number == 2: - self.smoothening_cutoff_frequency = 60.0 - elif self.atomic_number == 4: - self.smoothening_cutoff_frequency = 60.0 - elif self.atomic_number == 10: - self.smoothening_cutoff_frequency = 60.0 - elif self.atomic_number == 12: - self.smoothening_cutoff_frequency = 60.0 - elif self.atomic_number == 18: - self.smoothening_cutoff_frequency = 100.0 - assert isinstance(self.smoothening_cutoff_frequency, float), \ - SMOOTHENING_CUTOFF_FREQUENCY_NOT_FLOAT_ERROR.format(type(self.smoothening_cutoff_frequency)) - assert self.smoothening_cutoff_frequency >= 0., \ - SMOOTHENING_CUTOFF_FREQUENCY_NEGATIVE_ERROR.format(self.smoothening_cutoff_frequency) else: - if self.smoothening_cutoff_frequency is not None: - print(SMOOTHENING_CUTOFF_FREQUENCY_NOT_NONE_FOR_XC_FUNCTIONAL_OTHER_THAN_RPA_WARNING.format(self.xc_functional)) - self.smoothening_cutoff_frequency = None + if self.enable_parallelization is not None: + print(ENABLE_PARALLELIZATION_NOT_NONE_FOR_XC_FUNCTIONAL_OTHER_THAN_RPA_WARNING.format(self.xc_functional)) # double hybrid flag if self.double_hybrid_flag is None: @@ -638,23 +732,23 @@ def set_and_check_initial_parameters(self): # hybrid mixing parameter # Only validate for hybrid functionals (PBE0, HF) - if self.xc_functional in ['PBE0', 'HF']: + if self.xc_functional in ['PBE0', 'HF', 'EXX', 'RPA']: if self.hybrid_mixing_parameter is None: # Use default values based on functional if self.xc_functional == 'PBE0': self.hybrid_mixing_parameter = 0.25 # print(NO_HYBRID_MIXING_PARAMETER_PROVIDED_FOR_HYBRID_FUNCTIONAL_WARNING.format(self.xc_functional, 0.25)) - elif self.xc_functional == 'HF': + elif self.xc_functional in ['HF', 'EXX', 'RPA']: self.hybrid_mixing_parameter = 1.0 - + # If the hybrid mixing parameter is provided, check the type and value assert isinstance(self.hybrid_mixing_parameter, (float, int)), \ HYBRID_MIXING_PARAMETER_NOT_FLOAT_ERROR.format(type(self.hybrid_mixing_parameter)) assert 0.0 <= self.hybrid_mixing_parameter <= 1.0, \ HYBRID_MIXING_PARAMETER_NOT_IN_ZERO_ONE_ERROR.format(self.hybrid_mixing_parameter) - if self.xc_functional == "HF": + if self.xc_functional in ["HF", "EXX", "RPA"]: assert self.hybrid_mixing_parameter == 1.0, \ - HYBRID_MIXING_PARAMETER_NOT_ONE_FOR_HF_ERROR.format(self.hybrid_mixing_parameter) + HYBRID_MIXING_PARAMETER_NOT_ONE_ERROR.format(self.xc_functional, self.hybrid_mixing_parameter) else: # For non-hybrid functionals, hybrid_mixing_parameter is not used # Set it to None to avoid confusion @@ -670,6 +764,7 @@ def set_and_check_initial_parameters(self): assert 0.0 < self.density_mixing_parameter <= 1.0, \ DENSITY_MIXING_PARAMETER_NOT_IN_ZERO_ONE_ERROR.format(self.density_mixing_parameter) + # mesh spacing if self.mesh_spacing is None: self.mesh_spacing = 0.1 @@ -706,32 +801,31 @@ def print_input_parameters(self): print("\t\t INPUT PARAMETERS") print("=" * 60) - print("\t atomic_number : {}".format(self.atomic_number)) - print("\t domain_size : {}".format(self.domain_size)) - print("\t number_of_finite_elements : {}".format(self.number_of_finite_elements)) - print("\t polynomial_order : {}".format(self.polynomial_order)) - print("\t quadrature_point_number : {}".format(self.quadrature_point_number)) - print("\t xc_functional : {}".format(self.xc_functional)) - print("\t mesh_type : {}".format(self.mesh_type)) - print("\t mesh_concentration : {}".format(self.mesh_concentration)) - print("\t scf_tolerance : {}".format(self.scf_tolerance)) - print("\t all_electron_flag : {}".format(self.all_electron_flag)) - print("\t enable_domain_size_test : {}".format(self.enable_domain_size_test)) - print("\t psp_dir_path : {}".format(psp_path_display)) - print("\t psp_file_name : {}".format(self.psp_file_name)) - print("\t frequency_integration_point_number : {}".format(self.frequency_integration_point_number)) - print("\t eigentolerance_x : {}".format(self.eigentolerance_x)) - print("\t eigentolerance_c : {}".format(self.eigentolerance_c)) - print("\t l_max_quantum_number : {}".format(self.l_max_quantum_number)) - print("\t smoothening_cutoff_frequency : {}".format(self.smoothening_cutoff_frequency)) - print("\t double_hybrid_flag : {}".format(self.double_hybrid_flag)) - print("\t hybrid_mixing_parameter : {}".format(self.hybrid_mixing_parameter)) - print("\t mesh_spacing : {}".format(self.mesh_spacing)) - print("\t density_mixing_parameter : {}".format(self.density_mixing_parameter)) + print("\t atomic_number : {}".format(self.atomic_number)) + print("\t xc_functional : {}".format(self.xc_functional)) + print("\t domain_size : {}".format(self.domain_size)) + print("\t number_of_finite_elements : {}".format(self.number_of_finite_elements)) + print("\t polynomial_order : {}".format(self.polynomial_order)) + print("\t quadrature_point_number : {}".format(self.quadrature_point_number)) + print("\t oep_basis_number : {}".format(self.oep_basis_number)) + print("\t mesh_type : {}".format(self.mesh_type)) + print("\t mesh_concentration : {}".format(self.mesh_concentration)) + print("\t scf_tolerance : {}".format(self.scf_tolerance)) + print("\t all_electron_flag : {}".format(self.all_electron_flag)) + print("\t use_oep : {}".format(self.use_oep)) + print("\t psp_dir_path : {}".format(psp_path_display)) + print("\t psp_file_name : {}".format(self.psp_file_name)) + print("\t frequency_quadrature_point_number : {}".format(self.frequency_quadrature_point_number)) + print("\t angular_momentum_cutoff : {}".format(self.angular_momentum_cutoff)) + print("\t enable_parallelization : {}".format(self.enable_parallelization)) + print("\t double_hybrid_flag : {}".format(self.double_hybrid_flag)) + print("\t hybrid_mixing_parameter : {}".format(self.hybrid_mixing_parameter)) + print("\t mesh_spacing : {}".format(self.mesh_spacing)) + print("\t density_mixing_parameter : {}".format(self.density_mixing_parameter)) print() - def _initialize_grids(self) -> Tuple[GridData, GridData]: + def _initialize_grids(self) -> Tuple[GridData, GridData, Optional[GridData]]: """ Initialize finite element grids and quadrature. @@ -751,25 +845,26 @@ def _initialize_grids(self) -> Tuple[GridData, GridData]: # Generate mesh boundaries mesh1d = Mesh1D( - domain_radius = self.domain_size / 2.0, + domain_radius = self.domain_size / 2.0, finite_elements_num = self.number_of_finite_elements, mesh_type = self.mesh_type, clustering_param = self.mesh_concentration, exp_shift = getattr(self, 'exp_shift', None) ) + boundaries_nodes, _ = mesh1d.generate_mesh_nodes_and_width() # Generate standard FE nodes global_nodes = Mesh1D.generate_fe_nodes( boundaries_nodes = boundaries_nodes, - interp_nodes = interp_nodes_ref + interp_nodes = interp_nodes_ref ) # Generate refined FE nodes (for Hartree potential solver) refined_interp_nodes_ref = Mesh1D.refine_interpolation_nodes(interp_nodes_ref) refined_global_nodes = Mesh1D.generate_fe_nodes( boundaries_nodes = boundaries_nodes, - interp_nodes = refined_interp_nodes_ref + interp_nodes = refined_interp_nodes_ref ) # Generate Gauss-Legendre quadrature nodes and weights @@ -799,15 +894,37 @@ def _initialize_grids(self) -> Tuple[GridData, GridData]: quadrature_nodes = quadrature_nodes, quadrature_weights = quadrature_weights ) + + # For OEP method, extra set of grids are needed for solving the OEP equation + grid_data_oep : Optional[GridData] = None + + if self.use_oep: + # Generate Lobatto interpolation nodes for OEP basis + oep_interp_nodes_ref, _ = Quadrature1D.lobatto( + self.oep_basis_number + ) + + # Generate OEP basis nodes + oep_global_nodes = Mesh1D.generate_fe_nodes( + boundaries_nodes = boundaries_nodes, + interp_nodes = oep_interp_nodes_ref, + ) + + grid_data_oep = GridData( + number_of_finite_elements = self.number_of_finite_elements, + physical_nodes = oep_global_nodes, + quadrature_nodes = quadrature_nodes, + quadrature_weights = quadrature_weights + ) - return grid_data_standard, grid_data_dense + return grid_data_standard, grid_data_dense, grid_data_oep def _initialize_scf_components( self, - ops_builder_standard: RadialOperatorsBuilder, - grid_data_standard: GridData, - ops_builder_dense: RadialOperatorsBuilder, + ops_builder_standard : RadialOperatorsBuilder, + grid_data_standard : GridData, + ops_builder_dense : RadialOperatorsBuilder, ) -> None: """ Initialize all SCF components. @@ -845,7 +962,7 @@ def _initialize_scf_components( ) # SCF Driver (create first to get xc_calculator) - eigensolver = EigenSolver(xc_functional=self.xc_functional) + eigensolver = EigenSolver(xc_functional = self.xc_functional) mixer = Mixer( tol = self.scf_tolerance, alpha_lin = (self.density_mixing_parameter, self.density_mixing_parameter), @@ -855,29 +972,38 @@ def _initialize_scf_components( ) self.scf_driver = SCFDriver( - hamiltonian_builder = self.hamiltonian_builder, - density_calculator = self.density_calculator, - poisson_solver = self.poisson_solver, - eigensolver = eigensolver, - mixer = mixer, - occupation_info = self.occupation_info, - xc_functional = self.xc_functional, - hybrid_mixing_parameter = self.hybrid_mixing_parameter + hamiltonian_builder = self.hamiltonian_builder, + density_calculator = self.density_calculator, + poisson_solver = self.poisson_solver, + eigensolver = eigensolver, + mixer = mixer, + occupation_info = self.occupation_info, + xc_functional = self.xc_functional, + hybrid_mixing_parameter = self.hybrid_mixing_parameter, + use_oep = self.use_oep, + ops_builder_oep = self.ops_builder_oep, + oep_mixing_parameter = self.oep_mixing_parameter, + frequency_quadrature_point_number = self.frequency_quadrature_point_number, + angular_momentum_cutoff = self.angular_momentum_cutoff, + enable_parallelization = self.enable_parallelization, ) # Get XC calculator and HF calculator from scf_driver - xc_calculator = self.scf_driver.xc_calculator if hasattr(self.scf_driver, 'xc_calculator') else None - hf_calculator = self.scf_driver.hf_calculator if hasattr(self.scf_driver, 'hf_calculator') else None + xc_calculator = self.scf_driver.xc_calculator if hasattr(self.scf_driver, 'xc_calculator') else None + hf_calculator = self.scf_driver.hf_calculator if hasattr(self.scf_driver, 'hf_calculator') else None + oep_calculator = self.scf_driver.oep_calculator if hasattr(self.scf_driver, 'oep_calculator') else None # Energy calculator (uses standard grid data and ops_builder, but dense derivative matrix) self.energy_calculator = EnergyCalculator( + switches = self.scf_driver.switches, grid_data = grid_data_standard, occupation_info = self.occupation_info, ops_builder = ops_builder_standard, poisson_solver = self.poisson_solver, pseudo = self.pseudo, xc_calculator = xc_calculator, - hf_calculator = hf_calculator, # Pass HF calculator from SCFDriver + hf_calculator = hf_calculator, # Pass HF calculator from SCFDriver + oep_calculator = oep_calculator, # Pass OEP calculator from SCFDriver derivative_matrix = ops_builder_dense.derivative_matrix # Use dense grid derivative for accuracy ) @@ -899,7 +1025,7 @@ def _get_scf_settings(self, xc_functional: str) -> Dict[str, Any]: } # For functionals requiring outer loop (HF, OEP, RPA) - if xc_functional in ['HF', 'PBE0', 'OEPx', 'RPA']: + if xc_functional in ['HF', 'PBE0', 'EXX', 'RPA']: settings['outer_max_iter'] = 50 return settings @@ -1013,13 +1139,13 @@ def _get_initial_density_and_orbitals_with_warm_start( # Create temporary SCFDriver with specified xc_functional # Reuse existing hamiltonian_builder, density_calculator, poisson_solver in the main SCFDriver scf_driver_warm = SCFDriver( - hamiltonian_builder = self.scf_driver.hamiltonian_builder, - density_calculator = self.scf_driver.density_calculator, - poisson_solver = self.scf_driver.poisson_solver, - eigensolver = eigensolver_warm, - mixer = self.scf_driver.mixer, - occupation_info = self.scf_driver.occupation_info, - xc_functional = xc_functional, # Use specified functional + hamiltonian_builder = self.scf_driver.hamiltonian_builder, + density_calculator = self.scf_driver.density_calculator, + poisson_solver = self.scf_driver.poisson_solver, + eigensolver = eigensolver_warm, + mixer = self.scf_driver.mixer, + occupation_info = self.scf_driver.occupation_info, + xc_functional = xc_functional, # Use specified functional hybrid_mixing_parameter = self.scf_driver.hybrid_mixing_parameter ) @@ -1045,7 +1171,13 @@ def _get_initial_density_and_orbitals_with_warm_start( return rho_final, orbitals_final - def forward(self, orbitals) -> Dict[str, Any]: + def forward( + self, + orbitals : np.ndarray, + full_eigen_energies : Optional[np.ndarray] = None, + full_orbitals : Optional[np.ndarray] = None, + full_l_terms : Optional[np.ndarray] = None, + ) -> Dict[str, Any]: """ Forward pass of the atomic DFT solver. @@ -1056,11 +1188,15 @@ def forward(self, orbitals) -> Dict[str, Any]: Parameters ---------- - rho : np.ndarray - Electron density, shape (n_quad_points,) orbitals : np.ndarray Kohn-Sham orbitals (radial wavefunctions R_nl(r)) Shape: (n_states, n_quad_points) + full_eigen_energies : Optional[np.ndarray] + Full eigenvalues of the Kohn-Sham orbitals + full_orbitals : Optional[np.ndarray] + Full orbitals of the Kohn-Sham orbitals + full_l_terms : Optional[np.ndarray] + Full l terms of the Kohn-Sham orbitals Returns ------- @@ -1079,6 +1215,11 @@ def forward(self, orbitals) -> Dict[str, Any]: """ # Phase 1: Get XC functional requirements # Note: Grids and SCF components are already initialized in __init__ + switches = SwitchesFlags( + xc_functional = self.xc_functional, + use_oep = self.use_oep, + hybrid_mixing_parameter = self.hybrid_mixing_parameter, + ) xc_requirements = get_functional_requirements(self.xc_functional) # Phase 2: Calculate rho_nlcc (non-linear core correction for pseudopotentials) @@ -1095,9 +1236,30 @@ def forward(self, orbitals) -> Dict[str, Any]: rho_nlcc = rho_nlcc ) - # Phase 4: Compute XC potential data (using density_data with NLCC) - xc_potential : XCPotentialData = self.energy_calculator.compute_xc_potential(density_data_with_nlcc) + # Phase 4: Compute localized XC potential data (using density_data with NLCC) + n_grid = len(self.grid_data_standard.quadrature_nodes) + v_x = np.zeros(n_grid) + v_c = np.zeros(n_grid) + v_x_oep = np.zeros(n_grid) + v_c_oep = np.zeros(n_grid) + + if switches.use_xc_functional: + v_x, v_c = self.energy_calculator.compute_local_xc_potential(density_data_with_nlcc) + + if switches.use_oep: + v_x_oep, v_c_oep = self.oep_calculator.compute_oep_potentials( + full_eigen_energies = full_eigen_energies, + full_orbitals = full_orbitals, + full_l_terms = full_l_terms, + ) + if switches.hybrid_mixing_parameter is not None: + v_x_local = v_x * (1.0 - switches.hybrid_mixing_parameter) + v_x_oep * switches.hybrid_mixing_parameter + v_c_local = v_c * (1.0 - switches.hybrid_mixing_parameter) + v_c_oep * switches.hybrid_mixing_parameter + else: + v_x_local = v_x + v_x_oep + v_c_local = v_c + v_c_oep + # Phase 5: Create density_data without NLCC for energy calculation # Energy calculation uses valence density only (without NLCC) density_data_valence = self.density_calculator.create_density_data_from_orbitals( @@ -1110,38 +1272,65 @@ def forward(self, orbitals) -> Dict[str, Any]: # Phase 6: Compute final energy (using valence density only) energy_components : EnergyComponents = self.energy_calculator.compute_energy( - orbitals = orbitals, - density_data = density_data_valence, - mixing_parameter = self.hybrid_mixing_parameter + orbitals = orbitals, + density_data = density_data_valence, + mixing_parameter = self.hybrid_mixing_parameter, + full_eigen_energies = full_eigen_energies, + full_orbitals = full_orbitals, + full_l_terms = full_l_terms, + enable_parallelization = self.enable_parallelization, ) - # Phase 7: Evaluate basis functions on uniform grid (optional, for future use) - # TODO: Implement this when needed - # Note: _evaluate_basis_on_uniform_grid expects SCFResult, which we don't have in forward pass - # For now, we'll skip this step as it's not essential for the forward pass - # If needed in the future, we can create a minimal SCFResult-like object or modify the method + if self.print_debug: + energy_components.print_info(title = f"Total Energy ({self.xc_functional})") + print("="*60) + print("\t\t Forward Pass Complete") + print("="*60) + print() + + # Phase 7: Evaluate basis functions on uniform grid uniform_grid, orbitals_on_uniform_grid = self._evaluate_basis_on_uniform_grid( ops_builder_standard = self.ops_builder_standard, orbitals = orbitals ) - # Phase 10: Pack and return results + # evaluate local potentials on uniform grid + v_x_local_on_uniform_grid = self.ops_builder_standard.evaluate_single_orbital_on_given_grid( + given_grid = uniform_grid, + orbital = v_x_local, + ) + v_c_local_on_uniform_grid = self.ops_builder_standard.evaluate_single_orbital_on_given_grid( + given_grid = uniform_grid, + orbital = v_c_local, + ) + + + # Phase 8: Pack and return results final_result = { - 'eigen_energies' : None, # Not computed in forward pass - 'orbitals' : orbitals, - 'rho' : density_data_valence.rho, # Valence density - 'density_data' : density_data_valence, # Density data without NLCC - 'rho_nlcc' : rho_nlcc, - 'energy' : energy_components.total, - 'energy_components' : energy_components, - 'converged' : True, # Forward pass doesn't iterate, so always "converged" - 'iterations' : 0, # No iterations in forward pass - 'rho_residual' : 0.0, # No residual in forward pass - 'grid_data' : self.grid_data_standard, - 'occupation_info' : self.occupation_info, - 'xc_potential' : xc_potential, + 'eigen_energies' : None, # Not computed in forward pass + 'orbitals' : orbitals, # Input orbitals + 'rho' : density_data_valence.rho, # Valence density + 'density_data' : density_data_valence, # Density data without NLCC + 'rho_nlcc' : rho_nlcc, # Non-linear core correction density + 'energy' : energy_components.total, # Total energy + 'energy_components' : energy_components, # Energy components, instance of the classEnergyComponents + 'converged' : None, # Forward pass doesn't iterate, so always "converged" + 'iterations' : None, # No iterations in forward pass + 'rho_residual' : None, # No residual in forward pass + 'grid_data' : self.grid_data_standard, # Standard grid data + 'occupation_info' : self.occupation_info, # Occupation info + 'v_x_local' : v_x_local, # Local XC potential + 'v_c_local' : v_c_local, # Local XC potential + 'uniform_grid' : uniform_grid, # Uniform grid + 'orbitals_on_uniform_grid' : orbitals_on_uniform_grid, # Orbitals on uniform grid + 'v_x_local_on_uniform_grid' : v_x_local_on_uniform_grid, # Local XC potential on uniform grid + 'v_c_local_on_uniform_grid' : v_c_local_on_uniform_grid, # Local XC potential on uniform grid + 'full_eigen_energies' : full_eigen_energies, # Full eigenvalues + 'full_orbitals' : full_orbitals, # Full orbitals + 'full_l_terms' : full_l_terms, # Full l terms } - + + return final_result @@ -1165,7 +1354,7 @@ def solve(self) -> Dict[str, Any]: orbitals_initial = None # Warm start calculation for relatively expensive meta-GGA functionals - if self.xc_functional in ['SCAN', 'RSCAN', 'R2SCAN']: + if self.xc_functional in ['SCAN', 'RSCAN', 'R2SCAN'] or self.use_oep: rho_initial, orbitals_initial = self._get_initial_density_and_orbitals_with_warm_start( xc_functional = "GGA_PBE", rho_initial = rho_initial, @@ -1179,13 +1368,23 @@ def solve(self) -> Dict[str, Any]: ) # Phase 3: Compute final xc potential data - xc_potential : XCPotentialData = self.energy_calculator.compute_xc_potential(scf_result.density_data) + v_x_local, v_c_local = self.energy_calculator.compute_local_xc_potential( + density_data = scf_result.density_data, + full_eigen_energies = scf_result.full_eigen_energies, + full_orbitals = scf_result.full_orbitals, + full_l_terms = scf_result.full_l_terms, + enable_parallelization = self.enable_parallelization, + ) # Phase 4: Compute final energy energy_components : EnergyComponents = self.energy_calculator.compute_energy( - orbitals = scf_result.orbitals, - density_data = scf_result.density_data, - mixing_parameter = self.hybrid_mixing_parameter + orbitals = scf_result.orbitals, + density_data = scf_result.density_data, + mixing_parameter = self.hybrid_mixing_parameter, + full_eigen_energies = scf_result.full_eigen_energies, + full_orbitals = scf_result.full_orbitals, + full_l_terms = scf_result.full_l_terms, + enable_parallelization = self.enable_parallelization, ) # Phase 5: Evaluate basis functions on uniform grid @@ -1194,24 +1393,47 @@ def solve(self) -> Dict[str, Any]: orbitals = scf_result.orbitals ) + # evaluate local potentials on uniform grid + v_x_local_on_uniform_grid = self.ops_builder_standard.evaluate_single_orbital_on_given_grid( + given_grid = uniform_grid, + orbital = v_x_local, + ) + v_c_local_on_uniform_grid = self.ops_builder_standard.evaluate_single_orbital_on_given_grid( + given_grid = uniform_grid, + orbital = v_c_local, + ) + + # Print debug information + if self.print_debug: + energy_components.print_info(title = f"Total Energy ({self.xc_functional})") + print("="*60) + print("\t\t Calculation Complete") + print("="*60) + print() # Phase 6: Pack and return results final_result = { - 'eigen_energies' : scf_result.eigen_energies, - 'orbitals' : scf_result.orbitals, - 'rho' : scf_result.density_data.rho, # Interpolate over psi and calculate rho at that site - 'density_data' : scf_result.density_data, # - 'rho_nlcc' : rho_nlcc, - 'energy' : energy_components.total, - 'energy_components' : energy_components, - 'converged' : scf_result.converged, - 'iterations' : scf_result.iterations, - 'rho_residual' : scf_result.rho_residual, - 'grid_data' : self.grid_data_standard, - 'occupation_info' : self.occupation_info, - 'xc_potential' : xc_potential, - 'uniform_grid' : uniform_grid, - 'orbitals_on_uniform_grid' : orbitals_on_uniform_grid, + 'eigen_energies' : scf_result.eigen_energies, + 'orbitals' : scf_result.orbitals, + 'rho' : scf_result.density_data.rho, # Interpolate over psi and calculate rho at that site + 'density_data' : scf_result.density_data, + 'rho_nlcc' : rho_nlcc, + 'energy' : energy_components.total, + 'energy_components' : energy_components, + 'converged' : scf_result.converged, + 'iterations' : scf_result.iterations, + 'rho_residual' : scf_result.rho_residual, + 'grid_data' : self.grid_data_standard, + 'occupation_info' : self.occupation_info, + 'v_x_local' : v_x_local, + 'v_c_local' : v_c_local, + 'uniform_grid' : uniform_grid, + 'orbitals_on_uniform_grid' : orbitals_on_uniform_grid, + 'v_x_local_on_uniform_grid' : v_x_local_on_uniform_grid, # Local XC potential on uniform grid + 'v_c_local_on_uniform_grid' : v_c_local_on_uniform_grid, # Local XC potential on uniform grid + 'full_eigen_energies' : scf_result.full_eigen_energies, + 'full_orbitals' : scf_result.full_orbitals, + 'full_l_terms' : scf_result.full_l_terms, } @@ -1226,28 +1448,112 @@ def solve(self) -> Dict[str, Any]: # Gauge for energy density (double integral -> different gauges) # Non-locality + + return final_result + + + def solve_raw(self) -> Dict[str, Any]: + """ + Solve the Kohn-Sham equations for the given atomic number. + """ + # 1) Initialize grids + grid_data_standard, grid_data_dense = self._initialize_grids() if self.print_debug: - energy_components.print_info(title = f"Total Energy ({self.xc_functional})") - print("="*60) - print("\t\t Calculation Complete") - print("="*60) + print("=" * 60) + print("\t step 1) Grid initialization completed") + print("=" * 60) + print(" - standard grid nodes.shape : ", grid_data_standard.physical_nodes.shape) + print(" - dense grid nodes.shape : ", grid_data_dense.physical_nodes.shape) + print(" - quadrature_nodes.shape : ", grid_data_standard.quadrature_nodes.shape) + print(" - quadrature_weights.shape : ", grid_data_standard.quadrature_weights.shape) print() - return final_result + + # 2) Operators (Radial FE assembly) + rho_guess = self.pseudo.get_rho_guess(grid_data_standard.quadrature_nodes) + rho_nlcc = self.pseudo.get_rho_core_correction(grid_data_standard.quadrature_nodes) + + # Build operators using grid data + ops_builder = RadialOperatorsBuilder.from_grid_data(grid_data_standard) + ops_dense_builder = RadialOperatorsBuilder.from_grid_data(grid_data_dense) + + # kinetic term + H_kinetic = ops_builder.get_H_kinetic() + + # external potential term + if self.all_electron_flag: + # All-electron: use nuclear Coulomb potential V = -Z/r + V_external = ops_builder.get_nuclear_coulomb_potential(self.pseudo.z_nuclear) + else: + # Pseudopotential: use local pseudopotential component + V_external = self.pseudo.get_v_local_component_psp(grid_data_standard.quadrature_nodes) + + H_ext = ops_builder.build_potential_matrix(V_external) + + # angular momentum term, for solving the Schrödinger equation + H_r_inv_sq = ops_builder.get_H_r_inv_sq() + + # Inverse square root of the overlap matrix + S_inv_sqrt = ops_builder.get_S_inv_sqrt() + + + # dense operators + lagrange_basis_dense = ops_dense_builder.lagrange_basis + lagrange_basis_derivatives_dense = ops_dense_builder.lagrange_basis_derivatives + + + # dense laplacian + laplacian_dense = ops_dense_builder.laplacian + D_dense = ops_dense_builder.derivative_matrix + + if self.xc_functional in ['HF', 'PBE0', 'EXX']: + interpolation_matrix = ops_builder.global_interpolation_matrix + print("interpolation_matrix.shape = ", interpolation_matrix.shape) + np.savetxt("interpolation_matrix.txt", interpolation_matrix.reshape(-1,)) + + raise NotImplementedError("Not implemented") + + + # uniform grid + uniform_eval_grid = np.linspace(0, self.domain_size, + int(self.domain_size / self.mesh_spacing) + 1, + endpoint=True) + + # Evaluate basis functions on uniform grid (with proper padding handling) + lagrange_basis_uniform, uniform_grid_metadata = ops_builder.evaluate_basis_on_uniform_grid( + uniform_grid = uniform_eval_grid + ) + + # Compute non-local pseudopotential matrices (if using pseudopotential) + if not self.all_electron_flag: + # Initialize non-local pseudopotential calculator + nonlocal_calculator = NonLocalPseudopotential( + pseudo=self.pseudo, + ops_builder=ops_builder + ) + + # Compute non-local matrices for all l channels (pre-compute once) + nonlocal_psp_matrices : Dict[int, np.ndarray] = nonlocal_calculator.compute_all_nonlocal_matrices( + l_channels=self.occupation_info.unique_l_values + ) + + raise NotImplementedError("This function is only for testing purposes, should not be used in production code") + + if __name__ == "__main__": atomic_dft_solver = AtomicDFTSolver( atomic_number = 13, - print_debug = True, - xc_functional = "GGA_PBE", + xc_functional = "GGA_PBE", #'rSCAN', 'r2SCAN', "GGA_PBE", "LDA_SPW", "LDA_SVWN", "HFx", "PBE0" all_electron_flag = True, - ) + print_debug = True, + ) - results = atomic_dft_solver.solve() - rho = results['rho'] + results = atomic_dft_solver.solve() + rho = results['rho'] orbitals = results['orbitals'] - print(rho.shape) - print(orbitals.shape) + print("rho.shape = ", rho.shape) # (n_grid_points,) + print("orbitals.shape = ", orbitals.shape) # (n_grid_points, n_orbitals) diff --git a/utils/atom/testcase/__init__.py b/utils/atom/testcase/__init__.py deleted file mode 100644 index a5ff0c84..00000000 --- a/utils/atom/testcase/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Test case module for atomic_dft \ No newline at end of file diff --git a/utils/atom/testcase/mesh_builder_testcase.py b/utils/atom/testcase/mesh_builder_testcase.py deleted file mode 100644 index 59e5c500..00000000 --- a/utils/atom/testcase/mesh_builder_testcase.py +++ /dev/null @@ -1,241 +0,0 @@ - - -import os -import sys -import numpy as np -import matplotlib.pyplot as plt - - -sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) -from atom.mesh.builder import Quadrature1D, Mesh1D, LagrangeShapeFunctions - - - - -def test_function(x): - return 1 / (1 + x ** 2) - - -def print_test_passed(test_name : str): - print("\t {:<30} : test passed".format(test_name)) - - -def test_quadrature1d(n = 95, pnt=False): - # Test Function: Quadrature1D.gauss_legendre - # Test Purpose : Test the Gauss-Legendre quadrature nodes and weights - nodes, legendre_weights = Quadrature1D.gauss_legendre(n) - - integral = np.sum(test_function(nodes) * legendre_weights) - integral_ref = np.pi / 2 - assert np.isclose(integral, integral_ref, atol=1e-6) - print_test_passed("Quadrature1D.gauss_legendre") - - if pnt: - with open("outputs/mesh_builder_testcase/gauss_legendre_quadrature.txt", "w") as f: - f.write("legendre_weights: \n") - f.write(str(legendre_weights)) - f.write("\n") - f.write("nodes: \n") - f.write(str(nodes)) - f.write("\n") - - -def test_lobatto1d(n = 31, pnt=False): - # Test Function: Quadrature1D.lobatto - # Test Purpose : Test the Lobatto quadrature nodes and weights - - nodes, lobatto_weights = Quadrature1D.lobatto(n) - - integral = np.sum(test_function(nodes) * lobatto_weights) - integral_ref = np.pi / 2 - assert np.isclose(integral, integral_ref, atol=1e-6) - print_test_passed("Quadrature1D.lobatto") - - if pnt: - with open("outputs/mesh_builder_testcase/lobatto_quadrature.txt", "w") as f: - f.write("lobatto_weights: \n") - f.write(str(lobatto_weights)) - f.write("\n") - f.write("nodes: \n") - f.write(str(nodes)) - f.write("\n") - - -def test_mesh1d_grid(n = 17): - mesh = Mesh1D( - domain_radius = 10.0, - finite_elements_num = n, - mesh_type = "exponential", - clustering_param = 61.0, - exp_shift = 0.0) - mesh_nodes, mesh_width = mesh.generate_mesh_nodes_and_width() - with open("outputs/mesh_builder_testcase/mesh1d_grid.txt", "w") as f: - f.write("mesh_nodes: \n") - f.write(str(mesh_nodes)) - f.write("\n") - f.write("mesh_width: \n") - f.write(str(mesh_width)) - f.write("\n") - - -def test_mesh1d_fe_nodes(n = 17): - mesh = Mesh1D( - domain_radius = 10.0, - finite_elements_num = n, - mesh_type = "exponential", - clustering_param = 61.0, - exp_shift = 0.0) - mesh_nodes, _ = mesh.generate_mesh_nodes_and_width() - interp_nodes, _ = Quadrature1D.lobatto(31) - fe_nodes = Mesh1D.generate_fe_nodes(mesh_nodes, interp_nodes) - with open("outputs/mesh_builder_testcase/mesh1d_fe_nodes.txt", "w") as f: - f.write("len(fe_nodes): {}".format(len(fe_nodes))) - f.write("\n") - f.write("fe_nodes: \n") - f.write(str(fe_nodes)) - f.write("\n") - - -def test_fe_flat_to_block2d(): - """ - Test Function : fe_flat_to_block2d - Test Purpose : Verify reshaping 1D FE grids into (n_elem, m) blocks under - (1) shared-endpoint layout and (2) stacked layout; - also verify that invalid lengths raise AssertionError. - """ - - def print_pass(name): - print(f"\t {name} : passed") - - # ---------------------------- - # Case 1: endpoints_shared=True - # 3 elements, 4 points per element -> len(flat) = 3*(4-1)+1 = 10 - # rows should be windows with stride (m-1)=3 - # ---------------------------- - flat1 = np.arange(10, dtype=float) - n_elem1 = 3 - out1 = Mesh1D.fe_flat_to_block2d(flat1, n_elem1, endpoints_shared=True) - expected1 = np.array([ - [0, 1, 2, 3], - [3, 4, 5, 6], - [6, 7, 8, 9], - ], dtype=float) - - assert out1.shape == (3, 4) - np.testing.assert_array_equal(out1, expected1) - print_pass("fe_flat_to_block2d (endpoints_shared=True)") - - - # ---------------------------- - # Case 2: endpoints_shared=False - # 3 elements, 4 points per element -> len(flat) = 3*4 = 12 - # direct reshape into (3,4) - # ---------------------------- - flat2 = np.arange(12, dtype=float) - n_elem2 = 3 - out2 = Mesh1D.fe_flat_to_block2d(flat2, n_elem2, endpoints_shared=False) - expected2 = np.array([ - [0, 1, 2, 3], - [4, 5, 6, 7], - [8, 9, 10, 11], - ], dtype=float) - - assert out2.shape == (3, 4) - np.testing.assert_array_equal(out2, expected2) - print_pass("fe_flat_to_block2d (endpoints_shared=False)") - - - # ---------------------------- - # Case 3: invalid lengths should raise AssertionError - # For shared=True, length must be n_elem*(m-1)+1; here we break it - # For shared=False, length must be n_elem*m; here we break it - # ---------------------------- - bad_flat = np.arange(9) # length 9 doesn't fit either (3,4) rules above - try: - _ = Mesh1D.fe_flat_to_block2d(bad_flat, 3, endpoints_shared=True) - raise AssertionError("Expected AssertionError for shared=True did not occur") - except AssertionError: - pass - try: - _ = Mesh1D.fe_flat_to_block2d(bad_flat, 3, endpoints_shared=False) - raise AssertionError("Expected AssertionError for shared=False did not occur") - except AssertionError: - pass - print_pass("fe_flat_to_block2d (invalid lengths raise)") - - print("\t All fe_flat_to_block2d tests passed!") - - - -def test_lagrange_shape_functions_lagrange_basis_and_derivatives(n=17): - """ - Unit test for LagrangeShapeFunctions.lagrange_basis_and_derivatives. - - This function checks: - 1. Shape correctness of returned arrays. - 2. Partition of unity: sum_k L_k(x) == 1. - 3. Derivative consistency: sum_k dLdx_k(x) == 0. - 4. Interpolation property: L_k(x_j) = δ_kj at nodal points. - 5. Smoothness: dLdx finite and behaves as expected. - """ - - # --- Setup: simple node layout and evaluation points --- - nodes = np.array([0.0, 1.0, 2.0, 3.0]) # 4 equally spaced nodes - x_eval = np.linspace(0.0, 3.0, 31) # evaluation points in [0,3] - - # Wrap into element-batch form (1 element) - nodes = nodes[None, :] - x_eval = x_eval[None, :] - - # --- Call the function to test --- - L, dLdx = LagrangeShapeFunctions.lagrange_basis_and_derivatives(nodes, x_eval) - - # --- Check 1: shape correctness --- - assert L.shape == (1, x_eval.shape[1], nodes.shape[1]) - assert dLdx.shape == (1, x_eval.shape[1], nodes.shape[1]) - print("\t 1. Shape check passed:", L.shape, dLdx.shape) - - # --- Check 2: partition of unity --- - unity_error = np.max(np.abs(np.sum(L[0, :, :], axis=1) - 1.0)) - assert unity_error < 1e-12 - print("\t 2. Partition of unity check passed (max error = %.2e)" % unity_error) - - # --- Check 3: derivative consistency (sum of dLdx == 0) --- - deriv_sum_error = np.max(np.abs(np.sum(dLdx[0, :, :], axis=1))) - assert deriv_sum_error < 1e-10 - print("\t 3. Derivative consistency check passed (max error = %.2e)" % deriv_sum_error) - - # --- Check 4: interpolation property (nodal identity) --- - # Evaluate basis at nodal points - L_nodes, dLdx_nodes = LagrangeShapeFunctions.lagrange_basis_and_derivatives(nodes, nodes) - identity_error = np.max(np.abs(L_nodes[0, :, :] - np.eye(nodes.shape[1]))) - assert identity_error < 1e-12 - print("\t 4. Interpolation property check passed (max error = %.2e)" % identity_error) - - # --- Check 5: finite values --- - assert np.all(np.isfinite(L)) - assert np.all(np.isfinite(dLdx)) - print("\t 5. Finite value check passed") - - print("\t All Lagrange basis tests passed!") - - - - - - -if __name__ == "__main__": - if not os.path.exists("outputs/mesh_builder_testcase"): - os.makedirs("outputs/mesh_builder_testcase") - - print("Running tests...") - # test_quadrature1d(n = 95) - # test_lobatto1d(n = 31) - # test_mesh1d_grid(n = 17) - # test_mesh1d_fe_nodes(n = 17) - test_fe_flat_to_block2d() - # test_lagrange_shape_functions_lagrange_basis_and_derivatives(n = 17) - print("All tests passed") - - # Terminal command: - # python atomic_dft/testcase/mesh_builder_testcase.py \ No newline at end of file diff --git a/utils/atom/xc/evaluator.py b/utils/atom/xc/evaluator.py index bae7574f..d7cf5e57 100644 --- a/utils/atom/xc/evaluator.py +++ b/utils/atom/xc/evaluator.py @@ -671,11 +671,10 @@ def create_xc_evaluator( # hybrid functionals 'PBE0': GGA_PBE, - # 'HF': HF, - - # TODO: Add OEP and RPA - # 'OEPx': OEPx, - # 'RPA': RPA, + + # OEP functionals + 'EXX': GGA_PBE, + 'RPA': GGA_PBE, } if functional_name not in FUNCTIONAL_MAP: diff --git a/utils/atom/xc/functional_requirements.py b/utils/atom/xc/functional_requirements.py index c30aa375..984bd7e5 100644 --- a/utils/atom/xc/functional_requirements.py +++ b/utils/atom/xc/functional_requirements.py @@ -110,10 +110,10 @@ def is_hybrid(self) -> bool: ), # Optimized Effective Potential - 'OEPx': FunctionalRequirements( + 'EXX': FunctionalRequirements( needs_gradient=True, needs_tau=False, - functional_type='OEP', + functional_type='hybrid-GGA', needs_orbitals=True ), diff --git a/utils/atom/xc/hybrid.py b/utils/atom/xc/hybrid.py index ea655158..7050113c 100644 --- a/utils/atom/xc/hybrid.py +++ b/utils/atom/xc/hybrid.py @@ -9,13 +9,16 @@ from __future__ import annotations import numpy as np -from typing import Dict, Optional, Tuple, TYPE_CHECKING +from typing import Dict, Optional, Tuple, List, TYPE_CHECKING if TYPE_CHECKING: from ..utils.occupation_states import OccupationInfo from ..mesh.operators import RadialOperatorsBuilder # Error messages +FACTORIAL_N_MUST_BE_NON_NEGATIVE_INTEGER_ERROR = \ + "parameter n must be a non-negative integer, get {} instead" + L_VALUES_MUST_BE_INTEGERS_ERROR = \ "parameter l_values in class OccupationInfo must be integers, get type {} instead" @@ -24,11 +27,12 @@ ORBITALS_MUST_BE_A_2D_NUMPY_ARRAY_ERROR = \ "parameter orbitals must be a 2D numpy array, get dimension {} instead" ORBITALS_MUST_HAVE_N_GRID_ROWS_ERROR = \ - "parameter orbitals must have n_grid rows, get {} instead" + "parameter orbitals must have n_grid rows, get {} rows instead of {} rows" ORBITALS_MUST_HAVE_N_ORBITALS_COLUMNS_ERROR = \ - "parameter orbitals must have n_orbitals columns, get {} instead" - + "parameter orbitals must have n_orbitals columns, get {} columns instead of {} columns" +EXCHANGE_POTENTIAL_OUTPUT_SHAPE_ERROR = \ + "exchange potential must have shape (n_orbitals, n_grid), get shape {} instead" def factorial(n: int) -> int: """ @@ -38,7 +42,8 @@ def factorial(n: int) -> int: Uses lookup table for common values to avoid repeated computation. """ - assert n >= 0 and isinstance(n, int) + assert n >= 0 and isinstance(n, int), \ + FACTORIAL_N_MUST_BE_NON_NEGATIVE_INTEGER_ERROR.format(n) if n == 0: return 1 elif n == 1: return 1 @@ -57,44 +62,53 @@ def factorial(n: int) -> int: return result -def _wigner_3j_000(l1: int, l2: int, L: int) -> float: - """ - Wigner 3j symbol (l1 l2 L; 0 0 0) with built-in selection rules. - """ - J = l1 + l2 + L - # parity: l1 + l2 + L must be even - if (J & 1) == 1: - return 0.0 - # triangle inequalities - if l1 < abs(l2 - L) or l1 > l2 + L: - return 0.0 - if l2 < abs(l1 - L) or l2 > l1 + L: - return 0.0 - if L < abs(l1 - l2) or L > l1 + l2: - return 0.0 - - g = J // 2 - W = (-1)**g - W *= np.sqrt( - factorial(J - 2*l1) * factorial(J - 2*l2) * factorial(J - 2*L) - / factorial(J + 1) - ) - W *= factorial(g) / (factorial(g - l1) * factorial(g - l2) * factorial(g - L)) - return float(W) - - -def _radial_kernel(l: int, r_nodes: np.ndarray, r_weights: np.ndarray) -> np.ndarray: - """ - Compute kernel K^(l) with entries: - K_ij^(l) = [ r_<^l / r_>^(l+1) ] * (w_i w_j) / (2l + 1), - where r_< = min(r_i, r_j), r_> = max(r_i, r_j). - This term represents the radial part of the spherical harmonic expansion of the Coulomb interaction. - """ - r_min = np.minimum(r_nodes, r_nodes.reshape(-1, 1)) - r_max = np.maximum(r_nodes, r_nodes.reshape(-1, 1)) - - return ((r_min / r_max)**l / r_max) * (r_weights * r_weights.reshape(-1, 1)) / (2*l + 1) +class CoulombCouplingCalculator: + + @staticmethod + def radial_kernel(l: int, r_nodes: np.ndarray, r_weights: np.ndarray) -> np.ndarray: + """ + Compute kernel K^(l) with entries: + K_ij^(l) = [ r_<^l / r_>^(l+1) ] * (w_i w_j) / (2l + 1), + where r_< = min(r_i, r_j), r_> = max(r_i, r_j). + + This term represents the radial part of the spherical harmonic expansion of the Coulomb interaction. + """ + r_min = np.minimum(r_nodes, r_nodes.reshape(-1, 1)) + r_max = np.maximum(r_nodes, r_nodes.reshape(-1, 1)) + + return ((r_min / r_max)**l / r_max) * (r_weights * r_weights.reshape(-1, 1)) / (2*l + 1) + + + @staticmethod + def wigner_3j_000(l1: int, l2: int, L: int) -> float: + """ + Wigner 3j symbol (l1 l2 L; 0 0 0) with built-in selection rules. + """ + assert isinstance(l1, int), "l1 must be an integer, get type {} instead".format(type(l1)) + assert isinstance(l2, int), "l2 must be an integer, get type {} instead".format(type(l2)) + assert isinstance(L, int) , "L must be an integer, get type {} instead".format(type(L)) + J = l1 + l2 + L + # parity: l1 + l2 + L must be even + if (J & 1) == 1: + return 0.0 + # triangle inequalities + if l1 < abs(l2 - L) or l1 > l2 + L: + return 0.0 + if l2 < abs(l1 - L) or l2 > l1 + L: + return 0.0 + if L < abs(l1 - l2) or L > l1 + l2: + return 0.0 + + g = J // 2 + W = (-1)**g + W *= np.sqrt( + factorial(J - 2*l1) * factorial(J - 2*l2) * factorial(J - 2*L) + / factorial(J + 1) + ) + W *= factorial(g) / (factorial(g - l1) * factorial(g - l2) * factorial(g - L)) + return float(W) + class HartreeFockExchange: @@ -109,7 +123,7 @@ def __init__( self, ops_builder : 'RadialOperatorsBuilder', occupation_info: 'OccupationInfo' - ): + ): """ Initialize HF exchange calculator. @@ -129,14 +143,14 @@ def __init__( self.n_grid = len(self.quadrature_nodes) # Extract occupation data - self.l_values = occupation_info.l_values + self.l_values = occupation_info.l_values self.occupations = occupation_info.occupations - self.n_orbitals = len(self.l_values) + self.n_orbitals = len(self.l_values) assert self.l_values.dtype == int, \ L_VALUES_MUST_BE_INTEGERS_ERROR.format(self.l_values.dtype) - + def _compute_exchange_matrix( self, @@ -159,11 +173,11 @@ def _compute_exchange_matrix( for l_prime in l_coupling: # Angular part: compute vectorized alpha without in-place modification w3j_values = np.array([ - _wigner_3j_000(int(l_value), int(lj), int(l_prime)) for lj in self.l_values + CoulombCouplingCalculator.wigner_3j_000(int(l_value), int(lj), int(l_prime)) for lj in self.l_values ], dtype=float) # Radial part: compute the radial coupling kernel K^(L) - radial_kernel = _radial_kernel( + radial_kernel = CoulombCouplingCalculator.radial_kernel( int(l_prime), self.quadrature_nodes, self.quadrature_weights ) @@ -189,6 +203,67 @@ def _compute_exchange_matrix( return - H_hf_exchange_matrix + def compute_exchange_potentials( + self, + orbitals): + """ + Compute Hartree-Fock exchange potentials for all angular momentum channels. + + This function is useful for the OEP calculation. Here, the input orbitals should only contain the occupied orbitals. + + Parameters + ---------- + orbitals : np.ndarray + Kohn-Sham orbitals (radial wavefunctions) at quadrature points + Shape: (n_grid, n_orbitals) + + Returns + ------- + np.ndarray + Hartree-Fock exchange potential for all angular momentum channels + Shape: (len(l_values), n_grid) + """ + # Check Type and shape + assert isinstance(orbitals, np.ndarray), \ + ORBITALS_MUST_BE_A_NUMPY_ARRAY_ERROR.format(type(orbitals)) + assert orbitals.ndim == 2, \ + ORBITALS_MUST_BE_A_2D_NUMPY_ARRAY_ERROR.format(orbitals.ndim) + assert orbitals.shape[0] == self.n_grid, \ + ORBITALS_MUST_HAVE_N_GRID_ROWS_ERROR.format(self.n_grid, orbitals.shape[0]) + assert orbitals.shape[1] == self.n_orbitals, \ + ORBITALS_MUST_HAVE_N_ORBITALS_COLUMNS_ERROR.format(self.n_orbitals, orbitals.shape[1]) + + # Compute HF exchange matrices for all l channels + l_coupling = np.arange(0, 2 * np.max(self.l_values) + 1) + + # Compute exchange potential for each l channel + exchange_potential_l_contribution_list : List[np.ndarray] = [] + for l_prime in l_coupling: + _wigner_term = np.array([ + CoulombCouplingCalculator.wigner_3j_000(int(l1), int(l2), int(l_prime))**2 for l1 in self.l_values for l2 in self.l_values + ], dtype=float).reshape(len(self.l_values), len(self.l_values)) + + _exchange_potential_l_contribution = -0.5 * np.einsum( + 'ki,ji,ikj->kj', + _wigner_term * self.occupations, + orbitals, + np.einsum('li,lk,jl->ikj', + orbitals, + orbitals, + CoulombCouplingCalculator.radial_kernel(l_prime, self.quadrature_nodes, self.quadrature_weights) * (2 * l_prime + 1), + ), + optimize=True, + ) + + exchange_potential_l_contribution_list.append(_exchange_potential_l_contribution) + + exchange_potential = np.sum(exchange_potential_l_contribution_list, axis=0) + assert exchange_potential.shape == (self.n_orbitals, self.n_grid), \ + EXCHANGE_POTENTIAL_OUTPUT_SHAPE_ERROR.format(exchange_potential.shape, self.n_orbitals, self.n_grid) + + return exchange_potential + + def compute_exchange_matrices_dict( self, orbitals: np.ndarray @@ -277,13 +352,13 @@ def compute_exchange_energy( wigner_matrix = np.zeros((len(l_values), len(l_values))) for i1 in range(len(l_values)): for i2 in range(len(l_values)): - wigner_matrix[i1, i2] = _wigner_3j_000(int(l_values[i1]), int(l_values[i2]), int(l_coupling))**2 + wigner_matrix[i1, i2] = CoulombCouplingCalculator.wigner_3j_000(int(l_values[i1]), int(l_values[i2]), int(l_coupling))**2 # Create occupation matrix occ_matrix = occupations * occupations.reshape(-1, 1) # Compute radial kernel for this l coupling - r_kernel = _radial_kernel(l_coupling, self.quadrature_nodes, self.quadrature_weights) + r_kernel = CoulombCouplingCalculator.radial_kernel(l_coupling, self.quadrature_nodes, self.quadrature_weights) # Compute exchange energy contribution for this l coupling # This is the complex einsum from the reference code: @@ -308,3 +383,4 @@ def compute_exchange_energy( return E_HF + diff --git a/utils/atom/xc/oep.py b/utils/atom/xc/oep.py index e1a8cc30..3f3cca8a 100644 --- a/utils/atom/xc/oep.py +++ b/utils/atom/xc/oep.py @@ -1,12 +1,659 @@ -from __future__ import annotations -import numpy as np -from typing import Tuple, Any - -class OEPBuilder: - """Build OEP exchange/correlation potentials from eigenstates. Placeholder.""" - def __init__(self, lamda: float = 1.0): - self.lamda = float(lamda) - - def build(self, eigvals: np.ndarray, eigvecs: np.ndarray, meta: dict[str, Any]) -> Tuple[np.ndarray, np.ndarray]: - """Return (Vx_OEP, Vc_OEP) on quadrature grid.""" - raise NotImplementedError("OEP builder not implemented yet.") +from __future__ import annotations + +import scipy +from scipy.linalg import LinAlgWarning +import numpy as np + +import warnings +from typing import Tuple, List, Optional + +from .hybrid import CoulombCouplingCalculator, HartreeFockExchange +from .rpa import RPACorrelation +from ..utils.occupation_states import OccupationInfo +from ..mesh.operators import RadialOperatorsBuilder + + +from contextlib import nullcontext + +try: + # Optional dependency: used to limit BLAS/OpenMP threads during parallel sections + from threadpoolctl import threadpool_limits +except ImportError: + threadpool_limits = None # type: ignore + + + +# Error messages +USE_RPA_CORRELATION_NOT_BOOL_ERROR = \ + "Parameter use_rpa_correlation must be a bool, get type {} instead" +FREQUENCY_QUADRATURE_POINT_NUMBER_NOT_NONE_ERROR = \ + "Parameter frequency_quadrature_point_number must be not None, get None instead" +FREQUENCY_QUADRATURE_POINT_NUMBER_NOT_INTEGER_ERROR = \ + "Parameter frequency_quadrature_point_number must be an integer, get type {} instead" +FREQUENCY_QUADRATURE_POINT_NUMBER_NOT_GREATER_THAN_0_ERROR = \ + "Parameter frequency_quadrature_point_number must be greater than 0, get {} instead" +ANGULAR_MOMENTUM_CUTOFF_NOT_NONE_ERROR = \ + "Parameter angular_momentum_cutoff must be not None, get None instead" +ANGULAR_MOMENTUM_CUTOFF_NOT_INTEGER_ERROR = \ + "Parameter angular_momentum_cutoff must be an integer, get type {} instead" +ANGULAR_MOMENTUM_CUTOFF_NEGATIVE_ERROR = \ + "Parameter angular_momentum_cutoff must be non-negative, get {} instead" +ANGULAR_MOMENTUM_CUTOFF_NOT_NONE_ERROR_MESSAGE = \ + "Parameter angular_momentum_cutoff must be not None, get None instead" +OPS_BUILDER_NOT_RADIAL_OPERATORS_BUILDER_ERROR = \ + "Parameter ops_builder must be a RadialOperatorsBuilder instance, get type {} instead" +OPS_BUILDER_OEP_NOT_RADIAL_OPERATORS_BUILDER_ERROR = \ + "Parameter ops_builder_oep must be a RadialOperatorsBuilder instance, get type {} instead" +OPS_BUILDERS_NOT_CONSISTENT_AT_QUADRATURE_NODES_ERROR = \ + "Parameter ops_builder.quadrature_nodes must be equal to ops_builder_oep.quadrature_nodes, please check the grid data and the operators builder" +OPS_BUILDERS_NOT_CONSISTENT_AT_QUADRATURE_WEIGHTS_ERROR = \ + "Parameter ops_builder.quadrature_weights must be equal to ops_builder_oep.quadrature_weights, please check the grid data and the operators builder" +OCCUPATION_INFO_NOT_OCCUPATION_INFO_ERROR = \ + "Parameter occupation_info must be a OccupationInfo instance, get type {} instead" +OCCUPATION_INFO_L_TERMS_NOT_CONSISTENT_WITH_OCCUPATION_INFO_ERROR = \ + "Occupied l terms are not consistent with the occupation information, please check your inputs, got {} instead of {}" +FULL_EIGEN_ENERGIES_NOT_NUMPY_ARRAY_ERROR = \ + "Parameter full_eigen_energies must be a numpy array, get type {} instead" +FULL_ORBITALS_NOT_NUMPY_ARRAY_ERROR = \ + "Parameter full_orbitals must be a numpy array, get type {} instead" +FULL_L_TERMS_NOT_NUMPY_ARRAY_ERROR = \ + "Parameter full_l_terms must be a numpy array, get type {} instead" +FULL_EIGEN_ENERGIES_NOT_1D_ARRAY_ERROR = \ + "Parameter full_eigen_energies must be a 1D array, got ndim={}" +FULL_ORBITALS_NOT_2D_ARRAY_ERROR = \ + "Parameter full_orbitals must be a 2D array, got ndim={}" +FULL_L_TERMS_NOT_1D_ARRAY_ERROR = \ + "Parameter full_l_terms must be a 1D array, got ndim={}" +FULL_EIGEN_ENERGIES_AND_ORBITALS_SHAPE_ERROR = \ + "Parameter full_eigen_energies.shape[0] must equal full_orbitals.shape[1], got {} and {} instead" +FULL_EIGEN_ENERGIES_AND_L_TERMS_SHAPE_ERROR = \ + "Parameter full_eigen_energies.shape[0] must equal full_l_terms.shape[0], got {} and {} instead" +FULL_ORBITALS_AND_L_TERMS_SHAPE_ERROR = \ + "Parameter full_orbitals.shape[1] must equal full_l_terms.shape[0], got {} and {} instead" +FULL_L_TERMS_NON_NEGATIVE_ERROR = \ + "Parameter full_l_terms must be non-negative, got {} instead of all non-negative values" +INVALID_EX_TAG_ERROR = \ + "Parameter `ex_tag` must be 'exchange' or 'correlation', got {} instead" +PARENT_CLASS_RPACORRELATION_NOT_INITIALIZED_ERROR = \ + "Parent class RPACorrelation is not initialized, check your initialization at OEPCalculator class" +ENABLE_PARALLELIZATION_NOT_BOOL_ERROR = \ + "Parameter `parallelize` must be a bool, get type {} instead" + +# WARNING Messages +FREQUENCY_QUADRATURE_POINT_NUMBER_NOT_NONE_WHEN_RPA_CORRELATION_IS_NOT_USED_WARNING = \ + "WARNING: parameter `frequency_quadrature_point_number` is not None when RPA correlation is not used, so it will be ignored" +ANGULAR_MOMENTUM_CUTOFF_NOT_NONE_WHEN_RPA_CORRELATION_IS_NOT_USED_WARNING = \ + "WARNING: parameter `angular_momentum_cutoff` is not None when RPA correlation is not used, so it will be ignored" + + +class OEPCalculator(HartreeFockExchange, RPACorrelation): + + """Prepare and build OEP exchange/correlation potentials from eigenstates.""" + + def __init__( + self, + ops_builder : RadialOperatorsBuilder, + ops_builder_oep : RadialOperatorsBuilder, + occupation_info : OccupationInfo, + use_rpa_correlation : bool, + frequency_quadrature_point_number : Optional[int] = None, # parameters for RPA correlation potential + angular_momentum_cutoff : Optional[int] = None, # parameters for RPA functional only + ): + + """ + Parameters + ---------- + ops_builder : RadialOperatorsBuilder + Operators builder for the standard grid + ops_builder_oep : RadialOperatorsBuilder + Operators builder for the OEP grid + occupation_info : OccupationInfo + Occupation information + use_rpa_correlation : bool + Whether to use RPA correlation potential + If True, use RPA correlation potential, otherwise return zero correlation potential + """ + assert isinstance(ops_builder, RadialOperatorsBuilder), \ + OPS_BUILDER_NOT_RADIAL_OPERATORS_BUILDER_ERROR.format(type(ops_builder)) + assert isinstance(ops_builder_oep, RadialOperatorsBuilder), \ + OPS_BUILDER_OEP_NOT_RADIAL_OPERATORS_BUILDER_ERROR.format(type(ops_builder_oep)) + assert isinstance(occupation_info, OccupationInfo), \ + OCCUPATION_INFO_NOT_OCCUPATION_INFO_ERROR.format(type(occupation_info)) + assert isinstance(use_rpa_correlation, bool), \ + USE_RPA_CORRELATION_NOT_BOOL_ERROR.format(type(use_rpa_correlation)) + + # check if the two ops_builders are consistent at quadrature nodes + self._check_ops_builder_consistency_at_quadrature_nodes(ops_builder, ops_builder_oep) + + # initialize the Hartree-Fock exchange class + HartreeFockExchange.__init__( + self, + ops_builder = ops_builder, + occupation_info = occupation_info, + ) + + self.ops_builder_oep = ops_builder_oep + self.physical_nodes = ops_builder.physical_nodes + + # Some dimension information + self.n_quad : int = len(self.quadrature_nodes) + self.n_interior : int = len(self.physical_nodes) - 2 + + # Occupation information + self.occupations : np.ndarray = self.occupation_info.occupations + self.occ_l_values : np.ndarray = self.occupation_info.l_values + self.occ_n_values : np.ndarray = self.occupation_info.n_values + + # Ill_conditioned warning + self.ill_conditioned_warning_caught_times_for_exchange : int = 0 + self.ill_conditioned_warning_caught_times_for_correlation : int = 0 + self.rcond_list_for_exchange : List[float] = [] + self.rcond_list_for_correlation : List[float] = [] + + # Parameters for RPA correlation potential + self.use_rpa_correlation = use_rpa_correlation + + if use_rpa_correlation: + # check frequency quadrature point number + assert frequency_quadrature_point_number is not None, \ + FREQUENCY_QUADRATURE_POINT_NUMBER_NOT_NONE_ERROR.format(frequency_quadrature_point_number) + assert isinstance(frequency_quadrature_point_number, int), \ + FREQUENCY_QUADRATURE_POINT_NUMBER_NOT_INTEGER_ERROR.format(type(frequency_quadrature_point_number)) + assert frequency_quadrature_point_number > 0, \ + FREQUENCY_QUADRATURE_POINT_NUMBER_NOT_GREATER_THAN_0_ERROR.format(frequency_quadrature_point_number) + + # check angular momentum cutoff + assert angular_momentum_cutoff is not None, \ + ANGULAR_MOMENTUM_CUTOFF_NOT_NONE_ERROR.format(angular_momentum_cutoff) + assert isinstance(angular_momentum_cutoff, int), \ + ANGULAR_MOMENTUM_CUTOFF_NOT_INTEGER_ERROR.format(type(angular_momentum_cutoff)) + assert angular_momentum_cutoff >= 0, \ + ANGULAR_MOMENTUM_CUTOFF_NEGATIVE_ERROR.format(angular_momentum_cutoff) + + # initialize the RPA correlation class + RPACorrelation.__init__( + self, + ops_builder = ops_builder, + occupation_info = occupation_info, + frequency_quadrature_point_number = frequency_quadrature_point_number, + angular_momentum_cutoff = angular_momentum_cutoff, + ) + else: + if frequency_quadrature_point_number is not None: + print(FREQUENCY_QUADRATURE_POINT_NUMBER_NOT_NONE_WHEN_RPA_CORRELATION_IS_NOT_USED_WARNING) + if angular_momentum_cutoff is not None: + print(ANGULAR_MOMENTUM_CUTOFF_NOT_NONE_WHEN_RPA_CORRELATION_IS_NOT_USED_WARNING) + + + def _check_ops_builder_consistency_at_quadrature_nodes( + self, + ops_builder : RadialOperatorsBuilder, + ops_builder_oep: RadialOperatorsBuilder + ) -> None: + assert np.allclose(ops_builder.quadrature_nodes, ops_builder_oep.quadrature_nodes), \ + OPS_BUILDERS_NOT_CONSISTENT_AT_QUADRATURE_NODES_ERROR + assert np.allclose(ops_builder.quadrature_weights, ops_builder_oep.quadrature_weights), \ + OPS_BUILDERS_NOT_CONSISTENT_AT_QUADRATURE_WEIGHTS_ERROR + + # skip checking for now, will implement other consistency checks later (if needed) + pass + + + def reset(self, ): + """ + Reset the OEP calculator + """ + # reset the warning caught time + self.ill_conditioned_warning_caught_times_for_exchange = 0 + self.ill_conditioned_warning_caught_times_for_correlation = 0 + # Reset the rcond list for exchange and correlation + self.rcond_list_for_exchange.clear() + self.rcond_list_for_correlation.clear() + + + def compute_oep_potentials( + self, + full_eigen_energies : np.ndarray, + full_orbitals : np.ndarray, + full_l_terms : np.ndarray, + enable_parallelization : Optional[bool] = None, + ) -> Tuple[np.ndarray, np.ndarray]: + """ + Compute OEP potentials from full orbitals and eigenvalues. + + Parameters + ---------- + full_eigen_energies : np.ndarray + Full eigenvalues of the system, shape (n_total_orbitals,) + full_orbitals : np.ndarray + Full orbitals of the system, shape (n_grid, n_total_orbitals) + full_l_terms : np.ndarray + Specify the l value of each orbital, shape (n_total_orbitals,) + enable_parallelization : bool + Flag for parallelization of RPA calculations + + Returns + ------- + v_x_oep : np.ndarray + OEP exchange potential, shape (n_grid,) + v_c_oep : np.ndarray + OEP correlation potential, shape (n_grid,) + """ + + # Type check for required fields + self._validate_full_spectrum_inputs(full_eigen_energies, full_orbitals, full_l_terms) + + # Get occupation information + occ_orbitals = full_orbitals[:, :len(self.occ_l_values)] + + # get the global interpolation matrix + global_interpolation_matrix = self.ops_builder_oep.global_interpolation_matrix + + + ### =========================================== ### + ### Part 1: Compute OEP exchange potential ### + ### =========================================== ### + + # Compute exact exchange potentials + exact_exchange_potentials = self.compute_exchange_potentials(occ_orbitals) + + # Compute OEP exchange kernel and the exchange driving term + chi_0_kernel, exchange_driving_term = \ + self._compute_oep_kernel_and_exchange_driving_term( + full_eigen_energies = full_eigen_energies, + full_orbitals = full_orbitals, + full_l_terms = full_l_terms, + exchange_potentials = exact_exchange_potentials + ) + + + # Convert chi_0_kernel to sparser grid, + # Note: this matrix is shared while computing the RPA correlation potential + chi_0_kernel_sparser_grid = \ + self.convert_chi_0_kernel_to_sparser_grid( + chi_0_kernel = chi_0_kernel, + quadrature_weights = self.quadrature_weights, + global_interpolation_matrix = global_interpolation_matrix, + ) + + + # Convert exchange_driving_term to sparser grid + exchange_driving_term_sparser_grid = \ + self.convert_driving_term_to_sparser_grid( + driving_term = exchange_driving_term, + quadrature_weights = self.quadrature_weights, + global_interpolation_matrix = global_interpolation_matrix, + ) + + # solve for the OEP coefficient + oep_coefficient = self.solve_oep_coefficients( + chi_0_kernel = chi_0_kernel_sparser_grid, + driving_term = exchange_driving_term_sparser_grid, + ex_tag = 'exchange' + ) + + # compute the OEP exchange potential + v_x_oep = global_interpolation_matrix @ oep_coefficient + + + # zero point shift the OEP exchange potential + energy_correction = self._compute_zero_point_shift_correction_term( + quadrature_weights = self.quadrature_weights, + orbital_homo = occ_orbitals[:, -1], + exx_potential_homo = exact_exchange_potentials[-1, :], + oep_potential = v_x_oep, + ) + + # Add the zero point shift correction term to the OEP exchange potential + v_x_oep += energy_correction + + + ### =========================================== ### + ### Part 2: Compute OEP correlation potential ### + ### =========================================== ### + + # compute RPA correlation potential, if needed + if not self.use_rpa_correlation: + v_c_oep = np.zeros_like(v_x_oep) + else: + # Compute RPA correlation driving term + rpa_correlation_driving_term = self._compute_rpa_correlation_driving_term( + full_eigen_energies = full_eigen_energies, + full_orbitals = full_orbitals, + full_l_terms = full_l_terms, + enable_parallelization = enable_parallelization, + ) + + # Convert RPA correlation driving term to sparser grid + rpa_correlation_driving_term_sparser_grid = \ + self.convert_driving_term_to_sparser_grid( + driving_term = rpa_correlation_driving_term, + quadrature_weights = self.quadrature_weights, + global_interpolation_matrix = global_interpolation_matrix + ) + + # Solve for the RPA correlation coefficient + rpa_correlation_coefficient = self.solve_oep_coefficients( + chi_0_kernel = chi_0_kernel_sparser_grid, + driving_term = rpa_correlation_driving_term_sparser_grid, + ex_tag = 'correlation' + ) + + # shift the RPA correlation coefficient by the HOMO coefficient + rpa_correlation_coefficient -= rpa_correlation_coefficient[-1] + + # Compute the RPA correlation potential + v_c_oep = global_interpolation_matrix @ rpa_correlation_coefficient + + return v_x_oep, v_c_oep + + + def solve_oep_coefficients( + self, + chi_0_kernel : np.ndarray, + driving_term : np.ndarray, + ex_tag : str, + ) -> np.ndarray: + """ + Solve the OEP coefficients for exchange or correlation potential + + Parameters + ---------- + chi_0_kernel : np.ndarray + Chi_0 kernel, shape (n_quad, n_quad) + driving_term : np.ndarray + Driving term, shape (n_quad,) + ex_tag : str + Tag for the potential, must be 'exchange' or 'correlation' + This parameter is used to record the warning caught times and rcond_list for exchange or + correlation potential while solving the OEP coefficients. + + Returns + ------- + oep_coefficient : np.ndarray + OEP coefficients, shape (n_quad,) + """ + assert ex_tag in ['exchange', 'correlation'], \ + INVALID_EX_TAG_ERROR.format(ex_tag) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", LinAlgWarning) + oep_coefficient = scipy.linalg.solve(chi_0_kernel, driving_term) + + if caught: + cond = np.linalg.cond(chi_0_kernel) + rcond = 1.0 / cond if cond != 0 else np.inf + if ex_tag == 'exchange': + self.ill_conditioned_warning_caught_times_for_exchange += 1 + self.rcond_list_for_exchange.append(rcond) + elif ex_tag == 'correlation': + self.ill_conditioned_warning_caught_times_for_correlation += 1 + self.rcond_list_for_correlation.append(rcond) + + return oep_coefficient + + + @staticmethod + def _validate_full_spectrum_inputs( + full_eigen_energies : np.ndarray, + full_orbitals : np.ndarray, + full_l_terms : np.ndarray + ) -> None: + """ + Validate the inputs of full spectrum + """ + assert isinstance(full_eigen_energies, np.ndarray), \ + FULL_EIGEN_ENERGIES_NOT_NUMPY_ARRAY_ERROR.format(type(full_eigen_energies)) + assert isinstance(full_orbitals, np.ndarray), \ + FULL_ORBITALS_NOT_NUMPY_ARRAY_ERROR.format(type(full_orbitals)) + assert isinstance(full_l_terms, np.ndarray), \ + FULL_L_TERMS_NOT_NUMPY_ARRAY_ERROR.format(type(full_l_terms)) + assert full_eigen_energies.ndim == 1, \ + FULL_EIGEN_ENERGIES_NOT_1D_ARRAY_ERROR.format(full_eigen_energies.ndim) + assert full_orbitals.ndim == 2, \ + FULL_ORBITALS_NOT_2D_ARRAY_ERROR.format(full_orbitals.ndim) + assert full_l_terms.ndim == 1, \ + FULL_L_TERMS_NOT_1D_ARRAY_ERROR.format(full_l_terms.ndim) + assert full_eigen_energies.shape[0] == full_orbitals.shape[1], \ + FULL_EIGEN_ENERGIES_AND_ORBITALS_SHAPE_ERROR.format(full_eigen_energies.shape[0], full_orbitals.shape[1]) + assert full_eigen_energies.shape[0] == full_l_terms.shape[0], \ + FULL_EIGEN_ENERGIES_AND_L_TERMS_SHAPE_ERROR.format(full_eigen_energies.shape[0], full_l_terms.shape[0]) + assert full_orbitals.shape[1] == full_l_terms.shape[0], \ + FULL_ORBITALS_AND_L_TERMS_SHAPE_ERROR.format(full_orbitals.shape[1], full_l_terms.shape[0]) + assert np.all(full_l_terms >= 0), \ + FULL_L_TERMS_NON_NEGATIVE_ERROR.format(full_l_terms) + + + @staticmethod + def _compute_zero_point_shift_correction_term( + quadrature_weights : np.ndarray, + orbital_homo : np.ndarray, + exx_potential_homo : np.ndarray, + oep_potential : np.ndarray, + ) -> float: + """ + Compute the zero point shift correction term for the OEP exchange potential + """ + homo_oep_projection_energy = np.sum(oep_potential * (orbital_homo ** 2) * quadrature_weights) + homo_exact_exchange_energy = np.sum(exx_potential_homo * orbital_homo) + energy_correction = -homo_oep_projection_energy + homo_exact_exchange_energy + return energy_correction + + + @staticmethod + def convert_chi_0_kernel_to_sparser_grid( + chi_0_kernel : np.ndarray, + quadrature_weights : np.ndarray, + global_interpolation_matrix : np.ndarray, + ) -> np.ndarray: + """ + Convert chi_0_kernel and driving_term to sparser grid + """ + # convert chi_0_kernel and exchange_driving_term to sparser grid + chi_0_kernel_sparser_grid = np.einsum('i,ij,il,lk,l->jk', + quadrature_weights, + global_interpolation_matrix, + chi_0_kernel, + global_interpolation_matrix, + quadrature_weights, + optimize=True + ) + + return chi_0_kernel_sparser_grid + + + @staticmethod + def convert_driving_term_to_sparser_grid( + driving_term : np.ndarray, + quadrature_weights : np.ndarray, + global_interpolation_matrix : np.ndarray, + ) -> np.ndarray: + """ + Convert driving term to sparser grid + """ + driving_term_sparser_grid = np.einsum('i, ij, i->j', + quadrature_weights, + global_interpolation_matrix, + driving_term, + optimize=True + ) + return driving_term_sparser_grid + + + def _compute_oep_kernel_and_exchange_driving_term( + self, + full_eigen_energies : np.ndarray, + full_orbitals : np.ndarray, + full_l_terms : np.ndarray, + exchange_potentials : np.ndarray + ) -> np.ndarray: + """ + Compute OEP exchange kernel and the exchange driving term + """ + # get occupied orbitals + occ_orbitals = full_orbitals[:, :len(self.occ_l_values)] + + # get l channel indices for all orbitals + l_max = np.max(self.occ_l_values) + l_channel_orbital_indices = np.zeros((l_max + 1, self.n_interior), dtype=np.int32) + for l in range(l_max + 1): + l_channel_orbital_indices[l, :] = np.argwhere(full_l_terms == l)[:,0] + + # compute chi_0_kernel and the exchange driving term + chi_0_kernel = np.zeros((self.n_quad, self.n_quad)) + exchange_driving_term = np.zeros(self.n_quad) + + for idx in range(len(self.occ_l_values)): + # get l and n index + l_value = self.occ_l_values[idx] + n_value = self.occ_n_values[idx] - l_value - 1 + l_occ_num = len(np.argwhere(self.occ_l_values == l_value)[:,0]) + + # get all orbitals with indices of the same l value + unocc_orbitals_in_l_channel = full_orbitals[:, l_channel_orbital_indices[l_value, :]][:, l_occ_num:] + + # get the difference of eigenvalues + l_channel_eigenvalues = full_eigen_energies[l_channel_orbital_indices[l_value, :]] + diff_eigenvalues = l_channel_eigenvalues.reshape(-1, 1) - l_channel_eigenvalues.reshape(1, -1) + one_over_diff_eigenvalues = 1 / (diff_eigenvalues + np.eye(self.n_interior)) # avoid division by zero + one_over_diff_eigenvalues[np.arange(self.n_interior), np.arange(self.n_interior)] = 0 # set the diagonal to zero + + # get the green function block + _exchange_green_block = np.einsum('ji,ki,i->jk', + unocc_orbitals_in_l_channel, + unocc_orbitals_in_l_channel, + one_over_diff_eigenvalues[n_value, l_occ_num:], + optimize=True + ) + + # get the orbital and corresponding exchange potential inside this for loop + orbital = occ_orbitals[:, idx] + exchange_potential = exchange_potentials[idx] + + # update chi_0_kernel + chi_0_kernel += 4 * np.einsum('k,kj,j->kj', + orbital, + _exchange_green_block, + orbital, + optimize=True + ) * (2 * l_value + 1) + + # update exchange driving term + exchange_driving_term += 4 * np.einsum('k,kl,l->k', + orbital, + _exchange_green_block, + exchange_potential, + optimize=True + ) * (2 * l_value + 1) + + return chi_0_kernel, exchange_driving_term + + + + def _compute_rpa_correlation_driving_term( + self, + full_eigen_energies : np.ndarray, + full_orbitals : np.ndarray, + full_l_terms : np.ndarray, + enable_parallelization : bool = False, + ) -> np.ndarray: + """ + Compute RPA correlation driving term from full spectrum in parallel. + """ + assert hasattr(self, 'frequency_grid') and hasattr(self, 'frequency_weights'), \ + PARENT_CLASS_RPACORRELATION_NOT_INITIALIZED_ERROR + assert self.angular_momentum_cutoff is not None, \ + ANGULAR_MOMENTUM_CUTOFF_NOT_NONE_ERROR_MESSAGE + assert isinstance(enable_parallelization, bool), \ + ENABLE_PARALLELIZATION_NOT_BOOL_ERROR.format(type(enable_parallelization)) + self._validate_full_spectrum_inputs(full_eigen_energies, full_orbitals, full_l_terms) + + l_occ_max = np.max(self.occ_l_values) + l_unocc_max = np.max(full_l_terms) + l_couple_max = l_occ_max + l_unocc_max + + # Compute RPA Wigner symbols squared array + wigner_symbols_squared = self._compute_rpa_wigner_symbols_squared( + l_occ_max = np.max(self.occ_l_values), + l_unocc_max = np.max(full_l_terms), + ) + + # Compute RPA radial kernels dictionary + radial_kernels_dict = {} + for l_couple in range(l_couple_max + 1): + radial_kernels_dict[l_couple] = CoulombCouplingCalculator.radial_kernel( + l = l_couple, + r_nodes = self.quadrature_nodes, + r_weights = self.quadrature_weights, + ) + + q1c_term = np.zeros(self.n_quad) + q2c_term = np.zeros(self.n_quad) + + # if enable_parallelization is False, compute RPA correlation driving term at each frequency and sum them up + + if not enable_parallelization: + # Compute RPA correlation driving term at each frequency and sum them up + for frequency, frequency_weight in zip(self.frequency_grid, self.frequency_weights): + q1c_term_single, q2c_term_single = self._compute_rpa_correlation_driving_term_for_single_frequency( + frequency = frequency, + angular_momentum_cutoff = self.angular_momentum_cutoff, + occupation_info = self.occupation_info, + full_eigen_energies = full_eigen_energies, + full_orbitals = full_orbitals, + full_l_terms = full_l_terms, + wigner_symbols_squared = wigner_symbols_squared, + radial_kernels_dict = radial_kernels_dict, + ) + + q1c_term += 4 * q1c_term_single * frequency_weight # 4 comes from spin degeneracy + q2c_term += 2 * q2c_term_single * frequency_weight # 2 comes from spin degeneracy + else: + import multiprocessing as mp + from concurrent.futures import ThreadPoolExecutor + + + def _single_frequency_task(args): + idx, (frequency, frequency_weight) = args + q1c_term_single, q2c_term_single = self._compute_rpa_correlation_driving_term_for_single_frequency( + frequency = frequency, + angular_momentum_cutoff = self.angular_momentum_cutoff, + occupation_info = self.occupation_info, + full_eigen_energies = full_eigen_energies, + full_orbitals = full_orbitals, + full_l_terms = full_l_terms, + wigner_symbols_squared = wigner_symbols_squared, + radial_kernels_dict = radial_kernels_dict, + ) + return ( + idx, + 4 * q1c_term_single * frequency_weight, + 2 * q2c_term_single * frequency_weight, + ) + + n_workers = min(max(1, mp.cpu_count()), len(self.frequency_grid)) + + # limit BLAS/OpenMP threads during parallel sections + if threadpool_limits is not None: + blas_ctx = threadpool_limits(limits=1) + else: + blas_ctx = nullcontext() + + # run the parallel sections + with blas_ctx, ThreadPoolExecutor(max_workers=n_workers) as executor: + results = executor.map( + _single_frequency_task, + enumerate(zip(self.frequency_grid, self.frequency_weights)) + ) + for _, q1c_single_weighted, q2c_single_weighted in results: + q1c_term += q1c_single_weighted + q2c_term += q2c_single_weighted + + + assert q1c_term.shape == (self.n_quad,) + assert q2c_term.shape == (self.n_quad,) + + total_driving_term = (q1c_term + q2c_term) / (2 * np.pi) + + return total_driving_term + diff --git a/utils/atom/xc/rpa.py b/utils/atom/xc/rpa.py index 6173f1bf..da0548de 100644 --- a/utils/atom/xc/rpa.py +++ b/utils/atom/xc/rpa.py @@ -1,12 +1,719 @@ -from __future__ import annotations -import numpy as np -from typing import Any - -class RPACorrelation: - """Compute RPA correlation energy from eigenstates. Placeholder.""" - def __init__(self, q_omega: int, cutoff_freq: float = 0.0): - self.q_omega = int(q_omega) - self.cutoff_freq = float(cutoff_freq) - - def energy(self, eigvals: np.ndarray, eigvecs: np.ndarray, meta: dict[str, Any]) -> float: - raise NotImplementedError("RPA correlation energy not implemented yet.") +from __future__ import annotations + + +import numpy as np +from typing import Any, Tuple, List, Dict + +from .hybrid import CoulombCouplingCalculator +from ..mesh.builder import Quadrature1D +from ..mesh.operators import RadialOperatorsBuilder +from ..utils.occupation_states import OccupationInfo + +from contextlib import nullcontext + +try: + # Optional dependency: used to limit BLAS/OpenMP threads during parallel sections + from threadpoolctl import threadpool_limits +except ImportError: + threadpool_limits = None # type: ignore + +# Error messages +OPS_BUILDER_NOT_RADIAL_OPERATORS_BUILDER_ERROR = \ + "Parameter `ops_builder` must be a `RadialOperatorsBuilder` instance, get type `{}` instead" +OCCUPATION_INFO_NOT_OCCUPATION_INFO_ERROR = \ + "Parameter `occupation_info` must be a `OccupationInfo` instance, get type `{}` instead" +FREQUENCY_QUADRATURE_POINT_NUMBER_NOT_INTEGER_ERROR = \ + "Parameter `frequency_quadrature_point_number` must be an integer, get type {} instead" +FREQUENCY_QUADRATURE_POINT_NUMBER_NOT_GREATER_THAN_0_ERROR = \ + "Parameter `frequency_quadrature_point_number` must be greater than 0, get {} instead" +ANGULAR_MOMENTUM_CUTOFF_NOT_INTEGER_ERROR = \ + "Parameter `angular_momentum_cutoff` must be an integer, get type {} instead" +ANGULAR_MOMENTUM_CUTOFF_NEGATIVE_ERROR = \ + "Parameter `angular_momentum_cutoff` must be non-negative, get {} instead" +FREQUENCY_NOT_FLOAT_ERROR = \ + "Parameter `frequency` must be a float or a scaler numpy array, get type {} instead" + +ANGULAR_MOMENTUM_CUTOFF_NOT_NONE_ERROR_MESSAGE = \ + "Parameter `angular_momentum_cutoff` must be not None, get None instead" +OCCUPATION_INFO_L_TERMS_NOT_CONSISTENT_WITH_OCCUPATION_INFO_ERROR = \ + "Occupied l terms are not consistent with the occupation information, please check your inputs, got {} instead of {}" +PARENT_CLASS_RPACORRELATION_NOT_INITIALIZED_ERROR = \ + "Parent class `RPACorrelation` is not initialized, please initialize it first" + +L_OCC_MAX_NOT_INTEGER_ERROR = \ + "Parameter `l_occ_max` must be an integer, get type {} instead" +L_UNOCC_MAX_NOT_INTEGER_ERROR = \ + "Parameter `l_unocc_max` must be an integer, get type {} instead" +L_COUPLE_MAX_NOT_INTEGER_ERROR = \ + "Parameter `l_couple_max` must be an integer, get type {} instead" +ENABLE_PARALLELIZATION_NOT_BOOL_ERROR = \ + "Parameter `enable_parallelization` must be a bool, get type {} instead" + + + +class RPACorrelation: + """ + Compute RPA correlation energy from eigenstates. + """ + def __init__( + self, + ops_builder : 'RadialOperatorsBuilder', + occupation_info : 'OccupationInfo', + frequency_quadrature_point_number : int, + angular_momentum_cutoff : int + ): + """ + Parameters + ---------- + ops_builder : instance of RadialOperatorsBuilder + RadialOperatorsBuilder instance + occupation_info : instance of OccupationInfo + Occupation information + frequency_quadrature_point_number : int + Number of frequency quadrature points + angular_momentum_cutoff : int + Maximum angular momentum quantum number to include + """ + assert isinstance(ops_builder, RadialOperatorsBuilder), \ + OPS_BUILDER_NOT_RADIAL_OPERATORS_BUILDER_ERROR.format(type(ops_builder)) + assert isinstance(occupation_info, OccupationInfo), \ + OCCUPATION_INFO_NOT_OCCUPATION_INFO_ERROR.format(type(occupation_info)) + assert isinstance(frequency_quadrature_point_number, int), \ + FREQUENCY_QUADRATURE_POINT_NUMBER_NOT_INTEGER_ERROR.format(type(frequency_quadrature_point_number)) + assert frequency_quadrature_point_number > 0, \ + FREQUENCY_QUADRATURE_POINT_NUMBER_NOT_GREATER_THAN_0_ERROR.format(frequency_quadrature_point_number) + assert isinstance(angular_momentum_cutoff, int), \ + ANGULAR_MOMENTUM_CUTOFF_NOT_INTEGER_ERROR.format(type(angular_momentum_cutoff)) + assert angular_momentum_cutoff >= 0, \ + ANGULAR_MOMENTUM_CUTOFF_NEGATIVE_ERROR.format(angular_momentum_cutoff) + + + # Extract quadrature data from ops_builder + self.n_quad = len(ops_builder.quadrature_nodes) + self.quadrature_nodes = ops_builder.quadrature_nodes + self.quadrature_weights = ops_builder.quadrature_weights + + # initialize the frequency grid and weights + self.frequency_quadrature_point_number = frequency_quadrature_point_number + self.frequency_grid, self.frequency_weights = \ + self._initialize_frequency_grid_and_weights(frequency_quadrature_point_number) + + # occupation information + self.occupations : np.ndarray = occupation_info.occupations + self.occ_l_values : np.ndarray = occupation_info.l_values + self.occ_n_values : np.ndarray = occupation_info.n_values + + # angular momentum cutoff + self.angular_momentum_cutoff = angular_momentum_cutoff + + + @staticmethod + def _initialize_frequency_grid_and_weights(n: int) -> Tuple[np.ndarray, np.ndarray]: + """ + Initialize frequency grid and weights for RPA correlation energy calculations. + + Parameters + ---------- + n : int + Number of frequency quadrature points + + Reference: + https://journals.aps.org/prl/supplemental/10.1103/PhysRevLett.134.016402/scrpa4_SM.pdf + + Returns + ------- + frequency_grid : np.ndarray + Frequency grid + frequency_weights : np.ndarray + Frequency weights + """ + frequency_scale = 2.5 + + # Get Gauss-Legendre nodes on reference interval [-1, 1] + reference_nodes, reference_weights = Quadrature1D.gauss_legendre(n) + + # Transform to semi-infinite interval [0, ∞) + nodes = frequency_scale * (1 + reference_nodes) / (1 - reference_nodes) + weights = reference_weights * 2 * frequency_scale / (1 - reference_nodes)**2 + + return nodes, weights + + + @staticmethod + def _compute_rpa_correlation_driving_term_for_single_frequency( + frequency : float, + angular_momentum_cutoff : int, + occupation_info : OccupationInfo, + full_eigen_energies : np.ndarray, + full_orbitals : np.ndarray, + full_l_terms : np.ndarray, + wigner_symbols_squared : np.ndarray, + radial_kernels_dict : Dict[int, np.ndarray], + ) -> Tuple[np.ndarray, np.ndarray]: + """ + Compute RPA correlation driving term for at given frequency. + """ + try: + frequency = float(frequency) + except ValueError: + raise ValueError(FREQUENCY_NOT_FLOAT_ERROR.format(type(frequency))) + assert isinstance(occupation_info, OccupationInfo), \ + OCCUPATION_INFO_NOT_OCCUPATION_INFO_ERROR.format(type(occupation_info)) + + # get occupation information + occupations = occupation_info.occupations + occ_l_values = occupation_info.l_values + + + # get the number of occupied and unoccupied orbitals + occ_orbitals_num = len(occ_l_values) + total_orbitals_num = len(full_eigen_energies) + + # get the number of quadrature and n_interior points + n_quad = radial_kernels_dict[0].shape[0] + n_interior = len(np.argwhere(full_l_terms == 0)[:, 0]) + + # get occupied and unoccupied orbitals and energies + occ_orbitals = full_orbitals[:, :occ_orbitals_num] # shape: (n_grid, total_orbitals_num) + occ_energies = full_eigen_energies[:occ_orbitals_num] # shape: (total_orbitals_num,) + occ_l_terms = full_l_terms[:occ_orbitals_num] # shape: (total_orbitals_num,) + unocc_orbitals = full_orbitals[:, occ_orbitals_num:] # shape: (n_grid, unocc_orbitals_num) + unocc_energies = full_eigen_energies[occ_orbitals_num:] # shape: (unocc_orbitals_num,) + unocc_l_terms = full_l_terms[occ_orbitals_num:] # shape: (total_orbitals_num,) + + assert np.all(occ_l_terms == occ_l_values), \ + OCCUPATION_INFO_L_TERMS_NOT_CONSISTENT_WITH_OCCUPATION_INFO_ERROR.format(occ_l_terms, occ_l_values) + + ### ================================================ ### + ### Part 1: Compute the RPA correlation prefactors ### + ### ================================================ ### + + # Angular degeneracy factors f_p * (2l_q + 1) + # shape: (occ_num, unocc_num) + deg_factors = occupations[:, np.newaxis] * (2 * unocc_l_terms + 1)[np.newaxis, :] + + # Energy differences Δε_{pq} = ε_p - ε_q + # shape: (occ_num, unocc_num) + delta_eps = occ_energies[:, np.newaxis] - unocc_energies[np.newaxis, :] + + # Filter out zero entries (valid (p,q) pairs) + occ_idx, unocc_idx = np.argwhere((deg_factors != 0) & (delta_eps != 0)).T + deg_factors_valid = deg_factors[occ_idx, unocc_idx] + delta_eps_valid = delta_eps[occ_idx, unocc_idx] + + # Compute Lorentzian frequency response: Δε / (Δε² + ω²) + # shape: (n_valid_pairs) + lorentzian_factors = delta_eps_valid / (delta_eps_valid ** 2 + frequency ** 2) + + # Compute frequency derivative factors for Q2c: (Δε² - ω²) / (Δε² + ω²)² + # This arises from ∂χ₀,L/∂(iω), used in the Q2c driving term + # shape: (n_valid_pairs) + frequency_derivative_factors = (delta_eps_valid ** 2 - frequency ** 2) / (delta_eps_valid ** 2 + frequency ** 2) ** 2 + + # Combine angular and frequency factors (without Wigner 3j symbol yet) + # shape: (n_valid_pairs) + prefactors_q1c = deg_factors_valid * lorentzian_factors + prefactors_q2c = deg_factors_valid * frequency_derivative_factors + + # shape: (occ_num, unocc_num) + prefactors_self_energy = deg_factors * delta_eps / (delta_eps ** 2 + frequency ** 2) + + + ### ================================================== ### + ### Part 2: Compute the nonzero Wigner symbols ### + ### ================================================== ### + + l_couple_min = np.min(np.abs(occ_l_terms[occ_idx] - unocc_l_terms[unocc_idx])).astype(np.int32) + l_couple_max = np.max(occ_l_terms[occ_idx] + unocc_l_terms[unocc_idx]).astype(np.int32) + l_couple_range = np.arange(l_couple_min, l_couple_max + 1) + + + # Use advanced indexing with broadcasting - one operation instead of triple loop + # occ_l_terms[i] gives angular momentum of orbital i + # unocc_l_terms[j] gives angular momentum of orbital j + + wigner_symbols_squared_all = wigner_symbols_squared[ + occ_l_terms [:, np.newaxis, np.newaxis], # shape: (occ_orbitals_num, 1, 1) + unocc_l_terms [np.newaxis, :, np.newaxis], # shape: (1, unocc_orbitals_num, 1) + l_couple_range[np.newaxis, np.newaxis, :], # shape: (1, 1, l_couple_num) + ] + wigner_symbols_squared_valid = wigner_symbols_squared_all[occ_idx, unocc_idx, :] + + # Compute only the nonzero Wigner symbols for each l_couple channel + active_l_couple_idx_list : List[int] = [] + active_wigner_symbols_indices_list : List[np.ndarray] = [] + + for l_couple_idx, l_couple in enumerate(l_couple_range): + non_zero_wigner_symbols_indices = np.argwhere(wigner_symbols_squared_valid[:, l_couple_idx] != 0)[:, 0] + + # Collect active l_couple and corresponding nonzero Wigner symbols indices + if len(non_zero_wigner_symbols_indices) != 0: + active_l_couple_idx_list.append(l_couple_idx) + active_wigner_symbols_indices_list.append(non_zero_wigner_symbols_indices) + + + ### ================================================== ### + ### Part 3: Compute orbital products ### + ### ================================================== ### + + # orbital outer product: φ_p(r) ⊗ φ_q(r) for all (p,q) pairs + # shape: (n_grid, occ_num, unocc_num) + orbital_product_outer = np.einsum('li,lj->ijl', + occ_orbitals, + unocc_orbitals, + optimize = True, + ) + + # Orbital squared difference: φ_p²(r) - φ_q²(r) for valid (p,q) pairs + # shape: (n_grid, n_valid_pairs) + orbital_squared_diff = occ_orbitals[:, occ_idx] ** 2 - unocc_orbitals[:, unocc_idx] ** 2 + + # Orbital pair product: Φ_{pq}(r) = φ_p(r)φ_q(r) for valid (p,q) pairs + # shape: (n_grid, n_valid_pairs) + orbital_pair_product = occ_orbitals[:, occ_idx] * unocc_orbitals[:, unocc_idx] + + + # initialize the full self-energy potential + # shape: (total_orbitals_num, n_quad) + full_self_energy_potential = np.zeros((total_orbitals_num, n_quad)) + + + ### ================================================== ### + ### Part 4: Compute RPA correlation driving term ### + ### ================================================== ### + + # initialize the q1c and q2c terms + # shape: (n_quad,) + full_q1c_term = np.zeros(n_quad) + full_q2c_term = np.zeros(n_quad) + + + # Compute RPA correlation driving term for each l_couple channel + for active_l_couple_idx, active_wigner_symbols_indices in zip(active_l_couple_idx_list, active_wigner_symbols_indices_list): + + + active_l_couple = l_couple_range[active_l_couple_idx] + + # Get radial kernel + # shape: (n_quad, n_quad) + radial_kernel = radial_kernels_dict[active_l_couple] + + # Compute rpa_response_kernel + # shape: (n_quad, n_quad) + # Reference: https://stackoverflow.com/questions/17437523/python-fast-way-to-sum-outer-products + rpa_response_kernel = np.matmul( + radial_kernel, + 2 * np.einsum( + 'ij,ik->kj', + np.einsum('ji,i->ij', + orbital_pair_product[:, active_wigner_symbols_indices], + prefactors_q1c[active_wigner_symbols_indices], + optimize=True, + ), + np.einsum('i,ki->ik', + wigner_symbols_squared_valid[active_wigner_symbols_indices, active_l_couple], + orbital_pair_product[:, active_wigner_symbols_indices], + optimize=True, + ), + optimize=True, + ), + ) + + # Compute dyson solved response + dyson_solved_response = np.linalg.solve(np.eye(n_quad) - rpa_response_kernel, radial_kernel) - radial_kernel + + # Compute self-energy potential term + # shape: (total_orbitals_num, n_quad) + _self_energy_potential = np.zeros((total_orbitals_num, n_quad)) + + # Compute self-energy potential term for occupied states + for occ_l_index in range(occ_orbitals_num): + # Get nonzero unoccupied l indices for this occupied l + nonzero_unocc_indices = np.argwhere(wigner_symbols_squared_all[occ_l_index, :, active_l_couple] != 0)[:, 0] + _prefactor = prefactors_self_energy[occ_l_index, nonzero_unocc_indices] * wigner_symbols_squared_all[occ_l_index, nonzero_unocc_indices, active_l_couple] + + _self_energy_potential[occ_l_index, :] = \ + np.einsum('i,ki,il,lk->k', + _prefactor, + unocc_orbitals[:, nonzero_unocc_indices], + orbital_product_outer[occ_l_index, nonzero_unocc_indices, :], + dyson_solved_response, + optimize=True, + ) * (2 * active_l_couple + 1) + + + # Compute self-energy potential term for unoccupied states + self_energy_coupling_matrix = prefactors_self_energy * wigner_symbols_squared_all[:, :, active_l_couple] + coupled_unocc_indices = np.argwhere(~np.all(self_energy_coupling_matrix == 0, axis=0))[:, 0] + + _self_energy_potential[occ_orbitals_num:, :][coupled_unocc_indices, :] = \ + np.einsum('ji,kj,jik->ik', + self_energy_coupling_matrix[:, coupled_unocc_indices], + occ_orbitals, + np.einsum('jil,lk->jik', + orbital_product_outer[:, coupled_unocc_indices, :], + dyson_solved_response, + optimize=True, + ), + optimize=True, + ) * (2 * active_l_couple + 1) + + # Compute q2c term + _q2c_term = np.einsum('ki, i, i->k', + orbital_squared_diff[:, active_wigner_symbols_indices], + prefactors_q2c[active_wigner_symbols_indices] * wigner_symbols_squared_valid[active_wigner_symbols_indices, active_l_couple], + np.einsum('li,pi,pl->i', + orbital_pair_product[:, active_wigner_symbols_indices], + orbital_pair_product[:, active_wigner_symbols_indices], + dyson_solved_response, + optimize=True, + ), + optimize=True, + ) * (2 * active_l_couple + 1) + + # Update the full q2c term and self-energy potential + full_q2c_term += _q2c_term + full_self_energy_potential += _self_energy_potential + + + + # Compute q1c term + for l_value in range(angular_momentum_cutoff + 1): + + # Get total orbitals in this l_value channel + l_indices = np.argwhere(full_l_terms == l_value)[:, 0] + total_orbitals_in_l_channel = full_orbitals[:, l_indices] + self_energy_in_l_channel = full_self_energy_potential[l_indices, :] + eigenvalues_in_l_channel = full_eigen_energies[l_indices] + + # Compute difference of eigenvalues in this l_value channel + diff_eigenvalues = eigenvalues_in_l_channel.reshape(-1, 1) - eigenvalues_in_l_channel.reshape(1, -1) + one_over_diff_eigenvalues = 1 / (diff_eigenvalues + np.eye(n_interior)) # avoid division by zero + one_over_diff_eigenvalues[np.arange(n_interior), np.arange(n_interior)] = 0 # set the diagonal to zero + + + # Compute q1c term for this l_value channel + q1c_term_in_l_channel = \ + np.einsum('ki,ik->k', + total_orbitals_in_l_channel, + np.einsum('ij,kj,ij->ik', + one_over_diff_eigenvalues, + total_orbitals_in_l_channel, + np.einsum('ix,xj->ij', + self_energy_in_l_channel, + total_orbitals_in_l_channel, + optimize=True, + ), + optimize=True, + ), + optimize=True, + ) + + full_q1c_term -= q1c_term_in_l_channel + + assert full_q1c_term.shape == (n_quad,) + assert full_q2c_term.shape == (n_quad,) + + + return full_q1c_term, full_q2c_term + + + + @staticmethod + def _compute_rpa_wigner_symbols_squared( + l_occ_max : int, + l_unocc_max : int, + ) -> np.ndarray: + """ + Compute RPA Wigner symbols squared array. + + Parameters + ---------- + l_occ_max : int + Maximum angular momentum quantum number for occupied orbitals + l_unocc_max : int + Maximum angular momentum quantum number for unoccupied orbitals + + Returns + ------- + wigner_symbols_squared : np.ndarray + Wigner symbols squared array + shape: (l_occ_max + 1, l_unocc_max + 1, l_couple_max + 1), where l_couple_max = l_occ_max + l_unocc_max + """ + try: + l_occ_max = int(l_occ_max) + except ValueError: + raise ValueError(L_OCC_MAX_NOT_INTEGER_ERROR.format(type(l_occ_max))) + try: + l_unocc_max = int(l_unocc_max) + except ValueError: + raise ValueError(L_UNOCC_MAX_NOT_INTEGER_ERROR.format(type(l_unocc_max))) + assert l_occ_max >= 0 and l_unocc_max >= 0, \ + "All angular momentum quantum numbers must be non-negative" + + # Compute the maximum angular momentum quantum number for the coupled system + l_couple_max = l_occ_max + l_unocc_max + + # Initialize Wigner symbols squared array + wigner_symbols_squared = np.zeros((l_occ_max + 1, l_unocc_max + 1, l_couple_max + 1)) + + # Compute Wigner symbols squared for all (l_occ, l_unocc, l_couple) combinations + for l_occ in range(l_occ_max + 1): + for l_unocc in range(l_unocc_max + 1): + for l_couple in range(l_couple_max + 1): + wigner_symbols_squared[l_occ, l_unocc, l_couple] = \ + CoulombCouplingCalculator.wigner_3j_000(l_occ, l_unocc, l_couple) ** 2 + + return wigner_symbols_squared + + + @staticmethod + def _compute_correlation_energy_for_single_frequency( + frequency : float, + occupation_info : OccupationInfo, + full_eigen_energies : np.ndarray, + full_orbitals : np.ndarray, + full_l_terms : np.ndarray, + wigner_symbols_squared : np.ndarray, + radial_kernels_dict : Dict[int, np.ndarray], + ) -> float: + """ + Compute RPA correlation driving term for at given frequency. + """ + try: + frequency = float(frequency) + except ValueError: + raise ValueError(FREQUENCY_NOT_FLOAT_ERROR.format(type(frequency))) + assert isinstance(occupation_info, OccupationInfo), \ + OCCUPATION_INFO_NOT_OCCUPATION_INFO_ERROR.format(type(occupation_info)) + + # get occupation information + occupations = occupation_info.occupations + occ_l_values = occupation_info.l_values + + + # get the number of occupied and unoccupied orbitals + occ_orbitals_num = len(occ_l_values) + + # get the number of quadrature and n_interior points + n_quad = radial_kernels_dict[0].shape[0] + + # get occupied and unoccupied orbitals and energies + occ_orbitals = full_orbitals[:, :occ_orbitals_num] # shape: (n_grid, total_orbitals_num) + occ_energies = full_eigen_energies[:occ_orbitals_num] # shape: (total_orbitals_num,) + occ_l_terms = full_l_terms[:occ_orbitals_num] # shape: (total_orbitals_num,) + unocc_orbitals = full_orbitals[:, occ_orbitals_num:] # shape: (n_grid, unocc_orbitals_num) + unocc_energies = full_eigen_energies[occ_orbitals_num:] # shape: (unocc_orbitals_num,) + unocc_l_terms = full_l_terms[occ_orbitals_num:] # shape: (total_orbitals_num,) + + assert np.all(occ_l_terms == occ_l_values), \ + OCCUPATION_INFO_L_TERMS_NOT_CONSISTENT_WITH_OCCUPATION_INFO_ERROR.format(occ_l_terms, occ_l_values) + + + ### ================================================ ### + ### Part 1: Compute the RPA correlation prefactors ### + ### ================================================ ### + + # Angular degeneracy factors f_p * (2l_q + 1) + # shape: (occ_num, unocc_num) + deg_factors = occupations[:, np.newaxis] * (2 * unocc_l_terms + 1)[np.newaxis, :] + + # Energy differences Δε_{pq} = ε_p - ε_q + # shape: (occ_num, unocc_num) + delta_eps = occ_energies[:, np.newaxis] - unocc_energies[np.newaxis, :] + + # Filter out zero entries (valid (p,q) pairs) + occ_idx, unocc_idx = np.argwhere((deg_factors != 0) & (delta_eps != 0)).T + deg_factors_valid = deg_factors[occ_idx, unocc_idx] + delta_eps_valid = delta_eps[occ_idx, unocc_idx] + + # Compute Lorentzian frequency response: Δε / (Δε² + ω²) + # shape: (n_valid_pairs) + lorentzian_factors = delta_eps_valid / (delta_eps_valid ** 2 + frequency ** 2) + + # Combine angular and frequency factors (without Wigner 3j symbol yet) + # shape: (n_valid_pairs) + prefactors_q1c = deg_factors_valid * lorentzian_factors + + ### ================================================== ### + ### Part 2: Compute the nonzero Wigner symbols ### + ### ================================================== ### + + l_couple_min = np.min(np.abs(occ_l_terms[occ_idx] - unocc_l_terms[unocc_idx])).astype(np.int32) + l_couple_max = np.max(occ_l_terms[occ_idx] + unocc_l_terms[unocc_idx]).astype(np.int32) + l_couple_range = np.arange(l_couple_min, l_couple_max + 1) + + + # Use advanced indexing with broadcasting - one operation instead of triple loop + # occ_l_terms[i] gives angular momentum of orbital i + # unocc_l_terms[j] gives angular momentum of orbital j + + wigner_symbols_squared_all = wigner_symbols_squared[ + occ_l_terms [:, np.newaxis, np.newaxis], # shape: (occ_orbitals_num, 1, 1) + unocc_l_terms [np.newaxis, :, np.newaxis], # shape: (1, unocc_orbitals_num, 1) + l_couple_range[np.newaxis, np.newaxis, :], # shape: (1, 1, l_couple_num) + ] + wigner_symbols_squared_valid = wigner_symbols_squared_all[occ_idx, unocc_idx, :] + + # Compute only the nonzero Wigner symbols for each l_couple channel + active_l_couple_idx_list : List[int] = [] + active_wigner_symbols_indices_list : List[np.ndarray] = [] + + for l_couple_idx, l_couple in enumerate(l_couple_range): + non_zero_wigner_symbols_indices = np.argwhere(wigner_symbols_squared_valid[:, l_couple_idx] != 0)[:, 0] + + # Collect active l_couple and corresponding nonzero Wigner symbols indices + if len(non_zero_wigner_symbols_indices) != 0: + active_l_couple_idx_list.append(l_couple_idx) + active_wigner_symbols_indices_list.append(non_zero_wigner_symbols_indices) + + + ### ================================================== ### + ### Part 3: Compute the RPA correlation energy ### + ### ================================================== ### + + + # Orbital pair product: Φ_{pq}(r) = φ_p(r)φ_q(r) for valid (p,q) pairs + # shape: (n_grid, n_valid_pairs) + orbital_pair_product = occ_orbitals[:, occ_idx] * unocc_orbitals[:, unocc_idx] + + # initialize the q1c and q2c terms + # shape: (n_quad,) + full_correlation_energy_at_single_frequency = 0.0 + + # Compute RPA correlation driving term for each l_couple channel + for active_l_couple_idx, active_wigner_symbols_indices in zip(active_l_couple_idx_list, active_wigner_symbols_indices_list): + + active_l_couple = l_couple_range[active_l_couple_idx] + + # Get radial kernel + # shape: (n_quad, n_quad) + radial_kernel = radial_kernels_dict[active_l_couple] + + # Compute rpa_response_kernel + # shape: (n_quad, n_quad) + # Reference: https://stackoverflow.com/questions/17437523/python-fast-way-to-sum-outer-products + rpa_response_kernel = np.matmul( + radial_kernel, + 2 * np.einsum( + 'ij,ik->kj', + np.einsum('ji,i->ij', + orbital_pair_product[:, active_wigner_symbols_indices], + prefactors_q1c[active_wigner_symbols_indices], + optimize=True, + ), + np.einsum('i,ki->ik', + wigner_symbols_squared_valid[active_wigner_symbols_indices, active_l_couple], + orbital_pair_product[:, active_wigner_symbols_indices], + optimize=True, + ), + optimize=True, + ), + ) + + # Compute correlation energy for this l_couple channel + full_correlation_energy_at_single_frequency += \ + (2 * active_l_couple + 1) * (np.log(np.linalg.det(np.eye(n_quad) - rpa_response_kernel)) + np.trace(rpa_response_kernel)) + + return full_correlation_energy_at_single_frequency + + + + + def compute_correlation_energy( + self, + full_eigen_energies : np.ndarray, + full_orbitals : np.ndarray, + full_l_terms : np.ndarray, + enable_parallelization: bool = False, + ) -> float: + """ + Compute RPA correlation energy from eigenstates. + """ + assert hasattr(self, 'frequency_grid') and hasattr(self, 'frequency_weights'), \ + PARENT_CLASS_RPACORRELATION_NOT_INITIALIZED_ERROR + + assert isinstance(enable_parallelization, bool), \ + ENABLE_PARALLELIZATION_NOT_BOOL_ERROR.format(type(enable_parallelization)) + + + + self._validate_full_spectrum_inputs(full_eigen_energies, full_orbitals, full_l_terms) + + l_occ_max = np.max(self.occ_l_values) + l_unocc_max = np.max(full_l_terms) + l_couple_max = l_occ_max + l_unocc_max + + # Compute RPA Wigner symbols squared array + wigner_symbols_squared = self._compute_rpa_wigner_symbols_squared( + l_occ_max = np.max(self.occ_l_values), + l_unocc_max = np.max(full_l_terms), + ) + + # Compute RPA radial kernels dictionary + radial_kernels_dict = {} + for l_couple in range(l_couple_max + 1): + radial_kernels_dict[l_couple] = CoulombCouplingCalculator.radial_kernel( + l = l_couple, + r_nodes = self.quadrature_nodes, + r_weights = self.quadrature_weights, + ) + + # Compute RPA correlation energy at each frequency and sum them up + correlation_energy = 0.0 + + if not enable_parallelization: + for frequency, frequency_weight in zip(self.frequency_grid, self.frequency_weights): + correlation_energy_at_single_frequency = self._compute_correlation_energy_for_single_frequency( + frequency = frequency, + occupation_info = self.occupation_info, + full_eigen_energies = full_eigen_energies, + full_orbitals = full_orbitals, + full_l_terms = full_l_terms, + wigner_symbols_squared = wigner_symbols_squared, + radial_kernels_dict = radial_kernels_dict, + ) + + correlation_energy += correlation_energy_at_single_frequency * frequency_weight + else: + import multiprocessing as mp + from concurrent.futures import ThreadPoolExecutor + + def _single_frequency_task(args): + idx, (frequency, frequency_weight) = args + correlation_energy_at_single_frequency = self._compute_correlation_energy_for_single_frequency( + frequency = frequency, + occupation_info = self.occupation_info, + full_eigen_energies = full_eigen_energies, + full_orbitals = full_orbitals, + full_l_terms = full_l_terms, + wigner_symbols_squared = wigner_symbols_squared, + radial_kernels_dict = radial_kernels_dict, + ) + return ( + idx, + correlation_energy_at_single_frequency * frequency_weight, + ) + + n_workers = min(max(1, mp.cpu_count()), len(self.frequency_grid)) + + # limit BLAS/OpenMP threads during parallel sections + if threadpool_limits is not None: + blas_ctx = threadpool_limits(limits=1) + else: + blas_ctx = nullcontext() + + + with blas_ctx, ThreadPoolExecutor(max_workers=n_workers) as executor: + results = executor.map( + _single_frequency_task, + enumerate(zip(self.frequency_grid, self.frequency_weights)) + ) + for _, correlation_energy_single_weighted in results: + correlation_energy += correlation_energy_single_weighted + + + return correlation_energy / (2 * np.pi) + + + diff --git a/utils/pdos/README.md b/utils/pdos/README.md index 6b445200..bc395084 100644 --- a/utils/pdos/README.md +++ b/utils/pdos/README.md @@ -11,12 +11,12 @@ Before using the PDOS calculator, ensure you have the following Python packages - **Numpy** - Numerical computing library - **Scipy** - Scientific computing library (includes linalg, interpolate, sparse, special modules) - **PyYAML** - YAML configuration file parsing - +- **Pandas** - Data manipulation and analysis You can install the required packages using pip: ```bash -pip install numpy scipy pyyaml +pip install numpy scipy pandas pyyaml ``` diff --git a/utils/pdos/calculate_pdos.py b/utils/pdos/calculate_pdos.py index ee046631..34dd8010 100644 --- a/utils/pdos/calculate_pdos.py +++ b/utils/pdos/calculate_pdos.py @@ -48,6 +48,21 @@ # Regex patterns SCIENTIFIC_NOTATION_PATTERN = r'[+-]?\d+(?:\.\d+)?(?:[Ee][+-]?\d+)?' +# XC Functional valid list +# Valid XC Functional +VALID_XC_FUNCTIONAL_LIST = [ + 'LDA_PZ' , # LDA Perdew-Zunger + 'LDA_PW' , # LDA Perdew-Wang + 'GGA_PBE', # GGA Perdew-Burke-Ernzerhof + 'SCAN' , # SCAN functional, meta-GGA + 'RSCAN' , # RSCAN functional, meta-GGA + 'R2SCAN' , # R2SCAN functional, meta-GGA + 'HF' , # Hartree-Fock + 'PBE0' , # PBE0 Perdew-Burke-Ernzerhof, hybrid functional + 'EXX' , # Exact Exchange, using OEP method + 'RPA' , # Random Phase Approximation, with exact exchange +] + # Type check errors NATOM_TYPE_NOT_INTEGER_ERROR = \ "parameter natom_types must be an integer, get {} instead" @@ -270,7 +285,8 @@ "Warning: Parameter 'LATVEC' section not found in output file, trying to determine the lattice type from LATVEC_SCALE, perceed as the lattice is orthogonal" NOT_ORTHOGONALIZE_ATOMIC_ORBITALS_OVERLAP_MATRIX_NOT_PRINTED_WARNING = \ "WARNING: Since we are not to orthogonalize the atomic orbitals, the overlap matrix is not printed" - +XC_FUNCTIONAL_NOT_VALID_WARNING = \ + "WARNING: Currently, only the following XC functionals are supported: \n\t\t{}. \n\t Degrade to functional 'GGA_PBE' to generate the atomic orbitals for now.".format(', '.join(VALID_XC_FUNCTIONAL_LIST)) # PDOS output format pdos_output_header_msg = \ @@ -949,8 +965,11 @@ def get_default_generator_for_atomic_wave_function(xc_functional, psp_file_path) XC_functional: Exchange-correlation functional name psp_dir_path: Path to pseudopotential directory """ + xc_functional = "123" + if xc_functional not in VALID_XC_FUNCTIONAL_LIST: + print(XC_FUNCTIONAL_NOT_VALID_WARNING) + xc_functional = 'GGA_PBE' - assert xc_functional in ['LDA_PW', 'LDA_PZ', 'GGA_PBE', 'HF', 'PBE0', 'SCAN','RSCAN','R2SCAN'], XC_FUNCTIONAL_NOT_VALID_ERROR.format(xc_functional) assert os.path.exists(psp_file_path), "psp_file_path {} does not exist".format(psp_file_path) # Add parent directory to path in order to import AtomicDFTSolver diff --git a/utils/pdos/examples/Al_FCC/PDOS_output/PDOS.txt b/utils/pdos/examples/Al_FCC/PDOS_output/PDOS.txt index 299ef036..92034786 100644 --- a/utils/pdos/examples/Al_FCC/PDOS_output/PDOS.txt +++ b/utils/pdos/examples/Al_FCC/PDOS_output/PDOS.txt @@ -38,7 +38,7 @@ :TOT_ORBITALS: 16 :PROJ_ORBITALS: 16 :VOLUME: 448.293 -:CALCULATION_TIME: 0.383 +:CALCULATION_TIME: 0.374 :M_INDEX_SUMMED_OVER: False :ORTHOGONALIZE_ATOMIC_ORBITALS: False :ATOM_INDEX_FOR_PDOS: All