From f0a10a68e118ec814685b91a0482502d02230c54 Mon Sep 17 00:00:00 2001 From: m-baumgarten Date: Tue, 5 Mar 2024 00:08:17 -0500 Subject: [PATCH] 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}