Add gamma functions to the math module

pennylane/math/__init__.py

@@ -43,6 +43,7 @@
     einsum,
     eye,
     frobenius_inner_product,
+    gammainc,
     get_trainable_indices,
     ones_like,
     scatter,

pennylane/math/multi_dispatch.py

@@ -818,3 +818,37 @@ def expm(tensor, like=None):
     from scipy.linalg import expm as scipy_expm
 
     return scipy_expm(tensor)
+
+
+@multi_dispatch(argnum=[1])
+def gammainc(m, t, like=None):
+    r"""Return the lower incomplete Gamma function.
+
+    The lower incomplete Gamma function is defined in scipy as
+
+    .. math::
+
+        \gamma(m, t) = \frac{1}{\Gamma(m)} \int_{0}^{t} x^{m-1} e^{-x} dx,
+
+    where :math:`\Gamma` denotes the Gamma function.
+
+     Args:
+        m (float): exponent of the incomplete Gamma function
+        t (array[float]): upper limit of the incomplete Gamma function
+
+    Returns:
+        (array[float]): value of the incomplete Gamma function
+    """
+    if like == "jax":
+        from jax.scipy.special import gammainc
+
+        return gammainc(m, t)
+
+    if like == "autograd":
+        from autograd.scipy.special import gammainc
+
+        return gammainc(m, t)
+
+    import scipy
+
+    return scipy.special.gammainc(m, t)

pennylane/math/single_dispatch.py

@@ -87,6 +87,10 @@ def _cond(pred, true_fn, false_fn, args):
 ar.register_function("numpy", "cond", _cond)
 ar.register_function("builtins", "cond", _cond)
 
+ar.register_function("numpy", "gamma", lambda x: _i("scipy").special.gamma(x))
+
+ar.register_function("builtins", "gamma", lambda x: _i("scipy").special.gamma(x))
+
 # -------------------------------- Autograd --------------------------------- #
 
 
@@ -186,6 +190,8 @@ def _take_autograd(tensor, indices, axis=None):
 ar.register_function("autograd", "diagonal", lambda x, *args: _i("qml").numpy.diag(x))
 ar.register_function("autograd", "cond", _cond)
 
+ar.register_function("autograd", "gamma", lambda x: _i("autograd.scipy").special.gamma(x))
+
 
 # -------------------------------- TensorFlow --------------------------------- #
 
@@ -684,3 +690,7 @@ def _scatter_jax(indices, array, new_dimensions):
     "cond",
     lambda pred, true_fn, false_fn, args: _i("jax").lax.cond(pred, true_fn, false_fn, *args),
 )
+
+ar.register_function(
+    "jax", "gamma", lambda x: _i("jax").numpy.exp(_i("jax").scipy.special.gammaln(x))
+)

pennylane/pauli/utils.py

@@ -1301,7 +1301,8 @@ def simplify(h, cutoff=1.0e-12):
             c[o.index(op)] += h.coeffs[i]
 
     coeffs, ops = [], []
-    nonzero_ind = np.argwhere(abs(np.array(c)) > cutoff).flatten()
+    c = qml.math.convert_like(c, c[0])
+    nonzero_ind = qml.math.argwhere(abs(c) > cutoff).flatten()
     for i in nonzero_ind:
         coeffs.append(c[i])
         ops.append(qml.pauli.string_to_pauli_word(o[i], wire_map=wiremap))
@@ -1311,7 +1312,7 @@ def simplify(h, cutoff=1.0e-12):
     except ValueError:
         pass
 
-    return qml.Hamiltonian(np.array(coeffs), ops)
+    return qml.Hamiltonian(qml.math.array(coeffs), ops)
 
 
 def _pauli_mult(p1, p2):

pennylane/qchem/integrals.py

@@ -17,7 +17,6 @@
 # pylint: disable= unbalanced-tuple-unpacking, too-many-arguments
 import itertools as it
 
-import autograd.scipy as asp
 from scipy.special import factorial2 as fac2
 
 import pennylane as qml
@@ -172,8 +171,8 @@ def expansion(la, lb, ra, rb, alpha, beta, t):
     Args:
         la (integer): angular momentum component for the first Gaussian function
         lb (integer): angular momentum component for the second Gaussian function
-        ra (float): position component of the the first Gaussian function
-        rb (float): position component of the the second Gaussian function
+        ra (float): position component of the first Gaussian function
+        rb (float): position component of the second Gaussian function
         alpha (array[float]): exponent of the first Gaussian function
         beta (array[float]): exponent of the second Gaussian function
         t (integer): number of nodes in the Hermite Gaussian
