Global WF DM1 using BTensor library

examples/ewf/molecules/btensor-dm.py

@@ -0,0 +1,96 @@
+from time import perf_counter
+
+import numpy as np
+
+import pyscf
+import pyscf.gto
+import pyscf.scf
+import pyscf.cc
+import vayesta
+import vayesta.ewf
+
+mol = pyscf.gto.Mole()
+
+nwater = 4
+use_sym = True
+
+if nwater == 1:
+    mol.atom = """
+    O  0.0000   0.0000   0.1173
+    H  0.0000   0.7572  -0.4692
+    H  0.0000  -0.7572  -0.4692
+    """
+elif nwater == 2:
+    mol.atom = """
+    O  2.0000   0.0000   0.1173
+    H  2.0000   0.7572  -0.4692
+    H  2.0000  -0.7572  -0.4692
+    O  6.0000   0.0000   0.1173
+    H  6.0000   0.7572  -0.4692
+    H  6.0000  -0.7572  -0.4692
+    """
+elif nwater == 4:
+    mol.atom = """
+    O  0.0000   0.0000   0.1173
+    H  0.0000   0.7572  -0.4692
+    H  0.0000  -0.7572  -0.4692
+    O  2.0000   0.0000   0.1173
+    H  2.0000   0.7572  -0.4692
+    H  2.0000  -0.7572  -0.4692
+    O  6.0000   0.0000   0.1173
+    H  6.0000   0.7572  -0.4692
+    H  6.0000  -0.7572  -0.4692
+    O  8.0000   0.0000   0.1173
+    H  8.0000   0.7572  -0.4692
+    H  8.0000  -0.7572  -0.4692
+    """
+
+#mol.basis = 'cc-pVDZ'
+#mol.basis = 'cc-pVTZ'
+mol.basis = 'cc-pVQZ'
+mol.output = 'pyscf.txt'
+mol.build()
+
+# Hartree-Fock
+mf = pyscf.scf.RHF(mol)
+mf.kernel()
+
+# Embedded CCSD
+eta = 1e-6
+emb = vayesta.ewf.EWF(mf, bath_options=dict(threshold=eta))
+with emb.fragmentation() as f:
+    if nwater > 1 and use_sym:
+        with f.mirror_symmetry(axis='x', center=(4, 0, 0)):
+            for i in range(mol.natm//2):
+                f.add_atomic_fragment(i)
+    else:
+        f.add_all_atomic_fragments()
+emb.kernel()
+
+
+# Density-matrix
+
+#svd_tol = 1e-3
+svd_tol = None
+ovlp_tol = None
+with_t1 = True
+
+times = [perf_counter()]
+dm1 = emb._make_rdm1_ccsd_global_wf(svd_tol=svd_tol, ovlp_tol=ovlp_tol, with_t1=with_t1, use_sym=use_sym)
+times.append(perf_counter())
+
+print(f"Vayesta time 1= {times[1]-times[0]:.3f}")
+
+times = [perf_counter()]
+dm1_bt1 = emb._make_rdm1_ccsd_global_wf_btensor(with_t1=with_t1, svd_tol=svd_tol, ovlp_tol=ovlp_tol, use_sym=use_sym)
+times.append(perf_counter())
+dm1_bt1 = emb._make_rdm1_ccsd_global_wf_btensor(with_t1=with_t1, svd_tol=svd_tol, ovlp_tol=ovlp_tol, use_sym=use_sym)
+times.append(perf_counter())
+dm1_bt1 = emb._make_rdm1_ccsd_global_wf_btensor(with_t1=with_t1, svd_tol=svd_tol, ovlp_tol=ovlp_tol, use_sym=use_sym)
+times.append(perf_counter())
+
+print(f"BTensor time 1= {times[1]-times[0]:.3f}")
+print(f"BTensor time 2= {times[2]-times[1]:.3f}")
+print(f"BTensor time 3= {times[3]-times[2]:.3f}")
+
+print(f"Error in DM= {np.linalg.norm(dm1_bt1 - dm1)}")

vayesta/core/qemb/fragment.py

@@ -1,8 +1,10 @@
+from __future__ import annotations
 # --- Standard library
 import dataclasses
 import itertools
 import os.path
 import typing
+from typing import *
 
 # --- External
 import numpy as np
@@ -23,6 +25,7 @@
     time_string,
     timer,
     AbstractMethodError,
+    optional_import,
 )
 from vayesta.core import spinalg
 from vayesta.core.types import Cluster, Orbitals
