def test_libxc_gga_value():
# compare the calculated value of GGA potential
dtype = torch.float64
xc = get_libxc("gga_x_pbe")
assert xc.family == 2
assert xc.family == 2
torch.manual_seed(123)
n = 100
rho_u = torch.rand((n, ), dtype=dtype)
rho_d = torch.rand((n, ), dtype=dtype)
rho_tot = rho_u + rho_d
gradn_u = torch.rand((n, 3), dtype=dtype) * 0
gradn_d = torch.rand((n, 3), dtype=dtype) * 0
gradn_tot = gradn_u + gradn_d
densinfo_u = ValGrad(value=rho_u, grad=gradn_u)
densinfo_d = ValGrad(value=rho_d, grad=gradn_d)
densinfo_tot = densinfo_u + densinfo_d
densinfo = SpinParam(u=densinfo_u, d=densinfo_d)
# calculate the energy and compare with analytical expression
edens_unpol = xc.get_edensityxc(densinfo_tot)
edens_unpol_true = pbe_e_true(rho_tot, gradn_tot)
assert torch.allclose(edens_unpol, edens_unpol_true)
edens_pol = xc.get_edensityxc(densinfo)
edens_pol_true = 0.5 * (pbe_e_true(2 * rho_u, 2 * gradn_u) +
pbe_e_true(2 * rho_d, 2 * gradn_d))
assert torch.allclose(edens_pol, edens_pol_true)
def get_vxc(self, densinfo):
# densinfo.value: (*BD, nr)
# return:
# potentialinfo.value: (*BD, nr)
# polarized case
if not isinstance(densinfo, ValGrad):
assert _all_same_shape(densinfo.u, densinfo.d), ERRMSG
rho_u = densinfo.u.value
rho_d = densinfo.d.value
# calculate the dE/dn
dedn = CalcLDALibXCPol.apply(rho_u.reshape(-1), rho_d.reshape(-1),
1, self.libxc_pol) # (2, ninps)
dedn = dedn.reshape(N_VRHO, *rho_u.shape)
# split dE/dn into 2 different potential info
potinfo_u = ValGrad(value=dedn[0])
potinfo_d = ValGrad(value=dedn[1])
return SpinParam(u=potinfo_u, d=potinfo_d)
# unpolarized case
else:
rho = densinfo.value # (*BD, nr)
dedn = CalcLDALibXCUnpol.apply(rho.reshape(-1), 1,
self.libxc_unpol)
dedn = dedn.reshape(rho.shape)
potinfo = ValGrad(value=dedn)
return potinfo
def test_libxc_lda_value():
# check if the value is consistent
xc = get_libxc("lda_x")
assert xc.family == 1
assert xc.family == 1
torch.manual_seed(123)
n = 100
rho_u = torch.rand((n, ), dtype=torch.float64)
rho_d = torch.rand((n, ), dtype=torch.float64)
rho_tot = rho_u + rho_d
densinfo_u = ValGrad(value=rho_u)
densinfo_d = ValGrad(value=rho_d)
densinfo = SpinParam(u=densinfo_u, d=densinfo_d)
densinfo_tot = ValGrad(value=rho_tot)
# calculate the energy and compare with analytic
edens_unpol = xc.get_edensityxc(densinfo_tot)
edens_unpol_true = lda_e_true(rho_tot)
assert torch.allclose(edens_unpol, edens_unpol_true)
edens_pol = xc.get_edensityxc(densinfo)
edens_pol_true = 0.5 * (lda_e_true(2 * rho_u) + lda_e_true(2 * rho_d))
assert torch.allclose(edens_pol, edens_pol_true)
vxc_unpol = xc.get_vxc(densinfo_tot)
vxc_unpol_value_true = lda_v_true(rho_tot)
assert torch.allclose(vxc_unpol.value, vxc_unpol_value_true)
vxc_pol = xc.get_vxc(densinfo)
vxc_pol_u_value_true = lda_v_true(2 * rho_u)
vxc_pol_d_value_true = lda_v_true(2 * rho_d)
assert torch.allclose(vxc_pol.u.value, vxc_pol_u_value_true)
assert torch.allclose(vxc_pol.d.value, vxc_pol_d_value_true)
def test_libxc_ldac_value():
# check if the value of lda_c_pw is consistent
xc = get_libxc("lda_c_pw")
assert xc.family == 1
assert xc.family == 1
torch.manual_seed(123)
n = 100
rho_1 = torch.rand((n,), dtype=torch.float64)
rho_2 = torch.rand((n,), dtype=torch.float64)
rho_u = torch.maximum(rho_1, rho_2)
rho_d = torch.minimum(rho_1, rho_2)
rho_tot = rho_u + rho_d
xi = (rho_u - rho_d) / rho_tot
densinfo_u = ValGrad(value=rho_u)
densinfo_d = ValGrad(value=rho_d)
densinfo = SpinParam(u=densinfo_u, d=densinfo_d)
densinfo_tot = ValGrad(value=rho_tot)
# calculate the energy and compare with analytic
edens_unpol = xc.get_edensityxc(densinfo_tot)
edens_unpol_true = ldac_e_true(rho_tot, rho_tot * 0)
assert torch.allclose(edens_unpol, edens_unpol_true)
edens_pol = xc.get_edensityxc(densinfo)
edens_pol_true = ldac_e_true(rho_tot, xi)
assert torch.allclose(edens_pol, edens_pol_true)
def _postproc_libxc_voutput(densinfo: Union[SpinParam[ValGrad], ValGrad],
vrho: torch.Tensor,
vsigma: Optional[torch.Tensor] = None,
vlapl: Optional[torch.Tensor] = None,
vkin: Optional[torch.Tensor] = None) -> Union[SpinParam[ValGrad], ValGrad]:
# postprocess the output from libxc's 1st derivative into derivative
# suitable for valgrad
# densinfo.value: (..., nr)
# densinfo.grad: (..., ndim, nr)
# polarized case
if isinstance(densinfo, SpinParam):
# vrho: (2, *BD, nr)
# vsigma: (3, *BD, nr)
# vlapl: (2, *BD, nr)
# vkin: (2, *BD, nr)
vrho_u = vrho[0]
vrho_d = vrho[1]
# calculate the gradient potential
vgrad_u: Optional[torch.Tensor] = None
vgrad_d: Optional[torch.Tensor] = None
if vsigma is not None:
# calculate the grad_vxc
vgrad_u = 2 * vsigma[0].unsqueeze(-2) * densinfo.u.grad + \
vsigma[1].unsqueeze(-2) * densinfo.d.grad # (..., 3)
vgrad_d = 2 * vsigma[2].unsqueeze(-2) * densinfo.d.grad + \
vsigma[1].unsqueeze(-2) * densinfo.u.grad
vlapl_u: Optional[torch.Tensor] = None
vlapl_d: Optional[torch.Tensor] = None
if vlapl is not None:
vlapl_u = vlapl[0]
vlapl_d = vlapl[1]
vkin_u: Optional[torch.Tensor] = None
vkin_d: Optional[torch.Tensor] = None
if vkin is not None:
vkin_u = vkin[0]
vkin_d = vkin[1]
potinfo_u = ValGrad(value=vrho_u, grad=vgrad_u, lapl=vlapl_u, kin=vkin_u)
potinfo_d = ValGrad(value=vrho_d, grad=vgrad_d, lapl=vlapl_d, kin=vkin_d)
return SpinParam(u=potinfo_u, d=potinfo_d)
# unpolarized case
else:
# all are: (*BD, nr)
# calculate the gradient potential
if vsigma is not None:
vsigma = 2 * vsigma.unsqueeze(-2) * densinfo.grad # (*BD, ndim, nr)
potinfo = ValGrad(value=vrho, grad=vsigma, lapl=vlapl, kin=vkin)
return potinfo
def test_xc_default_vxc(xccls, libxcname):
# test if the default vxc implementation is correct, compared to libxc
dtype = torch.float64
xc = xccls()
libxc = get_libxc(libxcname)
torch.manual_seed(123)
n = 100
rho_u = torch.rand((n, ), dtype=dtype)
rho_d = torch.rand((n, ), dtype=dtype)
rho_tot = rho_u + rho_d
gradn_u = torch.rand((n, 3), dtype=dtype) * 0
gradn_d = torch.rand((n, 3), dtype=dtype) * 0
gradn_tot = gradn_u + gradn_d
densinfo_u = ValGrad(value=rho_u, grad=gradn_u)
densinfo_d = ValGrad(value=rho_d, grad=gradn_d)
densinfo_tot = densinfo_u + densinfo_d
densinfo = SpinParam(u=densinfo_u, d=densinfo_d)
def assert_valgrad(vg1, vg2):
assert torch.allclose(vg1.value, vg2.value)
assert (vg1.grad is None) == (vg2.grad is None)
assert (vg1.lapl is None) == (vg2.lapl is None)
if vg1.grad is not None:
assert torch.allclose(vg1.grad, vg2.grad)
if vg1.lapl is not None:
assert torch.allclose(vg1.lapl, vg2.lapl)
# check if the energy is the same (implementation check)
xc_edens_unpol = xc.get_edensityxc(densinfo_tot)
lxc_edens_unpol = libxc.get_edensityxc(densinfo_tot)
assert torch.allclose(xc_edens_unpol, lxc_edens_unpol)
xc_edens_pol = xc.get_edensityxc(densinfo)
lxc_edens_pol = libxc.get_edensityxc(densinfo)
assert torch.allclose(xc_edens_pol, lxc_edens_pol)
# calculate the potential and compare with analytical expression
xcpotinfo_unpol = xc.get_vxc(densinfo_tot)
lxcpotinfo_unpol = libxc.get_vxc(densinfo_tot)
assert_valgrad(xcpotinfo_unpol, lxcpotinfo_unpol)
xcpotinfo_pol = xc.get_vxc(densinfo)
lxcpotinfo_pol = libxc.get_vxc(densinfo)
# print(type(xcpotinfo_pol), type(lxcpotinfo_unpol))
assert_valgrad(xcpotinfo_pol.u, lxcpotinfo_pol.u)
assert_valgrad(xcpotinfo_pol.d, lxcpotinfo_pol.d)
def _dm2densinfo(self, dm, family):
# dm: (*BD, nao, nao)
# family: 1 for LDA, 2 for GGA, 3 for MGGA
# self.basis: (nao, ngrid)
if isinstance(dm, SpinParam):
res_u = self._dm2densinfo(dm.u, family)
res_d = self._dm2densinfo(dm.d, family)
return SpinParam(u=res_u, d=res_d)
else:
dens = self._get_dens_at_grid(dm)
# calculate the density gradient
gdens = None
if family >= 2: # requires gradient
# (*BD, ngrid, ndim)
# dm is multiplied by 2 because n(r) = sum (D_ij * phi_i * phi_j), thus
# d.n(r) = sum (D_ij * d.phi_i * phi_j + D_ij * phi_i * d.phi_j)
gdens = self._get_grad_dens_at_grid(dm)
# TODO: implement the density laplacian
# dens: (*BD, ngrid)
# gdens: (*BD, ngrid, ndim)
res = ValGrad(value=dens, grad=gdens)
return res
def test_vext(atomzs, dist):
# check if the external potential gives the correct density profile
# (energy is different because the energy of xc is not integral of potentials
# times density)
def get_dens(qc):
dm = qc.aodm()
rgrid = mol.get_grid().get_rgrid()
dens = mol.get_hamiltonian().aodm2dens(dm, rgrid)
return dens
with xt.enable_debug():
poss = torch.tensor([[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]],
dtype=dtype) * dist
mol = Mol((atomzs, poss), basis="3-21G", dtype=dtype)
xc = get_xc("lda_x")
qc = KS(mol, xc=xc).run()
# get the density profile
dens = get_dens(qc)
ldax_pot = xc.get_vxc(ValGrad(value=dens)).value
# ldax_pot2 = -0.7385587663820223 * (4.0 / 3) * dens ** (1.0 / 3)
# assert torch.allclose(ldax_pot, ldax_pot2)
# calculate the energy with the external potential
mol2 = Mol((atomzs, poss), basis="3-21G", dtype=dtype, vext=ldax_pot)
qc2 = KS(mol2, xc=None).run()
dens2 = get_dens(qc2)
# make sure the densities agree
assert torch.allclose(dens, dens2)
def get_vxc(self, densinfo):
# densinfo.value: (*BD, nr)
# densinfo.grad: (*BD, nr, ndim)
# return:
# potentialinfo.value: (*BD, nr)
# potentialinfo.grad: (*BD, nr, ndim)
# polarized case
if not isinstance(densinfo, ValGrad):
grad_u = densinfo.u.grad
grad_d = densinfo.d.grad
# calculate the dE/dn
vrho, vsigma = self._calc_pol(densinfo.u, densinfo.d,
1) # tuple of (nderiv, *BD, nr)
# calculate the grad_vxc
grad_vxc_u = 2 * vsigma[0].unsqueeze(-1) * grad_u + \
vsigma[1].unsqueeze(-1) * grad_d # (..., 3)
grad_vxc_d = 2 * vsigma[2].unsqueeze(-1) * grad_d + \
vsigma[1].unsqueeze(-1) * grad_u
# split dE/dn into 2 different potential info
potinfo_u = ValGrad(value=vrho[0], grad=grad_vxc_u)
potinfo_d = ValGrad(value=vrho[1], grad=grad_vxc_d)
return SpinParam(u=potinfo_u, d=potinfo_d)
# unpolarized case
else:
# calculate the derivative w.r.t density and grad density
vrho, vsigma = self._calc_unpol(densinfo, 1) # tuple of (*BD, nr)
# calculate the gradient potential
grad_vxc = 2 * vsigma.unsqueeze(
-1) * densinfo.grad # (*BD, nr, ndim)
potinfo = ValGrad(value=vrho, grad=grad_vxc)
return potinfo
def test_libxc_mgga_value():
# compare the calculated value of MGGA potential
dtype = torch.float64
xc = get_libxc("mgga_x_scan")
assert xc.family == 4
torch.manual_seed(123)
n = 100
rho_u = torch.rand((n,), dtype=dtype)
rho_d = torch.rand((n,), dtype=dtype)
rho_tot = rho_u + rho_d
gradn_u = torch.rand((3, n), dtype=dtype) * 0
gradn_d = torch.rand((3, n), dtype=dtype) * 0
gradn_tot = gradn_u + gradn_d
lapl_u = torch.rand((n,), dtype=torch.float64)
lapl_d = torch.rand((n,), dtype=torch.float64)
lapl_tot = lapl_u + lapl_d
tau_w_u = (torch.norm(gradn_u, dim=-2) ** 2 / (8 * rho_u))
tau_w_d = (torch.norm(gradn_d, dim=-2) ** 2 / (8 * rho_d))
kin_u = torch.rand((n,), dtype=torch.float64) + tau_w_u
kin_d = torch.rand((n,), dtype=torch.float64) + tau_w_d
kin_tot = kin_u + kin_d
densinfo_u = ValGrad(value=rho_u, grad=gradn_u, lapl=lapl_u, kin=kin_u)
densinfo_d = ValGrad(value=rho_d, grad=gradn_d, lapl=lapl_d, kin=kin_d)
densinfo_tot = densinfo_u + densinfo_d
densinfo = SpinParam(u=densinfo_u, d=densinfo_d)
# calculate the energy and compare with analytical expression
edens_unpol = xc.get_edensityxc(densinfo_tot)
edens_unpol_true = scan_e_true(rho_tot, gradn_tot, lapl_tot, kin_tot)
assert torch.allclose(edens_unpol, edens_unpol_true)
edens_pol = xc.get_edensityxc(densinfo)
edens_pol_true = 0.5 * (scan_e_true(2 * rho_u, 2 * gradn_u, 2 * lapl_u, 2 * kin_u) + \
scan_e_true(2 * rho_d, 2 * gradn_d, 2 * lapl_d, 2 * kin_d))
assert torch.allclose(edens_pol, edens_pol_true)
def get_vxc(self, densinfo):
"""
Returns the ValGrad for the xc potential given the density info
for unpolarized case.
"""
# This is the default implementation of vxc if there is no implementation
# in the specific class of XC.
# densinfo.value & lapl: (*BD, nr)
# densinfo.grad: (*BD, ndim, nr)
# return:
# potentialinfo.value & lapl: (*BD, nr)
# potentialinfo.grad: (*BD, ndim, nr)
# mark the densinfo components as requiring grads
with self._enable_grad_densinfo(densinfo):
with torch.enable_grad():
edensity = self.get_edensityxc(densinfo) # (*BD, nr)
grad_outputs = torch.ones_like(edensity)
grad_enabled = torch.is_grad_enabled()
if not isinstance(densinfo, ValGrad): # polarized case
if self.family == 1: # LDA
params = (densinfo.u.value, densinfo.d.value)
dedn_u, dedn_d = torch.autograd.grad(
edensity, params, create_graph=grad_enabled, grad_outputs=grad_outputs)
return SpinParam(u=ValGrad(value=dedn_u), d=ValGrad(value=dedn_d))
elif self.family == 2: # GGA
params = (densinfo.u.value, densinfo.d.value, densinfo.u.grad, densinfo.d.grad)
dedn_u, dedn_d, dedg_u, dedg_d = torch.autograd.grad(
edensity, params, create_graph=grad_enabled, grad_outputs=grad_outputs)
return SpinParam(
u=ValGrad(value=dedn_u, grad=dedg_u),
d=ValGrad(value=dedn_d, grad=dedg_d))
else:
raise NotImplementedError(
"Default polarized vxc for family %s is not implemented" % self.family)
else: # unpolarized case
if self.family == 1: # LDA
dedn, = torch.autograd.grad(
edensity, densinfo.value, create_graph=grad_enabled,
grad_outputs=grad_outputs)
return ValGrad(value=dedn)
elif self.family == 2: # GGA
dedn, dedg = torch.autograd.grad(
edensity, (densinfo.value, densinfo.grad), create_graph=grad_enabled,
grad_outputs=grad_outputs)
return ValGrad(value=dedn, grad=dedg)
else:
raise NotImplementedError("Default vxc for family %d is not implemented" % self.family)
def _dm2densinfo(self, dm: torch.Tensor) -> ValGrad:
# overloading from hcgto
# dm: (*BD, nkpts, nao, nao), Hermitian
# family: 1 for LDA, 2 for GGA, 3 for MGGA
# self.basis: (nkpts, nao, ngrid)
# dm @ ao will be used in every case
dmdmh = (dm + dm.transpose(-2, -1).conj()) * 0.5 # (*BD, nao, nao)
dmao = torch.matmul(dmdmh, self.basis.conj()) # (*BD, nao, nr)
dmao2 = 2 * dmao
# calculate the density
dens = torch.einsum("...kir,kir->...r", dmao, self.basis)
# calculate the density gradient
gdens: Optional[torch.Tensor] = None
if self.xcfamily == 2 or self.xcfamily == 4: # GGA or MGGA
if not self.is_grad_ao_set:
msg = "Please call `setup_grid(grid, gradlevel>=1)` to calculate the density gradient"
raise RuntimeError(msg)
gdens = torch.zeros((*dm.shape[:-3], 3, self.basis.shape[-1]),
dtype=dm.dtype,
device=dm.device) # (..., ndim, ngrid)
gdens[..., 0, :] = torch.einsum("...kir,kir->...r", dmao2,
self.grad_basis[0])
gdens[..., 1, :] = torch.einsum("...kir,kir->...r", dmao2,
self.grad_basis[1])
gdens[..., 2, :] = torch.einsum("...kir,kir->...r", dmao2,
self.grad_basis[2])
lapldens: Optional[torch.Tensor] = None
kindens: Optional[torch.Tensor] = None
# if self.xcfamily == 4: # TODO: to be completed
# # 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 = torch.einsum("...kir,kir->...r", dmao2, self.lapl_basis)
# grad_grad = torch.einsum("...kij,kir,kjr->...r", dmdmt, self.grad_basis[0], self.grad_basis[0].conj())
# grad_grad += torch.einsum("...kij,kir,kjr->...r", dmdmt, self.grad_basis[1], self.grad_basis[1].conj())
# grad_grad += torch.einsum("...kij,kir,kjr->...r", dmdmt, self.grad_basis[2], self.grad_basis[2].conj())
# lapldens = lapl_basis + 2 * grad_grad
# kindens = 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_pol(xc, rho_u, rho_d, grad_u, grad_d):
densinfo_u = ValGrad(value=rho_u, grad=grad_u)
densinfo_d = ValGrad(value=rho_d, grad=grad_d)
vxc = xc.get_vxc(SpinParam(u=densinfo_u, d=densinfo_d))
return vxc.u.value, vxc.d.value
def get_edens_pol(xc, rho_u, rho_d, grad_u, grad_d):
densinfo_u = ValGrad(value=rho_u, grad=grad_u)
densinfo_d = ValGrad(value=rho_d, grad=grad_d)
return xc.get_edensityxc(SpinParam(u=densinfo_u, d=densinfo_d))
def get_vxc_unpol(xc, rho, grad):
densinfo = ValGrad(value=rho, grad=grad)
return xc.get_vxc(densinfo).value
def get_edens_unpol(xc, rho, grad):
densinfo = ValGrad(value=rho, grad=grad)
return xc.get_edensityxc(densinfo)
def get_vxc_unpol(xc, rho, grad, lapl, kin):
densinfo = ValGrad(value=rho, grad=grad, lapl=lapl, kin=kin)
return xc.get_vxc(densinfo).value