def get_jk(agf2, eri, rdm1, with_j=True, with_k=True): ''' Get the J/K matrices. Args: eri : ndarray or H5 dataset Electronic repulsion integrals (NOT as _ChemistsERIs) rdm1 : 2D array Reduced density matrix Kwargs: with_j : bool Whether to compute J. Default value is True with_k : bool Whether to compute K. Default value is True Returns: tuple of ndarrays corresponding to J and K, if either are not requested then they are set to None. ''' if isinstance(eri, np.ndarray): vj, vk = _vhf.incore(eri, rdm1, with_j=with_j, with_k=with_k) else: nmo = rdm1.shape[0] npair = nmo * (nmo + 1) // 2 vj = vk = None if with_j: rdm1_tril = lib.pack_tril(rdm1 + np.tril(rdm1, k=-1)) vj = np.zeros_like(rdm1_tril) if with_k: vk = np.zeros_like(rdm1) blksize = _agf2.get_blksize(agf2.max_memory, (nmo * npair, nmo**3)) blksize = min(1, max(BLKMIN, blksize)) logger.debug1(agf2, 'blksize (ragf2.get_jk) = %d' % blksize) tril2sq = lib.square_mat_in_trilu_indices(nmo) for p0, p1 in lib.prange(0, nmo, blksize): idx = list(np.concatenate(tril2sq[p0:p1])) eri0 = eri[idx] # vj built in tril layout with scaled rdm1_tril if with_j: vj[idx] = np.dot(eri0, rdm1_tril) if with_k: eri0 = lib.unpack_tril(eri0, axis=-1) eri0 = eri0.reshape(p1 - p0, nmo, nmo, nmo) vk[p0:p1] = lib.einsum('ijkl,jk->il', eri0, rdm1) if with_j: vj = lib.unpack_tril(vj) return vj, vk
def _make_qmo_eris_outcore(agf2, eri, coeffs): ''' Returns H5 dataset ''' cput0 = (logger.process_clock(), logger.perf_counter()) log = logger.Logger(agf2.stdout, agf2.verbose) nmo = eri.nmo ci, cj, ca = coeffs ni = ci.shape[1] nj = cj.shape[1] na = ca.shape[1] npair = nmo * (nmo + 1) // 2 mask = get_frozen_mask(agf2) frozen = np.sum(~mask) # possible to have incore MO, outcore QMO if getattr(eri, 'feri', None) is None: eri.feri = lib.H5TmpFile() elif 'qmo' in eri.feri: del eri.feri['qmo'] eri.feri.create_dataset('qmo', (nmo - frozen, ni, nj, na), 'f8') blksize = _agf2.get_blksize(agf2.max_memory, (nmo * npair, nj * na, npair), (nmo * ni, nj * na)) blksize = min(nmo, max(BLKMIN, blksize)) log.debug1('blksize (ragf2._make_qmo_eris_outcore) = %d', blksize) tril2sq = lib.square_mat_in_trilu_indices(nmo) q1 = 0 for p0, p1 in lib.prange(0, nmo, blksize): if not np.any(mask[p0:p1]): # block is fully frozen continue inds = np.arange(p0, p1)[mask[p0:p1]] q0, q1 = q1, q1 + len(inds) idx = list(np.concatenate(tril2sq[inds])) buf = eri.eri[idx] # (blk, nmo, npair) buf = buf.reshape((q1 - q0) * nmo, -1) # (blk*nmo, npair) jasym, nja, cja, sja = ao2mo.incore._conc_mos(cj, ca, compact=True) buf = ao2mo._ao2mo.nr_e2(buf, cja, sja, 's2kl', 's1') buf = buf.reshape(q1 - q0, nmo, nj, na) buf = lib.einsum('xpja,pi->xija', buf, ci) eri.feri['qmo'][q0:q1] = np.asarray(buf, order='C') log.timer('QMO integral transformation', *cput0) return eri.feri['qmo']
def _make_qmo_eris_outcore(agf2, eri, coeffs_a, coeffs_b, spin=None): ''' Returns nested tuple of H5 dataset spin = None: ((aaaa, aabb), (bbaa, bbbb)) spin = 0: (aaaa, aabb) spin = 1: (bbbb, bbaa) ''' cput0 = (time.clock(), time.time()) log = logger.Logger(agf2.stdout, agf2.verbose) nmo = eri.nmo nmoa, nmob = nmo cxa = np.eye(nmoa) cxb = np.eye(nmob) mask = get_frozen_mask(agf2) frozena = np.sum(~mask[0]) frozenb = np.sum(~mask[1]) npaira, npairb = nmoa*(nmoa+1)//2, nmob*(nmob+1)//2 cia, cja, caa = coeffs_a cib, cjb, cab = coeffs_b nia, nja, naa = [x.shape[1] for x in coeffs_a] nib, njb, nab = [x.shape[1] for x in coeffs_b] # possible to have incore MO, outcore QMO if getattr(eri, 'feri', None) is None: eri.feri = lib.H5TmpFile() else: for key in ['aa', 'ab', 'ba', 'bb']: if 'qmo/%s'%key in eri.feri: del eri.feri['qmo/%s'%key] if spin is None or spin == 0: eri.feri.create_dataset('qmo/aa', (nmoa-frozena, nia, nja, naa), 'f8') eri.feri.create_dataset('qmo/ab', (nmoa-frozena, nia, njb, nab), 'f8') blksize = _agf2.get_blksize(agf2.max_memory, (nmoa**3, nmoa*nja*naa), (nmoa*nmob**2, nmoa*njb*nab)) blksize = min(nmoa, max(BLKMIN, blksize)) log.debug1('blksize (uagf2._make_qmo_eris_outcore) = %d', blksize) tril2sq = lib.square_mat_in_trilu_indices(nmoa) q1 = 0 for p0, p1 in lib.prange(0, nmoa, blksize): if not np.any(mask[0][p0:p1]): # block is fully frozen continue inds = np.arange(p0, p1)[mask[0][p0:p1]] q0, q1 = q1, q1 + len(inds) idx = list(np.concatenate(tril2sq[inds])) # aa buf = eri.eri_aa[idx] # (blk, nmoa, npaira) buf = buf.reshape((q1-q0)*nmoa, -1) # (blk*nmoa, npaira) jasym_aa, nja_aa, cja_aa, sja_aa = ao2mo.incore._conc_mos(cja, caa) buf = ao2mo._ao2mo.nr_e2(buf, cja_aa, sja_aa, 's2kl', 's1') buf = buf.reshape(q1-q0, nmoa, nja, naa) buf = lib.einsum('xpja,pi->xija', buf, cia) eri.feri['qmo/aa'][q0:q1] = np.asarray(buf, order='C') # ab buf = eri.eri_ab[idx] # (blk, nmoa, npairb) buf = buf.reshape((q1-q0)*nmob, -1) # (blk*nmoa, npairb) jasym_ab, nja_ab, cja_ab, sja_ab = ao2mo.incore._conc_mos(cjb, cab) buf = ao2mo._ao2mo.nr_e2(buf, cja_ab, sja_ab, 's2kl', 's1') buf = buf.reshape(q1-q0, nmoa, njb, nab) buf = lib.einsum('xpja,pi->xija', buf, cia) eri.feri['qmo/ab'][q0:q1] = np.asarray(buf, order='C') if spin is None or spin == 1: eri.feri.create_dataset('qmo/ba', (nmob-frozenb, nib, nja, naa), 'f8') eri.feri.create_dataset('qmo/bb', (nmob-frozenb, nib, njb, nab), 'f8') max_memory = agf2.max_memory - lib.current_memory()[0] blksize = int((max_memory/8e-6) / max(nmob**3+nmob*njb*nab, nmob*nmoa**2*nja*naa)) blksize = min(nmob, max(BLKMIN, blksize)) log.debug1('blksize (uagf2._make_qmo_eris_outcore) = %d', blksize) tril2sq = lib.square_mat_in_trilu_indices(nmob) q1 = 0 for p0, p1 in lib.prange(0, nmob, blksize): if not np.any(mask[1][p0:p1]): # block is fully frozen continue inds = np.arange(p0, p1)[mask[1][p0:p1]] q0, q1 = q1, q1 + len(inds) idx = list(np.concatenate(tril2sq[inds])) # ba buf = eri.eri_ba[idx] # (blk, nmob, npaira) buf = buf.reshape((q1-q0)*nmob, -1) # (blk*nmob, npaira) jasym_ba, nja_ba, cja_ba, sja_ba = ao2mo.incore._conc_mos(cja, caa) buf = ao2mo._ao2mo.nr_e2(buf, cja_ba, sja_ba, 's2kl', 's1') buf = buf.reshape(q1-q0, nmob, nja, naa) buf = lib.einsum('xpja,pi->xija', buf, cib) eri.feri['qmo/ba'][q0:q1] = np.asarray(buf, order='C') # bb buf = eri.eri_bb[idx] # (blk, nmob, npairb) buf = buf.reshape((q1-q0)*nmob, -1) # (blk*nmob, npairb) jasym_bb, nja_bb, cja_bb, sja_bb = ao2mo.incore._conc_mos(cjb, cab) buf = ao2mo._ao2mo.nr_e2(buf, cja_bb, sja_bb, 's2kl', 's1') buf = buf.reshape(q1-q0, nmob, njb, nab) buf = lib.einsum('xpja,pi->xija', buf, cib) eri.feri['qmo/bb'][q0:q1] = np.asarray(buf, order='C') if spin is None: qeri = ((eri.feri['qmo/aa'], eri.feri['qmo/ab']), (eri.feri['qmo/ba'], eri.feri['qmo/bb'])) elif spin == 0: qeri = (eri.feri['qmo/aa'], eri.feri['qmo/ab']) elif spin == 1: qeri = (eri.feri['qmo/bb'], eri.feri['qmo/ba']) log.timer('QMO integral transformation', *cput0) return qeri
def get_jk(agf2, eri, rdm1, with_j=True, with_k=True): ''' Get the J/K matrices. Args: eri : ndarray or H5 dataset Electronic repulsion integrals (NOT as _ChemistsERIs). In the case of no bra/ket symmetry, a tuple can be passed. rdm1 : 2D array Reduced density matrix Kwargs: with_j : bool Whether to compute J. Default value is True with_k : bool Whether to compute K. Default value is True Returns: tuple of ndarrays corresponding to J and K, if either are not requested then they are set to None. ''' nmo = rdm1.shape[0] npair = nmo * (nmo + 1) // 2 naux = agf2.with_df.get_naoaux() vj = vk = None if with_j: rdm1_tril = lib.pack_tril(rdm1 + np.tril(rdm1, k=-1)) vj = np.zeros((npair, )) if with_k: vk = np.zeros((nmo, nmo)) fdrv = ao2mo._ao2mo.libao2mo.AO2MOnr_e2_drv fmmm = ao2mo._ao2mo.libao2mo.AO2MOmmm_bra_nr_s2 ftrans = ao2mo._ao2mo.libao2mo.AO2MOtranse2_nr_s2 if isinstance(eri, tuple): bra, ket = eri else: bra = ket = eri blksize = _agf2.get_blksize(agf2.max_memory, (npair, npair, 1, nmo**2, nmo**2)) blksize = min(nmo, max(BLKMIN, blksize)) logger.debug1(agf2, 'blksize (dfragf2.get_jk) = %d' % blksize) buf = (np.empty((blksize, nmo, nmo)), np.empty((blksize, nmo, nmo))) for p0, p1 in mpi_helper.prange(0, naux, blksize): bra0 = bra[p0:p1] ket0 = ket[p0:p1] rho = np.dot(ket0, rdm1_tril) if with_j: vj += np.dot(rho, bra0) if with_k: buf1 = buf[0][:p1 - p0] fdrv(ftrans, fmmm, buf1.ctypes.data_as(ctypes.c_void_p), bra0.ctypes.data_as(ctypes.c_void_p), rdm1.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(p1 - p0), ctypes.c_int(nmo), (ctypes.c_int * 4)(0, nmo, 0, nmo), lib.c_null_ptr(), ctypes.c_int(0)) buf2 = lib.unpack_tril(ket0, out=buf[1]) buf1 = buf1.reshape(-1, nmo) buf2 = buf2.reshape(-1, nmo) vk = lib.dot(buf1.T, buf2, c=vk, beta=1) if with_j: mpi_helper.barrier() mpi_helper.allreduce_safe_inplace(vj) mpi_helper.barrier() vj = lib.unpack_tril(vj) if with_k: mpi_helper.barrier() mpi_helper.allreduce_safe_inplace(vk) return vj, vk