Revised DIIS code to be somewhat simpler and more general.

.github/workflows/CI.yaml

@@ -21,7 +21,8 @@ jobs:
     strategy:
       fail-fast: true
       matrix:
-        os: [macOS-latest, ubuntu-latest]
+        #os: [macOS-latest, ubuntu-latest]
+        os: [macOS-latest]
         python-version: ['3.11']
 
     steps:
@@ -60,11 +61,12 @@ jobs:
 
     - name: Install PyCC and Deps
       run: |
-        python -m pip install .
+        pip install -e .
         pip install pytest
 
     - name: Test and generate coverage report
       run: |
+        export KMP_DUPLICATE_LIB_OK=TRUE
         pip install coverage
         coverage run -m pytest
         coverage xml

pycc/cclambda.py

@@ -9,7 +9,7 @@
 import numpy as np
 import time
 from opt_einsum import contract
-from .utils import helper_diis
+from .utils import DIIS
 import torch
 from .cctriples import t3c_ijk, l3_ijk, l3_ijk_alt
 
@@ -95,6 +95,9 @@ def solve_lambda(self, e_conv=1e-7, r_conv=1e-7, maxiter=100, max_diis=8, start_
         ERI = self.ccwfn.H.ERI
         L = self.ccwfn.H.L
 
+        t1len = no*nv
+        t2len = no*no*nv*nv
+
         Hov = self.hbar.Hov
         Hvv = self.hbar.Hvv
         Hoo = self.hbar.Hoo
@@ -111,7 +114,8 @@ def solve_lambda(self, e_conv=1e-7, r_conv=1e-7, maxiter=100, max_diis=8, start_
 
         print("\nLCC Iter %3d: LCC PseudoE = %.15f  dE = % .5E" % (0, lecc, -lecc))
 
-        diis = helper_diis(l1, l2, max_diis, self.ccwfn.precision)
+        Lambda = np.hstack((l1.flatten(), l2.flatten()))
+        diis = DIIS(Lambda, max_diis, self.ccwfn.precision)
 
         contract = self.contract
 
@@ -268,9 +272,13 @@ def solve_lambda(self, e_conv=1e-7, r_conv=1e-7, maxiter=100, max_diis=8, start_
                     print("\nLambda-CC has converged in %.3f seconds.\n" % (time.time() - lambda_tstart))
                     return lecc
 
-            diis.add_error_vector(self.l1, self.l2)
+            Lambda = np.hstack((l1.flatten(), l2.flatten()))
+            e = np.hstack(((r1/Dia).flatten(), (r2/Dijab).flatten()))
+            diis.add_error_vector(Lambda, e)
             if niter >= start_diis:
-                self.l1, self.l2 = diis.extrapolate(self.l1, self.l2)
+                Lambda = diis.extrapolate(np.hstack((self.l1.flatten(), self.l2.flatten())))
+                self.l1 = np.reshape(Lambda[:t1len], (no, nv))
+                self.l2 = np.reshape(Lambda[t1len:], (no, no, nv, nv))
 
         if isinstance(r1, torch.Tensor):
             del Goo, Gvv, Hoo, Hvv, Hov, Hovvo, Hovov, Hvvvo, Hovoo, Hvovv, Hooov

pycc/ccresponse.py

@@ -7,7 +7,7 @@
 
 import numpy as np
 import time
-from .utils import helper_diis
+from .utils import DIIS
 
 class ccresponse(object):
     """
@@ -304,12 +304,18 @@ def solve_right(self, pertbar, omega, e_conv=1e-13, r_conv=1e-13, maxiter=200, m
         pseudo = self.pseudoresponse(pertbar, X1, X2)
         print(f"Iter {0:3d}: CC Pseudoresponse = {pseudo.real:.15f} dP = {pseudo.real:.5E}")
 
-        diis = helper_diis(X1, X2, max_diis)
+        X = np.hstack((X1.flatten(), X2.flatten()))
+        diis = DIIS(X, max_diis)
         contract = self.ccwfn.contract
 
         self.X1 = X1
         self.X2 = X2
 
+        no = self.ccwfn.no
+        nv = self.ccwfn.nv
+        t1len = no*nv
+        t2len = no*no*nv*nv
+
         for niter in range(1, maxiter+1):
             pseudo_last = pseudo
 
@@ -334,9 +340,13 @@ def solve_right(self, pertbar, omega, e_conv=1e-13, r_conv=1e-13, maxiter=200, m
                 print("\nPerturbed wave function converged in %.3f seconds.\n" % (time.time() - solver_start))
                 return self.X1, self.X2, pseudo
 
-            diis.add_error_vector(self.X1, self.X2)
+            X = np.hstack((self.X1.flatten(), self.X2.flatten()))
+            e = np.hstack(((r1/(Dia + omega)).flatten(), (r2/(Dijab + omega)).flatten()))
+            diis.add_error_vector(X, e)
             if niter >= start_diis:
-                self.X1, self.X2 = diis.extrapolate(self.X1, self.X2)
+                X = diis.extrapolate(X)
+                self.X1 = np.reshape(X[:t1len], (no, nv))
+                self.X2 = np.reshape(X[t1len:], (no, no, nv, nv))
 
     def r_X1(self, pertbar, omega):
         contract = self.contract

pycc/ccwfn.py

@@ -10,7 +10,7 @@
 import time
 import numpy as np
 import torch
-from .utils import helper_diis, cc_contract
+from .utils import DIIS, cc_contract
 from .hamiltonian import Hamiltonian
 from .local import Local
 from .cctriples import t_tjl, t3c_ijk, t3d_ijk, t3c_abc, t3d_abc
@@ -257,12 +257,17 @@ def solve_cc(self, e_conv=1e-7, r_conv=1e-7, maxiter=100, max_diis=8, start_diis
         Dia = self.Dia
         Dijab = self.Dijab
 
+        no = self.no
+        nv = self.nv
+        t1len = no*nv
+        t2len = no*no*nv*nv
+
         contract = self.contract
 
         ecc = self.cc_energy(o, v, F, L, self.t1, self.t2)
         print("CC Iter %3d: CC Ecorr = %.15f  dE = % .5E  MP2" % (0, ecc, -ecc))
 
-        diis = helper_diis(self.t1, self.t2, max_diis, self.precision)
+        diis = DIIS(np.hstack((self.t1.flatten(), self.t2.flatten())), max_diis, self.precision)
 
         for niter in range(1, maxiter+1):
 
@@ -321,9 +326,13 @@ def solve_cc(self, e_conv=1e-7, r_conv=1e-7, maxiter=100, max_diis=8, start_diis
                     print("E(TOT)  = %20.15f" % (ecc + self.eref))
                     return ecc
 
-            diis.add_error_vector(self.t1, self.t2)
+            T = np.hstack((self.t1.flatten(), self.t2.flatten()))
+            e = np.hstack(((r1/Dia).flatten(), (r2/Dijab).flatten()))
+            diis.add_error_vector(T, e)
             if niter >= start_diis:
-                self.t1, self.t2 = diis.extrapolate(self.t1, self.t2)
+                T = diis.extrapolate(T)
+                self.t1 = np.reshape(T[:t1len], (no, nv))
+                self.t2 = np.reshape(T[t1len:], (no, no, nv, nv))
 
     def residuals(self, F, t1, t2):
         """

pycc/utils.py

@@ -3,67 +3,47 @@
 import opt_einsum
 
 
-class helper_diis(object):
-    def __init__(self, t1, t2, max_diis, precision='DP'):
-        if isinstance(t1, torch.Tensor):
+class DIIS(object):
+    def __init__(self, T, max_diis, precision='DP'):
+        if isinstance(T, torch.Tensor):
             self.device0 = torch.device('cpu')
             self.device1 = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
-            self.oldt1 = t1.clone()
-            self.oldt2 = t2.clone()
-            self.diis_vals_t1 = [t1.clone()]
-            self.diis_vals_t2 = [t2.clone()]
+            self.diis_T = [T.clone()]
         else:
-            self.oldt1 = t1.copy()
-            self.oldt2 = t2.copy()
-            self.diis_vals_t1 = [t1.copy()]
-            self.diis_vals_t2 = [t2.copy()]
+            self.diis_T = [T.copy()]
 
         self.diis_errors = []
         self.diis_size = 0
         self.max_diis = max_diis
         self.precision = precision
 
-    def add_error_vector(self, t1, t2):
-        if isinstance(t1, torch.Tensor):
-            # Add DIIS vectors
-            self.diis_vals_t1.append(t1.clone())
-            self.diis_vals_t2.append(t2.clone())
-            # Add new error vectors
-            error_t1 = (self.diis_vals_t1[-1] - self.oldt1).ravel()
-            error_t2 = (self.diis_vals_t2[-1] - self.oldt2).ravel()
-            self.diis_errors.append(torch.cat((error_t1, error_t2)))
-            self.oldt1 = t1.clone()
-            self.oldt2 = t2.clone()
+    def add_error_vector(self, T, e):
+        if isinstance(T, torch.Tensor):
+            self.diis_T.append(T.clone())
+            self.diis_errors.append(e.ravel().clone())
         else:
-            # Add DIIS vectors
-            self.diis_vals_t1.append(t1.copy())
-            self.diis_vals_t2.append(t2.copy())
-            # Add new error vectors
-            error_t1 = (self.diis_vals_t1[-1] - self.oldt1).ravel()
-            error_t2 = (self.diis_vals_t2[-1] - self.oldt2).ravel()
-            self.diis_errors.append(np.concatenate((error_t1, error_t2)))
-            self.oldt1 = t1.copy()
-            self.oldt2 = t2.copy()
-
-    def extrapolate(self, t1, t2):
+            self.diis_T.append(T.copy())
+            self.diis_errors.append(e.ravel())
+
+    def extrapolate(self, T):
 
         if (self.max_diis == 0):
-            return t1, t2
+            return T
 
         # Limit size of DIIS vector
         if (len(self.diis_errors) > self.max_diis):
-            del self.diis_vals_t1[0]
-            del self.diis_vals_t2[0]
+            del self.diis_T[0]
             del self.diis_errors[0]
 
         self.diis_size = len(self.diis_errors)
 
-        if isinstance(t1, torch.Tensor):
+        if isinstance(T, torch.Tensor):
             # Build error matrix B
             if self.precision == 'DP':
-                B = torch.ones((self.diis_size + 1, self.diis_size + 1), dtype=torch.float64, device=self.device1) * -1
+                B = -1 * torch.ones((self.diis_size + 1, self.diis_size + 1), dtype=torch.float64, device=self.device1)
             elif self.precision == 'SP':
-                B = torch.ones((self.diis_size + 1, self.diis_size + 1), dtype=torch.float32, device=self.device1) * -1
+                B = -1 * torch.ones((self.diis_size + 1, self.diis_size + 1), dtype=torch.float32, device=self.device1)
+
             B[-1, -1] = 0
 
             for n1, e1 in enumerate(self.diis_errors):
@@ -87,15 +67,9 @@ def extrapolate(self, t1, t2):
             ci = torch.linalg.solve(B, resid)
 
             # Calculate new amplitudes
-            t1 = torch.zeros_like(self.oldt1)
-            t2 = torch.zeros_like(self.oldt2)
+            T *= 0
             for num in range(self.diis_size):
-                t1 += torch.real(ci[num] * self.diis_vals_t1[num + 1])
-                t2 += torch.real(ci[num] * self.diis_vals_t2[num + 1])
-
-            # Save extrapolated amplitudes to old_t amplitudes
-            self.oldt1 = t1.clone()
-            self.oldt2 = t2.clone()
+                T += torch.real(ci[num] * self.diis_T[num + 1])
 
         else:
             # Build error matrix B
@@ -126,17 +100,11 @@ def extrapolate(self, t1, t2):
             ci = np.linalg.solve(B, resid)
 
             # Calculate new amplitudes
-            t1 = np.zeros_like(self.oldt1)
-            t2 = np.zeros_like(self.oldt2)
+            T *= 0
             for num in range(self.diis_size):
-                t1 += ci[num] * self.diis_vals_t1[num + 1]
-                t2 += ci[num] * self.diis_vals_t2[num + 1]
-
-            # Save extrapolated amplitudes to old_t amplitudes
-            self.oldt1 = t1.copy()
-            self.oldt2 = t2.copy()
+                T += ci[num] * self.diis_T[num + 1]
 
-        return t1, t2
+        return T
 
 class cc_contract(object):
     """