def __dm2vhf(self, dm):
# from density matrix, returns the linear operator on electron-electron
# coulomb and exchange
elrep = self._hamilton.get_elrep(SpinParam.sum(dm))
exch = self._hamilton.get_exchange(dm)
vhf = SpinParam.apply_fcn(lambda exch: elrep + exch, exch)
return vhf
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 dm2params(dm: Union[torch.Tensor, SpinParam[torch.Tensor]]) -> \
Tuple[torch.Tensor, torch.Tensor]:
pc = SpinParam.apply_fcn(
lambda dm, norb: h.dm2ao_orb_params(SpinParam.sum(dm),
norb=norb), dm, norb)
p = SpinParam.apply_fcn(lambda pc: pc[0], pc)
c = SpinParam.apply_fcn(lambda pc: pc[1], pc)
params = self._engine.pack_aoparams(p)
coeffs = self._engine.pack_aoparams(c)
return params, coeffs
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(
"...kij,...kji,k->...", exc_mat.fullmatrix(), dm, self._wkpts),
exc_mat, dm)
enetot = SpinParam.sum(ene)
return enetot
def __dm2fock(self, dm):
elrep = self.hamilton.get_elrep(SpinParam.sum(dm)) # (..., nao, nao)
core_coul = self.knvext_linop + elrep
if self.xc is not None:
vxc = self.hamilton.get_vxc(
dm) # spin param or tensor (..., nao, nao)
return SpinParam.apply_fcn(lambda vxc_: vxc_ + core_coul, vxc)
else:
if isinstance(dm, SpinParam):
return SpinParam(u=core_coul, d=core_coul)
else:
return core_coul
def get_ene(orb_params, orb_coeffs):
if polarized:
orb_p = SpinParam(u=orb_params[..., :norb.u],
d=orb_params[..., norb.u:])
orb_c = SpinParam(u=orb_coeffs[..., :norb.u],
d=orb_coeffs[..., norb.u:])
else:
orb_p = orb_params
orb_c = orb_coeffs
dm2 = SpinParam.apply_fcn(
lambda orb_p, orb_c, orb_weights: h.ao_orb_params2dm(
orb_p, orb_c, orb_weights), orb_p, orb_c, orb_weights)
ene = qc.dm2energy(dm2)
return ene
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 get_vxc(self, densinfo):
avxc = self.a.get_vxc(densinfo)
if isinstance(densinfo, ValGrad):
return avxc * self.b
else:
return SpinParam(u=avxc.u * self.b, d=avxc.d * self.b)
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 __fock2dm(self, fock):
# diagonalize the fock matrix and obtain the density matrix
eigvals, eigvecs = self.diagonalize(fock, self._norb)
dm = SpinParam.apply_fcn(
lambda eivecs, orb_weights: self._hamilton.ao_orb2dm(
eivecs, orb_weights), eigvecs, self._orb_weight)
return dm
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 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 _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 get_vxc(self, dm):
# dm: (*BD, nao, nao)
assert self.xc is not None, "Please call .setup_grid with the xc object"
densinfo = self._dm2densinfo(dm, self.xc.family) # value: (*BD, nr)
potinfo = self.xc.get_vxc(densinfo) # value: (*BD, nr)
if isinstance(dm, torch.Tensor): # unpolarized case
# get the linear operator from the potential
vxc_linop = self.get_vext(potinfo.value)
if self.xcfamily >= 2: # GGA or MGGA
assert potinfo.grad is not None
vxc_linop = vxc_linop + self.get_grad_vext(potinfo.grad)
if self.xcfamily >= 3: # MGGA
assert potinfo.lapl is not None
vxc_linop = vxc_linop + self.get_lapl_vext(potinfo.lapl)
return vxc_linop
else: # polarized case
# get the linear operator from the potential
vxc_linop_u = self.get_vext(potinfo.u.value)
vxc_linop_d = self.get_vext(potinfo.d.value)
if self.xcfamily >= 2: # GGA or MGGA
assert potinfo.u.grad is not None
assert potinfo.d.grad is not None
vxc_linop_u = vxc_linop_u + self.get_grad_vext(potinfo.u.grad)
vxc_linop_d = vxc_linop_d + self.get_grad_vext(potinfo.d.grad)
if self.xcfamily >= 3: # MGGA
assert potinfo.u.lapl is not None
assert potinfo.d.lapl is not None
vxc_linop_u = vxc_linop_u + self.get_lapl_vext(potinfo.u.lapl)
vxc_linop_d = vxc_linop_d + self.get_lapl_vext(potinfo.d.lapl)
return SpinParam(u=vxc_linop_u, d=vxc_linop_d)
def __init__(self,
system: BaseSystem,
restricted: Optional[bool] = None,
build_grid_if_necessary: bool = False):
# decide if this is restricted or not
if restricted is None:
self._polarized = bool(system.spin != 0)
else:
self._polarized = not restricted
# construct the grid if the system requires it
if build_grid_if_necessary and system.requires_grid():
system.setup_grid()
system.get_hamiltonian().setup_grid(system.get_grid())
# build the basis
self._hamilton = system.get_hamiltonian().build()
self._system = system
# get the orbital info
self._orb_weight = system.get_orbweight(
polarized=self._polarized) # (norb,)
self._norb = SpinParam.apply_fcn(
lambda orb_weight: int(orb_weight.shape[-1]), self._orb_weight)
# set up the 1-electron linear operator
self._core1e_linop = self._hamilton.get_kinnucl(
) # kinetic and nuclear
def get_orbweight(
self,
polarized: bool = False
) -> Union[torch.Tensor, SpinParam[torch.Tensor]]:
if not polarized:
return self._orb_weights
else:
return SpinParam(u=self._orb_weights_u, d=self._orb_weights_d)
def get_vxc(self, densinfo):
avxc = self.a.get_vxc(densinfo)
bvxc = self.b.get_vxc(densinfo)
if isinstance(densinfo, ValGrad):
return avxc + bvxc
else:
return SpinParam(u=avxc.u + bvxc.u, d=avxc.d + bvxc.d)
def dm2energy(
self, dm: Union[torch.Tensor,
SpinParam[torch.Tensor]]) -> torch.Tensor:
# calculate the energy given the density matrix
dmtot = SpinParam.sum(dm)
e_core = self._hamilton.get_e_hcore(dmtot)
e_elrep = self._hamilton.get_e_elrep(dmtot)
e_exch = self._hamilton.get_e_exchange(dm)
return e_core + e_elrep + e_exch + self._system.get_nuclei_energy()
def unpack_aoparams(
self, aoparams: torch.Tensor
) -> Union[torch.Tensor, SpinParam[torch.Tensor]]:
# if polarized, then construct the SpinParam (reverting the pack_aoparams)
if isinstance(self._norb, SpinParam):
return SpinParam(u=aoparams[..., :self._norb.u],
d=aoparams[..., self._norb.u:])
else:
return aoparams
def diagonalize(self, fock, norb):
ovlp = self._hamilton.get_overlap()
if isinstance(fock, SpinParam):
assert isinstance(self._norb, SpinParam)
eivals_u, eivecs_u = xitorch.linalg.lsymeig(A=fock.u,
neig=norb.u,
M=ovlp,
**self.eigen_options)
eivals_d, eivecs_d = xitorch.linalg.lsymeig(A=fock.d,
neig=norb.d,
M=ovlp,
**self.eigen_options)
return SpinParam(u=eivals_u, d=eivals_d), SpinParam(u=eivecs_u,
d=eivecs_d)
else:
return xitorch.linalg.lsymeig(A=fock,
neig=norb,
M=ovlp,
**self.eigen_options)
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 aoparams2dm(self, aoparams: torch.Tensor, aocoeffs: torch.Tensor,
with_penalty: Optional[float] = None) -> \
Tuple[Union[torch.Tensor, SpinParam[torch.Tensor]], Optional[torch.Tensor]]:
# convert the aoparams to density matrix and penalty factor
aop = self.unpack_aoparams(aoparams) # tensor or SpinParam of tensor
aoc = self.unpack_aoparams(aocoeffs) # tensor or SpinParam of tensor
dm_penalty = SpinParam.apply_fcn(
lambda aop, aoc, orb_weight: self._hamilton.ao_orb_params2dm(
aop, aoc, orb_weight, with_penalty=with_penalty), aop, aoc,
self._orb_weight)
if with_penalty is not None:
dm = SpinParam.apply_fcn(lambda dm_penalty: dm_penalty[0],
dm_penalty)
penalty: Optional[torch.Tensor] = SpinParam.sum(
SpinParam.apply_fcn(lambda dm_penalty: dm_penalty[1],
dm_penalty))
else:
dm = dm_penalty
penalty = None
return dm, penalty
def scp2dm(
self,
scp: torch.Tensor) -> Union[torch.Tensor, SpinParam[torch.Tensor]]:
# scp is like KS, using the concatenated Fock matrix
if not self._polarized:
fock = xt.LinearOperator.m(_symm(scp), is_hermitian=True)
return self.__fock2dm(fock)
else:
fock_u = xt.LinearOperator.m(_symm(scp[0]), is_hermitian=True)
fock_d = xt.LinearOperator.m(_symm(scp[1]), is_hermitian=True)
return self.__fock2dm(SpinParam(u=fock_u, d=fock_d))
def __dm2fock(self, dm):
if isinstance(dm, torch.Tensor):
# construct the fock matrix from the density matrix
elrep = self.hamilton.get_elrep(dm) # (..., nao, nao)
vxc = self.hamilton.get_vxc(dm)
return self.knvext_linop + elrep + vxc
else:
elrep = self.hamilton.get_elrep(dm.u + dm.d) # (..., nao, nao)
vext_elrep = self.knvext_linop + elrep
vxc_ud = self.hamilton.get_vxc(dm)
return SpinParam(u=vext_elrep + vxc_ud.u, d=vext_elrep + vxc_ud.d)
def __fock2dm(self, fock):
# diagonalize the fock matrix and obtain the density matrix
eigvals, eigvecs = self.__diagonalize(fock)
if isinstance(eigvecs, torch.Tensor): # unpolarized
assert isinstance(self.orb_weight, torch.Tensor)
dm = self.hamilton.ao_orb2dm(eigvecs, self.orb_weight)
return dm
else: # polarized
assert isinstance(self.orb_weight, SpinParam)
dm_u = self.hamilton.ao_orb2dm(eigvecs.u, self.orb_weight.u)
dm_d = self.hamilton.ao_orb2dm(eigvecs.d, self.orb_weight.d)
return SpinParam(u=dm_u, d=dm_d)
def run(self, dm0: Optional[Union[torch.Tensor, SpinParam[torch.Tensor]]] = None, # type: ignore
eigen_options: Optional[Mapping[str, Any]] = None,
fwd_options: Optional[Mapping[str, Any]] = None,
bck_options: Optional[Mapping[str, Any]] = None) -> BaseQCCalc:
# setup the default options
if eigen_options is None:
eigen_options = {
"method": "exacteig"
}
if fwd_options is None:
fwd_options = {
"method": "broyden1",
"alpha": -0.5,
"maxiter": 50,
# "verbose": True,
}
if bck_options is None:
bck_options = {
# NOTE: it seems like in most cases the jacobian matrix is posdef
# if it is not the case, we can just remove the line below
"posdef": True,
# "verbose": True,
}
# save the eigen_options for use in diagonalization
self.engine.set_eigen_options(eigen_options)
# set up the initial self-consistent param guess
if dm0 is None:
if not self.polarized:
dm0 = torch.zeros(self.shape, dtype=self.dtype,
device=self.device)
else:
dm0_u = torch.zeros(self.shape, dtype=self.dtype,
device=self.device)
dm0_d = torch.zeros(self.shape, dtype=self.dtype,
device=self.device)
dm0 = SpinParam(u=dm0_u, d=dm0_d)
scp0 = self.engine.dm2scp(dm0)
# do the self-consistent iteration
scp = xitorch.optimize.equilibrium(
fcn=self.engine.scp2scp,
y0=scp0,
bck_options={**bck_options},
**fwd_options)
# post-process parameters
self._dm = self.engine.scp2dm(scp)
self.has_run = True
return self
def dm2energy(
self, dm: Union[torch.Tensor,
SpinParam[torch.Tensor]]) -> torch.Tensor:
# calculate the energy given the density matrix
dmtot = SpinParam.sum(dm)
e_core = self.hamilton.get_e_hcore(dmtot)
e_elrep = self.hamilton.get_e_elrep(dmtot)
if self.xc is not None:
e_xc: Union[torch.Tensor, float] = self.hamilton.get_e_xc(dm)
else:
e_xc = 0.0
return e_core + e_elrep + e_xc + self._system.get_nuclei_energy()
def scp2dm(self, scp: torch.Tensor) -> Union[torch.Tensor, SpinParam[torch.Tensor]]:
# convert the self-consistent parameter (scp) to the density matrix
def _symm(scp: torch.Tensor):
# forcely symmetrize the tensor
return (scp + scp.transpose(-2, -1)) * 0.5
if not self._polarized:
fock = xt.LinearOperator.m(_symm(scp), is_hermitian=True)
return self.__fock2dm(fock)
else:
fock_u = xt.LinearOperator.m(_symm(scp[0]), is_hermitian=True)
fock_d = xt.LinearOperator.m(_symm(scp[1]), is_hermitian=True)
return self.__fock2dm(SpinParam(u=fock_u, d=fock_d))
def _get_zero_dm(self) -> Union[SpinParam[torch.Tensor], torch.Tensor]:
# get the initial dm that are all zeros
if not self._polarized:
return torch.zeros(self._shape,
dtype=self.dtype,
device=self.device)
else:
dm0_u = torch.zeros(self._shape,
dtype=self.dtype,
device=self.device)
dm0_d = torch.zeros(self._shape,
dtype=self.dtype,
device=self.device)
return SpinParam(u=dm0_u, d=dm0_d)
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)