Source code for ASI_Embedding.embedding

# ~ Overall Embedding object
from abc import ABC, abstractmethod
from ASI_Embedding.parallel_utils import root_print
import time
import numpy as np

[docs] class EmbeddingBase(ABC): def __init__(self, atoms, embed_mask, calc_base_ll=None, calc_base_hl=None): import os self.asi_lib_path = os.environ['ASI_LIB_PATH'] self.embed_mask = embed_mask self.calculator_ll = calc_base_ll self.calculator_hl = calc_base_hl @property def scf_methods(self): return self._scf_methods @scf_methods.setter def scf_methods(self, val): if isinstance(self.embed_mask, int): assert len(val) == 2, \ "Invalid number of methods for given n_layers." elif isinstance(self.embed_mask, list): assert len(val) == len(set(self.embed_mask)), \ "Invalid number of methods for given n_layers." self._scf_methods = [] for scf in val: self._scf_methods.append(scf)
[docs] def set_layer(self, atoms, layer_name, calc, embed_mask, ghosts=0, no_scf=False): "Initialises the AtomsEmbed methods for a given method" from .atoms_embedding_asi import AtomsEmbed layer = AtomsEmbed(atoms, calc, embed_mask, outdir=layer_name, ghosts=ghosts, no_scf=no_scf) setattr(self, layer_name, layer)
@property def nlayers(self): return self._nlayers @nlayers.setter def nlayers(self, val): assert val == 2, \ "Only two layers currently valid for projection embedding." return self._nlayers
[docs] def select_atoms_basis_truncation(self, thresh): """_summary_ Returns a list of corresponding atoms for which the total contribution of an atoms constituent basis functions to the total charge of subsystem A, q^{A}: q^{A}_{/mu, /nu} = /gamma^{A}_{/mu, /nu} S_{/mu, /nu} exceeds the threshold, thresh: thresh < q_{/mu, /nu} This is unfortunately just Mulliken analysis. Args: thresh (float): _description_ Returns: active_mask (list): True/False mask for """ basis_charge = np.diag(self.AB_LL.density_matrices_out[0] @ self.AB_LL.overlap) atomic_charge = np.zeros(len(self.AB_LL.atoms)) for idx, charge in enumerate(basis_charge): atomic_charge[self.AB_LL.basis_atoms[idx]] += abs(charge) truncated_atom_list = [ charge > thresh for charge in atomic_charge ] return truncated_atom_list
[docs] def set_truncation_defaults(self, truncated_atom_list): """_summary_ Necessary to maintain consistency in values needed for matrix truncation /expansion between the AtomsEmbed objects (eg., self.basis_atoms, self.n_atoms). Failure to do so leads to unexpected behaviour. which Args: truncated_atom_list (list): A boolean mask asserting which atoms """ from ASI_Embedding.basis_info import Basis_info # Establish mapping corresponding to each new atom from the # truncated matrix active_atoms = np.array([ idx for idx, maskval in enumerate(truncated_atom_list) if maskval ]) # Remove non-active atoms from basis_atoms to form a truncated # analogue trunc_basis_atoms = np.array([atom for atom in self.AB_LL.basis_atoms if atom in active_atoms]) # Count number of basis functions included in truncation calculations new_nbasis = 0 for atom in self.AB_LL.basis_atoms: if truncated_atom_list[atom]: new_nbasis +=1 root_print(f" " ) root_print( f" ----------- Performing Truncation --------- " ) root_print(f" ") root_print(f" Number of atoms before truncation: {len(truncated_atom_list)}") root_print(f" Number of atoms after truncation: {len(active_atoms)}") root_print(f" ") root_print(f" Number of basis functions before truncation: {len(self.AB_LL.basis_atoms)}") root_print(f" Number of basis functions after truncation: {new_nbasis}") root_print(f" " ) root_print( f" ------------------------------------------- " ) root_print(f" " ) BasisInfo = Basis_info() BasisInfo.full_natoms = len(truncated_atom_list) BasisInfo.trunc_natoms = len(active_atoms) BasisInfo.active_atoms = active_atoms BasisInfo.full_basis_atoms = self.AB_LL.basis_atoms BasisInfo.trunc_basis_atoms = trunc_basis_atoms BasisInfo.full_nbasis = self.AB_LL.n_basis BasisInfo.trunc_nbasis = new_nbasis BasisInfo.set_basis_atom_indexes() self.AB_LL.truncate = False self.AB_LL.basis_info = BasisInfo self.A_LL.truncate = True self.A_LL.basis_info = BasisInfo self.A_HL.truncate = True self.A_HL.basis_info = BasisInfo self.A_HL_PP.truncate = True self.A_HL_PP.basis_info = BasisInfo self.AB_LL_PP.truncate = False self.AB_LL_PP.basis_info = BasisInfo
[docs] def calc_subsys_pop(self, overlap_matrix, density_matrix): """ Summary Calculates the overall electron population of a given subsystem through the following relation: P_{pop} = tr[S^{AB}/gamma] where S^{AB} is the overlap matrix for the supermolecular system and /gamma is the density matrix of a given subsystem. Args: overlap_matrix (np.ndarray): Supersystem overlap matrix in AO basis. density_matrix (np.ndarray): Subsystem density matrix in AO basis, Returns: population (int): Overall electronic population of subsystem """ population = np.trace(overlap_matrix @ (density_matrix)) return population
[docs] @abstractmethod def run(self): pass
[docs] class ProjectionEmbedding(EmbeddingBase): def __init__(self, atoms, embed_mask, calc_base_ll, calc_base_hl, frag_charge=0, post_scf=None, mu_val=1e+06, truncate_basis=False): """_summary_ A class controlling the interaction between two subsystems calculated at two levels of theory (low-level and high-level) with the Projection-Based Embedding (PbE) scheme of Manby et al.[1]. [1] Manby, F. R.; Stella, M.; Goodpaster, J. D.; Miller, T. F. I. A Simple, Exact Density-Functional-Theory Embedding Scheme. J. Chem. Theory Comput. 2012, 8 (8), 2564–2568. Methods _______ Args: atoms (ASE Atoms Object): Input ASE Atoms object used to pass structural information embed_mask (int OR list): _description_ calc_base_ll (ASE Calculator): _description_ calc_base_hl (ASE Calculator): _description_ frag_charge (int, optional): _description_. Defaults to 0. post_scf (_type_, optional): _description_. Defaults to None. mu_val (_type_, optional): _description_. Defaults to 1e+06 Ha. truncate_basis (bool, optional): _description_. Defaults to 1e+06 Ha. """ from copy import copy, deepcopy from mpi4py import MPI self.calc_names = ["AB_LL","A_LL","A_HL","A_HL_PP","AB_LL_PP"] super(ProjectionEmbedding, self).__init__(atoms, embed_mask, calc_base_ll, calc_base_hl) low_level_calculator_1 = deepcopy(self.calculator_ll) low_level_calculator_2 = deepcopy(self.calculator_ll) low_level_calculator_3 = deepcopy(self.calculator_ll) high_level_calculator_1 = deepcopy(self.calculator_hl) high_level_calculator_2 = deepcopy(self.calculator_hl) low_level_calculator_1.set(qm_embedding_calc = 1) self.set_layer(atoms, self.calc_names[0], low_level_calculator_1, embed_mask, ghosts=0, no_scf=False) low_level_calculator_3.set(qm_embedding_calc = 2) low_level_calculator_3.set(charge_mix_param = 0.) self.set_layer(atoms, self.calc_names[4], low_level_calculator_3, embed_mask, ghosts=0, no_scf=False) low_level_calculator_2.set(qm_embedding_calc = 2) low_level_calculator_2.set(charge_mix_param = 0.) low_level_calculator_2.set(charge = frag_charge) self.set_layer(atoms, self.calc_names[1], low_level_calculator_2, embed_mask, ghosts=2, no_scf=False) high_level_calculator_1.set(qm_embedding_calc = 3) high_level_calculator_1.set(charge = frag_charge) self.set_layer(atoms, self.calc_names[2], high_level_calculator_1, embed_mask, ghosts=2, no_scf=False) high_level_calculator_2.set(qm_embedding_calc = 2) high_level_calculator_2.set(charge_mix_param = 0.) high_level_calculator_2.set(charge = frag_charge) self.set_layer(atoms, self.calc_names[3], high_level_calculator_2, embed_mask, ghosts=2, no_scf=False) self.mu_val = mu_val self.rank = MPI.COMM_WORLD.Get_rank() self.ntasks = MPI.COMM_WORLD.Get_size() self.truncate_basis = truncate_basis
[docs] def calculate_levelshift_projector(self): """_summary_ Calculate the level-shift based projection operator from Manby et al.[1]: P^{B} = /mu S^{AB} D^{B} S^{AB} where S^{AB} is the overlap matrix for the supermolecular system, and the density matrix for subsystem B. [1] Manby, F. R.; Stella, M.; Goodpaster, J. D.; Miller, T. F. I. A Simple, Exact Density-Functional-Theory Embedding Scheme. J. Chem. Theory Comput. 2012, 8 (8), 2564–2568. """ self.P_b = self.mu_val * (self.AB_LL.overlap @ self.AB_LL.density_matrices_out[1] @ self.AB_LL.overlap)
[docs] def calculate_huzinaga_projector(self): self.P_b = self.AB_LL.hamiltonian_total @ self.AB_LL.density_matrices_out[1] @ self.AB_LL.overlap self.P_b = -0.5*(self.P_b + self.AB_LL.overlap @ self.AB_LL.density_matrices_out[1] @ self.AB_LL.hamiltonian_total)
[docs] def run(self): """ Summary The primary driver routine for performing QM-in-QM with a Projection-based embedding scheme. This scheme draws upon the work of Manby et al. [1, 2]. The embedding scheme uses the following total energy expression... Importing and exporting of density matrices and hamiltonians is performed with the ASI package [3]. The workflow operates as follows: 1) Calculate the KS-DFT energy of the combined subsystems A+B. Localised density matrices, /gamma^{A} and /gamma^{B} 2a) Extract the localised density matrices 2b) (Optional) Select atoms in subsystem B that contribute significantly to subsystem A (threshold 0.5 |e|) via Mulliken analysis: q^{A}_{/mu, /nu} = /gamma^{A}_{/mu, /nu} S_{/mu, /nu} Basis functions of said atoms within calculations of the embedded subsystem will be included as ghost atoms. Other basis functions will be removed (i.e., associated atomic centers not included in the QM calculation, and associated rows and columns in intermediate hamiltonians and density matrices deleted). 2) Calculate the total energy for subsystem A with the density matrix, /gamma^{A} 3) (For users of LaTeX, I am aware that a forward slash is used in place of the traditional backward slash for mathematical symbols - unfortunately using backslashes in these comment blocks produces ugly warnings within the comment blocks.) ... (1) Manby, F. R.; Stella, M.; Goodpaster, J. D.; Miller, T. F. I. A Simple, Exact Density-Functional-Theory Embedding Scheme. J. Chem. Theory Comput. 2012, 8 (8), 2564–2568. (2) Lee, S. J. R.; Welborn, M.; Manby, F. R.; Miller, T. F. Projection-Based Wavefunction-in-DFT Embedding. Acc. Chem. Res. 2019, 52 (5), 1359–1368. (3) TODO: REF """ import numpy as np root_print("Embedding calculation begun...") ''' Performs a single-point energy evaluation for a system composed of A and B. Returns localised density matrices for subsystems A and B, and the two-electron components of the hamiltonian (combined with nuclear-electron potential). ''' start = time.time() self.AB_LL.run() end = time.time() self.time_ab_lowlevel = end - start ''' Initialises the density matrix for subsystem A, and calculates the hamiltonian components for subsystem A at the low-level reference. ''' if self.truncate_basis: basis_mask = self.select_atoms_basis_truncation(0.0001) self.set_truncation_defaults(basis_mask) #else: # self.set_basis_info() self.A_LL.density_matrix_in = self.AB_LL.density_matrices_out[0] start = time.time() self.A_LL.run() end = time.time() self.time_a_lowlevel = end - start ''' Calculates the electron count for the combined (A+B) and separated subsystems (A and B). ''' self.AB_pop = self.calc_subsys_pop(self.AB_LL.overlap, (self.AB_LL.density_matrices_out[0] +self.AB_LL.density_matrices_out[1])) self.A_pop = self.calc_subsys_pop(self.AB_LL.overlap, self.AB_LL.density_matrices_out[0]) self.B_pop = self.calc_subsys_pop(self.AB_LL.overlap, self.AB_LL.density_matrices_out[1]) root_print(f" Population of Subsystem AB: {self.AB_pop}") root_print(f" Population of Subsystem A: {self.A_pop}") root_print(f" Population of Subsystem B: {self.B_pop}") ''' Initialises the density matrix for subsystem A, and calculated the hamiltonian components for subsystem A at the low-level reference. ''' self.calculate_levelshift_projector() #self.calculate_huzinaga_projector() ''' Calculate density matrix for subsystem A at the higher level of theory. Two terms are added to the hamiltonian matrix of the embedded subsystem to form the full embedded Fock-matrix, F^{A}: F^{A} = h^{core} + g^{high-level}[A] + v_{emb}^[A, B] + /mu P^{B} 1) v_{emb}^[A, B], the embedding potential, which introduces the Hartree and exchange-correlation potentials of the environment (in the case of FHI-aims, this includes the full electrostatic potential, i.e., the Hartree and nuclear-electron potentials of subsystem B acting on subsystem A). 2) The level-shift operator, /mu P^{B}, which orthogonalises the basis functions associated with the environment (subsystem B) from the embedded subsystem by adding a large energy penalty to hamiltonian components associated with subsystem B. Registered callbacks in ASI add the above components to the Fock-matrix at every SCF iteration. ''' self.A_HL.density_matrix_in = \ self.AB_LL.density_matrices_out[0] self.A_HL.fock_embedding_matrix = \ self.AB_LL.hamiltonian_electrostatic - \ self.A_LL.hamiltonian_electrostatic + self.P_b start = time.time() self.A_HL.run() end = time.time() self.time_a_highlevel = end - start ''' Calculate the total energy of the embedded subsystem A at the high level of theory without the associated embedding potential. ''' self.A_HL_PP.density_matrix_in = self.A_HL.density_matrices_out[0] start = time.time() self.A_HL_PP.run(ev_corr_scf=True) subsys_A_highlvl_totalen = self.A_HL_PP.ev_corr_total_energy end = time.time() self.time_a_highlevel_pp = end - start ''' A terrible cludge which requires improvement. ''' # Re-normalising charge for differing atomic solvers (bad cludge) root_print(f" Normalizing density matrix from high-level reference...") self.A_HL_pop = self.calc_subsys_pop(self.AB_LL.overlap, self.A_HL.density_matrices_out[0]) root_print(f" Population of Subystem A^[HL]: {self.A_HL_pop}") self.charge_renorm = (self.A_pop/self.A_HL_pop) root_print(f" Population of Subystem A^[HL] (post-norm): {self.calc_subsys_pop(self.AB_LL.overlap, self.charge_renorm*self.A_HL. density_matrices_out[0])}") self.charge_renorm = 1.0 # Calculate A low-level reference energy self.A_LL.density_matrix_in = self.charge_renorm * self.A_HL.density_matrices_out[0] self.A_LL.run(ev_corr_scf=True) start = time.time() subsys_A_lowlvl_totalen = self.A_LL.ev_corr_total_energy end = time.time() self.time_a_lowlevel_pp = end - start # Calculate AB low-level reference energy self.AB_LL_PP.density_matrix_in = (self.charge_renorm * self.A_HL.density_matrices_out[0]) + self.AB_LL.density_matrices_out[1] start = time.time() self.AB_LL_PP.run(ev_corr_scf=True) subsys_AB_lowlvl_totalen = self.AB_LL_PP.ev_corr_total_energy end = time.time() self.time_ab_lowlevel_pp = end - start # Calculate projected density correction to total energy self.PB_corr = (np.trace(self.P_b @ self.A_HL.density_matrices_out[0]) * 27.211384500) self.DFT_AinB_total_energy = subsys_A_highlvl_totalen - \ subsys_A_lowlvl_totalen + subsys_AB_lowlvl_totalen + self.PB_corr root_print( f" ----------- FINAL OUTPUTS --------- " ) root_print(f" ") root_print(f" Population Information:") root_print(f" Population of Subsystem AB: {self.AB_pop}") root_print(f" Population of Subsystem A: {self.A_pop}") root_print(f" Population of Subsystem B: {self.B_pop}") root_print(f" ") root_print(f" Intermediate Information:") root_print(f" WARNING: These are not faithful, ground-state KS total energies - ") root_print(f" In the case of low-level references, they are calculated using the ") root_print(f" density components of the high-level energy reference for fragment A. ") root_print(f" Do not naively use these energies unless you are comfortable with ") root_print(f" their true definition. ") root_print(f" Total Energy (A+B Low-Level): {subsys_AB_lowlvl_totalen} eV" ) root_print(f" Total Energy (A Low-Level): {subsys_A_lowlvl_totalen} eV" ) root_print(f" Total Energy (A High-Level): {subsys_A_highlvl_totalen} eV" ) root_print(f" Projection operator energy correction DM^(A_HL) @ Pb: {self.PB_corr} eV" ) root_print(f" " ) root_print(f" Final Energies Information:") root_print(f" Final total energy (Uncorrected): {self.DFT_AinB_total_energy - self.PB_corr} eV" ) root_print(f" Final total energy (Projection Corrected): {self.DFT_AinB_total_energy} eV" ) root_print(f" " ) root_print(f" -----------======================--------- " ) root_print(f" " )
# def test_spade(self): # # import copy # # total_c_mat = (np.loadtxt('./AB_LL/KS_region_1_1.dat') + np.loadtxt('./#AB_LL/KS_region_2_1.dat'))[:,:23] # # a_c_mat = copy.deepcopy(total_c_mat) # #a_c_mat = np.sqrt(self.AB_LL.overlap) @ a_c_mat # # active_atoms = [ idx for idx, maskval in enumerate(self.AB_LL.#embed_mask) if maskval==1 ] # # mask = [ bas_at in active_atoms for bas_at in self.AB_LL.basis_atoms ] # # for i in range(total_c_mat.shape[1]): # a_c_mat[:,i] = np.where( mask, a_c_mat[:,i], np.zeros(self.AB_LL.#n_basis) ) # # u, s, v = np.linalg.svd(a_c_mat) # # import matplotlib.pyplot as plt # plt.matshow(a_c_mat) # plt.matshow(total_c_mat @ v.T) # plt.show()