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/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/examples/19-1d_holstein/run_afqmc.py b/examples/19-1d_holstein/run_afqmc.py new file mode 100644 index 00000000..7223ad66 --- /dev/null +++ b/examples/19-1d_holstein/run_afqmc.py @@ -0,0 +1,103 @@ +# 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 ipie.config 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 +from ipie.propagation.propagator import Propagator +from ipie.addons.eph.propagation.holstein import HolsteinPropagator + +Propagator.update({HolsteinModel: HolsteinPropagator}) + +# 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 - 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.array( + [ + [-0.05171059, 0.18699386, 0.21710474, 0.14657979, 0.29875713], + [1.07606477, 0.28428281, 1.00976995, -0.10386843, 1.45999226], + [0.34467794, 0.39093646, 0.40150302, 0.33165806, 0.54705797], + [1.45939397, 1.13142463, -0.14783494, 1.58959742, -0.35513289], + ] +) + +# 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 = 10 +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/__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 new file mode 100644 index 00000000..e2aed039 --- /dev/null +++ b/ipie/addons/eph/__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/estimators/__init__.py b/ipie/addons/eph/estimators/__init__.py new file mode 100644 index 00000000..e2aed039 --- /dev/null +++ b/ipie/addons/eph/estimators/__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/estimators/energy.py b/ipie/addons/eph/estimators/energy.py new file mode 100644 index 00000000..613a9838 --- /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, +) -> xp.ndarray: + 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): + # 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 new file mode 100644 index 00000000..aab5589b --- /dev/null +++ b/ipie/addons/eph/estimators/local_energy_holstein.py @@ -0,0 +1,74 @@ +# 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 +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. + + Parameters + ---------- + system : :class:`Generic` + Generic object carrying information on number and spin of electrons + hamiltonian : :class:`HolsteinModel` + HolsteinModel object + walkers : :class:`EPhWalkers` + EPhWalkers object + trial : :class:`EPhTrialWavefunctionBase` + EPhTrialWavefunctionBase object + """ + + energy = xp.zeros((walkers.nwalkers, 4), dtype=xp.complex128) + + gf = trial.calc_greens_function(walkers) + 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[:, 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] *= 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 diff --git a/ipie/addons/eph/estimators/tests/test_estimators_eph.py b/ipie/addons/eph/estimators/tests/test_estimators_eph.py new file mode 100644 index 00000000..0eef56a0 --- /dev/null +++ b/ipie/addons/eph/estimators/tests/test_estimators_eph.py @@ -0,0 +1,31 @@ +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() diff --git a/ipie/addons/eph/hamiltonians/__init__.py b/ipie/addons/eph/hamiltonians/__init__.py new file mode 100644 index 00000000..e2aed039 --- /dev/null +++ b/ipie/addons/eph/hamiltonians/__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/hamiltonians/holstein.py b/ipie/addons/eph/hamiltonians/holstein.py new file mode 100644 index 00000000..a951fbd0 --- /dev/null +++ b/ipie/addons/eph/hamiltonians/holstein.py @@ -0,0 +1,70 @@ +# 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 + + +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 : :class:`float` + Electron-phonon coupling strength + t : :class:`float` + Electron hopping parameter + w0 : :class:`float` + Phonon frequency + nsites : :class:`int` + Length of the 1D Holstein chain + pbc : :class:``bool` + 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.0 * self.m * self.w0) + + 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) + + if self.pbc: + self.T[0, -1] = self.T[-1, 0] = 1.0 + + self.T *= -self.t + + self.T = [self.T.copy(), self.T.copy()] + + self.g_tensor = -self.g * numpy.eye(self.nsites) diff --git a/ipie/addons/eph/propagation/__init__.py b/ipie/addons/eph/propagation/__init__.py new file mode 100644 index 00000000..e2aed039 --- /dev/null +++ b/ipie/addons/eph/propagation/__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/propagation/holstein.py b/ipie/addons/eph/propagation/holstein.py new file mode 100644 index 00000000..7bd01c9b --- /dev/null +++ b/ipie/addons/eph/propagation/holstein.py @@ -0,0 +1,300 @@ +# 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 +from typing import Sequence + +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.utils.backend import synchronize +from ipie.propagation.operations import propagate_one_body +from ipie.propagation.continuous_base import PropagatorTimer + + +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. + + Parameters + ---------- + hamiltonian : + Hamiltonian caryying the one-body term as hamiltonian.T + dt : + Time step + + Returns + ------- + expH1 : + + """ + 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: + 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 + 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: HolsteinModel, + trial: EPhTrialWavefunctionBase = None, + walkers: EPhWalkers = None, + mpi_handler=None, + ) -> None: + """Necessary step before running the AFQMC procedure. + Sets required attributes. + + Parameters + ---------- + hamiltonian : + Holstein model + trial : + Trial object + 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.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: + 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.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.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: + 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() + synchronize() + 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 = 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: 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) + walkers.ovlp = ovlp + synchronize() + self.timer.tgf += time.time() - start_time + + # Update Walkers + # a) DMC for phonon degrees of freedom + self.propagate_phonons(walkers, hamiltonian, trial) + + # b) One-body propagation for electrons + self.propagate_electron(walkers, hamiltonian, trial) + + # 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) + walkers.ovlp = ovlp_new + 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) -> None: + walkers.weight *= ovlp_new / ovlp + + +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: 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 + 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)) + 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) + + 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) + + 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 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..a811c076 --- /dev/null +++ b/ipie/addons/eph/propagation/tests/test_holstein.py @@ -0,0 +1,2 @@ +import numpy +import pytest diff --git a/ipie/addons/eph/trial_wavefunction/__init__.py b/ipie/addons/eph/trial_wavefunction/__init__.py new file mode 100644 index 00000000..e2aed039 --- /dev/null +++ b/ipie/addons/eph/trial_wavefunction/__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/coherent_state.py b/ipie/addons/eph/trial_wavefunction/coherent_state.py new file mode 100644 index 00000000..5781b57f --- /dev/null +++ b/ipie/addons/eph/trial_wavefunction/coherent_state.py @@ -0,0 +1,233 @@ +# 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.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 + + .. 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. + w0 : + Phonon frequency + 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, 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 = w0 + self.m = 1 / w0 + self.nsites = self.nbasis + + 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 calc_energy(self, ham) -> float: + 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) + 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, + + .. 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 + """ + 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`. + + 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, ovlps=None) -> 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 + 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: + 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 + 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) + + 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: + """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) + ) + 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) + ) + 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 new file mode 100644 index 00000000..c076876c --- /dev/null +++ b/ipie/addons/eph/trial_wavefunction/eph_trial_base.py @@ -0,0 +1,84 @@ +# 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 typing import Tuple +from abc import ABCMeta, abstractmethod + + +class EPhTrialWavefunctionBase(metaclass=ABCMeta): + """Base class for electron-phonon trial wave functions. + + Parameters + ---------- + wavefunction : :class:`np.ndarray` + Concatenation of trial determinants of up and down spin spaces and beta + specifying the coherent state displacement. + num_elec : :class:`Tuple` + Tuple of numbers of up and down spins. + num_basis : :class:`int` + Number of sites of Holstein chain. + verbose : :class:`bool` + 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.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.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 + + def set_etrial(self, ham) -> None: + energy = self.calc_energy(ham) + self.energy = energy + + @abstractmethod + def calc_energy(self, ham) -> float: ... + + @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: ... 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 new file mode 100644 index 00000000..5316d67b --- /dev/null +++ b/ipie/addons/eph/trial_wavefunction/tests/test_variational.py @@ -0,0 +1,57 @@ +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.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") +def test_variational_energy_toyozawa(): + seed = 7 + numpy.random.seed(seed) + nelec = (2, 2) + nbasis = 4 + 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.skip("jax" not in sys.modules, reason="no jax") +def test_variational_energy_coherent_state(): + seed = 7 + numpy.random.seed(seed) + nelec = (2, 2) + nbasis = 4 + 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() diff --git a/ipie/addons/eph/trial_wavefunction/toyozawa.py b/ipie/addons/eph/trial_wavefunction/toyozawa.py new file mode 100644 index 00000000..aefe652a --- /dev/null +++ b/ipie/addons/eph/trial_wavefunction/toyozawa.py @@ -0,0 +1,416 @@ +# 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 typing import Tuple + +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.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 + + .. 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, + w0: float, + num_elec: Tuple[int, int], + num_basis: int, + verbose=False, + ): + super().__init__(wavefunction, w0, num_elec, num_basis, verbose=verbose) + self.perms = circ_perm(np.arange(self.nbasis)) + self.nperms = self.perms.shape[0] + + def calc_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 perm in 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`, + + .. 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 + ------- + 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.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 + 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 + + Returns + ------- + ovlp: :class:`np.ndarray` + Sum of product of electron and phonon overlap + """ + 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: + 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) + 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"""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) + grad *= -self.m * self.w0 + 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:: + \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 + arg2 = arg**2 + laplacian += (np.sum(arg2, axis=1) - self.nsites * self.m * self.w0) * ovlp + laplacian /= np.sum(ovlps, axis=1) + 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.ovlp_perm) + + 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) + + 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 + 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. + + 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: + 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.ovlp_perm.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) + if self.ndown > 0: + Gb = np.einsum("nij,n->nij", Gb, 1 / walkers.ovlp) + + return [Ga, Gb] 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/coherent_state_variational.py b/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py new file mode 100644 index 00000000..1da0c8c6 --- /dev/null +++ b/ipie/addons/eph/trial_wavefunction/variational/coherent_state_variational.py @@ -0,0 +1,141 @@ +# 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 + +from ipie.addons.eph.trial_wavefunction.variational.estimators import gab + +# pylint: disable=import-error +import jax + + +def local_energy( + X: np.ndarray, + G: np.ndarray, + m: float, + w: float, + g: float, + nsites: int, + T: np.ndarray, + nup: int, + ndown: int, +) -> np.ndarray: + + kinetic = np.sum(T[0] * G[0] + T[1] * G[1]) + rho = G[0].diagonal() + G[1].diagonal() + e_eph = -g * 2 * np.sum(rho * X) + e_ph = m * w**2 * np.sum(X * X) + + local_energy = kinetic + e_eph + e_ph + return local_energy + + +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] + + etot = local_energy(shift, G, m, w0, g, nbasis, 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]) + + maxiter = 500 + minimizer_kwargs = { + "jac": True, + "args": ( + hamiltonian.nsites, + hamiltonian.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 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..18202a14 --- /dev/null +++ b/ipie/addons/eph/trial_wavefunction/variational/estimators.py @@ -0,0 +1,25 @@ +# 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. + +# pylint: disable=import-error +import jax.numpy as npj +from jax.config import config + +config.update("jax_enable_x64", True) + + +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 new file mode 100644 index 00000000..8ee31b10 --- /dev/null +++ b/ipie/addons/eph/trial_wavefunction/variational/toyozawa_variational.py @@ -0,0 +1,154 @@ +# 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 +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 + + +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 + + 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 + + num_energy = 0.0 + denom = 0.0 + + beta0 = shift0 * npj.sqrt(m * w0 / 2.0) + + for perm in 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 variational_trial_toyozawa( + shift_init: np.ndarray, electron_init: np.ndarray, hamiltonian, system, verbose=2 +): + 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() + + res = minimize( + 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": verbose, + "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((system.nup, hamiltonian.nsites)).T + psib = res.x[hamiltonian.nsites * (system.nup + 1) :] + psib = psib.reshape((system.ndown, hamiltonian.nsites)).T + psi = np.column_stack([psia, psib]) + else: + psia = res.x[hamiltonian.nsites :].reshape((system.nup, hamiltonian.nsites)).T + psi = psia + + return etrial, beta_shift, psi 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..25d42deb --- /dev/null +++ b/ipie/addons/eph/utils/testing.py @@ -0,0 +1,56 @@ +import numpy + +from ipie.systems import Generic + +from ipie.addons.eph.hamiltonians.holstein import HolsteinModel +from ipie.addons.eph.walkers.eph_walkers import EPhWalkers +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 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. diff --git a/ipie/addons/eph/walkers/eph_walkers.py b/ipie/addons/eph/walkers/eph_walkers.py new file mode 100644 index 00000000..4177a372 --- /dev/null +++ b/ipie/addons/eph/walkers/eph_walkers.py @@ -0,0 +1,164 @@ +# 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 + + +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.phonon_disp. + + 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, + nup: int, + ndown: int, + nbasis: int, + nwalkers: int, + 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 + ) + 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)], + dtype=xp.complex128, + ) + + self.phib = xp.array( + [ + 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 + simulation. + + Parameters + ---------- + trial : + Trial wavefunction object. + """ + + if hasattr(trial, "nperms"): + 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.ovlp_perm = numpy.zeros(shape, dtype=numpy.complex128) + + 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) + + 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) + + def reortho(self): + """reorthogonalise walkers. This function is mostly from BaseWalkers, + with the exception that it adjusts all overlaps, possibly of numerous + coherent states. + + 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.el_ovlp[iw, ...] = self.el_ovlp[iw, ...] / detR[iw] + self.ovlp_perm[iw, ...] = self.ovlp_perm[iw, ...] / detR[iw] + self.ovlp[iw] = self.ovlp[iw] / detR[iw] + + synchronize() + return detR + + def reortho_batched(self): # gpu version + pass 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() 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/estimators/energy.py b/ipie/estimators/energy.py index b308a916..ee4f5861 100644 --- a/ipie/estimators/energy.py +++ b/ipie/estimators/energy.py @@ -161,7 +161,6 @@ def compute_estimator(self, system=None, walkers=None, hamiltonian=None, trial=N 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_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 ------- 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(), )