From bf4605f8af8f2fe884cef6b97bc809e1f9602e62 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Thu, 19 Dec 2024 15:39:31 -0500 Subject: [PATCH 01/67] put if __main__ guard to test --- tests/molbe_octane_get_rdms_test.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/molbe_octane_get_rdms_test.py b/tests/molbe_octane_get_rdms_test.py index 75bba445..2dde6db2 100644 --- a/tests/molbe_octane_get_rdms_test.py +++ b/tests/molbe_octane_get_rdms_test.py @@ -67,4 +67,5 @@ def test_rdm(): assert np.isclose(mybe.ebe_tot, -310.3311676424482) -test_rdm() +if __name__ == "__main__": + test_rdm() From eb2969d396647ee58465dbab6a0af1fbb9bc8552 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Thu, 19 Dec 2024 15:39:40 -0500 Subject: [PATCH 02/67] use h5py contextmanager nearly everywhere --- src/quemb/kbe/pbe.py | 16 +++++++--------- src/quemb/kbe/pfrag.py | 6 ++---- src/quemb/molbe/helper.py | 26 +++++++++++--------------- src/quemb/molbe/mbe.py | 16 +++++++--------- src/quemb/molbe/misc.py | 15 ++++++--------- src/quemb/molbe/pfrag.py | 12 +++++------- 6 files changed, 38 insertions(+), 53 deletions(-) diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py index 22dedad0..1a7d8bdc 100644 --- a/src/quemb/kbe/pbe.py +++ b/src/quemb/kbe/pbe.py @@ -794,11 +794,10 @@ def write_heff(self, heff_file="bepotfile.h5"): heff_file : str, optional Path to the file to store effective Hamiltonian, by default 'bepotfile.h5'. """ - filepot = h5py.File(heff_file, "w") - for fobj in self.Fobjs: - print(fobj.heff.shape, fobj.dname, flush=True) - filepot.create_dataset(fobj.dname, data=fobj.heff) - filepot.close() + with h5py.File(heff_file, "w") as filepot: + for fobj in self.Fobjs: + print(fobj.heff.shape, fobj.dname, flush=True) + filepot.create_dataset(fobj.dname, data=fobj.heff) def read_heff(self, heff_file="bepotfile.h5"): """ @@ -809,10 +808,9 @@ def read_heff(self, heff_file="bepotfile.h5"): heff_file : str, optional Path to the file storing effective Hamiltonian, by default 'bepotfile.h5'. """ - filepot = h5py.File(heff_file, "r") - for fobj in self.Fobjs: - fobj.heff = filepot.get(fobj.dname) - filepot.close() + with h5py.File(heff_file, "r") as filepot: + for fobj in self.Fobjs: + fobj.heff = filepot.get(fobj.dname) def initialize_pot(Nfrag, edge_idx): diff --git a/src/quemb/kbe/pfrag.py b/src/quemb/kbe/pfrag.py index 411f5207..b57a3193 100644 --- a/src/quemb/kbe/pfrag.py +++ b/src/quemb/kbe/pfrag.py @@ -421,10 +421,8 @@ def update_ebe_hf( else: jmax = self.TA.shape[1] if eri is None: - r = h5py.File(self.eri_file, "r") - eri = r[self.dname][()] - - r.close() + with h5py.File(self.eri_file, "r") as r: + eri = r[self.dname][()] e2 = numpy.zeros_like(e1) for i in range(self.nfsites): diff --git a/src/quemb/molbe/helper.py b/src/quemb/molbe/helper.py index 8b31b86b..e221c9f2 100644 --- a/src/quemb/molbe/helper.py +++ b/src/quemb/molbe/helper.py @@ -163,16 +163,14 @@ def get_eri(i_frag, Nao, symm=8, ignore_symm=False, eri_file="eri_file.h5"): Electron repulsion integrals, possibly restored with symmetry. """ # Open the HDF5 file and read the ERI for the specified fragment - r = h5py.File(eri_file, "r") - eri__ = numpy.array(r.get(i_frag)) + with h5py.File(eri_file, "r") as r: + eri__ = numpy.array(r.get(i_frag)) - # Optionally restore the symmetry of the ERI - if not ignore_symm: - # Set the number of threads for the library to 1 - lib.num_threads(1) - eri__ = ao2mo.restore(symm, eri__, Nao) - - r.close() + # Optionally restore the symmetry of the ERI + if not ignore_symm: + # Set the number of threads for the library to 1 + lib.num_threads(1) + eri__ = ao2mo.restore(symm, eri__, Nao) return eri__ @@ -279,9 +277,8 @@ def get_frag_energy( jmax = TA.shape[1] # Load the electron repulsion integrals from the HDF5 file - r = h5py.File(eri_file, "r") - eri = r[dname][()] - r.close() + with h5py.File(eri_file, "r") as r: + eri = r[dname][()] # Rotate the RDM2 into the MO basis rdm2s = numpy.einsum( @@ -413,9 +410,8 @@ def get_frag_energy_u( jmax = [TA[0].shape[1], TA[1].shape[1]] # Load ERIs from the HDF5 file - r = h5py.File(eri_file, "r") - Vs = [r[dname[0]][()], r[dname[1]][()], r[dname[2]][()]] - r.close() + with h5py.File(eri_file, "r") as r: + Vs = [r[dname[0]][()], r[dname[1]][()], r[dname[2]][()]] # Rotate the RDM2 into the MO basis rdm2s_k = [ diff --git a/src/quemb/molbe/mbe.py b/src/quemb/molbe/mbe.py index 0744cc2b..a0b5a31d 100644 --- a/src/quemb/molbe/mbe.py +++ b/src/quemb/molbe/mbe.py @@ -1048,11 +1048,10 @@ def write_heff(self, heff_file="bepotfile.h5"): heff_file : str, optional Path to the file to store effective Hamiltonian, by default 'bepotfile.h5'. """ - filepot = h5py.File(heff_file, "w") - for fobj in self.Fobjs: - print(fobj.heff.shape, fobj.dname, flush=True) - filepot.create_dataset(fobj.dname, data=fobj.heff) - filepot.close() + with h5py.File(heff_file, "w") as filepot: + for fobj in self.Fobjs: + print(fobj.heff.shape, fobj.dname, flush=True) + filepot.create_dataset(fobj.dname, data=fobj.heff) def read_heff(self, heff_file="bepotfile.h5"): """ @@ -1063,10 +1062,9 @@ def read_heff(self, heff_file="bepotfile.h5"): heff_file : str, optional Path to the file storing effective Hamiltonian, by default 'bepotfile.h5'. """ - filepot = h5py.File(heff_file, "r") - for fobj in self.Fobjs: - fobj.heff = filepot.get(fobj.dname) - filepot.close() + with h5py.File(heff_file, "r") as filepot: + for fobj in self.Fobjs: + fobj.heff = filepot.get(fobj.dname) def initialize_pot(Nfrag, edge_idx): diff --git a/src/quemb/molbe/misc.py b/src/quemb/molbe/misc.py index a8f48e19..b695f4dc 100644 --- a/src/quemb/molbe/misc.py +++ b/src/quemb/molbe/misc.py @@ -122,9 +122,8 @@ def be2fcidump(be_obj, fcidump_prefix, basis): """ for fidx, frag in enumerate(be_obj.Fobjs): # Read in eri - read = h5py.File(frag.eri_file, "r") - eri = read[frag.dname][()] # 2e in embedding basis - read.close() + with h5py.File(frag.eri_file, "r") as read: + eri = read[frag.dname][()] # 2e in embedding basis eri = ao2mo.restore(1, eri, frag.nao) if basis == "embedding": h1e = frag.fock @@ -173,9 +172,8 @@ def ube2fcidump(be_obj, fcidump_prefix, basis): """ for fidx, frag in enumerate(be_obj.Fobjs_a): # Read in eri - read = h5py.File(frag.eri_file, "r") - eri = read[frag.dname][()] # 2e in embedding basis - read.close() + with h5py.File(frag.eri_file, "r") as read: + eri = read[frag.dname][()] # 2e in embedding basis eri = ao2mo.restore(1, eri, frag.nao) if basis == "embedding": h1e = frag.fock @@ -208,9 +206,8 @@ def ube2fcidump(be_obj, fcidump_prefix, basis): for fidx, frag in enumerate(be_obj.Fobjs_b): # Read in eri - read = h5py.File(frag.eri_file, "r") - eri = read[frag.dname][()] # 2e in embedding basis - read.close() + with h5py.File(frag.eri_file, "r") as read: + eri = read[frag.dname][()] # 2e in embedding basis eri = ao2mo.restore(1, eri, frag.nao) if basis == "embedding": h1e = frag.fock diff --git a/src/quemb/molbe/pfrag.py b/src/quemb/molbe/pfrag.py index 6752fb37..456a739a 100644 --- a/src/quemb/molbe/pfrag.py +++ b/src/quemb/molbe/pfrag.py @@ -340,13 +340,11 @@ def update_ebe_hf( else: jmax = self.TA.shape[1] if eri is None: - r = h5py.File(self.eri_file, "r") - if isinstance(self.dname, list): - eri = [r[self.dname[0]][()], r[self.dname[1]][()]] - else: - eri = r[self.dname][()] - - r.close() + with h5py.File(self.eri_file, "r") as r: + if isinstance(self.dname, list): + eri = [r[self.dname[0]][()], r[self.dname[1]][()]] + else: + eri = r[self.dname][()] e2 = numpy.zeros_like(e1) for i in range(self.nfsites): From fa6cc27f2409fc00e5153a3d1a0064a674b7c393 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Thu, 19 Dec 2024 15:46:49 -0500 Subject: [PATCH 03/67] changed mbe.StoreBE into attrs defined class --- src/quemb/molbe/mbe.py | 54 ++++++++++++++++-------------------------- 1 file changed, 21 insertions(+), 33 deletions(-) diff --git a/src/quemb/molbe/mbe.py b/src/quemb/molbe/mbe.py index a0b5a31d..454b8fdc 100644 --- a/src/quemb/molbe/mbe.py +++ b/src/quemb/molbe/mbe.py @@ -5,6 +5,8 @@ import h5py import numpy +from attrs import define +from numpy import float64 from pyscf import ao2mo, scf from quemb.molbe._opt import BEOPT @@ -19,42 +21,28 @@ get_be_error_jacobian as _ext_get_be_error_jacobian, ) from quemb.shared.helper import copy_docstring +from quemb.shared.typing import Matrix +@define class storeBE: - def __init__( - self, - Nocc, - hf_veff, - hcore, - S, - C, - hf_dm, - hf_etot, - W, - lmo_coeff, - enuc, - E_core, - C_core, - P_core, - core_veff, - mo_energy, - ): - self.Nocc = Nocc - self.hf_veff = hf_veff - self.hcore = hcore - self.S = S - self.C = C - self.hf_dm = hf_dm - self.hf_etot = hf_etot - self.W = W - self.lmo_coeff = lmo_coeff - self.enuc = enuc - self.E_core = E_core - self.C_core = C_core - self.P_core = P_core - self.core_veff = core_veff - self.mo_energy = mo_energy + # TODO: some of the types are most likely wrong. + # this has to be checked in a review + Nocc: int + hf_veff: Matrix[float64] + hcore: Matrix[float64] + S: Matrix[float64] + C: Matrix[float64] + hf_dm: Matrix[float64] + hf_etot: float + W: Matrix[float64] + lmo_coeff: Matrix[float64] + enuc: float + ek: float + E_core: float + C_core: float + P_core: float + core_veff: float class BE(MixinLocalize): From be6176b478967400f2034d53139f41ea6e932af2 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Thu, 19 Dec 2024 15:51:44 -0500 Subject: [PATCH 04/67] ignore broken link to numpy.float64 in documentation --- docs/source/nitpick-exceptions | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/source/nitpick-exceptions b/docs/source/nitpick-exceptions index 357361e3..ccc310dc 100644 --- a/docs/source/nitpick-exceptions +++ b/docs/source/nitpick-exceptions @@ -1 +1,2 @@ py:class optional +py:class numpy.float64 From 86fa44b084d15482285a1208908b9ae7b557b8ee Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Thu, 19 Dec 2024 15:54:58 -0500 Subject: [PATCH 05/67] fixed a bug from wrong close statement --- src/quemb/molbe/mbe.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/quemb/molbe/mbe.py b/src/quemb/molbe/mbe.py index a0b5a31d..1f85160d 100644 --- a/src/quemb/molbe/mbe.py +++ b/src/quemb/molbe/mbe.py @@ -294,7 +294,6 @@ def __init__( with open(save_file, "wb") as rfile: pickle.dump(store_, rfile, pickle.HIGHEST_PROTOCOL) - rfile.close() if not restart: # Initialize fragments and perform initial calculations From 39423200ae0e3a7183ba63eeefd14f3aea4489b0 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Thu, 19 Dec 2024 15:59:52 -0500 Subject: [PATCH 06/67] testsuite don't wait for analysis --- .github/workflows/quemb_unittest.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/quemb_unittest.yml b/.github/workflows/quemb_unittest.yml index 733e54d5..ac60f76f 100644 --- a/.github/workflows/quemb_unittest.yml +++ b/.github/workflows/quemb_unittest.yml @@ -79,7 +79,6 @@ jobs: testsuite: runs-on: ubuntu-latest - needs: analysis strategy: matrix: python-version: ["3.10", "3.13"] From 499aa5b1edb1dfb8e7a93ad3151d541e0cbdf610 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Thu, 19 Dec 2024 16:01:09 -0500 Subject: [PATCH 07/67] added types for molbe BE object --- src/quemb/molbe/mbe.py | 77 ++++++++++++++++++++++-------------------- 1 file changed, 40 insertions(+), 37 deletions(-) diff --git a/src/quemb/molbe/mbe.py b/src/quemb/molbe/mbe.py index 5a35659f..1cc34934 100644 --- a/src/quemb/molbe/mbe.py +++ b/src/quemb/molbe/mbe.py @@ -12,6 +12,7 @@ from quemb.molbe._opt import BEOPT from quemb.molbe.be_parallel import be_func_parallel from quemb.molbe.eri_onthefly import integral_direct_DF +from quemb.molbe.fragment import fragpart from quemb.molbe.lo import MixinLocalize from quemb.molbe.misc import print_energy from quemb.molbe.pfrag import Frags @@ -21,7 +22,8 @@ get_be_error_jacobian as _ext_get_be_error_jacobian, ) from quemb.shared.helper import copy_docstring -from quemb.shared.typing import Matrix +from quemb.shared.manage_scratch import WorkDir +from quemb.shared.typing import Matrix, PathLike @define @@ -68,63 +70,64 @@ class BE(MixinLocalize): def __init__( self, - mf, - fobj, - eri_file="eri_file.h5", - lo_method="lowdin", - pop_method=None, - compute_hf=True, - restart=False, - save=False, - restart_file="storebe.pk", - save_file="storebe.pk", - hci_pt=False, - nproc=1, - ompnum=4, - scratch_dir=None, - hci_cutoff=0.001, - ci_coeff_cutoff=None, - select_cutoff=None, - integral_direct_DF=False, - auxbasis=None, - ): + mf: scf.hf.SCF, + fobj: fragpart, + eri_file: PathLike = "eri_file.h5", + lo_method: str = "lowdin", + pop_method: str | None = None, + compute_hf: bool = True, + restart: bool = False, + save: bool = False, + restart_file: PathLike = "storebe.pk", + save_file: PathLike = "storebe.pk", + hci_pt: bool = False, + nproc: int = 1, + ompnum: int = 4, + scratch_dir: WorkDir | None = None, + hci_cutoff: float = 0.001, + ci_coeff_cutoff: float | None = None, + select_cutoff: float | None = None, + integral_direct_DF: bool = False, + auxbasis: str | None = None, + ) -> None: """ Constructor for BE object. Parameters ---------- - mf : pyscf.scf.hf.SCF + mf : PySCF mean-field object. - fobj : quemb.molbe.fragment.fragpart + fobj : Fragment object containing sites, centers, edges, and indices. - eri_file : str, optional - Path to the file storing two-electron integrals, by default 'eri_file.h5'. - lo_method : str, optional + eri_file : + Path to the file storing two-electron integrals. + lo_method : Method for orbital localization, by default 'lowdin'. - compute_hf : bool, optional + compute_hf : Whether to compute Hartree-Fock energy, by default True. - restart : bool, optional + restart : Whether to restart from a previous calculation, by default False. - save : bool, optional + save : Whether to save intermediate objects for restart, by default False. - restart_file : str, optional + restart_file : Path to the file storing restart information, by default 'storebe.pk'. - save_file : str, optional + save_file : Path to the file storing save information, by default 'storebe.pk'. - nproc : int, optional + nproc : Number of processors for parallel calculations, by default 1. If set to >1, threaded parallel computation is invoked. - ompnum : int, optional + ompnum : Number of OpenMP threads, by default 4. - integral_direct_DF: bool, optional + integral_direct_DF: If mf._eri is None (i.e. ERIs are not saved in memory using incore_anyway), this flag is used to determine if the ERIs are computed integral-directly using density fitting; by default False. - auxbasis : str, optional + auxbasis : Auxiliary basis for density fitting, by default None (uses default auxiliary basis defined in PySCF). + solver_kwargs : + Keyword arguments to be passed on to the solver. """ - if restart: # Load previous calculation data from restart file with open(restart_file, "rb") as rfile: @@ -193,7 +196,7 @@ def __init__( self.cinv = None self.print_ini() - self.Fobjs = [] + self.Fobjs: list[Frags] = [] self.pot = initialize_pot(self.Nfrag, self.edge_idx) self.eri_file = eri_file self.scratch_dir = scratch_dir From 3492c66aab9dbf4f8f2f51f55d6c2e82b9c9f7f3 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Thu, 19 Dec 2024 16:20:47 -0500 Subject: [PATCH 08/67] fixed wrong close statements --- src/quemb/kbe/pbe.py | 2 -- src/quemb/molbe/mbe.py | 1 - 2 files changed, 3 deletions(-) diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py index 1a7d8bdc..5230df5d 100644 --- a/src/quemb/kbe/pbe.py +++ b/src/quemb/kbe/pbe.py @@ -105,7 +105,6 @@ def __init__( # Load previous calculation data from restart file with open(restart_file, "rb") as rfile: store_ = pickle.load(rfile) - rfile.close() self.Nocc = store_.Nocc self.hf_veff = store_.hf_veff self.hcore = store_.hcore @@ -323,7 +322,6 @@ def __init__( ) with open(save_file, "wb") as rfile: pickle.dump(store_, rfile, pickle.HIGHEST_PROTOCOL) - rfile.close() if not restart: self.initialize(compute_hf) diff --git a/src/quemb/molbe/mbe.py b/src/quemb/molbe/mbe.py index 1cc34934..51b6e57f 100644 --- a/src/quemb/molbe/mbe.py +++ b/src/quemb/molbe/mbe.py @@ -132,7 +132,6 @@ def __init__( # Load previous calculation data from restart file with open(restart_file, "rb") as rfile: store_ = pickle.load(rfile) - rfile.close() self.Nocc = store_.Nocc self.hf_veff = store_.hf_veff self.hcore = store_.hcore From 35adb0c370d49ba5d6c419f11081eee9f89df569 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Thu, 19 Dec 2024 16:28:03 -0500 Subject: [PATCH 09/67] fixed scratch-dir for molbe-be --- src/quemb/molbe/mbe.py | 21 +++++---------------- 1 file changed, 5 insertions(+), 16 deletions(-) diff --git a/src/quemb/molbe/mbe.py b/src/quemb/molbe/mbe.py index 51b6e57f..6c528c67 100644 --- a/src/quemb/molbe/mbe.py +++ b/src/quemb/molbe/mbe.py @@ -17,7 +17,6 @@ from quemb.molbe.misc import print_energy from quemb.molbe.pfrag import Frags from quemb.molbe.solver import be_func -from quemb.shared.config import settings from quemb.shared.external.optqn import ( get_be_error_jacobian as _ext_get_be_error_jacobian, ) @@ -197,24 +196,14 @@ def __init__( self.print_ini() self.Fobjs: list[Frags] = [] self.pot = initialize_pot(self.Nfrag, self.edge_idx) - self.eri_file = eri_file - self.scratch_dir = scratch_dir - # Set scratch directory - jobid = "" - if settings.CREATE_SCRATCH_DIR: - jobid = os.environ.get("SLURM_JOB_ID", "") - if settings.SCRATCH: - self.scratch_dir = settings.SCRATCH + str(jobid) - os.system("mkdir -p " + self.scratch_dir) - else: - self.scratch_dir = None - if not jobid: - self.eri_file = settings.SCRATCH + eri_file + if scratch_dir is None: + self.scratch_dir = WorkDir.from_environment() else: - self.eri_file = self.scratch_dir + "/" + eri_file + self.scratch_dir = scratch_dir + self.eri_file = self.scratch_dir / eri_file - self.frozen_core = False if not fobj.frozen_core else True + self.frozen_core = fobj.frozen_core self.ncore = 0 if not restart: self.E_core = 0 From b0d8817bfbebeae2a4a26ea4a604b051a4861b85 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Thu, 19 Dec 2024 16:32:53 -0500 Subject: [PATCH 10/67] added typing for molbe.BE.optimize - this fixed a likely bug in optimize. previously we had J0 = [[J0[-1, -1]]] now it is J0 = [[J0[-1][-1]]] since it is a list[list[float]] --- src/quemb/molbe/mbe.py | 55 +++++++++++++++++++++--------------------- 1 file changed, 27 insertions(+), 28 deletions(-) diff --git a/src/quemb/molbe/mbe.py b/src/quemb/molbe/mbe.py index 6c528c67..dda35ac1 100644 --- a/src/quemb/molbe/mbe.py +++ b/src/quemb/molbe/mbe.py @@ -22,7 +22,7 @@ ) from quemb.shared.helper import copy_docstring from quemb.shared.manage_scratch import WorkDir -from quemb.shared.typing import Matrix, PathLike +from quemb.shared.typing import KwargDict, Matrix, PathLike @define @@ -636,51 +636,50 @@ def compute_energy_full( def optimize( self, - solver="MP2", - method="QN", - only_chem=False, - conv_tol=1.0e-6, - relax_density=False, - J0=None, - nproc=1, - ompnum=4, - max_iter=500, - scratch_dir=None, - trust_region=False, - **solver_kwargs, - ): + solver: str = "MP2", + method: str = "QN", + only_chem: bool = False, + conv_tol: float = 1.0e-6, + relax_density: bool = False, + J0: list[list[float]] | None = None, + nproc: int = 1, + ompnum: int = 4, + max_iter: int = 500, + trust_region: bool = False, + solver_kwargs: KwargDict | None = None, + ) -> None: """BE optimization function Interfaces BEOPT to perform bootstrap embedding optimization. Parameters ---------- - solver : str, optional + solver : High-level solver for the fragment, by default 'MP2' - method : str, optional + method : Optimization method, by default 'QN' - only_chem : bool, optional + only_chem : If true, density matching is not performed -- only global chemical potential is optimized, by default False - conv_tol : float, optional + conv_tol : Convergence tolerance, by default 1.e-6 - relax_density : bool, optional + relax_density : Whether to use relaxed or unrelaxed densities, by default False This option is for using CCSD as solver. Relaxed density here uses Lambda amplitudes, whereas unrelaxed density only uses T amplitudes. c.f. See http://classic.chem.msu.su/cgi-bin/ceilidh.exe/gran/gamess/forum/?C34df668afbHW-7216-1405+00.htm for the distinction between the two - max_iter : int, optional + max_iter : Maximum number of optimization steps, by default 500 - nproc : int + nproc : Total number of processors assigned for the optimization. Defaults to 1. When nproc > 1, Python multithreading is invoked. - ompnum : int + ompnum : If nproc > 1, ompnum sets the number of cores for OpenMP parallelization. Defaults to 4 - J0 : list of list of float + J0 : Initial Jacobian. - trust_region : bool, optional + trust_region : Use trust-region based QN optimization, by default False """ # Check if only chemical potential optimization is required @@ -703,7 +702,7 @@ def optimize( hf_veff=self.hf_veff, nproc=nproc, ompnum=ompnum, - scratch_dir=scratch_dir, + scratch_dir=self.scratch_dir, max_space=max_iter, conv_tol=conv_tol, only_chem=only_chem, @@ -714,7 +713,7 @@ def optimize( hci_pt=self.hci_pt, solver=solver, ebe_hf=self.ebe_hf, - **solver_kwargs, + solver_kwargs=solver_kwargs, ) if method == "QN": @@ -722,7 +721,7 @@ def optimize( if only_chem: J0 = [[0.0]] J0 = self.get_be_error_jacobian(jac_solver="HF") - J0 = [[J0[-1, -1]]] + J0 = [[J0[-1][-1]]] else: J0 = self.get_be_error_jacobian(jac_solver="HF") @@ -737,7 +736,7 @@ def optimize( raise ValueError("This optimization method for BE is not supported") @copy_docstring(_ext_get_be_error_jacobian) - def get_be_error_jacobian(self, jac_solver="HF"): + def get_be_error_jacobian(self, jac_solver: str = "HF") -> list[list[float]]: return _ext_get_be_error_jacobian(self.Nfrag, self.Fobjs, jac_solver) def print_ini(self): From 1bb110572f7851315ccbd8df0eea10b96b7341da Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Thu, 19 Dec 2024 16:50:11 -0500 Subject: [PATCH 11/67] renamed _opt.py to opt.py finished BEOPT attrs refactoring --- mypy.ini | 2 +- src/quemb/kbe/pbe.py | 2 +- src/quemb/molbe/mbe.py | 2 +- src/quemb/molbe/{_opt.py => opt.py} | 99 +++++++++++++---------------- 4 files changed, 46 insertions(+), 59 deletions(-) rename src/quemb/molbe/{_opt.py => opt.py} (81%) diff --git a/mypy.ini b/mypy.ini index 80b2b678..9fb77851 100644 --- a/mypy.ini +++ b/mypy.ini @@ -6,7 +6,7 @@ # explicitly blacklist files, this means we can easily add in stricter type checks # by removing them from the blacklist -[mypy-quemb.molbe.ube,quemb.molbe.be_parallel,quemb.molbe.autofrag,quemb.molbe.eri_onthefly,quemb.molbe.fragment,quemb.molbe.helper,quemb.molbe.lchain,quemb.molbe.lo,quemb.molbe.mbe,quemb.molbe.misc,quemb.molbe._opt,quemb.molbe.pfrag,quemb.molbe.solver] +[mypy-quemb.molbe.ube,quemb.molbe.be_parallel,quemb.molbe.autofrag,quemb.molbe.eri_onthefly,quemb.molbe.fragment,quemb.molbe.helper,quemb.molbe.lchain,quemb.molbe.lo,quemb.molbe.mbe,quemb.molbe.misc,quemb.molbe.opt,quemb.molbe.pfrag,quemb.molbe.solver] disallow_untyped_defs = False check_untyped_defs = False diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py index 5230df5d..c34d458b 100644 --- a/src/quemb/kbe/pbe.py +++ b/src/quemb/kbe/pbe.py @@ -14,9 +14,9 @@ from quemb.kbe.lo import Mixin_k_Localize from quemb.kbe.misc import print_energy, storePBE from quemb.kbe.pfrag import Frags -from quemb.molbe._opt import BEOPT from quemb.molbe.be_parallel import be_func_parallel from quemb.molbe.helper import get_eri, get_scfObj, get_veff +from quemb.molbe.opt import BEOPT from quemb.molbe.solver import be_func from quemb.shared.config import settings from quemb.shared.external.optqn import ( diff --git a/src/quemb/molbe/mbe.py b/src/quemb/molbe/mbe.py index dda35ac1..a5e4ef1d 100644 --- a/src/quemb/molbe/mbe.py +++ b/src/quemb/molbe/mbe.py @@ -9,12 +9,12 @@ from numpy import float64 from pyscf import ao2mo, scf -from quemb.molbe._opt import BEOPT from quemb.molbe.be_parallel import be_func_parallel from quemb.molbe.eri_onthefly import integral_direct_DF from quemb.molbe.fragment import fragpart from quemb.molbe.lo import MixinLocalize from quemb.molbe.misc import print_energy +from quemb.molbe.opt import BEOPT from quemb.molbe.pfrag import Frags from quemb.molbe.solver import be_func from quemb.shared.external.optqn import ( diff --git a/src/quemb/molbe/_opt.py b/src/quemb/molbe/opt.py similarity index 81% rename from src/quemb/molbe/_opt.py rename to src/quemb/molbe/opt.py index 47c3df83..21763e53 100644 --- a/src/quemb/molbe/_opt.py +++ b/src/quemb/molbe/opt.py @@ -2,12 +2,18 @@ import numpy +from attrs import define +from numpy import float64 from quemb.molbe.be_parallel import be_func_parallel +from quemb.molbe.fragment import fragpart from quemb.molbe.solver import be_func from quemb.shared.external.optqn import FrankQN +from quemb.shared.manage_scratch import WorkDir +from quemb.shared.typing import KwargDict, Matrix +@define class BEOPT: """Perform BE optimization. @@ -18,84 +24,65 @@ class BEOPT: Parameters ---------- - pot : list + pot : List of initial BE potentials. The last element is for the global chemical potential. Fobjs : quemb.molbe.fragment.fragpart Fragment object - Nocc : int + Nocc : No. of occupied orbitals for the full system. - enuc : float + enuc : Nuclear component of the energy. - solver : str + solver : High-level solver in bootstrap embedding. 'MP2', 'CCSD', 'FCI' are supported. Selected CI versions, 'HCI', 'SHCI', & 'SCI' are also supported. Defaults to 'MP2' - only_chem : bool + only_chem : Whether to perform chemical potential optimization only. Refer to bootstrap embedding literatures. - nproc : int + nproc : Total number of processors assigned for the optimization. Defaults to 1. When nproc > 1, Python multithreading is invoked. - ompnum : int + ompnum : If nproc > 1, ompnum sets the number of cores for OpenMP parallelization. Defaults to 4 - max_space : int + max_space : Maximum number of bootstrap optimizaiton steps, after which the optimization is called converged. - conv_tol : float + conv_tol : Convergence criteria for optimization. Defaults to 1e-6 - ebe_hf : float + ebe_hf : Hartree-Fock energy. Defaults to 0.0 """ - def __init__( - self, - pot, - Fobjs, - Nocc, - enuc, - solver="MP2", - nproc=1, - ompnum=4, - only_chem=False, - hf_veff=None, - hci_pt=False, - hci_cutoff=0.001, - ci_coeff_cutoff=None, - select_cutoff=None, - max_space=500, - conv_tol=1.0e-6, - relax_density=False, - ebe_hf=0.0, - scratch_dir=None, - **solver_kwargs, - ): - # Initialize instance attributes - self.ebe_hf = ebe_hf - self.hf_veff = hf_veff - self.pot = pot - self.Fobjs = Fobjs - self.Nocc = Nocc - self.enuc = enuc - self.solver = solver - self.iter = 0 - self.err = 0.0 - self.Ebe = 0.0 - self.max_space = max_space - self.nproc = nproc - self.ompnum = ompnum - self.only_chem = only_chem - self.conv_tol = conv_tol - self.relax_density = relax_density - # HCI parameters - self.hci_cutoff = hci_cutoff - self.ci_coeff_cutoff = ci_coeff_cutoff - self.select_cutoff = select_cutoff - self.hci_pt = hci_pt - self.solver_kwargs = solver_kwargs - self.scratch_dir = scratch_dir + pot: list[float] + Fobjs: fragpart + Nocc: int + enuc: float + scratch_dir: WorkDir + solver: str = "MP2" + nproc: int = 1 + ompnum: int = 4 + only_chem: bool = False + hf_veff: Matrix[float64] | None = None + + max_space: int = 500 + conv_tol: float = 1.0e-6 + relax_density: bool = False + ebe_hf: float = 0.0 + + iter: int = 0 + err: float = 0.0 + Ebe: float = 0.0 + + solver_kwargs: KwargDict | None = None + + # HCI parameters + hci_cutoff: float = 0.001 + ci_coeff_cutoff: float | None = None + select_cutoff: float | None = None + hci_pt: bool = False def objfunc(self, xk): """ From b46479b766400b9de528a49135ff1933fb99fba7 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Thu, 19 Dec 2024 17:50:45 -0500 Subject: [PATCH 12/67] pass on solver_kwargs to be_func --- src/quemb/molbe/opt.py | 21 ++++++++++----------- src/quemb/molbe/solver.py | 2 +- 2 files changed, 11 insertions(+), 12 deletions(-) diff --git a/src/quemb/molbe/opt.py b/src/quemb/molbe/opt.py index 21763e53..6c54ada6 100644 --- a/src/quemb/molbe/opt.py +++ b/src/quemb/molbe/opt.py @@ -3,14 +3,14 @@ import numpy from attrs import define -from numpy import float64 +from numpy import array, float64 from quemb.molbe.be_parallel import be_func_parallel -from quemb.molbe.fragment import fragpart +from quemb.molbe.pfrag import Frags from quemb.molbe.solver import be_func from quemb.shared.external.optqn import FrankQN from quemb.shared.manage_scratch import WorkDir -from quemb.shared.typing import KwargDict, Matrix +from quemb.shared.typing import KwargDict, Matrix, Vector @define @@ -27,7 +27,7 @@ class BEOPT: pot : List of initial BE potentials. The last element is for the global chemical potential. - Fobjs : quemb.molbe.fragment.fragpart + Fobjs : Fragment object Nocc : No. of occupied orbitals for the full system. @@ -57,7 +57,7 @@ class BEOPT: """ pot: list[float] - Fobjs: fragpart + Fobjs: list[Frags] Nocc: int enuc: float scratch_dir: WorkDir @@ -74,7 +74,7 @@ class BEOPT: iter: int = 0 err: float = 0.0 - Ebe: float = 0.0 + Ebe: Matrix[float64] = array([[0.0]]) solver_kwargs: KwargDict | None = None @@ -84,7 +84,7 @@ class BEOPT: select_cutoff: float | None = None hci_pt: bool = False - def objfunc(self, xk): + def objfunc(self, xk: list[float]) -> Vector[float64]: """ Computes error vectors, RMS error, and BE energies. @@ -93,7 +93,7 @@ def objfunc(self, xk): Parameters ---------- - xk : list + xk : Current potentials in the BE optimization. Returns @@ -120,9 +120,8 @@ def objfunc(self, xk): ci_coeff_cutoff=self.ci_coeff_cutoff, select_cutoff=self.select_cutoff, hci_pt=self.hci_pt, - ebe_hf=self.ebe_hf, scratch_dir=self.scratch_dir, - **self.solver_kwargs, + solver_kwargs=self.solver_kwargs, ) else: err_, errvec_, ebe_ = be_func_parallel( @@ -143,7 +142,7 @@ def objfunc(self, xk): select_cutoff=self.select_cutoff, ebe_hf=self.ebe_hf, scratch_dir=self.scratch_dir, - **self.solver_kwargs, + solver_kwargs=self.solver_kwargs, ) # Update error and BE energy diff --git a/src/quemb/molbe/solver.py b/src/quemb/molbe/solver.py index e1c22b87..e36fac9f 100644 --- a/src/quemb/molbe/solver.py +++ b/src/quemb/molbe/solver.py @@ -26,6 +26,7 @@ def be_func( Nocc, solver, enuc, # noqa: ARG001 + solver_kwargs, hf_veff=None, only_chem=False, nproc=4, @@ -40,7 +41,6 @@ def be_func( return_vec=False, use_cumulant=True, scratch_dir=None, - **solver_kwargs, ): """ Perform bootstrap embedding calculations for each fragment. From bcabc40db2f146843165ac2705b3b76466bd7be1 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Thu, 19 Dec 2024 17:57:40 -0500 Subject: [PATCH 13/67] added types to be_func --- src/quemb/molbe/solver.py | 74 +++++++++++++++++++++------------------ 1 file changed, 39 insertions(+), 35 deletions(-) diff --git a/src/quemb/molbe/solver.py b/src/quemb/molbe/solver.py index e36fac9f..5cec6f87 100644 --- a/src/quemb/molbe/solver.py +++ b/src/quemb/molbe/solver.py @@ -3,11 +3,13 @@ import os import numpy +from numpy import float64 from numpy.linalg import multi_dot from pyscf import ao2mo, cc, fci, mcscf, mp from pyscf.cc.ccsd_rdm import make_rdm2 from quemb.molbe.helper import get_frag_energy, get_frag_energy_u +from quemb.molbe.pfrag import Frags from quemb.shared.config import settings from quemb.shared.external.ccsd_rdm import ( make_rdm1_ccsd_t1, @@ -18,29 +20,31 @@ from quemb.shared.external.uccsd_eri import make_eris_incore from quemb.shared.external.unrestricted_utils import make_uhf_obj from quemb.shared.helper import unused +from quemb.shared.manage_scratch import WorkDir +from quemb.shared.typing import KwargDict, Matrix def be_func( - pot, - Fobjs, - Nocc, - solver, - enuc, # noqa: ARG001 - solver_kwargs, - hf_veff=None, - only_chem=False, - nproc=4, - hci_pt=False, - hci_cutoff=0.001, - ci_coeff_cutoff=None, - select_cutoff=None, - eeval=False, - ereturn=False, - frag_energy=False, - relax_density=False, - return_vec=False, - use_cumulant=True, - scratch_dir=None, + pot: list[float], + Fobjs: list[Frags], + Nocc: int, + solver: str, + enuc: float, # noqa: ARG001 + solver_kwargs: KwargDict | None, + scratch_dir: WorkDir, + hf_veff: Matrix[float64] | None = None, + only_chem: bool = False, + nproc: int = 4, + hci_pt: bool = False, + hci_cutoff: float = 0.001, + ci_coeff_cutoff: float | None = None, + select_cutoff: float | None = None, + eeval: bool = False, + ereturn: bool = False, + frag_energy: bool = False, + relax_density: bool = False, + return_vec: bool = False, + use_cumulant: bool = True, ): """ Perform bootstrap embedding calculations for each fragment. @@ -50,35 +54,35 @@ def be_func( Parameters ---------- - pot : list + pot : List of potentials. Fobjs : list of quemb.molbe.fragment.fragpart List of fragment objects. - Nocc : int + Nocc : Number of occupied orbitals. - solver : str + solver : Quantum chemistry solver to use ('MP2', 'CCSD', 'FCI', 'HCI', 'SHCI', 'SCI'). - enuc : float + enuc : f Nuclear energy. - hf_veff : numpy.ndarray, optional + hf_veff : Hartree-Fock effective potential. Defaults to None. - only_chem : bool, optional + only_chem : Whether to only optimize the chemical potential. Defaults to False. - nproc : int, optional + nproc : Number of processors. Defaults to 4. This is only neccessary for 'SHCI' solver - eeval : bool, optional + eeval : Whether to evaluate the energy. Defaults to False. - ereturn : bool, optional + ereturn : Whether to return the energy. Defaults to False. - frag_energy : bool, optional + frag_energy : Whether to calculate fragment energy. Defaults to False. - relax_density : bool, optional + relax_density : Whether to relax the density. Defaults to False. - return_vec : bool, optional + return_vec : Whether to return the error vector. Defaults to False. - ebe_hf : float, optional + ebe_hf : Hartree-Fock energy. Defaults to 0. - use_cumulant : bool, optional + use_cumulant : Whether to use the cumulant-based energy expression. Defaults to True. Returns @@ -217,7 +221,7 @@ def be_func( rdm1_tmp, rdm2s = ci.make_rdm12(0, nmo, nelec) elif solver in ["block2", "DMRG", "DMRGCI", "DMRGSCF"]: - solver_kwargs_ = solver_kwargs.copy() + solver_kwargs_ = {} if solver_kwargs is None else solver_kwargs.copy() if scratch_dir is None and settings.CREATE_SCRATCH_DIR: tmp = os.path.join(settings.SCRATCH, str(os.getpid()), str(fobj.dname)) else: From 0517d470a6b64bba74e27fd15d9f747d951a99d2 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Thu, 19 Dec 2024 18:04:09 -0500 Subject: [PATCH 14/67] added delete multiple_files function --- src/quemb/shared/helper.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/quemb/shared/helper.py b/src/quemb/shared/helper.py index 1b1a4cd3..412fc5b0 100644 --- a/src/quemb/shared/helper.py +++ b/src/quemb/shared/helper.py @@ -1,5 +1,7 @@ +from collections.abc import Iterable from inspect import signature from itertools import islice +from pathlib import Path from typing import Any, Callable, TypeVar Function = TypeVar("Function", bound=Callable) @@ -88,3 +90,9 @@ def ncore_(z: int) -> int: else: raise ValueError("Ncore not computed in helper.ncore(), add it yourself!") return nc + + +def delete_multiple_files(*args: Iterable[Path]) -> None: + for files in args: + for file in files: + file.unlink() From a4583308a848a9970aee526622d902dfeaf8b6f6 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Thu, 19 Dec 2024 18:04:30 -0500 Subject: [PATCH 15/67] use the new scratch dir in be_func --- src/quemb/molbe/solver.py | 30 ++++++++++-------------------- 1 file changed, 10 insertions(+), 20 deletions(-) diff --git a/src/quemb/molbe/solver.py b/src/quemb/molbe/solver.py index 5cec6f87..813e7165 100644 --- a/src/quemb/molbe/solver.py +++ b/src/quemb/molbe/solver.py @@ -10,7 +10,6 @@ from quemb.molbe.helper import get_frag_energy, get_frag_energy_u from quemb.molbe.pfrag import Frags -from quemb.shared.config import settings from quemb.shared.external.ccsd_rdm import ( make_rdm1_ccsd_t1, make_rdm1_uccsd, @@ -19,7 +18,7 @@ ) from quemb.shared.external.uccsd_eri import make_eris_incore from quemb.shared.external.unrestricted_utils import make_uhf_obj -from quemb.shared.helper import unused +from quemb.shared.helper import delete_multiple_files, unused from quemb.shared.manage_scratch import WorkDir from quemb.shared.typing import KwargDict, Matrix @@ -169,14 +168,8 @@ def be_func( # pylint: disable-next=E0611,E0401 from pyscf.shciscf import shci # noqa: PLC0415 # shci is optional - if scratch_dir is None and settings.CREATE_SCRATCH_DIR: - tmp = os.path.join(settings.SCRATCH, str(os.getpid()), str(fobj.dname)) - elif scratch_dir is None: - tmp = settings.SCRATCH - else: - tmp = os.path.join(scratch_dir, str(os.getpid()), str(fobj.dname)) - if not os.path.isdir(tmp): - os.system("mkdir -p " + tmp) + frag_scratch = WorkDir(scratch_dir / fobj.dname) + nmo = fobj._mf.mo_coeff.shape[1] nelec = (fobj.nsocc, fobj.nsocc) @@ -222,24 +215,21 @@ def be_func( elif solver in ["block2", "DMRG", "DMRGCI", "DMRGSCF"]: solver_kwargs_ = {} if solver_kwargs is None else solver_kwargs.copy() - if scratch_dir is None and settings.CREATE_SCRATCH_DIR: - tmp = os.path.join(settings.SCRATCH, str(os.getpid()), str(fobj.dname)) - else: - tmp = os.path.join(scratch_dir, str(os.getpid()), str(fobj.dname)) - if not os.path.isdir(tmp): - os.system("mkdir -p " + tmp) + frag_scratch = WorkDir(scratch_dir / fobj.dname) try: rdm1_tmp, rdm2s = solve_block2( - fobj._mf, fobj.nsocc, frag_scratch=tmp, **solver_kwargs_ + fobj._mf, fobj.nsocc, frag_scratch=frag_scratch, **solver_kwargs_ ) except Exception as inst: raise inst finally: if solver_kwargs_.pop("force_cleanup", False): - os.system("rm -r " + os.path.join(tmp, "F.*")) - os.system("rm -r " + os.path.join(tmp, "FCIDUMP*")) - os.system("rm -r " + os.path.join(tmp, "node*")) + delete_multiple_files( + frag_scratch.path.glob("F.*"), + frag_scratch.path.glob("FCIDUMP*"), + frag_scratch.path.glob("node*"), + ) else: raise ValueError("Solver not implemented") From 5c384c0da8d32fd16deccdd1a9e07a7efe6e898f Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Thu, 19 Dec 2024 18:11:21 -0500 Subject: [PATCH 16/67] added types to molbe.BE.optimize + call self.scratch_dir.cleanup call self.scratch_dir.cleanup in both `optimize` and `oneshot` --- src/quemb/molbe/mbe.py | 39 +++++++++++++-------------------------- src/quemb/molbe/solver.py | 2 +- 2 files changed, 14 insertions(+), 27 deletions(-) diff --git a/src/quemb/molbe/mbe.py b/src/quemb/molbe/mbe.py index a5e4ef1d..a055dae0 100644 --- a/src/quemb/molbe/mbe.py +++ b/src/quemb/molbe/mbe.py @@ -1,6 +1,5 @@ # Author(s): Oinam Romesh Meitei -import os import pickle import h5py @@ -124,8 +123,6 @@ def __init__( auxbasis : Auxiliary basis for density fitting, by default None (uses default auxiliary basis defined in PySCF). - solver_kwargs : - Keyword arguments to be passed on to the solver. """ if restart: # Load previous calculation data from restart file @@ -735,6 +732,8 @@ def optimize( else: raise ValueError("This optimization method for BE is not supported") + self.scratch_dir.cleanup() + @copy_docstring(_ext_get_be_error_jacobian) def get_be_error_jacobian(self, jac_solver: str = "HF") -> list[list[float]]: return _ext_get_be_error_jacobian(self.Nfrag, self.Fobjs, jac_solver) @@ -899,35 +898,28 @@ def initialize(self, eri_, compute_hf, restart=False): def oneshot( self, - solver="MP2", - nproc=1, - ompnum=4, - calc_frag_energy=False, - clean_eri=False, - scratch_dir=None, - **solver_kwargs, + solver: str = "MP2", + nproc: int = 1, + ompnum: int = 4, + calc_frag_energy: bool = False, + solver_kwargs: KwargDict | None = None, ): """ Perform a one-shot bootstrap embedding calculation. Parameters ---------- - solver : str, optional + solver : High-level quantum chemistry method, by default 'MP2'. 'CCSD', 'FCI', and variants of selected CI are supported. - nproc : int, optional + nproc : Number of processors for parallel calculations, by default 1. If set to >1, multi-threaded parallel computation is invoked. - ompnum : int, optional + ompnum : Number of OpenMP threads, by default 4. - calc_frag_energy : bool, optional + calc_frag_energy : Whether to calculate fragment energies, by default False. - clean_eri : bool, optional - Whether to clean up ERI files after calculation, by default False. """ - self.scratch_dir = scratch_dir - self.solver_kwargs = solver_kwargs - print("Calculating Energy by Fragment? ", calc_frag_energy) if nproc == 1: rets = be_func( @@ -945,7 +937,7 @@ def oneshot( ereturn=True, eeval=True, scratch_dir=self.scratch_dir, - **self.solver_kwargs, + solver_kwargs=solver_kwargs, ) else: rets = be_func_parallel( @@ -993,12 +985,7 @@ def oneshot( if not calc_frag_energy: self.compute_energy_full(approx_cumulant=True, return_rdm=False) - if clean_eri: - try: - os.remove(self.eri_file) - os.rmdir(self.scratch_dir) - except (FileNotFoundError, TypeError): - print("Scratch directory not removed") + self.scratch_dir.cleanup() def update_fock(self, heff=None): """ diff --git a/src/quemb/molbe/solver.py b/src/quemb/molbe/solver.py index 813e7165..a86b51ac 100644 --- a/src/quemb/molbe/solver.py +++ b/src/quemb/molbe/solver.py @@ -24,7 +24,7 @@ def be_func( - pot: list[float], + pot: list[float] | None, Fobjs: list[Frags], Nocc: int, solver: str, From f4cb9fd0a927d45759387a74c17e012492d6dce2 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Thu, 19 Dec 2024 18:15:04 -0500 Subject: [PATCH 17/67] moved schmidt_decomposition to prevent circular import makes much more sense like this --- src/quemb/molbe/pfrag.py | 97 ++++++++++++++++++++++++++++++++++++++- src/quemb/molbe/solver.py | 96 -------------------------------------- 2 files changed, 96 insertions(+), 97 deletions(-) diff --git a/src/quemb/molbe/pfrag.py b/src/quemb/molbe/pfrag.py index 456a739a..d49fd732 100644 --- a/src/quemb/molbe/pfrag.py +++ b/src/quemb/molbe/pfrag.py @@ -6,7 +6,6 @@ from numpy.linalg import multi_dot from quemb.molbe.helper import get_eri, get_scfObj, get_veff -from quemb.molbe.solver import schmidt_decomposition class Frags: @@ -385,3 +384,99 @@ def update_ebe_hf( return (e_h1, e_coul, e1 + e2 + ec) else: return None + + +def schmidt_decomposition( + mo_coeff, nocc, Frag_sites, cinv=None, rdm=None, norb=None, return_orb_count=False +): + """ + Perform Schmidt decomposition on the molecular orbital coefficients. + + This function decomposes the molecular orbitals into fragment and environment parts + using the Schmidt decomposition method. It computes the transformation matrix (TA) + which includes both the fragment orbitals and the entangled bath. + + Parameters + ---------- + mo_coeff : numpy.ndarray + Molecular orbital coefficients. + nocc : int + Number of occupied orbitals. + Frag_sites : list of int + List of fragment sites (indices). + cinv : numpy.ndarray, optional + Inverse of the transformation matrix. Defaults to None. + rdm : numpy.ndarray, optional + Reduced density matrix. If not provided, it will be computed from the molecular + orbitals. Defaults to None. + norb : int, optional + Specifies number of bath orbitals. Used for UBE to make alpha and beta + spaces the same size. Defaults to None + return_orb_count : bool, optional + Return more information about the number of orbitals. Used in UBE. + Defaults to False + + Returns + ------- + numpy.ndarray + Transformation matrix (TA) including both fragment and entangled bath orbitals. + if return_orb_count: + numpy.ndarray, int, int + returns TA (above), number of orbitals in the fragment space, and number of + orbitals in bath space + """ + # Threshold for eigenvalue significance + thres = 1.0e-10 + + # Compute the reduced density matrix (RDM) if not provided + if mo_coeff is not None: + C = mo_coeff[:, :nocc] + if rdm is None: + Dhf = numpy.dot(C, C.T) + if cinv is not None: + Dhf = multi_dot((cinv, Dhf, cinv.conj().T)) + else: + Dhf = rdm + + # Total number of sites + Tot_sites = Dhf.shape[0] + + # Identify environment sites (indices not in Frag_sites) + Env_sites1 = numpy.array([i for i in range(Tot_sites) if i not in Frag_sites]) + Env_sites = numpy.array([[i] for i in range(Tot_sites) if i not in Frag_sites]) + Frag_sites1 = numpy.array([[i] for i in Frag_sites]) + + # Compute the environment part of the density matrix + Denv = Dhf[Env_sites, Env_sites.T] + + # Perform eigenvalue decomposition on the environment density matrix + Eval, Evec = numpy.linalg.eigh(Denv) + + # Identify significant environment orbitals based on eigenvalue threshold + Bidx = [] + + # Set the number of orbitals to be taken from the environment orbitals + # Based on an eigenvalue threshold ordering + if norb is not None: + n_frag_ind = len(Frag_sites1) + n_bath_ind = norb - n_frag_ind + ind_sort = numpy.argsort(numpy.abs(Eval)) + first_el = [x for x in ind_sort if x < 1.0 - thres][-1 * n_bath_ind] + for i in range(len(Eval)): + if numpy.abs(Eval[i]) >= first_el: + Bidx.append(i) + else: + for i in range(len(Eval)): + if thres < numpy.abs(Eval[i]) < 1.0 - thres: + Bidx.append(i) + + # Initialize the transformation matrix (TA) + TA = numpy.zeros([Tot_sites, len(Frag_sites) + len(Bidx)]) + TA[Frag_sites, : len(Frag_sites)] = numpy.eye(len(Frag_sites)) # Fragment part + TA[Env_sites1, len(Frag_sites) :] = Evec[:, Bidx] # Environment part + + if return_orb_count: + # return TA, norbs_frag, norbs_bath + return TA, Frag_sites1.shape[0], len(Bidx) + else: + return TA diff --git a/src/quemb/molbe/solver.py b/src/quemb/molbe/solver.py index a86b51ac..69938fba 100644 --- a/src/quemb/molbe/solver.py +++ b/src/quemb/molbe/solver.py @@ -978,99 +978,3 @@ def ao2mofn(moish): return (ucc, rdm1, rdm2) return (ucc, rdm1, None) return ucc - - -def schmidt_decomposition( - mo_coeff, nocc, Frag_sites, cinv=None, rdm=None, norb=None, return_orb_count=False -): - """ - Perform Schmidt decomposition on the molecular orbital coefficients. - - This function decomposes the molecular orbitals into fragment and environment parts - using the Schmidt decomposition method. It computes the transformation matrix (TA) - which includes both the fragment orbitals and the entangled bath. - - Parameters - ---------- - mo_coeff : numpy.ndarray - Molecular orbital coefficients. - nocc : int - Number of occupied orbitals. - Frag_sites : list of int - List of fragment sites (indices). - cinv : numpy.ndarray, optional - Inverse of the transformation matrix. Defaults to None. - rdm : numpy.ndarray, optional - Reduced density matrix. If not provided, it will be computed from the molecular - orbitals. Defaults to None. - norb : int, optional - Specifies number of bath orbitals. Used for UBE to make alpha and beta - spaces the same size. Defaults to None - return_orb_count : bool, optional - Return more information about the number of orbitals. Used in UBE. - Defaults to False - - Returns - ------- - numpy.ndarray - Transformation matrix (TA) including both fragment and entangled bath orbitals. - if return_orb_count: - numpy.ndarray, int, int - returns TA (above), number of orbitals in the fragment space, and number of - orbitals in bath space - """ - # Threshold for eigenvalue significance - thres = 1.0e-10 - - # Compute the reduced density matrix (RDM) if not provided - if mo_coeff is not None: - C = mo_coeff[:, :nocc] - if rdm is None: - Dhf = numpy.dot(C, C.T) - if cinv is not None: - Dhf = multi_dot((cinv, Dhf, cinv.conj().T)) - else: - Dhf = rdm - - # Total number of sites - Tot_sites = Dhf.shape[0] - - # Identify environment sites (indices not in Frag_sites) - Env_sites1 = numpy.array([i for i in range(Tot_sites) if i not in Frag_sites]) - Env_sites = numpy.array([[i] for i in range(Tot_sites) if i not in Frag_sites]) - Frag_sites1 = numpy.array([[i] for i in Frag_sites]) - - # Compute the environment part of the density matrix - Denv = Dhf[Env_sites, Env_sites.T] - - # Perform eigenvalue decomposition on the environment density matrix - Eval, Evec = numpy.linalg.eigh(Denv) - - # Identify significant environment orbitals based on eigenvalue threshold - Bidx = [] - - # Set the number of orbitals to be taken from the environment orbitals - # Based on an eigenvalue threshold ordering - if norb is not None: - n_frag_ind = len(Frag_sites1) - n_bath_ind = norb - n_frag_ind - ind_sort = numpy.argsort(numpy.abs(Eval)) - first_el = [x for x in ind_sort if x < 1.0 - thres][-1 * n_bath_ind] - for i in range(len(Eval)): - if numpy.abs(Eval[i]) >= first_el: - Bidx.append(i) - else: - for i in range(len(Eval)): - if thres < numpy.abs(Eval[i]) < 1.0 - thres: - Bidx.append(i) - - # Initialize the transformation matrix (TA) - TA = numpy.zeros([Tot_sites, len(Frag_sites) + len(Bidx)]) - TA[Frag_sites, : len(Frag_sites)] = numpy.eye(len(Frag_sites)) # Fragment part - TA[Env_sites1, len(Frag_sites) :] = Evec[:, Bidx] # Environment part - - if return_orb_count: - # return TA, norbs_frag, norbs_bath - return TA, Frag_sites1.shape[0], len(Bidx) - else: - return TA From 83e6773d3cfcbfcefb38d5a54bfb6061a816db7a Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Thu, 19 Dec 2024 18:28:06 -0500 Subject: [PATCH 18/67] added type hints to be_func_parallel --- src/quemb/molbe/be_parallel.py | 71 ++++++++++++++++++---------------- src/quemb/molbe/mbe.py | 2 +- 2 files changed, 38 insertions(+), 35 deletions(-) diff --git a/src/quemb/molbe/be_parallel.py b/src/quemb/molbe/be_parallel.py index b7b935ab..1ff0fb7a 100644 --- a/src/quemb/molbe/be_parallel.py +++ b/src/quemb/molbe/be_parallel.py @@ -4,6 +4,7 @@ from multiprocessing import Pool import numpy +from numpy import float64 from numpy.linalg import multi_dot from pyscf import ao2mo, fci, mcscf @@ -13,6 +14,7 @@ get_frag_energy_u, get_scfObj, ) +from quemb.molbe.pfrag import Frags from quemb.molbe.solver import ( make_rdm1_ccsd_t1, make_rdm2_urlx, @@ -24,6 +26,7 @@ from quemb.shared.external.ccsd_rdm import make_rdm1_uccsd, make_rdm2_uccsd from quemb.shared.external.unrestricted_utils import make_uhf_obj from quemb.shared.helper import unused +from quemb.shared.typing import Matrix def run_solver( @@ -389,24 +392,24 @@ def run_solver_u( def be_func_parallel( - pot, - Fobjs, - Nocc, - solver, - enuc, # noqa: ARG001 - hf_veff=None, - nproc=1, - ompnum=4, - only_chem=False, - relax_density=False, - use_cumulant=True, - eeval=False, - frag_energy=False, - hci_cutoff=0.001, - ci_coeff_cutoff=None, - select_cutoff=None, - return_vec=False, - writeh1=False, + pot: list[float] | None, + Fobjs: list[Frags], + Nocc: int, + solver: str, + enuc: float, # noqa: ARG001 + hf_veff: Matrix[float64] | None = None, + nproc: int = 1, + ompnum: int = 4, + only_chem: bool = False, + relax_density: bool = False, + use_cumulant: bool = True, + eeval: bool = False, + frag_energy: bool = False, + hci_cutoff: float = 0.001, + ci_coeff_cutoff: float | None = None, + select_cutoff: float | None = None, + return_vec: bool = False, + writeh1: bool = False, ): """ Embarrassingly Parallel High-Level Computation @@ -418,38 +421,36 @@ def be_func_parallel( Parameters ---------- - pot : list of float + pot : Potentials (local & global) that are added to the 1-electron Hamiltonian component. The last element in the list is the chemical potential. - Fobjs : list of quemb.molbe.fragment.fragpart + Fobjs : Fragment definitions. - Nocc : int + Nocc : Number of occupied orbitals for the full system. - solver : str + solver : High-level solver in bootstrap embedding. Supported values are 'MP2', 'CCSD', 'FCI', 'HCI', 'SHCI', and 'SCI'. - enuc : float + enuc : Nuclear component of the energy. - hf_veff : numpy.ndarray, optional + hf_veff : Hartree-Fock effective potential. - nproc : int, optional + nproc : Total number of processors assigned for the optimization. Defaults to 1. When nproc > 1, Python multithreading is invoked. - ompnum : int, optional + ompnum : If nproc > 1, sets the number of cores for OpenMP parallelization. Defaults to 4. - only_chem : bool, optional + only_chem : Whether to perform chemical potential optimization only. Refer to bootstrap embedding literature. Defaults to False. - eeval : bool, optional + eeval : Whether to evaluate energies. Defaults to False. - frag_energy : bool, optional + frag_energy : Whether to compute fragment energy. Defaults to False. - return_vec : bool, optional + return_vec : Whether to return the error vector. Defaults to False. - ebe_hf : float, optional - Hartree-Fock energy. Defaults to 0. - writeh1 : bool, optional + writeh1 : Whether to write the one-electron integrals. Defaults to False. Returns @@ -523,7 +524,9 @@ def be_func_parallel( results.append(result) # Collect results - [rdms.append(result.get()) for result in results] + for result in results: + rdms.append(result.get()) + pool_.close() if frag_energy: diff --git a/src/quemb/molbe/mbe.py b/src/quemb/molbe/mbe.py index a055dae0..4d86dc06 100644 --- a/src/quemb/molbe/mbe.py +++ b/src/quemb/molbe/mbe.py @@ -956,7 +956,7 @@ def oneshot( nproc=nproc, ompnum=ompnum, scratch_dir=self.scratch_dir, - **self.solver_kwargs, + solver_kwargs=solver_kwargs, ) print("-----------------------------------------------------", flush=True) From e24a5e3a64961a1d4a336d9e529a2b4dbfa81ec2 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 11:17:22 -0500 Subject: [PATCH 19/67] fixed typo --- src/quemb/molbe/solver.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/quemb/molbe/solver.py b/src/quemb/molbe/solver.py index 69938fba..e80826c0 100644 --- a/src/quemb/molbe/solver.py +++ b/src/quemb/molbe/solver.py @@ -61,7 +61,7 @@ def be_func( Number of occupied orbitals. solver : Quantum chemistry solver to use ('MP2', 'CCSD', 'FCI', 'HCI', 'SHCI', 'SCI'). - enuc : f + enuc : Nuclear energy. hf_veff : Hartree-Fock effective potential. Defaults to None. @@ -326,8 +326,7 @@ def be_func_u( use_cumulant=True, frozen=False, ): - """ - Perform bootstrap embedding calculations for each fragment with UCCSD. + """Perform bootstrap embedding calculations for each fragment with UCCSD. This function computes the energy and/or error for each fragment in a molecular system using various quantum chemistry solvers. From b5ea348f9e5dc9e4b39ebe8769985735860fa281 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 11:31:45 -0500 Subject: [PATCH 20/67] fixed several small errors in the code --- example/molbe_dmrg_block2.py | 42 +++++++++++++-------------- example/molbe_hexene_oneshot_uccsd.py | 2 +- src/quemb/kbe/pbe.py | 4 +-- src/quemb/molbe/misc.py | 4 +-- src/quemb/molbe/ube.py | 11 ++----- tests/ube-oneshot_test.py | 2 +- 6 files changed, 26 insertions(+), 39 deletions(-) diff --git a/example/molbe_dmrg_block2.py b/example/molbe_dmrg_block2.py index 8660ead2..f59f6473 100644 --- a/example/molbe_dmrg_block2.py +++ b/example/molbe_dmrg_block2.py @@ -2,8 +2,6 @@ # `block2` is a DMRG and sparse tensor network library developed by the # Garnet-Chan group at Caltech: https://block2.readthedocs.io/en/latest/index.html -import os - import matplotlib.pyplot as plt import numpy from pyscf import cc, fci, gto, scf @@ -15,9 +13,6 @@ seps = numpy.linspace(0.60, 1.6, num=num_points) fci_ecorr, ccsd_ecorr, ccsdt_ecorr, bedmrg_ecorr = [], [], [], [] -# Specify a scratch directory for fragment DMRG files: -scratch = os.getcwd() - for a in seps: # Hartree-Fock serves as the starting point for all BE calculations: mol = gto.M() @@ -57,9 +52,10 @@ # Next, run BE-DMRG with default parameters and maxM=100. mybe.oneshot( solver="block2", # or 'DMRG', 'DMRGSCF', 'DMRGCI' - scratch_dir=scratch, # Scratch dir for fragment DMRG - maxM=100, # Max fragment bond dimension - force_cleanup=True, # Remove all fragment DMRG tmpfiles + solver_kwargs=dict( + maxM=100, # Max fragment bond dimension + force_cleanup=True, # Remove all fragment DMRG tmpfiles + ), ) bedmrg_ecorr.append(mybe.ebe_tot - mf.e_tot) @@ -79,7 +75,7 @@ ax.plot(seps, bedmrg_ecorr, "o-", linewidth=1, label="BE1-DMRG") ax.legend() -plt.savefig(os.path.join(scratch, f"BEDMRG_H8_PES{num_points}.png")) +plt.savefig(f"BEDMRG_H8_PES{num_points}.png") # (See ../quemb/example/figures/BEDMRG_H8_PES20.png for an example.) @@ -102,17 +98,18 @@ mybe.optimize( solver="block2", # or 'DMRG', 'DMRGSCF', 'DMRGCI' - scratch=scratch, # Scratch dir for fragment DMRG - startM=20, # Initial fragment bond dimension (1st sweep) - maxM=200, # Maximum fragment bond dimension max_iter=60, # Max number of sweeps - twodot_to_onedot=50, # Sweep num to switch from two- to one-dot algo. - max_mem=40, # Max memory (in GB) allotted to fragment DMRG - max_noise=1e-3, # Max MPS noise introduced per sweep - min_tol=1e-8, # Tighest Davidson tolerance per sweep - block_extra_keyword=["fiedler"], # Specify orbital reordering algorithm - force_cleanup=True, # Remove all fragment DMRG tmpfiles only_chem=True, + solver_kwargs=dict( + startM=20, # Initial fragment bond dimension (1st sweep) + maxM=200, # Maximum fragment bond dimension + twodot_to_onedot=50, # Sweep num to switch from two- to one-dot algo. + max_mem=40, # Max memory (in GB) allotted to fragment DMRG + max_noise=1e-3, # Max MPS noise introduced per sweep + min_tol=1e-8, # Tighest Davidson tolerance per sweep + block_extra_keyword=["fiedler"], # Specify orbital reordering algorithm + force_cleanup=True, # Remove all fragment DMRG tmpfiles + ), ) # Or, alternatively, we can construct a full schedule by hand: @@ -126,11 +123,12 @@ # and pass it to the fragment solver through `schedule_kwargs`: mybe.optimize( solver="block2", - scratch=scratch, - schedule_kwargs=schedule, - block_extra_keyword=["fiedler"], - force_cleanup=True, only_chem=True, + solver_kwargs=dict( + schedule_kwargs=schedule, + block_extra_keyword=["fiedler"], + force_cleanup=True, + ), ) # To make sure the calculation is proceeding as expected, make sure to check diff --git a/example/molbe_hexene_oneshot_uccsd.py b/example/molbe_hexene_oneshot_uccsd.py index c54ba570..994ce0f8 100644 --- a/example/molbe_hexene_oneshot_uccsd.py +++ b/example/molbe_hexene_oneshot_uccsd.py @@ -32,4 +32,4 @@ # Perform one round of BE, without density or chemical potential matching, # and return the energy. # clean_eri will delete all of the ERI files from scratch -mybe.oneshot(solver="UCCSD", nproc=nproc, calc_frag_energy=True, clean_eri=True) +mybe.oneshot(solver="UCCSD", nproc=nproc, calc_frag_energy=True) diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py index c34d458b..3cc17e4a 100644 --- a/src/quemb/kbe/pbe.py +++ b/src/quemb/kbe/pbe.py @@ -678,9 +678,7 @@ def initialize(self, compute_hf, restart=False): fobj.udim = couti couti = fobj.set_udim(couti) - def oneshot( - self, solver="MP2", nproc=1, ompnum=4, calc_frag_energy=False, clean_eri=False - ): + def oneshot(self, solver="MP2", nproc=1, ompnum=4, calc_frag_energy=False): """ Perform a one-shot bootstrap embedding calculation. diff --git a/src/quemb/molbe/misc.py b/src/quemb/molbe/misc.py index b695f4dc..be15b78c 100644 --- a/src/quemb/molbe/misc.py +++ b/src/quemb/molbe/misc.py @@ -463,9 +463,7 @@ def be2puffin( # Run oneshot embedding and return system energy - mybe.oneshot( - solver=solver, nproc=nproc, ompnum=ompnum, calc_frag_energy=True, clean_eri=True - ) + mybe.oneshot(solver=solver, nproc=nproc, ompnum=ompnum, calc_frag_energy=True) return mybe.ebe_tot diff --git a/src/quemb/molbe/ube.py b/src/quemb/molbe/ube.py index 91c01446..d31059e4 100644 --- a/src/quemb/molbe/ube.py +++ b/src/quemb/molbe/ube.py @@ -403,9 +403,7 @@ def initialize(self, eri_, compute_hf): fobj.udim = couti couti = fobj.set_udim(couti) - def oneshot( - self, solver="UCCSD", nproc=1, ompnum=4, calc_frag_energy=False, clean_eri=False - ): + def oneshot(self, solver="UCCSD", nproc=1, ompnum=4, calc_frag_energy=False): if nproc == 1: E, E_comp = be_func_u( None, @@ -452,12 +450,7 @@ def oneshot( ) ) - if clean_eri: - try: - os.remove(self.eri_file) - os.rmdir(self.scratch_dir) - except (FileNotFoundError, TypeError): - print("Scratch directory not removed") + self.scratch_dir.cleanup() def initialize_pot(Nfrag, edge_idx): diff --git a/tests/ube-oneshot_test.py b/tests/ube-oneshot_test.py index 452efaeb..eb666597 100644 --- a/tests/ube-oneshot_test.py +++ b/tests/ube-oneshot_test.py @@ -109,7 +109,7 @@ def molecular_unrestricted_oneshot_test( mf.kernel() fobj = fragpart(frag_type="autogen", be_type=be_type, mol=mol, frozen_core=frz) mybe = UBE(mf, fobj) - mybe.oneshot(solver="UCCSD", nproc=1, calc_frag_energy=True, clean_eri=True) + mybe.oneshot(solver="UCCSD", nproc=1, calc_frag_energy=True) self.assertAlmostEqual( mybe.ebe_tot - mybe.uhf_full_e, exp_result, From 2fce4b90c2ca1e283382f1a9b140862bb34bb493 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 11:43:40 -0500 Subject: [PATCH 21/67] fixed be_func_parallel calls --- src/quemb/molbe/be_parallel.py | 34 +++++++++++++++++----------------- src/quemb/molbe/mbe.py | 2 -- src/quemb/molbe/opt.py | 2 -- 3 files changed, 17 insertions(+), 21 deletions(-) diff --git a/src/quemb/molbe/be_parallel.py b/src/quemb/molbe/be_parallel.py index 1ff0fb7a..0687ab81 100644 --- a/src/quemb/molbe/be_parallel.py +++ b/src/quemb/molbe/be_parallel.py @@ -26,6 +26,7 @@ from quemb.shared.external.ccsd_rdm import make_rdm1_uccsd, make_rdm2_uccsd from quemb.shared.external.unrestricted_utils import make_uhf_obj from quemb.shared.helper import unused +from quemb.shared.manage_scratch import WorkDir from quemb.shared.typing import Matrix @@ -397,6 +398,7 @@ def be_func_parallel( Nocc: int, solver: str, enuc: float, # noqa: ARG001 + scratch_dir: WorkDir, hf_veff: Matrix[float64] | None = None, nproc: int = 1, ompnum: int = 4, @@ -459,16 +461,14 @@ def be_func_parallel( Depending on the parameters, returns the error norm or a tuple containing the error norm, error vector, and the computed energy. """ - nfrag = len(Fobjs) # Create directories for fragments if required if writeh1 and solver == "SCI": - for nf in range(nfrag): - dname = Fobjs[nf].dname - os.system("mkdir " + dname) + for fobj in Fobjs: + _ = WorkDir(scratch_dir / fobj.dname) # Set the number of OpenMP threads os.system("export OMP_NUM_THREADS=" + str(ompnum)) - nprocs = int(nproc / ompnum) + nprocs = nproc // ompnum # Update the effective Hamiltonian with potentials if pot is not None: @@ -480,17 +480,17 @@ def be_func_parallel( rdms = [] # Run solver in parallel for each fragment - for nf in range(nfrag): - h1 = Fobjs[nf].fock + Fobjs[nf].heff - dm0 = Fobjs[nf].dm0.copy() - dname = Fobjs[nf].dname - nao = Fobjs[nf].nao - nocc = Fobjs[nf].nsocc - nfsites = Fobjs[nf].nfsites - efac = Fobjs[nf].efac - TA = Fobjs[nf].TA - h1_e = Fobjs[nf].h1 - veff0 = Fobjs[nf].veff0 + for fobj in Fobjs: + h1 = fobj.fock + fobj.heff + dm0 = fobj.dm0.copy() + dname = fobj.dname + nao = fobj.nao + nocc = fobj.nsocc + nfsites = fobj.nfsites + efac = fobj.efac + TA = fobj.TA + h1_e = fobj.h1 + veff0 = fobj.veff0 result = pool_.apply_async( run_solver, @@ -506,7 +506,7 @@ def be_func_parallel( hf_veff, h1_e, solver, - Fobjs[nf].eri_file, + fobj.eri_file, veff0, hci_cutoff, ci_coeff_cutoff, diff --git a/src/quemb/molbe/mbe.py b/src/quemb/molbe/mbe.py index 4d86dc06..a5e81138 100644 --- a/src/quemb/molbe/mbe.py +++ b/src/quemb/molbe/mbe.py @@ -950,13 +950,11 @@ def oneshot( hci_cutoff=self.hci_cutoff, ci_coeff_cutoff=self.ci_coeff_cutoff, select_cutoff=self.select_cutoff, - ereturn=True, eeval=True, frag_energy=calc_frag_energy, nproc=nproc, ompnum=ompnum, scratch_dir=self.scratch_dir, - solver_kwargs=solver_kwargs, ) print("-----------------------------------------------------", flush=True) diff --git a/src/quemb/molbe/opt.py b/src/quemb/molbe/opt.py index 6c54ada6..63bff1fe 100644 --- a/src/quemb/molbe/opt.py +++ b/src/quemb/molbe/opt.py @@ -140,9 +140,7 @@ def objfunc(self, xk: list[float]) -> Vector[float64]: relax_density=self.relax_density, ci_coeff_cutoff=self.ci_coeff_cutoff, select_cutoff=self.select_cutoff, - ebe_hf=self.ebe_hf, scratch_dir=self.scratch_dir, - solver_kwargs=self.solver_kwargs, ) # Update error and BE energy From f5255c96f2e06f0d8548827a7f7d54bd307bddc4 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 11:49:33 -0500 Subject: [PATCH 22/67] simplified be_func_parallel --- src/quemb/molbe/be_parallel.py | 87 +++++++++++++++------------------- 1 file changed, 37 insertions(+), 50 deletions(-) diff --git a/src/quemb/molbe/be_parallel.py b/src/quemb/molbe/be_parallel.py index 0687ab81..a7a5e0be 100644 --- a/src/quemb/molbe/be_parallel.py +++ b/src/quemb/molbe/be_parallel.py @@ -475,59 +475,46 @@ def be_func_parallel( for fobj in Fobjs: fobj.update_heff(pot, only_chem=only_chem) - pool_ = Pool(nprocs) - results = [] - rdms = [] - - # Run solver in parallel for each fragment - for fobj in Fobjs: - h1 = fobj.fock + fobj.heff - dm0 = fobj.dm0.copy() - dname = fobj.dname - nao = fobj.nao - nocc = fobj.nsocc - nfsites = fobj.nfsites - efac = fobj.efac - TA = fobj.TA - h1_e = fobj.h1 - veff0 = fobj.veff0 + with Pool(nprocs) as pool_: + results = [] + rdms = [] - result = pool_.apply_async( - run_solver, - [ - h1, - dm0, - dname, - nao, - nocc, - nfsites, - efac, - TA, - hf_veff, - h1_e, - solver, - fobj.eri_file, - veff0, - hci_cutoff, - ci_coeff_cutoff, - select_cutoff, - ompnum, - writeh1, - True, - True, - use_cumulant, - relax_density, - frag_energy, - ], - ) - - results.append(result) + # Run solver in parallel for each fragment + for fobj in Fobjs: + result = pool_.apply_async( + run_solver, + [ + fobj.fock + fobj.heff, + fobj.dm0.copy(), + fobj.dname, + fobj.nao, + fobj.nsocc, + fobj.nfsites, + fobj.efac, + fobj.TA, + hf_veff, + fobj.h1, + solver, + fobj.eri_file, + fobj.veff0, + hci_cutoff, + ci_coeff_cutoff, + select_cutoff, + ompnum, + writeh1, + True, + True, + use_cumulant, + relax_density, + frag_energy, + ], + ) - # Collect results - for result in results: - rdms.append(result.get()) + results.append(result) - pool_.close() + # Collect results + for result in results: + rdms.append(result.get()) if frag_energy: # Compute and return fragment energy From 17f443fd40cb19c31411930b7e8a983f5d27251a Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 12:02:31 -0500 Subject: [PATCH 23/67] added types to run_solver --- src/quemb/molbe/be_parallel.py | 84 +++++++++++++++++----------------- 1 file changed, 42 insertions(+), 42 deletions(-) diff --git a/src/quemb/molbe/be_parallel.py b/src/quemb/molbe/be_parallel.py index a7a5e0be..9d6a1723 100644 --- a/src/quemb/molbe/be_parallel.py +++ b/src/quemb/molbe/be_parallel.py @@ -31,74 +31,74 @@ def run_solver( - h1, - dm0, - dname, - nao, - nocc, - nfsites, - efac, - TA, - hf_veff, - h1_e, - solver="MP2", - eri_file="eri_file.h5", - veff0=None, - hci_cutoff=0.001, - ci_coeff_cutoff=None, - select_cutoff=None, - ompnum=4, - writeh1=False, - eeval=True, - return_rdm_ao=True, - use_cumulant=True, - relax_density=False, - frag_energy=False, + h1: Matrix[float64], + dm0: Matrix[float64], + dname: str, + nao: int, + nocc: int, + nfsites: int, + efac: float, + TA: Matrix[float64], + hf_veff: Matrix[float64], + h1_e: Matrix[float64], + solver: str = "MP2", + eri_file: str = "eri_file.h5", + veff0: Matrix[float64] | None = None, + hci_cutoff: float = 0.001, + ci_coeff_cutoff: float | None = None, + select_cutoff: float | None = None, + ompnum: int = 4, + writeh1: bool = False, + eeval: bool = True, + return_rdm_ao: bool = True, + use_cumulant: bool = True, + relax_density: bool = False, + frag_energy: bool = False, ): """ Run a quantum chemistry solver to compute the reduced density matrices. Parameters ---------- - h1 : numpy.ndarray + h1 : One-electron Hamiltonian matrix. - dm0 : numpy.ndarray + dm0 : Initial guess for the density matrix. - dname : str + dname : Directory name for storing intermediate files. - nao : int + nao : Number of atomic orbitals. - nocc : int + nocc : Number of occupied orbitals. - nfsites : int + nfsites : Number of fragment sites. - efac : float + efac : Scaling factor for the electronic energy. - TA : numpy.ndarray + TA : Transformation matrix for embedding orbitals. - hf_veff : numpy.ndarray + hf_veff : Hartree-Fock effective potential matrix. - h1_e : numpy.ndarray + h1_e : One-electron integral matrix. - solver : str, optional + solver : Solver to use for the calculation ('MP2', 'CCSD', 'FCI', 'HCI', 'SHCI', 'SCI'). Default is 'MP2'. - eri_file : str, optional + eri_file : Filename for the electron repulsion integrals. Default is 'eri_file.h5'. - ompnum : int, optional + ompnum : Number of OpenMP threads. Default is 4. - writeh1 : bool, optional + writeh1 : If True, write the one-electron integrals to a file. Default is False. - eeval : bool, optional + eeval : If True, evaluate the electronic energy. Default is True. - return_rdm_ao : bool, optional + return_rdm_ao : If True, return the reduced density matrices in the atomic orbital basis. Default is True. - use_cumulant : bool, optional + use_cumulant : If True, use the cumulant approximation for RDM2. Default is True. - frag_energy : bool, optional + frag_energy : If True, compute the fragment energy. Default is False. - relax_density : bool, optional + relax_density : If True, use CCSD relaxed density. Default is False Returns From f6f93c53b74d060dcbf9bc8a4a4b960accc4e9a9 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 12:06:20 -0500 Subject: [PATCH 24/67] use frag_scratch in be_func_parallel --- src/quemb/molbe/be_parallel.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/src/quemb/molbe/be_parallel.py b/src/quemb/molbe/be_parallel.py index 9d6a1723..16249f74 100644 --- a/src/quemb/molbe/be_parallel.py +++ b/src/quemb/molbe/be_parallel.py @@ -33,6 +33,7 @@ def run_solver( h1: Matrix[float64], dm0: Matrix[float64], + frag_scratch: WorkDir, dname: str, nao: int, nocc: int, @@ -64,8 +65,11 @@ def run_solver( One-electron Hamiltonian matrix. dm0 : Initial guess for the density matrix. + scratch_dir : + The scratch dir root. dname : Directory name for storing intermediate files. + frag_scratch: nao : Number of atomic orbitals. nocc : @@ -175,14 +179,14 @@ def run_solver( nao, nmo = mf_.mo_coeff.shape nelec = (nocc, nocc) - mch = shci.SHCISCF(mf_, nmo, nelec, orbpath=dname) + mch = shci.SHCISCF(mf_, nmo, nelec, orbpath=frag_scratch) mch.fcisolver.mpiprefix = "mpirun -np " + str(ompnum) mch.fcisolver.stochastic = True # this is for PT and doesnt add PT to rdm mch.fcisolver.nPTiter = 0 mch.fcisolver.sweep_iter = [0] mch.fcisolver.DoRDM = True mch.fcisolver.sweep_epsilon = [hci_cutoff] - mch.fcisolver.scratchDirectory = "/scratch/oimeitei/" + dname + mch.fcisolver.scratchDirectory = frag_scratch if not writeh1: mch.fcisolver.restart = True mch.mc1step() @@ -202,7 +206,7 @@ def run_solver( ) ci = cornell_shci.SHCI() - ci.runtimedir = dname + ci.runtimedir = frag_scratch ci.restart = True ci.config["var_only"] = True ci.config["eps_vars"] = [hci_cutoff] @@ -464,7 +468,7 @@ def be_func_parallel( # Create directories for fragments if required if writeh1 and solver == "SCI": for fobj in Fobjs: - _ = WorkDir(scratch_dir / fobj.dname) + frag_scratch = WorkDir(scratch_dir / fobj.dname) # Set the number of OpenMP threads os.system("export OMP_NUM_THREADS=" + str(ompnum)) @@ -486,6 +490,7 @@ def be_func_parallel( [ fobj.fock + fobj.heff, fobj.dm0.copy(), + frag_scratch, fobj.dname, fobj.nao, fobj.nsocc, From 6904c628443e6132e734ed22d742c3db70d2e699 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 12:33:52 -0500 Subject: [PATCH 25/67] use frag_scratch in run_solver --- src/quemb/molbe/be_parallel.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/src/quemb/molbe/be_parallel.py b/src/quemb/molbe/be_parallel.py index 16249f74..842620dc 100644 --- a/src/quemb/molbe/be_parallel.py +++ b/src/quemb/molbe/be_parallel.py @@ -33,7 +33,7 @@ def run_solver( h1: Matrix[float64], dm0: Matrix[float64], - frag_scratch: WorkDir, + scratch_dir: WorkDir, dname: str, nao: int, nocc: int, @@ -69,7 +69,9 @@ def run_solver( The scratch dir root. dname : Directory name for storing intermediate files. - frag_scratch: + scratch_dir : + The scratch directory. + Fragment files will be stored in :code:`scratch_dir / dname`. nao : Number of atomic orbitals. nocc : @@ -177,6 +179,8 @@ def run_solver( # pylint: disable-next=E0401,E0611 from pyscf.shciscf import shci # noqa: PLC0415 # shci is an optional module + frag_scratch = WorkDir(scratch_dir / dname) + nao, nmo = mf_.mo_coeff.shape nelec = (nocc, nocc) mch = shci.SHCISCF(mf_, nmo, nelec, orbpath=frag_scratch) @@ -196,6 +200,8 @@ def run_solver( # pylint: disable-next=E0611 from pyscf import cornell_shci # noqa: PLC0415 # optional module + frag_scratch = WorkDir(scratch_dir / dname) + nao, nmo = mf_.mo_coeff.shape nelec = (nocc, nocc) cas = mcscf.CASCI(mf_, nmo, nelec) @@ -465,11 +471,6 @@ def be_func_parallel( Depending on the parameters, returns the error norm or a tuple containing the error norm, error vector, and the computed energy. """ - # Create directories for fragments if required - if writeh1 and solver == "SCI": - for fobj in Fobjs: - frag_scratch = WorkDir(scratch_dir / fobj.dname) - # Set the number of OpenMP threads os.system("export OMP_NUM_THREADS=" + str(ompnum)) nprocs = nproc // ompnum @@ -490,7 +491,7 @@ def be_func_parallel( [ fobj.fock + fobj.heff, fobj.dm0.copy(), - frag_scratch, + scratch_dir, fobj.dname, fobj.nao, fobj.nsocc, From fe7bf55b6a0ff881ab776cfe25ba9cffd9a35d7f Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 12:40:52 -0500 Subject: [PATCH 26/67] added typehints to kbe BE class --- src/quemb/kbe/pbe.py | 95 ++++++++++++++++++++---------------------- src/quemb/molbe/mbe.py | 2 + 2 files changed, 48 insertions(+), 49 deletions(-) diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py index 3cc17e4a..bd2229b8 100644 --- a/src/quemb/kbe/pbe.py +++ b/src/quemb/kbe/pbe.py @@ -7,10 +7,11 @@ import h5py import numpy from libdmet.basis_transform.eri_transform import get_emb_eri_fast_gdf -from pyscf import ao2mo +from pyscf import ao2mo, pbc from pyscf.pbc import df, gto from pyscf.pbc.df.df_jk import _ewald_exxdiv_for_G0 +from quemb.kbe.fragment import fragpart from quemb.kbe.lo import Mixin_k_Localize from quemb.kbe.misc import print_energy, storePBE from quemb.kbe.pfrag import Frags @@ -18,11 +19,12 @@ from quemb.molbe.helper import get_eri, get_scfObj, get_veff from quemb.molbe.opt import BEOPT from quemb.molbe.solver import be_func -from quemb.shared.config import settings from quemb.shared.external.optqn import ( get_be_error_jacobian as _ext_get_be_error_jacobian, ) from quemb.shared.helper import copy_docstring +from quemb.shared.manage_scratch import WorkDir +from quemb.shared.typing import PathLike class BE(Mixin_k_Localize): @@ -47,59 +49,62 @@ class BE(Mixin_k_Localize): def __init__( self, - mf, - fobj, - eri_file="eri_file.h5", - lo_method="lowdin", - compute_hf=True, - restart=False, - save=False, - restart_file="storebe.pk", - save_file="storebe.pk", - hci_pt=False, - nproc=1, - ompnum=4, - hci_cutoff=0.001, - ci_coeff_cutoff=None, - select_cutoff=None, - iao_val_core=True, - exxdiv="ewald", - kpts=None, - cderi=None, - iao_wannier=False, + mf: pbc.scf.hf.SCF, + fobj: fragpart, + eri_file: PathLike = "eri_file.h5", + lo_method: str = "lowdin", + compute_hf: bool = True, + restart: bool = False, + save: bool = False, + restart_file: PathLike = "storebe.pk", + save_file: PathLike = "storebe.pk", + hci_pt: bool = False, + nproc: int = 1, + ompnum: int = 4, + hci_cutoff: float = 0.001, + ci_coeff_cutoff: float | None = None, + select_cutoff: float | None = None, + iao_val_core: bool = True, + exxdiv: str = "ewald", + kpts: list[list[float]] | None = None, + cderi: PathLike | None = None, + iao_wannier: bool = False, + scratch_dir: WorkDir | None = None, ): """ Constructor for BE object. Parameters ---------- - mf : pyscf.pbc.scf.hf.SCF + mf : PySCF periodic mean-field object. - fobj : quemb.kbe.fragment.fragpart + fobj : Fragment object containing sites, centers, edges, and indices. - kpts : list of list of float + kpts : k-points in the reciprocal space for periodic computation - eri_file : str, optional + eri_file : Path to the file storing two-electron integrals, by default 'eri_file.h5'. - lo_method : str, optional + lo_method : Method for orbital localization, by default 'lowdin'. - iao_wannier : bool, optional + iao_wannier : Whether to perform Wannier localization on the IAO space, by default False. - compute_hf : bool, optional + compute_hf : Whether to compute Hartree-Fock energy, by default True. - restart : bool, optional + restart : Whether to restart from a previous calculation, by default False. - save : bool, optional + save : Whether to save intermediate objects for restart, by default False. - restart_file : str, optional + restart_file : Path to the file storing restart information, by default 'storebe.pk'. - save_file : str, optional + save_file : Path to the file storing save information, by default 'storebe.pk'. - nproc : int, optional + nproc : Number of processors for parallel calculations, by default 1. If set to >1, multi-threaded parallel computation is invoked. - ompnum : int, optional + ompnum : Number of OpenMP threads, by default 4. + scratch_dir : + Scratch directory. """ if restart: # Load previous calculation data from restart file @@ -178,25 +183,17 @@ def __init__( self.lmo_coeff = None self.print_ini() - self.Fobjs = [] + self.Fobjs: list[Frags] = [] self.pot = initialize_pot(self.Nfrag, self.edge_idx) self.eri_file = eri_file self.cderi = cderi - # Set scratch directory - jobid = "" - if settings.CREATE_SCRATCH_DIR: - jobid = os.environ.get("SLURM_JOB_ID", "") - if settings.SCRATCH: - os.system("mkdir " + settings.SCRATCH + str(jobid)) - if not jobid: - self.eri_file = settings.SCRATCH + eri_file - if cderi: - self.cderi = settings.SCRATCH + cderi + if scratch_dir is None: + self.scratch_dir = WorkDir.from_environment() else: - self.eri_file = settings.SCRATCH + str(jobid) + "/" + eri_file - if cderi: - self.cderi = settings.SCRATCH + str(jobid) + "/" + cderi + self.scratch_dir = scratch_dir + self.eri_file = self.scratch_dir / eri_file + self.cderi = self.scratch_dir / cderi if cderi else None if exxdiv == "ewald": if not restart: diff --git a/src/quemb/molbe/mbe.py b/src/quemb/molbe/mbe.py index a5e81138..fbf634a2 100644 --- a/src/quemb/molbe/mbe.py +++ b/src/quemb/molbe/mbe.py @@ -116,6 +116,8 @@ def __init__( threaded parallel computation is invoked. ompnum : Number of OpenMP threads, by default 4. + scratch_dir : + Scratch directory. integral_direct_DF: If mf._eri is None (i.e. ERIs are not saved in memory using incore_anyway), this flag is used to determine if the ERIs are computed integral-directly From 9086c9c3b8276e614af9a72352267297b0e0758e Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 12:54:49 -0500 Subject: [PATCH 27/67] simplified a few boolean expressions --- src/quemb/kbe/pbe.py | 2 +- src/quemb/molbe/solver.py | 8 ++------ src/quemb/molbe/ube.py | 2 +- 3 files changed, 4 insertions(+), 8 deletions(-) diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py index bd2229b8..6222e7ca 100644 --- a/src/quemb/kbe/pbe.py +++ b/src/quemb/kbe/pbe.py @@ -214,7 +214,7 @@ def __init__( print("Energy may diverse.", flush=True) print(flush=True) - self.frozen_core = False if not fobj.frozen_core else True + self.frozen_core = fobj.frozen_core self.ncore = 0 if not restart: self.E_core = 0 diff --git a/src/quemb/molbe/solver.py b/src/quemb/molbe/solver.py index e80826c0..2fed3e58 100644 --- a/src/quemb/molbe/solver.py +++ b/src/quemb/molbe/solver.py @@ -367,9 +367,7 @@ def be_func_u( Depending on the options, it returns the norm of the error vector, the energy, or a combination of these values. """ - rdm_return = False - if relax_density: - rdm_return = True + rdm_return = relax_density E = 0.0 if frag_energy or eeval: total_e = [0.0, 0.0, 0.0] @@ -674,8 +672,6 @@ def solve_ccsd( rdm1a = mycc.make_rdm1(with_frozen=False) if rdm2_return: - if use_cumulant: - with_dm1 = False rdm2s = make_rdm2( mycc, mycc.t1, @@ -684,7 +680,7 @@ def solve_ccsd( mycc.l2, with_frozen=False, ao_repr=False, - with_dm1=with_dm1, + with_dm1=with_dm1 and not use_cumulant, ) return (t1, t2, rdm1a, rdm2s) return (t1, t2, rdm1a, mycc) diff --git a/src/quemb/molbe/ube.py b/src/quemb/molbe/ube.py index d31059e4..6d22cf88 100644 --- a/src/quemb/molbe/ube.py +++ b/src/quemb/molbe/ube.py @@ -98,7 +98,7 @@ def __init__( self.eri_file = eri_file self.ek = 0.0 - self.frozen_core = False if not fobj.frozen_core else True + self.frozen_core = fobj.frozen_core self.ncore = 0 self.E_core = 0 self.C_core = None From 925187398e7db0d51191b4b8cf5a8301c0134365 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 13:16:41 -0500 Subject: [PATCH 28/67] the tests should be working now --- src/quemb/kbe/misc.py | 55 +++++++++++----------------- src/quemb/kbe/pbe.py | 67 +++++++++++++++++++--------------- src/quemb/molbe/be_parallel.py | 3 +- src/quemb/molbe/solver.py | 3 +- 4 files changed, 63 insertions(+), 65 deletions(-) diff --git a/src/quemb/kbe/misc.py b/src/quemb/kbe/misc.py index 457b1137..023b4bcf 100644 --- a/src/quemb/kbe/misc.py +++ b/src/quemb/kbe/misc.py @@ -1,9 +1,14 @@ # Author(s): Oinam Romesh Meitei +from attrs import define from numpy import arange, exp, sqrt from pyscf.lib import cartesian_prod from pyscf.pbc import tools +from numpy import float64 + +from quemb.shared.typing import Matrix + def sgeom(cell, kmesh=None): """ @@ -30,41 +35,23 @@ def get_phase1(cell, kpts, kmesh): Ts = cartesian_prod((arange(kmesh[0]), arange(kmesh[1]), arange(kmesh[2]))) return exp(-1.0j * (Ts @ a_vec @ kpts.T)) - +@define class storePBE: - def __init__( - self, - Nocc, - hf_veff, - hcore, - S, - C, - hf_dm, - hf_etot, - W, - lmo_coeff, - enuc, - ek, - E_core, - C_core, - P_core, - core_veff, - ): - self.Nocc = Nocc - self.hf_veff = hf_veff - self.hcore = hcore - self.S = S - self.C = C - self.hf_dm = hf_dm - self.hf_etot = hf_etot - self.W = W - self.lmo_coeff = lmo_coeff - self.enuc = enuc - self.ek = ek - self.E_core = E_core - self.C_core = C_core - self.P_core = P_core - self.core_veff = core_veff + Nocc: int + hf_veff: Matrix[float64] + hcore: Matrix[float64] + S: Matrix[float64] + C: Matrix[float64] + hf_dm: Matrix[float64] + hf_etot: float + W: Matrix[float64] + lmo_coeff: Matrix[float64] + enuc: float + ek: float + E_core: float + C_core: float + P_core: float + core_veff: float def print_energy(ecorr, e_V_Kapprox, e_F_dg, e_hf, unitcell_nkpt): diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py index 6222e7ca..43821421 100644 --- a/src/quemb/kbe/pbe.py +++ b/src/quemb/kbe/pbe.py @@ -24,7 +24,7 @@ ) from quemb.shared.helper import copy_docstring from quemb.shared.manage_scratch import WorkDir -from quemb.shared.typing import PathLike +from quemb.shared.typing import KwargDict, PathLike class BE(Mixin_k_Localize): @@ -378,7 +378,6 @@ def optimize( else: pot = [0.0] - # Initialize the BEOPT object be_ = BEOPT( pot, self.Fobjs, @@ -387,6 +386,7 @@ def optimize( hf_veff=self.hf_veff, nproc=nproc, ompnum=ompnum, + scratch_dir=self.scratch_dir, max_space=max_iter, conv_tol=conv_tol, only_chem=only_chem, @@ -468,7 +468,7 @@ def ewald_sum(self): return e_.real - def initialize(self, compute_hf, restart=False): + def initialize(self, compute_hf: bool, restart: bool = False): """ Initialize the Bootstrap Embedding calculation. @@ -556,26 +556,26 @@ def initialize(self, compute_hf, restart=False): if self.nproc == 1: raise ValueError("If cderi is set, try again with nproc > 1") - nprocs = int(self.nproc / self.ompnum) - pool_ = Pool(nprocs) + nprocs = self.nproc // self.ompnum os.system("export OMP_NUM_THREADS=" + str(self.ompnum)) - results = [] - eris = [] - for frg in range(self.Nfrag): - result = pool_.apply_async( - eritransform_parallel, - [ - self.mf.cell.a, - self.mf.cell.atom, - self.mf.cell.basis, - self.kpts, - self.Fobjs[frg].TA, - self.cderi, - ], - ) - results.append(result) - [eris.append(result.get()) for result in results] - pool_.close() + with Pool(nprocs) as pool_: + results = [] + eris = [] + for frg in range(self.Nfrag): + result = pool_.apply_async( + eritransform_parallel, + [ + self.mf.cell.a, + self.mf.cell.atom, + self.mf.cell.basis, + self.kpts, + self.Fobjs[frg].TA, + self.cderi, + ], + ) + results.append(result) + for result in results: + eris.append(result.get()) for frg in range(self.Nfrag): file_eri.create_dataset(self.Fobjs[frg].dname, data=eris[frg]) @@ -675,23 +675,30 @@ def initialize(self, compute_hf, restart=False): fobj.udim = couti couti = fobj.set_udim(couti) - def oneshot(self, solver="MP2", nproc=1, ompnum=4, calc_frag_energy=False): + def oneshot( + self, + solver: str = "MP2", + nproc: int = 1, + ompnum: int = 4, + calc_frag_energy: bool = False, + solver_kwargs: KwargDict | None = None, + ) -> None: """ Perform a one-shot bootstrap embedding calculation. Parameters ---------- - solver : str, optional + solver : High-level quantum chemistry method, by default 'MP2'. 'CCSD', 'FCI', and variants of selected CI are supported. - nproc : int, optional + nproc : Number of processors for parallel calculations, by default 1. If set to >1, threaded parallel computation is invoked. - ompnum : int, optional + ompnum : Number of OpenMP threads, by default 4. - calc_frag_energy : bool, optional + calc_frag_energy : Whether to calculate fragment energies, by default False. - clean_eri : bool, optional + clean_eri : Whether to clean up ERI files after calculation, by default False. """ print("Calculating Energy by Fragment? ", calc_frag_energy) @@ -710,6 +717,8 @@ def oneshot(self, solver="MP2", nproc=1, ompnum=4, calc_frag_energy=False): frag_energy=calc_frag_energy, ereturn=True, eeval=True, + scratch_dir=self.scratch_dir, + solver_kwargs=solver_kwargs, ) else: rets = be_func_parallel( @@ -722,11 +731,11 @@ def oneshot(self, solver="MP2", nproc=1, ompnum=4, calc_frag_energy=False): hci_cutoff=self.hci_cutoff, ci_coeff_cutoff=self.ci_coeff_cutoff, select_cutoff=self.select_cutoff, - ereturn=True, eeval=True, frag_energy=calc_frag_energy, nproc=nproc, ompnum=ompnum, + scratch_dir=self.scratch_dir, ) print("-----------------------------------------------------", flush=True) diff --git a/src/quemb/molbe/be_parallel.py b/src/quemb/molbe/be_parallel.py index 842620dc..7ebfc57c 100644 --- a/src/quemb/molbe/be_parallel.py +++ b/src/quemb/molbe/be_parallel.py @@ -8,6 +8,7 @@ from numpy.linalg import multi_dot from pyscf import ao2mo, fci, mcscf +from quemb.kbe.pfrag import Frags as pFrags from quemb.molbe.helper import ( get_eri, get_frag_energy, @@ -404,7 +405,7 @@ def run_solver_u( def be_func_parallel( pot: list[float] | None, - Fobjs: list[Frags], + Fobjs: list[Frags] | list[pFrags], Nocc: int, solver: str, enuc: float, # noqa: ARG001 diff --git a/src/quemb/molbe/solver.py b/src/quemb/molbe/solver.py index 2fed3e58..6fbb1857 100644 --- a/src/quemb/molbe/solver.py +++ b/src/quemb/molbe/solver.py @@ -8,6 +8,7 @@ from pyscf import ao2mo, cc, fci, mcscf, mp from pyscf.cc.ccsd_rdm import make_rdm2 +from quemb.kbe.pfrag import Frags as pFrags from quemb.molbe.helper import get_frag_energy, get_frag_energy_u from quemb.molbe.pfrag import Frags from quemb.shared.external.ccsd_rdm import ( @@ -25,7 +26,7 @@ def be_func( pot: list[float] | None, - Fobjs: list[Frags], + Fobjs: list[Frags] | list[pFrags], Nocc: int, solver: str, enuc: float, # noqa: ARG001 From a329d326d62429d423732cbcc7d4b024c913c8ea Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 14:05:01 -0500 Subject: [PATCH 29/67] ensure scratch cleanup for kbe --- src/quemb/kbe/misc.py | 5 ++--- src/quemb/kbe/pbe.py | 9 +++------ 2 files changed, 5 insertions(+), 9 deletions(-) diff --git a/src/quemb/kbe/misc.py b/src/quemb/kbe/misc.py index 023b4bcf..78fd48c4 100644 --- a/src/quemb/kbe/misc.py +++ b/src/quemb/kbe/misc.py @@ -1,12 +1,10 @@ # Author(s): Oinam Romesh Meitei from attrs import define -from numpy import arange, exp, sqrt +from numpy import arange, exp, float64, sqrt from pyscf.lib import cartesian_prod from pyscf.pbc import tools -from numpy import float64 - from quemb.shared.typing import Matrix @@ -35,6 +33,7 @@ def get_phase1(cell, kpts, kmesh): Ts = cartesian_prod((arange(kmesh[0]), arange(kmesh[1]), arange(kmesh[2]))) return exp(-1.0j * (Ts @ a_vec @ kpts.T)) + @define class storePBE: Nocc: int diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py index 43821421..0427285c 100644 --- a/src/quemb/kbe/pbe.py +++ b/src/quemb/kbe/pbe.py @@ -421,6 +421,8 @@ def optimize( else: raise ValueError("This optimization method for BE is not supported") + self.scratch_dir.cleanup() + @copy_docstring(_ext_get_be_error_jacobian) def get_be_error_jacobian(self, jac_solver="HF"): return _ext_get_be_error_jacobian(self.Nfrag, self.Fobjs, jac_solver) @@ -764,12 +766,7 @@ def oneshot( if not calc_frag_energy: self.compute_energy_full(approx_cumulant=True, return_rdm=False) - if clean_eri: - try: - os.remove(self.eri_file) - os.rmdir(self.scratch_dir) - except (FileNotFoundError, TypeError): - print("Scratch directory not removed") + self.scratch_dir.cleanup() def update_fock(self, heff=None): """ From 28ac432c42b79248cba1db54e27b6e04dc134cfc Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 14:08:18 -0500 Subject: [PATCH 30/67] removed call to self.compute_energy_full (related to #35) --- src/quemb/kbe/pbe.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py index 0427285c..bb0c2ca4 100644 --- a/src/quemb/kbe/pbe.py +++ b/src/quemb/kbe/pbe.py @@ -763,9 +763,6 @@ def oneshot( self.ebe_tot = rets[0] - if not calc_frag_energy: - self.compute_energy_full(approx_cumulant=True, return_rdm=False) - self.scratch_dir.cleanup() def update_fock(self, heff=None): From 044707fa43482fa01e3dd42ad2fdd5e4d8224beb Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 14:16:15 -0500 Subject: [PATCH 31/67] write nicer Pool statements - used proper list comprehensions - use contextmanager for Pool --- src/quemb/kbe/pbe.py | 71 ++++++++++++++++------------------ src/quemb/molbe/be_parallel.py | 56 ++++++++++++--------------- 2 files changed, 57 insertions(+), 70 deletions(-) diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py index bb0c2ca4..562a10b9 100644 --- a/src/quemb/kbe/pbe.py +++ b/src/quemb/kbe/pbe.py @@ -562,7 +562,6 @@ def initialize(self, compute_hf: bool, restart: bool = False): os.system("export OMP_NUM_THREADS=" + str(self.ompnum)) with Pool(nprocs) as pool_: results = [] - eris = [] for frg in range(self.Nfrag): result = pool_.apply_async( eritransform_parallel, @@ -576,8 +575,7 @@ def initialize(self, compute_hf: bool, restart: bool = False): ], ) results.append(result) - for result in results: - eris.append(result.get()) + eris = [result.get() for result in results] for frg in range(self.Nfrag): file_eri.create_dataset(self.Fobjs[frg].dname, data=eris[frg]) @@ -585,25 +583,23 @@ def initialize(self, compute_hf: bool, restart: bool = False): file_eri.close() nprocs = int(self.nproc / self.ompnum) - pool_ = Pool(nprocs) - results = [] - veffs = [] - for frg in range(self.Nfrag): - result = pool_.apply_async( - parallel_fock_wrapper, - [ - self.Fobjs[frg].dname, - self.Fobjs[frg].nao, - self.hf_dm, - self.S, - self.Fobjs[frg].TA, - self.hf_veff, - self.eri_file, - ], - ) - results.append(result) - [veffs.append(result.get()) for result in results] - pool_.close() + with Pool(nprocs) as pool_: + results = [] + for frg in range(self.Nfrag): + result = pool_.apply_async( + parallel_fock_wrapper, + [ + self.Fobjs[frg].dname, + self.Fobjs[frg].nao, + self.hf_dm, + self.S, + self.Fobjs[frg].TA, + self.hf_veff, + self.eri_file, + ], + ) + results.append(result) + veffs = [result.get() for result in results] for frg in range(self.Nfrag): veff0, veff_ = veffs[frg] @@ -623,22 +619,21 @@ def initialize(self, compute_hf: bool, restart: bool = False): self.Fobjs[frg].scf(fs=True, dm0=self.Fobjs[frg].dm_init) else: nprocs = int(self.nproc / self.ompnum) - pool_ = Pool(nprocs) - os.system("export OMP_NUM_THREADS=" + str(self.ompnum)) - results = [] - mo_coeffs = [] - for frg in range(self.Nfrag): - nao = self.Fobjs[frg].nao - nocc = self.Fobjs[frg].nsocc - dname = self.Fobjs[frg].dname - h1 = self.Fobjs[frg].fock + self.Fobjs[frg].heff - result = pool_.apply_async( - parallel_scf_wrapper, - [dname, nao, nocc, h1, self.Fobjs[frg].dm_init, self.eri_file], - ) - results.append(result) - [mo_coeffs.append(result.get()) for result in results] - pool_.close() + with Pool(nprocs) as pool_: + os.system("export OMP_NUM_THREADS=" + str(self.ompnum)) + results = [] + for frg in range(self.Nfrag): + nao = self.Fobjs[frg].nao + nocc = self.Fobjs[frg].nsocc + dname = self.Fobjs[frg].dname + h1 = self.Fobjs[frg].fock + self.Fobjs[frg].heff + result = pool_.apply_async( + parallel_scf_wrapper, + [dname, nao, nocc, h1, self.Fobjs[frg].dm_init, self.eri_file], + ) + results.append(result) + mo_coeffs = [result.get() for result in results] + for frg in range(self.Nfrag): self.Fobjs[frg]._mo_coeffs = mo_coeffs[frg] diff --git a/src/quemb/molbe/be_parallel.py b/src/quemb/molbe/be_parallel.py index 7ebfc57c..6f2611b1 100644 --- a/src/quemb/molbe/be_parallel.py +++ b/src/quemb/molbe/be_parallel.py @@ -483,8 +483,6 @@ def be_func_parallel( with Pool(nprocs) as pool_: results = [] - rdms = [] - # Run solver in parallel for each fragment for fobj in Fobjs: result = pool_.apply_async( @@ -519,9 +517,7 @@ def be_func_parallel( results.append(result) - # Collect results - for result in results: - rdms.append(result.get()) + rdms = [result.get() for result in results] if frag_energy: # Compute and return fragment energy @@ -612,34 +608,30 @@ def be_func_parallel_u( """ # Set the number of OpenMP threads os.system("export OMP_NUM_THREADS=" + str(ompnum)) - nprocs = int(nproc / ompnum) - - pool_ = Pool(nprocs) - results = [] - energy_list = [] - - # Run solver in parallel for each fragment - for fobj_a, fobj_b in Fobjs: - result = pool_.apply_async( - run_solver_u, - [ - fobj_a, - fobj_b, - solver, - enuc, - hf_veff, - frag_energy, - relax_density, - frozen, - use_cumulant, - True, - ], - ) - results.append(result) + nprocs = nproc // ompnum + + with Pool(nprocs) as pool_: + results = [] + # Run solver in parallel for each fragment + for fobj_a, fobj_b in Fobjs: + result = pool_.apply_async( + run_solver_u, + [ + fobj_a, + fobj_b, + solver, + enuc, + hf_veff, + frag_energy, + relax_density, + frozen, + use_cumulant, + True, + ], + ) + results.append(result) - # Collect results - [energy_list.append(result.get()) for result in results] - pool_.close() + energy_list = [result.get() for result in results] if frag_energy: # Compute and return fragment energy From 12bba3b6e5d0878daaf09128fcdfc4b822a53f98 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 14:54:21 -0500 Subject: [PATCH 32/67] fixed naming error --- src/quemb/kbe/pfrag.py | 4 ++-- src/quemb/molbe/be_parallel.py | 20 ++++++++++++-------- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/src/quemb/kbe/pfrag.py b/src/quemb/kbe/pfrag.py index b57a3193..67035d6b 100644 --- a/src/quemb/kbe/pfrag.py +++ b/src/quemb/kbe/pfrag.py @@ -93,8 +93,8 @@ def __init__( self.udim = None self._rdm1 = None - self.__rdm1 = None - self.__rdm2 = None + self.rdm1__ = None + self.rdm2__ = None self._del_rdm1 = None self.rdm1 = None self.genvs = None diff --git a/src/quemb/molbe/be_parallel.py b/src/quemb/molbe/be_parallel.py index 6f2611b1..ec23f505 100644 --- a/src/quemb/molbe/be_parallel.py +++ b/src/quemb/molbe/be_parallel.py @@ -535,14 +535,18 @@ def be_func_parallel( e_1 = 0.0 e_2 = 0.0 e_c = 0.0 - for idx, fobj in enumerate(Fobjs): - e_1 += rdms[idx][0][0] - e_2 += rdms[idx][0][1] - e_c += rdms[idx][0][2] - fobj.mo_coeffs = rdms[idx][1] - fobj._rdm1 = rdms[idx][2] - fobj.rdm2__ = rdms[idx][3] - fobj.rdm1__ = rdms[idx][4] + + # I have to type ignore here, because of stupid behaviour of + # :code:`zip` and :code:`enumerate` + # https://stackoverflow.com/questions/74374059/correctly-specify-the-types-of-unpacked-zip + for fobj, rdm in zip(Fobjs, rdms): # type: ignore[assignment] + e_1 += rdm[0][0] + e_2 += rdm[0][1] + e_c += rdm[0][2] + fobj.mo_coeffs = rdm[1] + fobj._rdm1 = rdm[2] + fobj.rdm2__ = rdm[3] + fobj.rdm1__ = rdm[4] del rdms ernorm, ervec = solve_error(Fobjs, Nocc, only_chem=only_chem) From 182a0158d939e05667afd4fba1b887f05a177cf0 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 15:00:27 -0500 Subject: [PATCH 33/67] use more explicit way of freeing memory (it is still ugly... ) --- src/quemb/kbe/pbe.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py index 562a10b9..02231cd2 100644 --- a/src/quemb/kbe/pbe.py +++ b/src/quemb/kbe/pbe.py @@ -579,7 +579,7 @@ def initialize(self, compute_hf: bool, restart: bool = False): for frg in range(self.Nfrag): file_eri.create_dataset(self.Fobjs[frg].dname, data=eris[frg]) - eris = None + del eris file_eri.close() nprocs = int(self.nproc / self.ompnum) @@ -610,7 +610,7 @@ def initialize(self, compute_hf: bool, restart: bool = False): raise ValueError(f"Imaginary Veff {numpy.abs(veff_.imag).max()}") self.Fobjs[frg].fock = self.Fobjs[frg].h1 + veff_.real - veffs = None + del veffs # SCF parallelized if self.nproc == 1 and not transform_parallel: From 5f851b66b3c3bc4498ab17854b5146ecaa30c270 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 15:03:25 -0500 Subject: [PATCH 34/67] refactored expression --- src/quemb/kbe/pbe.py | 2 +- src/quemb/molbe/lo.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py index 02231cd2..7458c5df 100644 --- a/src/quemb/kbe/pbe.py +++ b/src/quemb/kbe/pbe.py @@ -582,7 +582,7 @@ def initialize(self, compute_hf: bool, restart: bool = False): del eris file_eri.close() - nprocs = int(self.nproc / self.ompnum) + nprocs = self.nproc // self.ompnum with Pool(nprocs) as pool_: results = [] for frg in range(self.Nfrag): diff --git a/src/quemb/molbe/lo.py b/src/quemb/molbe/lo.py index 381a3d2a..539c6282 100644 --- a/src/quemb/molbe/lo.py +++ b/src/quemb/molbe/lo.py @@ -35,7 +35,7 @@ def cano_orth(A, thr=1.0e-6, ovlp=None): def get_symm_orth_mat(A, thr=1.0e-6, ovlp=None): S = dot_gen(A, A, ovlp) e, u = eigh(S) - if int(numpy.sum(e < thr)) > 0: + if (e < thr).any(): raise ValueError( "Linear dependence is detected in the column space of A: " "smallest eigenvalue (%.3E) is less than thr (%.3E). " From be0c4032ea49281156eb2c8f8757b51a5409d0d3 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 15:06:34 -0500 Subject: [PATCH 35/67] use Tuple[int, ...] for numpy shapes :-( --- src/quemb/shared/typing.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/src/quemb/shared/typing.py b/src/quemb/shared/typing.py index f23668e9..802876ba 100644 --- a/src/quemb/shared/typing.py +++ b/src/quemb/shared/typing.py @@ -22,20 +22,25 @@ #: Type annotation of a generic covariant type. T_dtype_co = TypeVar("T_dtype_co", bound=np.generic, covariant=True) +# Currently we can define :code:`Matrix` and higher order tensors +# only with shape :code`Tuple[int, ...]` because of +# https://github.com/numpy/numpy/issues/27957 +# make the typechecks more strict over time, when shape checking finally comes to numpy. + #: Type annotation of a vector. Vector = np.ndarray[Tuple[int], np.dtype[T_dtype_co]] #: Type annotation of a matrix. -Matrix = np.ndarray[Tuple[int, int], np.dtype[T_dtype_co]] +Matrix = np.ndarray[Tuple[int, ...], np.dtype[T_dtype_co]] #: Type annotation of a tensor. -Tensor3D = np.ndarray[Tuple[int, int, int], np.dtype[T_dtype_co]] +Tensor3D = np.ndarray[Tuple[int, ...], np.dtype[T_dtype_co]] #: Type annotation of a tensor. -Tensor4D = np.ndarray[Tuple[int, int, int, int], np.dtype[T_dtype_co]] +Tensor4D = np.ndarray[Tuple[int, ...], np.dtype[T_dtype_co]] #: Type annotation of a tensor. -Tensor5D = np.ndarray[Tuple[int, int, int, int, int], np.dtype[T_dtype_co]] +Tensor5D = np.ndarray[Tuple[int, ...], np.dtype[T_dtype_co]] #: Type annotation of a tensor. -Tensor6D = np.ndarray[Tuple[int, int, int, int, int, int], np.dtype[T_dtype_co]] +Tensor6D = np.ndarray[Tuple[int, ...], np.dtype[T_dtype_co]] #: Type annotation of a tensor. -Tensor7D = np.ndarray[Tuple[int, int, int, int, int, int, int], np.dtype[T_dtype_co]] +Tensor7D = np.ndarray[Tuple[int, ...], np.dtype[T_dtype_co]] #: Type annotation of a tensor. Tensor = np.ndarray[Tuple[int, ...], np.dtype[T_dtype_co]] From b6a187cf1992ed194019a2e402af2b4dce96c270 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 18:06:10 -0500 Subject: [PATCH 36/67] refactored WorkDir to use atexit.register for cleanup --- src/quemb/kbe/pbe.py | 4 -- src/quemb/molbe/mbe.py | 4 -- src/quemb/molbe/ube.py | 2 - src/quemb/shared/manage_scratch.py | 105 ++++++++++++++++++----------- 4 files changed, 66 insertions(+), 49 deletions(-) diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py index 7458c5df..1516fd41 100644 --- a/src/quemb/kbe/pbe.py +++ b/src/quemb/kbe/pbe.py @@ -421,8 +421,6 @@ def optimize( else: raise ValueError("This optimization method for BE is not supported") - self.scratch_dir.cleanup() - @copy_docstring(_ext_get_be_error_jacobian) def get_be_error_jacobian(self, jac_solver="HF"): return _ext_get_be_error_jacobian(self.Nfrag, self.Fobjs, jac_solver) @@ -758,8 +756,6 @@ def oneshot( self.ebe_tot = rets[0] - self.scratch_dir.cleanup() - def update_fock(self, heff=None): """ Update the Fock matrix for each fragment with the effective Hamiltonian. diff --git a/src/quemb/molbe/mbe.py b/src/quemb/molbe/mbe.py index fbf634a2..319e7da2 100644 --- a/src/quemb/molbe/mbe.py +++ b/src/quemb/molbe/mbe.py @@ -734,8 +734,6 @@ def optimize( else: raise ValueError("This optimization method for BE is not supported") - self.scratch_dir.cleanup() - @copy_docstring(_ext_get_be_error_jacobian) def get_be_error_jacobian(self, jac_solver: str = "HF") -> list[list[float]]: return _ext_get_be_error_jacobian(self.Nfrag, self.Fobjs, jac_solver) @@ -985,8 +983,6 @@ def oneshot( if not calc_frag_energy: self.compute_energy_full(approx_cumulant=True, return_rdm=False) - self.scratch_dir.cleanup() - def update_fock(self, heff=None): """ Update the Fock matrix for each fragment with the effective Hamiltonian. diff --git a/src/quemb/molbe/ube.py b/src/quemb/molbe/ube.py index 6d22cf88..11b04aaf 100644 --- a/src/quemb/molbe/ube.py +++ b/src/quemb/molbe/ube.py @@ -450,8 +450,6 @@ def oneshot(self, solver="UCCSD", nproc=1, ompnum=4, calc_frag_energy=False): ) ) - self.scratch_dir.cleanup() - def initialize_pot(Nfrag, edge_idx): pot_ = [] diff --git a/src/quemb/shared/manage_scratch.py b/src/quemb/shared/manage_scratch.py index 87e64724..2d2565eb 100644 --- a/src/quemb/shared/manage_scratch.py +++ b/src/quemb/shared/manage_scratch.py @@ -1,6 +1,8 @@ from __future__ import annotations +import atexit import os +from functools import partial from pathlib import Path from shutil import rmtree from types import TracebackType @@ -12,8 +14,42 @@ from quemb.shared.typing import PathLike -def _to_abs_path(pathlike: PathLike) -> Path: - return Path(pathlike).resolve() +def _determine_path( + root: PathLike | None = None, + subdir_prefix: str | None = None, +) -> Path: + """Find a good path name for the scratch directory. + + The naming scheme is :python:`f"{root}/{subdir_prefix}{SLURM_JOB_ID}"` + on systems with :python:`SLURM`. + If :python:`SLURM` is not available, then the process ID is used instead. + """ + scratch_root = Path(root) if root else Path(settings.SCRATCH) + subdir_prefix = "QuEmb_" if subdir_prefix is None else subdir_prefix + if "SLURM_JOB_ID" in os.environ: + # we can safely assume that the SLURM_JOB_ID is unique + subdir = Path(f"{subdir_prefix}{os.environ['SLURM_JOB_ID']}/") + else: + # We cannot safely assume that PIDs are unique + id = os.getpid() + subdir = Path(f"{subdir_prefix}{id}/") + while (scratch_root / subdir).exists(): + id = id + 1 + subdir = Path(f"{subdir_prefix}{id}/") + return scratch_root / subdir + + +def _get_abs_path(pathlike: PathLike | None) -> Path: + """Return valid path names for the :class:`WorkDir` + + Ensure that absolute paths are returned. + if :class:`None` is given as argument, then the path name is automatically + determined. + """ + if pathlike is None: + return _determine_path().resolve() + else: + return Path(pathlike).resolve() @define(order=False) @@ -23,29 +59,31 @@ class WorkDir: Upon initialisation of the object the workdir :python:`path` is created, if it does not exist yet. If it already exists, it is ensured, that it is empty. - Internally the :python:`path` will be stored as absolute. + One can either enter a :python:`path` themselves, or if it is :class:`None`, + then the path is automatically determined by the environment, + i.e. are we on SLURM, where is the scratch etc. + Read more at :func:`from_environment` for more details. - If :python:`do_cleanup` is true, then the scratch area is deleted, - when :python:`self.cleanup` is called. - - Not that the :python:`/` is overloaded for this class and it can be used + Note that the :python:`/` is overloaded for this class and it can be used as :python:`pathlib.Path` in that regard, see example below. - + If :python:`cleanup_at_end` is true, + then :func:`cleanup` method is registered to be called when + the program terminates with no errors and deletes the scratch directory. The :python:`WorkDir` also exists as a ContextManager; then the cleanup is performed when leaving the ContextManager. - See an example below. + Again, assuming that :python:`cleanup_at_end` is true. Examples -------- - >>> with WorkDir('./test_dir') as scratch: + >>> with WorkDir('./test_dir', cleanup_at_end=True) as scratch: >>> with open(scratch / 'test.txt', 'w') as f: >>> f.write('hello world') './test_dir' does not exist anymore, if the outer contextmanager is left without errors. """ - path: Final[Annotated[Path, "An absolute path"]] = field(converter=_to_abs_path) + path: Final[Annotated[Path, "An absolute path"]] = field(converter=_get_abs_path) cleanup_at_end: Final[bool] = True # The __init__ is automatically created @@ -55,6 +93,8 @@ def __attrs_post_init__(self) -> None: self.path.mkdir(parents=True, exist_ok=True) if any(self.path.iterdir()): raise ValueError("scratch_area has to be empty.") + if self.cleanup_at_end: + atexit.register(partial(self.cleanup, ignore_error=True)) def __enter__(self) -> WorkDir: return self @@ -65,7 +105,7 @@ def __exit__( value: BaseException | None, traceback: TracebackType | None, ) -> Literal[False]: - if value is None: + if value is None and self.cleanup_at_end: self.cleanup() return False @@ -74,8 +114,8 @@ def from_environment( cls, *, user_defined_root: PathLike | None = None, - prefix: str = "QuEmb_", - do_cleanup: bool = True, + prefix: str | None = None, + cleanup_at_end: bool = True, ) -> WorkDir: """Create a WorkDir based on the environment. @@ -88,42 +128,29 @@ def from_environment( user_defined_root: The root directory where to create temporary directories e.g. :bash:`/tmp` or :bash:`/scratch`. - If :python:`None`, then the value from :python:`quemb.settings.SCRATCH` + If :class:`None`, then the :python:`SCRATCH` + value from :class:`quemb.shared.config.Settings` is taken. prefix: The prefix for the subdirectory. - do_cleanup: + cleanup_at_end: Perform cleanup when calling :python:`self.cleanup`. """ - scratch_root = ( - Path(user_defined_root) if user_defined_root else Path(settings.SCRATCH) - ) - - if "SLURM_JOB_ID" in os.environ: - # we can safely assume that the SLURM_JOB_ID is unique - subdir = Path(f"{prefix}{os.environ['SLURM_JOB_ID']}/") - else: - # We cannot safely assume that PIDs are unique - id = os.getpid() - subdir = Path(f"{prefix}{id}/") - while (scratch_root / subdir).exists(): - id = id + 1 - subdir = Path(f"{prefix}{id}/") - return cls(scratch_root / subdir, do_cleanup) - - def cleanup(self, force_cleanup: bool = False) -> None: + return cls(_determine_path(user_defined_root, prefix), cleanup_at_end) + + def cleanup(self, ignore_error: bool = False) -> None: """Conditionally cleanup the working directory. Parameters ---------- - force_cleanup: - If the instance was initialized with :python:`cleanup_at_end=True`, - or the argument :python:`force_cleanup` is given, then - the working directory is deleted. - Otherwise nothing happens. + ignore_errors : + Ignore :class:`FileNotFoundError`, and only that exception, when deleting. """ - if self.cleanup_at_end or force_cleanup: + try: rmtree(self.path) + except FileNotFoundError as e: + if not ignore_error: + raise e def __fspath__(self) -> str: return self.path.__fspath__() From c8cf15a5b3be527ad23d795bdeeab7f3c920675d Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 18:14:44 -0500 Subject: [PATCH 37/67] added better docstring for config --- src/quemb/shared/config.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/src/quemb/shared/config.py b/src/quemb/shared/config.py index b78fb8f8..6667e085 100644 --- a/src/quemb/shared/config.py +++ b/src/quemb/shared/config.py @@ -1,3 +1,22 @@ +"""Configure :python:`quemb` + +One can modify settings in one session or create an RC-file. +See examples below. + +Examples +-------- +>>> from quemb.shared.config import settings +>>> +>>> settings.SCRATCH = "/scratch" +Changes the default root for the scratch directory +for this python session. + +>>> from quemb.shared.config import dump_settings +>>> +>>> dump_settings() +Creates ~/.quembrc.yml file that allows changes to persist. +""" + from pathlib import Path from typing import Final From d37202265a306709447460c731e18a9e88659cfe Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 18:15:45 -0500 Subject: [PATCH 38/67] require static analysis to be ok for running test suite --- .github/workflows/quemb_unittest.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/quemb_unittest.yml b/.github/workflows/quemb_unittest.yml index ac60f76f..ef846291 100644 --- a/.github/workflows/quemb_unittest.yml +++ b/.github/workflows/quemb_unittest.yml @@ -15,6 +15,7 @@ permissions: jobs: analysis: runs-on: ubuntu-latest + needs: analysis strategy: matrix: python-version: ["3.10"] From d86d09b704e625c99db457dd24cabc3a1e26f8f1 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 18:19:37 -0500 Subject: [PATCH 39/67] renamed DMRG specific solver kwargs to its proper name --- example/molbe_dmrg_block2.py | 6 ++--- src/quemb/kbe/pbe.py | 4 ++-- src/quemb/molbe/mbe.py | 8 +++---- src/quemb/molbe/opt.py | 4 ++-- src/quemb/molbe/solver.py | 45 +++++++++++++++++++++--------------- 5 files changed, 37 insertions(+), 30 deletions(-) diff --git a/example/molbe_dmrg_block2.py b/example/molbe_dmrg_block2.py index f59f6473..0eaa5a42 100644 --- a/example/molbe_dmrg_block2.py +++ b/example/molbe_dmrg_block2.py @@ -52,7 +52,7 @@ # Next, run BE-DMRG with default parameters and maxM=100. mybe.oneshot( solver="block2", # or 'DMRG', 'DMRGSCF', 'DMRGCI' - solver_kwargs=dict( + DMRG_solver_kwargs=dict( maxM=100, # Max fragment bond dimension force_cleanup=True, # Remove all fragment DMRG tmpfiles ), @@ -100,7 +100,7 @@ solver="block2", # or 'DMRG', 'DMRGSCF', 'DMRGCI' max_iter=60, # Max number of sweeps only_chem=True, - solver_kwargs=dict( + DMRG_solver_kwargs=dict( startM=20, # Initial fragment bond dimension (1st sweep) maxM=200, # Maximum fragment bond dimension twodot_to_onedot=50, # Sweep num to switch from two- to one-dot algo. @@ -124,7 +124,7 @@ mybe.optimize( solver="block2", only_chem=True, - solver_kwargs=dict( + DMRG_solver_kwargs=dict( schedule_kwargs=schedule, block_extra_keyword=["fiedler"], force_cleanup=True, diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py index 1516fd41..d24aeb70 100644 --- a/src/quemb/kbe/pbe.py +++ b/src/quemb/kbe/pbe.py @@ -676,7 +676,7 @@ def oneshot( nproc: int = 1, ompnum: int = 4, calc_frag_energy: bool = False, - solver_kwargs: KwargDict | None = None, + DMRG_solver_kwargs: KwargDict | None = None, ) -> None: """ Perform a one-shot bootstrap embedding calculation. @@ -713,7 +713,7 @@ def oneshot( ereturn=True, eeval=True, scratch_dir=self.scratch_dir, - solver_kwargs=solver_kwargs, + DMRG_solver_kwargs=DMRG_solver_kwargs, ) else: rets = be_func_parallel( diff --git a/src/quemb/molbe/mbe.py b/src/quemb/molbe/mbe.py index 319e7da2..be61afa8 100644 --- a/src/quemb/molbe/mbe.py +++ b/src/quemb/molbe/mbe.py @@ -645,7 +645,7 @@ def optimize( ompnum: int = 4, max_iter: int = 500, trust_region: bool = False, - solver_kwargs: KwargDict | None = None, + DMRG_solver_kwargs: KwargDict | None = None, ) -> None: """BE optimization function @@ -712,7 +712,7 @@ def optimize( hci_pt=self.hci_pt, solver=solver, ebe_hf=self.ebe_hf, - solver_kwargs=solver_kwargs, + DMRG_solver_kwargs=DMRG_solver_kwargs, ) if method == "QN": @@ -902,7 +902,7 @@ def oneshot( nproc: int = 1, ompnum: int = 4, calc_frag_energy: bool = False, - solver_kwargs: KwargDict | None = None, + DMRG_solver_kwargs: KwargDict | None = None, ): """ Perform a one-shot bootstrap embedding calculation. @@ -937,7 +937,7 @@ def oneshot( ereturn=True, eeval=True, scratch_dir=self.scratch_dir, - solver_kwargs=solver_kwargs, + DMRG_solver_kwargs=DMRG_solver_kwargs, ) else: rets = be_func_parallel( diff --git a/src/quemb/molbe/opt.py b/src/quemb/molbe/opt.py index a0aba0cf..22be2b9e 100644 --- a/src/quemb/molbe/opt.py +++ b/src/quemb/molbe/opt.py @@ -76,7 +76,7 @@ class BEOPT: err: float = 0.0 Ebe: Matrix[float64] = array([[0.0]]) - solver_kwargs: KwargDict | None = None + DMRG_solver_kwargs: KwargDict | None = None # HCI parameters hci_cutoff: float = 0.001 @@ -121,7 +121,7 @@ def objfunc(self, xk: list[float]) -> Vector[float64]: select_cutoff=self.select_cutoff, hci_pt=self.hci_pt, scratch_dir=self.scratch_dir, - solver_kwargs=self.solver_kwargs, + DMRG_solver_kwargs=self.DMRG_solver_kwargs, ) else: err_, errvec_, ebe_ = be_func_parallel( diff --git a/src/quemb/molbe/solver.py b/src/quemb/molbe/solver.py index 6fbb1857..c7c5fe87 100644 --- a/src/quemb/molbe/solver.py +++ b/src/quemb/molbe/solver.py @@ -30,7 +30,7 @@ def be_func( Nocc: int, solver: str, enuc: float, # noqa: ARG001 - solver_kwargs: KwargDict | None, + DMRG_solver_kwargs: KwargDict | None, scratch_dir: WorkDir, hf_veff: Matrix[float64] | None = None, only_chem: bool = False, @@ -215,17 +215,22 @@ def be_func( rdm1_tmp, rdm2s = ci.make_rdm12(0, nmo, nelec) elif solver in ["block2", "DMRG", "DMRGCI", "DMRGSCF"]: - solver_kwargs_ = {} if solver_kwargs is None else solver_kwargs.copy() + DMRG_solver_kwargs_ = ( + {} if DMRG_solver_kwargs is None else DMRG_solver_kwargs.copy() + ) frag_scratch = WorkDir(scratch_dir / fobj.dname) try: rdm1_tmp, rdm2s = solve_block2( - fobj._mf, fobj.nsocc, frag_scratch=frag_scratch, **solver_kwargs_ + fobj._mf, + fobj.nsocc, + frag_scratch=frag_scratch, + **DMRG_solver_kwargs_, ) except Exception as inst: raise inst finally: - if solver_kwargs_.pop("force_cleanup", False): + if DMRG_solver_kwargs_.pop("force_cleanup", False): delete_multiple_files( frag_scratch.path.glob("F.*"), frag_scratch.path.glob("FCIDUMP*"), @@ -689,7 +694,7 @@ def solve_ccsd( return (t1, t2) -def solve_block2(mf, nocc, frag_scratch, **solver_kwargs): +def solve_block2(mf, nocc, frag_scratch, **DMRG_solver_kwargs): """DMRG fragment solver using the pyscf.dmrgscf wrapper. Parameters @@ -750,20 +755,22 @@ def solve_block2(mf, nocc, frag_scratch, **solver_kwargs): # pylint: disable-next=E0611 from pyscf import dmrgscf # noqa: PLC0415 # optional module - use_cumulant = solver_kwargs.pop("use_cumulant", True) - norb = solver_kwargs.pop("norb", mf.mo_coeff.shape[1]) - nelec = solver_kwargs.pop("nelec", mf.mo_coeff.shape[1]) - lo_method = solver_kwargs.pop("lo_method", None) - startM = solver_kwargs.pop("startM", 25) - maxM = solver_kwargs.pop("maxM", 500) - max_iter = solver_kwargs.pop("max_iter", 60) - max_mem = solver_kwargs.pop("max_mem", 100) - max_noise = solver_kwargs.pop("max_noise", 1e-3) - min_tol = solver_kwargs.pop("min_tol", 1e-8) - twodot_to_onedot = solver_kwargs.pop("twodot_to_onedot", int((5 * max_iter) // 6)) - root = solver_kwargs.pop("root", 0) - block_extra_keyword = solver_kwargs.pop("block_extra_keyword", ["fiedler"]) - schedule_kwargs = solver_kwargs.pop("schedule_kwargs", {}) + use_cumulant = DMRG_solver_kwargs.pop("use_cumulant", True) + norb = DMRG_solver_kwargs.pop("norb", mf.mo_coeff.shape[1]) + nelec = DMRG_solver_kwargs.pop("nelec", mf.mo_coeff.shape[1]) + lo_method = DMRG_solver_kwargs.pop("lo_method", None) + startM = DMRG_solver_kwargs.pop("startM", 25) + maxM = DMRG_solver_kwargs.pop("maxM", 500) + max_iter = DMRG_solver_kwargs.pop("max_iter", 60) + max_mem = DMRG_solver_kwargs.pop("max_mem", 100) + max_noise = DMRG_solver_kwargs.pop("max_noise", 1e-3) + min_tol = DMRG_solver_kwargs.pop("min_tol", 1e-8) + twodot_to_onedot = DMRG_solver_kwargs.pop( + "twodot_to_onedot", int((5 * max_iter) // 6) + ) + root = DMRG_solver_kwargs.pop("root", 0) + block_extra_keyword = DMRG_solver_kwargs.pop("block_extra_keyword", ["fiedler"]) + schedule_kwargs = DMRG_solver_kwargs.pop("schedule_kwargs", {}) if norb <= 2: block_extra_keyword = [ From ece5156cdaf033db19090d337813dd8f662b3868 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 18:26:25 -0500 Subject: [PATCH 40/67] better naming for DMRG stuff --- src/quemb/molbe/solver.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/quemb/molbe/solver.py b/src/quemb/molbe/solver.py index c7c5fe87..ad733d96 100644 --- a/src/quemb/molbe/solver.py +++ b/src/quemb/molbe/solver.py @@ -215,22 +215,22 @@ def be_func( rdm1_tmp, rdm2s = ci.make_rdm12(0, nmo, nelec) elif solver in ["block2", "DMRG", "DMRGCI", "DMRGSCF"]: - DMRG_solver_kwargs_ = ( - {} if DMRG_solver_kwargs is None else DMRG_solver_kwargs.copy() - ) frag_scratch = WorkDir(scratch_dir / fobj.dname) + DMRG_solver_kwargs = ( + {} if DMRG_solver_kwargs is None else DMRG_solver_kwargs.copy() + ) try: rdm1_tmp, rdm2s = solve_block2( fobj._mf, fobj.nsocc, frag_scratch=frag_scratch, - **DMRG_solver_kwargs_, + DMRG_solver_kwargs=DMRG_solver_kwargs, ) except Exception as inst: raise inst finally: - if DMRG_solver_kwargs_.pop("force_cleanup", False): + if DMRG_solver_kwargs.pop("force_cleanup", False): delete_multiple_files( frag_scratch.path.glob("F.*"), frag_scratch.path.glob("FCIDUMP*"), @@ -694,7 +694,7 @@ def solve_ccsd( return (t1, t2) -def solve_block2(mf, nocc, frag_scratch, **DMRG_solver_kwargs): +def solve_block2(mf, nocc, frag_scratch, DMRG_solver_kwargs: KwargDict): """DMRG fragment solver using the pyscf.dmrgscf wrapper. Parameters @@ -825,7 +825,7 @@ def solve_block2(mf, nocc, frag_scratch, **DMRG_solver_kwargs): mc.fcisolver.scratchDirectory = str(frag_scratch) mc.fcisolver.runtimeDir = str(frag_scratch) mc.fcisolver.memory = int(max_mem) - os.system("cd " + frag_scratch) + os.chdir(frag_scratch) mc.kernel(orbs) rdm1, rdm2 = dmrgscf.DMRGCI.make_rdm12(mc.fcisolver, root, norb, nelec) From 2a7514a35bc6e368e1420cb1a1f4d9c734ebd28e Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 18:29:01 -0500 Subject: [PATCH 41/67] added types to block2 DMRG function --- src/quemb/molbe/solver.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/quemb/molbe/solver.py b/src/quemb/molbe/solver.py index ad733d96..5cf81f16 100644 --- a/src/quemb/molbe/solver.py +++ b/src/quemb/molbe/solver.py @@ -7,6 +7,7 @@ from numpy.linalg import multi_dot from pyscf import ao2mo, cc, fci, mcscf, mp from pyscf.cc.ccsd_rdm import make_rdm2 +from pyscf.scf.hf import RHF from quemb.kbe.pfrag import Frags as pFrags from quemb.molbe.helper import get_frag_energy, get_frag_energy_u @@ -694,7 +695,9 @@ def solve_ccsd( return (t1, t2) -def solve_block2(mf, nocc, frag_scratch, DMRG_solver_kwargs: KwargDict): +def solve_block2( + mf: RHF, nocc: int, frag_scratch: WorkDir, DMRG_solver_kwargs: KwargDict +): """DMRG fragment solver using the pyscf.dmrgscf wrapper. Parameters From fc4dda7b2c6297cd0851eab70d11582878b8ded2 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 20 Dec 2024 18:44:59 -0500 Subject: [PATCH 42/67] refactor some DMRG arguments --- src/quemb/molbe/solver.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/src/quemb/molbe/solver.py b/src/quemb/molbe/solver.py index 5cf81f16..7c4c252a 100644 --- a/src/quemb/molbe/solver.py +++ b/src/quemb/molbe/solver.py @@ -227,6 +227,7 @@ def be_func( fobj.nsocc, frag_scratch=frag_scratch, DMRG_solver_kwargs=DMRG_solver_kwargs, + use_cumulant=use_cumulant, ) except Exception as inst: raise inst @@ -696,7 +697,11 @@ def solve_ccsd( def solve_block2( - mf: RHF, nocc: int, frag_scratch: WorkDir, DMRG_solver_kwargs: KwargDict + mf: RHF, + nocc: int, + frag_scratch: WorkDir, + DMRG_solver_kwargs: KwargDict, + use_cumulant: bool, ): """DMRG fragment solver using the pyscf.dmrgscf wrapper. @@ -758,9 +763,9 @@ def solve_block2( # pylint: disable-next=E0611 from pyscf import dmrgscf # noqa: PLC0415 # optional module - use_cumulant = DMRG_solver_kwargs.pop("use_cumulant", True) norb = DMRG_solver_kwargs.pop("norb", mf.mo_coeff.shape[1]) nelec = DMRG_solver_kwargs.pop("nelec", mf.mo_coeff.shape[1]) + lo_method = DMRG_solver_kwargs.pop("lo_method", None) startM = DMRG_solver_kwargs.pop("startM", 25) maxM = DMRG_solver_kwargs.pop("maxM", 500) @@ -782,7 +787,7 @@ def solve_block2( if lo_method is None: orbs = mf.mo_coeff - elif isinstance(lo_method, str): + else: raise NotImplementedError( "Localization within the fragment+bath subspace is currently not supported." ) From 4d04279c5f938d4b88dc0d2e96cd096a18aa8d74 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Sat, 21 Dec 2024 14:00:32 -0500 Subject: [PATCH 43/67] change behaviour of scratch contextmanager now it is ensured, that files are deleted even after an exception when using it as context manager and cleanup_at_end=True --- src/quemb/shared/manage_scratch.py | 2 +- tests/scratch_manager_test.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/quemb/shared/manage_scratch.py b/src/quemb/shared/manage_scratch.py index 2d2565eb..31db2397 100644 --- a/src/quemb/shared/manage_scratch.py +++ b/src/quemb/shared/manage_scratch.py @@ -105,7 +105,7 @@ def __exit__( value: BaseException | None, traceback: TracebackType | None, ) -> Literal[False]: - if value is None and self.cleanup_at_end: + if self.cleanup_at_end: self.cleanup() return False diff --git a/tests/scratch_manager_test.py b/tests/scratch_manager_test.py index 92c2603b..a67c8559 100644 --- a/tests/scratch_manager_test.py +++ b/tests/scratch_manager_test.py @@ -27,7 +27,7 @@ def test_keep_upon_error() -> None: with raises(ValueError): with WorkDir(my_tmp): raise ValueError - assert my_tmp.exists() + assert not my_tmp.exists() with WorkDir(my_tmp): pass From 9e7b26f80e976371f9c33885bde66f65c1cf1f55 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Sat, 21 Dec 2024 14:02:33 -0500 Subject: [PATCH 44/67] fixed the deadlock in the test requirements --- .github/workflows/quemb_unittest.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/quemb_unittest.yml b/.github/workflows/quemb_unittest.yml index ef846291..733e54d5 100644 --- a/.github/workflows/quemb_unittest.yml +++ b/.github/workflows/quemb_unittest.yml @@ -15,7 +15,6 @@ permissions: jobs: analysis: runs-on: ubuntu-latest - needs: analysis strategy: matrix: python-version: ["3.10"] @@ -80,6 +79,7 @@ jobs: testsuite: runs-on: ubuntu-latest + needs: analysis strategy: matrix: python-version: ["3.10", "3.13"] From 69755ef401fc06262c7b0518699e6e57fffe0389 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 23 Dec 2024 11:19:08 -0500 Subject: [PATCH 45/67] added type annotations to lo.py --- src/quemb/molbe/lo.py | 83 ++++++++++++++++++++++++++++--------------- 1 file changed, 55 insertions(+), 28 deletions(-) diff --git a/src/quemb/molbe/lo.py b/src/quemb/molbe/lo.py index 539c6282..8dcfaee9 100644 --- a/src/quemb/molbe/lo.py +++ b/src/quemb/molbe/lo.py @@ -2,37 +2,47 @@ # import numpy +from numpy import float64 from numpy.linalg import eigh, inv, multi_dot, norm, svd from pyscf.gto import intor_cross +from pyscf.gto.mole import Mole from quemb.shared.external.lo_helper import ( get_aoind_by_atom, reorder_by_atom_, ) from quemb.shared.helper import ncore_, unused +from quemb.shared.typing import Matrix, Tensor3D -def dot_gen(A, B, ovlp): +def dot_gen( + A: Matrix[float64], B: Matrix[float64], ovlp: Matrix[float64] | None = None +) -> Matrix[float64]: return A.T @ B if ovlp is None else A.T @ ovlp @ B -def get_cano_orth_mat(A, thr=1.0e-6, ovlp=None): +def get_cano_orth_mat( + A: Matrix[float64], thr: float = 1.0e-6, ovlp: Matrix[float64] | None = None +) -> Matrix[float64]: S = dot_gen(A, A, ovlp) e, u = eigh(S) if thr > 0: idx_keep = e / e[-1] > thr else: - idx_keep = list(range(e.shape[0])) - U = u[:, idx_keep] * e[idx_keep] ** -0.5 - return U + idx_keep = slice(0, e.shape[0]) + return u[:, idx_keep] * e[idx_keep] ** -0.5 -def cano_orth(A, thr=1.0e-6, ovlp=None): +def cano_orth( + A: Matrix[float64], thr: float = 1.0e-6, ovlp: Matrix[float64] | None = None +) -> Matrix[float64]: """Canonically orthogonalize columns of A""" return A @ get_cano_orth_mat(A, thr, ovlp) -def get_symm_orth_mat(A, thr=1.0e-6, ovlp=None): +def get_symm_orth_mat( + A: Matrix[float64], thr: float = 1.0e-6, ovlp: Matrix[float64] | None = None +) -> Matrix[float64]: S = dot_gen(A, A, ovlp) e, u = eigh(S) if (e < thr).any(): @@ -44,12 +54,16 @@ def get_symm_orth_mat(A, thr=1.0e-6, ovlp=None): return u @ numpy.diag(e**-0.5) @ u.T -def symm_orth(A, thr=1.0e-6, ovlp=None): +def symm_orth( + A: Matrix[float64], thr: float = 1.0e-6, ovlp: Matrix[float64] | None = None +) -> Matrix[float64]: """Symmetrically orthogonalize columns of A""" return A @ get_symm_orth_mat(A, thr, ovlp) -def remove_core_mo(Clo, Ccore, S, thr=0.5): +def remove_core_mo( + Clo: Matrix[float64], Ccore: Matrix[float64], S: Matrix[float64], thr: float = 0.5 +) -> Matrix[float64]: assert numpy.allclose(Clo.T @ S @ Clo, numpy.eye(Clo.shape[1])) assert numpy.allclose(Ccore.T @ S @ Ccore, numpy.eye(Ccore.shape[1])) @@ -63,14 +77,16 @@ def remove_core_mo(Clo, Ccore, S, thr=0.5): return symm_orth(Clo1[:, idx_keep], ovlp=S) -def get_xovlp(mol, basis="sto-3g"): +def get_xovlp( + mol: Mole, basis: str = "sto-3g" +) -> tuple[Matrix[float64] | Tensor3D[float64], Matrix[float64] | Tensor3D[float64]]: """Gets set of valence orbitals based on smaller (should be minimal) basis Parameters ---------- - mol : pyscf.gto.mole.Mole + mol : just need it for the working basis - basis : str + basis : the IAO basis, Knizia recommended 'minao' Returns @@ -90,7 +106,12 @@ def get_xovlp(mol, basis="sto-3g"): return S12, S22 -def get_iao(Co, S12, S1, S2=None): +def get_iao( + Co: Matrix[float64], + S12: Matrix[float64], + S1: Matrix[float64], + S2: Matrix[float64] | None = None, +) -> Matrix[float64]: """ Parameters @@ -136,19 +157,21 @@ def get_iao(Co, S12, S1, S2=None): return Ciao -def get_pao(Ciao, S, S12): +def get_pao( + Ciao: Matrix[float64], S: Matrix[float64], S12: Matrix[float64] +) -> Matrix[float64]: """ Parameters ---------- - Ciao: numpy.ndarray + Ciao: output of :func:`get_iao` - S: numpy.ndarray + S: ao ovlp matrix - S12: numpy.ndarray + S12: valence orbitals projected into ao basis Returns ------- - Cpao: numpy.ndarray + Cpao: (orthogonalized) """ n = Ciao.shape[0] @@ -161,12 +184,12 @@ def get_pao(Ciao, S, S12): Cpao = (numpy.eye(n) - Piao) @ nonval # project out IAOs from non-valence basis # begin canonical orthogonalization to get rid of redundant orbitals - Cpao = cano_orth(Cpao, ovlp=S) - - return Cpao + return cano_orth(Cpao, ovlp=S) -def get_pao_native(Ciao, S, mol, valence_basis): +def get_pao_native( + Ciao: Matrix[float64], S: Matrix[float64], mol: Mole, valence_basis: str +) -> Matrix[float64]: """ Parameters @@ -181,7 +204,7 @@ def get_pao_native(Ciao, S, mol, valence_basis): basis used for valence orbitals Returns ------- - Cpao: numpy.ndarray + Cpao: (symmetrically orthogonalized) """ @@ -217,7 +240,13 @@ def get_pao_native(Ciao, S, mol, valence_basis): return Cpao -def get_loc(mol, C, method, pop_method=None, init_guess=None): +def get_loc( + mol: Mole, + C: Matrix[float64], + method: str, + pop_method: str | None = None, + init_guess: Matrix[float64] | None = None, +) -> Mole: if method.upper() == "ER": from pyscf.lo import ER as Localizer # noqa: PLC0415 elif method.upper() == "PM": @@ -229,12 +258,10 @@ def get_loc(mol, C, method, pop_method=None, init_guess=None): mlo = Localizer(mol, C) if pop_method is not None: - mlo.pop_method = str(pop_method) + mlo.pop_method = pop_method mlo.init_guess = init_guess - C_ = mlo.kernel() - - return C_ + return mlo.kernel() class MixinLocalize: From f1d316b9e310b532efeebdb49de494ade505271f Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 23 Dec 2024 11:38:38 -0500 Subject: [PATCH 46/67] added new scratch dir also to ube --- src/quemb/kbe/pbe.py | 4 ++-- src/quemb/molbe/mbe.py | 12 ++++++------ src/quemb/molbe/ube.py | 41 +++++++++++++++++++---------------------- 3 files changed, 27 insertions(+), 30 deletions(-) diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py index d24aeb70..9c2b64f6 100644 --- a/src/quemb/kbe/pbe.py +++ b/src/quemb/kbe/pbe.py @@ -70,7 +70,7 @@ def __init__( cderi: PathLike | None = None, iao_wannier: bool = False, scratch_dir: WorkDir | None = None, - ): + ) -> None: """ Constructor for BE object. @@ -468,7 +468,7 @@ def ewald_sum(self): return e_.real - def initialize(self, compute_hf: bool, restart: bool = False): + def initialize(self, compute_hf: bool, restart: bool = False) -> None: """ Initialize the Bootstrap Embedding calculation. diff --git a/src/quemb/molbe/mbe.py b/src/quemb/molbe/mbe.py index be61afa8..6ba79183 100644 --- a/src/quemb/molbe/mbe.py +++ b/src/quemb/molbe/mbe.py @@ -758,7 +758,7 @@ def print_ini(self): print("-----------------------------------------------------------", flush=True) print(flush=True) - def initialize(self, eri_, compute_hf, restart=False): + def initialize(self, eri_, compute_hf, restart=False) -> None: """ Initialize the Bootstrap Embedding calculation. @@ -845,7 +845,7 @@ def initialize(self, eri_, compute_hf, restart=False): else: # Calculate ERIs on-the-fly to generate fragment ERIs # TODO: Future feature to be implemented # NOTE: Ideally, we want AO shell pair screening for this. - return NotImplementedError + raise NotImplementedError else: eri = None @@ -903,7 +903,7 @@ def oneshot( ompnum: int = 4, calc_frag_energy: bool = False, DMRG_solver_kwargs: KwargDict | None = None, - ): + ) -> None: """ Perform a one-shot bootstrap embedding calculation. @@ -983,7 +983,7 @@ def oneshot( if not calc_frag_energy: self.compute_energy_full(approx_cumulant=True, return_rdm=False) - def update_fock(self, heff=None): + def update_fock(self, heff: list[Matrix[float64]] | None = None) -> None: """ Update the Fock matrix for each fragment with the effective Hamiltonian. @@ -996,10 +996,10 @@ def update_fock(self, heff=None): for fobj in self.Fobjs: fobj.fock += fobj.heff else: - for idx, fobj in self.Fobjs: + for idx, fobj in enumerate(self.Fobjs): fobj.fock += heff[idx] - def write_heff(self, heff_file="bepotfile.h5"): + def write_heff(self, heff_file: str = "bepotfile.h5") -> None: """ Write the effective Hamiltonian to a file. diff --git a/src/quemb/molbe/ube.py b/src/quemb/molbe/ube.py index 11b04aaf..db83edc0 100644 --- a/src/quemb/molbe/ube.py +++ b/src/quemb/molbe/ube.py @@ -12,29 +12,33 @@ Add iterative UBE """ -import os +from pathlib import Path import h5py import numpy from pyscf import ao2mo +from pyscf.scf.uhf import UHF from quemb.molbe.be_parallel import be_func_parallel_u +from quemb.molbe.fragment import fragpart from quemb.molbe.mbe import BE from quemb.molbe.pfrag import Frags from quemb.molbe.solver import be_func_u -from quemb.shared.config import settings from quemb.shared.helper import unused +from quemb.shared.manage_scratch import WorkDir +from quemb.shared.typing import PathLike class UBE(BE): # 🍠 def __init__( self, - mf, - fobj, - eri_file="eri_file.h5", - lo_method="lowdin", - compute_hf=True, - ): + mf: UHF, + fobj: fragpart, + scratch_dir: WorkDir | None = None, + eri_file: PathLike = "eri_file.h5", + lo_method: PathLike = "lowdin", + compute_hf: bool = True, + ) -> None: """Initialize Unrestricted BE Object (ube🍠) .. note:: @@ -91,12 +95,12 @@ def __init__( self.print_ini() - self.Fobjs_a = [] - self.Fobjs_b = [] + self.Fobjs_a: list[Frags] = [] + self.Fobjs_b: list[Frags] = [] self.pot = initialize_pot(self.Nfrag, self.edge_idx) - self.eri_file = eri_file + self.eri_file = Path(eri_file) self.ek = 0.0 self.frozen_core = fobj.frozen_core self.ncore = 0 @@ -153,18 +157,11 @@ def __init__( valence_only=fobj.valence_only, ) - jobid = "" - if settings.CREATE_SCRATCH_DIR: - jobid = os.environ.get("SLURM_JOB_ID", "") - if settings.SCRATCH: - self.scratch_dir = settings.SCRATCH + str(jobid) - os.system("mkdir -p " + self.scratch_dir) + if scratch_dir is None: + self.scratch_dir = WorkDir.from_environment() else: - self.scratch_dir = None - if not jobid: - self.eri_file = settings.SCRATCH + eri_file - else: - self.eri_file = self.scratch_dir + "/" + eri_file + self.scratch_dir = scratch_dir + self.eri_file = self.scratch_dir / eri_file self.initialize(mf._eri, compute_hf) From e818f4569a496bbd037dd79fcd90e75831ad2064 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 23 Dec 2024 12:15:33 -0500 Subject: [PATCH 47/67] removed some run_solver args that are always true --- src/quemb/molbe/be_parallel.py | 112 ++++++++++++++------------------- 1 file changed, 46 insertions(+), 66 deletions(-) diff --git a/src/quemb/molbe/be_parallel.py b/src/quemb/molbe/be_parallel.py index ec23f505..cefbd587 100644 --- a/src/quemb/molbe/be_parallel.py +++ b/src/quemb/molbe/be_parallel.py @@ -51,8 +51,6 @@ def run_solver( select_cutoff: float | None = None, ompnum: int = 4, writeh1: bool = False, - eeval: bool = True, - return_rdm_ao: bool = True, use_cumulant: bool = True, relax_density: bool = False, frag_energy: bool = False, @@ -96,11 +94,6 @@ def run_solver( Number of OpenMP threads. Default is 4. writeh1 : If True, write the one-electron integrals to a file. Default is False. - eeval : - If True, evaluate the electronic energy. Default is True. - return_rdm_ao : - If True, return the reduced density matrices in the atomic orbital basis. - Default is True. use_cumulant : If True, use the cumulant approximation for RDM2. Default is True. frag_energy : @@ -120,10 +113,6 @@ def run_solver( eri = get_eri(dname, nao, eri_file=eri_file) # Initialize SCF object mf_ = get_scfObj(h1, eri, nocc, dm0=dm0) - rdm_return = False - - if relax_density: - rdm_return = True # Select solver if solver == "MP2": @@ -131,7 +120,7 @@ def run_solver( rdm1_tmp = mc_.make_rdm1() elif solver == "CCSD": - if not rdm_return: + if not relax_density: t1, t2 = solve_ccsd(mf_, mo_energy=mf_.mo_energy, rdm_return=False) rdm1_tmp = make_rdm1_ccsd_t1(t1) else: @@ -227,54 +216,51 @@ def run_solver( # Compute RDM1 rdm1 = multi_dot((mf_.mo_coeff, rdm1_tmp, mf_.mo_coeff.T)) * 0.5 - if eeval: - if solver == "CCSD" and not rdm_return: - with_dm1 = True - if use_cumulant: - with_dm1 = False - rdm2s = make_rdm2_urlx(t1, t2, with_dm1=with_dm1) - elif solver == "MP2": - rdm2s = mc_.make_rdm2() - elif solver == "FCI": - rdm2s = mc_.make_rdm2(civec, mc_.norb, mc_.nelec) - if use_cumulant: - hf_dm = numpy.zeros_like(rdm1_tmp) - hf_dm[numpy.diag_indices(nocc)] += 2.0 - del_rdm1 = rdm1_tmp.copy() - del_rdm1[numpy.diag_indices(nocc)] -= 2.0 - nc = ( - numpy.einsum("ij,kl->ijkl", hf_dm, hf_dm) - + numpy.einsum("ij,kl->ijkl", hf_dm, del_rdm1) - + numpy.einsum("ij,kl->ijkl", del_rdm1, hf_dm) - ) - nc -= ( - numpy.einsum("ij,kl->iklj", hf_dm, hf_dm) - + numpy.einsum("ij,kl->iklj", hf_dm, del_rdm1) - + numpy.einsum("ij,kl->iklj", del_rdm1, hf_dm) - ) * 0.5 - rdm2s -= nc - e_f = get_frag_energy( - mf_.mo_coeff, - nocc, - nfsites, - efac, - TA, - h1_e, - hf_veff, - rdm1_tmp, - rdm2s, - dname, - eri_file, - veff0, - ) - if frag_energy: - return e_f + if solver == "CCSD" and not relax_density: + with_dm1 = True + if use_cumulant: + with_dm1 = False + rdm2s = make_rdm2_urlx(t1, t2, with_dm1=with_dm1) - if return_rdm_ao: - return (e_f, mf_.mo_coeff, rdm1, rdm2s, rdm1_tmp) + elif solver == "MP2": + rdm2s = mc_.make_rdm2() + elif solver == "FCI": + rdm2s = mc_.make_rdm2(civec, mc_.norb, mc_.nelec) + if use_cumulant: + hf_dm = numpy.zeros_like(rdm1_tmp) + hf_dm[numpy.diag_indices(nocc)] += 2.0 + del_rdm1 = rdm1_tmp.copy() + del_rdm1[numpy.diag_indices(nocc)] -= 2.0 + nc = ( + numpy.einsum("ij,kl->ijkl", hf_dm, hf_dm) + + numpy.einsum("ij,kl->ijkl", hf_dm, del_rdm1) + + numpy.einsum("ij,kl->ijkl", del_rdm1, hf_dm) + ) + nc -= ( + numpy.einsum("ij,kl->iklj", hf_dm, hf_dm) + + numpy.einsum("ij,kl->iklj", hf_dm, del_rdm1) + + numpy.einsum("ij,kl->iklj", del_rdm1, hf_dm) + ) * 0.5 + rdm2s -= nc + e_f = get_frag_energy( + mf_.mo_coeff, + nocc, + nfsites, + efac, + TA, + h1_e, + hf_veff, + rdm1_tmp, + rdm2s, + dname, + eri_file, + veff0, + ) + if frag_energy: + return e_f - return (e_f, mf_.mo_coeff, rdm1, rdm2s) + return (e_f, mf_.mo_coeff, rdm1, rdm2s, rdm1_tmp) def run_solver_u( @@ -328,12 +314,8 @@ def run_solver_u( # Construct UHF object full_uhf, eris = make_uhf_obj(fobj_a, fobj_b, frozen=frozen) - rdm_return = False - if relax_density: - rdm_return = True - if solver == "UCCSD": - if rdm_return: + if relax_density: ucc, rdm1_tmp, rdm2s = solve_uccsd( full_uhf, eris, @@ -363,7 +345,7 @@ def run_solver_u( # Calculate Energies if ereturn: - if solver == "UCCSD" and not rdm_return: + if solver == "UCCSD" and not relax_density: with_dm1 = True if use_cumulant: with_dm1 = False @@ -400,7 +382,7 @@ def run_solver_u( ) return e_f else: - return NotImplementedError("Energy only calculated on a per-fragment basis") + raise NotImplementedError("Energy only calculated on a per-fragment basis") def be_func_parallel( @@ -507,8 +489,6 @@ def be_func_parallel( select_cutoff, ompnum, writeh1, - True, - True, use_cumulant, relax_density, frag_energy, From 57a2c11664f0aaaeec51e90ad05116f8a1eb48e1 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 23 Dec 2024 12:27:30 -0500 Subject: [PATCH 48/67] simplified boolean expressions in molbe --- src/quemb/molbe/be_parallel.py | 10 ++-------- src/quemb/molbe/misc.py | 8 +++----- src/quemb/molbe/solver.py | 28 +++++++++------------------- 3 files changed, 14 insertions(+), 32 deletions(-) diff --git a/src/quemb/molbe/be_parallel.py b/src/quemb/molbe/be_parallel.py index cefbd587..049abd00 100644 --- a/src/quemb/molbe/be_parallel.py +++ b/src/quemb/molbe/be_parallel.py @@ -218,10 +218,7 @@ def run_solver( rdm1 = multi_dot((mf_.mo_coeff, rdm1_tmp, mf_.mo_coeff.T)) * 0.5 if solver == "CCSD" and not relax_density: - with_dm1 = True - if use_cumulant: - with_dm1 = False - rdm2s = make_rdm2_urlx(t1, t2, with_dm1=with_dm1) + rdm2s = make_rdm2_urlx(t1, t2, with_dm1=not use_cumulant) elif solver == "MP2": rdm2s = mc_.make_rdm2() @@ -346,10 +343,7 @@ def run_solver_u( # Calculate Energies if ereturn: if solver == "UCCSD" and not relax_density: - with_dm1 = True - if use_cumulant: - with_dm1 = False - rdm2s = make_rdm2_uccsd(ucc, with_dm1=with_dm1) + rdm2s = make_rdm2_uccsd(ucc, with_dm1=not use_cumulant) else: raise NotImplementedError("RDM Return not Implemented") diff --git a/src/quemb/molbe/misc.py b/src/quemb/molbe/misc.py index be15b78c..204aaa30 100644 --- a/src/quemb/molbe/misc.py +++ b/src/quemb/molbe/misc.py @@ -360,8 +360,7 @@ def be2puffin( mol.incore_anyway = True if unrestricted: if use_df and jk is None: - print("UHF and df are incompatible: use_df = False") - use_df = False + raise ValueError("UHF and df are incompatible: use_df = False") if hcore is None: if pts_and_charges: print( @@ -393,12 +392,11 @@ def be2puffin( mf = qmmm.mm_charge( mf1, pts_and_charges[0], pts_and_charges[1], unit="bohr" ).newton() - print( + + raise ValueError( "Setting use_df to false and jk to none: have not tested DF and " "QM/MM from point charges at the same time" ) - use_df = False - jk = None elif use_df and jk is None: mf = scf.RHF(mol).density_fit(auxbasis=df_aux_basis) else: diff --git a/src/quemb/molbe/solver.py b/src/quemb/molbe/solver.py index 7c4c252a..3e427f2b 100644 --- a/src/quemb/molbe/solver.py +++ b/src/quemb/molbe/solver.py @@ -92,9 +92,6 @@ def be_func( Depending on the options, it returns the norm of the error vector, the energy, or a combination of these values. """ - rdm_return = False - if relax_density: - rdm_return = True if frag_energy or eeval: total_e = [0.0, 0.0, 0.0] @@ -113,7 +110,7 @@ def be_func( if solver == "MP2": fobj._mc = solve_mp2(fobj._mf, mo_energy=fobj._mf.mo_energy) elif solver == "CCSD": - if rdm_return: + if relax_density: fobj.t1, fobj.t2, rdm1_tmp, rdm2s = solve_ccsd( fobj._mf, mo_energy=fobj._mf.mo_energy, @@ -258,11 +255,8 @@ def be_func( ) if eeval or ereturn: - if solver == "CCSD" and not rdm_return: - with_dm1 = True - if use_cumulant: - with_dm1 = False - rdm2s = make_rdm2_urlx(fobj.t1, fobj.t2, with_dm1=with_dm1) + if solver == "CCSD" and not relax_density: + rdm2s = make_rdm2_urlx(fobj.t1, fobj.t2, with_dm1=not use_cumulant) elif solver == "MP2": rdm2s = fobj._mc.make_rdm2() elif solver == "FCI": @@ -375,7 +369,6 @@ def be_func_u( Depending on the options, it returns the norm of the error vector, the energy, or a combination of these values. """ - rdm_return = relax_density E = 0.0 if frag_energy or eeval: total_e = [0.0, 0.0, 0.0] @@ -388,7 +381,7 @@ def be_func_u( full_uhf, eris = make_uhf_obj(fobj_a, fobj_b, frozen=frozen) if solver == "UCCSD": - if rdm_return: + if relax_density: ucc, rdm1_tmp, rdm2s = solve_uccsd( full_uhf, eris, @@ -416,11 +409,8 @@ def be_func_u( ) if eeval or ereturn: - if solver == "UCCSD" and not rdm_return: - with_dm1 = True - if use_cumulant: - with_dm1 = False - rdm2s = make_rdm2_uccsd(ucc, with_dm1=with_dm1) + if solver == "UCCSD" and not relax_density: + rdm2s = make_rdm2_uccsd(ucc, with_dm1=not use_cumulant) fobj_a.rdm2__ = rdm2s[0].copy() fobj_b.rdm2__ = rdm2s[1].copy() if frag_energy: @@ -983,9 +973,9 @@ def ao2mofn(moish): if rdm_return: rdm1 = make_rdm1_uccsd(ucc, relax=relax) if rdm2_return: - if use_cumulant: - with_dm1 = False - rdm2 = make_rdm2_uccsd(ucc, relax=relax, with_dm1=with_dm1) + rdm2 = make_rdm2_uccsd( + ucc, relax=relax, with_dm1=with_dm1 and not use_cumulant + ) return (ucc, rdm1, rdm2) return (ucc, rdm1, None) return ucc From 07569c03a673adc341a0dfb2db22d88cb777b756 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 23 Dec 2024 12:31:44 -0500 Subject: [PATCH 49/67] simplified boolean expressions in kbe --- src/quemb/kbe/autofrag.py | 9 +---- src/quemb/kbe/lo.py | 81 ++++++++++++++++----------------------- 2 files changed, 34 insertions(+), 56 deletions(-) diff --git a/src/quemb/kbe/autofrag.py b/src/quemb/kbe/autofrag.py index a159fd67..384d5628 100644 --- a/src/quemb/kbe/autofrag.py +++ b/src/quemb/kbe/autofrag.py @@ -1924,10 +1924,7 @@ def autogen( ) w.close() - if valence_basis is not None: - pao = True - else: - pao = False + pao = valence_basis is not None if pao: cell2 = cell.copy() @@ -1994,10 +1991,8 @@ def autogen( centerf_idx = [] edge = [] - conmax = True nkcon = True - if gamma_2d or gamma_1d: - conmax = False + conmax = not (gamma_2d or gamma_1d) Ns = Nsite + 1 uNs = Ns * unitcell Sites = sites__.copy() diff --git a/src/quemb/kbe/lo.py b/src/quemb/kbe/lo.py index 488e45f1..89b85445 100644 --- a/src/quemb/kbe/lo.py +++ b/src/quemb/kbe/lo.py @@ -118,19 +118,15 @@ def localize( S12, S2 = get_xovlp_k(self.cell, self.kpts, basis=valence_basis) ciao_ = get_iao_k(Co, S12, self.S, S2=S2) - arrange_by_atom = True # tmp - aos are not rearrange and so below is not necessary - if arrange_by_atom: - nk, nao, nlo = ciao_.shape - Ciao_ = numpy.zeros((nk, nao, nlo), dtype=numpy.complex128) - for k in range(self.nkpt): - aoind_by_atom = get_aoind_by_atom(self.cell) - ctmp, iaoind_by_atom = reorder_by_atom_( - ciao_[k], aoind_by_atom, self.S[k] - ) - Ciao_[k] = ctmp - else: - Ciao_ = ciao_.copy() + nk, nao, nlo = ciao_.shape + Ciao_ = numpy.zeros((nk, nao, nlo), dtype=numpy.complex128) + for k in range(self.nkpt): + aoind_by_atom = get_aoind_by_atom(self.cell) + ctmp, iaoind_by_atom = reorder_by_atom_( + ciao_[k], aoind_by_atom, self.S[k] + ) + Ciao_[k] = ctmp # get_pao_k returns canonical orthogonalized orbitals # Cpao = get_pao_k(Ciao, self.S, S12, S2, self.cell) @@ -139,17 +135,14 @@ def localize( Ciao_, self.S, self.cell, valence_basis, self.kpts ) - if arrange_by_atom: - nk, nao, nlo = cpao_.shape - Cpao_ = numpy.zeros((nk, nao, nlo), dtype=numpy.complex128) - for k in range(self.nkpt): - aoind_by_atom = get_aoind_by_atom(self.cell) - ctmp, paoind_by_atom = reorder_by_atom_( - cpao_[k], aoind_by_atom, self.S[k] - ) - Cpao_[k] = ctmp - else: - Cpao_ = cpao_.copy() + nk, nao, nlo = cpao_.shape + Cpao_ = numpy.zeros((nk, nao, nlo), dtype=numpy.complex128) + for k in range(self.nkpt): + aoind_by_atom = get_aoind_by_atom(self.cell) + ctmp, paoind_by_atom = reorder_by_atom_( + cpao_[k], aoind_by_atom, self.S[k] + ) + Cpao_[k] = ctmp nk, nao, nlo = Ciao_.shape if self.frozen_core: @@ -225,29 +218,26 @@ def localize( for k in range(nk_): c_core_val[k] = numpy.hstack((ciao_core[k], Ciao_[k])) - arrange_by_atom = True # tmp - aos are not rearrange and so below is not necessary # (iaoind_by_atom is used to stack iao|pao later) - if arrange_by_atom: - nk, nao, nlo = c_core_val.shape - for k in range(self.nkpt): - aoind_by_atom = get_aoind_by_atom(self.cell) - ctmp, iaoind_by_atom = reorder_by_atom_( - c_core_val[k], aoind_by_atom, self.S[k] - ) + nk, nao, nlo = c_core_val.shape + for k in range(self.nkpt): + aoind_by_atom = get_aoind_by_atom(self.cell) + ctmp, iaoind_by_atom = reorder_by_atom_( + c_core_val[k], aoind_by_atom, self.S[k] + ) cpao_ = get_pao_native_k( c_core_val, self.S, self.cell, valence_basis, self.kpts, ortho=True ) - if arrange_by_atom: - nk, nao, nlo = cpao_.shape - Cpao_ = numpy.zeros((nk, nao, nlo), dtype=numpy.complex128) - for k in range(self.nkpt): - aoind_by_atom = get_aoind_by_atom(self.cell) - ctmp, paoind_by_atom = reorder_by_atom_( - cpao_[k], aoind_by_atom, self.S[k] - ) - Cpao_[k] = ctmp + nk, nao, nlo = cpao_.shape + Cpao_ = numpy.zeros((nk, nao, nlo), dtype=numpy.complex128) + for k in range(self.nkpt): + aoind_by_atom = get_aoind_by_atom(self.cell) + ctmp, paoind_by_atom = reorder_by_atom_( + cpao_[k], aoind_by_atom, self.S[k] + ) + Cpao_[k] = ctmp Cpao = Cpao_.copy() Ciao = Ciao_.copy() @@ -284,13 +274,8 @@ def localize( (self.nkpt, num_wann, num_wann), dtype=numpy.complex128 ) - i_init = True for k in range(self.nkpt): - if i_init: - A_matrix[k] = numpy.eye(num_wann, dtype=numpy.complex128) - else: - ovlp_ciao = Ciao[k].conj().T @ self.S[k] @ Ciao[k] - A_matrix[k] = ovlp_ciao + A_matrix[k] = numpy.eye(num_wann, dtype=numpy.complex128) A_matrix = A_matrix.transpose(1, 2, 0) w90.kernel(A_matrix=A_matrix) @@ -447,10 +432,8 @@ def localize( (self.nkpt, num_wann, num_wann), dtype=numpy.complex128 ) # Using A=I + lowdin orbital and A= + |psi> is the same - i_init = True for k in range(self.nkpt): - if i_init: - A_matrix[k] = numpy.eye(num_wann, dtype=numpy.complex128) + A_matrix[k] = numpy.eye(num_wann, dtype=numpy.complex128) A_matrix = A_matrix.transpose(1, 2, 0) From 0577e97e7bbaeb4f45510a70f851b3fee017827f Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 23 Dec 2024 13:45:11 -0500 Subject: [PATCH 50/67] introduced DMRG data class for solver args --- src/quemb/molbe/solver.py | 130 +++++++++++++++++++++++++++----------- 1 file changed, 94 insertions(+), 36 deletions(-) diff --git a/src/quemb/molbe/solver.py b/src/quemb/molbe/solver.py index 3e427f2b..a6c0d9ec 100644 --- a/src/quemb/molbe/solver.py +++ b/src/quemb/molbe/solver.py @@ -3,6 +3,7 @@ import os import numpy +from attrs import define from numpy import float64 from numpy.linalg import multi_dot from pyscf import ao2mo, cc, fci, mcscf, mp @@ -22,7 +23,78 @@ from quemb.shared.external.unrestricted_utils import make_uhf_obj from quemb.shared.helper import delete_multiple_files, unused from quemb.shared.manage_scratch import WorkDir -from quemb.shared.typing import KwargDict, Matrix +from quemb.shared.typing import Matrix + + +@define +class DMRG_ArgsUser: + #: Becomes mf.mo_coeff.shape[1] by default + norb: int | None = None + #: Becomes mf.mo_coeff.shape[1] by default + nelec: int | None = None + + startM: int = 25 + maxM: int = 500 + max_iter: int = 60 + max_mem: int = 100 + max_noise: float = 1e-3 + min_tol: float = 1e-8 + twodot_to_onedot: int = (5 * max_iter) // 6 + root: int = 0 + block_extra_keyword: list[str] = ["fiedler"] + schedule_kwargs: dict[str, list[int] | list[float]] = {} + force_cleanup: bool = False + + +@define +class DMRG_Args: + """Properly initialized DMRG arguments + + Some default values of :class:`DMRG_ArgsUser` can only be filled + later in the calculation. + Use :func:`from_user_input` to properly initialize. + """ + + norb: int + nelec: int + + startM: int + maxM: int + max_iter: int + max_mem: int + max_noise: float + min_tol: float + twodot_to_onedot: int + root: int + block_extra_keyword: list[str] + schedule_kwargs: dict[str, list[int] | list[float]] + force_cleanup: bool + + @classmethod + def from_user_input(cls, user_args: DMRG_ArgsUser, mf: RHF): + norb = mf.mo_coeff.shape[1] if user_args.norb is None else user_args.norb + nelec = mf.mo_coeff.shape[1] if user_args.nelec is None else user_args.nelec + if norb <= 2: + block_extra_keyword = [ + "noreorder" + ] # Other reordering algorithms explode if the network is too small. + else: + block_extra_keyword = user_args.block_extra_keyword + return cls( + norb=norb, + nelec=nelec, + startM=user_args.startM, + maxM=user_args.maxM, + max_iter=user_args.max_iter, + max_mem=user_args.max_mem, + max_noise=user_args.max_noise, + min_tol=user_args.min_tol, + twodot_to_onedot=user_args.twodot_to_onedot, + root=user_args.root, + block_extra_keyword=block_extra_keyword, + schedule_kwargs=user_args.schedule_kwargs, + force_cleanup=user_args.force_cleanup, + ) def be_func( @@ -31,7 +103,7 @@ def be_func( Nocc: int, solver: str, enuc: float, # noqa: ARG001 - DMRG_solver_kwargs: KwargDict | None, + user_DMRG_args: DMRG_ArgsUser, scratch_dir: WorkDir, hf_veff: Matrix[float64] | None = None, only_chem: bool = False, @@ -215,21 +287,20 @@ def be_func( elif solver in ["block2", "DMRG", "DMRGCI", "DMRGSCF"]: frag_scratch = WorkDir(scratch_dir / fobj.dname) - DMRG_solver_kwargs = ( - {} if DMRG_solver_kwargs is None else DMRG_solver_kwargs.copy() - ) + DMRG_args = DMRG_Args.from_user_input(user_DMRG_args, fobj._mf) + try: rdm1_tmp, rdm2s = solve_block2( fobj._mf, fobj.nsocc, frag_scratch=frag_scratch, - DMRG_solver_kwargs=DMRG_solver_kwargs, + DMRG_args=DMRG_args, use_cumulant=use_cumulant, ) except Exception as inst: raise inst finally: - if DMRG_solver_kwargs.pop("force_cleanup", False): + if DMRG_args.force_cleanup: delete_multiple_files( frag_scratch.path.glob("F.*"), frag_scratch.path.glob("FCIDUMP*"), @@ -690,7 +761,7 @@ def solve_block2( mf: RHF, nocc: int, frag_scratch: WorkDir, - DMRG_solver_kwargs: KwargDict, + DMRG_args: DMRG_ArgsUser, use_cumulant: bool, ): """DMRG fragment solver using the pyscf.dmrgscf wrapper. @@ -753,34 +824,21 @@ def solve_block2( # pylint: disable-next=E0611 from pyscf import dmrgscf # noqa: PLC0415 # optional module - norb = DMRG_solver_kwargs.pop("norb", mf.mo_coeff.shape[1]) - nelec = DMRG_solver_kwargs.pop("nelec", mf.mo_coeff.shape[1]) - - lo_method = DMRG_solver_kwargs.pop("lo_method", None) - startM = DMRG_solver_kwargs.pop("startM", 25) - maxM = DMRG_solver_kwargs.pop("maxM", 500) - max_iter = DMRG_solver_kwargs.pop("max_iter", 60) - max_mem = DMRG_solver_kwargs.pop("max_mem", 100) - max_noise = DMRG_solver_kwargs.pop("max_noise", 1e-3) - min_tol = DMRG_solver_kwargs.pop("min_tol", 1e-8) - twodot_to_onedot = DMRG_solver_kwargs.pop( - "twodot_to_onedot", int((5 * max_iter) // 6) - ) - root = DMRG_solver_kwargs.pop("root", 0) - block_extra_keyword = DMRG_solver_kwargs.pop("block_extra_keyword", ["fiedler"]) - schedule_kwargs = DMRG_solver_kwargs.pop("schedule_kwargs", {}) - - if norb <= 2: - block_extra_keyword = [ - "noreorder" - ] # Other reordering algorithms explode if the network is too small. - - if lo_method is None: - orbs = mf.mo_coeff - else: - raise NotImplementedError( - "Localization within the fragment+bath subspace is currently not supported." - ) + norb = DMRG_args.norb + nelec = DMRG_args.nelec + + startM = DMRG_args.startM + maxM = DMRG_args.maxM + max_iter = DMRG_args.max_iter + max_mem = DMRG_args.max_mem + max_noise = DMRG_args.max_noise + min_tol = DMRG_args.min_tol + twodot_to_onedot = DMRG_args.twodot_to_onedot + root = DMRG_args.root + schedule_kwargs = DMRG_args.schedule_kwargs + block_extra_keyword = DMRG_args.block_extra_keyword + + orbs = mf.mo_coeff mc = mcscf.CASCI(mf, norb, nelec) mc.fcisolver = dmrgscf.DMRGCI(mf.mol) From 9a7ea7a0435ab3c7d9802ff7fdb330ff5e847e2d Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 23 Dec 2024 15:15:17 -0500 Subject: [PATCH 51/67] Use a solver_args abstract class --- src/quemb/kbe/pbe.py | 8 ++++---- src/quemb/molbe/mbe.py | 12 ++++++------ src/quemb/molbe/opt.py | 8 ++++---- src/quemb/molbe/solver.py | 12 +++++++++--- 4 files changed, 23 insertions(+), 17 deletions(-) diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py index 9c2b64f6..33988d3e 100644 --- a/src/quemb/kbe/pbe.py +++ b/src/quemb/kbe/pbe.py @@ -18,13 +18,13 @@ from quemb.molbe.be_parallel import be_func_parallel from quemb.molbe.helper import get_eri, get_scfObj, get_veff from quemb.molbe.opt import BEOPT -from quemb.molbe.solver import be_func +from quemb.molbe.solver import UserSolverArgs, be_func from quemb.shared.external.optqn import ( get_be_error_jacobian as _ext_get_be_error_jacobian, ) from quemb.shared.helper import copy_docstring from quemb.shared.manage_scratch import WorkDir -from quemb.shared.typing import KwargDict, PathLike +from quemb.shared.typing import PathLike class BE(Mixin_k_Localize): @@ -676,7 +676,7 @@ def oneshot( nproc: int = 1, ompnum: int = 4, calc_frag_energy: bool = False, - DMRG_solver_kwargs: KwargDict | None = None, + solver_args: UserSolverArgs | None = None, ) -> None: """ Perform a one-shot bootstrap embedding calculation. @@ -713,7 +713,7 @@ def oneshot( ereturn=True, eeval=True, scratch_dir=self.scratch_dir, - DMRG_solver_kwargs=DMRG_solver_kwargs, + solver_args=solver_args, ) else: rets = be_func_parallel( diff --git a/src/quemb/molbe/mbe.py b/src/quemb/molbe/mbe.py index 6ba79183..21492552 100644 --- a/src/quemb/molbe/mbe.py +++ b/src/quemb/molbe/mbe.py @@ -15,13 +15,13 @@ from quemb.molbe.misc import print_energy from quemb.molbe.opt import BEOPT from quemb.molbe.pfrag import Frags -from quemb.molbe.solver import be_func +from quemb.molbe.solver import UserSolverArgs, be_func from quemb.shared.external.optqn import ( get_be_error_jacobian as _ext_get_be_error_jacobian, ) from quemb.shared.helper import copy_docstring from quemb.shared.manage_scratch import WorkDir -from quemb.shared.typing import KwargDict, Matrix, PathLike +from quemb.shared.typing import Matrix, PathLike @define @@ -645,7 +645,7 @@ def optimize( ompnum: int = 4, max_iter: int = 500, trust_region: bool = False, - DMRG_solver_kwargs: KwargDict | None = None, + solver_args: UserSolverArgs | None = None, ) -> None: """BE optimization function @@ -712,7 +712,7 @@ def optimize( hci_pt=self.hci_pt, solver=solver, ebe_hf=self.ebe_hf, - DMRG_solver_kwargs=DMRG_solver_kwargs, + solver_args=solver_args, ) if method == "QN": @@ -902,7 +902,7 @@ def oneshot( nproc: int = 1, ompnum: int = 4, calc_frag_energy: bool = False, - DMRG_solver_kwargs: KwargDict | None = None, + solver_args: UserSolverArgs | None = None, ) -> None: """ Perform a one-shot bootstrap embedding calculation. @@ -937,7 +937,7 @@ def oneshot( ereturn=True, eeval=True, scratch_dir=self.scratch_dir, - DMRG_solver_kwargs=DMRG_solver_kwargs, + solver_args=solver_args, ) else: rets = be_func_parallel( diff --git a/src/quemb/molbe/opt.py b/src/quemb/molbe/opt.py index 22be2b9e..95408b0a 100644 --- a/src/quemb/molbe/opt.py +++ b/src/quemb/molbe/opt.py @@ -7,10 +7,10 @@ from quemb.molbe.be_parallel import be_func_parallel from quemb.molbe.pfrag import Frags -from quemb.molbe.solver import be_func +from quemb.molbe.solver import UserSolverArgs, be_func from quemb.shared.external.optqn import FrankQN from quemb.shared.manage_scratch import WorkDir -from quemb.shared.typing import KwargDict, Matrix, Vector +from quemb.shared.typing import Matrix, Vector @define @@ -76,7 +76,7 @@ class BEOPT: err: float = 0.0 Ebe: Matrix[float64] = array([[0.0]]) - DMRG_solver_kwargs: KwargDict | None = None + solver_args: UserSolverArgs | None = None # HCI parameters hci_cutoff: float = 0.001 @@ -121,7 +121,7 @@ def objfunc(self, xk: list[float]) -> Vector[float64]: select_cutoff=self.select_cutoff, hci_pt=self.hci_pt, scratch_dir=self.scratch_dir, - DMRG_solver_kwargs=self.DMRG_solver_kwargs, + solver_args=self.solver_args, ) else: err_, errvec_, ebe_ = be_func_parallel( diff --git a/src/quemb/molbe/solver.py b/src/quemb/molbe/solver.py index a6c0d9ec..e45f39f0 100644 --- a/src/quemb/molbe/solver.py +++ b/src/quemb/molbe/solver.py @@ -1,6 +1,7 @@ # Author(s): Oinam Romesh Meitei, Leah Weisburn, Shaun Weatherly import os +from abc import ABC import numpy from attrs import define @@ -26,8 +27,12 @@ from quemb.shared.typing import Matrix +class UserSolverArgs(ABC): + pass + + @define -class DMRG_ArgsUser: +class DMRG_ArgsUser(UserSolverArgs): #: Becomes mf.mo_coeff.shape[1] by default norb: int | None = None #: Becomes mf.mo_coeff.shape[1] by default @@ -103,7 +108,7 @@ def be_func( Nocc: int, solver: str, enuc: float, # noqa: ARG001 - user_DMRG_args: DMRG_ArgsUser, + solver_args: UserSolverArgs | None, scratch_dir: WorkDir, hf_veff: Matrix[float64] | None = None, only_chem: bool = False, @@ -287,7 +292,8 @@ def be_func( elif solver in ["block2", "DMRG", "DMRGCI", "DMRGSCF"]: frag_scratch = WorkDir(scratch_dir / fobj.dname) - DMRG_args = DMRG_Args.from_user_input(user_DMRG_args, fobj._mf) + assert isinstance(solver_args, DMRG_ArgsUser) + DMRG_args = DMRG_Args.from_user_input(solver_args, fobj._mf) try: rdm1_tmp, rdm2s = solve_block2( From 997dc38e3a3a37f7c70f489395f87f0ff99d35ae Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 23 Dec 2024 15:47:48 -0500 Subject: [PATCH 52/67] finished SHCI solver args dataclass --- src/quemb/kbe/pbe.py | 17 +-------- src/quemb/molbe/be_parallel.py | 37 +++++++++--------- src/quemb/molbe/mbe.py | 21 +---------- src/quemb/molbe/opt.py | 14 +------ src/quemb/molbe/solver.py | 68 ++++++++++++++++++++++++++-------- 5 files changed, 75 insertions(+), 82 deletions(-) diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py index 33988d3e..4f8aa0ff 100644 --- a/src/quemb/kbe/pbe.py +++ b/src/quemb/kbe/pbe.py @@ -58,12 +58,8 @@ def __init__( save: bool = False, restart_file: PathLike = "storebe.pk", save_file: PathLike = "storebe.pk", - hci_pt: bool = False, nproc: int = 1, ompnum: int = 4, - hci_cutoff: float = 0.001, - ci_coeff_cutoff: float | None = None, - select_cutoff: float | None = None, iao_val_core: bool = True, exxdiv: str = "ewald", kpts: list[list[float]] | None = None, @@ -159,12 +155,6 @@ def __init__( self.nkpt = nkpts_ self.kpts = kpts - # HCI parameters - self.hci_cutoff = hci_cutoff - self.ci_coeff_cutoff = ci_coeff_cutoff - self.select_cutoff = select_cutoff - self.hci_pt = hci_pt - if not restart: self.mo_energy = mf.mo_energy mf.exxdiv = None @@ -705,9 +695,6 @@ def oneshot( solver, self.enuc, hf_veff=self.hf_veff, - hci_cutoff=self.hci_cutoff, - ci_coeff_cutoff=self.ci_coeff_cutoff, - select_cutoff=self.select_cutoff, nproc=ompnum, frag_energy=calc_frag_energy, ereturn=True, @@ -723,14 +710,12 @@ def oneshot( solver, self.enuc, hf_veff=self.hf_veff, - hci_cutoff=self.hci_cutoff, - ci_coeff_cutoff=self.ci_coeff_cutoff, - select_cutoff=self.select_cutoff, eeval=True, frag_energy=calc_frag_energy, nproc=nproc, ompnum=ompnum, scratch_dir=self.scratch_dir, + solver_args=solver_args, ) print("-----------------------------------------------------", flush=True) diff --git a/src/quemb/molbe/be_parallel.py b/src/quemb/molbe/be_parallel.py index 049abd00..ea9e8ce5 100644 --- a/src/quemb/molbe/be_parallel.py +++ b/src/quemb/molbe/be_parallel.py @@ -17,6 +17,9 @@ ) from quemb.molbe.pfrag import Frags from quemb.molbe.solver import ( + SHCI_Args, + SHCI_ArgsUser, + UserSolverArgs, make_rdm1_ccsd_t1, make_rdm2_urlx, solve_ccsd, @@ -46,14 +49,12 @@ def run_solver( solver: str = "MP2", eri_file: str = "eri_file.h5", veff0: Matrix[float64] | None = None, - hci_cutoff: float = 0.001, - ci_coeff_cutoff: float | None = None, - select_cutoff: float | None = None, ompnum: int = 4, writeh1: bool = False, use_cumulant: bool = True, relax_density: bool = False, frag_energy: bool = False, + solver_args: UserSolverArgs | None = None, ): """ Run a quantum chemistry solver to compute the reduced density matrices. @@ -141,19 +142,17 @@ def run_solver( # pylint: disable-next=E0611 from pyscf import hci # noqa: PLC0415 # hci is an optional module + assert isinstance(solver_args, SHCI_ArgsUser) + SHCI_args = SHCI_Args.from_user_input(solver_args) + nao, nmo = mf_.mo_coeff.shape eri = ao2mo.kernel(mf_._eri, mf_.mo_coeff, aosym="s4", compact=False).reshape( 4 * ((nmo),) ) ci_ = hci.SCI(mf_.mol) - if select_cutoff is None and ci_coeff_cutoff is None: - select_cutoff = hci_cutoff - ci_coeff_cutoff = hci_cutoff - elif select_cutoff is None or ci_coeff_cutoff is None: - raise ValueError - ci_.select_cutoff = select_cutoff - ci_.ci_coeff_cutoff = ci_coeff_cutoff + ci_.select_cutoff = SHCI_args.select_cutoff + ci_.ci_coeff_cutoff = SHCI_args.ci_coeff_cutoff nelec = (nocc, nocc) h1_ = multi_dot((mf_.mo_coeff.T, h1, mf_.mo_coeff)) @@ -171,6 +170,9 @@ def run_solver( frag_scratch = WorkDir(scratch_dir / dname) + assert isinstance(solver_args, SHCI_ArgsUser) + SHCI_args = SHCI_Args.from_user_input(solver_args) + nao, nmo = mf_.mo_coeff.shape nelec = (nocc, nocc) mch = shci.SHCISCF(mf_, nmo, nelec, orbpath=frag_scratch) @@ -179,7 +181,7 @@ def run_solver( mch.fcisolver.nPTiter = 0 mch.fcisolver.sweep_iter = [0] mch.fcisolver.DoRDM = True - mch.fcisolver.sweep_epsilon = [hci_cutoff] + mch.fcisolver.sweep_epsilon = [solver_args.hci_cutoff] mch.fcisolver.scratchDirectory = frag_scratch if not writeh1: mch.fcisolver.restart = True @@ -190,6 +192,9 @@ def run_solver( # pylint: disable-next=E0611 from pyscf import cornell_shci # noqa: PLC0415 # optional module + assert isinstance(solver_args, SHCI_ArgsUser) + SHCI_args = SHCI_Args.from_user_input(solver_args) + frag_scratch = WorkDir(scratch_dir / dname) nao, nmo = mf_.mo_coeff.shape @@ -205,7 +210,7 @@ def run_solver( ci.runtimedir = frag_scratch ci.restart = True ci.config["var_only"] = True - ci.config["eps_vars"] = [hci_cutoff] + ci.config["eps_vars"] = [solver_args.hci_cutoff] ci.config["get_1rdm_csv"] = True ci.config["get_2rdm_csv"] = True ci.kernel(h1, eri, nmo, nelec) @@ -386,6 +391,7 @@ def be_func_parallel( solver: str, enuc: float, # noqa: ARG001 scratch_dir: WorkDir, + solver_args: UserSolverArgs | None, hf_veff: Matrix[float64] | None = None, nproc: int = 1, ompnum: int = 4, @@ -394,9 +400,6 @@ def be_func_parallel( use_cumulant: bool = True, eeval: bool = False, frag_energy: bool = False, - hci_cutoff: float = 0.001, - ci_coeff_cutoff: float | None = None, - select_cutoff: float | None = None, return_vec: bool = False, writeh1: bool = False, ): @@ -478,14 +481,12 @@ def be_func_parallel( solver, fobj.eri_file, fobj.veff0, - hci_cutoff, - ci_coeff_cutoff, - select_cutoff, ompnum, writeh1, use_cumulant, relax_density, frag_energy, + solver_args, ], ) diff --git a/src/quemb/molbe/mbe.py b/src/quemb/molbe/mbe.py index 21492552..5c67cd4e 100644 --- a/src/quemb/molbe/mbe.py +++ b/src/quemb/molbe/mbe.py @@ -78,13 +78,9 @@ def __init__( save: bool = False, restart_file: PathLike = "storebe.pk", save_file: PathLike = "storebe.pk", - hci_pt: bool = False, nproc: int = 1, ompnum: int = 4, scratch_dir: WorkDir | None = None, - hci_cutoff: float = 0.001, - ci_coeff_cutoff: float | None = None, - select_cutoff: float | None = None, integral_direct_DF: bool = False, auxbasis: str | None = None, ) -> None: @@ -168,12 +164,6 @@ def __init__( self.ebe_hf = 0.0 self.ebe_tot = 0.0 - # HCI parameters - self.hci_cutoff = hci_cutoff - self.ci_coeff_cutoff = ci_coeff_cutoff - self.select_cutoff = select_cutoff - self.hci_pt = hci_pt - self.mf = mf if not restart: self.mo_energy = mf.mo_energy @@ -705,11 +695,7 @@ def optimize( max_space=max_iter, conv_tol=conv_tol, only_chem=only_chem, - hci_cutoff=self.hci_cutoff, - ci_coeff_cutoff=self.ci_coeff_cutoff, relax_density=relax_density, - select_cutoff=self.select_cutoff, - hci_pt=self.hci_pt, solver=solver, ebe_hf=self.ebe_hf, solver_args=solver_args, @@ -929,9 +915,6 @@ def oneshot( solver, self.enuc, hf_veff=self.hf_veff, - hci_cutoff=self.hci_cutoff, - ci_coeff_cutoff=self.ci_coeff_cutoff, - select_cutoff=self.select_cutoff, nproc=ompnum, frag_energy=calc_frag_energy, ereturn=True, @@ -947,14 +930,12 @@ def oneshot( solver, self.enuc, hf_veff=self.hf_veff, - hci_cutoff=self.hci_cutoff, - ci_coeff_cutoff=self.ci_coeff_cutoff, - select_cutoff=self.select_cutoff, eeval=True, frag_energy=calc_frag_energy, nproc=nproc, ompnum=ompnum, scratch_dir=self.scratch_dir, + solver_args=solver_args, ) print("-----------------------------------------------------", flush=True) diff --git a/src/quemb/molbe/opt.py b/src/quemb/molbe/opt.py index 95408b0a..db0bbe75 100644 --- a/src/quemb/molbe/opt.py +++ b/src/quemb/molbe/opt.py @@ -78,12 +78,6 @@ class BEOPT: solver_args: UserSolverArgs | None = None - # HCI parameters - hci_cutoff: float = 0.001 - ci_coeff_cutoff: float | None = None - select_cutoff: float | None = None - hci_pt: bool = False - def objfunc(self, xk: list[float]) -> Vector[float64]: """ Computes error vectors, RMS error, and BE energies. @@ -114,12 +108,8 @@ def objfunc(self, xk: list[float]) -> Vector[float64]: return_vec=True, hf_veff=self.hf_veff, only_chem=self.only_chem, - hci_cutoff=self.hci_cutoff, nproc=self.ompnum, relax_density=self.relax_density, - ci_coeff_cutoff=self.ci_coeff_cutoff, - select_cutoff=self.select_cutoff, - hci_pt=self.hci_pt, scratch_dir=self.scratch_dir, solver_args=self.solver_args, ) @@ -136,11 +126,9 @@ def objfunc(self, xk: list[float]) -> Vector[float64]: nproc=self.nproc, ompnum=self.ompnum, only_chem=self.only_chem, - hci_cutoff=self.hci_cutoff, relax_density=self.relax_density, - ci_coeff_cutoff=self.ci_coeff_cutoff, - select_cutoff=self.select_cutoff, scratch_dir=self.scratch_dir, + solver_args=self.solver_args, ) # Update error and BE energy diff --git a/src/quemb/molbe/solver.py b/src/quemb/molbe/solver.py index e45f39f0..07253371 100644 --- a/src/quemb/molbe/solver.py +++ b/src/quemb/molbe/solver.py @@ -102,6 +102,47 @@ def from_user_input(cls, user_args: DMRG_ArgsUser, mf: RHF): ) +@define +class SHCI_ArgsUser(UserSolverArgs): + hci_pt: bool = False + hci_cutoff: float = 0.001 + ci_coeff_cutoff: float | None = None + select_cutoff: float | None = None + + +@define +class SHCI_Args: + """Properly initialized SCHI arguments + + Some default values of :class:`SHCI_ArgsUser` can only be filled + later in the calculation. + Use :func:`from_user_input` to properly initialize. + """ + + hci_pt: bool + hci_cutoff: float + ci_coeff_cutoff: float + select_cutoff: float + + @classmethod + def from_user_input(cls, args: SHCI_ArgsUser): + if (args.select_cutoff is None) and (args.ci_coeff_cutoff is None): + select_cutoff = args.hci_cutoff + ci_coeff_cutoff = args.hci_cutoff + elif (args.select_cutoff is not None) and (args.ci_coeff_cutoff is not None): + ci_coeff_cutoff = args.ci_coeff_cutoff + select_cutoff = args.select_cutoff + else: + raise ValueError + + return cls( + hci_pt=args.hci_pt, + hci_cutoff=args.hci_cutoff, + ci_coeff_cutoff=ci_coeff_cutoff, + select_cutoff=select_cutoff, + ) + + def be_func( pot: list[float] | None, Fobjs: list[Frags] | list[pFrags], @@ -113,10 +154,6 @@ def be_func( hf_veff: Matrix[float64] | None = None, only_chem: bool = False, nproc: int = 4, - hci_pt: bool = False, - hci_cutoff: float = 0.001, - ci_coeff_cutoff: float | None = None, - select_cutoff: float | None = None, eeval: bool = False, ereturn: bool = False, frag_energy: bool = False, @@ -211,6 +248,9 @@ def be_func( # pylint: disable-next=E0611 from pyscf import hci # noqa: PLC0415 # optional module + assert isinstance(solver_args, SHCI_ArgsUser) + SHCI_args = SHCI_Args.from_user_input(solver_args) + nmo = fobj._mf.mo_coeff.shape[1] eri = ao2mo.kernel( @@ -218,14 +258,9 @@ def be_func( ).reshape(4 * ((nmo),)) ci_ = hci.SCI(fobj._mf.mol) - if select_cutoff is None and ci_coeff_cutoff is None: - select_cutoff = hci_cutoff - ci_coeff_cutoff = hci_cutoff - elif select_cutoff is None or ci_coeff_cutoff is None: - raise ValueError - ci_.select_cutoff = select_cutoff - ci_.ci_coeff_cutoff = ci_coeff_cutoff + ci_.select_cutoff = SHCI_args.select_cutoff + ci_.ci_coeff_cutoff = SHCI_args.ci_coeff_cutoff nelec = (fobj.nsocc, fobj.nsocc) h1_ = fobj.fock + fobj.heff @@ -244,6 +279,9 @@ def be_func( # pylint: disable-next=E0611,E0401 from pyscf.shciscf import shci # noqa: PLC0415 # shci is optional + assert isinstance(solver_args, SHCI_ArgsUser) + SHCI_args = SHCI_Args.from_user_input(solver_args) + frag_scratch = WorkDir(scratch_dir / fobj.dname) nmo = fobj._mf.mo_coeff.shape[1] @@ -251,9 +289,9 @@ def be_func( nelec = (fobj.nsocc, fobj.nsocc) mch = shci.SHCISCF(fobj._mf, nmo, nelec, orbpath=fobj.dname) mch.fcisolver.mpiprefix = "mpirun -np " + str(nproc) - if hci_pt: + if SHCI_args.hci_pt: mch.fcisolver.stochastic = False - mch.fcisolver.epsilon2 = hci_cutoff + mch.fcisolver.epsilon2 = SHCI_args.hci_cutoff else: mch.fcisolver.stochastic = ( True # this is for PT and doesnt add PT to rdm @@ -261,7 +299,7 @@ def be_func( mch.fcisolver.nPTiter = 0 mch.fcisolver.sweep_iter = [0] mch.fcisolver.DoRDM = True - mch.fcisolver.sweep_epsilon = [hci_cutoff] + mch.fcisolver.sweep_epsilon = [SHCI_args.hci_cutoff] mch.fcisolver.scratchDirectory = scratch_dir mch.mc1step() rdm1_tmp, rdm2s = mch.fcisolver.make_rdm12(0, nmo, nelec) @@ -283,7 +321,7 @@ def be_func( ci.runtimedir = fobj.dname ci.restart = True ci.config["var_only"] = True - ci.config["eps_vars"] = [hci_cutoff] + ci.config["eps_vars"] = [SHCI_args.hci_cutoff] ci.config["get_1rdm_csv"] = True ci.config["get_2rdm_csv"] = True ci.kernel(h1, eri, nmo, nelec) From 2e9c60e8e76867aa3bf31fa058625437b870c634 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 23 Dec 2024 15:53:08 -0500 Subject: [PATCH 53/67] fixed typing issues in examples --- example/molbe_dmrg_block2.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/example/molbe_dmrg_block2.py b/example/molbe_dmrg_block2.py index 0eaa5a42..2f06c52e 100644 --- a/example/molbe_dmrg_block2.py +++ b/example/molbe_dmrg_block2.py @@ -7,6 +7,7 @@ from pyscf import cc, fci, gto, scf from quemb.molbe import BE, fragpart +from quemb.molbe.solver import DMRG_ArgsUser # We'll consider the dissociation curve for a 1D chain of 8 H-atoms: num_points = 3 @@ -52,7 +53,7 @@ # Next, run BE-DMRG with default parameters and maxM=100. mybe.oneshot( solver="block2", # or 'DMRG', 'DMRGSCF', 'DMRGCI' - DMRG_solver_kwargs=dict( + solver_args=DMRG_ArgsUser( maxM=100, # Max fragment bond dimension force_cleanup=True, # Remove all fragment DMRG tmpfiles ), @@ -100,7 +101,7 @@ solver="block2", # or 'DMRG', 'DMRGSCF', 'DMRGCI' max_iter=60, # Max number of sweeps only_chem=True, - DMRG_solver_kwargs=dict( + solver_args=DMRG_ArgsUser( startM=20, # Initial fragment bond dimension (1st sweep) maxM=200, # Maximum fragment bond dimension twodot_to_onedot=50, # Sweep num to switch from two- to one-dot algo. @@ -113,7 +114,7 @@ ) # Or, alternatively, we can construct a full schedule by hand: -schedule = { +schedule: dict[str, list[int] | list[float]] = { "scheduleSweeps": [0, 10, 20, 30, 40, 50], # Sweep indices "scheduleMaxMs": [25, 50, 100, 200, 500, 500], # Sweep maxMs "scheduleTols": [1e-5, 1e-5, 1e-6, 1e-6, 1e-8, 1e-8], # Sweep Davidson tolerances @@ -124,7 +125,7 @@ mybe.optimize( solver="block2", only_chem=True, - DMRG_solver_kwargs=dict( + solver_args=DMRG_ArgsUser( schedule_kwargs=schedule, block_extra_keyword=["fiedler"], force_cleanup=True, From bce7ef71e6c3a173fff1946e0314524d3e7695e2 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 23 Dec 2024 16:04:11 -0500 Subject: [PATCH 54/67] made solver_args immutable and improved their docstrings --- src/quemb/molbe/solver.py | 165 +++++++++++++++++++------------------- 1 file changed, 83 insertions(+), 82 deletions(-) diff --git a/src/quemb/molbe/solver.py b/src/quemb/molbe/solver.py index 07253371..09805ca1 100644 --- a/src/quemb/molbe/solver.py +++ b/src/quemb/molbe/solver.py @@ -2,6 +2,7 @@ import os from abc import ABC +from typing import Final import numpy from attrs import define @@ -31,27 +32,62 @@ class UserSolverArgs(ABC): pass -@define +@define(frozen=True) class DMRG_ArgsUser(UserSolverArgs): + """ + + Parameters + ---------- + max_mem: + Maximum memory in GB. + root: + Number of roots to solve for. + startM: + Starting MPS bond dimension - where the sweep schedule begins. + maxM: + Maximum MPS bond dimension - where the sweep schedule terminates. + max_iter: + Maximum number of sweeps. + twodot_to_onedot: + Sweep index at which to transition to one-dot DMRG algorithm. + All sweeps prior to this will use the two-dot algorithm. + block_extra_keyword: + Other keywords to be passed to block2. + See: https://block2.readthedocs.io/en/latest/user/keywords.html + schedule_kwargs: + Dictionary containing DMRG scheduling parameters to be passed to block2. + + e.g. The default schedule used here would be equivalent to the following: + + .. code-block:: python + + schedule_kwargs = { + 'scheduleSweeps': [0, 10, 20, 30, 40, 50], + 'scheduleMaxMs': [25, 50, 100, 200, 500, 500], + 'scheduleTols': [1e-5,1e-5, 1e-6, 1e-6, 1e-8, 1e-8], + 'scheduleNoises': [0.01, 0.01, 0.001, 0.001, 1e-4, 0.0], + } + """ + #: Becomes mf.mo_coeff.shape[1] by default - norb: int | None = None + norb: Final[int | None] = None #: Becomes mf.mo_coeff.shape[1] by default - nelec: int | None = None - - startM: int = 25 - maxM: int = 500 - max_iter: int = 60 - max_mem: int = 100 - max_noise: float = 1e-3 - min_tol: float = 1e-8 - twodot_to_onedot: int = (5 * max_iter) // 6 - root: int = 0 - block_extra_keyword: list[str] = ["fiedler"] - schedule_kwargs: dict[str, list[int] | list[float]] = {} - force_cleanup: bool = False - - -@define + nelec: Final[int | None] = None + + startM: Final[int] = 25 + maxM: Final[int] = 500 + max_iter: Final[int] = 60 + max_mem: Final[int] = 100 + max_noise: Final[float] = 1e-3 + min_tol: Final[float] = 1e-8 + twodot_to_onedot: Final[int] = (5 * max_iter) // 6 + root: Final[int] = 0 + block_extra_keyword: Final[list[str]] = ["fiedler"] + schedule_kwargs: Final[dict[str, list[int] | list[float]]] = {} + force_cleanup: Final[bool] = False + + +@define(frozen=True) class DMRG_Args: """Properly initialized DMRG arguments @@ -60,20 +96,20 @@ class DMRG_Args: Use :func:`from_user_input` to properly initialize. """ - norb: int - nelec: int - - startM: int - maxM: int - max_iter: int - max_mem: int - max_noise: float - min_tol: float - twodot_to_onedot: int - root: int - block_extra_keyword: list[str] - schedule_kwargs: dict[str, list[int] | list[float]] - force_cleanup: bool + norb: Final[int] + nelec: Final[int] + + startM: Final[int] + maxM: Final[int] + max_iter: Final[int] + max_mem: Final[int] + max_noise: Final[float] + min_tol: Final[float] + twodot_to_onedot: Final[int] + root: Final[int] + block_extra_keyword: Final[list[str]] + schedule_kwargs: Final[dict[str, list[int] | list[float]]] + force_cleanup: Final[bool] @classmethod def from_user_input(cls, user_args: DMRG_ArgsUser, mf: RHF): @@ -102,15 +138,15 @@ def from_user_input(cls, user_args: DMRG_ArgsUser, mf: RHF): ) -@define +@define(frozen=True) class SHCI_ArgsUser(UserSolverArgs): - hci_pt: bool = False - hci_cutoff: float = 0.001 - ci_coeff_cutoff: float | None = None - select_cutoff: float | None = None + hci_pt: Final[bool] = False + hci_cutoff: Final[float] = 0.001 + ci_coeff_cutoff: Final[float | None] = None + select_cutoff: Final[float | None] = None -@define +@define(frozen=True) class SHCI_Args: """Properly initialized SCHI arguments @@ -119,10 +155,10 @@ class SHCI_Args: Use :func:`from_user_input` to properly initialize. """ - hci_pt: bool - hci_cutoff: float - ci_coeff_cutoff: float - select_cutoff: float + hci_pt: Final[bool] + hci_cutoff: Final[float] + ci_coeff_cutoff: Final[float] + select_cutoff: Final[float] @classmethod def from_user_input(cls, args: SHCI_ArgsUser): @@ -812,30 +848,16 @@ def solve_block2( Parameters ---------- - mf: pyscf.scf.hf.RHF + mf: Mean field object or similar following the data signature of the pyscf.RHF class. - nocc: int + nocc: Number of occupied MOs in the fragment, used for constructing the fragment 1- and 2-RDMs. - frag_scratch: os.PathLike, optional + frag_scratch: Fragment-level DMRG scratch directory. - max_mem: int, optional - Maximum memory in GB. - root: int, optional - Number of roots to solve for. - startM: int, optional - Starting MPS bond dimension - where the sweep schedule begins. - maxM: int, optional - Maximum MPS bond dimension - where the sweep schedule terminates. - max_iter: int, optional - Maximum number of sweeps. - twodot_to_onedot: int, optional - Sweep index at which to transition to one-dot DMRG algorithm. - All sweeps prior to this will use the two-dot algorithm. - block_extra_keyword: list(str), optional - Other keywords to be passed to block2. - See: https://block2.readthedocs.io/en/latest/user/keywords.html + use_cumulant: + Use the cumulant energy expression. Returns ------- @@ -843,27 +865,6 @@ def solve_block2( 1-Particle reduced density matrix for fragment. rdm2: numpy.ndarray 2-Particle reduced density matrix for fragment. - - Other Parameters - ---------------- - schedule_kwargs: dict, optional - Dictionary containing DMRG scheduling parameters to be passed to block2. - - e.g. The default schedule used here would be equivalent to the following: - - .. code-block:: python - - schedule_kwargs = { - 'scheduleSweeps': [0, 10, 20, 30, 40, 50], - 'scheduleMaxMs': [25, 50, 100, 200, 500, 500], - 'scheduleTols': [1e-5,1e-5, 1e-6, 1e-6, 1e-8, 1e-8], - 'scheduleNoises': [0.01, 0.01, 0.001, 0.001, 1e-4, 0.0], - } - - Raises - ------ - - """ # pylint: disable-next=E0611 from pyscf import dmrgscf # noqa: PLC0415 # optional module From 02e7ac2b50175f6043c981e4aa003e7f21deb2dd Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 23 Dec 2024 16:06:02 -0500 Subject: [PATCH 55/67] incremented the python version for type checking --- .github/workflows/quemb_unittest.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/quemb_unittest.yml b/.github/workflows/quemb_unittest.yml index 733e54d5..0ad7c185 100644 --- a/.github/workflows/quemb_unittest.yml +++ b/.github/workflows/quemb_unittest.yml @@ -17,7 +17,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: ["3.10"] + python-version: ["3.12"] steps: From 128fffb4eb41e09a4eb86a0a97ac050a4825d10c Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 23 Dec 2024 16:30:59 -0500 Subject: [PATCH 56/67] try floating instead of float64 --- src/quemb/molbe/lo.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/quemb/molbe/lo.py b/src/quemb/molbe/lo.py index 8dcfaee9..024d450d 100644 --- a/src/quemb/molbe/lo.py +++ b/src/quemb/molbe/lo.py @@ -2,7 +2,7 @@ # import numpy -from numpy import float64 +from numpy import float64, floating from numpy.linalg import eigh, inv, multi_dot, norm, svd from pyscf.gto import intor_cross from pyscf.gto.mole import Mole @@ -16,8 +16,8 @@ def dot_gen( - A: Matrix[float64], B: Matrix[float64], ovlp: Matrix[float64] | None = None -) -> Matrix[float64]: + A: Matrix[floating], B: Matrix[floating], ovlp: Matrix[floating] | None = None +) -> Matrix[floating]: return A.T @ B if ovlp is None else A.T @ ovlp @ B From cbc1d7d8a48eea512b2ff209edc76af3c93ea6fa Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 23 Dec 2024 16:34:52 -0500 Subject: [PATCH 57/67] don't overspecify the types --- src/quemb/molbe/lo.py | 49 ++++++++++++++++--------------------------- 1 file changed, 18 insertions(+), 31 deletions(-) diff --git a/src/quemb/molbe/lo.py b/src/quemb/molbe/lo.py index 8dcfaee9..755b03c5 100644 --- a/src/quemb/molbe/lo.py +++ b/src/quemb/molbe/lo.py @@ -2,7 +2,6 @@ # import numpy -from numpy import float64 from numpy.linalg import eigh, inv, multi_dot, norm, svd from pyscf.gto import intor_cross from pyscf.gto.mole import Mole @@ -15,15 +14,13 @@ from quemb.shared.typing import Matrix, Tensor3D -def dot_gen( - A: Matrix[float64], B: Matrix[float64], ovlp: Matrix[float64] | None = None -) -> Matrix[float64]: +def dot_gen(A: Matrix, B: Matrix, ovlp: Matrix | None = None) -> Matrix: return A.T @ B if ovlp is None else A.T @ ovlp @ B def get_cano_orth_mat( - A: Matrix[float64], thr: float = 1.0e-6, ovlp: Matrix[float64] | None = None -) -> Matrix[float64]: + A: Matrix, thr: float = 1.0e-6, ovlp: Matrix | None = None +) -> Matrix: S = dot_gen(A, A, ovlp) e, u = eigh(S) if thr > 0: @@ -33,16 +30,14 @@ def get_cano_orth_mat( return u[:, idx_keep] * e[idx_keep] ** -0.5 -def cano_orth( - A: Matrix[float64], thr: float = 1.0e-6, ovlp: Matrix[float64] | None = None -) -> Matrix[float64]: +def cano_orth(A: Matrix, thr: float = 1.0e-6, ovlp: Matrix | None = None) -> Matrix: """Canonically orthogonalize columns of A""" return A @ get_cano_orth_mat(A, thr, ovlp) def get_symm_orth_mat( - A: Matrix[float64], thr: float = 1.0e-6, ovlp: Matrix[float64] | None = None -) -> Matrix[float64]: + A: Matrix, thr: float = 1.0e-6, ovlp: Matrix | None = None +) -> Matrix: S = dot_gen(A, A, ovlp) e, u = eigh(S) if (e < thr).any(): @@ -54,16 +49,12 @@ def get_symm_orth_mat( return u @ numpy.diag(e**-0.5) @ u.T -def symm_orth( - A: Matrix[float64], thr: float = 1.0e-6, ovlp: Matrix[float64] | None = None -) -> Matrix[float64]: +def symm_orth(A: Matrix, thr: float = 1.0e-6, ovlp: Matrix | None = None) -> Matrix: """Symmetrically orthogonalize columns of A""" return A @ get_symm_orth_mat(A, thr, ovlp) -def remove_core_mo( - Clo: Matrix[float64], Ccore: Matrix[float64], S: Matrix[float64], thr: float = 0.5 -) -> Matrix[float64]: +def remove_core_mo(Clo: Matrix, Ccore: Matrix, S: Matrix, thr: float = 0.5) -> Matrix: assert numpy.allclose(Clo.T @ S @ Clo, numpy.eye(Clo.shape[1])) assert numpy.allclose(Ccore.T @ S @ Ccore, numpy.eye(Ccore.shape[1])) @@ -79,7 +70,7 @@ def remove_core_mo( def get_xovlp( mol: Mole, basis: str = "sto-3g" -) -> tuple[Matrix[float64] | Tensor3D[float64], Matrix[float64] | Tensor3D[float64]]: +) -> tuple[Matrix | Tensor3D, Matrix | Tensor3D]: """Gets set of valence orbitals based on smaller (should be minimal) basis Parameters @@ -107,11 +98,11 @@ def get_xovlp( def get_iao( - Co: Matrix[float64], - S12: Matrix[float64], - S1: Matrix[float64], - S2: Matrix[float64] | None = None, -) -> Matrix[float64]: + Co: Matrix, + S12: Matrix, + S1: Matrix, + S2: Matrix | None = None, +) -> Matrix: """ Parameters @@ -157,9 +148,7 @@ def get_iao( return Ciao -def get_pao( - Ciao: Matrix[float64], S: Matrix[float64], S12: Matrix[float64] -) -> Matrix[float64]: +def get_pao(Ciao: Matrix, S: Matrix, S12: Matrix) -> Matrix: """ Parameters ---------- @@ -187,9 +176,7 @@ def get_pao( return cano_orth(Cpao, ovlp=S) -def get_pao_native( - Ciao: Matrix[float64], S: Matrix[float64], mol: Mole, valence_basis: str -) -> Matrix[float64]: +def get_pao_native(Ciao: Matrix, S: Matrix, mol: Mole, valence_basis: str) -> Matrix: """ Parameters @@ -242,10 +229,10 @@ def get_pao_native( def get_loc( mol: Mole, - C: Matrix[float64], + C: Matrix, method: str, pop_method: str | None = None, - init_guess: Matrix[float64] | None = None, + init_guess: Matrix | None = None, ) -> Mole: if method.upper() == "ER": from pyscf.lo import ER as Localizer # noqa: PLC0415 From 6e7dc845a413b0f98ae24905ba4c33e2b759afb6 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 23 Dec 2024 16:41:46 -0500 Subject: [PATCH 58/67] added correct link to doc --- src/quemb/molbe/lo.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/quemb/molbe/lo.py b/src/quemb/molbe/lo.py index 755b03c5..01677099 100644 --- a/src/quemb/molbe/lo.py +++ b/src/quemb/molbe/lo.py @@ -160,7 +160,7 @@ def get_pao(Ciao: Matrix, S: Matrix, S12: Matrix) -> Matrix: valence orbitals projected into ao basis Returns ------- - Cpao: + Cpao: :class:`quemb.shared.typing.Matrix` (orthogonalized) """ n = Ciao.shape[0] @@ -191,7 +191,7 @@ def get_pao_native(Ciao: Matrix, S: Matrix, mol: Mole, valence_basis: str) -> Ma basis used for valence orbitals Returns ------- - Cpao: + Cpao: :class:`quemb.shared.typing.Matrix` (symmetrically orthogonalized) """ From e88d6e4d01ea3a222eaa499aa90b6a88fcef4042 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 23 Dec 2024 17:17:49 -0500 Subject: [PATCH 59/67] simplified DMRG args --- src/quemb/molbe/solver.py | 99 ++++++++++++++++++--------------------- 1 file changed, 46 insertions(+), 53 deletions(-) diff --git a/src/quemb/molbe/solver.py b/src/quemb/molbe/solver.py index 09805ca1..8254282e 100644 --- a/src/quemb/molbe/solver.py +++ b/src/quemb/molbe/solver.py @@ -5,7 +5,7 @@ from typing import Final import numpy -from attrs import define +from attrs import Factory, define, field from numpy import float64 from numpy.linalg import multi_dot from pyscf import ao2mo, cc, fci, mcscf, mp @@ -82,10 +82,40 @@ class DMRG_ArgsUser(UserSolverArgs): min_tol: Final[float] = 1e-8 twodot_to_onedot: Final[int] = (5 * max_iter) // 6 root: Final[int] = 0 - block_extra_keyword: Final[list[str]] = ["fiedler"] - schedule_kwargs: Final[dict[str, list[int] | list[float]]] = {} + block_extra_keyword: Final[list[str]] = Factory(lambda: ["fiedler"]) + schedule_kwargs: dict[str, list[int] | list[float]] = field() force_cleanup: Final[bool] = False + @schedule_kwargs.default + def _get_schedule_kwargs_default(self): + return { + "scheduleSweeps": [(i * self.max_iter) // 6 for i in range(1, 7)], + "scheduleMaxMs": [ + self.startM if (self.startM < self.maxM) else self.maxM, + self.startM * 2 if (self.startM * 2 < self.maxM) else self.maxM, + self.startM * 4 if (self.startM * 4 < self.maxM) else self.maxM, + self.startM * 8 if (self.startM * 8 < self.maxM) else self.maxM, + self.maxM, + self.maxM, + ], + "scheduleTols": [ + self.min_tol * 1e3, + self.min_tol * 1e3, + self.min_tol * 1e2, + self.min_tol * 1e1, + self.min_tol, + self.min_tol, + ], + "scheduleNoises": [ + self.max_noise, + self.max_noise, + self.max_noise / 10, + self.max_noise / 100, + self.max_noise / 100, + 0.0, + ], + } + @define(frozen=True) class DMRG_Args: @@ -869,67 +899,30 @@ def solve_block2( # pylint: disable-next=E0611 from pyscf import dmrgscf # noqa: PLC0415 # optional module - norb = DMRG_args.norb - nelec = DMRG_args.nelec - - startM = DMRG_args.startM - maxM = DMRG_args.maxM - max_iter = DMRG_args.max_iter - max_mem = DMRG_args.max_mem - max_noise = DMRG_args.max_noise - min_tol = DMRG_args.min_tol - twodot_to_onedot = DMRG_args.twodot_to_onedot - root = DMRG_args.root - schedule_kwargs = DMRG_args.schedule_kwargs - block_extra_keyword = DMRG_args.block_extra_keyword - orbs = mf.mo_coeff - mc = mcscf.CASCI(mf, norb, nelec) + mc = mcscf.CASCI(mf, DMRG_args.norb, DMRG_args.nelec) mc.fcisolver = dmrgscf.DMRGCI(mf.mol) # Sweep scheduling - mc.fcisolver.scheduleSweeps = schedule_kwargs.pop( - "scheduleSweeps", - [ - (1 * max_iter) // 6, - (2 * max_iter) // 6, - (3 * max_iter) // 6, - (4 * max_iter) // 6, - (5 * max_iter) // 6, - max_iter, - ], - ) - mc.fcisolver.scheduleMaxMs = schedule_kwargs.pop( - "scheduleMaxMs", - [ - startM if (startM < maxM) else maxM, - startM * 2 if (startM * 2 < maxM) else maxM, - startM * 4 if (startM * 4 < maxM) else maxM, - startM * 8 if (startM * 8 < maxM) else maxM, - maxM, - maxM, - ], - ) - mc.fcisolver.scheduleTols = schedule_kwargs.pop( - "scheduleTols", - [min_tol * 1e3, min_tol * 1e3, min_tol * 1e2, min_tol * 1e1, min_tol, min_tol], - ) - mc.fcisolver.scheduleNoises = schedule_kwargs.pop( - "scheduleNoises", - [max_noise, max_noise, max_noise / 10, max_noise / 100, max_noise / 100, 0.0], - ) + mc.fcisolver.scheduleSweeps = DMRG_args.schedule_kwargs["scheduleSweeps"] + mc.fcisolver.scheduleMaxMs = DMRG_args.schedule_kwargs["scheduleMaxMs"] + mc.fcisolver.scheduleTols = DMRG_args.schedule_kwargs["scheduleTols"] + mc.fcisolver.scheduleNoises = DMRG_args.schedule_kwargs["scheduleNoises"] + # Other DMRG parameters mc.fcisolver.threads = int(os.environ.get("OMP_NUM_THREADS", "8")) - mc.fcisolver.twodot_to_onedot = int(twodot_to_onedot) - mc.fcisolver.maxIter = int(max_iter) - mc.fcisolver.block_extra_keyword = list(block_extra_keyword) + mc.fcisolver.twodot_to_onedot = DMRG_args.twodot_to_onedot + mc.fcisolver.maxIter = DMRG_args.max_iter + mc.fcisolver.block_extra_keyword = DMRG_args.block_extra_keyword mc.fcisolver.scratchDirectory = str(frag_scratch) mc.fcisolver.runtimeDir = str(frag_scratch) - mc.fcisolver.memory = int(max_mem) + mc.fcisolver.memory = DMRG_args.max_mem os.chdir(frag_scratch) mc.kernel(orbs) - rdm1, rdm2 = dmrgscf.DMRGCI.make_rdm12(mc.fcisolver, root, norb, nelec) + rdm1, rdm2 = dmrgscf.DMRGCI.make_rdm12( + mc.fcisolver, DMRG_args.root, DMRG_args.norb, DMRG_args.nelec + ) # Subtract off non-cumulant contribution to correlated 2RDM. if use_cumulant: From 271b1b0d43ef2287a320ad000b40d37566452b4a Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 23 Dec 2024 17:19:02 -0500 Subject: [PATCH 60/67] use a factory for mutable default attributes --- src/quemb/molbe/opt.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/quemb/molbe/opt.py b/src/quemb/molbe/opt.py index db0bbe75..efcc243d 100644 --- a/src/quemb/molbe/opt.py +++ b/src/quemb/molbe/opt.py @@ -2,7 +2,7 @@ import numpy -from attrs import define +from attrs import Factory, define from numpy import array, float64 from quemb.molbe.be_parallel import be_func_parallel @@ -74,7 +74,7 @@ class BEOPT: iter: int = 0 err: float = 0.0 - Ebe: Matrix[float64] = array([[0.0]]) + Ebe: Matrix[float64] = Factory(lambda: array([[0.0]])) solver_args: UserSolverArgs | None = None From 4198d9b59c1fd26df52cf29bef299b06695b1b70 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 23 Dec 2024 17:20:09 -0500 Subject: [PATCH 61/67] removed unused setting and changed default for SCRATCH to /tmp --- src/quemb/shared/config.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/quemb/shared/config.py b/src/quemb/shared/config.py index 6667e085..a2944177 100644 --- a/src/quemb/shared/config.py +++ b/src/quemb/shared/config.py @@ -32,8 +32,7 @@ @define class Settings: PRINT_LEVEL: int = 5 - SCRATCH: str = "" - CREATE_SCRATCH_DIR: bool = False + SCRATCH: Path = Path("/tmp") INTEGRAL_TRANSFORM_MAX_MEMORY: float = 50 # in GB From 214b09186f17071beeaa3a18f197359fbe60752e Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Mon, 23 Dec 2024 17:55:59 -0500 Subject: [PATCH 62/67] added types to kbe.BE.optimize and fixed corresponding tests --- src/quemb/kbe/pbe.py | 29 +++++++++++++---------------- src/quemb/molbe/opt.py | 3 ++- 2 files changed, 15 insertions(+), 17 deletions(-) diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py index 4f8aa0ff..a39324cd 100644 --- a/src/quemb/kbe/pbe.py +++ b/src/quemb/kbe/pbe.py @@ -315,16 +315,16 @@ def __init__( def optimize( self, - solver="MP2", - method="QN", - only_chem=False, - conv_tol=1.0e-6, - relax_density=False, - J0=None, - nproc=1, - ompnum=4, - max_iter=500, - ): + solver: str = "MP2", + method: str = "QN", + only_chem: bool = False, + conv_tol: float = 1.0e-6, + relax_density: bool = False, + J0: list[list[float]] | None = None, + nproc: int = 1, + ompnum: int = 4, + max_iter: int = 500, + ) -> None: """BE optimization function Interfaces BEOPT to perform bootstrap embedding optimization. @@ -380,10 +380,7 @@ def optimize( max_space=max_iter, conv_tol=conv_tol, only_chem=only_chem, - hci_cutoff=self.hci_cutoff, - ci_coeff_cutoff=self.ci_coeff_cutoff, relax_density=relax_density, - select_cutoff=self.select_cutoff, solver=solver, ebe_hf=self.ebe_hf, ) @@ -393,7 +390,7 @@ def optimize( if only_chem: J0 = [[0.0]] J0 = self.get_be_error_jacobian(jac_solver="HF") - J0 = [[J0[-1, -1]]] + J0 = [[J0[-1][-1]]] else: J0 = self.get_be_error_jacobian(jac_solver="HF") @@ -412,10 +409,10 @@ def optimize( raise ValueError("This optimization method for BE is not supported") @copy_docstring(_ext_get_be_error_jacobian) - def get_be_error_jacobian(self, jac_solver="HF"): + def get_be_error_jacobian(self, jac_solver: str = "HF") -> list[list[float]]: return _ext_get_be_error_jacobian(self.Nfrag, self.Fobjs, jac_solver) - def print_ini(self): + def print_ini(self) -> None: """ Print initialization banner for the kBE calculation. """ diff --git a/src/quemb/molbe/opt.py b/src/quemb/molbe/opt.py index efcc243d..364ecf57 100644 --- a/src/quemb/molbe/opt.py +++ b/src/quemb/molbe/opt.py @@ -5,6 +5,7 @@ from attrs import Factory, define from numpy import array, float64 +from quemb.kbe.pfrag import Frags as pFrags from quemb.molbe.be_parallel import be_func_parallel from quemb.molbe.pfrag import Frags from quemb.molbe.solver import UserSolverArgs, be_func @@ -57,7 +58,7 @@ class BEOPT: """ pot: list[float] - Fobjs: list[Frags] + Fobjs: list[Frags] | list[pFrags] Nocc: int enuc: float scratch_dir: WorkDir From a29ec7ab09dbb1e55aae8b2e77b0b28b30564407 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Thu, 9 Jan 2025 17:00:34 -0500 Subject: [PATCH 63/67] fixed small error --- src/quemb/kbe/pbe.py | 1 - src/quemb/molbe/mbe.py | 1 - 2 files changed, 2 deletions(-) diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py index 5503a0a5..c7bde6ce 100644 --- a/src/quemb/kbe/pbe.py +++ b/src/quemb/kbe/pbe.py @@ -698,7 +698,6 @@ def oneshot( solver, self.enuc, nproc=ompnum, - ereturn=True, eeval=True, scratch_dir=self.scratch_dir, solver_args=solver_args, diff --git a/src/quemb/molbe/mbe.py b/src/quemb/molbe/mbe.py index 4af8bafa..129f2b2a 100644 --- a/src/quemb/molbe/mbe.py +++ b/src/quemb/molbe/mbe.py @@ -930,7 +930,6 @@ def oneshot( solver, self.enuc, nproc=ompnum, - ereturn=True, eeval=True, scratch_dir=self.scratch_dir, solver_args=solver_args, From 6d9bdb736a491a7ea7290c30bfbc0a5227c7fe12 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Thu, 9 Jan 2025 17:08:39 -0500 Subject: [PATCH 64/67] use tempfile.gettempdir() --- src/quemb/shared/config.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/quemb/shared/config.py b/src/quemb/shared/config.py index a2944177..de16ec6e 100644 --- a/src/quemb/shared/config.py +++ b/src/quemb/shared/config.py @@ -18,6 +18,7 @@ """ from pathlib import Path +from tempfile import gettempdir from typing import Final import yaml @@ -32,7 +33,7 @@ @define class Settings: PRINT_LEVEL: int = 5 - SCRATCH: Path = Path("/tmp") + SCRATCH: Path = Path(gettempdir()) INTEGRAL_TRANSFORM_MAX_MEMORY: float = 50 # in GB From 5601fa4a21a2bef29e15fc86469e802d8286c886 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Thu, 9 Jan 2025 17:08:47 -0500 Subject: [PATCH 65/67] fixed J0[-1, -1] also for kbe --- src/quemb/kbe/pbe.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py index c7bde6ce..c20c8185 100644 --- a/src/quemb/kbe/pbe.py +++ b/src/quemb/kbe/pbe.py @@ -7,6 +7,7 @@ import h5py import numpy from libdmet.basis_transform.eri_transform import get_emb_eri_fast_gdf +from numpy import array, floating from pyscf import ao2mo, pbc from pyscf.pbc import df, gto from pyscf.pbc.df.df_jk import _ewald_exxdiv_for_G0 @@ -24,7 +25,7 @@ ) from quemb.shared.helper import copy_docstring from quemb.shared.manage_scratch import WorkDir -from quemb.shared.typing import PathLike +from quemb.shared.typing import Matrix, PathLike class BE(Mixin_k_Localize): @@ -321,7 +322,7 @@ def optimize( use_cumulant: bool = True, conv_tol: float = 1.0e-6, relax_density: bool = False, - J0: list[list[float]] | None = None, + J0: Matrix[floating] | None = None, nproc: int = 1, ompnum: int = 4, max_iter: int = 500, @@ -391,9 +392,9 @@ def optimize( if method == "QN": # Prepare the initial Jacobian matrix if only_chem: - J0 = [[0.0]] + J0 = array([[0.0]]) J0 = self.get_be_error_jacobian(jac_solver="HF") - J0 = [[J0[-1][-1]]] + J0 = J0[-1:, -1:] else: J0 = self.get_be_error_jacobian(jac_solver="HF") @@ -416,7 +417,7 @@ def optimize( raise ValueError("This optimization method for BE is not supported") @copy_docstring(_ext_get_be_error_jacobian) - def get_be_error_jacobian(self, jac_solver: str = "HF") -> list[list[float]]: + def get_be_error_jacobian(self, jac_solver: str = "HF") -> Matrix[floating]: return _ext_get_be_error_jacobian(self.Nfrag, self.Fobjs, jac_solver) def print_ini(self) -> None: From 8f080e75e7268a344e754d9b0d4fb3670f8ea4b5 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 10 Jan 2025 16:06:54 -0500 Subject: [PATCH 66/67] adressed Shaun's comment --- src/quemb/molbe/solver.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/quemb/molbe/solver.py b/src/quemb/molbe/solver.py index aed9e21e..ee14db73 100644 --- a/src/quemb/molbe/solver.py +++ b/src/quemb/molbe/solver.py @@ -85,7 +85,7 @@ class DMRG_ArgsUser(UserSolverArgs): force_cleanup: Final[bool] = False @schedule_kwargs.default - def _get_schedule_kwargs_default(self): + def _get_schedule_kwargs_default(self) -> dict[str, list[int] | list[float]]: return { "scheduleSweeps": [(i * self.max_iter) // 6 for i in range(1, 7)], "scheduleMaxMs": [ @@ -197,7 +197,10 @@ def from_user_input(cls, args: SHCI_ArgsUser): ci_coeff_cutoff = args.ci_coeff_cutoff select_cutoff = args.select_cutoff else: - raise ValueError + raise ValueError( + "Solver args `ci_coeff_cutoff` and `select_cutoff` must both " + "be specified or both be `None`!" + ) return cls( hci_pt=args.hci_pt, From 65ec4b45d2baf9ab6271483bbad7ba24708d91bb Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 10 Jan 2025 16:08:49 -0500 Subject: [PATCH 67/67] properly hide _DMRG_Args and _SHCI_Args --- src/quemb/molbe/be_parallel.py | 8 ++++---- src/quemb/molbe/solver.py | 10 +++++----- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/quemb/molbe/be_parallel.py b/src/quemb/molbe/be_parallel.py index 65b4183f..e221823a 100644 --- a/src/quemb/molbe/be_parallel.py +++ b/src/quemb/molbe/be_parallel.py @@ -17,9 +17,9 @@ ) from quemb.molbe.pfrag import Frags from quemb.molbe.solver import ( - SHCI_Args, SHCI_ArgsUser, UserSolverArgs, + _SHCI_Args, make_rdm1_ccsd_t1, make_rdm2_urlx, solve_ccsd, @@ -149,7 +149,7 @@ def run_solver( from pyscf import hci # noqa: PLC0415 # hci is an optional module assert isinstance(solver_args, SHCI_ArgsUser) - SHCI_args = SHCI_Args.from_user_input(solver_args) + SHCI_args = _SHCI_Args.from_user_input(solver_args) nao, nmo = mf_.mo_coeff.shape eri = ao2mo.kernel(mf_._eri, mf_.mo_coeff, aosym="s4", compact=False).reshape( @@ -177,7 +177,7 @@ def run_solver( frag_scratch = WorkDir(scratch_dir / dname) assert isinstance(solver_args, SHCI_ArgsUser) - SHCI_args = SHCI_Args.from_user_input(solver_args) + SHCI_args = _SHCI_Args.from_user_input(solver_args) nao, nmo = mf_.mo_coeff.shape nelec = (nocc, nocc) @@ -199,7 +199,7 @@ def run_solver( from pyscf import cornell_shci # noqa: PLC0415 # optional module assert isinstance(solver_args, SHCI_ArgsUser) - SHCI_args = SHCI_Args.from_user_input(solver_args) + SHCI_args = _SHCI_Args.from_user_input(solver_args) frag_scratch = WorkDir(scratch_dir / dname) diff --git a/src/quemb/molbe/solver.py b/src/quemb/molbe/solver.py index ee14db73..77be5855 100644 --- a/src/quemb/molbe/solver.py +++ b/src/quemb/molbe/solver.py @@ -116,7 +116,7 @@ def _get_schedule_kwargs_default(self) -> dict[str, list[int] | list[float]]: @define(frozen=True) -class DMRG_Args: +class _DMRG_Args: """Properly initialized DMRG arguments Some default values of :class:`DMRG_ArgsUser` can only be filled @@ -175,7 +175,7 @@ class SHCI_ArgsUser(UserSolverArgs): @define(frozen=True) -class SHCI_Args: +class _SHCI_Args: """Properly initialized SCHI arguments Some default values of :class:`SHCI_ArgsUser` can only be filled @@ -312,7 +312,7 @@ def be_func( from pyscf import hci # noqa: PLC0415 # optional module assert isinstance(solver_args, SHCI_ArgsUser) - SHCI_args = SHCI_Args.from_user_input(solver_args) + SHCI_args = _SHCI_Args.from_user_input(solver_args) nmo = fobj._mf.mo_coeff.shape[1] @@ -343,7 +343,7 @@ def be_func( from pyscf.shciscf import shci # noqa: PLC0415 # shci is optional assert isinstance(solver_args, SHCI_ArgsUser) - SHCI_args = SHCI_Args.from_user_input(solver_args) + SHCI_args = _SHCI_Args.from_user_input(solver_args) frag_scratch = WorkDir(scratch_dir / fobj.dname) @@ -394,7 +394,7 @@ def be_func( frag_scratch = WorkDir(scratch_dir / fobj.dname) assert isinstance(solver_args, DMRG_ArgsUser) - DMRG_args = DMRG_Args.from_user_input(solver_args, fobj._mf) + DMRG_args = _DMRG_Args.from_user_input(solver_args, fobj._mf) try: rdm1_tmp, rdm2s = solve_block2(