Electron-phonon extension

examples/13-1d_holstein/run_afqmc.py

@@ -0,0 +1,87 @@
+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.propagation.holstein import HolsteinPropagatorImportance
+from ipie.addons.eph.estimators.energy import EnergyEstimator
+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)
+
+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)
+

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.

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.

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]

ipie/addons/eph/estimators/local_energy_holstein.py

@@ -0,0 +1,51 @@
+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
+) -> 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} 
+    |\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, 4), 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[:, 3] -= 0.5 * trial.calc_phonon_laplacian_locenergy(walkers) / hamiltonian.m
+
+    energy[:, 0] = np.sum(energy[:,1:], axis=1)
+
+    return energy
+
+
+

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.

ipie/addons/eph/hamiltonians/holstein.py

@@ -0,0 +1,57 @@
+import numpy
+from ipie.utils.backend import arraylib as xp
+
+
+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):
+        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):
+        """Constructs electronic hopping matrix."""
+        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/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.