Merge pull request #248 from JoonhoLee-Group/allow_noci

Allow for NOCI trial.

examples/11-trexio/run_afqmc.py

@@ -85,11 +85,9 @@
 
 wavefunction = (coeff, occa, occb)
 
-from ipie.trial_wavefunction.particle_hole import ParticleHoleWicks
+from ipie.trial_wavefunction.particle_hole import ParticleHole
 
-trial = ParticleHoleWicks(
-    wavefunction, (nocca, noccb), nbasis, num_dets_for_props=len(wavefunction[0])
-)
+trial = ParticleHole(wavefunction, (nocca, noccb), nbasis, num_dets_for_props=len(wavefunction[0]))
 trial.build()
 trial.half_rotate(ham, comm=comm)
 

ipie/estimators/energy.py

@@ -15,46 +15,33 @@
 # Author: Fionn Malone <[email protected]>
 #
 
+import plum
+
 from ipie.estimators.estimator_base import EstimatorBase
 from ipie.estimators.local_energy_batch import (
     local_energy_batch,
     local_energy_multi_det_trial_batch,
 )
-from ipie.utils.backend import arraylib as xp
-
-from ipie.trial_wavefunction.particle_hole import (
-    ParticleHoleNaive,
-    ParticleHoleWicks,
-    ParticleHoleWicksNonChunked,
-    ParticleHoleWicksSlow,
-)
-from ipie.trial_wavefunction.single_det import SingleDet
-
+from ipie.estimators.local_energy_noci import local_energy_noci
+from ipie.estimators.local_energy_sd import local_energy_single_det_uhf
 from ipie.estimators.local_energy_wicks import (
     local_energy_multi_det_trial_wicks_batch,
     local_energy_multi_det_trial_wicks_batch_opt,
     local_energy_multi_det_trial_wicks_batch_opt_chunked,
 )
-from ipie.estimators.local_energy_sd import (
-    local_energy_single_det_uhf,
-)
-
+from ipie.hamiltonians.generic import GenericComplexChol, GenericRealChol
 from ipie.systems.generic import Generic
-from ipie.hamiltonians.generic import GenericRealChol, GenericComplexChol
+from ipie.trial_wavefunction.noci import NOCI
+from ipie.trial_wavefunction.particle_hole import (
+    ParticleHole,
+    ParticleHoleNaive,
+    ParticleHoleNonChunked,
+    ParticleHoleSlow,
+)
+from ipie.trial_wavefunction.single_det import SingleDet
+from ipie.utils.backend import arraylib as xp
 from ipie.walkers.uhf_walkers import UHFWalkers
 
