Source code for dqc.hamilton.hcgto

from typing import List, Optional, Union, overload, Tuple, Type
import warnings
import torch
import xitorch as xt
import dqc.hamilton.intor as intor
from dqc.df.base_df import BaseDF
from dqc.df.dfmol import DFMol
from dqc.hamilton.base_hamilton import BaseHamilton
from dqc.hamilton.orbconverter import OrbitalOrthogonalizer, IdentityOrbConverter
from dqc.hamilton.orbparams import BaseOrbParams, QROrbParams, MatExpOrbParams
from dqc.utils.datastruct import AtomCGTOBasis, ValGrad, SpinParam, DensityFitInfo
from dqc.grid.base_grid import BaseGrid
from dqc.xc.base_xc import BaseXC
from dqc.utils.cache import Cache
from dqc.utils.mem import chunkify, get_dtype_memsize
from dqc.utils.config import config
from dqc.utils.misc import logger

[docs]class HamiltonCGTO(BaseHamilton): """ Hamiltonian object of contracted Gaussian type-orbital. This class orthogonalizes the basis by taking the weighted eigenvectors of the overlap matrix, i.e. the eigenvectors divided by square root of the eigenvalues. The advantage of doing this is making the overlap matrix in Roothan's equation identity and it could handle overcomplete basis. """ def __init__(self, atombases: List[AtomCGTOBasis], spherical: bool = True, df: Optional[DensityFitInfo] = None, efield: Optional[Tuple[torch.Tensor, ...]] = None, vext: Optional[torch.Tensor] = None, cache: Optional[Cache] = None, orthozer: bool = True, aoparamzer: str = "qr") -> 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 # set up the orbital converter ovlp = intor.overlap(self.libcint_wrapper) if orthozer: self._orthozer = OrbitalOrthogonalizer(ovlp) else: self._orthozer = IdentityOrbConverter(ovlp) # set up the orbital parameterized if aoparamzer == "qr": self._orbparam: Type[BaseOrbParams] = QROrbParams elif aoparamzer == "matexp": warnings.warn("Parametrization with matrix exponential is still at the experimental stage.") self._orbparam = MatExpOrbParams else: aoparam_opts = ["qr", "matexp"] raise RuntimeError( f"Unknown ao parameterizer: {aoparamzer}. Available options are: {aoparam_opts}") # set up the density matrix self._dfoptions = df if df is None: self._df: Optional[DFMol] = None else: self._df = DFMol(df, wrapper=self.libcint_wrapper, orthozer=self._orthozer) self._efield = efield self._vext = vext 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", "efield0"]) if self._df is None: self._cache.add_cacheable_params(["elrep"]) @property def nao(self) -> int: return self._orthozer.nao() @property def kpts(self) -> torch.Tensor: raise TypeError("Isolated molecule Hamiltonian does not have kpts property") @property def df(self) -> Optional[BaseDF]: return self._df ############# setups #############
[docs] def build(self) -> BaseHamilton: # get the matrices (all (nao, nao), except el_mat) # these matrices have already been normalized with self._cache.open(): # check the signature self._cache.check_signature({ "atombases": self.atombases, "spherical": self.spherical, "dfoptions": self._dfoptions, }) logger.log("Calculating the overlap matrix") self.olp_mat = self._cache.cache("overlap", lambda: intor.overlap(self.libcint_wrapper)) logger.log("Calculating the kinetic matrix") kin_mat = self._cache.cache("kinetic", lambda: intor.kinetic(self.libcint_wrapper)) logger.log("Calculating the nuclear attraction matrix") nucl_mat = self._cache.cache("nuclattr", lambda: intor.nuclattr(self.libcint_wrapper)) self.nucl_mat = nucl_mat self.kinnucl_mat = kin_mat + nucl_mat # electric field integral if self._efield is not None: # (ndim, nao, nao) fac: float = 1.0 for i in range(len(self._efield)): fac *= i + 1 intor_fcn = lambda: intor.int1e("r0" * (i + 1), self.libcint_wrapper) efield_mat_f = self._cache.cache(f"efield{i}", intor_fcn) efield_mat = torch.einsum("dab,d->ab", efield_mat_f, self._efield[i]) self.kinnucl_mat = self.kinnucl_mat + efield_mat / fac if self._df is None: logger.log("Calculating the electron repulsion matrix") self.el_mat = self._cache.cache("elrep", lambda: intor.elrep(self.libcint_wrapper)) # (nao^4) # TODO: decide whether to precompute the 2-eris in the new basis # based on the memory self.el_mat = self._orthozer.convert4(self.el_mat) else: logger.log("Building the density fitting matrices") self._df.build() self.is_built = True # orthogonalize the matrices self.olp_mat = self._orthozer.convert2(self.olp_mat) # (nao2, nao2) self.kinnucl_mat = self._orthozer.convert2(self.kinnucl_mat) self.nucl_mat = self._orthozer.convert2(self.nucl_mat) # external potential if self._vext is not None: vext_mat = self.get_vext(self._vext).fullmatrix() self.kinnucl_mat = self.kinnucl_mat + vext_mat logger.log("Setting up the Hamiltonian done") return self
[docs] def setup_grid(self, grid: BaseGrid, xc: Optional[BaseXC] = None) -> None: # save the family and save the xc self.xc = xc if xc is None: self.xcfamily = 1 else: self.xcfamily = xc.family # save the grid self.grid = grid self.rgrid = grid.get_rgrid() assert grid.coord_type == "cart" # setup the basis as a spatial function logger.log("Calculating the basis values in the grid") self.is_ao_set = True self.basis = intor.eval_gto(self.libcint_wrapper, self.rgrid, to_transpose=True) # (ngrid, nao) self.dvolume = self.grid.get_dvolume() self.basis_dvolume = self.basis * self.dvolume.unsqueeze(-1) # (ngrid, nao) if self.xcfamily == 1: # LDA return # setup the gradient of the basis logger.log("Calculating the basis gradient values in the grid") self.is_grad_ao_set = True # (ndim, nao, ngrid) self.grad_basis = intor.eval_gradgto(self.libcint_wrapper, self.rgrid, to_transpose=True) if self.xcfamily == 2: # GGA return # setup the laplacian of the basis self.is_lapl_ao_set = True logger.log("Calculating the basis laplacian values in the grid") self.lapl_basis = intor.eval_laplgto(self.libcint_wrapper, self.rgrid, to_transpose=True) # (nao, ngrid)
############ fock matrix components ############
[docs] def get_nuclattr(self) -> xt.LinearOperator: # nucl_mat: (nao, nao) # return: (nao, nao) return xt.LinearOperator.m(self.nucl_mat, is_hermitian=True)
[docs] def get_kinnucl(self) -> xt.LinearOperator: # kinnucl_mat: (nao, nao) # return: (nao, nao) return xt.LinearOperator.m(self.kinnucl_mat, is_hermitian=True)
[docs] def get_overlap(self) -> xt.LinearOperator: # olp_mat: (nao, nao) # return: (nao, nao) return xt.LinearOperator.m(self.olp_mat, is_hermitian=True)
[docs] def get_elrep(self, dm: torch.Tensor) -> xt.LinearOperator: # dm: (*BD, nao, nao) # elrep_mat: (nao, nao, nao, nao) # return: (*BD, nao, nao) if self._df is None: mat = torch.einsum("...ij,ijkl->...kl", dm, self.el_mat) mat = (mat + mat.transpose(-2, -1)) * 0.5 # reduce numerical instability return xt.LinearOperator.m(mat, is_hermitian=True) else: elrep = self._df.get_elrep(dm) return elrep
@overload def get_exchange(self, dm: torch.Tensor) -> xt.LinearOperator: ... @overload def get_exchange(self, dm: SpinParam[torch.Tensor]) -> SpinParam[xt.LinearOperator]: ...
[docs] def get_exchange(self, dm): # get the exchange operator # dm: (*BD, nao, nao) # el_mat: (nao, nao, nao, nao) # return: (*BD, nao, nao) if self._df is not None: raise RuntimeError("Exact exchange cannot be computed with density fitting") elif isinstance(dm, torch.Tensor): # the einsum form below is to hack PyTorch's bug #57121 # mat = -0.5 * torch.einsum("...jk,ijkl->...il", dm, self.el_mat) # slower mat = -0.5 * torch.einsum("...il,ijkl->...ijk", dm, self.el_mat).sum(dim=-3) # faster mat = (mat + mat.transpose(-2, -1)) * 0.5 # reduce numerical instability return xt.LinearOperator.m(mat, is_hermitian=True) else: # dm is SpinParam # using the spin-scaling property of exchange energy return SpinParam(u=self.get_exchange(2 * dm.u), d=self.get_exchange(2 * dm.d))
[docs] def get_vext(self, vext: torch.Tensor) -> xt.LinearOperator: # vext: (*BR, ngrid) if not self.is_ao_set: raise RuntimeError("Please call `setup_grid(grid, xc)` to call this function") mat = torch.einsum("...r,rb,rc->...bc", vext, self.basis_dvolume, self.basis) # (*BR, nao, nao) mat = self._orthozer.convert2(mat) mat = (mat + mat.transpose(-2, -1)) * 0.5 # ensure the symmetricity and reduce numerical instability return xt.LinearOperator.m(mat, is_hermitian=True)
@overload def get_vxc(self, dm: SpinParam[torch.Tensor]) -> SpinParam[xt.LinearOperator]: ... @overload def get_vxc(self, dm: torch.Tensor) -> xt.LinearOperator: ...
[docs] def get_vxc(self, dm): # dm: (*BD, nao, nao) assert self.xc is not None, "Please call .setup_grid with the xc object" densinfo = SpinParam.apply_fcn( lambda dm_: self._dm2densinfo(dm_), dm) # value: (*BD, nr) potinfo = self.xc.get_vxc(densinfo) # value: (*BD, nr) vxc_linop = SpinParam.apply_fcn( lambda potinfo_: self._get_vxc_from_potinfo(potinfo_), potinfo) return vxc_linop
############### interface to dm ###############
[docs] def ao_orb2dm(self, orb: torch.Tensor, orb_weight: torch.Tensor) -> torch.Tensor: # convert the atomic orbital to the density matrix # in CGTO, it is U.W.U^T # orb: (*BO, nao, norb) # orb_weight: (*BW, norb) # return: (*BOW, nao, nao) orb_w = orb * orb_weight.unsqueeze(-2) # (*BOW, nao, norb) return torch.matmul(orb, orb_w.transpose(-2, -1)) # (*BOW, nao, nao)
[docs] def aodm2dens(self, dm: torch.Tensor, xyz: torch.Tensor) -> torch.Tensor: # xyz: (*BR, ndim) # dm: (*BD, nao, nao) # returns: (*BRD) dm = self._orthozer.unconvert_dm(dm) nao = dm.shape[-1] xyzshape = xyz.shape # basis: (nao, *BR) basis = intor.eval_gto(self.libcint_wrapper, xyz.reshape(-1, xyzshape[-1])).reshape((nao, *xyzshape[:-1])) basis = torch.movedim(basis, 0, -1) # (*BR, nao) # torch.einsum("...ij,...i,...j->...", dm, basis, basis) dens = torch.matmul(dm, basis.unsqueeze(-1)) # (*BRD, nao, 1) dens = torch.matmul(basis.unsqueeze(-2), dens).squeeze(-1).squeeze(-1) # (*BRD) return dens
############### energy of the Hamiltonian ###############
[docs] def get_e_hcore(self, dm: torch.Tensor) -> torch.Tensor: # get the energy from one electron operator return torch.einsum("...ij,...ji->...", self.kinnucl_mat, dm)
[docs] def get_e_elrep(self, dm: torch.Tensor) -> torch.Tensor: # get the energy from two electron repulsion operator elrep_mat = self.get_elrep(dm).fullmatrix() return 0.5 * torch.einsum("...ij,...ji->...", elrep_mat, dm)
[docs] def get_e_exchange(self, dm: Union[torch.Tensor, SpinParam[torch.Tensor]]) -> torch.Tensor: # get the energy from two electron exchange operator exc_mat = self.get_exchange(dm) ene = SpinParam.apply_fcn( lambda exc_mat, dm: 0.5 * torch.einsum("...ij,...ji->...", exc_mat.fullmatrix(), dm), exc_mat, dm) enetot = SpinParam.sum(ene) return enetot
[docs] def get_e_xc(self, dm: Union[torch.Tensor, SpinParam[torch.Tensor]]) -> torch.Tensor: assert self.xc is not None, "Please call .setup_grid with the xc object" # obtain the energy density per unit volume densinfo = SpinParam.apply_fcn( lambda dm_: self._dm2densinfo(dm_), dm) # (spin) value: (*BD, nr) edens = self.xc.get_edensityxc(densinfo) # (*BD, nr) return torch.sum(self.grid.get_dvolume() * edens, dim=-1)
############### free parameters for variational method ############### @overload def ao_orb_params2dm(self, ao_orb_params: torch.Tensor, ao_orb_coeffs: torch.Tensor, orb_weight: torch.Tensor, with_penalty: None) -> torch.Tensor: ... @overload def ao_orb_params2dm(self, ao_orb_params: torch.Tensor, ao_orb_coeffs: torch.Tensor, orb_weight: torch.Tensor, with_penalty: float) -> Union[torch.Tensor, torch.Tensor]: ...
[docs] def ao_orb_params2dm(self, ao_orb_params, ao_orb_coeffs, orb_weight, with_penalty=None): # convert from atomic orbital parameters to density matrix # the atomic orbital parameter is the inverse QR of the orbital # ao_orb_params: (*BD, nao, norb) out = self._orbparam.params2orb(ao_orb_params, ao_orb_coeffs, with_penalty=with_penalty) if with_penalty is None: ao_orbq = out else: ao_orbq, penalty = out ao_orb = self._orthozer.convert_ortho_orb(ao_orbq) dm = self.ao_orb2dm(ao_orb, orb_weight) if with_penalty is None: return dm else: return dm, penalty
[docs] def dm2ao_orb_params(self, dm: torch.Tensor, norb: int) -> Tuple[torch.Tensor, torch.Tensor]: # convert back the density matrix to one solution in the parameters space # NOTE: this assumes that the orbital weights always decreasing in order mdmm = self._orthozer.unconvert_to_ortho_dm(dm) w, orbq = torch.linalg.eigh(mdmm) # w is ordered increasingly, so we take the last parts orbq_params = orbq[..., -norb:] # (nao, norb) orbq_params = torch.flip(orbq_params, dims=(-1,)) return self._orbparam.orb2params(orbq_params)
################ misc ################ def _dm2densinfo(self, dm: torch.Tensor) -> ValGrad: # dm: (*BD, nao, nao), Hermitian # family: 1 for LDA, 2 for GGA, 4 for MGGA # self.basis: (ngrid, nao) # self.grad_basis: (ndim, ngrid, nao) ngrid = self.basis.shape[-2] batchshape = dm.shape[:-2] # dm @ ao will be used in every case dmdmt = (dm + dm.transpose(-2, -1)) * 0.5 # (*BD, nao2, nao2) # convert it back to dm in the cgto basis dmdmt = self._orthozer.unconvert_dm(dmdmt) # prepare the densinfo components dens = torch.empty((*batchshape, ngrid), dtype=self.dtype, device=self.device) gdens: Optional[torch.Tensor] = None lapldens: Optional[torch.Tensor] = None kindens: Optional[torch.Tensor] = None if self.xcfamily == 2 or self.xcfamily == 4: # GGA or MGGA gdens = torch.empty((*dm.shape[:-2], 3, ngrid), dtype=self.dtype, device=self.device) # (..., ndim, ngrid) if self.xcfamily == 4: # MGGA lapldens = torch.empty((*batchshape, ngrid), dtype=self.dtype, device=self.device) kindens = torch.empty((*batchshape, ngrid), dtype=self.dtype, device=self.device) # It is faster to split into chunks than evaluating a single big chunk maxnumel = config.CHUNK_MEMORY // get_dtype_memsize(self.basis) for basis, ioff, iend in chunkify(self.basis, dim=0, maxnumel=maxnumel): # basis: (ngrid2, nao) dmao = torch.matmul(basis, dmdmt) # (ngrid2, nao) dens[..., ioff:iend] = torch.einsum("...ri,ri->...r", dmao, basis) if self.xcfamily == 2 or self.xcfamily == 4: # GGA or MGGA assert gdens is not None if not self.is_grad_ao_set: msg = "Please call `setup_grid(grid, gradlevel>=1)` to calculate the density gradient" raise RuntimeError(msg) # summing it 3 times is faster than applying the d-axis directly grad_basis0 = self.grad_basis[0, ioff:iend, :] # (ngrid2, nao) grad_basis1 = self.grad_basis[1, ioff:iend, :] grad_basis2 = self.grad_basis[2, ioff:iend, :] gdens[..., 0, ioff:iend] = torch.einsum("...ri,ri->...r", dmao, grad_basis0) * 2 gdens[..., 1, ioff:iend] = torch.einsum("...ri,ri->...r", dmao, grad_basis1) * 2 gdens[..., 2, ioff:iend] = torch.einsum("...ri,ri->...r", dmao, grad_basis2) * 2 if self.xcfamily == 4: assert lapldens is not None assert kindens is not None # calculate the laplacian of the density and kinetic energy density at the grid if not self.is_lapl_ao_set: msg = "Please call `setup_grid(grid, gradlevel>=2)` to calculate the density gradient" raise RuntimeError(msg) lapl_basis_cat = self.lapl_basis[ioff:iend, :] lapl_basis = torch.einsum("...ri,ri->...r", dmao, lapl_basis_cat) grad_grad = torch.einsum("...ri,ri->...r", torch.matmul(grad_basis0, dmdmt), grad_basis0) grad_grad += torch.einsum("...ri,ri->...r", torch.matmul(grad_basis1, dmdmt), grad_basis1) grad_grad += torch.einsum("...ri,ri->...r", torch.matmul(grad_basis2, dmdmt), grad_basis2) # pytorch's "...ij,ir,jr->...r" is really slow for large matrix # grad_grad = torch.einsum("...ij,ir,jr->...r", dmdmt, self.grad_basis[0], self.grad_basis[0]) # grad_grad += torch.einsum("...ij,ir,jr->...r", dmdmt, self.grad_basis[1], self.grad_basis[1]) # grad_grad += torch.einsum("...ij,ir,jr->...r", dmdmt, self.grad_basis[2], self.grad_basis[2]) lapldens[..., ioff:iend] = (lapl_basis + grad_grad) * 2 kindens[..., ioff:iend] = grad_grad * 0.5 # dens: (*BD, ngrid) # gdens: (*BD, ndim, ngrid) res = ValGrad(value=dens, grad=gdens, lapl=lapldens, kin=kindens) return res def _get_vxc_from_potinfo(self, potinfo: ValGrad) -> xt.LinearOperator: # obtain the vxc operator from the potential information # potinfo.value: (*BD, nr) # potinfo.grad: (*BD, ndim, nr) # potinfo.lapl: (*BD, nr) # potinfo.kin: (*BD, nr) # self.basis: (nr, nao) # self.grad_basis: (ndim, nr, nao) # prepare the fock matrix component from vxc nao = self.basis.shape[-1] mat = torch.zeros((*potinfo.value.shape[:-1], nao, nao), dtype=self.dtype, device=self.device) # Split the r-dimension into several parts, it is usually faster than # evaluating all at once maxnumel = config.CHUNK_MEMORY // get_dtype_memsize(self.basis) for basis, ioff, iend in chunkify(self.basis, dim=0, maxnumel=maxnumel): # basis: (nr, nao) vb = potinfo.value[..., ioff:iend].unsqueeze(-1) * basis # (*BD, nr, nao) if self.xcfamily in [2, 4]: # GGA or MGGA assert potinfo.grad is not None # (..., ndim, nr) vgrad = potinfo.grad[..., ioff:iend] * 2 grad_basis0 = self.grad_basis[0, ioff:iend, :] # (nr, nao) grad_basis1 = self.grad_basis[1, ioff:iend, :] grad_basis2 = self.grad_basis[2, ioff:iend, :] vb += torch.einsum("...r,ra->...ra", vgrad[..., 0, :], grad_basis0) vb += torch.einsum("...r,ra->...ra", vgrad[..., 1, :], grad_basis1) vb += torch.einsum("...r,ra->...ra", vgrad[..., 2, :], grad_basis2) if self.xcfamily == 4: # MGGA assert potinfo.lapl is not None # (..., nrgrid) assert potinfo.kin is not None lapl = potinfo.lapl[..., ioff:iend] kin = potinfo.kin[..., ioff:iend] vb += 2 * lapl.unsqueeze(-1) * self.lapl_basis[ioff:iend, :] # calculating the matrix from multiplication with the basis mat += torch.matmul(self.basis_dvolume[ioff:iend, :].transpose(-2, -1), vb) if self.xcfamily == 4: # MGGA assert potinfo.lapl is not None # (..., nrgrid) assert potinfo.kin is not None lapl_kin_dvol = (2 * lapl + 0.5 * kin) * self.dvolume[..., ioff:iend] mat += torch.einsum("...r,rb,rc->...bc", lapl_kin_dvol, grad_basis0, grad_basis0) mat += torch.einsum("...r,rb,rc->...bc", lapl_kin_dvol, grad_basis1, grad_basis1) mat += torch.einsum("...r,rb,rc->...bc", lapl_kin_dvol, grad_basis2, grad_basis2) # construct the Hermitian linear operator mat = self._orthozer.convert2(mat) mat = (mat + mat.transpose(-2, -1)) * 0.5 vxc_linop = xt.LinearOperator.m(mat, is_hermitian=True) return vxc_linop
[docs] def getparamnames(self, methodname: str, prefix: str = "") -> List[str]: if methodname == "get_kinnucl": return [prefix + "kinnucl_mat"] elif methodname == "get_nuclattr": return [prefix + "nucl_mat"] elif methodname == "get_overlap": return [prefix + "olp_mat"] elif methodname == "get_elrep": if self._df is None: return [prefix + "el_mat"] else: return self._df.getparamnames("get_elrep", prefix=prefix + "_df.") elif methodname == "get_exchange": return [prefix + "el_mat"] elif methodname == "ao_orb2dm": return [] elif methodname == "ao_orb_params2dm": return self.getparamnames("ao_orb2dm", prefix=prefix) + \ self._orthozer.getparamnames("convert_ortho_orb", prefix=prefix + "_orthozer.") elif methodname == "get_e_hcore": return [prefix + "kinnucl_mat"] elif methodname == "get_e_elrep": return self.getparamnames("get_elrep", prefix=prefix) elif methodname == "get_e_exchange": return self.getparamnames("get_exchange", prefix=prefix) elif methodname == "get_e_xc": assert self.xc is not None return self.getparamnames("_dm2densinfo", prefix=prefix) + \ self.xc.getparamnames("get_edensityxc", prefix=prefix + "xc.") + \ self.grid.getparamnames("get_dvolume", prefix=prefix + "grid.") elif methodname == "get_vext": return [prefix + "basis_dvolume", prefix + "basis"] + \ self._orthozer.getparamnames("convert2", prefix=prefix + "_orthozer.") elif methodname == "get_grad_vext": return [prefix + "basis_dvolume", prefix + "grad_basis"] elif methodname == "get_lapl_kin_vext": return [prefix + "dvolume", prefix + "basis", prefix + "grad_basis", prefix + "lapl_basis"] elif methodname == "get_vxc": assert self.xc is not None return self.getparamnames("_dm2densinfo", prefix=prefix) + \ self.getparamnames("_get_vxc_from_potinfo", prefix=prefix) + \ self.xc.getparamnames("get_vxc", prefix=prefix + "xc.") elif methodname == "_dm2densinfo": params = [prefix + "basis"] + \ self._orthozer.getparamnames("unconvert_dm", prefix=prefix + "_orthozer.") if self.xcfamily == 2 or self.xcfamily == 4: params += [prefix + "grad_basis"] if self.xcfamily == 4: params += [prefix + "lapl_basis"] return params elif methodname == "_get_vxc_from_potinfo": params = [prefix + "basis", prefix + "basis_dvolume"] + \ self._orthozer.getparamnames("convert2", prefix=prefix + "_orthozer.") if self.xcfamily in [2, 4]: params += [prefix + "grad_basis"] if self.xcfamily == 4: params += [prefix + "lapl_basis", prefix + "dvolume"] return params else: raise KeyError("getparamnames has no %s method" % methodname)
# TODO: complete this