14968442546

Committed 12 May 2025 09:14AM UTC coverage: 83.955%. First build

Build # 14968442546

Build Type

Pull #187

github

Committed by

web-flow

Commit Message

Merge a67f074c6 into 20fe7ebf9

Pull Request Pull Request #187: Clean up Main

Run Details

951 of 1326 branches covered (71.72%)

Branch coverage included in aggregate %.

287 of 362 new or added lines in 47 files covered. (79.28%)

4522 of 5193 relevant lines covered (87.08%)

0.87 hits per line

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

81.93

/qmctorch/wavefunction/jastrows/elec_elec/kernels/double_fully_connected_jastrow_kernel.py

import torch
from torch import nn
from .jastrow_kernel_electron_electron_base import JastrowKernelElectronElectronBase


class DoubleFullyConnectedJastrowKernel(JastrowKernelElectronElectronBase):
    def __init__(
        self,
        nup: int,
        ndown: int,
        cuda: bool,
        size1: int = 16,
        size2: int = 8,
        activation: torch.nn.Module = torch.nn.Sigmoid(),
    ) -> None:
        """Defines a fully connected jastrow factors.

        Args:
            nup (int): Number of spin up electrons.
            ndown (int): Number of spin down electrons.
            cuda (bool): Whether to use the GPU or not.
            size1 (int, optional): Number of neurons in the first hidden layer. Defaults to 16.
            size2 (int, optional): Number of neurons in the second hidden layer. Defaults to 8.
            activation (torch.nn.Module, optional): Activation function. Defaults to torch.nn.Sigmoid.
        """

        super().__init__(nup, ndown, cuda)

        self.fc1_same = nn.Linear(1, size1, bias=True)
        self.fc2_same = nn.Linear(size1, size2, bias=True)
        self.fc3_same = nn.Linear(size2, 1, bias=True)

        self.fc1_opp = nn.Linear(1, size1, bias=True)
        self.fc2_opp = nn.Linear(size1, size2, bias=True)
        self.fc3_opp = nn.Linear(size2, 1, bias=True)

        eps = 1e-3
        self.fc1_same.weight.data *= eps
        self.fc2_same.weight.data *= eps
        self.fc3_same.weight.data *= eps


        self.fc1_opp.weight.data *= eps
        self.fc2_opp.weight.data *= eps
        self.fc3_opp.weight.data *= eps

        eps = 1e-6
        self.fc1_same.bias.data *= eps
        self.fc2_same.bias.data *= eps
        self.fc3_same.bias.data *= eps


        self.fc1_opp.bias.data *= eps
        self.fc2_opp.bias.data *= eps
        self.fc3_opp.bias.data *= eps

        self.nl_func = activation

        self.get_idx_pair()

        self.requires_autograd = True

    def get_idx_pair(self) -> None:
        """define the variational weight."""

        nelec = self.nup + self.ndown

        same_idx_pair = []
        opp_idx_pair = []
        ipair = 0
        for i in range(nelec - 1):
            ispin = 0 if i < self.nup else 1
            for j in range(i + 1, nelec):
                jspin = 0 if j < self.nup else 1

                if ispin == jspin:
                    same_idx_pair.append(ipair)
                else:
                    opp_idx_pair.append(ipair)

                ipair += 1

        self.same_idx_pair = torch.as_tensor(same_idx_pair).to(self.device)
        self.opp_idx_pair = torch.as_tensor(opp_idx_pair).to(self.device)


    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = torch.zeros_like(x)
        if len(self.same_idx_pair) > 0:
            out[:, self.same_idx_pair] = self._fsame(x[:, self.same_idx_pair])
        if len(self.opp_idx_pair) > 0:
            out[:, self.opp_idx_pair] = self._fopp(x[:, self.opp_idx_pair])
        return out

    def _fsame(self, x: torch.Tensor) -> torch.Tensor:
        """Compute the kernel values of same spin pairs

        Args:
            x (torch.tensor): e-e distance Nbatch, Nele_pairs

        Returns:
            torch.tensor: values of the f_ij
        """
        nbatch, npairs = x.shape

        # reshape the input so that all elements are considered
        # independently of each other
        x = x.reshape(-1, 1)

        x = self.fc1_same(x)
        x = self.nl_func(x)
        x = self.fc2_same(x)
        x = self.nl_func(x)
        x = self.fc3_same(x)
        x = self.nl_func(x)

        # reshape to the original shape
        x = x.reshape(nbatch, npairs)

        return x

    def _fopp(self, x: torch.Tensor) -> torch.Tensor:
        """Compute the kernel values of opposite spin pairs

        Args:
            x (torch.tensor): e-e distance Nbatch, Nele_pairs

        Returns:
            torch.tensor: values of the f_ij
        """
        nbatch, npairs = x.shape

        # reshape the input so that all elements are considered
        # independently of each other
        x = x.reshape(-1, 1)

        x = self.fc1_opp(x)
        x = self.nl_func(x)
        x = self.fc2_opp(x)
        x = self.nl_func(x)
        x = self.fc3_opp(x)
        x = self.nl_func(x)

        # reshape to the original shape
        x = x.reshape(nbatch, npairs)

        return x

