Fix #309 for the factorial2 function

Handling of int float and arrays

Author: D-TheProgrammer

iodata/overlap.py

@@ -32,32 +32,52 @@
 
 
 def factorial2(n, exact=False):
-    """Wrap scipy.special.factorial2 to return 1.0 when the input is -1.
+    """Wrap scipy.special.factorial2 to return 1.0 when the input is -1 and handle float arrays.
 
     This is a temporary workaround while we wait for Scipy's update.
     To learn more, see https://github.com/scipy/scipy/issues/18409.
 
     Parameters
     ----------
-    n : int or np.ndarray
+    n : int, float, or np.ndarray
         Values to calculate n!! for. If n={0, -1}, the return value is 1.
         For n < -1, the return value is 0.
     """
-    # Scipy  1.11.x returns an integer when n is an integer, but 1.10.x returns an array,
-    # so np.array(n) is passed to make sure the output is always an array.
-    out = scipy.special.factorial2(np.array(n), exact=exact)
-    out[out <= 0] = 1.0
-    out[out <= -2] = 0.0
-    return out
+    # Handle integer inputs
+    if isinstance(n, (int, np.integer)):
+        if n == -1 or n == 0:
+            return 1.0
+        elif n < -1:
+            return 0.0
+        else:
+            return scipy.special.factorial2(n, exact=exact)
+    
+    # Handle float inputs
+    elif isinstance(n, float):
+        if n == -1.0 or n == 0.0:
+            return 1.0
+        elif n < -1.0:
+            return 0.0
+        else:
+            return scipy.special.factorial2(int(n), exact=exact)
+    
+    # Handle array inputs
+    elif isinstance(n, np.ndarray):
+        result = np.zeros_like(n, dtype=float)
+        for i, val in np.ndenumerate(n):
+            if val == -1.0 or val == 0.0:
+                result[i] = 1.0
+            elif val < -1.0:
+                result[i] = 0.0
+            else:
+                result[i] = scipy.special.factorial2(int(val), exact=exact)
+        return result
 
 
 # pylint: disable=too-many-nested-blocks,too-many-statements,too-many-branches
-def compute_overlap(
-    obasis0: MolecularBasis,
-    atcoords0: np.ndarray,
-    obasis1: Optional[MolecularBasis] = None,
-    atcoords1: Optional[np.ndarray] = None,
-) -> np.ndarray:
+def compute_overlap(obasis0: MolecularBasis, atcoords0: np.ndarray,
+                    obasis1: Optional[MolecularBasis] = None,
+                    atcoords1: Optional[np.ndarray] = None,) -> np.ndarray:
     r"""Compute overlap matrix for the given molecular basis set(s).
 
     .. math::
@@ -89,30 +109,28 @@ def compute_overlap(
         The matrix with overlap integrals, ``shape=(obasis0.nbasis, obasis1.nbasis)``.
 
     """
-    if obasis0.primitive_normalization != "L2":
-        raise ValueError("The overlap integrals are only implemented for L2 " "normalization.")
+    if obasis0.primitive_normalization != 'L2':
+        raise ValueError('The overlap integrals are only implemented for L2 '
+                         'normalization.')
 
     # Get a segmented basis, for simplicity
     obasis0 = obasis0.get_segmented()
 
     # Handle optional arguments
     if obasis1 is None:
         if atcoords1 is not None:
-            raise TypeError(
-                "When no second basis is given, no second second "
-                "array of atomic coordinates is expected."
-            )
+            raise TypeError("When no second basis is given, no second second "
+                            "array of atomic coordinates is expected.")
         obasis1 = obasis0
         atcoords1 = atcoords0
         identical = True
     else:
-        if obasis1.primitive_normalization != "L2":
-            raise ValueError("The overlap integrals are only implemented for L2 " "normalization.")
+        if obasis1.primitive_normalization != 'L2':
+            raise ValueError('The overlap integrals are only implemented for L2 '
+                             'normalization.')
         if atcoords1 is None:
-            raise TypeError(
-                "When a second basis is given, a second second "
-                "array of atomic coordinates is expected."
-            )
+            raise TypeError("When a second basis is given, a second second "
+                            "array of atomic coordinates is expected.")
         # Get a segmented basis, for simplicity
         obasis1 = obasis1.get_segmented()
         identical = False
@@ -122,13 +140,13 @@ def compute_overlap(
 
     # Compute the normalization constants of the Cartesian primitives, with the
     # contraction coefficients multiplied in.
-    scales0 = [_compute_cart_shell_normalizations(shell) * shell.coeffs for shell in obasis0.shells]
+    scales0 = [_compute_cart_shell_normalizations(shell) * shell.coeffs
+               for shell in obasis0.shells]
     if identical:
         scales1 = scales0
     else:
-        scales1 = [
-            _compute_cart_shell_normalizations(shell) * shell.coeffs for shell in obasis1.shells
-        ]
+        scales1 = [_compute_cart_shell_normalizations(shell) * shell.coeffs
+                   for shell in obasis1.shells]
 
     n_max = max(np.max(shell.angmoms) for shell in obasis0.shells)
     if not identical:
@@ -203,9 +221,9 @@ def compute_overlap(
                 # END of Cartesian coordinate system (if going to pure coordinates)
 
                 # cart to pure
-                if shell0.kinds[0] == "p":
+                if shell0.kinds[0] == 'p':
                     shell_overlap = np.dot(tfs[shell0.angmoms[0]], shell_overlap)
-                if shell1.kinds[0] == "p":
+                if shell1.kinds[0] == 'p':
                     shell_overlap = np.dot(shell_overlap, tfs[shell1.angmoms[0]].T)
 
                 # store lower triangular result
@@ -254,12 +272,12 @@ def compute_overlap_gaussian_1d(self, x1, x2, n1, n2, two_at):
             pf_i = self.binomials[n1][i] * x1 ** (n1 - i)
             for j in range(i % 2, n2 + 1, 2):
                 m = i + j
-                integ = self.facts[m] / two_at ** (m / 2)  # TODO // 2
+                integ = self.facts[m] / two_at ** (m / 2)   # TODO // 2
                 value += pf_i * self.binomials[n2][j] * x2 ** (n2 - j) * integ
         return value
 
 
-def _compute_cart_shell_normalizations(shell: "Shell") -> np.ndarray:
+def _compute_cart_shell_normalizations(shell: 'Shell') -> np.ndarray:
     """Return normalization constants for the primitives in a given shell.
 
     Parameters
@@ -274,7 +292,7 @@ def _compute_cart_shell_normalizations(shell: "Shell") -> np.ndarray:
         shell is pure.
 
     """
-    shell = attr.evolve(shell, kinds=["c"] * shell.ncon)
+    shell = attr.evolve(shell, kinds=['c'] * shell.ncon)
     result = []
     for angmom in shell.angmoms:
         for exponent in shell.exponents:
@@ -301,6 +319,5 @@ def gob_cart_normalization(alpha: np.ndarray, n: np.ndarray) -> np.ndarray:
         The normalization constant for the gaussian cartesian basis.
 
     """
-    return np.sqrt(
-        (4 * alpha) ** sum(n) * (2 * alpha / np.pi) ** 1.5 / np.prod(factorial2(2 * n - 1))
-    )
+    return np.sqrt((4 * alpha) ** sum(n) * (2 * alpha / np.pi) ** 1.5
+                   / np.prod(factorial2(2 * n - 1)))