mlmm-toolkit 0.2.2.dev0__py3-none-any.whl

pysisyphus/wavefunction/ints/self_ovlp3d.py

@@ -0,0 +1,198 @@
+"""
+Molecular integrals over Gaussian basis functions generated by sympleints.
+See https://github.com/eljost/sympleints for more information.
+
+sympleints version: 0.1.dev79+g63f1ef8.d20230515
+symppy version: 1.10.1
+
+sympleints was executed with the following arguments:
+	lmax = 4
+	lauxmax = 6
+	write = False
+	out_dir = devel_ints
+	keys = ['~2c2e', '~3c2e_sph']
+	sph = False
+	opt_basic = True
+	normalize = cgto
+"""
+
+import numpy
+
+
+def self_ovlp3d_00(ax, da, A, bx, db, B):
+    """Cartesian 3D (ss) self overlap.
+
+    Generated code; DO NOT modify by hand!"""
+
+    result = numpy.zeros((1, 1), dtype=float)
+
+    x0 = (ax + bx) ** (-1.0)
+
+    # 1 item(s)
+    result[0, 0] = numpy.sum(5.568327996831708 * da * db * x0**1.5)
+    return result
+
+
+def self_ovlp3d_11(ax, da, A, bx, db, B):
+    """Cartesian 3D (pp) self overlap.
+
+    Generated code; DO NOT modify by hand!"""
+
+    result = numpy.zeros((3, 3), dtype=float)
+
+    x0 = ax + bx
+    x1 = 2.784163998415854 * da * db * numpy.sqrt(x0 ** (-1.0)) / x0**2
+
+    # 3 item(s)
+    result[0, 0] = numpy.sum(x1)
+    result[0, 1] = numpy.sum(x1)
+    result[0, 2] = numpy.sum(x1)
+    return result
+
+
+def self_ovlp3d_22(ax, da, A, bx, db, B):
+    """Cartesian 3D (dd) self overlap.
+
+    Generated code; DO NOT modify by hand!"""
+
+    result = numpy.zeros((6, 6), dtype=float)
+
+    x0 = ax + bx
+    x1 = x0 ** (-1.0)
+    x2 = 1.392081999207927 * da * db
+    x3 = numpy.sqrt(x1) * x2 / x0**3
+    x4 = x1**1.5 * x2 / x0**2
+
+    # 6 item(s)
+    result[0, 0] = numpy.sum(x3)
+    result[0, 1] = numpy.sum(x4)
+    result[0, 2] = numpy.sum(x4)
+    result[0, 3] = numpy.sum(x3)
+    result[0, 4] = numpy.sum(x4)
+    result[0, 5] = numpy.sum(x3)
+    return result
+
+
+def self_ovlp3d_33(ax, da, A, bx, db, B):
+    """Cartesian 3D (ff) self overlap.
+
+    Generated code; DO NOT modify by hand!"""
+
+    result = numpy.zeros((10, 10), dtype=float)
+
+    x0 = ax + bx
+    x1 = x0 ** (-1.0)
+    x2 = 0.6960409996039635 * da * db
+    x3 = numpy.sqrt(x1) * x2 / x0**4
+    x4 = x1**1.5 * x2 / x0**3
+
+    # 10 item(s)
+    result[0, 0] = numpy.sum(x3)
+    result[0, 1] = numpy.sum(x4)
+    result[0, 2] = numpy.sum(x4)
+    result[0, 3] = numpy.sum(x4)
+    result[0, 4] = numpy.sum(x4)
+    result[0, 5] = numpy.sum(x4)
+    result[0, 6] = numpy.sum(x3)
+    result[0, 7] = numpy.sum(x4)
+    result[0, 8] = numpy.sum(x4)
+    result[0, 9] = numpy.sum(x3)
+    return result
+
+
+def self_ovlp3d_44(ax, da, A, bx, db, B):
+    """Cartesian 3D (gg) self overlap.
+
+    Generated code; DO NOT modify by hand!"""
+
+    result = numpy.zeros((15, 15), dtype=float)
+
+    x0 = 0.5 / (ax + bx)
+    x1 = (ax + bx) ** (-1.0)
+    x2 = (x1 * (ax + bx) - 1.0) * A[0]
+    x3 = 1.772453850905516 * numpy.sqrt(x1)
+    x4 = x2**2 * x3
+    x5 = x0 * x3
+    x6 = 3.0 * x5
+    x7 = x0 * (3.0 * x4 + x6)
+    x8 = x4 + x5
+    x9 = x2 * x8
+    x10 = x2 * x5
+    x11 = x2 * (2.0 * x10 + x9)
+    x12 = 5.0 * x0 * (x11 + x7)
+    x13 = 4.0 * x0 * (2.0 * x10 + x9)
+    x14 = x11 + x7
+    x15 = x14 * x2
+    x16 = x2 * (x13 + x15)
+    x17 = x12 + x16
+    x18 = da * db
+    x19 = 0.02991993003418851 * x1 * x18
+    x20 = (x1 * (ax + bx) - 1.0) * A[1]
+    x21 = x20**2 * x3
+    x22 = x21 + x5
+    x23 = x18 * x22
+    x24 = 0.06666666666666667 * x3
+    x25 = x17 * x24
+    x26 = (x1 * (ax + bx) - 1.0) * A[2]
+    x27 = x26**2 * x3
+    x28 = x27 + x5
+    x29 = x18 * x28
+    x30 = x0 * (3.0 * x21 + x6)
+    x31 = x20 * x22
+    x32 = 2.0 * x5
+    x33 = x20 * (x20 * x32 + x31)
+    x34 = x30 + x33
+    x35 = x18 * x34
+    x36 = 0.1111111111111111 * x14 * x3
+    x37 = 0.3333333333333333 * x28
+    x38 = x0 * (3.0 * x27 + x6)
+    x39 = x26 * x28
+    x40 = x26 * (x26 * x32 + x39)
+    x41 = x38 + x40
+    x42 = 5.0 * x0 * (x30 + x33)
+    x43 = 8.0 * x5
+    x44 = x0 * (x20 * x43 + 4.0 * x31)
+    x45 = x20 * x34
+    x46 = x20 * (x44 + x45)
+    x47 = x42 + x46
+    x48 = x24 * x47
+    x49 = x18 * x8
+    x50 = 5.0 * x0 * (x38 + x40)
+    x51 = x0 * (x26 * x43 + 4.0 * x39)
+    x52 = x26 * x41
+    x53 = x26 * (x51 + x52)
+    x54 = x50 + x53
+    x55 = x24 * x54
+
+    # 15 item(s)
+    result[0, 0] = numpy.sum(
+        x19 * (7.0 * x0 * (x12 + x16) + x2 * (6.0 * x0 * (x13 + x15) + x17 * x2))
+    )
+    result[0, 1] = numpy.sum(x23 * x25)
+    result[0, 2] = numpy.sum(x25 * x29)
+    result[0, 3] = numpy.sum(x35 * x36)
+    result[0, 4] = numpy.sum(x14 * x23 * x37)
+    result[0, 5] = numpy.sum(x18 * x36 * x41)
+    result[0, 6] = numpy.sum(x48 * x49)
+    result[0, 7] = numpy.sum(x35 * x37 * x8)
+    result[0, 8] = numpy.sum(0.3333333333333333 * x23 * x41 * x8)
+    result[0, 9] = numpy.sum(x49 * x55)
+    result[0, 10] = numpy.sum(
+        x19 * (7.0 * x0 * (x42 + x46) + x20 * (6.0 * x0 * (x44 + x45) + x20 * x47))
+    )
+    result[0, 11] = numpy.sum(x29 * x48)
+    result[0, 12] = numpy.sum(0.1111111111111111 * x3 * x35 * x41)
+    result[0, 13] = numpy.sum(x23 * x55)
+    result[0, 14] = numpy.sum(
+        x19 * (7.0 * x0 * (x50 + x53) + x26 * (6.0 * x0 * (x51 + x52) + x26 * x54))
+    )
+    return result
+
+
+self_ovlp3d = {
+    (0, 0): self_ovlp3d_00,
+    (1, 1): self_ovlp3d_11,
+    (2, 2): self_ovlp3d_22,
+    (3, 3): self_ovlp3d_33,
+    (4, 4): self_ovlp3d_44,
+}

pysisyphus/wavefunction/localization.py

@@ -0,0 +1,458 @@
+# [1] https://doi.org/10.1002/jcc.540140615
+#     Comparison of the Boys and Pipek–Mezey localizations in the local
+#     correlation approach and automatic virtual basis selection
+#     Boughton, Pulay, 1993
+# [2] https://doi.org/10.1063/1.2360264
+#     Fast noniterative orbital localization for large molecules
+#     Aquilante, Pedersen, 2006
+# [3] https://doi.org/10.1063/1.3042233
+#     Constructing diabatic states from adiabatic states: Extending
+#     generalized Mulliken–Hush to multiple charge centers with Boys localization
+#     Subotnik, Yeganeh, Cave, Ratner, 2008
+# [4] https://doi.org/10.1063/1.1681683
+#     Localized molecular orbitals for polyatomic molecules.
+#     I. A comparison of the Edmiston‐Ruedenberg and Boys localization methods
+#     Kleier, Halgren, Hall Jr., Lipscomb, 1974
+# [5] https://doi.org/10.1063/1.4894472
+#     Diabatization based on the dipole and quadrupole: The DQ method
+#     Hoyer, Xu, Ma, Gagliardi, Truhlar, 2014
+# [6] https://pubs.acs.org/doi/pdf/10.1021/acs.jctc.2c00261
+#     Implementation of Occupied and Virtual Edmiston−Ruedenberg
+#     Orbitals Using Cholesky Decomposed Integrals
+# [7] https://doi.org/10.1063/1.1790971
+#     An efficient method for calculating maxima of homogeneous
+#     functions of orthogonal matrices: Applications to localized
+#     occupied orbitals
+
+
+from dataclasses import dataclass
+from functools import singledispatch
+import itertools as it
+from typing import Callable, Optional
+import warnings
+
+import numpy as np
+from numpy.typing import NDArray
+
+from pysisyphus.helpers_pure import rms
+from pysisyphus.linalg import matrix_power, pivoted_cholesky
+from pysisyphus.wavefunction import logger, Wavefunction
+from pysisyphus.wavefunction.DIIS import DIIS
+
+
+PI_QUART = np.pi / 4
+
+
+@singledispatch
+def cholesky(C: NDArray[float]):
+    """Localization via pivoted Cholesky factorization.
+
+    See [2].
+
+    Parameters
+    ----------
+    C
+        Matrix of molecular orbital coefficients to be localized.
+        Shape is (naos, nmos).
+
+    Returns
+    -------
+    C_loc
+        Localized molecular orbital coefficient matrix of shape (naos, nmos).
+    """
+    naos, nmos = C.shape
+
+    # We can't use scipy's builtin Cholesky-factorization, as it does not
+    # support positive semi-definite matrices.
+    L, piv, _ = pivoted_cholesky(C @ C.T)
+    # Restore original ordering via permutation matrix.
+    P = np.zeros((naos, naos))
+    P[piv, np.arange(naos)] = 1
+    C_loc = P @ L[:, :nmos]
+    return C_loc
+
+
+@cholesky.register
+def _(wf: Wavefunction):
+    """Currently localizes only C_(α,occ) MOs."""
+    Cao, _ = wf.C_occ
+    return cholesky(Cao)
+
+
+def rot_inplace(mat, rad, i, j):
+    """Inplace rotation of matrix columns mat[:, i] and mat[:, j] by 'rad' radians."""
+    cos = np.cos(rad)
+    sin = np.sin(rad)
+    i_rot = cos * mat[:, i] + sin * mat[:, j]
+    j_rot = -sin * mat[:, i] + cos * mat[:, j]
+    mat[:, i] = i_rot
+    mat[:, j] = j_rot
+
+
+@dataclass
+class JacobiSweepResult:
+    is_converged: bool
+    cur_cycle: int
+    C: NDArray[float]
+    P: float
+
+
+def jacobi_sweeps(
+    C: NDArray[float],
+    cost_func: Callable,
+    ab_func: Callable,
+    callback: Optional[Callable] = None,
+    max_cycles: int = 100,
+    dP_thresh: float = 1e-8,
+) -> JacobiSweepResult:
+    """Wrapper for 2x2 Jacobi-sweeps as used in localization/diabatization.
+
+    Parameters
+    ----------
+    C
+        MO coefficient matrix (shape naos x nmos) or rotation matrix (nstates x nstates).
+    cost_func
+        Function to be maximized/minimized.
+    ab_func
+        Function that returns A & B values, used to calculate the angle
+        for the 2x2 rotation.
+    callback
+        Function that is called after the 2x2 rotation took place. It takes three arguments:
+        (gamma, j, k).
+    max_cycles
+        Maximum number of macro cycles.
+    dP_thresh
+        Indicate convergence when change in cost function is equal or below
+        this threshold.
+
+    Returns
+    -------
+    C_loc
+        Localized molecular orbital coefficient matrix of shape (naos, nmos).
+    """
+
+    assert max_cycles > 0
+    C = C.copy()
+    _, nmos = C.shape
+
+    if callback is None:
+        callback = lambda *args: None
+
+    P_prev = 0.0
+
+    logger.info(f"Starting Jacobi sweeps.")
+    for i in range(max_cycles):
+        # Loop over pairs of MO indices and do 2x2 rotations.
+        for j, k in it.combinations(range(nmos), 2):
+            A, B = ab_func(j, k, C)
+
+            if (A**2 + B**2) <= 1e-12:
+                continue
+
+            gamma = np.sign(B) * np.arccos(-A / np.sqrt(A**2 + B**2)) / 4
+            assert -PI_QUART <= gamma <= PI_QUART
+            # 2x2 MO rotations
+            rot_inplace(C, gamma, j, k)
+            callback(gamma, j, k)
+        # Outside of loop over orbital pairs
+
+        # Calculate the target cost function we want to maximize/minimize.
+        P = cost_func(C)
+        dP = P - P_prev
+        logger.info(f"{i:03d}: {P=: >12.8f} {dP=: >12.8f}")
+        if is_converged := (dP <= dP_thresh):
+            logger.info(f"Jacobi sweeps converged in {i+1} cycles.")
+            break
+        P_prev = P
+    # Outside macro cycles
+
+    result = JacobiSweepResult(
+        is_converged=is_converged,
+        cur_cycle=i,
+        C=C,
+        P=P,
+    )
+    return result
+
+
+@singledispatch
+def pipek_mezey(C, S, ao_center_map, **kwargs) -> JacobiSweepResult:
+    """Pipek-Mezey localization using Mulliken population analysis.
+
+    Python adaption of code found in orbloc.f90 of Multiwfn 3.8.
+    For now, only Mulliken population and localization exponent 2 is supported.
+
+    Parameters
+    ----------
+    C
+        Matrix of molecular orbital coefficients to be localized.
+        Shape is (naos, nmos).
+    S
+        Overlap matrix of shape (naos x naos).
+    ao_center_map
+        Mapping between atom indices and AOs, centered at the respective atom.
+
+    Returns
+    -------
+    C_loc
+        Localized molecular orbital coefficient matrix of shape (naos, nmos).
+    """
+    _, nmos = C.shape
+    centers = list(ao_center_map.keys())
+
+    SC = S @ C
+
+    def ab_func(j, k, C):
+        Q = SC * C
+        A = 0.0
+        B = 0.0
+        # Eq. (9) in [1]
+        for center in centers:
+            ao_inds = ao_center_map[center]
+            Qjk = (
+                C[ao_inds, j] * SC[ao_inds, k] + C[ao_inds, k] * SC[ao_inds, j]
+            ).sum() / 2
+            Qjj = Q[ao_inds, j].sum()
+            Qkk = Q[ao_inds, k].sum()
+            A += (Qjk**2) - ((Qjj - Qkk) ** 2) / 4
+            B += Qjk * (Qjj - Qkk)
+        return A, B
+
+    def callback(gamma, j, k):
+        # The MO rotations invalidate the SC matrix product. We update it too.
+        rot_inplace(SC, gamma, j, k)
+
+    def cost_func(C):
+        # We wan't to maximize P and we monitor the progress. Eq. (8) in [1].
+        P = 0.0
+        for l in range(nmos):
+            for center in centers:
+                ao_inds = ao_center_map[center]
+                Q = (C[ao_inds, l][:, None] * C[:, l] * S[ao_inds, :]).sum()
+                P += Q**2
+        return P
+
+    logger.info("Pipek-Mezey localization")
+    return jacobi_sweeps(C, cost_func, ab_func, callback, **kwargs)
+
+
+@pipek_mezey.register
+def _(wf: Wavefunction) -> JacobiSweepResult:
+    """Currently localizes only C_(α,occ) MOs."""
+    S = wf.S
+    Cao, _ = wf.C_occ
+    C_chol_loc = cholesky(Cao)
+    return pipek_mezey(C_chol_loc, S, wf.ao_center_map)
+
+
+def get_fb_contract(moments_ints):
+    def contract(mo_j, mo_k):
+        return np.einsum(
+            "xkl,k,l->x",
+            moments_ints,
+            mo_j,
+            mo_k,
+            optimize=["einsum_path", (0, 1), (0, 1)],
+        )
+
+    return contract
+
+
+def get_fb_ab_func(moments_ints):
+    contract = get_fb_contract(moments_ints)
+
+    def ab_func(j, k, C):
+        mo_j = C[:, j]
+        mo_k = C[:, k]
+        jrj = contract(mo_j, mo_j)
+        jrk = contract(mo_j, mo_k)
+        krk = contract(mo_k, mo_k)
+        A = (jrk**2).sum() - ((jrj - krk) ** 2).sum() / 4  # Eq. (9) in [4]
+        B = ((jrj - krk) * jrk).sum()  # Eq. (10) in [4]
+        return A, B
+
+    return ab_func
+
+
+def get_fb_cost_func(moments_ints):
+    def cost_func(C):
+        vals = np.einsum(
+            "xkl,ki,li->ix",
+            moments_ints,
+            C,
+            C,
+            optimize=["einsum_path", (0, 1), (0, 1)],
+        )
+        val = ((vals[:, None, :] - vals) ** 2).sum()
+        return val
+
+    return cost_func
+
+
+@singledispatch
+def foster_boys(
+    C: NDArray[float], dip_ints: NDArray[float], **kwargs
+) -> JacobiSweepResult:
+    """Foster-Boys localization.
+
+     nMO   nMO
+     ___   ___
+     ╲     ╲                         2
+      ╲     ╲   |                   |
+      ╱     ╱   | <i|R|i> - <j|R|j> |
+     ╱     ╱    |                   |
+     ‾‾‾   ‾‾‾
+    i = 1 j = 1
+
+    or similarily (see eq. (6) in [4] or the appendix of [5])
+
+     nMO
+     ___
+     ╲                2
+      ╲    |         |
+      ╱    | <i|R|i> |
+     ╱     |         |
+     ‾‾‾
+    i = 1
+
+    Parameters
+    ----------
+    dip_ints
+        Dipole moment integral matrices with shape (3, naos, naos).
+    C
+        Matrix of molecular orbital coefficients to be localized.
+        Shape is (naos, nmos).
+    kwargs
+        Additional keyword arguments that are passed jacobi_sweeps.
+
+    Returns
+    -------
+    C_loc
+        Localized molecular orbital coefficient matrix of shape (naos, nmos).
+    """
+
+    ab_func = get_fb_ab_func(dip_ints)
+    cost_func = get_fb_cost_func(dip_ints)
+    logger.info("Foster-Boys localization")
+    return jacobi_sweeps(C, cost_func, ab_func, **kwargs)
+
+
+@foster_boys.register
+def _(wf: Wavefunction) -> JacobiSweepResult:
+    """Currently localizes only C_(α,occ) MOs."""
+    Cao, _ = wf.C_occ
+    dip_ints = wf.dipole_ints()
+    C_chol_loc = cholesky(Cao)
+    return foster_boys(C_chol_loc, dip_ints)
+
+
+def edmiston_ruedenberg_cost_func(L):
+    """Edmiston-Ruedenberg cost function.
+
+    Eq. (2) in [7]."""
+    return np.einsum("Jpp,Jpp->", L, L, optimize="greedy")
+
+
+def edmiston_ruedenberg_grad(L):
+    """Edmiston-Ruedenberg gradient.
+
+    Eq. (5) in [7]."""
+    g1 = np.einsum("Jpq,Jqq->pq", L, L, optimize="greedy")
+    g2 = np.einsum("Jqp,Jpp->pq", L, L, optimize="greedy")
+    return -4 * (g1 - g2)
+
+
+def edmiston_ruedenberg(C, Lao, Sao, diis_cycles=5, max_cycles=100):
+    """Edmiston-Ruedenberg localization using DF integrals and DIIS.
+
+    Parameters
+    ----------
+    C
+        MO-coefficients of shape (naos, nmos); this means MOs should
+        be in columns.
+
+    Lao
+        Density fitting tensor obtained from contracting the 2e3c-matrix
+        with 2c2e**-0.5 in the AO basis. Must have shape (naux, nao, nao).
+        Same as in
+    Sao
+        Overlap matrix of shape (nao, nao) in the AO basis.
+    diis_cycles
+        Integer >= 0. If > 0, DIIS is employed to accelerate convergence.
+    max_cycles
+        Positive integer; maximum number of localization cycles.
+
+    Returns
+    -------
+    Crot
+        Edmiston-Ruedenberg localized orbitals.
+    """
+    nmos = C.shape[1]
+    # We can't keep more error vectors than MOs
+    if nmos < diis_cycles:
+        warnings.warn(
+            f"Can keep at most {nmos} DIIS error vectors, but "
+            f"{diis_cycles} were requested!."
+        )
+    diis_cycles = min(diis_cycles, nmos)
+    if diis_cycles == 1:
+        diis_cycles = 0
+    started_to_store_diis = False
+    specs = {
+        "err_vecs": np.zeros((diis_cycles, nmos**2)),
+        "R_mats": np.zeros((diis_cycles, nmos, nmos)),
+        "D_mats": np.zeros((diis_cycles, nmos, nmos)),
+    }
+    diis = DIIS(specs)
+
+    D = np.eye(nmos)
+    C0 = C.copy()
+    Crot = C.copy()
+    for i in range(max_cycles):
+        # Transform cholesky integrals to MO basis, step (2) in [7].
+        Lpq = np.einsum("Luv,up,vq->Lpq", Lao, Crot, Crot, optimize="greedy")
+        f = edmiston_ruedenberg_cost_func(Lpq)
+        g = edmiston_ruedenberg_grad(Lpq)
+        gnorm = np.linalg.norm(g)
+        grms = rms(g)
+        # Construct transformation, step (3) in [7].
+        # Rji = (Xj Xi Xi Xi)
+        R = np.einsum("Jji,Jii->ji", Lpq, Lpq, optimize="greedy")
+        # DIIS error; lack of symmetry in R.
+        err = (R - R.T).flatten()
+        errrms = rms(err)
+        print(
+            f"Cycle {i:02d} f={f:.8f}, |g|={np.linalg.norm(g):.4f}, rms(g)={grms:>8.4e} "
+            f"rms(err)={errrms:>8.4e}"
+        )
+        if converged := grms <= 1e-4:
+            print("Edmiston-Ruedenberg localization converged!")
+            break
+
+        U = R @ matrix_power(R.T @ R, -0.5)
+        D = D @ U
+        if diis_cycles and (started_to_store_diis or (gnorm <= 2.0e-1)):
+            diis.store(
+                {
+                    "err_vecs": err,
+                    "R_mats": R,
+                    "D_mats": D,
+                }
+            )
+            started_to_store_diis = True
+        if (diis_coeffs := diis.get_coeffs()) is not None:
+            D_diis = np.einsum("i,ijk->jk", diis_coeffs, diis.get("D_mats"))
+            C_diis = C0 @ D_diis
+            # DIIS-2 algorithm in [7]
+            R_diis = np.einsum("i,ijk->jk", diis_coeffs, diis.get("R_mats"))
+            S_diis = C_diis.T @ Sao @ C_diis
+            S_diis_inv = matrix_power(S_diis, -1.0)
+            # Generalized eta step
+            V = S_diis_inv @ R_diis @ matrix_power(R_diis.T @ S_diis_inv @ R_diis, -0.5)
+            # Recalculate D from DIIS results
+            D = D_diis @ V
+        # Update Crot, step (4) in [7]. This either uses the D matrix obtained from DIIS
+        # or the one previously calculated.
+        Crot = C0 @ D
+    # TODO: return JacobiSweepResult?
+    return Crot, f