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 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