[PBC Resources Estimates 3/4] Add computation of lambda

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.
+#   limitations under the License.

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

@@ -0,0 +1,106 @@
+# 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):
+    """Light container to store return values of compute_lambda function"""
+
+    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

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)

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)

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

@@ -0,0 +1,120 @@
+# 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):
+    """Light container to store return values of compute_lambda function"""
+
+    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)
+        del A
+        del B
+
+    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