Merge pull request #121 from BoothGroup/CCSDt_interface

CCSDt' interface (and general EBCC interface updates)

examples/ewf/molecules/26-ebcc-solvers.py

@@ -0,0 +1,53 @@
+import pyscf
+import pyscf.gto
+import pyscf.scf
+import pyscf.mcscf
+
+import vayesta
+import vayesta.ewf
+
+mol = pyscf.gto.Mole()
+mol.atom = ['N 0 0 0', 'N 0 0 2']
+mol.basis = 'aug-cc-pvdz'
+mol.output = 'pyscf.out'
+mol.build()
+
+# Hartree-Fock
+mf = pyscf.scf.RHF(mol)
+mf.kernel()
+
+# Reference CASCI
+casci = pyscf.mcscf.CASCI(mf, 8, 10)
+casci.kernel()
+
+# Reference CASSCF
+casscf = pyscf.mcscf.CASSCF(mf, 8, 10)
+casscf.kernel()
+
+def get_emb_result(ansatz, bathtype='full'):
+    # Uses fastest available solver for given ansatz; PySCF if available, otherwise ebcc.
+    emb = vayesta.ewf.EWF(mf, solver=ansatz, bath_options=dict(bathtype=bathtype),
+                          solver_options=dict(solve_lambda=False))
+    # Both these alternative specifications will always use an ebcc solver.
+    # Note that the capitalization of the solver name other than the ansatz is arbitrary.
+    #emb = vayesta.ewf.EWF(mf, solver=f'EB{ansatz}', bath_options=dict(bathtype=bathtype),
+    #                      solver_options=dict(solve_lambda=False))
+    #emb = vayesta.ewf.EWF(mf, solver='ebcc', bath_options=dict(bathtype=bathtype),
+    #                      solver_options=dict(solve_lambda=False, ansatz=ansatz))
+
+    with emb.iao_fragmentation() as f:
+        with f.rotational_symmetry(2, "y", center=(0, 0, 1)):
+            f.add_atomic_fragment(0)
+    emb.kernel()
+    return emb.e_tot
+
+e_ccsd = get_emb_result('CCSD', 'full')
+e_ccsdt = get_emb_result('CCSDT', 'dmet')
+e_ccsdtprime = get_emb_result("CCSDt'", 'full')
+
+print("E(HF)=                                           %+16.8f Ha" % mf.e_tot)
+print("E(CASCI)=                                        %+16.8f Ha" % casci.e_tot)
+print("E(CASSCF)=                                       %+16.8f Ha" % casscf.e_tot)
+print("E(CCSD, complete)=                               %+16.8f Ha" % e_ccsd)
+print("E(emb. CCSDT, DMET CAS)=                         %+16.8f Ha" % e_ccsdt)
+print("E(emb. CCSDt', complete+DMET active space)=      %+16.8f Ha" % e_ccsdtprime)

vayesta/core/qemb/qemb.py

@@ -103,7 +103,7 @@ class Options(OptionsBase):
             # EBFCI/EBCCSD
             max_boson_occ=2,
             # EBCC
-            ansatz=None, fermion_wf=False,
+            ansatz=None, store_as_ccsd=None, fermion_wf=False,
             # Dump
             dumpfile='clusters.h5',
             # MP2

vayesta/core/types/wf/ccsd.py

@@ -5,12 +5,13 @@
 # import pyscf
 # import pyscf.cc
 from vayesta.core import spinalg
-from vayesta.core.util import NotCalculatedError, Object, callif, einsum
+from vayesta.core.util import NotCalculatedError, Object, callif, einsum, dot
 from vayesta.core.types import wf as wf_types
 from vayesta.core.types.orbitals import SpatialOrbitals
 from vayesta.core.types.wf.project import (project_c1, project_c2, project_uc1, project_uc2, symmetrize_c2,
-                                           symmetrize_uc2)
+                                           symmetrize_uc2, transform_c1, transform_c2, transform_uc1, transform_uc2)
 from vayesta.core.helper import pack_arrays, unpack_arrays
+from scipy.linalg import block_diag
 
 
 def CCSD_WaveFunction(mo, t1, t2, **kwargs):
@@ -170,6 +171,42 @@ def as_ccsd(self):
     def as_fci(self):
         raise NotImplementedError
 
