# ~ Overall Embedding object
from abc import ABC, abstractmethod
from asiembedding.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 asiembedding.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()