[PBC Resources Estimates 1/4] Add k-point THC factorization

dev_tools/requirements/deps/resource_estimates.txt

@@ -1,3 +1,4 @@
 pyscf
 jax
 jaxlib
+ase

dev_tools/requirements/resource_estimates.env.txt

@@ -4,6 +4,8 @@
 #
 #    pip-compile --output-file=resource_estimates.env.txt deps/resource_estimates.txt pytest.env.txt
 #
+ase==3.22.1
+    # via -r deps/resource_estimates.txt
 attrs==23.1.0
     # via
     #   -r pytest.env.txt
@@ -83,6 +85,7 @@ kiwisolver==1.4.4
 matplotlib==3.7.1
     # via
     #   -r pytest.env.txt
+    #   ase
     #   cirq-core
 ml-dtypes==0.2.0
     # via
@@ -101,6 +104,7 @@ networkx==2.8.8
 numpy==1.23.5
     # via
     #   -r pytest.env.txt
+    #   ase
     #   cirq-core
     #   contourpy
     #   h5py
@@ -174,6 +178,7 @@ requests==2.31.0
 scipy==1.9.3
     # via
     #   -r pytest.env.txt
+    #   ase
     #   cirq-core
     #   jax
     #   jaxlib

src/openfermion/resource_estimates/pbc/__init__.py

@@ -0,0 +1,12 @@
+# 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.

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

@@ -0,0 +1,19 @@
+# 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 openfermion.resource_estimates import HAVE_DEPS_FOR_RESOURCE_ESTIMATES
+
+if HAVE_DEPS_FOR_RESOURCE_ESTIMATES:
+    from .hamiltonian import (build_hamiltonian,
+                              build_momentum_transfer_mapping,
+                              cholesky_from_df_ints)

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

