import h5py
import torch
from typing import Optional
from torch.autograd import Variable, grad
[docs]
class WaveFunction(torch.nn.Module):
def __init__(
self, nelec: int, ndim: int, kinetic: str = "auto", cuda: bool = False
):
"""
Base class for wave functions.
Args:
nelec (int): number of electrons
ndim (int): number of dimensions
kinetic (str): kinetic energy type. Defaults to "auto".
cuda (bool): move the model to GPU. Defaults to False.
Returns:
None
"""
super(WaveFunction, self).__init__()
self.ndim = ndim
self.nelec = nelec
self.ndim_tot = self.nelec * self.ndim
self.kinetic = kinetic
self.cuda = cuda
self.device = torch.device("cpu")
if self.cuda:
self.device = torch.device("cuda")
self.kinetic_energy = self.kinetic_energy_autograd
self.gradients = self.gradients_autograd
[docs]
def forward(self, x: torch.Tensor):
"""Compute the value of the wave function.
for a multiple conformation of the electrons
Args:
parameters : variational param of the wf
pos: position of the electrons
Returns: values of psi
"""
raise NotImplementedError()
[docs]
def electronic_potential(self, pos: torch.Tensor) -> torch.Tensor:
r"""Computes the electron-electron term
.. math:
V_{ee} = \sum_{e_1} \sum_{e_2} \\frac{1}{r_{e_1e_2}}
Args:
x (torch.tensor): sampling points (Nbatch, 3*Nelec)
Returns:
torch.tensor: values of the electon-electron energy at each sampling points
"""
pot = torch.zeros(pos.shape[0], device=self.device)
for ielec1 in range(self.nelec - 1):
epos1 = pos[:, ielec1 * self.ndim : (ielec1 + 1) * self.ndim]
for ielec2 in range(ielec1 + 1, self.nelec):
epos2 = pos[:, ielec2 * self.ndim : (ielec2 + 1) * self.ndim]
r = torch.sqrt(((epos1 - epos2) ** 2).sum(1)) # + 1E-12
pot += 1.0 / r
return pot.view(-1, 1)
[docs]
def nuclear_potential(self, pos: torch.Tensor) -> torch.Tensor:
r"""Computes the electron-nuclear term
.. math:
V_{en} = - \sum_e \sum_n \\frac{Z_n}{r_{en}}
Args:
x (torch.tensor): sampling points (Nbatch, 3*Nelec)
Returns:
torch.tensor: values of the electon-nuclear energy at each sampling points
"""
p = torch.zeros(pos.shape[0], device=self.device)
for ielec in range(self.nelec):
istart = ielec * self.ndim
iend = (ielec + 1) * self.ndim
pelec = pos[:, istart:iend]
for iatom in range(self.natom):
patom = self.ao.atom_coords[iatom, :]
Z = self.ao.atomic_number[iatom]
r = torch.sqrt(((pelec - patom) ** 2).sum(1)) # + 1E-12
p += -Z / r
return p.view(-1, 1)
[docs]
def nuclear_repulsion(self) -> torch.Tensor:
r"""Computes the nuclear-nuclear repulsion term
.. math:
V_{nn} = \sum_{n_1} \sum_{n_2} \\frac{Z_1Z_2}{r_{n_1n_2}}
Returns:
torch.tensor: values of the nuclear-nuclear energy at each sampling points
"""
vnn = 0.0
for at1 in range(self.natom - 1):
c0 = self.ao.atom_coords[at1, :]
Z0 = self.ao.atomic_number[at1]
for at2 in range(at1 + 1, self.natom):
c1 = self.ao.atom_coords[at2, :]
Z1 = self.ao.atomic_number[at2]
rnn = torch.sqrt(((c0 - c1) ** 2).sum())
vnn += Z0 * Z1 / rnn
return vnn
[docs]
def gradients_autograd(
self, pos: torch.Tensor, pdf: Optional[bool] = False
) -> torch.Tensor:
"""Computes the gradients of the wavefunction (or density)
w.r.t the values of the pos.
Args:
pos (torch.tensor): positions of the walkers
pdf (bool, optional) : if true compute the grads of the density
Returns:
torch.tensor: values of the gradients
"""
out = self.forward(pos)
# compute the grads
grads = grad(out, pos, grad_outputs=torch.ones_like(out), only_inputs=True)[0]
# if we return grad of pdf
if pdf:
grads = 2 * grads * out
return grads
[docs]
def kinetic_energy_autograd(self, pos: torch.Tensor) -> torch.Tensor:
"""Compute the kinetic energy through the 2nd derivative
w.r.t the value of the pos.
Note: this could be replaced by
>>> compute_batch_hessian = vmap(hessian(self.forward), argnums=0), in_dims=0)
>>> batch_hess = compute_batch_hessian(pos).squeeze()
>>> batch_hess = torch.diagonal(batch_hess, dim1=-2, dim2=-1).view(-1,1)
>>> return 0.5 * batch_hess / self.forward(pos)
However this approach seems to requires 10x more memory
Args:
pos (torch.tensor): positions of the walkers
Returns:
values of nabla^2 * Psi
"""
out = self.forward(pos)
# compute the jacobian
z = torch.ones_like(out)
jacob = grad(out, pos, grad_outputs=z, only_inputs=True, create_graph=True)[0]
# compute the diagonal element of the Hessian
z = Variable(torch.ones(jacob.shape[0])).to(self.device)
hess = torch.zeros(jacob.shape[0]).to(self.device)
for idim in range(jacob.shape[1]):
tmp = grad(
jacob[:, idim],
pos,
grad_outputs=z,
only_inputs=True,
create_graph=False,
retain_graph=True,
)[0]
hess += tmp[:, idim]
return -0.5 * hess.view(-1, 1) / out
[docs]
def local_energy(self, pos: torch.Tensor) -> torch.Tensor:
"""Computes the local energy
.. math::
E = K(R) + V_{ee}(R) + V_{en}(R) + V_{nn}
Args:
pos (torch.tensor): sampling points (Nbatch, 3*Nelec)
Returns:
[torch.tensor]: values of the local enrgies at each sampling points
Examples::
>>> mol = Molecule('h2.xyz', calculator='adf', basis = 'dzp')
>>> wf = SlaterJastrow(mol, configs='cas(2,2)')
>>> pos = torch.rand(500,6)
>>> vals = wf.local_energy(pos)
Note:
by default kinetic_energy refers to kinetic_energy_autograd
users can overwrite it to poit to any other methods
see kinetic_energy_jacobi in wf_orbital
"""
ke = self.kinetic_energy(pos)
return (
ke
+ self.nuclear_potential(pos)
+ self.electronic_potential(pos)
+ self.nuclear_repulsion()
)
[docs]
def energy(self, pos: torch.Tensor) -> torch.Tensor:
"""Total energy for the sampling points."""
return torch.mean(self.local_energy(pos))
[docs]
def variance(self, pos: torch.Tensor) -> torch.Tensor:
"""Variance of the energy at the sampling points."""
return torch.var(self.local_energy(pos))
[docs]
def sampling_error(self, eloc: torch.Tensor) -> torch.Tensor:
"""Compute the statistical uncertainty.
Assuming the samples are uncorrelated."""
Npts = eloc.shape[0]
return torch.sqrt(eloc.var() / Npts)
def _energy_variance(self, pos: torch.Tensor) -> torch.Tensor:
"""Return energy and variance."""
el = self.local_energy(pos)
return torch.mean(el), torch.var(el)
def _energy_variance_error(self, pos: torch.Tensor) -> torch.Tensor:
"""Return energy variance and sampling error."""
el = self.local_energy(pos)
return torch.mean(el), torch.var(el), self.sampling_error(el)
[docs]
def pdf(
self, pos: torch.Tensor, return_grad: Optional[bool] = False
) -> torch.Tensor:
"""density of the wave function."""
if return_grad:
return self.gradients(pos, pdf=True)
return (self.forward(pos) ** 2).reshape(-1)
[docs]
def get_number_parameters(self) -> int:
"""Computes the total number of parameters."""
nparam = 0
for _, param in self.named_parameters():
if param.requires_grad:
nparam += param.data.numel()
return nparam
[docs]
def load(
self,
filename: str,
group: Optional[str] = "wf_opt",
model: Optional[str] = "best",
):
"""Load trained parameters
Args:
filename (str): hdf5 filename
group (str, optional): group in the hdf5 file where the model is stored.
Defaults to 'wf_opt'.
model (str, optional): 'best' or ' last'. Defaults to 'best'.
"""
f5 = h5py.File(filename, "r")
grp = f5[group]["models"][model]
data = dict()
for name, val in grp.items():
data[name] = torch.as_tensor(val)
self.load_state_dict(data)
f5.close()