1#!/usr/bin/env python 2# Copyright 2014-2021 The PySCF Developers. All Rights Reserved. 3# 4# Licensed under the Apache License, Version 2.0 (the "License"); 5# you may not use this file except in compliance with the License. 6# You may obtain a copy of the License at 7# 8# http://www.apache.org/licenses/LICENSE-2.0 9# 10# Unless required by applicable law or agreed to in writing, software 11# distributed under the License is distributed on an "AS IS" BASIS, 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13# See the License for the specific language governing permissions and 14# limitations under the License. 15 16 17import numpy 18import h5py 19from pyscf import gto 20from pyscf import lib 21from pyscf.lib import logger 22from pyscf.ao2mo import _ao2mo 23from pyscf.ao2mo import incore 24from pyscf import __config__ 25 26IOBLK_SIZE = getattr(__config__, 'ao2mo_outcore_ioblk_size', 256) # 256 MB 27IOBUF_WORDS = getattr(__config__, 'ao2mo_outcore_iobuf_words', 1e8) # 800 MB 28IOBUF_ROW_MIN = getattr(__config__, 'ao2mo_outcore_row_min', 160) 29MAX_MEMORY = getattr(__config__, 'ao2mo_outcore_max_memory', 2000) # 2GB 30 31 32def full(mol, mo_coeff, erifile, dataname='eri_mo', 33 intor='int2e', aosym='s4', comp=None, 34 max_memory=MAX_MEMORY, ioblk_size=IOBLK_SIZE, verbose=logger.WARN, 35 compact=True): 36 r'''Transfer arbitrary spherical AO integrals to MO integrals for given orbitals 37 38 Args: 39 mol : :class:`Mole` object 40 AO integrals will be generated in terms of mol._atm, mol._bas, mol._env 41 mo_coeff : ndarray 42 Transform (ij|kl) with the same set of orbitals. 43 erifile : str or h5py File or h5py Group object 44 To store the transformed integrals, in HDF5 format. 45 46 Kwargs: 47 dataname : str 48 The dataset name in the erifile (ref the hierarchy of HDF5 format 49 http://www.hdfgroup.org/HDF5/doc1.6/UG/09_Groups.html). By assigning 50 different dataname, the existed integral file can be reused. If 51 the erifile contains the dataname, the new integrals data will 52 overwrite the old one. 53 intor : str 54 Name of the 2-electron integral. Ref to :func:`getints_by_shell` 55 for the complete list of available 2-electron integral names 56 aosym : int or str 57 Permutation symmetry for the AO integrals 58 59 | 4 or '4' or 's4': 4-fold symmetry (default) 60 | '2ij' or 's2ij' : symmetry between i, j in (ij|kl) 61 | '2kl' or 's2kl' : symmetry between k, l in (ij|kl) 62 | 1 or '1' or 's1': no symmetry 63 | 'a4ij' : 4-fold symmetry with anti-symmetry between i, j in (ij|kl) (TODO) 64 | 'a4kl' : 4-fold symmetry with anti-symmetry between k, l in (ij|kl) (TODO) 65 | 'a2ij' : anti-symmetry between i, j in (ij|kl) (TODO) 66 | 'a2kl' : anti-symmetry between k, l in (ij|kl) (TODO) 67 68 comp : int 69 Components of the integrals, e.g. int2e_ip_sph has 3 components. 70 max_memory : float or int 71 The maximum size of cache to use (in MB), large cache may **not** 72 improve performance. 73 ioblk_size : float or int 74 The block size for IO, large block size may **not** improve performance 75 verbose : int 76 Print level 77 compact : bool 78 When compact is True, depending on the four oribital sets, the 79 returned MO integrals has (up to 4-fold) permutation symmetry. 80 If it's False, the function will abandon any permutation symmetry, 81 and return the "plain" MO integrals 82 83 Returns: 84 None 85 86 Examples: 87 88 >>> from pyscf import gto 89 >>> from pyscf import ao2mo 90 >>> import h5py 91 >>> def view(h5file, dataname='eri_mo'): 92 ... f5 = h5py.File(h5file, 'r') 93 ... print('dataset %s, shape %s' % (str(f5.keys()), str(f5[dataname].shape))) 94 ... f5.close() 95 >>> mol = gto.M(atom='O 0 0 0; H 0 1 0; H 0 0 1', basis='sto3g') 96 >>> mo1 = numpy.random.random((mol.nao_nr(), 10)) 97 >>> ao2mo.outcore.full(mol, mo1, 'full.h5') 98 >>> view('full.h5') 99 dataset ['eri_mo'], shape (55, 55) 100 >>> ao2mo.outcore.full(mol, mo1, 'full.h5', dataname='new', compact=False) 101 >>> view('full.h5', 'new') 102 dataset ['eri_mo', 'new'], shape (100, 100) 103 >>> ao2mo.outcore.full(mol, mo1, 'full.h5', intor='int2e_ip1_sph', aosym='s1', comp=3) 104 >>> view('full.h5') 105 dataset ['eri_mo', 'new'], shape (3, 100, 100) 106 >>> ao2mo.outcore.full(mol, mo1, 'full.h5', intor='int2e_ip1_sph', aosym='s2kl', comp=3) 107 >>> view('full.h5') 108 dataset ['eri_mo', 'new'], shape (3, 100, 55) 109 ''' 110 general(mol, (mo_coeff,)*4, erifile, dataname, 111 intor, aosym, comp, max_memory, ioblk_size, verbose, compact) 112 return erifile 113 114def general(mol, mo_coeffs, erifile, dataname='eri_mo', 115 intor='int2e', aosym='s4', comp=None, 116 max_memory=MAX_MEMORY, ioblk_size=IOBLK_SIZE, verbose=logger.WARN, 117 compact=True): 118 r'''For the given four sets of orbitals, transfer arbitrary spherical AO 119 integrals to MO integrals on the fly. 120 121 Args: 122 mol : :class:`Mole` object 123 AO integrals will be generated in terms of mol._atm, mol._bas, mol._env 124 mo_coeffs : 4-item list of ndarray 125 Four sets of orbital coefficients, corresponding to the four 126 indices of (ij|kl) 127 erifile : str or h5py File or h5py Group object 128 To store the transformed integrals, in HDF5 format. 129 130 Kwargs 131 dataname : str 132 The dataset name in the erifile (ref the hierarchy of HDF5 format 133 http://www.hdfgroup.org/HDF5/doc1.6/UG/09_Groups.html). By assigning 134 different dataname, the existed integral file can be reused. If 135 the erifile contains the dataname, the new integrals data will 136 overwrite the old one. 137 intor : str 138 Name of the 2-electron integral. Ref to :func:`getints_by_shell` 139 for the complete list of available 2-electron integral names 140 aosym : int or str 141 Permutation symmetry for the AO integrals 142 143 | 4 or '4' or 's4': 4-fold symmetry (default) 144 | '2ij' or 's2ij' : symmetry between i, j in (ij|kl) 145 | '2kl' or 's2kl' : symmetry between k, l in (ij|kl) 146 | 1 or '1' or 's1': no symmetry 147 | 'a4ij' : 4-fold symmetry with anti-symmetry between i, j in (ij|kl) (TODO) 148 | 'a4kl' : 4-fold symmetry with anti-symmetry between k, l in (ij|kl) (TODO) 149 | 'a2ij' : anti-symmetry between i, j in (ij|kl) (TODO) 150 | 'a2kl' : anti-symmetry between k, l in (ij|kl) (TODO) 151 152 comp : int 153 Components of the integrals, e.g. int2e_ip_sph has 3 components. 154 max_memory : float or int 155 The maximum size of cache to use (in MB), large cache may **not** 156 improve performance. 157 ioblk_size : float or int 158 The block size for IO, large block size may **not** improve performance 159 verbose : int 160 Print level 161 compact : bool 162 When compact is True, depending on the four oribital sets, the 163 returned MO integrals has (up to 4-fold) permutation symmetry. 164 If it's False, the function will abandon any permutation symmetry, 165 and return the "plain" MO integrals 166 167 Returns: 168 None 169 170 Examples: 171 172 >>> from pyscf import gto 173 >>> from pyscf import ao2mo 174 >>> import h5py 175 >>> def view(h5file, dataname='eri_mo'): 176 ... f5 = h5py.File(h5file, 'r') 177 ... print('dataset %s, shape %s' % (str(f5.keys()), str(f5[dataname].shape))) 178 ... f5.close() 179 >>> mol = gto.M(atom='O 0 0 0; H 0 1 0; H 0 0 1', basis='sto3g') 180 >>> mo1 = numpy.random.random((mol.nao_nr(), 10)) 181 >>> mo2 = numpy.random.random((mol.nao_nr(), 8)) 182 >>> mo3 = numpy.random.random((mol.nao_nr(), 6)) 183 >>> mo4 = numpy.random.random((mol.nao_nr(), 4)) 184 >>> ao2mo.outcore.general(mol, (mo1,mo2,mo3,mo4), 'oh2.h5') 185 >>> view('oh2.h5') 186 dataset ['eri_mo'], shape (80, 24) 187 >>> ao2mo.outcore.general(mol, (mo1,mo2,mo3,mo3), 'oh2.h5') 188 >>> view('oh2.h5') 189 dataset ['eri_mo'], shape (80, 21) 190 >>> ao2mo.outcore.general(mol, (mo1,mo2,mo3,mo3), 'oh2.h5', compact=False) 191 >>> view('oh2.h5') 192 dataset ['eri_mo'], shape (80, 36) 193 >>> ao2mo.outcore.general(mol, (mo1,mo1,mo2,mo2), 'oh2.h5') 194 >>> view('oh2.h5') 195 dataset ['eri_mo'], shape (55, 36) 196 >>> ao2mo.outcore.general(mol, (mo1,mo1,mo1,mo1), 'oh2.h5', dataname='new') 197 >>> view('oh2.h5', 'new') 198 dataset ['eri_mo', 'new'], shape (55, 55) 199 >>> ao2mo.outcore.general(mol, (mo1,mo1,mo1,mo1), 'oh2.h5', intor='int2e_ip1_sph', aosym='s1', comp=3) 200 >>> view('oh2.h5') 201 dataset ['eri_mo', 'new'], shape (3, 100, 100) 202 >>> ao2mo.outcore.general(mol, (mo1,mo1,mo1,mo1), 'oh2.h5', intor='int2e_ip1_sph', aosym='s2kl', comp=3) 203 >>> view('oh2.h5') 204 dataset ['eri_mo', 'new'], shape (3, 100, 55) 205 ''' 206 if any(c.dtype == numpy.complex128 for c in mo_coeffs): 207 raise NotImplementedError('Integral transformation for complex orbitals') 208 209 time_0pass = (logger.process_clock(), logger.perf_counter()) 210 log = logger.new_logger(mol, verbose) 211 212 nmoi = mo_coeffs[0].shape[1] 213 nmoj = mo_coeffs[1].shape[1] 214 nmol = mo_coeffs[3].shape[1] 215 nao = mo_coeffs[0].shape[0] 216 217 intor, comp = gto.moleintor._get_intor_and_comp(mol._add_suffix(intor), comp) 218 assert(nao == mol.nao_nr('_cart' in intor)) 219 220 aosym = _stand_sym_code(aosym) 221 if aosym in ('s4', 's2kl'): 222 nao_pair = nao * (nao+1) // 2 223 else: 224 nao_pair = nao * nao 225 226 if (compact and iden_coeffs(mo_coeffs[0], mo_coeffs[1]) and 227 aosym in ('s4', 's2ij')): 228 nij_pair = nmoi*(nmoi+1) // 2 229 else: 230 nij_pair = nmoi*nmoj 231 232 klmosym, nkl_pair, mokl, klshape = \ 233 incore._conc_mos(mo_coeffs[2], mo_coeffs[3], 234 compact and aosym in ('s4', 's2kl')) 235 236# if nij_pair > nkl_pair: 237# log.warn('low efficiency for AO to MO trans!') 238 239 if isinstance(erifile, str): 240 if h5py.is_hdf5(erifile): 241 feri = h5py.File(erifile, 'a') 242 if dataname in feri: 243 del(feri[dataname]) 244 else: 245 feri = h5py.File(erifile, 'w') 246 else: 247 assert(isinstance(erifile, h5py.Group)) 248 feri = erifile 249 250 if comp == 1: 251 chunks = (nmoj, nmol) 252 shape = (nij_pair, nkl_pair) 253 else: 254 chunks = (1, nmoj, nmol) 255 shape = (comp, nij_pair, nkl_pair) 256 257 if nij_pair == 0 or nkl_pair == 0: 258 feri.create_dataset(dataname, shape, 'f8') 259 if isinstance(erifile, str): 260 feri.close() 261 return erifile 262 else: 263 h5d_eri = feri.create_dataset(dataname, shape, 'f8', chunks=chunks) 264 265 log.debug('MO integrals %s are saved in %s/%s', intor, erifile, dataname) 266 log.debug('num. MO ints = %.8g, required disk %.8g MB', 267 float(nij_pair)*nkl_pair*comp, nij_pair*nkl_pair*comp*8/1e6) 268 269# transform e1 270 fswap = lib.H5TmpFile() 271 half_e1(mol, mo_coeffs, fswap, intor, aosym, comp, max_memory, ioblk_size, 272 log, compact) 273 274 time_1pass = log.timer('AO->MO transformation for %s 1 pass'%intor, 275 *time_0pass) 276 277 def load(icomp, row0, row1, buf): 278 if icomp+1 < comp: 279 icomp += 1 280 else: # move to next row-block 281 row0, row1 = row1, min(nij_pair, row1+iobuflen) 282 icomp = 0 283 if row0 < row1: 284 _load_from_h5g(fswap['%d'%icomp], row0, row1, buf) 285 286 def save(icomp, row0, row1, buf): 287 if comp == 1: 288 h5d_eri[row0:row1] = buf[:row1-row0] 289 else: 290 h5d_eri[icomp,row0:row1] = buf[:row1-row0] 291 292 ioblk_size = max(max_memory*.1, ioblk_size) 293 iobuflen = guess_e2bufsize(ioblk_size, nij_pair, max(nao_pair,nkl_pair))[0] 294 buf = numpy.empty((iobuflen,nao_pair)) 295 buf_prefetch = numpy.empty_like(buf) 296 outbuf = numpy.empty((iobuflen,nkl_pair)) 297 buf_write = numpy.empty_like(outbuf) 298 299 log.debug('step2: kl-pair (ao %d, mo %d), mem %.8g MB, ioblock %.8g MB', 300 nao_pair, nkl_pair, iobuflen*nao_pair*8/1e6, 301 iobuflen*nkl_pair*8/1e6) 302 303 #klaoblks = len(fswap['0']) 304 ijmoblks = int(numpy.ceil(float(nij_pair)/iobuflen)) * comp 305 ao_loc = mol.ao_loc_nr('_cart' in intor) 306 ti0 = time_1pass 307 istep = 0 308 with lib.call_in_background(load) as prefetch: 309 with lib.call_in_background(save) as async_write: 310 _load_from_h5g(fswap['0'], 0, min(nij_pair, iobuflen), buf_prefetch) 311 312 for row0, row1 in prange(0, nij_pair, iobuflen): 313 nrow = row1 - row0 314 315 for icomp in range(comp): 316 istep += 1 317 log.debug1('step 2 [%d/%d], [%d,%d:%d], row = %d', 318 istep, ijmoblks, icomp, row0, row1, nrow) 319 320 buf, buf_prefetch = buf_prefetch, buf 321 prefetch(icomp, row0, row1, buf_prefetch) 322 _ao2mo.nr_e2(buf[:nrow], mokl, klshape, aosym, klmosym, 323 ao_loc=ao_loc, out=outbuf) 324 async_write(icomp, row0, row1, outbuf) 325 outbuf, buf_write = buf_write, outbuf # avoid flushing writing buffer 326 327 ti1 = (logger.process_clock(), logger.perf_counter()) 328 log.debug1('step 2 [%d/%d] CPU time: %9.2f, Wall time: %9.2f', 329 istep, ijmoblks, ti1[0]-ti0[0], ti1[1]-ti0[1]) 330 ti0 = ti1 331 332 fswap = None 333 if isinstance(erifile, str): 334 feri.close() 335 336 log.timer('AO->MO transformation for %s 2 pass'%intor, *time_1pass) 337 log.timer('AO->MO transformation for %s '%intor, *time_0pass) 338 return erifile 339 340 341# swapfile will be overwritten if exists. 342def half_e1(mol, mo_coeffs, swapfile, 343 intor='int2e', aosym='s4', comp=1, 344 max_memory=MAX_MEMORY, ioblk_size=IOBLK_SIZE, verbose=logger.WARN, 345 compact=True, ao2mopt=None): 346 r'''Half transform arbitrary spherical AO integrals to MO integrals 347 for the given two sets of orbitals 348 349 Args: 350 mol : :class:`Mole` object 351 AO integrals will be generated in terms of mol._atm, mol._bas, mol._env 352 mo_coeff : ndarray 353 Transform (ij|kl) with the same set of orbitals. 354 swapfile : str or h5py File or h5py Group object 355 To store the transformed integrals, in HDF5 format. The transformed 356 integrals are saved in blocks. 357 358 Kwargs 359 intor : str 360 Name of the 2-electron integral. Ref to :func:`getints_by_shell` 361 for the complete list of available 2-electron integral names 362 aosym : int or str 363 Permutation symmetry for the AO integrals 364 365 | 4 or '4' or 's4': 4-fold symmetry (default) 366 | '2ij' or 's2ij' : symmetry between i, j in (ij|kl) 367 | '2kl' or 's2kl' : symmetry between k, l in (ij|kl) 368 | 1 or '1' or 's1': no symmetry 369 | 'a4ij' : 4-fold symmetry with anti-symmetry between i, j in (ij|kl) (TODO) 370 | 'a4kl' : 4-fold symmetry with anti-symmetry between k, l in (ij|kl) (TODO) 371 | 'a2ij' : anti-symmetry between i, j in (ij|kl) (TODO) 372 | 'a2kl' : anti-symmetry between k, l in (ij|kl) (TODO) 373 374 comp : int 375 Components of the integrals, e.g. int2e_ip_sph has 3 components. 376 verbose : int 377 Print level 378 max_memory : float or int 379 The maximum size of cache to use (in MB), large cache may **not** 380 improve performance. 381 ioblk_size : float or int 382 The block size for IO, large block size may **not** improve performance 383 verbose : int 384 Print level 385 compact : bool 386 When compact is True, depending on the four oribital sets, the 387 returned MO integrals has (up to 4-fold) permutation symmetry. 388 If it's False, the function will abandon any permutation symmetry, 389 and return the "plain" MO integrals 390 ao2mopt : :class:`AO2MOpt` object 391 Precomputed data to improve perfomance 392 393 Returns: 394 None 395 396 ''' 397 if any(c.dtype == numpy.complex128 for c in mo_coeffs): 398 raise NotImplementedError('Integral transformation for complex orbitals') 399 400 intor = mol._add_suffix(intor) 401 time0 = (logger.process_clock(), logger.perf_counter()) 402 log = logger.new_logger(mol, verbose) 403 404 nao = mo_coeffs[0].shape[0] 405 aosym = _stand_sym_code(aosym) 406 if aosym in ('s4', 's2ij'): 407 nao_pair = nao * (nao+1) // 2 408 else: 409 nao_pair = nao * nao 410 411 ijmosym, nij_pair, moij, ijshape = \ 412 incore._conc_mos(mo_coeffs[0], mo_coeffs[1], 413 compact and aosym in ('s4', 's2ij')) 414 415 e1buflen, mem_words, iobuf_words, ioblk_words = \ 416 guess_e1bufsize(max_memory, ioblk_size, nij_pair, nao_pair, comp) 417 ioblk_size = ioblk_words * 8/1e6 418# The buffer to hold AO integrals in C code, see line (@) 419 aobuflen = max(int((mem_words - 2*comp*e1buflen*nij_pair) // (nao_pair*comp)), 420 IOBUF_ROW_MIN) 421 ao_loc = mol.ao_loc_nr('_cart' in intor) 422 shranges = guess_shell_ranges(mol, (aosym in ('s4', 's2kl')), e1buflen, 423 aobuflen, ao_loc) 424 if ao2mopt is None: 425 if intor == 'int2e_cart' or intor == 'int2e_sph': 426 ao2mopt = _ao2mo.AO2MOpt(mol, intor, 'CVHFnr_schwarz_cond', 427 'CVHFsetnr_direct_scf') 428 else: 429 ao2mopt = _ao2mo.AO2MOpt(mol, intor) 430 431 if isinstance(swapfile, h5py.Group): 432 fswap = swapfile 433 else: 434 fswap = lib.H5TmpFile(swapfile) 435 for icomp in range(comp): 436 fswap.create_group(str(icomp)) # for h5py old version 437 438 log.debug('step1: tmpfile %s %.8g MB', fswap.filename, nij_pair*nao_pair*8/1e6) 439 log.debug('step1: (ij,kl) = (%d,%d), mem cache %.8g MB, iobuf %.8g MB', 440 nij_pair, nao_pair, mem_words*8/1e6, iobuf_words*8/1e6) 441 nstep = len(shranges) 442 e1buflen = max([x[2] for x in shranges]) 443 444 e2buflen, chunks = guess_e2bufsize(ioblk_size, nij_pair, e1buflen) 445 def save(istep, iobuf): 446 for icomp in range(comp): 447 _transpose_to_h5g(fswap, '%d/%d'%(icomp,istep), iobuf[icomp], 448 e2buflen, None) 449 450 # transform e1 451 ti0 = log.timer('Initializing ao2mo.outcore.half_e1', *time0) 452 with lib.call_in_background(save) as async_write: 453 buf1 = numpy.empty((comp*e1buflen,nao_pair)) 454 buf2 = numpy.empty((comp*e1buflen,nij_pair)) 455 buf_write = numpy.empty_like(buf2) 456 fill = _ao2mo.nr_e1fill 457 f_e1 = _ao2mo.nr_e1 458 for istep,sh_range in enumerate(shranges): 459 log.debug1('step 1 [%d/%d], AO [%d:%d], len(buf) = %d', 460 istep+1, nstep, *(sh_range[:3])) 461 buflen = sh_range[2] 462 iobuf = numpy.ndarray((comp,buflen,nij_pair), buffer=buf2) 463 nmic = len(sh_range[3]) 464 p1 = 0 465 for imic, aoshs in enumerate(sh_range[3]): 466 log.debug2(' fill iobuf micro [%d/%d], AO [%d:%d], len(aobuf) = %d', 467 imic+1, nmic, *aoshs) 468 buf = fill(intor, aoshs, mol._atm, mol._bas, mol._env, 469 aosym, comp, ao2mopt, out=buf1).reshape(-1,nao_pair) 470 buf = f_e1(buf, moij, ijshape, aosym, ijmosym) 471 p0, p1 = p1, p1 + aoshs[2] 472 iobuf[:,p0:p1] = buf.reshape(comp,aoshs[2],nij_pair) 473 ti0 = log.timer_debug1('gen AO/transform MO [%d/%d]'%(istep+1,nstep), *ti0) 474 475 async_write(istep, iobuf) 476 buf2, buf_write = buf_write, buf2 477 478 fswap = None 479 return swapfile 480 481def _load_from_h5g(h5group, row0, row1, out=None): 482 nkeys = len(h5group) 483 dat = h5group['0'] 484 ncol = sum(h5group[str(key)].shape[-1] for key in range(nkeys)) 485 if dat.ndim == 2: 486 out = numpy.ndarray((row1-row0, ncol), dat.dtype, buffer=out) 487 col1 = 0 488 for key in range(nkeys): 489 dat = h5group[str(key)][row0:row1] 490 col0, col1 = col1, col1 + dat.shape[1] 491 out[:,col0:col1] = dat 492 else: # multiple components 493 out = numpy.ndarray((dat.shape[0], row1-row0, ncol), dat.dtype, buffer=out) 494 col1 = 0 495 for key in range(nkeys): 496 dat = h5group[str(key)][:,row0:row1] 497 col0, col1 = col1, col1 + dat.shape[2] 498 out[:,:,col0:col1] = dat 499 return out 500 501def _transpose_to_h5g(h5group, key, dat, blksize, chunks=None): 502 nrow, ncol = dat.shape 503 dset = h5group.create_dataset(key, (ncol,nrow), 'f8', chunks=chunks) 504 for col0, col1 in prange(0, ncol, blksize): 505 dset[col0:col1] = lib.transpose(dat[:,col0:col1]) 506 507def full_iofree(mol, mo_coeff, intor='int2e', aosym='s4', comp=None, 508 max_memory=MAX_MEMORY, ioblk_size=IOBLK_SIZE, 509 verbose=logger.WARN, compact=True): 510 r'''Transfer arbitrary spherical AO integrals to MO integrals for given orbitals 511 This function is a wrap for :func:`ao2mo.outcore.general`. It's not really 512 IO free. The returned MO integrals are held in memory. For backward compatibility, 513 it is used to replace the non-existed function direct.full_iofree. 514 515 Args: 516 mol : :class:`Mole` object 517 AO integrals will be generated in terms of mol._atm, mol._bas, mol._env 518 mo_coeff : ndarray 519 Transform (ij|kl) with the same set of orbitals. 520 erifile : str 521 To store the transformed integrals, in HDF5 format. 522 523 Kwargs 524 dataname : str 525 The dataset name in the erifile (ref the hierarchy of HDF5 format 526 http://www.hdfgroup.org/HDF5/doc1.6/UG/09_Groups.html). By assigning 527 different dataname, the existed integral file can be reused. If 528 the erifile contains the dataname, the new integrals data will 529 overwrite the old one. 530 intor : str 531 Name of the 2-electron integral. Ref to :func:`getints_by_shell` 532 for the complete list of available 2-electron integral names 533 aosym : int or str 534 Permutation symmetry for the AO integrals 535 536 | 4 or '4' or 's4': 4-fold symmetry (default) 537 | '2ij' or 's2ij' : symmetry between i, j in (ij|kl) 538 | '2kl' or 's2kl' : symmetry between k, l in (ij|kl) 539 | 1 or '1' or 's1': no symmetry 540 | 'a4ij' : 4-fold symmetry with anti-symmetry between i, j in (ij|kl) (TODO) 541 | 'a4kl' : 4-fold symmetry with anti-symmetry between k, l in (ij|kl) (TODO) 542 | 'a2ij' : anti-symmetry between i, j in (ij|kl) (TODO) 543 | 'a2kl' : anti-symmetry between k, l in (ij|kl) (TODO) 544 545 comp : int 546 Components of the integrals, e.g. int2e_ip_sph has 3 components. 547 verbose : int 548 Print level 549 max_memory : float or int 550 The maximum size of cache to use (in MB), large cache may **not** 551 improve performance. 552 ioblk_size : float or int 553 The block size for IO, large block size may **not** improve performance 554 verbose : int 555 Print level 556 compact : bool 557 When compact is True, depending on the four oribital sets, the 558 returned MO integrals has (up to 4-fold) permutation symmetry. 559 If it's False, the function will abandon any permutation symmetry, 560 and return the "plain" MO integrals 561 562 Returns: 563 2D/3D MO-integral array. They may or may not have the permutation 564 symmetry, depending on the given orbitals, and the kwargs compact. If 565 the four sets of orbitals are identical, the MO integrals will at most 566 have 4-fold symmetry. 567 568 Examples: 569 570 >>> from pyscf import gto 571 >>> from pyscf import ao2mo 572 >>> mol = gto.M(atom='O 0 0 0; H 0 1 0; H 0 0 1', basis='sto3g') 573 >>> mo1 = numpy.random.random((mol.nao_nr(), 10)) 574 >>> eri1 = ao2mo.outcore.full_iofree(mol, mo1) 575 >>> print(eri1.shape) 576 (55, 55) 577 >>> eri1 = ao2mo.outcore.full_iofree(mol, mo1, compact=False) 578 >>> print(eri1.shape) 579 (100, 100) 580 >>> eri1 = ao2mo.outcore.full_iofree(mol, mo1, intor='int2e_ip1_sph', aosym='s1', comp=3) 581 >>> print(eri1.shape) 582 (3, 100, 100) 583 >>> eri1 = ao2mo.outcore.full_iofree(mol, mo1, intor='int2e_ip1_sph', aosym='s2kl', comp=3) 584 >>> print(eri1.shape) 585 (3, 100, 55) 586 ''' 587 with lib.H5TmpFile() as feri: 588 general(mol, (mo_coeff,)*4, feri, dataname='eri_mo', 589 intor=intor, aosym=aosym, comp=comp, 590 max_memory=max_memory, ioblk_size=ioblk_size, 591 verbose=verbose, compact=compact) 592 return numpy.asarray(feri['eri_mo']) 593 594def general_iofree(mol, mo_coeffs, intor='int2e', aosym='s4', comp=None, 595 max_memory=MAX_MEMORY, ioblk_size=IOBLK_SIZE, 596 verbose=logger.WARN, compact=True): 597 r'''For the given four sets of orbitals, transfer arbitrary spherical AO 598 integrals to MO integrals on the fly. This function is a wrap for 599 :func:`ao2mo.outcore.general`. It's not really IO free. The returned MO 600 integrals are held in memory. For backward compatibility, it is used to 601 replace the non-existed function direct.general_iofree. 602 603 Args: 604 mol : :class:`Mole` object 605 AO integrals will be generated in terms of mol._atm, mol._bas, mol._env 606 mo_coeffs : 4-item list of ndarray 607 Four sets of orbital coefficients, corresponding to the four 608 indices of (ij|kl) 609 610 Kwargs 611 intor : str 612 Name of the 2-electron integral. Ref to :func:`getints_by_shell` 613 for the complete list of available 2-electron integral names 614 aosym : int or str 615 Permutation symmetry for the AO integrals 616 617 | 4 or '4' or 's4': 4-fold symmetry (default) 618 | '2ij' or 's2ij' : symmetry between i, j in (ij|kl) 619 | '2kl' or 's2kl' : symmetry between k, l in (ij|kl) 620 | 1 or '1' or 's1': no symmetry 621 | 'a4ij' : 4-fold symmetry with anti-symmetry between i, j in (ij|kl) (TODO) 622 | 'a4kl' : 4-fold symmetry with anti-symmetry between k, l in (ij|kl) (TODO) 623 | 'a2ij' : anti-symmetry between i, j in (ij|kl) (TODO) 624 | 'a2kl' : anti-symmetry between k, l in (ij|kl) (TODO) 625 626 comp : int 627 Components of the integrals, e.g. int2e_ip_sph has 3 components. 628 verbose : int 629 Print level 630 compact : bool 631 When compact is True, depending on the four oribital sets, the 632 returned MO integrals has (up to 4-fold) permutation symmetry. 633 If it's False, the function will abandon any permutation symmetry, 634 and return the "plain" MO integrals 635 636 Returns: 637 2D/3D MO-integral array. They may or may not have the permutation 638 symmetry, depending on the given orbitals, and the kwargs compact. If 639 the four sets of orbitals are identical, the MO integrals will at most 640 have 4-fold symmetry. 641 642 Examples: 643 644 >>> from pyscf import gto 645 >>> from pyscf import ao2mo 646 >>> import h5py 647 >>> def view(h5file, dataname='eri_mo'): 648 ... f5 = h5py.File(h5file, 'r') 649 ... print('dataset %s, shape %s' % (str(f5.keys()), str(f5[dataname].shape))) 650 ... f5.close() 651 >>> mol = gto.M(atom='O 0 0 0; H 0 1 0; H 0 0 1', basis='sto3g') 652 >>> mo1 = numpy.random.random((mol.nao_nr(), 10)) 653 >>> mo2 = numpy.random.random((mol.nao_nr(), 8)) 654 >>> mo3 = numpy.random.random((mol.nao_nr(), 6)) 655 >>> mo4 = numpy.random.random((mol.nao_nr(), 4)) 656 >>> eri1 = ao2mo.outcore.general_iofree(mol, (mo1,mo2,mo3,mo4)) 657 >>> print(eri1.shape) 658 (80, 24) 659 >>> eri1 = ao2mo.outcore.general_iofree(mol, (mo1,mo2,mo3,mo3)) 660 >>> print(eri1.shape) 661 (80, 21) 662 >>> eri1 = ao2mo.outcore.general_iofree(mol, (mo1,mo2,mo3,mo3), compact=False) 663 >>> print(eri1.shape) 664 (80, 36) 665 >>> eri1 = ao2mo.outcore.general_iofree(mol, (mo1,mo1,mo1,mo1), intor='int2e_ip1_sph', aosym='s1', comp=3) 666 >>> print(eri1.shape) 667 (3, 100, 100) 668 >>> eri1 = ao2mo.outcore.general_iofree(mol, (mo1,mo1,mo1,mo1), intor='int2e_ip1_sph', aosym='s2kl', comp=3) 669 >>> print(eri1.shape) 670 (3, 100, 55) 671 ''' 672 with lib.H5TmpFile() as feri: 673 general(mol, mo_coeffs, feri, dataname='eri_mo', 674 intor=intor, aosym=aosym, comp=comp, 675 max_memory=max_memory, ioblk_size=ioblk_size, 676 verbose=verbose, compact=compact) 677 return numpy.asarray(feri['eri_mo']) 678 679 680def iden_coeffs(mo1, mo2): 681 return (id(mo1) == id(mo2)) \ 682 or (mo1.shape==mo2.shape and numpy.allclose(mo1,mo2)) 683 684def prange(start, end, step): 685 for i in range(start, end, step): 686 yield i, min(i+step, end) 687 688def guess_e1bufsize(max_memory, ioblk_size, nij_pair, nao_pair, comp): 689 mem_words = max(1, max_memory * 1e6 / 8) 690# part of the max_memory is used to hold the AO integrals. The iobuf is the 691# buffer to temporary hold the transformed integrals before streaming to disk. 692# iobuf is then divided to small blocks (ioblk_words) and streamed to disk. 693 iobuf_words = max(int(mem_words//6), IOBUF_WORDS) 694 ioblk_words = int(min(ioblk_size*1e6/8, iobuf_words)) 695 696 e1buflen = int(mem_words*.66/(comp*(nij_pair*2+nao_pair))) 697 e1buflen = max(e1buflen, IOBUF_ROW_MIN) 698 return e1buflen, mem_words, iobuf_words, ioblk_words 699 700def guess_e2bufsize(ioblk_size, nrows, ncols): 701 e2buflen = int(min(ioblk_size*1e6/8/ncols, nrows)) 702 e2buflen = max(e2buflen//IOBUF_ROW_MIN, 1) * IOBUF_ROW_MIN 703 chunks = (IOBUF_ROW_MIN, ncols) 704 return e2buflen, chunks 705 706# based on the size of buffer, dynamic range of AO-shells for each buffer 707def guess_shell_ranges(mol, aosym, max_iobuf, max_aobuf=None, ao_loc=None, 708 compress_diag=True): 709 if ao_loc is None: ao_loc = mol.ao_loc_nr() 710 max_iobuf = max(1, max_iobuf) 711 712 dims = ao_loc[1:] - ao_loc[:-1] 713 dijs = (dims.reshape(-1,1) * dims) 714 nbas = dijs.shape[0] 715 716 if aosym: 717 if compress_diag: 718 #:for i in range(mol.nbas): 719 #: di = ao_loc[i+1] - ao_loc[i] 720 #: for j in range(i): 721 #: dj = ao_loc[j+1] - ao_loc[j] 722 #: lstdij.append(di*dj) 723 #: lstdij.append(di*(di+1)//2) 724 idx = numpy.arange(nbas) 725 dijs[idx,idx] = dims*(dims+1)//2 726 lstdij = dijs[numpy.tril_indices(nbas)] 727 else: 728 #:for i in range(mol.nbas): 729 #: di = ao_loc[i+1] - ao_loc[i] 730 #: for j in range(i+1): 731 #: dj = ao_loc[j+1] - ao_loc[j] 732 #: lstdij.append(di*dj) 733 lstdij = dijs[numpy.tril_indices(nbas)] 734 else: 735 #:for i in range(mol.nbas): 736 #: di = ao_loc[i+1] - ao_loc[i] 737 #: for j in range(mol.nbas): 738 #: dj = ao_loc[j+1] - ao_loc[j] 739 #: lstdij.append(di*dj) 740 lstdij = dijs.ravel() 741 742 dij_loc = numpy.append(0, numpy.cumsum(lstdij)) 743 ijsh_range = balance_partition(dij_loc, max_iobuf) 744 745 if max_aobuf is not None: 746 max_aobuf = max(1, max_aobuf) 747 def div_each_iobuf(ijstart, ijstop, buflen): 748 # to fill each iobuf, AO integrals may need to be fill to aobuf several times 749 return (ijstart, ijstop, buflen, 750 balance_partition(dij_loc, max_aobuf, ijstart, ijstop)) 751 ijsh_range = [div_each_iobuf(*x) for x in ijsh_range] 752 return ijsh_range 753 754def _stand_sym_code(sym): 755 if isinstance(sym, int): 756 return 's%d' % sym 757 elif 's' == sym[0] or 'a' == sym[0]: 758 return sym 759 else: 760 return 's' + sym 761 762def balance_segs(segs_lst, blksize, start_id=0, stop_id=None): 763 loc = numpy.append(0, numpy.cumsum(segs_lst)) 764 return balance_partition(loc, blksize, start_id, stop_id) 765def balance_partition(ao_loc, blksize, start_id=0, stop_id=None): 766 if stop_id is None: 767 stop_id = len(ao_loc) - 1 768 else: 769 stop_id = min(stop_id, start_id+len(ao_loc)-1) 770 displs = lib.misc._blocksize_partition(ao_loc[start_id:stop_id+1], blksize) 771 displs = [i+start_id for i in displs] 772 tasks = [] 773 for i0, i1 in zip(displs[:-1],displs[1:]): 774 tasks.append((i0, i1, ao_loc[i1]-ao_loc[i0])) 775 return tasks 776 777del(MAX_MEMORY) 778 779 780if __name__ == '__main__': 781 from pyscf import scf 782 from pyscf.ao2mo import addons 783 mol = gto.Mole() 784 mol.verbose = 5 785 #mol.output = 'out_outcore' 786 mol.atom = [ 787 ["O" , (0. , 0. , 0.)], 788 [1 , (0. , -0.757 , 0.587)], 789 [1 , (0. , 0.757 , 0.587)]] 790 791 mol.basis = {'H': 'cc-pvtz', 792 'O': 'cc-pvtz',} 793 mol.build() 794 nao = mol.nao_nr() 795 npair = nao*(nao+1)//2 796 797 rhf = scf.RHF(mol) 798 rhf.scf() 799 800 print(logger.process_clock()) 801 full(mol, rhf.mo_coeff, 'h2oeri.h5', max_memory=10, ioblk_size=5) 802 print(logger.process_clock()) 803 eri0 = incore.full(rhf._eri, rhf.mo_coeff) 804 feri = h5py.File('h2oeri.h5', 'r') 805 print('full', abs(eri0-feri['eri_mo']).sum()) 806 feri.close() 807 808 print(logger.process_clock()) 809 c = rhf.mo_coeff 810 general(mol, (c,c,c,c), 'h2oeri.h5', max_memory=10, ioblk_size=5) 811 print(logger.process_clock()) 812 feri = h5py.File('h2oeri.h5', 'r') 813 print('general', abs(eri0-feri['eri_mo']).sum()) 814 feri.close() 815 816 # set ijsame and klsame to False, then check 817 c = rhf.mo_coeff 818 n = c.shape[1] 819 general(mol, (c,c,c,c), 'h2oeri.h5', max_memory=10, ioblk_size=5, compact=False) 820 feri = h5py.File('h2oeri.h5', 'r') 821 eri1 = numpy.array(feri['eri_mo']).reshape(n,n,n,n) 822 eri1 = addons.restore(4, eri1, n) 823 print('general', abs(eri0-eri1).sum()) 824