qmctorch.wavefunction.jastrows.elec_elec_nuclei package
Subpackages
- qmctorch.wavefunction.jastrows.elec_elec_nuclei.kernels package
- Submodules
- qmctorch.wavefunction.jastrows.elec_elec_nuclei.kernels.boys_handy_jastrow_kernel module
- qmctorch.wavefunction.jastrows.elec_elec_nuclei.kernels.fully_connected_jastrow_kernel module
- qmctorch.wavefunction.jastrows.elec_elec_nuclei.kernels.jastrow_kernel_electron_electron_nuclei_base module
- Module contents
Submodules
qmctorch.wavefunction.jastrows.elec_elec_nuclei.jastrow_factor_electron_electron_nuclei module
- class qmctorch.wavefunction.jastrows.elec_elec_nuclei.jastrow_factor_electron_electron_nuclei.JastrowFactorElectronElectronNuclei(*args: Any, **kwargs: Any)[source]
Bases:
ModuleJastrow Factor of the elec-elec-nuc term:
\[J = \exp\left( \sum_A \sum_{i<j} K(R_{iA}, r_{jA}, r_{rij}) \right)\]- Parameters:
nup (int) – number of spin up electons
ndow (int) – number of spin down electons
atomic_pos (torch.tensor) – positions of the atoms
jastrow_kernel (kernel) – class of a electron-electron Jastrow kernel
kernel_kwargs (dict, optional) – keyword argument of the kernel. Defaults to {}.
cuda (bool, optional) – Turns GPU ON/OFF. Defaults to False.
- get_mask_tri_up() Tuple[torch.Tensor, torch.Tensor, torch.Tensor][source]
Get the mask to select the triangular up matrix
- Returns:
mask of the tri up matrix
- Return type:
torch.tensor
- extract_tri_up(inp: torch.Tensor) torch.Tensor[source]
extract the upper triangular elements
- Parameters:
input (torch.tensor) – input matrices (…, nelec, nelec)
- Returns:
triangular up element (…, nelec_pair)
- Return type:
torch.tensor
- extract_elec_nuc_dist(en_dist: torch.Tensor) torch.Tensor[source]
Organize the elec nuc distances
- Parameters:
en_dist (torch.tensor) – electron-nuclei distances nbatch x nelec x natom or nbatch x 3 x nelec x natom (dr)
- Returns:
nbatch x natom x nelec_pair x 2 or torch.tensor: nbatch x 3 x natom x nelec_pair x 2 (dr)
- Return type:
torch.tensor
- assemble_dist(pos: torch.Tensor) torch.Tensor[source]
Assemle the different distances for easy calculations
- Parameters:
pos (torch.tensor) – Positions of the electrons Size : Nbatch, Nelec x Ndim
- Returns:
nbatch, natom, nelec_pair, 3
- Return type:
torch.tensor
- assemble_dist_deriv(pos: torch.Tensor, derivative: int = 1) torch.Tensor[source]
- Assemle the different distances for easy calculations
the output has dimension nbatch, 3 x natom, nelec_pair, 3 the last dimension is composed of [r_{e_1n}, r_{e_2n}, r_{ee}]
- Parameters:
pos (torch.tensor) – Positions of the electrons Size : Nbatch, Nelec x Ndim
- Returns:
nbatch, 3 x natom, nelec_pair, 3
- Return type:
torch.tensor
- forward(pos: torch.Tensor, derivative: int = 0, sum_grad: bool = True) torch.Tensor[source]
Compute the Jastrow factors.
- Parameters:
pos (torch.tensor) – Positions of the electrons Size : Nbatch, Nelec x Ndim
derivative (int, optional) – order of the derivative (0,1,2,). Defaults to 0.
sum_grad (bool, optional) – Return the sum_grad (i.e. the sum of the derivatives) or the individual terms. Defaults to True. False only for derivative=1
- Returns:
- value of the jastrow parameter for all confs
derivative = 0 (Nmo) x Nbatch x 1 derivative = 1 (Nmo) x Nbatch x Nelec (for sum_grad = True) derivative = 1 (Nmo) x Nbatch x Ndim x Nelec (for sum_grad = False) derivative = 2 (Nmo) x Nbatch x Nelec
- Return type:
torch.tensor
- jastrow_factor_derivative(r: torch.Tensor, dr: torch.Tensor, jast: torch.Tensor, sum_grad: bool) torch.Tensor[source]
Compute the value of the derivative of the Jastrow factor
- Parameters:
r (torch.tensor) – ee distance matrix Nbatch x Nelec x Nelec
jast (torch.tensor) – values of the jastrow elements Nbatch x Nelec x Nelec
- Returns:
- gradient of the jastrow factors
Nbatch x Nelec x Ndim
- Return type:
torch.tensor
- jastrow_factor_second_derivative(r: torch.Tensor, dr: torch.Tensor, d2r: torch.Tensor, jast: torch.Tensor) torch.Tensor[source]
Compute the value of the pure 2nd derivative of the Jastrow factor
- Parameters:
r (torch.tensor) – ee distance matrix Nbatch x Nelec x Nelec
jast (torch.tensor) – values of the ajstrow elements Nbatch x Nelec x Nelec
- Returns:
- diagonal hessian of the jastrow factors
Nbatch x Nelec x Ndim
- Return type:
torch.tensor
- partial_derivative(djast: torch.Tensor) torch.Tensor[source]
Compute the partial derivative of the jastrow factor
- Parameters:
djast (torch.tensor) – derivative of the jastrow elements Nbatch x Ndim x Nelec x Nelec x 3
- Returns:
- diagonal hessian of the jastrow factors
Nbatch x Nelec x Ndim
- Return type:
torch.tensor
- jastrow_factor_second_derivative_auto(pos: torch.Tensor) torch.Tensor
Compute the second derivative of the Jastrow factor using automatic differentiation.
This function utilizes PyTorch’s automatic differentiation capabilities to compute the Hessian of the Jastrow factor with respect to the electron positions. It calculates the diagonal elements of the Hessian matrix, which correspond to the second derivatives along each dimension.
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) >>> return batch_hess.view(nbatch, self.nelec, 3).sum(2) However this approach seems to requires 10x more memory
- Parameters:
pos (torch.Tensor) – Positions of the electrons, with shape (Nbatch, Nelec * Ndim).
- Returns:
- The summed diagonal elements of the Hessian matrix, with shape
(Nbatch, Nelec), representing the second derivative of the Jastrow factor for each electron.
- Return type:
Module contents
- qmctorch.wavefunction.jastrows.elec_elec_nuclei.JastrowFactor
alias of
JastrowFactorElectronElectronNuclei
- class qmctorch.wavefunction.jastrows.elec_elec_nuclei.BoysHandyJastrowKernel(*args: Any, **kwargs: Any)[source]
Bases:
JastrowKernelElectronElectronNucleiBaseDefines a Boys Handy jastrow factors.
J.W. Moskowitz et. al Correlated Monte Carlo Wave Functions for Some Cations and Anions of the First Row Atoms Journal of Chemical Physics, 97, 3382-85 (1992)
\[\begin{split}\\text{K}(R_{iA}, R_{jA}, r_{ij) = \\sum_\\mu c_\\mu \\left(\\frac{a_{1_\\mu} R_{iA}}{1 + b_{1_\\mu}R_{iA}}\\right)^{u_\\mu} \\left(\\frac{a_{2_\\mu} R_{jA}}{1 + b_{2_\\mu}R_{iA}}\\right)^{v_\\mu} \\left(\\frac{a_{3_\\mu} r_{ij}}{1 + b_{3_\\mu}r_{ij}}\\right)^{w_\\mu}\end{split}\]We restrict the parameters of the two electron-nucleus distance to be equal otherwise the jastrow factor is not permutation invariant
- forward(x: torch.Tensor) torch.Tensor[source]
Compute the values of the kernel
- Parameters:
x (torch.tensor) – e-e and e-n distances distance (Nbatch, Natom, Nelec_pairs, 3) the last dimension holds the values [R_{iA}, R_{jA}, r_{ij}] where i,j are electron index in the pair and A the atom index.
- Returns:
values of the kernel (Nbatch, Natom, Nelec_pairs, 1)
- Return type:
torch.tensor
- class qmctorch.wavefunction.jastrows.elec_elec_nuclei.FullyConnectedJastrowKernel(*args: Any, **kwargs: Any)[source]
Bases:
JastrowKernelElectronElectronNucleiBaseDefines a fully connected jastrow factors.
- Parameters:
- forward(x: torch.Tensor) torch.Tensor[source]
Compute the values of the individual f_ij=f(r_ij)
- Parameters:
x (torch.tensor) – e-e distance Nbatch, Nele_pairs
- Returns:
values of the f_ij
- Return type:
torch.tensor