+    def rotate(self, t, inplace=False):
+        """Rotate wavefunction representation to another basis.
+        Only rotations which don't mix occupied and virtual orbitals are supported.
+        Assumes rotated orbitals have same occupancy ordering as originals.
+        """
+        o = self.mo.occ > 0
+        v = self.mo.occ == 0
+
+        to = t[np.ix_(o, o)]
+        tv = t[np.ix_(v, v)]
+        tov = t[np.ix_(o, v)]
+        tvo = t[np.ix_(v, o)]
+        if abs(tov).max() > 1e-12 or abs(tvo).max() > 1e-12:
+            raise ValueError("Provided rotation mixes occupied and virtual orbitals.")
+        return self.rotate_ov(to, tv, inplace=inplace)
+
+    def rotate_ov(self, to, tv, inplace=False):
+        """Rotate wavefunction representation to another basis.
+        Only rotations which don't mix occupied and virtual orbitals are supported.
+
+        Parameters
+        ----------
+        to : new occupied orbital coefficients in terms of current ones.
+        tv : new virtual orbital coefficients in terms of current ones.
+        inplace : Whether to transform in-place or return a new object.
+        """
+        wf = self if inplace else self.copy()
+        wf.mo.basis_transform(lambda c: dot(c, block_diag(to, tv)), inplace=True)
+        wf.t1 = transform_c1(wf.t1, to, tv)
+        wf.t2 = transform_c2(wf.t2, to, tv)
+        if wf.l1 is not None:
+            wf.l1 = transform_c1(wf.l1, to, tv)
+        if wf.l2 is not None:
+            wf.l2 = transform_c2(wf.l2, to, tv)
+        return wf
+
 
 class UCCSD_WaveFunction(RCCSD_WaveFunction):
 
@@ -296,6 +333,56 @@ def multiply(self, factor):
         if self.l2 is not None:
             self.l2 = spinalg.multiply(self.l2, len(self.l2)*[factor])
 
+    def rotate(self, t, inplace=False):
+        """Rotate wavefunction representation to another basis.
+        Only rotations which don't mix occupied and virtual orbitals are supported.
+        Assumes rotated orbitals have same occupancy ordering as originals.
+        """
+        # Allow support for same rotation for alpha and beta.
+        if isinstance(t, np.ndarray) and t.ndim == 2:
+            t = (t, t)
+        def get_components(tsp, occ):
+            o = occ > 0
+            v = occ == 0
+            tspo = tsp[np.ix_(o, o)]
+            tspv = tsp[np.ix_(v, v)]
+            tspov = tsp[np.ix_(o, v)]
+            tspvo = tsp[np.ix_(v, o)]
+            if abs(tspov).max() > 1e-12 or abs(tspvo).max() > 1e-12:
+                raise ValueError("Provided rotation mixes occupied and virtual orbitals.")
+            return tspo, tspv
+
+        toa, tva = get_components(t[0], self.mo.alpha.occ)
+        tob, tvb = get_components(t[1], self.mo.beta.occ)
+
+        return self.rotate_ov((toa, tob), (tva, tvb), inplace=inplace)
+
+    def rotate_ov(self, to, tv, inplace=False):
+        """Rotate wavefunction representation to another basis.
+        Only rotations which don't mix occupied and virtual orbitals are supported.
+
+        Parameters
+        ----------
+        to : new occupied orbital coefficients in terms of current ones.
+        tv : new virtual orbital coefficients in terms of current ones.
+        inplace : Whether to transform in-place or return a new object.
+        """
+        wf = self if inplace else self.copy()
+        if isinstance(to, np.ndarray) and len(to) == 2:
+            assert(isinstance(tv, np.ndarray) and len(tv) == 2)
+            trafo = lambda c: dot(c, block_diag(to, tv))
+        else:
+            trafo = [lambda c: dot(c, x) for x in (block_diag(to[0], tv[0]), block_diag(to[1], tv[1]))]
+
+        wf.mo.basis_transform(trafo, inplace=True)
+        wf.t1 = transform_uc1(wf.t1, to, tv)
+        wf.t2 = transform_uc2(wf.t2, to, tv)
+        if wf.l1 is not None:
+            wf.l1 = transform_uc1(wf.l1, to, tv)
+        if wf.l2 is not None:
+            wf.l2 = transform_uc2(wf.l2, to, tv)
+        return wf
+
     #def pack(self, dtype=float):
     #    """Pack into a single array of data type `dtype`.
 

vayesta/core/types/wf/project.py

@@ -1,6 +1,7 @@
 """Utility functions for projection of wave functions."""
 
 import numpy as np
+from vayesta.core.util import einsum, dot
 
 
 def project_c1(c1, p):
@@ -45,3 +46,27 @@ def symmetrize_uc2(c2, inplace=True):
     if len(c2) == 4:
         c2ab = (c2ab + c2[2].transpose(1,0,3,2))/2
     return (c2aa, c2ab, c2bb)
