14928183883

Committed 09 May 2025 11:44AM UTC coverage: 84.285%. First build

Build # 14928183883

Build Type

Pull #187

github

Committed by

web-flow

Commit Message

Merge 7318ce1b4 into 20fe7ebf9

Pull Request Pull Request #187: Clean up Main

Run Details

952 of 1326 branches covered (71.79%)

Branch coverage included in aggregate %.

224 of 290 new or added lines in 27 files covered. (77.24%)

4551 of 5203 relevant lines covered (87.47%)

0.87 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

83.02

/qmctorch/wavefunction/orbitals/backflow/kernels/backflow_kernel_base.py

import torch
from torch import nn
from torch.autograd import grad
from typing import Tuple, List, Union
from .....scf import Molecule
from .....utils import gradients, hessian

class BackFlowKernelBase(nn.Module):
    def __init__(self, mol: Molecule, cuda: bool):
        """Compute the back flow kernel, i.e. the function
        f(rij) where rij is the distance between electron i and j
        This kernel is used in the backflow transformation
        .. math:
            q_i = r_i + \\sum_{j\\neq i} f(r_{ij}) (r_i-r_j)
        """
        super().__init__()
        self.nelec = mol.nelec
        self.cuda = cuda
        self.device = torch.device("cpu")
        if self.cuda:
            self.device = torch.device("cuda")

    def forward(self, ree: torch.Tensor, derivative: int = 0) -> torch.Tensor:
        """Computes the desired values of the kernel
         Args:
            ree (torch.tensor): e-e distance Nbatch x Nelec x Nelec
            derivative (int): derivative requried 0, 1, 2

        Returns:
            torch.tensor : f(r) Nbatch x Nelec x Nelec
        """

        if derivative == 0:
            return self._backflow_kernel(ree)

        elif derivative == 1:
            return self._backflow_kernel_derivative(ree)

        elif derivative == 2:
            return self._backflow_kernel_second_derivative(ree)

        else:
            raise ValueError("derivative of the kernel must be 0, 1 or 2")

    def _backflow_kernel(self, ree: torch.Tensor) -> torch.Tensor:
        """Computes the kernel via autodiff

        Args:
            ree ([type]): [description]

        Returns:
            [type]: [description]
        """
        raise NotImplementedError("Please implement the backflow kernel")

    def _backflow_kernel_derivative(self, ree: torch.Tensor) -> torch.Tensor:
        """Computes the first derivative of the kernel via autodiff

        Args:
            ree ([type]): [description]

        Returns:
            [type]: [description]
        """
        if ree.requires_grad == False:
            ree.requires_grad = True

        with torch.enable_grad():
            kernel_val = self._backflow_kernel(ree)

        return gradients(kernel_val, ree)

    def _backflow_kernel_second_derivative(self, ree: torch.Tensor) -> torch.Tensor:
        """Computes the second derivative of the kernel via autodiff

        Args:
            ree ([type]): [description]

        Returns:
            [type]: [description]
        """
        if ree.requires_grad == False:
            ree.requires_grad = True

        with torch.enable_grad():
            kernel_val = self._backflow_kernel(ree)
            hess_val, _ = hessian(kernel_val, ree)

        # if the kernel is linear, hval is None
        if hess_val is None:
            hess_val = torch.zeros_like(ree)

        return hess_val

1	import torch	1✔
2	from torch import nn	1✔
3	from torch.autograd import grad	1✔
4	from typing import Tuple, List, Union	1✔
5	from .....scf import Molecule	1✔
6	from .....utils import gradients, hessian	1✔
7
8	class BackFlowKernelBase(nn.Module):	1✔
9	def __init__(self, mol: Molecule, cuda: bool):	1✔
10	"""Compute the back flow kernel, i.e. the function
11	f(rij) where rij is the distance between electron i and j
12	This kernel is used in the backflow transformation
13	.. math:
14	q_i = r_i + \\sum_{j\\neq i} f(r_{ij}) (r_i-r_j)
15	"""
16	super().__init__()	1✔
17	self.nelec = mol.nelec	1✔
18	self.cuda = cuda	1✔
19	self.device = torch.device("cpu")	1✔
20	if self.cuda:	1!
21	self.device = torch.device("cuda")	×
22
23	def forward(self, ree: torch.Tensor, derivative: int = 0) -> torch.Tensor:	1✔
24	"""Computes the desired values of the kernel
25	Args:
26	ree (torch.tensor): e-e distance Nbatch x Nelec x Nelec
27	derivative (int): derivative requried 0, 1, 2
28
29	Returns:
30	torch.tensor : f(r) Nbatch x Nelec x Nelec
31	"""
32
33	if derivative == 0:	1✔
34	return self._backflow_kernel(ree)	1✔
35
36	elif derivative == 1:	1✔
37	return self._backflow_kernel_derivative(ree)	1✔
38
39	elif derivative == 2:	1!
40	return self._backflow_kernel_second_derivative(ree)	1✔
41
42	else:
43	raise ValueError("derivative of the kernel must be 0, 1 or 2")	×
44
45	def _backflow_kernel(self, ree: torch.Tensor) -> torch.Tensor:	1✔
46	"""Computes the kernel via autodiff
47
48	Args:
49	ree ([type]): [description]
50
51	Returns:
52	[type]: [description]
53	"""
54	raise NotImplementedError("Please implement the backflow kernel")	×
55
56	def _backflow_kernel_derivative(self, ree: torch.Tensor) -> torch.Tensor:	1✔
57	"""Computes the first derivative of the kernel via autodiff
58
59	Args:
60	ree ([type]): [description]
61
62	Returns:
63	[type]: [description]
64	"""
65	if ree.requires_grad == False:	1✔
66	ree.requires_grad = True	1✔
67
68	with torch.enable_grad():	1✔
69	kernel_val = self._backflow_kernel(ree)	1✔
70
71	return gradients(kernel_val, ree)	1✔
72
73	def _backflow_kernel_second_derivative(self, ree: torch.Tensor) -> torch.Tensor:	1✔
74	"""Computes the second derivative of the kernel via autodiff
75
76	Args:
77	ree ([type]): [description]
78
79	Returns:
80	[type]: [description]
81	"""
82	if ree.requires_grad == False:	1!
83	ree.requires_grad = True	×
84
85	with torch.enable_grad():	1✔
86	kernel_val = self._backflow_kernel(ree)	1✔
87	hess_val, _ = hessian(kernel_val, ree)	1✔
88
89	# if the kernel is linear, hval is None
90	if hess_val is None:	1!
NEW 91	hess_val = torch.zeros_like(ree)	×
92
93	return hess_val	1✔

NLESC-JCER / QMCTorch / 14928183883

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous