Merge pull request #238 from tovrstra/olp-two-obasis

Overlap integrals for a pair of basis sets

iodata/overlap.py

@@ -18,6 +18,8 @@
 # --
 """Module for computing overlap of atomic orbital basis functions."""
 
+from typing import Optional
+
 import attr
 import numpy as np
 from scipy.special import binom, factorial2
@@ -29,46 +31,83 @@
 __all__ = ['OVERLAP_CONVENTIONS', 'compute_overlap', 'gob_cart_normalization']
 
 
-# pylint: disable=too-many-nested-blocks,too-many-statements
-def compute_overlap(obasis: MolecularBasis, atcoords: np.ndarray) -> np.ndarray:
-    r"""Compute overlap matrix for the given molecular basis set.
+# 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:
+    r"""Compute overlap matrix for the given molecular basis set(s).
 
     .. math::
         \braket{\psi_{i}}{\psi_{j}}
 
+    When only one basis set is given, the overlap matrix of that basis (with
+    itself) is computed. If a second basis set (with its atomic coordinates) is
+    provided, the overlap between the two basis sets is computed.
+
     This function takes into account the requested order of the basis functions
-    in ``obasis.conventions``. Note that only L2 normalized primitives are
-    supported at the moment.
+    in ``obasis0.conventions`` (and ``obasis1.conventions``). Note that only L2
+    normalized primitives are supported at the moment.
 
     Parameters
     ----------
-    obasis
+    obasis0
         The orbital basis set.
-    atcoords
+    atcoords0
         The atomic Cartesian coordinates (including those of ghost atoms).
+    obasis1
+        An optional second orbital basis set.
+    atcoords1
+        An optional second array with atomic Cartesian coordinates
+        (including those of ghost atoms).
 
     Returns
     -------
     overlap
