From f7310b86bf5e975515d8dee5e80cc2e058bf6092 Mon Sep 17 00:00:00 2001 From: m-baumgarten Date: Thu, 8 Feb 2024 21:38:28 -0500 Subject: [PATCH 01/61] fixed typo --- ipie/estimators/local_energy_noci.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ipie/estimators/local_energy_noci.py b/ipie/estimators/local_energy_noci.py index 4444ae3a..e1faa799 100644 --- a/ipie/estimators/local_energy_noci.py +++ b/ipie/estimators/local_energy_noci.py @@ -108,7 +108,7 @@ def local_energy_noci( walkers : WalkerBatch Walkers object. trial : trial object - Trial wavefunctioni. + Trial wavefunction. Returns ------- From 4ffbe5680a9e99cc415f795079f8e57d2713a08d Mon Sep 17 00:00:00 2001 From: m-baumgarten Date: Thu, 8 Feb 2024 21:41:33 -0500 Subject: [PATCH 02/61] holstein with toyozawa trial, coherent state trial todo --- ipie/estimators/energy.py | 12 ++ ipie/estimators/energy_holstein.py | 35 +++ ipie/estimators/local_energy_holstein.py | 38 ++++ ipie/hamiltonians/holstein.py | 31 +++ ipie/propagation/holstein.py | 199 ++++++++++++++++++ ipie/propagation/propagator.py | 4 +- .../holstein/coherent_state.py | 68 ++++++ .../holstein/eph_trial_base.py | 30 +++ ipie/trial_wavefunction/holstein/toyozawa.py | 110 ++++++++++ ipie/walkers/eph_walkers.py | 120 +++++++++++ 10 files changed, 646 insertions(+), 1 deletion(-) create mode 100644 ipie/estimators/energy_holstein.py create mode 100644 ipie/estimators/local_energy_holstein.py create mode 100644 ipie/hamiltonians/holstein.py create mode 100644 ipie/propagation/holstein.py create mode 100644 ipie/trial_wavefunction/holstein/coherent_state.py create mode 100644 ipie/trial_wavefunction/holstein/eph_trial_base.py create mode 100644 ipie/trial_wavefunction/holstein/toyozawa.py create mode 100644 ipie/walkers/eph_walkers.py diff --git a/ipie/estimators/energy.py b/ipie/estimators/energy.py index d8a6936b..883661a1 100644 --- a/ipie/estimators/energy.py +++ b/ipie/estimators/energy.py @@ -18,6 +18,7 @@ import plum from ipie.estimators.estimator_base import EstimatorBase +from ipie.estimators.local_energy_holstein import local_energy_holstein from ipie.estimators.local_energy_batch import ( local_energy_batch, local_energy_multi_det_trial_batch, @@ -30,6 +31,7 @@ local_energy_multi_det_trial_wicks_batch_opt_chunked, ) from ipie.hamiltonians.generic import GenericComplexChol, GenericRealChol +from ipie.hamiltonians.holstein import HolsteinModel from ipie.systems.generic import Generic from ipie.trial_wavefunction.noci import NOCI from ipie.trial_wavefunction.particle_hole import ( @@ -39,9 +41,19 @@ ParticleHoleSlow, ) from ipie.trial_wavefunction.single_det import SingleDet +from ipie.trial_wavefunction.holstein.eph_trial_base import EphTrialWavefunctionBase from ipie.utils.backend import arraylib as xp from ipie.walkers.uhf_walkers import UHFWalkers +from ipie.walkers.eph_walkers import EphWalkers +@plum.dispatch +def local_energy( + system: Generic, + hamiltonian: HolsteinModel, + walkers: EphWalkers, + trial: EphTrialWavefunctionBase +): + return local_energy_holstein(system, hamiltonian, walkers, trial) @plum.dispatch def local_energy( diff --git a/ipie/estimators/energy_holstein.py b/ipie/estimators/energy_holstein.py new file mode 100644 index 00000000..b59c09ba --- /dev/null +++ b/ipie/estimators/energy_holstein.py @@ -0,0 +1,35 @@ +from ipie.estimators.estimator_base import EstimatorBase +from ipie.hamiltonians.holstein import HolsteinModel +from ipie.systems.generic import Generic + + +def local_energy( + system: Generic, hamiltonian: HolsteinModel, walkers: EphWalkers, trial: EphTrial +): + return local_energy_holstein(system, hamiltonian, walkers, trial) + + +class EnergyEstimatorHolstein(EstimatorBase): + """""" + def __init__(self, system, hamiltonian, trial): + assert system is not None + assert ham is not None + assert trial is not None + super().__init__() + self.scalar_estimator = True + self._data = { + "ENumer": 0.0j, + "EDenom": 0.0j, + "ETotal": 0.0j + } + self._shape = (len(self.names),) + self._data_index = {k: i for i, k in enumerate(list(self._data.keys()))} + self.print_to_stdout = True + self.ascii_filename = filename + + def compute_estimator(self, system, walkers, hamiltonian, trial, istep)=1): + trial.calc_greens_function(walkers) + energy = local_energy(system, hamiltonian, walkers, trial) + self._data["ENumer"] = xp.sum(walkers.weight * energy[:, 0].real) + self._data["EDenom"] = xp.sum(walkers.weight) + diff --git a/ipie/estimators/local_energy_holstein.py b/ipie/estimators/local_energy_holstein.py new file mode 100644 index 00000000..97e609ab --- /dev/null +++ b/ipie/estimators/local_energy_holstein.py @@ -0,0 +1,38 @@ +import numpy as np + +from ipie.hamiltonians.holstein import HolsteinModel +from ipie.systems.generic import Generic +from ipie.trial_wavefunction.holstein.eph_trial_base import EphTrialWavefunctionBase +from ipie.utils.backend import arraylib as xp +from ipie.walkers.eph_walkers import EphWalkers + +#TODO the whole thing, call stuff correctly +def local_energy_holstein( + system: Generic, + hamiltonian: HolsteinModel, + walkers: EphWalkers, + trial: EphTrialWavefunctionBase +): + """""" + energy = xp.zeros((walkers.nwalkers, 5), dtype=xp.complex128) + + #get greens_function from estimators + greensfct = trial.calc_greens_function(walkers) + + #TODO make this nicer + for n in range(walkers.nwalkers): + gf = greensfct[n] + energy[n, 1] = np.einsum('ij->', hamiltonian.T[0] * gf) #NOTE 1e hack + + energy[:, 2] = hamiltonian.const * np.einsum('nii,ni->n', greensfct, walkers.x) + + energy[:, 3] = 0.5 * hamiltonian.m * hamiltonian.w0**2 * np.einsum('ni->n', walkers.x**2) + energy[:, 3] -= 0.5 * hamiltonian.nsites * hamiltonian.w0 + + energy[:, 4] = -0.5 * trial.calc_phonon_laplacian_locenergy(walkers) / hamiltonian.m + + energy[:, 0] = np.sum(energy[:,1:], axis=1) + return energy + + + diff --git a/ipie/hamiltonians/holstein.py b/ipie/hamiltonians/holstein.py new file mode 100644 index 00000000..c7baf097 --- /dev/null +++ b/ipie/hamiltonians/holstein.py @@ -0,0 +1,31 @@ +import numpy +from ipie.hamiltonians.generic_base import GenericBase +from ipie.utils.backend import arraylib as xp + + +class HolsteinModel(GenericBase): + """Class for Holsteib model carrying elph tensor and system parameters + """ + + def __init__(self, g: float, t: float, w0: float, nsites: int, pbc: bool): + self.g = g + self.t = t + self.w0 = w0 + self.m = 1/self.w0 + self.nsites = nsites + self.pbc = pbc + self.T = None + self.const = -self.g * numpy.sqrt(2. * self.m * self.w0) + + def build(self): + self.T = numpy.diag(numpy.ones(self.nsites-1), 1) + self.T += numpy.diag(numpy.ones(self.nsites-1), -1) + + if self.pbc: + self.T[0,-1] = self.T[-1,0] = 1. + + self.T *= -self.t + + self.T = [self.T.copy(), self.T.copy()] + + diff --git a/ipie/propagation/holstein.py b/ipie/propagation/holstein.py new file mode 100644 index 00000000..6fd77742 --- /dev/null +++ b/ipie/propagation/holstein.py @@ -0,0 +1,199 @@ +import numpy +import time +import scipy.linalg + +from ipie.propagation.phaseless_base import PhaselessBase +from ipie.hamiltonians.holstein import HolsteinModel +from ipie.propagation.operations import propagate_one_body #TODO check this +from ipie.utils.backend import synchronize, cast_to_device +from ipie.propagation.continuous_base import PropagatorTimer + +def construct_one_body_propagator(hamiltonian: HolsteinModel, dt: float): + """""" + H1 = hamiltonian.T + expH1 = numpy.array( + [scipy.linalg.expm(-0.5 * dt * H1[0]), scipy.linalg.expm(-0.5 * dt * H1[1])] + ) + return expH1 + + +class HolsteinPropagatorFree(): #for hubbard inherit form phaseless base + """""" + + def __init__(self, time_step, verbose=False): + #super().__init__(time_step, verbose=verbose) + self.dt = time_step + self.verbose = verbose + self.timer = PropagatorTimer() + + self.sqrt_dt = self.dt**0.5 + self.dt_ph = 0.5 * self.dt + self.mpi_handler = None + + def build(self, hamiltonian) : + self.expH1 = construct_one_body_propagator(hamiltonian, self.dt) + self.const = hamiltonian.g * numpy.sqrt(2. * hamiltonian.m * hamiltonian.w0) * self.dt + self.w0 = hamiltonian.w0 + self.m = hamiltonian.m + self.scale = numpy.sqrt(self.dt_ph / self.m) + self.nsites = hamiltonian.nsites + + def propagate_walkers_one_body(self, walkers): + start_time = time.time() + walkers.phia = propagate_one_body(walkers.phia, self.expH1[0]) + if walkers.ndown > 0 and not walkers.rhf: + walkers.phib = propagate_one_body(walkers.phib, self.expH1[1]) + synchronize() + self.timer.tgemm += time.time() - start_time + + def propagate_phonons(self, walkers): + start_time = time.time() + + pot = 0.25 * self.m * self.w0**2 * numpy.einsum('ni->n', walkers.x**2) + pot = numpy.real(pot) + walkers.weight *= numpy.exp(-self.dt_ph * pot) + + N = numpy.random.normal(loc=0.0, scale=self.scale, + size=(walkers.nwalkers, self.nsites)) + walkers.x = walkers.x + N + + pot = 0.25 * self.m * self.w0**2 * numpy.einsum('ni->n', walkers.x**2) + pot = numpy.real(pot) + walkers.weight *= numpy.exp(-self.dt_ph * pot) + + walkers.weight *= numpy.exp(self.dt_ph * self.nsites * self.w0 / 2) #doesnt matter for estimators + + synchronize() + self.timer.tgemm += time.time() - start_time + + def propagate_electron(self, walkers, trial): + start_time = time.time() + ovlp = trial.calc_greens_function(walkers) + synchronize() + self.timer.tgf += time.time() - start_time + + walkers.phia = propagate_one_body(walkers.phia, self.expH1[0]) + if walkers.ndown > 0: + walkers.phib = propagate_one_body(walkers.phib, self.expH1[1]) + + expEph = numpy.exp(self.const * walkers.x) + #NOTE wouldnt it be great to cut the additional index in propagate_one_body? + walkers.phia = numpy.einsum('ni,nie->nie', expEph, walkers.phia) +# walkers.phia = propagate_one_body(walkers.phia, numpy.diag(expEph), H1diag=True) + if walkers.ndown > 0: +# walkers.phib = propagate_one_body(walkers.phib, numpy.diag(expEph), H1diag=True) + walkers.phib = numpy.einsum('ni,nie->nie', expEph, walkers.phib) + + walkers.phia = propagate_one_body(walkers.phia, self.expH1[0]) + if walkers.ndown > 0: + walkers.phib = propagate_one_body(walkers.phib, self.expH1[1]) + + + def propagate_walkers(self, walkers, hamiltonian, trial, eshift=None): + #TODO where should I store the phonon overlaps? + synchronize() + start_time = time.time() + ovlp = trial.calc_overlap(walkers) + synchronize() + self.timer.tgf += time.time() - start_time + + # 2. Update Walkers + # 2.a DMC for phonon degrees of freedom + self.propagate_phonons(walkers) + + # 2.b One-body propagation for electrons + self.propagate_electron(walkers, trial) + + # 2.c DMC for phonon degrees of freedom + self.propagate_phonons(walkers) + + # Update weights (and later do phaseless for multi-electron) + start_time = time.time() + ovlp_new = trial.calc_overlap(walkers) + synchronize() + self.timer.tovlp += time.time() - start_time + + start_time = time.time() + self.update_weight(walkers, ovlp, ovlp_new) + synchronize() + self.timer.tupdate += time.time() - start_time + + + def update_weight(self, walkers, ovlp, ovlp_new): + walkers.weight *= ovlp_new / ovlp + + +class HolsteinPropagatorImportance(HolsteinPropagatorFree): + """""" + def __init__(self, time_step, verbose=False): + super().__init__(time_step, verbose=verbose) + + def propagate_phonons(self, walkers, hamiltonian, trial): + r"""Propagates phonons via DMC + + Parameters + ---------- + walkers : + hamiltonian : + trial : + """ + start_time = time.time() + + #no ZPE in pot -> cancels with ZPE of etrial, wouldn't affect estimators anyways + ph_ovlp_old = trial.calc_phonon_overlap(walkers) + + pot = 0.5 * hamiltonian.m * hamiltonian.w0**2 * numpy.einsum('ni->n',walkers.x**2) + pot -= 0.5 * trial.calc_phonon_laplacian_importance(walkers) / hamiltonian.m + pot = numpy.real(pot) + walkers.weight *= numpy.exp(-self.dt_ph * pot / 2) + + N = numpy.random.normal(loc=0.0, scale=self.scale, size=(walkers.nwalkers, self.nsites)) + drift = trial.calc_phonon_gradient(walkers) + walkers.x = walkers.x + N + self.dt_ph * drift / hamiltonian.m + + ph_ovlp_new = trial.calc_phonon_overlap(walkers) + + pot = 0.5 * hamiltonian.m * hamiltonian.w0**2 * numpy.einsum('ni->n', walkers.x**2) + pot -= 0.5 * trial.calc_phonon_laplacian_importance(walkers) / hamiltonian.m + pot = numpy.real(pot) + walkers.weight *= numpy.exp(-self.dt_ph * pot / 2) + + walkers.weight *= ph_ovlp_old / ph_ovlp_new + walkers.weight *= numpy.exp(self.dt_ph * trial.energy) + + synchronize() + self.timer.tgemm += time.time() - start_time + +# def propagate_walkers(self, walkers, hamiltonian, trial): +# pass + + def propagate_walkers(self, walkers, hamiltonian, trial, eshift=None): + #TODO where should I store the phonon overlaps? + synchronize() + start_time = time.time() + ovlp = trial.calc_overlap(walkers) + synchronize() + self.timer.tgf += time.time() - start_time + + # 2. Update Walkers + # 2.a DMC for phonon degrees of freedom + self.propagate_phonons(walkers, hamiltonian, trial) + + # 2.b One-body propagation for electrons + self.propagate_electron(walkers, trial) + + # 2.c DMC for phonon degrees of freedom + self.propagate_phonons(walkers, hamiltonian, trial) + + # Update weights (and later do phaseless for multi-electron) + start_time = time.time() + ovlp_new = trial.calc_overlap(walkers) + synchronize() + self.timer.tovlp += time.time() - start_time + + start_time = time.time() + self.update_weight(walkers, ovlp, ovlp_new) + synchronize() + self.timer.tupdate += time.time() - start_time + + diff --git a/ipie/propagation/propagator.py b/ipie/propagation/propagator.py index 0200cd11..2f56f21d 100644 --- a/ipie/propagation/propagator.py +++ b/ipie/propagation/propagator.py @@ -1,4 +1,6 @@ from ipie.hamiltonians.generic import GenericRealChol, GenericComplexChol +from ipie.hamiltonians.holstein import HolsteinModel from ipie.propagation.phaseless_generic import PhaselessGeneric +from ipie.propagation.holstein import HolsteinPropagatorImportance -Propagator = {GenericRealChol: PhaselessGeneric, GenericComplexChol: PhaselessGeneric} +Propagator = {GenericRealChol: PhaselessGeneric, GenericComplexChol: PhaselessGeneric, HolsteinModel: HolsteinPropagatorImportance} diff --git a/ipie/trial_wavefunction/holstein/coherent_state.py b/ipie/trial_wavefunction/holstein/coherent_state.py new file mode 100644 index 00000000..9dae69ac --- /dev/null +++ b/ipie/trial_wavefunction/holstein/coherent_state.py @@ -0,0 +1,68 @@ +import numpy as np +import scipy.linalg + +#TODO add greens_function_coherent_state in estimators +from ipie.estimators.greens_function_single_det import greens_function_single_det +from ipie.trial_wavefunction.eph_trial_base import EphTrialWavefunctionBase + +#TODO greensfunctions are in estimators +class CoherentState(EphTrialWavefunctionBase): + """""" + def __init__(self, wavefunction, num_elec, num_basis, verbose=False): + super().__init__(wavefunction, num_elec, num_basis, verbose=verbose) + self.num_elec = num_elec + self.nalpha, self.nbeta = self.num_elec + self.w0 = hamiltonian.w0 + self.m = hamiltonian.m + self.nsites = hamiltonian.nsites + + self.psia = wavefunction[:self.nalpha, :] + self.psia = self.psia[0] + #NOTE 1e hack + self.psib = wavefunction[self.nalpha:self.nalpha+self.nbeta, :] + self.beta_shift = np.squeeze(wavefunction[-1, :]) + + + def build(self, walkers) -> None: + walkers.ph_ovlp = np.zeros(walkers.nwalkers) + walkers.el_ovlp = np.zeros(walkers.nwalkers) + walkers.total_ovlp = np.zeros(self.nsites) + + def calculate_energy(self, system, hamiltonian): + #TODO variational_energy_coherent_state in ipie.estimators.local_energy + ... + + def calc_overlap(self, walkers) -> np.ndarray: + #TODO this will be a concoction of phonon and electronic overlap + + + def calc_phonon_overlap(self, walkers) -> np.ndarray: + walker.ph_ov = np.exp(-(self.m * self.w0 / 2) * (walkers.x - self.beta_shift)**2) + walker.ph_ov = np.prod(ph_ov, axis=1) + return ph_ov + + def calc_phonon_gradient(self, walkers) -> np.ndarray: + grad = np.einsum('ni,n->ni', (walkers.x - self.beta_shift), ovlp) + grad *= -self.m * self.w0 + return grad + + def calc_phonon_laplacian(self, walkers) -> np.ndarray: + arg = (walkers.x - self.beta_shift) * self.m * self.w0 + arg2 = arg**2 + laplacian = np.sum(arg2, axis=1) - self.nsites * self.m * self.w0 + return laplacian + + def calc_phonon_laplacian_importance(self, walkers) -> np.ndarray: + return self.calc_phonon_laplacian(walkers) + + def calc_phonon_laplacian_locenergy(self, walkers) -> np.ndarray: + return self.calc_phonon_laplacian(walkers) + + def calc_electronic_overlap(self, walkers) -> np.ndarray: + walkers.el_ovlp[ip] = np.einsum('i,nie->n', self.psia[perm].conj(), walkers.phia) + if self.nbeta > 0: + pass + + + def calc_greens_function(self, walkers) -> np.ndarray: + diff --git a/ipie/trial_wavefunction/holstein/eph_trial_base.py b/ipie/trial_wavefunction/holstein/eph_trial_base.py new file mode 100644 index 00000000..21219a0c --- /dev/null +++ b/ipie/trial_wavefunction/holstein/eph_trial_base.py @@ -0,0 +1,30 @@ +import numpy as np +from ipie.trial_wavefunction.wavefunction_base import TrialWavefunctionBase + +class EphTrialWavefunctionBase(): + def __init__(self, wavefunction, num_elec, num_basis, verbose=False): + #super().__init__(wavefunction, num_elec, num_basis, verbose=verbose) + self.nelec = num_elec + self.nbasis = num_basis + self.nalpha, self.nbeta = self.nelec + self.verbose = verbose + self._num_dets = 0 + self._max_num_dets = self._num_dets + self.ortho_expansion = False + self.optimized = True + + self.psia = wavefunction[:self.nalpha] + self.psib = wavefunction[self.nalpha:self.nalpha+self.nbeta] + self.beta_shift = wavefunction[self.nalpha+self.nbeta:] + + self.compute_trial_energy = False + self.energy = None + + def build(self) -> None: + pass + + def set_etrial(self, energy) -> None: + self.energy = energy + + + diff --git a/ipie/trial_wavefunction/holstein/toyozawa.py b/ipie/trial_wavefunction/holstein/toyozawa.py new file mode 100644 index 00000000..277a07fe --- /dev/null +++ b/ipie/trial_wavefunction/holstein/toyozawa.py @@ -0,0 +1,110 @@ +import numpy as np +import scipy.linalg + +#TODO add greens_function_coherent_state in estimators +#NOTE can use single_det for now +from ipie.estimators.greens_function_single_det import greens_function_single_det +from ipie.trial_wavefunction.holstein.eph_trial_base import EphTrialWavefunctionBase + + +#TODO greensfunctions are in estimators + +def circ_perm(lst): + """""" + cpy = lst[:] # take a copy because a list is a mutable object + yield cpy + for i in range(len(lst) - 1): + cpy = cpy[1:] + [cpy[0]] + yield cpy + +class ToyozawaTrial(EphTrialWavefunctionBase): + """""" + def __init__(self, wavefunction, hamiltonian, num_elec, num_basis, verbose=False): + super().__init__(wavefunction, num_elec, num_basis, verbose=verbose) + self.perms = list(circ_perm([i for i in range(self.nbasis)])) + self.nperms = len(self.perms) + self.num_elec = num_elec + self.nalpha, self.nbeta = self.num_elec + self.w0 = hamiltonian.w0 + self.m = hamiltonian.m + self.nsites = hamiltonian.nsites + + self.psia = wavefunction[:self.nalpha, :] + self.psia = self.psia[0] + #NOTE 1e hack + self.psib = wavefunction[self.nalpha:self.nalpha+self.nbeta, :] + self.beta_shift = np.squeeze(wavefunction[-1, :]) + + def build(self, walkers) -> None: + walkers.ph_ovlp = np.zeros((self.nperms, walkers.nwalkers)) + walkers.el_ovlp = np.zeros((self.nperms, walkers.nwalkers)) + walkers.total_ovlp = np.zeros(self.nsites) + + def calculate_energy(self, system, hamiltonian): + #TODO variational_energy_coherent_state in ipie.estimators.local_energy + pass + + def calc_overlap(self, walkers) -> np.ndarray: + #TODO this will be a concoction of phonon and electronic overlap + _ = self.calc_phonon_overlap(walkers) + _ = self.calc_electronic_overlap(walkers) + walkers.total_ovlp = np.einsum('pn,pn->pn', walkers.el_ovlp, walkers.ph_ovlp) + total_ovlp = np.sum(walkers.total_ovlp, axis=0) + return total_ovlp + + + def calc_phonon_overlap(self, walkers) -> np.ndarray: + """""" + for ip,perm in enumerate(self.perms): + ph_ov = np.exp(-(self.m * self.w0 / 2) * (walkers.x - self.beta_shift[perm])**2) + walkers.ph_ovlp[ip] = np.prod(ph_ov, axis=1) + ph_ovlp = np.sum(walkers.ph_ovlp, axis=0) + return ph_ovlp + + def calc_phonon_gradient(self, walkers) -> np.ndarray: + r"""No reevaluation of phonon overlap because it reuses the overlap from the previous + evaluation of the laplacian. The gradient only surfaces in the quantum force.""" + grad = np.zeros_like(walkers.x) + for ovlp, perm in zip(walkers.ph_ovlp, self.perms): + grad += np.einsum('ni,n->ni', (walkers.x - self.beta_shift[perm]), ovlp) + grad *= -self.m * self.w0 + grad = np.einsum('ni,n->ni', grad, 1/np.sum(walkers.ph_ovlp, axis=0)) + return grad + + def calc_phonon_laplacian(self, walkers, ovlps) -> np.ndarray: + r"""""" + laplacian = np.zeros(walkers.nwalkers, dtype=np.complex128) + for ovlp, perm in zip(ovlps, self.perms): + arg = (walkers.x - self.beta_shift[perm]) * self.m * self.w0 + arg2 = arg**2 + laplacian += (np.sum(arg2, axis=1) - self.nsites * self.m * self.w0) * ovlp + laplacian /= np.sum(ovlps, axis=0) + return laplacian + + def calc_phonon_laplacian_importance(self, walkers) -> np.ndarray: + return self.calc_phonon_laplacian(walkers, walkers.ph_ovlp) + + def calc_phonon_laplacian_locenergy(self, walkers) -> np.ndarray: + return self.calc_phonon_laplacian(walkers, walkers.total_ovlp) + + + def calc_electronic_overlap(self, walkers) -> np.ndarray: + """""" + for ip,perm in enumerate(self.perms): + walkers.el_ovlp[ip] = np.einsum('i,nie->n', self.psia[perm].conj(), walkers.phia) #this is single electron + if self.nbeta > 0: + pass #TODO -> adjust ovlps shape + el_ovlp = np.sum(walkers.el_ovlp, axis=0) #NOTE this was ph before?? + return el_ovlp + + def calc_greens_function(self, walkers) -> np.ndarray: + """""" + greensfct = np.zeros((walkers.nwalkers, self.nsites, self.nsites), dtype=np.complex128) + for ovlp, perm in zip(walkers.total_ovlp, self.perms): + overlap_inv = 1 / np.einsum('i,nie->n', self.psia[perm].conj(), walkers.phia) #NOTE psi currently hacked + greensfct += np.einsum('nie,n,j,n->nji', walkers.phia, overlap_inv, + self.psia[perm].conj(), ovlp) + greensfct = np.einsum('nij,n->nij', greensfct, 1 / np.sum(walkers.total_ovlp, axis=0)) + return greensfct + + diff --git a/ipie/walkers/eph_walkers.py b/ipie/walkers/eph_walkers.py new file mode 100644 index 00000000..42faf4bd --- /dev/null +++ b/ipie/walkers/eph_walkers.py @@ -0,0 +1,120 @@ +import numpy + +from ipie.config import config +from ipie.utils.backend import arraylib as xp +from ipie.utils.backend import cast_to_device, qr, qr_mode, synchronize +from ipie.walkers.base_walkers import BaseWalkers + +class EphWalkers(BaseWalkers): + """""" + def __init__( + self, + initial_walker: numpy.ndarray, + nup: int, + ndown: int, + nbasis: int, + nwalkers: int, + mpi_handler, + verbose: bool = False + ): + + self.nup = nup + self.ndown = ndown + self.nbasis = nbasis + self.mpi_handler = mpi_handler + + + super().__init__(nwalkers, verbose=verbose) + + #TODO is there a reason we dont use numpy tile for these? + self.phia = xp.array( + [initial_walker[:, : self.nup].copy() for iw in range(self.nwalkers)], + dtype=xp.complex128, + ) +# self.phia = numpy.squeeze(self.phia) #NOTE: 1e hack to work with 1e overlaps + + self.phib = xp.array( + [initial_walker[:, self.nup:self.nup+self.ndown].copy() for iw in range(self.nwalkers)], + dtype=xp.complex128, + ) + self.x = xp.array( + [initial_walker[:, self.nup+self.ndown:].copy() for iw in range(self.nwalkers)], + dtype=xp.complex128 + ) + self.x = numpy.squeeze(self.x) #NOTE: 1e hack to work with 1e overlaps + + self.buff_names += ["phia", "phib", "x"] + + self.buff_size = round(self.set_buff_size_single_walker() / float(self.nwalkers)) + self.walker_buffer = numpy.zeros(self.buff_size, dtype=numpy.complex128) + + def build(self, trial): + self.ovlp = trial.calc_greens_function(self) + + def cast_to_cupy(self, verbose=False): + cast_to_device(self, verbose) + + def reortho(self): + """reorthogonalise walkers. + + parameters + ---------- + """ + if config.get_option("use_gpu"): + return self.reortho_batched() + ndown = self.ndown + detR = [] + for iw in range(self.nwalkers): + (self.phia[iw], Rup) = qr(self.phia[iw], mode=qr_mode) + # TODO: FDM This isn't really necessary, the absolute value of the + # weight is used for population control so this shouldn't matter. + # I think this is a legacy thing. + # Wanted detR factors to remain positive, dump the sign in orbitals. + Rup_diag = xp.diag(Rup) + signs_up = xp.sign(Rup_diag) + self.phia[iw] = xp.dot(self.phia[iw], xp.diag(signs_up)) + + # include overlap factor + # det(R) = \prod_ii R_ii + # det(R) = exp(log(det(R))) = exp((sum_i log R_ii) - C) + # C factor included to avoid over/underflow + log_det = xp.sum(xp.log(xp.abs(Rup_diag))) + + if ndown > 0: + (self.phib[iw], Rdn) = qr(self.phib[iw], mode=qr_mode) + Rdn_diag = xp.diag(Rdn) + signs_dn = xp.sign(Rdn_diag) + self.phib[iw] = xp.dot(self.phib[iw], xp.diag(signs_dn)) + log_det += sum(xp.log(abs(Rdn_diag))) + + detR += [xp.exp(log_det - self.detR_shift[iw])] + self.log_detR[iw] += xp.log(detR[iw]) + self.detR[iw] = detR[iw] + self.ovlp[iw] = self.ovlp[iw] / detR[iw] + + synchronize() + return detR + + def reortho_batched(self): + """reorthogonalise walkers. + + parameters + ---------- + """ + assert config.get_option("use_gpu") + (self.phia, Rup) = qr(self.phia, mode=qr_mode) + Rup_diag = xp.einsum("wii->wi", Rup) + log_det = xp.einsum("wi->w", xp.log(abs(Rup_diag))) + + if self.ndown > 0: + (self.phib, Rdn) = qr(self.phib, mode=qr_mode) + Rdn_diag = xp.einsum("wii->wi", Rdn) + log_det += xp.einsum("wi->w", xp.log(abs(Rdn_diag))) + self.detR = xp.exp(log_det - self.detR_shift) + self.ovlp = self.ovlp / self.detR + + synchronize() + + return self.detR + + From fa6f13b67949e1a8456b2a4b19d9e51ae6bac7df Mon Sep 17 00:00:00 2001 From: m-baumgarten Date: Fri, 16 Feb 2024 11:14:32 -0500 Subject: [PATCH 03/61] working toyozawa implementation for 1 electron, 1D Holstein --- ipie/estimators/energy.py | 4 +- ipie/estimators/local_energy_holstein.py | 16 ++++-- ipie/hamiltonians/elph/generic.py | 11 ++++ ipie/hamiltonians/{ => elph}/holstein.py | 2 +- ipie/propagation/holstein.py | 50 ++++++++++++----- ipie/propagation/propagator.py | 2 +- ipie/trial_wavefunction/holstein/toyozawa.py | 56 ++++++++------------ ipie/walkers/eph_walkers.py | 28 ++++++++-- 8 files changed, 113 insertions(+), 56 deletions(-) create mode 100644 ipie/hamiltonians/elph/generic.py rename ipie/hamiltonians/{ => elph}/holstein.py (92%) diff --git a/ipie/estimators/energy.py b/ipie/estimators/energy.py index 883661a1..0ca2b4e2 100644 --- a/ipie/estimators/energy.py +++ b/ipie/estimators/energy.py @@ -31,7 +31,7 @@ local_energy_multi_det_trial_wicks_batch_opt_chunked, ) from ipie.hamiltonians.generic import GenericComplexChol, GenericRealChol -from ipie.hamiltonians.holstein import HolsteinModel +from ipie.hamiltonians.elph.holstein import HolsteinModel from ipie.systems.generic import Generic from ipie.trial_wavefunction.noci import NOCI from ipie.trial_wavefunction.particle_hole import ( @@ -151,7 +151,7 @@ def compute_estimator(self, system, walkers, hamiltonian, trial, istep=1): self._data["EDenom"] = xp.sum(walkers.weight) self._data["E1Body"] = xp.sum(walkers.weight * energy[:, 1].real) self._data["E2Body"] = xp.sum(walkers.weight * energy[:, 2].real) - + return self.data def get_index(self, name): diff --git a/ipie/estimators/local_energy_holstein.py b/ipie/estimators/local_energy_holstein.py index 97e609ab..9c3175f6 100644 --- a/ipie/estimators/local_energy_holstein.py +++ b/ipie/estimators/local_energy_holstein.py @@ -1,11 +1,13 @@ import numpy as np -from ipie.hamiltonians.holstein import HolsteinModel +from ipie.hamiltonians.elph.holstein import HolsteinModel from ipie.systems.generic import Generic from ipie.trial_wavefunction.holstein.eph_trial_base import EphTrialWavefunctionBase from ipie.utils.backend import arraylib as xp from ipie.walkers.eph_walkers import EphWalkers +verbose=False + #TODO the whole thing, call stuff correctly def local_energy_holstein( system: Generic, @@ -18,20 +20,28 @@ def local_energy_holstein( #get greens_function from estimators greensfct = trial.calc_greens_function(walkers) - + if verbose: + print('gf walker0: ', greensfct[0]) #TODO make this nicer for n in range(walkers.nwalkers): gf = greensfct[n] energy[n, 1] = np.einsum('ij->', hamiltonian.T[0] * gf) #NOTE 1e hack - energy[:, 2] = hamiltonian.const * np.einsum('nii,ni->n', greensfct, walkers.x) + if verbose: + print('const, phonons, elph: ', hamiltonian.const, walkers.x[0,0], energy[0, 2]) energy[:, 3] = 0.5 * hamiltonian.m * hamiltonian.w0**2 * np.einsum('ni->n', walkers.x**2) energy[:, 3] -= 0.5 * hamiltonian.nsites * hamiltonian.w0 + if verbose: + print('e_ph pot: ', energy[0, 3]) energy[:, 4] = -0.5 * trial.calc_phonon_laplacian_locenergy(walkers) / hamiltonian.m + if verbose: + print('e_ph kin: ', energy[0, 4]) energy[:, 0] = np.sum(energy[:,1:], axis=1) + if verbose: + print('energy: ', energy[0,0]) return energy diff --git a/ipie/hamiltonians/elph/generic.py b/ipie/hamiltonians/elph/generic.py new file mode 100644 index 00000000..84df1a80 --- /dev/null +++ b/ipie/hamiltonians/elph/generic.py @@ -0,0 +1,11 @@ +import numpy +from ipie.hamiltonians.generic_base import GenericBase +from ipie.utils.backend import arraylib as xp + + +class GenericElphBase(GenericBase): + """""" + + def __init__(self, ): + pass + diff --git a/ipie/hamiltonians/holstein.py b/ipie/hamiltonians/elph/holstein.py similarity index 92% rename from ipie/hamiltonians/holstein.py rename to ipie/hamiltonians/elph/holstein.py index c7baf097..46354b94 100644 --- a/ipie/hamiltonians/holstein.py +++ b/ipie/hamiltonians/elph/holstein.py @@ -4,7 +4,7 @@ class HolsteinModel(GenericBase): - """Class for Holsteib model carrying elph tensor and system parameters + """Class for Holstein model carrying elph tensor and system parameters """ def __init__(self, g: float, t: float, w0: float, nsites: int, pbc: bool): diff --git a/ipie/propagation/holstein.py b/ipie/propagation/holstein.py index 6fd77742..21ddd9ab 100644 --- a/ipie/propagation/holstein.py +++ b/ipie/propagation/holstein.py @@ -3,11 +3,13 @@ import scipy.linalg from ipie.propagation.phaseless_base import PhaselessBase -from ipie.hamiltonians.holstein import HolsteinModel +from ipie.hamiltonians.elph.holstein import HolsteinModel from ipie.propagation.operations import propagate_one_body #TODO check this from ipie.utils.backend import synchronize, cast_to_device from ipie.propagation.continuous_base import PropagatorTimer +verbose=False + def construct_one_body_propagator(hamiltonian: HolsteinModel, dt: float): """""" H1 = hamiltonian.T @@ -138,10 +140,9 @@ def propagate_phonons(self, walkers, hamiltonian, trial): trial : """ start_time = time.time() - #no ZPE in pot -> cancels with ZPE of etrial, wouldn't affect estimators anyways ph_ovlp_old = trial.calc_phonon_overlap(walkers) - + pot = 0.5 * hamiltonian.m * hamiltonian.w0**2 * numpy.einsum('ni->n',walkers.x**2) pot -= 0.5 * trial.calc_phonon_laplacian_importance(walkers) / hamiltonian.m pot = numpy.real(pot) @@ -157,33 +158,48 @@ def propagate_phonons(self, walkers, hamiltonian, trial): pot -= 0.5 * trial.calc_phonon_laplacian_importance(walkers) / hamiltonian.m pot = numpy.real(pot) walkers.weight *= numpy.exp(-self.dt_ph * pot / 2) - + walkers.weight *= ph_ovlp_old / ph_ovlp_new walkers.weight *= numpy.exp(self.dt_ph * trial.energy) - + synchronize() self.timer.tgemm += time.time() - start_time -# def propagate_walkers(self, walkers, hamiltonian, trial): -# pass def propagate_walkers(self, walkers, hamiltonian, trial, eshift=None): - #TODO where should I store the phonon overlaps? + """""" + synchronize() start_time = time.time() - ovlp = trial.calc_overlap(walkers) + + if verbose: + print('before ovlp: ', walkers.phia[0,0,0], walkers.x[0,0], walkers.weight[0]) + + ovlp = trial.calc_overlap(walkers).copy() + + if verbose: + print('ovlp: ', ovlp[0]) + synchronize() self.timer.tgf += time.time() - start_time # 2. Update Walkers # 2.a DMC for phonon degrees of freedom self.propagate_phonons(walkers, hamiltonian, trial) - + + if verbose: + print('weight in prop I: ', walkers.weight[0]) + # 2.b One-body propagation for electrons self.propagate_electron(walkers, trial) - + # 2.c DMC for phonon degrees of freedom self.propagate_phonons(walkers, hamiltonian, trial) + + if verbose: + print('weight in prop II: ', walkers.weight[0]) + +# print('electron & phonon & weight: ', walkers.phia[0], walkers.x, walkers.weight) # Update weights (and later do phaseless for multi-electron) start_time = time.time() @@ -192,8 +208,16 @@ def propagate_walkers(self, walkers, hamiltonian, trial, eshift=None): self.timer.tovlp += time.time() - start_time start_time = time.time() - self.update_weight(walkers, ovlp, ovlp_new) + + if verbose: + print('new and old overlaps: ', ovlp_new[0], ovlp[0]) + + self.update_weight(walkers, ovlp, ovlp_new) #walkers.previously ovlp_for_weight_update +# walkers.total_ovlp_for_weight_update = ovlp_new + + if verbose: + print('weight in prop III: ', walkers.weight[0]) + synchronize() self.timer.tupdate += time.time() - start_time - diff --git a/ipie/propagation/propagator.py b/ipie/propagation/propagator.py index 2f56f21d..a121122d 100644 --- a/ipie/propagation/propagator.py +++ b/ipie/propagation/propagator.py @@ -1,5 +1,5 @@ from ipie.hamiltonians.generic import GenericRealChol, GenericComplexChol -from ipie.hamiltonians.holstein import HolsteinModel +from ipie.hamiltonians.elph.holstein import HolsteinModel from ipie.propagation.phaseless_generic import PhaselessGeneric from ipie.propagation.holstein import HolsteinPropagatorImportance diff --git a/ipie/trial_wavefunction/holstein/toyozawa.py b/ipie/trial_wavefunction/holstein/toyozawa.py index 277a07fe..2bab7dcc 100644 --- a/ipie/trial_wavefunction/holstein/toyozawa.py +++ b/ipie/trial_wavefunction/holstein/toyozawa.py @@ -5,7 +5,9 @@ #NOTE can use single_det for now from ipie.estimators.greens_function_single_det import greens_function_single_det from ipie.trial_wavefunction.holstein.eph_trial_base import EphTrialWavefunctionBase +from ipie.trial_wavefunction.holstein.coherent_state import CoherentStateTrial +verbose=False #TODO greensfunctions are in estimators @@ -17,28 +19,12 @@ def circ_perm(lst): cpy = cpy[1:] + [cpy[0]] yield cpy -class ToyozawaTrial(EphTrialWavefunctionBase): +class ToyozawaTrial(CoherentStateTrial): """""" def __init__(self, wavefunction, hamiltonian, num_elec, num_basis, verbose=False): - super().__init__(wavefunction, num_elec, num_basis, verbose=verbose) + super().__init__(wavefunction, hamiltonian, num_elec, num_basis, verbose=verbose) self.perms = list(circ_perm([i for i in range(self.nbasis)])) self.nperms = len(self.perms) - self.num_elec = num_elec - self.nalpha, self.nbeta = self.num_elec - self.w0 = hamiltonian.w0 - self.m = hamiltonian.m - self.nsites = hamiltonian.nsites - - self.psia = wavefunction[:self.nalpha, :] - self.psia = self.psia[0] - #NOTE 1e hack - self.psib = wavefunction[self.nalpha:self.nalpha+self.nbeta, :] - self.beta_shift = np.squeeze(wavefunction[-1, :]) - - def build(self, walkers) -> None: - walkers.ph_ovlp = np.zeros((self.nperms, walkers.nwalkers)) - walkers.el_ovlp = np.zeros((self.nperms, walkers.nwalkers)) - walkers.total_ovlp = np.zeros(self.nsites) def calculate_energy(self, system, hamiltonian): #TODO variational_energy_coherent_state in ipie.estimators.local_energy @@ -48,37 +34,40 @@ def calc_overlap(self, walkers) -> np.ndarray: #TODO this will be a concoction of phonon and electronic overlap _ = self.calc_phonon_overlap(walkers) _ = self.calc_electronic_overlap(walkers) - walkers.total_ovlp = np.einsum('pn,pn->pn', walkers.el_ovlp, walkers.ph_ovlp) - total_ovlp = np.sum(walkers.total_ovlp, axis=0) - return total_ovlp + walkers.total_ovlp = np.einsum('np,np->np', walkers.el_ovlp, walkers.ph_ovlp) + walkers.ovlp = np.sum(walkers.total_ovlp, axis=1) + return walkers.ovlp def calc_phonon_overlap(self, walkers) -> np.ndarray: """""" + ph_ovlp = np.zeros(walkers.nwalkers, dtype=np.complex128) for ip,perm in enumerate(self.perms): ph_ov = np.exp(-(self.m * self.w0 / 2) * (walkers.x - self.beta_shift[perm])**2) - walkers.ph_ovlp[ip] = np.prod(ph_ov, axis=1) - ph_ovlp = np.sum(walkers.ph_ovlp, axis=0) + walkers.ph_ovlp[:, ip] = np.prod(ph_ov, axis=1) + ph_ovlp = np.sum(walkers.ph_ovlp, axis=1) + if verbose: + print('ph_ovlp: ', ph_ovlp[0]) return ph_ovlp def calc_phonon_gradient(self, walkers) -> np.ndarray: r"""No reevaluation of phonon overlap because it reuses the overlap from the previous evaluation of the laplacian. The gradient only surfaces in the quantum force.""" - grad = np.zeros_like(walkers.x) - for ovlp, perm in zip(walkers.ph_ovlp, self.perms): + grad = np.zeros_like(walkers.x, dtype=np.complex128) + for ovlp, perm in zip(walkers.ph_ovlp.T, self.perms): grad += np.einsum('ni,n->ni', (walkers.x - self.beta_shift[perm]), ovlp) grad *= -self.m * self.w0 - grad = np.einsum('ni,n->ni', grad, 1/np.sum(walkers.ph_ovlp, axis=0)) + grad = np.einsum('ni,n->ni', grad, 1/np.sum(walkers.ph_ovlp, axis=1)) return grad def calc_phonon_laplacian(self, walkers, ovlps) -> np.ndarray: r"""""" laplacian = np.zeros(walkers.nwalkers, dtype=np.complex128) - for ovlp, perm in zip(ovlps, self.perms): + for ovlp, perm in zip(ovlps.T, self.perms): arg = (walkers.x - self.beta_shift[perm]) * self.m * self.w0 arg2 = arg**2 laplacian += (np.sum(arg2, axis=1) - self.nsites * self.m * self.w0) * ovlp - laplacian /= np.sum(ovlps, axis=0) + laplacian /= np.sum(ovlps, axis=1) return laplacian def calc_phonon_laplacian_importance(self, walkers) -> np.ndarray: @@ -87,24 +76,25 @@ def calc_phonon_laplacian_importance(self, walkers) -> np.ndarray: def calc_phonon_laplacian_locenergy(self, walkers) -> np.ndarray: return self.calc_phonon_laplacian(walkers, walkers.total_ovlp) - def calc_electronic_overlap(self, walkers) -> np.ndarray: """""" for ip,perm in enumerate(self.perms): - walkers.el_ovlp[ip] = np.einsum('i,nie->n', self.psia[perm].conj(), walkers.phia) #this is single electron + walkers.el_ovlp[:, ip] = np.einsum('i,nie->n', self.psia[perm].conj(), walkers.phia) #this is single electron if self.nbeta > 0: pass #TODO -> adjust ovlps shape - el_ovlp = np.sum(walkers.el_ovlp, axis=0) #NOTE this was ph before?? + el_ovlp = np.sum(walkers.el_ovlp, axis=1) #NOTE this was ph before?? + if verbose: + print('el_ovlp: ', el_ovlp[0]) return el_ovlp def calc_greens_function(self, walkers) -> np.ndarray: """""" greensfct = np.zeros((walkers.nwalkers, self.nsites, self.nsites), dtype=np.complex128) - for ovlp, perm in zip(walkers.total_ovlp, self.perms): + for ovlp, perm in zip(walkers.total_ovlp.T, self.perms): overlap_inv = 1 / np.einsum('i,nie->n', self.psia[perm].conj(), walkers.phia) #NOTE psi currently hacked greensfct += np.einsum('nie,n,j,n->nji', walkers.phia, overlap_inv, self.psia[perm].conj(), ovlp) - greensfct = np.einsum('nij,n->nij', greensfct, 1 / np.sum(walkers.total_ovlp, axis=0)) + greensfct = np.einsum('nij,n->nij', greensfct, 1 / np.sum(walkers.total_ovlp, axis=1)) #these sums can be replaced by walkers.ovlp calls return greensfct diff --git a/ipie/walkers/eph_walkers.py b/ipie/walkers/eph_walkers.py index 42faf4bd..377a9af3 100644 --- a/ipie/walkers/eph_walkers.py +++ b/ipie/walkers/eph_walkers.py @@ -4,6 +4,7 @@ from ipie.utils.backend import arraylib as xp from ipie.utils.backend import cast_to_device, qr, qr_mode, synchronize from ipie.walkers.base_walkers import BaseWalkers +from ipie.trial_wavefunction.holstein.toyozawa import ToyozawaTrial class EphWalkers(BaseWalkers): """""" @@ -26,6 +27,8 @@ def __init__( super().__init__(nwalkers, verbose=verbose) + self.weight = numpy.ones(self.nwalkers, dtype=numpy.complex128) + #TODO is there a reason we dont use numpy tile for these? self.phia = xp.array( [initial_walker[:, : self.nup].copy() for iw in range(self.nwalkers)], @@ -41,15 +44,31 @@ def __init__( [initial_walker[:, self.nup+self.ndown:].copy() for iw in range(self.nwalkers)], dtype=xp.complex128 ) - self.x = numpy.squeeze(self.x) #NOTE: 1e hack to work with 1e overlaps - + self.x = numpy.squeeze(self.x) + self.buff_names += ["phia", "phib", "x"] self.buff_size = round(self.set_buff_size_single_walker() / float(self.nwalkers)) self.walker_buffer = numpy.zeros(self.buff_size, dtype=numpy.complex128) def build(self, trial): - self.ovlp = trial.calc_greens_function(self) + """NOTE: total_ovlp is the same as ovlp for coherent state trial, and just + serves the purpose of not recomputing overlaps for each permutation + but passing it to the pop_control and adjsuting it accordingly for the + Toyozawa trial.""" + + if isinstance(trial, ToyozawaTrial): + shape = (self.nwalkers, trial.nperms) + else: + shape = self.nwalkers + + self.ph_ovlp = numpy.zeros(shape, dtype=numpy.complex128) + self.el_ovlp = numpy.zeros(shape, dtype=numpy.complex128) + self.total_ovlp = numpy.zeros(shape, dtype=numpy.complex128) + + self.buff_names += ['total_ovlp'] #, 'el_ovlp', 'total_ovlp'] #not really necessary to bring 'el_ovlp', 'total_ovlp' along if computing overlap after normalization anyways. + self.buff_size = round(self.set_buff_size_single_walker() / float(self.nwalkers)) + self.walker_buffer = numpy.zeros(self.buff_size, dtype=numpy.complex128) def cast_to_cupy(self, verbose=False): cast_to_device(self, verbose) @@ -90,6 +109,9 @@ def reortho(self): detR += [xp.exp(log_det - self.detR_shift[iw])] self.log_detR[iw] += xp.log(detR[iw]) self.detR[iw] = detR[iw] + + self.el_ovlp[iw, :] = self.el_ovlp[iw, :] / detR[iw] + self.total_ovlp[iw, :] = self.total_ovlp[iw, :] / detR[iw] self.ovlp[iw] = self.ovlp[iw] / detR[iw] synchronize() From caf7a34bb44c28a228fe95ae339a825a2506f6cc Mon Sep 17 00:00:00 2001 From: m-baumgarten Date: Thu, 22 Feb 2024 23:33:21 -0500 Subject: [PATCH 04/61] multi-electron implementation. toyozawa very dodgy --- ipie/estimators/energy_holstein.py | 35 ---------- ipie/estimators/local_energy_holstein.py | 26 ++----- .../holstein/coherent_state.py | 69 ++++++++++++------- ipie/trial_wavefunction/holstein/toyozawa.py | 63 ++++++++++++----- ipie/walkers/eph_walkers.py | 22 +++--- 5 files changed, 110 insertions(+), 105 deletions(-) delete mode 100644 ipie/estimators/energy_holstein.py diff --git a/ipie/estimators/energy_holstein.py b/ipie/estimators/energy_holstein.py deleted file mode 100644 index b59c09ba..00000000 --- a/ipie/estimators/energy_holstein.py +++ /dev/null @@ -1,35 +0,0 @@ -from ipie.estimators.estimator_base import EstimatorBase -from ipie.hamiltonians.holstein import HolsteinModel -from ipie.systems.generic import Generic - - -def local_energy( - system: Generic, hamiltonian: HolsteinModel, walkers: EphWalkers, trial: EphTrial -): - return local_energy_holstein(system, hamiltonian, walkers, trial) - - -class EnergyEstimatorHolstein(EstimatorBase): - """""" - def __init__(self, system, hamiltonian, trial): - assert system is not None - assert ham is not None - assert trial is not None - super().__init__() - self.scalar_estimator = True - self._data = { - "ENumer": 0.0j, - "EDenom": 0.0j, - "ETotal": 0.0j - } - self._shape = (len(self.names),) - self._data_index = {k: i for i, k in enumerate(list(self._data.keys()))} - self.print_to_stdout = True - self.ascii_filename = filename - - def compute_estimator(self, system, walkers, hamiltonian, trial, istep)=1): - trial.calc_greens_function(walkers) - energy = local_energy(system, hamiltonian, walkers, trial) - self._data["ENumer"] = xp.sum(walkers.weight * energy[:, 0].real) - self._data["EDenom"] = xp.sum(walkers.weight) - diff --git a/ipie/estimators/local_energy_holstein.py b/ipie/estimators/local_energy_holstein.py index 9c3175f6..4b46a41d 100644 --- a/ipie/estimators/local_energy_holstein.py +++ b/ipie/estimators/local_energy_holstein.py @@ -6,9 +6,7 @@ from ipie.utils.backend import arraylib as xp from ipie.walkers.eph_walkers import EphWalkers -verbose=False -#TODO the whole thing, call stuff correctly def local_energy_holstein( system: Generic, hamiltonian: HolsteinModel, @@ -19,29 +17,19 @@ def local_energy_holstein( energy = xp.zeros((walkers.nwalkers, 5), dtype=xp.complex128) #get greens_function from estimators - greensfct = trial.calc_greens_function(walkers) - if verbose: - print('gf walker0: ', greensfct[0]) + gf = trial.calc_greens_function(walkers) + #TODO make this nicer for n in range(walkers.nwalkers): - gf = greensfct[n] - energy[n, 1] = np.einsum('ij->', hamiltonian.T[0] * gf) #NOTE 1e hack - energy[:, 2] = hamiltonian.const * np.einsum('nii,ni->n', greensfct, walkers.x) - if verbose: - print('const, phonons, elph: ', hamiltonian.const, walkers.x[0,0], energy[0, 2]) - + energy[n, 1] = np.einsum('ij->', hamiltonian.T[0] * gf[0][n] + hamiltonian.T[1] * gf[1][n]) + + #TODO this performs too many summations + energy[:, 2] = hamiltonian.const * np.einsum('nii,ni->n', gf[0] + gf[1], walkers.x) + energy[:, 3] = 0.5 * hamiltonian.m * hamiltonian.w0**2 * np.einsum('ni->n', walkers.x**2) energy[:, 3] -= 0.5 * hamiltonian.nsites * hamiltonian.w0 - if verbose: - print('e_ph pot: ', energy[0, 3]) - energy[:, 4] = -0.5 * trial.calc_phonon_laplacian_locenergy(walkers) / hamiltonian.m - if verbose: - print('e_ph kin: ', energy[0, 4]) - energy[:, 0] = np.sum(energy[:,1:], axis=1) - if verbose: - print('energy: ', energy[0,0]) return energy diff --git a/ipie/trial_wavefunction/holstein/coherent_state.py b/ipie/trial_wavefunction/holstein/coherent_state.py index 9dae69ac..9f57f620 100644 --- a/ipie/trial_wavefunction/holstein/coherent_state.py +++ b/ipie/trial_wavefunction/holstein/coherent_state.py @@ -3,46 +3,42 @@ #TODO add greens_function_coherent_state in estimators from ipie.estimators.greens_function_single_det import greens_function_single_det -from ipie.trial_wavefunction.eph_trial_base import EphTrialWavefunctionBase +from ipie.trial_wavefunction.holstein.eph_trial_base import EphTrialWavefunctionBase +from ipie.utils.backend import arraylib as xp #TODO greensfunctions are in estimators -class CoherentState(EphTrialWavefunctionBase): +class CoherentStateTrial(EphTrialWavefunctionBase): """""" - def __init__(self, wavefunction, num_elec, num_basis, verbose=False): + def __init__(self, wavefunction, hamiltonian, num_elec, num_basis, verbose=False): super().__init__(wavefunction, num_elec, num_basis, verbose=verbose) self.num_elec = num_elec - self.nalpha, self.nbeta = self.num_elec + self.nup, self.ndown = self.num_elec self.w0 = hamiltonian.w0 self.m = hamiltonian.m self.nsites = hamiltonian.nsites - self.psia = wavefunction[:self.nalpha, :] - self.psia = self.psia[0] - #NOTE 1e hack - self.psib = wavefunction[self.nalpha:self.nalpha+self.nbeta, :] - self.beta_shift = np.squeeze(wavefunction[-1, :]) - - - def build(self, walkers) -> None: - walkers.ph_ovlp = np.zeros(walkers.nwalkers) - walkers.el_ovlp = np.zeros(walkers.nwalkers) - walkers.total_ovlp = np.zeros(self.nsites) + self.beta_shift = np.squeeze(wavefunction[:, 0]) + self.psia = wavefunction[:, 1:self.nup+1] + self.psib = wavefunction[:, self.nup+1:self.nup+self.ndown+1] def calculate_energy(self, system, hamiltonian): #TODO variational_energy_coherent_state in ipie.estimators.local_energy ... def calc_overlap(self, walkers) -> np.ndarray: - #TODO this will be a concoction of phonon and electronic overlap - + _ = self.calc_phonon_overlap(walkers) + _ = self.calc_electronic_overlap(walkers) + walkers.total_ovlp = np.einsum('n,n->n', walkers.el_ovlp, walkers.ph_ovlp) + walkers.ovlp = walkers.total_ovlp + return walkers.ovlp def calc_phonon_overlap(self, walkers) -> np.ndarray: - walker.ph_ov = np.exp(-(self.m * self.w0 / 2) * (walkers.x - self.beta_shift)**2) - walker.ph_ov = np.prod(ph_ov, axis=1) - return ph_ov + ph_ovlp = np.exp(-(self.m * self.w0 / 2) * (walkers.x - self.beta_shift)**2) + walkers.ph_ovlp = np.prod(ph_ovlp, axis=1) + return walkers.ph_ovlp def calc_phonon_gradient(self, walkers) -> np.ndarray: - grad = np.einsum('ni,n->ni', (walkers.x - self.beta_shift), ovlp) + grad = walkers.x - self.beta_shift grad *= -self.m * self.w0 return grad @@ -59,10 +55,35 @@ def calc_phonon_laplacian_locenergy(self, walkers) -> np.ndarray: return self.calc_phonon_laplacian(walkers) def calc_electronic_overlap(self, walkers) -> np.ndarray: - walkers.el_ovlp[ip] = np.einsum('i,nie->n', self.psia[perm].conj(), walkers.phia) - if self.nbeta > 0: - pass +# walkers.el_ovlp = np.einsum('i,nie->n', self.psia[perm].conj(), walkers.phia) +# if self.nbeta > 0: +# pass + ovlp_a = xp.einsum("wmi,mj->wij", walkers.phia, self.psia.conj(), optimize=True) + sign_a, log_ovlp_a = xp.linalg.slogdet(ovlp_a) + + if self.ndown > 0: + ovlp_b = xp.einsum("wmi,mj->wij", walkers.phib, self.psib.conj(), optimize=True) + sign_b, log_ovlp_b = xp.linalg.slogdet(ovlp_b) + ot = sign_a * sign_b * xp.exp(log_ovlp_a + log_ovlp_b - walkers.log_shift) + else: + ot = sign_a * xp.exp(log_ovlp_a - walkers.log_shift) + + walkers.el_ovlp = ot + return walkers.el_ovlp def calc_greens_function(self, walkers) -> np.ndarray: +# overlap_inv = 1 / np.einsum('i,nie->n', self.psia.conj(), walkers.phia) +# greensfct = np.einsum('nie,n,j->nji', walkers.phia, overlap_inv, self.psia[perm].conj()) +# walkers.Ga = np.zeros((walkers.nwalkers, self.nsites, self.nsites), dtype=np.complex128) +# walkers.Gb = np.zeros_like(walkers.Ga) + + inv_Oa = xp.linalg.inv(xp.einsum('ie,nif->nef', self.psia, walkers.phia.conj())) + walkers.Ga = xp.einsum('nie,nef,jf->nji', walkers.phia, inv_Oa, self.psia.conj()) + + if self.ndown > 0: + inv_Ob = xp.linalg.inv(xp.einsum('ie,nif->nef', self.psib, walkers.phib.conj())) + walkers.Gb = xp.einsum('nie,nef,jf->nji', walkers.phib, inv_Ob, self.psib.conj()) + + return [walkers.Ga, walkers.Gb] diff --git a/ipie/trial_wavefunction/holstein/toyozawa.py b/ipie/trial_wavefunction/holstein/toyozawa.py index 2bab7dcc..a7364522 100644 --- a/ipie/trial_wavefunction/holstein/toyozawa.py +++ b/ipie/trial_wavefunction/holstein/toyozawa.py @@ -7,7 +7,10 @@ from ipie.trial_wavefunction.holstein.eph_trial_base import EphTrialWavefunctionBase from ipie.trial_wavefunction.holstein.coherent_state import CoherentStateTrial -verbose=False +from ipie.utils.backend import arraylib as xp +from ipie.propagation.overlap import calc_overlap_single_det_uhf +from ipie.estimators.greens_function_single_det import greens_function_single_det +from ipie.config import CommType, config, MPI #TODO greensfunctions are in estimators @@ -24,6 +27,7 @@ class ToyozawaTrial(CoherentStateTrial): def __init__(self, wavefunction, hamiltonian, num_elec, num_basis, verbose=False): super().__init__(wavefunction, hamiltonian, num_elec, num_basis, verbose=verbose) self.perms = list(circ_perm([i for i in range(self.nbasis)])) +# self.perms = [self.perms[0]] self.nperms = len(self.perms) def calculate_energy(self, system, hamiltonian): @@ -46,8 +50,6 @@ def calc_phonon_overlap(self, walkers) -> np.ndarray: ph_ov = np.exp(-(self.m * self.w0 / 2) * (walkers.x - self.beta_shift[perm])**2) walkers.ph_ovlp[:, ip] = np.prod(ph_ov, axis=1) ph_ovlp = np.sum(walkers.ph_ovlp, axis=1) - if verbose: - print('ph_ovlp: ', ph_ovlp[0]) return ph_ovlp def calc_phonon_gradient(self, walkers) -> np.ndarray: @@ -79,22 +81,51 @@ def calc_phonon_laplacian_locenergy(self, walkers) -> np.ndarray: def calc_electronic_overlap(self, walkers) -> np.ndarray: """""" for ip,perm in enumerate(self.perms): - walkers.el_ovlp[:, ip] = np.einsum('i,nie->n', self.psia[perm].conj(), walkers.phia) #this is single electron - if self.nbeta > 0: - pass #TODO -> adjust ovlps shape +# print('psia shape: ', self.psia.shape) +# print('walkers shape: ', walkers.phia.shape) + ovlp_a = xp.einsum("wmi,mj->wij", walkers.phia, self.psia[perm, :].conj(), optimize=True) + sign_a, log_ovlp_a = xp.linalg.slogdet(ovlp_a) + + if self.ndown > 0: + ovlp_b = xp.einsum("wmi,mj->wij", walkers.phib, self.psib[perm, :].conj(), optimize=True) + sign_b, log_ovlp_b = xp.linalg.slogdet(ovlp_b) + ot = sign_a * sign_b * xp.exp(log_ovlp_a + log_ovlp_b - walkers.log_shift) + else: + ot = sign_a * xp.exp(log_ovlp_a - walkers.log_shift) + + walkers.el_ovlp[:, ip] = ot + el_ovlp = np.sum(walkers.el_ovlp, axis=1) #NOTE this was ph before?? - if verbose: - print('el_ovlp: ', el_ovlp[0]) + +# if verbose: +# print('el_ovlp: ', el_ovlp[0]) + return el_ovlp - def calc_greens_function(self, walkers) -> np.ndarray: + def calc_greens_function(self, walkers, build_full=True) -> np.ndarray: """""" - greensfct = np.zeros((walkers.nwalkers, self.nsites, self.nsites), dtype=np.complex128) - for ovlp, perm in zip(walkers.total_ovlp.T, self.perms): - overlap_inv = 1 / np.einsum('i,nie->n', self.psia[perm].conj(), walkers.phia) #NOTE psi currently hacked - greensfct += np.einsum('nie,n,j,n->nji', walkers.phia, overlap_inv, - self.psia[perm].conj(), ovlp) - greensfct = np.einsum('nij,n->nij', greensfct, 1 / np.sum(walkers.total_ovlp, axis=1)) #these sums can be replaced by walkers.ovlp calls - return greensfct +# greensfct = np.zeros((walkers.nwalkers, self.nsites, self.nsites), dtype=np.complex128) + walkers.Ga = np.zeros((walkers.nwalkers, self.nsites, self.nsites), dtype=np.complex128) + walkers.Gb = np.zeros_like(walkers.Ga) + + for ovlp, perm in zip(walkers.total_ovlp.T, self.perms): + #TODO adjust this by just calling gab from exisiting extimators rubric TODO + #overlap_inv = 1 / np.einsum('i,nie->n', self.psia[perm].conj(), walkers.phia) #NOTE psi currently hacked + #greensfct += np.einsum('nie,n,j,n->nji', walkers.phia, overlap_inv, + # self.psia[perm].conj(), ovlp) + + inv_Oa = xp.linalg.inv(xp.einsum('ie,nif->nef', self.psia[perm,:], walkers.phia.conj())) + walkers.Ga += xp.einsum('nie,nef,jf,n->nji', walkers.phia, inv_Oa, self.psia[perm].conj(), ovlp) + + if self.ndown > 0: + inv_Ob = xp.linalg.inv(xp.einsum('ie,nif->nef', self.psib[perm,:], walkers.phib.conj())) + walkers.Gb += xp.einsum('nie,nef,jf,n->nji', walkers.phib, inv_Ob, self.psib[perm].conj(), ovlp) + #greensfct += greens_function_single_det(walkers, self, build_full=build_full) * ovlp + +# greensfct = np.einsum('nij,n->nij', greensfct, 1 / np.sum(walkers.total_ovlp, axis=1)) #these sums can be replaced by walkers.ovlp calls + walkers.Ga = np.einsum('nij,n->nij', walkers.Ga, 1 / np.sum(walkers.total_ovlp, axis=1)) + walkers.Gb = np.einsum('nij,n->nij', walkers.Gb, 1 / np.sum(walkers.total_ovlp, axis=1)) + + return [walkers.Ga, walkers.Gb] diff --git a/ipie/walkers/eph_walkers.py b/ipie/walkers/eph_walkers.py index 377a9af3..e622531f 100644 --- a/ipie/walkers/eph_walkers.py +++ b/ipie/walkers/eph_walkers.py @@ -29,23 +29,23 @@ def __init__( self.weight = numpy.ones(self.nwalkers, dtype=numpy.complex128) + self.x = xp.array( + [initial_walker[:,0].copy() for iw in range(self.nwalkers)], + dtype=xp.complex128 + ) + self.x = numpy.squeeze(self.x) + #TODO is there a reason we dont use numpy tile for these? self.phia = xp.array( - [initial_walker[:, : self.nup].copy() for iw in range(self.nwalkers)], + [initial_walker[:, 1 : self.nup+1].copy() for iw in range(self.nwalkers)], dtype=xp.complex128, ) -# self.phia = numpy.squeeze(self.phia) #NOTE: 1e hack to work with 1e overlaps self.phib = xp.array( - [initial_walker[:, self.nup:self.nup+self.ndown].copy() for iw in range(self.nwalkers)], + [initial_walker[:, self.nup+1 : self.nup+self.ndown+1].copy() for iw in range(self.nwalkers)], dtype=xp.complex128, ) - self.x = xp.array( - [initial_walker[:, self.nup+self.ndown:].copy() for iw in range(self.nwalkers)], - dtype=xp.complex128 - ) - self.x = numpy.squeeze(self.x) - + self.buff_names += ["phia", "phib", "x"] self.buff_size = round(self.set_buff_size_single_walker() / float(self.nwalkers)) @@ -110,8 +110,8 @@ def reortho(self): self.log_detR[iw] += xp.log(detR[iw]) self.detR[iw] = detR[iw] - self.el_ovlp[iw, :] = self.el_ovlp[iw, :] / detR[iw] - self.total_ovlp[iw, :] = self.total_ovlp[iw, :] / detR[iw] + self.el_ovlp[iw, ...] = self.el_ovlp[iw, ...] / detR[iw] + self.total_ovlp[iw, ...] = self.total_ovlp[iw, ...] / detR[iw] self.ovlp[iw] = self.ovlp[iw] / detR[iw] synchronize() From e7507f75c71c33ecfc45887916c880173b5f4ae8 Mon Sep 17 00:00:00 2001 From: m-baumgarten Date: Sun, 3 Mar 2024 16:29:23 -0500 Subject: [PATCH 05/61] Some einsum cleanup --- ipie/estimators/local_energy_holstein.py | 20 ++++---- .../holstein/coherent_state.py | 42 ++++++++++----- ipie/trial_wavefunction/holstein/toyozawa.py | 3 +- ipie/walkers/eph_walkers.py | 51 +++++++------------ 4 files changed, 58 insertions(+), 58 deletions(-) diff --git a/ipie/estimators/local_energy_holstein.py b/ipie/estimators/local_energy_holstein.py index 4b46a41d..03b0e319 100644 --- a/ipie/estimators/local_energy_holstein.py +++ b/ipie/estimators/local_energy_holstein.py @@ -13,20 +13,22 @@ def local_energy_holstein( walkers: EphWalkers, trial: EphTrialWavefunctionBase ): - """""" + """Computes the local energy for the Holstein model.""" energy = xp.zeros((walkers.nwalkers, 5), dtype=xp.complex128) - #get greens_function from estimators gf = trial.calc_greens_function(walkers) - #TODO make this nicer - for n in range(walkers.nwalkers): - energy[n, 1] = np.einsum('ij->', hamiltonian.T[0] * gf[0][n] + hamiltonian.T[1] * gf[1][n]) - + energy[:, 1] = np.sum(hamiltonian.T[0] * gf[0], axis=(1,2)) + if system.ndown > 0: + energy[:, 1] += np.sum(hamiltonian.T[1] * gf[1], axis=(1,2)) + #TODO this performs too many summations - energy[:, 2] = hamiltonian.const * np.einsum('nii,ni->n', gf[0] + gf[1], walkers.x) - - energy[:, 3] = 0.5 * hamiltonian.m * hamiltonian.w0**2 * np.einsum('ni->n', walkers.x**2) + energy[:, 2] = np.sum(np.diagonal(gf[0], axis1=1, axis2=2) * walkers.x, axis=1) + if system.ndown > 0: + energy[:, 2] += np.sum(np.diagonal(gf[1], axis1=1, axis2=2) * walkers.x, axis=1) + energy[:, 2] *= hamiltonian.const + + energy[:, 3] = 0.5 * hamiltonian.m * hamiltonian.w0**2 * np.sum(walkers.x**2, axis=1) energy[:, 3] -= 0.5 * hamiltonian.nsites * hamiltonian.w0 energy[:, 4] = -0.5 * trial.calc_phonon_laplacian_locenergy(walkers) / hamiltonian.m energy[:, 0] = np.sum(energy[:,1:], axis=1) diff --git a/ipie/trial_wavefunction/holstein/coherent_state.py b/ipie/trial_wavefunction/holstein/coherent_state.py index 9f57f620..b3066e61 100644 --- a/ipie/trial_wavefunction/holstein/coherent_state.py +++ b/ipie/trial_wavefunction/holstein/coherent_state.py @@ -28,7 +28,7 @@ def calculate_energy(self, system, hamiltonian): def calc_overlap(self, walkers) -> np.ndarray: _ = self.calc_phonon_overlap(walkers) _ = self.calc_electronic_overlap(walkers) - walkers.total_ovlp = np.einsum('n,n->n', walkers.el_ovlp, walkers.ph_ovlp) + walkers.total_ovlp = walkers.el_ovlp * walkers.ph_ovlp #np.einsum('n,n->n', walkers.el_ovlp, walkers.ph_ovlp) walkers.ovlp = walkers.total_ovlp return walkers.ovlp @@ -55,9 +55,18 @@ def calc_phonon_laplacian_locenergy(self, walkers) -> np.ndarray: return self.calc_phonon_laplacian(walkers) def calc_electronic_overlap(self, walkers) -> np.ndarray: -# walkers.el_ovlp = np.einsum('i,nie->n', self.psia[perm].conj(), walkers.phia) -# if self.nbeta > 0: -# pass + """Computes electronic overlap. + + Parameters + ---------- + walkers : class + EphWalkers class object + + Returns + ------- + walker.el_ovlp : np.ndarray + Electronic overlap + """ ovlp_a = xp.einsum("wmi,mj->wij", walkers.phia, self.psia.conj(), optimize=True) sign_a, log_ovlp_a = xp.linalg.slogdet(ovlp_a) @@ -73,17 +82,24 @@ def calc_electronic_overlap(self, walkers) -> np.ndarray: return walkers.el_ovlp def calc_greens_function(self, walkers) -> np.ndarray: -# overlap_inv = 1 / np.einsum('i,nie->n', self.psia.conj(), walkers.phia) -# greensfct = np.einsum('nie,n,j->nji', walkers.phia, overlap_inv, self.psia[perm].conj()) -# walkers.Ga = np.zeros((walkers.nwalkers, self.nsites, self.nsites), dtype=np.complex128) -# walkers.Gb = np.zeros_like(walkers.Ga) - - inv_Oa = xp.linalg.inv(xp.einsum('ie,nif->nef', self.psia, walkers.phia.conj())) - walkers.Ga = xp.einsum('nie,nef,jf->nji', walkers.phia, inv_Oa, self.psia.conj()) + """Computes Greens function. + + Parameters + ---------- + walkers : class + EphWalkers class object + + Returns + ------- + walkers.G : list + Greens function for each spin space + """ + inv_Oa = xp.linalg.inv(xp.einsum('ie,nif->nef', self.psia, walkers.phia.conj()), optimize=True) + walkers.Ga = xp.einsum('nie,nef,jf->nji', walkers.phia, inv_Oa, self.psia.conj(), optimize=True) if self.ndown > 0: - inv_Ob = xp.linalg.inv(xp.einsum('ie,nif->nef', self.psib, walkers.phib.conj())) - walkers.Gb = xp.einsum('nie,nef,jf->nji', walkers.phib, inv_Ob, self.psib.conj()) + inv_Ob = xp.linalg.inv(xp.einsum('ie,nif->nef', self.psib, walkers.phib.conj()), optimize=True) + walkers.Gb = xp.einsum('nie,nef,jf->nji', walkers.phib, inv_Ob, self.psib.conj(), optimize=True) return [walkers.Ga, walkers.Gb] diff --git a/ipie/trial_wavefunction/holstein/toyozawa.py b/ipie/trial_wavefunction/holstein/toyozawa.py index a7364522..384f5f8f 100644 --- a/ipie/trial_wavefunction/holstein/toyozawa.py +++ b/ipie/trial_wavefunction/holstein/toyozawa.py @@ -27,7 +27,6 @@ class ToyozawaTrial(CoherentStateTrial): def __init__(self, wavefunction, hamiltonian, num_elec, num_basis, verbose=False): super().__init__(wavefunction, hamiltonian, num_elec, num_basis, verbose=verbose) self.perms = list(circ_perm([i for i in range(self.nbasis)])) -# self.perms = [self.perms[0]] self.nperms = len(self.perms) def calculate_energy(self, system, hamiltonian): @@ -38,7 +37,7 @@ def calc_overlap(self, walkers) -> np.ndarray: #TODO this will be a concoction of phonon and electronic overlap _ = self.calc_phonon_overlap(walkers) _ = self.calc_electronic_overlap(walkers) - walkers.total_ovlp = np.einsum('np,np->np', walkers.el_ovlp, walkers.ph_ovlp) + walkers.total_ovlp = walkers.el_ovlp * walkers.ph_ovlp walkers.ovlp = np.sum(walkers.total_ovlp, axis=1) return walkers.ovlp diff --git a/ipie/walkers/eph_walkers.py b/ipie/walkers/eph_walkers.py index e622531f..e9e77f70 100644 --- a/ipie/walkers/eph_walkers.py +++ b/ipie/walkers/eph_walkers.py @@ -7,7 +7,8 @@ from ipie.trial_wavefunction.holstein.toyozawa import ToyozawaTrial class EphWalkers(BaseWalkers): - """""" + """Class tailored to el-ph models where keeping track of phonon overlaps is + required.""" def __init__( self, initial_walker: numpy.ndarray, @@ -15,14 +16,12 @@ def __init__( ndown: int, nbasis: int, nwalkers: int, - mpi_handler, verbose: bool = False ): self.nup = nup self.ndown = ndown self.nbasis = nbasis - self.mpi_handler = mpi_handler super().__init__(nwalkers, verbose=verbose) @@ -35,7 +34,6 @@ def __init__( ) self.x = numpy.squeeze(self.x) - #TODO is there a reason we dont use numpy tile for these? self.phia = xp.array( [initial_walker[:, 1 : self.nup+1].copy() for iw in range(self.nwalkers)], dtype=xp.complex128, @@ -52,10 +50,14 @@ def __init__( self.walker_buffer = numpy.zeros(self.buff_size, dtype=numpy.complex128) def build(self, trial): - """NOTE: total_ovlp is the same as ovlp for coherent state trial, and just - serves the purpose of not recomputing overlaps for each permutation - but passing it to the pop_control and adjsuting it accordingly for the - Toyozawa trial.""" + """Allocates memory for computation of overlaps throughout the + simulation. + + Parameters + ---------- + trial : class + Trial wavefunction object. + """ if isinstance(trial, ToyozawaTrial): shape = (self.nwalkers, trial.nperms) @@ -66,7 +68,7 @@ def build(self, trial): self.el_ovlp = numpy.zeros(shape, dtype=numpy.complex128) self.total_ovlp = numpy.zeros(shape, dtype=numpy.complex128) - self.buff_names += ['total_ovlp'] #, 'el_ovlp', 'total_ovlp'] #not really necessary to bring 'el_ovlp', 'total_ovlp' along if computing overlap after normalization anyways. + self.buff_names += ['total_ovlp'] self.buff_size = round(self.set_buff_size_single_walker() / float(self.nwalkers)) self.walker_buffer = numpy.zeros(self.buff_size, dtype=numpy.complex128) @@ -74,9 +76,11 @@ def cast_to_cupy(self, verbose=False): cast_to_device(self, verbose) def reortho(self): - """reorthogonalise walkers. + """reorthogonalise walkers. This function is mostly from BaseWalkers, + with the exception that it adjusts all overlaps, possibly of numerous + coherent states. - parameters + Parameters ---------- """ if config.get_option("use_gpu"): @@ -117,26 +121,5 @@ def reortho(self): synchronize() return detR - def reortho_batched(self): - """reorthogonalise walkers. - - parameters - ---------- - """ - assert config.get_option("use_gpu") - (self.phia, Rup) = qr(self.phia, mode=qr_mode) - Rup_diag = xp.einsum("wii->wi", Rup) - log_det = xp.einsum("wi->w", xp.log(abs(Rup_diag))) - - if self.ndown > 0: - (self.phib, Rdn) = qr(self.phib, mode=qr_mode) - Rdn_diag = xp.einsum("wii->wi", Rdn) - log_det += xp.einsum("wi->w", xp.log(abs(Rdn_diag))) - self.detR = xp.exp(log_det - self.detR_shift) - self.ovlp = self.ovlp / self.detR - - synchronize() - - return self.detR - - + def reortho_batched(self): # gpu version + pass From bf49a314445bfb8adf01b56e44ff077300112004 Mon Sep 17 00:00:00 2001 From: m-baumgarten Date: Mon, 4 Mar 2024 11:35:07 -0500 Subject: [PATCH 06/61] added variationcal code to ipie and example run file --- examples/13-1d_holstein/run_afqmc.py | 95 ++++++++++ .../variational/coherent_state_variational.py | 173 ++++++++++++++++++ .../variational/toyozawa_variational.py | 166 +++++++++++++++++ 3 files changed, 434 insertions(+) create mode 100644 examples/13-1d_holstein/run_afqmc.py create mode 100644 ipie/trial_wavefunction/holstein/variational/coherent_state_variational.py create mode 100644 ipie/trial_wavefunction/holstein/variational/toyozawa_variational.py diff --git a/examples/13-1d_holstein/run_afqmc.py b/examples/13-1d_holstein/run_afqmc.py new file mode 100644 index 00000000..64bf5d54 --- /dev/null +++ b/examples/13-1d_holstein/run_afqmc.py @@ -0,0 +1,95 @@ +import numpy as np + +from mpi4py import MPI + +from ipie.qmc.afqmc import AFQMC +from ipie.estimators.energy import EnergyEstimator +from ipie.systems import Generic +from ipie.hamiltonians.elph.holstein import HolsteinModel +from ipie.trial_wavefunction.holstein.toyozawa import ToyozawaTrial +from ipie.trial_wavefunction.holstein.variational.toyozawa_variational import variational_trial_toyozawa +from ipie.walkers.eph_walkers import EphWalkers +from ipie.propagation.holstein import HolsteinPropagatorImportance +from ipie.qmc.options import QMCParams + +#System Parameters +nup = 2 +ndown = 2 +nelec = [nup, ndown] + +#Hamiltonian Parameters +g = 2. +t = 1. +w0 = 4. +nsites = 4 +pbc = True + +#Walker Parameters & Setup +comm = MPI.COMM_WORLD +nwalkers = 1000 // comm.size + +#Setup initial guess for variational optimization +initial_electron = np.random.random((nsites, nup + ndown)) +initial_phonons = np.ones(nsites) * 0.1 + +#System and Hamiltonian setup +system = Generic(nelec) +ham = HolsteinModel(g=g, t=t, w0=w0, nsites=nsites, pbc=pbc) +ham.build() + +#Variational procedure +etrial, beta_shift, el_trial = variational_trial_toyozawa( + initial_phonons, initial_electron, ham, system +) + +wavefunction = np.column_stack([beta_shift, el_trial]) + +#Setup trial +trial = ToyozawaTrial( + wavefunction=wavefunction, + hamiltonian=ham, + num_elec=[nup, ndown], + num_basis=nsites +) +trial.set_etrial(etrial) + +#Setup walkers +walkers = EphWalkers( + initial_walker=wavefunction, + nup=nup, + ndown=ndown, + nbasis=nsites, + nwalkers=nwalkers +) +walkers.build(trial) + +timestep = 0.01 +propagator = HolsteinPropagatorImportance(timestep) +propagator.build(ham) + +num_steps_per_block = 10 +num_blocks = 10000 +add_est = { + "energy": EnergyEstimator( + system=system, ham=ham, trial=trial + ) +} + +stabilize_freq = 5 +pop_control_freq = 5 +seed = 125 +params = QMCParams( + num_walkers=nwalkers, + total_num_walkers=nwalkers * comm.size, + num_blocks=num_blocks, + num_steps_per_block=num_steps_per_block, + timestep=timestep, + num_stblz=stabilize_freq, + pop_control_freq=pop_control_freq, + rng_seed=seed, +) + +ephqmc = AFQMC(system, ham, trial, walkers, propagator, params) +trial.calc_overlap(walkers) #Sets ovlp, ph and el overlaps +ephqmc.run(additional_estimators=add_est, verbose=False) + diff --git a/ipie/trial_wavefunction/holstein/variational/coherent_state_variational.py b/ipie/trial_wavefunction/holstein/variational/coherent_state_variational.py new file mode 100644 index 00000000..6388c912 --- /dev/null +++ b/ipie/trial_wavefunction/holstein/variational/coherent_state_variational.py @@ -0,0 +1,173 @@ +import numpy as np +from scipy.optimize import basinhopping + +from ipie.legacy.trial_wavefunction.harmonic_oscillator import HarmonicOscillator +from ipie.systems.generic import Generic +from ipie.hamiltonians.elph.holstein import HolsteinModel + +import jax + +def local_energy( + X: np.ndarray, + G: np.ndarray, + m: float, + w: float, + g: float, + nsites: int, + Lap: np.ndarray, + T: np.ndarray, + nup: int, + ndown: int + ) -> np.ndarray: + + kinetic_contrib = jax.numpy.einsum('ij->', T[0] * G[0]) + if ndown > 0: + kinetic_contrib += jax.numpy.einsum('ij->', T[1] * G[1]) + + rho = G[0].diagonal() + G[1].diagonal() + el_ph_contrib = -g * np.sqrt(2 * m * w) * np.sum(rho * X) + + phonon_contrib = m * w**2 * np.sum(X * X) / 2 + phonon_contrib += -0.5 * np.sum(Lap) / m - 0.5 * w * nsites + + local_energy = kinetic_contrib + el_ph_contrib + phonon_contrib + return local_energy + +def gab(A: np.ndarray, B: np.ndarray): + inv_O = jax.numpy.linalg.inv((A.conj().T).dot(B)) + GAB = B.dot(inv_O.dot(A.conj().T)) + return GAB + +def get_T(nsites: int, t: float, pbc: bool, ndown: int) -> np.ndarray: + if nsites == 1: + return np.array([0.]) + + neighbor_map = np.zeros(shape=(nsites,nsites)) + for i in range(nsites): + neighbors = np.array([(i-1)%nsites, (i+1)%nsites]) + neighbor_map[i, neighbors] = 1. + + if not pbc: + neighbor_map[-1,0] = neighbor_map[0,-1] = 0. + T = -t * neighbor_map + + if ndown > 0: + T = [T.copy(), T.copy()] + else: + T = [T,] + return T + +def objective_function(x: np.ndarray, nbasis: int, T: np.ndarray, + g: float, m: float, w0: float, nup: int, ndown: int): + shift = x[0:nbasis].copy() + shift = shift.astype(np.float64) + + c0a = x[nbasis : (nup+1)*nbasis].copy() + c0a = jax.numpy.reshape(c0a, (nbasis, nup)) + c0a = c0a.astype(np.float64) + Ga = gab(c0a, c0a) + + if ndown > 0: + c0b = x[(nup+1)*nbasis : ].copy() + c0b = jax.numpy.reshape(c0b, (nbasis, ndown)) + c0b = c0b.astype(np.float64) + Gb = gab(c0b, c0b) + else: + Gb = jax.numpy.zeros_like(Ga, dtype=np.float64) + + G = [Ga, Gb] + + phi = HarmonicOscillator(m, w0, order=0, shift=shift) + Lap = phi.laplacian(shift) + + etot = local_energy(shift, G, m, w0, g, nbasis, Lap, T, nup, ndown) + return etot.real + + +def gradient(x: np.ndarray, nbasis: int, T: np.ndarray, + g: float, m: float, w0: float, nup: int, ndown: int): + grad = np.array( + jax.grad(objective_function)( + x, nbasis, T, g, m, w0, nup, ndown + ), + dtype=np.float64 + ) + return grad + +def func(x: np.ndarray, nbasis: int, T: np.ndarray, + g: float, m: float, w0: float, nup: int, ndown: int): + f = objective_function(x, nbasis, T, g, m, w0, nup, ndown) + df = gradient(x, nbasis, T, g, m, w0, nup, ndown) + return f, df + +def print_fun(x: np.ndarray, f: float, accepted: bool): + print("at minimum %.4f accepted %d" % (f, int(accepted))) + +def variational_trial(init_phonons: np.ndarray, init_electron: np.ndarray, hamiltonian, system): + + init_phonons = init_phonons.astype(np.float64) + init_electron = init_electron.astype(np.float64).ravel() + + x = np.hstack([init_phonons, init_electron]) + + #TODO this could definitely be hamiltonian.T +# T = get_T(hamiltonian.nsites, hamiltonian.t, hamiltonian.pbc, system.ndown) + T = hamiltonian.T + + maxiter = 500 + minimizer_kwargs = { + "jac": True, + "args": (hamiltonian.nsites, T, hamiltonian.g, hamiltonian.m, hamiltonian.w0, system.nup, system.ndown), + "options": { + "gtol": 1e-10, + "eps": 1e-10, + "maxiter": maxiter, + "disp": False, + }, + } + + res = basinhopping( + func, + x, + minimizer_kwargs=minimizer_kwargs, + callback=print_fun, + niter=maxiter, + niter_success=3, + ) + + etrial = res.fun + + beta_shift = res.x[:hamiltonian.nsites] + if system.ndown > 0: + psia = res.x[hamiltonian.nsites : hamiltonian.nsites*(system.nup+1)] + psia = psia.reshape((hamiltonian.nsites, system.nup)) + psib = res.x[hamiltonian.nsites*(system.nup+1) : ] + psib = psib.reshape((hamiltonian.nsites, system.ndown)) + psi = np.column_stack([psia, psib]) + else: + psia = res.x[hamiltonian.nsites:].reshape((hamiltonian.nsites, system.nup)) + psi = psia + + return etrial, beta_shift, psi + +if __name__ == '__main__': + + nup = 1 + ndown = 1 + system = Generic([nup, ndown]) + + g = 1. + t = 1. + w0 = 1. + nsites = 3 + pbc = True + + ham = HolsteinModel(g=g, t=t, w0=w0, nsites=nsites, pbc=pbc) + + init_phonons = np.array([0.1, 0.1, 0.1], dtype=np.float64) +# init_electron = np.array([1.,1.], dtype=np.float64) +# init_electron /= np.linalg.norm(init_electron) + init_electron = np.array([[1,0,0],[0,1,0]]) + + print(variational_trial(init_phonons, init_electron, ham, system)) + diff --git a/ipie/trial_wavefunction/holstein/variational/toyozawa_variational.py b/ipie/trial_wavefunction/holstein/variational/toyozawa_variational.py new file mode 100644 index 00000000..e75d239d --- /dev/null +++ b/ipie/trial_wavefunction/holstein/variational/toyozawa_variational.py @@ -0,0 +1,166 @@ +import numpy as np +from scipy.optimize import minimize, basinhopping +from ipie.systems import Generic +from ipie.hamiltonians.elph.holstein import HolsteinModel + +from jax.config import config +config.update("jax_enable_x64", True) +import jax +import jax.numpy as npj + +def gab(A,B): + inv_O = npj.linalg.inv((A.conj().T).dot(B)) + GAB = B.dot(inv_O.dot(A.conj().T)) + return GAB + +def gradient_toyozawa_mo(x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted): + grad = np.array(jax.grad(objective_function_toyozawa_mo)(x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted)) + return grad + +def objective_function_toyozawa_mo(x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted, zero_th=1e-12): + nbasis = int(round(nbasis)) + nup = int(round(nup)) + ndown = int(round(ndown)) + + shift0 = x[:nbasis] + + nbsf = nbasis + nocca = nup + noccb = ndown + nvira = nbasis - nocca + nvirb = nbasis - noccb + + psi0a = npj.array(x[nbsf:nbsf+nbsf*nocca],dtype=npj.float64) + psi0a = psi0a.reshape((nocca, nbsf)).T + + if (noccb>0): + psi0b = npj.array(x[nbsf+nbsf*nocca : ], dtype=npj.float64) + psi0b = psi0b.reshape((noccb, nbsf)).T + + nperms = len(perms) + + num_energy = 0. + denom = 0.0 + + beta0 = shift0 * npj.sqrt(m * w0 /2.0) + + for i,perm in enumerate(perms): + psia_j = psi0a[perm, :] + beta_j = beta0[npj.array(perm)] + + if (noccb > 0): + psib_j = psi0b[perm, :] + overlap = npj.linalg.det(psi0a.T.dot(psia_j)) * npj.linalg.det(psi0b.T.dot(psib_j)) * npj.prod(npj.exp(-0.5 * (beta0**2 + beta_j**2) + beta0*beta_j)) + else: + overlap = npj.linalg.det(psi0a.T.dot(psia_j)) * npj.prod(npj.exp (- 0.5 * (beta0**2 + beta_j**2) + beta0*beta_j)) + + if overlap < zero_th: + continue + + Ga_j = gab(psi0a, psia_j) + if (noccb > 0): + Gb_j = gab(psi0b, psib_j) + else: + Gb_j = npj.zeros_like(Ga_j) + + G_j = [Ga_j, Gb_j] + + rho = G_j[0].diagonal() + G_j[1].diagonal() + ke = npj.sum(T[0] * G_j[0] + T[1] * G_j[1]) + pe = U * npj.dot(G_j[0].diagonal(), G_j[1].diagonal()) + e_ph = w0 * npj.sum(beta0 * beta_j) + e_eph = - g * npj.dot(rho, beta0 + beta_j) + + num_energy += (ke + pe + e_ph + e_eph) * overlap # 2.0 comes from hermiticity + denom += overlap + + etot = num_energy / denom + return etot.real + +def circ_perm(lst): + cpy = lst[:] + yield cpy + for i in range(len(lst) - 1): + cpy = cpy[1:] + [cpy[0]] + yield cpy + +def func(x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted): + f = objective_function_mo(x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted) + df = gradient_mo(x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted) + return f, df + +def print_fun(x: np.ndarray, f: float, accepted: bool): + print("at minimum %.4f accepted %d" % (f, int(accepted))) + +def func_toyo(x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted): + f = objective_function_toyozawa_mo(x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted) + df = gradient_toyozawa_mo(x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted) + return f, df + +def variational_trial_toyozawa(shift_init: np.ndarray, electron_init: np.ndarray, hamiltonian, system): + psi = electron_init.T.real.ravel() + shift = shift_init.real + + perms = list(circ_perm([i for i in range(hamiltonian.nsites)])) + + x = np.zeros((system.nup + system.ndown + 1) * hamiltonian.nsites) + x[:hamiltonian.nsites] = shift.copy() + x[hamiltonian.nsites:] = psi.copy() #[:,0] + + res = minimize( + objective_function_toyozawa_mo, + x, + args = (float(hamiltonian.nsites), float(system.nup), float(system.ndown), + hamiltonian.T, 0., hamiltonian.g, hamiltonian.m, hamiltonian.w0, perms, False), + jac = gradient_toyozawa_mo, + tol = 1e-10, + method='L-BFGS-B', + options = { + 'maxls': 20, + 'iprint': 2, + 'gtol': 1e-10, + 'eps': 1e-10, + 'maxiter': 15000, + 'ftol': 1.0e-10, + 'maxcor': 1000, + 'maxfun': 15000, + 'disp':True + } + ) + + etrial = res.fun + beta_shift = res.x[:hamiltonian.nsites] + if system.ndown > 0: + psia = res.x[hamiltonian.nsites : hamiltonian.nsites*(system.nup+1)] + psia = psia.reshape((hamiltonian.nsites, system.nup)) + psib = res.x[hamiltonian.nsites*(system.nup+1) : ] + psib = psib.reshape((hamiltonian.nsites, system.ndown)) + psi = np.column_stack([psia, psib]) + else: + psia = res.x[hamiltonian.nsites:].reshape((hamiltonian.nsites, system.nup)) + psi = psia + + + return etrial, beta_shift, psi + +if __name__ == '__main__': + w0 = 1. + m = 1/w0 + g = 1. + t = 1. + nsites = 4 + pbc = True + + nup = 2 + ndown = 2 + + system = Generic([nup, ndown]) + + shift_init = np.array([2.78351996, 0.04490717, 0.04490717, 0.04490717]) + electron_init = np.random.random((nsites, nup + ndown)).astype(np.float32) + + hamiltonian = HolsteinModel(g=g, t=t, w0=w0, nsites=nsites, pbc=pbc) + hamiltonian.build() + + print(variational_trial_toyozawa(shift_init, electron_init, hamiltonian, system)) + From 2ef4fb8f5fc823d70572f7376378cd7a018339ef Mon Sep 17 00:00:00 2001 From: m-baumgarten Date: Mon, 4 Mar 2024 13:13:12 -0500 Subject: [PATCH 07/61] holstein propagator cleanup --- ipie/propagation/holstein.py | 69 ++++++++---------------------------- 1 file changed, 14 insertions(+), 55 deletions(-) diff --git a/ipie/propagation/holstein.py b/ipie/propagation/holstein.py index 21ddd9ab..796531fa 100644 --- a/ipie/propagation/holstein.py +++ b/ipie/propagation/holstein.py @@ -2,14 +2,11 @@ import time import scipy.linalg -from ipie.propagation.phaseless_base import PhaselessBase from ipie.hamiltonians.elph.holstein import HolsteinModel -from ipie.propagation.operations import propagate_one_body #TODO check this +from ipie.propagation.operations import propagate_one_body from ipie.utils.backend import synchronize, cast_to_device from ipie.propagation.continuous_base import PropagatorTimer -verbose=False - def construct_one_body_propagator(hamiltonian: HolsteinModel, dt: float): """""" H1 = hamiltonian.T @@ -19,11 +16,9 @@ def construct_one_body_propagator(hamiltonian: HolsteinModel, dt: float): return expH1 -class HolsteinPropagatorFree(): #for hubbard inherit form phaseless base +class HolsteinPropagatorFree(): """""" - def __init__(self, time_step, verbose=False): - #super().__init__(time_step, verbose=verbose) self.dt = time_step self.verbose = verbose self.timer = PropagatorTimer() @@ -51,7 +46,7 @@ def propagate_walkers_one_body(self, walkers): def propagate_phonons(self, walkers): start_time = time.time() - pot = 0.25 * self.m * self.w0**2 * numpy.einsum('ni->n', walkers.x**2) + pot = 0.25 * self.m * self.w0**2 * numpy.sum(walkers.x**2, axis=1) pot = numpy.real(pot) walkers.weight *= numpy.exp(-self.dt_ph * pot) @@ -59,7 +54,7 @@ def propagate_phonons(self, walkers): size=(walkers.nwalkers, self.nsites)) walkers.x = walkers.x + N - pot = 0.25 * self.m * self.w0**2 * numpy.einsum('ni->n', walkers.x**2) + pot = 0.25 * self.m * self.w0**2 * numpy.sum(walkers.x**2, axis=1) pot = numpy.real(pot) walkers.weight *= numpy.exp(-self.dt_ph * pot) @@ -74,25 +69,18 @@ def propagate_electron(self, walkers, trial): synchronize() self.timer.tgf += time.time() - start_time - walkers.phia = propagate_one_body(walkers.phia, self.expH1[0]) - if walkers.ndown > 0: - walkers.phib = propagate_one_body(walkers.phib, self.expH1[1]) - expEph = numpy.exp(self.const * walkers.x) - #NOTE wouldnt it be great to cut the additional index in propagate_one_body? + + walkers.phia = propagate_one_body(walkers.phia, self.expH1[0]) walkers.phia = numpy.einsum('ni,nie->nie', expEph, walkers.phia) -# walkers.phia = propagate_one_body(walkers.phia, numpy.diag(expEph), H1diag=True) - if walkers.ndown > 0: -# walkers.phib = propagate_one_body(walkers.phib, numpy.diag(expEph), H1diag=True) - walkers.phib = numpy.einsum('ni,nie->nie', expEph, walkers.phib) - walkers.phia = propagate_one_body(walkers.phia, self.expH1[0]) + if walkers.ndown > 0: walkers.phib = propagate_one_body(walkers.phib, self.expH1[1]) - + walkers.phib = numpy.einsum('ni,nie->nie', expEph, walkers.phib) + walkers.phib = propagate_one_body(walkers.phib, self.expH1[1]) def propagate_walkers(self, walkers, hamiltonian, trial, eshift=None): - #TODO where should I store the phonon overlaps? synchronize() start_time = time.time() ovlp = trial.calc_overlap(walkers) @@ -131,19 +119,12 @@ def __init__(self, time_step, verbose=False): super().__init__(time_step, verbose=verbose) def propagate_phonons(self, walkers, hamiltonian, trial): - r"""Propagates phonons via DMC - - Parameters - ---------- - walkers : - hamiltonian : - trial : - """ + """Propagates phonons via DMC""" start_time = time.time() #no ZPE in pot -> cancels with ZPE of etrial, wouldn't affect estimators anyways ph_ovlp_old = trial.calc_phonon_overlap(walkers) - pot = 0.5 * hamiltonian.m * hamiltonian.w0**2 * numpy.einsum('ni->n',walkers.x**2) + pot = 0.5 * hamiltonian.m * hamiltonian.w0**2 * numpy.sum(walkers.x**2, axis=1) pot -= 0.5 * trial.calc_phonon_laplacian_importance(walkers) / hamiltonian.m pot = numpy.real(pot) walkers.weight *= numpy.exp(-self.dt_ph * pot / 2) @@ -154,7 +135,7 @@ def propagate_phonons(self, walkers, hamiltonian, trial): ph_ovlp_new = trial.calc_phonon_overlap(walkers) - pot = 0.5 * hamiltonian.m * hamiltonian.w0**2 * numpy.einsum('ni->n', walkers.x**2) + pot = 0.5 * hamiltonian.m * hamiltonian.w0**2 * numpy.sum(walkers.x**2, axis=1) pot -= 0.5 * trial.calc_phonon_laplacian_importance(walkers) / hamiltonian.m pot = numpy.real(pot) walkers.weight *= numpy.exp(-self.dt_ph * pot / 2) @@ -168,18 +149,11 @@ def propagate_phonons(self, walkers, hamiltonian, trial): def propagate_walkers(self, walkers, hamiltonian, trial, eshift=None): """""" - synchronize() start_time = time.time() - if verbose: - print('before ovlp: ', walkers.phia[0,0,0], walkers.x[0,0], walkers.weight[0]) - ovlp = trial.calc_overlap(walkers).copy() - if verbose: - print('ovlp: ', ovlp[0]) - synchronize() self.timer.tgf += time.time() - start_time @@ -187,36 +161,21 @@ def propagate_walkers(self, walkers, hamiltonian, trial, eshift=None): # 2.a DMC for phonon degrees of freedom self.propagate_phonons(walkers, hamiltonian, trial) - if verbose: - print('weight in prop I: ', walkers.weight[0]) - # 2.b One-body propagation for electrons self.propagate_electron(walkers, trial) # 2.c DMC for phonon degrees of freedom self.propagate_phonons(walkers, hamiltonian, trial) - - if verbose: - print('weight in prop II: ', walkers.weight[0]) -# print('electron & phonon & weight: ', walkers.phia[0], walkers.x, walkers.weight) - - # Update weights (and later do phaseless for multi-electron) start_time = time.time() + ovlp_new = trial.calc_overlap(walkers) synchronize() self.timer.tovlp += time.time() - start_time start_time = time.time() - if verbose: - print('new and old overlaps: ', ovlp_new[0], ovlp[0]) - - self.update_weight(walkers, ovlp, ovlp_new) #walkers.previously ovlp_for_weight_update -# walkers.total_ovlp_for_weight_update = ovlp_new - - if verbose: - print('weight in prop III: ', walkers.weight[0]) + self.update_weight(walkers, ovlp, ovlp_new) synchronize() self.timer.tupdate += time.time() - start_time From f0a10a68e118ec814685b91a0482502d02230c54 Mon Sep 17 00:00:00 2001 From: m-baumgarten Date: Tue, 5 Mar 2024 00:08:17 -0500 Subject: [PATCH 08/61] moved files to adhere to addons infrastructure --- examples/13-1d_holstein/run_afqmc.py | 13 ++--- .../estimators/local_energy_holstein.py | 56 +++++++++++++++++++ .../ephqmc/hamiltonians}/generic.py | 0 .../ephqmc/hamiltonians}/holstein.py | 0 .../ephqmc}/propagation/holstein.py | 2 +- .../trial_wavefunction}/coherent_state.py | 4 +- .../trial_wavefunction}/eph_trial_base.py | 2 +- .../ephqmc/trial_wavefunction}/toyozawa.py | 37 +----------- .../variational/coherent_state_variational.py | 2 +- .../variational/toyozawa_variational.py | 2 +- .../ephqmc}/walkers/eph_walkers.py | 2 +- ipie/estimators/energy.py | 8 +-- ipie/estimators/local_energy_holstein.py | 38 ------------- ipie/propagation/propagator.py | 4 +- 14 files changed, 77 insertions(+), 93 deletions(-) create mode 100644 ipie/addons/ephqmc/estimators/local_energy_holstein.py rename ipie/{hamiltonians/elph => addons/ephqmc/hamiltonians}/generic.py (100%) rename ipie/{hamiltonians/elph => addons/ephqmc/hamiltonians}/holstein.py (100%) rename ipie/{ => addons/ephqmc}/propagation/holstein.py (98%) rename ipie/{trial_wavefunction/holstein => addons/ephqmc/trial_wavefunction}/coherent_state.py (94%) rename ipie/{trial_wavefunction/holstein => addons/ephqmc/trial_wavefunction}/eph_trial_base.py (90%) rename ipie/{trial_wavefunction/holstein => addons/ephqmc/trial_wavefunction}/toyozawa.py (70%) rename ipie/{trial_wavefunction/holstein => addons/ephqmc/trial_wavefunction}/variational/coherent_state_variational.py (98%) rename ipie/{trial_wavefunction/holstein => addons/ephqmc/trial_wavefunction}/variational/toyozawa_variational.py (98%) rename ipie/{ => addons/ephqmc}/walkers/eph_walkers.py (98%) delete mode 100644 ipie/estimators/local_energy_holstein.py diff --git a/examples/13-1d_holstein/run_afqmc.py b/examples/13-1d_holstein/run_afqmc.py index 64bf5d54..0d4ac668 100644 --- a/examples/13-1d_holstein/run_afqmc.py +++ b/examples/13-1d_holstein/run_afqmc.py @@ -1,15 +1,15 @@ import numpy as np - +np.random.seed(125) from mpi4py import MPI from ipie.qmc.afqmc import AFQMC from ipie.estimators.energy import EnergyEstimator from ipie.systems import Generic -from ipie.hamiltonians.elph.holstein import HolsteinModel -from ipie.trial_wavefunction.holstein.toyozawa import ToyozawaTrial -from ipie.trial_wavefunction.holstein.variational.toyozawa_variational import variational_trial_toyozawa -from ipie.walkers.eph_walkers import EphWalkers -from ipie.propagation.holstein import HolsteinPropagatorImportance +from ipie.addons.ephqmc.hamiltonians.holstein import HolsteinModel +from ipie.addons.ephqmc.trial_wavefunction.toyozawa import ToyozawaTrial +from ipie.addons.ephqmc.trial_wavefunction.variational.toyozawa_variational import variational_trial_toyozawa +from ipie.addons.ephqmc.walkers.eph_walkers import EphWalkers +from ipie.addons.ephqmc.propagation.holstein import HolsteinPropagatorImportance from ipie.qmc.options import QMCParams #System Parameters @@ -41,7 +41,6 @@ etrial, beta_shift, el_trial = variational_trial_toyozawa( initial_phonons, initial_electron, ham, system ) - wavefunction = np.column_stack([beta_shift, el_trial]) #Setup trial diff --git a/ipie/addons/ephqmc/estimators/local_energy_holstein.py b/ipie/addons/ephqmc/estimators/local_energy_holstein.py new file mode 100644 index 00000000..0de5ace9 --- /dev/null +++ b/ipie/addons/ephqmc/estimators/local_energy_holstein.py @@ -0,0 +1,56 @@ +import numpy as np + +from ipie.addons.ephqmc.hamiltonians.holstein import HolsteinModel +from ipie.addons.ephqmc.trial_wavefunction.eph_trial_base import EphTrialWavefunctionBase +from ipie.addons.ephqmc.walkers.eph_walkers import EphWalkers +from ipie.systems.generic import Generic +from ipie.utils.backend import arraylib as xp + + +def local_energy_holstein( + system: Generic, + hamiltonian: HolsteinModel, + walkers: EphWalkers, + trial: EphTrialWavefunctionBase +): + """Computes the local energy for the Holstein model.""" +# energy = xp.zeros((walkers.nwalkers, 5), dtype=xp.complex128) + +# gf = trial.calc_greens_function(walkers) + +# energy[:, 1] = np.sum(hamiltonian.T[0] * gf[0], axis=(1,2)) +# if system.ndown > 0: +# energy[:, 1] += np.sum(hamiltonian.T[1] * gf[1], axis=(1,2)) + +# energy[:, 2] = np.sum(np.diagonal(gf[0], axis1=1, axis2=2) * walkers.x, axis=1) +# if system.ndown > 0: +# energy[:, 2] += np.sum(np.diagonal(gf[1], axis1=1, axis2=2) * walkers.x, axis=1) +# energy[:, 2] *= hamiltonian.const + +# energy[:, 3] = 0.5 * hamiltonian.m * hamiltonian.w0**2 * np.sum(walkers.x**2, axis=1) +# energy[:, 3] -= 0.5 * hamiltonian.nsites * hamiltonian.w0 +# energy[:, 4] = -0.5 * trial.calc_phonon_laplacian_locenergy(walkers) / hamiltonian.m + +# energy[:, 0] = np.sum(energy[:,1:], axis=1) + ## + energy = xp.zeros((walkers.nwalkers, 5), dtype=xp.complex128) + + #get greens_function from estimators + gf = trial.calc_greens_function(walkers) + + #TODO make this nicer + for n in range(walkers.nwalkers): + energy[n, 1] = np.einsum('ij->', hamiltonian.T[0] * gf[0][n] + hamiltonian.T[1] * gf[1][n]) + + #TODO this performs too many summations + energy[:, 2] = hamiltonian.const * np.einsum('nii,ni->n', gf[0] + gf[1], walkers.x) + + energy[:, 3] = 0.5 * hamiltonian.m * hamiltonian.w0**2 * np.einsum('ni->n', walkers.x**2) + energy[:, 3] -= 0.5 * hamiltonian.nsites * hamiltonian.w0 + energy[:, 4] = -0.5 * trial.calc_phonon_laplacian_locenergy(walkers) / hamiltonian.m + energy[:, 0] = np.sum(energy[:,1:], axis=1) + + return energy + + + diff --git a/ipie/hamiltonians/elph/generic.py b/ipie/addons/ephqmc/hamiltonians/generic.py similarity index 100% rename from ipie/hamiltonians/elph/generic.py rename to ipie/addons/ephqmc/hamiltonians/generic.py diff --git a/ipie/hamiltonians/elph/holstein.py b/ipie/addons/ephqmc/hamiltonians/holstein.py similarity index 100% rename from ipie/hamiltonians/elph/holstein.py rename to ipie/addons/ephqmc/hamiltonians/holstein.py diff --git a/ipie/propagation/holstein.py b/ipie/addons/ephqmc/propagation/holstein.py similarity index 98% rename from ipie/propagation/holstein.py rename to ipie/addons/ephqmc/propagation/holstein.py index 796531fa..08434520 100644 --- a/ipie/propagation/holstein.py +++ b/ipie/addons/ephqmc/propagation/holstein.py @@ -2,7 +2,7 @@ import time import scipy.linalg -from ipie.hamiltonians.elph.holstein import HolsteinModel +from ipie.addons.ephqmc.hamiltonians.holstein import HolsteinModel from ipie.propagation.operations import propagate_one_body from ipie.utils.backend import synchronize, cast_to_device from ipie.propagation.continuous_base import PropagatorTimer diff --git a/ipie/trial_wavefunction/holstein/coherent_state.py b/ipie/addons/ephqmc/trial_wavefunction/coherent_state.py similarity index 94% rename from ipie/trial_wavefunction/holstein/coherent_state.py rename to ipie/addons/ephqmc/trial_wavefunction/coherent_state.py index b3066e61..809021cf 100644 --- a/ipie/trial_wavefunction/holstein/coherent_state.py +++ b/ipie/addons/ephqmc/trial_wavefunction/coherent_state.py @@ -1,9 +1,7 @@ import numpy as np import scipy.linalg -#TODO add greens_function_coherent_state in estimators -from ipie.estimators.greens_function_single_det import greens_function_single_det -from ipie.trial_wavefunction.holstein.eph_trial_base import EphTrialWavefunctionBase +from ipie.addons.ephqmc.trial_wavefunction.eph_trial_base import EphTrialWavefunctionBase from ipie.utils.backend import arraylib as xp #TODO greensfunctions are in estimators diff --git a/ipie/trial_wavefunction/holstein/eph_trial_base.py b/ipie/addons/ephqmc/trial_wavefunction/eph_trial_base.py similarity index 90% rename from ipie/trial_wavefunction/holstein/eph_trial_base.py rename to ipie/addons/ephqmc/trial_wavefunction/eph_trial_base.py index 21219a0c..1e242626 100644 --- a/ipie/trial_wavefunction/holstein/eph_trial_base.py +++ b/ipie/addons/ephqmc/trial_wavefunction/eph_trial_base.py @@ -1,9 +1,9 @@ import numpy as np from ipie.trial_wavefunction.wavefunction_base import TrialWavefunctionBase +#NOTE could inherit from TrialWavefunctionBase, but would need to redefine abstract methods.. class EphTrialWavefunctionBase(): def __init__(self, wavefunction, num_elec, num_basis, verbose=False): - #super().__init__(wavefunction, num_elec, num_basis, verbose=verbose) self.nelec = num_elec self.nbasis = num_basis self.nalpha, self.nbeta = self.nelec diff --git a/ipie/trial_wavefunction/holstein/toyozawa.py b/ipie/addons/ephqmc/trial_wavefunction/toyozawa.py similarity index 70% rename from ipie/trial_wavefunction/holstein/toyozawa.py rename to ipie/addons/ephqmc/trial_wavefunction/toyozawa.py index 384f5f8f..770520ad 100644 --- a/ipie/trial_wavefunction/holstein/toyozawa.py +++ b/ipie/addons/ephqmc/trial_wavefunction/toyozawa.py @@ -1,26 +1,9 @@ import numpy as np import scipy.linalg -#TODO add greens_function_coherent_state in estimators -#NOTE can use single_det for now -from ipie.estimators.greens_function_single_det import greens_function_single_det -from ipie.trial_wavefunction.holstein.eph_trial_base import EphTrialWavefunctionBase -from ipie.trial_wavefunction.holstein.coherent_state import CoherentStateTrial - +from ipie.addons.ephqmc.trial_wavefunction.coherent_state import CoherentStateTrial from ipie.utils.backend import arraylib as xp -from ipie.propagation.overlap import calc_overlap_single_det_uhf -from ipie.estimators.greens_function_single_det import greens_function_single_det -from ipie.config import CommType, config, MPI - -#TODO greensfunctions are in estimators - -def circ_perm(lst): - """""" - cpy = lst[:] # take a copy because a list is a mutable object - yield cpy - for i in range(len(lst) - 1): - cpy = cpy[1:] + [cpy[0]] - yield cpy +from ipie.addons.ephqmc.trial_wavefunction.variational.toyozawa_variational import circ_perm class ToyozawaTrial(CoherentStateTrial): """""" @@ -34,7 +17,6 @@ def calculate_energy(self, system, hamiltonian): pass def calc_overlap(self, walkers) -> np.ndarray: - #TODO this will be a concoction of phonon and electronic overlap _ = self.calc_phonon_overlap(walkers) _ = self.calc_electronic_overlap(walkers) walkers.total_ovlp = walkers.el_ovlp * walkers.ph_ovlp @@ -80,8 +62,6 @@ def calc_phonon_laplacian_locenergy(self, walkers) -> np.ndarray: def calc_electronic_overlap(self, walkers) -> np.ndarray: """""" for ip,perm in enumerate(self.perms): -# print('psia shape: ', self.psia.shape) -# print('walkers shape: ', walkers.phia.shape) ovlp_a = xp.einsum("wmi,mj->wij", walkers.phia, self.psia[perm, :].conj(), optimize=True) sign_a, log_ovlp_a = xp.linalg.slogdet(ovlp_a) @@ -94,25 +74,16 @@ def calc_electronic_overlap(self, walkers) -> np.ndarray: walkers.el_ovlp[:, ip] = ot - el_ovlp = np.sum(walkers.el_ovlp, axis=1) #NOTE this was ph before?? - -# if verbose: -# print('el_ovlp: ', el_ovlp[0]) + el_ovlp = np.sum(walkers.el_ovlp, axis=1) return el_ovlp def calc_greens_function(self, walkers, build_full=True) -> np.ndarray: """""" - -# greensfct = np.zeros((walkers.nwalkers, self.nsites, self.nsites), dtype=np.complex128) walkers.Ga = np.zeros((walkers.nwalkers, self.nsites, self.nsites), dtype=np.complex128) walkers.Gb = np.zeros_like(walkers.Ga) for ovlp, perm in zip(walkers.total_ovlp.T, self.perms): - #TODO adjust this by just calling gab from exisiting extimators rubric TODO - #overlap_inv = 1 / np.einsum('i,nie->n', self.psia[perm].conj(), walkers.phia) #NOTE psi currently hacked - #greensfct += np.einsum('nie,n,j,n->nji', walkers.phia, overlap_inv, - # self.psia[perm].conj(), ovlp) inv_Oa = xp.linalg.inv(xp.einsum('ie,nif->nef', self.psia[perm,:], walkers.phia.conj())) walkers.Ga += xp.einsum('nie,nef,jf,n->nji', walkers.phia, inv_Oa, self.psia[perm].conj(), ovlp) @@ -120,9 +91,7 @@ def calc_greens_function(self, walkers, build_full=True) -> np.ndarray: if self.ndown > 0: inv_Ob = xp.linalg.inv(xp.einsum('ie,nif->nef', self.psib[perm,:], walkers.phib.conj())) walkers.Gb += xp.einsum('nie,nef,jf,n->nji', walkers.phib, inv_Ob, self.psib[perm].conj(), ovlp) - #greensfct += greens_function_single_det(walkers, self, build_full=build_full) * ovlp -# greensfct = np.einsum('nij,n->nij', greensfct, 1 / np.sum(walkers.total_ovlp, axis=1)) #these sums can be replaced by walkers.ovlp calls walkers.Ga = np.einsum('nij,n->nij', walkers.Ga, 1 / np.sum(walkers.total_ovlp, axis=1)) walkers.Gb = np.einsum('nij,n->nij', walkers.Gb, 1 / np.sum(walkers.total_ovlp, axis=1)) diff --git a/ipie/trial_wavefunction/holstein/variational/coherent_state_variational.py b/ipie/addons/ephqmc/trial_wavefunction/variational/coherent_state_variational.py similarity index 98% rename from ipie/trial_wavefunction/holstein/variational/coherent_state_variational.py rename to ipie/addons/ephqmc/trial_wavefunction/variational/coherent_state_variational.py index 6388c912..47ccd03c 100644 --- a/ipie/trial_wavefunction/holstein/variational/coherent_state_variational.py +++ b/ipie/addons/ephqmc/trial_wavefunction/variational/coherent_state_variational.py @@ -3,7 +3,7 @@ from ipie.legacy.trial_wavefunction.harmonic_oscillator import HarmonicOscillator from ipie.systems.generic import Generic -from ipie.hamiltonians.elph.holstein import HolsteinModel +from ipie.addons.ephqmc.hamiltonians.holstein import HolsteinModel import jax diff --git a/ipie/trial_wavefunction/holstein/variational/toyozawa_variational.py b/ipie/addons/ephqmc/trial_wavefunction/variational/toyozawa_variational.py similarity index 98% rename from ipie/trial_wavefunction/holstein/variational/toyozawa_variational.py rename to ipie/addons/ephqmc/trial_wavefunction/variational/toyozawa_variational.py index e75d239d..95489a6a 100644 --- a/ipie/trial_wavefunction/holstein/variational/toyozawa_variational.py +++ b/ipie/addons/ephqmc/trial_wavefunction/variational/toyozawa_variational.py @@ -1,7 +1,7 @@ import numpy as np from scipy.optimize import minimize, basinhopping from ipie.systems import Generic -from ipie.hamiltonians.elph.holstein import HolsteinModel +from ipie.addons.ephqmc.hamiltonians.holstein import HolsteinModel from jax.config import config config.update("jax_enable_x64", True) diff --git a/ipie/walkers/eph_walkers.py b/ipie/addons/ephqmc/walkers/eph_walkers.py similarity index 98% rename from ipie/walkers/eph_walkers.py rename to ipie/addons/ephqmc/walkers/eph_walkers.py index e9e77f70..9146c760 100644 --- a/ipie/walkers/eph_walkers.py +++ b/ipie/addons/ephqmc/walkers/eph_walkers.py @@ -4,7 +4,7 @@ from ipie.utils.backend import arraylib as xp from ipie.utils.backend import cast_to_device, qr, qr_mode, synchronize from ipie.walkers.base_walkers import BaseWalkers -from ipie.trial_wavefunction.holstein.toyozawa import ToyozawaTrial +from ipie.addons.ephqmc.trial_wavefunction.toyozawa import ToyozawaTrial class EphWalkers(BaseWalkers): """Class tailored to el-ph models where keeping track of phonon overlaps is diff --git a/ipie/estimators/energy.py b/ipie/estimators/energy.py index 0ca2b4e2..64a83062 100644 --- a/ipie/estimators/energy.py +++ b/ipie/estimators/energy.py @@ -18,7 +18,7 @@ import plum from ipie.estimators.estimator_base import EstimatorBase -from ipie.estimators.local_energy_holstein import local_energy_holstein +from ipie.addons.ephqmc.estimators.local_energy_holstein import local_energy_holstein from ipie.estimators.local_energy_batch import ( local_energy_batch, local_energy_multi_det_trial_batch, @@ -31,7 +31,7 @@ local_energy_multi_det_trial_wicks_batch_opt_chunked, ) from ipie.hamiltonians.generic import GenericComplexChol, GenericRealChol -from ipie.hamiltonians.elph.holstein import HolsteinModel +from ipie.addons.ephqmc.hamiltonians.holstein import HolsteinModel from ipie.systems.generic import Generic from ipie.trial_wavefunction.noci import NOCI from ipie.trial_wavefunction.particle_hole import ( @@ -41,10 +41,10 @@ ParticleHoleSlow, ) from ipie.trial_wavefunction.single_det import SingleDet -from ipie.trial_wavefunction.holstein.eph_trial_base import EphTrialWavefunctionBase +from ipie.addons.ephqmc.trial_wavefunction.eph_trial_base import EphTrialWavefunctionBase from ipie.utils.backend import arraylib as xp from ipie.walkers.uhf_walkers import UHFWalkers -from ipie.walkers.eph_walkers import EphWalkers +from ipie.addons.ephqmc.walkers.eph_walkers import EphWalkers @plum.dispatch def local_energy( diff --git a/ipie/estimators/local_energy_holstein.py b/ipie/estimators/local_energy_holstein.py deleted file mode 100644 index 03b0e319..00000000 --- a/ipie/estimators/local_energy_holstein.py +++ /dev/null @@ -1,38 +0,0 @@ -import numpy as np - -from ipie.hamiltonians.elph.holstein import HolsteinModel -from ipie.systems.generic import Generic -from ipie.trial_wavefunction.holstein.eph_trial_base import EphTrialWavefunctionBase -from ipie.utils.backend import arraylib as xp -from ipie.walkers.eph_walkers import EphWalkers - - -def local_energy_holstein( - system: Generic, - hamiltonian: HolsteinModel, - walkers: EphWalkers, - trial: EphTrialWavefunctionBase -): - """Computes the local energy for the Holstein model.""" - energy = xp.zeros((walkers.nwalkers, 5), dtype=xp.complex128) - - gf = trial.calc_greens_function(walkers) - - energy[:, 1] = np.sum(hamiltonian.T[0] * gf[0], axis=(1,2)) - if system.ndown > 0: - energy[:, 1] += np.sum(hamiltonian.T[1] * gf[1], axis=(1,2)) - - #TODO this performs too many summations - energy[:, 2] = np.sum(np.diagonal(gf[0], axis1=1, axis2=2) * walkers.x, axis=1) - if system.ndown > 0: - energy[:, 2] += np.sum(np.diagonal(gf[1], axis1=1, axis2=2) * walkers.x, axis=1) - energy[:, 2] *= hamiltonian.const - - energy[:, 3] = 0.5 * hamiltonian.m * hamiltonian.w0**2 * np.sum(walkers.x**2, axis=1) - energy[:, 3] -= 0.5 * hamiltonian.nsites * hamiltonian.w0 - energy[:, 4] = -0.5 * trial.calc_phonon_laplacian_locenergy(walkers) / hamiltonian.m - energy[:, 0] = np.sum(energy[:,1:], axis=1) - return energy - - - diff --git a/ipie/propagation/propagator.py b/ipie/propagation/propagator.py index a121122d..a605718f 100644 --- a/ipie/propagation/propagator.py +++ b/ipie/propagation/propagator.py @@ -1,6 +1,6 @@ from ipie.hamiltonians.generic import GenericRealChol, GenericComplexChol -from ipie.hamiltonians.elph.holstein import HolsteinModel from ipie.propagation.phaseless_generic import PhaselessGeneric -from ipie.propagation.holstein import HolsteinPropagatorImportance +from ipie.addons.ephqmc.propagation.holstein import HolsteinPropagatorImportance +from ipie.addons.ephqmc.hamiltonians.holstein import HolsteinModel Propagator = {GenericRealChol: PhaselessGeneric, GenericComplexChol: PhaselessGeneric, HolsteinModel: HolsteinPropagatorImportance} From 582a24e07204efd5d3f7b492b1e8a3282bf804d3 Mon Sep 17 00:00:00 2001 From: m-baumgarten Date: Tue, 5 Mar 2024 01:07:10 -0500 Subject: [PATCH 09/61] changed folder name from ipie/addons/ephqmc to ipie/addons/eph --- examples/13-1d_holstein/run_afqmc.py | 10 +++++----- ipie/addons/ephqmc/estimators/local_energy_holstein.py | 6 +++--- ipie/addons/ephqmc/propagation/holstein.py | 2 +- .../addons/ephqmc/trial_wavefunction/coherent_state.py | 2 +- ipie/addons/ephqmc/trial_wavefunction/toyozawa.py | 4 ++-- .../variational/coherent_state_variational.py | 2 +- .../variational/toyozawa_variational.py | 2 +- ipie/addons/ephqmc/walkers/eph_walkers.py | 2 +- ipie/estimators/energy.py | 8 ++++---- ipie/propagation/propagator.py | 4 ++-- 10 files changed, 21 insertions(+), 21 deletions(-) diff --git a/examples/13-1d_holstein/run_afqmc.py b/examples/13-1d_holstein/run_afqmc.py index 0d4ac668..89c75f1c 100644 --- a/examples/13-1d_holstein/run_afqmc.py +++ b/examples/13-1d_holstein/run_afqmc.py @@ -5,11 +5,11 @@ from ipie.qmc.afqmc import AFQMC from ipie.estimators.energy import EnergyEstimator from ipie.systems import Generic -from ipie.addons.ephqmc.hamiltonians.holstein import HolsteinModel -from ipie.addons.ephqmc.trial_wavefunction.toyozawa import ToyozawaTrial -from ipie.addons.ephqmc.trial_wavefunction.variational.toyozawa_variational import variational_trial_toyozawa -from ipie.addons.ephqmc.walkers.eph_walkers import EphWalkers -from ipie.addons.ephqmc.propagation.holstein import HolsteinPropagatorImportance +from ipie.addons.eph.hamiltonians.holstein import HolsteinModel +from ipie.addons.eph.trial_wavefunction.toyozawa import ToyozawaTrial +from ipie.addons.eph.trial_wavefunction.variational.toyozawa_variational import variational_trial_toyozawa +from ipie.addons.eph.walkers.eph_walkers import EphWalkers +from ipie.addons.eph.propagation.holstein import HolsteinPropagatorImportance from ipie.qmc.options import QMCParams #System Parameters diff --git a/ipie/addons/ephqmc/estimators/local_energy_holstein.py b/ipie/addons/ephqmc/estimators/local_energy_holstein.py index 0de5ace9..bc3100c5 100644 --- a/ipie/addons/ephqmc/estimators/local_energy_holstein.py +++ b/ipie/addons/ephqmc/estimators/local_energy_holstein.py @@ -1,8 +1,8 @@ import numpy as np -from ipie.addons.ephqmc.hamiltonians.holstein import HolsteinModel -from ipie.addons.ephqmc.trial_wavefunction.eph_trial_base import EphTrialWavefunctionBase -from ipie.addons.ephqmc.walkers.eph_walkers import EphWalkers +from ipie.addons.eph.hamiltonians.holstein import HolsteinModel +from ipie.addons.eph.trial_wavefunction.eph_trial_base import EphTrialWavefunctionBase +from ipie.addons.eph.walkers.eph_walkers import EphWalkers from ipie.systems.generic import Generic from ipie.utils.backend import arraylib as xp diff --git a/ipie/addons/ephqmc/propagation/holstein.py b/ipie/addons/ephqmc/propagation/holstein.py index 08434520..591a4da2 100644 --- a/ipie/addons/ephqmc/propagation/holstein.py +++ b/ipie/addons/ephqmc/propagation/holstein.py @@ -2,7 +2,7 @@ import time import scipy.linalg -from ipie.addons.ephqmc.hamiltonians.holstein import HolsteinModel +from ipie.addons.eph.hamiltonians.holstein import HolsteinModel from ipie.propagation.operations import propagate_one_body from ipie.utils.backend import synchronize, cast_to_device from ipie.propagation.continuous_base import PropagatorTimer diff --git a/ipie/addons/ephqmc/trial_wavefunction/coherent_state.py b/ipie/addons/ephqmc/trial_wavefunction/coherent_state.py index 809021cf..323f6725 100644 --- a/ipie/addons/ephqmc/trial_wavefunction/coherent_state.py +++ b/ipie/addons/ephqmc/trial_wavefunction/coherent_state.py @@ -1,7 +1,7 @@ import numpy as np import scipy.linalg -from ipie.addons.ephqmc.trial_wavefunction.eph_trial_base import EphTrialWavefunctionBase +from ipie.addons.eph.trial_wavefunction.eph_trial_base import EphTrialWavefunctionBase from ipie.utils.backend import arraylib as xp #TODO greensfunctions are in estimators diff --git a/ipie/addons/ephqmc/trial_wavefunction/toyozawa.py b/ipie/addons/ephqmc/trial_wavefunction/toyozawa.py index 770520ad..32d14c10 100644 --- a/ipie/addons/ephqmc/trial_wavefunction/toyozawa.py +++ b/ipie/addons/ephqmc/trial_wavefunction/toyozawa.py @@ -1,9 +1,9 @@ import numpy as np import scipy.linalg -from ipie.addons.ephqmc.trial_wavefunction.coherent_state import CoherentStateTrial +from ipie.addons.eph.trial_wavefunction.coherent_state import CoherentStateTrial from ipie.utils.backend import arraylib as xp -from ipie.addons.ephqmc.trial_wavefunction.variational.toyozawa_variational import circ_perm +from ipie.addons.eph.trial_wavefunction.variational.toyozawa_variational import circ_perm class ToyozawaTrial(CoherentStateTrial): """""" diff --git a/ipie/addons/ephqmc/trial_wavefunction/variational/coherent_state_variational.py b/ipie/addons/ephqmc/trial_wavefunction/variational/coherent_state_variational.py index 47ccd03c..d7e48b94 100644 --- a/ipie/addons/ephqmc/trial_wavefunction/variational/coherent_state_variational.py +++ b/ipie/addons/ephqmc/trial_wavefunction/variational/coherent_state_variational.py @@ -3,7 +3,7 @@ from ipie.legacy.trial_wavefunction.harmonic_oscillator import HarmonicOscillator from ipie.systems.generic import Generic -from ipie.addons.ephqmc.hamiltonians.holstein import HolsteinModel +from ipie.addons.eph.hamiltonians.holstein import HolsteinModel import jax diff --git a/ipie/addons/ephqmc/trial_wavefunction/variational/toyozawa_variational.py b/ipie/addons/ephqmc/trial_wavefunction/variational/toyozawa_variational.py index 95489a6a..4a4bd175 100644 --- a/ipie/addons/ephqmc/trial_wavefunction/variational/toyozawa_variational.py +++ b/ipie/addons/ephqmc/trial_wavefunction/variational/toyozawa_variational.py @@ -1,7 +1,7 @@ import numpy as np from scipy.optimize import minimize, basinhopping from ipie.systems import Generic -from ipie.addons.ephqmc.hamiltonians.holstein import HolsteinModel +from ipie.addons.eph.hamiltonians.holstein import HolsteinModel from jax.config import config config.update("jax_enable_x64", True) diff --git a/ipie/addons/ephqmc/walkers/eph_walkers.py b/ipie/addons/ephqmc/walkers/eph_walkers.py index 9146c760..63ee05f5 100644 --- a/ipie/addons/ephqmc/walkers/eph_walkers.py +++ b/ipie/addons/ephqmc/walkers/eph_walkers.py @@ -4,7 +4,7 @@ from ipie.utils.backend import arraylib as xp from ipie.utils.backend import cast_to_device, qr, qr_mode, synchronize from ipie.walkers.base_walkers import BaseWalkers -from ipie.addons.ephqmc.trial_wavefunction.toyozawa import ToyozawaTrial +from ipie.addons.eph.trial_wavefunction.toyozawa import ToyozawaTrial class EphWalkers(BaseWalkers): """Class tailored to el-ph models where keeping track of phonon overlaps is diff --git a/ipie/estimators/energy.py b/ipie/estimators/energy.py index 64a83062..49e85846 100644 --- a/ipie/estimators/energy.py +++ b/ipie/estimators/energy.py @@ -18,7 +18,7 @@ import plum from ipie.estimators.estimator_base import EstimatorBase -from ipie.addons.ephqmc.estimators.local_energy_holstein import local_energy_holstein +from ipie.addons.eph.estimators.local_energy_holstein import local_energy_holstein from ipie.estimators.local_energy_batch import ( local_energy_batch, local_energy_multi_det_trial_batch, @@ -31,7 +31,7 @@ local_energy_multi_det_trial_wicks_batch_opt_chunked, ) from ipie.hamiltonians.generic import GenericComplexChol, GenericRealChol -from ipie.addons.ephqmc.hamiltonians.holstein import HolsteinModel +from ipie.addons.eph.hamiltonians.holstein import HolsteinModel from ipie.systems.generic import Generic from ipie.trial_wavefunction.noci import NOCI from ipie.trial_wavefunction.particle_hole import ( @@ -41,10 +41,10 @@ ParticleHoleSlow, ) from ipie.trial_wavefunction.single_det import SingleDet -from ipie.addons.ephqmc.trial_wavefunction.eph_trial_base import EphTrialWavefunctionBase +from ipie.addons.eph.trial_wavefunction.eph_trial_base import EphTrialWavefunctionBase from ipie.utils.backend import arraylib as xp from ipie.walkers.uhf_walkers import UHFWalkers -from ipie.addons.ephqmc.walkers.eph_walkers import EphWalkers +from ipie.addons.eph.walkers.eph_walkers import EphWalkers @plum.dispatch def local_energy( diff --git a/ipie/propagation/propagator.py b/ipie/propagation/propagator.py index a605718f..587586b3 100644 --- a/ipie/propagation/propagator.py +++ b/ipie/propagation/propagator.py @@ -1,6 +1,6 @@ from ipie.hamiltonians.generic import GenericRealChol, GenericComplexChol from ipie.propagation.phaseless_generic import PhaselessGeneric -from ipie.addons.ephqmc.propagation.holstein import HolsteinPropagatorImportance -from ipie.addons.ephqmc.hamiltonians.holstein import HolsteinModel +from ipie.addons.eph.propagation.holstein import HolsteinPropagatorImportance +from ipie.addons.eph.hamiltonians.holstein import HolsteinModel Propagator = {GenericRealChol: PhaselessGeneric, GenericComplexChol: PhaselessGeneric, HolsteinModel: HolsteinPropagatorImportance} From 69e1ef5b633a88f1caf7b8839fe3ba1cf5644cd8 Mon Sep 17 00:00:00 2001 From: m-baumgarten Date: Tue, 5 Mar 2024 15:15:46 -0500 Subject: [PATCH 10/61] changed directory name finally --- ipie/addons/{ephqmc => eph}/estimators/local_energy_holstein.py | 0 ipie/addons/{ephqmc => eph}/hamiltonians/generic.py | 0 ipie/addons/{ephqmc => eph}/hamiltonians/holstein.py | 0 ipie/addons/{ephqmc => eph}/propagation/holstein.py | 0 ipie/addons/{ephqmc => eph}/trial_wavefunction/coherent_state.py | 0 ipie/addons/{ephqmc => eph}/trial_wavefunction/eph_trial_base.py | 0 ipie/addons/{ephqmc => eph}/trial_wavefunction/toyozawa.py | 0 .../trial_wavefunction/variational/coherent_state_variational.py | 0 .../trial_wavefunction/variational/toyozawa_variational.py | 0 ipie/addons/{ephqmc => eph}/walkers/eph_walkers.py | 0 10 files changed, 0 insertions(+), 0 deletions(-) rename ipie/addons/{ephqmc => eph}/estimators/local_energy_holstein.py (100%) rename ipie/addons/{ephqmc => eph}/hamiltonians/generic.py (100%) rename ipie/addons/{ephqmc => eph}/hamiltonians/holstein.py (100%) rename ipie/addons/{ephqmc => eph}/propagation/holstein.py (100%) rename ipie/addons/{ephqmc => eph}/trial_wavefunction/coherent_state.py (100%) rename ipie/addons/{ephqmc => eph}/trial_wavefunction/eph_trial_base.py (100%) rename ipie/addons/{ephqmc => eph}/trial_wavefunction/toyozawa.py (100%) rename ipie/addons/{ephqmc => eph}/trial_wavefunction/variational/coherent_state_variational.py (100%) rename ipie/addons/{ephqmc => eph}/trial_wavefunction/variational/toyozawa_variational.py (100%) rename ipie/addons/{ephqmc => eph}/walkers/eph_walkers.py (100%) diff --git a/ipie/addons/ephqmc/estimators/local_energy_holstein.py b/ipie/addons/eph/estimators/local_energy_holstein.py similarity index 100% rename from ipie/addons/ephqmc/estimators/local_energy_holstein.py rename to ipie/addons/eph/estimators/local_energy_holstein.py diff --git a/ipie/addons/ephqmc/hamiltonians/generic.py b/ipie/addons/eph/hamiltonians/generic.py similarity index 100% rename from ipie/addons/ephqmc/hamiltonians/generic.py rename to ipie/addons/eph/hamiltonians/generic.py diff --git a/ipie/addons/ephqmc/hamiltonians/holstein.py b/ipie/addons/eph/hamiltonians/holstein.py similarity index 100% rename from ipie/addons/ephqmc/hamiltonians/holstein.py rename to ipie/addons/eph/hamiltonians/holstein.py diff --git a/ipie/addons/ephqmc/propagation/holstein.py b/ipie/addons/eph/propagation/holstein.py similarity index 100% rename from ipie/addons/ephqmc/propagation/holstein.py rename to ipie/addons/eph/propagation/holstein.py diff --git a/ipie/addons/ephqmc/trial_wavefunction/coherent_state.py b/ipie/addons/eph/trial_wavefunction/coherent_state.py similarity index 100% rename from ipie/addons/ephqmc/trial_wavefunction/coherent_state.py rename to ipie/addons/eph/trial_wavefunction/coherent_state.py diff --git a/ipie/addons/ephqmc/trial_wavefunction/eph_trial_base.py b/ipie/addons/eph/trial_wavefunction/eph_trial_base.py similarity index 100% rename from ipie/addons/ephqmc/trial_wavefunction/eph_trial_base.py rename to ipie/addons/eph/trial_wavefunction/eph_trial_base.py diff --git a/ipie/addons/ephqmc/trial_wavefunction/toyozawa.py b/ipie/addons/eph/trial_wavefunction/toyozawa.py similarity index 100% rename from ipie/addons/ephqmc/trial_wavefunction/toyozawa.py rename to ipie/addons/eph/trial_wavefunction/toyozawa.py diff --git a/ipie/addons/ephqmc/trial_wavefunction/variational/coherent_state_variational.py b/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py similarity index 100% rename from ipie/addons/ephqmc/trial_wavefunction/variational/coherent_state_variational.py rename to ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py diff --git a/ipie/addons/ephqmc/trial_wavefunction/variational/toyozawa_variational.py b/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py similarity index 100% rename from ipie/addons/ephqmc/trial_wavefunction/variational/toyozawa_variational.py rename to ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py diff --git a/ipie/addons/ephqmc/walkers/eph_walkers.py b/ipie/addons/eph/walkers/eph_walkers.py similarity index 100% rename from ipie/addons/ephqmc/walkers/eph_walkers.py rename to ipie/addons/eph/walkers/eph_walkers.py From ccc26493eb1913ffbe07752a43c3cbe556d1b811 Mon Sep 17 00:00:00 2001 From: m-baumgarten Date: Tue, 5 Mar 2024 16:31:52 -0500 Subject: [PATCH 11/61] Cleanup of HolsteinModel class plus some explanation of the model. --- ipie/addons/eph/hamiltonians/generic.py | 11 -------- ipie/addons/eph/hamiltonians/holstein.py | 32 +++++++++++++++++++++--- 2 files changed, 29 insertions(+), 14 deletions(-) delete mode 100644 ipie/addons/eph/hamiltonians/generic.py diff --git a/ipie/addons/eph/hamiltonians/generic.py b/ipie/addons/eph/hamiltonians/generic.py deleted file mode 100644 index 84df1a80..00000000 --- a/ipie/addons/eph/hamiltonians/generic.py +++ /dev/null @@ -1,11 +0,0 @@ -import numpy -from ipie.hamiltonians.generic_base import GenericBase -from ipie.utils.backend import arraylib as xp - - -class GenericElphBase(GenericBase): - """""" - - def __init__(self, ): - pass - diff --git a/ipie/addons/eph/hamiltonians/holstein.py b/ipie/addons/eph/hamiltonians/holstein.py index 46354b94..2456611c 100644 --- a/ipie/addons/eph/hamiltonians/holstein.py +++ b/ipie/addons/eph/hamiltonians/holstein.py @@ -1,10 +1,35 @@ import numpy -from ipie.hamiltonians.generic_base import GenericBase from ipie.utils.backend import arraylib as xp -class HolsteinModel(GenericBase): - """Class for Holstein model carrying elph tensor and system parameters +class HolsteinModel(): + r"""Class carrying parameters specifying a 1D Holstein chain. + + The Holstein model is described by the Hamiltonian + + .. math:: + \hat{H} = -t \sum_{\langle ij\rangle} \hat{a}_i^\dagger \hat{a}_j + - g \sqrt{2 w_0 m} \sum_i \hat{a}_i^\dagger \hat{a}_i \hat{X}_i + + \bigg(\sum_i \frac{m w_0^2}{2} \hat{X}_i^2 + \frac{1}{2m} \hat{P}_i^2 + - \frac{w_0}{2}\bigg), + + where :math:`t` is associated with the electronic hopping, :math:`g` with + the electron-phonon coupling strength, and :math:``w_0` with the phonon + frequency. + + Parameters + ---------- + g : + Electron-phonon coupling strength + t : + Electron hopping parameter + w0 : + Phonon frequency + nsites : + Length of the 1D Holstein chain + pbc : + Boolean specifying whether periodic boundary conditions should be + employed. """ def __init__(self, g: float, t: float, w0: float, nsites: int, pbc: bool): @@ -18,6 +43,7 @@ def __init__(self, g: float, t: float, w0: float, nsites: int, pbc: bool): self.const = -self.g * numpy.sqrt(2. * self.m * self.w0) def build(self): + """Constructs electronic hopping matrix.""" self.T = numpy.diag(numpy.ones(self.nsites-1), 1) self.T += numpy.diag(numpy.ones(self.nsites-1), -1) From ce14a30c427704ecd3ed6c355249ca8276c36e5c Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Tue, 5 Mar 2024 16:33:57 -0500 Subject: [PATCH 12/61] added __init__ files in new directories --- ipie/addons/eph/__init__.py | 17 +++++++++++++++++ ipie/addons/eph/estimators/__init__.py | 17 +++++++++++++++++ ipie/addons/eph/hamiltonians/__init__.py | 17 +++++++++++++++++ ipie/addons/eph/propagation/__init__.py | 17 +++++++++++++++++ ipie/addons/eph/trial_wavefunction/__init__.py | 17 +++++++++++++++++ ipie/addons/eph/walkers/__init__.py | 17 +++++++++++++++++ 6 files changed, 102 insertions(+) create mode 100644 ipie/addons/eph/__init__.py create mode 100644 ipie/addons/eph/estimators/__init__.py create mode 100644 ipie/addons/eph/hamiltonians/__init__.py create mode 100644 ipie/addons/eph/propagation/__init__.py create mode 100644 ipie/addons/eph/trial_wavefunction/__init__.py create mode 100644 ipie/addons/eph/walkers/__init__.py diff --git a/ipie/addons/eph/__init__.py b/ipie/addons/eph/__init__.py new file mode 100644 index 00000000..871770c1 --- /dev/null +++ b/ipie/addons/eph/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2022 The ipie Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Directory for additions to ipie which depend on the core ipie library. +# New features should mirror the ipie layout e.g. +# ipie/addons/finite_temperature/qmc/afqmc.py etc. diff --git a/ipie/addons/eph/estimators/__init__.py b/ipie/addons/eph/estimators/__init__.py new file mode 100644 index 00000000..871770c1 --- /dev/null +++ b/ipie/addons/eph/estimators/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2022 The ipie Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Directory for additions to ipie which depend on the core ipie library. +# New features should mirror the ipie layout e.g. +# ipie/addons/finite_temperature/qmc/afqmc.py etc. diff --git a/ipie/addons/eph/hamiltonians/__init__.py b/ipie/addons/eph/hamiltonians/__init__.py new file mode 100644 index 00000000..871770c1 --- /dev/null +++ b/ipie/addons/eph/hamiltonians/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2022 The ipie Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Directory for additions to ipie which depend on the core ipie library. +# New features should mirror the ipie layout e.g. +# ipie/addons/finite_temperature/qmc/afqmc.py etc. diff --git a/ipie/addons/eph/propagation/__init__.py b/ipie/addons/eph/propagation/__init__.py new file mode 100644 index 00000000..871770c1 --- /dev/null +++ b/ipie/addons/eph/propagation/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2022 The ipie Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Directory for additions to ipie which depend on the core ipie library. +# New features should mirror the ipie layout e.g. +# ipie/addons/finite_temperature/qmc/afqmc.py etc. diff --git a/ipie/addons/eph/trial_wavefunction/__init__.py b/ipie/addons/eph/trial_wavefunction/__init__.py new file mode 100644 index 00000000..871770c1 --- /dev/null +++ b/ipie/addons/eph/trial_wavefunction/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2022 The ipie Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Directory for additions to ipie which depend on the core ipie library. +# New features should mirror the ipie layout e.g. +# ipie/addons/finite_temperature/qmc/afqmc.py etc. diff --git a/ipie/addons/eph/walkers/__init__.py b/ipie/addons/eph/walkers/__init__.py new file mode 100644 index 00000000..871770c1 --- /dev/null +++ b/ipie/addons/eph/walkers/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2022 The ipie Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Directory for additions to ipie which depend on the core ipie library. +# New features should mirror the ipie layout e.g. +# ipie/addons/finite_temperature/qmc/afqmc.py etc. From d839b233b77da3d1a83a793518e60625f7b27ef8 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Tue, 5 Mar 2024 16:48:25 -0500 Subject: [PATCH 13/61] cleaned up comments and added docstring --- .../eph/estimators/local_energy_holstein.py | 51 +++++++++---------- 1 file changed, 23 insertions(+), 28 deletions(-) diff --git a/ipie/addons/eph/estimators/local_energy_holstein.py b/ipie/addons/eph/estimators/local_energy_holstein.py index bc3100c5..cf7c04d4 100644 --- a/ipie/addons/eph/estimators/local_energy_holstein.py +++ b/ipie/addons/eph/estimators/local_energy_holstein.py @@ -3,6 +3,7 @@ from ipie.addons.eph.hamiltonians.holstein import HolsteinModel from ipie.addons.eph.trial_wavefunction.eph_trial_base import EphTrialWavefunctionBase from ipie.addons.eph.walkers.eph_walkers import EphWalkers + from ipie.systems.generic import Generic from ipie.utils.backend import arraylib as xp @@ -12,44 +13,38 @@ def local_energy_holstein( hamiltonian: HolsteinModel, walkers: EphWalkers, trial: EphTrialWavefunctionBase -): - """Computes the local energy for the Holstein model.""" -# energy = xp.zeros((walkers.nwalkers, 5), dtype=xp.complex128) - -# gf = trial.calc_greens_function(walkers) +) -> np.ndarray: + r"""Computes the local energy for the Holstein model via -# energy[:, 1] = np.sum(hamiltonian.T[0] * gf[0], axis=(1,2)) -# if system.ndown > 0: -# energy[:, 1] += np.sum(hamiltonian.T[1] * gf[1], axis=(1,2)) - -# energy[:, 2] = np.sum(np.diagonal(gf[0], axis1=1, axis2=2) * walkers.x, axis=1) -# if system.ndown > 0: -# energy[:, 2] += np.sum(np.diagonal(gf[1], axis1=1, axis2=2) * walkers.x, axis=1) -# energy[:, 2] *= hamiltonian.const - -# energy[:, 3] = 0.5 * hamiltonian.m * hamiltonian.w0**2 * np.sum(walkers.x**2, axis=1) -# energy[:, 3] -= 0.5 * hamiltonian.nsites * hamiltonian.w0 -# energy[:, 4] = -0.5 * trial.calc_phonon_laplacian_locenergy(walkers) / hamiltonian.m - -# energy[:, 0] = np.sum(energy[:,1:], axis=1) - ## + .. math:: + \frac{\langle \Psi_\mathrm{T} | \hat{H} | \Phi_\mathrm{w}\rangle} + {\langle \Psi_\mathrm{T} | \Phi_\mathrm{w}\rangle}, + + where :math:`| \Phi_\mathrm{w}\rangle = \sum_{\tau \in cyclic perms} + |\phi(\tau(r))\rangle \otimes |\beta(\tau(X))\rangle`. In this ansatz for + the walkers :math:`|\beta\rangle` is a coherent state, which corresonds to a + by :math:`\beta` displaced vacuum state. + """ + energy = xp.zeros((walkers.nwalkers, 5), dtype=xp.complex128) - #get greens_function from estimators gf = trial.calc_greens_function(walkers) - #TODO make this nicer - for n in range(walkers.nwalkers): - energy[n, 1] = np.einsum('ij->', hamiltonian.T[0] * gf[0][n] + hamiltonian.T[1] * gf[1][n]) + energy[:, 1] = np.sum(hamiltonian.T[0] * gf[0], axis=(1,2)) + if system.ndown > 0: + energy[:, 1] += np.sum(hamiltonian.T[1] * gf[1], axis=(1,2)) - #TODO this performs too many summations - energy[:, 2] = hamiltonian.const * np.einsum('nii,ni->n', gf[0] + gf[1], walkers.x) + energy[:, 2] = np.sum(np.diagonal(gf[0], axis1=1, axis2=2) * walkers.x, axis=1) + if system.ndown > 0: + energy[:, 2] += np.sum(np.diagonal(gf[1], axis1=1, axis2=2) * walkers.x, axis=1) + energy[:, 2] *= hamiltonian.const - energy[:, 3] = 0.5 * hamiltonian.m * hamiltonian.w0**2 * np.einsum('ni->n', walkers.x**2) + energy[:, 3] = 0.5 * hamiltonian.m * hamiltonian.w0**2 * np.sum(walkers.x**2, axis=1) energy[:, 3] -= 0.5 * hamiltonian.nsites * hamiltonian.w0 energy[:, 4] = -0.5 * trial.calc_phonon_laplacian_locenergy(walkers) / hamiltonian.m - energy[:, 0] = np.sum(energy[:,1:], axis=1) + energy[:, 0] = np.sum(energy[:,1:], axis=1) + return energy From 00870e5d46c10016133d36fc2b3e4d1baacc63b6 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Tue, 5 Mar 2024 22:16:54 -0500 Subject: [PATCH 14/61] added energy estimator, more or less identical to the one in ipie/estimators/energy.py --- examples/13-1d_holstein/run_afqmc.py | 2 +- ipie/addons/eph/estimators/energy.py | 90 +++++++++++++++++++ .../eph/estimators/local_energy_holstein.py | 4 +- 3 files changed, 93 insertions(+), 3 deletions(-) create mode 100644 ipie/addons/eph/estimators/energy.py diff --git a/examples/13-1d_holstein/run_afqmc.py b/examples/13-1d_holstein/run_afqmc.py index 89c75f1c..3edb7849 100644 --- a/examples/13-1d_holstein/run_afqmc.py +++ b/examples/13-1d_holstein/run_afqmc.py @@ -3,13 +3,13 @@ from mpi4py import MPI from ipie.qmc.afqmc import AFQMC -from ipie.estimators.energy import EnergyEstimator from ipie.systems import Generic from ipie.addons.eph.hamiltonians.holstein import HolsteinModel from ipie.addons.eph.trial_wavefunction.toyozawa import ToyozawaTrial from ipie.addons.eph.trial_wavefunction.variational.toyozawa_variational import variational_trial_toyozawa from ipie.addons.eph.walkers.eph_walkers import EphWalkers from ipie.addons.eph.propagation.holstein import HolsteinPropagatorImportance +from ipie.addons.eph.estimators.energy import EnergyEstimator from ipie.qmc.options import QMCParams #System Parameters diff --git a/ipie/addons/eph/estimators/energy.py b/ipie/addons/eph/estimators/energy.py new file mode 100644 index 00000000..42be4bff --- /dev/null +++ b/ipie/addons/eph/estimators/energy.py @@ -0,0 +1,90 @@ +# Copyright 2022 The ipie Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import plum + +from ipie.estimators.estimator_base import EstimatorBase +from ipie.addons.eph.estimators.local_energy_holstein import local_energy_holstein +from ipie.addons.eph.hamiltonians.holstein import HolsteinModel +from ipie.systems.generic import Generic +from ipie.addons.eph.trial_wavefunction.eph_trial_base import EphTrialWavefunctionBase +from ipie.utils.backend import arraylib as xp +from ipie.addons.eph.walkers.eph_walkers import EphWalkers + +@plum.dispatch +def local_energy( + system: Generic, + hamiltonian: HolsteinModel, + walkers: EphWalkers, + trial: EphTrialWavefunctionBase +): + return local_energy_holstein(system, hamiltonian, walkers, trial) + +class EnergyEstimator(EstimatorBase): + def __init__( + self, + system=None, + ham=None, + trial=None, + filename=None, + ): + assert system is not None + assert ham is not None + assert trial is not None + super().__init__() + self._eshift = 0.0 + self.scalar_estimator = True + self._data = { + "ENumer": 0.0j, + "EDenom": 0.0j, + "ETotal": 0.0j, + "EEl": 0.0j, + "EElPh": 0.0j, + "EPh": 0.0j, + } + self._shape = (len(self.names),) + self._data_index = {k: i for i, k in enumerate(list(self._data.keys()))} + self.print_to_stdout = True + self.ascii_filename = filename + + def compute_estimator(self, system, walkers, hamiltonian, trial, istep=1): + trial.calc_greens_function(walkers) + # Need to be able to dispatch here + energy = local_energy(system, hamiltonian, walkers, trial) + self._data["ENumer"] = xp.sum(walkers.weight * energy[:, 0].real) + self._data["EDenom"] = xp.sum(walkers.weight) + self._data["EEl"] = xp.sum(walkers.weight * energy[:, 1].real) + self._data["EElPh"] = xp.sum(walkers.weight * energy[:, 2].real) + self._data["EPh"] = xp.sum(walkers.weight * energy[:, 3].real) + + return self.data + + def get_index(self, name): + index = self._data_index.get(name, None) + if index is None: + raise RuntimeError(f"Unknown estimator {name}") + return index + + def post_reduce_hook(self, data): + ix_proj = self._data_index["ETotal"] + ix_nume = self._data_index["ENumer"] + ix_deno = self._data_index["EDenom"] + data[ix_proj] = data[ix_nume] / data[ix_deno] + ix_nume = self._data_index["EEl"] + data[ix_nume] = data[ix_nume] / data[ix_deno] + ix_nume = self._data_index["EElPh"] + data[ix_nume] = data[ix_nume] / data[ix_deno] + ix_nume = self._data_index["EPh"] + data[ix_nume] = data[ix_nume] / data[ix_deno] diff --git a/ipie/addons/eph/estimators/local_energy_holstein.py b/ipie/addons/eph/estimators/local_energy_holstein.py index cf7c04d4..506a23e8 100644 --- a/ipie/addons/eph/estimators/local_energy_holstein.py +++ b/ipie/addons/eph/estimators/local_energy_holstein.py @@ -26,7 +26,7 @@ def local_energy_holstein( by :math:`\beta` displaced vacuum state. """ - energy = xp.zeros((walkers.nwalkers, 5), dtype=xp.complex128) + energy = xp.zeros((walkers.nwalkers, 4), dtype=xp.complex128) gf = trial.calc_greens_function(walkers) @@ -41,7 +41,7 @@ def local_energy_holstein( energy[:, 3] = 0.5 * hamiltonian.m * hamiltonian.w0**2 * np.sum(walkers.x**2, axis=1) energy[:, 3] -= 0.5 * hamiltonian.nsites * hamiltonian.w0 - energy[:, 4] = -0.5 * trial.calc_phonon_laplacian_locenergy(walkers) / hamiltonian.m + energy[:, 3] -= 0.5 * trial.calc_phonon_laplacian_locenergy(walkers) / hamiltonian.m energy[:, 0] = np.sum(energy[:,1:], axis=1) From e2e703be5b4fa6cad871b776e10fd5f782e5cd92 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Tue, 5 Mar 2024 22:21:40 -0500 Subject: [PATCH 15/61] removed addon imports and dispatch to local_energy_holstein --- ipie/estimators/energy.py | 15 +-------------- 1 file changed, 1 insertion(+), 14 deletions(-) diff --git a/ipie/estimators/energy.py b/ipie/estimators/energy.py index 49e85846..67113911 100644 --- a/ipie/estimators/energy.py +++ b/ipie/estimators/energy.py @@ -18,7 +18,6 @@ import plum from ipie.estimators.estimator_base import EstimatorBase -from ipie.addons.eph.estimators.local_energy_holstein import local_energy_holstein from ipie.estimators.local_energy_batch import ( local_energy_batch, local_energy_multi_det_trial_batch, @@ -31,7 +30,6 @@ local_energy_multi_det_trial_wicks_batch_opt_chunked, ) from ipie.hamiltonians.generic import GenericComplexChol, GenericRealChol -from ipie.addons.eph.hamiltonians.holstein import HolsteinModel from ipie.systems.generic import Generic from ipie.trial_wavefunction.noci import NOCI from ipie.trial_wavefunction.particle_hole import ( @@ -41,19 +39,9 @@ ParticleHoleSlow, ) from ipie.trial_wavefunction.single_det import SingleDet -from ipie.addons.eph.trial_wavefunction.eph_trial_base import EphTrialWavefunctionBase from ipie.utils.backend import arraylib as xp from ipie.walkers.uhf_walkers import UHFWalkers -from ipie.addons.eph.walkers.eph_walkers import EphWalkers -@plum.dispatch -def local_energy( - system: Generic, - hamiltonian: HolsteinModel, - walkers: EphWalkers, - trial: EphTrialWavefunctionBase -): - return local_energy_holstein(system, hamiltonian, walkers, trial) @plum.dispatch def local_energy( @@ -150,8 +138,7 @@ def compute_estimator(self, system, walkers, hamiltonian, trial, istep=1): self._data["ENumer"] = xp.sum(walkers.weight * energy[:, 0].real) self._data["EDenom"] = xp.sum(walkers.weight) self._data["E1Body"] = xp.sum(walkers.weight * energy[:, 1].real) - self._data["E2Body"] = xp.sum(walkers.weight * energy[:, 2].real) - + self._data["E2Body"] = xp.sum(walkers.weight * energy[:, 2].real) return self.data def get_index(self, name): From dab3687643b3838dca91f4bb5439c5c454b94fdf Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Tue, 5 Mar 2024 22:59:06 -0500 Subject: [PATCH 16/61] enabled AFQMC.build by adding some dummy parameters to HolsteinPropagator.build --- examples/13-1d_holstein/run_afqmc.py | 31 ++++++++++--------------- ipie/addons/eph/propagation/holstein.py | 23 +++++++----------- 2 files changed, 21 insertions(+), 33 deletions(-) diff --git a/examples/13-1d_holstein/run_afqmc.py b/examples/13-1d_holstein/run_afqmc.py index 3edb7849..e3294a38 100644 --- a/examples/13-1d_holstein/run_afqmc.py +++ b/examples/13-1d_holstein/run_afqmc.py @@ -15,7 +15,7 @@ #System Parameters nup = 2 ndown = 2 -nelec = [nup, ndown] +nelec = (nup, ndown) #Hamiltonian Parameters g = 2. @@ -62,10 +62,6 @@ ) walkers.build(trial) -timestep = 0.01 -propagator = HolsteinPropagatorImportance(timestep) -propagator.build(ham) - num_steps_per_block = 10 num_blocks = 10000 add_est = { @@ -74,21 +70,18 @@ ) } -stabilize_freq = 5 -pop_control_freq = 5 seed = 125 -params = QMCParams( - num_walkers=nwalkers, - total_num_walkers=nwalkers * comm.size, - num_blocks=num_blocks, - num_steps_per_block=num_steps_per_block, - timestep=timestep, - num_stblz=stabilize_freq, - pop_control_freq=pop_control_freq, - rng_seed=seed, -) -ephqmc = AFQMC(system, ham, trial, walkers, propagator, params) -trial.calc_overlap(walkers) #Sets ovlp, ph and el overlaps +# Note nwalkers specifies the number of walkers on each CPU +ephqmc = AFQMC.build( + num_elec=nelec, + hamiltonian=ham, + trial_wavefunction=trial, + walkers=walkers, + num_walkers=nwalkers, + seed=seed, + num_steps_per_block=num_steps_per_block, + num_blocks=num_blocks, +) ephqmc.run(additional_estimators=add_est, verbose=False) diff --git a/ipie/addons/eph/propagation/holstein.py b/ipie/addons/eph/propagation/holstein.py index 591a4da2..24537600 100644 --- a/ipie/addons/eph/propagation/holstein.py +++ b/ipie/addons/eph/propagation/holstein.py @@ -8,7 +8,9 @@ from ipie.propagation.continuous_base import PropagatorTimer def construct_one_body_propagator(hamiltonian: HolsteinModel, dt: float): - """""" + """Exponentiates the electronic hopping term to apply it later as + part of the trotterized algorithm. + """ H1 = hamiltonian.T expH1 = numpy.array( [scipy.linalg.expm(-0.5 * dt * H1[0]), scipy.linalg.expm(-0.5 * dt * H1[1])] @@ -16,9 +18,9 @@ def construct_one_body_propagator(hamiltonian: HolsteinModel, dt: float): return expH1 -class HolsteinPropagatorFree(): +class HolsteinPropagatorFree: """""" - def __init__(self, time_step, verbose=False): + def __init__(self, time_step: float, verbose: bool = False): self.dt = time_step self.verbose = verbose self.timer = PropagatorTimer() @@ -27,7 +29,7 @@ def __init__(self, time_step, verbose=False): self.dt_ph = 0.5 * self.dt self.mpi_handler = None - def build(self, hamiltonian) : + def build(self, hamiltonian: HolsteinModel, trial=None, walkers=None, mpi_handler=None): self.expH1 = construct_one_body_propagator(hamiltonian, self.dt) self.const = hamiltonian.g * numpy.sqrt(2. * hamiltonian.m * hamiltonian.w0) * self.dt self.w0 = hamiltonian.w0 @@ -35,14 +37,6 @@ def build(self, hamiltonian) : self.scale = numpy.sqrt(self.dt_ph / self.m) self.nsites = hamiltonian.nsites - def propagate_walkers_one_body(self, walkers): - start_time = time.time() - walkers.phia = propagate_one_body(walkers.phia, self.expH1[0]) - if walkers.ndown > 0 and not walkers.rhf: - walkers.phib = propagate_one_body(walkers.phib, self.expH1[1]) - synchronize() - self.timer.tgemm += time.time() - start_time - def propagate_phonons(self, walkers): start_time = time.time() @@ -119,9 +113,10 @@ def __init__(self, time_step, verbose=False): super().__init__(time_step, verbose=verbose) def propagate_phonons(self, walkers, hamiltonian, trial): - """Propagates phonons via DMC""" + """Propagates phonons via Diffusion MC.""" start_time = time.time() - #no ZPE in pot -> cancels with ZPE of etrial, wouldn't affect estimators anyways + + # No ZPE in pot -> cancels with ZPE of etrial, wouldn't affect estimators anyways ph_ovlp_old = trial.calc_phonon_overlap(walkers) pot = 0.5 * hamiltonian.m * hamiltonian.w0**2 * numpy.sum(walkers.x**2, axis=1) From ad5c2cdb1365a8fb2e32a75eeb863a8b266cd9bf Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Tue, 5 Mar 2024 23:00:16 -0500 Subject: [PATCH 17/61] added docstring to EphWalkers and moved initial trial.calc_overlap from run script to EphWalkers.build --- ipie/addons/eph/walkers/eph_walkers.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/ipie/addons/eph/walkers/eph_walkers.py b/ipie/addons/eph/walkers/eph_walkers.py index 63ee05f5..94796e47 100644 --- a/ipie/addons/eph/walkers/eph_walkers.py +++ b/ipie/addons/eph/walkers/eph_walkers.py @@ -8,7 +8,25 @@ class EphWalkers(BaseWalkers): """Class tailored to el-ph models where keeping track of phonon overlaps is - required.""" + required. Each walker carries along its Slater determinants a phonon + displacement vector, self.x. + + Parameters + ---------- + initial_walker : + Walker that we start the simulation from. Ideally chosen according to + the trial. + nup : + Number of electrons in up-spin space. + ndown : + Number of electrons in down-spin space. + nbasis : + Number of sites in the 1D Holstein chain. + nwalkers : + Number of walkers in the simulation. + verbose : + Print level. + """ def __init__( self, initial_walker: numpy.ndarray, @@ -72,6 +90,8 @@ def build(self, trial): self.buff_size = round(self.set_buff_size_single_walker() / float(self.nwalkers)) self.walker_buffer = numpy.zeros(self.buff_size, dtype=numpy.complex128) + trial.calc_overlap(self) + def cast_to_cupy(self, verbose=False): cast_to_device(self, verbose) From e782941d5bb7440e4c35bde04a72f2623db50597 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Thu, 7 Mar 2024 00:34:20 -0500 Subject: [PATCH 18/61] renamed Eph classes to EPh --- examples/13-1d_holstein/run_afqmc.py | 6 ++-- ipie/addons/eph/estimators/energy.py | 8 ++--- .../eph/estimators/local_energy_holstein.py | 8 ++--- .../eph/trial_wavefunction/coherent_state.py | 31 ++++++++++++++++--- .../eph/trial_wavefunction/eph_trial_base.py | 26 +++++++++++++--- .../addons/eph/trial_wavefunction/toyozawa.py | 26 ++++++++++------ ipie/addons/eph/walkers/eph_walkers.py | 6 ++-- 7 files changed, 77 insertions(+), 34 deletions(-) diff --git a/examples/13-1d_holstein/run_afqmc.py b/examples/13-1d_holstein/run_afqmc.py index e3294a38..9171cfaa 100644 --- a/examples/13-1d_holstein/run_afqmc.py +++ b/examples/13-1d_holstein/run_afqmc.py @@ -7,10 +7,8 @@ from ipie.addons.eph.hamiltonians.holstein import HolsteinModel from ipie.addons.eph.trial_wavefunction.toyozawa import ToyozawaTrial from ipie.addons.eph.trial_wavefunction.variational.toyozawa_variational import variational_trial_toyozawa -from ipie.addons.eph.walkers.eph_walkers import EphWalkers -from ipie.addons.eph.propagation.holstein import HolsteinPropagatorImportance +from ipie.addons.eph.walkers.eph_walkers import EPhWalkers from ipie.addons.eph.estimators.energy import EnergyEstimator -from ipie.qmc.options import QMCParams #System Parameters nup = 2 @@ -53,7 +51,7 @@ trial.set_etrial(etrial) #Setup walkers -walkers = EphWalkers( +walkers = EPhWalkers( initial_walker=wavefunction, nup=nup, ndown=ndown, diff --git a/ipie/addons/eph/estimators/energy.py b/ipie/addons/eph/estimators/energy.py index 42be4bff..cfe7b6bc 100644 --- a/ipie/addons/eph/estimators/energy.py +++ b/ipie/addons/eph/estimators/energy.py @@ -19,16 +19,16 @@ from ipie.addons.eph.estimators.local_energy_holstein import local_energy_holstein from ipie.addons.eph.hamiltonians.holstein import HolsteinModel from ipie.systems.generic import Generic -from ipie.addons.eph.trial_wavefunction.eph_trial_base import EphTrialWavefunctionBase +from ipie.addons.eph.trial_wavefunction.eph_trial_base import EPhTrialWavefunctionBase from ipie.utils.backend import arraylib as xp -from ipie.addons.eph.walkers.eph_walkers import EphWalkers +from ipie.addons.eph.walkers.eph_walkers import EPhWalkers @plum.dispatch def local_energy( system: Generic, hamiltonian: HolsteinModel, - walkers: EphWalkers, - trial: EphTrialWavefunctionBase + walkers: EPhWalkers, + trial: EPhTrialWavefunctionBase ): return local_energy_holstein(system, hamiltonian, walkers, trial) diff --git a/ipie/addons/eph/estimators/local_energy_holstein.py b/ipie/addons/eph/estimators/local_energy_holstein.py index 506a23e8..bce85d48 100644 --- a/ipie/addons/eph/estimators/local_energy_holstein.py +++ b/ipie/addons/eph/estimators/local_energy_holstein.py @@ -1,8 +1,8 @@ import numpy as np from ipie.addons.eph.hamiltonians.holstein import HolsteinModel -from ipie.addons.eph.trial_wavefunction.eph_trial_base import EphTrialWavefunctionBase -from ipie.addons.eph.walkers.eph_walkers import EphWalkers +from ipie.addons.eph.trial_wavefunction.eph_trial_base import EPhTrialWavefunctionBase +from ipie.addons.eph.walkers.eph_walkers import EPhWalkers from ipie.systems.generic import Generic from ipie.utils.backend import arraylib as xp @@ -11,8 +11,8 @@ def local_energy_holstein( system: Generic, hamiltonian: HolsteinModel, - walkers: EphWalkers, - trial: EphTrialWavefunctionBase + walkers: EPhWalkers, + trial: EPhTrialWavefunctionBase ) -> np.ndarray: r"""Computes the local energy for the Holstein model via diff --git a/ipie/addons/eph/trial_wavefunction/coherent_state.py b/ipie/addons/eph/trial_wavefunction/coherent_state.py index 323f6725..2423598a 100644 --- a/ipie/addons/eph/trial_wavefunction/coherent_state.py +++ b/ipie/addons/eph/trial_wavefunction/coherent_state.py @@ -1,12 +1,33 @@ import numpy as np import scipy.linalg -from ipie.addons.eph.trial_wavefunction.eph_trial_base import EphTrialWavefunctionBase +from ipie.addons.eph.trial_wavefunction.eph_trial_base import EPhTrialWavefunctionBase from ipie.utils.backend import arraylib as xp -#TODO greensfunctions are in estimators -class CoherentStateTrial(EphTrialWavefunctionBase): - """""" +class CoherentStateTrial(EPhTrialWavefunctionBase): + r"""Coherent state trial of the form + + .. math:: + |\Phi\rangle \otimes |\beta\rangle, + + where :math:`|\Phi\rangle` corresponds to the electronic wave function and + :math:`|\beta\rangle` to the bosonic wave function. This latter is a + coherent state, i.e. a vacuum state displaced by :math:`\beta`. + + Parameters + ---------- + wavefunction : + Concatenation of trial determinants of up and down spin spaces and beta + specifying the coherent state displacement. + hamiltonian : + Holstein model Hamiltonian + num_elec : + Tuple specifying number of up and down spins + num_basis : + Number of sites of Holstein chain. + verbose : + Print level + """ def __init__(self, wavefunction, hamiltonian, num_elec, num_basis, verbose=False): super().__init__(wavefunction, num_elec, num_basis, verbose=verbose) self.num_elec = num_elec @@ -20,7 +41,7 @@ def __init__(self, wavefunction, hamiltonian, num_elec, num_basis, verbose=False self.psib = wavefunction[:, self.nup+1:self.nup+self.ndown+1] def calculate_energy(self, system, hamiltonian): - #TODO variational_energy_coherent_state in ipie.estimators.local_energy + # TODO variational_energy_coherent_state in ipie.estimators.local_energy ... def calc_overlap(self, walkers) -> np.ndarray: diff --git a/ipie/addons/eph/trial_wavefunction/eph_trial_base.py b/ipie/addons/eph/trial_wavefunction/eph_trial_base.py index 1e242626..dc013cd1 100644 --- a/ipie/addons/eph/trial_wavefunction/eph_trial_base.py +++ b/ipie/addons/eph/trial_wavefunction/eph_trial_base.py @@ -1,9 +1,27 @@ import numpy as np from ipie.trial_wavefunction.wavefunction_base import TrialWavefunctionBase +from typing import Tuple -#NOTE could inherit from TrialWavefunctionBase, but would need to redefine abstract methods.. -class EphTrialWavefunctionBase(): - def __init__(self, wavefunction, num_elec, num_basis, verbose=False): +# NOTE could inherit from TrialWavefunctionBase, +# but would need to redefine abstract methods.. + +class EPhTrialWavefunctionBase: + """Base class for electron-phonon trial wave functions. + + Parameters + ---------- + wavefunction : + Concatenation of trial determinants of up and down spin spaces and beta + specifying the coherent state displacement. + num_elec : + Tuple of numbers of up and down spins. + num_basis : + Number of sites of Holstein chain. + verbose : + Print level + """ + def __init__(self, wavefunction: np.ndarray, num_elec: Tuple[int, int], + num_basis: int, verbose=False): self.nelec = num_elec self.nbasis = num_basis self.nalpha, self.nbeta = self.nelec @@ -23,7 +41,7 @@ def __init__(self, wavefunction, num_elec, num_basis, verbose=False): def build(self) -> None: pass - def set_etrial(self, energy) -> None: + def set_etrial(self, energy: float) -> None: self.energy = energy diff --git a/ipie/addons/eph/trial_wavefunction/toyozawa.py b/ipie/addons/eph/trial_wavefunction/toyozawa.py index 32d14c10..54b744fe 100644 --- a/ipie/addons/eph/trial_wavefunction/toyozawa.py +++ b/ipie/addons/eph/trial_wavefunction/toyozawa.py @@ -1,13 +1,17 @@ import numpy as np import scipy.linalg +from typing import Tuple +from ipie.addons.eph.hamiltonians.holstein import HolsteinModel +from ipie.addons.eph.walkers.eph_walkers import EPhWalkers from ipie.addons.eph.trial_wavefunction.coherent_state import CoherentStateTrial from ipie.utils.backend import arraylib as xp from ipie.addons.eph.trial_wavefunction.variational.toyozawa_variational import circ_perm class ToyozawaTrial(CoherentStateTrial): """""" - def __init__(self, wavefunction, hamiltonian, num_elec, num_basis, verbose=False): + def __init__(self, wavefunction: np.ndarray, hamiltonian: HolsteinModel, + num_elec: Tuple[int, int], num_basis: int, verbose=False): super().__init__(wavefunction, hamiltonian, num_elec, num_basis, verbose=verbose) self.perms = list(circ_perm([i for i in range(self.nbasis)])) self.nperms = len(self.perms) @@ -16,7 +20,7 @@ def calculate_energy(self, system, hamiltonian): #TODO variational_energy_coherent_state in ipie.estimators.local_energy pass - def calc_overlap(self, walkers) -> np.ndarray: + def calc_overlap(self, walkers: EPhWalkers) -> np.ndarray: _ = self.calc_phonon_overlap(walkers) _ = self.calc_electronic_overlap(walkers) walkers.total_ovlp = walkers.el_ovlp * walkers.ph_ovlp @@ -24,16 +28,18 @@ def calc_overlap(self, walkers) -> np.ndarray: return walkers.ovlp - def calc_phonon_overlap(self, walkers) -> np.ndarray: + def calc_phonon_overlap(self, walkers: EPhWalkers) -> np.ndarray: """""" ph_ovlp = np.zeros(walkers.nwalkers, dtype=np.complex128) for ip,perm in enumerate(self.perms): - ph_ov = np.exp(-(self.m * self.w0 / 2) * (walkers.x - self.beta_shift[perm])**2) + ph_ov = np.exp( + -(0.5 * self.m * self.w0) * (walkers.x - self.beta_shift[perm])**2 + ) walkers.ph_ovlp[:, ip] = np.prod(ph_ov, axis=1) ph_ovlp = np.sum(walkers.ph_ovlp, axis=1) return ph_ovlp - def calc_phonon_gradient(self, walkers) -> np.ndarray: + def calc_phonon_gradient(self, walkers: EPhWalkers) -> np.ndarray: r"""No reevaluation of phonon overlap because it reuses the overlap from the previous evaluation of the laplacian. The gradient only surfaces in the quantum force.""" grad = np.zeros_like(walkers.x, dtype=np.complex128) @@ -43,7 +49,7 @@ def calc_phonon_gradient(self, walkers) -> np.ndarray: grad = np.einsum('ni,n->ni', grad, 1/np.sum(walkers.ph_ovlp, axis=1)) return grad - def calc_phonon_laplacian(self, walkers, ovlps) -> np.ndarray: + def calc_phonon_laplacian(self, walkers: EPhWalkers, ovlps: np.ndarray) -> np.ndarray: r"""""" laplacian = np.zeros(walkers.nwalkers, dtype=np.complex128) for ovlp, perm in zip(ovlps.T, self.perms): @@ -53,13 +59,13 @@ def calc_phonon_laplacian(self, walkers, ovlps) -> np.ndarray: laplacian /= np.sum(ovlps, axis=1) return laplacian - def calc_phonon_laplacian_importance(self, walkers) -> np.ndarray: + def calc_phonon_laplacian_importance(self, walkers: EPhWalkers) -> np.ndarray: return self.calc_phonon_laplacian(walkers, walkers.ph_ovlp) - def calc_phonon_laplacian_locenergy(self, walkers) -> np.ndarray: + def calc_phonon_laplacian_locenergy(self, walkers: EPhWalkers) -> np.ndarray: return self.calc_phonon_laplacian(walkers, walkers.total_ovlp) - def calc_electronic_overlap(self, walkers) -> np.ndarray: + def calc_electronic_overlap(self, walkers: EPhWalkers) -> np.ndarray: """""" for ip,perm in enumerate(self.perms): ovlp_a = xp.einsum("wmi,mj->wij", walkers.phia, self.psia[perm, :].conj(), optimize=True) @@ -78,7 +84,7 @@ def calc_electronic_overlap(self, walkers) -> np.ndarray: return el_ovlp - def calc_greens_function(self, walkers, build_full=True) -> np.ndarray: + def calc_greens_function(self, walkers: EPhWalkers, build_full=True) -> np.ndarray: """""" walkers.Ga = np.zeros((walkers.nwalkers, self.nsites, self.nsites), dtype=np.complex128) walkers.Gb = np.zeros_like(walkers.Ga) diff --git a/ipie/addons/eph/walkers/eph_walkers.py b/ipie/addons/eph/walkers/eph_walkers.py index 94796e47..5a27274e 100644 --- a/ipie/addons/eph/walkers/eph_walkers.py +++ b/ipie/addons/eph/walkers/eph_walkers.py @@ -4,9 +4,9 @@ from ipie.utils.backend import arraylib as xp from ipie.utils.backend import cast_to_device, qr, qr_mode, synchronize from ipie.walkers.base_walkers import BaseWalkers -from ipie.addons.eph.trial_wavefunction.toyozawa import ToyozawaTrial +#from ipie.addons.eph.trial_wavefunction.toyozawa import ToyozawaTrial -class EphWalkers(BaseWalkers): +class EPhWalkers(BaseWalkers): """Class tailored to el-ph models where keeping track of phonon overlaps is required. Each walker carries along its Slater determinants a phonon displacement vector, self.x. @@ -77,7 +77,7 @@ def build(self, trial): Trial wavefunction object. """ - if isinstance(trial, ToyozawaTrial): + if hasattr(trial, "nperms"): shape = (self.nwalkers, trial.nperms) else: shape = self.nwalkers From 096ea3f93b96288f3446299ff3e6491931545045 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Thu, 7 Mar 2024 10:26:18 -0500 Subject: [PATCH 19/61] fix optimize position --- ipie/addons/eph/trial_wavefunction/coherent_state.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ipie/addons/eph/trial_wavefunction/coherent_state.py b/ipie/addons/eph/trial_wavefunction/coherent_state.py index 2423598a..1daa978d 100644 --- a/ipie/addons/eph/trial_wavefunction/coherent_state.py +++ b/ipie/addons/eph/trial_wavefunction/coherent_state.py @@ -113,11 +113,11 @@ def calc_greens_function(self, walkers) -> np.ndarray: walkers.G : list Greens function for each spin space """ - inv_Oa = xp.linalg.inv(xp.einsum('ie,nif->nef', self.psia, walkers.phia.conj()), optimize=True) + inv_Oa = xp.linalg.inv(xp.einsum('ie,nif->nef', self.psia, walkers.phia.conj(), optimize=True)) walkers.Ga = xp.einsum('nie,nef,jf->nji', walkers.phia, inv_Oa, self.psia.conj(), optimize=True) if self.ndown > 0: - inv_Ob = xp.linalg.inv(xp.einsum('ie,nif->nef', self.psib, walkers.phib.conj()), optimize=True) + inv_Ob = xp.linalg.inv(xp.einsum('ie,nif->nef', self.psib, walkers.phib.conj(), optimize=True)) walkers.Gb = xp.einsum('nie,nef,jf->nji', walkers.phib, inv_Ob, self.psib.conj(), optimize=True) return [walkers.Ga, walkers.Gb] From b425c448b4bbfbdd7706ab09bb52c4b5ca4fbc2e Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Thu, 7 Mar 2024 16:43:20 -0500 Subject: [PATCH 20/61] Docstrings and renaming of HolsteinPropagatorImportance --- ipie/addons/eph/propagation/holstein.py | 164 +++++++++++++++++------- ipie/propagation/propagator.py | 4 +- 2 files changed, 118 insertions(+), 50 deletions(-) diff --git a/ipie/addons/eph/propagation/holstein.py b/ipie/addons/eph/propagation/holstein.py index 24537600..69defa8f 100644 --- a/ipie/addons/eph/propagation/holstein.py +++ b/ipie/addons/eph/propagation/holstein.py @@ -10,6 +10,13 @@ def construct_one_body_propagator(hamiltonian: HolsteinModel, dt: float): """Exponentiates the electronic hopping term to apply it later as part of the trotterized algorithm. + + Parameters + ---------- + hamiltonian : + Hamiltonian caryying the one-body term as hamiltonian.T + dt : + Time step """ H1 = hamiltonian.T expH1 = numpy.array( @@ -19,7 +26,26 @@ def construct_one_body_propagator(hamiltonian: HolsteinModel, dt: float): class HolsteinPropagatorFree: - """""" + r"""Propagates walkers by trotterization, + .. math:: + \mathrm{e}^{-\Delta \tau \hat{H}} \approx \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{ph}} / 2} + \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{el}} / 2} \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{el-ph}}} + \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{el}} / 2} \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{ph}} / 2}, + + where propagation under :math:`\hat{H}_{\mathrm{ph}}` employs a generic + Diffucion MC procedure (notably without importance sampling). Propagation by + :math:`\hat{H}_{\mathrm{el}}` consists of a simple mat-vec. As + :math:`\hat{H}_{\mathrm{el-ph}}` is diagonal in bosonic position space we + can straightforwardly exponentiate the displacements and perform another + mat-vec with this diagonal matrix apllied to electronic degrees of freedom. + + Parameters + ---------- + time_step : + Time step + verbose : + Print level + """ def __init__(self, time_step: float, verbose: bool = False): self.dt = time_step self.verbose = verbose @@ -29,7 +55,22 @@ def __init__(self, time_step: float, verbose: bool = False): self.dt_ph = 0.5 * self.dt self.mpi_handler = None - def build(self, hamiltonian: HolsteinModel, trial=None, walkers=None, mpi_handler=None): + def build(self, hamiltonian: HolsteinModel, trial=None, + walkers=None, mpi_handler=None) -> None: + """Necessary step before running the AFQMC procedure. + Sets required attributes. + + Parameters + ---------- + hamiltonian : + Holstein model + trial : + Trial class + walkers : + Walkers class + mpi_handler : + MPIHandler specifying rank and size + """ self.expH1 = construct_one_body_propagator(hamiltonian, self.dt) self.const = hamiltonian.g * numpy.sqrt(2. * hamiltonian.m * hamiltonian.w0) * self.dt self.w0 = hamiltonian.w0 @@ -37,7 +78,27 @@ def build(self, hamiltonian: HolsteinModel, trial=None, walkers=None, mpi_handle self.scale = numpy.sqrt(self.dt_ph / self.m) self.nsites = hamiltonian.nsites - def propagate_phonons(self, walkers): + def propagate_phonons(self, walkers, hamiltonian, trial) -> None: + r"""Propagates phonon displacements by adjusting weigths according to + bosonic on-site energies and sampling the momentum contribution, again + by trotterizing the phonon propagator. + + .. math:: + \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{ph}} / 2} \approx + \mathrm{e}^{\Delta \tau N \omega / 4} + \mathrm{e}^{-\Delta \tau \sum_i m \omega \hat{X}_i^2 / 8} + \mathrm{e}^{-\Delta \tau \sum_i \hat{P}_i^2 / (4 \omega)} + \mathrm{e}^{-\Delta \tau \sum_i m \omega \hat{X}_i^2 / 8} + + One can obtain the sampling prescription by insertion of resolutions of + identity, :math:`\int dX |X\rangle \langleX|, and performin the resulting + Fourier transformation. + + Parameters + ---------- + walkers : + Walkers class + """ start_time = time.time() pot = 0.25 * self.m * self.w0**2 * numpy.sum(walkers.x**2, axis=1) @@ -52,12 +113,31 @@ def propagate_phonons(self, walkers): pot = numpy.real(pot) walkers.weight *= numpy.exp(-self.dt_ph * pot) - walkers.weight *= numpy.exp(self.dt_ph * self.nsites * self.w0 / 2) #doesnt matter for estimators + # Does not matter for estimators but helps with population control + walkers.weight *= numpy.exp(self.dt_ph * self.nsites * self.w0 / 2) synchronize() self.timer.tgemm += time.time() - start_time - def propagate_electron(self, walkers, trial): + def propagate_electron(self, walkers, hamiltonian, trial) -> None: + r"""Propagates electronic degrees of freedom via + + .. math:: + \mathrm{e}^{-\Delta \tau (\hat{H}_{\mathrm{el}} \otimes \hat{I}_{\mathrm{ph}} + \hat{H}_{\mathrm{el-ph}})} + \approx \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{el}} / 2} + \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{el-ph}}} + \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{el}} / 2}. + + This acts on walkers of the form :math:`|\phi(\tau)\rangle \otimes |X(\tau)\rangle`. + + + Parameters + ---------- + walkers : + Walkers class + trial : + Trial class + """ start_time = time.time() ovlp = trial.calc_greens_function(walkers) synchronize() @@ -81,15 +161,15 @@ def propagate_walkers(self, walkers, hamiltonian, trial, eshift=None): synchronize() self.timer.tgf += time.time() - start_time - # 2. Update Walkers - # 2.a DMC for phonon degrees of freedom - self.propagate_phonons(walkers) + # Update Walkers + # a) DMC for phonon degrees of freedom + self.propagate_phonons(walkers, hamiltonian, trial) - # 2.b One-body propagation for electrons - self.propagate_electron(walkers, trial) + # b) One-body propagation for electrons + self.propagate_electron(walkers, hamiltonian, trial) - # 2.c DMC for phonon degrees of freedom - self.propagate_phonons(walkers) + # c) DMC for phonon degrees of freedom + self.propagate_phonons(walkers, hamiltonian, trial) # Update weights (and later do phaseless for multi-electron) start_time = time.time() @@ -107,13 +187,34 @@ def update_weight(self, walkers, ovlp, ovlp_new): walkers.weight *= ovlp_new / ovlp -class HolsteinPropagatorImportance(HolsteinPropagatorFree): - """""" +class HolsteinPropagator(HolsteinPropagatorFree): + r"""Propagates walkers by trotterization, employing importance sampling for + the bosonic degrees of freedom. This results in a different weigth update, + and the additional displacement update by the drift term, + + .. math:: + D = \frac{\nabla_X \langle \Psi_\mathrm{T} | \psi(\tau), X(\tau)\rangle} + {\langle \Psi_\mathrm{T} | \psi(\tau), X(\tau)\rangle}, + + such that the revised displacement update reads + + .. math:: + X(\tau+\Delta\tau) = X(\tau) + + \mathcal{N}(\mu=0, \sigma = \sqrt{\frac{\Delta\tau}{m}}) + + \frac{\Delta\tau}{m} D. + + Parameters + ---------- + time_step : + Time step + verbose : + Print level + """ def __init__(self, time_step, verbose=False): super().__init__(time_step, verbose=verbose) def propagate_phonons(self, walkers, hamiltonian, trial): - """Propagates phonons via Diffusion MC.""" + """Propagates phonons via Diffusion MC including drift term.""" start_time = time.time() # No ZPE in pot -> cancels with ZPE of etrial, wouldn't affect estimators anyways @@ -142,36 +243,3 @@ def propagate_phonons(self, walkers, hamiltonian, trial): self.timer.tgemm += time.time() - start_time - def propagate_walkers(self, walkers, hamiltonian, trial, eshift=None): - """""" - synchronize() - start_time = time.time() - - ovlp = trial.calc_overlap(walkers).copy() - - synchronize() - self.timer.tgf += time.time() - start_time - - # 2. Update Walkers - # 2.a DMC for phonon degrees of freedom - self.propagate_phonons(walkers, hamiltonian, trial) - - # 2.b One-body propagation for electrons - self.propagate_electron(walkers, trial) - - # 2.c DMC for phonon degrees of freedom - self.propagate_phonons(walkers, hamiltonian, trial) - - start_time = time.time() - - ovlp_new = trial.calc_overlap(walkers) - synchronize() - self.timer.tovlp += time.time() - start_time - - start_time = time.time() - - self.update_weight(walkers, ovlp, ovlp_new) - - synchronize() - self.timer.tupdate += time.time() - start_time - diff --git a/ipie/propagation/propagator.py b/ipie/propagation/propagator.py index 587586b3..a8bc0f62 100644 --- a/ipie/propagation/propagator.py +++ b/ipie/propagation/propagator.py @@ -1,6 +1,6 @@ from ipie.hamiltonians.generic import GenericRealChol, GenericComplexChol from ipie.propagation.phaseless_generic import PhaselessGeneric -from ipie.addons.eph.propagation.holstein import HolsteinPropagatorImportance +from ipie.addons.eph.propagation.holstein import HolsteinPropagator from ipie.addons.eph.hamiltonians.holstein import HolsteinModel -Propagator = {GenericRealChol: PhaselessGeneric, GenericComplexChol: PhaselessGeneric, HolsteinModel: HolsteinPropagatorImportance} +Propagator = {GenericRealChol: PhaselessGeneric, GenericComplexChol: PhaselessGeneric, HolsteinModel: HolsteinPropagator} From 3a578a26ff162942b5ee93eac6aab98401ddb792 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Thu, 7 Mar 2024 17:16:36 -0500 Subject: [PATCH 21/61] varaitional clean up, removed comments and moved gab to estimators.py --- .../variational/coherent_state_variational.py | 51 +------------------ .../variational/estimators.py | 9 ++++ .../variational/toyozawa_variational.py | 30 +---------- 3 files changed, 12 insertions(+), 78 deletions(-) create mode 100644 ipie/addons/eph/trial_wavefunction/variational/estimators.py diff --git a/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py b/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py index d7e48b94..1159f4d7 100644 --- a/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py +++ b/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py @@ -4,6 +4,7 @@ from ipie.legacy.trial_wavefunction.harmonic_oscillator import HarmonicOscillator from ipie.systems.generic import Generic from ipie.addons.eph.hamiltonians.holstein import HolsteinModel +from ipie.addons.eph.trial_wavefunction.variational.estimators import gab import jax @@ -33,30 +34,6 @@ def local_energy( local_energy = kinetic_contrib + el_ph_contrib + phonon_contrib return local_energy -def gab(A: np.ndarray, B: np.ndarray): - inv_O = jax.numpy.linalg.inv((A.conj().T).dot(B)) - GAB = B.dot(inv_O.dot(A.conj().T)) - return GAB - -def get_T(nsites: int, t: float, pbc: bool, ndown: int) -> np.ndarray: - if nsites == 1: - return np.array([0.]) - - neighbor_map = np.zeros(shape=(nsites,nsites)) - for i in range(nsites): - neighbors = np.array([(i-1)%nsites, (i+1)%nsites]) - neighbor_map[i, neighbors] = 1. - - if not pbc: - neighbor_map[-1,0] = neighbor_map[0,-1] = 0. - T = -t * neighbor_map - - if ndown > 0: - T = [T.copy(), T.copy()] - else: - T = [T,] - return T - def objective_function(x: np.ndarray, nbasis: int, T: np.ndarray, g: float, m: float, w0: float, nup: int, ndown: int): shift = x[0:nbasis].copy() @@ -110,14 +87,10 @@ def variational_trial(init_phonons: np.ndarray, init_electron: np.ndarray, hamil x = np.hstack([init_phonons, init_electron]) - #TODO this could definitely be hamiltonian.T -# T = get_T(hamiltonian.nsites, hamiltonian.t, hamiltonian.pbc, system.ndown) - T = hamiltonian.T - maxiter = 500 minimizer_kwargs = { "jac": True, - "args": (hamiltonian.nsites, T, hamiltonian.g, hamiltonian.m, hamiltonian.w0, system.nup, system.ndown), + "args": (hamiltonian.nsites, hamiltonian.T, hamiltonian.g, hamiltonian.m, hamiltonian.w0, system.nup, system.ndown), "options": { "gtol": 1e-10, "eps": 1e-10, @@ -150,24 +123,4 @@ def variational_trial(init_phonons: np.ndarray, init_electron: np.ndarray, hamil return etrial, beta_shift, psi -if __name__ == '__main__': - - nup = 1 - ndown = 1 - system = Generic([nup, ndown]) - - g = 1. - t = 1. - w0 = 1. - nsites = 3 - pbc = True - - ham = HolsteinModel(g=g, t=t, w0=w0, nsites=nsites, pbc=pbc) - - init_phonons = np.array([0.1, 0.1, 0.1], dtype=np.float64) -# init_electron = np.array([1.,1.], dtype=np.float64) -# init_electron /= np.linalg.norm(init_electron) - init_electron = np.array([[1,0,0],[0,1,0]]) - - print(variational_trial(init_phonons, init_electron, ham, system)) diff --git a/ipie/addons/eph/trial_wavefunction/variational/estimators.py b/ipie/addons/eph/trial_wavefunction/variational/estimators.py new file mode 100644 index 00000000..1be8d8a9 --- /dev/null +++ b/ipie/addons/eph/trial_wavefunction/variational/estimators.py @@ -0,0 +1,9 @@ +import numpy as np +import jax.numpy as npj +from jax.config import config +config.update("jax_enable_x64", True) + +def gab(A,B): + inv_O = npj.linalg.inv((A.conj().T).dot(B)) + GAB = B.dot(inv_O.dot(A.conj().T)) + return GAB diff --git a/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py b/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py index 4a4bd175..1b1425f7 100644 --- a/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py +++ b/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py @@ -2,17 +2,11 @@ from scipy.optimize import minimize, basinhopping from ipie.systems import Generic from ipie.addons.eph.hamiltonians.holstein import HolsteinModel +from ipie.addons.eph.trial_wavefunction.variational.estimators import gab -from jax.config import config -config.update("jax_enable_x64", True) import jax import jax.numpy as npj -def gab(A,B): - inv_O = npj.linalg.inv((A.conj().T).dot(B)) - GAB = B.dot(inv_O.dot(A.conj().T)) - return GAB - def gradient_toyozawa_mo(x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted): grad = np.array(jax.grad(objective_function_toyozawa_mo)(x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted)) return grad @@ -142,25 +136,3 @@ def variational_trial_toyozawa(shift_init: np.ndarray, electron_init: np.ndarray return etrial, beta_shift, psi - -if __name__ == '__main__': - w0 = 1. - m = 1/w0 - g = 1. - t = 1. - nsites = 4 - pbc = True - - nup = 2 - ndown = 2 - - system = Generic([nup, ndown]) - - shift_init = np.array([2.78351996, 0.04490717, 0.04490717, 0.04490717]) - electron_init = np.random.random((nsites, nup + ndown)).astype(np.float32) - - hamiltonian = HolsteinModel(g=g, t=t, w0=w0, nsites=nsites, pbc=pbc) - hamiltonian.build() - - print(variational_trial_toyozawa(shift_init, electron_init, hamiltonian, system)) - From f79e95a3d7f4a51b4c1b22fbe01a5d9591cb0ebe Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Fri, 8 Mar 2024 11:34:40 -0500 Subject: [PATCH 22/61] docstrings and type specifications --- ipie/addons/eph/propagation/holstein.py | 49 +++++++++++++++++++------ 1 file changed, 38 insertions(+), 11 deletions(-) diff --git a/ipie/addons/eph/propagation/holstein.py b/ipie/addons/eph/propagation/holstein.py index 69defa8f..70fb4a68 100644 --- a/ipie/addons/eph/propagation/holstein.py +++ b/ipie/addons/eph/propagation/holstein.py @@ -3,6 +3,9 @@ import scipy.linalg from ipie.addons.eph.hamiltonians.holstein import HolsteinModel +from ipie.addons.eph.trial_wavefunction.eph_trial_base import EPhTrialWavefunctionBase +from ipie.addons.eph.walkers.eph_walkers import EPhWalkers + from ipie.propagation.operations import propagate_one_body from ipie.utils.backend import synchronize, cast_to_device from ipie.propagation.continuous_base import PropagatorTimer @@ -55,8 +58,9 @@ def __init__(self, time_step: float, verbose: bool = False): self.dt_ph = 0.5 * self.dt self.mpi_handler = None - def build(self, hamiltonian: HolsteinModel, trial=None, - walkers=None, mpi_handler=None) -> None: + def build(self, hamiltonian: HolsteinModel, + trial: EPhTrialWavefunctionBase = None, + walkers: EPhWalkers = None, mpi_handler = None) -> None: """Necessary step before running the AFQMC procedure. Sets required attributes. @@ -65,9 +69,9 @@ def build(self, hamiltonian: HolsteinModel, trial=None, hamiltonian : Holstein model trial : - Trial class + Trial object walkers : - Walkers class + Walkers object mpi_handler : MPIHandler specifying rank and size """ @@ -78,7 +82,9 @@ def build(self, hamiltonian: HolsteinModel, trial=None, self.scale = numpy.sqrt(self.dt_ph / self.m) self.nsites = hamiltonian.nsites - def propagate_phonons(self, walkers, hamiltonian, trial) -> None: + def propagate_phonons(self, walkers: EPhWalkers, + hamiltonian: HolsteinModel, + trial: EPhTrialWavefunctionBase) -> None: r"""Propagates phonon displacements by adjusting weigths according to bosonic on-site energies and sampling the momentum contribution, again by trotterizing the phonon propagator. @@ -119,7 +125,8 @@ def propagate_phonons(self, walkers, hamiltonian, trial) -> None: synchronize() self.timer.tgemm += time.time() - start_time - def propagate_electron(self, walkers, hamiltonian, trial) -> None: + def propagate_electron(self, walkers: EPhWalkers, hamiltonian: HolsteinModel, + trial: EPhTrialWavefunctionBase) -> None: r"""Propagates electronic degrees of freedom via .. math:: @@ -139,7 +146,6 @@ def propagate_electron(self, walkers, hamiltonian, trial) -> None: Trial class """ start_time = time.time() - ovlp = trial.calc_greens_function(walkers) synchronize() self.timer.tgf += time.time() - start_time @@ -154,7 +160,26 @@ def propagate_electron(self, walkers, hamiltonian, trial) -> None: walkers.phib = numpy.einsum('ni,nie->nie', expEph, walkers.phib) walkers.phib = propagate_one_body(walkers.phib, self.expH1[1]) - def propagate_walkers(self, walkers, hamiltonian, trial, eshift=None): + def propagate_walkers(self, + walkers: EPhWalkers, + hamiltonian: HolsteinModel, + trial: EPhTrialWavefunctionBase, + eshift: float = None + ) -> None: + r"""Propagates walkers by trotterized propagator. + + Parameters + ---------- + walkers : + EPhWalkers object + hamiltonian : + HolsteinModel object + trial : + EPhTrialWavefunctionBase object + eshift : + Only purpose is compatibility with AFQMC object, irrelevant for + propagation + """ synchronize() start_time = time.time() ovlp = trial.calc_overlap(walkers) @@ -183,7 +208,7 @@ def propagate_walkers(self, walkers, hamiltonian, trial, eshift=None): self.timer.tupdate += time.time() - start_time - def update_weight(self, walkers, ovlp, ovlp_new): + def update_weight(self, walkers, ovlp, ovlp_new) -> None: walkers.weight *= ovlp_new / ovlp @@ -213,11 +238,13 @@ class HolsteinPropagator(HolsteinPropagatorFree): def __init__(self, time_step, verbose=False): super().__init__(time_step, verbose=verbose) - def propagate_phonons(self, walkers, hamiltonian, trial): + def propagate_phonons(self, walkers: EPhWalkers, hamiltonian: HolsteinModel, + trial: EPhTrialWavefunctionBase) -> None: """Propagates phonons via Diffusion MC including drift term.""" start_time = time.time() - # No ZPE in pot -> cancels with ZPE of etrial, wouldn't affect estimators anyways + # No ZPE in pot -> cancels with ZPE of etrial, + # wouldn't affect estimators anyways ph_ovlp_old = trial.calc_phonon_overlap(walkers) pot = 0.5 * hamiltonian.m * hamiltonian.w0**2 * numpy.sum(walkers.x**2, axis=1) From 4ea65a67d3161c226a51a411222af164123c2e52 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Fri, 8 Mar 2024 11:52:42 -0500 Subject: [PATCH 23/61] added licensing header --- ipie/addons/__init__.py | 4 -- ipie/addons/eph/__init__.py | 4 -- ipie/addons/eph/estimators/__init__.py | 4 -- ipie/addons/eph/estimators/energy.py | 1 - .../eph/estimators/local_energy_holstein.py | 14 +++++ ipie/addons/eph/hamiltonians/__init__.py | 4 -- ipie/addons/eph/hamiltonians/holstein.py | 16 +++++- ipie/addons/eph/propagation/__init__.py | 4 -- ipie/addons/eph/propagation/holstein.py | 14 +++++ .../addons/eph/trial_wavefunction/__init__.py | 4 -- .../eph/trial_wavefunction/coherent_state.py | 16 +++++- .../eph/trial_wavefunction/eph_trial_base.py | 14 +++++ .../addons/eph/trial_wavefunction/toyozawa.py | 20 ++++++- .../variational/coherent_state_variational.py | 14 +++++ .../variational/estimators.py | 14 +++++ .../variational/toyozawa_variational.py | 52 +++++++++++-------- ipie/addons/eph/walkers/eph_walkers.py | 17 +++++- ipie/estimators/handler.py | 2 +- 18 files changed, 164 insertions(+), 54 deletions(-) diff --git a/ipie/addons/__init__.py b/ipie/addons/__init__.py index 871770c1..e2aed039 100644 --- a/ipie/addons/__init__.py +++ b/ipie/addons/__init__.py @@ -11,7 +11,3 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - -# Directory for additions to ipie which depend on the core ipie library. -# New features should mirror the ipie layout e.g. -# ipie/addons/finite_temperature/qmc/afqmc.py etc. diff --git a/ipie/addons/eph/__init__.py b/ipie/addons/eph/__init__.py index 871770c1..e2aed039 100644 --- a/ipie/addons/eph/__init__.py +++ b/ipie/addons/eph/__init__.py @@ -11,7 +11,3 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - -# Directory for additions to ipie which depend on the core ipie library. -# New features should mirror the ipie layout e.g. -# ipie/addons/finite_temperature/qmc/afqmc.py etc. diff --git a/ipie/addons/eph/estimators/__init__.py b/ipie/addons/eph/estimators/__init__.py index 871770c1..e2aed039 100644 --- a/ipie/addons/eph/estimators/__init__.py +++ b/ipie/addons/eph/estimators/__init__.py @@ -11,7 +11,3 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - -# Directory for additions to ipie which depend on the core ipie library. -# New features should mirror the ipie layout e.g. -# ipie/addons/finite_temperature/qmc/afqmc.py etc. diff --git a/ipie/addons/eph/estimators/energy.py b/ipie/addons/eph/estimators/energy.py index cfe7b6bc..4c190890 100644 --- a/ipie/addons/eph/estimators/energy.py +++ b/ipie/addons/eph/estimators/energy.py @@ -11,7 +11,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -# import plum diff --git a/ipie/addons/eph/estimators/local_energy_holstein.py b/ipie/addons/eph/estimators/local_energy_holstein.py index bce85d48..05402aa1 100644 --- a/ipie/addons/eph/estimators/local_energy_holstein.py +++ b/ipie/addons/eph/estimators/local_energy_holstein.py @@ -1,3 +1,17 @@ +# Copyright 2022 The ipie Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import numpy as np from ipie.addons.eph.hamiltonians.holstein import HolsteinModel diff --git a/ipie/addons/eph/hamiltonians/__init__.py b/ipie/addons/eph/hamiltonians/__init__.py index 871770c1..e2aed039 100644 --- a/ipie/addons/eph/hamiltonians/__init__.py +++ b/ipie/addons/eph/hamiltonians/__init__.py @@ -11,7 +11,3 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - -# Directory for additions to ipie which depend on the core ipie library. -# New features should mirror the ipie layout e.g. -# ipie/addons/finite_temperature/qmc/afqmc.py etc. diff --git a/ipie/addons/eph/hamiltonians/holstein.py b/ipie/addons/eph/hamiltonians/holstein.py index 2456611c..6fda9691 100644 --- a/ipie/addons/eph/hamiltonians/holstein.py +++ b/ipie/addons/eph/hamiltonians/holstein.py @@ -1,8 +1,22 @@ +# Copyright 2022 The ipie Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import numpy from ipie.utils.backend import arraylib as xp -class HolsteinModel(): +class HolsteinModel: r"""Class carrying parameters specifying a 1D Holstein chain. The Holstein model is described by the Hamiltonian diff --git a/ipie/addons/eph/propagation/__init__.py b/ipie/addons/eph/propagation/__init__.py index 871770c1..e2aed039 100644 --- a/ipie/addons/eph/propagation/__init__.py +++ b/ipie/addons/eph/propagation/__init__.py @@ -11,7 +11,3 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - -# Directory for additions to ipie which depend on the core ipie library. -# New features should mirror the ipie layout e.g. -# ipie/addons/finite_temperature/qmc/afqmc.py etc. diff --git a/ipie/addons/eph/propagation/holstein.py b/ipie/addons/eph/propagation/holstein.py index 70fb4a68..f4feed2f 100644 --- a/ipie/addons/eph/propagation/holstein.py +++ b/ipie/addons/eph/propagation/holstein.py @@ -1,3 +1,17 @@ +# Copyright 2022 The ipie Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import numpy import time import scipy.linalg diff --git a/ipie/addons/eph/trial_wavefunction/__init__.py b/ipie/addons/eph/trial_wavefunction/__init__.py index 871770c1..e2aed039 100644 --- a/ipie/addons/eph/trial_wavefunction/__init__.py +++ b/ipie/addons/eph/trial_wavefunction/__init__.py @@ -11,7 +11,3 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - -# Directory for additions to ipie which depend on the core ipie library. -# New features should mirror the ipie layout e.g. -# ipie/addons/finite_temperature/qmc/afqmc.py etc. diff --git a/ipie/addons/eph/trial_wavefunction/coherent_state.py b/ipie/addons/eph/trial_wavefunction/coherent_state.py index 1daa978d..5f91152d 100644 --- a/ipie/addons/eph/trial_wavefunction/coherent_state.py +++ b/ipie/addons/eph/trial_wavefunction/coherent_state.py @@ -1,3 +1,17 @@ +# Copyright 2022 The ipie Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import numpy as np import scipy.linalg @@ -47,7 +61,7 @@ def calculate_energy(self, system, hamiltonian): def calc_overlap(self, walkers) -> np.ndarray: _ = self.calc_phonon_overlap(walkers) _ = self.calc_electronic_overlap(walkers) - walkers.total_ovlp = walkers.el_ovlp * walkers.ph_ovlp #np.einsum('n,n->n', walkers.el_ovlp, walkers.ph_ovlp) + walkers.total_ovlp = walkers.el_ovlp * walkers.ph_ovlp walkers.ovlp = walkers.total_ovlp return walkers.ovlp diff --git a/ipie/addons/eph/trial_wavefunction/eph_trial_base.py b/ipie/addons/eph/trial_wavefunction/eph_trial_base.py index dc013cd1..124733b2 100644 --- a/ipie/addons/eph/trial_wavefunction/eph_trial_base.py +++ b/ipie/addons/eph/trial_wavefunction/eph_trial_base.py @@ -1,3 +1,17 @@ +# Copyright 2022 The ipie Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import numpy as np from ipie.trial_wavefunction.wavefunction_base import TrialWavefunctionBase from typing import Tuple diff --git a/ipie/addons/eph/trial_wavefunction/toyozawa.py b/ipie/addons/eph/trial_wavefunction/toyozawa.py index 54b744fe..38e33f23 100644 --- a/ipie/addons/eph/trial_wavefunction/toyozawa.py +++ b/ipie/addons/eph/trial_wavefunction/toyozawa.py @@ -1,3 +1,17 @@ +# Copyright 2022 The ipie Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import numpy as np import scipy.linalg from typing import Tuple @@ -13,8 +27,10 @@ class ToyozawaTrial(CoherentStateTrial): def __init__(self, wavefunction: np.ndarray, hamiltonian: HolsteinModel, num_elec: Tuple[int, int], num_basis: int, verbose=False): super().__init__(wavefunction, hamiltonian, num_elec, num_basis, verbose=verbose) - self.perms = list(circ_perm([i for i in range(self.nbasis)])) - self.nperms = len(self.perms) +# self.perms = list(circ_perm([i for i in range(self.nbasis)])) +# self.nperms = len(self.perms) + self.perms = circ_perm(np.arange(self.nbasis)) + self.nperms = self.perms.shape[0] def calculate_energy(self, system, hamiltonian): #TODO variational_energy_coherent_state in ipie.estimators.local_energy diff --git a/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py b/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py index 1159f4d7..60d591c0 100644 --- a/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py +++ b/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py @@ -1,3 +1,17 @@ +# Copyright 2022 The ipie Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import numpy as np from scipy.optimize import basinhopping diff --git a/ipie/addons/eph/trial_wavefunction/variational/estimators.py b/ipie/addons/eph/trial_wavefunction/variational/estimators.py index 1be8d8a9..6436364a 100644 --- a/ipie/addons/eph/trial_wavefunction/variational/estimators.py +++ b/ipie/addons/eph/trial_wavefunction/variational/estimators.py @@ -1,3 +1,17 @@ +# Copyright 2022 The ipie Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import numpy as np import jax.numpy as npj from jax.config import config diff --git a/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py b/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py index 1b1425f7..c872e3ab 100644 --- a/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py +++ b/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py @@ -1,3 +1,17 @@ +# Copyright 2022 The ipie Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import numpy as np from scipy.optimize import minimize, basinhopping from ipie.systems import Generic @@ -71,35 +85,30 @@ def objective_function_toyozawa_mo(x, nbasis, nup, ndown, T, U, g, m, w0, perms, etot = num_energy / denom return etot.real -def circ_perm(lst): - cpy = lst[:] - yield cpy - for i in range(len(lst) - 1): - cpy = cpy[1:] + [cpy[0]] - yield cpy - -def func(x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted): - f = objective_function_mo(x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted) - df = gradient_mo(x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted) - return f, df - -def print_fun(x: np.ndarray, f: float, accepted: bool): - print("at minimum %.4f accepted %d" % (f, int(accepted))) - -def func_toyo(x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted): - f = objective_function_toyozawa_mo(x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted) - df = gradient_toyozawa_mo(x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted) - return f, df +def circ_perm(lst: np.ndarray) -> np.ndarray: + """Returns a matrix which rows consist of all possible + cyclic permutations given an initial array lst. + + Parameters + ---------- + lst : + Initial array which is to be cyclically permuted + """ + circs = lst + for shift in range(1, len(lst)): + new_circ = np.roll(lst, -shift) + circs = np.vstack([circs, new_circ]) + return circs def variational_trial_toyozawa(shift_init: np.ndarray, electron_init: np.ndarray, hamiltonian, system): psi = electron_init.T.real.ravel() shift = shift_init.real - perms = list(circ_perm([i for i in range(hamiltonian.nsites)])) + perms = circ_perm(np.arange(hamiltonian.nsites)) x = np.zeros((system.nup + system.ndown + 1) * hamiltonian.nsites) x[:hamiltonian.nsites] = shift.copy() - x[hamiltonian.nsites:] = psi.copy() #[:,0] + x[hamiltonian.nsites:] = psi.copy() res = minimize( objective_function_toyozawa_mo, @@ -134,5 +143,4 @@ def variational_trial_toyozawa(shift_init: np.ndarray, electron_init: np.ndarray psia = res.x[hamiltonian.nsites:].reshape((hamiltonian.nsites, system.nup)) psi = psia - return etrial, beta_shift, psi diff --git a/ipie/addons/eph/walkers/eph_walkers.py b/ipie/addons/eph/walkers/eph_walkers.py index 5a27274e..91aef5b9 100644 --- a/ipie/addons/eph/walkers/eph_walkers.py +++ b/ipie/addons/eph/walkers/eph_walkers.py @@ -1,10 +1,23 @@ +# Copyright 2022 The ipie Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import numpy from ipie.config import config from ipie.utils.backend import arraylib as xp from ipie.utils.backend import cast_to_device, qr, qr_mode, synchronize from ipie.walkers.base_walkers import BaseWalkers -#from ipie.addons.eph.trial_wavefunction.toyozawa import ToyozawaTrial class EPhWalkers(BaseWalkers): """Class tailored to el-ph models where keeping track of phonon overlaps is @@ -73,7 +86,7 @@ def build(self, trial): Parameters ---------- - trial : class + trial : Trial wavefunction object. """ diff --git a/ipie/estimators/handler.py b/ipie/estimators/handler.py index a35dbfe1..c65f309f 100644 --- a/ipie/estimators/handler.py +++ b/ipie/estimators/handler.py @@ -233,7 +233,7 @@ def print_block(self, comm, block, walker_factors, div_factor=None): shift = None walker_factors.eshift = comm.bcast(shift) if comm.rank == 0: - self.output.push_to_chunk(self.global_estimates, f"data") +# self.output.push_to_chunk(self.global_estimates, f"data") self.output.increment() if comm.rank == 0: print(f"{block:>17d} " + output_string) From 7eec46beaccd1020d68325310f752cf617160b14 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Fri, 8 Mar 2024 12:12:32 -0500 Subject: [PATCH 24/61] renamed phonon displacements stored in EPhWalkers.x to EPhWalkers.phonon_disp --- .../eph/estimators/local_energy_holstein.py | 17 ++++++++++++++--- ipie/addons/eph/propagation/holstein.py | 14 +++++++------- .../eph/trial_wavefunction/coherent_state.py | 6 +++--- ipie/addons/eph/trial_wavefunction/toyozawa.py | 8 ++++---- ipie/addons/eph/walkers/eph_walkers.py | 8 ++++---- 5 files changed, 32 insertions(+), 21 deletions(-) diff --git a/ipie/addons/eph/estimators/local_energy_holstein.py b/ipie/addons/eph/estimators/local_energy_holstein.py index 05402aa1..81411e0f 100644 --- a/ipie/addons/eph/estimators/local_energy_holstein.py +++ b/ipie/addons/eph/estimators/local_energy_holstein.py @@ -38,6 +38,17 @@ def local_energy_holstein( |\phi(\tau(r))\rangle \otimes |\beta(\tau(X))\rangle`. In this ansatz for the walkers :math:`|\beta\rangle` is a coherent state, which corresonds to a by :math:`\beta` displaced vacuum state. + + Parameters + ---------- + system : + Generic object carrying information on number and spin of electrons + hamiltonian : + HolsteinModel object + walkers : + EPhWalkers object + trial : + EPhTrialWavefunctionBase object """ energy = xp.zeros((walkers.nwalkers, 4), dtype=xp.complex128) @@ -48,12 +59,12 @@ def local_energy_holstein( if system.ndown > 0: energy[:, 1] += np.sum(hamiltonian.T[1] * gf[1], axis=(1,2)) - energy[:, 2] = np.sum(np.diagonal(gf[0], axis1=1, axis2=2) * walkers.x, axis=1) + energy[:, 2] = np.sum(np.diagonal(gf[0], axis1=1, axis2=2) * walkers.phonon_disp, axis=1) if system.ndown > 0: - energy[:, 2] += np.sum(np.diagonal(gf[1], axis1=1, axis2=2) * walkers.x, axis=1) + energy[:, 2] += np.sum(np.diagonal(gf[1], axis1=1, axis2=2) * walkers.phonon_disp, axis=1) energy[:, 2] *= hamiltonian.const - energy[:, 3] = 0.5 * hamiltonian.m * hamiltonian.w0**2 * np.sum(walkers.x**2, axis=1) + energy[:, 3] = 0.5 * hamiltonian.m * hamiltonian.w0**2 * np.sum(walkers.phonon_disp**2, axis=1) energy[:, 3] -= 0.5 * hamiltonian.nsites * hamiltonian.w0 energy[:, 3] -= 0.5 * trial.calc_phonon_laplacian_locenergy(walkers) / hamiltonian.m diff --git a/ipie/addons/eph/propagation/holstein.py b/ipie/addons/eph/propagation/holstein.py index f4feed2f..8ae45c54 100644 --- a/ipie/addons/eph/propagation/holstein.py +++ b/ipie/addons/eph/propagation/holstein.py @@ -121,15 +121,15 @@ def propagate_phonons(self, walkers: EPhWalkers, """ start_time = time.time() - pot = 0.25 * self.m * self.w0**2 * numpy.sum(walkers.x**2, axis=1) + pot = 0.25 * self.m * self.w0**2 * numpy.sum(walkers.phonon_disp**2, axis=1) pot = numpy.real(pot) walkers.weight *= numpy.exp(-self.dt_ph * pot) N = numpy.random.normal(loc=0.0, scale=self.scale, size=(walkers.nwalkers, self.nsites)) - walkers.x = walkers.x + N + walkers.phonon_disp = walkers.phonon_disp + N - pot = 0.25 * self.m * self.w0**2 * numpy.sum(walkers.x**2, axis=1) + pot = 0.25 * self.m * self.w0**2 * numpy.sum(walkers.phonon_disp**2, axis=1) pot = numpy.real(pot) walkers.weight *= numpy.exp(-self.dt_ph * pot) @@ -163,7 +163,7 @@ def propagate_electron(self, walkers: EPhWalkers, hamiltonian: HolsteinModel, synchronize() self.timer.tgf += time.time() - start_time - expEph = numpy.exp(self.const * walkers.x) + expEph = numpy.exp(self.const * walkers.phonon_disp) walkers.phia = propagate_one_body(walkers.phia, self.expH1[0]) walkers.phia = numpy.einsum('ni,nie->nie', expEph, walkers.phia) @@ -261,18 +261,18 @@ def propagate_phonons(self, walkers: EPhWalkers, hamiltonian: HolsteinModel, # wouldn't affect estimators anyways ph_ovlp_old = trial.calc_phonon_overlap(walkers) - pot = 0.5 * hamiltonian.m * hamiltonian.w0**2 * numpy.sum(walkers.x**2, axis=1) + pot = 0.5 * hamiltonian.m * hamiltonian.w0**2 * numpy.sum(walkers.phonon_disp**2, axis=1) pot -= 0.5 * trial.calc_phonon_laplacian_importance(walkers) / hamiltonian.m pot = numpy.real(pot) walkers.weight *= numpy.exp(-self.dt_ph * pot / 2) N = numpy.random.normal(loc=0.0, scale=self.scale, size=(walkers.nwalkers, self.nsites)) drift = trial.calc_phonon_gradient(walkers) - walkers.x = walkers.x + N + self.dt_ph * drift / hamiltonian.m + walkers.phonon_disp = walkers.phonon_disp + N + self.dt_ph * drift / hamiltonian.m ph_ovlp_new = trial.calc_phonon_overlap(walkers) - pot = 0.5 * hamiltonian.m * hamiltonian.w0**2 * numpy.sum(walkers.x**2, axis=1) + pot = 0.5 * hamiltonian.m * hamiltonian.w0**2 * numpy.sum(walkers.phonon_disp**2, axis=1) pot -= 0.5 * trial.calc_phonon_laplacian_importance(walkers) / hamiltonian.m pot = numpy.real(pot) walkers.weight *= numpy.exp(-self.dt_ph * pot / 2) diff --git a/ipie/addons/eph/trial_wavefunction/coherent_state.py b/ipie/addons/eph/trial_wavefunction/coherent_state.py index 5f91152d..8ff33d16 100644 --- a/ipie/addons/eph/trial_wavefunction/coherent_state.py +++ b/ipie/addons/eph/trial_wavefunction/coherent_state.py @@ -66,17 +66,17 @@ def calc_overlap(self, walkers) -> np.ndarray: return walkers.ovlp def calc_phonon_overlap(self, walkers) -> np.ndarray: - ph_ovlp = np.exp(-(self.m * self.w0 / 2) * (walkers.x - self.beta_shift)**2) + ph_ovlp = np.exp(-(self.m * self.w0 / 2) * (walkers.phonon_disp - self.beta_shift)**2) walkers.ph_ovlp = np.prod(ph_ovlp, axis=1) return walkers.ph_ovlp def calc_phonon_gradient(self, walkers) -> np.ndarray: - grad = walkers.x - self.beta_shift + grad = walkers.phonon_disp - self.beta_shift grad *= -self.m * self.w0 return grad def calc_phonon_laplacian(self, walkers) -> np.ndarray: - arg = (walkers.x - self.beta_shift) * self.m * self.w0 + arg = (walkers.phonon_disp - self.beta_shift) * self.m * self.w0 arg2 = arg**2 laplacian = np.sum(arg2, axis=1) - self.nsites * self.m * self.w0 return laplacian diff --git a/ipie/addons/eph/trial_wavefunction/toyozawa.py b/ipie/addons/eph/trial_wavefunction/toyozawa.py index 38e33f23..490b3606 100644 --- a/ipie/addons/eph/trial_wavefunction/toyozawa.py +++ b/ipie/addons/eph/trial_wavefunction/toyozawa.py @@ -49,7 +49,7 @@ def calc_phonon_overlap(self, walkers: EPhWalkers) -> np.ndarray: ph_ovlp = np.zeros(walkers.nwalkers, dtype=np.complex128) for ip,perm in enumerate(self.perms): ph_ov = np.exp( - -(0.5 * self.m * self.w0) * (walkers.x - self.beta_shift[perm])**2 + -(0.5 * self.m * self.w0) * (walkers.phonon_disp - self.beta_shift[perm])**2 ) walkers.ph_ovlp[:, ip] = np.prod(ph_ov, axis=1) ph_ovlp = np.sum(walkers.ph_ovlp, axis=1) @@ -58,9 +58,9 @@ def calc_phonon_overlap(self, walkers: EPhWalkers) -> np.ndarray: def calc_phonon_gradient(self, walkers: EPhWalkers) -> np.ndarray: r"""No reevaluation of phonon overlap because it reuses the overlap from the previous evaluation of the laplacian. The gradient only surfaces in the quantum force.""" - grad = np.zeros_like(walkers.x, dtype=np.complex128) + grad = np.zeros_like(walkers.phonon_disp, dtype=np.complex128) for ovlp, perm in zip(walkers.ph_ovlp.T, self.perms): - grad += np.einsum('ni,n->ni', (walkers.x - self.beta_shift[perm]), ovlp) + grad += np.einsum('ni,n->ni', (walkers.phonon_disp - self.beta_shift[perm]), ovlp) grad *= -self.m * self.w0 grad = np.einsum('ni,n->ni', grad, 1/np.sum(walkers.ph_ovlp, axis=1)) return grad @@ -69,7 +69,7 @@ def calc_phonon_laplacian(self, walkers: EPhWalkers, ovlps: np.ndarray) -> np.nd r"""""" laplacian = np.zeros(walkers.nwalkers, dtype=np.complex128) for ovlp, perm in zip(ovlps.T, self.perms): - arg = (walkers.x - self.beta_shift[perm]) * self.m * self.w0 + arg = (walkers.phonon_disp - self.beta_shift[perm]) * self.m * self.w0 arg2 = arg**2 laplacian += (np.sum(arg2, axis=1) - self.nsites * self.m * self.w0) * ovlp laplacian /= np.sum(ovlps, axis=1) diff --git a/ipie/addons/eph/walkers/eph_walkers.py b/ipie/addons/eph/walkers/eph_walkers.py index 91aef5b9..b8d60131 100644 --- a/ipie/addons/eph/walkers/eph_walkers.py +++ b/ipie/addons/eph/walkers/eph_walkers.py @@ -22,7 +22,7 @@ class EPhWalkers(BaseWalkers): """Class tailored to el-ph models where keeping track of phonon overlaps is required. Each walker carries along its Slater determinants a phonon - displacement vector, self.x. + displacement vector, self.phonon_disp. Parameters ---------- @@ -59,11 +59,11 @@ def __init__( self.weight = numpy.ones(self.nwalkers, dtype=numpy.complex128) - self.x = xp.array( + self.phonon_disp = xp.array( [initial_walker[:,0].copy() for iw in range(self.nwalkers)], dtype=xp.complex128 ) - self.x = numpy.squeeze(self.x) + self.phonon_disp = numpy.squeeze(self.phonon_disp) self.phia = xp.array( [initial_walker[:, 1 : self.nup+1].copy() for iw in range(self.nwalkers)], @@ -75,7 +75,7 @@ def __init__( dtype=xp.complex128, ) - self.buff_names += ["phia", "phib", "x"] + self.buff_names += ["phia", "phib", "phonon_disp"] self.buff_size = round(self.set_buff_size_single_walker() / float(self.nwalkers)) self.walker_buffer = numpy.zeros(self.buff_size, dtype=numpy.complex128) From c01760f59b6005b56b846565fe605f83ba31701a Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Sun, 10 Mar 2024 18:37:35 -0400 Subject: [PATCH 25/61] docstrings with math --- .../addons/eph/trial_wavefunction/toyozawa.py | 267 ++++++++++++++++-- 1 file changed, 236 insertions(+), 31 deletions(-) diff --git a/ipie/addons/eph/trial_wavefunction/toyozawa.py b/ipie/addons/eph/trial_wavefunction/toyozawa.py index 490b3606..7d614808 100644 --- a/ipie/addons/eph/trial_wavefunction/toyozawa.py +++ b/ipie/addons/eph/trial_wavefunction/toyozawa.py @@ -23,41 +23,150 @@ from ipie.addons.eph.trial_wavefunction.variational.toyozawa_variational import circ_perm class ToyozawaTrial(CoherentStateTrial): - """""" + r"""The Toyozawa trial + + .. math:: + |\Psi(\kappa)\rangle = \sum_n e^{i \kappa n} \sum_{n_1} \alpha_{n_1}^{\kappa} + a_{n_1}^{\dagger} \exp(-\sum_{n_2} (\beta^\kappa_{n_2 - n} b_{n_2}^{\dagger} + - \beta^{\kappa^*}_{n_2 - n} b_{n_2}))|0\rangle + + developed by `Toyozawa `_ is translationally + invariant and reliable offers a good approximation to the polaron ground state + for most parameter regimes of the Holstein Model. Here, :math:`\alpha,\beta`are + varaitional parameters, and :math:`|0\rangle` is the total vacuum state. + For a 1D Holstein chain this reduces to a superposition of cyclically `CoherentState` + type trials. + More details may be found in `Zhao et al. `_. + + Attributes + ---------- + perms : :class:`np.ndarray` + Rows of this matrix corresponds to cyclic permutations of `range(nsites)` + nperms : :class:`int` + Number of permutations in `perms` + """ def __init__(self, wavefunction: np.ndarray, hamiltonian: HolsteinModel, num_elec: Tuple[int, int], num_basis: int, verbose=False): super().__init__(wavefunction, hamiltonian, num_elec, num_basis, verbose=verbose) -# self.perms = list(circ_perm([i for i in range(self.nbasis)])) -# self.nperms = len(self.perms) self.perms = circ_perm(np.arange(self.nbasis)) self.nperms = self.perms.shape[0] - def calculate_energy(self, system, hamiltonian): - #TODO variational_energy_coherent_state in ipie.estimators.local_energy - pass + def calc_overlap_perm(self, walkers: EPhWalkers) -> np.ndarray: + r"""Computes the product of electron and phonon overlaps for each + permutation :math:`\sigma`, + + .. math:: + \langle \psi_T(\sigma(r))|\psi(\tau)\rangle + \langle \phi(\sigma(\beta))|X_{\mathrm{w}(\tau)}\rangle. + + Parameters + ---------- + walkers : :class:`EPhWalkers` + EPhWalkers object + + Returns + ------- + total_ovlp : :class:`np.ndarray` + Product of electron and phonon overlap + """ + ph_ovlp_perm = self.calc_phonon_overlap_perms(walkers) + el_ovlp_perm = self.calc_electronic_overlap_perms(walkers) + walkers.total_ovlp = el_ovlp_perm * ph_ovlp_perm + return walkers.total_ovlp def calc_overlap(self, walkers: EPhWalkers) -> np.ndarray: - _ = self.calc_phonon_overlap(walkers) - _ = self.calc_electronic_overlap(walkers) - walkers.total_ovlp = walkers.el_ovlp * walkers.ph_ovlp - walkers.ovlp = np.sum(walkers.total_ovlp, axis=1) - return walkers.ovlp + r"""Sums product of electronic and phonon overlap for each permutation + over all permutations, + + .. math:: + \sum_\tau \langle \psi_T(\sigma(r))|\psi(\tau)\rangle + \langle \phi(\sigma(\beta))|X_{\mathrm{w}(\tau)}\rangle. + Used when evaluating local energy and when updating + weight. + + Parameters + ---------- + walkers : :class:`EPhWalkers` + EPhWalkers object - def calc_phonon_overlap(self, walkers: EPhWalkers) -> np.ndarray: - """""" - ph_ovlp = np.zeros(walkers.nwalkers, dtype=np.complex128) + Returns + ------- + ovlp: :class:`np.ndarray` + Sum of product of electron and phonon overlap + """ + total_ovlp = self.calc_overlap_perm(walkers) + ovlp = np.sum(total_ovlp, axis=1) + return ovlp + + def calc_phonon_overlap_perms(self, walkers: EPhWalkers) -> np.ndarray: + r"""Updates the walker phonon overlap with each permutation :math:`\tau`, + i.e. :math:`\langle\phi(\tau(\beta))|X_{\mathrm{w}}\rangle` and stores + it in `walkers.ph_ovlp`. + + Parameters + ---------- + walkers : :class:`EPhWalkers` + EPhWalkers object + + Returns + ------- + ph_ovlp_perm : :class:`np.ndarray` + Overlap of walker with permuted coherent states + """ for ip,perm in enumerate(self.perms): ph_ov = np.exp( -(0.5 * self.m * self.w0) * (walkers.phonon_disp - self.beta_shift[perm])**2 ) walkers.ph_ovlp[:, ip] = np.prod(ph_ov, axis=1) - ph_ovlp = np.sum(walkers.ph_ovlp, axis=1) + return walkers.ph_ovlp + + def calc_phonon_overlap(self, walkers: EPhWalkers) -> np.ndarray: + r"""Sums walker phonon overlaps with permuted coherent states over all + permutations, + + .. math:: + \sum_\tau \langle \phi(\tau(\beta)) | X_{\mathrm{w}} \rangle + + to get total phonon overlap. This is only used to correct + for the importance sampling in propagate_phonons. + + Parameters + ---------- + walkers : :class:`EPhWalkers` + EPhWalkers object + + Returns + ------- + ph_ovlp : :class:`np.ndarray` + Total walker phonon overlap + """ + ph_ovlp_perm = self.calc_phonon_overlap_perms(walkers) + ph_ovlp = np.sum(ph_ovlp_perm, axis=1) return ph_ovlp def calc_phonon_gradient(self, walkers: EPhWalkers) -> np.ndarray: - r"""No reevaluation of phonon overlap because it reuses the overlap from the previous - evaluation of the laplacian. The gradient only surfaces in the quantum force.""" + r"""Computes the phonon gradient, + + .. math:: + \sum_\sigma \frac{\nabla_X \langle \phi(\sigma(\beta)) | X(\tau) \rangle} + {\rangle \phi(\sigma(\beta)) | X(\tau) \rangle} + = \sum_\sigma -m \omega \frac{(X(\tau) - \sigma(\beta)) * \langle\phi(\sigma(\beta))|X(\tau)\rangle} + {\sum_\simga \langle\phi(\sigma(\beta))|X(\tau)\rangle}. + + This is only used when calculating the drift term for the importance + sampling DMC part of the algorithm. + + Parameters + ---------- + walkers : :class:`EPhWalkers` + EPhWalkers object + + Returns + ------- + grad : :class:`np.ndarray` + Phonon gradient + """ grad = np.zeros_like(walkers.phonon_disp, dtype=np.complex128) for ovlp, perm in zip(walkers.ph_ovlp.T, self.perms): grad += np.einsum('ni,n->ni', (walkers.phonon_disp - self.beta_shift[perm]), ovlp) @@ -66,7 +175,27 @@ def calc_phonon_gradient(self, walkers: EPhWalkers) -> np.ndarray: return grad def calc_phonon_laplacian(self, walkers: EPhWalkers, ovlps: np.ndarray) -> np.ndarray: - r"""""" + r"""Computes the phonon Laplacian, which weights coherent state laplacians + by overlaps :math:`o(\sigma, r, X, \tau)` passed to this function, + + .. math:: + \sum_\sigma \frac{\nabla_X \langle \phi(\sigma(\beta)) | X(\tau) \rangle} + {\rangle \phi(\sigma(\beta)) | X(\tau) \rangle} + = \frac{\sum_sigma ((\sum_i (m \omega (X_i(\tau) - \sigma(\beta)_i))^2) - N m \omega) o(\sigma, r, X, \tau)} + {\sum_\sigma o(\sigma, r, X, \tau)}. + + Parameters + ---------- + walkers : :class:`EPhWalkers` + EPhWalkers object + ovlps : :class:`np.ndarray` + Overlaps weighting contributions from permuted coherent states + + Returns + ------- + laplacian : :class:`np.ndarray` + Phonon Laplacian + """ laplacian = np.zeros(walkers.nwalkers, dtype=np.complex128) for ovlp, perm in zip(ovlps.T, self.perms): arg = (walkers.phonon_disp - self.beta_shift[perm]) * self.m * self.w0 @@ -76,13 +205,57 @@ def calc_phonon_laplacian(self, walkers: EPhWalkers, ovlps: np.ndarray) -> np.nd return laplacian def calc_phonon_laplacian_importance(self, walkers: EPhWalkers) -> np.ndarray: + r"""Computes phonon Laplacian via `calc_phonon_laplacian` with weighting + by pure phonon overlap. This is only utilized in the importance sampling + of the DMC procedure. + + Parameters + ---------- + walkers : :class:`EPhWalkers` + EPhWalkers object + + Returns + ------- + ph_lapl : :class:`np.ndarray` + Phonon Laplacian weigthed by phonon overlaps + """ return self.calc_phonon_laplacian(walkers, walkers.ph_ovlp) def calc_phonon_laplacian_locenergy(self, walkers: EPhWalkers) -> np.ndarray: + """Computes phonon Laplacian using total overlap weights as required in + local energy evaluation. + + Parameters + ---------- + walkers : :class:`EPhWalkers` + EPhWalkers object + + Returns + ------- + ph_lapl : :class:`np.ndarray` + Phonon Laplacian weigthed by total overlaps + """ return self.calc_phonon_laplacian(walkers, walkers.total_ovlp) - def calc_electronic_overlap(self, walkers: EPhWalkers) -> np.ndarray: - """""" + def calc_electronic_overlap_perms(self, walkers: EPhWalkers) -> np.ndarray: + r"""Calculates the electronic overlap of each walker with each permuted + Slater determinant :math:`|\Phi_T(\tau(r_i))\rangle` of the trial, + + .. math:: + \langle \Phi_T(\tau(r_i))|\psi_w\rangle = \mathrm{det(U^{\dagger}V)}, + + where :math:`U,V` parametrized the two Slater determinants. + + Parameters + ---------- + walkers : :class:`EPhWalkers` + EPhWalkers object + + Returns + ------- + el_ovlp_perm : :class:`np.ndarray` + Electronic overlap of each permuted Slater Determiant with walkers + """ for ip,perm in enumerate(self.perms): ovlp_a = xp.einsum("wmi,mj->wij", walkers.phia, self.psia[perm, :].conj(), optimize=True) sign_a, log_ovlp_a = xp.linalg.slogdet(ovlp_a) @@ -95,27 +268,59 @@ def calc_electronic_overlap(self, walkers: EPhWalkers) -> np.ndarray: ot = sign_a * xp.exp(log_ovlp_a - walkers.log_shift) walkers.el_ovlp[:, ip] = ot + return walkers.el_ovlp + + def calc_electronic_overlap(self, walkers: EPhWalkers) -> np.ndarray: + """Sums walkers.el_ovlp over permutations to obtain total electronic + overlap of trial with walkers. - el_ovlp = np.sum(walkers.el_ovlp, axis=1) - + Parameters + ---------- + walkers : :class:`EPhWalkers` + EPhWalkers object + + Returns + ------- + el_ovlp : :class:`np.ndarray` + Electronic overlap of trial with walkers + """ + el_ovlp_perms = self.calc_electronic_overlap_perms(walkers) + el_ovlp = np.sum(el_ovlp_perms, axis=1) return el_ovlp def calc_greens_function(self, walkers: EPhWalkers, build_full=True) -> np.ndarray: - """""" - walkers.Ga = np.zeros((walkers.nwalkers, self.nsites, self.nsites), dtype=np.complex128) - walkers.Gb = np.zeros_like(walkers.Ga) + r"""Calculates Greens functions by + + .. math:: + G^{\Phi \Psi}_{p\alpha, q\beta} + = \frac{\sum_{\tau} \delta_{\alpha\beta}(U_\alpha(V^{\dagger}_\alhpa(\tau) U_\alpha) V^{\dagger}_\alpha(\tau)) \langle\Phi_T(\tau(r_i))|\psi_w\rangle} + {\sum_{\tau} \langle\Phi_T(\tau(r_i))|\psi_w\rangle} + + Parameters + ---------- + walkers : :class:`EPhWalkers` + EPhWalkers object + + Returns + ------- + G : :class:`list` + List of Greens functions for :math:`\alpha,\beta` spin spaces. + """ + Ga = np.zeros((walkers.nwalkers, self.nsites, self.nsites), dtype=np.complex128) + Gb = np.zeros_like(Ga) for ovlp, perm in zip(walkers.total_ovlp.T, self.perms): inv_Oa = xp.linalg.inv(xp.einsum('ie,nif->nef', self.psia[perm,:], walkers.phia.conj())) - walkers.Ga += xp.einsum('nie,nef,jf,n->nji', walkers.phia, inv_Oa, self.psia[perm].conj(), ovlp) + Ga += xp.einsum('nie,nef,jf,n->nji', walkers.phia, inv_Oa, self.psia[perm].conj(), ovlp) if self.ndown > 0: inv_Ob = xp.linalg.inv(xp.einsum('ie,nif->nef', self.psib[perm,:], walkers.phib.conj())) - walkers.Gb += xp.einsum('nie,nef,jf,n->nji', walkers.phib, inv_Ob, self.psib[perm].conj(), ovlp) - - walkers.Ga = np.einsum('nij,n->nij', walkers.Ga, 1 / np.sum(walkers.total_ovlp, axis=1)) - walkers.Gb = np.einsum('nij,n->nij', walkers.Gb, 1 / np.sum(walkers.total_ovlp, axis=1)) + Gb += xp.einsum('nie,nef,jf,n->nji', walkers.phib, inv_Ob, self.psib[perm].conj(), ovlp) - return [walkers.Ga, walkers.Gb] + Ga = np.einsum('nij,n->nij', Ga, 1 / walkers.ovlp) + if self.ndown > 0: + Gb = np.einsum('nij,n->nij', Gb, 1 / walkers.ovlp) + + return [Ga, Gb] From e59ead634eae803f74f0d03c00e96909b6596155 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Sun, 10 Mar 2024 18:39:00 -0400 Subject: [PATCH 26/61] removed redundant calculation of greens function --- ipie/addons/eph/estimators/energy.py | 1 - 1 file changed, 1 deletion(-) diff --git a/ipie/addons/eph/estimators/energy.py b/ipie/addons/eph/estimators/energy.py index 4c190890..74ab6c18 100644 --- a/ipie/addons/eph/estimators/energy.py +++ b/ipie/addons/eph/estimators/energy.py @@ -59,7 +59,6 @@ def __init__( self.ascii_filename = filename def compute_estimator(self, system, walkers, hamiltonian, trial, istep=1): - trial.calc_greens_function(walkers) # Need to be able to dispatch here energy = local_energy(system, hamiltonian, walkers, trial) self._data["ENumer"] = xp.sum(walkers.weight * energy[:, 0].real) From a3c3675f6ee23f801e4201d09c4af2c3c7ea70dd Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Sun, 10 Mar 2024 18:39:48 -0400 Subject: [PATCH 27/61] added abstract functions, required for every new trial --- .../eph/trial_wavefunction/eph_trial_base.py | 48 +++++++++++++++---- 1 file changed, 39 insertions(+), 9 deletions(-) diff --git a/ipie/addons/eph/trial_wavefunction/eph_trial_base.py b/ipie/addons/eph/trial_wavefunction/eph_trial_base.py index 124733b2..81d8c475 100644 --- a/ipie/addons/eph/trial_wavefunction/eph_trial_base.py +++ b/ipie/addons/eph/trial_wavefunction/eph_trial_base.py @@ -13,13 +13,10 @@ # limitations under the License. import numpy as np -from ipie.trial_wavefunction.wavefunction_base import TrialWavefunctionBase from typing import Tuple +from abc import ABCMeta, abstractmethod -# NOTE could inherit from TrialWavefunctionBase, -# but would need to redefine abstract methods.. - -class EPhTrialWavefunctionBase: +class EPhTrialWavefunctionBase(metaclass=ABCMeta): """Base class for electron-phonon trial wave functions. Parameters @@ -48,15 +45,48 @@ def __init__(self, wavefunction: np.ndarray, num_elec: Tuple[int, int], self.psia = wavefunction[:self.nalpha] self.psib = wavefunction[self.nalpha:self.nalpha+self.nbeta] self.beta_shift = wavefunction[self.nalpha+self.nbeta:] - + self.compute_trial_energy = False self.energy = None - def build(self) -> None: - pass - def set_etrial(self, energy: float) -> None: self.energy = energy + # TODO This should be abstract method as well + def calculate_energy(self, system, hamiltonian): + # TODO variational_energy_coherent_state in ipie.estimators.local_energy + ... + + @abstractmethod + def calc_overlap(self, walkers) -> np.ndarray: + ... + + @abstractmethod + def calc_phonon_overlap(self, walkers) -> np.ndarray: + ... + + @abstractmethod + def calc_phonon_gradient(self, walkers) -> np.ndarray: + ... + + @abstractmethod + def calc_phonon_laplacian(self, walkers) -> np.ndarray: + ... + + @abstractmethod + def calc_phonon_laplacian_importance(self, walkers) -> np.ndarray: + ... + + @abstractmethod + def calc_phonon_laplacian_locenergy(self, walkers) -> np.ndarray: + ... + + @abstractmethod + def calc_electronic_overlap(self, walkers) -> np.ndarray: + ... + + @abstractmethod + def calc_greens_function(self, walkers) -> np.ndarray: + ... From 317dcf30367da7b2f02841850f55f91ba49b8b8e Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Sun, 10 Mar 2024 18:45:09 -0400 Subject: [PATCH 28/61] intializes greens functions new and sets overlap in build --- ipie/addons/eph/walkers/eph_walkers.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/ipie/addons/eph/walkers/eph_walkers.py b/ipie/addons/eph/walkers/eph_walkers.py index b8d60131..3baa78cd 100644 --- a/ipie/addons/eph/walkers/eph_walkers.py +++ b/ipie/addons/eph/walkers/eph_walkers.py @@ -103,7 +103,10 @@ def build(self, trial): self.buff_size = round(self.set_buff_size_single_walker() / float(self.nwalkers)) self.walker_buffer = numpy.zeros(self.buff_size, dtype=numpy.complex128) - trial.calc_overlap(self) + self.Ga = numpy.zeros((self.nwalkers, self.nbasis, self.nbasis), dtype=numpy.complex128) + self.Gb = numpy.zeros_like(self.Ga) + + self.ovlp = trial.calc_overlap(self) def cast_to_cupy(self, verbose=False): cast_to_device(self, verbose) From ee1470850ff9ec12c3538d51cdfc782504e31519 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Mon, 11 Mar 2024 00:19:28 -0400 Subject: [PATCH 29/61] extra layer for propagators from addons --- ipie/addons/propagator.py | 4 ++++ ipie/propagation/propagator.py | 6 +++--- 2 files changed, 7 insertions(+), 3 deletions(-) create mode 100644 ipie/addons/propagator.py diff --git a/ipie/addons/propagator.py b/ipie/addons/propagator.py new file mode 100644 index 00000000..2f0a14ca --- /dev/null +++ b/ipie/addons/propagator.py @@ -0,0 +1,4 @@ +from ipie.addons.eph.propagation.holstein import HolsteinPropagator +from ipie.addons.eph.hamiltonians.holstein import HolsteinModel + +PropagatorAddons = {HolsteinModel: HolsteinPropagator} diff --git a/ipie/propagation/propagator.py b/ipie/propagation/propagator.py index a8bc0f62..807a5412 100644 --- a/ipie/propagation/propagator.py +++ b/ipie/propagation/propagator.py @@ -1,6 +1,6 @@ from ipie.hamiltonians.generic import GenericRealChol, GenericComplexChol from ipie.propagation.phaseless_generic import PhaselessGeneric -from ipie.addons.eph.propagation.holstein import HolsteinPropagator -from ipie.addons.eph.hamiltonians.holstein import HolsteinModel +from ipie.addons.propagator import PropagatorAddons -Propagator = {GenericRealChol: PhaselessGeneric, GenericComplexChol: PhaselessGeneric, HolsteinModel: HolsteinPropagator} +Propagator = {GenericRealChol: PhaselessGeneric, GenericComplexChol: PhaselessGeneric} +Propagator.update(PropagatorAddons) From eeeb44c1b03400ba2e159d4bd9b83c6472bee63a Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Mon, 11 Mar 2024 00:24:02 -0400 Subject: [PATCH 30/61] reformatted code --- ipie/addons/eph/estimators/energy.py | 14 +- .../eph/estimators/local_energy_holstein.py | 42 +++-- ipie/addons/eph/hamiltonians/holstein.py | 36 ++-- ipie/addons/eph/propagation/holstein.py | 154 +++++++++-------- .../eph/trial_wavefunction/coherent_state.py | 123 +++++++++---- .../eph/trial_wavefunction/eph_trial_base.py | 53 +++--- .../addons/eph/trial_wavefunction/toyozawa.py | 162 ++++++++++-------- .../variational/coherent_state_variational.py | 83 +++++---- .../variational/estimators.py | 4 +- .../variational/toyozawa_variational.py | 126 ++++++++------ ipie/addons/eph/walkers/eph_walkers.py | 43 ++--- 11 files changed, 481 insertions(+), 359 deletions(-) diff --git a/ipie/addons/eph/estimators/energy.py b/ipie/addons/eph/estimators/energy.py index 74ab6c18..5aaf8bc9 100644 --- a/ipie/addons/eph/estimators/energy.py +++ b/ipie/addons/eph/estimators/energy.py @@ -18,19 +18,21 @@ from ipie.addons.eph.estimators.local_energy_holstein import local_energy_holstein from ipie.addons.eph.hamiltonians.holstein import HolsteinModel from ipie.systems.generic import Generic -from ipie.addons.eph.trial_wavefunction.eph_trial_base import EPhTrialWavefunctionBase +from ipie.addons.eph.trial_wavefunction.eph_trial_base import EPhTrialWavefunctionBase from ipie.utils.backend import arraylib as xp -from ipie.addons.eph.walkers.eph_walkers import EPhWalkers +from ipie.addons.eph.walkers.eph_walkers import EPhWalkers + @plum.dispatch def local_energy( - system: Generic, - hamiltonian: HolsteinModel, - walkers: EPhWalkers, - trial: EPhTrialWavefunctionBase + system: Generic, + hamiltonian: HolsteinModel, + walkers: EPhWalkers, + trial: EPhTrialWavefunctionBase, ): return local_energy_holstein(system, hamiltonian, walkers, trial) + class EnergyEstimator(EstimatorBase): def __init__( self, diff --git a/ipie/addons/eph/estimators/local_energy_holstein.py b/ipie/addons/eph/estimators/local_energy_holstein.py index 81411e0f..aab5589b 100644 --- a/ipie/addons/eph/estimators/local_energy_holstein.py +++ b/ipie/addons/eph/estimators/local_energy_holstein.py @@ -19,58 +19,56 @@ from ipie.addons.eph.walkers.eph_walkers import EPhWalkers from ipie.systems.generic import Generic -from ipie.utils.backend import arraylib as xp +from ipie.utils.backend import arraylib as xp def local_energy_holstein( - system: Generic, - hamiltonian: HolsteinModel, - walkers: EPhWalkers, - trial: EPhTrialWavefunctionBase + system: Generic, + hamiltonian: HolsteinModel, + walkers: EPhWalkers, + trial: EPhTrialWavefunctionBase, ) -> np.ndarray: r"""Computes the local energy for the Holstein model via - + .. math:: \frac{\langle \Psi_\mathrm{T} | \hat{H} | \Phi_\mathrm{w}\rangle} {\langle \Psi_\mathrm{T} | \Phi_\mathrm{w}\rangle}, - where :math:`| \Phi_\mathrm{w}\rangle = \sum_{\tau \in cyclic perms} + where :math:`| \Phi_\mathrm{w}\rangle = \sum_{\tau \in cyclic perms} |\phi(\tau(r))\rangle \otimes |\beta(\tau(X))\rangle`. In this ansatz for - the walkers :math:`|\beta\rangle` is a coherent state, which corresonds to a - by :math:`\beta` displaced vacuum state. + the walkers :math:`|\beta\rangle` is a coherent state, which corresonds to a + by :math:`\beta` displaced vacuum state. Parameters ---------- - system : + system : :class:`Generic` Generic object carrying information on number and spin of electrons - hamiltonian : + hamiltonian : :class:`HolsteinModel` HolsteinModel object - walkers : + walkers : :class:`EPhWalkers` EPhWalkers object - trial : + trial : :class:`EPhTrialWavefunctionBase` EPhTrialWavefunctionBase object """ energy = xp.zeros((walkers.nwalkers, 4), dtype=xp.complex128) gf = trial.calc_greens_function(walkers) - - energy[:, 1] = np.sum(hamiltonian.T[0] * gf[0], axis=(1,2)) + walkers.Ga, walkers.Gb = gf[0], gf[1] + + energy[:, 1] = np.sum(hamiltonian.T[0] * gf[0], axis=(1, 2)) if system.ndown > 0: - energy[:, 1] += np.sum(hamiltonian.T[1] * gf[1], axis=(1,2)) + energy[:, 1] += np.sum(hamiltonian.T[1] * gf[1], axis=(1, 2)) energy[:, 2] = np.sum(np.diagonal(gf[0], axis1=1, axis2=2) * walkers.phonon_disp, axis=1) if system.ndown > 0: - energy[:, 2] += np.sum(np.diagonal(gf[1], axis1=1, axis2=2) * walkers.phonon_disp, axis=1) + energy[:, 2] += np.sum(np.diagonal(gf[1], axis1=1, axis2=2) * walkers.phonon_disp, axis=1) energy[:, 2] *= hamiltonian.const energy[:, 3] = 0.5 * hamiltonian.m * hamiltonian.w0**2 * np.sum(walkers.phonon_disp**2, axis=1) energy[:, 3] -= 0.5 * hamiltonian.nsites * hamiltonian.w0 energy[:, 3] -= 0.5 * trial.calc_phonon_laplacian_locenergy(walkers) / hamiltonian.m - - energy[:, 0] = np.sum(energy[:,1:], axis=1) - - return energy - + energy[:, 0] = np.sum(energy[:, 1:], axis=1) + return energy diff --git a/ipie/addons/eph/hamiltonians/holstein.py b/ipie/addons/eph/hamiltonians/holstein.py index 6fda9691..56ad65b7 100644 --- a/ipie/addons/eph/hamiltonians/holstein.py +++ b/ipie/addons/eph/hamiltonians/holstein.py @@ -22,27 +22,27 @@ class HolsteinModel: The Holstein model is described by the Hamiltonian .. math:: - \hat{H} = -t \sum_{\langle ij\rangle} \hat{a}_i^\dagger \hat{a}_j + \hat{H} = -t \sum_{\langle ij\rangle} \hat{a}_i^\dagger \hat{a}_j - g \sqrt{2 w_0 m} \sum_i \hat{a}_i^\dagger \hat{a}_i \hat{X}_i - + \bigg(\sum_i \frac{m w_0^2}{2} \hat{X}_i^2 + \frac{1}{2m} \hat{P}_i^2 + + \bigg(\sum_i \frac{m w_0^2}{2} \hat{X}_i^2 + \frac{1}{2m} \hat{P}_i^2 - \frac{w_0}{2}\bigg), where :math:`t` is associated with the electronic hopping, :math:`g` with - the electron-phonon coupling strength, and :math:``w_0` with the phonon - frequency. + the electron-phonon coupling strength, and :math:``w_0` with the phonon + frequency. Parameters ---------- - g : + g : :class:`float` Electron-phonon coupling strength - t : + t : :class:`float` Electron hopping parameter - w0 : + w0 : :class:`float` Phonon frequency - nsites : + nsites : :class:`int` Length of the 1D Holstein chain - pbc : - Boolean specifying whether periodic boundary conditions should be + pbc : :class:``bool` + Boolean specifying whether periodic boundary conditions should be employed. """ @@ -50,22 +50,20 @@ def __init__(self, g: float, t: float, w0: float, nsites: int, pbc: bool): self.g = g self.t = t self.w0 = w0 - self.m = 1/self.w0 + self.m = 1 / self.w0 self.nsites = nsites self.pbc = pbc self.T = None - self.const = -self.g * numpy.sqrt(2. * self.m * self.w0) + self.const = -self.g * numpy.sqrt(2.0 * self.m * self.w0) - def build(self): + def build(self) -> None: """Constructs electronic hopping matrix.""" - self.T = numpy.diag(numpy.ones(self.nsites-1), 1) - self.T += numpy.diag(numpy.ones(self.nsites-1), -1) - + self.T = numpy.diag(numpy.ones(self.nsites - 1), 1) + self.T += numpy.diag(numpy.ones(self.nsites - 1), -1) + if self.pbc: - self.T[0,-1] = self.T[-1,0] = 1. + self.T[0, -1] = self.T[-1, 0] = 1.0 self.T *= -self.t self.T = [self.T.copy(), self.T.copy()] - - diff --git a/ipie/addons/eph/propagation/holstein.py b/ipie/addons/eph/propagation/holstein.py index 8ae45c54..a5c9c19d 100644 --- a/ipie/addons/eph/propagation/holstein.py +++ b/ipie/addons/eph/propagation/holstein.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import numpy +import numpy import time import scipy.linalg @@ -20,10 +20,11 @@ from ipie.addons.eph.trial_wavefunction.eph_trial_base import EPhTrialWavefunctionBase from ipie.addons.eph.walkers.eph_walkers import EPhWalkers +from ipie.utils.backend import synchronize from ipie.propagation.operations import propagate_one_body -from ipie.utils.backend import synchronize, cast_to_device from ipie.propagation.continuous_base import PropagatorTimer + def construct_one_body_propagator(hamiltonian: HolsteinModel, dt: float): """Exponentiates the electronic hopping term to apply it later as part of the trotterized algorithm. @@ -32,7 +33,7 @@ def construct_one_body_propagator(hamiltonian: HolsteinModel, dt: float): ---------- hamiltonian : Hamiltonian caryying the one-body term as hamiltonian.T - dt : + dt : Time step """ H1 = hamiltonian.T @@ -44,25 +45,26 @@ def construct_one_body_propagator(hamiltonian: HolsteinModel, dt: float): class HolsteinPropagatorFree: r"""Propagates walkers by trotterization, - .. math:: - \mathrm{e}^{-\Delta \tau \hat{H}} \approx \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{ph}} / 2} - \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{el}} / 2} \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{el-ph}}} + .. math:: + \mathrm{e}^{-\Delta \tau \hat{H}} \approx \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{ph}} / 2} + \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{el}} / 2} \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{el-ph}}} \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{el}} / 2} \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{ph}} / 2}, - where propagation under :math:`\hat{H}_{\mathrm{ph}}` employs a generic - Diffucion MC procedure (notably without importance sampling). Propagation by - :math:`\hat{H}_{\mathrm{el}}` consists of a simple mat-vec. As - :math:`\hat{H}_{\mathrm{el-ph}}` is diagonal in bosonic position space we + where propagation under :math:`\hat{H}_{\mathrm{ph}}` employs a generic + Diffucion MC procedure (notably without importance sampling). Propagation by + :math:`\hat{H}_{\mathrm{el}}` consists of a simple mat-vec. As + :math:`\hat{H}_{\mathrm{el-ph}}` is diagonal in bosonic position space we can straightforwardly exponentiate the displacements and perform another mat-vec with this diagonal matrix apllied to electronic degrees of freedom. Parameters ---------- - time_step : + time_step : Time step - verbose : + verbose : Print level """ + def __init__(self, time_step: float, verbose: bool = False): self.dt = time_step self.verbose = verbose @@ -72,51 +74,55 @@ def __init__(self, time_step: float, verbose: bool = False): self.dt_ph = 0.5 * self.dt self.mpi_handler = None - def build(self, hamiltonian: HolsteinModel, - trial: EPhTrialWavefunctionBase = None, - walkers: EPhWalkers = None, mpi_handler = None) -> None: - """Necessary step before running the AFQMC procedure. - Sets required attributes. - + def build( + self, + hamiltonian: HolsteinModel, + trial: EPhTrialWavefunctionBase = None, + walkers: EPhWalkers = None, + mpi_handler=None, + ) -> None: + """Necessary step before running the AFQMC procedure. + Sets required attributes. + Parameters ---------- - hamiltonian : + hamiltonian : Holstein model trial : Trial object - walkers : + walkers : Walkers object mpi_handler : MPIHandler specifying rank and size """ self.expH1 = construct_one_body_propagator(hamiltonian, self.dt) - self.const = hamiltonian.g * numpy.sqrt(2. * hamiltonian.m * hamiltonian.w0) * self.dt + self.const = hamiltonian.g * numpy.sqrt(2.0 * hamiltonian.m * hamiltonian.w0) * self.dt self.w0 = hamiltonian.w0 self.m = hamiltonian.m self.scale = numpy.sqrt(self.dt_ph / self.m) self.nsites = hamiltonian.nsites - def propagate_phonons(self, walkers: EPhWalkers, - hamiltonian: HolsteinModel, - trial: EPhTrialWavefunctionBase) -> None: + def propagate_phonons( + self, walkers: EPhWalkers, hamiltonian: HolsteinModel, trial: EPhTrialWavefunctionBase + ) -> None: r"""Propagates phonon displacements by adjusting weigths according to bosonic on-site energies and sampling the momentum contribution, again by trotterizing the phonon propagator. - - .. math:: - \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{ph}} / 2} \approx - \mathrm{e}^{\Delta \tau N \omega / 4} + + .. math:: + \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{ph}} / 2} \approx + \mathrm{e}^{\Delta \tau N \omega / 4} + \mathrm{e}^{-\Delta \tau \sum_i m \omega \hat{X}_i^2 / 8} + \mathrm{e}^{-\Delta \tau \sum_i \hat{P}_i^2 / (4 \omega)} \mathrm{e}^{-\Delta \tau \sum_i m \omega \hat{X}_i^2 / 8} - \mathrm{e}^{-\Delta \tau \sum_i \hat{P}_i^2 / (4 \omega)} - \mathrm{e}^{-\Delta \tau \sum_i m \omega \hat{X}_i^2 / 8} One can obtain the sampling prescription by insertion of resolutions of identity, :math:`\int dX |X\rangle \langleX|, and performin the resulting - Fourier transformation. + Fourier transformation. Parameters ---------- - walkers : + walkers : Walkers class """ start_time = time.time() @@ -125,38 +131,38 @@ def propagate_phonons(self, walkers: EPhWalkers, pot = numpy.real(pot) walkers.weight *= numpy.exp(-self.dt_ph * pot) - N = numpy.random.normal(loc=0.0, scale=self.scale, - size=(walkers.nwalkers, self.nsites)) - walkers.phonon_disp = walkers.phonon_disp + N + N = numpy.random.normal(loc=0.0, scale=self.scale, size=(walkers.nwalkers, self.nsites)) + walkers.phonon_disp = walkers.phonon_disp + N pot = 0.25 * self.m * self.w0**2 * numpy.sum(walkers.phonon_disp**2, axis=1) pot = numpy.real(pot) walkers.weight *= numpy.exp(-self.dt_ph * pot) - + # Does not matter for estimators but helps with population control walkers.weight *= numpy.exp(self.dt_ph * self.nsites * self.w0 / 2) synchronize() self.timer.tgemm += time.time() - start_time - def propagate_electron(self, walkers: EPhWalkers, hamiltonian: HolsteinModel, - trial: EPhTrialWavefunctionBase) -> None: + def propagate_electron( + self, walkers: EPhWalkers, hamiltonian: HolsteinModel, trial: EPhTrialWavefunctionBase + ) -> None: r"""Propagates electronic degrees of freedom via - .. math:: - \mathrm{e}^{-\Delta \tau (\hat{H}_{\mathrm{el}} \otimes \hat{I}_{\mathrm{ph}} + \hat{H}_{\mathrm{el-ph}})} + .. math:: + \mathrm{e}^{-\Delta \tau (\hat{H}_{\mathrm{el}} \otimes \hat{I}_{\mathrm{ph}} + \hat{H}_{\mathrm{el-ph}})} \approx \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{el}} / 2} \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{el-ph}}} \mathrm{e}^{-\Delta \tau \hat{H}_{\mathrm{el}} / 2}. This acts on walkers of the form :math:`|\phi(\tau)\rangle \otimes |X(\tau)\rangle`. - + Parameters ---------- - walkers : + walkers : Walkers class - trial : + trial : Trial class """ start_time = time.time() @@ -164,31 +170,32 @@ def propagate_electron(self, walkers: EPhWalkers, hamiltonian: HolsteinModel, self.timer.tgf += time.time() - start_time expEph = numpy.exp(self.const * walkers.phonon_disp) - + walkers.phia = propagate_one_body(walkers.phia, self.expH1[0]) - walkers.phia = numpy.einsum('ni,nie->nie', expEph, walkers.phia) + walkers.phia = numpy.einsum("ni,nie->nie", expEph, walkers.phia) walkers.phia = propagate_one_body(walkers.phia, self.expH1[0]) - + if walkers.ndown > 0: walkers.phib = propagate_one_body(walkers.phib, self.expH1[1]) - walkers.phib = numpy.einsum('ni,nie->nie', expEph, walkers.phib) + walkers.phib = numpy.einsum("ni,nie->nie", expEph, walkers.phib) walkers.phib = propagate_one_body(walkers.phib, self.expH1[1]) - def propagate_walkers(self, - walkers: EPhWalkers, - hamiltonian: HolsteinModel, - trial: EPhTrialWavefunctionBase, - eshift: float = None - ) -> None: + def propagate_walkers( + self, + walkers: EPhWalkers, + hamiltonian: HolsteinModel, + trial: EPhTrialWavefunctionBase, + eshift: float = None, + ) -> None: r"""Propagates walkers by trotterized propagator. - + Parameters ---------- - walkers : + walkers : EPhWalkers object - hamiltonian : + hamiltonian : HolsteinModel object - trial : + trial : EPhTrialWavefunctionBase object eshift : Only purpose is compatibility with AFQMC object, irrelevant for @@ -197,6 +204,7 @@ def propagate_walkers(self, synchronize() start_time = time.time() ovlp = trial.calc_overlap(walkers) + walkers.ovlp = ovlp synchronize() self.timer.tgf += time.time() - start_time @@ -213,6 +221,7 @@ def propagate_walkers(self, # Update weights (and later do phaseless for multi-electron) start_time = time.time() ovlp_new = trial.calc_overlap(walkers) + walkers.ovlp = ovlp_new synchronize() self.timer.tovlp += time.time() - start_time @@ -221,56 +230,57 @@ def propagate_walkers(self, synchronize() self.timer.tupdate += time.time() - start_time - def update_weight(self, walkers, ovlp, ovlp_new) -> None: walkers.weight *= ovlp_new / ovlp class HolsteinPropagator(HolsteinPropagatorFree): - r"""Propagates walkers by trotterization, employing importance sampling for + r"""Propagates walkers by trotterization, employing importance sampling for the bosonic degrees of freedom. This results in a different weigth update, and the additional displacement update by the drift term, - + .. math:: D = \frac{\nabla_X \langle \Psi_\mathrm{T} | \psi(\tau), X(\tau)\rangle} {\langle \Psi_\mathrm{T} | \psi(\tau), X(\tau)\rangle}, such that the revised displacement update reads - .. math:: - X(\tau+\Delta\tau) = X(\tau) - + \mathcal{N}(\mu=0, \sigma = \sqrt{\frac{\Delta\tau}{m}}) + .. math:: + X(\tau+\Delta\tau) = X(\tau) + + \mathcal{N}(\mu=0, \sigma = \sqrt{\frac{\Delta\tau}{m}}) + \frac{\Delta\tau}{m} D. Parameters ---------- - time_step : + time_step : Time step verbose : Print level """ + def __init__(self, time_step, verbose=False): super().__init__(time_step, verbose=verbose) - def propagate_phonons(self, walkers: EPhWalkers, hamiltonian: HolsteinModel, - trial: EPhTrialWavefunctionBase) -> None: + def propagate_phonons( + self, walkers: EPhWalkers, hamiltonian: HolsteinModel, trial: EPhTrialWavefunctionBase + ) -> None: """Propagates phonons via Diffusion MC including drift term.""" start_time = time.time() - - # No ZPE in pot -> cancels with ZPE of etrial, + + # No ZPE in pot -> cancels with ZPE of etrial, # wouldn't affect estimators anyways ph_ovlp_old = trial.calc_phonon_overlap(walkers) - + pot = 0.5 * hamiltonian.m * hamiltonian.w0**2 * numpy.sum(walkers.phonon_disp**2, axis=1) pot -= 0.5 * trial.calc_phonon_laplacian_importance(walkers) / hamiltonian.m pot = numpy.real(pot) walkers.weight *= numpy.exp(-self.dt_ph * pot / 2) - N = numpy.random.normal(loc=0.0, scale=self.scale, size=(walkers.nwalkers, self.nsites)) + N = numpy.random.normal(loc=0.0, scale=self.scale, size=(walkers.nwalkers, self.nsites)) drift = trial.calc_phonon_gradient(walkers) walkers.phonon_disp = walkers.phonon_disp + N + self.dt_ph * drift / hamiltonian.m - ph_ovlp_new = trial.calc_phonon_overlap(walkers) + ph_ovlp_new = trial.calc_phonon_overlap(walkers) pot = 0.5 * hamiltonian.m * hamiltonian.w0**2 * numpy.sum(walkers.phonon_disp**2, axis=1) pot -= 0.5 * trial.calc_phonon_laplacian_importance(walkers) / hamiltonian.m @@ -282,5 +292,3 @@ def propagate_phonons(self, walkers: EPhWalkers, hamiltonian: HolsteinModel, synchronize() self.timer.tgemm += time.time() - start_time - - diff --git a/ipie/addons/eph/trial_wavefunction/coherent_state.py b/ipie/addons/eph/trial_wavefunction/coherent_state.py index 8ff33d16..a2c88a18 100644 --- a/ipie/addons/eph/trial_wavefunction/coherent_state.py +++ b/ipie/addons/eph/trial_wavefunction/coherent_state.py @@ -13,35 +13,36 @@ # limitations under the License. import numpy as np -import scipy.linalg from ipie.addons.eph.trial_wavefunction.eph_trial_base import EPhTrialWavefunctionBase from ipie.utils.backend import arraylib as xp + class CoherentStateTrial(EPhTrialWavefunctionBase): - r"""Coherent state trial of the form - + r"""Coherent state trial of the form + .. math:: |\Phi\rangle \otimes |\beta\rangle, - + where :math:`|\Phi\rangle` corresponds to the electronic wave function and :math:`|\beta\rangle` to the bosonic wave function. This latter is a coherent state, i.e. a vacuum state displaced by :math:`\beta`. - + Parameters ---------- - wavefunction : + wavefunction : Concatenation of trial determinants of up and down spin spaces and beta specifying the coherent state displacement. - hamiltonian : + hamiltonian : Holstein model Hamiltonian - num_elec : + num_elec : Tuple specifying number of up and down spins - num_basis : - Number of sites of Holstein chain. - verbose : + num_basis : + Number of sites of Holstein chain. + verbose : Print level """ + def __init__(self, wavefunction, hamiltonian, num_elec, num_basis, verbose=False): super().__init__(wavefunction, num_elec, num_basis, verbose=verbose) self.num_elec = num_elec @@ -51,14 +52,25 @@ def __init__(self, wavefunction, hamiltonian, num_elec, num_basis, verbose=False self.nsites = hamiltonian.nsites self.beta_shift = np.squeeze(wavefunction[:, 0]) - self.psia = wavefunction[:, 1:self.nup+1] - self.psib = wavefunction[:, self.nup+1:self.nup+self.ndown+1] - - def calculate_energy(self, system, hamiltonian): - # TODO variational_energy_coherent_state in ipie.estimators.local_energy - ... + self.psia = wavefunction[:, 1 : self.nup + 1] + self.psib = wavefunction[:, self.nup + 1 : self.nup + self.ndown + 1] def calc_overlap(self, walkers) -> np.ndarray: + r"""Computes the product of electron and phonon overlaps, + + .. math:: + \langle \Psi_T(r)|\Phi(\tau)\rangle \langle \phi(\beta)|X_{\mathrm{w}(\tau)}\rangle. + + Parameters + ---------- + walkers : :class:`EPhWalkers` + EPhWalkers object + + Returns + ------- + ovlp : :class:`np.ndarray` + Product of electron and phonon overlap + """ _ = self.calc_phonon_overlap(walkers) _ = self.calc_electronic_overlap(walkers) walkers.total_ovlp = walkers.el_ovlp * walkers.ph_ovlp @@ -66,16 +78,61 @@ def calc_overlap(self, walkers) -> np.ndarray: return walkers.ovlp def calc_phonon_overlap(self, walkers) -> np.ndarray: - ph_ovlp = np.exp(-(self.m * self.w0 / 2) * (walkers.phonon_disp - self.beta_shift)**2) + r"""Computes phonon overlap, :math:`\langle \phi(\beta)|X_{\mathrm{w}(\tau)}\rangle`. + + Parameters + ---------- + walkers : :class:`EPhWalkers` + EPhWalkers object + + Returns + ------- + ph_ovlp : :class:`np.ndarray` + Overlap of walekr position states with coherent trial state. + """ + ph_ovlp = np.exp(-(self.m * self.w0 / 2) * (walkers.phonon_disp - self.beta_shift) ** 2) walkers.ph_ovlp = np.prod(ph_ovlp, axis=1) return walkers.ph_ovlp def calc_phonon_gradient(self, walkers) -> np.ndarray: + r"""Computes the gradient of phonon overlaps, + + .. math:: + \frac{\nabla_X \langle\phi(\beta)|X_\mathrm{w}(\tau)\rangle} + {\langle\phi(\beta)|X_\mathrm{w}(\tau)\rangle} = -m \omega_0 (X(\tau) - \beta). + + Parameters + ---------- + walkers : :class:`EPhWalkers` + EPhWalkers object + + Returns + ------- + grad : :class:`np.ndarray` + Gradient of phonon overlap + """ grad = walkers.phonon_disp - self.beta_shift grad *= -self.m * self.w0 return grad def calc_phonon_laplacian(self, walkers) -> np.ndarray: + r"""Computes the Laplacian of phonon overlaps, + + .. math:: + \frac{\nabla^2_X \langle\phi(\beta)|X_\mathrm{w}(\tau)\rangle} + {\langle\phi(\beta)|X_\mathrm{w}(\tau)\rangle} + = - N m \omega_0 + \sum_i ((X_i(\tau) - \beta_i) m \omega_0)^2 + + Parameters + ---------- + walkers : :class:`EPhWalkers` + EPhWalkers object + + Returns + ------- + laplacian: :class:`np.ndarray` + Laplacian of phonon overlap + """ arg = (walkers.phonon_disp - self.beta_shift) * self.m * self.w0 arg2 = arg**2 laplacian = np.sum(arg2, axis=1) - self.nsites * self.m * self.w0 @@ -84,17 +141,20 @@ def calc_phonon_laplacian(self, walkers) -> np.ndarray: def calc_phonon_laplacian_importance(self, walkers) -> np.ndarray: return self.calc_phonon_laplacian(walkers) - def calc_phonon_laplacian_locenergy(self, walkers) -> np.ndarray: + def calc_phonon_laplacian_locenergy(self, walkers) -> np.ndarray: return self.calc_phonon_laplacian(walkers) def calc_electronic_overlap(self, walkers) -> np.ndarray: - """Computes electronic overlap. - + r"""Computes electronic overlap, + + .. math:: + \langle \Psi_T(r)|\Phi(\tau)\rangle = \prod_\sigma \mathrm{det}(V^{\dagger}_{\sigma}U_\sigam) + Parameters ---------- walkers : class EphWalkers class object - + Returns ------- walker.el_ovlp : np.ndarray @@ -116,23 +176,26 @@ def calc_electronic_overlap(self, walkers) -> np.ndarray: def calc_greens_function(self, walkers) -> np.ndarray: """Computes Greens function. - + Parameters ---------- walkers : class EphWalkers class object - + Returns ------- walkers.G : list Greens function for each spin space """ - inv_Oa = xp.linalg.inv(xp.einsum('ie,nif->nef', self.psia, walkers.phia.conj(), optimize=True)) - walkers.Ga = xp.einsum('nie,nef,jf->nji', walkers.phia, inv_Oa, self.psia.conj(), optimize=True) + inv_Oa = xp.linalg.inv( + xp.einsum("ie,nif->nef", self.psia, walkers.phia.conj(), optimize=True) + ) + Ga = xp.einsum("nie,nef,jf->nji", walkers.phia, inv_Oa, self.psia.conj(), optimize=True) if self.ndown > 0: - inv_Ob = xp.linalg.inv(xp.einsum('ie,nif->nef', self.psib, walkers.phib.conj(), optimize=True)) - walkers.Gb = xp.einsum('nie,nef,jf->nji', walkers.phib, inv_Ob, self.psib.conj(), optimize=True) - - return [walkers.Ga, walkers.Gb] + inv_Ob = xp.linalg.inv( + xp.einsum("ie,nif->nef", self.psib, walkers.phib.conj(), optimize=True) + ) + Gb = xp.einsum("nie,nef,jf->nji", walkers.phib, inv_Ob, self.psib.conj(), optimize=True) + return [Ga, Gb] diff --git a/ipie/addons/eph/trial_wavefunction/eph_trial_base.py b/ipie/addons/eph/trial_wavefunction/eph_trial_base.py index 81d8c475..f535d4e0 100644 --- a/ipie/addons/eph/trial_wavefunction/eph_trial_base.py +++ b/ipie/addons/eph/trial_wavefunction/eph_trial_base.py @@ -16,23 +16,26 @@ from typing import Tuple from abc import ABCMeta, abstractmethod + class EPhTrialWavefunctionBase(metaclass=ABCMeta): """Base class for electron-phonon trial wave functions. - + Parameters ---------- - wavefunction : + wavefunction : :class:`np.ndarray` Concatenation of trial determinants of up and down spin spaces and beta specifying the coherent state displacement. - num_elec : + num_elec : :class:`Tuple` Tuple of numbers of up and down spins. - num_basis : + num_basis : :class:`int` Number of sites of Holstein chain. - verbose : + verbose : :class:`bool` Print level """ - def __init__(self, wavefunction: np.ndarray, num_elec: Tuple[int, int], - num_basis: int, verbose=False): + + def __init__( + self, wavefunction: np.ndarray, num_elec: Tuple[int, int], num_basis: int, verbose=False + ): self.nelec = num_elec self.nbasis = num_basis self.nalpha, self.nbeta = self.nelec @@ -41,10 +44,10 @@ def __init__(self, wavefunction: np.ndarray, num_elec: Tuple[int, int], self._max_num_dets = self._num_dets self.ortho_expansion = False self.optimized = True - - self.psia = wavefunction[:self.nalpha] - self.psib = wavefunction[self.nalpha:self.nalpha+self.nbeta] - self.beta_shift = wavefunction[self.nalpha+self.nbeta:] + + self.psia = wavefunction[: self.nalpha] + self.psib = wavefunction[self.nalpha : self.nalpha + self.nbeta] + self.beta_shift = wavefunction[self.nalpha + self.nbeta :] self.compute_trial_energy = False self.energy = None @@ -58,35 +61,25 @@ def calculate_energy(self, system, hamiltonian): ... @abstractmethod - def calc_overlap(self, walkers) -> np.ndarray: - ... + def calc_overlap(self, walkers) -> np.ndarray: ... @abstractmethod - def calc_phonon_overlap(self, walkers) -> np.ndarray: - ... + def calc_phonon_overlap(self, walkers) -> np.ndarray: ... @abstractmethod - def calc_phonon_gradient(self, walkers) -> np.ndarray: - ... + def calc_phonon_gradient(self, walkers) -> np.ndarray: ... @abstractmethod - def calc_phonon_laplacian(self, walkers) -> np.ndarray: - ... + def calc_phonon_laplacian(self, walkers) -> np.ndarray: ... @abstractmethod - def calc_phonon_laplacian_importance(self, walkers) -> np.ndarray: - ... + def calc_phonon_laplacian_importance(self, walkers) -> np.ndarray: ... @abstractmethod - def calc_phonon_laplacian_locenergy(self, walkers) -> np.ndarray: - ... - - @abstractmethod - def calc_electronic_overlap(self, walkers) -> np.ndarray: - ... + def calc_phonon_laplacian_locenergy(self, walkers) -> np.ndarray: ... @abstractmethod - def calc_greens_function(self, walkers) -> np.ndarray: - ... - + def calc_electronic_overlap(self, walkers) -> np.ndarray: ... + @abstractmethod + def calc_greens_function(self, walkers) -> np.ndarray: ... diff --git a/ipie/addons/eph/trial_wavefunction/toyozawa.py b/ipie/addons/eph/trial_wavefunction/toyozawa.py index 7d614808..83c886a9 100644 --- a/ipie/addons/eph/trial_wavefunction/toyozawa.py +++ b/ipie/addons/eph/trial_wavefunction/toyozawa.py @@ -13,28 +13,28 @@ # limitations under the License. import numpy as np -import scipy.linalg from typing import Tuple from ipie.addons.eph.hamiltonians.holstein import HolsteinModel from ipie.addons.eph.walkers.eph_walkers import EPhWalkers from ipie.addons.eph.trial_wavefunction.coherent_state import CoherentStateTrial -from ipie.utils.backend import arraylib as xp from ipie.addons.eph.trial_wavefunction.variational.toyozawa_variational import circ_perm +from ipie.utils.backend import arraylib as xp + class ToyozawaTrial(CoherentStateTrial): - r"""The Toyozawa trial - + r"""The Toyozawa trial + .. math:: |\Psi(\kappa)\rangle = \sum_n e^{i \kappa n} \sum_{n_1} \alpha_{n_1}^{\kappa} a_{n_1}^{\dagger} \exp(-\sum_{n_2} (\beta^\kappa_{n_2 - n} b_{n_2}^{\dagger} - \beta^{\kappa^*}_{n_2 - n} b_{n_2}))|0\rangle - developed by `Toyozawa `_ is translationally - invariant and reliable offers a good approximation to the polaron ground state + developed by `Toyozawa `_ is translationally + invariant and reliable offers a good approximation to the polaron ground state for most parameter regimes of the Holstein Model. Here, :math:`\alpha,\beta`are varaitional parameters, and :math:`|0\rangle` is the total vacuum state. - For a 1D Holstein chain this reduces to a superposition of cyclically `CoherentState` + For a 1D Holstein chain this reduces to a superposition of cyclically `CoherentState` type trials. More details may be found in `Zhao et al. `_. @@ -45,8 +45,15 @@ class ToyozawaTrial(CoherentStateTrial): nperms : :class:`int` Number of permutations in `perms` """ - def __init__(self, wavefunction: np.ndarray, hamiltonian: HolsteinModel, - num_elec: Tuple[int, int], num_basis: int, verbose=False): + + def __init__( + self, + wavefunction: np.ndarray, + hamiltonian: HolsteinModel, + num_elec: Tuple[int, int], + num_basis: int, + verbose=False, + ): super().__init__(wavefunction, hamiltonian, num_elec, num_basis, verbose=verbose) self.perms = circ_perm(np.arange(self.nbasis)) self.nperms = self.perms.shape[0] @@ -54,16 +61,16 @@ def __init__(self, wavefunction: np.ndarray, hamiltonian: HolsteinModel, def calc_overlap_perm(self, walkers: EPhWalkers) -> np.ndarray: r"""Computes the product of electron and phonon overlaps for each permutation :math:`\sigma`, - + .. math:: - \langle \psi_T(\sigma(r))|\psi(\tau)\rangle + \langle \psi_T(\sigma(r))|\psi(\tau)\rangle \langle \phi(\sigma(\beta))|X_{\mathrm{w}(\tau)}\rangle. - + Parameters ---------- walkers : :class:`EPhWalkers` - EPhWalkers object - + EPhWalkers object + Returns ------- total_ovlp : :class:`np.ndarray` @@ -75,16 +82,16 @@ def calc_overlap_perm(self, walkers: EPhWalkers) -> np.ndarray: return walkers.total_ovlp def calc_overlap(self, walkers: EPhWalkers) -> np.ndarray: - r"""Sums product of electronic and phonon overlap for each permutation + r"""Sums product of electronic and phonon overlap for each permutation over all permutations, - + .. math:: - \sum_\tau \langle \psi_T(\sigma(r))|\psi(\tau)\rangle + \sum_\tau \langle \psi_T(\sigma(r))|\psi(\tau)\rangle \langle \phi(\sigma(\beta))|X_{\mathrm{w}(\tau)}\rangle. Used when evaluating local energy and when updating weight. - + Parameters ---------- walkers : :class:`EPhWalkers` @@ -100,40 +107,40 @@ def calc_overlap(self, walkers: EPhWalkers) -> np.ndarray: return ovlp def calc_phonon_overlap_perms(self, walkers: EPhWalkers) -> np.ndarray: - r"""Updates the walker phonon overlap with each permutation :math:`\tau`, + r"""Updates the walker phonon overlap with each permutation :math:`\tau`, i.e. :math:`\langle\phi(\tau(\beta))|X_{\mathrm{w}}\rangle` and stores - it in `walkers.ph_ovlp`. - + it in `walkers.ph_ovlp`. + Parameters ---------- - walkers : :class:`EPhWalkers` + walkers : :class:`EPhWalkers` EPhWalkers object Returns ------- ph_ovlp_perm : :class:`np.ndarray` - Overlap of walker with permuted coherent states + Overlap of walker with permuted coherent states """ - for ip,perm in enumerate(self.perms): + for ip, perm in enumerate(self.perms): ph_ov = np.exp( - -(0.5 * self.m * self.w0) * (walkers.phonon_disp - self.beta_shift[perm])**2 + -(0.5 * self.m * self.w0) * (walkers.phonon_disp - self.beta_shift[perm]) ** 2 ) walkers.ph_ovlp[:, ip] = np.prod(ph_ov, axis=1) return walkers.ph_ovlp def calc_phonon_overlap(self, walkers: EPhWalkers) -> np.ndarray: - r"""Sums walker phonon overlaps with permuted coherent states over all + r"""Sums walker phonon overlaps with permuted coherent states over all permutations, - - .. math:: - \sum_\tau \langle \phi(\tau(\beta)) | X_{\mathrm{w}} \rangle - + + .. math:: + \sum_\tau \langle \phi(\tau(\beta)) | X_{\mathrm{w}} \rangle + to get total phonon overlap. This is only used to correct for the importance sampling in propagate_phonons. Parameters ---------- - walkers : :class:`EPhWalkers` + walkers : :class:`EPhWalkers` EPhWalkers object Returns @@ -145,16 +152,16 @@ def calc_phonon_overlap(self, walkers: EPhWalkers) -> np.ndarray: ph_ovlp = np.sum(ph_ovlp_perm, axis=1) return ph_ovlp - def calc_phonon_gradient(self, walkers: EPhWalkers) -> np.ndarray: + def calc_phonon_gradient(self, walkers: EPhWalkers) -> np.ndarray: r"""Computes the phonon gradient, - + .. math:: \sum_\sigma \frac{\nabla_X \langle \phi(\sigma(\beta)) | X(\tau) \rangle} - {\rangle \phi(\sigma(\beta)) | X(\tau) \rangle} + {\rangle \phi(\sigma(\beta)) | X(\tau) \rangle} = \sum_\sigma -m \omega \frac{(X(\tau) - \sigma(\beta)) * \langle\phi(\sigma(\beta))|X(\tau)\rangle} {\sum_\simga \langle\phi(\sigma(\beta))|X(\tau)\rangle}. - - This is only used when calculating the drift term for the importance + + This is only used when calculating the drift term for the importance sampling DMC part of the algorithm. Parameters @@ -169,18 +176,18 @@ def calc_phonon_gradient(self, walkers: EPhWalkers) -> np.ndarray: """ grad = np.zeros_like(walkers.phonon_disp, dtype=np.complex128) for ovlp, perm in zip(walkers.ph_ovlp.T, self.perms): - grad += np.einsum('ni,n->ni', (walkers.phonon_disp - self.beta_shift[perm]), ovlp) + grad += np.einsum("ni,n->ni", (walkers.phonon_disp - self.beta_shift[perm]), ovlp) grad *= -self.m * self.w0 - grad = np.einsum('ni,n->ni', grad, 1/np.sum(walkers.ph_ovlp, axis=1)) + grad = np.einsum("ni,n->ni", grad, 1 / np.sum(walkers.ph_ovlp, axis=1)) return grad - + def calc_phonon_laplacian(self, walkers: EPhWalkers, ovlps: np.ndarray) -> np.ndarray: r"""Computes the phonon Laplacian, which weights coherent state laplacians by overlaps :math:`o(\sigma, r, X, \tau)` passed to this function, - - .. math:: + + .. math:: \sum_\sigma \frac{\nabla_X \langle \phi(\sigma(\beta)) | X(\tau) \rangle} - {\rangle \phi(\sigma(\beta)) | X(\tau) \rangle} + {\rangle \phi(\sigma(\beta)) | X(\tau) \rangle} = \frac{\sum_sigma ((\sum_i (m \omega (X_i(\tau) - \sigma(\beta)_i))^2) - N m \omega) o(\sigma, r, X, \tau)} {\sum_\sigma o(\sigma, r, X, \tau)}. @@ -190,14 +197,14 @@ def calc_phonon_laplacian(self, walkers: EPhWalkers, ovlps: np.ndarray) -> np.nd EPhWalkers object ovlps : :class:`np.ndarray` Overlaps weighting contributions from permuted coherent states - + Returns ------- laplacian : :class:`np.ndarray` Phonon Laplacian """ laplacian = np.zeros(walkers.nwalkers, dtype=np.complex128) - for ovlp, perm in zip(ovlps.T, self.perms): + for ovlp, perm in zip(ovlps.T, self.perms): arg = (walkers.phonon_disp - self.beta_shift[perm]) * self.m * self.w0 arg2 = arg**2 laplacian += (np.sum(arg2, axis=1) - self.nsites * self.m * self.w0) * ovlp @@ -208,7 +215,7 @@ def calc_phonon_laplacian_importance(self, walkers: EPhWalkers) -> np.ndarray: r"""Computes phonon Laplacian via `calc_phonon_laplacian` with weighting by pure phonon overlap. This is only utilized in the importance sampling of the DMC procedure. - + Parameters ---------- walkers : :class:`EPhWalkers` @@ -221,10 +228,10 @@ def calc_phonon_laplacian_importance(self, walkers: EPhWalkers) -> np.ndarray: """ return self.calc_phonon_laplacian(walkers, walkers.ph_ovlp) - def calc_phonon_laplacian_locenergy(self, walkers: EPhWalkers) -> np.ndarray: + def calc_phonon_laplacian_locenergy(self, walkers: EPhWalkers) -> np.ndarray: """Computes phonon Laplacian using total overlap weights as required in local energy evaluation. - + Parameters ---------- walkers : :class:`EPhWalkers` @@ -238,13 +245,13 @@ def calc_phonon_laplacian_locenergy(self, walkers: EPhWalkers) -> np.ndarray: return self.calc_phonon_laplacian(walkers, walkers.total_ovlp) def calc_electronic_overlap_perms(self, walkers: EPhWalkers) -> np.ndarray: - r"""Calculates the electronic overlap of each walker with each permuted - Slater determinant :math:`|\Phi_T(\tau(r_i))\rangle` of the trial, + r"""Calculates the electronic overlap of each walker with each permuted + Slater determinant :math:`|\Phi_T(\tau(r_i))\rangle` of the trial, - .. math:: + .. math:: \langle \Phi_T(\tau(r_i))|\psi_w\rangle = \mathrm{det(U^{\dagger}V)}, - - where :math:`U,V` parametrized the two Slater determinants. + + where :math:`U,V` parametrized the two Slater determinants. Parameters ---------- @@ -256,24 +263,28 @@ def calc_electronic_overlap_perms(self, walkers: EPhWalkers) -> np.ndarray: el_ovlp_perm : :class:`np.ndarray` Electronic overlap of each permuted Slater Determiant with walkers """ - for ip,perm in enumerate(self.perms): - ovlp_a = xp.einsum("wmi,mj->wij", walkers.phia, self.psia[perm, :].conj(), optimize=True) + for ip, perm in enumerate(self.perms): + ovlp_a = xp.einsum( + "wmi,mj->wij", walkers.phia, self.psia[perm, :].conj(), optimize=True + ) sign_a, log_ovlp_a = xp.linalg.slogdet(ovlp_a) if self.ndown > 0: - ovlp_b = xp.einsum("wmi,mj->wij", walkers.phib, self.psib[perm, :].conj(), optimize=True) + ovlp_b = xp.einsum( + "wmi,mj->wij", walkers.phib, self.psib[perm, :].conj(), optimize=True + ) sign_b, log_ovlp_b = xp.linalg.slogdet(ovlp_b) ot = sign_a * sign_b * xp.exp(log_ovlp_a + log_ovlp_b - walkers.log_shift) - else: + else: ot = sign_a * xp.exp(log_ovlp_a - walkers.log_shift) walkers.el_ovlp[:, ip] = ot return walkers.el_ovlp def calc_electronic_overlap(self, walkers: EPhWalkers) -> np.ndarray: - """Sums walkers.el_ovlp over permutations to obtain total electronic + """Sums walkers.el_ovlp over permutations to obtain total electronic overlap of trial with walkers. - + Parameters ---------- walkers : :class:`EPhWalkers` @@ -290,16 +301,16 @@ def calc_electronic_overlap(self, walkers: EPhWalkers) -> np.ndarray: def calc_greens_function(self, walkers: EPhWalkers, build_full=True) -> np.ndarray: r"""Calculates Greens functions by - + .. math:: - G^{\Phi \Psi}_{p\alpha, q\beta} + G^{\Phi \Psi}_{p\alpha, q\beta} = \frac{\sum_{\tau} \delta_{\alpha\beta}(U_\alpha(V^{\dagger}_\alhpa(\tau) U_\alpha) V^{\dagger}_\alpha(\tau)) \langle\Phi_T(\tau(r_i))|\psi_w\rangle} {\sum_{\tau} \langle\Phi_T(\tau(r_i))|\psi_w\rangle} - + Parameters ---------- walkers : :class:`EPhWalkers` - EPhWalkers object + EPhWalkers object Returns ------- @@ -310,17 +321,22 @@ def calc_greens_function(self, walkers: EPhWalkers, build_full=True) -> np.ndarr Gb = np.zeros_like(Ga) for ovlp, perm in zip(walkers.total_ovlp.T, self.perms): - - inv_Oa = xp.linalg.inv(xp.einsum('ie,nif->nef', self.psia[perm,:], walkers.phia.conj())) - Ga += xp.einsum('nie,nef,jf,n->nji', walkers.phia, inv_Oa, self.psia[perm].conj(), ovlp) - - if self.ndown > 0: - inv_Ob = xp.linalg.inv(xp.einsum('ie,nif->nef', self.psib[perm,:], walkers.phib.conj())) - Gb += xp.einsum('nie,nef,jf,n->nji', walkers.phib, inv_Ob, self.psib[perm].conj(), ovlp) - - Ga = np.einsum('nij,n->nij', Ga, 1 / walkers.ovlp) + + inv_Oa = xp.linalg.inv( + xp.einsum("ie,nif->nef", self.psia[perm, :], walkers.phia.conj()) + ) + Ga += xp.einsum("nie,nef,jf,n->nji", walkers.phia, inv_Oa, self.psia[perm].conj(), ovlp) + + if self.ndown > 0: + inv_Ob = xp.linalg.inv( + xp.einsum("ie,nif->nef", self.psib[perm, :], walkers.phib.conj()) + ) + Gb += xp.einsum( + "nie,nef,jf,n->nji", walkers.phib, inv_Ob, self.psib[perm].conj(), ovlp + ) + + Ga = np.einsum("nij,n->nij", Ga, 1 / walkers.ovlp) if self.ndown > 0: - Gb = np.einsum('nij,n->nij', Gb, 1 / walkers.ovlp) + Gb = np.einsum("nij,n->nij", Gb, 1 / walkers.ovlp) return [Ga, Gb] - diff --git a/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py b/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py index 60d591c0..dd7a6eb4 100644 --- a/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py +++ b/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py @@ -22,22 +22,23 @@ import jax + def local_energy( - X: np.ndarray, - G: np.ndarray, - m: float, - w: float, - g: float, - nsites: int, - Lap: np.ndarray, - T: np.ndarray, - nup: int, - ndown: int - ) -> np.ndarray: - - kinetic_contrib = jax.numpy.einsum('ij->', T[0] * G[0]) + X: np.ndarray, + G: np.ndarray, + m: float, + w: float, + g: float, + nsites: int, + Lap: np.ndarray, + T: np.ndarray, + nup: int, + ndown: int, +) -> np.ndarray: + + kinetic_contrib = jax.numpy.einsum("ij->", T[0] * G[0]) if ndown > 0: - kinetic_contrib += jax.numpy.einsum('ij->', T[1] * G[1]) + kinetic_contrib += jax.numpy.einsum("ij->", T[1] * G[1]) rho = G[0].diagonal() + G[1].diagonal() el_ph_contrib = -g * np.sqrt(2 * m * w) * np.sum(rho * X) @@ -48,24 +49,26 @@ def local_energy( local_energy = kinetic_contrib + el_ph_contrib + phonon_contrib return local_energy -def objective_function(x: np.ndarray, nbasis: int, T: np.ndarray, - g: float, m: float, w0: float, nup: int, ndown: int): + +def objective_function( + x: np.ndarray, nbasis: int, T: np.ndarray, g: float, m: float, w0: float, nup: int, ndown: int +): shift = x[0:nbasis].copy() shift = shift.astype(np.float64) - - c0a = x[nbasis : (nup+1)*nbasis].copy() + + c0a = x[nbasis : (nup + 1) * nbasis].copy() c0a = jax.numpy.reshape(c0a, (nbasis, nup)) c0a = c0a.astype(np.float64) Ga = gab(c0a, c0a) if ndown > 0: - c0b = x[(nup+1)*nbasis : ].copy() + c0b = x[(nup + 1) * nbasis :].copy() c0b = jax.numpy.reshape(c0b, (nbasis, ndown)) c0b = c0b.astype(np.float64) Gb = gab(c0b, c0b) else: Gb = jax.numpy.zeros_like(Ga, dtype=np.float64) - + G = [Ga, Gb] phi = HarmonicOscillator(m, w0, order=0, shift=shift) @@ -75,25 +78,27 @@ def objective_function(x: np.ndarray, nbasis: int, T: np.ndarray, return etot.real -def gradient(x: np.ndarray, nbasis: int, T: np.ndarray, - g: float, m: float, w0: float, nup: int, ndown: int): +def gradient( + x: np.ndarray, nbasis: int, T: np.ndarray, g: float, m: float, w0: float, nup: int, ndown: int +): grad = np.array( - jax.grad(objective_function)( - x, nbasis, T, g, m, w0, nup, ndown - ), - dtype=np.float64 + jax.grad(objective_function)(x, nbasis, T, g, m, w0, nup, ndown), dtype=np.float64 ) return grad -def func(x: np.ndarray, nbasis: int, T: np.ndarray, - g: float, m: float, w0: float, nup: int, ndown: int): + +def func( + x: np.ndarray, nbasis: int, T: np.ndarray, g: float, m: float, w0: float, nup: int, ndown: int +): f = objective_function(x, nbasis, T, g, m, w0, nup, ndown) df = gradient(x, nbasis, T, g, m, w0, nup, ndown) return f, df + def print_fun(x: np.ndarray, f: float, accepted: bool): print("at minimum %.4f accepted %d" % (f, int(accepted))) + def variational_trial(init_phonons: np.ndarray, init_electron: np.ndarray, hamiltonian, system): init_phonons = init_phonons.astype(np.float64) @@ -104,7 +109,15 @@ def variational_trial(init_phonons: np.ndarray, init_electron: np.ndarray, hamil maxiter = 500 minimizer_kwargs = { "jac": True, - "args": (hamiltonian.nsites, hamiltonian.T, hamiltonian.g, hamiltonian.m, hamiltonian.w0, system.nup, system.ndown), + "args": ( + hamiltonian.nsites, + hamiltonian.T, + hamiltonian.g, + hamiltonian.m, + hamiltonian.w0, + system.nup, + system.ndown, + ), "options": { "gtol": 1e-10, "eps": 1e-10, @@ -121,20 +134,18 @@ def variational_trial(init_phonons: np.ndarray, init_electron: np.ndarray, hamil niter=maxiter, niter_success=3, ) - + etrial = res.fun - beta_shift = res.x[:hamiltonian.nsites] + beta_shift = res.x[: hamiltonian.nsites] if system.ndown > 0: - psia = res.x[hamiltonian.nsites : hamiltonian.nsites*(system.nup+1)] + psia = res.x[hamiltonian.nsites : hamiltonian.nsites * (system.nup + 1)] psia = psia.reshape((hamiltonian.nsites, system.nup)) - psib = res.x[hamiltonian.nsites*(system.nup+1) : ] + psib = res.x[hamiltonian.nsites * (system.nup + 1) :] psib = psib.reshape((hamiltonian.nsites, system.ndown)) psi = np.column_stack([psia, psib]) else: - psia = res.x[hamiltonian.nsites:].reshape((hamiltonian.nsites, system.nup)) + psia = res.x[hamiltonian.nsites :].reshape((hamiltonian.nsites, system.nup)) psi = psia return etrial, beta_shift, psi - - diff --git a/ipie/addons/eph/trial_wavefunction/variational/estimators.py b/ipie/addons/eph/trial_wavefunction/variational/estimators.py index 6436364a..29e1bb3c 100644 --- a/ipie/addons/eph/trial_wavefunction/variational/estimators.py +++ b/ipie/addons/eph/trial_wavefunction/variational/estimators.py @@ -15,9 +15,11 @@ import numpy as np import jax.numpy as npj from jax.config import config + config.update("jax_enable_x64", True) -def gab(A,B): + +def gab(A: npj.ndarray, B: npj.ndarray) -> npj.ndarray: inv_O = npj.linalg.inv((A.conj().T).dot(B)) GAB = B.dot(inv_O.dot(A.conj().T)) return GAB diff --git a/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py b/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py index c872e3ab..5dafdd1c 100644 --- a/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py +++ b/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py @@ -13,7 +13,7 @@ # limitations under the License. import numpy as np -from scipy.optimize import minimize, basinhopping +from scipy.optimize import minimize, basinhopping from ipie.systems import Generic from ipie.addons.eph.hamiltonians.holstein import HolsteinModel from ipie.addons.eph.trial_wavefunction.variational.estimators import gab @@ -21,15 +21,23 @@ import jax import jax.numpy as npj + def gradient_toyozawa_mo(x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted): - grad = np.array(jax.grad(objective_function_toyozawa_mo)(x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted)) + grad = np.array( + jax.grad(objective_function_toyozawa_mo)( + x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted + ) + ) return grad -def objective_function_toyozawa_mo(x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted, zero_th=1e-12): + +def objective_function_toyozawa_mo( + x, nbasis, nup, ndown, T, U, g, m, w0, perms, restricted, zero_th=1e-12 +): nbasis = int(round(nbasis)) nup = int(round(nup)) ndown = int(round(ndown)) - + shift0 = x[:nbasis] nbsf = nbasis @@ -37,36 +45,42 @@ def objective_function_toyozawa_mo(x, nbasis, nup, ndown, T, U, g, m, w0, perms, noccb = ndown nvira = nbasis - nocca nvirb = nbasis - noccb - - psi0a = npj.array(x[nbsf:nbsf+nbsf*nocca],dtype=npj.float64) + + psi0a = npj.array(x[nbsf : nbsf + nbsf * nocca], dtype=npj.float64) psi0a = psi0a.reshape((nocca, nbsf)).T - if (noccb>0): - psi0b = npj.array(x[nbsf+nbsf*nocca : ], dtype=npj.float64) + if noccb > 0: + psi0b = npj.array(x[nbsf + nbsf * nocca :], dtype=npj.float64) psi0b = psi0b.reshape((noccb, nbsf)).T - + nperms = len(perms) - num_energy = 0. + num_energy = 0.0 denom = 0.0 - - beta0 = shift0 * npj.sqrt(m * w0 /2.0) - for i,perm in enumerate(perms): + beta0 = shift0 * npj.sqrt(m * w0 / 2.0) + + for i, perm in enumerate(perms): psia_j = psi0a[perm, :] beta_j = beta0[npj.array(perm)] - - if (noccb > 0): + + if noccb > 0: psib_j = psi0b[perm, :] - overlap = npj.linalg.det(psi0a.T.dot(psia_j)) * npj.linalg.det(psi0b.T.dot(psib_j)) * npj.prod(npj.exp(-0.5 * (beta0**2 + beta_j**2) + beta0*beta_j)) + overlap = ( + npj.linalg.det(psi0a.T.dot(psia_j)) + * npj.linalg.det(psi0b.T.dot(psib_j)) + * npj.prod(npj.exp(-0.5 * (beta0**2 + beta_j**2) + beta0 * beta_j)) + ) else: - overlap = npj.linalg.det(psi0a.T.dot(psia_j)) * npj.prod(npj.exp (- 0.5 * (beta0**2 + beta_j**2) + beta0*beta_j)) + overlap = npj.linalg.det(psi0a.T.dot(psia_j)) * npj.prod( + npj.exp(-0.5 * (beta0**2 + beta_j**2) + beta0 * beta_j) + ) if overlap < zero_th: continue Ga_j = gab(psi0a, psia_j) - if (noccb > 0): + if noccb > 0: Gb_j = gab(psi0b, psib_j) else: Gb_j = npj.zeros_like(Ga_j) @@ -76,8 +90,8 @@ def objective_function_toyozawa_mo(x, nbasis, nup, ndown, T, U, g, m, w0, perms, rho = G_j[0].diagonal() + G_j[1].diagonal() ke = npj.sum(T[0] * G_j[0] + T[1] * G_j[1]) pe = U * npj.dot(G_j[0].diagonal(), G_j[1].diagonal()) - e_ph = w0 * npj.sum(beta0 * beta_j) - e_eph = - g * npj.dot(rho, beta0 + beta_j) + e_ph = w0 * npj.sum(beta0 * beta_j) + e_eph = -g * npj.dot(rho, beta0 + beta_j) num_energy += (ke + pe + e_ph + e_eph) * overlap # 2.0 comes from hermiticity denom += overlap @@ -85,13 +99,14 @@ def objective_function_toyozawa_mo(x, nbasis, nup, ndown, T, U, g, m, w0, perms, etot = num_energy / denom return etot.real + def circ_perm(lst: np.ndarray) -> np.ndarray: - """Returns a matrix which rows consist of all possible + """Returns a matrix which rows consist of all possible cyclic permutations given an initial array lst. - + Parameters ---------- - lst : + lst : Initial array which is to be cyclically permuted """ circs = lst @@ -100,47 +115,60 @@ def circ_perm(lst: np.ndarray) -> np.ndarray: circs = np.vstack([circs, new_circ]) return circs -def variational_trial_toyozawa(shift_init: np.ndarray, electron_init: np.ndarray, hamiltonian, system): + +def variational_trial_toyozawa( + shift_init: np.ndarray, electron_init: np.ndarray, hamiltonian, system +): psi = electron_init.T.real.ravel() shift = shift_init.real perms = circ_perm(np.arange(hamiltonian.nsites)) x = np.zeros((system.nup + system.ndown + 1) * hamiltonian.nsites) - x[:hamiltonian.nsites] = shift.copy() - x[hamiltonian.nsites:] = psi.copy() + x[: hamiltonian.nsites] = shift.copy() + x[hamiltonian.nsites :] = psi.copy() res = minimize( - objective_function_toyozawa_mo, - x, - args = (float(hamiltonian.nsites), float(system.nup), float(system.ndown), - hamiltonian.T, 0., hamiltonian.g, hamiltonian.m, hamiltonian.w0, perms, False), - jac = gradient_toyozawa_mo, - tol = 1e-10, - method='L-BFGS-B', - options = { - 'maxls': 20, - 'iprint': 2, - 'gtol': 1e-10, - 'eps': 1e-10, - 'maxiter': 15000, - 'ftol': 1.0e-10, - 'maxcor': 1000, - 'maxfun': 15000, - 'disp':True - } + objective_function_toyozawa_mo, + x, + args=( + float(hamiltonian.nsites), + float(system.nup), + float(system.ndown), + hamiltonian.T, + 0.0, + hamiltonian.g, + hamiltonian.m, + hamiltonian.w0, + perms, + False, + ), + jac=gradient_toyozawa_mo, + tol=1e-10, + method="L-BFGS-B", + options={ + "maxls": 20, + "iprint": 2, + "gtol": 1e-10, + "eps": 1e-10, + "maxiter": 15000, + "ftol": 1.0e-10, + "maxcor": 1000, + "maxfun": 15000, + "disp": True, + }, ) etrial = res.fun - beta_shift = res.x[:hamiltonian.nsites] + beta_shift = res.x[: hamiltonian.nsites] if system.ndown > 0: - psia = res.x[hamiltonian.nsites : hamiltonian.nsites*(system.nup+1)] + psia = res.x[hamiltonian.nsites : hamiltonian.nsites * (system.nup + 1)] psia = psia.reshape((hamiltonian.nsites, system.nup)) - psib = res.x[hamiltonian.nsites*(system.nup+1) : ] + psib = res.x[hamiltonian.nsites * (system.nup + 1) :] psib = psib.reshape((hamiltonian.nsites, system.ndown)) psi = np.column_stack([psia, psib]) else: - psia = res.x[hamiltonian.nsites:].reshape((hamiltonian.nsites, system.nup)) + psia = res.x[hamiltonian.nsites :].reshape((hamiltonian.nsites, system.nup)) psi = psia - + return etrial, beta_shift, psi diff --git a/ipie/addons/eph/walkers/eph_walkers.py b/ipie/addons/eph/walkers/eph_walkers.py index 3baa78cd..c09f4243 100644 --- a/ipie/addons/eph/walkers/eph_walkers.py +++ b/ipie/addons/eph/walkers/eph_walkers.py @@ -19,71 +19,74 @@ from ipie.utils.backend import cast_to_device, qr, qr_mode, synchronize from ipie.walkers.base_walkers import BaseWalkers + class EPhWalkers(BaseWalkers): """Class tailored to el-ph models where keeping track of phonon overlaps is - required. Each walker carries along its Slater determinants a phonon + required. Each walker carries along its Slater determinants a phonon displacement vector, self.phonon_disp. - + Parameters ---------- initial_walker : - Walker that we start the simulation from. Ideally chosen according to + Walker that we start the simulation from. Ideally chosen according to the trial. - nup : + nup : Number of electrons in up-spin space. ndown : Number of electrons in down-spin space. nbasis : Number of sites in the 1D Holstein chain. - nwalkers : + nwalkers : Number of walkers in the simulation. - verbose : + verbose : Print level. """ + def __init__( - self, + self, initial_walker: numpy.ndarray, nup: int, ndown: int, nbasis: int, nwalkers: int, - verbose: bool = False + verbose: bool = False, ): self.nup = nup self.ndown = ndown self.nbasis = nbasis - super().__init__(nwalkers, verbose=verbose) self.weight = numpy.ones(self.nwalkers, dtype=numpy.complex128) self.phonon_disp = xp.array( - [initial_walker[:,0].copy() for iw in range(self.nwalkers)], - dtype=xp.complex128 + [initial_walker[:, 0].copy() for iw in range(self.nwalkers)], dtype=xp.complex128 ) self.phonon_disp = numpy.squeeze(self.phonon_disp) - + self.phia = xp.array( - [initial_walker[:, 1 : self.nup+1].copy() for iw in range(self.nwalkers)], + [initial_walker[:, 1 : self.nup + 1].copy() for iw in range(self.nwalkers)], dtype=xp.complex128, ) self.phib = xp.array( - [initial_walker[:, self.nup+1 : self.nup+self.ndown+1].copy() for iw in range(self.nwalkers)], + [ + initial_walker[:, self.nup + 1 : self.nup + self.ndown + 1].copy() + for iw in range(self.nwalkers) + ], dtype=xp.complex128, ) - + self.buff_names += ["phia", "phib", "phonon_disp"] self.buff_size = round(self.set_buff_size_single_walker() / float(self.nwalkers)) self.walker_buffer = numpy.zeros(self.buff_size, dtype=numpy.complex128) def build(self, trial): - """Allocates memory for computation of overlaps throughout the + """Allocates memory for computation of overlaps throughout the simulation. - + Parameters ---------- trial : @@ -99,7 +102,7 @@ def build(self, trial): self.el_ovlp = numpy.zeros(shape, dtype=numpy.complex128) self.total_ovlp = numpy.zeros(shape, dtype=numpy.complex128) - self.buff_names += ['total_ovlp'] + self.buff_names += ["total_ovlp"] self.buff_size = round(self.set_buff_size_single_walker() / float(self.nwalkers)) self.walker_buffer = numpy.zeros(self.buff_size, dtype=numpy.complex128) @@ -113,7 +116,7 @@ def cast_to_cupy(self, verbose=False): def reortho(self): """reorthogonalise walkers. This function is mostly from BaseWalkers, - with the exception that it adjusts all overlaps, possibly of numerous + with the exception that it adjusts all overlaps, possibly of numerous coherent states. Parameters @@ -149,7 +152,7 @@ def reortho(self): detR += [xp.exp(log_det - self.detR_shift[iw])] self.log_detR[iw] += xp.log(detR[iw]) self.detR[iw] = detR[iw] - + self.el_ovlp[iw, ...] = self.el_ovlp[iw, ...] / detR[iw] self.total_ovlp[iw, ...] = self.total_ovlp[iw, ...] / detR[iw] self.ovlp[iw] = self.ovlp[iw] / detR[iw] From 93e6a49c9451f9f7a0c1d025a8216e88220a3d24 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Mon, 11 Mar 2024 12:55:50 -0400 Subject: [PATCH 31/61] renamed EPhWalkers.total_ovlp to EPhWalkers.ovlp_perm --- ipie/addons/eph/trial_wavefunction/toyozawa.py | 17 ++++++++--------- ipie/addons/eph/walkers/eph_walkers.py | 6 +++--- 2 files changed, 11 insertions(+), 12 deletions(-) diff --git a/ipie/addons/eph/trial_wavefunction/toyozawa.py b/ipie/addons/eph/trial_wavefunction/toyozawa.py index 83c886a9..02a55669 100644 --- a/ipie/addons/eph/trial_wavefunction/toyozawa.py +++ b/ipie/addons/eph/trial_wavefunction/toyozawa.py @@ -73,13 +73,13 @@ def calc_overlap_perm(self, walkers: EPhWalkers) -> np.ndarray: Returns ------- - total_ovlp : :class:`np.ndarray` - Product of electron and phonon overlap + ovlp_perm : :class:`np.ndarray` + Product of electron and phonon overlap for each permutation """ ph_ovlp_perm = self.calc_phonon_overlap_perms(walkers) el_ovlp_perm = self.calc_electronic_overlap_perms(walkers) - walkers.total_ovlp = el_ovlp_perm * ph_ovlp_perm - return walkers.total_ovlp + walkers.ovlp_perm = el_ovlp_perm * ph_ovlp_perm + return walkers.ovlp_perm def calc_overlap(self, walkers: EPhWalkers) -> np.ndarray: r"""Sums product of electronic and phonon overlap for each permutation @@ -102,8 +102,8 @@ def calc_overlap(self, walkers: EPhWalkers) -> np.ndarray: ovlp: :class:`np.ndarray` Sum of product of electron and phonon overlap """ - total_ovlp = self.calc_overlap_perm(walkers) - ovlp = np.sum(total_ovlp, axis=1) + ovlp_perm = self.calc_overlap_perm(walkers) + ovlp = np.sum(ovlp_perm, axis=1) return ovlp def calc_phonon_overlap_perms(self, walkers: EPhWalkers) -> np.ndarray: @@ -242,7 +242,7 @@ def calc_phonon_laplacian_locenergy(self, walkers: EPhWalkers) -> np.ndarray: ph_lapl : :class:`np.ndarray` Phonon Laplacian weigthed by total overlaps """ - return self.calc_phonon_laplacian(walkers, walkers.total_ovlp) + return self.calc_phonon_laplacian(walkers, walkers.ovlp_perm) def calc_electronic_overlap_perms(self, walkers: EPhWalkers) -> np.ndarray: r"""Calculates the electronic overlap of each walker with each permuted @@ -320,8 +320,7 @@ def calc_greens_function(self, walkers: EPhWalkers, build_full=True) -> np.ndarr Ga = np.zeros((walkers.nwalkers, self.nsites, self.nsites), dtype=np.complex128) Gb = np.zeros_like(Ga) - for ovlp, perm in zip(walkers.total_ovlp.T, self.perms): - + for ovlp, perm in zip(walkers.ovlp_perm.T, self.perms): inv_Oa = xp.linalg.inv( xp.einsum("ie,nif->nef", self.psia[perm, :], walkers.phia.conj()) ) diff --git a/ipie/addons/eph/walkers/eph_walkers.py b/ipie/addons/eph/walkers/eph_walkers.py index c09f4243..4177a372 100644 --- a/ipie/addons/eph/walkers/eph_walkers.py +++ b/ipie/addons/eph/walkers/eph_walkers.py @@ -100,9 +100,9 @@ def build(self, trial): self.ph_ovlp = numpy.zeros(shape, dtype=numpy.complex128) self.el_ovlp = numpy.zeros(shape, dtype=numpy.complex128) - self.total_ovlp = numpy.zeros(shape, dtype=numpy.complex128) + self.ovlp_perm = numpy.zeros(shape, dtype=numpy.complex128) - self.buff_names += ["total_ovlp"] + self.buff_names += ["ovlp_perm"] self.buff_size = round(self.set_buff_size_single_walker() / float(self.nwalkers)) self.walker_buffer = numpy.zeros(self.buff_size, dtype=numpy.complex128) @@ -154,7 +154,7 @@ def reortho(self): self.detR[iw] = detR[iw] self.el_ovlp[iw, ...] = self.el_ovlp[iw, ...] / detR[iw] - self.total_ovlp[iw, ...] = self.total_ovlp[iw, ...] / detR[iw] + self.ovlp_perm[iw, ...] = self.ovlp_perm[iw, ...] / detR[iw] self.ovlp[iw] = self.ovlp[iw] / detR[iw] synchronize() From 3edb5e867921139db9f0d94dce5c9331b8f9eb13 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Mon, 11 Mar 2024 13:13:15 -0400 Subject: [PATCH 32/61] calc_greens_function and calc_overlap are now "purer", calc_overlap still sets walkers.ph_ovlp and walkers.el_ovlp --- ipie/addons/eph/trial_wavefunction/coherent_state.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/ipie/addons/eph/trial_wavefunction/coherent_state.py b/ipie/addons/eph/trial_wavefunction/coherent_state.py index a2c88a18..d41cdf72 100644 --- a/ipie/addons/eph/trial_wavefunction/coherent_state.py +++ b/ipie/addons/eph/trial_wavefunction/coherent_state.py @@ -71,11 +71,10 @@ def calc_overlap(self, walkers) -> np.ndarray: ovlp : :class:`np.ndarray` Product of electron and phonon overlap """ - _ = self.calc_phonon_overlap(walkers) - _ = self.calc_electronic_overlap(walkers) - walkers.total_ovlp = walkers.el_ovlp * walkers.ph_ovlp - walkers.ovlp = walkers.total_ovlp - return walkers.ovlp + ph_ovlp = self.calc_phonon_overlap(walkers) + el_ovlp = self.calc_electronic_overlap(walkers) + ovlp = el_ovlp * ph_ovlp + return ovlp def calc_phonon_overlap(self, walkers) -> np.ndarray: r"""Computes phonon overlap, :math:`\langle \phi(\beta)|X_{\mathrm{w}(\tau)}\rangle`. From d66d467f8670c15ddf6780a3120c6889dc6321af Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Mon, 11 Mar 2024 15:14:23 -0400 Subject: [PATCH 33/61] energy estimator unit test --- .../eph/estimators/tests/test_estimators.py | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) create mode 100644 ipie/addons/eph/estimators/tests/test_estimators.py diff --git a/ipie/addons/eph/estimators/tests/test_estimators.py b/ipie/addons/eph/estimators/tests/test_estimators.py new file mode 100644 index 00000000..92a6925c --- /dev/null +++ b/ipie/addons/eph/estimators/tests/test_estimators.py @@ -0,0 +1,30 @@ +import tempfile + +import pytest + +from ipie.addons.eph.estimators.energy import EnergyEstimator +from ipie.addons.eph.utils.testing import gen_random_test_instances + +@pytest.mark.unit +def test_energy_estimator(): + pbc = True + nbasis = 4 + nelec = (2,2) + trial_type = "toyozawa" + nwalkers = 500 + sys, ham, walkers, trial = gen_random_test_instances(nelec, nbasis, nwalkers, trial_type) + estim = EnergyEstimator(sys, ham, trial) + estim.compute_estimator(sys, walkers, ham, trial) + assert len(estim.names) == 6 + assert estim["ENumer"].real == pytest.approx(-3136.7469620055163) + assert estim["ETotal"] == pytest.approx(0.0) + tmp = estim.data.copy() + estim.post_reduce_hook(tmp) + assert tmp[estim.get_index("ETotal")].real == pytest.approx(-6.273493924011032) + assert estim.print_to_stdout + assert estim.ascii_filename == None + assert estim.shape == (6,) + +if __name__ == "__main__": + test_energy_estimator() + From 20d86a253ecc03a4a45124fb6d402fcc3306ca4d Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Mon, 11 Mar 2024 15:15:00 -0400 Subject: [PATCH 34/61] no more hamiltonian needed for trial initialisation. pass w0 instead --- ipie/addons/eph/trial_wavefunction/coherent_state.py | 12 ++++++------ ipie/addons/eph/trial_wavefunction/toyozawa.py | 5 ++--- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/ipie/addons/eph/trial_wavefunction/coherent_state.py b/ipie/addons/eph/trial_wavefunction/coherent_state.py index d41cdf72..64e709a8 100644 --- a/ipie/addons/eph/trial_wavefunction/coherent_state.py +++ b/ipie/addons/eph/trial_wavefunction/coherent_state.py @@ -33,8 +33,8 @@ class CoherentStateTrial(EPhTrialWavefunctionBase): wavefunction : Concatenation of trial determinants of up and down spin spaces and beta specifying the coherent state displacement. - hamiltonian : - Holstein model Hamiltonian + w0 : + Phonon frequency num_elec : Tuple specifying number of up and down spins num_basis : @@ -43,13 +43,13 @@ class CoherentStateTrial(EPhTrialWavefunctionBase): Print level """ - def __init__(self, wavefunction, hamiltonian, num_elec, num_basis, verbose=False): + def __init__(self, wavefunction, w0, num_elec, num_basis, verbose=False): super().__init__(wavefunction, num_elec, num_basis, verbose=verbose) self.num_elec = num_elec self.nup, self.ndown = self.num_elec - self.w0 = hamiltonian.w0 - self.m = hamiltonian.m - self.nsites = hamiltonian.nsites + self.w0 = w0 + self.m = 1 / w0 + self.nsites = self.nbasis self.beta_shift = np.squeeze(wavefunction[:, 0]) self.psia = wavefunction[:, 1 : self.nup + 1] diff --git a/ipie/addons/eph/trial_wavefunction/toyozawa.py b/ipie/addons/eph/trial_wavefunction/toyozawa.py index 02a55669..47eb5f53 100644 --- a/ipie/addons/eph/trial_wavefunction/toyozawa.py +++ b/ipie/addons/eph/trial_wavefunction/toyozawa.py @@ -15,7 +15,6 @@ import numpy as np from typing import Tuple -from ipie.addons.eph.hamiltonians.holstein import HolsteinModel from ipie.addons.eph.walkers.eph_walkers import EPhWalkers from ipie.addons.eph.trial_wavefunction.coherent_state import CoherentStateTrial from ipie.addons.eph.trial_wavefunction.variational.toyozawa_variational import circ_perm @@ -49,12 +48,12 @@ class ToyozawaTrial(CoherentStateTrial): def __init__( self, wavefunction: np.ndarray, - hamiltonian: HolsteinModel, + w0: float, num_elec: Tuple[int, int], num_basis: int, verbose=False, ): - super().__init__(wavefunction, hamiltonian, num_elec, num_basis, verbose=verbose) + super().__init__(wavefunction, w0, num_elec, num_basis, verbose=verbose) self.perms = circ_perm(np.arange(self.nbasis)) self.nperms = self.perms.shape[0] From 24a99d54ef1481decd35ee923136a8e4944f1093 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Mon, 11 Mar 2024 19:02:10 -0400 Subject: [PATCH 35/61] fixed bug with final reshaping causing dimensions to be mixed up --- .../trial_wavefunction/variational/toyozawa_variational.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py b/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py index 5dafdd1c..15aba236 100644 --- a/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py +++ b/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py @@ -163,12 +163,12 @@ def variational_trial_toyozawa( beta_shift = res.x[: hamiltonian.nsites] if system.ndown > 0: psia = res.x[hamiltonian.nsites : hamiltonian.nsites * (system.nup + 1)] - psia = psia.reshape((hamiltonian.nsites, system.nup)) + psia = psia.reshape((system.nup, hamiltonian.nsites)).T psib = res.x[hamiltonian.nsites * (system.nup + 1) :] - psib = psib.reshape((hamiltonian.nsites, system.ndown)) + psib = psib.reshape((system.ndown, hamiltonian.nsites)).T psi = np.column_stack([psia, psib]) else: - psia = res.x[hamiltonian.nsites :].reshape((hamiltonian.nsites, system.nup)) + psia = res.x[hamiltonian.nsites :].reshape((system.nup, hamiltonian.nsites)).T psi = psia return etrial, beta_shift, psi From dfb3a4dd530aaef10cc29d9d9fcbc59c0a2929ad Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Mon, 11 Mar 2024 19:04:35 -0400 Subject: [PATCH 36/61] added variational_energy to toyozawa --- .../eph/trial_wavefunction/eph_trial_base.py | 8 +-- .../addons/eph/trial_wavefunction/toyozawa.py | 61 +++++++++++++++++++ 2 files changed, 65 insertions(+), 4 deletions(-) diff --git a/ipie/addons/eph/trial_wavefunction/eph_trial_base.py b/ipie/addons/eph/trial_wavefunction/eph_trial_base.py index f535d4e0..1ebff9f0 100644 --- a/ipie/addons/eph/trial_wavefunction/eph_trial_base.py +++ b/ipie/addons/eph/trial_wavefunction/eph_trial_base.py @@ -38,16 +38,16 @@ def __init__( ): self.nelec = num_elec self.nbasis = num_basis - self.nalpha, self.nbeta = self.nelec + self.nup, self.ndown = self.nelec self.verbose = verbose self._num_dets = 0 self._max_num_dets = self._num_dets self.ortho_expansion = False self.optimized = True - self.psia = wavefunction[: self.nalpha] - self.psib = wavefunction[self.nalpha : self.nalpha + self.nbeta] - self.beta_shift = wavefunction[self.nalpha + self.nbeta :] + self.psia = wavefunction[: self.nup] + self.psib = wavefunction[self.nup : self.nup + self.ndown] + self.beta_shift = wavefunction[self.nup + self.ndown :] self.compute_trial_energy = False self.energy = None diff --git a/ipie/addons/eph/trial_wavefunction/toyozawa.py b/ipie/addons/eph/trial_wavefunction/toyozawa.py index 47eb5f53..174c32f1 100644 --- a/ipie/addons/eph/trial_wavefunction/toyozawa.py +++ b/ipie/addons/eph/trial_wavefunction/toyozawa.py @@ -19,6 +19,7 @@ from ipie.addons.eph.trial_wavefunction.coherent_state import CoherentStateTrial from ipie.addons.eph.trial_wavefunction.variational.toyozawa_variational import circ_perm from ipie.utils.backend import arraylib as xp +from ipie.estimators.greens_function_single_det import gab_mod_ovlp class ToyozawaTrial(CoherentStateTrial): @@ -57,6 +58,66 @@ def __init__( self.perms = circ_perm(np.arange(self.nbasis)) self.nperms = self.perms.shape[0] + def variational_energy(self, ham, zero_th=1e-12): + r"""Computes the variational energy of the trial, i.e. + + .. math:: + E_T = \frac{\langle\Psi_T|\hat{H}|\Psi_T\rangle}{\langle\Psi_T|\Psi_T\rangle}. + + As the Toyozawa trial wavefunction is a superposition of coherent state trials + the evaluation of :math:`E_T` a naive implementation would scale quadratically + with the number of sites. Here, we exploit the translational symmetry of the + wavefunction to obtain linear scaling. + + Parameters + ---------- + ham: + Hamiltonian + + Returns + ------- + etrial : :class:`float` + Trial energy + """ + num_energy = 0.0 + denom = 0.0 + beta0 = self.beta_shift * np.sqrt(0.5 * ham.m * ham.w0) + for i, perm in enumerate(self.perms): + psia_i = self.psia[perm, :] + beta_i = beta0[perm] + + if self.ndown > 0: + psib_i = self.psib[perm, :] + ov = ( + np.linalg.det(self.psia.T.dot(psia_i)) + * np.linalg.det(self.psib.T.dot(psib_i)) + * np.prod(np.exp(-0.5 * (beta0**2 + beta_i**2) + beta0 * beta_i)) + ) + else: + ov = np.linalg.det(self.psia.T.dot(psia_i)) * np.prod( + np.exp(-0.5 * (beta0**2 + beta_i**2) + beta0 * beta_i) + ) + + if ov < zero_th: + continue + + Ga_i, _, _ = gab_mod_ovlp(self.psia, psia_i) + if self.ndown > 0: + Gb_i, _, _ = gab_mod_ovlp(self.psib, psib_i) + else: + Gb_i = np.zeros_like(Ga_i) + G_i = [Ga_i, Gb_i] + + kinetic = np.sum(ham.T[0] * G_i[0] + ham.T[1] * G_i[1]) + e_ph = ham.w0 * np.sum(beta0 * beta_i) + rho = ham.g_tensor * (G_i[0] + G_i[1]) + e_eph = np.sum(np.dot(rho, beta0 + beta_i)) + num_energy += (kinetic + e_ph + e_eph) * ov + denom += ov + + etrial = num_energy / denom + return etrial + def calc_overlap_perm(self, walkers: EPhWalkers) -> np.ndarray: r"""Computes the product of electron and phonon overlaps for each permutation :math:`\sigma`, From a0cfab51ff56f8dbe045b1eb9e62e40f05bd71d2 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Mon, 11 Mar 2024 22:53:53 -0400 Subject: [PATCH 37/61] added calc_energy to coherent state trial and made it abstractmethod --- .../eph/trial_wavefunction/coherent_state.py | 20 ++++++++++++++++++- .../eph/trial_wavefunction/eph_trial_base.py | 9 ++++----- .../addons/eph/trial_wavefunction/toyozawa.py | 2 +- 3 files changed, 24 insertions(+), 7 deletions(-) diff --git a/ipie/addons/eph/trial_wavefunction/coherent_state.py b/ipie/addons/eph/trial_wavefunction/coherent_state.py index 64e709a8..b440b0d2 100644 --- a/ipie/addons/eph/trial_wavefunction/coherent_state.py +++ b/ipie/addons/eph/trial_wavefunction/coherent_state.py @@ -16,7 +16,7 @@ from ipie.addons.eph.trial_wavefunction.eph_trial_base import EPhTrialWavefunctionBase from ipie.utils.backend import arraylib as xp - +from ipie.estimators.greens_function_single_det import gab_mod_ovlp class CoherentStateTrial(EPhTrialWavefunctionBase): r"""Coherent state trial of the form @@ -55,6 +55,24 @@ def __init__(self, wavefunction, w0, num_elec, num_basis, verbose=False): self.psia = wavefunction[:, 1 : self.nup + 1] self.psib = wavefunction[:, self.nup + 1 : self.nup + self.ndown + 1] + def calc_energy(self, ham) -> float: + r"""""" + Ga, _, _ = gab_mod_ovlp(self.psia, self.psia) + if self.ndown > 0: + Gb, _, _ = gab_mod_ovlp(self.psib, self.psib) + else: + Gb = np.zeros_like(Ga) + G = [Ga, Gb] + + kinetic = np.sum(ham.T[0] * G[0] + ham.T[1] * G[1]) + + e_ph = ham.w0 * np.sum(self.beta_shift ** 2) + rho = ham.g_tensor * (G[0] + G[1]) + e_eph = np.sum(np.dot(rho, 2 * self.beta_shift)) + + etrial = kinetic + e_ph + e_eph + return etrial + def calc_overlap(self, walkers) -> np.ndarray: r"""Computes the product of electron and phonon overlaps, diff --git a/ipie/addons/eph/trial_wavefunction/eph_trial_base.py b/ipie/addons/eph/trial_wavefunction/eph_trial_base.py index 1ebff9f0..e062e169 100644 --- a/ipie/addons/eph/trial_wavefunction/eph_trial_base.py +++ b/ipie/addons/eph/trial_wavefunction/eph_trial_base.py @@ -52,13 +52,12 @@ def __init__( self.compute_trial_energy = False self.energy = None - def set_etrial(self, energy: float) -> None: + def set_etrial(self, ham) -> None: + energy = self.calc_energy(ham) self.energy = energy - # TODO This should be abstract method as well - def calculate_energy(self, system, hamiltonian): - # TODO variational_energy_coherent_state in ipie.estimators.local_energy - ... + @abstractmethod + def calc_energy(self, hamiltonian) -> float: ... @abstractmethod def calc_overlap(self, walkers) -> np.ndarray: ... diff --git a/ipie/addons/eph/trial_wavefunction/toyozawa.py b/ipie/addons/eph/trial_wavefunction/toyozawa.py index 174c32f1..7e2f0d3c 100644 --- a/ipie/addons/eph/trial_wavefunction/toyozawa.py +++ b/ipie/addons/eph/trial_wavefunction/toyozawa.py @@ -58,7 +58,7 @@ def __init__( self.perms = circ_perm(np.arange(self.nbasis)) self.nperms = self.perms.shape[0] - def variational_energy(self, ham, zero_th=1e-12): + def calc_energy(self, ham, zero_th=1e-12): r"""Computes the variational energy of the trial, i.e. .. math:: From cf8f379425b1f048b8ec45bdd6236d6d32f4004a Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Mon, 11 Mar 2024 22:55:53 -0400 Subject: [PATCH 38/61] added unittests for propagation and variational --- .../eph/propagation/tests/test_holstein.py | 3 ++ .../tests/test_variational.py | 52 +++++++++++++++++++ 2 files changed, 55 insertions(+) create mode 100644 ipie/addons/eph/propagation/tests/test_holstein.py create mode 100644 ipie/addons/eph/trial_wavefunction/tests/test_variational.py diff --git a/ipie/addons/eph/propagation/tests/test_holstein.py b/ipie/addons/eph/propagation/tests/test_holstein.py new file mode 100644 index 00000000..3990a898 --- /dev/null +++ b/ipie/addons/eph/propagation/tests/test_holstein.py @@ -0,0 +1,3 @@ +import numpy +import pytest + diff --git a/ipie/addons/eph/trial_wavefunction/tests/test_variational.py b/ipie/addons/eph/trial_wavefunction/tests/test_variational.py new file mode 100644 index 00000000..a027d033 --- /dev/null +++ b/ipie/addons/eph/trial_wavefunction/tests/test_variational.py @@ -0,0 +1,52 @@ +import pytest +import numpy + +from ipie.addons.eph.utils.testing import ( + build_random_trial, + get_random_sys_holstein +) +from ipie.addons.eph.trial_wavefunction.toyozawa import ToyozawaTrial +from ipie.addons.eph.trial_wavefunction.variational.toyozawa_variational import variational_trial_toyozawa +from ipie.addons.eph.trial_wavefunction.coherent_state import CoherentStateTrial +from ipie.addons.eph.trial_wavefunction.variational.coherent_state_variational import variational_trial + +@pytest.mark.unit +def test_variational_energy_toyozawa(): + seed = 7 + numpy.random.seed(seed) + nelec = (2,2) + nbasis = 4 + trial_type = "toyozawa" + pbc = True + sys, ham = get_random_sys_holstein(nelec, nbasis, pbc) + ie = numpy.random.random((nbasis, nelec[0] + nelec[1])) + ip = numpy.random.random(nbasis) + etrial, p, e = variational_trial_toyozawa(ip, ie, ham, sys, verbose=-1) + + wfn = numpy.column_stack([p, e]) + trial = ToyozawaTrial(wavefunction=wfn, w0=ham.w0, num_elec=nelec, num_basis=nbasis) + trial.set_etrial(ham) + assert etrial == pytest.approx(trial.energy) + +@pytest.mark.unit +def test_variational_energy_coherent_state(): + seed = 7 + numpy.random.seed(seed) + nelec = (2,2) + nbasis = 4 + trial_type = "coherent_state" + pbc = True + sys, ham = get_random_sys_holstein(nelec, nbasis, pbc) + ie = numpy.random.random((nbasis, nelec[0] + nelec[1])) + ip = numpy.random.random(nbasis) + etrial, p, e = variational_trial(ip, ie, ham, sys) + + wfn = numpy.column_stack([p, e]) + trial = CoherentStateTrial(wavefunction=wfn, w0=ham.w0, num_elec=nelec, num_basis=nbasis) + trial.set_etrial(ham) + assert etrial == pytest.approx(trial.energy) + + +if __name__ == "__main__": + test_variational_energy_toyozawa() + test_variational_energy_coherent_state() From 28d40df5e0f305683634d39c4e193fa1f5893129 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Mon, 11 Mar 2024 22:57:08 -0400 Subject: [PATCH 39/61] utils for unittests --- ipie/addons/eph/utils/__init__.py | 13 +++++++ ipie/addons/eph/utils/testing.py | 58 +++++++++++++++++++++++++++++++ 2 files changed, 71 insertions(+) create mode 100644 ipie/addons/eph/utils/__init__.py create mode 100644 ipie/addons/eph/utils/testing.py diff --git a/ipie/addons/eph/utils/__init__.py b/ipie/addons/eph/utils/__init__.py new file mode 100644 index 00000000..e2aed039 --- /dev/null +++ b/ipie/addons/eph/utils/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2022 The ipie Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/ipie/addons/eph/utils/testing.py b/ipie/addons/eph/utils/testing.py new file mode 100644 index 00000000..6d1c6157 --- /dev/null +++ b/ipie/addons/eph/utils/testing.py @@ -0,0 +1,58 @@ +import numpy + +from ipie.systems import Generic +from ipie.qmc.afqmc import AFQMC + +from ipie.addons.eph.hamiltonians.holstein import HolsteinModel +from ipie.addons.eph.propagation.holstein import HolsteinPropagator +from ipie.addons.eph.walkers.eph_walkers import EPhWalkers +from ipie.addons.eph.trial_wavefunction.eph_trial_base import EPhTrialWavefunctionBase +from ipie.addons.eph.trial_wavefunction.toyozawa import ToyozawaTrial +from ipie.addons.eph.trial_wavefunction.coherent_state import CoherentStateTrial + +def get_random_sys_holstein(nelec, nbasis, pbc): + sys = Generic(nelec=nelec) + g = numpy.random.rand() + t = numpy.random.rand() + w0 = numpy.random.rand() + ham = HolsteinModel(g=g, t=t, w0=w0, nsites=nbasis, pbc=pbc) + ham.build() + return sys, ham + +def get_random_wavefunction(nelec, nbasis): + init = numpy.random.random((nbasis, (nelec[0] + nelec[1] + 1))) + return init + +def build_random_toyozawa_trial(nelec, nbasis, w0): + wfn = get_random_wavefunction(nelec, nbasis) + trial = ToyozawaTrial( + wavefunction=wfn, w0=w0, num_elec=nelec, num_basis=nbasis + ) + return trial + +def build_random_coherent_state_trial(nelec, nbasis, w0): + wfn = get_random_wavefunction(nelec, nbasis) + trial = CoherentStateTrial( + wavefunction=wfn, w0=w0, num_elec=nelec, num_basis=nbasis + ) + return trial + +def build_random_trial(nelec, nbasis, w0, trial_type): + if trial_type == "coherent_state": + return build_random_coherent_state_trial(nelec, nbasis, w0) + elif trial_type == "toyozawa": + return build_random_toyozawa_trial(nelec, nbasis, w0) + else: + raise ValueError(f"Unkown trial type: {trial_type}") + +def gen_random_test_instances(nelec, nbasis, nwalkers, trial_type, seed=7): + numpy.random.seed(seed) + + wfn = get_random_wavefunction(nelec, nbasis) + sys, ham = get_random_sys_holstein(nelec, nbasis, True) + trial = build_random_trial(nelec, nbasis, ham.w0, trial_type) + walkers = EPhWalkers(wfn, nelec[0], nelec[1], nbasis, nwalkers) + walkers.build(trial) + + return sys, ham, walkers, trial + From 298ad97827d70d8a938c3758ec538740633079b8 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Mon, 11 Mar 2024 22:58:46 -0400 Subject: [PATCH 40/61] reformatting --- .../eph/propagation/tests/test_holstein.py | 1 - .../eph/trial_wavefunction/coherent_state.py | 9 ++++--- .../tests/test_variational.py | 25 +++++++++++-------- 3 files changed, 19 insertions(+), 16 deletions(-) diff --git a/ipie/addons/eph/propagation/tests/test_holstein.py b/ipie/addons/eph/propagation/tests/test_holstein.py index 3990a898..a811c076 100644 --- a/ipie/addons/eph/propagation/tests/test_holstein.py +++ b/ipie/addons/eph/propagation/tests/test_holstein.py @@ -1,3 +1,2 @@ import numpy import pytest - diff --git a/ipie/addons/eph/trial_wavefunction/coherent_state.py b/ipie/addons/eph/trial_wavefunction/coherent_state.py index b440b0d2..1f29108d 100644 --- a/ipie/addons/eph/trial_wavefunction/coherent_state.py +++ b/ipie/addons/eph/trial_wavefunction/coherent_state.py @@ -18,6 +18,7 @@ from ipie.utils.backend import arraylib as xp from ipie.estimators.greens_function_single_det import gab_mod_ovlp + class CoherentStateTrial(EPhTrialWavefunctionBase): r"""Coherent state trial of the form @@ -33,7 +34,7 @@ class CoherentStateTrial(EPhTrialWavefunctionBase): wavefunction : Concatenation of trial determinants of up and down spin spaces and beta specifying the coherent state displacement. - w0 : + w0 : Phonon frequency num_elec : Tuple specifying number of up and down spins @@ -63,10 +64,10 @@ def calc_energy(self, ham) -> float: else: Gb = np.zeros_like(Ga) G = [Ga, Gb] - + kinetic = np.sum(ham.T[0] * G[0] + ham.T[1] * G[1]) - - e_ph = ham.w0 * np.sum(self.beta_shift ** 2) + + e_ph = ham.w0 * np.sum(self.beta_shift**2) rho = ham.g_tensor * (G[0] + G[1]) e_eph = np.sum(np.dot(rho, 2 * self.beta_shift)) diff --git a/ipie/addons/eph/trial_wavefunction/tests/test_variational.py b/ipie/addons/eph/trial_wavefunction/tests/test_variational.py index a027d033..9e72e298 100644 --- a/ipie/addons/eph/trial_wavefunction/tests/test_variational.py +++ b/ipie/addons/eph/trial_wavefunction/tests/test_variational.py @@ -1,20 +1,22 @@ import pytest -import numpy +import numpy -from ipie.addons.eph.utils.testing import ( - build_random_trial, - get_random_sys_holstein -) +from ipie.addons.eph.utils.testing import build_random_trial, get_random_sys_holstein from ipie.addons.eph.trial_wavefunction.toyozawa import ToyozawaTrial -from ipie.addons.eph.trial_wavefunction.variational.toyozawa_variational import variational_trial_toyozawa +from ipie.addons.eph.trial_wavefunction.variational.toyozawa_variational import ( + variational_trial_toyozawa, +) from ipie.addons.eph.trial_wavefunction.coherent_state import CoherentStateTrial -from ipie.addons.eph.trial_wavefunction.variational.coherent_state_variational import variational_trial +from ipie.addons.eph.trial_wavefunction.variational.coherent_state_variational import ( + variational_trial, +) + @pytest.mark.unit def test_variational_energy_toyozawa(): seed = 7 numpy.random.seed(seed) - nelec = (2,2) + nelec = (2, 2) nbasis = 4 trial_type = "toyozawa" pbc = True @@ -22,17 +24,18 @@ def test_variational_energy_toyozawa(): ie = numpy.random.random((nbasis, nelec[0] + nelec[1])) ip = numpy.random.random(nbasis) etrial, p, e = variational_trial_toyozawa(ip, ie, ham, sys, verbose=-1) - + wfn = numpy.column_stack([p, e]) trial = ToyozawaTrial(wavefunction=wfn, w0=ham.w0, num_elec=nelec, num_basis=nbasis) trial.set_etrial(ham) assert etrial == pytest.approx(trial.energy) - + + @pytest.mark.unit def test_variational_energy_coherent_state(): seed = 7 numpy.random.seed(seed) - nelec = (2,2) + nelec = (2, 2) nbasis = 4 trial_type = "coherent_state" pbc = True From 51b48a150845b122e101257f9758364620c7233d Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Mon, 11 Mar 2024 23:01:20 -0400 Subject: [PATCH 41/61] added verbosity for toyozawa_variational --- examples/13-1d_holstein/run_afqmc.py | 85 ------------------- .../variational/__init__.py | 13 +++ .../variational/toyozawa_variational.py | 4 +- 3 files changed, 15 insertions(+), 87 deletions(-) delete mode 100644 examples/13-1d_holstein/run_afqmc.py create mode 100644 ipie/addons/eph/trial_wavefunction/variational/__init__.py diff --git a/examples/13-1d_holstein/run_afqmc.py b/examples/13-1d_holstein/run_afqmc.py deleted file mode 100644 index 9171cfaa..00000000 --- a/examples/13-1d_holstein/run_afqmc.py +++ /dev/null @@ -1,85 +0,0 @@ -import numpy as np -np.random.seed(125) -from mpi4py import MPI - -from ipie.qmc.afqmc import AFQMC -from ipie.systems import Generic -from ipie.addons.eph.hamiltonians.holstein import HolsteinModel -from ipie.addons.eph.trial_wavefunction.toyozawa import ToyozawaTrial -from ipie.addons.eph.trial_wavefunction.variational.toyozawa_variational import variational_trial_toyozawa -from ipie.addons.eph.walkers.eph_walkers import EPhWalkers -from ipie.addons.eph.estimators.energy import EnergyEstimator - -#System Parameters -nup = 2 -ndown = 2 -nelec = (nup, ndown) - -#Hamiltonian Parameters -g = 2. -t = 1. -w0 = 4. -nsites = 4 -pbc = True - -#Walker Parameters & Setup -comm = MPI.COMM_WORLD -nwalkers = 1000 // comm.size - -#Setup initial guess for variational optimization -initial_electron = np.random.random((nsites, nup + ndown)) -initial_phonons = np.ones(nsites) * 0.1 - -#System and Hamiltonian setup -system = Generic(nelec) -ham = HolsteinModel(g=g, t=t, w0=w0, nsites=nsites, pbc=pbc) -ham.build() - -#Variational procedure -etrial, beta_shift, el_trial = variational_trial_toyozawa( - initial_phonons, initial_electron, ham, system -) -wavefunction = np.column_stack([beta_shift, el_trial]) - -#Setup trial -trial = ToyozawaTrial( - wavefunction=wavefunction, - hamiltonian=ham, - num_elec=[nup, ndown], - num_basis=nsites -) -trial.set_etrial(etrial) - -#Setup walkers -walkers = EPhWalkers( - initial_walker=wavefunction, - nup=nup, - ndown=ndown, - nbasis=nsites, - nwalkers=nwalkers -) -walkers.build(trial) - -num_steps_per_block = 10 -num_blocks = 10000 -add_est = { - "energy": EnergyEstimator( - system=system, ham=ham, trial=trial - ) -} - -seed = 125 - -# Note nwalkers specifies the number of walkers on each CPU -ephqmc = AFQMC.build( - num_elec=nelec, - hamiltonian=ham, - trial_wavefunction=trial, - walkers=walkers, - num_walkers=nwalkers, - seed=seed, - num_steps_per_block=num_steps_per_block, - num_blocks=num_blocks, -) -ephqmc.run(additional_estimators=add_est, verbose=False) - diff --git a/ipie/addons/eph/trial_wavefunction/variational/__init__.py b/ipie/addons/eph/trial_wavefunction/variational/__init__.py new file mode 100644 index 00000000..e2aed039 --- /dev/null +++ b/ipie/addons/eph/trial_wavefunction/variational/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2022 The ipie Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py b/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py index 15aba236..91d01dfc 100644 --- a/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py +++ b/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py @@ -117,7 +117,7 @@ def circ_perm(lst: np.ndarray) -> np.ndarray: def variational_trial_toyozawa( - shift_init: np.ndarray, electron_init: np.ndarray, hamiltonian, system + shift_init: np.ndarray, electron_init: np.ndarray, hamiltonian, system, verbose=2 ): psi = electron_init.T.real.ravel() shift = shift_init.real @@ -148,7 +148,7 @@ def variational_trial_toyozawa( method="L-BFGS-B", options={ "maxls": 20, - "iprint": 2, + "iprint": verbose, "gtol": 1e-10, "eps": 1e-10, "maxiter": 15000, From 6fb14af3935db6f30574b3824164509ff6c8359d Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Mon, 11 Mar 2024 23:02:21 -0400 Subject: [PATCH 42/61] fixed bug with local_energy, where coherent state expectation value of X was wrong by factor of sqrt(2) --- .../variational/coherent_state_variational.py | 22 ++++++------------- 1 file changed, 7 insertions(+), 15 deletions(-) diff --git a/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py b/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py index dd7a6eb4..a97fa06f 100644 --- a/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py +++ b/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py @@ -15,7 +15,6 @@ import numpy as np from scipy.optimize import basinhopping -from ipie.legacy.trial_wavefunction.harmonic_oscillator import HarmonicOscillator from ipie.systems.generic import Generic from ipie.addons.eph.hamiltonians.holstein import HolsteinModel from ipie.addons.eph.trial_wavefunction.variational.estimators import gab @@ -30,23 +29,17 @@ def local_energy( w: float, g: float, nsites: int, - Lap: np.ndarray, T: np.ndarray, nup: int, ndown: int, ) -> np.ndarray: - kinetic_contrib = jax.numpy.einsum("ij->", T[0] * G[0]) - if ndown > 0: - kinetic_contrib += jax.numpy.einsum("ij->", T[1] * G[1]) - + kinetic = np.sum(T[0] * G[0] + T[1] * G[1]) rho = G[0].diagonal() + G[1].diagonal() - el_ph_contrib = -g * np.sqrt(2 * m * w) * np.sum(rho * X) - - phonon_contrib = m * w**2 * np.sum(X * X) / 2 - phonon_contrib += -0.5 * np.sum(Lap) / m - 0.5 * w * nsites + e_eph = -g * 2 * np.sqrt(m * w) * np.sum(rho * X) + e_ph = m * w**2 * np.sum(X * X) - local_energy = kinetic_contrib + el_ph_contrib + phonon_contrib + local_energy = kinetic + e_eph + e_ph return local_energy @@ -71,10 +64,8 @@ def objective_function( G = [Ga, Gb] - phi = HarmonicOscillator(m, w0, order=0, shift=shift) - Lap = phi.laplacian(shift) - - etot = local_energy(shift, G, m, w0, g, nbasis, Lap, T, nup, ndown) + etot = local_energy(shift, G, m, w0, g, nbasis, T, nup, ndown) + return etot.real @@ -137,6 +128,7 @@ def variational_trial(init_phonons: np.ndarray, init_electron: np.ndarray, hamil etrial = res.fun + beta_shift = res.x[: hamiltonian.nsites] if system.ndown > 0: psia = res.x[hamiltonian.nsites : hamiltonian.nsites * (system.nup + 1)] From 5a0a70aa1115afa698a242c85cf9c4c7bb712404 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Mon, 11 Mar 2024 23:02:50 -0400 Subject: [PATCH 43/61] added g_tensor in anticipation of generalization --- ipie/addons/eph/hamiltonians/holstein.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/ipie/addons/eph/hamiltonians/holstein.py b/ipie/addons/eph/hamiltonians/holstein.py index 56ad65b7..c188bd9e 100644 --- a/ipie/addons/eph/hamiltonians/holstein.py +++ b/ipie/addons/eph/hamiltonians/holstein.py @@ -67,3 +67,5 @@ def build(self) -> None: self.T *= -self.t self.T = [self.T.copy(), self.T.copy()] + + self.g_tensor = -self.g * numpy.eye(self.nsites) From 9dfd155a40f763a9accbb7a3c395bbfa9f499f6c Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Mon, 11 Mar 2024 23:04:19 -0400 Subject: [PATCH 44/61] unittests for ephwalkers --- .../eph/walkers/tests/test_ephwalkers.py | 40 +++++++++++++++++++ 1 file changed, 40 insertions(+) create mode 100644 ipie/addons/eph/walkers/tests/test_ephwalkers.py diff --git a/ipie/addons/eph/walkers/tests/test_ephwalkers.py b/ipie/addons/eph/walkers/tests/test_ephwalkers.py new file mode 100644 index 00000000..e3477153 --- /dev/null +++ b/ipie/addons/eph/walkers/tests/test_ephwalkers.py @@ -0,0 +1,40 @@ +import numpy +import pytest +from ipie.addons.eph.walkers.eph_walkers import EPhWalkers +from ipie.addons.eph.utils.testing import get_random_wavefunction, build_random_trial + + +@pytest.mark.unit +def test_ephwalkers(): + nwalkers = 100 + nelec = (3, 3) + nbasis = 6 + wfn = get_random_wavefunction(nelec, nbasis) + walkers = EPhWalkers(wfn, nelec[0], nelec[1], nbasis, nwalkers) + assert numpy.allclose(walkers.phonon_disp, wfn[:, 0]) + assert numpy.allclose(walkers.phia, wfn[:, 1 : 1 + nelec[0]].reshape(nbasis, nelec[0])) + assert numpy.allclose(walkers.phib, wfn[:, 1 + nelec[0] :].reshape(nbasis, nelec[1])) + + +@pytest.mark.unit +def test_overlap_init(): + nwalkers = 100 + nelec = (3, 3) + nbasis = 6 + seed = 7 + numpy.random.seed(7) + w0 = numpy.random.rand() + trial_type = "coherent_state" + trial = build_random_trial(nelec, nbasis, w0, trial_type) + wfn = get_random_wavefunction(nelec, nbasis) + walkers = EPhWalkers(wfn, nelec[0], nelec[1], nbasis, nwalkers) + walkers.build(trial) + assert len(walkers.buff_names) == 11 + assert walkers.ovlp[0].real == pytest.approx(0.0007324519852172784) + assert walkers.ph_ovlp[0].real == pytest.approx(0.7141097126634587) + assert walkers.el_ovlp[0].real == pytest.approx(0.0010256855105434813) + + +if __name__ == "__main__": + test_ephwalkers() + test_overlap_init() From c355ffa0f2a31bb5d7aaadade0f0454417fdd644 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Tue, 12 Mar 2024 00:13:40 -0400 Subject: [PATCH 45/61] unittests for trial methods --- .../tests/test_coherent_state.py | 34 ++++++++++++++++++ .../trial_wavefunction/tests/test_toyozawa.py | 36 +++++++++++++++++++ .../tests/test_variational.py | 4 +-- 3 files changed, 71 insertions(+), 3 deletions(-) create mode 100644 ipie/addons/eph/trial_wavefunction/tests/test_coherent_state.py create mode 100644 ipie/addons/eph/trial_wavefunction/tests/test_toyozawa.py diff --git a/ipie/addons/eph/trial_wavefunction/tests/test_coherent_state.py b/ipie/addons/eph/trial_wavefunction/tests/test_coherent_state.py new file mode 100644 index 00000000..a6134a5e --- /dev/null +++ b/ipie/addons/eph/trial_wavefunction/tests/test_coherent_state.py @@ -0,0 +1,34 @@ +import pytest +import numpy + +from ipie.addons.eph.utils.testing import build_random_trial, get_random_wavefunction +from ipie.addons.eph.walkers.eph_walkers import EPhWalkers + + +@pytest.mark.unit +def test_trial_methods(): + seed = 7 + numpy.random.seed(seed) + nelec = (2, 2) + nbasis = 4 + trial_type = "coherent_state" + pbc = True + w0 = numpy.random.rand() + nwalkers = 100 + + trial = build_random_trial(nelec, nbasis, w0, trial_type) + + wfn = get_random_wavefunction(nelec, nbasis) + walkers = EPhWalkers(wfn, nelec[0], nelec[1], nbasis, nwalkers) + walkers.build(trial) + + assert trial.calc_overlap(walkers)[0] == pytest.approx(0.0008183683599516786) + assert trial.calc_phonon_overlap(walkers)[0] == pytest.approx(0.8042312422253469) + assert trial.calc_phonon_gradient(walkers)[0, 0] == pytest.approx(-0.170210708173531) + assert trial.calc_phonon_laplacian(walkers)[0] == pytest.approx(-3.5642631271035823) + assert trial.calc_electronic_overlap(walkers)[0] == pytest.approx(0.0010175784239458462) + assert trial.calc_greens_function(walkers)[0][0, 0, 0] == pytest.approx(2.759966097679107) + + +if __name__ == "__main__": + test_trial_methods() diff --git a/ipie/addons/eph/trial_wavefunction/tests/test_toyozawa.py b/ipie/addons/eph/trial_wavefunction/tests/test_toyozawa.py new file mode 100644 index 00000000..9b457f2c --- /dev/null +++ b/ipie/addons/eph/trial_wavefunction/tests/test_toyozawa.py @@ -0,0 +1,36 @@ +import pytest +import numpy + +from ipie.addons.eph.utils.testing import build_random_trial, get_random_wavefunction +from ipie.addons.eph.walkers.eph_walkers import EPhWalkers + + +@pytest.mark.unit +def test_trial_methods(): + seed = 9 + numpy.random.seed(seed) + nelec = (2, 2) + nbasis = 4 + trial_type = "toyozawa" + pbc = True + w0 = numpy.random.rand() + nwalkers = 100 + + trial = build_random_trial(nelec, nbasis, w0, trial_type) + + wfn = get_random_wavefunction(nelec, nbasis) + walkers = EPhWalkers(wfn, nelec[0], nelec[1], nbasis, nwalkers) + walkers.build(trial) + + assert trial.calc_overlap(walkers)[0] == pytest.approx(0.05496697699720597) + assert trial.calc_phonon_overlap(walkers)[0] == pytest.approx(3.3244959521904325) + assert trial.calc_phonon_gradient(walkers)[0, 0] == pytest.approx(-0.028056870680617366) + assert trial.calc_phonon_laplacian(walkers, walkers.ovlp_perm)[0] == pytest.approx( + -3.672010193624128 + ) + assert trial.calc_electronic_overlap(walkers)[0] == pytest.approx(0.06480736261172242) + assert trial.calc_greens_function(walkers)[0][0, 0, 0] == pytest.approx(0.3341677951334292) + + +if __name__ == "__main__": + test_trial_methods() diff --git a/ipie/addons/eph/trial_wavefunction/tests/test_variational.py b/ipie/addons/eph/trial_wavefunction/tests/test_variational.py index 9e72e298..8eb2fc09 100644 --- a/ipie/addons/eph/trial_wavefunction/tests/test_variational.py +++ b/ipie/addons/eph/trial_wavefunction/tests/test_variational.py @@ -1,7 +1,7 @@ import pytest import numpy -from ipie.addons.eph.utils.testing import build_random_trial, get_random_sys_holstein +from ipie.addons.eph.utils.testing import get_random_sys_holstein from ipie.addons.eph.trial_wavefunction.toyozawa import ToyozawaTrial from ipie.addons.eph.trial_wavefunction.variational.toyozawa_variational import ( variational_trial_toyozawa, @@ -18,7 +18,6 @@ def test_variational_energy_toyozawa(): numpy.random.seed(seed) nelec = (2, 2) nbasis = 4 - trial_type = "toyozawa" pbc = True sys, ham = get_random_sys_holstein(nelec, nbasis, pbc) ie = numpy.random.random((nbasis, nelec[0] + nelec[1])) @@ -37,7 +36,6 @@ def test_variational_energy_coherent_state(): numpy.random.seed(seed) nelec = (2, 2) nbasis = 4 - trial_type = "coherent_state" pbc = True sys, ham = get_random_sys_holstein(nelec, nbasis, pbc) ie = numpy.random.random((nbasis, nelec[0] + nelec[1])) From 40c5ed534b1ffe9183f6909b7c8ea00d34c8eedd Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Tue, 12 Mar 2024 00:14:06 -0400 Subject: [PATCH 46/61] docstring --- .../eph/trial_wavefunction/coherent_state.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/ipie/addons/eph/trial_wavefunction/coherent_state.py b/ipie/addons/eph/trial_wavefunction/coherent_state.py index 1f29108d..b2a43aee 100644 --- a/ipie/addons/eph/trial_wavefunction/coherent_state.py +++ b/ipie/addons/eph/trial_wavefunction/coherent_state.py @@ -57,7 +57,21 @@ def __init__(self, wavefunction, w0, num_elec, num_basis, verbose=False): self.psib = wavefunction[:, self.nup + 1 : self.nup + self.ndown + 1] def calc_energy(self, ham) -> float: - r"""""" + r"""Computes the variational energy of the trial, + + .. math:: + E = \frac{\langle \Psi_T |\hat{H}|\Psi_T\rangle}{\langle \Psi_T |\Psi_T\rangle}. + + Parameters + ---------- + ham : :class:`HolsteinModel` + Holstein model + + Returns + ------- + etrial : :class:`float` + Variational trial energy + """ Ga, _, _ = gab_mod_ovlp(self.psia, self.psia) if self.ndown > 0: Gb, _, _ = gab_mod_ovlp(self.psib, self.psib) From 3b1256fc14f24500f1c2d1c8ffd86e831c6eb77f Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Tue, 12 Mar 2024 00:35:32 -0400 Subject: [PATCH 47/61] reformatting --- .../eph/estimators/tests/test_estimators.py | 5 +++-- ipie/addons/eph/hamiltonians/holstein.py | 1 - .../eph/trial_wavefunction/coherent_state.py | 6 ++--- .../eph/trial_wavefunction/eph_trial_base.py | 2 +- .../addons/eph/trial_wavefunction/toyozawa.py | 2 +- .../variational/coherent_state_variational.py | 5 +---- .../variational/estimators.py | 1 - .../variational/toyozawa_variational.py | 10 ++------- ipie/addons/eph/utils/testing.py | 22 +++++++++---------- 9 files changed, 21 insertions(+), 33 deletions(-) diff --git a/ipie/addons/eph/estimators/tests/test_estimators.py b/ipie/addons/eph/estimators/tests/test_estimators.py index 92a6925c..0eef56a0 100644 --- a/ipie/addons/eph/estimators/tests/test_estimators.py +++ b/ipie/addons/eph/estimators/tests/test_estimators.py @@ -5,11 +5,12 @@ from ipie.addons.eph.estimators.energy import EnergyEstimator from ipie.addons.eph.utils.testing import gen_random_test_instances + @pytest.mark.unit def test_energy_estimator(): pbc = True nbasis = 4 - nelec = (2,2) + nelec = (2, 2) trial_type = "toyozawa" nwalkers = 500 sys, ham, walkers, trial = gen_random_test_instances(nelec, nbasis, nwalkers, trial_type) @@ -25,6 +26,6 @@ def test_energy_estimator(): assert estim.ascii_filename == None assert estim.shape == (6,) + if __name__ == "__main__": test_energy_estimator() - diff --git a/ipie/addons/eph/hamiltonians/holstein.py b/ipie/addons/eph/hamiltonians/holstein.py index c188bd9e..a951fbd0 100644 --- a/ipie/addons/eph/hamiltonians/holstein.py +++ b/ipie/addons/eph/hamiltonians/holstein.py @@ -13,7 +13,6 @@ # limitations under the License. import numpy -from ipie.utils.backend import arraylib as xp class HolsteinModel: diff --git a/ipie/addons/eph/trial_wavefunction/coherent_state.py b/ipie/addons/eph/trial_wavefunction/coherent_state.py index b2a43aee..6cf963fc 100644 --- a/ipie/addons/eph/trial_wavefunction/coherent_state.py +++ b/ipie/addons/eph/trial_wavefunction/coherent_state.py @@ -58,10 +58,10 @@ def __init__(self, wavefunction, w0, num_elec, num_basis, verbose=False): def calc_energy(self, ham) -> float: r"""Computes the variational energy of the trial, - - .. math:: + + .. math:: E = \frac{\langle \Psi_T |\hat{H}|\Psi_T\rangle}{\langle \Psi_T |\Psi_T\rangle}. - + Parameters ---------- ham : :class:`HolsteinModel` diff --git a/ipie/addons/eph/trial_wavefunction/eph_trial_base.py b/ipie/addons/eph/trial_wavefunction/eph_trial_base.py index e062e169..c076876c 100644 --- a/ipie/addons/eph/trial_wavefunction/eph_trial_base.py +++ b/ipie/addons/eph/trial_wavefunction/eph_trial_base.py @@ -57,7 +57,7 @@ def set_etrial(self, ham) -> None: self.energy = energy @abstractmethod - def calc_energy(self, hamiltonian) -> float: ... + def calc_energy(self, ham) -> float: ... @abstractmethod def calc_overlap(self, walkers) -> np.ndarray: ... diff --git a/ipie/addons/eph/trial_wavefunction/toyozawa.py b/ipie/addons/eph/trial_wavefunction/toyozawa.py index 7e2f0d3c..d248a164 100644 --- a/ipie/addons/eph/trial_wavefunction/toyozawa.py +++ b/ipie/addons/eph/trial_wavefunction/toyozawa.py @@ -82,7 +82,7 @@ def calc_energy(self, ham, zero_th=1e-12): num_energy = 0.0 denom = 0.0 beta0 = self.beta_shift * np.sqrt(0.5 * ham.m * ham.w0) - for i, perm in enumerate(self.perms): + for perm in self.perms: psia_i = self.psia[perm, :] beta_i = beta0[perm] diff --git a/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py b/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py index a97fa06f..6e413f23 100644 --- a/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py +++ b/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py @@ -15,8 +15,6 @@ import numpy as np from scipy.optimize import basinhopping -from ipie.systems.generic import Generic -from ipie.addons.eph.hamiltonians.holstein import HolsteinModel from ipie.addons.eph.trial_wavefunction.variational.estimators import gab import jax @@ -65,7 +63,7 @@ def objective_function( G = [Ga, Gb] etot = local_energy(shift, G, m, w0, g, nbasis, T, nup, ndown) - + return etot.real @@ -128,7 +126,6 @@ def variational_trial(init_phonons: np.ndarray, init_electron: np.ndarray, hamil etrial = res.fun - beta_shift = res.x[: hamiltonian.nsites] if system.ndown > 0: psia = res.x[hamiltonian.nsites : hamiltonian.nsites * (system.nup + 1)] diff --git a/ipie/addons/eph/trial_wavefunction/variational/estimators.py b/ipie/addons/eph/trial_wavefunction/variational/estimators.py index 29e1bb3c..f97f7c39 100644 --- a/ipie/addons/eph/trial_wavefunction/variational/estimators.py +++ b/ipie/addons/eph/trial_wavefunction/variational/estimators.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import numpy as np import jax.numpy as npj from jax.config import config diff --git a/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py b/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py index 91d01dfc..98d46469 100644 --- a/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py +++ b/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py @@ -13,9 +13,7 @@ # limitations under the License. import numpy as np -from scipy.optimize import minimize, basinhopping -from ipie.systems import Generic -from ipie.addons.eph.hamiltonians.holstein import HolsteinModel +from scipy.optimize import minimize from ipie.addons.eph.trial_wavefunction.variational.estimators import gab import jax @@ -43,8 +41,6 @@ def objective_function_toyozawa_mo( nbsf = nbasis nocca = nup noccb = ndown - nvira = nbasis - nocca - nvirb = nbasis - noccb psi0a = npj.array(x[nbsf : nbsf + nbsf * nocca], dtype=npj.float64) psi0a = psi0a.reshape((nocca, nbsf)).T @@ -53,14 +49,12 @@ def objective_function_toyozawa_mo( psi0b = npj.array(x[nbsf + nbsf * nocca :], dtype=npj.float64) psi0b = psi0b.reshape((noccb, nbsf)).T - nperms = len(perms) - num_energy = 0.0 denom = 0.0 beta0 = shift0 * npj.sqrt(m * w0 / 2.0) - for i, perm in enumerate(perms): + for perm in perms: psia_j = psi0a[perm, :] beta_j = beta0[npj.array(perm)] diff --git a/ipie/addons/eph/utils/testing.py b/ipie/addons/eph/utils/testing.py index 6d1c6157..25d42deb 100644 --- a/ipie/addons/eph/utils/testing.py +++ b/ipie/addons/eph/utils/testing.py @@ -1,15 +1,13 @@ import numpy from ipie.systems import Generic -from ipie.qmc.afqmc import AFQMC from ipie.addons.eph.hamiltonians.holstein import HolsteinModel -from ipie.addons.eph.propagation.holstein import HolsteinPropagator from ipie.addons.eph.walkers.eph_walkers import EPhWalkers -from ipie.addons.eph.trial_wavefunction.eph_trial_base import EPhTrialWavefunctionBase from ipie.addons.eph.trial_wavefunction.toyozawa import ToyozawaTrial from ipie.addons.eph.trial_wavefunction.coherent_state import CoherentStateTrial + def get_random_sys_holstein(nelec, nbasis, pbc): sys = Generic(nelec=nelec) g = numpy.random.rand() @@ -19,24 +17,24 @@ def get_random_sys_holstein(nelec, nbasis, pbc): ham.build() return sys, ham + def get_random_wavefunction(nelec, nbasis): init = numpy.random.random((nbasis, (nelec[0] + nelec[1] + 1))) return init + def build_random_toyozawa_trial(nelec, nbasis, w0): wfn = get_random_wavefunction(nelec, nbasis) - trial = ToyozawaTrial( - wavefunction=wfn, w0=w0, num_elec=nelec, num_basis=nbasis - ) + trial = ToyozawaTrial(wavefunction=wfn, w0=w0, num_elec=nelec, num_basis=nbasis) return trial + def build_random_coherent_state_trial(nelec, nbasis, w0): wfn = get_random_wavefunction(nelec, nbasis) - trial = CoherentStateTrial( - wavefunction=wfn, w0=w0, num_elec=nelec, num_basis=nbasis - ) + trial = CoherentStateTrial(wavefunction=wfn, w0=w0, num_elec=nelec, num_basis=nbasis) return trial + def build_random_trial(nelec, nbasis, w0, trial_type): if trial_type == "coherent_state": return build_random_coherent_state_trial(nelec, nbasis, w0) @@ -45,14 +43,14 @@ def build_random_trial(nelec, nbasis, w0, trial_type): else: raise ValueError(f"Unkown trial type: {trial_type}") + def gen_random_test_instances(nelec, nbasis, nwalkers, trial_type, seed=7): numpy.random.seed(seed) - + wfn = get_random_wavefunction(nelec, nbasis) sys, ham = get_random_sys_holstein(nelec, nbasis, True) trial = build_random_trial(nelec, nbasis, ham.w0, trial_type) walkers = EPhWalkers(wfn, nelec[0], nelec[1], nbasis, nwalkers) walkers.build(trial) - return sys, ham, walkers, trial - + return sys, ham, walkers, trial From 27b53659185ed1a69f5f2ce30f49bf64be7a63f7 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Tue, 12 Mar 2024 11:50:29 -0400 Subject: [PATCH 48/61] light reformatting --- ipie/addons/eph/estimators/energy.py | 2 +- ipie/addons/eph/propagation/holstein.py | 9 ++++++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/ipie/addons/eph/estimators/energy.py b/ipie/addons/eph/estimators/energy.py index 5aaf8bc9..4d4195cb 100644 --- a/ipie/addons/eph/estimators/energy.py +++ b/ipie/addons/eph/estimators/energy.py @@ -29,7 +29,7 @@ def local_energy( hamiltonian: HolsteinModel, walkers: EPhWalkers, trial: EPhTrialWavefunctionBase, -): +) -> np.ndarray: return local_energy_holstein(system, hamiltonian, walkers, trial) diff --git a/ipie/addons/eph/propagation/holstein.py b/ipie/addons/eph/propagation/holstein.py index a5c9c19d..cfc2fce9 100644 --- a/ipie/addons/eph/propagation/holstein.py +++ b/ipie/addons/eph/propagation/holstein.py @@ -25,7 +25,9 @@ from ipie.propagation.continuous_base import PropagatorTimer -def construct_one_body_propagator(hamiltonian: HolsteinModel, dt: float): +def construct_one_body_propagator( + hamiltonian: HolsteinModel, dt: float +) -> Sequence[np.ndarray, np.ndarray]: """Exponentiates the electronic hopping term to apply it later as part of the trotterized algorithm. @@ -35,6 +37,11 @@ def construct_one_body_propagator(hamiltonian: HolsteinModel, dt: float): Hamiltonian caryying the one-body term as hamiltonian.T dt : Time step + + Returns + ------- + expH1 : + """ H1 = hamiltonian.T expH1 = numpy.array( From a801b80b63b32e23aaf3de9638eb4bc3d5b91265 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Tue, 12 Mar 2024 21:07:34 -0400 Subject: [PATCH 49/61] removed redundant sqrt(m*w) --- .../variational/coherent_state_variational.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py b/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py index 6e413f23..e879e13a 100644 --- a/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py +++ b/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py @@ -34,7 +34,7 @@ def local_energy( kinetic = np.sum(T[0] * G[0] + T[1] * G[1]) rho = G[0].diagonal() + G[1].diagonal() - e_eph = -g * 2 * np.sqrt(m * w) * np.sum(rho * X) + e_eph = -g * 2 * np.sum(rho * X) e_ph = m * w**2 * np.sum(X * X) local_energy = kinetic + e_eph + e_ph From 42faa914c7da2384b97a5aded5be323f6d1f458d Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Tue, 12 Mar 2024 21:09:11 -0400 Subject: [PATCH 50/61] minor bug fixes --- ipie/addons/eph/estimators/energy.py | 2 +- ipie/addons/eph/propagation/holstein.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/ipie/addons/eph/estimators/energy.py b/ipie/addons/eph/estimators/energy.py index 4d4195cb..613a9838 100644 --- a/ipie/addons/eph/estimators/energy.py +++ b/ipie/addons/eph/estimators/energy.py @@ -29,7 +29,7 @@ def local_energy( hamiltonian: HolsteinModel, walkers: EPhWalkers, trial: EPhTrialWavefunctionBase, -) -> np.ndarray: +) -> xp.ndarray: return local_energy_holstein(system, hamiltonian, walkers, trial) diff --git a/ipie/addons/eph/propagation/holstein.py b/ipie/addons/eph/propagation/holstein.py index cfc2fce9..84411d88 100644 --- a/ipie/addons/eph/propagation/holstein.py +++ b/ipie/addons/eph/propagation/holstein.py @@ -15,6 +15,7 @@ import numpy import time import scipy.linalg +from typing import Sequence from ipie.addons.eph.hamiltonians.holstein import HolsteinModel from ipie.addons.eph.trial_wavefunction.eph_trial_base import EPhTrialWavefunctionBase @@ -27,7 +28,7 @@ def construct_one_body_propagator( hamiltonian: HolsteinModel, dt: float -) -> Sequence[np.ndarray, np.ndarray]: +) -> Sequence[numpy.ndarray]: """Exponentiates the electronic hopping term to apply it later as part of the trotterized algorithm. From 5d7c560ec6334ccfffe3b739bc1a7e44d27b5db8 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Wed, 20 Mar 2024 13:59:29 -0400 Subject: [PATCH 51/61] example for holstein model 4 electrons --- examples/14-1d_holstein/run_afqmc.py | 92 ++++++++++++++++++++++++++++ 1 file changed, 92 insertions(+) create mode 100644 examples/14-1d_holstein/run_afqmc.py diff --git a/examples/14-1d_holstein/run_afqmc.py b/examples/14-1d_holstein/run_afqmc.py new file mode 100644 index 00000000..cca22eaa --- /dev/null +++ b/examples/14-1d_holstein/run_afqmc.py @@ -0,0 +1,92 @@ +# Copyright 2022 The ipie Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np + +np.random.seed(125) +from mpi4py import MPI + +from ipie.qmc.afqmc import AFQMC +from ipie.systems import Generic +from ipie.addons.eph.hamiltonians.holstein import HolsteinModel +from ipie.addons.eph.trial_wavefunction.toyozawa import ToyozawaTrial +from ipie.addons.eph.trial_wavefunction.variational.toyozawa_variational import ( + variational_trial_toyozawa, +) +from ipie.addons.eph.walkers.eph_walkers import EPhWalkers +from ipie.addons.eph.estimators.energy import EnergyEstimator + +# System Parameters +nup = 2 +ndown = 2 +nelec = (nup, ndown) + +# Hamiltonian Parameters +g = 2.0 +t = 1.0 +w0 = 4.0 +nsites = 4 +pbc = True + +# Walker Parameters & Setup +comm = MPI.COMM_WORLD +nwalkers = 1000 // comm.size + +# Setup initial guess for variational optimization +initial_electron = np.random.random((nsites, nup + ndown)) +initial_phonons = np.ones(nsites) * 0.1 + +# System and Hamiltonian setup +system = Generic(nelec) +ham = HolsteinModel(g=g, t=t, w0=w0, nsites=nsites, pbc=pbc) +ham.build() + +# Variational procedure +_, beta_shift, el_trial = variational_trial_toyozawa( + initial_phonons, initial_electron, ham, system +) +wavefunction = np.column_stack([beta_shift, el_trial]) + +# Setup trial +trial = ToyozawaTrial( + wavefunction=wavefunction, w0=ham.w0, num_elec=[nup, ndown], num_basis=nsites +) +trial.set_etrial(ham) + +# Setup walkers +walkers = EPhWalkers( + initial_walker=wavefunction, nup=nup, ndown=ndown, nbasis=nsites, nwalkers=nwalkers +) +walkers.build(trial) + +num_steps_per_block = 10 +num_blocks = 10000 +add_est = { + "energy": EnergyEstimator(system=system, ham=ham, trial=trial), +} + +seed = 125 + +# Note nwalkers specifies the number of walkers on each CPU +ephqmc = AFQMC.build( + num_elec=nelec, + hamiltonian=ham, + trial_wavefunction=trial, + walkers=walkers, + num_walkers=nwalkers, + seed=seed, + num_steps_per_block=num_steps_per_block, + num_blocks=num_blocks, +) +ephqmc.run(additional_estimators=add_est, verbose=False) From 92e3620aad7379330469e0f7172e953191b99087 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Tue, 16 Apr 2024 18:02:58 -0400 Subject: [PATCH 52/61] jax imports in tests guarded --- .../tests/test_variational.py | 24 ++++++++++--------- .../addons/eph/trial_wavefunction/toyozawa.py | 15 +++++++++++- .../variational/toyozawa_variational.py | 16 +------------ 3 files changed, 28 insertions(+), 27 deletions(-) diff --git a/ipie/addons/eph/trial_wavefunction/tests/test_variational.py b/ipie/addons/eph/trial_wavefunction/tests/test_variational.py index 8eb2fc09..ffb0bca9 100644 --- a/ipie/addons/eph/trial_wavefunction/tests/test_variational.py +++ b/ipie/addons/eph/trial_wavefunction/tests/test_variational.py @@ -1,18 +1,20 @@ import pytest import numpy - +import sys from ipie.addons.eph.utils.testing import get_random_sys_holstein from ipie.addons.eph.trial_wavefunction.toyozawa import ToyozawaTrial -from ipie.addons.eph.trial_wavefunction.variational.toyozawa_variational import ( - variational_trial_toyozawa, -) from ipie.addons.eph.trial_wavefunction.coherent_state import CoherentStateTrial -from ipie.addons.eph.trial_wavefunction.variational.coherent_state_variational import ( - variational_trial, -) - - -@pytest.mark.unit +try: + from ipie.addons.eph.trial_wavefunction.variational.toyozawa_variational import ( + variational_trial_toyozawa, + ) + from ipie.addons.eph.trial_wavefunction.variational.coherent_state_variational import ( + variational_trial, + ) +except ImportError: + pass + +@pytest.mark.skip('jax' not in sys.modules, reason="no jax") def test_variational_energy_toyozawa(): seed = 7 numpy.random.seed(seed) @@ -30,7 +32,7 @@ def test_variational_energy_toyozawa(): assert etrial == pytest.approx(trial.energy) -@pytest.mark.unit +@pytest.mark.skip('jax' not in sys.modules, reason="no jax") def test_variational_energy_coherent_state(): seed = 7 numpy.random.seed(seed) diff --git a/ipie/addons/eph/trial_wavefunction/toyozawa.py b/ipie/addons/eph/trial_wavefunction/toyozawa.py index d248a164..d5ac9b98 100644 --- a/ipie/addons/eph/trial_wavefunction/toyozawa.py +++ b/ipie/addons/eph/trial_wavefunction/toyozawa.py @@ -17,10 +17,23 @@ from ipie.addons.eph.walkers.eph_walkers import EPhWalkers from ipie.addons.eph.trial_wavefunction.coherent_state import CoherentStateTrial -from ipie.addons.eph.trial_wavefunction.variational.toyozawa_variational import circ_perm from ipie.utils.backend import arraylib as xp from ipie.estimators.greens_function_single_det import gab_mod_ovlp +def circ_perm(lst: np.ndarray) -> np.ndarray: + """Returns a matrix which rows consist of all possible + cyclic permutations given an initial array lst. + + Parameters + ---------- + lst : + Initial array which is to be cyclically permuted + """ + circs = lst + for shift in range(1, len(lst)): + new_circ = np.roll(lst, -shift) + circs = np.vstack([circs, new_circ]) + return circs class ToyozawaTrial(CoherentStateTrial): r"""The Toyozawa trial diff --git a/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py b/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py index 98d46469..c0a0d7ce 100644 --- a/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py +++ b/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py @@ -15,6 +15,7 @@ import numpy as np from scipy.optimize import minimize from ipie.addons.eph.trial_wavefunction.variational.estimators import gab +from ipie.addons.eph.trial_wavefunction.toyozawa import circ_perm import jax import jax.numpy as npj @@ -94,21 +95,6 @@ def objective_function_toyozawa_mo( return etot.real -def circ_perm(lst: np.ndarray) -> np.ndarray: - """Returns a matrix which rows consist of all possible - cyclic permutations given an initial array lst. - - Parameters - ---------- - lst : - Initial array which is to be cyclically permuted - """ - circs = lst - for shift in range(1, len(lst)): - new_circ = np.roll(lst, -shift) - circs = np.vstack([circs, new_circ]) - return circs - def variational_trial_toyozawa( shift_init: np.ndarray, electron_init: np.ndarray, hamiltonian, system, verbose=2 From 9bcb208f021adab5f398662cf3e3382e4b4e3bbf Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Tue, 16 Apr 2024 18:48:11 -0400 Subject: [PATCH 53/61] removed run_test in holstein example and renamed test_estimators --- examples/14-fqe-wavefunction/run_afqmc.py | 218 ------------------ ...t_estimators.py => test_estimators_eph.py} | 0 2 files changed, 218 deletions(-) delete mode 100644 examples/14-fqe-wavefunction/run_afqmc.py rename ipie/addons/eph/estimators/tests/{test_estimators.py => test_estimators_eph.py} (100%) diff --git a/examples/14-fqe-wavefunction/run_afqmc.py b/examples/14-fqe-wavefunction/run_afqmc.py deleted file mode 100644 index 08cd9bb6..00000000 --- a/examples/14-fqe-wavefunction/run_afqmc.py +++ /dev/null @@ -1,218 +0,0 @@ -# Copyright 2022 The ipie Developers. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# Author: Fionn Malone -# -"""Convert an FQE wavefunction to ipie and vice-versa. - -Play around with various thresholds to see how it affects the energy. -""" -import sys -from typing import List, Tuple, Union - -try: - import fqe -except (ImportError, ModuleNotFoundError): - print("fqe required") - sys.exit(0) -try: - import pyscf -except (ImportError, ModuleNotFoundError): - print("pyscf required") - sys.exit(0) -import numpy as np -from pyscf import ao2mo, gto, mcscf, scf - -from ipie.hamiltonians.generic import GenericRealChol as GenericHam -from ipie.systems.generic import Generic as GenericSys -from ipie.trial_wavefunction.particle_hole import ParticleHole -from ipie.utils.from_pyscf import generate_hamiltonian - - -def get_occa_occb_coeff_from_fqe_wfn( - fqe_wf: fqe.Wavefunction, threshold: float = 0.0 -) -> Tuple[List[np.ndarray], ...]: - """Generate occlists from fqe wavefunction.""" - - def _get_sector_data(sector, threshold, occa_list, occb_list, coeffs): - for inda in range(sector._core.lena()): - alpha_str = sector._core.string_alpha(inda) - for indb in range(sector._core.lenb()): - if np.abs(sector.coeff[inda, indb]) > threshold: - alpha_str = sector._core.string_alpha(inda) - beta_str = sector._core.string_beta(indb) - coeff = sector.coeff[inda, indb] - - occa_list.append(fqe.bitstring.integer_index(alpha_str)) - occb_list.append(fqe.bitstring.integer_index(beta_str)) - coeffs.append(coeff) - - occa_list: List[np.ndarray] = [] - occb_list: List[np.ndarray] = [] - coeffs: List[np.ndarray] = [] - - for sector_key in fqe_wf.sectors(): - sector = fqe_wf.sector(sector_key) - _get_sector_data(sector, threshold, occa_list, occb_list, coeffs) - - return (np.asarray(coeffs), np.asarray(occa_list), np.asarray(occb_list)) - - -def get_fqe_wfn_from_occ_coeff( - coeffs: np.ndarray, - occa: np.ndarray, - occb: np.ndarray, - n_elec: int, - n_orb: int, - ms: int = 0, - threshold: float = 0.0, -) -> fqe.Wavefunction: - """A helper function to map an AFQMC wavefunction to FQE. - - Args: - coeffs: The ci coefficients - occa: The alpha occupation strings. - occb: The beta occupation strings. - n_elec: Number of electrons. - n_orb: number of orbitals. - ms: spin polarization. - threshold: ci coefficient threshold. A coefficient whose absolute value - below this value is considered zero. - """ - - def _set_sector_data(sector, threshold, occa_list, occb_list, coeffs): - fqe_graph = sector.get_fcigraph() - for idet, (occa, occb) in enumerate(zip(occa_list, occb_list)): - alpha_str = fqe.bitstring.reverse_integer_index(occa) - beta_str = fqe.bitstring.reverse_integer_index(occb) - inda = fqe_graph.index_alpha(alpha_str) - indb = fqe_graph.index_alpha(beta_str) - if np.abs(coeffs[idet]) > threshold: - sector.coeff[inda, indb] = coeffs[idet] - - # ensure it is normalized - _coeffs = coeffs / np.dot(coeffs.conj(), coeffs) ** 0.5 - fqe_wf = fqe.Wavefunction([[n_elec, ms, n_orb]]) - - for sector_key in fqe_wf.sectors(): - sector = fqe_wf.sector(sector_key) - _set_sector_data(sector, threshold, occa, occb, _coeffs) - - return fqe_wf - - -def get_fqe_variational_energy( - ecore: float, h1e: np.ndarray, eris: np.ndarray, wfn: fqe.Wavefunction -) -> float: - """Compute FQE variational energy from ERIs and FQE wavefunction.""" - # get integrals into openfermion order - of_eris = np.transpose(eris, (0, 2, 3, 1)) - # ... and then into FQE format - fqe_ham = fqe.restricted_hamiltonian.RestrictedHamiltonian( - (h1e, np.einsum("ijlk", -0.5 * of_eris)), e_0=ecore - ) - return wfn.expectationValue(fqe_ham).real - - -def build_ipie_wavefunction_from_pyscf( - fcivec: np.ndarray, mc: Union[pyscf.mcscf.CASCI, pyscf.mcscf.CASSCF], tol: float = 1e-12 -) -> ParticleHole: - """Build ipie wavefunction in the full space (i.e. with "melting cores")""" - coeff, occa, occb = zip( - *pyscf.fci.addons.large_ci(fcivec, mc.ncas, mc.nelecas, tol=tol, return_strs=False) - ) - ix = np.argsort(np.abs(coeff))[::-1] - nelec = mc._scf.mol.nelec - nmo = mc._scf.mo_coeff.shape[-1] - return ParticleHole((np.array(coeff)[ix], np.array(occa)[ix], np.array(occb)[ix]), nelec, nmo) - - -def strip_melting_cores( - occa: np.ndarray, occb: np.ndarray, n_melting: int -) -> Tuple[np.ndarray, np.ndarray]: - """Strip any melting cores from ipie wavefunction.""" - occa_new = [] - occb_new = [] - for oa, ob in zip(occa, occb): - # ipie typically builds the cas wavefunction in the full space by inserting "melting" cores - # To map back to the active space you need to strip these and then shift - # the orbital indices back by the number of melting cores (n_melting) - occa_new.append(np.array([o - n_melting for o in oa[n_melting:]])) - occb_new.append(np.array([o - n_melting for o in ob[n_melting:]])) - - return np.array(occa_new), np.array(occb_new) - - -def build_ipie_sys_ham_from_pyscf( - mc: Union[pyscf.mcscf.CASCI, pyscf.mcscf.CASSCF], chol_cut: float = 1e-6 -) -> Tuple[GenericSys, GenericHam]: - """Build ipie system and hamiltonian from MCSCF object.""" - ham = generate_hamiltonian( - mc._scf.mol, - mc.mo_coeff, - mc._scf.get_hcore(), - mc.mo_coeff, - chol_cut=chol_cut, - num_frozen_core=0, - verbose=False, - ) - nelec = (mol.nelec[0], mol.nelec[1]) - return GenericSys(nelec), ham - - -if __name__ == "__main__": - mol = gto.Mole(atom=[("N", (0.0, 0.0, 0.0)), ("N", (0.0, 0.0, 1.45))], spin=0, basis="sto-3g") - mol.build() - mf = scf.RHF(mol) - mf.kernel() - - nalpha, nbeta = mol.nelec - nmo = mf.mo_coeff.shape[1] - - ncas, nelecas = (6, 6) - mc = mcscf.CASSCF(mf, nelecas, ncas) - - e_tot, e_cas, fcivec, mo, mo_energy = mc.kernel() - print(f"DIM(H) = {fcivec.ravel().shape}") - # Get the active space ERIs in the CASSCF MO basis - h1e, e0 = mc.get_h1eff(mc.mo_coeff) - eris = ao2mo.restore("1", mc.get_h2eff(mc.mo_coeff), ncas).reshape((ncas,) * 4) - # you can check how truncating the wavefunction affects the energy - wfn = build_ipie_wavefunction_from_pyscf(fcivec, mc, tol=0.0) - print(f"Length of truncated CI expansion: {wfn.num_dets}") - # you check how truncating the Cholesky dimension affects the energy - sys, ham = build_ipie_sys_ham_from_pyscf(mc, chol_cut=1e-12) - - ipie_energy = wfn.calculate_energy(sys, ham)[0] - msg = f"{ipie_energy.real:.10f}" - print(f"ipie energy: {msg}") - assert np.isclose(e_tot, ipie_energy, atol=1e-8), f"{e_tot} != {msg}" - - # Convert to FQE and check the energy - occa_fqe, occb_fqe = strip_melting_cores(wfn.occa, wfn.occb, wfn.nmelting) - fqe_wfn = get_fqe_wfn_from_occ_coeff( - wfn.coeffs, occa_fqe, occb_fqe, nelecas, ncas, ms=0, threshold=0.0 - ) - fqe_energy = get_fqe_variational_energy(e0, h1e, eris, fqe_wfn) - msg = f"{fqe_energy.real:.10f}" - print(f"FQE energy: {msg}") - assert np.isclose(e_tot, fqe_energy, atol=1e-8), f"{e_tot} != {msg}" - - # round trip back to ipie - coeff, occa, occb = get_occa_occb_coeff_from_fqe_wfn(fqe_wfn, threshold=0.0) - wfn_round_trip = ParticleHole((coeff, occa, occb), mc._scf.mol.nelec, mc.mo_coeff.shape[-1]) - ipie_energy = wfn_round_trip.calculate_energy(sys, ham)[0] - msg = f"{ipie_energy.real:.10f}" - print(f"ipie energy from round trip: {msg}") - assert np.isclose(e_tot, ipie_energy, atol=1e-8), f"{e_tot} != {msg}" diff --git a/ipie/addons/eph/estimators/tests/test_estimators.py b/ipie/addons/eph/estimators/tests/test_estimators_eph.py similarity index 100% rename from ipie/addons/eph/estimators/tests/test_estimators.py rename to ipie/addons/eph/estimators/tests/test_estimators_eph.py From 705d1e6dd4e2f84f23613461eed8f987a6db680c Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Sun, 4 Aug 2024 18:23:32 -0400 Subject: [PATCH 54/61] removed jax dependency in example --- examples/14-1d_holstein/run_afqmc.py | 17 +++++++++-------- examples/14-1d_holstein/wavefunction.npy | Bin 0 -> 288 bytes 2 files changed, 9 insertions(+), 8 deletions(-) create mode 100644 examples/14-1d_holstein/wavefunction.npy diff --git a/examples/14-1d_holstein/run_afqmc.py b/examples/14-1d_holstein/run_afqmc.py index cca22eaa..c4aff76c 100644 --- a/examples/14-1d_holstein/run_afqmc.py +++ b/examples/14-1d_holstein/run_afqmc.py @@ -21,9 +21,9 @@ from ipie.systems import Generic from ipie.addons.eph.hamiltonians.holstein import HolsteinModel from ipie.addons.eph.trial_wavefunction.toyozawa import ToyozawaTrial -from ipie.addons.eph.trial_wavefunction.variational.toyozawa_variational import ( - variational_trial_toyozawa, -) +# from ipie.addons.eph.trial_wavefunction.variational.toyozawa_variational import ( +# variational_trial_toyozawa, +# ) from ipie.addons.eph.walkers.eph_walkers import EPhWalkers from ipie.addons.eph.estimators.energy import EnergyEstimator @@ -52,11 +52,12 @@ ham = HolsteinModel(g=g, t=t, w0=w0, nsites=nsites, pbc=pbc) ham.build() -# Variational procedure -_, beta_shift, el_trial = variational_trial_toyozawa( - initial_phonons, initial_electron, ham, system -) -wavefunction = np.column_stack([beta_shift, el_trial]) +# Variational procedure - If Jax provided +# _, beta_shift, el_trial = variational_trial_toyozawa( +# initial_phonons, initial_electron, ham, system +# ) +# wavefunction = np.column_stack([beta_shift, el_trial]) +wavefunction = np.load('wavefunction.npy') # Setup trial trial = ToyozawaTrial( diff --git a/examples/14-1d_holstein/wavefunction.npy b/examples/14-1d_holstein/wavefunction.npy new file mode 100644 index 0000000000000000000000000000000000000000..dab7802f47c4e67d7e537a325e725fda19993311 GIT binary patch literal 288 zcmbR27wQ`j$;eQ~P_3SlTAW;@Zl$1ZlV+i=qoAIaUsO_*m=~X4l#&V(4=E~51qv5u zBo?Fsxf&)q3Z|Mm3bhL41Fm1wSLdIvT(v*CIDS*V`A7T9d*=z8@L#jv|G=+yZRB_R zC5&4&vfdxJ|2a{ob%Wt0`;RZ*#7ME;w2!fG*%~1A$$r^cqiJHNPTSwI-PXaP@xi`R z>gyKEO*ici>1(ApzB#l%^>2B!!r?>qa{GD}wnvs%?bm%#+4?u{r~QlK kNfy`SF59=*PQ0NI^WFaIVW;i&jSub1 Date: Sun, 4 Aug 2024 18:36:08 -0400 Subject: [PATCH 55/61] accommodates new Propagator dict --- examples/14-1d_holstein/run_afqmc.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/examples/14-1d_holstein/run_afqmc.py b/examples/14-1d_holstein/run_afqmc.py index c4aff76c..e99cf671 100644 --- a/examples/14-1d_holstein/run_afqmc.py +++ b/examples/14-1d_holstein/run_afqmc.py @@ -26,6 +26,9 @@ # ) from ipie.addons.eph.walkers.eph_walkers import EPhWalkers from ipie.addons.eph.estimators.energy import EnergyEstimator +from ipie.propagation.propagator import Propagator +from ipie.addons.eph.propagation.holstein import HolsteinPropagator +Propagator.update({HolsteinModel: HolsteinPropagator}) # System Parameters nup = 2 From e02282ca6dde6314465207a1ada8d7179d96eed9 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Mon, 5 Aug 2024 11:17:23 -0400 Subject: [PATCH 56/61] jax workaround like torch --- dev/jax.txt | 1 + setup.py | 1 + 2 files changed, 2 insertions(+) create mode 100644 dev/jax.txt diff --git a/dev/jax.txt b/dev/jax.txt new file mode 100644 index 00000000..1ae27bc0 --- /dev/null +++ b/dev/jax.txt @@ -0,0 +1 @@ +jax >= 0.4.13 diff --git a/setup.py b/setup.py index 06e25a3b..f74b0f2c 100644 --- a/setup.py +++ b/setup.py @@ -80,6 +80,7 @@ def main() -> None: "mpi": load_requirements("dev/mpi.txt"), "dev": load_requirements("dev/dev.txt"), "torch": load_requirements("dev/torch.txt"), + "jax": load_requirements("dev/jax.txt"), }, long_description=open("README.rst").read(), ) From 44505b5799a5f09c15177eda6b1dd89839f4e2e5 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Fri, 6 Sep 2024 18:02:15 -0400 Subject: [PATCH 57/61] fixed pytests --- examples/14-1d_holstein/run_afqmc.py | 2 +- ipie/estimators/energy.py | 2 +- ipie/estimators/handler.py | 2 +- ipie/propagation/propagator.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/examples/14-1d_holstein/run_afqmc.py b/examples/14-1d_holstein/run_afqmc.py index e99cf671..5b6edab1 100644 --- a/examples/14-1d_holstein/run_afqmc.py +++ b/examples/14-1d_holstein/run_afqmc.py @@ -75,7 +75,7 @@ walkers.build(trial) num_steps_per_block = 10 -num_blocks = 10000 +num_blocks = 10 add_est = { "energy": EnergyEstimator(system=system, ham=ham, trial=trial), } diff --git a/ipie/estimators/energy.py b/ipie/estimators/energy.py index 33292937..08364517 100644 --- a/ipie/estimators/energy.py +++ b/ipie/estimators/energy.py @@ -150,7 +150,7 @@ def compute_estimator(self, system=None, walkers=None, hamiltonian=None, trial=N self._data["ENumer"] = xp.sum(walkers.weight * energy[:, 0].real) self._data["EDenom"] = xp.sum(walkers.weight) self._data["E1Body"] = xp.sum(walkers.weight * energy[:, 1].real) - self._data["E2Body"] = xp.sum(walkers.weight * energy[:, 2].real) + self._data["E2Body"] = xp.sum(walkers.weight * energy[:, 2].real) return self.data def get_index(self, name): diff --git a/ipie/estimators/handler.py b/ipie/estimators/handler.py index a068e9bd..cbe01fe9 100644 --- a/ipie/estimators/handler.py +++ b/ipie/estimators/handler.py @@ -235,7 +235,7 @@ def print_block(self, comm, block, walker_factors, div_factor=None): shift = None walker_factors.eshift = comm.bcast(shift) if comm.rank == 0: -# self.output.push_to_chunk(self.global_estimates, f"data") + self.output.push_to_chunk(self.global_estimates, f"data") self.output.increment() if comm.rank == 0: print(f"{block:>17d} " + output_string) diff --git a/ipie/propagation/propagator.py b/ipie/propagation/propagator.py index db2ef0ed..cc92576b 100644 --- a/ipie/propagation/propagator.py +++ b/ipie/propagation/propagator.py @@ -7,4 +7,4 @@ GenericRealChol: PhaselessGeneric, GenericComplexChol: PhaselessGeneric, GenericRealCholChunked: PhaselessGenericChunked, -} \ No newline at end of file +} From fab267b4341a42075570f18c7749326c07c7dd51 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Fri, 6 Sep 2024 18:06:20 -0400 Subject: [PATCH 58/61] fixed pylint errors --- .../trial_wavefunction/variational/coherent_state_variational.py | 1 + ipie/addons/eph/trial_wavefunction/variational/estimators.py | 1 + .../eph/trial_wavefunction/variational/toyozawa_variational.py | 1 + 3 files changed, 3 insertions(+) diff --git a/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py b/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py index e879e13a..1da0c8c6 100644 --- a/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py +++ b/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py @@ -17,6 +17,7 @@ from ipie.addons.eph.trial_wavefunction.variational.estimators import gab +# pylint: disable=import-error import jax diff --git a/ipie/addons/eph/trial_wavefunction/variational/estimators.py b/ipie/addons/eph/trial_wavefunction/variational/estimators.py index f97f7c39..18202a14 100644 --- a/ipie/addons/eph/trial_wavefunction/variational/estimators.py +++ b/ipie/addons/eph/trial_wavefunction/variational/estimators.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +# pylint: disable=import-error import jax.numpy as npj from jax.config import config diff --git a/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py b/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py index c0a0d7ce..5bc6c999 100644 --- a/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py +++ b/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py @@ -17,6 +17,7 @@ from ipie.addons.eph.trial_wavefunction.variational.estimators import gab from ipie.addons.eph.trial_wavefunction.toyozawa import circ_perm +# pylint: disable=import-error import jax import jax.numpy as npj From 03f7b054d1fc5d3687304c8eb8daa6b885414145 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Fri, 6 Sep 2024 19:16:20 -0400 Subject: [PATCH 59/61] fixed remaining health errors --- examples/14-1d_holstein/run_afqmc.py | 2 +- ipie/addons/eph/trial_wavefunction/coherent_state.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/14-1d_holstein/run_afqmc.py b/examples/14-1d_holstein/run_afqmc.py index 5b6edab1..b85cefee 100644 --- a/examples/14-1d_holstein/run_afqmc.py +++ b/examples/14-1d_holstein/run_afqmc.py @@ -15,7 +15,7 @@ import numpy as np np.random.seed(125) -from mpi4py import MPI +from ipie.config import MPI from ipie.qmc.afqmc import AFQMC from ipie.systems import Generic diff --git a/ipie/addons/eph/trial_wavefunction/coherent_state.py b/ipie/addons/eph/trial_wavefunction/coherent_state.py index 6cf963fc..5781b57f 100644 --- a/ipie/addons/eph/trial_wavefunction/coherent_state.py +++ b/ipie/addons/eph/trial_wavefunction/coherent_state.py @@ -147,7 +147,7 @@ def calc_phonon_gradient(self, walkers) -> np.ndarray: grad *= -self.m * self.w0 return grad - def calc_phonon_laplacian(self, walkers) -> np.ndarray: + def calc_phonon_laplacian(self, walkers, ovlps=None) -> np.ndarray: r"""Computes the Laplacian of phonon overlaps, .. math:: From b73fdd73d40fec9d6b5befb67f50f2725191e3b3 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Sat, 7 Sep 2024 12:55:32 -0400 Subject: [PATCH 60/61] fixed examples and black --- ipie/addons/eph/propagation/holstein.py | 6 ++---- .../eph/trial_wavefunction/tests/test_variational.py | 10 ++++++---- ipie/addons/eph/trial_wavefunction/toyozawa.py | 2 ++ .../variational/toyozawa_variational.py | 1 - 4 files changed, 10 insertions(+), 9 deletions(-) diff --git a/ipie/addons/eph/propagation/holstein.py b/ipie/addons/eph/propagation/holstein.py index 84411d88..7bd01c9b 100644 --- a/ipie/addons/eph/propagation/holstein.py +++ b/ipie/addons/eph/propagation/holstein.py @@ -26,9 +26,7 @@ from ipie.propagation.continuous_base import PropagatorTimer -def construct_one_body_propagator( - hamiltonian: HolsteinModel, dt: float -) -> Sequence[numpy.ndarray]: +def construct_one_body_propagator(hamiltonian: HolsteinModel, dt: float) -> Sequence[numpy.ndarray]: """Exponentiates the electronic hopping term to apply it later as part of the trotterized algorithm. @@ -42,7 +40,7 @@ def construct_one_body_propagator( Returns ------- expH1 : - + """ H1 = hamiltonian.T expH1 = numpy.array( diff --git a/ipie/addons/eph/trial_wavefunction/tests/test_variational.py b/ipie/addons/eph/trial_wavefunction/tests/test_variational.py index ffb0bca9..5316d67b 100644 --- a/ipie/addons/eph/trial_wavefunction/tests/test_variational.py +++ b/ipie/addons/eph/trial_wavefunction/tests/test_variational.py @@ -4,17 +4,19 @@ from ipie.addons.eph.utils.testing import get_random_sys_holstein from ipie.addons.eph.trial_wavefunction.toyozawa import ToyozawaTrial from ipie.addons.eph.trial_wavefunction.coherent_state import CoherentStateTrial + try: from ipie.addons.eph.trial_wavefunction.variational.toyozawa_variational import ( variational_trial_toyozawa, - ) + ) from ipie.addons.eph.trial_wavefunction.variational.coherent_state_variational import ( variational_trial, - ) + ) except ImportError: pass -@pytest.mark.skip('jax' not in sys.modules, reason="no jax") + +@pytest.mark.skip("jax" not in sys.modules, reason="no jax") def test_variational_energy_toyozawa(): seed = 7 numpy.random.seed(seed) @@ -32,7 +34,7 @@ def test_variational_energy_toyozawa(): assert etrial == pytest.approx(trial.energy) -@pytest.mark.skip('jax' not in sys.modules, reason="no jax") +@pytest.mark.skip("jax" not in sys.modules, reason="no jax") def test_variational_energy_coherent_state(): seed = 7 numpy.random.seed(seed) diff --git a/ipie/addons/eph/trial_wavefunction/toyozawa.py b/ipie/addons/eph/trial_wavefunction/toyozawa.py index d5ac9b98..aefe652a 100644 --- a/ipie/addons/eph/trial_wavefunction/toyozawa.py +++ b/ipie/addons/eph/trial_wavefunction/toyozawa.py @@ -20,6 +20,7 @@ from ipie.utils.backend import arraylib as xp from ipie.estimators.greens_function_single_det import gab_mod_ovlp + def circ_perm(lst: np.ndarray) -> np.ndarray: """Returns a matrix which rows consist of all possible cyclic permutations given an initial array lst. @@ -35,6 +36,7 @@ def circ_perm(lst: np.ndarray) -> np.ndarray: circs = np.vstack([circs, new_circ]) return circs + class ToyozawaTrial(CoherentStateTrial): r"""The Toyozawa trial diff --git a/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py b/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py index 5bc6c999..8ee31b10 100644 --- a/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py +++ b/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py @@ -96,7 +96,6 @@ def objective_function_toyozawa_mo( return etot.real - def variational_trial_toyozawa( shift_init: np.ndarray, electron_init: np.ndarray, hamiltonian, system, verbose=2 ): From fb6958448eaccfffddff4800c15b407f99a093d9 Mon Sep 17 00:00:00 2001 From: Moritz Baumgarten Date: Sat, 7 Sep 2024 13:51:23 -0400 Subject: [PATCH 61/61] fixed example --- examples/14-1d_holstein/wavefunction.npy | Bin 288 -> 0 bytes .../run_afqmc.py | 15 +++++++++++---- 2 files changed, 11 insertions(+), 4 deletions(-) delete mode 100644 examples/14-1d_holstein/wavefunction.npy rename examples/{14-1d_holstein => 19-1d_holstein}/run_afqmc.py (86%) diff --git a/examples/14-1d_holstein/wavefunction.npy b/examples/14-1d_holstein/wavefunction.npy deleted file mode 100644 index dab7802f47c4e67d7e537a325e725fda19993311..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 288 zcmbR27wQ`j$;eQ~P_3SlTAW;@Zl$1ZlV+i=qoAIaUsO_*m=~X4l#&V(4=E~51qv5u zBo?Fsxf&)q3Z|Mm3bhL41Fm1wSLdIvT(v*CIDS*V`A7T9d*=z8@L#jv|G=+yZRB_R zC5&4&vfdxJ|2a{ob%Wt0`;RZ*#7ME;w2!fG*%~1A$$r^cqiJHNPTSwI-PXaP@xi`R z>gyKEO*ici>1(ApzB#l%^>2B!!r?>qa{GD}wnvs%?bm%#+4?u{r~QlK kNfy`SF59=*PQ0NI^WFaIVW;i&jSub1