moved files to adhere to addons infrastructure

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

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
+
+
+

ipie/propagation/holstein.py → 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

ipie/trial_wavefunction/holstein/coherent_state.py → 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 

ipie/trial_wavefunction/holstein/eph_trial_base.py → 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

ipie/trial_wavefunction/holstein/toyozawa.py → 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,35 +74,24 @@ 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) 
 
             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))
 

ipie/trial_wavefunction/holstein/variational/coherent_state_variational.py → 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
 

ipie/trial_wavefunction/holstein/variational/toyozawa_variational.py → 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)

ipie/walkers/eph_walkers.py → 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

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(

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}