diff --git a/examples/16-ft_afqmc/run_afqmc.py b/examples/16-ft_afqmc/run_afqmc.py new file mode 100644 index 00000000..272da7e4 --- /dev/null +++ b/examples/16-ft_afqmc/run_afqmc.py @@ -0,0 +1,87 @@ +import json +import numpy + +from ueg import UEG +from ipie.config import MPI +from ipie.addons.thermal.qmc.calc import build_thermal_afqmc_driver +from ipie.analysis.extraction import extract_observable +from ipie.analysis.autocorr import reblock_by_autocorr + +comm = MPI.COMM_WORLD + +verbose = False if (comm.rank != 0) else True + +# 1. Generate UEG integrals. +ueg_opts = { + "nup": 1, + "ndown": 1, + "rs": 3, + "ecut": 0.5, + "thermal": True, + "write_integrals": True + } + +ueg = UEG(ueg_opts, verbose=verbose) + +if comm.rank == 0: + ueg.build(verbose=verbose) + +comm.barrier() + +# 2. Build thermal AFQMC driver. +options = { + 'trial': { + 'name': 'one_body', + }, + + 'walkers': { + 'lowrank': False, + }, + + 'qmc': { + 'mu': 0.133579, + 'beta': 10., + 'timestep': 0.5, + 'nwalkers': 12 // comm.size, + 'stack_size': 10, + 'seed': 7, + 'nblocks': 20, + }, + } + +afqmc = build_thermal_afqmc_driver( + comm, + nelec=ueg.nelec, + hamiltonian_file='ueg_integrals.h5', + seed=7, + options=options, + verbosity=verbose + ) + +if verbose: + print(f'\nThermal AFQMC options: \n{json.dumps(options, indent=4)}\n') + print(afqmc.params) # Inspect the qmc options. + +# 3. Run thermal AFQMC calculation. +afqmc.run(verbose=verbose) +afqmc.finalise() +afqmc.estimators.compute_estimators(afqmc.hamiltonian, afqmc.trial, afqmc.walkers) + +if comm.rank == 0: + energy_data = extract_observable(afqmc.estimators.filename, "energy") + number_data = extract_observable(afqmc.estimators.filename, "nav") + + print(f'filename: {afqmc.estimators.filename}') + print(f'\nenergy_data: \n{energy_data}\n') + print(f'number_data: \n{number_data}\n') + + y = energy_data["ETotal"] + y = y[1:] # Discard first 1 block. + df = reblock_by_autocorr(y, verbose=verbose) + print(df) + print() + + y = number_data["Nav"] + y = y[1:] # Discard first 1 block. + df = reblock_by_autocorr(y, verbose=verbose) + print(df) diff --git a/examples/16-ft_afqmc/ueg.py b/examples/16-ft_afqmc/ueg.py new file mode 100644 index 00000000..8c08db48 --- /dev/null +++ b/examples/16-ft_afqmc/ueg.py @@ -0,0 +1,564 @@ +import numpy +import scipy.sparse +from ipie.utils.io import write_qmcpack_sparse + + +class UEG(object): + """UEG system class (integrals read from fcidump) + + Parameters + ---------- + nup : int + Number of up electrons. + + ndown : int + Number of down electrons. + + rs : float + Density parameter. + + ecut : float + Scaled cutoff energy. + + ktwist : :class:`numpy.ndarray` + Twist vector. + + verbose : bool + Print extra information. + + Attributes + ---------- + T : :class:`numpy.ndarray` + One-body part of the Hamiltonian. This is diagonal in plane wave basis. + + ecore : float + Madelung contribution to the total energy. + + h1e_mod : :class:`numpy.ndarray` + Modified one-body Hamiltonian. + + nfields : int + Number of field configurations per walker for back propagation. + + basis : :class:`numpy.ndarray` + Basis vectors within a cutoff. + + kfac : float + Scale factor (2pi/L). + """ + + def __init__(self, options, verbose=False): + if verbose: + print("# Parsing input options.") + + self.name = "UEG" + self.nup = options.get("nup") + self.ndown = options.get("ndown") + self.nelec = (self.nup, self.ndown) + self.rs = options.get("rs") + self.ecut = options.get("ecut") + self.ktwist = numpy.array(options.get("ktwist", [0, 0, 0])).reshape(3) + + self.thermal = options.get("thermal", False) + self._alt_convention = options.get("alt_convention", False) + self.write_ints = options.get("write_integrals", False) + + self.sparse = True + self.control_variate = False + self.diagH1 = True + + # Total # of electrons. + self.ne = self.nup + self.ndown + # Spin polarisation. + self.zeta = (self.nup - self.ndown) / self.ne + # Density. + self.rho = ((4.0 * numpy.pi) / 3.0 * self.rs**3.0) ** (-1.0) + # Box Length. + self.L = self.rs * (4.0 * self.ne * numpy.pi / 3.0) ** (1 / 3.0) + # Volume + self.vol = self.L**3.0 + # k-space grid spacing. + self.kfac = 2 * numpy.pi / self.L + # Fermi Wavevector (infinite system). + self.kf = (3 * (self.zeta + 1) * numpy.pi**2 * self.ne / self.L**3) ** (1 / 3.0) + # Fermi energy (inifinite systems). + self.ef = 0.5 * self.kf**2 + # Core energy. + self.ecore = 0.5 * self.ne * self.madelung() + + if verbose: + if self.thermal: + print("# Thermal UEG activated.") + + print(f"# Number of spin-up electrons: {self.nup:d}") + print(f"# Number of spin-down electrons: {self.ndown:d}") + print(f"# rs: {self.rs:6.4e}") + print(f"# Spin polarisation (zeta): {self.zeta:6.4e}") + print(f"# Electron density (rho): {self.rho:13.8e}") + print(f"# Box Length (L): {self.L:13.8e}") + print(f"# Volume: {self.vol:13.8e}") + print(f"# k-space factor (2pi/L): {self.kfac:13.8e}") + + + def build(self, verbose=False): + # Get plane wave basis vectors and corresponding eigenvalues. + self.sp_eigv, self.basis, self.nmax = self.sp_energies( + self.ktwist, self.kfac, self.ecut) + self.shifted_nmax = 2 * self.nmax + self.imax_sq = numpy.dot(self.basis[-1], self.basis[-1]) + self.create_lookup_table() + + for i, k in enumerate(self.basis): + assert i == self.lookup_basis(k) + + # Number of plane waves. + self.nbasis = len(self.sp_eigv) + self.nactive = self.nbasis + self.ncore = 0 + self.nfv = 0 + self.mo_coeff = None + + # --------------------------------------------------------------------- + T = numpy.diag(self.sp_eigv) + h1e_mod = self.mod_one_body(T) + self.H1 = numpy.array([T, T]) # Making alpha and beta. + self.h1e_mod = numpy.array([h1e_mod, h1e_mod]) + + # --------------------------------------------------------------------- + # Allowed momentum transfers (4*ecut). + eigs, qvecs, self.qnmax = self.sp_energies(self.ktwist, self.kfac, 4 * self.ecut) + + # Omit Q = 0 term. + self.qvecs = numpy.copy(qvecs[1:]) + self.vqvec = numpy.array([self.vq(self.kfac * q) for q in self.qvecs]) + + # Number of momentum transfer vectors / auxiliary fields. + # Can reduce by symmetry but be stupid for the moment. + self.nchol = len(self.qvecs) + self.nfields = 2 * len(self.qvecs) + self.get_momentum_transfers() + + if verbose: + print(f"# Number of plane waves: {self.nbasis:d}") + print(f"# Number of Cholesky vectors: {self.nchol:d}.") + print(f"# Number of auxiliary fields: {self.nfields:d}.") + print("# Constructing two-body potentials incore.") + + # --------------------------------------------------------------------- + self.chol_vecs, self.iA, self.iB = self.two_body_potentials_incore() + + if self.write_ints: + self.write_integrals() + + if verbose: + print("# Approximate memory required for " + "two-body potentials: {:13.8e} GB.".format((3 * self.iA.nnz * 16 / (1024**3)))) + print("# Finished constructing two-body potentials.") + print("# Finished building UEG object.") + + + def sp_energies(self, ks, kfac, ecut): + """Calculate the allowed kvectors and resulting single particle eigenvalues (basically kinetic energy) + which can fit in the sphere in kspace determined by ecut. + + Parameters + ---------- + kfac : float + kspace grid spacing. + + ecut : float + energy cutoff. + + Returns + ------- + spval : :class:`numpy.ndarray` + Array containing sorted single particle eigenvalues. + + kval : :class:`numpy.ndarray` + Array containing basis vectors, sorted according to their + corresponding single-particle energy. + """ + + # Scaled Units to match with HANDE. + # So ecut is measured in units of 1/kfac^2. + nmax = int(numpy.ceil(numpy.sqrt((2 * ecut)))) + + spval = [] + vec = [] + kval = [] + + for ni in range(-nmax, nmax + 1): + for nj in range(-nmax, nmax + 1): + for nk in range(-nmax, nmax + 1): + spe = 0.5 * (ni**2 + nj**2 + nk**2) + + if spe <= ecut: + kijk = [ni, nj, nk] + + # Reintroduce 2 \pi / L factor. + ek = 0.5 * numpy.dot(numpy.array(kijk) + ks, numpy.array(kijk) + ks) + kval.append(kijk) + spval.append(kfac**2 * ek) + + # Sort the arrays in terms of increasing energy. + spval = numpy.array(spval) + ix = numpy.argsort(spval, kind="mergesort") + spval = spval[ix] + kval = numpy.array(kval)[ix] + return spval, kval, nmax + + + def create_lookup_table(self): + basis_ix = [] + for k in self.basis: + basis_ix.append(self.map_basis_to_index(k)) + + self.lookup = numpy.zeros(max(basis_ix) + 1, dtype=int) + + for i, b in enumerate(basis_ix): + self.lookup[b] = i + + self.max_ix = max(basis_ix) + + + def lookup_basis(self, vec): + if numpy.dot(vec, vec) <= self.imax_sq: + ix = self.map_basis_to_index(vec) + + if ix >= len(self.lookup): + ib = None + + else: + ib = self.lookup[ix] + + return ib + + else: + ib = None + + + def map_basis_to_index(self, k): + return ((k[0] + self.nmax) + + self.shifted_nmax * (k[1] + self.nmax) + + self.shifted_nmax * self.shifted_nmax * (k[2] + self.nmax)) + + + def get_momentum_transfers(self): + """Get arrays of plane wave basis vectors connected by momentum transfers Q.""" + nlimit = self.nup + if self.thermal: + nlimit = self.nbasis + + self.ikpq_i = [] + self.ikpq_kpq = [] + + for iq, q in enumerate(self.qvecs): + idxkpq_list_i = [] + idxkpq_list_kpq = [] + + for i, k in enumerate(self.basis[0:nlimit]): + kpq = k + q + idxkpq = self.lookup_basis(kpq) + + if idxkpq is not None: + idxkpq_list_i += [i] + idxkpq_list_kpq += [idxkpq] + + self.ikpq_i += [idxkpq_list_i] + self.ikpq_kpq += [idxkpq_list_kpq] + + # --------------------------------------------------------------------- + self.ipmq_i = [] + self.ipmq_pmq = [] + + for iq, q in enumerate(self.qvecs): + idxpmq_list_i = [] + idxpmq_list_pmq = [] + + for i, p in enumerate(self.basis[0:nlimit]): + pmq = p - q + idxpmq = self.lookup_basis(pmq) + + if idxpmq is not None: + idxpmq_list_i += [i] + idxpmq_list_pmq += [idxpmq] + + self.ipmq_i += [idxpmq_list_i] + self.ipmq_pmq += [idxpmq_list_pmq] + + for iq, q in enumerate(self.qvecs): + self.ikpq_i[iq] = numpy.array(self.ikpq_i[iq], dtype=numpy.int64) + self.ikpq_kpq[iq] = numpy.array(self.ikpq_kpq[iq], dtype=numpy.int64) + self.ipmq_i[iq] = numpy.array(self.ipmq_i[iq], dtype=numpy.int64) + self.ipmq_pmq[iq] = numpy.array(self.ipmq_pmq[iq], dtype=numpy.int64) + + + def madelung(self): + """Use expression in Schoof et al. (PhysRevLett.115.130402) for the + Madelung contribution to the total energy fitted to L.M. Fraser et al. + Phys. Rev. B 53, 1814. + + Parameters + ---------- + rs : float + Wigner-Seitz radius. + + ne : int + Number of electrons. + + Returns + ------- + v_M: float + Madelung potential (in Hartrees). + """ + c1 = -2.837297 + c2 = (3.0 / (4.0 * numpy.pi)) ** (1.0 / 3.0) + return c1 * c2 / (self.ne ** (1.0 / 3.0) * self.rs) + + + def mod_one_body(self, T): + """Absorb the diagonal term of the two-body Hamiltonian to the one-body term. + Essentially adding the third term in Eq.(11b) of Phys. Rev. B 75, 245123. + + Parameters + ---------- + T : float + one-body Hamiltonian (i.e. kinetic energy) + + Returns + ------- + h1e_mod: float + modified one-body Hamiltonian + """ + h1e_mod = numpy.copy(T) + + fac = 1.0 / (2.0 * self.vol) + for i, ki in enumerate(self.basis): + for j, kj in enumerate(self.basis): + if i != j: + q = self.kfac * (ki - kj) + h1e_mod[i, i] = h1e_mod[i, i] - fac * self.vq(q) + + return h1e_mod + + + def vq(self, q): + """The typical 3D Coulomb kernel + + Parameters + ---------- + q : float + a plane-wave vector + + Returns + ------- + v_M: float + 3D Coulomb kernel (in Hartrees) + """ + return 4 * numpy.pi / numpy.dot(q, q) + + + def density_operator(self, iq): + """Density operator as defined in Eq.(6) of Phys. Rev. B 75, 245123. + + Parameters + ---------- + q : float + a plane-wave vector + + Returns + ------- + rho_q: float + density operator + """ + nnz = self.rho_ikpq_kpq[iq].shape[0] # Number of non-zeros + ones = numpy.ones((nnz), dtype=numpy.complex128) + rho_q = scipy.sparse.csc_matrix( + (ones, (self.rho_ikpq_kpq[iq], self.rho_ikpq_i[iq])), + shape=(self.nbasis, self.nbasis), + dtype=numpy.complex128) + return rho_q + + + def scaled_density_operator_incore(self, transpose): + """Density operator as defined in Eq.(6) of PRB(75)245123 + + Parameters + ---------- + q : float + a plane-wave vector + + Returns + ------- + rho_q: float + density operator + """ + rho_ikpq_i = [] + rho_ikpq_kpq = [] + + for iq, q in enumerate(self.qvecs): + idxkpq_list_i = [] + idxkpq_list_kpq = [] + + for i, k in enumerate(self.basis): + kpq = k + q + idxkpq = self.lookup_basis(kpq) + + if idxkpq is not None: + idxkpq_list_i += [i] + idxkpq_list_kpq += [idxkpq] + + rho_ikpq_i += [idxkpq_list_i] + rho_ikpq_kpq += [idxkpq_list_kpq] + + for iq, q in enumerate(self.qvecs): + rho_ikpq_i[iq] = numpy.array(rho_ikpq_i[iq], dtype=numpy.int64) + rho_ikpq_kpq[iq] = numpy.array(rho_ikpq_kpq[iq], dtype=numpy.int64) + + nq = len(self.qvecs) + nnz = 0 + for iq in range(nq): + nnz += rho_ikpq_kpq[iq].shape[0] + + col_index = [] + row_index = [] + values = [] + + if transpose: + for iq in range(nq): + qscaled = self.kfac * self.qvecs[iq] + # Due to the HS transformation, we have to do pi / 2*vol as opposed to 2*pi / vol + piovol = numpy.pi / (self.vol) + factor = (piovol / numpy.dot(qscaled, qscaled)) ** 0.5 + + for innz, kpq in enumerate(rho_ikpq_kpq[iq]): + row_index += [rho_ikpq_kpq[iq][innz] + rho_ikpq_i[iq][innz] * self.nbasis] + col_index += [iq] + values += [factor] + + else: + for iq in range(nq): + qscaled = self.kfac * self.qvecs[iq] + # Due to the HS transformation, we have to do pi / 2*vol as opposed to 2*pi / vol + piovol = numpy.pi / (self.vol) + factor = (piovol / numpy.dot(qscaled, qscaled)) ** 0.5 + + for innz, kpq in enumerate(rho_ikpq_kpq[iq]): + row_index += [rho_ikpq_kpq[iq][innz] * self.nbasis + rho_ikpq_i[iq][innz]] + col_index += [iq] + values += [factor] + + rho_q = scipy.sparse.csc_matrix( + (values, (row_index, col_index)), + shape=(self.nbasis * self.nbasis, nq), + dtype=numpy.complex128) + return rho_q + + + def two_body_potentials_incore(self): + """Calculate A and B of Eq.(13) of PRB(75)245123 for a given plane-wave vector q + + Returns + ------- + iA : numpy array + Eq.(13a) + + iB : numpy array + Eq.(13b) + """ + rho_q = self.scaled_density_operator_incore(False) + rho_qH = self.scaled_density_operator_incore(True) + iA = 1j * (rho_q + rho_qH) + iB = -(rho_q - rho_qH) + return (rho_q, iA, iB) + + + def hijkl(self, i, j, k, l): + """Compute = (ik|jl) = 1/Omega * 4pi/(kk-ki)**2 + + Checks for momentum conservation k_i + k_j = k_k + k_k, or + k_k - k_i = k_j - k_l. + + Parameters + ---------- + i, j, k, l : int + Orbital indices for integral (ik|jl) = . + + Returns + ------- + integral : float + (ik|jl) + """ + q1 = self.basis[k] - self.basis[i] + q2 = self.basis[j] - self.basis[l] + + if numpy.dot(q1, q1) > 1e-12 and numpy.dot(q1 - q2, q1 - q2) < 1e-12: + return 1.0 / self.vol * self.vq(self.kfac * q1) + + else: + return 0.0 + + + def compute_real_transformation(self): + U22 = numpy.zeros((2, 2), dtype=numpy.complex128) + U22[0, 0] = 1.0 / numpy.sqrt(2.0) + U22[0, 1] = 1.0 / numpy.sqrt(2.0) + U22[1, 0] = -1.0j / numpy.sqrt(2.0) + U22[1, 1] = 1.0j / numpy.sqrt(2.0) + + U = numpy.zeros((self.nbasis, self.nbasis), dtype=numpy.complex128) + + for i, b in enumerate(self.basis): + if numpy.sum(b * b) == 0: + U[i, i] = 1.0 + + else: + mb = -b + diff = numpy.einsum("ij->i", (self.basis - mb) ** 2) + idx = numpy.argwhere(diff == 0) + assert idx.ravel().shape[0] == 1 + + if i < idx: + idx = idx.ravel()[0] + U[i, i] = U22[0, 0] + U[i, idx] = U22[0, 1] + U[idx, i] = U22[1, 0] + U[idx, idx] = U22[1, 1] + + else: + continue + + U = U.T.copy() + return U + + + def eri_4(self): + eri_chol = 4 * self.chol_vecs.dot(self.chol_vecs.T) + eri_chol = ( + eri_chol.toarray().reshape((self.nbasis, self.nbasis, self.nbasis, self.nbasis)).real) + eri_chol = eri_chol.transpose(0, 1, 3, 2) + return eri_chol + + + def eri_8(self): + """Compute 8-fold symmetric integrals. Useful for running standard + quantum chemistry methods,""" + eri = self.eri_4() + U = self.compute_real_transformation() + eri0 = numpy.einsum("mp,mnls->pnls", U.conj(), eri, optimize=True) + eri1 = numpy.einsum("nq,pnls->pqls", U, eri0, optimize=True) + eri2 = numpy.einsum("lr,pqls->pqrs", U.conj(), eri1, optimize=True) + eri3 = numpy.einsum("st,pqrs->pqrt", U, eri2, optimize=True).real + return eri3 + + + def write_integrals(self, filename="ueg_integrals.h5"): + write_qmcpack_sparse( + self.H1[0], + 2 * self.chol_vecs.toarray(), + self.nelec, + self.nbasis, + #enuc=self.ecore, + enuc=0., + filename=filename) + diff --git a/ipie/addons/thermal/__init__.py b/ipie/addons/thermal/__init__.py new file mode 100644 index 00000000..f91ef518 --- /dev/null +++ b/ipie/addons/thermal/__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. +# +# Authors: Fionn Malone +# Joonho Lee +# diff --git a/ipie/addons/thermal/analysis/__init__.py b/ipie/addons/thermal/analysis/__init__.py new file mode 100644 index 00000000..f91ef518 --- /dev/null +++ b/ipie/addons/thermal/analysis/__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. +# +# Authors: Fionn Malone +# Joonho Lee +# diff --git a/ipie/addons/thermal/analysis/extraction.py b/ipie/addons/thermal/analysis/extraction.py new file mode 100644 index 00000000..913cafb5 --- /dev/null +++ b/ipie/addons/thermal/analysis/extraction.py @@ -0,0 +1,64 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + + +from ipie.utils.misc import get_from_dict + +def set_info(frame, md): + ncols = len(frame.columns) + system = md.get("system") + hamiltonian = md.get("hamiltonian") + trial = md.get("trial") + qmc = md.get("params") + fp = get_from_dict(md, ["propagators", "free_projection"]) + bp = get_from_dict(md, ["estimates", "estimates", "back_prop"]) + + br = qmc.get("beta_scaled") + + ints = system.get("integral_file") + chol = system.get("threshold") + + frame["nup"] = system.get("nup") + frame["ndown"] = system.get("ndown") + frame["mu"] = qmc.get("mu") + frame["beta"] = qmc.get("beta") + frame["dt"] = qmc.get("timestep") + frame["ntot_walkers"] = qmc.get("total_num_walkers", 0) + frame["nbasis"] = hamiltonian.get("nbasis", 0) + + if trial is not None: + frame["mu_T"] = trial.get("mu") + frame["Nav_T"] = trial.get("nav") + + if fp is not None: + frame["free_projection"] = fp + + if bp is not None: + frame["tau_bp"] = bp["tau_bp"] + + if br is not None: + frame["beta_red"] = br + + if ints is not None: + frame["integrals"] = ints + + if chol is not None: + frame["cholesky_treshold"] = chol + + return list(frame.columns[ncols:]) + diff --git a/ipie/addons/thermal/analysis/thermal_analysis.py b/ipie/addons/thermal/analysis/thermal_analysis.py new file mode 100644 index 00000000..b9cba603 --- /dev/null +++ b/ipie/addons/thermal/analysis/thermal_analysis.py @@ -0,0 +1,163 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +#!/usr/bin/env python + +import sys +import argparse + +import glob +import numpy +import scipy.optimize +import pandas as pd + +from ipie.analysis.extraction import ( + extract_observable, + get_metadata, + get_sys_param + ) + +from ipie.addons.thermal.analysis.extraction import set_info + + +def parse_args(args): + """Parse command-line arguments. + + Parameters + ---------- + args : list of strings + command-line arguments. + + Returns + ------- + options : :class:`argparse.ArgumentParser` + Command line arguments. + """ + + parser = argparse.ArgumentParser(description = __doc__) + parser.add_argument('-c', '--chem-pot', dest='fit_chem_pot', + action='store_true', default=False, + help='Estimate optimal chemical potential') + parser.add_argument('-n', '--nav', dest='nav', type=float, + help='Target electron density.') + parser.add_argument('-o', '--order', dest='order', type=int, + default=3, help='Order polynomial to fit.') + parser.add_argument('-p', '--plot', dest='plot', action='store_true', + help='Plot density vs. mu.') + parser.add_argument('-f', nargs='+', dest='filenames', + help='Space-separated list of files to analyse.') + + options = parser.parse_args(args) + + if not options.filenames: + parser.print_help() + sys.exit(1) + + return options + + +def analyse(files, block_idx=1): + sims = [] + files = sorted(files) + + for f in files: + data_energy = extract_observable(f, name='energy', block_idx=block_idx) + data_nav = extract_observable(f, name='nav', block_idx=block_idx) + data = pd.concat([data_energy, data_nav['Nav']], axis=1) + md = get_metadata(f) + keys = set_info(data, md) + sims.append(data[1:]) + + full = pd.concat(sims).groupby(keys, sort=False) + + analysed = [] + for i, g in full: + cols = ["ETotal", "E1Body", "E2Body", "Nav"] + averaged = pd.DataFrame(index=[0]) + + for c in cols: + mean = numpy.real(g[c].values).mean() + error = scipy.stats.sem(numpy.real(g[c].values), ddof=1) + averaged[c] = [mean] + averaged[c + "_error"] = [error] + + for k, v in zip(full.keys, i): + averaged[k] = v + + analysed.append(averaged) + + return pd.concat(analysed).reset_index(drop=True).sort_values(by=keys) + + +def nav_mu(mu, coeffs): + return numpy.polyval(coeffs, mu) + + +def find_chem_pot(data, target, vol, order=3, plot=False): + print(f"# System volume: {vol}.") + print(f"# Target number of electrons: {vol * target}.") + nav = data.Nav.values / vol + nav_error = data.Nav_error.values / vol + # Half filling special case where error bar is zero. + zeros = numpy.where(nav_error == 0)[0] + nav_error[zeros] = 1e-8 + mus = data.mu.values + delta = nav - target + s = 0 + e = len(delta) + rmin = None + +def main(args): + """Run reblocking and data analysis on PAUXY output. + + Parameters + ---------- + args : list of strings + command-line arguments. + + Returns + ------- + None. + """ + + options = parse_args(args) + if '*' in options.filenames[0]: + files = glob.glob(options.filenames[0]) + + else: + files = options.filenames + + data = analyse(files) + + if options.fit_chem_pot: + name = get_sys_param(files[0], 'name') + vol = 1. + mu = find_chem_pot(data, options.nav, vol, + order=options.order, plot=options.plot) + + if mu is not None: + print("# Optimal chemical potential found to be: {}.".format(mu)) + + else: + print("# Failed to find chemical potential.") + + print(data.to_string(index=False)) + + +if __name__ == '__main__': + main(sys.argv[1:]) diff --git a/ipie/addons/thermal/estimators/__init__.py b/ipie/addons/thermal/estimators/__init__.py new file mode 100644 index 00000000..f91ef518 --- /dev/null +++ b/ipie/addons/thermal/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. +# +# Authors: Fionn Malone +# Joonho Lee +# diff --git a/ipie/addons/thermal/estimators/energy.py b/ipie/addons/thermal/estimators/energy.py new file mode 100644 index 00000000..7ebe8637 --- /dev/null +++ b/ipie/addons/thermal/estimators/energy.py @@ -0,0 +1,57 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +from typing import Union + +from ipie.utils.backend import arraylib as xp +from ipie.hamiltonians.generic import GenericComplexChol, GenericRealChol +from ipie.estimators.energy import EnergyEstimator + +from ipie.addons.thermal.walkers.uhf_walkers import UHFThermalWalkers +from ipie.addons.thermal.estimators.thermal import one_rdm_from_G +from ipie.addons.thermal.estimators.generic import local_energy_generic_cholesky + + +def local_energy( + hamiltonian: Union[GenericRealChol, GenericComplexChol], + walkers: UHFThermalWalkers): + energies = xp.zeros((walkers.nwalkers, 3), dtype=xp.complex128) + + for iw in range(walkers.nwalkers): + # Want the full Green's function when calculating observables. + walkers.calc_greens_function(iw, slice_ix=walkers.stack[iw].nslice) + P = one_rdm_from_G(xp.array([walkers.Ga[iw], walkers.Gb[iw]])) + energy = local_energy_generic_cholesky(hamiltonian, P) + energies[iw] = energy + + return energies + + +class ThermalEnergyEstimator(EnergyEstimator): + def __init__(self, system=None, hamiltonian=None, trial=None, filename=None): + super().__init__(system=system, ham=hamiltonian, trial=trial, filename=filename) + + def compute_estimator(self, walkers, hamiltonian, trial, istep=1): + # Need to be able to dispatch here. + # Re-calculated Green's function in `local_energy`. + energy = local_energy(hamiltonian, walkers) + self._data["ENumer"] = xp.sum(walkers.weight * energy[:, 0].real) + 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 diff --git a/ipie/addons/thermal/estimators/generic.py b/ipie/addons/thermal/estimators/generic.py new file mode 100644 index 00000000..0ad971fa --- /dev/null +++ b/ipie/addons/thermal/estimators/generic.py @@ -0,0 +1,138 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import plum +import numpy +from ipie.hamiltonians.generic import GenericRealChol, GenericComplexChol + + +@plum.dispatch +def local_energy_generic_cholesky(hamiltonian: GenericRealChol, P): + r"""Calculate local for generic two-body hamiltonian. + + This uses the cholesky decomposed two-electron integrals. + + Parameters + ---------- + hamiltonian : :class:`Generic` + ab-initio hamiltonian information + P : :class:`numpy.ndarray` + Walker's density matrix. + + Returns + ------- + (E, T, V): tuple + Local, kinetic and potential energies. + """ + # Element wise multiplication. + e1b = numpy.sum(hamiltonian.H1[0] * P[0]) + numpy.sum(hamiltonian.H1[1] * P[1]) + nbasis = hamiltonian.nbasis + nchol = hamiltonian.nchol + Pa, Pb = P[0], P[1] + + # Ecoul. + Xa = hamiltonian.chol.T.dot(Pa.real.ravel()) + 1.0j * hamiltonian.chol.T.dot(Pa.imag.ravel()) + Xb = hamiltonian.chol.T.dot(Pb.real.ravel()) + 1.0j * hamiltonian.chol.T.dot(Pb.imag.ravel()) + X = Xa + Xb + ecoul = 0.5 * numpy.dot(X, X) + + # Ex. + PaT = Pa.T.copy() + PbT = Pb.T.copy() + T = numpy.zeros((nbasis, nbasis), dtype=numpy.complex128) + exx = 0.0j # we will iterate over cholesky index to update Ex energy for alpha and beta + + for x in range(nchol): # Write a numba function that calls BLAS for this. + Lmn = hamiltonian.chol[:, x].reshape((nbasis, nbasis)) + T[:, :].real = PaT.real.dot(Lmn) + T[:, :].imag = PaT.imag.dot(Lmn) + exx += numpy.trace(T.dot(T)) + T[:, :].real = PbT.real.dot(Lmn) + T[:, :].imag = PbT.imag.dot(Lmn) + exx += numpy.trace(T.dot(T)) + + exx *= 0.5 + e2b = ecoul - exx + return (e1b + e2b + hamiltonian.ecore, e1b + hamiltonian.ecore, e2b) + + +@plum.dispatch +def local_energy_generic_cholesky(hamiltonian: GenericComplexChol, P): + r"""Calculate local for generic two-body hamiltonian. + + This uses the cholesky decomposed two-electron integrals. + + Parameters + ---------- + hamiltonian : :class:`Generic` + ab-initio hamiltonian information + P : :class:`numpy.ndarray` + Walker's density matrix. + + Returns + ------- + (E, T, V): tuple + Local, kinetic and potential energies. + """ + # Element wise multiplication. + e1b = numpy.sum(hamiltonian.H1[0] * P[0]) + numpy.sum(hamiltonian.H1[1] * P[1]) + nbasis = hamiltonian.nbasis + nchol = hamiltonian.nchol + Pa, Pb = P[0], P[1] + + # Ecoul. + XAa = hamiltonian.A.T.dot(Pa.ravel()) + XAb = hamiltonian.A.T.dot(Pb.ravel()) + XA = XAa + XAb + + XBa = hamiltonian.B.T.dot(Pa.ravel()) + XBb = hamiltonian.B.T.dot(Pb.ravel()) + XB = XBa + XBb + + ecoul = 0.5 * (numpy.dot(XA, XA) + numpy.dot(XB, XB)) + + # Ex. + PaT = Pa.T.copy() + PbT = Pb.T.copy() + TA = numpy.zeros((nbasis, nbasis), dtype=numpy.complex128) + TB = numpy.zeros((nbasis, nbasis), dtype=numpy.complex128) + exx = 0.0j # we will iterate over cholesky index to update Ex energy for alpha and beta + + for x in range(nchol): # write a cython function that calls blas for this. + Amn = hamiltonian.A[:, x].reshape((nbasis, nbasis)) + Bmn = hamiltonian.B[:, x].reshape((nbasis, nbasis)) + TA[:, :] = PaT.dot(Amn) + TB[:, :] = PaT.dot(Bmn) + exx += numpy.trace(TA.dot(TA)) + numpy.trace(TB.dot(TB)) + + TA[:, :] = PbT.dot(Amn) + TB[:, :] = PbT.dot(Bmn) + exx += numpy.trace(TA.dot(TA)) + numpy.trace(TB.dot(TB)) + + exx *= 0.5 + e2b = ecoul - exx + return (e1b + e2b + hamiltonian.ecore, e1b + hamiltonian.ecore, e2b) + + +def fock_generic(hamiltonian, P): + nbasis = hamiltonian.nbasis + nchol = hamiltonian.nchol + hs_pot = hamiltonian.chol.T.reshape(nchol, nbasis, nbasis) + mf_shift = 1j * numpy.einsum("lpq,spq->l", hs_pot, P) + VMF = 1j * numpy.einsum("lpq,l->pq", hs_pot, mf_shift) + return hamiltonian.h1e_mod - VMF diff --git a/ipie/addons/thermal/estimators/greens_function.py b/ipie/addons/thermal/estimators/greens_function.py new file mode 100644 index 00000000..c39abd19 --- /dev/null +++ b/ipie/addons/thermal/estimators/greens_function.py @@ -0,0 +1,136 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import numpy +import scipy.linalg + +def greens_function(A): + r"""Construct Greens function from density matrix. + + .. math:: + G_{ij} = \langle c_{i} c_j^{\dagger} \rangle \\ + = \left[\frac{1}{1+A}\right]_{ij} + + Uses stable algorithm from White et al. (1988) + + Parameters + ---------- + A : :class:`numpy.ndarray` + Density matrix (product of B matrices). + + Returns + ------- + G : :class:`numpy.ndarray` + Thermal Green's function. + """ + G = numpy.zeros(A.shape, dtype=A.dtype) + (U1, S1, V1) = scipy.linalg.svd(A) + T = numpy.dot(U1.conj().T, V1.conj().T) + numpy.diag(S1) + (U2, S2, V2) = scipy.linalg.svd(T) + U3 = numpy.dot(U1, U2) + D3 = numpy.diag(1.0 / S2) + V3 = numpy.dot(V2, V1) + G = (V3.conj().T).dot(D3).dot(U3.conj().T) + return G + + +def greens_function_qr_strat(walkers, iw, slice_ix=None, inplace=True): + """Compute the Green's function for walker with index `iw` at time + `slice_ix`. Uses the Stratification method (DOI 10.1109/IPDPS.2012.37) + """ + stack_iw = walkers.stack[iw] + + if slice_ix == None: + slice_ix = stack_iw.time_slice + + bin_ix = slice_ix // stack_iw.stack_size + # For final time slice want first block to be the rightmost (for energy + # evaluation). + if bin_ix == stack_iw.nstack: + bin_ix = -1 + + Ga_iw, Gb_iw = None, None + if not inplace: + Ga_iw = numpy.zeros(walkers.Ga[iw].shape, walkers.Ga.dtype) + Gb_iw = numpy.zeros(walkers.Gb[iw].shape, walkers.Gb.dtype) + + for spin in [0, 1]: + # Need to construct the product A(l) = B_l B_{l-1}..B_L...B_{l+1} in + # stable way. Iteratively construct column pivoted QR decompositions + # (A = QDT) starting from the rightmost (product of) propagator(s). + B = stack_iw.get((bin_ix + 1) % stack_iw.nstack) + + (Q1, R1, P1) = scipy.linalg.qr(B[spin], pivoting=True, check_finite=False) + # Form D matrices + D1 = numpy.diag(R1.diagonal()) + D1inv = numpy.diag(1.0 / R1.diagonal()) + T1 = numpy.einsum("ii,ij->ij", D1inv, R1) + # permute them + T1[:, P1] = T1[:, range(walkers.nbasis)] + + for i in range(2, stack_iw.nstack + 1): + ix = (bin_ix + i) % stack_iw.nstack + B = stack_iw.get(ix) + C2 = numpy.dot(numpy.dot(B[spin], Q1), D1) + (Q1, R1, P1) = scipy.linalg.qr(C2, pivoting=True, check_finite=False) + # Compute D matrices + D1inv = numpy.diag(1.0 / R1.diagonal()) + D1 = numpy.diag(R1.diagonal()) + tmp = numpy.einsum("ii,ij->ij", D1inv, R1) + tmp[:, P1] = tmp[:, range(walkers.nbasis)] + T1 = numpy.dot(tmp, T1) + + # G^{-1} = 1+A = 1+QDT = Q (Q^{-1}T^{-1}+D) T + # Write D = Db^{-1} Ds + # Then G^{-1} = Q Db^{-1}(Db Q^{-1}T^{-1}+Ds) T + Db = numpy.zeros(B[spin].shape, B[spin].dtype) + Ds = numpy.zeros(B[spin].shape, B[spin].dtype) + for i in range(Db.shape[0]): + absDlcr = abs(Db[i, i]) + if absDlcr > 1.0: + Db[i, i] = 1.0 / absDlcr + Ds[i, i] = numpy.sign(D1[i, i]) + else: + Db[i, i] = 1.0 + Ds[i, i] = D1[i, i] + + T1inv = scipy.linalg.inv(T1, check_finite=False) + # C = (Db Q^{-1}T^{-1}+Ds) + C = numpy.dot(numpy.einsum("ii,ij->ij", Db, Q1.conj().T), T1inv) + Ds + Cinv = scipy.linalg.inv(C, check_finite=False) + + # Then G = T^{-1} C^{-1} Db Q^{-1} + # Q is unitary. + if inplace: + if spin == 0: + walkers.Ga[iw] = numpy.dot( + numpy.dot(T1inv, Cinv), numpy.einsum("ii,ij->ij", Db, Q1.conj().T)) + else: + walkers.Gb[iw] = numpy.dot( + numpy.dot(T1inv, Cinv), numpy.einsum("ii,ij->ij", Db, Q1.conj().T)) + + else: + if spin == 0: + Ga_iw = numpy.dot( + numpy.dot(T1inv, Cinv), numpy.einsum("ii,ij->ij", Db, Q1.conj().T)) + + else: + Gb_iw = numpy.dot( + numpy.dot(T1inv, Cinv), numpy.einsum("ii,ij->ij", Db, Q1.conj().T)) + + return Ga_iw, Gb_iw diff --git a/ipie/addons/thermal/estimators/handler.py b/ipie/addons/thermal/estimators/handler.py new file mode 100644 index 00000000..8ae4b8b7 --- /dev/null +++ b/ipie/addons/thermal/estimators/handler.py @@ -0,0 +1,176 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +"""Routines and classes for estimation of observables.""" + +import os +from typing import Tuple, Union + +from ipie.config import config, MPI +from ipie.estimators.handler import EstimatorHandler + +from ipie.addons.thermal.estimators.energy import ThermalEnergyEstimator +from ipie.addons.thermal.estimators.particle_number import ThermalNumberEstimator + +# Some supported (non-custom) estimators +_predefined_estimators = { + "energy": ThermalEnergyEstimator, + "nav": ThermalNumberEstimator, +} + + +class ThermalEstimatorHandler(EstimatorHandler): + """Container for qmc options of observables. + + Parameters + ---------- + comm : MPI.COMM_WORLD + MPI Communicator. + hamiltonian : :class:`ipie.hamiltonian.X' object + Hamiltonian describing the system. + trial : :class:`ipie.trial_wavefunction.X' object + Trial wavefunction class. + walker_state : :class:`WalkerAccumulator` object + WalkerAccumulator class. + verbose : bool + If true we print out additional setup information. + filename : str + .h5 file name for saving data. + basename : str + .h5 base name for saving data. + overwrite : bool + Whether to overwrite .h5 files. + observables : tuple + Tuple listing observables to be calculated. + + Attributes + ---------- + estimators : dict + Dictionary of estimator objects. + """ + + def __init__( + self, + comm, + hamiltonian, + trial, + walker_state=None, + verbose: bool = False, + filename: Union[str, None] = None, + basename: str = "estimates", + overwrite=True, + observables: Tuple[str] = ("energy", "nav"), # TODO: Use factory method! + ): + if verbose: + print("# Setting up estimator object.") + if comm.rank == 0: + self.basename = basename + self.filename = filename + self.index = 0 + if self.filename is None: + self.filename = f"{self.basename}.{self.index}.h5" + while os.path.isfile(self.filename) and not overwrite: + self.index = int(self.filename.split(".")[1]) + self.index = self.index + 1 + self.filename = f"{self.basename}.{self.index}.h5" + if verbose: + print(f"# Writing estimator data to {self.filename}") + else: + self.filename = None + self.buffer_size = config.get_option("estimator_buffer_size") + if walker_state is not None: + self.num_walker_props = walker_state.size + self.walker_header = walker_state.names + else: + self.num_walker_props = 0 + self.walker_header = "" + self._estimators = {} + self._shapes = [] + self._offsets = {} + self.json_string = "{}" + # TODO: Replace this, should be built outside + for obs in observables: + try: + est = _predefined_estimators[obs]( + hamiltonian=hamiltonian, + trial=trial) + self[obs] = est + except KeyError: + raise RuntimeError(f"unknown observable: {obs}") + if verbose: + print("# Finished settting up estimator object.") + + def compute_estimators(self, hamiltonian, trial, walker_batch): + """Update estimators with bached walkers. + + Parameters + ---------- + hamiltonian : :class:`ipie.hamiltonian.X' object + Hamiltonian describing the system. + trial : :class:`ipie.trial_wavefunction.X' object + Trial wavefunction class. + walker_batch : :class:`UHFThermalWalkers' object + Walkers class. + """ + # Compute all estimators + # For the moment only consider estimators compute per block. + # TODO: generalize for different block groups (loop over groups) + offset = self.num_walker_props + for k, e in self.items(): + e.compute_estimator(walker_batch, hamiltonian, trial) + start = offset + self.get_offset(k) + end = start + int(self[k].size) + self.local_estimates[start:end] += e.data + + def print_time_slice(self, comm, time_slice, walker_state): + """Print estimators at a time slice of the imgainary time propagation. + + Parameters + ---------- + comm : MPI.COMM_WORLD + MPI Communicator. + time_slice : int + Time slice. + walker_state : :class:`WalkerAccumulator` object + WalkerAccumulator class. + """ + comm.Reduce(self.local_estimates, self.global_estimates, op=MPI.SUM) + # Get walker data. + offset = walker_state.size + + if comm.rank == 0: + k = 'energy' + e = self[k] + start = offset + self.get_offset(k) + end = start + int(self[k].size) + estim_data = self.global_estimates[start:end] + e.post_reduce_hook(estim_data) + etotal = estim_data[e.get_index("ETotal")] + + k = 'nav' + e = self[k] + start = offset + self.get_offset(k) + end = start + int(self[k].size) + estim_data = self.global_estimates[start:end] + e.post_reduce_hook(estim_data) + nav = estim_data[e.get_index("Nav")] + + print(f"cut : {time_slice} {nav.real} {etotal.real}") + + self.zero() + diff --git a/ipie/addons/thermal/estimators/local_energy.py b/ipie/addons/thermal/estimators/local_energy.py new file mode 100644 index 00000000..34a06775 --- /dev/null +++ b/ipie/addons/thermal/estimators/local_energy.py @@ -0,0 +1,52 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import numpy +from typing import Union + +from ipie.hamiltonians.generic import GenericComplexChol, GenericRealChol +from ipie.addons.thermal.trial.one_body import OneBody +from ipie.addons.thermal.trial.mean_field import MeanField +from ipie.addons.thermal.estimators.generic import local_energy_generic_cholesky +from ipie.addons.thermal.estimators.thermal import one_rdm_from_G + +def local_energy_from_density_matrix( + hamiltonian: Union[GenericRealChol, GenericComplexChol], + trial: Union[OneBody, MeanField], + P: numpy.ndarray): + """Compute local energy from a given density matrix P. + + Parameters + ---------- + hamiltonian : hamiltonian object + Hamiltonian being studied. + trial : trial wavefunction object + Trial wavefunction. + P : np.ndarray + Walker density matrix. + + Returns: + ------- + local_energy : tuple / array + Total, one-body and two-body energies. + """ + assert len(P) == 2 + return local_energy_generic_cholesky(hamiltonian, P) + +def local_energy(hamiltonian, walker, trial): + return local_energy_from_density_matrix(hamiltonian, trial, one_rdm_from_G(walker.G)) diff --git a/ipie/addons/thermal/estimators/particle_number.py b/ipie/addons/thermal/estimators/particle_number.py new file mode 100644 index 00000000..63fa1f60 --- /dev/null +++ b/ipie/addons/thermal/estimators/particle_number.py @@ -0,0 +1,89 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import numpy +from ipie.utils.backend import arraylib as xp +from ipie.estimators.estimator_base import EstimatorBase +from ipie.addons.thermal.estimators.thermal import one_rdm_from_G + + +def particle_number(dmat: numpy.ndarray): + """Compute average particle number from the thermal 1RDM. + + Parameters + ---------- + dmat : :class:`numpy.ndarray` + Thermal 1RDM. + + Returns + ------- + nav : float + Average particle number. + """ + nav = dmat[0].trace() + dmat[1].trace() + return nav + + +class ThermalNumberEstimator(EstimatorBase): + def __init__(self, hamiltonian=None, trial=None, filename=None): + # We define a dictionary to contain whatever we want to compute. + # Note we typically want to separate the numerator and denominator of + # the estimator. + # We require complex valued buffers for accumulation + self._data = { + "NavNumer": 0.0j, + "NavDenom": 0.0j, + "Nav": 0.0j, + } + + # We also need to specify the shape of the desired estimator + self._shape = (len(self.names),) + + # Optional but good to know (we can redirect to custom filepath (ascii) + # and / or print to stdout but we shouldnt do this for non scalar + # quantities. + self._data_index = {k: i for i, k in enumerate(list(self._data.keys()))} + self.print_to_stdout = True + self.ascii_filename = filename + + # Must specify that we're dealing with array valued estimator. + self.scalar_estimator = True + + def compute_estimator(self, walkers, hamiltonian, trial): + for iw in range(walkers.nwalkers): + # Want the full Green's function when calculating observables. + walkers.calc_greens_function(iw, slice_ix=walkers.stack[iw].nslice) + nav_iw = particle_number(one_rdm_from_G( + xp.array([walkers.Ga[iw], walkers.Gb[iw]]))) + self._data["NavNumer"] += walkers.weight[iw] * nav_iw.real + + self._data["NavDenom"] = sum(walkers.weight) + + 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["Nav"] + ix_nume = self._data_index["NavNumer"] + ix_deno = self._data_index["NavDenom"] + data[ix_proj] = data[ix_nume] / data[ix_deno] diff --git a/ipie/addons/thermal/estimators/tests/__init__.py b/ipie/addons/thermal/estimators/tests/__init__.py new file mode 100644 index 00000000..381108b3 --- /dev/null +++ b/ipie/addons/thermal/estimators/tests/__init__.py @@ -0,0 +1,14 @@ +# 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/thermal/estimators/tests/test_estimators.py b/ipie/addons/thermal/estimators/tests/test_estimators.py new file mode 100644 index 00000000..15c9ebb7 --- /dev/null +++ b/ipie/addons/thermal/estimators/tests/test_estimators.py @@ -0,0 +1,172 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import pytest +import tempfile +import numpy +from typing import Tuple, Union + +from ipie.config import MPI +from ipie.hamiltonians.generic import Generic as HamGeneric +from ipie.hamiltonians.generic import GenericRealChol, GenericComplexChol + +from ipie.addons.thermal.estimators.energy import ThermalEnergyEstimator +from ipie.addons.thermal.estimators.particle_number import ThermalNumberEstimator +from ipie.addons.thermal.estimators.handler import ThermalEstimatorHandler +from ipie.addons.thermal.utils.testing import build_generic_test_case_handlers + +# System params. +nup = 5 +ndown = 5 +nelec = (nup, ndown) +ne = nup + ndown +nbasis = 10 + +# Thermal AFQMC params. +mu = -10. +beta = 0.1 +timestep = 0.01 +nwalkers = 10 +lowrank = False + +mf_trial = True +complex_integrals = False +debug = True +verbose = True +seed = 7 +numpy.random.seed(seed) + +@pytest.mark.unit +def test_energy_estimator(): + # Test. + objs = build_generic_test_case_handlers( + nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, + mf_trial=mf_trial, complex_integrals=complex_integrals, debug=debug, + seed=seed, verbose=verbose) + trial = objs['trial'] + hamiltonian = objs['hamiltonian'] + walkers = objs['walkers'] + + assert isinstance(hamiltonian, GenericRealChol) + chol = hamiltonian.chol + + # GenericRealChol. + re_estim = ThermalEnergyEstimator(hamiltonian=hamiltonian, trial=trial) + re_estim.compute_estimator(walkers, hamiltonian, trial) + assert len(re_estim.names) == 5 + assert re_estim["ENumer"].real == pytest.approx(24.66552451455761) + assert re_estim["ETotal"] == pytest.approx(0.0) + tmp = re_estim.data.copy() + re_estim.post_reduce_hook(tmp) + assert tmp[re_estim.get_index("ETotal")] == pytest.approx(2.4665524514557613) + assert re_estim.print_to_stdout + assert re_estim.ascii_filename == None + assert re_estim.shape == (5,) + header = re_estim.header_to_text + data_to_text = re_estim.data_to_text(tmp) + assert len(data_to_text.split()) == 5 + + # GenericComplexChol. + cx_chol = numpy.array(chol, dtype=numpy.complex128) + cx_hamiltonian = HamGeneric( + numpy.array(hamiltonian.H1, dtype=numpy.complex128), cx_chol, + hamiltonian.ecore, verbose=False) + + assert isinstance(cx_hamiltonian, GenericComplexChol) + + cx_estim = ThermalEnergyEstimator(hamiltonian=cx_hamiltonian, trial=trial) + cx_estim.compute_estimator(walkers, cx_hamiltonian, trial) + assert len(cx_estim.names) == 5 + assert cx_estim["ENumer"].real == pytest.approx(24.66552451455761) + assert cx_estim["ETotal"] == pytest.approx(0.0) + tmp = cx_estim.data.copy() + cx_estim.post_reduce_hook(tmp) + assert tmp[cx_estim.get_index("ETotal")] == pytest.approx(2.4665524514557613) + assert cx_estim.print_to_stdout + assert cx_estim.ascii_filename == None + assert cx_estim.shape == (5,) + header = cx_estim.header_to_text + data_to_text = cx_estim.data_to_text(tmp) + assert len(data_to_text.split()) == 5 + + numpy.testing.assert_allclose(re_estim.data, cx_estim.data) + + +@pytest.mark.unit +def test_number_estimator(): + # Test. + objs = build_generic_test_case_handlers( + nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, + mf_trial=mf_trial, complex_integrals=True, debug=debug, + seed=seed, verbose=verbose) + trial = objs['trial'] + hamiltonian = objs['hamiltonian'] + walkers = objs['walkers'] + + estim = ThermalNumberEstimator(hamiltonian=hamiltonian, trial=trial) + estim.compute_estimator(walkers, hamiltonian, trial) + assert len(estim.names) == 3 + assert estim["NavNumer"].real == pytest.approx(ne * nwalkers) + assert estim["Nav"] == pytest.approx(0.0) + tmp = estim.data.copy() + estim.post_reduce_hook(tmp) + assert tmp[estim.get_index("Nav")] == pytest.approx(ne) + assert estim.print_to_stdout + assert estim.ascii_filename == None + assert estim.shape == (3,) + header = estim.header_to_text + data_to_text = estim.data_to_text(tmp) + assert len(data_to_text.split()) == 3 + + +@pytest.mark.unit +def test_estimator_handler(): + with tempfile.NamedTemporaryFile() as tmp1, tempfile.NamedTemporaryFile() as tmp2: + # Test. + objs = build_generic_test_case_handlers( + nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, + mf_trial=mf_trial, complex_integrals=True, debug=debug, + seed=seed, verbose=verbose) + trial = objs['trial'] + hamiltonian = objs['hamiltonian'] + walkers = objs['walkers'] + + estim = ThermalEnergyEstimator(hamiltonian=hamiltonian, trial=trial, + filename=tmp1.name) + estim.print_to_stdout = False + + comm = MPI.COMM_WORLD + handler = ThermalEstimatorHandler( + comm, + hamiltonian, + trial, + observables=("energy",), + filename=tmp2.name) + handler["energy1"] = estim + handler.json_string = "" + handler.initialize(comm) + handler.compute_estimators(hamiltonian, trial, walkers) + + +if __name__ == "__main__": + test_energy_estimator() + test_number_estimator() + test_estimator_handler() + + + diff --git a/ipie/addons/thermal/estimators/tests/test_generic.py b/ipie/addons/thermal/estimators/tests/test_generic.py new file mode 100644 index 00000000..4205ab58 --- /dev/null +++ b/ipie/addons/thermal/estimators/tests/test_generic.py @@ -0,0 +1,91 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import numpy +import pytest +from typing import Tuple, Union + +try: + from ipie.addons.thermal.utils.legacy_testing import build_legacy_generic_test_case_handlers + _no_cython = False + +except ModuleNotFoundError: + _no_cython = True + +from ipie.config import MPI +from ipie.addons.thermal.estimators.thermal import one_rdm_from_G +from ipie.addons.thermal.estimators.generic import local_energy_generic_cholesky +from ipie.addons.thermal.utils.testing import build_generic_test_case_handlers + +from ipie.legacy.estimators.thermal import one_rdm_from_G as legacy_one_rdm_from_G +from ipie.legacy.estimators.generic import local_energy_generic_cholesky as legacy_local_energy_generic_cholesky + +comm = MPI.COMM_WORLD + +@pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") +@pytest.mark.unit +def test_local_energy_cholesky(mf_trial=False): + # System params. + nup = 5 + ndown = 5 + nelec = (nup, ndown) + nbasis = 10 + + # Thermal AFQMC params. + mu = -10. + beta = 0.1 + timestep = 0.01 + nwalkers = 12 + lowrank = False + + mf_trial = True + complex_integrals = False + debug = True + verbose = True + seed = 7 + numpy.random.seed(seed) + + # Test. + objs = build_generic_test_case_handlers( + nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, + mf_trial=mf_trial, complex_integrals=complex_integrals, debug=debug, + seed=seed, verbose=verbose) + trial = objs['trial'] + hamiltonian = objs['hamiltonian'] + P = one_rdm_from_G(trial.G) + eloc = local_energy_generic_cholesky(hamiltonian, P) + + # Legacy. + legacy_objs = build_legacy_generic_test_case_handlers( + hamiltonian, comm, nelec, mu, beta, timestep, nwalkers=nwalkers, + lowrank=lowrank, mf_trial=mf_trial, seed=seed, verbose=verbose) + legacy_system = legacy_objs['system'] + legacy_trial = legacy_objs['trial'] + legacy_hamiltonian = legacy_objs['hamiltonian'] + + legacy_P = legacy_one_rdm_from_G(legacy_trial.G) + legacy_eloc = legacy_local_energy_generic_cholesky( + legacy_system, legacy_hamiltonian, legacy_P) + + numpy.testing.assert_allclose(trial.G, legacy_trial.G, atol=1e-10) + numpy.testing.assert_allclose(P, legacy_P, atol=1e-10) + numpy.testing.assert_allclose(eloc, legacy_eloc, atol=1e-10) + + +if __name__ == '__main__': + test_local_energy_cholesky(mf_trial=True) diff --git a/ipie/addons/thermal/estimators/tests/test_generic_complex.py b/ipie/addons/thermal/estimators/tests/test_generic_complex.py new file mode 100644 index 00000000..2902f61b --- /dev/null +++ b/ipie/addons/thermal/estimators/tests/test_generic_complex.py @@ -0,0 +1,122 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import numpy +import pytest +from typing import Tuple, Union + +from ipie.hamiltonians.generic import Generic as HamGeneric +from ipie.hamiltonians.generic import GenericRealChol, GenericComplexChol +from ipie.addons.thermal.utils.testing import build_generic_test_case_handlers +from ipie.addons.thermal.estimators.generic import local_energy_generic_cholesky +from ipie.addons.thermal.estimators.thermal import one_rdm_from_G + +# System params. +nup = 5 +ndown = 5 +nelec = (nup, ndown) +nbasis = 10 + +# Thermal AFQMC params. +mu = -10. +beta = 0.1 +timestep = 0.01 +nwalkers = 12 +lowrank = False + +mf_trial = True +complex_integrals = False +debug = True +verbose = True +seed = 7 +numpy.random.seed(seed) + +@pytest.mark.unit +def test_local_energy_vs_real(): + # Test. + objs = build_generic_test_case_handlers( + nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, + mf_trial=mf_trial, complex_integrals=complex_integrals, debug=debug, + seed=seed, verbose=verbose) + trial = objs['trial'] + walkers = objs['walkers'] + hamiltonian = objs['hamiltonian'] + + assert isinstance(hamiltonian, GenericRealChol) + + chol = hamiltonian.chol + cx_chol = numpy.array(chol, dtype=numpy.complex128) + cx_hamiltonian = HamGeneric( + numpy.array(hamiltonian.H1, dtype=numpy.complex128), cx_chol, + hamiltonian.ecore, verbose=False) + + assert isinstance(cx_hamiltonian, GenericComplexChol) + + for iw in range(walkers.nwalkers): + P = one_rdm_from_G(numpy.array([walkers.Ga[iw], walkers.Gb[iw]])) + energy = local_energy_generic_cholesky(hamiltonian, P) + cx_energy = local_energy_generic_cholesky(cx_hamiltonian, P) + numpy.testing.assert_allclose(energy, cx_energy, atol=1e-10) + + +@pytest.mark.unit +def test_local_energy_vs_eri(): + # Test. + objs = build_generic_test_case_handlers( + nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, + mf_trial=mf_trial, debug=debug, complex_integrals=True, with_eri=True, + seed=seed, verbose=verbose) + trial = objs['trial'] + walkers = objs['walkers'] + hamiltonian = objs['hamiltonian'] + assert isinstance(hamiltonian, GenericComplexChol) + eri = objs['eri'].reshape(nbasis, nbasis, nbasis, nbasis) + + chol = hamiltonian.chol.copy() + nchol = chol.shape[1] + chol = chol.reshape(nbasis, nbasis, nchol) + + # Check if chol and eri are consistent. + eri_chol = numpy.einsum('mnx,slx->mnls', chol, chol.conj()) + numpy.testing.assert_allclose(eri, eri_chol, atol=1e-10) + + for iw in range(walkers.nwalkers): + P = one_rdm_from_G(numpy.array([walkers.Ga[iw], walkers.Gb[iw]])) + Pa, Pb = P + Ptot = Pa + Pb + etot, e1, e2 = local_energy_generic_cholesky(hamiltonian, P) + + # Test 1-body term. + h1e = hamiltonian.H1[0] + e1ref = numpy.einsum('ij,ij->', h1e, Ptot) + numpy.testing.assert_allclose(e1, e1ref, atol=1e-10) + + # Test 2-body term. + ecoul = 0.5 * numpy.einsum('ijkl,ij,kl->', eri, Ptot, Ptot) + exx = -0.5 * numpy.einsum('ijkl,il,kj->', eri, Pa, Pa) + exx -= 0.5 * numpy.einsum('ijkl,il,kj->', eri, Pb, Pb) + e2ref = ecoul + exx + numpy.testing.assert_allclose(e2, e2ref, atol=1e-10) + + etotref = e1ref + e2ref + numpy.testing.assert_allclose(etot, etotref, atol=1e-10) + + +if __name__ == '__main__': + test_local_energy_vs_real() + test_local_energy_vs_eri() diff --git a/ipie/addons/thermal/estimators/thermal.py b/ipie/addons/thermal/estimators/thermal.py new file mode 100644 index 00000000..198f965c --- /dev/null +++ b/ipie/addons/thermal/estimators/thermal.py @@ -0,0 +1,87 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import numpy +import scipy.linalg + +def one_rdm_from_G(G): + r"""Compute one-particle reduced density matrix from Green's function. + + .. math:: + rho_{ij} = \langle c_{i}^{\dagger} c_{j} \rangle \\ + = 1 - G_{ji} + Parameters + ---------- + G : :class:`numpy.ndarray` + Thermal Green's function. + + Returns + ------- + P : :class:`numpy.ndarray` + Thermal 1RDM. + """ + I = numpy.identity(G.shape[-1]) + return numpy.array([I - G[0].T, I - G[1].T], dtype=numpy.complex128) + +def one_rdm_stable(BT, num_slices): + nbasis = BT.shape[-1] + G = [] + for spin in [0, 1]: + # Need to construct the product A(l) = B_l B_{l-1}..B_L...B_{l+1} in + # stable way. Iteratively construct column pivoted QR decompositions + # (A = QDT) starting from the rightmost (product of) propagator(s). + (Q1, R1, P1) = scipy.linalg.qr(BT[spin], pivoting=True, check_finite=False) + # Form D matrices + D1 = numpy.diag(R1.diagonal()) + D1inv = numpy.diag(1.0 / R1.diagonal()) + T1 = numpy.einsum("ii,ij->ij", D1inv, R1) + # permute them + T1[:, P1] = T1[:, range(nbasis)] + + for i in range(0, num_slices - 1): + C2 = numpy.dot(numpy.dot(BT[spin], Q1), D1) + (Q1, R1, P1) = scipy.linalg.qr(C2, pivoting=True, check_finite=False) + # Compute D matrices + D1inv = numpy.diag(1.0 / R1.diagonal()) + D1 = numpy.diag(R1.diagonal()) + tmp = numpy.einsum("ii,ij->ij", D1inv, R1) + tmp[:, P1] = tmp[:, range(nbasis)] + T1 = numpy.dot(tmp, T1) + # G^{-1} = 1+A = 1+QDT = Q (Q^{-1}T^{-1}+D) T + # Write D = Db^{-1} Ds + # Then G^{-1} = Q Db^{-1}(Db Q^{-1}T^{-1}+Ds) T + Db = numpy.zeros(BT[spin].shape, BT[spin].dtype) + Ds = numpy.zeros(BT[spin].shape, BT[spin].dtype) + for i in range(Db.shape[0]): + absDlcr = abs(Db[i, i]) + if absDlcr > 1.0: + Db[i, i] = 1.0 / absDlcr + Ds[i, i] = numpy.sign(D1[i, i]) + else: + Db[i, i] = 1.0 + Ds[i, i] = D1[i, i] + + T1inv = scipy.linalg.inv(T1, check_finite=False) + # C = (Db Q^{-1}T^{-1}+Ds) + C = numpy.dot(numpy.einsum("ii,ij->ij", Db, Q1.conj().T), T1inv) + Ds + Cinv = scipy.linalg.inv(C, check_finite=False) + + # Then G = T^{-1} C^{-1} Db Q^{-1} + # Q is unitary. + G.append(numpy.dot(numpy.dot(T1inv, Cinv), numpy.einsum("ii,ij->ij", Db, Q1.conj().T))) + return one_rdm_from_G(numpy.array(G)) diff --git a/ipie/addons/thermal/propagation/__init__.py b/ipie/addons/thermal/propagation/__init__.py new file mode 100644 index 00000000..f91ef518 --- /dev/null +++ b/ipie/addons/thermal/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. +# +# Authors: Fionn Malone +# Joonho Lee +# diff --git a/ipie/addons/thermal/propagation/force_bias.py b/ipie/addons/thermal/propagation/force_bias.py new file mode 100644 index 00000000..47fd2138 --- /dev/null +++ b/ipie/addons/thermal/propagation/force_bias.py @@ -0,0 +1,72 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import plum +import numpy + +from ipie.hamiltonians.generic import GenericRealChol, GenericComplexChol +from ipie.addons.thermal.estimators.thermal import one_rdm_from_G +from ipie.utils.backend import arraylib as xp + +@plum.dispatch +def construct_force_bias(hamiltonian: GenericRealChol, walkers): + r"""Compute optimal force bias. + + Parameters + ---------- + G: :class:`numpy.ndarray` + Walker's 1RDM: . + + Returns + ------- + xbar : :class:`numpy.ndarray` + Force bias. + """ + vbias = xp.empty((walkers.nwalkers, hamiltonian.nchol), dtype=walkers.Ga.dtype) + + for iw in range(walkers.nwalkers): + P = one_rdm_from_G(numpy.array([walkers.Ga[iw], walkers.Gb[iw]])) + vbias[iw] = hamiltonian.chol.T.dot(P[0].ravel()) + hamiltonian.chol.T.dot(P[1].ravel()) + + return vbias + + +@plum.dispatch +def construct_force_bias(hamiltonian: GenericComplexChol, walkers): + r"""Compute optimal force bias. + + Parameters + ---------- + G: :class:`numpy.ndarray` + Walker's 1RDM: . + + Returns + ------- + xbar : :class:`numpy.ndarray` + Force bias. + """ + nchol = hamiltonian.nchol + vbias = xp.empty((walkers.nwalkers, hamiltonian.nfields), dtype=walkers.Ga.dtype) + + for iw in range(walkers.nwalkers): + P = one_rdm_from_G(numpy.array([walkers.Ga[iw], walkers.Gb[iw]])) + vbias[iw, :nchol] = hamiltonian.A.T.dot(P[0].ravel()) + hamiltonian.A.T.dot(P[1].ravel()) + vbias[iw, nchol:] = hamiltonian.B.T.dot(P[0].ravel()) + hamiltonian.B.T.dot(P[1].ravel()) + + return vbias + diff --git a/ipie/addons/thermal/propagation/operations.py b/ipie/addons/thermal/propagation/operations.py new file mode 100644 index 00000000..661f96b1 --- /dev/null +++ b/ipie/addons/thermal/propagation/operations.py @@ -0,0 +1,45 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +from ipie.utils.backend import arraylib as xp + +def apply_exponential(VHS, exp_nmax): + """Apply exponential propagator of the HS transformation + + Parameters + ---------- + phi : numpy array + a state + VHS : numpy array + HS transformation potential + + Returns + ------- + phi : numpy array + Exp(VHS) * phi + """ + # Temporary array for matrix exponentiation. + phi = xp.identity(VHS.shape[-1], dtype=xp.complex128) + Temp = xp.zeros(phi.shape, dtype=phi.dtype) + xp.copyto(Temp, phi) + + for n in range(1, exp_nmax + 1): + Temp = VHS.dot(Temp) / n + phi += Temp + + return phi # Shape (nbasis, nbasis). diff --git a/ipie/addons/thermal/propagation/phaseless_base.py b/ipie/addons/thermal/propagation/phaseless_base.py new file mode 100644 index 00000000..c5fac4dc --- /dev/null +++ b/ipie/addons/thermal/propagation/phaseless_base.py @@ -0,0 +1,348 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import time +import plum +import math +import cmath +import numpy +import scipy.linalg + +from abc import abstractmethod +from ipie.utils.backend import arraylib as xp +from ipie.propagation.continuous_base import ContinuousBase +from ipie.propagation.operations import apply_exponential +from ipie.hamiltonians.generic import GenericRealChol, GenericComplexChol + +from ipie.addons.thermal.estimators.thermal import one_rdm_from_G +from ipie.addons.thermal.propagation.force_bias import construct_force_bias + +# TODO: Add lowrank implementation. See: https://github.com/JoonhoLee-Group/ipie/issues/302 +# Ref: 10.1103/PhysRevB.80.214116 for bounds. + +@plum.dispatch +def construct_mean_field_shift(hamiltonian: GenericRealChol, trial): + r"""Compute mean field shift. + + .. math:: + + \bar{v}_n = \sum_{ik\sigma} v_{(ik),n} P_{ik\sigma} + + """ + # hamiltonian.chol has shape (nbasis^2, nchol). + P = one_rdm_from_G(trial.G) + P = (P[0] + P[1]).ravel() + tmp_real = numpy.dot(hamiltonian.chol.T, P.real) + tmp_imag = numpy.dot(hamiltonian.chol.T, P.imag) + mf_shift = 1.0j * tmp_real - tmp_imag + return mf_shift # Shape (nchol,). + + +@plum.dispatch +def construct_mean_field_shift(hamiltonian: GenericComplexChol, trial): + r"""Compute mean field shift. + + .. math:: + + \bar{v}_n = \sum_{ik\sigma} v_{(ik),n} P_{ik\sigma} + + """ + # hamiltonian.chol has shape (nbasis^2, nchol). + P = one_rdm_from_G(trial.G) + P = (P[0] + P[1]).ravel() + nchol = hamiltonian.nchol + mf_shift = numpy.zeros(hamiltonian.nfields, dtype=hamiltonian.chol.dtype) + mf_shift[:nchol] = 1j * numpy.dot(hamiltonian.A.T, P.ravel()) + mf_shift[nchol:] = 1j * numpy.dot(hamiltonian.B.T, P.ravel()) + return mf_shift # Shape (nchol,). + + +class PhaselessBase(ContinuousBase): + """A base class for generic continuous HS transform FT-AFQMC propagators.""" + + def __init__(self, timestep, mu, lowrank=False, exp_nmax=6, verbose=False): + super().__init__(timestep, verbose=verbose) + self.mu = mu + self.sqrt_dt = self.dt**0.5 + self.isqrt_dt = 1j * self.sqrt_dt + + self.nfb_trig = 0 # number of force bias triggered + self.ebound = (2.0 / self.dt) ** 0.5 # energy bound range + self.fbbound = 1.0 + self.mpi_handler = None + self.lowrank = lowrank + self.exp_nmax = exp_nmax + + + def build(self, hamiltonian, trial=None, walkers=None, mpi_handler=None, verbose=False): + # dt/2 one-body propagator + start = time.time() + self.mf_shift = construct_mean_field_shift(hamiltonian, trial) + + if verbose: + print(f"# Time to mean field shift: {time.time() - start} s") + print( + "# Absolute value of maximum component of mean field shift: " + "{:13.8e}.".format(numpy.max(numpy.abs(self.mf_shift))) + ) + + # Construct one-body propagator + self.BH1 = self.construct_one_body_propagator(hamiltonian) + + # Allocate force bias (we don't need to do this here - it will be allocated when it is needed) + self.vbias = None + + # Legacy attributes. + self.mf_core = hamiltonian.ecore + 0.5 * numpy.dot(self.mf_shift, self.mf_shift) + self.mf_const_fac = cmath.exp(-self.dt * self.mf_core) + + + @plum.dispatch + def construct_one_body_propagator(self, hamiltonian: GenericRealChol): + r"""Construct mean-field shifted one-body propagator. + + .. math:: + + H1 \rightarrow H1 - v0 + v0_{ik} = \sum_n v_{(ik),n} \bar{v}_n + + Parameters + ---------- + hamiltonian : hamiltonian class + Generic hamiltonian object. + dt : float + Timestep. + """ + nb = hamiltonian.nbasis + shift = 1j * numpy.einsum("mx,x->m", hamiltonian.chol, self.mf_shift).reshape(nb, nb) + muN = self.mu * numpy.identity(nb, dtype=hamiltonian.H1.dtype) + H1 = hamiltonian.h1e_mod - numpy.array([shift + muN, shift + muN]) + expH1 = numpy.array([ + scipy.linalg.expm(-0.5 * self.dt * H1[0]), + scipy.linalg.expm(-0.5 * self.dt * H1[1])]) + return expH1 # Shape (nbasis, nbasis). + + + @plum.dispatch + def construct_one_body_propagator(self, hamiltonian: GenericComplexChol): + r"""Construct mean-field shifted one-body propagator. + + .. math:: + + H1 \rightarrow H1 - v0 + v0_{ik} = \sum_n v_{(ik),n} \bar{v}_n + + Parameters + ---------- + hamiltonian : hamiltonian class + Generic hamiltonian object. + dt : float + Timestep. + """ + nb = hamiltonian.nbasis + nchol = hamiltonian.nchol + shift = xp.zeros((nb, nb), dtype=hamiltonian.chol.dtype) + shift = 1j * numpy.einsum("mx,x->m", hamiltonian.A, self.mf_shift[:nchol]).reshape(nb, nb) + shift += 1j * numpy.einsum("mx,x->m", hamiltonian.B, self.mf_shift[nchol:]).reshape(nb, nb) + muN = self.mu * numpy.identity(nb, dtype=hamiltonian.H1.dtype) + H1 = hamiltonian.h1e_mod - numpy.array([shift + muN, shift + muN]) + expH1 = numpy.array([ + scipy.linalg.expm(-0.5 * self.dt * H1[0]), + scipy.linalg.expm(-0.5 * self.dt * H1[1])]) + return expH1 # Shape (nbasis, nbasis). + + + def construct_two_body_propagator(self, walkers, hamiltonian, trial, debug=False): + r"""Construct two-body propagator. + + .. math:: + \bar{x}_n &= \sqrt{\Delta\tau} \bar{v}_n \\ + x_\mathrm{shifted}_n &= x_n - \bar{x}_n \\ + C_{MF} &= -\sqrt{\Delta\tau} \sum_{n} x_\mathrm{shifted}_n \bar{v}_n \\ + &= -\sqrt{\Delta\tau} \sum_{n} (x_n - \sqrt{\Delta\tau} \bar{v}_n) \bar{v}_n \\ + &= -\sqrt{\Delta\tau} \sum_{n} x_n \bar{v}_n + \Delta\tau \sum_{n} \bar{v}_n^2. + + Parameters + ---------- + walkers: walker class + UHFThermalWalkers object. + hamiltonian : hamiltonian class + Generic hamiltonian object. + trial : trial class + Trial dnsity matrix. + """ + # Optimal force bias + xbar = xp.zeros((walkers.nwalkers, hamiltonian.nfields)) + start_time = time.time() + self.vbias = construct_force_bias(hamiltonian, walkers) + xbar = -self.sqrt_dt * (1j * self.vbias - self.mf_shift) + self.timer.tfbias += time.time() - start_time + + # Force bias bounding + xbar = self.apply_bound_force_bias(xbar, self.fbbound) + + # Normally distrubted auxiliary fields. + xi = xp.random.normal(0.0, 1.0, hamiltonian.nfields * walkers.nwalkers).reshape( + walkers.nwalkers, hamiltonian.nfields) + + if debug: self.xi = xi # For debugging. + xshifted = xi - xbar # Shape (nwalkers, nfields). + + # Constant factor arising from force bias and mean field shift + cmf = -self.sqrt_dt * xp.einsum("wx,x->w", xshifted, self.mf_shift) # Shape (nwalkers,). + # Constant factor arising from shifting the propability distribution. + cfb = xp.einsum("wx,wx->w", xi, xbar) - 0.5 * xp.einsum("wx,wx->w", xbar, xbar) # Shape (nwalkers,). + + xshifted = xshifted.T.copy() # Shape (nfields, nwalkers). + VHS = self.construct_VHS(hamiltonian, xshifted) # Shape (nwalkers, nbasis, nbasis). + return cmf, cfb, xshifted, VHS + + def propagate_walkers_one_body(self, walkers): + pass + + def propagate_walkers_two_body(self, walkers, hamiltonian, trial): + pass + + def propagate_walkers(self, walkers, hamiltonian, trial, eshift=0., debug=False): + start_time = time.time() + cmf, cfb, xshifted, VHS = self.construct_two_body_propagator( + walkers, hamiltonian, trial, debug=debug) + assert walkers.nwalkers == xshifted.shape[-1] + self.timer.tvhs += time.time() - start_time + assert len(VHS.shape) == 3 + + start_time = time.time() + for iw in range(walkers.nwalkers): + stack = walkers.stack[iw] + phi = xp.identity(VHS[iw].shape[-1], dtype=xp.complex128) + BV = apply_exponential(phi, VHS[iw], self.exp_nmax) # Shape (nbasis, nbasis). + B = numpy.array([BV.dot(self.BH1[0]), BV.dot(self.BH1[1])]) + B = numpy.array([self.BH1[0].dot(B[0]), self.BH1[1].dot(B[1])]) + + # Compute determinant ratio det(1+A')/det(1+A). + # 1. Current walker's Green's function. + tix = stack.nslice + start_time = time.time() + G = walkers.calc_greens_function(iw, slice_ix=tix, inplace=False) + self.timer.tgf += time.time() - start_time + + # 2. Compute updated Green's function. + start_time = time.time() + stack.update_new(B) + walkers.calc_greens_function(iw, slice_ix=tix, inplace=True) + + # 3. Compute det(G/G') + # Now apply phaseless approximation. + # Use legacy thermal weight update for now. + self.update_weight_legacy(walkers, iw, G, cfb, cmf, eshift) + #self.update_weight(walkers, iw, G, cfb, cmf, eshift) + + self.timer.tupdate += time.time() - start_time + + + def update_weight(self, walkers, iw, G, cfb, cmf, eshift): + """Update weight for walker `iw`. + """ + M0a = scipy.linalg.det(G[0], check_finite=False) + M0b = scipy.linalg.det(G[1], check_finite=False) + Mnewa = scipy.linalg.det(walkers.Ga[iw], check_finite=False) + Mnewb = scipy.linalg.det(walkers.Gb[iw], check_finite=False) + + # ovlp = det( G^{-1} ) + ovlp_ratio = (M0a * M0b) / (Mnewa * Mnewb) # ovlp_new / ovlp_old + hybrid_energy = -(xp.log(ovlp_ratio) + cfb[iw] + cmf[iw]) / self.dt # Scalar. + hybrid_energy = self.apply_bound_hybrid(hybrid_energy, eshift) + importance_function = xp.exp( + -self.dt * (0.5 * (hybrid_energy + walkers.hybrid_energy) - eshift)) + + # Splitting w_k = |I(x, \bar{x}, |phi_k>)| e^{i theta_k}, where `k` + # labels the time slice. + magn = xp.abs(importance_function) + walkers.hybrid_energy = hybrid_energy + dtheta = (-self.dt * hybrid_energy - cfb[iw]).imag # Scalar. + cosine_fac = xp.amax([0., xp.cos(dtheta)]) + walkers.weight[iw] *= magn * cosine_fac + walkers.M0a[iw] = Mnewa + walkers.M0b[iw] = Mnewb + + + def update_weight_legacy(self, walkers, iw, G, cfb, cmf, eshift): + """Update weight for walker `iw` using legacy code. + """ + #M0a = walkers.M0a[iw] + #M0b = walkers.M0b[iw] + M0a = scipy.linalg.det(G[0], check_finite=False) + M0b = scipy.linalg.det(G[1], check_finite=False) + Mnewa = scipy.linalg.det(walkers.Ga[iw], check_finite=False) + Mnewb = scipy.linalg.det(walkers.Gb[iw], check_finite=False) + _cfb = cfb[iw] + _cmf = cmf[iw] + + try: + # Could save M0 rather than recompute. + oratio = (M0a * M0b) / (Mnewa * Mnewb) + # Might want to cap this at some point. + hybrid_energy = cmath.log(oratio) + _cfb + _cmf + Q = cmath.exp(hybrid_energy) + expQ = self.mf_const_fac * Q + (magn, _) = cmath.polar(expQ) + + if not math.isinf(magn): + # Determine cosine phase from Arg(det(1+A'(x))/det(1+A(x))). + # Note this doesn't include exponential factor from shifting + # proability distribution. + dtheta = cmath.phase(cmath.exp(hybrid_energy - _cfb)) + cosine_fac = max(0, math.cos(dtheta)) + walkers.weight[iw] *= magn * cosine_fac + walkers.M0a[iw] = Mnewa + walkers.M0b[iw] = Mnewb + + else: + walkers.weight[iw] = 0. + + except ZeroDivisionError: + walkers.weight[iw] = 0. + + def apply_bound_force_bias(self, xbar, max_bound=1.0): + absxbar = xp.abs(xbar) + idx_to_rescale = absxbar > max_bound + nonzeros = absxbar > 1e-13 + xbar_rescaled = xbar.copy() + xbar_rescaled[nonzeros] = xbar_rescaled[nonzeros] / absxbar[nonzeros] + xbar = xp.where(idx_to_rescale, xbar_rescaled, xbar) + self.nfb_trig += xp.sum(idx_to_rescale) + return xbar + + + def apply_bound_hybrid(self, ehyb, eshift): # Shift is a number but ehyb is not + # For initial steps until first estimator communication, `eshift` will be + # zero and hybrid energy can be incorrect. So just avoid capping for + # first block until reasonable estimate of `eshift` can be computed. + if abs(eshift) < 1e-10: + return ehyb + + emax = eshift.real + self.ebound + emin = eshift.real - self.ebound + return xp.minimum(emax, xp.maximum(ehyb, emin)) + + + # Form VHS. + @abstractmethod + def construct_VHS(self, hamiltonian, xshifted): + pass + diff --git a/ipie/addons/thermal/propagation/phaseless_generic.py b/ipie/addons/thermal/propagation/phaseless_generic.py new file mode 100644 index 00000000..6f196ed8 --- /dev/null +++ b/ipie/addons/thermal/propagation/phaseless_generic.py @@ -0,0 +1,52 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import plum + +from ipie.hamiltonians.generic import GenericRealChol, GenericComplexChol +from ipie.addons.thermal.propagation.phaseless_base import PhaselessBase +from ipie.utils.backend import arraylib as xp + + +class PhaselessGeneric(PhaselessBase): + """A class for performing phaseless propagation with real, generic, hamiltonian.""" + + def __init__(self, time_step, mu, exp_nmax=6, lowrank=False, verbose=False): + super().__init__(time_step, mu, lowrank=lowrank, exp_nmax=exp_nmax, verbose=verbose) + + @plum.dispatch + def construct_VHS(self, hamiltonian: GenericRealChol, xshifted: xp.ndarray): + """Includes `nwalkers`. + """ + nwalkers = xshifted.shape[-1] # Shape (nfields, nwalkers). + VHS = hamiltonian.chol.dot(xshifted) # Shape (nbasis^2, nwalkers). + VHS = self.isqrt_dt * VHS.T.reshape(nwalkers, hamiltonian.nbasis, hamiltonian.nbasis) + return VHS # Shape (nwalkers, nbasis, nbasis). + + @plum.dispatch + def construct_VHS(self, hamiltonian: GenericComplexChol, xshifted: xp.ndarray): + """Includes `nwalkers`. + """ + nwalkers = xshifted.shape[-1] + nchol = hamiltonian.nchol + VHS = self.isqrt_dt * ( + hamiltonian.A.dot(xshifted[:nchol]) + hamiltonian.B.dot(xshifted[nchol:]) + ) + VHS = VHS.T.copy() + VHS = VHS.reshape(nwalkers, hamiltonian.nbasis, hamiltonian.nbasis) + return VHS diff --git a/ipie/addons/thermal/propagation/propagator.py b/ipie/addons/thermal/propagation/propagator.py new file mode 100644 index 00000000..ab9b4600 --- /dev/null +++ b/ipie/addons/thermal/propagation/propagator.py @@ -0,0 +1,22 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +from ipie.hamiltonians.generic import GenericRealChol, GenericComplexChol +from ipie.addons.thermal.propagation.phaseless_generic import PhaselessGeneric + +Propagator = {GenericRealChol: PhaselessGeneric, GenericComplexChol: PhaselessGeneric} diff --git a/ipie/addons/thermal/propagation/tests/__init__.py b/ipie/addons/thermal/propagation/tests/__init__.py new file mode 100644 index 00000000..f91ef518 --- /dev/null +++ b/ipie/addons/thermal/propagation/tests/__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. +# +# Authors: Fionn Malone +# Joonho Lee +# diff --git a/ipie/addons/thermal/propagation/tests/test_prop_generic.py b/ipie/addons/thermal/propagation/tests/test_prop_generic.py new file mode 100644 index 00000000..c2e7377d --- /dev/null +++ b/ipie/addons/thermal/propagation/tests/test_prop_generic.py @@ -0,0 +1,204 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import numpy +import pytest + +try: + from ipie.addons.thermal.utils.legacy_testing import build_legacy_generic_test_case_handlers + from ipie.addons.thermal.utils.legacy_testing import legacy_propagate_walkers + _no_cython = False + +except ModuleNotFoundError: + _no_cython = True + +from ipie.config import MPI +from ipie.addons.thermal.estimators.generic import local_energy_generic_cholesky +from ipie.addons.thermal.estimators.thermal import one_rdm_from_G +from ipie.addons.thermal.utils.testing import build_generic_test_case_handlers + +from ipie.legacy.estimators.generic import local_energy_generic_cholesky as legacy_local_energy_generic_cholesky +from ipie.legacy.estimators.thermal import one_rdm_from_G as legacy_one_rdm_from_G + +comm = MPI.COMM_WORLD + +# System params. +nup = 5 +ndown = 5 +nelec = (nup, ndown) +nbasis = 10 + +# Thermal AFQMC params. +mu = -10. +beta = 0.1 +timestep = 0.01 +nwalkers = 12 +nblocks = 12 +lowrank = False + +mf_trial = True +complex_integrals = False +debug = True +verbose = True +seed = 7 +numpy.random.seed(seed) + + +@pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") +@pytest.mark.unit +def test_mf_shift(): + # Test. + objs = build_generic_test_case_handlers( + nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, + mf_trial=mf_trial, complex_integrals=complex_integrals, debug=debug, + seed=seed, verbose=verbose) + hamiltonian = objs['hamiltonian'] + propagator = objs['propagator'] + + # Legacy. + legacy_objs = build_legacy_generic_test_case_handlers( + hamiltonian, comm, nelec, mu, beta, timestep, + nwalkers=nwalkers, lowrank=lowrank, mf_trial=mf_trial, + seed=seed, verbose=verbose) + legacy_propagator = legacy_objs['propagator'] + + numpy.testing.assert_almost_equal(legacy_propagator.propagator.mf_shift, + propagator.mf_shift, decimal=10) + + +@pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") +@pytest.mark.unit +def test_BH1(): + # Test. + objs = build_generic_test_case_handlers( + nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, + mf_trial=mf_trial, complex_integrals=complex_integrals, debug=debug, + seed=seed, verbose=verbose) + hamiltonian = objs['hamiltonian'] + propagator = objs['propagator'] + + # Legacy. + legacy_objs = build_legacy_generic_test_case_handlers( + hamiltonian, comm, nelec, mu, beta, timestep, + nwalkers=nwalkers, lowrank=lowrank, mf_trial=mf_trial, + seed=seed, verbose=verbose) + legacy_propagator = legacy_objs['propagator'] + + numpy.testing.assert_almost_equal(legacy_propagator.propagator.BH1, + propagator.BH1, decimal=10) + + +@pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") +@pytest.mark.unit +def test_construct_two_body_propagator(): + # Test. + objs = build_generic_test_case_handlers( + nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, + mf_trial=mf_trial, complex_integrals=complex_integrals, debug=debug, + seed=seed, verbose=verbose) + trial = objs['trial'] + hamiltonian = objs['hamiltonian'] + walkers = objs['walkers'] + propagator = objs['propagator'] + + # Legacy. + legacy_objs = build_legacy_generic_test_case_handlers( + hamiltonian, comm, nelec, mu, beta, timestep, + nwalkers=nwalkers, lowrank=lowrank, mf_trial=mf_trial, + seed=seed, verbose=verbose) + legacy_trial = legacy_objs['trial'] + legacy_hamiltonian = legacy_objs['hamiltonian'] + legacy_walkers = legacy_objs['walkers'] + legacy_propagator = legacy_objs['propagator'] + + cmf, cfb, xshifted, VHS = propagator.construct_two_body_propagator( + walkers, hamiltonian, trial, debug=True) + + legacy_cmf = [] + legacy_cfb = [] + legacy_xshifted = [] + legacy_VHS = [] + + for iw in range(walkers.nwalkers): + _cmf, _cfb, _xshifted, _VHS = legacy_propagator.two_body_propagator( + legacy_walkers.walkers[iw], legacy_hamiltonian, + legacy_trial, xi=propagator.xi[iw]) + legacy_cmf.append(_cmf) + legacy_cfb.append(_cfb) + legacy_xshifted.append(_xshifted) + legacy_VHS.append(_VHS) + + legacy_xshifted = numpy.array(legacy_xshifted).T + + numpy.testing.assert_almost_equal(legacy_cmf, cmf, decimal=10) + numpy.testing.assert_almost_equal(legacy_cfb, cfb, decimal=10) + numpy.testing.assert_almost_equal(legacy_xshifted, xshifted, decimal=10) + numpy.testing.assert_almost_equal(legacy_VHS, VHS, decimal=10) + + +@pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") +@pytest.mark.unit +def test_phaseless_generic_propagator(): + # Test. + objs = build_generic_test_case_handlers( + nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, + mf_trial=mf_trial, complex_integrals=complex_integrals, debug=debug, + seed=seed, verbose=verbose) + trial = objs['trial'] + hamiltonian = objs['hamiltonian'] + walkers = objs['walkers'] + propagator = objs['propagator'] + + # Legacy. + legacy_objs = build_legacy_generic_test_case_handlers( + hamiltonian, comm, nelec, mu, beta, timestep, + nwalkers=nwalkers, lowrank=lowrank, mf_trial=mf_trial, + seed=seed, verbose=verbose) + legacy_system = legacy_objs['system'] + legacy_trial = legacy_objs['trial'] + legacy_hamiltonian = legacy_objs['hamiltonian'] + legacy_walkers = legacy_objs['walkers'] + legacy_propagator = legacy_objs['propagator'] + + for t in range(walkers.stack[0].nslice): + for iw in range(walkers.nwalkers): + P = one_rdm_from_G(numpy.array([walkers.Ga[iw], walkers.Gb[iw]])) + eloc = local_energy_generic_cholesky(hamiltonian, P) + + legacy_P = legacy_one_rdm_from_G(numpy.array(legacy_walkers.walkers[iw].G)) + legacy_eloc = legacy_local_energy_generic_cholesky( + legacy_system, legacy_hamiltonian, legacy_P) + + numpy.testing.assert_almost_equal(legacy_eloc, eloc, decimal=10) + numpy.testing.assert_allclose(legacy_walkers.walkers[iw].G[0], walkers.Ga[iw]) + numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].G[1], walkers.Gb[iw], decimal=10) + numpy.testing.assert_almost_equal(legacy_P, P, decimal=10) + numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].stack.ovlp[0], walkers.stack[iw].ovlp[0], decimal=10) + numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].stack.ovlp[1], walkers.stack[iw].ovlp[1], decimal=10) + + propagator.propagate_walkers(walkers, hamiltonian, trial, debug=True) + legacy_walkers = legacy_propagate_walkers( + legacy_hamiltonian, legacy_trial, legacy_walkers, + legacy_propagator, xi=propagator.xi) + + +if __name__ == "__main__": + test_mf_shift() + test_BH1() + test_construct_two_body_propagator() + test_phaseless_generic_propagator() diff --git a/ipie/addons/thermal/propagation/tests/ueg/test_prop_ueg.py b/ipie/addons/thermal/propagation/tests/ueg/test_prop_ueg.py new file mode 100644 index 00000000..c7a788c5 --- /dev/null +++ b/ipie/addons/thermal/propagation/tests/ueg/test_prop_ueg.py @@ -0,0 +1,125 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import pytest +import numpy + +try: + from ipie.addons.thermal.utils.legacy_testing import build_legacy_ueg_test_case_handlers + from ipie.addons.thermal.utils.legacy_testing import legacy_propagate_walkers + from ipie.legacy.estimators.ueg import local_energy_ueg as legacy_local_energy_ueg + _no_cython = False + +except ModuleNotFoundError: + _no_cython = True + +from ipie.config import MPI +from ipie.addons.thermal.estimators.generic import local_energy_generic_cholesky +from ipie.addons.thermal.estimators.thermal import one_rdm_from_G +from ipie.addons.thermal.utils.testing import build_ueg_test_case_handlers + +from ipie.legacy.estimators.thermal import one_rdm_from_G as legacy_one_rdm_from_G + +comm = MPI.COMM_WORLD + + +@pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") +@pytest.mark.unit +def test_phaseless_ueg_propagator(): + # UEG params. + nup = 7 + ndown = 7 + nelec = (nup, ndown) + rs = 1. + ecut = 1. + + # Thermal AFQMC params. + mu = -1. + beta = 0.1 + timestep = 0.01 + nwalkers = 1 + lowrank = False + + debug = True + verbose = False if (comm.rank != 0) else True + seed = 7 + numpy.random.seed(seed) + + # Test. + objs = build_ueg_test_case_handlers( + nelec, rs, ecut, mu, beta, timestep, nwalkers=nwalkers, + lowrank=lowrank, debug=debug, seed=seed, verbose=verbose) + trial = objs['trial'] + hamiltonian = objs['hamiltonian'] + walkers = objs['walkers'] + propagator = objs['propagator'] + + # Legacy. + legacy_objs = build_legacy_ueg_test_case_handlers( + comm, nelec, rs, ecut, mu, beta, timestep, nwalkers=nwalkers, + lowrank=lowrank, seed=seed, verbose=verbose) + legacy_system = legacy_objs['system'] + legacy_trial = legacy_objs['trial'] + legacy_hamiltonian = legacy_objs['hamiltonian'] + legacy_walkers = legacy_objs['walkers'] + legacy_propagator = legacy_objs['propagator'] + + h1e = legacy_hamiltonian.H1[0] + eri = legacy_hamiltonian.eri_4() + + for t in range(walkers.stack[0].nslice): + for iw in range(walkers.nwalkers): + P = one_rdm_from_G(numpy.array([walkers.Ga[iw], walkers.Gb[iw]])) + eloc = local_energy_generic_cholesky(hamiltonian, P) + + legacy_P = legacy_one_rdm_from_G(numpy.array(legacy_walkers.walkers[iw].G)) + legacy_eloc = legacy_local_energy_ueg(legacy_system, legacy_hamiltonian, legacy_P) + + legacy_Pa, legacy_Pb = legacy_P + legacy_Ptot = legacy_Pa + legacy_Pb + ref_e1 = numpy.einsum('ij,ij->', h1e, legacy_Ptot) + + Ptot = legacy_Ptot + Pa = legacy_Pa + Pb = legacy_Pb + + ecoul = 0.5 * numpy.einsum('ijkl,ij,kl->', eri, Ptot, Ptot) + exx = -0.5 * numpy.einsum('ijkl,il,kj->', eri, Pa, Pa) + exx -= 0.5 * numpy.einsum('ijkl,il,kj->', eri, Pb, Pb) + ref_e2 = ecoul + exx + ref_eloc = (ref_e1 + ref_e2, ref_e1, ref_e2) + + numpy.testing.assert_almost_equal(legacy_P, P, decimal=10) + numpy.testing.assert_almost_equal(legacy_trial.dmat, trial.dmat, decimal=10) + numpy.testing.assert_allclose(eloc, ref_eloc, atol=1e-10) + numpy.testing.assert_allclose(legacy_eloc, ref_eloc, atol=1e-10) + numpy.testing.assert_almost_equal(legacy_eloc, eloc, decimal=10) + + numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].G[0], walkers.Ga[iw], decimal=10) + numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].G[1], walkers.Gb[iw], decimal=10) + numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].stack.ovlp[0], walkers.stack[iw].ovlp[0], decimal=10) + numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].stack.ovlp[1], walkers.stack[iw].ovlp[1], decimal=10) + + propagator.propagate_walkers(walkers, hamiltonian, trial, debug=True) + legacy_walkers = legacy_propagate_walkers( + legacy_hamiltonian, legacy_trial, legacy_walkers, + legacy_propagator, xi=propagator.xi) + + +if __name__ == "__main__": + test_phaseless_ueg_propagator() diff --git a/ipie/addons/thermal/qmc/__init__.py b/ipie/addons/thermal/qmc/__init__.py new file mode 100644 index 00000000..f91ef518 --- /dev/null +++ b/ipie/addons/thermal/qmc/__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. +# +# Authors: Fionn Malone +# Joonho Lee +# diff --git a/ipie/addons/thermal/qmc/calc.py b/ipie/addons/thermal/qmc/calc.py new file mode 100644 index 00000000..659ce515 --- /dev/null +++ b/ipie/addons/thermal/qmc/calc.py @@ -0,0 +1,149 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +"""Helper Routines for setting up a calculation""" + +from ipie.config import MPI +from ipie.utils.mpi import MPIHandler +from ipie.utils.io import get_input_value + +from ipie.systems.utils import get_system +from ipie.hamiltonians.utils import get_hamiltonian + +from ipie.addons.thermal.trial.utils import get_trial_density_matrix +from ipie.addons.thermal.walkers.uhf_walkers import UHFThermalWalkers +from ipie.addons.thermal.propagation.propagator import Propagator +from ipie.addons.thermal.qmc.options import ThermalQMCOpts, ThermalQMCParams +from ipie.addons.thermal.qmc.thermal_afqmc import ThermalAFQMC + + +def get_driver(options: dict, comm: MPI.COMM_WORLD) -> ThermalAFQMC: + verbosity = options.get("verbosity", 1) + qmc_opts = get_input_value(options, "qmc", default={}, alias=["qmc_options"]) + + sys_opts = get_input_value( + options, "system", default={}, alias=["model"], verbose=verbosity > 1 + ) + ham_opts = get_input_value(options, "hamiltonian", default={}, verbose=verbosity > 1) + # backward compatibility with previous code (to be removed) + for item in sys_opts.items(): + if item[0].lower() == "name" and "name" in ham_opts.keys(): + continue + ham_opts[item[0]] = item[1] + + tdm_opts = get_input_value( + options, "trial", default={}, alias=["trial_density_matrix"], verbose=verbosity > 1 + ) + + wlk_opts = get_input_value( + options, "walkers", default={}, alias=["walker", "walker_opts"], verbose=verbosity > 1 + ) + + if comm.rank != 0: + verbosity = 0 + lowrank = get_input_value(wlk_opts, "lowrank", default=False, alias=["low_rank"], verbose=verbosity) + batched = get_input_value(qmc_opts, "batched", default=False, verbose=verbosity) + debug = get_input_value(qmc_opts, "debug", default=False, verbose=verbosity) + + if (lowrank == True) or (batched == True): + raise ValueError("Option not supported in thermal code.") + else: + qmc = ThermalQMCOpts(qmc_opts, verbose=0) + mpi_handler = MPIHandler(nmembers=qmc_opts.get("nmembers", 1), verbose=verbosity) + system = get_system( + sys_opts, verbose=verbosity, comm=comm + ) # Have to deal with shared comm in the future. I think we will remove this... + ham_file = get_input_value(ham_opts, "integrals", None, verbose=verbosity) + if ham_file is None: + raise ValueError("Hamiltonian filename not specified.") + pack_chol = get_input_value( + ham_opts, "symmetry", True, alias=["pack_chol", "pack_cholesky"], verbose=verbosity + ) + hamiltonian = get_hamiltonian( + ham_file, mpi_handler.scomm, pack_chol=pack_chol, verbose=verbosity + ) + num_elec = (system.nup, system.ndown) + trial = get_trial_density_matrix( + hamiltonian, + num_elec, + qmc.beta, + qmc.dt, + options=tdm_opts, + comm=comm, + verbose=verbosity, + ) + stack_size = get_input_value(wlk_opts, 'stack_size', default=10, verbose=verbosity) + lowrank_thresh = get_input_value(wlk_opts, 'lowrank_thresh', default=1e-6, alias=["low_rank_thresh"], verbose=verbosity) + walkers = UHFThermalWalkers( + trial, hamiltonian.nbasis, qmc.nwalkers, stack_size=stack_size, + lowrank=lowrank, lowrank_thresh=lowrank_thresh, verbose=verbosity) + + if (comm.rank == 0) and (qmc.nsteps > 1): + print("Only num_steps_per_block = 1 allowed in thermal code! Resetting to value of 1.") + + # pylint: disable = no-value-for-parameter + params = ThermalQMCParams( + mu=qmc.mu, + beta=qmc.beta, + num_walkers=qmc.nwalkers, + total_num_walkers=qmc.nwalkers * comm.size, + num_blocks=qmc.nblocks, + timestep=qmc.dt, + num_stblz=qmc.nstblz, + pop_control_freq=qmc.npop_control, + pop_control_method=qmc.pop_control_method, + rng_seed=qmc.rng_seed, + ) + propagator = Propagator[type(hamiltonian)](params.timestep, params.mu) + propagator.build(hamiltonian, trial, walkers, mpi_handler) + afqmc = ThermalAFQMC( + system, + hamiltonian, + trial, + walkers, + propagator, + mpi_handler, + params, + debug=debug, + verbose=verbosity, + ) + + return afqmc + + +def build_thermal_afqmc_driver( + comm, + nelec: tuple, + hamiltonian_file: str = "hamiltonian.h5", + seed: int = None, + options: dict = None, + verbosity: int = 0, +): + if comm.rank != 0: + verbosity = 0 + + sys_opts = {"nup": nelec[0], "ndown": nelec[1]} + ham_opts = {"integrals": hamiltonian_file} + qmc_opts = {"rng_seed": seed} + + options["system"] = sys_opts + options["hamiltonian"] = ham_opts + options["qmc"].update(qmc_opts) + options["verbosity"] = verbosity + + return get_driver(options, comm) diff --git a/ipie/addons/thermal/qmc/options.py b/ipie/addons/thermal/qmc/options.py new file mode 100644 index 00000000..671bedab --- /dev/null +++ b/ipie/addons/thermal/qmc/options.py @@ -0,0 +1,122 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +from dataclasses import dataclass +from typing import ClassVar, Optional + +from ipie.utils.io import get_input_value +from ipie.qmc.options import QMCOpts, QMCParams + + +class ThermalQMCOpts(QMCOpts): + r"""Input options and certain constants / parameters derived from them. + + Initialised from a dict containing the following options, not all of which + are required. + + Attributes + ---------- + batched : bool + Whether to do batched calculations. + nwalkers : int + Number of walkers to propagate in a simulation. + dt : float + Timestep. + nsteps : int + Number of steps per block. + nblocks : int + Number of blocks. Total number of iterations = nblocks * nsteps. + nstblz : int + Frequency of Gram-Schmidt orthogonalisation steps. + npop_control : int + Frequency of population control. + pop_control_method : str + Population control method. + eqlb_time : float + Time scale of equilibration phase. Only used to fix local + energy bound when using phaseless approximation. + neqlb : int + Number of time steps for the equilibration phase. Only used to fix the + local energy bound when using phaseless approximation. + rng_seed : int + The random number seed. + mu : float + Chemical potential. + beta : float + Inverse temperature. + """ + # pylint: disable=dangerous-default-value + # TODO: Remove this class / replace with dataclass + def __init__(self, inputs={}, verbose=False): + super().__init__(inputs, verbose) + + self.mu = get_input_value( + inputs, + "mu", + default=None, + verbose=verbose, + ) + self.beta = get_input_value( + inputs, + "beta", + default=None, + verbose=verbose, + ) + + +@dataclass +class ThermalQMCParams(QMCParams): + r"""Input options and certain constants / parameters derived from them. + + Attributes + ---------- + mu : float + Chemical potential. + beta : float + Inverse temperature. + num_walkers : int + Number of walkers **per** core / task / computational unit. + total_num_walkers : int + The total number of walkers in the simulation. + timestep : float + The timestep delta_t + num_steps_per_block : int + Number of steps of propagation before estimators are evaluated. + num_blocks : int + Number of blocks. Total number of iterations = num_blocks * num_steps_per_block. + num_stblz : int + Number of steps before QR stabilization of walkers is performed. + pop_control_freq : int + Frequency at which population control occurs. + rng_seed : int + The random number seed. If run in parallel the seeds on other cores / + threads are determined from this. + """ + # Due to structure of FT algorithm, `num_steps_per_block` is fixed at 1. + # Overide whatever input for backward compatibility. + num_steps_per_block: ClassVar[int] = 1 + mu: Optional[float] = None + beta: Optional[float] = None + pop_control_method: str = 'pair_branch' + + def __post_init__(self): + if self.mu is None: + raise TypeError("__init__ missing 1 required argument: 'mu'") + if self.beta is None: + raise TypeError("__init__ missing 1 required argument: 'beta'") + diff --git a/ipie/addons/thermal/qmc/tests/__init__.py b/ipie/addons/thermal/qmc/tests/__init__.py new file mode 100644 index 00000000..f91ef518 --- /dev/null +++ b/ipie/addons/thermal/qmc/tests/__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. +# +# Authors: Fionn Malone +# Joonho Lee +# diff --git a/ipie/addons/thermal/qmc/tests/test_afqmc_generic.py b/ipie/addons/thermal/qmc/tests/test_afqmc_generic.py new file mode 100644 index 00000000..88e138b1 --- /dev/null +++ b/ipie/addons/thermal/qmc/tests/test_afqmc_generic.py @@ -0,0 +1,168 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import json +import tempfile +import h5py +import uuid +import pytest +import numpy +from typing import Union + +try: + from ipie.addons.thermal.utils.legacy_testing import build_legacy_driver_generic_test_instance + _no_cython = False + +except ModuleNotFoundError: + _no_cython = True + +from ipie.config import MPI +from ipie.analysis.extraction import ( + extract_test_data_hdf5, + extract_data, + extract_observable, + extract_mixed_estimates) +from ipie.addons.thermal.utils.testing import build_driver_generic_test_instance + +comm = MPI.COMM_WORLD +serial_test = comm.size == 1 + +# Unique filename to avoid name collision when running through CI. +if comm.rank == 0: + test_id = str(uuid.uuid1()) + +else: + test_id = None + +test_id = comm.bcast(test_id, root=0) + + +@pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") +@pytest.mark.unit +def test_thermal_afqmc(): + # System params. + nup = 5 + ndown = 5 + nelec = (nup, ndown) + nbasis = 10 + + # Thermal AFQMC params. + mu = -10.0 + beta = 0.1 + timestep = 0.01 + nwalkers = 32 // comm.size + nblocks = 12 + stabilize_freq = 10 + pop_control_freq = 1 + pop_control_method = 'pair_branch' + #pop_control_method = 'comb' + lowrank = False + + verbose = 0 if (comm.rank != 0) else 1 + # Local energy evaluation in legacy code seems wrong. + complex_integrals = False + debug = True + seed = 7 + numpy.random.seed(seed) + + with tempfile.NamedTemporaryFile() as tmpf1, tempfile.NamedTemporaryFile() as tmpf2: + # --------------------------------------------------------------------- + # Test. + # --------------------------------------------------------------------- + afqmc = build_driver_generic_test_instance( + nelec, nbasis, mu, beta, timestep, nblocks, nwalkers=nwalkers, + lowrank=lowrank, pop_control_method=pop_control_method, + stabilize_freq=stabilize_freq, pop_control_freq=pop_control_freq, + complex_integrals=complex_integrals, debug=debug, seed=seed, + verbose=verbose) + afqmc.run(verbose=verbose, estimator_filename=tmpf1.name) + afqmc.finalise() + afqmc.estimators.compute_estimators(afqmc.hamiltonian, afqmc.trial, afqmc.walkers) + + test_energy_data = None + test_energy_numer = None + test_energy_denom = None + test_number_data = None + + if comm.rank == 0: + test_energy_data = extract_observable(afqmc.estimators.filename, "energy") + test_energy_numer = afqmc.estimators["energy"]["ENumer"] + test_energy_denom = afqmc.estimators["energy"]["EDenom"] + test_number_data = extract_observable(afqmc.estimators.filename, "nav") + + # --------------------------------------------------------------------- + # Legacy. + # --------------------------------------------------------------------- + legacy_afqmc = build_legacy_driver_generic_test_instance( + afqmc.hamiltonian, comm, nelec, mu, beta, timestep, + nblocks, nwalkers=nwalkers, lowrank=lowrank, + stabilize_freq=stabilize_freq, + pop_control_freq=pop_control_freq, + pop_control_method=pop_control_method, seed=seed, + estimator_filename=tmpf2.name, verbose=verbose) + legacy_afqmc.run(comm=comm) + legacy_afqmc.finalise(verbose=False) + legacy_afqmc.estimators.estimators["mixed"].update( + legacy_afqmc.qmc, + legacy_afqmc.system, + legacy_afqmc.hamiltonian, + legacy_afqmc.trial, + legacy_afqmc.walk, + 0, + legacy_afqmc.propagators.free_projection) + + legacy_mixed_data = None + enum = None + legacy_energy_numer = None + legacy_energy_denom = None + + if comm.rank == 0: + legacy_mixed_data = extract_mixed_estimates(legacy_afqmc.estimators.filename) + enum = legacy_afqmc.estimators.estimators["mixed"].names + legacy_energy_numer = legacy_afqmc.estimators.estimators["mixed"].estimates[enum.enumer] + legacy_energy_denom = legacy_afqmc.estimators.estimators["mixed"].estimates[enum.edenom] + + # Check. + assert test_energy_numer.real == pytest.approx(legacy_energy_numer.real) + assert test_energy_denom.real == pytest.approx(legacy_energy_denom.real) + assert test_energy_numer.imag == pytest.approx(legacy_energy_numer.imag) + assert test_energy_denom.imag == pytest.approx(legacy_energy_denom.imag) + + assert numpy.mean(test_energy_data.WeightFactor.values[1:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.WeightFactor.values[1:-1].real)) + assert numpy.mean(test_energy_data.Weight.values[1:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.Weight.values[1:-1].real)) + assert numpy.mean(test_energy_data.ENumer.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.ENumer.values[:-1].real)) + assert numpy.mean(test_energy_data.EDenom.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.EDenom.values[:-1].real)) + assert numpy.mean(test_energy_data.ETotal.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.ETotal.values[:-1].real)) + assert numpy.mean(test_energy_data.E1Body.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.E1Body.values[:-1].real)) + assert numpy.mean(test_energy_data.E2Body.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.E2Body.values[:-1].real)) + assert numpy.mean(test_energy_data.HybridEnergy.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.EHybrid.values[:-1].real)) + assert numpy.mean(test_number_data.Nav.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.Nav.values[:-1].real)) + + +if __name__ == '__main__': + test_thermal_afqmc() + diff --git a/ipie/addons/thermal/qmc/tests/ueg/test_afqmc_ueg.py b/ipie/addons/thermal/qmc/tests/ueg/test_afqmc_ueg.py new file mode 100644 index 00000000..a8906a84 --- /dev/null +++ b/ipie/addons/thermal/qmc/tests/ueg/test_afqmc_ueg.py @@ -0,0 +1,545 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import os +import sys +import json +import pprint +import tempfile +import h5py +import uuid +import pytest +import numpy +from typing import Union + +try: + from ipie.addons.thermal.utils.legacy_testing import build_legacy_driver_ueg_test_instance + _no_cython = False + +except ModuleNotFoundError: + _no_cython = True + +from ipie.config import MPI +from ipie.analysis.extraction import ( + get_metadata, + extract_test_data_hdf5, + extract_data, + extract_observable, + extract_mixed_estimates) +from ipie.addons.thermal.utils.testing import build_driver_ueg_test_instance + +comm = MPI.COMM_WORLD +serial_test = comm.size == 1 + +# Unique filename to avoid name collision when running through CI. +if comm.rank == 0: + test_id = str(uuid.uuid1()) + +else: + test_id = None + +test_id = comm.bcast(test_id, root=0) + + +def compare_test_data(ref_data, test_data): + comparison = {} + + for k, v in ref_data.items(): + alias = [k] + + if k == "sys_info": + continue + + elif k == "EHybrid": + alias.append("HybridEnergy") + + err = 0 + ref = ref_data[k] + + for a in alias: + try: + test = test_data[a] + comparison[k] = ( + numpy.array(ref), + numpy.array(test), + numpy.max(numpy.abs(numpy.array(ref) - numpy.array(test))) < 1e-10) + + except KeyError: + err += 1 + + if err == len(alias): + print(f"# Issue with test data key {k}") + + return comparison + + +@pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") +@pytest.mark.unit +def test_thermal_afqmc_1walker(against_ref=False): + # UEG params. + nup = 7 + ndown = 7 + nelec = (nup, ndown) + rs = 1. + ecut = 1. + + # Thermal AFQMC params. + mu = -1. + beta = 0.1 + timestep = 0.01 + nwalkers = 1 + nblocks = 11 + + stabilize_freq = 10 + pop_control_freq = 1 + # `pop_control_method` doesn't matter for 1 walker. + pop_control_method = "pair_branch" + #pop_control_method = "comb" + lowrank = False + + verbose = False if (comm.rank != 0) else True + debug = True + seed = 7 + numpy.random.seed(seed) + + with tempfile.NamedTemporaryFile() as tmpf1, tempfile.NamedTemporaryFile() as tmpf2: + # --------------------------------------------------------------------- + # Test. + # --------------------------------------------------------------------- + afqmc = build_driver_ueg_test_instance( + nelec, rs, ecut, mu, beta, timestep, nblocks, nwalkers=nwalkers, + lowrank=lowrank, pop_control_method=pop_control_method, + stabilize_freq=stabilize_freq, pop_control_freq=pop_control_freq, + debug=debug, seed=seed, verbose=verbose) + afqmc.run(verbose=verbose, estimator_filename=tmpf1.name) + afqmc.finalise() + afqmc.estimators.compute_estimators(afqmc.hamiltonian, afqmc.trial, afqmc.walkers) + + test_energy_data = None + test_energy_numer = None + test_energy_denom = None + test_number_data = None + + if comm.rank == 0: + test_energy_data = extract_observable(afqmc.estimators.filename, "energy") + test_energy_numer = afqmc.estimators["energy"]["ENumer"] + test_energy_denom = afqmc.estimators["energy"]["EDenom"] + test_number_data = extract_observable(afqmc.estimators.filename, "nav") + + # --------------------------------------------------------------------- + # Legacy. + # --------------------------------------------------------------------- + legacy_afqmc = build_legacy_driver_ueg_test_instance( + comm, nelec, rs, ecut, mu, beta, timestep, nblocks, + nwalkers=nwalkers, lowrank=lowrank, + stabilize_freq=stabilize_freq, + pop_control_freq=pop_control_freq, + pop_control_method=pop_control_method, seed=seed, + estimator_filename=tmpf2.name, verbose=verbose) + legacy_afqmc.run(comm=comm) + legacy_afqmc.finalise(verbose=False) + legacy_afqmc.estimators.estimators["mixed"].update( + legacy_afqmc.qmc, + legacy_afqmc.system, + legacy_afqmc.hamiltonian, + legacy_afqmc.trial, + legacy_afqmc.walk, + 0, + legacy_afqmc.propagators.free_projection) + + legacy_mixed_data = None + enum = None + legacy_energy_numer = None + legacy_energy_denom = None + + if comm.rank == 0: + legacy_mixed_data = extract_mixed_estimates(legacy_afqmc.estimators.filename) + enum = legacy_afqmc.estimators.estimators["mixed"].names + legacy_energy_numer = legacy_afqmc.estimators.estimators["mixed"].estimates[enum.enumer] + legacy_energy_denom = legacy_afqmc.estimators.estimators["mixed"].estimates[enum.edenom] + + # Check. + assert test_energy_numer.real == pytest.approx(legacy_energy_numer.real) + assert test_energy_denom.real == pytest.approx(legacy_energy_denom.real) + assert test_energy_numer.imag == pytest.approx(legacy_energy_numer.imag) + assert test_energy_denom.imag == pytest.approx(legacy_energy_denom.imag) + + assert numpy.mean(test_energy_data.WeightFactor.values[1:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.WeightFactor.values[1:-1].real)) + assert numpy.mean(test_energy_data.Weight.values[1:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.Weight.values[1:-1].real)) + assert numpy.mean(test_energy_data.ENumer.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.ENumer.values[:-1].real)) + assert numpy.mean(test_energy_data.EDenom.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.EDenom.values[:-1].real)) + assert numpy.mean(test_energy_data.ETotal.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.ETotal.values[:-1].real)) + assert numpy.mean(test_energy_data.E1Body.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.E1Body.values[:-1].real)) + assert numpy.mean(test_energy_data.E2Body.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.E2Body.values[:-1].real)) + assert numpy.mean(test_energy_data.HybridEnergy.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.EHybrid.values[:-1].real)) + assert numpy.mean(test_number_data.Nav.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.Nav.values[:-1].real)) + + # --------------------------------------------------------------------- + # Test against reference data. + if against_ref: + _data_dir = os.path.abspath(os.path.dirname(__file__)).split("qmc")[0] + "/reference_data/" + _legacy_test_dir = "ueg" + _legacy_test = _data_dir + _legacy_test_dir + "/reference_1walker.json" + + test_name = _legacy_test_dir + with open(_legacy_test, "r") as f: + ref_data = json.load(f) + + skip_val = ref_data.get("extract_skip_value", 10) + _test_energy_data = test_energy_data[::skip_val].to_dict(orient="list") + _test_number_data = test_number_data[::skip_val].to_dict(orient="list") + energy_comparison = compare_test_data(ref_data, _test_energy_data) + number_comparison = compare_test_data(ref_data, _test_number_data) + + print('\nenergy comparison:') + pprint.pprint(energy_comparison) + print('\nnumber comparison:') + pprint.pprint(number_comparison) + + local_err_count = 0 + + for k, v in energy_comparison.items(): + if not v[-1]: + local_err_count += 1 + print(f"\n *** FAILED *** : mismatch between benchmark and test run: {test_name}") + print(f" name = {k}\n ref = {v[0]}\n test = {v[1]}\n delta = {v[0]-v[1]}\n") + + for k, v in number_comparison.items(): + if not v[-1]: + local_err_count += 1 + print(f"\n *** FAILED *** : mismatch between benchmark and test run: {test_name}") + print(f" name = {k}\n ref = {v[0]}\n test = {v[1]}\n delta = {v[0]-v[1]}\n") + + if local_err_count == 0: + print(f"\n*** PASSED : {test_name} ***\n") + + +@pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") +@pytest.mark.unit +def test_thermal_afqmc(against_ref=False): + # UEG params. + nup = 7 + ndown = 7 + nelec = (nup, ndown) + rs = 1. + ecut = 1. + + # Thermal AFQMC params. + mu = -1. + beta = 0.1 + timestep = 0.01 + nwalkers = 32 + # Must be fixed at 1 for Thermal AFQMC--legacy code overides whatever input! + nsteps_per_block = 1 + nblocks = 10 + stabilize_freq = 10 + pop_control_freq = 1 + pop_control_method = "pair_branch" + #pop_control_method = "comb" + lowrank = False + + verbose = False if (comm.rank != 0) else True + debug = True + seed = 7 + numpy.random.seed(seed) + + with tempfile.NamedTemporaryFile() as tmpf1, tempfile.NamedTemporaryFile() as tmpf2: + # --------------------------------------------------------------------- + # Test. + # --------------------------------------------------------------------- + afqmc = build_driver_ueg_test_instance( + nelec, rs, ecut, mu, beta, timestep, nblocks, nwalkers=nwalkers, + lowrank=lowrank, pop_control_method=pop_control_method, + stabilize_freq=stabilize_freq, pop_control_freq=pop_control_freq, + debug=debug, seed=seed, verbose=verbose) + afqmc.run(verbose=verbose, estimator_filename=tmpf1.name) + afqmc.finalise() + afqmc.estimators.compute_estimators(afqmc.hamiltonian, afqmc.trial, afqmc.walkers) + + test_energy_data = None + test_energy_numer = None + test_energy_denom = None + test_number_data = None + + if comm.rank == 0: + test_energy_data = extract_observable(afqmc.estimators.filename, "energy") + test_energy_numer = afqmc.estimators["energy"]["ENumer"] + test_energy_denom = afqmc.estimators["energy"]["EDenom"] + test_number_data = extract_observable(afqmc.estimators.filename, "nav") + + # --------------------------------------------------------------------- + # Legacy. + # --------------------------------------------------------------------- + legacy_afqmc = build_legacy_driver_ueg_test_instance( + comm, nelec, rs, ecut, mu, beta, timestep, nblocks, + nwalkers=nwalkers, lowrank=lowrank, + stabilize_freq=stabilize_freq, + pop_control_freq=pop_control_freq, + pop_control_method=pop_control_method, seed=seed, + estimator_filename=tmpf2.name, verbose=verbose) + legacy_afqmc.run(comm=comm) + legacy_afqmc.finalise(verbose=False) + legacy_afqmc.estimators.estimators["mixed"].update( + legacy_afqmc.qmc, + legacy_afqmc.system, + legacy_afqmc.hamiltonian, + legacy_afqmc.trial, + legacy_afqmc.walk, + 0, + legacy_afqmc.propagators.free_projection) + + legacy_mixed_data = None + enum = None + legacy_energy_numer = None + legacy_energy_denom = None + + if comm.rank == 0: + legacy_mixed_data = extract_mixed_estimates(legacy_afqmc.estimators.filename) + enum = legacy_afqmc.estimators.estimators["mixed"].names + legacy_energy_numer = legacy_afqmc.estimators.estimators["mixed"].estimates[enum.enumer] + legacy_energy_denom = legacy_afqmc.estimators.estimators["mixed"].estimates[enum.edenom] + + # Check. + assert test_energy_numer.real == pytest.approx(legacy_energy_numer.real) + assert test_energy_denom.real == pytest.approx(legacy_energy_denom.real) + assert test_energy_numer.imag == pytest.approx(legacy_energy_numer.imag) + assert test_energy_denom.imag == pytest.approx(legacy_energy_denom.imag) + + assert numpy.mean(test_energy_data.WeightFactor.values[1:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.WeightFactor.values[1:-1].real)) + assert numpy.mean(test_energy_data.Weight.values[1:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.Weight.values[1:-1].real)) + assert numpy.mean(test_energy_data.ENumer.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.ENumer.values[:-1].real)) + assert numpy.mean(test_energy_data.EDenom.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.EDenom.values[:-1].real)) + assert numpy.mean(test_energy_data.ETotal.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.ETotal.values[:-1].real)) + assert numpy.mean(test_energy_data.E1Body.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.E1Body.values[:-1].real)) + assert numpy.mean(test_energy_data.E2Body.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.E2Body.values[:-1].real)) + assert numpy.mean(test_energy_data.HybridEnergy.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.EHybrid.values[:-1].real)) + assert numpy.mean(test_number_data.Nav.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.Nav.values[:-1].real)) + + # --------------------------------------------------------------------- + # Test against reference data. + if against_ref: + _data_dir = os.path.abspath(os.path.dirname(__file__)).split("qmc")[0] + "/reference_data/" + _legacy_test_dir = "ueg" + _legacy_test = _data_dir + _legacy_test_dir + "/reference_nompi.json" + + test_name = _legacy_test_dir + with open(_legacy_test, "r") as f: + ref_data = json.load(f) + + skip_val = ref_data.get("extract_skip_value", 10) + _test_energy_data = test_energy_data[::skip_val].to_dict(orient="list") + _test_number_data = test_number_data[::skip_val].to_dict(orient="list") + energy_comparison = compare_test_data(ref_data, _test_energy_data) + number_comparison = compare_test_data(ref_data, _test_number_data) + + print('\nenergy comparison:') + pprint.pprint(energy_comparison) + print('\nnumber comparison:') + pprint.pprint(number_comparison) + + local_err_count = 0 + + for k, v in energy_comparison.items(): + if not v[-1]: + local_err_count += 1 + print(f"\n *** FAILED *** : mismatch between benchmark and test run: {test_name}") + print(f" name = {k}\n ref = {v[0]}\n test = {v[1]}\n delta = {v[0]-v[1]}\n") + + for k, v in number_comparison.items(): + if not v[-1]: + local_err_count += 1 + print(f"\n *** FAILED *** : mismatch between benchmark and test run: {test_name}") + print(f" name = {k}\n ref = {v[0]}\n test = {v[1]}\n delta = {v[0]-v[1]}\n") + + if local_err_count == 0: + print(f"\n*** PASSED : {test_name} ***\n") + + +@pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") +@pytest.mark.mpi +def test_thermal_afqmc_mpi(against_ref=False): + # UEG params. + nup = 7 + ndown = 7 + nelec = (nup, ndown) + rs = 1. + ecut = 1. + + # Thermal AFQMC params. + mu = -1. + beta = 0.1 + timestep = 0.01 + nwalkers = 32 // comm.size + # Must be fixed at 1 for Thermal AFQMC--legacy code overides whatever input! + nsteps_per_block = 1 + nblocks = 10 + stabilize_freq = 10 + pop_control_freq = 1 + pop_control_method = "pair_branch" + #pop_control_method = "comb" + lowrank = False + + verbose = False if (comm.rank != 0) else True + debug = True + seed = 7 + numpy.random.seed(seed) + + with tempfile.NamedTemporaryFile() as tmpf1, tempfile.NamedTemporaryFile() as tmpf2: + # --------------------------------------------------------------------- + # Test. + # --------------------------------------------------------------------- + afqmc = build_driver_ueg_test_instance( + nelec, rs, ecut, mu, beta, timestep, nblocks, nwalkers=nwalkers, + lowrank=lowrank, pop_control_method=pop_control_method, + stabilize_freq=stabilize_freq, pop_control_freq=pop_control_freq, + debug=debug, seed=seed, verbose=verbose) + afqmc.run(verbose=verbose, estimator_filename=tmpf1.name) + afqmc.finalise() + afqmc.estimators.compute_estimators(afqmc.hamiltonian, afqmc.trial, afqmc.walkers) + + test_energy_data = None + test_energy_numer = None + test_energy_denom = None + test_number_data = None + + if comm.rank == 0: + test_energy_data = extract_observable(afqmc.estimators.filename, "energy") + test_energy_numer = afqmc.estimators["energy"]["ENumer"] + test_energy_denom = afqmc.estimators["energy"]["EDenom"] + test_number_data = extract_observable(afqmc.estimators.filename, "nav") + + # --------------------------------------------------------------------- + # Legacy. + # --------------------------------------------------------------------- + legacy_afqmc = build_legacy_driver_ueg_test_instance( + comm, nelec, rs, ecut, mu, beta, timestep, nblocks, + nwalkers=nwalkers, lowrank=lowrank, + stabilize_freq=stabilize_freq, + pop_control_freq=pop_control_freq, + pop_control_method=pop_control_method, seed=seed, + estimator_filename=tmpf2.name, verbose=verbose) + legacy_afqmc.run(comm=comm) + legacy_afqmc.finalise(verbose=False) + legacy_afqmc.estimators.estimators["mixed"].update( + legacy_afqmc.qmc, + legacy_afqmc.system, + legacy_afqmc.hamiltonian, + legacy_afqmc.trial, + legacy_afqmc.walk, + 0, + legacy_afqmc.propagators.free_projection) + + legacy_mixed_data = None + enum = None + legacy_energy_numer = None + legacy_energy_denom = None + + if comm.rank == 0: + legacy_mixed_data = extract_mixed_estimates(legacy_afqmc.estimators.filename) + enum = legacy_afqmc.estimators.estimators["mixed"].names + legacy_energy_numer = legacy_afqmc.estimators.estimators["mixed"].estimates[enum.enumer] + legacy_energy_denom = legacy_afqmc.estimators.estimators["mixed"].estimates[enum.edenom] + + # Check. + assert test_energy_numer.real == pytest.approx(legacy_energy_numer.real) + assert test_energy_denom.real == pytest.approx(legacy_energy_denom.real) + assert test_energy_numer.imag == pytest.approx(legacy_energy_numer.imag) + assert test_energy_denom.imag == pytest.approx(legacy_energy_denom.imag) + + assert numpy.mean(test_energy_data.WeightFactor.values[1:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.WeightFactor.values[1:-1].real)) + assert numpy.mean(test_energy_data.Weight.values[1:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.Weight.values[1:-1].real)) + assert numpy.mean(test_energy_data.ENumer.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.ENumer.values[:-1].real)) + assert numpy.mean(test_energy_data.EDenom.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.EDenom.values[:-1].real)) + assert numpy.mean(test_energy_data.ETotal.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.ETotal.values[:-1].real)) + assert numpy.mean(test_energy_data.E1Body.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.E1Body.values[:-1].real)) + assert numpy.mean(test_energy_data.E2Body.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.E2Body.values[:-1].real)) + assert numpy.mean(test_energy_data.HybridEnergy.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.EHybrid.values[:-1].real)) + assert numpy.mean(test_number_data.Nav.values[:-1].real) == pytest.approx( + numpy.mean(legacy_mixed_data.Nav.values[:-1].real)) + + # --------------------------------------------------------------------- + # Test against reference data. + if against_ref: + _data_dir = os.path.abspath(os.path.dirname(__file__)).split("qmc")[0] + "/reference_data/" + _legacy_test_dir = "ueg" + _legacy_test = _data_dir + _legacy_test_dir + "/reference.json" + + test_name = _legacy_test_dir + with open(_legacy_test, "r") as f: + ref_data = json.load(f) + + skip_val = ref_data.get("extract_skip_value", 10) + _test_energy_data = test_energy_data[::skip_val].to_dict(orient="list") + _test_number_data = test_number_data[::skip_val].to_dict(orient="list") + energy_comparison = compare_test_data(ref_data, _test_energy_data) + number_comparison = compare_test_data(ref_data, _test_number_data) + + print('\nenergy comparison:') + pprint.pprint(energy_comparison) + print('\nnumber comparison:') + pprint.pprint(number_comparison) + + local_err_count = 0 + + for k, v in energy_comparison.items(): + if not v[-1]: + local_err_count += 1 + print(f"\n *** FAILED *** : mismatch between benchmark and test run: {test_name}") + print(f" name = {k}\n ref = {v[0]}\n test = {v[1]}\n delta = {v[0]-v[1]}\n") + + for k, v in number_comparison.items(): + if not v[-1]: + local_err_count += 1 + print(f"\n *** FAILED *** : mismatch between benchmark and test run: {test_name}") + print(f" name = {k}\n ref = {v[0]}\n test = {v[1]}\n delta = {v[0]-v[1]}\n") + + if local_err_count == 0: + print(f"\n*** PASSED : {test_name} ***\n") + + +if __name__ == '__main__': + test_thermal_afqmc_1walker(against_ref=True) + test_thermal_afqmc(against_ref=True) + #test_thermal_afqmc_mpi(against_ref=True) + diff --git a/ipie/addons/thermal/qmc/thermal_afqmc.py b/ipie/addons/thermal/qmc/thermal_afqmc.py new file mode 100644 index 00000000..3cea31ee --- /dev/null +++ b/ipie/addons/thermal/qmc/thermal_afqmc.py @@ -0,0 +1,339 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +"""Driver to perform Thermal AFQMC calculation""" + +import numpy +import time +import json +from typing import Dict, Optional, Tuple + +from ipie.addons.thermal.walkers.pop_controller import ThermalPopController +from ipie.addons.thermal.walkers.uhf_walkers import UHFThermalWalkers +from ipie.addons.thermal.propagation.propagator import Propagator +from ipie.addons.thermal.estimators.handler import ThermalEstimatorHandler +from ipie.addons.thermal.qmc.options import ThermalQMCParams + +from ipie.utils.io import to_json +from ipie.utils.backend import arraylib as xp +from ipie.utils.backend import synchronize +from ipie.utils.mpi import MPIHandler +from ipie.systems.generic import Generic +from ipie.estimators.estimator_base import EstimatorBase +from ipie.walkers.base_walkers import WalkerAccumulator +from ipie.qmc.afqmc import AFQMC + + +class ThermalAFQMC(AFQMC): + """Thermal AFQMC driver. + + Parameters + ---------- + hamiltonian : + Hamiltonian describing the system. + trial : + Trial density matrix. + walkers : + Walkers used for open ended random walk. + propagator : + Class describing how to propagate walkers. + params : + Parameters of simulation. See ThermalQMCParams for description. + verbose : bool + How much information to print. + + Attributes + ---------- + _parallel_rng_seed : int + Seed deduced from params.rng_seed which is generally different on each + MPI process. + """ + def __init__(self, + system, # For compatibility with 0T AFQMC code. + hamiltonian, + trial, + walkers, + propagator, + mpi_handler, + params: ThermalQMCParams, + debug: bool = False, + verbose: int = 0): + super().__init__(system, hamiltonian, trial, walkers, propagator, mpi_handler, params, verbose) + self.debug = debug + + if self.debug and verbose: + print('# Using legacy `update_weights`.') + + + @staticmethod + def build( + nelec: Tuple[int, int], + mu: float, + beta: float, + hamiltonian, + trial, + nwalkers: int = 100, + stack_size: int = 10, + seed: int = None, + nblocks: int = 100, + timestep: float = 0.005, + stabilize_freq: int = 5, + pop_control_freq: int = 5, + pop_control_method: str = 'pair_branch', + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + debug: bool = False, + verbose: int = 0, + mpi_handler=None,) -> "Thermal AFQMC": + """Factory method to build thermal AFQMC driver from hamiltonian and trial density matrix. + + Parameters + ---------- + nelec : tuple(int, int) + Number of alpha and beta electrons. + mu : float + Chemical potential. + beta : float + Inverse temperature. + hamiltonian : + Hamiltonian describing the system. + trial : + Trial density matrix. + nwalkers : int + Number of walkers per MPI process used in the simulation. The TOTAL + number of walkers is nwalkers * number of processes. + nblocks : int + Number of blocks to perform. + timestep : float + Imaginary timestep. Default 0.005. + stabilize_freq : float + Frequency at which to perform QR factorization of walkers (in units + of steps.) Default 25. + pop_control_freq : int + Frequency at which to perform population control (in units of + steps.) Default 25. + lowrank : bool + Low-rank algorithm for thermal propagation. Doesn't work for now! + lowrank_thresh : bool + Threshold for low-rank algorithm. + verbose : bool + Log verbosity. Default True i.e. print information to stdout. + """ + if mpi_handler is None: + mpi_handler = MPIHandler() + comm = mpi_handler.comm + + else: + comm = mpi_handler.comm + + # pylint: disable = no-value-for-parameter + params = ThermalQMCParams( + mu=mu, + beta=beta, + num_walkers=nwalkers, + total_num_walkers=nwalkers * comm.size, + num_blocks=nblocks, + timestep=timestep, + num_stblz=stabilize_freq, + pop_control_freq=pop_control_freq, + pop_control_method=pop_control_method, + rng_seed=seed) + + system = Generic(nelec) + walkers = UHFThermalWalkers( + trial, hamiltonian.nbasis, nwalkers, stack_size=stack_size, + lowrank=lowrank, lowrank_thresh=lowrank_thresh, + mpi_handler=mpi_handler, verbose=verbose) + propagator = Propagator[type(hamiltonian)]( + timestep, mu, lowrank=lowrank, verbose=verbose) + propagator.build(hamiltonian, trial=trial, walkers=walkers, + mpi_handler=mpi_handler, verbose=verbose) + return ThermalAFQMC( + system, + hamiltonian, + trial, + walkers, + propagator, + mpi_handler, + params, + debug=debug, + verbose=verbose) + + + def run(self, + walkers = None, + verbose: bool = True, + estimator_filename = None, + additional_estimators: Optional[Dict[str, EstimatorBase]] = None, + print_time_slice: bool = False): + """Perform Thermal AFQMC simulation on state object using open-ended random walk. + + Parameters + ---------- + state : :class:`pie.state.State` object + Model and qmc parameters. + + walkers: :class:`pie.walker.Walkers` object + Initial wavefunction / distribution of walkers. + + estimator_filename : str + File to write estimates to. + + additional_estimators : dict + Dictionary of additional estimators to evaluate. + """ + # Setup. + self.setup_timers() + ft_setup = time.time() + eshift = 0. + + if walkers is not None: + self.walkers = walkers + + self.pcontrol = ThermalPopController( + self.params.num_walkers, + self.params.num_steps_per_block, + self.mpi_handler, + self.params.pop_control_method, + verbose=self.verbose) + + self.get_env_info() + self.setup_estimators(estimator_filename, additional_estimators=additional_estimators) + + synchronize() + comm = self.mpi_handler.comm + self.tsetup += time.time() - ft_setup + + # Propagate. + total_steps = self.params.num_steps_per_block * self.params.num_blocks + # TODO: This magic value of 2 is pretty much never controlled on input. + # Moreover I'm not convinced having a two stage shift update actually + # matters at all. + neqlb_steps = 2.0 / self.params.timestep + nslices = numpy.rint(self.params.beta / self.params.timestep).astype(int) + + for step in range(1, total_steps + 1): + synchronize() + start_step = time.time() + + for t in range(nslices): + if self.verbosity >= 2 and comm.rank == 0: + print(" # Timeslice %d of %d." % (t, nslices)) + + start = time.time() + self.propagator.propagate_walkers( + self.walkers, self.hamiltonian, self.trial, eshift, debug=self.debug) + + self.tprop_fbias = self.propagator.timer.tfbias + self.tprop_update = self.propagator.timer.tupdate + self.tprop_vhs = self.propagator.timer.tvhs + self.tprop_gemm = self.propagator.timer.tgemm + + start_clip = time.time() + if t > 0: + wbound = self.pcontrol.total_weight * 0.10 + xp.clip(self.walkers.weight, a_min=-wbound, a_max=wbound, + out=self.walkers.weight) # In-place clipping. + + synchronize() + self.tprop_clip += time.time() - start_clip + + start_barrier = time.time() + if t % self.params.pop_control_freq == 0: + comm.Barrier() + + self.tprop_barrier += time.time() - start_barrier + self.tprop += time.time() - start + + if (t > 0) and (t % self.params.pop_control_freq == 0): + start = time.time() + self.pcontrol.pop_control(self.walkers, comm) + synchronize() + self.tpopc += time.time() - start + self.tpopc_send = self.pcontrol.timer.send_time + self.tpopc_recv = self.pcontrol.timer.recv_time + self.tpopc_comm = self.pcontrol.timer.communication_time + self.tpopc_non_comm = self.pcontrol.timer.non_communication_time + + # Print estimators at each time slice. + if print_time_slice: + self.estimators.compute_estimators( + self.hamiltonian, self.trial, self.walkers) + self.estimators.print_time_slice(comm, t, self.accumulators) + + # Accumulate weight, hybrid energy etc. across block. + start = time.time() + self.accumulators.update(self.walkers) + self.testim += time.time() - start + + # Calculate estimators. + start = time.time() + if step % self.params.num_steps_per_block == 0: + self.estimators.compute_estimators( + self.hamiltonian, self.trial, self.walkers) + + self.estimators.print_block( + comm, step // self.params.num_steps_per_block, self.accumulators) + self.accumulators.zero() + + synchronize() + self.testim += time.time() - start + + if step < neqlb_steps: + eshift = self.accumulators.eshift + + else: + eshift += self.accumulators.eshift - eshift + + self.walkers.reset(self.trial) # Reset stack, weights, phase. + + synchronize() + self.tstep += time.time() - start_step + + def setup_estimators( + self, + filename, + additional_estimators: Optional[Dict[str, EstimatorBase]] = None): + self.accumulators = WalkerAccumulator( + ["Weight", "WeightFactor", "HybridEnergy"], self.params.num_steps_per_block) + comm = self.mpi_handler.comm + self.estimators = ThermalEstimatorHandler( + self.mpi_handler.comm, + self.hamiltonian, + self.trial, + walker_state=self.accumulators, + verbose=(comm.rank == 0 and self.verbose), + filename=filename) + + if additional_estimators is not None: + for k, v in additional_estimators.items(): + self.estimators[k] = v + + # TODO: Move this to estimator and log uuid etc in serialization + json.encoder.FLOAT_REPR = lambda o: format(o, ".6f") + json_string = to_json(self) + self.estimators.json_string = json_string + self.estimators.initialize(comm) + + # Calculate estimates for initial distribution of walkers. + self.estimators.compute_estimators(self.hamiltonian, + self.trial, self.walkers) + self.accumulators.update(self.walkers) + self.estimators.print_block(comm, 0, self.accumulators) + self.accumulators.zero() + diff --git a/ipie/addons/thermal/reference_data/generic/generic_integrals.h5 b/ipie/addons/thermal/reference_data/generic/generic_integrals.h5 new file mode 100644 index 00000000..cd0fa383 Binary files /dev/null and b/ipie/addons/thermal/reference_data/generic/generic_integrals.h5 differ diff --git a/ipie/addons/thermal/reference_data/generic/generic_ref.json b/ipie/addons/thermal/reference_data/generic/generic_ref.json new file mode 100644 index 00000000..d042dba9 --- /dev/null +++ b/ipie/addons/thermal/reference_data/generic/generic_ref.json @@ -0,0 +1 @@ +{"WeightFactor": [1.0, 1.0], "Weight": [1.0, 0.00031475080684714826], "ENumer": [0.12706480956961563, 0.00057748930951224], "EDenom": [1.0, 0.00031475080684714826], "ETotal": [0.12706480956961563, 1.834750847176334], "E1Body": [-12.035567865348977, -1.6611431481738972], "E2Body": [12.162632674918592, 3.495893995350231], "EHybrid": [0.0, 0.0], "Overlap": [1.0, 1.0], "Nav": [9.864020634242591, 5.424380018458521]} \ No newline at end of file diff --git a/ipie/addons/thermal/reference_data/ueg/input.json b/ipie/addons/thermal/reference_data/ueg/input.json new file mode 100644 index 00000000..defb4831 --- /dev/null +++ b/ipie/addons/thermal/reference_data/ueg/input.json @@ -0,0 +1,32 @@ +{ + "system": { + "name": "UEG", + "nup": 7, + "ndown": 7, + "rs": 1.0, + "mu": -1.0, + "ecut": 1.0 + }, + "qmc": { + "dt": 0.01, + "nwalkers": 32, + "blocks": 10, + "nsteps": 1, + "beta": 0.1, + "rng_seed": 7, + "pop_control_freq": 1, + "stabilise_freq": 10, + "batched": false + }, + "trial": { + "name": "one_body" + }, + "walkers": { + "population_control": "pair_branch" + }, + "estimators": { + "mixed": { + "one_rdm": true + } + } +} diff --git a/ipie/addons/thermal/reference_data/ueg/input_1walker.json b/ipie/addons/thermal/reference_data/ueg/input_1walker.json new file mode 100644 index 00000000..425e79ed --- /dev/null +++ b/ipie/addons/thermal/reference_data/ueg/input_1walker.json @@ -0,0 +1,32 @@ +{ + "system": { + "name": "UEG", + "nup": 7, + "ndown": 7, + "rs": 1.0, + "mu": -1.0, + "ecut": 1.0 + }, + "qmc": { + "dt": 0.01, + "nwalkers": 1, + "blocks": 10, + "nsteps": 1, + "beta": 0.1, + "rng_seed": 7, + "pop_control_freq": 1, + "stabilise_freq": 10, + "batched": false + }, + "trial": { + "name": "one_body" + }, + "walkers": { + "population_control": "pair_branch" + }, + "estimators": { + "mixed": { + "one_rdm": true + } + } +} diff --git a/ipie/addons/thermal/reference_data/ueg/input_nompi.json b/ipie/addons/thermal/reference_data/ueg/input_nompi.json new file mode 100644 index 00000000..defb4831 --- /dev/null +++ b/ipie/addons/thermal/reference_data/ueg/input_nompi.json @@ -0,0 +1,32 @@ +{ + "system": { + "name": "UEG", + "nup": 7, + "ndown": 7, + "rs": 1.0, + "mu": -1.0, + "ecut": 1.0 + }, + "qmc": { + "dt": 0.01, + "nwalkers": 32, + "blocks": 10, + "nsteps": 1, + "beta": 0.1, + "rng_seed": 7, + "pop_control_freq": 1, + "stabilise_freq": 10, + "batched": false + }, + "trial": { + "name": "one_body" + }, + "walkers": { + "population_control": "pair_branch" + }, + "estimators": { + "mixed": { + "one_rdm": true + } + } +} diff --git a/ipie/addons/thermal/reference_data/ueg/reference.json b/ipie/addons/thermal/reference_data/ueg/reference.json new file mode 100644 index 00000000..d1cfd38b --- /dev/null +++ b/ipie/addons/thermal/reference_data/ueg/reference.json @@ -0,0 +1 @@ +{"WeightFactor": [32.0, 47.947639102045635], "Weight": [32.0, 31.999999999999993], "ENumer": [853.4128425513718, 986.7978362646822], "EDenom": [32.0, 31.999999999999993], "ETotal": [26.66915132973037, 30.837432383271327], "E1Body": [28.374994808285745, 33.217171356971804], "E2Body": [-1.705843478555375, -2.379738973700476], "EHybrid": [0.0, 0.0], "Overlap": [1.0, 1.0], "Nav": [14.000000381209672, 16.37587194751124]} \ No newline at end of file diff --git a/ipie/addons/thermal/reference_data/ueg/reference_1walker.json b/ipie/addons/thermal/reference_data/ueg/reference_1walker.json new file mode 100644 index 00000000..46a6ca8f --- /dev/null +++ b/ipie/addons/thermal/reference_data/ueg/reference_1walker.json @@ -0,0 +1 @@ +{"WeightFactor": [1.0, 1.0], "Weight": [1.0, 0.1], "ENumer": [26.669151329730372, 3.0880611634447632], "EDenom": [1.0, 0.1], "ETotal": [26.669151329730372, 30.880611634447632], "E1Body": [28.374994808285745, 33.2449965368494], "E2Body": [-1.7058434785553742, -2.364384902401771], "EHybrid": [0.0, 0.0], "Overlap": [1.0, 1.0], "Nav": [14.000000381209672, 16.388846234869657], "sys_info": {"nranks": 1, "branch": "ft_cleanup", "sha1": "3e25fb5948014be236679907b176cb8087863e3e-dirty", "numpy": {"version": "1.24.4", "path": "/Users/jiang/opt/anaconda3/envs/ipie_ftclean/lib/python3.8/site-packages/numpy", "BLAS": {"lib": "openblas64_ openblas64_", "path": "/usr/local/lib"}}, "scipy": {"version": "1.10.1", "path": "/Users/jiang/opt/anaconda3/envs/ipie_ftclean/lib/python3.8/site-packages/scipy"}, "h5py": {"version": "3.9.0", "path": "/Users/jiang/opt/anaconda3/envs/ipie_ftclean/lib/python3.8/site-packages/h5py"}}} \ No newline at end of file diff --git a/ipie/addons/thermal/reference_data/ueg/reference_nompi.json b/ipie/addons/thermal/reference_data/ueg/reference_nompi.json new file mode 100644 index 00000000..d7a15e47 --- /dev/null +++ b/ipie/addons/thermal/reference_data/ueg/reference_nompi.json @@ -0,0 +1 @@ +{"WeightFactor": [32.0, 47.78867518157943], "Weight": [32.0, 32.00000000000001], "ENumer": [853.4128425513711, 986.9845115919738], "EDenom": [32.0, 32.00000000000001], "ETotal": [26.669151329730347, 30.843265987249175], "E1Body": [28.374994808285724, 33.2211214029656], "E2Body": [-1.7058434785553744, -2.3778554157164185], "EHybrid": [0.0, 0.0], "Overlap": [1.0, 1.0], "Nav": [14.000000381209661, 16.3778298504611]} \ No newline at end of file diff --git a/ipie/addons/thermal/trial/__init__.py b/ipie/addons/thermal/trial/__init__.py new file mode 100644 index 00000000..f91ef518 --- /dev/null +++ b/ipie/addons/thermal/trial/__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. +# +# Authors: Fionn Malone +# Joonho Lee +# diff --git a/ipie/addons/thermal/trial/chem_pot.py b/ipie/addons/thermal/trial/chem_pot.py new file mode 100644 index 00000000..85244d40 --- /dev/null +++ b/ipie/addons/thermal/trial/chem_pot.py @@ -0,0 +1,87 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import numpy + +from ipie.addons.thermal.estimators.particle_number import particle_number +from ipie.addons.thermal.estimators.thermal import one_rdm_stable +from ipie.utils.io import format_fixed_width_floats, format_fixed_width_strings + + +def find_chemical_potential(alt_convention, rho, beta, num_bins, target, + deps=1e-6, max_it=1000, verbose=False): + """Find the chemical potential to match . + """ + # TODO: some sort of generic starting point independent of + # system/temperature + dmu1 = dmu2 = 1 + mu1 = -1 + mu2 = 1 + sign = -1 if alt_convention else 1 + if verbose: + print(f"# Finding chemical potential to match = {target:13.8e}") + while numpy.sign(dmu1) * numpy.sign(dmu2) > 0: + rho1 = compute_rho(rho, mu1, beta, sign=sign) + dmat = one_rdm_stable(rho1, num_bins) + dmu1 = delta_nav(dmat, target) + rho2 = compute_rho(rho, mu2, beta, sign=sign) + dmat = one_rdm_stable(rho2, num_bins) + dmu2 = delta_nav(dmat, target) + if numpy.sign(dmu1) * numpy.sign(dmu2) < 0: + if verbose: + print(f"# Chemical potential lies within range of [{mu1:f},{mu2:f}]") + print(f"# delta_mu1 = {dmu1.real:f}, delta_mu2 = {dmu2.real:f}") + break + else: + mu1 -= 2 + mu2 += 2 + if verbose: + print(f"# Increasing chemical potential search to [{mu1:f},{mu2:f}]") + found_mu = False + if verbose: + print("# " + format_fixed_width_strings(["iteration", "mu", "Dmu", ""])) + for i in range(0, max_it): + mu = 0.5 * (mu1 + mu2) + rho_mu = compute_rho(rho, mu, beta, sign=sign) + dmat = one_rdm_stable(rho_mu, num_bins) + dmu = delta_nav(dmat, target).real + if verbose: + out = [i, mu, dmu, particle_number(dmat).real] + print("# " + format_fixed_width_floats(out)) + if abs(dmu) < deps: + found_mu = True + break + else: + if dmu * dmu1 > 0: + mu1 = mu + elif dmu * dmu2 > 0: + mu2 = mu + if found_mu: + return mu + else: + print("# Error chemical potential not found") + return None + + +def delta_nav(dm, nav): + return particle_number(dm) - nav + + +def compute_rho(rho, mu, beta, sign=1): + return numpy.einsum( + "ijk,k->ijk", rho, numpy.exp(sign * beta * mu * numpy.ones(rho.shape[-1]))) diff --git a/ipie/addons/thermal/trial/mean_field.py b/ipie/addons/thermal/trial/mean_field.py new file mode 100644 index 00000000..4ba6fa78 --- /dev/null +++ b/ipie/addons/thermal/trial/mean_field.py @@ -0,0 +1,118 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import numpy +import scipy.linalg + +from ipie.addons.thermal.estimators.generic import fock_generic +from ipie.addons.thermal.estimators.particle_number import particle_number +from ipie.addons.thermal.estimators.thermal import one_rdm_stable +from ipie.addons.thermal.estimators.greens_function import greens_function +from ipie.addons.thermal.trial.chem_pot import compute_rho, find_chemical_potential +from ipie.addons.thermal.trial.one_body import OneBody + +class MeanField(OneBody): + def __init__(self, hamiltonian, nelec, beta, dt, options=None, alt_convention=False, H1=None, verbose=False): + if options is None: + options = {} + + super().__init__(hamiltonian, nelec, beta, dt, options=options, + alt_convention=alt_convention, H1=H1, verbose=verbose) + if verbose: + print("# Building THF density matrix.") + + self.alpha = options.get("alpha", 0.75) + self.max_scf_it = options.get("max_scf_it", self.max_it) + self.max_macro_it = options.get("max_macro_it", self.max_it) + self.find_mu = options.get("find_mu", True) + self.P, HMF, self.mu = self.thermal_hartree_fock(hamiltonian, beta) + muN = self.mu * numpy.eye(hamiltonian.nbasis, dtype=self.G.dtype) + self.dmat = numpy.array([scipy.linalg.expm(-dt * (HMF[0] - muN)), + scipy.linalg.expm(-dt * (HMF[1] - muN))]) + self.dmat_inv = numpy.array([scipy.linalg.inv(self.dmat[0], check_finite=False), + scipy.linalg.inv(self.dmat[1], check_finite=False)]) + self.G = numpy.array([greens_function(self.dmat[0]), + greens_function(self.dmat[1])]) + self.nav = particle_number(self.P).real + + def thermal_hartree_fock(self, hamiltonian, beta): + dt = self.dtau + mu_old = self.mu + P = self.P.copy() + + if self.verbose: + print("# Determining Thermal Hartree-Fock Density Matrix.") + + for it in range(self.max_macro_it): + if self.verbose: + print(f"\n# Macro iteration: {it}") + + HMF = self.scf(hamiltonian, beta, mu_old, P) + rho = numpy.array([scipy.linalg.expm(-dt * HMF[0]), + scipy.linalg.expm(-dt * HMF[1])]) + if self.find_mu: + mu = find_chemical_potential( + self.alt_convention, rho, dt, self.nstack, self.nav, + deps=self.deps, max_it=self.max_it, verbose=self.verbose) + + else: + mu = self.mu + + rho_mu = compute_rho(rho, mu_old, dt) + P = one_rdm_stable(rho_mu, self.nstack) + dmu = abs(mu - mu_old) + + if self.verbose: + print(f"# New mu: {mu:13.8e} Old mu: {mu_old:13.8e} Dmu: {dmu:13.8e}") + + if dmu < self.deps: + break + + mu_old = mu + + return P, HMF, mu + + def scf(self, hamiltonian, beta, mu, P): + # Compute HMF + HMF = fock_generic(hamiltonian, P) + dt = self.dtau + muN = mu * numpy.eye(hamiltonian.nbasis, dtype=self.G.dtype) + rho = numpy.array([scipy.linalg.expm(-dt * (HMF[0] - muN)), + scipy.linalg.expm(-dt * (HMF[1] - muN))]) + Pold = one_rdm_stable(rho, self.nstack) + + if self.verbose: + print("# Running Thermal SCF.") + + for it in range(self.max_scf_it): + HMF = fock_generic(hamiltonian, Pold) + rho = numpy.array([scipy.linalg.expm(-dt * (HMF[0] - muN)), + scipy.linalg.expm(-dt * (HMF[1] - muN))]) + Pnew = (1 - self.alpha) * one_rdm_stable(rho, self.nstack) + self.alpha * Pold + change = numpy.linalg.norm(Pnew - Pold) + + if change < self.deps: + break + + Pold = Pnew.copy() + + if self.verbose: + N = particle_number(P).real + print(f"# Average particle number: {N:13.8e}") + + return HMF diff --git a/ipie/addons/thermal/trial/one_body.py b/ipie/addons/thermal/trial/one_body.py new file mode 100644 index 00000000..c982d452 --- /dev/null +++ b/ipie/addons/thermal/trial/one_body.py @@ -0,0 +1,137 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import numpy +import scipy.linalg + +from ipie.addons.thermal.estimators.particle_number import particle_number +from ipie.addons.thermal.estimators.thermal import one_rdm_stable +from ipie.addons.thermal.estimators.greens_function import greens_function +from ipie.addons.thermal.trial.chem_pot import compute_rho, find_chemical_potential +from ipie.utils.misc import update_stack + + +class OneBody: + def __init__(self, hamiltonian, nelec, beta, dt, options=None, + alt_convention=False, H1=None, verbose=False): + if options is None: + options = {} + + self.name = "thermal" + self.compute_trial_energy = False + self.verbose = verbose + self.alt_convention = alt_convention + + if H1 is None: + try: + self.H1 = hamiltonian.H1 + + except AttributeError: + self.H1 = hamiltonian.h1e + + else: + self.H1 = H1 + + if verbose: + print("# Building OneBody density matrix.") + print(f"# beta in OneBody: {beta}") + print(f"# dt in OneBody: {dt}") + + dmat_up = scipy.linalg.expm(-dt * (self.H1[0])) + dmat_down = scipy.linalg.expm(-dt * (self.H1[1])) + self.dmat = numpy.array([dmat_up, dmat_down]) + cond = numpy.linalg.cond(self.dmat[0]) + + if verbose: + print(f"# condition number of BT: {cond: 10e}") + + self.nelec = nelec + self.nav = options.get("nav", None) + + if self.nav is None: + self.nav = numpy.sum(self.nelec) + + if verbose: + print(f"# Target average electron number: {self.nav}") + + self.max_it = options.get("max_it", 1000) + self.deps = options.get("threshold", 1e-6) + self.mu = options.get("mu", None) + + self.nslice = int(beta / dt) + self.stack_size = options.get("stack_size", None) + + if self.stack_size == None: + if verbose: + print("# Estimating stack size from BT.") + + self.cond = numpy.linalg.cond(self.dmat[0]) + # We will end up multiplying many BTs together. Can roughly determine + # safe stack size from condition number of BT as the condition number of + # the product will scale roughly as cond(BT)^(number of products). + # We can determine a conservative stack size by requiring that the + # condition number of the product does not exceed 1e3. + self.stack_size = min(self.nslice, int(3.0 / numpy.log10(self.cond))) + + if verbose: + print("# Initial stack size, # of slices: {}, {}".format( + self.stack_size, self.nslice)) + + # Adjust stack size + self.stack_size = update_stack(self.stack_size, self.nslice, verbose=verbose) + self.nstack = int(beta / (self.stack_size * dt)) + + if verbose: + print(f"# Number of stacks: {self.nstack}") + + sign = 1 + if self.alt_convention: + if verbose: + print("# Using alternate sign convention for chemical potential.") + + sign = -1 + + self.dtau = self.stack_size * dt + + if self.mu is None: + self.rho = numpy.array([scipy.linalg.expm(-self.dtau * (self.H1[0])), + scipy.linalg.expm(-self.dtau * (self.H1[1]))]) + self.mu = find_chemical_potential( + self.alt_convention, self.rho, self.dtau, self.nstack, + self.nav, deps=self.deps, max_it=self.max_it, verbose=verbose) + + else: + self.rho = numpy.array([scipy.linalg.expm(-self.dtau * (self.H1[0])), + scipy.linalg.expm(-self.dtau * (self.H1[1]))]) + + if self.verbose: + print(f"# Chemical potential in trial density matrix: {self.mu: .10e}") + + self.P = one_rdm_stable(compute_rho(self.rho, self.mu, self.dtau, sign=sign), self.nstack) + self.nav = particle_number(self.P).real + + if self.verbose: + print(f"# Average particle number in trial density matrix: {self.nav}") + + self.dmat = compute_rho(self.dmat, self.mu, dt, sign=sign) + self.dmat_inv = numpy.array([scipy.linalg.inv(self.dmat[0], check_finite=False), + scipy.linalg.inv(self.dmat[1], check_finite=False)]) + + self.G = numpy.array([greens_function(self.dmat[0]), greens_function(self.dmat[1])]) + self.error = False + self.init = numpy.array([0]) diff --git a/ipie/addons/thermal/trial/tests/__init__.py b/ipie/addons/thermal/trial/tests/__init__.py new file mode 100644 index 00000000..f91ef518 --- /dev/null +++ b/ipie/addons/thermal/trial/tests/__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. +# +# Authors: Fionn Malone +# Joonho Lee +# diff --git a/ipie/addons/thermal/trial/tests/test_chem_pot.py b/ipie/addons/thermal/trial/tests/test_chem_pot.py new file mode 100644 index 00000000..7fcbcff9 --- /dev/null +++ b/ipie/addons/thermal/trial/tests/test_chem_pot.py @@ -0,0 +1,52 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import numpy +import scipy.linalg +import pytest + +from ipie.addons.thermal.trial.chem_pot import find_chemical_potential +from ipie.legacy.trial_density_matrices.chem_pot import find_chemical_potential as legacy_find_chemical_potential + + +@pytest.mark.unit +def test_find_chemical_potential(): + dt = 0.01 + beta = 1 + stack_size = 3 + nstack = 20 + nav = 7 + nbsf = 14 + alt_convention = False + + dtau = dt * stack_size + h1e = numpy.random.random((nbsf, nbsf)) + rho = numpy.array([scipy.linalg.expm(-dtau * h1e), + scipy.linalg.expm(-dtau * h1e)]) + + mu = find_chemical_potential(alt_convention, rho, dt, nstack, nav) + legacy_mu = legacy_find_chemical_potential(alt_convention, rho, dt, nstack, nav) + + numpy.testing.assert_allclose(mu, legacy_mu) + + +if __name__ == '__main__': + test_find_chemical_potential() + + + diff --git a/ipie/addons/thermal/trial/tests/test_mean_field.py b/ipie/addons/thermal/trial/tests/test_mean_field.py new file mode 100644 index 00000000..8c3aaf6d --- /dev/null +++ b/ipie/addons/thermal/trial/tests/test_mean_field.py @@ -0,0 +1,100 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import numpy +import pytest + +try: + from ipie.legacy.trial_density_matrices.mean_field import MeanField as LegacyMeanField + _no_cython = False + +except ModuleNotFoundError: + _no_cython = True + +from ipie.systems.generic import Generic +from ipie.utils.testing import generate_hamiltonian +from ipie.hamiltonians.generic import Generic as HamGeneric +from ipie.addons.thermal.trial.mean_field import MeanField +from ipie.legacy.hamiltonians._generic import Generic as LegacyHamGeneric + + +@pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") +@pytest.mark.unit +def test_mean_field(): + nup = 5 + ndown = 5 + nelec = (nup, ndown) + nbasis = 10 + + mu = -10. + beta = 0.1 + timestep = 0.01 + + alt_convention = False + sparse = False + complex_integrals = True + verbose = True + + sym = 8 + if complex_integrals: sym = 4 + + # Test. + system = Generic(nelec) + h1e, chol, _, eri = generate_hamiltonian(nbasis, nelec, cplx=complex_integrals, + sym=sym, tol=1e-10) + hamiltonian = HamGeneric(h1e=numpy.array([h1e, h1e]), + chol=chol.reshape((-1, nbasis**2)).T.copy(), + ecore=0) + trial = MeanField(hamiltonian, nelec, beta, timestep, verbose=verbose) + + # Lgeacy. + legacy_system = Generic(nelec, verbose=verbose) + legacy_system.mu = mu + legacy_hamiltonian = LegacyHamGeneric( + h1e=hamiltonian.H1, + chol=hamiltonian.chol, + ecore=hamiltonian.ecore, verbose=verbose) + legacy_hamiltonian.hs_pot = numpy.copy(hamiltonian.chol) + legacy_hamiltonian.hs_pot = legacy_hamiltonian.hs_pot.T.reshape( + (hamiltonian.nchol, hamiltonian.nbasis, hamiltonian.nbasis)) + legacy_hamiltonian.mu = mu + legacy_hamiltonian._alt_convention = alt_convention + legacy_hamiltonian.sparse = sparse + legacy_trial = LegacyMeanField(legacy_system, legacy_hamiltonian, beta, + timestep, verbose=verbose) + + assert trial.nelec == nelec + numpy.testing.assert_almost_equal(trial.nav, numpy.sum(nelec), decimal=5) + assert trial.rho.shape == (2, nbasis, nbasis) + assert trial.dmat.shape == (2, nbasis, nbasis) + assert trial.P.shape == (2, nbasis, nbasis) + assert trial.G.shape == (2, nbasis, nbasis) + + numpy.testing.assert_allclose(trial.mu, legacy_trial.mu) + numpy.testing.assert_allclose(trial.nav, legacy_trial.nav) + numpy.testing.assert_allclose(trial.P, legacy_trial.P) + numpy.testing.assert_allclose(trial.G, legacy_trial.G) + numpy.testing.assert_allclose(trial.dmat, legacy_trial.dmat) + numpy.testing.assert_allclose(trial.dmat_inv, legacy_trial.dmat_inv) + + +if __name__ == '__main__': + test_mean_field() + + + diff --git a/ipie/addons/thermal/trial/tests/test_one_body.py b/ipie/addons/thermal/trial/tests/test_one_body.py new file mode 100644 index 00000000..bdc7dce8 --- /dev/null +++ b/ipie/addons/thermal/trial/tests/test_one_body.py @@ -0,0 +1,66 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import numpy +import pytest + +from ipie.systems.generic import Generic +from ipie.utils.testing import generate_hamiltonian +from ipie.hamiltonians.generic import Generic as HamGeneric +from ipie.addons.thermal.trial.one_body import OneBody + + +@pytest.mark.unit +def test_one_body(): + nup = 5 + ndown = 5 + nelec = (nup, ndown) + nbasis = 10 + + mu = -1. + beta = 0.1 + timestep = 0.01 + + complex_integrals = True + verbose = True + + sym = 8 + if complex_integrals: sym = 4 + + # Test. + system = Generic(nelec) + h1e, chol, _, eri = generate_hamiltonian(nbasis, nelec, cplx=complex_integrals, + sym=sym, tol=1e-10) + hamiltonian = HamGeneric(h1e=numpy.array([h1e, h1e]), + chol=chol.reshape((-1, nbasis**2)).T.copy(), + ecore=0) + trial = OneBody(hamiltonian, nelec, beta, timestep, verbose=verbose) + + assert trial.nelec == nelec + numpy.testing.assert_almost_equal(trial.nav, numpy.sum(nelec), decimal=6) + assert trial.rho.shape == (2, nbasis, nbasis) + assert trial.dmat.shape == (2, nbasis, nbasis) + assert trial.P.shape == (2, nbasis, nbasis) + assert trial.G.shape == (2, nbasis, nbasis) + + +if __name__ == '__main__': + test_one_body() + + + diff --git a/ipie/addons/thermal/trial/utils.py b/ipie/addons/thermal/trial/utils.py new file mode 100644 index 00000000..60808a15 --- /dev/null +++ b/ipie/addons/thermal/trial/utils.py @@ -0,0 +1,69 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +from ipie.addons.thermal.trial.mean_field import MeanField +from ipie.addons.thermal.trial.one_body import OneBody + + +def get_trial_density_matrix(hamiltonian, nelec, beta, dt, options=None, + comm=None, verbose=False): + """Wrapper to select trial wavefunction class. + + Parameters + ---------- + + Returns + ------- + trial : class or None + Trial density matrix class. + """ + if options is None: + options = {} + + trial_type = options.get("name", "one_body") + alt_convention = options.get("alt_convention", False) + if comm is None or comm.rank == 0: + if trial_type == "one_body_mod": + trial = OneBody( + hamiltonian, + nelec, + beta, + dt, + options=options, + H1=hamiltonian.h1e_mod, + verbose=verbose, + ) + + elif trial_type == "one_body": + trial = OneBody(hamiltonian, nelec, beta, dt, options=options, + alt_convention=alt_convention, verbose=verbose) + + elif trial_type == "thermal_hartree_fock": + trial = MeanField(hamiltonian, nelec, beta, dt, options=options, + alt_convention=alt_convention, verbose=verbose) + + else: + trial = None + + else: + trial = None + + if comm is not None: + trial = comm.bcast(trial) + + return trial diff --git a/ipie/addons/thermal/utils/__init__.py b/ipie/addons/thermal/utils/__init__.py new file mode 100644 index 00000000..f91ef518 --- /dev/null +++ b/ipie/addons/thermal/utils/__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. +# +# Authors: Fionn Malone +# Joonho Lee +# diff --git a/ipie/addons/thermal/utils/legacy_testing.py b/ipie/addons/thermal/utils/legacy_testing.py new file mode 100644 index 00000000..2e2c829d --- /dev/null +++ b/ipie/addons/thermal/utils/legacy_testing.py @@ -0,0 +1,492 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import numpy +from typing import Tuple, Union + +from ipie.systems.generic import Generic +from ipie.utils.mpi import MPIHandler + +from ipie.addons.thermal.qmc.options import ThermalQMCOpts + +from ipie.legacy.systems.ueg import UEG as LegacyUEG +from ipie.legacy.hamiltonians.ueg import UEG as LegacyHamUEG +from ipie.legacy.hamiltonians._generic import Generic as LegacyHamGeneric +from ipie.legacy.trial_density_matrices.onebody import OneBody as LegacyOneBody +from ipie.legacy.trial_density_matrices.mean_field import MeanField as LegacyMeanField +from ipie.legacy.walkers.handler import Walkers +from ipie.legacy.thermal_propagation.continuous import Continuous +from ipie.legacy.thermal_propagation.planewave import PlaneWave +from ipie.legacy.qmc.thermal_afqmc import ThermalAFQMC as LegacyThermalAFQMC + + +def legacy_propagate_walkers(legacy_hamiltonian, legacy_trial, legacy_walkers, legacy_propagator, xi=None): + if xi is None: + xi = [None] * len(legacy_walkers) + + for iw, walker in enumerate(legacy_walkers.walkers): + legacy_propagator.propagate_walker( + legacy_hamiltonian, walker, legacy_trial, xi=xi[iw]) + + return legacy_walkers + + +def build_legacy_generic_test_case_handlers( + hamiltonian, + comm, + nelec: Tuple[int, int], + mu: float, + beta: float, + timestep: float, + nwalkers: int = 100, + stack_size: int = 10, + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + stabilize_freq: int = 5, + pop_control_freq: int = 5, + pop_control_method: str = 'pair_branch', + alt_convention: bool = False, + sparse: bool = False, + mf_trial: bool = True, + propagate: bool = False, + seed: Union[int, None] = None, + verbose: int = 0): + numpy.random.seed(seed) + + legacy_options = { + "walkers": { + "stack_size": stack_size, + "low_rank": lowrank, + "low_rank_thresh": lowrank_thresh, + "pop_control": pop_control_method + }, + + "propagator": { + "optimised": False + }, + } + + # 1. Build system. + legacy_system = Generic(nelec, verbose=verbose) + legacy_system.mu = mu + + # 2. Build Hamiltonian. + legacy_hamiltonian = LegacyHamGeneric( + h1e=hamiltonian.H1, + chol=hamiltonian.chol, + ecore=hamiltonian.ecore) + legacy_hamiltonian.hs_pot = numpy.copy(hamiltonian.chol) + legacy_hamiltonian.hs_pot = legacy_hamiltonian.hs_pot.T.reshape( + (hamiltonian.nchol, hamiltonian.nbasis, hamiltonian.nbasis)) + legacy_hamiltonian.mu = mu + legacy_hamiltonian._alt_convention = alt_convention + legacy_hamiltonian.sparse = sparse + + # 3. Build trial. + legacy_trial = LegacyOneBody(legacy_system, legacy_hamiltonian, beta, + timestep, verbose=verbose) + if mf_trial: + legacy_trial = LegacyMeanField(legacy_system, legacy_hamiltonian, beta, + timestep, verbose=verbose) + # 4. Build walkers. + qmc_opts = ThermalQMCOpts() + qmc_opts.nwalkers = nwalkers + qmc_opts.ntot_walkers = nwalkers + qmc_opts.beta = beta + qmc_opts.nsteps = 1 + qmc_opts.dt = timestep + qmc_opts.nstblz = stabilize_freq + qmc_opts.npop_control = pop_control_freq + qmc_opts.pop_control_method = pop_control_method + qmc_opts.seed = seed + + legacy_walkers = Walkers(legacy_system, legacy_hamiltonian, legacy_trial, + qmc_opts, walker_opts=legacy_options['walkers'], + verbose=verbose, comm=comm) + + # 5. Build propagator. + legacy_propagator = Continuous( + legacy_options["propagator"], qmc_opts, legacy_system, + legacy_hamiltonian, legacy_trial, verbose=verbose, + lowrank=lowrank) + + if propagate: + for t in range(legacy_walkers[0].stack.ntime_slices): + for iw, walker in enumerate(legacy_walkers): + legacy_propagator.propagate_walker( + legacy_hamiltonian, walker, legacy_trial) + + legacy_objs = {'system': legacy_system, + 'trial': legacy_trial, + 'hamiltonian': legacy_hamiltonian, + 'walkers': legacy_walkers, + 'propagator': legacy_propagator} + return legacy_objs + + +def build_legacy_generic_test_case_handlers_mpi( + hamiltonian, + mpi_handler: MPIHandler, + nelec: Tuple[int, int], + mu: float, + beta: float, + timestep: float, + nwalkers: int = 100, + stack_size: int = 10, + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + stabilize_freq: int = 5, + pop_control_freq: int = 5, + pop_control_method: str = 'pair_branch', + alt_convention: bool = False, + sparse: bool = False, + mf_trial: bool = True, + propagate: bool = False, + seed: Union[int, None] = None, + verbose: int = 0): + numpy.random.seed(seed) + comm = mpi_handler.comm + + legacy_options = { + "walkers": { + "stack_size": stack_size, + "low_rank": lowrank, + "low_rank_thresh": lowrank_thresh, + "pop_control": pop_control_method + }, + + "propagator": { + "optimised": False + }, + } + + # 1. Build system. + legacy_system = Generic(nelec, verbose=verbose) + legacy_system.mu = mu + + # 2. Build Hamiltonian. + legacy_hamiltonian = LegacyHamGeneric( + h1e=hamiltonian.H1, + chol=hamiltonian.chol, + ecore=hamiltonian.ecore) + legacy_hamiltonian.hs_pot = numpy.copy(hamiltonian.chol) + legacy_hamiltonian.hs_pot = legacy_hamiltonian.hs_pot.T.reshape( + (hamiltonian.nchol, hamiltonian.nbasis, hamiltonian.nbasis)) + legacy_hamiltonian.mu = mu + legacy_hamiltonian._alt_convention = alt_convention + legacy_hamiltonian.sparse = sparse + + # 3. Build trial. + legacy_trial = LegacyOneBody(legacy_system, legacy_hamiltonian, beta, + timestep, verbose=verbose) + if mf_trial: + legacy_trial = LegacyMeanField(legacy_system, legacy_hamiltonian, beta, + timestep, verbose=verbose) + # 4. Build walkers. + qmc_opts = ThermalQMCOpts() + qmc_opts.nwalkers = nwalkers + qmc_opts.ntot_walkers = nwalkers * comm.size + qmc_opts.beta = beta + qmc_opts.nsteps = 1 + qmc_opts.dt = timestep + qmc_opts.nstblz = stabilize_freq + qmc_opts.npop_control = pop_control_freq + qmc_opts.pop_control_method = pop_control_method + qmc_opts.seed = seed + + legacy_walkers = Walkers(legacy_system, legacy_hamiltonian, legacy_trial, + qmc_opts, walker_opts=legacy_options['walkers'], + verbose=verbose, comm=comm) + + # 5. Build propagator. + legacy_propagator = Continuous( + legacy_options["propagator"], qmc_opts, legacy_system, + legacy_hamiltonian, legacy_trial, verbose=verbose, + lowrank=lowrank) + + if propagate: + for t in range(legacy_walkers[0].stack.ntime_slices): + for iw, walker in enumerate(legacy_walkers): + legacy_propagator.propagate_walker( + legacy_hamiltonian, walker, legacy_trial) + + legacy_objs = {'system': legacy_system, + 'trial': legacy_trial, + 'hamiltonian': legacy_hamiltonian, + 'walkers': legacy_walkers, + 'propagator': legacy_propagator} + return legacy_objs + + +def build_legacy_driver_generic_test_instance( + hamiltonian, + comm, + nelec: Tuple[int, int], + mu: float, + beta: float, + timestep: float, + nblocks: int, + nwalkers: int = 100, + stack_size: int = 10, + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + stabilize_freq: int = 5, + pop_control_freq: int = 5, + pop_control_method: str = 'pair_branch', + alt_convention: bool = False, + sparse: bool = False, + seed: Union[int, None] = None, + estimator_filename: Union[str, None] = None, + verbose: int = 0): + nup, ndown = nelec + numpy.random.seed(seed) + + legacy_options = { + "qmc": { + "dt": timestep, + # Input of `nwalkers` refers to the total number of walkers in + # legacy `ThermalAFQMC`. + "nwalkers": nwalkers * comm.size, + "blocks": nblocks, + "nsteps": 1, + "beta": beta, + "stabilise_freq": stabilize_freq, + "pop_control_freq": pop_control_freq, + "pop_control_method": pop_control_method, + "rng_seed": seed, + "batched": False + }, + + "propagator": { + "optimised": False + }, + + "walkers": { + "stack_size": stack_size, + "low_rank": lowrank, + "low_rank_thresh": lowrank_thresh, + "pop_control": pop_control_method + }, + + "system": { + "name": "Generic", + "nup": nup, + "ndown": ndown, + "mu": mu + }, + + "estimators": { + "filename": estimator_filename, + }, + } + + legacy_system = Generic(nelec) + legacy_system.mu = mu + legacy_hamiltonian = LegacyHamGeneric( + h1e=hamiltonian.H1, + chol=hamiltonian.chol, + ecore=hamiltonian.ecore) + legacy_hamiltonian.hs_pot = numpy.copy(hamiltonian.chol) + legacy_hamiltonian.hs_pot = legacy_hamiltonian.hs_pot.T.reshape( + (hamiltonian.nchol, hamiltonian.nbasis, hamiltonian.nbasis)) + legacy_hamiltonian.mu = mu + legacy_hamiltonian._alt_convention = alt_convention + legacy_hamiltonian.sparse = sparse + legacy_trial = LegacyMeanField(legacy_system, legacy_hamiltonian, beta, timestep) + + afqmc = LegacyThermalAFQMC(comm, legacy_options, legacy_system, + legacy_hamiltonian, legacy_trial, verbose=verbose) + return afqmc + + +def build_legacy_ueg_test_case_handlers( + comm, + nelec: Tuple[int, int], + rs: float, + ecut: float, + mu: float, + beta: float, + timestep: float, + nwalkers: int = 100, + stack_size: int = 10, + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + propagate: bool = False, + stabilize_freq: int = 5, + pop_control_freq: int = 5, + pop_control_method: str = 'pair_branch', + alt_convention: bool = False, + sparse: bool = False, + seed: Union[int, None] = None, + verbose: int = 0): + numpy.random.seed(seed) + nup, ndown = nelec + legacy_options = { + "ueg": { + "nup": nup, + "ndown": ndown, + "rs": rs, + "ecut": ecut, + "thermal": True, + "write_integrals": False, + "low_rank": lowrank + }, + + "propagator": { + "optimised": False + }, + + "walkers": { + "stack_size": stack_size, + "low_rank": lowrank, + "low_rank_thresh": lowrank_thresh, + "pop_control": pop_control_method + }, + } + + # 1. Build out system. + legacy_system = LegacyUEG(options=legacy_options['ueg']) + legacy_system.mu = mu + + # 2. Build Hamiltonian. + legacy_hamiltonian = LegacyHamUEG(legacy_system, options=legacy_options['ueg']) + legacy_hamiltonian.mu = mu + legacy_hamiltonian._alt_convention = alt_convention + + # 3. Build trial. + legacy_trial = LegacyOneBody(legacy_system, legacy_hamiltonian, beta, + timestep, verbose=verbose) + + # 4. Build walkers. + qmc_opts = ThermalQMCOpts() + qmc_opts.nwalkers = nwalkers + qmc_opts.ntot_walkers = nwalkers * comm.size + qmc_opts.beta = beta + qmc_opts.nsteps = 1 + qmc_opts.dt = timestep + qmc_opts.nstblz = stabilize_freq + qmc_opts.npop_control = pop_control_freq + qmc_opts.pop_control_method = pop_control_method + qmc_opts.seed = seed + + legacy_walkers = Walkers(legacy_system, legacy_hamiltonian, legacy_trial, + qmc_opts, walker_opts=legacy_options['walkers'], + verbose=verbose, comm=comm) + + # 5. Build propagator. + legacy_propagator = PlaneWave(legacy_system, legacy_hamiltonian, legacy_trial, + qmc_opts, options=legacy_options["propagator"], + lowrank=lowrank, verbose=verbose) + + if propagate: + for t in range(legacy_walkers[0].stack.ntime_slices): + for iw, walker in enumerate(legacy_walkers): + legacy_propagator.propagate_walker( + legacy_hamiltonian, walker, legacy_trial) + + legacy_objs = {'system': legacy_system, + 'trial': legacy_trial, + 'hamiltonian': legacy_hamiltonian, + 'walkers': legacy_walkers, + 'propagator': legacy_propagator} + return legacy_objs + + +def build_legacy_driver_ueg_test_instance( + comm, + nelec: Tuple[int, int], + rs: float, + ecut: float, + mu: float, + beta: float, + timestep: float, + nblocks: int, + nwalkers: int = 100, + stack_size: int = 10, + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + stabilize_freq: int = 5, + pop_control_freq: int = 5, + pop_control_method: str = 'pair_branch', + alt_convention: bool = False, + sparse: bool = False, + seed: Union[int, None] = None, + estimator_filename: Union[str, None] = None, + verbose: int = 0): + numpy.random.seed(seed) + nup, ndown = nelec + legacy_options = { + "ueg": { + "nup": nup, + "ndown": ndown, + "rs": rs, + "ecut": ecut, + "thermal": True, + "write_integrals": False, + "low_rank": lowrank + }, + + "qmc": { + "dt": timestep, + # Input of `nwalkers` refers to the total number of walkers in + # legacy `ThermalAFQMC`. + "nwalkers": nwalkers * comm.size, + "blocks": nblocks, + "nsteps": 1, + "beta": beta, + "stabilise_freq": stabilize_freq, + "pop_control_freq": pop_control_freq, + "pop_control_method": pop_control_method, + "rng_seed": seed, + "batched": False + }, + + "propagator": { + "optimised": False + }, + + "walkers": { + "stack_size": stack_size, + "low_rank": lowrank, + "low_rank_thresh": lowrank_thresh, + "pop_control": pop_control_method + }, + + "estimators": { + "filename": estimator_filename, + }, + } + + # 1. Build out system. + legacy_system = LegacyUEG(options=legacy_options['ueg']) + legacy_system.mu = mu + + # 2. Build Hamiltonian. + legacy_hamiltonian = LegacyHamUEG(legacy_system, options=legacy_options['ueg']) + legacy_hamiltonian.mu = mu + legacy_hamiltonian._alt_convention = alt_convention + + # 3. Build trial. + legacy_trial = LegacyOneBody(legacy_system, legacy_hamiltonian, beta, + timestep, verbose=verbose) + + # 4. Build Thermal AFQMC. + afqmc = LegacyThermalAFQMC(comm, legacy_options, legacy_system, + legacy_hamiltonian, legacy_trial, verbose=verbose) + return afqmc + diff --git a/ipie/addons/thermal/utils/testing.py b/ipie/addons/thermal/utils/testing.py new file mode 100644 index 00000000..0e25b695 --- /dev/null +++ b/ipie/addons/thermal/utils/testing.py @@ -0,0 +1,346 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import numpy +from typing import Tuple, Union + +from ipie.utils.mpi import MPIHandler +from ipie.utils.testing import generate_hamiltonian +from ipie.hamiltonians.generic import Generic as HamGeneric + +from ipie.addons.thermal.utils.ueg import UEG +from ipie.addons.thermal.trial.one_body import OneBody +from ipie.addons.thermal.trial.mean_field import MeanField +from ipie.addons.thermal.walkers.uhf_walkers import UHFThermalWalkers +from ipie.addons.thermal.propagation.phaseless_generic import PhaselessGeneric +from ipie.addons.thermal.qmc.thermal_afqmc import ThermalAFQMC + + +def build_generic_test_case_handlers( + nelec: Tuple[int, int], + nbasis: int, + mu: float, + beta: float, + timestep: float, + nwalkers: int = 100, + stack_size: int = 10, + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + diagonal: bool = False, + mf_trial: bool = True, + propagate: bool = False, + complex_integrals: bool = False, + debug: bool = False, + with_eri: bool = False, + seed: Union[int, None] = None, + verbose: int = 0): + sym = 8 + if complex_integrals: sym = 4 + numpy.random.seed(seed) + + # 1. Generate random integrals. + h1e, chol, _, eri = generate_hamiltonian(nbasis, nelec, cplx=complex_integrals, + sym=sym, tol=1e-10) + + if diagonal: + h1e = numpy.diag(numpy.diag(h1e)) + + # 2. Build Hamiltonian. + hamiltonian = HamGeneric(h1e=numpy.array([h1e, h1e]), + chol=chol.reshape((-1, nbasis**2)).T.copy(), + ecore=0) + + # 3. Build trial. + trial = OneBody(hamiltonian, nelec, beta, timestep, verbose=verbose) + + if mf_trial: + trial = MeanField(hamiltonian, nelec, beta, timestep, verbose=verbose) + + # 4. Build walkers. + walkers = UHFThermalWalkers( + trial, nbasis, nwalkers, stack_size=stack_size, lowrank=lowrank, + lowrank_thresh=lowrank_thresh, verbose=verbose) + + # 5. Build propagator. + propagator = PhaselessGeneric(timestep, mu, lowrank=lowrank, verbose=verbose) + propagator.build(hamiltonian, trial=trial, walkers=walkers, verbose=verbose) + + if propagate: + for t in range(walkers.stack[0].nslice): + propagator.propagate_walkers(walkers, hamiltonian, trial, debug=debug) + + objs = {'trial': trial, + 'hamiltonian': hamiltonian, + 'walkers': walkers, + 'propagator': propagator} + + if with_eri: + objs['eri'] = eri + + return objs + + +def build_generic_test_case_handlers_mpi( + nelec: Tuple[int, int], + nbasis: int, + mu: float, + beta: float, + timestep: float, + mpi_handler: MPIHandler, + nwalkers: int = 100, + stack_size: int = 10, + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + diagonal: bool = False, + mf_trial: bool = True, + propagate: bool = False, + complex_integrals: bool = False, + debug: bool = False, + seed: Union[int, None] = None, + verbose: int = 0): + sym = 8 + if complex_integrals: sym = 4 + numpy.random.seed(seed) + + # 1. Generate random integrals. + h1e, chol, _, _ = generate_hamiltonian(nbasis, nelec, cplx=complex_integrals, + sym=sym, tol=1e-10) + + if diagonal: + h1e = numpy.diag(numpy.diag(h1e)) + + # 2. Build Hamiltonian. + hamiltonian = HamGeneric(h1e=numpy.array([h1e, h1e]), + chol=chol.reshape((-1, nbasis**2)).T.copy(), + ecore=0) + + # 3. Build trial. + trial = OneBody(hamiltonian, nelec, beta, timestep, verbose=verbose) + + if mf_trial: + trial = MeanField(hamiltonian, nelec, beta, timestep, verbose=verbose) + + # 4. Build walkers. + walkers = UHFThermalWalkers( + trial, nbasis, nwalkers, stack_size=stack_size, lowrank=lowrank, + lowrank_thresh=lowrank_thresh, mpi_handler=mpi_handler, verbose=verbose) + + # 5. Build propagator. + propagator = PhaselessGeneric(timestep, mu, lowrank=lowrank, verbose=verbose) + propagator.build(hamiltonian, trial=trial, walkers=walkers, + mpi_handler=mpi_handler, verbose=verbose) + + if propagate: + for t in range(walkers.stack[0].nslice): + propagator.propagate_walkers(walkers, hamiltonian, trial, debug=debug) + + objs = {'trial': trial, + 'hamiltonian': hamiltonian, + 'walkers': walkers, + 'propagator': propagator} + return objs + + +def build_driver_generic_test_instance( + nelec: Tuple[int, int], + nbasis: int, + mu: float, + beta: float, + timestep: float, + nblocks: int, + nwalkers: int = 100, + stack_size: int = 10, + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + stabilize_freq: int = 5, + pop_control_freq: int = 5, + pop_control_method: str = 'pair_branch', + diagonal: bool = False, + complex_integrals: bool = False, + debug: bool = False, + seed: Union[int, None] = None, + verbose: int = 0): + sym = 8 + if complex_integrals: sym = 4 + numpy.random.seed(seed) + + # 1. Generate random integrals. + h1e, chol, _, _ = generate_hamiltonian(nbasis, nelec, cplx=complex_integrals, + sym=sym, tol=1e-10) + + if diagonal: + h1e = numpy.diag(numpy.diag(h1e)) + + # 2. Build Hamiltonian. + hamiltonian = HamGeneric(h1e=numpy.array([h1e, h1e]), + chol=chol.reshape((-1, nbasis**2)).T.copy(), + ecore=0) + + # 3. Build trial. + trial = MeanField(hamiltonian, nelec, beta, timestep) + + # 4. Build Thermal AFQMC driver. + afqmc = ThermalAFQMC.build( + nelec, mu, beta, hamiltonian, trial, nwalkers=nwalkers, + stack_size=stack_size, seed=seed, nblocks=nblocks, timestep=timestep, + stabilize_freq=stabilize_freq, pop_control_freq=pop_control_freq, + pop_control_method=pop_control_method, lowrank=lowrank, + lowrank_thresh=lowrank_thresh, debug=debug, verbose=verbose) + return afqmc + + +def build_ueg_test_case_handlers( + nelec: Tuple[int, int], + rs: float, + ecut: float, + mu: float, + beta: float, + timestep: float, + nwalkers: int = 100, + stack_size: int = 10, + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + propagate: bool = False, + debug: bool = False, + seed: Union[int, None] = None, + verbose: int = 0): + nup, ndown = nelec + ueg_opts = { + "nup": nup, + "ndown": ndown, + "rs": rs, + "ecut": ecut, + "thermal": True, + "write_integrals": False, + "low_rank": lowrank + } + + numpy.random.seed(seed) + + # 1. Generate UEG integrals. + ueg = UEG(ueg_opts, verbose=verbose) + ueg.build(verbose=verbose) + nbasis = ueg.nbasis + nchol = ueg.nchol + + if verbose: + print(f"# nbasis = {nbasis}") + print(f"# nchol = {nchol}") + print(f"# nup = {nup}") + print(f"# ndown = {ndown}") + + h1 = ueg.H1[0] + chol = 2. * ueg.chol_vecs.toarray().copy() + ecore = 0. + + # 2. Build Hamiltonian. + hamiltonian = HamGeneric( + numpy.array([h1, h1], dtype=numpy.complex128), + numpy.array(chol, dtype=numpy.complex128), + ecore, + verbose=verbose) + + # 3. Build trial. + trial = OneBody(hamiltonian, nelec, beta, timestep, verbose=verbose) + + # 4. Build walkers. + walkers = UHFThermalWalkers( + trial, nbasis, nwalkers, stack_size=stack_size, lowrank=lowrank, + lowrank_thresh=lowrank_thresh, verbose=verbose) + + # 5. Build propagator. + propagator = PhaselessGeneric(timestep, mu, lowrank=lowrank, verbose=verbose) + propagator.build(hamiltonian, trial=trial, walkers=walkers, verbose=verbose) + + if propagate: + for t in range(walkers.stack[0].nslice): + propagator.propagate_walkers(walkers, hamiltonian, trial, debug=debug) + + objs = {'trial': trial, + 'hamiltonian': hamiltonian, + 'walkers': walkers, + 'propagator': propagator} + return objs + + +def build_driver_ueg_test_instance( + nelec: Tuple[int, int], + rs: float, + ecut: float, + mu: float, + beta: float, + timestep: float, + nblocks: int, + nwalkers: int = 100, + stack_size: int = 10, + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + stabilize_freq: int = 5, + pop_control_freq: int = 5, + pop_control_method: str = 'pair_branch', + debug: bool = False, + seed: Union[int, None] = None, + verbose: int = 0): + nup, ndown = nelec + ueg_opts = { + "nup": nup, + "ndown": ndown, + "rs": rs, + "ecut": ecut, + "thermal": True, + "write_integrals": False, + "low_rank": lowrank + } + + numpy.random.seed(seed) + + # 1. Generate UEG integrals. + ueg = UEG(ueg_opts, verbose=verbose) + ueg.build(verbose=verbose) + nbasis = ueg.nbasis + nchol = ueg.nchol + + if verbose: + print(f"# nbasis = {nbasis}") + print(f"# nchol = {nchol}") + print(f"# nup = {nup}") + print(f"# ndown = {ndown}") + + h1 = ueg.H1[0] + chol = 2. * ueg.chol_vecs.toarray().copy() + ecore = 0. + + # 2. Build Hamiltonian. + hamiltonian = HamGeneric( + numpy.array([h1, h1], dtype=numpy.complex128), + numpy.array(chol, dtype=numpy.complex128), + ecore, + verbose=verbose) + + # 3. Build trial. + trial = OneBody(hamiltonian, nelec, beta, timestep, verbose=verbose) + + # 4. Build Thermal AFQMC driver. + afqmc = ThermalAFQMC.build( + nelec, mu, beta, hamiltonian, trial, nwalkers=nwalkers, + stack_size=stack_size, seed=seed, nblocks=nblocks, timestep=timestep, + stabilize_freq=stabilize_freq, pop_control_freq=pop_control_freq, + pop_control_method=pop_control_method, lowrank=lowrank, + lowrank_thresh=lowrank_thresh, debug=debug, verbose=verbose) + return afqmc + diff --git a/ipie/addons/thermal/utils/ueg.py b/ipie/addons/thermal/utils/ueg.py new file mode 100644 index 00000000..43286bfa --- /dev/null +++ b/ipie/addons/thermal/utils/ueg.py @@ -0,0 +1,557 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import numpy +import scipy.sparse +from ipie.utils.io import write_qmcpack_sparse + +class UEG: + """UEG system class (integrals read from fcidump) + + Parameters + ---------- + nup : int + Number of up electrons. + + ndown : int + Number of down electrons. + + rs : float + Density parameter. + + ecut : float + Scaled cutoff energy. + + ktwist : :class:`numpy.ndarray` + Twist vector. + + verbose : bool + Print extra information. + + Attributes + ---------- + T : :class:`numpy.ndarray` + One-body part of the Hamiltonian. This is diagonal in plane wave basis. + + ecore : float + Madelung contribution to the total energy. + + h1e_mod : :class:`numpy.ndarray` + Modified one-body Hamiltonian. + + nfields : int + Number of field configurations per walker for back propagation. + + basis : :class:`numpy.ndarray` + Basis vectors within a cutoff. + + kfac : float + Scale factor (2pi/L). + """ + + def __init__(self, options, verbose=False): + if verbose: + print("# Parsing input options.") + + self.name = "UEG" + self.nup = options.get("nup") + self.ndown = options.get("ndown") + self.nelec = (self.nup, self.ndown) + self.rs = options.get("rs") + self.ecut = options.get("ecut") + self.ktwist = numpy.array(options.get("ktwist", [0, 0, 0])).reshape(3) + + self.thermal = options.get("thermal", False) + self._alt_convention = options.get("alt_convention", False) + self.write_ints = options.get("write_integrals", False) + + self.sparse = True + self.control_variate = False + self.diagH1 = True + + # Total # of electrons. + self.ne = self.nup + self.ndown + # Spin polarisation. + self.zeta = (self.nup - self.ndown) / self.ne + # Density. + self.rho = ((4.0 * numpy.pi) / 3.0 * self.rs**3.0) ** (-1.0) + # Box Length. + self.L = self.rs * (4.0 * self.ne * numpy.pi / 3.0) ** (1 / 3.0) + # Volume + self.vol = self.L**3.0 + # k-space grid spacing. + self.kfac = 2 * numpy.pi / self.L + # Fermi Wavevector (infinite system). + self.kf = (3 * (self.zeta + 1) * numpy.pi**2 * self.ne / self.L**3) ** (1 / 3.0) + # Fermi energy (inifinite systems). + self.ef = 0.5 * self.kf**2 + # Core energy. + self.ecore = 0.5 * self.ne * self.madelung() + + if verbose: + if self.thermal: + print("# Thermal UEG activated.") + + print(f"# Number of spin-up electrons: {self.nup:d}") + print(f"# Number of spin-down electrons: {self.ndown:d}") + print(f"# rs: {self.rs:6.4e}") + print(f"# Spin polarisation (zeta): {self.zeta:6.4e}") + print(f"# Electron density (rho): {self.rho:13.8e}") + print(f"# Box Length (L): {self.L:13.8e}") + print(f"# Volume: {self.vol:13.8e}") + print(f"# k-space factor (2pi/L): {self.kfac:13.8e}") + + + def build(self, verbose=False): + # Get plane wave basis vectors and corresponding eigenvalues. + self.sp_eigv, self.basis, self.nmax = self.sp_energies( + self.ktwist, self.kfac, self.ecut) + self.shifted_nmax = 2 * self.nmax + self.imax_sq = numpy.dot(self.basis[-1], self.basis[-1]) + self.create_lookup_table() + + for i, k in enumerate(self.basis): + assert i == self.lookup_basis(k) + + # Number of plane waves. + self.nbasis = len(self.sp_eigv) + self.nactive = self.nbasis + self.ncore = 0 + self.nfv = 0 + self.mo_coeff = None + + # --------------------------------------------------------------------- + T = numpy.diag(self.sp_eigv) + h1e_mod = self.mod_one_body(T) + self.H1 = numpy.array([T, T]) # Making alpha and beta. + self.h1e_mod = numpy.array([h1e_mod, h1e_mod]) + + # --------------------------------------------------------------------- + # Allowed momentum transfers (4*ecut). + _, qvecs, self.qnmax = self.sp_energies(self.ktwist, self.kfac, 4 * self.ecut) + + # Omit Q = 0 term. + self.qvecs = numpy.copy(qvecs[1:]) + self.vqvec = numpy.array([self.vq(self.kfac * q) for q in self.qvecs]) + + # Number of momentum transfer vectors / auxiliary fields. + # Can reduce by symmetry but be stupid for the moment. + self.nchol = len(self.qvecs) + self.nfields = 2 * len(self.qvecs) + self.get_momentum_transfers() + + if verbose: + print(f"# Number of plane waves: {self.nbasis:d}") + print(f"# Number of Cholesky vectors: {self.nchol:d}.") + print(f"# Number of auxiliary fields: {self.nfields:d}.") + print("# Constructing two-body potentials incore.") + + # --------------------------------------------------------------------- + self.chol_vecs, self.iA, self.iB = self.two_body_potentials_incore() + + if self.write_ints: + self.write_integrals() + + if verbose: + print("# Approximate memory required for " + "two-body potentials: {:13.8e} GB.".format((3 * self.iA.nnz * 16 / (1024**3)))) + print("# Finished constructing two-body potentials.") + print("# Finished building UEG object.") + + + def sp_energies(self, ks, kfac, ecut): + """Calculate the allowed kvectors and resulting single particle eigenvalues (basically kinetic energy) + which can fit in the sphere in kspace determined by ecut. + + Parameters + ---------- + kfac : float + kspace grid spacing. + + ecut : float + energy cutoff. + + Returns + ------- + spval : :class:`numpy.ndarray` + Array containing sorted single particle eigenvalues. + + kval : :class:`numpy.ndarray` + Array containing basis vectors, sorted according to their + corresponding single-particle energy. + """ + + # Scaled Units to match with HANDE. + # So ecut is measured in units of 1/kfac^2. + nmax = int(numpy.ceil(numpy.sqrt((2 * ecut)))) + + spval = [] + kval = [] + + for ni in range(-nmax, nmax + 1): + for nj in range(-nmax, nmax + 1): + for nk in range(-nmax, nmax + 1): + spe = 0.5 * (ni**2 + nj**2 + nk**2) + + if spe <= ecut: + kijk = [ni, nj, nk] + + # Reintroduce 2 \pi / L factor. + ek = 0.5 * numpy.dot(numpy.array(kijk) + ks, numpy.array(kijk) + ks) + kval.append(kijk) + spval.append(kfac**2 * ek) + + # Sort the arrays in terms of increasing energy. + spval = numpy.array(spval) + ix = numpy.argsort(spval, kind="mergesort") + spval = spval[ix] + kval = numpy.array(kval)[ix] + return spval, kval, nmax + + + def create_lookup_table(self): + basis_ix = [] + for k in self.basis: + basis_ix.append(self.map_basis_to_index(k)) + + self.lookup = numpy.zeros(max(basis_ix) + 1, dtype=int) + + for i, b in enumerate(basis_ix): + self.lookup[b] = i + + self.max_ix = max(basis_ix) + + + def lookup_basis(self, vec): + if numpy.dot(vec, vec) <= self.imax_sq: + ix = self.map_basis_to_index(vec) + + if ix >= len(self.lookup): + ib = None + + else: + ib = self.lookup[ix] + + return ib + + else: + ib = None + + + def map_basis_to_index(self, k): + return ((k[0] + self.nmax) + + self.shifted_nmax * (k[1] + self.nmax) + + self.shifted_nmax * self.shifted_nmax * (k[2] + self.nmax)) + + + def get_momentum_transfers(self): + """Get arrays of plane wave basis vectors connected by momentum transfers Q.""" + nlimit = self.nup + if self.thermal: + nlimit = self.nbasis + + self.ikpq_i = [] + self.ikpq_kpq = [] + + for iq, q in enumerate(self.qvecs): + idxkpq_list_i = [] + idxkpq_list_kpq = [] + + for i, k in enumerate(self.basis[0:nlimit]): + kpq = k + q + idxkpq = self.lookup_basis(kpq) + + if idxkpq is not None: + idxkpq_list_i += [i] + idxkpq_list_kpq += [idxkpq] + + self.ikpq_i += [idxkpq_list_i] + self.ikpq_kpq += [idxkpq_list_kpq] + + # --------------------------------------------------------------------- + self.ipmq_i = [] + self.ipmq_pmq = [] + + for iq, q in enumerate(self.qvecs): + idxpmq_list_i = [] + idxpmq_list_pmq = [] + + for i, p in enumerate(self.basis[0:nlimit]): + pmq = p - q + idxpmq = self.lookup_basis(pmq) + + if idxpmq is not None: + idxpmq_list_i += [i] + idxpmq_list_pmq += [idxpmq] + + self.ipmq_i += [idxpmq_list_i] + self.ipmq_pmq += [idxpmq_list_pmq] + + for iq, q in enumerate(self.qvecs): + self.ikpq_i[iq] = numpy.array(self.ikpq_i[iq], dtype=numpy.int64) + self.ikpq_kpq[iq] = numpy.array(self.ikpq_kpq[iq], dtype=numpy.int64) + self.ipmq_i[iq] = numpy.array(self.ipmq_i[iq], dtype=numpy.int64) + self.ipmq_pmq[iq] = numpy.array(self.ipmq_pmq[iq], dtype=numpy.int64) + + + def madelung(self): + """Use expression in Schoof et al. (PhysRevLett.115.130402) for the + Madelung contribution to the total energy fitted to L.M. Fraser et al. + Phys. Rev. B 53, 1814. + + Parameters + ---------- + rs : float + Wigner-Seitz radius. + + ne : int + Number of electrons. + + Returns + ------- + v_M: float + Madelung potential (in Hartrees). + """ + c1 = -2.837297 + c2 = (3.0 / (4.0 * numpy.pi)) ** (1.0 / 3.0) + return c1 * c2 / (self.ne ** (1.0 / 3.0) * self.rs) + + + def mod_one_body(self, T): + """Absorb the diagonal term of the two-body Hamiltonian to the one-body term. + Essentially adding the third term in Eq.(11b) of Phys. Rev. B 75, 245123. + + Parameters + ---------- + T : float + one-body Hamiltonian (i.e. kinetic energy) + + Returns + ------- + h1e_mod: float + modified one-body Hamiltonian + """ + h1e_mod = numpy.copy(T) + + fac = 1.0 / (2.0 * self.vol) + for i, ki in enumerate(self.basis): + for j, kj in enumerate(self.basis): + if i != j: + q = self.kfac * (ki - kj) + h1e_mod[i, i] = h1e_mod[i, i] - fac * self.vq(q) + + return h1e_mod + + + def vq(self, q): + """The typical 3D Coulomb kernel + + Parameters + ---------- + q : float + a plane-wave vector + + Returns + ------- + v_M: float + 3D Coulomb kernel (in Hartrees) + """ + return 4 * numpy.pi / numpy.dot(q, q) + + + def scaled_density_operator_incore(self, transpose): + """Density operator as defined in Eq.(6) of PRB(75)245123 + + Parameters + ---------- + q : float + a plane-wave vector + + Returns + ------- + rho_q: float + density operator + """ + rho_ikpq_i = [] + rho_ikpq_kpq = [] + + for iq, q in enumerate(self.qvecs): + idxkpq_list_i = [] + idxkpq_list_kpq = [] + + for i, k in enumerate(self.basis): + kpq = k + q + idxkpq = self.lookup_basis(kpq) + + if idxkpq is not None: + idxkpq_list_i += [i] + idxkpq_list_kpq += [idxkpq] + + rho_ikpq_i += [idxkpq_list_i] + rho_ikpq_kpq += [idxkpq_list_kpq] + + for iq, q in enumerate(self.qvecs): + rho_ikpq_i[iq] = numpy.array(rho_ikpq_i[iq], dtype=numpy.int64) + rho_ikpq_kpq[iq] = numpy.array(rho_ikpq_kpq[iq], dtype=numpy.int64) + + nq = len(self.qvecs) + nnz = 0 + for iq in range(nq): + nnz += rho_ikpq_kpq[iq].shape[0] + + col_index = [] + row_index = [] + values = [] + + if transpose: + for iq in range(nq): + qscaled = self.kfac * self.qvecs[iq] + # Due to the HS transformation, we have to do pi / 2*vol as opposed to 2*pi / vol + piovol = numpy.pi / (self.vol) + factor = (piovol / numpy.dot(qscaled, qscaled)) ** 0.5 + + for innz, kpq in enumerate(rho_ikpq_kpq[iq]): + row_index += [rho_ikpq_kpq[iq][innz] + rho_ikpq_i[iq][innz] * self.nbasis] + col_index += [iq] + values += [factor] + + else: + for iq in range(nq): + qscaled = self.kfac * self.qvecs[iq] + # Due to the HS transformation, we have to do pi / 2*vol as opposed to 2*pi / vol + piovol = numpy.pi / (self.vol) + factor = (piovol / numpy.dot(qscaled, qscaled)) ** 0.5 + + for innz, kpq in enumerate(rho_ikpq_kpq[iq]): + row_index += [rho_ikpq_kpq[iq][innz] * self.nbasis + rho_ikpq_i[iq][innz]] + col_index += [iq] + values += [factor] + + rho_q = scipy.sparse.csc_matrix( + (values, (row_index, col_index)), + shape=(self.nbasis * self.nbasis, nq), + dtype=numpy.complex128) + return rho_q + + + def two_body_potentials_incore(self): + """Calculate A and B of Eq.(13) of PRB(75)245123 for a given plane-wave vector q + + Returns + ------- + iA : numpy array + Eq.(13a) + + iB : numpy array + Eq.(13b) + """ + rho_q = self.scaled_density_operator_incore(False) + rho_qH = self.scaled_density_operator_incore(True) + iA = 1j * (rho_q + rho_qH) + iB = -(rho_q - rho_qH) + return (rho_q, iA, iB) + + + def hijkl(self, i, j, k, l): + """Compute = (ik|jl) = 1/Omega * 4pi/(kk-ki)**2 + + Checks for momentum conservation k_i + k_j = k_k + k_k, or + k_k - k_i = k_j - k_l. + + Parameters + ---------- + i, j, k, l : int + Orbital indices for integral (ik|jl) = . + + Returns + ------- + integral : float + (ik|jl) + """ + q1 = self.basis[k] - self.basis[i] + q2 = self.basis[j] - self.basis[l] + + if numpy.dot(q1, q1) > 1e-12 and numpy.dot(q1 - q2, q1 - q2) < 1e-12: + return 1.0 / self.vol * self.vq(self.kfac * q1) + + else: + return 0.0 + + + def compute_real_transformation(self): + U22 = numpy.zeros((2, 2), dtype=numpy.complex128) + U22[0, 0] = 1.0 / numpy.sqrt(2.0) + U22[0, 1] = 1.0 / numpy.sqrt(2.0) + U22[1, 0] = -1.0j / numpy.sqrt(2.0) + U22[1, 1] = 1.0j / numpy.sqrt(2.0) + + U = numpy.zeros((self.nbasis, self.nbasis), dtype=numpy.complex128) + + for i, b in enumerate(self.basis): + if numpy.sum(b * b) == 0: + U[i, i] = 1.0 + + else: + mb = -b + diff = numpy.einsum("ij->i", (self.basis - mb) ** 2) + idx = numpy.argwhere(diff == 0) + assert idx.ravel().shape[0] == 1 + + if i < idx: + idx = idx.ravel()[0] + U[i, i] = U22[0, 0] + U[i, idx] = U22[0, 1] + U[idx, i] = U22[1, 0] + U[idx, idx] = U22[1, 1] + + else: + continue + + U = U.T.copy() + return U + + + def eri_4(self): + eri_chol = 4 * self.chol_vecs.dot(self.chol_vecs.T) + eri_chol = ( + eri_chol.toarray().reshape((self.nbasis, self.nbasis, self.nbasis, self.nbasis)).real) + eri_chol = eri_chol.transpose(0, 1, 3, 2) + return eri_chol + + + def eri_8(self): + """Compute 8-fold symmetric integrals. Useful for running standard + quantum chemistry methods,""" + eri = self.eri_4() + U = self.compute_real_transformation() + eri0 = numpy.einsum("mp,mnls->pnls", U.conj(), eri, optimize=True) + eri1 = numpy.einsum("nq,pnls->pqls", U, eri0, optimize=True) + eri2 = numpy.einsum("lr,pqls->pqrs", U.conj(), eri1, optimize=True) + eri3 = numpy.einsum("st,pqrs->pqrt", U, eri2, optimize=True).real + return eri3 + + + def write_integrals(self, filename="ueg_integrals.h5"): + write_qmcpack_sparse( + self.H1[0], + 2 * self.chol_vecs.toarray(), + self.nelec, + self.nbasis, + enuc=0.0, + filename=filename) + diff --git a/ipie/addons/thermal/walkers/__init__.py b/ipie/addons/thermal/walkers/__init__.py new file mode 100644 index 00000000..f91ef518 --- /dev/null +++ b/ipie/addons/thermal/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. +# +# Authors: Fionn Malone +# Joonho Lee +# diff --git a/ipie/addons/thermal/walkers/pop_controller.py b/ipie/addons/thermal/walkers/pop_controller.py new file mode 100644 index 00000000..a0aed19a --- /dev/null +++ b/ipie/addons/thermal/walkers/pop_controller.py @@ -0,0 +1,440 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import numpy + +from ipie.config import MPI +from ipie.utils.backend import arraylib as xp +from ipie.walkers.pop_controller import PopController, PopControllerTimer + + +class ThermalPopController(PopController): + def __init__( + self, + num_walkers_local, + num_steps, + mpi_handler=None, + pop_control_method="pair_branch", + min_weight=0.1, + max_weight=4, + reconfiguration_freq=50, + verbose=False, + ): + super().__init__( + num_walkers_local, num_steps, mpi_handler, pop_control_method, + min_weight, max_weight, reconfiguration_freq, verbose) + + def pop_control(self, walkers, comm): + self.timer.start_time() + if self.ntot_walkers == 1: + return + weights = numpy.abs(xp.array(walkers.weight)) + global_weights = numpy.empty(len(weights) * comm.size) + self.timer.add_non_communication() + self.timer.start_time() + if self.method == "comb": + comm.Allgather(weights, global_weights) + total_weight = sum(global_weights) + else: + sum_weights = numpy.sum(weights) + total_weight = numpy.empty(1, dtype=numpy.float64) + if hasattr(sum_weights, "get"): + sum_weights = sum_weights.get() + comm.Reduce(sum_weights, total_weight, op=MPI.SUM, root=0) + comm.Bcast(total_weight, root=0) + total_weight = total_weight[0] + + self.timer.add_communication() + self.timer.start_time() + + # Rescale weights to combat exponential decay/growth. + scale = total_weight / self.target_weight + if total_weight < 1e-8: + if comm.rank == 0: + print(f"# Warning: Total weight is {total_weight:13.8e}") + print("# Something is seriously wrong.") + raise ValueError + self.total_weight = total_weight + # Todo: Just standardise information we want to send between routines. + walkers.unscaled_weight = walkers.weight + walkers.weight = walkers.weight / scale + if self.method == "comb": + global_weights = global_weights / scale + self.timer.add_non_communication() + comb(walkers, comm, global_weights, self.target_weight, self.timer) + elif self.method == "pair_branch": + pair_branch(walkers, comm, self.max_weight, self.min_weight, self.timer) + elif self.method == "stochastic_reconfiguration": + self.reconfiguration_counter += 1 + if self.reconfiguration_counter % self.reconfiguration_freq == 0: + stochastic_reconfiguration(walkers, comm, self.timer) + self.reconfiguration_counter = 0 + else: + if comm.rank == 0: + print("Unknown population control method.") + + +def get_buffer(walkers, iw): + """Get iw-th walker buffer for MPI communication + iw : int + the walker index of interest + Returns + ------- + buff : dict + Relevant walker information for population control. + """ + s = 0 + buff = xp.zeros(walkers.buff_size, dtype=numpy.complex128) + for d in walkers.buff_names: + data = walkers.__dict__[d] + if (data is None) or isinstance(data, (int, float, complex, numpy.float64, numpy.complex128)): + continue + assert data.size % walkers.nwalkers == 0 # Only walker-specific data is being communicated + if isinstance(data[iw], (xp.ndarray)): + buff[s : s + data[iw].size] = xp.array(data[iw].ravel()) + s += data[iw].size + elif isinstance(data[iw], list): # when data is list + for l in data[iw]: + if isinstance(l, (xp.ndarray)): + buff[s : s + l.size] = xp.array(l.ravel()) + s += l.size + elif isinstance(l, (int, float, complex, numpy.float64, numpy.complex128)): + buff[s : s + 1] = l + s += 1 + else: + buff[s : s + 1] = xp.array(data[iw]) + s += 1 + + stack_buff = walkers.stack[iw].get_buffer() + buff = numpy.concatenate((buff, stack_buff)) + return buff + + +def set_buffer(walkers, iw, buff): + """Set walker buffer following MPI communication + Parameters + ------- + buff : dict + Relevant walker information for population control. + """ + s = 0 + for d in walkers.buff_names: + data = walkers.__dict__[d] + if (data is None) or isinstance(data, (int, float, complex, numpy.float64, numpy.complex128)): + continue + assert data.size % walkers.nwalkers == 0 # Only walker-specific data is being communicated + if isinstance(data[iw], xp.ndarray): + walkers.__dict__[d][iw] = xp.array( + buff[s : s + data[iw].size].reshape(data[iw].shape).copy()) + s += data[iw].size + elif isinstance(data[iw], list): + for ix, l in enumerate(data[iw]): + if isinstance(l, (xp.ndarray)): + walkers.__dict__[d][iw][ix] = xp.array( + buff[s : s + l.size].reshape(l.shape).copy()) + s += l.size + elif isinstance(l, (int, float, complex)): + walkers.__dict__[d][iw][ix] = buff[s] + s += 1 + else: + if isinstance(walkers.__dict__[d][iw], (int, numpy.int64)): + walkers.__dict__[d][iw] = int(buff[s].real) + elif isinstance(walkers.__dict__[d][iw], (float, numpy.float64)): + walkers.__dict__[d][iw] = buff[s].real + else: + walkers.__dict__[d][iw] = buff[s] + s += 1 + + walkers.stack[iw].set_buffer(buff[walkers.buff_size:]) + + +def comb(walkers, comm, weights, target_weight, timer=PopControllerTimer()): + """Apply the comb method of population control / branching. + + See Booth & Gubernatis PRE 80, 046704 (2009). + + Parameters + ---------- + comm : MPI communicator + """ + # Need make a copy to since the elements in psi are only references to + # walker objects in memory. We don't want future changes in a given + # element of psi having unintended consequences. + # todo : add phase to walker for free projection + timer.start_time() + if comm.rank == 0: + parent_ix = numpy.zeros(len(weights), dtype="i") + else: + parent_ix = numpy.empty(len(weights), dtype="i") + if comm.rank == 0: + total_weight = sum(weights) + cprobs = numpy.cumsum(weights) + r = numpy.random.random() + comb = [(i + r) * (total_weight / target_weight) for i in range(target_weight)] + iw = 0 + ic = 0 + while ic < len(comb): + if comb[ic] < cprobs[iw]: + parent_ix[iw] += 1 + ic += 1 + else: + iw += 1 + data = {"ix": parent_ix} + else: + data = None + + timer.add_non_communication() + + timer.start_time() + data = comm.bcast(data, root=0) + timer.add_communication() + timer.start_time() + parent_ix = data["ix"] + # Keep total weight saved for capping purposes. + # where returns a tuple (array,), selecting first element. + kill = numpy.where(parent_ix == 0)[0] + clone = numpy.where(parent_ix > 1)[0] + reqs = [] + # First initiate non-blocking sends of walkers. + timer.add_non_communication() + timer.start_time() + comm.barrier() + timer.add_communication() + for i, (c, k) in enumerate(zip(clone, kill)): + # Sending from current processor? + if c // walkers.nwalkers == comm.rank: + timer.start_time() + # Location of walker to clone in local list. + clone_pos = c % walkers.nwalkers + # copying walker data to intermediate buffer to avoid issues + # with accessing walker data during send. Might not be + # necessary. + dest_proc = k // walkers.nwalkers + buff = get_buffer(walkers, clone_pos) + timer.add_non_communication() + timer.start_time() + reqs.append(comm.Isend(buff, dest=dest_proc, tag=i)) + timer.add_send_time() + # Now receive walkers on processors where walkers are to be killed. + for i, (c, k) in enumerate(zip(clone, kill)): + # Receiving to current processor? + if k // walkers.nwalkers == comm.rank: + timer.start_time() + # Processor we are receiving from. + source_proc = c // walkers.nwalkers + # Location of walker to kill in local list of walkers. + kill_pos = k % walkers.nwalkers + buffer = walkers.walker_buffer + buffer = numpy.concatenate((walkers.walker_buffer, walkers.stack[0].stack_buffer)) + timer.add_non_communication() + timer.start_time() + comm.Recv(buffer, source=source_proc, tag=i) + # with h5py.File('walkers_recv.h5', 'w') as fh5: + # fh5['walk_{}'.format(k)] = walkers.walker_buffer.copy() + timer.add_recv_time() + timer.start_time() + set_buffer(walkers, kill_pos, buffer) + timer.add_non_communication() + # with h5py.File('after_{}.h5'.format(comm.rank), 'a') as fh5: + # fh5['walker_{}_{}_{}'.format(c,k,comm.rank)] = walkers.walkers[kill_pos].get_buffer() + timer.start_time() + # Complete non-blocking send. + for rs in reqs: + rs.wait() + # Necessary? + # if len(kill) > 0 or len(clone) > 0: + # sys.exit() + comm.Barrier() + timer.add_communication() + # Reset walker weight. + # TODO: check this. + # for w in walkers.walkers: + # w.weight = 1.0 + timer.start_time() + walkers.weight.fill(1.0) + timer.add_non_communication() + + +def pair_branch(walkers, comm, max_weight, min_weight, timer=PopControllerTimer()): + timer.start_time() + walker_info_0 = xp.array(xp.abs(walkers.weight)) + timer.add_non_communication() + + timer.start_time() + glob_inf = None + glob_inf_0 = None + glob_inf_1 = None + glob_inf_2 = None + glob_inf_3 = None + if comm.rank == 0: + glob_inf_0 = numpy.empty([comm.size, walkers.nwalkers], dtype=numpy.float64) + glob_inf_1 = numpy.empty([comm.size, walkers.nwalkers], dtype=numpy.int64) + glob_inf_1.fill(1) + glob_inf_2 = numpy.array( + [[r for i in range(walkers.nwalkers)] for r in range(comm.size)], + dtype=numpy.int64) + glob_inf_3 = numpy.array( + [[r for i in range(walkers.nwalkers)] for r in range(comm.size)], + dtype=numpy.int64) + + timer.add_non_communication() + + timer.start_time() + if hasattr(walker_info_0, "get"): + walker_info_0 = walker_info_0.get() + comm.Gather( + walker_info_0, glob_inf_0, root=0 + ) # gather |w_i| from all processors (comm.size x nwalkers) + timer.add_communication() + + # Want same random number seed used on all processors + timer.start_time() + if comm.rank == 0: + # Rescale weights. + glob_inf = numpy.zeros((walkers.nwalkers * comm.size, 4), dtype=numpy.float64) + glob_inf[:, 0] = glob_inf_0.ravel() # contains walker |w_i| + glob_inf[:, 1] = glob_inf_1.ravel() # all initialized to 1 when it becomes 2 then it will be "branched" + glob_inf[:, 2] = glob_inf_2.ravel() # contain processor+walker indices (initial) (i.e., where walkers live) + glob_inf[:, 3] = glob_inf_3.ravel() # contain processor+walker indices (final) (i.e., where walkers live) + sort = numpy.argsort(glob_inf[:, 0], kind="mergesort") + isort = numpy.argsort(sort, kind="mergesort") + glob_inf = glob_inf[sort] + s = 0 + e = len(glob_inf) - 1 + tags = [] + # go through walkers pair-wise + while s < e: + if glob_inf[s][0] < min_weight or glob_inf[e][0] > max_weight: + # sum of paired walker weights + wab = glob_inf[s][0] + glob_inf[e][0] + r = numpy.random.rand() + if r < glob_inf[e][0] / wab: + # clone large weight walker + glob_inf[e][0] = 0.5 * wab + glob_inf[e][1] = 2 + # Processor we will send duplicated walker to + glob_inf[e][3] = glob_inf[s][2] + send = glob_inf[s][2] + # Kill small weight walker + glob_inf[s][0] = 0.0 + glob_inf[s][1] = 0 + glob_inf[s][3] = glob_inf[e][2] + else: + # clone small weight walker + glob_inf[s][0] = 0.5 * wab + glob_inf[s][1] = 2 + # Processor we will send duplicated walker to + glob_inf[s][3] = glob_inf[e][2] + send = glob_inf[e][2] + # Kill small weight walker + glob_inf[e][0] = 0.0 + glob_inf[e][1] = 0 + glob_inf[e][3] = glob_inf[s][2] + tags.append([send]) + s += 1 + e -= 1 + else: + break + nw = walkers.nwalkers + glob_inf = glob_inf[isort].reshape((comm.size, nw, 4)) + else: + data = None + glob_inf = None + timer.add_non_communication() + timer.start_time() + + data = numpy.empty([walkers.nwalkers, 4], dtype=numpy.float64) + # 0 = weight, 1 = status (live, branched, die), 2 = initial index, 3 = final index + comm.Scatter(glob_inf, data, root=0) + + timer.add_communication() + # Keep total weight saved for capping purposes. + reqs = [] + for iw, walker in enumerate(data): + if walker[1] > 1: + timer.start_time() + tag = comm.rank * walkers.nwalkers + walker[3] + walkers.weight[iw] = walker[0] + buff = get_buffer(walkers, iw) + timer.add_non_communication() + timer.start_time() + reqs.append(comm.Isend(buff, dest=int(round(walker[3])), tag=tag)) + timer.add_send_time() + for iw, walker in enumerate(data): + if walker[1] == 0: + timer.start_time() + tag = walker[3] * walkers.nwalkers + comm.rank + buffer = walkers.walker_buffer + buffer = numpy.concatenate((walkers.walker_buffer, walkers.stack[0].stack_buffer)) + timer.add_non_communication() + timer.start_time() + comm.Recv(buffer, source=int(round(walker[3])), tag=tag) + timer.add_recv_time() + timer.start_time() + set_buffer(walkers, iw, buffer) + timer.add_non_communication() + timer.start_time() + for r in reqs: + r.wait() + timer.add_communication() + + +def stochastic_reconfiguration(walkers, comm, timer=PopControllerTimer()): + # gather all walker information on the root + timer.start_time() + nwalkers = walkers.nwalkers + local_buffer = xp.array([get_buffer(walkers, i) for i in range(nwalkers)]) + walker_len = local_buffer[0].shape[0] + global_buffer = None + if comm.rank == 0: + global_buffer = numpy.zeros((comm.size, nwalkers, walker_len), dtype=numpy.complex128) + timer.add_non_communication() + + timer.start_time() + comm.Gather(local_buffer, global_buffer, root=0) + timer.add_communication() + + # perform sr on the root + new_global_buffer = None + timer.start_time() + if comm.rank == 0: + new_global_buffer = numpy.zeros((comm.size, nwalkers, walker_len), dtype=numpy.complex128) + cumulative_weights = numpy.cumsum(abs(global_buffer[:, :, 0])) + total_weight = cumulative_weights[-1] + new_average_weight = total_weight / nwalkers / comm.size + zeta = numpy.random.rand() + for i in range(comm.size * nwalkers): + z = (i + zeta) / nwalkers / comm.size + new_i = numpy.searchsorted(cumulative_weights, z * total_weight) + new_global_buffer[i // nwalkers, i % nwalkers] = global_buffer[ + new_i // nwalkers, new_i % nwalkers + ] + new_global_buffer[i // nwalkers, i % nwalkers, 0] = new_average_weight + + timer.add_non_communication() + + # distribute information of newly selected walkers + timer.start_time() + comm.Scatter(new_global_buffer, local_buffer, root=0) + timer.add_communication() + + # set walkers using distributed information + timer.start_time() + for i in range(nwalkers): + set_buffer(walkers, i, local_buffer[i]) + timer.add_non_communication() diff --git a/ipie/addons/thermal/walkers/stack.py b/ipie/addons/thermal/walkers/stack.py new file mode 100644 index 00000000..af048b77 --- /dev/null +++ b/ipie/addons/thermal/walkers/stack.py @@ -0,0 +1,411 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +# TODO: Incorporate the `stack` buffer in the `walkers` buffer for MPI. See: +# https://github.com/JoonhoLee-Group/ipie/issues/301 + +import numpy +import scipy.linalg + +from ipie.utils.misc import get_numeric_names + +class PropagatorStack: + def __init__( + self, + stack_size, + nslice, + nbasis, + dtype, + BT=None, + BTinv=None, + diagonal=False, + averaging=False, + lowrank=True, + thresh=1e-6, + ): + self.time_slice = 0 + self.stack_size = stack_size + self.nslice = nslice + self.nstack = self.nslice // self.stack_size + self.nbasis = nbasis + self.diagonal_trial = diagonal + self.averaging = averaging + self.thresh = thresh + self.lowrank = lowrank + self.ovlp = numpy.asarray([1.0, 1.0]) + self.reortho = 1 + + if self.lowrank: + assert diagonal + + if self.nstack * self.stack_size < self.nslice: + print("stack_size must divide the total path length") + assert self.nstack * self.stack_size == self.nslice + + self.dtype = dtype + self.BT = BT + self.BTinv = BTinv + self.counter = 0 + self.block = 0 + + self.stack = numpy.zeros((self.nstack, 2, nbasis, nbasis), dtype=dtype) + self.left = numpy.zeros((self.nstack, 2, nbasis, nbasis), dtype=dtype) + self.right = numpy.zeros((self.nstack, 2, nbasis, nbasis), dtype=dtype) + + self.G = numpy.asarray([numpy.eye(self.nbasis, dtype=dtype), # Ga + numpy.eye(self.nbasis, dtype=dtype)]) # Gb + + if self.lowrank: + self.update_new = self.update_low_rank + else: + self.update_new = self.update_full_rank + + # Global block matrix + if self.lowrank: + self.Ql = numpy.zeros((2, nbasis, nbasis), dtype=dtype) + self.Dl = numpy.zeros((2, nbasis), dtype=dtype) + self.Tl = numpy.zeros((2, nbasis, nbasis), dtype=dtype) + + self.Qr = numpy.zeros((2, nbasis, nbasis), dtype=dtype) + self.Dr = numpy.zeros((2, nbasis), dtype=dtype) + self.Tr = numpy.zeros((2, nbasis, nbasis), dtype=dtype) + + self.CT = numpy.zeros((2, nbasis, nbasis), dtype=dtype) + self.theta = numpy.zeros((2, nbasis, nbasis), dtype=dtype) + self.mT = nbasis + + self.buff_names, self.buff_size = get_numeric_names(self.__dict__) + self.stack_buffer = numpy.zeros(self.buff_size, dtype=numpy.complex128) + + # Set all entries to be the identity matrix + self.reset() + + def get(self, ix): + return self.stack[ix] + + def get_buffer(self): + s = 0 + buff = numpy.zeros(self.buff_size, dtype=numpy.complex128) + for d in self.buff_names: + data = self.__dict__[d] + if isinstance(data, (numpy.ndarray)): + buff[s : s + data.size] = data.ravel() + s += data.size + else: + buff[s : s + 1] = data + s += 1 + return buff + + def set_buffer(self, buff): + s = 0 + for d in self.buff_names: + data = self.__dict__[d] + if isinstance(data, numpy.ndarray): + self.__dict__[d] = buff[s : s + data.size].reshape(data.shape).copy() + dsize = data.size + else: + if isinstance(self.__dict__[d], int): + self.__dict__[d] = int(buff[s].real) + elif isinstance(self.__dict__[d], float): + self.__dict__[d] = float(buff[s].real) + else: + self.__dict__[d] = buff[s] + dsize = 1 + s += dsize + + def set_all(self, BT): + # Diagonal = True assumes BT is diagonal and left is also diagonal + if self.diagonal_trial: + for i in range(0, self.nslice): + ix = i // self.stack_size # bin index + # Commenting out these two. It is only useful for Hubbard + self.left[ix, 0] = numpy.diag( + numpy.multiply(BT[0].diagonal(), self.left[ix, 0].diagonal()) + ) + self.left[ix, 1] = numpy.diag( + numpy.multiply(BT[1].diagonal(), self.left[ix, 1].diagonal()) + ) + self.stack[ix, 0] = self.left[ix, 0].copy() + self.stack[ix, 1] = self.left[ix, 1].copy() + else: + for i in range(0, self.nslice): + ix = i // self.stack_size # bin index + self.left[ix, 0] = numpy.dot(BT[0], self.left[ix, 0]) + self.left[ix, 1] = numpy.dot(BT[1], self.left[ix, 1]) + self.stack[ix, 0] = self.left[ix, 0].copy() + self.stack[ix, 1] = self.left[ix, 1].copy() + + if self.lowrank: + self.initialize_left() + for s in [0, 1]: + self.Qr[s] = numpy.identity(self.nbasis, dtype=self.dtype) + self.Dr[s] = numpy.ones(self.nbasis, dtype=self.dtype) + self.Tr[s] = numpy.identity(self.nbasis, dtype=self.dtype) + + def reset(self): + self.time_slice = 0 + self.block = 0 + for i in range(0, self.nstack): + self.stack[i, 0] = numpy.identity(self.nbasis, dtype=self.dtype) + self.stack[i, 1] = numpy.identity(self.nbasis, dtype=self.dtype) + self.right[i, 0] = numpy.identity(self.nbasis, dtype=self.dtype) + self.right[i, 1] = numpy.identity(self.nbasis, dtype=self.dtype) + self.left[i, 0] = numpy.identity(self.nbasis, dtype=self.dtype) + self.left[i, 1] = numpy.identity(self.nbasis, dtype=self.dtype) + + if self.lowrank: + for s in [0, 1]: + self.Qr[s] = numpy.identity(self.nbasis, dtype=self.dtype) + self.Dr[s] = numpy.ones(self.nbasis, dtype=self.dtype) + self.Tr[s] = numpy.identity(self.nbasis, dtype=self.dtype) + + # Form BT product for i = 1, ..., nslices - 1 (i.e., skip i = 0) + # \TODO add non-diagonal version of this + def initialize_left(self): + assert self.diagonal_trial + for spin in [0, 1]: + # We will assume that B matrices are all diagonal for left.... + # B = self.stack[1] + B = self.stack[0] + self.Dl[spin] = B[spin].diagonal() + self.Ql[spin] = numpy.identity(B[spin].shape[0]) + self.Tl[spin] = numpy.identity(B[spin].shape[0]) + + # for ix in range(2, self.nstack): + for ix in range(1, self.nstack): + B = self.stack[ix] + C2 = numpy.einsum("ii,i->i", B[spin], self.Dl[spin]) + self.Dl[spin] = C2 + + def update(self, B): + if self.counter == 0: + self.stack[self.block, 0] = numpy.identity(B.shape[-1], dtype=B.dtype) + self.stack[self.block, 1] = numpy.identity(B.shape[-1], dtype=B.dtype) + self.stack[self.block, 0] = B[0].dot(self.stack[self.block, 0]) + self.stack[self.block, 1] = B[1].dot(self.stack[self.block, 1]) + self.time_slice += 1 + self.block = self.time_slice // self.stack_size + self.counter = (self.counter + 1) % self.stack_size + + def update_full_rank(self, B): + # Diagonal = True assumes BT is diagonal and left is also diagonal + if self.counter == 0: + self.right[self.block, 0] = numpy.identity(B.shape[-1], dtype=B.dtype) + self.right[self.block, 1] = numpy.identity(B.shape[-1], dtype=B.dtype) + + if self.diagonal_trial: + self.left[self.block, 0] = numpy.diag( + numpy.multiply(self.left[self.block, 0].diagonal(), self.BTinv[0].diagonal()) + ) + self.left[self.block, 1] = numpy.diag( + numpy.multiply(self.left[self.block, 1].diagonal(), self.BTinv[1].diagonal()) + ) + else: + self.left[self.block, 0] = self.left[self.block, 0].dot(self.BTinv[0]) + self.left[self.block, 1] = self.left[self.block, 1].dot(self.BTinv[1]) + + self.right[self.block, 0] = B[0].dot(self.right[self.block, 0]) + self.right[self.block, 1] = B[1].dot(self.right[self.block, 1]) + + if self.diagonal_trial: + self.stack[self.block, 0] = numpy.einsum( + "ii,ij->ij", self.left[self.block, 0], self.right[self.block, 0] + ) + self.stack[self.block, 1] = numpy.einsum( + "ii,ij->ij", self.left[self.block, 1], self.right[self.block, 1] + ) + else: + self.stack[self.block, 0] = self.left[self.block, 0].dot(self.right[self.block, 0]) + self.stack[self.block, 1] = self.left[self.block, 1].dot(self.right[self.block, 1]) + + self.time_slice += 1 # Count the time slice + self.block = self.time_slice // self.stack_size # Move to the next block if necessary + self.counter = (self.counter + 1) % self.stack_size # Counting within a stack + + def update_low_rank(self, B): + assert not self.averaging + # Diagonal = True assumes BT is diagonal and left is also diagonal + assert self.diagonal_trial + + if self.counter == 0: + for s in [0, 1]: + self.Tl[s] = self.left[self.block, s] + + mR = B.shape[-1] # initial mR + mL = B.shape[-1] # initial mR + mT = B.shape[-1] # initial mR + next_block = (self.time_slice + 1) // self.stack_size # move to the next block if necessary + # print("next_block", next_block) + # print("self.block", self.block) + if next_block > self.block: # Do QR and update here? + for s in [0, 1]: + mR = len(self.Dr[s][numpy.abs(self.Dr[s]) > self.thresh]) + self.Dl[s] = numpy.einsum("i,ii->i", self.Dl[s], self.BTinv[s]) + mL = len(self.Dl[s][numpy.abs(self.Dl[s]) > self.thresh]) + + self.Qr[s][:, :mR] = B[s].dot(self.Qr[s][:, :mR]) # N x mR + self.Qr[s][:, mR:] = 0.0 + + Ccr = numpy.einsum("ij,j->ij", self.Qr[s][:, :mR], self.Dr[s][:mR]) # N x mR + (Qlcr, Rlcr, Plcr) = scipy.linalg.qr(Ccr, pivoting=True, check_finite=False) + Dlcr = Rlcr[:mR, :mR].diagonal() # mR + + self.Dr[s][:mR] = Dlcr + self.Dr[s][mR:] = 0.0 + self.Qr[s] = Qlcr + + Dinv = 1.0 / Dlcr # mR + tmp = numpy.einsum("i,ij->ij", Dinv[:mR], Rlcr[:mR, :mR]) # mR, mR x mR -> mR x mR + tmp[:, Plcr] = tmp[:, range(mR)] + Tlcr = numpy.dot(tmp, self.Tr[s][:mR, :]) # mR x N + + self.Tr[s][:mR, :] = Tlcr + + # assume left stack is all diagonal (i.e., QDT = diagonal -> Q and T are identity) + Clcr = numpy.einsum( + "i,ij->ij", + self.Dl[s][:mL], + numpy.einsum("ij,j->ij", Qlcr[:mL, :mR], Dlcr[:mR]), + ) # mL x mR + + (Qlcr, Rlcr, Plcr) = scipy.linalg.qr( + Clcr, pivoting=True, check_finite=False + ) # mL x mL, min(mL,mR) x min(mL,mR), mR x mR + Dlcr = Rlcr.diagonal()[: min(mL, mR)] + Dinv = 1.0 / Dlcr + + mT = len(Dlcr[numpy.abs(Dlcr) > self.thresh]) + + assert mT <= mL and mT <= mR + + tmp = numpy.einsum("i,ij->ij", Dinv[:mT], Rlcr[:mT, :]) + tmp[:, Plcr] = tmp[:, range(mR)] # mT x mR + Tlcr = numpy.dot(tmp, Tlcr) # mT x N + + Db = numpy.zeros(mT, B[s].dtype) + Ds = numpy.zeros(mT, B[s].dtype) + for i in range(mT): + absDlcr = abs(Dlcr[i]) + if absDlcr > 1.0: + Db[i] = 1.0 / absDlcr + Ds[i] = numpy.sign(Dlcr[i]) + else: + Db[i] = 1.0 + Ds[i] = Dlcr[i] + Dbinv = 1.0 / Db + + TQ = Tlcr[:, :mL].dot(Qlcr[:mL, :mT]) # mT x mT + TQinv = scipy.linalg.inv(TQ, check_finite=False) + tmp = numpy.einsum("ij,j->ij", TQinv, Db) + numpy.diag(Ds) # mT x mT + + M = numpy.einsum("ij,j->ij", tmp, Dbinv).dot(TQ) + # self.ovlp[s] = 1.0 / scipy.linalg.det(M, check_finite=False) + self.ovlp[s] = scipy.linalg.det(M, check_finite=False) + + tmp = scipy.linalg.inv(tmp, check_finite=False) + A = numpy.einsum("i,ij->ij", Db, tmp.dot(TQinv)) # mT x mT + Qlcr_pad = numpy.zeros((self.nbasis, self.nbasis), dtype=B[s].dtype) + Qlcr_pad[:mL, :mT] = Qlcr[:, :mT] + + # self.G[s] = numpy.eye(self.nbasis, dtype=B[s].dtype) - Qlcr_pad[:,:mT].dot(numpy.diag(Dlcr[:mT])).dot(A).dot(Tlcr) + + self.CT[s][:, :] = 0.0 + self.CT[s][:, :mT] = (A.dot(Tlcr)).T.conj() + self.theta[s][:, :] = 0.0 + self.theta[s][:mT, :] = Qlcr_pad[:, :mT].dot(numpy.diag(Dlcr[:mT])).T + # self.G[s] = numpy.eye(self.nbasis, dtype=B[s].dtype) - self.CT[s][:,:mT].dot(self.theta[s][:mT,:]) + self.G[s] = numpy.eye(self.nbasis, dtype=B[s].dtype) -\ + self.theta[s][:mT, :].T.dot(self.CT[s][:, :mT].T.conj()) + # self.CT[s][:,:mT] = self.CT[s][:,:mT].conj() + + # print("# mL, mR, mT = {}, {}, {}".format(mL, mR, mT)) + else: # don't do QR and just update + for s in [0, 1]: + mR = len(self.Dr[s][numpy.abs(self.Dr[s]) > self.thresh]) + + self.Dl[s] = numpy.einsum("i,ii->i", self.Dl[s], self.BTinv[s]) + mL = len(self.Dl[s][numpy.abs(self.Dl[s]) > self.thresh]) + + self.Qr[s][:, :mR] = B[s].dot(self.Qr[s][:, :mR]) # N x mR + self.Qr[s][:, mR:] = 0.0 + + Ccr = numpy.einsum("ij,j->ij", self.Qr[s][:, :mR], self.Dr[s][:mR]) # N x mR + Clcr = numpy.einsum("i,ij->ij", self.Dl[s][:mL], Ccr[:mL, :mR]) # mL x mR + + (Qlcr, Rlcr, Plcr) = scipy.linalg.qr( + Clcr, pivoting=True, check_finite=False + ) # mL x mL, min(mL,mR) x min(mL,mR), mR x mR + Dlcr = Rlcr.diagonal()[: min(mL, mR)] + Dinv = 1.0 / Dlcr + + mT = len(Dlcr[numpy.abs(Dlcr) > self.thresh]) + + assert mT <= mL and mT <= mR + + tmp = numpy.einsum("i,ij->ij", Dinv[:mT], Rlcr[:mT, :]) + tmp[:, Plcr] = tmp[:, range(mR)] # mT x mR + Tlcr = numpy.dot(tmp, self.Tr[s][:mR, :]) # mT x N + + Db = numpy.zeros(mT, B[s].dtype) + Ds = numpy.zeros(mT, B[s].dtype) + for i in range(mT): + absDlcr = abs(Dlcr[i]) + if absDlcr > 1.0: + Db[i] = 1.0 / absDlcr + Ds[i] = numpy.sign(Dlcr[i]) + else: + Db[i] = 1.0 + Ds[i] = Dlcr[i] + Dbinv = 1.0 / Db + + TQ = Tlcr[:, :mL].dot(Qlcr[:mL, :mT]) # mT x mT + TQinv = scipy.linalg.inv(TQ, check_finite=False) + tmp = numpy.einsum("ij,j->ij", TQinv, Db) + numpy.diag(Ds) # mT x mT + + M = numpy.einsum("ij,j->ij", tmp, Dbinv).dot(TQ) + # self.ovlp[s] = 1.0 / scipy.linalg.det(M, check_finite=False) + self.ovlp[s] = scipy.linalg.det(M, check_finite=False) + + tmp = scipy.linalg.inv(tmp, check_finite=False) + A = numpy.einsum("i,ij->ij", Db, tmp.dot(TQinv)) # mT x mT + Qlcr_pad = numpy.zeros((self.nbasis, self.nbasis), dtype=B[s].dtype) + Qlcr_pad[:mL, :mT] = Qlcr[:, :mT] + + # self.CT[s][:,:] = 0.0 + # self.CT[s][:,:mT] = Qlcr_pad[:,:mT].dot(numpy.diag(Dlcr[:mT])) + # self.theta[s][:,:] = 0.0 + # self.theta[s][:mT,:] = A.dot(Tlcr) + # self.G[s] = numpy.eye(self.nbasis, dtype=B[s].dtype) - self.CT[s][:,:mT].dot(self.theta[s][:mT,:]) + # self.CT[s][:,:mT] = self.CT[s][:,:mT].conj() + self.CT[s][:, :] = 0.0 + self.CT[s][:, :mT] = (A.dot(Tlcr)).T.conj() + self.theta[s][:, :] = 0.0 + self.theta[s][:mT, :] = Qlcr_pad[:, :mT].dot(numpy.diag(Dlcr[:mT])).T + # self.G[s] = numpy.eye(self.nbasis, dtype=B[s].dtype) - self.CT[s][:,:mT].dot(self.theta[s][:mT,:]) + self.G[s] = numpy.eye(self.nbasis, dtype=B[s].dtype) -\ + self.theta[s][:mT, :].T.dot(self.CT[s][:, :mT].T.conj()) + + # self.CT = numpy.zeros(shape=(2, nbasis, nbasis),dtype=dtype) + # self.theta = numpy.zeros(shape=(2, nbasis, nbasis),dtype=dtype) + # print("# mL, mR, mT = {}, {}, {}".format(mL, mR, mT)) + + # print("ovlp = {}".format(self.ovlp)) + self.mT = mT + self.time_slice += 1 # Count the time slice + self.block = self.time_slice // self.stack_size # move to the next block if necessary + self.counter = (self.counter + 1) % self.stack_size # Counting within a stack diff --git a/ipie/addons/thermal/walkers/tests/__init__.py b/ipie/addons/thermal/walkers/tests/__init__.py new file mode 100644 index 00000000..381108b3 --- /dev/null +++ b/ipie/addons/thermal/walkers/tests/__init__.py @@ -0,0 +1,14 @@ +# 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/thermal/walkers/tests/test_population_control.py b/ipie/addons/thermal/walkers/tests/test_population_control.py new file mode 100644 index 00000000..9250e5d5 --- /dev/null +++ b/ipie/addons/thermal/walkers/tests/test_population_control.py @@ -0,0 +1,358 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import numpy +import pytest +from typing import Union + +try: + from ipie.addons.thermal.utils.legacy_testing import build_legacy_generic_test_case_handlers_mpi + from ipie.addons.thermal.utils.legacy_testing import legacy_propagate_walkers + _no_cython = False + +except ModuleNotFoundError: + _no_cython = True + +from ipie.config import MPI +from ipie.utils.mpi import MPIHandler +from ipie.addons.thermal.walkers.pop_controller import ThermalPopController +from ipie.addons.thermal.utils.testing import build_generic_test_case_handlers_mpi + +comm = MPI.COMM_WORLD + +@pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") +@pytest.mark.unit +def test_pair_branch_batch(): + mpi_handler = MPIHandler() + + # System params. + nup = 5 + ndown = 5 + nelec = (nup, ndown) + nbasis = 10 + + # Thermal AFQMC params. + mu = -10. + beta = 0.1 + timestep = 0.01 + nwalkers = 12 + nblocks = 3 + # Must be fixed at 1 for Thermal AFQMC--legacy code overides whatever input! + nsteps_per_block = 1 + pop_control_method = 'pair_branch' + lowrank = False + + mf_trial = True + complex_integrals = False + debug = True + verbose = False if (comm.rank != 0) else True + seed = 7 + numpy.random.seed(seed) + + # Test. + objs = build_generic_test_case_handlers_mpi( + nelec, nbasis, mu, beta, timestep, mpi_handler, nwalkers=nwalkers, + lowrank=lowrank, mf_trial=mf_trial, complex_integrals=complex_integrals, + debug=debug, seed=seed, verbose=verbose) + trial = objs['trial'] + hamiltonian = objs['hamiltonian'] + walkers = objs['walkers'] + propagator = objs['propagator'] + pcontrol = ThermalPopController(nwalkers, nsteps_per_block, mpi_handler, + pop_control_method, verbose=verbose) + + # Legacy. + legacy_objs = build_legacy_generic_test_case_handlers_mpi( + hamiltonian, mpi_handler, nelec, mu, beta, timestep, + nwalkers=nwalkers, lowrank=lowrank, mf_trial=mf_trial, + seed=seed, verbose=verbose) + legacy_system = legacy_objs['system'] + legacy_trial = legacy_objs['trial'] + legacy_hamiltonian = legacy_objs['hamiltonian'] + legacy_walkers = legacy_objs['walkers'] + legacy_propagator = legacy_objs['propagator'] + + for block in range(nblocks): + for t in range(walkers.stack[0].nslice): + propagator.propagate_walkers(walkers, hamiltonian, trial, debug=True) + legacy_walkers = legacy_propagate_walkers( + legacy_hamiltonian, legacy_trial, legacy_walkers, + legacy_propagator, xi=propagator.xi) + + if t > 0: + pcontrol.pop_control(walkers, mpi_handler.comm) + legacy_walkers.pop_control(mpi_handler.comm) + + walkers.reset(trial) # Reset stack, weights, phase. + legacy_walkers.reset(legacy_trial) + + for iw in range(walkers.nwalkers): + assert numpy.allclose(walkers.Ga[iw], legacy_walkers.walkers[iw].G[0]) + assert numpy.allclose(walkers.Gb[iw], legacy_walkers.walkers[iw].G[1]) + assert numpy.allclose(walkers.weight[iw], legacy_walkers.walkers[iw].weight) + assert numpy.allclose(walkers.unscaled_weight[iw], legacy_walkers.walkers[iw].unscaled_weight) + +# TODO: Lowrank code is WIP. See: https://github.com/JoonhoLee-Group/ipie/issues/302 +#@pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") +#@pytest.mark.unit +def test_pair_branch_batch_lowrank(): + mpi_handler = MPIHandler() + + # System params. + nup = 5 + ndown = 5 + nelec = (nup, ndown) + nbasis = 10 + + # Thermal AFQMC params. + mu = -10. + beta = 0.1 + timestep = 0.01 + nwalkers = 12 + nblocks = 3 + # Must be fixed at 1 for Thermal AFQMC--legacy code overides whatever input! + nsteps_per_block = 1 + pop_control_method = 'pair_branch' + lowrank = True + + mf_trial = False + diagonal = True + complex_integrals = False + debug = True + verbose = False if (comm.rank != 0) else True + seed = 7 + numpy.random.seed(seed) + + options = { + 'nelec': nelec, + 'nbasis': nbasis, + 'mu': mu, + 'beta': beta, + 'timestep': timestep, + 'nwalkers': nwalkers, + 'seed': seed, + 'nsteps_per_block': nsteps_per_block, + 'nblocks': nblocks, + 'stabilize_freq': stabilize_freq, + 'pop_control_freq': pop_control_freq, + 'pop_control_method': pop_control_method, + 'lowrank': lowrank, + 'complex_integrals': complex_integrals, + 'mf_trial': mf_trial, + 'propagate': propagate, + 'diagonal': diagonal, + } + + # Test. + objs = build_generic_test_case_handlers_mpi( + nelec, nbasis, mu, beta, timestep, mpi_handler, nwalkers=nwalkers, + lowrank=lowrank, mf_trial=mf_trial, complex_integrals=complex_integrals, + diagonal=diagonal, debug=debug, seed=seed, verbose=verbose) + trial = objs['trial'] + hamiltonian = objs['hamiltonian'] + walkers = objs['walkers'] + propagator = objs['propagator'] + pcontrol = ThermalPopController(nwalkers, nsteps_per_block, mpi_handler, + pop_control_method=pop_control_method, verbose=verbose) + + # Legacy. + legacy_objs = build_legacy_generic_test_case_handlers_mpi( + hamiltonian, mpi_handler, nelec, mu, beta, timestep, + nwalkers=nwalkers, lowrank=lowrank, mf_trial=mf_trial, + seed=seed, verbose=verbose) + legacy_system = legacy_objs['system'] + legacy_trial = legacy_objs['trial'] + legacy_hamiltonian = legacy_objs['hamiltonian'] + legacy_walkers = legacy_objs['walkers'] + legacy_propagator = legacy_objs['propagator'] + + for block in range(nblocks): + for t in range(walkers.stack[0].nslice): + propagator.propagate_walkers(walkers, hamiltonian, trial, debug=True) + legacy_walkers = legacy_propagate_walkers( + legacy_hamiltonian, legacy_trial, legacy_walkers, + legacy_propagator, xi=propagator.xi) + + if t > 0: + pcontrol.pop_control(walkers, mpi_handler.comm) + legacy_walkers.pop_control(mpi_handler.comm) + + walkers.reset(trial) # Reset stack, weights, phase. + legacy_walkers.reset(legacy_trial) + + for iw in range(walkers.nwalkers): + assert numpy.allclose(walkers.Ga[iw], legacy_walkers.walkers[iw].G[0]) + assert numpy.allclose(walkers.Gb[iw], legacy_walkers.walkers[iw].G[1]) + assert numpy.allclose(walkers.weight[iw], legacy_walkers.walkers[iw].weight) + assert numpy.allclose(walkers.unscaled_weight[iw], legacy_walkers.walkers[iw].unscaled_weight) + + +@pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") +@pytest.mark.unit +def test_comb_batch(): + mpi_handler = MPIHandler() + # System params. + nup = 5 + ndown = 5 + nelec = (nup, ndown) + nbasis = 10 + + # Thermal AFQMC params. + mu = -10. + beta = 0.1 + timestep = 0.01 + nwalkers = 12 + nblocks = 3 + # Must be fixed at 1 for Thermal AFQMC--legacy code overides whatever input! + nsteps_per_block = 1 + pop_control_method = 'comb' + lowrank = False + + mf_trial = True + complex_integrals = False + debug = True + verbose = False if (comm.rank != 0) else True + seed = 7 + numpy.random.seed(seed) + + # Test. + objs = build_generic_test_case_handlers_mpi( + nelec, nbasis, mu, beta, timestep, mpi_handler, nwalkers=nwalkers, + lowrank=lowrank, mf_trial=mf_trial, complex_integrals=complex_integrals, + debug=debug, seed=seed, verbose=verbose) + trial = objs['trial'] + hamiltonian = objs['hamiltonian'] + walkers = objs['walkers'] + propagator = objs['propagator'] + pcontrol = ThermalPopController(nwalkers, nsteps_per_block, mpi_handler, + pop_control_method=pop_control_method, verbose=verbose) + + # Legacy. + legacy_objs = build_legacy_generic_test_case_handlers_mpi( + hamiltonian, mpi_handler, nelec, mu, beta, timestep, + nwalkers=nwalkers, lowrank=lowrank, mf_trial=mf_trial, + pop_control_method=pop_control_method, seed=seed, + verbose=verbose) + legacy_system = legacy_objs['system'] + legacy_trial = legacy_objs['trial'] + legacy_hamiltonian = legacy_objs['hamiltonian'] + legacy_walkers = legacy_objs['walkers'] + legacy_propagator = legacy_objs['propagator'] + + for block in range(nblocks): + for t in range(walkers.stack[0].nslice): + propagator.propagate_walkers(walkers, hamiltonian, trial, debug=True) + legacy_walkers = legacy_propagate_walkers( + legacy_hamiltonian, legacy_trial, legacy_walkers, + legacy_propagator, xi=propagator.xi) + + if t > 0: + pcontrol.pop_control(walkers, mpi_handler.comm) + legacy_walkers.pop_control(mpi_handler.comm) + + walkers.reset(trial) # Reset stack, weights, phase. + legacy_walkers.reset(legacy_trial) + + for iw in range(walkers.nwalkers): + assert numpy.allclose(walkers.Ga[iw], legacy_walkers.walkers[iw].G[0]) + assert numpy.allclose(walkers.Gb[iw], legacy_walkers.walkers[iw].G[1]) + assert numpy.allclose(walkers.weight[iw], legacy_walkers.walkers[iw].weight) + assert numpy.allclose(walkers.unscaled_weight[iw], legacy_walkers.walkers[iw].unscaled_weight) + + +# TODO: Lowrank code is WIP. See: https://github.com/JoonhoLee-Group/ipie/issues/302 +#@pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") +#@pytest.mark.unit +def test_comb_batch_lowrank(): + mpi_handler = MPIHandler() + + # System params. + nup = 5 + ndown = 5 + nelec = (nup, ndown) + nbasis = 10 + + # Thermal AFQMC params. + mu = -10. + beta = 0.1 + timestep = 0.01 + nwalkers = 12 + nblocks = 3 + # Must be fixed at 1 for Thermal AFQMC--legacy code overides whatever input! + nsteps_per_block = 1 + pop_control_method = 'comb' + lowrank = True + + mf_trial = False + diagonal = True + complex_integrals = False + debug = True + verbose = False if (comm.rank != 0) else True + seed = 7 + numpy.random.seed(seed) + + # Test. + objs = build_generic_test_case_handlers_mpi( + nelec, nbasis, mu, beta, timestep, mpi_handler, nwalkers=nwalkers, + lowrank=lowrank, mf_trial=mf_trial, complex_integrals=complex_integrals, + diagonal=diagonal, debug=debug, seed=seed, verbose=verbose) + trial = objs['trial'] + hamiltonian = objs['hamiltonian'] + walkers = objs['walkers'] + propagator = objs['propagator'] + pcontrol = ThermalPopController(nwalkers, nsteps_per_block, mpi_handler, + pop_control_method=pop_control_method, verbose=verbose) + + # Legacy. + legacy_objs = build_legacy_generic_test_case_handlers_mpi( + hamiltonian, mpi_handler, nelec, mu, beta, timestep, + nwalkers=nwalkers, lowrank=lowrank, mf_trial=mf_trial, + seed=seed, verbose=verbose) + legacy_system = legacy_objs['system'] + legacy_trial = legacy_objs['trial'] + legacy_hamiltonian = legacy_objs['hamiltonian'] + legacy_walkers = legacy_objs['walkers'] + legacy_propagator = legacy_objs['propagator'] + + for block in range(nblocks): + for t in range(walkers.stack[0].nslice): + propagator.propagate_walkers(walkers, hamiltonian, trial, debug=True) + legacy_walkers = legacy_propagate_walkers( + legacy_hamiltonian, legacy_trial, legacy_walkers, + legacy_propagator, xi=propagator.xi) + + if t > 0: + pcontrol.pop_control(walkers, mpi_handler.comm) + legacy_walkers.pop_control(mpi_handler.comm) + + walkers.reset(trial) # Reset stack, weights, phase. + legacy_walkers.reset(legacy_trial) + + for iw in range(walkers.nwalkers): + assert numpy.allclose(walkers.Ga[iw], legacy_walkers.walkers[iw].G[0]) + assert numpy.allclose(walkers.Gb[iw], legacy_walkers.walkers[iw].G[1]) + assert numpy.allclose(walkers.weight[iw], legacy_walkers.walkers[iw].weight) + assert numpy.allclose(walkers.unscaled_weight[iw], legacy_walkers.walkers[iw].unscaled_weight) + + +if __name__ == "__main__": + test_pair_branch_batch() + test_comb_batch() + + #test_pair_branch_batch_lowrank() + #test_comb_batch_lowrank() diff --git a/ipie/addons/thermal/walkers/tests/test_thermal_walkers.py b/ipie/addons/thermal/walkers/tests/test_thermal_walkers.py new file mode 100644 index 00000000..a02da99c --- /dev/null +++ b/ipie/addons/thermal/walkers/tests/test_thermal_walkers.py @@ -0,0 +1,156 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import numpy +import pytest +from typing import Union + +try: + from ipie.addons.thermal.utils.legacy_testing import build_legacy_generic_test_case_handlers + from ipie.addons.thermal.utils.legacy_testing import legacy_propagate_walkers + _no_cython = False + +except ModuleNotFoundError: + _no_cython = True + +from ipie.config import MPI +from ipie.addons.thermal.estimators.generic import local_energy_generic_cholesky +from ipie.addons.thermal.estimators.thermal import one_rdm_from_G +from ipie.addons.thermal.utils.testing import build_generic_test_case_handlers + +from ipie.legacy.estimators.generic import local_energy_generic_cholesky as legacy_local_energy_generic_cholesky +from ipie.legacy.estimators.thermal import one_rdm_from_G as legacy_one_rdm_from_G + +comm = MPI.COMM_WORLD + +@pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") +@pytest.mark.unit +def test_thermal_walkers_fullrank(): + # System params. + nup = 5 + ndown = 5 + nelec = (nup, ndown) + nbasis = 10 + + # Thermal AFQMC params. + mu = -10. + beta = 0.1 + timestep = 0.01 + nwalkers = 10 + lowrank = False + + mf_trial = True + complex_integrals = False + debug = True + verbose = True + seed = 7 + numpy.random.seed(seed) + + # Test. + objs = build_generic_test_case_handlers( + nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, + mf_trial=mf_trial, complex_integrals=complex_integrals, debug=debug, + seed=seed, verbose=verbose) + trial = objs['trial'] + hamiltonian = objs['hamiltonian'] + walkers = objs['walkers'] + + # Legacy. + legacy_objs = build_legacy_generic_test_case_handlers( + hamiltonian, comm, nelec, mu, beta, timestep, nwalkers=nwalkers, + lowrank=lowrank, mf_trial=mf_trial, seed=seed, verbose=verbose) + legacy_system = legacy_objs['system'] + legacy_trial = legacy_objs['trial'] + legacy_hamiltonian = legacy_objs['hamiltonian'] + legacy_walkers = legacy_objs['walkers'] + + for iw in range(walkers.nwalkers): + P = one_rdm_from_G(numpy.array([walkers.Ga[iw], walkers.Gb[iw]])) + eloc = local_energy_generic_cholesky(hamiltonian, P) + + legacy_P = legacy_one_rdm_from_G(numpy.array(legacy_walkers.walkers[iw].G)) + legacy_eloc = legacy_local_energy_generic_cholesky( + legacy_system, legacy_hamiltonian, legacy_P) + + numpy.testing.assert_almost_equal(legacy_eloc, eloc, decimal=10) + numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].G[0], walkers.Ga[iw], decimal=10) + numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].G[1], walkers.Gb[iw], decimal=10) + numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].stack.ovlp[0], walkers.stack[iw].ovlp[0], decimal=10) + numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].stack.ovlp[1], walkers.stack[iw].ovlp[1], decimal=10) + + +# TODO: Lowrank code is WIP. +#@pytest.mark.skipif(_no_cython, reason="Need to build cython modules.") +#@pytest.mark.unit +def test_thermal_walkers_lowrank(): + # System params. + nup = 5 + ndown = 5 + nelec = (nup, ndown) + nbasis = 10 + + # Thermal AFQMC params. + mu = -10. + beta = 0.1 + timestep = 0.01 + nwalkers = 10 + lowrank = True + + mf_trial = False + diagonal = True + complex_integrals = False + debug = True + verbose = True + seed = 7 + numpy.random.seed(seed) + + # Test. + objs = build_generic_test_case_handlers( + nelec, nbasis, mu, beta, timestep, nwalkers=nwalkers, lowrank=lowrank, + mf_trial=mf_trial, complex_integrals=complex_integrals, debug=debug, + seed=seed, verbose=verbose) + trial = objs['trial'] + hamiltonian = objs['hamiltonian'] + walkers = objs['walkers'] + + # Legacy. + legacy_objs = build_legacy_generic_test_case_handlers( + hamiltonian, comm, nelec, mu, beta, timestep, nwalkers=nwalkers, + lowrank=lowrank, mf_trial=mf_trial, seed=seed, verbose=verbose) + legacy_system = legacy_objs['system'] + legacy_trial = legacy_objs['trial'] + legacy_hamiltonian = legacy_objs['hamiltonian'] + legacy_walkers = legacy_objs['walkers'] + + for iw in range(walkers.nwalkers): + P = one_rdm_from_G(numpy.array([walkers.Ga[iw], walkers.Gb[iw]])) + eloc = local_energy_generic_cholesky(hamiltonian, P) + + legacy_P = legacy_one_rdm_from_G(numpy.array(legacy_walkers.walkers[iw].G)) + legacy_eloc = legacy_local_energy_generic_cholesky( + legacy_system, legacy_hamiltonian, legacy_P) + + numpy.testing.assert_almost_equal(legacy_eloc, eloc, decimal=10) + numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].G[0], walkers.Ga[iw], decimal=10) + numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].G[1], walkers.Gb[iw], decimal=10) + numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].stack.ovlp[0], walkers.stack[iw].ovlp[0], decimal=10) + numpy.testing.assert_almost_equal(legacy_walkers.walkers[iw].stack.ovlp[1], walkers.stack[iw].ovlp[1], decimal=10) + +if __name__ == "__main__": + test_thermal_walkers_fullrank() + #test_thermal_walkers_lowrank() diff --git a/ipie/addons/thermal/walkers/uhf_walkers.py b/ipie/addons/thermal/walkers/uhf_walkers.py new file mode 100644 index 00000000..e3cac522 --- /dev/null +++ b/ipie/addons/thermal/walkers/uhf_walkers.py @@ -0,0 +1,158 @@ +# 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. +# +# Authors: Fionn Malone +# Joonho Lee +# + +import numpy +import scipy.linalg + +from ipie.addons.thermal.estimators.particle_number import particle_number +from ipie.addons.thermal.estimators.thermal import one_rdm_from_G +from ipie.addons.thermal.estimators.greens_function import greens_function_qr_strat +from ipie.addons.thermal.walkers.stack import PropagatorStack +from ipie.utils.misc import update_stack +from ipie.walkers.base_walkers import BaseWalkers +from ipie.addons.thermal.trial.one_body import OneBody + +class UHFThermalWalkers(BaseWalkers): + def __init__( + self, + trial: OneBody, + nbasis: int, + nwalkers: int, + stack_size = None, + lowrank: bool = False, + lowrank_thresh: float = 1e-6, + mpi_handler = None, + verbose: bool = False, + ): + """UHF style walker. + """ + assert isinstance(trial, OneBody) + super().__init__(nwalkers, verbose=verbose) + + self.nbasis = nbasis + self.mpi_handler = mpi_handler + self.nslice = trial.nslice + self.stack_size = stack_size + + if self.stack_size == None: + self.stack_size = trial.stack_size + + if (self.nslice // self.stack_size) * self.stack_size != self.nslice: + if verbose: + print("# Input stack size does not divide number of slices.") + self.stack_size = update_stack(self.stack_size, self.nslice, verbose) + + if self.stack_size > trial.stack_size: + if verbose: + print("# Walker stack size differs from that estimated from trial density matrix.") + print(f"# Be careful. cond(BT)**stack_size: {trial.cond ** self.stack_size:10.3e}.") + + self.stack_length = self.nslice // self.stack_size + self.lowrank = lowrank + self.lowrank_thresh = lowrank_thresh + + self.Ga = numpy.zeros( + shape=(self.nwalkers, self.nbasis, self.nbasis), + dtype=numpy.complex128) + self.Gb = numpy.zeros( + shape=(self.nwalkers, self.nbasis, self.nbasis), + dtype=numpy.complex128) + self.Ghalf = None + + max_diff_diag = numpy.linalg.norm( + (numpy.diag( + trial.dmat[0].diagonal()) - trial.dmat[0])) + + if max_diff_diag < 1e-10: + self.diagonal_trial = True + if verbose: + print("# Trial density matrix is diagonal.") + else: + self.diagonal_trial = False + if verbose: + print("# Trial density matrix is not diagonal.") + + if verbose: + print(f"# Walker stack size: {self.stack_size}") + print(f"# Using low rank trick: {self.lowrank}") + + self.stack = [PropagatorStack( + self.stack_size, + self.nslice, + self.nbasis, + numpy.complex128, + trial.dmat, + trial.dmat_inv, + diagonal=self.diagonal_trial, + lowrank=self.lowrank, + thresh=self.lowrank_thresh, + ) for iw in range(self.nwalkers)] + + # Initialise all propagators to the trial density matrix. + for iw in range(self.nwalkers): + self.stack[iw].set_all(trial.dmat) + greens_function_qr_strat(self, iw) + self.stack[iw].G[0] = self.Ga[iw] + self.stack[iw].G[1] = self.Gb[iw] + + # Shape (nwalkers,). + self.M0a = numpy.array([ + scipy.linalg.det(self.Ga[iw], check_finite=False) for iw in range(self.nwalkers)]) + self.M0b = numpy.array([ + scipy.linalg.det(self.Gb[iw], check_finite=False) for iw in range(self.nwalkers)]) + + for iw in range(self.nwalkers): + self.stack[iw].ovlp = numpy.array([1.0 / self.M0a[iw], 1.0 / self.M0b[iw]]) + + self.hybrid_energy = 0.0 + if verbose: + for iw in range(self.nwalkers): + G = numpy.array([self.Ga[iw], self.Gb[iw]]) + P = one_rdm_from_G(G) + nav = particle_number(P) + print(f"# Trial electron number for {iw}-th walker: {nav}") + + self.buff_names += ["Ga", "Gb"] + 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 calc_greens_function(self, iw, slice_ix=None, inplace=True): + """Return the Green's function for walker `iw`. + """ + if self.lowrank: + return self.stack[iw].G # G[0] = Ga, G[1] = Gb + + else: + return greens_function_qr_strat(self, iw, slice_ix=slice_ix, inplace=inplace) + + + def reset(self, trial): + self.weight = numpy.ones(self.nwalkers) + self.phase = numpy.ones(self.nwalkers, dtype=numpy.complex128) + + for iw in range(self.nwalkers): + self.stack[iw].reset() + self.stack[iw].set_all(trial.dmat) + self.calc_greens_function(iw) + + # For compatibiltiy with BaseWalkers class. + def reortho(self): + pass + + def reortho_batched(self): + pass diff --git a/ipie/analysis/extraction.py b/ipie/analysis/extraction.py index f9103afd..ec33ac8b 100755 --- a/ipie/analysis/extraction.py +++ b/ipie/analysis/extraction.py @@ -230,7 +230,10 @@ def extract_test_data_hdf5(filename, skip=10): data = extract_mixed_estimates(filename) # use list so can json serialise easily. data = data.drop(["Iteration", "Time"], axis=1)[::skip].to_dict(orient="list") - data["sys_info"] = get_metadata(filename)["sys_info"] + try: + data["sys_info"] = get_metadata(filename)["sys_info"] + except KeyError: + print("\n# No 'sys_info' metadata!") try: mrdm = extract_rdm(filename, est_type="mixed", rdm_type="one_rdm") except (KeyError, TypeError, AttributeError): diff --git a/ipie/estimators/energy.py b/ipie/estimators/energy.py index 9d163bd1..23b5652b 100644 --- a/ipie/estimators/energy.py +++ b/ipie/estimators/energy.py @@ -118,9 +118,6 @@ def __init__( 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 diff --git a/ipie/estimators/estimator_base.py b/ipie/estimators/estimator_base.py index 13a9e12c..c78c9322 100644 --- a/ipie/estimators/estimator_base.py +++ b/ipie/estimators/estimator_base.py @@ -15,13 +15,13 @@ # Author: Fionn Malone # -from abc import ABCMeta, abstractmethod +from abc import abstractmethod, ABCMeta import numpy as np +from ipie.utils.io import format_fixed_width_strings, format_fixed_width_floats from ipie.utils.backend import arraylib as xp from ipie.utils.backend import to_host -from ipie.utils.io import format_fixed_width_floats, format_fixed_width_strings class EstimatorBase(metaclass=ABCMeta): @@ -89,7 +89,8 @@ def shape(self, shape) -> tuple: self._shape = shape @abstractmethod - def compute_estimator(self, system, walkers, hamiltonian, trial) -> np.ndarray: ... + def compute_estimator(self, system, walkers, hamiltonian, trial) -> np.ndarray: + ... @property def names(self): @@ -141,4 +142,5 @@ def zero(self): else: self._data[k] = 0.0j - def post_reduce_hook(self, data) -> None: ... + def post_reduce_hook(self, data) -> None: + ... diff --git a/ipie/estimators/generic.py b/ipie/estimators/generic.py index 469abc3c..6f69d085 100644 --- a/ipie/estimators/generic.py +++ b/ipie/estimators/generic.py @@ -33,7 +33,9 @@ def local_energy_generic_opt(system, G, Ghalf=None, eri=None): assert eri is not None vipjq_aa = eri[0, : na**2 * M**2].reshape((na, M, na, M)) - vipjq_bb = eri[0, na**2 * M**2 : na**2 * M**2 + nb**2 * M**2].reshape((nb, M, nb, M)) + vipjq_bb = eri[0, na**2 * M**2 : na**2 * M**2 + nb**2 * M**2].reshape( + (nb, M, nb, M) + ) vipjq_ab = eri[0, na**2 * M**2 + nb**2 * M**2 :].reshape((na, M, nb, M)) Ga, Gb = Ghalf[0], Ghalf[1] diff --git a/ipie/estimators/tests/test_estimators.py b/ipie/estimators/tests/test_estimators.py index 5c68cad4..59adc89f 100644 --- a/ipie/estimators/tests/test_estimators.py +++ b/ipie/estimators/tests/test_estimators.py @@ -74,8 +74,3 @@ def test_estimator_handler(): handler.initialize(comm) handler.compute_estimators(comm, system, ham, trial, walker_batch) handler.compute_estimators(comm, system, ham, trial, walker_batch) - - -if __name__ == "__main__": - test_energy_estimator() - test_estimator_handler() diff --git a/ipie/hamiltonians/generic.py b/ipie/hamiltonians/generic.py index 137ccdb0..d00873c5 100644 --- a/ipie/hamiltonians/generic.py +++ b/ipie/hamiltonians/generic.py @@ -30,7 +30,7 @@ def construct_h1e_mod(chol, h1e, h1e_mod): # Subtract one-body bit following reordering of 2-body operators. - # Eqn (17) of [Motta17]_ + # Eqn (17) of [Motta17]. nbasis = h1e.shape[-1] nchol = chol.shape[-1] chol_view = chol.reshape((nbasis, nbasis * nchol)) @@ -109,7 +109,7 @@ def __init__(self, h1e, chol, ecore=0.0, verbose=False): self.chol = numpy.array(chol, dtype=numpy.complex128) # [M^2, nchol] self.nchol = self.chol.shape[-1] - self.nfields = self.nchol * 2 + self.nfields = 2 * self.nchol assert self.nbasis**2 == chol.shape[0] self.chunked = False diff --git a/ipie/legacy/estimators/generic.py b/ipie/legacy/estimators/generic.py index 1682e1e9..910a5f92 100644 --- a/ipie/legacy/estimators/generic.py +++ b/ipie/legacy/estimators/generic.py @@ -190,6 +190,37 @@ def _exx_compute_batch(rchol_a, rchol_b, GaT_stacked, GbT_stacked, lwalker): return exx_vec_b + exx_vec_a +# FDM: deprecated remove? +def local_energy_generic_opt(system, G, Ghalf=None, eri=None): + """Compute local energy using half-rotated eri tensor.""" + + na = system.nup + nb = system.ndown + M = system.nbasis + assert eri is not None + + vipjq_aa = eri[0, : na**2 * M**2].reshape((na, M, na, M)) + vipjq_bb = eri[0, na**2 * M**2 : na**2 * M**2 + nb**2 * M**2].reshape( + (nb, M, nb, M) + ) + vipjq_ab = eri[0, na**2 * M**2 + nb**2 * M**2 :].reshape((na, M, nb, M)) + + Ga, Gb = Ghalf[0], Ghalf[1] + # Element wise multiplication. + e1b = numpy.sum(system.H1[0] * G[0]) + numpy.sum(system.H1[1] * G[1]) + # Coulomb + eJaa = 0.5 * numpy.einsum("irjs,ir,js", vipjq_aa, Ga, Ga) + eJbb = 0.5 * numpy.einsum("irjs,ir,js", vipjq_bb, Gb, Gb) + eJab = numpy.einsum("irjs,ir,js", vipjq_ab, Ga, Gb) + + eKaa = -0.5 * numpy.einsum("irjs,is,jr", vipjq_aa, Ga, Ga) + eKbb = -0.5 * numpy.einsum("irjs,is,jr", vipjq_bb, Gb, Gb) + + e2b = eJaa + eJbb + eJab + eKaa + eKbb + + return (e1b + e2b + system.ecore, e1b + system.ecore, e2b) + + def local_energy_generic_cholesky_opt_batched( system, ham, diff --git a/ipie/legacy/estimators/local_energy.py b/ipie/legacy/estimators/local_energy.py index 8de3ce93..53225747 100644 --- a/ipie/legacy/estimators/local_energy.py +++ b/ipie/legacy/estimators/local_energy.py @@ -5,10 +5,10 @@ from ipie.legacy.estimators.ueg import local_energy_ueg except ImportError as e: print(e) -from ipie.estimators.generic import local_energy_generic_opt from ipie.legacy.estimators.ci import get_hmatel from ipie.legacy.estimators.generic import ( local_energy_generic, + local_energy_generic_opt, local_energy_generic_cholesky, local_energy_generic_cholesky_opt, local_energy_generic_cholesky_opt_stochastic, diff --git a/ipie/legacy/estimators/ueg.py b/ipie/legacy/estimators/ueg.py index dde01443..210bb0ac 100644 --- a/ipie/legacy/estimators/ueg.py +++ b/ipie/legacy/estimators/ueg.py @@ -25,7 +25,7 @@ def coulomb_greens_function(nq, kpq_i, kpq, pmq_i, pmq, Gkpq, Gpmq, G): Gpmq[iq] += G[i, idxpmq] -def local_energy_ueg(system, ham, G, Ghalf=None, two_rdm=None): +def local_energy_ueg(system, ham, G, Ghalf=None, two_rdm=None, debug=False): """Local energy computation for uniform electron gas Parameters ---------- @@ -69,20 +69,31 @@ def local_energy_ueg(system, ham, G, Ghalf=None, two_rdm=None): if two_rdm is None: two_rdm = numpy.zeros((2, 2, len(ham.qvecs)), dtype=numpy.complex128) - two_rdm[0, 0] = numpy.multiply(Gkpq[0], Gpmq[0]) - Gprod[0] - essa = (1.0 / (2.0 * ham.vol)) * ham.vqvec.dot(two_rdm[0, 0]) + two_rdm[0, 0] = numpy.multiply(Gkpq[0], Gpmq[0]) - Gprod[0] two_rdm[1, 1] = numpy.multiply(Gkpq[1], Gpmq[1]) - Gprod[1] - essb = (1.0 / (2.0 * ham.vol)) * ham.vqvec.dot(two_rdm[1, 1]) - two_rdm[0, 1] = numpy.multiply(Gkpq[0], Gpmq[1]) two_rdm[1, 0] = numpy.multiply(Gkpq[1], Gpmq[0]) - eos = (1.0 / (2.0 * ham.vol)) * ham.vqvec.dot(two_rdm[0, 1]) + ( - 1.0 / (2.0 * ham.vol) - ) * ham.vqvec.dot(two_rdm[1, 0]) + essa = (1.0 / (2.0 * ham.vol)) * ham.vqvec.dot(two_rdm[0, 0]) + essb = (1.0 / (2.0 * ham.vol)) * ham.vqvec.dot(two_rdm[1, 1]) + eos = (1.0 / (2.0 * ham.vol)) * ( + ham.vqvec.dot(two_rdm[0, 1]) + ham.vqvec.dot(two_rdm[1, 0])) pe = essa + essb + eos + if debug: + ecoul, exx = 0., 0. + + for s1 in range(2): + exx -= (1.0 / (2.0 * ham.vol)) * ham.vqvec.dot(Gprod[s1]) + + for s2 in range(2): + ecoul += (1.0 / (2.0 * ham.vol)) * ham.vqvec.dot(numpy.multiply(Gkpq[s1], Gpmq[s2])) + + print(f'\n# ueg ecoul = {ecoul}') + print(f'# ueg exx = {exx}') + print(f'# ueg e2 = {(ecoul + exx)}') + return (ke + pe, ke, pe) diff --git a/ipie/legacy/hamiltonians/_generic.py b/ipie/legacy/hamiltonians/_generic.py index e361deb5..6b1a2ef8 100644 --- a/ipie/legacy/hamiltonians/_generic.py +++ b/ipie/legacy/hamiltonians/_generic.py @@ -140,7 +140,7 @@ def __init__( if self.verbose: print("# mixed_precision is used for the propagation") - if isrealobj(self.chol_vecs.dtype): + if isrealobj(self.chol_vecs): if verbose: print("# Found real Choleksy integrals.") self.cplx_chol = False @@ -314,7 +314,7 @@ def construct_h1e_mod(chol, h1e, h1e_mod): chol_view = chol.reshape((nbasis, nbasis * nchol)) # assert chol_view.__array_interface__['data'][0] == chol.__array_interface__['data'][0] v0 = 0.5 * numpy.dot( - chol_view, chol_view.T + chol_view, chol_view.T.conj() # Conjugate added to account for complex integrals ) # einsum('ikn,jkn->ij', chol_3, chol_3, optimize=True) h1e_mod[0, :, :] = h1e[0] - v0 h1e_mod[1, :, :] = h1e[1] - v0 diff --git a/ipie/legacy/propagation/pw.py b/ipie/legacy/propagation/pw.py index 379d4f9f..87445311 100644 --- a/ipie/legacy/propagation/pw.py +++ b/ipie/legacy/propagation/pw.py @@ -5,7 +5,7 @@ import numpy import scipy -from ipie.estimators.utils import convolve, scipy_fftconvolve +from ipie.legacy.estimators.utils import convolve, scipy_fftconvolve from ipie.legacy.propagation.operations import kinetic_real diff --git a/ipie/legacy/propagation/tests/test_planewave.py b/ipie/legacy/propagation/tests/test_planewave.py index e05b2604..ed12be56 100644 --- a/ipie/legacy/propagation/tests/test_planewave.py +++ b/ipie/legacy/propagation/tests/test_planewave.py @@ -1,5 +1,3 @@ -import os - import numpy import pytest diff --git a/ipie/legacy/qmc/thermal_afqmc.py b/ipie/legacy/qmc/thermal_afqmc.py index f4e60a9c..2fb15e54 100644 --- a/ipie/legacy/qmc/thermal_afqmc.py +++ b/ipie/legacy/qmc/thermal_afqmc.py @@ -116,10 +116,10 @@ def __init__( self.sha1, self.branch, self.local_mods = get_git_info() else: self.sha1 = "None" - if verbose: - self.sys_info = print_env_info( - self.sha1, self.branch, self.local_mods, self.uuid, comm.size - ) + #if verbose: + # self.sys_info = print_env_info( + # self.sha1, self.branch, self.local_mods, self.uuid, comm.size + # ) # Hack - this is modified later if running in parallel on # initialisation. self.root = comm.rank == 0 diff --git a/ipie/legacy/thermal_propagation/continuous.py b/ipie/legacy/thermal_propagation/continuous.py index a97c6dc4..cfc6dc7e 100644 --- a/ipie/legacy/thermal_propagation/continuous.py +++ b/ipie/legacy/thermal_propagation/continuous.py @@ -84,7 +84,7 @@ def __init__(self, options, qmc, system, hamiltonian, trial, verbose=False, lowr if verbose: print("# Finished setting up propagator.") - def two_body_propagator(self, walker, system, trial): + def two_body_propagator(self, walker, system, trial, xi=None): r"""Continuous Hubbard-Statonovich transformation. Parameters @@ -97,7 +97,9 @@ def two_body_propagator(self, walker, system, trial): Trial wavefunction object. """ # Normally distrubted auxiliary fields. - xi = numpy.random.normal(0.0, 1.0, system.nfields) + if xi is None: # For debugging. + xi = numpy.random.normal(0.0, 1.0, system.nfields) + if self.force_bias: P = one_rdm_from_G(walker.G) xbar = self.propagator.construct_force_bias(system, P, trial) @@ -157,7 +159,7 @@ def exponentiate(self, VHS, debug=False): print(f"DIFF: {(c2 - phi).sum() / c2.size: 10.8e}") return phi - def propagate_walker_free(self, system, walker, trial, eshift=0): + def propagate_walker_free(self, system, walker, trial, eshift=0, xi=None): r"""Free projection for continuous HS transformation. .. Warning:: @@ -173,7 +175,7 @@ def propagate_walker_free(self, system, walker, trial, eshift=0): state : :class:`state.State` Simulation state. """ - (cmf, cfb, xmxbar, VHS) = self.two_body_propagator(walker, system, trial) + (cmf, cfb, xmxbar, VHS) = self.two_body_propagator(walker, system, trial, xi=xi) BV = self.exponentiate(VHS) B = numpy.array([BV.dot(self.BH1[0]), BV.dot(self.BH1[1])]) @@ -206,7 +208,7 @@ def propagate_walker_free(self, system, walker, trial, eshift=0): except ZeroDivisionError: walker.weight = 0.0 - def propagate_walker_phaseless(self, system, walker, trial, eshift=0): + def propagate_walker_phaseless(self, system, walker, trial, eshift=0, xi=None): r"""Propagate walker using phaseless approximation. Uses importance sampling and the hybrid method. @@ -223,7 +225,7 @@ def propagate_walker_phaseless(self, system, walker, trial, eshift=0): Trial wavefunction object. """ - (cmf, cfb, xmxbar, VHS) = self.two_body_propagator(walker, system, trial) + (cmf, cfb, xmxbar, VHS) = self.two_body_propagator(walker, system, trial, xi=xi) BV = self.exponentiate(VHS) B = numpy.array([BV.dot(self.BH1[0]), BV.dot(self.BH1[1])]) @@ -232,11 +234,15 @@ def propagate_walker_phaseless(self, system, walker, trial, eshift=0): # Compute determinant ratio det(1+A')/det(1+A). # 1. Current walker's green's function. tix = walker.stack.ntime_slices + G = walker.greens_function(None, slice_ix=tix, inplace=False) # 2. Compute updated green's function. walker.stack.update_new(B) walker.greens_function(None, slice_ix=tix, inplace=True) # 3. Compute det(G/G') - M0 = walker.M0 + M0 = [ + scipy.linalg.det(G[0], check_finite=False), + scipy.linalg.det(G[1], check_finite=False) + ] Mnew = [ scipy.linalg.det(walker.G[0], check_finite=False), scipy.linalg.det(walker.G[1], check_finite=False), diff --git a/ipie/legacy/thermal_propagation/generic.py b/ipie/legacy/thermal_propagation/generic.py index 7454b5df..9c4c63de 100644 --- a/ipie/legacy/thermal_propagation/generic.py +++ b/ipie/legacy/thermal_propagation/generic.py @@ -23,7 +23,7 @@ class GenericContinuous(object): qmc : :class:`pie.qmc.options.QMCOpts` QMC options. system : :class:`pie.system.System` - System object. + System object is actually HAMILTONIAN!!! trial : :class:`pie.trial_wavefunctioin.Trial` Trial wavefunction object. verbose : bool @@ -91,7 +91,10 @@ def construct_mean_field_shift(self, system, P): mf_shift = 1j * P[0].ravel() * system.chol_vecs mf_shift += 1j * P[1].ravel() * system.chol_vecs else: + # Need to reshape `chol_vecs` just to run the einsum lol. + system.chol_vecs = system.chol_vecs.T.reshape(system.nchol, system.nbasis, system.nbasis) mf_shift = 1j * numpy.einsum("lpq,spq->l", system.chol_vecs, P) + system.chol_vecs = system.chol_vecs.reshape(system.nchol, system.nbasis**2).T return mf_shift def construct_one_body_propagator(self, system, dt): diff --git a/ipie/legacy/thermal_propagation/planewave.py b/ipie/legacy/thermal_propagation/planewave.py index 8b258c38..a0a92871 100644 --- a/ipie/legacy/thermal_propagation/planewave.py +++ b/ipie/legacy/thermal_propagation/planewave.py @@ -105,7 +105,6 @@ def construct_one_body_propagator(self, hamiltonian, dt): """ H1 = hamiltonian.h1e_mod I = numpy.identity(H1[0].shape[0], dtype=H1.dtype) - print(f"hamiltonian.mu = {hamiltonian.mu}") # No spin dependence for the moment. self.BH1 = numpy.array( [ @@ -228,7 +227,7 @@ def propagate_greens_function(self, walker, B, Binv): walker.G[0] = B[0].dot(walker.G[0]).dot(Binv[0]) walker.G[1] = B[1].dot(walker.G[1]).dot(Binv[1]) - def two_body_propagator(self, walker, hamiltonian, force_bias=True): + def two_body_propagator(self, walker, hamiltonian, force_bias=True, xi=None): """It appliese the two-body propagator Parameters ---------- @@ -249,7 +248,8 @@ def two_body_propagator(self, walker, hamiltonian, force_bias=True): """ # Normally distrubted auxiliary fields. - xi = numpy.random.normal(0.0, 1.0, hamiltonian.nfields) + if xi is None: + xi = numpy.random.normal(0.0, 1.0, hamiltonian.nfields) # Optimal force bias. xbar = numpy.zeros(hamiltonian.nfields) @@ -472,7 +472,7 @@ def propagate_walker_free_low_rank(self, system, walker, trial, eshift=0, force_ except ZeroDivisionError: walker.weight = 0.0 - def propagate_walker_phaseless_full_rank(self, hamiltonian, walker, trial, eshift=0): + def propagate_walker_phaseless_full_rank(self, hamiltonian, walker, trial, eshift=0, xi=None): # """Phaseless propagator # Parameters # ---------- @@ -486,7 +486,7 @@ def propagate_walker_phaseless_full_rank(self, hamiltonian, walker, trial, eshif # ------- # """ - (cmf, cfb, xmxbar, VHS) = self.two_body_propagator(walker, hamiltonian, True) + (cmf, cfb, xmxbar, VHS) = self.two_body_propagator(walker, hamiltonian, True, xi=xi) BV = self.exponentiate(VHS) # could use a power-series method to build this B = numpy.array( @@ -525,7 +525,10 @@ def propagate_walker_phaseless_full_rank(self, hamiltonian, walker, trial, eshif walker.greens_function(None, slice_ix=tix, inplace=True) # 3. Compute det(G/G') - M0 = walker.M0 + M0 = [ + scipy.linalg.det(G[0], check_finite=False), + scipy.linalg.det(G[1], check_finite=False) + ] Mnew = numpy.array( [ scipy.linalg.det(walker.G[0], check_finite=False), @@ -554,7 +557,7 @@ def propagate_walker_phaseless_full_rank(self, hamiltonian, walker, trial, eshif except ZeroDivisionError: walker.weight = 0.0 - def propagate_walker_phaseless_low_rank(self, hamiltonian, walker, trial, eshift=0): + def propagate_walker_phaseless_low_rank(self, hamiltonian, walker, trial, eshift=0, xi=None): # """Phaseless propagator # Parameters # ---------- @@ -567,7 +570,7 @@ def propagate_walker_phaseless_low_rank(self, hamiltonian, walker, trial, eshift # Returns # ------- # """ - (cmf, cfb, xmxbar, VHS) = self.two_body_propagator(walker, hamiltonian, True) + (cmf, cfb, xmxbar, VHS) = self.two_body_propagator(walker, hamiltonian, True, xi=xi) BV = self.exponentiate(VHS) # could use a power-series method to build this B = numpy.array( diff --git a/ipie/legacy/trial_density_matrices/mean_field.py b/ipie/legacy/trial_density_matrices/mean_field.py index 46c01bb4..1fe8e1e2 100644 --- a/ipie/legacy/trial_density_matrices/mean_field.py +++ b/ipie/legacy/trial_density_matrices/mean_field.py @@ -41,7 +41,7 @@ def thermal_hartree_fock(self, system, beta): mu_old = self.mu P = self.P.copy() if self.verbose: - print("# Determining Thermal Hartree--Fock Density Matrix.") + print("# Determining Thermal Hartree-Fock Density Matrix.") for it in range(self.max_macro_it): if self.verbose: print(f"# Macro iteration: {it}") @@ -49,7 +49,7 @@ def thermal_hartree_fock(self, system, beta): rho = numpy.array([scipy.linalg.expm(-dt * HMF[0]), scipy.linalg.expm(-dt * HMF[1])]) if self.find_mu: mu = find_chemical_potential( - system, + system._alt_convention, rho, dt, self.num_bins, @@ -96,16 +96,16 @@ def scf(self, system, beta, mu, P): change = numpy.linalg.norm(Pnew - Pold) if change < self.deps: break - if self.verbose: - N = particle_number(P).real - E = local_energy(system, P)[0].real - S = entropy(beta, mu, HMF) - omega = E - mu * N - 1.0 / beta * S - print( - " # Iteration: {:4d} dP: {:13.8e} Omega: {:13.8e}".format( - it, change, omega.real - ) - ) + #if self.verbose: + # N = particle_number(P).real + # E = local_energy(system, P)[0].real + # S = entropy(beta, mu, HMF) + # omega = E - mu * N - 1.0 / beta * S + # print( + # " # Iteration: {:4d} dP: {:13.8e} Omega: {:13.8e}".format( + # it, change, omega.real + # ) + # ) Pold = Pnew.copy() if self.verbose: N = particle_number(P).real diff --git a/ipie/legacy/trial_wavefunction/multi_slater.py b/ipie/legacy/trial_wavefunction/multi_slater.py index 1553c9f2..d198a5c0 100644 --- a/ipie/legacy/trial_wavefunction/multi_slater.py +++ b/ipie/legacy/trial_wavefunction/multi_slater.py @@ -30,6 +30,7 @@ def __init__( nbasis=None, options={}, init=None, + cplx=False, verbose=False, orbs=None, ): @@ -48,7 +49,7 @@ def __init__( else: self.psi = wfn[1] imag_norm = numpy.sum(self.psi.imag.ravel() * self.psi.imag.ravel()) - if imag_norm <= 1e-8: + if (not cplx) and (imag_norm <= 1e-8): # print("# making trial wavefunction MO coefficient real") self.psi = numpy.array(self.psi.real, dtype=numpy.float64) self.coeffs = numpy.array(wfn[0], dtype=numpy.complex128) diff --git a/ipie/qmc/options.py b/ipie/qmc/options.py index be352469..437a5422 100644 --- a/ipie/qmc/options.py +++ b/ipie/qmc/options.py @@ -45,60 +45,32 @@ class QMCOpts(object): Initialised from a dict containing the following options, not all of which are required. - Parameters + Attributes ---------- - method : string - Which auxiliary field method are we using? Currently only CPMC is - implemented. + batched : bool + Whether to do batched calculations. nwalkers : int Number of walkers to propagate in a simulation. dt : float Timestep. nsteps : int Number of steps per block. - nmeasure : int - Frequency of energy measurements. + nblocks : int + Number of blocks. Total number of iterations = nblocks * nsteps. nstblz : int Frequency of Gram-Schmidt orthogonalisation steps. npop_control : int Frequency of population control. - temp : float - Temperature. Currently not used. - nequilibrate : int - Number of steps used for equilibration phase. Only used to fix local + pop_control_method : str + Population control method. + eqlb_time : float + Time scale of equilibration phase. Only used to fix local energy bound when using phaseless approximation. - importance_sampling : boolean - Are we using importance sampling. Default True. - hubbard_statonovich : string - Which hubbard stratonovich transformation are we using. Currently the - options are: - - - discrete : Use the discrete Hirsch spin transformation. - - opt_continuous : Use the continuous transformation for the Hubbard - model. - - generic : Use the generic transformation. To be used with Generic - system class. - - ffts : boolean - Use FFTS to diagonalise the kinetic energy propagator? Default False. - This may speed things up for larger lattices. - - Attributes - ---------- - cplx : boolean - Do we require complex wavefunctions? - mf_shift : float - Mean field shift for continuous Hubbard-Stratonovich transformation. - iut_fac : complex float - Stores i*(U*dt)**0.5 for continuous Hubbard-Stratonovich transformation. - ut_fac : float - Stores (U*dt) for continuous Hubbard-Stratonovich transformation. - mf_nsq : float - Stores M * mf_shift for continuous Hubbard-Stratonovich transformation. - local_energy_bound : float - Energy pound for continuous Hubbard-Stratonovich transformation. - mean_local_energy : float - Estimate for mean energy for continuous Hubbard-Stratonovich transformation. + neqlb : int + Number of time steps for the equilibration phase. Only used to fix the + local energy bound when using phaseless approximation. + rng_seed : int + The random number seed. """ # pylint: disable=dangerous-default-value @@ -133,6 +105,13 @@ def __init__(self, inputs={}, verbose=False): alias=["npop_control", "pop_control"], verbose=verbose, ) + self.pop_control_method = get_input_value( + inputs, + "pop_control_method", + default="pair_branch", + alias=["pop_control", "population_control"], + verbose=verbose, + ) self.eqlb_time = get_input_value( inputs, "equilibration_time", @@ -160,18 +139,25 @@ def __str__(self, verbose=0): class QMCParams: r"""Input options and certain constants / parameters derived from them. - Args: - num_walkers: number of walkers **per** core / task / computational unit. - total_num_walkers: The total number of walkers in the simulation. - timestep: The timestep delta_t - num_steps_per_block: Number of steps of propagation before estimators - are evaluated. - num_blocks: The number of blocks. Total number of iterations = - num_blocks * num_steps_per_block. - num_stblz: number of steps before QR stabilization of walkers is performed. - pop_control_freq: Frequency at which population control occurs. - rng_seed: The random number seed. If run in parallel the seeds on other - cores / threads are determined from this. + Attributes + ---------- + num_walkers : int + Number of walkers **per** core / task / computational unit. + total_num_walkers : int + The total number of walkers in the simulation. + timestep : float + The timestep delta_t + num_steps_per_block : int + Number of steps of propagation before estimators are evaluated. + num_blocks : int + Number of blocks. Total number of iterations = num_blocks * num_steps_per_block. + num_stblz : int + Number of steps before QR stabilization of walkers is performed. + pop_control_freq : int + Frequency at which population control occurs. + rng_seed : int + The random number seed. If run in parallel the seeds on other cores / + threads are determined from this. """ num_walkers: int diff --git a/ipie/qmc/tests/reference_data/ft_ueg_ecut1.0_rs1.0/reference.json b/ipie/qmc/tests/reference_data/ft_ueg_ecut1.0_rs1.0/reference.json index c278179f..d1cfd38b 100644 --- a/ipie/qmc/tests/reference_data/ft_ueg_ecut1.0_rs1.0/reference.json +++ b/ipie/qmc/tests/reference_data/ft_ueg_ecut1.0_rs1.0/reference.json @@ -1 +1 @@ -{"WeightFactor": [32.0, 47.94831398616111], "Weight": [32.0, 31.999999999999996], "ENumer": [853.4128425513718, 986.7912119123343], "EDenom": [32.0, 31.999999999999996], "ETotal": [26.66915132973037, 30.837225372260455], "E1Body": [28.374994808285745, 33.21707786343164], "E2Body": [-1.705843478555375, -2.379852491171182], "EHybrid": [0.0, 0.0], "Overlap": [1.0, 1.0], "Nav": [14.000000381209672, 16.375812925843057], "sys_info": {"nranks": 4, "branch": "hubbard_updates", "sha1": "618262bc7511a252e2f2bb3f23cc96fa4e8b9eb5", "numpy": {"version": "1.18.4", "path": "/usr/local/lib/python3.7/site-packages/numpy", "BLAS": {"lib": "openblas openblas", "path": "/usr/local/lib"}}, "scipy": {"version": "1.4.1", "path": "/usr/local/lib/python3.7/site-packages/scipy"}, "h5py": {"version": "2.10.0", "path": "/usr/local/lib/python3.7/site-packages/h5py"}, "mpi4py": {"version": "3.0.1", "path": "/usr/local/lib/python3.7/site-packages/mpi4py", "mpicc": "/usr/local/bin/mpicc"}}} \ No newline at end of file +{"WeightFactor": [32.0, 47.947639102045635], "Weight": [32.0, 31.999999999999993], "ENumer": [853.4128425513718, 986.7978362646822], "EDenom": [32.0, 31.999999999999993], "ETotal": [26.66915132973037, 30.837432383271327], "E1Body": [28.374994808285745, 33.217171356971804], "E2Body": [-1.705843478555375, -2.379738973700476], "EHybrid": [0.0, 0.0], "Overlap": [1.0, 1.0], "Nav": [14.000000381209672, 16.37587194751124]} \ No newline at end of file diff --git a/ipie/systems/utils.py b/ipie/systems/utils.py index d1fe1912..6822f2c4 100644 --- a/ipie/systems/utils.py +++ b/ipie/systems/utils.py @@ -44,11 +44,12 @@ def get_system(sys_opts=None, verbose=0, comm=None): if comm.rank == 0: print("# Error: Number of electrons not specified.") sys.exit() + nelec = (nup, ndown) system = Generic(nelec, verbose) else: if comm.rank == 0: print(f"# Error: unrecognized system name {sys_type}.") - raise ValueError + raise ValueError return system diff --git a/ipie/trial_wavefunction/single_det.py b/ipie/trial_wavefunction/single_det.py index abc043fd..3af4e335 100644 --- a/ipie/trial_wavefunction/single_det.py +++ b/ipie/trial_wavefunction/single_det.py @@ -40,7 +40,7 @@ def __init__(self, wavefunction, num_elec, num_basis, handler=MPIHandler(), verb self._max_num_dets = 1 imag_norm = numpy.sum(self.psi.imag.ravel() * self.psi.imag.ravel()) if imag_norm <= 1e-8: - # print("# making trial wavefunction MO coefficient real") + #print("# making trial wavefunction MO coefficient real") self.psi = numpy.array(self.psi.real, dtype=numpy.float64) self.psi0a = self.psi[:, : self.nalpha] diff --git a/ipie/trial_wavefunction/tests/test_noci.py b/ipie/trial_wavefunction/tests/test_noci.py index 1db2004b..d7acf4d2 100644 --- a/ipie/trial_wavefunction/tests/test_noci.py +++ b/ipie/trial_wavefunction/tests/test_noci.py @@ -33,3 +33,7 @@ def test_noci(): assert trial._rH1a.shape == (trial.num_dets, nalpha, nbasis) assert trial._rH1b.shape == (trial.num_dets, nbeta, nbasis) trial.calculate_energy(sys, ham) + + +if __name__ == '__main__': + test_noci() diff --git a/ipie/trial_wavefunction/tests/test_wavefunction_base.py b/ipie/trial_wavefunction/tests/test_wavefunction_base.py index bd7d168f..94e47cff 100644 --- a/ipie/trial_wavefunction/tests/test_wavefunction_base.py +++ b/ipie/trial_wavefunction/tests/test_wavefunction_base.py @@ -16,3 +16,7 @@ def test_wavefunction_base(): ) assert trial.nelec == (nalpha, nbeta) assert trial.nbasis == num_basis + + +if __name__ == '__main__': + test_wavefunction_base() diff --git a/ipie/utils/testing.py b/ipie/utils/testing.py index be0c37bb..c141766e 100644 --- a/ipie/utils/testing.py +++ b/ipie/utils/testing.py @@ -48,9 +48,11 @@ def generate_hamiltonian(nmo, nelec, cplx=False, sym=8, tol=1e-3): h1e = numpy.random.random((nmo, nmo)) if cplx: h1e = h1e + 1j * numpy.random.random((nmo, nmo)) + eri = numpy.random.normal(scale=0.01, size=(nmo, nmo, nmo, nmo)) if cplx: eri = eri + 1j * numpy.random.normal(scale=0.01, size=(nmo, nmo, nmo, nmo)) + # Restore symmetry to the integrals. if sym >= 4: # (ik|jl) = (jl|ik) diff --git a/ipie/walkers/pop_controller.py b/ipie/walkers/pop_controller.py index de9cf8ef..15f7efbf 100644 --- a/ipie/walkers/pop_controller.py +++ b/ipie/walkers/pop_controller.py @@ -5,7 +5,6 @@ from ipie.config import MPI from ipie.utils.backend import arraylib as xp - class PopControllerTimer: def __init__(self): self.start_time_const = 0.0 @@ -155,6 +154,7 @@ def get_buffer(walkers, iw): else: buff[s : s + 1] = xp.array(data[iw]) s += 1 + return buff @@ -173,15 +173,13 @@ def set_buffer(walkers, iw, buff): assert data.size % walkers.nwalkers == 0 # Only walker-specific data is being communicated if isinstance(data[iw], xp.ndarray): walkers.__dict__[d][iw] = xp.array( - buff[s : s + data[iw].size].reshape(data[iw].shape).copy() - ) + buff[s : s + data[iw].size].reshape(data[iw].shape).copy()) s += data[iw].size elif isinstance(data[iw], list): for ix, l in enumerate(data[iw]): if isinstance(l, (xp.ndarray)): walkers.__dict__[d][iw][ix] = xp.array( - buff[s : s + l.size].reshape(l.shape).copy() - ) + buff[s : s + l.size].reshape(l.shape).copy()) s += l.size elif isinstance(l, (int, float, complex)): walkers.__dict__[d][iw][ix] = buff[s] @@ -318,12 +316,10 @@ def pair_branch(walkers, comm, max_weight, min_weight, timer=PopControllerTimer( glob_inf_1.fill(1) glob_inf_2 = numpy.array( [[r for i in range(walkers.nwalkers)] for r in range(comm.size)], - dtype=numpy.int64, - ) + dtype=numpy.int64) glob_inf_3 = numpy.array( [[r for i in range(walkers.nwalkers)] for r in range(comm.size)], - dtype=numpy.int64, - ) + dtype=numpy.int64) timer.add_non_communication() @@ -341,15 +337,9 @@ def pair_branch(walkers, comm, max_weight, min_weight, timer=PopControllerTimer( # Rescale weights. glob_inf = numpy.zeros((walkers.nwalkers * comm.size, 4), dtype=numpy.float64) glob_inf[:, 0] = glob_inf_0.ravel() # contains walker |w_i| - glob_inf[:, 1] = ( - glob_inf_1.ravel() - ) # all initialized to 1 when it becomes 2 then it will be "branched" - glob_inf[:, 2] = ( - glob_inf_2.ravel() - ) # contain processor+walker indices (initial) (i.e., where walkers live) - glob_inf[:, 3] = ( - glob_inf_3.ravel() - ) # contain processor+walker indices (final) (i.e., where walkers live) + glob_inf[:, 1] = glob_inf_1.ravel() # all initialized to 1 when it becomes 2 then it will be "branched" + glob_inf[:, 2] = glob_inf_2.ravel() # contain processor+walker indices (initial) (i.e., where walkers live) + glob_inf[:, 3] = glob_inf_3.ravel() # contain processor+walker indices (final) (i.e., where walkers live) sort = numpy.argsort(glob_inf[:, 0], kind="mergesort") isort = numpy.argsort(sort, kind="mergesort") glob_inf = glob_inf[sort] diff --git a/ipie/walkers/tests/test_population_control.py b/ipie/walkers/tests/test_population_control.py index 4dcd6374..1ee26d5f 100644 --- a/ipie/walkers/tests/test_population_control.py +++ b/ipie/walkers/tests/test_population_control.py @@ -68,6 +68,10 @@ def test_pair_branch_batch(): batched_data.walkers.weight[iw], legacy_data.walker_handler.walkers[iw].weight, ) + assert numpy.allclose( + batched_data.walkers.unscaled_weight[iw], + legacy_data.walker_handler.walkers[iw].unscaled_weight, + ) assert pytest.approx(batched_data.walkers.weight[0]) == 0.2571750688329709 assert pytest.approx(batched_data.walkers.weight[1]) == 1.0843219322894988