1	import torch	1✔
2	from torch import nn	1✔
3	from .jastrow_kernel_electron_electron_base import JastrowKernelElectronElectronBase	1✔
4
5
6	class DoubleFullyConnectedJastrowKernel(JastrowKernelElectronElectronBase):	1✔
7	def __init__(	1✔
8	self,
9	nup: int,
10	ndown: int,
11	cuda: bool,
12	size1: int = 16,
13	size2: int = 8,
14	activation: torch.nn.Module = torch.nn.Sigmoid(),
15	) -> None:
16	"""Defines a fully connected jastrow factors.
17
18	Args:
19	nup (int): Number of spin up electrons.
20	ndown (int): Number of spin down electrons.
21	cuda (bool): Whether to use the GPU or not.
22	size1 (int, optional): Number of neurons in the first hidden layer. Defaults to 16.
23	size2 (int, optional): Number of neurons in the second hidden layer. Defaults to 8.
24	activation (torch.nn.Module, optional): Activation function. Defaults to torch.nn.Sigmoid.
25	"""
26
27	super().__init__(nup, ndown, cuda)	1✔
28
29	self.fc1_same = nn.Linear(1, size1, bias=True)	1✔
30	self.fc2_same = nn.Linear(size1, size2, bias=True)	1✔
31	self.fc3_same = nn.Linear(size2, 1, bias=True)	1✔
32
33	self.fc1_opp = nn.Linear(1, size1, bias=True)	1✔
34	self.fc2_opp = nn.Linear(size1, size2, bias=True)	1✔
35	self.fc3_opp = nn.Linear(size2, 1, bias=True)	1✔
36
37	eps = 1e-3	1✔
38	self.fc1_same.weight.data *= eps	1✔
39	self.fc2_same.weight.data *= eps	1✔
40	self.fc3_same.weight.data *= eps	1✔
41
42
43	self.fc1_opp.weight.data *= eps	1✔
44	self.fc2_opp.weight.data *= eps	1✔
45	self.fc3_opp.weight.data *= eps	1✔
46
47	eps = 1e-6	1✔
48	self.fc1_same.bias.data *= eps	1✔
49	self.fc2_same.bias.data *= eps	1✔
50	self.fc3_same.bias.data *= eps	1✔
51
52
53	self.fc1_opp.bias.data *= eps	1✔
54	self.fc2_opp.bias.data *= eps	1✔
55	self.fc3_opp.bias.data *= eps	1✔
56
57	self.nl_func = activation	1✔
58
59	self.get_idx_pair()	1✔
60
61	self.requires_autograd = True	1✔
62
63	def get_idx_pair(self) -> None:	1✔
64	"""define the variational weight."""
65
66	nelec = self.nup + self.ndown	1✔
67
68	same_idx_pair = []	1✔
69	opp_idx_pair = []	1✔
70	ipair = 0	1✔
71	for i in range(nelec - 1):	1✔
72	ispin = 0 if i < self.nup else 1	1✔
73	for j in range(i + 1, nelec):	1✔
74	jspin = 0 if j < self.nup else 1	1✔
75
76	if ispin == jspin:	1!
NEW 77	same_idx_pair.append(ipair)	×
78	else:
79	opp_idx_pair.append(ipair)	1✔
80
81	ipair += 1	1✔
82
83	self.same_idx_pair = torch.as_tensor(same_idx_pair).to(self.device)	1✔
84	self.opp_idx_pair = torch.as_tensor(opp_idx_pair).to(self.device)	1✔
85
86
87	def forward(self, x: torch.Tensor) -> torch.Tensor:	1✔
88	out = torch.zeros_like(x)	1✔
89	if len(self.same_idx_pair) > 0:	1!
NEW 90	out[:, self.same_idx_pair] = self._fsame(x[:, self.same_idx_pair])	×
91	if len(self.opp_idx_pair) > 0:	1!
92	out[:, self.opp_idx_pair] = self._fopp(x[:, self.opp_idx_pair])	1✔
93	return out	1✔
94
95	def _fsame(self, x: torch.Tensor) -> torch.Tensor:	1✔
96	"""Compute the kernel values of same spin pairs
97
98	Args:
99	x (torch.tensor): e-e distance Nbatch, Nele_pairs
100
101	Returns:
102	torch.tensor: values of the f_ij
103	"""
NEW 104	nbatch, npairs = x.shape	×
105
106	# reshape the input so that all elements are considered
107	# independently of each other
NEW 108	x = x.reshape(-1, 1)	×
109
NEW 110	x = self.fc1_same(x)	×
NEW 111	x = self.nl_func(x)	×
NEW 112	x = self.fc2_same(x)	×
NEW 113	x = self.nl_func(x)	×
NEW 114	x = self.fc3_same(x)	×
NEW 115	x = self.nl_func(x)	×
116
117	# reshape to the original shape
NEW 118	x = x.reshape(nbatch, npairs)	×
119
NEW 120	return x	×
121
122	def _fopp(self, x: torch.Tensor) -> torch.Tensor:	1✔
123	"""Compute the kernel values of opposite spin pairs
124
125	Args:
126	x (torch.tensor): e-e distance Nbatch, Nele_pairs
127
128	Returns:
129	torch.tensor: values of the f_ij
130	"""
131	nbatch, npairs = x.shape	1✔
132
133	# reshape the input so that all elements are considered
134	# independently of each other
135	x = x.reshape(-1, 1)	1✔
136
137	x = self.fc1_opp(x)	1✔
138	x = self.nl_func(x)	1✔
139	x = self.fc2_opp(x)	1✔
140	x = self.nl_func(x)	1✔
141	x = self.fc3_opp(x)	1✔
142	x = self.nl_func(x)	1✔
143
144	# reshape to the original shape
145	x = x.reshape(nbatch, npairs)	1✔
146
147	return x	1✔

NLESC-JCER / QMCTorch / 14968442546

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