Improve performance of computing molecular integrals and Hamiltonian

pennylane/hf/hamiltonian.py

@@ -380,14 +380,14 @@ def simplify(h, cutoff=1.0e-12):
 
     c = []
     o = []
-    for i, op in enumerate(h.terms()[1]):
+    for i, op in enumerate(h.ops):
         op = qml.operation.Tensor(op).prune()
         op = qml.grouping.pauli_word_to_string(op, wire_map=wiremap)
         if op not in o:
-            c.append(h.terms()[0][i])
+            c.append(h.coeffs[i])
             o.append(op)
         else:
-            c[o.index(op)] += h.terms()[0][i]
+            c[o.index(op)] += h.coeffs[i]
 
     coeffs = []
     ops = []

pennylane/hf/integrals.py

@@ -488,15 +488,18 @@ def _boys(n, t):
 
     Args:
         n (float): order of the Boys function
-        t (float): exponent of the Boys function
+        t (array[float]): exponent of the Boys function
 
     Returns:
         float: value of the Boys function
     """
-    if t == 0.0:
-        return 1 / (2 * n + 1)
-
-    return asp.special.gammainc(n + 0.5, t) * asp.special.gamma(n + 0.5) / (2 * t ** (n + 0.5))
+    return anp.where(
+        t == 0,
+        1 / (2 * n + 1),
+        asp.special.gammainc(n + 0.5, t + (t == 0))
+        * asp.special.gamma(n + 0.5)
+        / (2 * (t + (t == 0)) ** (n + 0.5)),
+    )  # (t == 0) is added to avoid divide by zero
 
 
 def _hermite_coulomb(t, u, v, n, p, dr):
@@ -543,10 +546,7 @@ def _hermite_coulomb(t, u, v, n, p, dr):
     r = 0
 
     if t == u == v == 0:
-        f = []
-        for term in T.flatten():
-            f.append(_boys(n, term))
-        return ((-2 * p) ** n) * anp.array(f).reshape(T.shape)
+        return ((-2 * p) ** n) * _boys(n, T)
 
     if t == u == 0:
         if v > 1:
@@ -733,22 +733,19 @@ def electron_repulsion(la, lb, lc, ld, ra, rb, rc, rd, alpha, beta, gamma, delta
         + delta * rd[:, anp.newaxis, anp.newaxis, anp.newaxis, anp.newaxis]
     ) / (gamma + delta)
 
+    g_t = [expansion(l1, l2, ra[0], rb[0], alpha, beta, t) for t in range(l1 + l2 + 1)]
+    g_u = [expansion(m1, m2, ra[1], rb[1], alpha, beta, u) for u in range(m1 + m2 + 1)]
+    g_v = [expansion(n1, n2, ra[2], rb[2], alpha, beta, v) for v in range(n1 + n2 + 1)]
+    g_r = [expansion(l3, l4, rc[0], rd[0], gamma, delta, r) for r in range(l3 + l4 + 1)]
+    g_s = [expansion(m3, m4, rc[1], rd[1], gamma, delta, s) for s in range(m3 + m4 + 1)]
+    g_w = [expansion(n3, n4, rc[2], rd[2], gamma, delta, w) for w in range(n3 + n4 + 1)]
+
     g = 0.0
     lengths = [l1 + l2 + 1, m1 + m2 + 1, n1 + n2 + 1, l3 + l4 + 1, m3 + m4 + 1, n3 + n4 + 1]
     for t, u, v, r, s, w in it.product(*[range(length) for length in lengths]):
-        g = g + expansion(l1, l2, ra[0], rb[0], alpha, beta, t) * expansion(
-            m1, m2, ra[1], rb[1], alpha, beta, u
-        ) * expansion(n1, n2, ra[2], rb[2], alpha, beta, v) * expansion(
-            l3, l4, rc[0], rd[0], gamma, delta, r
-        ) * expansion(
-            m3, m4, rc[1], rd[1], gamma, delta, s
-        ) * expansion(
-            n3, n4, rc[2], rd[2], gamma, delta, w
-        ) * (
+        g = g + g_t[t] * g_u[u] * g_v[v] * g_r[r] * g_s[s] * g_w[w] * (
             (-1) ** (r + s + w)
-        ) * _hermite_coulomb(
-            t + r, u + s, v + w, 0, (p * q) / (p + q), p_ab - p_cd
-        )
+        ) * _hermite_coulomb(t + r, u + s, v + w, 0, (p * q) / (p + q), p_ab - p_cd)
 
     g = g * 2 * (anp.pi**2.5) / (p * q * anp.sqrt(p + q))
 

qchem/pennylane_qchem/qchem/hf/integrals.py

@@ -692,15 +692,18 @@ def _boys(n, t):
 
     Args:
         n (float): order of the Boys function
-        t (float): exponent of the Boys function
+        t (array[float]): exponent of the Boys function
 
     Returns:
-        float: value of the Boys function
+        (array[float]): value of the Boys function
     """
