Skip to content

Commit 78bba46

Browse files
fdmalonefdmalone
authored
[PBC Resources Estimates 3/4] Add computation of lambda (#823)
* Add k-point THC code. * Add k-thc notebook. * Add utilities. * Add reference data. * Add missing __init__ * Add missing init / skipifs. * Fix formatting. * Resolve review comments. * Address comments. * Remove utils. * No more utils. * More review comments. * Fix formatting. * Formatting + add map to gvec_logic. * Fix import issues. * Add utility classes to handler different PBC Hamiltonian factorizations. Add inits. * No more utils. * Tidyup. * Fix formatting. * Add functionality to compute PBC lambda values. * Fix formatting and docstrings. * Fix whitespace. * Update tests. * Fix import and mark slow. * Fix test failures. * Fix docstrings. * Revert change. * Address review comments. * Fix format.
1 parent 1700e24 commit 78bba46

10 files changed

Lines changed: 796 additions & 2 deletions

File tree

‎src/openfermion/resource_estimates/pbc/df/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,4 +8,4 @@
88
# distributed under the License is distributed on an "AS IS" BASIS,
99
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1010
# See the License for the specific language governing permissions and
11-
# limitations under the License.
11+
# limitations under the License.

‎src/openfermion/resource_estimates/pbc/df/compute_lambda_df.py‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
# coverage: ignore
2+
# Licensed under the Apache License, Version 2.0 (the "License");
3+
# you may not use this file except in compliance with the License.
4+
# You may obtain a copy of the License at
5+
#
6+
# http://www.apache.org/licenses/LICENSE-2.0
7+
#
8+
# Unless required by applicable law or agreed to in writing, software
9+
# distributed under the License is distributed on an "AS IS" BASIS,
10+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
# See the License for the specific language governing permissions and
12+
# limitations under the License.
13+
from dataclasses import dataclass
14+
import numpy as np
15+
import numpy.typing as npt
16+
17+
from openfermion.resource_estimates.pbc.df.df_integrals import (
18+
DFABKpointIntegrals,)
19+
from openfermion.resource_estimates.pbc.hamiltonian import (
20+
HamiltonianProperties,)
21+
22+
23+
@dataclass
24+
class DFHamiltonianProperties(HamiltonianProperties):
25+
"""Store for return values of compute_lambda function
26+
27+
Extension of HamiltonianProperties dataclass to also hold the number of
28+
retained eigenvalues (num_eig).
29+
"""
30+
31+
num_eig: int
32+
33+
34+
def compute_lambda(hcore: npt.NDArray,
35+
df_obj: DFABKpointIntegrals) -> DFHamiltonianProperties:
36+
"""Compute lambda for double-factorized Hamiltonian.
37+
38+
one-body term h_pq(k) = hcore_{pq}(k)
39+
- 0.5 * sum_{Q}sum_{r}(pkrQ|rQqk)
40+
+ 0.5 sum_{Q}sum_{r}(pkqk|rQrQ)
41+
The first term is the kinetic energy + pseudopotential (or electron-nuclear)
42+
second term is from rearranging two-body operator into chemist charge-charge
43+
type notation, and the third is from the one body term obtained when
44+
squaring the two-body A and B operators.
45+
46+
Arguments:
47+
hcore: List len(kpts) long of nmo x nmo complex hermitian arrays
48+
df_obj: DFABKpointIntegrals integral helper.
49+
50+
Returns:
51+
ham_props: A HamiltonianProperties instance containing Lambda values for
52+
DF hamiltonian.
53+
"""
54+
kpts = df_obj.kmf.kpts
55+
nkpts = len(kpts)
56+
one_body_mat = np.empty((len(kpts)), dtype=object)
57+
lambda_one_body = 0.0
58+
59+
for kidx in range(len(kpts)):
60+
# matrices for - 0.5 * sum_{Q}sum_{r}(pkrQ|rQqk)
61+
# and + 0.5 sum_{Q}sum_{r}(pkqk|rQrQ)
62+
h1_pos = np.zeros_like(hcore[kidx])
63+
h1_neg = np.zeros_like(hcore[kidx])
64+
for qidx in range(len(kpts)):
65+
# - 0.5 * sum_{Q}sum_{r}(pkrQ|rQqk)
66+
eri_kqqk_pqrs = df_obj.get_eri_exact([kidx, qidx, qidx, kidx])
67+
h1_neg -= (np.einsum("prrq->pq", eri_kqqk_pqrs, optimize=True) /
68+
nkpts)
69+
# + 0.5 sum_{Q}sum_{r}(pkqk|rQrQ)
70+
eri_kkqq_pqrs = df_obj.get_eri_exact([kidx, kidx, qidx, qidx])
71+
h1_pos += np.einsum("pqrr->pq", eri_kkqq_pqrs) / nkpts
72+
73+
one_body_mat[kidx] = hcore[kidx] + 0.5 * h1_neg + h1_pos
74+
one_eigs, _ = np.linalg.eigh(one_body_mat[kidx])
75+
lambda_one_body += np.sum(np.abs(one_eigs))
76+
77+
lambda_two_body = 0.0
78+
num_eigs = 0
79+
for qidx in range(len(kpts)):
80+
for nn in range(df_obj.naux):
81+
first_number_to_square = 0
82+
second_number_to_square = 0
83+
# sum up p,k eigenvalues
84+
for kidx in range(len(kpts)):
85+
# A and B are W
86+
if df_obj.amat_lambda_vecs[kidx, qidx, nn] is None:
87+
continue
88+
eigs_a_fixed_n_q = df_obj.amat_lambda_vecs[kidx, qidx,
89+
nn] / np.sqrt(nkpts)
90+
eigs_b_fixed_n_q = df_obj.bmat_lambda_vecs[kidx, qidx,
91+
nn] / np.sqrt(nkpts)
92+
first_number_to_square += np.sum(np.abs(eigs_a_fixed_n_q))
93+
num_eigs += len(eigs_a_fixed_n_q)
94+
if eigs_b_fixed_n_q is not None:
95+
second_number_to_square += np.sum(np.abs(eigs_b_fixed_n_q))
96+
num_eigs += len(eigs_b_fixed_n_q)
97+
98+
lambda_two_body += first_number_to_square**2
99+
lambda_two_body += second_number_to_square**2
100+
101+
lambda_two_body *= 0.25
102+
103+
lambda_tot = lambda_one_body + lambda_two_body
104+
df_data = DFHamiltonianProperties(
105+
lambda_total=lambda_tot,
106+
lambda_one_body=lambda_one_body,
107+
lambda_two_body=lambda_two_body,
108+
num_eig=num_eigs,
109+
)
110+
return df_data

‎src/openfermion/resource_estimates/pbc/df/compute_lambda_df_test.py‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
# coverage: ignore
2+
# Licensed under the Apache License, Version 2.0 (the "License");
3+
# you may not use this file except in compliance with the License.
4+
# You may obtain a copy of the License at
5+
#
6+
# http://www.apache.org/licenses/LICENSE-2.0
7+
#
8+
# Unless required by applicable law or agreed to in writing, software
9+
# distributed under the License is distributed on an "AS IS" BASIS,
10+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
# See the License for the specific language governing permissions and
12+
# limitations under the License.
13+
from functools import reduce
14+
import numpy as np
15+
import pytest
16+
17+
from openfermion.resource_estimates import HAVE_DEPS_FOR_RESOURCE_ESTIMATES
18+
19+
if HAVE_DEPS_FOR_RESOURCE_ESTIMATES:
20+
from pyscf.pbc import mp
21+
22+
from openfermion.resource_estimates.pbc.df.compute_lambda_df import (
23+
compute_lambda,)
24+
from openfermion.resource_estimates.pbc.df.df_integrals import (
25+
DFABKpointIntegrals,)
26+
from openfermion.resource_estimates.pbc.hamiltonian import (
27+
cholesky_from_df_ints,)
28+
from openfermion.resource_estimates.pbc.testing import (
29+
make_diamond_113_szv,)
30+
31+
32+
@pytest.mark.skipif(not HAVE_DEPS_FOR_RESOURCE_ESTIMATES,
33+
reason='pyscf and/or jax not installed.')
34+
def test_lambda_calc():
35+
mf = make_diamond_113_szv()
36+
mymp = mp.KMP2(mf)
37+
Luv = cholesky_from_df_ints(mymp)
38+
helper = DFABKpointIntegrals(cholesky_factor=Luv, kmf=mf)
39+
helper.double_factorize(thresh=1.0e-13)
40+
41+
hcore_ao = mf.get_hcore()
42+
hcore_mo = np.asarray([
43+
reduce(np.dot, (mo.T.conj(), hcore_ao[k], mo))
44+
for k, mo in enumerate(mf.mo_coeff)
45+
])
46+
47+
lambda_data = compute_lambda(hcore_mo, helper)
48+
assert np.isclose(lambda_data.lambda_total, 179.62240330857406)
49+
50+
lambda_two_body = 0
51+
lambda_two_body_v2 = 0
52+
nkpts = len(mf.kpts)
53+
for qidx in range(nkpts):
54+
aval_to_square = np.zeros((helper.naux), dtype=np.complex128)
55+
bval_to_square = np.zeros((helper.naux), dtype=np.complex128)
56+
57+
aval_to_square_v2 = np.zeros((helper.naux), dtype=np.complex128)
58+
bval_to_square_v2 = np.zeros((helper.naux), dtype=np.complex128)
59+
60+
for kidx in range(nkpts):
61+
Amats, Bmats = helper.build_A_B_n_q_k_from_chol(qidx, kidx)
62+
Amats /= np.sqrt(nkpts)
63+
Bmats /= np.sqrt(nkpts)
64+
wa, _ = np.linalg.eigh(Amats)
65+
wb, _ = np.linalg.eigh(Bmats)
66+
aval_to_square += np.einsum("npq->n", np.abs(Amats)**2)
67+
bval_to_square += np.einsum("npq->n", np.abs(Bmats)**2)
68+
69+
aval_to_square_v2 += np.sum(np.abs(wa)**2, axis=-1)
70+
bval_to_square_v2 += np.sum(np.abs(wb)**2, axis=-1)
71+
assert np.allclose(
72+
np.sum(np.abs(wa)**2, axis=-1),
73+
np.einsum("npq->n",
74+
np.abs(Amats)**2),
75+
)
76+
77+
lambda_two_body += np.sum(aval_to_square)
78+
lambda_two_body += np.sum(bval_to_square)
79+
80+
lambda_two_body_v2 += np.sum(aval_to_square_v2)
81+
lambda_two_body_v2 += np.sum(bval_to_square_v2)
82+
83+
assert np.isclose(lambda_two_body, lambda_two_body_v2)

‎src/openfermion/resource_estimates/pbc/hamiltonian/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,4 +16,4 @@
1616
if HAVE_DEPS_FOR_RESOURCE_ESTIMATES:
1717
from .hamiltonian import (build_hamiltonian,
1818
build_momentum_transfer_mapping,
19-
cholesky_from_df_ints)
19+
cholesky_from_df_ints, HamiltonianProperties)

