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 __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 __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 __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_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 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 __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 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
def __init__(self,
system: BaseSystem,
xc: Union[str, BaseXC, None],
restricted: Optional[bool] = None):
# get the xc object
if isinstance(xc, str):
self.xc: Optional[BaseXC] = get_xc(xc)
elif isinstance(xc, BaseXC):
self.xc = xc
else:
self.xc = xc
# system = self.hf_engine.get_system()
self._system = system
# build and setup basis and grid
self.hamilton = system.get_hamiltonian()
if self.xc is not None or system.requires_grid():
system.setup_grid()
self.hamilton.setup_grid(system.get_grid(), self.xc)
# get the HF engine and build the hamiltonian
# no need to rebuild the grid because it has been constructed
self.hf_engine = _HFEngine(system,
restricted=restricted,
build_grid_if_necessary=False)
self._polarized = self.hf_engine.polarized
# 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 vext linear operator
self.knvext_linop = self.hamilton.get_kinnucl(
) # kinetic, nuclear, and external potential
def __dm2fock(self, dm):
vhf = self.__dm2vhf(dm)
fock = SpinParam.apply_fcn(lambda vhf: self._core1e_linop + vhf, vhf)
return fock
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 lowest_eival_orb_hessian(qc: BaseQCCalc) -> torch.Tensor:
"""
Get the lowest eigenvalue of the orbital Hessian
Arguments
---------
qc: BaseQCCalc
The qc calc object that has been executed.
Returns
-------
torch.Tensor
A single-element tensor representing the lowest eigenvalue of the
Hessian of energy with respect to orbital parameters.
It is useful to check the convergence stability whether it ends up
in a ground state or an excited state.
"""
# check if the orbital is in the ground state
dm = qc.aodm()
polarized = isinstance(dm, SpinParam)
system = qc.get_system()
h = system.get_hamiltonian()
# (nao, norb)
orb_weights = system.get_orbweight(polarized=polarized)
norb = SpinParam.apply_fcn(lambda orb_weights: len(orb_weights),
orb_weights)
norb_max = SpinParam.reduce(norb, max)
orb_pc = SpinParam.apply_fcn(
lambda dm, norb: h.dm2ao_orb_params(dm, norb=norb), dm,
norb) # (*, nao, norb1), (*, nao, norb2)
orb_p = SpinParam.apply_fcn(lambda orb_pc: orb_pc[0], orb_pc)
orb_c = SpinParam.apply_fcn(lambda orb_pc: orb_pc[1], orb_pc)
# concatenate the parameters in -1 dim if it is polarized
if isinstance(orb_p, SpinParam):
orb_params = torch.cat((orb_p.u, orb_p.d),
dim=-1).detach().requires_grad_()
orb_coeffs = torch.cat((orb_c.u, orb_c.d),
dim=-1).detach().requires_grad_()
else:
orb_params = orb_p.detach().requires_grad_()
orb_coeffs = orb_c.detach().requires_grad_()
# now reconstruct the orbital from the orbital parameters (just to construct
# the graph)
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
# construct the hessian of the energy w.r.t. orb_params
hess = xt.grad.hess(get_ene, (orb_params, orb_coeffs), idxs=0)
assert isinstance(hess, xt.LinearOperator)
# get the lowest eigenvalue
eival, eivec = xt.linalg.symeig(hess, neig=1, mode="lowest")
return eival