def _dip_nuc(mol: gto.Mole) -> np.ndarray: """ this function returns the nuclear contribution to the molecular dipole moment """ # coordinates and charges of nuclei coords = mol.atom_coords() charges = mol.atom_charges() return contract('i,ix->ix', charges, coords)
def setUpClass(cls): cls.mol = mol = Mole() mol.atom = [ [8, (0., 0., 0.)], [1, (0., -0.757, 0.587)], [1, (0., 0.757, 0.587)]] mol.basis = 'cc-pvdz' mol.verbose = 5 mol.build() cls.model_rks = model_rks = RKS(mol) # model_rks.xc = 'hf' model_rks.kernel() cls.td_model_rks = td_model_rks = TDRKS(model_rks) td_model_rks.nroots = 4 td_model_rks.kernel() cls.td_proxy_model = td_proxy_model = TDProxy(model_rks, "dft") td_proxy_model.nroots = td_model_rks.nroots td_proxy_model.kernel() testing.assert_allclose(td_model_rks.e, td_proxy_model.e, atol=1e-6) cls.gw = gw = GW(model_rks, td_model_rks) gw.kernel()
def _e_nuc(mol: gto.Mole, mm_mol: Union[None, gto.Mole]) -> np.ndarray: """ this function returns the nuclear repulsion energy """ # coordinates and charges of nuclei coords = mol.atom_coords() charges = mol.atom_charges() # internuclear distances (with self-repulsion removed) dist = gto.inter_distance(mol) dist[np.diag_indices_from(dist)] = 1e200 e_nuc = contract('i,ij,j->i', charges, 1. / dist, charges) * .5 # possible interaction with mm sites if mm_mol is not None: mm_coords = mm_mol.atom_coords() mm_charges = mm_mol.atom_charges() for j in range(mol.natm): q2, r2 = charges[j], coords[j] r = lib.norm(r2 - mm_coords, axis=1) e_nuc[j] += q2 * np.sum(mm_charges / r) return e_nuc
def _mm_pot(mol: gto.Mole, mm_mol: gto.Mole) -> np.ndarray: """ this function returns the full mm potential (adapted from: qmmm/itrf.py:get_hcore() in PySCF) """ # settings coords = mm_mol.atom_coords() charges = mm_mol.atom_charges() blksize = BLKSIZE # integrals intor = 'int3c2e_cart' if mol.cart else 'int3c2e_sph' cintopt = gto.moleintor.make_cintopt(mol._atm, mol._bas, mol._env, intor) # compute interaction potential mm_pot = 0 for i0, i1 in lib.prange(0, charges.size, blksize): fakemol = gto.fakemol_for_charges(coords[i0:i1]) j3c = df.incore.aux_e2(mol, fakemol, intor=intor, aosym='s2ij', cintopt=cintopt) mm_pot += np.einsum('xk,k->x', j3c, -charges[i0:i1]) mm_pot = lib.unpack_tril(mm_pot) return mm_pot
def make_molecule(self): mol = Mole() mol.atom = """ O 0 0 0 H 1 0 0 H 0 1 0 """ mol.basis = "sto-3g" mol.build() return mol
def setUpClass(cls): cls.mol = mol = Mole() mol.atom = [ [8, (0., 0., 0.)], [1, (0., -0.757, 0.587)], [1, (0., 0.757, 0.587)]] mol.basis = 'cc-pvdz' mol.verbose = 5 mol.build() cls.model_rhf = model_rhf = RHF(mol) model_rhf.kernel() cls.td_model_rhf = td_model_rhf = TDHF(model_rhf) td_model_rhf.nroots = 4 td_model_rhf.kernel() cls.ref_m = retrieve_m(td_model_rhf)
def setUpClass(cls): cls.mol = mol = Mole() mol.atom = [ [8, (0., 0., 0.)], [1, (0., -0.757, 0.587)], [1, (0., 0.757, 0.587)]] mol.basis = 'cc-pvdz' mol.verbose = 5 mol.build() cls.model_rks = model_rks = RKS(mol) model_rks.kernel() cls.td_model_rks = td_model_rks = TDDFT(model_rks) td_model_rks.nroots = 4 td_model_rks.kernel() e = model_rks.mo_energy nocc = int(sum(model_rks.mo_occ) // 2) cls.ref_m = mk_make_canonic(retrieve_m(td_model_rks), e[:nocc], e[nocc:])
def setUpModule(): global mol, mf mol = Mole() mol.atom = ''' C 0.681068338 0.605116159 0.307300799 C -0.733665805 0.654940451 -0.299036438 C -1.523996730 -0.592207689 0.138683275 H 0.609941801 0.564304456 1.384183068 H 1.228991034 1.489024155 0.015946420 H -1.242251083 1.542928348 0.046243898 H -0.662968178 0.676527364 -1.376503770 H -0.838473936 -1.344174292 0.500629028 H -2.075136399 -0.983173387 -0.703807608 H -2.212637905 -0.323898759 0.926200671 O 1.368219958 -0.565620846 -0.173113101 H 2.250134219 -0.596689848 0.204857736 ''' mol.basis = 'STO-3G' mol.verbose = 0 mol.output = None mol.build() mf = RHF(mol) mf.conv_tol = 1.0e-12 mf.kernel()
def setUpClass(cls): cls.mol = mol = Mole() mol.atom = [ [8, (0., 0., 0.)], [1, (0., -0.757, 0.587)], [1, (0., 0.757, 0.587)]] mol.basis = 'cc-pvdz' mol.verbose = 5 mol.build() cls.model_rhf = model_rhf = RHF(mol) model_rhf.kernel() cls.model_rks = model_rks = RKS(mol) model_rks.xc = 'hf' model_rks.kernel() testing.assert_allclose(model_rhf.mo_energy, model_rks.mo_energy) assert_vectors_close(model_rhf.mo_coeff.T, model_rks.mo_coeff.T) model_rks.mo_coeff = model_rhf.mo_coeff cls.td_model_rks = td_model_rks = TDRHF(model_rks) td_model_rks.nroots = 4 td_model_rks.kernel() cls.td_model_rhf_slow = td_model_rhf_slow = TDRHF_slow(model_rhf) td_model_rhf_slow.nroots = td_model_rks.nroots td_model_rhf_slow.kernel() cls.td_proxy_model = td_proxy_model = TDProxy(model_rks, "dft") td_proxy_model.nroots = td_model_rks.nroots td_proxy_model.kernel() cls.gw = gw = GW(model_rks, td_model_rks) gw.kernel()
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
# # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import unittest import numpy from pyscf.gto import Mole from pyscf.scf import RHF from pyscf.lo.cholesky import cholesky_mos mol = Mole() mol.atom = ''' C 0.681068338 0.605116159 0.307300799 C -0.733665805 0.654940451 -0.299036438 C -1.523996730 -0.592207689 0.138683275 H 0.609941801 0.564304456 1.384183068 H 1.228991034 1.489024155 0.015946420 H -1.242251083 1.542928348 0.046243898 H -0.662968178 0.676527364 -1.376503770 H -0.838473936 -1.344174292 0.500629028 H -2.075136399 -0.983173387 -0.703807608 H -2.212637905 -0.323898759 0.926200671 O 1.368219958 -0.565620846 -0.173113101 H 2.250134219 -0.596689848 0.204857736 ''' mol.basis = 'STO-3G'
#!/usr/bin/env python # ''' DF-MP2 natural orbitals for the allyl radical ''' from pyscf.gto import Mole from pyscf.scf import UHF from pyscf.tools import molden from pyscf.mp.dfump2_native import DFMP2 mol = Mole() mol.atom = ''' C -1.1528 -0.1151 -0.4645 C 0.2300 -0.1171 -0.3508 C 0.9378 0.2246 0.7924 H 0.4206 0.5272 1.7055 H 2.0270 0.2021 0.8159 H -1.6484 -0.3950 -1.3937 H -1.7866 0.1687 0.3784 H 0.8086 -0.4120 -1.2337 ''' mol.basis = 'def2-TZVP' mol.spin = 1 mol.build() mf = UHF(mol).run() # MP2 natural occupation numbers and natural orbitals natocc, natorb = DFMP2(mf).make_natorbs() # store the natural orbitals in a molden file
#!/usr/bin/env python # ''' Calculating densities with DF-MP2, demonstrated for the dipole moment of CH3Cl. ''' from numpy.linalg import norm from pyscf.gto import Mole from pyscf.scf import RHF from pyscf.mp.dfmp2_native import DFMP2 mol = Mole() mol.atom = ''' C 0.000000 0.000000 0.000000 Cl 0.000000 0.000000 1.785000 H 1.019297 0.000000 -0.386177 H -0.509649 0.882737 -0.386177 H -0.509649 -0.882737 -0.386177 ''' mol.basis = 'aug-cc-pVTZ' mol.build() mf = RHF(mol).run() pt = DFMP2(mf).run() # The unrelaxed density always has got natural occupation numbers between 2 and 0. # However, it is inaccurate for properties. dm_ur = pt.make_rdm1_unrelaxed(ao_repr=True) # The relaxed density is more accurate for properties when MP2 is well-behaved, # whereas the natural occupation numbers can be above 2 or below 0 for ill-behaved systems.
def nr_uks(self, mol: gto.Mole, grids: dft.Grids, xc_code: str, dms: Union[Sequence[np.ndarray], Sequence[Sequence[np.ndarray]]], relativity: int = 0, hermi: int = 0, max_memory: float = 20000, verbose=None) -> Tuple[np.ndarray, float, np.ndarray]: """Calculates UKS XC functional and potential matrix on a given grid. Args: mol: PySCF molecule. grids: grid on which to evaluate the functional. xc_code: XC code. Unused. NeuralNumInt hard codes the XC functional based upon the functional argument given to the constructor. dms: the density matrix or sequence of density matrices for each spin channel. Multiple density matrices for each spin channel are not currently supported. Each density matrix is shape (nao, nao), where nao is the number of atomic orbitals. relativity: Unused. (pyscf.dft.numint.NumInt.nr_rks does not currently use this argument.) hermi: 0 if the density matrix is Hermitian, 1 if the density matrix is non-Hermitian. max_memory: the maximum cache to use, in MB. verbose: verbosity level. Unused. (PySCF currently does not handle the verbosity level passed in here.) Returns: nelec, excsum, vmat, where nelec is the number of alpha, beta electrons obtained by numerical integration of the density matrix as an array of size 2. excsum is the functional's XC energy. vmat is the functional's XC potential matrix, shape (2, nao, nao), where vmat[0] and vmat[1] are the potential matrices for the alpha and beta spin channels respectively. Raises: NotImplementedError: if multiple density matrices for each spin channel are supplied. """ # Wrap nr_uks so we can store internal variables required to evaluate the # contribution to the XC potential from local Hartree-Fock features. # See pyscf.dft.numint.nr_uks for more details. if isinstance(dms, np.ndarray) and dms.ndim == 2: # RHF DM ndms = _get_number_of_density_matrices(dms) else: ndms = _get_number_of_density_matrices(dms[0]) if ndms > 1: raise NotImplementedError( 'NeuralNumInt does not support multiple density matrices. ' 'Only ground state DFT calculations are currently implemented.') nao = mol.nao_nr() self._vmat_hf = np.zeros((2, nao, nao)) self._system_state = _SystemState(mol=mol, dms=dms) nelec, excsum, vmat = super().nr_uks( mol=mol, grids=grids, xc_code=xc_code, dms=dms, relativity=relativity, hermi=hermi, max_memory=max_memory, verbose=verbose) vmat[0] += self._vmat_hf[0] + self._vmat_hf[0].T vmat[1] += self._vmat_hf[1] + self._vmat_hf[1].T # Clear internal state to prevent accidental re-use. self._system_state = None self._grid_state = None self._vmat_hf = None return nelec, excsum, vmat
https://doi.org/10.1063/1.2360264 This example performs separate (split) localizations of 1) the occupied orbitals and 2) the virtual orbitals, and produces a molden output containing both orbital sets. ''' import numpy from pyscf.gto import Mole from pyscf.scf import RHF from pyscf.lo import cholesky_mos from pyscf.tools import molden # Set up of the molecule (C3H7OH) mol = Mole() mol.atom = ''' C 0.681068338 0.605116159 0.307300799 C -0.733665805 0.654940451 -0.299036438 C -1.523996730 -0.592207689 0.138683275 H 0.609941801 0.564304456 1.384183068 H 1.228991034 1.489024155 0.015946420 H -1.242251083 1.542928348 0.046243898 H -0.662968178 0.676527364 -1.376503770 H -0.838473936 -1.344174292 0.500629028 H -2.075136399 -0.983173387 -0.703807608 H -2.212637905 -0.323898759 0.926200671 O 1.368219958 -0.565620846 -0.173113101 H 2.250134219 -0.596689848 0.204857736 ''' mol.basis = 'def2-SVP'
# Superfluous columns are cropped out. P = np.zeros((nao, nao)) P[piv, np.arange(nao)] = 1 mo_loc = np.dot(P, L[:, :nmo]) return mo_loc if __name__ == "__main__": import numpy from pyscf.gto import Mole from pyscf.scf import RHF from pyscf.tools.mo_mapping import mo_comps mol = Mole() mol.atom = ''' C 0.681068338 0.605116159 0.307300799 C -0.733665805 0.654940451 -0.299036438 C -1.523996730 -0.592207689 0.138683275 H 0.609941801 0.564304456 1.384183068 H 1.228991034 1.489024155 0.015946420 H -1.242251083 1.542928348 0.046243898 H -0.662968178 0.676527364 -1.376503770 H -0.838473936 -1.344174292 0.500629028 H -2.075136399 -0.983173387 -0.703807608 H -2.212637905 -0.323898759 0.926200671 O 1.368219958 -0.565620846 -0.173113101 H 2.250134219 -0.596689848 0.204857736 ''' mol.basis = 'STO-3G'
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