Electron-phonon extension

examples/13-1d_holstein/run_afqmc.py

@@ -0,0 +1,94 @@
+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.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
+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)
+

ipie/addons/ephqmc/estimators/local_energy_holstein.py

@@ -0,0 +1,56 @@
+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.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/addons/ephqmc/hamiltonians/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
+

ipie/addons/ephqmc/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 Holstein 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()]
+
+

ipie/addons/ephqmc/propagation/holstein.py

@@ -0,0 +1,182 @@
+import numpy 
+import time
+import scipy.linalg
+
+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
+
+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():
+    """"""
+    def __init__(self, time_step, verbose=False):
+        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.sum(walkers.x**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 
+
+        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)
+
+        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
+
+        expEph = numpy.exp(self.const * walkers.x)
+
+        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, 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):
+        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):
+        """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.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)
+
+        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.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)
+
+        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, 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
+