def _h_core(mol: gto.Mole, mm_mol: Union[None, gto.Mole]) -> Tuple[np.ndarray, np.ndarray, \
np.ndarray, Union[None, np.ndarray]]:
"""
this function returns the components of the core hamiltonian
"""
# kinetic integrals
kin = mol.intor_symmetric('int1e_kin')
# coordinates and charges of nuclei
coords = mol.atom_coords()
charges = mol.atom_charges()
# individual atomic potentials
sub_nuc = np.zeros([mol.natm, mol.nao_nr(), mol.nao_nr()], dtype=np.float64)
for k in range(mol.natm):
with mol.with_rinv_origin(coords[k]):
sub_nuc[k] = -1. * mol.intor('int1e_rinv') * charges[k]
# total nuclear potential
nuc = np.sum(sub_nuc, axis=0)
# possible mm potential
if mm_mol is not None:
mm_pot = _mm_pot(mol, mm_mol)
else:
mm_pot = None
return kin, nuc, sub_nuc, mm_pot
def prop_tot(mol: gto.Mole, mf: Union[scf.hf.SCF, dft.rks.KohnShamDFT], \
mo_coeff: Tuple[np.ndarray, np.ndarray], mo_occ: Tuple[np.ndarray, np.ndarray], \
ref: str, pop: str, prop_type: str, part: str, multiproc: bool, \
**kwargs: Any) -> Dict[str, Union[np.ndarray, List[np.ndarray]]]:
"""
this function returns atom-decomposed mean-field properties
"""
# declare nested kernel functions in global scope
global prop_atom
global prop_eda
global prop_bonds
# dft logical
dft_calc = isinstance(mf, dft.rks.KohnShamDFT)
# ao dipole integrals with specified gauge origin
if prop_type == 'dipole':
with mol.with_common_origin(kwargs['dipole_origin']):
ao_dip = mol.intor_symmetric('int1e_r', comp=3)
else:
ao_dip = None
# compute total 1-RDM (AO basis)
rdm1_tot = np.array([make_rdm1(mo_coeff[0], mo_occ[0]), make_rdm1(mo_coeff[1], mo_occ[1])])
# mol object projected into minao basis
if pop == 'iao':
pmol = lo.iao.reference_mol(mol)
else:
pmol = mol
# effective atomic charges
if 'weights' in kwargs:
weights = kwargs['weights']
charge_atom = pmol.atom_charges() - np.sum(weights[0] + weights[1], axis=0)
else:
charge_atom = 0.
# possible mm region
mm_mol = getattr(mf, 'mm_mol', None)
# cosmo/pcm solvent model
if getattr(mf, 'with_solvent', None):
e_solvent = _solvent(mol, np.sum(rdm1_tot, axis=0), mf.with_solvent)
else:
e_solvent = None
# nuclear repulsion property
if prop_type == 'energy':
prop_nuc_rep = _e_nuc(pmol, mm_mol)
elif prop_type == 'dipole':
prop_nuc_rep = _dip_nuc(pmol)
# core hamiltonian
kin, nuc, sub_nuc, mm_pot = _h_core(mol, mm_mol)
# fock potential
vj, vk = mf.get_jk(mol=mol, dm=rdm1_tot)
# calculate xc energy density
if dft_calc:
# xc-type and ao_deriv
xc_type, ao_deriv = _xc_ao_deriv(mf.xc)
# update exchange operator wrt range-separated parameter and exact exchange components
vk = _vk_dft(mol, mf, mf.xc, rdm1_tot, vk)
# ao function values on given grid
ao_value = _ao_val(mol, mf.grids.coords, ao_deriv)
# grid weights
grid_weights = mf.grids.weights
# compute all intermediates
c0_tot, c1_tot, rho_tot = _make_rho(ao_value, rdm1_tot, ref, xc_type)
# evaluate xc energy density
eps_xc = dft.libxc.eval_xc(mf.xc, rho_tot, spin=0 if isinstance(rho_tot, np.ndarray) else -1)[0]
# nlc (vv10)
if mf.nlc.upper() == 'VV10':
nlc_pars = dft.libxc.nlc_coeff(mf.xc)
ao_value_nlc = _ao_val(mol, mf.nlcgrids.coords, 1)
grid_weights_nlc = mf.nlcgrids.weights
c0_vv10, c1_vv10, rho_vv10 = _make_rho(ao_value_nlc, np.sum(rdm1_tot, axis=0), ref, 'GGA')
eps_xc_nlc = numint._vv10nlc(rho_vv10, mf.nlcgrids.coords, rho_vv10, \
grid_weights_nlc, mf.nlcgrids.coords, nlc_pars)[0]
else:
eps_xc_nlc = None
else:
xc_type = ''
grid_weights = grid_weights_nlc = None
ao_value = ao_value_nlc = None
eps_xc = eps_xc_nlc = None
c0_tot = c1_tot = None
c0_vv10 = c1_vv10 = None
# dimensions
alpha = mol.alpha
beta = mol.beta
if part == 'eda':
ao_labels = mol.ao_labels(fmt=None)
def prop_atom(atom_idx: int) -> Dict[str, Any]:
"""
this function returns atom-wise energy/dipole contributions
"""
# init results
if prop_type == 'energy':
res = {comp_key: 0. for comp_key in COMP_KEYS}
else:
res = {'el': np.zeros(3, dtype=np.float64)}
# atom-specific rdm1
rdm1_atom = np.zeros_like(rdm1_tot)
# loop over spins
for i, spin_mo in enumerate((alpha, beta)):
# loop over spin-orbitals
for m, j in enumerate(spin_mo):
# get orbital(s)
orb = mo_coeff[i][:, j].reshape(mo_coeff[i].shape[0], -1)
# orbital-specific rdm1
rdm1_orb = make_rdm1(orb, mo_occ[i][j])
# weighted contribution to rdm1_atom
rdm1_atom[i] += rdm1_orb * weights[i][m][atom_idx]
# coulumb & exchange energy associated with given atom
if prop_type == 'energy':
res['coul'] += _trace(np.sum(vj, axis=0), rdm1_atom[i], scaling = .5)
res['exch'] -= _trace(vk[i], rdm1_atom[i], scaling = .5)
# common energy contributions associated with given atom
if prop_type == 'energy':
res['kin'] += _trace(kin, np.sum(rdm1_atom, axis=0))
res['nuc_att_glob'] += _trace(sub_nuc[atom_idx], np.sum(rdm1_tot, axis=0), scaling = .5)
res['nuc_att_loc'] += _trace(nuc, np.sum(rdm1_atom, axis=0), scaling = .5)
if mm_pot is not None:
res['solvent'] += _trace(mm_pot, np.sum(rdm1_atom, axis=0))
if e_solvent is not None:
res['solvent'] += e_solvent[atom_idx]
# additional xc energy contribution
if dft_calc:
# atom-specific rho
_, _, rho_atom = _make_rho(ao_value, np.sum(rdm1_atom, axis=0), ref, xc_type)
# energy from individual atoms
res['xc'] += _e_xc(eps_xc, grid_weights, rho_atom)
# nlc (vv10)
if eps_xc_nlc is not None:
_, _, rho_atom_vv10 = _make_rho(ao_value_nlc, np.sum(rdm1_atom, axis=0), ref, 'GGA')
res['xc'] += _e_xc(eps_xc_nlc, grid_weights_nlc, rho_atom_vv10)
elif prop_type == 'dipole':
res['el'] -= _trace(ao_dip, np.sum(rdm1_atom, axis=0))
# sum up electronic contributions
if prop_type == 'energy':
for comp_key in COMP_KEYS[:-2]:
res['el'] += res[comp_key]
return res
def prop_eda(atom_idx: int) -> Dict[str, Any]:
"""
this function returns EDA energy/dipole contributions
"""
# init results
if prop_type == 'energy':
res = {comp_key: 0. for comp_key in COMP_KEYS}
else:
res = {'el': np.zeros(3, dtype=np.float64)}
# get AOs on atom k
select = np.where([atom[0] == atom_idx for atom in ao_labels])[0]
# common energy contributions associated with given atom
if prop_type == 'energy':
# loop over spins
for i, _ in enumerate((alpha, beta)):
res['coul'] += _trace(np.sum(vj, axis=0)[select], rdm1_tot[i][select], scaling = .5)
res['exch'] -= _trace(vk[i][select], rdm1_tot[i][select], scaling = .5)
res['kin'] += _trace(kin[select], np.sum(rdm1_tot, axis=0)[select])
res['nuc_att_glob'] += _trace(sub_nuc[atom_idx], np.sum(rdm1_tot, axis=0), scaling = .5)
res['nuc_att_loc'] += _trace(nuc[select], np.sum(rdm1_tot, axis=0)[select], scaling = .5)
if mm_pot is not None:
res['solvent'] += _trace(mm_pot[select], np.sum(rdm1_tot, axis=0)[select])
if e_solvent is not None:
res['solvent'] += e_solvent[atom_idx]
# additional xc energy contribution
if dft_calc:
# atom-specific rho
rho_atom = _make_rho_interm2(c0_tot[:, select], \
c1_tot if c1_tot is None else c1_tot[:, :, select], \
ao_value[:, :, select], xc_type)
# energy from individual atoms
res['xc'] += _e_xc(eps_xc, grid_weights, rho_atom)
# nlc (vv10)
if eps_xc_nlc is not None:
rho_atom_vv10 = _make_rho_interm2(c0_vv10[:, select], \
c1_vv10 if c1_vv10 is None else c1_vv10[:, :, select], \
ao_value_nlc[:, :, select], 'GGA')
res['xc'] += _e_xc(eps_xc_nlc, grid_weights_nlc, rho_atom_vv10)
elif prop_type == 'dipole':
res['el'] -= _trace(ao_dip[:, select], np.sum(rdm1_tot, axis=0)[select])
# sum up electronic contributions
if prop_type == 'energy':
for comp_key in COMP_KEYS[:-2]:
res['el'] += res[comp_key]
return res
def prop_bonds(spin_idx: int, orb_idx: int) -> Dict[str, Any]:
"""
this function returns bond-wise energy/dipole contributions
"""
# init res
if prop_type == 'energy':
res = {comp_key: 0. for comp_key in COMP_KEYS}
else:
res = {}
# get orbital(s)
orb = mo_coeff[spin_idx][:, orb_idx].reshape(mo_coeff[spin_idx].shape[0], -1)
# orbital-specific rdm1
rdm1_orb = make_rdm1(orb, mo_occ[spin_idx][orb_idx])
# total energy or dipole moment associated with given spin-orbital
if prop_type == 'energy':
res['coul'] += _trace(np.sum(vj, axis=0), rdm1_orb, scaling = .5)
res['exch'] -= _trace(vk[spin_idx], rdm1_orb, scaling = .5)
res['kin'] += _trace(kin, rdm1_orb)
res['nuc_att'] += _trace(nuc, rdm1_orb)
if mm_pot is not None:
res['solvent'] += _trace(mm_pot, rdm1_orb)
# additional xc energy contribution
if dft_calc:
# orbital-specific rho
_, _, rho_orb = _make_rho(ao_value, rdm1_orb, ref, xc_type)
# xc energy from individual orbitals
res['xc'] += _e_xc(eps_xc, grid_weights, rho_orb)
# nlc (vv10)
if eps_xc_nlc is not None:
_, _, rho_orb_vv10 = _make_rho(ao_value_nlc, rdm1_orb, ref, 'GGA')
res['xc'] += _e_xc(eps_xc_nlc, grid_weights_nlc, rho_orb_vv10)
elif prop_type == 'dipole':
res['el'] = -_trace(ao_dip, rdm1_orb)
# sum up electronic contributions
if prop_type == 'energy':
for comp_key in COMP_KEYS[:-2]:
res['el'] += res[comp_key]
return res
# perform decomposition
if part in ['atoms', 'eda']:
# init atom-specific energy or dipole arrays
if prop_type == 'energy':
prop = {comp_key: np.zeros(pmol.natm, dtype=np.float64) for comp_key in COMP_KEYS}
elif prop_type == 'dipole':
prop = {comp_key: np.zeros([pmol.natm, 3], dtype=np.float64) for comp_key in COMP_KEYS[-2:]}
# domain
domain = np.arange(pmol.natm)
# execute kernel
if multiproc:
n_threads = min(domain.size, lib.num_threads())
with mp.Pool(processes=n_threads) as pool:
res = pool.map(prop_atom if part == 'atoms' else prop_eda, domain) # type:ignore
else:
res = list(map(prop_atom if part == 'atoms' else prop_eda, domain)) # type:ignore
# collect results
for k, r in enumerate(res):
for key, val in r.items():
prop[key][k] = val
prop['struct'] = prop_nuc_rep
else: # bonds
# get rep_idx
rep_idx = kwargs['rep_idx']
# init orbital-specific energy or dipole array
if prop_type == 'energy':
prop = {comp_key: [np.zeros(len(rep_idx[0]), dtype=np.float64), np.zeros(len(rep_idx[1]), dtype=np.float64)] for comp_key in COMP_KEYS}
elif prop_type == 'dipole':
prop = {comp_key: [np.zeros([len(rep_idx[0]), 3], dtype=np.float64), np.zeros([len(rep_idx[1]), 3], dtype=np.float64)] for comp_key in COMP_KEYS}
# domain
domain = np.array([(i, j) for i, _ in enumerate((mol.alpha, mol.beta)) for j in rep_idx[i]])
# execute kernel
if multiproc:
n_threads = min(domain.size, lib.num_threads())
with mp.Pool(processes=n_threads) as pool:
res = pool.starmap(prop_bonds, domain) # type:ignore
else:
res = list(starmap(prop_bonds, domain)) # type:ignore
# collect results
for k, r in enumerate(res):
for key, val in r.items():
prop[key][0 if k < len(rep_idx[0]) else 1][k % len(rep_idx[0])] = val
prop['struct'] = prop_nuc_rep
return {**prop, 'charge_atom': charge_atom}
def main(mol: gto.Mole, decomp: DecompCls, \
mf: Union[None, scf.hf.SCF, dft.rks.KohnShamDFT], \
dipole_origin: Union[List[float], np.ndarray] = [0.] * 3) -> Dict[str, Any]:
"""
main decodense program
"""
# sanity check
sanity_check(mol, decomp)
# init time
time = MPI.Wtime()
# mf calculation
mo_coeff, mo_occ, ref = format_mf(mf)
# molecular dimensions
mol.alpha, mol.beta = dim(mol, mo_occ)
# overlap matrix
s = mol.intor_symmetric('int1e_ovlp')
# compute localized molecular orbitals
if decomp.loc != '':
mo_coeff = loc_orbs(mol, mo_coeff, s, ref, decomp.loc)
# inter-atomic distance array
dist = gto.mole.inter_distance(mol) * lib.param.BOHR
# decompose property
if decomp.part in ['atoms', 'eda']:
weights = assign_rdm1s(mol, s, mo_coeff, mo_occ, ref, decomp.pop, \
decomp.part, decomp.multiproc, decomp.verbose)[0]
decomp.res = prop_tot(mol, mf, mo_coeff, mo_occ, ref, decomp.pop, \
decomp.prop, decomp.part, decomp.multiproc, \
weights = weights, dipole_origin = dipole_origin)
elif decomp.part == 'bonds':
rep_idx, centres = assign_rdm1s(mol, s, mo_coeff, mo_occ, ref, decomp.pop, \
decomp.part, decomp.multiproc, decomp.verbose, \
thres = decomp.thres)
decomp.res = prop_tot(mol, mf, mo_coeff, mo_occ, ref, decomp.pop, \
decomp.prop, decomp.part, decomp.multiproc, \
rep_idx = rep_idx, dipole_origin = dipole_origin)
# write cube files
if decomp.cube:
write_cube(mol, decomp.part, mo_coeff, mo_occ, \
weights if decomp.part == 'atoms' else None, \
rep_idx if decomp.part == 'bonds' else None)
# determine spin
decomp.res['ss'], decomp.res['s'] = scf.uhf.spin_square((mo_coeff[0][:, mol.alpha], \
mo_coeff[1][:, mol.beta]), s)
# collect time
decomp.res['time'] = MPI.Wtime() - time
# collect centres & dist
if decomp.part == 'bonds':
decomp.res['centres'] = centres
decomp.res['dist'] = dist
return decomp.res