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 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 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 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 __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 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