Merge pull request #81 from BoothGroup/more_diis

DIIS improvements
BoothGroup · Sep 21, 2024 · 439dc88 · 439dc88
2 parents 3508eb5 + f372c34
commit 439dc88
Show file tree

Hide file tree

Showing 8 changed files with 306 additions and 167 deletions.
diff --git a/ebcc/cc/base.py b/ebcc/cc/base.py
@@ -15,6 +15,7 @@
 from ebcc.core.dump import Dump
 from ebcc.core.logging import ANSI
 from ebcc.core.precision import astype, types
+from ebcc.util import _BaseOptions
 
 if TYPE_CHECKING:
     from typing import Any, Callable, Literal, Optional, Union
@@ -23,6 +24,7 @@
     from numpy.typing import NDArray
     from pyscf.scf.hf import SCF
 
+    from ebcc.core.damping import BaseDamping
     from ebcc.core.logging import Logger
     from ebcc.ham.base import BaseElectronBoson, BaseERIs, BaseFock
     from ebcc.opt.base import BaseBruecknerEBCC
@@ -41,14 +43,15 @@
 
 
 @dataclass
-class BaseOptions:
+class BaseOptions(_BaseOptions):
     """Options for EBCC calculations.
 
     Args:
         e_tol: Threshold for convergence in the correlation energy.
         t_tol: Threshold for convergence in the amplitude norm.
         max_iter: Maximum number of iterations.
         diis_space: Number of amplitudes to use in DIIS extrapolation.
+        diis_min_space: Minimum number of amplitudes to use in DIIS extrapolation.
         damping: Damping factor for DIIS extrapolation.
         shift: Shift the boson operators such that the Hamiltonian is normal-ordered with respect
             to a coherent state. This removes the bosonic coupling to the static mean-field
@@ -58,7 +61,8 @@ class BaseOptions:
     e_tol: float = 1e-8
     t_tol: float = 1e-8
     max_iter: int = 200
-    diis_space: int = 12
+    diis_space: int = 9
+    diis_min_space: int = 1
     damping: float = 0.0
     shift: bool = True
 
@@ -68,6 +72,7 @@ class BaseEBCC(ABC):
 
     # Types
     Options: type[BaseOptions] = BaseOptions
+    Damping: type[BaseDamping] = DIIS
     ERIs: type[BaseERIs]
     Fock: type[BaseFock]
     CDERIs: type[BaseERIs]
@@ -195,6 +200,7 @@ def __init__(
         self.log.info(f" > t_tol:  {ANSI.y}{self.options.t_tol}{ANSI.R}")
         self.log.info(f" > max_iter:  {ANSI.y}{self.options.max_iter}{ANSI.R}")
         self.log.info(f" > diis_space:  {ANSI.y}{self.options.diis_space}{ANSI.R}")
+        self.log.info(f" > diis_min_space:  {ANSI.y}{self.options.diis_min_space}{ANSI.R}")
         self.log.info(f" > damping:  {ANSI.y}{self.options.damping}{ANSI.R}")
         self.log.debug("")
         self.log.info(f"{ANSI.B}Ansatz{ANSI.R}: {ANSI.m}{self.ansatz}{ANSI.R}")
@@ -244,18 +250,16 @@ def kernel(self, eris: Optional[ERIsInputType] = None) -> float:
         self.log.info(f"{0:4d} {e_cc:16.10f} {e_cc + self.mf.e_tot:18.10f}")
 
         if not self.ansatz.is_one_shot:
-            # Set up DIIS:
-            diis = DIIS()
-            diis.space = self.options.diis_space
-            diis.damping = self.options.damping
+            # Set up damping:
+            damping = self.Damping(options=self.options)
 
             converged = False
             for niter in range(1, self.options.max_iter + 1):
                 # Update the amplitudes, extrapolate with DIIS and calculate change:
                 amplitudes_prev = amplitudes
                 amplitudes = self.update_amps(amplitudes=amplitudes, eris=eris)
                 vector = self.amplitudes_to_vector(amplitudes)
-                vector = diis.update(vector)
+                vector = damping(vector)
                 amplitudes = self.vector_to_amplitudes(vector)
                 dt = np.linalg.norm(
                     np.abs(vector - self.amplitudes_to_vector(amplitudes_prev)), ord=np.inf
@@ -345,10 +349,8 @@ def solve_lambda(
         self.log.debug("")
         self.log.info(f"{ANSI.B}{'Iter':>4s} {'Δ(Ampl.)':>13s}{ANSI.R}")
 
-        # Set up DIIS:
-        diis = DIIS()
-        diis.space = self.options.diis_space
-        diis.damping = self.options.damping
+        # Set up damping:
+        damping = self.Damping(options=self.options)
 
         converged = False
         for niter in range(1, self.options.max_iter + 1):
@@ -361,7 +363,7 @@ def solve_lambda(
                 eris=eris,
             )
             vector = self.lambdas_to_vector(lambdas)
-            vector = diis.update(vector)
+            vector = damping(vector)
             lambdas = self.vector_to_lambdas(vector)
             dl = np.linalg.norm(np.abs(vector - self.lambdas_to_vector(lambdas_prev)), ord=np.inf)