@@ -0,0 +1,183 @@
+# 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, asdict
+from typing import Tuple
+import h5py
+import numpy as np
+import numpy.typing as npt
+
+from pyscf import lib
+from pyscf.ao2mo import _ao2mo
+from pyscf.lib import logger
+from pyscf.pbc.df import df
+from pyscf.pbc.lib.kpts_helper import gamma_point
+from pyscf.pbc.mp.kmp2 import _add_padding
+from pyscf.pbc import mp, scf, gto
+
+
+@dataclass
+class HamiltonianProperties:
+    """Lighweight descriptive data class to hold return values from
+    compute_lambda functions.
+
+    Attributes:
+        lambda_total: Total lambda value (norm) of Hamiltonian.
+        lambda_one_body: One-body lambda value (norm) of Hamiltonian.
+        lambda_two_body: Two-body lambda value (norm) of Hamiltonian.
+    """
+
+    lambda_total: float
+    lambda_one_body: float
+    lambda_two_body: float
+
+    dict = asdict
+
+
+def build_hamiltonian(mf: "scf.KRHF") -> Tuple[npt.NDArray, npt.NDArray]:
+    """Utility function to build one- and two-electron matrix elements from mean
+    field object.
+
+    Arguments:
+        mf: pyscf KRHF object.
+
+    Returns:
+        hcore_mo: one-body Hamiltonian in MO basis.
+        chol: 3-index RSGDF density fitted integrals.
+    """
+    # Build temporary mp2 object so MO coeffs can potentially be padded if mean
+    # field solution yields different number of MOs per k-point.
+    tmp_mp2 = mp.KMP2(mf)
+    mo_coeff_padded = _add_padding(tmp_mp2, tmp_mp2.mo_coeff,
+                                   tmp_mp2.mo_energy)[0]
+    hcore_mo = np.asarray([
+        C.conj().T @ hk @ C for (C, hk) in zip(mo_coeff_padded, mf.get_hcore())
+    ])
+    chol = cholesky_from_df_ints(tmp_mp2)
+    return hcore_mo, chol
+
+
+def cholesky_from_df_ints(mp2_inst, pad_mos_with_zeros=True) -> npt.NDArray:
+    """Compute 3-center electron repulsion integrals, i.e. (L|ov),
+    where `L` denotes DF auxiliary basis functions and `o` and `v` occupied and
+    virtual canonical crystalline orbitals. Note that `o` and `v` contain kpt
+    indices `ko` and `kv`, and the third kpt index `kL` is determined by
+    the conservation of momentum.
+
+    Note that if the number of mos differs at each k-point then this function
+    will pad MOs with zeros to ensure contiguity.
+
+    Args:
+        mp2_inst: pyscf KMP2 instance.
+
+    Returns:
+        Lchol: 3-center DF ints, with shape (nkpts, nkpts, naux, nmo, nmo)
+    """
+
+    log = logger.Logger(mp2_inst.stdout, mp2_inst.verbose)
+
+    if mp2_inst._scf.with_df._cderi is None:
+        mp2_inst._scf.with_df.build()
+
+    cell = mp2_inst._scf.cell
+    if cell.dimension == 2:
+        # 2D ERIs are not positive definite. The 3-index tensors are stored in
+        # two part. One corresponds to the positive part and one corresponds
+        # to the negative part. The negative part is not considered in the
+        # DF-driven CCSD implementation.
+        raise NotImplementedError
+
+    # nvir = nmo - nocc
+    nao = cell.nao_nr()
+
+    mo_coeff = mp2_inst._scf.mo_coeff
+    kpts = mp2_inst.kpts
+    if pad_mos_with_zeros:
+        mo_coeff = _add_padding(mp2_inst, mp2_inst.mo_coeff,
+                                mp2_inst.mo_energy)[0]
+        nmo = mp2_inst.nmo
+    else:
+        nmo = nao
+        num_mo_per_kpt = np.array([C.shape[-1] for C in mo_coeff])
+        if not (num_mo_per_kpt == nmo).all():
+            log.info("Number of MOs differs at each k-point or is not the same "
+                     "as the number of AOs.")
+    nkpts = len(kpts)
+    if gamma_point(kpts):
+        dtype = np.double
+    else:
+        dtype = np.complex128
+    dtype = np.result_type(dtype, *mo_coeff)
+    Lchol = np.empty((nkpts, nkpts), dtype=object)
+
+    cput0 = (logger.process_clock(), logger.perf_counter())
+
+    bra_start = 0
+    bra_end = nmo
+    ket_start = nmo
+    ket_end = 2 * nmo
+    with h5py.File(mp2_inst._scf.with_df._cderi, "r") as f:
+        kptij_lst = f["j3c-kptij"][:]
+        tao = []
+        ao_loc = None
+        for ki, kpti in enumerate(kpts):
+            for kj, kptj in enumerate(kpts):
+                kpti_kptj = np.array((kpti, kptj))
+                Lpq_ao = np.asarray(df._getitem(f, "j3c", kpti_kptj, kptij_lst))
+
+                mo = np.hstack((mo_coeff[ki], mo_coeff[kj]))
+                mo = np.asarray(mo, dtype=dtype, order="F")
+                if dtype == np.double:
+                    out = _ao2mo.nr_e2(
+                        Lpq_ao,
+                        mo,
+                        (bra_start, bra_end, ket_start, ket_end),
+                        aosym="s2",
+                    )
+                else:
+                    # Note: Lpq.shape[0] != naux if linear dependency is found
+                    # in auxbasis
+                    if Lpq_ao[0].size != nao**2:  # aosym = 's2'
+                        Lpq_ao = lib.unpack_tril(Lpq_ao).astype(np.complex128)
+                    out = _ao2mo.r_e2(
+                        Lpq_ao,
+                        mo,
+                        (bra_start, bra_end, ket_start, ket_end),
+                        tao,
+                        ao_loc,
+                    )
+                Lchol[ki, kj] = out.reshape(-1, nmo, nmo)
+
+    log.timer_debug1("transforming DF-AO integrals to MO", *cput0)
+
+    return Lchol
+
+
+def build_momentum_transfer_mapping(cell: gto.Cell,
+                                    kpoints: np.ndarray) -> np.ndarray:
+    # Define mapping momentum_transfer_map[Q][k1] = k2 that satisfies
+    # k1 - k2 + G = Q.
+    a = cell.lattice_vectors() / (2 * np.pi)
+    delta_k1_k2_Q = (kpoints[:, None, None, :] - kpoints[None, :, None, :] -
+                     kpoints[None, None, :, :])
+    delta_k1_k2_Q += kpoints[0][None, None, None, :]  # shift to center
+    delta_dot_a = np.einsum("wx,kpQx->kpQw", a, delta_k1_k2_Q)
+    int_delta_dot_a = np.rint(delta_dot_a)
+    # Should be zero if transfer is statisfied (2*pi*n)
+    mapping = np.where(
+        np.sum(np.abs(delta_dot_a - int_delta_dot_a), axis=3) < 1e-10)
+    num_kpoints = len(kpoints)
+    momentum_transfer_map = np.zeros((num_kpoints,) * 2, dtype=np.int32)
+    # Note index flip due to Q being first index in map but broadcasted last..
+    momentum_transfer_map[mapping[1], mapping[0]] = mapping[2]
+
+    return momentum_transfer_map

