From 0f3c6c2d7eea70d4a6fb2bdb4c292b8d7bfe8068 Mon Sep 17 00:00:00 2001 From: Minsik Date: Wed, 25 Sep 2024 09:04:07 -0400 Subject: [PATCH] Feat: Add on-the-fly density-fitted fragment ERI transform (#38) * Feat: Add on-the-fly density-fitted fragment ERI transform Co-authored-by: hongzhouye * Add test for on-the-fly DF ERI transformation * Add note about HF-in-HF error for int direct DF --------- Co-authored-by: hongzhouye --- molbe/be_var.py | 1 + molbe/eri_onthefly.py | 92 ++++++++++++++++++++++++++++++++++++++ molbe/mbe.py | 21 ++++++++- tests/eri_onthefly_test.py | 28 ++++++++++++ 4 files changed, 140 insertions(+), 2 deletions(-) create mode 100644 molbe/eri_onthefly.py create mode 100644 tests/eri_onthefly_test.py diff --git a/molbe/be_var.py b/molbe/be_var.py index 9fdb8fa..fc97dc1 100644 --- a/molbe/be_var.py +++ b/molbe/be_var.py @@ -3,3 +3,4 @@ SOLVER_CALL = 0 SCRATCH='' CREATE_SCRATCH_DIR=False +INTEGRAL_TRANSFORM_MAX_MEMORY=50 # in GB diff --git a/molbe/eri_onthefly.py b/molbe/eri_onthefly.py new file mode 100644 index 0000000..687f76c --- /dev/null +++ b/molbe/eri_onthefly.py @@ -0,0 +1,92 @@ +# Author(s): Minsik Cho, Hong-Zhou Ye + +import numpy, os + +from pyscf import lib +from pyscf.df.addons import make_auxmol +from pyscf.gto.moleintor import getints3c, make_loc, make_cintopt +from pyscf.gto import mole +from pyscf.ao2mo.addons import restore +from scipy.linalg import cholesky, solve_triangular +from . import be_var + +def integral_direct_DF(mf, Fobjs, file_eri, auxbasis=None): + """ Calculate AO density-fitted 3-center integrals on-the-fly and transform to Schmidt space for given fragment objects + + Parameters + ---------- + mf : pyscf.scf.RHF + Mean-field object for the chemical system (typically BE.mf) + Fobjs : list of BE.Frags + List containing fragment objects (typically BE.Fobjs) + The MO coefficients are taken from Frags.TA and the transformed ERIs are stored in Frags.dname as h5py datasets. + file_eri : h5py.File + HDF5 file object to store the transformed fragment ERIs + auxbasis : string, optional + Auxiliary basis used for density fitting. If not provided, use pyscf's default choice for the basis set used to construct mf object; by default None + """ + + def calculate_pqL(aux_range): + """ Internal function to calculate the 3-center integrals for a given range of auxiliary indices + + Parameters + ---------- + aux_range : tuple of int + (start index, end index) of the auxiliary basis functions to calculate the 3-center integrals, i.e. (pq|L) with L ∈ [start, end) is returned + """ + if be_var.PRINT_LEVEL > 10: print('Start calculating (μν|P) for range', aux_range, flush=True) + p0, p1 = aux_range + shls_slice = (0, mf.mol.nbas, 0, mf.mol.nbas, mf.mol.nbas + p0, mf.mol.nbas + p1) + ints = getints3c(mf.mol._add_suffix('int3c2e'), atm, bas, env, shls_slice, 1, 's1', ao_loc, cintopt, out=None) + if be_var.PRINT_LEVEL > 10: print('Finish calculating (μν|P) for range', aux_range, flush=True) + return ints + + def block_step_size(nfrag, naux, nao): + """ Internal function to calculate the block step size for the 3-center integrals calculation + + Parameters + ---------- + nfrag : int + Number of fragments + naux : int + Number of auxiliary basis functions + nao : int + Number of atomic orbitals + """ + return max(1, int( + be_var.INTEGRAL_TRANSFORM_MAX_MEMORY * 1e9 / 8 / nao / nao / naux / nfrag + )) #max(int(500*.24e6/8/nao),1) + + if be_var.PRINT_LEVEL > 2: + print('Evaluating fragment ERIs on-the-fly using density fitting...', flush=True) + print('In this case, note that HF-in-HF error includes DF error on top of numerical error from embedding.', flush=True) + auxmol = make_auxmol(mf.mol, auxbasis = auxbasis) + j2c = auxmol.intor(mf.mol._add_suffix('int2c2e'), hermi=1) # (L|M) + low = cholesky(j2c, lower=True) + pqL_frag = [numpy.zeros((auxmol.nao, fragobj.nao, fragobj.nao)) for fragobj in Fobjs] # place to store fragment (pq|L) + end = 0 + atm, bas, env = mole.conc_env(mf.mol._atm, mf.mol._bas, mf.mol._env, auxmol._atm, auxmol._bas, auxmol._env) + ao_loc = make_loc(bas, mf.mol._add_suffix('int3c2e')); cintopt = make_cintopt(atm, bas, env, mf.mol._add_suffix('int3c2e')) + blockranges = [(x, y) for x, y in lib.prange(0, auxmol.nbas, block_step_size(len(Fobjs), auxmol.nbas, mf.mol.nao))] + if be_var.PRINT_LEVEL > 4: + print('Aux Basis Block Info: ', blockranges, flush=True) + + for idx, ints in enumerate(lib.map_with_prefetch(calculate_pqL, blockranges)): + if be_var.PRINT_LEVEL > 4: print('Calculating pq|L block #', idx, blockranges[idx], flush=True) + # Transform pq (AO) to fragment space (ij) + start = end; end += ints.shape[2] + for fragidx in range(len(Fobjs)): + if be_var.PRINT_LEVEL > 10: print('(μν|P) -> (ij|P) for frag #', fragidx, flush=True) + Lqp = numpy.transpose(ints, axes=(2,1,0)) + Lqi = Lqp @ Fobjs[fragidx].TA + Liq = numpy.moveaxis(Lqi, 2, 1) + pqL_frag[fragidx][start:end, :, :] = Liq @ Fobjs[fragidx].TA + # Fit to get B_{ij}^{L} + for fragidx in range(len(Fobjs)): + if be_var.PRINT_LEVEL > 10: print('Fitting B_{ij}^{L} for frag #', fragidx, flush=True) + b = pqL_frag[fragidx].reshape(auxmol.nao, -1) + bb = solve_triangular(low, b, lower=True, overwrite_b=True, check_finite=False) + if be_var.PRINT_LEVEL > 10: print('Finished obtaining B_{ij}^{L} for frag #', fragidx, flush=True) + eri_nosym = bb.T @ bb + eri = restore('4', eri_nosym, Fobjs[fragidx].nao) + file_eri.create_dataset(Fobjs[fragidx].dname, data=eri) diff --git a/molbe/mbe.py b/molbe/mbe.py index 2696d84..4d0d449 100644 --- a/molbe/mbe.py +++ b/molbe/mbe.py @@ -55,7 +55,8 @@ def __init__(self, mf, fobj, eri_file='eri_file.h5', mo_energy = None, save_file='storebe.pk',hci_pt=False, nproc=1, ompnum=4, - hci_cutoff=0.001, ci_coeff_cutoff = None, select_cutoff=None): + hci_cutoff=0.001, ci_coeff_cutoff = None, select_cutoff=None, + integral_direct_DF=False, auxbasis = None): """ Constructor for BE object. @@ -85,6 +86,10 @@ def __init__(self, mf, fobj, eri_file='eri_file.h5', Number of processors for parallel calculations, by default 1. If set to >1, threaded parallel computation is invoked. ompnum : int, optional Number of OpenMP threads, by default 4. + integral_direct_DF: bool, optional + 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 + Auxiliary basis for density fitting, by default None (uses default auxiliary basis defined in PySCF). """ if restart: @@ -111,6 +116,8 @@ def __init__(self, mf, fobj, eri_file='eri_file.h5', self.unrestricted = False self.nproc = nproc self.ompnum = ompnum + self.integral_direct_DF = integral_direct_DF + self.auxbasis = auxbasis # Fragment information from fobj self.frag_type=fobj.frag_type @@ -303,7 +310,8 @@ def initialize(self, eri_,compute_hf, restart=False): # Yes -- ao2mo, incore version # No -- Do we have (ij|P) from density fitting? # Yes -- ao2mo, outcore version, using saved (ij|P) - assert (not eri_ is None) or (hasattr(self.mf, 'with_df')), "Input mean-field object is missing ERI (mf._eri) or DF (mf.with_df) object. You may want to ensure that incore_anyway was set for non-DF calculations." + # No -- if integral_direct_DF is requested, invoke on-the-fly routine + assert (not eri_ is None) or (hasattr(self.mf, 'with_df')) or (self.integral_direct_DF), "Input mean-field object is missing ERI (mf._eri) or DF (mf.with_df) object AND integral direct DF routine was not requested. Please check your inputs." if not eri_ is None: # incore ao2mo using saved eri from mean-field calculation for I in range(self.Nfrag): eri = ao2mo.incore.full(eri_, self.Fobjs[I].TA, compact=True) @@ -313,6 +321,15 @@ def initialize(self, eri_,compute_hf, restart=False): for I in range(self.Nfrag): eri = self.mf.with_df.ao2mo(self.Fobjs[I].TA, compact=True) file_eri.create_dataset(self.Fobjs[I].dname, data=eri) + else: + # If ERIs are not saved on memory, compute fragment ERIs integral-direct + if self.integral_direct_DF: # Use density fitting to generate fragment ERIs on-the-fly + from .eri_onthefly import integral_direct_DF + integral_direct_DF(self.mf, self.Fobjs, file_eri, auxbasis=self.auxbasis) + 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 else: eri=None diff --git a/tests/eri_onthefly_test.py b/tests/eri_onthefly_test.py new file mode 100644 index 0000000..bd442f6 --- /dev/null +++ b/tests/eri_onthefly_test.py @@ -0,0 +1,28 @@ +""" +This script tests the on-the-fly DF-based ERI transformation routine for molecular systems. + +Author(s): Minsik Cho +""" + +import os, unittest +from pyscf import gto, scf +from molbe import fragpart, BE + +class TestDF_ontheflyERI(unittest.TestCase): + @unittest.skipIf(os.getenv("GITHUB_ACTIONS") == "true", "Skip expensive tests on Github Actions") + def test_octane_BE2(self): + # Octane, cc-pvtz + mol = gto.M() + mol.atom = os.path.join(os.path.dirname(__file__), 'xyz/octane.xyz') + mol.basis = 'cc-pvtz' + mol.charge = 0.; mol.spin = 0. + mol.build() + mf = scf.RHF(mol); mf.direct_scf = True; mf.kernel() + fobj = fragpart(frag_type='autogen', be_type='be2', mol = mol) + mybe = BE(mf, fobj, integral_direct_DF=True) + self.assertAlmostEqual(mybe.ebe_hf, mf.e_tot, + msg = "HF-in-HF energy for Octane (BE2) does not match the HF energy!", + delta = 1e-6) + +if __name__ == '__main__': + unittest.main()