def trans_e1_outcore(mol, mo, ncore, ncas, erifile, max_memory=None, level=1, verbose=logger.WARN): time0 = (time.clock(), time.time()) if isinstance(verbose, logger.Logger): log = verbose else: log = logger.Logger(mol.stdout, verbose) log.debug1("trans_e1_outcore level %d max_memory %d", level, max_memory) nao, nmo = mo.shape nao_pair = nao * (nao + 1) // 2 nocc = ncore + ncas _tmpfile1 = tempfile.NamedTemporaryFile() faapp_buf = h5py.File(_tmpfile1.name) feri = h5py.File(erifile, "w") mo_c = numpy.asarray(mo, order="C") mo = numpy.asarray(mo, order="F") pashape = (0, nmo, ncore, ncas) papa_buf = numpy.zeros((nao, ncas, nmo * ncas)) j_pc = numpy.zeros((nmo, ncore)) k_pc = numpy.zeros((nmo, ncore)) mem_words = int(max(2000, max_memory - papa_buf.nbytes / 1e6) * 1e6 / 8) aobuflen = mem_words // (nao_pair + nocc * nmo) + 1 shranges = outcore.guess_shell_ranges(mol, aobuflen, aobuflen, "s4") ao2mopt = _ao2mo.AO2MOpt(mol, "cint2e_sph", "CVHFnr_schwarz_cond", "CVHFsetnr_direct_scf") ao_loc = numpy.array(mol.ao_loc_nr(), dtype=numpy.int32) log.debug("mem cache %.8g MB", mem_words * 8 / 1e6) ti0 = log.timer("Initializing trans_e1_outcore", *time0) nstep = len(shranges) paapp = 0 maxbuflen = max([x[2] for x in shranges]) bufs1 = numpy.empty((maxbuflen, nao_pair)) bufs2 = numpy.empty((maxbuflen, nmo * ncas)) bufs3 = numpy.empty((maxbuflen, nao * ncore)) # fmmm, ftrans, fdrv for level 1 fmmm = _fpointer("MCSCFhalfmmm_nr_s2_ket") ftrans = _fpointer("AO2MOtranse1_nr_s4") fdrv = getattr(libmcscf, "AO2MOnr_e2_drv") for istep, sh_range in enumerate(shranges): log.debug("[%d/%d], AO [%d:%d], len(buf) = %d", istep + 1, nstep, *(sh_range[:3])) buf = bufs1[: sh_range[2]] _ao2mo.nr_e1fill_("cint2e_sph", sh_range[:3], mol._atm, mol._bas, mol._env, "s4", 1, ao2mopt, buf) if log.verbose >= logger.DEBUG1: ti1 = log.timer("AO integrals buffer", *ti0) bufpa = bufs2[: sh_range[2]] _ao2mo.nr_e1_(buf, mo, pashape, "s4", "s1", vout=bufpa) # jc_pp, kc_pp if level == 1: # ppaa, papa and vhf, jcp, kcp if log.verbose >= logger.DEBUG1: ti1 = log.timer("buffer-pa", *ti1) buf1 = bufs3[: sh_range[2]] fdrv( ftrans, fmmm, buf1.ctypes.data_as(ctypes.c_void_p), buf.ctypes.data_as(ctypes.c_void_p), mo.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(sh_range[2]), ctypes.c_int(nao), ctypes.c_int(0), ctypes.c_int(nao), ctypes.c_int(0), ctypes.c_int(ncore), ctypes.POINTER(ctypes.c_void_p)(), ctypes.c_int(0), ) p0 = 0 for ij in range(sh_range[0], sh_range[1]): i, j = _ao2mo._extract_pair(ij) i0 = ao_loc[i] j0 = ao_loc[j] i1 = ao_loc[i + 1] j1 = ao_loc[j + 1] di = i1 - i0 dj = j1 - j0 if i == j: dij = di * (di + 1) // 2 buf = numpy.empty((di, di, nao * ncore)) idx = numpy.tril_indices(di) buf[idx] = buf1[p0 : p0 + dij] buf[idx[1], idx[0]] = buf1[p0 : p0 + dij] buf = buf.reshape(di, di, nao, ncore) mo1 = mo_c[i0:i1] tmp = numpy.einsum("uvpc,pc->uvc", buf, mo[:, :ncore]) tmp = pyscf.lib.dot(mo1.T, tmp.reshape(di, -1)) j_pc += numpy.einsum("vp,pvc->pc", mo1, tmp.reshape(nmo, di, ncore)) tmp = numpy.einsum("uvpc,uc->vcp", buf, mo1[:, :ncore]) tmp = pyscf.lib.dot(tmp.reshape(-1, nmo), mo).reshape(di, ncore, nmo) k_pc += numpy.einsum("vp,vcp->pc", mo1, tmp) else: dij = di * dj buf = buf1[p0 : p0 + dij].reshape(di, dj, nao, ncore) mo1 = mo_c[i0:i1] mo2 = mo_c[j0:j1] tmp = numpy.einsum("uvpc,pc->uvc", buf, mo[:, :ncore]) tmp = pyscf.lib.dot(mo1.T, tmp.reshape(di, -1)) j_pc += numpy.einsum("vp,pvc->pc", mo2, tmp.reshape(nmo, dj, ncore)) * 2 tmp = numpy.einsum("uvpc,uc->vcp", buf, mo1[:, :ncore]) tmp = pyscf.lib.dot(tmp.reshape(-1, nmo), mo).reshape(dj, ncore, nmo) k_pc += numpy.einsum("vp,vcp->pc", mo2, tmp) tmp = numpy.einsum("uvpc,vc->ucp", buf, mo2[:, :ncore]) tmp = pyscf.lib.dot(tmp.reshape(-1, nmo), mo).reshape(di, ncore, nmo) k_pc += numpy.einsum("up,ucp->pc", mo1, tmp) p0 += dij if log.verbose >= logger.DEBUG1: ti1 = log.timer("j_cp and k_cp", *ti1) if log.verbose >= logger.DEBUG1: ti1 = log.timer("half transformation of the buffer", *ti1) # ppaa, papa faapp_buf[str(istep)] = bufpa.reshape(sh_range[2], nmo, ncas)[:, ncore:nocc].reshape(-1, ncas ** 2).T p0 = 0 for ij in range(sh_range[0], sh_range[1]): i, j = _ao2mo._extract_pair(ij) i0 = ao_loc[i] j0 = ao_loc[j] i1 = ao_loc[i + 1] j1 = ao_loc[j + 1] di = i1 - i0 dj = j1 - j0 if i == j: dij = di * (di + 1) // 2 buf1 = numpy.empty((di, di, nmo * ncas)) idx = numpy.tril_indices(di) buf1[idx] = bufpa[p0 : p0 + dij] buf1[idx[1], idx[0]] = bufpa[p0 : p0 + dij] else: dij = di * dj buf1 = bufpa[p0 : p0 + dij].reshape(di, dj, -1) mo1 = mo[j0:j1, ncore:nocc].copy() for i in range(di): pyscf.lib.dot(mo1.T, buf1[i], 1, papa_buf[i0 + i], 1) mo1 = mo[i0:i1, ncore:nocc].copy() buf1 = pyscf.lib.dot(mo1.T, buf1.reshape(di, -1)) papa_buf[j0:j1] += buf1.reshape(ncas, dj, -1).transpose(1, 0, 2) p0 += dij if log.verbose >= logger.DEBUG1: ti1 = log.timer("ppaa and papa buffer", *ti1) ti0 = log.timer("gen AO/transform MO [%d/%d]" % (istep + 1, nstep), *ti0) buf = buf1 = bufs1 = bufs2 = bufs3 = bufpa = None time1 = log.timer("mc_ao2mo pass 1", *time0) log.debug1("Half transformation done. Current memory %d", pyscf.lib.current_memory()[0]) nblk = int(max(8, min(nmo, max(2000, max_memory * 1e6 / 8 - papa_buf.size) / (ncas ** 2 * nmo)))) log.debug1("nblk for papa = %d", nblk) dset = feri.create_dataset("papa", (nmo, ncas, nmo, ncas), "f8") for i0, i1 in prange(0, nmo, nblk): tmp = pyscf.lib.dot(mo[:, i0:i1].T, papa_buf.reshape(nao, -1)) dset[i0:i1] = tmp.reshape(i1 - i0, ncas, nmo, ncas) papa_buf = tmp = None time1 = log.timer("papa pass 2", *time1) tmp = numpy.empty((ncas ** 2, nao_pair)) p0 = 0 for istep, sh_range in enumerate(shranges): tmp[:, p0 : p0 + sh_range[2]] = faapp_buf[str(istep)] p0 += sh_range[2] nblk = int(max(8, min(nmo, max(2000, max_memory * 1e6 / 8 - tmp.size) / (ncas ** 2 * nmo) - 1))) log.debug1("nblk for ppaa = %d", nblk) dset = feri.create_dataset("ppaa", (nmo, nmo, ncas, ncas), "f8") for i0, i1 in prange(0, nmo, nblk): tmp1 = _ao2mo.nr_e2_(tmp, mo, (i0, i1 - i0, 0, nmo), "s4", "s1", ao_loc=ao_loc) tmp1 = tmp1.reshape(ncas, ncas, i1 - i0, nmo) for j in range(i1 - i0): dset[i0 + j] = tmp1[:, :, j].transpose(2, 0, 1) tmp = tmp1 = None time1 = log.timer("ppaa pass 2", *time1) faapp_buf.close() feri.close() _tmpfile1 = None time0 = log.timer("mc_ao2mo", *time0) return j_pc, k_pc
def half_e1(mol, mo_coeffs, swapfile, intor='cint2e_sph', aosym='s4', comp=1, max_memory=2000, ioblk_size=256, verbose=logger.WARN, compact=True, ao2mopt=None): r'''Half transform arbitrary spherical AO integrals to MO integrals for the given two sets of orbitals Args: mol : :class:`Mole` object AO integrals will be generated in terms of mol._atm, mol._bas, mol._env mo_coeff : ndarray Transform (ij|kl) with the same set of orbitals. swapfile : str or h5py File or h5py Group object To store the transformed integrals, in HDF5 format. The transformed integrals are saved in blocks. Kwargs intor : str Name of the 2-electron integral. Ref to :func:`getints_by_shell` for the complete list of available 2-electron integral names aosym : int or str Permutation symmetry for the AO integrals | 4 or '4' or 's4': 4-fold symmetry (default) | '2ij' or 's2ij' : symmetry between i, j in (ij|kl) | '2kl' or 's2kl' : symmetry between k, l in (ij|kl) | 1 or '1' or 's1': no symmetry | 'a4ij' : 4-fold symmetry with anti-symmetry between i, j in (ij|kl) (TODO) | 'a4kl' : 4-fold symmetry with anti-symmetry between k, l in (ij|kl) (TODO) | 'a2ij' : anti-symmetry between i, j in (ij|kl) (TODO) | 'a2kl' : anti-symmetry between k, l in (ij|kl) (TODO) comp : int Components of the integrals, e.g. cint2e_ip_sph has 3 components. verbose : int Print level max_memory : float or int The maximum size of cache to use (in MB), large cache may **not** improve performance. ioblk_size : float or int The block size for IO, large block size may **not** improve performance verbose : int Print level compact : bool When compact is True, depending on the four oribital sets, the returned MO integrals has (up to 4-fold) permutation symmetry. If it's False, the function will abandon any permutation symmetry, and return the "plain" MO integrals ao2mopt : :class:`AO2MOpt` object Precomputed data to improve perfomance Returns: None ''' time0 = (time.clock(), time.time()) if isinstance(verbose, logger.Logger): log = verbose else: log = logger.Logger(mol.stdout, verbose) ijsame = compact and iden_coeffs(mo_coeffs[0], mo_coeffs[1]) nmoi = mo_coeffs[0].shape[1] nmoj = mo_coeffs[1].shape[1] nao = mo_coeffs[0].shape[0] aosym = _stand_sym_code(aosym) if aosym in ('s4', 's2ij'): nao_pair = nao * (nao+1) // 2 else: nao_pair = nao * nao if compact and ijsame and aosym in ('s4', 's2ij'): log.debug('i-mo == j-mo') ijmosym = 's2' nij_pair = nmoi*(nmoi+1) // 2 moij = numpy.asarray(mo_coeffs[0], order='F') ijshape = (0, nmoi, 0, nmoi) else: ijmosym = 's1' nij_pair = nmoi*nmoj moij = numpy.asarray(numpy.hstack((mo_coeffs[0],mo_coeffs[1])), order='F') ijshape = (0, nmoi, nmoi, nmoj) e1buflen, mem_words, iobuf_words, ioblk_words = \ guess_e1bufsize(max_memory, ioblk_size, nij_pair, nao_pair, comp) # The buffer to hold AO integrals in C code, see line (@) aobuflen = int((mem_words - iobuf_words) // (nao_pair*comp)) shranges = guess_shell_ranges(mol, e1buflen, aobuflen, aosym) if ao2mopt is None: if intor == 'cint2e_sph': ao2mopt = _ao2mo.AO2MOpt(mol, intor, 'CVHFnr_schwarz_cond', 'CVHFsetnr_direct_scf') else: ao2mopt = _ao2mo.AO2MOpt(mol, intor) log.debug('step1: tmpfile %.8g MB', nij_pair*nao_pair*8/1e6) log.debug('step1: (ij,kl) = (%d,%d), mem cache %.8g MB, iobuf %.8g MB', nij_pair, nao_pair, mem_words*8/1e6, iobuf_words*8/1e6) if isinstance(swapfile, str): fswap = h5py.File(swapfile, 'w') else: fswap = swapfile for icomp in range(comp): g = fswap.create_group(str(icomp)) # for h5py old version # transform e1 ti0 = log.timer('Initializing ao2mo.outcore.half_e1', *time0) nstep = len(shranges) maxbuflen = max([x[2] for x in shranges]) bufs1 = numpy.empty((comp*maxbuflen,nao_pair)) bufs2 = numpy.empty((comp*maxbuflen,nij_pair)) for istep,sh_range in enumerate(shranges): log.debug('step 1 [%d/%d], AO [%d:%d], len(buf) = %d', \ istep+1, nstep, *(sh_range[:3])) buflen = sh_range[2] iobuf = bufs2[:comp*buflen].reshape(comp,buflen,nij_pair) nmic = len(sh_range[3]) p0 = 0 for imic, aoshs in enumerate(sh_range[3]): log.debug1(' fill iobuf micro [%d/%d], AO [%d:%d], len(aobuf) = %d', \ imic+1, nmic, *aoshs) buf = bufs1[:comp*aoshs[2]] # (@) _ao2mo.nr_e1fill_(intor, aoshs, mol._atm, mol._bas, mol._env, aosym, comp, ao2mopt, out=buf) buf = _ao2mo.nr_e1_(buf, moij, ijshape, aosym, ijmosym) iobuf[:,p0:p0+aoshs[2]] = buf.reshape(comp,aoshs[2],-1) p0 += aoshs[2] ti2 = log.timer('gen AO/transform MO [%d/%d]'%(istep+1,nstep), *ti0) e2buflen, chunks = guess_e2bufsize(ioblk_size, nij_pair, buflen) for icomp in range(comp): _transpose_to_h5g(fswap, '%d/%d'%(icomp,istep), iobuf[icomp], e2buflen, None) ti0 = log.timer('transposing to disk', *ti2) bufs1 = bufs2 = None if isinstance(swapfile, str): fswap.close() return swapfile
def light_e1_outcore(mol, mo, ncore, ncas, max_memory=None, approx=1, verbose=logger.WARN): time0 = (time.clock(), time.time()) if isinstance(verbose, logger.Logger): log = verbose else: log = logger.Logger(mol.stdout, verbose) nao, nmo = mo.shape nao_pair = nao*(nao+1)//2 nocc = ncore + ncas aapp_buf = numpy.empty((nao_pair,ncas,ncas)) appa_buf = numpy.zeros((ncas,nao,nmo*ncas)) max_memory -= (aapp_buf.nbytes+appa_buf.nbytes) / 1e6 mo = numpy.asarray(mo, order='F') nao, nmo = mo.shape pashape = (0, nmo, ncore, ncas) if approx == 1: jc = numpy.empty((nao,nao,ncore)) kc = numpy.zeros((nao,nao,ncore)) else: dm_core = numpy.dot(mo[:,:ncore], mo[:,:ncore].T) * 2 jc = numpy.zeros((nao,nao)) kc = numpy.zeros((nao,nao)) max_memory -= (jc.nbytes+kc.nbytes) / 1e6 mem_words = int(max(1000,max_memory)*1e6/8) aobuflen = mem_words//(nao_pair+nocc*nmo) + 1 shranges = outcore.guess_shell_ranges(mol, aobuflen, aobuflen, 's4') ao2mopt = _ao2mo.AO2MOpt(mol, 'cint2e_sph', 'CVHFnr_schwarz_cond', 'CVHFsetnr_direct_scf') ao_loc = numpy.array(mol.ao_loc_nr(), dtype=numpy.int32) log.debug('mem cache %.8g MB', mem_words*8/1e6) ti0 = log.timer('Initializing light_e1_outcore', *time0) nstep = len(shranges) paapp = 0 maxbuflen = max([x[2] for x in shranges]) bufs1 = numpy.empty((maxbuflen, nao_pair)) bufs2 = numpy.empty((maxbuflen, pashape[1]*pashape[3])) bufs3 = numpy.empty((maxbuflen, nao*ncore)) for istep,sh_range in enumerate(shranges): log.debug('[%d/%d], AO [%d:%d], len(buf) = %d', istep+1, nstep, *(sh_range[:3])) buf = bufs1[:sh_range[2]] _ao2mo.nr_e1fill_('cint2e_sph', sh_range[:3], mol._atm, mol._bas, mol._env, 's4', 1, ao2mopt, buf) if log.verbose >= logger.DEBUG1: ti1 = log.timer('AO integrals buffer', *ti0) bufpa = bufs2[:sh_range[2]] # jc_pp, kc_pp if approx == 1: # aapp, appa and vhf, jcp, kcp _ao2mo.nr_e1_(buf, mo, pashape, 's4', 's1', vout=bufpa) if log.verbose >= logger.DEBUG1: ti1 = log.timer('buffer-pa', *ti1) buf1 = bufs3[:sh_range[2]] fmmm = _fpointer('MCSCFhalfmmm_nr_s2_ket') ftrans = _fpointer('AO2MOtranse1_nr_s4') fdrv = getattr(libmcscf, 'AO2MOnr_e2_drv') fdrv(ftrans, fmmm, buf1.ctypes.data_as(ctypes.c_void_p), buf.ctypes.data_as(ctypes.c_void_p), mo.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(sh_range[2]), ctypes.c_int(nao), ctypes.c_int(0), ctypes.c_int(nao), ctypes.c_int(0), ctypes.c_int(ncore), ctypes.POINTER(ctypes.c_void_p)(), ctypes.c_int(0)) p0 = 0 for ij in range(sh_range[0], sh_range[1]): i,j = _ao2mo._extract_pair(ij) i0 = ao_loc[i] j0 = ao_loc[j] i1 = ao_loc[i+1] j1 = ao_loc[j+1] di = i1 - i0 dj = j1 - j0 if i == j: dij = di * (di+1) // 2 buf = numpy.empty((di,di,nao*ncore)) idx = numpy.tril_indices(di) buf[idx] = buf1[p0:p0+dij] buf[idx[1],idx[0]] = buf1[p0:p0+dij] buf = buf.reshape(di,di,nao,ncore) jc[i0:i1,j0:j1] = numpy.einsum('uvpc,pc->uvc', buf, mo[:,:ncore]) kc[j0:j1] += numpy.einsum('uvpc,uc->vpc', buf, mo[i0:i1,:ncore]) else: dij = di * dj buf = buf1[p0:p0+dij].reshape(di,dj,nao,ncore) jc[i0:i1,j0:j1] = numpy.einsum('uvpc,pc->uvc', buf, mo[:,:ncore]) jc[j0:j1,i0:i1] = jc[i0:i1,j0:j1].transpose(1,0,2) kc[j0:j1] += numpy.einsum('uvpc,uc->vpc', buf, mo[i0:i1,:ncore]) kc[i0:i1] += numpy.einsum('uvpc,vc->upc', buf, mo[j0:j1,:ncore]) p0 += dij if log.verbose >= logger.DEBUG1: ti1 = log.timer('jc and kc buffer', *ti1) elif approx == 2: # aapp, appa, vhf fdrv = libmcscf.MCSCFnrs4_aapp_jk fdrv(buf.ctypes.data_as(ctypes.c_void_p), bufpa.ctypes.data_as(ctypes.c_void_p), mo.ctypes.data_as(ctypes.c_void_p), dm_core.ctypes.data_as(ctypes.c_void_p), jc.ctypes.data_as(ctypes.c_void_p), kc.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(sh_range[0]), ctypes.c_int(sh_range[1]-sh_range[0]), ctypes.c_int(pashape[0]), ctypes.c_int(pashape[1]), ctypes.c_int(pashape[2]), ctypes.c_int(pashape[3]), ao_loc.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.nbas)) if log.verbose >= logger.DEBUG1: ti1 = log.timer('bufpa + core vj and vk', *ti1) else: # aapp, appa _ao2mo.nr_e1_(buf, mo, pashape, 's4', 's1', vout=bufpa) # aapp, appa aapp_buf[paapp:paapp+sh_range[2]] = \ bufpa.reshape(sh_range[2],nmo,ncas)[:,ncore:nocc] paapp += sh_range[2] p0 = 0 for ij in range(sh_range[0], sh_range[1]): i,j = _ao2mo._extract_pair(ij) i0 = ao_loc[i] j0 = ao_loc[j] i1 = ao_loc[i+1] j1 = ao_loc[j+1] di = i1 - i0 dj = j1 - j0 if i == j: dij = di * (di+1) // 2 buf1 = numpy.empty((di,di,nmo*ncas)) idx = numpy.tril_indices(di) buf1[idx] = bufpa[p0:p0+dij] buf1[idx[1],idx[0]] = bufpa[p0:p0+dij] buf1 = buf1.reshape(di,-1) else: dij = di * dj buf1 = bufpa[p0:p0+dij].reshape(di,dj,-1) mo1 = mo[j0:j1,ncore:nocc].copy() for i in range(di): appa_buf[:,i0+i] += pyscf.lib.dot(mo1.T, buf1[i]) buf1 = bufpa[p0:p0+dij].reshape(di,-1) mo1 = mo[i0:i1,ncore:nocc].copy() appa_buf[:,j0:j1] += pyscf.lib.dot(mo1.T, buf1).reshape(ncas,dj,-1) p0 += dij if log.verbose >= logger.DEBUG1: ti1 = log.timer('aapp and appa buffer', *ti1) ti0 = log.timer('gen AO/transform MO [%d/%d]'%(istep+1,nstep), *ti0) bufs1 = bufs2 = bufs3 = None aapp_buf = pyscf.lib.transpose(aapp_buf.reshape(nao_pair,-1)) aapp = _ao2mo.nr_e2_(aapp_buf, mo, (0,nmo,0,nmo), 's4', 's1', ao_loc=ao_loc) aapp = aapp.reshape(ncas,ncas,nmo,nmo) aapp_buf = None if nao == nmo: appa = appa_buf else: appa = numpy.empty((ncas,nao,nmo*ncas)) for i in range(ncas): appa[i] = numpy.dot(mo.T, appa_buf[i].reshape(nao,-1)) appa = appa.reshape(ncas,nmo,nmo,ncas) appa_buf = None if approx == 1: vhf_c = numpy.einsum('ijc->ij', jc)*2 - numpy.einsum('ijc->ij', kc) vhf_c = reduce(numpy.dot, (mo.T, vhf_c, mo)) j_cp = numpy.dot(mo.T, jc.reshape(nao,-1)).reshape(nao,nao,ncore) j_cp = numpy.einsum('pj,jpi->ij', mo, j_cp) k_cp = numpy.dot(mo.T, kc.reshape(nao,-1)).reshape(nao,nao,ncore) k_cp = numpy.einsum('pj,jpi->ij', mo, k_cp) elif approx == 2: jc = pyscf.lib.hermi_triu(jc, hermi=1, inplace=True) kc = pyscf.lib.hermi_triu(kc, hermi=1, inplace=True) vhf_c = reduce(numpy.dot, (mo.T, jc-kc*.5, mo)) j_cp = k_cp = None elif approx == 3: vhf_c = j_cp = k_cp = None time0 = log.timer('mc_ao2mo', *time0) return vhf_c, j_cp, k_cp, aapp, appa