Update existing EDMET code to use new wavefunction objects.

vayesta/core/qemb/qemb.py

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

vayesta/core/types/ebwf/ebcc.py

@@ -36,7 +36,7 @@ class REBCC_WaveFunction(EBWavefunction, RCCSD_WaveFunction):
     _spin_type = "R"
     _driver = ebcc.REBCC
 
-    def __init__(self, mo, ansatz, amplitudes, lambdas=None, mbos=None, projector=None):
+    def __init__(self, mo, ansatz, amplitudes, lambdas=None, mbos=None, projector=None, xi=None):
         super().__init__(mo, mbos, projector)
         self.amplitudes = amplitudes
         if lambdas is not None and len(lambdas) == 0:
@@ -47,6 +47,11 @@ def __init__(self, mo, ansatz, amplitudes, lambdas=None, mbos=None, projector=No
         else:
             self.ansatz = ebcc.Ansatz.from_string(ansatz)
         self._eqns = self.ansatz._get_eqns(self._spin_type)
+        self.xi = xi
+
+    @property
+    def options(self):
+        return ebcc.util.Namespace(shift=self.xi is not None)
 
     @property
     def name(self):
@@ -115,22 +120,32 @@ def _pack_codegen_kwargs(self, *extra_kwargs, eris=False):
                 kwargs.update(kw)
         return kwargs
 
-    def make_rdm1(self,  t_as_lambda=False, with_mf=True, ao_basis=False):
+    def make_rdm1(self,  t_as_lambda=False, with_mf=True, ao_basis=False, hermitise=True, **kwargs):
         assert(not t_as_lambda and with_mf and not ao_basis)
-        return self._driver.make_rdm1_f(self, eris=False, amplitudes=self.amplitudes, lambdas=self.lambdas, hermitise=True)
+        return self._driver.make_rdm1_f(self, eris=False, amplitudes=self.amplitudes, lambdas=self.lambdas,
+                                        hermitise=True, **kwargs)
 
-    def make_rdm2(self, t_as_lambda=False, with_dm1=True, ao_basis=False, approx_cumulant=True):
+    def make_rdm2(self, t_as_lambda=False, with_dm1=True, ao_basis=False, approx_cumulant=False, hermitise=True,
+                  **kwargs):
         assert(not t_as_lambda and with_dm1 and not ao_basis and not approx_cumulant)
-        return self._driver.make_rdm2_f(self, eris=False, amplitudes=self.amplitudes, lambdas=self.lambdas, hermitise=True)
+        return self._driver.make_rdm2_f(self, eris=False, amplitudes=self.amplitudes, lambdas=self.lambdas,
+                                        hermitise=hermitise, **kwargs)
+
+    def make_rdm1_b(self, hermitise=True, **kwargs):
+        return self._driver.make_rdm1_b(self, eris=False, amplitudes=self.amplitudes, lambdas=self.lambdas,
+                                        hermitise=hermitise, **kwargs)
 
-    def make_rdm1_b(self, *args, **kwargs):
-        return self._driver.make_rdm1_b(self, eris=False, amplitudes=self.amplitudes, lambdas=self.lambdas, hermitise=True)
+    def make_sing_b_dm(self, hermitise=True, **kwargs):
+        return self._driver.make_sing_b_dm(self, eris=False, amplitudes=self.amplitudes, lambdas=self.lambdas,
+                                           hermitise=hermitise, **kwargs)
 
-    def make_sing_b_dm(self, *args, **kwargs):
-        return self._driver.make_sing_b_dm(self, eris=False, amplitudes=self.amplitudes, lambdas=self.lambdas, hermitise=True)
+    def make_rdm_eb(self, hermitise=True, **kwargs):
+        dmeb = self._driver.make_eb_coup_rdm(self, eris=False, amplitudes=self.amplitudes, lambdas=self.lambdas,
+                                             hermitise=hermitise, **kwargs)
+        return (dmeb[0].transpose((1, 2, 0)) / 2, dmeb[0].transpose((1, 2, 0)) / 2)
 
-    def make_eb_coup_rdm(self, *args, **kwargs):
-        return self._driver.make_eb_coup_rdm(self, eris=False, amplitudes=self.amplitudes, lambdas=self.lambdas, hermitise=True)
+    make_rdm1_f = make_rdm1
+    make_rdm2_f = make_rdm2
 
     def copy(self):
         proj = callif(spinalg.copy, self.projector)
@@ -293,3 +308,9 @@ def make_rdm1(self, *args, **kwargs):
     def make_rdm2(self, *args, **kwargs):
         dm2 = super().make_rdm2(*args, **kwargs)
         return dm2.aaaa, dm2.aabb, dm2.bbbb
+
+    def make_rdm_eb(self, hermitise=True, **kwargs):
+        dmeb = self._driver.make_eb_coup_rdm(self, eris=False, amplitudes=self.amplitudes, lambdas=self.lambdas,
+                                             hermitise=hermitise, **kwargs)
+
+        return (dmeb.aa[0].transpose(1, 2, 0), dmeb.bb[0].transpose(1, 2, 0))

vayesta/edmet/fragment.py

@@ -688,7 +688,7 @@ def kernel(self, solver=None, eris=None, construct_bath=False,
         dm2 = wf.make_rdm2()
         if self.nbos > 0:
             self.check_qba_approx(dm1)
-        dm_eb = wf.make_rdmeb()
+        dm_eb = wf.make_rdm_eb()
         self._results = results = self.Results(fid=self.id, n_active=self.cluster.norb_active,
                 converged=True, wf=wf, dm1=dm1, dm2=dm2, dm_eb=dm_eb)
         results.e1, results.e2, results.e_fb = self.get_edmet_energy_contrib()

vayesta/ewf/ewf.py

@@ -27,8 +27,6 @@ class Options(Embedding.Options):
     # --- Bath settings
     bath_options: dict = Embedding.Options.change_dict_defaults('bath_options',
             bathtype='mp2', threshold=1e-8)
-    solver_options: dict = Embedding.Options.change_dict_defaults('solver_options',
-            fermion_wf=True)
     #ewdmet_max_order: int = 1
     # If multiple bno thresholds are to be calculated, we can project integrals and amplitudes from a previous larger cluster:
     project_eris: bool = False          # Project ERIs from a pervious larger cluster (corresponding to larger eta), can result in a loss of accuracy especially for large basis sets!

vayesta/solver/ebcc.py

@@ -27,7 +27,7 @@ def kernel(self):
 
         mf_clus.mo_coeff, space = self.get_space(self.hamil.cluster.c_active, mf_clus.mo_occ, frozen=frozen)
 
-        mycc = ebcc.EBCC(mf_clus, log=self.log, ansatz=self.opts.ansatz, space=space, **self.get_nonnull_solver_opts())
+        mycc = ebcc.EBCC(mf_clus, log=self.log, ansatz=self.opts.ansatz, space=space, shift=False, **self.get_nonnull_solver_opts())
         mycc.kernel()
         self.converged = mycc.converged
         if self.opts.solve_lambda:
@@ -69,7 +69,6 @@ def construct_wavefunction(self, mycc, mo):
             except TypeError:
                 self.wf = CCSD_WaveFunction(mo, mycc.t1, mycc.t2)
         else:
-
             self.wf = EBCC_WaveFunction(mo, mycc.ansatz, mycc.amplitudes, mycc.lambdas, None)
 
         # Need to rotate wavefunction back into original cluster active space.
@@ -209,6 +208,20 @@ def get_couplings(self):
         # EBCC wants contribution  g_{xpq} p^\\dagger q b; need to transpose to get this contribution.
         return self.hamil.couplings.transpose(0, 2, 1)
 
+    def construct_wavefunction(self, mycc, mo):
+
+        nbosons = len(self.hamil.bos_freqs)
+
+        class dummy_mbos:
+            @property
+            def nbos(self):
+                return nbosons
+
+        self.wf = EBCC_WaveFunction(mo, mycc.ansatz, mycc.amplitudes, mycc.lambdas, mbos=dummy_mbos(),
+                                    xi=self.hamil.polaritonic_shift)
+        self.wf.rotate(t=mycc.mo_coeff.T, inplace=True)
+
+
 class EB_UEBCC_Solver(EB_REBCC_Solver, UEBCC_Solver):
     @dataclasses.dataclass
     class Options(UEBCC_Solver.Options, EB_REBCC_Solver.Options):
@@ -218,6 +231,19 @@ def get_couplings(self):
         # EBCC wants contribution  g_{xpq} p^\\dagger q b; need to transpose to get this contribution.
         return tuple([x.transpose(0, 2, 1) for x in self.hamil.couplings])
 
+    def construct_wavefunction(self, mycc, mo):
+
+        nbosons = len(self.hamil.bos_freqs)
+
+        class dummy_mbos:
+            @property
+            def nbos(self):
+                return nbosons
+
+        self.wf = EBCC_WaveFunction(mo, mycc.ansatz, mycc.amplitudes, mycc.lambdas, mbos=dummy_mbos(),
+                                    xi=self.hamil.polaritonic_shift)
+        self.wf.rotate(t=[x.T for x in mycc.mo_coeff], inplace=True)
+
 
 def gen_space(c_occ, c_vir, co_active=None, cv_active=None, frozen_orbs=None):
     """Given the occupied and virtual orbital coefficients in the local cluster basis, which orbitals are frozen, and

vayesta/solver/ebfci.py

@@ -26,7 +26,7 @@ def kernel(self):
         self.wf = WaveFunction(self.hamil.mo)
         self.wf.make_rdm1 = lambda *args, **kwargs: solver.make_rdm1(*args, **kwargs)
         self.wf.make_rdm2 = lambda *args, **kwargs: solver.make_rdm2(*args, **kwargs)
-        self.wf.make_rdmeb = lambda *args, **kwargs: np.array(solver.make_rdm_eb(*args, **kwargs)) + \
+        self.wf.make_rdm_eb = lambda *args, **kwargs: np.array(solver.make_rdm_eb(*args, **kwargs)) + \
                                                      np.array(
                                                          self.hamil.get_eb_dm_polaritonic_shift(self.wf.make_rdm1()))
         self.wf.make_dd_moms = lambda max_mom, *args, **kwargs: solver.make_dd_moms(max_mom, *args, **kwargs)

vayesta/solver/hamiltonian.py

@@ -748,6 +748,8 @@ def __init__(self, *args, **kwargs):
         self.bos_freqs = self._fragment.bos_freqs
         if self.opts.polaritonic_shift:
             self.set_polaritonic_shift(self.bos_freqs, self.unshifted_couplings)
+        else:
+            self._polaritonic_shift = None
 
     @property
     def polaritonic_shift(self):