-        The matrix with overlap integrals, shape=(obasis.nbasis, obasis.nbasis).
+        The matrix with overlap integrals, ``shape=(obasis0.nbasis, obasis1.nbasis)``.
 
     """
-    if obasis.primitive_normalization != 'L2':
+    if obasis0.primitive_normalization != 'L2':
         raise ValueError('The overlap integrals are only implemented for L2 '
                          'normalization.')
 
-    # Initialize result
-    overlap = np.zeros((obasis.nbasis, obasis.nbasis))
-
     # Get a segmented basis, for simplicity
-    obasis = obasis.get_segmented()
+    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.")
+        obasis1 = obasis0
+        atcoords1 = atcoords0
+        identical = True
+    else:
+        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.")
+        # Get a segmented basis, for simplicity
+        obasis1 = obasis1.get_segmented()
+        identical = False
+
+    # Initialize result
+    overlap = np.zeros((obasis0.nbasis, obasis1.nbasis))
 
     # Compute the normalization constants of the Cartesian primitives, with the
     # contraction coefficients multiplied in.
-    scales = [_compute_cart_shell_normalizations(shell) * shell.coeffs
-              for shell in obasis.shells]
-
-    n_max = np.max([np.max(shell.angmoms) for shell in obasis.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]
+
+    n_max = max(np.max(shell.angmoms) for shell in obasis0.shells)
+    if not identical:
+        n_max = max(n_max, max(np.max(shell.angmoms) for shell in obasis1.shells))
     go = GaussianOverlap(n_max)
 
     # define a python ufunc (numpy function) for broadcasted calling over angular momentums
@@ -78,21 +117,25 @@ def compute_overlap(obasis: MolecularBasis, atcoords: np.ndarray) -> np.ndarray:
     begin0 = 0
 
     # pylint: disable=too-many-nested-blocks
-    for i0, shell0 in enumerate(obasis.shells):
-        r0 = atcoords[shell0.icenter]
+    for i0, shell0 in enumerate(obasis0.shells):
+        r0 = atcoords0[shell0.icenter]
         end0 = begin0 + shell0.nbasis
 
         # Loop over shell1 (lower triangular only, including diagonal)
         begin1 = 0
-        for i1, shell1 in enumerate(obasis.shells[:i0 + 1]):
-            r1 = atcoords[shell1.icenter]
+        if identical:
+            nshell1 = i0 + 1
+        else:
+            nshell1 = len(obasis1.shells)
+        for i1, shell1 in enumerate(obasis1.shells[:nshell1]):
+            r1 = atcoords1[shell1.icenter]
             end1 = begin1 + shell1.nbasis
 
             # prepare some constants to save FLOPS later on
             rij = r0 - r1
             rij_norm_sq = np.dot(rij, rij)
 
-            # Check if the result is going to signifcant
+            # Check if the result is going to significant.
             a0_min = np.min(shell0.exponents)
             a1_min = np.min(shell1.exponents)
             prefactor_max = np.exp(-a0_min * a1_min * rij_norm_sq / (a0_min + a1_min))
@@ -102,14 +145,14 @@ def compute_overlap(obasis: MolecularBasis, atcoords: np.ndarray) -> np.ndarray:
                 # arrays of angular momentums [[2, 0, 0], [0, 2, 0], ..., [0, 1, 1]]
                 n0 = np.array(list(iter_cart_alphabet(shell0.angmoms[0])))
                 n1 = np.array(list(iter_cart_alphabet(shell1.angmoms[0])))
-                result = np.zeros((n0.shape[0], n1.shape[0]))
+                shell_overlap = np.zeros((n0.shape[0], n1.shape[0]))
 
                 # Loop over primitives in shell0 (Cartesian)
-                for scales0, a0 in zip(scales[i0], shell0.exponents):
+                for shell_scales0, a0 in zip(scales0[i0], shell0.exponents):
                     a0_r0 = a0 * r0
 
                     # Loop over primitives in shell1 (Cartesian)
-                    for scales1, a1 in zip(scales[i1], shell1.exponents):
+                    for shell_scales1, a1 in zip(scales1[i1], shell1.exponents):
                         at = a0 + a1
                         prefactor = np.exp(-a0 * a1 / at * rij_norm_sq)
                         if prefactor < 1e-15:
@@ -127,27 +170,35 @@ def compute_overlap(obasis: MolecularBasis, atcoords: np.ndarray) -> np.ndarray:
                         # array operations.
                         vs = compute_overlap_1d(rn_0, rn_1, n0[:, None, :], n1[None, :, :], two_at)
                         v = np.prod(vs.astype(float), axis=2)
-                        result += v * prefactor * scales0[:, None] * scales1[None, :]
+                        v *= prefactor
+                        v *= shell_scales0[:, None]
+                        v *= shell_scales1[None, :]
+                        shell_overlap += v
 
                 # END of Cartesian coordinate system (if going to pure coordinates)
 
                 # cart to pure
                 if shell0.kinds[0] == 'p':
-                    result = np.dot(tfs[shell0.angmoms[0]], result)
+                    shell_overlap = np.dot(tfs[shell0.angmoms[0]], shell_overlap)
                 if shell1.kinds[0] == 'p':
-                    result = np.dot(result, tfs[shell1.angmoms[0]].T)
+                    shell_overlap = np.dot(shell_overlap, tfs[shell1.angmoms[0]].T)
 
                 # store lower triangular result
-                overlap[begin0:end0, begin1:end1] = result
-                # store upper triangular result
-                overlap[begin1:end1, begin0:end0] = result.T
+                overlap[begin0:end0, begin1:end1] = shell_overlap
+                if identical:
+                    # store upper triangular result
+                    overlap[begin1:end1, begin0:end0] = shell_overlap.T
 
             begin1 = end1
         begin0 = end0
 
-    permutation, signs = convert_conventions(obasis, OVERLAP_CONVENTIONS, reverse=True)
-    overlap = overlap[permutation] * signs.reshape(-1, 1)
-    overlap = overlap[:, permutation] * signs
+    permutation0, signs0 = convert_conventions(obasis0, OVERLAP_CONVENTIONS, reverse=True)
+    overlap = overlap[permutation0] * signs0.reshape(-1, 1)
+    if identical:
+        permutation1, signs1 = permutation0, signs0
+    else:
+        permutation1, signs1 = convert_conventions(obasis1, OVERLAP_CONVENTIONS, reverse=True)
+    overlap = overlap[:, permutation1] * signs1
     return overlap
 
 
@@ -176,7 +227,7 @@ 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)
+                integ = self.facts[m] / two_at ** (m / 2)   # TODO // 2
                 value += pf_i * self.binomials[n2][j] * x2 ** (n2 - j) * integ
         return value
 

iodata/test/test_overlap.py

@@ -18,13 +18,15 @@
 # --
 """Test iodata.overlap & iodata.overlap_accel modules."""
 
+import itertools
+
 import attr
 import numpy as np
 from numpy.testing import assert_allclose
-from pytest import raises
+import pytest
 
 from ..api import load_one
-from ..basis import MolecularBasis, Shell
+from ..basis import MolecularBasis, Shell, convert_conventions
 from ..overlap import compute_overlap, OVERLAP_CONVENTIONS
 
 try:
@@ -83,5 +85,62 @@ def test_load_fchk_o2_cc_pvtz_cart_num():
 def test_overlap_l1():
     dbasis = MolecularBasis([], {}, 'L1')
     atcoords = np.zeros((1, 3))
-    with raises(ValueError):
+    with pytest.raises(ValueError):
         _ = compute_overlap(dbasis, atcoords)
+    obasis = MolecularBasis([], {}, 'L2')
+    with pytest.raises(ValueError):
+        _ = compute_overlap(obasis, atcoords, dbasis, atcoords)
+
+
+def test_overlap_two_basis_exceptions():
+    with path('iodata.test.data', 'hf_sto3g.fchk') as fn_fchk:
+        data = load_one(fn_fchk)
+    with pytest.raises(TypeError):
+        compute_overlap(data.obasis, data.atcoords, data.obasis, None)
+    with pytest.raises(TypeError):
+        compute_overlap(data.obasis, data.atcoords, None, data.atcoords)
+
+
+FNS_TWO_BASIS = [
+    "h_sto3g.fchk",
+    "hf_sto3g.fchk",
+    "2h-azirine-cc.fchk",
+    "water_ccpvdz_pure_hf_g03.fchk"
+]
+
+
+@pytest.mark.parametrize("fn", FNS_TWO_BASIS)
+def test_overlap_two_basis_same(fn):
+    with path('iodata.test.data', fn) as pth:
+        mol = load_one(pth)
+    olp_a = compute_overlap(mol.obasis, mol.atcoords, mol.obasis, mol.atcoords)
+    olp_b = compute_overlap(mol.obasis, mol.atcoords)
+    assert_allclose(olp_a, olp_b, rtol=0, atol=1e-14)
+
+
+@pytest.mark.parametrize("fn0,fn1", itertools.combinations_with_replacement(FNS_TWO_BASIS, 2))
+def test_overlap_two_basis_different(fn0, fn1):
+    with path('iodata.test.data', fn0) as pth0:
+        mol0 = load_one(pth0)
+    with path('iodata.test.data', fn1) as pth1:
+        mol1 = load_one(pth1)
+    # Direct computation of the off-diagonal block.
+    olp_a = compute_overlap(mol0.obasis, mol0.atcoords, mol1.obasis, mol1.atcoords)
+    # Poor-man's approach: combine two molecules into one and compute its
+    # overlap matrix.
+    atcoords = np.concatenate([mol0.atcoords, mol1.atcoords])
+    shells = mol0.obasis.shells + [
+        attr.evolve(shell, icenter=shell.icenter + mol0.natom)
+        for shell in mol1.obasis.shells
+    ]
+    obasis = MolecularBasis(shells, OVERLAP_CONVENTIONS, "L2")
+    olp_big = compute_overlap(obasis, atcoords)
+    # Get the off-diagonal block and reorder.
+    olp_b = olp_big[:olp_a.shape[0], olp_a.shape[0]:]
+    assert olp_a.shape == olp_b.shape
+    permutation0, signs0 = convert_conventions(mol0.obasis, OVERLAP_CONVENTIONS, reverse=True)
+    olp_b = olp_b[permutation0] * signs0.reshape(-1, 1)
+    permutation1, signs1 = convert_conventions(mol1.obasis, OVERLAP_CONVENTIONS, reverse=True)
+    olp_b = olp_b[:, permutation1] * signs1
+    # Finally compare the numbers.
+    assert_allclose(olp_a, olp_b, rtol=0, atol=1e-14)