Remove excessive basis set normalisation (#2850)

doc/releases/changelog-dev.md

@@ -268,6 +268,11 @@ of operators. [(#2622)](https://github.com/PennyLaneAI/pennylane/pull/2622)
 
 <h3>Improvements</h3>
 
+* The efficiency of the Hartree-Fock workflow is improved by removing the repetitive basis set
+  normalisation steps and modifying how the permutational symmetries are applied to avoid repetitive
+  electron repulsion integral calculations.
+  [(#2850)](https://github.com/PennyLaneAI/pennylane/pull/2850)
+
 * The coefficients of the non-differentiable molecular Hamiltonians generated with openfermion have
   `requires_grad = False` by default.
   [(#2865)](https://github.com/PennyLaneAI/pennylane/pull/2865)

pennylane/qchem/integrals.py

@@ -256,12 +256,13 @@ def gaussian_overlap(la, lb, ra, rb, alpha, beta):
     return s
 
 
-def overlap_integral(basis_a, basis_b):
+def overlap_integral(basis_a, basis_b, normalize=True):
     r"""Return a function that computes the overlap integral for two contracted Gaussian functions.
 
     Args:
         basis_a (~qchem.basis_set.BasisFunction): first basis function
         basis_b (~qchem.basis_set.BasisFunction): second basis function
+        normalize (bool): if True, the basis functions get normalized
 
     Returns:
         function: function that computes the overlap integral
@@ -291,11 +292,13 @@ def _overlap_integral(*args):
         alpha, ca, ra = _generate_params(basis_a.params, args_a)
         beta, cb, rb = _generate_params(basis_b.params, args_b)
 
-        ca = ca * primitive_norm(basis_a.l, alpha)
-        cb = cb * primitive_norm(basis_b.l, beta)
-
-        na = contracted_norm(basis_a.l, alpha, ca)
-        nb = contracted_norm(basis_b.l, beta, cb)
+        if basis_a.params[1].requires_grad or normalize:
+            ca = ca * primitive_norm(basis_a.l, alpha)
+            cb = cb * primitive_norm(basis_b.l, beta)
+            na = contracted_norm(basis_a.l, alpha, ca)
+            nb = contracted_norm(basis_b.l, beta, cb)
+        else:
+            na = nb = 1.0
 
         return (
             na
@@ -426,7 +429,7 @@ def gaussian_moment(li, lj, ri, rj, alpha, beta, order, r):
     return s
 
 
-def moment_integral(basis_a, basis_b, order, idx):
+def moment_integral(basis_a, basis_b, order, idx, normalize=True):
     r"""Return a function that computes the multipole moment integral for two contracted Gaussians.
 
     The multipole moment integral for two primitive Gaussian functions is computed as
@@ -451,6 +454,7 @@ def moment_integral(basis_a, basis_b, order, idx):
         basis_b (~qchem.basis_set.BasisFunction): right basis function
         order (integer): exponent of the position component
         idx (integer): index determining the dimension of the multipole moment integral
+        normalize (bool): if True, the basis functions get normalized
 
     Returns:
         function: function that computes the multipole moment integral
@@ -484,11 +488,13 @@ def _moment_integral(*args):
         alpha, ca, ra = _generate_params(basis_a.params, args_a)
         beta, cb, rb = _generate_params(basis_b.params, args_b)
 
-        ca = ca * primitive_norm(basis_a.l, alpha)
-        cb = cb * primitive_norm(basis_b.l, beta)
-
-        na = contracted_norm(basis_a.l, alpha, ca)
-        nb = contracted_norm(basis_b.l, beta, cb)
+        if basis_a.params[1].requires_grad or normalize:
+            ca = ca * primitive_norm(basis_a.l, alpha)
+            cb = cb * primitive_norm(basis_b.l, beta)
+            na = contracted_norm(basis_a.l, alpha, ca)
+            nb = contracted_norm(basis_b.l, beta, cb)
+        else:
+            na = nb = 1.0
 
         p = alpha[:, anp.newaxis] + beta
         q = anp.sqrt(anp.pi / p)
@@ -614,12 +620,13 @@ def gaussian_kinetic(la, lb, ra, rb, alpha, beta):
     return -0.5 * (t1 + t2 + t3)
 
 
-def kinetic_integral(basis_a, basis_b):
+def kinetic_integral(basis_a, basis_b, normalize=True):
     r"""Return a function that computes the kinetic integral for two contracted Gaussian functions.
 
     Args:
         basis_a (~qchem.basis_set.BasisFunction): first basis function
         basis_b (~qchem.basis_set.BasisFunction): second basis function
+        normalize (bool): if True, the basis functions get normalized
 
     Returns:
         function: function that computes the kinetic integral
@@ -650,11 +657,13 @@ def _kinetic_integral(*args):
         alpha, ca, ra = _generate_params(basis_a.params, args_a)
         beta, cb, rb = _generate_params(basis_b.params, args_b)
 
-        ca = ca * primitive_norm(basis_a.l, alpha)
-        cb = cb * primitive_norm(basis_b.l, beta)
-
-        na = contracted_norm(basis_a.l, alpha, ca)
-        nb = contracted_norm(basis_b.l, beta, cb)
+        if basis_a.params[1].requires_grad or normalize:
+            ca = ca * primitive_norm(basis_a.l, alpha)
+            cb = cb * primitive_norm(basis_b.l, beta)
+            na = contracted_norm(basis_a.l, alpha, ca)
+            nb = contracted_norm(basis_b.l, beta, cb)
+        else:
+            na = nb = 1.0
 
         return (
             na
@@ -815,14 +824,15 @@ def nuclear_attraction(la, lb, ra, rb, alpha, beta, r):
     return a
 
 
-def attraction_integral(r, basis_a, basis_b):
+def attraction_integral(r, basis_a, basis_b, normalize=True):
     r"""Return a function that computes the nuclear attraction integral for two contracted Gaussian
     functions.
 
     Args:
         r (array[float]): position vector of nucleus
         basis_a (~qchem.basis_set.BasisFunction): first basis function
         basis_b (~qchem.basis_set.BasisFunction): second basis function
+        normalize (bool): if True, the basis functions get normalized
 
     Returns:
         function: function that computes the electron-nuclear attraction integral
@@ -862,11 +872,13 @@ def _attraction_integral(*args):
         alpha, ca, ra = _generate_params(basis_a.params, args_a)
         beta, cb, rb = _generate_params(basis_b.params, args_b)
 
-        ca = ca * primitive_norm(basis_a.l, alpha)
-        cb = cb * primitive_norm(basis_b.l, beta)
-
-        na = contracted_norm(basis_a.l, alpha, ca)
-        nb = contracted_norm(basis_b.l, beta, cb)
+        if basis_a.params[1].requires_grad or normalize:
+            ca = ca * primitive_norm(basis_a.l, alpha)
+            cb = cb * primitive_norm(basis_b.l, beta)
+            na = contracted_norm(basis_a.l, alpha, ca)
+            nb = contracted_norm(basis_b.l, beta, cb)
+        else:
+            na = nb = 1.0
 
         v = (
             na
@@ -956,7 +968,7 @@ def electron_repulsion(la, lb, lc, ld, ra, rb, rc, rd, alpha, beta, gamma, delta
     return g
 
 
-def repulsion_integral(basis_a, basis_b, basis_c, basis_d):
+def repulsion_integral(basis_a, basis_b, basis_c, basis_d, normalize=True):
     r"""Return a function that computes the electron-electron repulsion integral for four contracted
     Gaussian functions.
 
@@ -965,6 +977,8 @@ def repulsion_integral(basis_a, basis_b, basis_c, basis_d):
         basis_b (~qchem.basis_set.BasisFunction): second basis function
         basis_c (~qchem.basis_set.BasisFunction): third basis function
         basis_d (~qchem.basis_set.BasisFunction): fourth basis function
+        normalize (bool): if True, the basis functions get normalized
+
     Returns:
         function: function that computes the electron repulsion integral
 
@@ -1003,15 +1017,18 @@ def _repulsion_integral(*args):
         gamma, cc, rc = _generate_params(basis_c.params, args_c)
         delta, cd, rd = _generate_params(basis_d.params, args_d)
 
-        ca = ca * primitive_norm(basis_a.l, alpha)
-        cb = cb * primitive_norm(basis_b.l, beta)
-        cc = cc * primitive_norm(basis_c.l, gamma)
-        cd = cd * primitive_norm(basis_d.l, delta)
+        if basis_a.params[1].requires_grad or normalize:
+            ca = ca * primitive_norm(basis_a.l, alpha)
+            cb = cb * primitive_norm(basis_b.l, beta)
+            cc = cc * primitive_norm(basis_c.l, gamma)
+            cd = cd * primitive_norm(basis_d.l, delta)
 
-        n1 = contracted_norm(basis_a.l, alpha, ca)
-        n2 = contracted_norm(basis_b.l, beta, cb)
-        n3 = contracted_norm(basis_c.l, gamma, cc)
-        n4 = contracted_norm(basis_d.l, delta, cd)
+            n1 = contracted_norm(basis_a.l, alpha, ca)
+            n2 = contracted_norm(basis_b.l, beta, cb)
+            n3 = contracted_norm(basis_c.l, gamma, cc)
+            n4 = contracted_norm(basis_d.l, delta, cd)
+        else:
+            n1 = n2 = n3 = n4 = 1.0
 
         e = (
             n1

pennylane/qchem/matrices.py

@@ -96,7 +96,7 @@ def overlap(*args):
             args_ab = []
             if args:
                 args_ab.extend(arg[[i, j]] for arg in args)
-            integral = overlap_integral(a, b)(*args_ab)
+            integral = overlap_integral(a, b, normalize=False)(*args_ab)
 
             o = anp.zeros((n, n))
             o[i, j] = o[j, i] = 1.0
@@ -148,7 +148,7 @@ def _moment_matrix(*args):
             args_ab = []
             if args:
                 args_ab.extend(arg[[i, j]] for arg in args)
-            integral = moment_integral(a, b, order, idx)(*args_ab)
+            integral = moment_integral(a, b, order, idx, normalize=False)(*args_ab)
 
             o = anp.zeros((n, n))
             o[i, j] = o[j, i] = 1.0
@@ -196,7 +196,7 @@ def kinetic(*args):
             args_ab = []
             if args:
                 args_ab.extend(arg[[i, j]] for arg in args)
-            integral = kinetic_integral(a, b)(*args_ab)
+            integral = kinetic_integral(a, b, normalize=False)(*args_ab)
 
             o = anp.zeros((n, n))
             o[i, j] = o[j, i] = 1.0
@@ -255,12 +255,16 @@ def attraction(*args):
                 for k, c in enumerate(r):
                     if c.requires_grad:
                         args_ab = [args[0][k]] + args_ab
-                    integral = integral - charges[k] * attraction_integral(c, a, b)(*args_ab)
+                    integral = integral - charges[k] * attraction_integral(
+                        c, a, b, normalize=False
+                    )(*args_ab)
                     if c.requires_grad:
                         args_ab = args_ab[1:]
             else:
                 for k, c in enumerate(r):
-                    integral = integral - charges[k] * attraction_integral(c, a, b)()
+                    integral = (
+                        integral - charges[k] * attraction_integral(c, a, b, normalize=False)()
+                    )
 
             o = anp.zeros((n, n))
             o[i, j] = o[j, i] = 1.0
@@ -310,14 +314,14 @@ def repulsion(*args):
         """
         n = len(basis_functions)
         tensor = anp.zeros((n, n, n, n))
-        e_calc = []
+        e_calc = anp.full((n, n, n, n), anp.nan)
 
         for (i, a), (j, b), (k, c), (l, d) in it.product(enumerate(basis_functions), repeat=4):
-            if (i, j, k, l) not in e_calc:
+            if anp.isnan(e_calc[(i, j, k, l)]):
                 args_abcd = []
                 if args:
                     args_abcd.extend(arg[[i, j, k, l]] for arg in args)
-                integral = repulsion_integral(a, b, c, d)(*args_abcd)
+                integral = repulsion_integral(a, b, c, d, normalize=False)(*args_abcd)
 
                 permutations = [
                     (i, j, k, l),
@@ -333,8 +337,8 @@ def repulsion(*args):
                 o = anp.zeros((n, n, n, n))
                 for perm in permutations:
                     o[perm] = 1.0
+                    e_calc[perm] = 1.0
                 tensor = tensor + integral * o
-                e_calc = e_calc + permutations
         return tensor
 
     return repulsion

pennylane/qchem/molecule.py

@@ -47,6 +47,7 @@ class Molecule:
         alpha (array[float]): exponents of the primitive Gaussian functions
         coeff (array[float]): coefficients of the contracted Gaussian functions
         r (array[float]): positions of the Gaussian functions
+        normalize (bool): if True, the basis functions get normalized
 
     **Example**
 
@@ -68,6 +69,7 @@ def __init__(
         l=None,
         alpha=None,
         coeff=None,
+        normalize=True,
     ):
 
         if basis_name not in ["sto-3g", "STO-3G", "6-31g", "6-31G"]:
@@ -92,6 +94,11 @@ def __init__(
 
         if coeff is None:
             coeff = [np.array(i[2], requires_grad=False) for i in self.basis_data]
+            if normalize:
+                coeff = [
+                    np.array(c * primitive_norm(l[i], alpha[i]), requires_grad=False)
+                    for i, c in enumerate(coeff)
+                ]
 
         r = list(
             itertools.chain(
@@ -138,7 +145,6 @@ def atomic_orbital(self, index):
         coeff = self.basis_set[index].coeff
         r = self.basis_set[index].r
 
-        coeff = coeff * primitive_norm(l, alpha)
         coeff = coeff * contracted_norm(l, alpha, coeff)
 
         lx, ly, lz = l

tests/qchem/test_integrals.py

@@ -400,7 +400,7 @@ def test_moment_integral(self, symbols, geometry, e, idx, ref):
         basis_a = mol.basis_set[0]
         basis_b = mol.basis_set[1]
         args = [p for p in [geometry] if p.requires_grad]
-        s = qchem.moment_integral(basis_a, basis_b, e, idx)(*args)
+        s = qchem.moment_integral(basis_a, basis_b, e, idx, normalize=False)(*args)
 
         assert np.allclose(s, ref)
 

tests/qchem/test_molecule.py

@@ -170,7 +170,7 @@ def test_basisset(self, symbols, geometry, l, alpha, coeff, r):
         r"""Test that the molecule object contains the correct basis set and non-default basis data
         for a given molecule.
         """
-        mol = qchem.Molecule(symbols, geometry)
+        mol = qchem.Molecule(symbols, geometry, normalize=False)
 
         assert set(map(type, mol.basis_set)) == {qchem.BasisFunction}
         assert mol.l == l