From a6cad3107179064fd1ea453a4048b250b821fc57 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Fri, 10 Jan 2025 16:38:25 -0500 Subject: [PATCH] Dataclass for solver kwargs (#77) - Introduced a dataclass for solver kwargs. This is makes it much more explicit what arguments are supported by DMRG and SHCI - Could simplify the interface of BE by removing a couple of arguments --- example/molbe_dmrg_block2.py | 9 +- src/quemb/kbe/pbe.py | 69 +++---- src/quemb/molbe/be_parallel.py | 66 +++--- src/quemb/molbe/mbe.py | 41 +--- src/quemb/molbe/opt.py | 41 ++-- src/quemb/molbe/solver.py | 362 +++++++++++++++++++++------------ src/quemb/shared/config.py | 4 +- 7 files changed, 326 insertions(+), 266 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, diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py index 63ecf6bb..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 @@ -18,13 +19,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 Matrix, PathLike class BE(Mixin_k_Localize): @@ -58,12 +59,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 +156,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 @@ -325,17 +316,17 @@ def __init__( def optimize( self, - solver="MP2", - method="QN", - only_chem=False, - use_cumulant=True, - 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, + use_cumulant: bool = True, + conv_tol: float = 1.0e-6, + relax_density: bool = False, + J0: Matrix[floating] | None = None, + nproc: int = 1, + ompnum: int = 4, + max_iter: int = 500, + ) -> None: """BE optimization function Interfaces BEOPT to perform bootstrap embedding optimization. @@ -393,10 +384,7 @@ def optimize( conv_tol=conv_tol, only_chem=only_chem, use_cumulant=use_cumulant, - 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, ) @@ -404,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") @@ -429,10 +417,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") -> Matrix[floating]: 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. """ @@ -683,7 +671,7 @@ def oneshot( use_cumulant: bool = True, nproc: int = 1, ompnum: int = 4, - DMRG_solver_kwargs: KwargDict | None = None, + solver_args: UserSolverArgs | None = None, ) -> None: """ Perform a one-shot bootstrap embedding calculation. @@ -711,14 +699,11 @@ def oneshot( solver, self.enuc, nproc=ompnum, - use_cumulant=use_cumulant, eeval=True, - return_vec=False, - hci_cutoff=self.hci_cutoff, - ci_coeff_cutoff=self.ci_coeff_cutoff, - select_cutoff=self.select_cutoff, scratch_dir=self.scratch_dir, - DMRG_solver_kwargs=DMRG_solver_kwargs, + solver_args=solver_args, + use_cumulant=use_cumulant, + return_vec=False, ) else: rets = be_func_parallel( @@ -727,15 +712,13 @@ def oneshot( self.Nocc, solver, self.enuc, + eeval=True, nproc=nproc, ompnum=ompnum, + scratch_dir=self.scratch_dir, + solver_args=solver_args, use_cumulant=use_cumulant, - eeval=True, return_vec=False, - hci_cutoff=self.hci_cutoff, - ci_coeff_cutoff=self.ci_coeff_cutoff, - select_cutoff=self.select_cutoff, - scratch_dir=self.scratch_dir, ) print("-----------------------------------------------------", flush=True) @@ -759,8 +742,6 @@ def oneshot( flush=True, ) - self.ebe_tot = rets[0] - def update_fock(self, heff=None): """ Update the Fock matrix for each fragment with the effective Hamiltonian. diff --git a/src/quemb/molbe/be_parallel.py b/src/quemb/molbe/be_parallel.py index b56bf21a..e221823a 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_ArgsUser, + UserSolverArgs, + _SHCI_Args, make_rdm1_ccsd_t1, make_rdm2_urlx, solve_ccsd, @@ -46,15 +49,13 @@ def run_solver( eri_file: str = "eri_file.h5", veff: Matrix[float64] | None = None, 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, ret_vec: bool = False, use_cumulant: bool = True, relax_density: bool = False, + solver_args: UserSolverArgs | None = None, ): """ Run a quantum chemistry solver to compute the reduced density matrices. @@ -67,9 +68,12 @@ def run_solver( Initial guess for the density matrix. scratch_dir : The scratch dir root. - Fragment files will be stored in :code:`scratch_dir / dname`. dname : Directory name for storing intermediate files. + Fragment files will be stored in :code:`scratch_dir / dname`. + scratch_dir : + The scratch directory. + Fragment files will be stored in :code:`scratch_dir / dname`. nao : Number of atomic orbitals. nocc : @@ -95,12 +99,12 @@ def run_solver( Number of OpenMP threads. Default is 4. writeh1 : If True, write the one-electron integrals to a file. Default is False. + use_cumulant : + If True, use the cumulant approximation for RDM2. Default is True. eeval : If True, evaluate the electronic energy. Default is True. ret_vec : If True, return vector with error and rdms. Default is True. - use_cumulant : - If True, use the cumulant approximation for RDM2. Default is True. relax_density : If True, use CCSD relaxed density. Default is False @@ -144,19 +148,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)) @@ -174,6 +176,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) @@ -182,7 +187,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 @@ -193,6 +198,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 @@ -208,7 +216,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) @@ -382,16 +390,14 @@ def be_func_parallel( solver: str, enuc: float, # noqa: ARG001 scratch_dir: WorkDir, - only_chem: bool = False, + solver_args: UserSolverArgs | None, nproc: int = 1, ompnum: int = 4, + only_chem: bool = False, relax_density: bool = False, use_cumulant: bool = True, - eeval: bool = True, - return_vec: bool = True, - hci_cutoff: float = 0.001, - ci_coeff_cutoff: float | None = None, - select_cutoff: float | None = None, + eeval: bool = False, + return_vec: bool = False, writeh1: bool = False, ): """ @@ -416,21 +422,21 @@ def be_func_parallel( 'FCI', 'HCI', 'SHCI', and 'SCI'. enuc : Nuclear component of the energy. - scratch_dir : - Scratch directory root - only_chem : - Whether to perform chemical potential optimization only. - Refer to bootstrap embedding literature. Defaults to False. nproc : Total number of processors assigned for the optimization. Defaults to 1. When nproc > 1, Python multithreading is invoked. ompnum : If nproc > 1, sets the number of cores for OpenMP parallelization. Defaults to 4. - use_cumulant : - Use cumulant energy expression. Defaults to True + only_chem : + Whether to perform chemical potential optimization only. + Refer to bootstrap embedding literature. Defaults to False. eeval : Whether to evaluate energies. Defaults to False. + scratch_dir : + Scratch directory root + use_cumulant : + Use cumulant energy expression. Defaults to True return_vec : Whether to return the error vector. Defaults to False. writeh1 : @@ -472,15 +478,13 @@ def be_func_parallel( fobj.eri_file, fobj.veff if not use_cumulant else None, fobj.veff0, - hci_cutoff, - ci_coeff_cutoff, - select_cutoff, ompnum, writeh1, eeval, return_vec, use_cumulant, relax_density, + solver_args, ], ) diff --git a/src/quemb/molbe/mbe.py b/src/quemb/molbe/mbe.py index c1b28e4e..129f2b2a 100644 --- a/src/quemb/molbe/mbe.py +++ b/src/quemb/molbe/mbe.py @@ -15,13 +15,13 @@ from quemb.molbe.misc import print_energy_cumulant, print_energy_noncumulant 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 @@ -79,10 +79,6 @@ def __init__( nproc: int = 1, ompnum: int = 4, scratch_dir: WorkDir | None = None, - hci_pt: bool = False, - 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: @@ -169,12 +165,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 @@ -647,7 +637,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 @@ -709,14 +699,10 @@ def optimize( conv_tol=conv_tol, only_chem=only_chem, use_cumulant=use_cumulant, - 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, - DMRG_solver_kwargs=DMRG_solver_kwargs, + solver_args=solver_args, ) if method == "QN": @@ -918,7 +904,7 @@ def oneshot( use_cumulant: bool = True, nproc: int = 1, ompnum: int = 4, - DMRG_solver_kwargs: KwargDict | None = None, + solver_args: UserSolverArgs | None = None, ) -> None: """ Perform a one-shot bootstrap embedding calculation. @@ -944,14 +930,11 @@ def oneshot( solver, self.enuc, nproc=ompnum, - use_cumulant=use_cumulant, eeval=True, - return_vec=False, - hci_cutoff=self.hci_cutoff, - ci_coeff_cutoff=self.ci_coeff_cutoff, - select_cutoff=self.select_cutoff, scratch_dir=self.scratch_dir, - DMRG_solver_kwargs=DMRG_solver_kwargs, + solver_args=solver_args, + use_cumulant=use_cumulant, + return_vec=False, ) else: rets = be_func_parallel( @@ -960,15 +943,13 @@ def oneshot( self.Nocc, solver, self.enuc, + eeval=True, nproc=nproc, ompnum=ompnum, + scratch_dir=self.scratch_dir, + solver_args=solver_args, use_cumulant=use_cumulant, - eeval=True, return_vec=False, - hci_cutoff=self.hci_cutoff, - ci_coeff_cutoff=self.ci_coeff_cutoff, - select_cutoff=self.select_cutoff, - scratch_dir=self.scratch_dir, ) print("-----------------------------------------------------", flush=True) diff --git a/src/quemb/molbe/opt.py b/src/quemb/molbe/opt.py index 5db491c9..dea701f0 100644 --- a/src/quemb/molbe/opt.py +++ b/src/quemb/molbe/opt.py @@ -2,15 +2,16 @@ import numpy -from attrs import define +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 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 @@ -39,6 +40,9 @@ class BEOPT: 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 : + Whether to perform chemical potential optimization only. + Refer to bootstrap embedding literatures. nproc : Total number of processors assigned for the optimization. Defaults to 1. When nproc > 1, Python multithreading @@ -46,9 +50,6 @@ class BEOPT: ompnum : If nproc > 1, ompnum sets the number of cores for OpenMP parallelization. Defaults to 4 - only_chem : - Whether to perform chemical potential optimization only. - Refer to bootstrap embedding literatures. max_space : Maximum number of bootstrap optimizaiton steps, after which the optimization is called converged. @@ -59,14 +60,14 @@ class BEOPT: """ pot: list[float] - Fobjs: list[Frags] + Fobjs: list[Frags] | list[pFrags] Nocc: int enuc: float scratch_dir: WorkDir + solver: str = "MP2" nproc: int = 1 ompnum: int = 4 only_chem: bool = False - solver: str = "MP2" use_cumulant: bool = True max_space: int = 500 @@ -76,15 +77,9 @@ class BEOPT: iter: int = 0 err: float = 0.0 - Ebe: Matrix[float64] = array([[0.0]]) - - DMRG_solver_kwargs: KwargDict | None = None + Ebe: Matrix[float64] = Factory(lambda: array([[0.0]])) - # HCI parameters - hci_cutoff: float = 0.001 - ci_coeff_cutoff: float | None = None - select_cutoff: float | None = None - hci_pt: bool = False + solver_args: UserSolverArgs | None = None def objfunc(self, xk: list[float]) -> Vector[float64]: """ @@ -115,15 +110,11 @@ def objfunc(self, xk: list[float]) -> Vector[float64]: only_chem=self.only_chem, nproc=self.ompnum, relax_density=self.relax_density, + scratch_dir=self.scratch_dir, + solver_args=self.solver_args, use_cumulant=self.use_cumulant, eeval=True, return_vec=True, - hci_cutoff=self.hci_cutoff, - ci_coeff_cutoff=self.ci_coeff_cutoff, - select_cutoff=self.select_cutoff, - hci_pt=self.hci_pt, - scratch_dir=self.scratch_dir, - DMRG_solver_kwargs=self.DMRG_solver_kwargs, ) else: err_, errvec_, ebe_ = be_func_parallel( @@ -136,13 +127,11 @@ def objfunc(self, xk: list[float]) -> Vector[float64]: nproc=self.nproc, ompnum=self.ompnum, relax_density=self.relax_density, + scratch_dir=self.scratch_dir, + solver_args=self.solver_args, use_cumulant=self.use_cumulant, eeval=True, return_vec=True, - hci_cutoff=self.hci_cutoff, - ci_coeff_cutoff=self.ci_coeff_cutoff, - select_cutoff=self.select_cutoff, - scratch_dir=self.scratch_dir, ) # Update error and BE energy diff --git a/src/quemb/molbe/solver.py b/src/quemb/molbe/solver.py index 667273f6..77be5855 100644 --- a/src/quemb/molbe/solver.py +++ b/src/quemb/molbe/solver.py @@ -1,8 +1,11 @@ # Author(s): Oinam Romesh Meitei, Leah Weisburn, Shaun Weatherly import os +from abc import ABC +from typing import Final import numpy +from attrs import Factory, define, field from numpy.linalg import multi_dot from pyscf import ao2mo, cc, fci, mcscf, mp from pyscf.cc.ccsd_rdm import make_rdm2 @@ -21,7 +24,190 @@ 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 + + +class UserSolverArgs(ABC): + pass + + +@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: Final[int | None] = None + #: Becomes mf.mo_coeff.shape[1] by default + 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]] = 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) -> dict[str, list[int] | list[float]]: + 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: + """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: 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): + 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, + ) + + +@define(frozen=True) +class SHCI_ArgsUser(UserSolverArgs): + 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(frozen=True) +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: Final[bool] + hci_cutoff: Final[float] + ci_coeff_cutoff: Final[float] + select_cutoff: Final[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( + "Solver args `ci_coeff_cutoff` and `select_cutoff` must both " + "be specified or both be `None`!" + ) + + 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( @@ -30,18 +216,14 @@ def be_func( Nocc: int, solver: str, enuc: float, # noqa: ARG001 + solver_args: UserSolverArgs | None, scratch_dir: WorkDir, - DMRG_solver_kwargs: KwargDict | None, only_chem: bool = False, nproc: int = 4, - relax_density: bool = False, - use_cumulant: bool = True, eeval: bool = False, + relax_density: bool = False, return_vec: bool = False, - hci_pt: bool = False, - hci_cutoff: float = 0.001, - ci_coeff_cutoff: float | None = None, - select_cutoff: float | None = None, + use_cumulant: bool = True, ): """ Perform bootstrap embedding calculations for each fragment. @@ -65,8 +247,14 @@ def be_func( Whether to only optimize the chemical potential. Defaults to False. nproc : Number of processors. Defaults to 4. This is only neccessary for 'SHCI' solver + eeval : + Whether to evaluate the energy. Defaults to False. + ereturn : + Whether to return the energy. Defaults to False. relax_density : Whether to relax the density. Defaults to False. + return_vec : + Whether to return the error vector. Defaults to False. use_cumulant : Whether to use the cumulant-based energy expression. Defaults to True. @@ -123,6 +311,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( @@ -130,14 +321,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 @@ -156,6 +342,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] @@ -163,9 +352,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 @@ -173,7 +362,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) @@ -195,7 +384,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) @@ -204,21 +393,21 @@ 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() - ) + assert isinstance(solver_args, DMRG_ArgsUser) + DMRG_args = _DMRG_Args.from_user_input(solver_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*"), @@ -666,37 +855,23 @@ 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. 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 ------- @@ -704,105 +879,34 @@ 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 - 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." - ) + 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: diff --git a/src/quemb/shared/config.py b/src/quemb/shared/config.py index 6667e085..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,8 +33,7 @@ @define class Settings: PRINT_LEVEL: int = 5 - SCRATCH: str = "" - CREATE_SCRATCH_DIR: bool = False + SCRATCH: Path = Path(gettempdir()) INTEGRAL_TRANSFORM_MAX_MEMORY: float = 50 # in GB