-import plum
-
-
-# Single dispatch
-_dispatcher = {
-    ParticleHoleNaive: local_energy_multi_det_trial_batch,
-    ParticleHoleWicks: local_energy_multi_det_trial_wicks_batch_opt_chunked,
-    ParticleHoleWicksNonChunked: local_energy_multi_det_trial_wicks_batch_opt,
-    ParticleHoleWicksSlow: local_energy_multi_det_trial_wicks_batch,
-    SingleDet: local_energy_batch,
-}
-
 
 @plum.dispatch
 def local_energy(
@@ -88,7 +75,7 @@ def local_energy(
     system: Generic,
     hamiltonian: GenericRealChol,
     walkers: UHFWalkers,
-    trial: ParticleHoleWicks,
+    trial: ParticleHole,
 ):
     return local_energy_multi_det_trial_wicks_batch_opt_chunked(system, hamiltonian, walkers, trial)
 
@@ -98,7 +85,7 @@ def local_energy(
     system: Generic,
     hamiltonian: GenericRealChol,
     walkers: UHFWalkers,
-    trial: ParticleHoleWicksNonChunked,
+    trial: ParticleHoleNonChunked,
 ):
     return local_energy_multi_det_trial_wicks_batch_opt(system, hamiltonian, walkers, trial)
 
@@ -108,11 +95,16 @@ def local_energy(
     system: Generic,
     hamiltonian: GenericRealChol,
     walkers: UHFWalkers,
-    trial: ParticleHoleWicksSlow,
+    trial: ParticleHoleSlow,
 ):
     return local_energy_multi_det_trial_wicks_batch(system, hamiltonian, walkers, trial)
 
 
+@plum.dispatch
+def local_energy(system: Generic, hamiltonian: GenericRealChol, walkers: UHFWalkers, trial: NOCI):
+    return local_energy_noci(system, hamiltonian, walkers, trial)
+
+
 class EnergyEstimator(EstimatorBase):
     def __init__(
         self,

ipie/estimators/greens_function.py

@@ -17,18 +17,18 @@
 #          linusjoonho <[email protected]>
 #
 
-from ipie.utils.misc import is_cupy
-
-from ipie.trial_wavefunction.particle_hole import ParticleHoleWicks
-from ipie.trial_wavefunction.noci import NOCI
-from ipie.trial_wavefunction.single_det import SingleDet
-
+from ipie.estimators.greens_function_multi_det import (
+    greens_function_multi_det_wicks_opt,
+    greens_function_noci,
+)
 from ipie.estimators.greens_function_single_det import (
-    greens_function_single_det_batch,
     greens_function_single_det,
+    greens_function_single_det_batch,
 )
-from ipie.estimators.greens_function_multi_det import greens_function_multi_det
-from ipie.estimators.greens_function_multi_det import greens_function_multi_det_wicks_opt
+from ipie.trial_wavefunction.noci import NOCI
+from ipie.trial_wavefunction.particle_hole import ParticleHole
+from ipie.trial_wavefunction.single_det import SingleDet
+from ipie.utils.misc import is_cupy
 
 
 def compute_greens_function(walker_batch, trial):
@@ -59,8 +59,8 @@ def get_greens_function(trial):
         else:
             compute_greens_function = greens_function_single_det
     elif isinstance(trial, NOCI):
-        compute_greens_function = greens_function_multi_det
-    elif isinstance(trial, ParticleHoleWicks):
+        compute_greens_function = greens_function_noci
+    elif isinstance(trial, ParticleHole):
         compute_greens_function = greens_function_multi_det_wicks_opt
     else:
         compute_greens_function = None

ipie/estimators/greens_function_multi_det.py

@@ -2,15 +2,16 @@
 import scipy.linalg
 from numba import jit
 
+from ipie.estimators.kernels.cpu import wicks as wk
+from ipie.legacy.estimators.greens_function import gab_mod
 from ipie.propagation.overlap import (
     compute_determinants_batched,
-    get_overlap_one_det_wicks,
     get_cofactor_matrix_batched,
+    get_det_matrix_batched,
+    get_overlap_one_det_wicks,
+    reduce_to_CI_tensor,
 )
 from ipie.utils.linalg import minor_mask
-from ipie.propagation.overlap import get_det_matrix_batched, reduce_to_CI_tensor
-from ipie.estimators.kernels.cpu import wicks as wk
-from ipie.legacy.estimators.greens_function import gab_mod
 
 
 def greens_function_multi_det(walker_batch, trial, build_full=False):
@@ -72,6 +73,64 @@ def greens_function_multi_det(walker_batch, trial, build_full=False):
     return tot_ovlps
 
 
+def greens_function_noci(walker_batch, trial, build_full=False):
+    """Compute walker's green's function.
+
+    Parameters
+    ----------
+    walker_batch : object
+        MultiDetTrialWalkerBatch object.
+    trial : object
+        Trial wavefunction object.
+    Returns
+    -------
+    det : float64 / complex128
+        Determinant of overlap matrix.
+    """
+    nup = walker_batch.nup
+    walker_batch.Ga.fill(0.0)
+    walker_batch.Gb.fill(0.0)
+    tot_ovlps = numpy.zeros(walker_batch.nwalkers, dtype=numpy.complex128)
+    for iw in range(walker_batch.nwalkers):
+        for ix, detix in enumerate(trial.psi):
+            Oup = numpy.dot(walker_batch.phia[iw].T, detix[:, :nup].conj())
+            sign_a, logdet_a = numpy.linalg.slogdet(Oup)
+            ovlpa = sign_a * numpy.exp(logdet_a)
+            walker_batch.det_ovlpas[iw, ix] = ovlpa
+            if abs(ovlpa) < 1e-16:
+                continue
+
+            Odn = numpy.dot(walker_batch.phib[iw].T, detix[:, nup:].conj())
+            sign_b, logdet_b = numpy.linalg.slogdet(Odn)
+            ovlpb = sign_b * numpy.exp(logdet_b)
+            ovlp = ovlpa * ovlpb
+            walker_batch.det_ovlpbs[iw, ix] = ovlpb
+            if abs(ovlp) < 1e-16:
+                continue
+
+            inv_ovlp = scipy.linalg.inv(Oup)
+            walker_batch.Ghalfa[ix, iw, :, :] = numpy.dot(inv_ovlp, walker_batch.phia[iw].T)
+            walker_batch.Gia[ix, iw, :, :] = numpy.dot(
+                detix[:, :nup].conj(), walker_batch.Ghalfa[ix, iw, :, :]
+            )
+
+            inv_ovlp = scipy.linalg.inv(Odn)
+            walker_batch.Ghalfb[ix, iw, :, :] = numpy.dot(inv_ovlp, walker_batch.phib[iw].T)
+            walker_batch.Gib[ix, iw, :, :] = numpy.dot(
+                detix[:, nup:].conj(), walker_batch.Ghalfb[ix, iw, :, :]
+            )
+
+            tot_ovlps[iw] += trial.coeffs[ix].conj() * ovlp
+
+            walker_batch.Ga[iw] += walker_batch.Gia[ix, iw, :, :] * ovlp * trial.coeffs[ix].conj()
+            walker_batch.Gb[iw] += walker_batch.Gib[ix, iw, :, :] * ovlp * trial.coeffs[ix].conj()
+
+        walker_batch.Ga[iw] /= tot_ovlps[iw]
+        walker_batch.Gb[iw] /= tot_ovlps[iw]
+
+    return tot_ovlps
+
+
 def greens_function_multi_det_wicks(walker_batch, trial, build_full=False):
     """Compute walker's green's function using Wick's theorem.
 
@@ -1123,8 +1182,10 @@ def greens_function_multi_det_wicks_opt(walker_batch, trial, build_full=False):
     walker_batch.Ga += numpy.einsum("w,wpq->wpq", ovlps, G0a, optimize=True)
     walker_batch.Gb += numpy.einsum("w,wpq->wpq", ovlps, G0b, optimize=True)
     # intermediates for contribution 2 (connected diagrams)
-    build_CI_single_excitation_opt(walker_batch, trial, c_phasea_ovlpb, c_phaseb_ovlpa)
-    build_CI_double_excitation_opt(walker_batch, trial, c_phasea_ovlpb, c_phaseb_ovlpa)
+    if trial.max_excite >= 1:
+        build_CI_single_excitation_opt(walker_batch, trial, c_phasea_ovlpb, c_phaseb_ovlpa)
+    if trial.max_excite >= 2:
+        build_CI_double_excitation_opt(walker_batch, trial, c_phasea_ovlpb, c_phaseb_ovlpa)
     if trial.max_excite >= 3:
         build_CI_triple_excitation_opt(walker_batch, trial, c_phasea_ovlpb, c_phaseb_ovlpa)
     for iexcit in range(4, trial.max_excite + 1):

ipie/estimators/greens_function_single_det.py

@@ -2,6 +2,45 @@
 from ipie.utils.backend import synchronize
 
 
+def gab_mod_ovlp(A, B):
+    r"""One-particle Green's function.
+
+    This actually returns 1-G since it's more useful, i.e.,
+
+    .. math::
+        \langle \phi_A|c_i^{\dagger}c_j|\phi_B\rangle =
+        [B(A^{\dagger}B)^{-1}A^{\dagger}]_{ji}
+
+    where :math:`A,B` are the matrices representing the Slater determinants
+    :math:`|\psi_{A,B}\rangle`.
+
+    For example, usually A would represent (an element of) the trial wavefunction.
+
+    .. warning::
+        Assumes A and B are not orthogonal.
+
+    Parameters
+    ----------
+    A : :class:`numpy.ndarray`
+        Matrix representation of the bra used to construct G.
+    B : :class:`numpy.ndarray`
+        Matrix representation of the ket used to construct G.
+
+    Returns
+    -------
+    GAB : :class:`numpy.ndarray`
+        the green's function.
+    Ghalf : :class:`numpy.ndarray`
+        the half rotated green's function.
+    inv_O : :class:`numpy.ndarray`
+        Inverse of the overlap matrix.
+    """
+    inv_O = xp.linalg.inv(xp.dot(B.T, A.conj()))
+    GHalf = xp.dot(inv_O, B.T)
+    G = xp.dot(A.conj(), GHalf)
+    return (G, GHalf, inv_O)
+
+
 def greens_function_single_det_ghf(walkers, trial):
     det = []
     for iw in range(walkers.nwalkers):

ipie/estimators/kernels/cpu/wicks.py

@@ -258,7 +258,7 @@ def build_det_matrix(cre, anh, mapping, offset, G0, det_mat):
 # Green's function
 
 
-@jit(nopython=False, fastmath=False)
+@jit(nopython=True, fastmath=True)
 def reduce_CI_singles(cre, anh, mapping, phases, CI):
     """Reduction to CI intermediate for singles.
 

ipie/estimators/local_energy.py

@@ -17,10 +17,9 @@
 #
 
 from ipie.estimators.generic import local_energy_cholesky_opt, local_energy_generic_cholesky
+from ipie.estimators.greens_function_single_det import gab_mod_ovlp
 from ipie.legacy.estimators.ci import get_hmatel
-
-# from ipie.legacy.estimators.local_energy import local_energy_G as legacy_local_energy_G
-# from ipie.hamiltonians.generic import Generic
+from ipie.utils.backend import arraylib as xp
 
 
 def local_energy_G(system, hamiltonian, trial, G, Ghalf):
@@ -58,8 +57,6 @@ def local_energy_G(system, hamiltonian, trial, G, Ghalf):
         )
     else:
         return local_energy_generic_cholesky(system, hamiltonian, G)
-    # else:
-    # return legacy_local_energy_G(system, hamiltonian, trial, G, Ghalf)
 
 
 def local_energy(system, hamiltonian, walker, trial):
@@ -108,3 +105,22 @@ def variational_energy_ortho_det(system, ham, occs, coeffs):
                 one_body += e1b.conj()
                 two_body += e2b.conj()
     return evar / denom, one_body / denom, two_body / denom
+
+
+def variational_energy_noci(system, hamiltonian, trial):
+    weight = 0
+    energies = 0
+    denom = 0
+    for i, (Ba, Bb) in enumerate(zip(trial.psia, trial.psib)):
+        for j, (Aa, Ab) in enumerate(zip(trial.psia, trial.psib)):
+            # construct "local" green's functions for each component of A
+            Gup, _, inv_O_up = gab_mod_ovlp(Ba, Aa)
+            Gdn, _, inv_O_dn = gab_mod_ovlp(Bb, Ab)
+            ovlp = 1.0 / (xp.linalg.det(inv_O_up) * xp.linalg.det(inv_O_dn))
+            weight = (trial.coeffs[i].conj() * trial.coeffs[j]) * ovlp
+            G = xp.array([Gup, Gdn])
+            # Ghalf = [Ghalfa, Ghalfb]
+            e = xp.array(local_energy_G(system, hamiltonian, trial, G, Ghalf=None))
+            energies += weight * e
+            denom += weight
+    return tuple(energies / denom)