def __init__(self, atombases: List[AtomCGTOBasis], spherical: bool = True, df: Optional[DensityFitInfo] = None, efield: Optional[torch.Tensor] = None, cache: Optional[Cache] = None) -> None: self.atombases = atombases self.spherical = spherical self.libcint_wrapper = intor.LibcintWrapper(atombases, spherical) self.dtype = self.libcint_wrapper.dtype self.device = self.libcint_wrapper.device if df is None: self._df: Optional[DFMol] = None else: self._df = DFMol(df, wrapper=self.libcint_wrapper) self._efield = efield if efield is not None: assert efield.ndim == 1 and efield.numel() == 3 self.is_grid_set = False self.is_ao_set = False self.is_grad_ao_set = False self.is_lapl_ao_set = False self.xc: Optional[BaseXC] = None self.xcfamily = 1 self.is_built = False # initialize cache self._cache = cache if cache is not None else Cache.get_dummy() self._cache.add_cacheable_params( ["overlap", "kinetic", "nuclattr", "efield"]) if self._df is None: self._cache.add_cacheable_params(["elrep"])
def build(self) -> BaseDF: self._is_built = True # construct the matrix used to calculate the electron repulsion for # density fitting method method = self.dfinfo.method auxbasiswrapper = intor.LibcintWrapper(self.dfinfo.auxbases, spherical=self.wrapper.spherical) basisw, auxbw = intor.LibcintWrapper.concatenate(self.wrapper, auxbasiswrapper) if method == "coulomb": logger.log("Calculating the 2e2c integrals") j2c = intor.coul2c(auxbw) # (nxao, nxao) logger.log("Calculating the 2e3c integrals") j3c = intor.coul3c(basisw, other1=basisw, other2=auxbw) # (nao, nao, nxao) elif method == "overlap": j2c = intor.overlap(auxbw) # (nxao, nxao) # TODO: implement overlap3c raise NotImplementedError( "Density fitting with overlap minimization is not implemented") self._j2c = j2c # (nxao, nxao) self._j3c = j3c # (nao, nao, nxao) logger.log("Precompute matrix for density fittings") self._inv_j2c = torch.inverse(j2c) # if the memory is too big, then don't precompute elmat if get_memory(j3c) > config.THRESHOLD_MEMORY: self._precompute_elmat = False else: self._precompute_elmat = True self._el_mat = torch.matmul(j3c, self._inv_j2c) # (nao, nao, nxao) logger.log("Density fitting done") return self
def build(self) -> BaseDF: self._is_built = True # construct the matrix used to calculate the electron repulsion for # density fitting method method = self.dfinfo.method auxbasiswrapper = intor.LibcintWrapper( self.dfinfo.auxbases, spherical=self.wrapper.spherical) basisw, auxbw = intor.LibcintWrapper.concatenate( self.wrapper, auxbasiswrapper) if method == "coulomb": j2c = intor.coul2c(auxbw) # (nxao, nxao) j3c = intor.coul3c(basisw, other1=basisw, other2=auxbw) # (nao, nao, nxao) elif method == "overlap": j2c = intor.overlap(auxbw) # (nxao, nxao) # TODO: implement overlap3c raise NotImplementedError( "Density fitting with overlap minimization is not implemented") self._j2c = j2c # (nxao, nxao) self._j3c = j3c # (nao, nao, nxao) self._inv_j2c = torch.inverse(j2c) self._el_mat = torch.matmul(j3c, self._inv_j2c) # (nao, nao, nxao) return self
def __init__( self, atombases: List[AtomCGTOBasis], latt: Lattice, *, kpts: Optional[torch.Tensor] = None, wkpts: Optional[ torch.Tensor] = None, # weights of k-points to get the density spherical: bool = True, df: Optional[DensityFitInfo] = None, lattsum_opt: Optional[Union[intor.PBCIntOption, Dict]] = None, cache: Optional[Cache] = None) -> None: self._atombases = atombases self._spherical = spherical self._lattice = latt # alpha for the compensating charge # TODO: calculate eta properly or put it in lattsum_opt self._eta = 0.2 self._eta = 0.46213127322256375 # temporary to follow pyscf.df # lattice sum integral options self._lattsum_opt = intor.PBCIntOption.get_default(lattsum_opt) self._basiswrapper = intor.LibcintWrapper(atombases, spherical=spherical, lattice=latt) self.dtype = self._basiswrapper.dtype self.cdtype = get_complex_dtype(self.dtype) self.device = self._basiswrapper.device # set the default k-points and their weights self._kpts = kpts if kpts is not None else \ torch.zeros((1, 3), dtype=self.dtype, device=self.device) nkpts = self._kpts.shape[0] # default weights are just 1/nkpts (nkpts,) self._wkpts = wkpts if wkpts is not None else \ torch.ones((nkpts,), dtype=self.dtype, device=self.device) / nkpts assert self._wkpts.shape[0] == self._kpts.shape[0] assert self._wkpts.ndim == 1 assert self._kpts.ndim == 2 # initialize cache self._cache = cache if cache is not None else Cache.get_dummy() self._cache.add_cacheable_params(["overlap", "kinetic", "nuclattr"]) if df is None: self._df: Optional[BaseDF] = None else: self._df = DFPBC(dfinfo=df, wrapper=self._basiswrapper, kpts=self._kpts, wkpts=self._wkpts, eta=self._eta, lattsum_opt=self._lattsum_opt, cache=self._cache.add_prefix("df")) self._is_built = False
def _calc_nucl_attr(self) -> torch.Tensor: # calculate the nuclear attraction matrix # this follows the equation (31) in Sun, et al., J. Chem. Phys. 147 (2017) # construct the fake nuclei atombases for nuclei # (in this case, we assume each nucleus is a very sharp s-type orbital) nucl_atbases = self._create_fake_nucl_bases(alpha=1e16, chargemult=1) # add a compensating charge cnucl_atbases = self._create_fake_nucl_bases(alpha=self._eta, chargemult=-1) # real charge + compensating charge nucl_atbases_all = nucl_atbases + cnucl_atbases nucl_wrapper = intor.LibcintWrapper(nucl_atbases_all, spherical=self._spherical, lattice=self._lattice) cnucl_wrapper = intor.LibcintWrapper(cnucl_atbases, spherical=self._spherical, lattice=self._lattice) natoms = nucl_wrapper.nao() // 2 # construct the k-points ij # duplicating kpts to have shape of (nkpts, 2, ndim) kpts_ij = self._kpts.unsqueeze(-2) * torch.ones( (2, 1), dtype=self.dtype, device=self.device) ############# 1st part of nuclear attraction: short range ############# # get the 1st part of the nuclear attraction: the charge and compensating charge # nuc1: (nkpts, nao, nao, 2 * natoms) # nuc1 is not hermitian basiswrapper1, nucl_wrapper1 = intor.LibcintWrapper.concatenate( self._basiswrapper, nucl_wrapper) nuc1_c = intor.pbc_coul3c(basiswrapper1, other1=basiswrapper1, other2=nucl_wrapper1, kpts_ij=kpts_ij, options=self._lattsum_opt) nuc1 = -nuc1_c[..., :natoms] + nuc1_c[..., natoms:] nuc1 = torch.sum(nuc1, dim=-1) # (nkpts, nao, nao) # add vbar for 3 dimensional cell # vbar is the interaction between the background charge and the # compensating function. # https://github.com/pyscf/pyscf/blob/c9aa2be600d75a97410c3203abf35046af8ca615/pyscf/pbc/df/aft.py#L239 nucbar = sum([-atb.atomz / self._eta for atb in self._atombases]) nuc1_b = -nucbar * np.pi / self._lattice.volume() * self._olp_mat nuc1 = nuc1 + nuc1_b ############# 2nd part of nuclear attraction: long range ############# # get the 2nd part from the Fourier Transform # get the G-points, choosing min because the two FTs are multiplied gcut = get_gcut(self._lattsum_opt.precision, wrappers=[cnucl_wrapper, self._basiswrapper], reduce="min") # gvgrids: (ngv, ndim), gvweights: (ngv,) gvgrids, gvweights = self._lattice.get_gvgrids(gcut) # the compensating charge's Fourier Transform # TODO: split gvgrids and gvweights to reduce the memory usage cnucl_ft = intor.eval_gto_ft(cnucl_wrapper, gvgrids) # (natoms, ngv) # overlap integral of the electron basis' Fourier Transform cbas_ft = intor.pbcft_overlap( self._basiswrapper, gvgrid=-gvgrids, kpts=self._kpts, options=self._lattsum_opt) # (nkpts, nao, nao, ngv) # coulomb kernel Fourier Transform coul_ft = unweighted_coul_ft(gvgrids) * gvweights # (ngv,) coul_ft = coul_ft.to(cbas_ft.dtype) # cast to complex # optimized by opt_einsum # nuc2 = -torch.einsum("tg,kabg,g->kab", cnucl_ft, cbas_ft, coul_ft) nuc2_temp = torch.einsum("g,tg->g", coul_ft, cnucl_ft) nuc2 = -torch.einsum("g,kabg->kab", nuc2_temp, cbas_ft) # (nkpts, nao, nao) # print((nuc2 - nuc2.conj().transpose(-2, -1)).abs().max()) # check hermitian-ness # get the total contribution from the short range and long range nuc = nuc1 + nuc2 # symmetrize for more stable numerical calculation nuc = (nuc + nuc.conj().transpose(-2, -1)) * 0.5 return nuc
def build(self) -> BaseDF: self._is_built = True df = self._dfinfo # calculate the matrices required to calculate the electron repulsion operator # i.e. the 3-centre 2-electron integrals (short + long range) and j3c @ (j2c^-1) method = df.method.lower() df_auxbases = _renormalize_auxbases(df.auxbases) aux_comp_bases = self._create_compensating_bases(df_auxbases, eta=self._eta) fuse_aux_bases = df_auxbases + aux_comp_bases fuse_aux_wrapper = intor.LibcintWrapper( fuse_aux_bases, spherical=self._wrapper.spherical, lattice=self._lattice) aux_comp_wrapper = intor.LibcintWrapper( aux_comp_bases, spherical=self._wrapper.spherical, lattice=self._lattice) aux_wrapper = intor.LibcintWrapper(df_auxbases, spherical=self._wrapper.spherical, lattice=self._lattice) nxcao = aux_comp_wrapper.nao( ) # number of aux compensating basis wrapper nxao = fuse_aux_wrapper.nao() - nxcao # number of aux basis wrapper assert nxcao == nxao # only gaussian density fitting is implemented at the moment if method != "gdf": raise NotImplementedError( "Density fitting %s is not implemented (only gdf)" % df.method) # get the k-points needed for the integrations nkpts = self._kpts.shape[0] kpts_ij = _combine_kpts_to_kpts_ij(self._kpts) # (nkpts_ij, 2, ndim) kpts_reduce = _reduce_kpts_ij(kpts_ij) # (nkpts_ij, ndim) nkpts_ij = kpts_ij.shape[0] kpts_j = kpts_ij[..., 1, :] # (nkpts_ij, ndim) def _calc_integrals(): ######################## short-range integrals ######################## ############# 3-centre 2-electron integral ############# _basisw, _fusew = intor.LibcintWrapper.concatenate( self._wrapper, fuse_aux_wrapper) # (nkpts_ij, nao, nao, nxao+nxcao) j3c_short_f = intor.pbc_coul3c(_basisw, other2=_fusew, kpts_ij=kpts_ij, options=self._lattsum_opt) j3c_short = j3c_short_f[..., :nxao] - j3c_short_f[ ..., nxao:] # (nkpts_ij, nao, nao, nxao) ############# 2-centre 2-electron integrals ############# # (nkpts_unique, nxao+nxcao, nxao+nxcao) j2c_short_f = intor.pbc_coul2c(fuse_aux_wrapper, kpts=kpts_reduce, options=self._lattsum_opt) # j2c_short: (nkpts_unique, nxao, nxao) j2c_short = j2c_short_f[..., :nxao, :nxao] + j2c_short_f[..., nxao:, nxao:] \ - j2c_short_f[..., :nxao, nxao:] - j2c_short_f[..., nxao:, :nxao] ######################## long-range integrals ######################## # only use the compensating wrapper as the gcut gcut = get_gcut(self._lattsum_opt.precision, [aux_comp_wrapper]) # gvgrids: (ngv, ndim), gvweights: (ngv,) gvgrids, gvweights = self._lattice.get_gvgrids(gcut) ngv = gvgrids.shape[0] gvk = gvgrids.unsqueeze(-2) + kpts_reduce # (ngv, nkpts_ij, ndim) gvk = gvk.view(-1, gvk.shape[-1]) # (ngv * nkpts_ij, ndim) # get the fourier transform variables # TODO: iterate over ngv axis # ft of the compensating basis comp_ft = intor.eval_gto_ft(aux_comp_wrapper, gvk) # (nxcao, ngv * nkpts_ij) comp_ft = comp_ft.view(-1, ngv, nkpts_ij) # (nxcao, ngv, nkpts_ij) # ft of the auxiliary basis auxb_ft_c = intor.eval_gto_ft(aux_wrapper, gvk) # (nxao, ngv * nkpts_ij) auxb_ft_c = auxb_ft_c.view(-1, ngv, nkpts_ij) # (nxao, ngv, nkpts_ij) auxb_ft = auxb_ft_c - comp_ft # (nxao, ngv, nkpts_ij) # ft of the overlap integral of the basis (nkpts_ij, nao, nao, ngv) aoao_ft = self._get_pbc_overlap_with_kpts_ij( gvgrids, kpts_reduce, kpts_j) # ft of the coulomb kernel coul_ft = unweighted_coul_ft(gvk) # (ngv * nkpts_ij,) coul_ft = coul_ft.to(comp_ft.dtype).view( ngv, nkpts_ij) * gvweights.unsqueeze(-1) # (ngv, nkpts_ij) # 1: (nkpts_ij, nxao, nxao) pattern = "gi,xgi,ygi->ixy" j2c_long = torch.einsum(pattern, coul_ft, comp_ft.conj(), auxb_ft) # 2: (nkpts_ij, nxao, nxao) j2c_long += torch.einsum(pattern, coul_ft, auxb_ft.conj(), comp_ft) # 3: (nkpts_ij, nxao, nxao) j2c_long += torch.einsum(pattern, coul_ft, comp_ft.conj(), comp_ft) # calculate the j3c long-range patternj3 = "gi,xgi,iyzg->iyzx" # (nkpts_ij, nao, nao, nxao) j3c_long = torch.einsum(patternj3, coul_ft, comp_ft.conj(), aoao_ft) # get the average potential auxbar_f = self._auxbar( kpts_reduce, fuse_aux_wrapper) # (nkpts_ij, nxao + nxcao) auxbar = auxbar_f[:, :nxao] - auxbar_f[:, nxao:] # (nkpts_ij, nxao) auxbar = auxbar.reshape(nkpts, nkpts, auxbar.shape[-1]) # (nkpts, nkpts, nxao) olp_mat = intor.pbc_overlap( self._wrapper, kpts=self._kpts, options=self._lattsum_opt) # (nkpts, nao, nao) j3c_bar = auxbar[:, :, None, None, :] * olp_mat[ ..., None] # (nkpts, nkpts, nao, nao, nxao) j3c_bar = j3c_bar.reshape( -1, *j3c_bar.shape[2:]) # (nkpts_ij, nao, nao, nxao) ######################## combining integrals ######################## j2c = j2c_short + j2c_long # (nkpts_ij, nxao, nxao) j3c = j3c_short + j3c_long - j3c_bar # (nkpts_ij, nao, nao, nxao) el_mat = torch.einsum("kxy,kaby->kabx", torch.inverse(j2c), j3c) # (nkpts_ij, nao, nao, nxao) return j2c, j3c, el_mat with self._cache.open(): # check the signature self._cache.check_signature({ "dfinfo": self._dfinfo, "kpts": self._kpts.detach(), "wkpts": self._wkpts.detach(), "atombases": self._wrapper.atombases, "alattice": self._lattice.lattice_vectors().detach(), }) j2c, j3c, el_mat = self._cache.cache_multi( ["j2c", "j3c", "el_mat"], _calc_integrals) self._j2c = j2c self._j3c = j3c self._el_mat = el_mat return self