+
+def transform_c1(c1, to, tv):
+    if c1 is None: return None
+    return dot(to.T, c1, tv)
+
+def transform_c2(c2, to, tv, to2=None, tv2=None):
+    if c2 is None: return None
+    if to2 is None: to2 = to
+    if tv2 is None: tv2 = tv
+    # Use einsum for now- tensordot would be faster but less readable.
+    return einsum("ijab,iI,jJ,aA,bB->IJAB", c2, to, to2, tv, tv2)
+
+def transform_uc1(c1, to, tv):
+    if c1 is None: return None
+    return (transform_c1(c1[0], to[0], tv[0]),
+            transform_c1(c1[1], to[1], tv[1]))
+
+def transform_uc2(c2, to, tv):
+    if c2 is None: return None
+    c2ba = (c2[2] if len(c2) == 4 else c2[1].transpose(1,0,3,2))
+    return (transform_c2(c2[0], to[0], tv[0]),
+            transform_c2(c2[1], to[0], tv[0], to[1], tv[1]),
+            transform_c2(c2ba, to[1], tv[1], to[0], tv[0]),
+            transform_c2(c2[-1], to[1], tv[1]))

vayesta/core/types/wf/wf.py

@@ -65,6 +65,12 @@ def make_rdm1(self, *args, **kwargs):
     def make_rdm2(self, *args, **kwargs):
         raise AbstractMethodError
 
+    def rotate_ov(self, *args, **kwargs):
+        raise AbstractMethodError
+
+    def rotate(self, *args, **kwargs):
+        raise AbstractMethodError
+
     @staticmethod
     def from_pyscf(obj, **kwargs):
         # HF

vayesta/ewf/ewf.py

@@ -229,7 +229,7 @@ def t1_diagnostic(self, warntol=0.02):
     # Defaults
 
     def make_rdm1(self, *args, **kwargs):
-        if self.solver.lower() == 'ccsd':
+        if "cc" in self.solver.lower():
             return self._make_rdm1_ccsd_global_wf(*args, **kwargs)
         if self.solver.lower() == 'mp2':
             return self._make_rdm1_mp2_global_wf(*args, **kwargs)

vayesta/ewf/fragment.py

@@ -100,8 +100,13 @@ def __init__(self, *args, **kwargs):
         # For self-consistent mode
         self.solver_results = None
 
+    def _reset(self, *args, **kwargs):
+        super()._reset(*args, **kwargs)
+        # Need to unset these so can be regenerated each iteration.
+        self.opts.c_cas_occ = self.opts.c_cas_vir = None
+
     def set_cas(self, iaos=None, c_occ=None, c_vir=None, minao='auto', dmet_threshold=None):
-        """Set complete active space for tailored CCSD"""
+        """Set complete active space for tailored CCSD and active-space CC methods."""
         if dmet_threshold is None:
             dmet_threshold = 2*self.opts.bath_options['dmet_threshold']
         if iaos is not None:
@@ -282,8 +287,9 @@ def get_solver_options(self, solver):
             self.log.debugv("Passing fragment option %s to solver.", attr)
             solver_opts[attr] = getattr(self.opts, attr)
 
-        if solver.upper() == 'TCCSD':
-            solver_opts['tcc'] = True
+        has_actspace = ((solver == "TCCSD") or ("CCSDt'" in solver) or
+                        (solver.upper() == "EBCC" and self.opts.solver_options['ansatz'] == "CCSDt'"))
+        if has_actspace:
             # Set CAS orbitals
             if self.opts.c_cas_occ is None:
                 self.log.warning("Occupied CAS orbitals not set. Setting to occupied DMET cluster orbitals.")
@@ -293,7 +299,8 @@ def get_solver_options(self, solver):
                 self.opts.c_cas_vir = self._dmet_bath.c_cluster_vir
             solver_opts['c_cas_occ'] = self.opts.c_cas_occ
             solver_opts['c_cas_vir'] = self.opts.c_cas_vir
-            solver_opts['tcc_fci_opts'] = self.opts.tcc_fci_opts
+            if solver == "TCCSD":
+                solver_opts['tcc_fci_opts'] = self.opts.tcc_fci_opts
         elif solver.upper() == 'DUMP':
             solver_opts['filename'] = self.opts.solver_options['dumpfile']
         solver_opts['external_corrections'] = self.flags.external_corrections

vayesta/ewf/ufragment.py

@@ -11,8 +11,14 @@
 
 class Fragment(RFragment, BaseFragment):
 