@@ -193,8 +192,7 @@ def expansion(la, lb, ra, rb, alpha, beta, t):
     array([1.])
     """
     p = alpha + beta
-    q = alpha * beta / p
-    q = np.array(q, requires_grad=True)
+    q = qml.math.array(alpha * beta / p)
     r = ra - rb
 
     if la == lb == t == 0:
@@ -536,8 +534,8 @@ def _diff2(i, j, ri, rj, alpha, beta):
     Args:
         i (integer): angular momentum component for the first Gaussian function
         j (integer): angular momentum component for the second Gaussian function
-        ri (array[float]): position component of the the first Gaussian function
-        rj (array[float]): position component of the the second Gaussian function
+        ri (array[float]): position component of the first Gaussian function
+        rj (array[float]): position component of the second Gaussian function
         alpha (array[float]): exponent of the first Gaussian function
         beta (array[float]): exponent of the second Gaussian function
 
@@ -576,8 +574,8 @@ def gaussian_kinetic(la, lb, ra, rb, alpha, beta):
     Args:
         la (tuple[int]): angular momentum for the first Gaussian function
         lb (tuple[int]): angular momentum for the second Gaussian function
-        ra (array[float]): position vector of the the first Gaussian function
-        rb (array[float]): position vector of the the second Gaussian function
+        ra (array[float]): position vector of the first Gaussian function
+        rb (array[float]): position vector of the second Gaussian function
         alpha (array[float]): exponent of the first Gaussian function
         beta (array[float]): exponent of the second Gaussian function
 
@@ -711,11 +709,11 @@ def _boys(n, t):
     Returns:
         (array[float]): value of the Boys function
     """
