diff --git a/src/openfermion/resource_estimates/pbc/df/__init__.py b/src/openfermion/resource_estimates/pbc/df/__init__.py index f808bf29..53f003f0 100644 --- a/src/openfermion/resource_estimates/pbc/df/__init__.py +++ b/src/openfermion/resource_estimates/pbc/df/__init__.py @@ -8,4 +8,4 @@ # 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. \ No newline at end of file +# limitations under the License. diff --git a/src/openfermion/resource_estimates/pbc/df/compute_lambda_df.py b/src/openfermion/resource_estimates/pbc/df/compute_lambda_df.py new file mode 100644 index 00000000..875dd4da --- /dev/null +++ b/src/openfermion/resource_estimates/pbc/df/compute_lambda_df.py @@ -0,0 +1,110 @@ +# coverage: ignore +# 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. +from dataclasses import dataclass +import numpy as np +import numpy.typing as npt + +from openfermion.resource_estimates.pbc.df.df_integrals import ( + DFABKpointIntegrals,) +from openfermion.resource_estimates.pbc.hamiltonian import ( + HamiltonianProperties,) + + +@dataclass +class DFHamiltonianProperties(HamiltonianProperties): + """Store for return values of compute_lambda function + + Extension of HamiltonianProperties dataclass to also hold the number of + retained eigenvalues (num_eig). + """ + + num_eig: int + + +def compute_lambda(hcore: npt.NDArray, + df_obj: DFABKpointIntegrals) -> DFHamiltonianProperties: + """Compute lambda for double-factorized Hamiltonian. + + one-body term h_pq(k) = hcore_{pq}(k) + - 0.5 * sum_{Q}sum_{r}(pkrQ|rQqk) + + 0.5 sum_{Q}sum_{r}(pkqk|rQrQ) + The first term is the kinetic energy + pseudopotential (or electron-nuclear) + second term is from rearranging two-body operator into chemist charge-charge + type notation, and the third is from the one body term obtained when + squaring the two-body A and B operators. + + Arguments: + hcore: List len(kpts) long of nmo x nmo complex hermitian arrays + df_obj: DFABKpointIntegrals integral helper. + + Returns: + ham_props: A HamiltonianProperties instance containing Lambda values for + DF hamiltonian. + """ + kpts = df_obj.kmf.kpts + nkpts = len(kpts) + one_body_mat = np.empty((len(kpts)), dtype=object) + lambda_one_body = 0.0 + + for kidx in range(len(kpts)): + # matrices for - 0.5 * sum_{Q}sum_{r}(pkrQ|rQqk) + # and + 0.5 sum_{Q}sum_{r}(pkqk|rQrQ) + h1_pos = np.zeros_like(hcore[kidx]) + h1_neg = np.zeros_like(hcore[kidx]) + for qidx in range(len(kpts)): + # - 0.5 * sum_{Q}sum_{r}(pkrQ|rQqk) + eri_kqqk_pqrs = df_obj.get_eri_exact([kidx, qidx, qidx, kidx]) + h1_neg -= (np.einsum("prrq->pq", eri_kqqk_pqrs, optimize=True) / + nkpts) + # + 0.5 sum_{Q}sum_{r}(pkqk|rQrQ) + eri_kkqq_pqrs = df_obj.get_eri_exact([kidx, kidx, qidx, qidx]) + h1_pos += np.einsum("pqrr->pq", eri_kkqq_pqrs) / nkpts + + one_body_mat[kidx] = hcore[kidx] + 0.5 * h1_neg + h1_pos + one_eigs, _ = np.linalg.eigh(one_body_mat[kidx]) + lambda_one_body += np.sum(np.abs(one_eigs)) + + lambda_two_body = 0.0 + num_eigs = 0 + for qidx in range(len(kpts)): + for nn in range(df_obj.naux): + first_number_to_square = 0 + second_number_to_square = 0 + # sum up p,k eigenvalues + for kidx in range(len(kpts)): + # A and B are W + if df_obj.amat_lambda_vecs[kidx, qidx, nn] is None: + continue + eigs_a_fixed_n_q = df_obj.amat_lambda_vecs[kidx, qidx, + nn] / np.sqrt(nkpts) + eigs_b_fixed_n_q = df_obj.bmat_lambda_vecs[kidx, qidx, + nn] / np.sqrt(nkpts) + first_number_to_square += np.sum(np.abs(eigs_a_fixed_n_q)) + num_eigs += len(eigs_a_fixed_n_q) + if eigs_b_fixed_n_q is not None: + second_number_to_square += np.sum(np.abs(eigs_b_fixed_n_q)) + num_eigs += len(eigs_b_fixed_n_q) + + lambda_two_body += first_number_to_square**2 + lambda_two_body += second_number_to_square**2 + + lambda_two_body *= 0.25 + + lambda_tot = lambda_one_body + lambda_two_body + df_data = DFHamiltonianProperties( + lambda_total=lambda_tot, + lambda_one_body=lambda_one_body, + lambda_two_body=lambda_two_body, + num_eig=num_eigs, + ) + return df_data diff --git a/src/openfermion/resource_estimates/pbc/df/compute_lambda_df_test.py b/src/openfermion/resource_estimates/pbc/df/compute_lambda_df_test.py new file mode 100644 index 00000000..7835598e --- /dev/null +++ b/src/openfermion/resource_estimates/pbc/df/compute_lambda_df_test.py @@ -0,0 +1,83 @@ +# coverage: ignore +# 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. +from functools import reduce +import numpy as np +import pytest + +from openfermion.resource_estimates import HAVE_DEPS_FOR_RESOURCE_ESTIMATES + +if HAVE_DEPS_FOR_RESOURCE_ESTIMATES: + from pyscf.pbc import mp + + from openfermion.resource_estimates.pbc.df.compute_lambda_df import ( + compute_lambda,) + from openfermion.resource_estimates.pbc.df.df_integrals import ( + DFABKpointIntegrals,) + from openfermion.resource_estimates.pbc.hamiltonian import ( + cholesky_from_df_ints,) + from openfermion.resource_estimates.pbc.testing import ( + make_diamond_113_szv,) + + +@pytest.mark.skipif(not HAVE_DEPS_FOR_RESOURCE_ESTIMATES, + reason='pyscf and/or jax not installed.') +def test_lambda_calc(): + mf = make_diamond_113_szv() + mymp = mp.KMP2(mf) + Luv = cholesky_from_df_ints(mymp) + helper = DFABKpointIntegrals(cholesky_factor=Luv, kmf=mf) + helper.double_factorize(thresh=1.0e-13) + + hcore_ao = mf.get_hcore() + hcore_mo = np.asarray([ + reduce(np.dot, (mo.T.conj(), hcore_ao[k], mo)) + for k, mo in enumerate(mf.mo_coeff) + ]) + + lambda_data = compute_lambda(hcore_mo, helper) + assert np.isclose(lambda_data.lambda_total, 179.62240330857406) + + lambda_two_body = 0 + lambda_two_body_v2 = 0 + nkpts = len(mf.kpts) + for qidx in range(nkpts): + aval_to_square = np.zeros((helper.naux), dtype=np.complex128) + bval_to_square = np.zeros((helper.naux), dtype=np.complex128) + + aval_to_square_v2 = np.zeros((helper.naux), dtype=np.complex128) + bval_to_square_v2 = np.zeros((helper.naux), dtype=np.complex128) + + for kidx in range(nkpts): + Amats, Bmats = helper.build_A_B_n_q_k_from_chol(qidx, kidx) + Amats /= np.sqrt(nkpts) + Bmats /= np.sqrt(nkpts) + wa, _ = np.linalg.eigh(Amats) + wb, _ = np.linalg.eigh(Bmats) + aval_to_square += np.einsum("npq->n", np.abs(Amats)**2) + bval_to_square += np.einsum("npq->n", np.abs(Bmats)**2) + + aval_to_square_v2 += np.sum(np.abs(wa)**2, axis=-1) + bval_to_square_v2 += np.sum(np.abs(wb)**2, axis=-1) + assert np.allclose( + np.sum(np.abs(wa)**2, axis=-1), + np.einsum("npq->n", + np.abs(Amats)**2), + ) + + lambda_two_body += np.sum(aval_to_square) + lambda_two_body += np.sum(bval_to_square) + + lambda_two_body_v2 += np.sum(aval_to_square_v2) + lambda_two_body_v2 += np.sum(bval_to_square_v2) + + assert np.isclose(lambda_two_body, lambda_two_body_v2) diff --git a/src/openfermion/resource_estimates/pbc/hamiltonian/__init__.py b/src/openfermion/resource_estimates/pbc/hamiltonian/__init__.py index f0d52f8e..b3d33813 100644 --- a/src/openfermion/resource_estimates/pbc/hamiltonian/__init__.py +++ b/src/openfermion/resource_estimates/pbc/hamiltonian/__init__.py @@ -16,4 +16,4 @@ if HAVE_DEPS_FOR_RESOURCE_ESTIMATES: from .hamiltonian import (build_hamiltonian, build_momentum_transfer_mapping, - cholesky_from_df_ints) + cholesky_from_df_ints, HamiltonianProperties) diff --git a/src/openfermion/resource_estimates/pbc/sf/compute_lambda_sf.py b/src/openfermion/resource_estimates/pbc/sf/compute_lambda_sf.py new file mode 100644 index 00000000..d1029d13 --- /dev/null +++ b/src/openfermion/resource_estimates/pbc/sf/compute_lambda_sf.py @@ -0,0 +1,122 @@ +# coverage: ignore +# 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. +from dataclasses import dataclass +import numpy as np +import numpy.typing as npt +from openfermion.resource_estimates.pbc.hamiltonian import ( + HamiltonianProperties,) + +from openfermion.resource_estimates.pbc.sf.sf_integrals import ( + SingleFactorization,) + + +@dataclass +class SFHamiltonianProperties(HamiltonianProperties): + """Store for return values of compute_lambda function + + Extension of HamiltonianProperties dataclass to also hold the number of + retained cholesky vectors (num_aux). + """ + + num_aux: int + + +def compute_lambda(hcore: npt.NDArray, + sf_obj: SingleFactorization) -> SFHamiltonianProperties: + """Lambda for single-factorized Hamiltonian. + + Compute one-body and two-body lambda for qubitization of + single-factorized Hamiltonian. + + one-body term h_pq(k) = hcore_{pq}(k) + - 0.5 * sum_{Q}sum_{r}(pkrQ|rQqk) + + sum_{Q}sum_{r}(pkqk|rQrQ) + The first term is the kinetic energy + pseudopotential + (or electron-nuclear), second term is from rearranging two-body operator + into chemist charge-charge type notation, and the third is from the one body + term obtained when squaring the two-body A and B operators. + + two-body term V = 0.5 sum_{Q}sum_{n}(A_{n}(Q)^2 +_ B_{n}(Q)^2) + or V = 0.5 sum_{Q}sum_{n'}W_{n}(Q)^{2} where n' is twice the range of n. + lambda is 0.5sum_{Q}sum_{n'}(sum_{p,q}^{N_{k}N/2}|Re[W_{p,q}(Q)^{n}]| + + |Im[W_{pq}(Q)^{n}]|)^{2} + + Args: + hcore: List len(kpts) long of nmo x nmo complex hermitian arrays + sf_obj: SingleFactorization integral helper object. + + Returns: + ham_props: A HamiltonianProperties instance containing Lambda values for + SF hamiltonian. + """ + kpts = sf_obj.kmf.kpts + nkpts = len(kpts) + one_body_mat = np.empty((len(kpts)), dtype=object) + lambda_one_body = 0.0 + + old_naux = sf_obj.naux # need to reset naux for one-body computation + sf_obj.naux = sf_obj.chol[0, 0].shape[0] + + for kidx in range(len(kpts)): + # matrices for - 0.5 * sum_{Q}sum_{r}(pkrQ|rQqk) + # and + 0.5 sum_{Q}sum_{r}(pkqk|rQrQ) + h1_pos = np.zeros_like(hcore[kidx]) + h1_neg = np.zeros_like(hcore[kidx]) + for qidx in range(len(kpts)): + # - 0.5 * sum_{Q}sum_{r}(pkrQ|rQqk) + eri_kqqk_pqrs = sf_obj.get_eri([kidx, qidx, qidx, kidx]) + h1_neg -= (np.einsum("prrq->pq", eri_kqqk_pqrs, optimize=True) / + nkpts) + # + sum_{Q}sum_{r}(pkqk|rQrQ) + eri_kkqq_pqrs = sf_obj.get_eri([kidx, kidx, qidx, qidx]) + h1_pos += np.einsum("pqrr->pq", eri_kkqq_pqrs) / nkpts + + one_body_mat[kidx] = hcore[kidx] + 0.5 * h1_neg + h1_pos + lambda_one_body += np.sum( + np.abs(one_body_mat[kidx].real) + np.abs(one_body_mat[kidx].imag)) + + ############################################################################ + # + # \lambda_{V} = \frac 12 \sum_{\Q}\sum_{n}^{M}\left + # (\sum_{\K,pq}(|\Rea[L_{p\K,q\K-\Q,n}]| + + # |\Ima[L_{p\K,q\K-\Q,n}]|)\right)^{2} + # + # chol = [nkpts, nkpts, naux, nao, nao] + # + ############################################################################ + sf_obj.naux = old_naux # reset naux to original value + # this part needs to change + lambda_two_body = 0.0 + for qidx in range(len(kpts)): + # A and B are W + A, B = sf_obj.build_AB_from_chol(qidx) # [naux, nao * nk, nao * nk] + A /= np.sqrt(nkpts) + B /= np.sqrt(nkpts) + # sum_q sum_n (sum_{pq} |Re{A_{pq}^n}| + |Im{A_{pq}^n|)^2 + lambda_two_body += np.sum( + np.einsum("npq->n", + np.abs(A.real) + np.abs(A.imag))**2) + lambda_two_body += np.sum( + np.einsum("npq->n", + np.abs(B.real) + np.abs(B.imag))**2) + + lambda_two_body *= 0.5 + + lambda_tot = lambda_one_body + lambda_two_body + sf_data = SFHamiltonianProperties( + lambda_total=lambda_tot, + lambda_one_body=lambda_one_body, + lambda_two_body=lambda_two_body, + num_aux=sf_obj.naux, + ) + return sf_data diff --git a/src/openfermion/resource_estimates/pbc/sf/compute_lambda_sf_test.py b/src/openfermion/resource_estimates/pbc/sf/compute_lambda_sf_test.py new file mode 100644 index 00000000..75aa204c --- /dev/null +++ b/src/openfermion/resource_estimates/pbc/sf/compute_lambda_sf_test.py @@ -0,0 +1,96 @@ +# coverage: ignore +# 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. +from functools import reduce +import numpy as np +import pytest + +from openfermion.resource_estimates import HAVE_DEPS_FOR_RESOURCE_ESTIMATES +if HAVE_DEPS_FOR_RESOURCE_ESTIMATES: + from ase.build import bulk + + from pyscf.pbc import gto, scf, mp + from pyscf.pbc.tools import pyscf_ase + + from openfermion.resource_estimates.pbc.sf.compute_lambda_sf import ( + compute_lambda,) + from openfermion.resource_estimates.pbc.sf.sf_integrals import ( + SingleFactorization,) + from openfermion.resource_estimates.pbc.hamiltonian import ( + build_hamiltonian, + cholesky_from_df_ints, + ) + from openfermion.resource_estimates.pbc.testing import ( + make_diamond_113_szv,) + + +@pytest.mark.skipif(not HAVE_DEPS_FOR_RESOURCE_ESTIMATES, + reason='pyscf and/or jax not installed.') +def test_lambda_calc(): + mf = make_diamond_113_szv() + hcore, Luv = build_hamiltonian(mf) + helper = SingleFactorization(cholesky_factor=Luv, kmf=mf) + lambda_data = compute_lambda(hcore, helper) + assert np.isclose(lambda_data.lambda_total, 2123.4342903006627) + + +@pytest.mark.skipif(not HAVE_DEPS_FOR_RESOURCE_ESTIMATES, + reason='pyscf and/or jax not installed.') +def test_padding(): + ase_atom = bulk("H", "bcc", a=2.0, cubic=True) + cell = gto.Cell() + cell.atom = pyscf_ase.ase_atoms_to_pyscf(ase_atom) + cell.a = ase_atom.cell[:].copy() + cell.basis = "gth-szv" + cell.pseudo = "gth-hf-rev" + cell.verbose = 0 + cell.build() + + kmesh = [1, 2, 2] + kpts = cell.make_kpts(kmesh) + nkpts = len(kpts) + mf = scf.KRHF(cell, kpts).rs_density_fit() + mf.with_df.mesh = cell.mesh + mf.kernel() + + from pyscf.pbc.mp.kmp2 import _add_padding + + mymp = mp.KMP2(mf) + Luv_padded = cholesky_from_df_ints(mymp) + mo_coeff_padded = _add_padding(mymp, mymp.mo_coeff, mymp.mo_energy)[0] + helper = SingleFactorization(cholesky_factor=Luv_padded, kmf=mf) + assert mf.mo_coeff[0].shape[-1] != mo_coeff_padded[0].shape[-1] + + hcore_ao = mf.get_hcore() + hcore_no_padding = np.asarray([ + reduce(np.dot, (mo.T.conj(), hcore_ao[k], mo)) + for k, mo in enumerate(mf.mo_coeff) + ]) + hcore_padded = np.asarray([ + reduce(np.dot, (mo.T.conj(), hcore_ao[k], mo)) + for k, mo in enumerate(mo_coeff_padded) + ]) + assert hcore_no_padding[0].shape != hcore_padded[0].shape + assert np.isclose(np.sum(hcore_no_padding), np.sum(hcore_padded)) + Luv_no_padding = cholesky_from_df_ints(mymp, pad_mos_with_zeros=False) + for k1 in range(nkpts): + for k2 in range(nkpts): + assert np.isclose(np.sum(Luv_padded[k1, k2]), + np.sum(Luv_no_padding[k1, k2])) + + helper_no_padding = SingleFactorization(cholesky_factor=Luv_no_padding, + kmf=mf) + lambda_data_pad = compute_lambda(hcore_no_padding, helper_no_padding) + helper = SingleFactorization(cholesky_factor=Luv_padded, kmf=mf) + lambda_data_no_pad = compute_lambda(hcore_padded, helper) + assert np.isclose(lambda_data_pad.lambda_total, + lambda_data_no_pad.lambda_total) diff --git a/src/openfermion/resource_estimates/pbc/sparse/compute_lambda_sparse.py b/src/openfermion/resource_estimates/pbc/sparse/compute_lambda_sparse.py new file mode 100644 index 00000000..a192f491 --- /dev/null +++ b/src/openfermion/resource_estimates/pbc/sparse/compute_lambda_sparse.py @@ -0,0 +1,99 @@ +# coverage: ignore +# 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. +from dataclasses import dataclass +import numpy as np +import numpy.typing as npt + +from openfermion.resource_estimates.pbc.hamiltonian import ( + HamiltonianProperties,) +from openfermion.resource_estimates.pbc.sparse.sparse_integrals import ( + SparseFactorization,) + + +@dataclass +class SparseHamiltonianProperties(HamiltonianProperties): + """Store for return values of compute_lambda function + + Extension of HamiltonianProperties dataclass to also hold the number of + retained matrix elements (num_sym_unique). + """ + + num_sym_unique: int + + +def compute_lambda(hcore: npt.NDArray, sparse_int_obj: SparseFactorization + ) -> SparseHamiltonianProperties: + """Compute lambda value for sparse method + + Arguments: + hcore: array of hcore(k) by kpoint. k-point order + is pyscf order generated for this problem. + sparse_int_obj: The sparse integral object that is used + to compute eris and the number of unique + terms. + + Returns: + ham_props: A SparseHamiltonianProperties instance containing Lambda + values for the sparse hamiltonian. + """ + kpts = sparse_int_obj.kmf.kpts + nkpts = len(kpts) + one_body_mat = np.empty((len(kpts)), dtype=object) + lambda_one_body = 0.0 + + import time + + for kidx in range(len(kpts)): + # matrices for - 0.5 * sum_{Q}sum_{r}(pkrQ|rQqk) + # and + 0.5 sum_{Q}sum_{r}(pkqk|rQrQ) + h1_pos = np.zeros_like(hcore[kidx]) + h1_neg = np.zeros_like(hcore[kidx]) + for qidx in range(len(kpts)): + # - 0.5 * sum_{Q}sum_{r}(pkrQ|rQqk) + eri_kqqk_pqrs = sparse_int_obj.get_eri_exact( + [kidx, qidx, qidx, kidx]) + h1_neg -= (np.einsum("prrq->pq", eri_kqqk_pqrs, optimize=True) / + nkpts) + # + sum_{Q}sum_{r}(pkqk|rQrQ) + eri_kkqq_pqrs = sparse_int_obj.get_eri_exact( + [kidx, kidx, qidx, qidx]) + + h1_pos += np.einsum("pqrr->pq", eri_kkqq_pqrs) / nkpts + + one_body_mat[kidx] = hcore[kidx] + 0.5 * h1_neg + h1_pos + lambda_one_body += np.sum(np.abs(one_body_mat[kidx].real)) + np.sum( + np.abs(one_body_mat[kidx].imag)) + + lambda_two_body = 0.0 + nkpts = len(kpts) + # recall (k, k-q|k'-q, k') + for kidx in range(nkpts): + for kpidx in range(nkpts): + for qidx in range(nkpts): + kmq_idx = sparse_int_obj.k_transfer_map[qidx, kidx] + kpmq_idx = sparse_int_obj.k_transfer_map[qidx, kpidx] + test_eri_block = ( + sparse_int_obj.get_eri([kidx, kmq_idx, kpmq_idx, kpidx]) / + nkpts) + lambda_two_body += np.sum(np.abs(test_eri_block.real)) + np.sum( + np.abs(test_eri_block.imag)) + + lambda_tot = lambda_one_body + lambda_two_body + sparse_data = SparseHamiltonianProperties( + lambda_total=lambda_tot, + lambda_one_body=lambda_one_body, + lambda_two_body=lambda_two_body, + num_sym_unique=sparse_int_obj.get_total_unique_terms_above_thresh( + return_nk_counter=False), + ) + return sparse_data diff --git a/src/openfermion/resource_estimates/pbc/sparse/compute_lambda_sparse_test.py b/src/openfermion/resource_estimates/pbc/sparse/compute_lambda_sparse_test.py new file mode 100644 index 00000000..b0a4e1ce --- /dev/null +++ b/src/openfermion/resource_estimates/pbc/sparse/compute_lambda_sparse_test.py @@ -0,0 +1,45 @@ +# coverage: ignore +# 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. +from functools import reduce + +import numpy as np +import pytest + +from openfermion.resource_estimates import HAVE_DEPS_FOR_RESOURCE_ESTIMATES + +if HAVE_DEPS_FOR_RESOURCE_ESTIMATES: + from pyscf.pbc import mp + from openfermion.resource_estimates.pbc.sparse.\ + compute_lambda_sparse import compute_lambda + from openfermion.resource_estimates.pbc.sparse.sparse_integrals import ( + SparseFactorization,) + from openfermion.resource_estimates.pbc.hamiltonian import ( + cholesky_from_df_ints,) + from openfermion.resource_estimates.pbc.testing import ( + make_diamond_113_szv,) + + +@pytest.mark.skipif(not HAVE_DEPS_FOR_RESOURCE_ESTIMATES, + reason='pyscf and/or jax not installed.') +def test_lambda_sparse(): + mf = make_diamond_113_szv() + mymp = mp.KMP2(mf) + Luv = cholesky_from_df_ints(mymp) + helper = SparseFactorization(cholesky_factor=Luv, kmf=mf) + + hcore_ao = mf.get_hcore() + hcore_mo = np.asarray([ + reduce(np.dot, (mo.T.conj(), hcore_ao[k], mo)) + for k, mo in enumerate(mf.mo_coeff) + ]) + compute_lambda(hcore_mo, helper) diff --git a/src/openfermion/resource_estimates/pbc/thc/compute_lambda_thc.py b/src/openfermion/resource_estimates/pbc/thc/compute_lambda_thc.py new file mode 100644 index 00000000..5240b97a --- /dev/null +++ b/src/openfermion/resource_estimates/pbc/thc/compute_lambda_thc.py @@ -0,0 +1,150 @@ +# coverage: ignore +# 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. +from dataclasses import dataclass +from typing import Tuple +import numpy as np +import numpy.typing as npt + +from openfermion.resource_estimates.pbc.thc.thc_integrals import ( + KPTHCDoubleTranslation,) +from openfermion.resource_estimates.pbc.hamiltonian import ( + HamiltonianProperties,) + + +@dataclass +class THCHamiltonianProperties(HamiltonianProperties): + """Store for return values of compute_lambda function + + Extension of HamiltonianProperties dataclass to also hold the THC dimension + (num_sym_unique). + """ + + thc_dim: int + + +def compute_lambda_real( + h1: npt.NDArray, + etaPp: npt.NDArray, + MPQ: npt.NDArray, + chol: npt.NDArray, +) -> Tuple[float, float, float]: + """Compute lambda assuming real THC factors (molecular way) + + Avoids the need of a molecular object as in the molecular code. + + Args: + h1: one-body hamiltonian + etaPp: THC leaf tensor + MPQ: THC central tensor. + chol: Cholesky factors. + + Returns: + lambda_tot: Total lambda + lambda_one_body: One-body lambda + lambda_two_body: Two-body lambda + """ + + # projecting into the THC basis requires each THC factor mu to be nrmlzd. + # we roll the normalization constant into the central tensor zeta + SPQ = etaPp.dot( + etaPp.T) # (nthc x norb) x (norb x nthc) -> (nthc x nthc) metric + cP = np.diag(np.diag( + SPQ)) # grab diagonal elements. equivalent to np.diag(np.diagonal(SPQ)) + # no sqrts because we have two normalized THC vectors (index by mu and nu) + # on each side. + MPQ_normalized = cP.dot(MPQ).dot(cP) # get normalized zeta in Eq. 11 & 12 + + lambda_z = np.sum(np.abs(MPQ_normalized)) * 0.5 # Eq. 13 + # NCR: originally Joonho's code add np.einsum('llij->ij', eri_thc) + # NCR: I don't know how much this matters. + T = (h1 - 0.5 * np.einsum("nil,nlj->ij", chol, chol, optimize=True) + + np.einsum("nll,nij->ij", chol, chol, optimize=True)) # Eq. 3 + Eq. 18 + # e, v = np.linalg.eigh(T) + e = np.linalg.eigvalsh(T) # only need eigenvalues + lambda_T = np.sum( + np.abs(e)) # Eq. 19. NOTE: sum over spin orbitals removes 1/2 factor + + lambda_tot = lambda_z + lambda_T # Eq. 20 + + return lambda_tot, lambda_T, lambda_z + + +def compute_lambda(hcore: npt.NDArray, + thc_obj: KPTHCDoubleTranslation) -> HamiltonianProperties: + """Compute one-body and two-body lambda for qubitization the THC LCU. + + Args: + hcore: List len(kpts) long of nmo x nmo complex hermitian arrays + thc_obj: Object of KPTHCDoubleTranslation + + Returns: + ham_props: A HamiltonianProperties instance containing Lambda values for + THC hamiltonian. + """ + kpts = thc_obj.kmf.kpts + nkpts = len(kpts) + one_body_mat = np.empty((len(kpts)), dtype=object) + lambda_one_body = 0.0 + + for kidx in range(len(kpts)): + # matrices for - 0.5 * sum_{Q}sum_{r}(pkrQ|rQqk) + # and + 0.5 sum_{Q}sum_{r}(pkqk|rQrQ) + h1_neg = np.zeros_like(hcore[kidx]) + for qidx in range(len(kpts)): + # - 0.5 * sum_{Q}sum_{r}(pkrQ|rQqk) + eri_kqqk_pqrs = thc_obj.get_eri_exact([kidx, qidx, qidx, kidx]) + h1_neg -= (np.einsum("prrq->pq", eri_kqqk_pqrs, optimize=True) / + nkpts) + + one_body_mat[kidx] = hcore[kidx] + 0.5 * h1_neg # + h1_pos + one_eigs, _ = np.linalg.eigh(one_body_mat[kidx]) + lambda_one_body += np.sum(np.abs(one_eigs)) + + # projecting into the THC basis requires each THC factor mu to be + # normalized. we roll the normalization constant into the central tensor + # zeta # no sqrts because we have two normalized thc vectors + # (index by mu and nu) on each side. + norm_kP = (np.einsum("kpP,kpP->kP", + thc_obj.chi.conj(), + thc_obj.chi, + optimize=True)**0.5) + lambda_two_body = 0 + for iq, zeta_Q in enumerate(thc_obj.zeta): + # xy einsum subscript indexes G index. + for ik in range(nkpts): + ik_minus_q = thc_obj.k_transfer_map[iq, ik] + gpq = thc_obj.g_mapping[iq, ik] + for ik_prime in range(nkpts): + ik_prime_minus_q = thc_obj.k_transfer_map[iq, ik_prime] + gsr = thc_obj.g_mapping[iq, ik_prime] + norm_left = norm_kP[ik] * norm_kP[ik_minus_q] + norm_right = norm_kP[ik_prime_minus_q] * norm_kP[ik_prime] + MPQ_normalized = (np.einsum( + "P,PQ,Q->PQ", + norm_left, + zeta_Q[gpq, gsr], + norm_right, + optimize=True, + ) / nkpts) + lambda_two_body += np.sum(np.abs(MPQ_normalized.real)) + lambda_two_body += np.sum(np.abs(MPQ_normalized.imag)) + lambda_two_body *= 2 + + lambda_tot = lambda_one_body + lambda_two_body + lambda_data = THCHamiltonianProperties( + lambda_total=lambda_tot, + lambda_one_body=lambda_one_body, + lambda_two_body=lambda_two_body, + thc_dim=thc_obj.chi.shape[-1], + ) + return lambda_data diff --git a/src/openfermion/resource_estimates/pbc/thc/compute_lambda_thc_test.py b/src/openfermion/resource_estimates/pbc/thc/compute_lambda_thc_test.py new file mode 100644 index 00000000..67e50d20 --- /dev/null +++ b/src/openfermion/resource_estimates/pbc/thc/compute_lambda_thc_test.py @@ -0,0 +1,89 @@ +# coverage: ignore +# 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. +from functools import reduce + +import numpy as np +import pytest + +from openfermion.resource_estimates import HAVE_DEPS_FOR_RESOURCE_ESTIMATES +if HAVE_DEPS_FOR_RESOURCE_ESTIMATES: + from pyscf.pbc import gto, mp, scf + from openfermion.resource_estimates.pbc.hamiltonian import ( + cholesky_from_df_ints,) + from openfermion.resource_estimates.pbc.thc.factorizations.thc_jax import ( + kpoint_thc_via_isdf,) + from openfermion.resource_estimates.pbc.thc.compute_lambda_thc import ( + compute_lambda,) + from openfermion.resource_estimates.pbc.thc.thc_integrals import ( + KPTHCDoubleTranslation,) + + +@pytest.mark.skipif(not HAVE_DEPS_FOR_RESOURCE_ESTIMATES, + reason='pyscf and/or jax not installed.') +@pytest.mark.slow +def test_kpoint_thc_lambda(): + cell = gto.Cell() + cell.atom = """ + C 0.000000000000 0.000000000000 0.000000000000 + C 1.685068664391 1.685068664391 1.685068664391 + """ + cell.basis = "gth-szv" + cell.pseudo = "gth-hf-rev" + cell.a = """ + 0.000000000, 3.370137329, 3.370137329 + 3.370137329, 0.000000000, 3.370137329 + 3.370137329, 3.370137329, 0.000000000""" + cell.unit = "B" + cell.verbose = 0 + cell.mesh = [11] * 3 + cell.build() + + kmesh = [1, 1, 2] + kpts = cell.make_kpts(kmesh) + mf = scf.KRHF(cell, kpts) + mf.kernel() + # + # Build kpoint THC eris + # + # + rsmf = scf.KRHF(mf.cell, mf.kpts).rs_density_fit() + # Force same MOs as FFTDF at least + rsmf.mo_occ = mf.mo_occ + rsmf.mo_coeff = mf.mo_coeff + rsmf.mo_energy = mf.mo_energy + rsmf.with_df.mesh = mf.cell.mesh + mymp = mp.KMP2(rsmf) + Luv = cholesky_from_df_ints(mymp) + cthc = 4 + num_thc = cthc * mf.mo_coeff[0].shape[-1] + np.random.seed(7) + kpt_thc, _ = kpoint_thc_via_isdf( + mf, + Luv, + num_thc, + perform_adagrad_opt=False, + perform_bfgs_opt=True, + bfgs_maxiter=10, + verbose=False, + ) + hcore_ao = mf.get_hcore() + hcore_mo = np.asarray([ + reduce(np.dot, (mo.T.conj(), hcore_ao[k], mo)) + for k, mo in enumerate(mf.mo_coeff) + ]) + helper = KPTHCDoubleTranslation(kpt_thc.chi, kpt_thc.zeta, mf) + lambda_data = compute_lambda( + hcore_mo, + helper, + ) + assert np.isclose(lambda_data.lambda_total, 93.84613761765415)