Skip to content

Commit

Permalink
fix f2 issue
Browse files Browse the repository at this point in the history
  • Loading branch information
NicoRenaud committed Dec 5, 2024
1 parent 5a30d0e commit 69bb397
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 19 deletions.
15 changes: 14 additions & 1 deletion qmctorch/utils/algebra_utils.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import torch

from scipy.special import factorial2 as f2

def btrace(M):
"""Computes the trace of batched matrices
Expand Down Expand Up @@ -43,6 +43,19 @@ def bdet2(M):
return M[..., 0, 0] * M[..., 1, 1] - M[..., 0, 1] * M[..., 1, 0]


def double_factorial(input):
"""Computes the double factorial of an array of int
Args:
input (List): input numbers
Returns:
List: values of the double factorial
"""
output = f2(input)
return [1 if o==0 else o for o in output]


class BatchDeterminant(torch.autograd.Function):

@staticmethod
Expand Down
28 changes: 10 additions & 18 deletions qmctorch/wavefunction/orbitals/norm_orbital.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import torch
import numpy as np

from ...utils.algebra_utils import double_factorial

def atomic_orbital_norm(basis):
"""Computes the norm of the atomic orbitals
Expand Down Expand Up @@ -77,16 +77,13 @@ def norm_gaussian_spherical(bas_n, bas_exp):
Returns:
torch.tensor: normalization factor
"""

from scipy.special import factorial2 as f2

bas_n = torch.tensor(bas_n)
bas_n = bas_n + 1.
exp1 = 0.25 * (2. * bas_n + 1.)

A = torch.tensor(bas_exp)**exp1
B = 2**(2. * bas_n + 3. / 2)
C = torch.as_tensor(f2(2 * bas_n.int() - 1) * np.pi **
C = torch.as_tensor(double_factorial(2 * bas_n.int() - 1) * np.pi **
0.5).type(torch.get_default_dtype())

return torch.sqrt(B / C) * A
Expand All @@ -106,22 +103,20 @@ def norm_slater_cartesian(a, b, c, n, exp):
Returns:
torch.tensor: normalization factor
"""
from scipy.special import factorial2 as f2

lvals = a + b + c + n + 1.

lfact = torch.as_tensor([np.math.factorial(int(2 * i))
for i in lvals]).type(torch.get_default_dtype())

prefact = 4 * np.pi * lfact / ((2 * exp)**(2 * lvals + 1))

num = torch.as_tensor(f2(2 * a.astype('int') - 1) *
f2(2 * b.astype('int') - 1) *
f2(2 * c.astype('int') - 1)
num = torch.as_tensor(double_factorial(2 * a.astype('int') - 1) *
double_factorial(2 * b.astype('int') - 1) *
double_factorial(2 * c.astype('int') - 1)
).type(torch.get_default_dtype())

denom = torch.as_tensor(
f2((2 * a + 2 * b + 2 * c + 1).astype('int')
double_factorial((2 * a + 2 * b + 2 * c + 1).astype('int')
)).type(torch.get_default_dtype())

return torch.sqrt(1. / (prefact * num / denom))
Expand All @@ -135,22 +130,19 @@ def norm_gaussian_cartesian(a, b, c, exp):
a (torch.tensor): exponent of x
b (torch.tensor): exponent of y
c (torch.tensor): exponent of z
exp (torch.tensor): Sater exponent
exp (torch.tensor): Slater exponent
Returns:
torch.tensor: normalization factor
"""

from scipy.special import factorial2 as f2

pref = torch.as_tensor((2 * exp / np.pi)**(0.75))
am1 = (2 * a - 1).astype('int')
x = (4 * exp)**(a / 2) / torch.sqrt(torch.as_tensor(f2(am1)))
x = (4 * exp)**(a / 2) / torch.sqrt(torch.as_tensor(double_factorial(am1)))

bm1 = (2 * b - 1).astype('int')
y = (4 * exp)**(b / 2) / torch.sqrt(torch.as_tensor(f2(bm1)))
y = (4 * exp)**(b / 2) / torch.sqrt(torch.as_tensor(double_factorial(bm1)))

cm1 = (2 * c - 1).astype('int')
z = (4 * exp)**(c / 2) / torch.sqrt(torch.as_tensor(f2(cm1)))
z = (4 * exp)**(c / 2) / torch.sqrt(torch.as_tensor(double_factorial(cm1)))

return (pref * x * y * z).type(torch.get_default_dtype())

0 comments on commit 69bb397

Please sign in to comment.