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