-    if t == 0.0:
-        return 1 / (2 * n + 1)
-
-    return asp.special.gammainc(n + 0.5, t) * asp.special.gamma(n + 0.5) / (2 * t ** (n + 0.5))
+    return anp.where(
+        t == 0,
+        1 / (2 * n + 1),
+        asp.special.gammainc(n + 0.5, t + (t == 0))
+        * asp.special.gamma(n + 0.5)
+        / (2 * (t + (t == 0)) ** (n + 0.5)),
+    )  # (t == 0) is added to avoid divide by zero
 
 
 def _hermite_coulomb(t, u, v, n, p, dr):
@@ -747,10 +750,7 @@ def _hermite_coulomb(t, u, v, n, p, dr):
     r = 0
 
     if t == u == v == 0:
-        f = []
-        for term in T.flatten():
-            f.append(_boys(n, term))
-        return ((-2 * p) ** n) * anp.array(f).reshape(T.shape)
+        return ((-2 * p) ** n) * _boys(n, T)
 
     if t == u == 0:
         if v > 1:
@@ -937,22 +937,19 @@ def electron_repulsion(la, lb, lc, ld, ra, rb, rc, rd, alpha, beta, gamma, delta
         + delta * rd[:, anp.newaxis, anp.newaxis, anp.newaxis, anp.newaxis]
     ) / (gamma + delta)
 
+    g_t = [expansion(l1, l2, ra[0], rb[0], alpha, beta, t) for t in range(l1 + l2 + 1)]
+    g_u = [expansion(m1, m2, ra[1], rb[1], alpha, beta, u) for u in range(m1 + m2 + 1)]
+    g_v = [expansion(n1, n2, ra[2], rb[2], alpha, beta, v) for v in range(n1 + n2 + 1)]
+    g_r = [expansion(l3, l4, rc[0], rd[0], gamma, delta, r) for r in range(l3 + l4 + 1)]
+    g_s = [expansion(m3, m4, rc[1], rd[1], gamma, delta, s) for s in range(m3 + m4 + 1)]
+    g_w = [expansion(n3, n4, rc[2], rd[2], gamma, delta, w) for w in range(n3 + n4 + 1)]
+
     g = 0.0
     lengths = [l1 + l2 + 1, m1 + m2 + 1, n1 + n2 + 1, l3 + l4 + 1, m3 + m4 + 1, n3 + n4 + 1]
     for t, u, v, r, s, w in it.product(*[range(length) for length in lengths]):
-        g = g + expansion(l1, l2, ra[0], rb[0], alpha, beta, t) * expansion(
-            m1, m2, ra[1], rb[1], alpha, beta, u
-        ) * expansion(n1, n2, ra[2], rb[2], alpha, beta, v) * expansion(
-            l3, l4, rc[0], rd[0], gamma, delta, r
-        ) * expansion(
-            m3, m4, rc[1], rd[1], gamma, delta, s
-        ) * expansion(
-            n3, n4, rc[2], rd[2], gamma, delta, w
-        ) * (
+        g = g + g_t[t] * g_u[u] * g_v[v] * g_r[r] * g_s[s] * g_w[w] * (
             (-1) ** (r + s + w)
-        ) * _hermite_coulomb(
-            t + r, u + s, v + w, 0, (p * q) / (p + q), p_ab - p_cd
-        )
+        ) * _hermite_coulomb(t + r, u + s, v + w, 0, (p * q) / (p + q), p_ab - p_cd)
 
     g = g * 2 * (anp.pi**2.5) / (p * q * anp.sqrt(p + q))
 

qchem/tests/hf/test_integrals.py

@@ -818,7 +818,7 @@ def test_gradient_repulsion(symbols, geometry, alpha, coeff):
 
 @pytest.mark.parametrize(
     ("n", "t", "f_ref"),
-    [(1.25, 0.00, 0.2857142857142857), (2.75, 1.23, 0.061750771828252976)],
+    [(2.75, np.array([0.0, 1.23]), np.array([0.15384615384615385, 0.061750771828252976]))],
 )
 def test_boys(n, t, f_ref):
     r"""Test that the Boys function is evaluated correctly."""

tests/hf/test_integrals.py

@@ -664,7 +664,7 @@ def test_gradient_repulsion(symbols, geometry, alpha, coeff):
 
 @pytest.mark.parametrize(
     ("n", "t", "f_ref"),
-    [(1.25, 0.00, 0.2857142857142857), (2.75, 1.23, 0.061750771828252976)],
+    [(2.75, np.array([0.0, 1.23]), np.array([0.15384615384615385, 0.061750771828252976]))],
 )
 def test_boys(n, t, f_ref):
     r"""Test that the Boys function is evaluated correctly."""