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 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 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 _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, 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 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 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 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_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_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 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 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 __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 __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 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 __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_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)
def dqcelrepxc(atom: str, spin: int = 0, xc: str = "lda_x", basis: str = BASIS):
# returns the electron repulsion and xc operator using DQC
mol = Mol(atom, spin=spin, basis=basis, dtype=DTYPE, grid=4)
qc = KS(mol, xc=xc)
hamilt = mol.get_hamiltonian()
if spin == 0:
# set dm to be an identity matrix
dm = torch.eye(hamilt.nao, dtype=DTYPE)
velrepxc = hamilt.get_vxc(dm) + hamilt.get_elrep(dm)
return velrepxc.fullmatrix()
else:
dmu = torch.eye(hamilt.nao, dtype=DTYPE)
dm = SpinParam(u=dmu, d=dmu)
vxc = hamilt.get_vxc(dm)
elrep = hamilt.get_elrep(dm.u + dm.d)
return torch.cat(((vxc.u + elrep).fullmatrix().unsqueeze(0),
(vxc.d + elrep).fullmatrix().unsqueeze(0)), dim=0)
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):
# 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 __init__(self, system: BaseSystem, xc: Union[str, BaseXC],
vext: Optional[torch.Tensor] = None,
restricted: Optional[bool] = None):
# decide if this is restricted or not
if restricted is None:
self._polarized = bool(system.spin != 0)
else:
self._polarized = not restricted
# get the xc object
if isinstance(xc, str):
self.xc: BaseXC = get_xc(xc)
else:
self.xc = xc
self.system = system
# build and setup basis and grid
self.system.setup_grid()
self.hamilton = system.get_hamiltonian().build()
self.hamilton.setup_grid(system.get_grid(), self.xc)
# get the orbital info
self.orb_weight = system.get_orbweight(polarized=self._polarized) # (norb,)
if self._polarized:
assert isinstance(self.orb_weight, SpinParam)
self.norb: Union[int, SpinParam[int]] = SpinParam(
u=self.orb_weight.u.shape[-1],
d=self.orb_weight.d.shape[-1])
else:
assert isinstance(self.orb_weight, torch.Tensor)
self.norb = self.orb_weight.shape[-1]
# set up the vext linear operator
self.knvext_linop = self.hamilton.get_kinnucl() # kinetic, nuclear, and external potential
if vext is not None:
assert vext.shape[-1] == system.get_grid().get_rgrid().shape[-2]
self.knvext_linop = self.knvext_linop + self.hamilton.get_vext(vext)
def test_instability_check():
# check if is_orb_min returns False for a known excited state
atomzs = [8, 8] # O2
atomposs = torch.tensor([[-1.0, 0., 0.], [1.0, 0., 0.]], dtype=dtype)
mol = Mol(moldesc=(atomzs, atomposs), basis="3-21G", dtype=dtype, spin=2)
# known excited state of O2
dm0 = SpinParam(u=torch.tensor([[ 1.4374e-02, -6.5933e-17, -1.3556e-17, 1.1258e-17, -1.8834e-03,
2.1021e-16, -2.3638e-13, -2.6444e-12, 2.2637e-02, -4.9366e-17,
-5.9629e-02, 1.8549e-17, 5.0367e-17, -1.5571e-16, -5.6398e-12,
-6.9681e-12, -1.0048e-01, 1.2257e-15],
[-6.5933e-17, 2.3236e-02, -2.1167e-12, -5.2740e-13, -7.3181e-17,
-1.0873e-02, -8.3590e-17, 1.2056e-17, -9.3082e-17, 8.0595e-02,
3.6757e-16, 9.1094e-12, -3.9485e-13, 1.2539e-01, -9.1129e-18,
-2.8306e-17, -7.9300e-17, -1.8937e-02],
[-1.3556e-17, -2.1167e-12, 4.2572e-02, 8.2444e-04, -1.2280e-16,
1.5572e-12, -1.5306e-17, 1.7276e-16, 2.7191e-18, 6.4115e-13,
1.1101e-16, -1.9820e-01, -3.8413e-02, -1.7302e-12, -7.2395e-16,
-5.6557e-16, 1.6237e-17, 2.3304e-13],
[ 1.1258e-17, -5.2740e-13, 8.2444e-04, 3.9468e-02, 6.1344e-17,
3.8795e-13, 1.2361e-16, -6.1389e-18, 8.2929e-17, 1.5994e-13,
-1.2242e-16, -4.0719e-02, 1.9040e-01, -4.3109e-13, 2.8388e-17,
-3.7710e-16, 9.6793e-17, 5.8088e-14],
[-1.8834e-03, -7.3181e-17, -1.2280e-16, 6.1344e-17, 2.6433e-04,
-1.9302e-19, -1.0905e-13, -1.2203e-12, -1.8316e-03, -1.1234e-16,
4.4601e-03, 5.1546e-16, 4.1761e-16, -5.2628e-16, -7.2761e-12,
-8.9879e-12, 1.5410e-02, -1.5680e-16],
[ 2.1021e-16, -1.0873e-02, 1.5572e-12, 3.8795e-13, -1.9302e-19,
1.9720e-02, -4.5050e-18, 4.1017e-17, 5.5749e-17, -4.4163e-02,
-3.6973e-16, -6.9926e-12, 3.0281e-13, -3.5499e-02, 5.1310e-16,
-3.6489e-17, -2.1794e-16, 1.2650e-01],
[-2.3638e-13, -8.3590e-17, -1.5306e-17, 1.2361e-16, -1.0905e-13,
-4.5050e-18, 1.7231e-02, -1.7505e-04, -9.2770e-14, -3.9129e-16,
3.4332e-13, -4.4355e-18, 8.8915e-16, -4.0986e-16, -9.3919e-02,
9.0076e-02, 1.0851e-12, 1.3598e-16],
[-2.6444e-12, 1.2056e-17, 1.7276e-16, -6.1389e-18, -1.2203e-12,
4.1017e-17, -1.7505e-04, 1.5289e-02, -1.0377e-12, 1.2434e-17,
3.8408e-12, -1.3267e-15, -1.7575e-16, 1.1638e-16, 8.5865e-02,
8.7648e-02, 1.2149e-11, 4.9825e-17],
[ 2.2637e-02, -9.3082e-17, 2.7191e-18, 8.2929e-17, -1.8316e-03,
5.5749e-17, -9.2770e-14, -1.0377e-12, 1.0891e-01, 5.4901e-16,
-3.1041e-01, -2.1589e-16, 9.7090e-17, -2.0236e-16, 2.5871e-12,
3.1957e-12, -1.3275e-02, -3.7240e-16],
[-4.9366e-17, 8.0595e-02, 6.4115e-13, 1.5994e-13, -1.1234e-16,
-4.4163e-02, -3.9129e-16, 1.2434e-17, 5.4901e-16, 9.9128e-01,
6.4760e-16, -9.9216e-13, 4.2760e-14, -1.2043e-02, 1.5942e-16,
-1.0898e-16, 1.7759e-16, 7.3878e-03],
[-5.9629e-02, 3.6757e-16, 1.1101e-16, -1.2242e-16, 4.4601e-03,
-3.6973e-16, 3.4332e-13, 3.8408e-12, -3.1041e-01, 6.4760e-16,
8.8722e-01, -4.5809e-17, 1.7987e-16, -2.7821e-16, 2.1268e-13,
2.6265e-13, -1.1616e-02, -1.0344e-16],
[ 1.8549e-17, 9.1094e-12, -1.9820e-01, -4.0719e-02, 5.1546e-16,
-6.9926e-12, -4.4355e-18, -1.3267e-15, -2.1589e-16, -9.9216e-13,
-4.5809e-17, 9.5723e-01, 1.5197e-04, -1.8852e-12, 1.6341e-16,
1.0968e-16, -7.3105e-17, 1.1978e-12],
[ 5.0367e-17, -3.9485e-13, -3.8413e-02, 1.9040e-01, 4.1761e-16,
3.0281e-13, 8.8915e-16, -1.7575e-16, 9.7090e-17, 4.2760e-14,
1.7987e-16, 1.5197e-04, 9.6073e-01, 8.1569e-14, -4.7992e-16,
3.4136e-16, 3.3857e-17, -5.1807e-14],
[-1.5571e-16, 1.2539e-01, -1.7302e-12, -4.3109e-13, -5.2628e-16,
-3.5499e-02, -4.0986e-16, 1.1638e-16, -2.0236e-16, -1.2043e-02,
-2.7821e-16, -1.8852e-12, 8.1569e-14, 9.8251e-01, 1.5091e-16,
-1.8316e-16, -1.7641e-16, 7.2007e-03],
[-5.6398e-12, -9.1129e-18, -7.2395e-16, 2.8388e-17, -7.2761e-12,
5.1310e-16, -9.3919e-02, 8.5865e-02, 2.5871e-12, 1.5942e-16,
2.1268e-13, 1.6341e-16, -4.7992e-16, 1.5091e-16, 9.8353e-01,
9.6527e-04, -1.3954e-12, 5.5578e-17],
[-6.9681e-12, -2.8306e-17, -5.6557e-16, -3.7710e-16, -8.9879e-12,
-3.6489e-17, 9.0076e-02, 8.7648e-02, 3.1957e-12, -1.0898e-16,
2.6265e-13, 1.0968e-16, 3.4136e-16, -1.8316e-16, 9.6527e-04,
9.8395e-01, -1.7236e-12, 1.2375e-16],
[-1.0048e-01, -7.9300e-17, 1.6237e-17, 9.6793e-17, 1.5410e-02,
-2.1794e-16, 1.0851e-12, 1.2149e-11, -1.3275e-02, 1.7759e-16,
-1.1616e-02, -7.3105e-17, 3.3857e-17, -1.7641e-16, -1.3954e-12,
-1.7236e-12, 9.8924e-01, -1.2411e-16],
[ 1.2257e-15, -1.8937e-02, 2.3304e-13, 5.8088e-14, -1.5680e-16,
1.2650e-01, 1.3598e-16, 4.9825e-17, -3.7240e-16, 7.3878e-03,
-1.0344e-16, 1.1978e-12, -5.1807e-14, 7.2007e-03, 5.5578e-17,
1.2375e-16, -1.2411e-16, 9.8325e-01]], dtype=torch.float64),
d=torch.tensor([[ 2.2958e-02, 5.9133e-17, -4.3465e-17, -3.5203e-17, -5.2141e-03,
1.9909e-16, 1.6717e-14, 1.8684e-13, 3.9446e-02, 2.7848e-16,
-6.6246e-02, 1.2336e-17, -2.6465e-16, 8.6160e-16, -2.8860e-11,
-3.5653e-11, -1.2829e-01, 1.5006e-15],
[ 5.9133e-17, 4.6415e-03, -1.3965e-10, -3.4794e-11, -2.6531e-17,
-7.5184e-03, 6.8007e-17, -1.8251e-17, 1.2508e-16, 4.6059e-02,
-5.7129e-18, 6.4528e-10, -2.7957e-11, -7.4002e-03, -7.1789e-16,
5.1820e-16, -9.0260e-16, -4.8860e-02],
[-4.3465e-17, -1.3965e-10, 4.3859e-02, 1.0928e-02, -6.8315e-17,
6.4583e-11, 1.4620e-17, 1.7630e-16, -1.1049e-16, -2.1608e-10,
2.5101e-16, -2.0430e-01, 8.8512e-03, -7.6096e-10, -7.9051e-16,
-5.0478e-16, 1.2440e-16, -8.9813e-11],
[-3.5203e-17, -3.4794e-11, 1.0928e-02, 2.7228e-03, -1.0420e-17,
1.6091e-11, 3.7273e-17, 1.7772e-17, -2.7800e-17, -5.3837e-11,
5.6712e-18, -5.0902e-02, 2.2053e-03, -1.8960e-10, -7.8317e-16,
8.0795e-17, 2.4562e-16, -2.2377e-11],
[-5.2141e-03, -2.6531e-17, -6.8315e-17, -1.0420e-17, 1.2314e-03,
-5.6992e-18, -1.4157e-13, -1.5837e-12, -7.1148e-03, 7.0004e-17,
9.4083e-03, 3.5543e-16, 4.2740e-17, -2.7677e-16, -2.1264e-12,
-2.6263e-12, 3.2613e-02, 1.1833e-18],
[ 1.9909e-16, -7.5184e-03, 6.4583e-11, 1.6091e-11, -5.6992e-18,
1.6029e-02, -2.4000e-17, 2.0307e-17, 1.2229e-16, -3.2365e-02,
-2.4677e-16, -2.8977e-10, 1.2554e-11, -5.3257e-03, 2.8397e-16,
-1.7180e-16, -8.3203e-17, 1.2099e-01],
[ 1.6717e-14, 6.8007e-17, 1.4620e-17, 3.7273e-17, -1.4157e-13,
-2.4000e-17, 3.7952e-03, 7.4301e-04, -4.3281e-13, 4.4101e-16,
-1.6173e-12, -4.2060e-17, 2.1463e-16, 2.1636e-16, -4.1059e-02,
4.5765e-02, -2.8999e-12, 1.0801e-16],
[ 1.8684e-13, -1.8251e-17, 1.7630e-16, 1.7772e-17, -1.5837e-12,
2.0307e-17, 7.4301e-04, 1.2040e-02, -4.8418e-12, -1.2562e-16,
-1.8091e-11, -1.2667e-15, -4.7435e-17, -1.7249e-16, 7.2754e-02,
8.1250e-02, -3.2434e-11, 8.1554e-17],
[ 3.9446e-02, 1.2508e-16, -1.1049e-16, -2.7800e-17, -7.1148e-03,
1.2229e-16, -4.3281e-13, -4.8418e-12, 1.3971e-01, 1.1579e-15,
-3.3375e-01, -4.4814e-17, -5.2377e-16, 1.4905e-15, -7.4596e-11,
-9.2152e-11, -8.4825e-02, 5.5948e-16],
[ 2.7848e-16, 4.6059e-02, -2.1608e-10, -5.3837e-11, 7.0004e-17,
-3.2365e-02, 4.4101e-16, -1.2562e-16, 1.1579e-15, 9.2048e-01,
-1.0153e-16, 9.8356e-10, -4.2613e-11, -2.6337e-01, -1.3755e-15,
6.0548e-16, -1.5088e-15, -2.5713e-02],
[-6.6246e-02, -5.7129e-18, 2.5101e-16, 5.6712e-18, 9.4083e-03,
-2.4677e-16, -1.6173e-12, -1.8091e-11, -3.3375e-01, -1.0153e-16,
8.6354e-01, 9.1119e-17, 9.7805e-16, -3.0709e-15, -3.0487e-11,
-3.7662e-11, -4.4392e-02, -2.2345e-16],
[ 1.2336e-17, 6.4528e-10, -2.0430e-01, -5.0902e-02, 3.5543e-16,
-2.8977e-10, -4.2060e-17, -1.2667e-15, -4.4814e-17, 9.8356e-10,
9.1119e-17, 9.5163e-01, -4.1229e-02, 3.5411e-09, 3.4000e-16,
-7.3179e-17, 4.4148e-17, 5.0171e-10],
[-2.6465e-16, -2.7957e-11, 8.8512e-03, 2.2053e-03, 4.2740e-17,
1.2554e-11, 2.1463e-16, -4.7435e-17, -5.2377e-16, -4.2613e-11,
9.7805e-16, -4.1229e-02, 1.7862e-03, -1.5342e-10, -3.2950e-15,
1.6944e-15, 1.3384e-15, -2.1736e-11],
[ 8.6160e-16, -7.4002e-03, -7.6096e-10, -1.8960e-10, -2.7677e-16,
-5.3257e-03, 2.1636e-16, -1.7249e-16, 1.4905e-15, -2.6337e-01,
-3.0709e-15, 3.5411e-09, -1.5342e-10, 8.9643e-02, -5.2073e-15,
2.4346e-15, -5.8143e-15, -1.1027e-01],
[-2.8860e-11, -7.1789e-16, -7.9051e-16, -7.8317e-16, -2.1264e-12,
2.8397e-16, -4.1059e-02, 7.2754e-02, -7.4596e-11, -1.3755e-15,
-3.0487e-11, 3.4000e-16, -3.2950e-15, -5.2073e-15, 9.9295e-01,
-4.0971e-03, -9.4472e-12, -6.9572e-16],
[-3.5653e-11, 5.1820e-16, -5.0478e-16, 8.0795e-17, -2.6263e-12,
-1.7180e-16, 4.5765e-02, 8.1250e-02, -9.2152e-11, 6.0548e-16,
-3.7662e-11, -7.3179e-17, 1.6944e-15, 2.4346e-15, -4.0971e-03,
9.9121e-01, -1.1671e-11, 5.1711e-16],
[-1.2829e-01, -9.0260e-16, 1.2440e-16, 2.4562e-16, 3.2613e-02,
-8.3203e-17, -2.8999e-12, -3.2434e-11, -8.4825e-02, -1.5088e-15,
-4.4392e-02, 4.4148e-17, 1.3384e-15, -5.8143e-15, -9.4472e-12,
-1.1671e-11, 9.7256e-01, -8.6388e-16],
[ 1.5006e-15, -4.8860e-02, -8.9813e-11, -2.2377e-11, 1.1833e-18,
1.2099e-01, 1.0801e-16, 8.1554e-17, 5.5948e-16, -2.5713e-02,
-2.2345e-16, 5.0171e-10, -2.1736e-11, -1.1027e-01, -6.9572e-16,
5.1711e-16, -8.6388e-16, 9.6920e-01]], dtype=torch.float64))
qc = HF(mol).run(dm0=dm0)
ene = qc.energy()
assert not is_orb_min(qc)
# known ground state of O2
dm1 = SpinParam(u=torch.tensor([[ 2.0683e-02, -1.4420e-16, 1.8209e-17, -8.9218e-18, -1.1292e-03,
2.4994e-16, 1.3424e-17, -1.3778e-17, 3.1936e-02, -2.0260e-16,
-6.6341e-02, 5.2691e-17, -3.3777e-17, -6.2702e-16, -1.9655e-17,
-1.3381e-15, -1.2179e-01, 1.1911e-15],
[-1.4420e-16, 2.5205e-02, 1.5440e-18, -1.5470e-18, -1.2490e-16,
-1.1032e-02, -4.4534e-17, -4.0071e-17, -1.5877e-16, 8.2001e-02,
3.1857e-16, 3.3331e-16, -2.5398e-16, 1.3203e-01, -1.0306e-16,
1.1040e-17, -7.2944e-17, -1.7112e-02],
[ 1.8209e-17, 1.5440e-18, 3.6441e-02, 1.0879e-13, -1.1019e-16,
-6.7960e-17, 2.3446e-17, 1.4381e-16, -1.3948e-17, 6.9656e-17,
6.2990e-17, -1.8362e-01, -3.7391e-02, 4.0401e-16, -7.8922e-16,
-5.2022e-16, -2.4977e-16, -2.3482e-17],
[-8.9218e-18, -1.5470e-18, 1.0879e-13, 3.6441e-02, 5.4484e-17,
-1.0500e-16, 8.8368e-17, -3.8825e-17, 3.4381e-17, 2.4618e-16,
-1.2742e-16, -3.7391e-02, 1.8362e-01, 2.4770e-16, -2.9739e-16,
-2.9106e-16, 9.1167e-17, 5.5659e-17],
[-1.1292e-03, -1.2490e-16, -1.1019e-16, 5.4484e-17, 7.1554e-05,
1.1112e-17, 2.0896e-17, -8.7962e-18, -8.9524e-04, -1.6020e-16,
1.0596e-03, 4.8917e-16, 3.8244e-16, -8.4877e-16, -1.6487e-16,
1.5251e-16, 8.2674e-03, -2.0908e-16],
[ 2.4994e-16, -1.1032e-02, -6.7960e-17, -1.0500e-16, 1.1112e-17,
1.9013e-02, -3.8712e-18, -2.4730e-18, 1.3092e-16, -4.3091e-02,
-3.6917e-16, 3.2346e-16, -3.8502e-16, -3.7068e-02, 2.5626e-17,
-2.1922e-16, -3.1125e-16, 1.2369e-01],
[ 1.3424e-17, -4.4534e-17, 2.3446e-17, 8.8368e-17, 2.0896e-17,
-3.8712e-18, 1.7462e-02, -1.4887e-11, 4.7461e-18, -8.8927e-17,
-4.0357e-17, -1.1873e-17, 5.0543e-16, -3.6232e-16, -9.4065e-02,
9.1151e-02, -1.0645e-15, 1.8543e-17],
[-1.3778e-17, -4.0071e-17, 1.4381e-16, -3.8825e-17, -8.7962e-18,
-2.4730e-18, -1.4887e-11, 1.7462e-02, -1.2512e-17, -1.1642e-16,
5.8123e-17, -1.3728e-15, -2.0690e-16, -1.4420e-16, 9.1151e-02,
9.4065e-02, -9.7118e-16, 4.9315e-17],
[ 3.1936e-02, -1.5877e-16, -1.3948e-17, 3.4381e-17, -8.9524e-04,
1.3092e-16, 4.7461e-18, -1.2512e-17, 1.2198e-01, 5.5052e-16,
-3.2192e-01, 7.7566e-18, -6.8423e-17, -6.6652e-16, 1.3687e-16,
-5.6171e-16, -4.9446e-02, -1.5290e-16],
[-2.0260e-16, 8.2001e-02, 6.9656e-17, 2.4618e-16, -1.6020e-16,
-4.3091e-02, -8.8927e-17, -1.1642e-16, 5.5052e-16, 9.9113e-01,
6.4477e-16, -1.2044e-16, -2.2427e-16, -1.2837e-02, 1.0462e-16,
-6.5168e-17, 8.9549e-17, 6.9992e-03],
[-6.6341e-02, 3.1857e-16, 6.2990e-17, -1.2742e-16, 1.0596e-03,
-3.6917e-16, -4.0357e-17, 5.8123e-17, -3.2192e-01, 6.4477e-16,
8.7574e-01, 9.2656e-17, 1.0997e-16, -4.7293e-16, 1.4328e-17,
-2.7222e-16, -2.8003e-02, -4.2238e-17],
[ 5.2691e-17, 3.3331e-16, -1.8362e-01, -3.7391e-02, 4.8917e-16,
3.2346e-16, -1.1873e-17, -1.3728e-15, 7.7566e-18, -1.2044e-16,
9.2656e-17, 9.6356e-01, 1.9387e-14, -1.1383e-16, -3.7055e-16,
2.4158e-16, -8.2073e-17, -9.3030e-17],
[-3.3777e-17, -2.5398e-16, -3.7391e-02, 1.8362e-01, 3.8244e-16,
-3.8502e-16, 5.0543e-16, -2.0690e-16, -6.8423e-17, -2.2427e-16,
1.0997e-16, 1.9387e-14, 9.6356e-01, -1.2755e-16, -4.3096e-16,
2.4102e-16, -3.6390e-17, 9.8892e-17],
[-6.2702e-16, 1.3203e-01, 4.0401e-16, 2.4770e-16, -8.4877e-16,
-3.7068e-02, -3.6232e-16, -1.4420e-16, -6.6652e-16, -1.2837e-02,
-4.7293e-16, -1.1383e-16, -1.2755e-16, 9.8060e-01, 5.0343e-17,
3.6493e-17, -2.1165e-16, 7.1881e-03],
[-1.9655e-17, -1.0306e-16, -7.8922e-16, -2.9739e-16, -1.6487e-16,
2.5626e-17, -9.4065e-02, 9.1151e-02, 1.3687e-16, 1.0462e-16,
1.4328e-17, -3.7055e-16, -4.3096e-16, 5.0343e-17, 9.8254e-01,
8.2092e-11, -1.4172e-18, 2.2854e-17],
[-1.3381e-15, 1.1040e-17, -5.2022e-16, -2.9106e-16, 1.5251e-16,
-2.1922e-16, 9.1151e-02, 9.4065e-02, -5.6171e-16, -6.5168e-17,
-2.7222e-16, 2.4158e-16, 2.4102e-16, 3.6493e-17, 8.2092e-11,
9.8254e-01, 4.9448e-18, 9.8734e-17],
[-1.2179e-01, -7.2944e-17, -2.4977e-16, 9.1167e-17, 8.2674e-03,
-3.1125e-16, -1.0645e-15, -9.7118e-16, -4.9446e-02, 8.9549e-17,
-2.8003e-02, -8.2073e-17, -3.6390e-17, -2.1165e-16, -1.4172e-18,
4.9448e-18, 9.8153e-01, -1.4542e-16],
[ 1.1911e-15, -1.7112e-02, -2.3482e-17, 5.5659e-17, -2.0908e-16,
1.2369e-01, 1.8543e-17, 4.9315e-17, -1.5290e-16, 6.9992e-03,
-4.2238e-17, -9.3030e-17, 9.8892e-17, 7.1881e-03, 2.2854e-17,
9.8734e-17, -1.4542e-16, 9.8405e-01]], dtype=torch.float64),
d=torch.tensor([[ 1.2641e-02, -1.5358e-16, 4.7752e-17, -8.9111e-18, -4.4876e-03,
2.2773e-16, 6.1698e-18, -7.4570e-18, 2.0856e-02, -6.6447e-17,
-5.7084e-02, -4.3042e-17, -5.2418e-17, -4.8881e-16, -1.3178e-17,
-1.0353e-15, -9.3633e-02, 1.2411e-15],
[-1.5358e-16, 3.0229e-02, 6.8437e-17, 3.5949e-17, -1.0507e-16,
-1.5192e-02, -1.7095e-17, -2.0281e-17, -2.3740e-16, 8.8960e-02,
5.2239e-16, 2.5002e-17, -8.2051e-17, 1.4194e-01, 1.1475e-16,
1.5330e-17, -3.2950e-17, -3.1999e-02],
[ 4.7752e-17, 6.8437e-17, 7.0466e-31, 2.7196e-31, -2.0268e-17,
-3.0557e-17, 6.9564e-18, -2.8971e-17, -1.2805e-17, 1.5775e-16,
5.3677e-17, 8.9575e-31, 3.5820e-31, 3.5525e-16, -4.2190e-16,
-2.6774e-16, -5.3815e-16, -4.5207e-17],
[-8.9111e-18, 3.5949e-17, 2.7196e-31, 3.5707e-31, 4.3021e-18,
-1.4450e-17, 8.7953e-18, -3.0556e-17, 1.6753e-17, 1.8210e-16,
-5.2250e-17, 9.0062e-31, 5.3400e-31, 1.3106e-16, -4.6254e-16,
-2.6536e-16, 1.2935e-16, 4.9788e-18],
[-4.4876e-03, -1.0507e-16, -2.0268e-17, 4.3021e-18, 1.6888e-03,
-6.6376e-18, 1.3986e-18, 3.8055e-18, -4.7627e-03, -4.8232e-16,
1.2498e-02, 1.2663e-17, 1.4987e-17, -5.3503e-16, -2.2757e-17,
4.8352e-16, 3.8560e-02, -3.3090e-16],
[ 2.2773e-16, -1.5192e-02, -3.0557e-17, -1.4450e-17, -6.6376e-18,
1.9762e-02, 1.6945e-17, 2.4527e-17, 1.1483e-16, -4.5236e-02,
-3.5267e-16, -2.8982e-18, 4.2041e-17, -4.5713e-02, -8.0265e-17,
1.0546e-18, -2.5558e-16, 1.2249e-01],
[ 6.1698e-18, -1.7095e-17, 6.9564e-18, 8.7953e-18, 1.3986e-18,
1.6945e-17, 3.5428e-03, 4.3824e-11, 1.0206e-17, -2.4135e-17,
-2.5997e-17, -9.4414e-18, 9.8201e-17, -6.4590e-17, -4.2669e-02,
4.1347e-02, -4.9636e-16, 7.9868e-17],
[-7.4570e-18, -2.0281e-17, -2.8971e-17, -3.0556e-17, 3.8055e-18,
2.4527e-17, 4.3824e-11, 3.5428e-03, -1.7468e-17, -1.2992e-16,
4.3756e-17, -1.0961e-16, -4.9473e-17, -4.8268e-17, 4.1347e-02,
4.2669e-02, -4.2779e-16, 1.5287e-16],
[ 2.0856e-02, -2.3740e-16, -1.2805e-17, 1.6753e-17, -4.7627e-03,
1.1483e-16, 1.0206e-17, -1.7468e-17, 1.0739e-01, 3.0990e-16,
-3.0878e-01, -1.4332e-16, -1.8655e-16, -6.0195e-16, -1.9554e-17,
-1.0795e-16, -7.5289e-03, -1.6864e-16],
[-6.6447e-17, 8.8960e-02, 1.5775e-16, 1.8210e-16, -4.8232e-16,
-4.5236e-02, -2.4135e-17, -1.2992e-16, 3.0990e-16, 9.8962e-01,
6.2003e-16, -4.3189e-16, -5.5535e-16, -1.5301e-02, 1.3215e-16,
-6.0640e-17, 9.2815e-17, 8.7875e-03],
[-5.7084e-02, 5.2239e-16, 5.3677e-17, -5.2250e-17, 1.2498e-02,
-3.5267e-16, -2.5997e-17, 4.3756e-17, -3.0878e-01, 6.2003e-16,
8.8879e-01, 4.0697e-16, 5.3094e-16, -4.3659e-16, -2.9932e-17,
-6.8804e-17, -9.2814e-03, -2.3331e-17],
[-4.3042e-17, 2.5002e-17, 8.9575e-31, 9.0062e-31, 1.2663e-17,
-2.8982e-18, -9.4414e-18, -1.0961e-16, -1.4332e-16, -4.3189e-16,
4.0697e-16, 4.0132e-30, 1.6971e-30, 4.3752e-16, -1.1655e-15,
-1.4303e-15, 1.7323e-16, -1.6309e-17],
[-5.2418e-17, -8.2051e-17, 3.5820e-31, 5.3400e-31, 1.4987e-17,
4.2041e-17, 9.8201e-17, -4.9473e-17, -1.8655e-16, -5.5535e-16,
5.3094e-16, 1.6971e-30, 4.1255e-30, -1.9731e-16, -1.7601e-15,
5.5023e-16, 1.8679e-16, 4.7534e-17],
[-4.8881e-16, 1.4194e-01, 3.5525e-16, 1.3106e-16, -5.3503e-16,
-4.5713e-02, -6.4590e-17, -4.8268e-17, -6.0195e-16, -1.5301e-02,
-4.3659e-16, 4.3752e-16, -1.9731e-16, 9.7688e-01, 1.6012e-16,
1.7160e-17, -1.5571e-16, 1.0700e-02],
[-1.3178e-17, 1.1475e-16, -4.2190e-16, -4.6254e-16, -2.2757e-17,
-8.0265e-17, -4.2669e-02, 4.1347e-02, -1.9554e-17, 1.3215e-16,
-2.9932e-17, -1.1655e-15, -1.7601e-15, 1.6012e-16, 9.9646e-01,
-2.4166e-10, 1.4456e-17, 2.1470e-17],
[-1.0353e-15, 1.5330e-17, -2.6774e-16, -2.6536e-16, 4.8352e-16,
1.0546e-18, 4.1347e-02, 4.2669e-02, -1.0795e-16, -6.0640e-17,
-6.8804e-17, -1.4303e-15, 5.5023e-16, 1.7160e-17, -2.4166e-10,
9.9646e-01, -6.1813e-17, 5.6203e-17],
[-9.3633e-02, -3.2950e-17, -5.3815e-16, 1.2935e-16, 3.8560e-02,
-2.5558e-16, -4.9636e-16, -4.2779e-16, -7.5289e-03, 9.2815e-17,
-9.2814e-03, 1.7323e-16, 1.8679e-16, -1.5571e-16, 1.4456e-17,
-6.1813e-17, 9.8949e-01, -1.6997e-16],
[ 1.2411e-15, -3.1999e-02, -4.5207e-17, 4.9788e-18, -3.3090e-16,
1.2249e-01, 7.9868e-17, 1.5287e-16, -1.6864e-16, 8.7875e-03,
-2.3331e-17, -1.6309e-17, 4.7534e-17, 1.0700e-02, 2.1470e-17,
5.6203e-17, -1.6997e-16, 9.8351e-01]], dtype=torch.float64))
qc = HF(mol).run(dm0=dm1)
ene = qc.energy()
assert is_orb_min(qc)
def run(
self,
dm0: Optional[
Union[str, torch.Tensor,
SpinParam[torch.Tensor]]] = "1e", # type: ignore
eigen_options: Optional[Dict[str, Any]] = None,
fwd_options: Optional[Dict[str, Any]] = None,
bck_options: Optional[Dict[str, Any]] = None) -> BaseQCCalc:
# get default options
if not self._variational:
fwd_defopt = {
"method": "broyden1",
"alpha": -0.5,
"maxiter": 50,
"verbose": config.VERBOSE > 0,
}
else:
fwd_defopt = {
"method": "gd",
"step": 1e-2,
"maxiter": 5000,
"f_rtol": 1e-10,
"x_rtol": 1e-10,
"verbose": config.VERBOSE > 0,
}
bck_defopt = {
# 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,
}
# setup the default options
if eigen_options is None:
eigen_options = {"method": "exacteig"}
if fwd_options is None:
fwd_options = {}
if bck_options is None:
bck_options = {}
fwd_options = set_default_option(fwd_defopt, fwd_options)
bck_options = set_default_option(bck_defopt, bck_options)
# 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:
dm = self._get_zero_dm()
elif isinstance(dm0, str):
if dm0 == "1e": # initial density based on 1-electron Hamiltonian
dm = self._get_zero_dm()
scp0 = self._engine.dm2scp(dm)
dm = self._engine.scp2dm(scp0)
else:
raise RuntimeError("Unknown dm0: %s" % dm0)
else:
dm = SpinParam.apply_fcn(lambda dm0: dm0.detach(), dm0)
# making it spin param for polarized and tensor for nonpolarized
if isinstance(dm, torch.Tensor) and self._polarized:
dm_u = dm * 0.5
dm_d = dm * 0.5
dm = SpinParam(u=dm_u, d=dm_d)
elif isinstance(dm, SpinParam) and not self._polarized:
dm = dm.u + dm.d
if not self._variational:
scp0 = self._engine.dm2scp(dm)
# 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)
else:
system = self.get_system()
h = system.get_hamiltonian()
orb_weights = system.get_orbweight(polarized=self._polarized)
norb = SpinParam.apply_fcn(lambda orb_weights: len(orb_weights),
orb_weights)
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 params2dm(params: torch.Tensor, coeffs: torch.Tensor) \
-> Union[torch.Tensor, SpinParam[torch.Tensor]]:
p: Union[
torch.Tensor,
SpinParam[torch.Tensor]] = self._engine.unpack_aoparams(
params)
c: Union[
torch.Tensor,
SpinParam[torch.Tensor]] = self._engine.unpack_aoparams(
coeffs)
dm = SpinParam.apply_fcn(
lambda p, c, orb_weights: h.ao_orb_params2dm(
p, c, orb_weights, with_penalty=None), p, c,
orb_weights)
return dm
params0, coeffs0 = dm2params(dm)
params0 = params0.detach()
coeffs0 = coeffs0.detach()
min_params0: torch.Tensor = xitorch.optimize.minimize(
fcn=self._engine.aoparams2ene,
# random noise to add the chance of it gets to the minimum, not
# a saddle point
y0=params0 +
torch.randn_like(params0) * 0.03 / params0.numel(),
params=(
coeffs0,
None,
), # coeffs & with_penalty
bck_options={
**bck_options
},
**fwd_options).detach()
if torch.is_grad_enabled():
# If the gradient is required, then put it through the minimization
# one more time with penalty on the parameters.
# The penalty is to keep the Hamiltonian invertible, stabilizing
# inverse.
# Without the penalty, the Hamiltonian could have 0 eigenvalues
# because of the overparameterization of the aoparams.
min_dm = params2dm(min_params0, coeffs0)
params0, coeffs0 = dm2params(min_dm)
min_params0 = xitorch.optimize.minimize(
fcn=self._engine.aoparams2ene,
y0=params0,
params=(
coeffs0,
1e-1,
), # coeffs & with_penalty
bck_options={**bck_options},
method="gd",
step=0,
maxiter=0)
self._dm = params2dm(min_params0, coeffs0)
self._has_run = True
return self
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_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