diff --git a/iodata/overlap.py b/iodata/overlap.py index dd7e4b5b..06725cd6 100644 --- a/iodata/overlap.py +++ b/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 diff --git a/iodata/test/test_overlap.py b/iodata/test/test_overlap.py index 797d7de4..c444bb6c 100644 --- a/iodata/test/test_overlap.py +++ b/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,50 @@ 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) + + +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)