@@ -50,6 +53,9 @@
 from vayesta.mpi import mpi
 from vayesta.solver import get_solver_class, check_solver_config, ClusterHamiltonian
 from vayesta.core.screening.screening_crpa import get_frag_W
+# Optional
+btensor = optional_import('btensor')
+
 
 # Get MPI rank of fragment
 get_fragment_mpi_rank = lambda *args: args[0].mpi_rank
@@ -809,6 +815,37 @@ def get_yield(fragment, arrays):
                     intermediates, axes=axes, symtree=grandchildren, maxgen=(maxgen - 1)
                 )
 
+    def _loop_symmetry_children_operations(self, symtree=None, maxgen=1000):
+        """Loop over all symmetry related fragments and return together with symmetry operations."""
+        if maxgen == 0:
+            return
+        if symtree is None:
+            symtree = self.get_symmetry_tree()
+        for child, grandchildren in symtree:
+            yield child, (child.sym_op,)
+            if grandchildren and maxgen > 1:
+                yield from [(x[0], (child.sym_op, *x[1])) for x in child._loop_symmetry_children_operations(
+                    symtree=grandchildren, maxgen=maxgen - 1)]
+
+    def loop_ao_symmetry_bases(self, ao: 'btensor.Basis') -> Iterator['btensor.Basis']:
+        """Returns an iterator over symmetry related AO bases.
+
+        Args:
+            ao: AO basis.
+
+        Yields:
+            Iterator over bases, which are symmetry related to the original AO basis.
+        """
+        for fx2, sym_ops in self._loop_symmetry_children_operations():
+            # Simple translation only require a reordering of the AOs:
+            if len(sym_ops) == 1 and isinstance(sym_ops[0], SymmetryTranslation):
+                trafo = sym_ops[0].ao_reorder
+            else:
+                trafo = np.eye(self.base.nao)
+                for sym_op in sym_ops:
+                    trafo = sym_op(trafo)
+            yield ao.make_subbasis(trafo, name=f'ao-sym[{fx2.id}]')
+
     @property
     def n_symmetry_children(self):
         """Includes children of children, etc."""

vayesta/core/qemb/qemb.py

@@ -5,7 +5,7 @@
 import itertools
 import os
 import os.path
-from typing import Optional
+from typing import *
 
 import numpy as np
 
@@ -31,6 +31,7 @@
     log_method,
     log_time,
     with_doc,
+    optional_import,
 )
 from vayesta.core import spinalg, eris
 from vayesta.core.scmf import PDMET, Brueckner
@@ -67,6 +68,8 @@
 # from . import helper
 from vayesta.core.qemb.rdm import make_rdm1_demo_rhf
 from vayesta.core.qemb.rdm import make_rdm2_demo_rhf
+# Optional
+btensor = optional_import('btensor')
 
 
 @dataclasses.dataclass
@@ -1233,6 +1236,30 @@ def svd(cx, cy):
                 fx._occ_bath_factory = None
                 fx._vir_bath_factory = None
 
