Merge pull request #118 from BoothGroup/variational_embedding

Variational embedding

.github/workflows/ci.yaml

@@ -35,7 +35,7 @@ jobs:
           python -m pip install wheel --user
           python -m pip install setuptools --upgrade --user
           python -m pip install https://github.com/BoothGroup/dyson/archive/master.zip
-          python -m pip install -e .[dmet,ebcc] --user
+          python -m pip install -e .[dmet,ebcc] --user 
       - name: Run unit tests
         run: |
           python -m pip install pytest pytest-cov --user

examples/variational_embedding/01-simple-vembedding.py

@@ -0,0 +1,161 @@
+import vayesta.ewf
+
+from pyscf import gto, scf, lib, cc, fci
+from vayesta.misc.variational_embedding import variational_params
+from vayesta.misc.molecules import ring
+
+import numpy as np
+
+
+def run_ewf(natom, r, n_per_frag=1, bath_options={"bathtype":"dmet"}):
+    if abs(natom / n_per_frag - natom // n_per_frag) > 1e-6:
+        raise ValueError(f"Atoms per fragment ({n_per_frag}) doesn't exactly divide natoms ({natom})")
+
+    nfrag = natom // n_per_frag
+
+    mol = gto.M()
+    mol.atom = ring("H", natom, bond_length=r)
+    mol.basis = "sto-3g"
+    mol.spin = 0
+    mol.charge = 0
+    mol.build()
+
+    rmf = scf.RHF(mol).density_fit();
+    rmf.conv_tol = 1e-12;
+    rmf.kernel()
+
+    out = rmf.stability(external=True)
+
+    rewf = vayesta.ewf.EWF(rmf, solver="FCI", bath_options=bath_options)
+    with rewf.iao_fragmentation() as f:
+        with f.rotational_symmetry(nfrag, axis="z"):
+            f.add_atomic_fragment(atoms=list(range(n_per_frag)))
+    rewf.kernel()
+    return rewf, rmf
+
+
+def get_wf_composite(emb, inc_mf=False):
+    """Compute energy resulting from generalised eigenproblem between all local cluster wavefunctions."""
+    h, s, dm = variational_params.get_wf_couplings(emb, inc_mf=inc_mf)
+    w_bare, v, seig = lib.linalg_helper.safe_eigh(h, s, lindep=1e-12)
+    # Return lowest eigenvalue.
+    return w_bare[0]
+
+
+def get_density_projected(emb, inc_mf=False):
+    p_frags = [x.get_fragment_projector(x.cluster.c_active) / emb.mf.mol.nelectron for x in emb.fragments]
+    barewfs = [x.results.wf for x in emb.fragments]
+    wfs = [x.project(y) for x, y in zip(barewfs, p_frags)]
+    h, s, dm = variational_params.get_wf_couplings(emb, emb.fragments, wfs, inc_mf=inc_mf)
+    sum_energy = sum(h.reshape(-1)) / sum(s.reshape(-1))
+    w, v, seig = lib.linalg_helper.safe_eigh(h, s, lindep=1e-12)
+    return sum_energy, w[0]
+
+
+def get_occ_projected(emb):
+    p_frags = [x.get_overlap('frag|cluster-occ') for x in emb.fragments]
+    p_frags = [np.dot(x.T, x) for x in p_frags]
+    wfs = [x.results.wf.project_occ(y) for x, y in zip(emb.fragments, p_frags)]
+    h, s, dm = variational_params.get_wf_couplings(emb, wfs=wfs, inc_mf=True)
+    sum_energy = sum(h.reshape(-1)) / sum(s.reshape(-1))
+    w, v, seig = lib.linalg_helper.safe_eigh(h, s, lindep=1e-12)
+    return sum_energy, w[0]
+
+
+def gen_comp_graph(fname):
+    import matplotlib.pyplot as plt
+    fig, (ax0, ax1) = plt.subplots(2, 1, sharex=True)
+    fig.set_tight_layout(True)
+    plot_results(fname, False, ax=ax0)
+    plot_results(fname, True, ax=ax1)
+    fig.set_size_inches(9, 6)
+    fig.set_tight_layout(True)
+    plt.draw()
+
+
+def draw_comparison_error(fname1, fname2):
+    import matplotlib.pyplot as plt
+    fig, (ax0, ax1) = plt.subplots(1, 2, sharex=True, sharey=True)
+    fig.set_tight_layout(True)
+    plot_results(fname1, True, ax0)
+    plot_results(fname2, True, ax1)
+    fig.set_size_inches(9, 6)
+    plt.show(block=False)
+    plt.draw()
+
+
+def plot_results(fname="results.txt", vsfci=False, ax=None, nodmenergy=True):
+    import matplotlib.pyplot as plt
+    if ax is None:
+        ax = plt.subplots(1, 1)[1]
+    res = np.genfromtxt(fname)
+    labs = ["$r_{HH}/\AA$", "HF", "FCI", "CCSD", "EWF-proj-wf", "EWF-dm-wf", "NO-CAS-CI", "NO-dproj", "NO-dproj-CAS-CI",
+            "NO-oproj", "NO-oproj-CAS-CI", "var-NO-FCI"]
+
+    def remove_ind(results, labels, i):
+        labels = labels[:i] + labels[i + 1:]
+        results = np.concatenate([results[:, :i], results[:, i + 1:]], axis=1)
+        return results, labels
+
+    if nodmenergy:
+        res, labs = remove_ind(res, labs, 5)
+
+    if vsfci:
+        fcires = res[:, 2]
+        res, labs = remove_ind(res, labs, 2)
+        res[:, 1:] = (res[:, 1:].T - fcires).T
+    for i in range(1, res.shape[-1]):
+        ax.plot(res[:, 0], res[:, i], label=labs[i])
+    ax.set_xlabel(labs[0])
+
+    leg = ax.legend().set_draggable(True)
+
+    if vsfci:
+        ax.set_ylabel("E-$E_{FCI}/E_h$")
+    else:
+        ax.set_ylabel("E/$E_h$")
+
+    plt.show(block=False)
+
+
+if __name__ == "__main__":
+    import sys
+
+    nat = int(sys.argv[1])
+    n_per_frag = int(sys.argv[2])
+    for r in list(np.arange(0.6, 2.0, 0.1)) + list(np.arange(2.5, 10.0, 0.5)):
+        emb, mf = run_ewf(nat, r, n_per_frag)
+        # These calculate the standard EWF energy estimators.
+        eewf_wf = emb.get_wf_energy()
+        eewf_dm = emb.get_dm_energy()
+        # This calculates the energy of the variationally optimal combination of the local wavefunctions in each case.
+        # This uses the bare local wavefunctions...
+        e_barewf = get_wf_composite(emb)
+        # This uses the density projected local wavefunctions, and also returns the energy of a sum of these
+        # wavefunctions.
+        e_dense_proj, e_dense_opt = get_density_projected(emb)
+        # This does the same, but with the local wavefunction projected via the occupied projector (as in standard EWF).
+        e_occ_proj, e_occ_opt = get_occ_projected(emb)
+        # This variationally optimises all coefficients in the local wavefunctions simulanteously; the value of
+        # emb.fragments[x].results.wf is updated to the local portion of the global wavefunction.
+        e_opt = variational_params.optimise_full_varwav(emb)
+
+        # This computes the FCI and CCSD energies for comparison.
+        myfci = fci.FCI(mf)
+        efci = myfci.kernel()[0]
+        mycc = cc.CCSD(mf)
+        try:
+            mycc.kernel()
+
+            if mycc.converged:
+                ecc = mycc.e_tot
+            else:
+                ecc = np.nan
+        except np.linalg.LinAlgError:
+            ecc = np.nan
+
+        res = (mf.e_tot, efci, ecc, eewf_wf, eewf_dm, e_barewf, e_dense_proj, e_dense_opt, e_occ_proj, e_occ_opt, e_opt)
+
+        with open("results.txt", "a") as f:
+
+            f.write((f" {r:4.2f} ") + ("   {:12.8f}  " * len(res)).format(*res) + "\n")

pyproject.toml

@@ -68,6 +68,11 @@ dyson = [
 ebcc = [
     "ebcc"
 ]
+
+pygnme = [
+    "pygnme @ git+https://github.com/BoothGroup/pygnme@master" 
+]
+
 dev = [
     "cvxpy>=1.1",
     "mpi4py>=3.0.0",

vayesta/__init__.py

@@ -95,6 +95,8 @@ def import_package(name, required=True):
 import_package('cvxpy', False)
 dyson = import_package('dyson', False)
 ebcc = import_package('ebcc', False)
+import_package('pygnme', False)
+
 
 # --- Git hashes
 

vayesta/core/types/wf/fci.py

@@ -1,8 +1,10 @@
 import numpy as np
 import pyscf
 import pyscf.fci
-from vayesta.core.util import decompress_axes, dot, einsum, tril_indices_ndim, replace_attr
+from vayesta.core.util import decompress_axes, dot, einsum, tril_indices_ndim, callif, replace_attr
 from vayesta.core.types import wf as wf_types
+import scipy.sparse.linalg
+from vayesta.core import spinalg
 
 def FCI_WaveFunction(mo, ci, **kwargs):
     if mo.nspin == 1:
@@ -18,6 +20,10 @@ def __init__(self, mo, ci, projector=None):
         super().__init__(mo, projector=projector)
         self.ci = ci
 
+    @property
+    def nfci(self):
+        return self.ci.size
+
     def make_rdm1(self, ao_basis=False, with_mf=True):
         dm1 = pyscf.fci.direct_spin1.make_rdm1(self.ci, self.norb, self.nelec)
         if not with_mf:
@@ -47,7 +53,53 @@ def make_rdm2(self, ao_basis=False, with_dm1=True, approx_cumulant=True):
         return einsum('ijkl,ai,bj,ck,dl->abcd', dm2, *(4*[self.mo.coeff]))
 
     def project(self, projector, inplace=False):
-        raise NotImplementedError
+        """Apply one-body projector to the FCI wavefunction using pyscf.
+        This is assumed to  indicate a one-body
+        """
+        wf = self if inplace else self.copy()
+        # Apply one-body operator of projector to ci string.
+        wf.ci = self._apply_onebody(projector)
+        return wf
+
+    def project_occ(self, projector, inplace=False):
+        """Apply projector onto the occupied indices of all CI coefficient tensors.
+        Note that `projector` is nocc x nocc.
+
+        Action of occupied projector can be written as
+        P^{x}_{occ} = P_{ij}^{x}
+        """
+        c0 = self.c0
+        # Get result of applying bare projector; need to keep original ci string just in case
+        ci0 = self.ci
+        # Pad projector from occupied to full orbital space.
+        projector = np.pad(projector, ((0, self.nvir),)).T
+
+        wf = self.project(projector, inplace)
+        wf.ci = (2*sum(np.diag(projector)) * ci0 - wf.ci) / c0
+
+        # Now just have to divide each coefficient by its excitation level; this corresponds to action of
+        #    R^{-1} = (1 + \sum_{i\in occ} i i^+)^{-1} = (N_{elec} + 1 - \sum_{i\in occ} i^+ i)^{-1}
+        # So we seek x to solve
+        #           x = R^{-1} a
+        # which could be obtained straightforwardly by solving
+        #           Rx = a.
+        # In practice, it is more stable to just compute the excitation level of each state and divide by it.
+        mf_vdensity_op = np.eye(self.norb) - np.pad(np.eye(self.nocc), ((0, self.nvir),))
+
+        # Set up sparse LinearOperator object to apply the hole counting operator to the FCI string.
+        def myop(ci):
+            return self._apply_onebody(mf_vdensity_op, ci)
+
+        cishape = wf.ci.shape
+        # Calculate excitation level+1 for all states using this operation.
+        ex_lvl = myop(np.full_like(wf.ci, fill_value=1.0).reshape(-1)).reshape(cishape)
+        ex_lvl[0,0] = 1.0
+        wf.ci = einsum("pq,pq->pq", ex_lvl**(-1), wf.ci)
+        return wf
+
+    def _apply_onebody(self, proj, ci=None):
+        ci = self.ci if ci is None else ci
+        return pyscf.fci.direct_spin1.contract_1e(proj, ci, self.norb, self.nelec)
 
     def restore(self, projector=None, inplace=False):
         raise NotImplementedError
@@ -56,6 +108,10 @@ def restore(self, projector=None, inplace=False):
     def c0(self):
         return self.ci[0,0]
 
+    def copy(self):
+        proj = callif(spinalg.copy, self.projector)
+        return type(self)(self.mo.copy(), spinalg.copy(self.ci), projector=proj)
+
     def as_unrestricted(self):
         mo = self.mo.to_spin_orbitals()
         return UFCI_WaveFunction(mo, self.ci)
@@ -213,6 +269,47 @@ def make_rdm2(self, ao_basis=False, with_dm1=True, approx_cumulant=True):
                 einsum('ijkl,ai,bj,ck,dl->abcd', dm2[1], *[moa, moa, mob, mob]),
                 einsum('ijkl,ai,bj,ck,dl->abcd', dm2[2], *(4*[mob])))
 
+    def project_occ(self, projector, inplace=False):
+        """Apply projector onto the occupied indices of all CI coefficient tensors.
+        Note that `projector` is nocc x nocc.
+
+        Action of occupied projector can be written as
+        P^{x}_{occ} = P_{ij}^{x}
+        """
+        c0 = self.c0
+        # Get result of applying bare projector; need to keep original ci string just in case
+        ci0 = self.ci
+        # Pad projector from occupied to full orbital space.
+        projector = (np.pad(projector[0], ((0, self.nvir[0]),)).T, np.pad(projector[1], ((0, self.nvir[1]),)).T)
+
+        wf = self.project(projector, inplace)
+        wf.ci = ((projector[0].trace() + projector[1].trace()) * ci0 - wf.ci) / c0
+
+        # Now just have to divide each coefficient by its excitation level; this corresponds to action of
+        #    R^{-1} = (1 + \sum_{i\in occ} i i^+)^{-1} = (N_{elec} + 1 - \sum_{i\in occ} i^+ i)^{-1}
+        # So we seek x to solve
+        #           x = R^{-1} a
+        # which could be obtained straightforwardly by solving
+        #           Rx = a.
+        # In practice, it is more stable to just compute the excitation level of each state and divide by it.
+        mf_vdensity_op = tuple([np.eye(self.norb[i]) - np.pad(np.eye(self.nocc[i]), ((0, self.nvir[i]),)) for i in [0, 1]])
+
+        # Set up sparse LinearOperator object to apply the hole counting operator to the FCI string.
+        def myop(ci):
+            return self._apply_onebody(mf_vdensity_op, ci)
+
+        cishape = wf.ci.shape
+        # Calculate excitation level+1 for all states using this operation.
+        ex_lvl = myop(np.full_like(wf.ci, fill_value=1.0).reshape(-1)).reshape(cishape)
+        ex_lvl[0,0] = 1.0
+        wf.ci = einsum("pq,pq->pq", ex_lvl**(-1), wf.ci)
+        return wf
+
+    def _apply_onebody(self, proj, ci=None):
+        ci = self.ci if ci is None else ci
+        assert(self.norb[0] == self.norb[1])
+        return pyscf.fci.direct_uhf.contract_1e(proj, ci, self.norb[0], self.nelec)
+
     def as_cisd(self, c0=None):
         if self.projector is not None:
             raise NotImplementedError

vayesta/misc/variational_embedding/__init__.py

@@ -0,0 +1,4 @@
+try:
+    import pygnme
+except ModuleNotFoundError as e:
+    raise ModuleNotFoundError("pygnme is required for variational embedding approaches.") from e