-    return np.where(
+    return qml.math.where(
         t == 0.0,
         1 / (2 * n + 1),
-        asp.special.gammainc(n + 0.5, t + (t == 0.0))
-        * asp.special.gamma(n + 0.5)
+        qml.math.gammainc(n + 0.5, t + (t == 0.0))
+        * qml.math.gamma(n + 0.5)
         / (2 * (t + (t == 0.0)) ** (n + 0.5)),
     )  # (t == 0.0) is added to avoid divide by zero
 
@@ -803,8 +801,8 @@ def nuclear_attraction(la, lb, ra, rb, alpha, beta, r):
     Args:
         la (tuple[int]): angular momentum for the first Gaussian function
         lb (tuple[int]): angular momentum for the second Gaussian function
-        ra (array[float]): position vector of the the first Gaussian function
-        rb (array[float]): position vector of the the second Gaussian function
+        ra (array[float]): position vector of the first Gaussian function
+        rb (array[float]): position vector of the second Gaussian function
         alpha (array[float]): exponent of the first Gaussian function
         beta (array[float]): exponent of the second Gaussian function
         r (array[float]): position vector of nucleus
@@ -922,10 +920,10 @@ def electron_repulsion(la, lb, lc, ld, ra, rb, rc, rd, alpha, beta, gamma, delta
         lb (tuple[int]): angular momentum for the second Gaussian function
         lc (tuple[int]): angular momentum for the third Gaussian function
         ld (tuple[int]): angular momentum for the forth Gaussian function
-        ra (array[float]): position vector of the the first Gaussian function
-        rb (array[float]): position vector of the the second Gaussian function
-        rc (array[float]): position vector of the the third Gaussian function
-        rd (array[float]): position vector of the the forth Gaussian function
+        ra (array[float]): position vector of the first Gaussian function
+        rb (array[float]): position vector of the second Gaussian function
+        rc (array[float]): position vector of the third Gaussian function
+        rd (array[float]): position vector of the forth Gaussian function
         alpha (array[float]): exponent of the first Gaussian function
         beta (array[float]): exponent of the second Gaussian function
         gamma (array[float]): exponent of the third Gaussian function

pennylane/qchem/observable_hf.py

@@ -42,7 +42,7 @@ def fermionic_observable(constant, one=None, two=None, cutoff=1.0e-12):
     >>> ops
     [[], [0, 0], [0, 2], [1, 1], [1, 3], [2, 0], [2, 2], [3, 1], [3, 3]]
     """
-    coeffs = np.array([])
+    coeffs = qml.math.array([])
 
     if constant != np.array([0.0]):
         coeffs = qml.math.concatenate((coeffs, constant))
@@ -55,11 +55,12 @@ def fermionic_observable(constant, one=None, two=None, cutoff=1.0e-12):
         # up-up + down-down terms
         operators_one = (indices_one * 2).tolist() + (indices_one * 2 + 1).tolist()
         coeffs_one = qml.math.tile(one[abs(one) >= cutoff], 2)
+        coeffs = qml.math.convert_like(coeffs, one)
         coeffs = qml.math.concatenate((coeffs, coeffs_one))
         operators = operators + operators_one
 
     if two is not None:
-        indices_two = qml.math.argwhere(abs(two) >= cutoff)
+        indices_two = np.array(qml.math.argwhere(abs(two) >= cutoff))
         n = len(indices_two)
         operators_two = (
             [(indices_two[i] * 2).tolist() for i in range(n)]  # up-up-up-up
@@ -73,6 +74,8 @@ def fermionic_observable(constant, one=None, two=None, cutoff=1.0e-12):
         operators = operators + operators_two
 
     indices_sort = [operators.index(i) for i in sorted(operators)]
+    if len(indices_sort) != 0:
+        indices_sort = qml.math.array(indices_sort)
 
     return coeffs[indices_sort], sorted(operators)
 

tests/math/test_multi_dispatch.py

@@ -168,3 +168,30 @@ def test_kron():
     assert np.allclose(expected_result, fn.kron(m1, m2, like="scipy"))
     with pytest.warns(DeprecationWarning):
         _ = s.kron(m1, m2)
+
+
+@pytest.mark.parametrize(
+    ("n", "t", "gamma_ref"),
+    [
+        (
+            0.1,
+            jnp.array([0.2, 0.3, 0.4]),
+            jnp.array([0.87941963, 0.90835799, 0.92757383]),
+        ),
+        (
+            0.1,
+            np.array([0.2, 0.3, 0.4]),
+            np.array([0.87941963, 0.90835799, 0.92757383]),
+        ),
+        (
+            0.1,
+            onp.array([0.2, 0.3, 0.4]),
+            onp.array([0.87941963, 0.90835799, 0.92757383]),
+        ),
+    ],
+)
+def test_gammainc(n, t, gamma_ref):
+    """Test that the lower incomplete Gamma function is computed correctly."""
+    gamma = fn.gammainc(n, t)
+
+    assert np.allclose(gamma, gamma_ref)

tests/qchem/test_hamiltonians.py

@@ -322,3 +322,53 @@ def circuit(*args):
     grad_finitediff = (e_2 - e_1) / 0.0002
 
     assert np.allclose(grad_qml[0][0], grad_finitediff)
+
+
+class TestJax:
+    @pytest.mark.jax
+    def test_gradient_expvalH(self):
+        r"""Test that the gradient of expval(H) computed with ``jax.grad`` is equal to the value
+        obtained with the finite difference method."""
+        import jax
+
+        symbols = ["H", "H"]
+        geometry = (
+            np.array([[0.0, 0.0, -0.3674625962], [0.0, 0.0, 0.3674625962]], requires_grad=False)
+            / 0.529177210903
+        )
+        alpha = np.array(
+            [[3.42525091, 0.62391373, 0.1688554], [3.42525091, 0.62391373, 0.1688554]],
+            requires_grad=True,
+        )
+
+        mol = qchem.Molecule(symbols, geometry, alpha=alpha)
+        args = [jax.numpy.array(alpha)]
+        dev = qml.device("default.qubit", wires=4)
+
+        def energy(mol):
+            @qml.qnode(dev, interface="jax")
+            def circuit(*args):
+                qml.PauliX(0)
+                qml.PauliX(1)
+                qml.DoubleExcitation(0.22350048111151138, wires=[0, 1, 2, 3])
+                h_qubit = qchem.diff_hamiltonian(mol)(*args)
+                return qml.expval(h_qubit)
+
+            return circuit
+
+        grad_jax = jax.grad(energy(mol), argnums=0)(*args)
+
+        alpha_1 = np.array(
+            [[3.42515091, 0.62391373, 0.1688554], [3.42525091, 0.62391373, 0.1688554]],
+        )  # alpha[0][0] -= 0.0001
+
+        alpha_2 = np.array(
+            [[3.42535091, 0.62391373, 0.1688554], [3.42525091, 0.62391373, 0.1688554]],
+        )  # alpha[0][0] += 0.0001
+
+        e_1 = energy(mol)(*[alpha_1])
+        e_2 = energy(mol)(*[alpha_2])
+
+        grad_finitediff = (e_2 - e_1) / 0.0002
+
+        assert np.allclose(grad_jax[0][0], grad_finitediff, rtol=1e-02)

tests/qchem/test_hartree_fock.py

@@ -239,3 +239,29 @@ def test_nuclear_energy_gradient(symbols, geometry, g_ref):
     args = [mol.coordinates]
     g = qml.grad(qchem.nuclear_energy(mol.nuclear_charges, mol.coordinates))(*args)
     assert np.allclose(g, g_ref)
+
+
+class TestJax:
+    @pytest.mark.parametrize(
+        ("symbols", "geometry", "g_ref"),
+        [
+            (
+                ["H", "H"],
+                np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]], requires_grad=True),
+                # HF gradient computed with pyscf using rnuc_grad_method().kernel()
+                np.array([[0.0, 0.0, 0.3650435], [0.0, 0.0, -0.3650435]]),
+            ),
+        ],
+    )
+    @pytest.mark.jax
+    def test_hf_energy_gradient(self, symbols, geometry, g_ref):
+        r"""Test that the gradient of the Hartree-Fock energy wrt differentiable parameters is
+        correct."""
+        import jax
+
+        mol = qchem.Molecule(symbols, geometry)
+        args = [jax.numpy.array(mol.coordinates)]
+        g = jax.grad(qchem.hf_energy(mol))(*args)
+        g_ref = jax.numpy.array(g_ref)
+
+        assert np.allclose(g, g_ref)