src/openfermion/resource_estimates/pbc/hamiltonian/hamiltonian_test.py

@@ -0,0 +1,90 @@
+# 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.
+import itertools
+
+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 cc, mp
+
+    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_build_hamiltonian():
+    mf = make_diamond_113_szv()
+    nmo = mf.mo_coeff[0].shape[-1]
+    naux = 108
+    hcore, chol = build_hamiltonian(mf)
+    nkpts = len(mf.mo_coeff)
+    assert hcore.shape == (nkpts, nmo, nmo)
+    assert chol.shape == (nkpts, nkpts)
+    assert chol[0, 0].shape == (naux, nmo, nmo)
+
+
+@pytest.mark.skipif(not HAVE_DEPS_FOR_RESOURCE_ESTIMATES,
+                    reason='pyscf and/or jax not installed.')
+def test_pyscf_chol_from_df():
+    mf = make_diamond_113_szv()
+    mymp = mp.KMP2(mf)
+    nmo = mymp.nmo
+    nocc = mymp.nocc
+    nvir = nmo - nocc
+    Luv = cholesky_from_df_ints(mymp)
+
+    # 1. Test that the DF integrals give the correct SCF energy (oo block)
+    mf.exxdiv = None  # exclude ewald exchange correction
+    Eref = mf.energy_elec()[1]
+    Eout = 0.0j
+    nkpts = len(mf.mo_coeff)
+    for ik, jk in itertools.product(range(nkpts), repeat=2):
+        Lii = Luv[ik, ik][:, :nocc, :nocc]
+        Ljj = Luv[jk, jk][:, :nocc, :nocc]
+        Lij = Luv[ik, jk][:, :nocc, :nocc]
+        Lji = Luv[jk, ik][:, :nocc, :nocc]
+        oooo_d = np.einsum("Lij,Lkl->ijkl", Lii, Ljj) / nkpts
+        oooo_x = np.einsum("Lij,Lkl->ijkl", Lij, Lji) / nkpts
+        Eout += 2.0 * np.einsum("iijj->", oooo_d)
+        Eout -= np.einsum("ijji->", oooo_x)
+    assert abs(Eout / nkpts - Eref) < 1e-12
+
+    # 2. Test that the DF integrals agree with those from MP2 (ov block)
+    from pyscf.pbc.mp.kmp2 import _init_mp_df_eris
+
+    Ltest = _init_mp_df_eris(mymp)
+    for ik, jk in itertools.product(range(nkpts), repeat=2):
+        assert np.allclose(Luv[ik, jk][:, :nocc, nocc:],
+                           Ltest[ik, jk],
+                           atol=1e-12)
+
+    # 3. Test that the DF integrals have correct vvvv block (vv)
+    Ivvvv = np.zeros((nkpts, nkpts, nkpts, nvir, nvir, nvir, nvir),
+                     dtype=np.complex128)
+    for ik, jk, kk in itertools.product(range(nkpts), repeat=3):
+        lk = mymp.khelper.kconserv[ik, jk, kk]
+        Lij = Luv[ik, jk][:, nocc:, nocc:]
+        Lkl = Luv[kk, lk][:, nocc:, nocc:]
+        Imo = np.einsum("Lij,Lkl->ijkl", Lij, Lkl)
+        Ivvvv[ik, jk, kk] = Imo / nkpts
+
+    mycc = cc.KRCCSD(mf)
+    eris = mycc.ao2mo()
+    assert np.allclose(eris.vvvv,
+                       Ivvvv.transpose(0, 2, 1, 3, 5, 4, 6),
+                       atol=1e-12)