From 0a8ec41fd3d28e64325193bdc603dd9dba9854c8 Mon Sep 17 00:00:00 2001 From: Fionn Malone Date: Sat, 9 Dec 2023 17:29:17 -0800 Subject: [PATCH] Remove dot_real_cmplx function from wicks kernels. (#276) This function assumed the cholesky intermediate was real in an attempt to save some multiplications. This assumption is incorrect in general. This fortunately did not lead to incorrect results as the formula held for both real and complex input A, which explains why the unit tests passed. Removes this in favour of * for element wise multiplication. Also force more complex walkers in the tests to exercise the code path. Fixes #275 --- ipie/estimators/kernels/cpu/wicks.py | 95 ++++----------------- ipie/estimators/tests/test_generic_phmsd.py | 24 +++--- 2 files changed, 28 insertions(+), 91 deletions(-) diff --git a/ipie/estimators/kernels/cpu/wicks.py b/ipie/estimators/kernels/cpu/wicks.py index ee4e9e0c..1761f676 100644 --- a/ipie/estimators/kernels/cpu/wicks.py +++ b/ipie/estimators/kernels/cpu/wicks.py @@ -18,36 +18,6 @@ import numpy from numba import jit - -# element wise -@jit(nopython=True, fastmath=True) -def dot_real_cplx( - A, - B_real, - B_cplx, -): - """Element wise multiplication of a real number with a complex one. - - C = A * B - - Numba complains if types aren't matched so split it up. - - Parameters - ---------- - A : float - Real number / array. - B : complex - Complex number / array. - - Returns - ------- - C : complex - result - """ - - return A * B_real + 1j * (A * B_cplx) - - # Overlap # Note mapping arrays account for occupied indices not matching compressed @@ -551,8 +521,6 @@ def fill_os_doubles(cre, anh, mapping, offset, G0, chol_factor, spin_buffer, det ndets = cre.shape[0] nwalkers = G0.shape[0] for iw in range(nwalkers): - G0_real = G0[iw].real.copy() - G0_imag = G0[iw].imag.copy() for idet in range(ndets): p = mapping[cre[idet, 0]] q = anh[idet, 0] @@ -563,10 +531,10 @@ def fill_os_doubles(cre, anh, mapping, offset, G0, chol_factor, spin_buffer, det ro = cre[idet, 1] + offset so = anh[idet, 1] + offset spin_buffer[iw, start + idet, :] = ( - dot_real_cplx(chol_factor[iw, q, p, :], G0_real[ro, so], G0_imag[ro, so]) - - dot_real_cplx(chol_factor[iw, s, p, :], G0_real[ro, qo], G0_imag[ro, qo]) - - dot_real_cplx(chol_factor[iw, q, r, :], G0_real[po, so], G0_imag[po, so]) - + dot_real_cplx(chol_factor[iw, s, r, :], G0_real[po, qo], G0_imag[po, qo]) + chol_factor[iw, q, p, :] * G0[iw, ro, so] + - chol_factor[iw, s, p, :] * G0[iw, ro, qo] + - chol_factor[iw, q, r, :] * G0[iw, po, so] + + chol_factor[iw, s, r, :] * G0[iw, po, qo] ) @@ -616,41 +584,23 @@ def fill_os_triples(cre, anh, mapping, offset, G0w, chol_factor, spin_buffer, de to = cre[idet, 2] + offset uo = anh[idet, 2] + offset cofac = G0[ro, so] * G0[to, uo] - G0[ro, uo] * G0[to, so] - spin_buffer[iw, start + idet, :] = dot_real_cplx( - chol_factor[iw, q, p], cofac.real, cofac.imag - ) + spin_buffer[iw, start + idet, :] = chol_factor[iw, q, p] * cofac cofac = G0[ro, qo] * G0[to, uo] - G0[ro, uo] * G0[to, qo] - spin_buffer[iw, start + idet, :] -= dot_real_cplx( - chol_factor[iw, s, p], cofac.real, cofac.imag - ) + spin_buffer[iw, start + idet, :] -= chol_factor[iw, s, p] * cofac cofac = G0[ro, qo] * G0[to, so] - G0[ro, so] * G0[to, qo] - spin_buffer[iw, start + idet, :] += dot_real_cplx( - chol_factor[iw, u, p], cofac.real, cofac.imag - ) + spin_buffer[iw, start + idet, :] += chol_factor[iw, u, p] * cofac cofac = G0[po, so] * G0[to, uo] - G0[to, so] * G0[po, uo] - spin_buffer[iw, start + idet, :] -= dot_real_cplx( - chol_factor[iw, q, r], cofac.real, cofac.imag - ) + spin_buffer[iw, start + idet, :] -= chol_factor[iw, q, r] * cofac cofac = G0[po, qo] * G0[to, uo] - G0[to, qo] * G0[po, uo] - spin_buffer[iw, start + idet, :] += dot_real_cplx( - chol_factor[iw, s, r], cofac.real, cofac.imag - ) + spin_buffer[iw, start + idet, :] += chol_factor[iw, s, r] * cofac cofac = G0[po, qo] * G0[to, so] - G0[to, qo] * G0[po, so] - spin_buffer[iw, start + idet, :] -= dot_real_cplx( - chol_factor[iw, u, r], cofac.real, cofac.imag - ) + spin_buffer[iw, start + idet, :] -= chol_factor[iw, u, r] * cofac cofac = G0[po, so] * G0[ro, uo] - G0[ro, so] * G0[po, uo] - spin_buffer[iw, start + idet, :] += dot_real_cplx( - chol_factor[iw, q, t], cofac.real, cofac.imag - ) + spin_buffer[iw, start + idet, :] += chol_factor[iw, q, t] * cofac cofac = G0[po, qo] * G0[ro, uo] - G0[ro, qo] * G0[po, uo] - spin_buffer[iw, start + idet, :] -= dot_real_cplx( - chol_factor[iw, s, t], cofac.real, cofac.imag - ) + spin_buffer[iw, start + idet, :] -= chol_factor[iw, s, t] * cofac cofac = G0[po, qo] * G0[ro, so] - G0[ro, qo] * G0[po, so] - spin_buffer[iw, start + idet, :] += dot_real_cplx( - chol_factor[iw, u, t], cofac.real, cofac.imag - ) + spin_buffer[iw, start + idet, :] += chol_factor[iw, u, t] * cofac @jit(nopython=True, fastmath=True) @@ -824,16 +774,11 @@ def reduce_os_spin_factor(ps, qs, mapping, phase, cof_mat, chol_factor, spin_buf nwalker = chol_factor.shape[0] ndet = cof_mat.shape[1] start = det_sls.start - # assert ndet == det_sls.end - det_sls.start for iw in range(nwalker): for idet in range(ndet): det_cofactor = phase * numpy.linalg.det(cof_mat[iw, idet]) p = mapping[ps[idet]] - spin_buffer[iw, start + idet] += dot_real_cplx( - chol_factor[iw, qs[idet], p], - det_cofactor.real, - det_cofactor.imag, - ) + spin_buffer[iw, start + idet] += chol_factor[iw, qs[idet], p] * det_cofactor @jit(nopython=True, fastmath=True) @@ -920,19 +865,11 @@ def reduce_ss_spin_factor( chol_a = chol_factor[iw, ss[idet], r] chol_b = chol_factor[iw, qs[idet], p] cont_ab = numpy.dot(chol_a, chol_b) - spin_buffer[iw, start + idet] += dot_real_cplx( - cont_ab, - det_cofactor.real, - det_cofactor.imag, - ) + spin_buffer[iw, start + idet] += cont_ab * det_cofactor chol_c = chol_factor[iw, qs[idet], r] chol_d = chol_factor[iw, ss[idet], p] cont_cd = numpy.dot(chol_c, chol_d) - spin_buffer[iw, start + idet] -= dot_real_cplx( - cont_cd, - det_cofactor.real, - det_cofactor.imag, - ) + spin_buffer[iw, start + idet] -= cont_cd * det_cofactor @jit(nopython=True, fastmath=True) diff --git a/ipie/estimators/tests/test_generic_phmsd.py b/ipie/estimators/tests/test_generic_phmsd.py index e05a1ea1..3774ca46 100644 --- a/ipie/estimators/tests/test_generic_phmsd.py +++ b/ipie/estimators/tests/test_generic_phmsd.py @@ -46,7 +46,12 @@ ) from ipie.utils.misc import dotdict from ipie.utils.mpi import MPIHandler -from ipie.utils.testing import generate_hamiltonian, get_random_phmsd, get_random_phmsd_opt +from ipie.utils.testing import ( + generate_hamiltonian, + get_random_phmsd, + get_random_phmsd_opt, + shaped_normal, +) from ipie.walkers.uhf_walkers import UHFWalkersParticleHole from ipie.walkers.walkers_dispatch import UHFWalkersTrial @@ -337,7 +342,7 @@ def test_phmsd_local_energy(): # Test PH type wavefunction. # wfn, init = get_random_phmsd(system.nup, system.ndown, ham.nbasis, ndet=5, init=True) wfn, init = get_random_phmsd( - system.nup, system.ndown, ham.nbasis, ndet=3000, init=True, cmplx=False + system.nup, system.ndown, ham.nbasis, ndet=3000, init=True, cmplx=True ) ci, oa, ob = wfn wfn_2 = (ci[::50], oa[::50], ob[::50]) # Get high excitation determinants too @@ -476,10 +481,9 @@ def test_kernels_energy(): test = numpy.zeros((nwalkers, ndets, nchol), dtype=numpy.complex128) slices_alpha, slices_beta = trial.slices_alpha, trial.slices_beta nbasis = ham.nbasis - from ipie.utils.testing import shaped_normal - Laa = shaped_normal((nwalkers, nbasis, system.nup, nchol)) - Lbb = shaped_normal((nwalkers, nbasis, system.ndown, nchol)) + Laa = shaped_normal((nwalkers, nbasis, system.nup, nchol), cmplx=True) + Lbb = shaped_normal((nwalkers, nbasis, system.ndown, nchol), cmplx=True) # 1. fill_opp_spin_factors_batched_singles( trial.cre_ex_b[1], @@ -820,8 +824,7 @@ def test_kernels_gf_active_space(): prop = PhaselessGeneric(qmc["dt"]) prop.build(ham, trial_ref) - I = numpy.eye(nmo) - init = numpy.hstack([I[:, : nelec[0]], I[:, : nelec[1]]]) + init = shaped_normal((nmo, system.ne), cmplx=True) walker_batch_ref = UHFWalkersTrial( trial_ref, init, system.nup, system.ndown, ham.nbasis, nwalkers, MPIHandler() @@ -1057,9 +1060,8 @@ def test_kernels_energy_active_space(): ) walker_batch_ref.CIa.fill(0.0 + 0.0j) walker_batch_ref.CIb.fill(0.0 + 0.0j) - from ipie.utils.testing import shaped_normal - Lbb = shaped_normal((nwalkers, nmo, system.ndown, nchol)) + Lbb = shaped_normal((nwalkers, nmo, system.ndown, nchol), cmplx=True) slices_alpha, slices_beta = trial_test.slices_alpha, trial_test.slices_beta assert trial_ref.nfrozen != trial_test.nfrozen # 1. @@ -1356,10 +1358,9 @@ def test_phmsd_local_energy_active_space_polarised(): ecore=0, # options={"symmetry": False}, ) - from ipie.utils.testing import get_random_phmsd_opt, shaped_normal wfn, init = get_random_phmsd_opt(7, 5, nact, ndet=100, init=True) - init = shaped_normal((nmo, system.ne)) + init = shaped_normal((nmo, system.ne), cmplx=True) ci, occa, occb = wfn core = [0, 1] with_core_a = numpy.array( @@ -1465,7 +1466,6 @@ def test_phmsd_local_energy_active_space_non_aufbau(): ecore=0, # options={"symmetry": False}, ) - from ipie.utils.testing import get_random_phmsd_opt, shaped_normal wfn, init = get_random_phmsd_opt(7, 7, nact, ndet=100, init=True, cmplx_coeffs=False) init = shaped_normal((nmo, system.ne))