def test_unpack_row(self):
row = numpy.arange(28.)
ref = lib.unpack_tril(row)[4]
self.assertTrue(numpy.array_equal(ref, lib.unpack_row(row, 4)))
row = numpy.arange(28, dtype=numpy.int32)
ref = lib.unpack_tril(row)[4]
a = lib.unpack_row(row, 4)
self.assertTrue(numpy.array_equal(ref, a))
self.assertTrue(a.dtype == numpy.int32)
def test_unpack_row(self):
row = numpy.arange(28.)
ref = lib.unpack_tril(row)[4]
self.assertTrue(numpy.array_equal(ref, lib.unpack_row(row, 4)))
row = numpy.arange(28, dtype=numpy.int32)
ref = lib.unpack_tril(row)[4]
a = lib.unpack_row(row, 4)
self.assertTrue(numpy.array_equal(ref, a))
self.assertTrue(a.dtype == numpy.int32)
def general(eri, mo_coeffs, erifile, dataname='eri_mo',
ioblk_size=IOBLK_SIZE, compact=True, verbose=logger.NOTE):
'''For the given four sets of orbitals, transfer arbitrary spherical AO
integrals to MO integrals on disk.
Args:
eri : 8-fold reduced eri vector
mo_coeffs : 4-item list of ndarray
Four sets of orbital coefficients, corresponding to the four
indices of (ij|kl)
erifile : str or h5py File or h5py Group object
To store the transformed integrals, in HDF5 format.
Kwargs
dataname : str
The dataset name in the erifile (ref the hierarchy of HDF5 format
http://www.hdfgroup.org/HDF5/doc1.6/UG/09_Groups.html). By assigning
different dataname, the existed integral file can be reused. If
the erifile contains the dataname, the new integrals data will
overwrite the old one.
ioblk_size : float or int
The block size for IO, large block size may **not** improve performance
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
Pseudocode / algorithm:
u = mu
v = nu
l = lambda
o = sigma
Assume eri's are 8-fold reduced.
nij/nkl_pair = npair or i*j/k*l if only transforming a subset
First half transform:
Initialize half_eri of size (nij_pair,npair)
For lo = 1 -> npair
Unpack row lo
Unpack row lo to matrix E_{uv}^{lo}
Transform C_ui^+*E*C_nj -> E_{ij}^{lo}
Ravel or pack E_{ij}^{lo}
Save E_{ij}^{lo} -> half_eri[:,lo]
Second half transform:
Initialize h5d_eri of size (nij_pair,nkl_pair)
For ij = 1 -> nij_pair
Load and unpack half_eri[ij,:] -> E_{lo}^{ij}
Transform C_{lk}E_{lo}^{ij}C_{ol} -> E_{kl}^{ij}
Repack E_{kl}^{ij}
Save E_{kl}^{ij} -> h5d_eri[ij,:]
Each matrix is indexed by the composite index ij x kl, where ij/kl is
either npair or ixj/kxl, if only a subset of MOs are being transformed.
Since entire rows or columns need to be read in, the arrays are chunked
such that IOBLK_SIZE = row/col x chunking col/row. For example, for the
first half transform, we would save in nij_pair x IOBLK_SIZE/nij_pair,
then load in IOBLK_SIZE/nkl_pair x npair for the second half transform.
------ kl ----->
|jxl
|
ij
|
|
v
As a first guess, the chunking size is jxl. If the super-rows/cols are
larger than IOBLK_SIZE, then the chunk rectangle jxl is trimmed
accordingly. The pathological limiting case is where the dimensions
nao_pair, nij_pair, or nkl_pair are so large that the arrays are
chunked 1x1, in which case IOBLK_SIZE needs to be increased.
'''
log = logger.new_logger(None, verbose)
log.info('******** ao2mo disk, custom eri ********')
nmoi = mo_coeffs[0].shape[1]
nmoj = mo_coeffs[1].shape[1]
nmok = mo_coeffs[2].shape[1]
nmol = mo_coeffs[3].shape[1]
nao = mo_coeffs[0].shape[0]
nao_pair = nao*(nao+1) // 2
if compact and iden_coeffs(mo_coeffs[0], mo_coeffs[1]):
ij_red = False
nij_pair = nmoi*(nmoi+1) // 2
else:
ij_red = True
nij_pair = nmoi*nmoj
if compact and iden_coeffs(mo_coeffs[2], mo_coeffs[3]):
kl_red = False
nkl_pair = nmok*(nmok+1) // 2
else:
kl_red = True
nkl_pair = nmok*nmol
chunks_half = (max(1, numpy.minimum(int(ioblk_size//(nao_pair*f8_size)),nmoj)),
max(1, numpy.minimum(int(ioblk_size//(nij_pair*f8_size)),nmol)))
'''
ideally, the final transformed eris should have a chunk of nmoj x nmol to
optimize read operations. However, I'm chunking the row size so that the
write operations during the transform can be done as fast as possible.
'''
chunks_full = (numpy.minimum(int(ioblk_size//(nkl_pair*f8_size)),nmoj),nmol)
if isinstance(erifile, str):
if h5py.is_hdf5(erifile):
feri = h5py.File(erifile)
if dataname in feri:
del(feri[dataname])
else:
feri = h5py.File(erifile,'w',libver='latest')
else:
assert(isinstance(erifile, h5py.Group))
feri = erifile
h5d_eri = feri.create_dataset(dataname,(nij_pair,nkl_pair),'f8',chunks=chunks_full)
feri_swap = lib.H5TmpFile(libver='latest')
half_eri = feri_swap.create_dataset(dataname,(nij_pair,nao_pair),'f8',chunks=chunks_half)
log.debug('Memory information:')
log.debug(' IOBLK_SIZE (MB): {}'.format(ioblk_size))
log.debug(' jxl {}x{}, half eri chunk dim {}x{}'.format(nmoj,nmol,chunks_half[0],chunks_half[1]))
log.debug(' jxl {}x{}, full eri chunk dim {}x{}'.format(nmoj,nmol,chunks_full[0],chunks_full[1]))
log.debug(' Final disk eri size (MB): {:.3g}, chunked {:.3g}'
.format(nij_pair*nkl_pair*f8_size,numpy.prod(chunks_full)*f8_size))
log.debug(' Half transformed eri size (MB): {:.3g}, chunked {:.3g}'
.format(nij_pair*nao_pair*f8_size,numpy.prod(chunks_half)*f8_size))
log.debug(' RAM buffer for half transform (MB): {:.3g}'
.format(nij_pair*chunks_half[1]*f8_size*2))
log.debug(' RAM buffer for full transform (MB): {:.3g}'
.format(f8_size*chunks_full[0]*nkl_pair*2 + chunks_half[0]*nao_pair*f8_size*2))
def save1(piece,buf):
start = piece*chunks_half[1]
stop = (piece+1)*chunks_half[1]
if stop > nao_pair:
stop = nao_pair
half_eri[:,start:stop] = buf[:,:stop-start]
return
def load2(piece):
start = piece*chunks_half[0]
stop = (piece+1)*chunks_half[0]
if stop > nij_pair:
stop = nij_pair
if start >= nij_pair:
start = stop - 1
return half_eri[start:stop,:]
def prefetch2(piece):
start = piece*chunks_half[0]
stop = (piece+1)*chunks_half[0]
if stop > nij_pair:
stop = nij_pair
if start >= nij_pair:
start = stop - 1
buf_prefetch[:stop-start,:] = half_eri[start:stop,:]
return
def save2(piece,buf):
start = piece*chunks_full[0]
stop = (piece+1)*chunks_full[0]
if stop > nij_pair:
stop = nij_pair
h5d_eri[start:stop,:] = buf[:stop-start,:]
return
# transform \mu\nu -> ij
cput0 = time.clock(), time.time()
Cimu = mo_coeffs[0].conj().transpose()
buf_write = numpy.empty((nij_pair,chunks_half[1]))
buf_out = numpy.empty_like(buf_write)
wpiece = 0
with lib.call_in_background(save1) as async_write:
for lo in range(nao_pair):
if lo % chunks_half[1] == 0 and lo > 0:
#save1(wpiece,buf_write)
buf_out, buf_write = buf_write, buf_out
async_write(wpiece,buf_out)
wpiece += 1
buf = lib.unpack_row(eri,lo)
uv = lib.unpack_tril(buf)
uv = Cimu.dot(uv).dot(mo_coeffs[1])
if ij_red:
ij = numpy.ravel(uv) # grabs by row
else:
ij = lib.pack_tril(uv)
buf_write[:,lo % chunks_half[1]] = ij
# final write operation & cleanup
save1(wpiece,buf_write)
log.timer('(uv|lo) -> (ij|lo)', *cput0)
uv = None
ij = None
buf = None
# transform \lambda\sigma -> kl
cput1 = time.clock(), time.time()
Cklam = mo_coeffs[2].conj().transpose()
buf_write = numpy.empty((chunks_full[0],nkl_pair))
buf_out = numpy.empty_like(buf_write)
buf_read = numpy.empty((chunks_half[0],nao_pair))
buf_prefetch = numpy.empty_like(buf_read)
rpiece = 0
wpiece = 0
with lib.call_in_background(save2,prefetch2) as (async_write,prefetch):
buf_read = load2(rpiece)
prefetch(rpiece+1)
for ij in range(nij_pair):
if ij % chunks_full[0] == 0 and ij > 0:
#save2(wpiece,buf_write)
buf_out, buf_write = buf_write, buf_out
async_write(wpiece,buf_out)
wpiece += 1
if ij % chunks_half[0] == 0 and ij > 0:
#buf_read = load2(rpiece)
buf_read, buf_prefetch = buf_prefetch, buf_read
rpiece += 1
prefetch(rpiece+1)
lo = lib.unpack_tril(buf_read[ij % chunks_half[0],:])
lo = Cklam.dot(lo).dot(mo_coeffs[3])
if kl_red:
kl = numpy.ravel(lo)
else:
kl = lib.pack_tril(lo)
buf_write[ij % chunks_full[0],:] = kl
save2(wpiece,buf_write)
log.timer('(ij|lo) -> (ij|kl)', *cput1)
if isinstance(erifile, str):
feri.close()
return erifile
def general(eri,
mo_coeffs,
erifile,
dataname='eri_mo',
ioblk_size=IOBLK_SIZE,
compact=True,
verbose=logger.NOTE):
'''For the given four sets of orbitals, transfer arbitrary spherical AO
integrals to MO integrals on disk.
Args:
eri : 8-fold reduced eri vector
mo_coeffs : 4-item list of ndarray
Four sets of orbital coefficients, corresponding to the four
indices of (ij|kl)
erifile : str or h5py File or h5py Group object
To store the transformed integrals, in HDF5 format.
Kwargs
dataname : str
The dataset name in the erifile (ref the hierarchy of HDF5 format
http://www.hdfgroup.org/HDF5/doc1.6/UG/09_Groups.html). By assigning
different dataname, the existed integral file can be reused. If
the erifile contains the dataname, the new integrals data will
overwrite the old one.
ioblk_size : float or int
The block size for IO, large block size may **not** improve performance
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
Pseudocode / algorithm:
u = mu
v = nu
l = lambda
o = sigma
Assume eri's are 8-fold reduced.
nij/nkl_pair = npair or i*j/k*l if only transforming a subset
First half transform:
Initialize half_eri of size (nij_pair,npair)
For lo = 1 -> npair
Unpack row lo
Unpack row lo to matrix E_{uv}^{lo}
Transform C_ui^+*E*C_nj -> E_{ij}^{lo}
Ravel or pack E_{ij}^{lo}
Save E_{ij}^{lo} -> half_eri[:,lo]
Second half transform:
Initialize h5d_eri of size (nij_pair,nkl_pair)
For ij = 1 -> nij_pair
Load and unpack half_eri[ij,:] -> E_{lo}^{ij}
Transform C_{lk}E_{lo}^{ij}C_{ol} -> E_{kl}^{ij}
Repack E_{kl}^{ij}
Save E_{kl}^{ij} -> h5d_eri[ij,:]
Each matrix is indexed by the composite index ij x kl, where ij/kl is
either npair or ixj/kxl, if only a subset of MOs are being transformed.
Since entire rows or columns need to be read in, the arrays are chunked
such that IOBLK_SIZE = row/col x chunking col/row. For example, for the
first half transform, we would save in nij_pair x IOBLK_SIZE/nij_pair,
then load in IOBLK_SIZE/nkl_pair x npair for the second half transform.
------ kl ----->
|jxl
|
ij
|
|
v
As a first guess, the chunking size is jxl. If the super-rows/cols are
larger than IOBLK_SIZE, then the chunk rectangle jxl is trimmed
accordingly. The pathological limiting case is where the dimensions
nao_pair, nij_pair, or nkl_pair are so large that the arrays are
chunked 1x1, in which case IOBLK_SIZE needs to be increased.
'''
log = logger.new_logger(None, verbose)
log.info('******** ao2mo disk, custom eri ********')
nmoi = mo_coeffs[0].shape[1]
nmoj = mo_coeffs[1].shape[1]
nmok = mo_coeffs[2].shape[1]
nmol = mo_coeffs[3].shape[1]
nao = mo_coeffs[0].shape[0]
nao_pair = nao * (nao + 1) // 2
if compact and iden_coeffs(mo_coeffs[0], mo_coeffs[1]):
ij_red = False
nij_pair = nmoi * (nmoi + 1) // 2
else:
ij_red = True
nij_pair = nmoi * nmoj
if compact and iden_coeffs(mo_coeffs[2], mo_coeffs[3]):
kl_red = False
nkl_pair = nmok * (nmok + 1) // 2
else:
kl_red = True
nkl_pair = nmok * nmol
chunks_half = (max(
1, numpy.minimum(int(ioblk_size // (nao_pair * f8_size)), nmoj)),
max(
1,
numpy.minimum(int(ioblk_size // (nij_pair * f8_size)),
nmol)))
'''
ideally, the final transformed eris should have a chunk of nmoj x nmol to
optimize read operations. However, I'm chunking the row size so that the
write operations during the transform can be done as fast as possible.
'''
chunks_full = (numpy.minimum(int(ioblk_size // (nkl_pair * f8_size)),
nmoj), nmol)
if isinstance(erifile, str):
if h5py.is_hdf5(erifile):
feri = h5py.File(erifile)
if dataname in feri:
del (feri[dataname])
else:
feri = h5py.File(erifile, 'w', libver='latest')
else:
assert (isinstance(erifile, h5py.Group))
feri = erifile
h5d_eri = feri.create_dataset(dataname, (nij_pair, nkl_pair),
'f8',
chunks=chunks_full)
feri_swap = lib.H5TmpFile(libver='latest')
half_eri = feri_swap.create_dataset(dataname, (nij_pair, nao_pair),
'f8',
chunks=chunks_half)
log.debug('Memory information:')
log.debug(' IOBLK_SIZE (MB): {}'.format(ioblk_size))
log.debug(' jxl {}x{}, half eri chunk dim {}x{}'.format(
nmoj, nmol, chunks_half[0], chunks_half[1]))
log.debug(' jxl {}x{}, full eri chunk dim {}x{}'.format(
nmoj, nmol, chunks_full[0], chunks_full[1]))
log.debug(' Final disk eri size (MB): {:.3g}, chunked {:.3g}'.format(
nij_pair * nkl_pair * f8_size,
numpy.prod(chunks_full) * f8_size))
log.debug(
' Half transformed eri size (MB): {:.3g}, chunked {:.3g}'.format(
nij_pair * nao_pair * f8_size,
numpy.prod(chunks_half) * f8_size))
log.debug(' RAM buffer for half transform (MB): {:.3g}'.format(
nij_pair * chunks_half[1] * f8_size * 2))
log.debug(' RAM buffer for full transform (MB): {:.3g}'.format(
f8_size * chunks_full[0] * nkl_pair * 2 +
chunks_half[0] * nao_pair * f8_size * 2))
def save1(piece, buf):
start = piece * chunks_half[1]
stop = (piece + 1) * chunks_half[1]
if stop > nao_pair:
stop = nao_pair
half_eri[:, start:stop] = buf[:, :stop - start]
return
def load2(piece):
start = piece * chunks_half[0]
stop = (piece + 1) * chunks_half[0]
if stop > nij_pair:
stop = nij_pair
if start >= nij_pair:
start = stop - 1
return half_eri[start:stop, :]
def prefetch2(piece):
start = piece * chunks_half[0]
stop = (piece + 1) * chunks_half[0]
if stop > nij_pair:
stop = nij_pair
if start >= nij_pair:
start = stop - 1
buf_prefetch[:stop - start, :] = half_eri[start:stop, :]
return
def save2(piece, buf):
start = piece * chunks_full[0]
stop = (piece + 1) * chunks_full[0]
if stop > nij_pair:
stop = nij_pair
h5d_eri[start:stop, :] = buf[:stop - start, :]
return
# transform \mu\nu -> ij
cput0 = time.clock(), time.time()
Cimu = mo_coeffs[0].conj().transpose()
buf_write = numpy.empty((nij_pair, chunks_half[1]))
buf_out = numpy.empty_like(buf_write)
wpiece = 0
with lib.call_in_background(save1) as async_write:
for lo in range(nao_pair):
if lo % chunks_half[1] == 0 and lo > 0:
#save1(wpiece,buf_write)
buf_out, buf_write = buf_write, buf_out
async_write(wpiece, buf_out)
wpiece += 1
buf = lib.unpack_row(eri, lo)
uv = lib.unpack_tril(buf)
uv = Cimu.dot(uv).dot(mo_coeffs[1])
if ij_red:
ij = numpy.ravel(uv) # grabs by row
else:
ij = lib.pack_tril(uv)
buf_write[:, lo % chunks_half[1]] = ij
# final write operation & cleanup
save1(wpiece, buf_write)
log.timer('(uv|lo) -> (ij|lo)', *cput0)
uv = None
ij = None
buf = None
# transform \lambda\sigma -> kl
cput1 = time.clock(), time.time()
Cklam = mo_coeffs[2].conj().transpose()
buf_write = numpy.empty((chunks_full[0], nkl_pair))
buf_out = numpy.empty_like(buf_write)
buf_read = numpy.empty((chunks_half[0], nao_pair))
buf_prefetch = numpy.empty_like(buf_read)
rpiece = 0
wpiece = 0
with lib.call_in_background(save2, prefetch2) as (async_write, prefetch):
buf_read = load2(rpiece)
prefetch(rpiece + 1)
for ij in range(nij_pair):
if ij % chunks_full[0] == 0 and ij > 0:
#save2(wpiece,buf_write)
buf_out, buf_write = buf_write, buf_out
async_write(wpiece, buf_out)
wpiece += 1
if ij % chunks_half[0] == 0 and ij > 0:
#buf_read = load2(rpiece)
buf_read, buf_prefetch = buf_prefetch, buf_read
rpiece += 1
prefetch(rpiece + 1)
lo = lib.unpack_tril(buf_read[ij % chunks_half[0], :])
lo = Cklam.dot(lo).dot(mo_coeffs[3])
if kl_red:
kl = numpy.ravel(lo)
else:
kl = lib.pack_tril(lo)
buf_write[ij % chunks_full[0], :] = kl
save2(wpiece, buf_write)
log.timer('(ij|lo) -> (ij|kl)', *cput1)
if isinstance(erifile, str):
feri.close()
return erifile
def general(eri, mo_coeffs, erifile, dataname='eri_mo',
ioblk_size=IOBLK_SIZE, compact=True, verbose=logger.NOTE):
'''For the given four sets of orbitals, transfer arbitrary spherical AO
integrals to MO integrals on disk.
Args:
eri : 8-fold reduced eri vector
mo_coeffs : 4-item list of ndarray
Four sets of orbital coefficients, corresponding to the four
indices of (ij|kl)
erifile : str or h5py File or h5py Group object
To store the transformed integrals, in HDF5 format.
Kwargs
dataname : str
The dataset name in the erifile (ref the hierarchy of HDF5 format
http://www.hdfgroup.org/HDF5/doc1.6/UG/09_Groups.html). By assigning
different dataname, the existed integral file can be reused. If
the erifile contains the dataname, the new integrals data will
overwrite the old one.
ioblk_size : float or int
The block size for IO, large block size may **not** improve performance
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
Pseudocode / algorithm:
u = mu
v = nu
l = lambda
o = sigma
Assume eri's are 8-fold reduced.
nij/nkl_pair = npair or i*j/k*l if only transforming a subset
First half transform:
Initialize half_eri of size (nij_pair,npair)
For lo = 1 -> npair
Unpack row lo
Unpack row lo to matrix E_{uv}^{lo}
Transform C_ui^+*E*C_nj -> E_{ij}^{lo}
Ravel or pack E_{ij}^{lo}
Save E_{ij}^{lo} -> half_eri[:,lo]
Second half transform:
Initialize h5d_eri of size (nij_pair,nkl_pair)
For ij = 1 -> nij_pair
Load and unpack half_eri[ij,:] -> E_{lo}^{ij}
Transform C_{lk}E_{lo}^{ij}C_{ol} -> E_{kl}^{ij}
Repack E_{kl}^{ij}
Save E_{kl}^{ij} -> h5d_eri[ij,:]
Each matrix is indexed by the composite index ij x kl, where ij/kl is
either npair or ixj/kxl, if only a subset of MOs are being transformed.
Since entire rows or columns need to be read in, the arrays are chunked
such that IOBLK_SIZE = row/col x chunking col/row. For example, for the
first half transform, we would save in nij_pair x IOBLK_SIZE/nij_pair,
then load in IOBLK_SIZE/nkl_pair x npair for the second half transform.
------ kl ----->
|jxl
|
ij
|
|
v
As a first guess, the chunking size is jxl. If the super-rows/cols are
larger than IOBLK_SIZE, then the chunk rectangle jxl is trimmed
accordingly. The pathological limiting case is where the dimensions
nao_pair, nij_pair, or nkl_pair are so large that the arrays are
chunked 1x1, in which case IOBLK_SIZE needs to be increased.
'''
log = logger.new_logger(None, verbose)
log.info('******** ao2mo disk, custom eri ********')
eri_ao = numpy.asarray(eri, order='C')
nao, nmoi = mo_coeffs[0].shape
nmoj = mo_coeffs[1].shape[1]
nao_pair = nao*(nao+1)//2
ijmosym, nij_pair, moij, ijshape = _conc_mos(mo_coeffs[0], mo_coeffs[1], compact)
klmosym, nkl_pair, mokl, klshape = _conc_mos(mo_coeffs[2], mo_coeffs[3], compact)
ijshape = (ijshape[0], ijshape[1]-ijshape[0],
ijshape[2], ijshape[3]-ijshape[2])
dtype = numpy.result_type(eri, *mo_coeffs)
typesize = dtype.itemsize/1e6 # in MB
if nij_pair == 0:
return numpy.empty((nij_pair,nkl_pair))
ij_red = ijmosym == 's1'
kl_red = klmosym == 's1'
if isinstance(erifile, str):
if h5py.is_hdf5(erifile):
feri = h5py.File(erifile, 'a')
if dataname in feri:
del(feri[dataname])
else:
feri = h5py.File(erifile,'w',libver='latest')
else:
assert(isinstance(erifile, h5py.Group))
feri = erifile
h5d_eri = feri.create_dataset(dataname,(nij_pair,nkl_pair), dtype.char)
feri_swap = lib.H5TmpFile(libver='latest')
chunk_size = min(nao_pair, max(4, int(ioblk_size*1e6/8/nao_pair)))
log.debug('Memory information:')
log.debug(' IOBLK_SIZE (MB): {} chunk_size: {}'
.format(ioblk_size, chunk_size))
log.debug(' Final disk eri size (MB): {:.3g}'
.format(nij_pair*nkl_pair*typesize))
log.debug(' Half transformed eri size (MB): {:.3g}'
.format(nij_pair*nao_pair*typesize))
log.debug(' RAM buffer (MB): {:.3g}'
.format(nij_pair*IOBLK_SIZE*typesize*2))
if eri_ao.size == nao_pair**2: # 4-fold symmetry
# half_e1 first transforms the indices which are contiguous in memory
# transpose the 4-fold integrals to make ij the contiguous indices
eri_ao = lib.transpose(eri_ao)
ftrans = _ao2mo.libao2mo.AO2MOtranse1_incore_s4
elif eri_ao.size == nao_pair*(nao_pair+1)//2:
ftrans = _ao2mo.libao2mo.AO2MOtranse1_incore_s8
else:
raise NotImplementedError
if ijmosym == 's2':
fmmm = _ao2mo.libao2mo.AO2MOmmm_nr_s2_s2
elif nmoi <= nmoj:
fmmm = _ao2mo.libao2mo.AO2MOmmm_nr_s2_iltj
else:
fmmm = _ao2mo.libao2mo.AO2MOmmm_nr_s2_igtj
fdrv = getattr(_ao2mo.libao2mo, 'AO2MOnr_e1incore_drv')
def save(piece, buf):
feri_swap[str(piece)] = buf.T
# transform \mu\nu -> ij
cput0 = time.clock(), time.time()
with lib.call_in_background(save) as async_write:
for istep, (p0, p1) in enumerate(lib.prange(0, nao_pair, chunk_size)):
if dtype == numpy.double:
buf = numpy.empty((p1-p0, nij_pair))
fdrv(ftrans, fmmm,
buf.ctypes.data_as(ctypes.c_void_p),
eri_ao.ctypes.data_as(ctypes.c_void_p),
moij.ctypes.data_as(ctypes.c_void_p),
ctypes.c_int(p0), ctypes.c_int(p1-p0),
ctypes.c_int(nao),
ctypes.c_int(ijshape[0]), ctypes.c_int(ijshape[1]),
ctypes.c_int(ijshape[2]), ctypes.c_int(ijshape[3]))
else: # complex
tmp = numpy.empty((p1-p0, nao_pair))
if eri_ao.size == nao_pair**2: # 4-fold symmetry
tmp = eri_ao[p0:p1]
else: # 8-fold symmetry
for i in range(p0, p1):
tmp[i-p0] = lib.unpack_row(eri_ao, i)
tmp = lib.unpack_tril(tmp, filltriu=lib.SYMMETRIC)
buf = lib.einsum('xpq,pi,qj->xij', tmp, mo_coeffs[0].conj(), mo_coeffs[1])
if ij_red:
buf = buf.reshape(p1-p0,-1) # grabs by row
else:
buf = lib.pack_tril(buf)
async_write(istep, buf)
log.timer('(uv|lo) -> (ij|lo)', *cput0)
# transform \lambda\sigma -> kl
cput1 = time.clock(), time.time()
Cklam = mo_coeffs[2].conj()
buf_read = numpy.empty((chunk_size,nao_pair), dtype=dtype)
buf_prefetch = numpy.empty_like(buf_read)
def load(start, stop, buf):
if start < stop:
_load_from_h5g(feri_swap, start, stop, buf)
def save(start, stop, buf):
if start < stop:
h5d_eri[start:stop] = buf[:stop-start]
with lib.call_in_background(save,load) as (async_write, prefetch):
for p0, p1 in lib.prange(0, nij_pair, chunk_size):
if p0 == 0:
load(p0, p1, buf_prefetch)
buf_read, buf_prefetch = buf_prefetch, buf_read
prefetch(p1, min(p1+chunk_size, nij_pair), buf_prefetch)
lo = lib.unpack_tril(buf_read[:p1-p0], filltriu=lib.SYMMETRIC)
lo = lib.einsum('xpq,pi,qj->xij', lo, Cklam, mo_coeffs[3])
if kl_red:
kl = lo.reshape(p1-p0,-1)
else:
kl = lib.pack_tril(lo)
async_write(p0, p1, kl)
log.timer('(ij|lo) -> (ij|kl)', *cput1)
if isinstance(erifile, str):
feri.close()
return erifile