+    def _build_basis_dict(self, include_symmetry: bool = False) -> Dict[str, 'btensor.Basis']:
+        """Build dictionary containing BTensor basis objects."""
+        # Mean-field bases:
+        ao = btensor.Basis(self.mol.nao, metric=self.get_ovlp(), name='ao')
+        mo = ao.make_subbasis(self.mo_coeff, name='mo', orthonormal=True)
+        occ = mo.make_subbasis(self.mo_occ > 0, name='occ', orthonormal=True)
+        vir = mo.make_subbasis(self.mo_occ == 0, name='vir', orthonormal=True)
+        basis_dict = dict(ao=ao, mo=mo, occ=occ, vir=vir)
+        # Fragment specific bases:
+        frag_filter = {} if include_symmetry else dict(sym_parent=None)
+        for fx in self.get_fragments(**frag_filter):
+            # Fragment orbitals
+            c_frag = fx.get_overlap('mo|frag')
+            name_frag = f'frag[{fx.id}]'
+            basis_dict[name_frag] = mo.make_subbasis(c_frag, name=name_frag, orthonormal=True)
+            # Occupied cluster orbitals
+            c_occ = fx.get_overlap('mo[occ]|cluster[occ]')
+            name_occ = f'occ[{fx.id}]'
+            basis_dict[name_occ] = occ.make_subbasis(c_occ, name=name_occ, orthonormal=True)
+            c_vir = fx.get_overlap('mo[vir]|cluster[vir]')
+            name_vir = f'vir[{fx.id}]'
+            basis_dict[name_vir] = vir.make_subbasis(c_vir, name=name_vir, orthonormal=True)
+        return basis_dict
+
     # Results
     # -------
 

vayesta/core/util.py

@@ -1,6 +1,10 @@
+from __future__ import annotations
+
+import types
 from contextlib import contextmanager
 from copy import deepcopy
 import itertools
+import importlib
 import dataclasses
 import functools
 import logging
@@ -9,6 +13,7 @@
 import string
 import sys
 from timeit import default_timer
+from typing import *
 
 
 try:
@@ -62,6 +67,7 @@
     "getif",
     "callif",
     "permutations_with_signs",
+    "MissingModule",
 ]
 
 
@@ -503,7 +509,7 @@ def time_string(seconds, show_zeros=False):
     elif seconds >= 60:
         tstr = "%.0f min %.0f s" % (m, s)
     else:
-        tstr = "%.1f s" % s
+        tstr = "%.5f s" % s
     if sign == -1:
         tstr = "-%s" % tstr
     return tstr
@@ -763,3 +769,20 @@ def _permutations(seq):
         return items
 
     return [(item, -1 if i % 2 else 1) for i, item in enumerate(_permutations(list(seq)))]
+
+
+class MissingModule:
+
+    def __init__(self, module_name: str) -> None:
+        self.module_name = module_name
+
+    def __getattr__(self, item: Any) -> NoReturn:
+        raise ModuleNotFoundError(f"module {self.module_name} required, but not found")
+
+
+def optional_import(module_name: str) -> types.ModuleType | MissingModule:
+    """Attempt import, return module if successful and MissingModule object if not."""
+    try:
+        return importlib.import_module(module_name)
+    except ModuleNotFoundError:
+        return MissingModule(module_name)

vayesta/ewf/ewf.py

@@ -28,6 +28,7 @@
 from vayesta.ewf.rdm import make_rdm2_ccsd_global_wf
 from vayesta.ewf.rdm import make_rdm1_ccsd_proj_lambda
 from vayesta.ewf.rdm import make_rdm2_ccsd_proj_lambda
+from vayesta.ewf.rdm import make_rdm1_ccsd_global_wf_btensor
 from vayesta.ewf.icmp2 import get_intercluster_mp2_energy_rhf
 
 
@@ -273,6 +274,15 @@ def _make_rdm1_ccsd_global_wf(self, *args, ao_basis=False, with_mf=True, **kwarg
             dm1 = dot(self.mo_coeff, dm1, self.mo_coeff.T)
         return dm1
 
+    @log_method()
+    def _make_rdm1_ccsd_global_wf_btensor(self, *args, ao_basis=False, with_mf=True, **kwargs):
+        dm1 = make_rdm1_ccsd_global_wf_btensor(self, *args, **kwargs)
+        if with_mf:
+            dm1[np.diag_indices(self.nocc)] += 2
+        if ao_basis:
+            dm1 = dot(self.mo_coeff, dm1, self.mo_coeff.T)
+        return dm1
+
     @cache(copy=True)
     def _make_rdm1_ccsd_global_wf_cached(self, *args, **kwargs):
         return make_rdm1_ccsd_global_wf(self, *args, **kwargs)