import h5py
import numpy
try:
import mpi4py
mpi4py.rc.recv_mprobe = False
from mpi4py import MPI
mpi_sum = MPI.SUM
except ImportError:
mpi_sum = None
import sys
from pauxy.estimators.utils import H5EstimatorHelper
from pauxy.estimators.greens_function import gab
from pauxy.estimators.mixed import local_energy
from pauxy.estimators.ekt import ekt_1p_fock_opt, ekt_1h_fock_opt
from pauxy.propagation.generic import back_propagate_generic
from pauxy.propagation.planewave import back_propagate_planewave
import pauxy.propagation.hubbard
[docs]class BackPropagation(object):
"""Class for computing back propagated estimates.
Parameters
----------
bp : dict
Input options for BP estimates.
root : bool
True if on root/master processor.
h5f : :class:`h5py.File`
Output file object.
qmc : :class:`pauxy.state.QMCOpts` object.
Container for qmc input options.
system : system object
System object.
trial : :class:`pauxy.trial_wavefunction.X' object
Trial wavefunction class.
dtype : complex or float
Output type.
BT2 : :class:`numpy.ndarray`
One-body propagator for back propagation.
Attributes
----------
nmax : int
Max number of measurements.
header : int
Output header.
rdm : bool
True if output BP RDM to file.
nreg : int
Number of regular estimates (exluding iteration).
G : :class:`numpy.ndarray`
One-particle RDM.
estimates : :class:`numpy.ndarray`
Store for mixed estimates per processor.
global_estimates : :class:`numpy.ndarray`
Store for mixed estimates accross all processors.
output : :class:`pauxy.estimators.H5EstimatorHelper`
Class for outputting data to HDF5 group.
rdm_output : :class:`pauxy.estimators.H5EstimatorHelper`
Class for outputting rdm data to HDF5 group.
"""
def __init__(self, bp, root, filename, qmc, system, trial, dtype, BT2):
self.tau_bp = bp.get('tau_bp', 0)
self.nmax = int(self.tau_bp/qmc.dt)
self.header = ['E', 'E1b', 'E2b']
self.calc_one_rdm = bp.get('one_rdm', True)
self.calc_two_rdm = bp.get('two_rdm', None)
self.init_walker = bp.get('init_walker', False)
self.nsplit = bp.get('nsplit', 1)
self.splits = numpy.array([(i+1)*(self.nmax//self.nsplit)
for i in range(self.nsplit)])
self.nreg = len(self.header)
self.accumulated = False
self.eval_energy = bp.get('evaluate_energy', False)
self.eval_ekt = bp.get('evaluate_ekt', False)
self.G = numpy.zeros(trial.G.shape, dtype=trial.G.dtype)
self.nstblz = qmc.nstblz
self.BT2 = BT2
self.restore_weights = bp.get('restore_weights', None)
if root:
print("# restore_weights = {}".format(self.restore_weights))
self.dt = qmc.dt
dms_size = self.G.size
# Abuse of language for the moment. Only accumulates S(k) for UEG.
# TODO: Add functionality to accumulate 2RDM?
if self.calc_two_rdm is not None:
if self.calc_two_rdm == "structure_factor":
two_rdm_shape = (2,2,len(system.qvecs),)
else:
two_rdm_shape = (system.nbasis, system.nbasis, system.nbasis, system.nbasis)
self.two_rdm = numpy.zeros(two_rdm_shape,
dtype=numpy.complex128)
dms_size += self.two_rdm.size
else:
self.two_rdm = None
if self.eval_ekt:
self.ekt_fock_1h = numpy.zeros_like(self.G[0])
self.ekt_fock_1p = numpy.zeros_like(self.G[0])
dms_size += self.ekt_fock_1h.size
dms_size += self.ekt_fock_1p.size
self.estimates = numpy.zeros(self.nreg+1+dms_size, dtype=dtype)
self.global_estimates = numpy.zeros(self.nreg+1+dms_size,
dtype=dtype)
self.key = {
'ETotal': "BP estimate for total energy.",
'E1B': "BP estimate for one-body energy.",
'E2B': "BP estimate for two-body energy."
}
if root:
self.setup_output(filename)
if trial.type == 'GHF':
self.update = self.update_ghf
self.back_propagate = pauxy.propagation.hubbard.back_propagate_ghf
else:
self.update = self.update_uhf
if system.name == "Generic":
self.back_propagate = back_propagate_generic
elif system.name == "UEG":
self.back_propagate = back_propagate_planewave
else:
self.back_propagate = pauxy.propagation.hubbard.back_propagate
[docs] def update_uhf(self, system, qmc, trial, psi, step, free_projection=False):
"""Calculate back-propagated estimates for RHF/UHF walkers.
Parameters
----------
system : system object in general.
Container for model input options.
qmc : :class:`pauxy.state.QMCOpts` object.
Container for qmc input options.
trial : :class:`pauxy.trial_wavefunction.X' object
Trial wavefunction class.
psi : :class:`pauxy.walkers.Walkers` object
CPMC wavefunction.
step : int
Current simulation step
free_projection : bool
True if doing free projection.
"""
buff_ix = psi.walkers[0].field_configs.step
if buff_ix not in self.splits:
return
nup = system.nup
for i, wnm in enumerate(psi.walkers):
if self.init_walker:
phi_bp = trial.init.copy()
else:
phi_bp = trial.psi.copy()
# TODO: Fix for ITCF.
self.back_propagate(phi_bp, wnm.field_configs, system,
self.nstblz, self.BT2, self.dt)
self.G[0] = gab(phi_bp[:,:nup], wnm.phi_old[:,:nup]).T
self.G[1] = gab(phi_bp[:,nup:], wnm.phi_old[:,nup:]).T
if self.eval_energy:
eloc = local_energy(system, self.G, opt=False,
two_rdm=self.two_rdm)
energies = numpy.array(list(eloc))
else:
energies = numpy.zeros(3)
if self.calc_two_rdm is not None and self.calc_two_rdm is not "structure_factor":
# <p^+ q^+ s r> = G(p, r, q, s) also spin-summed
self.two_rdm = numpy.einsum("pr,qs->prqs",self.G[0], self.G[0], optimize=True)\
- numpy.einsum("ps,qr->prqs",self.G[0], self.G[0], optimize=True)
self.two_rdm += numpy.einsum("pr,qs->prqs",self.G[1], self.G[1], optimize=True)\
- numpy.einsum("ps,qr->prqs",self.G[1], self.G[1], optimize=True)
self.two_rdm += numpy.einsum("pr,qs->prqs",self.G[0], self.G[1], optimize=True)\
+ numpy.einsum("pr,qs->prqs",self.G[1], self.G[0], optimize=True)
if self.eval_ekt:
if (system.name == 'UEG'):
# there needs to be a factor of 2.0 here to account for the convention of cholesky vectors in the system class
chol_vecs = 2.0 * system.chol_vecs.toarray().T.reshape((system.nchol, system.nbasis, system.nbasis))
self.ekt_fock_1p = ekt_1p_fock_opt(system.H1[0],chol_vecs, self.G[0], self.G[1])
self.ekt_fock_1h = ekt_1h_fock_opt(system.H1[0],chol_vecs, self.G[0], self.G[1])
else:
self.ekt_fock_1p = ekt_1p_fock_opt(system.H1[0],system.chol_vecs, self.G[0], self.G[1])
self.ekt_fock_1h = ekt_1h_fock_opt(system.H1[0],system.chol_vecs, self.G[0], self.G[1])
if self.restore_weights is not None:
cosine_fac, ph_fac = wnm.field_configs.get_wfac()
if self.restore_weights == "full":
# BP-Pres
wfac = ph_fac / cosine_fac
else:
# BP-PRes (partial)
wfac = ph_fac
weight = wnm.weight * wfac
else:
# BP-PhL
weight = wnm.weight
self.estimates[:self.nreg] += weight*energies
self.estimates[self.nreg] += weight
start = self.nreg + 1
end = start + self.G.size
self.estimates[start:end] += weight*self.G.flatten()
if self.calc_two_rdm is not None:
start = end
end = end + self.two_rdm.size
self.estimates[start:end] += weight*self.two_rdm.flatten()
if self.eval_ekt:
start = end
end = end + self.ekt_fock_1p.size
self.estimates[start:end] += weight*self.ekt_fock_1p.flatten()
start = end
end = end + self.ekt_fock_1h.size
self.estimates[start:end] += weight*self.ekt_fock_1h.flatten()
if buff_ix == self.splits[-1]:
wnm.field_configs.reset()
if buff_ix == self.splits[-1]:
psi.copy_historic_wfn()
self.accumulated = True
self.buff_ix = buff_ix
[docs] def update_ghf(self, system, qmc, trial, psi, step, free_projection=False):
"""Calculate back-propagated estimates for GHF walkers.
Parameters
----------
system : system object in general.
Container for model input options.
qmc : :class:`pauxy.state.QMCOpts` object.
Container for qmc input options.
trial : :class:`pauxy.trial_wavefunction.X' object
Trial wavefunction class.
psi : :class:`pauxy.walkers.Walkers` object
CPMC wavefunction.
step : int
Current simulation step
free_projection : bool
True if doing free projection.
"""
if step % self.nmax != 0:
return
print(" ***** Back Propagation with GHF is broken.")
sys.exit()
psi_bp = self.back_propagate(system, psi.walkers, trial,
self.nstblz, self.BT2,
self.dt)
denominator = sum(wnm.weight for wnm in psi.walkers)
nup = system.nup
for i, (wnm, wb) in enumerate(zip(psi.walkers, psi_bp)):
construct_multi_ghf_gab(wb.phi, wnm.phi_old, wb.weights, wb.Gi, wb.ots)
# note that we are abusing the weights variable from the multighf
# walker to store the reorthogonalisation factors.
weights = wb.weights * trial.coeffs * wb.ots
denom = sum(weights)
energies = numpy.array(list(local_energy_ghf(system, wb.Gi, weights, denom)))
self.G = numpy.einsum('i,ijk->jk', weights, wb.Gi) / denom
self.estimates[1:]= (
self.estimates[1:] + wnm.weight*numpy.append(energies,self.G.flatten())
)
self.estimates[0] += denominator
psi.copy_historic_wfn()
psi.copy_bp_wfn(psi_bp)
[docs] def print_step(self, comm, nprocs, step, nsteps=1, free_projection=False):
"""Print back-propagated estimates to file.
Parameters
----------
comm :
MPI communicator.
nprocs : int
Number of processors.
step : int
Current iteration number.
nmeasure : int
Number of steps between measurements.
"""
if not self.accumulated:
return
comm.Reduce(self.estimates, self.global_estimates, op=mpi_sum)
if comm.rank == 0:
weight = self.global_estimates[self.nreg]
self.output.push(numpy.array([weight]),
'denominator_{:d}'.format(self.buff_ix))
if self.eval_energy:
if free_projection:
self.output.push(self.global_estimates[:self.nreg],
'energies_{:d}'.format(self.buff_ix))
else:
self.output.push(self.global_estimates[:self.nreg]/weight,
'energies_{:d}'.format(self.buff_ix))
if self.calc_one_rdm:
start = self.nreg + 1
end = self.nreg + 1 + self.G.size
rdm = self.global_estimates[start:end].reshape(self.G.shape)
self.output.push(rdm, 'one_rdm_{:d}'.format(self.buff_ix))
if self.calc_two_rdm:
start = self.nreg + 1 + self.G.size
end = start + self.two_rdm.size
rdm = self.global_estimates[start:end].reshape(self.two_rdm.shape)
self.output.push(rdm, 'two_rdm_{:d}'.format(self.buff_ix))
if self.eval_ekt:
start = self.nreg + 1
if self.calc_one_rdm:
start += self.G.size
if self.calc_two_rdm:
start += self.two_rdm.size
end = start + self.ekt_fock_1p.size
fock = self.global_estimates[start:end].reshape(self.ekt_fock_1p.shape)
self.output.push(fock, 'fock_1p_{:d}'.format(self.buff_ix))
start = end
end = end + self.ekt_fock_1h.size
fock = self.global_estimates[start:end].reshape(self.ekt_fock_1h.shape)
self.output.push(fock, 'fock_1h_{:d}'.format(self.buff_ix))
if self.buff_ix == self.splits[-1]:
self.output.increment()
self.accumulated = False
self.zero()
[docs] def zero(self):
"""Zero (in the appropriate sense) various estimator arrays."""
self.estimates[:] = 0
self.global_estimates[:] = 0
[docs] def setup_output(self, filename):
est_name = 'back_propagated'
if self.eval_energy:
with h5py.File(filename, 'a') as fh5:
fh5[est_name+'/headers'] = numpy.array(self.header).astype('S')
self.output = H5EstimatorHelper(filename, est_name)