import ast
import h5py
import numpy
import sys
import scipy.linalg
import time
from scipy.sparse import csr_matrix
from pauxy.utils.linalg import modified_cholesky
from pauxy.utils.io import (
from_qmcpack_sparse,
from_qmcpack_dense,
write_qmcpack_sparse,
write_qmcpack_dense,
)
from pauxy.estimators.generic import (
local_energy_generic, core_contribution,
local_energy_generic_cholesky, core_contribution_cholesky
)
[docs]class Generic(object):
"""Generic system defined by ab-initio Hamiltonian.
Can be created by either passing the one and two electron integrals directly
or initialised from integrals stored in QMCPACK hdf5 format. If initialising
from file the `inputs' optional dictionary should be populated.
Parameters
----------
nelec : tuple
Number of alpha and beta electrons.
h1e : :class:`numpy.ndarray'
One-electron integrals. Optional. Default: None.
chol : :class:`numpy.ndarray'
Factorized 2-electron integrals of shape (naux, nbasis, nbasis).
Optional. Default: None.
ecore : float
Core energy.
inputs : dict
Input options defined below.
nup : int
Number of up electrons.
ndown : int
Number of down electrons.
integrals : string
Path to file containing one- and two-electron integrals in QMCPACK
format.
threshold : float
Cutoff for cholesky decomposition or minimum eigenvalue.
verbose : bool
Print extra information.
Attributes
----------
H1 : :class:`numpy.ndarray`
One-body part of the Hamiltonian. Spin-dependent by default.
ecore : float
Core contribution to the total energy.
h1e_mod : :class:`numpy.ndarray`
Modified one-body Hamiltonian.
chol_vecs : :class:`numpy.ndarray`
Cholesky vectors.
nchol : int
Number of cholesky vectors.
nfields : int
Number of auxiliary fields required.
sparse_cutoff : float
Screen out integrals below this threshold. Optional. Default 0.
cplx_chol : bool
Force setting of interpretation of cholesky decomposition. Optional.
Default False, i.e. real/complex factorization determined from cholesky
integrals.
"""
def __init__(self, nelec=None, h1e=None, chol=None, ecore=None, inputs={}, verbose=False):
if verbose:
print("# Parsing input options.")
self.name = "Generic"
self.atom = inputs.get('atom', None)
self.verbose = verbose
if nelec is None:
self.nup = inputs['nup']
self.ndown = inputs['ndown']
self.nelec = (self.nup, self.ndown)
else:
self.nup, self.ndown = nelec
self.nelec = nelec
self.ne = self.nup + self.ndown
self.integral_file = inputs.get('integrals')
self.cutoff = inputs.get('sparse_cutoff', None)
self.sparse = inputs.get('sparse', False)
self.half_rotated_integrals = inputs.get('integral_tensor', False)
self._opt = self.sparse
self.cplx_chol = inputs.get('complex_cholesky', False)
self.mu = inputs.get('mu', None)
if verbose:
print("# Reading integrals from %s." % self.integral_file)
if chol is not None:
self.chol_vecs = chol
self.ecore = ecore
if numpy.max(numpy.abs(self.chol_vecs.imag)) > 1e-6:
self.cplx_chol = True
if verbose:
print("# Found complex integrals.")
print("# Using Hermitian Cholesky decomposition.")
else:
if verbose:
print("# Using real symmetric Cholesky decomposition.")
self.cplx_chol = False
else:
start = time.time()
h1e, self.chol_vecs, self.ecore = self.read_integrals()
if numpy.max(numpy.abs(self.chol_vecs.imag)) > 1e-6:
self.cplx_chol = True
if verbose:
print("# Found complex integrals.")
print("# Using Hermitian Cholesky decomposition.")
else:
if verbose:
print("# Using real symmetric Cholesky decomposition.")
self.cplx_chol = False
if verbose:
print("# Time to read integrals: {:.6f} "
"s".format(time.time()-start))
self.H1 = numpy.array([h1e,h1e])
self.nbasis = h1e.shape[0]
self._alt_convention = False
mem = self.chol_vecs.nbytes / (1024.0**3)
if verbose:
print("# Number of orbitals: %d"%self.nbasis)
print("# Number of electrons: (%d, %d)"%(self.nup, self.ndown))
print("# Approximate memory required by Cholesky vectors %f GB"%mem)
self.nchol = self.chol_vecs.shape[0]
start = time.time()
self.construct_h1e_mod()
if verbose:
print("# Time to construct H1e_mod: "
"{:.6f} s".format(time.time()-start))
self.ktwist = numpy.array([None])
# For consistency
self.vol = 1.0
start = time.time()
if self.cplx_chol:
self.nfields = 2 * self.nchol
self.hs_pot = numpy.zeros(shape=(self.nfields,self.nbasis,self.nbasis),
dtype=numpy.complex128)
for (n,cn) in enumerate(self.chol_vecs):
vplus = 0.5*(cn+cn.conj().T)
vminus = 0.5j*(cn-cn.conj().T)
self.hs_pot[n] = vplus
self.hs_pot[self.nchol+n] = vminus
else:
self.hs_pot = self.chol_vecs
self.nfields = self.nchol
if verbose:
print("# Number of Cholesky vectors: %d"%(self.nchol))
print("# Number of fields: %d"%(self.nfields))
print("# Time to construct Hubbard--Stratonovich potentials: "
"%f s"%(time.time()-start))
if self.sparse:
if verbose:
print("# Using sparse linear algebra.")
if self.cutoff is not None:
self.hs_pot[numpy.abs(self.hs_pot) < self.cutoff] = 0
tmp = numpy.transpose(self.hs_pot, axes=(1,2,0))
tmp = tmp.reshape(self.nbasis*self.nbasis, self.nfields)
self.hs_pot = csr_matrix(tmp)
else:
self.hs_pot = numpy.transpose(self.hs_pot, axes=(1,2,0))
self.hs_pot = self.hs_pot.reshape(self.nbasis*self.nbasis, self.nfields)
write_ints = inputs.get('write_integrals', None)
if write_ints is not None:
self.write_integrals()
if verbose:
print("# Finished setting up Generic system object.")
[docs] def read_integrals(self):
try:
(h1e, schol_vecs, ecore, nbasis, nup, ndown) = (
from_qmcpack_sparse(self.integral_file)
)
chol_vecs = schol_vecs.toarray().T.reshape((-1,nbasis,nbasis))
except KeyError:
(h1e, chol_vecs, ecore, nbasis, nup, ndown) = (
from_qmcpack_dense(self.integral_file)
)
chol_vecs = chol_vecs.T.reshape((-1,nbasis,nbasis))
except OSError:
print("# Unknown Hamiltonian file {}.".format(self.integral_file))
except:
print("# Unknown Hamiltonian file format.")
if ((nup != self.nup) or ndown != self.ndown):
print("# Warning: Number of electrons differs from integral file.")
print("# file: %d %d vs. input: %d %d"%(nup, ndown, self.nup, self.ndown))
return h1e, chol_vecs, ecore
[docs] def construct_h1e_mod(self):
# Subtract one-body bit following reordering of 2-body operators.
# Eqn (17) of [Motta17]_
v0 = 0.5 * numpy.einsum('lik,ljk->ij', self.chol_vecs,
self.chol_vecs.conj(), optimize='optimal')
self.h1e_mod = numpy.array([self.H1[0]-v0, self.H1[1]-v0])
[docs] def construct_integral_tensors_real(self, trial):
# Half rotated cholesky vectors (by trial wavefunction).
# Assuming nup = ndown here
M = self.nbasis
na = self.nup
nb = self.ndown
if self.verbose:
print("# Constructing half rotated Cholesky vectors.")
# rup = numpy.zeros(shape=(self.nchol, na, M),
# dtype=numpy.complex128)
# rdn = numpy.zeros(shape=(self.nchol, nb, M),
# dtype=numpy.complex128)
if self.sparse:
self.hs_pot = self.hs_pot.toarray().reshape(M,M,self.nfields)
else:
self.hs_pot = self.hs_pot.reshape(M,M,self.nfields)
start = time.time()
# rrup = numpy.einsum('ia,ikn->akn',
# trial.psi[:,:na].conj(),
# self.hs_pot,
# optimize='greedy')
# rdn = numpy.einsum('ia,ikn->akn',
# trial.psi[:,na:].conj(),
# self.hs_pot,
# optimize='greedy')
rup = numpy.tensordot(trial.psi[:,:na].conj(),
self.hs_pot,
axes=((0),(0)))
rdn = numpy.tensordot(trial.psi[:,na:].conj(),
self.hs_pot,
axes=((0),(0)))
trot = time.time() - start
# This is much faster than einsum.
# for l in range(self.nchol):
# rup[l] = numpy.dot(trial.psi[:,:na].conj().T, self.chol_vecs[l])
# rdn[l] = numpy.dot(trial.psi[:,na:].conj().T, self.chol_vecs[l])
if self.half_rotated_integrals:
start = time.time()
if self.verbose:
print("# Constructing half rotated V_{(ab)(kl)}.")
vakbl_a = (numpy.einsum('akn,bln->akbl', rup, rup, optimize='greedy') -
numpy.einsum('bkn,aln->akbl', rup, rup, optimize='greedy'))
vakbl_b = (numpy.einsum('akn,bln->akbl', rdn, rdn, optimize='greedy') -
numpy.einsum('bkn,aln->akbl', rdn, rdn, optimize='greedy'))
tvakbl = time.time() - start
if self.cutoff is not None:
rup[numpy.abs(rup) < self.cutoff] = 0.0
rdn[numpy.abs(rdn) < self.cutoff] = 0.0
if self.half_rotated_integrals:
vakbl_a[numpy.abs(vakbl_a) < self.cutoff] = 0.0
vakbl_b[numpy.abs(vakbl_b) < self.cutoff] = 0.0
if self.half_rotated_integrals:
self.vakbl = [csr_matrix(vakbl_a.reshape((M*na, M*na))),
csr_matrix(vakbl_b.reshape((M*nb, M*nb)))]
if self.sparse:
if self.cutoff is not None:
self.hs_pot[numpy.abs(self.hs_pot) < self.cutoff] = 0
self.hs_pot = self.hs_pot.reshape((M*M,-1))
self.hs_pot = csr_matrix(self.hs_pot)
self.rot_hs_pot = [csr_matrix(rup.reshape((M*na, -1))),
csr_matrix(rdn.reshape((M*nb, -1)))]
else:
self.rot_hs_pot = [rup.reshape((M*na, -1)), rdn.reshape((M*nb, -1))]
self.hs_pot = self.hs_pot.reshape((M*M,-1))
self.rchol_vecs = self.rot_hs_pot
if self.verbose:
print("# Time to construct half-rotated Cholesky: %f s"%trot)
if self.sparse:
nnz = self.rchol_vecs[0].nnz
print("# Number of non-zero elements in rotated cholesky: %d"%nnz)
nelem = self.rchol_vecs[0].shape[0] * self.rchol_vecs[0].shape[1]
print("# Sparsity: %f"%(1-float(nnz)/nelem))
mem = (2*nnz*16/(1024.0**3))
else:
mem = self.rchol_vecs[0].nbytes + self.rchol_vecs[1].nbytes
mem /= 1024.0**3
print("# Approximate memory required for half-rotated Cholesky: "
"{:.6f} GB".format(mem))
if self.half_rotated_integrals:
print("# Time to construct V_{(ak)(bl)}: %f"%tvakbl)
nnz = self.vakbl[0].nnz
print("# Number of non-zero elements in V_{(ak)(bl)}: %d"%nnz)
mem = (2*nnz*16/(1024.0**3))
print("# Approximate memory used %f GB"%mem)
nelem = self.vakbl[0].shape[0] * self.vakbl[0].shape[1]
print("# Sparsity: %f"%(1-float(nnz)/nelem))
[docs] def construct_integral_tensors_cplx(self, trial):
# Half rotated cholesky vectors (by trial wavefunction).
# Assuming nup = ndown here
M = self.nbasis
na = self.nup
nb = self.ndown
if self.verbose:
print("# Constructing complex half rotated HS Potentials.")
rup = numpy.zeros(shape=(self.nfields, na, M),
dtype=numpy.complex128)
rdn = numpy.zeros(shape=(self.nfields, nb, M),
dtype=numpy.complex128)
# rup = numpy.einsum('ia,lik->lak',
# trial.psi[:,:na].conj(),
# self.hs_pot)
# rdn = numpy.einsum('ia,lik->lak',
# trial.psi[:,na:].conj(),
# self.hs_pot)
# This is much faster than einsum.
start = time.time()
if self.sparse:
self.hs_pot = self.hs_pot.toarray().reshape(M,M,self.nfields)
self.hs_pot = self.hs_pot.transpose(2,0,1)
for (n,cn) in enumerate(self.hs_pot):
rup[n] = numpy.dot(trial.psi[:,:na].conj().T, self.hs_pot[n])
rdn[n] = numpy.dot(trial.psi[:,na:].conj().T, self.hs_pot[n])
self.rot_hs_pot = [csr_matrix(rup.reshape((-1,M*na)).T),
csr_matrix(rdn.reshape((-1,M*nb)).T)]
if self.verbose:
print("# Time to construct half-rotated HS potentials: "
"%f s"%(time.time()-start))
nnz = self.rot_hs_pot[0].nnz
print("# Number of non-zero elements in rotated potentials: %d"%nnz)
nelem = self.rot_hs_pot[0].shape[0] * self.rot_hs_pot[0].shape[1]
print("# Sparsity: %f"%(1-float(nnz)/nelem))
mem = (2*nnz*16/(1024.0**3))
print("# Approximate memory required %f" " GB"%mem)
print("# Constructing half rotated V_{(ab)(kl)}.")
# This is also much faster than einsum.
Qak = numpy.zeros((self.nchol, M*na), dtype=numpy.complex128)
Rbl = numpy.zeros((self.nchol, M*na), dtype=numpy.complex128)
start = time.time()
for (n,cn) in enumerate(self.chol_vecs):
Qak[n] = numpy.dot(trial.psi[:,:na].conj().T, cn).ravel()
Rbl[n] = numpy.dot(trial.psi[:,:na].conj().T, cn.conj()).ravel()
if self.verbose:
print("# Time to construct Qak, Rbl: %f s"%(time.time()-start))
Makbl = numpy.dot(Qak.T,Rbl)
vakbl_a = (
Makbl -
Makbl.reshape(na,M,na,M).transpose((2,1,0,3)).reshape(na*M,na*M)
)
Qak = numpy.zeros((self.nchol, M*nb), dtype=numpy.complex128)
Rbl = numpy.zeros((self.nchol, M*nb), dtype=numpy.complex128)
for (n,cn) in enumerate(self.chol_vecs):
Qak[n] = numpy.dot(trial.psi[:,na:].conj().T, cn).ravel()
Rbl[n] = numpy.dot(trial.psi[:,na:].conj().T, cn.conj()).ravel()
Makbl = numpy.dot(Qak.T,Rbl)
vakbl_b = (
Makbl -
Makbl.reshape(nb,M,nb,M).transpose((2,1,0,3)).reshape(nb*M,nb*M)
)
self.vakbl = [csr_matrix(vakbl_a.reshape((M*na, M*na))),
csr_matrix(vakbl_b.reshape((M*nb, M*nb)))]
tvakbl = time.time() - start
# TODO: Stop converting hs pot to dense
if self.sparse:
if self.cutoff is not None:
self.hs_pot[numpy.abs(self.hs_pot) < self.cutoff] = 0
tmp = numpy.transpose(self.hs_pot, axes=(1,2,0))
tmp = tmp.reshape(self.nbasis*self.nbasis, self.nfields)
self.hs_pot = csr_matrix(tmp)
else:
self.hs_pot = numpy.transpose(self.hs_pot, axes=(1,2,0))
self.hs_pot = self.hs_pot.reshape(self.nbasis*self.nbasis, self.nfields)
if self.verbose:
print("# Time to construct V_{(ak)(bl)}: %f s"%(tvakbl))
nnz = self.vakbl[0].nnz
mem = (2*nnz*16/(1024.0**3))
print("# Number of non-zero elements in V_{(ak)(bl)}: %d"%nnz)
print("# Approximate memory used %f GB"%mem)
nelem = self.vakbl[0].shape[0] * self.vakbl[0].shape[1]
print("# Sparsity: %f"%(1-float(nnz)/nelem))
[docs] def hijkl(self, i, j, k, l):
return numpy.dot(self.chol_vecs[:,i,k], self.chol_vecs[:,j,l])
[docs] def write_integrals(self, filename='hamil.h5'):
if self.sparse:
write_qmcpack_sparse(self.H1[0],
self.chol_vecs.reshape((-1,self.nbasis*self.nbasis)).T.copy(),
self.nelec, self.nbasis,
ecuc=self.ecore, filename=filename)
else:
write_qmcpack_dense(self.H1[0],
self.chol_vecs.reshape((-1,self.nbasis*self.nbasis)).T.copy(),
self.nelec, self.nbasis,
enuc=self.ecore, filename=filename,
real_chol=not self.cplx_chol)