‎src/openfermion/resource_estimates/pbc/sf/compute_lambda_sf.py‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
# coverage: ignore
2+
# Licensed under the Apache License, Version 2.0 (the "License");
3+
# you may not use this file except in compliance with the License.
4+
# You may obtain a copy of the License at
5+
#
6+
# http://www.apache.org/licenses/LICENSE-2.0
7+
#
8+
# Unless required by applicable law or agreed to in writing, software
9+
# distributed under the License is distributed on an "AS IS" BASIS,
10+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
# See the License for the specific language governing permissions and
12+
# limitations under the License.
13+
from dataclasses import dataclass
14+
import numpy as np
15+
import numpy.typing as npt
16+
from openfermion.resource_estimates.pbc.hamiltonian import (
17+
HamiltonianProperties,)
18+
19+
from openfermion.resource_estimates.pbc.sf.sf_integrals import (
20+
SingleFactorization,)
21+
22+
23+
@dataclass
24+
class SFHamiltonianProperties(HamiltonianProperties):
25+
"""Store for return values of compute_lambda function
26+
27+
Extension of HamiltonianProperties dataclass to also hold the number of
28+
retained cholesky vectors (num_aux).
29+
"""
30+
31+
num_aux: int
32+
33+
34+
def compute_lambda(hcore: npt.NDArray,
35+
sf_obj: SingleFactorization) -> SFHamiltonianProperties:
36+
"""Lambda for single-factorized Hamiltonian.
37+
38+
Compute one-body and two-body lambda for qubitization of
39+
single-factorized Hamiltonian.
40+
41+
one-body term h_pq(k) = hcore_{pq}(k)
42+
- 0.5 * sum_{Q}sum_{r}(pkrQ|rQqk)
43+
+ sum_{Q}sum_{r}(pkqk|rQrQ)
44+
The first term is the kinetic energy + pseudopotential
45+
(or electron-nuclear), second term is from rearranging two-body operator
46+
into chemist charge-charge type notation, and the third is from the one body
47+
term obtained when squaring the two-body A and B operators.
48+
49+
two-body term V = 0.5 sum_{Q}sum_{n}(A_{n}(Q)^2 +_ B_{n}(Q)^2)
50+
or V = 0.5 sum_{Q}sum_{n'}W_{n}(Q)^{2} where n' is twice the range of n.
51+
lambda is 0.5sum_{Q}sum_{n'}(sum_{p,q}^{N_{k}N/2}|Re[W_{p,q}(Q)^{n}]| +
52+
|Im[W_{pq}(Q)^{n}]|)^{2}
53+
54+
Args:
55+
hcore: List len(kpts) long of nmo x nmo complex hermitian arrays
56+
sf_obj: SingleFactorization integral helper object.
57+
58+
Returns:
59+
ham_props: A HamiltonianProperties instance containing Lambda values for
60+
SF hamiltonian.
61+
"""
62+
kpts = sf_obj.kmf.kpts
63+
nkpts = len(kpts)
64+
one_body_mat = np.empty((len(kpts)), dtype=object)
65+
lambda_one_body = 0.0
66+
67+
old_naux = sf_obj.naux # need to reset naux for one-body computation
68+
sf_obj.naux = sf_obj.chol[0, 0].shape[0]
69+
70+
for kidx in range(len(kpts)):
71+
# matrices for - 0.5 * sum_{Q}sum_{r}(pkrQ|rQqk)
72+
# and + 0.5 sum_{Q}sum_{r}(pkqk|rQrQ)
73+
h1_pos = np.zeros_like(hcore[kidx])
74+
h1_neg = np.zeros_like(hcore[kidx])
75+
for qidx in range(len(kpts)):
76+
# - 0.5 * sum_{Q}sum_{r}(pkrQ|rQqk)
77+
eri_kqqk_pqrs = sf_obj.get_eri([kidx, qidx, qidx, kidx])
78+
h1_neg -= (np.einsum("prrq->pq", eri_kqqk_pqrs, optimize=True) /
79+
nkpts)
80+
# + sum_{Q}sum_{r}(pkqk|rQrQ)
81+
eri_kkqq_pqrs = sf_obj.get_eri([kidx, kidx, qidx, qidx])
82+
h1_pos += np.einsum("pqrr->pq", eri_kkqq_pqrs) / nkpts
83+
84+
one_body_mat[kidx] = hcore[kidx] + 0.5 * h1_neg + h1_pos
85+
lambda_one_body += np.sum(
86+
np.abs(one_body_mat[kidx].real) + np.abs(one_body_mat[kidx].imag))
87+
88+
############################################################################
89+
#
90+
# \lambda_{V} = \frac 12 \sum_{\Q}\sum_{n}^{M}\left
91+
# (\sum_{\K,pq}(|\Rea[L_{p\K,q\K-\Q,n}]| +
92+
# |\Ima[L_{p\K,q\K-\Q,n}]|)\right)^{2}
93+
#
94+
# chol = [nkpts, nkpts, naux, nao, nao]
95+
#
96+
############################################################################
97+
sf_obj.naux = old_naux # reset naux to original value
98+
# this part needs to change
99+
lambda_two_body = 0.0
100+
for qidx in range(len(kpts)):
101+
# A and B are W
102+
A, B = sf_obj.build_AB_from_chol(qidx) # [naux, nao * nk, nao * nk]
103+
A /= np.sqrt(nkpts)
104+
B /= np.sqrt(nkpts)
105+
# sum_q sum_n (sum_{pq} |Re{A_{pq}^n}| + |Im{A_{pq}^n|)^2
106+
lambda_two_body += np.sum(
107+
np.einsum("npq->n",
108+
np.abs(A.real) + np.abs(A.imag))**2)
109+
lambda_two_body += np.sum(
110+
np.einsum("npq->n",
111+
np.abs(B.real) + np.abs(B.imag))**2)
112+
113+
lambda_two_body *= 0.5
114+
115+
lambda_tot = lambda_one_body + lambda_two_body
116+
sf_data = SFHamiltonianProperties(
117+
lambda_total=lambda_tot,
118+
lambda_one_body=lambda_one_body,
119+
lambda_two_body=lambda_two_body,
120+
num_aux=sf_obj.naux,
121+
)
122+
return sf_data

0 commit comments

Comments
 (0)