-    def set_cas(self, *args, **kwargs):
-        raise NotImplementedError()
+    def set_cas(self, iaos=None, c_occ=None, c_vir=None, minao='auto', dmet_threshold=None):
+        """Set complete active space for tailored CCSD and active-space CC methods."""
+        if iaos is not None:
+                raise NotImplementedError("Unrestricted IAO-based CAS not implemented yet.")
+
+        self.opts.c_cas_occ = c_occ
+        self.opts.c_cas_vir = c_vir
+        return c_occ, c_vir
 
     def get_fragment_energy(self, c1, c2, hamil=None, fock=None, axis1='fragment', c2ba_order='ba'):
         """Calculate fragment correlation energy contribution from projected C1, C2.

vayesta/solver/__init__.py

@@ -2,14 +2,19 @@
 from vayesta.solver.cisd import RCISD_Solver, UCISD_Solver
 from vayesta.solver.coupled_ccsd import coupledRCCSD_Solver
 from vayesta.solver.dump import DumpSolver
-from vayesta.solver.ebcc import REBCC_Solver, UEBCC_Solver, EB_REBCC_Solver, EB_UEBCC_Solver
 from vayesta.solver.ebfci import EB_EBFCI_Solver, EB_UEBFCI_Solver
 from vayesta.solver.ext_ccsd import extRCCSD_Solver, extUCCSD_Solver
 from vayesta.solver.fci import FCI_Solver, UFCI_Solver
 from vayesta.solver.hamiltonian import is_ham, is_uhf_ham, is_eb_ham, ClusterHamiltonian
 from vayesta.solver.mp2 import RMP2_Solver, UMP2_Solver
 from vayesta.solver.tccsd import TRCCSD_Solver
 
+try:
+    from vayesta.solver.ebcc import REBCC_Solver, UEBCC_Solver, EB_REBCC_Solver, EB_UEBCC_Solver
+except ImportError:
+    _has_ebcc = False
+else:
+    _has_ebcc = True
 
 def get_solver_class(ham, solver):
     assert (is_ham(ham))
@@ -36,7 +41,6 @@ def _get_solver_class(is_uhf, is_eb, solver, log):
 
 
 def _get_solver_class_internal(is_uhf, is_eb, solver):
-    solver = solver.upper()
     # First check if we have a CC approach as implemented in pyscf.
     if solver == "CCSD" and not is_eb:
         # Use pyscf solvers.
@@ -48,21 +52,25 @@ def _get_solver_class_internal(is_uhf, is_eb, solver):
         if is_uhf or is_eb:
             raise ValueError("TCCSD is not implemented for unrestricted or electron-boson calculations!")
         return TRCCSD_Solver
-    if solver == "EXTCCSD":
+    if solver == "extCCSD":
         if is_eb:
             raise ValueError("extCCSD is not implemented for electron-boson calculations!")
         if is_uhf:
             return extUCCSD_Solver
         return extRCCSD_Solver
-    if solver == "COUPLEDCCSD":
+    if solver == "coupledCCSD":
         if is_eb:
             raise ValueError("coupledCCSD is not implemented for electron-boson calculations!")
         if is_uhf:
             raise ValueError("coupledCCSD is not implemented for unrestricted calculations!")
         return coupledRCCSD_Solver
 
     # Now consider general CC ansatzes; these are solved via EBCC.
-    if "CC" in solver:
+    # Note that we support all capitalisations of `ebcc`, but need `CC` to be capitalised when also using this to
+    # specify an ansatz.
+    if "CC" in solver.upper():
+        if not _has_ebcc:
+            raise ImportError(f"{solver} solver is only accessible via ebcc. Please install ebcc.")
         if is_uhf:
             if is_eb:
                 solverclass = EB_UEBCC_Solver
@@ -73,23 +81,24 @@ def _get_solver_class_internal(is_uhf, is_eb, solver):
                 solverclass = EB_REBCC_Solver
             else:
                 solverclass = REBCC_Solver
-        if solver == "EBCC":
+        if solver.upper() == "EBCC":
             # Default to `opts.ansatz`.
             return solverclass
-        if solver[:2] == "EB":
+        if solver[:2].upper() == "EB":
             solver = solver[2:]
         if solver == "CCSD" and is_eb:
             # Need to specify CC level for coupled electron-boson model; throw an error rather than assume.
             raise ValueError(
                 "Please specify a coupled electron-boson CC ansatz as a solver, for example CCSD-S-1-1,"
                 "rather than CCSD")
-
+        # This is just a wrapper to allow us to use the solver option as the ansatz kwarg in this case.
         def get_right_CC(*args, **kwargs):
-
-            if kwargs.get("ansatz", None) is not None:
-                raise ValueError(
-                    "Desired CC ansatz specified differently in solver and solver_options.ansatz."
-                    "Please use only specify via one approach, or ensure they agree.")
+            setansatz = kwargs.get("ansatz", None)
+            if setansatz is not None:
+                if setansatz != solver:
+                    raise ValueError(
+                        "Desired CC ansatz specified differently in solver and solver_options.ansatz."
+                        "Please use only specify via one approach, or ensure they agree.")
             kwargs["ansatz"] = solver
             return solverclass(*args, **kwargs)